PaperlayerView exact revision on arXiv

Source-linked AI summary

Puzzle: Distillation-Based NAS for Inference-Optimized LLMs

Akhiad Bercovich, Tomer Ronen, Talor Abramovich, Nir Ailon, Nave Assaf, Mohammad Dabbah, Ido Galil, Amnon Geifman, Yonatan Geifman, Izhak Golan, Netanel Haber, Ehud 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

arXiv:2411.19146v5cs.LG

TL;DR

High inference costs and the expense of evaluating many LLM architectures constrain practical deployment. Puzzle combines blockwise distillation and hardware-aware architecture search; its Nemotron derivatives target single-H100 inference with substantially less training than the parent model. The paper argues that model choice should reflect inference conditions, not parameter count alone.

  • Problem

    High inference costs restrict LLM accessibility and scalability, and evaluating many candidates through full-model distillation is prohibitively costly.

    Source: §1 · pp. 1–2

  • Method

    Puzzle trains block alternatives in parallel, scores them, selects hardware-constrained architectures with mixed-integer programming, and applies global distillation to the reassembled model.

    Source: §1 · p. 2

  • Results

    Nemotron-51B and Nemotron-49B are optimized for single-H100 inference, and Puzzle’s Llama-3.1-70B-Instruct run used under 50B tokens versus over 15T to train the parent.

    Source: §1 · pp. 2–3

  • Takeaways & Limitations

    Model selection should account for hardware, inference engine, quantization, budget, and usage profile rather than parameter count alone.

    Source: §1 · p. 2

    Blockwise local distillation does not ensure compatibility between blocks, so errors may propagate and compound; global distillation is added to mitigate this.

    Source: §3 · p. 5

Abstract

from arXiv · show

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.

1 Introduction

Puzzle addresses the high inference costs and costly candidate evaluation that limit LLM deployment by searching for hardware- and scenario-specific architectures. Its distillation and optimization pipeline produces Nemotron models for single-H100 inference using far fewer training tokens than the parent model.

1 Introduction

  • High inference costs restrict LLM accessibility and scalability, while full-model distillation for many candidates in a large search space remains prohibitively costly.Puzzle searches after blockwise distillation and reserves global distillation for the reassembled final architecture.Source: §1 · pp. 1–3
  • Puzzle trains block variants in parallel, scores replacement options, uses mixed-integer programming to select architectures under inference constraints, then applies global knowledge distillation.The resulting models can use non-uniform layer computation allocation.Source: §1 · p. 2
  • Under 50B tokens versus over 15T for parent training, Puzzle produces Nemotron-51B and Nemotron-49B optimized for single-H100 inference.The under-50B-token figure is for the full Puzzle training run on Llama-3.1-70B-Instruct.Source: §1 · pp. 2–3
  • Puzzle is designed to create multiple child models from one parent for different hardware and inference requirements, using only the parent’s weights rather than its training data.The framework supports constraints including throughput, latency, and memory usage.Source: §1 · pp. 2–3
  • Nemotron-49B preserves its parent’s 128K context via short uptraining and, with RLHF alignment, outperforms Nemotron-51B on some tasks.Puzzle is also applied with varied constraints, datasets, budgets, and parent models.Source: §1 · p. 3

2 Search Space

Puzzle searches layer-by-layer combinations of attention and feed-forward alternatives, including reduced or skipped operations. These choices trade representational power against memory and computation, while the enormous resulting architecture space makes exhaustive evaluation impractical.

2 Search Space

  • Attention choices include MHA, GQA with varying key-value head counts, a linear layer, or a no-op; FFN choices include reduced dimensions, linear layers, or a no-op.The search space pairs one attention and one FFN option for each parent layer.Source: §2 · p. 3
  • Reducing key-value heads can lower attention computation and KV-cache memory, while reduced FFN dimensions or linear layers lower computation and full dimensions retain greater representational power.Lower KV-cache use can help meet memory constraints or enable larger batches.Source: §2 · p. 4
  • All child subblocks in this work retain their parent layer’s input and output dimensions, although varying embedding dimensions are identified as a possible future scheme.An LLM block typically means a transformer layer composed of attention and FFN subblocks.Source: §2 · pp. 3–4
  • Each parent transformer layer pairs an attention option with an FFN option, creating a large architecture space whose representative subset is computationally infeasible to evaluate.For Llama 3.1-70B, the described setup has 80 layers and 54 possible configurations per layer.Source: §2 · p. 4

