AWS walkthrough: training Qwen3 on Trainium with torchtitan
An AWS Machine Learning Blog post walks through training Qwen3 on a single trn2.48xlarge Trainium instance using torchtitan, saying TorchNeuron registers Trainium as a native PyTorch device so torchtitan can use it through standard PyTorch interfaces. The post describes three configurations that share the same torchrun launch pattern and config structure: Hugging Face Transformers v5 as a modeling backend running Hub models through AutoModelForCausalLM with published safetensors, torchtitan's native Qwen3 selected by changing config.model_spec, and an NKI mixture-of-experts path using kernels from the torchtitan-neuron recipes attached via torchtitan's override registry. AWS says the setups support both pre-training from random initialization and fine-tuning published weights, and pins versions including torch 2.14.0, torch_neuronx 2.14.3, transformers 5.16.0 and torchtitan-neuron release-2.33.
The post documents Trainium as a supported torchtitan device across all three configurations, giving PyTorch users a vendor-described path to run Qwen3 pre-training or fine-tuning on Trainium without leaving standard PyTorch interfaces.
Using torchtitan on AWS Trainium means using it the way you always have. TorchNeuron registers Trainium as a native PyTorch device, which lets torchtitan use it through standard PyTorch interfaces. You can use these configurations to pre-train from random initialization or fine-tune published weights. You can also customize the model implementation and kernels for your workload. In this post, you train Qwen3 three ways on a single trn2.48xlarge.
The first method uses Transformers v5 as a modeling backend in torchtitan. It runs Hub models through AutoModelForCausalLM and reads and writes published safetensors directly. The second method switches to torchtitan’s native Qwen3 by changing config.model_spec. The third adds Neuron Kernel Interface (NKI) mixture-of-experts (MoE) kernels from the torchtitan-neuron recipes and then uses torchtitan’s override registry to attach your own kernel. All three configurations use the same torchrun launch pattern and config structure.
Prerequisites
To follow this walkthrough, use a single trn2.48xlarge with either the AWS Deep Learning AMI for Neuron or a Neuron PyTorch training container. Both environments provide the Neuron driver, runtime, and compiler required to train on Trainium.
#AWS#Trainium#trn2.48xlarge#Qwen3#torchtitan#torchtitan-neuron#PyTorch#TorchNeuron#Hugging Face Transformers#Neuron Kernel Interface (NKI)#Qwen3-4B-Instruct-2507