Nemotron 3.5 Lightning: Оптимизация NVFP4 через QAD

В эпоху, когда большие языковые модели (LLM) становятся всё более ресурсоемкими, разработчики и исследователи сталкиваются с фундаментальной дилеммой: как сохранить высокую точность и интеллектуальные способности модели, одновременно drastically снижая требования к памяти и вычислительной мощности? Ответ на этот вопрос лежит в плоскости квантования — процесса преобразования весов модели из высокоточных форматов (например, BF16) в более компактные. Однако простое сжатие часто приводит к необратимой потере качества. Именно здесь на сцену выходит передовая методика, представленная командой NVIDIA в рамках семейства моделей Nemotron 3.5 Lightning.

Новый чекпоинт Nemotron 3.5 Lightning NVFP4 демонстрирует прорывной подход к оптимизации. Используя технику квантования NVFP4, модель сжимается с 66 ГБ до всего 22 ГБ, что позволяет увеличить пропускную способность до 4 раз. Но главное достижение — это сохранение точности, близкой к базовому уровню BF16. Секрет этого успеха кроется не в простом постобучающем квантовании (PTQ), а в сложном двухэтапном процессе, включающем дистилляцию, Aware квантования (QAD). В этой статье мы подробно разберем, как работает этот пайплайн, какие технические решения применяются и как разработчики могут воспроизвести эти результаты для своих проектов.

01Что такое дистилляция, Aware квантования (QAD)?

Чтобы понять, почему QAD превосходит традиционные методы, нужно рассмотреть архитектуру процесса. QAD — это не просто алгоритм сжатия, а полноценный процесс обучения, в котором участвуют две модели: «учитель» (teacher) и «ученик» (student). «Учитель» — это исходная полноточная модель в формате BF16, которая сохраняет все нюансы своих знаний. «Ученик» — это квантованная версия модели, которая должна научиться имитировать поведение учителя, несмотря на ограничения низкой разрядности.

Процесс состоит из двух ключевых этапов. На первом этапе применяется постобучающее квантование (PTQ). Это быстрый метод, при котором веса модели преобразуются в низкоточный формат (например, W4A16 или NVFP4) без дополнительного обучения. Однако агрессивное квантование на этом этапе неизбежно вносит «шум» и снижает точность. На втором этапе, который и является сутью QAD, квантованная модель (ученик) проходит через процесс дистилляции. Она обучается на данных, используя функцию потерь KL-divergence, чтобы сравнить логиты (вероятностные распределения) ученика и учителя. При этом модель учителя остается замороженной (frozen), выступая в роли эталона.

Этот подход позволяет модели-ученику не просто предсказывать следующее слово, но и воспроизводить сложное поведение полноточной модели. Даже если квантование на первом этапе было очень агрессивным и привело к значительной потере точности, второй этап QAD способен восстановить почти все потерянные качества. Это открывает возможность использовать более жесткие параметры квантования, которые ранее считались неприемлемыми, обеспечивая максимальную экономию памяти и вычислений.

Почему QAD лучше простого PTQ?

Традиционное постобучающее квантование (PTQ) часто ограничивается целью восстановления медианной точности на уровне 99% и выше. Если вы используете только PTQ, вы вынуждены искать компромисс: либо сохранить точность, пожертвовав степенью сжатия, либо сильно сжать модель, но получить неприемлемый спад качества. QAD меняет эту парадигму. Поскольку мы знаем, что дистилляция восстановит часть точности, на этапе PTQ можно позволить себе более агрессивное квантование. Целевой показатель восстановления точности после PTQ снижается до 95-99%. Этот «запас прочности» в виде небольшого падения точности подтверждает, что мы максимально использовали потенциал сжатия, а QAD затем закрывает этот разрыв.

02Техническая реализация: Пайплайн NVIDIA Model Optimizer

Для реализации этого сложного процесса NVIDIA предоставляет инструменты в составе NVIDIA Model Optimizer. Рассмотрим пошаговый процесс разработки чекпоинта Nemotron 3.5 Lightning NVFP4. Исходной точкой выступает базовая модель NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16, которая служит учителем.

Шаг 1: Получение чекпоинта PTQ

Первый шаг — создание модели-ученика путем применения PTQ к базовой модели. В случае с Nemotron 3.5 Lightning команда исследователей протестировала несколько рецептов квантования, различающихся методами калибровки и степенью агрессивности квантования слоев Mamba и KV-кэша. Особое внимание уделялось слоям Mamba: вместо более безопасного формата FP8, они были квантованы до W4A16. Это решение позволило значительно повысить пропускную способность без критического падения точности, которое впоследствии восстанавливалось через QAD.

Калибровка проводилась на 1000 образцах данных с использованием последовательности длиной от 8k до 128k. Исследования показали, что длина последовательности 32K обеспечивает наилучшие результаты для рецепта «four_over_six». Важно отметить, что все рецепты квантовали слой lm_head до W4A16 (так называемый «верный lm_head»), в то время как слои проекции внимания оставались в BF16. Это балансировало между эффективностью и точностью.

Ниже приведены пять основных рецептов PTQ, которые были протестированы. Обратите внимание на различия в методах калибровки (dynamic vs static) и форматах кэша KV.

Для воспроизведения этого шага на своей модели можно использовать следующий код на Python с библиотекой Model Optimizer:

import modelopt.torch.quantization as mtq

