针对大模型蒸馏,提出分层分区优化方法,显著提升HPC系统训练效率。
Optimizing Teacher-Student Partitioning for Scalable Knowledge Distillation on HPC Systems

- 分离师生模型分区策略,结合垂直与水平切分
- 相较TRL最高提升67%样本处理速度
- 适合大规模分布式训练场景的工程优化
知识蒸馏(KD)通过大模型指导小模型训练,广泛采用的TRL库对师生模型采取对称处理,忽略了二者在内存占用和通信需求上的显著差异。本文提出一种面向高性能计算(HPC)系统的蒸馏优化方法,高效解耦师生模型的分区策略。该方法通过避免教师模型不必要的数据结构并选择最优分割方案,相比TRL最高实现67%的样本每秒吞吐量提升。研究结合模型的垂直与水平分区,推导出解析表达式,识别出不同分割模式间的拐点。结果表明,通过拓扑感知的并行策略利用师生模型的不对称性,可显著加速生产级HPC集群上的通用知识蒸馏(GKD)训练。
原文摘要 · Abstract (English)
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
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。