Forwarded from Yandex Research
У нас вышла новая статья FLAMIA: FLAsh MIcroscaling Attention for Fast and Causally Safe Training on Blackwell про ускорение attention на GPU NVIDIA Blackwell. FlashAttention-4 уже хорошо оптимизирован под архитектуру Blackwell, но его реализация в BF16 не использует весь потенциал вычислений в низкой точности FP4/FP8. Мы разработали FLAMIA — набор attention-кернелов, которые позволяют ускорить вычисления при обучении и решить ряд проблем, возникающих при квантизации. Что сделали? Подобрали точность вычислений. Выяснили, что квантовать всё в 4 бита не всегда выгодно: затраты на квантование могут перекрывать ускорение матричных умножений. Поэтому реализовали два режима — MXFP8 и смешанный NVFP4/FP8. Во втором используем NVFP4 для QKᵀ, а FP8-битные форматы — для остальных матричных умножений. Квантование блоками позволяет повторно использовать представления на обратном проходе. Исследовали causal leakage. При низкобитном attention общие масштабы квантования могут учитывать будущие токены, из-за чего информация о них влияет на предыдущие позиции, несмотря на causal mask. Мы предложили считать масштабы, специфичные для каждого ключа, чтобы устранить эту проблему. Оптимизировали varlen attention и упаковку последовательностей. Реализовали обработку документов разной длины без дополнительной перепаковки данных. Что получилось? MXFP8 ускоряет attention с учётом квантования в 1,10–1,40 раза на B200 и 1,06–1,42 раза на B300 по сравнению с самыми быстрыми из протестированных реализаций: FlashAttention-4 в BF16, cuDNN, LP-FA4 и FlashInfer NVFP4. Проверили FLAMIA на обучении Llama-подобной модели с 1B параметров на 100B токенов. Обучение оставалось стабильным, а итоговые метрики были близки к BF16. статья · код · Hugging Face ⏩Yandex Research
Open in Telegram