TorchNeuron: running HF models on AWS Trainium is easier than ever

Summary

TorchNeuron is a new PyTorch plugin that adds a native Neuron backend for AWS Trainium, supporting eager mode, torch.compile, and distributed APIs like FSDP and TP. It will be available as a public beta starting October 20, 2026, and can be installed with pip install torch-neuronx. Hugging Face libraries such as Transformers, Diffusers, Accelerate, PEFT, and TRL now run directly on Trainium without needing the Optimum Neuron wrapper, using the same classes and scripts as on any other accelerator. For text generation, 89% of the top 100 most-used model architectures are supported out of the box, compared to 46% with Optimum Neuron v0.4.5. Training uses standard TRL scripts with only a few Neuron-specific settings like static shapes and torch.compile, and the resulting checkpoints are hardware-agnostic.

Takeaways

1. 1

   TorchNeuron integrates as a native PyTorch device, supporting eager mode, torch.compile, and distributed primitives like FSDP and TP, eliminating the need for the Optimum Neuron wrapper.

2. 2

   With TorchNeuron, 89% of the top 100 most-used text-generation model architectures are supported out of the box, up from 46% with Optimum Neuron v0.4.5.

3. 3

   Training on Trainium uses standard Hugging Face libraries (Transformers, Accelerate, PEFT, TRL) with only a few Neuron-specific settings such as static shapes and torch.compile, and produces hardware-agnostic checkpoints.

FAQ