This document is relevant for: Inf2, Trn1, Trn2

PyTorch Neuron (torch-neuronx) Setup

Install PyTorch with Neuron support for training and inference on Inf2, Trn1, Trn2, and Trn3 instances. Choose from a pre-configured DLAMI, a Docker container, or a manual pip installation.

For the full setup guide with all options, see Install PyTorch for Neuron.

🚀 DLAMI Installation

Pre-configured environment with all dependencies. Recommended for most users.

🚀 Multi-Framework DLAMI

Pre-configured AMI with PyTorch, JAX, and vLLM virtual environments ready to use.

� Deep Learning Container

Pre-configured Docker images from AWS ECR for containerized deployments.

🔧 Manual Installation

Install on bare OS AMIs or existing systems with full control over dependencies.

Rocky Linux 9

Install on Rocky Linux 9 using the Rocky-9-EC2-Base AMI.

This document is relevant for: Inf2, Trn1, Trn2