Hugging Face libraries now run natively on AWS Trainium via TorchNeuron

Hugging Face has announced that TorchNeuron, a new native PyTorch backend for AWS Trainium, is supported directly by its Transformers, Diffusers, Accelerate, PEFT, TRL and Kernels libraries. TorchNeuron is a PyTorch plugin that adds a Trainium backend with support for eager mode, torch.compile and native distributed APIs such as FSDP and TP. Starting Oct. 20, 2026, it will be available as a public beta, and researchers and developers can install it with pip install torch-neuronx.

Until now, the Hugging Face experience on Trainium and Inferentia, AWS chips built for cost-effective training and inference at scale, came through the optimum-neuron library. That library relies on the torch-neuronx package, which is based on PyTorch XLA. The post calls it the natural choice at the time, but says it meant thinking like a compiler: lazy tensors instead of eager execution, static shapes only, and a recompilation for every graph break. Optimum Neuron existed to hide that complexity and maintained custom code, including modeling classes and trainers, to do so. TorchNeuron instead exposes neuron as a native PyTorch device, with eager execution, PyTorch's own distributed primitives (DTensor, FSDP2) and torch.compile. For the Hugging Face libraries, Trainium becomes a device like any other, running the same classes and scripts used on other accelerators. Hugging Face now recommends that Neuron users use these libraries directly instead of going through Optimum Neuron.

On the Transformers side, text generation is natively supported for a wide range of models. With TorchNeuron, 89% of the top 100 most-used text-generation model architectures are supported out of the box, compared with 46% with Optimum Neuron v0.4.5. Moving a model to the hardware takes the device_map=neuron argument. The post's inference example runs Qwen/Qwen3.8-27B with tensor parallelism across the available devices. It uses a static cache, which the post says is strongly recommended on Neuron to prevent frequent recompilations, along with chunked prefill of 512 tokens and a maximum sequence length of 1024.

For training, the post says there is no longer a need for a dedicated trainer or a Neuron-specific fork of the model code. A standard TRL SFTTrainer fine-tuning script runs with only a few Neuron-specific settings. Parallelism is handled by Transformers: a DistributedConfig with tensor-parallel and FSDP sizes goes into model_init_kwargs, the model is sharded as its weights load, and no --fsdp flag or Accelerate FSDP config is needed. Three Trainium-specific points are listed. First, static shapes: the script disables Dynamo's automatic dynamic shapes (the Neuron backend cannot compile them yet, per a code comment linked to a TRL issue) and uses wrapped packing to produce 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, and compiled graphs are reused for later steps with the same shapes. Second, passing --torch_compile true is enough, because torch.compile detects the Trainium device and picks the Neuron backend automatically. Third, Transformers' built-in asynchronous weight loading is not supported on Neuron, so it is disabled for now with HF_DEACTIVATE_ASYNC_LOAD=1.

The example launch fine-tunes Qwen/Qwen3-8B on HuggingFaceH4/ultrachat_200k for one epoch in bfloat16 using torchrun, which starts one process per NeuronCore. TP_SIZE=4 and FSDP_SIZE=8 give 4 x 8 = 32 processes, matching --nproc_per_node=32. 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 authors ran it on 8 Trainium2 chips (32 NeuronCores), half a trn2.48xlarge instance. Training ran to completion, and trainer.save_model wrote a standard Transformers checkpoint that can be pushed to the Hub as is and loaded on other hardware.

For Diffusers, many pipelines are compatible with TorchNeuron in eager and compiled mode, and FLUX.1-dev, FLUX.2, Qwen-Image, Qwen-Image-2.1 and H3 also work with tensor parallelism. The example generates a 1024 x 1024 image with Qwen/Qwen-Image-2.1 over 40 inference steps with TP=8, keeping the text encoder on the CPU, launched with NEURON_LOGICAL_NC_CONFIG=2 torchrun --nproc_per_node=8. The post says more pipelines will gain tensor-parallel support and points to a guide for adding it to a pipeline that lacks it.

