XProf ব্যবহার করে JAX গণনার প্রোফাইলিং

XProf হলো আপনার প্রোগ্রামের পারফরম্যান্স ট্রেস ও প্রোফাইল সংগ্রহ এবং ভিজ্যুয়ালাইজ করার একটি চমৎকার উপায়, যার মধ্যে GPU এবং TPU-এর কার্যকলাপও অন্তর্ভুক্ত থাকে। চূড়ান্ত ফলাফলটি দেখতে অনেকটা এইরকম:

প্রোগ্রাম্যাটিক ক্যাপচার

আপনি jax.profiler.start_trace এবং jax.profiler.stop_trace মেথডগুলোর মাধ্যমে JAX কোডের জন্য একটি প্রোফাইলার ট্রেস ক্যাপচার করতে আপনার কোডকে ইন্সট্রুমেন্ট করতে পারেন। ট্রেস ফাইলগুলো লেখার জন্য ডিরেক্টরিটি উল্লেখ করে jax.profiler.start_trace মেথডটি কল করুন। এটি XProf শুরু করার জন্য ব্যবহৃত --logdir ডিরেক্টরির মতোই হওয়া উচিত। এরপর, ট্রেসগুলো দেখার জন্য আপনি 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 UI ব্যবহার করে সেটি দেখতে পারেন।

আপনি আপনার লগ ডিরেক্টরি নির্দেশ করে স্বতন্ত্র XProf কমান্ড ব্যবহার করে সরাসরি প্রোফাইলার UI চালু করতে পারেন:

$ 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 কী ব্যবহার করতে পারেন, এবং আরও বিস্তারিত তথ্যের জন্য ইভেন্ট নির্বাচন করতে ক্লিক বা ড্র্যাগ করতে পারেন। ট্রেস ভিউয়ার ব্যবহারের বিষয়ে আরও বিস্তারিত জানতে ট্রেস ভিউয়ার টুলের ডকুমেন্টেশন দেখুন।

XProf এর মাধ্যমে ম্যানুয়াল ক্যাপচার

চলমান কোনো প্রোগ্রাম থেকে ম্যানুয়ালি চালু করা N-সেকেন্ডের ট্রেস ক্যাপচার করার নির্দেশাবলী নিচে দেওয়া হলো।

1. একটি XProf সার্ভার চালু করুন:

   xprof --logdir /tmp/profile-data/
   

   আপনি <http://localhost:8791/> -এ XProf লোড করতে পারবেন। আপনি --port ফ্ল্যাগ ব্যবহার করে একটি ভিন্ন পোর্ট নির্দিষ্ট করতে পারেন।

2. আপনি যে পাইথন প্রোগ্রাম বা প্রসেসটির প্রোফাইল তৈরি করতে চান, সেটির শুরুর দিকে কোথাও নিম্নলিখিতটি যোগ করুন:

   import jax.profiler
   jax.profiler.start_server(9999)
   

   এটি প্রোফাইলার সার্ভারটি চালু করে, যার সাথে XProf সংযোগ স্থাপন করে। পরবর্তী ধাপে যাওয়ার আগে প্রোফাইলার সার্ভারটি অবশ্যই চালু থাকতে হবে। সার্ভার ব্যবহার করা শেষ হলে, আপনি jax.profiler.stop_server() কল করে এটিকে বন্ধ করে দিতে পারেন।

   আপনি যদি কোনো দীর্ঘ প্রোগ্রামের একটি অংশের (যেমন একটি দীর্ঘ ট্রেনিং লুপ) প্রোফাইলিং করতে চান, তবে আপনি এটি প্রোগ্রামের শুরুতে রাখতে পারেন এবং আপনার প্রোগ্রামটি স্বাভাবিকভাবে শুরু করতে পারেন। আপনি যদি একটি ছোট প্রোগ্রামের (যেমন একটি মাইক্রোবেঞ্চমার্ক) প্রোফাইলিং করতে চান, তবে একটি উপায় হলো IPython শেলে প্রোফাইলার সার্ভারটি চালু করা এবং পরবর্তী ধাপে ক্যাপচার শুরু করার পর %run ব্যবহার করে ছোট প্রোগ্রামটি চালানো। আরেকটি উপায় হলো প্রোগ্রামের শুরুতে প্রোফাইলার সার্ভারটি চালু করা এবং ক্যাপচার শুরু করার জন্য যথেষ্ট সময় পেতে time.sleep() ব্যবহার করা।

