3、通信原语:AllReduce、AllGather、ReduceScatter、Broadcast的原理与适用场景

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

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

3.1 Broadcast:一卡发,多卡收

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

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

核心特点:单源广播,数据流向是单向的。

3.2 Reduce:多卡数据汇总到一张卡

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

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

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

3.3 AllReduce:多卡数据汇总,每卡都有完整结果

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

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

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

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

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

3.4 ReduceScatter:先Reduce,再Scatter

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

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

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

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

3.5 AllGather:每卡收集所有卡的数据

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

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

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

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

3.6 它们之间的关系:一张图说清楚

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

通信原语关系图 Broadcast Reduce AllReduce ReduceScatter AllGather 单源 汇总 组合 = Ring AllReduce 注:虚线表示组合关系,实线表示派生关系

3.7 实际选型建议

好了,原理讲完了。那实际项目中怎么选?我直接给结论:

场景 推荐原语 原因
梯度同步(小模型) AllReduce 简单直接,NCCL优化好
梯度同步(大模型) ReduceScatter + AllGather 通信量减半,带宽利用率高
参数广播 Broadcast 一卡发多卡收,效率最高
全局统计量 AllGather 需要每卡持有完整数据
模型并行中的分片 ReduceScatter 只拿自己需要的部分
一句话总结:AllReduce最通用,但数据量大时用ReduceScatter+AllGather组合更优。Broadcast和AllGather用于数据分发,Reduce和ReduceScatter用于数据汇总。

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

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

想要更多资料,联系可可