Дистилляция знаний — техника, при которой меньшая модель (студент) обучается воспроизводить поведение более крупной модели (учителя), — стала стандартным этапом сжатия больших языковых моделей. С ростом числа открытых моделей, таких как gpt-oss, Qwen, GLM или Kimi, дистилляция вновь оказалась в центре внимания: например, модель Kimi-K3 содержит 2,8 триллиона параметров и требует около 3 ТБ VRAM только для загрузки. Сжатие таких моделей до меньших размеров с последующим восстановлением качества через дистилляцию — обычная практика, которую применяют Nvidia (Nemotron 3 Puzzle 75B) и Multiverse Computing (Hypernova 60B). Однако сам этап дистилляции часто оказывается самым дорогим: он требует одновременного удержания в памяти и учителя, и студента, а также вычисления распределения вероятностей по всему словарю для каждого токена. Это приводит к необходимости сотен GPU и сложных стратегий тензорного параллелизма.
Исследователи Hugging Face предложили два системных изменения, которые радикально снижают стоимость дистилляции. Первое — офлайн-дистилляция: вместо пересчёта учителя на каждом шаге, его выходные данные вычисляются один раз, и кэшируются top-100 наиболее вероятных токенов для каждой позиции. Учитель больше не нужен в памяти во время обучения, а кэш можно переиспользовать для множества экспериментов. Второе — новый fused chunked KL loss, который обрабатывает данные порциями и избегает материализации полной матрицы размером «словарь × длина последовательности». Это снижает пиковое потребление VRAM с примерно 250 ГБ до 128 ГБ, что позволяет проводить обучение на одной видеокарте H200 (141 ГБ VRAM) вместо кластера из сотен GPU.
| Метод | Пиковое потребление VRAM | Скорость |
|---|---|---|
| Dense KL | ~250 ГБ | Медленный |
| Forward-chunked KL | ~128 ГБ | Быстрый |
| Fused chunked KL | ~128 ГБ | Средний |
Для понимания масштаба проблемы: стандартный подход — онлайн-дистилляция с использованием KL-дивергенции — требует одновременного удержания учителя и студента, а также полного тензора вероятностей учителя. Для модели gpt-oss-120b с словарём в 201 088 токенов, при длине последовательности 32K и batch size 4, тензор вероятностей учителя имеет форму 4 × 201 088 × 32 768, что в bfloat16 занимает около 50 ГБ только для одного тензора. С учётом градиентов, активаций и весов модели, пиковое потребление достигает 250 ГБ, что превышает возможности даже H200 или B200. Новый метод решает эту проблему, разбивая вычисления на части и используя разреженное представление учителя.
Пиковое потребление VRAM снижается с 250 ГБ до 128 ГБ, что умещается в одну H200.
Авторы сравнили три способа вычисления KL-лосса: плотный (dense), прямой с разбиением на чанки (forward-chunked) и fused chunked. Плотный метод — это базовый подход, который строит полную плотную сетку вероятностей учителя и сравнивает её с плотной сеткой студента, что требует двойного объёма памяти. Forward-chunked метод сохраняет учителя разреженным (только top-100 логитов) и вычисляет лосс по частям, что оказалось самым быстрым, но всё ещё требует плотного представления логитов студента. Fused chunked метод, предложенный в статье, объединяет вычисления и полностью избегает материализации больших матриц, достигая пикового потребления 128 ГБ.
Эти изменения делают дистилляцию доступной для широкого круга исследователей и компаний, которые не имеют доступа к большим кластерам GPU. Возможность проводить эксперименты на одной видеокарте значительно ускоряет итерации и позволяет исследовать длинные контексты, что ранее было невозможно из-за ограничений памяти. Статья опубликована на Hugging Face Blog, и авторы предоставили детали, достаточные для воспроизведения результатов. Это важный шаг к демократизации технологий сжатия моделей, который может повлиять на развитие open-source экосистемы ИИ.
