סקירה כללית על Cloud TPU Multislice

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

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

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

זרימת נתונים מרובת-פלחים

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

מושגים

סוג המאיץ
הצורה של כל פרוסת TPU שמרכיבה ריבוי-פרוסות. כל פרוסה בבקשה של ריבוי-פרוסות היא מאותו סוג של מאיץ. סוג של מאיץ מורכב מסוג TPU (גרסה 4 ואילך) ואחריו מספר ליבות TensorCore. לדוגמה, v5litepod-128 מציין TPU v5e עם 128 TensorCores.
תיקון אוטומטי
כשפרוסת TPU נתקלת באירוע תחזוקה, בהפסקה זמנית או בכשל חומרה, Cloud TPU יוצר פרוסת TPU חדשה. אם אין מספיק משאבים ליצירת פרוסה חדשה, היצירה לא תושלם עד שחומרה תהיה זמינה. אחרי שיוצרים את הפרוסה החדשה, כל הפרוסות האחרות בסביבת Multislice מופעלות מחדש כדי שאפשר יהיה להמשיך את האימון. אם סקריפט ההפעלה מוגדר בצורה נכונה, סקריפט האימון יכול להפעיל מחדש באופן אוטומטי את האימון ללא התערבות המשתמש, לטעון את נקודת הבדיקה האחרונה ולהמשיך ממנה.
Data Center Networking (DCN)
רשת עם זמן אחזור ארוך יותר וקצב העברת נתונים נמוך יותר (בהשוואה ל-ICI) שמחברת פרוסות TPU בהגדרה של Multislice.
תזמון של קבוצות
כשכל פרוסות ה-TPU מוקצות יחד, בו-זמנית, כך שמובטח שכולן יוקצו בהצלחה או שאף אחת מהן לא תוקצה.
Interchip interconnect (ICI)
קישורים פנימיים מהירים עם זמן אחזור נמוך שמקשרים בין יחידות TPU בתוך TPU Pod.
Multislice
שני חלקי שבב TPU או יותר שיכולים לתקשר באמצעות DCN.
Node
בהקשר של Multislice, צומת מתייחס לTPU slice יחיד. לכל פרוסת TPU ב-Multislice מוקצה מזהה צומת.
סקריפט לטעינה בזמן ההפעלה
סקריפט הפעלה של 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. In the Google Cloud console, activate Cloud Shell.

    Activate Cloud Shell

    At the bottom of the Google Cloud console, a Cloud Shell session starts and displays a command-line prompt. Cloud Shell is a shell environment with the Google Cloud CLI already installed and with values already set for your current project. It can take a few seconds for the session to initialize.

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

    דוגמה להודעה מתפרצת

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

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

    מקביליות של ICI

    ‫ICI הוא קיצור של Interconnect Chip Interface, ממשק שבבי קישוריות מהיר שמקשר בין יחידות ה-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 וירטואלית באמצעות 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, כמו תגים. למידע נוסף, קראו את המאמר יצירת קודי QR.

      המסוף

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

        כניסה לדף TPUs

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

      3. בשדה שם, מזינים שם ל-TPU.

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

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

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

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

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

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

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

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

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

      $ bash setup.sh

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

      $ 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.

      gcloud

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

      המסוף

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

        כניסה לדף TPUs

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

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

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

    חלוקת נתונים (sharding) למספר חלקים באמצעות מקביליות של 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 וירטואלית. כאן אנחנו משתמשים ב-multihost_runner.py. השלבים הבאים דומים מאוד לאלה שמופיעים במאמר תחילת העבודה: ניסויים מהירים בכמה פלחים ממאגר MaxText, אלא שכאן מריצים את shardings.py במקום את ה-LLM המורכב יותר ב-train.py.

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

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

    ‫multihost_runner.py מתחבר אוטומטית לעובדי TPU באמצעות 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. יוצרים סביבת Multislice באמצעות הפקודה queued-resources create או באמצעות מסוף Google Cloud .

      gcloud

      הפקודה הבאה מראה איך ליצור 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. לוחצים על יצירת TPU.

      3. בשדה שם, מזינים שם ל-TPU.

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

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

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

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

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

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

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

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

    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

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

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

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

    בחירת גודל הפרוסות בסביבות Multislice

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

    הפעלת העבודה בכמה פרוסות

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

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

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

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

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

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

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

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

    גישה ידנית

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

    ניהול סביבת 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 VM שרוצים להשתמש בה.
      • ‫your-network-name: אופציונלי, רשת שאליה יצורף המשאב בתור
      • ‫your-subnetwork-name (אופציונלי): רשת משנה שאליה יצורף המשאב שנוסף לתור
      • ‫example-tag-1: מחרוזת תג שרירותית (אופציונלי)
      • your-startup-script: סקריפט לטעינה בזמן ההפעלה שיפעל כשהמשאב בתור יוקצה
      • ‫slice-count: מספר פרוסות ה-TPU בסביבת Multislice
      • ‫your-queued-resource-id: המזהה שסופק על ידי המשתמש למשאב שנמצא בתור

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

      כדי להשתמש בקיבולת של 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. לוחצים על יצירת TPU.

    3. בשדה שם, מזינים שם ל-TPU.

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

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

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

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

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

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

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

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

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

    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. לוחצים על הכרטיסייה משאבים בתור.

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

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

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

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

    במאמר הרצת קוד JAX בפרוסות TPU מוסבר איך להתחבר למכונות וירטואליות של TPU באמצעות SSH בפרוסה אחת. כדי להתחבר לכל מכונות ה-TPU הווירטואליות בסביבת Multislice באמצעות 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()) שווה למספר הצ'יפים בסביבת Multislice. כדי לעשות את זה, מריצים את הפקודה הבאה בכל פרוסה:

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

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

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

    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. לוחצים על הכרטיסייה משאבים בתור.

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

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

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

    איפוס תגובות מהירות

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

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

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

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

    שחזור אוטומטי של כשלים

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

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

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

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

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

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

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

    כל שבב TPU מדגם v4 מבצע עד ‎275 * 1012 FLOPS בשנייה.

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

    כלומר, העוצמה האריתמטית היא ‎4 * 275 * 1012 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

    אם משתמשים ביותר משני רכיבי Slice: Operational intensity = 4 * B/64

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

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

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

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

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

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

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

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

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

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

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

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

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

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

    טעינת נתונים

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

    Checkpointing

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

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

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

    צורות

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

    תזמור

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

    ‫Frameworks

    התכונה Multislice תומכת רק בעומסי עבודה של JAX ו-PyTorch.

    מקביליות

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

    תמיכה ומשוב

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