Intel / article

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

Source: huggingface.co
Date: Published Oct 19, 2026

Context

Hugging Face's point is that its own libraries can now treat AWS Trainium as an ordinary PyTorch device, so the separate Trainium layer it used to maintain is no longer the recommended path (Hugging Face). Trainium and Inferentia are Amazon's chips for training and for . Until now, Hugging Face models reached them through Optimum Neuron, which relied on a PyTorch XLA-based package. That meant working the way a compiler does: operations were recorded lazily into a graph instead of running as the code reached them, input shapes had to stay fixed, and every graph break forced a recompilation. Optimum Neuron hid this by maintaining its own modeling classes and trainers.

TorchNeuron is a PyTorch plugin that exposes the chip as a native device with eager execution, meaning each operation runs as it is called, plus PyTorch's own distributed tools and torch.compile. Transformers, Diffusers, Accelerate, PEFT and TRL now run on it with the same classes and scripts used on other accelerators, and Hugging Face recommends using them directly instead of Optimum Neuron. The post says 89% of the 100 most-used text-generation architectures work out of the box, against 46% with Optimum Neuron v0.4.5. A standard TRL script ran to completion on 8 Trainium2 chips and saved an ordinary Transformers checkpoint that loads on any hardware.

Fixed shapes still matter. The examples use a static cache for generation to avoid frequent recompilation, TRL pads Neuron batches to a multiple of 128 , automatic dynamic shapes are switched off because the backend cannot compile them yet, and Transformers' asynchronous weight loading is not supported. TorchNeuron becomes a public beta on October 20, 2026.

Key takeaways · AI-distilled

TorchNeuron is a PyTorch plugin that exposes AWS Trainium as a native neuron device with eager execution, torch.compile, FSDP2 and DTensor, replacing a PyTorch XLA path built on lazy tensors, static shapes and recompiles at graph breaks.
Hugging Face now recommends Neuron users run Transformers, Diffusers, Accelerate, PEFT and TRL directly instead of Optimum Neuron, and says 89% of the top 100 text-generation architectures work out of the box versus 46% with Optimum Neuron v0.4.5.
A standard TRL SFTTrainer script fine-tuned Qwen3-8B on 8 Trainium2 chips with a few Neuron settings (automatic dynamic shapes off, async weight loading off, a TP and FSDP DistributedConfig), and saved a checkpoint that loads on any hardware.
Static shapes still matter: the post recommends a static cache for generation to avoid frequent recompilations, and TRL 1.16 and later pads Neuron batches to multiples of 128 tokens so compiled graphs are reused.

Terms in this piece · Glossary

inference — Running a trained model to get answers — the phase where AI is actually used, as opposed to trained.
fine-tuning — Taking a trained model and training it a bit more on your own examples so it gets better at one specific job.
token — The chunk of text a model reads and writes in — roughly three-quarters of a word — and the unit AI usage is billed in.

Read the source huggingface.co