跨节点部署的数据并行深度神经网络（DNN）训练系统已广泛应用于大模型训练场景，而多节点间梯度同步的通信开销，是限制系统训练性能的主要瓶颈。Top-k稀疏化压缩是缓解通信瓶颈的主流高效方法，但传统Top-k算法存在显著性能缺陷，制约了模型训练效果与系统吞吐量。传统Top-k算法存在两大核心问题：一是DNN各层梯度为多维张量，算法筛选的最大k个梯度元素仅集中在部分维度，存在严重的维度缺失问题，丢失大量维度信息，造成模型收敛精度下降；二是传统方案需对梯度元素进行全局排序，无法充分发挥GPU内核并行能力，导致训练吞吐量降低。

实验室毕业博士明章强在华中科技大学胡燏翀老师和香港中文大学Patrick P.C. Lee老师的共同指导下，针对上述问题，设计了全维度Top-k稀疏化压缩方案ADTopk，用于高性能分布式DNN训练系统。该方案对梯度所有维度进行稀疏化压缩，彻底解决维度缺失问题，有效提升模型收敛精度；同时采用多重局部排序方式，提升GPU内核并行性，显著优化训练吞吐量。在期刊扩展研究中，提出进一步优化方案，扩展了交替压缩策略，实现收敛精度与训练速度的有效权衡，同时设计高效阈值估计算法，有效降低稀疏压缩带来的计算开销，完善了整体方案的训练效率。多组集群实验结果表明，相较于现有主流稀疏化方法，ADTopk可实现2.75%-22.31%的收敛精度提升，训练吞吐量提升幅度达20.8%-268.0%，综合性能优势显著。

图1 不同训练任务上的端到端收敛精度对比

图2 不同训练任务上的训练吞吐量对比

该研究成果以“Enabling Efficient All-Dimension Top-k Sparsification for High-Performance Distributed DNN Training Systems”为题，被计算机体系结构领域的CCF-A类期刊ACM TACO全文录用。本次录用的期刊论文是对前期CCF-A类系统结构顶级会议HPDC 2024成果的深度拓展，实现了研究成果的完善与升级。该研究工作得到了国家自然科学基金面上项目（No.62272185）、国家重点研发计划项目（No. 2022YFB4501300）和信息存储系统教育部重点实验室的支持。