通过 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 环境页面中的说明操作。这些 说明将引导您完成创建 Google Cloud 项目并将其配置为使用 Cloud TPU 所需的步骤。您也可以使用现有 Google Cloud 项目。如果您选择这样做,可以跳过创建 Google Cloud 项目步骤,直接从设置环境以使用 Cloud TPU开始。

您需要拥有 Hugging Face 访问令牌才能使用本教程。您可以在 Hugging Face 上注册 免费账号。拥有账号后,请生成访问令牌:

  1. 欢迎使用 Hugging Face页面上, 点击您的账号头像,然后选择访问令牌
  2. 访问令牌 页面上,点击创建新令牌
  3. 选择读取 令牌类型,然后为令牌输入名称。
  4. 系统会显示您的访问令牌。请将令牌保存在安全的位置。

设置环境

已排入队列的资源

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

export PROJECT_ID=<PROJECT>
export TPU_NAME=<TPU_NAME>
export ZONE=<ZONE>
export QR_ID=<QR_ID>
export TPU_TYPE=<TPU_TYPE>

设置变量:

  • PROJECT - 您的项目的名称。
  • TPU_NAME - 您将创建的 TPU 虚拟机机器的名称。
  • ZONE - 您在其中创建新虚拟机的云可用区。
  • TPU_TYPE - 您创建的 TPU 虚拟机的类型。例如:v6e-1v6e-4
  • QR_ID - 您创建的已排入队列的资源的名称。

创建 Queued Resources 请求:

gcloud alpha compute tpus queued-resources create $QR_ID \
 --node-id $TPU_NAME \
 --project $PROJECT_ID \
 --zone $ZONE \
 --accelerator-type $TPU_TYPE \
 --runtime-version v2-alpha-tpuv6e

检查以确保您的 TPU 虚拟机已准备就绪。

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

例如,当状态为 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

预订

使用预订创建 Cloud TPU v6e 虚拟机。对于 qwen2-7b-instruct,我们建议使用 v6e-1 TPU。首先设置环境变量:

export PROJECT_ID="YOUR_PROJECT_ID"
export TPU_NAME="TPU_MACHINE_NAME"
export ZONE="ZONE_NAME"
export TPU_TYPE=v6e-8
export HF_TOKEN="YOUR_HF_TOKEN"
export RESERVATION="RESERVATION_NAME"

设置变量:

  • PROJECT - 您的项目的名称。
  • TPU_NAME - 您将创建的 TPU 虚拟机机器的名称。
  • ZONE - 您在其中创建新虚拟机的云可用区。
  • TPU_TYPE - 您创建的 TPU 虚拟机的类型。例如:v6e-1v6e-4
  • RESERVATION - 包含 TPU 的预订的名称。

使用预订创建 TPU 虚拟机:

gcloud alpha compute tpus tpu-vm create $TPU_NAME \
    --zone=$ZONE \
    --project $PROJECT_ID \
    --accelerator-type=$TPU_TYPE \
    --version=v2-alpha-tpuv6e \
    --provisioning-model=reservation-bound \
    --reservation=$RESERVATION

连接到 TPU 虚拟机。

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

使用 Qwen2-7B-instruct 运行 vLLM

  1. 设置 Hugging Face 令牌和模型名称变量。

      export HF_TOKEN="YOUR_HF_TOKEN"
      export MODEL_NAME="Qwen/Qwen2-7B-Instruct"
    
  2. 在 TPU 虚拟机内,以分离模式运行 vLLM Docker 容器,并启动 vLLM 服务器。此命令使用 10 GB 的共享内存大小。

    export DOCKER_URI="vllm/vllm-tpu:v0.18.0"
    export CONTAINER_NAME="${USER}-vllm"
    export MAX_MODEL_LEN=4096
    export TP=1 # number of chips
    
    sudo docker run -d --name "${CONTAINER_NAME}" \
        --privileged --net=host \
        -v /dev/shm:/dev/shm \
        --shm-size 10gb \
        -e "HF_HOME=/dev/shm" \
        -e "HF_TOKEN=${HF_TOKEN}" \
        -p 8000:8000 "${DOCKER_URI}" \
            vllm serve ${MODEL_NAME} \
                --seed 42 \
                --gpu-memory-utilization 0.98 \
                --max-num-batched-tokens 1024 \
                --max-num-seqs 128 \
                --tensor-parallel-size $TP \
                --max-model-len $MAX_MODEL_LEN
  3. 检查服务器日志以确认服务器正在运行。

    sudo docker logs -f "${CONTAINER_NAME}"

    vLLM 服务器运行时,您会看到类似于以下内容的输出。 输出显示后,按 CTRL+C 返回到终端。

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

发送推理请求

vLLM 服务器运行后,您可以向 API 发送请求。如需了解详情,请参阅 vLLM API 参考文档

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

    sudo docker exec "${CONTAINER_NAME}" \
      curl http://localhost:8000/v1/completions \
        -H "Content-Type: application/json" \
        -d '{
            "prompt": "The future of AI is",
            "max_tokens": 200,
            "temperature": 0
          }'

响应以 JSON 格式返回。

运行基准测试工作负载

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

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

    sudo docker exec "${CONTAINER_NAME}" pip install datasets
  2. 在容器内,运行 vllm bench serve 命令。

    sudo docker exec "${CONTAINER_NAME}" \
        vllm bench serve \
            --backend vllm \
            --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. 在终端中,输入 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 虚拟机。使用以下命令终止虚拟机, 将 TPU 释放回您的预订。

gcloud compute tpus tpu-vm delete $TPU_NAME --zone $ZONE --project $PROJECT_ID --quiet

后续步骤