Chunked KL loss, running Knowledge Distillation locally in less <6GB VRAM

A small, self-contained CUDA benchmark for comparing the memory and throughput cost of three knowledge-distillation losses:

Full Dense KL— reconstructs a dense teacher and evaluates the standard dense forward KL; this is the correctness baseline.** Forward-Chunked Loss**— keeps the top-K teacher sparse and chunks the loss computation, but retains the full student-logit tensor for backward.Full Chunked KL— fuses the output projection into the loss and recomputes logits by chunk during backward, so full-sequence logits are never stored.

Each experiment runs in an isolated torchrun

subprocess. An out-of-memory error or kernel failure is recorded as FAILED

, and the remaining benchmark matrix continues.

For every method, sequence length, and applicable chunk size, the benchmark reports:

Synthetic hidden states, output weights, and sparse teacher targets are generated deterministically and reused across methods so that every loss sees equivalent inputs.

• Python 3.10+
• PyTorch with CUDA support
• One or more NVIDIA GPUs

Install PyTorch using the command recommended for your CUDA version at pytorch.org. The benchmark itself uses only PyTorch and the Python standard library.

Run the default benchmark:

python benchmark.py

Choose the sequence lengths, chunk sizes, tensor-parallel size, iterations, and batch size:

python benchmark.py \
  --sequence-lengths 4096,8192,16384,32768,65536,131072,262144 \
  --chunk-sizes 8192,4096,2048,1024,512 \
  --tp 2 \
  --iterations 10 \
  --topk 100 \
  --batch-size 1

The parent process launches one isolated distributed job for each row in the matrix. This isolation is intentional: CUDA OOM failures do not poison the process used by later experiments.

Use python benchmark.py --help

for the complete set of defaults and implementation-specific options.

Run the CPU equivalence tests before benchmarking:

python test_losses.py

The tests compare loss values and hidden-state gradients across the three implementations.

All three methods evaluate the same forward KL against cached top-K teacher probabilities. For teacher support

L_KL(p, z) = sum_{v in S} p_v log(p_v) - sum_{v in S} p_v z_v + M logsumexp(z)

Only the log-normalizer spans the full vocabulary; the teacher-entropy and cross terms require just the K cached entries.

The reference implementation scatters the cached top-K probabilities into a dense teacher tensor and compares it with the student's dense log-softmax. It materializes the dense teacher and student log-probabilities in addition to the student logits, giving

This variant keeps the teacher sparse and evaluates the objective directly in contiguous sequence chunks, avoiding both a dense teacher and a dense log-softmax. However, the language-model head still produces the full student logits, which autograd retains for backward. Peak memory therefore remains

This variant fuses the vocabulary projection into the loss. The forward pass projects one sequence chunk at a time, accumulates the log-normalizer and sparse loss terms, and immediately discards the logits. The backward pass recomputes each chunk and applies the closed-form gradient

To regenerate the results figure, also install Matplotlib:

python -m pip install matplotlib

For a SVG:

python plot_peak_vram.py \
  --input kd_benchmark_results.csv \
  --output assets/peak_vram.svg \
  --chunk-size 4096

For a PNG:

python plot_peak_vram.py \
  --input kd_benchmark_results.csv \
  --output assets/peak_vram.png \
  --chunk-size 4096 \
  --dpi 180

This repository accompanies Distillation Tricks for Compact LLMs: Efficient Offline and Chunked Knowledge Distillation by Bakbergen Ryskulov, Iker García-Ferrero, David Montero, David Jansen, Ali Hashemi, Jezabel R Garcia, Antonio Tiene, and Román Orús.

Link to the paper: https://arxiv.org/abs/2608.03796

@misc{ryskulov2026efficientknowledgedistillationllms,
      title={Efficient Knowledge Distillation for LLMs: Offline Top-K Logits and a Fused Chunked KL Loss}, 
      author={Bakbergen Ryskulov and Iker García-Ferrero and David Montero and David Jansen and Ali Hashemi and Jezabel R. Garcia and Antonio Tiene and Román Orús},
      year={2026},
      eprint={2608.03796},
      archivePrefix={arXiv},
      primaryClass={cs.CL},
      url={https://arxiv.org/abs/2608.03796}, 
}