Amazon's TorchNeuron lets PyTorch code run on Trainium chips with one line

For years the strongest argument for Nvidia was not the chip. It was the software. Almost every AI team writes its models in PyTorch, and PyTorch has been tuned for Nvidia's CUDA platform so thoroughly that running on anything else usually meant extra work, strange errors and weeks of porting. On 19 October 2026, Amazon Web Services announced a tool that is meant to remove that excuse for its own chips. It is called TorchNeuron, and AWS says it lets existing PyTorch code run on Amazon's Trainium chips with a single configuration change and no rewrite.

This article is based on the official PyTorch blog post "Introducing TorchNeuron", Amazon's own announcement on About Amazon, the AWS Neuron 2.33.0 release notes and reporting by TechCrunch. The code is public on GitHub under aws-neuron/torch-neuronx with an Apache 2.0 license. Performance claims below come from AWS and its customers. No independent benchmark has been published yet.

What Amazon actually released

TorchNeuron is a native backend for PyTorch. In plain words, it teaches PyTorch that a Trainium chip is just another device, in the same way PyTorch already knows about a CPU or an Nvidia GPU. Developers move their model with `.to('neuron')` instead of `.to('cuda')`, and the rest of the training loop stays as written. Code that already uses `torch.accelerator`, the device neutral API that the PyTorch project now recommends, runs without any change at all.

AWS announced the tool on Monday 19 October and released it as open source and as a public beta on Tuesday 20 October, timed with the PyTorch Conference North America. It ships as part of the AWS Neuron SDK version 2.33.0 and installs with `pip install torch-neuronx`. It supports Trainium2 and Trainium3, the two current generations of Amazon's in house AI chips, which are designed by Annapurna Labs, the chip unit Amazon bought in 2015.

Amazon says about 85 percent of AI researchers and developers write their code in PyTorch. That number explains why this matters. If the framework that almost everyone uses treats Trainium as a first class device, then the cost of trying Amazon's chips drops from a porting project to an afternoon experiment.

Why this is aimed straight at Nvidia

Nvidia's dominance in AI hardware rests on two legs. One is raw performance. The other is CUDA, the software layer Nvidia has built since 2007, and the huge amount of PyTorch code, libraries and kernels that assume CUDA is there. TechCrunch put it bluntly: Nvidia's hardware is now deeply embedded in PyTorch, essentially by default. Even when a rival chip is cheaper on paper, the engineering time needed to move a real model over can wipe out the savings.

Amazon has had a PyTorch path to Trainium before, through older Neuron tools, but it required Neuron specific software and configuration. Hugging Face says exactly that in its quote for the launch: until now, using Trainium meant adapting models with Neuron specific code. TorchNeuron instead plugs into PrivateUse1, the official extension point PyTorch created for custom hardware. That means the standard tools, from `torch.nn` to `torch.distributed` to `torch.profiler`, work on the `neuron` device as they are.

Colin Brace, VP of Technology Software at AWS Annapurna Labs, framed the launch as a fight against lock in. "The real story is removing the lock-in that kept developers from having a real choice at all," he told TechCrunch. Amazon also says it is contributing changes upstream to PyTorch to remove hard coded assumptions about specific chips. The PyTorch blog lists more than 25 pull requests to PyTorch and TorchTitan, several of them removing CUDA assumptions from the test suite.

AWS is not the first to try this. PyTorch users can already target AMD GPUs and Apple's chips with small configuration changes. What makes the Amazon move different is scale: AWS is the largest cloud provider, it designs its own chips, and it can bundle them with storage, networking and every other service a startup needs.

How it works under the hood

TorchNeuron supports both ways PyTorch runs code. In eager mode, each operation runs immediately, so Python control flow, print statements and breakpoints behave as usual. This is the mode researchers use to debug. On Trainium, each eager operation compiles its own small executable, launches it and moves its results through the chip's high bandwidth memory. That works, but it is slow, so AWS added an experimental "adaptive eager" mode that tries to fuse operations together without the developer changing anything.

For production, developers use `torch.compile`. The model is captured as a graph, the backward pass is traced ahead of time, and the result is lowered through Torch MLIR into StableHLO before the Neuron compiler produces a NEFF file, the executable format Trainium runs.

The release also covers the parts of PyTorch that matter for large models. Distributed training works with DDP, FSDP2, tensor parallelism, pipeline parallelism and combinations of them, and collective operations run over Trainium's dedicated collective cores inside a server and over Amazon's Elastic Fabric Adapter between servers. For engineers who want maximum speed, the Neuron Kernel Interface (NKI) lets them write custom kernels and register them as normal PyTorch operators. AWS also open sourced a library of ready made kernels, including FlashAttention and mixture of experts kernels, and the 2.33.0 release adds 25 new experimental kernels for models such as DeepSeek V3.2, Qwen3.5, Gemma4 and GPT OSS.

On the library side, Trainium is now a supported device upstream in Hugging Face Transformers, Diffusers, TRL and the kernels library. Loading a model is a matter of writing `device_map="neuron"`, and `device_map="auto"` finds the chip on its own. For serving, vLLM Neuron handles language models and vLLM Omni Neuron handles image and video diffusion models. TorchTitan, PyTorch's native library for large scale training, also runs on Trainium through a companion package.

The numbers AWS shared, and the ones it did not

