Google utilise la technologie IA pour traduire le contenu dans votre langue préférée. Les traductions générées par IA peuvent contenir des erreurs.

Diffuser Gemma 4 26B sur GKE avec vLLM TorchTPU multi-hôte sur des TPU Trillium

Ce tutoriel vous explique comment déployer et diffuser le modèle de mélange d'experts multimodal Gemma 4 26B (google/gemma-4-26B-A4B-it) sur Google Kubernetes Engine (GKE) à l'aide d'accélérateurs Cloud TPU v6e (Trillium) multi-hôtes et de vLLM TorchTPU. Étant donné que la charge de travail s'étend sur plusieurs hôtes TPU dans une topologie 2x4, vous utilisez l'API Kubernetes LeaderWorkerSet et Ray pour coordonner l'inférence distribuée en parallèle de tenseurs sur huit puces TPU.

Objectifs

Configurez votre cluster GKE et votre pool de nœuds TPU v6e multi-hôte.
Déployez le serveur vLLM multi-hôte avec LeaderWorkerSet.
Envoyer des requêtes d'inférence textuelles et multimodales
Exécutez une charge de travail de benchmark de diffusion.
Effectuer un nettoyage.

Coûts

Dans ce document, vous utilisez les composants facturables suivants deGoogle Cloud :

Obtenez une estimation des coûts en fonction de votre utilisation prévue à l'aide du simulateur de coût.

Une fois que vous avez terminé les tâches décrites dans ce document, vous pouvez éviter de continuer à payer des frais en supprimant les ressources que vous avez créées. Pour en savoir plus, consultez la section Effectuer un nettoyage.

Avant de commencer

Vous avez besoin d'un jeton d'accès Hugging Face pour suivre ce tutoriel. Vous pouvez créer un compte sans frais sur Hugging Face. Une fois que vous disposez d'un compte, générez un jeton d'accès :

