XProf की मदद से JAX कंप्यूटेशन की प्रोफ़ाइलिंग करना

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

XProf का उदाहरण

प्रोग्रामैटिक कैप्चर

jax.profiler.start_trace और jax.profiler.stop_trace तरीकों का इस्तेमाल करके, JAX कोड के लिए प्रोफ़ाइलर ट्रेस कैप्चर करने के लिए अपने कोड को इंस्ट्रुमेंट किया जा सकता है. कॉल jax.profiler.start_trace उस डायरेक्ट्री के साथ जिसमें ट्रेस फ़ाइलें लिखनी हैं. यह वही --logdir डायरेक्ट्री होनी चाहिए जिसका इस्तेमाल XProf को शुरू करने के लिए किया गया था. इसके बाद, ट्रेस देखने के लिए XProf का इस्तेमाल किया जा सकता है.

उदाहरण के लिए, प्रोफ़ाइलर ट्रेस लेने के लिए:

import jax

jax.profiler.start_trace("/tmp/profile-data")

# Run the operations to be profiled
key = jax.random.key(0)
x = jax.random.normal(key, (5000, 5000))
y = x @ x
y.block_until_ready()

jax.profiler.stop_trace()

jax.block_until_ready कॉल को नोट करें. हम इसका इस्तेमाल यह पक्का करने के लिए करते हैं कि डिवाइस पर किए गए काम को ट्रेस में कैप्चर किया गया हो. इस बारे में ज़्यादा जानने के लिए, एसिंक्रोनस डिस्पैच देखें.

start_trace और stop_trace के विकल्प के तौर पर, jax.profiler.trace कॉन्टेक्स्ट मैनेजर का इस्तेमाल भी किया जा सकता है:

import jax

with jax.profiler.trace("/tmp/profile-data"):
  key = jax.random.key(0)
  x = jax.random.normal(key, (5000, 5000))
  y = x @ x
  y.block_until_ready()

ट्रेस देखना

ट्रेस कैप्चर करने के बाद, XProf यूज़र इंटरफ़ेस (यूआई) का इस्तेमाल करके इसे देखा जा सकता है.

XProf की स्टैंडअलोन कमांड का इस्तेमाल करके, सीधे तौर पर प्रोफ़ाइलर यूज़र इंटरफ़ेस (यूआई) लॉन्च किया जा सकता है. इसके लिए, आपको इसे अपनी लॉग डायरेक्ट्री पर पॉइंट करना होगा:

$ xprof --port=8791 /tmp/profile-data
Attempting to start XProf server:
  Log Directory: /tmp/profile-data
  Port: 8791
  Worker Service Address: 0.0.0.0:50051
  Hide Capture Button: False
XProf at http://localhost:8791/ (Press CTRL+C to quit)

