Hot AI News

Hugging Face drops its Trainium middleman: Transformers now runs on Neuron directly

Hugging Face says its main libraries now treat AWS Trainium as an ordinary PyTorch device, and it now recommends them over Optimum Neuron.

Hugging Face drops its Trainium middleman: Transformers now runs on Neuron directly
Image: Hugging Face

Hugging Face says its Transformers, Diffusers, Accelerate, PEFT and TRL libraries now run natively on AWS Trainium through TorchNeuron, and that Neuron users should now use them directly instead of the Optimum Neuron library. Hugging Face reports that 89% of the top 100 text-generation architectures work out of the box.

The details come from a Hugging Face blog post published October 19, 2026, written with the AWS Neuron team. TorchNeuron is a PyTorch plugin that exposes Trainium as a native neuron device, with eager execution, torch.compile and PyTorch’s own distributed tools such as FSDP2 and DTensor. We covered the Amazon side of the launch in our earlier report.

Hugging Face drops its Trainium middleman: Transformers now runs on Neuron directly
Image: Hugging Face

Text generation coverage jumps from 46% to 89%

Until now, Hugging Face support on Trainium and Inferentia went through optimum-neuron, which sits on top of the older torch-neuronx package built on PyTorch XLA. Hugging Face describes that route as requiring compiler-style thinking: lazy tensors, static shapes only and a recompilation at each graph break. Optimum Neuron hid this with its own modeling classes and trainers.

With TorchNeuron, Hugging Face says 89% of the 100 most-used text-generation architectures are supported, against 46% with Optimum Neuron v0.4.5. Loading a model onto the chips takes the argument device_map=neuron. For multi-chip inference, a DistributedConfig sets the tensor-parallel size, and the post recommends a static cache to avoid frequent recompilation.

Training uses the same stack as on any GPU. The post shows a standard TRL SFTTrainer script that fine-tunes Qwen3-8B on the UltraChat 200k dataset, launched with torchrun across 32 NeuronCores in a 4 x 8 tensor-parallel and FSDP grid. Hugging Face says it ran on 8 Trainium2 chips, half of a trn2.48xlarge instance, and finished, producing a standard Transformers checkpoint that can be pushed to the Hub and loaded on other hardware. No throughput, cost or loss figures are given.

For image models, Hugging Face lists FLUX.1-dev, FLUX.2, Qwen-Image, Qwen-Image-2.1 and H3 as compatible with tensor parallelism in Diffusers, with an example running at TP=8. It also points to Neuron kernels on the Hub, written with the Neuron Kernel Interface (NKI), including one it says speeds up vision transformers by up to 2.4x.

Advertisements

Rough edges the post admits

Several caveats are in the post itself. Asynchronous weight loading in Transformers is not supported on Neuron and must be disabled with HF_DEACTIVATE_ASYNC_LOAD=1. The training script turns off Dynamo’s automatic dynamic shapes because the backend cannot compile them yet, and batches are padded to multiples of 128 tokens with TRL 1.16 or later to limit distinct shapes. The Diffusers example keeps the text encoder on the CPU. Only some pipelines support tensor parallelism so far.

Availability is also narrower than the headline suggests. Hugging Face says TorchNeuron becomes a public beta on October 20, 2026, installable with pip install torch-neuronx. AWS’s Neuron documentation had until recently described the native backend as available only through a closed beta program, so this is the first broad opening. Our earlier coverage reports that Amazon plans to open-source the software under Apache 2.0 at the PyTorch Conference on the same date. The post does not mention pricing or regions.

Why a native device matters more than the percentage

The technical change is where the work moves. Under XLA, the hardware vendor and Hugging Face both had to maintain Trainium-specific copies of model code, because the graph had to be traced and compiled ahead of time. As a regular PyTorch backend, the chip runs ordinary model classes, and new architectures added to Transformers can work without a separate port. That helps explain the jump in architecture coverage, though the 89% and 46% figures are Hugging Face’s own and were measured against different software stacks.

Static shapes remain the hard part. Trainium compiles graphs ahead of time, so varying sequence lengths trigger recompiles, which is why the examples pad inputs and use fixed caches. TorchNeuron makes the programming model familiar; it does not remove that constraint.

The hardware is also moving. Amazon’s materials, as relayed by analyst firm Introl, put a Trainium3 chip at 2.52 petaflops of FP8 compute with 144 GB of HBM3e memory, and a 144-chip UltraServer at 362 petaflops. Those are vendor figures with no workload details. The Hugging Face example runs on Trainium2, and the post does not say how it performs on Trainium3.

What the source does not establish: no independent benchmarks, no cost comparison with GPU instances, no list of the 11% of architectures that remain unsupported, and no timeline for retiring Optimum Neuron. Hugging Face recommends moving away from it but does not say it will stop maintaining it. Anyone with an existing Optimum Neuron pipeline should check model coverage before switching.

Sign Up for the Daily AI Pulse

One email a day. All the stories that matter.

Alex CarterAI News & Big Tech Correspondent

Alex Carter writes for Aiholics, keeping readers updated on the fast-paced world of AI and Big Tech.

Leave a Reply

Your email address will not be published. Required fields are marked *

You may also like