Новый метод в 15 раз снижает стоимость сжатия больших моделей ИИ
Исследователи нашли способ обучать меньшие и дешёвые модели ИИ на основе гигантских, используя лишь часть памяти, которая требовалась раньше. Это открывает возможность проводить эксперименты, для которых когда-то нужны были целые склады оборудования.
Ключевые моменты
- Дистилляция знаний — процесс копирования поведения большой модели ИИ в меньшую — обычно требует сотни специализированных чипов и тщательной координации.
- В новом методе выходные данные большой модели кэшируются один раз в начале, так что она больше не требуется во время обучения, что значительно снижает потребление памяти.
- Вторая техника, называемая fused chunked KL loss, обрабатывает обучающие данные небольшими частями вместо всего сразу, поддерживая пиковое потребление памяти почти неизменным даже при очень длинных входных последовательностях.
- При 32 000 токенов контекста пиковое потребление памяти падает с 85,2 гигабайт до 5,45 гигабайт — сокращение в 15,6 раз.
- Полный код энергоэффективной функции потерь доступен на GitHub.
Самые большие модели ИИ в мире имеют практически невообразимый размер. Kimi-K3, выпущенная в начале этого года, содержит 2,8 триллиона параметров (параметры можно представить как миллионы крошечных регуляторов внутри модели, которые настраиваются во время обучения) и требует примерно три терабайта специализированной памяти только для загрузки. Это больше, чем объём памяти двадцати высокопроизводительных графических карт для потребителей вместе взятых.
Поскольку запуск моделей такого размера чрезвычайно дорогостоящ, исследователи обычно сжимают их в меньшие версии, которые дешевле в эксплуатации. Стандартный метод называется дистилляцией знаний: вы запускаете большую «учительскую» модель на куче текста, записываете, как она бы ответила, а затем обучаете меньшую «студенческую» модель как можно точнее копировать эти ответы. Студенческая модель в итоге намного дешевле в работе, сохраняя при этом большую часть качества учительской модели.
Почему это так дорого делать?
Расходы возникают из двух источников. Во-первых, стандартный подход держит в памяти одновременно и учительскую, и студенческую модель. Во-вторых, сравнение их выходов требует построения огромной матрицы: одна строка для каждого слова в словаре модели (часто более 100 000 слов), один столбец для каждой позиции на входе. Умножьте это на длинный текст и вы быстро заполните всю доступную память на чипах.
Конкретный пример: одна популярная открытая модель gpt-oss-120b имеет словарь из 201 088 слов. При длине последовательности 32 000 токенов (примерно текст небольшого романа) и скромном размере пакета из четырёх примеров одновременно, только тензор выхода учителя занимает около 50 гигабайт. Добавьте студенческую модель, градиенты и всё остальное, что нужно процессу обучения, и один шаг обучения может потребовать примерно 250 гигабайт — больше, чем может вместить H200, один из самых мощных чипов ИИ, доступных на сегодня.
Что изменили исследователи?
Две вещи, работающие вместе. Обе описаны в статье, опубликованной командой CompactifAI, впервые поделившейся ею на Hugging Face.
Первое изменение — автономная дистилляция. Вместо того чтобы перезапускать учителя на каждом шаге обучения, исследователи запускают его один раз и сохраняют только его 100 наиболее вероятных прогнозов слов на каждой позиции. Этот кэш записывается на диск. Затем учитель полностью отключается, а студент обучается только на основе сохранённого кэша. Поскольку кэш можно использовать повторно в разных экспериментах, дорогостоящий запуск учителя становится единовременными затратами.
Второе изменение — fused chunked KL loss. KL-дивергенция — это математическая мера того, насколько различны две распределения вероятностей, в данном случае — насколько отличаются прогнозы слов студента от прогнозов учителя. Стандартный способ вычисления строит полную матрицу словарь-на-последовательность в один раз и держит всё это в памяти. Новый метод разбивает последовательность на маленькие куски, вычисляет и отбрасывает каждый кусок перед переходом к следующему, а при обратном проходе пересчитывает необходимые значения вместо их сохранения.
| Метод | Пиковое потребление памяти (8K контекста) | Время итерации |
|---|---|---|
| Онлайн дистилляция | 102,8 ГБ | 25,9 с |
| Автономная, плотная KL | 78,3 ГБ | 18,5 с |
| Автономная, forward-chunked KL | 61,8 ГБ | 18,4 с |
| Автономная, fused chunked KL | 58,3 ГБ | 20,2 с |
Все четыре метода дают практически идентичные результаты обучения, что означает, что потеря качества из-за использования кэшированных 100 лучших предсказаний вместо полного распределения учителя минимальна.
Имеет ли это значение для более длинных документов?
Да, и разница растёт быстро. Преимущество fused chunked метода в плане памяти скромно при 8 000 токенов, но драматично при больших длинах. При 32 000 токенов пиковое потребление памяти падает с 85,2 гигабайт при стандартном подходе до 5,45 гигабайт с новым. Плотная функция потерь полностью отказывает при 64 000 токенов. При 256 000 токенов fused chunked метод использует 11,6 гигабайт против 134,2 гигабайт для следующего по эффективности варианта.
Проще говоря: эксперименты, которые раньше требовали стойку с сотнями чипов, теперь могут работать на одной высокопроизводительной карте. Это делает реалистичным для небольших исследовательских групп и компаний создавать сжатые модели на основе самых крупных доступных систем ИИ с открытым исходным кодом.
Частые вопросы
Меньшая модель получается хуже оригинальной?
В представленных тестах потеря при обучении студенческой модели была практически одинаковой для всех четырёх методов, включая новый экономичный подход. Качество осталось на уровне несмотря на использование только 100 лучших предсказаний слов на позицию вместо полного выхода учителя.
Может ли кто-нибудь использовать эту технику?
Код chunked loss имеет открытый исходный код и опубликован на GitHub компанией CompactifAI. В принципе, любой, у кого есть один современный ИИ-чип и копия большой модели с открытым исходным кодом, может проводить собственные эксперименты по дистилляции с его помощью.
Кто выиграет в первую очередь?
Небольшие исследовательские группы и компании в области ИИ, которые хотят сжимать большие модели с открытым исходным кодом для более дешёвого развёртывания, почувствуют разницу в первую очередь. Конечные пользователи могут выиграть позже, когда более доступные сжатые модели появятся в продуктах и сервисах.
