在 GKE 上使用 vLLM TorchTPU 在 TPU7x 加速器上部署 Qwen3.5-397B

本教程介绍了如何使用单主机 TPU7x (Ironwood) 加速器和 vLLM TPU 服务框架（搭配 TorchTPU [torchtpu]），在 Google Kubernetes Engine (GKE) 上部署和应用混合专家模型 Qwen3.5-397B (Qwen/Qwen3.5-397B-A17B-FP8)。由于 Qwen3.5-397B 至少需要四个 TPU7x 芯片，因此您可以在启用数据并行和专家并行的 2x2x1 拓扑中的单个 tpu7x-standard-4t 节点上运行工作负载。

目标

  • 设置 GKE 集群和单主机 TPU7x 节点池。
  • 在 GKE 上部署单主机 vLLM TorchTPU 服务器。
  • 向 vLLM TorchTPU 服务器发送推理请求。
  • 运行服务基准工作负载。
  • 清理。

费用

在本文档中，您将使用Google Cloud的以下收费组件：

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

完成本文档中描述的任务后，您可以通过删除所创建的资源来避免继续计费。如需了解详情，请参阅清理。

准备工作

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

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

  • 在 Google Cloud 项目中启用 Kubernetes Engine API：

    启用 Kubernetes Engine API

  • 在 Google Cloud 项目中启用 Artifact Registry API：

    启用 Artifact Registry API

  • 安装并初始化 Google Cloud CLI，然后安装 kubectl。如果您使用 Cloud Shell，则这两个工具均已预安装。

  • 确保您的项目在 TPU 容量所在的区域中有足够的 SSD 配额和外部 IP 地址配额：

    1. 在 Google Cloud 控制台中，前往配额和系统限制页面。

      进入“配额和系统限制”

    2. 按永久性磁盘 SSD (GB) 进行过滤，并验证您是否有至少 1,500 GB 的可用配额用于模型缓存卷。

    3. 按正在使用的区域级外部 IPv4 地址进行过滤，并验证目标区域中是否有足够的 IP 地址配额。

如需获得完成本教程所需的权限，请让您的管理员为您授予项目的以下 IAM 角色：

如需详细了解如何授予角色，请参阅管理对项目、文件夹和组织的访问权限。

您也可以通过自定义角色或其他预定义角色来获取所需的权限。

设置环境变量

设置身份验证：

gcloud auth login
gcloud auth application-default login

设置环境变量：

export PROJECT_ID="YOUR_PROJECT_ID"
export REGION="YOUR_REGION"
export ZONE="YOUR_ZONE"
export CLUSTER_NAME="YOUR_CLUSTER_NAME"
export MACHINE_TYPE="tpu7x-standard-4t"
export NODE_POOL_NAME="np-0"
export RESERVATION_NAME="YOUR_RESERVATION_NAME"
export HF_TOKEN="YOUR_HF_TOKEN"
export NETWORK="YOUR_NETWORK_NAME"
export SUBNETWORK="YOUR_SUBNETWORK_NAME"
export VLLM_IMAGE="YOUR_VLLM_TPU_IMAGE"

gcloud config set project "${PROJECT_ID}"
gcloud config set billing/quota_project "${PROJECT_ID}"

替换以下内容：

  • YOUR_PROJECT_ID：您的 Google Cloud项目的 ID。
  • YOUR_REGION：您要在其中部署集群的区域，例如 us-central1。
  • YOUR_ZONE：您要部署 TPU7x 节点池的可用区，例如 us-central1-c。
  • YOUR_CLUSTER_NAME：GKE 集群的名称，例如 tpu7x-tutorial-cluster。
  • YOUR_RESERVATION_NAME：如果您使用的是预留，则为 TPU7x 预留的名称。
  • YOUR_HF_TOKEN：您的 Hugging Face 访问令牌。
  • YOUR_NETWORK_NAME：Google Cloud 网络的名称。如果您创建了自定义网络，请指定该网络。否则，请指定 default。
  • YOUR_SUBNETWORK_NAME：Google Cloud 子网的名称。如果您创建了自定义子网，请指定该子网。否则，请指定 default。
  • YOUR_VLLM_TPU_IMAGE：要使用的 vLLM TorchTPU 映像。如需了解公开映像，请参阅 vLLM TorchTPU 代码库。

创建 GKE 集群

