Dump HLO Computations

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

• HloProto: प्रोटोकॉल बफ़र फ़ाइलें. ये ज़्यादा स्ट्रक्चर्ड और मशीन के पढ़ने लायक फ़ॉर्मैट में होती हैं.
• HloSnapshot: HLO मॉड्यूल और उसके इनपुट. एचएलओ को फिर से चलाने के लिए, कभी-कभी आपको किसी कंप्यूटेशन में दिए गए असली इनपुट की ज़रूरत होती है. इसके लिए, आपको रैंडम डेटा की ज़रूरत नहीं होती.

डंप पाने और उन्हें तय करने के लिए, XLA फ़्लैग का इस्तेमाल किया जा सकता है. ज़्यादातर मामलों में, इसे एनवायरमेंट वैरिएबल के साथ सेट किया जा सकता है. JAX, HLO डंप को प्रिंट करने का प्रोग्रामैटिक तरीका भी उपलब्ध कराता है.

लोकल एक्ज़ीक्यूशन

एनवायरमेंट वैरिएबल का इस्तेमाल करना

डंप पाने के लिए, ज़रूरी फ़्लैग के साथ XLA_FLAGS एनवायरमेंट वैरिएबल सेट किया जा सकता है. यह JAX, TensorFlow, और PyTorch/XLA के लिए काम करता है.

एचएलओ मॉड्यूल और डीबग करने से जुड़ी अन्य जानकारी को किसी डायरेक्ट्री में डंप करने के लिए, अपने प्रोग्राम को --xla_dump_to फ़्लैग के साथ चलाएं:

XLA_FLAGS="--xla_dump_to=DIRECTORY_PATH"

उदाहरण के लिए, पाथ के तौर पर /tmp या /tmp/xladump का इस्तेमाल किया जा सकता है.

डिफ़ॉल्ट रूप से, यह HLO मॉड्यूल को टेक्स्ट के तौर पर डंप करता है. ऐसा ऑप्टिमाइज़ेशन पाइपलाइन की शुरुआत और आखिर में किया जाता है.

फ़ॉर्मैट को साफ़ तौर पर भी बताया जा सकता है:

1. टेक्स्ट डंप

XLA_FLAGS="--xla_dump_hlo_as_text --xla_dump_to=DIRECTORY_PATH"

1. एचएलओ प्रोटो

XLA_FLAGS="--xla_dump_hlo_as_proto --xla_dump_to=DIRECTORY_PATH"

1. एचएलओ स्नैपशॉट

XLA_FLAGS="--xla_dump_hlo_snapshots --xla_dump_to=DIRECTORY_PATH"

1. Graphviz सर्वर की मदद से ग्राफ़ रेंडर करना (यह सिर्फ़ छोटे ग्राफ़ के लिए सही तरीके से काम करता है)

XLA_FLAGS="--xla_dump_hlo_as_url --xla_dump_to=DIRECTORY_PATH"

1. ग्राफ़ को एचटीएमएल फ़ाइल में रेंडर करना (यह सिर्फ़ छोटे ग्राफ़ के लिए अच्छी तरह से काम करता है)

XLA_FLAGS="--xla_dump_hlo_as_html --xla_dump_to=DIRECTORY_PATH"

बड़े ग्राफ़ के कुछ हिस्सों को विज़ुअलाइज़ करने के लिए, interactive_graphviz का इस्तेमाल किया जा सकता है.

खास इंटरमीडिएट पास डंप करें

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

XLA_FLAGS="--xla_dump_hlo_pass_re=regex --xla_dump_to=DIRECTORY_PATH"

रेगुलर एक्सप्रेशन (रेगुलर एक्सप्रेशन) से मेल खाने वाले नामों वाले पास के लिए, HLO मॉड्यूल डंप किए जाएंगे. उदाहरण के लिए, एसपीएमडी पार्टीशनिंग से जुड़े पास से मिलने वाले एचएलओ को इन तरीकों से देखा जा सकता है:

XLA_FLAGS="--xla_dump_to=DIRECTORY_PATH --xla_dump_hlo_pass_re=spmd|propagation"

हर XLA पास के बाद नतीजे को डंप करने के लिए (इससे कई फ़ाइलें बनेंगी), यह सेट किया जा सकता है:

XLA_FLAGS="--xla_dump_to=DIRECTORY_PATH --xla_dump_hlo_pass_re=.*"

JAX के लिए खास विकल्प

JAX में प्रोग्राम के हिसाब से, अपने-आप होने वाली प्रोसेस

फ़्लैग या एनवायरमेंट वैरिएबल पास करने के बजाय, JAX के lower और compile एपीआई का इस्तेमाल करके, प्रोग्राम के हिसाब से HLO को डंप किया जा सकता है.

ऑप्टिमाइज़ नहीं किए गए ओरिजनल लोअर एचएलओ को स्थानीय तौर पर फ़ेच करने के लिए:

jax.jit(f).lower(*args).as_text('hlo')

एचएलओ कंपाइलेशन पास के दौरान फ़ाइलों में डंप करने के लिए, यह विकल्प सेट करें:

compilation_args = {
    'xla_dump_to': DIRECTORY_PATH,
    'xla_dump_hlo_pass_re': 'spmd|propagation', # or some other pass filter
    ...
    }

jax.jit(f).lower(*args).compile(compilation_args)

Dump jaxprs

jaxprs, प्रोग्राम ट्रेस के लिए JAX का इंटरमीडिएट रिप्रेज़ेंटेशन है. इसे डंप करने के लिए, एनवायरमेंट वैरिएबल सेट करें:

