来自商汤和南洋理工的研究人员,实现了仅用1.5分钟训练ImageNet/AlexNet。尽管多种物体识别方法在ImageNet上已经达到了很高的准确率,但模型训练过程仍然非常耗时,缩短模型训练时间是AI研究人员必须面临的工程挑战。为减少训练时间,作者开发了用于分布式数据通讯后端GradientFlow,辅以一系列的网络通讯优化技巧实现惊人的效率提升。
论文地址:
https://arxiv.org/abs/1902.06855
主要贡献
本文主要介绍多个在分布式训练中常用的网络优化技巧,包括多层次环形allreduce、混合精度训练、懒惰allreduce、计算和通信重叠等。
环形 allreduce
环形allreduce通过持续通信提升性能,在大规模GPU集群中,可能无法充分利用带宽。作者提出了分层思路解决带宽利用问题,将N个GPU分为M组,分为三个阶段完成:1)每组将中间结果存储到主GPU上。2)主GPU启动常规环形allreduce操作。3)主GPU在组内广播,将最终结果传送到所有GPU。
混合精度训练
在深度学习训练中,参数和梯度可以使用半精度(FP16)浮点数替代单精度(FP32)浮点数,显著降低GPU显存占用和带宽压力。为了避免准确率的下降,本文在训练中使用了混合精度技术,前向、后向计算阶段使用低精度的FP16,模型更新阶段则使用高精度的FP32。
懒惰allreduce
作者发现每次只传输小的张量,NCCL不能充分利用网络带宽。懒惰allreduce将多个张量融合为一个,基于生成的顺序放在内存池中,等待张量的显存占用达到一定阈值才执行allreduce操作,直接在融合内存池执行操作,不会带来额外的内存复制开销。
计算和通信重叠
在分布式训练中,作者通过下层来重叠上层的所有操作,来减少通信时间。当完成反向计算时,layer i层的梯度张量立即生成,并不会通过反向计算来改变layer j。因此,作者强制执行每一层一旦生成梯度张量,就进行通信。这一层执行allreduce的时间可能与反向传播的实际重叠,文中叫做计算和通信的重叠。
实验验证
作者采用imageNet-1k作为实验数据,测试了AlexNet和ResNet-50。系统设计结合了环 allreduce、混合精度、计算和通信的交叠、懒惰 allreduce。每个GPU 批大小为128。在Alexnet完成训练95epochs,Resnet50训练90个epochs。实验结果显示,本系统在保持精度几乎没有下降的前提下,实现了训练时间由原来的2.6分钟缩短为1.5分钟。
结论
作者提出了GradientFlow显著提高了网络的训练性能,融合环形 allreduce、混合精度、计算和通信的交叠等技术。作者发现真实的系统应用这些方法和理论效果还有一定差距。为了进一步减少损失,作者提出了lazy allreduce来通过融合多个通信操作来提高吞吐量,同时设计只传输重要梯度块的粗粒度稀疏通信技术来减少网络障碍。实验结果显示,GradientFlow可以显著的减少分布式训练的时间。使用512GPUs训练ImageNet/Alexnet,实验可以达到410.2的加速比,完成95epochs只需要1.5分钟,刷新当时最好成绩。