Amazon Announces TorchNeuron, a Native PyTorch Backend for Its Trainium AI Chips

Amazon on October 19, 2026 announced TorchNeuron, an open source backend that lets code written in PyTorch run on its Trainium AI chips as a standard PyTorch device. The software goes into public beta on October 20 under the Apache 2.0 license and supports Trainium2 and Trainium3, the company said.

PyTorch is the open source framework used by many AI developers, and Nvidia's lead in AI chips rests partly on how much of that code was built around its CUDA software, GeekWire notes. Until now, running PyTorch on Trainium meant calling interfaces specific to Amazon's chips, Colin Brace, vice president of technology software at Amazon's chip unit Annapurna Labs, told GeekWire. TorchNeuron has been in private beta since late 2025, with more than 30 companies and internal Amazon teams using it, according to GeekWire.

What exactly did Amazon announce?

Amazon said TorchNeuron registers Trainium through PyTorch's own extension point for custom hardware, so a model moves to the chip with .to('neuron') and the rest of a training loop stays the same. The backend will be published on GitHub and installable with pip from October 20, timed to PyTorch Conference North America in San Jose, California (GeekWire, The Next Web).

According to the PyTorch blog post written by AWS engineers, both of PyTorch's execution modes work. Eager mode runs each operation immediately, which suits research and debugging. torch.compile turns a model into a graph that the Neuron compiler optimizes for production. Standard distributed training tools such as FSDP2, tensor parallelism and pipeline parallelism also run as written, AWS says.

Developers who want more speed can write custom kernels, small hand tuned programs for one operation, in the Neuron Kernel Interface (NKI), and AWS has released a library of such kernels for attention and mixture of experts models.

"If you know PyTorch, then you know Trainium," Brace told GeekWire.

What changes for developers already on Trainium?

The main change is that Trainium no longer needs its own version of the PyTorch stack, according to Hugging Face. Its earlier Trainium library, Optimum Neuron, sat on a package built on PyTorch XLA, which meant lazy tensors instead of eager execution, static shapes only and a recompilation for every graph break, Hugging Face wrote in its blog post.

With TorchNeuron, Hugging Face's Transformers, Diffusers, Accelerate, PEFT and TRL libraries treat Trainium like any other device, and the company now recommends that Neuron users call them directly instead of going through Optimum Neuron.

Hugging Face notes that the asynchronous weight loading in Transformers does not work on Neuron yet.

How well does it work so far?

AWS's own testing shows high but incomplete coverage. In a sweep of about 100,000 Hugging Face models run in reduced configurations as of September 2026, 95.6% passed in eager mode (96,618 of 101,076 runs) and 96.6% with torch.compile (95,097 of 98,420), according to charts in the PyTorch blog post.

On PyTorch's own upstream test suites in August 2026, AWS reports pass rates of 94.7% for ATen operators, 96.6% for the Dynamo compiler, 94.7% for distributed tests and 100% for the profiler, counted against the tests it enabled. That run skipped roughly 69,000 of about 154,000 test outcomes. About 30,000 were skips PyTorch applies itself, such as CUDA only tests, and about 39,000 were deliberate Neuron skips, the largest group being some 22,000 tests for operators that do not yet run natively on the chip.

AWS says it chose not to support several features in the first release, including sparse tensors, NestedTensor, forward mode autograd and dynamic and symbolic shapes. FP8 and the TorchAO quantization library are in private beta, with general availability planned for the next release.

Customer results are company claims. Deepgram, a speech AI startup, says that in the benchmarks it evaluated, Trainium processed about 60% more input audio seconds per dollar than its existing GPU baselines, CTO Adam Sypniewski said in the PyTorch post. Deepgram now runs Trainium in production, Brace told GeekWire.

Why does Amazon want this now?

Amazon wants Trainium used well beyond the AI labs it has financial ties to, and Brace called TorchNeuron "an accelerant" toward that goal, according to GeekWire. Trainium's biggest users so far are Anthropic, which uses the chips to train and run its Claude models, and OpenAI, which in February agreed to 2 gigawatts of Trainium capacity as part of a deal in which Amazon is investing $50 billion, according to GeekWire.

The chip business already carries weight at Amazon. CEO Andy Jassy wrote in his shareholder letter on April 9, 2026 that Amazon's custom chips generated more than $20 billion in annualized revenue and would be worth about $50 billion a year as a standalone chip company (The Next Web). Amazon now says the business has passed a $25 billion annual revenue run rate.

Rivals are taking the same route. Google is building TorchTPU to run PyTorch natively on its TPUs (GeekWire), and Microsoft said on October 12 that its Maia chips will become its main inference system. One TorchNeuron partner named in the PyTorch post is Decart, whose software helps developers move models between AI chips and which Amazon is in talks to buy at a valuation of about $7 billion (WSJ via The Next Web).

What we don't know yet

• How TorchNeuron performs against Nvidia GPUs on the same models in independent benchmarks.
• When dynamic shapes and the other features left out of the first release will be supported.
• Whether the public beta comes with any change to Trainium instance pricing or availability.
• Whether Amazon's reported talks to buy Decart will lead to a deal.