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-07-20
💡 一句话要点
提出KDFlow以提升大语言模型蒸馏效率
🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)
关键词: 知识蒸馏 大语言模型 高效推理 解耦架构 零拷贝传输 API设计 模型压缩
📋 核心要点
- 现有知识蒸馏框架对学生和教师模型使用同质训练后端，导致训练效率低下。
- KDFlow框架采用解耦架构，利用SGLang进行教师推理，优化了训练和推理效率。
- 实验结果显示KDFlow在速度上相比现有框架提升1.44至6.36倍，降低了工程开销。
📝 摘要（中文）
知识蒸馏（KD）是将大型语言模型（LLMs）压缩为更小模型的重要技术。然而，现有框架通常对学生模型和教师模型使用同质训练后端，导致训练效率低下。本文提出了一种新颖的LLM蒸馏框架KDFlow，采用解耦架构并使用SGLang进行教师推理。KDFlow有效结合了FSDP2的训练效率和SGLang的推理效率，充分利用两者优势。此外，框架仅传输教师的隐藏状态，重新计算学生端的logits，从而平衡通信成本与KD性能。实验表明，KDFlow相比现有KD框架可实现1.44倍至6.36倍的加速，帮助研究人员快速原型和扩展LLM蒸馏，工程开销最小。
🔬 方法详解
问题定义：本文旨在解决现有知识蒸馏框架在训练效率上的不足，尤其是对学生模型和教师模型使用同质训练后端的问题，这导致了资源的浪费和性能的低下。
核心思路：KDFlow通过解耦架构设计，分别优化教师模型的推理和学生模型的训练，采用SGLang进行高效推理，从而提升整体蒸馏效率。
技术框架：KDFlow的整体架构包括教师模型和学生模型的解耦训练，教师模型使用SGLang进行推理，学生模型则通过零拷贝数据传输接收教师的隐藏状态，并在本地重新计算logits。
关键创新：KDFlow的主要创新在于其解耦架构和零拷贝数据传输机制，这与传统框架的同质训练方式形成鲜明对比，显著提升了训练和推理的效率。
关键设计：框架支持离线和在线蒸馏，设计了高度可扩展的API以支持跨tokenizer的蒸馏，同时在参数设置和损失函数上进行了优化，以确保蒸馏过程的高效性和准确性。
🖼️ 关键图片
📊 实验亮点
实验结果表明，KDFlow在蒸馏速度上实现了1.44倍至6.36倍的显著提升，相较于现有的知识蒸馏框架，展示了更高的训练效率和更低的工程开销，极大地促进了大语言模型的快速原型开发。
🎯 应用场景
KDFlow框架在大语言模型的蒸馏过程中具有广泛的应用潜力，能够有效降低模型的计算资源需求，适用于需要快速部署和迭代的自然语言处理任务。其用户友好的API设计使得研究人员和开发者能够轻松集成和使用，推动了大语言模型在实际应用中的普及。
📄 摘要（原文）
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:this https URL