AWS Shows Qwen3 Training on Trainium via torchtitan
A recipe for fine-tuning Qwen3 on Trainium2 without rewriting training code
The facts
- Walkthrough uses a single trn2.48xlarge instance with 16 Trainium2 accelerators and 64 logical NeuronCores
- Three training paths shown: Transformers v5 AutoModelForCausalLM, torchtitan's native Qwen3Model, and custom NKI mixture-of-experts kernels
- Fine-tuning Qwen3-4B-Instruct-2507 from published weights gave a step-1 loss of 6.54 versus 12.43 from random initialization, per AWS
- Exported checkpoints reload as standard Hugging Face safetensors with no missing or unexpected keys besides the tied lm_head.weight
- Pinned versions: torch/torch_neuronx 2.14.0/2.14.3, transformers 5.16.0, torchtitan-neuron release-2.33
AWS has published a walkthrough for training Qwen3 models on its Trainium2 chips rather than Nvidia GPUs. It uses torchtitan, PyTorch's native training library. The guide runs on a single trn2.48xlarge instance, with 16 Trainium2 accelerators and 64 logical NeuronCores.
It shows three ways to train the model. One is through Hugging Face Transformers v5's AutoModelForCausalLM. Another is through torchtitan's own native Qwen3Model. A third swaps in custom mixture-of-experts kernels written in Neuron Kernel Interface, or NKI, AWS's kernel language for Trainium. All three use the same torchrun launch command and config structure, which AWS says means no workflow changes are needed to move onto its silicon.
The post pins specific versions for reproducibility: torch and torch_neuronx 2.14.0/2.14.3, transformers 5.16.0, and torchtitan-neuron release-2.33. For fine-tuning from published weights, AWS reports a first-step training loss of 6.54 when loading Qwen/Qwen3-4B-Instruct-2507 from Hugging Face. That compares with 12.43 from random initialisation, a sanity check that the published checkpoint loaded correctly. The exported checkpoint then round-trips back into a standard Hugging Face format. There are no missing or unexpected keys, other than the tied lm_head.weight.
The MoE example runs Qwen3-30B-A3B trimmed to four layers and 64 experts, routing each token to four experts, so the full model fits on one instance. AWS is explicit that this demonstrates kernel registration and autograd wiring rather than a measured speed gain.
The final section walks through attaching a custom NKI kernel for rotary position embeddings, using the rope_hf kernel from nkilib. This includes handling the tensor-layout mismatch between torchtitan's storage order and what the kernel expects, and registering a backward pass via register_autograd. AWS frames this as the extensibility case: teams can keep the standard Hub or torchtitan model code but drop in their own optimised kernel for a specific operation.
Why it matters
This gives teams already using torchtitan a documented, drop-in path to Trainium2 without rewriting training code. That matters as Nvidia GPU capacity stays tight and priced accordingly. The custom-kernel example signals AWS wants Trainium to host serious kernel-level optimisation, not just serve as a cheaper commodity substitute. Anyone evaluating Trainium now has a concrete cost comparison to run, not just a marketing claim.
Primary source
AWS Machine Learning — aws.amazon.comWritten from the primary source above by The Neural Net’s automated desk (claude-sonnet-5, theneuralnet/pipeline@1) without human review. How this works. Corrections: check the source first — if we got it wrong, the source link is the record.