MaskLLM: Esparsidad Semi-Estructurada Aprendible para Modelos de Lenguaje Grandes

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

Resumen

Los Modelos de Lenguaje Grandes (LLMs) se distinguen por sus enormes recuentos de parámetros, que generalmente resultan en una redundancia significativa. Este trabajo presenta MaskLLM, un método de poda aprendible que establece una Esparsidad Semiestructurada (o "N:M") en LLMs, con el objetivo de reducir la sobrecarga computacional durante la inferencia. En lugar de desarrollar un nuevo criterio de importancia, MaskLLM modela explícitamente patrones N:M como una distribución aprendible a través de muestreo Gumbel Softmax. Este enfoque facilita el entrenamiento de extremo a extremo en conjuntos de datos a gran escala y ofrece dos ventajas notables: 1) Máscaras de alta calidad: nuestro método escala eficazmente a conjuntos de datos grandes y aprende máscaras precisas; 2) Transferibilidad: la modelización probabilística de la distribución de máscaras permite el aprendizaje de transferencia de la esparsidad entre dominios o tareas. Evaluamos MaskLLM utilizando una esparsidad de 2:4 en varios LLMs, incluidos LLaMA-2, Nemotron-4 y GPT-3, con tamaños que van desde 843M hasta 15B parámetros, y nuestros resultados empíricos muestran mejoras sustanciales sobre los métodos de vanguardia. Por ejemplo, en Wikitext, enfoques líderes logran una perplejidad (PPL) de 10 o más en comparación con el 5.12 PPL del modelo denso, pero MaskLLM logra un PPL significativamente menor de 6.72 únicamente aprendiendo las máscaras con pesos congelados. Además, la naturaleza aprendible de MaskLLM permite máscaras personalizadas para la aplicación sin pérdida de una esparsidad de 2:4 en tareas o dominios posteriores. El código está disponible en 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.