אופטימיזציה של הביצועים ב-TPU7x ‏ (Ironwood)

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

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

אימון ברמת דיוק נמוכה באמצעות FP8

‫FP8 (נקודה צפה של 8 ביט) הוא פורמט נתונים מספרי יעיל שמשמש בעיקר להאצת אימון מודלים והסקת מסקנות. באמצעות ייצוג מספרים ב-8 ביטים – במקום בפורמטים סטנדרטיים של 16 ביטים (FP16 או BF16) ו-32 ביטים (FP32) – יחידות TPU יכולות לעבד נתונים הרבה יותר מהר ולהשתמש בפחות זיכרון.

‫TPU7x תומך בהאצת חומרה מובנית לסוגי נתונים של FP8, ומציע ביצועים תיאורטיים מקסימליים של 4,614 TFLOPS לכל שבב. היכולת הזו יכולה לקצר באופן משמעותי את משך האימון מקצה לקצה. בפעולות תואמות, במיוחד בפעולות של כפל מטריצות צפופות שמשותפות לעומסי עבודה של AI, שימוש ב-FP8 יכול להניב שיפורים בביצועים של פי 1.3 לעומת אימון BF16 רגיל. בהשוואה ל-BF16, ‏ FP8 מכפיל את מספר ה-FLOPs המקסימלי ומקטין בחצי את נפח הזיכרון שבשימוש של המשקלים וההפעלות. ‫FP8 צריך להיות כלי מרכזי לכוונון עומסי עבודה שמוגבלים על ידי יכולת החישוב, וגם תרחישים שמוגבלים על ידי קיבולת הזיכרון או רוחב הפס.

השימוש ב-FP8 מציע את היתרונות הבאים מבחינת ביצועים:

• הפחתת העומס על זיכרון ברוחב פס גבוה (HBM): הזיכרון שבשימוש קטן יותר מאפשר למודלים גדולים יותר, או למודלים עם מטמון KV גדול יותר במהלך היקש, להתאים באופן מלא לזיכרון HBM בנפח 192GB. כך נמנעת העברה יקרה לזיכרון מארח איטי יותר.
• הגדלת גודל האצווה האפקטיבי: באמצעות הקטנת הזיכרון שנדרש להפעלות, FP8 מאפשרת שימוש בגדלים גדולים יותר של אצווה. כך משפרים את המקביליות של הנתונים, מה שיכול להוביל לתפוקה גבוהה יותר ולניצול טוב יותר של יחידות החישוב.
• דרישות נמוכות יותר לרוחב פס של הזיכרון: העברת חצי מכמות הנתונים לכל פעולה מצמצמת את הדרישה לנתיב הנתונים מ-HBM ל-MXU. במערכות שבהן תנועת הנתונים היא צוואר בקבוק נפוץ, התכונה הזו עוזרת לשמור על עומס עבודה גבוה ביחידות ה-MXU.

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

• גרנולריות של שינוי קנה מידה: מתחילים עם שינוי קנה מידה לכל טנסור כערך בסיסי. אם יש בעיות באיכות או בביצועים, כדאי לעבור להגדלת קנה מידה לכל ציר. יכול להיות שאין צורך בהתאמת הגודל של ערוץ המשנה.
• מצב שינוי גודל: שינוי גודל דינמי, שמחשב את גורמי שינוי הגודל בזמן הריצה, הוא ברירת מחדל טובה לשמירה על האיכות. הגדלה סטטית יכולה לשפר משמעותית את הביצועים כי היא מבטלת חישובים, אבל היא מחייבת פרופיל מדויק כדי לקבוע את גורמי ההגדלה הנכונים, והיא לא מתאימה לכל תרחישי השימוש, במיוחד כשמשנים את הגדרות המודל. לעומת זאת, מודלים ותצורות חזקים מסוימים יכולים לתקן את קנה המידה למגבלת FP8 עבור משקלים או הפעלות, וכך לאפשר לכם לצמצם את התקורה של הכימות תוך שמירה על הדיוק ושיפור הביצועים.
• פורמטים של FP8‏ (E4M3 ו-E5M2): גישה נפוצה ויעילה היא להשתמש בשילוב של פורמטים של FP8. לדוגמה, אפשר להשתמש ב-E4M3 למשקלים ולהפעלות בהעברה קדימה כדי לנצל את הדיוק הגבוה יותר של E4M3, ולהשתמש ב-E5M2 לשיפועים בהעברה לאחור כדי להתאים לטווח הדינמי הרחב יותר של השיפועים.
• עיגול: שימוש בשיטת העיגול 'עיגול למספר הזוגי הקרוב ביותר' (RNE) במקום עיגול סטוכסטי לגרדיאנטים יכול לשמור על האיכות ולספק ביצועים טובים יותר ושחזור טוב יותר.
• הפעלת FP8 ב-MaxText: ‫MaxText תומך באימון FP8 באמצעות ספריית הכימות QWIX. כדי להפעיל קוונטיזציה, מגדירים את הדגל הבא בהגדרות: use_qwix_quantization=true.

