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

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

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

1.1 数据并行

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

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

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

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 进行运算。

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

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

为了更好的实现 scaling law 扩大规模，如上图所示，在真实的训练中，我们一般采用多种模式混合并行。那么模型的训练中也处处充满集合通信。

2. 集合通信与 NCCL 的出现

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

2.1 什么是MPI

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

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

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

MPI 程序一般包含，变量定义，MPI 环境初始化, 执行进程间通信和退出MPI 环境。

MPI系统的通信方式都建立在点对点通信之上，包括阻塞式点对点通信和非阻塞式点对点通信等。

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

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

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

2.2 常见集合通信算子介绍

MPI_Scatter与MPI_Bcast的区别

MPI_Scatter与MPI_Bcast非常相似，都是一对多的通信方式，不同的是后者的0号进程将相同的信息发送给所有的进程，而前者则是将一段array 的不同部分发送给所有的进程，其区别可以用上图概括。

AllGather 与 Allreduce 的区别

当数据分布在所有的进程中时，MPI_Allgather 是将所有的数据聚合到每个进程中。而MPI_Allreduce是将所有的数据执行规约（sum、min 等）后到每个进程中。

RingAllreduce

RingAllreduce 是一种基于环（Ring）拓扑结构的 Allreduce 操作。在这种算法中，每个参与的进程（或GPU）只与它的两个邻居通信。数据在环中的每个节点之间传递，通过一系列的归约（Reduce）和数据交换操作来聚合数据。这种算法的优点是它只需要局部通信，因此可以减少通信延迟和网络拥塞。

Ring算法默认把每个节点的数据切分成N份。当然，这要求数据块中的元素个数count =S/sizeof(element)大于N，否则要退化为使用其他算法。
第一阶段通过(N-1)步，让每个节点都得到1/N的完整数据块。每一步的通信耗时是α+S/(NB)，计算耗时是(S/N) * C。这一阶段也可视为scatter-reduce。
第二阶段通过(N-1)步，让所有节点的每个1/N数据块都变得完整。每一步的通信耗时也是α+S/(NB)，没有计算。这一阶段也可视为allgather。

整体耗时大概是2 (N-1)_ [α+S/(NB)] + (N-1) [(S/N) C]

Alltoall 与 Allgather

AlltoAll 操作允许每个进程与所有其他进程交换数据，但与 Allgather 不同的是，AlltoAll 通常涉及到每个进程向每个其他进程发送和接收不同的数据。这个操作可以看作是一种更通用的数据交换机制，它允许更灵活的数据重新分配。

当需要重新分配数据以平衡负载或改变数据的分布时, 或者在需要每个进程都有来自所有其他进程的部分数据时，可以采用 AlltoAll。

如上图所示，Allgather 操作后，每个进程拥有完整的数据集；AlltoAll 操作后，每个进程可能拥有不同的数据集，取决于发送和接收的数据。

2.3 NCCL的出现

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

这是因为一方面因为MPI基本没有考虑过GPU系统架构，在 MPI 设计中各个工作节点基本视为等同，并没有考虑节点间latency和带宽的不同。

而在 GPU 中，P40, V100， A100 的性能相差很大，不同设备存在异构性，需要重新设计策略来考虑异构场景下的硬件性能。

而 Nccl 就是为了更贴合 GPU 硬件。NCCL1.x只能在单机内部进行通信，这时多机间的通信还需依赖 MPI，从NCCL2.0开始支持多机间通信。

2.4 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 提供了容错机制，确保了计算任务在出现硬件故障时仍能继续执行。

3. NCCL 的原理

NCCL是专为NVIDIA GPU设计的集合通信库，它和 MPI 一样支持多种高效的集体通信操作，如广播、归约、全收集等。在GPU之间的通信可以通过以下几种方式实现：

GPU Shared Memory

在这种方式中，GPU之间的数据传输需要经过CPU的主机内存（Host Memory）。这意味着数据首先需要从一块GPU拷贝到CPU的主机内存，然后再从CPU的主机内存拷贝到另一块GPU。这种方式涉及到CPU和PCIe总线，可能会导致较高的通信延迟和额外的性能开销。

