一种面向分布式深度学习的轻量级聚合通信库*

王笑雨,董德尊

(国防科技大学计算机学院,湖南 长沙 410073)

1 引言

深度学习技术在许多领域应用广泛，特别是在计算机视觉和自然语言处理领域。随着处理的问题更加复杂，同时为进一步提高模型精度，深度神经网络DNN(Deep Neural Network)层数随之增多，数据集规模也逐渐变大，导致在单节点上进行相关模型的训练变得格外耗时。为缩短模型训练时间，分布式训练成为主流深度学习框架(如MXNet[1]、TensorFlow[2]和PyTorch[3])常用的解决办法。分布式训练利用多个节点的计算资源来加速大规模深度学习模型的训练，通常分为数据并行和模型并行2种方式。其中数据并行方式更为主流，其将数据集拆分为若干份并分配给各节点，各节点在拆分后的数据集上训练本地的全局模型副本，并通过在节点间交换训练过程中各自产生的梯度信息来更新全局模型的参数，从而完成在整个数据集上对模型的训练。这一方法在使得更多的计算资源可以被利用的同时也引入了通信开销。而随着节点的增多，节点间的梯度交换变得频繁，由此产生的通信开销占据了整个训练过程中的较大部分，减少了分布式并行化所带来的优势，成为加速分布式训练的瓶颈。

聚合通信操作是目前主流的用于梯度信息同步的方式之一，被主流的深度学习框架所采用。为进一步缩短分布式训练耗时，许多研究工作针对这一操作进行了优化。比如，通过调整聚合通信操作的执行顺序来利用计算开销隐藏部分通信开销，甚至通过对其进行拆分和合并来使得计算和通信尽可能重叠[4,5]。还有通过梯度量化[6,7]和稀疏化[8]等方法来减少需要交换的梯度信息的大小，从而减少通信开销。这些工作都将聚合通信操作视作黑匣子，并从如何有效使用的角度来进行优化，在一定程度上缩短了训练时间。

最近一些研究工作采用了新的改进思路——根据分布式训练的特征来调整聚合通信操作本身。eager-SGD(Stochastic Gradient Descent)[9]考虑到训练过程的鲁棒性提出了部分聚合通信操作,与原先操作不同的是，其在语义逻辑上并不要求所有节点都参与梯度同步过程，减少了等待掉队者的开销。此外，考虑到分布式训练中后产生的梯度信息会先被用到，文献[10]提出了抢占式AllReduce，使得后产生的具有更高优先级的梯度信息在逻辑上可以中断低优先级的梯度同步过程，以更快地进行同步。聚合通信操作在高性能计算领域已经使用了快30年，而其近期才被引入深度学习领域，还未完全针对这一应用场景进行优化。上述这些工作考虑到了传统聚合通信操作的语义和深度学习应用的特点不匹配的现象，并试图缓解这一矛盾。此外，还有一些将分布式训练特点与网络结合进行优化的工作改变了聚合通信操作的实现[11 - 13]。这些工作都说明研究专用于分布式训练的聚合通信操作，并进行应用和传输的协同设计对减少通信开销来说很有必要。

对于现有去中心化的分布式训练框架来说，进行相关的协同设计较为困难。其通常采用MPI(Message Passing Interface)标准规定的聚合通信接口来完成实际的梯度同步操作。MPI具有不同的实现版本，比如OpenMPI和MPICH。这些版本为了保证鲁棒性、普适性和扩展性，同时为了便于开发，通常会具有复杂的架构和接口调用，多种不同功能的通信函数以及庞大的代码量，这使得尽管其在工业生产中稳定健壮，但对在其上进行研究和改进并不友好。

为解决上述问题，本文设计并实现了一个针对分布式训练中聚合通信操作的轻量级通信库，使得相关的实验更易实现。该轻量级通信库整体架构简单但提供了分布式训练中必要的聚合通信操作，且向上支持主流的深度学习框架，向下支持集群中常见的网络环境。该库使用较少的封装来将聚合通信操作的执行细节暴露出来，以便于进一步分析和改动。此外，本文希望为研究分布式训练中的聚合通信操作提供一个统一的实验环境。通过该通信库，相关研究人员可以简单高效地实现自己的想法并在更多的环境配置下进行实验，进而获得更广泛的影响。而在此库基础上开发的算法也可以更轻松地应用到多个深度学习框架中。最后本文从聚合通信操作和深度学习应用2方面对该库进行了评估，并将其与常用的聚合通信库进行了对比。最终实验结果显示，该库在部分情况下具有优势。

