通信库入门：NCCL、MPI、Gloo 与通信原语

各位同学，今天我们来聊聊多卡训练里最绕不开的话题——通信库。

说实话，我刚入行那会儿，觉得多卡训练就是把模型扔到几张卡上，然后跑就完了。结果第一次跑 8 卡，发现速度还不如单卡快。嗯，那时候我才意识到，通信才是真正的瓶颈。

这一节，我带你从底层把 NCCL、MPI、Gloo 这三个主流通信库摸清楚，再搞懂 AllReduce、AllGather 这些通信原语到底在干什么。

1. 为什么需要通信库？

多卡训练的本质，就是把一个大任务拆成小任务，分给多张 GPU 并行计算。但计算完了，结果得汇总吧？梯度得同步吧？

这个「汇总」和「同步」的过程，就是通信。

你想想看，如果每张卡都自己算自己的，最后模型参数各不一样，那训练还有什么意义？所以，我们需要一个高效、可靠的通信机制，把各张卡的计算结果统一起来。

通信库就是干这个的。它封装了底层硬件（比如 NVLink、InfiniBand、以太网）的细节，给你提供一套标准的 API。你只需要调用几个函数，就能完成复杂的多卡数据交换。

2. NCCL：NVIDIA 的亲儿子

NCCL，全称 NVIDIA Collective Communications Library。说白了，这是 NVIDIA 专门为自家 GPU 打造的通信库。

我个人习惯，只要是用 NVIDIA GPU，首选 NCCL。为什么？因为它对 NVLink、NVSwitch 这些硬件的优化是其他库没法比的。

我记得有一次，我在一个 8 卡 A100 的节点上做测试，用 NCCL 做 AllReduce，带宽能跑到 600 GB/s 以上。换成别的库，直接腰斩。

NCCL 的特点

• 硬件感知： 自动检测 GPU 拓扑（比如哪些卡在同一个 NVSwitch 域内），选择最优通信路径。
• 高性能： 针对 NVIDIA GPU 做了深度汇编级优化，延迟极低。
• 易用性： 提供简单的 C API，PyTorch 等框架直接集成。

3. MPI：老牌劲旅，通用性强

MPI，Message Passing Interface，是高性能计算领域的老大哥。它不挑硬件，CPU、GPU、甚至不同厂商的 GPU 都能用。

但说实话，在纯 GPU 训练场景下，MPI 的通信效率不如 NCCL。因为 MPI 最初是为 CPU 集群设计的，对 GPU 显存直通、NVLink 这些特性支持得不够好。

不过，MPI 有一个巨大的优势：跨平台。如果你的集群里既有 NVIDIA 卡，又有 AMD 卡，或者你需要在 CPU 节点之间通信，MPI 就是唯一的选择。

MPI 的常用原语

• MPI_Allreduce：所有进程的数据归约后，结果分发给所有进程。
• MPI_Allgather：每个进程的数据收集后，分发给所有进程。
• MPI_Bcast：一个进程的数据广播给所有进程。

4. Gloo：Facebook 的轻量级选择

Gloo 是 Facebook（现 Meta）开发的通信库，最初是为了配合 Caffe2 和 PyTorch 使用。

Gloo 的特点是轻量、易用。它不像 MPI 那么庞大，也不像 NCCL 那么依赖特定硬件。它支持 TCP、InfiniBand 等多种后端，适合快速原型开发。

不过，Gloo 在大规模集群上的性能一般。我个人的经验是，单机 4 卡以内，Gloo 和 NCCL 差距不大；但超过 8 卡，Gloo 的带宽就开始明显下降。

5. 通信原语：AllReduce、AllGather 等

通信原语，就是通信库提供的基本操作。就像盖房子用的砖块，你不需要自己烧砖，直接用就行。

5.1 AllReduce

这是分布式训练里最常用的原语。它的作用是把所有 GPU 上的数据（比如梯度）进行某种归约操作（通常是求和），然后把结果分发给所有 GPU。

你想想看，每张卡算完梯度后，梯度是分散的。AllReduce 就是把这些梯度加起来，然后让每张卡都拿到这个总和。这样，每张卡就能用相同的梯度去更新模型参数。

常见的 AllReduce 算法有：

• Ring AllReduce： 把 GPU 连成一个环，数据在环上依次传递。带宽利用率高，但延迟随 GPU 数量线性增加。
• Tree AllReduce： 用树形结构做归约和广播。延迟低，但带宽利用率不如 Ring。
• NVLink 直连： 利用 NVLink 的高带宽，做全对全通信。延迟极低，但只适用于同一节点内的 GPU。

5.2 AllGather

AllGather 的作用是把每个 GPU 上的数据收集起来，然后分发给所有 GPU。和 AllReduce 的区别是，AllGather 不做归约操作，只是简单地把数据拼在一起。

举个例子：假设你有 4 张卡，每张卡上有一个 1MB 的张量。AllGather 之后，每张卡都会拿到一个 4MB 的张量，里面包含了所有卡的数据。

AllGather 在模型并行、数据并行混合的场景中很常见。比如，当每张卡负责模型的不同部分时，需要把各部分的输出收集起来，才能计算下一层。

5.3 ReduceScatter

ReduceScatter 是 AllReduce 的变体。它先做归约，然后把结果分散到各张卡上。也就是说，每张卡只拿到归约结果的一部分。

这个原语在 ZeRO 优化器中用得很多。ZeRO 把模型状态（优化器状态、梯度、参数）分散到各张卡上，ReduceScatter 正好用来做梯度的归约和分散。

5.4 Broadcast

Broadcast 是最简单的原语：一个进程把数据发给所有其他进程。在训练开始时，常用 Broadcast 把初始模型参数从 rank 0 分发给所有卡。

6. 知识体系总览

下面这张图，我帮你把本章的核心知识点串起来。你可以看到，通信库和通信原语之间的关系，以及它们在实际训练中的位置。

7. 实战：如何选择通信库？

说了这么多，到底该用哪个？我总结了一个简单的选择表：

好了，通信库的基础知识就讲到这里。下一节，我们会深入 NCCL 的调优参数，看看怎么通过环境变量把通信性能压榨到极致。