Training Qwen3 on AWS Trainium with torchtitan and custom kernels

AWS has released a detailed technical walkthrough demonstrating how to train the Qwen3 large language model on its Trainium2 accelerators using the torchtitan framework. Published in October 2026, the guide outlines three distinct configuration paths for engineers looking to leverage native PyTorch interfaces on specialized hardware. The documentation serves as a practical reference for machine learning teams aiming to optimize pre-training and fine-tuning workflows on Amazon Web Services infrastructure.

What happened

The article provides a step-by-step implementation guide for running Qwen3 models on a single trn2.48xlarge instance, which contains 16 Trainium2 accelerators and 128 physical NeuronCores. The authors demonstrate that developers can use torchtitan on Trainium much like they would on standard GPU clusters, thanks to TorchNeuron registering Trainium as a native PyTorch device. This allows the framework to detect the hardware at startup and compile operations either eagerly or via torch.compile.

The walkthrough covers three specific training configurations. The first method utilizes Hugging Face Transformers v5 as the modeling backend, allowing users to load models directly from the Hub using AutoModelForCausalLM. The second approach switches to torchtitan’s native Qwen3 implementation, offering more control over the model architecture and sharding strategies. The third method introduces advanced customization by integrating Neuron Kernel Interface (NKI) mixture-of-experts kernels, enabling developers to attach custom operators for performance-critical layers.

Each configuration uses the same torchrun launch pattern, ensuring consistency across different experimental setups. The guide emphasizes reproducibility by specifying exact software versions, including PyTorch 2.14.0, Transformers 5.16.0, and torchtitan-neuron release-2.33. It also details how to manage checkpoints, verifying that weights saved in Hugging Face safetensors format can be loaded back without missing keys, confirming the integrity of the training pipeline.

How it works

The core mechanism relies on the integration between torchtitan and the AWS Neuron stack. When torch_neuronx is installed, it exposes Trainium as a neuron device. Torchtitan then uses standard PyTorch distributed training primitives, such as Fully Sharded Data Parallel (FSDP) and tensor parallelism, to distribute the workload across the 64 logical NeuronCores available on the instance. The framework manages the sharding of model parameters, gradients, and optimizer states, while sequence parallelism handles activation sharding along the sequence dimension.

For custom kernel integration, the guide demonstrates how to use torchtitan’s override registry. Developers can write kernels in NKI, the Neuron kernel language, and wrap them as PyTorch operators using wrap_nki and nki_op. These operators are then registered to replace specific modules within the model graph. For example, a custom rotary embedding kernel can replace the default PyTorch implementation. The system validates these overrides at startup, ensuring that each replacement targets the correct fully qualified names within the model structure.

The training loop reports key metrics such as loss, gradient norm, high-bandwidth memory usage, and model FLOPs utilization. This visibility allows engineers to monitor resource consumption and adjust parallelism settings dynamically. The compilation process occurs during the first training step, which takes longer than subsequent steps, but enables optimized execution for the remainder of the run.

Key details

• The setup requires a trn2.48xlarge instance with 64 logical NeuronCores configured at LNC=2.
• Software dependencies include PyTorch 2.14.0, Transformers 5.16.0, and torchtitan-neuron 0.2.2+neuron.1.
• Hugging Face checkpoint loading verifies correctness by checking for a step 1 loss near 6.54 for the Qwen3-4B-Instruct-2507 model.
• Custom NKI kernels must adhere to specific sequence length constraints, such as being a multiple of 128 * LNC.
• Tensor parallelism shards layer weights across ranks, while data parallelism shards parameters and optimizer states.
• The override registry prevents conflicts by stopping if two overrides claim the same node and warning if no nodes match.

Why it matters

For software engineers and ML practitioners, this guide lowers the barrier to entry for training large models on non-NVIDIA hardware. By maintaining compatibility with standard PyTorch interfaces and Hugging Face ecosystems, it reduces the need for extensive code refactoring. Teams can leverage existing model repositories and checkpoint formats, minimizing the operational overhead associated with adopting new accelerator architectures.

The ability to inject custom NKI kernels is particularly significant for performance optimization. As models grow in complexity, generic implementations may not fully utilize hardware capabilities. Allowing developers to replace specific operations, such as attention mechanisms or rotary embeddings, with hand-tuned kernels enables finer control over computational efficiency. This flexibility is crucial for achieving competitive throughput and cost-effectiveness in large-scale training jobs.

Furthermore, the detailed verification steps provide confidence in the reliability of the training process. Confirming that exported checkpoints load correctly ensures that models trained on Trainium can be deployed seamlessly in production environments that rely on standard inference stacks. This end-to-end validation supports robust MLOps practices and facilitates collaboration across teams using different hardware backends.

What you can do

• Set up a trn2.48xlarge instance with the AWS Deep Learning AMI for Neuron or a Neuron PyTorch container.
• Install the specified versions of torchtitan-neuron and dependencies using the provided requirements.txt.
• Download the Qwen3 model artifacts from Hugging Face Hub to a local directory for training.
• Configure the training script to use data parallel sharding across all 64 ranks by setting data_parallel_shard_degree = -1.
• Verify custom kernel outputs against baseline implementations using float32 inputs to ensure numerical accuracy.
• Monitor per-rank high-bandwidth memory usage to adjust batch sizes and sequence lengths appropriately.