MaskLLM: Apprendimento della sparità semi-strutturata apprendibile per modelli di linguaggio di grandi dimensioni

MaskLLM: Learnable Semi-Structured Sparsity for Large Language Models

September 26, 2024

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

cs.AI

Abstract

I Large Language Models (LLM) si distinguono per il loro massiccio conteggio dei parametri, che tipicamente comporta una significativa ridondanza. Questo lavoro introduce MaskLLM, un metodo di potatura apprendibile che stabilisce una Sparità Semi-strutturata (o "N:M") nei LLM, mirata a ridurre il carico computazionale durante l'inferezza. Invece di sviluppare un nuovo criterio di importanza, MaskLLM modella esplicitamente i pattern N:M come una distribuzione apprendibile attraverso il campionamento Gumbel Softmax. Questo approccio facilita l'addestramento end-to-end su dataset su larga scala e offre due vantaggi notevoli: 1) Maschere di alta qualità - il nostro metodo scala efficacemente su dataset di grandi dimensioni e apprende maschere accurate; 2) Trasferibilità - la modellazione probabilistica della distribuzione delle maschere consente il trasferimento dell'apprendimento della sparità tra domini o compiti. Abbiamo valutato MaskLLM utilizzando una sparità 2:4 su vari LLM, inclusi LLaMA-2, Nemotron-4 e GPT-3, con dimensioni che vanno da 843M a 15B parametri, e i nostri risultati empirici mostrano miglioramenti sostanziali rispetto ai metodi all'avanguardia. Ad esempio, approcci principali raggiungono una perplessità (PPL) di 10 o superiore su Wikitext rispetto al 5.12 PPL del modello denso, ma MaskLLM raggiunge un significativamente inferiore 6.72 PPL semplicemente apprendendo le maschere con pesi congelati. Inoltre, la natura apprendibile di MaskLLM consente maschere personalizzate per un'applicazione senza perdite della sparità 2:4 a compiti o domini successivi. Il codice è disponibile su 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.