CompactifAI/Full-Chunked-KL-Loss

CompactifAI/Full-Chunked-KL-Loss

Description: Fully chunked KL-Loss kernel, that massively reduces VRAM usage

Language: Python

Stars: 0

Forks: 0

Open issues: 0

Created: 2026-08-03T13:09:50Z

Pushed: 2026-08-04T15:07:49Z

Default branch: main

Fork: no

Archived: no

README:

PyTorch KL Loss Benchmark

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.


What it measures

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.

Requirements

• 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.

Quick start

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.

Command-line options

| 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.

Validate the losses

Run the CPU equivalence tests before benchmarking:

python test_losses.py

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

Methods

All three methods evaluate the same forward KL against cached top-K teacher probabilities. For teacher support $S$, retained mass $M=\sum_{v\in S}p_v$ (which is not renormalized), and student logits $z$, the per-token objective is

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.

Full Dense KL

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 $O(SBV)$ peak memory for sequence length $S$, batch size $B$, and vocabulary size $V$. It is the implementation closest to conventional online distillation and serves as the correctness baseline.

Forward-Chunked Loss

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 $O(SBV)$, although with a much smaller constant than the dense baseline. It is the fastest of the three methods in the paper's short-context profiling.

Full Chunked KL

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 $M\,\mathrm{softmax}(z)-p$. Vocabulary-sized intermediate memory is bounded by the chunk size rather than the full sequence length; the trade-off is one extra output projection during backward. At long contexts, the lower memory requirement can also reduce model parallelism and inter-node communication.

Get Results Figure

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

Reference

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. arXiv link forthcoming.

License