ניפוי באגים בשגיאות OOM באמצעות XProf

שגיאות Out Of Memory (OOM) מתרחשות כשנפח הזיכרון (HBM) של המאיץ (GPU או TPU) מוצה. במאמרים E1000 – Compile Time HBM OOM documentation ו-JAX documentation on GPU memory allocation מפורטות כמה סיבות נפוצות לבעיות שקשורות ל-OOM וטכניקות לניפוי באגים.

בדף הזה מוסבר איך להשתמש בכלי Memory Viewer של XProf כדי להציג באופן חזותי את השימוש בזיכרון של תוכנית JAX, לזהות מקרים של שימוש שיא ולנפות באגים בשגיאות OOM. התהליך כולל את השלבים הבאים:

  1. מריצים את התוכנית עם jax.profiler.trace כדי לצלם את הפרופיל.
  2. מריצים את XProf ברקע ומשתמשים בכלי לבדיקת הזיכרון כדי לראות את פרטי השימוש בזיכרון.

תוכנית לדוגמה

התוכנית הבאה של JAX מובילה לשגיאת OOM:

‫
import jax
from jax import random
import jax.numpy as jnp


@jax.profiler.trace("/tmp/xprof")
@jax.jit
def oom():
    a = random.normal(random.PRNGKey(1), (327680, 327680), dtype=jnp.bfloat16)
    return a @ a


if __name__ == "__main__":
    oom()

במכונת TPU, התוכנית הזו נכשלת עם השגיאה:

XlaRuntimeError: RESOURCE_EXHAUSTED: Allocation (size=107374182400) would exceed memory (size=17179869184) :: #allocation7 [shape = 'u8[327680,327680]{1,0:T(8,128)(4,1)}', space=hbm, size = 0xffffffffffffffff, tag = 'output of xor_convert_fusion@{}'] :: <no-hlo-instruction>

(במחשב עם GPU, השגיאה נראית כך: XlaRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 214748364800 bytes.)

הפעלת XProf

מתקינים את xprof (pip install xprof) ומפעילים מכונת XProf עם ציון הספרייה שבה הפרופיל מאוחסן:

xprof --logdir=/tmp/xprof/ --port=6006

עוברים למופע (במחשב מקומי, בכתובת http://localhost:6006). בתפריט הנפתח Tools, בוחרים באפשרות Memory Viewer, ובחלון הכלי Memory Viewer, בוחרים באפשרות HBM בתפריט הנפתח Memory Types (בדרך כלל האפשרות הזו נבחרת כברירת מחדל).

הדף XProf Memory Viewer עבור תוכנית הדוגמה שלמעלה

במסמכי הכלי XProf: Memory Viewer מפורטים רכיבי הכלי והמידע שמוצג בו.

מתמקדים בקטע HLO Ops at Peak Memory Allocation שבו מוצגים שלושה תרשימי מאגר בנקודת השימוש המקסימלית בזיכרון. המאגר כולל: * קלט ופלט של תוכניות: אצוות אימון, מצבי אופטימיזציה וכו'. * נתונים זמניים של TensorCore ו-SparseCore: זיכרון דינמי שנדרש לחישובים ביניים (כמו הפעלות, שיפועים וכו').

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

טכניקות ספציפיות לניפוי באגים מפורטות במאמר E1000: ניפוי באגים.