कैटगरी: कंपाइल का समय: स्कोप किए गए Vmem OOM
इस गड़बड़ी का मतलब है कि प्रोग्राम को, तय की गई स्कोप की गई वेक्टर मेमोरी (वीएमएम) से ज़्यादा मेमोरी की ज़रूरत है.
गड़बड़ी के मैसेज के उदाहरण:
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 बैकएंड: टीपीयू
खास जानकारी
TPU में वेक्टर मेमोरी (वीएमईएम) होती है. यह एक लोकल स्क्रैचपैड मेमोरी होती है, जिसका इस्तेमाल सिर्फ़ TensorCore (टीसी) करता है. कंपाइलर, अलग-अलग तरह के ऐलोकेशन के लिए Vmem को मैनेज करता है:
- निर्देश के दायरे में आने वाले असाइनमेंट: किसी एक एचएलओ निर्देश को लागू करते समय, Vmem में कुछ समय के लिए स्टोरेज. इसमें ऑपरेंड स्पैन बफ़र (जैसे, डबल बफ़रिंग के लिए) और रजिस्टर स्पिल शामिल हैं.
- प्रोग्राम के स्कोप में आने वाले असाइनमेंट: ऐसे असाइनमेंट जो एक HLO निर्देश के स्कोप से बाहर होते हैं. ये आम तौर पर, एचएलओ के टेंपररी और इंटरमीडिएट नतीजे होते हैं. ये एचएलओ निर्देशों के इनपुट और/या आउटपुट होते हैं.
कंपाइल टाइम स्कोप वाले Vmem OOM की समस्या तब होती है, जब निर्देश के स्कोप वाले असाइनमेंट, उस निर्देश के लिए तय की गई असाइनमेंट की सीमा से ज़्यादा हो जाते हैं. इस सीमा को कंट्रोल किया जाता है
- --xla_tpu_scoped_vmem_limit_kib फ़्लैग का इस्तेमाल करके, पूरे प्रोग्राम के लिए ग्लोबल तौर पर
और
- vmem_limit_bytes पैरामीटर के ज़रिए, हर कस्टम कर्नल के लिए सेट किया जा सकता है.
आम तौर पर, ये गड़बड़ियां कंपाइलर में मौजूद किसी बग या कस्टम कर्नल के लिए तय की गई सीमा से ज़्यादा मेमोरी इस्तेमाल करने की वजह से होती हैं.
डीबग करना
गड़बड़ी के मैसेज का ध्यान से विश्लेषण करें. इससे यह पता चलेगा कि गड़बड़ी कस्टम कर्नल या स्टैंडर्ड एचएलओ की वजह से हुई है या नहीं. कस्टम कर्नल की वजह से होने वाली गड़बड़ी का सिग्नेचर ऐसा होना चाहिए:
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" ...
- कस्टम कर्नेल स्कोप किया गया Vmem OOM: अगर गड़बड़ी कस्टम कर्नेल की वजह से हुई है, तो → कर्नेल को फिर से ट्यून करें पर जाएं.
- नॉन-कर्नल Vmem से जुड़ी समस्याएं: अगर Vmem OOM, कस्टम-कर्नल के अलावा किसी अन्य ओप की वजह से होता है, तो हो सकता है कि यह कंपाइलर से जुड़ी कोई गड़बड़ी हो. कृपया XLA पर गड़बड़ी की रिपोर्ट दर्ज करें. साथ ही, HLO डंप भी शामिल करें.
कर्नेल को फिर से ट्यून करना
अगर गड़बड़ी कस्टम कर्नल की वजह से हो रही है, तो कर्नल के लिए मेमोरी की ज़रूरत को कम करने के लिए, यहां दिए गए तरीकों का इस्तेमाल करें:
- ब्लॉक के साइज़ में बदलाव करें: स्कोप किए गए वीएमईएम के इस्तेमाल को कम करने के लिए, अपने कर्नल कॉन्फ़िगरेशन में ब्लॉक के साइज़ (टाइल के साइज़) को कम करें.
- हर कर्नल के लिए, वीएमईएम की सीमाएं सेट करना: vmem_limit_bytes पैरामीटर का इस्तेमाल करके, उस खास कर्नल के लिए ज़रूरी मेमोरी का अनुरोध करें
- मेमोरी कलरिंग में बदलाव करना: pallas.tpu.with_memory_space_constraint का इस्तेमाल करके, कर्नल के इनपुट/आउटपुट को Vmem के हिसाब से रंग/सीमित करें. Vmem के लिए बहुत ज़्यादा इनपुट/आउटपुट को कलर न करें, क्योंकि इससे Vmem OOM की समस्या हो सकती है.
- वीएमएम की सीमा में बदलाव करें: अगर कर्नल के हिसाब से रीट्यूनिंग करना मुश्किल है या समस्या कई कर्नल पर असर डालती है, तो --xla_tpu_scoped_vmem_limit_kib फ़्लैग का इस्तेमाल करके, ग्लोबल वीएमएम की सीमा में बदलाव किया जा सकता है.