Rätsel: Distillationsbasierte NAS zur Inferenz-optimierten LLMs

Puzzle: Distillation-Based NAS for Inference-Optimized LLMs

November 28, 2024

Autoren: Akhiad Bercovich, Tomer Ronen, Talor Abramovich, Nir Ailon, Nave Assaf, Mohammad Dabbah, Ido Galil, Amnon Geifman, Yonatan Geifman, Izhak Golan, Netanel Haber, Ehud Karpas, Itay Levy, Shahar Mor, Zach Moshe, Najeeb Nabwani, Omri Puny, Ran Rubin, Itamar Schen, Ido Shahaf, Oren Tropp, Omer Ullman Argov, Ran Zilberstein, Ran El-Yaniv

cs.AI

Zusammenfassung

Große Sprachmodelle (LLMs) haben bemerkenswerte Fähigkeiten gezeigt, aber ihre Übernahme wird durch hohe Rechenkosten während der Inferenz eingeschränkt. Während eine Erhöhung der Parameteranzahl die Genauigkeit verbessert, vergrößert sie auch die Kluft zwischen den modernsten Fähigkeiten und der praktischen Einsatzfähigkeit. Wir stellen Puzzle vor, ein Framework zur Beschleunigung der LLM-Inferenz auf spezifischer Hardware, während ihre Fähigkeiten erhalten bleiben. Durch eine innovative Anwendung der neuronalen Architektursuche (NAS) in einem beispiellosen Maßstab optimiert Puzzle systematisch Modelle mit zig Milliarden Parametern unter Hardwarebeschränkungen. Unser Ansatz nutzt blockweise lokale Wissensvermittlung (BLD) für parallele Architekturerkundung und setzt gemischt-ganzzahlige Programmierung für präzise Beschränkungsoptimierung ein. Wir zeigen den realen Einfluss unseres Frameworks durch Llama-3.1-Nemotron-51B-Instruct (Nemotron-51B), ein öffentlich verfügbares Modell, abgeleitet von Llama-3.1-70B-Instruct. Nemotron-51B erreicht eine 2,17-fache Beschleunigung der Inferenzdurchsatzgeschwindigkeit, passt auf eine einzelne NVIDIA H100 GPU und erhält dabei 98,4% der Fähigkeiten des Originalmodells. Nemotron-51B gilt derzeit als das genaueste Sprachmodell, das Inferenz auf einer einzelnen GPU mit großen Batch-Größen ermöglicht. Bemerkenswerterweise erforderte diese Transformation nur 45B Trainings-Token im Vergleich zu über 15T Token, die für das 70B-Modell verwendet wurden, von dem es abgeleitet wurde. Dies etabliert ein neues Paradigma, bei dem leistungsstarke Modelle für eine effiziente Bereitstellung optimiert werden können, mit nur vernachlässigbaren Kompromissen bei ihren Fähigkeiten, und zeigt, dass die Inferenzleistung, nicht allein die Parameteranzahl, die Modellauswahl leiten sollte. Mit der Veröffentlichung von Nemotron-51B und der Präsentation des Puzzle-Frameworks bieten wir Praktikern sofortigen Zugang zu modernsten Sprachmodellierungsfähigkeiten zu deutlich reduzierten Rechenkosten.

English

Large language models (LLMs) have demonstrated remarkable capabilities, but their adoption is limited by high computational costs during inference. While increasing parameter counts enhances accuracy, it also widens the gap between state-of-the-art capabilities and practical deployability. We present Puzzle, a framework to accelerate LLM inference on specific hardware while preserving their capabilities. Through an innovative application of neural architecture search (NAS) at an unprecedented scale, Puzzle systematically optimizes models with tens of billions of parameters under hardware constraints. Our approach utilizes blockwise local knowledge distillation (BLD) for parallel architecture exploration and employs mixed-integer programming for precise constraint optimization. We demonstrate the real-world impact of our framework through Llama-3.1-Nemotron-51B-Instruct (Nemotron-51B), a publicly available model derived from Llama-3.1-70B-Instruct. Nemotron-51B achieves a 2.17x inference throughput speedup, fitting on a single NVIDIA H100 GPU while preserving 98.4% of the original model's capabilities. Nemotron-51B currently stands as the most accurate language model capable of inference on a single GPU with large batch sizes. Remarkably, this transformation required just 45B training tokens, compared to over 15T tokens used for the 70B model it was derived from. This establishes a new paradigm where powerful models can be optimized for efficient deployment with only negligible compromise of their capabilities, demonstrating that inference performance, not parameter count alone, should guide model selection. With the release of Nemotron-51B and the presentation of the Puzzle framework, we provide practitioners immediate access to state-of-the-art language modeling capabilities at significantly reduced computational costs.