Puzzle: NAS basata sulla distillazione per LLM ottimizzati per l'inferenza

Puzzle: Distillation-Based NAS for Inference-Optimized LLMs

November 28, 2024

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

Abstract

I modelli linguistici di grandi dimensioni (LLM) hanno dimostrato capacità notevoli, ma la loro adozione è limitata dai costi computazionali elevati durante l'inferenza. Aumentare il numero di parametri migliora la precisione, ma allarga anche il divario tra le capacità all'avanguardia e la praticità della distribuzione. Presentiamo Puzzle, un framework per accelerare l'inferenza dei LLM su hardware specifici preservandone le capacità. Attraverso un'applicazione innovativa della ricerca dell'architettura neurale (NAS) su una scala senza precedenti, Puzzle ottimizza sistematicamente modelli con decine di miliardi di parametri sotto vincoli hardware. Il nostro approccio utilizza la distillazione della conoscenza locale a blocchi (BLD) per l'esplorazione dell'architettura parallela e impiega la programmazione mista intera per un'ottimizzazione precisa dei vincoli. Dimostriamo l'impatto concreto del nostro framework attraverso Llama-3.1-Nemotron-51B-Instruct (Nemotron-51B), un modello disponibile pubblicamente derivato da Llama-3.1-70B-Instruct. Nemotron-51B ottiene un aumento della velocità di elaborazione dell'inferenza del 2,17x, adattandosi a una singola GPU NVIDIA H100 preservando il 98,4% delle capacità del modello originale. Attualmente, Nemotron-51B è il modello linguistico più accurato in grado di effettuare inferenze su una singola GPU con batch di grandi dimensioni. Sorprendentemente, questa trasformazione ha richiesto solo 45B di token di addestramento, rispetto ai oltre 15T di token utilizzati per il modello da 70B da cui è derivato. Questo stabilisce un nuovo paradigma in cui modelli potenti possono essere ottimizzati per una distribuzione efficiente con solo compromessi trascurabili delle loro capacità, dimostrando che le prestazioni dell'inferenza, non solo il numero di parametri, dovrebbero guidare la selezione del modello. Con il rilascio di Nemotron-51B e la presentazione del framework Puzzle, offriamo agli operatori l'accesso immediato a capacità di modellazione linguistica all'avanguardia a costi computazionali significativamente ridotti.

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.