Paper deep dive
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
Intelligence
Status: not_run | Model: - | Prompt: - | Confidence: 0%
Entities (0)
Relation Signals (0)
No relation signals yet.
Cypher Suggestions (0)
No Cypher suggestions yet.
Abstract
Abstract: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: this https URL.
Tags
Links
- Source: https://arxiv.org/abs/2608.03796v1
- Canonical: https://arxiv.org/abs/2608.03796v1
Trouble viewing inline? Open PDF directly →
Full Text
36,294 characters extracted from source content.
Expand or collapse full text
Efficient Knowledge Distillation for LLMs: Offline Top-K Logits and a Fused Chunked KL Loss Bakbergen Ryskulov1, Iker García-Ferrero2, David Montero, David Jansen, Ali Hashemi, Jezabel R. Garcia, Antonio Tiene, Román Orús Abstract 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 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. 11footnotetext: bakbergen.ryskulov@multiversecomputing.com22footnotetext: iker.garcia@multiversecomputing.com 1 Introduction Figure 1: The fused chunked KL loss unlocks long-context healing. GPU memory over one training iteration at a 32K context length, broken down by component. The dense KL loss materialises a vocabulary-sized logit/teacher tensor (hatched, extrapolated) that spikes peak memory near 250250 GB, exceeding a single H200’s 141141 GB capacity, whereas the fused chunked loss never forms that tensor and peaks at 128128 GB. Only the loss/logits component differs; the remaining components are an estimated split of the measured total. Large language models are increasingly deployed under hard production constraints: tight latency budgets, per-token cost ceilings, and on-premises serving where the largest models do not fit. The common response is to deploy a compact model derived from a capable teacher, whose quality is then recovered through knowledge distillation, training the student to match the teacher’s behaviour. This recovery step decides most of the practical cost and the final quality, yet it is under-documented relative to its impact. We report guidance from an extensive distillation campaign on a compact (∼ 3.2B) student derived from Llama 3.1 8B Instruct (Llama Team, AI @ Meta 2024). The method used to obtain the compact initialisation is independent of this work (Muralidharan et al. 2024); the distillation recipe applies to any compact model initialised from a larger teacher. This is a practically driven contribution in the strict sense: no new algorithm, but deployment-driven choices measured at scale and reported with their trade-offs, including where they fail. The paper makes two efficiency contributions, which we deliberately keep separate because they address different bottlenecks. Offline distillation. Computing the teacher on the fly (online KD) keeps both models resident and recomputes the teacher forward pass at every step. We instead compute the teacher logits once, cache the top-K per token, and train the student against the cache. This matches online quality at near-identical loss while removing the teacher from memory and the teacher forward pass from the loop, which lowers cost and, crucially for a research campaign, lets us run tens of ablations against the same cached targets (§4.1). A fused, chunked KL loss. The binding constraint on long-context healing is not the transformer body but the logit tensor produced by the language-model head and its loss, whose vocabulary-sized footprint spikes peak memory and caps the trainable sequence length (Wijmans et al. 2025; Hsu et al. 2024). Memory-efficient cross-entropy losses that fuse the output projection and chunk the sequence are by now established (Wijmans et al. 2025; Hsu et al. 2024); we bring the technique to a knowledge-distillation objective, where the target is a sparse top-K teacher distribution with partially retained mass rather than a one-hot label, giving a forward KL with a different closed-form gradient (§3.2). The result makes peak memory linear in the sequence length (Figure 1), removing the spike that caps context and unlocking the long-context healing a compact student needs to serve long inputs (Liu et al. 2024; Gao et al. 2025). This implementation is not available in current libraries and we will release it. Contributions. • An offline top-K logit-distillation pipeline that matches online quality while cutting memory and raising throughput, which is what makes large-scale ablation studies practical (§4.1). • A fused, chunked KL loss that extends memory-efficient cross-entropy kernels (Wijmans et al. 2025; Hsu et al. 2024) to a sparse top-K teacher, making peak memory linear in the sequence length and unlocking long-context healing on a single GPU, with the loss-kernel scaling isolated in a controlled output-head benchmark (§3.2, §4.2, §4.3). • Supporting ablations enabled by the efficient setup, (loss design and sequence packing) that round out a reproducible recipe (§4.4). 2 Related Work Knowledge distillation. Distilling a large teacher into a smaller student is a long-standing route to compression, from the original logit-matching formulation (Hinton et al. 2015) to task-agnostic distillation of pretrained transformers (Sanh et al. 2019; Wang et al. 2020). At the scale of instruction-tuned language models, distillation is most often used as the recovery step after structured pruning: Minitron prunes a large model and heals it with distillation (Muralidharan et al. 2024), and the Gemma reports distil from top-K teacher logits, with K of order a few hundred (Gemma Team 2024), which motivates our study of how many cached logits are actually needed. The closest work in framing is the industry-track comparison of task-agnostic distillation methods of Udagawa et al. (2023); we extend that line to instruction-tuned LLM scale, an offline and memory-efficient training pipeline, and long-context healing with an explicit deployment framing. Our aim is a practitioner’s recovery recipe rather than a new distillation algorithm, so we report the choices and their trade-offs, including where they fail. Long context and memory-efficient losses. Long-context ability is now a first-class capability: models are asked to reason over long documents, retrieved passages, and agentic histories, and use that context unevenly enough to need careful evaluation (Liu et al. 2024), so we adopt HELMET (Yen et al. 2025) with its Ruler and retrieval-augmented generation tasks (Hsieh et al. 2024) to locate where our student regresses. The same capability must be acquired in training, where the binding cost is not the transformer body but the logit tensor materialised by the language-model head and its loss (Wijmans et al. 2025; Gao et al. 2025). Cut Cross-Entropy (Wijmans et al. 2025) and the Liger kernels (Hsu et al. 2024) remove this cost for the cross-entropy objective by chunking the sequence and fusing the output projection so the full logits are never materialised. Our fused chunked loss (§3.2) brings the same technique to a knowledge-distillation objective these libraries do not support: the target is a sparse top-K teacher with retained mass M≤1M≤ 1, so the loss is a forward KL divergence and the closed-form gradient is a dense “M⋅softmaxM·softmax” term minus a sparse teacher correction (Equation 3), rather than a single ground-truth subtraction. This adaptation is the part absent from current libraries, and we will release it. 3 Types of Knowledge Distillation We distil a large teacher into a smaller student by matching their output distributions with a forward Kullback–Leibler (KL) objective. For a single token, let z∈ℝVz ^V be the student logits over a vocabulary of size V and qv=exp(zv)/Zq_v= (z_v)/Z with Z=∑uexp(zu)Z= _u (z_u) the student probability of token v, so that logqv=zv−logZ q_v=z_v- Z. With p∈ℝVp ^V the teacher distribution, the per-token loss is ℒKL(p,z)=∑v=1Vpv(logpv−logqv),L_KL(p,z)= _v=1^Vp_v\,( p_v- q_v), (1) averaged over all next-token-shifted, loss-masked positions. We consider two regimes for obtaining p: online, where the teacher is resident in memory, and offline, where only the teacher’s top-K probabilities are precomputed and cached. 3.1 Online Distillation In the online setting both teacher and student are loaded simultaneously. A teacher forward pass produces the dense distribution p∈ℝVp ^V for every position, and (1) is evaluated directly against the student’s dense log-softmax. This regime is the most expressive (the full teacher distribution is available) but the most memory- and compute-intensive: it holds both models and materialises two dense ℝVR^V tensors per position (teacher probabilities and student log-probabilities), on top of recomputing the teacher at every step. 3.2 Offline Distillation with a Top-K Teacher To remove the teacher from memory, we precompute and cache only its K=100K=100 largest probabilities per position. Let ⊂1,…,VS⊂\1,…,V\, ||=K|S|=K, denote this support, with pv=0p_v=0 for v∉v . The retained mass M=∑v∈pv≤1M= _v p_v≤ 1 may be strictly below one because of truncation; we do not renormalise, and the formulation below accounts for the partial mass exactly. Substituting logqv=zv−logZ q_v=z_v- Z into (1) and restricting to the support yields the identity that underlies all three offline implementations: ℒKL(p,z)=∑v∈pvlogpv⏟H−∑v∈pvzv⏟C+MlogZ.L_KL(p,z)= [] _v ^p_v p_v_H- [] _v ^p_vz_v_C+M Z. (2) The teacher-entropy term H, the cross term C, and the mass M depend only on the K support entries, so they require gathering just K student logits per position. Only the log-normaliser logZ Z depends on the entire vocabulary, but it is a scalar per position, a reduction, not an ℝVR^V tensor. The three offline methods below are mathematically equivalent evaluations of (2); they differ only in how they handle logZ Z and the student logits. We use the untempered objective (τ=1τ=1) throughout. Full Dense KL Computation The simplest approach reconstructs the dense teacher: the cached top-K values are scattered into a dense tensor p∈ℝB×S×Vp ^B× S× V and the objective is evaluated against the student’s dense log-softmax, exactly as in the online case. It therefore materialises two vocabulary-sized tensors (the reconstructed top-K teacher and the student log-probabilities) on top of the student logits, so peak memory is O(SBV)O(SBV). It serves as a correctness baseline: it is the offline computation closest to online distillation. Sparse KL and Forward-Chunked Loss This variant keeps the teacher sparse and evaluates (2) directly, never forming a dense teacher or a dense log-softmax. The student logits z (produced by the language-model head) are processed in contiguous chunks of CsC_s sequence positions: each chunk computes a numerically stable maximum and exponential sum over the vocabulary and writes the scalar logZ Z for its positions, while the sparse terms H, C, and M are accumulated by scatter-add over the K retained entries. Because the chunking removes only the auxiliary dense tensors, the student logits and their gradient are still held in full, so peak memory retains a vocabulary-sized O(SBV)O(SBV) term and the method does not on its own enable longer sequences. It is, however, the fastest variant in our profiling, since it keeps the standard output projection but removes the dense teacher, dense log-softmax, and dense KL arithmetic. Full Chunked KL Computation This is our main contribution. The loss fuses the output projection into the loss, so the full [S,B,V][S,B,V] logit tensor is never materialised, in neither the forward pass nor as a stored gradient, only a transient chunk of logits exists at a time. Given hidden states h∈ℝS×B×dh ^S× B× d and the output projection W∈ℝV×dW ^V× d, the forward pass processes the sequence chunk by chunk: it projects z[s0:s1]=h[s0:s1]W⊤z_[s_0:s_1]=h_[s_0:s_1]W , accumulates logZ Z and the sparse terms of (2), and immediately discards each chunk of logits. It retains for the backward pass only h, the per-position scalars logZ Z and M (each ℝS×BR^S× B), and the sparse teacher entries. The backward pass recomputes the logits chunk by chunk, rebuilds q[s0:s1]=exp(z[s0:s1]−logZ[s0:s1])q_[s_0:s_1]= (z_[s_0:s_1]- Z_[s_0:s_1]) from the saved normaliser, and forms the logit gradient in closed form, ∂ℒKL∂zv=Mqv−pv, _KL∂ z_v=M\,q_v-p_v, (3) i.e. a dense “M⋅softmaxM·softmax” term minus the sparse teacher correction at the K support positions. Each chunk gradient is projected back to accumulate ∂ℒ/∂h /∂ h and ∂ℒ/∂W /∂ W, after which the chunk is freed (Algorithm 1). The vocabulary-sized cost is thus confined to a single chunk and is independent of sequence length; the only quantity that grows with S is the hidden-state activation, O(SBd)O(SBd) with d≪Vd V. Peak memory is therefore linear in the sequence length, in contrast to the O(SBV)O(SBV) footprint of the other two variants, at the cost of one extra output projection per chunk in the backward pass, a gradient-checkpointing trade-off on the output head. The formulation is compatible with vocabulary sharding: the normaliser and sparse top-K terms are reduced across shards, and each rank stores and differentiates only its local vocabulary slice. Algorithm 1 Fused chunked KL (output projection fused into the loss) 1:hidden states h∈ℝS×B×dh ^S× B× d; projection W∈ℝV×dW ^V× d; support (t,b,v,p)(t,b,v,p); chunk size CsC_s 2:Forward (no gradient tracked): 3:for each chunk [s0,s1)[s_0,s_1) do ⊳ Pass 1: normaliser 4: z←h[s0:s1]W⊤z← h_[s_0:s_1]W ; m←maxvzm← _vz ⊳ distributed max if sharded 5: σ[s0:s1]←∑vexp(z−m) _[s_0:s_1]← _v (z-m); discard z 6:end for 7:logZ←m+logσ Z← m+ σ ⊳ reduce σ if sharded 8:for each chunk [s0,s1)[s_0,s_1) do ⊳ Pass 2: loss 9: z←h[s0:s1]W⊤z← h_[s_0:s_1]W 10: M←∑pM←Σ p; H←∑plogpH←Σ p p; C←∑pzt,b,vC←Σ p\,z_t,b,v 11: ℒ[s0:s1]←H−C+M⊙logZ[s0:s1]L_[s_0:s_1]← H-C+M Z_[s_0:s_1]; discard z 12:end for 13:save h,W,logZ,M,(t,b,v,p)h,\,W,\, Z,\,M,\,(t,b,v,p); return ℒL 14:Backward given g=∂ℒ/∂ℓg= /∂ : 15:for each chunk [s0,s1)[s_0,s_1) do 16: z←h[s0:s1]W⊤z← h_[s_0:s_1]W ⊳ recompute logits 17: q←exp(z−logZ[s0:s1])q← (z- Z_[s_0:s_1]) ⊳ rebuild softmax 18: G←(M⊙g)[s0:s1]⋅qG←(M g)_[s_0:s_1]· q 19: Gt,b,v-=p⋅gt,bG_t,b,v -=p· g_t,b ⊳ Eq. (3) 20: ∂h[s0:s1]←GW∂ h_[s_0:s_1]← G\,W; ∂W+=G⊤h[s0:s1]∂ W +=G h_[s_0:s_1] 21:end for 22:reduce ∂h∂ h across shards; return ∂h,∂W∂ h,\ ∂ W 4 Experiments We organise the experiments around the two contributions and a set of supporting ablations, and describe the relevant setup inline in each subsection rather than in a separate section. Unless stated otherwise the teacher is Llama 3.1 8B Instruct and the student is a compact ∼ 3.2B model; distillation uses NVIDIA Megatron-Bridge with the Megatron-LM backend and NVIDIA ModelOpt, and efficiency is measured with the PyTorch memory profiler, Megatron-Bridge profiling, and NVIDIA Nsight Systems. The complete training configuration and software stack are listed in Appendix A. The controlled experiment in §4.3 is the sole exception: it uses a toy output-projection network and synthetic tensors to isolate the loss kernel, rather than a pretrained or end-to-end language model. Figure 2: All methods compared at 8K context on a single H200: online distillation and the three offline implementations (dense, forward-chunked, and fused chunked KL). Top-left: training loss per step is near-identical across all methods, including the offline runs that use only the top-100100 cached logits. Top-right: iteration time (bar height) and throughput (TFLOP/s, labelled inside each bar). Bottom: peak GPU memory by component; offline removes the resident teacher (red), and the loss/logits term (blue) shrinks as the vocabulary-sized tensor is chunked away, effectively vanishing for the fused loss. 4.1 Online vs. Offline (Full Dense KL) Setup. We profile a single training step on one H200 GPU at a sequence length of 8,192 on the SmolTalk supervised fine-tuning data (Allal et al. 2025), comparing online distillation against offline distillation with the full dense KL loss. The offline run caches the teacher’s top-100100 logits per token; correctness is assessed by matched training loss, and efficiency by peak memory, throughput, and seconds per iteration. Results. The two regimes reach near-identical training-loss curves, even though the offline run trains against only the top-100100 cached logits (Figure 2, top-left). Offline distillation is meanwhile markedly cheaper: it removes the teacher from memory, lowering peak memory from about 103103 to 7878 GB, runs about 29%29\% faster per iteration (25.9→18.525.9→ 18.5 s), and raises throughput from 237237 to 331331 TFLOP/s (∼40% 40\%). The quality match at lower cost is what makes offline distillation attractive at scale: the teacher need only be run once, after which the cache can be reused across many ablations, and the savings grow with teacher size (e.g. a 70B teacher need not sit in the training loop). 4.2 Adding the Chunked KL Loss Setup. Using the same single-H200, 8K, SmolTalk profiling setup, we compare the three offline implementations of the same objective: the full dense KL baseline (§3.2), the forward-chunked loss (§3.2), and the fused chunked loss (§3.2). Results. All three implementations produce the same training-loss curve (Figure 2, top-left), confirming they evaluate the same objective. They differ in cost. The fused chunked loss reduces peak memory furthest, from 7878 GB (dense) to 6262 GB (forward-chunked) to 5858 GB (fused), by never materialising the vocabulary-sized logit tensor, a reduction not available in current libraries such as ModelOpt. Because peak memory is now linear in sequence length rather than O(SBV)O(SBV), the fused loss removes the memory spike that otherwise caps context: on the same single H200 it trains at 32,76832,768 tokens, roughly four times the context that fits with the dense loss. The forward-chunked loss is the fastest per iteration, while the fused loss pays a small recompute cost in the backward pass; in exchange it is the only variant that unlocks long context. We stress that under this setting, a single H200 at 8K context, the fused chunked loss is not the fastest approach: its extra output projection in the backward pass makes the forward-chunked loss lead on iteration time. Scaling to larger models and longer context. The advantage of the fused loss grows sharply with model and context size, where the freed memory turns into a throughput win. In further experiments distilling GPT-OSS-20B at a context length of 32,76832,768 on 8×8×H200 nodes, the memory it frees lets the model drop from four nodes (tensor parallel 44, pipeline parallel 44, expert parallel 22) to a single node (tensor parallel 22, pipeline parallel 11, expert parallel 44), removing most of the inter-node communication. Step time then falls from 57.057.0 to 12.2312.23 seconds (∼ 5× faster) and throughput rises from 74.274.2 to 345.7345.7 TFLOP/s per GPU: never instantiating the O(SBV)O(S\,B\,V) logit tensor both removes the memory bottleneck and pushes GPU utilisation far higher. The peak-memory reductions reported above (78→62→5878\!→\!62\!→\!58 GB) are measured at 8K context, and the gap widens with length. At 32,76832,768 tokens the dense loss peaks at roughly 250250 GB, beyond a single H200’s capacity, so it does not fit, whereas the fused chunked loss peaks at about 128128 GB. 4.3 Isolating Loss-Kernel Scaling Figure 3: Controlled loss-only benchmark Left: peak memory per GPU. Right: forward-backward iteration rate on a logarithmic scale. Dense KL fails from 64K onward; forward-chunked is fastest through 32K, while fully chunked is faster at longer contexts and uses substantially less memory. Batch size 11, hidden size 4,0964,096, vocabulary 131,072131,072, top-K=100K=100, tensor parallelism 22, and chunk size 4,0964,096. The preceding results measure real LLM training and therefore mix the cost of the KL implementation with transformer layers, attention, optimiser state, data movement, and framework overhead. To isolate the mechanism behind the memory reduction, we additionally run a controlled microbenchmark using a toy neural network. The results are depicted in Figure 3. This experiment is deliberately not an LLM benchmark and its absolute memory and iteration-rate values should not be compared directly with Figures 2 and 1. Setup. The toy network contains only a vocabulary output projection: synthetic hidden states h are multiplied by a synthetic weight matrix W, then one of the three KL implementations executes its forward and backward passes. There are no transformer blocks, attention layers, optimiser step, data loader, or resident teacher model. Deterministic hidden states, projection weights, and sparse top-100100 teacher targets are reused across methods. The plotted run uses hidden size 4,0964,096, vocabulary size 131,072131,072, batch size 11, bfloat16, tensor parallelism 22, and a 4,0964,096-token chunk for both chunked variants, over sequence lengths from 4K to 256K. Each configuration is launched in a fresh distributed subprocess so an out-of-memory failure cannot contaminate later measurements. Peak allocated CUDA memory and mean forward-plus-backward iteration time are reduced by the maximum over the two ranks. Separate deterministic CPU tests verify agreement of the per-token loss and hidden-state gradients to 10−410^-4 tolerance. Full details appear in Appendix B. Results. At 32K tokens, peak memory is 85.285.2 GiB for dense KL, 17.717.7 GiB for the forward-chunked loss, and 5.455.45 GiB for the fully chunked loss (Figure 3, left), a 15.6×15.6× reduction from dense to fully chunked. Dense KL then fails at 64K. At 256K, the forward-chunked loss still holds the full logits and reaches 134.2134.2 GiB per GPU, whereas the fully chunked loss uses 11.611.6 GiB. The timing panel (Figure 3, right) exposes the recomputation trade-off: at 32K, forward-chunked leads at 5.465.46 iterations/s versus 5.045.04 for fully chunked, but the ranking reverses at 64K. At 256K the fully chunked loss reaches 0.6300.630 iterations/s versus 0.1900.190 for forward-chunked, a 3.3×3.3× advantage in this isolated workload. The fully chunked loss uses 15.6×15.6× less memory than dense KL at 32K and 11.6×11.6× less than forward-chunked at 256K. Although forward-chunked is faster at smaller contexts, fully chunked overtakes it from 64K onward and is 3.3×3.3× faster at 256K. Thus the microbenchmark validates the intended kernel-level scaling: full-sequence logits dominate the other implementations, while the fused implementation bounds vocabulary-sized storage by the chunk. 4.4 Additional Ablations The efficient offline setup made several smaller studies cheap to run. We summarise the two that bear directly on the recipe. Figure 4: Loss design ablation. Best MMLU and GSM8K for a compact student healed under different losses, with the teacher for reference. An intermediate (feature) loss alone collapses; logit KL is indispensable; and logit KL plus a hidden-state feature loss is best. Loss design. Holding the student, teacher, and data budget fixed and varying only the loss, the choice of loss is the dominant driver of recovery (Figure 4). An intermediate-layer feature loss applied on its own collapses the student (MMLU (Hendrycks et al. 2021) near 28%28\%, GSM8K (Cobbe et al. 2021) near 4%4\%): logit-level KL is indispensable, recovering MMLU to 59.9%59.9\% and GSM8K to 65.9%65.9\%. Adding a hidden-state feature loss on top of logit KL gives a small, consistent gain, reaching 60.6%60.6\% MMLU and 67.5%67.5\% GSM8K (mean-squared-error variant; a cosine variant is comparable). The recommendation is therefore simple: always include logit KL, and add a hidden-state feature loss for a reliable improvement. Figure 5: Sequence packing. MMLU after supervised fine-tuning for non-packed vs. naively packed (all-ones mask) training; packing costs about one point. Sequence packing. Packing concatenates short examples into one long sequence to keep the affordable long context full of useful tokens. Packing with a naive all-ones attention mask, which permits attention across example boundaries, costs only about one point of MMLU relative to non-packed training (Figure 5). The teacher’s KL signal appears to compensate for the missing per-example block mask, so cheap naive packing is a reasonable default for distillation. Figure 6: Short-context accuracy of the compact student against the teacher. The student retains most of the teacher’s BoolQ and HellaSwag accuracy and stays within about nine points on MMLU, with larger gaps on WinoGrande and GSM8K. Figure 6 summarises the resulting compact student against its teacher on BoolQ (Clark et al. 2019), WinoGrande (Sakaguchi et al. 2020), MMLU, HellaSwag (Zellers et al. 2019), and GSM8K: the student retains most of the teacher’s short-context accuracy at less than half the size. 5 Limitations Our study focuses on a single teacher–student pair: an 8B instruction-tuned teacher and a compact ∼ 3.2B student. Although the recipe is intended to be broadly applicable, we do not evaluate different model families, compression methods, or student sizes, so the extent to which the recommendations transfer to substantially different architectures remains open. The 4K–256K loss-kernel sweep intentionally uses a toy output-projection network with synthetic inputs. It isolates the asymptotic memory and timing of the loss implementations, but it does not measure end-to-end training speed, model quality, convergence, or interactions with attention and optimiser state at those sequence lengths. We therefore use it as mechanistic evidence alongside, not as a replacement for, the real-LLM experiments. Our systems results are obtained with Megatron-Bridge and ModelOpt on H200 GPUs. The fused chunked KL formulation is generic, but its efficiency characteristics on other hardware and frameworks remain to be validated. 6 Conclusion We presented a practitioner’s recipe for efficient knowledge-distillation recovery of a compact LLM. Two systems choices carry most of the benefit: distil offline from cached top-K teacher logits to match online quality at lower memory and higher throughput, and use a fused chunked KL loss so that peak memory is linear in sequence length and long-context healing fits on a single GPU. A controlled loss-only benchmark isolates this mechanism: the fully chunked implementation remains within 11.611.6 GiB per GPU at 256K tokens and overtakes the forward-chunked implementation in iteration rate at long sequence lengths, while we explicitly separate these toy-network results from end-to-end LLM throughput. Supporting ablations recommend combining logit KL with a hidden-state feature loss and show that cheap naive sequence packing costs only about a point of MMLU. The net result is most of the teacher’s short-context quality at a fraction of the size and training cost, with known long-context limits. We made our chunked-loss implementation open-source. References L. B. Allal, A. Lozhkov, E. Bakouch, G. M. Blázquez, G. Penedo, L. Tunstall, A. Marafioti, H. Kydlíček, A. P. Lajarín, V. Srivastav, et al. (2025) SmolLM2: when smol goes big – data-centric training of a small language model. arXiv preprint arXiv:2502.02737. Cited by: §4.1. C. Clark, K. Lee, M. Chang, T. Kwiatkowski, M. Collins, and K. Toutanova (2019) BoolQ: exploring the surprising difficulty of natural yes/no questions. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics (NAACL), Cited by: §4.4. K. Cobbe, V. Kosaraju, M. Bavarian, M. Chen, H. Jun, L. Kaiser, M. Plappert, J. Tworek, J. Hilton, R. Nakano, C. Hesse, and J. Schulman (2021) Training verifiers to solve math word problems. arXiv preprint arXiv:2110.14168. Cited by: §4.4. T. Gao, A. Wettig, H. Yen, and D. Chen (2025) How to train long-context language models (effectively). In Proceedings of the 63rd Annual Meeting of the Association for Computational Linguistics (ACL), Cited by: §1, §2. Gemma Team (2024) Gemma: open models based on gemini research and technology. arXiv preprint arXiv:2403.08295. Cited by: §2. D. Hendrycks, C. Burns, S. Basart, A. Zou, M. Mazeika, D. Song, and J. Steinhardt (2021) Measuring massive multitask language understanding. In International Conference on Learning Representations (ICLR), Cited by: §4.4. G. Hinton, O. Vinyals, and J. Dean (2015) Distilling the knowledge in a neural network. arXiv preprint arXiv:1503.02531. Cited by: §2. C. Hsieh, S. Sun, S. Kriman, S. Acharya, D. Rekesh, F. Jia, and B. Ginsburg (2024) RULER: what’s the real context size of your long-context language models?. In Conference on Language Modeling (COLM), Cited by: §2. P. Hsu, Y. Dai, V. Kothapalli, Q. Song, S. Tang, S. Zhu, S. Shimizu, S. Sahni, H. Ning, and Y. Chen (2024) Liger kernel: efficient triton kernels for LLM training. arXiv preprint arXiv:2410.10989. Cited by: 2nd item, §1, §2. N. F. Liu, K. Lin, J. Hewitt, A. Paranjape, M. Bevilacqua, F. Petroni, and P. Liang (2024) Lost in the middle: how language models use long contexts. Transactions of the Association for Computational Linguistics 12, p. 157–173. Cited by: §1, §2. Llama Team, AI @ Meta (2024) The llama 3 herd of models. arXiv preprint arXiv:2407.21783. Cited by: §1. S. Muralidharan, S. T. Sreenivas, R. Joshi, M. Chochowski, M. Patwary, M. Shoeybi, B. Catanzaro, J. Kautz, and P. Molchanov (2024) Compact language models via pruning and knowledge distillation. In Advances in Neural Information Processing Systems (NeurIPS), Cited by: §1, §2. K. Sakaguchi, R. L. Bras, C. Bhagavatula, and Y. Choi (2020) WinoGrande: an adversarial winograd schema challenge at scale. In Proceedings of the AAAI Conference on Artificial Intelligence, Cited by: §4.4. V. Sanh, L. Debut, J. Chaumond, and T. Wolf (2019) DistilBERT, a distilled version of bert: smaller, faster, cheaper and lighter. arXiv preprint arXiv:1910.01108. Cited by: §2. T. Udagawa, A. Trivedi, M. Merler, and B. Bhattacharjee (2023) A comparative analysis of task-agnostic distillation methods for compressing transformer language models. In Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing: Industry Track, Cited by: §2. W. Wang, F. Wei, L. Dong, H. Bao, N. Yang, and M. Zhou (2020) MiniLM: deep self-attention distillation for task-agnostic compression of pre-trained transformers. In Advances in Neural Information Processing Systems (NeurIPS), Cited by: §2. E. Wijmans, B. Huval, A. Hertzberg, V. Koltun, and P. Krähenbühl (2025) Cut your losses in large-vocabulary language models. In International Conference on Learning Representations (ICLR), Cited by: 2nd item, §1, §2. H. Yen, T. Gao, M. Hou, K. Ding, D. Fleischer, P. Izsak, M. Wasserblat, and D. Chen (2025) HELMET: how to evaluate long-context language models effectively and thoroughly. In International Conference on Learning Representations (ICLR), Cited by: §2. R. Zellers, A. Holtzman, Y. Bisk, A. Farhadi, and Y. Choi (2019) HellaSwag: can a machine really finish your sentence?. In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics (ACL), Cited by: §4.4. Appendix A Experimental Configuration Table 1 lists the configuration shared by all profiling runs in §4.1 and §4.2, taken from the merged configuration dumps recorded in the run logs. The runs vary only the KL-loss implementation and the sequence length (8,1928,192 or 32,76832,768); all other settings are held fixed. Teacher top-K logits for the offline runs were precomputed with SGLang. Profiling used 15-iteration runs with NVIDIA Nsight Systems capture over iterations 10–13 and per-step CUDA memory-history snapshots. Student architecture (3.233.23B parameters) Layers / hidden / FFN size 2828 / 2,8162,816 / 7,1687,168 Heads (GQA groups) 3232 (88), 128128 KV channels Normalisation RMSNorm (ϵ=10−5ε=10^-5) Positions / activation RoPE / SwiGLU, no biases Dropout (attention / hidden) 0.00.0 / 0.00.0 Optimisation Optimiser Adam (0.90.9, 0.9990.999, 10−810^-8) Weight decay / grad clip 0.10.1 / 1.01.0 Learning rate (cosine) 2×10−6→2×10−72\!×\!10^-6\!→\!2\!×\!10^-7 Warmup 1010 iterations Batch size (global / micro) 3232 / 11 Random seed 12341234 Distillation Objective forward KL, τ=1τ=1 Teacher support top-K, K=100K=100 Precision and parallelism Precision bf16 mixed, FP32 grad reduction Attention backend FlashAttention TP / P / CP / EP 11 / 11 / 11 / 11 (single GPU) Software stack Base container NVIDIA NGC PyTorch 26.01 (Ubuntu Linux, CUDA 12) Training framework Megatron-Bridge v0.3.0 (Megatron-Core, -FSDP) ModelOpt Kernels Transformer Engine 2.11 Multi-node comm. NCCL 2.29, AWS EFA (aws-ofi-nccl 1.18.0) GDRCopy 2.5.1 Other libraries Transformers 4.57.6 Teacher-logit precompute SGLang Table 1: Training configuration shared by all profiling runs. TP/P/CP/EP: tensor, pipeline, context, and expert parallelism. Appendix B Toy Loss-Kernel Benchmark Configuration Table 2 records the controlled configuration used for Figure 3. The benchmark code constructs random but deterministic hidden states, output-projection weights, and sparse teacher targets once per configuration and presents equivalent inputs to all three losses. It measures only the output projection and KL forward/backward path; no transformer body or optimiser step is present. The plotted CSV selects a 4,0964,096-token chunk from a sweep that also contains other chunk sizes. Component Value Network scope Vocabulary output projection only Hidden / vocabulary size 4,0964,096 / 131,072131,072 Batch size / teacher support 11 / top-100100 Precision / temperature bfloat16 / 1.01.0 Tensor parallelism 22 H200 GPUs Plotted chunk size 4,0964,096 tokens Sequence lengths 4K, 8K, 16K, 32K, 64K, 128K, 256K Timing Mean forward + backward after warm-up Memory Peak allocated CUDA memory per GPU Isolation Fresh distributed process per setting Table 2: Configuration for the controlled loss-only benchmark. Timing and memory report the maximum rank value under tensor parallelism.