MaskLLM: 대규모 언어 모델을 위한 학습 가능한 반구조 희소성

MaskLLM: Learnable Semi-Structured Sparsity for Large Language Models

September 26, 2024

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

cs.AI

초록

대형 언어 모델 (LLM)은 일반적으로 상당한 중복을 초래하는 대규모 매개변수 개수로 특징 지어집니다. 본 연구는 추론 중의 계산 오버헤드를 줄이기 위해 LLM에 반구조화 (또는 "N:M") 희소성을 정립하는 학습 가능한 가지치기 방법인 MaskLLM을 소개합니다. 새로운 중요도 기준을 개발하는 대신 MaskLLM은 N:M 패턴을 Gumbel Softmax 샘플링을 통해 학습 가능한 분포로 명시적으로 모델링합니다. 이 접근 방식은 대규모 데이터셋에서의 종단 간 훈련을 용이하게 하며 두 가지 주목할만한 이점을 제공합니다: 1) 고품질 마스크 - 우리의 방법은 대규모 데이터셋으로 효과적으로 확장되며 정확한 마스크를 학습합니다; 2) 전이성 - 마스크 분포의 확률적 모델링은 도메인이나 작업 간의 희소성 전이 학습을 가능하게 합니다. 우리는 LLaMA-2, Nemotron-4, GPT-3을 포함한 다양한 LLM에 843M에서 15B 매개변수까지 크기가 다른 2:4 희소성을 적용하여 MaskLLM을 평가했으며 실험 결과는 최신 기술 방법에 비해 상당한 개선을 보여줍니다. 예를 들어, 선도적인 접근 방식은 Wikitext에서 밀도 모델의 5.12 PPL에 비해 10 이상의 헷갈림도 (PPL)를 달성하지만 MaskLLM은 얼려진 가중치로 마스크를 학습함으로써 6.72 PPL로 상당히 낮은 값을 달성합니다. 더 나아가, MaskLLM의 학습 가능한 특성은 하류 작업이나 도메인에 2:4 희소성을 손실 없이 적용하기 위한 사용자 정의 마스크를 가능하게 합니다. 코드는 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.