В эпоху, когда размер языковых моделей измеряется триллионами параметров, а стоимость их развертывания может достигать миллионов долларов, индустрия искусственного интеллекта столкнулась с фундаментальной проблемой: как сделать мощные модели доступными для широкого круга разработчиков и компаний? Ответ лежит в области сжатия моделей, и одним из самых перспективных методов здесь является дистилляция знаний. Этот процесс позволяет обучить компактную «студенческую» модель так, чтобы она имитировала поведение гораздо более крупной и сложной «учительской» модели. Однако традиционные подходы к дистилляции требуют огромных вычислительных ресурсов, часто превышающих возможности даже самых современных видеокарт.
Команда Multiverse Computing представила революционный подход к решению этой проблемы. В своей новой работе они демонстрируют, как можно радикально снизить требования к памяти (VRAM) и вычислительной мощности, необходимые для дистилляции больших языковых моделей (LLM). Их метод позволяет проводить обучение на одном GPU, что делает процесс доступным для небольших команд и исследовательских групп. Мы подробно разберем, как работают эти технические инновации, почему старые методы были слишком дорогими и как новые алгоритмы меняют правила игры в индустрии.
01Почему дистилляция LLM стала такой дорогой?
Чтобы понять масштаб проблемы, нужно заглянуть под капот процесса дистилляции. Классический подход, известный как онлайн-дистилляция, предполагает одновременную загрузку в память обеих моделей: учителя и студента. На каждом шаге обучения учительская модель выполняет полный прямой проход (forward pass), чтобы сгенерировать распределение вероятностей для всего словаря. Затем студент обучается подстраиваться под это распределение, используя расхождение Кульбака-Лейблера (KL-divergence loss).
Этот метод обладает высокой выразительной силой, так как студент видит полную картину предсказаний учителя. Но цена этой точности — колоссальное потребление видеопамяти. Давайте рассмотрим конкретный пример. Возьмем модель gpt-oss-120b, которая имеет словарь из 201 088 токенов. Если мы установим длину последовательности (sequence length) на уровне 32K токенов и размер батча (batch size) равным 4, то только тензор вероятностей учителя будет иметь размерность 4 × 201,088 × 32,768. В формате bfloat16, который часто используется для оптимизации памяти, этот один тензор занимает около 50 ГБ видеопамяти.
Но это только начало. Помимо самого тензора, необходимо хранить градиенты, активации, веса модели и состояния оптимизатора. В совокупности одна итерация обучения может требовать до 250 ГБ VRAM. Для сравнения: даже флагманские видеокарты NVIDIA H200 или B200 имеют объем памяти около 141 ГБ. Таким образом, стандартный подход к дистилляции таких моделей физически невозможен на одном устройстве и требует кластеров из сотен GPU с сложной стратегией тензорного параллелизма.
02Два системных решения для экономии памяти
Исследователи из Multiverse Computing предложили два ключевых изменения в архитектуре процесса обучения, которые позволяют обойти эти ограничения. Первое изменение касается того, как мы взаимодействуем с учителем, а второе — того, как мы вычисляем функцию потерь.
1. Оффлайн-дистилляция с кэшированием
Вместо того чтобы пересчитывать выход учителя на каждом шаге обучения, команда предлагает вычислить его один раз. Учительская модель запускается через датасет, и для каждой позиции в последовательности сохраняются только топ-K наиболее вероятных токенов (в их экспериментах K=100). Этот кэш создается заранее и сохраняется на диске или в оперативной памяти.
Во время обучения студента учительская модель вообще не загружается в GPU. Студент обучается, подстраиваясь под этот заранее подготовленный кэш. Это решает сразу две проблемы: во-первых, учителю не нужно занимать место в VRAM во время обучения студента, а во-вторых, кэш можно переиспользовать для множества экспериментов (абляций) без повторных дорогостоящих вычислений учителя. Это превращает дистилляцию из процесса, требующего синхронной работы двух гигантских моделей, в процесс, где активна только одна модель — студент.
2. Слиянная кусочковая функция потерь KL (Fused Chunked KL Loss)
Даже с кэшированием учителя вычисление функции потерь KL остается узким местом. Функция KL требует сравнения распределения студента с распределением учителя для каждого токена. Если представить это как сетку, где строки — это токены словаря, а столбцы — позиции в последовательности, то для словаря в 100K+ слов и длинной последовательности эта сетка будет огромной. Стандартные библиотеки (например, PyTorch или NVIDIA Megatron-Bridge) обычно создают эту полную сетку в памяти перед тем, как вычислить одно числовое значение потери.
Авторы предлагают три способа вычисления этой потери, сравнивая их эффективность:
- Dense KL (Плотная KL): Классический подход. Он восстанавливает полную плотную сетку вероятностей учителя из кэшированных топ-K логитов и сравнивает ее с плотной сеткой лог-вероятностей студента. Это самый точный метод, но он требует хранения полной сетки «словарь × последовательность» в памяти дважды (для учителя и студента). Как мы видели выше, это приводит к пиковому потреблению памяти в ~250 ГБ.
- Forward-chunked KL (Кусочковая KL по прямому проходу): Этот метод сохраняет распределение учителя разреженным (топ-K логиты), никогда не расширяя их в плотную сетку. Потеря вычисляется по частям, по одному срезу позиций последовательности за раз. Это убирает плотную сетку учителя, но сетка студента (логиты от выходного слоя модели) все еще вычисляется полностью и хранится для обратного прохода (backward pass). Память все еще растет линейно с длиной последовательности, но пиковое значение ниже.
- Fused chunked KL (Слиянная кусочковая KL): Главное нововведение. Этот метод сливает проекцию выходного слоя модели студента непосредственно в вычисление функции потерь. Он никогда не создает полную сетку логитов студента. Вместо этого он обрабатывает последовательность по одному чанку (куску) за раз: проецирует скрытые состояния в логиты для этого чанка, добавляет результат в накопленную потерю и немедленно выбрасывает чанк из памяти. Обратный проход пересчитывает каждый чанк на лету. Хотя это требует двойного вычисления проекции (один раз вперед, один раз назад), пиковая память растет только линейно с длиной последовательности, а не взрывается от размера словаря.
03Практические результаты: сравнение методов
Чтобы оценить эффективность предложенных методов, команда провела серию бенчмарков на одной видеокарте NVIDIA H200. В качестве учителя использовалась модель Llama 3.1 8B Instruct, а в качестве студента — модель Llama размером 3.2B. Длина контекста была установлена на 8K токенов. Результаты показали, что все четыре метода (онлайн-дистилляция и три варианта оффлайн-дистилляции) достигают практически идентичных кривых потерь (loss curves). Это подтверждает, что использование кэшированных топ-100 логитов не приводит к потере качества модели.
Однако различия в потреблении ресурсов были значительными:
- Онлайн-дистилляция: Пиковая память 102.8 ГБ. Время итерации: 25.9 секунд. Пропускная способность: 237 TFLOP/s.
- Оффлайн, Dense KL: Пиковая память 78.3 ГБ. Время итерации: 18.5 секунд. Пропускная способность: 331 TFLOP/s.
- Оффлайн, Forward-chunked KL: Пиковая память 61.8 ГБ. Время итерации: 18.4 секунд. Пропускная способность: 335 TFLOP/s.
- Оффлайн, Fused chunked KL: Пиковая память 58.3 ГБ. Время итерации: 20.2 секунды. Пропускная способность: 304 TFLOP/s.
Как видно из данных, оффлайн-подходы значительно экономят память по сравнению с онлайн-методом. Метод Fused chunked KL демонстрирует самое низкое потребление памяти (58.3 ГБ), что оставляет большой запас для других процессов. Хотя он немного медленнее на коротких контекстах из-за накладных расходов на пересчет в обратном проходе, его преимущества раскрываются полностью при увеличении длины контекста.
04Масштабирование на длинные контексты
Настоящая сила предложенного метода проявляется при работе с длинными последовательностями. В изолированном бенчмарке на тестовой сети (без тела трансформера, только ядро функции потерь) при длине контекста 32K токенов пиковая память упала с 85.2 ГБ (для плотной потери) до 5.45 ГБ (для полностью кусочковой версии). Это снижение в 15.6 раза! Более того, плотная потеря полностью отказывала при длине 64K токенов из-за нехватки памяти.
При длине контекста 256K токенов полностью кусочковая потеря использовала всего 11.6 ГБ против 134.2 ГБ у следующего лучшего варианта. Кроме того, она оказалась примерно в 3.3 раза быстрее на итерацию при такой длине. Это открывает двери для дистилляции моделей, работающих с очень длинными документами, книгами или кодовыми базами, что ранее было практически невозможно без кластера GPU.
Команда также провела дистилляцию модели GPT-OSS 20B с контекстом 32,768 токенов. Благодаря экономии памяти от слиянной кусочковой потери, конфигурация сократилась с четырех узлов GPU до одного. Время шага обучения упало с 57.0 до 12.23 секунд (ускорение в 5 раз), а пропускная способность на GPU выросла с 74.2 до 345.7 TFLOP/s. Это не просто экономия денег — это возможность проводить итеративное обучение там, где раньше это было технически невозможно.
05Качество полученной модели
Главный вопрос: не потеряли ли мы в качестве из-за таких агрессивных оптимизаций? Результаты тестирования студента, дистиллированного из Llama 3.1 8B Instruct (до ~3.2B параметров), показали отличные результаты. Модель сохранила большую часть точности учителя на задачах BoolQ и HellaSwag. На бенчмарке MMLU (Massive Multitask Language Understanding) отставание составило всего около 9 пунктов, что является отличным результатом для модели, имеющей менее чем половину параметров учителя.
Это доказывает, что оффлайн-дистилляция с кэшированием топ-K логитов является «без потерь» по сравнению с онлайн-дистилляцией в плане качества конечной модели. Студент учится эффективно использовать информацию, заложенную в распределении учителя, даже если видит только топ-100 токенов, а не весь словарь.
06Что это значит на практике
Для разработчиков и исследователей в России и мире эти новости означают несколько важных вещей. Во-первых, барьер для входа в область тонкой настройки и дистилляции больших моделей снижается. Теперь не нужно искать доступ к кластерам из десятков A100/H100, чтобы попробовать дистиллировать модель под свои задачи. Один мощный GPU (вроде RTX 4090, A6000 или даже H200, если он доступен) становится достаточной платформой для экспериментов.
Во-вторых, это ускоряет цикл разработки. Возможность запускать дистилляцию на одном узле позволяет быстрее перебирать гиперпараметры, тестировать новые датасеты и проводить абляционные исследования. То, что раньше занимало дни настройки распределенного обучения, теперь можно сделать за часы локально.
В-третьих, это открывает путь к созданию специализированных моделей для нишевых задач. Компании могут дистиллировать общие большие модели в компактные версии, оптимизированные для конкретных доменов (медицина, право, код), используя собственные закрытые данные, не раскрывая их публично и не тратя миллионы на облачные вычисления.
Код реализации кусочковой функции потерь уже открыт в репозитории github.com/CompactifAI/Full-Chunked-KL-Loss. Исследователи также опубликовали полную статью с математическими выкладками, включая замкнутую форму градиента для слиянной кусочковой потери. Это дает возможность любому желающему интегрировать эти методы в свои пайплайны обучения. В условиях, когда доступ к передовым вычислительным ресурсам может быть ограничен, такие алгоритмические оптимизации становятся критически важным инструментом для поддержания конкурентоспособности в области ИИ.
Дистилляция знаний перестает быть элитарным процессом, доступным только гигантам индустрии. Благодаря умным алгоритмическим решениям, таким как кэширование логитов и кусочковое вычисление потерь, мы движемся к эпохе, где мощные модели можно создавать и адаптировать на гораздо более скромном оборудовании. Это демократизирует доступ к технологиям и ускоряет инновации во всем мире.
Источник: Hugging Face ↗