3. <http://localhost:8791/> খুলুন এবং উপরের বাম দিকের "ক্যাপচার প্রোফাইল" বোতামে ক্লিক করুন। প্রোফাইল সার্ভিস ইউআরএল হিসেবে "localhost:9999" লিখুন (এটি আগের ধাপে আপনার চালু করা প্রোফাইলার সার্ভারের ঠিকানা)। আপনি যত মিলিসেকেন্ডের জন্য প্রোফাইল করতে চান, সেই সংখ্যাটি লিখুন এবং "ক্যাপচার" বোতামে ক্লিক করুন।

4. যে কোডটির প্রোফাইলিং করতে চান, সেটি যদি আগে থেকে চালু না থাকে (যেমন, আপনি যদি পাইথন শেলে প্রোফাইলার সার্ভারটি চালু করে থাকেন), তাহলে ক্যাপচার চলার সময় সেটি চালান।

5. ক্যাপচার শেষ হওয়ার পর, XProf স্বয়ংক্রিয়ভাবে রিফ্রেশ হয়ে যাবে। (XProf-এর সব প্রোফাইলিং ফিচার JAX-এর সাথে সংযুক্ত নয়, তাই প্রাথমিকভাবে মনে হতে পারে যে কিছুই ক্যাপচার করা হয়নি।) বাম দিকে "Tools"-এর নিচে, "Trace Viewer" নির্বাচন করুন।

এখন আপনি এক্সিকিউশনের একটি টাইমলাইন দেখতে পাবেন। ট্রেসটি নেভিগেট করার জন্য আপনি WASD কী ব্যবহার করতে পারেন, এবং নীচে আরও বিস্তারিত দেখতে ইভেন্টগুলি নির্বাচন করার জন্য ক্লিক বা ড্র্যাগ করতে পারেন। ট্রেস ভিউয়ার ব্যবহারের বিষয়ে আরও বিস্তারিত জানতে ট্রেস ভিউয়ার টুলের ডকুমেন্টেশন দেখুন।

ক্রমাগত প্রোফাইলিং স্ন্যাপশট ক্যাপচার করুন

যেকোনো সময়ে একটি অবিচ্ছিন্ন প্রোফাইলিং স্ন্যাপশট ক্যাপচার করার নির্দেশাবলী নিচে দেওয়া হলো।

1. আপনি যে পাইথন প্রোগ্রাম বা প্রসেসটির প্রোফাইল তৈরি করতে চান, সেটির শুরুর দিকে কোথাও নিম্নলিখিতটি যোগ করুন:

import jax.profiler
jax.profiler.start_server(9999)

1. আপনার প্রোগ্রাম থেকে স্বাধীনভাবে নিরবচ্ছিন্ন প্রোফাইলিং শুরু করুন:

from xprof.api import continuous_profiling_snapshot
continuous_profiling_snapshot.start_continuous_profiling('localhost:9999', {})

1. এক ঝলক দেখুন:

from xprof.api import continuous_profiling_snapshot
continuous_profiling_snapshot.get_snapshot('localhost:9999', '/tmp/profile-data/')

1. ক্রমাগত প্রোফাইলিং বন্ধ করুন:

from xprof.api import continuous_profiling_snapshot
continuous_profiling_snapshot.stop_continuous_profiling('localhost:9999')

1. এক্সপ্রফ শুরু করুন:

xprof --port=8791 /tmp/profile-data

এখন আপনি এক্সিকিউশনের একটি টাইমলাইন দেখতে পাবেন। ট্রেসটি নেভিগেট করার জন্য আপনি WASD কী ব্যবহার করতে পারেন, এবং নীচে আরও বিস্তারিত দেখতে ইভেন্টগুলি নির্বাচন করার জন্য ক্লিক বা ড্র্যাগ করতে পারেন। ট্রেস ভিউয়ার ব্যবহারের বিষয়ে আরও বিস্তারিত জানতে ট্রেস ভিউয়ার টুলের ডকুমেন্টেশন দেখুন।

