🔬 Метод

Предложенный NVIDIA рецепт обучения в FP4 включает следующие ключевые составляющие:
📌 Адамаровы вращения для борьбы с выбросами и большей стабильности.
📌 Стохастическое округление градиентов по параметрам.
📌 2d скейлинги в блоках 16x16, которые позволяют нивелировать нестыковку между поведением модели на прямом и обратном проходе.
📌 Несколько первых блоков в начале и конце модели держат в bfloat16 (в сумме 15% от размера модели). Если их квантизовать - обучение нестабильно.

В конце, подобно статье от Intel, предлагают учить модель небольшое число итераций в bf16, чтобы сократить разрыв по качеству между bf16 и fp4 обучением. Конечная модель уже не такая компактная и быстрая на инференсе, но типа сэкономили на обучении.

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

Обучают очередной Nemotron на некоей смеси данных с большой долей синты с warmup-stable decay расписанием.

Большую часть обучения лосс FP4 не сильно отстает от FP8, но разрыв немного увеличивается в конце, при маленьких шагах обучения.

По бенчам +/- на уровне бейзлайна.

Все компоненты метода важны. Аблейтят довольно странно, стартуя с 3Т токенов убирают одну за другую составляющую. Не очевидно, что эффект будет таким же, если с нуля запустить несколько обучений с отказом от компонент по одной.

MXFP менее эффективен на обучении и достигает того же лосса, что и NVFP, при на 36% большем количестве данных.

💡 Выводы

Рецепт выглядит рабочим. Но по ощущениям все же требуется немало свистоплясок. Кроме того, необходимость держать часть слоев в bf16 несколько осложняет интеграцию и ограничивает максимально достижимое ускорение. Не хватает сравнительного анализа квантизации трансформерных и Mamba-2 слоев, хотя это кажется важным и интересным. Есть мнение, что SSMки тяжелее квантизуются.