Modelos de Lenguaje Compactos mediante Poda y Destilación de Conocimiento

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

Resumen

Los modelos de lenguaje grandes (LLMs) dirigidos a diferentes escalas y tamaños de implementación actualmente se producen mediante el entrenamiento de cada variante desde cero; esto es extremadamente intensivo en cómputo. En este documento, investigamos si podar un LLM existente y luego volver a entrenarlo con una fracción (<3%) de los datos de entrenamiento originales puede ser una alternativa adecuada al entrenamiento completo repetido. Con este fin, desarrollamos un conjunto de prácticas de compresión prácticas y efectivas para LLMs que combinan la poda de profundidad, ancho, atención y MLP con el reentrenamiento basado en destilación de conocimientos; llegamos a estas mejores prácticas a través de una exploración empírica detallada de estrategias de poda para cada eje, métodos para combinar ejes, estrategias de destilación y técnicas de búsqueda para llegar a arquitecturas comprimidas óptimas. Utilizamos esta guía para comprimir la familia de LLMs Nemotron-4 en un factor de 2-4 veces, y comparamos su rendimiento con modelos de tamaño similar en una variedad de tareas de modelado de lenguaje. Derivar modelos de 8B y 4B a partir de un modelo de 15B ya preentrenado utilizando nuestro enfoque requiere hasta 40 veces menos tokens de entrenamiento por modelo en comparación con el entrenamiento desde cero; esto resulta en ahorros de costos computacionales de 1.8x para el entrenamiento de la familia completa de modelos (15B, 8B y 4B). Los modelos Minitron muestran una mejora de hasta un 16% en las puntuaciones de MMLU en comparación con el entrenamiento desde cero, se desempeñan de manera comparable a otros modelos comunitarios como Mistral 7B, Gemma 7B y Llama-3 8B, y superan a las técnicas de compresión de vanguardia de la literatura. Hemos hecho públicos los pesos del modelo Minitron en Huggingface, con material suplementario correspondiente que incluye código de ejemplo disponible en 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.