如需在 Standard 模式下创建 GKE 集群并挂接单主机 TPU7x 节点池，请完成以下步骤：

  1. 创建区域级 GKE 集群：

    gcloud container clusters create "${CLUSTER_NAME}" \
    --project="${PROJECT_ID}" \
    --location="${REGION}" \
    --node-locations="${ZONE}" \
    --network="${NETWORK}" \
    --subnetwork="${SUBNETWORK}"

    输出类似于以下内容：

    NAME                    LOCATION       MASTER_VERSION      MASTER_IP      MACHINE_TYPE    NODE_VERSION        NUM_NODES  STATUS
tpu7x-tutorial-cluster  us-central1    1.35.7-gke.1027000  34.63.222.246  n2-standard-16  1.35.7-gke.1027000  3          RUNNING

  2. 使用预留创建一个具有一个 tpu7x-standard-4t 节点（2x2x1 拓扑中的四个 TPU7x 芯片）的单主机 TPU7x 节点池：

    gcloud container node-pools create "${NODE_POOL_NAME}" \
    --cluster="${CLUSTER_NAME}" \
    --machine-type="${MACHINE_TYPE}" \
    --location="${REGION}" \
    --node-locations="${ZONE}" \
    --project="${PROJECT_ID}" \
    --num-nodes=1 \
    --reservation-affinity=specific \
    --reservation="${RESERVATION_NAME}"

    您也可以不使用特定预留，而是通过传递以下标志，使用 Spot 容量或灵活启动型虚拟机创建节点池：

    • Spot 虚拟机：将 --reservation-affinity 和 --reservation 标志替换为 --spot，以使用 Spot 容量。使用 Spot 容量的实例在节点池创建后可能需要一些时间才能创建，并且可能会在工作负载运行期间被抢占。如果您的 Pod 需要一段时间才能创建，或者工作负载在运行过程中停止，则可能是这个原因。
    • 灵活启动型虚拟机：传递 --flex-start 标志，以使用灵活启动使用选项预配节点池。如需了解详情，请参阅使用 TPU 和灵活启动虚拟机运行小型批量工作负载。

    输出类似于以下内容：

    NAME  MACHINE_TYPE       DISK_SIZE_GB  NODE_VERSION
np-0  tpu7x-standard-4t  100           1.35.7-gke.1027000

  3. 配置 kubectl 凭据以与您的集群通信：

    gcloud container clusters get-credentials "${CLUSTER_NAME}" \
    --location="${REGION}" \
    --project="${PROJECT_ID}"

配置凭据

创建一个名为 hf-secret 的 Kubernetes Secret，用于存储您的 Hugging Face 访问令牌：

kubectl create secret generic hf-secret \
    --from-literal=hf_api_token="${HF_TOKEN}" \
    --dry-run=client -o yaml | kubectl apply -f -

部署 vLLM TorchTPU 服务器

工作负载清单 (qwen3_5_server.yaml) 定义了以下资源：

  • 一个 StorageClass (hyperdisk-balanced-qwen3-5) 和一个 1,500 GiB PersistentVolumeClaim (hd-claim-qwen3-5)，由 Google Cloud Hyperdisk Balanced 存储空间提供支持，用于存储下载的模型权重、XLA 编译缓存和 PyTorch 编译缓存。
  • 在具有 2x2x1 拓扑 (google.com/tpu: '4') 的单个 TPU7x 节点上调度的 Deployment (vllm-qwen3-5)。容器会针对 Qwen/Qwen3.5-397B-A17B-FP8 运行 vllm serve，并使用 --data-parallel-size=8、--tensor-parallel-size=1、--enable-expert-parallel、--kv-cache-dtype=fp8、--max-model-len=9236 和 --max-num-seqs=64。
  • 一个公开端口 8000 的 ClusterIP Service (vllm-qwen3-5-service)。

如需部署和监控工作负载，请完成以下步骤：

  1. 提取服务器清单：

    curl -O https://raw.githubusercontent.com/GoogleCloudSamples/AIHypercomputerSamples/refs/heads/main/tpu/inference/qwen3-5-397b/qwen3_5_server.yaml

    或者，您也可以选择创建一个名为 qwen3_5_server.yaml 的文件，其中包含以下内容：

    apiVersion: storage.k8s.io/v1
kind: StorageClass
metadata:
  name: hyperdisk-balanced-qwen3-5
provisioner: pd.csi.storage.gke.io
parameters:
  type: hyperdisk-balanced
reclaimPolicy: Delete
volumeBindingMode: WaitForFirstConsumer
allowVolumeExpansion: true
---
apiVersion: v1
kind: PersistentVolumeClaim
metadata:
  name: hd-claim-qwen3-5
  namespace: default