GPU Direct P2P

GPU Direct Peer-to-Peer（P2P）技术允许同一节点上的GPU直接相互通信，无需通过CPU的主机内存。这种直接点对点（P2P）通信减少了数据传输的延迟，并且可以显著提高多GPU之间的通信效率。GPU Direct P2P依赖于PCIe架构，允许GPU共享内存资源，从而实现高效的数据交换。

NVLink

NVLink是NVIDIA开发的一种高速互连技术，它提供了比传统PCIe更高的带宽和更低的延迟。通过NVLink，GPU之间的数据传输不再通过PCIe总线，而是直接通过NVLink连接。NVLink通过NVSwitch设备实现多GPU之间的全互联，这对于高性能计算和深度学习应用中的大规模并行处理尤为重要。

3.1 Nccl 基本架构

NCCL的架构与MPI基本一致，它将每个进程称为一个“rank”，每个rank都有一个唯一的标识符，即rank ID。这些rank的集合构成了一个“communicator”，它定义了一组可以相互通信的进程。

在NCCL中，一个进程可以属于多个communicators，并且在不同的communicators中可能有不同的rank ID。为了进行通信，每个设备上都需要创建一个NCCL Communicator对象。

3.1.1 初始化Communicator

NCCL通过以下函数来初始化communicator：

ncclCommInitRank()：初始化指定rank的communicator。
ncclCommInitAll()：同时初始化所有rank的communicator。

在创建communicator之前，root rank（通常是rank 0）需要使用ncclGetUniqueId()生成一个唯一的ID，然后将这个ID广播给所有参与通信的进程。这个ID相当于一个标识符，它确保每个进程都能够识别自己是属于某个communicator的一部分，并开始进行集体通信。

3.1.2 NCCL初始化流程梳理

initTransportsRank函数是NCCL初始化过程中的一个关键步骤，它在ncclCommInitRank函数中被调用。这个函数执行以下关键任务：

检测设备和拓扑结构：initTransportsRank首先检测系统中可用的GPU设备以及这些设备之间的拓扑结构，包括它们是否支持直接点对点（P2P）通信。
计算通信结构：基于检测到的拓扑结构，initTransportsRank计算出最佳的通信路径，这可能包括环状（RING）、树状（TREE）或集合网络（COLLNET）等不同的通信结构。
建立设备连接：根据计算出的通信结构，initTransportsRank建立设备之间的连接。这可能涉及到设置跨GPU的P2P连接、共享内存连接，或者在跨主机通信时建立网络连接。

最终经过ncclTopoFillGpu函数可以建立包含GPU的XML树结构。NCCL中可通过设置环境变量NCCL_TOPO_DUMP_FILE来书输出XML文件，并通过该XML文件来查看机器的拓扑结构。

当然你也可以通过 Nvidia-smi topo -m 来查看单机内的拓扑结构。

下面是 V100 的拓扑网络：

从上面也可以看出，GPU 通信并非仅构造硬件或软件就完事了，它也属于软硬一体优化的。

3.2 GPU拓扑结构

3.2.1 机内互联

上图为单机 8 卡 A100 的拓扑构图，从图中我们可以获得以下信息：

节点上 2 个 CPU NUMA Socket，之间通过 UPI 互联。
每个 CPU Socket 作为 PCIe Root Complex 连接 PCIe Switch 进一步互联外部设备，包括 A100 GPU、高速网卡和 NVMe 存储等。
8 个 A100 GPU 卡通过6 个 NvSwitch 实现全互联的操作。

通过 Nvidia-smi topo -m 可以看到GPU 卡的拓扑结构完全使用 Nv12 进行全互联。

从上面的拓扑结构信息，我们可以看到。虽然近几年 GPU 发展突飞猛进，但GPU 外设依然是以Pcle连接到以 CPU 为中心的互联体系中的。

PCIe是由 Intel 主导的，PCIe Gen 5 x16 双向带宽也才 128GB/s，已经远远落后于 NVLink。

3.2.2 机间互联

在 DGX A100 系列，每个节点上 8 张 GPU 通过 NVLink 和 NVSwitch 互联，机间直接用 200G IB HDR 网络互联。

