הצגת Gemma 4 26B ב-GKE עם vLLM TorchTPU בכמה מארחים ב-Trillium TPU

במדריך הזה נסביר איך לפרוס ולהפעיל את מודל שילוב המומחים (MoE) המולטימודלי Gemma 4 26B (google/gemma-4-26B-A4B-it) ב-Google Kubernetes Engine ‏(GKE) באמצעות מאיצי Cloud TPU v6e (Trillium) עם כמה מארחים ו-vLLM TorchTPU. מכיוון שעומס העבודה מתפרס על פני כמה מארחי TPU בטופולוגיה של 2x4, אתם משתמשים ב-Kubernetes LeaderWorkerSet API וב-Ray כדי לתאם הסקת מסקנות מבוזרת במקביל לטנזור על פני שמונה שבבי TPU.

מטרות

• מגדירים את אשכול GKE ואת מאגר הצמתים של TPU v6e עם כמה מארחים.
• פורסים את שרת ה-vLLM עם כמה מארחים באמצעות LeaderWorkerSet.
• לשלוח בקשות להסקת מסקנות מטקסט וממולטי-מודאליות.
• מריצים עומס עבודה להערכת ביצועים של מילוי בקשות.
• לפנות מקום.

עלויות

במסמך הזה משתמשים ברכיבים הבאים שלGoogle Cloud, והשימוש בהם כרוך בתשלום:

• GKE
• Cloud TPU
• Compute Engine

כדי להעריך את העלויות בהתאם לשימוש הצפוי, אתם יכולים להשתמש במחשבון העלויות.

כשמסיימים את המשימות שמתוארות במסמך הזה אפשר למחוק את המשאבים שיצרתם כדי להימנע מחיובים נוספים. מידע נוסף זמין בקטע הסרת המשאבים.

חשוב: כדי להשלים את המדריך הזה, צריכה להיות לכם גישה ל-v6e-8 TPUs. במדריך הזה אנחנו מניחים שכבר יש לכם מקום שמור ל-TPU האלה בפרויקט ב-Google Cloud , כדי לוודא שהקיבולת זמינה כשאתם צריכים אותה. במאמר מידע על הקצאות מקום ל-Cloud TPU תוכלו לקרוא איך מקצים מקום ל-Cloud TPU לתאריך ולשעה בעתיד.

לפני שמתחילים

כדי להשתמש במדריך הזה, אתם צריכים טוקן גישה ל-Hugging Face. אפשר להירשם לחשבון בחינם ב-Hugging Face. אחרי שיוצרים חשבון, יוצרים אסימון גישה:

1. בדף Welcome to Hugging Face, לוחצים על הדמות של החשבון ובוחרים באפשרות טוקן גישה.
2. בדף טוקני גישה, לוחצים על יצירת טוקן חדש.
3. בוחרים את סוג הטוקן Read ומזינים שם לטוקן.
4. טוקן הגישה שלכם יוצג. שומרים את הטוקן במקום בטוח.

• מתקינים ומפעילים את Google Cloud CLI, ומתקינים את kubectl. אם אתם משתמשים ב-Cloud Shell, שני הכלים מותקנים מראש.

כדי לקבל את ההרשאות שדרושות להשלמת המדריך, בקשו מהאדמין להקצות לכם את התפקידים הבאים ב-IAM בפרויקט:

• אדמין ב-Kubernetes Engine ‏ (roles/container.admin)
• אדמין ב-Compute ‏ (roles/compute.admin)
• משתמש בחשבון שירות (roles/iam.serviceAccountUser)

להסבר על מתן תפקידים, ראו איך מנהלים את הגישה ברמת הפרויקט, התיקייה והארגון.

יכול להיות שאפשר לקבל את ההרשאות הנדרשות גם באמצעות תפקידים בהתאמה אישית או תפקידים מוגדרים מראש.

הגדרת משתני הסביבה

מגדירים את משתני הסביבה:

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

מחליפים את מה שכתוב בשדות הבאים:

• ‫YOUR_PROJECT_ID: מזהה הפרויקט ב- Google Cloud.
• ‫YOUR_RESERVATION_NAME: שם ההזמנה.
• ‫YOUR_REGION: האזור שבו רוצים לפרוס את אשכול GKE.
• ‫YOUR_ZONE: האזור שבו רוצים לפרוס את מאגר הצמתים של TPU v6e.
• ‫YOUR_CLUSTER_NAME: השם של אשכול GKE.
• ‫YOUR_HF_TOKEN: טוקן הגישה ל-Hugging Face.
• ‫YOUR_VLLM_TPU_IMAGE: קובץ האימג' של vLLM TPU שבו רוצים להשתמש. לגבי קובצי אימג' ציבוריים, עיינו במאגר vLLM TorchTPU.

יצירת אשכול GKE

כדי ליצור את אשכול GKE ולצרף מאגר של צומתי TPU v6e עם כמה מארחים:

1. יוצרים אשכול GKE אזורי שמשתמש ב-e2-standard-4 צמתים למאגר הצמתים שמוגדר כברירת מחדל, ובמאגר צמתים של TPU v6e עם שני צמתים מסוג ct6e-standard-4t שמסודרים בטופולוגיה של 2x4 (סה"כ שמונה שבבי TPU). סוג המכונה שמוגדר כברירת מחדל ב-GKE‏ (e2-medium) לא מספק מספיק משאבי מעבד (CPU) לבקר LeaderWorkerSet, ולכן במדריך הזה מצוין 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}"

   במקום להשתמש בהזמנה ספציפית, אתם יכולים גם ליצור את מאגר הצמתים באמצעות קיבולת של Spot או מכונות וירטואליות מסוג Flex-start, על ידי העברת הדגלים הבאים:

   • VM במודל Spot: מחליפים את הדגלים --reservation-affinity ו---reservation בדגל --spot כדי להשתמש בקיבולת של Spot. יכול להיות שייקח זמן ליצור מכונות שמשתמשות בקיבולת של Spot אחרי שנוצר מאגר הצמתים, ויכול להיות שהן יופסקו במהלך הרצת עומסי העבודה. אם לוקח זמן ליצור את ה-Pods, או שעומסי העבודה מפסיקים באמצע הריצה, יכול להיות שזו הסיבה.
   • ‫Flex-start VMs: מעבירים את הערך --flex-start כדי להקצות את מאגר הצמתים באמצעות אפשרות הצריכה Flex-start. מידע נוסף זמין במאמר הרצת עומס עבודה קטן של אצווה באמצעות יחידות TPU ומכונות וירטואליות מסוג Flex-start.

2. מגדירים את פרטי הכניסה של kubectl כדי לתקשר עם האשכול:

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

הגדרת פרטי כניסה והתקנת LeaderWorkerSet

כדי להכין את האשכול לפריסה בכמה מארחים, מבצעים את השלבים הבאים:

1. מייצאים את הטוקן של Hugging Face ויוצרים Kubernetes Secret בשם hf-secret:

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

2. מתקינים את ה-CustomResourceDefinition והבקר של LeaderWorkerSet באשכול:

   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

פריסת שרת vLLM עם כמה מארחים

במניפסט של עומס העבודה (gemma4-26b-multihost-v6e-2x4.yaml) מוגדרים שני משאבים:

• משאב LeaderWorkerSet‏ (vllm-tpu-gemma4-26b) עם size: 2 (Pod אחד של Leader ו-Pod אחד של Worker) מתוזמן בשני הצמתים של ct6e-standard-4t בפרוסת ה-TPU של 2x4. כל Pod מריץ InitContainer שמאחזר מראש את משקלי המודל google/gemma-4-26B-A4B-it מ-Hugging Face לנפח בזיכרון לפני הפעלת Ray ושרת ה-API של vLLM OpenAI עם הדגלים --tensor-parallel-size=8,‏ --kv-cache-dtype=fp8,‏ --max-model-len=5248 ו---max-num-seqs=128.
• שירות ClusterIP‏ (vllm-service) שחושף את יציאה 8000 ב-Pod הראשי.

כדי לפרוס את עומס העבודה ולעקוב אחריו:

1. מורידים את המניפסט של LeaderWorkerSet מ-GitHub:

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

   אפשר גם ליצור קובץ בשם gemma4-26b-multihost-v6e-2x4.yaml עם התוכן הבא:

   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. מחילים את המניפסט של 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. עוקבים אחרי הסטטוס והיומנים של ה-Pod הראשי עד שהאפליקציה מסיימת את תהליך ההפעלה:

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

   ההפעלה הראשונית של השירות נמשכת כ-9 דקות, בזמן שה-Pods מאחזרים מראש את משקלי המודל, יוצרים את אשכול Ray עם שני הצמתים בכל שמונת שבבי ה-TPU ומקמפלים את המודל. לפני ששולחים תנועה, צריך להמתין עד שמופיע ביומנים הדיווח Application startup complete.

4. בסשן מסוף שני, מעבירים את היציאה המקומית 8000 אל vllm-service:

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

   משאירים את ההעברה ליציאה אחרת פועלת בסשן הטרמינל הזה.

שליחת בקשות להסקת מסקנות

חוזרים לסשן הראשון בטרמינל ובודקים את המודל שנפרס באמצעות בקשות להשלמת צ'אט עם טקסט ובקשות מולטי-מודאליות:

1. שולחים בקשה להשלמת צ'אט טקסט לנקודת הקצה (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. שליחת בקשה מולטי-מודאלית שמכילה גם כתובת URL של תמונה וגם פרומפט טקסטואלי:

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

הרצת עומס עבודה להערכת ביצועים

כדי לבצע השוואה לשוק של המודל שנפרס, מריצים את הפקודה הבאה:

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

הסרת המשאבים

כדי להימנע מחיובים שוטפים נוספים בחשבון Google Cloud על המשאבים שבהם השתמשתם במדריך הזה, אתם צריכים למחוק את עומסי העבודה ב-Kubernetes, את מאגר הצמתים של ה-TPU ואת אשכול ה-GKE:

1. מוחקים את פריסת שרת ה-vLLM ואת הסוד של 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. מוחקים את מאגר הצמתים של TPU v6e ואת אשכול 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. חוזרים לסשן השני בטרמינל ומפסיקים את העברה ליציאה אחרת.

המאמרים הבאים

• מידע נוסף על vLLM TorchTPU
• מידע נוסף על Cloud TPU