कैटगरी: रनटाइम: बफ़र ऐलोकेशन में गड़बड़ी
इस गड़बड़ी से पता चलता है कि 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ऑब्जेक्ट को बनाए रखने से, प्रोग्राम के कंपाइल होने के बाद भी, मेमोरी अपने-आप खाली नहीं हो सकती.
- पक्का करें कि
टूलिंग
उस tpu_log_allocations_on_oom फ़्लैग को चालू करें जिसके लिए ऐलोकेटर, ओओएम होने पर सभी मौजूदा ऐलोकेशन की पूरी जानकारी वाली रिपोर्ट डंप करेगा. यह रिपोर्ट डीबग करने के लिए बहुत काम की हो सकती है.