集合通信及其通信原语

1. 前言

大模型训练离不开集合通信

大模型分布式训练往往需要上千乃至上万 GPU 卡进行超大规模并行训练，是典型的计算密集型和通信密集型场景。

在真实的场景中，为了高效的训练，我们一般采用多种并行策略混合的方式。常见的包括数据并行，流水并行，张量并行，序列并行，专家并行，其中张量并行和流水线并行都属于模型并行。

1.1 数据并行

在数据并行模式下，每个 GPU 都运行相同的模型代码，而数据集被拆分为多份分配给不同的 GPU 进行训练。每轮迭代完成后，需要通过 all-reduce 操作进行同步。随着模型规模的增加，单个 GPU 的通信量可以达到 10 GB 以上，每个step都需要进行一次通信。

一种优化方法是将数据并行中的 all-reduce 操作拆分为 reduce-scatter 和 all-gather 两个部分。

在训练过程中，每个 GPU 只存储模型的一部分：

• Forward Pass：通过 all-gather 将模型参数聚集到每个 GPU，然后进行前向计算。
• Backward Pass：同样通过 all-gather 将模型参数聚集到每个 GPU，计算出本地梯度后，通过 reduce-scatter 将平均后的梯度分摊到各自的 GPU 上，然后进行本地的权重更新。

通过这种方法有效地降低了通信开销，提高了数据并行训练效率。

1.2 模型并行

在模型并行模式下，每个 GPU 负责模型的一部分参数和计算。流水并行通过按层切分来实现，而张量并行则通过分解模型的张量来进行。

例如，多头自注意力（MHA）非常适合张量并行。Q、K、V 矩阵按列切分，每个头的计算独立分配到不同的 GPU 上。在张量并行中，虽然需要通过 AllReduce 操作来同步矩阵乘法的结果，但通信的数据量与batch size大小相关，矩阵大小可以达到 GB 级别，一个step可能需要进行几十次通信。

流水并行，就是将输入的 batch size 切分为多个mini-batch 的数据，通过划分为多个阶段进行并行计算，掩盖计算过程中的“气泡”时间。

例如，Megatron通过将模型的层被均匀分配到多个 GPU 上，每个 GPU 负责一部分层，将每个 pipeline stage 进一步划分为多个虚拟阶段并行计算，降低气泡比率，提高资源利用率。在流水并行中，需要点对点的集合通信，其通信的数据量一般在 MB 级别，一个 step 几十次通信。

1.3 序列并行

序列并行，是针对张量并行的进一步补充，对于那些需要全局统计信息不能按 Tensor 维度拆分的，序列并行则按照序列进行切分，使得计算被进一步分摊，减少每个 GPU 的显存占用。

例如，在 Megatron 中，序列并行通过将 LayerNorm 和 Dropout 的输入在序列维度上切分，切分数量等于张量并行的数量。这使得每个 GPU 只需处理部分 LayerNorm 和 Dropout 计算，整体上保持数学等价性。

使用序列并行后：

• 前向传播: 需要使用 all-gather 来合并序列切分的部分，以确保完整输入用于后续计算。然后使用 reduce-scatter 来分散计算结果，从而减少 Dropout 的显存占用。
• 反向传播: 需要使用 reduce-scatter 来分散梯度计算，再使用 all-gather 来合并梯度信息。

虽然序列并行改变了通信模式，但总体通信量没有增加，因为 all-reduce 可以等效为 reduce-scatter 和 all-gather 的组合。这些优化提升了计算效率，同时有效管理了显存使用。

1.4 专家并行

专家并行，就是通过选择性的激活一部分参数来处理不同的数据，从而解决模型规模增大训练成本成平方级别的增长的问题。

作为基于 Transformer 的 MoE 模型，主要由以下两部分组成：

