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

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

כדי לשפר את הביצועים באמצעות 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 גדולים יותר במהלך ההסקה לתוך 192GiB של HBM. כך נמנעת העברה יקרה לזיכרון מארח איטי יותר.
• גודל אצווה אפקטיבי גדול יותר: הפחתת הזיכרון שנדרש להפעלות מאפשרת להשתמש בגודלי אצווה גדולים יותר ב-FP8. כך משפרים את מקביליות הנתונים, מה שיכול להוביל לשיפור התפוקה ולניצול טוב יותר של יחידות המחשוב.
• דרישות נמוכות יותר לרוחב פס של זיכרון: העברת חצי מכמות הנתונים בכל פעולה מפחיתה את העומס על נתיב הנתונים מ-HBM ל-MXU. במערכות שבהן תנועת הנתונים היא צוואר בקבוק נפוץ, זה עוזר לשמור על יחידות ה-MXU עמוסות בעבודה.

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

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

חלוקה לפלחים (sharding) והקבלה (parallelism)

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

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

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

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

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

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

• ‫Fully Sharded Data Parallelism (FSDP): זו אסטרטגיית ברירת המחדל המועדפת להקבלה של נתונים. ‫FSDP מפצל את משקלי המודל, את הגרדיאנטים ואת מצבי האופטימיזציה בין המכשירים שמקבילים לנתונים. במהלך החישוב, כל מכשיר מבצע פעולת All-Gather כדי לאחזר את המשקלים המלאים שנדרשים ל-microbatch המקומי שלו. ‫FSDP יעיל מאוד כל עוד גודל האצווה לכל מכשיר גדול מספיק כדי להסתיר את זמן האחזור של התקשורת מסוג All-Gather. במודלים של שילוב מומחים (MoE), חישוב העוצמה האריתמטית צריך להביא בחשבון את הדלילות.
• ‫Tensor Parallelism (TP): חלוקה של טנזורים נפרדים בין מכשירים. בדרך כלל, הטנזורים הם מטריצות משקל בפרספטרון רב-שכבתי (MLP) ובבלוקים של תשומת לב. העוצמה האריתמטית הגבוהה של החומרה (11.5k) מטילה דרישה גבוהה מאוד על ממדי המודל כדי להפוך את ה-TP לכדאי ב-ICI, וניסיון להשתמש ב-TP עלול לגרום לכך שהמערכת תהיה מוגבלת לתקשורת.
• ‫Expert Parallelism (EP): האסטרטגיה הסטנדרטית וההכרחית לאימון מודלים של MoE. ב-EP, שכבות ה "מומחה" מפוצלות בין קבוצת מכשירים, וקולקטיב תקשורת All-to-All משמש לניתוב טוקנים למכשיר המומחה הייעודי שלהם. השיטה הזו יכולה להיות יעילה אם ממד ה-MLP של המודל גדול מספיק כדי להתקרב ל-roofline.
• הקבלה לפי הקשר (CP): אסטרטגיה מיוחדת שחיונית לאימון מודלים עם רצפים ארוכים מאוד. התפקיד העיקרי שלו הוא לנהל את צריכת הזיכרון של ההפעלות, שגדלה באופן ריבועי עם אורך הרצף ויכולה לחרוג מהקיבולת של HBM. ‫CP מפצל את מאפיין הרצף של טנזורי ההפעלה, מה שמאפשר להשתמש בגודל אצווה חלקי לכל מכשיר. מכיוון ש-CP כולל יותר תקשורת מאשר FSDP, הכלל הכללי הוא להשתמש בדרגה המינימלית של CP שנדרשת כדי לעמוד במגבלות הזיכרון ולוודא שפיצול ציר ה-batch נשאר מספר שלם.

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

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

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

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

הפעלה מחדש של הפעלה

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

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

ליבות Tokamax

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

• Splash attention: מנגנון Splash attention משמש כהטמעה העיקרית של מנגנון תשומת הלב כדי למנוע את צוואר הבקבוק של HBM במנגנון תשומת הלב הרגיל, והוא משתמש בהטמעה היעילה ביותר של מנגנון תשומת הלב ב-TPU.
• ‫Megablox Grouped Matrix Multiplication (GMM): בעומסי עבודה של MoE,‏ Megablox מטפל ביעילות במכפלות מטריצה מקובצות על ידי חישוב על ייצוג ההפעלות הלא אחידות. הוא ממפה ביעילות את הממד הלא אחיד, מחשב כפל מטריצות בין קבוצות לא אחידות של שורות ב-LHS, ואת מטריצת המומחה המתאימה, וכך נמנע הצורך לרפד קבוצות לגודל קבוע.
• כוונון אמפירי באמצעות tune-jax: בספרייה tune-jax יש כלים לביצוע חיפושים אמפיריים של גודלי בלוקים אופטימליים. לעתים קרובות, גודלי ברירת המחדל של הליבה הם לא אופטימליים. הכוונון מאפשר לבחור גודלי אריחים של VMEM שמתאימים לחומרה, כדי למקסם את השימוש בחומרה.
• הערכת לוגיטים מקסימלית: אפשר לבצע אופטימיזציה נוספת של ליבת הקשב Tokamax Splash על ידי הגדרת ערך ל-max_logit_const. אם הוא מוגדר, הוא מחליף את חישוב ההפחתה של הלוגיט המקסימלי במהלך פעולת ה-softmax של הקשב (softmax(Q * KT)), ומפחית חלק מהתקורה של החישוב והסנכרון. ב-MaxText, הוא מיושם על ידי ההגדרה use_max_logits_estimate, שאפשר להגדיר אותה ל-None (מושבת) או לערך של נקודה צפה (floating point). כדי למנוע גלישה מספרית, ודאו שטווח הלוגיט של המודל הספציפי שלכם תואם לאומדן. אם הערך הזה מוגדר, מומלץ לבצע בדיקת התכנסות.