All posts

Blog

Amazon releases TorchNeuron so PyTorch code runs on its Trainium AI chips

FindAmNow
AmazonAWSTrainiumPyTorchAI chipsopen source

AWS released TorchNeuron, an open-source native PyTorch backend for its Trainium2 and Trainium3 chips, in public beta with Neuron SDK 2.33 (19 October 2026), and is open-sourcing it under Apache 2.0 at PyTorch Conference on 20 October. AWS says most existing PyTorch code needs only a device change. Coverage gaps remain in the beta.

Amazon Web Services (AWS) has released TorchNeuron, software that lets code written in PyTorch, the most widely used framework for building AI models, run directly on Amazon's own Trainium AI chips. TorchNeuron went into public beta with version 2.33 of AWS's Neuron software kit on 19 October 2026. AWS says it is being released as open source under the Apache 2.0 licence at PyTorch Conference North America on 20 October.

The problem it tries to solve

Most AI software is written with Nvidia GPUs in mind. Moving that code to a different chip has often meant rewriting and retesting parts of it. Until now, PyTorch support on Trainium came through PyTorch/XLA, a layer that turns code into graphs for a compiler. Hugging Face describes that approach as "thinking like a compiler": no step-by-step execution, fixed input shapes only, and a recompile whenever the graph changes.

TorchNeuron instead registers Trainium as a normal PyTorch device, called neuron, through PyTorch's official extension point for new hardware. According to AWS's documentation, moving a model usually means changing .to('cuda') to .to('neuron'). Code written with PyTorch's device-neutral torch.accelerator interface needs no change at all, AWS says.

What it supports

According to the PyTorch blog post by the AWS team and AWS's documentation:

  • Two ways of running code: eager mode, which runs each step immediately and suits research and debugging, and torch.compile, which turns the model into an optimised graph for production.
  • Training across many chips using standard PyTorch tools such as FSDP2, tensor parallelism and pipeline parallelism.
  • Custom kernels, small hand-tuned programs for the chip, written with AWS's Neuron Kernel Interface.
  • Profiling with the usual PyTorch profiler, linked to AWS's Neuron Explorer tool.
  • Hardware: Trainium2 and Trainium3 instances only.

Hugging Face says its Transformers, Diffusers, Accelerate, PEFT and TRL libraries now run on Trainium directly. It says 89% of the top 100 text-generation model architectures work out of the box, compared with 46% through its older Optimum Neuron library, and it now recommends the direct route.

Limits of the beta

AWS is open about gaps. In its August run of PyTorch's own test suite, AWS says about 22,000 tests were skipped because the operations do not yet have a native Neuron version. Features not supported in this first release include sparse tensors, nested tensors, dynamic shapes and some attention options. FP8 support is in private beta. AWS says it plans to close the gaps in later releases.

Older routes are being retired. AWS's release notes say PyTorch/XLA inference on Trainium2 and Trainium3 ends with Neuron 2.34, and PyTorch 2.9 is the last version on PyTorch/XLA.

Customer claims

AWS's announcement includes quotes from early users. Deepgram, a voice AI company, says that in its own benchmarks Trainium processed about 60% more audio seconds per dollar than its existing GPU setup. Cartesia says it ran its inference pipeline "with minimal changes". These are company statements, not independent tests.

Why it matters

Nvidia's lead in AI hardware rests partly on software: most AI code already works on its chips. Cloud providers building their own chips need to make switching cheap. Amazon says about 85% of AI researchers and developers write code in PyTorch, so first-class PyTorch support is a direct way to lower that barrier. AWS says it has also sent changes to the main PyTorch project to remove assumptions that code runs on Nvidia's CUDA, which could help other chipmakers too.

Whether developers move will depend on real-world speed, price and how quickly the remaining gaps close.

Sources

Source