大模型训练常见通信源语解释
背景:最近在对ctr场景的MMOE结构做分布式训练过程中发现对很多通信源语不是很了解,查阅了一些资料将目前常用的一些通信源语进行了总结,欢迎大家批评指正。
目前大模型分布式训练里面主要涉及到的通信源于有Broadcast、Scatter、Gather、Reduce,这四种操作实现的功能通常是一对多或多对一的数据通信操作。实现多对多通信操作,则得通过AllGather、AllReduce、ReduceScatter、All-To-All。

Scatter 操作:一种一对多操作,主要是将主节点数据分成若干份发送到不同的 work 节点,数据有切分与 gather 是逆向操作。主要应用场景是在对进行模型进行分布式并行训练过程中,需要将模型拆分成若干份并放到不同的卡上,该操作可通过 scatter 实现。
Gather 操作:一种多对一操作,主要是将若干个节点的数据汇聚到主节点,与 scatter 是逆向操作。主要应用场景是在 ReduceScatter 组合里的 Scatter 操作,这个在后面介绍 ReduceScatter 时候介绍。
Reduce 操作:一种多对一操作,主要是将若干个节点的数据规约运算到一个主节点上,常用的规约操作符有:求累加和 SUM、求累乘积 PROD、求最大值 MAX、求最小值 MIN、逻辑与 LAND、按位与 BAND、逻辑或 LOR、按位或 BOR、逻辑异或 LXOR、按位异或 BOXR、求最大值和最小大的位置 MAXLOC、求最小值和最小值的位置 MINLOC 等,这些规约运算也需要加速卡支持对应的算子才能生效(图中是求和)。
Broadcast 操作:一种一对多操作,实现一对多的传输,将 root 节点的数据广播到其他 rank。应用场景一个是数据并行的参数初始化,确保每张卡上的初始参数是一致的。另一个场景是在 allReduce 实现里的会先通过 broadcast 操作后再进行 reduce。还有一个应用场景是在分布式训练 parameter server 参数服务器结构里的 master 节点 broadcast 数据到 worker 节点,再从 worker 节点 reduce 数据回 master 节点里的 broadcast 操作。

AllReduce 操作:多对多的通信原语,在集群中的每一个节点上都执行 Reduce 原语(也可以理解为单节点的 reduce + broadcast,也等价于 reducescatter + allgather)。
ReduceScatter 操作:多对多的通信原语,将数据在分片基础上进行规约,可以理解为 reduce + scatter。
AllGather 操作:Allgather 操作是将多个设备上的数据进行聚合然后同步到多个设备上,该操作可通过 Gather+Broadcast 操作实现,该操作与 ReduceScatter 互为反向操作。该操作可以被用在模型并行中,这是因为模型并行里前向计算过程中需要对所有参数全同步,需要用 allgather 把模型并行里将切分到不同的卡上的参数全同步到一张卡上才能进行前向计算。该操作相当于单节点的 gather+broadcast。
All-To-All 操作:多对多的通信原语,对分布式数据的全收集,但是不同节点收集的维度不同,相当于进行了转置。
评论