הרצת עומס עבודה של אצווה באמצעות Pathways

המדריך הזה משתמש בעומס עבודה של אימון JAX שפועל עד להשלמה באותו אשכול כמו רכיבי הבקר של Pathways, כדי להדגים איך להריץ עומס עבודה באצווה.

אתם יכולים להריץ עומסי עבודה באצווה באשכול של Google Kubernetes Engine (GKE) או באשכול של Cluster Director עם Slurm:

• ב-GKE, בסיום עומס העבודה של JAX, רכיבי אשכול Pathways מסתיימים.
• ב-Cluster Director, הרכיבים של בקר Pathways פועלים כמשימת Slurm נפרדת, ולכן הם ממשיכים לפעול אחרי שעומס העבודה של JAX מסתיים, עד שאתם עוצרים אותם.

בדף הזה יש הוראות לשתי הפלטפורמות. בכל קטע, בוחרים בכרטיסייה GKE או Cluster Director.

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

• צריך לבחור את הכרטיסייה הרלוונטית לאופן שבו תכננתם להשתמש בדוגמאות בדף הזה:

  gcloud

  מפעילים את Cloud Shell במסוף Google Cloud .

  הפעלת Cloud Shell

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

  REST

  כדי להשתמש בסביבת פיתוח מקומית בדוגמאות של API בארכיטקטורת REST שבדף הזה, צריך להשתמש בפרטי הכניסה שאתם נותנים ל-CLI של gcloud.

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

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

  מידע נוסף מופיע במאמר אימות לשימוש ב-REST במסמכי האימות של Google Cloud .

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

GKE

מבצעים את הפעולות הבאות:

• יוצרים אשכול GKE.
• הגדרת Cluster Toolkit.
• התקנת כלי Kubernetes.
• מפעילים את Kubernetes Engine API.
• יוצרים קטגוריה של Cloud Storage לפריטי המידע ולפלט של Pathways.

Cluster Director

מבצעים את הפעולות הבאות:

• יוצרים אשכול ב-Cluster Director ומגדירים את סביבת Pathways.
• יוצרים קטגוריה של Cloud Storage שאפשר לכתוב בה נתונים זמניים ונקודות ביקורת של Pathways.

יצירת קובץ אימג' לאימון באמצעות MaxText

‫MaxText הוא פרויקט קוד פתוח של מודל שפה גדול (LLM) שפותח על ידי Google. הוא כתוב ב-JAX ומיועד לפעול ביעילות ב-TPU וב-GPU. Google Cloud

GKE

כדי ליצור קובץ אימג' של MaxText ב-Docker באמצעות הגרסה היציבה האחרונה של JAX ממאגר GitHub בקוד פתוח, מריצים את הפקודה הבאה:

git clone https://github.com/AI-Hypercomputer/maxtext
cd maxtext
gcloud config set project PROJECT_ID
bash src/dependencies/scripts/docker_build_dependency_image.sh MODE=stable
gcloud auth configure-docker
# This script needs bash version >= 4.2 to execute.
bash src/dependencies/scripts/docker_upload_runner.sh \
    CLOUD_IMAGE_NAME=USER_runner

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

• ‫PROJECT_ID: מזהה הפרויקט ב-Google Cloud .
• ‫USER: מזהה המשתמש ב-Google Cloud .

הפקודה הזו מעבירה את קובץ האימג' של MaxText Kubernetes אל gcr.io/PROJECT_ID/USER_runner. אתם יכולים להשתמש בקובץ האימג' הזה של Docker כדי להריץ אימון ב-TPU באמצעות הקצה העורפי של Pathways.

Cluster Director

לא צריך ליצור קובץ אימג' לאימון: setup_env.sh התקינה את MaxText בסביבת Pathways ב-/home/pathways, ו-run_maxtext.sh מריצה אותו.

הרצת עומס עבודה של Batch

GKE

1. אם הקובץ הבינארי gcluster לא נמצא במשתנה הסביבה PATH, מוסיפים את הספרייה שמכילה את gcluster ל-PATH:

   export PATH="$PATH:CLUSTER_TOOLKIT_PATH"

   מחליפים את CLUSTER_TOOLKIT_PATH בנתיב המוחלט לספרייה שמכילה את הקובץ הבינארי gcluster, למשל $HOME/cluster-toolkit.

