Google utiliza tecnología de IA para traducir contenido a tu idioma preferido. Las traducciones realizadas con IA pueden contener errores.

Entrega Gemma 4 26B en GKE con vLLM TorchTPU de varios hosts en TPU Trillium

En este instructivo, se muestra cómo implementar y entregar el modelo de mezcla de expertos multimodal Gemma 4 26B (google/gemma-4-26B-A4B-it) en Google Kubernetes Engine (GKE) con aceleradores de Cloud TPU v6e (Trillium) de varios hosts y vLLM TorchTPU. Debido a que la carga de trabajo abarca varios hosts de TPU en una topología 2x4, usas la API de LeaderWorkerSet de Kubernetes y Ray para coordinar la inferencia distribuida paralela de tensores en ocho chips TPU.

Objetivos

Configura tu clúster de GKE y el grupo de nodos de TPU v6e de varios hosts.
Implementa el servidor de vLLM de varios hosts con LeaderWorkerSet.
Enviar solicitudes de inferencia multimodales y de texto
Ejecuta una carga de trabajo de comparativas de entrega.
Realizar una limpieza

Costos

En este documento, usarás los siguientes componentes facturables de Google Cloud:

Para generar una estimación de costos en función del uso previsto, usa la calculadora de precios.

Cuando finalices las tareas que se describen en este documento, puedes borrar los recursos que creaste para evitar que continúe la facturación. Para obtener más información, consulta Cómo realizar una limpieza.

Antes de comenzar

Necesitas un token de acceso de Hugging Face para usar este instructivo. Puedes registrarte para obtener una cuenta gratuita en Hugging Face. Después de tener una cuenta, genera un token de acceso:

En la página Welcome to Hugging Face, haz clic en el avatar de tu cuenta y selecciona Tokens de acceso.
En la página Tokens de acceso, haz clic en Crear token nuevo.
Selecciona el tipo de token Lectura y, luego, ingresa un nombre para el token.
Se mostrará tu token de acceso. Guarda el token en un lugar seguro.

Instala y, luego, inicializa Google Cloud CLI, y, luego, instala kubectl. Si usas Cloud Shell, ambas herramientas están preinstaladas.

Si quieres obtener los permisos que necesitas para completar este instructivo, pídele a tu administrador que te otorgue los siguientes roles de IAM en tu proyecto:

Administrador de Kubernetes Engine (roles/container.admin)
Administrador de Compute (roles/compute.admin)
Usuario de la cuenta de servicio (roles/iam.serviceAccountUser)

Para obtener más información sobre cómo otorgar roles, consulta Administra el acceso a proyectos, carpetas y organizaciones.

También puedes obtener los permisos necesarios a través de roles personalizados o cualquier otro rol predefinido.

Configura tus variables de entorno

Configura las variables de entorno:

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

Reemplaza lo siguiente:

YOUR_PROJECT_ID: Es el ID de tu proyecto de Google Cloud.
YOUR_RESERVATION_NAME: Es el nombre de tu reserva.
YOUR_REGION: Es la región en la que deseas implementar tu clúster de GKE.
YOUR_ZONE: Es la zona en la que deseas implementar tu grupo de nodos de TPU v6e.
YOUR_CLUSTER_NAME: el nombre del clúster de GKE.
YOUR_HF_TOKEN: Es tu token de acceso de Hugging Face.
YOUR_VLLM_TPU_IMAGE: Es la imagen de TPU de vLLM que se usará. Para obtener imágenes públicas, consulta el repositorio de vLLM TorchTPU.

Crea el clúster de GKE

Para crear el clúster de GKE y conectar un grupo de nodos de TPU v6e de varios hosts, completa los siguientes pasos:

