This document is relevant for: Inf1, Inf2, Trn1, Trn2, Trn3
PyTorch Support on Neuron#
PyTorch running on Neuron unlocks high-performance and cost-effective deep learning acceleration on AWS Trainium-based and AWS Inferentia-based Amazon EC2 instances.
The PyTorch plugin for Neuron architecture enables native PyTorch models to be accelerated on Neuron devices, so you can use your existing framework application and get started easily with minimal code changes.
For help selecting a framework type for inference, see Comparison of torch-neuron (Inf1) versus torch-neuronx (Inf2 & Trn1) for Inference.
Introducing TorchNeuron, a native backend for AWS Trainium
At re:Invent ‘25, AWS Neuron announced their new PyTorch package, “TorchNeuron”, which includes the torch-neuronx library and initial support for a native PyTorch backend (TorchDynamo) with eager execution, torch.compile, and standard distributed APIs.
For more details on what is coming with TorchNeuron and PyTorch eager mode support, see Native PyTorch for AWS Trainium.
PyTorch NeuronX#
PyTorch Neuron#
This document is relevant for: Inf1, Inf2, Trn1, Trn2, Trn3