Making Knowledge Distillation Cheap Enough to Run at Scale

Summary

Knowledge distillation trains a smaller student to match a larger teacher's capabilities, but conventional online training keeps both models loaded and materializes full vocabulary distributions at every step, making long-context recovery expensive. The proposed system caches the teacher's top-100 logits once and introduces a fused, chunked KL-divergence loss that avoids storing the full vocabulary-by-sequence grid or the student's full logits. In benchmarks with an 8K-token context on one H200, offline fused chunked KL used 58.3GB peak memory versus 102.8GB for online training, while all four methods reached near-identical training loss. At 32K tokens, the fully chunked loss reduced isolated kernel memory from 85.2 GiB to 5.45 GiB; distilling GPT-OSS 20B at that context cut step time from 57.0 to 12.23 seconds and reduced the setup from four GPU nodes to one.

Takeaways

1. 1

   Caching the teacher's top-100 logits removes the teacher from student training and allows one cache to be reused across multiple ablations.

2. 2

   Fused chunked KL avoids materializing the student's full vocabulary-by-sequence logits grid, reducing memory growth with sequence length; it trades this reduction for an extra output projection during backward computation.

3. 3

   At 32K tokens, fully chunked loss lowered isolated loss-kernel memory from 85.2 GiB to 5.45 GiB, while GPT-OSS 20B distillation fell from four GPU nodes to one and improved step time about fivefold.