🔬 Метод

Форматы FP4

Первым делом исследуют конфигурации форматов FP4. Напомню, что MXFP4 квантизует веса группами по 32 и квантизует скейлы в E8M0, а NVFP4 группами по 16 и скейлы в E4M3. Авторы фиксируют размер группы 16 и перебирают варианты квантизации скейлов от E1M6 до E8M0. Обучают Llama-like LLM на 350M параметров и замечают, что при E4M3/E3M4 скейлах достигается минимальный лосс (при фиксированном числе итераций). Из всех конфигураций расходится только E1M6 (c cамым узким диапазоном). В дальнейшем везде используют E4M3. Блоки размера 16 выбирают так как при больших лосс сходится хуже, а меньшие уже не дают профита.

Стохастическая квантизация

Квантизовать можно к ближайшему значению, а можно стохастически - вверх или вниз, с вероятностью, зависящей от расстояния до соседа.

Ребята из интела перебирают разные варианты детерминистического и стохастического квантования для весов/градиентов и активаций и получают, что лучше всего сходится вариант с детерминированной квантизацией весов и активаций на прямом проходе, и стохастической для градиентов и активаций на обратном проходе, Роль стохастики в квантизации - уменьшить bias, возникающий из-за округления тензоров.

В ходе оптимизации сигнал от градиента постепенно убывает и с какого-то момента перекрывается шумом оптимизации. Не мудрствуя лукаво, авторы предлагают обучать небольшое время с градиентами в более высокой (bf16) точности (на прямом проходе все еще FP4). И это позволяет сойтись до уровня half-precision обучения за то же суммарное число итераций.

🧪Эксперименты

Обучают семейство моделей архитектуры Llama-2 на датасете Красная Пижама. В главном эксперименте учат модель размера 7B на 1Т токенах причем не абы на чем, а на ускорителях Intel Gaudi2 (сыр 🧀 тут ни при чем, это в честь архитектора)

Обучение идет без спайков, лосс отстает несколько от bf16 бейзлайна, но нагоняет после короткой фазы с более точными градиентами (QAF).

0-шоты без QAF чуть хуже безйлайна, с QAF - такие же примерно. Впрочем, все равно оно лишь чуть лучше рандома)

💡 Выводы

Выглядит как очередной аргумент перейти на обучение LLM в FP4. Сам по себе метод выглядит не шибко изощренно, хотя необходимость QAF для лучших результатов несколько противоречит названию статьи (надо было назвать FP4 Most the way). Quartet в этом отношении по изящнее. Интересно, кто из крупных игроков выложит первый техрепорт про полное обучение серьезной модели в FP4? Ставлю либо на Нвидию, либо на Moonshot.