Artificial Intelligence

TorchNeuron Brings Native PyTorch Support to AWS Trainium

A new PyTorch plugin and backend allows developers to run Hugging Face models on AWS Trainium using standard libraries and native interfaces.

Radwan AbdulhadiEditor and technology writer

Published: October 19, 2026 · Updated: October 19, 2026 · 3 min read

TorchNeuron Brings Native PyTorch Support to AWS Trainium — Image: Hugging Face Blog

Introduction to TorchNeuron

TorchNeuron provides a new native backend for PyTorch on AWS Trainium, designed to make running machine learning models on Trainium hardware and sharing them on the Hugging Face hub easier than ever. Detailed announcements and updates can be found directly on the Hugging Face Blog.

The PyTorch plugin brings a new backend for Trainium that enables full support for eager execution mode, torch.compile, and native distributed APIs such as FSDP and tensor parallelism. Researchers and developers can install the package starting October 20, 2026, by using pip install torch-neuronx.

Transitioning Away From Optimum Neuron

Historically, running Hugging Face tools on AWS Trainium and Inferentia relied on the optimum-neuron library, which was built on PyTorch XLA. While effective, that approach treated computations like a compiler by enforcing lazy tensors, static shapes, and full recompilations on graph breaks. Optimum Neuron was necessary to mask that operational complexity through custom modeling classes and trainers.

TorchNeuron shifts this paradigm by integrating directly into PyTorch to expose neuron as a native device. This architecture supports eager execution, PyTorch native distributed primitives like DTensor and FSDP2, and torch.compile. Consequently, Hugging Face libraries now treat Trainium like any other hardware accelerator. Core libraries including Transformers, Diffusers, Accelerate, PEFT, and TRL run natively using standard scripts, leading Hugging Face to recommend that Neuron users transition directly to these libraries.

Expanded Text Generation Support

Text generation now features native support across a wide range of model architectures within Transformers. TorchNeuron supports 89% of the top 100 most-used text-generation model architectures out of the box, a substantial increase compared to the 46% coverage seen in Optimum Neuron v0.4.5.

To utilize the hardware, developers simply pass the device_map=neuron argument within their code to transfer models directly to Neuron devices. The integration heavily recommends utilizing a static cache during generation to prevent frequent graph recompilations.

Streamlined Training and Fine-Tuning

Training workloads on Trainium now utilize the exact same Hugging Face ecosystem libraries used for other accelerators, eliminating the need for dedicated trainers or hardware-specific model code forks. Standard fine-tuning scripts utilizing tools like TRL's SFTTrainer execute smoothly on Trainium instances with minimal configuration adjustments.

Parallelism is managed natively by Transformers through a DistributedConfig passing tensor-parallel and FSDP sizes in model_init_kwargs. This setup allows the SFTTrainer to delegate sharding during weight loading without requiring manual flags or complex Accelerate configurations. Optimization steps include disabling automatic dynamic shapes in torch.compile, implementing wrapped packing for sequence lengths up to 2048 tokens, and managing batch padding in multiples of 128 tokens.

Diffuser Pipelines and Custom Kernels

Beyond text generation and training, multiple pipelines from the Diffusers library achieve full compatibility with TorchNeuron in both eager and compiled modes. Prominent models such as FLUX.1-dev, FLUX.2, Qwen-Image, Qwen-Image-2.1, and H3 also incorporate tensor parallelism compatibility to maximize hardware utilization.

Additionally, developers can leverage specialized kernels from the Neuron team to accelerate vision transformers and other operations. Written via the Neuron Kernel Interface, these low-control custom operations can be shared on the hub to deliver significant speed improvements.