随着 NVIDIA DGX H100 的推出，NVIDIA 将机内的 NVSwitch 技术扩展到了交换机上，创造了名为 NVLink Switch22 的创新产品。

NVSwitch 通常指的是机内连接，而 NVLink Switch 则是指机间的交换机，类似于在 DGX A100 中替代 InfiniBand的技术。NVLink Switch 能够支持多达 256 个 GPU, 也被称为NVIDIA 的 SuperPod，相当于 32 台 DGX 服务器的规模。

在 H100 的时代，要建设千卡级别的集群，那么由 4 个 DGX H100 SuperPod 组成，之间再通过 Infiniband NDR 网络连接即可。

IB 网络最高也就 200GB/s 的带宽，而 Nvlink 可以达到 900GB/s 的带宽，实现 Nvlink完全替代 IB 可以说一直在路上。

果然，在最新发布的NVIDIA GB200 NVL72 中，英伟达引入了第五代 NVLink，可在单个 NVLink 域中连接多达 576 个 GPU，总带宽超过 1 PB/s。

GB200 构建一个千卡的集群，那么只需要 2 个GB200 NVL72 SuperPod 组成即可。

3.2.3 拓扑优化

NCCL 2.12 中引入的PXN拓扑优化，称为 PCI × NVLink，可以加速一倍的机间的数据传输效率。

如上图所示，在传统的数据中心网络架构中，跨节点的GPU间通信往往需要经过多个网络交换机。例如，一台服务器上的GPU0想要与另一台服务器上的GPU3进行数据交换，数据包必须先从GPU0传输到叶交换机（L0），然后再拷贝到目标交换机 L3 共同连接的 S1 父交换机上，然后拷贝到 L3 交换机，最终拷贝到另一台 GPU3 上。

PXN 则利用节点内 GPU 之间的 NVIDIA NVSwitch 连接，首先将 GPU 0上的数据包移动到与目的地相同的轨道的本机的GPU3上，然后在不跨越轨道的情况下通过 L3 交换机将其发送到目的地机的 GPU3 上，这可以实现消息聚合和网络流量优化。

3.3 Nccl通信算法

NCCL通过智能地利用网络拓扑结构，能够为多GPU环境选择最佳的通信策略。在最新的NCCL版本中，库不仅支持传统的RING、TREE、COLLNET等通信算法，还引入了更多先进的算法，这些算法的细节可以在源码中的ncclTopoTuneModel方法中找到。

3.3.1 NCCL的Ring算法

Ring算法是NCCL中用于实现多GPU间通信的一种基础算法。它通过构建一个环形网络，使得每个GPU仅与其直接相邻的两个邻居进行数据交换。这种设计使得数据在环中传递，直到最终回到发送者。Ring算法以其简洁性和对等性质，为小规模GPU集群提供了有效的通信机制。

然而，Ring算法并不总是最优选择。在实际应用中，我们需要在通信延迟和带宽利用率之间做出权衡。例如，在执行基于Ring的AllReduce操作时，对于小规模数据传输，算法可能会表现出较高的延迟和较低的效率。

此外，GPU之间的物理连接方式，即拓扑结构，存在显著的异构性。在单个节点内部，GPU通常通过高速的NVLink（双向带宽可达300GBps）相连；而在节点之间，则可能依赖于相对较慢的InfiniBand网络（每个网络接口卡NIC的带宽为12.5-25GBps）。不同设备供应商可能采用不同的拓扑设计，这进一步增加了通信优化的复杂性。

3.3.2 NCCL的Tree算法

2019年上半年NCCL2.4提出double binary tree算法, 其主要思想是利用二叉树中大约一半节点是叶子节点的特性，通过将叶子节点变换为非叶子节点，得到两颗二叉树，每个节点在其中一颗二叉树上是叶子节点，在另一颗二叉树上是非叶子节点。

这种方法理论上是能够提供比ring算法更低的延迟（log2N < N）。

节点内通信：Ring 算法适用于节点内的高速互联环境，能够高效利用带宽和拓扑结构。
节点间通信：Tree 算法适用于节点间的网络通信环境，能够减少延迟并具有良好的扩展性。