2 背景介绍

2.1 分布式训练和框架

分布式训练主要用于解决在单节点上训练大规模DNN耗时的问题，其通过将训练任务进行拆分并部署到多个节点上同时进行训练来加速模型的训练过程。在现有的解决方案中，数据并行是较为流行的训练方式之一，被主流深度学习框架广泛使用。在数据并行的方式中每一个worker(对应一个进程，可能在CPU上也可能在GPU一类的加速器上)都拥有一份全局模型的副本和一部分数据集，其在拥有的数据集上独立训练本地全局模型副本，并周期性地与其他的worker同步各自的梯度信息，以更新全局模型的参数，从而实现在多个节点上基于整个数据集来并行地训练模型。在更新全局模型参数的过程中，所有worker交换各自训练过程中在反向传播阶段产生的梯度，并将这些梯度进行归约(一般是求和或求平均值)，以使用随机梯度下降SGD的变种算法来更新全局和本地的模型参数，更新后的参数会在下一轮训练的前向传播阶段使用。值得注意的是，梯度产生和对应模型参数使用的顺序通常是相反的，也就是说,在反向传播阶段越晚产生的梯度，基于其更新的模型参数会在前向传播中越早被用到。

前述主流的深度学习框架都为用户使用分布式训练提供了便捷的方式，通常只需要在原有单节点训练的代码上进行简单的修改即可。但是，这些框架具有不同的编程语义和性能，比如，TensorFlow采用的是声明式的语义，用户只需要声明需要执行的操作，这些操作由TensorFlow进行定义；PyTorch则采用了命令式的语义，需要用户规定操作具体怎么执行；MXNet则结合了这2种语义。对于梯度同步过程，这些框架也采用了不同的通信架构，PyTorch采用的是AllReduce架构，而TensorFlow和MXNet虽然采用的都是参数服务器PS(Parameter Server)架构，但实现方式不同，分别是Rendezvous和KVStore,如图1所示。这些差异使得不同的深度学习框架具有各自不同的优势，比如，MXNet对初学者更加友好；PyTorch在使用上更加灵活；而TensorFlow则对在CPU上训练大规模深度神经网络具有更好的性能。这些具有不同特性的框架在学术界和工业界都被广泛应用，对于分布式训练的研究也非常重要。

Figure 1 Architecture of the proposed design图1 本文设计整体架构

2.2 聚合通信及相关工作

聚合通信操作是高性能计算领域的经典技术，但若将其应用于分布式训练还有改进的空间。在梯度同步过程中，聚合通信操作AllReduce最常被使用。它将每个节点上的梯度数据通过指定的操作(如求和或求平均值)进行归约，并将归约结果同步到各个节点。从AllReduce的语义可以看出，其要求所有的节点都参与该过程。在现有的算法中，Ring AllReduce最常被用于分布式训练，其可以充分利用每个节点的网络带宽资源，因而也被用于缓解PS架构中server节点带宽瓶颈的问题。在此基础上，一些工作尝试从结合上层深度神经网络训练应用特点的角度合理且高效地使用该操作来减少通信开销[4,5,10]；还有一些工作则试图从底层聚合通信算法和网络传输的角度来优化该操作[11]，但很少有工作是结合两者进行优化。前面提到的部分聚合操作工作就有类似的优化思路，其尝试将上层的深度学习应用的鲁棒性特点与底层聚合通信操作的语义相结合来进行改进，放宽聚合通信操作要求所有节点全部参加的条件，从而避免等待拖尾节点，减少了梯度同步的时间，并且不影响模型的收敛。可见，减少深度学习应用的需求和聚合通信操作的语义间的不匹配问题对加速分布式训练至关重要。

2.3 网络和通信库

