קטגוריה: משך הזמן לקימפול: חריגה מזיכרון וירטואלי (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 עם פריקת HLO.
התאמה מחדש של הליבה
אם השגיאה נובעת מליבת מערכת הפעלה בהתאמה אישית, אפשר להשתמש בטכניקות הבאות כדי להקטין את דרישת הזיכרון של ליבת מערכת ההפעלה:
- שינוי גודל החסימות: כדי להקטין את השימוש ב-Scoped Vmem, צריך להקטין את גודל החסימות (גודל המשבצות) בהגדרות של ליבת המערכת.
- הגדרת מגבלות של זיכרון וירטואלי (Vmem) בהיקף של ליבה: אפשר לבקש במפורש את כמות הזיכרון הנדרשת עבור הליבה הספציפית הזו באמצעות הפרמטר vmem_limit_bytes.
- שינוי הצביעה של הזיכרון: צביעה או הגבלה מפורשת של הקלט/פלט של ליבת ה-VMEM באמצעות pallas.tpu.with_memory_space_constraint. אבל חשוב להיזהר ולא לצבוע יותר מדי קלט/פלט ב-Vmem, כי זה עלול לגרום ל-VMEM OOM באופן כללי.
- אם קשה לבצע כוונון מחדש ספציפי לליבת המערכת או אם הבעיה משפיעה על הרבה ליבות, אפשר לשנות את מגבלת ה-Vmem הגלובלית באמצעות הדגל --xla_tpu_scoped_vmem_limit_kib.