This document is relevant for: Inf1, Inf2, Trn1, Trn1n

PyTorch Neuron#

PyTorch Neuron unlocks high-performance and cost-effective deep learning acceleration on AWS Trainium-based and Inferentia-based Amazon EC2 instances.

PyTorch Neuron plugin architecture enables native PyTorch models to be accelerated on Neuron devices, so you can use your existing framework application and get started easily with minimal code changes.

Pytorch Neuron Setup
Install Drivers and Tools
# Configure Linux for Neuron repository updates
sudo tee /etc/yum.repos.d/neuron.repo > /dev/null <<EOF
[neuron]
name=Neuron YUM Repository
baseurl=https://yum.repos.neuron.amazonaws.com
enabled=1
metadata_expire=0
EOF
sudo rpm --import https://yum.repos.neuron.amazonaws.com/GPG-PUB-KEY-AMAZON-AWS-NEURON.PUB

# Update OS packages 
sudo yum update -y

# Install OS headers 
sudo yum install kernel-devel-$(uname -r) kernel-headers-$(uname -r) -y

# Install git 
sudo yum install git -y

# install Neuron Driver
sudo yum install aws-neuronx-dkms-2.* -y

# Install Neuron Tools 
sudo yum install aws-neuronx-tools-2.* -y

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

# 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 Tools 
sudo apt-get install aws-neuronx-tools=2.* -y

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

Note

Neuron Driver installed on Deep Learning AMI (DLAMI) with Conda does not support Inf2\Trn1.

If you want to use DLAMI with Conda, please make sure to uninstall aws-neuron-dkms and install aws-neuronx-dkms before using Neuron on DLAMI with Conda.

Note

If you are facing a timeout 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.

Neuron SDK team highly recommends to launch the instances using DLAMIs since they already have a fix for any such problems.

Tutorials

Note

To use Jupyter Notebook see:

Additional Examples
API Reference Guide
Developer Guide
Misc

Note

For help selecting a framework type, see:

Comparison of torch-neuron (Inf1) versus torch-neuronx (Inf2 & Trn1) for Inference

Tutorials (torch-neuronx)

Note

To use Jupyter Notebook see:

Additional Examples (torch-neuronx)
API Reference Guide (torch-neuronx)
Developer Guide (torch-neuronx)
Misc (torch-neuronx)
Transformers Neuron (transformers-neuronx)
Tutorials (torch-neuron)

Note

To use Jupyter Notebook see:

Additional Examples (torch-neuron)
API Reference Guide (torch-neuron)
Developer Guide (torch-neuron)
Misc (torch-neuron)

This document is relevant for: Inf1, Inf2, Trn1, Trn1n