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

Install PyTorch for Neuron#

Install PyTorch with AWS Neuron support for training and inference on Inferentia and Trainium instances.

Supported Instances: Inf2, Trn1, Trn2, Trn3

PyTorch Version: 2.9+ with Native Neuron backend

Warning

The torch-neuronx library is no longer included on Neuron DLAMIs and DLCs starting with Neuron release v2.32.0. The setup instructions in this document assume your environment uses Neuron v2.31.0 or earlier. If you are using torch-neuronx for training your models, pin your environment to Neuron v2.31.0.

The NeuronX Distributed (NxD) library for training (neuronx_distributed_training) is no longer included on Neuron DLAMIs and DLCs as of release v2.31.0. To manually configure your environment to use neuronx_distributed_training, see Install PyTorch via manual installation.

Choose installation method#

Method	Best for	Considerations
DLAMI	Getting started quickly, prototyping, single-user development	Pre-configured with tested dependency versions; launch a new AMI to update
DLC	Production deployments, CI/CD pipelines, multi-tenant environments	Requires Docker and Neuron driver on host; portable across EC2, ECS, EKS
Manual	Custom OS images, shared clusters, integrating into existing environments	Full control over versions and dependencies; requires manual dependency management

🚀 AWS Deep Learning AMI

Recommended for most users

Pre-configured environment with all dependencies

✅ All dependencies included

✅ Tested configurations

✅ Multiple Python versions

⏱️ Setup time: ~5 minutes

🐳 Deep Learning Container

For containerized deployments

Pre-configured Docker images from AWS ECR

✅ Docker-based isolation

✅ Training and inference images

✅ vLLM-ready images available

⏱️ Setup time: ~10 minutes

🔧 Manual Installation

For custom environments

Install on existing systems or custom setups

✅ Existing system integration

✅ Custom Python versions

✅ Full control over dependencies

⏱️ Setup time: ~15 minutes

Prerequisites#

Before installing, ensure you have:

Requirement	Details
Instance Type	Inf2, Trn1, Trn2, or Trn3 instance
Operating System	Ubuntu 24.04, Ubuntu 22.04, or Amazon Linux 2023
Python Version	Python 3.11, 3.12 or 3.13
AWS Account	With EC2 launch permissions
SSH Access	Key pair for instance connection

What You’ll Get#

After installation, you’ll have:

PyTorch 2.9+ with Native Neuron backend
torch-neuronx package for Neuron-specific operations
neuronx-cc compiler for model optimization
Neuron Runtime for model execution
Neuron Tools for profiling and debugging

Version Information#

Component	Version
PyTorch	2.9.0+
torch-neuronx	2.9.0+
neuronx-cc	2.15.0+
Python	3.11, 3.12, 3.13

Next Steps#

After installation:

Verify Installation: Run verification commands in the installation guide
Try a Tutorial:
- Inference: Quickstart: Launch an OpenAI-compatible API server with vLLM Neuron
- Training: MLP training tutorial [archived content]
Read the torch-neuronx Programming Guide: Training developer guide [archived content]
Explore Examples: PyTorch Support on Neuron

Update an existing installation#

Already have PyTorch Neuron installed and need to update to a newer PyTorch version or Neuron SDK release? Select the guide that matches your installation method.

🔄 Update DLAMI

Update PyTorch and drivers on an existing Deep Learning AMI

🔄 Update DLC

Pull the latest container image and update the host driver

🔄 Update manual install

Update PyTorch packages and drivers on a manual installation