📚 课程目录第3章通信原语

3. 通信原语

好，咱们进入正题。这一章聊的是通信原语，说白了就是多卡之间怎么“说话”的几种固定套路。你想想看，八张卡一起算一个模型，中间肯定要交换梯度、同步参数吧？那这些交换动作，就封装成了几个标准操作：AllReduce、AllGather、ReduceScatter、Broadcast。

我个人习惯，先把这几个家伙的关系理清楚。它们不是孤立的，而是可以互相组合的。比如ReduceScatter + AllGather，就能拼出一个AllReduce。嗯，这个后面细说。

Broadcast是最简单的原语。一张卡把数据发给其他所有卡。其他卡只收不发。

适用场景：初始化参数广播。比如你在一张卡上加载了模型权重，然后通过Broadcast把权重分发给其他卡。我在项目中遇到过，用PyTorch DDP时，其实底层就帮你做了这个，但如果你自己写通信逻辑，就得手动调一下。

核心特点：单源广播，数据流向是单向的。

Reduce和Broadcast正好相反。多张卡各自有数据，通过某种操作（比如求和、求最大值）汇总到一张卡上。其他卡的数据就没了。

适用场景：梯度汇总。但注意，Reduce只把结果放到一张卡上，其他卡拿不到。所以实际训练中，我们很少直接用Reduce，而是用AllReduce。

小提示：Reduce操作可以是SUM、MAX、MIN、PROD等。最常用的是SUM，因为梯度累加就是求和。

AllReduce是Reduce的升级版。它把多张卡的数据汇总后，再把结果分发给每一张卡。最终每张卡都持有相同的汇总结果。

说白了，就是“大家把数据交出来，我算完，再每人发一份”。

适用场景：分布式训练中的梯度同步。每张卡算完自己的梯度后，调用AllReduce把所有梯度加起来，然后每张卡都拿到总梯度，再各自更新参数。这是最最常见的用法。

我曾经踩过一个坑：AllReduce默认是求和操作，但如果你不小心传了平均值，那结果就全乱了。所以一定要确认你的操作类型。

注意：AllReduce的通信量是O(N)，其中N是每张卡的数据量。数据量越大，通信开销越高。

ReduceScatter是个混合体。它把多张卡的数据先做Reduce（求和），然后把结果分片（Scatter）到各张卡上。每张卡只拿到一部分结果。

举个例子：4张卡，每张卡有4个元素。ReduceScatter会把所有卡的第1个元素求和，结果给卡0；第2个元素求和，结果给卡1；以此类推。最终每张卡只拿到1个元素。

适用场景：当数据量很大时，ReduceScatter可以降低单卡的通信量。它常和AllGather配合，组成Ring AllReduce算法。

  核心优势：通信量从O(N)降到O(N/K)，K是卡数。数据量越大，优势越明显。

AllGather和ReduceScatter正好相反。每张卡把自己的一份数据发给所有卡，最终每张卡都持有所有卡的完整数据。

说白了，就是“每人交一份，每人拿全套”。

适用场景：需要全局信息的场景，比如Batch Normalization中的全局统计量计算。或者在做模型并行时，需要把各层的输出收集起来。

小技巧：AllGather的通信量和卡数成正比。卡越多，通信越重。所以能用ReduceScatter的地方，尽量别用AllGather。

下面这张SVG图，把四个原语的关系和组合方式画出来了。你看一眼就明白。

好了，原理讲完了。那实际项目中怎么选？我直接给结论：

  一句话总结：AllReduce最通用，但数据量大时用ReduceScatter+AllGather组合更优。Broadcast和AllGather用于数据分发，Reduce和ReduceScatter用于数据汇总。

嗯，这一章就到这。记住，选原语的时候，先想清楚你的数据流向：是汇总还是分发？是全部还是分片？想明白了，代码写起来就顺了。

避坑指南：我曾经在8卡机器上直接用AllReduce同步一个1GB的梯度，结果通信耗时比计算还长。后来换成ReduceScatter+AllGather，通信时间直接砍半。所以，数据量大的时候，别偷懒，用组合方案。