NeuZip: метод сжатия весов для обучения и инференса

NeuZip (https://arxiv.org/pdf/2410.20650) - алгоритм сжатия весов нейронных сетей, разработанный для снижения требований к памяти во время обучения и инференса без ущерба для производительности.

В основе NeuZip лежит принцип энтропии чисел с плавающей запятой в нейронных сетях. Веса нейронных сетей концентрируются вокруг нуля, что приводит к низкой энтропии битов экспоненты, а значит, почему бы не сжимать биты экспоненты без потерь с помощью асимметричной системы счисления (ANS)?

ANS — это алгоритм сжатия без потерь, который обеспечивает высокую пропускную способность на параллельных вычислительных устройствах, например, на GPU.

Для обучения используется вариант NeuZip без потерь, который сжимает только биты экспоненты, сохраняя полную точность представления чисел. В процессе обучения веса хранятся в сжатом виде, а декомпрессия происходит послойно, непосредственно перед вычислениями . Это позволяет избежать дублирования памяти и снизить ее пиковое потребление. При этом backpropagation не затрагивается, так как градиенты вычисляются с использованием декомпрессированных весов.

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

Эксперименты, проведенные на различных архитектурах (GPT, Llama, T5) и задачах (языковое моделирование, генерация SQL), подтвердили эффективность NeuZip. В частности, при обучении модели Llama-3 8B удалось сократить потребление памяти с 31 ГБ до менее 16 ГБ без изменения динамики обучения.

В задачах инференса NeuZip демонстрирует достижение >50% сокращения памяти при сохранении практически идентичной производительности по сравнению с QLoRA и современными методами квантования.

Код экспериментов из пейпера в задачах обучения и инференса с Neuzip доступен (https://github.com/BorealisAI/neuzip/blob/main/examples) в /examples репозитория проекта (https://github.com/BorealisAI/neuzip) на Github.

Установка и использование:

# Install from PyPI  
 pip install neuzip 
 
 # Use Neuzip for Pytorch model 
 
 model: torch.nn.Module = # your model 
 
 + manager = neuzip.Manager() 
 
 + model = manager.convert(model)

Arxiv (https://arxiv.org/pdf/2410.20650)

GitHub (https://github.com/BorealisAI/neuzip)

Источник: github.com

Комментарии: