BitStack: Эффективное управление памятью для сжатия больших языковых моделей

С ростом популярности больших языковых моделей (LLMs) их применение охватывает широкий спектр задач, от генерации текста до разработки программного обеспечения. Однако, несмотря на их впечатляющие возможности, развертывание LLMs на локальных устройствах часто сталкивается с серьезными ограничениями по памяти. Эта проблема становится особенно актуальной в условиях переменной доступности памяти, когда ресурсы могут колебаться в зависимости от других запущенных приложений. В этой статье мы рассмотрим новый подход к сжатию весов, предложенный в работе "BitStack", который позволяет динамически управлять размером модели, обеспечивая при этом высокую производительность.

Проблема доступности памяти

С увеличением размеров LLMs, таких как Llama3, требования к оперативной памяти (RAM) становятся критическими. Например, для загрузки весов модели Llama3 с 18 миллиардами параметров требуется около 15 ГБ RAM в формате FP16, что не учитывает дополнительных ресурсов, необходимых для активаций во время вывода. Это создает значительные трудности для развертывания на устройствах с ограниченной памятью, таких как смартфоны и персональные компьютеры.

Традиционные методы сжатия, такие как квантование и обрезка, часто требуют заранее определенных коэффициентов сжатия и отдельного процесса сжатия для каждой настройки, что усложняет развертывание в условиях переменной памяти. Поэтому необходима стратегия, позволяющая динамически адаптировать размер модели в зависимости от доступных ресурсов.

Подход BitStack

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

Декомпозиция весов

Процесс декомпозиции начинается с применения сингулярного разложения (SVD) к весовым матрицам. Формально, весовая матрица ( W ) может быть представлена как:

[ W = U \Sigma V^T ]

где ( U ) и ( V ) — матрицы левых и правых сингулярных векторов соответственно, а ( \Sigma ) — диагональная матрица с сингулярными значениями. Важно отметить, что в BitStack используется не только само разложение, но и дополнительные шаги для учета значимости весов.

Активация с учетом важности

Одной из ключевых особенностей BitStack является использование активации с учетом важности. Это означает, что перед декомпозицией весовая матрица масштабируется с использованием вектора масштабирования, который вычисляется на основе L2-нормы активаций. Это позволяет сохранить точность восстановления значимых весов и минимизировать ошибки сжатия.

Итеративное абсолютное разложение

Для снижения ошибки приближения в каждом шаге декомпозиции применяется абсолютное значение разложения, что позволяет сохранить информацию о знаках весов. Это позволяет эффективно упаковывать знаковые матрицы в типы данных, поддерживаемые графическими процессорами (GPU), что снижает требования к памяти.

В результате, в каждом шаге декомпозиции достигается уменьшение объема памяти примерно на 1 бит на параметр, что делает BitStack очень эффективным для больших моделей.

Динамическое управление памятью

Одной из самых сильных сторон BitStack является возможность динамически загружать и выгружать остаточные блоки весов в зависимости от текущей доступности памяти. Это позволяет моделям адаптироваться к изменяющимся условиям и эффективно использовать доступные ресурсы. Когда память становится доступной, можно загружать дополнительные остаточные блоки из заранее отсортированного стека, а при нехватке памяти — выгружать их обратно на устройства хранения.

Эксперименты и результаты

BitStack был протестирован на различных моделях, включая Llama2 и Llama3, с размерами от 7 до 70 миллиардов параметров. Результаты экспериментов показывают, что несмотря на возможность сжатия, BitStack consistently matches or surpasses the performance of widely adopted compression methods like GPTQ and AWQ, especially at extreme compression ratios.

Оценка производительности

В ходе экспериментов были вычислены показатели перплексии на тестовом наборе WikiText2 и точность на различных задачах нулевого выстрела. BitStack продемонстрировал выдающиеся результаты, показывая, что даже при высоких коэффициентах сжатия, он сохраняет производительность на уровне, сопоставимом с традиционными методами сжатия.

Выводы

BitStack предлагает новый подход к сжатию LLMs, который не только решает проблемы доступности памяти, но и обеспечивает высокую производительность в условиях переменной памяти. Это делает его идеальным решением для развертывания больших языковых моделей на локальных устройствах, где ресурсы могут быть ограничены.

Заключение

С учетом растущих требований к памяти и вычислительным ресурсам, подход BitStack может стать важным шагом вперед в области сжатия и оптимизации больших языковых моделей. Он открывает новые возможности для использования LLMs в приложениях, требующих высокой производительности при ограниченных ресурсах.