大模型训练（7）：集合通信与通信原语

0 背景

分布式训练过程中设计到许多通信上的操作，每个操作有其不同的术语并且有所区别，这里将其用简单的例子和描述总结一下，方便理解。

集合通信（Collective Communications）是一个进程组的所有进程都参与的全局通信操作，其最为基础的操作有发送send、接收receive、复制copy、组内进程栅障同步Barrier以及节点间进程同步，这几个最基本的操作经过组合构成了一组通信模板也叫通信原语，比如：

1对多的广播broadcast
多对1的规约reduce
多对多的规约all-reduce
多对1的收集gather
多对多的收集all-gather
1对多的发散scatter
组合的规约与发散reduce-scatter
多对多的all-to-all等

集合通信的难点在于通信效率以及网络硬件连接拓扑结构的最佳适用。其设备间的物理连接方式可以是私有物理互联协议，比如CXL、NVLINK，也可以是PCIe、InfiniBand、Ethernet等，本文将以此物理拓扑结构描述集合通信中常用的几组通信原语。

1 Broadcast（广播）

将一个设备上的数据广播到各个设备上，其他设备对该设备只进行读数据。

Broadcast属于1对多的通信原语，一个数据发送者，多个数据接收者，可以在集群内把一个节点自身的数据广播到其他节点上。如上图所示，GPU1~4表示集群中的训练加速卡节点，小方块表示内存空间，最后都变成了a1b1c1d1。当主节点 1 执行Broadcast时，数据即从主节点0被广播至其他节点。

Broadcast是数据的1对多的同步，它将一张XPU卡上的数据同步到其他所有的XPU卡上，其应用场景有：

数据并行的参数初始化，确保每张卡上的初始参数是一致的
All-Reduce里 broadcast + reduce组合操作里的broadcast部分
分布式训练 parameter server 参数服务器结构里的 master 节点 broadcast 数据到 worker 节点，再从 worker 节点 reduce 数据回 master 节点里的broadcast操作

2 Reduce（规约）

Reduce是函数式编程的经典概念。数据规约涉及通过函数将一组数字缩减为一个较小的集合。例如，假设我们有一个数字列表[1,2,3,4,5]。用sum函数缩减这个数字列表将产生sum（[1,2,3,4,5]）= 15。同样，乘法规约将产生乘法([1, 2, 3, 4, 5]) = 120。

在上图的例子中，各设备上相同位置的元素进行加和，并将结果呈现在一个设备上（图中的例子为设备1）。

Reduce属于多对1的通信原语，具有多个数据发送者，一个数据接收者，可以在集群内把多个节点的数据规约运算到一个主节点上，常用的规约操作符有：求累加和SUM、求累乘积PROD、求最大值MAX、求最小值MIN、逻辑与 LAND、按位与BAND、逻辑或LOR、按位或BOR、逻辑异或LXOR、按位异或BOXR、求最大值和最小大的位置MAXLOC、求最小值和最小值的位置MINLOC等，这些规约运算也需要加速卡支持对应的算子才能生效。

Reduce是数据的多对1的规约运算，它将所有张XPU卡上的数据规约（比如SUM求和）到1张XPU卡上，其应用场景有：

All-Reduce里的 broadcast + reduce组合里的reduce操作
Reduce-Scatter组合里的 reduce操作
分布式训练 parameter server 参数服务器结构里的 master节点 broadcast 数据到worker节点，再从worker节点reduce数据回master节点里的reduce操作；

3 All Reduce（全集规约）

相当于Reduce之后再来了一个Broadcast。

All-Reduce属于多对多的通信原语，具有多个数据发送者，多个数据接收者，其在集群内的所有节点上都执行相同的Reduce操作，可以将集群内所有节点的数据规约运算得到的结果发送到所有的节点上。All-Reduce操作可通过在主节点上执行Reduce + Broadcast或Reduce-Scatter + All-Gather实现。每个节点的数值就都保持一致。

All-Reduce 是数据的多对多的规约运算，它将所有的XPU卡上的数据规约（比如SUM求和）到集群内每张XPU卡上，其应用场景有：

AllReduce应用于数据并行
数据并行各种通信拓扑结构比如Ring，All-Reduce、Tree All-Reduce里的 All-Reduce操作

4 Gather（收集）

收集（即把东西放到一起，并不做运算），与Reduce不同的地方是，Gather只是将数据汇总到一起，而Reduce需要按照指定的映射函数进行运算。

Gather操作属于多对1的通信原语，具有多个数据发送者，一个数据接收者，可以在集群内把多个节点的数据收集到一个节点上，如上图所示，不同数据块都汇聚到了GPU1上。

Gather是数据的多对1的收集，它将多张XPU卡上的数据收集到1张XPU卡上，他的反向操作对应Scatter。

5 All Gather（全集收集）

