Source-linked AI summary
Efficient Knowledge Distillation for LLMs: Offline Top-K Logits and a Fused Chunked KL Loss
Bakbergen Ryskulov, Iker García-Ferrero, David Montero, David Jansen, Ali Hashemi, Jezabel R. Garcia, Antonio Tiene, Román Orús
TL;DR
Knowledge distillation is costly even though it is central to recovering compact language models for constrained deployment. This paper studies cached top-K teacher targets and a fused chunked KL loss, finding comparable online quality with lower memory use and efficient long-context training.
Problem
Knowledge distillation recovers compact models for constrained deployment, but this recovery step is costly and under-documented relative to its practical impact.
Method
The paper caches top-K teacher logits offline and applies a fused, chunked KL objective whose peak memory is linear in sequence length.
Results
Offline distillation matches online quality, while the chunked KL implementation stays within 11.6 GiB per GPU at 256K tokens in a controlled loss-only benchmark.
Takeaways & Limitations
Together, these systems choices make large-scale healing and repeated distillation ablations more affordable while enabling long-context healing on a single GPU.
Takeaways & Limitations
The study evaluates only one 8B teacher–3.2B student pair and does not test whether its recommendations transfer across substantially different architectures or sizes.
Abstract
from arXiv · showhide
Small language models are often the only option for deployment under tight latency, cost, and on-premises constraints, but they are rarely trained from scratch: a compressed model is usually recovered through knowledge distillation (KD). This recovery step largely decides the final quality, yet it is expensive. We present a practitioner's study of how to make distillation training efficient, organised around two systems contributions. First, we show that offline KD (caching the teacher's top-$K$ logits once and training the student against the cache) matches online distillation at near-identical training loss while removing the teacher from memory, running about 29\% faster per iteration, and reaching up to 41\% higher throughput on a single H200 GPU. Second, we introduce a \emph{fused, chunked KL loss} that never materialises the full vocabulary-sized logit tensor, making peak memory linear in the sequence length. This removes the memory spike that otherwise caps context length and lets us train at four times the context (32{,}768 tokens) on a single GPU. A separate output-head-only toy benchmark isolates the loss kernel and confirms its memory and iteration-rate scaling from 4K to 256K tokens. Together these make large-scale healing and hundreds of ablations affordable. We also report supporting ablations on loss design and sequence packing. We release our chunked-loss implementation: https://github.com/CompactifAI/Full-Chunked-KL-Loss.
1 Introduction
The paper studies deployment-driven knowledge distillation for compact language models under latency, cost, and on-premises constraints. It separates two efficiency contributions: offline top-K teacher-logit caching and a fused, chunked KL loss for long-context training.
- Study scope: The study is a practically driven campaign on a compact ∼3.2B student derived from Llama 3.1 8B Instruct, with deployment-scale choices and reported trade-offs.The compact initialization method is independent of this work, and the distillation recipe is intended to apply to compact models initialized from larger teachers.
- Fused, chunked KL loss: 128 GB is the fused chunked loss’s peak memory at 32K context, versus nearly 250 GB for dense KL, exceeding a single H200’s 141 GB capacity.The difference comes from avoiding the vocabulary-sized logit/teacher tensor; other components are unchanged in the comparison.
- Offline distillation: Offline distillation caches the teacher’s top-K logits once, removes the teacher from memory and the forward-pass loop, and matches online quality at near-identical loss.This makes repeated ablations against the same cache practical.
- Fused, chunked KL loss: The fused, chunked KL loss extends memory-efficient cross-entropy techniques to sparse top-K teacher distributions with partially retained mass.It avoids materialising the vocabulary-sized tensor and computes a forward KL objective for knowledge distillation.
2 Related Work
Prior work frames knowledge distillation as a longstanding compression and recovery method, especially after structured pruning. Related research also emphasizes long-context evaluation and memory-efficient losses that avoid materializing full vocabulary-sized logits.
- Knowledge distillation: Knowledge distillation has long compressed large teachers into smaller students and commonly heals structured-pruned instruction-tuned language models.The cited lineage runs from logit matching to task-agnostic pretrained-transformer distillation and post-pruning recovery.
- Long context and memory-efficient losses: Long-context capability is evaluated over documents, retrieved passages, and agentic histories, motivating HELMET, Ruler, and retrieval-augmented generation tasks.These evaluations help identify where a student regresses.
- Long context and memory-efficient losses: Prior memory-efficient cross-entropy methods avoid full logit materialization by chunking sequences and fusing the output projection.The passage identifies Cut Cross-Entropy and Liger kernels as examples addressing the language-model head and loss bottleneck.
3 Types of Knowledge Distillation
The section defines forward-KL distillation in online and offline regimes, where offline training caches only the teacher’s top-K probabilities. Its main contribution is a fused chunked loss that avoids materializing the full vocabulary-sized logit tensor and keeps peak memory linear in sequence length.
- Online distillation: Online distillation keeps teacher and student in memory and materializes dense teacher probabilities and student log-probabilities, making it the most memory- and compute-intensive regime.The teacher produces a dense distribution for every position, while both models remain loaded simultaneously.
- Offline distillation: Offline distillation caches the teacher’s K=100 largest probabilities, retains unnormalized partial mass, and gathers only K student logits while computing scalar log Z over the vocabulary.The teacher-entropy, cross, and mass terms depend only on support entries; only log Z requires a full-vocabulary reduction.
- Offline implementations: The dense-reconstruction offline variant scatters cached top-K values into a dense teacher tensor and has peak memory O(SBV), serving as a correctness baseline closest to online distillation.It materializes both the reconstructed teacher and student log-probabilities in addition to student logits.
- Offline implementations: The sparse chunked variant evaluates the offline identity directly, processing sequence chunks while computing stable vocabulary reductions and accumulating sparse terms without dense teacher or log-softmax tensors.Each chunk writes scalar log Z values and uses scatter-add over the K retained entries.
- Fused chunked KL: The fused chunked loss integrates the output projection, discards each transient logit chunk, and reduces peak memory to O(SBd), linear in sequence length because d ≪ V.Its trade-off is one extra output projection per chunk during backward, analogous to output-head gradient checkpointing.
4 Experiments
Experiments show that offline top-100-logit distillation matches online training loss while reducing memory and iteration cost, and that fused chunking enables substantially longer contexts. Supporting ablations find logit KL essential for recovery, hidden-state features mildly beneficial, and naive packing a practical default.
- Offline distillation: Offline distillation matches online training-loss curves using only top-100 cached logits while lowering peak memory from about 103 to 78 GB.It also removes the resident teacher from memory.
- Fused chunked KL: At 8K context, fused chunking reduces peak memory from 78 GB for dense KL to 62 GB for forward-chunked and 58 GB for fused chunked.All three offline implementations produce the same training-loss curve; forward-chunked is fastest per iteration in this setting, while fused chunking unlocks long context.
- Controlled loss-only benchmark: At 32K tokens, fully chunked KL uses 5.45 GiB versus 85.2 GiB for dense KL, a 15.6× reduction, while dense KL fails at 64K.At 256K, fully chunked uses 11.6 GiB versus 134.2 GiB for forward-chunked.
- Controlled loss-only benchmark: At 256K, fully chunked reaches 0.630 iterations/s versus 0.190 for forward-chunked, a 3.3× advantage, after overtaking it from 64K onward.Forward-chunked leads at 32K with 5.46 iterations/s versus 5.04 for fully chunked.
- Supporting ablations: Naive all-ones sequence packing costs only about one point of MMLU relative to non-packed training, making it a reasonable default for distillation.The teacher’s KL signal appears to compensate for missing per-example block masking.
5 Limitations
The study evaluates one 8B teacher–∼3.2B student pair and does not establish transfer across substantially different architectures or compression settings. Its toy loss-kernel sweep isolates implementation scaling rather than end-to-end training, while hardware and framework portability remains unvalidated.
- Scope: The study evaluates only an 8B instruction-tuned teacher with a compact ∼3.2B student, leaving transfer to other architectures, compression methods, and student sizes open.The paper does not evaluate different model families, compression methods, or student sizes.
- Benchmark limitations: The 4K–256K sweep uses a toy output-projection network with synthetic inputs, so it isolates loss memory and timing rather than end-to-end training behavior.It does not measure training speed, model quality, convergence, or interactions with attention and optimiser state, and is presented alongside—not instead of—the real-LLM experiments.
- Systems portability: The systems results use Megatron-Bridge and ModelOpt on H200 GPUs, leaving the fused chunked KL formulation’s efficiency on other hardware and frameworks to be validated.The formulation is generic, but its efficiency characteristics beyond this software and hardware setup remain untested.
6 Conclusion
The paper’s efficient distillation recipe rests on offline cached top-K teacher logits and a fused chunked KL loss. Together, these choices reduce memory, improve throughput, and enable long-context healing on a single GPU.
- Conclusion: Offline distillation from cached top-K teacher logits matches online quality with lower memory and higher throughput.This removes the teacher from the training loop after caching its logits.
- Conclusion: The fused chunked KL loss keeps peak memory linear in sequence length, allowing long-context healing on a single GPU.The loss avoids materialising the full vocabulary-sized logit tensor.
- Conclusion: A controlled loss-only benchmark isolates the chunked-loss mechanism.The benchmark evaluates the fully chunked implementation independently of the broader distillation system.
A Experimental Configuration
All profiling runs share one fixed configuration, varying only the KL-loss implementation and sequence length. Offline teacher top-K logits were precomputed, and profiling combined Nsight Systems capture with CUDA memory snapshots.
- Configuration: All profiling runs use a shared configuration, with only the KL-loss implementation and sequence length varied.The tested sequence lengths are 8,192 and 32,768; all other settings remain fixed.
- Configuration: Offline runs use teacher top-K logits precomputed with SGLang.
- Profiling: Profiling uses 15-iteration runs, Nsight Systems capture over iterations 10–13, and per-step CUDA memory-history snapshots.
B Toy Loss-Kernel Benchmark Configuration
The controlled benchmark isolates the KL loss forward/backward path using deterministic equivalent inputs across three losses. Timing and memory are reported for the maximum tensor-parallel rank, with plotted results selecting a 4,096-token chunk from a broader sweep.
- Controlled configuration: Deterministic random hidden states, output-projection weights, and sparse teacher targets are generated once per configuration and reused across all three losses.This ensures equivalent inputs for the controlled comparison.
- Controlled configuration: The benchmark measures only the output projection and KL forward/backward path, excluding the transformer body and optimizer step.It therefore isolates the loss kernel rather than end-to-end training.
- Benchmark reporting: Timing and memory report the maximum rank value under tensor parallelism, while Figure 3 uses a 4,096-token chunk from a sweep containing other chunk sizes.The plotted CSV selects the 4,096-token configuration from the broader chunk-size sweep.