OneFlow深度学习框架的个人空间

分布式、高并发、多线程，似乎是程序员永远逃离不了的3个关键词，只要脱离了单机/单节点，涉及到2个以上的设备，就会碰到分布式。深度学习领域也一样，当你拥有海量数据/巨大模型的训练需求时，即使一张是TESLA V100也可能满足不了性能需求。这时就会扩展为单机多卡、多机多卡的分布式训练....

不管是理论研究还是工程一线，已经出现了很多（只）适合分布式训练的深度学习模型。比如：CV领域，为了将训练ImageNet的时间压缩至最短，腾讯团队曾在2018年，使用了2048块Tesla P40，将ResNet50在ImageNet上的训练时间压缩至6.6分钟，详见论文《Highly Scalable Deep Learning Training System withMixed-Precision: Training ImageNet in Four Minutes》，知乎也有相关报道：4分钟训练ImageNet！腾讯机智创造AI训练世界纪录

NLP领域，更是大模型频出，如BERT、GTP系列。为了训练GPT-2模型，用了256个Google Cloud TPU v3，据说GPT-3的训练更是耗费了N多显卡和1200万美金，知乎上也有相关文章：如何评价1700亿参数的GPT-3？

天下武功，唯快不破，要想快，就必须走分布式训练的路子。现在，各大深度学习框架基本都支持了深度学习模型的分布式训练，那么问题来了：

1. 深度学习的分布式训练的有哪些基本原理？

2. 深度学习的分布式训练有哪些现成的技术和框架？

3. 各大框架的分布式训练上手难易程度如何，训练孰优孰劣？

本文将对以上问题进行粗浅的回答和总结，权当抛砖引玉，欢迎大家关注和交流！本篇文章重点梳理深度学习分布式训练领域常用的一些技术及概念。如有疏漏和不足之处，还请多指点。

1.数据并行or模型并行

单机单卡情况下，信息都在一台机器上，无所谓分发。而分布式训练中，信息是要被“分发”的，分发的不同方式，常被称为“并行方式”。通常，习惯上将分发方分为“数据并行”(Data Parallelism)和“模型并行”(Model Parallelism)两种。

在数据并行中，将样本数据进行切分，切分后的数据被送至各个训练节点，与完整的模型进行运算，最后将多个节点的信息进行合并，如下图所示：

在模型并行中，将模型进行切分，完整的数据被送至各个训练节点，与切分后的模型进行运算，最后将多个节点的运算结果合并，如下图所示：

1.1数据并行

什么是数据并行？

以GPU的维度来看，数据并行简单来说就是在并行训练的设备上，对完整训练数据进行分片训练，同一个训练的时间间隔内，不同GPU设备上用各自分片的数据对模型进行训练，其后再进行模型梯度的汇总更新和各GPU间的状态同步。这样做的结果就是在一个训练的时间间隔内，各个GPU设备可以并行地用各自分片的数据进行模型训练，从而大大加速了整体模型的训练。

从上图可见，数据并行时，每个GPU设备上保持了同样的模型数据，且一次完整的训练过程包括以下3步：

1. CPU负责将不同的训练数据（mini-batch）分别喂给GPU0和GPU1设备；

2. 不同的显卡设备上存储了完全一致的模型，通过mini-batch数据进行了前向和反向传播；

3. 位于不同GPU设备上的模型进行权重同步和更新

为什么需要数据并行？

简单来说，就是加速训练过程，当我们的数据规模越来越大，训练一个完整模型所需的时间越来越长，为了加速训练，我们希望通过分布式训练的方式，通过拓展设备数量来压缩训练时间，达到近乎线性的加速比。数据并行是最常见也最早出现的一种分布式方式，各分布式框架均有容易上手的数据并行接口（但不一定支持下文介绍的模型并行）。

那……

举个栗子吧！

已知：

假设训练集为128万张的ImageNet；模型为ResNet50；单GPU的显存能支持的最大batch size为128；迭代1个batch（完成数据加载+前向+反向梯度更新）需要的时间为7.2秒。

求：

单个GPU训练100epoch所需时间？

答案：

完成一个epoch需要的时间为1280000/128 * 7.2  = 72000(s)，即20小时，迭代100epoch，则在单张GPU上所需的时间是20*100=2000小时，如果用单GPU的方式训练模型，需要耗时间近3个月（头发都掉光了～）

而如果是线性加速比，如果用4台8卡的机器训练，则时间更是缩短至原来的1/32，即2000/32=62.5小时（还是可以接受的）。

