Train Qwen3 on AWS Trainium with torchtitan

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.

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.

Start with transformers_modeling_backend_full() . It leaves the fields that map to Hugging Face configuration keys unset, so AutoConfig.from_pretrained() loads the model architecture from the Hub.

To fine-tune from published weights, enable checkpoint loading and point initial_load_path to ASSETS . The state-dictionary adapter adjusts the leading model. prefix and handles tied embeddings. It saves the final checkpoint as Hugging Face safetensors.

Use the first training-step loss as a quick check that torchtitan loaded the published weights. 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.

Figure 1: The published checkpoint starts with a loss of 6.54, and the exported weights load without missing or unexpected keys

Before you install torchtitan-neuron , check whether your environment already includes it. The following command prints the installed version and package path:

This walkthrough uses the versions in the following table. To reproduce the examples, pin these versions. The Transformers v5 integration and Trainium-specific override modules can change between torchtitan releases.

When torch_neuronx is installed, TorchNeuron registers Trainium as the neuron device and torchtitan detects it at startup. In eager mode, Neuron compiles operations as they execute. With torch.compile , Neuron compiles the captured graph. This same launch pattern works across all three configurations.