Puzzle: Distillation-Based NAS for Inference-Optimized LLMs

Thank you for your positive feedback!,

"For the throughput comparison of table 2, can the authors clarify what batch sizes they used for the numbers?"

This 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?

"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?"

You 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. We 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.

"For Figure 5, can the authors clarify if they used MMLU or MMLU Chat for the accuracy computations?"

We used MMLU (and not MMLU Chat).

"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"

You make a good point, and we are working on including these relevant comparisons in the revision. In 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. ShortGPT considers layer removal (similar to how Puzzle allows "no-op" layers), using a cosine similarity score to identify redundant layers. SlimGPT 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.

Finally, thank you for mentioning LayerNAS, we will also include it as a reference to related literature in the revision.

Thank you for your constructive feedback,

"no compression results for a 7B model" "experiments only focus on compressing a 70B model to 51B"

We demonstrate Puzzle's robustness by applying it 11 times with varied constraints, datasets, and budgets:

(1) 4 derivatives of Llama-3.1-70B (including Puzzle-51B).

(2) 6 derivatives of Llama-3.3-8B.

(3) A derivative of Llama-3.3-70B: "Puzzle-49B" (Sec. 5).

Moreover, post-submission, we applied Puzzle in two additional scenarios:

(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).

(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).

We'll include these in the revision to highlight Puzzle's practical value.

"Although it attains a 2x inference speedup, the performance still falls short of the original level even after additional training."

We 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.

"The additional training cost further diminishes the trade-off, making it less appealing for real-world applications."

Indeed 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:

(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.

(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).

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.

"compresses a 70B model to 51B, achieving a very limited compression rate."

As 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.

"comparisons with other compression methods, such as low-rank approximation"

The 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.

Still, 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.

Below 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.

For Wanda [3], which doesn't include additional training, distillation post-pruning yielded marginal gains (MMLU 73.69; MT-Bench unchanged). We are also working to evaluate Sheared Llama [2] as you suggested.

Finally, we're exploring integrating these methods within Puzzle, demonstrating their complementary strength, which we'll aim to present clearly in the revision.

[3] Sun et al. A Simple and Effective Pruning Approach for Large Language Models, ICLR

[4] Khodak et al. Initialization and Regularization of Factorized Neural Layers, ICLR

"The distillation-based approach...lacks innovation and new insights..."

We respectfully disagree; limited space prevents elaboration but we are happy to clarify if needed.

Thank you for your feedback and appreciation of the ablation studies,

"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"

Appendix 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. We 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).

"the author only evaluates the method on one model (llama-3.1-70B), which raises concern on its generalizability."

The 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:

(1) 4 derivatives of Llama-3.1-70B: a) Puzzle-51B, b) 2 other 51B derivatives with a different BLD token budget (Appendix F.1.3., Table 9), c) A Puzzle derivative from a different dataset (Gutenberg dataset, Appendix F.1.2., Table 8), d) A limited search space variant (Appendix F.1.5., Table 11)

(2) 6 derivatives of Llama-3.3-8B: a) A "coupled BLD" derivative (appearing both in the main paper at Table 5 and in Appendix F.1.4., Table 7), b) 2 derivatives with LM loss block scoring (Figure 7 in Appendix F.1.4.), c) 2 derivatives with decoupled BLD (Figure 7 in Appendix F.1.4.), d) A derivative with a downstream, "Half-MMLU" block score (Table 10 in Appendix F.1.4.)

(3) A derivative of Llama-3.3-70B: "Puzzle-49B" (Section 5).

We 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:

(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).

(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).

We will make sure the revision includes these examples to emphasize Puzzle's robustness across a variety of constraints and models.

"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."

We 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:

(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).

(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.

"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"

Thank 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.

In 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?

Thank you for your feedback and suggestion,

"Key baselines are missing, which 1) uses distillation on a fully random architecture 2) uses distillation on a random-from-block-library architecture."

We 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. Additionally, 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. To ensure fairness, we allocated the same 10B token budget for training each of these models:

We will include these baselines in the revision.

We 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.