分布式训练 – 第 3 篇 - 集合通信及其通信原语
ref: [1]. https://zhuanlan.zhihu.com/p/493092647
概述
集合通信(Collective Communications)是一个进程组的所有进程都参与的全局通信操作,其最为基础的操作有 发送 send、接收 receive、复制 copy、组内进程栅障同步 Barrier 以及节点间进程同步 (signal+wait),这几个最基本的操作经过组合构成了一组通信模板也叫通信原语,比如:1 对多的广播 broadcast、多对 1 的收集 gather、多对多的收集 all-gather、1 对多的发散 scatter、多对 1 的规约 reduce、多对多的规约 all-reduce、组合的规约与发散 reduce-scatter、多对多的 all-to-all 等,集合通信的难点在于通信效率以及网络硬件连接拓扑结构的最佳适用。
通信原语
以一台集成了 4 张训练加速卡的服务器为例,如下图,服务器内四张训练加速卡是全连接的,物理连接方式可以是私有物理互联协议,比如 CXL、NVLINK,也可以是 PCIe、InfiniBand、Ethernet 等,本文将以此物理拓扑结构描述集合通信中常用的几组通信原语。
Broadcast
Broadcast 属于 1 对多的通信原语,一个数据发送者,多个数据接收者,可以在集群内把一个节点自身的数据广播到其他节点上。如下图所示,圈圈表示集群中的训练加速卡节点,相同的颜色的小方块则代表相同的数据。当主节点 0 执行 Broadcast 时,数据即从主节点 0 被广播至其他节点。
Broadcast 是数据的 1 对多的同步,它将一张 XPU 卡上的数据同步到其他所有的 XPU 卡上,其应用场景有:
1)数据并行的参数初始化,确保每张卡上的初始参数是一致的;
2)allReduce 里的 broadcast + reduce 组合里的 broadcast 操作;
3)分布式训练 parameter server 参数服务器结构里的 master 节点 broadcast 数据到 worker 节点,再从 worker 节点 reduce 数据回 master 节点里的 broadcast 操作;
Scatter
同 Broadcast 一样,Scatter 也是一个 1 对多的通信原语,也是一个数据发送者,多个数据接收者,可以在集群内把一个节点自身的数据发散到其他节点上。与 Broadcast 不同的是 Broadcast 把主节点 0 的数据发送给所有节点,而 Scatter 则是将数据的进行切片再分发给集群内所有的节点,如下图所示,不相同的颜色的小方块代表不相同的数据,主节点 0 将数据分为四份分发到了节点 0-3。
Scatter 是数据的 1 对多的分发,它将一张 XPU 卡上的数据进行分片再分发到其他所有的 XPU 卡上,他的反向操作对应 Gather,其应用场景有: 1)ReduceScatter 组合里的 Scatter 操作; 2)模型并行里初始化时将模型 scatter 到不同的 XPU 上;
Gather
Gather 操作属于多对 1 的通信原语,具有多个数据发送者,一个数据接收者,可以在集群内把多个节点的数据收集到一个节点上,如下图所示,不相同的颜色的小方块代表不相同的数据。
Gather 是数据的多对 1 的收集,它将多张 XPU 卡上的数据收集到 1 张 XPU 卡上,他的反向操作对应 Scatter,其应用场景有:
1)ReduceScatter 组合里的 Scatter 操作;
AllGather
AllGather 属于多对多的通信原语,具有多个数据发送者,多个数据接收者,可以在集群内把多个节点的数据收集到一个主节点上(Gather),再把这个收集到的数据分发到其他节点上(broadcast),即收集集群内所有的数据到所有的节点上。
AllGather 是数据的多对多的同步全收集,它将多张 XPU 卡上的数据收集到多张 XPU 卡上,可以看做 Gather + Broadcast 的操作组合,它的反向操作对应 ReduceScatter,其最应用场景有:
1) AllGather 可应用于模型并行;
2)模型并行里前向计算里的参数全同步,需要用 allgather 把模型并行里将切分到不同的 XPU 上的参数全同步到一张 XPU 上才能进行前向计算。
Reduce
Reduce 属于多对 1 的通信原语,具有多个数据发送者,一个数据接收者,可以在集群内把多个节点的数据规约运算到一个主节点上,常用的规约操作符有:求累加和 SUM、求累乘积 PROD、求最大值 MAX、求最小值 MIN、逻辑与 LAND、按位与 BAND、逻辑或 LOR、按位或 BOR、逻辑异或 LXOR、按位异或 BOXR、求最大值和最小大的位置 MAXLOC、求最小值和最小值的位置 MINLOC 等,这些规约运算也需要加速卡支持对应的算子才能生效。
Reuduce 操作从集群内每个节点上获取一个输入数据,通过规约运算操作后,得到精简数据,如下图的 SUM 求累加和:节点 0 数值 5、节点 1 数值 6、节点 2 数值 7、节点 3 数值 8,经过 SUM 运算后 累积和为 26,即得到更为精简的数值,在 reduce 原语里回会去调用 reduce SUM 算子来完成这个求和累加。
Reduce 是数据的多对 1 的规约运算,它将所有张 XPU 卡上的数据规约(比如 SUM 求和)到 1 张 XPU 卡上,其应用场景有:
1)AllReduce 里的 broadcast + reduce 组合里的 reduce 操作;
2)ReduceScatter 组合里的 reduce 操作;
3)分布式训练 parameter server 参数服务器结构里的 master 节点 broadcast 数据到 worker 节点,再从 worker 节点 reduce 数据回 master 节点里的 reduce 操作;
ReduceScatter
ReduceScatter 属于多对多的通信原语,具有多个数据发送者,多个数据接收者,其在集群内的所有节点上都按维度执行相同的 Reduce 规约运算,再将结果发散到集群内所有的节点上,Reduce-scatter 等价于节点个数次的 reduce 规约运算操作,再后面执行节点个数的 scatter 次操作,其反向操作是 AllGather。
如下图所示,先 reduce 操作 XPU 0-3 的数据 reduce 为 A (A0+A1+A2+A3) + B (B0 + B1 +B2 + B3) + C (C0 + C1 + C2 + C3) + D (D0 + D1 + D2 + D3 ) 到一张 XPU 上,再进行分片 scatter 到集群内所有的 XPU 卡上。
ReduceScatter 是数据的多对多的 reduce + scatter 运算,它将所有的 XPU 卡上的数据先规约(比如 SUM 求和)到 1 张 XPU 卡上,再进行 scatter,其应用场景有:
1)ReduceScatter 即可应用于数据并行也可应用于模型并行;
2)数据并行 allReduce 里的 ReduceScatter+ Allgather 组合里的 ReduceScatter 操作;
3)模型并行里在前向 allgather 后的反向计算里的 ReduceScatter;
AllReduce
AllReduce 属于多对多的通信原语,具有多个数据发送者,多个数据接收者,其在集群内的所有节点上都执行相同的 Reduce 操作,可以将集群内所有节点的数据规约运算得到的结果发送到所有的节点上。AllReduce 操作可通过在主节点上执行 Reduce + Broadcast 或 ReduceScatter + AllGather 实现,如下图所示:先在主节点上执行 reduce 得到规约累加和 26,再把这个累加和 26 broadcast 到其他的节点,这样整个集群内,每个节点的数值就都保持一致。
AllReduce 是数据的多对多的规约运算,它将所有的 XPU 卡上的数据规约(比如 SUM 求和)到集群内每张 XPU 卡上,其应用场景有:
1) AllReduce 应用于数据并行;
2)数据并行各种通信拓扑结构比如 Ring allReduce、Tree allReduce 里的 allReduce 操作;
All-To-All
All-To-All 操作每一个节点的数据会 scatter 到集群内所有节点上,同时每一个节点也会 Gather 集群内所有节点的数据。ALLTOALL 是对 ALLGATHER 的扩展,区别是 ALLGATHER 操作中,不同节点向某一节点收集到的数据是相同的,而在 ALLTOALL 中,不同的节点向某一节点收集到的数据是不同的,如下图所示:
AllToAll 是数据的多对多的转置,它将所有张 XPU 卡上的数据转置到所有的 XPU 卡上,其主要应用场景有:
1) AllToAll 应用于模型并行;
2)模型并行里的矩阵转置;
3)数据并行到模型并行的矩阵转置;
Send 与 Receive
数据或参数在不同 XPU 之间的发送与接收。
Barrier
BARRIER 同步操作会阻塞所有的调用者直到所有的组内成员都调用了它, 用于一个集合通信子中所有进程的同步,调用函数时进程将处于等待状态,直到通信子中所有进程 都调用了该函数后才继续执行。
Signal 与 Wait
Signal 与 Wait 属于记录型信号量机制: wait (s),signal (s) 可用于解决进程间的同步问题,在通信原语里从一个节点发送一个数据到另外一个节点时,会同时 signal 一个 event 值到对端,对端的 wait 操作接收到这个 event 时会返回一个确认给 signal,这样保证在节点的进程间进行数据的同步操作。
小结
在分布式训练过程中,深度学习训练框架不会去直接操作底层的通信网络,而是通过使用网络通信库来完成数据的集合通信,各家 AI 芯片加速卡厂家都会提供私有的网络通信库比如:xxx-AWARE OpenMPI 或 xCCL 来完成这个底层通信硬件的屏蔽与抽象。在分布式训练集群里网络通信硬件连接样式多种多样,可以是 Ethernet、InfiniBand 、RoCE v2/v1 等也可以是 CXL、NVLINK 等私有协议,这就要求在通信的后端层根据各个厂家的自己的 SDK 开发库接口,根据实际情况实现 各自的网络通信库,比如 cuda-aware MPI、NCCL、NVSHMEM,以及根据实际的网络拓扑组合完成对应的最有效的网络拓扑算法。
本文讲述了分布式训练里的集合通信原语,这些原语是集合通信拓扑算法的基本组成单元,后续的文章里会讲述如何组合这些通信原语以完成合适的通信拓扑算法。

