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

Install PyTorch via manual installation#

Install PyTorch with Neuron support on a bare OS AMI or existing system.

⏱️ Estimated time: 15 minutes

Note

For a faster setup, consider using the DLAMI-based installation instead.

Note

  • Instructions in this page only apply to setting up Neuron components on Linux host running Ubuntu or Amazon Linux AMI.

  • When launching a Trn1/Trn2/Trn3 instance, you must adjust your primary EBS volume size to a minimum of 512GB.

Prerequisites#

Requirement

Details

Instance Type

Inf2, Trn1, Trn2, or Trn3

Operating System

Ubuntu 24.04, Ubuntu 22.04, or Amazon Linux 2023

Python Version

Python 3.10, 3.11, or 3.12

AWS Account

With EC2 permissions

SSH Key Pair

For instance access

Installation steps#

Step 1: Launch instance

Step 2: Install drivers and tools

# Configure Linux for Neuron repository updates
. /etc/os-release
sudo tee /etc/apt/sources.list.d/neuron.list > /dev/null <<EOF
deb https://apt.repos.neuron.amazonaws.com ${VERSION_CODENAME} main
EOF
wget -qO - https://apt.repos.neuron.amazonaws.com/GPG-PUB-KEY-AMAZON-AWS-NEURON.PUB | sudo apt-key add -

# Update OS packages 
sudo apt-get update -y

# Install OS headers 
sudo apt-get install linux-headers-$(uname -r) -y

# Install git 
sudo apt-get install git -y

# install Neuron Driver
sudo apt-get install aws-neuronx-dkms=2.* -y

# Install Neuron Runtime 
sudo apt-get install aws-neuronx-collectives=2.* -y
sudo apt-get install aws-neuronx-runtime-lib=2.* -y

# Install Neuron Tools 
sudo apt-get install aws-neuronx-tools=2.* -y

# Add PATH
export PATH=/opt/aws/neuron/bin:$PATH

Step 3: Install EFA (Trn1/Trn1n/Trn2/Trn3 only)

# Install EFA Driver (only required for multi-instance training)
curl -O https://efa-installer.amazonaws.com/aws-efa-installer-latest.tar.gz 
wget https://efa-installer.amazonaws.com/aws-efa-installer.key && gpg --import aws-efa-installer.key 
cat aws-efa-installer.key | gpg --fingerprint 
wget https://efa-installer.amazonaws.com/aws-efa-installer-latest.tar.gz.sig && gpg --verify ./aws-efa-installer-latest.tar.gz.sig 
tar -xvf aws-efa-installer-latest.tar.gz 
cd aws-efa-installer && sudo bash efa_installer.sh --yes 
cd 
sudo rm -rf aws-efa-installer-latest.tar.gz aws-efa-installer

# U24 2.9.0 Pytorch Neuronx Upgrade

Step 4: Install PyTorch and Neuron packages

# Install Python venv 
sudo apt-get install -y python3.12-venv g++ 

# Create Python venv
python3.12 -m venv aws_neuron_venv_pytorch 

# Activate Python venv 
source aws_neuron_venv_pytorch/bin/activate 
python -m pip install -U pip 

# Install Jupyter notebook kernel
pip install ipykernel 
python3.12 -m ipykernel install --user --name aws_neuron_venv_pytorch --display-name "Python (torch-neuronx)"
pip install jupyter notebook
pip install environment_kernels

# Set pip repository pointing to the Neuron repository 
python -m pip config set global.extra-index-url https://pip.repos.neuron.amazonaws.com

# Install wget, awscli 
python -m pip install wget 
python -m pip install awscli 

# Install Neuron Compiler and Framework
python -m pip install neuronx-cc==2.* torch-neuronx==2.9.* torchvision

