Я посмотрел, как Яндекс FLAMIA ускоряет вычисление attention при обучении трансформеров на ускорителях NVIDIA Blackwell. Сравнение с FlashAttention-4 впечатляет, но главный результат относится к конкретным низкоточным режимам, а не ко всему обучению модели.
Насколько FLAMIA быстрее FlashAttention-4 на Blackwell
FLAMIA ускоряет полный вызов attention, включая прямой и обратный проходы, в 1,10–1,40 раза на B200 и в 1,06–1,42 раза на B300. В отдельном обратном проходе с NVFP4 прирост достигает 2,01 раза на B200 и 2,00 раза на B300. Эти показатели авторы сравнивают с самым быстрым протестированным базовым решением для каждой задачи.
| Решение | Платформа и формат | Измерение | Результат |
|---|---|---|---|
| FlashAttention-4 | B200, BF16 | Прямой проход | до 1613 TFLOPs/s |
| FLAMIA | B200, MXFP8 | Прямой и обратный проходы | 1,10–1,40x |
| FLAMIA | B300, MXFP8 | Прямой и обратный проходы | 1,06–1,42x |
| FLAMIA | B200, NVFP4 | Обратный проход | 1,22–2,01x |
Точность FLAMIA и проверка на данных FineWeb-Edu
FLAMIA использует MXFP8 и смешанный режим NVFP4/FP8. Для этих режимов авторы отдельно решали проблемы накладных расходов на квантизацию, утечки причинной информации через общие масштабы и несовпадений прямых и обратных вычислений. Это важная часть результата: низкая точность полезна, только если ускорение не нарушает ход обучения.
Проверку провели на Llama-подобной модели с 1 млрд параметров, обученной на 100 млрд токенов FineWeb-Edu. Итоговый loss составил 2,8876 у базовой версии BF16 FlashAttention-4, 2,8884 у MXFP8 и 2,8885 у смешанного режима. Разница меньше 0,001, но эти цифры относятся только к данному эксперименту, а не к любому масштабу модели и набору данных.
Код FLAMIA опубликовали в открытом репозитории Yandex Research под лицензией BSD 3-Clause. Поддерживаются B200 и B300. Переменные по длине упакованные документы можно обрабатывать без предварительной переупаковки. Для компаний в СНГ практический эффект зависит от доступа к таким ускорителям и совместимости программного стека. Данных о снижении стоимости обучения или замерах на других GPU публикация не приводит. FLAMIA может подойти командам, которые уже обучают модели на Blackwell и готовы проверить свои нагрузки на опубликованном коде.
