在 GKE 上使用多主机 vLLM TorchTPU 在 Trillium TPU 上提供 Gemma 4 26B

设置环境变量

设置环境变量：

export PROJECT_ID="YOUR_PROJECT_ID"
export RESERVATION_NAME="YOUR_RESERVATION_NAME"
export REGION="YOUR_REGION"
export ZONE="YOUR_ZONE"
export CLUSTER_NAME="YOUR_CLUSTER_NAME"
export HF_TOKEN="YOUR_HF_TOKEN"
export NODE_POOL_NAME="tpuv6e-pool"
export VLLM_IMAGE="vllm/vllm-torchtpu:v0.29.0"

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

替换以下内容：

• YOUR_PROJECT_ID：您的 Google Cloud项目的 ID。
• YOUR_RESERVATION_NAME：预留的名称。
• YOUR_REGION：您要在其中部署 GKE 集群的区域。
• YOUR_ZONE：要在其中部署 TPU v6e 节点池的可用区。
• YOUR_CLUSTER_NAME：GKE 集群的名称。
• YOUR_HF_TOKEN：您的 Hugging Face 访问令牌。
• YOUR_VLLM_TPU_IMAGE：要使用的 vLLM TPU 映像。对于公共映像，请参阅 vLLM TorchTPU 代码库。

创建 GKE 集群

如需创建 GKE 集群并附加多主机 TPU v6e 节点池，请完成以下步骤：

1. 创建一个区域级 GKE 集群，该集群使用 e2-standard-4 节点作为默认节点池，并使用一个 TPU v6e 节点池，该节点池包含两个 ct6e-standard-4t 节点，这些节点采用 2x4 拓扑（总共 8 个 TPU 芯片）。默认 GKE 机器类型 (e2-medium) 无法为 LeaderWorkerSet 控制器提供足够的 CPU 资源，因此本教程指定了 e2-standard-4：

   gcloud container clusters create "${CLUSTER_NAME}" \
       --project="${PROJECT_ID}" \
       --location="${REGION}" \
       --machine-type=e2-standard-4 \
   
   gcloud container node-pools create "${NODE_POOL_NAME}" \
       --project="${PROJECT_ID}" \
       --location="${REGION}" \
       --cluster="${CLUSTER_NAME}" \
       --node-locations="${ZONE}" \
       --machine-type=ct6e-standard-4t \
       --tpu-topology=2x4 \
       --num-nodes=2 \
       --disk-size=500 \
       --reservation-affinity=specific \
       --reservation="${RESERVATION_NAME}"

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

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

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

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

配置凭证并安装 LeaderWorkerSet

如需为多主机部署准备集群，请完成以下步骤：

1. 导出 Hugging Face 令牌并创建名为 hf-secret 的 Kubernetes Secret：

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

2. 在集群上安装 LeaderWorkerSet CustomResourceDefinition 和控制器：

   kubectl apply --server-side -k "github.com/kubernetes-sigs/lws/config/default?ref=main"
   echo "Waiting for LWS controller to be ready..."
   kubectl wait --for=condition=Available --timeout=300s deployment/lws-controller-manager -n lws-system

部署多主机 vLLM 服务器

工作负载清单 (gemma4-26b-multihost-v6e-2x4.yaml) 定义了两个资源：

• 一个 LeaderWorkerSet 资源 (vllm-tpu-gemma4-26b)，其中包含 size: 2（一个主 Pod 和一个工作器 Pod），调度到 2x4 TPU 切片中的两个 ct6e-standard-4t 节点上。每个 Pod 都会运行一个 InitContainer，该容器会先将 google/gemma-4-26B-A4B-it 模型权重从 Hugging Face 预提取到内存卷中，然后再使用 --tensor-parallel-size=8、--kv-cache-dtype=fp8、--max-model-len=5248 和 --max-num-seqs=128 标志启动 Ray 和 vLLM OpenAI API 服务器。
• 一个 ClusterIP Service (vllm-service)，用于在 Leader Pod 上公开端口 8000。

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

