This document is relevant for: Inf1, Inf2, Trn1, Trn2, Trn3
Update PyTorch on a Deep Learning AMI#
Update PyTorch and Neuron components on an existing DLAMI to the latest release.
Update PyTorch on Ubuntu 24.04#
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.
# Activate Python venv
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
python -m pip install --upgrade neuronx-cc==2.* torch-neuronx==2.9.* torchvision
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.
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
Update PyTorch on Ubuntu 22.04#
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.
# Activate Python venv
source aws_neuron_venv_pytorch/bin/activate
# 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
# Update Neuron Compiler and Framework
python -m pip install --upgrade neuronx-cc==2.* torch-neuronx==2.9.* torchvision
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.
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
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.
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
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.
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
Update Neuron driver and runtime#
Update the Neuron driver, runtime, and tools on your DLAMI 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
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:
source /opt/aws_neuronx_venv_pytorch/bin/activate
And verify the update:
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 |
+--------+--------+--------+-----------+
Previous releases#
To install a specific previous Neuron SDK release:
This document is relevant for: Inf1, Inf2, Trn1, Trn2, Trn3