Menyajikan Gemma 4 26B di GKE dengan TorchTPU vLLM multi-host di TPU Trillium

Tutorial ini menunjukkan cara men-deploy dan menyajikan model mixture of experts multimodal Gemma 4 26B (google/gemma-4-26B-A4B-it) di Google Kubernetes Engine (GKE) menggunakan akselerator Cloud TPU v6e (Trillium) multi-host dan vLLM TorchTPU. Karena workload mencakup beberapa host TPU dalam topologi 2x4, Anda menggunakan LeaderWorkerSet API Kubernetes dan Ray untuk mengoordinasikan inferensi tensor-paralel terdistribusi di delapan chip TPU.

Tujuan

  • Siapkan cluster GKE dan node pool TPU v6e multi-host.
  • Deploy server vLLM multi-host dengan LeaderWorkerSet.
  • Mengirim permintaan inferensi multimodal dan teks.
  • Menjalankan workload benchmark inferensi.
  • Jalankan pembersihan.

Biaya

Dalam dokumen ini, Anda akan menggunakan komponenGoogle Cloudyang dapat ditagih berikut:

Untuk membuat perkiraan biaya berdasarkan proyeksi penggunaan Anda, gunakan kalkulator harga.

Setelah menyelesaikan tugas yang dijelaskan dalam dokumen ini, Anda dapat menghindari penagihan berkelanjutan dengan menghapus resource yang Anda buat. Untuk mengetahui informasi selengkapnya, lihat Pembersihan.

Sebelum memulai

Anda memerlukan token akses Hugging Face untuk menggunakan tutorial ini. Anda dapat mendaftar untuk mendapatkan akun gratis di Hugging Face. Setelah Anda memiliki akun, buat token akses:

  1. Di halaman Welcome to Hugging Face, klik avatar akun Anda, lalu pilih Token akses.
  2. Di halaman Token akses, klik Create new token.
  3. Pilih jenis token Read dan masukkan nama untuk token Anda.
  4. Token akses Anda akan ditampilkan. Simpan token di tempat yang aman.

Untuk mendapatkan izin yang Anda perlukan untuk menyelesaikan tutorial ini, minta administrator untuk memberi Anda peran IAM berikut di project Anda:

Untuk mengetahui informasi selengkapnya tentang pemberian peran, lihat Mengelola akses ke project, folder, dan organisasi.

Anda mungkin juga bisa mendapatkan izin yang diperlukan melalui peran khusus atau peran bawaan lainnya.

Menyiapkan variabel lingkungan Anda

Siapkan variabel lingkungan Anda:

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

Ganti kode berikut:

  • YOUR_PROJECT_ID: ID project Google Cloud Anda.
  • YOUR_RESERVATION_NAME: nama pemesanan Anda.
  • YOUR_REGION: region tempat Anda ingin men-deploy cluster GKE.
  • YOUR_ZONE: zona tempat Anda ingin men-deploy node pool TPU v6e.
  • YOUR_CLUSTER_NAME: nama cluster GKE Anda.
  • YOUR_HF_TOKEN: token akses Hugging Face Anda.
  • YOUR_VLLM_TPU_IMAGE: image TPU vLLM yang akan digunakan. Untuk image publik, lihat repositori vLLM TorchTPU.

Membuat cluster GKE

