我国科学家发明混合并行新算法,训练 AI 速度近乎翻倍
中国科学报今天(5月29日)报道,我国科学家联合发明混合并行新算法GroPipe,首次融合流水线模型并行与数据并行,训练AI速度近乎实现翻倍。
IT之家援引博文介绍,该算法由西北农林科技大学信息工程学院智能计算与农业信息系统团队刘斌教授带领,携手美国纽约州立大学和云南大学,相关论文联合发表在国际计算机体系结构领域权威期刊《IEEE Transactions on Computers》(TC,CCF A类期刊)上。
西北农林科技大学信息工程学院教授刘斌担任第一作者,纪泽宇老师与云南大学副教授何臻力共同担任通讯作者,纽约州立大学教授李克勤(欧洲科学院院士,并行与分布式计算领域公认的顶尖学者,国家特聘教授)作为重要合作者。
研究指出,在训练大型深度卷积神经网络(DCNNs)方面,使用越来越大的数据集来提高模型精度已变得非常耗时,数据并行(DP)和流水线模型并行(PMP)等分布式训练方法虽然提供了可行解决方案,但面临负载不平衡和显著的通信开销等挑战。
为此团队提出了GroPipe算法架构,协同整合PMP和DP,采用了基于性能预测技术的自动模型分割算法,确保负载平衡并便于在PMP中进行定量性能评估。
GroPipe方法整体框架图。西北农林科技大学供图
该方法首次将流水线模型并行与数据并行相融合,构建“组内流水线+组间数据并行”的分层训练架构,并通过自动模型划分算法(AMPA)实现计算负载的动态均衡调度,大幅提升GPU资源利用率。
在一台8-GPU服务器上广泛测试后,发现在ImageNet数据集上,GroPipe方法相较于主流方案(如DP、Torchgpipe、DAPPLE和DeepSpeed)表现出卓越优势:ResNet系列平均加速比达42.2%,VGG系列达79.2%;在BERT-base模型训练中,性能提升最高可达51%。
