Проблема: стоимость дистилляции превышает возможности GPU
Дистилляция знаний (Knowledge Distillation) — процесс обучения компактной модели-студента на основе выходов большой модели-учителя — стала стандартом для сжатия LLM, таких как Kimi-K3 (2.8 трлн параметров) или Qwen. Однако традиционный онлайн-режим требует одновременного удержания в памяти обеих моделей и вычисления полного распределения вероятностей по всему словарю (до 200k+ токенов) для каждого шага.
Для модели gpt-oss-120b при длине контекста 32K и батче 4 только тензор вероятностей учителя занимает ~50 ГБ в формате bfloat16. С учетом градиентов и оптимизаторов пиковое потребление VRAM достигает 250 ГБ, что невозможно для одной GPU H200 (лимит 141 ГБ) и требует сложных стратегий тензорного параллелизма на сотнях устройств.
Решение: Офлайн-кэширование и Chunked KL Loss
Команда Multiverse Computing (авторы статьи Efficient Knowledge Distillation for LLMs) предлагает два системных изменения, позволяющих провести дистилляцию на одном устройстве:
- Offline Top-K Caching: Выходы учителя вычисляются один раз, кэшируются только топ-100 наиболее вероятных токенов для каждой позиции. Учитель не участвует в обучении, что исключает его весовые данные из памяти.
- Fused Chunked KL Loss: Новый алгоритм вычисления расхождения Кульбака-Лейблера. Вместо создания полной матрицы «словарь × контекст», он обрабатывает последовательность чанками (кусками). Прямой проход проецирует скрытые состояния в логиты, вычисляет часть потери, и данные сразу выбрасываются. Обратный проход пересчитывает чанк на лету. Это устраняет пиковые скачки памяти.
Результаты: экономия VRAM и сохранение качества
Эксперименты проводились на GPU NVIDIA H200 с моделью Llama 3.1 8B Instruct в качестве учителя и 3.2B Llama в качестве студента. Метрика качества (training loss) осталась идентичной для всех методов, но затраты ресурсов радикально изменились.
| Метод (контекст 8K) | Пиковая VRAM | Время итерации | Пропускная способность |
|---|---|---|---|
| Online distillation | 102.8 GB | 25.9 s | 237 TFLOP/s |
| Offline, dense KL | 78.3 GB | 18.5 s | 331 TFLOP/s |
| Offline, forward-chunked KL | 61.8 GB | 18.4 s | 335 TFLOP/s |
| Offline, fused chunked KL | 58.3 GB | 20.2 s | 304 TFLOP/s |
Главное преимущество Fused Chunked KL раскрывается на длинных контекстах. При длине 32K токенов пиковая память падает с 85.2 ГБ (dense) до 5.45 ГБ (chunked) — снижение в 15.6 раза. Dense-метод вообще падает с ошибкой OOM при 64K токенах, тогда как chunked-версия стабильно работает на 256K токенах, потребляя всего 11.6 ГБ против 134.2 ГБ у аналогов.
Значение для индустрии
Открытый код реализации доступен на GitHub (CompactifAI/Full-Chunked-KL-Loss). Этот подход делает экономически целесообразной дистилляцию огромных моделей (вроде Kimi-K3 или Nemotron 3 Puzzle 75B) на доступном оборудовании, устраняя необходимость в кластерах из сотен GPU для этапа обучения студентов.
Источник: Hugging Face blog ↗