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

PyTorch NeuronX Application Notes#

This section contains application notes specific to PyTorch NeuronX (torch-neuronx) for Trn1 and Inf2 instances. These guides cover PyTorch version migrations, advanced features, optimization techniques, and best practices for training and inference on AWS Trainium and Inferentia2.

PyTorch Version Support#

PyTorch 2.9 Support

New features and migration guide for PyTorch 2.9 on Neuron

PyTorch 2.8 Support

New features and migration guide for PyTorch 2.8 on Neuron

PyTorch 2.7 Support

Features and improvements introduced with PyTorch 2.7 support

PyTorch 2.x Overview

General guide to PyTorch 2.x series support and features

Advanced Features#

Graph Partitioner

Advanced graph partitioning strategies for distributed training and inference

Data Parallel Inference

Scale inference workloads using torch.neuronx.DataParallel for multi-core execution

XLA Migration Guide

Migrate from deprecated XLA environment variables to PyTorch mixed-precision options

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