এক্সপ্রফ এবং টেনসরবোর্ড

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 এর একটি ইনস্ট্যান্স হওয়া উচিত।

উদাহরণস্বরূপ, সমস্ত পাইথন ও হোস্ট ট্রেস নিষ্ক্রিয় করতে:

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()

সাধারণ বিকল্পগুলি

1. host_tracer_level : হোস্ট-সাইডের কার্যকলাপের জন্য ট্রেস লেভেল নির্ধারণ করে।

   সমর্থিত মানসমূহ:

   • 0 : হোস্ট (সিপিইউ) ট্রেসিং সম্পূর্ণরূপে নিষ্ক্রিয় করে।
   • 1 : শুধুমাত্র ব্যবহারকারী-সৃষ্ট TraceMe ইভেন্টগুলোর ট্রেসিং সক্ষম করে।
   • 2 : লেভেল ১ ট্রেসের পাশাপাশি উচ্চ-স্তরের প্রোগ্রাম নির্বাহের বিবরণ, যেমন ব্যয়বহুল XLA অপারেশন অন্তর্ভুক্ত থাকে (ডিফল্ট)।
   • 3 : এতে লেভেল ২ ট্রেসের পাশাপাশি আরও বিশদ, নিম্ন-স্তরের প্রোগ্রাম নির্বাহের বিবরণ অন্তর্ভুক্ত থাকে, যেমন সস্তা XLA অপারেশন।

2. device_tracer_level : ডিভাইস ট্রেসিং চালু থাকবে কিনা তা নিয়ন্ত্রণ করে।

   সমর্থিত মানসমূহ:

   • 0 : ডিভাইস ট্রেসিং নিষ্ক্রিয় করে।
   • 1 : ডিভাইস ট্রেসিং সক্ষম করে (ডিফল্ট)।

3. python_tracer_level : পাইথন ট্রেসিং চালু থাকবে কিনা তা নিয়ন্ত্রণ করে।

   সমর্থিত মানসমূহ:

   • 0 : পাইথন ফাংশন কল ট্রেসিং নিষ্ক্রিয় করে (ডিফল্ট)।
   • 1 : পাইথন ট্রেসিং সক্ষম করে।

উন্নত কনফিগারেশন বিকল্পগুলি

টিপিইউ বিকল্পগুলি

1. tpu_trace_mode : টিপিইউ ট্রেসিং-এর মোড নির্দিষ্ট করে।

   সমর্থিত মানসমূহ:

   • TRACE_ONLY_HOST : এর অর্থ হলো শুধুমাত্র হোস্ট-সাইড (CPU) কার্যকলাপ ট্রেস করা হয় এবং কোনো ডিভাইস (TPU/GPU) ট্রেস সংগ্রহ করা হয় না।
   • TRACE_ONLY_XLA : এর অর্থ হলো, ডিভাইসে শুধুমাত্র XLA-স্তরের অপারেশনগুলোই ট্রেস করা হবে।
   • TRACE_COMPUTE : এটি ডিভাইসে পরিচালিত গণনা কার্যক্রমগুলো ট্র্যাক করে।
   • TRACE_COMPUTE_AND_SYNC : এটি ডিভাইসে সংঘটিত কম্পিউট অপারেশন এবং সিনক্রোনাইজেশন ইভেন্ট উভয়েরই সন্ধান রাখে।

   যদি "tpu_trace_mode" প্রদান করা না হয়, তাহলে trace_mode ডিফল্টরূপে TRACE_ONLY_XLA হয়ে যায়।

2. tpu_num_sparse_cores_to_trace : টিপিইউ-তে ট্রেস করার জন্য স্পার্স কোরের সংখ্যা নির্দিষ্ট করে।

3. tpu_num_sparse_core_tiles_to_trace : টিপিইউ-তে প্রতিটি স্পার্স কোরের মধ্যে কতগুলো টাইল ট্রেস করা হবে তা নির্দিষ্ট করে।

