AWS Opens TorchNeuron Public Beta for Trainium2 and Trainium3

Amazon Web Services (AWS) opened the public beta of TorchNeuron on Tuesday, Oct. 20, giving developers a new way to run PyTorch workloads on Trainium2 and Trainium3 chips.

Trainium2 and Trainium3 are Amazon's AI chips, while PyTorch is a framework developers use to build and run AI models. TorchNeuron is an open-source backend that connects PyTorch workloads with those chips. It registers a neuron device and supports eager execution, torch.compile, distributed APIs and custom kernels through the Neuron Kernel Interface. The supported features give developers several ways to run workloads through PyTorch as they evaluate the beta.

AWS recommends TorchNeuron for new workloads on PyTorch 2.14. The older XLA-based torch-neuronx package remains a separate software line. The public beta follows AWS's earlier TorchNeuron announcement, making the backend available for public testing on both chip generations.

The beta has compatibility limits. For compiled models, input shapes must remain static. The beta also lacks support for APIs including torch.sparse, torch.func, torch.nested and torch.distributions. Teams evaluating a model will need to account for both its shape requirements and the operations it uses.

Colin Brace, vice president of software at Annapurna Labs, said more than 30 early customers were using TorchNeuron. “This has been about meeting customers where they’re at,” Brace said.

Deepgram evaluated its Nova-3 speech-to-text model on a single Trainium2 trn2.3xlarge instance for offline, throughput-focused batch inference. In that test, the model processed 11% to 70% more audio per dollar than compute options Deepgram had used. The comparison used prices from the evaluation period, and the result applies to that test setup.

Developers can now test workloads on either chip generation in the public beta, checking whether their model operations fit the supported API list and static-shape requirement.