MaskLLM: Lernbare halbstrukturierte Sparsamkeit für große Sprachmodelle

MaskLLM: Learnable Semi-Structured Sparsity for Large Language Models

September 26, 2024

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

cs.AI

Zusammenfassung

Große Sprachmodelle (LLMs) zeichnen sich durch ihre massiven Parameteranzahlen aus, die in der Regel zu erheblicher Redundanz führen. Diese Arbeit stellt MaskLLM vor, eine erlernbare Beschneidungsmethode, die Semi-strukturierte (oder "N:M") Sparsamkeit in LLMs einführt, um den Rechenaufwand während der Inferenz zu reduzieren. Anstatt ein neues Wichtigkeitskriterium zu entwickeln, modelliert MaskLLM N:M-Muster explizit als erlernbare Verteilung durch Gumbel-Softmax-Abtastung. Dieser Ansatz erleichtert das End-to-End-Training auf groß angelegten Datensätzen und bietet zwei bemerkenswerte Vorteile: 1) Hochwertige Masken - unsere Methode skaliert effektiv auf große Datensätze und lernt präzise Masken; 2) Übertragbarkeit - die probabilistische Modellierung der Maskenverteilung ermöglicht das Transferlernen von Sparsamkeit über Domänen oder Aufgaben hinweg. Wir haben MaskLLM unter Verwendung von 2:4-Sparsamkeit auf verschiedenen LLMs, einschließlich LLaMA-2, Nemotron-4 und GPT-3, mit Größen von 843M bis 15B Parametern, bewertet, und unsere empirischen Ergebnisse zeigen signifikante Verbesserungen gegenüber state-of-the-art Methoden. Beispielsweise erreichen führende Ansätze eine Perplexität (PPL) von 10 oder mehr auf Wikitext im Vergleich zu den 5,12 PPL des dichten Modells, aber MaskLLM erreicht allein durch das Erlernen der Masken mit eingefrorenen Gewichten eine deutlich niedrigere PPL von 6,72. Darüber hinaus ermöglicht die erlernbare Natur von MaskLLM maßgeschneiderte Masken für den verlustfreien Einsatz von 2:4-Sparsamkeit in nachgelagerten Aufgaben oder Domänen. Der Code ist unter https://github.com/NVlabs/MaskLLM verfügbar.

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.