Puzzle: Distillation-Based NAS for Inference-Optimized LLMs (2411.19146v5)

Abstract: 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.

• The paper introduces Puzzle, a framework using BLD and MIP for distillation-based NAS to optimize LLMs for inference efficiency.
• Puzzle optimized Llama-3.1-70B, yielding a model with 2.17x inference speedup on one H100 GPU and 98.4% of original capabilities.
• Puzzle provides a practical method to optimize LLMs for specific hardware constraints, significantly cutting inference costs and improving real-world deployability.

Insights on "Puzzle: Distillation-Based NAS for Inference-Optimized LLMs"

The paper "Puzzle: Distillation-Based NAS for Inference-Optimized LLMs" presents an innovative approach addressing the computational challenges faced during the deployment of LLMs. The authors introduce Puzzle, a framework that employs neural architecture search (NAS) combined with a distillation process to systematically optimize LLMs with tens of billions of parameters, specifically tailored for hardware constraints.

Problem Statement

LLMs have shown impressive prowess in various applications, yet their real-world adoption is hindered by substantial inference costs. These costs arise mainly from increased parameter counts, which improve accuracy but inflate computational requirements. The gap between state-of-the-art capabilities and practical deployability necessitates solutions like Puzzle, which aim to enhance inference efficiency while preserving the models' core functionalities.

Methodology

The Puzzle framework integrates NAS with an efficient distillation process. At the heart of this approach is blockwise local knowledge distillation (BLD) and mixed-integer programming (MIP). BLD aids in rapidly training model blocks by isolating their operations from the rest of the model, ensuring quick convergence with minimal data. The MIP is then used to navigate the vast search space of architectural configurations effectively, selecting the most suitable models under specific hardware and task constraints.

Key Results

The paper demonstrates the applicability of the Puzzle framework through the deployment of Llama-3.1-Nemotron-51B-Instruct, an optimized variant of Llama-3.1-70B-Instruct. This model showcases a 2.17x speedup in inference throughput, enabling it to operate on a single NVIDIA H100 GPU while retaining 98.4% of the original model's capabilities. It achieves this efficiency by reducing the training dataset requirements to just 45 billion tokens, a stark decrease from the 15 trillion tokens needed by the parent model.

Technical Contributions

1. Blockwise Local Distillation (BLD): This component facilitates parallel architecture exploration by training numerous block variants independently, significantly reducing computational overhead.
2. Mixed-Integer Programming (MIP): By framing the NAS problem as a MIP, this method allows efficient traversal of a massive search space, identifying configurations satisfying predefined inference constraints.
3. Enhanced TensorRT-LLM: The framework introduces modifications to TensorRT-LLM to support 'non-uniform' model architectures, accommodating variations in attention mechanisms across layers, a crucial development for deploying puzzle-derived architectures.

Implications

The introduction of Puzzle marks a significant step in optimizing LLM deployment. The framework's ability to drastically reduce inference costs without substantial accuracy losses makes it a practical tool for enhancing real-world applicability. Moreover, the provision of a detailed empirical analysis enhances our understanding of how architectural choices impact hardware efficiency, a crucial insight for designing future hardware-aware LLM architectures.

Future Directions

While the current implementation of Puzzle focuses on optimizing LLMs to fit specific hardware configurations, future enhancements could extend its range to support multimodal tasks or deploy on constraints like energy consumption. Additionally, exploring reinforcement learning or adaptive algorithms for NAS could yield even more efficient architectures and broader deployment capabilities.

In essence, Puzzle sets a precedence in the field of AI model optimization, pushing the boundaries of what can be achieved with limited computational resources while maintaining impressive model performance. This balance of efficiency and capability is pivotal in making advanced AI technologies more accessible and applicable across diverse contexts.