PyTorch Neuron (“torch-neuronx”) Setup on Ubuntu 22 — AWS Neuron Documentation (original) (raw)
Contents
- Get Started with Latest Release of PyTorch Neuron (torch-neuronx)
- Update to latest PyTorch NeuronX
- Install Previous PyTorch NeuronX Releases for Ubuntu
This document is relevant for: Inf2, Trn1, Trn2
Select a Different Framework or Platform for Setup
PyTorch Neuron (“torch-neuronx”) Setup on Ubuntu 22#
Table of contents
- Get Started with Latest Release of PyTorch Neuron (torch-neuronx)
- Update to latest PyTorch NeuronX
- Install Previous PyTorch NeuronX Releases for Ubuntu
Get Started with Latest Release of PyTorch Neuron (torch-neuronx)#
This section provide links that will assist you to quickly start with a fresh installation of PyTorch Neuron for both Inference and Training.
Note
If you are facing a connectivity issue during the model loading process on a Trn1 instance with Ubuntu, that could probably be because of Ubuntu limitations with multiple interfaces. To solve this problem, please follow the steps mentioned here.
Users are highly encouraged to use DLAMI to launch the instances, since DLAMIs come with the required fix.
- Please follow the instructions at launch an Amazon EC2 Instance to launch an instance. When choosing the instance type at the EC2 console, please make sure to select the correct instance type.
- To get more information about instances sizes and pricing see: Trn1 web page, Inf2 web page
- Select Ubuntu Server 22 AMI
- When launching a Trn1, please adjust your primary EBS volume size to a minimum of 512GB.
- After launching the instance, follow the instructions in Connect to your instance to connect to the instance
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
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
PyTorch 2.6.0
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 torchvision
PyTorch 2.5.1
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.5.* torchvision
Visit PyTorch Neuron(torch-neuronx) for Inference section
Visit PyTorch Neuron(torch-neuronx) for Training section
Update to latest PyTorch NeuronX#
If you already have a previous Neuron release installed, this section provide links that will assist you to update to latest Neuron release.
PyTorch 2.6.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, please 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 torchvision
PyTorch 2.5.1
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, please 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.5.* torchvision
Install Previous PyTorch NeuronX Releases for Ubuntu#
This section will assist you to install previous Neuron releases.
Neuron 2.22.0
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.17.194.0 torch-neuronx==2.5.1.2.6.0
Neuron 2.21.0
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.16.345.0 torch-neuronx==2.5.1.2.4.0
Neuron 2.20.0
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.15.128.0 torch-neuronx==2.1.2.2.3.0
This document is relevant for: Inf2, Trn1, Trn2