Rompecabezas: NAS basado en destilación para LLMs optimizados para inferencia

Puzzle: Distillation-Based NAS for Inference-Optimized LLMs

November 28, 2024

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

Resumen

Los modelos de lenguaje de gran tamaño (LLMs) han demostrado capacidades notables, pero su adopción se ve limitada por los altos costos computacionales durante la inferencia. Aunque aumentar el número de parámetros mejora la precisión, también amplía la brecha entre las capacidades de vanguardia y la desplegabilidad práctica. Presentamos Puzzle, un marco para acelerar la inferencia de LLM en hardware específico mientras se preservan sus capacidades. A través de una aplicación innovadora de búsqueda de arquitectura neuronal (NAS) a una escala sin precedentes, Puzzle optimiza sistemáticamente modelos con decenas de miles de millones de parámetros bajo restricciones de hardware. Nuestro enfoque utiliza destilación de conocimiento local por bloques (BLD) para exploración de arquitectura paralela y emplea programación entera mixta para una optimización precisa de las restricciones. Demostramos el impacto real de nuestro marco a través de Llama-3.1-Nemotron-51B-Instruct (Nemotron-51B), un modelo disponible públicamente derivado de Llama-3.1-70B-Instruct. Nemotron-51B logra una aceleración de 2.17 veces en el rendimiento de la inferencia, ajustándose en una sola GPU NVIDIA H100 mientras preserva el 98.4% de las capacidades del modelo original. Actualmente, Nemotron-51B es el modelo de lenguaje más preciso capaz de inferencia en una sola GPU con tamaños de lote grandes. De manera notable, esta transformación solo requirió 45B de tokens de entrenamiento, en comparación con más de 15T de tokens utilizados para el modelo de 70B del que se derivó. Esto establece un nuevo paradigma donde los modelos potentes pueden optimizarse para un despliegue eficiente con solo compromisos insignificantes en sus capacidades, demostrando que el rendimiento de la inferencia, no solo el número de parámetros, debería guiar la selección del modelo. Con el lanzamiento de Nemotron-51B y la presentación del marco Puzzle, brindamos a los profesionales acceso inmediato a capacidades de modelado de lenguaje de vanguardia a costos computacionales significativamente reducidos.

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.