This document is relevant for: Trn2, Trn3

Install PyTorch via Deep Learning AMI

Install PyTorch with Neuron support using pre-configured AWS Deep Learning AMIs.

⏱️ Estimated time: 5 minutes

Note

Want to read about Neuron’s Deep Learning machine images (DLAMIs) before diving in? Check out the Neuron DLAMI User Guide.

---

Prerequisites

Requirement	Details
Instance Type	Trn2 or Trn3
AWS Account	With EC2 permissions
SSH Key Pair	For instance access
AWS CLI	Configured with credentials (optional)

Installation steps

Step 1: Find the latest AMI

List the Neuron DLAMIs for Neuron SDK 2.33.0 with the AWS CLI. In the output, find the multi-framework DLAMI for your operating system and record its ImageId:

aws ec2 describe-images \
  --owners amazon \
  --filters "Name=name,Values=Deep Learning *Neuron*" \
            "Name=description,Values=*Neuron SDK 2.33.0*" \
  --query 'Images[].{Name:Name, ImageId:ImageId}' \
  --output table \
  --region us-east-1

Or use the SSM parameter to resolve the latest image directly:

aws ssm get-parameter \
  --region us-east-1 \
  --name /aws/service/neuron/dlami/multi-framework/ubuntu-26.04/latest/image_id \
  --query "Parameter.Value" \
  --output text

You can also use the AWS EC2 parameter store to find the ID of a DLAMI. See Find a DLAMI ID for details. Record the ID (image-id) for the next step.

Step 2: Launch instance

Launch a Trn2 or Trn3 instance with the AMI using the AWS CLI:

aws ec2 run-instances \
  --image-id ami-xxxxxxxxxxxxxxxxx \
  --instance-type trn2.48xlarge \
  --key-name your-key-pair \
  --security-group-ids sg-xxxxxxxxx \
  --subnet-id subnet-xxxxxxxxx

Replace:

• ami-xxxxxxxxxxxxxxxxx with AMI ID from Step 1

• your-key-pair with your SSH key pair name

• sg-xxxxxxxxx with your security group ID

• subnet-xxxxxxxxx with your subnet ID

You can also launch your DLAMI through the AWS EC2 web console, which also provides hints for security group and subnet IDs. For more details, see Launch a DLAMI.

Step 3: Connect to instance

ssh -i your-key-pair.pem ubuntu@<instance-public-ip>

Step 4: Activate environment

The multi-framework DLAMI includes pre-configured virtual environments. Activate the vLLM environment for LLM inference:

source /opt/aws_neuronx_venv_pytorch_inference_vllm_0_29_0_1_2_0/bin/activate

Or activate the JAX environment:

source /opt/aws_neuronx_venv_jax_0_10/bin/activate

For a full list of available virtual environments, see Virtual Environments pre-installed.

Step 5: Verify installation

neuron-ls
pip list | grep -E "vllm|neuron|torch"

You should see output similar to this (the versions, instance IDs, and details should match your expected ones, not the ones in this example):

Expected output:

+--------+--------+--------+-----------+
| DEVICE | CORES  | MEMORY | CONNECTED |
+--------+--------+--------+-----------+
| 0      | 2      | 32 GB  | Yes       |
| 1      | 2      | 32 GB  | Yes       |
+--------+--------+--------+-----------+

⚠️ Troubleshooting: Module not found

If you see ModuleNotFoundError when importing packages:

1. Verify virtual environment is activated:

   which python
   # Should show a path under /opt/aws_neuronx_venv_*
   

2. Check Python version:

   python --version
   # Should be 3.12 or higher
   

3. Verify installed packages:

   pip list | grep -E "neuron|vllm|torch"
   

⚠️ Troubleshooting: No Neuron devices found

If neuron-ls shows no devices:

1. Verify instance type:

   curl http://169.254.169.254/latest/meta-data/instance-type
   # Should show trn2.* or trn3.*
   

