KDFlow：面向大语言模型的用户友好且高效的知识蒸馏框架

# 面向大语言模型的友好高效知识蒸馏框架
来源：https://arxiv.org/html/2603.01875
宋明璋1,2,3，薛 张1,2，佟 张3，博杰 胡3，玉峰 陈1,2，金安 徐1,2
1北京交通大学计算机科学与技术学院，北京，中国
2交通领域大数据与人工智能重点实验室（北京交通大学），教育部
3微信，腾讯公司，中国
\{smzhang22,zhang\_xue,chenyf,jaxu\}@bjtu\.edu\.cn

###### 摘要

知识蒸馏（KD）对于大语言模型（LLMs）的模型压缩和后期训练至关重要，然而现有框架通常对教师模型和学生模型使用统一的训练后端，导致效率欠佳。本文提出 KDFlow，一种新颖的 LLM 蒸馏框架，其采用解耦架构，并利用 SGLang 进行教师推理。通过桥接 FSDP2 的训练效率和 SGLang 的推理效率，KDFlow 在统一的蒸馏流水线中充分发挥了两套系统的互补优势。此外，我们的框架通过零拷贝数据传输传递教师的隐藏状态，并在学生端重新计算 logits，从而实现了高效的**全词汇蒸馏**。更进一步，我们的框架同时支持离策略和在线策略蒸馏，并通过高度可扩展且用户友好的 API 融合了跨分词器算法。实验表明，KDFlow 相比现有 KD 框架实现了 **1.44×至 6.36×的加速**，使研究人员能够快速进行 LLM 蒸馏的原型设计与扩展，而工程开销极低。代码和文档已公开提供。

KDFlow: 面向大语言模型的友好高效知识蒸馏框架

宋明璋1,2,3††thanks：通讯作者。††thanks：在腾讯实习期间完成工作。，薛 张1,2，佟 张3，博杰 胡3，玉峰 陈1,2，金安 徐1,2
1北京交通大学计算机科学与技术学院，北京，中国
2交通领域大数据与人工智能重点实验室（北京交通大学），教育部
3微信，腾讯公司，中国
\{smzhang22,zhang\_xue,chenyf,jaxu\}@bjtu\.edu\.cn

## 1 引言

