Amazon Opens Trainium to PyTorch Developers With TorchNeuron

Amazon Opens Trainium to PyTorch Developers With TorchNeuron

Amazon and the PyTorch project are releasing TorchNeuron, an open-source software layer that lets developers run PyTorch AI code on AWS Trainium chips with minimal changes. TorchNeuron gives PyTorch developers a native software path to Amazon’s AI hardware.

Trainium joins PyTorch

TorchNeuron makes Trainium appear as a standard PyTorch device. It supports eager execution for research and debugging, torch.compile for production workloads, and custom kernels through the Neuron Kernel Interface. The software supports AWS Trainium2 and Trainium3 chips.

PyTorch says TorchNeuron will become a public beta on October 20 under the Apache 2.0 license. The source code is available in the AWS Neuron GitHub repository. The Neuron documentation describes its supported features and tools.

Hardware choice expands

The project is designed to preserve common PyTorch APIs, distributed training tools and model libraries. PyTorch says TorchNeuron works with tools such as TorchTitan and Hugging Face’s Transformers and Diffusers libraries.

That design could reduce the software work needed to move AI workloads onto Trainium. It also gives cloud customers another accelerator option within the PyTorch ecosystem.

Early results show savings

AWS says speech company Deepgram used TorchNeuron to run its Nova-3 speech-to-text model on Trainium2. In a published benchmark, the workload processed 11% to 70% more audio per dollar than the compute options Deepgram had previously used.

The result applied to Deepgram’s offline, throughput-focused batch tests. The first release also has limits, including unsupported or limited support for some sparse operations, dynamic and symbolic shapes, and certain quantization features.

References

Bshara, N., Hundt, R., Diamant, R., Brace, C., Solanki, R., Suleiman, M., ElGalaind, K., Kahsai, T., Bortolato, C., Thangakrishnan, I., Zahirazami, S., Ozarde, S., Singh, A., Wang, H., Lakhotia, E., & Lang, J. (2026, October 19). Introducing TorchNeuron: Native PyTorch backend for AWS Trainium. PyTorch. https://pytorch.org/blog/introducing-torchneuron-native-pytorch-backend-for-aws-trainium/

Amazon. (n.d.). Amazon opens its AWS Trainium AI chips to every PyTorch developer, no rewrite needed. https://www.aboutamazon.com/news/innovation-at-amazon/torchneuron-pytorch-aws-trainium-chips-ai

Lienhart, P., Gomes, A., Zhang, E., An, H., Chaudhari, N., Thekkanal, R., & Wang, V. (2026, October 19). How Deepgram optimized voice AI at scale with AWS Trainium. AWS Machine Learning Blog. https://aws.amazon.com/blogs/machine-learning/how-deepgram-optimized-voice-ai-at-scale-with-aws-trainium/

AWS Neuron. (2026). TorchNeuron: Native PyTorch backend for AWS Trainium. GitHub. https://github.com/aws-neuron/torch-neuronx