Сделать дистилляцию знаний достаточно дешёвой для масштабного применения

Antonio Tiene

AntonioTN

MultiverseComputingCAI

Iker García-Ferrero

Iker

MultiverseComputingCAI

Дистилляция знаний — хорошо известный в машинном обучении метод, при котором меньшую модель-ученика обучают повторять производительность большей модели-учителя. Благодаря недавней волне открытых больших языковых моделей, таких как gpt-oss, Qwen, GLM или Kimi, эта тема вновь стала актуальным направлением исследований. Развёртывание таких очень крупных моделей обходится дорого: недавно выпущенная модель Kimi-K3 содержит 2,8 триллиона параметров и требует около 3 ТБ VRAM только для загрузки. Поэтому сжатие моделей до меньшего размера и восстановление их исходных возможностей с помощью дистилляции знаний стали стандартной практикой; такие компании, как Nvidia (Nemotron 3 Puzzle 75B) и Multiverse Computing (Hypernova 60B), недавно выпустили высококачественные сжатые модели.

Именно этап дистилляции в наибольшей степени определяет итоговое качество, но обычно он же является самой дорогой частью конвейера. Одновременная загрузка учителя и ученика, а также построение распределения вероятностей по всему словарю для каждого токена требуют огромного объёма VRAM и обычно возможны только при использовании сотен GPU и тщательно продуманных стратегий тензорного параллелизма. В нашей последней статье Efficient Knowledge Distillation for LLMs: Offline Top-K Logits and a Fused Chunked KL Loss мы решаем эту проблему с помощью двух системных изменений: однократного кэширования top-K логитов учителя, благодаря чему учителю больше не нужно находиться в памяти вместе с учеником, и нового эффективного по памяти критерия KL-дивергенции, который позволяет никогда не материализовывать матрицу полного размера «словарь × длина последовательности». Это снижает использование VRAM значительно сильнее, чем стандартные реализации в библиотеках вроде PyTorch или NVIDIA Megatron-Bridge. Вместе эти два изменения настолько сокращают стоимость обучения, что восстановление моделей с длинным контекстом становится возможным на одном GPU, а масштабные эксперименты — практически доступными.

Почему восстановление с помощью дистилляции обходится дорого

Стандартная конфигурация — онлайн-дистилляция с использованием критерия расхождения Кульбака — Лейблера (KL-критерия) — одновременно удерживает в памяти и учителя, и ученика. На каждом шаге обучения учитель выполняет полный прямой проход для получения выходного распределения, а ученик обучается ему соответствовать. Это наиболее выразительная конфигурация, поскольку доступно полное распределение учителя, но она же требует больше всего памяти и вычислений: для каждой позиции токена необходимо хранить два тензора полного размера словаря, а учителя приходится пересчитывать на каждом шаге, хотя его поведение не меняется в процессе обучения.

В качестве практического примера рассмотрим gpt-oss-120b со словарём из 201 088 токенов. При длине последовательности 32K и размере батча 4 один только тензор вероятностей учителя имеет форму 4 × 201,088 × 32,768; в формате bfloat16 это уже около 50 ГБ VRAM для одного тензора. Если добавить градиенты, активации, веса модели и состояния оптимизатора, одна итерация дистилляционного обучения может достигать пикового использования примерно 250 ГБ VRAM — больше, чем может предоставить даже GPU H200 или B200. В этой статье мы показываем, что переформулировка KL-критерия для обработки данных по фрагментам снижает эти затраты почти до нуля.

Плотный KL-критерий достигает примерно 250 ГБ, превышая ёмкость одного H200, равную 141 ГБ. При использовании объединённого фрагментированного критерия такой пик не возникает, а пиковое значение составляет около 128 ГБ. Источник: рисунок 1 статьи.

Два системных изменения

Офлайн-дистилляция. Вместо пересчёта учителя на каждом шаге мы вычисляем его выход один раз, кэшируем 100 наиболее вероятных токенов для каждой позиции и обучаем ученика на основе этого кэша. Во время обучения учителю не нужно находиться в памяти и повторно запускаться после создания кэша, поэтому один и тот же кэш можно использовать в многочисленных абляционных экспериментах.

Объединённый фрагментированный KL-критерий. Чтобы понять, почему сам критерий требует больших затрат, представьте, что именно он строит: для каждой позиции токена в последовательности и каждого слова в словаре критерий должен содержать число, описывающее степень расхождения предсказания ученика с предсказанием учителя. Если представить это в виде таблицы, получится одна строка на каждый элемент словаря и один столбец на каждую позицию последовательности. Для словаря из более чем 100 тысяч слов и длинной последовательности такая таблица огромна, а стандартный способ вычисления KL-критерия строит её целиком, прежде чем получить единственное число.

Мы сравниваем три способа вычисления этого же критерия, математически эквивалентных:

