‫Google משתמשת בטכנולוגיית AI כדי לתרגם תוכן לשפה המועדפת עליך. בתרגומים כאלו עשויות להיות שגיאות.

הצגת מודלי LLM ביחידות TPU עם כמה מארחים באמצעות Ray ו-TorchTPU

Autopilot Standard

במדריך הזה נסביר איך להפעיל מודל שפה גדול (LLM) מסוג שילוב מומחים (MoE) ‏Gemma 4 26B-A4B-it עם מקביליות של טנזורים ומקביליות של מומחים בפרוסות Cloud TPU בכמה מארחים ב-Google Kubernetes Engine ‏ (GKE) באמצעות Ray Serve LLM,‏ vLLM ו-vLLM-TorchTPU.

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

לפני שמתחילים את המדריך, חשוב לוודא שאתם מכירים את הנושאים הבאים:

רקע

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

יחידות TPU

יחידות Tensor Processing Unit‏ (TPU) הן מאיצים ש-Google פיתחה בהתאמה אישית לעומסי עבודה של למידת מכונה. המדריך הזה תומך גם ב-TPU Trillium (v6e) וגם ב-Ironwood (TPU7x). בפרוסת TPU עם כמה מארחים, כמה מארחי TPU מתקשרים ישירות דרך רשת Inter-Chip Interconnect‏ (ICI) במהירות גבוהה. כך הם מספקים את הזיכרון המצטבר ברוחב פס גבוה (HBM) ואת התקשורת הקולקטיבית עם זמן אחזור קצר שנדרשים כדי להפעיל מודלים שחורגים מקיבולת הזיכרון של שבב יחיד או של מארח TPU יחיד.

‫vLLM ו-TorchTPU

‫vLLM הוא מנוע בקוד פתוח להסקת מסקנות ולפרסום של מודלים גדולים של שפה (LLM). ‫TorchTPU הוא קצה עורפי של PyTorch בקוד פתוח ל-Cloud TPU, שמספק ניהול מכשירים ותקשורת קולקטיבית דרך ICI.

‫vLLM-TorchTPU רושם את הפלאגין של פלטפורמת ה-TPU ב-vLLM ומקמפל את המעברים קדימה של מודל PyTorch ל-Cloud TPU באמצעות TorchTPU והקומפיילר Accelerated Linear Algebra‏ (XLA). הוא תומך במקביליות מבוזרת של טנזורים ובמקביליות מומחים בפרוסות TPU עם כמה מארחים באמצעות Ray.

‫Ray Serve LLM ב-GKE

‫Ray Serve LLM מספק Ingress מסוג HTTP שתואם ל-OpenAI לנקודות הקצה /v1/chat/completions, ‏ /v1/completions ו-/v1/models, ומתזמר פריסות vLLM מבוזרות ב-GKE. כשמפעילים את התוסף Ray on GKE, מפעיל KubeRay וה-webhook של TPU שלו מוסיפים מטא-נתונים של טופולוגיית פרוסות, שמות מארחים של Worker ומשתני תיאום של TorchTPU לכל Pod של Worker ב-TPU עם כמה מארחים.

‫Ray Serve LLM מספק את היכולות הבאות לאירוח מודלים ב-TPU בכמה מארחים:

מיקום פרוסות בהתאם לטופולוגיה שומר פרוסת TPU שלמה עם כמה מארחים לכל רפליקה של מודל כשמציינים טופולוגיית TPU בשדה accelerator_config, ומתזמן וורקר vLLM אחד לכל מכשיר TPU בפרוסה.
ניתוב לפי תחילית מפנה בקשות עם הנחיות מערכת משותפות או היסטוריית שיחות עם כמה שלבים לרפליקה שכבר שומרת במטמון את הבלוקים של מפתח-ערך (KV) של התחילית ב-TPU HBM, כשמציגים בכמה רפליקות של פרוסות TPU.
העברת טוקנים ישירה אוספת אצוות של טוקנים יוצאים בבקשות מקבילות ומעבירה תשובות ישירות מרפליקות של המודל כדי לקצר את הזמן עד לטוקן הראשון (TTFT) ואת זמן האחזור בין טוקנים (ITL).
התאמה לעומס (auto-scaling) של רפליקות על סמך בו-זמניות משנה את קנה המידה של רפליקות של מודלים במרווחים של פרוסות שלמות על סמך בו-זמניות של בקשות מתמשכות.

מטרות

במדריך הזה מוסבר איך:

יוצרים קובץ אימג' של קונטיינר שמוסיף שכבה של Ray Serve מעל קובץ האימג' הבסיסי vLLM-TorchTPU, ודוחפים אותו ל-Artifact Registry.
יוצרים אשכול GKE במצב Autopilot או Standard ב-TPU Trillium (v6e) או Ironwood (TPU7x).
פורסים משאב בהתאמה אישית של RayCluster עם כמה מארחים ומאמתים את הגדרת הטופולוגיה של ה-TPU.
פורסים אפליקציית Ray Serve LLM כדי להציג את המודל google/gemma-4-26B-A4B-it עם מקביליות טנזור ומקביליות מומחים בפרוסת TPU עם כמה מארחים, או פורסים גם את האשכול וגם את האפליקציה באופן הצהרתי באמצעות משאב מותאם אישית של RayService.
שליחת בקשות להשלמת צ'אט ולגילוי מודלים שתואמות ל-OpenAI אל המודל שמוצג.

עלויות

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

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

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

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

לפני שמתחילים, חשוב לוודא שביצעתם את המשימות הבאות:

מפעילים את Artifact Registry API ואת Google Kubernetes Engine API.

הפעלת ממשקי ה-API

כדי להשתמש ב-Google Cloud CLI למשימה הזו, מתקינים ואז מאתחלים את ה-CLI של gcloud. אם כבר התקנתם את ה-CLI של gcloud, הריצו את הפקודה gcloud components update כדי לקבל את הגרסה העדכנית. יכול להיות שבגרסאות קודמות של ה-CLI של gcloud אי אפשר להריץ את הפקודות שבמסמך הזה.
הערה: אם ה-CLI של gcloud כבר מותקן, צריך להגדיר את המאפיין compute/region. אם אתם משתמשים בעיקר באשכולות אזוריים, הגדירו במקום זאת את compute/zone. אם מגדירים מיקום ברירת מחדל, אפשר להימנע משגיאות ב-CLI של gcloud כמו: One of [--zone, --region] must be supplied: Please specify location. יכול להיות שתצטרכו לציין את המיקום בפקודות מסוימות אם המיקום של האשכול שונה מברירת המחדל שהגדרתם.

מוודאים שיש לפרויקט מכסה מספיקה לקיבולת של TPU Trillium (v6e) או Ironwood (TPU7x) באזור שנבחר. מידע נוסף זמין במאמר מכסות של Cloud TPU.
מוודאים שיש לכם חשבון Hugging Face, מאשרים את תנאי הרישיון בדף המודל של google/gemma-4-26B-A4B-it ויוצרים טוקן גישה ל-Hugging Face עם הרשאות Read.
מוודאים שלחשבון שלכם יש את התפקידים הבאים ב-IAM:
- אדמין ב-Kubernetes Engine ‏ (roles/container.admin)
- אדמין ב-Artifact Registry ‏ (roles/artifactregistry.admin)

הכנת הסביבה

