אימון מודל LLM באמצעות JAX,‏ Ray Train ו-TPU Trillium ב-GKE

טייס אוטומטי רגילה

במדריך הזה נסביר איך לאמן את מודל השפה הגדול (LLM) ‏Llama 3 8B ב-Google Kubernetes Engine ‏ (GKE) באמצעות MaxText,‏ Ray Train ויחידות TPU.

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

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

רקע

השילוב של GKE,‏ KubeRay,‏ MaxText ו-TPU מספק פלטפורמה חזקה וניתנת להרחבה לאימון מודלים בקנה מידה גדול. בקטע הזה מתוארות הטכנולוגיות העיקריות שמופיעות במדריך הזה:

JAX

‫JAX היא ספריית Python לחישוב מערכים ולהמרת תוכניות שמותאמת למאיצים, ומיועדת לחישובים מספריים עם ביצועים גבוהים וללמידת מכונה בקנה מידה גדול.

‫JAX מספקת מערכת ניתנת להרחבה להמרת פונקציות מספריות כמו jax.grad,‏ jax.jit ו-jax.vmap, תוך שימוש בקומפיילר XLA ליצירת קוד שעבר אופטימיזציה גבוהה וניתן להרחבה ביעילות במאיצים כמו GPU ו-TPU. היתרון העיקרי של JAX הוא הקומפוזביליות שלו, שמאפשרת למשתמשים לשלב את הטרנספורמציות האלה כדי ליצור תוכניות מספריות מורכבות עם ביצועים גבוהים להרצה מבוזרת.

MaxText

‫MaxText הוא מודל שפה גדול (LLM) בקוד פתוח עם ביצועים גבוהים, שנועד להתאמה אישית ולשינוי קנה מידה. ‫MaxText מבוסס על JAX ומותאם להפעלה יעילה ב-Cloud TPU וב-GPU.

TPUs

יחידות לעיבוד טנסורים (TPU) הן מאיצים שנוצרו על ידי Google בהתאמה אישית כדי לבצע אופטימיזציה של עומסי עבודה של למידת מכונה. בניגוד למעבדים מרכזיים לשימוש כללי או למעבדים גרפיים לעיבוד מקבילי, יחידות TPU הן מאוד ייעודיות לחישובים של מטריצות וטנסורים גדולים, שהם הבסיס ללמידה עמוקה, ולכן הן יעילות במשימה הספציפית הזו. היתרון העיקרי של מעבדי TPU הוא הביצועים בקנה מידה נרחב.

במדריך הזה נעשה שימוש ב-TPU Trillium, שהוא הדור השישי של TPUs. מידע נוסף זמין במאמר היתרונות של שימוש ב-TPU Trillium.

KubeRay

‫KubeRay הוא אופרטור של Kubernetes שמספק דרך מאוחדת לפריסה, לניהול ולניטור של אפליקציות Ray ב-Kubernetes. אפשר להתקין את האופרטור KubeRay ולנהל אותו באמצעות התוסף Ray on GKE. זו הדרך המומלצת לפרוס ולנהל אשכולות Ray ב-GKE.

מטרות

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

1. מגדירים אשכול GKE עם מאגר צומתי TPU בכמה מארחים.
2. מגדירים את KubeRay לניהול סביבת האימון המבוזרת.
3. יצירת קובץ אימג' של Docker בהתאמה אישית שמכיל תלות ב-MaxText, ב-Ray וב-JAX.
4. יוצרים סקריפט אימון ב-Python שמשתמש ב-JaxTrainer של Ray Train כדי לתזמן את לולאת האימון של MaxText בפריסת ה-TPU.
5. מגדירים RayCluster משאב בהתאמה אישית כדי להקצות את צומתי ה-head וה-worker עם משאבי ה-TPU הנדרשים.
6. שולחים את משימת האימון אל RayCluster ועוקבים אחרי ההתקדמות שלה.
7. שימוש ב-Cloud Storage לאחסון של נקודות ביקורת של מודלים.

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

נכנסים לחשבון Google Cloud . אנחנו ממליצים למשתמשים חדשים ב- Google Cloud ליצור חשבון כדי שיוכלו להעריך את הביצועים של המוצרים שלנו בתרחישים מהעולם האמיתי. לקוחות חדשים מקבלים בחינם גם קרדיט בשווי 300$ להרצה, לבדיקה ולפריסה של עומסי העבודה.

התקינו את ה-CLI של Google Cloud.

