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.
PyTorch Neuron is available in two versions to support different AWS ML accelerator architectures:
PyTorch NeuronX (torch-neuronx): The next-generation PyTorch integration supporting NeuronCores v2 architecture (Trn1, Trn2, Inf2, Trn1n). This version provides enhanced capabilities for both training and inference workloads with support for the latest PyTorch features.
PyTorch Neuron (torch-neuron): The original PyTorch integration supporting NeuronCores v1 architecture (Inf1). This version is optimized for inference workloads on Inf1 instances.
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.
This document is relevant for: Inf1, Inf2, Trn1, Trn2, Trn3