AWS announces TorchNeuron: Native PyTorch Support for AWS Trainium (Public Beta)

Today, AWS launched the public beta of TorchNeuron, an open source native PyTorch backend for AWS Trainium. Trainium is a portfolio of AI chips designed to deliver high performance and cost efficiency for deep learning AI training and inference workloads. TorchNeuron enables AI scientists and researchers to experiment and build models with PyTorch natively on Trainium.

With TorchNeuron, developers can enable their PyTorch code on Trainium simply by switching to .to('neuron') while maintaining full PyTorch compatibility with eager execution, torch.compile, and standard PyTorch APIs. Additionally, TorchNeuron integrates with Neuron Kernel Interface (NKI) to enable custom operators and performance optimization for Trainium hardware.

TorchNeuron is available as open source through the TorchNeuron github repo for Amazon EC2 Trainium2 instances and Amazon EC2 Trainium3 instances.

To learn more about TorchNeuron and AWS Trainium accelerators, visit the AWS Neuron page.