即将所有人的所有数据都让所有人知道，做到所有数据透明。和Broadcast的区别是，这里是针对全局所有机器的数据来说，而Broadcast是针对一个机器，让其他所有人都知道你的数据。也就是All Gather数据传输的量更大。

All-Gather是数据的多对多的同步全收集，它将多张XPU卡上的数据收集到多张XPU卡上，可以看做Gather + Broadcast的操作组合，它的反向操作对应Reduce-Scatter，其最应用场景有：

All-Gather可应用于模型并行
模型并行里前向计算里的参数全同步，需要用All-Gather把模型并行里将切分到不同的XPU上的参数全同步到一张XPU上才能进行前向计算

6 Scatter（离散）

中文意思为离散，扩散。即将一个机器上的不同数据分别给到不同机器。

同Broadcast一样，Scatter也是一个1对多的通信原语，也是一个数据发送者，多个数据接收者，可以在集群内把一个节点自身的数据发散到其他节点上。与Broadcast不同的是Broadcast把主节点0的数据发送给所有节点，而Scatter则是将数据的进行切片再分发给集群内所有的节点，如上图所示，主节点 1 将数据分为3份分发到了节点1-3。

Scatter是数据的1对多的分发，它将一张XPU卡上的数据进行分片再分发到其他所有的XPU卡上，他的反向操作对应Gather，其应用场景有：

Reduce-Scatter 组合里的 Scatter 操作
模型并行里初始化时将模型 Scatter 到不同的XPU上

7 Reduce-Scatter（规约离散）

Reduce_scatter最终呈现效果为：每个GPU上有一部分完成所有设备粒度规约的数据，但是只有一块，其他块在其他的设备上。他和All reduce的区别在于，All reduce是所有设备上都有完整规约的数据。

这是一个不常见的通信函数，在深度学习显存优化ZeRO中经常被用到，如下图所示。

Reduce-Scatter 属于多对多的通信原语，具有多个数据发送者，多个数据接收者，其在集群内的所有节点上都按维度执行相同的Reduce规约运算，再将结果发散到集群内其他的节点上，Reduce-Scatter等价于节点个数次的reduce规约运算操作，再后面执行节点个数的scatter次操作，其反向操作是All-Gather。

ReduceScatter是数据的多对多的reduce + scatter运算，它将所有的XPU卡上的数据先规约（比如SUM求和）到1张XPU卡上，再进行scatter，其应用场景有：

Reduce-Scatter即可应用于数据并行也可应用于模型并行
数据并行All-Reduce里的 Reduce-Scatter + All-Gather组合里的Reduce-Scatter操作
模型并行里在前向All-Gather后的反向计算里的Reduce-Scatter

8 All-To-All

All-To-All 操作每一个节点的数据会 Scatter 到集群内所有节点上，同时每一个节点也会 Gather 集群内所有节点的数据。

All-To-All 是对 All-Gather 的扩展，区别是 All-Gather 操作中，不同节点向某一节点收集到的数据是相同的，而在 All-To-All 中，不同的节点向某一节点收集到的数据是不同的，如下图所示

All-To-All是数据的多对多的转置，它将所有张XPU卡上的数据转置到所有的XPU卡上，其主要应用场景有：

All-To-All应用于模型并行；
模型并行里的矩阵转置；
数据并行到模型并行的矩阵转置；

9 Send 与 Receive

数据或参数在不同XPU之间的发送与接收。

10 Barrier

Barrier同步操作会阻塞所有的调用者直到所有的组内成员都调用了它，用于一个集合通信子中所有进程的同步，调用函数时进程将处于等待状态，直到通信子中所有进程都调用了该函数后才继续执行。

11 Signal与Wait

Signal与Wait属于记录型信号量机制： wait(s)，signal(s)可用于解决进程间的同步问题，在通信原语里从一个节点发送一个数据到另外一个节点时，会同时signal一个event值到对端，对端的wait操作接收到这个event时会返回一个确认给signal，这样保证在节点的进程间进行数据的同步操作。

12 相互关系与总结

在分布式训练过程中，深度学习训练框架不会去直接操作底层的通信网络，而是通过使用网络通信库来完成数据的集合通信，各家AI芯片加速卡厂家都会提供私有的网络通信库比如：xxx-AWARE OpenMPI或xCCL来完成这个底层通信硬件的屏蔽与抽象。在分布式训练集群里网络通信硬件连接样式多种多样，可以是Ethernet、InfiniBand 、RoCE v2/v1 等也可以是CXL、NVLINK等私有协议，这就要求在通信的后端层根据各个厂家的自己的SDK开发库接口，根据实际情况实现各自的网络通信库，比如cuda-aware MPI、NCCL、NVSHMEM，以及根据实际的网络拓扑组合完成对应的最有效的网络拓扑算法。

本文讲述了分布式训练里的集合通信原语，这些原语是集合通信拓扑算法的基本组成单元，后续的文章里会讲述如何组合这些通信原语以完成合适的通信拓扑算法。