गड़बड़ी का कोड: E1000

कैटगरी: कंपाइल का समय: HBM OOM

इस गड़बड़ी से पता चलता है कि प्रोग्राम को टीपीयू डिवाइस पर उपलब्ध एचबीएम से ज़्यादा एचबीएम की ज़रूरत है.

गड़बड़ी के मैसेज के उदाहरण:

RESOURCE_EXHAUSTED: TPU TensorCore Hbm usage: 34.82G, SparseCore Hbm usage 174.10G, exceeding available bytes: 95.74G

RESOURCE_EXHAUSTED: XLA:TPU compile permanent error. Ran out of memory in memory space hbm. Used 49.34G of 32.00G hbm. Exceeded hbm capacity by 17.34G.

XLA बैकएंड: टीपीयू

खास जानकारी

XLA, यह पक्का करने के लिए जांच करता है कि सभी ज़रूरी स्टैटिक ऐलोकेशन का कुल साइज़, डिवाइस के एचबीएम में फ़िट हो.

कंपाइलर, टीपीयू की तय की गई एचबीएम क्षमता को कई तरह के एलॉकेशन के लिए मैनेज करता है:

• प्रोग्राम के इनपुट और आउटपुट: ट्रेनिंग बैच, ऑप्टिमाइज़र की स्थितियां वगैरह.
• टीपीयू टेंपररी: इंटरमीडिएट कैलकुलेशन के लिए डाइनैमिक मेमोरी की ज़रूरत होती है (जैसे, ऐक्टिवेशन, ग्रेडिएंट वगैरह).
• कंपाइल किया गया बाइनरी: यह TensorCore (टीसी) और SparseCore (एससी), दोनों के लिए मशीन कोड होता है.
• सिस्टम ओवरहेड: XLA रनटाइम के लिए रिज़र्व की गई जगह. उदाहरण के लिए, टीपीयू की पुरानी जनरेशन पर इनफ़ीड बफ़र.
• स्थिरांक: एचएलओ आईआर में एम्बेड की गई कॉन्सटेंट वैल्यू, एचबीएम पर असाइन की जाती हैं.
• कंपाइलर इंटरनल: प्रोग्राम लेवल और हर एचएलओ के लिए मेमोरी का बंटवारा (जैसे, मेश में नोड के लिए राउटिंग की जानकारी).

यह गड़बड़ी तब होती है, जब XLA कंपाइलर, ऊपर दिए गए सभी असाइनमेंट को डिवाइस के एचबीएम में फ़िट नहीं कर पाता.

डीबग करना

गड़बड़ी के मैसेज और लॉग का ध्यान से विश्लेषण करें. इससे यह पता चलेगा कि नीचे दी गई एचबीएम ओओएम की कौनसी कैटगरी आपकी गड़बड़ी के बारे में सबसे अच्छी तरह बताती है:

• "टीसी एचबीएम का इस्तेमाल: X, एससी एचबीएम का इस्तेमाल Y": अगर गड़बड़ी में एचबीएम के इस्तेमाल के बारे में साफ़ तौर पर बताया गया है, तो इसका मतलब है कि कुल TensorCore (टीसी) + SparseCore (एससी) का इस्तेमाल, एचबीएम की सीमा से ज़्यादा है. → पहली स्थिति पर जाएं. टीसी और एससी एचबीएम के इस्तेमाल को बैलेंस करें.
• "Ran out of memory in memory space HBM": HBM पर सबसे बड़े असाइनमेंट की गिनती के लिए, लॉग देखें.
  • अगर एक या इससे ज़्यादा बड़े टेंसर (जैसे, एचबीएम की सीमा का > 50%) मौजूद हैं, तो दूसरा परिदृश्य पर जाएं. उम्मीद से ज़्यादा मेमोरी इस्तेमाल होने की वजह से, मेमोरी खत्म हो गई.
  • अगर लॉग में अचानक बड़े टेंसर मौजूद नहीं हैं → तीसरे चरण पर जाएं. कुल मिलाकर किए गए स्टोरेज के बंटवारे की वजह से, मेमोरी खत्म हो गई है.

पहली स्थिति. टीसी और एससी एचबीएम के इस्तेमाल के बीच संतुलन बनाए रखना

अगर गड़बड़ी में, इस्तेमाल की जानकारी साफ़ तौर पर दी गई है, जैसे कि "टीसी एचबीएम का इस्तेमाल: X, एससी एचबीएम का इस्तेमाल Y", तो इसका मतलब है कि कुल TensorCore (टीसी) + SparseCore (एससी) का इस्तेमाल, एचबीएम की सीमा से ज़्यादा है. इन दोनों वैल्यू की तुलना करके, परफ़ॉर्मेंस में आने वाली रुकावट का पता लगाएं:

