Efficient Knowledge Distillation for LLMs: Offline Top-K Logits and a Fused Chunked KL Loss

2026-08-04

AI Summary

This work addresses the challenge of deploying large language models in resource-constrained settings, where conventional knowledge distillation incurs prohibitive memory and computational costs. The authors propose an efficient offline knowledge distillation approach that caches the teacher model’s top-K logits and introduces a block-wise fused KL divergence loss to substantially reduce peak training memory. By integrating sequence packing with GPU-specific optimizations, the method achieves a 29% speedup and 41% higher throughput on a single H200 GPU, supports context lengths up to 32,768 tokens, and exhibits linear memory scaling with respect to sequence length. The core implementation has been made publicly available.

Abstract

Small language models are often the only option for deployment under tight latency, cost, and on-premises constraints, but they are rarely trained from scratch: a compressed model is usually recovered through knowledge distillation (KD). This recovery step largely decides the final quality, yet it is expensive. We present a practitioner's study of how to make distillation training efficient, organised around two systems contributions. First, we show that offline KD (caching the teacher's top-$K$ logits once and training the student against the cache) matches online distillation at near-identical training loss while removing the teacher from memory, running about 29\% faster per iteration, and reaching up to 41\% higher throughput on a single H200 GPU. Second, we introduce a \emph{fused, chunked KL loss} that never materialises the full vocabulary-sized logit tensor, making peak memory linear in the sequence length. This removes the memory spike that otherwise caps context length and lets us train at four times the context (32{,}768 tokens) on a single GPU. A separate output-head-only toy benchmark isolates the loss kernel and confirms its memory and iteration-rate scaling from 4K to 256K tokens. Together these make large-scale healing and hundreds of ablations affordable. We also report supporting ablations on loss design and sequence packing. We release our chunked-loss implementation: https://github.com/CompactifAI/Full-Chunked-KL-Loss.

Problem

• Knowledge Distillation
• Large Language Models
• Training Efficiency
• Memory Optimization
• KL Loss

Innovation

• knowledge distillation
• offline top-K logits
• fused chunked KL loss
• memory efficiency
• large language models