कैटगरी: रनटाइम: प्रोग्राम असाइन करने में गड़बड़ी
इस गड़बड़ी से पता चलता है कि टीपीयू डिवाइस पर XLA रनटाइम, कंपाइल किए गए XLA प्रोग्राम के एक्ज़ीक्यूटेबल को टीपीयू के एचबीएम में लोड नहीं कर सका.
गड़बड़ी के मैसेज का उदाहरण:
XlaRuntimeError: RESOURCE_EXHAUSTED: Error loading program 'jit_embedding_pipeline_step_fn': Attempting to reserve 29.49G at the bottom of memory. That was not possible. There are 147.64M free, 0B reserved, and 147.64M reservable. Scope: unknown..: while running replica 0 and partition 34 of a replicated computation (other replicas may have failed as well).
XLA बैकएंड: टीपीयू
खास जानकारी
आम तौर पर, यह गड़बड़ी इनमें से किसी एक वजह से होती है:
- प्रोग्राम का साइज़, उपलब्ध एचबीएम से ज़्यादा है: कंपाइल किया गया XLA प्रोग्राम, उसके निर्देशों, स्टैटिक डेटा, और एम्बेड किए गए किसी भी कॉन्स्टेंट के साथ, उस टीपीयू कोर पर फ़िलहाल उपलब्ध कुल मुफ़्त एचबीएम से बड़ा है जिस पर प्रोग्राम लोड किया जा रहा है.
- एचबीएम फ़्रैगमेंटेशन: ऐसा हो सकता है कि डिवाइस पर कुल खाली एचबीएम, प्रोग्राम के लिए काफ़ी हो. हालांकि, यह एक ही लगातार ब्लॉक में उपलब्ध नहीं होता है, जो पूरे प्रोग्राम को फ़िट करने के लिए काफ़ी बड़ा हो.
यह समझना ज़रूरी है कि टीपीयू रनटाइम, मेमोरी को कैसे प्राथमिकता देता है. ज़्यादा लोड वाले प्रोग्राम के मुकाबले, बफ़र के लिए ज़्यादा जगह दी जाती है. अगर बफ़र असाइन नहीं किया जा सका, तो रनटाइम, पहले से लोड किए गए प्रोग्राम को HBM से हटा देगा, ताकि जगह खाली की जा सके. इस वजह से, ऐसा हो सकता है कि पहले लोड होने वाला कोई प्रोग्राम अब ओओएम गड़बड़ी की वजह से लोड न हो पाए. ऐसा इसलिए होता है, क्योंकि अब एचबीएम में ज़्यादा डेटा बफ़र मौजूद हैं.
डीबग करना
- बफ़र मेमोरी फ़ुटप्रिंट कम करें: डेटा बफ़र के लिए इस्तेमाल की गई मेमोरी को खाली करने से, प्रोग्राम के लिए ज़्यादा जगह मिल जाएगी:
- बैच का साइज़ कम करें: यह ऐक्टिवेशन के लिए इस्तेमाल की गई मेमोरी को कम करने का सबसे असरदार तरीका है.
- पैरामीटर शार्डिंग: बहुत बड़े मॉडल के लिए, मॉडल पैरललिज़्म या शार्डिंग तकनीकों (जैसे, FSDP या Megascale) का इस्तेमाल करें. इससे मॉडल के पैरामीटर और कंप्यूटेशन को कई टीपीयू कोर या होस्ट में डिस्ट्रिब्यूट किया जा सकता है.
- सीक्वेंस/कॉन्टेक्स्ट की लंबाई कम करें: क्रम से लगे डेटा को प्रोसेस करने वाले मॉडल के लिए (जैसे, एनएलपी मॉडल) के लिए, सीक्वेंस की लंबाई कम करने से मेमोरी का इस्तेमाल काफ़ी कम हो सकता है.
- Buffer Donation: फ़्रेमवर्क की सुविधाओं का इस्तेमाल करें. जैसे,
jax.jit(..., donate_argnums=...)) का इस्तेमाल किया जाता है, ताकि XLA, आउटपुट को सेव करने के लिए इनपुट बफ़र की मेमोरी का फिर से इस्तेमाल कर सके. इससे मेमोरी के इस्तेमाल की ज़्यादा से ज़्यादा सीमा कम हो जाती है.
- प्रोग्राम के लिए, कुछ समय के लिए इस्तेमाल होने वाले डेटा के लिए मेमोरी की ज़रूरत कम करें:
tpu_shared_memory_percentफ़्लैग का इस्तेमाल करके, प्रोग्राम के लिए कुछ समय के लिए मेमोरी के इस्तेमाल को कम करें. ध्यान दें कि इससे परफ़ॉर्मेंस पर बुरा असर पड़ सकता है.
- विज्ञापन दिखाने की रणनीति को ऑप्टिमाइज़ करें/विज्ञापन दिखाने का लोड कम करें:
- प्रोग्राम लोडिंग मैनेज करना: अगर एक साथ कई फ़ंक्शन JIT-कंपाइल किए जा रहे हैं, तो ध्यान रखें कि हर फ़ंक्शन की वजह से एक प्रोग्राम लोड हो सकता है. अपने वर्कलोड को इस तरह से व्यवस्थित करें कि एक साथ लोड होने वाले प्रोग्राम की संख्या कम हो.
- पक्का करें कि कोई मेमोरी लीक न हो:
- पक्का करें कि
jax.Arrayऑब्जेक्ट के रेफ़रंस को तय समय से ज़्यादा समय तक न रोका गया हो.jax.Arrayऑब्जेक्ट को बनाए रखने से, प्रोग्राम के कंपाइल होने के बाद भी, मेमोरी अपने-आप खाली नहीं हो सकती.
- पक्का करें कि
टूलिंग
- उस
tpu_log_allocations_on_oomफ़्लैग को चालू करें जिसके लिए ऐलोकेटर, ओओएम होने पर सभी मौजूदा ऐलोकेशन की पूरी रिपोर्ट डंप करेगा. यह रिपोर्ट, डीबग करने के लिए बहुत काम की हो सकती है. - अपने प्रोग्राम की प्रोफ़ाइल बनाएं: JAX मेमोरी प्रोफ़ाइलर या TensorFlow प्रोफ़ाइलर का इस्तेमाल करके, समय के साथ अपने प्रोग्राम की मेमोरी के इस्तेमाल के बारे में ज़्यादा जानकारी पाएं. इससे मेमोरी के इस्तेमाल में अचानक होने वाली बढ़ोतरी का पता लगाया जा सकता है.