跳到主要内容

用 Neuron 安装

从 vLLM 0.3.3 版本开始,支持在 AWS Trainium/Inferentia 上使用 Neuron SDK 进行模型推理和服务。目前,Neuron SDK 不支持 Paged Attention,但在 transformers-neuronx 中支持简单的连续批处理。Neuron SDK 当前支持的数据类型是 FP16 和 BF16。

环境依赖

  • OS: Linux
  • Python: 3.8 – 3.11
  • Accelerator: NeuronCore_v2 (in trn1/inf2 instances)
  • Pytorch 2.0.1/2.1.1
  • AWS Neuron SDK 2.16/2.17 (在python 3.8已验证)

从源码构建

以下说明适用于 Neuron SDK 2.16 及以上版本。

第1步、 启动 Trn1/Inf2 实例

下面是启动 trn1/inf2 实例的步骤,为了安装 PyTorch Neuron (“torch-neuronx”) Setup on Ubuntu 22.04 LTS.

  • 请按照说明操作在 Amazon EC2 上启动一个实例。在选择 EC2 控制台中的实例类型时,请确保选择正确的实例类型。
  • 要了解有关实例大小和定价的更多信息,请查看: Trn1 网页, Inf2 网页
  • 选择 Ubuntu Server 22.04 TLS AMI
  • 当启动 Trn1/Inf2 时, 请将主 EBS 卷大小调整至至少 512GB。
  • 在启动实例后,按照说明进行操作以连接到该实例。

第2步、安装驱动和工具

如果没有安装Deep Learning AMI Neuronis,请按照以下操作安装驱动和工具:

# 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 Runtime
sudo apt-get install aws-neuronx-collectives=2.* -y
sudo apt-get install aws-neuronx-runtime-lib=2.* -y

# Install Neuron Tools
sudo apt-get install aws-neuronx-tools=2.* -y

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

第3步、安装transformers-neuronx及其依赖

transformers-neuronx 作为后端,以支持在 trn1/inf2 实例上进行推理。请按照以下步骤安装 transformer-neuronx 包及其依赖项。

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

# Create Python venv
python3.10 -m venv aws_neuron_venv_pytorch

# Activate Python venv
source aws_neuron_venv_pytorch/bin/activate

# Install Jupyter notebook kernel
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

# 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

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

第4步、从源码安装vLLM

安装好 neuronx-cc 和 transformers-neuronx 包,我们就可以根据以下步骤安装 vllm

cd vllm
pip install -U -r requirements-neuron.txt
pip install .

如果安装过程中正确检测到了 neuron 包,将会安装 vllm-0.3.0+neuron212。