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

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

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

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

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.

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

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

खास जानकारी

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

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

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

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

डीबग करना

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

TensorCore (TC) + SparseCore (SC) HBM Usage Exceeds Limit: अगर गड़बड़ी में, इस्तेमाल की जानकारी साफ़ तौर पर दी गई है, जैसे कि "TC Hbm usage: X, SC Hbm usage Y". → सेक्शन 1. टीसी और एससी एचबीएम के इस्तेमाल को बैलेंस करें.
अचानक ज़्यादा मेमोरी इस्तेमाल होना: अगर गड़बड़ी का मैसेज "एचबीएम मेमोरी स्पेस में मेमोरी खत्म हो गई" दिखता है, तो एचबीएम पर सबसे ज़्यादा मेमोरी इस्तेमाल करने वाले लॉग देखें. अगर एक या उससे ज़्यादा ऐसे टेंसर मौजूद हैं जिनका साइज़ उम्मीद से ज़्यादा है (जैसे, एचबीएम की सीमा का 50% से ज़्यादा) → दूसरे सेक्शन पर जाएं. अनचाहे तौर पर ज़्यादा मेमोरी का इस्तेमाल.
कुल आवंटन, एचबीएम की सीमा से ज़्यादा है: अगर गड़बड़ी का मैसेज "एचबीएम में मेमोरी स्पेस खत्म हो गया है" दिखता है, लेकिन लॉग में कोई बड़ा टेंसर मौजूद नहीं है, तो तीसरे सेक्शन पर जाएं. Aggregate Allocations Exceed HBM Limit.

सेक्शन 1. टीसी और एससी एचबीएम के इस्तेमाल में संतुलन बनाए रखना

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

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 के लिए ज़्यादा वैल्यू सेट करने से, एचबीएम हीप में पहले से फ़ेच किए गए इनपुट डेटा की मात्रा बढ़ जाती है. इस वैल्यू को कम करने से, काफ़ी मेमोरी खाली हो सकती है.
- शार्डिंग की पुष्टि करें: पक्का करें कि एम्बेड की गई टेबल, सभी चिप में सही तरीके से मोड-शार्ड की गई हों. सीमाएं, टेबल पर कैसे लागू होती हैं लेख पढ़ें.
- ज़्यादा आइडिया पाने के लिए, SC: परफ़ॉर्मेंस और मेमोरी से जुड़ी समस्याएँ देखें.
TensorCore का ज़्यादा इस्तेमाल:
- Section 2 पर जाएं.
संतुलित
- अगर दोनों में से कोई भी वैल्यू अलग-अलग बहुत ज़्यादा नहीं है, लेकिन दोनों का योग बहुत ज़्यादा है, तो इसका मतलब है कि चिप की क्षमता पूरी हो गई है. आपको दोनों कॉम्पोनेंट के इस्तेमाल को कम करना होगा. तीनों सेक्शन में दिए गए सुझावों का पालन करें.

सेक्शन 2. अचानक से बहुत ज़्यादा मेमोरी इस्तेमाल होना

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

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

सेक्शन 3. कुल आवंटन, एचबीएम की सीमा से ज़्यादा है

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

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

A. टेंसर पैडिंग और अलाइनमेंट की जांच करना

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

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

B. कॉन्फ़िगरेशन में बदलाव करना

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

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

C. आर्किटेक्चर और शार्डिंग को ऑप्टिमाइज़ करना

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

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

डी॰ मेमोरी पर असर डालने वाले मुख्य XLA फ़्लैग को ट्यून करें:

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

ई॰ Tune XLA Rematerialization Pass / Manual Checkpointing

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

झंडा	ब्यौरा	असर / ट्रेड-ऑफ़
`--xla_tpu_max_hbm_size_mib`	यह विकल्प, रीमटेरियलाइज़ेशन पास के लिए इस्तेमाल किए गए एचबीएम के साइज़ की सीमा को मैन्युअल तरीके से सेट करता है.	यह विकल्प, कंपाइलर को प्रोग्राम को एचबीएम की तय सीमा से कम में फ़िट करने के लिए मजबूर करता है.
`--xla_tpu_rematerialization_algo=PEAK_PRIORITY`	मेमोरी का सबसे ज़्यादा इस्तेमाल होने पर, इस विकल्प का इस्तेमाल किया जाता है.	यह डिफ़ॉल्ट एल्गोरिदम की तुलना में, मेमोरी को कम करने के लिए ज़्यादा असरदार हो सकता है.
`--xla_tpu_rematerialization_max_block_size_limit=32`	यह विकल्प, किसी ब्लॉक में मौजूद निर्देशों की ज़्यादा से ज़्यादा संख्या को कंट्रोल करता है. इन निर्देशों को एक साथ फिर से बनाया जा सकता है.	इसे बढ़ाने से मेमोरी की बचत होती है. हालांकि, इससे कंपाइल होने में लगने वाला समय काफ़ी बढ़ जाता है.
`--xla_tpu_rematerialization_block_effort_factor=10.0`	इससे यह तय होता है कि फिर से तैयार किए जाने वाले ब्लॉक को खोजने में कितना समय (कंपाइल करने में लगने वाला समय) लगेगा.	ज़्यादा वैल्यू से, कंपाइल होने में ज़्यादा समय लगता है. हालांकि, इससे मेमोरी को ज़्यादा से ज़्यादा सेव किया जा सकता है.
`--xla_tpu_pre_fusion_remat=true`	इस विकल्प को चालू करने पर, फ़्यूज़न पास से पहले एक और रीमटेरियलाइज़ेशन पास चालू हो जाता है.	इससे ज़्यादा मेमोरी बचाई जा सकती है. हालांकि, इससे कंपाइल होने में ज़्यादा समय लगता है और संख्यात्मक स्थिरता पर असर पड़ सकता है.

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

F. ऐडवांस प्रोफ़ाइलिंग टूल का इस्तेमाल करना

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

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