此页面由 Cloud Translation API 翻译。

在 v6e TPU 上进行 vLLM 推理

本教程介绍了如何在 v6e TPU 上运行 vLLM 推理。还介绍了如何针对 Meta Llama-3.1 8B 模型运行基准测试脚本。

如需开始在 v6e TPU 上使用 vLLM，请参阅 vLLM 快速入门。

如果您使用的是 GKE，请参阅 GKE 教程。

准备工作

您必须签署同意协议，才能使用 HuggingFace 代码库中的 Llama3 系列模型。前往 meta-llama/Llama-3.1-8B，填写同意协议，然后等待批准。

准备预配具有 4 个芯片的 TPU v6e：

按照设置 Cloud TPU 环境指南设置 Google Cloud 项目、配置 Google Cloud CLI、启用 Cloud TPU API，并确保您有权使用 Cloud TPU。

使用 Google Cloud 进行身份验证，并为 Google Cloud CLI 配置默认项目和区域。

gcloud auth login
gcloud config set project PROJECT_ID
gcloud config set compute/zone ZONE

保障容量

当您准备好预订 TPU 容量时，请参阅 Cloud TPU 配额，详细了解 Cloud TPU 配额。如果您对如何确保容量还有其他疑问，请与您的 Cloud TPU 销售团队或客户支持团队联系。

预配 Cloud TPU 环境

您可以使用 GKE、GKE 和 XPK 预配 TPU 虚拟机，也可以将其作为队列化资源预配。

前提条件

验证您的项目是否有足够的 TPUS_PER_TPU_FAMILY 配额，该配额指定您可以在Google Cloud 项目中访问的芯片数量上限。
验证您的项目是否有足够的 TPU 配额：
- TPU 虚拟机配额
- IP 地址配额
- Hyperdisk Balanced 配额
用户项目权限
- 如果您将 GKE 与 XPK 搭配使用，请参阅用户账号或服务账号的 Cloud 控制台权限，了解运行 XPK 所需的权限。

预配 TPU v6e

   gcloud alpha compute tpus queued-resources create QUEUED_RESOURCE_ID \
      --node-id TPU_NAME \
       --project PROJECT_ID \
       --zone ZONE \
       --accelerator-type v6e-4 \
       --runtime-version v2-alpha-tpuv6e \
       --service-account SERVICE_ACCOUNT

命令标志说明

变量	说明
NODE_ID	在队列中的资源请求分配时创建的 TPU 的用户分配 ID。
PROJECT_ID	Google Cloud 项目名称。使用现有项目或创建新项目。
ZONE	如需了解支持的区域，请参阅 TPU 区域和可用区文档。
ACCELERATOR_TYPE	如需了解支持的加速器类型，请参阅加速器类型文档。
RUNTIME_VERSION	`v2-alpha-tpuv6e`
SERVICE_ACCOUNT	这是您的服务账号的电子邮件地址，您可以在 Google Cloud 控制台 -> IAM -> 服务账号中找到。例如：tpu-service-account@<your_project_ID>.iam.gserviceaccount.com.com

使用 list 或 describe 命令查询队列中资源的状态。

gcloud alpha compute tpus queued-resources describe QUEUED_RESOURCE_ID  \
   --project PROJECT_ID --zone ZONE

如需查看已加入队列的资源请求状态的完整列表，请参阅已加入队列的资源文档。

使用 SSH 连接到 TPU

  gcloud compute tpus tpu-vm ssh TPU_NAME

安装依赖项

为 Miniconda 创建一个目录：
```
mkdir -p ~/miniconda3
```

下载 Miniconda 安装程序脚本：

wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh -O ~/miniconda3/miniconda.sh

安装 Miniconda：

bash ~/miniconda3/miniconda.sh -b -u -p ~/miniconda3

移除 Miniconda 安装程序脚本：
```
rm -rf ~/miniconda3/miniconda.sh
```

将 Miniconda 添加到 PATH 变量：

export PATH="$HOME/miniconda3/bin:$PATH"

重新加载 ~/.bashrc 以将更改应用于 PATH 变量：
```
source ~/.bashrc
```
创建 Conda 环境：

注意：如果这是您首次运行 Conda，则需要先运行 conda init 并重新加载 shell，然后再运行 conda activate vllm。
```
conda create -n vllm python=3.11 -y
conda activate vllm
```

克隆 vLLM 仓库并前往 vllm 目录：

git clone https://github.com/vllm-project/vllm.git && cd vllm

清理现有的 torch 和 torch-xla 软件包：
```
pip uninstall torch torch-xla -y
```
注意：如果您未安装这些软件包，则会看到一条错误消息，提示您未安装这些软件包。您可以放心地忽略这些错误消息。
安装其他构建依赖项：
```
pip install -r requirements/tpu.txt
VLLM_TARGET_DEVICE="tpu" python -m pip install --editable .
sudo apt-get install libopenblas-base libopenmpi-dev libomp-dev
```
注意：如果您看到类似以下错误：

- E: Unable to locate package libopenblas-base
- E: Unable to locate package libopenmpi-dev
- E: Package 'libomp-dev' has no installation candidate

请先运行 sudo apt-get update，然后再运行安装命令。

获取对模型的访问权限

如果您还没有 Hugging Face 令牌，请生成一个新令牌：

点击您的个人资料 > 设置 > 访问令牌。
选择新建令牌 (New Token)。
指定您选择的名称和一个至少具有 Read 权限的角色。
选择生成令牌。
将生成的令牌复制到剪贴板，将其设置为环境变量，然后使用 huggingface-cli 进行身份验证：
```
export TOKEN=YOUR_TOKEN
git config --global credential.helper store
huggingface-cli login --token $TOKEN
```

启动 vLLM 服务器

以下命令会将模型权重从 Hugging Face 模型中心下载到 TPU 虚拟机的 /tmp 目录，预编译一系列输入形状，并将模型编译结果写入 ~/.cache/vllm/xla_cache。

如需了解详情，请参阅 vLLM 文档。

cd ~/vllm
vllm serve "meta-llama/Llama-3.1-8B" --download_dir /tmp --num-scheduler-steps 4 --swap-space 16 --disable-log-requests --tensor_parallel_size=4 --max-model-len=2048 &> serve.log &

运行 vLLM 基准测试

运行 vLLM 基准测试脚本：

export MODEL = <var>your-model</var>
python benchmarks/benchmark_serving.py \
   --backend vllm \
   --model $MODEL  \
   --dataset-name random \
   --random-input-len 1820 \
   --random-output-len 128 \
   --random-prefix-len 0

清理

删除 TPU：

gcloud compute tpus queued-resources delete QUEUED_RESOURCE_ID \
    --project PROJECT_ID \
    --zone ZONE \
    --force \
    --async