Compact Language Models via Pruning and Knowledge Distillation
Large language models (LLMs) targeting different deployment scales and sizes
are currently produced by training each variant from scratch; this is extremely
compute-intensive. In this paper, we investigate if pruning an existing LLM and
then re-training it with a fraction (<3%) of the original training data can be
a suitable alternative to repeated, full retraining. To this end, we develop a
set of practical and effective compression best practices for LLMs that combine
depth, width, attention and MLP pruning with knowledge distillation-based
retraining; we arrive at these best practices through a detailed empirical
exploration of pruning strategies for each axis, methods to combine axes,
distillation strategies, and search techniques for arriving at optimal
compressed architectures. We use this guide to compress the Nemotron-4 family
of LLMs by a factor of 2-4x, and compare their performance to similarly-sized
models on a variety of language modeling tasks. Deriving 8B and 4B models from
an already pretrained 15B model using our approach requires up to 40x fewer
training tokens per model compared to training from scratch; this results in
compute cost savings of 1.8x for training the full model family (15B, 8B, and
4B). Minitron models exhibit up to a 16% improvement in MMLU scores compared to
training from scratch, perform comparably to other community models such as
Mistral 7B, Gemma 7B and Llama-3 8B, and outperform state-of-the-art
compression techniques from the literature. We have open-sourced Minitron model
weights on Huggingface, with corresponding supplementary material including
example code available on GitHub.
669ef80d65049e06488040ac
Saurav Muralidharan
claimed_verified
2024-07-24T07:28:34.804Z
64caca9b4ca74090b97a8795
/avatars/310f7703ff3c07d283c02aca56fe86df.svg
Saurav Muralidharan
user
srvm
669ef80d65049e06488040ad
Sharath Turuvekere Sreenivas
claimed_verified
2024-07-24T07:27:28.676Z
66a01db4d370743fd12531c0
https://cdn-avatars.huggingface.co/v1/production/uploads/66a01db4d370743fd12531c0/de8kFO3PY7nZ2MJd2F_w2.jpeg
Sharath Turuvekere Sreenivas
user
sharathts
669ef80d65049e06488040ae
Raviraj Joshi
669ef80d65049e06488040af
Marcin Chochowski
669ef80d65049e06488040b0
Mostofa Patwary
669ef80d65049e06488040b1
Mohammad Shoeybi
669ef80d65049e06488040b2
Bryan Catanzaro
669ef80d65049e06488040b3
Jan Kautz
669ef80d65049e06488040b4
Pavlo Molchanov
claimed_verified
2024-07-24T07:28:31.780Z
646d0c1c534e52f8c30500a6
https://cdn-avatars.huggingface.co/v1/production/uploads/646d0c1c534e52f8c30500a6/75VH8ClbRaP75BU2ONfXE.png
Pavlo Molchanov
user
pmolchanov
Compact Language Models via Pruning and Knowledge Distillation
Large language models (LLMs) targeting different deployment scales and sizes
are currently produced by training each variant from scratch; this is extremely
compute-intensive. In this paper, we investigate if pruning an existing LLM and
then re-training it with a fraction (<3%) of the original training data can be
a suitable alternative to repeated, full retraining. To this end, we develop a
set of practical and effective compression best practices for LLMs that combine
depth, width, attention and MLP pruning with knowledge distillation-based
retraining; we arrive at these best practices through a detailed empirical
exploration of pruning strategies for each axis, methods to combine axes,
distillation strategies, and search techniques for arriving at optimal
compressed architectures. We use this guide to compress the Nemotron-4 family
of LLMs by a factor of 2-4x, and compare their performance to similarly-sized
models on a variety of language modeling tasks. Deriving 8B and 4B models from
an already pretrained 15B model using our approach requires up to 40x fewer
training tokens per model compared to training from scratch; this results in
compute cost savings of 1.8x for training the full model family (15B, 8B, and
4B). Minitron models exhibit up to a 16% improvement in MMLU scores compared to
training from scratch, perform comparably to other community models such as
Mistral 7B, Gemma 7B and Llama-3 8B, and outperform state-of-the-art
compression techniques from the literature. We have open-sourced Minitron model
weights on Huggingface, with corresponding supplementary material including
example code available on GitHub.
Compressing large language models through strategic pruning and knowledge distillation reduces training compute costs and achieves performance comparable or better than similarly sized models trained from scratch.
pruning depth pruning width pruning attention pruning MLP pruning knowledge distillation Nemotron-4 compressed architectures MMLU scores Minitron Mistral 7B Gemma 7B Llama-3 8B