Optimizing Teacher-Student Partitioning for Scalable Knowledge Distillation on HPC Systems

📄 arXiv: 2606.27797v1 📥 PDF

作者: Adrian P. Dieguez, Victor Conchello Vendrell, Alex Batlle, Vinnam Kim, Jordi Ros-Giralt, Harris Teague

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

发布日期: 2026-06-26

💡 一句话要点

提出HPC感知的知识蒸馏优化方法以提升训练效率

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

关键词: 知识蒸馏 高性能计算 模型优化 深度学习 并行计算

📋 核心要点

  1. 现有的知识蒸馏方法未能充分利用教师和学生模型在内存和通信需求上的不对称性，导致效率低下。
  2. 本文提出了一种HPC感知的知识蒸馏方法，通过解耦教师和学生的分区来优化训练过程。
  3. 实验结果显示，该方法在样本处理速度上比现有TRL库提高了67%，显著提升了训练效率。

📝 摘要（中文）

知识蒸馏（KD）使得在大型教师模型指导下训练小型学生模型成为可能，而广泛采用的TRL库对这两种模型进行了对称处理，未能充分利用它们在内存占用和通信需求上的显著不对称性。本文提出了一种HPC感知的KD方法，能够高效地解耦教师和学生的分区。通过避免不必要的教师模型数据结构并选择最佳分割策略，我们的方法在样本处理速度上比TRL提高了67%。我们结合了模型的垂直和水平分区，推导出识别分割模式之间拐点的分析表达式。这些结果表明，通过拓扑感知的并行性利用教师-学生不对称性，显著加速了我们公司生产HPC集群上的GKD训练。

🔬 方法详解

问题定义：本文旨在解决现有知识蒸馏方法对教师和学生模型对称处理的问题，导致在内存占用和通信效率上的不足。

核心思路：提出了一种HPC感知的知识蒸馏方法，通过有效解耦教师和学生的分区，利用它们的内存和通信不对称性来提升训练效率。

技术框架：整体方法包括教师和学生模型的垂直和水平分区，首先分析模型的内存需求和通信需求，然后选择最佳的分割策略。

关键创新：最重要的创新在于提出了一个分析表达式，用于识别不同分割模式之间的拐点，从而优化模型的分区策略，这与现有方法的对称处理形成鲜明对比。

关键设计：在参数设置上，选择了适合HPC环境的分区策略，并设计了相应的损失函数以平衡教师和学生模型的训练过程。

🖼️ 关键图片

fig_0
fig_1

📊 实验亮点

实验结果表明，所提出的方法在样本处理速度上比TRL库提高了67%，显著提升了训练效率。这一提升不仅优化了资源利用，还加速了模型的训练过程，具有重要的应用价值。

🎯 应用场景

该研究的潜在应用领域包括大规模深度学习模型的训练，尤其是在高性能计算（HPC）环境下。通过优化知识蒸馏过程，可以在资源受限的情况下实现更高效的模型训练，具有重要的实际价值和广泛的应用前景。

📄 摘要（原文）

Knowledge Distillation (KD) enables training smaller student models under the guidance of larger teacher models, and the widely adopted TRL library implements it. Yet, TRL treats both models symmetrically, missing opportunities to exploit their pronounced asymmetry in memory footprint, and communication requirements. This paper presents an HPC-aware methodology for KD that decouples teacher and student partitioning efficiently. Our approach achieves up to 67% higher samples-per-second than TRL by avoiding unnecessary teacher-model data structures and selecting the best split strategy. We combine vertical and horizontal partitioning of models, deriving an analytical expression that identifies the existence of inflection points between splitting regimes. These results showed that exploiting teacher--student asymmetry through topology-aware parallelism notably accelerated GKD training on production HPC clusters at our company