QAD от NVIDIA: разбираюсь, почему 4-битная квантизация перестала всё ломать

На прошлой неделе NVIDIA выложила отчёт про QAD и я его проигнорировал. Потому что каждый месяц кто-то "решает квантизацию" и каждый раз на практике всё не так радужно.

Но потом коллега скинул табличку с AIME и я залип.

Контекст такой. У нас на проде крутится модель на FP8, всё более-менее ок, но менеджмент хочет запихнуть 49B модель туда, где сейчас живёт 20B. Потому что "ну там же Blackwell, там же FP4, давайте". Я месяц назад попробовал наивно квантизовать в 4 бита — модель начала нести пургу на математике. Не сильно, но заметно. Закрыл тему, сказал что FP4 не готов.

И тут этот отчёт.

Что вообще происходит

NVFP4 — это 4-битный формат NVIDIA. E2M1 если кому интересно (2 бита экспонента, 1 бит мантисса, знак). На Blackwell обещают 2-3x throughput по сравнению с FP8. Звучит красиво.

Проблема: при 4 битах модели начинают тупить. Особенно на задачах где нужна точность — математика, код, логические цепочки. Чем меньше модель, тем хуже.

Стандартное решение — QAT, quantization-aware training. Берёшь квантизованную модель, дообучаешь, профит. Работает со времён мобильных свёрточек, должно работать и тут.

Не работает.

Вот цифры которые меня зацепили (Nemotron 3 Nano на AIME25):

• BF16: 89.1

• Просто квантизовали (PTQ): 85.0

• QAT: 83.3

• QAD: 87.9

QAT хуже чем ничего не делать. Это как?

Почему QAT ломается

Я сначала думал что они накосячили с гиперпараметрами. Или данные не те. Но нет, там есть объяснение и оно имеет смысл.

Современные LLM готовят сложно. Pretrain, потом SFT, потом RL (GRPO, RLHF, что угодно). Каждый этап меняет модель. RL особенно — он учит модель думать определённым образом, выстраивать цепочки рассуждений.

Когда ты делаешь QAT, ты берёшь какие-то данные и минимизируешь cross-entropy. Это по сути ещё один раунд SFT. И он переписывает то, чему модель научилась на RL.

У меня не было такой интуиции раньше. Я думал — ну данные те же, лосс тот же, всё должно восстановиться. Но нет. Cross-entropy это про "правильные токены", а не про "как именно модель приходит к ответу".

Что делает QAD

Идея тупая до безобразия: вместо того чтобы учить квантизованную модель правильным ответам, учим её копировать поведение оригинала.

Берём BF16 модель (учитель), квантизованную модель (ученик), прогоняем один и тот же текст через обе. Минимизируем KL-дивергенцию между распределениями.

L = KL(p_teacher || p_student)

Не рокет саенс. Дистилляция известна сто лет. Но почему-то для квантизации LLM это не было стандартом.

Самое интересное — табличка из статьи, которая всё объясняет:

Смотри: QAT достигает того же cross-entropy что и оригинал. По этой метрике идеально. Но KL-дивергенция 0.311 — распределение выходов совсем другое.

QAD наоборот: cross-entropy чуть хуже, зато модель ведёт себя практически идентично оригиналу.

Цифры на разных моделях

Llama Nemotron Super 49B (после SFT):

На AIME25 разница между QAT и QAD — 4 пункта. Это много, там каждая задача на вес золота.

AceReason 7B (после RL — вот тут самое мясо):

QAT уронил модель на 17 пунктов. Семнадцать. Это катастрофа. QAD почти восстановил до оригинала.

Вот это меня убило

Дальше в статье есть эксперимент, от которого у меня немного поехала крыша.

AceReason обучена на математике и коде. Логично делать QAD на математике и коде. А что если взять только код?

QAD на данных только по коду восстанавливает математику почти так же хорошо, как QAD на полных данных.

Ладно, может код и математика как-то связаны. Дальше хуже:

На рандомных токенах. РАНДОМНЫХ. И модель не развалилась.

Для сравнения — PTQ без всего даёт 58.7. То есть QAD на мусоре лучше чем ничего.

Я пока не понимаю почему это работает. Есть гипотеза что KL-лосс как-то иначе гонит градиенты, но надо думать.

Практические заметки

Из статьи, плюс мои догадки:

Learning rate. Для моделей после SFT — низкий, 1e-6 или меньше. Для моделей после RL — выше, 1e-5. Контринтуитивно, но видимо RL-модели более устойчивы к изменениям.

Данные. Llama Nemotron 49B обучали на ~0.3B токенов. Это смешно мало по сравнению с оригинальным post-training. AceReason 7B — ~0.8B токенов.

Температура. T=1 для обоих. Не выпендривайтесь.

Что квантизовать. Все линейные слои в attention и MLP. Если у вас гибрид с Mamba — attention лучше оставить в BF16, SSM слои плохо переносят квантизацию.

Учитель побольше. Не работает. Пробовали 12B учителя для 9B ученика — хуже чем просто 9B на себя.

Что я об этом думаю

Честно — настроен скептически к 4 битам для серьёзных задач. Слишком много раз видел как квантизация убивает тонкие способности модели. Отладил, работает, квантизовал, сломалось, откатил.

Но эти результаты выглядят правдоподобно. Не "мы решили квантизацию", а "мы нашли способ не ломать RL при квантизации". Это конкретная проблема с конкретным решением.

Что меня напрягает:

1. Все эксперименты на моделях NVIDIA. Хочу увидеть это на Llama/Qwen/Mistral.

2. NVFP4 специфичен для их железа. На десктопе это не запустишь.

3. Бенчмарки — это бенчмарки. На реальных задачах может быть иначе.

Но идея с KL вместо CE — это я точно попробую. Даже если не FP4, для FP8 тоже должно помочь.

Что выложили

Чекпоинты для нескольких моделей, код для Megatron-LM и HuggingFace. Ссылки в оригинальном отчёте, не буду копипастить.

Жду когда кто-то портирует на llama.cpp. Вот тогда будет понятно, насколько это реально работает вне экосистемы NVIDIA.

Если кто-то уже пробовал — напишите. Интересно сравнить результаты.

https://t.me/tokensaway — иногда пишу про такое, иногда о том, как LLM думают, или просто притворяются.