במדריך הזה תשתמשו ב-Cloud Shell או בתחנת עבודה מקומית כדי לנהל משאבים שמתארחים ב- Google Cloud. ב-Cloud Shell מותקנים מראש כלי שורת הפקודה שמשמשים במדריך הזה, כולל kubectl,‏ docker,‏ python3 ו-CLI של gcloud.

כדי להגדיר את הסביבה באמצעות Cloud Shell:

במסוף Google Cloud , לוחצים על Activate Cloud Shell במסוףGoogle Cloud כדי להפעיל סשן של Cloud Shell. הפעולה הזו מפעילה סשן בחלונית התחתונה של מסוף Google Cloud .

משכפלים את המאגר kubernetes-engine-samples ועוברים לספריית הדוגמאות של המדריך הזה:

git clone https://github.com/GoogleCloudPlatform/kubernetes-engine-samples.git
cd kubernetes-engine-samples/ai-ml/gke-ray/rayserve/llm/torchtpu

יוצרים ומפעילים סביבה וירטואלית מקומית של Python עם Ray CLI מותקן:
```
python3 -m venv ~/ray-env
source ~/ray-env/bin/activate
pip install "ray[default]==2.58.0"
```
מגדירים את משתני הסביבה לפרויקט ול-cluster ב- Google Cloud :
```
export PROJECT_ID=PROJECT_ID
export CLUSTER_NAME=CLUSTER_NAME
export REGION=REGION
export ZONE=ZONE
export ARTIFACT_REGISTRY=ARTIFACT_REGISTRY
export HF_TOKEN=HF_TOKEN
export IMAGE_URI=${REGION}-docker.pkg.dev/${PROJECT_ID}/${ARTIFACT_REGISTRY}/ray-torchtpu-vllm:latest
gcloud config set project ${PROJECT_ID}
```
מחליפים את מה שכתוב בשדות הבאים:
- ‫PROJECT_ID: מזהה הפרויקט ב- Google Cloud .
- ‫CLUSTER_NAME: השם של אשכול GKE, למשל ray-serve-torchtpu-cluster.
- ‫REGION: האזור שבו נמצאים מאגר Artifact Registry וקיבולת ה-TPU, למשל us-east5.
- ‫ZONE: האזור שבו קיבולת ה-TPU זמינה, למשל us-east5-b. מידע נוסף זמין במאמר אימות הזמינות של TPU ב-GKE.
- ARTIFACT_REGISTRY: השם של מאגר Docker ב-Artifact Registry, למשל ray-torchtpu.
- ‫HF_TOKEN: טוקן הגישה ל-Hugging Face עם הרשאות Read.

יצירת קובץ אימג' של קונטיינר בהתאמה אישית

יוצרים קובץ אימג' של קונטיינר שמתקין את החבילה ray[default,serve] על גבי קובץ האימג' הבסיסי של vLLM-TorchTPU, שכולל את TorchTPU,‏ vLLM ואת הפלאגין vLLM-TorchTPU.

בודקים את Dockerfile בספרייה kubernetes-engine-samples/ai-ml/gke-ray/rayserve/llm/torchtpu:

FROM vllm/vllm-torchtpu:v0.29.0

ENV HF_HOME=/dev/shm/huggingface
ENV RAY_SERVE_DEFAULT_HTTP_HOST=0.0.0.0