1. 从 GitHub 提取 LeaderWorkerSet 清单：

   curl -LO https://raw.githubusercontent.com/GoogleCloudSamples/AIHypercomputerSamples/main/tpu/inference/gemma-4-26b/gemma4-26b-multihost-v6e-2x4.yaml

   或者，创建一个名为 gemma4-26b-multihost-v6e-2x4.yaml 的文件，其中包含以下内容：

   apiVersion: leaderworkerset.x-k8s.io/v1
   kind: LeaderWorkerSet
   metadata:
     name: vllm-tpu-gemma4-26b
   spec:
     replicas: 1
     leaderWorkerTemplate:
       # One pod per TPU host: a 2x4 v6e slice has 2 ct6e-standard-4t hosts.
       size: 2
       restartPolicy: RecreateGroupOnPodRestart
       leaderTemplate:
         metadata:
           labels:
             role: leader
         spec:
           # Download the model into the pod-local cache before vLLM starts, so the
           # serving ranks can run offline and never hit Hugging Face during init.
           initContainers:
           - name: prefetch-model
             image: $VLLM_IMAGE
             command: ["hf", "download", "google/gemma-4-26B-A4B-it"]
             env:
             - name: HF_TOKEN
               valueFrom:
                 secretKeyRef:
                   name: hf-secret
                   key: hf_api_token
             - name: HF_HOME
               value: "/data/huggingface_cache"
             volumeMounts:
             - name: model-cache
               mountPath: /data
           containers:
           - name: vllm-leader
             image: $VLLM_IMAGE
             command:
               - sh
               - -c
               - |
                 set -e
                 mkdir -p "${TMPDIR}"
                 ray start --head --port=6379
                 echo "Waiting for all 8 TPU chips to join the Ray cluster..."
                 until ray status 2>/dev/null | grep -E -q "8(\.0)? TPU"; do
                   sleep 5
                 done
                 vllm serve google/gemma-4-26B-A4B-it \
                   --tensor-parallel-size 8 \
                   --distributed-executor-backend ray \
                   --host 0.0.0.0 \
                   --port 8000 \
                   --max-model-len 5248 \
                   --max-num-seqs 128 \
                   --max-num-batched-tokens 4096 \
                   --kv-cache-dtype fp8 \
                   --gpu-memory-utilization 0.85 \
                   --disable-chunked-mm-input
             resources:
               requests:
                 cpu: "128"
                 memory: "384Gi"
                 google.com/tpu: 4
               limits:
                 cpu: "170"
                 memory: "650Gi"
                 google.com/tpu: 4
             ports:
             - containerPort: 8000
             - containerPort: 6379
             env:
             - name: TPU_MULTIHOST_BACKEND
               value: "ray"
             - name: VLLM_ENGINE_READY_TIMEOUT_S
               value: "3600"
             # Same cache the init container populated; offline mode guarantees no
             # Hugging Face requests from the serving ranks.
             - name: HF_HOME
               value: "/data/huggingface_cache"
             - name: HF_HUB_OFFLINE
               value: "1"
             - name: TMPDIR
               value: "/data/tmp"
             volumeMounts:
             - mountPath: /dev/shm
               name: dshm
             - mountPath: /data
               name: model-cache
             # Model loading and compilation can take a long time; allow up to 60 minutes.
             startupProbe:
               httpGet:
                 path: /health
                 port: 8000
               periodSeconds: 10
               failureThreshold: 360
             livenessProbe:
               httpGet:
                 path: /health
                 port: 8000
               periodSeconds: 10
             readinessProbe:
               httpGet:
                 path: /health
                 port: 8000
               periodSeconds: 5
           nodeSelector:
             topology.kubernetes.io/zone: "$ZONE"
             cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
             cloud.google.com/gke-tpu-topology: 2x4
             cloud.google.com/reservation-name: $RESERVATION_NAME
           volumes:
           - name: dshm
             emptyDir:
               medium: Memory
               sizeLimit: 32Gi
           - name: model-cache
             emptyDir:
               medium: Memory
               sizeLimit: 350Gi
       workerTemplate:
         spec:
           # Each pod has its own cache, so the worker downloads the model too.
           initContainers:
           - name: prefetch-model
             image: $VLLM_IMAGE
             command: ["hf", "download", "google/gemma-4-26B-A4B-it"]
             env:
             - name: HF_TOKEN
               valueFrom:
                 secretKeyRef:
                   name: hf-secret
                   key: hf_api_token
             - name: HF_HOME
               value: "/data/huggingface_cache"
             volumeMounts:
             - name: model-cache
               mountPath: /data
           containers:
           - name: vllm-worker
             image: $VLLM_IMAGE
             command:
               - sh
               - -c
               - |
                 set -e
                 mkdir -p "${TMPDIR}"
                 echo "Waiting for the Ray head at ${LWS_LEADER_ADDRESS}:6379..."
                 until ray status --address="${LWS_LEADER_ADDRESS}:6379" 2>/dev/null | grep -q "node_"; do
                   sleep 3
                 done
                 ray start --address="${LWS_LEADER_ADDRESS}:6379" --block
             resources:
               requests:
                 cpu: "128"
                 memory: "384Gi"
                 google.com/tpu: 4
               limits:
                 cpu: "170"
                 memory: "650Gi"
                 google.com/tpu: 4
             env:
             - name: TPU_MULTIHOST_BACKEND
               value: "ray"
             - name: VLLM_ENGINE_READY_TIMEOUT_S
               value: "3600"
             - name: HF_HOME
               value: "/data/huggingface_cache"
             - name: HF_HUB_OFFLINE
               value: "1"
             - name: TMPDIR
               value: "/data/tmp"
             volumeMounts:
             - mountPath: /dev/shm
               name: dshm
             - mountPath: /data
               name: model-cache
           nodeSelector:
             topology.kubernetes.io/zone: "$ZONE"
             cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
             cloud.google.com/gke-tpu-topology: 2x4
             cloud.google.com/reservation-name: $RESERVATION_NAME
           volumes:
           - name: dshm
             emptyDir:
               medium: Memory
               sizeLimit: 32Gi
           - name: model-cache
             emptyDir:
               medium: Memory
               sizeLimit: 350Gi
   ---
   apiVersion: v1
   kind: Service
   metadata:
     name: vllm-service
   spec:
     selector:
       leaderworkerset.sigs.k8s.io/name: vllm-tpu-gemma4-26b
       role: leader
     type: ClusterIP
     ports:
       - protocol: TCP
         port: 8000
         targetPort: 8000