On Kernels, the post says more Neuron kernels from the Neuron team will become available on the Hub. They can be written with the Neuron Kernel Interface (NKI), which gives lower-level control for optimizing operations, so users can run and share their own custom operations. One such kernel is cited as giving up to a x2.4 speedup for vision transformers on Trainium. The authors close by thanking the AWS team for their work on TorchNeuron and for responding quickly to feedback on the Hugging Face integration.

Key facts

TorchNeuron is a PyTorch plugin that adds a native Trainium backend with eager mode, torch.compile, FSDP2 and DTensor; the public beta starts Oct. 20, 2026 and installs with pip install torch-neuronx.
Hugging Face now recommends that Neuron users use Transformers, Diffusers, Accelerate, PEFT and TRL directly instead of Optimum Neuron, which relied on PyTorch XLA.
With TorchNeuron, 89% of the top 100 most-used text-generation architectures are supported out of the box, versus 46% with Optimum Neuron v0.4.5.
A standard TRL SFT script fine-tuned Qwen3-8B with TP 4 x FSDP 8 on 8 Trainium2 chips (32 NeuronCores), half a trn2.48xlarge instance, and saved a normal Transformers checkpoint.
Diffusers pipelines run in eager and compiled mode, with tensor parallelism for FLUX.1-dev, FLUX.2, Qwen-Image, Qwen-Image-2.1 and H3; a Neuron team kernel is cited at up to a x2.4 vision-transformer speedup.

Why it matters

Running Hugging Face models on AWS's own chips used to mean going through a separate library, Optimum Neuron, built on PyTorch XLA with lazy tensors, static shapes and recompilation at every graph break. TorchNeuron makes Trainium a regular PyTorch device, so the same Transformers, TRL and Diffusers code used on other accelerators runs there. The post puts the coverage gain at 89% of the top 100 text-generation architectures out of the box, against 46% with Optimum Neuron v0.4.5.

Who it affects

People who train or serve models on AWS Trainium (and the Inferentia chips Optimum Neuron also covered) and use Hugging Face libraries are the main audience: Hugging Face now recommends they use the libraries directly. Researchers and developers who want to try the backend can install it from the public beta. Authors of custom kernels and Diffusers pipelines are also addressed, through NKI and the tensor-parallel guide.

How to use it

Starting Oct. 20, 2026, install the beta with pip install torch-neuronx. For inference in Transformers, pass device_map=neuron; for multi-device runs, use a DistributedConfig with tp_size and a DeviceMesh on the neuron backend, and use a static cache for generation. For training, take a normal TRL SFTTrainer script, set TP_SIZE and FSDP_SIZE through DistributedConfig, launch with torchrun using one process per NeuronCore, pass --torch_compile true, and set HF_DEACTIVATE_ASYNC_LOAD=1. For Diffusers, use TensorParallelConfig with a DeviceMesh and launch with torchrun; the text encoder stays on the CPU in the example. Neuron kernels are published on the Hub and can be written with NKI. The post does not state pricing or instance availability.

How solid is it

This is a first-party announcement from Hugging Face, with the AWS team thanked for building TorchNeuron. The code examples are concrete, and the authors report that the Qwen3-8B fine-tune ran to completion and produced a standard checkpoint. The numbers are thin, though. The 89% versus 46% figure counts supported architectures, and the post does not list which ones are covered. No throughput, latency, cost or training-speed comparisons with GPUs or with Optimum Neuron are given, no training quality results are reported, and the baseline for the up-to-x2.4 vision-transformer speedup is not stated.

Risks and caveats

TorchNeuron is a public beta, and the post gives no general-availability date. Known gaps appear in the examples: the Neuron backend cannot yet compile Dynamo's automatic dynamic shapes, so the training script turns them off, and Transformers' asynchronous weight loading is unsupported and must be disabled for now. Static shapes still matter: generation should use a static cache to avoid frequent recompilations, and batches on Neuron are padded to multiples of 128 tokens with TRL 1.16 or later. Hugging Face recommends rather than requires leaving Optimum Neuron, and the post does not say Optimum Neuron is deprecated. Tensor parallelism for Diffusers covers only some pipelines so far.

“For the Hugging Face libraries, Trainium becomes a device like any other”

— Hugging Face blog post