• Плотный KL-критерий — классический подход. Он заново строит полную плотную таблицу вероятностей учителя на основе кэшированных top-100 логитов и сравнивает её с собственной плотной таблицей логарифмов вероятностей ученика. Эта версия наиболее близка к тому, как уже работает онлайн-дистилляция, поэтому мы используем её как эталон корректности, однако она дважды удерживает в памяти всю таблицу «словарь × последовательность».
• KL-критерий с фрагментацией прямого прохода сохраняет разреженность учителя (только его кэшированные top-100 логитов для каждой позиции, без расширения до плотной таблицы) и вычисляет критерий по частям, по одному фрагменту позиций последовательности за раз. Это устраняет плотного учителя и плотное сравнение и, как показывают наши тесты, оказывается самым быстрым из трёх методов. Однако у него остаётся одна проблема: собственные логиты ученика — таблица, создаваемая выходным слоем модели, — по-прежнему вычисляются полностью и сохраняются для обратного прохода, поэтому объём памяти всё ещё резко растёт с длиной последовательности.
• Объединённый фрагментированный KL-критерий, наш основной вклад, делает ещё один шаг и объединяет выходную проекцию модели непосредственно с вычислением критерия. Он вообще не создаёт полную таблицу логитов ученика: обрабатывает по одному фрагменту последовательности за раз от начала до конца, преобразует скрытые состояния в логиты для этого фрагмента, добавляет результат к текущему значению критерия и удаляет фрагмент перед переходом к следующему. При обратном проходе каждый фрагмент пересчитывается на лету вместо хранения. Цена этого решения — выполнение проекции дважды: один раз в прямом проходе и один раз в обратном. Зато пиковое потребление памяти растёт только линейно с длиной последовательности, а не резко увеличивается пропорционально полному размеру «словарь × последовательность».

GIF ниже показывает разницу между плотным и объединённым фрагментированным подходами: первый строит всю таблицу сравнения и удерживает её целиком, а второй строит и удаляет по одному фрагменту за раз, поэтому память никогда не превышает объём одного фрагмента.

Мы опубликовали исходный код реализации фрагментированного критерия: github.com/CompactifAI/Full-Chunked-KL-Loss

Что это меняет на практике

В таблице ниже сопоставлены все четыре конфигурации: онлайн-дистилляция и три описанные реализации офлайн-критерия. При сравнении на одном GPU H200 с использованием Llama 3.1 8B Instruct в качестве учителя и модели Llama на 3,2 млрд параметров в качестве ученика при контексте в 8K токенов все четыре метода достигают практически идентичного значения потерь обучения, несмотря на то, что в офлайн-запусках используются только кэшированные top-100 логитов для каждого токена.

Кривые потерь почти полностью совпадают у всех четырёх методов, подтверждая, что офлайн-дистилляция с кэшированными логитами top-100 не уступает онлайн-дистилляции по качеству. Источник: рисунок 2 статьи. При такой длине последовательности объединённый фрагментированный критерий пока не является самым быстрым вариантом: дополнительная проекция в обратном проходе немного снижает скорость, но его реальное преимущество проявляется только с ростом длины контекста, что показано в следующем разделе.

Масштабирование до длинных контекстов

Чтобы нагляднее увидеть закономерность масштабирования, мы провели изолированный тест на небольшой сети выходной проекции (без трансформерного корпуса, только ядро функции потерь). При длине 32K токенов пиковое потребление памяти снижается с 85.2 ГиБ для плотного критерия до 5.45 ГиБ для полностью фрагментированной версии — в 15,6 раза; начиная с 64K токенов плотный критерий полностью выходит из строя. При длине 256K токенов полностью фрагментированный критерий использует 11.6 ГиБ против 134.2 ГиБ у следующего по эффективности фрагментированного варианта и примерно в 3,3 раза быстрее выполняется на каждой итерации.

При дистилляции модели GPT-OSS 20B с контекстом в 32 768 токенов высвободившаяся благодаря объединённому критерию память позволила сократить конфигурацию с четырёх GPU-узлов до одного. Время шага уменьшилось с 57.0 до 12.23 секунды, то есть примерно в 5 раз, а пропускная способность на GPU выросла с 74.2 до 345.7 TFLOP/s.

Получившаяся модель-ученик

Именно эффективная офлайн-конфигурация изначально сделала крупномасштабную кампанию дистилляции доступной по стоимости. Получившаяся компактная модель-ученик, дистиллированная из Llama 3.1 8B Instruct до примерно 3,2 млрд параметров, сохраняет большую часть точности учителя на BoolQ и HellaSwag, а на MMLU отстаёт от него примерно на девять пунктов, имея менее половины числа параметров.

Модель-ученик сохраняет большую часть точности учителя при коротком контексте и имеет менее половины его размера. Источник: рисунок 6 статьи.

Эта работа является частью продолжающихся исследований Multiverse Computing, направленных на то, чтобы сделать дистилляцию и восстановление моделей практически применимыми в масштабах, а не просто одноразовым рецептом, но процессом, который команды могут недорого многократно улучшать. В статье также рассматриваются дополнительные абляционные эксперименты, например влияние выбора функции потерь и упаковки последовательностей на качество восстановления.

Хотите ознакомиться с полными техническими подробностями, включая градиент в замкнутой форме для объединённого фрагментированного критерия и полную конфигурацию обучения? Прочитайте полную статью или свяжитесь с нашей командой, чтобы обсудить применение этого подхода к вашим собственным конвейерам дистилляции.

Мы также опубликовали исходный код реализации фрагментированного критерия: github.com/CompactifAI/Full-Chunked-KL-Loss

Модели, упомянутые в этой статье 6

Статьи, упомянутые в этой статье 1

Коллекции, упомянутые в этой статье 1