TorchNeuron brings native PyTorch support for HF models to AWS Trainium

A native PyTorch backend for Trainium

TorchNeuron is a PyTorch plugin that adds a backend for AWS Trainium. According to Hugging Face, it supports eager mode, torch.compile and native distributed APIs such as FSDP and TP. Researchers and developers will be able to install it with pip install torch-neuronx once the public beta starts on Oct. 20, 2026.

Until now, the Hugging Face experience on AWS Trainium and Inferentia, chips designed respectively for cost-effective training and inference at scale, has been delivered through the optimum-neuron library. That library relies on the torch-neuronx package, which is based on PyTorch XLA. Hugging Face describes this as a choice that made sense at the time, but one that required "thinking like a compiler": lazy tensors instead of eager execution, static shapes only, and a recompilation for every graph break. Optimum Neuron maintained custom code, including modeling classes and trainers, to hide that complexity.

Trainium as a device like any other

TorchNeuron integrates directly into PyTorch and exposes neuron as a native device, with eager execution, PyTorch's own distributed primitives (DTensor, FSDP2) and torch.compile. For the Hugging Face libraries, this means Transformers, Diffusers, Accelerate, PEFT and TRL run on Trainium directly, using the same classes and scripts as on any other accelerator.

Hugging Face now recommends that Neuron users work with these libraries directly instead of going through Optimum Neuron.

For inference, text generation is natively supported in Transformers. The company says 89% of the top 100 most-used text-generation model architectures are supported out of the box with TorchNeuron, compared to 46% with Optimum Neuron v0.4.5. Users pass the device_map=neuron argument and the model is moved to the Neuron devices. The example in the post uses tensor parallelism and a static cache, which Hugging Face strongly recommends on Neuron to prevent frequent recompilations.

Training with the standard TRL workflow

Training on Trainium now uses the same libraries as other accelerators: Transformers, Accelerate, PEFT and TRL. According to the post, there is no longer a need for a dedicated trainer or a Neuron-specific fork of the model code. A standard TRL fine-tuning script runs on a Trainium instance with only a few Neuron-specific settings.

Parallelism is handled natively by Transformers. A DistributedConfig with tensor-parallel and FSDP sizes is placed in the model initialization arguments, and the model is sharded as its weights are loaded. No --fsdp flag or Accelerate FSDP config is needed, and nothing in the example script is tied to Trainium.

Hugging Face lists a few points to know when running on Trainium:

• Static shapes: the script disables automatic dynamic shapes in torch.compile and uses wrapped packing for sequences of up to 2048 tokens. With TRL 1.16 or later, batches on Neuron are padded to a multiple of 128 tokens, which limits the number of input shapes. Compiled graphs are reused for later steps with the same shapes.
• torch.compile: passing --torch_compile true is enough, as it detects the device and picks the Neuron backend automatically.
• Async weight loading: the built-in asynchronous loading in Transformers is not supported on Neuron, so it is disabled for now with HF_DEACTIVATE_ASYNC_LOAD=1.

The example launches training with torchrun, one process per NeuronCore. With TP_SIZE set to 4 and FSDP_SIZE set to 8, 32 processes form a 4 x 8 grid. Tensor parallelism splits each layer across the 4 cores of a TP group, and FSDP shards parameters, gradients and optimizer state across the 8 TP groups. The example fine-tunes Qwen/Qwen3-8B on the HuggingFaceH4/ultrachat_200k dataset, and Hugging Face says it ran on 8 Trainium2 chips (32 NeuronCores).

What it means for Trainium users

For developers already working with Hugging Face tools, the change reduces the Neuron-specific code they need to write. Scripts written for other accelerators should carry over to Trainium with a small number of settings, and models can be shared on the Hugging Face hub. TorchNeuron is a public beta, and the post notes a known limitation: the Neuron backend cannot yet compile Dynamo's automatic dynamic shapes, so some workarounds remain necessary.