दिए गए यूआरएल पर जाएं (जैसे, http://localhost:8791/) पर क्लिक करें.

उपलब्ध ट्रेस, बाईं ओर मौजूद "सेशन" ड्रॉपडाउन मेन्यू में दिखते हैं. अपनी पसंद का सेशन चुनें. इसके बाद, "टूल" ड्रॉपडाउन में जाकर, "ट्रेस व्यूअर" चुनें. अब आपको एक्ज़ीक्यूशन की टाइमलाइन दिखेगी. ट्रेस में नेविगेट करने के लिए, WASD कुंजियों का इस्तेमाल किया जा सकता है. साथ ही, ज़्यादा जानकारी के लिए, इवेंट चुनने के लिए क्लिक या खींचें और छोड़ें. ट्रेस व्यूअर का इस्तेमाल करने के बारे में ज़्यादा जानकारी के लिए, Trace Viewer टूल का दस्तावेज़ देखें.

XProf की मदद से मैन्युअल तरीके से कैप्चर करना

यहां चल रहे प्रोग्राम से, मैन्युअल तरीके से ट्रिगर किए गए N-सेकंड के ट्रेस को कैप्चर करने के निर्देश दिए गए हैं.

XProf सर्वर चालू करें:
```
xprof --logdir /tmp/profile-data/
```
आपको <http://localhost:8791/> पर XProf को लोड करने का विकल्प दिखना चाहिए. --port फ़्लैग का इस्तेमाल करके, कोई दूसरा पोर्ट तय किया जा सकता है.
आपको जिस Python प्रोग्राम या प्रोसेस की प्रोफ़ाइल बनानी है उसमें यहां दिया गया कोड जोड़ें. इसे प्रोग्राम की शुरुआत में जोड़ें:
```
import jax.profiler
jax.profiler.start_server(9999)
```
इससे प्रोफ़ाइलर सर्वर शुरू होता है, जिससे XProf कनेक्ट होता है. अगले चरण पर जाने से पहले, प्रोफ़ाइलर सर्वर चालू होना चाहिए. सर्वर का इस्तेमाल पूरा हो जाने पर, इसे बंद करने के लिए jax.profiler.stop_server() पर कॉल करें.

अगर आपको लंबे समय तक चलने वाले प्रोग्राम (जैसे कि लंबा ट्रेनिंग लूप) के किसी स्निपेट की प्रोफ़ाइल बनानी है, तो इसे प्रोग्राम की शुरुआत में रखें और प्रोग्राम को सामान्य तरीके से शुरू करें. अगर आपको किसी छोटे प्रोग्राम (जैसे कि माइक्रोबेंचमार्क) की प्रोफ़ाइल बनानी है, तो एक विकल्प यह है कि IPython शेल में प्रोफ़ाइलर सर्वर शुरू करें. इसके बाद, अगले चरण में कैप्चर शुरू करने के बाद, %run के साथ छोटा प्रोग्राम चलाएं. एक और विकल्प यह है कि प्रोग्राम की शुरुआत में प्रोफ़ाइलर सर्वर शुरू करें और time.sleep() का इस्तेमाल करें, ताकि आपको कैप्चर शुरू करने के लिए काफ़ी समय मिल सके.
<http://localhost:8791/> खोलें. इसके बाद, सबसे ऊपर बाईं ओर मौजूद "प्रोफ़ाइल कैप्चर करें" बटन पर क्लिक करें. प्रोफ़ाइल सेवा के यूआरएल के तौर पर "localhost:9999" डालें. यह प्रोफ़ाइलर सर्वर का पता है, जिसे आपने पिछले चरण में शुरू किया था. आपको जितने मिलीसेकंड के लिए प्रोफ़ाइल बनानी है उतने मिलीसेकंड डालें. इसके बाद, "CAPTURE" पर क्लिक करें.
अगर आपको जिस कोड की प्रोफ़ाइल बनानी है वह पहले से नहीं चल रहा है (उदाहरण के लिए, अगर आपने Python शेल में प्रोफ़ाइलर सर्वर शुरू किया है), तो कैप्चरिंग के दौरान उसे चलाएं.
कैप्चर पूरा होने के बाद, XProf अपने-आप रीफ़्रेश होना चाहिए. (XProf की सभी प्रोफ़ाइलिंग सुविधाएं, JAX से जुड़ी नहीं हैं. इसलिए, ऐसा हो सकता है कि शुरुआत में आपको लगे कि कुछ भी कैप्चर नहीं किया गया है.) बाईं ओर मौजूद "टूल" में जाकर, "ट्रेस व्यूअर" चुनें.

अब आपको एक्ज़ीक्यूशन की टाइमलाइन दिखेगी. ट्रेस में नेविगेट करने के लिए, WASD कुंजियों का इस्तेमाल किया जा सकता है. साथ ही, इवेंट चुनने के लिए क्लिक या खींचें. इससे आपको सबसे नीचे ज़्यादा जानकारी दिखेगी. ट्रेस व्यूअर का इस्तेमाल करने के बारे में ज़्यादा जानने के लिए, Trace Viewer टूल का दस्तावेज़ देखें.

XProf और Tensorboard

XProf, Tensorboard में प्रोफ़ाइलिंग और ट्रेस कैप्चर करने की सुविधा देने वाला टूल है. जब तक xprof इंस्टॉल है, तब तक Tensorboard में "प्रोफ़ाइल" टैब दिखेगा. इसका इस्तेमाल करना, XProf को अलग से लॉन्च करने जैसा ही है. हालांकि, इसे उसी लॉग डायरेक्ट्री की ओर पॉइंट करके लॉन्च किया जाना चाहिए. इसमें प्रोफ़ाइल कैप्चर करने, उसका विश्लेषण करने, और उसे देखने की सुविधा शामिल है. XProf, tensorboard_plugin_profile की जगह लेता है. पहले इस सुविधा का इस्तेमाल करने का सुझाव दिया जाता था.

$ tensorboard --logdir=/tmp/profile-data
[...]
Serving TensorBoard on localhost; to expose to the network, use a proxy or pass --bind_all
TensorBoard 2.19.0 at http://localhost:6006/ (Press CTRL+C to quit)

कस्टम ट्रेस इवेंट जोड़ना

डिफ़ॉल्ट रूप से, ट्रेस व्यूअर में मौजूद इवेंट, ज़्यादातर लो-लेवल वाले इंटरनल JAX फ़ंक्शन होते हैं. अपने कोड में jax.profiler.TraceAnnotation और jax.profiler.annotate_function का इस्तेमाल करके, अपने इवेंट और फ़ंक्शन जोड़े जा सकते हैं.

प्रोफ़ाइलर के विकल्पों को कॉन्फ़िगर करना

start_trace तरीके में एक वैकल्पिक profiler_options पैरामीटर होता है. इससे प्रोफ़ाइलर के व्यवहार को बेहतर तरीके से कंट्रोल किया जा सकता है. यह पैरामीटर, jax.profiler.ProfileOptions का एक इंस्टेंस होना चाहिए.

उदाहरण के लिए, सभी Python और होस्ट ट्रेस बंद करने के लिए:

import jax

options = jax.profiler.ProfileOptions()
options.python_tracer_level = 0
options.host_tracer_level = 0
jax.profiler.start_trace("/tmp/profile-data", profiler_options=options)

# Run the operations to be profiled
key = jax.random.key(0)
x = jax.random.normal(key, (5000, 5000))
y = x @ x
y.block_until_ready()

jax.profiler.stop_trace()

सामान्य विकल्प

host_tracer_level: यह होस्ट-साइड की गतिविधियों के लिए ट्रेस लेवल सेट करता है.

इस्तेमाल की जा सकने वाली वैल्यू:
- 0: इससे होस्ट (सीपीयू) ट्रेसिंग पूरी तरह से बंद हो जाती है.
- 1: इससे सिर्फ़ उपयोगकर्ता के इंस्ट्रुमेंट किए गए TraceMe इवेंट को ट्रेस किया जा सकता है.
- 2: इसमें लेवल 1 के ट्रेस के साथ-साथ प्रोग्राम के एक्ज़ीक्यूशन की ज़्यादा जानकारी भी शामिल होती है. जैसे, XLA की महंगी कार्रवाइयां (डिफ़ॉल्ट रूप से).
- 3: इसमें लेवल 2 के ट्रेस के साथ-साथ, प्रोग्राम के एक्ज़ीक्यूशन की ज़्यादा जानकारी भी शामिल होती है. जैसे, XLA के कम लागत वाले ऑपरेशन.
device_tracer_level: इससे यह कंट्रोल किया जाता है कि डिवाइस की पहचान करने की सुविधा चालू है या नहीं.

इस्तेमाल की जा सकने वाली वैल्यू:
- 0: यह डिवाइस की पहचान करने की सुविधा बंद करता है.
- 1: डिवाइस ट्रेसिंग की सुविधा चालू करता है (डिफ़ॉल्ट).
python_tracer_level: इससे यह कंट्रोल किया जाता है कि Python ट्रेसिंग की सुविधा चालू है या नहीं.

इस्तेमाल की जा सकने वाली वैल्यू:
- 0: यह Python फ़ंक्शन कॉल ट्रेसिंग को बंद करता है (डिफ़ॉल्ट रूप से चालू होता है).
- 1: इससे Python ट्रेसिंग चालू होती है.

कॉन्फ़िगरेशन के बेहतर विकल्प

TPU के विकल्प

tpu_trace_mode: टीपीयू ट्रेसिंग के मोड के बारे में बताता है.

इस्तेमाल की जा सकने वाली वैल्यू:
- TRACE_ONLY_HOST: इसका मतलब है कि सिर्फ़ होस्ट-साइड (सीपीयू) की गतिविधियों को ट्रैक किया जाता है. साथ ही, डिवाइस (टीपीयू/जीपीयू) के कोई भी ट्रेस इकट्ठा नहीं किए जाते.
- TRACE_ONLY_XLA: इसका मतलब है कि डिवाइस पर सिर्फ़ XLA-लेवल के ऑपरेशनों को ट्रेस किया जाता है.
- TRACE_COMPUTE: इससे डिवाइस पर कंप्यूटिंग से जुड़ी कार्रवाइयों का पता चलता है.
- TRACE_COMPUTE_AND_SYNC: यह डिवाइस पर कंप्यूट ऑपरेशन और सिंक्रनाइज़ेशन इवेंट, दोनों को ट्रैक करता है.
अगर "tpu_trace_mode" नहीं दिया गया है, तो trace_mode डिफ़ॉल्ट रूप से TRACE_ONLY_XLA पर सेट होता है.
tpu_num_sparse_cores_to_trace: इससे पता चलता है कि टीपीयू पर कितने स्पार्स कोर को ट्रेस करना है.
tpu_num_sparse_core_tiles_to_trace: इससे यह तय होता है कि टीपीयू पर ट्रेस करने के लिए, हर स्पार्स कोर में कितने टाइल होंगे.
tpu_num_chips_to_profile_per_task: इससे हर टास्क के लिए, प्रोफ़ाइल किए जाने वाले टीपीयू चिप की संख्या तय की जाती है.

GPU के विकल्प

जीपीयू की प्रोफ़ाइलिंग के लिए, ये विकल्प उपलब्ध हैं:

gpu_max_callback_api_events: यह CUPTI कॉलबैक एपीआई से इकट्ठा किए गए इवेंट की ज़्यादा से ज़्यादा संख्या सेट करता है. यह डिफ़ॉल्ट रूप से 2*1024*1024 पर सेट होता है.
gpu_max_activity_api_events: यह CUPTI activity API से इकट्ठा किए गए इवेंट की ज़्यादा से ज़्यादा संख्या सेट करता है. यह डिफ़ॉल्ट रूप से 2*1024*1024 पर सेट होता है.
gpu_max_annotation_strings: इससे एनोटेशन स्ट्रिंग की ज़्यादा से ज़्यादा संख्या सेट की जाती है. यह डिफ़ॉल्ट रूप से 1024*1024 पर सेट होता है.
gpu_enable_nvtx_tracking: इससे CUPTI में NVTX ट्रैकिंग चालू होती है. डिफ़ॉल्ट रूप से, यह False पर सेट होता है.
gpu_enable_cupti_activity_graph_trace: इससे CUDA ग्राफ़ के लिए, CUPTI गतिविधि ग्राफ़ ट्रेसिंग चालू की जा सकती है. यह डिफ़ॉल्ट रूप से False पर सेट होता है.
gpu_pm_sample_counters: कॉमा लगाकर अलग की गई, GPU परफ़ॉर्मेंस की निगरानी करने वाली मेट्रिक की स्ट्रिंग.इनका इस्तेमाल CUPTI की पीएम सैंपलिंग सुविधा का इस्तेमाल करके किया जाता है. उदाहरण के लिए, "sm__cycles_active.avg.pct_of_peak_sustained_elapsed". पीएम सैंपलिंग डिफ़ॉल्ट रूप से बंद होती है. उपलब्ध मेट्रिक के लिए, NVIDIA का CUPTI दस्तावेज़ देखें.
gpu_pm_sample_interval_us: यह CUPTI PM सैंपलिंग के लिए, माइक्रोसेकंड में सैंपलिंग इंटरवल सेट करता है. यह डिफ़ॉल्ट रूप से 500 पर सेट होता है.
gpu_pm_sample_buffer_size_per_gpu_mb: यह CUPTI PM सैंपलिंग के लिए, हर डिवाइस में सिस्टम मेमोरी बफ़र का साइज़ सेट करता है. यह साइज़ एमबी में होता है. डिफ़ॉल्ट रूप से, यह 64 एमबी पर सेट होती है. ज़्यादा से ज़्यादा 4 जीबी की वैल्यू इस्तेमाल की जा सकती है.
gpu_num_chips_to_profile_per_task: इससे यह तय किया जाता है कि हर टास्क के लिए, कितने जीपीयू डिवाइसों की प्रोफ़ाइल बनानी है. अगर इसे तय नहीं किया जाता है, 0 पर सेट किया जाता है या अमान्य वैल्यू पर सेट किया जाता है, तो सभी उपलब्ध जीपीयू की प्रोफ़ाइल बनाई जाएगी. इसका इस्तेमाल, ट्रेस कलेक्शन के साइज़ को कम करने के लिए किया जा सकता है.
gpu_dump_graph_node_mapping: इस विकल्प को चालू करने पर, CUDA ग्राफ़ नोड मैपिंग की जानकारी को ट्रेस में डंप कर दिया जाता है. यह डिफ़ॉल्ट रूप से False पर सेट होता है.

उदाहरण के लिए:

options = ProfileOptions()
options.advanced_configuration = {"tpu_trace_mode" : "TRACE_ONLY_HOST", "tpu_num_sparse_cores_to_trace" : 2}

अगर पहचान न की गई कोई कुंजी या विकल्प वैल्यू मिलती है, तो यह फ़ंक्शन InvalidArgumentError दिखाता है.