Eroga Gemma 4 26B su GKE con vLLM TorchTPU multi-host su TPU Trillium

Questo tutorial mostra come eseguire il deployment e l'inferenza del modello Mixture of Experts multimodale Gemma 4 26B (google/gemma-4-26B-A4B-it) su Google Kubernetes Engine (GKE) utilizzando acceleratori Cloud TPU v6e (Trillium) multi-host e vLLM TorchTPU. Poiché il workload si estende su più host TPU in una topologia 2x4, utilizzi l'API LeaderWorkerSet di Kubernetes e Ray per coordinare l'inferenza distribuita con parallelismo dei tensori su otto chip TPU.

Obiettivi

  • Configura il cluster GKE e il pool di nodi TPU v6e multi-host.
  • Esegui il deployment del server vLLM multi-host con LeaderWorkerSet.
  • Inviare richieste di inferenza di testo e multimodali.
  • Esegui un workload di benchmark di distribuzione.
  • Eseguire la pulizia.

Costi

In questo documento vengono utilizzati i seguenti componenti fatturabili diGoogle Cloud:

Per generare una stima dei costi in base all'utilizzo previsto, utilizza il calcolatore prezzi.

Al termine delle attività descritte in questo documento, puoi evitare l'addebito di ulteriori costi eliminando le risorse che hai creato. Per saperne di più, consulta Esegui la pulizia.

Prima di iniziare

Per utilizzare questo tutorial, devi disporre di un token di accesso a Hugging Face. Puoi creare un account senza costi su Hugging Face. Dopo aver creato un account, genera un token di accesso:

  1. Nella pagina Welcome to Hugging Face, fai clic sull'avatar del tuo account e seleziona Token di accesso.
  2. Nella pagina Token di accesso, fai clic su Crea nuovo token.
  3. Seleziona il tipo di token Lettura e inserisci un nome per il token.
  4. Viene visualizzato il token di accesso. Salva il token in un luogo sicuro.

Per ottenere le autorizzazioni necessarie per completare questo tutorial, chiedi all'amministratore di concederti i seguenti ruoli IAM nel progetto:

Per saperne di più sulla concessione dei ruoli, consulta Gestisci l'accesso a progetti, cartelle e organizzazioni.

Potresti anche riuscire a ottenere le autorizzazioni richieste tramite i ruoli personalizzati o altri ruoli predefiniti.

Configura le variabili di ambiente

Imposta le variabili di ambiente:

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

Sostituisci quanto segue:

  • YOUR_PROJECT_ID: l'ID del tuo progetto Google Cloud.
  • YOUR_RESERVATION_NAME: il nome della tua prenotazione.
  • YOUR_REGION: la regione in cui vuoi eseguire il deployment del cluster GKE.
  • YOUR_ZONE: la zona in cui vuoi eseguire il deployment del pool di nodi TPU v6e.
  • YOUR_CLUSTER_NAME: il nome del tuo cluster GKE.
  • YOUR_HF_TOKEN: il tuo token di accesso a Hugging Face.
  • YOUR_VLLM_TPU_IMAGE: l'immagine TPU vLLM da utilizzare. Per le immagini pubbliche, consulta il repository vLLM TorchTPU.

Crea il cluster GKE

Per creare il cluster GKE e collegare un pool di nodi TPU v6e multi-host, completa i seguenti passaggi:

  1. Crea un cluster GKE a livello di regione che utilizza nodi e2-standard-4 per il pool di nodi predefinito e un pool di nodi TPU v6e con due nodi ct6e-standard-4t disposti in una topologia 2x4 (otto chip TPU in totale). Il tipo di macchina GKE predefinito (e2-medium) non fornisce risorse CPU sufficienti per il controller LeaderWorkerSet, pertanto questo tutorial specifica 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}"

    Anziché utilizzare una prenotazione specifica, puoi anche creare il pool di nodi utilizzando la capacità Spot o le Flex-start VM passando i seguenti flag:

    • VM spot: sostituisci i flag --reservation-affinity e --reservation con --spot per utilizzare la capacità spot. Le istanze che utilizzano la capacità Spot potrebbero richiedere del tempo per essere create dopo la creazione del pool di nodi e potrebbero essere prerilasciate durante l'esecuzione dei workload. Se i pod richiedono tempo per essere creati o i workload si interrompono durante l'esecuzione, questa potrebbe essere la causa.
    • VM con avvio flessibile: passa --flex-start per eseguire il provisioning del node pool utilizzando l'opzione di consumo Avvio flessibile. Per saperne di più, consulta Esegui un workload batch di piccole dimensioni con TPU e VM Flex-start.

  2. Configura le credenziali kubectl per comunicare con il tuo cluster:

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

Configura le credenziali e installa LeaderWorkerSet

Per preparare il cluster per il deployment multi-host, completa i seguenti passaggi:

  1. Esporta il tuo token Hugging Face e crea un secret di Kubernetes denominato hf-secret:

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

  2. Installa CustomResourceDefinition e il controller LeaderWorkerSet sul tuo 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

Esegui il deployment del server vLLM multi-host

Il manifest del workload (gemma4-26b-multihost-v6e-2x4.yaml) definisce due risorse:

  • Una risorsa LeaderWorkerSet (vllm-tpu-gemma4-26b) con size: 2 (un pod leader e un pod worker) pianificata tra i due nodi ct6e-standard-4t nella sezione TPU 2x4. Ogni pod esegue un InitContainer che precarica i pesi del modello google/gemma-4-26B-A4B-it da Hugging Face in un volume in memoria prima di avviare Ray e il server API vLLM OpenAI con i flag --tensor-parallel-size=8, --kv-cache-dtype=fp8, --max-model-len=5248 e --max-num-seqs=128.
  • Un servizio ClusterIP (vllm-service) che espone la porta 8000 sul pod leader.

Per eseguire il deployment e monitorare il workload, completa i seguenti passaggi:

  1. Recupera il manifest LeaderWorkerSet da GitHub:

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

    In alternativa, crea un file denominato gemma4-26b-multihost-v6e-2x4.yaml con i seguenti contenuti:

    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. Applica il manifest 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. Monitora lo stato e i log del pod leader fino al completamento dell'avvio dell'applicazione:

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

    L'avvio iniziale del servizio richiede circa 9 minuti, durante i quali i pod eseguono il precaricamento dei pesi del modello, formano il cluster Ray a due nodi su tutti e otto i chip TPU e compilano il modello. Attendi che i log segnalino Application startup complete prima di inviare traffico.

  4. In una seconda sessione del terminale, inoltra la porta locale 8000 a vllm-service:

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

    Lascia il port forwarding in esecuzione in questa sessione del terminale.

Invia richieste di inferenza

Torna alla prima sessione del terminale e testa il modello di cui è stato eseguito il deployment con richieste di completamento della chat sia di testo che multimodali:

  1. Invia una richiesta di completamento della chat di testo all'endpoint /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. Invia una richiesta multimodale contenente sia l'URL di un'immagine sia un prompt testuale:

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

Esegui un workload di benchmark

Per eseguire il benchmark del modello di cui è stato eseguito il deployment, esegui questo comando:

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

Esegui la pulizia

Per evitare che al tuo account Google Cloud vengano addebitati costi aggiuntivi relativi alle risorse utilizzate in questo tutorial, elimina i workload Kubernetes, il pool di nodi TPU e il cluster GKE:

  1. Elimina il deployment del server vLLM e il secret di 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. Elimina il pool di nodi TPU v6e e il cluster 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. Torna alla seconda sessione del terminale e interrompi il port forwarding.

Passaggi successivi