Эффективная дистилляция знаний: как уменьшить стоимость обучения ИИ моделей до масштаба одного GPU
Распространение открытых больших языковых моделей (LLM) сделало технику дистилляции знаний — обучение компактной модели по образцу гигантского учителя — стандартной практикой. Однако традиционные методы требуют колоссальных объемов оперативной памяти видеокарт, делая их применение непрактичным для больших контекстов и масштабируемых экспериментов. Команда Multiverse Computing представила систему, которая кардинально снижает эти затраты, используя оффлайн-кэширование топ-K логов и скомпилированную поточно-блочную функцию потерь КЛ-дивергенции. Эти нововведения позволяют проводить качественную дистилляцию на одном устройстве с поддержкой контекста до 32K токенов, устраняя необходимость в сотнях графических процессоров.
# Эффективная дистилляция знаний: как уменьшить стоимость обучения ИИ моделей до масштаба одного GPU
В эпоху экспоненциального роста параметров нейросетей задача сокращения их размера без потери интеллекта стала критически важной. Метод дистилляции знаний (knowledge distillation) традиционно позволяет обучить небольшую «ученическую» модель так, чтобы она воспроизводила поведение огромной «учительской». Однако, если учителем служит модель с сотнями миллиардов параметров, процесс обучения становится экстремально дорогим и энергозатратным.
Команда исследователей из Multiverse Computing представила решение, которое снижает стоимость этой операции настолько, что теперь её можно выполнять на стандартном оборудовании с поддержкой длинных контекстов.
Почему традиционная дистилляция требует огромных ресурсов?
Стандартный подход к дистилляции в реальном времени (online distillation) требует одновременного размещения в памяти обеих моделей: учителя и ученика. На каждом шаге обучения учитель генерирует полный распределение вероятностей для всего словаря (вocabulary) каждого токена в последовательности.
Для современных моделей это создаёт непреодолимое препятствие для памяти видеокарт. Например, у модели gpt-oss-120b словарь содержит 201 088 токенов. При длине последовательности в 32 000 токенов и микро-пакете (batch size) в 4 элемента, тензор вероятностей учителя уже занимает около 50 ГБ оперативной памяти. При добавлении градиентов, активаций и состояний оптимизатора, пиковая нагрузка достигает 250 ГБ. Для сравнения: флагманская видеокарта NVIDIA H200 имеет 141 ГБ памяти. Это означает, что для одного шага дистилляции такого масштаба требуется объединенное вычислительное ядро, эквивалентное десяткам или сотням GPU, настроенных через сложные стратегии параллельности тензоров.
Две системные инновации для экономии памяти
Чтобы преодолеть эти ограничения, исследователи в своей работе «Efficient Knowledge Distillation for LLMs» внедрили две ключевые модификации:
1. Оффлайн-кэширование топ-K логов
Вместо того чтобы заново вычислять пробы в реальном времени, система вычисляет их заранее и сохраняет в кэш. Вместо хранения полного распределения (которое может быть размером с весь диск), система хранит только топ-K наиболее вероятных токенов с их логом вероятности.
Это позволяет избежать необходимости хранить огромные тензоры с малыми вероятностями, которые редко используются в процессе обучения, но требуют много памяти для хранения.
2. Скомпилированная поточно-блочная функция потерь (Fused Chunked KL Loss)
Второй метод заключается в изменении математики функции потерь. Традиционно она требует вычисления градиентов для всей последовательности сразу, что создаёт огромную нагрузку на память GPU. Новое решение разделяет последовательность на небольшие непересекающиеся блоки (chunks). Для каждого блока вычисляется локальная функция потерь, затем результаты агрегируются.
Это позволяет системе не хранить состояние всей длинной последовательности в памяти одновременно, значительно снижая пиковый объём используемой памяти.
Результаты: от сотен GPU до одного
Благодаря этим модификациям, команда достигла прорыва в масштабируемости дистилляции. В эксперименте с моделью gpt-oss-120b при контексте в 32K токенов:
* Объем памяти, необходимый для дистилляции, сократился с 85,2 ГБ до 5,45 ГБ. Это более чем 15-кратное снижение. * Вместо использования четырёх узлов с GPU для выполнения дистилляции, теперь достаточно всего одного современного GPU. * Скорость выполнения одного шага обучения увеличилась с 57 секунд до 12,23 секунд (в 5 раз быстрее). * Пропускная способность устройства выросла с 74,2 до 345,7 TFLOP/s.
Практические результаты для конечных моделей
Эффективность метода была проверена на реальной дистилляции модели Llama 3.1 8B в компактный вариант объёмом около 3,2 миллиардов параметров.
Результаты показали, что «малыш» сохранил высокую точность на стандартных бенчмарках: * На задачах BoolQ и HellaSwag точность почти полностью сохранилась на уровне учителя. * На MMLU (тест на общие знания) отставание составило менее 9 баллов от оригинала. * При этом модель сохранила высокое качество работы в кратких контекстах.
Таким образом, сложная архитектура в 2,5 раза меньше оригинала сохранила её интеллектуальный потенциал, что делает её идеальной для мобильных устройств или облачных сервисов с низкими задержками.
Почему это важно для индустрии?
Доступ к таким методам позволяет компаниям проводить исследования и разработки по созданию новых моделей значительно дешевле и быстрее. Вместо закупки сотен дорогих видеокарт, можно масштабировать эксперименты, используя стандартные конфигурации, не жертвуя качеством результатов. Это ускоряет цикл разработки и снижает порог входа для команд, занимающихся созданием и оптимизацией языковых моделей.
Исследователи также открыли исходный код реализации метода под открытым доступом на GitHub, чтобы ускорить внедрение этих практик в индустрию.
Источники: 1. [Paper: Efficient Knowledge Distillation for LLMs](https://arxiv.org/abs/2608.03796) 2. [Репо с кодом на GitHub](https://github.com/CompactifAI/Full-Chunked-KL-Loss)