不同的集群通常有不同的网络环境，而一个集群内通常也有多套网络。集群的网络环境会从多方面影响分布式训练中的通信性能，包括带宽、协议及拓扑结构等。相较于以太网，InfiniBand网络具有高带宽和低时延的特点，这得益于RDMA(Remote Direct Memory Access)可以直接绕过内核，相较TCP开销更小，可以从网络传输速率上加速梯度同步[14,15]。此外，Bcube比Fat-Tree的拓扑结构更适合AllReduce架构[16]。同时使用集群中的多套网络设备显然也能加速通信。鉴于集群环境的多样性，充分利用集群的网络资源和特性对加速通信十分重要。然而不同的网络设备通常使用各自的接口库来进行驱动，这些库的使用方式和调用机制不同，给实现上述改进增加了困难。比如，InfiniBand使用的ibverbs库在调用实际的通信操作前需要对许多对象进行初始化和配置，如context和QP(Queue Pair)，并且需要通信双方交换诸如各自地址的控制信息，相比之下，以太网常用的socket库则仅需创建socket对象并调用个别函数即可建立连接，而对于研究分布式训练聚合通信操作则不需要考虑这些细节。

主流深度学习框架并不是直接使用不同网络的底层库进行点对点通信，而是调用诸如MPI和NCCL等聚合通信库，来完成梯度同步操作的，由这些库负责实现并选择聚合通信算法并对底层网络的接口进行封装。其中，MPI实际上是一个编程接口标准，被广泛应用于大规模并行编程中，有多种实现版本，比如OpenMPI和MPICH等。NCCL则是由NVIDIA公司开发的用于其GPU间交换数据的聚合通信库。这些库提供了便捷的使用方式，对底层网络接口进行封装，但对分析和优化这些聚合操作的细节来说较为麻烦。

2.4 动机

如前文所述，通过综合考虑深度学习应用特征和底层网络特点来进行协同设计，可以减少其中存在的不匹配现象。然而,在现有的聚合通信库基础上分析和优化分布式深度学习中的通信操作较为复杂。期望有一个架构简单、调用逻辑清晰且代码量较少的聚合通信库。此外，相关的工作大部分都是针对于某种特定环境，比如某个框架或是某种网络，想要移植到其他环境较为困难。为了满足实验的不同需要以及使相关研究便于拓展移植，本文希望聚合通信库能对上层的深度学习框架和下层的集群网络有较为广泛的支持。

3 轻量级聚合通信库

本文的目标是提供一个可以用于简化聚合通信操作算法相关研究和实现的轻量级通信库。为实现这一目标，该库需要支持主流的框架及网络，以满足在其上进行相关研究的需求。本节将阐述该通信库的整体架构和关键细节。

3.1 整体架构

首先，不同的分布式训练框架通常通过其各自的框架和通信库来进行梯度同步。比如，如图1所示，MXNet使用基于ZMQ(ZeroMQ)库的PS架构，而TensorFlow则采用了gRPC库。为了实现对主流深度学习框架支持这一目标，本文选择Horovod[17]为通信中间件之一。Horovod对分布式训练中的参数同步逻辑进行了抽象，并据此针对主流框架分别定义了各自用于通信的接口以支持这些框架，并在底层统一实现。在通信的具体实现上，Horovod选择了AllReduce作为其参数同步的框架，如图1所示，底层使用MPI、NCCL和Gloo[18]作为执行实际聚合通信操作的通信库。Gloo是Horovod自带的一个轻量级的通信库，为分布式深度学习应用提供了一系列基本的聚合通信算法,使Horovod可以在缺少MPI的环境中运行。对于Horovod来说，MPI中聚合通信的内部执行细节是透明的，仅需要按照MPI标准进行调用，具体与MPI的实现版本有关，而Gloo则是可以被Horovod感知的。此外，相较各种MPI实现版本，Gloo的架构简单，代码量少，更易于进一步分析和改进。而 NCCL则是专用于GPU间通信的，不支持在CPU上进行训练。因此本文选择了Gloo作为通信库的基础，从而借助Horovod来确保对上层框架的支持。

其次，集群中通常具有多种网络设施，并且不同集群的网络环境也不同。通过支持多种网络来使得该库具有更广泛的适用性，具备充分利用网络资源的能力。本文在对Gloo进行分析后发现，其对网络环境的支持不如MPI。比如，在InfiniBand网络模式下，Horovod无法通过使用Gloo完成AllReduce操作。因此需要对Gloo的底层网络进行扩展。Gloo的扩展方式是使用各自网络的接口库来分别实现符合Gloo中点对点通信语义的传输层子类，从而为其算法层提供通信传输的接口。本文也可以按照这一思路来对其他网络进行扩展，但这对不同网络的了解和掌握程度提出了要求，工作量较大，且后续的扩展也不能得到保障。此外，使用网络接口库实现点对点通信的具体实现细节对研究分布式训练中的聚合通信来说也不是必要的。本文最终使用UCX(Unified Communication X)[19]作为通信中间件来解决Gloo的网络扩展问题。UCX整合并封装了多种底层网络的接口，提供了一套统一的高性能的点对点通信接口，在屏蔽不同网络接口库实现细节的同时，还通过提供多种语义的操作来体现不同网络的特性。

