Bootstrap

AI的分布式通信操作reduce/gather/broadcast/scatter

分布式训练过程中设计到许多通信上的操作, 每个操作有其不同的术语并且有所区别,这里将其用简单的例子和描述总结一下,方便理解。

一、Broadcast

这个概念比较好理解,即广播,将一个设备上的数据广播到各个设备上,其他设备只进行读数据。

二、Reduce

Reduce是函数式编程的经典概念。数据规约涉及通过函数将一组数字缩减为一个较小的集合。例如,假设我们有一个数字列表[1,2,3,4,5]。用sum函数缩减这个数字列表将产生sum([1,2,3,4,5])= 15。同样,乘法规约将产生乘法([1, 2, 3, 4, 5]) = 120

在上图的例子中,各设备上相同位置的元素进行加和,并将结果呈现在一个设备上(图中的例子为设备1)。

Reduce例子1

Reduce例子2

三、All Reduce

相当于Reduce之后再来了一个Broadcast。

四、Gather

Gather的中文叫做收集(即把东西放到一起,并不做运算),与Reduce不同的地方是,Gather只是将数据汇总到一起,而Reduce需要“按照指定的映射函数进行运算”。

五、All Gather

即将所有人的所有数据都让所有人知道,做到所有数据透明。和Broadcast的区别是,这里是针对全局所有机器的数据来说,而Broadcast是针对一个机器,让其他所有人都知道你的数据。也就是All Gather数据传输的量更大。

六、Scatter

中文意思为离散,扩散。即将一个机器上的不同数据分别给到不同机器。而广播的含义是将一个机器上的数据全部传输给其他机器。

七、Reduce Scatter

Reduce_scatter最终呈现效果为:每个GPU上有一块完整加和后的数据。他和All reduce的区别在于,All reduce是所有完整加和的数据。

这是一个不常见的通信函数,在深度学习显存优化ZeRO中经常被用到,如下图所示。

七、相互关系与总结

All Reduce包括了All Gather的概念,即将加和(或者其他操作)后的数据再分发到各个设备上。

总结(方便大家梳理各个概念的区别):我们要记住,只有Reduce相关的操作要做计算,其余操作都不涉及计算。而Broadcast是针对单台机器对多台机器(1->N),Gather是多台机器对单台机器(N->1),All相关的操作是多台机器对多台(N->N),没有All的操作仅仅将数据汇总到一台设备。

八、参考文献

Megatron + zero_我想静静,的博客-CSDN博客

https://docs.nvidia.com/deeplea

;