TorchNeuron is Amazon Web Services’ new open-source software that lets code written in PyTorch run natively on AWS’ own Trainium AI chips with close to a one-line change. Announced on October 19, 2026 and opened as a public beta on October 20, TorchNeuron targets the biggest reason most AI teams still train on Nvidia hardware: the software. Instead of rewriting code to call Amazon-specific chip interfaces, developers can now point their existing PyTorch models at a device simply called neuron. If you know PyTorch, Amazon’s pitch goes, you now know Trainium.
For a market where the hardware is the easy part and the software stack decides who wins, this is one of the more consequential infrastructure releases of the year. Here is what TorchNeuron is, how it works, what it costs to run, and what it still cannot do.
What TorchNeuron actually is
TorchNeuron is a native PyTorch backend for AWS Trainium2 and Trainium3 chips, developed by Annapurna Labs, Amazon’s in-house chip design unit. It is released as open source under the Apache 2.0 license on GitHub (aws-neuron/torch-neuronx) and installed with a simple pip install torch-neuronx from the AWS Neuron package repository.
Before TorchNeuron, running PyTorch code on Trainium meant modifying that code to call programming interfaces specific to Amazon’s chips. That friction kept most teams on Nvidia’s CUDA platform, which the industry’s AI software was largely built around. TorchNeuron removes that step: PyTorch registers Trainium as a normal device, so the line model.to('cuda') effectively becomes model.to('neuron'), and the same code runs.
Under the hood, TorchNeuron plugs into PyTorch’s official PrivateUse1 extension mechanism, the standard way for new hardware to register itself as a device. Because of that, standard PyTorch APIs such as torch.nn, torch.autocast, torch.distributed and torch.profiler operate on the neuron device as written. Code already written against PyTorch’s device-agnostic torch.accelerator API runs unmodified. Compiled models use backend="neuron" with torch.compile.
What it supports today
The beta covers the pieces a serious training or serving setup needs:
- Execution modes: eager mode for interactive development and debugging, an experimental “adaptive eager” mode, and
torch.compilewith a dedicated Neuron backend for production speed. - Distributed training: DDP, FSDP2, DTensor and DeviceMesh, plus tensor, pipeline and context parallelism.
- Custom kernels: the Neuron Kernel Interface (NKI) plus the open-source NKI-lib, for teams that need to write low-level operators tuned for the hardware.
- Ecosystem integrations: Hugging Face Transformers, Diffusers and TRL, TorchTitan tutorials that cover pre-training Llama, DeepSeek and GPT-OSS models, and vLLM Neuron plugins for model serving.
- Hardware: AWS Trainium2 and Trainium3 instances. Note that FP8 and MXFP8 precision formats need Trainium3.
Amazon says 85 percent of AI researchers and developers write their code in PyTorch, which explains why the company targeted that framework rather than building another proprietary toolchain. TorchNeuron had been in private beta since late 2025 after an initial reveal at re:Invent that December; the October 20 release at PyTorch Conference North America is the first version any developer can try.
How much does running Trainium cost?
The software itself is free. You pay for Trainium compute time the same way you pay for GPU instances. For a sense of the economics, one third-party breakdown priced a Trainium2 chip-hour at $2.235 in EC2 Capacity Blocks against $5.970 for an Nvidia H100. Amazon’s own announcement materials cite Deepgram, a voice AI company, saying Trainium handles about 60 percent more audio per dollar than its GPU baselines after a multi-month port done with AWS’ Generative AI Innovation Center.
Those numbers come from Amazon and its partners rather than independent benchmarks, so treat them as directionally encouraging rather than verified. The real cost comparison for any team depends on how much custom tuning their workload needs, which leads to the biggest caveat.
The catch: moving code is easy, getting speed is not
TorchNeuron genuinely removes the rewrite step, but peak performance on Trainium still comes from kernels written specifically for the chip, using Amazon’s NKI language. Those kernels do not run anywhere else. So while the porting friction has dropped sharply, a form of lock-in has simply moved one layer down, into the performance-critical code.
The beta also has real limits in this first release:
- No dynamic shapes yet. Compiled models currently expect static input shapes; a changed shape means recompiling, and native dynamic-shape support is still in development.
- Not all operators are covered. Some PyTorch operations have no Trainium lowering yet, so
torch.compilecan fail on code paths it cannot handle. - Trainium2/3 only. Users on older Trainium1 or Inferentia chips stay on the previous Neuron stack.
- It is a beta. APIs, behavior and performance can change before a stable release.
In practice, the sensible move for most teams right now is a pilot: port one workload, measure end-to-end cost and latency on your own data, and ask AWS which parts of your model would need hand-tuned kernels and who writes them, before committing anything serious.
Why AWS is doing this now
Nvidia’s strongest advantage was never only its chips. It is the software developers already know how to use. Every cloud giant is now trying to break that grip while remaining among Nvidia’s biggest customers. TorchNeuron is Amazon’s answer, and it is not alone:
- Google detailed TorchTPU, a similar project for running PyTorch natively on its TPU chips, and released it on October 21, 2026, one day after TorchNeuron’s beta.
- Microsoft said it plans to run most of its AI services on its homegrown Maia chips rather than Nvidia hardware.
- Challenger chipmakers keep raising serious money to take on the incumbents, as Tenstorrent’s recent $1.5 billion Series E shows. Read our Tenstorrent Series E breakdown for why investors are betting on alternatives.
Nvidia, for its part, keeps shipping new hardware of its own, including the recently launched RTX Spark. The chip race is increasingly a software race, and TorchNeuron is Amazon’s attempt to make its hardware speak the language developers already use.
It also fits AWS’ broader infrastructure push. The company continues to pour money into data centers and custom silicon, the same program we covered in our Amazon Built Together explainer. Cheaper, more efficient AI compute is central to keeping AWS’ largest customers on its platform.
Who should try TorchNeuron, and how to start
TorchNeuron matters most for teams that already train or serve their own models on AWS and are paying Nvidia GPU prices. If you only consume AI through APIs, the chip question barely affects you; the benefits will reach you indirectly through cheaper inference, if at all.
Getting started is straightforward:
- Spin up a Trainium2 or Trainium3 instance (EC2 Capacity Blocks are the main route for reserved Trainium access).
- Install from the AWS Neuron pip repository:
pip install torch-neuronx. - Point your model at the device:
model = model.to('neuron'). - For production throughput, compile with
torch.compile(model, backend="neuron")and keep input shapes static. - Use the TorchNeuron docs and TorchTitan tutorials to walk through a full pre-training or fine-tuning run.
Voice AI company Cartesia, quoted in Amazon’s announcement, ported its code with minimal changes, a plausible best case for models built on standard Transformer architectures. Expect more friction with exotic architectures or heavy custom operators.
Bottom line
TorchNeuron does exactly what Amazon claims: it makes Trainium programmable through PyTorch with roughly a one-line change, free and open source. That is a genuine reduction in switching costs and a real threat to Nvidia’s software moat. But “the code runs” is not the same as “the code runs fast”, and the speed still lives in Trainium-specific kernels that lock you right back in. Treat the public beta as an invitation to benchmark, not as a finished migration path, and make AWS earn your workload with your own numbers.
Frequently asked questions
What is TorchNeuron?
TorchNeuron is AWS’ open-source PyTorch backend for its Trainium AI chips. It registers Trainium as a standard PyTorch device called neuron, so existing PyTorch code runs on the chips with a one-line change instead of a rewrite. It entered public beta on October 20, 2026.
Is TorchNeuron free?
Yes, the software is free and open source under the Apache 2.0 license. You only pay for Trainium compute time on AWS, similar to how you pay for GPU instances.
Which chips does TorchNeuron support?
TorchNeuron supports AWS Trainium2 and Trainium3. It does not support the older Trainium1 or Inferentia chips, which remain on the previous Neuron software stack.
Does TorchNeuron replace Nvidia CUDA?
Not entirely. It removes the need to rewrite PyTorch code when moving to Trainium, which weakens CUDA’s lock-in. But peak performance still requires Trainium-specific kernels written in Amazon’s NKI language, and the beta has limits such as no dynamic-shape support yet.
How do I install TorchNeuron?
From the AWS Neuron pip repository: pip install torch-neuronx. Full setup instructions, hardware requirements and tutorials are in the TorchNeuron documentation on GitHub (aws-neuron/torch-neuronx).