What Is TorchNeuron? Amazon's Free Tool to Run PyTorch on Trainium (and Take On Nvidia)

TorchNeuron is a "native backend" for PyTorch, the most popular framework for building AI models. Per the official PyTorch blog post written by AWS engineers, it registers Trainium as a standard PyTorch device, so normal PyTorch tools, from training loops to distributed training to the profiler, work on it as written. It supports Trainium2 and Trainium3, and the code goes to github.com/aws-neuron/torch-neuronx.

What is CUDA lock-in? (The non-technical version)

Think of AI chips like game consoles. Nvidia's console is CUDA, its software platform. For over a decade, almost every AI tool, tutorial and research paper has been written assuming you own that console.

Amazon builds its own chips (Trainium) and says they're cheaper to run. But if your model code, your libraries and your engineers' know-how all assume Nvidia, moving costs months of rewriting and testing. The Next Web describes exactly that switching cost as what TorchNeuron is meant to cut.

PyTorch sits on top of the chips. If PyTorch treats Trainium as a first-class citizen, then most developers never touch CUDA directly anyway, and the "console" matters less. That's the bet.

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

This is the confusing bit, because the package name hasn't changed.

Before, PyTorch on Trainium went through PyTorch/XLA, a layer that traced your model into a graph and compiled it. It worked, but it felt different from normal PyTorch: debugging was awkward and code often needed Neuron-specific changes. AWS's own docs say PyTorch 2.9 was the last XLA-based release and that the move to native PyTorch was first announced at re:Invent in December 2025.

Does it really work with a single code change?

For the simple case, yes, according to AWS. This is the example from the PyTorch blog:

# Run a training loop on Trainium
model = MyModel().to('neuron')   # the only change
optimizer = torch.optim.AdamW(model.parameters())
for batch in dataloader:
    optimizer.zero_grad()
    loss = criterion(model(batch), targets)
    loss.backward()
    optimizer.step()

Code written against torch.accelerator, PyTorch's device-agnostic API, is said to run with no changes at all. For compiled models you name the backend once with torch.compiler.set_default_backend("neuron"). Hugging Face Transformers can load models straight onto the chip with device_map="neuron".

Editor's view: "One line" is true for the demo, and it's a big step. Getting a model to run is different from getting it to run fast. Expect real teams to still spend time on profiling and custom kernels for the parts that matter.

The technical deep dive: what's under the hood?

Here's how the pieces fit, based on the PyTorch blog and AWS docs.

Eager mode. Every operation runs immediately, so Python breakpoints and prints work. Each op compiles to its own Neuron Executable File Format (NEFF) and launches on the chip, which is flexible but costs time per launch. The new adaptive eager mode (experimental) tries to fuse operations and reduce launch overhead without code changes.

torch.compile. TorchDynamo captures an FX graph, AOTAutograd traces the backward pass, the Neuron backend lowers it through Torch-MLIR to StableHLO, and the Neuron compiler emits a NEFF. FlexAttention is supported natively. AWS's vLLM Neuron plugin for serving LLMs is built on this same path.

Ops and data types. AWS claims broad ATen operator coverage across text, vision, video, state-space and world models. BF16, FP32 and FP16 work, with torch.autocast. FP8 and TorchAO are in private beta.

Distributed. DDP, FSDP2, SimpleFSDP, tensor and pipeline parallelism, ZeRO, DTensor and DeviceMesh, plus standard collectives (all_reduce, all_gather, reduce_scatter) and send/recv. Inside a server, collectives run on Trainium's dedicated collective cores; between servers, over Elastic Fabric Adapter.

Streams. PyTorch streams and events map onto Trainium hardware queues. One catch: non-default streams share a single hardware queue, so they keep ordering semantics but don't each get their own engine. Errors also surface at the next sync point, not the op that caused them.

Custom kernels (NKI). The Neuron Kernel Interface lets you write low-level kernels and register them as ordinary PyTorch custom ops with autograd. AWS also publishes NKI-lib, with kernels such as FlashAttention and MoE blockwise ops.

Profiling. Standard torch.profiler with a Neuron config, viewable in Neuron Explorer, which can also export Parquet for automated analysis.

What doesn't TorchNeuron support yet?

AWS is unusually upfront here. In its August 2026 run of PyTorch's upstream test suites (about 154,000 outcomes), around 39,000 tests were deliberately skipped for Neuron reasons. The biggest group, roughly 22,000, covers operators without a native Neuron lowering yet.

Features AWS says it chose not to support in the first release:

• Dynamic and symbolic shapes
• Sparse tensors and NestedTensor
• Forward-mode autograd (JVP)
• Arbitrary attention masks in native multi-head attention (no-mask and causal work)
• Per-channel quantisation
• Bilinear/bicubic upsampling and torch.signal windows
• Fused LSTM/GRU cells (in the operator-gap group)

