Amazon Web Services (AWS) has published a technical guide detailing how to train Qwen3 models on its Trainium hardware using torchtitan, a library designed for large-scale model training. The documentation illustrates how developers can leverage TorchNeuron to integrate Trainium accelerators with standard PyTorch interfaces, enabling both pre-training and fine-tuning workloads.
Training Configurations
The guide outlines three distinct methods for training Qwen3 on a trn2.48xlarge instance, which features 16 Trainium2 accelerators and 128 physical NeuronCores:
- Transformers v5 Integration: This approach uses the Transformers library as the modeling backend. It enables the direct use of Hugging Face Hub models through
AutoModelForCausalLM, allowing for the direct reading and writing of safetensors. - Native torchtitan Models: By adjusting the model configuration, users can utilize the native Qwen3 implementation provided within
torchtitan, offering an alternative to standard Hugging Face workflows. - Mixture-of-Experts (MoE) with NKI Kernels: This configuration incorporates Neuron Kernel Interface (NKI) recipes to handle mixture-of-experts layers. Developers can use the
torchtitanoverride registry to attach custom NKI operators, providing flexibility for specialized workload optimization.
Technical Implementation
For these configurations, AWS utilizes torchrun to manage 64 training processes across the logical NeuronCores. The implementation supports both eager mode and torch.compile, with TorchNeuron handling the registration of Trainium as a native PyTorch device. The guide also details how to create and register custom NKI kernels, such as a specialized rotary embedding kernel, to optimize performance for specific tensor shapes and sequence lengths.
AWS notes that these configurations maintain consistent launch patterns and structural requirements, allowing users to move between different training approaches without altering their underlying infrastructure setup. The documentation emphasizes that the override registry allows for granular control over model components, enabling developers to replace standard PyTorch operations with high-performance NKI-based kernels while maintaining autograd compatibility.
As of the publication, these methods rely on specific versions of torch, transformers, and the torchtitan-neuron package to ensure compatibility with the Neuron driver, runtime, and compiler environments.
Source: aws.amazon.com
Original article: https://aws.amazon.com/blogs/machine-learning/train-qwen3-on-aws-trainium-with-torchtitan/