קטגוריה: זמן ריצה: כשל בהקצאת מאגר נתונים זמני
השגיאה הזו מציינת שהקצאת הזיכרון של XLA:TPU runtime לא הצליחה למצוא בלוק זיכרון מתאים ב-HBM של המאיץ להקצאה המבוקשת.
הודעת שגיאה לדוגמה:
ValueError: RESOURCE_EXHAUSTED: Error allocating device buffer: Attempting to allocate 8.00M. That was not possible. There are 6.43M free.; (0x0x1_HBM0)
XLA backends: TPU
סקירה כללית
השגיאה הזו מוצגת במקרים הבאים:
- כשלים בהקצאת מאגרים שיזם המשתמש דרך
jax.device_put
או
- כשלים בהקצאות של פלט תוכנית שנקבעה על ידי המשתמש.
הסיבות הנפוצות לכשלים האלה הן:
- חריגה מזיכרון (OOM): המשתמש מנסה להקצות נתח זיכרון שגדול מהכמות הכוללת של הזיכרון הפנוי שזמין ב-HBM של TPU.
- פיצול הזיכרון: ההקצאה נכשלת כי אין בלוק פנוי רציף אחד בנפח הזיכרון שגדול מספיק כדי לעמוד בדרישות הגודל. הכמות הכוללת של הזיכרון הפנוי מספיקה להקצאה, אבל היא מפוזרת במרחב הזיכרון בבלוקים קטנים ולא רציפים.
לסביבת זמן הריצה של TPU יש מספר מנגנונים שמנסים שוב להקצות משאבים במקרה של כשל בהקצאה, כולל:
- אם יש הקצאות שמועברות לתור, זמן הריצה מנסה שוב הקצאות שנכשלו.
- במקרים של שגיאות OOM שנגרמות בגלל פיצול, זמן הריצה יכול להפעיל באופן אוטומטי איחוי וניסיון חוזר.
- זמן הריצה של TPU נותן עדיפות להקצאות של מאגרים על פני שמירה של תוכניות בטעינה. אם הקצאת מאגר נכשלת בגלל חוסר ב-HBM, המערכת תסלק תוכניות TPU שנטענו עד שיהיה מספיק זיכרון למאגר.
לכן, אם נתקלים בשגיאה אחרי שמבצעים את הפעולות שצוינו למעלה, בדרך כלל נדרשת פעולה מצד המשתמש.
ניפוי באגים
- הפחתת השימוש בזיכרון של המודל
- הקטנת גודל האצווה: הקטנת גודל האצווה מקטינה ישירות את השימוש בזיכרון.
- חלוקת פרמטרים: במודלים גדולים מאוד, כדאי להשתמש בטכניקות כמו מקביליות של מודלים או חלוקה כדי להפיץ פרמטרים ב-HBM של כמה ליבות TPU או מארחים.
- קיצור אורך הרצף או חלון ההקשר: במודלים שפועלים על רצפים (כמו מודלים של שפה), קיצור אורך רצף הקלט יכול להקטין באופן משמעותי את הזיכרון שבשימוש.
- תרומת מאגר: שימוש בתכונות של המסגרת (למשל:
jax.jit(..., donate_argnums=...)) כדי לציין ל-XLA שאפשר לדרוס מאגרי קלט מסוימים ולעשות בהם שימוש חוזר לפלט. פרטים נוספים זמינים במאמר בנושא תרומה של זמן אחסון. - אופטימיזציה של אסטרטגיית נקודות ביקורת: במקום לשמור את כל מצב המודל בבת אחת, כדאי לשמור רק את משקלי המודל או להשתמש באסטרטגיית נקודות ביקורת עם חלוקה לשברים.
- פריסת זיכרון של כתובות וריפוד
- הזיכרון של TPU מוקצה בחלקים, וריפוד יכול להגדיל את הגודל בפועל של טנסורים.
- לוודא שאין דליפות זיכרון
- חשוב לוודא שההפניות לאובייקטים מסוג
jax.Arrayלא נשמרות למשך זמן ארוך מהמתוכנן. שמירה של אובייקטים מסוגjax.Arrayעשויה למנוע את ביטול ההקצאה האוטומטי גם אחרי שהקומפילציה של התוכנית מסתיימת.
- חשוב לוודא שההפניות לאובייקטים מסוג
אפשר לעיין גם בקוד השגיאה: E1000 כדי לקבל מידע על אסטרטגיות נוספות שאפשר להשתמש בהן כדי לצמצם את כמות הזיכרון שכל תוכנה צורכת.
כלים
מפעילים את הדגל tpu_log_allocations_on_oom כדי שהקצאת הזיכרון תייצר דוח מפורט של כל ההקצאות הנוכחיות כשמתרחשת שגיאת OOM. הדוח הזה יכול להיות מאוד שימושי לניפוי באגים.