Disponibilizar o Gemma 4 26B no GKE com o vLLM TorchTPU de vários hosts em TPUs Trillium

Neste tutorial, mostramos como implantar e disponibilizar o modelo multimodal de mix de especialistas Gemma 4 26B (google/gemma-4-26B-A4B-it) no Google Kubernetes Engine (GKE) usando aceleradores Cloud TPU v6e (Trillium) de vários hosts e vLLM TorchTPU. Como a carga de trabalho abrange vários hosts de TPU em uma topologia 2x4, use a API LeaderWorkerSet do Kubernetes e o Ray para coordenar a inferência distribuída em paralelo de tensor em oito chips de TPU.

Objetivos

• Configure o cluster do GKE e o pool de nós de TPU v6e de vários hosts.
• Implantar o servidor vLLM de vários hosts com o LeaderWorkerSet.
• Enviar solicitações de inferência de texto e multimodais.
• Executar uma carga de trabalho de comparativo de mercado de exibição.
• Fazer a limpeza.

Custos

Neste documento, você usará os seguintes componentes faturáveis do Google Cloud:

• GKE
• Cloud TPU
• Compute Engine

Para gerar uma estimativa de custo baseada na projeção de uso deste tutorial, use a calculadora de preços.

Ao concluir as tarefas descritas neste documento, é possível evitar o faturamento contínuo excluindo os recursos criados. Para mais informações, consulte Limpeza.

Importante: para concluir este tutorial, você precisa ter acesso às TPUs v6e-8. Este tutorial pressupõe que você já tenha uma reserva para essas TPUs no seu projeto doGoogle Cloud para garantir que tenha a capacidade disponível quando precisar. Para saber mais sobre as opções de reserva de capacidade do Cloud TPU para uma data e hora futuras, consulte Sobre as reservas do Cloud TPU.

Antes de começar

Você precisa de um token de acesso do Hugging Face para usar este tutorial. Você pode se inscrever para uma conta sem custo financeiro no Hugging Face. Depois de criar uma conta, gere um token de acesso:

1. Na página Welcome to Hugging Face, clique no avatar da sua conta e selecione token de acesso.
2. Na página Tokens de acesso, clique em Criar novo token.
3. Selecione o tipo de token Read e insira um nome para ele.
4. O token de acesso será exibido. Salve o token em um local seguro.

• Instale e inicialize a Google Cloud CLI e instale kubectl. Se você usa o Cloud Shell, as duas ferramentas já estão pré-instaladas.

Para conseguir as permissões necessárias para concluir o tutorial, peça ao administrador para conceder a você os seguintes papéis do IAM no seu projeto:

• Administrador do Kubernetes Engine (roles/container.admin)
• Administrador do Compute (roles/compute.admin)
• Usuário da conta de serviço (roles/iam.serviceAccountUser)

Para mais informações sobre a concessão de papéis, consulte Gerenciar o acesso a projetos, pastas e organizações.

Também é possível conseguir as permissões necessárias usando papéis personalizados ou outros papéis predefinidos.

Configurar as variáveis de ambiente

Configure as variáveis de 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}"

Substitua:

• YOUR_PROJECT_ID: ID do seu projeto do Google Cloud.
• YOUR_RESERVATION_NAME: o nome da sua reserva.
• YOUR_REGION: a região em que você quer implantar o cluster do GKE.
• YOUR_ZONE: a zona em que você quer implantar o pool de nós da TPU v6e.
• YOUR_CLUSTER_NAME: o nome do cluster do GKE.
• YOUR_HF_TOKEN: seu token de acesso do Hugging Face.
• YOUR_VLLM_TPU_IMAGE: a imagem de TPU do vLLM a ser usada. Para imagens públicas, consulte o repositório vLLM TorchTPU.

Criar o cluster do GKE

Para criar o cluster do GKE e anexar um pool de nós de TPU v6e de vários hosts, siga estas etapas:

1. Crie um cluster regional do GKE que use nós e2-standard-4 para o pool de nós padrão e um pool de nós da TPU v6e com dois nós ct6e-standard-4t organizados em uma topologia 2x4 (oito chips de TPU no total). O tipo de máquina padrão do GKE (e2-medium) não fornece recursos de CPU suficientes para o controlador LeaderWorkerSet. Por isso, este tutorial 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}"

   Em vez de usar uma reserva específica, também é possível criar o pool de nós usando a capacidade spot ou VMs flex-start transmitindo as seguintes flags:

   • VMs spot: substitua as flags --reservation-affinity e --reservation por --spot para usar a capacidade de spot. As instâncias que usam a capacidade do Spot podem levar algum tempo para serem criadas após a criação do pool de nós e podem ser interrompidas durante as execuções da carga de trabalho. Se os pods demorarem para ser criados ou as cargas de trabalho pararem no meio da execução, essa pode ser a causa.
   • VMs flex-start: transmita --flex-start para provisionar o pool de nós usando a opção de consumo de início flexível. Para mais informações, consulte Executar uma carga de trabalho em lote pequena com TPUs e VMs de início flexível.

2. Configure as credenciais de kubectl para se comunicar com o cluster:

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

Configurar credenciais e instalar o LeaderWorkerSet

Para preparar o cluster para a implantação de vários hosts, siga estas etapas:

1. Exporte o token do Hugging Face e crie um secret do Kubernetes chamado hf-secret:

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

2. Instale o CustomResourceDefinition e o controlador do LeaderWorkerSet no 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

Implantar o servidor vLLM de vários hosts

O manifesto da carga de trabalho (gemma4-26b-multihost-v6e-2x4.yaml) define dois recursos:

• Um recurso LeaderWorkerSet (vllm-tpu-gemma4-26b) com size: 2 (um pod líder e um pod de worker) programado nos dois nós ct6e-standard-4t na fração de TPU 2x4. Cada pod executa um InitContainer que busca previamente os pesos do modelo google/gemma-4-26B-A4B-it do Hugging Face em um volume na memória antes de iniciar o Ray e o servidor de API OpenAI do vLLM com as flags --tensor-parallel-size=8, --kv-cache-dtype=fp8, --max-model-len=5248 e --max-num-seqs=128.
• Um serviço ClusterIP (vllm-service) que expõe a porta 8000 no pod líder.

Para implantar e monitorar a carga de trabalho, siga estas etapas:

1. Busque o manifesto do LeaderWorkerSet no GitHub:

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

   Como alternativa, crie um arquivo chamado gemma4-26b-multihost-v6e-2x4.yaml com o seguinte conteúdo:

   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. Aplique o manifesto do 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. Monitore o status e os registros do pod líder até que a inicialização do aplicativo seja concluída:

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

   A inicialização do serviço leva aproximadamente 9 minutos, enquanto os pods buscam os pesos do modelo, formam o cluster do Ray de dois nós em todos os oito chips de TPU e compilam o modelo. Aguarde até que os registros informem Application startup complete antes de enviar tráfego.

4. Em uma segunda sessão de terminal, encaminhe a porta local 8000 para vllm-service:

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

   Deixe o encaminhamento de portas em execução nessa sessão do terminal.

Enviar solicitações de inferência

Volte para a primeira sessão do terminal e teste o modelo implantado com solicitações de conclusão de chat de texto e multimodais:

1. Envie uma solicitação de conclusão de chat de texto para o 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. Envie uma solicitação multimodal que contenha um URL de imagem e um comando 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
       }'

Executar uma carga de trabalho de comparativo de mercado

Para comparar o modelo implantado, execute o seguinte 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

Limpar

Para evitar cobranças extras na sua conta do Google Cloud pelos recursos usados neste tutorial, exclua as cargas de trabalho do Kubernetes, o pool de nós da TPU e o cluster do GKE:

1. Exclua a implantação do servidor vLLM e o secret do 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. Exclua o pool de nós da TPU v6e e o cluster do 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. Volte para a segunda sessão do terminal e interrompa o encaminhamento de portas.

A seguir

• Saiba mais sobre o vLLM TorchTPU.
• Saiba mais sobre o Cloud TPU.