סקירה כללית על ריבוי-פרוסות (Multislice) ב-Cloud TPU

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

חשוב: אפשר להשתמש ב-Multislice רק ב-Cloud TPU v4 ובגרסאות TPU מתקדמות יותר.

‫Cloud TPU Multislice היא טכנולוגיה מלאה להרחבת הביצועים שמאפשרת למשימת אימון להשתמש בכמה פרוסות TPU בתוך פוד יחיד או בפרוסות בכמה פודים עם מקביליות נתונים רגילה. ב-TPU v4, אפשר להשתמש ביותר מ-4,096 שבבים להרצת אימון אחת. במשימות אימון שדורשות פחות מ-4,096 צ'יפים, פרוסה אחת יכולה להציע את הביצועים הטובים ביותר. עם זאת, יש יותר זמינות של כמה פרוסות קטנות יותר, מה שמאפשר זמן הפעלה מהיר יותר כשמשתמשים ב-Multislice עם פרוסות קטנות יותר.

כשפריסת ה-TPU מתבצעת בתצורות Multislice, שבבי ה-TPU בכל slice מתקשרים באמצעות inter-chip-interconnect ‏ (ICI). שבבי TPU בפרוסות שונות מתקשרים באמצעות העברת נתונים למעבדים (מארחים), שמעבירים את הנתונים דרך רשת מרכז הנתונים (DCN). מידע נוסף על שינוי גודל באמצעות Multislice זמין במאמר איך לשנות את גודל האימון של AI לעשרות אלפי שבבי Cloud TPU באמצעות Multislice.

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

מושגים

סוג המאיץ

הצורה של כל פרוסת TPU שמרכיבה Multislice. כל פרוסה בבקשה לריבוי פרוסות היא מאותו סוג מאיץ. סוג המאיץ מורכב מסוג TPU (v4 ואילך) ואחריו מספר ה-TensorCores. לדוגמה, v5litepod-128 מציין TPU v5e עם 128 TensorCores.

תיקון אוטומטי

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

Data Center Networking (DCN)

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

תזמון קבוצתי

כשכל פרוסות ה-TPU מוקצות יחד, באותו זמן, מובטח שכל הפרוסות או אף אחת מהן לא יוקצו בהצלחה.

חיבור בין שבבים (ICI)

קישורים פנימיים במהירות גבוהה וזמן אחזור קצר שמחברים בין יחידות TPU בתוך TPU Pod.

ריבוי-פרוסות (Multislice)

שני פרוסות של שבבי TPU או יותר שיכולים לתקשר דרך DCN.

Node

בהקשר של ריבוי-פרוסות (Multislice), צומת מתייחס לפרוסת TPU אחת. לכל פרוסת TPU בריבוי-פרוסות מוקצה מזהה צומת.

סקריפט לטעינה בזמן ההפעלה

סקריפט סטנדרטי להפעלה של Compute Engine שפועל בכל פעם שמפעילים או מאתחלים מכונה וירטואלית. במקרה של Multislice, הוא מצוין בבקשה ליצירת קוד ה-QR. מידע נוסף על סקריפטים לטעינה בזמן ההפעלה של Cloud TPU זמין במאמר בנושא ניהול משאבי TPU.

Tensor

מבנה נתונים שמשמש לייצוג נתונים רב-ממדיים במודל למידת מכונה.

סוגי הקיבולת של Cloud TPU

אפשר ליצור יחידות TPU מסוגים שונים של קיבולת (ראו את האפשרויות לשימוש במאמר איך מתבצע התמחור של TPU):

• הזמנה: כדי להשתמש בהזמנה, צריך להיות לכם הסכם הזמנה עם Google. משתמשים בדגל --reserved כשיוצרים את המשאבים.

• Spot: מכוון למכסות זמניות באמצעות VM במודל Spot. יכול להיות שהמשאבים שלכם יופקעו כדי לפנות מקום לבקשות של משימה בעדיפות גבוהה יותר. משתמשים בדגל --spot כשיוצרים את המשאבים.

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

קדימה, מתחילים

1. הגדרת סביבת Cloud TPU.

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

   הפעלת Cloud Shell

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

כדי להשתמש ב-Multislice, צריך לנהל את משאבי ה-TPU כמשאבים בתור.

דוגמה ראשונית

במדריך הזה נעשה שימוש בקוד ממאגר GitHub של MaxText. ‫MaxText הוא מודל בסיסי של LLM בקוד פתוח, עם ביצועים גבוהים, יכולת התאמה לעומס (autoscaling) שרירותית ויכולת מוכחת, שנכתב ב-Python וב-Jax. ‫MaxText תוכנן לאימון יעיל ב-Cloud TPU.

הקוד ב-shardings.py נועד לעזור לכם להתחיל להתנסות באפשרויות שונות של מקביליות. לדוגמה, מקביליות נתונים, מקביליות נתונים עם חלוקה מלאה (FSDP) ומקביליות טנסור. הקוד ניתן להרחבה מסביבות של פרוסה אחת לסביבות של Multislice.

מקביליות של ICI

‫ICI הוא קיצור של Interchip Interconnect, כלומר חיבור בין-שבבי במהירות גבוהה שמחבר את יחידות ה-TPU בפרוסת TPU אחת. חלוקה ל-ICI תואמת לחלוקה לפרוסות. ‫shardings.py מספק שלושה פרמטרים של מקביליות ICI:

• ici_data_parallelism
• ici_fsdp_parallelism
• ici_tensor_parallelism

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

צריך להגביל את הקלט כך ש-ici_data_parallelism * ici_fsdp_parallelism * ici_tensor_parallelism יהיה שווה למספר הצ'יפים בפלח.

בטבלה הבאה מוצגות דוגמאות לקלט של משתמשים להקבלה של ICI לארבעת הצ'יפים שזמינים ב-v4-8:

	ici_data_parallelism	ici_fsdp_parallelism	ici_tensor_parallelism
4-way FSDP	1	4	1
מקביליות טנסור עם 4 כיוונים	1	1	4
‫2-way FSDP + 2-way Tensor parallelism	1	2	2

