分布式训练硬核技术——通讯原语

作者：ZOMI酱

文章来源：https://zhuanlan.zhihu.com/p/465967735

针对分布式训练服务器的集群进行架构设计，是为了更好地解决机器学习中分布式训练中的通讯问题。目前机器学习中主要由两种分布式架构：

参数服务器架构（Parameter Server，PS）
去中心化架构（Decentralized Network）

其中，分布式训练通常在计算集群上进行，集群的每个节点分别执行一部分计算。不同节点的计算之间有数据依赖和共享，需要将数据在不同节点间传输，这就是通信。分布式的通信一般有两大类：

集合通信（Collective communication，CC）：在一组节点内进行通信
点对点通信（Point to point communication，P2P）：在两个节点之间进行通信

深度学习训练过程中因为需要传输大量的网络模型权重参数和训练过程中产生的大量临时变量等，因此主要使用集合通信的方式。可以理解为，机器学习/深度学习的分布式训练，主要是采用在PS架构下的集合通讯模式；而在大模型的分布式训练中，因为减少跟单点参数服务器统一更新，更多直接采用纯集合通讯模式。

在深度学习框架中，分布式训练的通讯原语和通讯实现方式对AI框架分布式训练起着非常重要的作用，如果想要训练大模型（Foundation Model）肯定离不开进行通讯操作，下面橙色标签的是分布式训练中通讯原语在AI框架中的的位置。

preview

为什么需要对通讯操作

当将神经网络的训练并行化到集群中不同的节点时，必须选择如何将不同的计算操作分配到集群中可用的节点，这就离不开对分布式集群的通讯进行操作。

数据并行通讯问题

通常神经网络的训练学习过程中，使用随机梯度下降算法SGD进行训练。

以分布式数据并行为例，在数据并行训练中，每块NPU都有整个神经网络模型的完整副本，每次进行迭代的时候，只分配了小批次Batch中数据样本的子集，每块NPU利用该Batch样本数据进行网络模型的前向计算，然后计算损失误差进行反向传播，以计算损耗相对于网络参数的梯度。

最后，NPU 之间进行相互通信，计算由不同 NPU 得到的梯度的平均值，将平均梯度应用于总体权重以获得网络模型中本次迭代的最新权重。任何一块 NPU 都在锁定步骤的迭代中前进，一旦某块 NPU 完成了本轮迭代，必须等待集群中所有 NPU 完成各自的迭代，这样神经网络的总体权重才能被正确更新。简单地理解，相当于在单个 NPU 上执行SGD，但是通过在多个 NPU 之间分发数据，通过并行执行计算来获得加速。这就是简单的分布式数据并行（DP）方式。

当只有两个 NPU 和以兆字节数据衡量网络模型参数的时候，NPU 的通信方式可能并不重要。然而，当网络模型的参数量去到十亿甚至万亿个参数时，仅仅是计算梯度所产生的临时变量就需要数十亿字节的 NPU 内存空间，并且需要协调数十个 NPU 共同工作，这个时候通信机制变得至关重要。

例如，考虑上图中最直接的数据并行通信方式：每一个 NPU 都计算其小Batch上的梯度，然后每个 NPU 将其梯度发送到 CPU 上，该 CPU 取所有梯度的平均值，并将平均值发送回所有其他 NPU。

直接从单个 CPU 发送和接收数据的机制中，CPU 必须从所有 NPU 中接收参数，并将所有参数逐一发送到所有 NPU中，当分布式集群系统中的 NPU 节点越多，通信成本就越大。

现在来以上述的通信策略评估在真实网络模型上运行的带宽效果。以三亿个可训练参数的网络模型为例：假设每个参数四个字节，三亿个参数约1.2千兆字节，分布式系统的网络硬件可以支持每秒1千兆字节的带宽。在这种情况下，将系统并行化到两个 NPU 上将使每次迭代减慢1.2秒，将训练并行化到10个 NPU 将使每次迭代减慢10.8秒。随着 NPU 数量的增长，每次迭代所需的时间将会呈线性增长。

需要发送的数据越多，通讯的时间就越长。每个通信通道都有一个最大的吞吐量，即网络带宽。例如，一个Internet连接可以提供每秒15兆字节的带宽，而千兆以太网连接可以提供每秒125兆字节的带宽，HPC集群上的专用网络硬件Infiniband可以在节点之间提供每秒GB级的带宽。即使每次迭代花费几秒钟，通信成本的线性增长也会使得分布式并行，并没有带来实际的优化，反而会降低训练效率。