当然 Nccl 是实现了自动化的选择机制，其大致的原理为预估算法不同拓扑下的通信时间：

`time = latency + nBytes / algo_bw
`

algo_bw 为算法的带宽。

我们可以通过下载 nccl-test 来进行不同算法的性能测试：

`git clone git@github.com:NVIDIA/nccl-tests.git

NCCL_ALGO=Tree ./build/all_reduce_perf -b 1M -e 2048M -f 2 -g 8

Avg bus bandwidth : 114.143

NCCL_ALGO=Ring ./build/all_reduce_perf -b 1M -e 2048M -f 2 -g 8

Avg bus bandwidth : 141.226

经过测试发现在机内Ring算法确实带宽相比 Tree 算法要高一点的。

3.3.3 NCCL的CollNet算法

NCCL 2.6 引入了一种创新的通信算法——CollNet算法，它是建立在SHArP（Scalable Hierarchical Aggregation and Reduction Protocol）基础之上的，专为与InfiniBand（IB）网络配合使用而设计。

SHArP，也被称为NCCL Plugin或NCCL-RDMA-SHARP插件，是提升通信性能的关键工具，它通过优化数据在网络中的传输方式，显著提高了大规模GPU集群的通信效率。

CollNet算法的核心优势在于其将计算任务卸载到网络交换机的能力。这种技术允许交换机直接参与到数据的聚合（Reduce）操作中，从而减少了GPU之间的直接通信需求。这种方法不仅减轻了GPU的计算负担，还优化了网络资源的使用，提高了整体的通信性能。

在传统的Ring或Tree Allreduce算法中，数据在reduce scatter阶段需要在所有参与的GPU（rank）之间传递，随后在allgather阶段再次进行传递。这个过程涉及到多次的数据传输，增加了通信延迟和复杂性。

相比之下，SHArP算法通过将数据的规约操作卸载到IB交换机，每个节点仅需要发送一次数据，然后由交换机完成数据的聚合。之后，每个节点再接收一次完整的聚合结果，大大简化了通信流程。

4.总结

在数据中心和超级计算的领域，网络互联技术的每一次飞跃，都像是在数据宇宙中架设了一座座桥梁，连接着计算的孤岛，让信息的河流得以自由流动。

从施乐帕洛阿尔托研究中心的以太网到InfiniBand联盟的高速网络，再到NVIDIA的NVLink和NVLink Network，每一次技术的演进都标志着对更快、更高效通信方式的追求。

二十五年前，受限于 PCI 总线带宽瓶颈，由微软、Intel、IBM、Sun、Compaq、HP 等巨头公司组成的 InfiniBand 联盟 IBTA 提出 Infiniband 互联架构，试图取代 PCI、以太网等成为统一的互联方案。InfiniBand 最突出的一个优势，就是率先引入了 RDMA，然而随之而来的互联网泡沫，使其野心并未实现。

十年前，Intel主导的PCIe标准限制了包括GPU在内的PCIe设备间的直接数据传输，这些设备通常需要通过CPU进行通信，这限制了它们的通信带宽。为了突破这一限制，NVIDIA的CEO黄仁勋推出了NVLink技术，允许GPU之间直接互联，绕过CPU，显著提高了数据传输速度。

如今，随着生成式AI的兴起，我们再次站在了一个新的十字路口。Meta的LLama3模型，通过同时采用RoCEv2和IB网络技术，以期拓展自己网络通信领域技术。在国内，例如华为和阿里也在不断探索Ethernet的潜力，以期在未来的网络世界中占据一席之地。

尽管 NVIDIA 现在还卖着 NVLink + Infiniband，NVIDIA并未停止前进的脚步。Bill Dally在HotI会议上的发言，预示着NVIDIA正在探索构建下一代以NVLink为核心的互联技术，这不仅是技术的迭代，更是对未来的一次抉择。

作者：Pulsar planet
来源：Tim在路上

推荐阅读

欢迎大家点赞留言，更多Arm技术文章动态请关注极术社区嵌入式AI专栏欢迎添加极术小姐姐微信（id:aijishu20)加入技术交流群，请备注研究方向。

一文讲清 NCCL 集合通信原理与优化