3 Blockwise Local Distillation

Blockwise local distillation trains child blocks independently to mimic parent blocks, while decoupled BLD reduces the cost of constructing a library of alternatives. Training-free initialization techniques support reduced or replaced subblocks, and global distillation addresses BLD’s block-compatibility limitation.

3 Blockwise Local Distillation

  • BLD trains each child block independently to mimic its corresponding parent block, using parent activations transferred between layers.Independent training permits parallelism across GPUs and focuses each child subblock on mimicking a single parent subblock.Source: §3 · pp. 4–5, Figure 2 · p. 4
  • Child-block optimization uses normalized MSE between the outputs of the original parent block and the modified child block.The inputs to the current block are parent activations.Source: §3 · p. 5
  • BLD does not ensure compatibility between blocks because each is trained on inputs from preceding parent blocks, so errors may propagate through the child model.The framework adds global knowledge distillation as a final training phase to mitigate this limitation.Source: §3 · p. 5

3.2 Initialization of Alternative Subblocks

  • Channel Contribution ranks FFN intermediate channels by the output change caused by pruning, prioritizing the lowest-contribution channels during initialization.Per-token channel contributions are averaged across tokens in the calibration dataset.Source: §3.2 · pp. 5–6

3.1 Building a Block Library with Decoupled Blockwise Local Distillation

  • For m = n = 20 across 80 layers, decoupled BLD trains 3,200 blocks versus 32,000 with coupled BLD.The library-construction comparison is illustrated in Figure 3.Source: §3.1 · p. 5, Figure 3 · p. 6

3.2 Initialization of Alternative Subblocks

  • For linear-layer initialization, the method combines FFN up- and down-projection matrices, while attention variants use mean-pooled key/value projections or value/output products.The FFN initialization ignores gating; the attention linear-layer initialization simulates each token attending only to itself.Source: §3.2 · p. 6

4 Decomposed NAS Search Algorithm for LLMs

Puzzle builds a block library, estimates block costs and quality, and searches for architectures under deployment constraints. It measures resources on target hardware, scores variants by single-block replacements, and uses mixed-integer programming to select architectures.

4 Decomposed NAS Search Algorithm for LLMs

  • Puzzle’s decomposed NAS workflow builds a block library, estimates resource needs, scores block variants by location, and searches for architectures under constraints.The framework adapts these NAS steps to the scale of multi-billion-parameter LLMs.Source: §4 · p. 6

4.1 Estimating Resource Requirements

  • Resource requirements are measured on target hardware across sequence lengths, batch sizes, and prefill and generation phases.These measurements determine the constraints used in the MIP optimization.Source: §4.1 · p. 7

4.2 Scoring Architecture Solutions

  • Replace-1-block scores measure whole-model performance after replacing one parent block, then estimate candidate quality by summing the selected blocks’ scores.Each block variant is scored once, rather than evaluating candidate models during search.Source: §4.2 · p. 7
  • KL divergence is highlighted as an effective block-scoring metric for decomposed NAS.The paper compares downstream accuracy, language-model loss, and KL divergence as potential scores.Source: §4.2 · pp. 7–8

4.3 Search Algorithm: Mixed-integer Programming

  • The MIP search maximizes selected block quality scores while meeting memory, throughput, and latency constraints and selecting exactly one variant per layer.The formulation treats layers as groups of block alternatives with quality scores and resource costs.Source: §4.3 · pp. 8–9
  • A diversity constraint limits similarity to prior solutions, while repeated solves at different batch sizes explore runtime-memory trade-offs.After block libraries and scores are built, Puzzle can generate architectures for different deployment scenarios without additional training.Source: §4.3 · p. 9

5 Post-Puzzle Inter-Block Uptraining