本文的设计思路如图1所示，将Gloo的底层网络用UCX进行扩展，从而实现支持多种框架和网络的目标。在这一架构中，Horovod整合并接管了不同深度学习框架的通信部分，UCX用于覆盖尽可能多的网络支持，Gloo作为聚合通信的算法逻辑层，为上层应用和下层网络提供交互条件，充分整合各层信息。与图1中使用MPI的框架相比，这种设计可以将聚合通信操作的过程加以分解，从而暴露算法的更多细节。这一改变使得对分布式训练中诸如AllReduce一类的通信操作的研究和改进变得便捷高效。此外，由于该库仅提供了分布式训练必要的聚合通信操作，降低了该库的复杂性。

3.2 Gloo

如图1所示，Gloo是Horovod和UCX间的桥梁。本文将Gloo的架构拆分为算法层和传输层,其架构如图2所示。

Figure 2 Architecture of Gloo图2 Gloo架构

上层算法层主要包括聚合操作的算法逻辑，其基于点对点语义的通信逻辑实现。此外，该层提供了一些分布式训练必要操作的算法实现，比如AllReduce和Broadcast,用户可以参考这些算法实现来实现自定义算法。底层的传输层主要提供点对点通信的调用接口以及不同网络环境下对应的具体实现。Gloo本身仅支持tcp、uv及ibverbs 3种库，其中对于ibverbs的支持还不完善。此外，对其点对点通信进行扩展需要用网络对应的接口重写其基本类和函数，但这对用户了解网络设备硬件及熟练使用对应接口提出了更高的要求，并且对于进行聚合操作相关的研究来说引入了更多不必要的细节。

如图2所示，Buffer是Gloo中的关键类，其既存储了需要传输的数据又是调用通信函数的对象，是算法层和传输层的连接。在算法层中，通过Buffer对象来调用对应的点对点通信函数，进而与对端进行数据的发送和接收，以完成算法逻辑。在传输层中，Buffer则通过Pair对象调用实际的通信函数。

为了在实现UCX的扩展时更符合Gloo的语义，本文以Horovod调用AllReduce操作为例，分析了调用流程和类的功能。在实际调用AllReduce之前，Horovod首先通过调用CreateDevice、CreateContext和CreatePair来初始化Gloo。其中，CreateDevice配置网络设备，包括接口和协议等；CreateContext设置通信上下文，包括节点数和节点编号等；CreatePair创建Pair对象，该对象代表与另一个对等节点的连接，并配置其通信语义，包括同步或异步等。需要注意的是，一个Context对象可用于创建多个Pair对象，同样一个Device对象可以创建多个Context对象。然后，Horovod通过调用connect函数来建立全局连接，这样所有节点都可以相互通信。初始化工作完成后，Horovod在需要进行通信时可以调用AllReduce操作。这一过程中将创建Buffer对象，用于存储通信数据并调用send和recv函数。值得注意的是，send和recv是节点之间通信实际发生的位置，其语义是非阻塞的，这意味着函数在实际通信完成之前就已返回，因此需要调用waitSend和waitRecv以确保数据已传输完成。

3.3 UCX

UCX是一组用于高吞吐量计算的网络API，为下一代应用程序和系统实现了高性能和可扩展的网络堆栈。UCX框架包含3个主要组件：UCS、UCT和UCP，其中每个组件都提供一套API，并且可以作为独立的库来使用。UCS主要是为UCX内部提供支持;UCT是一个抽象了各种硬件体系结构之间差异的传输层，并提供了一个可实现通信协议的底层API;UCP通过使用UCT层的功能来实现较高级的协议传输功能，包括标签匹配、流传输等多种通信语义。

UCP相较于UCT的优势在于其使用方式简单且功能多样。前文提到的Horovod接管了主流深度学习框架的通信部分，并调用Gloo的聚合通信操作来完成，而UCX封装了各种网络API，本文将前者的底层传输替换为后者，以在支持多框架的基础下也能对网络提供更好的支持。

3.4 实现细节