The announcement is long on features and short on hard performance data. The clearest number comes from Deepgram, a speech to text company. Its CTO Adam Sypniewski says in the PyTorch blog that, in the benchmarks Deepgram evaluated, Trainium processes about 60 percent more input audio seconds per dollar than its existing GPU baseline. In TechCrunch, AWS's Colin Brace cited the same customer with a higher figure, 70 percent more audio per dollar at production grade accuracy. The two numbers may come from different tests or different points in time, but neither is an independent measurement, and Deepgram's port took a multi month engagement with the AWS Generative AI Innovation Center. That is a long way from a one line change.

AWS does deserve credit for publishing its test coverage honestly. It runs a "blind sweep" of 100,000 Hugging Face models, which it says cover 99.9 percent of real world usage by download count, and checks whether each one loads, runs forward and backward passes and compiles. It did not give a headline pass rate in the text, only a chart and a promise to close the remaining gaps in future releases.

The PyTorch upstream test results are more concrete. The August 2026 run produced about 154,000 test outcomes. Roughly 69,000 were skipped. About 30,000 of those are skipped by PyTorch itself because they only apply to CUDA or CPU. The other roughly 39,000 are deliberate Neuron skips, and the biggest group, about 22,000, covers operators that do not yet have a native Neuron implementation, such as corner cases of elementwise, reduction and advanced indexing operations and fused LSTM and GRU cells.

AWS also lists features it chose not to support in the first release because customers did not ask for them: sparse tensors, NestedTensor, forward mode autograd, arbitrary attention masks in native multi head attention, per channel quantization, bilinear and bicubic upsampling, and dynamic or symbolic shapes. FP8 and TorchAO are only in private beta. For many mainstream transformer workloads none of that matters. For research code that leans on dynamic shapes or unusual operators, it might.

Who is already using it

The announcement names three customers with quotes. Cartesia, which builds real time voice models on state space architectures like Mamba, says it ran its standard inference pipeline with minimal changes and could focus its own kernel work on the layers that drive latency. Deepgram describes the cost gains above. Hugging Face says TorchNeuron makes Trainium a native backend for Transformers and Diffusers, which it calls a practical expression of its commitment to silicon neutrality.

Beyond that, AWS lists partners such as Anyscale, Decart, Modular and Luminal, and thanks a long list of organizations for feedback, including Anthropic, Meta, Databricks, ByteDance, Salesforce, Rivian and several universities.

The Anthropic name is worth a moment. Anthropic is Amazon's biggest AI partner and one of the largest users of Trainium through Project Rainier, the giant Trainium cluster AWS built for it. Being thanked for feedback is not the same as running production work on TorchNeuron, and the blog does not claim that. But it does show that the teams training frontier models on Trainium were part of shaping the tool.

What it means for developers

For a startup or research team, the practical change is the cost of an experiment. Before, testing Trainium meant learning a separate toolchain. Now a team can take an existing PyTorch training script, switch the device, and see what happens. If performance per dollar is good, they can dig deeper with `torch.compile` and custom kernels. If it is not, they have lost little time.

There are also caveats. A public beta is a public beta. The model sweep and test numbers show that coverage is broad but not complete. Eager mode on Trainium is slower by design, so the "one line" story is mostly about getting code to run, while getting it to run fast still means compiling and tuning. And all of this only works on AWS, since Trainium is not sold as a standalone chip. Moving from Nvidia to Trainium reduces dependence on one chip vendor while increasing dependence on one cloud.

AWS argues that is a feature, not a bug. Brace told TechCrunch that on AWS "the infrastructure to go from prototype to production is already there," and that using the same chip for training and inference improves the economics of a full workload, even if AWS does not offer the cheapest price per token. Amazon's own post adds that large customers can shift capacity from inference to training during off peak hours instead of leaving specialized hardware idle.

What it means for Nvidia

In the short term, very little changes. Nvidia's GPUs remain the default for most frontier training, its supply is still tight, and CUDA has years of tuning that no backend can match overnight. But the long term threat is real. If PyTorch becomes truly hardware neutral, then the decision about which chip to use becomes a question of price and availability rather than a question of how much code must be rewritten. That is a fight Nvidia has avoided for a long time, and it is the fight every cloud with its own chips, from Amazon to Google to Microsoft, wants to have.

TechCrunch raised a fair counterpoint: one of the points of PyTorch is to hide hardware details from developers. Pushing hardware choice deeper into the framework could give developers more options, or it could give them more configuration headaches. The answer will come from independent benchmarks and from how many teams outside Amazon's partner list actually move workloads.

Dany's take

This is the right attack on Nvidia. Faster chips alone never broke the CUDA moat, because the moat was always the code. By plugging into PyTorch's official device extension point and open sourcing the whole backend, Amazon is going after the switching cost, which is the thing that really keeps teams on Nvidia. I like that AWS published skipped test counts and a list of unsupported features instead of pretending everything works. What I do not like is that the only customer number comes in two versions, 60 and 70 percent, from sources tied to AWS, and that the star example needed months of help from AWS engineers. Before anyone calls this a CUDA killer, I want to see independent training and inference benchmarks on Trainium3 against current Nvidia hardware, with real models and real prices.

The bottom line

AWS released TorchNeuron, an open source native PyTorch backend that lets existing PyTorch code run on Trainium2 and Trainium3 with a one line device change, and published it as a public beta on 20 October 2026. It supports eager mode, `torch.compile`, distributed training, custom kernels and the Hugging Face libraries, and it is the most direct move yet against the software advantage that keeps most AI developers on Nvidia. The first release still skips about 39,000 PyTorch tests on purpose, leaves some features out and offers only AWS linked performance claims. If independent tests confirm the price performance story, choosing an AI chip could start to look a lot more like choosing a cloud instance than like rewriting your code.