Hugging Face adds native PyTorch support for AWS Trainium chips

Hugging Face has introduced TorchNeuron, a new native backend that allows PyTorch models to run directly on AWS Trainium hardware. Released as a public beta on October 20, 2026, this update removes the need for specialized compilation workflows, letting developers use standard PyTorch APIs for training and inference on Amazon’s AI accelerators.

What happened

For several years, running Hugging Face models on AWS Trainium and Inferentia chips required using Optimum Neuron. This library acted as a bridge between PyTorch and the hardware, relying on the torch-neuronx package which was built on PyTorch XLA. While functional, this approach forced developers to adapt to a compiler-centric mindset. They had to deal with lazy tensors instead of immediate execution, restrict themselves to static shapes, and manage frequent recompilations whenever the computation graph changed. Optimum Neuron maintained custom code, including specific modeling classes and trainers, to hide this complexity from users.

TorchNeuron changes this architecture by integrating directly into PyTorch itself. It exposes neuron as a native device, similar to how developers interact with CUDA GPUs or CPU devices. This shift enables eager execution, meaning operations are performed immediately as they are called, rather than being queued for later compilation. It also supports native distributed APIs such as Fully Sharded Data Parallel (FSDP) and Tensor Parallelism (TP), along with PyTorch’s own distributed primitives like DTensor and FSDP2. Additionally, it works seamlessly with torch.compile, allowing for performance optimizations without leaving the standard PyTorch workflow.

As a result, major Hugging Face libraries including Transformers, Diffusers, Accelerate, PEFT, and TRL now run directly on Trainium. Developers can use the same classes and scripts they would use on any other accelerator. Hugging Face now recommends that users migrate away from Optimum Neuron and use these libraries directly with the new backend. The update aims to simplify the developer experience and make sharing models on the Hugging Face Hub easier for those using AWS infrastructure.

How it works

The core mechanism behind TorchNeuron is its integration as a native PyTorch backend. Instead of translating PyTorch operations through an external compiler layer like XLA, TorchNeuron allows PyTorch to recognize neuron as a first-class device. This enables eager mode execution, where each operation is executed immediately on the hardware. This contrasts with the previous lazy evaluation model, which required building a complete computational graph before execution could begin.

By supporting native distributed primitives, TorchNeuron leverages PyTorch’s built-in tools for scaling models across multiple chips. Features like FSDP2 allow for efficient memory usage during training by sharding model parameters across devices. For inference, the backend supports tensor parallelism, which splits model layers across multiple accelerators to speed up processing. The system also encourages the use of static caches for generation tasks to minimize recompilation overhead, ensuring that performance remains high even with the flexibility of eager execution.

Key details

• TorchNeuron becomes available as a public beta starting October 20, 2026.
• Developers can install the package using the command pip install torch-neuronx.
• The new backend supports eager execution, torch.compile, and native distributed APIs like FSDP and TP.
• In Transformers, 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.
• Users can move models to Neuron devices by passing the device_map=neuron argument in their code.
• Hugging Face recommends using Transformers, Diffusers, and other libraries directly instead of going through Optimum Neuron.

Why it matters

This update significantly lowers the barrier to entry for engineers wanting to use AWS Trainium. Previously, the need to think like a compiler and manage static shapes created a steep learning curve. By aligning the Trainium experience with standard PyTorch workflows, Hugging Face allows teams to reuse existing code and expertise. This reduces the engineering effort required to port models to AWS hardware and minimizes the risk of bugs associated with custom compilation steps.

The increase in supported model architectures also expands the utility of Trainium for production workloads. With 89% of top text-generation models supported natively, organizations can deploy a wider range of state-of-the-art models without waiting for custom integrations. This parity with other accelerators means that choice of hardware can be driven more by cost and availability rather than software compatibility constraints.

What you can do

• Install the new backend by running pip install torch-neuronx in your environment.
• Update your inference scripts to use device_map=neuron when loading models with Transformers.
• Replace Optimum Neuron dependencies with direct calls to Hugging Face libraries like Transformers and Diffusers.
• Implement static caches for text generation tasks to prevent frequent recompilations and improve performance.
• Test your existing distributed training setups using native PyTorch primitives like FSDP2 and DTensor on Trainium instances.
• Review the PyTorch blog for additional technical details on how TorchNeuron integrates with the broader PyTorch ecosystem.