# Puzzle: Distillation-Based NAS for Inference-Optimized LLMs

arXiv: https://arxiv.org/abs/2411.19146 · HTML: https://mlsanity.com/paper/2411.19146/

## Key facts

- Cited by 17 papers in the map (data through 14 Oct 2026).
- Builds on 52 of its 65 references that are in the map.
- First posted on arXiv 28 Nov 2024.

## Builds on (most cited first, 30 of 52)

- [GPT-4 Technical Report](https://mlsanity.com/paper/2303.08774/) (2023)
- [LoRA: Low-Rank Adaptation of Large Language Models](https://mlsanity.com/paper/2106.09685/) (2021)
- [Chain-of-Thought Prompting Elicits Reasoning in Large Language Models](https://mlsanity.com/paper/2201.11903/) (2022)
- [Distilling the Knowledge in a Neural Network](https://mlsanity.com/paper/1503.02531/) (2015)
- [Training Verifiers to Solve Math Word Problems](https://mlsanity.com/paper/2110.14168/) (2021)
- [Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks](https://mlsanity.com/paper/2005.11401/) (2020)
- [Judging LLM-as-a-Judge with MT-Bench and Chatbot Arena](https://mlsanity.com/paper/2306.05685/) (2023)
- [Evaluating Large Language Models Trained on Code](https://mlsanity.com/paper/2107.03374/) (2021)
- [Measuring Massive Multitask Language Understanding](https://mlsanity.com/paper/2009.03300/) (2020)
- [Visual Instruction Tuning](https://mlsanity.com/paper/2304.08485/) (2023)
- [Scaling Laws for Neural Language Models](https://mlsanity.com/paper/2001.08361/) (2020)
- [Efficient Memory Management for Large Language Model Serving with PagedAttention](https://mlsanity.com/paper/2309.06180/) (2023)
- [Think you have Solved Question Answering? Try ARC, the AI2 Reasoning Challenge](https://mlsanity.com/paper/1803.05457/) (2018)
- [Mamba: Linear-Time Sequence Modeling with Selective State Spaces](https://mlsanity.com/paper/2312.00752/) (2023)
- [HellaSwag: Can a Machine Really Finish Your Sentence?](https://mlsanity.com/paper/1905.07830/) (2019)
- [Longformer: The Long-Document Transformer](https://mlsanity.com/paper/2004.05150/) (2020)
- [Tree of Thoughts: Deliberate Problem Solving with Large Language Models](https://mlsanity.com/paper/2305.10601/) (2023)
- [Training Compute-Optimal Large Language Models](https://mlsanity.com/paper/2203.15556/) (2022)
- [TruthfulQA: Measuring How Models Mimic Human Falsehoods](https://mlsanity.com/paper/2109.07958/) (2021)
- [Efficiently Modeling Long Sequences with Structured State Spaces](https://mlsanity.com/paper/2111.00396/) (2021)
- [Efficient Streaming Language Models with Attention Sinks](https://mlsanity.com/paper/2309.17453/) (2023)
- [Neural Architecture Search with Reinforcement Learning](https://mlsanity.com/paper/1611.01578/) (2016)
- [GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints](https://mlsanity.com/paper/2305.13245/) (2023)
- [RULER: What's the Real Context Size of Your Long-Context Language Models?](https://mlsanity.com/paper/2404.06654/) (2024)
- [The FineWeb Datasets: Decanting the Web for the Finest Text Data at Scale](https://mlsanity.com/paper/2406.17557/) (2024)
- [Reconciling modern machine learning practice and the bias-variance trade-off](https://mlsanity.com/paper/1812.11118/) (2018)
- [A Simple and Effective Pruning Approach for Large Language Models](https://mlsanity.com/paper/2306.11695/) (2023)
- [A Convergence Theory for Deep Learning via Over-Parameterization](https://mlsanity.com/paper/1811.03962/) (2018)
- [Fast Transformer Decoding: One Write-Head is All You Need](https://mlsanity.com/paper/1911.02150/) (2019)
- [Intrinsic Dimensionality Explains the Effectiveness of Language Model Fine-Tuning](https://mlsanity.com/paper/2012.13255/) (2020)

## Built on by (most cited first, 17 of 17)

- [OctoThinker: Mid-training Incentivizes Reinforcement Learning Scaling](https://mlsanity.com/paper/2506.20512/) (2025)
- [Nemotron-H: A Family of Accurate and Efficient Hybrid Mamba-Transformer Models](https://mlsanity.com/paper/2504.03624/) (2025)
- [Jet-Nemotron: Efficient Language Model with Post Neural Architecture Search](https://mlsanity.com/paper/2508.15884/) (2025)
- [Gumiho: A Hybrid Architecture to Prioritize Early Tokens in Speculative Decoding](https://mlsanity.com/paper/2503.10135/) (2025)
- [RAD: Redundancy-Aware Distillation for Hybrid Models via Self-Speculative Decoding](https://mlsanity.com/paper/2505.22135/) (2025)
- [Pangu Light: Weight Re-Initialization for Pruning and Accelerating LLMs](https://mlsanity.com/paper/2505.20155/) (2025)
- [Distill-then-Replace: Efficient Task-Specific Hybrid Attention Model Construction](https://mlsanity.com/paper/2601.11667/) (2026)
- [Hankel Singular Value Regularization for Highly Compressible State Space Models](https://mlsanity.com/paper/2510.22951/) (2025)
- [Attention Editing: A Versatile Framework for Cross-Architecture Attention Conversion](https://mlsanity.com/paper/2604.05688/) (2026)
- [FFN Fusion: Rethinking Sequential Computation in Large Language Models](https://mlsanity.com/paper/2503.18908/) (2025)
- [TraceNAS: Zero-shot LLM Pruning via Gradient Trace Correlation](https://mlsanity.com/paper/2602.02891/) (2026)
- [Nemotron-Labs-3-Puzzle-75B-A9B: Compressing Hybrid MoE LLMs](https://mlsanity.com/paper/2607.04371/) (2026)
- [HALO: Memory-Centric Heterogeneous Accelerator with 2.5D Integration for Low-Batch LLM Inference](https://mlsanity.com/paper/2510.02675/) (2025)
- [Scaling Laws for Task-Specific LLM Distillation](https://mlsanity.com/paper/2606.24747/) (2026)
- [MOSAIC: Adaptive Inter-layer Composition for Efficient Heterogeneous Vision-Language Models](https://mlsanity.com/paper/2607.09029/) (2026)
- [Quantization-Aware Healing: A Practical Recipe for Recovering Compressed, 4-Bit LLMs](https://mlsanity.com/paper/2608.20953/) (2026)
- [Optimal Model Activation Policies for Inference Networks of Large Language Models](https://mlsanity.com/paper/2609.15992/) (2026)

## Abstract (arXiv)

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.