Multiverse Computing предложила два системных приёма — офлайн-кэширование top-K логитов учителя и фьюзed KL-лосс с чанкованием, — которые сокращают потребление VRAM при дистилляции LLM до уровня одной GPU.
Содержание
- Почему дистилляция стала самой дорогой частью конвейера
- Два системных изменения вместо сотен GPU
- Три способа посчитать одну и ту же функцию потерь
- Лицом к лицу: четыре конфигурации на одной H200
- Где метод раскрывается по-настоящему: длинный контекст
- Что получилось у студента весом 3,2 млрд параметров
- Что это меняет для индустрии и российских команд
Испанская компания Multiverse Computing опубликовала исследование, которое может заметно снизить порог входа в сжатие больших языковых моделей. В статье «Efficient Knowledge Distillation for LLMs: Offline Top-K Logits and a Fused Chunked KL Loss» команда показала, как два системных изменения — однократное кэширование top-K логитов модели-учителя и новая экономичная реализация KL-дивергенции — позволяют проводить дистилляцию с длинным контекстом на одной видеокарте вместо четырёх GPU-узлов, ускоряя итерацию обучения примерно в пять раз.
- Почему дистилляция стала самой дорогой частью конвейера
- Два системных изменения вместо сотен GPU
- Три способа посчитать одну и ту же функцию потерь
- Лицом к лицу: четыре конфигурации на одной H200
- Где метод раскрывается по-настоящему: длинный контекст
- Что получилось у студента весом 3,2 млрд параметров
- Что это меняет для индустрии и российских команд
Почему дистилляция стала самой дорогой частью конвейера
Дистилляция знаний — обучение компактной модели-студента подражать поведению большой модели-учителя — техника давно известная, но волна открытых LLM вроде gpt-oss, Qwen, GLM и Kimi вернула ей статус мейнстрима. Причина проста: развёртывание гигантов обходится неприлично дорого. Свежая модель Kimi-K3 насчитывает 2,8 трлн параметров, и только для её загрузки требуется порядка 3 ТБ видеопамяти. Поэтому сжатие моделей с последующим восстановлением их возможностей через дистилляцию превратилось в стандартную практику — качественные сжатые модели недавно выпустили Nvidia (Nemotron 3 Puzzle 75B) и сама Multiverse Computing (Hypernova 60B).
Именно этап дистилляции определяет большую часть итогового качества, но он же обычно и самый затратный во всём конвейере. Классическая схема — онлайн-дистилляция с KL-дивергенцией — требует одновременно держать в памяти и учителя, и студента. На каждом шаге учитель выполняет полный прямой проход и выдаёт распределение вероятностей по всему словарю, а студент учится его воспроизводить. Проблема в том, что поведение учителя в ходе обучения не меняется, но пересчитывать его приходится снова и снова.
Масштаб беды хорошо виден на конкретном примере. У gpt-oss-120b словарь насчитывает 201 088 токенов. При длине последовательности 32K и размере батча 4 один только тензор вероятностей учителя имеет форму 4 × 201 088 × 32 768 — в формате bfloat16 это уже около 50 ГБ видеопамяти на единственный тензор. Добавьте градиенты, активации, веса моделей и состояния оптимизатора — и одна итерация дистилляции достигает пика примерно в 250 ГБ VRAM. Это больше, чем способны предоставить даже H200 или B200.
Два системных изменения вместо сотен GPU
Первое изменение — переход к офлайн-дистилляции. Вместо того чтобы пересчитывать учителя на каждом шаге, команда вычисляет его выход один раз, кэширует top-100 наиболее вероятных токенов для каждой позиции и обучает студента уже по этому кэшу. Учителю больше не нужно находиться в памяти во время тренировки, а готовый кэш можно повторно использовать в десятках экспериментов и абляций — для исследовательских команд это принципиальная экономия.
Второе изменение — фьюзed-реализация KL-лосса с чанкованием. Чтобы понять, почему сам лосс так дорог, достаточно представить, что он строит: для каждой позиции в последовательности и каждого слова в словаре нужно число, описывающее расхождение предсказания студента с предсказанием учителя. В виде сетки это строка на каждую из 100 тысяч с лишним словарных единиц и столбец на каждую позицию последовательности. Стандартный способ вычисления KL-дивергенции строит всю эту гигантскую сетку целиком, прежде чем выдать хотя бы одно число. Именно здесь и кроется основной расход памяти, а не в весах моделей.
Три способа посчитать одну и ту же функцию потерь
Исследователи сравнили три математически эквивалентных способа вычислить один и тот же лосс. Первый — плотный KL (Dense KL), учебниковый подход: из кэшированных top-100 логитов восстанавливается полная плотная сетка вероятностей учителя и сравнивается с такой же плотной сеткой лог-вероятностей студента. Этот вариант ближе всего к тому, как работает онлайн-дистилляция, поэтому его взяли за эталон корректности, но платить приходится двойным хранением полной сетки «словарь × последовательность».
Второй способ — forward-chunked KL — оставляет учителя разреженным: его top-100 логиты никогда не разворачиваются в плотную сетку, а лосс считается по частям, срез за срезом позиций последовательности. Это убирает плотного учителя и плотное сравнение, и в бенчмарках команды именно этот метод оказался самым быстрым из трёх. Но у него остаётся слепое пятно: собственные логиты студента, которые выдаёт выходной слой модели, по-прежнему вычисляются полностью и хранятся ради обратного прохода, так что память всё равно круто растёт с длиной контекста.
Третий способ — fused chunked KL, главный вклад авторов. Здесь выходная проекция модели встроена прямо в вычисление лосса: полная сетка логитов студента не материализуется вообще. Последовательность обрабатывается чанками от начала до конца — скрытые состояния проецируются в логиты только для текущего чанка, результат сворачивается в накопленное значение лосса, а чанк выбрасывается до перехода к следующему. Обратный проход пересчитывает каждый чанк на лету вместо того, чтобы хранить его. Плата за это — двойная проекция (один раз вперёд, один раз назад), но взамен пиковое потребление памяти растёт лишь линейно с длиной последовательности, а не взмывает пропорционально произведению словаря на контекст. Реализация выложена в открытый доступ в репозитории CompactifAI/Full-Chunked-KL-Loss на GitHub.
Лицом к лицу: четыре конфигурации на одной H200
Все четыре схемы — онлайн-дистилляция и три офлайн-реализации лосса — команда сравнила на одной видеокарте H200. В роли учителя выступила Llama 3.1 8B Instruct, студентом стала модель семейства Llama на 3,2 млрд параметров, контекст — 8K токенов. Важнейший результат: все четыре варианта сошлись к практически идентичному значению обучающего лосса, хотя офлайн-прогоны обучались только по кэшированным top-100 логитам. Кривые обучения накладываются друг на друга почти идеально — то есть офлайн-дистилляция с top-100 кэшем по качеству неотличима от онлайн-версии.
| Метод (контекст 8K, одна H200) | Пиковая память | Время итерации | Пропускная способность |
|---|---|---|---|
| Онлайн-дистилляция | 102,8 ГБ | 25,9 с | 237 TFLOP/s |
| Офлайн, плотный 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 |
На контексте 8K фьюзed-вариант пока не самый быстрый — дополнительная проекция на обратном проходе отнимает немного скорости. Но его настоящее преимущество проявляется только с ростом длины контекста, что и показал следующий эксперимент.
Где метод раскрывается по-настоящему: длинный контекст
Чтобы наглядно увидеть закономерность масштабирования, авторы прогнали изолированный бенчмарк на игрушечной сети выходной проекции — без тела трансформера, только ядро лосса. На 32K токенов пиковая память падает с 85,2 ГиБ у плотного лосса до 5,45 ГиБ у полностью чанкованной версии — сокращение в 15,6 раза. А начиная с 64K токенов плотный лосс просто отказывается работать, упираясь в нехватку памяти.
На 256K токенов разрыв становится драматичным: полностью чанкованный лосс использует 11,6 ГиБ против 134,2 ГиБ у ближайшего чанкованного конкурента и при этом примерно в 3,3 раза быстрее на итерацию. Для сравнения: при дистилляции gpt-oss-120b на контексте 32K и батче 4 классическая плотная схема выстреливает пиком около 250 ГБ — выше возможностей одной H200 с её 141 ГБ, тогда как фьюзed-лосс в принципе не формирует этот пик и останавливается примерно на 128 ГБ.
Практическая проверка подтвердила цифры синтетических тестов. Дистилляция модели GPT-OSS 20B на контексте 32 768 токенов благодаря освобождённой памяти сжалась с четырёх GPU-узлов до одного. Время шага сократилось с 57,0 до 12,23 секунды — почти впятеро, а пропускная способность в пересчёте на одну GPU выросла с 74,2 до 345,7 TFLOP/s. Именно это и позволило авторам говорить о «лечении» длинного контекста на единственной видеокарте.
Что получилось у студента весом 3,2 млрд параметров
Эффективная офлайн-схема — то, что вообще сделало масштабную дистилляционную кампанию финансово осуществимой. Итоговый компактный студент, дистиллированный из Llama 3.1 8B Instruct примерно до 3,2 млрд параметров, сохраняет бо́льшую часть точности учителя на бенчмарках BoolQ и HellaSwag и отстаёт от него лишь примерно на девять пунктов на MMLU — при менее чем половинном числе параметров.
Иными словами, экономия на железе не обернулась заметной потерей качества: студент удерживает основную часть краткоконтекстной точности учителя при существенно меньшем размере. Для команды это доказательство, что top-100 кэшированных логитов несут достаточно сигнала для полноценного восстановления способностей модели.
Что это меняет для индустрии и российских команд
Работа вписывается в продолжающееся исследование Multiverse Computing по демократизации дистилляции — и приходит очень вовремя. Пока флагманские модели растут к триллионам параметров, а их обслуживание требует серверных стоек, индустрия всё активнее ищет способы упаковать те же возможности в модели, которые помещаются на одну-две карты. Методы вроде описанного смещают бутылочное горлышко с закупки железа на качество данных и грамотный подбор гиперпараметров.
Для российских команд значимость такой методики трудно переоценить. Доступ к современным ускорителям ограничен, очереди на вычислительные мощности в облаках растянуты, а стоимость аренды H100-класса карт кусается даже для крупных компаний. Возможность проводить полноценную дистилляцию длинноконтекстных моделей на одной GPU превращает задачу из «нужен кластер на сотни карт» в «нужна одна машина и терпение». Это открывает дорогу к сжатию открытых моделей под специфические задачи — от юридических ассистентов до промышленной аналитики — не только лабораториям уровня Яндекса или Сбера, но и небольшим продуктовым командам и университетским исследователям.
Отдельно стоит отметить открытость результата: реализация чанкованного лосса выложена на GitHub, а сама статья опубликована в блоге Hugging Face, так что воспроизвести и встроить метод в собственные конвейеры можно уже сейчас. Если практика подтвердит заявленные цифры на других архитектурах, офлайн-дистилляция с кэшированными top-K логитами имеет все шансы стать новым стандартом де-факто для сжатия LLM — так же, как когда-то стандартным стал сам приём дистилляции.
Источник: HF