arXiv:2603.01875cs.CLcs.AI2026-03被引 5

让大模型蒸馏更快更省力,速度最高提升6.36倍。

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

  • 拆分师生模型训练架构,用SGLang加速教师推理
  • 只传隐状态不传完整输出,通信量减少且性能不降
  • 支持多种蒸馏方式,接口友好适合快速实验

知识蒸馏是将大语言模型压缩为小型模型的关键技术。然而,现有框架通常对师生模型使用相同的训练后端(如FSDP和DeepSpeed),导致训练效率不佳。本文提出一种新型大模型蒸馏框架KDFlow,采用解耦架构并使用SGLang进行教师模型推理。通过结合FSDP2的训练效率与SGLang的推理效率,实现统一系统中的双重优势。此外,本框架仅通过零拷贝传输教师模型的隐藏状态,学生端重新计算输出概率,有效平衡通信开销与蒸馏效果。框架还支持离策略与在线策略蒸馏,并通过高度可扩展的API实现跨分词器蒸馏。实验表明,KDFlow相比现有框架提速1.44×至6.36×,使研究人员能以极低工程成本快速原型与规模化部署。代码已开源:https://github.com/songmzhang/KDFlow。

原文摘要 · Abstract (English)

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

大模型压缩知识蒸馏高效训练LLM

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。