קוד שגיאה: 1000

קטגוריה: משך הזמן לקימפול: HBM OOM

השגיאה הזו מציינת שהתוכנית דורשת יותר זיכרון HBM (זיכרון ברוחב פס גבוה) ממה שזמין פיזית במכשיר TPU.

הודעות שגיאה לדוגמה:

RESOURCE_EXHAUSTED: XLA:TPU compile permanent error. Ran out of memory in memory space hbm. Used 49.34G of 32.00G hbm. Exceeded hbm capacity by 17.34G.

RESOURCE_EXHAUSTED: TPU TensorCore Hbm usage: 34.82G, SparseCore Hbm usage 174.10G, exceeding available bytes: 95.74G

‫XLA Backends: TPU

סקירה כללית

ה-XLA מבצע בדיקות כדי לוודא שהגודל הכולל של כל ההקצאות הסטטיות הדרושות מתאים ל-HBM של המכשיר.

הקומפיילר מנהל את קיבולת ה-HBM הקבועה של ה-TPU עבור כמה סוגים של הקצאות:

• קלט ופלט של התוכנית: אצוות אימון, מצבי אופטימיזציה וכו'.
• ‫TensorCore + SparseCore Temporaries: זיכרון דינמי שנדרש לחישובים ביניים (למשל, הפעלות, שיפועים וכו').
• קובץ בינארי שעבר קומפילציה: קוד המכונה של TensorCore ‏ (TC) ו-SparseCore ‏(SC).
• תקורה של המערכת: מקום שמור לזמן הריצה של XLA (למשל, מאגרי נתונים זמניים בפיד בדורות ישנים יותר של TPU).
• קבועים: ערכים קבועים שמוטמעים ב-HLO IR מוקצים ב-HBM.
• ‫Compiler Internals: הקצאות ברמת התוכנית וברמת HLO (למשל, routing info לצמתים ברשת)

השגיאה הזו מתרחשת כשמהדר ה-XLA לא יכול להקצות את כל ההקצאות שלמעלה ל-HBM של המכשיר.

ניפוי באגים

צריך לנתח בקפידה את הודעת השגיאה ואת היומנים כדי לקבוע איזו קטגוריה של HBM OOM שמופיעה בהמשך מתארת בצורה הכי טובה את השגיאה:

• TensorCore (TC) + SparseCore (SC) HBM Usage Exceeds Limit: אם השגיאה מפרטת את השימוש, למשל: ‫"TC Hbm usage: X, SC Hbm usage Y". → מעבר אל קטע 1. איזון השימוש ב-HBM של נתוני שקיפות והסכמה ושל נתוני שקיפות ופרטיות.
• הקצאות גדולות באופן לא צפוי: אם השגיאה היא "Ran out of memory in memory space HBM", צריך לבדוק ביומנים את רשימת ההקצאות הגדולות ביותר ב-HBM. במקרה שבו יש טנסורים גדולים באופן לא צפוי (לדוגמה, יותר מ-50% ממגבלת ה-HBM) ← עוברים אל קטע 2. הקצאות גדולות באופן לא צפוי.
• הקצאות מצטברות חורגות ממגבלת HBM: אם השגיאה היא "Ran out of memory in memory space HBM" אבל אין טנסורים גדולים באופן לא צפוי ביומנים → עוברים אל קטע 3. ההקצאות המצטברות חורגות ממגבלת ה-HBM.

חלק 1. איזון השימוש ב-HBM של TC ו-SC

אם השגיאה מפרטת את השימוש, למשל: ‫"TC Hbm usage: X, SC Hbm usage Y" השוואה בין שני הערכים כדי לזהות את צוואר הבקבוק

• שימוש גבוה ב-SparseCore:
  • אופטימיזציה של השימוש ב-HBM Stack: צריכת הזיכרון של HBM Stack גדלה ביחס ל-feature_width, max_unique_nz_per_row ו-logical_replica_count. כדי לצמצם את השימוש המקסימלי במחסנית, אפשר לשנות את הערך של הדגל --xla_sc_num_serialized_tables_to_optimize_hbm, שמבצע סדרתי של עיבוד הטבלאות. החיסרון הוא שרמת המקביליות יורדת.
  • בדיקת התקורה של הריפוד:‏ SparseCore מיישר טבלאות הטמעה ל-32B (8 מספרים ממשיים). טבלאות עם רוחב תכונות קטן (למשל, ‫< 8 floats) incur significant padding overhead, wasting HBM.
  • הפחתת השימוש ב-Heap: ערכים גבוהים של maximum_parallel_iterations מגדילים את כמות נתוני הקלט שנטענים מראש ל-Heap של HBM. הורדת הערך הזה יכולה לפנות כמות משמעותית של זיכרון.
  • אימות של חלוקת נתונים: מוודאים שטבלאות ההטמעה מחולקות בצורה נכונה בין כל השבבים. איך המגבלות מתורגמות לטבלאות
  • במאמר SC: Performance and memory bottlenecks יש רעיונות נוספים.
• שימוש גבוה ב-TensorCore:
  • עוברים אל חלק 2.
• מאוזן
  • אם אף אחד מהם לא חורג מהמגבלה בנפרד אבל הסכום גבוה מדי, הגעתם למגבלת הקיבולת של השבב. צריך לנסות לצמצם את השימוש בשני הרכיבים. מומלץ לפעול לפי ההמלצות בכל שלושת החלקים.

חלק 2. הקצאות גדולות באופן לא צפוי

אם ביומנים מופיעות הקצאות גדולות באופן לא צפוי (מעל 50% ממגבלת ה-HBM), כמעט אף פעם לא מדובר בבעיה בקיבולת החומרה. בדרך כלל מדובר בשגיאת הגדרה. בודקים את התווית XLA (אם היא קיימת) של ההקצאות הגדולות, כדי לקבל רמזים לגבי קוד המקור שלהן ב-JAX.

• הסרת פריטי מידע שנוצרו בתהליך הפיתוח (Artifact) לניפוי באגים:
  • השימוש ב-jax.debug.print() בהרצות בקנה מידה גדול יכול לחייב את הקומפיילר להפוך את הטנסור המלא למוחשי ב-HBM כדי להעביר אותו ל-CPU, מה שגורם לביטול המיזוג ולהגדלת השימוש בזיכרון בשיא. מסירים את כל התגים jax.debug.print() שנותרו.
• תיקון של צורות רשת לא יעילות או של חלוקה לא אחידה:
  • צורות רשת לא תקינות או הערות חסרות של חלוקה לשברים עלולות לגרום לקומפיילר להשתמש בברירת המחדל Replication, וכך הקומפיילר ינסה להתאים טנסורים גדולים מאוד לשבב יחיד
  • בודקים את הצורות של ההקצאות הגדולות ומוודאים שהחלוקה לשברי לוגיקה (sharding) מוגדרת ומופצת בצורה נכונה על ידי XLA.

סעיף 3. הקצאות מצטברות חורגות ממגבלת HBM

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

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

א. בדיקת הריפוד והיישור של טנסור

צורות טנסור לא יעילות הן סיבה נפוצה ושקטה לשגיאות OOM ב-TPU. כדי להשיג ביצועים אופטימליים ב-TPU, ‏ XLA מוסיף אפסים לממדי הטנסור – בדרך כלל למכפלות של 128 עבור הממד הכי קטן ו-8 עבור הממד השני הכי קטן. הריפוד הזה משפיע גם על מערכי הקלט וגם על טנסורים ביניים (HLO זמניים), ויכול להגדיל באופן משמעותי את השימוש בזיכרון, במיוחד עם גדלים קטנים של מימדים. מידע נוסף על פריסות של מערכים

• בדיקת צורות של מאגרי נתונים גדולים: (ב-TPU v5 עם פריסות ברירת מחדל)
  • כשמציבים את הסמן מעל מאגר בXprof Memory Viewer, מוצג כרטיס עם פרטי המאגר, כולל פרטים על הריווח.
  • דוגמה: יכול להיות שצורה של (129, 1024) תרופד ל-(256, 1024), וכך יבוזבז כמעט 50% מהזיכרון.
  • תיקון: לצורה (128, 1024) לא נדרש ריפוד, והיא לא גורמת לבזבוז זיכרון.
• התאמת המימדים: מוודאים שכל המימדים של טנסור גדול (גודל אצווה, מימד הטמעה, גודל נסתר) הם כפולות של 128.

ב. שינוי ההגדרות

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

• הקטנת גודל האצווה: הזיכרון שנדרש להפעלות ולגרדיאנטים ביניים הוא ביחס ישר לגודל האצווה. הקטנת גודל האצווה יכולה לעזור להקטין את השימוש בזיכרון.
• Donate Input Buffers: כשמשתמשים ב-jax.jit, צריך לציין את donate_argnums לפרמטרים של המודל. כך XLA יכולה להחליף את הזיכרון של הקלט בזיכרון של הפלט.
• הפעלה של דיוק מעורב (bfloat16): אם ארכיטקטורת המודל ודרישות האיכות מאפשרות, כדאי להשתמש ב-bfloat16 או בקוונטיזציה (int8 וכו') עבור הטנסורים הגדולים ביותר בתוכנית.

ג. אופטימיזציה של הארכיטקטורה והחלוקה למקטעים

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

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

ד. התאמת זיכרון המפתחות שמשפיע על סימוני XLA:

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

ה. Tune XLA Rematerialization Pass / Manual Checkpointing

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

אפשרות נוספת היא להשתמש ב-decorator‏ jax.checkpoint עם jax.grad כדי לשלוט באופן ידני באילו ערכי ביניים נשמרים במעבר קדימה לעומת אילו ערכים מחושבים מחדש במעבר אחורה, וכך להמיר מחזורי חישוב ל-HBM.

ו. שימוש בכלים מתקדמים ליצירת פרופילים

במאמר Debug OOM errors with XProf יש מדריך לשימוש בXProf Memory Viewer כדי להמחיש את תצוגת הקומפיילר של השימוש ב-HBM.

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