异步 vs 同步通讯

异步 vs 同步通信属于网络模型参数更新的策略。

数据并行中，在各个 NPU 分别根据各自获得的batch，前向计算获得损失，进而反向传播计算梯度。计算好梯度后，就涉及到一个梯度同步的问题：每个节点都有根据自己的数据计算的梯度，如何在不同节点之间维护模型的不同副本之间的一致性？

如果不同的模型以某种方式最终获得不同的权重，则权重更新将变得不一致，最终导致网络模型参数收敛迭代收敛情况前后不一致，这称为模型一致性问题（consistency of model），如何同步更新网络模型参数，这是分布式机器学习系统的一个核心问题。

为了解决模型一致性问题，分布式训练的梯度同步策略可分为异步（Asynchronous）梯度更新和同步（Synchronous）梯度更新机制。

同步指的是所有的节点都采用相同的模型参数来训练，等待所有设备的batch训练完成后，收集它们的梯度然后取均值，最后所有网络模型的参数一次性更新。

由于每一轮迭代结束时，计算的比较快的节点都需等待算得慢的节点信息进行同步，再进行下一轮迭代。类似于木桶效应，所以同步方式的性能取决于集群通讯最慢的节点。

异步通讯中，各节点完成一个batch数据的训练之后，不需要等待其它节点，直接更新Server中的网络模型参数，因为是异步更新参数的因此不需要阻塞性等待单节点的数据，使得总体会训练速度会快很多。

异步通讯的问题是梯度失效（stale gradients），刚开始所有节点采用相同的参数来训练，但是异步情况下，某节点完成 T 步迭代训练后，可能发现模型参数其实已经被其它设备在 T+n 步更新过了，如果当前节点把 T 步迭代训练参数更新到 Server中，会造成 n 步内的梯度更新无效。由于梯度失效问题，异步训练虽然速度快，但是可能陷入次优解（sub-optimal training performance）。

回到上面三亿个可训练参数的网络模型案例中，另一种参数更新的选择是异步通讯，虽然异步更新会使得数据同步不需要阻塞性地等待，减少通讯的延时，但是消除了参数的同步更新后会使得SGD很难调试和收敛，因此不考虑异步更新的策略。

总的来说，在训练大模型的时候离不开分布式并行操作，需要利用分布式集群中的带宽通讯优化操作，优化深度学习参数同步策略来解决通信瓶颈。

通讯原语操作

集合通讯中包含多个sender和多个receiver，一般的通信原语包括broadcast、gather、all-gather、scatter、reduce、all-reduce、reduce-scatter、all-to-all等通信操作进行数据传输，下面将会分别介绍其具体含义。

Broadcast

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

如图所示，圆圈表示分布式系统中的独立节点，一共4个节点，小方块则代表了数据，颜色相同表示数据一样。Broadcast代表广播行为，执行Broadcast时，数据从主节点0广播至其他各个指定的节点（0~3）。

Broadcast操作是将某节点的输入广播到其他节点上，分布式机器学习中常用于网络参数的初始化。如图中，从单个sender数据发送到其他节点上，将0卡大小为1xN的Tensor进行广播，最终每张卡输出均为[1xN]的矩阵。

Scatter

Scatter操作表示一种散播行为，将主节点的数据进行划分并散布至其他指定的节点。

实际上，Scatter与Broadcast非常相似，都是一对多的通信方式，不同的是Broadcast的0号节点将相同的信息发送给所有的节点，而Scatter则是将数据的不同部分，按需发送给所有的节点。

如图所示，从单个sender数据发送到其他节点上。

Reduce

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

在NCCL中的Reduce，从多个sender那里接收数据，最终combine到一个节点上。

All Reduce

Reduce是一系列简单运算操作的统称，All Reduce则是在所有的节点上都应用同样的Reduce操作。以All Reduce Sum为例。

All Reduce操作可通过单节点上Reduce + Broadcast操作完成。在NCCL中的All Reduce中，则是从多个sender那里接收数据，最终合并和分发到每一个节点上。

Gather

Gather操作将多个sender上的数据收集到单个节点上，Gather可以理解为反向的Scatter。

Gather操作会从多个节点里面收集数据到一个节点上面，而不是从一个节点分发数据到多个节点。这个机制对很多平行算法很有用，比如并行的排序和搜索。

All Gather

很多时候发送多个元素到多个节点也很有用，即在多对多通信模式的场景。这个时候就需要 All Gather操作。

