Amazon's TorchNeuron: Your PyTorch Code Now Runs Natively on AWS Trainium, No Rewrite Needed
Amazon has released TorchNeuron, a native PyTorch backend for its Trainium AI chips, as an open-source public beta at PyTorch Conference on October 20, 2026. Existing PyTorch code moves to Trainium with a single device change, Amazon tests it nightly against 100,000 Hugging Face models, and Deepgram says it gets about 60% more audio per dollar than on GPUs. Here is how it works, what is still missing, and why it is aimed squarely at Nvidia's software lock-in.
What happened
On October 19, 2026, Amazon announced TorchNeuron, an open-source, native PyTorch backend for its AWS Trainium AI chips, in a post on the PyTorch Foundation blog and on About Amazon. It goes live as a public beta under the Apache 2.0 licence on October 20, timed with PyTorch Conference North America in San Jose, at github.com/aws-neuron/torch-neuronx.
The pitch is simple: the PyTorch code developers already write runs on Trainium with one configuration change and no rewrite. As Colin Brace, a vice president at Amazon's chip unit Annapurna Labs, told GeekWire: "If you know PyTorch, then you know Trainium." This page is the long-form briefing behind the BSH Technologies Instagram carousel.
1. The one-line switch: .to('neuron')
- The only change. Moving a model to Trainium is
model = MyModel().to('neuron'). Code written against PyTorch's device-agnostictorch.acceleratorAPI runs unmodified. Compiled models name the backend once withtorch.compiler.set_default_backend("neuron"). - How it plugs in. TorchNeuron registers Trainium through PrivateUse1, PyTorch's official extension point for custom hardware. That puts a
neurondevice next to CPU and GPU, sotorch.nn,torch.autocast,torch.distributedandtorch.profilerwork as written. - Both execution modes. Eager mode runs operations immediately, so Python control flow, prints and breakpoints work for research and debugging; a new experimental adaptive eager mode fuses operations to cut launch overhead.
torch.compilecaptures the model into a graph (TorchDynamo, AOTAutograd, Torch-MLIR, StableHLO) that the Neuron compiler optimises for production. - Install.
pip install torch-neuronxfrom the AWS Neuron pip repository. It supports Trainium2 and Trainium3.
2. Why it matters: 85% of AI code is PyTorch
- The reach. Amazon says 85% of AI researchers and developers write their code in PyTorch. Until now, moving that code to a new chip meant rewriting and revalidating work that took months to perfect.
- Before TorchNeuron. Per GeekWire, developers who wanted to run PyTorch on Trainium had to modify their code to call Amazon-specific programming interfaces. That friction kept many teams on Nvidia GPUs.
- Upstream fixes. Amazon says it is contributing changes to PyTorch itself to remove code that hard-wires the framework to specific chips, and lists 26 of its pull requests to PyTorch and TorchTitan.
3. The test suite: 100,000 Hugging Face models every night
- Blind sweep. Amazon continuously runs about 100,000 Hugging Face models, which it says cover 99.9% of real-world usage by download count, in eager and
torch.compilemodes, checking whether each loads, runs forward and backward passes, compiles, and which operators fall back to CPU. - Native in Hugging Face. Trainium is now a supported device upstream in Transformers, Diffusers, TRL and the kernels library. In Transformers v5,
device_map="neuron"(or"auto") loads a model straight onto Trainium. Hugging Face's Jeff Boudier calls it "a practical expression of our commitment to silicon neutrality". - Scale-out. Distributed training uses the standard APIs: DDP, FSDP2, SimpleFSDP, tensor and pipeline parallelism, DTensor and DeviceMesh, with collectives running over Trainium's collective cores and AWS's Elastic Fabric Adapter. TorchTitan trains on Trainium today, and vLLM Neuron serves LLMs.
- Custom kernels. When performance matters, the Neuron Kernel Interface (NKI) lets developers write custom operators, and Amazon has open-sourced NKI-lib with kernels such as FlashAttention and mixture-of-experts blocks.
4. The payoff: what early users report
- Deepgram. The speech-AI company says that in the benchmarks it evaluated, Trainium processes about 60% more input audio seconds per dollar than its existing GPU baselines, after a multi-month port with AWS's Generative AI Innovation Center. It now runs Trainium in production, per GeekWire. (AWS's Brace cited a 70% figure to TechCrunch; we use Deepgram's own 60%.)
- Cartesia. The voice-AI startup behind Sonic and Ink ran its standard inference pipeline with minimal changes to its PyTorch code, then focused its hand-written NKI kernels on the layers that drive latency.
- Others. More than 30 companies and Amazon teams used the private beta, per GeekWire. Prime Video used it to build Prime 360, which reconstructs NBA plays in 3D. The Lieber Institute for Brain Development uses it for drug-discovery models.
5. The bigger picture: Nvidia's software moat
- Lock-in lives in software. Nvidia's dominance comes partly from years of AI software built around its CUDA platform; adapting that code to other hardware can take months of engineering. "The real story is removing the lock-in that kept developers from having a real choice at all," Brace told TechCrunch.
- Not alone. Google is working on TorchTPU to run PyTorch natively on its TPUs, and Microsoft says it plans to run most of its AI services on its own Maia chips, per GeekWire. PyTorch users can already target AMD GPUs and Apple silicon with a configuration change.
- Amazon's chip business. It traces back to the 2015 acquisition of Annapurna Labs. Amazon has since built four chip families, and CEO Andy Jassy wrote in April that its chip business, including Trainium and Graviton, has an annual revenue run rate above $20 billion. Anthropic trains and runs Claude on Trainium.
The catch: what is not there yet
- It is a beta. In Amazon's August 2026 run of PyTorch's upstream test suites, about 39,000 tests were deliberately skipped for Neuron, the largest group (~22,000) for operators without a native Neuron lowering yet.
- Not supported in the first release. Sparse tensors, NestedTensor, forward-mode autograd, dynamic and symbolic shapes, arbitrary attention masks in native multi-head attention, per-channel quantization, and bilinear or bicubic upsampling. TorchAO and FP8 are still in private beta.
- Speed still takes work. "No rewrite" means your code runs; getting top performance can still mean writing NKI kernels, as Cartesia did. There are no independent benchmarks against CUDA systems yet, and the Deepgram figure is the company's own.
What it means if you are learning AI engineering
The safest skill bet just got safer: learn PyTorch itself well, including device-agnostic code (torch.accelerator, no hard-coded "cuda" strings), torch.compile, and distributed training with FSDP. Those habits now pay off on Nvidia, AMD, Apple, Google TPUs and Amazon Trainium alike. For teams, a credible second chip option means better prices and less dependence on a single supplier, so cost-aware model deployment is becoming as valuable as model building.
Sources
- PyTorch Foundation blog: Introducing TorchNeuron, native PyTorch backend for AWS Trainium (Oct 19, 2026)
- About Amazon: Amazon opens its AWS Trainium AI chips to every PyTorch developer, no rewrite needed (Oct 19, 2026)
- AWS Neuron documentation: TorchNeuron, native PyTorch on AWS Neuron
- GitHub: aws-neuron/torch-neuronx
- GeekWire: Amazon challenges Nvidia's software advantage with new tool for its own AI chips
- TechCrunch: Amazon is taking its biggest shot at Nvidia yet with new TorchNeuron dev tool
Frequently asked questions
What is TorchNeuron?
TorchNeuron is Amazon's open-source, native PyTorch backend for its AWS Trainium AI chips. Announced October 19, 2026 and released as a public beta under Apache 2.0 at PyTorch Conference on October 20, it registers Trainium as a standard PyTorch device called neuron, so existing PyTorch code runs with a single device change.
How do I run PyTorch code on Trainium with TorchNeuron?
Install it with pip install torch-neuronx from the AWS Neuron pip repository, then move your model with .to('neuron'). For compiled models, call torch.compiler.set_default_backend('neuron') once or pass backend='neuron' to torch.compile. In Hugging Face Transformers v5, device_map='neuron' loads a model onto Trainium.
Which chips and PyTorch features does TorchNeuron support?
Trainium2 and Trainium3. It supports eager mode, torch.compile, autocast and mixed precision (BF16, FP16, FP32), DDP, FSDP2, tensor and pipeline parallelism, DTensor, torch.profiler and custom NKI kernels. Sparse tensors, dynamic shapes and forward-mode autograd are not supported in the first release, and FP8 is in private beta.
Is Trainium cheaper than GPUs with TorchNeuron?
Early users say it can be. Deepgram reports about 60% more input audio seconds per dollar than its GPU baselines in its own benchmarks. There are no independent benchmarks yet, and top performance can still require custom NKI kernels.
Why is TorchNeuron a challenge to Nvidia?
Much of Nvidia's advantage comes from AI software built around its CUDA platform, which makes switching chips costly. Because about 85% of AI developers write in PyTorch, letting that code run unchanged on Trainium lowers the cost of trying an alternative to Nvidia.
From the blog
View all posts
Anthropic's Own Data: Junior Tech Hires Fell From 27% to 9% as AI Agents Hit 6x Human Hours
On October 16, 2026, Anthropic published a rare inside look at how AI has changed its own workforce: a much smaller share of junior technical hires, AI agents logging six hours for every human hour, five times more code per engineer, and human attention as the new bottleneck. Two days later it pledged 5% of its shares to a new Public Benefit Fund, with AI's impact on jobs as a named focus. Here is what the data says and what it means if you are starting a tech career.
DeepSeek V4.1-Flash-1017: The Open MIT Model That Tops Opus 5 on DeepSWE, Plus a $0.75 Search API
On October 17, 2026, DeepSeek upgraded its Flash model to V4.1-Flash-1017, an open-weight MIT-licensed update that scores 77.0 on DeepSWE v1.1, ahead of Claude Opus 5 in DeepSeek's own table, while finishing tasks faster with shorter outputs. It also opened its in-house search engine to developers at $0.75 per 1,000 calls. Here is the model, the benchmarks, the memory trick, the new API and the price.