TorchNeuron beta brings native PyTorch to AWS Trainium

TorchNeuron, a new way to run PyTorch code on Amazon's Trainium chips, entered public beta on October 20, 2026. It matters because most popular open models now run on Trainium without changes. Hugging Face says 89% of the 100 most-used text-generation model designs work out of the box, up from 46% with the older tooling.

What TorchNeuron is

Trainium is the AI chip Amazon designs and rents out through AWS, its cloud business. PyTorch is the most widely used library for building and training AI models. Until now, PyTorch code reached Trainium through a translation layer, and many models needed extra work to run.

TorchNeuron is a PyTorch plugin that adds Trainium as a native device. AWS's documentation says it "registers Trainium as a standard PyTorch device." In practice, a developer moves a model to the device named neuron, the same way they would move it to a GPU with cuda.

The project is open source under the Apache-2.0 license, according to its GitHub page. AWS lists it as a public beta in the Neuron 2.33 release, its software kit for the chips.

What changed from the old path

The older route went through Optimum Neuron, a Hugging Face library, which relied on PyTorch XLA. XLA is a compiler layer that turns a model into a fixed graph before it runs. That approach made many models hard to support, because each one needed its own adaptation.

TorchNeuron drops that layer. Hugging Face's own libraries now run on Trainium directly: Transformers for language models, Diffusers for image models, plus Accelerate, PEFT and TRL for training and fine-tuning. The table below sums up the features, from the AWS documentation and the Hugging Face post.

Feature	What it means
Eager mode	Code runs line by line, as normal PyTorch does
torch.compile	Optional compiling for speed, using the neuron backend
Distributed training	FSDP, DDP, DTensor, tensor and pipeline parallelism
Custom kernels	Hand-written chip code through AWS's NKI language
Chips	Trainium 2 and Trainium 3; FP8 and MXFP8 need Trainium 3

FSDP and DDP are standard PyTorch ways to split training across many chips. FP8 is a smaller number format that saves memory and time.

Models and speed

Hugging Face's post shows a fine-tuning run of Qwen3-8B with TRL, split across chips with tensor parallelism and FSDP. It also runs image pipelines such as FLUX.1-dev, FLUX.2 and Qwen-Image with tensor parallelism.

AWS's documentation lists mixture-of-experts models, which route each input to a few specialist sub-networks. DeepSeek, GPT-OSS and Qwen3 are supported. Training those models goes through a separate tool, TorchTitan Neuron.

The only speed figure is narrow. A vision transformer kernel from AWS's Neuron team gives "up to a x2.4 speedup on Trainium," the Hugging Face post says. It does not say what that is measured against.

How to install it

The GitHub page gives a pinned install, run after the Neuron driver and runtime packages are in place:

pip install torch==2.14.0 --index-url https://download.pytorch.org/whl/cpu
pip install "torch-neuronx==2.14.*" "neuronx-cc==2.*" "nki>=0.6" --extra-index-url https://pip.repos.neuron.amazonaws.com

The pins matter. TorchNeuron ships under the same package name, torch-neuronx, as the older XLA-based product. Earlier versions install the old one, the repository warns. Running python -m torch_neuronx.collect_env confirms the device is usable from PyTorch. AWS's ready-made machine images and containers have it preinstalled.

What this means for developers

TorchNeuron turns Trainium from a special target into one more PyTorch device. That is the step that makes price comparisons with GPUs practical for ordinary teams.

• Test your own model first. 89% coverage still leaves about one in ten popular designs that need work. Run a short training or inference job before planning a move.
• Pin versions in every environment. An unpinned torch-neuronx in a requirements file can pull the old XLA product. Check collect_env output in CI.
• Check the chip generation. Only Trainium 2 and 3 instances are supported. If your plan depends on FP8 training, you need Trainium 3.
• Treat it as beta. None of the three sources calls it ready for production. Keep a GPU fallback for anything with a deadline.
• Benchmark cost per job, not per hour. The 2.4x figure covers one kernel with no stated baseline. Measure your own tokens per dollar before you commit budget.