MaskLLM : Sparsité Semi-Structurée Apprenable pour les Grands Modèles de Langage

MaskLLM: Learnable Semi-Structured Sparsity for Large Language Models

September 26, 2024

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

cs.AI

Résumé

Les grands modèles de langage (LLM) se distinguent par leurs énormes nombres de paramètres, qui entraînent généralement une redondance significative. Ce travail présente MaskLLM, une méthode de taille adaptable qui établit une Sparsité Semi-structurée (ou "N:M") dans les LLM, visant à réduire la charge computationnelle lors de l'inférence. Au lieu de développer un nouveau critère d'importance, MaskLLM modélise explicitement les motifs N:M comme une distribution apprenable à travers un échantillonnage Gumbel Softmax. Cette approche facilite l'entraînement de bout en bout sur des ensembles de données à grande échelle et offre deux avantages notables : 1) Masques de haute qualité - notre méthode s'adapte efficacement aux grands ensembles de données et apprend des masques précis ; 2) Transférabilité - la modélisation probabiliste de la distribution des masques permet le transfert d'apprentissage de la sparsité entre domaines ou tâches. Nous avons évalué MaskLLM en utilisant une sparsité de 2:4 sur différents LLM, y compris LLaMA-2, Nemotron-4 et GPT-3, avec des tailles allant de 843M à 15B de paramètres, et nos résultats empiriques montrent des améliorations substantielles par rapport aux méthodes de pointe. Par exemple, les approches principales obtiennent une perplexité (PPL) de 10 ou plus sur Wikitext par rapport au modèle dense à 5,12 PPL, mais MaskLLM atteint significativement 6,72 PPL uniquement en apprenant les masques avec des poids figés. De plus, la nature apprenable de MaskLLM permet des masques personnalisés pour une application sans perte de la sparsité de 2:4 aux tâches ou domaines en aval. Le code est disponible sur 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.