2. Check Neuron driver:

   lsmod | grep neuron
   # Should show neuron driver loaded
   

3. Restart Neuron runtime:

   sudo systemctl restart neuron-monitor
   neuron-ls
   

Use a single-framework PyTorch DLAMI

To use a single-framework PyTorch DLAMI instead of the multi-framework DLAMI, follow the steps for your operating system:

• Ubuntu 26.04 and Ubuntu 24.04: use the Deep Learning AMI Neuron torch-neuronx 2.14.0 DLAMI. It provides TorchNeuron (torch-neuronx) for PyTorch 2.14, which supports Trn2 and Trn3 instances.

• Amazon Linux 2023: the torch-neuronx 2.14.0 DLAMI is not available. The Amazon Linux 2023 tab covers the PyTorch 2.9 (PyTorch/XLA) DLAMI from Neuron SDK 2.31.0 and earlier.

Ubuntu 26.04

Step 1: Find the latest AMI

List the Neuron DLAMIs for Neuron SDK 2.33.0 with the AWS CLI. In the output, find the torch-neuronx 2.14.0 DLAMI for Ubuntu 26.04 and record its ImageId:

aws ec2 describe-images \
  --owners amazon \
  --filters "Name=name,Values=Deep Learning *Neuron*" \
            "Name=description,Values=*Neuron SDK 2.33.0*" \
  --query 'Images[].{Name:Name, ImageId:ImageId}' \
  --output table \
  --region us-east-1

You can also use the AWS EC2 parameter store to find the ID of a DLAMI. See Find a DLAMI ID for details. Record the ID (image-id) for the next step.

Step 2: Launch instance

Launch a Trn2 or Trn3 instance with the AMI using the AWS CLI:

aws ec2 run-instances \
  --image-id ami-xxxxxxxxxxxxxxxxx \
  --instance-type trn2.48xlarge \
  --key-name your-key-pair \
  --security-group-ids sg-xxxxxxxxx \
  --subnet-id subnet-xxxxxxxxx

Replace:

• ami-xxxxxxxxxxxxxxxxx with AMI ID from Step 1

• your-key-pair with your SSH key pair name

• sg-xxxxxxxxx with your security group ID

• subnet-xxxxxxxxx with your subnet ID

Step 3: Connect to instance

ssh -i your-key-pair.pem ubuntu@<instance-public-ip>

Step 4: Activate environment

The DLAMI includes a pre-configured virtual environment:

source /opt/aws_neuronx_venv_torch_neuronx_2_14_0/bin/activate

Step 5: Verify installation

python3 -c "import torch; import torch_neuronx; print(f'PyTorch {torch.__version__}, torch-neuronx {torch_neuronx.__version__}')"
neuron-ls

You should see output similar to this (the versions, instance IDs, and details should match your expected ones, not the ones in this example):

Expected output:

PyTorch 2.14.0+cpu, torch-neuronx 2.14.3.0

+--------+--------+--------+-----------+
| DEVICE | CORES  | MEMORY | CONNECTED |
+--------+--------+--------+-----------+
| 0      | 2      | 32 GB  | Yes       |
| 1      | 2      | 32 GB  | Yes       |
+--------+--------+--------+-----------+

⚠️ Troubleshooting: Module not found

If you see ModuleNotFoundError: No module named 'torch_neuronx':

1. Verify virtual environment is activated

2. Check Python version: python --version (should be 3.12 or higher)

3. Reinstall: pip install --force-reinstall torch-neuronx

⚠️ Troubleshooting: No Neuron devices found

If neuron-ls shows no devices:

1. Verify instance type

2. Check Neuron driver: lsmod | grep neuron

3. Restart runtime: sudo systemctl restart neuron-monitor

Ubuntu 24.04

Step 1: Find the latest AMI

List the Neuron DLAMIs for Neuron SDK 2.33.0 with the AWS CLI. In the output, find the torch-neuronx 2.14.0 DLAMI for Ubuntu 24.04 and record its ImageId:

aws ec2 describe-images \
  --owners amazon \
  --filters "Name=name,Values=Deep Learning *Neuron*" \
            "Name=description,Values=*Neuron SDK 2.33.0*" \
  --query 'Images[].{Name:Name, ImageId:ImageId}' \
  --output table \
  --region us-east-1

You can also use the AWS EC2 parameter store to find the ID of a DLAMI. See Find a DLAMI ID for details. Record the ID (image-id) for the next step.

Step 2: Launch instance

Launch a Trn2 or Trn3 instance with the AMI using the AWS CLI:

aws ec2 run-instances \
  --image-id ami-xxxxxxxxxxxxxxxxx \
  --instance-type trn2.48xlarge \
  --key-name your-key-pair \
  --security-group-ids sg-xxxxxxxxx \
  --subnet-id subnet-xxxxxxxxx

Replace:

• ami-xxxxxxxxxxxxxxxxx with AMI ID from Step 1

• your-key-pair with your SSH key pair name

• sg-xxxxxxxxx with your security group ID

• subnet-xxxxxxxxx with your subnet ID

You can also launch your DLAMI through the AWS EC2 web console, which also provides hints for security group and subnet IDs. For more details, see Launch a DLAMI.

Step 3: Connect to instance

ssh -i your-key-pair.pem ubuntu@<instance-public-ip>

Step 4: Activate environment

The DLAMI includes a pre-configured virtual environment:

source /opt/aws_neuronx_venv_torch_neuronx_2_14_0/bin/activate

Step 5: Verify installation

python3 -c "import torch; import torch_neuronx; print(f'PyTorch {torch.__version__}, torch-neuronx {torch_neuronx.__version__}')"
neuron-ls

You should see output similar to this (the versions, instance IDs, and details should match your expected ones, not the ones in this example):

Expected output:

PyTorch 2.14.0+cpu, torch-neuronx 2.14.3.0

+--------+--------+--------+-----------+
| DEVICE | CORES  | MEMORY | CONNECTED |
+--------+--------+--------+-----------+
| 0      | 2      | 32 GB  | Yes       |
| 1      | 2      | 32 GB  | Yes       |
+--------+--------+--------+-----------+

⚠️ Troubleshooting: Module not found

If you see ModuleNotFoundError: No module named 'torch_neuronx':

1. Verify virtual environment is activated:

   which python
   # Should show /opt/aws_neuronx_venv_torch_neuronx_2_14_0/bin/python
   

2. Check Python version:

   python --version
   # Should be 3.12 or higher
   

3. Reinstall torch-neuronx:

   pip install --force-reinstall torch-neuronx
   

⚠️ Troubleshooting: No Neuron devices found

If neuron-ls shows no devices:

1. Verify instance type:

   curl http://169.254.169.254/latest/meta-data/instance-type
   # Should show trn2.* or trn3.*
   

2. Check Neuron driver:

   lsmod | grep neuron
   # Should show neuron driver loaded
   

3. Restart Neuron runtime:

   sudo systemctl restart neuron-monitor
   neuron-ls
   

Amazon Linux 2023

Note

The Deep Learning AMI Neuron torch-neuronx 2.14.0 DLAMI is available on Ubuntu 24.04 and Ubuntu 26.04. The following steps are for the PyTorch 2.9 (PyTorch/XLA) DLAMI, which is published for Neuron SDK 2.31.0 and earlier only.

Step 1: Find the latest AMI

List the Neuron DLAMIs for Neuron SDK 2.31.0 with the AWS CLI. In the output, find the PyTorch 2.9 DLAMI for Amazon Linux 2023 and record its ImageId:

aws ec2 describe-images \
  --owners amazon \
  --filters "Name=name,Values=Deep Learning *Neuron*" \
            "Name=description,Values=*Neuron SDK 2.31.0*" \
  --query 'Images[].{Name:Name, ImageId:ImageId}' \
  --output table \
  --region us-east-1