Also note: it's Trainium2 and Trainium3 only, and it's a beta. Some AWS doc pages still describe a closed beta, which the 20 October public release supersedes.

Editor's view: Dynamic shapes is the one to watch. Lots of real inference code (variable-length inputs, some agent workloads) leans on it.

Amazon vs Nvidia AI chips: does this change the fight?

It narrows one gap, not all of them. Nvidia's lead also comes from raw performance, mature tooling and the fact that every cloud sells its GPUs.

What TorchNeuron changes is the cost of trying Trainium. Amazon's evidence so far comes from customers it chose to quote, so treat it as vendor-supplied:

• Deepgram CTO Adam Sypniewski said in the PyTorch post that, in their benchmarks, Trainium processed "about 60% more input audio seconds per dollar" than their GPU baselines.
• Cartesia co-founder Arjun Desai said they ran their standard inference pipeline "with minimal changes to our PyTorch code."
• Hugging Face VP Product Jeff Boudier said TorchNeuron makes Trainium "a native backend for Transformers and Diffusers."
• Meta's Jana van Greunen, Director of PyTorch Engineering, said native Trainium support "expands that choice."

The commercial context is big. The Next Web reports Amazon says its chip business has passed a $25 billion annual run rate, and that Andy Jassy suggested in April it could be worth $50 billion. If you're new to Amazon's AI stack, our Bedrock explainer shows where Trainium fits for people who don't write PyTorch at all.

What are developers saying?

At the time of writing I couldn't find substantive discussion on r/MachineLearning, r/aws or Hacker News; the open-source repo only goes public on 20 October. The Next Web also noted it found no independent analysts quoted. So here's what developers are likely to ask once they can try it:

• How does real throughput compare with an H100/B200 on my model, not a showcase?
• How many ops silently fall back to CPU, and how do I spot them?
• Will it be upstreamed into PyTorch itself? (AWS docs say that's the goal.)
• Does it work with vLLM and SGLang out of the box?
• Is it truly portable, or will I end up writing NKI kernels anyway?

How do I try TorchNeuron?

1. Get access to a Trainium2 or Trainium3 instance on AWS (for example trn2).
2. Install with pip install torch-neuronx from the AWS Neuron pip repository, following the Neuron setup docs.
3. Change your device to neuron, run, then profile.

New to the cloud side? Start with what AWS is and how it bills.

Why this launched the same week as Amazon's nuclear news

Amazon's chip push and its power push are the same story. More Trainium means more data centres, which means more electricity. On the same day, Amazon announced funding for 12 small nuclear reactors in Indiana. We break that down in Amazon to Fund 12 Small Nuclear Reactors in Indiana. Investors may hear both themes on the Q3 2026 earnings call.

FAQ

What is TorchNeuron?

An open-source native PyTorch backend from AWS that lets PyTorch code run on Trainium2 and Trainium3 chips with standard PyTorch APIs.

Is TorchNeuron free?

Yes, the software is open source under Apache 2.0. You still pay for the Trainium instances you run it on.

Do I need to rewrite my PyTorch code for Trainium?

For many models, no. The basic change is .to('neuron') instead of .to('cuda'). Models using unsupported features, like dynamic shapes or sparse tensors, may need work.

Is TorchNeuron the same as torch-neuronx?

It uses the same package name, but it's a new native backend that replaces the older PyTorch/XLA-based approach.

Which chips does TorchNeuron support?

Trainium2 and Trainium3. It does not cover older Trainium1 or Inferentia chips, per the current docs.

Does TorchNeuron support torch.compile and FSDP?

Yes. It supports torch.compile with backend="neuron", and distributed APIs including DDP, FSDP2, tensor parallelism and pipeline parallelism.

Is TorchNeuron faster than Nvidia GPUs?

There's no independent benchmark yet. The performance and cost claims so far come from AWS and customers it quoted, such as Deepgram.

When is TorchNeuron available?

As a public beta from 20 October 2026, launched at PyTorch Conference North America.

Sources

• PyTorch blog: Introducing TorchNeuron: Native PyTorch backend for AWS Trainium
• AWS Neuron docs: TorchNeuron overview
• AWS Neuron docs: PyTorch on Neuron, about
• GitHub: aws-neuron/torch-neuronx
• The Next Web: Amazon's TorchNeuron lets PyTorch code run on AWS Trainium chips
• Techzine: TorchNeuron brings PyTorch code to AWS Trainium without modifications
• AWS: Trainium

Featured image: AWS Trainium3 chip by Annapurna Labs. Amazon. Source