Skip to content
audiotale.net
Sign in
NewsNews

Amazon open-sources TorchNeuron, a native PyTorch backend for its Trainium AI chips

AWS says existing PyTorch code runs on Trainium2 and Trainium3 with a device change; the Apache 2.0 code goes public on 20 October, with gaps it admits.

audiotale.net team

· 8 min read

Written by AI from the sources linked below. How we use AI

On this page
  1. Key points
  2. What happened
  3. Background
  4. The numbers
  5. What it means
  6. What we don't know yet
  7. What to watch next
  8. Sources

Amazon has announced TorchNeuron, an open-source native PyTorch backend for its Trainium AI chips, and says existing PyTorch code can run on Trainium by changing the device to neuron. The announcement is on the PyTorch project's blog, dated 19 October 2026; the code, the pip packages and a public beta are due on 20 October at PyTorch Conference North America. It matters because the stickiest part of Nvidia's lead has long been software, not silicon: most AI code is written against PyTorch, and PyTorch has historically meant "CUDA first". A well-behaved, open backend would make a second chip family a realistic option for more teams, if it delivers what the blog post describes.

Key points

  • TorchNeuron is a native PyTorch backend for AWS Trainium2 and Trainium3, to be released under the Apache 2.0 licence at github.com/aws-neuron/torch-neuronx as a public beta on 20 October 2026, according to the PyTorch blog.
  • Trainium registers as a standard PyTorch device through the PrivateUse1 extension point, so model.to('neuron') is the only change in the blog's training-loop example; torch.compile(model, backend="neuron") handles production graphs.
  • The first release deliberately omits some features: roughly 22,000 upstream PyTorch tests are skipped because the operators have no native Neuron lowering yet, and dynamic and symbolic shapes, sparse tensors and NestedTensor are not supported.
  • Deepgram, a customer, says Trainium processes about 60% more input audio seconds per dollar than its GPU baselines; that is a vendor-reported figure, and AWS's own case study gives a range of 11 to 70% for a different measurement.
  • Hugging Face says Transformers, Diffusers and TRL now treat Trainium as a native device, with no Neuron-specific code needed to load a model.

What happened

On 19 October 2026 the PyTorch blog published "Introducing TorchNeuron: Native PyTorch backend for AWS Trainium", written by AWS engineers. The post says TorchNeuron is "an open source native PyTorch backend for Trainium", so that model code, libraries and tools "you already use" run on the chip. Amazon published its own announcement, and TechCrunch and The Next Web covered it.

The timing needs care. The blog post says the source code goes live as a public beta "on Oct. 20" at PyTorch Conference North America, and lists the GitHub repository, the documentation and the pip package among resources that "will go live on Oct. 20". We could not open the repository when we wrote this, so everything below about behaviour comes from AWS's and PyTorch's own descriptions, not from our own testing. Amazon's announcement page describes the release as already open-sourced at the conference; the PyTorch post, which is the more precise document, uses the future tense.

The technical idea is simple to state. PyTorch lets hardware vendors register a new device through PrivateUse1 and the torch.accelerator APIs. TorchNeuron uses that route, so torch.nn, torch.autocast, torch.distributed and torch.profiler operate on a neuron device without a fork of PyTorch. The post gives three ways to run:

  • Eager mode, where each operation runs immediately, which suits research and debugging. Each operation compiles its own Neuron Executable File Format (NEFF) file and moves intermediates through the chip's high-bandwidth memory, which is slow. The post introduces an experimental "adaptive eager" mode to cut launch cost and fuse operations without code changes.
  • torch.compile, where TorchDynamo captures the model as an FX graph, AOTAutograd traces the backward pass, and the Neuron backend lowers the result through Torch-MLIR to StableHLO before the Neuron compiler produces a NEFF. FlexAttention is supported natively.
  • The Neuron Kernel Interface (NKI), for writing custom operators by hand. AWS also open-sources its optimised kernels as NKI-lib, including FlashAttention and mixture-of-experts kernels.

Distributed training uses the standard APIs: DDP, FSDP2, SimpleFSDP, tensor parallelism, pipeline parallelism, DTensor and DeviceMesh. Collectives run over Trainium's dedicated collective cores inside a node and over the Elastic Fabric Adapter between nodes. Inference runs through two vLLM plugins, one for language models and one for diffusion models. BF16, FP16 and FP32 are supported; TorchAO and FP8 are in private beta and promised for the next release.

Background

TorchNeuron is not new as a name. AWS listed "native PyTorch support through TorchNeuron (private beta)" in its Neuron SDK 2.27.0 release on 22 December 2025. A Neuron announcement dated 16 December 2025 told existing users that "PyTorch 2.9 will be the last version based on PyTorch/XLA", and that PyTorch 2.10 and later would get native support through TorchNeuron.

That is the real story behind this week's news. Until now, running PyTorch on Trainium meant going through the XLA route, a translation layer originally built for Google's TPUs, which lacks the eager, debuggable feel of PyTorch on a GPU. Moving to the same extension mechanism used by other accelerators removes a layer of indirection and, AWS hopes, a reason for engineers to say "it works on CUDA, we will port it later".

Amazon frames the strategic goal plainly. Its announcement says about 85% of AI researchers and developers write their code in PyTorch (an Amazon claim we have not verified), and that it is contributing changes upstream to remove hard-coded CUDA assumptions. The PyTorch post lists 26 pull requests to PyTorch and TorchTitan, among them changes to test suites that assumed CUDA. The Next Web, citing Amazon, reports that the company's chip business has passed a $25 billion annual revenue run rate and that Andy Jassy said in April it could be worth $50 billion; those are company statements, not audited figures.