• SparseCore का ज़्यादा इस्तेमाल
  • एचबीएम स्टैक के इस्तेमाल को ऑप्टिमाइज़ करें: एचबीएम स्टैक की मेमोरी का इस्तेमाल feature_width, max_unique_nz_per_row, और logical_replica_count के साथ बढ़ता है. टेबल की प्रोसेसिंग को क्रम से करने वाले --xla_sc_num_serialized_tables_to_optimize_hbm फ़्लैग को ट्यून करके, स्टैक के इस्तेमाल को कम किया जा सकता है. हालांकि, इससे पैरललिज़्म कम हो जाता है.
  • पैडिंग ओवरहेड की जांच करें: SparseCore, एम्बेडिंग टेबल को 32B (8 फ़्लोट) के साथ अलाइन करता है. कम चौड़ाई वाली सुविधाओं वाली टेबल (जैसे, < 8 फ़्लोट) में, पैडिंग काफ़ी ज़्यादा होती है. इससे एचबीएम का इस्तेमाल सही तरीके से नहीं हो पाता.
  • हीप के इस्तेमाल को कम करना: maximum_parallel_iterations की ज़्यादा वैल्यू से, एचबीएम हीप में पहले से फ़ेच किए गए इनपुट डेटा की मात्रा बढ़ जाती है. इस वैल्यू को कम करने से, काफ़ी मेमोरी खाली हो सकती है.
  • शार्डिंग की पुष्टि करें: पक्का करें कि एम्बेड की गई टेबल को सभी चिप में सही तरीके से mod-शार्ड किया गया हो. सीमाएं, टेबल पर कैसे लागू होती हैं लेख पढ़ें.
  • ज़्यादा आइडिया पाने के लिए, SC: परफ़ॉर्मेंस और मेमोरी की समस्याएं देखें.
• TensorCore का ज़्यादा इस्तेमाल
  • दूसरी स्थिति पर जाएं.
• संतुलित
  • अगर दोनों में से कोई भी वैल्यू अलग-अलग बहुत ज़्यादा नहीं है, लेकिन दोनों का योग बहुत ज़्यादा है, तो इसका मतलब है कि चिप की क्षमता पूरी हो गई है. आपको दोनों कॉम्पोनेंट के इस्तेमाल को कम करना होगा. तीनों सेक्शन में दिए गए सुझावों का पालन करें.

परिदृश्य 2. ज़रूरत से ज़्यादा मेमोरी इस्तेमाल होने की वजह से, मेमोरी से बाहर

अगर आपको गड़बड़ी का यह मैसेज दिखता है: "एचबीएम मेमोरी स्पेस में मेमोरी खत्म हो गई है" और लॉग में एक या उससे ज़्यादा बड़े असाइनमेंट मौजूद हैं (एचबीएम की सीमा का > 50%), तो यह समस्या हार्डवेयर की क्षमता से जुड़ी नहीं होती. आम तौर पर, यह कॉन्फ़िगरेशन से जुड़ी गड़बड़ी होती है. बड़े असाइनमेंट के XLA लेबल (अगर मौजूद है) की जांच करें, ताकि उनके JAX सोर्स कोड के बारे में जानकारी मिल सके.

• डीबग करने से जुड़ी आर्टफ़ैक्ट फ़ाइलें हटाना
  • बड़े पैमाने पर रन करने के दौरान, jax.debug.print() का इस्तेमाल करने से कंपाइलर, पूरे टेंसर को एचबीएम में ट्रांसफ़र करने के लिए मजबूर हो सकता है. इससे फ़्यूज़न टूट जाता है और मेमोरी का इस्तेमाल बढ़ जाता है. बचे हुए jax.debug.print() हटाएं.
• मेष के आकार या शार्डिंग से जुड़ी समस्याओं को ठीक करना
  • मेश के गलत आकार या शार्डिंग के एनोटेशन मौजूद न होने की वजह से, कंपाइलर डिफ़ॉल्ट रूप से रेप्लिकेशन का इस्तेमाल कर सकता है. इससे कंपाइलर को एक ही चिप पर बहुत बड़े टेंसर फ़िट करने पड़ सकते हैं.
  • बड़े ऐलोकेशन के साइज़ की जांच करें. साथ ही, पुष्टि करें कि XLA ने शार्डिंग को सही तरीके से तय किया है और उसे लागू किया है.

परिदृश्य 3. कुल मेमोरी खत्म हो गई है

