TorchNeuron explained: Hugging Face says 89% of top text-generation architectures now run natively on AWS Trainium
A new PyTorch backend replaces the XLA-based Optimum Neuron path. Here is what changes for practitioners, what the numbers actually say, and the static-shape catch.
TorchNeuron provides a new native backend for PyTorch on AWS Trainium. From now on, we recommend that Neuron users use these libraries directly instead of going through Optimum Neuron. Nothing in the script is tied to Trainium: the same code runs on any hardware.
Comprehensive take
Switch from Optimum Neuron to TorchNeuron if you are on Trainium; GPU users should wait for throughput and cost numbers before migrating.
Use cases
SFT of 8B-class models on Trainium2
The post's reference run fine-tunes Qwen3-8B on ultrachat_200k with TP=4, FSDP=8 across 32 NeuronCores using an unmodified TRL SFTTrainer script.
Tensor-parallel inference of 27B models
The inference example runs Qwen3.8-27B with a static cache, max sequence 1024 and prefill chunk 512, using the neuron distributed backend.
Diffusion image generation with TP
FLUX.1-dev, FLUX.2, Qwen-Image, Qwen-Image-2.1 and H3 support tensor parallelism; the example generates 1024x1024 images in 40 steps at TP=8.
Specs
| Architecture coverage (top 100 text-gen) | 89% vs 46% (Optimum Neuron v0.4.5) source |
|---|---|
| Public beta start | 20 October 2026 source |
| Install | pip install torch-neuronx source |
| Example training hardware | 8 Trainium2 chips / 32 NeuronCores source |
| Vision transformer kernel speedup | up to 2.4x source |
| Padding multiple (TRL >= 1.16) | 128 tokens source |
Insights
- The 89% coverage figure is about models loading, not running faster; no throughput or cost is published.
- Static shapes are still mandatory in practice because Dynamo dynamic shapes cannot compile on Neuron yet.
- Hugging Face is effectively deprecating Optimum Neuron in favour of direct library use.
Axion Scorecard
Overall: 7.0
89% of top-100 text-gen architectures run out of the box, up from 46%; speed unpublished. source
Standard open HF libraries and Hub-hosted kernels; the Neuron compiler stack remains AWS-controlled. source
Unmodified Transformers/TRL/Diffusers scripts, but static shapes and Neuron env flags remain. source
Not yet scored — no cited data
Limitations (dynamic shapes, async load, TP coverage) are stated plainly; performance numbers are absent. source
Replacing an XLA wrapper with a native PyTorch backend is a meaningful architectural shift for non-NVIDIA accelerators. source
The problem TorchNeuron solves
Until now, running Hugging Face models on AWS Trainium meant going through Optimum Neuron, which sat on torch-neuronx, which in turn was built on PyTorch XLA. That stack imposed lazy tensors, static shapes only, and a recompilation at every graph break. It also required custom modelling and trainer code, so the Trainium path diverged from the code you ran everywhere else.
TorchNeuron, announced by Hugging Face on 19 October with public beta from 20 October, is a PyTorch plugin that adds `neuron` as a native device. It supports eager mode, `torch.compile`, and PyTorch distributed APIs including FSDP and tensor parallelism. Because it plugs into PyTorch itself rather than wrapping it, Transformers, Diffusers, Accelerate, PEFT and TRL run with their standard classes and scripts. Hugging Face's own line is blunt: it now recommends Neuron users 'use these libraries directly instead of going through Optimum Neuron.'
The numbers, exactly as published
| Metric | Value |
|---|---|
| Top-100 text-generation architectures working out of the box | 89% (TorchNeuron) vs 46% (Optimum Neuron v0.4.5) |
| Training hardware in the example | 8 Trainium2 chips, 32 NeuronCores (half a trn2.48xlarge) |
| Training layout | TP=4, FSDP=8, 32 processes |
| Training model / dataset | Qwen/Qwen3-8B on HuggingFaceH4/ultrachat_200k |
| Training hyperparameters | max length 2048, lr 2e-5, warmup 0.03, 1 epoch, per-device batch 4, grad accum 4, bf16 |
| Inference example | Qwen/Qwen3.8-27B, max seq 1024, prefill chunk 512, 128 new tokens |
| Diffusers example | TP=8, 1024x1024, 40 steps, Qwen-Image-2.1 |
| Kernel speedup | up to 2.4x for vision transformers via a Hub kernel |
What is not published: throughput, latency, or dollar cost. The 89% figure is about coverage, not speed. Read it as 'your model probably loads', not 'your model is faster'.
How you actually use it
Install is `pip install torch-neuronx`. For Transformers inference, pass `device_map=neuron`, or initialise a process group with `dist.init_process_group(backend="neuron")` for tensor parallelism with a static cache. For training, the post uses an unmodified TRL `SFTTrainer` script with a `DistributedConfig` for TP and FSDP, launched as `TP_SIZE=4 FSDP_SIZE=8 HF_DEACTIVATE_ASYNC_LOAD=1 uv run torchrun --nproc_per_node=32 sft_neuron_example.py`. Diffusers runs with `NEURON_LOGICAL_NC_CONFIG=2 torchrun --nproc_per_node=8 run.py`. The quotable claim is that 'nothing in the script is tied to Trainium: the same code runs on any hardware.'
The catches Hugging Face states
This is a beta, and the post is honest about the edges. Automatic dynamic shapes in Dynamo cannot yet be compiled on Neuron, so static shapes are still required in practice and a static KV cache is recommended to avoid recompilation. With TRL 1.16 or later, Neuron batches pad to a multiple of 128 tokens, which costs you some throughput on short sequences. Transformers' built-in asynchronous weight loading is unsupported and must be disabled with `HF_DEACTIVATE_ASYNC_LOAD=1`. On the Diffusers side, tensor parallelism is confirmed for FLUX.1-dev, FLUX.2, Qwen-Image, Qwen-Image-2.1 and H3; other pipelines work in eager and compiled mode but not yet with TP.
So the headline 'no custom code' is true for model definitions and trainers, but not for the environment: you still carry Neuron-specific launch flags and a static-shape discipline that GPU users do not.
Why this matters beyond AWS
The architectural shift is the interesting part for practitioners. A native PyTorch backend means the Trainium path now follows upstream PyTorch and Transformers releases instead of waiting for a wrapper library to catch up. The jump from 46% to 89% architecture coverage is exactly what you would expect from removing a translation layer. It also means kernels can be distributed through the Hub (the AWS Neuron team's NKI kernels are already there), which is the same distribution model the GPU world has converged on.
Axion verdict
If you already run on Trainium through Optimum Neuron, start migrating now; Hugging Face has effectively deprecated that path. If you are on GPUs and Trainium-curious, the missing throughput and cost numbers mean you cannot yet make the business case from this post alone. Benchmark your own model with the static-shape constraints before committing capacity.