הערה: אם התקנתם את ה-CLI של gcloud, השתמשו בפקודה gcloud components update כדי לבדוק אם הגרסה העדכנית מותקנת.

אם אתם משתמשים בספק זהויות חיצוני (IdP), קודם אתם צריכים להיכנס ל-CLI של gcloud באמצעות המאגר המאוחד לניהול זהויות.

כדי לאתחל את ה-CLI של gcloud, הריצו את הפקודה הבאה:

gcloud init

יוצרים או בוחרים Google Cloud פרויקט.

תפקידים שנדרשים כדי לבחור או ליצור פרויקט

• בחירת פרויקט: כדי לבחור פרויקט לא נדרש תפקיד IAM ספציפי – אתם יכולים לבחור כל פרויקט שקיבלתם בו תפקיד.
• יצירת פרויקט: כדי ליצור פרויקט, צריך את התפקיד Project Creator (roles/resourcemanager.projectCreator), שכולל את ההרשאה resourcemanager.projects.create. איך מקצים תפקידים

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

• יוצרים Google Cloud פרויקט:

  gcloud projects create PROJECT_ID

  מחליפים את PROJECT_ID בשם של פרויקט Google Cloud שיוצרים.

• בוחרים את הפרויקט שיצרתם: Google Cloud

  gcloud config set project PROJECT_ID

  מחליפים את PROJECT_ID בשם הפרויקט ב- Google Cloud .

מוודאים שהחיוב מופעל בפרויקט ב- Google Cloud .

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

תפקידים שנדרשים להפעלת ממשקי API

כדי להפעיל ממשקי API, צריך את ההרשאה serviceusage.services.enable. אם יצרתם את הפרויקט, סביר להניח שכבר יש לכם את ההרשאה הזו דרך התפקיד 'בעלים' (roles/owner). אחרת, תוכלו לקבל את ההרשאה הזו דרך התפקיד 'אדמין בממשק Service Usage' (roles/serviceusage.serviceUsageAdmin). איך מקצים תפקידים

gcloud services enable container.googleapis.com cloudbuild.googleapis.com

התקינו את ה-CLI של Google Cloud.

הערה: אם התקנתם את ה-CLI של gcloud, השתמשו בפקודה gcloud components update כדי לבדוק אם הגרסה העדכנית מותקנת.

אם אתם משתמשים בספק זהויות חיצוני (IdP), קודם אתם צריכים להיכנס ל-CLI של gcloud באמצעות המאגר המאוחד לניהול זהויות.

כדי לאתחל את ה-CLI של gcloud, הריצו את הפקודה הבאה:

gcloud init

יוצרים או בוחרים Google Cloud פרויקט.

תפקידים שנדרשים כדי לבחור או ליצור פרויקט

• בחירת פרויקט: כדי לבחור פרויקט לא נדרש תפקיד IAM ספציפי – אתם יכולים לבחור כל פרויקט שקיבלתם בו תפקיד.
• יצירת פרויקט: כדי ליצור פרויקט, צריך את התפקיד Project Creator (roles/resourcemanager.projectCreator), שכולל את ההרשאה resourcemanager.projects.create. איך מקצים תפקידים

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

• יוצרים Google Cloud פרויקט:

  gcloud projects create PROJECT_ID

  מחליפים את PROJECT_ID בשם של פרויקט Google Cloud שיוצרים.

• בוחרים את הפרויקט שיצרתם: Google Cloud

  gcloud config set project PROJECT_ID

  מחליפים את PROJECT_ID בשם הפרויקט ב- Google Cloud .

מוודאים שהחיוב מופעל בפרויקט ב- Google Cloud .

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

תפקידים שנדרשים להפעלת ממשקי API

כדי להפעיל ממשקי API, צריך את ההרשאה serviceusage.services.enable. אם יצרתם את הפרויקט, סביר להניח שכבר יש לכם את ההרשאה הזו דרך התפקיד 'בעלים' (roles/owner). אחרת, תוכלו לקבל את ההרשאה הזו דרך התפקיד 'אדמין בממשק Service Usage' (roles/serviceusage.serviceUsageAdmin). איך מקצים תפקידים

gcloud services enable container.googleapis.com cloudbuild.googleapis.com

• נותנים תפקידים לחשבון המשתמש. מריצים את הפקודה הבאה לכל אחד מהתפקידים הבאים ב-IAM: roles/container.admin, roles/iam.serviceAccountAdmin, roles/cloudbuild.builds.editor

  gcloud projects add-iam-policy-binding PROJECT_ID --member="user:USER_IDENTIFIER" --role=ROLE

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

  • ‫PROJECT_ID: מזהה הפרויקט.
  • ‫USER_IDENTIFIER: המזהה של חשבון המשתמש . לדוגמה, [e-mail].
  • ‫ROLE: תפקיד ה-IAM שאתם מקצים לחשבון המשתמש.