שימו לב: ברוב המקרים צריך להשאיר את הערך של ici_data_parallelism כ-1, כי רשת ה-ICI מהירה מספיק כדי ש-FSDP תמיד יהיה עדיף על מקביליות נתונים.

בדוגמה הזו אנחנו מניחים שאתם יודעים איך להריץ קוד בפרוסת TPU אחת, כמו שמוסבר במאמר הרצת חישוב ב-Cloud TPU VM באמצעות JAX. בדוגמה הזו אפשר לראות איך מריצים את shardings.py על פרוסה אחת.

1. מגדירים את הסביבה:

   $ gcloud auth login
   $ export QR_ID=your-queued-resource-id
   $ export TPU_NAME=your-tpu-name
   $ export PROJECT=your-project-name
   $ export ZONE=us-central1-a
   $ export NETWORK_NAME=your-network-name
   $ export SUBNETWORK_NAME=your-subnetwork-name
   $ export RUNTIME_VERSION=v2-alpha-tpuv5-lite
   $ export ACCELERATOR_TYPE=v5litepod-16
   $ export EXAMPLE_TAG_1=your-tag-1
   $ export EXAMPLE_TAG_2=your-tag-2
   $ export SLICE_COUNT=4
   $ export STARTUP_SCRIPT='#!/bin/bash\n'

   תיאורי משתנים

   • ‫QR_ID: המזהה שהמשתמש הקצה למשאב בתור.
   • ‫TPU_NAME: השם שהמשתמש הקצה ל-TPU.
   • ‫PROJECT:שם הפרויקט ב- Google Cloud
   • ‫ZONE: מציין את האזור שבו רוצים ליצור את המשאבים.
   • ‫NETWORK_NAME: השם של רשת ה-VPC.
   • ‫SUBNETWORK_NAME: השם של תת-הרשת ברשת ה-VPC
   • ‫RUNTIME_VERSION: גרסת התוכנה של Cloud TPU.
   • ‫ACCELERATOR_TYPE: סוג המאיץ מציין את הגרסה והגודל של ה-Cloud TPU שרוצים ליצור.
   • ‫EXAMPLE_TAG_1, EXAMPLE_TAG_2 …: תגים שמשמשים לזיהוי מקורות או יעדים חוקיים של חומות אש ברשת
   • ‫SLICE_COUNT: מספר הפלחים. מוגבל ל-256 פרוסות לכל היותר.
   • ‫STARTUP_SCRIPT: אם מציינים סקריפט לטעינה בזמן ההפעלה, הסקריפט יפעל כשפרוסת ה-TPU תוקצה או תופעל מחדש.

2. יוצרים מפתחות SSH ל-gcloud. מומלץ להשאיר סיסמה ריקה (צריך להקיש על מקש Enter פעמיים אחרי הרצת הפקודה הבאה). אם מופיעה הודעה שהקובץ google_compute_engine כבר קיים, מחליפים את הגרסה הקיימת.

   $ ssh-keygen -f ~/.ssh/google_compute_engine

3. מקצים את יחידות ה-TPU:

   gcloud

   $ gcloud compute tpus queued-resources \
       create ${QR_ID} \
       --accelerator-type=${ACCELERATOR_TYPE} \
       --runtime-version=${RUNTIME_VERSION} \
       --node-id=${TPU_NAME} \
       --zone=${ZONE} \
       [--reserved |--spot]

   ב-Google Cloud CLI אין תמיכה בכל האפשרויות ליצירת בקשות איכות, כמו תגים. למידע נוסף, אפשר לעיין במאמר בנושא יצירת קודי QR.

   המסוף

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

      כניסה לדף TPUs

   2. לוחצים על Create TPU.

   3. בשדה Name, נותנים שם ל-TPU.

   4. בתיבה Zone (אזור), בוחרים את האזור שבו רוצים ליצור את ה-TPU.

   5. בתיבה TPU type, בוחרים סוג של מאיץ. סוג המאיץ מציין את הגרסה והגודל של ה-Cloud TPU שרוצים ליצור. מידע נוסף על סוגי המאיצים הנתמכים לכל גרסת TPU זמין במאמר גרסאות TPU.

   6. בתיבה TPU software version, בוחרים גרסת תוכנה. כשיוצרים מכונת TPU וירטואלית ב-Cloud TPU, גרסת התוכנה של ה-TPU מציינת את גרסת זמן הריצה של ה-TPU שמותקנת. מידע נוסף זמין במאמר בנושא גרסאות תוכנה של TPU.

   7. לוחצים על המתג הפעלת הוספה לתור.

   8. בשדה Queued resource name, מזינים שם לבקשת הקצאת המשאבים.

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

4. מחכים עד שהמשאב בתור יהיה במצב ACTIVE, כלומר שצומתי העובדים יהיו במצב READY. אחרי שההקצאה של המשאב בתור מתחילה, יכול להיות שיעברו בין דקה ל-5 דקות עד שהיא תסתיים, בהתאם לגודל של המשאב בתור. אתם יכולים לבדוק את הסטטוס של בקשה להקצאת משאבים בתור באמצעות ה-CLI של gcloud או מסוף Google Cloud :

   gcloud

   $ gcloud compute tpus queued-resources \
       list --filter=${QR_ID} --zone=${ZONE}

   המסוף

   1. נכנסים לדף TPUs במסוף Google Cloud :

      כניסה לדף TPUs

   2. לוחצים על הכרטיסייה משאבים בתור.

   3. לוחצים על השם של בקשת המשאב שנמצאת בתור.

5. מתחברים ל-TPU VM באמצעות SSH:

   $ gcloud compute tpus tpu-vm ssh ${TPU_NAME} --zone=${ZONE}

6. משכפלים את MaxText (שכולל את shardings.py) ל-TPU VM:

   $ git clone https://github.com/AI-Hypercomputer/maxtext && cd maxtext

7. מתקינים את Python 3.10:

   $ sudo apt-get update
   $ sudo apt install python3.10
   $ sudo apt install python3.10-venv

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

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

   $ python3 -m venv your-venv-name
   $ source your-venv-name/bin/activate

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

   $ bash setup.sh

