קוד שגיאה: 0101

קטגוריה: זמן ריצה: כשל בהקצאת תוכנית

השגיאה הזו מציינת שסביבת זמן הריצה של XLA במכשיר TPU לא הצליחה לטעון קובץ הפעלה של תוכנית XLA שעברה קומפילציה ל-HBM של ה-TPU.

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

XlaRuntimeError: RESOURCE_EXHAUSTED: Error loading program 'jit_embedding_pipeline_step_fn': Attempting to reserve 29.49G at the bottom of memory. That was not possible. There are 147.64M free, 0B reserved, and 147.64M reservable. Scope: unknown..: while running replica 0 and partition 34 of a replicated computation (other replicas may have failed as well).

‫XLA backends: TPU

סקירה כללית

השגיאה הזו נגרמת בדרך כלל מאחת מהסיבות הבאות:

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

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

ניפוי באגים

• הקטנת הזיכרון שבשימוש במאגר הנתונים הזמני: פינוי הזיכרון שמשמש מאגרי נתונים זמניים יאפשר יותר מקום לתוכנה עצמה:
  • הקטנת גודל האצווה: זו אחת הדרכים היעילות ביותר לצמצם את כמות הזיכרון שמשמשת להפעלות.
  • חלוקת פרמטרים: במודלים גדולים מאוד, כדאי להשתמש במקביליות של מודלים או בטכניקות חלוקה (כמו FSDP או Megascale) כדי לחלק את הפרמטרים והחישובים של המודל בין ליבות TPU או מארחים מרובים.
  • קיצור אורך הרצף או ההקשר: למודלים שמבצעים עיבוד של נתונים רציפים (למשל, מודלים של NLP), צמצום אורך הרצף יכול להפחית באופן משמעותי את השימוש בזיכרון.
  • תרומה למאגר נתונים זמני: שימוש בתכונות של המסגרת (לדוגמה, ‫jax.jit(..., donate_argnums=...)) כדי לאפשר ל-XLA לעשות שימוש חוזר בזיכרון של מאגרי קלט לאחסון פלט, וכך להקטין את השימוש בזיכרון בשיא.
• צריך להקטין את דרישות הזיכרון של התוכנית עבור קבצים זמניים:
  • כדי להקטין את נפח הזיכרון שנדרש לתוכניות בשביל קבצים זמניים, משתמשים בדגל tpu_shared_memory_percent. שימו לב שהדבר עלול להשפיע לרעה על הביצועים.
• אופטימיזציה של אסטרטגיית הביצוע או צמצום עומס הצגת המודעות:
  • ניהול טעינת התוכנית: אם אתם מבצעים קומפילציה של כמה פונקציות בשיטת JIT, חשוב לדעת שכל פונקציה יכולה לגרום לטעינה של תוכנית. כדאי לנסות לבנות את עומס העבודה כך שמספר התוכנות שנטענות בו-זמנית יהיה מינימלי.
• מוודאים שאין דליפות זיכרון:
  • חשוב לוודא שההפניות לאובייקטים מסוג jax.Array לא נשמרות מעבר לזמן הרצוי. החזקה של אובייקטים מסוג jax.Array עשויה למנוע את ביטול ההקצאה האוטומטי גם אחרי שהקומפילציה של התוכנית מסתיימת.

כלים

• מפעילים את הדגל tpu_log_allocations_on_oom כדי שהקצאת הזיכרון תציג דוח מפורט של כל ההקצאות הנוכחיות כשמתרחשת שגיאת OOM. הדוח הזה יכול להיות שימושי מאוד לניפוי באגים.
• יצירת פרופיל של התוכנית: אפשר להשתמש בכלי ליצירת פרופיל של זיכרון JAX או בכלי ליצירת פרופיל של TensorFlow כדי לקבל תצוגה מפורטת של השימוש בזיכרון של התוכנית לאורך זמן. הנתונים האלה יכולים לעזור לזהות שיאים לא צפויים בצריכת הזיכרון.