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 loss is also ~3× faster at long context lengths.
This unlocks the possibility of performing Knowledge Distillation locally, allowing you to train small student models from large teacher models without needing access to an H200 node.
One caveat is that, because the implementation requires patching the model’s forward pass to chunk the LM-head computation, it cannot be used as a direct drop-in replacement for PyTorch’s KL loss. However, this patch can easily be implemented with Claude Code, Kimi K3, GLM, etc. in just a few minutes.
The loss is intended to work with cached top-k logits: you pre-compute the logits from the teacher model and store only the top-100 logits. In the paper, we show that this produces an almost identical loss to using the full distribution, while requiring significantly less memory and compute.
If you’re interested in the details of how the loss works, we’ve also published a paper on arXiv:
https://arxiv.org/abs/2608.03796
ikergarcia1996•1h ago
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 loss is also ~3× faster at long context lengths.
This unlocks the possibility of performing Knowledge Distillation locally, allowing you to train small student models from large teacher models without needing access to an H200 node.
The code is open-source: https://github.com/CompactifAI/Full-Chunked-KL-Loss/
One caveat is that, because the implementation requires patching the model’s forward pass to chunk the LM-head computation, it cannot be used as a direct drop-in replacement for PyTorch’s KL loss. However, this patch can easily be implemented with Claude Code, Kimi K3, GLM, etc. in just a few minutes.
The loss is intended to work with cached top-k logits: you pre-compute the logits from the teacher model and store only the top-100 logits. In the paper, we show that this produces an almost identical loss to using the full distribution, while requiring significantly less memory and compute.
If you’re interested in the details of how the loss works, we’ve also published a paper on arXiv: https://arxiv.org/abs/2608.03796