अगर आपको गड़बड़ी का यह मैसेज दिखता है: "एचबीएम मेमोरी स्पेस में मेमोरी खत्म हो गई" और लॉग में कोई भी बड़ा टेंसर मौजूद नहीं है, तो इसका मतलब है कि एचबीएम की सीमा से ज़्यादा मेमोरी इस्तेमाल की गई है. ऐसे में, मेमोरी प्रोफ़ाइल को विज़ुअलाइज़ करना अक्सर मददगार होता है. इससे, पीक यूसेज में योगदान देने वाले बफ़र की पहचान की जा सकती है. मेमोरी की सबसे ज़्यादा खपत करने वाले कॉम्पोनेंट की पहचान करने के लिए, सिलसिलेवार निर्देशों वाली गाइड देखने के लिए, XProf की मदद से ओओएम गड़बड़ियों को डीबग करना लेख पढ़ें.

सबसे ज़्यादा योगदान देने वाले कुछ लोगों की पहचान करने के बाद, मेमोरी फ़ुटप्रिंट को ऑप्टिमाइज़ करने के लिए यह तरीका अपनाएं.

तीसरी स्थिति.A कॉन्फ़िगरेशन में बदलाव करना

कॉन्फ़िगरेशन में ये बदलाव करके, अक्सर ओओएम की समस्याओं को ठीक किया जा सकता है:

• बैच का साइज़ कम करें: इंटरमीडिएट ऐक्टिवेशन और ग्रेडिएंट के लिए ज़रूरी मेमोरी, बैच के साइज़ के हिसाब से तय होती है. बैच का साइज़ कम करने से, अक्सर मेमोरी का इस्तेमाल कम करने में मदद मिलती है.
• दान किए गए इनपुट बफ़र: jax.jit का इस्तेमाल करते समय, अपने मॉडल पैरामीटर के लिए donate_argnums तय करें. इससे XLA को इनपुट मेमोरी को आउटपुट से बदलने की अनुमति मिलती है.
• मिक्सड प्रिसिशन (bfloat16) चालू करें: अगर मॉडल आर्किटेक्चर और क्वालिटी से जुड़ी ज़रूरी शर्तें पूरी होती हैं, तो प्रोग्राम में सबसे बड़े टेंसर के लिए bfloat16 या क्वॉन्टाइज़ेशन (int8 वगैरह) का इस्तेमाल करें. ध्यान दें कि इस बदलाव से मॉडल के व्यवहार पर असर पड़ सकता है. इसलिए, इस पर सावधानी से विचार करना चाहिए.

तीसरा चरण.बी आर्किटेक्चर और शार्डिंग को ऑप्टिमाइज़ करना

अगर कॉन्फ़िगरेशन में किए गए बदलाव काफ़ी नहीं हैं, तो हो सकता है कि मॉडल टोपोलॉजी, मौजूदा हार्डवेयर सेटअप के लिए बहुत बड़ी हो.

• टीपीयू की नई जनरेशन का इस्तेमाल करें: आम तौर पर, नए टीपीयू में हर चिप के लिए ज़्यादा एचबीएम उपलब्ध होता है; अगर टीपीयू की नई जनरेशन उपलब्ध है, तो उसका इस्तेमाल करें.
• बड़े चिप टोपोलॉजी पर चलाएं: अगर मॉडल के वेट, मौजूदा टोपोलॉजी के लिए बहुत बड़े हैं, तो उन्हें ज़्यादा चिप पर शार्ड करने की कोशिश करें.
• शार्डिंग की ऐडवांस तकनीकों को लागू करें:
  • डेटा, टेंसर या पाइपलाइन पैरललिज़्म के ज़्यादा बेहतर तरीकों के बारे में जानें.
  • इंटरमीडिएट वैल्यू और आउटपुट के लिए, शार्डिंग के सुझाव दें.
• JAX होस्ट ऑफलोडिंग का इस्तेमाल करें: होस्ट ऑफलोडिंग की तकनीकों की मदद से, उपयोगकर्ता बड़े टेंसर को होस्ट सीपीयू मेमोरी में ऑफलोड कर सकता है. जैसे, ऐक्टिवेशन ऑफलोडिंग और ऑप्टिमाइज़र स्टेट ऑफलोडिंग.

तीसरा उदाहरण.सी: टेंसर पैडिंग और अलाइनमेंट की जांच करना

