Профилирование TorchInductor на GPU

В этом разделе представлены полезные команды и рабочие процессы, которые помогут вам углубиться в производительность модели в TorchInductor. Если модель работает не так быстро, как ожидалось, вы можете проверить отдельные ядра модели. Обычно самые интересные ядра — это те, которые занимают большую часть времени работы GPU. После этого вы также можете запустить отдельные ядра напрямую и проверить их производительность. PyTorch предоставляет инструменты для всего вышеперечисленного.

Релевантные переменные окружения

Вы можете использовать следующие переменные окружения в вашем анализе:

• TORCHINDUCTOR_UNIQUE_KERNEL_NAMES

  • По умолчанию TorchInductor называет ядро Triton как ‘triton\_’. При включении этой переменной окружения inductor создаёт более информативное имя ядра в трассировке, например, triton_poi_fused_cat_155, которое содержит категорию ядра (poi для точечных операций) и исходный оператор ATen. Данная настройка отключена по умолчанию, чтобы повысить вероятность попадания в кеш компиляции.

• TORCHINDUCTOR_BENCHMARK_KERNEL

  • Включение этой переменной заставит инструмент кодогенерации inductor проводить бенчмаркинг отдельных ядер Triton.

• TORCHINDUCTOR_MAX_AUTOTUNE

  • Автотюнер inductor проведет бенчмаркинг большего количества triton.Configs и выберет вариант с лучшими результатами производительности. Это увеличит время компиляции с надеждой на улучшение производительности.

Разбиение времени работы модели на GPU

Ниже приведены шаги по разбиению времени выполнения модели на отдельные ядра. Мы возьмем mixnet_l в качестве примера.

1. Запустите скрипт бенчмаркинга для модели:

   TORCHINDUCTOR_UNIQUE_KERNEL_NAMES=1 TORCHINDUCTOR_BENCHMARK_KERNEL=1
   python -u benchmarks/dynamo/timm_models.py –backend inductor –amp
   –performance –dashboard –only mixnet_l –disable-cudagraphs –training
   

   Примечание

   Инструмент полагается на имя ядра для определения его категории. Включение TORCHINDUCTOR_UNIQUE_KERNEL_NAMES имеет решающее значение для этого.

2. В выходном журнале найдите строки:

   **Compiled module path:
   /tmp/torchinductor_shunting/qz/cqz7hvhood7y3psp7fy6msjxsxyli7qiwiybizdwtjw6ffyq5wwd.py**
   

Каждая скомпилированная модуль имеет свою строку. Если нет дополнительных разрывов графа, в журнале будет видно 2 таких строки, одна для прямого графа, а другая для обратного.

Для нашей примерной команды мы получим следующие скомпилированные модули для прямых и обратных графов соответственно:

• https://gist.github.com/shunting314/c2a4d8a28b00fcb5586d0e9d9bf77f9f
• https://gist.github.com/shunting314/48efc83b12ec3ead950052e4a0220b10

3. Теперь мы можем углубиться в производительность каждого отдельного скомпилированного модуля. Для иллюстрации выберем модуль прямого графа. Для удобства назовем его fwd.py. Запустите его напрямую с аргументом -p.

   **> python fwd.py -p**
   

Полный выходной лог см. в этом примере gist.

В выводе вы можете заметить следующее:

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

Трассировка Chrome для профиля записана в /tmp/compiled_module_profile.json

Загрузка трассировки в Chrome (перейдите по адресу chrome://tracing в браузере Chrome и загрузите файл, как предложит интерфейс) покажет интерфейс следующим образом:

Вы можете приближать и отдалять масштаб, чтобы проверить профиль.

• Мы сообщаем процент времени, в течение которого GPU занят, через строку журнала, например:

  Процент времени, когда GPU занят: 102,88%

  Иногда вы можете увидеть значение, большее 100%. Причина в том, что PyTorch использует время выполнения ядра с включённым профилированием, а при отключённом профилировании — время стен. Профилирование может немного исказить время выполнения ядра. Но в целом это не должно быть большой проблемой.

  Если мы запустим модель как densenet121 с небольшим размером пакета, мы увидим низкий процент времени, когда GPU занят:

  (Forward graph) Percent of time when GPU is busy: 32.69%
  

  Это означает, что модель имеет большой накладные расходы на CPU. Это согласуется с тем фактом, что включение cudagraphs значительно улучшает производительность densenet121.

• Мы можем разбить время работы GPU на разные категории ядер. В примере mixnet_l мы видим

  • точечные ядра занимают 28,58%
  • ядра сокращения занимают 13,85%
  • постоянные ядра сокращения занимают 3,89%
  • остальные ядра cutlass/cudnn для mm/conv занимают 56,57%

  Эта информация может быть найдена в строке сводки (последняя строка) отчёта для каждой категории ядер.

• Мы также можем перейти к определённой категории ядер. Например, давайте проверим ядра сокращения:

  

  Мы можем увидеть упорядоченную таблицу времени выполнения для каждого отдельного ядра сокращения. Мы также видим, сколько раз выполняется ядро. Это полезно по нескольким причинам:

  • Если ядро занимает лишь незначительное время, например 0,1%, то улучшение его принесёт не более 0,1% общего прироста. На это не стоит тратить много усилий.
  • Если ядро занимает 2% времени, то улучшение его в 2 раза принесёт 1% общего прироста, что оправдывает затраченные усилия.

Бенчмаркинг отдельного ядра Triton

Допустим, мы хотим более подробно изучить triton_red_fused\__native_batch_norm_legit_functional_16, которое является самым дорогим ядром сокращения и занимает 2,19% от общего времени стен для прямого графа.

Мы можем найти имя ядра в fwd.py, и найти комментарий, например:

# путь к ядру: /tmp/torchinductor_shunting/jk/cjk2vm3446xrk7rth7hr6pun7xxo3dnzubwcn6ydrpifal4eykrz.py

Для удобства переименуем его в k.py. Вот содержимое файла: ссылка.

k.py — это автономный Python-модуль, содержащий код ядра и его бенчмаркинг.

Прямой запуск k.py выдаст время выполнения и пропускную способность:

Мы можем проверить, помогает ли max-autotune этому ядру, запустив:

**TORCHINDUCTOR_MAX_AUTOTUNE=1 python /tmp/k.py**

Мы также можем временно добавить больше эвристик сокращения и снова запустить скрипт, чтобы проверить, как это помогает с ядром.