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