线性加速比是利用分布式做横向扩展的理想状态，实际上，如果框架设计或实现上不合理，会使得加速比大大低于理论值，造成“大幅度增加了机器但是并没有大幅度提高训练速度”的尴尬局面。

1.2 模型并行

什么是模型并行？

模型并行和数据并行类似，将整个模型的不同网络层(或者某一层)的参数矩阵切分至不同的GPU设备上，进行模型训练的过程。

从上图可知，模型并行时，完整的模型网络切分到了不同设备：GPU0和GPU1，且训练过程分为如下几步：

1. mini-batch喂送给GPU0；

2. 数据在GPU0所在网络上进行前向过程；

3. 上一步的数据继续喂给GPU1并在GPU1的网络上继续进行前向；

4. GPU1进行反向；

5. 反向数据回传给GPU0，继续进行反向

可见，模型并行时，并不需要进行各个设备上模型权重参数的同步更新，而是会有中间数据在各个GPU上的模型之间流动。

为什么需要模型并行？

某些情况下，模型规模特别巨大，参数特别多以至于单个GPU的显存塞不下（譬如某些分类网络/人脸模型由于num_classes特别大，导致最后FC全连接层的参数量巨大），于是只能通过模型并行的方式进行训练，即将模型的各网络层甚至是某一层的参数矩阵划分至多张GPU上进行训练。

当然，除了常见的数据并行、模型并行以外，还有数据-模型混合并行，流水并行等其他并行方式，本文侧重数据并行模型训练。

2. 分布式训练和集合通信

分布式框架采用的常见底层支撑库，可以归类为以下三类：

• 集合通信库：用于分布式训练时，多个计算设备之间的集合通信，常见的有 Open MPI、NCCL、Gloo等

• 数据加载与预处理库：分布式训练需要处理海量数据，这使得单机单卡时代不需要考虑的数据加载问题，在分布式时代很容易成为瓶颈，为此，通常需要对分布式训练中的数据即预处理做相关优化。如：现在市面上比较流行的Nvidia提供的DALI库。

• 分布式通信调度库：虽然分布式深度学习框架的通信大都采用的是以上提及的NCCL、Open MPI一类的通信库，但是因为深度学习框架设计的历史包袱以及实现质量的参差不齐，各框架原生的通信接口的效果并不稳定，也不易上手。在实际应用中，反而是第三方的Horovod成为了各框架的“标配第三方库”。Horovod 对底层通信库（MPI、NCCL、Gloo等）进行了封装，针对TensorFlow，Keras，PyTorch和Apache MXNet的分布式深度学习训练框架，提供稳定高效的通信能力。Horovod的提供的上层接口也更通用和易于用户上手，通常降低了学习门槛。

本文将主要介绍分布式深度学习中集合通信相关的概念以及常用的集合通信库(Open MPI、NCCL、Gloo)。

何为集合通信(Collective communication)？

要说集合通信，首先得了解P2P点对点通信(Point-to-point)。P2P通信通常为两个不同进程间的通信，是1对1的；在MPI规范中，既有同步阻塞的P2P接口：MPI_send和MPI_Recv接口，也定义了非阻塞的P2P接口如：MPI_Isend、MPI_Irecve。和P2P通信相对应，集合通信则是1对多或是多对多的。在分布式系统中，各个节点间往往存在大量的集合通信需求，而我们可以用消息传递接口(Message Passing Interface, MPI)来定义一些比较底层的消息通信行为譬如Reduce、AllReduce、Scatter、Gather、Allgather等。

MPI作为高性能计算领域的元老和通信标准，定义了一系列的通信接口，其上层可以由多种编程语音实现(如c/c++、fortran、java……)，有一些比较流行的通信库实现如：MPICH2、OpenMPI，这些通信库用不同的代码/算法实现了MPI的接口定义的通信模式。其中常用的通信模式有：

• Broadcast

• Scatter

• Gather

• Reduce

• AllReduce

• All gather

下面，我们通过文字和示意图对其中的一些作简单讲解。

Broadcast&Scatter

在集合通信中，如果某个节点想把自身的数据发送到集群中的其他节点，那么就可以使用广播Broadcast的操作(对应MPI_Bcast接口)。

下面我们通过示意图看一下“广播”的具体过程：圆圈表示分布式系统中的独立节点（进程），如下图的0~3，共4个节点；小方块则代表了数据，颜色相同表示数据一样。

Broadcast代表了一种广播的行为，执行Broadcast时，数据从主节点广播至其他各个指定的节点；和Broadcast类似，Scatter表示一种散播行为，将主节点的数据划分散布至其他指定的节点。

Reduce

Reduce称为规约运算，是一系列运算操作的统称，细分来说包括SUM、MIN、MAX、PROD、LOR等。Reduce意为减少/精简，因为其操作在每个进程上获取一个输入元素数组，通过执行操作后，将得到精简的更少的元素。例如下面的Reduce Sum：

AllReduce

Reduce是一系列操作的统称，AllReduce则是在所有的节点进程上都应用同样的Reduce操作。

AllReduce Sum：

从图中可以看出，AllReduce操作可通过单节点上Reduce+Broadcast操作完成。

更多mpi相关的教程请参考：

点击获取更多mpi教程

3. 集合通信库

分布式系统中因为面临大量的信息同步、更新需求，因此传统的点对点的通信方式不能很好的满足需求。上世纪90年代针对HPC领域定义了一套集合通信相关的接口标准，称为 MPI，其中常用的集合通信模式(broadcast、gather、AllReduce……)深刻地影响着后续的分布式通信库的实现。

时至今日，作为分布式深度学习框架的底层支撑库，几大集合通信库（如：Open MPI、Gloo、NCCL等）都在MPI的基础上，对各种集合通信的模式和算法作了各自的实现。

Open MPI

借用官网描述：Open MPI项目是一个开源MPI（消息传递接口 ）实现，由学术，研究和行业合作伙伴联盟开发和维护。因此，Open MPI可以整合高性能计算社区中所有专家，技术和资源，以构建可用的最佳MPI库。

Gloo

Gloo是Facebook开源的一套集体通信库，提供了对机器学习中有用的一些集合通信算法如Bbarrier, Broadcast, AllReduce。

NCCL

NCCL是英伟达基于NCIDIA-GPU的一套开源的集合通信库，如其官网描述：NVIDIA集合通信库（NCCL）实现了针对NVIDIA GPU性能优化的多GPU和多节点集合通信原语。NCCL提供了诸如All-Gather, All Reduce, Broadcast, Reduce, Reduce-Scatter等实现，这些实现优化后可以通过PCIe和NVLink等高速互联，从而实现高带宽和低延迟。

因为NCCL则是NVIDIA基于自身硬件定制的，能做到更有针对性且更方便优化，故在英伟达硬件上，NCCL的效果往往比其它的通信库更好。

4.各通信库&AllReduce

分布式深度学习框架的开发者和使用者，可能最关心是集合通信库里的 AllReduce 操作，因为它的使用频率最高，也最必须。我们先介绍下为何 AllReduce 操作对于深度学习框架如此重要：

在1.1节，我们知道了数据并行训练的主要过程大致分为3步：

• 数据划分：不同GPU设备上划分出不同的mini-batch，作为训练的数据集

• 前向+反向：不同GPU设备上用相同的模型，用各自接收到的mini-batch数据进行训练（前向和反向传播)

• 梯度同步更新：每个GPU设备得到了mini-batch训练后的权重值，这些值需要汇总然后更新至每一个GPU设备，保证每一次迭代后，每个GPU设备上的模型完全一致。

我们可以看见，第3步梯度同步更新包含从各节点收集梯度、汇总、更新至每一节点的全部过程，这些组合起来就是一个AllReduce的过程，具体点说是AllReduce sum操作。通过AllReduce sum，各自节点更新的梯度值汇总后，再更新至每一个节点。由此可见，深度学习的分布式训练和AllReduce的关系是十分紧密的。

那以上介绍的通信框架，他们的AllReduce 表现各自如何呢？请看下文介绍。

AllReduce哪家强？

OpenMPI

Open MPI的实现里，有各种各样的all reduce算法，在最新的OpenMPI-4.0.5的代码中(openmpi-4.0.5/ompi/mca/coll/tuned/coll_tuned_allreduce_decision.c)，我们可以看到有7种不同的all reduce算法实现：

    {0, "ignore"},

    {1, "basic_linear"},

    {2, "nonoverlapping"},

    {3, "recursive_doubling"},

    {4, "ring"},

    {5, "segmented_ring"},

    {6, "rabenseifner"},

其中，Ring AllReduce算法是在深度学习分布式训练环境下比较优秀的一种all reduce算法，关于这些算法的更具体的比较和分析，可以参考：点击阅读AllReduce算法的前世今生

NCCL Allreduce

