Your PyTorch Code on Trainium: What Actually Breaks

AWS Trainium chips offer lower costs per training run compared to equivalent GPU instances, but due to historical limitations, rewriting training loops around the Neuron SDK used to be necessary. TorchNeuron simplifies this process by registering a native PyTorch backend, requiring only a few compile steps for most code porting needs. However, the catch lies in the fact that no modifications mean only the forward pass is supported, leaving the rest of the training job on CPU.

This article provides a step-by-step guide to compiling a model for Trainium, identifying which operations may fall back to CPU, maintaining a smooth training loop, and debugging failures encountered during XLA graph compilation.

To get started with Trainium, ensure that the training workload runs on EC2 trn1 instances, and note that the software stack is version-sensitive. Install the AWS Neuron PyTorch wheel rather than PyTorch from PyPI, and refer to the Neuron setup documentation for the matching driver and runtime versions before initiating the process.

Install the torch-neuronx package within a fresh virtual environment, as mixing a stock PyTorch build with torch-neuronx can lead to confusing import errors that appear as model issues but are actually dependency conflicts.

Compiling the model involves importing torch and torch_neuronx, defining your model, and evaluating it with a sample input. The traced model is then saved for later use. This compilation process freezes the shape of the input, and any subsequent calls with different batch sizes or sequence lengths will trigger recompilation at runtime, leading to a drop in throughput until the graph is rebuilt. To prevent this, ensure that input shapes are fixed, and if working with transformers, bucket sequence lengths accordingly.

Operations not supported by the Neuron compiler include dynamic control flow, sparse tensors, certain fused optimizers, and unusual einsum patterns. When a compile fails, the error will specify the unsupported operation directly, such as a particular aten::index_put_ variant that the compiler cannot lower. Three practical responses to this issue include replacing the operation with a supported alternative, moving the operation outside the traced region to run on CPU between compiled stages, or checking the Neuron compiler release notes for potential fixes.

As newer releases improve support for more operations, some previously unsupported operations may now compile successfully.