spec:
  storageClassName: hyperdisk-balanced-qwen3-5
  accessModes:
    - ReadWriteOnce
  resources:
    requests:
      storage: 1500Gi
---
apiVersion: apps/v1
kind: Deployment
metadata:
  name: vllm-qwen3-5
  namespace: default
spec:
  replicas: 1
  selector:
    matchLabels:
      app: vllm-qwen3-5
  template:
    metadata:
      labels:
        app: vllm-qwen3-5
    spec:
      nodeSelector:
        cloud.google.com/gke-tpu-accelerator: tpu7x
        cloud.google.com/gke-tpu-topology: 2x2x1
      tolerations:
      - key: "cloud.google.com/gke-spot"
        operator: "Equal"
        value: "true"
        effect: "NoSchedule"
      containers:
      - name: vllm-tpu
        image: $VLLM_IMAGE
        command: ["vllm", "serve"]
        args:
        - --host=0.0.0.0
        - --port=8000
        - --model=Qwen/Qwen3.5-397B-A17B-FP8
        - --tokenizer=Qwen/Qwen3.5-397B-A17B-FP8
        - --max-model-len=9236
        - --max-num-batched-tokens=1024
        - --max-num-seqs=64
        - --no-enable-prefix-caching
        - --gpu-memory-utilization=0.9
        - --data-parallel-size=8
        - --tensor-parallel-size=1
        - --async-scheduling
        - --language-model-only
        - --enable-auto-tool-choice
        - --tool-call-parser=qwen3_coder
        - --reasoning-parser=qwen3
        - '--limit-mm-per-prompt={"image": 0, "video": 0}'
        - --kv-cache-dtype=fp8
        - --enable-expert-parallel
        - --download-dir=/data
        - --block-size=128
        env:
        - name: MODEL_IMPL_TYPE
          value: "vllm"
        - name: DP_SCHED_ENABLED
          value: "1"
        - name: ONEHOT_MOE_PERMUTE_THRESHOLD
          value: "32768"
        - name: VLLM_ENGINE_READY_TIMEOUT_S
          value: "9600"
        - name: VLLM_ENGINE_ITERATION_TIMEOUT_S
          value: "1800"
        - name: USE_MOE_SPARSE_CORE
          value: "1"
        - name: TPU_MOE_SKIP_PADDED_TOKENS
          value: "1"
        - name: VLLM_CACHE_ROOT
          value: "/data/vllm_cache"
        - name: VLLM_TORCH_COMPILE_CACHE_DIR
          value: "/data/torch_compile_cache"
        - name: XLA_CACHE_DIR
          value: "/data/xla_cache"
        - name: HF_HOME
          value: /data
        - name: HF_TOKEN
          valueFrom:
            secretKeyRef:
              name: hf-secret
              key: hf_api_token
        ports:
        - containerPort: 8000
        resources:
          limits:
            google.com/tpu: '4'
          requests:
            google.com/tpu: '4'
        readinessProbe:
          tcpSocket:
            port: 8000
          initialDelaySeconds: 15
          periodSeconds: 10
        volumeMounts:
        - mountPath: "/data"
          name: data-volume
        - mountPath: /dev/shm
          name: dshm
      volumes:
      - emptyDir:
          medium: Memory
        name: dshm
      - name: data-volume
        persistentVolumeClaim:
          claimName: hd-claim-qwen3-5
---
apiVersion: v1
kind: Service
metadata:
  name: vllm-qwen3-5-service
  namespace: default
spec:
  selector:
    app: vllm-qwen3-5
  type: ClusterIP
  ports:
  - name: http
    protocol: TCP
    port: 8000
    targetPort: 8000

  2. 将 vLLM TorchTPU 服务器清单应用到您的集群：

    envsubst '$VLLM_IMAGE' < "qwen3_5_server.yaml" | kubectl apply -f -

  3. 监控服务器部署日志，直到应用启动完成：

    kubectl wait --for=jsonpath='{.status.phase}'=Running pod -l app=vllm-qwen3-5 --timeout=900s