英伟达于2015年公开发布NCCL,一个开源的、基于自身硬件的开源的集合通信库实现。其算法基本实现原理，和mpi的实现是基本类似的，由于其完全基于自家硬件，可以进行充分的优化，所以基于nvidia-gpu时，使用NCCL性能是很强的，真香！

NCCL的实现参考了MPI接口，定义了一系列名称不同但功能类似的通信方法。针对GPU部分常用的接口做了比较大的优化，一些其他的则没有实现，有点类似于MPI的一个对GPU通信支持的很强的子集。具体可以参考英伟达自己的说明。点击阅读该说明

简单总结就是：

• Open MPI 的Allreduce 算法实现比较丰富，并且与 MPI 接口完全一致，比较适合做 MPI 研究的同学

• NCCL 可以轻松与 MPI结合使用。将MPI用于CPU到CPU的通信，将NCCL用于GPU到GPU的通信

• NCCL 因为是硬件厂商英伟达针对自家产品推出的通信库，性能优化方面有得天独厚的优势

随着NVIDIA-GPU的普及以及NCCL出色的性能，各大深度学习框架基本都添加了对NCCL的支持，并推荐NCCL作为分布式训练中的集合通信后端。

Gloo Allreduce

Gloo是Facebook开源的用于机器学习任务中的集合通信库，其实现了一些常见的AllReduce算法：

• AllReduce_ring

• AllReduce_ring_chunked

• AllReduce_halving_doubling

• AllReducube_bcube

除此之外，Gloo还有一系列基于CUDA-aware为⽬的的AllReduce实现：

• cuda_AllReduce_ring

• cuda_AllReduce_ring_chunked

• cuda_AllReduce_halving_doubling

• cuda_AllReduce_halving_doubling_pipelined

截至目前，Tensorflow(2.3)中尚未添加对Gloo作为集合通信后端的支持，Pytorch(1.7)支持Gloo作为集合通信后端的，但仅推荐在分布式CPU环境下使用，在分布式GPU环境下，Pytorch官方推荐使用NCCL作为后端，详见：DISTRIBUTED COMMUNICATION PACKAGE - TORCH.DISTRIBUTED

总结

分布式深度学习的通信库选择目前比较通用的是基于MPI的实现(如OpenMPI)、NVIDIA的NCCL实现，如果是在使用NVIDIA-硬件的情况下，主流的选择是NVIDIA自家的NCCL。

关于NCCL、OpenMPI、Gloo的Allreduce对比，可以参考开源项目mlbench-benchmarks。该项目使用pytorch框架，在相同的GPU及软硬件环境一致的情况下，对NCCL、OpenMPI、Gloo三种通信库的Allreduce作了性能测试。

各节点间Allreduce通信时间和发送Tensor尺寸的曲线关系图：

• 对于小张量，OpenMPI的性能优于NCCL（尤其是当workers数量增大时）

• 无论workers数量如何，NCCL在大张量情况下都优于MPI和GLOO。

更多介绍和对比可以参考：Communication Backends, Raw performance benchmarking

其他

除了上文介绍的通信库中的AllReduce外，还有其它的通信库，也实现了 AllReduce 算法，虽然相对小众，但是在特定的方面有较大的性能提升。在此也一并介绍。

尽管在MPI的各种实现中（譬如OpenMPI），很早就有了优秀的Ring AllReduce算法，不过将其引入到深度学习中，还是百度首创的。百度2016年在论文：Bringing HPC Techniques to Deep Learning中介绍了一种来自高性能分布式计算中的概念——Ring AllReduce，并将其引入了深度学习（给Tensorflow贡献了代码，增加了基于mpi源语实现的ring all reduce），且获得了显著的性能提升。

可参考知乎：[翻译] Bringing HPC Techniques to Deep Learning

代码：https://github.com/baidu-research/baidu-allreduce

基于double binary tree的AllReduce

《Two-Tree Algorithms for Full BandwidthBroadcast, Reduction and Scan》

double binary tree于2009年在MPI中引入，并随后在NCCL2.4中也引入了此实现：https://developer.nvidia.com/blog/massively-scale-deep-learning-training-nccl-2-4/#ref3

分层ring AllReduce

《Highly Scalable Deep Learning Training System with Mixed-Precision: Training ImageNet in Four Minutes》

基于spanning tree的AllReduce

《Blink: Fast and Generic Collectives for Distributed ML》

基于异构环境下、拓扑感知的AllReduce框架BlueConnect

《BlueConnect: Decomposing All-Reduce for Deep Learning on Heterogeneous Network Hierarchy》

撰文：赵露阳

2020/11/05