MaskLLM: Esparsidade Semiestruturada Aprendível para Modelos de Linguagem de Grande Escala

MaskLLM: Learnable Semi-Structured Sparsity for Large Language Models

September 26, 2024

Autores: Gongfan Fang, Hongxu Yin, Saurav Muralidharan, Greg Heinrich, Jeff Pool, Jan Kautz, Pavlo Molchanov, Xinchao Wang

cs.AI

Resumo

Os Modelos de Linguagem de Grande Escala (LLMs) se destacam por seus enormes números de parâmetros, que geralmente resultam em uma redundância significativa. Este trabalho apresenta o MaskLLM, um método de poda aprendível que estabelece a Esparsidade Semi-estruturada (ou "N:M") em LLMs, com o objetivo de reduzir a sobrecarga computacional durante a inferência. Em vez de desenvolver um novo critério de importância, o MaskLLM modela explicitamente padrões N:M como uma distribuição aprendível por meio de amostragem Gumbel Softmax. Essa abordagem facilita o treinamento de ponta a ponta em conjuntos de dados em grande escala e oferece duas vantagens notáveis: 1) Máscaras de alta qualidade - nosso método escala efetivamente para grandes conjuntos de dados e aprende máscaras precisas; 2) Transferibilidade - a modelagem probabilística da distribuição de máscaras permite a aprendizagem de transferência de esparsidade entre domínios ou tarefas. Avaliamos o MaskLLM usando uma esparsidade de 2:4 em vários LLMs, incluindo LLaMA-2, Nemotron-4 e GPT-3, com tamanhos variando de 843M a 15B parâmetros, e nossos resultados empíricos mostram melhorias substanciais em relação aos métodos de ponta. Por exemplo, abordagens líderes alcançam uma perplexidade (PPL) de 10 ou mais no Wikitext em comparação com os 5,12 PPL do modelo denso, mas o MaskLLM alcança um PPL significativamente menor de 6,72 apenas aprendendo as máscaras com pesos congelados. Além disso, a natureza aprendível do MaskLLM permite máscaras personalizadas para a aplicação sem perdas de uma esparsidade de 2:4 em tarefas ou domínios subsequentes. O código está disponível em 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.