KDFlow: A User-Friendly and Efficient Knowledge Distillation Framework for Large Language Models
作者: Songming Zhang, Xue Zhang, Tong Zhang, Bojie Hu, Yufeng Chen, Jinan Xu
分类: cs.CL, cs.AI, cs.LG
发布日期: 2026-03-02
备注: 8 pages, 4 figures, 3 tables, code is available at: https://github.com/songmzhang/KDFlow
🔗 代码/项目: GITHUB
💡 一句话要点
提出KDFlow框架，通过解耦架构和SGLang加速大语言模型知识蒸馏。
🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)
关键词: 知识蒸馏 大语言模型 模型压缩 SGLang 解耦架构 零拷贝数据传输 高效训练 用户友好
📋 核心要点
- 现有知识蒸馏框架对教师和学生模型采用同构训练后端，忽略了两者差异，导致训练效率低下。
- KDFlow采用解耦架构，教师模型使用SGLang进行高效推理，学生模型使用FSDP进行高效训练，充分利用各自优势。
- 实验表明，KDFlow相比现有框架可实现1.44倍至6.36倍的加速，降低了LLM蒸馏的工程成本。
📝 摘要（中文）
知识蒸馏（KD）是将大型语言模型（LLM）压缩为较小模型的重要技术。然而，尽管学生模型和教师模型在KD中作用不同，但现有框架大多对两者使用同构训练后端（如FSDP和DeepSpeed），导致训练效率欠佳。本文提出了一种新的LLM蒸馏框架KDFlow，其特点是解耦架构，并采用SGLang进行教师模型推理。通过桥接FSDP2的训练效率和SGLang的推理效率，KDFlow在一个统一的系统中充分利用了两者的优势。此外，该框架仅使用零拷贝数据传输教师模型的隐藏状态，并在学生端重新计算logits，从而有效平衡了通信成本和KD性能。该框架还支持离策略和在线策略蒸馏，并通过高度可扩展且用户友好的API集成了用于跨分词器KD的KD算法。实验表明，与现有KD框架相比，KDFlow可实现1.44倍至6.36倍的加速，使研究人员能够以最小的工程开销快速构建和扩展LLM蒸馏。
🔬 方法详解
问题定义：现有知识蒸馏框架在训练大型语言模型时，通常采用同构的训练后端（例如FSDP或DeepSpeed）来训练教师模型和学生模型。这种方法忽略了教师模型和学生模型在知识蒸馏过程中的不同角色，导致训练效率低下，无法充分利用硬件资源。尤其是在教师模型推理阶段，同构训练后端效率较低，成为瓶颈。
核心思路：KDFlow的核心思路是解耦教师模型和学生模型的训练过程，针对性地选择最适合的训练/推理框架。具体来说，教师模型采用SGLang进行高效推理，而学生模型则继续使用FSDP等框架进行训练。通过这种解耦，可以充分利用SGLang在推理方面的优势和FSDP在训练方面的优势，从而提高整体的知识蒸馏效率。
技术框架：KDFlow的整体架构包含教师模型推理模块和学生模型训练模块。教师模型推理模块使用SGLang进行推理，并将隐藏状态通过零拷贝数据传输到学生模型训练模块。学生模型训练模块接收到教师模型的隐藏状态后，重新计算logits，并使用知识蒸馏损失函数进行训练。该框架支持离策略和在线策略蒸馏，并提供了可扩展的API，方便用户集成不同的知识蒸馏算法。
关键创新：KDFlow最重要的技术创新点在于其解耦架构。与现有框架不同，KDFlow不再使用同构的训练后端，而是根据教师模型和学生模型的不同角色，分别选择最适合的框架。此外，KDFlow还采用了零拷贝数据传输技术，减少了通信开销。
关键设计：KDFlow的关键设计包括：1) 使用SGLang进行教师模型推理，充分利用其高效的推理能力；2) 使用零拷贝数据传输技术，减少教师模型和学生模型之间的通信开销；3) 提供可扩展的API，方便用户集成不同的知识蒸馏算法和损失函数；4) 支持跨分词器的知识蒸馏，提高了框架的通用性。
🖼️ 关键图片
📊 实验亮点
实验结果表明，KDFlow框架相比于现有的知识蒸馏框架，能够实现1.44倍到6.36倍的加速。这一显著的性能提升使得研究人员能够更快地迭代模型，并以更低的成本训练出高性能的小型语言模型。此外，KDFlow框架的用户友好型API也降低了LLM蒸馏的工程门槛。
🎯 应用场景
KDFlow框架可应用于各种需要将大型语言模型压缩为小型模型的场景，例如移动设备上的自然语言处理、边缘计算设备上的智能应用、以及对延迟敏感的在线服务。该框架能够加速模型蒸馏过程，降低计算成本，并提高部署效率，具有广泛的应用前景。
📄 摘要（原文）
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