2. שולחים את קובץ האימג' של MaxText Docker שנוצר מראש באמצעות הפקודה gcluster job submit:

   gcluster job submit \
       --pathways \
       --pathways-gcs-location="gs://BUCKET_NAME/pathways-artifacts" \
       --name=WORKLOAD \
       --cluster=CLUSTER \
       --project=PROJECT_ID \
       --location=LOCATION \
       --num-slices=WORKLOAD_NODEPOOL_COUNT \
       --compute-type=COMPUTE_TYPE \
       --topology=TOPOLOGY \
       --image="gcr.io/PROJECT_ID/USER_runner" \
       --command="python3 -m maxtext.trainers.pre_train.train /deps/src/maxtext/configs/base.yml base_output_directory=gs://BUCKET_NAME per_device_batch_size=1 enable_checkpointing=false remat_policy=full global_parameter_scale=1 steps=20 max_target_length=2048 use_iota_embed=true reuse_example_batch=1 dataset_type=synthetic attention=flash gcs_metrics=True enable_single_controller=True run_name=RUN_NAME-pathways-job"

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

   • ‫BUCKET_NAME: הקטגוריה של Cloud Storage לפלט.
   • ‫WORKLOAD: שם ייחודי לזיהוי עומס העבודה (עד 22 תווים, בגלל מגבלות התוויות של DNS).
   • ‫CLUSTER: השם של אשכול GKE.
   • ‫PROJECT_ID: מזהה הפרויקט ב-Google Cloud .
   • ‫LOCATION: האזור או התחום של אשכול GKE.
   • ‫WORKLOAD_NODEPOOL_COUNT: מספר מאגרי הצמתים של פרוסות ה-TPU.
   • ‫COMPUTE_TYPE: סוג המכונה של ה-TPU, לדוגמה ct6e-standard-4t. למידע נוסף על סוגי ה-TPU שנתמכים בכל גרסת TPU, קראו את המאמר גרסאות TPU.
   • ‫TOPOLOGY: טופולוגיית המיקום של ה-TPU, לדוגמה 2x4.
   • ‫USER: מזהה המשתמש ב-Google Cloud .
   • ‫RUN_NAME: שם שהמשתמש הקצה כדי לזהות את הרצת תהליך העבודה.

   מידע נוסף על האפשרויות לשליחת משימות זמין במדריך לשליחת משימות ב-Cluster Toolkit.

3. עוקבים אחרי ההתקדמות של עומס העבודה באמצעות הפקודה gcluster job logs:

   gcluster job logs WORKLOAD \
       --cluster=CLUSTER \
       --project=PROJECT_ID \
       --location=LOCATION \
       --main-only \
       --follow

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

   completed step: 0, seconds: 1.259, TFLOP/s/device: 10.792, Tokens/s/device: 1626.636, total_weights: 65536, loss: 10.876, lm_loss: 10.876, perplexity: 52906.090
   completed step: 1, seconds: 3.366, TFLOP/s/device: 4.036, Tokens/s/device: 608.366, total_weights: 65536, loss: 10.876, lm_loss: 10.876, perplexity: 52906.090
   completed step: 2, seconds: 0.264, TFLOP/s/device: 51.415, Tokens/s/device: 7749.885, total_weights: 65536, loss: 10.563, lm_loss: 10.563, perplexity: 38665.316
   completed step: 3, seconds: 1.025, TFLOP/s/device: 13.257, Tokens/s/device: 1998.187, total_weights: 65536, loss: 9.983, lm_loss: 9.983, perplexity: 21646.125
   completed step: 4, seconds: 1.194, TFLOP/s/device: 11.376, Tokens/s/device: 1714.750, total_weights: 65536, loss: 9.462, lm_loss: 9.462, perplexity: 12864.218
   completed step: 5, seconds: 1.185, TFLOP/s/device: 11.465, Tokens/s/device: 1728.203, total_weights: 65536, loss: 8.961, lm_loss: 8.961, perplexity: 7791.373
   

4. אופציונלי: כדי לבטל את עומס העבודה לפני שהוא מסתיים, משתמשים בפקודה gcluster job cancel:

   gcluster job cancel WORKLOAD \
       --cluster=CLUSTER \
       --project=PROJECT_ID \
       --location=LOCATION

Cluster Director

