Gemma 4 26B in der GKE mit Multi-Host-vLLM-TorchTPU auf Trillium-TPUs bereitstellen

In dieser Anleitung erfahren Sie, wie Sie das multimodale Mixture of Experts-Modell Gemma 4 26B (google/gemma-4-26B-A4B-it) in Google Kubernetes Engine (GKE) mit Cloud TPU v6e (Trillium)-Beschleunigern mit mehreren Hosts und vLLM TorchTPU bereitstellen und bereitstellen. Da die Arbeitslast mehrere TPU-Hosts in einer 2x4-Topologie umfasst, verwenden Sie die Kubernetes LeaderWorkerSet API und Ray, um die verteilte tensorparallele Inferenz über acht TPU-Chips hinweg zu koordinieren.

Ziele

  • Richten Sie Ihren GKE-Cluster und den TPU v6e-Knotenpool mit mehreren Hosts ein.
  • Stellen Sie den vLLM-Server mit mehreren Hosts mit LeaderWorkerSet bereit.
  • Text- und multimodale Inferenzanfragen senden
  • Führen Sie eine Benchmark-Arbeitslast für die Bereitstellung aus.
  • bereinigen.

Kosten

In diesem Dokument verwenden Sie die folgenden kostenpflichtigen Komponenten vonGoogle Cloud:

Mit dem Preisrechner können Sie eine Kostenschätzung Ihrer voraussichtlichen Nutzung vornehmen.

Nach Abschluss der in diesem Dokument beschriebenen Aufgaben können Sie weitere Kosten vermeiden, indem Sie die erstellten Ressourcen löschen. Weitere Informationen finden Sie unter Bereinigen.

Hinweis

Für diese Anleitung benötigen Sie ein Hugging Face-Zugriffstoken. Sie können sich bei Hugging Face für ein kostenloses Konto registrieren. Nachdem Sie ein Konto erstellt haben, generieren Sie ein Zugriffstoken:

  1. Klicken Sie auf der Seite Willkommen bei Hugging Face auf Ihren Konto-Avatar und wählen Sie Zugriffstokens aus.
  2. Klicken Sie auf der Seite Zugriffstokens auf Neues Token erstellen.
  3. Wählen Sie den Tokentyp Lesen aus und geben Sie einen Namen für das Token ein.
  4. Ihr Zugriffstoken wird angezeigt. Speichern Sie das Token an einem sicheren Ort.

Bitten Sie Ihren Administrator, Ihnen die folgenden IAM-Rollen für Ihr Projekt zuzuweisen, um die Berechtigungen zu erhalten, die Sie zum Ausführen dieser Anleitung benötigen:

Weitere Informationen zum Zuweisen von Rollen finden Sie unter Zugriff auf Projekte, Ordner und Organisationen verwalten.

Sie können die erforderlichen Berechtigungen auch über benutzerdefinierte Rollen oder andere vordefinierte Rollen erhalten.

Umgebungsvariablen einrichten

Richten Sie die Umgebungsvariablen ein:

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}"

Ersetzen Sie Folgendes:

  • YOUR_PROJECT_ID: die ID Ihres Google Cloud-Projekts.
  • YOUR_RESERVATION_NAME: der Name Ihrer Reservierung.
  • YOUR_REGION: die Region, in der Sie Ihren GKE-Cluster bereitstellen möchten.
  • YOUR_ZONE: die Zone, in der Sie Ihren TPU v6e-Knotenpool bereitstellen möchten.
  • YOUR_CLUSTER_NAME: Name Ihres GKE-Clusters.
  • YOUR_HF_TOKEN: Ihr Hugging Face-Zugriffstoken.
  • YOUR_VLLM_TPU_IMAGE: das zu verwendende vLLM-TPU-Image. Informationen zu öffentlichen Images finden Sie im vLLM TorchTPU-Repository.

GKE-Cluster erstellen

