深度学习（DL）已经被广泛用于许多人工智能（AI）应用。在深度学习中，训练一个令人满意的模型需要大量的计算，因此人们部署GPU集群并应用分布式训练技术来加速训练过程。
数据并行是一种广泛使用的分布式训练方法。在这种方法中，训练数据集被拆分并分配给所有机器，因此每台机器都可以使用分配的数据片来训练模型（“计算阶段”）。 在他们完成本地训练后，机器交换梯度（“通信阶段”）以保证所有机器拥有一致的深度神经网络（DNN）模型。整个训练任务一直重复这个过程，直到训练出一个令人满意的全局模型。然而，迭代中的通信阶段比较耗时，影响训练任务的完成时间。
图1: 三种隐藏通信开销的流水线设计。（a）WFBP指无等待反向传播方式，其中模型每一层的通信和反向计算并行；（b）WFBP+Compression 指无等待反向传播和压缩的混合方式，其中模型每一层的通信和计算并行，反向计算和压缩计算是串行的；（c）PipeCompress 是我们的方法，其中每一层通信，反向计算和压缩计算互相并行
当前主要有两种技术来应对通信挑战：基于压缩的技术（也称为压缩方法）和基于流水线的技术（也称为流水线方法）。压缩方法在发送之前将梯度压缩为低比特值，或者仅交换一小部分梯度而不是全部梯度，因此减少了交换所需的数据量大小。
流水线方法试图通过将通信和反向传播计算流水线并行来隐藏通信时间。 具体来说，每次迭代中的计算阶段包括两个步骤：前向计算和反向传播计算，并且这两个步骤都是逐层的。在流水线方法中，机器在一层完成反向传播计算后发送该层的梯度，而不等待后续层的反向传播，如图1（a）所示。 因此，这种方法也称为“无等待反向传播”（WFBP）。通过这种方式，可以通过反向传播计算隐藏部分通信时间。
此外，人们还提出了一种混合方法，将无等待反向传播（WFBP）和梯度压缩结合使用，以进一步降低通信时间成本，如图1（b）所示。
应用压缩技术后，通信时间明显减少，但压缩梯度非常耗时，成为训练效率的新瓶颈。
当前的无等待反向传播（WFBP）和压缩混合方法旨在隐藏通信时间，但压缩产生的时间成本仍然存在，如图1（b）所示。如果压缩的时间成本太高，就会抵消减少通信流量带来的收益。我们测了一些DNN模型在训练过程中压缩梯度所花费的时间。我们在1Gbps以太网连接的两台服务器上使用混合方法训练这些模型，每台机器都有两个GPU（NVIDIA Titan XP）。我们观察了两种典型的梯度压缩方法的性能：量化和稀疏化。我们将训练时间分解，观察每次迭代中压缩时间所占的比例。 表1给出了结果，表明梯度压缩时间占整个训练过程的21%-65%。因此，压缩时间不可忽略，有必要减少或隐藏它以加速模型训练。
表1: 一轮迭代训练中的压缩时间占比
当我们查看图1（b）时，很自然地会考虑增加并行度以缩短计算的时间线。在本文中，我们考虑那些具有分层结构的DNN模型，在这些模型中，压缩和反向传播相互独立，这意味着更浅层的反向传播计算可以在一个层完成反向传播计算后立即进行，而不用等待该层的压缩结果。图1（c）表明了我们的想法，从中我们可以看到并行度增加，迭代时间可以进一步减少。然而，实现上仍存在一些挑战：压缩和反向传播计算都需要占用GPU资源，这意味着我们应该让GPU同时运行这两个操作，而不是分时间片进行。否则，增加并行度是没有意义的。 在本文中，我们提出了PipeCompress，这是一种两级流水线，用于加速DNN模型的分布式训练过程。我们利用GPU共享功能来进行流水线压缩和反向传播操作，并将它们打包在单个GPU中，以真正的并行方式运行它们。这样既可以隐藏压缩时间，又可以提高GPU利用率。
为了验证二级流水线的可行性以及最大化反向传播、压缩和通信之间的重叠并提高并行性的机会，我们首先分析反向传播、压缩和通信之间的依赖关系，然后展示了如何利用GPU共享功能来实现这一目标并展示了PipeCompress的架构。
在分布式训练过程的每次迭代中，前向、反向传播、压缩和通信操作之间可能存在内在依赖关系。例如，第i层的压缩操作的开始依赖于同层反向传播操作的完成。另一方面，某些操作之间没有内在依赖性，因此它们永远不会相互阻塞，并且这些操作可以流水并行。 例如，通信操作的开始永远不会被更浅层的反向传播操作阻塞。无等待反向传播（WFBP）利用这一特性，使通信操作和反向传播操作重叠以减少迭代时间。
我们注意到，在当前的WFBP和压缩混合方法中，第i层的压缩操作往往会阻塞更浅层的反向传播操作。 我们注意到这些操作之间没有内在依赖性，更浅层的反向传播操作不需要第i层的压缩结果输出。 换句话说，这两个操作可以由两个进程同时进行而不会相互阻塞，如图1（c) 所示。 我们利用这些操作之间的独立性将它们解耦为两个进程。 这样，这些操作可以在时间轴上重叠，提高并行度，从而可以减少训练时间。
在二级流水线中，压缩和反向传播可以分两个进程运行，从而提高并行度，减少迭代时间。 然而，压缩和反向传播都需要计算资源，这意味着如果两个进程都在同一个GPU上运行，则两个进程之间可能存在资源竞争。如果不能妥善解决竞争，将它们解耦为两个过程是没有意义的。 这与无等待反向传播（WFBP）不同，无等待反向传播（WFBP）解耦后的两个进程需要不同类型的资源，即计算过程需要计算资源（GPU），而通信过程不消耗计算资源。
图2: PipeCompress中关闭和开启GPU共享功能比较
如图2所示，当两个进程以传统方式在同一个GPU上运行时，进程是以独占整个GPU的方式使用GPU资源的。因此，这两个进程会竞争GPU的时间片，即分时间片利用GPU。 也就是说，每个时间片内只能运行一个进程。 这不是真正的并行系统，并没有真正提高并行度。PipeCompress需要一个精心设计的实现来让反向计算和压缩计算并行。事实上，一个GPU有数百个计算单元（Stream Multiprocessor）。我们可以让这些计算单元在同一个时间片内被2个任务共享，即反向计算任务和压缩计算任务都占用一部分计算单元进行计算，这样我们就能解决反向计算和压缩计算在资源调度上的竞争问题。具体实现上，PipeCompress利用NVIDIA GPU提供的MPS（Multiple Process Sharing）接口，将反向计算进程和压缩计算进程包裹在一个CUDA上下文中，这样在同一个时间片内2个进程就能同时共享这些计算单元，这种共享模式也被称为packing。
我们基于深度学习框架PyTorch实现了PipeCompress。如图3所示，PipeCompress架构中的每个worker中都有一个计算引擎、压缩管理器和通信管理器。计算引擎负责执行前向计算和反向传播以生成梯度，压缩管理器负责从计算中接收梯度并压缩梯度，通信管理器负责从压缩管理器接收合并的梯度并发送梯度。在训练开始时，PipeCompress启动MPS将计算引擎和压缩管理器打包到同一个GPU中，以便反向传播操作和压缩操作可以共享GPU。为了减轻计算引擎、压缩管理器和通信管理器之间消息传递的开销，我们采用了一种具有共享内存的避免复制消息机制。为了避免共享内存上的读写冲突，这些组件仅在从其发送方接收到来自控制流的信号时才访问共享内存。
图3: PipeCompress架构
吞吐量：
这4种DNN模型训练任务都用于图像分类，因此，我们定义吞吐量为每个 GPU每秒可以处理的图像数量。实验结果如图4所示。我们可以看到PipeCompress的吞吐量总是能比无等待反向传播（WFBP）+压缩混合方法的吞吐量更高
图4: 吞吐量(上：量化，下：稀疏化)
时间分解：
在上文的讨论中，我们观测了两种方法的吞吐量指标。吞吐量指标在一定程度上反映了一轮迭代的端到端时间。为了更好地理解PipeCompress对性能提升的收益来源，我们将一轮迭代时间分解为前向计算、反向计算、压缩和通信开销这4个部分。 实验结果如图5所示，我们可以看到无论是否使用量化压缩或稀疏化压缩方法， PipeCompress总是能减少所有四个模型训练任务的迭代时间。此外，我们还观察到以下现象：
图5: 时间分解(上：量化，下：稀疏化)
参考文献
[1] R. Mayer and H.-A. Jacobsen, “Scalable deep learning on distributed infrastructures: Challenges, techniques, and tools,” ACM Computing Surveys (CSUR), vol. 53, no. 1, pp. 1–37, 2020.
[2] S. Shi, Q. Wang, X. Chu, B. Li, Y. Qin, R. Liu, and X. Zhao, “Communication-efficient distributed deep learning with merged gradient sparsification on gpus,” in IEEE INFOCOM 2020-IEEE Conference on Computer Communications. IEEE, 2020, pp. 406–415.
[3] S. Shi, Z. Tang, Q. Wang, K. Zhao, and X. Chu, “Layer-wise adaptive gradient sparsification for distributed deep learning with convergence guarantees,” arXiv preprint arXiv:1911.08727, 2019.
[4] D. Alistarh, D. Grubic, J. Li, R. Tomioka, and M. Vojnovic, “QSGD: Communication-efficient sgd via gradient quantization and encoding,” in Advances in Neural Information Processing Systems, 2017.
[5] 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 Advances in Neural Information Processing Systems, 2017.
[6] F. Seide, H. Fu, J. Droppo, G. Li, and D. Yu, “1-bit stochastic gradient descent and its application to data-parallel distributed training of speech dnns,” in Fifteenth Annual Conference of the International Speech Communication Association, 2014.
[7] S. U. Stich, J.-B. Cordonnier, and M. Jaggi, “Sparsified SGD with memory,” in Advances in Neural Information Processing Systems, 2018.
[8] D. Alistarh, T. Hoefler, M. Johansson, S. Khirirat, N. Konstantinov, and C. Renggli, “The convergence of sparsified gradient methods,” in Advances in Neural Information Processing Systems, 2018.
[9] Y. Lin, S. Han, H. Mao, Y. Wang, and W. J. Dally, “Deep gradient compression: Reducing the communication bandwidth for distributed training,” in International Conference on Learning Representations, 2018.
[10] H. Zhang, Z. Zheng, S. Xu, W. Dai, Q. Ho, X. Liang, Z. Hu, J. Wei, P. Xie, and E. P. Xing, “Poseidon: An efficient communication architecture for distributed deep learning on GPU clusters,” in USENIX Annual Technical Conference, 2017, pp. 181–193.
[11] A. Jayarajan, J. Wei, G. Gibson, A. Fedorova, and G. Pekhimenko, “Priority-based parameter propagation for distributed dnn training,” in 2nd Conference on Systems and Machine Learning (SysML), SysML, 2019.
[12] Y. Peng, Y. Zhu, Y. Chen, Y. Bao, B. Yi, C. Lan, C. Wu, and C. Guo, “A generic communication scheduler for distributed dnn training acceleration,” in Proceedings of the 27th ACM Symposium on Operating Systems Principles, 2019, pp. 16–29.
[13] S. Shi, X. Chu, and B. Li, “MG-WFBP: Efficient data communication for distributed synchronous sgd algorithms,” in IEEE INFOCOM 2019- IEEE Conference on Computer Communications. IEEE, 2019, pp. 172– 180.
[14] Y. Bao, Y. Peng, Y. Chen, and C. Wu, “Preemptive all-reduce scheduling for expediting distributed dnn training,” in IEEE INFOCOM 2020-IEEE Conference on Computer Communications. IEEE, 2020, pp. 626–635.
[15] NVIDIA, “Cuda multi-process service,” 2020,https://docs.nvidia.com/deploy/pdf/CUDA Multi Process Service Overview.pdf.
[16] M. Jeon, S. Venkataraman, A. Phanishayee, J. Qian, W. Xiao, and F. Yang, “Analysis of large-scale multi-tenant GPU clusters for DNN training workloads,” in 2019 USENIX Annual Technical Conference (USENIX ATC 19), 2019, pp. 947–960.
[17] K. Simonyan and A. Zisserman, “Very deep convolutional networks for large-scale image recognition,” in International Conference on Learning Representations, 2015.
[18] K. He, X. Zhang, S. Ren, and J. Sun, “Deep residual learning for image recognition,” in Proceedings of the IEEE conference on computer vision and pattern recognition, 2016, pp. 770–778.
[19]C. Szegedy, W. Liu, Y. Jia, P. Sermanet, S. Reed, D. Anguelov, D. Erhan, V. Vanhoucke, and A. Rabinovich, “Going deeper with convolutions,” in Proceedings of the IEEE conference on computer vision and pattern recognition, 2015, pp. 1–9.
原文始发于微信公众号（风眼实验室）：PipeCompress：一种加速分布式训练通信的二阶流水并行架构
