This document is relevant for: Trn2, Trn3

Update a manual PyTorch installation#

Update PyTorch and Neuron components on an existing manual installation to the latest release.

Update PyTorch on Ubuntu 24.04#

Select the PyTorch version that your existing installation uses.

TorchNeuron (torch-neuronx) for PyTorch 2.14 supports Trn2 and Trn3 instances. Activate your existing virtual environment, then update TorchNeuron, the Neuron Compiler, and NKI:

source aws_neuron_venv_pytorch_2_14/bin/activate

python -m pip install --upgrade "torch-neuronx==2.14.*" \
  "neuronx-cc==2.*" nki \
  --extra-index-url https://pip.repos.neuron.amazonaws.com \
  --extra-index-url https://download.pytorch.org/whl/cpu

If you already have a previous Neuron release installed, select the PyTorch version tab below to get the update commands for your environment.

Note

These commands update the earlier PyTorch/XLA-based torch-neuronx for PyTorch 2.9, which last shipped in Neuron 2.31, to the Neuron 2.31 package versions. Starting with Neuron 2.33.0, torch-neuronx for PyTorch 2.13 and 2.14 is TorchNeuron, the native PyTorch backend for Trainium. To update TorchNeuron, use the PyTorch 2.14 tab.

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.

# Activate your existing Python virtual environment
source aws_neuron_venv_pytorch/bin/activate

# 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

# Update Neuron Compiler and Framework (Neuron 2.31 versions)
python -m pip install --upgrade neuronx-cc==2.26.6360.0 torch-neuronx==2.9.0.2.15.32035

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.

Note

PyTorch versions 2.7 and 2.8 are no longer supported on Neuron. If you are looking for setup instructions specific to PyTorch 2.7 and 2.8 on Amazon Linux 2023, Ubuntu 24.04, or Ubuntu 22.04, see the Neuron release 2.28.0 version of the setup docs.

Update PyTorch on Ubuntu 26.04#

On Ubuntu 26.04, Neuron supports TorchNeuron (torch-neuronx) for PyTorch 2.14 on Trn2 and Trn3 instances. Activate your existing virtual environment, then update TorchNeuron, the Neuron Compiler, and NKI:

source aws_neuron_venv_pytorch_2_14/bin/activate

python -m pip install --upgrade "torch-neuronx==2.14.*" \
  "neuronx-cc==2.*" nki \
  --extra-index-url https://pip.repos.neuron.amazonaws.com \
  --extra-index-url https://download.pytorch.org/whl/cpu

Update PyTorch on Amazon Linux 2023#

If you already have a previous Neuron release installed, select the PyTorch version tab below to get the update commands for your environment.

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.

Note

PyTorch versions 2.7 and 2.8 are no longer supported on Neuron. If you are looking for setup instructions specific to PyTorch 2.7 and 2.8 on Amazon Linux 2023, Ubuntu 24.04, or Ubuntu 22.04, see the Neuron release 2.28.0 version of the setup docs.

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.

Note

PyTorch versions 2.7 and 2.8 are no longer supported on Neuron. If you are looking for setup instructions specific to PyTorch 2.7 and 2.8 on Amazon Linux 2023, Ubuntu 24.04, or Ubuntu 22.04, see the Neuron release 2.28.0 version of the setup docs.

Update Neuron driver and runtime#

Update the Neuron driver, runtime, and tools on your host. This is recommended when updating to a new Neuron SDK release.

sudo apt-get update
sudo apt-get install -y aws-neuronx-dkms
sudo apt-get install -y aws-neuronx-runtime-lib
sudo apt-get install -y aws-neuronx-collectives
sudo apt-get install -y aws-neuronx-tools

sudo apt-get update
sudo apt-get install -y aws-neuronx-dkms
sudo apt-get install -y aws-neuronx-runtime-lib
sudo apt-get install -y aws-neuronx-collectives
sudo apt-get install -y aws-neuronx-tools

Note

PyTorch versions 2.9 and earlier are not supported by Neuron for Amazon Linux 2023. If you are looking for setup instructions specific to PyTorch 2.9 on Amazon Linux 2023, see the Neuron release 2.28.0 version of the setup docs.

sudo dnf install -y aws-neuronx-dkms
sudo dnf install -y aws-neuronx-runtime-lib
sudo dnf install -y aws-neuronx-collectives
sudo dnf install -y aws-neuronx-tools

