{"slug": "chunked-kl-loss-running-knowledge-distillation-locally-in-less-6gb-vram", "title": "Chunked KL loss, running Knowledge Distillation locally in less <6GB VRAM", "summary": "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.", "body_md": "A small, self-contained CUDA benchmark for comparing the memory and throughput cost of three knowledge-distillation losses:\n\n**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.\n\nEach experiment runs in an isolated `torchrun`\n\nsubprocess. An out-of-memory\nerror or kernel failure is recorded as `FAILED`\n\n, and the remaining benchmark\nmatrix continues.\n\nFor every method, sequence length, and applicable chunk size, the benchmark reports:\n\n| Metric | Meaning |\n|---|---|\n`peak_vram_gib` |\nPeak CUDA memory allocated per GPU |\n`iteration_ms` |\nMean forward + backward iteration time |\n`tflops_per_gpu` |\nEstimated useful throughput per GPU |\n`status` |\n`OK` or `FAILED` |\n`error` |\nCaptured exception for failed experiments |\n\nSynthetic hidden states, output weights, and sparse teacher targets are generated deterministically and reused across methods so that every loss sees equivalent inputs.\n\n- Python 3.10+\n- PyTorch with CUDA support\n- One or more NVIDIA GPUs\n\nInstall PyTorch using the command recommended for your CUDA version at\n[pytorch.org](https://pytorch.org/get-started/locally/). The benchmark itself\nuses only PyTorch and the Python standard library.\n\nRun the default benchmark:\n\n```\npython benchmark.py\n```\n\nChoose the sequence lengths, chunk sizes, tensor-parallel size, iterations, and batch size:\n\n```\npython benchmark.py \\\n  --sequence-lengths 4096,8192,16384,32768,65536,131072,262144 \\\n  --chunk-sizes 8192,4096,2048,1024,512 \\\n  --tp 2 \\\n  --iterations 10 \\\n  --topk 100 \\\n  --batch-size 1\n```\n\nThe 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.\n\n| Option | Description |\n|---|---|\n`--sequence-lengths` |\nComma-separated sequence lengths to benchmark |\n`--chunk-sizes` |\nComma-separated chunk sizes for chunked methods |\n`--tp` |\nTensor-parallel world size |\n`--iterations` |\nNumber of measured iterations per experiment |\n`--batch-size` |\nPer-experiment batch size |\n`--topk` |\nNumber of top-K teacher logits to use |\n\nUse `python benchmark.py --help`\n\nfor the complete set of defaults and\nimplementation-specific options.\n\nRun the CPU equivalence tests before benchmarking:\n\n```\npython test_losses.py\n```\n\nThe tests compare loss values and hidden-state gradients across the three implementations.\n\nAll three methods evaluate the same forward KL against cached top-K teacher\nprobabilities. For teacher support\n\n```\nL_KL(p, z) = sum_{v in S} p_v log(p_v) - sum_{v in S} p_v z_v + M logsumexp(z)\n```\n\nOnly the log-normalizer spans the full vocabulary; the teacher-entropy and cross terms require just the K cached entries.\n\nThe reference implementation scatters the cached top-K probabilities into a\ndense teacher tensor and compares it with the student's dense log-softmax. It\nmaterializes the dense teacher and student log-probabilities in addition to the\nstudent logits, giving\n\nThis variant keeps the teacher sparse and evaluates the objective directly in\ncontiguous sequence chunks, avoiding both a dense teacher and a dense\nlog-softmax. However, the language-model head still produces the full student\nlogits, which autograd retains for backward. Peak memory therefore remains\n\nThis variant fuses the vocabulary projection into the loss. The forward pass\nprojects one sequence chunk at a time, accumulates the log-normalizer and sparse\nloss terms, and immediately discards the logits. The backward pass recomputes\neach chunk and applies the closed-form gradient\n\nTo regenerate the results figure, also install Matplotlib:\n\n```\npython -m pip install matplotlib\n```\n\nFor a SVG:\n\n```\npython plot_peak_vram.py \\\n  --input kd_benchmark_results.csv \\\n  --output assets/peak_vram.svg \\\n  --chunk-size 4096\n```\n\nFor a PNG:\n\n```\npython plot_peak_vram.py \\\n  --input kd_benchmark_results.csv \\\n  --output assets/peak_vram.png \\\n  --chunk-size 4096 \\\n  --dpi 180\n```\n\nThis repository accompanies *Distillation Tricks for Compact LLMs: Efficient\nOffline and Chunked Knowledge Distillation* by Bakbergen Ryskulov, Iker\nGarcía-Ferrero, David Montero, David Jansen, Ali Hashemi, Jezabel R Garcia,\nAntonio Tiene, and Román Orús.\n\nLink to the paper: [https://arxiv.org/abs/2608.03796](https://arxiv.org/abs/2608.03796)\n\n```\n@misc{ryskulov2026efficientknowledgedistillationllms,\n      title={Efficient Knowledge Distillation for LLMs: Offline Top-K Logits and a Fused Chunked KL Loss}, \n      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},\n      year={2026},\n      eprint={2608.03796},\n      archivePrefix={arXiv},\n      primaryClass={cs.CL},\n      url={https://arxiv.org/abs/2608.03796}, \n}\n```\n\n", "url": "https://wpnews.pro/news/chunked-kl-loss-running-knowledge-distillation-locally-in-less-6gb-vram", "canonical_source": "https://github.com/CompactifAI/Full-Chunked-KL-Loss", "published_at": "2026-08-11 12:24:48+00:00", "updated_at": "2026-08-11 12:43:02.973659+00:00", "lang": "en", "topics": ["machine-learning", "large-language-models", "ai-infrastructure", "developer-tools"], "entities": ["PyTorch", "NVIDIA", "CUDA"], "alternates": {"html": "https://wpnews.pro/news/chunked-kl-loss-running-knowledge-distillation-locally-in-less-6gb-vram", "markdown": "https://wpnews.pro/news/chunked-kl-loss-running-knowledge-distillation-locally-in-less-6gb-vram.md", "text": "https://wpnews.pro/news/chunked-kl-loss-running-knowledge-distillation-locally-in-less-6gb-vram.txt", "jsonld": "https://wpnews.pro/news/chunked-kl-loss-running-knowledge-distillation-locally-in-less-6gb-vram.jsonld"}}