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

Get Started with TensorFlow Neuron#

This page provide links that will assist you to quickly start with TensorFlow Neuron.

Note

Below instructions are for Ubuntu20, if you looking for complete setup instructions for different platforms, please Check Here.

Launch the Instance

  • 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, Inf1 web page

  • Select your Amazon Machine Image (AMI) of choice, please note that Neuron supports Amazon Linux 2 AMI(HVM) - Kernel 5.10.

  • 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

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.

Install Drivers and Tools
# 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
Install EFA (Applicable only for Trn1)
# 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
Install Tensorflow Neuron (tensorflow-neuronx)
# 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_tensorflow 

# Activate Python venv 
source aws_neuron_venv_tensorflow/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_tensorflow --display-name "Python (tensorflow-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.* tensorflow-neuronx
Get Started with Inference (Inf2 & Trn1)

Roberta-Base

Visit TensorFlow Neuron section for more
Install TensorFlow Neuron (tensorflow-neuron)
# Install Python venv 
sudo apt-get install -y python3.8-venv g++ 

# Create Python venv
python3.8 -m venv aws_neuron_venv_tensorflow 

# Activate Python venv 
source aws_neuron_venv_tensorflow/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_tensorflow --display-name "Python (tensorflow-neuron)"
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 TensorFlow Neuron
python -m pip install tensorflow-neuron[cc] "protobuf"

# Install Neuron TensorBoard
python -m pip install tensorboard-plugin-neuron

# Optional: Install Tensorflow Neuron model server
sudo apt-get install tensorflow-model-server-neuronx=2.10.1.2.8.1.0 -y
Get Started with Inference (Inf1)
Visit TensorFlow Neuron section for more

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