JAX_DUMP_IR_TO="DIRECTORY_PATH" JAX_DUMP_IR_MODES=jaxpr

स्टेज किए गए कंप्यूटेशन को एक्सपोर्ट और क्रम से लगाने के बारे में ज़्यादा जानें: डीबग करना. इसके लिए, JAX के दस्तावेज़ पढ़ें.

Google Colab

एनवायरमेंट वैरिएबल

अपनी नोटबुक की पहली सेल में, ऊपर दी गई XLA_FLAGS को os.environ के साथ जोड़ें. ऐसा इसलिए, क्योंकि एनवायरमेंट वैरिएबल और कमांड-लाइन फ़्लैग को आम तौर पर सिर्फ़ एक बार प्रोसेस किया जाता है. जैसे, मॉड्यूल इंपोर्ट करने के समय या XLA बैकएंड शुरू करने के समय. उदाहरण के लिए:

import os
os.environ['XLA_FLAGS'] = "--xla_dump_to=DIRECTORY_PATH"

इससे कंप्यूटेशन को DIRECTORY_PATH में डंप कर दिया जाएगा. उदाहरण के लिए, /tmp. Colab में, इस डायरेक्ट्री को देखने और ऐक्सेस करने के लिए, बाईं ओर मौजूद साइडबार में "Files" ब्राउज़र पर जाएं.

लोकल एक्ज़ीक्यूशन सेक्शन में बताए गए सभी फ़्लैग इस्तेमाल किए जा सकते हैं.

JAX के लिए खास विकल्प

यह लोकल एक्ज़ीक्यूशन की तरह ही होता है. लाइव और इंटरैक्टिव इंट्रॉस्पेक्शन के लिए, सीधे तौर पर कंप्यूटेशन के प्री-ऑप्टिमाइज़ किए गए एचएलओ को प्रिंट किया जा सकता है:

def f(x):
    return jax.numpy.sin(jax.numpy.cos(x))

c = jax.jit(f).lower(3.).compiler_ir('hlo')

print(c.as_hlo_text())

आपके पास सीधे तौर पर, कंप्यूटेशन के ऑप्टिमाइज़ किए गए एचएलओ को प्रिंट करने का विकल्प भी है:

def optimized_HLO(f, *args, platform=None):
    print(jax.jit(f).lower(*args).compile().as_text())

def f(x):
    return jax.numpy.sin(jax.numpy.cos(x))

optimized_HLO(f, 1.0)

सभी/कुछ कंप्यूटेशन को डंप करना

अगर आपको डंप में सब कुछ देखना है, जिसमें सभी छोटे कंपाइलेशन शामिल हैं, तो JAX एनवायरमेंट वैरिएबल को सेट करें:

JAX_COMPILER_DETAILED_LOGGING_MIN_OPS=0

मोज़ेक

Mosaic, Pallas TPU बैकएंड और एक्सपेरिमेंट के तौर पर उपलब्ध Pallas GPU बैकएंड के लिए कंपाइलर है. मोज़ेक कंप्यूटेशन को डंप करने के लिए, यह फ़्लैग सेट करें:

--xla_mosaic_dump_to=/tmp/mosaic_dumps

इसके अलावा, टीपीयू को शुरू करने के आर्ग्युमेंट को एनवायरमेंट वैरिएबल के तौर पर सेट करें:

export LIBTPU_INIT_ARGS="--xla_mosaic_dump_to=/tmp/mosaic_dumps"

ज़्यादा जानने के लिए, Pallas और Mosaic के बारे में JAX का दस्तावेज़ देखें.

HLO Dumps के अन्य वीडियो

सही कैलकुलेशन का पता लगाना

आम तौर पर, कई कंप्यूटेशन डंप हो जाते हैं. डंप की गई फ़ाइलों का नाम, JAX, TensorFlow या PyTorch/XLA "कंप्यूटेशन नेम” के हिसाब से रखा जाता है. ये नाम, लॉग में भी दिखते हैं. इससे काम की HLO फ़ाइलों की पहचान करना आसान हो जाता है. उदाहरण के लिए:

1624325116260738.module_0065.pmap__unnamed_wrapped_function_.186875.before_optimizations.txt

इसके अलावा, ripgrep का इस्तेमाल करके यह तुरंत पता लगाया जा सकता है कि किस मॉड्यूल में खास सिंबल या हिसाब-किताब मौजूद है.

सलाह: गड़बड़ी की रिपोर्ट में, बफ़र असाइनमेंट से पहले/बाद/बफ़र की गई तीन फ़ाइलें शामिल करें.

एचएलओ कन्वर्ज़न

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

इसका इस्तेमाल करने का तरीका जानें: XLA टूलिंग से जुड़ा दस्तावेज़: hlo-opt.

फिर से चलाएं

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

यहां दी गई कमांड में फ़र्ज़ी डेटा का इस्तेमाल किया गया है. अगर आपने HLO स्नैपशॉट सेव किए हैं, तो उनके डेटा का इस्तेमाल किया जा सकता है. स्नैपशॉट चलाते समय नकली डेटा का इस्तेमाल जारी रखने के लिए, --force_fake_data फ़्लैग पास करें.

सीपीयू बैकएंड:

bazel run -c opt //xla/hlo/tools:run_hlo_module -- --platform=cpu
 /tmp/xladump/module_4561.before_optimizations.txt

जीपीयू बैकएंड:

bazel run -c opt //xla/hlo/tools:run_hlo_module -- --platform=CUDA
 /tmp/xladump/module_4561.before_optimizations.txt