מריצים את הפקודות הבאות בצומת ההתחברות של האשכול מ-~/pathways. הסקריפטים מחפשים אחד את השני וכותבים את היומנים שלהם בספרייה הזו.

1. מפעילים את שירות Pathways:

   cd ~/pathways
   SLICE_SHAPE=TOPOLOGY \
       MIN_SLICES=MIN_SLICES \
       MAX_SLICES=MAX_SLICES \
       GCS_BUCKET=gs://BUCKET_NAME \
       sbatch start_pathways.sh

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

   • ‫TOPOLOGY: הטופולוגיה של TPU בכל פרוסה. הוא צריך להתאים ל-acceleratorTopology של קבוצת הצמתים של ה-TPU, לדוגמה 2x4.
   • ‫MIN_SLICES: המספר המינימלי של פרוסות פעילות שנדרשות כדי להתחיל את האימון ולהמשיך אותו.
   • ‫MAX_SLICES: מספר הפרוסות המקסימלי לשימוש, בדרך כלל מספר הפרוסות בקבוצת הצמתים.
   • ‫BUCKET_NAME: הקטגוריה של Cloud Storage לנתוני טיוטה ולנקודות ביקורת של Pathways.

   אם מחיצת ה-TPU היא לא מחיצת ברירת המחדל של האשכול, צריך גם להגדיר את PARTITION לשם שלה.

   השירות מוכן כשהיומן שלו, pathways-service-JOB_ID.out (כאשר JOB_ID הוא מזהה המשימה ב-Slurm), מציג פלט שדומה לזה:

   Pathways Runtime Service is READY!
    Discovery file    : /home/USER/pathways/cluster.env
    IFRT Proxy Target : grpc://RM_IP:29000
    Running TPU Slices : 2 / 2 (min=2)
   

2. שולחים את עומס העבודה לאימון MaxText. אפשר לשלוח אותה מיד אחרי משימת השירות: היא ממתינה עד שעה עד שהשירות יהיה מוכן.

   sbatch run_maxtext.sh
   

   משימת MaxText ממתינה עד שבשירות Pathways יש לפחות MIN_SLICES פרוסות פועלות לפני שהיא מתחילה את האימון. עומס העבודה מאמן את Llama 3.1 8B על נתונים סינתטיים. הגדרות הפרוסה והקטגוריה מגיעות מהשירות, אז אל תגדירו אותן במשימה הזו. אפשר להגדיר את המשתנים האופציונליים הבאים:

   • ‫STEPS: מספר שלבי האימון. ערך ברירת המחדל הוא 30.
   • ‫CHECKPOINT_PERIOD: מספר השלבים בין נקודות הביקורת. ערך ברירת המחדל הוא 5.
   • ‫BASE_OUTPUT_DIR: הנתיב ב-Cloud Storage שבו MaxText כותב נקודות ביקורת. ערך ברירת המחדל הוא gs://BUCKET_NAME/maxtext-runs.
   • ‫REDUCED_LAYERS=1: אימון של וריאנט עם 8 שכבות במקום 32 שכבות מלאות, לפרוסות 2x2.

   עוקבים אחרי ההתקדמות של עומס העבודה (מחליפים את JOB_ID במזהה המשימה של MaxText ב-Slurm). ‫MaxText כותב את היומן שלו לשגיאה סטנדרטית (stderr):

   tail -f maxtext-llama31-JOB_ID.err

   כל שלב אימון מדפיס שורה שדומה לזו:

   completed step: 24, seconds: 0.326, live slice count: 2, TFLOP/s/device: 88.199, Tokens/s/device: 6289.27, total_weights: 32768, loss: 6.238
   

   אפשר לשלוח שוב את run_maxtext.sh מאוחר יותר בלי להפעיל מחדש את השירות.

3. בסיום, עוצרים את עומס העבודה ואז את שירות Pathways:

   scancel -u ${USER} -n maxtext-llama31
   scancel -u ${USER} -n pathways-service
   

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

נתקלתם בבעיות? קראו את המאמר פתרון בעיות ב-Pathways ב-Cloud.

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

• יצירת אשכול GKE באמצעות Pathways
• יצירת אשכול Cluster Director עם Pathways
• הרצת עומס עבודה אינטראקטיבי באמצעות Pathways
• העברת עומסי עבודה של JAX ל-Pathways
• אימון עמיד באמצעות Pathways
• פתרון בעיות ב-Pathways on Cloud