From the source
Amazon's walkthrough shows how to train Qwen3 on a single AWS Trainium trn2.48xlarge instance with torchtitan, using TorchNeuron to register Trainium as a native PyTorch device.
- It covers three configurations that share the same torchrun launch pattern and config structure: Transformers v5 as a modeling backend, torchtitan's native Qwen3 model, and NKI mixture-of-experts kernels attached through torchtitan's override registry.
- The guide says the setup supports pre-training from random initialization or fine-tuning published weights, and that the same launch pattern works across all three configurations.
- For Qwen/Qwen3-4B-Instruct-2507, the verified run produced a step 1 loss of 6.54 from published weights and 12.43 from random initialization, which the post offers as a quick check that the weights loaded.
- Exported checkpoints load back with AutoModelForCausalLM without missing or unexpected keys other than lm_head.weight, which Qwen3-4B ties to the token embeddings.
- The walkthrough pins specific versions because the Transformers v5 integration and Trainium-specific override modules can change between torchtitan releases.