Post-BLD uptraining uses selected losses to address a mismatch between independently trained blocks and the assembled child model. The reported ablation finds that LM loss is unnecessary for optimal uptraining and often harms downstream tasks.

5 Post-Puzzle Inter-Block Uptraining

  • End-to-end global knowledge distillation (GKD) addresses BLD’s block-input mismatch by training the assembled child against the parent’s outputs or representations.During BLD, blocks receive parent-model predecessor outputs; after assembly, their input distribution can differ, motivating short GKD training.Source: §5 · p. 9
  • The evaluated loss components include supervised LM cross-entropy, cosine similarity between parent and child hidden states, and KLD alignment of their logits.The KLD distributions come from parent and child logits, respectively.Source: §5 · pp. 9–10
  • The cosine-loss notation compares child and parent hidden states by transformer layer, but one-to-one teacher–student block correspondence is not required.The paper assumes one-to-one correspondence for the notation while stating that the framework does not require it.Source: §5 · p. 10
  • Loss-combination experiments used approximately 5B tokens on post-BLD Nemotron-51B, while the final model used 45B tokens with the highest-scoring combination.The ablation table identifies the models as Nemotron-51B derivatives of Llama-3.1-70B-Instruct.Source: §5 · p. 10, Table 1 · p. 10
  • LM loss was unnecessary for optimal uptraining and harmed downstream tasks in most cases.Table 1 reports the loss-combination ablation for Nemotron-51B; its models were generally trained for approximately 5B tokens.Source: §5 · p. 10, Table 1 · p. 10

6 Supporting Fast Inference for Variable-Block Architectures in TensorRT-LLM

TensorRT-LLM supports fast LLM inference on NVIDIA GPUs. The authors revised its paged KV-cache strategy and runtime support so Puzzle’s variable-block architectures can run with FP8 precision.

6 Supporting Fast Inference for Variable-Block Architectures in TensorRT-LLM

  • TensorRT-LLM is an optimized NVIDIA GPU runtime for accelerating LLM inference.It supports custom paged-attention kernels for KV caching, FP8 quantization, and scheduling policies.Source: §6 · p. 10
  • Revising the paged KV-cache strategy enabled variable GQA ratios, removing the assumption that all attention layers have the same number of key and value heads.The revised runtime supports architectures from Puzzle’s search space.Source: §6 · p. 11
  • TensorRT-LLM supports Puzzle architectures with variable blocks and linear replacements, using FP8 for weights, activations, and KV cache.Source: §6 · p. 11

7 Main Results

Nemotron-51B retains most of its parent’s benchmark accuracy while improving inference throughput. Additional evaluations show comparable human preference, a favorable accuracy-throughput tradeoff, and derivatives spanning longer contexts and consumer-grade hardware.

7 Main Results

  • Blind human evaluation found comparable performance between Nemotron-51B and Llama-3.1-70B-Instruct on general-purpose tasks.The test included reasoning, longform generation, and knowledge tasks.Source: §7 · p. 11, Figure 4 · p. 12
  • 2.17x maximum throughput speedup and 98.4% average parent-accuracy retention characterize Nemotron-51B’s main results.Throughput is compared across scenarios on H100 GPUs; Table 2 compares accuracy across benchmarks.Source: §7 · p. 11, Table 2 · p. 11, Table 3 · p. 12
  • Nemotron-51B is reported to push beyond the accuracy-throughput frontier and provide the best accuracy-per-dollar tradeoff.Accuracy combines MT-Bench and MMLU as (MT-Bench × 10 + MMLU) / 2.Source: §7 · p. 12, Figure 5 · p. 13
  • Runtime analysis identifies savings in early and late layers while retaining full computation in layers 16–42.Early layers reduce attention and FFN computation; later layers apply more aggressive attention optimization.Source: §7 · p. 12, Figure 6 · p. 14

