Chunked KL loss: Reducing 32K Context to 5GB

Chunked KL loss divides computation into chunks similarly to Flash Attention and fuses forward and backward passes, reducing KL loss memory usage from quadratic to linear.

• 32K context VRAM: approx. 85GB → 5GB
• Approx. 3x faster processing compared to existing implementations for long contexts
• Mathematically equivalent to standard KL loss in PyTorch
• Supports local training for distilling knowledge from large teacher models to small student models

It also supports caching and using top-100 logits instead of the full logits from the teacher model; the paper states this yields a loss nearly identical to the full distribution. However, it is not a direct drop-in replacement for PyTorch's KL loss because the model's lm-head forward pass must be patched to support chunking.

The code is available on GitHub, and implementation details along with top-k logits experiment results are documented in the arXiv paper.