Three ways to train Qwen3 on Trainium with torchtitan

Training Qwen3 on Trainium now supports three distinct torchtitan setups that all use the standard PyTorch device registration provided by TorchNeuron. The change lets engineers pick a high‑level Transformers v5 backend, switch to torchtitan’s native Qwen3 model spec, or attach custom mixture‑of‑experts kernels, all on a single trn2.48xlarge instance.

Prerequisites and environment

All three paths require a trn2.48xlarge with either the AWS Deep Learning AMI for Neuron or a Neuron‑enabled PyTorch training container. The environment must already contain the Neuron driver, runtime, and compiler. Verify the torchtitan installation with a simple version check:

python3 -c "import torchtitan, torchtitan_neuron; print(torchtitan.__version__, torchtitan_neuron.__file__)"

Pinning the component versions listed in the source ensures reproducibility: torch 2.14.0, torch_neuronx 2.14.3, transformers 5.16.0, and torchtitan‑neuron release‑2.33 (0.2.2+neuron.1). Install torchtitan‑neuron without pulling new dependencies to keep the Neuron‑specific torch build:

git clone https://github.com/aws-neuron/torchtitan-neuron
cd torchtitan-neuron
pip install -r requirements.txt --no-deps

Download the Qwen3 model artifacts from Hugging Face Hub into a local directory that the training config can reference:

hf download Qwen/Qwen3-4B-Instruct-2507 \
  --local-dir "$PWD/assets/hf/Qwen3-4B-Instruct-2507"

Configuration paths for Qwen3 on Trainium

Each approach builds a Python config function that feeds torchtitan’s train module. The base config is created with transformers_modeling_backend_full() and then mutated for the specific scenario.

1. Transformers v5 backend – The model is loaded via AutoModelForCausalLM using the Hub ID. No explicit model spec is set, so the Hub definition supplies the architecture.

import os
from torchtitan.experiments.transformers_modeling_backend.config_registry import (
    transformers_modeling_backend_full,
)
from torchtitan.experiments.transformers_modeling_backend.configs import (
    TransformersBackendConfig,
)
ASSETS = os.path.abspath("./assets/hf/Qwen3-4B-Instruct-2507")

def v5_qwen3_4b() -> TransformersBackendConfig:
    config = transformers_modeling_backend_full()
    config.hf_model = "Qwen/Qwen3-4B-Instruct-2507"
    config.hf_assets_path = ASSETS
    config.dataloader.dataset = "c4_test"
    config.training.steps = 10
    config.training.local_batch_size = 1
    config.parallelism.data_parallel_shard_degree = -1
    return config

Launch with the standard torchrun pattern, which creates 64 processes – one per logical NeuronCore – and shards the model across them:

torchrun --nproc_per_node=64 --rdzv_backend c10d --rdzv_endpoint localhost:0 \
  --local-ranks-filter 0 --role rank --tee 3 \
  -m torchtitan.train --module neuron_qwen3 --config v5_qwen3_4b

2. Native Qwen3 spec – By setting config.model_spec to torchtitan’s built‑in Qwen3 definition, the training loop bypasses the Hugging Face backend and uses the optimized native implementation.

3. Custom NKI MoE kernels – The third variant adds Neuron Kernel Interface mixture‑of‑experts kernels from the torchtitan‑neuron recipes and registers them via torchtitan’s override mechanism. This demonstrates how to inject workload‑specific kernels without altering the overall launch command.

Operational implications

All three configurations share the same launch pattern, meaning operational tooling (e.g., job schedulers, monitoring agents) can treat them uniformly. Because TorchNeuron registers the neuron device at import time, any PyTorch code that respects device arguments will automatically target Trainium. In eager mode, each operation compiles on first use; with torch.compile, the entire graph is compiled ahead of execution, offering a performance knob that practitioners can experiment with.

Running 64 logical cores on a single instance simplifies networking – there is no inter‑node traffic – but it also requires careful resource pinning to avoid oversubscription. The data‑parallel shard degree of -1 tells torchtitan to spread the model evenly across all ranks, which is the default for a single‑node, full‑core deployment.

Related CloudNinjas coverage: AWS.

What This Means For Practitioners

• Engineers can adopt Trainium without rewriting existing PyTorch pipelines; the device registration is transparent.
• Choosing between a high‑level Transformers backend and a native spec lets teams balance ease of use against potential performance gains.
• Custom kernel overrides are supported out‑of‑the‑box, enabling workload‑specific optimizations without changing the training script.
• Version pinning is critical – mismatched torchtitan, transformers, or torch_neuronx releases can break the compile path.
• Future evaluations should track torchtitan releases, new NKI kernel recipes, and scaling beyond a single trn2.48xlarge to multi‑node Trainium clusters.