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

MaskLLM: Learnable Semi-Structured Sparsity for Large Language Models

September 26, 2024

Авторы: Gongfan Fang, Hongxu Yin, Saurav Muralidharan, Greg Heinrich, Jeff Pool, Jan Kautz, Pavlo Molchanov, Xinchao Wang

cs.AI

Аннотация

Большие языковые модели (LLM) отличаются своими огромными параметрами, что обычно приводит к значительной избыточности. В данной работе представлена MaskLLM, обучаемый метод обрезки, который устанавливает Полуструктурированную (или "N:M") Разреженность в LLM с целью снижения вычислительной нагрузки во время вывода. Вместо разработки нового критерия важности, MaskLLM явно моделирует образцы N:M как обучаемое распределение с помощью выборки Gumbel Softmax. Этот подход облегчает обучение на масштабных наборах данных и предлагает два примечательных преимущества: 1) Маски высокого качества - наш метод эффективно масштабируется на большие наборы данных и изучает точные маски; 2) Переносимость - вероятностное моделирование распределения масок позволяет переносить обучение разреженности между областями или задачами. Мы оценили MaskLLM, используя разреженность 2:4 на различных LLM, включая LLaMA-2, Nemotron-4 и GPT-3, с размерами параметров от 843 миллионов до 15 миллиардов, и наши эмпирические результаты показывают существенные улучшения по сравнению с современными методами. Например, ведущие подходы достигают перплексии (PPL) 10 или более на Wikitext по сравнению с плотной моделью, у которой PPL 5,12, но MaskLLM достигает значительно более низкой PPL 6,72 только за счет обучения масок с зафиксированными весами. Более того, обучаемая природа MaskLLM позволяет создавать настраиваемые маски для без потерь применения разреженности 2:4 к последующим задачам или областям. Код доступен по адресу https://github.com/NVlabs/MaskLLM.

English

Large Language Models (LLMs) are distinguished by their massive parameter counts, which typically result in significant redundancy. This work introduces MaskLLM, a learnable pruning method that establishes Semi-structured (or ``N:M'') Sparsity in LLMs, aimed at reducing computational overhead during inference. Instead of developing a new importance criterion, MaskLLM explicitly models N:M patterns as a learnable distribution through Gumbel Softmax sampling. This approach facilitates end-to-end training on large-scale datasets and offers two notable advantages: 1) High-quality Masks - our method effectively scales to large datasets and learns accurate masks; 2) Transferability - the probabilistic modeling of mask distribution enables the transfer learning of sparsity across domains or tasks. We assessed MaskLLM using 2:4 sparsity on various LLMs, including LLaMA-2, Nemotron-4, and GPT-3, with sizes ranging from 843M to 15B parameters, and our empirical results show substantial improvements over state-of-the-art methods. For instance, leading approaches achieve a perplexity (PPL) of 10 or greater on Wikitext compared to the dense model's 5.12 PPL, but MaskLLM achieves a significantly lower 6.72 PPL solely by learning the masks with frozen weights. Furthermore, MaskLLM's learnable nature allows customized masks for lossless application of 2:4 sparsity to downstream tasks or domains. Code is available at https://github.com/NVlabs/MaskLLM.