Multiverse Computing сократила видеопамять для дистилляции нейросетей до 15,6 раза

Компания Multiverse Computing опубликовала статью и пост в блоге на Hugging Face «Efficient Knowledge Distillation for LLMs: Offline Top-K Logits and a Fused Chunked KL Loss» с двумя изменениями в процессе дистилляции, переноса знаний от большой модели-учителя в компактную модель-ученика. Изменения резко снижают расход видеопамяти при обучении, и код метода выложен в открытый доступ.

Проблема, которую решает статья: в стандартной («онлайн») дистилляции учитель и ученик одновременно держатся в памяти, и на каждом шаге обучения учитель заново считает распределение вероятностей по всему словарю для каждого токена. Пример из поста: у модели gpt-oss-120b словарь на 201 088 токенов; при длине последовательности 32K и размере батча 4 один только тензор вероятностей учителя имеет форму 4×201 088×32 768 и весит около 50 ГБ в формате bfloat16. С учётом градиентов, активаций, весов модели и состояний оптимизатора один шаг обучения дистилляции может доходить до пика примерно в 250 ГБ видеопамяти, это больше, чем даёт даже одна карта H200 (141 ГБ) или B200.

Авторы предлагают два системных изменения. Первое, офлайн-дистилляция: вместо пересчёта учителя на каждом шаге его выход считается один раз, для каждой позиции кэшируются топ-100 самых вероятных токенов, и ученик обучается уже по этому кэшу; учителю больше не нужно постоянно находиться в памяти, а кэш можно переиспользовать в разных экспериментах. Второе, объединённая поблочная функция потерь на основе KL-дивергенции (fused chunked KL loss): она обрабатывает последовательность по частям, «вплавляя» финальную проекцию модели прямо в вычисление функции потерь, так что полная матрица логитов ученика по всему словарю никогда не строится целиком; при обратном проходе каждый блок пересчитывается заново, а не хранится в памяти. В итоге пиковая память растёт линейно с длиной последовательности, а не скачком, как при полной матрице словарь×длина. На практике это снижает пик с ~250 ГБ у обычной («плотной») KL до ~128 ГБ у объединённой поблочной версии.

Корректность подхода проверена отдельно: на контрольном сравнении (одна карта H200, учитель Llama 3.1 8B Instruct, ученик, модель Llama на 3.2 млрд параметров, контекст 8K токенов) все четыре способа, онлайн-дистилляция и три офлайн-варианта функции потерь, дали почти идентичную кривую обучения; авторы делают вывод, что офлайн-дистилляция с кэшированными топ-100 логитами не уступает по качеству онлайн-варианту. На изолированном тесте самого механизма потерь (без модели, только ядро вычислений) масштабирование на длинный контекст выглядит так: при 32K токенах пиковая память падает с 85,2 ГиБ у плотной версии до 5,45 ГиБ у полностью поблочной, сокращение до 15,6 раза; плотная версия перестаёт работать вовсе начиная с 64K токенов; при 256K токенах полностью поблочная версия использует 11,6 ГиБ против 134,2 ГиБ у следующей по эффективности поблочной альтернативы и работает примерно в 3,3 раза быстрее на шаг.

На реальной задаче, дистилляции модели GPT-OSS 20B с длиной контекста 32 768 токенов, высвобожденная память позволила сократить установку с четырёх узлов GPU до одного. Время шага обучения упало с 57,0 до 12,23 секунды (примерно в 5 раз), а пропускная способность на один GPU выросла с 74,2 до 345,7 TFLOP/с. Итоговая компактная модель, полученная дистилляцией из учителя Llama 3.1 8B Instruct в ученика на ~3,2 млрд параметров, сохраняет большую часть точности учителя на тестах BoolQ и HellaSwag и отстаёт от него примерно на девять пунктов на MMLU, при размере меньше половины от учителя.

