October 20, 2026 | AI | By Mahesh
Amazon Web Services released TorchNeuron on October 20, 2026, an open-source native PyTorch backend that lets standard PyTorch code run on its Trainium2 and Trainium3 AI chips, in public beta under the Apache 2.0 license [1][3]. In eager mode the move from a GPU is a one-line change to the neuron device, and installation is pip install torch-neuronx [1]. The only performance number comes from a customer, Deepgram, which reports about 60% more input audio seconds per dollar than its GPU baselines.
- TorchNeuron registers AWS Trainium as a native PyTorch device, so existing code moves to Trainium2 or Trainium3 with a one-line change.
- The software is free under Apache 2.0 and installs with pip in public beta from October 20, while Trainium instance pricing is not part of the announcement.
- The only performance figure is Deepgram's own claim of 60% more audio seconds per dollar, and dynamic shapes, sparse tensors and arbitrary attention masks are unsupported.
What TorchNeuron is and what is new
TorchNeuron is a backend that plugs Trainium into PyTorch through its PrivateUse1 extension point and the torch.accelerator APIs. That makes the neuron device work with torch.nn, torch.autocast, torch.distributed and torch.profiler [1]. It supports eager mode, with an experimental adaptive eager mode, and torch.compile. Trainium previously needed AWS's own software path, and Amazon now describes the chip as a first-class PyTorch device [1][3].
The post is written by Nafea Bshara and co-authors at AWS, and the code is public on GitHub [1][4]. Hugging Face says Neuron users should now work with Transformers, Diffusers, Accelerate, PEFT and TRL directly, in place of its earlier Optimum Neuron library [2].
"PyTorch's vision is straightforward: developers should have a consistent experience across hardware platforms."Jana van Greunen, Director of PyTorch Engineering, Meta [1]
Under the hood: how the backend works
Code changes are small. In eager mode a model moves with MyModel().to('neuron'), and compiled training calls torch.compiler.set_default_backend('neuron') once before compiling. Transformers models load with device_map='neuron' [1]. Custom kernels are written in the Neuron Kernel Interface, and an open-source library called NKI-lib supplies optimized kernels such as FlashAttention and MoE blockwise kernels [1].
Profiling uses the standard torch.profiler extended with a NeuronConfig, viewable in Neuron Explorer with Parquet export [1]. Hugging Face's guide fine-tunes Qwen3-8B with tensor parallelism of 4 and FSDP of 8 across 32 NeuronCores on 8 Trainium2 chips, which is half of a trn2.48xlarge instance [2]. Parallelism settings are passed through a DistributedConfig.
Evidence: what has and has not been measured
Deepgram's chief technology officer, Adam Sypniewski, says Trainium processes about 60% more input audio seconds per dollar than the company's GPU baselines, after a multi-month port with the AWS Generative AI Innovation Center. Deepgram measured this and AWS did not [1]. The Next Web does not repeat the figure and describes Amazon's evidence as a single demo in which the same training script ran on a standard chip and on Trainium and the model learned at the same rate [3].
Hugging Face reports that 89% of the 100 most-used text-generation architectures run out of the box, against 46% with Optimum Neuron version 0.4.5, and one vision transformer kernel reaches up to a 2.4 times speedup [2]. No independent benchmarks, throughput tables or cost comparisons against Nvidia hardware were published, and AWS says only qualitatively that a Llama 3 8B training run on 8 trn2.48xlarge instances behaves as expected [1][3].
Pricing, availability and how to access it
TorchNeuron is free to use under Apache 2.0. Public beta began October 20 at PyTorch Conference North America in San Jose, and the package installs from the AWS Neuron pip repository [1][3]. It runs on Trainium2 and Trainium3 instances such as trn2.48xlarge, and none of the sources state instance prices. Named early users are Cartesia, Deepgram and Hugging Face, with partners including Anyscale, Modular and Luminal [1].
Limitations, trade-offs and alternatives
About 22,000 upstream PyTorch tests are skipped because the operators lack a native Neuron lowering. The first release does not support sparse tensors, NestedTensor, forward-mode autograd, per-channel quantization, dynamic or symbolic shapes or arbitrary attention masks, since only no-mask and causal are handled [1]. TorchAO and FP8 are in private beta, and errors surface at the next synchronization point and not at the failing operation [1].
Neuron also favors static shapes, so Hugging Face pads training batches to multiples of 128 tokens and recommends a static cache for generation to avoid repeated recompiles [2]. Nvidia's chips and CUDA still dominate AI computing, as Depth Grid noted in its coverage of Nvidia's talks over Reflection AI, and TorchNeuron lowers the porting cost without proving price or speed parity. Teams should run their own workload on a trn2 instance before committing.
Sources
- AWS and PyTorch. "Introducing TorchNeuron: Native PyTorch backend for AWS Trainium," PyTorch Blog, October 19, 2026. pytorch.org
- Hugging Face. "TorchNeuron: running HF models on AWS Trainium is easier than ever." huggingface.co
- The Next Web. "Amazon's TorchNeuron lets PyTorch code run on AWS Trainium chips." thenextweb.com
- AWS Neuron. torch-neuronx source repository, GitHub. github.com
Read More on Depth Grid
- Nvidia in Talks to Deepen Stake in or Buy Reflection AI
- Mistral Large 4 Preview: 1.05T Open-Weight MoE, Specs and Pricing
- Claude Haiku 5.5 Launch: $0.10 Input, 1M Context, Benchmarks
- White House AI Incident Reporting and Anthropic
Article by Mahesh | Depth Grid


Discussion (0)
Comments are disabled for this post.