在 GKE 上使用多主機 vLLM TorchTPU，透過 Trillium TPU 提供 Gemma 4 26B

本教學課程說明如何使用多主機 Cloud TPU v6e (Trillium) 加速器和 vLLM TorchTPU，在 Google Kubernetes Engine (GKE) 部署及提供多模態專家混合模型 Gemma 4 26B (google/gemma-4-26B-A4B-it)。由於工作負載會跨越 2x4 拓撲中的多個 TPU 主機，因此您可以使用 Kubernetes LeaderWorkerSet API 和 Ray，在八個 TPU 晶片之間協調分散式張量平行推論。

目標

  • 設定 GKE 叢集和多主機 TPU v6e 節點集區。
  • 使用 LeaderWorkerSet 部署多主機 vLLM 伺服器。
  • 傳送文字和多模態推論要求。
  • 執行推論基準測試工作負載。
  • 清除所用資源。

費用

在本文件中，您會使用下列Google Cloud計費元件：

您可以使用 Pricing Calculator 根據預測用量估算費用。

完成本文所述工作後，您可以刪除建立的資源，避免繼續計費，詳情請參閱「清除所用資源」一節。

事前準備

您需要 Hugging Face 存取權杖，才能使用本教學課程。您可以前往 Hugging Face 註冊免費帳戶。取得帳戶後，請產生存取權杖：

  1. 在「Welcome to Hugging Face」頁面，點選帳戶顯示圖片，然後選取「存取權杖」。
  2. 在「存取權杖」頁面，按一下「建立新權杖」。
  3. 選取「Read」權杖類型，並輸入權杖名稱。
  4. 系統會顯示存取權杖。將權杖儲存在安全的地方。
  • 安裝並初始化 Google Cloud CLI，然後安裝 kubectl。如果使用 Cloud Shell，這兩項工具都已預先安裝。

如要取得完成本教學課程所需的權限，請要求管理員在專案中授予您下列 IAM 角色：

如要進一步瞭解如何授予角色，請參閱「管理專案、資料夾和組織的存取權」。

您或許也能透過自訂角色或其他預先定義的角色，取得必要權限。

設定環境變數

設定環境變數：

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 拓撲排列 (共八個 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 運算資源或彈性啟動型 VM 建立節點集區：

    • Spot VM：將 --reservation-affinity 和 --reservation 旗標替換為 --spot，即可使用 Spot 容量。建立節點集區後，使用 Spot 容量的執行個體可能需要一段時間才能建立，且在工作負載執行期間可能會遭到先占。如果 Pod 需要一段時間才能建立，或工作負載在執行期間停止，可能就是這個原因。
    • 彈性啟動型 VM：傳遞 --flex-start，使用彈性啟動用量方案佈建節點集區。詳情請參閱「使用 TPU 和彈性啟動型 VM 執行小型批次工作負載」。

  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 (一個 leader Pod 和一個 worker Pod)，排定在 2x4 TPU 切片的兩個 ct6e-standard-4t 節點上。每個 Pod 都會執行 InitContainer，先將 Hugging Face 的 google/gemma-4-26B-A4B-it 模型權重預先擷取至記憶體內磁碟區，再使用 --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 會預先擷取模型權重、在所有八個 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

清除所用資源

如要避免系統向您的 Google Cloud 帳戶持續收取本教學課程所用資源的額外費用，請刪除 Kubernetes 工作負載、TPU 節點集區和 GKE 叢集：

  1. 刪除 vLLM 伺服器部署作業和 Hugging Face 密鑰：

    envsubst '$RESERVATION_NAME $ZONE $VLLM_IMAGE' < gemma4-26b-multihost-v6e-2x4.yaml | kubectl delete -f - --ignore-not-found || true
kubectl delete secret hf-secret --ignore-not-found || true

  2. 刪除 TPU v6e 節點集區和 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

  3. 返回第二個終端機工作階段，停止通訊埠轉送。

後續步驟