此页面由 Cloud Translation API 翻译。

在 v6e TPU 上进行 vLLM 推理

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

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

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

准备工作

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

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

登录您的 Google 账号。如果您还没有 Google 账号，请注册新账号。
在 Google Cloud 控制台中，从项目选择器页面选择或创建 Google Cloud 项目。
为您的 Google Cloud 项目启用结算功能。所有 Google Cloud 使用情况都需要结算。
安装 gcloud alpha 组件。
运行以下命令以安装最新版本的 gcloud 组件。
```
gcloud components update
```
使用 Cloud Shell 通过以下 gcloud 命令启用 TPU API。您也可以从 Google Cloud 控制台启用。
```
gcloud services enable tpu.googleapis.com
```

为 TPU 虚拟机创建服务身份。

gcloud alpha compute tpus tpu-vm service-identity create --zone=ZONE

创建 TPU 服务账号并授予对 Google Cloud 服务的访问权限。

通过服务账号， Google Cloud TPU 服务可以访问其他 Google Cloud服务。建议使用用户代管式服务账号。请按照以下指南创建和授予角色。您需要拥有以下角色：
- TPU 管理员：创建 TPU 所需
- Storage Admin：需要此角色才能访问 Cloud Storage
- 日志写入器：需要使用 Logging API 写入日志
- Monitoring Metric Writer：用于将指标写入 Cloud Monitoring

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

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

保障容量

请与您的 Cloud TPU 销售团队或客户支持团队联系，申请 TPU 配额并咨询容量方面的任何问题。

预配 Cloud TPU 环境

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

前提条件

验证您的项目是否有足够的 TPUS_PER_TPU_FAMILY 配额，该配额指定您可以在Google Cloud 项目中访问的芯片数量上限。
本教程使用以下配置进行了测试：
- Python 3.10 or later
- 每夜软件版本：
  - 每夜 JAX 0.4.32.dev20240912
  - 每夜 LibTPU 0.1.dev20240912+nightly
- 稳定版软件版本：
  - v0.4.35 的 JAX + JAX 库
验证您的项目是否有足够的 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.10 -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 setup.py develop
sudo apt-get install libopenblas-base libopenmpi-dev libomp-dev

获取对模型的访问权限

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

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

下载基准比较数据

创建一个 /data 目录，然后从 Hugging Face 下载 ShareGPT 数据集。

mkdir ~/data && cd ~/data
wget https://huggingface.co/datasets/anon8231489123/ShareGPT_Vicuna_unfiltered/resolve/main/ShareGPT_V3_unfiltered_cleaned_split.json

启动 vLLM 服务器

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

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

   cd ~/vllm
   vllm serve "meta-llama/Meta-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 基准测试脚本：

   python benchmarks/benchmark_serving.py \
       --backend vllm \
       --model "meta-llama/Meta-Llama-3.1-8B"  \
       --dataset-name sharegpt \
       --dataset-path ~/data/ShareGPT_V3_unfiltered_cleaned_split.json  \
       --num-prompts 1000

清理

删除 TPU：

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