Amazon TorchNeuron: a native PyTorch backend for AWS Trainium, now open source and in public beta

Amazon's TorchNeuron, a native PyTorch backend for AWS Trainium chips, became publicly available on October 19, 2026. TorchNeuron lets PyTorch code run on Trainium2 and Trainium3 by moving models to a neuron device, instead of going through the older PyTorch/XLA compiler stack. AWS shipped it as a public beta in Neuron SDK 2.33.0 on October 19, announced it on the PyTorch blog the same day, and open-sourced it under Apache 2.0 at PyTorch Conference North America in San Jose on October 20. It's Amazon's bid to make the chips in its AI datacenters easier to program.

This story is based on the PyTorch blog announcement, AWS Neuron documentation and release notes, Amazon's news post and SiliconANGLE's report, checked October 20, 2026. For how big Amazon's Trainium clusters have become, see our gigawatt datacenter analysis.

The short version

• What: a PyTorch backend that registers Trainium as a native device, so standard PyTorch APIs work on it.
• Chips: Trainium2 and Trainium3 (Trn2 and Trn3 instances). Not Inferentia, and not first-generation Trainium.
• Status: public beta in Neuron 2.33.0, released October 19, 2026 per AWS's release notes, for PyTorch 2.13 and 2.14. The PyTorch blog dates the public beta from October 20, alongside the open-source release. Open source under Apache 2.0.
• Install: pip install torch-neuronx from the AWS Neuron pip repository, the same package name as the older library.
• Not entirely new: the Neuron docs say AWS first announced TorchNeuron at re:Invent in December 2025 and shipped an initial version in Neuron 2.27.0. Until now it was a private beta.

What is TorchNeuron?

TorchNeuron is "an open source native PyTorch backend for Trainium, so the model code, the libraries, and the tools you already use run on Trainium," the October 19 PyTorch blog post says.

It plugs into PyTorch through PrivateUse1, which the post calls "PyTorch's official extension point for custom hardware." It supports three ways of running code:

• Eager mode, where operations run immediately, so "Python control flow, prints, and breakpoints all work with no compile or trace step." It includes an experimental "adaptive eager" mode.
• torch.compile, for production performance.
• NKI, the Neuron Kernel Interface, for writing custom kernels.

Per the release notes, distributed training uses a new neuron c10d backend that supports DDP, FSDP2, DTensor, pipeline parallelism and context parallelism. The blog says TorchTitan, PyTorch's training library, "trains on Trainium today," and that Hugging Face Transformers, Diffusers, trl and kernels support Trainium upstream.

How is it different from torch-neuronx and the Neuron SDK?

The Neuron SDK is AWS's whole software stack for its chips: compiler, runtime, tools and framework plugins. torch-neuronx is its PyTorch package. What changes is how that package connects to PyTorch.

Last verified: October 20, 2026

Sources: the Neuron docs' "About PyTorch on Neuron" page (now archived), AWS's December 16, 2025 and May 20, 2026 Neuron announcements, and the Neuron 2.33.0 TorchNeuron release notes. The same data is in data/D28-torchneuron-vs-xla.csv. AWS documents give different starting versions: the archived overview says PyTorch 2.10 or later, the May 2026 announcement says 2.12 with Neuron 2.31, and the current "What is TorchNeuron?" page recommends it for PyTorch 2.14. The table uses the October 19 release notes, which list 2.13 and 2.14.

The practical upshot: if you run Inferentia2 or Trn1, TorchNeuron isn't for you. Those instances stay on the PyTorch/XLA path, which AWS still supports for inference but no longer for training.

What do developers change in their code?

Per the docs and the PyTorch blog, the main change is the device:

model = MyModel().to('neuron')   # the only change

The blog says code written against torch.accelerator, PyTorch's device-agnostic API, "runs unmodified," and Hugging Face's device_map="auto" detects Neuron automatically. For compiled mode, you name the backend:

torch.compile(model, backend="neuron")

The release notes list conditions the announcement doesn't stress:

• Moving off XLA is a breaking change. Code written for the XLA-based library must replace torch_xla and xm.xla_device() with the neuron device.
• torch.compile needs static shapes. It "requires dynamic=False, and each new input shape can require another compilation."
• Missing PyTorch features: no torch.sparse, torch.func, torch.nested or torch.distributions support, and no TorchAO or MXFP8 integration yet. The blog says TorchAO and FP8 are in private beta.
• Errors arrive late. In eager mode, the blog says, "an error surfaces at the next synchronization point rather than at the operation that caused it."

AWS says about 100,000 Hugging Face models ran as of September 2026, covering what it calls "99.9% of real-world usage (based on the download count)." That's AWS's own claim. AWS's upstream PyTorch test run in August 2026 skipped about 39,000 tests for Neuron-specific reasons, including about 22,000 operators without a native Neuron implementation.

Who's using it?

The blog quotes three companies. Deepgram's CTO says that "in the benchmarks we evaluated, Trainium processes about 60% more input audio seconds per dollar than our existing GPU baselines." That's a vendor-reported customer figure with no published method. Cartesia said it ran its inference pipeline "with minimal changes to our PyTorch code," and Hugging Face is also quoted. Listed partners include Anyscale, Modular and Decart, the startup Amazon was reported in talks to buy (our Decart story).

Where does Anthropic fit?

Anthropic is Trainium's biggest known user, but nothing published says it uses TorchNeuron. The PyTorch post lists Anthropic only in its acknowledgments, among dozens of organizations thanked "for their feedback in shaping TorchNeuron's design and capabilities." Amazon's own post doesn't mention Anthropic.

The scale is real, though. In its April 20, 2026 compute deal with Amazon, Anthropic said "we currently use over one million Trainium2 chips to train and serve Claude." It expects "nearly 1GW total of Trainium2 and Trainium3 capacity coming online by the end of 2026," and committed more than $100 billion over ten years to AWS technologies (Anthropic explained). Every big cloud now designs its own accelerators, as Microsoft's Maia 300 shows, and software decides how easily customers can use them.

What's confirmed, reported and uncertain

• Confirmed (PyTorch blog, Neuron docs, Amazon): what TorchNeuron is, Trainium2 and Trainium3 support, public beta in Neuron 2.33.0, Apache 2.0 license, the code changes and limitations above.
• Vendor-reported: Deepgram's 60% figure, AWS's 100,000-model coverage and Amazon's claim of "a single configuration change and no rewrite." Amazon's side-by-side training comparison is its own.
• Uncertain: when TorchNeuron leaves beta, when TorchAO and FP8 become generally available, and whether Anthropic or other large customers will move production training to it.