10. מריצים את הפקודה הבאה כדי להפעיל את shardings.py בפלח ה-TPU.

    $ python3 -m pedagogical_examples.shardings \
      --ici_fsdp_parallelism 4 \
      --batch_size 131072 \
      --embedding_dimension 2048

    התוצאות מופיעות ביומנים. יחידות ה-TPU אמורות להגיע ל-260 TFLOP לשנייה, או לניצול מרשים של יותר מ-90%FLOP. במקרה הזה, בחרנו בערך שקרוב לגודל המקסימלי של אצווה שאפשר להכניס לזיכרון ברוחב פס גבוה (HBM) של ה-TPU.

11. אתם יכולים לבדוק אסטרטגיות אחרות של חלוקה לשברים ב-ICI, למשל, אפשר לנסות את השילוב הבא:

    $ python3 -m pedagogical_examples.shardings \
      --ici_tensor_parallelism 4 \
      --batch_size 131072 \
      --embedding_dimension 2048

12. בסיום, מוחקים את המשאב בתור ואת פרוסת ה-TPU. צריך להריץ את שלבי הניקוי האלה מהסביבה שבה הגדרתם את הפרוסה (קודם מריצים את exit כדי לצאת מסשן ה-SSH). תהליך המחיקה יימשך בין שתי לחמש דקות. אם אתם משתמשים ב-CLI של gcloud, אתם יכולים להריץ את הפקודה הזו ברקע עם הדגל האופציונלי --async.

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

    gcloud

    $ gcloud compute tpus queued-resources \
        delete ${QR_ID} --force (--async)

    המסוף

    1. נכנסים לדף TPUs במסוף Google Cloud :

       כניסה לדף TPUs

    2. לוחצים על הכרטיסייה Queued resources.

    3. מסמנים את התיבה לצד בקשת המשאב בתור.

    4. לוחצים על מחיקה.

פיצול לפרוסות מרובות באמצעות מקביליות DCN

סקריפט shardings.py מקבל שלושה פרמטרים שמציינים את המקביליות של DCN, בהתאם למספר הרסיסים של כל סוג של מקביליות נתונים:

• dcn_data_parallelism
• dcn_fsdp_parallelism
• dcn_tensor_parallelism

הערכים של הפרמטרים האלה צריכים להיות מוגבלים כך שdcn_data_parallelism * dcn_fsdp_parallelism * dcn_tensor_parallelism יהיה שווה למספר הפלחים.

לדוגמה, כדי ליצור שני פלחים, משתמשים ב---dcn_data_parallelism = 2.

	dcn_data_parallelism	dcn_fsdp_parallelism	dcn_tensor_parallelism	מספר הפלחים
הקבלה דו-כיוונית של נתונים	2	1	1	2

הערך של dcn_tensor_parallelism תמיד צריך להיות 1 כי ה-DCN לא מתאים לחלוקה כזו לשברים. בעומסי עבודה אופייניים של LLM בשבבי v4, צריך להגדיר את dcn_fsdp_parallelism ל-1, ולכן צריך להגדיר את dcn_data_parallelism למספר הפרוסות, אבל זה תלוי באפליקציה.

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

הפעלת shardings.py בסביבת Multislice

אפשר להריץ את shardings.py בסביבת ריבוי-פרוסות (Multislice) באמצעות multihost_runner.py או על ידי הרצת shardings.py בכל TPU VM. כאן אנחנו משתמשים ב-multihost_runner.py. השלבים הבאים דומים מאוד לשלבים שמפורטים בקטע Getting Started: Quick Experiments on Multiple slices במאגר MaxText, אבל כאן אנחנו מריצים את shardings.py במקום את ה-LLM המורכב יותר ב-train.py.

הכלי multihost_runner.py מותאם לניסויים מהירים, שבהם נעשה שימוש חוזר באותם מעבדי TPU. הסקריפט multihost_runner.py מסתמך על חיבורי SSH לטווח ארוך, ולכן לא מומלץ להשתמש בו למשימות שפועלות לאורך זמן. אם אתם רוצים להריץ משימה ארוכה יותר (למשל, שעות או ימים), מומלץ להשתמש ב-multihost_job.py.

במדריך הזה, המונח runner מציין את המכונה שבה מריצים את הסקריפט multihost_runner.py. אנחנו משתמשים במונח workers כדי לציין את מכונות ה-TPU הווירטואליות שמרכיבות את הפרוסות. אפשר להריץ את multihost_runner.py במכונה מקומית או בכל מכונה וירטואלית של Compute Engine באותו פרויקט שבו נמצאות הפרוסות. אי אפשר להריץ את multihost_runner.py ב-worker.

‫multihost_runner.py מתחבר באופן אוטומטי ל-TPU workers באמצעות SSH.

בדוגמה הזו, מריצים את shardings.py בשתי פרוסות v5e-16, סך הכל ארבע מכונות וירטואליות ו-16 שבבי TPU. אפשר לשנות את הדוגמה כדי להריץ אותה על יותר יחידות TPU.

מגדירים את הסביבה

1. משכפלים את MaxText במכונת ה-runner:

   $ git clone https://github.com/AI-Hypercomputer/maxtext

2. עוברים לספריית המאגר.

   $ cd maxtext

3. יוצרים מפתחות SSH עבור gcloud. מומלץ להשאיר סיסמה ריקה (מקישים על Enter פעמיים אחרי שמריצים את הפקודה הבאה). אם מופיעה הודעה שהקובץ google_compute_engine כבר קיים, בוחרים לא לשמור את הגרסה הקיימת.

     $ ssh-keygen -f ~/.ssh/google_compute_engine
     

4. מוסיפים משתנה סביבה כדי להגדיר את מספר פרוסות ה-TPU ל-2.

     $ export SLICE_COUNT=2
     

