Hugging Face adds native TorchNeuron path for AWS Trainium

Hugging Face has added native TorchNeuron support to its main model-development libraries, giving supported AWS Trainium workflows a direct PyTorch device path rather than routing them through the Optimum Neuron abstraction. Hugging Face says the TorchNeuron public beta is scheduled to become available October 20 through AWS’s torch-neuronx package.

The practical change is that Transformers, Diffusers, Accelerate, PEFT and TRL can use Trainium through their standard APIs. TorchNeuron registers the accelerators as PyTorch’s neuron device through the framework’s PrivateUse1 extension, with support for eager execution, torch.compile and standard distributed interfaces. AWS documents the beta for training and inference on Trainium2 and Trainium3.

Optimum Neuron takes a different route, using PyTorch XLA’s lazy execution model and providing Neuron-specific trainers and model classes. Hugging Face now recommends using its libraries directly for workflows supported by the native backend. That does not amount to a general retirement of Optimum Neuron: its public repository still covers Trainium (trn1) and Inferentia2 (inf2), while the native beta’s stated hardware scope is narrower.

Hugging Face says the new backend works out of the box with 89% of the 100 most-used text-generation model architectures, compared with 46% for Optimum Neuron v0.4.5. This is a company compatibility claim; the published material does not include the complete model list, test harness or failure criteria needed to reproduce it independently.

Hugging Face also says a standard TRL supervised fine-tuning workflow can run without a dedicated Neuron trainer or model fork. In its example, the company fine-tuned Qwen3-8B on eight Trainium2 chips, equal to 32 NeuronCores, using tensor parallelism of four and fully sharded data parallelism of eight, then saved a standard Transformers checkpoint.

Separately, Hugging Face says a Neuron vision-transformer kernel can deliver up to a 2.4x speedup on Trainium. This remains a company-presented performance claim and has not been independently reproduced.

The beta has material constraints. Dynamic shapes are not natively supported with torch.compile, so users must set dynamic=False; a new input shape can cause another compilation, making padding or bucketing necessary. AWS also lists torch.sparse, torch.func, torch.nested and torch.distributions as unsupported, while torch.fft has limited coverage. The first eager operation or compiled graph includes compilation time when no matching cached artifact exists; later matching runs can reuse the cache.