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 support is available at three levels:

  • TorchNeuron Native (recommended): The newest native PyTorch backend providing eager execution, torch.compile, and standard distributed APIs (FSDP, DTensor, DDP, Tensor Parallelism) for Trainium and Inferentia. This is the recommended starting point for new workloads.

  • PyTorch NeuronX (torch-neuronx) (supported): The XLA-based PyTorch integration supporting NeuronCores v2 architecture (Trn1, Trn2, Inf2, Trn1n). Provides full capabilities for both training and inference workloads.

  • PyTorch Neuron (torch-neuron) (archived): The legacy PyTorch integration for NeuronCores v1 architecture (Inf1 only). This package is no longer actively developed. See PyTorch Neuron (torch-neuron) — Archived for reference documentation.

Which Neuron framework for PyTorch should I select?

For help selecting a framework type for inference, see: * About PyTorch on AWS Neuron * Comparison of torch-neuron (Inf1) versus torch-neuronx (Inf2 & Trn1) for Inference

Get Started#

TorchNeuron Native Backend Overview

Recommended for new workloads — Learn about the native PyTorch backend with eager execution, torch.compile support, and standard distributed APIs for Trainium and Inferentia.

Setup Guide

Install and configure PyTorch NeuronX for your environment.

Training & Inference#

Training on Trn1 and Trn2

Train models using PyTorch NeuronX on Trainium instances.

Inference on Inf2, Trn1, and Trn2

Deploy inference workloads using PyTorch NeuronX.

Release Notes#

PyTorch Neuron Component Release Notes

Review the PyTorch Neuron release notes for all versions of the Neuron SDK.

Note

Looking for torch-neuron (Inf1) documentation? The torch-neuron package has been archived. See PyTorch Neuron (torch-neuron) — Archived for legacy Inf1 documentation.

This document is relevant for: Inf1, Inf2, Trn1, Trn2, Trn3