Untuk membuat cluster GKE dan memasang node pool TPU v6e multi-host, selesaikan langkah-langkah berikut:

  1. Buat cluster GKE regional yang menggunakan node e2-standard-4 untuk node pool default, dan node pool TPU v6e dengan dua node ct6e-standard-4t yang disusun dalam topologi 2x4 (total delapan chip TPU). Jenis mesin GKE default (e2-medium) tidak menyediakan resource CPU yang cukup untuk pengontrol LeaderWorkerSet, sehingga tutorial ini menentukan 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}"

    Alih-alih menggunakan pemesanan tertentu, Anda juga dapat membuat node pool menggunakan kapasitas Spot atau VM Flex-start dengan meneruskan flag berikut:

    • Spot VM: ganti flag --reservation-affinity dan --reservation dengan --spot untuk menggunakan kapasitas Spot. Instance yang menggunakan kapasitas Spot mungkin memerlukan waktu untuk dibuat setelah node pool dibuat, dan mungkin di-preempt selama workload berjalan. Jika Pod Anda memerlukan waktu untuk dibuat, atau workload berhenti di tengah-tengah proses, hal ini bisa menjadi penyebabnya.
    • Flex-start VM: Teruskan --flex-start untuk menyediakan node pool menggunakan opsi konsumsi Flex-start. Untuk mengetahui informasi selengkapnya, lihat Menjalankan workload batch kecil dengan TPU dan VM Flex-start.

  2. Konfigurasi kredensial kubectl untuk berkomunikasi dengan cluster Anda:

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

Mengonfigurasi kredensial dan menginstal LeaderWorkerSet

Untuk menyiapkan cluster Anda untuk deployment multi-host, selesaikan langkah-langkah berikut:

  1. Ekspor token Hugging Face Anda dan buat Secret Kubernetes bernama hf-secret:

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

  2. Instal CustomResourceDefinition dan pengontrol LeaderWorkerSet di cluster Anda:

    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

Men-deploy server vLLM multi-host

Manifes workload (gemma4-26b-multihost-v6e-2x4.yaml) menentukan dua resource:

  • Resource LeaderWorkerSet (vllm-tpu-gemma4-26b) dengan size: 2 (satu Pod leader dan satu Pod worker) yang dijadwalkan di dua node ct6e-standard-4t dalam slice TPU 2x4. Setiap Pod menjalankan InitContainer yang mengambil bobot model google/gemma-4-26B-A4B-it dari Hugging Face ke dalam volume dalam memori sebelum memulai Ray dan server API OpenAI vLLM dengan flag --tensor-parallel-size=8, --kv-cache-dtype=fp8, --max-model-len=5248, dan --max-num-seqs=128.
  • Service ClusterIP (vllm-service) yang mengekspos port 8000 di Pod leader.

Untuk men-deploy dan memantau workload, selesaikan langkah-langkah berikut:

  1. Ambil manifes LeaderWorkerSet dari GitHub:

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

    Atau, buat file bernama gemma4-26b-multihost-v6e-2x4.yaml dengan konten berikut:

    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. Terapkan manifes 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. Pantau status dan log Pod pemimpin hingga startup aplikasi selesai:

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

    Startup layanan awal memerlukan waktu sekitar 9 menit saat Pod mengambil bobot model terlebih dahulu, membentuk cluster Ray dua node di delapan chip TPU, dan mengompilasi model. Tunggu hingga log melaporkan Application startup complete sebelum mengirim traffic.

  4. Di sesi terminal kedua, teruskan port lokal 8000 ke vllm-service:

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

    Biarkan penerusan port berjalan di sesi terminal ini.

Mengirim permintaan inferensi

Kembali ke sesi terminal pertama Anda dan uji model yang di-deploy dengan permintaan penyelesaian chat teks dan multimodal:

  1. Kirim permintaan penyelesaian chat teks ke 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. Kirim permintaan multimodal yang berisi URL gambar dan perintah teks:

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

Menjalankan workload benchmark

Untuk melakukan benchmark pada model yang di-deploy, jalankan perintah berikut:

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

Pembersihan

Agar tidak menimbulkan biaya berkelanjutan tambahan pada akun Google Cloud Anda untuk resource yang digunakan dalam tutorial ini, hapus workload Kubernetes, TPU node pool, dan cluster GKE:

  1. Hapus deployment server vLLM dan 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

  2. Hapus node pool TPU v6e dan 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. Kembali ke sesi terminal kedua Anda dan hentikan penerusan port.

Langkah berikutnya