GradientFlow：90秒训练ImageNet/AlexNet的分布式通讯库

来自商汤和南洋理工的研究人员,实现了仅用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分钟，刷新当时最好成绩。