2. 应用 LeaderWorkerSet 清单：

   MANIFEST="gemma4-26b-multihost-v6e-2x4.yaml"
   echo "Applying LeaderWorkerSet (retrying if webhook is not fully ready)..."
   MAX_ATTEMPTS=12
   attempt=1
   until envsubst '$RESERVATION_NAME $ZONE $VLLM_IMAGE' < "${MANIFEST}" | kubectl apply -f -; do
     if [ "${attempt}" -ge "${MAX_ATTEMPTS}" ]; then
       echo "Error: failed to apply ${MANIFEST} after ${MAX_ATTEMPTS} attempts." >&2
       exit 1
     fi
     echo "Webhook not ready yet, retrying in 5 seconds (attempt ${attempt}/${MAX_ATTEMPTS})..."
     attempt=$((attempt + 1))
     sleep 5
   done

3. 监控主 Pod 状态和日志，直到应用启动完成：

   LIMIT=60
   count=0
   until kubectl get pod -l role=leader -o name | grep -q "vllm-tpu"; do
     if [ "${count}" -ge "${LIMIT}" ]; then
       echo "Error: timeout waiting for leader pod to be created." >&2
       exit 1
     fi
     echo "Waiting for leader pod to be created..."
     sleep 5
     count=$((count + 1))
   done
   
   # Wait up to 20 minutes for the pod to start running (TPU nodes take a while to boot)
   echo "Waiting for pod to be scheduled and running..."
   if ! kubectl wait --for=jsonpath='{.status.phase}'=Running --timeout=1200s pod -l role=leader; then
     echo "Pod failed to start running within 20 minutes! Here are the events:" >&2
     kubectl describe pod -l role=leader
     exit 1
   fi
   
   echo "Pod is running. Streaming logs while waiting for it to become Ready (this includes model download)..."
   # Stream logs in the background so the user can see what's happening
   kubectl logs -f -l role=leader &
   LOG_PID=$!
   trap 'kill "${LOG_PID}" 2>/dev/null || true' EXIT
   
   if ! kubectl wait --for=condition=Ready --timeout=3600s pod -l role=leader; then
     echo "Pod failed to become Ready within 60 minutes!" >&2
     kubectl describe pod -l role=leader
     exit 1
   fi
   
   echo "Pod is Ready!"

   初始服务启动大约需要 9 分钟，在此期间，Pod 会预提取模型权重，在所有 8 个 TPU 芯片上形成双节点 Ray 集群，并编译模型。请等到日志报告 Application startup complete 后再发送流量。

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

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

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

发送推理请求

返回到第一个终端会话，并使用文本和多模态对话补全请求测试已部署的模型：

1. 向 /v1/chat/completions 端点发送文本对话补全请求：

   curl http://localhost:8000/v1/chat/completions \
       -H "Content-Type: application/json" \
       -d '{
           "model": "google/gemma-4-26B-A4B-it",
           "messages": [{"role": "user", "content": "Explain Mixture of Experts in 2 sentences."}],
           "max_tokens": 128
       }'

2. 发送包含图片网址和文本提示词的多模态请求：

   curl http://localhost:8000/v1/chat/completions \
       -H "Content-Type: application/json" \
       -d '{
           "model": "google/gemma-4-26B-A4B-it",
           "messages": [
               {
                   "role": "user",
                   "content": [
                       {
                           "type": "image_url",
                           "image_url": {
                               "url": "https://images.unsplash.com/photo-1501594907352-04cda38ebc29?auto=format&fit=crop&w=800&q=80"
                           }
                       },
                       {
                           "type": "text",
                           "text": "Describe the image above"
                       }
                   ]
               }
           ],
           "max_tokens": 300,
           "temperature": 0.0,
           "top_p": 1.0
       }'

运行基准工作负载

如需对已部署的模型进行基准测试，请运行以下命令：

kubectl exec vllm-tpu-gemma4-26b-0 -- vllm bench serve \
  --backend openai-chat \
  --endpoint /v1/chat/completions \
  --model google/gemma-4-26B-A4B-it \
  --host localhost \
  --port 8000 \
  --dataset-name random \
  --num-prompts 250 \
  --random-input-len 1024 \
  --random-output-len 1024 \
  --max-concurrency 128