パズル：推論最適化LLM向けの蒸留ベースNAS

Puzzle: Distillation-Based NAS for Inference-Optimized LLMs

November 28, 2024

著者: 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

要旨

大規模言語モデル（LLMs）は驚異的な能力を示していますが、推論時の高い計算コストによってその採用が制限されています。パラメータ数を増やすことが精度を向上させる一方で、最先端の能力と実用的な展開可能性との間にギャップを広げることにもなります。本研究では、Puzzleというフレームワークを提案し、LLMの推論を特定のハードウェア上で加速させながらその能力を維持します。前例のないスケールでニューラルアーキテクチャサーチ（NAS）を革新的に適用することで、Puzzleは数百億のパラメータを持つモデルをハードウェアの制約下で系統的に最適化します。当アプローチは、ブロック単位のローカル知識蒸留（BLD）を並列アーキテクチャ探索に利用し、精密な制約最適化のために混合整数計画法を採用しています。 当フレームワークの実世界への影響を示すために、Llama-3.1-Nemotron-51B-Instruct（Nemotron-51B）というLlama-3.1-70B-Instructから派生した公開モデルを取り上げます。Nemotron-51Bは、オリジナルモデルの能力の98.4%を維持しながら、単一のNVIDIA H100 GPUに適合し、推論スループットを2.17倍高速化しています。Nemotron-51Bは、大規模なバッチサイズで推論が可能な最も正確な言語モデルとして現在の地位にあります。驚くべきことに、この変換にはわずか45Bのトレーニングトークンが必要であり、それは派生元の70Bモデルに使用された15Tのトークンと比較しています。これは、強力なモデルが、能力のほとんどを犠牲にすることなく効率的な展開のために最適化できる新たなパラダイムを確立し、パラメータ数だけでなく推論パフォーマンスがモデル選択を導くべきであることを示しています。Nemotron-51BのリリースとPuzzleフレームワークの発表により、実務家は著しく削減された計算コストで最先端の言語モデリング能力に直ちにアクセスできます。

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.