Neuron Support

You can use Aphrodite to run models on AWS Trainium1/Inferentia2 chips with the Neuron SDK. Keep in mind that the Neuron implementation does not use Paged Attention, so it won't be as effective.

Requirements

• OS: Linux
• Python: 3.8 - 3.11
• Accelerate: NeuronCore_v2
• PyTorch 2.0.1/2.1.1
• AWS Neuron SDK 2.16/2.17

Building

Launch Trainium1/Inferentia2 instances

• Follow the instructions here to launch an instance. When choosing the instance type at the EC2 console, please make sure to select the correct instance type.
• For information on pricing, see these pages: Trainium1, Inferentia2
• Select Ubuntu Server 22.04 TLS AMI.
• When launching an instance, adjust your primary EBS volume to a minimum of 512 GiB.
• After launching the instance, follow the instructions here to connect to your instance.

Install Drivers

This step is optional, as systems with Deep Learning AMI Neuron will have these pre-installed. If they don't for any reason, run these:

. /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 -


sudo apt-get update -y

sudo apt-get install linux-headers-$(uname -r) -y

sudo apt-get install git -y

sudo apt-get install aws-neuronx-dkms=2.* aws-neuronx-collectives=2.* aws-neuronx-runtime-lib=2.* aws-neuronx-tools=2.* -y

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

Install Transformers-Neuronx

sudo apt-get install -y python3.10-venv g++

python3.10 -m venv aws_neuron_venv_pytorch

source aws_neuron_venv_pytorch/bin/activate

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

python -m pip config set global.extra-index-url https://pip.repos.neuron.amazonaws.com

python -m pip install wget awscli

python -m pip install --upgrade neuronx-cc==2.* --pre torch-neuronx==2.1.* torchvision transformers-neuronx

Build Aphrodite

git clone https://github.com/PygmalionAI/aphrodite-engine.git && cd aphrodite-engine

pip install -U -r requirements-neuron.txt

pip install -e .