Training Qwen3 on AWS Trainium with torchtitan – a hands‑on look

According to the AWS Machine Learning blog, the team showed three ways to train Qwen3 on a single trn2.48xlarge using torchtitan. The post matters because it demonstrates that large‑language‑model (LLM) pre‑training can be done on AWS's custom Trainium chips with the same PyTorch code you already know.

Three configurations on a single Trainium instance

The walkthrough runs on trn2.48xlarge – 16 Trainium2 accelerators, 128 physical NeuronCores, exposed as 64 logical cores (LNC=2). The software stack is pinned to:

All three runs use the same launch pattern (torchrun with 64 processes) and a tiny training loop of 10 steps, local batch size 1, global batch size 64. The only differences are the model backend and the kernel set.

1. Transformers v5 backend

• Uses AutoModelForCausalLM to pull the Qwen3‑4B‑Instruct model from Hugging Face.
• Reads and writes standard safetensors checkpoints.
• Loss after the first step is 6.54 when loading the published weights, 12.43 from a random start.
• Minimal code change – just point config.hf_model at the Hub ID.

2. torchtitan native Qwen3Model

• Switches config.model_spec to model_registry("4B", attn_backend="sdpa").
• Bypasses the Hugging Face transformer library; torchtitan provides its own Qwen3 implementation.
• Keeps the same loss numbers because the same weights are loaded via the torchtitan checkpoint adapter.
• Requires a different config builder (qwen3_32b() then replace spec) but otherwise identical launch.

3. NKI mixture‑of‑experts (MoE) kernels

• Adds Neuron Kernel Interface (NKI) kernels via qwen3_nki_moe_config.
• Sets num_experts=64, top_k=4, and reduces layers to 4 so the 30‑B‑A3B variant fits on one instance.
• Runs in eager mode (CompileConfig(enable=False)) for simplicity; can be switched to torch.compile later.
• Demonstrates how to register a custom NKI operator (the example rewrites rotary‑position embeddings).

How torchtitan talks to Trainium

TorchNeuron registers Trainium as a native PyTorch device (neuron). When torch_neuronx is installed, torchtitan detects the device at start‑up and routes all tensor operations through the Neuron compiler. Two execution modes exist:

• Eager – each operation is compiled the first time it runs. This adds a one‑time latency but gives immediate feedback.
• Graph compile – torch.compile captures a full forward‑backward graph and compiles it in one go, reducing per‑step overhead at the cost of a larger initial compile time.

Both modes work with the same torchrun launch, meaning you can swap them without changing your training script.

Trade‑offs nobody spells out

Performance vs. flexibility

*Throughput is not reported in the source; the table reflects the qualitative impact of each choice.

• Cost – All three runs stay on a single trn2.48xlarge, so hourly price is the same. The only extra cost is developer time for custom kernels.
• Scalability – The launch uses data‑parallel sharding across 64 logical cores. Moving to multiple instances would require a more involved torchrun mesh and careful expert‑parallel handling for MoE.
• Reliability – The first two configurations produce standard Hugging Face checkpoints that load without missing keys. The NKI example is a proof‑of‑concept; it registers operators but does not guarantee a performance win.

What actually changes for a typical data‑science team

The headline here is you can keep your existing PyTorch code and run it on Trainium with no rewrites, but you gain two optional levers:

1. Native torchtitan models give you a path to tighter hardware integration without abandoning the PyTorch API.
2. NKI kernels let you replace hot‑spot operations with custom code that runs directly on the Neuron compiler, which can matter for very large MoE models.

In practice this means most teams will start with the Transformers v5 backend, verify loss numbers, then decide whether the marginal speed boost of native kernels justifies the extra config work. Only groups that need MoE scaling should invest in NKI development.

Quick start you can run today

1. Launch a trn2.48xlarge using the AWS Deep Learning AMI for Neuron or the Neuron PyTorch container.
2. Clone the torchtitan‑neuron repo and pin the versions listed in the table.
3. Install without pulling newer torch builds:

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

4. Download the Qwen3‑4B‑Instruct assets:

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

5. Create neuron_qwen3.py with the v5_qwen3_4b() config from the blog (just copy‑paste).
6. Run the training command:

   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
   

7. Watch the first‑step loss. If you see 6.54, the weights loaded correctly; if it’s higher, double‑check the asset path.

From there you can experiment with the native or NKI configs by swapping the config function name. The same launch script works, so the barrier to trying the next level is just a few lines of Python.

Sources

• Train Qwen3 on AWS Trainium with torchtitan