कैटगरी: रनटाइम: बफ़र ऐलोकेशन में गड़बड़ी
इस गड़बड़ी से पता चलता है कि XLA:TPU रनटाइम का मेमोरी ऐलोकेटर, अनुरोध किए गए ऐलोकेशन के लिए, ऐक्सलरेटर के एचबीएम पर मेमोरी का सही ब्लॉक नहीं ढूंढ सका.
गड़बड़ी के मैसेज का उदाहरण:
ValueError: RESOURCE_EXHAUSTED: Error allocating device buffer: Attempting to allocate 8.00M. That was not possible. There are 6.43M free.; (0x0x1_HBM0)
XLA बैकएंड: टीपीयू
खास जानकारी
यह गड़बड़ी इन पर दिखती है:
- उपयोगकर्ता की ओर से शुरू किए गए बफ़र के असाइनमेंट में आने वाली समस्याएं. ये समस्याएं
jax.device_putके ज़रिए होती हैं
या
- उपयोगकर्ता के शेड्यूल किए गए प्रोग्राम के आउटपुट के लिए, मेमोरी का बंटवारा नहीं हो सका.
आम तौर पर, ऐसा इन वजहों से होता है:
- आउट ऑफ़ मेमोरी (OOM): उपयोगकर्ता, मेमोरी का ऐसा हिस्सा असाइन करने की कोशिश कर रहा है जो टीपीयू के एचबीएम पर उपलब्ध कुल खाली मेमोरी से बड़ा है.
- मेमोरी फ़्रैगमेंटेशन: मेमोरी स्पेस में कोई भी सिंगल कंटीग्यूअस फ़्री ब्लॉक, अनुरोध किए गए साइज़ के हिसाब से बड़ा नहीं है. इसलिए, मेमोरी स्पेस को असाइन नहीं किया जा सका. कुल खाली मेमोरी, मेमोरी के लिए काफ़ी है. हालांकि, यह मेमोरी स्पेस में छोटे-छोटे, अलग-अलग ब्लॉक में बंटी हुई है.
टीपीयू रनटाइम में, टीपीयू को असाइन करने में होने वाली गड़बड़ियों को ठीक करने के लिए कई तरीके मौजूद हैं. इनमें ये शामिल हैं:
- अगर डी-अलॉकेट करने के लिए कोई अनुरोध लाइन में लगा है, तो रनटाइम, अलॉट नहीं किए जा सके अनुरोधों को फिर से प्रोसेस करने की कोशिश करता है.
- फ़्रैगमेंटेशन की वजह से होने वाली ओओएम की समस्याओं के लिए, रनटाइम अपने-आप डीफ़्रैगमेंटेशन और फिर से कोशिश करने की प्रोसेस को ट्रिगर कर सकता है.
- टीपीयू रनटाइम, प्रोग्राम को लोड रखने के बजाय बफ़र के लिए मेमोरी को प्राथमिकता देता है. अगर एचबीएम की कमी की वजह से बफ़र असाइन नहीं किया जा सका, तो सिस्टम लोड किए गए टीपीयू प्रोग्राम को तब तक हटाता रहेगा, जब तक बफ़र के लिए ज़रूरी मेमोरी उपलब्ध नहीं हो जाती.
इसलिए, ऊपर दिए गए तरीकों को आज़माने के बाद भी अगर कोई गड़बड़ी होती है, तो आम तौर पर उपयोगकर्ता को कार्रवाई करनी पड़ती है.
डीबग करना
- अपने मॉडल के मेमोरी फ़ुटप्रिंट को कम करना
- बैच का साइज़ कम करें: बैच का साइज़ कम करने से, मेमोरी का इस्तेमाल सीधे तौर पर कम हो जाता है.
- पैरामीटर शार्डिंग: बहुत बड़े मॉडल के लिए, मॉडल पैरललिज़्म या शार्डिंग जैसी तकनीकों का इस्तेमाल करें. इससे, पैरामीटर को कई टीपीयू कोर या होस्ट के एचबीएम में डिस्ट्रिब्यूट किया जा सकता है.
- सीक्वेंस/कॉन्टेक्स्ट की लंबाई कम करें: सीक्वेंस पर काम करने वाले मॉडल (जैसे, भाषा मॉडल) के लिए, इनपुट सीक्वेंस की लंबाई कम करने से मेमोरी फ़ुटप्रिंट में काफ़ी कमी आ सकती है.
- बफ़र डोनेशन: फ़्रेमवर्क की सुविधाओं (जैसे:
jax.jit(..., donate_argnums=...)) का इस्तेमाल करके, XLA को यह सिग्नल दें कि कुछ इनपुट बफ़र को फिर से लिखा जा सकता है और आउटपुट के लिए फिर से इस्तेमाल किया जा सकता है. ज़्यादा जानकारी के लिए, बफ़र डोनेशन पढ़ें. - चेकपॉइंट की रणनीति को ऑप्टिमाइज़ करें: पूरे मॉडल की स्थिति को एक साथ सेव करने के बजाय, सिर्फ़ मॉडल के वेट सेव करें या चेकपॉइंट की शार्ड की गई रणनीति का इस्तेमाल करें.
- मेमोरी लेआउट और पैडिंग को मैनेज करना
- टीपीयू मेमोरी को हिस्सों में बांटा जाता है. साथ ही, पैडिंग से टेंसर का साइज़ बढ़ सकता है.
- पक्का करें कि कोई मेमोरी लीक न हो
- पक्का करें कि
jax.Arrayऑब्जेक्ट के रेफ़रंस को तय समय से ज़्यादा समय तक न रोका गया हो.jax.Arrayऑब्जेक्ट को बनाए रखने से, प्रोग्राम के कंपाइल होने के बाद भी, मेमोरी अपने-आप खाली नहीं होती.
- पक्का करें कि
हर प्रोग्राम के लिए इस्तेमाल होने वाली मेमोरी की मात्रा को कम करने के लिए, गड़बड़ी कोड: E1000 देखें.
टूलिंग
उस tpu_log_allocations_on_oom फ़्लैग को चालू करें जिसके लिए ऐलोकेटर, ओओएम होने पर सभी मौजूदा ऐलोकेशन की पूरी जानकारी वाली रिपोर्ट डंप करेगा. यह रिपोर्ट डीबग करने के लिए बहुत काम की हो सकती है.