टेंसर के गलत शेप की वजह से, टीपीयू पर ओओएम की समस्या आम तौर पर होती है. टीपीयू पर सबसे अच्छी परफ़ॉर्मेंस पाने के लिए, XLA पैड, टेंसर डाइमेंशन को पैड करता है. आम तौर पर, सबसे छोटे डाइमेंशन के लिए 128 के मल्टीपल और दूसरे सबसे छोटे डाइमेंशन के लिए 8 के मल्टीपल. इस पैडिंग का असर, इनपुट ऐरे और इंटरमीडिएट टेंसर (एचएलओ टेंपररी) दोनों पर पड़ता है. इससे मेमोरी का इस्तेमाल काफ़ी बढ़ सकता है. ऐसा खास तौर पर, डाइमेंशन के छोटे साइज़ के साथ होता है. ऐरे लेआउट देखें.

• बड़े बफ़र के साइज़ की जांच करें: (डिफ़ॉल्ट लेआउट के साथ TPU v5 पर)
  • Xprof Memory Viewer में किसी बफ़र पर कर्सर घुमाने से, बफ़र की जानकारी वाला कार्ड दिखता है. इसमें बफ़र की जानकारी के साथ-साथ पैडिंग की जानकारी भी शामिल होती है.
  • उदाहरण: (129, 1024) के आकार को (256, 1024) में बदला जा सकता है. इससे करीब 50% मेमोरी बर्बाद हो जाती है.
  • सुधार: (128, 1024) के लिए पैडिंग की ज़रूरत नहीं होती है और इसमें 0% मेमोरी बर्बाद होती है.
• डाइमेंशन अलाइन करें: पक्का करें कि सभी बड़े टेंसर डाइमेंशन (बैच साइज़, एम्बेडिंग डाइमेंशन, हिडन साइज़) 128 के गुणज हों. ध्यान दें कि इस बदलाव से मॉडल के व्यवहार पर असर पड़ सकता है. इसलिए, इस पर सावधानी से विचार करना चाहिए.

तीसरा उदाहरण.डी XLA फ़्लैग पर असर डालने वाली मुख्य मेमोरी को ट्यून करना

मेमोरी के मुख्य फ़्लैग को इस तरह से ट्यून किया जा सकता है कि मेमोरी का इस्तेमाल कम हो और परफ़ॉर्मेंस पर भी असर न पड़े. हालांकि, इस रणनीति का इस्तेमाल आखिरी विकल्प के तौर पर किया जाना चाहिए, क्योंकि इससे परफ़ॉर्मेंस पर बुरा असर पड़ सकता है.

तीसरा उदाहरण.ई Tune XLA rematerialization पास/मैन्युअल चेकपॉइंटिंग

अगर मॉडल मेमोरी में फ़िट होने के करीब है, तो jax.grad के साथ jax.checkpoint डेकोरेटर का इस्तेमाल किया जा सकता है. इससे यह मैन्युअल तरीके से कंट्रोल किया जा सकता है कि फ़ॉरवर्ड पास के दौरान कौनसे इंटरमीडिएट सेव किए जाएं और बैकवर्ड पास के दौरान कौनसे फिर से कंप्यूट किए जाएं. इससे कंप्यूट साइकल को एचबीएम के लिए ट्रेड किया जा सकता है.

इसके अलावा, XLA::Rematerialization पास को मेमोरी बचाने के लिए मजबूर किया जा सकता है. इससे कंपाइल होने में ज़्यादा समय लग सकता है:

ध्यान दें कि XLA फ़्लैग में बदलाव करने का विकल्प, आखिरी विकल्प के तौर पर इस्तेमाल किया जाना चाहिए. इसकी वजह यह है कि इससे परफ़ॉर्मेंस पर बुरा असर पड़ सकता है.

तीसरा चरण.एफ़: बेहतर प्रोफ़ाइलिंग टूल का इस्तेमाल करना

XProf की मदद से ओओएम गड़बड़ियों को डीबग करना लेख में, XProf Memory Viewer का इस्तेमाल करने के बारे में ट्यूटोरियल दिया गया है. इसकी मदद से, कंपाइलर के हिसाब से एचबीएम के इस्तेमाल को विज़ुअलाइज़ किया जा सकता है.

इस टूल की मदद से, मेमोरी के ज़्यादा से ज़्यादा इस्तेमाल और बफ़र के लाइफ़टाइम को देखा जा सकता है. इससे यह समझने में मदद मिलती है कि मेमोरी के ज़्यादा से ज़्यादा इस्तेमाल के समय, HBM का इस्तेमाल किस वजह से हो रहा है. प्रोफ़ाइलिंग के सामान्य सेटअप के लिए, Xprof का इस्तेमाल शुरू करना और TensorBoard की प्रोफ़ाइलिंग सुविधा देखें.