5. יוצרים סביבת ריבוי-פרוסות באמצעות הפקודה queued-resources create או מסוף Google Cloud .

   gcloud

   הערה: ודאו שיש לכם את המכסה המתאימה לפני שאתם בוחרים באפשרות --reserved או --spot, או במכסת ברירת המחדל על פי דרישה. למידע על סוגי המכסות, קראו את מדיניות המכסות.

   הפקודה הבאה מראה איך ליצור TPU מסוג v5e Multislice. כדי להשתמש בגרסה אחרת של TPU, צריך לציין accelerator-type ו-runtime-version אחרים.

   $ gcloud compute tpus queued-resources \
       create ${QR_ID} \
       --accelerator-type=${ACCELERATOR_TYPE} \
       --runtime-version=${RUNTIME_VERSION} \
       --node-count=${SLICE_COUNT} \
       --node-prefix=${TPU_NAME} \
       --zone=${ZONE} \
       [--reserved|--spot]

   המסוף

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

      כניסה לדף TPUs

   2. לוחצים על Create TPU.

   3. בשדה Name, נותנים שם ל-TPU.

   4. בתיבה Zone (אזור), בוחרים את האזור שבו רוצים ליצור את ה-TPU.

   5. בתיבה TPU type, בוחרים סוג מאיץ. סוג המאיץ מציין את הגרסה והגודל של ה-Cloud TPU שרוצים ליצור. האפשרות Multislice נתמכת רק ב-Cloud TPU v4 ובגרסאות TPU חדשות יותר. למידע נוסף על גרסאות TPU, ראו גרסאות TPU.

   6. בתיבה TPU software version, בוחרים גרסת תוכנה. כשיוצרים מכונת TPU וירטואלית ב-Cloud, גרסת התוכנה של ה-TPU מציינת את הגרסה של סביבת זמן הריצה של ה-TPU שצריך להתקין במכונות ה-TPU הווירטואליות. מידע נוסף זמין במאמר בנושא גרסאות תוכנה של TPU.

   7. לוחצים על המתג הפעלת הוספה לתור.

   8. נותנים לבקשה שם בשדה Queued resource name.

   9. לוחצים על תיבת הסימון Make this a Multislice TPU (הגדרת TPU מרובה-פרוסות).

   10. בשדה מספר הפרוסות, מזינים את מספר הפרוסות שרוצים ליצור.

   11. לוחצים על Create כדי ליצור את הבקשה להקצאת משאבים בתור.

6. ההקצאה של המשאב בתור יכולה להימשך עד חמש דקות, בהתאם לגודל המשאב. מחכים עד שהמשאב בתור יהיה במצב ACTIVE. אתם יכולים לבדוק את הסטטוס של בקשה להקצאת משאבים בתור באמצעות ה-CLI של gcloud או מסוף Google Cloud :

   gcloud

   $ gcloud compute tpus queued-resources list \
       --filter=${QR_ID} --zone=${ZONE} --project=${PROJECT}

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

   NAME        ZONE           NODE_COUNT  ACCELERATOR_TYPE  STATE
   ...
   que-res-id  us-central2-b  4           v5litepod-16             ACTIVE
   ...

   המסוף

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

      כניסה לדף TPUs

   2. לוחצים על הכרטיסייה משאבים בתור.

   3. לוחצים על השם של הבקשה להקצאת משאבים בתור.

   אם הסטטוס של קוד ה-QR הוא WAITING_FOR_RESOURCES או PROVISIONING במשך יותר מ- 15 דקות, צריך לפנות Google Cloud לאיש הקשר האחראי לחשבון.

7. מתקינים את יחסי התלות.

   $ python3 multihost_runner.py \
       --TPU_PREFIX=${TPU_NAME} \
       --ZONE=${ZONE} \
       --COMMAND="bash setup.sh"

8. מריצים את shardings.py בכל עובד באמצעות multihost_runner.py.

   $ python3 multihost_runner.py \
       --TPU_PREFIX=${TPU_NAME} \
       --ZONE=${ZONE} \
       --COMMAND="python3 -m pedagogical_examples.shardings \
       --dcn_data_parallelism ${SLICE_COUNT} \
       --ici_fsdp_parallelism 16 \
       --batch_size 131072 \
       --embedding_dimension 2048"

   בקובצי היומן תראו ביצועים של כ-230 TFLOPs לשנייה.

   מידע נוסף על הגדרת מקביליות זמין במאמרים בנושא Multislice sharding using DCN parallelism ו-shardings.py.

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

הרחבת עומס העבודה ל-Multislice

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

• כשיוצרים את הרשת, משתמשים ב-jax.experimental.mesh_utils.create_hybrid_device_mesh במקום ב-jax.experimental.mesh_utils.create_device_mesh.

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

כדי לוודא שלקוד יש גישה לכל המכשירים, אתם יכולים לבדוק אם הערך של len(jax.devices()) שווה למספר הצ'יפים בסביבת ריבוי-הפרוסות. לדוגמה, אם משתמשים בארבעה פלחים של v4-16, יש שמונה שבבים לכל פלח * 4 פלחים, ולכן הפונקציה len(jax.devices()) צריכה להחזיר 32.

בחירת גודלי פרוסות לסביבות עם ריבוי-פרוסות (Multislice)

כדי להאיץ את התהליך באופן לינארי, מוסיפים פרוסות חדשות באותו גודל של הפרוסה הקיימת. לדוגמה, אם משתמשים בפרוסת v4-512, ‏ Multislice תשיג בערך פי שניים ביצועים טובים יותר על ידי הוספה של פרוסת v4-512 שנייה והכפלת גודל האצווה הגלובלי. מידע נוסף זמין במאמר Sharding With Multislice for Maximum Performance.

הפעלת המשימה בכמה פרוסות

יש שלוש גישות שונות להרצת עומס עבודה בהתאמה אישית בסביבה של ריבוי-פרוסות:

1. שימוש בסקריפט להרצת ניסויים, multihost_runner.py
2. באמצעות סקריפט ההרצה של הייצור, multihost_job.py
3. שימוש בגישה ידנית

סקריפט להרצת ניסויים

