КПД

MaskLLM: Learnable Semi-Structured Sparsity for Large Language Models

[Статья] [Страница проекта] [Код][Пост на Machine Learning]

Введение

2:4 (она же semistructured sparsity) дает какое-никакое ускорение на GPU от Ampere и новее. Однако, просадка при прунинге обычно слишком велика для LLMок, дабы быть интересной на практике.

В этой статье предлагают метод обучения хороших 2:4 масок через Gumbel-Softmax.

Метод

Маска суть дискретная сущность потому ее просто так не отпизируешь градиентным спуском, и авторы предлагают моделировать распределение масок через Gumbel-Softmax с 6 = binom(2, 4) вариантам. На обучении оптимизируются логиты вероятности сэмплирования одного из вариантов масок (т.е маска есть взвешенная сумма возможных вариантов), а на инференсе берется наиболее вероятный. Обучение суть просто оптимизация кросс-энтропии (как на pretrain). Веса при этом заморожены.

Если какой-то вес зануляется или близок к нулю, то логиты маски почти не получают градиентов, потому авторы добавляют регуляризационный член как weight_decay, но со знаком ➖, чтобы расталкивать веса от нуля, тем самым поддерживая не нулевую норму у немаскированных весов.

Кроме того, маски полученные условным SparseGPT/Wanda являются хорошей инициализацией для масок и позволяют чуть улучшить результат.

Эксперименты

Метод валидируют на 🦙-2, Nemotron-4 15B и двух маленьких проприетарных GPT-3. Замеряют по классике перплексию на Wikitext и 0-шоты.

По метрикам опережают уверенно все бейзлайны (SparseGPT, Wanda, Magnitude). SparseGPT, правда, можно завести и получше. В отличие от алгоритмов one-shot прунинга, которые быстро насыщаются от количества данных, MaskLLM продолжает улучшаться при большем и большем количестве данных, что неудивительно ибо это есть по сути метод оптимизации с большим количеством обучаемых параметров.

Ablations:
1️⃣ Инициализация маской от one-shot прунера накидывает в конечном качестве.
2️⃣ Достаточная степень стохастичности сэмплирования важна для хорошего качества, дабы модель могла “попробовать” разные варианты масок.
3️⃣ Анти-weight decay не то чтобы сильно, но улучшает качество.
4️⃣ Кроме того, полученную маску можно оптимизировать на downstream и даже временами оверфитнуться улучшить перплексию по сравнению с floating-point моделью.

Вывод

Вполне годная стратегия для обучения 2:4, но требующая определенных вычислительных затрат (т.е прилично дороже чем прогнать SparseGPT). Результат достойный, но все же просадка остается довольно заметной - больше чем у SOTA методов 2-битной квантизации. Вероятно, если еще оптимизировать веса вместе с масками - можно выжать больше.

👍6❤2

www.tgoop.com/quant_prune_distill/338

1.62K viewsedited  Oct 15, 2024 at 21:11

MaskLLM: Learnable Semi-Structured Sparsity for Large Language Models

[Статья] [Страница проекта] [Код][Пост на Machine Learning]

Введение

2:4 (она же semistructured sparsity) дает какое-никакое ускорение на GPU от Ampere и новее. Однако, просадка при прунинге обычно слишком велика для LLMок, дабы быть интересной на практике.

В этой статье предлагают метод обучения хороших 2:4 масок через Gumbel-Softmax.

Метод

Маска суть дискретная сущность потому ее просто так не отпизируешь градиентным спуском, и авторы предлагают моделировать распределение масок через Gumbel-Softmax с 6 = binom(2, 4) вариантам. На обучении оптимизируются логиты вероятности сэмплирования одного из вариантов масок (т.е маска есть взвешенная сумма возможных вариантов), а на инференсе берется наиболее вероятный. Обучение суть просто оптимизация кросс-энтропии (как на pretrain). Веса при этом заморожены.

Если какой-то вес зануляется или близок к нулю, то логиты маски почти не получают градиентов, потому авторы добавляют регуляризационный член как weight_decay, но со знаком ➖, чтобы расталкивать веса от нуля, тем самым поддерживая не нулевую норму у немаскированных весов.

Кроме того, маски полученные условным SparseGPT/Wanda являются хорошей инициализацией для масок и позволяют чуть улучшить результат.

Эксперименты

Метод валидируют на 🦙-2, Nemotron-4 15B и двух маленьких проприетарных GPT-3. Замеряют по классике перплексию на Wikitext и 0-шоты.

По метрикам опережают уверенно все бейзлайны (SparseGPT, Wanda, Magnitude). SparseGPT, правда, можно завести и получше. В отличие от алгоритмов one-shot прунинга, которые быстро насыщаются от количества данных, MaskLLM продолжает улучшаться при большем и большем количестве данных, что неудивительно ибо это есть по сути метод оптимизации с большим количеством обучаемых параметров.

Ablations:
1️⃣ Инициализация маской от one-shot прунера накидывает в конечном качестве.
2️⃣ Достаточная степень стохастичности сэмплирования важна для хорошего качества, дабы модель могла “попробовать” разные варианты масок.
3️⃣ Анти-weight decay не то чтобы сильно, но улучшает качество.
4️⃣ Кроме того, полученную маску можно оптимизировать на downstream и даже временами оверфитнуться улучшить перплексию по сравнению с floating-point моделью.

Вывод

Вполне годная стратегия для обучения 2:4, но требующая определенных вычислительных затрат (т.е прилично дороже чем прогнать SparseGPT). Результат достойный, но все же просадка остается довольно заметной - больше чем у SOTA методов 2-битной квантизации. Вероятно, если еще оптимизировать веса вместе с масками - можно выжать больше.

BY КПД