# define forward loop with dataset
def forward_loop(model):
    for batch in calib_dataloader:
        model(batch)

# Quantize base model to NVFP4 to create the PTQ student checkpoint
# Example uses W4A16_NVFP4_CFG for quantization
model = mtq.quantize(model, mtq.W4A16_NVFP4_CFG, forward_loop=forward_loop)

Шаг 2: Дистилляция, Aware квантования (QAD)

После получения чекпоинта PTQ начинается самый важный этап — обучение. Здесь критически важны два параметра: длина последовательности и набор данных. Для сохранения производительности на длинных контекстах была выбрана длина последовательности 522K токенов (хотя для первоначальных абляций использовалось 256K). В качестве данных рекомендовано использовать открытые наборы Nemotron-Post-Training v1 и v2, которые хорошо репрезентируют распределение данных, использованное в оригинальной работе.

Процесс дистилляции происходит следующим образом: на каждом шаге один и тот же бат данных пропускается через замороженную модель учителя (BF16) и модель ученика (NVFP4). Ученик обучается минимизировать расхождение KL-divergence между логитами учителя и ученика. Обучение проводилось с постоянной скоростью обучения 5e-6, без warmup, с отключенным dropout и clipping градиентов на уровне 1.0. Распределение вычислений осуществлялось на двух узлах по 8 GPU (TP=2, EP=4).

03Управление масштабами квантования: Dynamic vs Frozen

Одним из самых тонких моментов в процессе QAD является выбор стратегии управления масштабами квантования (quantization scales) во время обучения. Существует два основных подхода, и выбор между ними зависит от того, какой рецепт PTQ был использован на первом этапе.

Dynamic Scale QAD

Этот подход применяется, если чекпоинт PTQ был откалиброван с использованием метода «max» (например, рецепты max или mamba_fp8_max). В этом случае веса и активации имеют динамические масштабы. Во время обучения QAD масштабы пересчитываются на каждом шаге на основе текущих тензоров. Это позволяет как весам, так и масштабам адаптироваться по мере обучения модели. Это гибкий подход, но он требует больше вычислительных ресурсов на каждом шаге.

Frozen Scale QAD

Этот подход используется для чекпоинтов, откалиброванных с помощью метода MSE (среднеквадратичной ошибки), таких как MSE, four_over_six или four_over_six + NVFP4 KV. Метод MSE находит оптимальные масштабы, минимизирующие ошибку квантования, но этот процесс поиска слишком ресурсоемок для повторения на каждом шаге обучения. Поэтому масштабы, найденные на этапе PTQ, «замораживаются» (freeze). Во время QAD обновляются только веса модели, а масштабы остаются фиксированными на значениях, полученных при калибровке. Это делает процесс обучения более стабильным и предсказуемым.

04Результаты и оценка качества

Эксперименты проводились на промежуточных чекпоинтах A (SFT) и B (RL). Результаты демонстрируют, что QAD consistently outperforms PTQ в медианном восстановлении точности и особенно в агентных и кодинговых бенчмарках. Важно отметить, что как PTQ, так и QAD чекпоинты использовали одинаковый агрессивный формат W4A16 и занимали одинаковый объем памяти (21.19 ГБ), в то время как базовая модель BF16 занимала 65.85 ГБ. Любое улучшение качества в QAD обусловлено исключительно методом дистилляции, а не увеличением размера модели.

Ниже приведены результаты сравнения для чекпоинта A:

Как видно из таблицы, QAD не только восстанавливает точность, но и в некоторых случаях (например, AIME 2025 и SciCode Subtask) превосходит базовую модель BF16. Это доказывает, что агрессивное квантование в сочетании с правильной дистилляцией может привести к улучшению обобщающей способности модели, а не просто к восстановлению исходных показателей.

05Что это значит на практике

Для разработчиков, работающих с LLM в условиях ограниченных ресурсов (например, при локальном запуске на потребительских GPU или при развертывании высоконагруженных сервисов), результаты работы с Nemotron 3.5 Lightning и QAD имеют прямое практическое значение. Сжатие модели в 3 раза (с 66 ГБ до 22 ГБ) при сохранении или даже улучшении качества означает, что сложные модели теперь можно запускать на более доступном оборудовании. Увеличение пропускной способности в 4 раза позволяет обслуживать больше запросов в секунду, что критически важно для реального времени.

Использование NVIDIA Model Optimizer и Megatron-Bridge предоставляет end-to-end рабочий процесс, который делает эти передовые техники доступными для широкой аудитории. Разработчикам больше не нужно изобретать велосипед: они могут использовать готовые рецепты, такие как four_over_six с NVFP4 KV, чтобы достичь максимальных результатов. Важно помнить, что выбор между Dynamic и Frozen scale должен быть осознанным и зависеть от выбранного метода калибровки PTQ. Экспериментируя с разными рецептами и стратегиями масштабов, можно найти оптимальную конфигурацию для каждой конкретной задачи, будь то генерация текста, кодинг или анализ данных.

В заключение, QAD представляет собой мощный инструмент в арсенале разработчика AI. Он позволяет преодолеть компромисс между эффективностью и точностью, открывая путь к созданию более быстрых, дешевых и доступных больших языковых моделей. С появлением таких инструментов, как NVIDIA Model Optimizer, барьер для внедрения передовых методов оптимизации значительно снижается, что способствует демократизации доступа к мощным AI-решениям.

Источник: NVIDIA Developer ↗