Puzzle: Distillation-Based NAS for Inference-Optimized LLMs
Fuente:
arXiv
Saved in:
| Main Authors: | , , , , , , , , , , , , , , , , , , , , , , , , , |
|---|---|
| Format: | Preprint |
| Published: |
2024
|
| Subjects: | |
| Online Access: | |
| Tags: |
Add Tag
No Tags, Be the first to tag this record!
|
| _version_ | 1866916773475909632 |
|---|---|
| author | Bercovich, Akhiad Ronen, Tomer Abramovich, Talor Ailon, Nir Assaf, Nave Dabbah, Mohammad Galil, Ido Geifman, Amnon Geifman, Yonatan Golan, Izhak Haber, Netanel Karpas, Ehud Koren, Roi Levy, Itay Molchanov, Pavlo Mor, Shahar Moshe, Zach Nabwani, Najeeb Puny, Omri Rubin, Ran Schen, Itamar Shahaf, Ido Tropp, Oren Argov, Omer Ullman Zilberstein, Ran El-Yaniv, Ran |
| author_facet | Bercovich, Akhiad Ronen, Tomer Abramovich, Talor Ailon, Nir Assaf, Nave Dabbah, Mohammad Galil, Ido Geifman, Amnon Geifman, Yonatan Golan, Izhak Haber, Netanel Karpas, Ehud Koren, Roi Levy, Itay Molchanov, Pavlo Mor, Shahar Moshe, Zach Nabwani, Najeeb Puny, Omri Rubin, Ran Schen, Itamar Shahaf, Ido Tropp, Oren Argov, Omer Ullman Zilberstein, Ran El-Yaniv, Ran |
| contents | Large language models (LLMs) offer remarkable capabilities, yet their high inference costs restrict wider adoption. While increasing parameter counts improves accuracy, it also broadens the gap between state-of-the-art capabilities and practical deployability. We present Puzzle, a hardware-aware framework that accelerates the inference of LLMs while preserving their capabilities. Using neural architecture search (NAS) at a large-scale, Puzzle optimizes models with tens of billions of parameters. Our approach utilizes blockwise local knowledge distillation (BLD) for parallel architecture exploration and employs mixed-integer programming for precise constraint optimization.
We showcase our framework's impact via Llama-3.1-Nemotron-51B-Instruct (Nemotron-51B) and Llama-3.3-Nemotron-49B, two publicly available models derived from Llama-70B-Instruct. Both models achieve a 2.17x inference throughput speedup, fitting on a single NVIDIA H100 GPU while retaining 98.4% of the original model's benchmark accuracies. These are the most accurate models supporting single H100 GPU inference with large batch sizes, despite training on 45B tokens at most, far fewer than the 15T used to train Llama-70B. Lastly, we show that lightweight alignment on these derived models allows them to surpass the parent model in specific capabilities. Our work establishes that powerful LLM models can be optimized for efficient deployment with only negligible loss in quality, underscoring that inference performance, not parameter count alone, should guide model selection. |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2411_19146 |
| institution | arXiv |
| publishDate | 2024 |
| record_format | arxiv |
| spellingShingle | Puzzle: Distillation-Based NAS for Inference-Optimized LLMs Bercovich, Akhiad Ronen, Tomer Abramovich, Talor Ailon, Nir Assaf, Nave Dabbah, Mohammad Galil, Ido Geifman, Amnon Geifman, Yonatan Golan, Izhak Haber, Netanel Karpas, Ehud Koren, Roi Levy, Itay Molchanov, Pavlo Mor, Shahar Moshe, Zach Nabwani, Najeeb Puny, Omri Rubin, Ran Schen, Itamar Shahaf, Ido Tropp, Oren Argov, Omer Ullman Zilberstein, Ran El-Yaniv, Ran Machine Learning Large language models (LLMs) offer remarkable capabilities, yet their high inference costs restrict wider adoption. While increasing parameter counts improves accuracy, it also broadens the gap between state-of-the-art capabilities and practical deployability. We present Puzzle, a hardware-aware framework that accelerates the inference of LLMs while preserving their capabilities. Using neural architecture search (NAS) at a large-scale, Puzzle optimizes models with tens of billions of parameters. Our approach utilizes blockwise local knowledge distillation (BLD) for parallel architecture exploration and employs mixed-integer programming for precise constraint optimization. We showcase our framework's impact via Llama-3.1-Nemotron-51B-Instruct (Nemotron-51B) and Llama-3.3-Nemotron-49B, two publicly available models derived from Llama-70B-Instruct. Both models achieve a 2.17x inference throughput speedup, fitting on a single NVIDIA H100 GPU while retaining 98.4% of the original model's benchmark accuracies. These are the most accurate models supporting single H100 GPU inference with large batch sizes, despite training on 45B tokens at most, far fewer than the 15T used to train Llama-70B. Lastly, we show that lightweight alignment on these derived models allows them to surpass the parent model in specific capabilities. Our work establishes that powerful LLM models can be optimized for efficient deployment with only negligible loss in quality, underscoring that inference performance, not parameter count alone, should guide model selection. |
| title | Puzzle: Distillation-Based NAS for Inference-Optimized LLMs |
| topic | Machine Learning |
| url | https://arxiv.org/abs/2411.19146 |