Puzzle: Distillation-Based NAS for Inference-Optimized LLMs
Authors: Akhiad Bercovich
Tomer Ronen
Talor Abramovich
Nir Ailon
Nave Assaf
Mohammed Dabbah
Ido Galil
Amnon Geifman
Yonatan Geifman
Izhak Golan
Netanel Haber
Ehud Dov Karpas
Roi Koren
Itay Levy
Pavlo Molchanov
Shahar Mor
Zach Moshe
Najeeb Nabwani
Omri Puny
Ran Rubin
Itamar Schen
Ido Shahaf
Oren Tropp
Omer Ullman Argov
Ran Zilberstein
Ran El-Yaniv
Venue: ICML 2025 poster (2025)
Keywords: Large Language Models
Neural Architecture Search with Distillation
Architecture Designed for Efficient Inference
Blockwise local Distillation
Hardware-Aware Optimization with Distillation
Primary area: deep_learning->large_language_models
TL;DR: Puzzle uses large-scale NAS to optimize LLM inference for hardware, reducing costs and retaining nearly all performance. Demonstrated on an H100 GPU, achieving 2.17x throughput with minimal accuracy loss.

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.

Puzzle: DistillationBassed NAS for Inference-Optimizee LLMs
Akhiad Bercovich" * Tomer Ronen"
Talor Abramovich Nir Ailon Nave Assaf Mohammad Dabbah Ido Galil Netanel Haber
Amnon Geifman Yonatan Geifman Izhak Golan
Roi Koren Itay Levy Pavlo Molchanov Shahar Mor Zach Moshe Najeeb Nabwani Omri Puny Ran Rubin Itamar Schen
Ehud Karpas
Io Shahaf Oren Tropp Omer Ullman Argov Ran Zilberstein Ran El-Yaniy
Abstract
1. Introduction
With remarkable advancements in the capability and accu- racy of LLMs, these models are increasingly adopted in various domains, ranging from virtual assistants to sophis- ticated enterprise solutions. This adoption trend accom- panied by a growing appetite for larger and more powerful models, as evidenced by the industry's ' push toward increas- ingly larger-scale LLMs (OpenAI, 2023; Anil al., 2023; cla; Dubey al., 2024) and inference-tine compute scal- ing (GPT-o1, (Wei et al., 2022; Yao et al., 2023; Hardt & Sun, 2024)). However, the high computational costs associ- ated with these models and their projected future iterations particularly during inference - restrict their accessibility and scalability, thus presenting a significant challenge for widespread personal and commercial applications.
Large language models (LLMs) offer remark- able capabilities, yet their high inference costs restrict wider adoption. While increasing parame- counts improves accuracy, it also broadens the gap between state-of-the-an capabilities and prac- tical 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 archi- tecture exploration and employs mixed-integer programming for precise constraint optimization.
LLMs require substantial amount of parameters for their training process to converge easily and achieve better gen- eralization (Kaplan et al., 2020; Hoffmann et al., 2024; Allen-Zhu et al., 2019; Belkin et al., 2018). This over- parameterization not only facilitates optimization, but also provides greater capacity store knowledge and learn com- plex patterns across diverse tasks, explaining why larger models consistently demonstrate superior performance (Ka- plan et al., 2020; Hoffmann et al., 2024). However, once trained, many parameters and computations turn out to be redundant for inference, as evidenced by the success of com- putational efficiency techniques (He et al., 2024; Beltagy et al., 2020; Hu et al., 2022). Yet, LLM architectures re- main largely uniform, comprising repeated identical layers, with little consideration given balancing each block's computational cost against its contribution to overall model predictive performance a design choice primarily driven by training stability and case of scaling rather than inference efficiency. This work addresses how to transform a trained LLM from a structure suited for training into one optimized for efficient inference on specific hardware (such as H100), while preserving its accumulated knowledge and predictive performance. Given "parent model". approach ex- plores large search space of architecture configurations to identify efficient options tailored to meet specific hard- ware and task-related constraints. This exploration requires
We showcase our framework's impact via Llama- 3. 1--eemotron-511--Instuu (Nemotron-51B) and two publicly available models derived from 7a--ma-------nntteuee Both models achieve a 2.17x inference throughput speedup, fitting on single NVIDIA H100 GPU while retaining 98.4% of the original model's benchmark accuracies. These are the most ac- curate models supporting single H100 GPU infer- ence 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.
Llama-3.3-Nemotnee
Equal contribution Correspondence Akhiad Bercovich, Tomer Ronen, Ran El-Yaniv <(abercovich, tronen, relyaniv| (elyanivvenviille @nvidia. com?.
Proceedings of the 4nnd International Conference Machine Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Puzzle: Distillation- RistillationBBased NAS for Inference-Optimmize LLMs
0
Teacher
Step 1: Crafting the "puzzle pieces" Applying block-wise local distillation every altarativv subblock replacementin in panallal and scoring its quaiity and inference build "Ubrary" blacks.
Step 2: Assembling the puzzle architecture; Utiliding Minlege-rroggmmmnn to sesemble
Step Uptraining
The resesembled architecture is
hetoroganoouss architacture that optimizes quality constraints such as Uhroughput, latency and memory usage.
traiied with gloibal Knowlldgge- Distillation to strengthen inter- block compatibllit.
Figure 1. An overview of the three stages of our Puzzle framework.
method to reliably estimate the performance of each po- tential configuration, allowing us to identify models that balance efficiency and accuracy for deployment.
"open- closed-data" scenarios where training data of the parent model is not publicly available. This allows prac- titioners to take any aveilalbe model and tailor it to their specific hardware or use case. To demonstrate the effec- tiveness of our framework, we present Llama-3. 1.ama-3.1-Nemoote I-Nemotron- derived from the Llama-70B
In this paper, we introduce the Puzzle framework, summa- rized Figure 1, which pioneers a decomposed neural architecture search (NAS) strategy for LLMs explore and optimize a vast search space of possible model architectures for the target hardware. Inspired by methods in computer vision NAS, and LANA (Molchanov et al., 2022) especially, we define design spaces that include alternative attention network (FFN) layers of varying efficiency degrees, up to a complete layer skip in the extreme case. We then use a blockwise local distillation (BLD) (See Section 3) framework train all these block variants for all layers of the parent LLM in parallel. Next, we efficiently score each alternative replacement "puzzle piece" and search an enor- mous design space for the most accurate models, while adhering to a set of inference constraints (e.g.. memory size, latency, and throughput). This is done by using a Mixed- (MIP) algorithm. Lastly, the reassem- bled model is trained with Global Knowledge Distillation (GKD) (Hinton et al., 2015). Unlike traditional uniform transformer architectures, our NAS framework produces non non-uniform models with adapted computation allocation, optimizing each layer's 's expressivity based on the model's overall requirements to focus resources where they matter most. This leads significant gains in efficiency without compromising model expressivity. By focusing on parent models with SOTA performance, we derive child models pushing the efficient frontier (see Figure 5 and Table 5), e.g., models which provide the best accuracy per dollar.
51B-Instruct (Nemotron- parent model using Puzzle. Nemotron-51 breaks the effi- cient frontier of LLMs on single NVIDIA H100 GPU, es- tablishing a new state-of-the-art for commercial applications by achieving unmatched throughput and memory efficiency on this hardware. Interestingly, the Nemotron-51 result- ing architecture is unique and irregular, with many layers featuring reduced or skipped attention and FFN operations. This design enhances NVIDIA H100 utilization under FP8 quantization while preserving accuracy. We also introduce derivative of Llama-3.1 1--aa--.1-8B-Inntteeee Instrucc,,, which breaks the efficient frontier for its throughput slice. While Nemotron- 51B also leads within its parameter range, we argue that categorizing models solely by parameter count-such as 50B or 70B-1 inadequate for real-worle applications. In- stead, inference performance under specific hardware, in- ference engine, quantization levels, budget constraints, and usage profiles-such such varying sequence lengths and batch sizes-shoule guide model selection.
and
Conventional approaches to designing LLMs for inference, such as training from scratch or knowledge distillation, present significant challenges. Training a model from scratch is impractical for evaluating multiple configurations. Knowledge distillation, while generally faster due to guid- ance from teacher model, remains prohibitively costly when evaluating multiple candidates from a large search space. Puzzle circumvents these limitations by conduct- ing architecture search immediately after BLD, during the search stage (MIP, Stage 2 in Figure 1). These stages effi- ciently identify promising configurations for multiple levels (slices) of inference constraints, without requiring GKD for each candidate. The computationally intensive GKD process is reserved for the final stage, after the optimal ar- chitectures have been reassembled. This enables Puzzle to focus resources on refining a single optimized model for each slice while maintaining low overall costs.
Our framework offers several advantages. First, it enjoys ex- tremely low costs relative to training a model from scratch. For instance, the entire training- BLD before the MIP stage and GKD afterward rewuired 45B tokens to run on L.ama-3.1--08 Instructt (henceforth Llama-70B) com- pared to more than 15T tokens used to train the parent model. Even smaller budgets are possible-down to 4B while still preserving strong performance (see Ta- ble 6). Additionally, our method requires only the parent model's 's weights-no not its training data making it ideal for
2

---

Puzzle: Distillation-Ban NAS for Inference-Optimizzed LLMs
Our contributions: (1) We introduce Puzzle, framework that applies decomposed NAS to distill an L.M into hard- ware and inference scenario optimized models. Our frame- work flexibly supports optimization across ranges of multi- ple constraints, including throughput, latency, and memory usage. Our work pioneers the large-scale use of blockwise distillation and MIP-based architecture search for LLMs, successfully scaling these techniques to tens of billions of parameters while requiring only fraction of the original training compute. We are the first to present NAS-derived models at the scale of tens of billions of parameters. Puzzle is designed to be low-cost, enabling the efficient creation of multiple child models from a single parent LLM, each tai- lored to different hardware and inference requirements. This scalability makes it feasible to publish optimized variants for diverse use cases and deployment environments. A demo of Puzzle is available at: http:///tuuub.oo////llbbb/b/////
key-value (KV) heads (Shazeer, 2019; Ainslie al., 2023) and many more. Given the high computational and mone- tary cost associated with running LLMs, optimizing these models to eliminate inefficiencies becomes a critical goal. NAS methods are defined by the search space, the search strategy and evaluation metric of candidate architectures. NAS offers a systematic approach to exploring architectural changes that balances performance and resource constraints.
We define vast search space that encompasses different operations for each layer of the parent LLM model. A block is composed of smaller components called subblocks. While blocks are user-defined, in LLMs, a block typically refers to a transformer layer, with the subblocks being the attention module and the feed-forward network (FFN). For each trans- former layer i, the search space combines options for the attention subblock (denoted A. = 1a where is the number of possible attention subblocks) and FFN subblock (denoted F: = {fx li=1 for subblocks). The attention subblock options could include mechanisms such as stan- dard multi-head attention (MHA), grouped query attention (GQA) (Ainslie et al., 2023) with varying numbers of key- value heads, replacing the attention layer with a single linear layer, no-op (i.e., entirely skipping the subblock). The FFN options include full reduced intermediate dimen- sions (which could be obtained with pruning techniques), linear layers, or no-ops. The combined search space for each transformer layer (or parent block), represented as A, x Fi, captures all possible pairings of attention and FFN configu- rations. In this work, each parent transformer layer can be replaced by single corresponding child transformer layer (block), although theoretically, multiple parent layers could be grouped and replaced by different number of child blocks (i.c., P parent layers to L child blocks). Specifically: (1) Attention subblocks: For Ai, we consider different variants of GQA with varying numbers of key-value heads (8, 4, 2, and 1). We also include the option to replace the entire attention subblock with a single linear layer or skip it entirely using no-op operation. (2) FFN subblocks: For expansion
(2) Using Puzzle, we introduce Nemotron-51B and Nemotron-49B, both optimized for a single H100 GPU, thus setting a new benchmark for commercial applications. We show Nemotron-49B preserves its parent's s 128K con- text via short uptraining and benefits from RLHF align- nent-ooutperfrmming Nemotron-51 B on some tasks. Puz- zle's robustness is demonstrated throughout the paper by applying it numerous times with varied constraints, datasets, budgets and parent models.
(3) Our method focuses on optimization for real world sce- narios, running on actual inference engines and efficient quantization levels (FP8). We therefore enhance TensorRT- LLM to efficiently support non-uniform blocks and atten- tion mechanisms with varying numbers of key-value heads across layers. This results in models that can directly im- prove the costs and usability of running I.IM inference in practical applications (see Appendix C).
(4) We provide a comprehensive empirical analysis of the relationship between architectural choices and hardware efficiency, offering insights to guide the design of future hardware-aware LLM architectures.
F. we consider the full cetemeddiate-iimennse factor of the parent model, along with reduced dimensions of approximately: 87%, 75%, 50%, 25% 20% and 10% of the original intermediate size. Furthermore, linear layers and no-op options are also included.
2. Search Space
The motivation for applying NAS in our work lies in the re- dundancy found in many trained LLMs. Numerous studies have shown that a significant portion of computations and parameters in LLMs become redundant post-training, during inference (Sanyal et al., 2024; Hu et al., 2022: Aghajanyan et al., 2021). Prior studies tried to solve these issues with techniques such as pruning (Muralidharan et al., 2024; Men et al., 2024; Ma et al., 2023; Xia et al., 2024; 2022), remov- ing entire layers (He et al., 2024; Gromov et al., 2024), local attention methods (e.g., window and sink attention) (Belt- agy et al., 2020; Xiao et al., 2024), reducing the number of
These variants offer a tradeoff between memory efficiency, computational cost, and representational power, allowing for flexibility based on specific resource constraints and performance needs. For example, reducing the number of key-value heads (by using GQA with fewer heads) not only speeds up the attention computation, but also helps decrease memory usage by reducing the KV-cache size, which can be crucial for meeting memory constraints or enabling larger batch sizes for better throughput (as GPUs operate more efficiently with larger batches). Using linear layers or re-
3

---

Puzzle: Distillation-Basse NAS for Inference-Optimizzed LLMs
duced FFN dimensions can similarly lower computational requirements, whereas keeping full dimensions maintains higher representational power for better accuracy. Lastly, in this work, we require that all subblocks within a layer have the same input and output dimensions as their parent. However, a scheme with subblocks of varying embedding dimensions could be developed in future work.
Our approach decomposes the crafting process operate individual blocks rather than full child models, which drastically decreases the computational cost. Since gradi- ents do not need to propagate across blocks, each child block can be trained independently and in parallel to lo- cally mimic its corresponding parent block. Only the parent activations are transferred between layers, as illustrated in Figure 2. This local distillation approach offers several ad- vantages. First, because each child block relies solely on its corresponding parent block, activations and gradients are isolated from other child blocks. This independence enables training blocks separately, leveraging pipeline paral- lelism across multiple GPUs. Second, each child subblock is trained to mimic a relatively simple function-a single parent subblock-makinn and more stable than training an entire child model. This focused training facilitates faster convergence and allows higher learning rates compared to standard language mod- eling or GKD methods. Additionally, we find that this approach requires only a small dataset (approximately one billion tokens). Third, each child subblock benefits from high-quality outputs from its preceding parent subblock, rather than the lower-quality outputs typical in global model training, which further enhances convergence speed.
To illustrate the scale of the search space, consider Llama 3.1-70B (Dubey et al., 2024), a model with 80 transformer layers. For the specific instantiation of our framework pre- sented in this work, we defined each transformer layer to have 6 potential alternatives for the attention subblock and 9 alternatives for the FFN subblock, resulting in 54 possible configurations per layer. Consequently, the total number of potential child model architectures is 5- which is approx- imately 10838
the process considerably simpler
138 different architectures. This number greatly exceeds the estimated number of atoms in the observable universe (10$2). Given such an immense number of pos- sibilities, and considering the costs of partially training or even measuring the accuracy of a single LLM architecture, evaluating any representative subset of the search space is computationally infeasible. Therefore, designing a tradi- tional search strategy and evaluation scheme for NAS in such vast space is challenging. To address this, we devised an efficient decomposed local distillation and evaluation framework, and our decomposed search strategy (described in Section 3 and Section 4). These strategies allow feasible navigation of the search space to find configurations that bal- ance expressivity and efficiency with practical constraints like latency, memory usage, and throughput.
To optimize the performance of each child block, we feed parent activations into the current block and compute a normalized mean squared error (MSE) loss (Kurtic et al., MSE(op.o.)
2023). Specifically, we define the loss as L M)) where Op and 0g represent the outputs of the original parent block and the modified child block, respectively.
A primary limitation of BLD is that it does not ensure com- patibility between different blocks. This issue arises because each block is trained with inputs from the preceding parent blocks rather than from the outputs of its predecessor blocks within the child model. This prevents later blocks from adapting to the errors of earlier blocks, and may lead er- rors propagating and compounding through the child model. To mitigate this, we introduce GKD as a final training phase in our framework (see Section 4.3). Nonetheless, empirical results show that the BLD stage alone recovers much of the parent model's performance (see Table 16).
3. Blockwise Local Distillation
Parent
Child
100%
loSs
To ensure broad coverage of diverse data domains within limited training schedules, we curated a dataset mixture, termed Distillation Mix, for all our distillation training runs. This mixture includes source code repositories, Wikipedia articles, books, news websites, and several other domains. The dataset comprises 224 billion tokens collected from three public datasets: Web (Penedo et al., 2024), Dolma (Soldaini et al., 2024), and Buzz-V1.2 (Hive-Digital- Technologies). In our BLD experiments, we used billion training tokens. We discuss the effect of varying BLD train- ing lengths on downstream tasks in Appendix F.1.3.
Figure 2. BLD: blocks are trained in paralicl and independently.
To create capable "puzzle pieces"-a - set of trained block vari- ants forming a block library for architectural exploration-we need to set effective weights for each child block. Our method involves efficient and training-free initializations for these blocks (see Appendix A). While these initialization techniques are beneficial for low- low-budget experiments, perfor- mance can be significantly improved by distilling knowledge from each parent block to its corresponding child block.
4

---

Puzzle: Distillation- BBsedd NAS for Inference-Optimizzed LLMs
3.1. Building a Block L Library with Decoupled Blockwise Local Distillation
its quality relative to the parent block in that location. (4) Search architectures: use search algorithm construct "Puzzle architectures" that have the best estimated quality under specified runtime and memory constraints. While these steps mirror established NAS approaches, their ap- plication to LLMs presents unique challenges due to the massive scale of the models. Our innovations in scaling these techniques to multi- multi-billion palaio--rraaeee models are de- tailed in the following subsections.
The first stage of our decomposed NAS framework (further discussed in Section 4) is building "library" of trained blocks. In order to cover the entire search space defined in Section 2, we need to obtain trained weights for each atten- tion variant aj and each FFN variant fú in each transformer layer i. We consider two methods to train the block library: coupled BLD and decoupled BLD. For each transformer layer coupled BLD constructs each possible block variant [a,, fo]s and trains it to emulate the corresponding parent block [pparents f faarent i. Given the significantly higher compu- tational costs of LLMs compared to CV models, and noting the inherent structure of the transformer layer, we propose decoupled BLD to drastically reduce the cost of building a block library: training [aj, fgureet"))) ); and (sparemt) : fk]s sep- arately to emulate [Aparent, parent] while freezing the child subblock that is identical to the parent, and composing the trained subblocks into a full block Su]; after training.
4.1. Estimating Resource Requirements
Accurate estimation of computational costs is crucial for optimizing LLM architectures for real-world deployment. While theoretical metrics like FLOPs or parameter count are commonly used to approximate cost, they often fail to capture the complexities of hardware acceleration. The actual runtime of neural network operations depends on nu- merous hardware-specifie factors: the number of streaming multiprocessors (SMs), tensor cores, memory bandwidth, and I/O patterns all significantly impact performance. For instance, operations with fewer FLOPs might run slower in practice due to poor hardware utilization. This disparity between theoretical and actual performance makes direct measurement on target hardware essential for optimization.
Given m attention variants, FFN variants, and trans- former layers, coupled BLD requires training m vari- ants, while decoupled BLD requires only (m++)) variants, significantly speeding up library construction, visualized in Figure 3. This efficiency becomes especially critical with a large toolbox for A, and F. For instance, with m = n = 20 and 1 = 80 layers, decoupled BLD would require training 3,200 blocks compared to 32,000 for coupled BLD. Decou- pled BLD enables us to explore large search space and produce high quality models. In Appendix F.1.1 we present a technique to combine coupled BLD and decoupled BLD.
Memory requirements during inference comprise two dis- tinct components with different scaling behaviors. Parame- ter memory, while substantial, remains constant regardless of the input size. In contrast, the key-value cache memory scales linearly with both batch size and sequence length, often becoming the dominant factor in long-sequence sce- narios. For example, in a model with 32 attention heads and 128-dimensional head size, each token requires 8KB of KV-cache per layer using FP16 precision. For an 8K- sequence with batch size 64, this amounts 4GB per layer for the KV-cache, possibly exceeding parameter memory.
#Actention Options
NAttention Options
HPPN Options
AREN Options
Coupleed touning mmin ewary cootkoottttooootoooooooooooo
rogether Mock Binck scoree are ponooeeeedddeeee tioooo, just they arnin coucled travingg
Deosspried out tlrck owtry
then
LLM inference has two phases with different performance characteristics. The prefill phase processes the input with high parallelization, using fewer passes than decoding. While efficient, the computational cost is significant for long contexts. In contrast, the generation phase processes one query token time auto-regressively, requiring re- peated forward passes, and often employing paged attention to manage memory efficiently. This difference between phases makes it crucial to actually measure both prefill and generation runtime in the target scenarios (see Table 2).
Figure 3. Coupled BLD requires training [A.
x [.l variants per transformer layer, while decoupled BLD requires only [A.] [Fi variants per layer, significantly speeding up library construction.
4. Decomposed Search Algorithm for LLMs
Our decomposed NAS framework is similar to earlier de- composed NAS methods used in computer vision (CV) such as DNA (Li al.. 2019), DONNA (Moons et al., 2020), and LANA (Molchanov et al., 2022): (1) Build block library: construct a diverse block library using BLD (see Section 3.1). (2) Estimate block resource requirements: estimate the runtime and memory requirements of each block variant across different scenarios. (3) Score blocks: score each block variant across the network to estimate
Batch size plays a critical role hardware efficiency, espe- cially during the generation phase. During prefill, even small batches process substantial amounts of data due to sequence length, allowing efficient hardware utilization. However, generation with small batch sizes processes minimal data per layer while still performing all the 10 needed to load the layer parameters, causing severe hardware under under-utilzzation
5

---

Puzzle: Distillation- Basedd NAS for Inference-Optimizzed LLMs
Increasing the batch size creates larger tensors that better utilize the GPU, often significantly improving throughput.
tokens in a validation corpus. KL divergence is widely used in KD setups statistical distance measure but has not been explored for decomposed NAS scoring before. Our analysis in Appendix F.1.4 highlights its effectiveness.
Given these complexities, our approach measures resource requirements directly on target hardware across various see- narios. For each block variant, we collect prefill and genera- tion latencies across different sequence lengths and batch sizes at a chosen quantization level. These measurements define the constraints in our MIP optimization (Appendix B), enabling the search algorithm to find architectures optimized for actual deployment rather than in theory.
Mixed-integer Programming. We treat the process of choosing a block variant for each trans- former layer as a grouped Knapsack problem, imposing constraints on memory (parameters plus key-value cache), throughput (tokens generated per second), and latency (time per batch). Specifically, we employ mixed-integer pro- gramming (MIP) formulation that maximizes a global qual- ity score while satisfying these hardware-driven require- ments. This approach allows us to efficiently navigate enor- mous search spaces (on the order of 1013s possible archi- tectures) and find solutions tailored to the usage needs and target hardware. Optimization details, including the formal problem definition, are provided in Appendix B.
Search Algorithm:
4.2. Scoring Architecture Solutions
A key advantage of decomposed NAS is its ability to es- timate the quality of an assembled model based on met- rics gathered from its individual blocks. This capability allows search algorithms to explore an enormous search space very efficiently, as the quality of each candidate ar- chitecture encountered during the search can be estimated within less than a second, instead of having to actually re- alize the candidate model and calculate some measure that requires performing forward passes on the entire model such as validation accuracy. Traditional NAS methods in CV, such as those in (Zoph& & Le, 2016) and (Real et al., 2017), typically rely on full model evaluation, which is com- putationally prohibitive for L.L.ss, where both search spaces and inference costs are substantially larger.
4.3. Post- Inter-Block Uptraining
As mentioned in Section 3, our BLD step trains each child block on parent-produced distributions, not on the outputs of the preceding child block. Consequently, block compat- ibility may be suboptimal post-BLD. To address this, we perform a short end-to-end global knowledge distillation (GKD) phase, also called Knowledge Distillation, which the child (student) aligns with the parent (teacher) (Hinton et al., 2015; Muralidharan et al., 2024; Lu et al., 2022). In particular, minimize:
We score each block variant at each network location by measuring the impact of replacing only that specific block in the parent model. To do this, we construct a model identical to the parent but with a single block replaced by trained block from our library, then calculate a performance mea- sure on the entire model. For efficient I/O, when scoring multiple variants, we load onto the GPU only the blocks that differ from the previously evaluated model. We call scores. During architecture
N ps log 25
COKD =
+
(1)
(1) Cosine Similarity Loss
(2) KL Divergence Loss
where (1) measures how well the student matches the
these scores replace-I-blan search, we estimate the quality of a constructed architecture as the sum of its individual replace- -block scores (see Ap- pendix B). Note that scoring is not performed on candidate models during the search phase - we estimate the quality of each block only once, then use its replace- -lloc score to estimate the quality of any candidate that contains it.
teacher's representations, and (2) aligns output distributions.
We find that explicitly including language-modeling loss here is not beneficial, as it can degrade downstream tasks (see Appendix 1.3. The effectiveness of this GKD step further demonstrated in Appendix F.3.
5. Main Results
We consider several metrics potential pollace-l-bloce scores: (1) downstream accuracy: accuracy on downstream tasks such as MMLU (Hendrycks et al., 2021). While down- stream accuracy is popular in CV applications, it can be very costly to measure in LLMs, especially given the number of block variants requiring scoring. (2) LM loss: causal lan- guage modeling loss, defined as the average log- likelihooood of validation corpus under the model's next-token predic- tion distribution. (3) KL divergence: Bullback-Leible di- vergence between the next-token prediction distributions of the evaluated model and the parent model, averaged across
Using our Puzzle framework, we generated Nemotron-511 as a child derivative of the Llama-70B model. Nemotron- 51B achieves a significant improvement in inference effi- ciency while retaining nearly all the accuracy of its parent, demonstrating the effectiveness of our approach.
To evaluate Puzzle- derived child models like Nemotron-51B,, two performance metrics are of primary interest: 1) Accuracy Preservation: This measures how much of the parent model's accuracy is retained by the child. A high retention percentage indicates
Evaluating Model Performance:
6

---

Puzzle: Distillation-Ban NAS for Inference-Optimizzed LLMs
that Puzzle produces High-performing children. 2) Compu- tational Efficiency: This reflects the child model's ability to adhere to the constraints it was optimized for. In our case, the focus is on rhroughput, showing how we improved the model's suitability for reducing inference cost.
selected the optimal TP and batch size to maximize through- put per GPU. The inference engine handled this selection dynamically for each run. For example, Nemotron-51 IB achieved optimal throughput with TP=1 and batch size 256, while Llama-3.11700 performed best with TP=4 and batch size 384.
These metrics demonstrate the balance between model qual- ity and achieving optimization for deployment needs.
Table 2. Throughput comparison of Nemotron-51B and Llama- 70B across various scenarios. Throughput is measured in tokens per second per GPU (NVIDIA H100). TP# indicates the num- ber of GPUs used in tensor parallelism. Note: Results were ob- tained on NVIDIA H100 SXM GPUs with FP8 quantization for weights, activations and KV cache using TensorRT-LLM. Optimal tensor parallelism was used for each model. Input/output sequence lengths indicate the prefill (input) and decode (output) operations performed by the LLM. * TTP=1 is not the optimal configuration for Llama-3. 1-00BB and is included for gqual-resource comparison.
Accuracy comparison: Table compares the accuracy of Nemotron-51B with its parent across several benchmarks. On average, Nemotron-51 B retains 98.4% of its parent's s accuracy, even exceeding it on certain benchmarks such as MT-Bench and GSM8K Chat. This underscores the robust- ness of the Puzzle framework in maintaining model quality despite significant architectural modifications.
Accuracy comparison of Nemotron-511 with Llama-70B across several benchmarks. Accuracy preserved is the ratio of child to parent accuracy. *Chat prompt as defined in Adler ct al. (2024). versioo by CodeParrot.
Table 1.
Semurie
EmpurCorpet
Nemotron-518 (TP1)
Liamo M1 700 Nil--astradt (TP):
Speedup
Chatbor
1251:38
5478 (TP1)
2545 (TP)
2.07
Gremian
138/1034
6472
2015 (TT4)/ (TT1
Lony Test Generitien
12822045
CEPS)
2786 (TP4)
17h
Inferece iime compas
1384495
3855 (EP2)
1928 (TP4)
2048/126
653 (IPU
(TP4:/ (TPL)
1/2/2117
Suesa Test
2045/2048
3622 (TP2)
1336 (TP4)
1.96
Ncavtzen-588
Pracenal
Winegraede DSskggadi st.l, 2000)
15.05
$4.55
15.35
ABC Chaleage (Clakk ct 200h)
71.31
19.31
The tradeoff between accuracy and efficiency is key for model selection, impact- ing deployment costs. Nemotron-51 is designed to balance the two and push beyond the current efficient frontier. Be- cause throughput directly affects cost, Nemotron-511 pro- vides the best accuracy per dollar, as shown in Figure 5. To account for both knowledge and conversational capabilities, accuracy is measured as a weighted combination of MMLU and -TT-Bench scores: (MT-Bende) x 10 + MMLU) / 2.
Accuracy vs. throughput frontier:
Belliing (Tilbes 2014)
$6.11
85%
GBMB (Cobbe al. 30211
12.04
91.43
19.34
Diziesly (Lan 2022
20.35
$1.63
9134
Xi.im (lllukillaannnn NP1)
31.6
GOMBE Char
81.58
Intunt Banaaiao (e-2n Chard a. 2021
75.55
73.54
05.35
ME O2kek22eeggii ctal, 2023)
6.99
10067
We complement the above accuracy performance benchmarks with a human evaluation compar- ing the parent and child models. blind test comparison of the models was conducted on a general purpose test set that includes tasks such reasoning, longform text gen- eration, knowledge and more. Results show comparable performance between the models (Figure 4), strengthening the claim that the accuracy degradation is minimal. For more details about the evaluation see Appendix D.
Human evaluation:
86
Liama: 3,1- instruct
3.
Nemotron-S18
83
Midmal 63228
82
Netther good (23%)
Liama 3.1
-70B-
Instruct (22%)
78
78
Nemotron -518 (21%)
Llama-3.1-g0 Instruct
Both good (34%)
2000
4000
6030
8000
10000
12000
Thrnughput
Figure 4. In a blind test, human annotators rated Llama-70B and Nemotron-51B comparably.
Figure 5. Accuracy vs. Throughput performance of Nemotron-51B
compared to state-of-the-art models. Throughput measured on NVIDIA H100 GPUs with optimal TP setting, running in FPS on "text generation" scenario (see Table 2). The red line represents the efficient frontier, highlighting models with the best accuracy- to-throughput tradeoff. Accuracy=(MT-Bend x 10 + MMLU)/ /
Table 2 specifies the throughput performance of Nemotron-51 against its parent across diverse input-output sequence lengths. Nemotron-51B achieves speedups up to 2.17x, enabling larger workloads per GPU and making it highly efficient for deployment. For each model and hardware configuration, we automatically
Throughput comparison:
The Nemotron-51 5IB architecture achieves substantial compu- tational savings through strategic reduction of computation
7

---

Puzzle: Distillation- Based NAS for Inference-Optimizze LLMs
across many layers, as shown in Figure 6. Observing the figure it is evident that our framework discovered signifi-
Table 3. Comparison of L (parent) and
Nemotron-4y9 Base (child) on RULER for context lengths up to 128K.
cant optimization opportunities in both early layers (0-15) and later layers (45-70), with computed savings shown in
Context
Parent Average Score
Child Average Score
Accuracy Preserved (%)
green. Different regions exhibit distinct optimization pat-
1K
96,77
97.40
100.65
terns: early layers show balanced reduction in attention and FFN components, while later layers demonstrate more aggressive attention optimization. Notably, the framework identifies central region (layers 16-42) that maintains full computational capacity, suggesting these middle layers are critical for performance. This automatically discovered structure demonstrates that large efficiency gains are achiev- able through careful targeting of computational reduction, while maintaining computation where it matters most.
8K
96.46
96.59
100.13
16K
95.98
96.09
100.11
32K
94,70
94,30
99.57
87.39
128K
52.25
49.62
94.96
Table 4. The performance of Nemotron-49B-RR B-Bsse hefore and after alignment, even surpassing its parent certain henchmarks.
Model
MMLJU
MT--Bench
Arena Hard (1. et al., 2024)
Nemotma-408 Rose alignment)
80.98
Nemotron 19D Baze (before alignment)
81.03
8.77
65.8
Llama: 3.3 70B Instract (parent)
81.66
8.4
71.70
Abrrtcn
paeve)
Compact model: We applied Puzzle to produce a child derivative of Llama-3.1 15.1-8B-Inntruu optimized for through- put specifically on an RTX 4090 GPU. This model breaks the efficient frontier for its throughput range, demonstrating Puzzle's 's ability to deliver highly efficient architectures also on consumer-grade hardware, while preserving the balance between accuracy and performance. Table highlights this model's superior tradeoff in accuracy and efficiency.
23
Figure 6. The runtime of the attention and FFN subblocks of Nemotron-5IB, relative to the original Llama-70B subblocks.
5.1. Additional Evaluations and Puzzle derivatives
Table 5. Accuracy and throughput of our high-throughput child derivative of Llama-3. 1.nna-3..--B---- 1-8--Instruct, which achieves equivalent throughput Llama-3.2-33 --Istuuct and far better accuracy. Throughput is estimated via the sum of measured block runtimes on single NVIDIA RTX 4090 GPU, measured with an input-output sequence length of 1024 tokens each, the scenario for which this model was optimized. Accuracy = (MT-Bench x 10 + MMLU) / 2.
Long-Context Performance: We evaluated Nemotron-511 on subset of the RULER benchmark (Hsieh et al., 2024). Notably, although Nemotron-51B was trained only on se- quences up to 8K tokens, it retained over 96% of its parent's accuracy at 16K tokens, showing Puzzle's 's ability to preserve performance beyond the direct training length. Performance degraded beyond 64K tokens, suggesting that fine tuning on longer contexts could extend the effective context length.
Model
Throughput*
Accuracy
Ours (child)
5856
73.98
3.2ma-3.2-3B-Isstte
5737
70.36
3lamm-3.1-88--mee (purent)
3385
76.40
To demonstrate this extension, we produced Llama- Settootoon-9--Supeee--re
(henceforth Nemotron-49B- Base, the base version of Llama-3. J.Jama--.--Nemotrron-- (tteemoorroo-SSS--4993333 from -----....------tttttt with identical constraints to Nemotron-51 and uptrained on an additional 5B tokens at 64K context and 5B at 128K. Table 3 shows that Nemotron- 49B-Base maintains exceeds its parent's performance up to 16K tokens, retains over 98% at 64K, and remains above 94% at 128K. Full performance details are provided in Appendix E.
Beyond the aforementioned models, Puzzle was also used in the development of the following derivatives: (1) Puzzle was applied to Llama- 3. 3.1 1-40SBBInnstruct with constraints requiring a 1.5x latency speedup and compatibility with a single NVIDIA 8x H100 node (640 GB) or single B100 GPU (192 GB). The re- sulting model retained 99.5% of the parent model's accu- racy (averaged across MMLU, MT-Bench, MMLU-Pro, Hu- manEval, and Arena Hard). FFN Fusion (Bercovich et al., 2025) was subsequently applied to further improve latency, followed by continued pretraining-resultinng in Llama-3. Nemotron- UlltUllatt------- 253B-CPT, which serves as the hase for the publicly released Llama-3.1-Nemotre J--ee-ttoon------- 2 2, (2) In Blakeman al. (2025), variation of Puzzle referred to as
Additional Puzzle derivatives:
Alignment: We further aligned Nemotron-499-i Base by following the RLHF and instruction-tuning recipe of Wang al. (2024), using 25M tokens. Table compares results before and after alignment, showing that Puzzle derivatives can undergo based alignment to boost their accuracy.
1ttpps: htttts:://ugginng htthhhugggnggacceccc wttp://ggaac...ccciiiiiiaaaaaaae wwt.c///ivaaaa.laaaaaaaa..aaaaa. N-n--emtrrn--uueee 498-v1
Pntepa://huggingfaace
ama-3_
8

---

Puzzle: Distillation-aan NAS for Inference-Optimizzed LLMs
"miniPuzzle" was applied to Sanooroo--H-56BBBaaa under constraints targeting RTX 5090 deployment with 1M con- text window. The resulting 47B model retained 99. 99.4%% of the parent model's accuracy (averaged over MMLU, MML.UPro, GSM8K, and HellaSwag). The key distinction between miniPuzzle and Puzzle is that miniPuzzle enforces a homogeneous block design across all layers.
level supervision and indicating diminishing returns past moderate data sizes (Appendices F.1.2, F.1.3).
(3) Block Scoring Metrics: We evaluated various scoring metrics to guide architecture search. KL divergence gener- ally yields the best results across tasks, outperforming LM loss and task-specific metrics. However, in specialized set- tings, such as instruction tuning benchmarks, task gookggrddded scores may offer an advantage. These results emphasize the importance of choosing scoring functions aligned with target deployment scenarios (Appendix F.1.4).
The GKD phase described ear- lier utilized 45B tokens, which might exceed practical bud- gets for some users. However, in practice, our method can achieve substantial accuracy recovery with significantly fewer tokens. Our experiments, summarized in Table 6, demonstrate that notable accuracy recovery can be real- ized with reduced token usage. After only 3.7B tokens of GKD, Nemotron-51B recovered 98.8% of its parent's ac- curacy MMLU and MT MT-Beecch benchmarks. Similarly, Nemotron-49B regained 99.63% of its parent's accuracy after only 8.68B tokens, and even 98.47% after just 2.9B tokens. These results indicate that while 45B tokens are modest compared to full model training from scratch, users can flexibly adjust GKD token counts based on available resources without severely compromising accuracy.
Training Token Budget:
(4) Search Space Diversity: We examined the impact of restricting the search space. Replacing all layer blocks with their original counterparts or limiting substitutions to no-op variants significantly degrades performance. This confirms that diversity in architectural choices critical for discover- ing high-quality configurations and validates the necessity of block-level flexibility in Puzzle (Appendix F.1.5).
(5) MIP Global Search vs. Greedy and Random Base- lines: We compared MIP-based global search with a simpler greedy algorithm that selects blocks layer-wise under fixed budgets. Despite meeting the same throughput constraints, greedy search yields significantly worse architectures, high- lighting the need for joint optimization across layers. Ad- ditional baselines, such as random- iemoom-from-iiie from-library selection, show that naive strategies even using trained blocks fall well short of MIP (Appendix F.2.2, F2.44.
Table Effectiveness of GKD training with reduced token bud-
gets. The table shows the performance recovery relative the parent model for Nemotron-5 SIB and Nemotron-49B using various
amounts of tokens GKD.
Mudd
G&D Tokano (8)
Perfrmaaae: OMMLU MMI-Baddo
Accuncce Proseval (96
Semmame 40t8-00 Bas
MUZ2 84023
99-17
6. Conclusions and Future Directions
Nemwonce 450 Bas
0.
$0.588.9 9.99
91.99
In this work we present Puzzle, a framework that modifies LLMs from their wer-parameterrize training configurations to optimized, inference-efficient architectures tailored for specific hardware. Puzzle achieves these improvements with remarkable efficiency in training resources. Requiring fewer than 50B tokens- to the trillions needed to train models from scratch-Puzzie produces models at a fraction of the usual cost. This extreme search and training efficiency still results in drastic inference per- formance improvements of Puzzle optimized models.
5.2. Ablation Studies Highlights:
We performed a series of ablations to examine how local distillation, block search strategies, data composition, and search space design each contribute to Puzzle's 's accuracy- efficiency tradeoffs. See Appendix F for full details. Below we summarize some of the most interesting findings:
(1) Decoupled BLD: A decoupled construction of the block library-training block alternatives indepen- dentll--signfficantly reduces cost by converting the search space from multiplicative to additive. This enables the con- struction of a large block library at fraction of the compute cost. Interestingly, adding a coupled training pass over a narrowed subspace of promising combinations provides an additional performance boost, suggesting hybrid BLD strategies can be beneficial (Appendix F.1.
The success of Puzzle opens several promising directions for future research. Our introduction of decoupled BLD, which is significantly more efficient than coupled BLD, makes it feasible evaluate a much larger set of potential blocks within a single optimization run. This efficiency enables exploration of novel operations as alternative Puzzle blocks, such as variable window attention mechanisms (Beltagy et al., 2020), state-space models (Gu et al., 2022; Gu & Dao, 2024), or other architectural innovations. The framework could also be extended to optimize models for specific capa- bilities, such as Thain-of-Thouge reasoning or multimodal tasks, including vision-language models (Liu et al., 2023) and eetrieval-augmente generation (Lewis et al., 2020).
(2) Data Composition and Scale:
We studied the effects of dataset scale and composition on BLD quality. Accuracy improves with more tokens, but saturates beyond 0.5B to- kens. Moreover, comain-diverse data improves robustness; however, training on even a narrow domain retains over 93% of full performance, highlighting the robustness of block-
9

---

Puzzle: Distillation- Based NAS for Inference-Optimizze LLMs
Impact Statement
ral Language Processing (Volume 1: Long Papers), pp. 7319-7328, Online, August 2021. Association for Com- putational Linguistics. doi: 10. 10.18653/ 568. URL https:/ httpps////aaaaachclg.....oo orgp://ooooo..g/2/20/ ac - 11w.1oo..5000 5588.
The Puzzle framework's ability to adapt state-of-the-art models hardware constraints, while maintaining high performance, can help make powerful AI systems more af- fordable and broadly deployable across diverse use cases. This broader adoption could expand the societal benefits of generative AI but also amplify its inherent risks (Solaiman et al., 2023; Capraro et al., 2024; Baldassarre et al., 2023). At the same time, by reducing energy consumption and hardware requirements needed for inference, Puzzle may alleviate some of the resource-intensive challenges of large- scale AI deployments, contributing to more sustainable approach to generative AI.
acl-
Ainslie, J., Lee-Thorp. J., de Jong, M., Zemlyanskiy, Y., Lebrón, F., and Sanghai, S. GQA: training generalized multi-query transformer models from multi-head check- points. In Bouamor, H., Pino, and Bali, K. (eds.), Pro- ceedings of the 2023 Conference on Empirical Methods in Natural Language Processing, EMNLP 2023, Singa- pore, December 6-10, 2023, pp. 189954900 Association for Computational Linguistics, 2023. doi: 10.18653/V1/ 2023.EMNLP-MMIN.D2. URL https://://do org/ 10. 10.18653/v1/200
Acknowledgments
We thank Saurav Muralidharan and Sharath Turuvekere Sreenivas for fruitful discussions.
Allen-Zhu, Z., Li, Y., and Song, Z. A convergence the- ory for deep learning via over wwrrpaaraetteiizatinn In Chaudhuri, K. and Salakhutdinov, R. (eds.), Proceed- ings of the 36th International Conference on Ma- chine Learning. ICML 2019, 9-15 June 2019, Long Beach, Califorria, USA, volume 97 of Proceedings of Machine Learning Research, pp. 242-252 PMLR, 2019, URL http: httt::::/pproceeeee http//pooeednnggeee mlr.m.l..... ppress/ v97 /al/aallnn-22e zhunn--huul3 9a. html.
References
The claude model family: Opus, sonnet, haiku. URL https:/ /appi
http:///ap...eettisssssccmc semanticschol org CorpusIDD:262234999
Adler, B., Agarwal, N., Aithal, A., Anh, D. H., Bhattacharya, P., Brundyn, A., Casper, J., Catanzaro, B., Clay, S., Co- hen, J. M.. Das, S., Dattagupta, A., Delalleau, O., Der- czynski, L., Dong, Y., Egert, D., Evans, E., Ficek, A., Fridman, D., Ghosh, S., Ginsburg, B., Gitman, I., Grze- gorzek, T., Hero, R., Huang, J., Jawa, V., Jennings, J., Jhunjhunwala, A., Kamalu, J., Khan, S., Kuchaiev, O., LeGresley, P., Li, H., Liu, J., Liu, Z., Long, E., Maha- baleshwarkar, A. S., Majumdar, S., Maki, J., Martinez, M., de Melo, M. R., Moshkov, I., Narayanan, D., Naren- thiran, S., Navarro, J., Nguyen, P., Nitski, O., Noroozi, V., Nutheti, G., Parisien, C., Parmar, J., Patwary, M., Pawelec, K., Ping, W., Prabhumoye, S., Roy, R., Saar, T., Sabavat, V. R. N., Satheesh, S., Scowcroft, J. P., Sewall, J., Shamis, P., Shen, G., Shocybi, M., Sizer, D., Smelyanskiy, M., Soares, F., Sreedhar, M. N., Su, D., Subramanian, S., Sun, S., Toshniwal, S., Wang, H., Wang, Z., You, J., Zeng, J., Zhang, J., Zhang, J., Zhang, V., Zhang. Y., and Zhu, C. Nemotron-4 340b technical report. CoRR, abs/2406. 1 11744, 2024. doi: 10.48500ARXX1.246611100 URL https:/ //doi. org/10. 48550/ www. 2406. 11704.
Anil, R., Borgeaud, S., Wu, Y., Alayrac, J., Yu, J.. Sori- cut, R., Schalkwyk, J., Dai, A. M., Hauth, A., Millican, K., Silver, D., Petrov, S., Johnson, M., Antonoglou, I., Schritwieser, J., Glaese, A., Chen, J., Pitler. E., Lilli- crap, T. P., Lazaridou, A., Firat, O., Molloy, J., Isard, M., Barham, P. R., Hennigan, T., Lee, B., Viola, F., Reynolds, M., Xu, Y., Doherty, R., Collins, E., Meyer, C., Ruther- ford, E., Moreira, E., Ayoub, K., Goel, M., Tucker, G., Piqueras, E., Krikun, M., Barr, I., Savinov, N., Danihelka, I., Roelofs, B., White, A., Andreassen, A., von Glehn, T., Yagati, I.., Kazemi, M., Gonzalez, I.., Khalman, M., Sygnowski, J., and al. Gemini: family of highly capable multimodal models. CoRR, abs/2312. 11805, 2023. doi: 10. 10.48550/RXIV.2312.111 18805 URL https: 2312 118.1.....11..........1111111111
//doi.org/// 11..rg/110.0
Baldassarre, M. T., Caivano, D., Nieto, B. F., Gigante, D., and Ragone, A. The social impact of generative Al: analysis chatgpt. In Proceedings of the 2023 ACM Conference on Information Technology for So- cial Good, GoodIT 2023, Lisbon, Portugal, Septem- ber 6-8, 2023, pp. 363-373. ACM, 2023. doi: 10. 0155335825115360955555 URL httpp::/ddat .org/ 10. ND.11145/358255 3609555..
Aghajanyan, A., Gupta, S., and Zettlemoyer, L. Intrin- sic dimensionality explains the effectiveness of language model fine-tuning. In Zong, C., Xia, F., Li, W., and Navigli, R. (eds.), Proceedings of the 59th Annual Meet- ing of the Association for Computational Linguistics and the I 1tth International Joint Conference on Natu-
Belkin, M., Hsu, D. J., Ma, S., and Mandal, S. Rec- onciling modern machine-learning practice and the classical Plas-variane trade-off. Proceedings of the National Academy of Sciences, 116:15849 15854,
10

---

Puzzle: Distillation- Based NAS for Inference-Optimizze LLMs
2018. URL https: httpss//aai..smmmttccchhh... org/CorpusTD::188 198496504.
Chen, M., Tworek, J., Jun, H., Yuan, Q. de Oliveira Pinto, H. P., Kaplan, J., Edwards, H., Burda, Y., Joseph, N., Brockman, G., Ray, A., Puri, R., Krueger, G., Petrov, M., Khlaaf, H., Sastry, G., Mishkin, P., Chan, B., Gray, S., Ryder, N., Pavlov, M., Power, A., Kaiser, I.,, Bavar- ian, M., Winter, C., Tillet, P., Such, F. P., Cummings, D., Plappert, M., Chantzis, F,, Barnes, E., Herbert- Voss, A., Guss, W. H., Nichol, A., Paino, A., Tezak, N., Tang, J., Babuschkin, I., Balaji, S., Jain, S., Saun- ders, W. Hesse, C., Carr, A. N., Leike, J., Achiam, J., Misra, V., Morikawa, E., Radford, A., Knight, M., Brundage, M., Murati, M., Mayer, K., Welinder, P., Mc- Grew, B., Amodei, D., McCandlish, S., Sutskever, I., and Zaremba, W. Evaluating large language models trained on code. CoRR, abs/2107.0333 033744, 2021. URL https: http///axvvvv... .rg. www.aiiww....aaeeeeeeeee
Beltagy. I., Peters, M. E., and Cohan, A. Longformer: The long-document transformer. CoRR, abs/2004.05150, 2020. URL https: / http/arxiiiiv.//// wttp:///rggg/////////// htttss//arxiiv.oo///a///000 05150.
Bercovich, A., Dabbah, M., Puny, O., Galil, I., Geif- man, A., Geifman, Y., Golan, I., Karpas, E., Levy, I,, Moshe, Z., Nabwani, N., Ronen, T., Schen, I., Se- gal, E., Shahaf, I., Tropp. O., Zilberstein, R., and EI- Yaniv, R. FFN fusion: Rethinking sequential computa- tion in large language models. CoRR, abs/2503. 18908, 2025. doi: 10.48550/ARXIV22500111 //doi org.i.org/a /tc.oo/ggoooo.ooooooooo 485444555////
URL https: 18908.
033744.
Blakeman, A., Basant, A., Khattar, A., Renduchintala, A.,
Bercovich, A., Ficek, A., Bjorlin, A., Taghibakhshi, A., Deshmukh, A. S., Mahabaleshwarkar, A. S., Tao, A., Shors, A., Aithal, A., Poojary, A., Dattagupta, A., Bud- dharaju, B., Chen, B., Ginsburg, B., Wang, B., Norick, B., Butterfield, B., Catanzaro, B., del Mundo, C., Dong, C., Harvey, C., Parisien, C., Su, D., Korzekwa, D., Yin, D., Gitman, D., Mosallanezhad, D., Narayanan, D., Frid- man, D., Rekesh, D., Ma, D., Pykhtar, D., Ahn, D., Riach, D., Stosic, D., Long, F.,, Segal, E., Evans, E., Chung, E., Galinkin, E., Bakhturina, E., Dobrowolska, E., Jia, F., Liu, F., Prasad, G., Shen, G., Liu, G., Chen, G. Qian, H., Ngo. H., Liu, H., Li, H., Gitman, I., Kar- manov, I., Moshkov, I., Golan, I., Kautz, J., Scowcroft, J. P. Casper, J., Seppänen, J., Lu, J., Sewall, J., Zeng, J., You, J., Zhang, J., Zhang, J., Huang, J., Xue, J., Huang, J. Conway, J., Kamalu, J., Barker, J., Cohen, J. M., Jen- nings, J., Parmar, J., Sapra, K., Briski, K. Chumachenko, K., Luna, K., Santhanam, K., Kong, K., Sivamani, K., Pawelec, K., Anik, K., Li, K., McAfee, L., Derczynski, L., Pavao, L., Vega, L.. Voegtle, L., Bala, M., de Melo, M. R., Sreedhar, M. N., Chochowski, M., and Kliegl, M. Nemotron-h: A family of accurate and efficient hy- models. CoRR, abs/2504.03624
Clark, P., Cowhey, I., Etzioni, O., Khot, T., Sabharwal, A.. Schoenick, C., and Tafjord, O. Think you have solved question answering? try arc, the A12 reasoning challenge. CoRR, abss/180.00557, 2018. URL http:://axxx/ org/ absabbss/a /1003.10.000 05457.
Cobbe, K., Kosaraju, V., Bavarian, M., Chen, M., Jun, H., Kaiser, L., Plappert, M.. Tworek, J., Hilton, J., Nakano, R., Hesse, C., and Schulman, J. Training verifiers to solve math word problems. CoRR, wwwp/111...41... 2021. URL https: / arxiv. wwwpa/s..orr//ooooeoao.o.ooo..sa /aaps////1110 2110. 14168.
Dubey, A., Jauhri, A., Pandey, A., Kadian, A., Al-Dahle, A., Letman, A., Mathur, A., Schelten, A., Yang. A., Fan, A., Goyal, A., Hartshorn, A., Yang, A., Mitra, A., Sravanku- mar, A., Korenev, A., Hinsvark, A., Rao, A., Zhang, A., Rodriguez, A., Gregerson, A., Spataru, A., Rozière, B., Biron, B., Tang, B., Chern B., Caucheteux, C., Nayak, C., Bi, C., Marra, C., McConelll, C., Keller, C., Touret, C., Wu, C., Wong, C., Ferrer, C. C., Nikolaidis, C.,, Allonsius, D., Song. D., Pintz, D., Livshits, D.,, Esiobu, D., Choud- hary. D., Mahajan, D., Starcia-Olane D., Perino, D., Hup- kes, D Lakomkin, E., AlBadawy, E.,, Lobanova, E., Di- nan, E., Smith, E. M., Radenovic, F., Zhang, F., Synnaeve, G., Lee, G., Anderson, G. L., Nail, G., Mialon, G., Pang, G., Cucurell, G., Nguyen, H., Korevaar, H., Xu, H., Tou- vron, H., Zarov, I., Ibarra, I. A., Kloumann, I. M., Misra, I., Evtimov, I., Copet, J., Lee, J., Geffert, J., Vranes, J., Park, J., Mahadeokar, J., Shah, J., van der Linde, J., Bil- lock, J., Hong. J., Lee, J., Fu, J., Chi, J., Huang, J., Liu, J., Wang, J., Yu, J., Bitton, J., Spisak, J., Park, J., Rocca, J., Johnstun, J., Saxe, J., Jia. J., Alwala, K. V., Upasani, K., Plawiak, K., Li, K., Hcafield, K., Stone, K., and et al. The llama 3 herd of models. CORR, abs/2407.21788. 2024. doi: 10.48550/ARXXI..440..117 / t/dpi. 000.0..88555000 48550885555111 2407 2407.21783.
brid mamba-
2025. doi: 10.48550/ARXIV.2500..6555 // //doi..co www.o/gg/////// 10. wtlp://....855///
URL https: 2504. ..244.33662 03624.
Capraro, V. Lentsch, A., Acemoglu, D., Akgün, S., Akhme- dova, A., Bilancini, E.,, Bonnefon, J., Brañas-Gazza P., Butera, L., Douglas, K. M., Everett, J. A. C., Gigerenzer, G., Greenhow, C., Hashimoto, D. A., Holt-Lunstad, J.,
Jetten, J., Johnson, S., Longoni, C., Lunn, P., Natale, S.,
Rahwan, I., Selwyn, N., Singh, V. Suri, S., Sutcliffe, J.,
Tomlinson, J., van der Linden, S., Lange, P. A. M. V. Wall, F., Bavel, J. J. V., and Viale, R. The impact of generative artificial intelligence on socioeconomic in- equalities and policy making. CORR, 250/2401.05377, URL ht https:
URL https:
2024. doi: 10.48550/ARXI.22410.0333 //doi. .rooi.orge weooi.org/1000 http:/10..455/////
Gromov, A., Tirumala, K., Shapourian, H., Glorioso, P., and Roberts, D. A. The unreasonable ineffectiveness of

---

Puzzle: Distillation-Basee NAS for Inference-Optimizzed LLMs
the deeper layers. CoRR, abs/2403. 17887, 2024, doi:
L. A., Welbl, J., Clark, A., Hennigan, T., Noland, E., Millican, K., den Driessche, G., Damoc, B., Guy, A., Osindero, S., Simonyan, K., Elsen, E., Vinyals, O., Rae, J. W., and Sifre, L. Training compute-optimat large lan- guage models. In Proceedings of the 36th International Conference on Neural Information Processing Systems, NIPS '22, Red Hook, NY, USA, 2024. Curran Associates Inc. ISBN 9781713871088.
10.4850/ARXXV.2400.7755 URL https: //ttp:://od... org/10. 08550/55055500 240.:/.....0...888..
Gu, A. and Dao, T. Mamba: Linear-time sequence modeling with selective state spaces. In First Conference on Lan- guage Modeling, 2024. URL https: httttt:://pppeeeaaaai http///pppeeeeeeeeeeeeeeee.eeeee wtt.::iiittttsttttx.
Gu, A., Goel, K., and Re, C. Efficiently modeling long sequences with structured state spaces. In In- ternational Conference on Learning Representations, 2022. URL htttt:://ppereeeiiiee id-uYLFozIvlAC.
Hsich, C., Sun, S., Kriman, S., Acharya, S., Rekesh, D., Jia, F., Zhang, Y., and Ginsburg, B. RULER: what's the real context size of your long-context language models? CoRR, abs/2404.06654, 2024. doi: 10.48550/ARXIV. 2404.06654. URL https: http//oi..og/o/////// 48550/ arXiv. 2404 2404. 06654.
met/forure
Hardt, M. and Sun, Y. Test-time training on nearest neigh- bors for large language models. In The Twelfth Interna- tional Conference on Learning Representations, ICLR 2024, Vienna, Austria, May 7-11, 2024. OpenReview.net, 2024. URL ftttps://operreviieet wet.iiewwweeeeeeeeeeeeeeeeeeeeee wwt.torunffrtto i dd=oNN2bbk44r.
Hu, E. J., Shen, Y., Wallis, P., Allen-Zhu, Z., Li, Y., Wang, S., Wang, L., and Chen, W. Lora: Low-rank adapta- tion of large language models. In The Tenth Interna- tional Conference
Learning Representations, ICLR 2022, Virtual Event, April 25-29, 2022. OpenReview.net, 2022. URL https:/tt////. http:////ppeeeeeeeeeeeeee htwpeeeww.eeeooooooeoeeeeeeee forum? id=nZeVKeeFYf9,
Hasan, T., Bhattacharjee, A., Islam, M. S., Mubasshir, K. S., Li, Y., Kang, Y., Rahman, M. S., and Shahriyar, R. XI- sum: Large-scale multilingual abstractive summariza- tion for 44 languages. In Zong, C., Xia, F., Li, W., and Navigli, R. (eds.), Findings of the Association for Computational Linguisties: ACL/ICNLP 2021, Online Event, August 1-6, 2021, volume ACL/ICCLPP 2021 of Findings of ACL, pp. 4693-4703. Association for Com- putational Linguistics, 2021. doi: 10.18653/V1/2021.
Inc., C.-O. F. Python-mip: collection of python tools for the modeling and solution of mixed-integer lin- ear programs, 2023. URL Noin-or/pythone
com/
Kaplan, J., McCandlish, S., Henighan, T., Brown, T. B., Chess, B., Child, R., Gray, S., Radford, A., Wu, J., and Amodei, D. Scaling laws for neural language models. CoRR, 0ttx/2001.086//1 2020. https:/ //axxiv. org/ /rg/abs////eeooccoeo /2t...001..00 htse/abb2200...o 08361.
FINDINGS-ACLA11. URL https: wordoi.ooog/..... 18653/v1/202. fei.iiiininnnnn acl. 413.
He, S., Sun, G., Shen, Z., and Li, A. What mat- ters in transformers? not all attention is needed. CoRR, his/2406.15786 2024. doi: 10.48550/ARXIV. 2406.15786. URL https: www.i..oo//...o 48550/ arXiv. 2406. 15786.
Khodak, M., Tenenholtz, N. A., Mackey, L., and Fusi, N. Initialization and regularization of factorized neu- ral layers. In 9th International Conference on Learn- ing Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net 2021. URL https: / //openreview.com httt://oorrmmmooo1111ooooooooooo
Hendrycks, D., Burns, C.. Basart, S., Zou, A., Mazeika, M., Song, D., and Steinhardt, J. Measuring massive multitask language understanding. In 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OvenReview.net, 2021. URL https::t///// htttps://ope.revviie http:/eeet/orrrrr 1d=67KB Jmmmmm
Kurtic, E., Kuznedelev, D., Frantar, E., Goin, M., and Alistarh, D. Sparse fine-tuning for inference accelera- tion of large language models. CoRR, 10.0227, 2023. doi: (0448500/ARX1V2310000990 //doi oro.ioorg/000 48550 201...2................. 06927.
URL https:
Hinton, G. E., Vinyals, O., and Dean, J. Distilling the knowl- edge in a neural network. CoRR, abs/1503.02531, 2015.
Kwon, W., Li, Z., Zhuang, S., Sheng. Y., Zheng, L., Yu, C. H., Gonzalez, J. E., Zhang, H., and Stoica, I. Er- ficient memory management for large language model serving with pagedattention, 2023. URL https:// arxiv.org//ao abs////b///a 200.//680......... 06180.
URL
Hive-Digitll-Technlloiee,
Hoffmann, J., Borgeaud, S., Mensch, A., Buchatskaya, E., Cai, T., Rutherford, E., de Las Casas, D., Hendricks,
Lewis, P., Perez, E., Piktus, A., Petroni, F., Karpukhin, V., Goyal, N., Kuttler, H., Lewis, M., tau Yih, W.
12

---

Puzzle: Distillation- Based NAS for Inference-Optimizze LLMs
Rocktäschel, T., Riedel, S., and Kiela, D. Retrieval- augmented generation for knowledge-intensive nlp tasks. ArXiv, abs/2005.11401, 2020. URL https: wew.llssmanticcccheee 21 218869575.
Levine, S. (eds.), Advances in Neural Information Pro- cessing Systems 36: Annual Conference Neural In- formation Processing Systems 2023, NeurIPS 2023, New Orleans, LA, USA, December 10 16. 2023, 2023.
org/CorpusID:
Men, X., Xu, M., Zhang, Q. Wang, B., Lin, H., Lu, Y., Han, X., and Chen, W. Shortgpt: Layers in large language models are more redundant than you expect. CoRR, abs/2403.0320 033853, 2024. doi: 10.48550/ARXIV. 2403.03853, URL https::///// /ttp/////....//./0////533 arXiv. 03853.
Li, C., Peng, J., Yuan, L., Wang, G., Liang. X., Lin, L., and Chang, neural architecture search with knowledge distilla- tion. 2020 IEEF/CVF Conference on Computer Vi- sion and Pattern Recognition (CVPR), pp. 1986-1995, 2019. URL https: : seraeissmatntiiciiciciia org/Corpus ID: 20851333
Block-wisely supervised
Molchanov, P., Hall, J., Yin, H., Kautz, J., Fusi, N., and Vahdat, A. LANA: latency aware network acceleration. In Avidan, S., Brostow, G. J., Cissé, M., Farinella, G. M., and Hassner, T. (eds.), Computer Vision ECCV 2022 17th European Conference, Tel Aviv, Israel, October 23- 27, 2022, Proceedings, Part XII, volume 13672 of Lecture Notes in Computer Science, pp. 137-156. Springer, 2022, doi: 10. 10.1007/978-3-031-117777 00t.111077997111
Li, T., Chiang, W., Frick, E., Dunlap, L., Wu, T., Zhu, B., Gonzalez, J. E., and Stoica, From crowdsourced data to high-quality benchmarks: Arena-hard and benchbuilder pipeline. CoRR, abs/2406. 1939,, 2024. doi: 10.48550/ ARXIV.2406. 11w.22406.119399 1939. URL https: http:ii..ggg//////// 48550/ 21550/arXiv.. 2406. 11939.
9. URL https:/
doi
Lin, S., Hilton, J., and Evans, O. Truthfulqa: Measuring how models mimic human falschoods. In Muresan, S., Nakov, P., and Villavicencio, (eds.), Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), ACL 2022, Dublin, Ireland, May 22-27, 2022, pp. 3214-3252. Association for Computational Linguistics, 2022. doi: 10.18653/ V1/002.ACLLLLON..2222 URLhttps: /ooi.org/ 10..86533v1/
Moons. B., Noorzad, P., Skliar, A., Mariani, G., Mehta, D., Lott, C., and Blankevoort, T. Distill- ing optimal neural networks: Rapid search in di- verse spaces. 2021 IEEF/CVF International Confer- ence on Computer Vision (ICCV), 12209-12218, 2020. URL https: //pii www.nticschole.......... org/CorpusTD::22222 http:a//ppp...ssaatiiiccceeeee
1299
Muralidharan, S., Sreenivas, S. T., Joshi, R., Chochowski, M., Patwary, M., Shoeybi, M., Catanzaro, B., Kautz, J.,
Ling, G., Wang. Z.,, Yan, Y. and Liu, Q. Slimgpt: Layer-wise structured pruning for large language models. In Globersons, A., Mackey, L., Belgrave, D., Fan, A.. Paquet, U., Tomezak, J. M., and Zhang, C. (eds.), Advances in Neural Information Processing Systems 38: Annual Conference on Neural Information Processing Systems 2024, NeurIPS 2024, Vancouver, BC, Canada, December 10- 15, 2024, 2024. URL http: http::tttt::psrt nips. cc http:paperr.eee wttp://wteeeeeeeeeeeeeeeeeeeeeee el c1c4 4 4344444355eee00099444eeeeeee html.
and Molchanov, P. Compact language models via pruning and knowledge distillation. CoRR, abs/2407. 14679, 2024. doi: 10.48550/ARXXV22200..4666 URL https:/ doi doi.org//i /
OpenAL. GPT-4 technical report. CoRR, abs/23330077777 2023. doi: 10.4550/ARXIV.2300..0088 URL https: L/doi orggi.org/100 Confe 0t..1....855/// 4004448550////
.08774.
Penedo, G., Kydlícek, H., Allal, L. B., Lozhkov, A., Mitchell, M., Raffel, C., von Werra, L., and Wolf, T. The fineweb datasets: Decanting the web for the finest text data at scale. CoRR, 0tw/2400.175555 2024. doi: 10.48550/ARXI.2406.115 org/10
Liu, H., Li, C., Wu, Q. and Lee, Y. J. Vi- sual instruction tuning. ArXiv, abs/2304.08485. 2023. URL http:s/app...mattiiiooeo semanticscholes.
URL https:// //doi.
Lu, C., Zhang, J., Chu, Y., Chen, Z., Zhou, J., Wu, F., Chen, H., and Yang, H. Knowledge distillation of transformer-based language models revisited. arXiv preprint arXiv:2206. 14366, 2022.
Project Gutenberg. https: pestps://ww..o.. http:://wwwggttbbbg.gee org.
Real, E.,, Moore, S., Selle, A., Saxena, S., Suematsu, Y. L., Tan, J., Le. Q. V., and Kurakin, A. Large-scale evolu- tion of image classifiers. In International Conference on Machine Learning. 2017. URL https: : /api. s semant.ccsccolar.orggg/
Ma, X., Fang, G., and Wang. X. Lim-pruner: On the struc- tural pruning of large language models. In Oh, A., Nau- mann, T., Globerson, A., Saenko, K., Hardt, M., and
13

---

Puzzle: IisiillationBaasee NAS for Inferenee-Optimizze LLMS
Sakaguchi, K., Bras, R. L.., Bhagavatula, C., and Choi, Y. Winogrande: An adversarial winograd schema challenge at scale. In The Thirty-Fourth AAAI Conference on Ar- rificial Intelligence, AAAI 2020, The Thirty-Second In- novative Applications of Artificial Intelligence Confer- ence, IAAI 2020, The Tenth AAAI Symposium on Educa- tional Advances in Artificial Intelligence, EAAI 2020, New York, NY. USA. February 7-12. 2020, pp. 8732- 8740. AAAI Press, 2020. doi: 10. 10.1609//AAAVVV4105. 6399. URL https: /ttp:///ooooooo http:////oo...r//....00 160.9//0..60/9//0// aaai.
Wang, Z., Bukharin, A., Delalleau, O., Egert, D., Shen, G., Zeng, J., Kuchaiev, O., and Dong, Y. Helpsteer2- preference: Complementing ratings with preferences. CoRR, abs/2410.012577, 2024. doi: 10.48550/ARXIV. 2410.01257. URL https::///// arXiv, Dexii.2410..0 2410. 01257.
Wei, J., Wang, X., Schuurmans, D., Bosma, M., Ichter, B., Xia, F., Chi, E. H., Le, Q. V., and Zhou, D. Chain-of- thought prompting elicits reasoning in large language models. In Koyejo, S., Mohamed, S., Agarwal, A., Bel- grave, D., Cho, K., and Oh, A. (eds.), Advances in Neu- ral Information Processing Systems 35: Annual Confer- ence on Neural Information Processing Systems 2022, NeurIPS 2022, New Orleans, LA, USA, November 28- December 9, 2022, 2022.
v341 05. 6399.
Sanyal, S., Shwartz-Ziv, R., Dimakis, A. G., and Sanghavi, S. Inheritune: Training smaller yet more attentive lan- guage models, 2024. URL https: /ttps:///axvv... wtw.ii..iiii................ abs/2404 08634.
Shazeer, N. Fast transformer decoding: One write-head is all you need. CoRR, abs/191 1 1.02150, 2019. URL http:// /tap:://axxtee wttp/v.xvvvvvo/e wrw..brggab............. /1 9111,,9111111111111111111 .02150.
Xia, M., Zhong, Z., and Chen, D. Structured pruning learns compact and accurate models. In Muresan, S., Nakov, P., and Villavicencio, A. (eds.), Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), ACL 2022, Dublin, Ireland, May 22-27, 2022, pp- 5513-1528. Association for Computational Linguistics, 2022. doi: 10.18653/ V1/0022ACCLLLON..000 107. URL httpt://////ttt //doi::///.... htt///oo..rogg org/ 10.18653/v1/2022. acl -oo...n..... 107.
Solaiman, I. Talat, Z., Agnew, W., Ahmad, L., Baker, D. K., Blodgett, S. L., III, H. D., Dodge, J., Evans, E., Hooker, S., Jernite, Y., Luccioni, A. S., Lusoli, A., Mitchell, M., Newman, J., Png, M., Strait, A., and Vas- silev, A. Evaluating the social impact of generative AI systems in systems and society. CoRR, ths/2200600599p3
Xia, M., Gao, T., Zeng. Z., and Chen, D. Sheared llama: Accelerating language model pre-training via structured pruning. In The Twelfth International Conference on Learning Representations, ICLR 2024, Vienna, Austria, May 7-11, 2024. OpenReview.net, 2024. URL https: / //openrevil ew net/ / forum? d=
2023. doi: 10048550/ARXXI12306000555 orgilorr//100 00w.444450//iix..
URL https: 05949.
Soldaini, I.,, Kinney, R., Bhagia, A., Schwenk, D., Atkin- son, D., Authur, R., Bogin, B., Chandu, K. R., Dumas, J., Elazar, Y., Hofmann, V., Jha, A. H., Kumar, S., Lucy, L., Lyu. X., Lambert, N., Magnusson. I., Morrison, J., Muennighoff, N., Naik, A., Nam, C., Peters, M. E., Ravichander, A., Richardson, K., Shen, Z., Strubell, E., Subramani, N., Tafjord, O., Walsh, E. P., Zettlemoyer, I., Smith, N. A., Hajishirzi, H., Beltagy, I., Groen- eveld, D., Dodge, J., and Lo, K. Dolma: corpus of three trillion tokens for language model pre- training research. In Ku, L., Martins, A., and Sriku- mar, V. (eds.), Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Vol- une 1: Long Papers), ACI. 2024, Bangkok, Thailand, August 11-16, 2024, pp. 15725-15788. Association for Computational Linguistics, 2024. doi: 10.18653/V1/ 2024.ACI-LLON..AA5 URL https:///// / adei.org. 840.
Xiao, G., Tian, Y., Chen, B., Han, S., and Lewis, M. Effi- cient streaming language models with attention sinks. In The Twelfth International Conference on Learning Repre- sentations, ICLR 2024, Vienna, Austria, May 7-11, 2024. OpenReview.net, 2024. URL htttt::://ppppee//// net
open
Yao, S., Yu, D., Zhao, J., Shafran, I., Griffiths, T., Cao, Y., and Narasimhan, K. Tree of thoughts: Deliberate problem solving with large language models. In Oh, A., Naumann, T., Globerson, A., Saenko, K., Hardt, M., and Levine, S. (eds.), Advances in Neural Information Processing Systems 36: Annual Conference on Neural Information Processing Systems 2023, NeurIPS 2023, New Orleans, LA, USA, December 10 16, 2023, 2023.
Sun, M., Liu, Z., Bair, A. and Kolter, J. Z. A sim- ple and effective pruning approach for large language models. In The Twelfth International Conference on Learning Representations, ICLR 2024, Vienna, Austria, May 7-11, 2024. OpenReview.net, 2024. URL https: //openreviews net/ formm? wtt.orrmmmidooooottt3t30
Zellers, R., Holtzman, A., Bisk, Y., Farhadi, A., and Choi, Y. Hellaswag: Can machine really finish your sentence? In Korhonen, A., Traum, D. R., and Mårquez, L. (eds.), Proceedings of the 57th Conference of the Association for Computational Linguistics, ACL 2019, Florence, Italy, July 28- August 2, 2019, Volume I: Long Papers, pp-
14

---

Puzzle: Distillation-aass NAS for Inference-Optimizze LLMs
4791-4800. Association for Computational Linguistics, 2019. doi: 10.18653/V1/119 1472. URL https:// doi. www.oogg/..... 1tt.1.11111//1111111111 1472.
A. Initialization of Alternative Subblocks
accelerate the distillation process, we introduce training- free initialization techniques for alternative subblocks. We propose method reduce the intermediate dimensions of FFN subblocks selectively pruning channels based on their contribution. Our approach, called Channel Contribu- tion, uses an activation- baade strategy to estimate channel contribution during forward passes over a calibration dataset, similar to (Muralidharan et al., 2024). To guide pruning, we rank intermediate channels based on their impact on the FFN output. This is quantified as the distance between the original FFN output and the output after pruning. The channels with the lowest contributions are prioritized for pruning during subblock initialization.
Zheng, L.,, Chiang, W., Sheng, Y., Zhuang, S., Wu, Z.,, Zhuang, Y., Lin, Z., Li, Z., Li, D., Xing, P., Zhang, H., Gonzalez, J. E.,, and Stoica, I. Judging llm-as-a-judge with mt-bench and chatbot arena. In Oh, A., Naumann, T., Globerson, A., Saenko, K., Hardt, M., and Levine, S. (eds.), Advances in Neural Information Processing Systems 36: Annual Conference on Neural Information Processing Systems 2023, NeurlPS 2023, New Orleans, LA, USA, December 10 16, 2023, 2023.
Zoph, B. and Le, Q. V. Neural architecture search with reinforcement learning. ArXiv, abs/1611.01578, 2016. URL https: http//pp...aaaiiccc.......eeeeee ort//oorpusDD:11710022
To formalize our method, let us denote the hidden dimen- sion by H, the FFN intermediate dimension by L, the FFN down projection matrix by Wdown R Let X R represent the FFN's intermediate activations for single token. The output of the FFN, Y E R is then given by
Y (Wddww) T X Σ XgWadon
We define the per-token contribution of channel as:
)
C((X) =
X Wedwwn
X W srowww
k.:
= =|xal- ||
We then compute the average contribution each channel across all tokens in the calibration dataset.
When replacing FFN subblock with linear layer, initialize it by computing the product of the up and down projection matrices for FFNs, thereby effectively ignoring the gating mechanism.
In attention subblocks with a reduced number of key-value heads, we initialize the projection matrices for the key and value heads by mean-pooling them into single projection matrices, following the approach used in (Ainslie al., 2023). When replacing an attention subblock with linear layer, we use the product of the value and output projection matrices, which simulates the scenario where each token attends only to itself.
B. Search Algorithm: Mixed-integer Programming
The search space for LLM architecture optimization enor- mous (we consider 10138 possibilities), as discussed in
15

---

Puzzle: Distillation- Based NAS for Inferenee-Optimizzed LLMs
Section 2. Efficiently navigating this space requires two key components: fast method to estimate candidate quality, which we addressed in Section 4.2, and an efficient algo- rithm to maximize the estimated quality under hardware- specific constraints.
rin b) is the runtime of block variant j in layer i when processing batch size b.
Memory max is the maximum allowed total memory, specified to fit the target GPU(s).
The structure of transformer models naturally frames our optimization problem as a grouped variant of the classical Knapsack Problem. Each layer of the model represents group, containing various block alternatives (attention and FFN variants) as items. Each block alternative has asso- ciated value (its quality score) and multiple costs (parameter memory, KV-cache memory, and runtime characteristics). The objective is to select exactly one block variant from each layer while maximizing the total score and satisfying deployment constraints.
Throughput is the minimum required throughput (tokens per second).
Latencymax is the maximum allowed latency per batch,
The objective maximizes the sum of quality scores across all selected block variants. The first constraint ensures the total memory usage stays within limits, accounting for both parameter memory (shared across batches) and K V--cchhe memory (which scales linearly with batch size as each se- quenee requires its own cache). The second constraint en- forces a minimum throughput requirement: for batch size b, we process b see len tokens within the total runtime, which must meet or exceed Throughputmin tokens per second. The third constraint ensures the total processing time for a batch does not exceed the maximum allowed latency. The fourth constraint guarantees exactly one variant is selected for each layer. Note that we do not impose any constraint on the num- ber of model parameters they are reduced naturally by the search algorithm due to the true scenario constraints.
Following (Molchanov et al., 2022), we formulate this as a Mixed Integer Programming (MIP) problem. Let Fig be a binary decision variable indicating whether block variant j is selected for layer i. For a model with L layers and K; variants per layer, the optimization problem is:
maximize
subject to
memparams (1.5) + h- (4.))) S Memorrymx
b. Jen
Throughputmin
Since the batch size b is not variable in this optimization, we solve the MIP problem multiple times with different values of b to explore the runtime- trade-off space. Larger batch sizes typically result higher throughput for all block operations, but also higher latency and more mem- ory for KV-cache storage, which forces a reduction mem- ory meet the constraints (such as reducing the number of KV heads). For each set of deployment constraints, we choose the batch size that produced the highest quality archi- tectures. If the target scenario specifies a maximum batch size, such as the typical number of active users for a chat bot, the search can be capped accordingly.
nuntime (1. b)
Ki
runtime(1, j. D)m,, S Latency
Viee Iy
Ju E (0.1) W.J.
where:
score(i, is the quality score of block variant j in layer i. measuring how well it maintains the parent model's performance. If the block scores represent a negative impact (e.g. LM loss or KL divergence) we minimize the sum of scores instead of maximizing it.
While MIP problems are NP-complete, modern solvers can efficiently handle instances of our size. Using the open- source pon-ho package (Inc., 2023), we obtain high- quality solutions within seconds. This efficiency enables us to explore multiple architecturally diverse solutions by adding a diversity constraint:
m1 is the parameter memory required for block variant j in layer i, which is shared across all sequences in a batch.
K: Lig Sa. Vy
memw. is the key-value cache memory required for a single sequence in layer i with block variant j.
b is the batch size- the number of sequences processed in parallel during inference.
where Y is the set of previous solutions and E [0, controls the maximum allowed similarity. For example, with a = 0.8, each new solution must differ from previous solutions in at least 20% of its layer choices. This constraint helps discover meaningfully different architectures.
seq_een is the total sequence length, including both prefill and generation.
16

---

Puzzle: Distillation- Based NAS for Inference-Optimizzed LLMs
A key feature of our approach is its ability generate so- lutions precisely tailored for specific hardware platforms. For example, in platforms with limited memory intended for batch size 1, the algorithm strongly favors semory-ssving techniques such as FFN pruning, while de-prioritizing K V- cache optimizations which have minimal impact at batch The hardware specificity extends to architectural features of
inference engines and awareness to runtime considerations when applying on NAS scheme contributes greatly to the usability of resulting models, such as Nemotron-51B.
D. Human Evaluation
A blind-test comparison between Nemotron-55 51B and Llama- 70B was conducted in the following manner:
the inference devices- for example, on H100 GPUs, FP8 quantization offers ~22 acceleration and can be used ag- gressively, while A100 GPUs where FP8 is unavailable, different optimization strategies must be employed. Even the difference in iPter-GPU bandwidth between H100 PCI-E and NVLink configurations influences the optimal architec- ture by affecting tensor parallel synchronization costs.
set of 169 samples was sent.
The evaluation was done by Nvidia's 's data factory team.
Three annotators annotated each sample independently, resulting total of 169 507 annotations.
This flexibility enables a powerful "train once, adapt on demand" methodology that requires minimal human inter- vention. After building a block library and computing block scores once, we can efficiently generate different architec- tures optimized for various deployment scenarios without additional training or manual tuning. The user needs only to specify the available block configurations for each platform the sardware-specific measurements of block variants nat- urally guide the optimization toward hafform-appropriatt solutions. This approach makes Puzzle particularly valuable for real-world deployment, where the same parent model might need to be optimized differently across diverse hard- ware configurations and deployment constraints.
Annotators saw [prompt, completion 1, completion 2] and had choose between options:
Completion is better.
Completion is better.
Both are good.
Neither is good.
The order of the completions randomized to avoid positional bias.
The test set was curated by the project's product team and Nvidia's data factory and included the following tasks and subtasks listed in brackets: The test set was curated and included the following tasks and subtasks listed in brackets:
C. Supporting Fast Inference for Variable-Block Architectures in TensorRT-LLM
Long form text generation (write blog, write story, other).
TensorRT-LLM is highly optimized LLM runtime de- signed to accelerate inference performance on NVIDIA GPUs. It provides industry-leading performance for both latency and throughput-oriented workloads. The runtime enables LLMs to run on GPUs while utilizing custom paged attention (Kwon al., 2023) kernels to efficiently man- age KV caching across sequences in a batch. Further- more, TensorRT-LLLM supports FP8 quantization and var- ious scheduling policies that allow LLM deployments to optimally utilize the underlying hardware.
Long inputs (Open book QA, Multi-hop questions, text to headline)
Grounded text generation (table of contents to blog, paraphrasing).
Multi-condition instructions (3-5 conditions, 6- con- ditions).
Knowledge and trivia.
Summarization (full document summary, summarize bullet points, summarize paragraph to sentence).
To enable the use of Puzzle-generated architectures from the search space (see Section 2) major underlying assumption of TenSSoRR--LLM had to be revised: that all attention layers contain the same number of key and value heads. We de- vised changes the paged KV cache strategy that enabled variable GQA ratios within model. TensorRT-LLM now supports running any architecture from our search space including variable blocks and linear replacements, using FP8 precision for weights, activations and KV cache. Our NAS framework is designed to generate models that run effi- ciently in real inference scenarios. Therefore full support in
Reasoning (temporal reasoning, cause and effect, navi- gation, general reasoning).
Semantic extraction.
E. RULER Benchmark Performance Tables
This section provides complete performance tables for our parent-child pairs on a subset of the RULER benchmark across all context lengths evaluated.
17

---

Puzzle: Distillation-aan NAS for Inference-Optimizzed LLMs
E.I. N vs. Llama-3.1-70B-Insst
verse Distillation Mix outperformed those trained on the limited-domain Project Gutenberg, but Gutenberg- trained models still retained 93% of performance, showcasing Puzzle's 's robustness (see Appendix F.1.2).
Table 19 shows the results of the parent (Llama-3.1-70B- Instruct) and child (Nemotron-51BB for all context lengths. As noted in the main text, Nemotron-511 was trained on sequences up to only 8K tokens yet retains more than 96% of its parent's performance at 16K. The child's performance degrades at 64K and beyond, which is unsurprising given its training horizon. Nevertheless, these results underscore that a large fraction of the parent's long-context capabilities can remain intact even without explicit training such long sequences.
Training Dataset Size:
BLD achieves strong perfor- mance even with smaller token budgets, with diminish- ing returns beyond 0.5B tokens (see Appendix F.1.3).
KL divergence scoring out- performed LM loss and task-specific downstream scor- ing, demonstrating better balance between general- ity and accuracy, although task-specific scoring pro- vides an advantage for customized target tasks (see Appendix F.1.4).
Block Scoring Metrics:
E.2. Nemotron-49B-Ran vs. 1a0mm---3--70BInstttur
In Table 7, we present the extended context-length results for Nemotron-49B- Base (uptrained with sequences up to 128K tokens) alongside its parent (Llama-3.3 Innsttttt. We rename columns to make them more consistent with the style ahove, although some tasks differ from those used for Nemotron-51B. Nemotron-49B- Base preserves 98% or more of its parent's 's performance up to 64K tokens and re- mains above 94% 128K. This highlights that adding short uptraining phase on longer contexts effectively ex- tend the context range of Puzzle-derived models. Additional details and per-task results appear below.
Reduced Search Space: Constraining the search space to no-op alternatives simplifies optimization and eliminates BLD but results in lower accuracy (75.4 vs. 78.39 MMLU), highlighting the value of diverse block replacements for optimal performance (see Ap- pendix F.1.5).
F.1.1. COMBINING COUPLED BLD AND DECOUPLED BLD
We investigate the effects of coupling in BLD (see Sec- tion 3.1) on Puzzle derivatives of Llama-3.1-88 1-8.1--8-IInsiie Instruct, which has 32 layers. For each layer, the search space we consider contains |A; = 6 variants of the attention subblock and JF = 12 variants of the FFN subblock. With coupled BLD, this would amount to training 6 12 - 32 = 2304 blocks. Decoupled BLD reduces the training requirements to only (4 + 10) 32 448 subblocks, which is consid- erably more resource-efficient. Note that besides moving from multiplicative composition to additive composition, with decoupled BLD we also do not need to train the no-op and parent variants in A; and F,-
F. In-Depth Analysis and Ablation Studies
To better understand the key components and design choices of the Puzzle framework, we conduct a series of detailed analyses and ablation studies. We evaluate the importance of global knowledge distillation, investigate the impact of training dataset size and composition, and analyze how our MIP solver adaptively chooses architectures under vary- ing constraints. These studies not only validate our design decisions but also provide insights into the fundamental trade-off's in LLM architecture optimization and the relative importance of different architectural components.
We propose a two-stage technique to reap the benefits of coupled BLD while still keeping the computational cost reasonable. First, run the full Puzzle framework with de- coupled BLD. Then, we analyze the architectures produced by the search algorithm, and identify the most prominent choices of subblock variants. We shrink the search space to include only these choices, then run the full Puzzle frame- work with coupled BLD on the reduced search space. In our case, we were able shrink |A, from 6 to 4. and [F|| from 12 to 3, resulting in a total number of 4 4-3-32 to train in coupled BLD, which is similar to the number we trained in decoupled BLD for the larger search space. This approach produced a higher-accuracy architecture. Note, however, that the one created with decoupled BLD was al- ready significantly above the efficient frontier. See results in Table 8. Both models underwent short GKD uptraining.
F.. Block Library Construction Ablation Studies
In Appendices 1. to F.1.5, we analyze the impact of crit- ical decisions in block library construction. These studies examine the fundamental trade-offs between computational cost, dataset characteristics, and model performance. Our key findings are:
Coupled vs. Decoupled BLD: Decoupled BLD re- duces training cost by transforming the search space from multiplicative to additive. Combining decoupled BLD for subspace narrowing with coupled BLD for refinement improves accuracy while maintaining com- putational costs (see Appendix F.I.I).
384 blocks
Dataset Composition: Models trained on the di-
18

---

Puzzle: Distillation-Basee NAS for Inference-Optimizzed LLMs
Table Performance comparison of the parent 1.3-am----0----smmmmma and child WNcmotron-49B-aae on subset of the RULER
benchmark. 'Accuracy Preserved' is (child / parent) x 100.
Centent
seglaarehh #ngkoredii.)) wellawek.) refl.wedff.1
vellawele)
vedi.vete rulluper vetelrreeeec. cwwenwwkkuiiide Irgreemhniirele
pupe eJlepet
frengr
Aner Proveved US)
Perel: Hire:
16
18
(8)
191
103
100
71
103
13
191
90.1
15
1533
33
65
1130
455
08.1
128k
13
98.25
15.35
33
1633
$3355
$2.25
ClAbs
46.
18
18)
191
100
100
160
5831
15
18868
146
188
(8)
181
100
9035
100
124
16
18811
64R
108
100
191
9035
365
73
9130
03.0
across domains regardless of the training data composition.
Table 8. The effect of coupled BLD decoupled BLD on high- throughput child derivatives of Llama-3. 1-8B-Innsruun We found a relevant subspace of the scarch space using a decoupled BLD Puzzle. then trained coupled BLD on this subspace and ran separate Puzzle, leading to additional improvement. Throughput is estimated via the sum of measured block runtimes on single NVIDIA RTX 4090 GPU. Accuracy = (MT-Bench x + MMLU)
Table 9, Benehmark results on Llama-3. 3.0--....7-00--0 1-1---8--nnstrree Instructt derivatives obtained from Puzzle without uptraining applied with different datasets.
Model
MT-Bench
MMLU
MMLU-STEM
-ruinberg-Trriied Trained
7.98
74.84
64.5
72
DesiilatooMiix-rriined Trained
8.61
78.39
70.35
Model
Throughput*
Accuracy
Puzzle with Coupled BL.D
5856
73.98
Puzzle with Decouple BLD
5834
73.10
-lama-3.2-3B-Insttt
5737
70.36
F.1.3. IMPACT OF BLOCKWISE LOCAL DISTILLATION TRAINING DATASET SIZE
To evaluate the efficiency of BLD under varying training dataset sizes, we conducted experiments using different token budgets. Specifically, we trained the same set of block variants using 0.25B, 0.5B, and 1.0B tokens from the Distillation Mix dataset (see Section 3). After completing the BLD stage for each token budget, we used the MIP optimization stage to generate optimized architectures under similar constraints those used to produce Nemotron-51B. and evaluated their performance downstream tasks.
F.1.2. IMPACT OF DATASET COMPOSITION ON DERIVED MODELS
We evaluate the robustness of the Puzzle framework to training data composition by comparing performances up through the BLD stage (prior to GKD). For this analysis, we contrast two datasets: our domain-diverse Distillation Mix (described in Section 3) and the English subset of Project Gutenberg (Project Gutenberg). The latter, dataset pre- dominantly comprising literary works, which lacks diverse coverage of technical, conversational and STEM-specific content, making it an interesting test case for framework robustness.
In Table 10 we present the performance of models trained with different BLD token budgets. Notably, while all models achieved comparable MMLU scores, the MT MT-Bench results indicate a more pronounced performance improvement as the token budget increases, particularly in multi-turn con- versational tasks. However, the improvements diminish as the token budget grows (e.g., the boost from 0.25B to 0.5B tokens is larger than that from 0.5B to 1.OB), suggesting that BLD facilitates rapid recovery of the parent model's performance even with moderate training budgets. These findings imply that longer BLD training can yield better block libraries but with diminishing returns larger scales.
As shown in Table 9, models derived using Project Guten- berg data (for training, pruning and block scoring) demon- strate strong performance despite the dataset's limitations. On general benchmarks like MT-Bench and MMLU, the Gutenberg-traineg trained model achieves 92.7% and 95.5% of the performance obtained with Distillation Mix, respectively. Even on STEM categories within MMLU, where the train- ing data's limitations are most relevant, the model maintains 91.7% of the performance (64.5 vs 70.35).
F.1.4. IMPACT OF DIFFERENT BLOCK SCORING METRICS
These results demonstrate that the Puzzle framework can effectively transfer knowledge from the parent model even when the training data provides limited coverage of specific domains. This robustness is particularly noteworthy given that no GKD uptraining was performed, suggesting that our BLD approach effectively preserves model capabilities
In Section 4.2, we presented three possible metrics for replace- 1-blckk scores: downstream accuracy, LM loss and KL divergence. We investigate the effects of using different replace---bloce scores when applying the Puzzle frame- work to Llama-3,1 --18-BB--BB 1--1--8BInnne In--rrutt with large search space
19

---

Puzzle: Distillation-Base NAS for Inference-Optimizzed LLMs
world knowledge, conversational abilities-are likely concentrated in different parts of the model. Due to the com- putational expense of calculating downstream accuracy, we use the reduced search space from Appendix F.1.1 to keep the experiment feasible. To evaluate downstream accuracy, we use the MMLU benchmark, splitting its 57 tasks into two nearly equal-sized sets stratified by the MMLU cate- gories {STEM, Social Sciences, Humanities, Other). These subsets are referred to Half-MMLUU with one serving as "train" set for block quality scoring and the other as "test" set for evaluation. Table 11 shows that, even with the same library of trained blocks, the block selection process can be customized to construct architectures optimized for specific target tasks. We note that this customization is not without cost: the architecture constructed using MaLf scores achieves the best accuracy on the target task, but its MT- Bench accuracy of 7.57 is lower than the 8.06 MT-Bench accuracy of the architecture constructed according to the more general KL divergence scores, Both models were con- structed using decoupled BLD and underwent short GKD uptraining.
Table 10. Performance comparison of Puzzle-optimized architec- tures trained with varying BLD token budgets. Metrics include MT-Bench and MMLU scores.
BLD Token Budget
MT-Bench
MMLU
1.OB Tokens
8.98
78.54
0.5B Tokens
8.86
78.44
0.25B Tokens
8.51
78.27
- - 6, - 12). We compare the use of the model's loss on validation set (LM loss in our case), a common choice in earlier decomposed NAS methods used in CV, with our proposed method based KL divergence. LM loss aims to capture the general quality degradation induced by changing a specific parent block, while KL divergence aims to capture the distance from the parent model induced by this change. As illustrated in Figure 7, KL divergence scoring results in better Puzzle architecture choices than scoring with LM loss. All models were constructed using decoupled BLD and underwent short GKD uptraining.
75
Div
Table II. The effect of task-oriented block scoring on high- throughput child derivatives of Llama-3.1 1-1----I-nstree We split the tasks in MMLU into two equal-sized sets and use one of them for block quality scoring and the other for evaluation, showing that even with the same library of trained blocks, block selection can be customized to build architectures that fit desired target task. Throughput is estimated via the sum of measured block runtimes on a single NVIDIA RTX 4090 GPU.
74.5
74
Accuracy
73.5
LM Loss
KL Div
73
Model
Threoghput*
MalE-MMLA Accuracy (Tes Seb
Pioale sccead with Hull-MMI Al accuracy (Imin sat)
S818
56:24
Paool: scceal wit Kl. dieeggace
S814
64.4
72.5
Liaoo 3.2 Instrct
5737
60006
LM Loss
72
5500
5900
F. 1.. EFFECTS OF LIMITED SEARCH SPACE DIVERSITY
Throughput
To explore the impact of reducing the search space com- plexity on the Puzzle framework, we constrained alternative child blocks to only allow replacing parent model subblocks with no-op operations. This eliminates the need for BLD as additional block variants that require training are consid- ered, further reducing the computational costs.
Figure 7. Accuracy vs. Throughput performance of Llama-3.1 1--B- Instruct child derivatives, some constructed using LM loss as replace-l-block scores, and some constructed using KL divergence as replace-I-biee 1-block scores. LM loss aims to capture the general quality degradation induced by changing a specific parent block, while KL divergence aims to capture the distance from the parent model induced by this change. KI. divergence results in better architecture choices. Accuracy = (MT-Bench x MMLU) Throughput is estimated via the sum of measured block runtimes on a single NVIDIA RTX 4090 GPU.
The resulting architecture, optimized using the same MIP approach but limited to no-ops, was evaluated against pre- uptraining Nemotron-51B, which was derived using the same Puzzle pipeline but with a more diverse block vari- ants. As shown in Table 12, the no-op-only model retained high throughput but exhibited noticeable drop in MMLU accuracy compared to Nemotron-51B (75.4 vs. 78.39).
Next, we explore the use of downstream accuracy replace- 1-block scores to customize the algorithm's 's block selection for specific target tasks. Our hypothesis is that, within the same budget, different architectures may excel at differ- ent tasks because distinct capabilities- such as reasoning,
These results illustrate the flexibility of the Puzzle frame- work in balancing resource constraints and model perfor- mance. Although limiting the search space simplifies the optimization process and reduces training costs even further,
20

---

Puzzle: Distillation- Based NAS for Inference-Optimizzed LLMs
sacrifices the fine-grained architectural customization en- abled by a broader range of block alternatives. This demon- strates that more diverse search space leads to architectures that achieve better accuracy.
architecture optimized for specific throughput target, with darker colors indicating higher computational cost relative to the parent model. The architecture of Nemotron-51 corresponding to a throughput target of 5500 tokens per second, is marked in green. Several intriguing patterns are evident:
Table 12. Comparison of pre-uptraining Nemotron-51B (derived using the full scarch space) and no-op-only variant.
Counter-intuitive local optimizations: While stricter throughput constraints generally lead faster blocks, we observe surprising inversions of this trend. For example, in layers 1-4, the MIP sometimes chooses computationally heavier blocks for higher throughput targets. This counter-intuitive choice suggests that local slowdowns can enable better global optimization, with other layers compensating to meet the overall throughput constraint.
Model
MMLU
Throughput (tokens/sec)
Puzzle (No-op only)
75.40
5604.18
Puzzle (Full search space)
78.39
5500.25
E22 Search Algorithm Ablation Studies
In Appendices 1.2.1 I to F.2.3, we analyze the MIP optimiza- tion process and alternative approaches within the Puzzle framework. These studies focus on how throughput con- straints and scoring strategies influence architecture selec- tion and model performance. Our key findings are:
Asymmetric treatment of components: The MIP treats attention and FFN components quite differently, While attention mechanisms are completely eliminated in some layers even under lenient throughput con- straints, FFN components are never entirely skipped. This suggests that FFN layers might play more fun- damental role in maintaining model capabilities, while attention mechanisms offer more flexibility for opti- mization.
MIP Optimization and Throughput Constraints: The MIP solver adapts architectures to meet various throughput targets, revealing nuanced trade-offs be- tween global and local efficiency. Notably, FFN com- ponents are preserved even under strict constraints, highlighting their critical role in maintaining model accuracy (see Appendix F2.1).
The heatmap reveals distinct "phases" across layer depths. Shallow layers (0-15) show high variability in both attention and FFN, mid- dle layers (16-42) maintain more consistent computa- tion, and deeper layers (43-80) show different patterns for attention versus FFN optimization. This suggests different layers serve distinct roles in the network's information processing.
Architectural phases:
Greedy Algorithm Alternative: A budget- constrained greedy algorithm, while simpler, resulted in significantly lower accuracy (70.74% MMLU) compared to MIP-based optimization (78.39%). This underscores the importance of global optimization for achieving superior accuracy-efficiency trade-offs (see Appendix F.2.2).
Throughput-dependent transitions: The solution patterns show clear transitions as throughput require- ments increase, but these transitions are not uniform across layers. Some layers maintain consistent compu- tation across different throughput targets while others show sharp transitions, indicating varying sensitivity to throughput constraints.
Data-free Scoring: Maximizing parameter count as a heuristic produced architectures with sharp perfor- mance declines (23.12% MMLU), emphasizing the necessity of data-driven scoring mechanisms for effec- tive architecture optimization (see Appendix F.2.3).
Random Architecture Baselines: Nemotron-51 51B out- performs both fully random and random-from
These patterns demonstrate the sophisticated optimization strategies discovered by the MIP solver, revealing that opti- mal architectures often require nuanced tradeoffs between local and global efficiency. The preservation of FFN com- putation even under striet constraints provides empirical evidence for the relative importance of different architec- tural components in such transformer models.
architectures under the same training and speed con- straints, underscoring the role of the MIP solver in selecting effective architectural compositions (see Ap- pendix F.2.4).
1.2.1. MIP SOLUTION ANALYSIS UNDER VARYING THROUGHPUT CONSTRAINTS
F.2.2. GREEDY SEARCH ALGORITHM ALTERNATIVE
Figure provides an interesting window into how our MIP solver adapts model architecture to different throughput requirements. Each row in the heatmaps represents distinct
To further evaluate the impact of the MIP optimization ap- proach, we implemented a budget-constrained greedy search
21

---

Puzzle: Distillation- Easillation-Based NAS for Inference-Optimizzed LLMs
Adont Rurtine Throughou
algorithm achieves a substantially lower MMLU accuracy of 70.74%; compared to 78.39% for the MIP-derived model, despite both architectures meeting the same throughput con- straint.
Table 13, Comparison of the budget-constrained greedy algorithu and MIP search algorithms for Puzzle. Results are shown for pre- uptraining Nemotron-51 B under identical throughput constraints.
Optimization Method
MMLU
Throughput (tokens/sec)
Greedy Algorithm MIP
70.74
5500.30
78.39
5500.25
Figure 8. Heatmaps showing how attention and FFN runtime pat- terns vary with throughput constraints across model layers. Dark colors indicate higher computational cost relative to the parent model. Each row represents an architecture optimized for spe- cific throughput target, with Nemotron-5 511's configuration (5500 tokens/sec) marked in green.
These findings underscore the critical role of global opti- mization the Puzzle framework. By considering jointly optimizing block selection, MIP achieves significantly better balance between efficiency and accuracy, making it indispensable for extracting the full potential of the Puzzle framework.
algorithm as an alternative. This algorithm prioritizes sim- plicity by using heuristic-based layer-wise selection pro- cess, contrasting with the global optimization provided by MIP.
F.2.3. IMPORTANCE OF DATA-DRIVEN QUALITY ESTIMATION IN ARCHITECTURE SCORING
further understand the importance of quality estimation in search space optimization, we conducted an experiment using simple heuristic: scoring blocks based their pa- rameter count. While not data-driven approach, parameter count serves as a straightforward metric often associated with improved performance in LLMs. This experiment aimed at assessing whether such a basic metric could pro- vide sufficient guidance the search space, shedding light on the role of more nuanced scoring mechanisms.
The greedy algorithm operates as follows:
At initialization, the runtime and memory budgets which are derived from the required constraints, are split equally across layers.
Layer scoring: Each layer is assigned a score based on heuristic metric designed to estimate how easily it can be replaced with minimal performance degradation, In our implementation, this metric the mean replace- 1-block KL divergence score across all block variants for the layer. Layers with lower average scores were deemed easier to optimize.
Under this method, the search algorithm is simplified to selecting the block variant with the largest number of pa- rameters that satisfied throughput and memory constraints. This resulted in an architecture composed uniformly of high-parameter blocks across all layers, without considering layer-specific block quality and representational needs.
Sequential replacement: Layers are processed as- cending order of their scores. For each layer, the algo- rithm selects the block variant with the lowest replace- 1-block KL divergence score that satisfies the layer's runtime and memory budget.
As shown in Table 14, maximizing parameter count leads to sharp drop in MMLU accuracy compared to the archi- tecture derived using the full Puzzle framework with MIP optimization. Despite having similar throughput, the sim- plistic parameter-maxiiizzaion approach fails to achieve competitive results, underscoring the necessity of quality- aware block scoring, and different layers require different computational budgets for optimal performance.
After selecting a block for the current layer, the remaining runtime and memory sav- ings are added to the next layer's budget. This allows for more dynamic algorithm that allocates more re- sources for layers that are harder to replace.
Constraint adjustment:
F.2.4. EVALUATING RANDOM ARCHITECTURES AS SEARCH BASELINES
Table 13 compares the performance of the greedy algorithm with the MIP-derived pre-uptraining Nemotron-511 The results show that the greedy algorithm leads to a significant drop in accuracy, highlighting the importance of global opti- mization. Specifically, the model derived using the greedy
To further validate the effectiveness of our search algorithm, we evaluated several baselines where block selection was performed randomly. All models were trained with 10B tokens.
22

---

Puzzle: Distillation-Basee NAS for Inference-Optimizzed LLMs
F.3.1. Loss COMPOSITION ABLATION STUDIES FOR UPTRAINING
Table 14. Comparison of maximizing parameter count with Puz- zle's ' MIP- based optimization as search algorithms, Results are shown for pre uetrraiiining Nemotron-55 under identical through- put constraints.
Recall from Section 4.3 that our final GKD objective com- prises cosine similarity and KL divergence losses:
Optimization Method
MMLU
Throughput (tokens/sec)
h
h
Maximizing Parameters
23.12
5727.0
Leosine
Σ
(2)
pre-uptraining Nemotron-51B
78.39
5500.25
1=1
N
The first baseline samples block variants from the full block library while satisfying the same throughput constraint as Nemotron- SIB, ignoring their block scores. The second ("Fully Random") uses completely random architecture that is not constrained to trained block variants, but adheres to the same speed constraints. Finally, we added a third Reveen-Randmmzeddwwww
randomly
CILD =
Σ
Ps
log
Pr
(3)
9:
In this subsection, we provide detailed ablation of different loss combinations, including the language modeling (LM) loss:
N Σ yi log(ū:),
LLM
(4)
evaluates Llama- 3.1-70B with randomized weights but no architectural mod- ifications.
where y, is the ground-truth label for the i-th token and it its predicted probability. We paired LLM with cosine and KL divergence losses in various configurations to see which composition best recovers the parent model's 's accuracy.
As shown in Table 15, Puzzle significantly outperforms all baselines. Notably, the random- conddo--room-lle from-library baseline achieves only 86.6% of Nemotron-5 spite being constructed from the same trained blocks. This experiment further emphasizes the role of MIP in selecting high-quality architectural compositions from the library.
Experimental Setup: We used Nemotron-51B (the imme- diate child after BLD) and trained with 5B tokens for each combination, measuring accuracy on MMLU and MT- Bench. Table 17 shows that adding LL often harms down- stream performance, consistent with Muralidharan et al. (2024). In contrast, combining Leosine and LKLD yields the best results. We thus adopt Equation (I):
performance, de-
Table 15. Comparison of Nemotron 51B with randomly con- structed architectures, all trained for 10B tokens. The random- librrry variant uses trained blocks, while the fully random variant ignores the block library. The Randomzeed model uses Ilama-70BB 70B's ' architecture with random weights.
LOKD Leosine + LKLD
as our final uptraining loss, ultimately running it over 45B tokens to produce Nemotron-51 IB-Ittrrutt.
Model
MMLLU
MT-Bench
Asg. Accurncy
Relative to Liuma- 7UB
Nemommon-51 SIB (10B tokens)
79.7
8.89
54.80
08.15
Randun- Samau---foo--nd--itooooo
74.01
58%
Table 17. Ablation study for different combinations of LM loss, block (hidden activations) loss, and logits KLD loss. All mod- cls (Nemotron-51 51B, derived from L 1-l.1-70---oe --0ttc-nstttc were trained for 5B tokens. First row did not undergo uptraining. Adjacent rows with the same color differ only in the LiM compo- nent. "During the KD process for this combination, the validation LXLD consistently increased. Trained for 45B tokens using LOKD defined in Equation (1).
Random
23.13
0.89
16:00
18.3%%
Parent-e Randemined
23:42
0.95
16.46
19 19.25%
Hame-3 1-708
81
8.93
85:48
100%
E3. Global Knowledge Distillation Uptraining
We assess the significance of the final GKD uptraining phase, in enhancing the accuracy of child models derived from Llama-3.1-70 1..1-700-IIstuuuc and results presented in Table 16 demonstrate that this stage contributes to improvements in both the MMLU and MT- Bench benchmark scores.
1-8-a33--8B-Insttucct The
LIM (4) Cosn (2) Cain (3)
MMLU
MT-Bench
Average
Validation CRED
X
78.5 39.39
8.67
82.55
0.19
78.55
7.71
77.83
0.31
83.88
x
x
79.33
8.68
83.07
0.10
79.04
78.52
0.30
79.40
8.74
83.40
0.16
79.45
5.65
83.03
0.14
79.61
$4.16
Table 16. Impact of global knowledge distillation uptraining on
Llama-3. 3.1-70B Instruct (parent)
81.66
8:93
$5.48
0.00
MMLU and MT-B BeTcBenchbee benchmark scores for child models derived
Nematron-S11 10---51----nnnee (childy
80.30
85.10
0.08
from Llama-3, 3...a.3.117000 1-70B-0---- Instruct and Llama-3, 3.J 1-B---BBBB 8B-Instruce
Model Nome
GKD Uptraining
MMLU
Milllench
Average
We hypothesize that LIM may cause overfitting to the up- training dataset, thus reducing the child's ability to mir- ror the parent's 's distribution across downstream tasks. Re- moving this term mitigates open- weights,, closed-data mis- matches, leading to better knowledge transfer.
Liuma 708 Instract
58.48
Numotnon--51 B-lnerus (child)
$2.55
L.amo:3. 1-915- (pasent)
derivative Chan-- 1-14-- (child)
65.46
73.08
23

---

Puzzle: Distillation-Basse NAS for Inference-Optimizze LLMs
F.4. Comparison with Related Work
There are many different techniques to prune compress models, most of which could be complementary to Puz- zle. Some pruning methods, for example, could be used to construct alternative blocks for Puzzle's block library. Nevertheless, in this section we compare several established methods directly against the basic Puzzle framework, with- out integrating them.
We compare Puzzle against structured sparsity (Wanda (Sun et al., 2024)) and low-rank approximation (similar to (Khodak et al., 2021) with subsequent distillation) ap- plied to 3lama-3.11708 under similar throughput constraints. Wanda applied 2:4 structured sparsity without additional training, while the low-rank method used factorized layers followed by distillation. Nemotron-55 significantly out- performed both methods, achieving 99.49% of the parent's 's average accuracy (MMLU and MT-Bench), compared to 92.23% for Wanda and 88.96% for low-rank approximation (see Table 18). Subsequent distillation post-pruning with Wanda slightly improved MMLU (73.69) without impact- ing MT-Bench. Moreover, since both structured spansity and low-rank approximations represent subsets of Puzzle's broader search space, integrating these approaches into Puz- zle could further enhance performance.
Table 18. Comparison of Puzzle, Wanda (structured sparsity), and low-rank approximation methods on Llama 7.B derivatives under similar throughput constraints.
Model
MMLU
MT-Bench
Average Accuracy
Accurscy Preserved (%)
Nemetren-S1R
90.49%
Wands (Son et al., 2024)
72.99
8.39
78.44
92.23%
Law-man
22.87
8:01
76:05
88 994
Iama-3. 1-7708 (Pornn))
81.66
8.99
85.48
1876
Other methods such as Minitron (Muralidharan et al., 2024), ShortGPT (Men et al., 2024), and SlimGPT (Ling et al., 2024) share conceptual similarities with Puzzle, but each represents constrained subset of Puzzle's broader opti- mization space. For the most part, Minitron restricts mod- ifications to homogeneous block replacements across all layers, ShortGPT focuses exclusively on redundant layer removal, and SlimGPT employs incremental pruning ratios via a fixed heuristic. Puzzle generalizes and extends these approaches, allowing for heterogeneous, layer-specific mod- ifications, diverse block alternatives including no-op layers, and customizable pruning ratios optimized globally through MIP-based optimization.
24

---

Puzzle: Distillation-Basee NAS for Inference-Optimizzed LLMs
Table 19. Full performance comparison of the parent (Llama-3.1 -i0ma---.-7OB-Insttric and child (Nemotron-51 on a subset of the RULER benchmark across all context lengths. 'Accuracy Preserved' (child parent) 100. Benchmark names refer the implementatiionssttinng in the official RULER repository. Varies depending on context length.
Contest Length
words estraction"
variable tracking.
vorfable tracking
mords extroction..
Presersed
Purent: Liumn 3.1-700BIItsrrut
N/A
N/A
100.00
100.00
100.00
99:40
99.67
99.81
2K
NIA
88.40
100.00
100.00
100.00
99.53
99.87
97.97
67.40
99.10
67.00
85380
90-96
103.00
99.40
07.67
99.91
92.08
16K
61.00
77.25
10300
95.53
100.00
89.39
72.00
94.53
99.57
56.85
Child: Nemotren 51B
N/A
NIA
10000
100.00
99.40
59.55
79
59/90
N/A
86.00
99.96
99.67
98.40
09.00
97.17
99
63.40
55.00
98.65
50.80
10300
97.79
1SK
$3.40
75.60
99.72
99.80
32K
70.60
98.28
78.54
$7.86
25


Reviews:
[{"review_id": "xC9h9WtoRv", "reviewer": "Reviewer_GBJ9", "summary": "This paper is concerned with model compression, which aims to compress the scales of LLMs. This paper proposes a NAS framework named Puzzle to conduct easy-to-achieve NAS. The Puzzle framework firstly trains decoupled blocks for each layer via block-wise local distillation, then searches best-fit plan for architecture, and finally uptrains the searched architecture for preserved performance. The experimental results show that the NAS-based architecture is more efficient than original architecture and is competitive with the original model in a wide range of tasks.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your feedback and suggestion,\n\n*\"Key baselines are missing, which 1) uses distillation on a fully random architecture 2) uses distillation on a random-from-block-library architecture.\"*\n\nWe conducted evaluations on the baselines you suggested, namely (1) a fully random architecture and (2) a random-from-block-library architecture. Both were constructed to adhere to the same speed constraints as Puzzle-51B.\nAdditionally, we extended (1) with an extra baseline: (3) using Llama-3.1-70B itself with randomized weights (an experiment we dub \"Parent-Randomized\"). This baseline explores whether increased capacity might have contributed to better performance.\nTo ensure fairness, we allocated the same 10B token budget for training each of these models:\n\n| Model                       | MMLU  | MT-Bench | Average Accuracy | Relative to Llama-70B |\n|----------------------------|:-----:|:-------:|:---------------:|:--------------------:|\n| Puzzle-51B (10B tokens)    | 79.7  | 8.89    | 84.3           | 98.61%              |\n| Random-from-block-library  | 66.02 | 8.2     | 74.01          | 86.58%              |\n| Fully Random               | 23.13 | 0.89    | 16.015         | 18.73%              |\n| Parent-Randomized          | 23.42 | 0.95    | 16.46          | 19.25%              |\n| Llama-3.1-70B              | 81.66 | 8.93    | 85.48          | 100%                |\n\nWe will include these baselines in the revision.\n\nWe believe the \"Random-from-block-library\" experiment is particularly informative, as it highlights the value of the MIP algorithm in selecting high-quality blocks from the block library. Additionally, our paper examines another baseline to the MIP algorithm in Appendix F.2.2., where we use a greedy algorithm to select blocks from the block library.", "claims_and_evidence": "The claims are supported by clear and convincing evidence. However, I still have several concerns:\n1) The comparison is sub-optimal and lacks critical baselines concerning about distillation, making the claims seem to be overclaimed.", "methods_and_evaluation_criteria": "The proposed methods are clear and the evaluation criteria is mostly adequate. However, I still have several concerns:\n1) Key baselines are missing, which 1) uses distillation on a fully random architecture 2) uses distillation on a random-from-block-library architecture.", "theoretical_claims": "The proofs for theoretical claims are correctly justified.", "experimental_designs_or_analyses": "The experimental designs and analyses are valid. However,I still have several concerns:\n1) It would be much better to integrate the above-mentioned baselines as performance references.", "relation_to_broader_scientific_literature": "N/A", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "N/A", "other_comments_or_suggestions": "N/A", "questions_for_authors": "N/A", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "6enJ3n9jEv", "reviewer": "Reviewer_bqJm", "summary": "The paper proposes a NAS pipeline for pruning a pre-trained large language model. The search space includes pruning the attention heads for the attention module and pruning FFN columns (intermediate size) for the FFN module. The pipeline includes three pieces: 1) blockwise local distillation: by training each pruned module with distillation loss to recover the original module output, this part produces a library of pruned module. 2) block scoring: the pruned modules are evaluated by their quality relative to the original modules 3) searching the best combination of pruning strategies: by formulating the pruning problem as a mixted-integer programming problem (constraint is the desired efficiency after pruning), we can search for the best set of pruned modules based on their scores.\n\nAfterwards, the pruned model is trained with distillation loss end-to-end again to additional heal the gap. The entire procedure utilizes 45B training tokens and is able to get a 51B model out of llama-3.1-70B-Instruct with almost no performance lost.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your feedback and appreciation of the ablation studies,\n\n*\"There are several metrics considered for the replace-1-block score, but I don't find where/if the author conducts ablation studies on them and which one is the best\"*\n\nAppendix F.1.4. examines the impact of different replace-1-block scores. In short, for general use, it is best to use KL divergence as the metric. We'll make sure to state this conclusion more clearly in the revision.\nWe also found that if a particular downstream task is prioritized, using data similar to that task for block scoring will produce results that outperform the KL divergence *on this specific task*, but underperform KL Divergence solutions across different tasks. See the Half-MMLU experiment in Appendix F.1.4. for more details. Additionally, LM Loss always underperforms KL divergence as a block score (as shown in Figure 7).\n\n*\"the author only evaluates the method on one model (llama-3.1-70B), which raises concern on its generalizability.\"*\n\nThe paper shows the application of the Puzzle method in crafting 11 different models with various constraints, datasets and budgets to show the robustness of the method:\n\n(1) 4 derivatives of Llama-3.1-70B:\na) Puzzle-51B,\nb) 2 other 51B derivatives with a different BLD token budget (Appendix F.1.3., Table 9),\nc) A Puzzle derivative from a different dataset (Gutenberg dataset, Appendix F.1.2., Table 8),\nd) A limited search space variant (Appendix F.1.5., Table 11)\n\n(2) 6 derivatives of Llama-3.3-8B:\na) A \"coupled BLD\" derivative (appearing both in the main paper at Table 5 and in Appendix F.1.4., Table 7),\nb) 2 derivatives with LM loss block scoring (Figure 7 in Appendix F.1.4.),\nc) 2 derivatives with decoupled BLD (Figure 7 in Appendix F.1.4.),\nd) A derivative with a downstream, \"Half-MMLU\" block score (Table 10 in Appendix F.1.4.)\n\n(3) A derivative of Llama-3.3-70B: \"Puzzle-49B\" (Section 5).\n\nWe agree that applying Puzzle to produce a large variety greatly strengthens the paper. That is why, after submission, we applied Puzzle in two additional scenarios with different parent models:\n\n(4) Using Puzzle, we created a 253B derivative of Llama-3.1-405B with Puzzle constraints to fit a single H100 node at a 1.5X latency. The resulting model retains 99.5% of the parent model performance (averaged on MMLU, MT-Bench, MMLU-Pro, HumanEval and Arena Hard).\n\n(5) Experimenting with a novel Mamba-hybrid model (consisting of more than 50B parameters) as a parent, we crafted a Puzzle derivative while constraining it to fit a single RTX 5090 with a 1M context length. The resulting model retains a 99.94% of the parent's performance (average on MMLU, MMLU-Pro, GSM8K and HellaSwag).\n\nWe will make sure the revision includes these examples to emphasize Puzzle's robustness across a variety of constraints and models.\n\n\n*\"The global distillation phase requires a large amount of tokens. 45B tokens in total is not a small number and hinder the applicability of the method.\"*\n\nWe agree 45B tokens might not be an applicable budget for every user, even if it is much lower than the trillions of tokens necessary to train LLMs from scratch. However:\n\n(1) GKD is meant to squeeze extra performance when budget allows. Puzzle derivatives remain surprisingly competitive even without GKD (see Table 14 in Appendix F.3., where Puzzle derivatives retain 96.5% or 90% of the parent's performance without applying GKD at all). \n\n(2) While we didn't mention it in the paper, even a *significantly* shorter GKD training may suffice for a substantial increase in performance, and we will emphasize this in the revision. After only 3.7B tokens invested in GKD, Puzzle-51B had already recovered 98.8% of its parent's performance on MMLU and MT-Bench (MMLU 79.38, MT-Bench 8.96). Puzzle-49B underwent a GKD of only 8.68 tokens (prior to the long context KD described in Section 5), at which stage it already recovered 99.63% of its parent's performance (80.73 MMLU, MT-Bench 8.87). Even after just 2.9B tokens for GKD, Puzzle-49B had already recovered 98.47% (MMLU 80.72, MT-Bench 8.675). We agree that it is important to include these results in the revision to clarify how GKD length can be adjusted based on the available budget.\n\n\n*\"the author had a good amount of ablation in the appendix and the author should select some of them to be put in the main text\"*\n\nThank you for your positive feedback, the ablation studies were of utmost importance for us to objectively conclude the best configurations for using Puzzle in a robust way. We intend to move several ablations into the main body.\n\nIn particular, since our core question in the ablation studies was to find the best configuration for applying Puzzle, we believe the ablations presented in F.1.1., F.1.3. and F.1.4. are the most important for practitioners (with F.1.2. also contributing to data preparation, if space constraints in the main body allow us to add it as well). What is your opinion on this selection?", "claims_and_evidence": "Yes, both the speedup and the accuracy of the model support the efficacy of the method.", "methods_and_evaluation_criteria": "The accuracy evaluation setup makes sense, but the author only evaluates the method on one model (llama-3.1-70B), which raises concern on its generalizability.", "theoretical_claims": "Not applicable for this paper.", "experimental_designs_or_analyses": "The main expeirmental design makes sense. There are not much ablations/analysis in the main paper though, which can be an area for further improvement.", "relation_to_broader_scientific_literature": "There are two branches of LLM pruning papers: structural pruning and unstructural pruning. This paper falls under the category of structural pruning and proposes a systematic NAS method to prune the model which is not found in previous literature. The methodology (recovery error, distillation) isn't super new but in combined showed promising performance. The main strength of this paper is its empirical success, where most of the unstructured pruning have either worse performance or an insufficient evaluation setup.", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "Strengths\n- Comprehensive evaluation and strong empirical results.\n\nWeakness\n- There are several metrics considered for the replace-1-block score, but I don't find where/if the author conducts ablation studies on them and which one is the best (there is only one analysis in the appendix which doesn't fully address this question).\n- The global distillation phase requires a large amount of tokens. 45B tokens in total is not a small number and hinder the applicability of the method.\n- The method is only tested on llama-3.1-70B, it is unclear if the methodology can transfer to other model sizes/families.", "other_comments_or_suggestions": "TLDR: put more ablation results in the main text instead of leaving them in the appendix\n\nWith a lot of design choices and experiments covered in the paper, I think the paper lacks a clear outline of the core research question/methodology it is investigating and the takeaways from the experiments. As the paper proposes a new pruning methodology for a fairly well-known pipeline, the focus should be on the ablation of various method design choices. I saw the author had a good amount of ablation in the appendix and the author should select some of them to be put in the main text along with relevant discussion to provide the reader more insights on what is important for pruning the model.", "questions_for_authors": "N/A", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "jRW3LIas01", "reviewer": "Reviewer_PN1H", "summary": "The paper introduces Puzzle, a hardware-aware framework that optimizes LLM inference efficiency using neural architecture search (NAS), blockwise local knowledge distillation (BLD), and mixed-integer programming. The authors demonstrate its effectiveness with Puzzle-51B, a 51B-parameter model derived from Llama-3.1-70B, achieving 2.17× inference speedup on a single H100 GPU while retaining 98.4% accuracy despite training on only 45B tokens.", "questions": "", "limitations": "", "rating": 2, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your constructive feedback,\n\n*\"no compression results for a 7B model\"*\n*\"experiments only focus on compressing a 70B model to 51B\"*\n\nWe demonstrate Puzzle's robustness by applying it 11 times with varied constraints, datasets, and budgets:\n\n(1) 4 derivatives of Llama-3.1-70B (including Puzzle-51B).\n\n(2) 6 derivatives of Llama-3.3-8B.\n\n(3) A derivative of Llama-3.3-70B: \"Puzzle-49B\" (Sec. 5).\n\nMoreover, post-submission, we applied Puzzle in two additional scenarios:\n\n(4) A 253B derivative of Llama-3.1-405B, constrained for a single H100 node at 1.5X latency, retaining 99.5\\% of parent performance (benchmarks: MMLU, MT-Bench, MMLU-Pro, HumanEval, Arena Hard).\n\n(5) A novel 50B+ Mamba-hybrid derivative, constrained for RTX 5090 with 1M context length, retaining 99.94\\% parent performance (benchmarks: MMLU, MMLU-Pro, GSM8K, HellaSwag).\n\nWe'll include these in the revision to highlight Puzzle's practical value.\n\n*\"Although it attains a 2x inference speedup, the performance still falls short of the original level even after additional training.\"*\n\nWe believe a 98\\% accuracy retention is impressive. Moreover, while Puzzle-49B already retains high accuracy, we show that a lightweight alignment phase further boosts its performance, leading it to outperform its parent model, Llama-3.3-70B, at 105.5\\% relative accuracy.\n\n*\"The additional training cost further diminishes the trade-off, making it less appealing for real-world applications.\"*\n\nIndeed 45B tokens might be expensive for some users, even if it is much lower than the trillions of tokens necessary to train LLMs from scratch. However:\n\n(1) GKD is meant to squeeze extra performance when budget allows. Puzzle derivatives remain competitive even without GKD (Table 14, Appendix F.3), retaining 96.5\\% or 90\\% parent performance without GKD.\n\n(2) While we didn't mention it in the paper, even a *significantly* shorter GKD training may suffice for a substantial increase in performance. After 3.7B tokens, Puzzle-51B reached 98.8\\% parent performance on MMLU and MT-Bench (MMLU 79.38, MT-Bench 8.96). Puzzle-49B, after 8.68B tokens GKD (pre long-context KD in Sec. 5), reached 99.63\\% parent performance (80.73 MMLU, MT-Bench 8.87). Even just 2.9B tokens GKD recovered 98.47\\% for Puzzle-49B (MMLU 80.72, MT-Bench 8.675).\n\nWe agree that it is important to include these results in the revision to clarify how GKD length can be adjusted based on the available budget.\n\n*\"compresses a 70B model to 51B, achieving a very limited compression rate.\"*\n\nAs noted in the paper, we argue that categorizing models by parameter count alone (50B vs. 70B) is less meaningful. Real-world choices should depend on hardware, budget constraints, and usage profiles (sequence length, batch size). Hypothetically, a good and resource-efficient 80B model that is faster than an 8B model is preferable to it. We note that even for parameter reduction alone, our accuracy retention remains significant.\n\n*\"comparisons with other compression methods, such as low-rank approximation\"*\n\nThe Puzzle framework is complementary to techniques such as structured sparsity and low-rank approximation, and can be incorporated within its search space. Thus, these techniques enhance rather than compete with Puzzle.\n\nStill, we agree comparisons with methods like [1] and [2] would indeed benefit the paper. We've initiated such evaluations for the revision. Instead of [1], we've chosen Wanda [3], a newer structured sparsity method with good results, which we hope you find acceptable.\n\nBelow are preliminary results comparing Puzzle, Wanda, and low-rank approximation. Wanda pruned Llama-3.1-70B (2:4 structured sparsity) targeting similar speedups as Puzzle-51B. The low-rank approximation resembles [4], with subsequent distillation.\n\n| Model             | MMLU  | MT-Bench | Average Accuracy | Accuracy Preserved |\n|-------------------|-------|----------|------------------|--------------------:|\n| Puzzle-51B       | 80.20 | 8.99     | 85.05            | 99.49              |\n| Wanda            | 72.99 | 8.39     | 78.44            | 92.23              |\n| Low-rank         | 72.87 | 8.01     | 76.05            | 88.96              |\n| Llama-3.1-70B    | 81.66 | 8.93     | 85.48            | 100                |\n\nFor Wanda [3], which doesn't include additional training, distillation post-pruning yielded marginal gains (MMLU 73.69; MT-Bench unchanged).\nWe are also working to evaluate Sheared Llama [2] as you suggested.\n\nFinally, we're exploring integrating these methods within Puzzle, demonstrating their complementary strength, which we'll aim to present clearly in the revision.\n\n[3] Sun et al. A Simple and Effective Pruning Approach for Large Language Models, ICLR\n\n[4] Khodak et al. Initialization and Regularization of Factorized Neural Layers, ICLR\n\n*\"The distillation-based approach...lacks innovation and new insights...\"*\n\nWe respectfully disagree; limited space prevents elaboration but we are happy to clarify if needed.", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "This paper compresses a 70B model to 51B, achieving a very limited compression rate. Although it attains a 2x inference speedup, the performance still falls short of the original level even after additional training. Furthermore, there are no compression results for a 7B model, significantly limiting its practical value.", "theoretical_claims": "None", "experimental_designs_or_analyses": "Although comparisons were conducted on some common benchmarks, the advantages of the proposed method are not clearly demonstrated. The paper lacks comparisons with widely-used compression techniques, such as SparseGPT[1] and Sheared LLaMA[2], which limits the ability to assess its relative effectiveness and innovation. \n\n[1]Frantar, Elias, and Dan Alistarh. \"Sparsegpt: Massive language models can be accurately pruned in one-shot.\" International Conference on Machine Learning. PMLR, 2023.  \n[2]Xia, Mengzhou, et al. \"Sheared LLaMA: Accelerating Language Model Pre-training via Structured Pruning.\" The Twelfth International Conference on Learning Representations.", "relation_to_broader_scientific_literature": "None", "essential_references_not_discussed": "[1]Frantar, Elias, and Dan Alistarh. \"Sparsegpt: Massive language models can be accurately pruned in one-shot.\" International Conference on Machine Learning. PMLR, 2023.\n\n[2]Xia, Mengzhou, et al. \"Sheared LLaMA: Accelerating Language Model Pre-training via Structured Pruning.\" The Twelfth International Conference on Learning Representations.", "other_strengths_and_weaknesses": "**Weaknesses**:\n1. The distillation-based approach for model compression lacks innovation and new insights, as it relies on well-established techniques without introducing novel methodologies or deeper understanding.  \n\n2. The experiments only focus on compressing a 70B model to 51B, which has limited practical value. The additional training cost further diminishes the trade-off, making it less appealing for real-world applications.  \n\n3. The paper lacks discussion and comparisons with other compression methods, such as **low-rank approximation** and **unstructured/structured sparsity**, which are critical for evaluating the proposed method's competitiveness and effectiveness in the broader context of model compression.", "other_comments_or_suggestions": "None", "questions_for_authors": "See Weaknesses", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Ar011tYrct", "reviewer": "Reviewer_Fna8", "summary": "The paper introduces Puzzle, a distillation-based NAS approach for extracting inference-optimized LLMs for existing trained models such as Llama. The authors first introduce the search space for their NAS-based optimizations - including the different attention and FFN subblocks to use, followed by the number of combinations, and introduce a decoupled block distillation algorithm to reduce the number of total combinations to explore for search. Following this, they define their overall algorithm for Puzzle:\n- Search constraints focused on real-time memory and throughput on H100 GPUs\n- Scoring method for each subblock\n- Mixed Integer Programming for selecting the best solution with their constraints\n- Additional global training to improve overall model quality \n\nFor results, the authors present the 51B model, which is distilled from Llama 70B, including evaluations on evaluation benchmarks, blind test comparisons, training a 49B model with similar memory requirements but with longer context length support, and a final smaller model distilled from the 8b model.\n\n-----------\n\n### After Rebuttal\nBased on the provided rebuttal and the other reviews / addressed comments, I've decided to retain my score. All concerns have been addressed.", "questions": "", "limitations": "", "rating": 5, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your positive feedback!,\n\n*\"For the throughput comparison of table 2, can the authors clarify what batch sizes they used for the numbers?\"*\n\nThis is a good question. For every model and hardware setting we selected the optimal batch size to get the best throughput per GPU. This was done automatically by the inference engine, selecting for each run the optimal batch size to use. For example, Puzzle-51B's optimal batch size for TP=1 was 256, and for Llama-3.1-70B at TP=4, it was 384. We will make sure to note this in our revision. Do you think it is beneficial to include the full list of batch sizes as well?\n\n*\"For the same table, when using TP=1 for the 128/1024 scenario (entry 2) - the slow throughput for Llama-70B almost seems to be an issue of hitting the memory + optimization wall very quickly with an H100 80GB instance rather than the model not being \"efficient\" here. Can the authors comment on this? If so, comparing with say TP=2 might have been a better comparison point vs TP=4 here?\"*\n\nYou are right about TP=1. As described in the caption, for each model the optimal TP was chosen for speed measurements, which is why we selected TP=4 and not TP=2 for Llama-70B: TP=4 is more flattering to the parent's throughput for this scenario.\nWe also included the speed measurements on TP=1 to also present a \"fair\" comparison to Puzzle-51B for a single GPU setting. We will clarify this point explicitly in the revision.\n\n*\"For Figure 5, can the authors clarify if they used MMLU or MMLU Chat for the accuracy computations?\"*\n\nWe used MMLU (and not MMLU Chat).\n\n*\"This is similar to other work done previously to find smaller networks such as Minitron [1], but also other pruning approaches such as ShortGPT [2] and SlimGPT [3]. However, none of these approaches are explicitly compared in the paper\"*\n\nYou make a good point, and we are working on including these relevant comparisons in the revision.\nIn short, some of these methods share similar aspects to Puzzle, but their solutions remain a subset of Puzzle's search space. For example, Minitron only considers homogeneous solutions (i.e., any modified block is applied across all layers), whereas Puzzle allows for heterogeneous, layer-specific configurations.\nShortGPT considers layer removal (similar to how Puzzle allows \"no-op\" layers), using a cosine similarity score to identify redundant layers.\nSlimGPT introduces an extension of the \"Optimal Brain Surgeon\" method, called \"Batched Greedy Pruning\", that could also be used by Puzzle to prune individual blocks. SlimGPT sets an \"Incremental Pruning Ratio\" heuristic that follows a fixed logarithmic curve. Puzzle, on the other hand, can consider the pruning ratio for each layer in a custom manner, which could also consider assigning higher pruning ratios to later layers like SlimGPT does.\n\nFinally, thank you for mentioning LayerNAS, we will also include it as a reference to related literature in the revision.", "claims_and_evidence": "Yes, the claims made in the submission have sufficient evidence for each aspect. Just highlighting some of them below:\n\n- The authors present a series of ablations for aspects of the algorithm, for example, the scoring method for each sub-block (using KL vs cross entropy vs accuracy on downstream tasks), which datasets they used for their BLD, and how long to train each sub-block for during exploration. The ablations are empirically supported by results from reasonable downstream evaluations, measuring throughput, etc. \n- For the mixed integer programming algorithm, they add a diversity constraint, and show graphs representing how different blocks have different sub-blocks selected (see Fig. 6 for example).", "methods_and_evaluation_criteria": "Yes, most of the methods and the equivalent evaluation criteria make sense for the problem the authors are solving.", "theoretical_claims": "N/A", "experimental_designs_or_analyses": "Yes, I did check the design and associated analysis. Here are some follow up questions:\n\n- For the throughput comparison of table 2, can the authors clarify what batch sizes they used for the numbers?\n- For the same table, when using TP=1 for the 128/1024 scenario (entry 2) - the slow throughput for Llama-70B almost seems to be an issue of hitting the memory + optimization wall very quickly with an H100 80GB instance rather than the model not being \"efficient\" here. Can the authors comment on this? If so, comparing with say TP=2 might have been a better comparison point vs TP=4 here? \n- For Figure 5, can the authors clarify if they used MMLU or MMLU Chat for the accuracy computations?", "relation_to_broader_scientific_literature": "The paper presents a NAS+distialltion based approach to get best in class LLMs for training. Given the complexity of using NAS for LLMs, the authors present new approaches to reduce complexity for NAS, especially dealing with sub-block optimizations that are initially greedy (optimizing for that block w/ inputs only from previous blocks). This is similar to approaches like LayerNAS [1] which have been previously explored for CV architectures. \n\n\n[1] LayerNAS: https://arxiv.org/abs/2304.11517", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "**Strengths**\n- The paper does thorough ablations of each decision choice in the creation of the algorithm\n- For the NAS objective to minimize, the authors consider real-world deployments over theoretical measures such as parameter count or FLOPs\n- The approach identifies potential algorithm pieces that may introduce extra complexity (such as the large search space) and focuses on solving those through feedback from ablations. \n\n\n**Weakness**\nThe paper discusses an approach to find smaller networks from an existing trained LLM using the suggested NAS-based approach for efficiency. This is similar to other work done previously to find smaller networks such as Minitron [1], but also other pruning approaches such as ShortGPT [2] and SlimGPT [3]. However, none of these approaches are explicitly compared in the paper, beyond trying to recover the accuracy of the original trained model. \n\n[1] Minitron: https://arxiv.org/abs/2408.11796\n[2] ShortGPT: https://arxiv.org/abs/2403.03853\n[3] SlimGPT: https://arxiv.org/abs/2412.18110", "other_comments_or_suggestions": "N/A", "questions_for_authors": "N/A", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {"ethical_review_concerns": "N/A"}}]

Author rebuttal:

Meta-review: This paper presents Puzzle, a distillation-based NAS framework for optimizing LLMs for efficient inference. It achieves strong empirical results, including a 2.17× throughput speedup on an H100 GPU with minimal accuracy loss. The method is well-engineered, robust across multiple model sizes, and grounded in practical deployment constraints.

While some reviewers noted limited theoretical novelty and missing comparisons, the authors provided thorough rebuttals, added key baselines, and committed to improving the revision. Given the overall strength of the work and its practical relevance, it is a solid contribution to the field.
Decision: Accept (poster)