arXiv:2411.19146cs.LG2024-11ICML被引 17

用知识蒸馏和搜索优化大模型,让超大模型跑得更快更省显存。

Puzzle: Distillation-Based NAS for Inference-Optimized LLMs

  • 用分块蒸馏与混合整数规划,高效搜索适合硬件的轻量模型结构。
  • 在单张H100上实现2.17倍加速,精度保留98.4%且支持大批次推理。
  • 仅用450亿训练数据就超越原版700亿参数模型部分能力,适合部署落地。

大语言模型(LLMs)虽能力强,但推理成本高,制约实际应用。我们提出Puzzle,一个面向硬件的框架,通过大规模神经架构搜索(NAS)优化数十亿参数的模型,实现高效推理。采用块级局部知识蒸馏(BLD)实现并行架构探索,并用混合整数规划进行精确约束优化。在两个基于Llama-70B-Instruct公开发布的模型——Nemotron-51B和Nemotron-49B上验证:两者在单张NVIDIA H100 GPU上实现2.17倍推理吞吐量提升,保持98.4%的原始基准准确率,是当前最准确、支持大批次单卡部署的超大模型。尽管训练仅用450亿词元,远少于原版15万亿词元,但经轻量对齐后,部分能力已超越父模型。证明高性能模型可高效部署,推理效率应优先于参数量选择。

原文摘要 · Abstract (English)

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.

大模型优化知识蒸馏硬件适配推理加速

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。