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

arXiv:2603.01875v4 Announce Type: replace-cross Abstract: Knowledge distillation (KD) is widely used to compress and post-train large language models (LLMs), yet many existing frameworks execute teacher inference with the same training-oriented backend as student optimization, leading to suboptimal efficiency. In this paper, we propose KDFlow, a novel framework for LLM distillation that features a decoupled architecture and employs SGLang for teacher inference. KDFlow combines SGLang for teacher inference with PyTorch FSDP2 for student optimization, allowing each model to run on a backend tailored to its workload. To enable efficient full-vocabulary distillation in this decoupled architecture, KDFlow transfers the teacher's final hidden states via Ray's object store and recomputes teacher logits on each student worker using a frozen copy of the teacher's output head. Furthermore, our framework supports both off-policy and on-policy distillation and incorporates cross-tokenizer algorithms through highly extensible and user-friendly APIs. Experiments show that KDFlow achieves a 1.44imes to 6.36imes speedup over MS-SWIFT in off-policy distillation and a 1.43imes to 1.75imes speedup over verl in on-policy distillation. KDFlow further scales to 64 GPUs, achieving 3.68imes and 2.52imes strong-scaling speedups in two representative model configurations. The code and documentation are publicly available.

Source: arXiv cs.LG | 2026-09-09