Verify the update#

After updating, activate your virtual environment and verify:

source aws_neuron_venv_pytorch_2_14/bin/activate

python3 << EOF
import torch
import torch_neuronx

print(f"PyTorch version: {torch.__version__}")
print(f"torch-neuronx version: {torch_neuronx.__version__}")

import subprocess
result = subprocess.run(['neuron-ls'], capture_output=True, text=True)
print(result.stdout)
EOF

You should see output similar to the following example from a trn2.48xlarge instance. Your versions, instance type, and device details depend on your installation and instance.

PyTorch version: 2.14.0+cpu
torch-neuronx version: 2.14.3.0
instance-type: trn2.48xlarge
logical-neuroncore-config: 2
+--------+--------+----------+--------+------+
| NEURON | NEURON |  NEURON  | NEURON | NUMA |
| DEVICE | CORES  | CORE IDS | MEMORY | NODE |
+--------+--------+----------+--------+------+
| 0      | 4      | 0-3      | 96 GB  | 1    |
| ...    | ...    | ...      | ...    | ...  |
| 15     | 4      | 60-63    | 96 GB  | 1    |
+--------+--------+----------+--------+------+

Install previous releases#

If you need to install older Neuron releases on your instance, follow the instructions below. Neuron supports Ubuntu 26.04 starting with Neuron 2.33.0, so there are no previous releases for Ubuntu 26.04.

Install previous releases on Ubuntu 24.04#

Use the tabs below to install a specific previous Neuron SDK release of PyTorch NeuronX on Ubuntu 24.04. Select the Neuron version you need.

# 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.26.6360.0 torch-neuronx==2.9.0.2.15.32035

# 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.25.3371.0 torch-neuronx==2.9.0.2.14.27725

# 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.24.8799.0 torch-neuronx==2.9.0.2.13.26312

Install previous releases on Ubuntu 22.04#

Important

Ubuntu 22.04 has reached end-of-support on Neuron. Neuron no longer provides Ubuntu 22.04 DLAMIs or container images. New deployments should use Ubuntu 24.04. See Neuron no longer includes Ubuntu 22.04 DLAMIs and DLCs starting this release.

Use the tabs below to install a specific previous Neuron SDK release of PyTorch NeuronX on Ubuntu 22.04. Select the Neuron version you need.

# Install Python (PyTorch 2.9 requires Python 3.11-3.13, newer than the Ubuntu 22.04 default)
sudo add-apt-repository ppa:deadsnakes/ppa
sudo apt-get install python3.13

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

# Create Python venv
python3.13 -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.13 -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.26.6360.0 torch-neuronx==2.9.0.2.15.32035

# Install Python (PyTorch 2.9 requires Python 3.11-3.13, not available by default on Ubuntu 22.04)
sudo add-apt-repository ppa:deadsnakes/ppa
sudo apt-get install python3.13

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

# Create Python venv
python3.13 -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.13 -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.25.3371.0 torch-neuronx==2.9.0.2.14.27725

# 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.24.8799.0 torch-neuronx==2.9.0.2.13.26312

Install previous releases on Amazon Linux 2023#

Use the tabs below to install a specific previous Neuron SDK release of PyTorch NeuronX on Amazon Linux 2023. Select the Neuron version you need.

Note

Currently, PyTorch 2.9 is not available on Amazon Linux 2023 and PyTorch 2.7 and 2.8 are no longer supported for Neuron. Use Ubuntu 24.04 for PyTorch 2.9 support. If you are using Neuron 2.28.0, see the Amazon Linux 2023 setup documentation in the 2.28.0 version of the Neuron docs.

Note

Currently, PyTorch 2.9 is not available on Amazon Linux 2023 and PyTorch 2.7 and 2.8 are no longer supported for Neuron. Use Ubuntu 24.04 for PyTorch 2.9 support. If you are using Neuron 2.28.0, see the Amazon Linux 2023 setup documentation in the 2.28.0 version of the Neuron docs.

Note

Currently, PyTorch 2.9 is not available on Amazon Linux 2023 and PyTorch 2.7 and 2.8 are no longer supported for Neuron. Use Ubuntu 24.04 for PyTorch 2.9 support. If you are using Neuron 2.28.0, see the Amazon Linux 2023 setup documentation in the 2.28.0 version of the Neuron docs.

This document is relevant for: Trn2, Trn3