KDFlow: A User-Friendly and Efficient Knowledge Distillation Framework for Large Language Models

TL;DR

A novel framework for LLM distillation, termed KDFlow, which features a decoupled architecture and employs SGLang for teacher inference and incorporates KD algorithms for cross-tokenizer KD through highly extensible and user-friendly APIs.

Abstract

Knowledge distillation (KD) is an essential technique to compress large language models (LLMs) into smaller ones. However, despite the distinct roles of the student model and the teacher model in KD, most existing frameworks still use a homogeneous training backend (e.g., FSDP and DeepSpeed) for both models, leading to suboptimal training efficiency. In this paper, we present a novel framework for LLM distillation, termed \textbf{KDFlow}, which features a decoupled architecture and employs SGLang for teacher inference. By bridging the training efficiency of FSDP2 and the inference efficiency of SGLang, KDFlow achieves full utilization of both advantages in a unified system. Moreover, instead of transferring full logits across different processes, our framework only transmits the teacher's hidden states using zero-copy data transfer and recomputes the logits on the student side, effectively balancing the communication cost and KD performance. Furthermore, our framework supports both off-policy and on-policy distillation and incorporates KD algorithms for cross-tokenizer KD through highly extensible and user-friendly APIs. Experiments show that KDFlow can achieve \textbf{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: https://github.com/songmzhang/KDFlow

Figures (4)

• Figure 1: Training time per step and teacher forward time occupation under different distillation configurations. The teacher's MoE architecture poses challenges for FSDP, while being well-supported by SGLang.
• Figure 2: The overview of KDFlow. The whole framework is built based on Ray moritz2018ray and decouples the distillation pipeline by allocating the teacher model to SGLang and the student model to FSDP2. Solid and dashed arrows illustrate the data flow for off-policy and on-policy distillation, respectively. Notably, KDFlow transfers compact hidden states from the teacher rather than full logits to reduce communication overhead.
• Figure 3: Comparisons between different decoupled distillation approaches.
• Figure 4: Loss curves of KDFlow and the pure FSDP implementation when distilling Qwen3-30B-A3B to Qwen3-4B.