Knowledge Distillation (KD) is a crucial technique for compressing Large Language Models (LLMs) into smaller ones. However, despite the distinct roles of the student and teacher models in KD, most existing frameworks still employ a homogeneous training backend (e.g., FSDP and DeepSpeed) for both models, resulting in suboptimal training efficiency. This paper presents a novel LLM distillation framework called KDFlow, featuring a decoupled architecture and utilizing SGLang for teacher inference. By bridging the training efficiency of FSDP2 and the inference efficiency of SGLang, KDFlow fully leverages both advantages in a unified system.
Moreover, instead of transferring full logits across different processes, our framework transmits only the teacher's hidden states using zero-copy data transfer and recomputes the logits on the student side, effectively balancing communication costs and KD performance. Furthermore, our framework supports both off-policy and on-policy distillation and integrates KD algorithms for cross-tokenizer KD through highly extensible and user-friendly APIs. Experiments demonstrate that KDFlow achieves 1.44$\times$ to 6.36$\times$ speedup compared to current KD frameworks, enabling researchers to rapidly prototype and scale LLM distillation with minimal engineering overhead. Code is available at GitHub.
Blogger's Review: KDFlow significantly enhances the efficiency of knowledge distillation through its decoupled architecture and efficient data transfer strategies. Its flexible API design and support for various distillation methods provide great convenience for researchers, especially in large-scale model training, greatly reducing time and resource consumption. Overall, the release of KDFlow is set to propel advancements in the LLM domain.