# Install wget for KubeRay health probes and pin Ray Serve and transformers
# versions compatible with vLLM 0.29.0 and multi-host TPUs.
RUN apt-get update && apt-get install -y --no-install-recommends wget \
    && rm -rf /var/lib/apt/lists/* \
    && pip install --no-cache-dir \
       "ray[default,serve] @ https://s3-us-west-2.amazonaws.com/ray-wheels/master/6b2b414bb7a5b4e2705ed9f4fdab9787b92f3b08/ray-3.0.0.dev0-cp312-cp312-manylinux2014_x86_64.whl" \
       "transformers>=5.18.0"

כדי להגיש מודל LLM ב-Ray Serve באמצעות TPU עם כמה מארחים, נדרשת גרסת Ray 2.56.0 ואילך ב-TPU Trillium (v6e) וגרסת Ray 2.58.0 ואילך לתזמון לכל צ'יפלט ב-Ironwood (TPU7x).

יוצרים את מאגר Docker ב-Artifact Registry, יוצרים את קובץ האימג' של הקונטיינר ודוחפים אותו ל-Artifact Registry:

gcloud artifacts repositories create ${ARTIFACT_REGISTRY} \
    --repository-format=docker \
    --location=${REGION} \
    --description="Ray and vLLM-TorchTPU container images"

gcloud auth configure-docker ${REGION}-docker.pkg.dev,us-docker.pkg.dev
docker build -t ${IMAGE_URI} .
docker push ${IMAGE_URI}

יצירת אשכול GKE

אתם יכולים להריץ מודלי LLM בפרוסות TPU עם כמה מארחים באשכול GKE Autopilot או באשכול רגיל, כשהתוסף המנוהל Ray operator מופעל.

יוצרים את אשכול GKE ומגדירים את פרטי הכניסה לכלי שורת הפקודה kubectl:
טייס אוטומטי
1. יוצרים אשכול GKE Autopilot אזורי עם תוסף Ray operator מופעל:
  
  gcloud container clusters create-auto ${CLUSTER_NAME} \ --project=${PROJECT_ID} \ --region=${REGION} \ --cluster-version=1.36 \ --enable-ray-operator
  
  כשפורסים את המשאב המותאם אישית RayCluster או RayService בשלב מאוחר יותר, GKE Autopilot מקצה באופן אוטומטי את צומת ה-CPU הראשי ואת פרוסת ה-TPU עם כמה מארחים שתואמים לבקשות המשאבים של Pod ולבוררי הצמתים.
2. מגדירים את כלי שורת הפקודה kubectl כדי לתקשר עם האשכול:
  
  gcloud container clusters get-credentials ${CLUSTER_NAME} \ --project=${PROJECT_ID} \ --region=${REGION}
רגילה
1. יוצרים אשכול GKE Standard אזורי עם מאגר צומתי מעבד (CPU) לאירוח ה-Pod הראשי של Ray, עם הזרמת תמונות ועם GKE Dataplane V2 מופעל:
  
  gcloud container clusters create ${CLUSTER_NAME} \ --project=${PROJECT_ID} \ --zone=${ZONE} \ --cluster-version=1.36 \ --machine-type=n2-standard-8 \ --num-nodes=1 \ --addons=RayOperator \ --enable-image-streaming \ --enable-dataplane-v2
  הערה: אם אתם מנהלים בעצמכם את KubeRay בקוד פתוח במקום להשתמש בתוסף RayOperator של GKE, התקינו את KubeRay ואת KubeRay TPU webhook בגרסה 1.4.0 ואילך כדי להחדיר את משתנה הסביבה TorchTPU.
2. מגדירים את כלי שורת הפקודה kubectl כדי לתקשר עם האשכול:
  
  gcloud container clusters get-credentials ${CLUSTER_NAME} \ --project=${PROJECT_ID} \ --zone=${ZONE}
3. יוצרים מאגר של פרוסות TPU עם מספר מארחים ו-16 מכשירים, עם טופולוגיית 4x4 בארבעה מארחים ב-TPU Trillium (v6e) או טופולוגיית 2x2x2 בשני מארחים ב-Ironwood (TPU7x):
  
  Trillium (v6e)
  gcloud container node-pools create tpu-multi-host-pool \ --project=${PROJECT_ID} \ --cluster=${CLUSTER_NAME} \ --zone=${ZONE} \ --machine-type=ct6e-standard-4t \ --tpu-topology=4x4 \ --num-nodes=4
  
  ‫Ironwood‏ (TPU7x)
  gcloud container node-pools create tpu-multi-host-pool \ --project=${PROJECT_ID} \ --cluster=${CLUSTER_NAME} \ --zone=${ZONE} \ --machine-type=tpu7x-standard-4t \ --tpu-topology=2x2x2 \ --num-nodes=2
  
  טיפ: כדי להשתמש בהזמנה ספציפית של TPU, מוסיפים את הדגלים --reservation-affinity=specific ו---reservation=RESERVATION_NAME לפקודה node-pools create, כאשר RESERVATION_NAME הוא השם של הזמנת ה-TPU. כדי להשתמש במכונות וירטואליות במודל Spot, מוסיפים את הדגל --spot.
יוצרים סוד של Kubernetes עם הטוקן של Hugging Face כדי ש-Ray Pods יוכלו לבצע אימות באמצעות Hugging Face כשמורידים משקלי מודלים:
```
kubectl create secret generic hf-secret \
    --from-literal=hf_api_token=${HF_TOKEN}
```

פריסת המשאב המותאם אישית של RayCluster עם כמה מארחים

פורסים משאב בהתאמה אישית של RayCluster שמקצה Pod ראשי של מעבד (CPU), נפח /dev/shm בזיכרון למאגר האובייקטים של הזיכרון המשותף של Ray וקבוצת וורקרים של TPU עם כמה מארחים.

בודקים את המניפסט של RayCluster לגרסת ה-TPU במאגר לדוגמה:

Trillium (v6e)

apiVersion: ray.io/v1
kind: RayCluster
metadata:
  name: ray-serve-torchtpu
spec:
  headGroupSpec:
    rayStartParams:
      num-cpus: "0"
    template:
      spec:
        containers:
        - name: ray-head
          image: ${IMAGE_URI}
          env:
          - name: HF_TOKEN
            valueFrom:
              secretKeyRef:
                name: hf-secret
                key: hf_api_token
          ports:
          - containerPort: 6379
            name: gcs
          - containerPort: 8265
            name: dashboard
          - containerPort: 10001
            name: client
          - containerPort: 8000
            name: serve
          resources:
            limits:
              cpu: "4"
              memory: "16Gi"
            requests:
              cpu: "4"
              memory: "16Gi"
  workerGroupSpecs:
  - groupName: v6e-16
    replicas: 1
    numOfHosts: 4
    template:
      spec:
        containers:
        - name: ray-worker
          image: ${IMAGE_URI}
          env:
          - name: HF_TOKEN
            valueFrom:
              secretKeyRef:
                name: hf-secret
                key: hf_api_token
          resources:
            limits:
              cpu: "64"
              memory: "200Gi"
              google.com/tpu: "4"
            requests:
              cpu: "64"
              memory: "200Gi"
              google.com/tpu: "4"
          volumeMounts:
          - mountPath: /dev/shm
            name: dshm
        volumes:
        - name: dshm
          emptyDir:
            medium: Memory
            sizeLimit: 100Gi
        nodeSelector:
          cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
          cloud.google.com/gke-tpu-topology: 4x4

‫Ironwood‏ (TPU7x)

apiVersion: ray.io/v1
kind: RayCluster
metadata:
  name: ray-serve-torchtpu
spec:
  headGroupSpec:
    rayStartParams:
      num-cpus: "0"
    template:
      spec:
        containers:
        - name: ray-head
          image: ${IMAGE_URI}
          env:
          - name: HF_TOKEN
            valueFrom:
              secretKeyRef:
                name: hf-secret
                key: hf_api_token
          - name: RAY_TPU_RESOURCE_PER_CHIP
            value: "2"
          ports:
          - containerPort: 6379
            name: gcs
          - containerPort: 8265
            name: dashboard
          - containerPort: 10001
            name: client
          - containerPort: 8000
            name: serve
          resources:
            limits:
              cpu: "4"
              memory: "16Gi"
            requests:
              cpu: "4"
              memory: "16Gi"
  workerGroupSpecs:
  - groupName: tpu7x-16
    replicas: 1
    numOfHosts: 2
    resources:
      # Advertises 8 logical TPU cores (2 chiplets per physical chip) per host.
      TPU: "8"
    template:
      spec:
        containers:
        - name: ray-worker
          image: ${IMAGE_URI}
          env:
          - name: HF_TOKEN
            valueFrom:
              secretKeyRef:
                name: hf-secret
                key: hf_api_token
          - name: RAY_TPU_RESOURCE_PER_CHIP
            value: "2"
          resources:
            limits:
              cpu: "64"
              memory: "250Gi"
              google.com/tpu: "4"
            requests:
              cpu: "64"
              memory: "250Gi"
              google.com/tpu: "4"
          volumeMounts:
          - mountPath: /dev/shm
            name: dshm
        volumes:
        - name: dshm
          emptyDir:
            medium: Memory
            sizeLimit: 128Gi
        nodeSelector:
          cloud.google.com/gke-tpu-accelerator: tpu7x
          cloud.google.com/gke-tpu-topology: 2x2x2

ב-Ironwood (TPU7x), כל שבב פיזי מכיל שני צ'יפלטים, שמספקים שמונה ליבות TPU לוגיות לכל מארח עם ארבעה שבבים. המניפסט מגדיר את הערך של השדה TPU ל-"8" בקטע resources, ומגדיר את משתנה הסביבה RAY_TPU_RESOURCE_PER_CHIP ל-"2" כדי לפרסם ולתזמן שמונה משאבי TPU לוגיים לכל מארח.

מחילים את המניפסט כדי ליצור את המשאב בהתאמה אישית של RayCluster:

Trillium (v6e)

envsubst < ray-cluster.tpu-v6e-multihost.yaml | kubectl apply -f -

‫Ironwood‏ (TPU7x)

envsubst < ray-cluster.tpu-v7x-multihost.yaml | kubectl apply -f -

מחכים עד שה-Pod הראשי וכל ה-Pods של ה-TPU worker יגיעו למצבים Running ו-1/1 READY:

kubectl get pods -l ray.io/cluster=ray-serve-torchtpu

הפלט אמור להיראות כך:

Trillium (v6e)

NAME                                     READY   STATUS    RESTARTS   AGE
ray-serve-torchtpu-head-5cj7g            1/1     Running   0          35s
ray-serve-torchtpu-v6e-16-worker-6x8cg   1/1     Running   0          35s
ray-serve-torchtpu-v6e-16-worker-79j5l   1/1     Running   0          35s
ray-serve-torchtpu-v6e-16-worker-psxbl   1/1     Running   0          35s
ray-serve-torchtpu-v6e-16-worker-rz7dv   1/1     Running   0          35s

‫Ironwood‏ (TPU7x)

NAME                                       READY   STATUS    RESTARTS   AGE
ray-serve-torchtpu-head-w7b2v              1/1     Running   0          2m30s
ray-serve-torchtpu-tpu7x-16-worker-6zlnj   1/1     Running   0          2m30s
ray-serve-torchtpu-tpu7x-16-worker-cb4zz   1/1     Running   0          2m30s

מוודאים שה-webhook של KubeRay TPU הזריק את משתני הסביבה של פרוסת המארחים המרובים, תיאום ה-worker ו-TorchTPU לפודים של ה-worker:

kubectl get pods -l ray.io/node-type=worker \
    -o=jsonpath='{range .items[0].spec.containers[0].env[*]}{.name}={.value}{"\n"}{end}' \
    | grep -E "TPU_WORKER_ID|TPU_WORKER_HOSTNAMES|TPU_NAME|TORCH_TPU_TOPOLOGY|TORCH_TPU_SLICEBUILDER_ADDRESSES"

הפלט אמור להיראות כך:

Trillium (v6e)

TPU_WORKER_HOSTNAMES=v6e-16-0-0.ray-serve-torchtpu-headless,v6e-16-0-1.ray-serve-torchtpu-headless,v6e-16-0-2.ray-serve-torchtpu-headless,v6e-16-0-3.ray-serve-torchtpu-headless
TORCH_TPU_TOPOLOGY=4,4,1
TORCH_TPU_SLICEBUILDER_ADDRESSES=v6e-16-0-0.ray-serve-torchtpu-headless:8471,v6e-16-0-0.ray-serve-torchtpu-headless:8472,v6e-16-0-0.ray-serve-torchtpu-headless:8473,v6e-16-0-0.ray-serve-torchtpu-headless:8474,v6e-16-0-1.ray-serve-torchtpu-headless:8471,v6e-16-0-1.ray-serve-torchtpu-headless:8472,v6e-16-0-1.ray-serve-torchtpu-headless:8473,v6e-16-0-1.ray-serve-torchtpu-headless:8474,v6e-16-0-2.ray-serve-torchtpu-headless:8471,v6e-16-0-2.ray-serve-torchtpu-headless:8472,v6e-16-0-2.ray-serve-torchtpu-headless:8473,v6e-16-0-2.ray-serve-torchtpu-headless:8474,v6e-16-0-3.ray-serve-torchtpu-headless:8471,v6e-16-0-3.ray-serve-torchtpu-headless:8472,v6e-16-0-3.ray-serve-torchtpu-headless:8473,v6e-16-0-3.ray-serve-torchtpu-headless:8474
TPU_WORKER_ID=1
TPU_NAME=v6e-16-0

‫Ironwood‏ (TPU7x)

TPU_WORKER_HOSTNAMES=tpu7x-16-0-0.ray-serve-torchtpu-headless,tpu7x-16-0-1.ray-serve-torchtpu-headless
TORCH_TPU_TOPOLOGY=2,2,2,2
TORCH_TPU_SLICEBUILDER_ADDRESSES=tpu7x-16-0-0.ray-serve-torchtpu-headless:8471,tpu7x-16-0-0.ray-serve-torchtpu-headless:8472,tpu7x-16-0-0.ray-serve-torchtpu-headless:8473,tpu7x-16-0-0.ray-serve-torchtpu-headless:8474,tpu7x-16-0-0.ray-serve-torchtpu-headless:8475,tpu7x-16-0-0.ray-serve-torchtpu-headless:8476,tpu7x-16-0-0.ray-serve-torchtpu-headless:8477,tpu7x-16-0-0.ray-serve-torchtpu-headless:8478,tpu7x-16-0-1.ray-serve-torchtpu-headless:8471,tpu7x-16-0-1.ray-serve-torchtpu-headless:8472,tpu7x-16-0-1.ray-serve-torchtpu-headless:8473,tpu7x-16-0-1.ray-serve-torchtpu-headless:8474,tpu7x-16-0-1.ray-serve-torchtpu-headless:8475,tpu7x-16-0-1.ray-serve-torchtpu-headless:8476,tpu7x-16-0-1.ray-serve-torchtpu-headless:8477,tpu7x-16-0-1.ray-serve-torchtpu-headless:8478
TPU_WORKER_ID=1
TPU_NAME=tpu7x-16-0

אפשר לבדוק את צומתי העובדים של ה-TPU ואת התוויות של ray.io/tpu-* שזוהו באופן אוטומטי בכרטיסייה Cluster ב-Ray Dashboard.

מעבירים את היציאות המקומיות למרכז הבקרה של Ray ביציאה 8265 ולנקודת הקצה של Ray Serve HTTP ביציאה 8000:
```
kubectl port-forward svc/ray-serve-torchtpu-head-svc 8265:8265 8000:8000 >/dev/null 2>&1 &
```

יצירה ופריסה של אפליקציית Ray Serve LLM

פורסים את אפליקציית Ray Serve LLM במשאב המותאם אישית של RayCluster כדי להציג את המודל google/gemma-4-26B-A4B-it בפרוסת ה-TPU עם כמה מארחים.

בודקים את הסקריפט serve_vllm_torchtpu.py במאגר לדוגמה:

from ray.serve.deployment import Application
from ray.serve.llm import (
    LLMConfig,
    LLMServingArgs,
    ModelLoadingConfig,
    build_openai_app,
)


def app(args: dict[str, str]) -> Application:
    llm_config = LLMConfig(
        model_loading_config=ModelLoadingConfig(
            model_id=args.get("model_id", "google/gemma-4-26B-A4B-it"),
        ),
        accelerator_type=args.get("accelerator_type", "TPU-V6E"),
        accelerator_config={
            "kind": "tpu",
            "topology": args.get("topology", "4x4"),
        },
        placement_group_config={
            "bundle_per_worker": {"TPU": 1},
        },
        engine_kwargs={
            "tensor_parallel_size": int(args.get("tensor_parallel_size", "16")),
            "enable_expert_parallel": True,
            "enable_prefix_caching": True,
            "language_model_only": True,
            "disable_hybrid_kv_cache_manager": True,
            "max_model_len": int(args.get("max_model_len", "8192")),
            "max_num_batched_tokens": int(
                args.get("max_num_batched_tokens", "4096")
            ),
        },
        experimental_configs={
            "stream_batching_interval_ms": 10,
        },
        runtime_env={
            "env_vars": {
                "TPU_MULTIHOST_BACKEND": "ray",
                # Execute Gemma 4 routed expert layers on TPU TensorCores instead of SparseCore.
                "USE_MOE_SPARSE_CORE": "0",
                "VLLM_RAY_EXTRA_ENV_VARS_TO_COPY": "USE_MOE_SPARSE_CORE",
            }
        },
    )
    return build_openai_app(LLMServingArgs(llm_configs=[llm_config]))

הסקריפט מגדיר את האובייקט LLMConfig עם ההגדרות הבאות להצגת TPU בכמה מארחים:

השדות accelerator_config ו-placement_group_config שומרים קבוצת מיקום בפרוסת TPU אחת ומקצים ליבת TPU אחת לכל וורקר של vLLM בפרוסה.
השדה engine_kwargs מחלק את המודל לכל 16 ליבות ה-TPU באמצעות מקביליות של טנזורים ומקביליות של מומחים, ומפעיל שמירת קידומות במטמון. הוא גם משבית קלטים מולטי-מודאליים כדי לשמור את ה-HBM של ה-TPU להסקת טקסט, משתמש בפריסת מטמון KV מאוחדת לשכבות הקשב של Gemma 4, ומגביל את אורך ההקשר ל-8,192 אסימונים ואת גודל האצווה למילוי מראש ל-4,096 אסימונים.
השדה experimental_configs אוסף אצוות של טוקנים יוצאים כל 10 אלפיות השנייה כדי לצמצם את התקורה של הסטרימינג בבקשות מקבילות.
השדה runtime_env מפעיל את הקצה העורפי של Ray עם כמה מארחים ל-vLLM-TorchTPU (TPU_MULTIHOST_BACKEND: "ray"), מריץ שכבות מומחה של Gemma 4 MoE ב-TPU TensorCores במקום ב-SparseCore, ומפיץ את ההגדרה הזו לכל 16 ה-worker actors.

פורסים את אפליקציית Ray Serve במשאב המותאם אישית של RayCluster באמצעות הפקודה serve run. הפקודה מעבירה בסטרימינג את יומני הפריסה עד לסיום הטעינה של משקלי המודל, השלמת הידור התרשים של XLA בכל 16 מכשירי ה-TPU והגעה של אפליקציית Serve למצב RUNNING:

Trillium (v6e)

ray job submit \
    --address http://localhost:8265 \
    --working-dir . \
    --runtime-env-json '{"excludes": ["ray-env", ".git", "*.whl"]}' \
    -- serve run --non-blocking serve_vllm_torchtpu:app accelerator_type=TPU-V6E topology=4x4

‫Ironwood‏ (TPU7x)

ray job submit \
    --address http://localhost:8265 \
    --working-dir . \
    --runtime-env-json '{"excludes": ["ray-env", ".git", "*.whl"]}' \
    -- serve run --non-blocking serve_vllm_torchtpu:app accelerator_type=TPU-V7X topology=2x2x2

הפלט אמור להיראות כך:

Trillium (v6e)

(ServeController pid=1637) INFO 2026-10-13 10:23:05,160 controller 1637 -- Deploying new version of Deployment(name='LLMServer:google--gemma-4-26B-A4B-it', app='default') (initial target replicas: 1).
(ServeController pid=1637) INFO 2026-10-13 10:23:05,161 controller 1637 -- Deploying new version of Deployment(name='OpenAiIngress', app='default') (initial target replicas: 1).
(ServeReplica:default:LLMServer:google--gemma-4-26B-A4B-it pid=515, ip=10.72.3.21) INFO 10-13 10:24:06 [__init__.py:318] TPU info: node_name=v6e-16-0 | tpu_type=v6e-16 | worker_id=1 | num_chips=4 | num_cores_per_chip=1
(ServeReplica:default:LLMServer:google--gemma-4-26B-A4B-it pid=515, ip=10.72.3.21) INFO 10-13 10:24:06 [__init__.py:238] Platform plugin torchtpu is activated
(ServeReplica:default:LLMServer:google--gemma-4-26B-A4B-it pid=515, ip=10.72.3.21) (EngineCore pid=1041) INFO 10-13 10:24:08 [core.py:112] Initializing a V1 LLM engine (v0.22.1) with config: model='google/gemma-4-26B-A4B-it', speculative_config=None, tokenizer='google/gemma-4-26B-A4B-it', skip_tokenizer_init=False, tokenizer_mode=auto, revision=None, tokenizer_revision=None, trust_remote_code=False, dtype=torch.bfloat16, max_seq_len=8192, download_dir=None, load_format=auto, tensor_parallel_size=16, pipeline_parallel_size=1, data_parallel_size=1, decode_context_parallel_size=1, dcp_comm_backend=ag_rs, disable_custom_all_reduce=True, quantization=None, quantization_config=None, enforce_eager=False, enable_return_routed_experts=False, kv_cache_dtype=auto, device_config=tpu, ...
(ServeReplica:default:LLMServer:google--gemma-4-26B-A4B-it pid=515, ip=10.72.3.21) (EngineCore pid=1041) (RayWorkerWrapper pid=1017, ip=10.72.14.21) INFO 10-13 10:24:30 [expert_map_manager.py:245] [EP Rank 4/16] Expert parallelism is enabled. Expert placement strategy: linear. Local/global number of experts: 8/128. Experts local to global index map: 0->32, 1->33, 2->34, 3->35, 4->36, 5->37, 6->38, 7->39.
(ServeReplica:default:LLMServer:google--gemma-4-26B-A4B-it pid=515, ip=10.72.3.21) (EngineCore pid=1041) (RayWorkerWrapper pid=856, ip=10.72.6.21) INFO 10-13 10:27:55 [default_loader.py:397] Loading weights took 12.76 seconds
(ServeReplica:default:LLMServer:google--gemma-4-26B-A4B-it pid=515, ip=10.72.3.21) (EngineCore pid=1041) (RayWorkerWrapper pid=856, ip=10.72.6.21) INFO 10-13 10:32:53 [tpu_worker.py:376] Memory statistics | total_hbm_limit_gb=31.25GiB | total_hbm_limit_cap_gb=28.75GiB | total_hbm_used_gb=3.69GiB | kv_cache_headroom_gb=5.0GiB | kv_connector_hbm_reserve_gb=0.0GiB | total_hbm_avail_gb=20.06GiB
(ServeReplica:default:LLMServer:google--gemma-4-26B-A4B-it pid=515, ip=10.72.3.21) (EngineCore pid=1041) INFO 10-13 10:32:54 [kv_cache_utils.py:1733] GPU KV cache size: 600,832 tokens
(ServeReplica:default:LLMServer:google--gemma-4-26B-A4B-it pid=515, ip=10.72.3.21) (EngineCore pid=1041) INFO 10-13 10:32:54 [kv_cache_utils.py:1734] Maximum concurrency for 8,192 tokens per request: 73.34x
(ServeReplica:default:LLMServer:google--gemma-4-26B-A4B-it pid=515, ip=10.72.3.21) INFO 2026-10-13 10:36:59,851 default_LLMServer:google--gemma-4-26B-A4B-it b2grwv9c -- Started vLLM engine.
INFO 2026-10-13 10:37:05,997 serve 1456 -- Application 'default' is ready at http://127.0.0.1:8000/.
2026-10-13 10:37:11,447 SUCC cli.py:66 -- ------------------------------------------
2026-10-13 10:37:11,447 SUCC cli.py:67 -- Job 'raysubmit_94DTwVzR2pvGQVwR' succeeded
2026-10-13 10:37:11,447 SUCC cli.py:68 -- ------------------------------------------

‫Ironwood‏ (TPU7x)

(ServeController pid=2042) INFO 2026-10-13 10:24:38,092 controller 2042 -- Deploying new version of Deployment(name='LLMServer:google--gemma-4-26B-A4B-it', app='default') (initial target replicas: 1).
(ServeController pid=2042) INFO 2026-10-13 10:24:38,093 controller 2042 -- Deploying new version of Deployment(name='OpenAiIngress', app='default') (initial target replicas: 1).
(ServeReplica:default:LLMServer:google--gemma-4-26B-A4B-it pid=316, ip=172.16.7.5) INFO 10-13 10:24:49 [__init__.py:318] TPU info: node_name=tpu7x-16-0 | tpu_type=tpu7x-16 | worker_id=0 | num_chips=8 | num_cores_per_chip=2
(ServeReplica:default:LLMServer:google--gemma-4-26B-A4B-it pid=316, ip=172.16.7.5) INFO 10-13 10:24:49 [__init__.py:238] Platform plugin torchtpu is activated
(ServeReplica:default:LLMServer:google--gemma-4-26B-A4B-it pid=316, ip=172.16.7.5) (EngineCore pid=719) INFO 10-13 10:25:23 [core.py:112] Initializing a V1 LLM engine (v0.22.1) with config: model='google/gemma-4-26B-A4B-it', speculative_config=None, tokenizer='google/gemma-4-26B-A4B-it', skip_tokenizer_init=False, tokenizer_mode=auto, revision=None, tokenizer_revision=None, trust_remote_code=False, dtype=torch.bfloat16, max_seq_len=8192, download_dir=None, load_format=auto, tensor_parallel_size=16, pipeline_parallel_size=1, data_parallel_size=1, decode_context_parallel_size=1, dcp_comm_backend=ag_rs, disable_custom_all_reduce=True, quantization=None, quantization_config=None, enforce_eager=False, enable_return_routed_experts=False, kv_cache_dtype=auto, device_config=tpu, ...
(ServeReplica:default:LLMServer:google--gemma-4-26B-A4B-it pid=316, ip=172.16.7.5) (EngineCore pid=719) (RayWorkerWrapper pid=602, ip=172.16.10.6) INFO 10-13 10:25:44 [expert_map_manager.py:245] [EP Rank 8/16] Expert parallelism is enabled. Expert placement strategy: linear. Local/global number of experts: 8/128. Experts local to global index map: 0->64, 1->65, 2->66, 3->67, 4->68, 5->69, 6->70, 7->71.
(ServeReplica:default:LLMServer:google--gemma-4-26B-A4B-it pid=316, ip=172.16.7.5) (EngineCore pid=719) (RayWorkerWrapper pid=762) INFO 10-13 10:29:06 [default_loader.py:397] Loading weights took 11.69 seconds
(ServeReplica:default:LLMServer:google--gemma-4-26B-A4B-it pid=316, ip=172.16.7.5) (EngineCore pid=719) (RayWorkerWrapper pid=604, ip=172.16.10.6) INFO 10-13 10:34:51 [tpu_worker.py:376] Memory statistics | total_hbm_limit_gb=94.75GiB | total_hbm_limit_cap_gb=87.17GiB | total_hbm_used_gb=3.74GiB | kv_cache_headroom_gb=5.0GiB | kv_connector_hbm_reserve_gb=0.0GiB | total_hbm_avail_gb=78.43GiB
(ServeReplica:default:LLMServer:google--gemma-4-26B-A4B-it pid=316, ip=172.16.7.5) (EngineCore pid=719) INFO 10-13 10:34:58 [kv_cache_utils.py:1733] GPU KV cache size: 2,349,568 tokens
(ServeReplica:default:LLMServer:google--gemma-4-26B-A4B-it pid=316, ip=172.16.7.5) (EngineCore pid=719) INFO 10-13 10:34:58 [kv_cache_utils.py:1734] Maximum concurrency for 8,192 tokens per request: 286.81x
(ServeReplica:default:LLMServer:google--gemma-4-26B-A4B-it pid=316, ip=172.16.7.5) INFO 2026-10-13 10:40:13,586 default_LLMServer:google--gemma-4-26B-A4B-it a15k5cv3 -- Started vLLM engine.
INFO 2026-10-13 10:40:18,415 serve 1953 -- Application 'default' is ready at http://127.0.0.1:8000/.
2026-10-13 10:40:24,372 SUCC cli.py:66 -- ------------------------------------------
2026-10-13 10:40:24,372 SUCC cli.py:67 -- Job 'raysubmit_mzWQDQGDM9QSgcMN' succeeded
2026-10-13 10:40:24,372 SUCC cli.py:68 -- ------------------------------------------

שליחת שאילתה למודל שמוצג

אחרי שאפליקציית Ray Serve מגיעה למצב RUNNING, מאמתים את נקודת הקצה של המודל הרשום, שולחים בקשת /v1/chat/completions שתואמת ל-OpenAI דרך יציאה 8000, ואם רוצים, מודדים את קצב הנתונים והזמן שלוקח להם להגיע.

מוודאים שהמודל רשום בנקודת הקצה (endpoint) של /v1/models:

curl http://localhost:8000/v1/models

הפלט אמור להיראות כך:

{"data":[{"id":"google/gemma-4-26B-A4B-it","object":"model","owned_by":"organization-owner","permission":[],"metadata":{"model_id":"google/gemma-4-26B-A4B-it","input_modality":"text","max_request_context_length":8192}}],"object":"list"}

שולחים הנחיה להשלמת צ'אט לנקודת הקצה (endpoint) של /v1/chat/completions:

curl -X POST http://localhost:8000/v1/chat/completions \
    -H "Content-Type: application/json" \
    -d '{
        "model": "google/gemma-4-26B-A4B-it",
        "messages": [
            {"role": "system", "content": "You are a helpful AI assistant."},
            {"role": "user", "content": "Explain in one sentence how tensor parallelism and expert parallelism accelerate Mixture-of-Experts LLM inference on TPUs."}
        ],
        "max_tokens": 128,
        "temperature": 0.0
    }'

הפלט ב-Ironwood (TPU7x) דומה לפלט הבא:

{"id":"chatcmpl-ec30e61c-011d-49cd-989a-9ae9f59448a8","object":"chat.completion","created":1791887860,"model":"google/gemma-4-26B-A4B-it","choices":[{"index":0,"message":{"role":"assistant","content":"Tensor parallelism accelerates inference by partitioning large weight matrices across TPU cores to reduce per-device memory and computation, while expert parallelism distributes the specialized MoE layers across different devices to enable massive model scaling and parallelize the processing of different tokens through different experts.","refusal":null,"annotations":null,"audio":null,"function_call":null,"tool_calls":[],"reasoning":null},"logprobs":null,"finish_reason":"stop","stop_reason":106,"token_ids":null,"routed_experts":null}],"service_tier":null,"system_fingerprint":"vllm-0.22.1-tp16-ep-aaf76e1a","usage":{"prompt_tokens":49,"total_tokens":100,"completion_tokens":51,"prompt_tokens_details":null},"prompt_logprobs":null,"prompt_token_ids":null,"prompt_text":null,"kv_transfer_params":null}

אופציונלי: בודקים את קצב העיבוד והזמן שלוקח להציג את התשובות באמצעות הפקודה vllm bench serve. הפקודה הבאה מריצה את לקוח מחולל העומסים של HTTP‏ vllm bench serve מ-Pod של Ray TPU worker, ושולחת 64 בקשות בו-זמנית ל-Proxy ל-HTTP של Ray Serve ביציאה 8000:

WORKER_POD=$(kubectl get pods -l ray.io/node-type=worker -o jsonpath='{.items[0].metadata.name}')
kubectl exec -it ${WORKER_POD} -- vllm bench serve \
    --backend openai-chat \
    --endpoint /v1/chat/completions \
    --model google/gemma-4-26B-A4B-it \
    --dataset-name random \
    --random-input-len 1024 \
    --random-output-len 512 \
    --max-concurrency 32 \
    --num-prompts 64 \
    --num-warmups 4 \
    --temperature 0.0 \
    --ignore-eos \
    --host ray-serve-torchtpu-head-svc \
    --port 8000

הפלט אמור להיראות כך:

Trillium (v6e)

============ Serving Benchmark Result ============
Successful requests:                     64
Failed requests:                         0
Maximum request concurrency:             32
Benchmark duration (s):                  8.41
Total input tokens:                      66392
Total generated tokens:                  32768
Request throughput (req/s):              7.61
Output token throughput (tok/s):         3894.26
Peak output token throughput (tok/s):    4409.00
Peak concurrent requests:                64.00
Total token throughput (tok/s):          11784.50
---------------Time to First Token----------------
Mean TTFT (ms):                          414.57
Median TTFT (ms):                        348.94
P99 TTFT (ms):                           640.26
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          7.37
Median TPOT (ms):                        7.29
P99 TPOT (ms):                           8.24
---------------Inter-token Latency----------------
Mean ITL (ms):                           7.37
Median ITL (ms):                         0.00
P99 ITL (ms):                            25.94
==================================================

‫Ironwood‏ (TPU7x)

============ Serving Benchmark Result ============
Successful requests:                     64
Failed requests:                         0
Maximum request concurrency:             32
Benchmark duration (s):                  9.65
Total input tokens:                      66392
Total generated tokens:                  32768
Request throughput (req/s):              6.64
Output token throughput (tok/s):         3397.33
Peak output token throughput (tok/s):    4574.00
Peak concurrent requests:                64.00
Total token throughput (tok/s):          10280.73
---------------Time to First Token----------------
Mean TTFT (ms):                          654.11
Median TTFT (ms):                        522.30
P99 TTFT (ms):                           1288.23
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          8.13
Median TPOT (ms):                        8.08
P99 TPOT (ms):                           9.69
---------------Inter-token Latency----------------
Mean ITL (ms):                           8.12
Median ITL (ms):                         9.31
P99 ITL (ms):                            26.85
==================================================

פריסה באופן הצהרתי באמצעות משאב בהתאמה אישית של RayService

בפריסות בסביבת ייצור, אתם יכולים לנהל גם את אשכול ה-Ray עם כמה מארחים וגם את אפליקציית ה-LLM של Ray Serve כמשאב מותאם אישית הצהרתי יחיד של RayService, במקום לשלוח משימה באופן ידני. מכיוון שה-LLM builder של Ray Serve מקבל את הגדרות המודל ישירות בהגדרות של Serve, אין צורך בסקריפט נפרד של Python. ‫KubeRay מקצה את האשכול, פורס את אפליקציית Serve, עוקב אחרי תקינות הרפליקה ויוצר שירות ClusterIP של Kubernetes ביציאה 8000 כשהאפליקציה מוכנה לשרת תנועה.

אם המשאב בהתאמה אישית של RayCluster מהקטע הקודם עדיין פועל, צריך להפסיק את העברה ליציאה אחרת ולמחוק את המשאב בהתאמה אישית של RayCluster כדי לשחרר את פרוסת ה-TPU:
```
pkill -f "kubectl port-forward" || true
kubectl delete raycluster ray-serve-torchtpu --ignore-not-found
```

בודקים את המניפסט של RayService לגרסת ה-TPU במאגר לדוגמה:

Trillium (v6e)

apiVersion: ray.io/v1
kind: RayService
metadata:
  name: ray-serve-torchtpu
spec:
  serveConfigV2: |
    applications:
    - name: default
      route_prefix: /
      import_path: ray.serve.llm:build_openai_app
      args:
        llm_configs:
        - model_loading_config:
            model_id: google/gemma-4-26B-A4B-it
          accelerator_type: TPU-V6E
          accelerator_config:
            kind: tpu
            topology: 4x4
          placement_group_config:
            bundle_per_worker:
              TPU: 1
          engine_kwargs:
            tensor_parallel_size: 16
            enable_expert_parallel: true
            enable_prefix_caching: true
            language_model_only: true
            disable_hybrid_kv_cache_manager: true
            max_model_len: 8192
            max_num_batched_tokens: 4096
          experimental_configs:
            stream_batching_interval_ms: 10
          runtime_env:
            env_vars:
              TPU_MULTIHOST_BACKEND: "ray"
              # Execute Gemma 4 routed expert layers on TPU TensorCores instead of SparseCore.
              USE_MOE_SPARSE_CORE: "0"
              VLLM_RAY_EXTRA_ENV_VARS_TO_COPY: "USE_MOE_SPARSE_CORE"
  rayClusterConfig:
    headGroupSpec:
      rayStartParams:
        num-cpus: "0"
      template:
        spec:
          containers:
          - name: ray-head
            image: ${IMAGE_URI}
            env:
            - name: HF_TOKEN
              valueFrom:
                secretKeyRef:
                  name: hf-secret
                  key: hf_api_token
            ports:
            - containerPort: 6379
              name: gcs
            - containerPort: 8265
              name: dashboard
            - containerPort: 10001
              name: client
            - containerPort: 8000
              name: serve
            resources:
              limits:
                cpu: "4"
                memory: "16Gi"
              requests:
                cpu: "4"
                memory: "16Gi"
    workerGroupSpecs:
    - groupName: v6e-16
      replicas: 1
      numOfHosts: 4
      template:
        spec:
          containers:
          - name: ray-worker
            image: ${IMAGE_URI}
            env:
            - name: HF_TOKEN
              valueFrom:
                secretKeyRef:
                  name: hf-secret
                  key: hf_api_token
            resources:
              limits:
                cpu: "64"
                memory: "200Gi"
                google.com/tpu: "4"
              requests:
                cpu: "64"
                memory: "200Gi"
                google.com/tpu: "4"
            volumeMounts:
            - mountPath: /dev/shm
              name: dshm
          volumes:
          - name: dshm
            emptyDir:
              medium: Memory
              sizeLimit: 100Gi
          nodeSelector:
            cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
            cloud.google.com/gke-tpu-topology: 4x4

‫Ironwood‏ (TPU7x)

apiVersion: ray.io/v1
kind: RayService
metadata:
  name: ray-serve-torchtpu
spec:
  serveConfigV2: |
    applications:
    - name: default
      route_prefix: /
      import_path: ray.serve.llm:build_openai_app
      args:
        llm_configs:
        - model_loading_config:
            model_id: google/gemma-4-26B-A4B-it
          accelerator_type: TPU-V7X
          accelerator_config:
            kind: tpu
            topology: 2x2x2
          placement_group_config:
            bundle_per_worker:
              TPU: 1
          engine_kwargs:
            tensor_parallel_size: 16
            enable_expert_parallel: true
            enable_prefix_caching: true
            language_model_only: true
            disable_hybrid_kv_cache_manager: true
            max_model_len: 8192
            max_num_batched_tokens: 4096
          experimental_configs:
            stream_batching_interval_ms: 10
          runtime_env:
            env_vars:
              TPU_MULTIHOST_BACKEND: "ray"
              # Execute Gemma 4 routed expert layers on TPU TensorCores instead of SparseCore.
              USE_MOE_SPARSE_CORE: "0"
              VLLM_RAY_EXTRA_ENV_VARS_TO_COPY: "USE_MOE_SPARSE_CORE"
  rayClusterConfig:
    headGroupSpec:
      rayStartParams:
        num-cpus: "0"
      template:
        spec:
          containers:
          - name: ray-head
            image: ${IMAGE_URI}
            env:
            - name: HF_TOKEN
              valueFrom:
                secretKeyRef:
                  name: hf-secret
                  key: hf_api_token
            - name: RAY_TPU_RESOURCE_PER_CHIP
              value: "2"
            ports:
            - containerPort: 6379
              name: gcs
            - containerPort: 8265
              name: dashboard
            - containerPort: 10001
              name: client
            - containerPort: 8000
              name: serve
            resources:
              limits:
                cpu: "4"
                memory: "16Gi"
              requests:
                cpu: "4"
                memory: "16Gi"
    workerGroupSpecs:
    - groupName: tpu7x-16
      replicas: 1
      numOfHosts: 2
      resources:
        # Advertises 8 logical TPU cores (2 chiplets per physical chip) per host.
        TPU: "8"
      template:
        spec:
          containers:
          - name: ray-worker
            image: ${IMAGE_URI}
            env:
            - name: HF_TOKEN
              valueFrom:
                secretKeyRef:
                  name: hf-secret
                  key: hf_api_token
            - name: RAY_TPU_RESOURCE_PER_CHIP
              value: "2"
            resources:
              limits:
                cpu: "64"
                memory: "250Gi"
                google.com/tpu: "4"
              requests:
                cpu: "64"
                memory: "250Gi"
                google.com/tpu: "4"
            volumeMounts:
            - mountPath: /dev/shm
              name: dshm
          volumes:
          - name: dshm
            emptyDir:
              medium: Memory
              sizeLimit: 128Gi
          nodeSelector:
            cloud.google.com/gke-tpu-accelerator: tpu7x
            cloud.google.com/gke-tpu-topology: 2x2x2

מחילים את המניפסט של RayService:

Trillium (v6e)

envsubst < ray-service.tpu-v6e-multihost.yaml | kubectl apply -f -

‫Ironwood‏ (TPU7x)

envsubst < ray-service.tpu-v7x-multihost.yaml | kubectl apply -f -

מחכים עד שהמשאב המותאם אישית של RayService יסיים לטעון את משקלי המודל ולקמפל את תרשימי ה-XLA בפרוסת ה-TPU, ויגיע למצב Running:

kubectl get rayservice ray-serve-torchtpu
kubectl get svc ray-serve-torchtpu-serve-svc

הפלט אמור להיראות כך:

NAME                 SERVICE STATUS   NUM SERVE ENDPOINTS
ray-serve-torchtpu   Running          3
NAME                           TYPE        CLUSTER-IP       EXTERNAL-IP   PORT(S)    AGE
ray-serve-torchtpu-serve-svc   ClusterIP   34.118.239.213   <none>        8000/TCP   29s

מעבירים את היציאה המקומית 8000 לשירות ray-serve-torchtpu-serve-svc ושולחים שאילתה למודל באמצעות הפקודות curl בקטע שליחת שאילתה למודל שמוצג:
```
kubectl port-forward svc/ray-serve-torchtpu-serve-svc 8000:8000 >/dev/null 2>&1 &
```

מעקב אחרי ביצועי המודל

כדי לעקוב אחרי השימוש ב-TPU TensorCore, ב-HBM ובמדדים של אשכול Ray ב-Cloud Monitoring, קראו את המאמרים איסוף והצגה של יומנים ומדדים של אשכולות Ray ב-GKE והצגת מדדי TPU במרכז הבקרה של Ray.

גישה ללוח הבקרה של Ray

כדי לבדוק את הסטטוס של פריסות Ray Serve, שחקני הרפליקה, תוויות הצמתים וזמני האחזור של השאילתות בלוח הבקרה המובנה של Ray:

אם העברה ליציאה אחרת לשירות הראשי של Ray עדיין לא פועלת, מפעילים אותה ברקע:
```
kubectl port-forward svc/ray-serve-torchtpu-head-svc 8265:8265 8000:8000 >/dev/null 2>&1 &
```
ב-Cloud Shell, לוחצים על Web Preview ובוחרים באפשרות Preview on port 8265, או פותחים את http://localhost:8265 בדפדפן המקומי.
לוחצים על הכרטיסיות Serve,‏ Cluster ו-Metrics כדי לראות את פריסות ה-vLLM, את תקינות הרפליקה, את תוויות הצומת ray.io/tpu-* שזוהו באופן אוטומטי, את מדדי זמן האחזור של הבקשות ואת התרשים TPU Utilization by Node.

הסרת המשאבים

כדי להימנע מחיובים בחשבון Google Cloud על המשאבים שבהם השתמשתם במדריך הזה, אתם צריכים למחוק את המשאב המותאם אישית של RayService או RayCluster, את אשכול GKE ואת המאגר ב-Artifact Registry:

עוצרים את תהליך העברת היציאות ברקע ומוחקים את המשאבים בהתאמה אישית של Ray ואת הסוד של Kubernetes:

pkill -f "kubectl port-forward" || true
kubectl delete rayservice ray-serve-torchtpu --ignore-not-found
kubectl delete raycluster ray-serve-torchtpu --ignore-not-found
kubectl delete secret hf-secret --ignore-not-found

מוחקים את אשכול ה-GKE:

טייס אוטומטי

gcloud container clusters delete ${CLUSTER_NAME} \
    --region=${REGION}

רגילה

gcloud container clusters delete ${CLUSTER_NAME} \
    --zone=${ZONE}

מוחקים את המאגר ב-Artifact Registry:

gcloud artifacts repositories delete ${ARTIFACT_REGISTRY} \
    --location=${REGION}

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

מידע נוסף על PyTorch ב-Cloud TPU
למסמכי התיעוד של TorchTPU
איך מכווננים מודלים ב-TPU עם כמה מארחים באמצעות Ray ו-TorchTPU
איך מפעילים מודלים פתוחים של Gemma באמצעות יחידות TPU עם כמה מארחים ב-GKE עם Ray
איך מנהלים פריסות של Ray Serve באופן הצהרתי באמצעות משאב בהתאמה אישית של KubeRay RayService.
מידע על יחידות TPU ב-GKE
למסמכי התיעוד של Ray Serve LLM

הצגת מודלי LLM ביחידות TPU עם כמה מארחים באמצעות Ray ו-TorchTPU קל לארגן דפים בעזרת אוספים אפשר לשמור ולסווג תוכן על סמך ההעדפות שלך.

רקע

יחידות TPU

‫vLLM ו-TorchTPU

‫Ray Serve LLM ב-GKE

מטרות

עלויות

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

הכנת הסביבה

יצירת קובץ אימג' של קונטיינר בהתאמה אישית

יצירת אשכול GKE

טייס אוטומטי

רגילה

Trillium (v6e)

‫Ironwood‏ (TPU7x)

פריסת המשאב המותאם אישית של RayCluster עם כמה מארחים

Trillium (v6e)

‫Ironwood‏ (TPU7x)

Trillium (v6e)

‫Ironwood‏ (TPU7x)

Trillium (v6e)

‫Ironwood‏ (TPU7x)

Trillium (v6e)

‫Ironwood‏ (TPU7x)

יצירה ופריסה של אפליקציית Ray Serve LLM

Trillium (v6e)

‫Ironwood‏ (TPU7x)

Trillium (v6e)

‫Ironwood‏ (TPU7x)

שליחת שאילתה למודל שמוצג

Trillium (v6e)

‫Ironwood‏ (TPU7x)

פריסה באופן הצהרתי באמצעות משאב בהתאמה אישית של RayService

Trillium (v6e)

‫Ironwood‏ (TPU7x)

Trillium (v6e)

‫Ironwood‏ (TPU7x)

מעקב אחרי ביצועי המודל

גישה ללוח הבקרה של Ray

הסרת המשאבים

טייס אוטומטי

רגילה

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

הצגת מודלי LLM ביחידות TPU עם כמה מארחים באמצעות Ray ו-TorchTPU