This document is relevant for: Inf1, Inf2, Trn1, 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.

Warning

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.

Prerequisites#

Requirement

Details

Instance Type

Inf2, Trn1, 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

Get the latest Neuron Multi-Framework DLAMI for Ubuntu 24.04 using the AWS CLI:

aws ec2 describe-images \
  --owners amazon \
  --filters "Name=name,Values=Deep Learning AMI Neuron (Ubuntu 24.04)*" \
  --query 'Images | sort_by(@, &CreationDate) | [-1].ImageId' \
  --output text

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-24.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 Trn1 or Inf2 instance with the AMI using the AWS CLI:

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

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_24_0_1_1_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 trn1.*, trn2.*, trn3.*, or inf2.*

  2. Check Neuron driver:

    lsmod | grep neuron
# Should show neuron driver loaded

  3. Restart Neuron runtime:

    sudo systemctl restart neuron-monitor
neuron-ls

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 (Ubuntu 24.04) and activate the pre-configured environment:

source /opt/aws_neuronx_venv_pytorch_inference_vllm_0_24_0_1_1_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:

  3. Read Documentation:

  4. Explore Tools:

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

Additional resources#

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