集合通信原语：AllReduce、AllGather、ReduceScatter、Broadcast的原理与开销模型

多卡训练，说白了就是让一堆GPU一起干活。但有个问题——它们得互相传数据。怎么传？传多少？快不快？这就引出了我们今天要聊的集合通信原语。

我刚开始做分布式训练那会儿，总觉得这些通信原语差不多，随便调一个就行。结果呢？训练速度慢得像蜗牛爬。后来踩了坑才明白，选错通信原语，性能差距能到好几倍。

Broadcast是最直观的通信模式。一个节点把数据发给所有其他节点。比如参数初始化时，主卡算好了初始权重，得告诉所有卡。

它的开销模型很简单：

你想想看，如果每个节点都单独发一次，那效率太低了。实际实现用的是树形结构或环状结构。我见过有人用最朴素的点对点通信实现Broadcast，结果8卡训练比单卡还慢——这就是没理解底层原理。

核心要点：Broadcast适合小数据量的场景。参数初始化、学习率广播这些，用Broadcast没问题。但大数据量就别用了。

AllReduce可以说是分布式训练的基石。每个节点都有数据，最终每个节点都要拿到所有数据的聚合结果。说白了就是「各算各的，最后汇总，人人有份」。

我项目中遇到最多的是梯度同步场景。每个GPU算完梯度后，用AllReduce把梯度加起来，然后每个GPU都拿到完整的梯度，再各自更新参数。

AllReduce的实现方式有好几种：

开销模型上，Ring AllReduce的通信量跟节点数无关，只跟数据量有关。这一点很关键。所以在大规模集群上，Ring AllReduce是主流选择。

我的经验：8卡以内，Tree AllReduce可能更快。超过8卡，Ring AllReduce优势明显。我曾经在32卡集群上测试过，Ring比Tree快了将近30%。

AllGather跟AllReduce有点像，但区别在于：AllReduce做聚合运算（求和、求平均），AllGather只是单纯收集数据，不做运算。

举个例子：每个GPU有1GB的特征数据，你想让所有GPU都拥有全部8GB的数据。这时候用AllGather。

它的开销模型：

嗯，这里要注意：AllGather的通信量比AllReduce小。因为AllReduce既要通信又要计算，而AllGather只通信不计算。但AllGather的结果是数据量变大了，内存占用会翻倍。

ReduceScatter是AllReduce的「前一半」。它把数据分散到各个节点上做归约，每个节点只保留一部分结果。

我刚开始学的时候，总觉得ReduceScatter和AllReduce差不多。后来才明白，ReduceScatter是Ring AllReduce的第一步。先ReduceScatter把数据分散归约，再用AllGather把结果收集回来，合起来就是完整的AllReduce。

开销模型：

避坑指南：我曾经在模型并行场景中误用了AllReduce代替ReduceScatter。结果每个节点都存了完整的梯度，内存直接爆了。后来换成ReduceScatter，内存占用降了8倍。记住：模型并行时，用ReduceScatter更省内存。

说白了，选哪个原语取决于你的场景：

我个人习惯是：先看数据量大小。数据量小（< 1MB），Broadcast或AllReduce都行。数据量大（> 100MB），优先考虑Ring AllReduce或ReduceScatter。

还有一个细节：NCCL的实现对Ring AllReduce做了深度优化。如果你用NVIDIA GPU，直接用NCCL的AllReduce就好，不用自己折腾。

我画了一张图，帮你理解这四种原语的数据流动方式：

这张图里，红色箭头表示数据流动方向。Broadcast是单源广播，AllReduce是全员聚合，AllGather是全员收集，ReduceScatter是分散归约。每种模式的数据流向和结果都不一样。

总结一下：理解这四种原语，关键是搞清楚「谁有数据」「谁要数据」「要不要做运算」。选对了原语，通信效率能提升好几倍。选错了，GPU就在那干等数据，白白浪费算力。

想要更多资料，联系可可