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

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

KubeRay

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

מטרות

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

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

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

• Select a project: כדי לבחור פרויקט לא צריך תפקיד 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, צריך את תפקיד ה-IAM 'אדמין של Service Usage' (roles/serviceusage.serviceUsageAdmin), שכולל את ההרשאה serviceusage.services.enable. איך מקצים תפקידים

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 פרויקט.

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

• Select a project: כדי לבחור פרויקט לא צריך תפקיד 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, צריך את תפקיד ה-IAM 'אדמין של Service Usage' (roles/serviceusage.serviceUsageAdmin), שכולל את ההרשאה serviceusage.services.enable. איך מקצים תפקידים

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. בחלק העליון של חלון המסוף, לוחצים על הלחצן Activate Cloud Shell . Google Cloud

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

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 מקצה מאגר צמתים שמורכב מארבע מכונות וירטואליות של TPU Trillium ‏ (v6e), שמוגדרות יחד כפרוסת TPU מרובת מארחים, עם טופולוגיה של 4x4, שמוכנה לעומסי עבודה של אימון מבוזר.

באשכול GKE שמופעל בו Ray operator, המערכת מתקינה אוטומטית את KubeRay ואת KubeRay TPU webhook באשכול.

הגדרת קטגוריה של 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 slice מרובה-מארחים ומריץ את משימת האימון 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 מבצע התאמה לעומס של כל ארבעת העובדים באופן אטומי, ומשתני הסביבה הנדרשים של 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, שקורא לקוד JaxTrainer Ray אל 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
• מידע נוסף על TPUs ב-GKE