הסקריפט multihost_runner.py מפיץ קוד לסביבה קיימת של ריבוי-פרוסות, מריץ את הפקודה בכל מארח, מעתיק את היומנים בחזרה ועוקב אחרי סטטוס השגיאה של כל פקודה. הסקריפט multihost_runner.py מתועד ב-MaxText README.

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

סקריפט להרצה בסביבת ייצור

לגבי משימות ייצור שצריכות להיות עמידות בפני כשלים בחומרה ושיבושים אחרים, מומלץ לבצע אינטגרציה ישירות עם ה-API של Create Queued Resource. אפשר להשתמש ב-multihost_job.py כדוגמה פעילה שמפעילה את הקריאה ל-API של Created Queued Resource עם הסקריפט המתאים לטעינה בזמן ההפעלה, כדי להריץ את האימון ולהמשיך אותו אחרי הפקעה. הסקריפט multihost_job.py מתועד בקובץ MaxText README.

מכיוון ש-multihost_job.py צריך להקצות משאבים לכל הרצה, מחזור האיטרציה שלו לא מהיר כמו של multihost_runner.py.

הגישה הידנית

מומלץ להשתמש ב-multihost_runner.py או ב-multihost_job.py, או להתאים אותם, כדי להריץ את עומס העבודה בהתאמה אישית בהגדרת ריבוי-הפרוסות. עם זאת, אם אתם מעדיפים להקצות ולנהל את הסביבה באמצעות פקודות QR ישירות, תוכלו לקרוא את המאמר ניהול סביבה עם ריבוי-פרוסות.

ניהול סביבת Multislice

כדי להקצות ולנהל באופן ידני את ה-QR בלי להשתמש בכלים שסופקו במאגר MaxText, קראו את הקטעים הבאים.

יצירת משאבים בתור

gcloud

1. יוצרים בקשה להקצאת משאבים בתור באמצעות הפקודה הבאה:

   $ gcloud compute tpus queued-resources \
       create ${QR_ID} \
       --project=${PROJECT} \
       --zone=${ZONE} \
       --node-count=${SLICE_COUNT} \
       --accelerator-type=${ACCELERATOR_TYPE} \
       --runtime-version=${RUNTIME_VERSION} \
       --network=${NETWORK_NAME} \
       --subnetwork=${SUBNETWORK_NAME} \
       --tags=${EXAMPLE_TAG_1},${EXAMPLE_TAG_2} \
       --metadata=startup-script="${STARTUP_SCRIPT}" \
       [--reserved|--spot]

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

curl

1. יוצרים קובץ בשם queued-resource-req.json ומעתיקים אליו את קוד ה-JSON הבא.

   {
   "guaranteed": { "reserved": true },
   "tpu": {
       "node_spec": [
       {
       "parent": "projects/your-project-number/locations/your-zone",
           "node": {
           "accelerator_type": "accelerator-type",
           "runtime_version": "tpu-vm-runtime-version",
           "network_config": {
               "network": "your-network-name",
               "subnetwork": "your-subnetwork-name",
               "enable_external_ips": true
           },
           "tags" : ["example-tag-1"]
           "metadata": {
               "startup-script": "your-startup-script"
           }
       },
       "multi_node_params": {
           "node_count": slice-count,
           "node_id_prefix": "your-queued-resource-id"
       }
       }
       ]
   }
   }

   מחליפים את הערכים הבאים:

   • ‫your-project-number: מספר הפרויקט ב- Google Cloud
   • ‫your-zone: האזור שבו רוצים ליצור את המשאב בתור
   • ‫accelerator-type: הגרסה והגודל של פרוסה אחת. התכונה Multislice נתמכת רק ב-Cloud TPU v4 ובגרסאות TPU מאוחרות יותר.
   • ‫tpu-vm-runtime-version: גרסת זמן הריצה של מכונת ה-TPU הווירטואלית שרוצים להשתמש בה.
   • ‫your-network-name: אופציונלי, רשת שאליה יצורף המשאב בתור
   • ‫your-subnetwork-name: אופציונלי, רשת משנה שאליה יצורף המשאב בתור
   • ‫example-tag-1: אופציונלי, מחרוזת תג שרירותית
   • your-startup-script: סקריפט לטעינה בזמן ההפעלה שיפעל כשהמשאב בתור יוקצה
   • ‫slice-count: מספר פרוסות ה-TPU בסביבת ריבוי הפרוסות
   • ‫your-queued-resource-id: המזהה שסופק על ידי המשתמש למשאב בתור

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

   הערה: למרות ש-node_spec הוא רשימה, יש תמיכה רק ברשומה אחת של node_spec.

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

   "guaranteed": { "reserved": true } עם "spot": {}

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

2. שולחים את הבקשה ליצירת משאב שהוכנסה לתור עם מטען ייעודי (payload) בפורמט JSON:

   $ curl -X POST -H "Authorization: Bearer $(gcloud auth print-access-token)" \
   -H "Content-Type: application/json" \
   -d @queuedresourcereq.json \
   https://tpu.googleapis.com/v2alpha1/projects/your-project-id/locations/your-zone/queuedResources\?queued_resource_id\=your-queued-resource-id

   מחליפים את הערכים הבאים:

   • ‫your-project-id: מזהה הפרויקט ב- Google Cloud
   • ‫your-zone: האזור שבו רוצים ליצור את המשאב בתור
   • ‫your-queued-resource-id: המזהה שסופק על ידי המשתמש למשאב בתור

התגובה אמורה להיראות כך:

{
"name": "projects/<your-project-id>/locations/<your-zone>/operations/operation-<your-qr-guid>",
"metadata": {
    "@type": "type.googleapis.com/google.cloud.common.OperationMetadata",
    "createTime": "2023-11-01T00:17:05.742546311Z",
    "target": "projects/<your-project-id>/locations/<your-zone>/queuedResources/<your-qa-id>",
    "verb": "create",
    "cancelRequested": false,
    "apiVersion": "v2alpha1"
},
"done": false
}

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

המסוף

1. נכנסים לדף TPUs במסוף Google Cloud :

   כניסה לדף TPUs

2. לוחצים על Create TPU (יצירת TPU).

3. בשדה Name, נותנים שם ל-TPU.

