TorchNeuron lets PyTorch code run on AWS Trainium unmodified

AWS and the PyTorch team released TorchNeuron today, an open-source native PyTorch backend for AWS Trainium chips, timed to go live as a public beta at the PyTorch Conference in San Jose. It registers Trainium through PyTorch’s own PrivateUse1 extension point, the same mechanism used for other accelerators, rather than requiring Neuron-specific code.

The change needed is small: swap a model’s device to 'neuron' for eager execution, or set the compile backend once for production. Standard APIs, torch.nn, torch.distributed, DDP, FSDP2, tensor and pipeline parallelism, run as written on Trainium2 and Trainium3. Hugging Face says Transformers v5, Diffusers and trl now treat Trainium as a native device too. It’s Apache 2.0, installed with pip install torch-neuronx, and the code lives at github.com/aws-neuron/torch-neuronx.

model = MyModel().to('neuron')  # the only change

Why it matters: the usual cost of trying alternative AI chips is rewriting model code against a vendor-specific stack. TorchNeuron removes that step for Trainium, so teams can benchmark it against CUDA GPUs using the PyTorch code they already have.

The caveat: this is a public beta. AWS’s own upstream test run skipped roughly 22,000 PyTorch tests for operators without a Neuron lowering yet, and dynamic or symbolic input shapes aren’t supported in this release.