گزینه‌های پیشرفته پروفایلر

این سند، فهرستی از پرچم‌های پیشرفته‌ی پروفایلینگ را ارائه می‌دهد که در XProf موجود هستند اما ممکن است در راهنمای اصلی پروفایلینگ JAX مستند نشده باشند. این پرچم‌ها معمولاً برای کنترل دقیق، نظارت بر توان و نمونه‌برداری دقیق از شمارنده‌ی عملکرد در TPUها استفاده می‌شوند.

این پرچم‌ها از طریق پارامتر profiler_options در jax.profiler.start_trace به عنوان بخشی از دیکشنری advanced_configuration ارسال می‌شوند.

مثال:

options = jax.profiler.ProfileOptions()
options.advanced_configuration = {
    "tpu_power_trace_level": 1,  # Integer
    "tpu_perf_counters": True,  # Boolean
    "tpu_cpu_perf_counter_profile_events": "context-switches,page-faults",  # String
}
jax.profiler.start_trace("/tmp/profile-data", profiler_options=options)

رویدادهای برق و حرارتی

این پرچم‌ها، جمع‌آوری رویدادهای توان و حرارت را در TPUها کنترل می‌کنند.

  • tpu_power_trace_level (عدد صحیح): سطح ردیابی توان را کنترل می‌کند. مقادیر پشتیبانی شده:
    • 0 : POWER_TRACE_NONE (پیش‌فرض)
    • 1 : POWER_TRACE_NORMAL
    • 2 : POWER_TRACE_SPI (فعال کردن ردیابی توان SPI، verbose)
  • tpu_e2e_enable_fw_throttle_event (بولی): رویدادهای مربوط به کنترل فریمور را فعال می‌کند.
  • tpu_e2e_enable_fw_power_level_event (بولی): رویدادهای سطح توان میان‌افزار را فعال می‌کند.
  • tpu_e2e_enable_fw_thermal_event (بولی): رویدادهای حرارتی میان‌افزار را فعال می‌کند.
  • tpu_enable_fw_pcie_utilization_event (بولی): رویدادهای استفاده از PCIe میان‌افزار را فعال می‌کند.

رویدادهای برق و حرارتی

پیکربندی Tracemark

  • tpu_tracemark_lower (عدد صحیح): کران پایین برای tracemark.
  • tpu_tracemark_upper (عدد صحیح): کران بالا برای tracemark.

گزینه‌های نمونه‌برداری شمارنده دوره‌ای

این گزینه‌ها نمونه‌برداری دوره‌ای از شمارنده‌های عملکرد مختلف روی TPU را پیکربندی می‌کنند.

  • tpu_enable_periodic_counter_sampling (بولی): نمونه‌برداری شمارنده دوره‌ای را فعال می‌کند.
  • tpu_tc_perf_counter_sampling_options (رشته): گزینه‌هایی برای نمونه‌برداری شمارنده‌ی TC perf. انتظار یک رشته‌ی متنی از xprof::XprofRequest::PeriodicCounterSamplingOptions دارد.
  • tpu_scs_perf_counter_sampling_options (رشته): گزینه‌هایی برای نمونه‌برداری شمارنده‌ی پرفکت SCS.
  • tpu_sctc_perf_counter_sampling_options (رشته): گزینه‌هایی برای نمونه‌برداری شمارنده‌ی پرفکت SCTC.
  • tpu_sctd_perf_counter_sampling_options (رشته): گزینه‌هایی برای نمونه‌برداری شمارنده‌ی پرفکت SCTD.
  • tpu_cmn_perf_counter_sampling_options (رشته): گزینه‌هایی برای نمونه‌برداری شمارنده‌ی پرفکت CMN.
  • tpu_icr_perf_counter_sampling_options (رشته): گزینه‌هایی برای نمونه‌برداری شمارنده‌ی پرفکت ICR.

سایر گزینه‌های پیشرفته TPU

  • tpu_perf_counters (مقدار بولی): مقدار بولی برای فعال/غیرفعال کردن شمارنده‌های عملکرد TPU.
  • max_trace_buffers (عدد صحیح): حداکثر اندازه بافرهای ردیابی را کنترل می‌کند.
  • tpu_circular_buffer_tracing (بولی): ردیابی بافر دایره‌ای را فعال می‌کند.
  • tpu_watched_sync_flag_number (عدد صحیح): شماره پرچم همگام‌سازی مشاهده‌شده.
  • tpu_watched_sync_flag_mask (عدد صحیح): ماسک پرچم همگام‌سازیِ مشاهده‌شده.
  • tpu_sc_dma (بولی): ردیابی DMA از طریق TPU SC را کنترل می‌کند.