Step 2: Launch instance

aws ec2 run-instances \
  --image-id ami-xxxxxxxxxxxxxxxxx \
  --instance-type trn1.2xlarge \
  --key-name your-key-pair \
  --security-group-ids sg-xxxxxxxxx \
  --subnet-id subnet-xxxxxxxxx

Step 3: Connect to instance

ssh -i your-key-pair.pem ec2-user@<instance-public-ip>

Note

Amazon Linux 2023 uses ec2-user instead of ubuntu.

Step 4: Activate environment

source /opt/aws_neuronx_venv_pytorch_2_9/bin/activate

Step 5: Verify installation

python3 -c "import torch; import torch_neuronx; print(f'PyTorch {torch.__version__}, torch-neuronx {torch_neuronx.__version__}')"
neuron-ls

You should see output similar to this (the versions, instance IDs, and details should match your expected ones, not the ones in this example):

Expected output:

PyTorch 2.9.0+cpu, torch-neuronx 2.9.0.1.0

+--------+--------+--------+-----------+
| DEVICE | CORES  | MEMORY | CONNECTED |
+--------+--------+--------+-----------+
| 0      | 2      | 32 GB  | Yes       |
| 1      | 2      | 32 GB  | Yes       |
+--------+--------+--------+-----------+

⚠️ Troubleshooting: Module not found

If you see ModuleNotFoundError: No module named 'torch_neuronx':

1. Verify virtual environment is activated

2. Check Python version: python --version (should be 3.12 or later)

3. Reinstall: pip install --force-reinstall torch-neuronx

⚠️ Troubleshooting: No Neuron devices found

If neuron-ls shows no devices:

1. Verify instance type

2. Check Neuron driver: lsmod | grep neuron

3. Restart runtime: sudo systemctl restart neuron-monitor

Update an existing installation

To update PyTorch versions or Neuron drivers on an existing DLAMI, see Update PyTorch on a Deep Learning AMI.

Tip

vLLM for LLM inference

Neuron provides a dedicated vLLM DLAMI with vLLM and the vLLM-Neuron Plugin pre-installed. Launch the Deep Learning AMI Neuron PyTorch Inference vLLM 0.29.0.1.2.0 (Ubuntu 26.04) and activate the pre-configured environment:

source /opt/aws_neuronx_venv_pytorch_inference_vllm_0_29_0_1_2_0/bin/activate

vLLM provides an OpenAI-compatible API, continuous batching, and supports models like Llama 2/3.1/3.3/4, Qwen 2.5/3, and multimodal models with quantization support (INT8/FP8).

The vLLM environment is also available in the multi-framework DLAMI. For more details on available DLAMIs and SSM parameters, see Neuron DLAMI User Guide.

Next steps

Now that PyTorch is installed:

1. Try a Quick Example:

   import torch
   import torch_neuronx
   
   # Simple tensor operation on Neuron
   x = torch.randn(3, 3)
   model = torch.nn.Linear(3, 3)
   
   # Compile for Neuron
   trace = torch_neuronx.trace(model, x)
   print(trace(x))
   

2. Follow Tutorials:

   • TorchNeuron tutorials

   • vLLM Neuron online serving quickstart

3. Read Documentation:

   • TorchNeuron quickstart

   • TorchNeuron: Native PyTorch on AWS Neuron

4. Explore Tools:

   • Neuron Explorer

   • Neuron Top User Guide

5. Deploy LLM inference: Neuron DLAMI User Guide (vLLM on Neuron)

Additional resources

• Neuron DLAMI User Guide - DLAMI documentation

• Deploy on AWS - Container-based deployment

• Installation Troubleshooting - Common issues and solutions

• AWS Neuron SDK Release Notes - Version compatibility information

This document is relevant for: Trn2, Trn3