В посте объясняется, зачем это вообще нужно: развёртывание крупнейших открытых моделей дорого, например, модели Kimi-K3 с 2,8 трлн параметров нужно около 3 ТБ видеопамяти просто для загрузки. Поэтому сжатие таких моделей с восстановлением качества через дистилляцию стало обычной практикой: компании называют в пример модель Nemotron 3 Puzzle 75B от Nvidia и собственную Hypernova 60B от Multiverse Computing (без подробностей качества). Эта работа, часть более широкого исследования компании о том, как сделать дистилляцию и «лечение» моделей после сжатия практичными в масштабе, а не разовым трюком. Реализация поблочной функции потерь выложена в открытый доступ: github.com/CompactifAI/Full-Chunked-KL-Loss.

Ключевые факты

Multiverse Computing выпустила статью и открыла код метода снижения расхода видеопамяти при дистилляции больших языковых моделей.
Два изменения: офлайн-кэширование топ-100 логитов учителя вместо постоянного хранения учителя в памяти + объединённая поблочная функция потерь на KL-дивергенции, которая не строит полную матрицу словарь×длина последовательности.
На изолированном тесте при 32K токенах пик памяти упал с 85,2 ГиБ (плотная версия) до 5,45 ГиБ (поблочная), в 15,6 раза; плотная версия отказывает вовсе начиная с 64K токенов.
На реальной дистилляции GPT-OSS 20B с контекстом 32 768 токенов установка сократилась с четырёх узлов GPU до одного, шаг обучения ускорился примерно в 5 раз (с 57,0 до 12,23 секунды), пропускная способность на GPU выросла с 74,2 до 345,7 TFLOP/с.
Дистиллированная модель на ~3,2 млрд параметров (из учителя Llama 3.1 8B Instruct) сохраняет большую часть точности учителя на BoolQ и HellaSwag и отстаёт лишь на девять пунктов на MMLU при размере меньше половины от учителя.

Почему это важно

Дистилляция, стандартный способ сжать гигантскую модель в компактную, сохранив качество, но именно этот шаг чаще всего упирается в видеопамять: нужно одновременно держать в памяти и учителя, и ученика, и распределение вероятностей по всему словарю для каждого токена. Метод Multiverse Computing убирает оба узких места разом, учитель больше не обязан находиться в памяти во время обучения, а функция потерь никогда не строит полную матрицу словарь×длина последовательности. В изолированном тесте это дало сокращение пиковой памяти до 15,6 раза при 32K токенах, а на реальной задаче дистилляции GPT-OSS 20B перевело обучение с четырёх узлов GPU на один.

Кому это важно

Инженерам и исследователям, которые сжимают крупные открытые модели (gpt-oss, Qwen, GLM, Kimi и похожие) для развёртывания на ограниченном железе, а также командам, которым нужна дистилляция с длинным контекстом, раньше она требовала сотен GPU с тщательным тензорным параллелизмом, теперь, по данным поста, помещается в память одной карты.

Как это применить

Реализация поблочной функции потерь на KL-дивергенции выложена в открытый доступ на GitHub (CompactifAI/Full-Chunked-KL-Loss), так что метод можно подключить к собственному пайплайну дистилляции напрямую. Компания также предлагает связаться с её командой по вопросам применения метода к своим пайплайнам; конкретных условий, цен или лицензии в источнике не приведено.

Можно ли доверять

Это собственный пост компании-разработчика метода в её блоге на Hugging Face, а не независимая рецензируемая публикация, и Multiverse Computing одновременно продаёт свои сжатые модели (Hypernova 60B), то есть заинтересована в продвижении инструмента. При этом ключевые цифры сопровождаются графиками из статьи, на которую даётся ссылка, а корректность офлайн-подхода проверена напрямую: кривые обучения всех четырёх сравниваемых методов на контрольном тесте совпадают почти точно. Код выложен в открытый доступ, что позволяет проверить результаты независимо.

Риски и подводные камни

Все ключевые цифры, самоотчёт компании, независимой проверки нет. Источник не называет ни авторов статьи, ни дату публикации, ни тип GPU, использованный в примере «с четырёх узлов до одного» (в отличие от H200, указанного для остальных тестов), сравнить условия целиком не получится. Стоимость обучения в деньгах не приводится, только память, время шага и пропускная способность. Числа для конкурентных сжатых моделей, Nemotron 3 Puzzle 75B от Nvidia и собственная Hypernova 60B, в посте не приводятся, они упомянуты лишь как примеры, без сравнения качества.