本文通过将Gloo的传输层替换为UCX来实现通信库的设计。根据对Gloo的分析和UCP的使用方法，本文按照如下方式对其进行整合。首先，分别将ucp_context、ucp_worker和ucp_ep对象分别映射到Device、Context和Pair，这样做是因为它们之间具有一致的对应关系。然后，在Gloo的相应初始化函数CreateDevice和CreateContext中调用ucp_init和ucp_worker_create，以初始化UCX，但在connect函数中才调用ucp_ep_create，因为其在UCP中的语义代表连接的建立。再次，根据UCP的语义及其使用方法，本文选择ucp_worker的地址而不是节点的IP地址作为连接标识，它包括该节点所有网络设备的相关信息。最后，在send和recv中分别调用ucp_tag_send_nb和ucp_tag_recv_nb与对等方进行通信。其中，tag的设置综合了集合操作的标识符和发送者的rank。此外，本文创建了请求队列，以维护这些操作的句柄，因为这些函数是非阻塞的，需要通过句柄来确认其传输是否完成，以确保waitSend和waitRecv的Buffer中数据的正确性。

Figure 3 Bandwidth of AllReduce in three collective operation library with different data size图3 AllReduce操作在不同数据大小下的带宽

4 实验评估

本节将对本文的轻量级聚合通信库中分布式训练时最常使用的操作AllReduce进行评估 ，并与MPI的实现及原始的Gloo进行比较。此外，本文还分别基于该库和MPI的实现在MNIST和CIFAR-10数据集上对模型进行了训练，并测量和对比了其训练速度。

4.1 实验环境

实验主要在由并行超算云提供的集群上进行，集群具有4个节点，每个节点有1个24核的E5-2678的CPU及64 GB的主存，集群网络环境为56 Gb/s的InfiniBand和1 000 Mb/s的以太网。此外，集群上安装有OpenMPI 4.0.1、Horovod 0.19.1、Python 3.7.0和MXNet 1.6.0。

4.2 聚合通信操作

本文在InfiniBand网络的以太网(ib0)和IB网(mlx5_0)模式下分别测试了所设计的轻量级聚合通信库、OpenMPI和原始Gloo的AllReduce操作的带宽，这里带宽指的是AllReduce处理的数据量和所花时间的比值。从图3所示的实验结果可以发现，在以太网模式下，该库和OpenMPI及原始的Gloo的性能相近，而在高速网模式下，该库与OpenMPI有较为明显的性能差异，在数据小于2 MB时OpenMPI性能更好，而大于2 MB时则该库更有优势，产生这样现象的原因可能与底层网络接口库的使用方式不同以及MPI对AllReduce操作的算法选择有关，还需要后续进一步实验分析。需要注意的是，原始Gloo仅支持使用以太网模式。

4.3 深度学习应用

本文还在上述这些库的基础上在MXNet框架中进行了分布式训练的实际应用，分别在MNIST和CIFAR-10数据集上训练了LeNet5和ResNet18 2个模型，测量其在不同批大小下的训练速度。图4和图5分别为MNIST和CIFAR-10的实验结果。从图4和图5可以发现,这些库训练速度较为接近，而本文提出的库相较OpenMPI在大部分情况下具有较小优势。结合图3中AllReduce性能的对比，较为可能的原因是每次同步的梯度较小，不足以使AllReduce性能呈现出明显差异。

Figure 4 Speed of training MNIST dataset with different batch sizes图4 MNIST数据集在不同批大小下的训练速度

Figure 5 Speed of training CIFAR-10 dataset with different batch sizes图5 CIFAR-10数据集在不同批大小下的训练速度

5 结束语

本文提出并实现了一个用于分布式训练的轻量级聚合通信库。该库具有较为简洁的调用架构和较少的代码量，并且提供了分布式训练中基本的聚合操作，对主流的深度学习框架和集群网络环境有较好的支持。相较于分布式训练中传统的基于MPI等的通信架构，该库的调用逻辑简单，能更清晰地展示聚合操作过程的细节，方便进一步分析和改进分布式训练中的聚合通信操作实现。并且得益于其对框架和网络的支持，用户可以在更多的环境配置下进行实验以获得更广的影响。此外，本文分别从聚合通信操作和深度学习应用方面对该库的性能进行了评估，其在实际的深度学习应用方面与常用的OpenMPI性能相近，可以作为分析和研究分布式深度学习中梯度同步的聚合通信库。