Führen Sie die folgenden Schritte aus, um den GKE-Cluster zu erstellen und einen TPU v6e-Knotenpool mit mehreren Hosts anzuhängen:

  1. Erstellen Sie einen regionalen GKE-Cluster, der e2-standard-4-Knoten für den Standardknotenpool verwendet, und einen TPU v6e-Knotenpool mit zwei ct6e-standard-4t-Knoten, die in einer 2x4-Topologie angeordnet sind (insgesamt acht TPU-Chips). Der GKE-Standardmaschinentyp (e2-medium) bietet nicht genügend CPU-Ressourcen für den LeaderWorkerSet-Controller. Daher wird in dieser Anleitung e2-standard-4 angegeben:

    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}"

    Anstatt eine bestimmte Reservierung zu verwenden, können Sie den Knotenpool auch mit Spot-Kapazität oder Flex-Start-VMs erstellen. Übergeben Sie dazu die folgenden Flags:

    • Spot-VMs: Ersetzen Sie die Flags --reservation-affinity und --reservation durch --spot, um Spot-Kapazität zu verwenden. Es kann einige Zeit dauern, bis Instanzen, die Spot-Kapazität verwenden, nach dem Erstellen des Knotenpools erstellt werden. Außerdem können sie während der Ausführung von Arbeitslasten vorzeitig beendet werden. Wenn die Erstellung Ihrer Pods lange dauert oder Arbeitslasten während der Ausführung beendet werden, kann dies die Ursache sein.
    • Flex-Start-VMs: Übergeben Sie --flex-start, um den Knotenpool mit der Verbrauchsoption Flex-Start bereitzustellen. Weitere Informationen finden Sie unter Kleine Batcharbeitslast mit TPUs und Flex-Start-VMs ausführen.

  2. Konfigurieren Sie kubectl-Anmeldedaten für die Kommunikation mit Ihrem Cluster:

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

Anmeldedaten konfigurieren und LeaderWorkerSet installieren

Führen Sie die folgenden Schritte aus, um Ihren Cluster für die Bereitstellung auf mehreren Hosts vorzubereiten:

  1. Exportieren Sie Ihr Hugging Face-Token und erstellen Sie ein Kubernetes-Secret mit dem Namen hf-secret:

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

  2. Installieren Sie die LeaderWorkerSet-CustomResourceDefinition und den Controller in Ihrem Cluster:

    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-Server mit mehreren Hosts bereitstellen

Das Arbeitslastmanifest (gemma4-26b-multihost-v6e-2x4.yaml) definiert zwei Ressourcen:

  • Eine LeaderWorkerSet-Ressource (vllm-tpu-gemma4-26b) mit size: 2 (ein Leader-Pod und ein Worker-Pod), die auf den beiden ct6e-standard-4t-Knoten im 2x4-TPU-Slice geplant ist. In jedem Pod wird ein InitContainer ausgeführt, der die google/gemma-4-26B-A4B-it-Modellgewichtungen von Hugging Face in ein In-Memory-Volume vorabruft, bevor Ray und der vLLM OpenAI API-Server mit den Flags --tensor-parallel-size=8, --kv-cache-dtype=fp8, --max-model-len=5248 und --max-num-seqs=128 gestartet werden.
  • Ein ClusterIP-Dienst (vllm-service), der Port 8000 auf dem Leader-Pod verfügbar macht.

Führen Sie die folgenden Schritte aus, um die Arbeitslast bereitzustellen und zu überwachen:

  1. Rufen Sie das LeaderWorkerSet-Manifest von GitHub ab:

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

    Alternativ können Sie eine Datei mit dem Namen gemma4-26b-multihost-v6e-2x4.yaml und folgendem Inhalt erstellen:

    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. Wenden Sie das LeaderWorkerSet-Manifest an:

    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. Überwachen Sie den Status und die Logs des Leader-Pods, bis die Anwendung gestartet wurde:

    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!"

    Der erste Dienststart dauert etwa 9 Minuten, während die Pods die Modellgewichtungen vorab abrufen, den Ray-Cluster mit zwei Knoten über alle acht TPU-Chips hinweg bilden und das Modell kompilieren. Warten Sie, bis in den Logs Application startup complete angezeigt wird, bevor Sie Traffic senden.

  4. Leiten Sie in einer zweiten Terminalsitzung den lokalen Port 8000 an vllm-service weiter:

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

    Lassen Sie die Portweiterleitung in dieser Terminalsitzung laufen.

Inferenzanfragen senden

Kehren Sie zu Ihrer ersten Terminalsitzung zurück und testen Sie das bereitgestellte Modell mit Text- und multimodalen Chatausführungsanfragen:

  1. Senden Sie eine Anfrage zur Textchat-Vervollständigung an den /v1/chat/completions-Endpunkt:

    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. Senden Sie eine multimodale Anfrage, die sowohl eine Bild-URL als auch einen Text-Prompt enthält:

    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
    }'

Benchmark-Arbeitslast ausführen

Führen Sie den folgenden Befehl aus, um das bereitgestellte Modell zu benchmarken:

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

Bereinigen

Löschen Sie die Kubernetes-Arbeitslasten, den TPU-Knotenpool und den GKE-Cluster, um zu vermeiden, dass Ihrem Google Cloud Konto die in dieser Anleitung verwendeten Ressourcen in Rechnung gestellt werden:

  1. Löschen Sie das vLLM-Server-Deployment und das Hugging Face-Secret:

    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. Löschen Sie den TPU v6e-Knotenpool und den GKE-Cluster:

    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. Kehren Sie zur zweiten Terminalsitzung zurück und beenden Sie die Portweiterleitung.

Nächste Schritte