חלוקה למקטעים ומקביליות

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

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

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

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

מסגרת לקבלת החלטות לגבי אסטרטגיית חלוקה

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

• Fully Sharded Data Parallelism (FSDP): זוהי אסטרטגיית ברירת המחדל המועדפת לביצוע מקבילי של נתונים. ‫FSDP מחלק את משקלי המודל, הגרדיאנטים ומצבי האופטימיזציה בין המכשירים המקבילים לנתונים. במהלך החישוב, כל מכשיר מבצע פעולת All-Gather כדי לאחזר את המשקלים המלאים הנדרשים עבור המיקרו-אצווה המקומית שלו. השימוש ב-FSDP יעיל מאוד כל עוד גודל האצווה לכל מכשיר גדול מספיק כדי להסתיר את זמן האחזור של התקשורת הזו מסוג All-Gather. במודלים של Mixture-of-Experts‏ (MoE), החישוב של עוצמת הפעולה האריתמטית צריך להתחשב בדלילות.
• Tensor Parallelism (TP): TP shards individual tensors across devices. בדרך כלל, הטנסורים הם מטריצות משקל בפרספטרון רב-שכבתי (MLP) ובבלוקים של תשומת לב. העוצמה האריתמטית הגבוהה של החומרה (11.5k) מחייבת שהמימדים של המודל יהיו גבוהים מאוד כדי שיהיה אפשר להשתמש ב-TP ב-ICI, וניסיון להשתמש ב-TP עלול לגרום לכך שהמערכת תהיה מוגבלת על ידי התקשורת.
• מקביליות מומחה (EP): זוהי האסטרטגיה הסטנדרטית והנדרשת לאימון מודלים של MoE. ה-EP מחלק את השכבות של ה-expert בין קבוצת מכשירים, ומשתמש בקולקטיב של תקשורת All-to-All כדי לנתב את הטוקנים למכשיר ה-expert הייעודי שלהם. השימוש ב-EP יכול להיות יעיל אם המימד של ה-MLP במודל גדול מספיק כדי להתקרב לקו הגג.
• מקביליות הקשר (CP): CP היא אסטרטגיה מיוחדת שחיונית לאימון מודלים עם רצפים ארוכים מאוד. התפקיד העיקרי שלו הוא לנהל את צריכת הזיכרון של ההפעלות, שגדלה באופן ריבועי עם אורך הרצף ויכולה לחרוג מהקיבולת של HBM. ‫CP מחלק את מימד הרצף של טנסורים של הפעלה, מה שמאפשר שימוש בגודל אצווה חלקי לכל מכשיר. מכיוון ש-CP מוסיף יותר תקשורת מ-FSDP, הכלל הכללי הוא להשתמש בדרגת ה-CP המינימלית שנדרשת כדי לעמוד במגבלות הזיכרון ולהבטיח שה-shard של ציר ה-batch יישאר מספר שלם.

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