7.1 Additional Evaluations and Puzzle derivatives

  • After training only on sequences up to 8K tokens, Nemotron-51B retained over 96% of parent accuracy at 16K tokens, with performance degrading beyond 64K.The authors suggest longer-context fine-tuning could extend effective context length.Source: §7.1 · p. 12
  • Nemotron-49B-Base was uptrained with additional 5B-token stages at 64K and 128K context, then aligned using 25M tokens.The alignment comparison reports that Puzzle derivatives can undergo RLHF-based alignment to boost accuracy.Source: §7.1 · pp. 12–13, Table 5 · p. 13
  • A Puzzle derivative of Llama-3.1-8B-Instruct achieved a superior accuracy-efficiency tradeoff on an RTX 4090 consumer GPU.Table 6 compares its accuracy and throughput with other models, including Llama-3.2-3B-Instruct.Source: §7.1 · pp. 13–14, Table 6 · p. 14
  • Other derivatives include a 405B model retaining 99.5% of parent accuracy and a 47B miniPuzzle model retaining 99.94%; reduced GKD budgets also recovered substantial accuracy.Nemotron-51B recovered 98.8% of parent accuracy on MMLU and MT-Bench after 3.7B GKD tokens; Nemotron-49B recovered 99.63% after 8.68B tokens and 98.47% after 2.9B.Source: §7.1 · p. 14

8 In-Depth Analysis and Ablation Studies

The ablations examine how block training, data, scoring, and search choices shape Puzzle architectures. Results favor diverse block alternatives and globally optimized, quality-aware selection, while showing tradeoffs in training cost, task specialization, and constraints.

8.1.1 Combining Coupled BLD and Decoupled BLD

  • A two-stage decoupled-then-coupled BLD strategy narrows training from a multiplicative search space, then improves accuracy through refinement.For the 8B model, decoupled training used 448 subblocks; coupled refinement on a narrowed space used 384 blocks.Source: §8.1.1 · p. 15, Table 8 · p. 15

8.1.2 Impact of Dataset Composition on Puzzle-Derived Models

  • Gutenberg-trained models achieved 92.7% and 95.5% of Distillation Mix performance on MT-Bench and MMLU, respectively, without GKD uptraining.They retained 91.7% on MMLU STEM categories (64.5 vs 70.35), despite the corpus’s limited domain coverage.Source: §8.1.2 · p. 16, Table 9 · p. 16

8.1.3 Impact of Blockwise Local Distillation Training Dataset Size

  • Larger BLD token budgets improved MT-Bench more than MMLU, with diminishing returns as training increased from 0.5B to 1.0B tokens.The gains were particularly pronounced in multi-turn conversational tasks.Source: §8.1.3 · p. 16, Table 10 · p. 16

8.1.4 Impact of Different Block Scoring Metrics

  • KL-divergence scoring produced better general architecture choices than LM-loss scoring, while task-specific scoring favored its target at an MT-Bench cost.Half-MMLU scoring achieved 7.57 MT-Bench accuracy, versus 8.06 for KL-divergence scoring.Source: §8.1.4 · pp. 16–17, Figure 7 · p. 17, Table 11 · p. 17

8.1.5 Effects of Limited Search Space Diversity

  • Restricting alternatives to no-op blocks retained throughput but reduced MMLU accuracy to 75.4 versus 78.39 for the more diverse search.The restricted design removes the need for BLD but sacrifices fine-grained architectural customization.Source: §8.1.5 · pp. 17–18, Table 12 · p. 18

8.2.1 MIP Solution Analysis Under Varying Throughput Constraints

  • MIP adapts computation unevenly across layers and components: attention may be skipped, while FFN components remain, even under strict throughput constraints.The heatmaps also show local heavier-block choices at higher throughput targets, with other layers compensating globally.Source: §8.2.1 · pp. 18–19, §8.2 · p. 18, Figure 8 · p. 19

8.2.2 Greedy Search Algorithm Alternative; 8.2.3 Importance of Data-Driven Quality Estimation in Architecture Scoring

  • At identical throughput constraints, greedy search scored 70.74% MMLU versus 78.39% for MIP; parameter-count maximization also sharply reduced accuracy.The latter result supports quality-aware, layer-specific scoring over selecting uniformly high-parameter blocks.Source: §8.2.2 · pp. 19–20, Table 13 · p. 20, §8.2.3 · p. 20

