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

> Source: <https://github.com/CompactifAI/Full-Chunked-KL-Loss>
> Published: 2026-08-11 12:24:48+00:00

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}, 
}
```