Оптимизация больших языковых моделей: удаление блоков трансформеров как метод глубокой компрессии


Ускорение работы больших языковых моделей (LLM) можно достичь различными способами, одним из самых простых и эффективных является удаление целых блоков трансформеров. Этот процесс, известный как глубинная обрезка (depth pruning), позволяет не только сократить размер модели, но и значительно повысить скорость вывода, а также сэкономить память. Однако, несмотря на очевидные преимущества, задача выбора блоков для удаления оказывается сложной. Неправильный выбор может привести к ухудшению качества модели, а взаимодействие между удаляемыми блоками делает эту задачу комбинаторной, а не простой. В нашей последней работе "LLM Compression by Block Removal with Constrained Binary Optimization" мы рассматриваем эту задачу через призму физики спиновых систем.

Проблема выбора блоков

Существующие методы удаления блоков обычно оценивают каждый блок отдельно, а затем удаляют те, которые кажутся наименее важными. Это похоже на среднеполевые методы в физике, где влияние соседей заменяется усредненным полем. Однако блоки в трансформерах не независимы, и удаление одного блока может повлиять на эффективность других. Например, удаление блока 20 зависит от того, удален ли блок 19 или 24. Это взаимодействие между блоками делает задачу выбора блоков сложной и требует учета всех возможных комбинаций.

Оптимизация выбора блоков

Мы предлагаем рассматривать выбор блоков как задачу оптимизации с ограничениями (CBO), которая моделируется как система Изинга. В этой модели каждому блоку присваивается бинарная переменная: 0 — сохранить блок, 1 — удалить. Мы используем второе производное (Гессиан) потерь модели относительно этих переменных, чтобы оценить, как удаление блоков повлияет на производительность модели. Это позволяет нам минимизировать энергию системы, что в свою очередь соответствует высокому качеству модели.

Практическое применение и результаты

Расчет Гессиана требует лишь одного прохода по небольшому набору данных, после чего оценка любой конфигурации становится быстрой и дешевой. Мы можем проверять десятки миллиардов конфигураций, используя один GPU. Например, для удаления 8 из 80 блоков модели Llama-3.3-70B это заняло около двух дней. Когда конфигурационное пространство становится слишком большим, мы можем использовать оптимизированные классические и квантовые решатели, которые эффективно работают с такими задачами.

Энергетический спектр и качество модели

Энергия системы Изинга служит надежным прокси для качества модели, но не является идеальной. Поэтому мы исследуем не только основное состояние, но и низкоэнергетические возбуждения, что позволяет нам получить несколько качественных кандидатов на обрезку. Например, в случае модели Llama-3.1-8B-Instruct, 17-е возбужденное состояние предложило удалить блок в начале модели и после легкой дообучения превзошло основное состояние по нескольким метрикам.

Применение к современным архитектурам

Метод удаления блоков также применим к современным гетерогенным архитектурам, где различные типы блоков переплетены. Мы протестировали метод на модели NVIDIA-Nemotron-3-Nano-30B-A3B-FP8, которая сочетает в себе различные слои, и результаты подтвердили, что некоторые слои более подвержены удалению, чем другие. Это подчеркивает важность поиска в связанном конфигурационном пространстве.

Заключение

Переформулирование задачи машинного обучения как задачи Гамильтона Изинга и использование классических и квантовых методов оптимизации является частью философии Multiverse Computing. Удаление блоков компонуется с другими методами, такими как квантование и сжатие по низкому рангу, создавая целостный подход к оптимизации языковых моделей. Для получения более подробной информации о методах и результатах, вы можете ознакомиться с полной версией работы на Hugging Face или связаться с нашей командой для обсуждения применения этих методов к вашим моделям. Код доступен на github.com/CompactifAI/Block_removal_through_constrained_binary_optimization.

Комментарии (0)

Войдите, чтобы оставлять комментарии.
Пока нет комментариев. Будьте первым!