Kompakte Sprachmodelle durch Beschneiden und Wissensvermittlung

Compact Language Models via Pruning and Knowledge Distillation

July 19, 2024

Autoren: Saurav Muralidharan, Sharath Turuvekere Sreenivas, Raviraj Joshi, Marcin Chochowski, Mostofa Patwary, Mohammad Shoeybi, Bryan Catanzaro, Jan Kautz, Pavlo Molchanov

cs.AI

Zusammenfassung

Große Sprachmodelle (LLMs), die auf verschiedene Bereitstellungsskalen und -größen abzielen, werden derzeit durch das Training jeder Variante von Grund auf erstellt; dies ist äußerst rechenintensiv. In diesem Papier untersuchen wir, ob das Beschneiden eines vorhandenen LLM und anschließendes Neutrainieren mit einem Bruchteil (<3%) der ursprünglichen Trainingsdaten eine geeignete Alternative zum wiederholten vollständigen Neutrainieren sein kann. Zu diesem Zweck entwickeln wir einen Satz praktischer und effektiver Komprimierungsempfehlungen für LLMs, die Tiefe, Breite, Aufmerksamkeit und MLP-Beschneidung mit auf Wissensvermittlung basierendem Neutrainieren kombinieren; wir gelangen zu diesen Empfehlungen durch eine detaillierte empirische Exploration von Beschneidungsstrategien für jede Achse, Methoden zur Kombination von Achsen, Vermittlungsstrategien und Suchtechniken zur Ermittlung optimaler komprimierter Architekturen. Wir verwenden diesen Leitfaden, um die Nemotron-4-Familie von LLMs um den Faktor 2-4x zu komprimieren und vergleichen deren Leistung mit ähnlich großen Modellen in einer Vielzahl von Sprachmodellierungsaufgaben. Die Ableitung von 8B- und 4B-Modellen aus einem bereits vortrainierten 15B-Modell mit unserem Ansatz erfordert bis zu 40-mal weniger Trainings-Token pro Modell im Vergleich zum Training von Grund auf; dies führt zu einer Rechenkosteneinsparung von 1,8x für das Training der gesamten Modellfamilie (15B, 8B und 4B). Minitron-Modelle zeigen eine bis zu 16%ige Verbesserung der MMLU-Werte im Vergleich zum Training von Grund auf, performen vergleichbar mit anderen Community-Modellen wie Mistral 7B, Gemma 7B und Llama-3 8B und übertreffen modernste Komprimierungstechniken aus der Literatur. Wir haben die Minitron-Modellgewichte auf Huggingface als Open Source veröffentlicht, zusammen mit entsprechendem ergänzendem Material, einschließlich Beispielcode auf 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.