使用 Neuron 安装

使用 Neuron 安装#

从 vLLM 0.3.3 开始，支持在 AWS Trainium/Inferentia 上使用 Neuron SDK 进行模型推理和服务的连续批处理。分页注意力和分块预填充功能目前正在开发中，即将推出。Neuron SDK 目前支持的数据类型是 FP16 和 BF16。

要求#

  • 操作系统: Linux

  • Python: 3.8 – 3.11

  • 加速器：NeuronCore_v2（在 trn1/inf2 实例中）

  • Pytorch 2.0.1/2.1.1

  • AWS Neuron SDK 2.16/2.17 (在 python 3.8 上验证)

安装步骤：

备注

目前支持的用于 Neuron 的 Pytorch 版本安装了 triton 版本 2.1.0。这与 vLLM >= 0.5.3 不兼容。您可能会看到错误 cannot import name ‘default_dump_dir…’。要解决这个问题，请在安装 vLLM wheel 后运行 pip install –upgrade triton==3.0.0。

从源码构建#

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

步骤 0. 启动 Trn1/Inf2 实例#

以下是启动 trn1/inf2 实例的步骤，以便在 Ubuntu 22.04 LTS 上安装 PyTorch Neuron (“torch-neuronx”)。

  • 请按照 启动 Amazon EC2 实例 的说明来启动一个实例。在 EC2 控制台选择实例类型时，请确保选择正确的实例类型。

  • 要获取更多关于实例大小和定价的信息，请参阅：Trn1 网页，Inf2 网页

  • 选择 Ubuntu Server 22.04 TLS AMI

  • 在启动 Trn1/Inf2 时，请将您的主 EBS 卷大小调整为最小 512GB。

  • 启动实例后，按照 连接到您的实例 中的说明连接到实例

步骤 1. 安装驱动程序和工具#

如果安装了 Deep Learning AMI Neuron ，则无需安装驱动程序和工具。如果操作系统上未安装驱动程序和工具，请按照以下步骤操作：

# 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

步骤 2. 安装 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

步骤 3. 从源代码安装 vLLM#

一旦安装了 neuronx-cc 和 transformers-neuronx 包，我们将能够如下安装 vllm：

$ git clone https://github.com/vllm-project/vllm.git
$ cd vllm
$ pip install -U -r requirements-neuron.txt
$ VLLM_TARGET_DEVICE="neuron" pip install .

如果在安装过程中正确检测到神经元包，将安装 vllm-0.3.0+neuron212。