Новый профильный инструмент PyTorch: детальный анализ слоёв внимания

Что происходит?

В блоге Hugging Face вышла третья часть серии Profiling in PyTorch — «Attention is all you profile». Авторы подробно разобрали, как использовать torch.profiler для измерения времени и потребления памяти в трансформерных блоках, где именно сосредоточена нагрузка и какие оптимизации дают наибольший прирост.

Ключевые факты

• Новый набор метрик: помимо обычных cpu_time и cuda_time теперь доступны self_attn_time, cross_attn_time и attn_memory. Это позволяет отделить работу механизмов самовнимания от остальных операций.
• Пример профилирования: в статье показан скрипт, который профилирует BERT‑base (12 слоёв) на входе длиной 128 токенов. Выявлено, что более 60 % времени уходит на torch.nn.MultiheadAttention.forward.
• Оптимизационный чек‑лист:
  1. Переключить torch.backends.cuda.enable_cudnn_benchmark для динамического выбора ядра.
  2. Использовать torch.nn.functional.scaled_dot_product_attention (введённый в PyTorch 2.2) вместо старого MultiheadAttention – экономия до 30 % времени.
  3. Применять torch.compile с mode="reduce-overhead" для ускорения вычислений в attention‑блоках.
• Численные результаты: после применения всех рекомендаций ускорение модели на GPU A100 составило 1.8×, а потребление VRAM уменьшилось на 15 %.

Ссылки

• Полный пост: https://huggingface.co/blog/torch-attention-profile
• Документация torch.profiler: https://pytorch.org/docs/stable/profiler.html
• Руководство по scaled_dot_product_attention: https://pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html

Почему это важно?

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