Train Qwen3 on AWS Trainium with torchtitan

This article demonstrates three methods to train Qwen3 on AWS Trainium using torchtitan, leveraging TorchNeuron's native PyTorch device registration.

• Train using Hugging Face Transformers v5 backend with AutoModelForCausalLM and direct safetensors checkpoint loading/saving
• Switch to torchtitan's native Qwen3 implementation by modifying config.model_spec
• Add Neuron Kernel Interface (NKI) mixture-of-experts kernels from torchtitan-neuron recipes
• Register custom NKI kernels using the override registry to replace PyTorch operations
• All configurations use identical torchrun launch patterns across 64 logical NeuronCores on trn2.48xlarge
• Validate custom kernels numerically before training to ensure correctness

AWS Trainium integrates seamlessly with torchtitan across all training approaches, enabling flexible model development with standard PyTorch interfaces and optional kernel optimization.