סוג עומס העבודה	חלוקה מומלצת של נתונים	חלוקה משנית	צווארי בקבוק מרכזיים	הסבר
מודל צפוף – רצף קצר	FSDP	לא רלוונטי	Rematerialization, FF Matmuls	השיטה FSDP מספקת את האיזון הטוב ביותר. בסדרות קצרות, זיכרון ההפעלה לא אמור להיות בעיה משמעותית. הפתרון הוא להשתמש באצווה גלובלית גדולה מספיק כדי להסתיר את הפעולה All-Gather של FSDP. ככל שגודל האצווה גדל, גודל ההפעלה גדל, ונדרשת מדיניות מתאימה של חומרים מחדש כדי להבטיח שההגדרה הזו לא תגרום לחוסר זיכרון.
מודל צפוף – רצף ארוך	FSDP	CP	Flash attention, activation memory	הזיכרון של ההפעלה הופך לאילוץ העיקרי. נדרש CP כדי להפעיל גדלים חלקיים של אצווה לכל מכשיר ולמנוע בעיות של חוסר זיכרון (OOM). התכונה Flash attention היא המקור העיקרי לחישובים ולבזבוז זמן.
מודל MoE – רצף קצר	FSDP + EP	לא רלוונטי	העברה מכל מקום לכל מקום (ניתוב מומחה), יצירה מחדש של חומרים	מודלים של MoE דורשים EP כדי לפצל את המומחים. התקשורת מסוג All-to-All (העברה מכל אחד לכל אחד) לניתוב טוקנים היא צוואר בקבוק משמעותי שצריך לחפוף אותו. גם תהליך החזרה של חומרים למצב גולמי הוא מקור משמעותי לבזבוז.
מודל MoE – בקנה מידה גדול מאוד	FSDP + EP + PP	מקביליות של מודלים (MP)	כל צווארי הבקבוק שצוינו קודם, וגם בועות בצינור	במודלים שחורגים מהזיכרון של אשכול יחיד, נדרש PP כדי לפצל שכבות בין אשכולות. כך מוצגים תקורה של בועות בצינור ובקשר של DCN. זו הגדרה מורכבת מאוד שדורשת כוונון מדויק.

אופטימיזציה של תקשורת

המנגנון העיקרי לתקשורת ולחישוב חופפים ב-TPU7x נקרא SparseCore Collective Offloading. ארכיטקטורת Ironwood כוללת יחידות SparseCore ייעודיות, שפועלות כשרשורי בקרה עצמאיים שיכולים לנהל את תנועת הנתונים במארג ה-ICI. כך אפשר לבצע פעולות תקשורת קולקטיביות (כמו All-Gather או Reduce-Scatter) במקביל לחישובים העיקריים שמתבצעים ב-TensorCores. זו השיטה המומלצת לקולקטיבים אסינכרוניים ב-TPU7x. כדי להפעיל את העברת הנתונים לקבצים עבור הקולקטיבים הנפוצים ביותר, משתמשים בדגלים המומלצים.

הפעלה מחדש של חומרים

הפעלה מחדש של חומרים, שנקראת גם gradient checkpointing, היא טכניקה בסיסית לצמצום השימוש ב-HBM של מודל. במקום לאחסן את כל ההפעלות הביניים מהמעבר קדימה ב-HBM כדי להשתמש בהן במהלך המעבר אחורה, הוא שומר רק כמה הפעלות מפתח (נקודות ביקורת) ומחשב מחדש את השאר לפי דרישה במהלך המעבר אחורה. השיטה הזו חוסכת כמות משמעותית של זיכרון, אבל היא דורשת יותר חישובים (בערך 25-30% יותר פעולות נקודה צפה לשנייה (FLOPs) עבור בלוק טרנספורמציה רגיל).

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

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

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

שינוי מדיניות השחזור ב-MaxText

‫MaxText מספק שליטה מפורטת על יצירה מחדש של טקסט באמצעות קבוצה של מדיניות מוגדרת מראש ומותאמת אישית, שמוגדרת באמצעות הדגל remat_policy.

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

ב-MaxText יש את כללי המדיניות המובנים הבאים:

• ‫full: המדיניות הכי אגרסיבית, שיוצרת מחדש כמעט הכול. האפשרות הזו מצמצמת את השימוש ב-HBM אבל ממקסמת את התקורה של החישוב מחדש. אידיאלי לתרחישים של רצפים ארוכים עם מגבלות זיכרון חמורות.
• ‫minimal: המדיניות הכי פחות אגרסיבית, ששומרת את רוב ההפעלות. האפשרות הזו ממקסמת את השימוש ב-HBM אבל מצמצמת את החישוב מחדש. הכי מתאים לעומסי עבודה קצרים שמוגבלים על ידי מחשוב, שבהם הזיכרון לא מהווה בעיה.
• מדיניות ביניים: אפשרויות כמו save_dot_with_context_except_mlp,‏ save_qkv_proj ו-save_out_proj מספקות פשרות שונות על ידי יצירת נקודות ביקורת באופן סלקטיבי בפלט של פעולות יקרות של מכפלה סקלרית, תוך יצירה מחדש של פעולות זולות יותר של איבר-איבר.

מדיניות בהתאמה אישית

