XProf की मदद से, OOM से जुड़ी गड़बड़ियां डीबग करना

आउट ऑफ़ मेमोरी (ओओएम) की गड़बड़ियां तब होती हैं, जब ऐक्सलरेटर (जीपीयू या टीपीयू) की हाई बैंडविड्थ मेमोरी (एचबीएम) की क्षमता खत्म हो जाती है. ओओएम की समस्याओं की कुछ सामान्य वजहों और डीबग करने के तरीकों के बारे में, E1000 - कंपाइल टाइम एचबीएम ओओएम दस्तावेज़ और जीपीयू मेमोरी के बंटवारे के बारे में JAX दस्तावेज़ में बताया गया है.

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

1. प्रोफ़ाइल कैप्चर करने के लिए, अपने प्रोग्राम को jax.profiler.trace के साथ चलाएं.
2. XProf को बैकग्राउंड में शुरू करें. इसके बाद, मेमोरी व्यूअर टूल का इस्तेमाल करके, मेमोरी के इस्तेमाल की जानकारी देखें.

प्रोग्राम का उदाहरण

इस JAX प्रोग्राम की वजह से, OOM गड़बड़ी होती है:

import jax
from jax import random
import jax.numpy as jnp


@jax.profiler.trace("/tmp/xprof")
@jax.jit
def oom():
    a = random.normal(random.PRNGKey(1), (327680, 327680), dtype=jnp.bfloat16)
    return a @ a


if __name__ == "__main__":
    oom()

टीपीयू-मशीन पर, यह प्रोग्राम इस मैसेज के साथ काम नहीं करता:

XlaRuntimeError: RESOURCE_EXHAUSTED: Allocation (size=107374182400) would exceed memory (size=17179869184) :: #allocation7 [shape = 'u8[327680,327680]{1,0:T(8,128)(4,1)}', space=hbm, size = 0xffffffffffffffff, tag = 'output of xor_convert_fusion@{}'] :: <no-hlo-instruction>

(जीपीयू-मशीन पर, गड़बड़ी इस तरह दिखती है: XlaRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 214748364800 bytes.)

XProf चलाएं

xprof (pip install xprof) इंस्टॉल करें. इसके बाद, XProf का एक इंस्टेंस शुरू करें. इसमें उस डायरेक्ट्री के बारे में बताएं जहां प्रोफ़ाइल सेव की गई है:

xprof --logdir=/tmp/xprof/ --port=6006

लोकल मशीन पर, http://localhost:6006 पर मौजूद इंस्टेंस पर जाएं. टूल ड्रॉपडाउन में, मेमोरी व्यूअर चुनें. इसके बाद, मेमोरी व्यूअर टूल विंडो में, मेमोरी टाइप ड्रॉपडाउन में एचबीएम चुनें. आम तौर पर, यह डिफ़ॉल्ट रूप से चुना जाता है.

XProf: मेमोरी व्यूअर टूल के दस्तावेज़ में, टूल के कॉम्पोनेंट और दिखाई गई जानकारी के बारे में बताया गया है.

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

बफ़र चार्ट पर कर्सर घुमाकर, ओप के बारे में ज़्यादा जानकारी पाई जा सकती है. जैसे, इसका साइज़, शेप, असाइनमेंट का टाइप वगैरह. इससे आपको उन ऑप्स की पहचान करने और उनका आकलन करने में मदद मिल सकती है जिनमें ज़्यादा या लंबे समय तक चलने वाले टेंपररी हो सकते हैं. साथ ही, ऐसे बड़े इनपुट/इंटरमीडिएट/आउटपुट टेंसर की पहचान करने में मदद मिल सकती है जिनमें खराब पैडिंग होती है. ये सभी चीज़ें, पीक मेमोरी में योगदान देती हैं और इन्हें अडजस्ट या ऑप्टिमाइज़ करने की ज़रूरत होती है.

E1000: डीबग करना में, डीबग करने की खास तकनीकें जानें.