应对大规模训练中的通信挑战：压缩及调度

摘要

随着各个领域数据的不断积累和规范化，越来越多的研究者开始着手解决更大规模的问题，不可避免地会提出或使用更多大规模的模型，分布式训练几乎成为了一个必选项。分布式系统面临的一大挑战是通信问题，特别是对于AI训练任务，使用的是高算力的设备，如果互联是低速的，那么多台设备组成的分布式系统性能甚至会低于单台设备。随着越来越多的大规模训练需求的产生，如何更好地应对分布式训练系统的通信问题已成为研究热点。我们在之前的文章中介绍了分布式训练中的通信问题，并从减少用于同步的通信和采用更利于通信的拓扑结构两个方面介绍了应对大规模训练中通信问题的方法。本文首先回顾分布式训练中的通信问题，然后进一步介绍对通信数据进行压缩以减少通信数据量，以及通过调度来使通信过程和计算过程尽量并行的相关技术。

分布式训练中的通信问题

我们首先回顾一下模型训练的主要流程。模型训练的典型工作流如图1，一次迭代的主要步骤为1）采样小批次数据（mini-batch）；2）读入数据，进行前向过程（Feed-Forward），计算目标函数的损失值；3）进行反向过程（Backward Propagation），为模型参数计算梯度；4）根据优化器算法更新（Update）模型参数。在分布式系统中，模型梯度或参数会分布在不同的设备上，故在更新模型参数前，需对不同设备上的梯度或参数进行聚合操作（Gradient/Model Aggregation）。

图1 模型训练的典型工作流(修改自[1])

由于模型训练使用的都是具有强大算力的设备，分布式训练过程中用于维持不同设备模型同步的通信需求是十分高频的需求，当前互联技术所能提供的带宽很难直接满足训练过程中设备间的数据传输。为应对这一通信挑战，我们可以从四个方面入手：1）放宽同步要求，减少用于同步的通信；2）采用更利于通信的拓扑结构；3）对通信数据进行压缩，减少通信数据量；4）通过调度来使通信过程和计算过程尽量并行。我们在应对大规模训练中的通信挑战：同步与拓扑一文中介绍了前两个方面的相关技术，本文将进一步介绍后两个方面的相关技术。

减少通信数据量

为降低分布式训练过程中通信压力，一个直接的办法是减少需要传输的数据量，常见的两种技术如图2，一方面可以进行量化(Quantization)，使用更少的比特位来表示每个数据，另一方面可以采用稀疏化(Sparsification)的手段，仅传输部分数据。这两种技术都会带来数据传输过程中的信息损失，因此相关算法设计的核心在于保障关键信息的传输，以保证模型训练的收敛性。

图2 量化(Quantization)和稀疏化(Sparsification)[1]

下面主要以梯度通信为例，分别介绍使用两种技术的代表性工作。

1. 量化(Quantization)

使用梯度量化的SGD算法典型流程如图3，图中Encode和Decode即为量化的编码和解码过程。梯度量化算法的核心在于设计Encode过程，通常需要保证量化后的梯度是原始梯度的无偏估计，而Decode过程通常不需要严格恢复原始信息。

图3 使用梯度量化的SGD算法流程[2]

这里以QSGD[2]为例来介绍Encode过程，QSGD使用一种随机量化方法，对于梯度  ，量化函数定义如下，

这一设计可以保证  ，进一步可以推出  ，即量化的梯度是原始梯度的无偏估计。

使用QSGD进行模型训练的效果如图4，图中每个bar底部深色区域为通信时间，上层浅色区域为计算时间，可以看到量化后通信时间大大减少。

图4 QSGD训练效果[2]

2. 稀疏化(Sparsification)

训练过程中产生的梯度通常是32位的浮点数，使用量化最多只能将通信量减低为1/32，要进一步降低通信量，就需要考虑稀疏化技术，即仅传输部分梯度。这一技术的理论基础在于，每次参数更新时仅有部分梯度影响较大，只使用这部分梯度进行参数更新即可获得比较好的收敛效果[5]。

这里以TopK-SGD[6]为例来介绍稀疏化过程，其伪代码描述如图5。每次迭代过程中，每个节点获取相应数据，计算梯度，然后对梯度的绝对值进行排序，找到排序位置在s%处的梯度，将其绝对值作为阈值，筛选出绝对值大于阈值的梯度，将这些梯度量化之后通过All-reduce进行汇总并用于模型参数更新。绝对值低于阈值的梯度则会在每个节点内进行累积。

图5 TopK-SGD算法流程[6]

表1为TopK-SGD在ImageNet上分别训练AlexNet和ResNet-50的结果，稀疏化技术配合量化技术在梯度通信过程中可以达到很高的压缩率，并且能保持模型的最终性能。

表1 在ImageNet数据集上使用TopK-SGD的训练结果[6]

TopK-SGD 带来的速度提升效果如图6，图中Deep Gradient Compression代表TopK-SGD的结果，可以看到通信占比较高(图中节点数更多的情况)以及通信带宽受限(图中1Gbps Ethernet的情况)时，TopK-SGD带来的速度提升效果更明显。

图6 TopK-SGD带来的速度提升情况[6]

TopK-SGD可以有效降低通信时间，但每个节点topk排序过程较为耗时，且需要使用AllGather对稀疏表达的梯度-索引进行各节点间的收集。对于这两个方面的优化可以进一步提升训练速度，具体实现可以参考MSTopK[7]。

调度通信过程和计算过程

