This document is relevant for: Inf2, Trn1, Trn1n

Install PyTorch Neuron (torch-neuronx)#

Develop on AWS ML accelerator instance#

Note

  • For a successful installation or update, execute each line of the instructions below separately or copy the contents of the code block into a script file and source its contents.

  • When launching a Trn1, please adjust your primary EBS volume size to a minimum of 512GB.

  • While launching the instance, please use the AMI with the name Deep Learning Base Neuron AMI (Amazon Linux 2) <Latest_Date>.

  • To launch an instance using a specific AMI, please refer to the instructions mentioned here.

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

# Install Python venv 
sudo yum install -y python3.7-venv gcc-c++ 

# Create Python venv
python3.7 -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.7 -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

Note

  • For a successful installation or update, execute each line of the instructions below separately or copy the contents of the code block into a script file and source its contents.

  • When launching a Trn1, please adjust your primary EBS volume size to a minimum of 512GB.

  • While launching the instance, please use the AMI with the name Deep Learning Base Neuron AMI (Ubuntu 20.04) <Latest_Date>.

  • To launch an instance using a specific AMI, please refer to the instructions mentioned here.

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

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

# Create Python venv
python3.8 -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.8 -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

Note

  • When launching a Trn1, please adjust your primary EBS volume size to a minimum of 512GB.

  • While launching the instance, please use the AMI with the name Deep Learning AMI Neuron PyTorch 1.13.0 (Amazon Linux 2) <Latest_Date>.

  • To launch an instance using a specific AMI, please refer to the instructions mentioned here.

# Update Neuron Runtime
sudo yum update aws-neuronx-collectives-2.* -y
sudo yum update aws-neuronx-runtime-lib-2.* -y

# Activate Python venv 
source /opt/aws_neuron_venv_pytorch/bin/activate 

# Install Jupyter notebook kernel
pip install ipykernel 
python3.7 -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

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

Note

  • When launching a Trn1, please adjust your primary EBS volume size to a minimum of 512GB.

  • While launching the instance, please use the AMI with the name Deep Learning AMI Neuron PyTorch 1.13.0 (Ubuntu 20.04) <Latest_Date>.

  • To launch an instance using a specific AMI, please refer to the instructions mentioned here.

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

# Activate Python venv 
source /opt/aws_neuron_venv_pytorch/bin/activate 

# Install Jupyter notebook kernel
pip install ipykernel 
python3.8 -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

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

Note

  • Please refer to the instructions under the tab Amazon Linux 2 DLAMI Base.

Note

  • Please refer to the instructions under the tab Ubuntu 20 DLAMI Base.

This document is relevant for: Inf2, Trn1, Trn1n