Quartet II — обучение в NVFP4 с точными несмещёнными градиентами

В последние годы наблюдается тренд на снижение точности, в которой обучают модели. Если ещё недавно использовали преимущественно BF16, то в 2024-м, например в DeepSeek-V3 — уже FP8, а в 2025-м переходят и на FP4. Снижение точности, однако, требует применения трюков. Сегодня разберём статью, в которой описывается новый метод обучения в NVFP4.

Прошлым летом мы придумали Quartet — первый метод обучения в FP4. Он использовал формат данных MXFP4, который поддерживается на уровне железа в GPU, вроде NVIDIA Blackwell. Quartet использовал случайные вращения тензоров для борьбы с аутлаерами и показал, что FP4 обучение даёт выигрыш при фиксированном бюджете относительно FP8.

Но есть и другой формат, основанный на FP4, — NVFP4, который NVIDIA создала специально для LLM. Он использует более точные FP8-скейлы. На его основе прошлой осенью NVIDIA предложила собственный, более точный рецепт обучения в FP4. Одной из его фишек было использование стохастического округления (stochastic rounding, SR) на backward pass — методики, которая делает оценки градиентов несмещёнными ценой более высокого уровня их шума. Мы задались вопросом: есть ли способ достичь несмещённых оценок без шумного SR в FP4?

Оказывается, вращения, которые помогали нам с аутлаерами, могут помочь и с несмещённостью. Квантизуя случайно повёрнутый тензор, мы усредняем ошибку по всем возможным направлениям, и получившаяся оценка будет коллинеарена оригинальному тензору. Отнормировав квантизованный тензор, коллинеарность можно превратить в несмещённость. Но FP8-скейлы не позволят просто так отнормировать квантизованный тензор — они слишком грубые, чтобы впитать точную перенормировку порядка 0,94~1,06.

Какой выход? Вшить эти скейлы в FP8 с помощью SR. Обратите внимание, что делается это не на каждый FP4-элемент, как в случае с методом NVIDIA, а только в FP8-скейлы, что даёт ошибку в сто раз меньше. В итоге, полная ошибка получается вдвое меньше, чем у простого SR. Мы объединили метод с более точной квантизацией на forward pass и получили Quartet II — SOTA обучения в NVFP4.

Таким образом, случайные вращения не просто борются с аутлаерами, а служат важным источником несмещённости в оценке градиентов, который уменьшает ошибку и ускоряет сходимость FP4-обучения. К статье прилагаются кернелы для 5090 и B200, чтобы можно было попробовать NVFP4 на своих моделях и задачах.

Разбор подготовил Андрей Панфёров, автор канала black_samorez

Душный NLP