Расширенные параметры профилировщика

В этом документе перечислены расширенные флаги профилирования, доступные в 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_power_trace_level (целое число): Управляет уровнем трассировки энергопотребления. Поддерживаемые значения:
    • 0 : POWER_TRACE_NONE (По умолчанию)
    • 1 : POWER_TRACE_NORMAL
    • 2 : POWER_TRACE_SPI (Включить трассировку питания SPI, подробный вывод)
  • 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 микропрограммой.

Энергетические и тепловые события

Конфигурация меток трассировки

  • tpu_tracemark_lower (Integer): Нижняя граница для параметра tracemark.
  • tpu_tracemark_upper (Integer): Верхняя граница для параметра tracemark.

Варианты периодического контрольного отбора проб

Эти параметры позволяют настроить периодическую выборку различных счетчиков производительности на TPU.

  • tpu_enable_periodic_counter_sampling (логическое значение): Включает периодическую выборку счетчика.
  • tpu_tc_perf_counter_sampling_options (String): Параметры для выборки счетчика производительности TC. Ожидается текстовая протострока типа xprof::XprofRequest::PeriodicCounterSamplingOptions .
  • tpu_scs_perf_counter_sampling_options (String): Параметры для выборки счетчика производительности SCS.
  • tpu_sctc_perf_counter_sampling_options (String): Параметры для выборки счетчика производительности SCTC.
  • tpu_sctd_perf_counter_sampling_options (String): Параметры для выборки счетчика производительности SCTD.
  • tpu_cmn_perf_counter_sampling_options (String): Параметры для выборки счетчика производительности CMN.
  • tpu_icr_perf_counter_sampling_options (String): Параметры для выборки счетчика производительности ICR.

Другие расширенные параметры TPU

  • tpu_perf_counters (логическое значение): логическое значение для включения/отключения счетчиков производительности TPU.
  • max_trace_buffers (целое число): Управляет максимальным размером буферов трассировки.
  • tpu_circular_buffer_tracing (логическое значение): Включает трассировку кольцевого буфера.
  • tpu_watched_sync_flag_number (Integer): Номер отслеживаемого флага синхронизации.
  • tpu_watched_sync_flag_mask (Integer): Маска флага синхронизации, за которой ведется наблюдение.
  • tpu_sc_dma (логическое значение): управляет трассировкой TPU SC DMA.