8.2.4 Evaluating Random Architectures as Search Baselines; 8.3 Global Knowledge Distillation Uptraining Importance; 8.4 Comparison with Related Work

  • Puzzle beats random baselines, Wanda, and low-rank methods; its heterogeneous global search generalizes constrained related methods, while GKD improves both benchmarks.The random-from-library baseline achieved 86.6% of Nemotron-51B’s performance; GKD improved MMLU and MT-Bench scores.Source: §8.2 · p. 18, §8.2.4 · p. 20, §8.4 · p. 21

9 Conclusions and Future Directions

Puzzle modifies over-parameterized LLMs into hardware-tailored inference-efficient architectures using fewer than 50B tokens. The authors identify decoupled BLD as a basis for broader architectural and capability-specific exploration.

9 Conclusions and Future Directions

  • Fewer than 50B tokens let Puzzle produce high-performing inference-optimized models, compared with the trillions needed to train models from scratch.The conclusion describes these models as achieving drastic inference performance improvements.Source: §9 · p. 22
  • Decoupled BLD is significantly more efficient than coupled BLD, making it feasible to evaluate more potential blocks in one optimization run.Future block options include variable-window attention [9] and state-space models [18] [17].Source: §9 · p. 22
  • Puzzle could be extended to optimize models for Chain-of-Thought reasoning, vision-language models [38], and retrieval-augmented generation [33].Source: §9 · p. 22

A.1 Human Evaluation

The human evaluation used 169 samples across a curated set of writing, long-input, instruction-following, knowledge, summarization, reasoning, and extraction tasks. Three annotators independently compared randomized completion pairs using four response options.

A.1 Human Evaluation

  • 169 samples received three independent annotations each, yielding 507 annotations.Source: §A.1 · p. 27
  • Annotators compared a prompt and two completions, choosing which completion was better, whether both were good, or whether neither was good; completion order was randomized.Source: §A.1 · p. 27
  • The curated test set covered long-form and grounded text generation, long-input tasks, multi-condition instructions, and knowledge and trivia.Source: §A.1 · pp. 27–28

B RULER Benchmark Performance Tables

This section provides parent–child performance tables for a subset of the RULER benchmark across the evaluated context lengths.

B RULER Benchmark Performance Tables

B.1 Nemotron-51B vs. Llama-3.1-70B-Instruct

Nemotron-51B retains more than 96% of its parent’s performance at 16K context despite training only up to 8K tokens. Its performance degrades at 64K and beyond.

B.1 Nemotron-51B vs. Llama-3.1-70B-Instruct

  • More than 96% of Llama-3.1-70B-Instruct’s performance was retained by Nemotron-51B at 16K context after training on sequences up to only 8K tokens.Source: §B.1 · p. 28
  • Nemotron-51B’s performance degrades at 64K context and beyond, consistent with its training horizon of up to 8K tokens.Source: §B.1 · p. 28
  • The results indicate that a large fraction of the parent’s long-context capabilities can remain intact without explicit training on such long sequences.Source: §B.1 · p. 28

B.2 Nemotron-49B-Base vs. Llama-3.3-70B-Instruct

Nemotron-49B-Base preserves at least 98% of its parent’s performance through 64K tokens and remains above 94% at 128K, extending the context range of a Puzzle-derived model.

B.2 Nemotron-49B-Base vs. Llama-3.3-70B-Instruct

  • At 128K tokens, Nemotron-49B-Base remains above 94% of its parent’s performance.The model was uptrained with sequences up to 128K tokens.Source: §B.2 · p. 28
  • A short uptraining phase on longer contexts can effectively extend the context range of Puzzle-derived models.The reported comparison is for Nemotron-49B-Base against Llama-3.3-70B-Instruct.Source: §B.2 · p. 28
  • Nemotron-49B-Base preserves 98% or more of Llama-3.3-70B-Instruct’s performance up to 64K tokens.Table 19 compares the parent and child on a subset of the RULER benchmark.Source: §B.2 · p. 28, Table 19 · p. 29

Related papers

Loading 2411.19146v5…