קטגוריה: משך הזמן לקימפול: HBM OOM
השגיאה הזו מציינת שהתוכנית דורשת יותר זיכרון HBM (זיכרון ברוחב פס גבוה) ממה שזמין פיזית במכשיר TPU.
הודעות שגיאה לדוגמה:
RESOURCE_EXHAUSTED: TPU TensorCore Hbm usage: 34.82G, SparseCore Hbm usage 174.10G, exceeding available bytes: 95.74G
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.
XLA backends: TPU
סקירה כללית
ה-XLA מבצע בדיקות כדי לוודא שהגודל הכולל של כל ההקצאות הסטטיות הדרושות מתאים ל-HBM של המכשיר.
הקומפיילר מנהל את קיבולת ה-HBM הקבועה של ה-TPU עבור כמה סוגים של הקצאות:
- קלט ופלט של התוכנית: אצוות אימון, מצבי אופטימיזציה וכו'.
- TPU temporaries: זיכרון דינמי שנדרש לחישובים ביניים (למשל, הפעלות, שיפועים וכו').
- קובץ בינארי שעבר קומפילציה: שפת המכונה עבור TensorCore (TC) ו-SparseCore (SC).
- תקורה של המערכת: מקום שמור לזמן הריצה של XLA (למשל, מאגרי נתונים של פידים בדורות ישנים יותר של TPU).
- קבועים: ערכים קבועים שמוטמעים ב-IR של HLO מוקצים ב-HBM.
- רכיבים פנימיים של קומפיילר: הקצאות ברמת התוכנית וברמת HLO (למשל, routing info לצמתים ברשת).
השגיאה הזו מתרחשת כשמהדר ה-XLA לא יכול להקצות את כל ההקצאות שלמעלה ל-HBM של המכשיר.
ניפוי באגים
צריך לנתח בקפידה את הודעת השגיאה ואת היומנים כדי לקבוע איזו קטגוריה של HBM OOM שמופיעה בהמשך מתארת בצורה הכי טובה את השגיאה:
- "TC Hbm usage: X, SC Hbm usage Y": אם השגיאה מפרטת את השימוש ב-HBM, סימן שהשימוש המצטבר ב-TensorCore (TC) + SparseCore (SC) חורג מהמגבלה של HBM. → מעבר אל תרחיש 1. איזון בין השימוש ב-TC לבין השימוש ב-SC ב-HBM.
- "Ran out of memory in memory space HBM" (נגמר הזיכרון במרחב הזיכרון HBM): בודקים את היומנים כדי לראות את ההקצאות הגדולות ביותר ב-HBM.
- במקרה שקיימים טנסורים גדולים באופן לא צפוי (לדוגמה, > 50% ממגבלת ה-HBM) ← עוברים אל תרחיש 2. נגמר הזיכרון בגלל הקצאות גדולות באופן לא צפוי.
- אם אין טנסורים גדולים באופן לא צפוי ביומנים → עוברים אל תרחיש 3. נגמר הזיכרון בגלל הקצאות מצטברות.
תרחיש 1. איזון השימוש ב-HBM של TC ו-SC
אם השגיאה מפרטת את השימוש, למשל "TC Hbm usage: X, SC Hbm usage Y", המשמעות היא שהשימוש המצטבר ב-TensorCore (TC) + SparseCore (SC) חורג מהמגבלה של HBM. משווים בין שני הערכים כדי לזהות את צוואר הבקבוק:
- שימוש גבוה ב-SparseCore
- אופטימיזציה של השימוש במערך HBM: צריכת הזיכרון של מערך HBM גדלה עם
feature_width,max_unique_nz_per_rowו-logical_replica_count. אפשר להפחית את השימוש המקסימלי בסטאק על ידי כוונון הדגל--xla_sc_num_serialized_tables_to_optimize_hbm, שגורם לסריאליזציה של עיבוד הטבלאות. החיסרון הוא שרמת המקביליות יורדת. - בדיקת התקורה של הריפוד: SparseCore מיישר טבלאות הטמעה ל-32B (8 מספרים ממשיים). טבלאות עם רוחבי תכונות קטנים (למשל, פחות מ-8 ערכים מסוג float) גורמות לתקורה משמעותית של ריווח פנימי, ולבזבוז של HBM.
- הפחתת השימוש ב-heap: ערכים גבוהים של
maximum_parallel_iterationsמגדילים את כמות נתוני הקלט שנטענים מראש ל-heap של HBM. הורדת הערך הזה יכולה לפנות כמות משמעותית של זיכרון. - אימות של חלוקת נתונים: מוודאים שטבלאות ההטמעה מחולקות בצורה נכונה בין כל השבבים. איך המגבלות מתורגמות לטבלאות
- במאמר SC: צווארי בקבוק בביצועים ובזיכרון יש רעיונות נוספים.
- אופטימיזציה של השימוש במערך HBM: צריכת הזיכרון של מערך HBM גדלה עם
- שימוש גבוה ב-TensorCore
- ממשיכים לתרחיש 2.
- מאוזן
- אם אף אחד מהם לא גבוה מדי בנפרד, אבל הסכום שלהם גבוה מדי, הגעתם לקיבולת של השבב. צריך לנסות לצמצם את השימוש בשני הרכיבים. חשוב לפעול לפי ההמלצות בכל שלושת החלקים.
תרחיש 2. אין מספיק זיכרון בגלל הקצאות גדולות באופן לא צפוי
אם מופיעה הודעת השגיאה "Ran out of memory in memory space HBM" (הזיכרון אזל במרחב הזיכרון HBM) ואחת או יותר הקצאות גדולות באופן לא צפוי מופיעות ביומנים (> 50% ממגבלת ה-HBM), כמעט אף פעם לא מדובר בבעיה בקיבולת החומרה. בדרך כלל מדובר בשגיאת הגדרה. בודקים את התווית XLA (אם היא קיימת) של ההקצאות הגדולות כדי לקבל רמזים לגבי קוד המקור שלהן ב-JAX.
- הסרת ארטיפקטים של ניפוי באגים
- שימוש ב-jax.debug.print() בהרצות בקנה מידה גדול עלול לגרום לקומפיילר לממש את הטנסור המלא ב-HBM כדי להעביר אותו למעבד, מה שיגרום להפרעה במיזוג ולהגדלת השימוש בזיכרון בשיא. מסירים את כל התגים
jax.debug.print()שנותרו.
- שימוש ב-jax.debug.print() בהרצות בקנה מידה גדול עלול לגרום לקומפיילר לממש את הטנסור המלא ב-HBM כדי להעביר אותו למעבד, מה שיגרום להפרעה במיזוג ולהגדלת השימוש בזיכרון בשיא. מסירים את כל התגים
- תיקון של צורות רשת לא יעילות או של חלוקה למקטעים
- צורות לא נכונות של רשתות או הערות חסרות של חלוקה לשברים עלולות לגרום להידור להשתמש בשכפול כברירת מחדל – מה שיגרום להידור לנסות להתאים טנסורים גדולים מאוד לשבב יחיד.
- בודקים את הצורות של ההקצאות הגדולות ומוודאים ש-XLA מציין ומפיץ את ה-sharding בצורה נכונה.
תרחיש 3. אין זיכרון פנוי בגלל הקצאות מצטברות
אם מופיעה הודעת השגיאה "Ran out of memory in memory space HBM" ולא מופיעים ביומנים טנסורים גדולים באופן לא צפוי, התוכנית חורגת מהקיבולת בגלל שהסכום הכולל של ההקצאות חורג מהמגבלה של HBM. במקרה כזה, כדאי להציג את פרופיל הזיכרון כדי לזהות את המאגרים הספציפיים שגורמים לשימוש המקסימלי. במאמר ניפוי באגים של שגיאות OOM באמצעות XProf מופיע מדריך מפורט לזיהוי הגורמים העיקריים לשימוש בזיכרון.
אחרי שמזהים חלק מהגורמים העיקריים לשימוש בזיכרון, אפשר לפעול לפי השלבים הבאים כדי לצמצם את הזיכרון שבשימוש.
תרחיש 3.א: שינוי ההגדרה
לרוב אפשר לפתור בעיות של OOM באמצעות שינויים בהגדרות האלה:
- הקטנת גודל האצווה: הזיכרון שנדרש להפעלות ולשיפועים ביניים הוא ביחס ישר לגודל האצווה. הקטנת גודל האצווה יכולה לעזור להפחית את השימוש בזיכרון.
- מאגרי קלט לתרומות: כשמשתמשים ב-
jax.jit, צריך לציין את donate_argnums לפרמטרים של המודל. כך XLA יכול להחליף את הזיכרון של הקלט בזיכרון של הפלט. - הפעלת דיוק מעורב (bfloat16): אם ארכיטקטורת המודל ודרישות האיכות מאפשרות, כדאי להשתמש ב-bfloat16 או בקוונטיזציה (int8 וכו') עבור הטנסורים הגדולים ביותר בתוכנית. חשוב לשים לב שהשינוי הזה יכול להשפיע על התנהגות המודל, ולכן צריך לשקול אותו בקפידה.
תרחיש 3.ב': אופטימיזציה של הארכיטקטורה והחלוקה
אם שינויים בהגדרות לא מספיקים, יכול להיות שהטופולוגיה של המודל גדולה מדי בשביל הגדרת החומרה הנוכחית.
- שימוש בדורות חדשים יותר של TPU: בדרך כלל, דורות חדשים יותר של TPU מציעים יותר HBM לכל שבב. אם יש לכם אפשרות, כדאי לעבור לדורות חדשים יותר של TPU.
- הפעלה בטופולוגיית שבבים גדולה יותר: אם משקלי המודל גדולים מדי עבור הטופולוגיה הקיימת, אפשר לנסות לפצל אותם בין יותר שבבים.
- הטמעה של טכניקות מתקדמות של חלוקה למקטעים:
- אפשר לעיין בגישות מתקדמות יותר של מקביליות נתונים, טנסורים או צינורות.
- מציינים רמזים לפיצול לערכי ביניים ולפלט.
- שימוש ב-JAX host offloading: טכניקות של host offloading מאפשרות למשתמש להעביר טנסורים גדולים לזיכרון המעבד המרכזי של המארח (לדוגמה, activation offloading ו-optimizer state offloading).
תרחיש 3.C בדיקת ריפוד ויישור של טנסורים
צורות טנסור לא יעילות הן סיבה נפוצה ושקטה לשגיאות OOM ב-TPU. כדי להשיג ביצועים אופטימליים ב-TPU, XLA מוסיף אפסים לממדי הטנסור – בדרך כלל למכפלות של 128 עבור הממד הכי קטן ו-8 עבור הממד השני הכי קטן. הריפוד הזה משפיע גם על מערכי הקלט וגם על טנסורים ביניים (HLO זמניים), ויכול להגדיל באופן משמעותי את השימוש בזיכרון, במיוחד עם גדלים קטנים של מימדים. מידע נוסף על פריסות של מערכים
- בדיקת צורות של מאגרי נתונים גדולים: (ב-TPU v5 עם פריסות ברירת מחדל)
- כשמעבירים את העכבר מעל מאגר ב-Xprof Memory Viewer, מוצג כרטיס עם פרטי המאגר, כולל פרטים על הריווח.
- דוגמה: יכול להיות שצורה של
(129, 1024)תרופד ל-(256, 1024), וכתוצאה מכך כמעט 50% מהזיכרון יבוזבז. - תיקון: לצורה
(128, 1024)לא נדרש ריפוד, והיא לא גורמת לבזבוז זיכרון.
- התאמת המימדים: מוודאים שכל המימדים של טנסור גדול (גודל אצווה, מימד הטמעה, גודל נסתר) הם כפולות של 128. חשוב לזכור שהשינוי הזה יכול להשפיע על התנהגות המודל, ולכן צריך לשקול אותו בקפידה.
תרחיש 3.D Tune key memory impacting XLA flags
אפשר לשנות את הערכים של דגלים מרכזיים של זיכרון כדי לשפר את הביצועים על חשבון שימוש נמוך יותר בזיכרון. עם זאת, מומלץ להשתמש באסטרטגיה הזו רק כמוצא אחרון, כי היא עלולה להשפיע לרעה על הביצועים.
תרחיש 3.E: כוונון של XLA rematerialization pass/manual checkpointing
אם המודל קרוב להיכנס לזיכרון, אפשר להשתמש במעצב jax.checkpoint עם jax.grad כדי לשלוט באופן ידני באילו נתונים ביניים נשמרים בהעברה קדימה לעומת חישוב מחדש בהעברה אחורה, וכך להחליף מחזורי חישוב ב-HBM.
אפשרות אחרת היא להכריח את המעבר XLA::Rematerialization כדי לתעדף חיסכון בזיכרון, אבל יכול להיות שהחיסכון הזה יבוא על חשבון מהירות ההידור:
| דגל | תיאור | השפעה / פשרה |
|---|---|---|
--xla_tpu_max_hbm_size_mib |
הגדרה ידנית של המגבלה על גודל ה-HBM שמשמש למעבר Rematerialization. | האפשרות הזו גורמת לקומפיילר לעבוד קשה יותר כדי להתאים את התוכנית למגבלה קטנה יותר מ-HBM הפיזי בפועל. |
--xla_tpu_rematerialization_algo=PEAK_PRIORITY |
מתמקדים בנקודות שבהן השימוש בזיכרון מגיע לשיא. | יכול להיות יעיל יותר מהאלגוריתם שמוגדר כברירת מחדל לצמצום משמעותי של הזיכרון. |
--xla_tpu_rematerialization_max_block_size_limit=32 |
ההגדרה הזו קובעת את המספר המקסימלי של הוראות בבלוק שאפשר ליצור מחדש בבת אחת. | הגדלת הערך הזה מאפשרת חיסכון בזיכרון, אבל משך הזמן לקימפול מתארך משמעותית. |
--xla_tpu_rematerialization_block_effort_factor=10.0 |
הגדרה של כמות המאמץ (משך הזמן לקימפול) שמושקעת בחיפוש בלוקים כדי ליצור אותם מחדש. | ערכים גבוהים יותר מאפשרים חיפוש מקיף יותר של חיסכון בזיכרון, אבל זמני ההידור מתארכים. |
--xla_tpu_pre_fusion_remat=true |
האפשרות הזו מאפשרת להפעיל שלב נוסף של יצירה מחדש של חומרים לפני שלב המיזוג. | אפשר לחסוך יותר זיכרון, אבל זמני ההידור מתארכים וייתכן שתהיה השפעה על היציבות המספרית. |
שימו לב: שינויים בדגלים של XLA צריכים להיעשות כמוצא אחרון, כי הם עלולים להשפיע לרעה על הביצועים.
תרחיש 3.F שימוש בכלים מתקדמים ליצירת פרופילים
במאמר Debug OOM errors with XProf יש מדריך לשימוש בXProf Memory Viewer כדי להמחיש את תצוגת הקומפיילר של השימוש ב-HBM.
הכלי הזה מאפשר לכם לראות את הקצאת הזיכרון המקסימלית ואת משך החיים של המאגר, וזה חיוני כדי להבין בדיוק מה צורך את ה-HBM בנקודת השימוש המקסימלית. למידע על הגדרת פרופילים באופן כללי, אפשר לעיין במאמרים תחילת העבודה עם Xprof וTensorBoard Profiling.