Traceback (most recent call last):
  File "/home/docs/checkouts/readthedocs.org/user_builds/awsdocs-neuron/checkouts/latest/src/helperscripts/n2-helper.py", line 1151, in <module>
    print(n2_manifest.generate_script(args))
  File "/home/docs/checkouts/readthedocs.org/user_builds/awsdocs-neuron/checkouts/latest/src/helperscripts/n2-helper.py", line 138, in generate_script
    str_python = self.set_python_venv(args)
  File "/home/docs/checkouts/readthedocs.org/user_builds/awsdocs-neuron/checkouts/latest/src/helperscripts/n2-helper.py", line 575, in set_python_venv
    packages_supporting_python_versions = self.get_pip_packages_supporting_python_versions(args)
  File "/home/docs/checkouts/readthedocs.org/user_builds/awsdocs-neuron/checkouts/latest/src/helperscripts/n2-helper.py", line 82, in get_pip_packages_supporting_python_versions
    framework_python_versions = df_framework.loc[
IndexError: index 0 is out of bounds for axis 0 with size 0

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       |
+--------+--------+--------+-----------+

Step 1: Launch instance

Step 2: Install drivers and tools

# Configure Linux for Neuron repository updates
. /etc/os-release
sudo tee /etc/apt/sources.list.d/neuron.list > /dev/null <<EOF
deb https://apt.repos.neuron.amazonaws.com ${VERSION_CODENAME} main
EOF
wget -qO - https://apt.repos.neuron.amazonaws.com/GPG-PUB-KEY-AMAZON-AWS-NEURON.PUB | sudo apt-key add -

# Update OS packages 
sudo apt-get update -y

# Install OS headers 
sudo apt-get install linux-headers-$(uname -r) -y

# Install git 
sudo apt-get install git -y

# install Neuron Driver
sudo apt-get install aws-neuronx-dkms=2.* -y

# Install Neuron Runtime 
sudo apt-get install aws-neuronx-collectives=2.* -y
sudo apt-get install aws-neuronx-runtime-lib=2.* -y

# Install Neuron Tools 
sudo apt-get install aws-neuronx-tools=2.* -y

# Add PATH
export PATH=/opt/aws/neuron/bin:$PATH

Step 3: Install EFA (Trn1/Trn1n/Trn2/Trn3 only)

# Install EFA Driver (only required for multi-instance training)
curl -O https://efa-installer.amazonaws.com/aws-efa-installer-latest.tar.gz 
wget https://efa-installer.amazonaws.com/aws-efa-installer.key && gpg --import aws-efa-installer.key 
cat aws-efa-installer.key | gpg --fingerprint 
wget https://efa-installer.amazonaws.com/aws-efa-installer-latest.tar.gz.sig && gpg --verify ./aws-efa-installer-latest.tar.gz.sig 
tar -xvf aws-efa-installer-latest.tar.gz 
cd aws-efa-installer && sudo bash efa_installer.sh --yes 
cd 
sudo rm -rf aws-efa-installer-latest.tar.gz aws-efa-installer

Step 4: Install PyTorch and Neuron packages

# Install Python venv 
sudo apt-get install -y python3.10-venv g++ 

# Create Python venv
python3.10 -m venv aws_neuron_venv_pytorch 

# Activate Python venv 
source aws_neuron_venv_pytorch/bin/activate 
python -m pip install -U pip 

# Install Jupyter notebook kernel
pip install ipykernel 
python3.10 -m ipykernel install --user --name aws_neuron_venv_pytorch --display-name "Python (torch-neuronx)"
pip install jupyter notebook
pip install environment_kernels

# Set pip repository pointing to the Neuron repository 
python -m pip config set global.extra-index-url https://pip.repos.neuron.amazonaws.com

# Install wget, awscli 
python -m pip install wget 
python -m pip install awscli 

# Install Neuron Compiler and Framework
python -m pip install neuronx-cc==2.* torch-neuronx==2.9.* torchvision

Traceback (most recent call last):
  File "/home/docs/checkouts/readthedocs.org/user_builds/awsdocs-neuron/checkouts/latest/src/helperscripts/n2-helper.py", line 1151, in <module>
    print(n2_manifest.generate_script(args))
  File "/home/docs/checkouts/readthedocs.org/user_builds/awsdocs-neuron/checkouts/latest/src/helperscripts/n2-helper.py", line 138, in generate_script
    str_python = self.set_python_venv(args)
  File "/home/docs/checkouts/readthedocs.org/user_builds/awsdocs-neuron/checkouts/latest/src/helperscripts/n2-helper.py", line 575, in set_python_venv
    packages_supporting_python_versions = self.get_pip_packages_supporting_python_versions(args)
  File "/home/docs/checkouts/readthedocs.org/user_builds/awsdocs-neuron/checkouts/latest/src/helperscripts/n2-helper.py", line 82, in get_pip_packages_supporting_python_versions
    framework_python_versions = df_framework.loc[
IndexError: index 0 is out of bounds for axis 0 with size 0

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       |
+--------+--------+--------+-----------+

Step 1: Launch instance

Step 2: Install drivers and tools

Traceback (most recent call last):
  File "/home/docs/checkouts/readthedocs.org/user_builds/awsdocs-neuron/checkouts/latest/src/helperscripts/n2-helper.py", line 1151, in <module>
    print(n2_manifest.generate_script(args))
  File "/home/docs/checkouts/readthedocs.org/user_builds/awsdocs-neuron/checkouts/latest/src/helperscripts/n2-helper.py", line 138, in generate_script
    str_python = self.set_python_venv(args)
  File "/home/docs/checkouts/readthedocs.org/user_builds/awsdocs-neuron/checkouts/latest/src/helperscripts/n2-helper.py", line 575, in set_python_venv
    packages_supporting_python_versions = self.get_pip_packages_supporting_python_versions(args)
  File "/home/docs/checkouts/readthedocs.org/user_builds/awsdocs-neuron/checkouts/latest/src/helperscripts/n2-helper.py", line 82, in get_pip_packages_supporting_python_versions
    framework_python_versions = df_framework.loc[
IndexError: index 0 is out of bounds for axis 0 with size 0

Step 3: Install EFA (Trn1/Trn1n/Trn2/Trn3 only)

Traceback (most recent call last):
  File "/home/docs/checkouts/readthedocs.org/user_builds/awsdocs-neuron/checkouts/latest/src/helperscripts/n2-helper.py", line 1151, in <module>
    print(n2_manifest.generate_script(args))
  File "/home/docs/checkouts/readthedocs.org/user_builds/awsdocs-neuron/checkouts/latest/src/helperscripts/n2-helper.py", line 138, in generate_script
    str_python = self.set_python_venv(args)
  File "/home/docs/checkouts/readthedocs.org/user_builds/awsdocs-neuron/checkouts/latest/src/helperscripts/n2-helper.py", line 575, in set_python_venv
    packages_supporting_python_versions = self.get_pip_packages_supporting_python_versions(args)
  File "/home/docs/checkouts/readthedocs.org/user_builds/awsdocs-neuron/checkouts/latest/src/helperscripts/n2-helper.py", line 82, in get_pip_packages_supporting_python_versions
    framework_python_versions = df_framework.loc[
IndexError: index 0 is out of bounds for axis 0 with size 0

Step 4: Install PyTorch and Neuron packages

Traceback (most recent call last):
  File "/home/docs/checkouts/readthedocs.org/user_builds/awsdocs-neuron/checkouts/latest/src/helperscripts/n2-helper.py", line 1151, in <module>
    print(n2_manifest.generate_script(args))
  File "/home/docs/checkouts/readthedocs.org/user_builds/awsdocs-neuron/checkouts/latest/src/helperscripts/n2-helper.py", line 138, in generate_script
    str_python = self.set_python_venv(args)
  File "/home/docs/checkouts/readthedocs.org/user_builds/awsdocs-neuron/checkouts/latest/src/helperscripts/n2-helper.py", line 575, in set_python_venv
    packages_supporting_python_versions = self.get_pip_packages_supporting_python_versions(args)
  File "/home/docs/checkouts/readthedocs.org/user_builds/awsdocs-neuron/checkouts/latest/src/helperscripts/n2-helper.py", line 82, in get_pip_packages_supporting_python_versions
    framework_python_versions = df_framework.loc[
IndexError: index 0 is out of bounds for axis 0 with size 0

Traceback (most recent call last):
  File "/home/docs/checkouts/readthedocs.org/user_builds/awsdocs-neuron/checkouts/latest/src/helperscripts/n2-helper.py", line 1151, in <module>
    print(n2_manifest.generate_script(args))
  File "/home/docs/checkouts/readthedocs.org/user_builds/awsdocs-neuron/checkouts/latest/src/helperscripts/n2-helper.py", line 138, in generate_script
    str_python = self.set_python_venv(args)
  File "/home/docs/checkouts/readthedocs.org/user_builds/awsdocs-neuron/checkouts/latest/src/helperscripts/n2-helper.py", line 575, in set_python_venv
    packages_supporting_python_versions = self.get_pip_packages_supporting_python_versions(args)
  File "/home/docs/checkouts/readthedocs.org/user_builds/awsdocs-neuron/checkouts/latest/src/helperscripts/n2-helper.py", line 82, in get_pip_packages_supporting_python_versions
    framework_python_versions = df_framework.loc[
IndexError: index 0 is out of bounds for axis 0 with size 0

Note

PyTorch 2.9 is not yet available on Amazon Linux 2023. Use Ubuntu 24.04 for PyTorch 2.9 support.

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):

PyTorch version: 2.9.1+cu128, torch-neuronx version: 2.9.0.2.13.23887+8e870898
$ neuron-ls
instance-type: trn1.2xlarge
instance-id: i-0bea223b1afb7e159
+--------+--------+----------+--------+--------------+----------+------+
| NEURON | NEURON |  NEURON  | NEURON |     PCI      |   CPU    | NUMA |
| DEVICE | CORES  | CORE IDS | MEMORY |     BDF      | AFFINITY | NODE |
+--------+--------+----------+--------+--------------+----------+------+
| 0      | 2      | 0-1      | 32 GB  | 0000:00:1e.0 | 0-7      | -1   |
+--------+--------+----------+--------+--------------+----------+------+

Tip

vLLM for LLM inference

After completing the manual installation, you can add vLLM for inference serving using the vllm-neuron plugin:

git clone https://github.com/vllm-project/vllm-neuron.git
cd vllm-neuron
pip install --extra-index-url=https://pip.repos.neuron.amazonaws.com -e .

Or use the pre-configured vLLM DLC image for a containerized deployment. See vLLM on Neuron for all deployment options.

Update an existing installation#

To update PyTorch versions or Neuron drivers on an existing manual installation, see Update a manual PyTorch installation.

Next steps#

Advanced#

Additional resources#

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