Modelos de Linguagem Compactos via Poda e Destilação de Conhecimento
Compact Language Models via Pruning and Knowledge Distillation
July 19, 2024
Autores: Saurav Muralidharan, Sharath Turuvekere Sreenivas, Raviraj Joshi, Marcin Chochowski, Mostofa Patwary, Mohammad Shoeybi, Bryan Catanzaro, Jan Kautz, Pavlo Molchanov
cs.AI
Resumo
Modelos de linguagem grandes (LLMs) direcionados a diferentes escalas e tamanhos de implantação são atualmente produzidos treinando cada variante do zero; isso é extremamente intensivo em termos de computação. Neste artigo, investigamos se a poda de um LLM existente e, em seguida, o seu retratamento com uma fração (<3%) dos dados de treinamento originais pode ser uma alternativa adequada ao treinamento completo repetido. Para isso, desenvolvemos um conjunto de práticas de compressão práticas e eficazes para LLMs que combinam poda de profundidade, largura, atenção e MLP com retratamento baseado em destilação de conhecimento; chegamos a essas melhores práticas por meio de uma exploração empírica detalhada de estratégias de poda para cada eixo, métodos para combinar eixos, estratégias de destilação e técnicas de busca para chegar a arquiteturas comprimidas ótimas. Utilizamos este guia para comprimir a família de LLMs Nemotron-4 por um fator de 2-4x, e comparamos seu desempenho com modelos de tamanho semelhante em uma variedade de tarefas de modelagem de linguagem. Derivar modelos de 8B e 4B a partir de um modelo pré-treinado de 15B usando nossa abordagem requer até 40x menos tokens de treinamento por modelo em comparação com o treinamento do zero; isso resulta em economia de custos computacionais de 1,8x para treinar a família completa de modelos (15B, 8B e 4B). Os modelos Minitron apresentam até 16% de melhoria nos escores de MMLU em comparação com o treinamento do zero, têm desempenho comparável a outros modelos comunitários como Mistral 7B, Gemma 7B e Llama-3 8B, e superam técnicas de compressão de ponta da literatura. Disponibilizamos os pesos do modelo Minitron no Huggingface, com material suplementar correspondente, incluindo código de exemplo disponível no GitHub.
English
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.