4. בתיבה Zone (אזור), בוחרים את האזור שבו רוצים ליצור את ה-TPU.

5. בתיבה TPU type, בוחרים סוג של מאיץ. סוג המאיץ מציין את הגרסה והגודל של Cloud TPU שרוצים ליצור. אפשר להשתמש בריבוי-פרוסות רק ב-Cloud TPU v4 ובגרסאות TPU חדשות יותר. מידע נוסף על סוגי המאיצים הנתמכים בכל גרסת TPU זמין במאמר בנושא גרסאות TPU.

6. בתיבה TPU software version, בוחרים גרסת תוכנה. כשיוצרים Cloud TPU VM, גרסת התוכנה של ה-TPU מציינת את גרסת סביבת זמן הריצה של ה-TPU שצריך להתקין. מידע נוסף זמין במאמר בנושא גרסאות תוכנה של TPU.

7. לוחצים על המתג הפעלת הוספה לתור.

8. נותנים לבקשה שם בשדה Queued resource name.

9. לוחצים על תיבת הסימון Make this a Multislice TPU (הגדרת TPU מרובה-פרוסות).

10. בשדה מספר הפרוסות, מזינים את מספר הפרוסות שרוצים ליצור.

11. לוחצים על Create כדי ליצור את הבקשה להקצאת משאבים בתור.

אחזור הסטטוס של משאב בתור

gcloud

$ gcloud compute tpus queued-resources describe ${QR_ID} --zone=${ZONE}

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

...
state:
    state: ACTIVE
...

curl

$ curl -X GET -H "Authorization: Bearer $(gcloud auth print-access-token)" -H "Content-Type: application/json" https://tpu.googleapis.com/v2/projects/your-project-id/locations/your-zone/queuedResources/${YOUR_QR_ID}

אם המשאב בתור נמצא במצב ACTIVE, הפלט נראה כך:

{
"name": your-queued-res,
"tpu": {
    "nodeSpec": [
    {
        ... // node 1
    },
    {
        ... // node 2
    },
    ...
    ]
},
...
"state": "ACTIVE"
}

המסוף

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

   כניסה לדף TPUs

2. לוחצים על הכרטיסייה Queued resources.

3. לוחצים על השם של בקשת המשאב שנמצאת בתור.

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

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

‫SSH והתקנת יחסי תלות

במאמר הרצת קוד JAX בפרוסות TPU מוסבר איך להתחבר למכונות וירטואליות של TPU באמצעות SSH בפרוסה אחת. כדי להתחבר לכל ה-TPU VMs בסביבת ריבוי-הפרוסות באמצעות SSH ולהתקין יחסי תלות, משתמשים בפקודה הבאה של gcloud:

  $ gcloud compute tpus queued-resources ssh ${QR_ID} \
        --zone=${ZONE} \
        --node=all \
        --worker=all \
        --command="command-to-run" \
        --batch-size=4

הפקודה gcloud שולחת את הפקודה שצוינה לכל העובדים והצמתים ב-QR באמצעות SSH. הפקודה נשלחת בו-זמנית בקבוצות של ארבעה. קבוצת הפקודות הבאה נשלחת כשהקבוצה הנוכחית מסיימת את הביצוע. אם אחת מהפקודות נכשלת, העיבוד נפסק ולא נשלחים עוד אצווה. מידע נוסף זמין במאמרי העזרה של ה-API של המשאבים בתור. אם מספר הפרוסות שבהן אתם משתמשים חורג ממגבלת השרשורים של המחשב המקומי (שנקראת גם מגבלת האצווה), תיתקלו במבוי סתום. לדוגמה, נניח שמגבלת האצווה במכונה המקומית היא 64. אם תנסו להריץ סקריפט אימון על יותר מ-64 פרוסות, למשל 100, פקודת ה-SSH תחלק את הפרוסות לקבוצות. הסקריפט יופעל על הקבוצה הראשונה של 64 פרוסות וימתין עד שהסקריפטים יסיימו את הפעולה לפני הפעלת הסקריפט על הקבוצה הנותרת של 36 פרוסות. עם זאת, אי אפשר להשלים את 64 הפרוסות הראשונות עד ש-36 הפרוסות הנותרות יתחילו להריץ את הסקריפט, מה שגורם לקיפאון.

כדי למנוע את התרחיש הזה, אתם יכולים להריץ את סקריפט האימון ברקע בכל VM על ידי הוספת סימן אמפרסנד (&) לפקודת הסקריפט שאתם מציינים באמצעות הדגל --command. כשעושים את זה, אחרי שמפעילים את סקריפט האימון על קבוצת הפרוסות הראשונה, השליטה חוזרת מיד לפקודת ה-SSH. אחרי שהפקודה SSH מסתיימת, אפשר להתחיל להריץ את סקריפט האימון על 36 הפרוסות שנותרו. כשמריצים את הפקודות ברקע, צריך להפנות את הזרמים stdout ו-stderr בצורה מתאימה. כדי להגדיל את המקביליות באותו קוד QR, אפשר לבחור פרוסות ספציפיות באמצעות הפרמטר --node.

הגדרת הרשת

כדי לוודא שפרוסות ה-TPU יכולות לתקשר זו עם זו, מבצעים את השלבים הבאים. מתקינים את JAX בכל אחד מהפלחים. מידע נוסף זמין במאמר הרצת קוד JAX בפרוסות TPU. מוודאים שהערך של len(jax.devices()) שווה למספר הצ'יפים בסביבת ריבוי-הפרוסות. כדי לעשות את זה, מריצים את הפקודה הבאה בכל פרוסה:

  $ python3 -c 'import jax; print(jax.devices())'

אם מריצים את הקוד הזה על ארבע פרוסות של v4-16, יש שמונה צ'יפים לכל פרוסה וארבע פרוסות, אז jax.devices() אמור להחזיר סך של 32 צ'יפים (מכשירים).

הצגת רשימה של המשאבים בתור

gcloud

כדי לראות את המצב של המשאבים בתור, אפשר להשתמש בפקודה queued-resources list:

