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

Tue, 11 Aug 2026 12:24:48 GMT

Hello everyone! I’ve been working on an efficient implementation of the KL-divergence loss that reduces its VRAM usage from quadratic to linear, using an approach very similar to Flash Attention: chunking and fusing the forward and backward passes. The loss is mathematically equivalent to the standard KL-divergence loss available in PyTorch. Until now, Knowledge Distillation required a huge amount of VRAM and was practically impossible to run locally at meaningful context lengths. With this implementation, the KL loss goes from requiring ~85 GB of VRAM at a 32K context length to ~5 GB. The…

Paper and discussion history. A mention is not an endorsement.

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

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