大型深度学习模型的分布式训练常常会出现故障,因此 checkpointing(检查点机制)被广泛用于故障恢复。当前关于快速恢复的前沿研究主要集中在频繁检查点机制上,但这种方式会生成大量检查点,导致成本显著增加,进而降低训练性能。最近,差分检查点机制被提出以降低成本,但其应用仅限于推荐系统,在通用分布式训练系统中的应用尚未得到探索。
实验室即将入学攻读博士的本科生姚晨炫和博士生刘飞帆在胡燏翀老师的指导下,共同提出了一种高效的频繁检查点框架 LowDiff。该框架通过复用通用分布式训练中常用的压缩梯度作为差分检查点,从而降低成本。此外,LowDiff 设计了批量梯度写入优化,并通过调整检查点频率和批处理大小来优化检查点性能。在多种模型上的实验表明,与最先进的方法相比,LowDiff 最多可减少 89.2% 的训练时间。

图1 LowDiff复用梯度示意图

图2 LowDiff整体结构图
该研究成果被SC 2025录用。SC是系统结构领域最重要的国际会议之一,也是中国计算机学会推荐的A类会议。该研究工作得到了国家自然科学基金(No. 62272185)和国家重点研发计划青年科学家项目(No.2021YFB0301400)的资助。