$ gcloud compute tpus queued-resources list --zone=${ZONE}

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

NAME        ZONE           NODE_COUNT  ACCELERATOR_TYPE  STATE
...
que-res-id  us-central1-a  4           v5litepod-16             ACTIVE
...

המסוף

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

   כניסה לדף TPUs

2. לוחצים על הכרטיסייה Queued resources.

הפעלת המשימה בסביבה שהוקצתה

אתם יכולים להריץ עומסי עבודה באופן ידני על ידי התחברות לכל המארחים בכל פרוסה באמצעות SSH והרצת הפקודה הבאה בכל המארחים.

$ gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
    --zone=${ZONE} \
    --worker=all \
    --command="command-to-run"

איפוס קודי QR

אפשר להשתמש ב-API‏ ResetQueuedResource כדי לאפס את כל המכונות הווירטואליות ב-ACTIVE QR. איפוס המכונות הווירטואליות מוחק את הזיכרון של המכונה ומאפס את המכונה הווירטואלית למצב הראשוני שלה. כל הנתונים שמאוחסנים באופן מקומי יישארו ללא שינוי, וסקריפט ההפעלה יופעל אחרי האיפוס. ‫ResetQueuedResource API יכול להיות שימושי כשרוצים להפעיל מחדש את כל ה-TPU. לדוגמה, כשהאימון נתקע וקל יותר לאפס את כל המכונות הווירטואליות מאשר לנפות באגים.

האיפוסים של כל המכונות הווירטואליות מתבצעים במקביל, ופעולת ResetQueuedResource נמשכת בין דקה לשתי דקות. כדי להפעיל את ה-API, משתמשים בפקודה הבאה:

$ gcloud compute tpus queued-resources reset ${QR_ID} --zone=${ZONE}

מחיקת משאבים שממתינים בתור

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

gcloud

$ gcloud compute tpus queued-resources \
    delete ${QR_ID} --zone=${ZONE} --force [--async]

המסוף

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

   כניסה לדף TPUs

2. לוחצים על הכרטיסייה Queued resources.

3. מסמנים את התיבה לצד בקשת המשאב בתור.

4. לוחצים על מחיקה.

התאוששות אוטומטית מתקלות

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

כדי להמשיך את האימון באופן אוטומטי אחרי הפסקה, צריך לציין סקריפט לטעינה בזמן ההפעלה שבודק וטוען את נקודות הבדיקה האחרונות שנשמרו. הסקריפט לטעינה בזמן ההפעלה פועל באופן אוטומטי בכל פעם שפרוסה מוקצית מחדש או שהמכונה הווירטואלית מאופסת. מציינים סקריפט הפעלה במטען הייעודי (payload) של JSON ששולחים ל-API של בקשת יצירת קוד QR.

הסקריפט הבא לטעינה בזמן ההפעלה (שמשמש ביצירת קודי QR) מאפשר לכם להתאושש אוטומטית מכשלים ולהמשיך את האימון מנקודות ביקורת שמאוחסנות בקטגוריה של Cloud Storage במהלך אימון MaxText:

{
 "tpu": {
   "node_spec": [
     {
      ...
         "metadata": {
               "startup-script": "#! /bin/bash \n pwd \n runuser -l user1 -c 'cd /home/user1/MaxText && python3 -m MaxText.train MaxText/configs/base.yml run_name=run_test_failure_recovery dcn_data_parallelism=4 ici_fsdp_parallelism=8 steps=10000 save_period=10 base_output_directory='gs://user1-us-central2'' EOF"
         }
     ...
     }
   ]
 }
}

לפני שמנסים את זה, צריך לשכפל את מאגר MaxText.

יצירת פרופילים וניפוי באגים

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

אופטימיזציה של האימון

בקטעים הבאים מוסבר איך לבצע אופטימיזציה של אימון Multislice.

חלוקה למקטעים באמצעות Multislice לביצועים מקסימליים

כדי להשיג ביצועים מקסימליים בסביבות של ריבוי-פרוסות, צריך להחליט איך לפצל את הנתונים בין הפרוסות השונות. בדרך כלל יש שלוש אפשרויות (מקביליות נתונים, מקביליות נתונים עם חלוקה מלאה למקטעים ומקביליות צינור עיבוד נתונים). לא מומלץ לפצל הפעלות בין ממדי המודל (לפעמים נקרא tensor parallelism) כי זה דורש יותר מדי רוחב פס בין הפרוסות. בכל האסטרטגיות האלה, אפשר להשתמש באותה אסטרטגיית חלוקה לשברים (sharding) בתוך פרוסה שעבדה בשבילכם בעבר.

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

מתי כדאי להשתמש במקביליות נתונים

מקביליות נתונים טהורה תעבוד טוב במקרים שבהם יש לכם עומס עבודה שפועל היטב, אבל אתם רוצים לשפר את הביצועים שלו על ידי התאמה לעומס (scaling) בכמה פרוסות.

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

כל שבב TPU v4 פועל בשיא של ‎275 * 10¹²‎ פעולות FLOPS בשנייה.

יש ארבעה שבבים לכל מארח TPU, ולכל מארח יש רוחב פס מקסימלי של 50Gbps ברשת.

כלומר, העוצמה האריתמטית היא ‎4 * 275 * 10¹² FLOPS / 50 Gbps = 22000 FLOPS / bit.

המודל ישתמש ב-32 עד 64 ביטים של רוחב פס של DCN לכל פרמטר בכל שלב. אם משתמשים בשני פרוסות, המודל ישתמש ב-32 ביטים של רוחב פס של DCN. אם משתמשים ביותר משני פלחים, הקומפיילר יבצע פעולת צמצום מלאה של ערבוב אקראי, ותשתמשו בעד 64 ביטים של רוחב פס DCN לכל פרמטר לכל שלב. כמות ה-FLOPS שנדרשת לכל פרמטר משתנה בהתאם למודל. באופן ספציפי, במודלי שפה שמבוססים על טרנספורמרים, מספר פעולות ה-FLOPS שנדרשות למעבר קדימה ואחורה הוא בערך ‎6 * B * P, כאשר:

