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

A new CUDA benchmark compares three knowledge-distillation loss implementations—Full Dense KL, Forward-Chunked Loss, and Full Chunked KL—showing that the Full Chunked KL method fuses the output projection into the loss and recomputes logits during backward, enabling full-sequence logits to never be stored and running in less than 6GB VRAM. The benchmark, which uses PyTorch and supports tensor parallelism, reports peak memory, iteration time, and throughput for sequence lengths up to 262144 and chunk sizes down to 512, with CPU equivalence tests provided.

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: | Metric | Meaning | |---|---| peak vram gib | Peak CUDA memory allocated per GPU | iteration ms | Mean forward + backward iteration time | tflops per gpu | Estimated useful throughput per GPU | status | OK or FAILED | error | Captured exception for failed experiments | 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 https://pytorch.org/get-started/locally/ . 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. | Option | Description | |---|---| --sequence-lengths | Comma-separated sequence lengths to benchmark | --chunk-sizes | Comma-separated chunk sizes for chunked methods | --tp | Tensor-parallel world size | --iterations | Number of measured iterations per experiment | --batch-size | Per-experiment batch size | --topk | Number of top-K teacher logits to use | 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 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}, }