大语言模型（LLMs）已在多种任务中展现出卓越能力，但其庞大的参数量给资源受限环境下的部署带来了巨大挑战。知识蒸馏（KD）通过将知识从大型教师模型迁移至紧凑的学生模型，为应对这一挑战提供了有效途径——Hinton 等人（2015 (https://arxiv.org/html/2603.01875#bib.bib1)）。

图 1：不同蒸馏配置下每步训练时间及教师前向时间占比。教师的 MoE 架构给 FSDP 带来挑战，而 SGLang 能很好地支持该架构。
尽管 KD 的研究和应用已相当广泛，LLM 蒸馏的基础设施仍不理想。在典型的 KD 过程中，教师模型与学生模型扮演不同角色：教师仅执行前向传播（推理），而学生需要前向和反向传播（训练）。然而，现有框架（如 TRL——von Werra 等人 (2020 (https://arxiv.org/html/2603.01875#bib.bib3)) 或 MS-SWIFT——Zhao 等人 (2025 (https://arxiv.org/html/2603.01875#bib.bib2))）大多使用统一的训练后端来运行两个模型。这就造成了结构上的不匹配，因为一个专为梯度计算和优化器状态管理设计的训练引擎，并不适合教师这种以推理为主的工作负载。结果，LLM 蒸馏的吞吐量常常受限于低效的教师执行，尤其是在使用混合专家（MoE）教师时（见图 1 (https://arxiv.org/html/2603.01875#S1.F1)）。

特性TRLMS-SWIFTEasyDistillROLLSlimeverlKDFlow（我们的）KD 原生设计✓✓解耦后端✓✓✓✓离策略蒸馏✓✓✓✓✓在线策略蒸馏✓✓✓✓✓✓自蒸馏✓✓✓跨分词器蒸馏✓✓多教师蒸馏✓✓✓✓✓Logits/Logprobs全词表全词表全词表top-k词表采样词采样词top-k词表采样词全词表散度指标FKL,RKL,JSDFKL,RKL,JSDFKL,RKLoff-policy: FKL, RKL, JSD,倾斜 F(R)KL, AKL,on-policy: RKLRKLRKLFKL, RKL, JSD,倾斜 F(R)KL,AKL, TVD

表 1：KDFlow 与现有框架的对比。
为解决这一问题，我们提出 KDFlow，一个专为 LLM 蒸馏设计的高性能框架。KDFlow 将教师和学生后端解耦：学生使用 PyTorch FSDP2，教师使用高吞吐推理引擎 SGLang（Zheng 等人 (2024 (https://arxiv.org/html/2603.01875#bib.bib4))）。该架构的一个关键挑战是如何将完整的教师 logits 从 SGLang 进程传递到 FSDP 进程：直接传输 logits 由于规模巨大而不可行111对于 128 个序列，长度 4096，来自 Qwen3 模型的完整 BF16 logits 占用 128×4096×151936×2字节≈160GB 的内存。，而只传输 top-k logits 则会破坏损失函数的数学等价性。KDFlow 通过从 SGLang 收集紧凑的教师隐藏状态，并在学生端重新计算完整 logit 分布来解决这一问题，从而在降低通信开销的同时保留了标准 KD 的数学形式。与 TRL、MS-SWIFT 等主流框架相比，KDFlow 在离策略蒸馏中实现了 **1.44×至 6.36×的训练加速**，在混合专家（MoE）教师上获益更大。

尽管采用了解耦设计，KDFlow 仍将分布式通信进行抽象，并与标准的 Hugging Face 模型格式集成，用户只需几行配置即可启动 KD。

总体而言，KDFlow 的贡献包括：

- **高效架构**：KDFlow 通过使用 SGLang 为教师提供服务，并在进程间传输紧凑的隐藏状态，将教师推理与学生训练解耦，相比 TRL 和 MS-SWIFT 等统一后端框架，实现了 1.44× 至 6.36× 的加速。
- **全面性**：KDFlow 支持全面的 KD 特性（例如离/在线策略、跨分词器和多教师蒸馏），并提供多种内置散度指标和算法，使其成为 LLM 蒸馏的即开即用工具包。
- **友好设计**：KDFlow 是一个基于 FSDP2 的轻量级框架，将算法与整个蒸馏流水线解耦。

## 2 相关工作

### 2.1 大语言模型的知识蒸馏

知识蒸馏（KD）最早由 Hinton 等人（2015 (https://arxiv.org/html/2603.01875#bib.bib1)）提出，用于将大型模型压缩为小型模型。对于 LLMs，KD 通常分为黑盒蒸馏（Kim 和 Rush (2016 (https://arxiv.org/html/2603.01875#bib.bib5))）和白盒蒸馏（Zhang 等人 (2023 (https://arxiv.org/html/2603.01875#bib.bib32)); Gu 等人 (2023 (https://arxiv.org/html/2603.01875#bib.bib6)); Agarwal 等人 (2024 (https://arxiv.org/html/2603.01875#bib.bib7)); Ko 等人 (2024 (https://arxiv.org/html/2603.01875#bib.bib30)); Wu 等人 (2025 (https://arxiv.org/html/2603.01875#bib.bib31))）。白盒 KD 使用散度指标对齐学生和教师的输出分布，提供更丰富的监督信号，通常性能更优。该范式也已从基于静态数据集的离策略蒸馏扩展到在线策略蒸馏，即学生从其自身生成的数据中学习（Gu 等人 (2023 (https://arxiv.org/html/2603.01875#bib.bib6)); Agarwal 等人 (2024 (https://arxiv.org/html/2603.01875#bib.bib7)); Xiao 等人 (2026 (https://arxiv.org/html/2603.01875#bib.bib8))）。近期工作进一步研究了跨分词器 KD，以处理教师和学生模型之间的词表不匹配（Wan 等人 (2024 (https://arxiv.org/html/2603.01875#bib.bib15)); Boizard 等人 (2024 (https://arxiv.org/html/2603.01875#bib.bib12)); Zhang 等人 (2024 (https://arxiv.org/html/2603.01875#bib.bib18),2025 (https://arxiv.org/html/2603.01875#bib.bib17)); Cui 等人 (2025 (https://arxiv.org/html/2603.01875#bib.bib16)); Chen 等人 (2025 (https://arxiv.org/html/2603.01875#bib.bib13)); Minixhofer 等人 (2025 (https://arxiv.org/html/2603.01875#bib.bib14))）。尽管算法层面不断进步，KD 研究仍然缺乏灵活高效的框架，这正是本工作的动机所在。

图 2：KDFlow 概览。该框架基于 Ray（Moritz 等人 (2018 (https://arxiv.org/html/2603.01875#bib.bib27))），通过将教师模型分配给 SGLang，学生模型分配给 FSDP2，从而解耦蒸馏流水线。实线和虚线箭头分别表示离策略和在线策略蒸馏的数据流。值得注意的是，KDFlow 从教师端传输紧凑的隐藏状态而非完整 logits，以降低通信开销。
### 2.2 现有框架

现有的 LLM 训练与蒸馏框架提供了有用支持，但仍存在重要局限。TRL（von Werra 等人 (2020 (https://arxiv.org/html/2603.01875#bib.bib3))）和 EasyDistill（Wang 等人 (2025a (https://arxiv.org/html/2603.01875#bib.bib9))）在 KD 中使用同构的训练引擎，这导致教师前向传播期间硬件利用率不足。Slime（Zhu 等人 (2025 (https://arxiv.org/html/2603.01875#bib.bib11))）和 verl（Sheng 等人 (2025 (https://arxiv.org/html/2603.01875#bib.bib10))）使用 vLLM/SGLang 等引擎将推理和训练解耦，但它们并非专为 KD 设计，仅支持 logit 信息不完整的在线策略蒸馏。MS-SWIFT（Zhao 等人 (2025 (https://arxiv.org/html/2603.01875#bib.bib2))）同时支持离策略和在线策略蒸馏，但其实现相对较重。相比之下，KDFlow 是一个轻量级、高效的 KD 框架，覆盖完整的 KD 场景，详见表 1 (https://arxiv.org/html/2603.01875#S1.T1)。

## 3 KDFlow 的设计

KDFlow 的整体框架如图 2 (https://arxiv.org/html/2603.01875#S2.F2) 所示。为了解决使用单一同构引擎同时进行推理和训练的低效问题，我们采用了自上而下的解耦设计。本节将介绍系统架构、核心通信机制、蒸馏工作流以及算法抽象。

### 3.1 系统架构

KD 过程通常涉及多个阶段：教师推理、学生训练和学生 rollout，这些阶段在训练时需要不同的后端。因此，KDFlow 基于 Ray 构建，以高效管理分布式进程。如图 2 (https://arxiv.org/html/2603.01875#S2.F2) 中间部分所示，该架构由一个单一控制器（Trainer）和三个功能独立的 actor 组构成：

- **Trainer（单一控制器）**：Trainer 是中央协调器，管理数据集、控制训练循环，并组织不同 actor 组之间的数据流。KDFlow 支持 OffPolicyKDTrainer 和 OnPolicyKDTrainer。
- **RolloutActorGroup**：RolloutActorGroup 用于在线策略蒸馏期间学生模型的 rollout 过程。遵循 Slime（Zhu 等人 (2025 (https://arxiv.org/html/2603.01875#bib.bib11))）等前沿 RL 框架的做法，它使用 SGLang Router 连接多个 SGLang HTTP 服务器，实现负载均衡的推理。此外，KDFlow 采用同址（colocated）模式，并通过 CUDA 进程间通信（IPC）更新 SGLang 中的模型权重。
- **TeacherActorGroup**：TeacherActorGroup 管理多个 TeacherRayActor 实例，执行教师的前向传播。具体而言，我们在每个 TeacherRayActor 中初始化一个 SGLang 引擎222SGLang 引擎相比 SGLang HTTP 服务器与 NumPy ndarray 对象的兼容性更好。，以获得教师的隐藏状态。这使得 KDFlow 能够利用 SGLang 的高吞吐量和灵活并行策略进行教师模型推理。
- **StudentActorGroup**：StudentActorGroup 部署了 PyTorch FSDP2。它处理学生的标准训练过程，包括前向传播、反向传播和优化器状态管理，然后返回训练状态（如损失值）用于日志记录。

这种解耦架构确保了推理密集型的教师模型和训练密集型的学生模型分别在其最优化后端上运行，显著提高了硬件利用率。

图 3：不同解耦蒸馏方法的对比。
### 3.2 通过隐藏状态实现高效通信

在解耦架构中，教师和学生通常运行在单独的进程中，使得教师知识的传递成为关键瓶颈。直接传输完整的 logit 分布由于数据量巨大而变得不可行，而仅传输 top-k logits 则降低了带宽，但代价是破坏了蒸馏的数学等价性并降低了性能。为解决这一问题，KDFlow 采用隐藏状态传输和 logit 重计算，如图 3 (https://arxiv.org/html/2603.01875#S3.F3) 所示。TeacherActorGroup 不发送完整 logits，而是仅输出教师的最终隐藏状态，其维度（例如 4096）远小于词表大小（例如 151936）。KDFlow 进一步利用共享内存和 Ray 的共享对象机制，实现跨进程的零拷贝传输。每个 StudentRayActor 在接收到教师的隐藏状态后，使用教师的语言模型头部在本地重新计算完整的 logit 分布。这种设计大幅减少了通信量，同时保持了与基于 logit 的标准 KD 的数学等价性。

### 3.3 蒸馏工作流

在单一控制器（Trainer）的调度下，KDFlow 无缝支持两种主要的蒸馏工作流，展示了解耦 actor 之间灵活的数据路由：

**离策略蒸馏**（图 2 (https://arxiv.org/html/2603.01875#S2.F2) 中的实线）：学生从静态数据集中学习。Trainer 将提示和对应回答直接发送给 TeacherActorGroup，以获得教师的隐藏状态。这些隐藏状态连同输入随后被传递给 StudentActorGroup，用于计算蒸馏损失并更新学生权重。

**在线策略蒸馏**（图 2 (https://arxiv.org/html/2603.01875#S2.F2) 中的虚线）：学生从自身生成的数据中学习。首先，Trainer 将提示发送给 RolloutActorGroup 以生成回答。接着，这些提示-回答对被发送给 TeacherActorGroup，以获得教师的隐藏状态。然后，数据和隐藏状态流入 StudentActorGroup 进行梯度更新。最后，学生模型更新后的权重同步回 RolloutActorGroup，以确保下一次生成步骤使用最新的策略。

### 3.4 全面的抽象与算法

为了提供用户友好、开箱即用的工具包，KDFlow 将底层系统流水线与蒸馏算法严格分离。如图 2 (https://arxiv.org/html/2603.01875#S2.F2) 底部所示，KDFlow 内置支持多种 KD 算法和散度指标，包括前向 KL（FKL）、反向 KL（RKL）、Jensen-Shannon 散度（JSD）和总变差距离（TVD）。用户可以用极少的代码轻松实现自定义的蒸馏损失或算法，无需理解复杂的分布式通信逻辑。此外，KDFlow 原生支持跨分词器蒸馏。当教师和学生模型的词表不同时，KDFlow 会在计算散度之前自动进行 logit 对齐。这种抽象使得研究人员可以专注于算法创新，而非系统工程的细节。