В последние годы дистилляция знаний, процесс обучения меньшей модели-студента для достижения производительности более крупной модели-учителя, вновь стала актуальной темой в области машинного обучения. С появлением открытых больших языковых моделей, таких как gpt-oss, Qwen, GLM и Kimi, исследование методов дистилляции приобрело новую волну интереса. Развертывание таких крупных моделей требует значительных ресурсов: например, модель Kimi-K3 имеет 2.8 триллиона параметров и требует около 3 ТБ видеопамяти только для загрузки.

Проблемы с дистилляцией

Сжатие больших моделей в меньшие и восстановление оригинальных возможностей через дистилляцию знаний стало стандартной практикой. Компании, такие как Nvidia с моделью Nemotron 3 Puzzle 75B и Multiverse Computing с Hypernova 60B, недавно выпустили качественные сжатые модели. Однако этап дистилляции часто оказывается самым дорогим в процессе, поскольку требует одновременной загрузки как модели-учителя, так и модели-студента, что приводит к огромным затратам видеопамяти. В стандартной настройке, использующей потерю Кульбака-Лейблера (KL loss), модель-учитель выполняет полный проход вперед на каждом шаге обучения, что требует значительных вычислительных ресурсов. Например, модель gpt-oss-120b имеет словарь из 201088 токенов, и при длине последовательности 32K и размере батча 4, тензор вероятностей учителя занимает около 50 ГБ видеопамяти. В результате, полная итерация дистилляции может достигать 250 ГБ видеопамяти, что превышает возможности даже самых мощных графических процессоров.

Новые подходы к дистилляции

В нашей последней работе, "Efficient Knowledge Distillation for LLMs: Offline Top-K Logits and a Fused Chunked KL Loss", мы предлагаем два изменения в системе: 1. Кэширование топ-K логитов учителя: вместо того чтобы пересчитывать выходные данные учителя на каждом шаге, мы вычисляем их один раз и кэшируем 100 наиболее вероятных токенов для каждой позиции. Это позволяет избежать необходимости держать учителя в памяти во время обучения. 2. Эффективная потеря KL: новая, экономичная по памяти версия потери KL, которая не требует создания полной матрицы размером словарь × длина последовательности, значительно снижает использование видеопамяти. Эти два изменения позволяют сократить затраты на обучение до такой степени, что дистилляция становится возможной на одном графическом процессоре, что делает масштабные эксперименты более доступными.

Сравнение методов

Мы провели сравнение четырех подходов к дистилляции: онлайн-дистилляция и три варианта оффлайн-потерь. Все они были протестированы на одном графическом процессоре H200 с моделью Llama 3.1 8B Instruct в качестве учителя и моделью Llama 3.2B в качестве студента при контексте в 8K токенов. Результаты показали, что все четыре метода достигли почти идентичных значений потерь, несмотря на то, что оффлайн-методы обучались только на кэшированных логитах. | Метод | Пиковая память | Время итерации | Пропускная способность | |--------------------------|----------------|----------------|-----------------------| | Онлайн-дистилляция | 102.8 ГБ | 25.9 с | 237 TFLOP/s | | Оффлайн, плотная KL | 78.3 ГБ | 18.5 с | 331 TFLOP/s | | Оффлайн, прямоугольная KL| 61.8 ГБ | 18.4 с | 335 TFLOP/s | | Оффлайн, объединенная KL | 58.3 ГБ | 20.2 с | 304 TFLOP/s |

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

Для более наглядного понимания масштабирования мы провели отдельные тесты на упрощенной сети проекции выходных данных. При длине последовательности 32K пиковая память упала с 85.2 ГБ при плотной потере до 5.45 ГБ с полностью чанкованной версией, что составляет сокращение в 15.6 раз. При длине 256K токенов полностью чанкованная потеря использовала 11.6 ГБ против 134.2 ГБ у следующего лучшего варианта. В результате дистилляции модели GPT-OSS 20B с контекстом в 32,768 токенов, освобожденная память позволила сократить настройки с четырех графических процессоров до одного. Время на шаг уменьшилось с 57.0 до 12.23 секунд, а пропускная способность на графический процессор возросла с 74.2 до 345.7 TFLOP/s.

Заключение

Эффективная оффлайн-настройка сделала возможной крупномасштабную кампанию по дистилляции, сохраняя при этом большую часть точности учителя при меньшем количестве параметров. Результирующая компактная модель-студент, дистиллированная из Llama 3.1 8B Instruct до примерно 3.2B параметров, сохраняет большинство характеристик учителя на задачах BoolQ и HellaSwag, оставаясь в пределах девяти пунктов на MMLU. Эта работа является частью продолжающегося исследования компании Multiverse Computing, направленного на упрощение процесса дистилляции и его применение в масштабах, позволяя командам эффективно экспериментировать.