• 稀疏 MoE 层：它将 FFN 拆成多个子层，每一个子层被称为Expert。一般来说，这些 Expert 都是 FFN，但是也可以是更复杂的网络，甚至是 MoE 本身。
• Router：也被称为 Gating Network，这部分用于决定将哪些 token 被发送到哪些 Expert。

专家并行的思路是将不同的专家分配到不同的 GPU 上，这有助于减少内存消耗并提高训练效率。计算前需要根据路有规则将 Token 通过 All-to-All 通信发送给不同的 Experts 所在的 GPU 进行运算。

并行训练中使用的集合通信

从下表可以看出，不同的并行模式下都离不开集合通信技术，所以说大模型的分布式训练是离不开集合通信的。

2. Collective Communication 和 NCCL 的出现

说起集合通信，就不得不提到 MPI。MPI 是集合通信中的元老，是消息传递的平台，是一种标准和规范，而非具体并行语言。真正的实现是在开源实现库里，例如openMPI。

2.1 What is MPI

我们知道，一旦我们的程序运行起来，它就会成为独立的进程。进程拥有独立的执行环境（内存、寄存器、程序计数器等），是操作系统中独立存在的可执行的基本程序单位。

进程间相互独立（内存空间不相交），但是进程间可以相互交换信息。消息传递是指这些信息在进程间的相互交换，是实现进程间通信的唯一方式。

最基本的消息传递操作包括发送消息send、接受消息receive、进程同 步barrier、归约reduction等。

在 OpenMPI 中，有多种集合通信算法的实现,例如 Allreduce, AlltoAll, AllGather等，且每种也都有不同的实现。

MPI 的使用示例（Allreduce）：

​ MPI 中实现了多种 AllReduce 算法，例如传统的 reduce+broadcast 的方式、或者 butterfly、Ring AllReduce 和 Segmented Ring 的方式。

一个 Allreduce 有多种实现，可以方便在不同的情况下采用不同的 Allreduce 的实现，来加快集合通信的计算。

集合通信（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等，集合通信的难点在于通信效率以及网络硬件连接拓扑结构的最佳适用。

3. 通信原语（常见集合通信算子介绍）

以一台集成了4张训练加速卡的服务器为例，如下图，服务器内四张训练加速卡是全连接的，物理连接方式可以是私有物理互联协议，比如Ethernet、NVLINK，也可以是PCIe、InfiniBand、ROCE等，本文将以此物理拓扑结构描述集合通信中常用的几组通信原语。

3.1 Broadcast

Broadcast属于1对多的通信原语，一个数据发送者，多个数据接收者，可以在集群内把一个节点自身的数据广播到其他节点上。如下图所示，圈圈表示集群中的训练加速卡节点，相同的颜色的小方块则代表相同的数据。当主节点 0 执行Broadcast时，数据即从主节点0被广播至其他节点。

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

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

3.2 Scatter

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

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

• ReduceScatter组合里的Scatter操作；
• 模型并行里初始化时将模型scatter到不同的XPU上；

3.3 Gather

Gather操作属于多对1的通信原语，具有多个数据发送者，一个数据接收者，可以在集群内把多个节点的数据收集到一个节点上，如下图所示，不相同的颜色的小方块代表不相同的数据。

Gather是数据的多对1的收集，它将多张XPU卡上的数据收集到1张XPU卡上，他的反向操作对应Scatter，其应用场景有：

• ReduceScatter组合里的Scatter操作；

3.4 AllGather

AllGather属于多对多的通信原语，具有多个数据发送者，多个数据接收者，可以在集群内把多个节点的数据收集到一个主节点上（Gather），再把这个收集到的数据分发到其他节点上（broadcast），即收集集群内所有的数据到所有的节点上。

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

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

3.5 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算子来完成这个求和累加。

3.6 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，其应用场景有：

• ReduceScatter即可应用于数据并行也可应用于模型并行；
• 数据并行allReduce里的 ReduceScatter+ Allgather组合里的ReduceScatter操作；
• 模型并行里在前向allgather后的反向计算里的ReduceScatter；

