This document is relevant for: Trn2, Trn3

Install PyTorch for Neuron#

Install TorchNeuron, the native PyTorch backend for AWS Trainium, for training and inference on Trn2 and Trn3 instances.

Supported Instances: Trn2, Trn3

PyTorch Version: 2.13 or 2.14 (torch-neuronx 2.13.3.0 or 2.14.3.0)

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	Trn2 or Trn3 instance
Operating System	Ubuntu 26.04, Ubuntu 24.04, or Amazon Linux 2023. For the tested Linux kernel and GLIBC versions, see Neuron Runtime and Driver compatibility reference.
Python Version	Python 3.12, 3.13, or 3.14
AWS Account	With EC2 launch permissions
SSH Access	Key pair for instance connection

What You’ll Get#

After installation, you’ll have:

PyTorch 2.13 or 2.14 with the TorchNeuron native backend
torch-neuronx package (TorchNeuron) 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.13.0 or 2.14.0
torch-neuronx	2.13.3.0 or 2.14.3.0
neuronx-cc	2.28.0+
Python	3.12, 3.13, 3.14

Next Steps#

After installation:

Verify Installation: Run verification commands in the installation guide
Try a Tutorial:
- Build your first model: TorchNeuron quickstart
- Inference: vLLM Neuron online serving quickstart
Read the TorchNeuron documentation: TorchNeuron: Native PyTorch on AWS Neuron
Explore tutorials: TorchNeuron tutorials

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