4. tpu_num_chips_to_profile_per_task : প্রতি টাস্কে প্রোফাইল করার জন্য TPU চিপের সংখ্যা নির্দিষ্ট করে।

5. tpu_cpu_perf_counter_profile_events : নাম অনুসারে প্রোফাইল করার জন্য সিপিইউ পারফরম্যান্স কাউন্টার ইভেন্টগুলির একটি কমা-বিভক্ত তালিকা নির্দিষ্ট করে (যেমন, "context-switches,page-faults" )। সিপিইউ ইভেন্ট প্রোফাইল করার জন্য এটিই প্রস্তাবিত পদ্ধতি, কারণ স্ট্রিং নামগুলি স্বয়ংক্রিয়ভাবে হার্ডওয়্যার/সফ্টওয়্যার প্রকারে রূপান্তরিত হয়।

6. tpu_cpu_perf_counter_configs : সিপিইউ পারফরম্যান্স কাউন্টার কনফিগারেশনগুলির একটি কমা-বিভক্ত তালিকা নির্দিষ্ট করে। প্রতিটি কনফিগারেশন config:type:name আকারে গঠিত, যেখানে config এবং type সরাসরি লিনাক্স perf_event_attr স্ট্রাকচারের ইন্টিজার ফিল্ডের সাথে ম্যাপ করা থাকে এবং name হলো একটি কাস্টম স্ট্রিং লেবেল। এটি RAW ইভেন্ট বা নামে স্বীকৃত নয় এমন কাস্টম PMU কনফিগারেশন পাস করার জন্য একটি উন্নত ফলব্যাক বিকল্প। উদাহরণস্বরূপ: "3:1:context-switches" ।

জিপিইউ বিকল্পগুলি

জিপিইউ প্রোফাইলিংয়ের জন্য নিম্নলিখিত বিকল্পগুলি উপলব্ধ রয়েছে:

• gpu_max_callback_api_events : CUPTI কলব্যাক API দ্বারা সংগৃহীত ইভেন্টের সর্বোচ্চ সংখ্যা নির্ধারণ করে। ডিফল্ট মান হলো 2*1024*1024 ।
• gpu_max_activity_api_events : CUPTI অ্যাক্টিভিটি 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 : CUPTI-এর PM স্যাম্পলিং বৈশিষ্ট্য ব্যবহার করে সংগ্রহ করার জন্য GPU পারফরম্যান্স মনিটরিং মেট্রিক্সের একটি কমা-বিভক্ত স্ট্রিং (যেমন "sm__cycles_active.avg.pct_of_peak_sustained_elapsed" )। PM স্যাম্পলিং ডিফল্টরূপে নিষ্ক্রিয় থাকে। উপলব্ধ মেট্রিক্সের জন্য, NVIDIA-এর CUPTI ডকুমেন্টেশন দেখুন।
• gpu_pm_sample_interval_us : CUPTI PM স্যাম্পলিং-এর জন্য স্যাম্পলিং ব্যবধান মাইক্রোসেকেন্ডে নির্ধারণ করে। ডিফল্ট মান 500 ।
• gpu_pm_sample_buffer_size_per_gpu_mb : CUPTI PM স্যাম্পলিং-এর জন্য প্রতি ডিভাইসের সিস্টেম মেমরি বাফার সাইজ মেগাবাইটে (MB) নির্ধারণ করে। ডিফল্ট মান ৬৪ মেগাবাইট। সর্বোচ্চ সমর্থিত মান হলো ৪ গিগাবাইট।
• gpu_num_chips_to_profile_per_task : প্রতি টাস্কে প্রোফাইল করার জন্য GPU ডিভাইসের সংখ্যা নির্দিষ্ট করে। যদি এটি নির্দিষ্ট না করা হয়, 0 সেট করা হয়, বা একটি অবৈধ মান সেট করা হয়, তাহলে সমস্ত উপলব্ধ GPU প্রোফাইল করা হবে। ট্রেস সংগ্রহের আকার কমানোর জন্য এটি ব্যবহার করা যেতে পারে।
• 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 রিটার্ন করে।