FeatherAttn: ускорение RDNA3 в 4 раза для LLM-тренинга

Прорыв в оптимизации для AMD RDNA3

Разработчик woct0rdho представил пакет FeatherAttn, который решает критическую проблему производительности при обучении больших языковых моделей (LLM) на графических процессорах архитектуры RDNA3 и RDNA 3.5. В отличие от стандартных решений, новое ядро фокусируется на максимальной эффективности операции обратного распространения (attention backward), которая часто становится узким местом в цикле обучения.

Конкретные метрики производительности

Тестирование проводилось на платформе Strix Halo. Результаты демонстрируют значительное превосходство над оптимизированным ядром AITER FlashAttn Triton bwd kernel. Ключевые показатели скорости ускорения (speedup) зависят от размерности головы (head dimension):

Интеграция и совместимость

FeatherAttn позиционируется как готовое решение для разработчиков. Пакет предоставляет обучаемую функцию внимания, которая является drop-in replacement (прямой заменой) для стандартной функции flash_attn_func. Архитектура гибридная:

• Прямое распространение (FWD): использует ядро AITER.
• Обратное распространение (BWD): использует высокопроизводительное ядро FeatherAttn.
• Fallback: при отсутствии поддержки специфических фич в кастомном ядре BWD, система автоматически откатывается к стандартному ядру AITER, обеспечивая стабильность.

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

Философия «Feather» и планы на FP8

Название проекта отсылает к идее быстрого выполнения операций матричного умножения (matmul) в формате FP8 на архитектурах RDNA3/3.5, которые не имеют нативной поддержки этого формата (см. проект ComfyUI-FeatherOps). Однако текущая версия FeatherAttn для attention-механизмов не использует квантование FP8 ни в прямом, ни в обратном проходе. Работа над внедрением FP8 для оптимизации памяти и вычислений в обоих направлениях находится в стадии активной разработки.

Источник: Github ↗