סקירה כללית על 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 שמרכיבה Multislice. כל פרוסה בבקשה של כמה פרוסות היא מאותו סוג של מאיץ. סוג המאיץ מורכב מסוג 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.
צומת
בהקשר של 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. במסוף Google Cloud , מפעילים את Cloud Shell.

    הפעלת Cloud Shell

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

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

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

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

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

מקביליות של ICI

‫ICI הוא קיצור של Interchip Interconnect, חיבור מהיר בין שבבים שמקשר בין יחידות ה-TPU בפרוסת TPU אחת. חלוקת נתונים (sharding) ב-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