通过 vLLM 在 TPU 上部署 Qwen2-7B-Instruct

本教程使用 v6e TPU 虚拟机上的 vLLM TPU 服务框架来部署 Qwen/Qwen2-7B-Instruct 模型。

目标

1. 设置环境。
2. 使用 Qwen2-7B-Instruct 运行 vLLM。
3. 发送推理请求。
4. 运行基准测试工作负载。
5. 清理。

费用

本教程使用 Google Cloud 的以下收费组件：

• Cloud TPU
• Compute Engine

如需根据您的预计使用量来估算费用，请使用价格计算器。

准备工作

在学习本教程之前，请按照设置 Cloud TPU 环境页面中的说明操作。这些说明将引导您完成创建 Google Cloud项目并将其配置为使用 Cloud TPU 所需的步骤。您也可以使用现有的Google Cloud 项目。如果您选择这样做，可以跳过创建Google Cloud 项目的步骤，直接从设置使用 Cloud TPU 的环境开始。

您需要拥有 Hugging Face 访问令牌才能使用本教程。您可以前往 Hugging Face 注册免费账号。拥有账号后，生成访问令牌：

1. 在 Welcome to Hugging Face 页面上，点击您的账号头像，然后选择 Access tokens。
2. 在访问令牌页面上，点击创建新令牌。
3. 选择 Read 令牌类型，然后输入令牌的名称。
4. 系统会显示您的访问令牌。将令牌保存在安全的位置。

设置环境

1. 使用 Queued Resources API 创建 Cloud TPU v6e 虚拟机。对于 qwen2-7b-instruct，我们建议使用 v6e-1 TPU。

   export PROJECT_ID=YOUR_PROJECT_ID
   export TPU_NAME=qwen2-7b-instruct-tutorial
   export ZONE=us-east5-a
   export QR_ID=qwen2-7b-instruct-qr
   
   gcloud alpha compute tpus queued-resources create $QR_ID \
    --node-id $TPU_NAME \
    --project $PROJECT_ID \
    --zone $ZONE \
    --accelerator-type v6e-1 \
    --runtime-version v2-alpha-tpuv6e
   

2. 检查以确保 TPU 虚拟机已准备就绪。

   gcloud compute tpus queued-resources describe $QR_ID \
     --project $PROJECT_ID \
     --zone $ZONE
   

   创建 TPU 虚拟机后，已排队的资源请求的状态将设置为 ACTIVE。例如：

   name: projects/your-project-id/locations/your-zone/queuedResources/your-queued-resource-id
     state:
     state: ACTIVE
     tpu:
     nodeSpec:
     - node:
         acceleratorType: v6e-1
         bootDisk: {}
         networkConfig:
             enableExternalIps: true
         queuedResource: projects/your-project-number/locations/your-zone/queuedResources/your-queued-resource-id
         runtimeVersion: v2-alpha-tpuv6e
         schedulingConfig: {}
         serviceAccount: {}
         shieldedInstanceConfig: {}
         useTpuVm: true
         nodeId: your-node-id
         parent: projects/your-project-number/locations/your-zone
   

3. 连接到 TPU 虚拟机。

     gcloud compute tpus tpu-vm ssh $TPU_NAME \
       --project $PROJECT_ID \
       --zone $ZONE
   

使用 Qwen2-7B-instruct 运行 vLLM

1. 在 TPU 虚拟机中，运行 vLLM Docker 容器。此命令使用 10 GB 的共享内存大小。

     export DOCKER_URI=vllm/vllm-tpu:latest
   
     sudo docker run -it --rm --name $USER-vllm --privileged --net=host \
       -v /dev/shm:/dev/shm \
       --shm-size 10gb \
       -p 8000:8000 \
       --entrypoint /bin/bash ${DOCKER_URI}
   

2. 在容器内，设置您的 Hugging Face 令牌。将 YOUR_HF_TOKEN 替换为您的 Hugging Face 令牌。

   export HF_HOME=/dev/shm
   export HF_TOKEN=YOUR_HF_TOKEN
   