Crea un clúster de GKE regional que use nodos e2-standard-4 para el grupo de nodos predeterminado y un grupo de nodos TPU v6e con dos nodos ct6e-standard-4t organizados en una topología 2x4 (ocho chips TPU en total). El tipo de máquina predeterminado de GKE (e2-medium) no proporciona suficientes recursos de CPU para el controlador LeaderWorkerSet, por lo que este instructivo especifica 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}"
```
En lugar de usar una reserva específica, también puedes crear el grupo de nodos con capacidad Spot o VMs de inicio flexible pasando las siguientes marcas:
- VMs Spot: Reemplaza las marcas --reservation-affinity y --reservation por --spot para usar la capacidad de Spot. Las instancias que usan la capacidad de Spot pueden tardar en crearse después de que se crea el grupo de nodos y se pueden interrumpir durante las ejecuciones de las cargas de trabajo. Si tus Pods tardan en crearse o las cargas de trabajo se detienen en medio de la ejecución, esta podría ser la causa.
- VMs de inicio flexible: Pasa --flex-start para aprovisionar el grupo de nodos con la opción de consumo Flex-start. Para obtener más información, consulta Ejecuta una carga de trabajo por lotes pequeña con TPU y VMs de inicio flexible.

Configura las credenciales de kubectl para comunicarte con tu clúster:

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

Configura las credenciales e instala LeaderWorkerSet

Para preparar tu clúster para la implementación de varios hosts, completa los siguientes pasos:

Exporta tu token de Hugging Face y crea un Secret de Kubernetes llamado hf-secret:

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

Instala la CustomResourceDefinition y el controlador de LeaderWorkerSet en tu clúster:

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

Implementa el servidor de vLLM de varios hosts

El manifiesto de la carga de trabajo (gemma4-26b-multihost-v6e-2x4.yaml) define dos recursos:

Un recurso LeaderWorkerSet (vllm-tpu-gemma4-26b) con size: 2 (un Pod líder y un Pod trabajador) programado en los dos nodos ct6e-standard-4t en la porción de TPU 2x4. Cada Pod ejecuta un InitContainer que recupera previamente los pesos del modelo google/gemma-4-26B-A4B-it de Hugging Face en un volumen en la memoria antes de iniciar Ray y el servidor de la API de OpenAI de vLLM con las marcas --tensor-parallel-size=8, --kv-cache-dtype=fp8, --max-model-len=5248 y --max-num-seqs=128.
Un Service ClusterIP (vllm-service) que expone el puerto 8000 en el Pod líder.

Para implementar y supervisar la carga de trabajo, completa los siguientes pasos:

Recupera el manifiesto de LeaderWorkerSet de GitHub:

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

Como alternativa, crea un archivo llamado gemma4-26b-multihost-v6e-2x4.yaml con el siguiente contenido:

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

Aplica el manifiesto de 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

Supervisa el estado y los registros del Pod líder hasta que se complete el inicio de la aplicación:

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

El inicio inicial del servicio tarda alrededor de 9 minutos mientras los Pods recuperan previamente los pesos del modelo, forman el clúster de Ray de dos nodos en los ocho chips TPU y compilan el modelo. Espera hasta que los registros informen Application startup complete antes de enviar tráfico.

En una segunda sesión de terminal, reenvía el puerto local 8000 a vllm-service:
```
kubectl port-forward service/vllm-service 8000:8000
```
Deja la redirección de puertos en ejecución en esta sesión de terminal.

Envía solicitudes de inferencia

Regresa a tu primera sesión de terminal y prueba el modelo implementado con solicitudes de finalización de chat de texto y multimodales:

Envía una solicitud de finalización de chat de texto al extremo /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
    }'

Envía una solicitud multimodal que contenga una URL de imagen y una instrucción de texto:

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

Ejecuta una carga de trabajo de comparativas

Para comparar el modelo implementado, ejecuta el siguiente 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

Realiza una limpieza

Para evitar que se apliquen cargos adicionales continuos a tu cuenta de Google Cloud por los recursos que se usaron en este instructivo, borra las cargas de trabajo de Kubernetes, el grupo de nodo TPU y el clúster de GKE:

Borra la implementación del servidor de vLLM y el secreto de 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

Borra el grupo de nodos TPU v6e y el clúster de 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

Regresa a tu segunda sesión de terminal y detén la redirección de puertos.

¿Qué sigue?

Obtén más información sobre vLLM TorchTPU.
Más información sobre Cloud TPU.