Pytorch Neuron Setup
This document is relevant for: Inf1
, Inf2
, Trn1
, Trn1n
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.
This document is relevant for: Inf1
, Inf2
, Trn1
, Trn1n