3. 使用 vllm serve 命令启动 vLLM 服务器。

   export MAX_MODEL_LEN=4096
   export TP=1 # number of chips
   
   vllm serve Qwen/Qwen2-7B-Instruct \
       --seed 42 \
       --disable-log-requests \
       --gpu-memory-utilization 0.98 \
       --max-num-batched-tokens 1024 \
       --max-num-seqs 128 \
       --tensor-parallel-size $TP \
       --max-model-len $MAX_MODEL_LEN
   

   当 vLLM 服务器运行时，您将看到如下所示的输出：

   (APIServer pid=7) INFO:     Started server process [7]
   (APIServer pid=7) INFO:     Waiting for application startup.
   (APIServer pid=7) INFO:     Application startup complete.
   

发送推理请求

vLLM 服务器运行后，您可以通过新 shell 向其发送请求。

1. 打开新 shell 并连接到 TPU 虚拟机。

     export PROJECT_ID=YOUR_PROJECT_ID
     export TPU_NAME=qwen2-7b-instruct-tutorial
     export ZONE=us-east5-a
   
     gcloud compute tpus tpu-vm ssh $TPU_NAME \
       --project $PROJECT_ID \
       --zone=$ZONE
   

2. 打开一个 shell 进入正在运行的 Docker 容器。

     sudo docker exec -it $USER-vllm /bin/bash
   

3. 使用 curl 向服务器发送测试请求。

     curl http://localhost:8000/v1/completions \
       -H "Content-Type: application/json" \
       -d '{
           "model": "Qwen/Qwen2-7B-Instruct",
           "prompt": "The future of AI is",
           "max_tokens": 200,
           "temperature": 0
         }'
   

响应以 JSON 格式返回。

运行基准工作负载

您可以在第二个终端中针对正在运行的服务器运行基准测试。

1. 在容器内，安装 datasets 库。

   pip install datasets
   

2. 运行 vllm bench serve 命令。

   export HF_HOME=/dev/shm
   cd /workspace/vllm
   
   vllm bench serve \
       --backend vllm \
       --model "Qwen/Qwen2-7B-Instruct"  \
       --dataset-name random \
       --num-prompts 1000 \
       --seed 100
   

基准比较结果如下所示：

============ Serving Benchmark Result ============
Successful requests:                     1000
Benchmark duration (s):                  45.35
Total input tokens:                      1024000
Total generated tokens:                  126848
Request throughput (req/s):              22.05
Output token throughput (tok/s):         2797.15
Peak output token throughput (tok/s):    4258.00
Peak concurrent requests:                1000.00
Total Token throughput (tok/s):          25377.57
---------------Time to First Token----------------
Mean TTFT (ms):                          21332.46
Median TTFT (ms):                        21330.37
P99 TTFT (ms):                           42436.47
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          37.36
Median TPOT (ms):                        38.56
P99 TPOT (ms):                           38.69
---------------Inter-token Latency----------------
Mean ITL (ms):                           37.35
Median ITL (ms):                         38.55
P99 ITL (ms):                            39.43
==================================================

清理

为避免因本教程中使用的资源导致您的 Google Cloud 账号产生费用，请删除包含这些资源的项目，或者保留项目但删除各个资源。

1. 在第二个 shell 中，输入 exit 以退出 vLLM 容器。
2. 在第二个 shell 中，输入 exit 命令以关闭终端。
3. 在第一个 shell 中，输入 Ctrl+C 以停止 vLLM 服务器。
4. 在第一个 shell 中，输入 exit 以退出 vLLM 容器。
5. 在第一个 shell 中，输入 exit 以断开与 TPU 虚拟机的连接。

删除您的资源

您可以删除项目，这样会删除所有资源；也可以保留项目，但删除资源。

删除项目

如需删除 Google Cloud 项目和所有关联的资源，请运行以下命令：

    gcloud projects delete $PROJECT_ID

删除 TPU 资源

删除 Cloud TPU 资源。以下命令使用 --force 参数同时删除已排队的资源请求和 TPU 虚拟机。

  gcloud alpha compute tpus queued-resources delete $QR_ID \
    --project=$PROJECT_ID \
    --zone=$ZONE \
    --force

后续步骤

• 详细了解 Cloud TPU 上的 vLLM。
• 详细了解 Cloud TPU。