• במדריך הזה נעשה שימוש ב-TPU Trillium ‏ (v6e), לכן צריך לבחור אזור או אזור משנה שבהם הוא זמין. מידע נוסף זמין במאמר בנושא מכסות של Cloud TPU.

הכנת הסביבה

במדריך הזה משתמשים ב-Cloud Shell. ב-Cloud Shell מותקנים מראש כלי שורת הפקודה gcloud, helm ו-kubectl שבהם נשתמש במדריך הזה.

1. נכנסים למסוףGoogle Cloud .

2. בחלק העליון של חלון המסוף Google Cloud , לוחצים על הלחצן Activate Cloud Shell .

   ב Google Cloud מסוף ייפתח סשן של Cloud Shell בתוך מסגרת חדשה ותופיע הנחיה של שורת הפקודה.

3. יוצרים ומפעילים סביבה וירטואלית של Python:

   python3 -m venv ray-env
   source ray-env/bin/activate
   

4. מתקינים את Ray CLI ואת יחסי התלות האחרים:

   pip install "ray[default]==2.49.1"
   

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

   export PROJECT_ID=$(gcloud config get project)
   export PROJECT_NUMBER=$(gcloud projects describe ${PROJECT_ID} --format="value(projectNumber)")
   export GS_BUCKET=GS_BUCKET
   export KSA_NAME=KSA_NAME
   export NAMESPACE=default
   export CLUSTER_NAME=CLUSTER_NAME
   export REGION=REGION
   export ZONE=ZONE
   export ARTIFACT_REGISTRY=ARTIFACT_REGISTRY
   

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

   • ‫GS_BUCKET: שם הקטגוריה ב-Cloud Storage.
   • ‫KSA_NAME: השם של חשבון השירות של Kubernetes.
   • ‫CLUSTER_NAME: השם של האשכול החדש.
   • ‫REGION: האזור שבו קיבולת TPU Trillium זמינה.
   • ‫ZONE: האזור שבו קיבולת TPU Trillium זמינה. מידע נוסף זמין במאמר זמינות של TPU ב-GKE.
   • ‫ARTIFACT_REGISTRY: השם של מאגר Artifact Registry.

יצירת אשכול GKE

אתם יכולים להגדיר את KubeRay ב-TPU באשכול GKE Autopilot או באשכול רגיל. מומלץ להשתמש באשכול Autopilot כדי ליהנות מחוויית Kubernetes מנוהלת באופן מלא. כדי לבחור את מצב הפעולה של GKE שהכי מתאים לעומסי העבודה שלכם, אפשר לעיין במאמר מידע על מצבי הפעולה של GKE.

טייס אוטומטי

1. ב-Cloud Shell, מריצים את הפקודה הבאה:

   gcloud container clusters create-auto $CLUSTER_NAME \
       --enable-ray-operator \
       --machine-type=n1-standard-16 \
       --location=$REGION
   

2. כדי לתקשר עם האשכול, מגדירים את kubectl :

   gcloud container clusters get-credentials CLUSTER_NAME \
       --location=$ZONE
   

רגילה

1. ב-Cloud Shell, יוצרים אשכול Standard שמופעל בו התוסף Ray operator באמצעות הפקודה הבאה:

   gcloud container clusters create $CLUSTER_NAME \
       --addons=RayOperator \
       --addons GcsFuseCsiDriver \
       --machine-type=n1-standard-16 \
       --workload-pool=$PROJECT_ID.svc.id.goog \
       --location=$ZONE
   

   הפקודה הזו גם מפעילה את GcsFuseCsiDriver, שמאפשר ל-Pods לטעון קטגוריות של Cloud Storage כמערכות קבצים מקומיות. יצירת האשכול עשויה להימשך כמה דקות.

2. כדי לתקשר עם האשכול, מגדירים את kubectl:

   gcloud container clusters get-credentials CLUSTER_NAME \
       --location=LOCATION
   

3. יוצרים מאגר צמתים של פרוסת TPU עם כמה מארחים:

   gcloud container node-pools create v6e-16 \
       --location=$ZONE \
       --cluster=$CLUSTER_NAME \
       --machine-type=ct6e-standard-4t \
       --threads-per-core=1 \
       --tpu-topology=4x4 \
       --num-nodes=4
   

