模型訓練
Offline Top-K Distillation Removes the Need for a Resident Teacher Model, Increasing Throughput by Up to 41% on a Single H200
New research first caches a small subset of the teacher model’s logits, then uses a fused, chunked KL kernel to avoid materializing the full-vocabulary tensor. This expands the distillation context length from 8K to 32K on a single GPU, though the speed and quality findings are still based on a limited set of model and hardware configurations.
The CompactifAI team has unveiled a method for reducing the cost of knowledge distillation for large language models by decoupling teacher inference from student training. Conventional online distillation requires the teacher and student to reside on the GPU simultaneously and regenerates the teacher distribution at every training step. The new workflow instead precomputes and stores the teacher’s Top-K logits at each position. Student training reads directly from this cache, freeing the GPU memory and compute resources occupied by the teacher.
The authors report that offline distillation achieves nearly identical training loss, runs about 29% faster per iteration, and increases throughput by up to 41% on a single NVIDIA H200. The Top-K cache requires additional disk space, but it can be reused across multiple student training runs and hyperparameter experiments, making it particularly suitable for model compression, post-quantization capability recovery, and large-scale ablation studies. It also turns expensive teacher inference into a schedulable, one-time data-generation task.
The second change is a fused, chunked KL divergence loss. A conventional implementation materializes the full student logits in a tensor whose shape is approximately sequence length multiplied by vocabulary size, then compares it with the teacher distribution. This can cause GPU memory spikes with long contexts. The new kernel computes the output projection, softmax, and KL divergence in chunks without materializing the full-vocabulary tensor, allowing peak memory usage to grow more smoothly with sequence length. As a result, the researchers expanded the training context length from approximately 8,192 tokens to 32,768 tokens on a single GPU. They also used output-head tests to evaluate memory usage and iteration-time scaling from 4K to 256K tokens.
The team has released the Full-Chunked-KL-Loss implementation, although the paper notes that related patent applications are pending. At this stage, kernel microbenchmarks should not be treated as evidence of end-to-end acceleration for a complete model. Top-K truncation may also discard low-probability signals from the teacher. Engineering teams should next examine different vocabulary sizes, model families, K values, cache I/O behavior, and whether communication in multi-GPU training replaces GPU memory as the new bottleneck.