Пользовательские вызовы XLA позволяют выполнять пользовательские ядра или операции, которые изначально не поддерживаются XLA. Чтобы получить представление о производительности этих пользовательских вызовов в средстве просмотра трассировки , можно использовать специальные флаги XLA для включения подробной трассировки и отладочной информации LLO (низкоуровневого оптимизатора).
Включение настраиваемой видимости звонков
Для включения пользовательского профилирования вызовов необходимо установить следующие флаги 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 для вашей конкретной реализации пользовательского вызова.