Руководство пользователя ядра XProf

XProf Kernel — это набор программ, разработанный для создания ядра Pallas и оптимизации его производительности. Он повышает наглядность и понимание производительности «пользовательских вызовов» во время выполнения — областей, которые обычно остаются «черными ящиками» в трассировке событий XProf.

В настоящее время XProf предлагает такие функции, как пользовательское профилирование вызовов , которое отслеживает статические пакеты TPU для сопоставления инструкций низкоуровневого оптимизатора (LLO) с отслеживаемыми событиями. Хотя это обеспечивает легковесную трассировку, ей не хватает видимости детального поведения различных компонентов TPU во время выполнения. Ядро XProf восполняет этот пробел, позволяя профилировать высокоточные счетчики производительности во время выполнения. В этом руководстве объясняется, как запускать, собирать и визуализировать эти счетчики в формате временных рядов в средстве просмотра трассировки.

Включение детализированных счетчиков производительности

Включите счетчики производительности в расширенных параметрах конфигурации в настройках JAX Profiler TPU, как показано в примере ниже:

options = jax.profiler.ProfileOptions()
options.advanced_configuration = {
    "tpu_enable_periodic_counter_sampling" : True,
    "tpu_tc_perf_counter_sampling_options" : (
        'interval_us:1 scaling:0 counter_size_bits:1 indices:10 indices:11 indices:56 indices:57 indices:58'
    ),
}

• tpu_enable_periodic_counter_sampling : Установите этот флаг, чтобы указать, что следует собирать подробные данные о производительности.
• tpu_<component>_perf_counter_sampling_options : Используйте эти флаги для настройки параметров выборки и указания индексов счетчиков для профилирования конкретного компонента (компоненты перечислены в разделе «Определение индексов счетчиков» ). Поддерживаемые настройки подробно описаны в таблице ниже; для настроек, не определенных явно, будут использоваться значения по умолчанию.

Определение контр-индексов

Используйте инструмент «Счетчики производительности» , чтобы определить конкретные индексы для компонента, который вы хотите профилировать. Выполните следующие шаги:

1. Перейдите к счетчикам производительности в любом окне трассировки TensorBoard, поддерживающем счетчики производительности.
2. Включите параметр "Показывать нулевые значения", чтобы получить полный список профилируемых счетчиков.
3. Найдите нужный компонент . Используйте выделенные жирным шрифтом ключевые слова из таблицы ниже в качестве поисковых запросов, чтобы найти соответствующие счетчики:

1. Найдите свой счетчик в полученной таблице. Вычислите индекс как ( индекс строки - 1 ). Ниже приведен пример инструмента счетчиков производительности и соответствующих счетчиков TC.

   

Визуализация счетчиков времени выполнения

В окне просмотра трассировки собранные счетчики отображаются в виде линий, основанных на указанных индексах и частоте. Выбор точек счетчика открывает нижнее модальное окно, отображающее их значение, включая приращение с момента последней точки данных и соответствующую метку времени.

Расширенные возможности: выборка счетчика по событию

Помимо периодического сбора данных со счетчиков производительности, вы можете настроить счетчики для сбора данных на основе внешних событий, таких как выполнение пользовательских вызовов или конкретных инструкций TPU.

Этот событийно-ориентированный подход обеспечивает задержку захвата менее микросекунды и высокоточную идентификацию. Сравнивая производительность во время выполнения со статическими моделями выполнения, вы можете более эффективно выявлять узкие места на уровне пакета инструкций или самой инструкции.

Чтобы включить сбор данных по событию, замените параметр interval_us на is_external_trigger:true в вашей конфигурации. Пример:

options = jax.profiler.ProfileOptions()

# Example request for externally triggered collection
options.advanced_configuration = {
    "tpu_enable_periodic_counter_sampling": True,
    "tpu_tc_perf_counter_sampling_options": (
        "is_external_trigger:true scaling:0 counter_size_bits:1 indices:10 indices:11 indices:56 indices:57 indices:58"
    ),
}