Sur la page Welcome to Hugging Face (Bienvenue sur Hugging Face), cliquez sur l'avatar de votre compte, puis sélectionnez Access tokens (Jetons d'accès).
Sur la page Access tokens (Jetons d'accès), cliquez sur Create new token (Créer un jeton).
Sélectionnez le type de jeton Read (Lecture) et saisissez un nom pour votre jeton.
Votre jeton d'accès s'affiche. Enregistrez le jeton dans un endroit sûr.

Installez et initialisez la Google Cloud CLI, puis installez kubectl. Si vous utilisez Cloud Shell, les deux outils sont préinstallés.

Pour obtenir les autorisations nécessaires pour suivre ce tutoriel, demandez à votre administrateur de vous accorder les rôles IAM suivants sur votre projet :

Administrateur Kubernetes Engine (roles/container.admin)
Administrateur de Compute (roles/compute.admin)
Utilisateur du compte de service (roles/iam.serviceAccountUser)

Pour en savoir plus sur l'attribution de rôles, consultez la page Gérer l'accès aux projets, aux dossiers et aux organisations.

Vous pouvez également obtenir les autorisations requises avec des rôles personnalisés ou d'autres rôles prédéfinis.

Configurer vos variables d'environnement

Configurez vos variables d'environnement :

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

Remplacez les éléments suivants :

YOUR_PROJECT_ID : ID de votre projet Google Cloud.
YOUR_RESERVATION_NAME : nom de votre réservation.
YOUR_REGION : région dans laquelle vous souhaitez déployer votre cluster GKE.
YOUR_ZONE : zone dans laquelle vous souhaitez déployer votre pool de nœuds TPU v6e.
YOUR_CLUSTER_NAME : nom de votre cluster GKE.
YOUR_HF_TOKEN : votre jeton d'accès Hugging Face.
YOUR_VLLM_TPU_IMAGE : image TPU vLLM à utiliser. Pour les images publiques, consultez le dépôt vLLM TorchTPU.

Créer le cluster GKE

Pour créer le cluster GKE et associer un pool de nœuds TPU v6e multi-hôte, procédez comme suit :

Créez un cluster GKE régional qui utilise des nœuds e2-standard-4 pour le pool de nœuds par défaut, et un pool de nœuds TPU v6e avec deux nœuds ct6e-standard-4t organisés selon une topologie 2x4 (huit puces TPU au total). Le type de machine GKE par défaut (e2-medium) ne fournit pas suffisamment de ressources de processeur pour le contrôleur LeaderWorkerSet. Ce tutoriel spécifie donc 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}"
```
Au lieu d'utiliser une réservation spécifique, vous pouvez également créer le pool de nœuds à l'aide de la capacité Spot ou de VM à démarrage flexible en transmettant les options suivantes :
- VM Spot : remplacez les options --reservation-affinity et --reservation par --spot pour utiliser la capacité Spot. La création des instances qui utilisent la capacité Spot peut prendre du temps après la création du pool de nœuds, et elles peuvent être préemptées pendant l'exécution des charges de travail. Si vos pods mettent du temps à être créés ou si les charges de travail s'arrêtent en cours d'exécution, cela peut en être la cause.
- VM à démarrage flexible : transmettez --flex-start pour provisionner le pool de nœuds à l'aide de l'option de consommation Démarrage flexible. Pour en savoir plus, consultez Exécuter une petite charge de travail par lot avec des TPU et des VM à démarrage flexible.

Configurez les identifiants kubectl pour communiquer avec votre cluster :

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

Configurer les identifiants et installer LeaderWorkerSet

Pour préparer votre cluster au déploiement multi-hôte, procédez comme suit :

Exportez votre jeton Hugging Face et créez un secret Kubernetes nommé hf-secret :

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

Installez la CustomResourceDefinition et le contrôleur LeaderWorkerSet sur votre 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

Déployer le serveur vLLM multi-hôte

Le fichier manifeste de la charge de travail (gemma4-26b-multihost-v6e-2x4.yaml) définit deux ressources :

Une ressource LeaderWorkerSet (vllm-tpu-gemma4-26b) avec size: 2 (un pod leader et un pod de nœud de calcul) planifiée sur les deux nœuds ct6e-standard-4t de la tranche TPU 2x4. Chaque pod exécute un InitContainer qui précharge les pondérations du modèle google/gemma-4-26B-A4B-it depuis Hugging Face dans un volume en mémoire avant de démarrer Ray et le serveur d'API OpenAI vLLM avec les options --tensor-parallel-size=8, --kv-cache-dtype=fp8, --max-model-len=5248 et --max-num-seqs=128.
Un service ClusterIP (vllm-service) qui expose le port 8000 sur le pod leader.

Pour déployer et surveiller la charge de travail, procédez comme suit :

Récupérez le fichier manifeste LeaderWorkerSet depuis GitHub :

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

Vous pouvez également créer un fichier nommé gemma4-26b-multihost-v6e-2x4.yaml avec le contenu suivant :

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

Appliquez le fichier manifeste 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

Surveillez l'état et les journaux du pod leader jusqu'à ce que le démarrage de l'application soit terminé :

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

Le démarrage initial du service prend environ neuf minutes, pendant lesquelles les pods préchargent les pondérations du modèle, forment le cluster Ray à deux nœuds sur les huit puces TPU et compilent le modèle. Attendez que les journaux signalent Application startup complete avant d'envoyer du trafic.

Dans une deuxième session de terminal, transférez le port local 8000 vers vllm-service :
```
kubectl port-forward service/vllm-service 8000:8000
```
Laissez le transfert de port s'exécuter dans cette session de terminal.

Envoyer des requêtes d'inférence

Revenez à votre première session de terminal et testez le modèle déployé avec des requêtes de complétion de chat textuelles et multimodales :

Envoyez une requête de complétion de chat textuel au point de terminaison /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
    }'

Envoyez une requête multimodale contenant à la fois l'URL d'une image et un prompt textuel :

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

Exécuter une charge de travail de benchmark

Pour évaluer le modèle déployé, exécutez la commande suivante :

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

Effectuer un nettoyage

Pour éviter que des frais supplémentaires ne soient facturés sur votre compte Google Cloud pour les ressources utilisées dans ce tutoriel, supprimez les charges de travail Kubernetes, le pool de nœuds TPU et le cluster GKE :

Supprimez le déploiement du serveur vLLM et le secret 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

Supprimez le pool de nœuds TPU v6e et le 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

Revenez à votre deuxième session de terminal et arrêtez le transfert de port.

Étapes suivantes

En savoir plus sur vLLM TorchTPU.
En savoir plus sur Cloud TPU