3.7 AllReduce

AllReduce属于多对多的通信原语，具有多个数据发送者，多个数据接收者，其在集群内的所有节点上都执行相同的Reduce操作，可以将集群内所有节点的数据规约运算得到的结果发送到所有的节点上。AllReduce操作可通过在主节点上执行Reduce + Broadcast或ReduceScatter + AllGather实现，如下图所示：先在主节点上执行reduce得到规约累加和26，再把这个累加和26 broadcast到其他的节点，这样整个集群内，每个节点的数值就都保持一致。

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

• AllReduce应用于数据并行；
• 数据并行各种通信拓扑结构比如Ring allReduce、Tree allReduce里的 allReduce操作；

3.8 AlltoAll

All-To-All操作每一个节点的数据会scatter到集群内所有节点上，同时每一个节点也会Gather集群内所有节点的数据。ALLTOALL是对ALLGATHER的扩展，区别是ALLGATHER 操作中，不同节点向某一节点收集到的数据是相同的，而在ALLTOALL中，不同的节点向某一节点收集到的数据是不同的，如下图所示：

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

• AllToAll应用于模型并行；
• 模型并行里的矩阵转置；
• 数据并行到模型并行的矩阵转置；

3.9 Send & Receive

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

3.10 Barrier

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

3.11 Signal & Wait

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

4. Why NCCL was created ?

在 2009年openMPI的算法已经完全成熟，但为什么 2015 年英伟达又发布了 Nccl 呢？

这是因为一方面因为MPI基本没有考虑过GPU系统架构，在MPI设计中各个工作节点基本视为等同，并没有考虑节点间latency和带宽的不同。而在GPU中，P40, V100，A100的性能相差很大，不同设备存在异构性，需要重新设计策略来考虑异构场景下的硬件性能。而Nccl就是为了更贴合GPU硬件。NCCL1.x只能在单机内部进行通信，这时多机间的通信还需依赖MPI，从NCCL2.0开始支持多机间通信。

5. NCCL 的特点

NCCL 是专为 NVIDIA GPU 设计的通信库，它充分利用了 NVIDIA 硬件的特性，包括但不限于：

• NVLink 互连：NCCL 优化了通过 NVLink 进行的 GPU 间通信，这是 NVIDIA 高端 GPU（如 Tesla V100 和 A100）之间的高速互连技术。
• GPU 直接通信：NCCL 允许 GPU 之间直接交换数据，绕过了传统的 CPU 中转，显著减少了通信延迟。
• CUDA 集成：NCCL 与 CUDA 紧密集成，使得开发者能够在 CUDA 程序中无缝地使用 NCCL 进行高效的数据传输和同步。
• 自适应拓扑：NCCL 能够自动检测并适应不同的硬件拓扑，无论是单节点多GPU还是跨节点的多GPU环境。

在异构计算环境中，CPU、GPU、网络和其他加速器需要高效协同工作。NCCL 在这种环境中展现出显著的性能优势：

• 跨平台支持：NCCL 不仅支持 NVIDIA GPU，还能够与 CPU 和其他硬件平台协同工作，提供了跨平台的通信解决方案。
• 可扩展性：NCCL 设计用于大规模并行计算，能够支持数千个 GPU 的集群，这对于处理大规模数据集和模型至关重要。
• 低延迟和高吞吐量：NCCL 优化了数据传输路径，减少了通信延迟，同时提供了高吞吐量的数据处理能力。
• 容错性：在异构环境中，硬件故障是常见的问题。NCCL 提供了容错机制，确保了计算任务在出现硬件故障时仍能继续执行。

6. 小结

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

下个章节中将介绍NCCL/RCCL 的实现原理。

reference

一文讲清 NCCL 集合通信原理与优化 - 极术社区 - 连接开发者与智能计算生态
分布式训练——集合通信及其通信原语