Puzzle : NAS basée sur la distillation pour les LLM optimisés pour l'inférence

Puzzle: Distillation-Based NAS for Inference-Optimized LLMs

November 28, 2024

Auteurs: 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

Résumé

Les grands modèles de langage (LLM) ont démontré des capacités remarquables, mais leur adoption est limitée par des coûts computationnels élevés lors de l'inférence. Alors que l'augmentation du nombre de paramètres améliore la précision, elle élargit également l'écart entre les capacités de pointe et le déploiement pratique. Nous présentons Puzzle, un cadre pour accélérer l'inférence des LLM sur un matériel spécifique tout en préservant leurs capacités. À travers une application innovante de la recherche en architecture neuronale (NAS) à une échelle sans précédent, Puzzle optimise systématiquement des modèles avec des dizaines de milliards de paramètres sous contraintes matérielles. Notre approche utilise la distillation de connaissances locales par blocs (BLD) pour l'exploration d'architecture parallèle et utilise la programmation mixte entière pour une optimisation précise des contraintes. Nous démontrons l'impact réel de notre cadre à travers Llama-3.1-Nemotron-51B-Instruct (Nemotron-51B), un modèle disponible publiquement dérivé de Llama-3.1-70B-Instruct. Nemotron-51B atteint une accélération de 2,17 fois du débit d'inférence, s'adaptant à une seule GPU NVIDIA H100 tout en préservant 98,4% des capacités du modèle original. Nemotron-51B est actuellement le modèle de langage le plus précis capable d'inférence sur une seule GPU avec de grandes tailles de lots. De manière remarquable, cette transformation a nécessité seulement 45 milliards de jetons d'entraînement, comparé aux plus de 15 billions de jetons utilisés pour le modèle 70B dont il est dérivé. Cela établit un nouveau paradigme où des modèles puissants peuvent être optimisés pour un déploiement efficace avec seulement un compromis négligeable de leurs capacités, démontrant que la performance d'inférence, et non seulement le nombre de paramètres, devrait guider la sélection des modèles. Avec la sortie de Nemotron-51B et la présentation du cadre Puzzle, nous offrons aux praticiens un accès immédiat aux capacités de modélisation de langage de pointe à des coûts computationnels considérablement réduits.

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.