• ‫B הוא גודל האצווה בטוקנים
• ‫P הוא מספר הפרמטרים

מספר ה-FLOPS לכל פרמטר הוא 6 * B, ומספר ה-FLOPS לכל פרמטר במהלך המעבר לאחור הוא 4 * B.

כדי להבטיח שינוי קנה מידה יעיל בכמה פלחים, צריך לוודא שעוצמת הפעולה עולה על העוצמה האריתמטית של חומרת ה-TPU. כדי לחשב את עוצמת הפעולה, מחלקים את מספר ה-FLOPS לכל פרמטר במהלך המעבר לאחור ברוחב הפס של הרשת (בביטים) לכל פרמטר לכל שלב: Operational Intensity = FLOPSbackwards_pass / DCN bandwidth

לכן, במודל שפה שמבוסס על טרנספורמרים, אם משתמשים בשני פלחים: Operational intensity = 4 * B / 32

אם משתמשים ביותר משני פלחים: Operational intensity = 4 * B/64

לכן, גודל האצווה המינימלי המומלץ למודלים של שפה שמבוססים על Transformer הוא בין 176,000 ל-352,000. מכיוון שרשת ה-DCN יכולה להשמיט חבילות לזמן קצר, מומלץ לשמור על מרווח משמעותי לשגיאות, ולפרוס מקביליות נתונים רק אם גודל האצווה לכל Pod הוא לפחות 350,000 (שני Pods) עד 700,000 (הרבה Pods).

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

מתי כדאי להשתמש ב-FSDP

מקביליות נתונים עם חלוקה מלאה (FSDP) משלבת מקביליות נתונים (חלוקת הנתונים בין הצמתים) עם חלוקת המשקלים בין הצמתים. בכל פעולה במעברים קדימה ואחורה, כל המשקלים נאספים כך שלכל פרוסה יש את המשקלים שהיא צריכה. במקום לסנכרן את הגרדיאנטים באמצעות all-reduce, הגרדיאנטים מפוזרים (reduce-scattered) כשהם נוצרים. כך כל פרוסה מקבלת רק את הגרדיאנטים של המשקלים שהיא אחראית עליהם.

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

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

הפעולות all-gather ו-all-reduce ב-FSDP דומות לפעולות ב-DP, כך שאפשר לקבוע אם עומס העבודה של FSDP מוגבל על ידי ביצועי ה-DCN באותה דרך שמתוארת בקטע הקודם.

מתי כדאי להשתמש בהרצה מקבילית של צינורות

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

1. היא גורמת ל'בועת צינור' שבה השבבים לא פעילים כי הם מחכים לנתונים.
2. הוא דורש עיבוד של קבוצות קטנות של נתונים, מה שמקטין את גודל הקבוצה האפקטיבי, את עוצמת החישוב ובסופו של דבר את ניצול ה-FLOP של המודל.

כדאי להשתמש במקביליות צינור עיבוד נתונים רק אם אסטרטגיות המקביליות האחרות דורשות גודל אצווה גלובלי גדול מדי. לפני שמנסים להשתמש ב-pipeline parallelism, כדאי לערוך ניסוי כדי לבדוק באופן אמפירי אם ההתכנסות לכל דגימה מאטה בגודל האצווה שנדרש כדי להשיג ביצועים גבוהים של FSDP. בדרך כלל, FSDP מאפשרת ניצול גבוה יותר של FLOP במודל, אבל אם ההתכנסות לכל דגימה מאטה ככל שגודל האצווה גדל, יכול להיות שעדיין עדיף להשתמש במקביליות של צינורות. רוב עומסי העבודה יכולים להתמודד עם גדלים גדולים מספיק של אצווה, כך שהם לא נהנים מביצוע מקביל של צינורות, אבל יכול להיות שעומס העבודה שלכם שונה.

אם צריך להשתמש במקביליות של צינור עיבוד נתונים, מומלץ לשלב אותה עם מקביליות של נתונים או עם FSDP. כך תוכלו לצמצם את עומק צינור הנתונים ולהגדיל את גודל האצווה לכל צינור נתונים, עד שזמן האחזור של DCN יהפוך לגורם פחות משמעותי ברוחב הפס. לדוגמה, אם יש לכם N פרוסות, כדאי להשתמש בצינורות (pipeline) בעומק 2 וב-N/2 עותקים של מקביליות נתונים, ואז בצינורות בעומק 4 וב-N/4 עותקים של מקביליות נתונים, וכך הלאה, עד שגודל האצווה לכל צינור יהיה מספיק גדול כדי שהפעולות הקולקטיביות של DCN יוכלו להיות מוסתרות מאחורי החישובים במעבר לאחור. כך אפשר לצמצם את ההאטה שנגרמת בגלל מקביליות בצינור עיבוד הנתונים, וגם להגדיל את המודל מעבר למגבלה של גודל האצווה הגלובלי.

שיטות מומלצות לשימוש ב-Multislice

בקטעים הבאים מפורטות שיטות מומלצות לאימון Multislice.

טעינת נתונים

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

Checkpointing

ספריית נקודות הבדיקה Orbax מספקת פרימיטיבים לנקודות בדיקה של JAX PyTrees לאחסון מקומי או לאחסון ב- Google Cloud . אנחנו מספקים שילוב של קובץ עזר עם נקודות ביקורת סינכרוניות ב-MaxText ב-checkpointing.py.

הגדרות נתמכות

בקטעים הבאים מתוארים הצורות, התיאום, המסגרות והמקביליות של הפרוסות שנתמכות ב-Multislice.

צורות

כל הפרוסות חייבות להיות באותה צורה (לדוגמה, אותו AcceleratorType). אין תמיכה בפרוסות בצורות שונות.

תזמור

יש תמיכה בתזמור ב-GKE. למידע נוסף, קראו את המאמר יחידות TPU ב-GKE.

‫Frameworks

אפשר להשתמש בריבוי-פרוסות רק בעומסי עבודה של JAX ו-PyTorch.

מקביליות

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

תמיכה ומשוב

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