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


翻译:知识蒸馏(KD)是将大型语言模型(LLM)压缩为更小模型的关键技术。然而,尽管学生模型与教师模型在知识蒸馏中扮演着截然不同的角色,现有大多数框架仍对两者采用同质化的训练后端(如FSDP和DeepSpeed),导致训练效率欠佳。本文提出了一种新型LLM蒸馏框架,命名为\textbf{KDFlow},该框架采用解耦架构,并利用SGLang进行教师模型推理。通过桥接FSDP2的训练效率与SGLang的推理效率,KDFlow在统一系统中实现了两者优势的充分融合。此外,该框架无需跨进程传输完整logits,而是利用零拷贝数据传输仅传递教师隐状态,并在学生端重新计算logits,有效平衡了通信开销与知识蒸馏性能。进一步地,KDFlow支持离线策略蒸馏与在线策略蒸馏,并通过高度可扩展且用户友好的API整合了跨分词器知识蒸馏算法。实验表明,相比现有知识蒸馏框架,KDFlow可实现\textbf{1.44倍至6.36倍}的加速,使研究人员能够以极低的工程开销快速原型验证与规模化部署LLM蒸馏。代码开源地址:https://github.com/songmzhang/KDFlow

0
下载
关闭预览

相关内容

大语言模型同策略蒸馏研究综述
专知会员服务
20+阅读 · 4月5日
大型语言模型的知识蒸馏综述:方法、评估与应用
专知会员服务
79+阅读 · 2024年7月4日
KnowledGPT:基于知识库的检索和存储访问增强大型语言模型
专知会员服务
27+阅读 · 2021年10月12日
模型压缩 | 知识蒸馏经典解读
AINLP
11+阅读 · 2020年5月31日
AI新视野 | 数据蒸馏Dataset Distillation
人工智能前沿讲习班
31+阅读 · 2019年6月14日
国家自然科学基金
18+阅读 · 2017年12月31日
国家自然科学基金
10+阅读 · 2015年12月31日
国家自然科学基金
9+阅读 · 2015年12月31日
国家自然科学基金
12+阅读 · 2015年12月31日
国家自然科学基金
1+阅读 · 2015年12月31日
国家自然科学基金
12+阅读 · 2014年12月31日
国家自然科学基金
5+阅读 · 2014年12月31日
国家自然科学基金
4+阅读 · 2014年12月31日
国家自然科学基金
11+阅读 · 2012年12月31日
VIP会员
相关主题
最新内容
综述 | 3D场景图:开放挑战与未来方向
专知会员服务
7+阅读 · 6月22日
21世纪的无人机战争
专知会员服务
4+阅读 · 6月22日
《量子技术的军事任务技术适配与利用》
专知会员服务
5+阅读 · 6月22日
美国从乌克兰无人机战争中学习经验
专知会员服务
7+阅读 · 6月21日
ICML 2026 | 面向视觉语言模型的语义鲁棒性认证
专知会员服务
5+阅读 · 6月21日
相关基金
国家自然科学基金
18+阅读 · 2017年12月31日
国家自然科学基金
10+阅读 · 2015年12月31日
国家自然科学基金
9+阅读 · 2015年12月31日
国家自然科学基金
12+阅读 · 2015年12月31日
国家自然科学基金
1+阅读 · 2015年12月31日
国家自然科学基金
12+阅读 · 2014年12月31日
国家自然科学基金
5+阅读 · 2014年12月31日
国家自然科学基金
4+阅读 · 2014年12月31日
国家自然科学基金
11+阅读 · 2012年12月31日
Top
微信扫码咨询专知VIP会员