Amazon’s TorchNeuron Puts Trainium One Device String Away From Your PyTorch Code

The best number in Amazon’s TorchNeuron launch comes from a customer who spent months, with AWS’s own specialists, to earn it. Deepgram CTO Adam Sypniewski says Trainium processes about 60% more input audio seconds per dollar than the company’s existing GPU baselines, after what he calls “a multi-month engagement” with the AWS Generative AI Innovation Center to port its model. Amazon announced TorchNeuron on October 19 as a native PyTorch backend for Trainium2 and Trainium3, and on October 20, the opening day of PyTorch Conference North America in San Jose, the code went public as an Apache 2.0 beta in the aws-neuron/torch-neuronx repository.

Amazon’s headline promise is that existing PyTorch training and inference code runs on Trainium after a configuration change, with no rewrite. “If you know PyTorch, then you know Trainium,” Colin Brace, vice president of technology software at Annapurna Labs, told GeekWire. Both stories can be true at once, and the space between them is what operators should measure before they move a job off NVIDIA, especially now that Google shipped its own native PyTorch backend for TPUs the same day.

One Device String Instead of a Second Codebase

Until now, running PyTorch on Trainium meant calling interfaces specific to Amazon’s chips, Brace told GeekWire, and Hugging Face’s Jeff Boudier says its models needed Neuron-specific software and configuration. TorchNeuron registers Trainium through PrivateUse1, PyTorch’s official extension point for out-of-tree hardware, so neuron shows up as an ordinary device the same way cuda does. The PyTorch blog deep dive shows a training loop where the only edit is model.to('neuron'), and the repo’s quick start never imports torch_neuronx at all, because PyTorch loads the backend automatically once the package is installed.

Eager mode dispatches each operation as Python reaches it, so prints and breakpoints behave like they do on a GPU, but each op compiles its own Neuron executable and moves its intermediates through HBM. Amazon offers an adaptive eager mode that fuses queued operations to claw back speed, and the PyTorch blog says it switches off when you call torch.use_deterministic_algorithms(). Production code goes through torch.compile with the Neuron backend, where TorchDynamo and AOTAutograd capture the graphs, Torch-MLIR lowers them to StableHLO, and the Neuron compiler emits the executable without TorchInductor in the loop.

Distributed jobs use the stock APIs, with DDP, FSDP, DTensor and tensor parallelism running collectives over Trainium’s dedicated collective cores inside a node and over Elastic Fabric Adapter between nodes. Hugging Face has made Trainium a supported device upstream in Transformers, Diffusers and TRL, so device_map="neuron" loads weights straight onto the chip, and Amazon’s vLLM Neuron plugin builds on the same compile support.

No Rewrite Still Means Version Pins and Core Counts

The install is the first catch. The repo README wants a Trn2 or Trn3 instance, Python 3.12 through 3.14, PyTorch 2.14 and torch-neuronx 2.14 from Amazon’s Neuron pip repository, and it warns that earlier torch-neuronx versions are the older XLA-based product with the same package name. The README tells you to keep those pins, so put them in your requirements file before an unpinned install surprises someone on your team.

Sharding is the second catch. Amazon’s Neuron docs make the NeuronCore, not the chip, the unit of distribution, so tensor parallel jobs need their core counts set for Trainium while FSDP needs no configuration change. The docs now list pipeline parallelism through torch.distributed.pipelining and mixed precision through torch.autocast(device_type="neuron"), and FP8 and MXFP8 need Trainium3, while the PyTorch blog says TorchAO and FP8 stay in private beta until the next release.

The coverage gaps are spelled out too. Amazon’s August run of PyTorch’s upstream suites produced about 154,000 test outcomes, and roughly 39,000 were deliberate Neuron skips, with the biggest group being about 22,000 tests for operators that have no native Neuron lowering yet. Features Amazon left out at launch include sparse tensors, NestedTensor, forward-mode autograd and dynamic or symbolic shapes. The docs have caught up with the launch since Monday, and they now call TorchNeuron a public beta in the Neuron 2.33 release and describe a compile-only mode that builds executables on a host with no Trainium attached, which should let CI precompile without renting a Trn2, even though the README still lists a Trainium instance as the install requirement.

Deepgram’s 60% Came After a Multi-Month Port

More than 30 companies and internal Amazon teams used TorchNeuron during a private beta that GeekWire dates to late last year. Deepgram’s figure is the only price-performance number with a named customer attached, and Brace told GeekWire the company now runs Trainium in production. Sypniewski qualifies it as “in the benchmarks we evaluated” without naming the GPU, instance pricing or workload, and our team at TechForward has not seen those benchmarks or run TorchNeuron ourselves.

Cartesia cofounder Arjun Desai describes a lighter first step, running the company’s standard inference pipeline “with minimal changes” and then hand-writing NKI kernels only for the layers that drive end-to-end latency. Put the two quotes side by side and the first run looks cheap, while the engineering hours go into performance work, either NKI kernels at Cartesia or an AWS-assisted port at Deepgram.

Amazon says it continuously runs 100,000 Hugging Face models, which it claims cover 99.9% of real-world usage by download count, and records which operators fell back to CPU. The docs explain what that column tells you. In eager mode an operator without a Neuron implementation runs on the host with a warning, so a model can pass a compatibility sweep while burning Trainium hours you pay for, and under torch.compile the same gap fails the compile by default.

Google Shipped the Same Exit From CUDA the Same Day

Google open-sourced TorchTPU on October 20 at the same conference, with nearly the same design, eager mode for debugging, torch.compile for production, upstream Hugging Face support through device_map="tpu", and TorchTitan recipes. The PyTorch/XLA repository now calls itself deprecated and says it will be archived on November 15, 2026. Meta’s Jana van Greunen and Hugging Face’s Jeff Boudier are quoted in both launches, so Meta’s PyTorch team and Hugging Face are lending their names to both backends.

The conference program billed AWS’s Maen Suleiman for a five-minute sponsored keynote on “Trainium’s Journey to Native PyTorch” on Tuesday morning, covering eager mode, torch.compile, TorchTitan, TorchAO and Transformers v5, and Google’s Bill Jia also held a keynote slot on workload fungibility. Microsoft said on October 12 that Maia will become the primary inference system for its products, with multi-gigawatt deployments starting in 2027, which puts the third big cloud on the same road.

Amazon needs that software gap closed more than most, because its biggest Trainium users so far are labs with financial ties to it. GeekWire notes that Anthropic trains and serves Claude on Trainium and that OpenAI agreed in February to use 2 gigawatts of Trainium capacity. My write-up of Tenstorrent’s developer day opened on research that stalls the moment it leaves CUDA, and Trainium2 clusters were already behind AWS AI Factories in our re:Invent 2025 roundup.

Run Your Own Loop on Trn2 Before Anyone Quotes You 60%

Deepgram got its 60% after months of porting with AWS in the room. TorchNeuron says the next team starts from its own unmodified training loop instead, and Google made the same promise for TPUs that day. That pairing is why this launch stuck with our team at TechForward.

The fine print is public, with operator gaps, NeuronCore sharding and version pins, and the customer numbers are Amazon’s to defend. The code is live now. Change the device string on your current loop, profile it with torch.profiler, and check what fell back to CPU before you compare dollars.

Clone torch-neuronx and read Amazon’s TorchNeuron launch post, then the PyTorch blog deep dive and the Neuron docs for the exact changes. Grab torchtitan-neuron if you want Amazon’s Llama 3 8B recipe on a single trn2.48xlarge, and keep Google’s TorchTPU launch open so you can compare the two backends side by side.

spot_img

More from this stream

Recomended