对于分发在所有节点上的一组数据来说，All Gather会收集所有数据到所有节点上。从最基础的角度来看，All Gather相当于一个Gather操作之后跟着一个Bcast操作。下面的示意图显示了All Gather调用之后数据是如何分布的。

Reduce Scatter

Reduce Scatter操作会将个节点的输入先进行求和，然后在第0维度按卡数切分，将数据分发到对应的卡上。例如上图所示，每卡的输入均为4x1的Tensor。Reduce Scatter先对输入求和得到[0, 4, 8, 12]的Tensor，然后进行分发，每卡获得1x1大小的Tensor。例如卡0对应的输出结果为[[0.0]]，卡1对应的输出结果为[[4.0]]。

All to All

All to All作为全交换操作，通过All to All通信，可以让每个节点都获取其他节点的值。

在使用 All to All 时，每一个节点都会向任意一个节点发送消息，每一个节点也都会接收到任意一个节点的消息。每个节点的接收缓冲区和发送缓冲区都是一个分为若干个数据块的数组。All to All 的具体操作是：将节点i的发送缓冲区中的第j块数据发送给节点j，节点j将接收到的来自节点i的数据块放在自身接收缓冲区的第i块位置。

All to All 与 All Gather 相比较，区别在于：All Gather 操作中，不同节点向某一节点收集到的数据是完全相同的，而在 All to All 中，不同的节点向某一节点收集到的数据是不同的。在每个节点的发送缓冲区中，为每个节点都单独准备了一块数据。

AI框架中的通信实现

分布式集群的网络硬件多种多样，可以是Ethernet、InfiniBand 等，深度学习框架通常不直接操作硬件，而是使用通信库。之所以采用通信库屏，是因为其蔽了底层硬件细节，提供了统一封装的通信接口。其中MPI和NCCL是最常用的通讯库，MPI专注于CPU的并行通讯，NCCL则专注于GPU的通讯。

Message Passing Interface (MPI)

MPI 信息传递接口，是一个用于编写并行计算程序的编程接口。它提供了丰富全面的通信功能。

MPI 常用于在计算集群、超算上编写程序，比如很多传统科学计算的并行程序。MPI 接口的兼容性好，通信功能丰富，在深度学习框架中主要用于 CPU 数据的通信。

MPI是一个开放接口，有多种实现的库，一种广泛使用的开源实现是 Open MPI。一些硬件厂商也提供针对硬件优化的实现。

NVIDIA Collective Communication Library (NCCL)

NCCL 英伟达集合通信库，是一个专用于多个 GPU 乃至多个节点间通信的实现。它专为英伟达的计算卡和网络优化，能带来更低的延迟和更高的带宽。

NCCL 也提供了较丰富的通信功能，接口形式上与 MPI 相似，可满足大多数深度学习任务的通信需求。它在深度学习框架中专用于 GPU 数据的通信。因为NCCL则是NVIDIA基于自身硬件定制的，能做到更有针对性且更方便优化，故在英伟达硬件上，NCCL的效果往往比其它的通信库更好。

MPI和NCCL的关系

openMPI的通讯算法和通讯操作原语最晚在2009年就都已经成熟并开源了，而Nvidia在2015年下半年首次公开发布NCCL。

既然openMPI已经实现了这么多All Reduce算法，为什么英伟达还要开发NCCL？是不是从此只要NCCL，不再需要MPI了呢？

从openMPI的源码里能看到，其完全没有考虑过深度学习的场景，基本没有考虑过GPU系统架构。很明显的一点，MPI中各个工作节点基本视为等同，并没有考虑节点间latency和带宽的不同，所以并不能充分发挥异构场景下的硬件性能。

Nvidia的策略还是比较聪明，不和MPI竞争，只结合硬件做MPI没做好的通信性能优化。在多机多卡分布式训练中，MPI还是广泛用来做节点管理，NCCL只做GPU的实际规约通信。NCCL可以轻松与MPI结合使用，将MPI用于CPU到CPU的通信，将NCCL用于GPU到GPU的通信。

而NCCL的优势就在于完全贴合英伟达自己的硬件，能充分发挥性能。但是基本的算法原理其实相比openMPI里实现的算法是没有太大变化。

NCCL1.x只能在单机内部进行通信，NCCL2.0开始支持多节点（2017年Q2）。所以在NCCL2之前大家还会依赖MPI来进行集合通信。

画图不易，肝一篇太难，希望可以支持一下ZOMI酱，给他个关注和点赞。