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

What problem does this paper attempt to address?

This paper attempts to solve the problem of high computational cost in the reasoning process of large - scale language models (LLMs). Although increasing the number of parameters can improve the accuracy of the model, it also widens the gap between the state - of - the - art capabilities and practical deployment at the same time. The paper proposes a framework named Puzzle, which aims to accelerate the LLM reasoning process on specific hardware while retaining the capabilities of the model. By innovatively applying the neural architecture search (NAS) technique, Puzzle systematically optimizes models with billions of parameters to adapt to hardware constraints. Specifically, the Puzzle framework achieves this goal through the following steps: 1. **Construct "puzzle pieces"**: Use block - level local distillation (BLD) to train each alternative sub - block in parallel and evaluate its quality and reasoning cost, thereby constructing a "block library". 2. **Assemble the puzzle architecture**: Utilize the mixed - integer programming (MIP) algorithm to assemble a heterogeneous architecture to optimize quality under given constraints such as throughput, latency, and memory usage. 3. **Retrain**: The reassembled architecture is trained through global knowledge distillation (GKD) to enhance the compatibility between blocks. The paper demonstrates the practical effect of the Puzzle framework through the Llama - 3.1 - Nemotron - 51B - Instruct (Nemotron - 51B) model. Nemotron - 51B is a model running on a single NVIDIA H100 GPU. Compared with the original model Llama - 3.1 - 70B - Instruct, it achieves a 2.17 - fold acceleration in reasoning throughput while retaining 98.4% of the original model capabilities. In addition, the training of Nemotron - 51B only requires 45 billion training tokens, far less than the 15 trillion tokens required by the original model. In conclusion, the Puzzle framework provides an efficient method to optimize the reasoning efficiency of large - scale language models without significantly sacrificing model performance, making it more suitable for practical applications.

Puzzle: Distillation-Based NAS for Inference-Optimized LLMs

Enabling High-Sparsity Foundational Llama Models with Efficient Pretraining and Deployment

Task Scheduling for Efficient Inference of Large Language Models on Single Moderate GPU Systems

Large Language Model Compression with Neural Architecture Search

LLM in a flash: Efficient Large Language Model Inference with Limited Memory

Distributed Inference and Fine-tuning of Large Language Models Over The Internet

LLaMA-NAS: Efficient Neural Architecture Search for Large Language Models

LLM-Inference-Bench: Inference Benchmarking of Large Language Models on AI Accelerators

Harnessing Your DRAM and SSD for Sustainable and Accessible LLM Inference with Mixed-Precision and Multi-level Caching

PowerInfer-2: Fast Large Language Model Inference on a Smartphone

Fast and Efficient 2-bit LLM Inference on GPU: 2/4/16-bit in a Weight Matrix with Asynchronous Dequantization

Inference Performance Optimization for Large Language Models on CPUs

Distributed Inference Performance Optimization for LLMs on CPUs

PUZZLE: Efficiently Aligning Large Language Models Through Light-Weight Context Switch.

SpikeLLM: Scaling up Spiking Neural Network to Large Language Models via Saliency-based Spiking

Petals: Collaborative Inference and Fine-tuning of Large Models

SqueezeLLM: Dense-and-Sparse Quantization

A Hardware Evaluation Framework for Large Language Model Inference

Accelerating LLaMA Inference by Enabling Intermediate Layer Decoding via Instruction Tuning with LITE

Hardware-Aware Parallel Prompt Decoding for Memory-Efficient Acceleration of LLM Inference