כדי לקבל רמת שליטה גבוהה יותר, אפשר להגדיר את remat_policy ל-custom. כך אפשר לציין את ההתנהגות של שכבות נפרדות במודול הפענוח של המודל. אפשר להקצות לכל שכבה אחת משלוש התנהגויות:

• ‫device: ההפעלה מאוחסנת ב-HBM במכשיר ה-TPU.
• ‫remat: ההפעלה נפסלת ותופעל מחדש במהלך המעבר לאחור.
• ‫offload: ההפעלה מועברת מ-HBM לזיכרון של מארח ה-CPU, וכך מפנה מקום ב-HBM על חשבון זמן האחזור של ההעברה ב-PCIe.

כוונון של VMEM בהיקף מוגבל

ביצועי הליבה, כמו flash attention, תלויים בגודל האריחים שנבחרו בליבה, שהגודל שלה מוגבל על ידי זיכרון הווקטור (VMEM) הזמין. לכל אחד משני ליבות TensorCore בשבב TPU7x יש 64MiB של זיכרון וקטורי (VMEM). אפשר לחלק את הקיבולת הזו של ה-VMEM בין ההיקף הנוכחי (VMEM בהיקף) לבין אחזור מראש של משקלים עתידיים. הגדלת ה-VMEM בהיקף מוגבל מאפשרת להגדיל את גודל הכפתורים בקרנל, מה שעשוי לצמצם את העיכובים בזיכרון ולשפר את הביצועים של הקרנלים. אפשר לשנות את הגודל של ה-VMEM המוגבל על ידי הגדרת xla_tpu_scoped_vmem_limit_kib (ב-LIBTPU_INIT_ARGS), שאפשר להשתמש בו כדי לבדוק את ביצועי הליבה וגם את מגבלות הביצועים מקצה לקצה. אופטימיזציה של גודל ה-VMEM המוגבל יכולה להשפיע באופן עקיף על הביצועים של ליבת Pallas בהתאמה אישית, כי הגדלת ה-VMEM המוגבל מאפשרת מרחב חיפוש גדול יותר של היפרפרמטרים לגדלים של משבצות בתוך הליבה.

ליבות Tokamax

‫Tokamax, ספרייה של ליבות JAX עם ביצועים גבוהים, כוללת ליבות TPU רבות שעברו אופטימיזציה גבוהה, ומטפלת בכמה צווארי בקבוק נפוצים שספציפיים לחומרה:

• תשומת לב מסוג Splash: תשומת לב מסוג Splash משמשת כיישום העיקרי של תשומת לב כדי לבטל את צוואר הבקבוק של HBM של תשומת לב רגילה, והיא משתמשת ביישום היעיל ביותר של תשומת לב ב-TPU.
• Megablox Grouped Matrix Multiplication (GMM): For MoE workloads, Megablox efficiently handles grouped matrix multiplications by computing over the ragged activations representation. היא ממפה ביעילות את הממד הלא אחיד, ומבצעת כפל מטריצות בין קבוצות לא אחידות של שורות ב-LHS, לבין מטריצת המומחים המתאימה, בלי הצורך להוסיף ריפוד לאצוות לגודל קבוע.
• כוונון אמפירי באמצעות tune-jax: בספריית tune-jax יש כלי עזר לביצוע חיפושים אמפיריים של גדלי בלוקים אופטימליים. גודלי הליבה שמוגדרים כברירת מחדל הם לרוב לא אופטימליים. כוונון מאפשר לבחור גדלים של משבצות VMEM שמתאימים לחומרה כדי למקסם את השימוש בחומרה.
• הערכה של לוגיטים מקסימליים: אפשר לבצע אופטימיזציה נוספת של ליבת הקשב Tokamax Splash על ידי הגדרת ערך ל-max_logit_const. אם המאפיין מוגדר, הוא מחליף את חישוב ההפחתה של הלוגיט המקסימלי במהלך פעולת ה-softmax של מנגנון תשומת הלב (softmax(Q * KT)), וכך מצמצם חלק מהתקורה של החישוב והסנכרון. ב-MaxText, הוא מיושם על ידי ההגדרה config use_max_logits_estimate, שאפשר להגדיר לה את הערך None (מושבת) או ערך של נקודה צפה. כדי למנוע הצפה מספרית, צריך לוודא שטווח הלוגיט של המודל הספציפי שלכם תואם לאומדן. מומלץ לבצע בדיקות של התכנסות אם הערך הזה מוגדר.