通过 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. 在 Welcome to Hugging Face 页面上，点击您的账号头像，然后选择 Access tokens。
  2. 在访问令牌页面上，点击创建新令牌。
  3. 选择读取令牌类型，然后输入令牌的名称。
  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

    例如，当状态为 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. 设置您的 Hugging Face 令牌。

      export HF_TOKEN="YOUR_HF_TOKEN"

  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 Qwen/Qwen2-7B-Instruct \
            --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 -ti "${CONTAINER_NAME}" \
    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 库。

    sudo docker exec -it "${CONTAINER_NAME}" pip install datasets

  2. 在容器内，运行 vllm bench serve 命令。

    sudo docker exec -it "${CONTAINER_NAME}" \
    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. 在终端中，输入 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

后续步骤