Amazon's TorchNeuron Lets PyTorch Code Run on Trainium With One Config Change

Amazon's Annapurna Labs just took its most direct shot at NVIDIA yet: TorchNeuron, an open-source tool that registers AWS Trainium chips as a native PyTorch device, so existing PyTorch code runs on Amazon's custom AI silicon with a single configuration change instead of months of CUDA porting. The TorchNeuron PyTorch Trainium integration was unveiled on October 19 and enters public beta on October 20 at the PyTorch Conference in San Jose, released under the permissive Apache 2.0 license. Wall Street noticed — Amazon shares climbed as much as 5% intraday, with Morgan Stanley naming the stock a top pick.

TorchNeuron for PyTorch on Trainium: how the switch works

Until now, pointing a PyTorch workload at Trainium meant real engineering work: modifying PyTorch code to use Amazon-specific programming interfaces. TorchNeuron replaces that with a device setting. It builds directly on PyTorch's own accelerator-agnostic extension points — the PrivateUse1 device backend and the torch.accelerator APIs — so the model code, libraries, and tools developers already use run on Trainium without modification. Installation is a pip command: pip install torch-neuronx from the AWS Neuron package repository, with the project repo going live at github.com/aws-neuron/torch-neuronx.

It supports both Trainium2 and Trainium3 across training and inference, and covers the features developers actually expect: eager execution, PyTorch's distributed training APIs, and integration with TorchTitan — the popular open-source training engine — including fully sharded data parallelism. For the performance-obsessed, there's the Neuron Kernel Interface (NKI), Amazon's custom language for low-level kernel optimization, plus a set of pre-packaged NKI operators: building blocks for mixture-of-experts architectures, memory reduction, and quantization.

Notably, Amazon also contributed work back into PyTorch itself — RFCs and pull requests that strip out hard-coded CUDA assumptions from the framework. That makes PyTorch more portable for everyone, not just Trainium users, and it's the kind of upstream contribution that buys goodwill with the developer community.

Why NVIDIA should be paying attention

Around 85% of AI researchers and developers write their code in PyTorch, according to Amazon. For more than a decade, NVIDIA has invested heavily in CUDA — the software layer that makes PyTorch code run fast on NVIDIA GPUs. Developers rarely write CUDA themselves; they write PyTorch, and NVIDIA's software handles what happens underneath. The result is a default: PyTorch code effectively assumes NVIDIA hardware.

Colin Brace, vice president of software for Annapurna Labs, framed TorchNeuron as a direct attack on that friction. "This has been about meeting customers where they're at," he told Reuters — fewer steps between an AI developer and Amazon's chips. He says more than 30 early customers, spanning robotics, world models, voice AI, speech, drug discovery, and finance, are already using the software and reporting steep cost savings. Named adopters include Cartesia AI, a voice-agent startup running inference on Trainium, and the Lieber Institute for Brain Development.

Amazon isn't the only cloud giant chipping at the CUDA default — Google is building better PyTorch tooling for its TPUs, and PyTorch users can already move workloads to AMD GPUs or Apple silicon with a configuration change. But Amazon's pitch bundles the chips with the whole AWS cloud, and the economics argument is the sharp end: if Trainium's price-performance holds, "try our chip" stops being a science project and becomes a line item.

The silicon behind the software

Trainium3 is the flagship: eight compute cores, a 154GB pool of onboard memory, running in custom AWS servers that pack up to 144 accelerators apiece — a single system good for up to 362 quadrillion calculations per second in the MXFP8 format AI models favor. In his April shareholder letter, CEO Andy Jassy said Amazon's chip business, Trainium plus Graviton, runs at more than $20 billion in annual revenue — and would be worth roughly $50 billion as a standalone chip company. He's even floated selling racks of Trainium to third parties. GeekWire reports Amazon is also in talks to acquire Decart, a startup whose software helps move NVIDIA-built models to alternative chips, in a deal that could value it near $7 billion.

The honest caveats

A configuration change gets you started, not finished. Model teams still have to validate behavior, operator coverage, distributed execution, debugging, and real-world cost on their own workloads — the one-line claim describes the on-ramp, not the whole deployment. This is a public beta, and Trainium's biggest users so far are labs with deep financial ties to Amazon: Anthropic trains and runs Claude on the chips, and OpenAI agreed in February to take 2 gigawatts of Trainium capacity as part of Amazon's $50 billion investment.

Even so, the direction is unmistakable. TorchNeuron lands in a busy season for Python-based AI tooling — GitHub's AI code scanning just went GA on pull requests, and Python 3.15 is out as 3.10 hits end of life. For years, "try our AI chip" meant "rewrite your stack." TorchNeuron turns it into a config flag. If the cost savings hold at scale, NVIDIA's software moat just sprang its first serious leak.

Key facts

• Unveiled Oct 19, 2026; public beta Oct 20 at the PyTorch Conference, San Jose
• Apache 2.0 open source; repo at github.com/aws-neuron/torch-neuronx
• Supports Trainium2/Trainium3 for training and inference
• Install: pip install torch-neuronx
• 30+ early customers; Amazon shares up ~5% intraday on the news