DisaggKD: Efficient Knowledge Distillation through Teacher-Student Disaggregation and a Fused Loss

📄 arXiv: 2610.17644v1 📥 PDF

作者: Ruixuan Tu, Yuyang Dong

分类: cs.LG, cs.CL, cs.DC

发布日期: 2026-10-13

💡 一句话要点

提出DisaggKD以解决知识蒸馏中的内存限制问题

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 知识蒸馏 模型分散 深度学习 内存优化 流式传输

📋 核心要点

  1. 现有的知识蒸馏方法将教师和学生模型放在同一GPU上，导致内存限制影响教师模型的规模。
  2. 论文提出DisaggKD，通过将教师和学生分散到不同节点，利用流式传输和分片损失来提高训练效率。
  3. 实验结果表明，分散布局在损失和速度上均优于同位置运行，且在更大的模型对比中表现出更好的内存效率。

📝 摘要（中文）

知识蒸馏（KD）框架通常将教师和学生模型放置在同一GPU上，因此受限于每个GPU的内存限制，无法容纳更大的教师模型。DisaggKD通过将教师和学生分散到不同节点上，教师副本向学生流式传输隐藏状态，利用行窗口和词汇分片的KD损失实现与全词汇损失相同的效果，而无需形成完整的词汇logits。此外，预取环节用于协调生产者和训练者的速度。通过闭合形式的内存和流量模型，确定每种布局的内存限制。在Qwen3.5-35B-A3B与9B的实验中，分散布局的损失与同位置运行的损失相匹配且速度更快；我们的融合损失也允许同位置运行在更少的节点上以更快的速度训练，而当学生跨多个节点时，分割方式则更为内存高效。在更大的模型对比中，使用融合损失的同位置运行仍然会出现内存溢出，而分割训练则不会。

🔬 方法详解

问题定义：论文旨在解决知识蒸馏中教师和学生模型因内存限制而无法使用更大教师模型的问题。现有方法在同一GPU上运行，导致内存不足，限制了模型的规模和训练效率。

核心思路：DisaggKD的核心思路是将教师和学生模型分散到不同的计算节点上，通过流式传输教师的隐藏状态到学生，并使用行窗口和词汇分片的KD损失来避免形成完整的词汇logits，从而提高内存利用率和训练速度。

技术框架：整体架构包括教师模型的分散副本、流式隐藏状态传输、KD损失计算模块以及预取环节。教师模型在不同节点上运行，学生模型接收流式数据并进行训练，同时通过预取机制协调数据流和训练过程。

关键创新：论文的主要创新在于提出了分散布局的知识蒸馏方法，能够在不形成完整词汇logits的情况下实现与同位置运行相同的损失，并且在速度和内存效率上均有显著提升。

关键设计：在设计中，采用了闭合形式的内存和流量模型来确定每种布局的内存限制，同时融合损失函数的设计使得同位置运行能够在更少的节点上以更快的速度进行训练。

🖼️ 关键图片

img_0
img_1
img_2

📊 实验亮点

实验结果显示，DisaggKD在Qwen3.5-35B-A3B与9B的对比中，分散布局的损失与同位置运行相匹配且速度更快。此外，在更大的模型对比中，使用融合损失的同位置运行出现内存溢出，而分散训练则成功避免了这一问题，展现了显著的内存效率提升。

🎯 应用场景

该研究的潜在应用领域包括大规模语言模型的训练、深度学习模型的优化以及资源受限环境下的模型部署。通过提高知识蒸馏的效率，DisaggKD可以帮助研究人员和工程师在有限的计算资源下训练更强大的模型，推动自然语言处理和计算机视觉等领域的发展。

📄 摘要（原文）

Knowledge distillation (KD) frameworks place teacher and student on the same GPUs, so per-GPU memory limits the teachers that fit. DisaggKD separates them across nodes: teacher replicas stream hidden states to the student, a KD loss over row windows and vocabulary shards equals the full-vocabulary loss in exact arithmetic without forming the full-vocabulary logits, and a prefetch ring paces producer and trainer. Closed-form memory and traffic models locate the memory limits of each layout. On the Qwen3.5-35B-A3B to 9B pair, the disaggregated layout matches the LM loss of the colocated run with the same loss and is faster than that run; our fused loss also lets a colocated run train on as few nodes at faster speed, while the split is memory efficient when the student spans several nodes. On the larger pairs, colocated runs with a fused loss still run OOM where the split trains, because teacher weights and student states fill the budget.