kubectl logs deployment/vllm-qwen3-5 -c vllm-tpu -f

    在 Pod 下载模型权重并编译模型时，初始服务启动可能需要大约 20 到 30 分钟。

    如果您使用的是 Spot 或灵活启动等抢占式容量，则在创建节点池后，TPU 切片节点可能需要更长时间才能创建，并且在工作负载运行期间可能会被抢占。如果 Pod 的启动时间超出预期，或者工作负载中途停止，则可能是这个原因。如需了解详情，请参阅抢占 Spot 虚拟机。如需检查 TPU 切片节点是否已创建，请参阅检查节点的状态和条件。

    请等待日志报告 Application startup complete，然后再发送流量：

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

  4. 打开第二个终端会话。

  5. 在第二个终端会话中，将本地端口 8000 转发到 vllm-qwen3-5-service：

    kubectl port-forward service/vllm-qwen3-5-service 8000:8000

    让端口转发在此终端会话中保持运行。

发送推理请求

返回到第一个终端会话，然后向 /v1/chat/completions 端点发送对话补全请求：

curl http://localhost:8000/v1/chat/completions \
    -H "Content-Type: application/json" \
    -d '{
        "model": "Qwen/Qwen3.5-397B-A17B-FP8",
        "messages": [
            {"role": "user", "content": "Hello! Tell me a joke."}
        ],
        "max_tokens": 512,
        "temperature": 0.7
    }' | jq .

运行基准工作负载

如需使用 InferenceX 基准套件，通过 8,192 个词元的输入和 1,024 个词元的输出 (8k/1k) 工作负载对已部署的 Qwen3.5 397B 模型进行基准测试，请完成以下步骤：

  1. 提取基准清单：

    curl -O https://raw.githubusercontent.com/GoogleCloudSamples/AIHypercomputerSamples/refs/heads/main/tpu/inference/qwen3-5-397b/qwen3_5-benchmark-8k1k.yaml

    或者，创建一个名为 qwen3_5-benchmark-8k1k.yaml 的文件，其中包含以下内容：

    apiVersion: v1
kind: Pod
metadata:
  name: qwen3-5-bench-8k1k
  namespace: default
spec:
  restartPolicy: OnFailure
  terminationGracePeriodSeconds: 60
  containers:
  - name: vllm-bench
    image: $VLLM_IMAGE
    command: ["/bin/bash", "-c"]
    args:
    - |
      while ! curl http://vllm-qwen3-5-service:8000/ping; do sleep 30 && echo 'Waiting for server...'; done
      apt-get update && apt-get install -y git && \
      git clone https://github.com/SemiAnalysisAI/InferenceX.git /ubench/inferencex && \
      cd /ubench/inferencex && \
      git checkout 89ce6098ef2bc4576a735c43f39c7d972b091cfc && \
      python3 /ubench/inferencex/utils/bench_serving/benchmark_serving.py \
        --backend=vllm \
        --percentile-metrics='ttft,tpot,itl,e2el' \
        --host=vllm-qwen3-5-service \
        --port=8000 \
        --model=Qwen/Qwen3.5-397B-A17B-FP8 \
        --tokenizer=Qwen/Qwen3.5-397B-A17B-FP8 \
        --dataset-name=random \
        --random-input-len=8192 \
        --random-output-len=1024 \
        --random-range-ratio=1.0 \
        --num-prompts=2560 \
        --num-warmups=512 \
        --max-concurrency=256 \
        --ignore-eos
    env:
    - name: HUGGING_FACE_HUB_TOKEN
      valueFrom:
        secretKeyRef:
          key: hf_api_token
          name: hf-secret

  2. 应用基准 Pod 清单：

    envsubst '$VLLM_IMAGE' < "qwen3_5-benchmark-8k1k.yaml" | kubectl apply -f -

  3. 监控基准测试日志：

    kubectl logs -f qwen3-5-bench-8k1k

    基准测试大约需要 25 到 30 分钟才能完成。

清理

为避免因本教程中使用的资源导致您的 Google Cloud 账号产生额外的持续费用，请删除 Kubernetes 工作负载、TPU 节点池和 GKE 集群：

  1. 删除 vLLM TorchTPU 服务器 Deployment、Service、PersistentVolumeClaim 和基准 Pod：

    kubectl delete -f qwen3_5_server.yaml --ignore-not-found || true
kubectl delete secret hf-secret --ignore-not-found || true

  2. 删除 TPU7x 节点池和 GKE 集群：

    gcloud container node-pools delete "${NODE_POOL_NAME}" \
    --cluster="${CLUSTER_NAME}" \
    --location="${REGION}" \
    --project="${PROJECT_ID}" \
    --quiet || true

gcloud container clusters delete "${CLUSTER_NAME}" \
    --location="${REGION}" \
    --project="${PROJECT_ID}" \
    --quiet || true

后续步骤