퍼즐: 추론 최적화 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

초록

대형 언어 모델 (LLM)은 놀라운 능력을 보여주었지만 추론 중에 높은 계산 비용으로 인해 채택이 제한되고 있습니다. 매개 변수 수를 증가시키면 정확도가 향상되지만 최첨단 능력과 실용적인 배포 가능성 사이의 격차도 커집니다. 우리는 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는 현재 단일 GPU에서 대형 배치 크기로 추론이 가능한 가장 정확한 언어 모델로서 성능을 발휘합니다. 놀랍게도, 이 변환에는 70B 모델로부터 파생된 모델에 사용된 15조 토큰 대비 단지 450억 훈련 토큰이 필요했습니다. 이는 강력한 모델이 능력의 거의 희생 없이 효율적인 배포를 위해 최적화될 수 있는 새로운 패러다임을 수립하며, 매개 변수 수만이 아닌 추론 성능이 모델 선택을 이끌어야 함을 입증합니다. 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.