KDFlow: Fast, Efficient System for Compressing Large Language Models

Researchers have introduced KDFlow, a framework designed to streamline the process of converting large language models into smaller, more efficient versions through knowledge distillation. The work addresses a key bottleneck in model compression: most existing systems treat the student (smaller) and teacher (larger) models the same way, which wastes computational resources.

How KDFlow Works#

KDFlow separates the two models' workflows. The teacher model runs through SGLang, a specialized inference system optimized for speed, while the student model trains using FSDP2, a framework built for efficient distributed training. This decoupled approach allows each model to use the technology best suited to its job.

The framework also cuts down on data transfer between processes. Rather than sending the teacher's full output logits—which can consume significant bandwidth—KDFlow transmits only the teacher's internal representations (hidden states) and recalculates the final outputs on the student side. This "zero-copy" technique reduces communication overhead without sacrificing compression quality.

Flexibility and Performance#

KDFlow supports both off-policy distillation (where the teacher's responses are pre-generated) and on-policy distillation (where the teacher adapts in real time). The system also handles cases where the student and teacher use different tokenizers—a constraint that often complicates distillation work.

In experiments, KDFlow delivered speedups ranging from 1.44× to 6.36× faster than competing knowledge distillation frameworks. By lowering engineering barriers, it enables researchers to test and scale model compression techniques with minimal overhead.

Code for the project is publicly available, making the approach accessible to the research community.