קטגוריה: משך הזמן לקימפול: שגיאת OOM של Vmem בהיקף מוגבל
השגיאה הזו מציינת שהתוכנית דורשת יותר זיכרון וקטורי בהיקף (Vmem) ממה שהוקצה לה.
הודעות שגיאה לדוגמה:
RESOURCE_EXHAUSTED: Ran out of memory in memory space vmem while allocating on stack for %my-custom-kernel = bf16[2048,4096]{1,0:T(8,128)(2,1)} custom-call(...) ...
XLA Backends: TPU
סקירה כללית
ל-TPU יש זיכרון וקטורי (VMEM) שהוא זיכרון מקומי זמני שמשמש באופן בלעדי את TensorCore (TC). הקומפיילר מנהל את ה-Vmem עבור סוגים שונים של הקצאות:
- הקצאות בהיקף הוראה: אחסון זמני ב-Vmem במהלך ביצוע של הוראת HLO אחת. השיטות האלה כוללות מאגר זמני של טווח אופרנדים (למשל, לשימוש באגירת נתונים כפולה) וזיכרון מטמון של רגיסטרים.
- הקצאות בהיקף התוכנית: הקצאות שחלות על יותר מהוראה אחת של HLO. בדרך כלל מדובר בתוצאות זמניות ותוצאות ביניים של HLO, שהן קלט או פלט של הוראות HLO.
שגיאת חריגה מזיכרון וירטואלי (Vmem) בטווח של זמן קומפילציה מתרחשת כשההקצאות בטווח של ההוראה חורגות ממגבלת ההקצאה של ההוראה הזו. ההגבלה הזו נקבעת
- באופן גלובלי לכל התוכנית באמצעות הדגל --xla_tpu_scoped_vmem_limit_kib
וגם
- לכל ליבת מותאמת אישית באמצעות הפרמטר vmem_limit_bytes.
בדרך כלל השגיאות האלה נגרמות בגלל באג פנימי בקומפיילר או בגלל ליבת מערכת הפעלה בהתאמה אישית שחורגת ממגבלת ההקצאה שלה.
ניפוי באגים
צריך לנתח בקפידה את הודעת השגיאה כדי לזהות אם השגיאה נובעת מליבת לינוקס מותאמת אישית או מ-HLO רגיל. שגיאה שנובעת מליבת מערכת הפעלה מותאמת אישית צריכה לכלול את החתימה הבאה:
Ran out of memory in memory space vmem while allocating on stack for %my-custom-call = <output-shape> custom-call(<params>), custom_call_target="tpu_custom_call" ...
- Custom kernel scoped Vmem OOM: אם השגיאה מצביעה על ליבת מותאמת אישית → עוברים אל Retune the Kernel.
- בעיות ב-Vmem שאינן קשורות לליבת המערכת: אם שגיאת ה-Vmem OOM מתרחשת בגלל פעולה שאינה קשורה לליבת מערכת מותאמת אישית, סביר להניח שמדובר בבאג פנימי בקומפיילר. מומלץ לשלוח דוח על באג ב-XLA עם קובץ dump של HLO.
הגדרה מחדש של הליבה
אם השגיאה נובעת מליבת מערכת הפעלה בהתאמה אישית, אפשר להשתמש בטכניקות הבאות כדי להקטין את דרישת הזיכרון של ליבת מערכת ההפעלה:
- שינוי גודל החסימות: כדי להקטין את השימוש ב-Scoped Vmem, צריך להקטין את גודל החסימות (גודל המשבצות) בהגדרות של ליבת המערכת.
- הגדרת מגבלות של זיכרון וירטואלי לכל ליבת מערכת הפעלה: צריך לבקש במפורש את כמות הזיכרון הנדרשת לליבת מערכת הפעלה ספציפית באמצעות הפרמטר vmem_limit_bytes.
- שינוי הצביעה של הזיכרון: צביעה או הגבלה מפורשת של הקלט/הפלט של ליבת ה-Kernel ל-Vmem באמצעות pallas.tpu.with_memory_space_constraint. חשוב להיזהר ולא לצבוע יותר מדי כניסות ויציאות ל-Vmem, כי זה עלול לגרום ל-Vmem OOM כללי.
- שינוי מגבלת ה-Vmem: אם קשה לבצע כוונון ספציפי לליבה או אם הבעיה משפיעה על הרבה ליבות, אפשר לשנות את מגבלת ה-Vmem הגלובלית באמצעות הדגל --xla_tpu_scoped_vmem_limit_kib.