The numbers

AWS publishes two kinds of evidence for coverage, and both are worth reading closely.

The first is a "blind sweep" of Hugging Face models. The post says AWS continuously runs 100,000 models, covering 99.9% of real-world usage by download count, on Trainium in a reduced configuration, and records whether each loads, runs forward and backward, compiles, and whether operators ran on the device or fell back to the CPU. The results appear only as a graph (as of September 2026) that we cannot reproduce here, and AWS says it is "working on closing the remaining gaps". We therefore cannot tell readers what share of models run end to end.

The second is the PyTorch upstream test suite, and here the post gives numbers. The August 2026 run produced about 154,000 test outcomes, of which roughly 69,000 were skipped with a stated reason:

type: bar
title: How the August 2026 PyTorch upstream test run on TorchNeuron breaks down
unit: thousand test outcomes (approximate)
source: PyTorch blog, "Introducing TorchNeuron" (AWS authors), 19 October 2026
highlight: Skipped: no native Neuron lowering
All outcomes | 154
Skipped (total) | 69
Skipped by PyTorch's own checks (CUDA-only, CPU-only) | 30
Skipped deliberately for Neuron | 39
Skipped: no native Neuron lowering | 22

About 30,000 of the skips come from PyTorch's own eligibility checks, such as tests that only apply to CUDA or CPU, and say nothing about Neuron. The other roughly 39,000 are deliberate Neuron skips. The largest group, about 22,000, covers operators without a native Neuron lowering: corner cases in elementwise, reduction, foreach, advanced-indexing and linear-algebra operators, and fused LSTM and GRU cells. The rest are features TorchNeuron "explicitly chose not to support in the first release as customers didn't request them": sparse tensors, NestedTensor, forward-mode autograd, arbitrary attention masks in native multi-head attention (no-mask and causal are supported), per-channel quantisation, bilinear and bicubic upsampling, and dynamic and symbolic shapes.

AWS does not publish a pass rate for the tests that were not skipped, so we cannot say how many of the remaining roughly 85,000 outcomes passed.

Performance evidence is thinner and comes from customers. Deepgram's chief technology officer, Adam Sypniewski, says in the PyTorch post that "in the benchmarks we evaluated" Trainium processes about 60% more input audio seconds per dollar than its existing GPU baselines. AWS's separate Deepgram case study, published the same day, reports that offline batch workloads on the Nova-3 speech model processed "11–70 percent more audio per dollar than the compute options Deepgram had been running", using Trainium2 trn2.3xlarge instances, after several months of porting work. TechCrunch reported the top of that range as 70%. The two AWS-published figures differ, neither names the GPU baseline, and the case study says the numbers reflect instance pricing at the time of evaluation. Treat all of them as vendor-reported.

What it means

For developers already on PyTorch, the promise is a lower porting cost, not zero. The blog's own example is a one-line device change, but it is also candid that the first release lacks dynamic shapes, which many production inference stacks lean on, and that custom code may need NKI kernels to reach good speed. Cartesia, a voice-AI start-up, says in the post that its standard inference pipeline ran "with minimal changes", while its team now concentrates on hand-writing NKI kernels for the layers that dominate latency. That is probably the realistic pattern: easy start, tuned hot paths.

For Nvidia, the threat is incremental rather than sudden. CUDA's advantage includes a decade of kernels, libraries and debugging tools. A native backend shrinks the gap in the framework layer, and Hugging Face's support matters because so many open models start there. Hugging Face's Jeff Boudier calls it "a practical expression of our commitment to silicon neutrality". Models that depend on CUDA-only code paths will still need work.

For AWS customers, it adds negotiating leverage and a cost option. AWS says the same Trainium capacity can serve training and inference, so teams could move capacity between them. Whether the price-performance holds outside the customers AWS chose to feature is the open question.

For the open-source ecosystem, an Apache 2.0 backend that registers through the public extension point, plus upstream fixes for CUDA assumptions, benefits other accelerator makers too. The PyTorch post thanks the PyTorch team at Meta and the PyTorch Foundation, and Meta's director of PyTorch engineering is quoted supporting the work.

What we don't know yet

  • Whether the repository and packages match the blog's description. They were not public when we wrote this.
  • What share of the 100,000-model sweep runs end to end, and how many operators still fall back to the CPU.
  • The pass rate of the upstream tests that were not skipped.
  • Which GPUs Deepgram compared against, at what prices, and why AWS's two published figures differ.
  • When dynamic shapes, FP8 and the adaptive eager mode will be stable.
  • How fast TorchNeuron tracks new PyTorch releases, which was the weak point of earlier out-of-tree backends.

What to watch next

Watch the 20 October release itself: the licence file, the issue tracker, how quickly the first bug reports are answered, and whether independent users post benchmarks. Watch for third-party results from teams outside the AWS launch list, and for the next release's promised FP8 and TorchAO support. Also watch the upstream PyTorch pull requests; if the CUDA-assumption cleanups merge, other backends benefit as well. For background on the hardware race this sits inside, see our coverage of Intel's new data-centre GPU and Nvidia's new open models.

Sources

SaveStart the discussion

Written by

audiotale.net team

The official account of the people who run audiotale.net: announcements, conversation starters and the newsroom's drafts. Posts written with AI say so. Every other account here is a real person.

Read next