‫GKE מקצה מאגר צמתים שמורכב מ-4 מכונות וירטואליות של TPU Trillium (v6e), שמוגדרות יחד כפרוסת TPU עם כמה מארחים, עם טופולוגיה של 4x4, שמוכנה לעומסי עבודה של אימון מבוזר.

האופרטור של Ray באשכול GKE מתקין באופן אוטומטי את KubeRay ואת ה-webhook של KubeRay TPU באשכול.

הגדרת קטגוריה של Cloud Storage וחשבון שירות

1. יוצרים קטגוריה של Cloud Storage לנקודות ביקורת משותפות בין צמתי ה-TPU עם כמה מארחים.

   gsutil mb -p ${PROJECT_ID} -c STANDARD -l ${REGION} gs://${GS_BUCKET}
   

2. כדי להפעיל גישה לקטגוריה של Cloud Storage, יוצרים חשבון שירות של Kubernetes:

   kubectl create serviceaccount ${KSA_NAME} --namespace ${NAMESPACE}
   

3. כדי להעניק גישה לקטגוריה של Cloud Storage, מוסיפים לחשבון השירות את קישורי המדיניות הנדרשים ב-IAM:

   gcloud storage buckets add-iam-policy-binding gs://${GS_BUCKET} \
       --member "principal://iam.googleapis.com/projects/${PROJECT_NUMBER}/locations/global/workloadIdentityPools/${PROJECT_ID}.svc.id.goog/subject/ns/${NAMESPACE}/sa/${KSA_NAME}" \
       --role "roles/storage.objectUser"
   

יצירת סקריפט האימון

הסקריפט הבא משתמש ב-JaxTrainer של Ray Train כדי להריץ משימת אימון מבוזרת של MaxText. הסקריפט מגדיר את סביבת האימון עבור מאגר צמתים של פרוסת TPU מרובת-מארחים ומריץ את משימת האימון של MaxText בכל צומת עובד. הפונקציה train_loop_per_worker עוטפת את נקודת הכניסה הראשית של MaxText, ומשתמשת במתזמן המבוזר של Ray כדי להפעיל את כלי ההדרכה של MaxText בפרוסת TPU מרובת מארחים.

1. שומרים את סקריפט Python הבא בשם maxtext_ray_trainer.py:

   import os
   from absl import app
   import logging
   from typing import Sequence
   import ray
   from ray.train.v2.api.config import ScalingConfig, RunConfig
   from ray.train.v2.jax import JaxTrainer
   
   def train_loop_per_worker(config):
       from MaxText.train import main as maxtext_main
   
       argv = config["argv"]
       maxtext_main(argv)
   
   def main(argv: Sequence[str]):
       trainer = JaxTrainer(
           train_loop_per_worker=train_loop_per_worker,
           train_loop_config={"argv": argv},
           scaling_config=ScalingConfig(
               use_tpu=True,
               num_workers=4,
               topology="4x4",
               accelerator_type="TPU-V6E",
               resources_per_worker={"TPU": 4},
               placement_strategy="SPREAD",
           ),
           run_config=RunConfig(
               name="maxtext_jaxtrainer",
               worker_runtime_env={
                   "env_vars": {
                       "JAX_PLATFORMS": "tpu",
                       "ENABLE_PJRT_COMPATIBILITY": "true",
                       "TPU_SLICE_BUILDER_DUMP_CHIP_FORCE": "true",
                       "TPU_SLICE_BUILDER_DUMP_ICI": "true",
                       "XLA_FLAGS": "--xla_dump_to=/tmp/xla_dump_file --xla_dump_hlo_as_proto",
                   }
               },
           ),
       )
       result = trainer.fit()
       logging.info("Training complete!")
       ray.shutdown()
   
   if __name__ == "__main__":
       app.run(main)

2. כדי לארח את התמונה המותאמת אישית, יוצרים מאגר ב-Artifact Registry:

   gcloud artifacts repositories create ${ARTIFACT_REGISTRY} \
       --repository-format=docker --location=${REGION} && \
   gcloud auth configure-docker ${REGION}-docker.pkg.dev
   

