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

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

XProf উদাহরণ

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

আপনি 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)

প্রোফাইলটি দেখতে আপনার ব্রাউজারে প্রদত্ত URL (যেমন, http://localhost:8791/ ) এ যান।

Available traces appear in the "Sessions" dropdown menu on the left. Select the session you're interested in, and then under the "Tools" dropdown, select "Trace Viewer". You should now see a timeline of the execution. You can use the WASD keys to navigate the trace, and click or drag to select events for more details. See the Trace Viewer Tool documentation for more details on using the trace viewer.

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() কল করতে পারেন।

    If you'd like to profile a snippet of a long-running program (eg a long training loop), you can put this at the beginning of the program and start your program as usual. If you'd like to profile a short program (eg a microbenchmark), one option is to start the profiler server in an IPython shell, and run the short program with %run after starting the capture in the next step. Another option is to start the profiler server at the beginning of the program and use time.sleep() to give you enough time to start the capture.

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

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

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

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

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

XProf is the underlying tool that powers the profiling and trace capturing functionality in Tensorboard. As long as xprof is installed, a "Profile" tab will be present within Tensorboard. Using this is identical to launching XProf independently, as long as it is launched pointing to the same log directory. This includes profile capture, analysis, and viewing functionality. XProf supplants the tensorboard_plugin_profile functionality that was previously recommended.

$ 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 : হোস্ট (CPU) ট্রেসিং সম্পূর্ণরূপে অক্ষম করে।
    • 1 : শুধুমাত্র ব্যবহারকারী-উপকরণযুক্ত TraceMe ইভেন্টগুলির ট্রেসিং সক্ষম করে।
    • 2 : লেভেল ১ ট্রেস এবং উচ্চ-স্তরের প্রোগ্রাম এক্সিকিউশনের বিবরণ যেমন ব্যয়বহুল XLA অপারেশন (ডিফল্ট) অন্তর্ভুক্ত করে।
    • 3 : লেভেল ২ ট্রেস এবং আরও ভার্বোজ, নিম্ন-স্তরের প্রোগ্রাম এক্সিকিউশনের বিবরণ যেমন সস্তা XLA অপারেশন অন্তর্ভুক্ত করে।

  2. device_tracer_level : ডিভাইস ট্রেসিং সক্ষম কিনা তা নিয়ন্ত্রণ করে।

    সমর্থিত মান:

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

  3. python_tracer_level : পাইথন ট্রেসিং সক্ষম কিনা তা নিয়ন্ত্রণ করে।

    সমর্থিত মান:

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

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

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

  1. tpu_trace_mode : TPU ট্রেসিংয়ের মোড নির্দিষ্ট করে।

    সমর্থিত মান:

    • 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 : TPU তে ট্রেস করার জন্য স্পার্স কোরের সংখ্যা নির্দিষ্ট করে।

  3. tpu_num_sparse_core_tiles_to_trace : TPU তে ট্রেস করার জন্য প্রতিটি স্পার্স কোরের মধ্যে টাইলের সংখ্যা নির্দিষ্ট করে।

  4. tpu_num_chips_to_profile_per_task : প্রতিটি টাস্কের প্রোফাইলে কতগুলি TPU চিপ ব্যবহার করা হবে তা নির্দিষ্ট করে।

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

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

  • 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 স্যাম্পলিং বৈশিষ্ট্য (যেমন "sm__cycles_active.avg.pct_of_peak_sustained_elapsed" ) ব্যবহার করে সংগ্রহ করার জন্য GPU পারফরম্যান্স মনিটরিং মেট্রিক্সের একটি কমা-বিভাজিত স্ট্রিং। PM স্যাম্পলিং ডিফল্টরূপে অক্ষম করা থাকে। উপলব্ধ মেট্রিক্সের জন্য, NVIDIA এর CUPTI ডকুমেন্টেশন দেখুন।
  • gpu_pm_sample_interval_us : CUPTI PM স্যাম্পলিংয়ের জন্য স্যাম্পলিং ব্যবধান মাইক্রোসেকেন্ডে সেট করে। ডিফল্ট 500 ।
  • gpu_pm_sample_buffer_size_per_gpu_mb : CUPTI PM স্যাম্পলিং এর জন্য প্রতি ডিভাইসে সিস্টেম মেমোরি বাফার সাইজ MB তে সেট করে। ডিফল্ট হিসেবে 64MB। সর্বাধিক সমর্থিত মান হল 4GB।
  • 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 প্রদান করে।