Puzzle: Distillation-Based NAS for Inference-Optimized LLMs

TL;DR

Puzzle is a hardware-aware neural architecture search framework that distills large language models to optimize inference speed on single GPUs while maintaining high accuracy, enabling more practical deployment of large models.

Contribution

The paper introduces Puzzle, a novel NAS framework utilizing blockwise local knowledge distillation and mixed-integer programming to optimize large language models for inference efficiency.

Findings

Achieves 2.17x inference speedup on Llama models on a single GPU.

Retains 98.4% of original model accuracy after optimization.

Supports large batch inference with high accuracy on limited hardware.

Abstract

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.