3. כדי ליצור קובץ אימג' שכולל יחסי תלות של Ray ו-MaxText לאימון, יוצרים Dockerfile:

   # Start from a Ray base image which includes JaxTrainer API.
   # Maxtext with TPU requires Python 3.12.
   FROM rayproject/ray:2.49.1-py312
   
   USER root
   RUN groupadd -r ray 2>/dev/null || true && usermod -g ray ray
   
   RUN sudo apt-get update -y \
     && sudo apt-get install --no-install-recommends -y git \
     && sudo rm -rf /var/lib/apt/lists/*
   
   WORKDIR /app
   
   # Clone the Maxtext repo and build from source, installing TPU dependencies.
   RUN git clone https://github.com/AI-Hypercomputer/maxtext.git
   
   RUN pip install --no-cache-dir uv
   
   RUN cd maxtext && \
       uv pip install --no-cache --system -e .[tpu] --resolution=lowest && \
       install_maxtext_github_deps
   
   # Copy the Ray Maxtext trainer to run on the remote container.
   COPY maxtext_ray_trainer.py .
   
   RUN chown -R ray:ray .
   ENV PYTHONPATH=/app/maxtext/src:/app/maxtext:/app
   USER ray

   הערה: במדריך הזה נעשה שימוש בתמונה בהתאמה אישית כי התלות גדולה מדי ואי אפשר להתקין אותה באמצעות הפקודה runtime_env של Ray.

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

   export DOCKER_IMAGE=${REGION}-docker.pkg.dev/${PROJECT_ID}/${ARTIFACT_REGISTRY}/ray-maxtext:latest
   gcloud builds submit --tag ${DOCKER_IMAGE}
   

אימון המודל

1. שומרים את קובץ המניפסט לדוגמה הבא בשם maxtext-tpu-cluster.yaml:

   apiVersion: ray.io/v1
   kind: RayCluster
   metadata:
     name: maxtext-tpu-cluster
   spec:
     headGroupSpec:
       rayStartParams: {}
       template:
         metadata:
           annotations:
             gke-gcsfuse/volumes: "true"
             gke-gcsfuse/cpu-limit: "0"
             gke-gcsfuse/memory-limit: "0"
             gke-gcsfuse/ephemeral-storage-limit: "0"
         spec:
           serviceAccountName: ${KSA_NAME}
           containers:
             - name: ray-head
               image: ${DOCKER_IMAGE}
               imagePullPolicy: IfNotPresent
               ports:
               - containerPort: 6379
                 name: gcs-server
               - containerPort: 8265
                 name: dashboard
               - containerPort: 10001
                 name: client
               resources:
                 limits:
                   memory: "16Gi"
                 requests:
                   cpu: "8"
                   memory: "16Gi"
               volumeMounts:
               - name: gcs-fuse-csi-ephemeral
                 mountPath: /data
               - name: dshm
                 mountPath: /dev/shm
           volumes:
           - name: gcs-fuse-cache
             emptyDir:
               medium: Memory
           - name: dshm
             emptyDir:
               medium: Memory
           - name: gcs-fuse-csi-ephemeral
             csi:
               driver: gcsfuse.csi.storage.gke.io
               volumeAttributes:
                 bucketName: ${GS_BUCKET}
                 mountOptions: "implicit-dirs"
     workerGroupSpecs:
       - replicas: 1
         numOfHosts: 4
         groupName: tpu-group
         rayStartParams: {}
         template:
           metadata:
             annotations:
               gke-gcsfuse/volumes: "true"
               gke-gcsfuse/cpu-limit: "0"
               gke-gcsfuse/memory-limit: "0"
               gke-gcsfuse/ephemeral-storage-limit: "0"
           spec:
             serviceAccountName: ${KSA_NAME}
             containers:
               - name: ray-worker
                 image: ${DOCKER_IMAGE}
                 imagePullPolicy: IfNotPresent
                 resources:
                   limits:
                     memory: 200G
                     google.com/tpu: "4"
                   requests:
                     cpu: "8"
                     memory: 200G
                     google.com/tpu: "4"
                 env:
                   - name: JAX_PLATFORMS
                     value: tpu
                   - name: ENABLE_PJRT_COMPATIBILITY
                     value: "true"
                 volumeMounts:
                 - name: gcs-fuse-csi-ephemeral
                   mountPath: /data
                 - name: dshm
                   mountPath: /dev/shm
             volumes:
             - name: gcs-fuse-cache
               emptyDir:
                 medium: Memory
             - name: dshm
               emptyDir:
                 medium: Memory
             - name: gcs-fuse-csi-ephemeral
               csi:
                 driver: gcsfuse.csi.storage.gke.io
                 volumeAttributes:
                   bucketName: ${GS_BUCKET}
                   mountOptions: "implicit-dirs"
             nodeSelector:
               cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
               cloud.google.com/gke-tpu-topology: 4x4

   הגדרת ה-RayCluster שמופיעה למעלה יוצרת קבוצת עובדים של TPU עם ארבעה עובדים (numOfHosts: 4) לכל רפליקה. כל עובד מבקש ארבעה שבבי TPU ‏(google.com/tpu: "4"). העובדים יתוזמנו בצומת שמריץ TPU Trillium ‏ (tpu-v6e-slice), שהוא חלק מאותו פלח מרובה מארחים שנמצא באותו מיקום. ‫KubeRay משנה את קנה המידה של כל ארבעת ה-workers באופן אטומי, ומשתני הסביבה הנדרשים של JAX, כמו גם Pod Affinities לתזמון, מאותחלים על ידי GKE באמצעות webhook משתנה.

2. כדי להגדיר את הערכים הנדרשים בקובץ ה-YAML, יוצרים את ה-RayCluster באמצעות envsubst:

   envsubst < maxtext-tpu-cluster.yaml | kubectl apply -f -
   

3. מוודאים שהאשכול מוכן ופועל:

   kubectl get rayclusters maxtext-tpu-cluster
   

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

   NAME                  DESIRED WORKERS   AVAILABLE WORKERS   CPUS   MEMORY        GPUS   STATUS   AGE
   maxtext-tpu-cluster   4                 4                   40     798027216Ki   0      ready    11m
   

4. כדי לגשת ללוח הבקרה של Ray דרך שירות ה-Ray head, צריך ליצור סשן של העברת פורטים:

   kubectl port-forward svc/maxtext-tpu-cluster-head-svc 8265:8265 2>&1 >/dev/null &
   

5. מוודאים שאפשר לגשת ל-RayCluster מהסביבה המקומית:

   ray list nodes --address http://localhost:8265
   

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

   ======== List: 2025-09-13 03:53:16.988269 ========
   Stats:
   ------------------------------
   Total: 5
   Table:
   ------------------------------
       NODE_ID                                                   NODE_IP    IS_HEAD_NODE    STATE    STATE_MESSAGE    NODE_NAME    RESOURCES_TOTAL                  LABELS
   0  92c79d04c34b659c1e3044f7642ad3fd47eb16f290785237149fab56  10.84.0.9
   (...)
   

6. שולחים את הסקריפט JaxTrainer ל-RayCluster ובודקים שה-RayJob הושלם בהצלחה:

   ray job submit \
     --address http://localhost:8265 \
     -- python /app/maxtext_ray_trainer.py \
         /app/maxtext/src/MaxText/configs/base.yml \
          base_output_directory=/data/ \
         dataset_type=synthetic \
         per_device_batch_size=1 \
         max_target_length=4096 \
         model_name=llama3-8b \
         steps=100 \
         ici_fsdp_parallelism=4 \
         ici_tensor_parallelism=4 \
         run_name=rayjob-8b-4096-tp4-4x4
   

   הפקודה הקודמת שולחת את סקריפט Python, שקורא לקוד JaxTrainerRay אל RayCluster. הפקודה ray job submit כוללת כמה ארגומנטים ספציפיים ל-MaxText שמועברים להגדרת המודל.

   הערה: יכול להיות שיחלפו כמה דקות עד שהסקריפט יושלם.

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

   (RayTrainWorker pid=21663, ip=10.12.3.6) completed step: 99, seconds: 1.100, TFLOP/s/device: 179.739, Tokens/s/device: 3725.218, total_weights: 65536, loss: 0.000 [repeated 3x across cluster]
   
   ------------------------------------------
   Job 'raysubmit_zCrJcWnuymMQv4C3' succeeded
   ------------------------------------------
   

הצלחה: השלמתם את ההדרכה באמצעות JaxTrainer ו-MaxText של Ray Train במאגרי צמתים מבוזרים של TPU מרובי מארחים.

הסרת המשאבים

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

1. מוחקים את RayCluster:

   kubectl delete raycluster maxtext-tpu-cluster
   

2. מוחקים את אשכול GKE:

   gcloud container clusters delete $CLUSTER_NAME --zone=$ZONE
   

3. מוחקים את הקטגוריה של Cloud Storage:

   gsutil rm -r gs://${GS_BUCKET}
   

4. מחיקת מאגר Artifact Registry:

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

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

• מידע על Ray ב-Kubernetes
• איך מריצים vLLM ב-GKE עם TPU
• איך מפעילים את SDXL ב-GKE עם TPU
• מידע נוסף על TPU ב-GKE