Пользовательские вызовы XLA позволяют выполнять пользовательские ядра или операции, которые изначально не поддерживаются XLA. Чтобы получить представление о производительности этих пользовательских вызовов в средстве просмотра трассировки , можно использовать специальные флаги XLA для включения подробной трассировки и отладочной информации LLO (низкоуровневого оптимизатора).
Как включить трассировку
Установите следующий флаг XLA для перекомпиляции рабочей нагрузки с использованием инструментирования:
--xla_xprof_enable_custom_call_tracing=true
Если вы используете Cloud TPU с параметром LIBTPU_INIT_ARGS , вы можете передать его следующим образом:
LIBTPU_INIT_ARGS="--xla_xprof_enable_custom_call_tracing=true --xla_xprof_register_llo_debug_info=true" python your_jax_workload.py
Пример средства просмотра трассировки
Вот пример того, как выглядят трассировки LLO в программе просмотра трассировок Xprof:
Расширенные параметры (Обработка пропущенных событий)
Если в Xprof вы видите пропуски событий или переполнение буфера, это означает, что точки трассировки запускаются слишком часто, перегружая буферы трассировки оборудования. Вы можете настроить частоту вставки трассировки LLO с помощью расширенных параметров.
Эти параметры настраиваются через xla_tpu_bundle_instrumentation_options . Вы можете контролировать, как часто трассировки упаковываются в пакеты инструкций.
Ключевые параметры:
-
trace_best_effort_frequency(По умолчанию: 10): Целевой интервал (в пакетах) для вставки трассировок, упакованных в существующие пакеты. Компилятор будет пытаться вставить трассировку с такой частотой, но не будет создавать для нее новые пакеты. -
trace_guaranteed_frequency(По умолчанию: 10): Максимальное количество пакетов, разрешенных между двумя трассировками. Это гарантированное значение. Если мы не можем обеспечить это значение путем упаковки трассировок в существующие пакеты, мы создадим новый пакет и поместим трассировку туда (отдельно).
Как настроить:
- Если вы видите сообщение "Пропуск событий" : увеличьте значения (например, установите значение 50 или 100), чтобы трассировка выполнялась реже , что уменьшит объем генерируемых данных трассировки.
- Если вам нужна более высокая детализация : уменьшите количество значений, которые необходимо отслеживать чаще (ценой увеличения накладных расходов и потенциального переполнения буфера).
Как рассчитывается количество циклов выполнения инструкций
Поскольку точки трассировки вводятся выборочно, а не при каждой отдельной инструкции, промежуточные временные метки интерполируются на основе оценочных затрат на аппаратные циклы.
Компилятор вычисляет внутреннюю стоимость аппаратных циклов каждой инструкции LLO на основе целевого поколения TPU и исполнительного блока, который её обрабатывает. Эти значения количества циклов представляют собой пропускную способность выполнения и задержки.
Поток высокого уровня
- Инструкция по разбору LLO : Определите код операции и метаданные.
- Получение базовых аппаратных циклов : определение циклов на основе поколения TPU (v5e/v5p, v6e/v7x и т. д.).
- Преобразование в такты глобального счетчика таймера (GTC) : Преобразуйте циклы в такты глобального счетчика таймера (GTC), используя формулу:
Cycles * (GTC_Freq * 16) / TC_Freq. - Создать временной интервал : равномерно интерполировать промежуточные события между известными границами трассировки.
Расчеты циклов по энергоблокам и источникам генерации
Ниже приведены примеры моделирования базовых аппаратных циклов для различных исполнительных блоков:
Блок матричного умножения (MXU)
Показатель количества циклов MXU отражает пропускную способность в зависимости от плотности типов данных.
| Категория инструкций | Подтип / Формат | (v5e/v5p) | (v6e/v7x) |
|---|---|---|---|
| Вектор Матмул | Ф32 | 8 | 8 |
| Предварительная обработка матричных уравнений (F8 до BF16) | 4 | 4 | |
| Упакованный BF16 | 2 | 2 | |
| Целочисленные форматы (U8, S8, U4, S4) | 1 | 1 | |
| Векторные защелки | Транспонированный F32 | 4 | 4 |
| Транспонированный BF16 | 8 | 8 | |
| Нетранспонированный F32 | 2 | 2 | |
| Нетранспонированный BF16 | 4 | 4 | |
| Matprep / Dwg | Все | 1 | 1 |
Транспонированный блок (XLU)
Количество циклов отражает задержки транспонирования в памяти и задержки кроссбара.
| Категория инструкций | Подтип / Формат | (v5e/v5p) | (v6e/v7x) |
|---|---|---|---|
| Упакованный транспон | Все | 17 | 4 |
| Стандартный транспонирование | B32 Транспонирование | 9 | 4 |
| B16 Транспонирование (сегментированное/сжатое) | 17 | 4 |
Пул исполнительных блоков (EUP)
Инструкции EUP представляют собой векторные математические функции (например, tanh , log , exp ).
| Категория инструкций | (v5e/v5p) | (v6e/v7x) |
|---|---|---|
Векторная математика ( tanh , exp и т. д.) | 2 | 1 |
(СТАРАЯ ВЕРСИЯ) Включение настраиваемой видимости звонков
Для включения пользовательского профилирования вызовов необходимо установить следующие флаги XLA при запуске рабочей нагрузки:
-
--xla_enable_custom_call_region_trace=true: Этот флаг включает трассировку для регионов, содержащих пользовательские вызовы. -
--xla_xprof_register_llo_debug_info=true: Этот флаг регистрирует отладочную информацию LLO, что позволяет XProf отображать подробную статистику использования для пользовательского вызова.
Пример :
LIBTPU_INIT_ARGS="--xla_enable_custom_call_region_trace=true --xla_xprof_register_llo_debug_info=true" python your_jax_workload.py
При включении этих флагов в окне просмотра трассировки для каждого ядра или устройства TPU, выполняющего пользовательский вызов, появится новая строка, описывающая использование LLO .
Линия использования LLO
Линия использования LLO предоставляет визуализацию того, как используются аппаратные ресурсы во время выполнения пользовательского вызова. Это особенно полезно для выявления узких мест в пользовательских ядрах (например, написанных на Pallas или Mosaic).
Передовые методы
- Включайте только при необходимости : эти флаги могут увеличить размер захваченного профиля и незначительно повлиять на производительность во время сбора данных. Используйте их в основном для отладки и оптимизации пользовательских вызовов.
- Проверьте наличие информации LLO : если вы включили эти флаги, но не видите строку использования LLO, убедитесь, что ваш компилятор поддерживает регистрацию отладочной информации LLO для вашей конкретной реализации пользовательского вызова.