MaskLLM: Leerbaar Semi-Gestructureerde Spaarzaamheid voor Grote Taalmodellen

Samenvatting

Grote Taalmodellen (LLM's) onderscheiden zich door hun enorme aantal parameters, wat doorgaans leidt tot aanzienlijke redundantie. Dit werk introduceert MaskLLM, een leerbare snoeimethode die Semi-gestructureerde (of "N:M") Spaarzaamheid in LLM's tot stand brengt, met als doel de rekenkundige overhead tijdens inferentie te verminderen. In plaats van een nieuwe belangrijkheidscriterium te ontwikkelen, modelleert MaskLLM N:M-patronen expliciet als een leerzame distributie via Gumbel Softmax-sampling. Deze benadering vergemakkelijkt end-to-end training op grootschalige datasets en biedt twee opmerkelijke voordelen: 1) Hoogwaardige Maskers - onze methode schaalt effectief naar grote datasets en leert nauwkeurige maskers; 2) Overdraagbaarheid - het probabilistisch modelleren van maskerdistributie maakt de transfer van spaarzaamheid tussen domeinen of taken mogelijk. We hebben MaskLLM geëvalueerd met 2:4 spaarzaamheid op verschillende LLM's, waaronder LLaMA-2, Nemotron-4 en GPT-3, met groottes variërend van 843M tot 15B parameters, en onze empirische resultaten tonen aanzienlijke verbeteringen ten opzichte van state-of-the-art methoden. Bijvoorbeeld, toonaangevende benaderingen behalen een perplexiteit (PPL) van 10 of meer op Wikitext in vergelijking met de 5.12 PPL van het dichte model, maar MaskLLM behaalt een aanzienlijk lagere 6.72 PPL alleen door de maskers met bevroren gewichten te leren. Bovendien maakt de leerbaarheid van MaskLLM aangepaste maskers mogelijk voor verliesloze toepassing van 2:4 spaarzaamheid op downstreamtaken of domeinen. De code is beschikbaar op 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.

MaskLLM: Leerbaar Semi-Gestructureerde Spaarzaamheid voor Grote Taalmodellen

MaskLLM: Learnable Semi-Structured Sparsity for Large Language Models

Samenvatting

Support