神经网络模型通常是很多个基础模块堆叠的形式，比如ResNet由多个bottleneck block堆叠而成，每个block包含多个残差单元，每个残差单元又包含多个卷积层，再如BERT由多个transformer encoder堆叠而成，每个encoder又包含注意力层和全连接层。对于这种层级堆叠的结构，我们可以找到没有依赖关系的通信过程和计算过程，通过调度来使两者并行执行。典型的工作如Wait-Free Backward Propagation[8]，如图7所示，Scheduling的蓝色块代表梯度计算过程，橙色块代表通信过程，其中颜色较深的部分为通信启动时间，颜色较浅的部分为数据传输时间。假设此时已完成当前层(l)的梯度计算，则可以将当前层(l)的梯度通信过程与下一层(l-1)的梯度计算过程并行执行。

图7 Wait-Free Backward Propagation(修改自[1])

在使用WFBP时通常会遇到三种情况，如图8所示，理想情况为Case 1，通信时间与梯度计算时间相比较短，通信过程基本上可以和梯度计算过程完全并行。在Case 2的情况下，虽然有部分通信过程耗时较长，在单次调度中无法被很好地覆盖掉，但整体上没有产生额外的通信开销。在Case 3的情况下，由于有更多耗时较长的通信过程，整体上产生了额外的通信开销。考虑到每个通信过程都有启动时间，Case 3的情况是比较常见的，所以就有了MG-WFBP(Merged-Gradient WFBP)[9]，即将多个通信过程合并为一个，由此减少启动时间占比，从而减少整体的通信时间。

图8 WFBP与MG-WFBP[9]

结合上述思路，可以将前向计算过程以及稀疏化过程也作为调度对象参与优化，感兴趣的读者可以参考相关论文[10,11]。

总结与思考

本文从两个方面讨论了应对通信瓶颈的方法。一是对通信数据进行压缩，减少通信数据量，具体压缩方法又可分为量化和稀疏化，其中量化方法通过减少每个数据所占的比特位来实现压缩，稀疏化方法则通过选择部分数据进行传输来实现压缩，两种方法配合使用可进一步提升压缩效果。这两种方法都会带来数据传输过程中的信息损失，因此在设计算法的过程中需要进行大量的收敛性验证。二是通过调度来使通信过程和计算过程尽量并行，主要介绍了WFBP，利用模型逐层计算的特点，设计梯度计算过程和梯度通信过程的并行方案，并针对通信过程启动时间较长的情况设计了MG-WFBP，将多次通信过程合并，以进一步减少通信时间。

随着模型规模和各领域数据的不断增长，应对模型训练过程中通信问题的方案层出不穷，我们会持续跟踪这一方向的最新进展。

参考文献

[1] Austin Derrow-Pinion, Jennifer She, David Wong, et al. ETA Predictionwith Graph Neural Networks in Google Maps. 2021

[1] Tang, Zhenheng, et al. "Communication-efficient distributed deep learning: A comprehensive survey." arXiv preprint arXiv:2003.06307 (2020).

[2] D. Alistarh, et al. "Qsgd: Randomized quantization for communication-optimal stochastic gradient descent,” ArXiv, vol. abs/1610.02132, 2016.

[3] W. Wen, C. Xu, F. Yan, C. Wu, Y. Wang, Y. Chen, and H. Li, "Terngrad: Ternary gradients to reduce communication in distributed deep learning,” in Proceedings of the 31st International Conference on Neural Information Processing Systems, ser. NIPS’17. USA: Curran Associates Inc., 2017, pp. 1508–1518.

[4] K. Mishchenko, E. A. Gorbunov, M. Tak´ac, and P. Richt´arik, "Distributed learning with compressed gradient differences,” ArXiv, vol. abs/1901.09269, 2019.

[5] Y. Lin, S. Han, H. Mao, Y. Wang, and B. Dally, "Deep gradient compression: Reducing the communication bandwidth for distributed training,” in International Conference on Learning Representations, 2018.

[6] Lin, Yujun, et al. "Deep gradient compression: Reducing the communication bandwidth for distributed training." arXiv preprint arXiv:1712.01887 (2017).

[7] Shi, Shaohuai, et al. "Towards scalable distributed training of deep learning on public cloud clusters." Proceedings of Machine Learning and Systems 3 (2021): 401-412.

[8] H. Zhang, et al. "Poseidon: An efficient communication architecture for distributed deep learning on GPU clusters,” in 2017 USENIX Annual Technical Conference (USENIX ATC 17), 2017, pp. 181–193.

[9] S. Shi, et al. "MG-WFBP: Efficient data communication for distributed synchronous SGD algorithms,” in IEEE INFOCOM 2019-IEEE Conference on Computer Communications. IEEE, 2019, pp. 172–180.

[10] Y. Peng, et al.  "A generic communication scheduler for distributed DNN training acceleration,” in Proceedings of the 27th ACM Symposium on Operating Systems Principles. ACM, 2019, pp. 16–29.

[11] S. Shi et al., "Communication-Efficient Distributed Deep Learning with Merged Gradient Sparsification on GPUs," IEEE INFOCOM 2020 - IEEE Conference on Computer Communications, 2020, pp. 406-415, doi: 10.1109/INFOCOM41043.2020.9155269.

 往期推荐

关于壁仞科技研究院

壁仞科技研究院作为壁仞科技的前沿研究部门，旨在研究新型智能计算系统的关键技术，重点关注新型架构，先进编译技术和设计方法学，并将逐渐拓展研究方向，探索未来智能系统的各种可能。壁仞科技研究院秉持开放的原则，将积极投入各类产学研合作并参与开源社区的建设，为相关领域的技术进步做出自己的贡献。

扫码关注我们