大模型训练(4):AllReduce详解

英文示意

  • Reduce(规约):规约运算,就是将多个进程中的数据按照指定的映射函数(举例:求和)进行运算得到最后的结果,并存在一个进程
  • Scatter(发散):是1对多的通信原语,是一个数据发送者,多个数据接收者,可以在集群内把一个节点自身的不同数据切片发散到其他节点上,与Broadcast不同的是Broadcast把主节点的数据发送给所有节点,而Scatter则是将数据先进行切片再分发给集群内所有的节点
  • Reduce-Scatter(规约发散):属于多对多的通信原语,具有多个数据发送者,多个数据接收者,其在集群内的所有节点上都按节点维度执行相同的Reduce规约运算,再将切片结果发散到集群内所有的节点上,Reduce-scatter由于scatter的切片数量为节点数量,等价于要执行节点数量次的reduce规约运算操作(即对不同切片分别规约),再后面执行节点数量次即切片数量个数的scatter操作,其反向操作是All-Gather
  • Gather(收集):是多对1的通信原语,具有多个数据发送者,一个数据接收者,可以在集群内把多个节点的数据收集到一个节点上
  • All-Gather(全收集):属于多对多的通信原语,具有多个数据发送者,多个数据接收者,可以在集群内把多个节点的数据收集到一个主节点上(Gather),再把这个收集到的数据分发到其他节点上(broadcast),即收集集群内所有的数据到所有的节点上

0 背景

分布式深度学习里的通信严重依赖于规则的集群通信,诸如 all-reduce, reduce-scatter, all-gather 等,因此,实现高度优化的集群通信,以及根据任务特点和通信拓扑选择合适的集群通信算法至关重要。

本文以数据并行经常使用的 all-reduce 为例来展示集群通信操作的数学性质。

1 All-Reduce的实现

image-20250112231652435

如上图所示,一共4个设备,每个设备上有一个矩阵(为简单起见,特意让每一行就一个元素),All-Reduce 操作的目的是,让每个设备上的矩阵里的每一个位置的数值都是所有设备上对应位置的数值之和。

image-20250112231741572

如上图所示, All-Reduce 可以通过 reduce-scatter 和 all-gather 这两个更基本的集群通信操作来实现。基于 ring 状通信可以高效的实现 reduce-scatter 和 all-gather,下面我们分别用示意图展示其过程。

2 Reduce-Scatter 的实现和性质

image-20250114012723283

上图为环状通信实现(reduce-scatter阶段),reduce-scatter 的结果是每个设备保存一部分reduce 完(归并完)的结果。为了方便讨论,我们先定义一些符号。

  • 假设有 N N N个设备(上面的例子中 N = 4 N=4 N=4
  • 假设整个矩阵大小是 Φ \Phi Φ,每个完整矩阵被划分为 N N N块,每一块的大小为 Φ N \frac{\Phi}{N} NΦ,那么 reduce-scatter 后,每个设备上完成归并数据块大小为 Φ N \frac{\Phi}{N} NΦ
  • 假设卡和卡之间通信带宽是 B B B,而且是双工通信(duplex),即每个设备出口和入口带宽可以同时达到 B B B,所有设备的入口带宽总和与所有设备的出口带宽总和均为 N B NB NB
  • 环状 reduce-scatter 一共需要 N − 1 N-1 N1步才能完成

下面介绍每一步的具体操作:

  • Step1:每个设备都负责某一块 Φ N \frac{\Phi}{N} NΦ 的数据,并向左边的设备发送这块数据,譬如在上图中
    • GPU1负责将第 2 片数据 b 1 b1 b1向GPU4发送
    • GPU2负责将第 3 片数据 c 2 c2 c2向GPU1发送
    • GPU3负责将第 4 片数据 d 3 d3 d3向GPU2发送
    • GPU4负责将第 1 片数据 a 4 a4 a4向GPU3发送
    • 每个设备收到右边设备的数据后,就把收到的数据累加到本地对应位置的数据上去(通过逐渐变深的颜色表示数据累加的次数更多)。注意,在这样的安排下,每个设备的入口带宽和出口带宽都被用上了,而且不会发生争抢带宽的事情
  • Step2:
    • GPU1负责将累加后的第 3 片数据 c 1 + c 2 c1+c2 c1+c2向GPU4发送
    • GPU2负责将累加后的第 4 片数据 d 2 + d 3 d2+d3 d2+d3向GPU4发送
    • GPU3负责将累加后的第 1 片数据 a 3 + a 4 a3+a4 a3+a4向GPU4发送
    • GPU4负责将累加后的第 2 片数据 b 1 + b 4 b1+b4 b1+b4向GPU4发送
    • 每个设备收到右边设备发过来的数据后,就把收到的数据累加到本地对应位置的数据上去(累加后颜色更深)
  • Step3:
    • GPU1负责将累加后的第 4 片数据 d 1 + d 2 + d 3 d1+d2+d3 d1+d2+d3向GPU4发送
    • GPU2负责将累加后的第 1 片数据 a 2 + a 3 + a 4 a2+a3+a4 a2+a3+a4向GPU1发送
    • GPU3负责将累加后的第 2 片数据 b 1 + b 3 + b 4 b1+b3+b4 b1+b3+b4向GPU2发送
    • GPU4负责将累加后的第 3 片数据 c 1 + c 2 + c 4 c1+c2+c4 c1+c2+c4向GPU3发送
    • 每个设备收到右边设备发过来的数据后,就把收到的数据累加到本地对应位置的数据上去(累加后颜色更深)

经过 N − 1 N-1 N1步之后,每个设备上都有了一片所有设备上对应位置数据 reduce 之后的数据。整个过程中,每个设备向外发送了 ( N − 1 ) Φ N \frac{(N-1)\Phi}{N} N(N1)Φ 大小的数据,也收到了 ( N − 1 ) Φ N \frac{(N-1)\Phi}{N} N(N1)Φ 大小的数据,因为每个设备的出口或入口带宽是 B B B,所以整个过程需要的时间是 ( N − 1 ) Φ N B \frac{(N-1)\Phi}{NB} NB(N1)Φ ,如果 N N N 足够大,完成时间近似为 Φ B \frac{\Phi}{B} BΦ ,神奇的是,这个时间和设备数 N N N 无关。当然,在所有设备间传递的数据量是 ( N − 1 ) Φ (N-1)\Phi (N1)Φ,和设备数 N N N 成正比。

让我们强调一下:基于环状通信的集群通信算法执行时间几乎和设备数无关,但总通信量和设备数成正比。

3 all-gather 的实现和性质

reduce-scatter 执行结束之后,再通过 all-gather 过程就可以实现 all-reduce,其中 all-gather 也可以通过环状通信算法来实现。

img

上图给出了环状 all-gather 的实现过程,就不详细描述了,值得注意的是,它的通信时间和通信量的分析与 reduce-scatter 一模一样:整个过程需要的时间是 ( N − 1 ) Φ N \frac{(N-1)\Phi}{N} N(N1)Φ ,如果 N N N 足够大,完成时间近似为 Φ B \frac{\Phi}{B} BΦ ,这个时间和设备数 N N N 无关,当然,在所有设备间传递的数据量是 ( N − 1 ) Φ (N-1)\Phi (N1)Φ ,和设备数 N N N 成正比。不过,请注意在 reduce-scatter 里, Φ \Phi Φ 都是完整矩阵的数据量,即 reduce-scatter 输入矩阵的数据量和 all-gather 输出矩阵的数据量。

4 通信量和冗余显存之间的关系

上文只分析了通信量,而没有分析对设备显存的占用量。第2节中的图为例,每个设备上的输入矩阵大小是 Φ \Phi Φ ,但经过 reduce-scatter 之后每个设备上只需要 Φ N \frac{\Phi}{N} NΦ 大小的显存,也就是 ( N − 1 ) Φ N \frac{(N-1)\Phi}{N} N(N1)Φ 的空间是冗余的,因为一共有 N N N 个设备,所以整个集群中有 ( N − 1 ) Φ (N-1)\Phi (N1)Φ 的显存是可以节省下来的。注意,每个设备冗余的显存恰好和每个设备的通信量一致,所有设备冗余的显存和所有设备的总通信量一致。

第3节中的图为例,每个设备上的输入矩阵大小是 Φ N \frac{\Phi}{N} NΦ ,但经过 all-gather 之后每个设备上需要的显存是 Φ \Phi Φ ,而且每个设备上的矩阵的大小和数值都完全一样,也就是经过 all-gather 之后,在设备之间产生了冗余,不同的设备存储了一些完全一样的数据。同样,每个设备冗余的显存恰好和每个设备的通信量一致,所有设备冗余的显存和所有设备的总通信量一致。

当然,冗余量和通信量之间的等价关系不是偶然的,正是因为这些通信才造成了设备之间数据的冗余。因此,当保持 Φ \Phi Φ 不变时,增大设备数 N N N (我们不妨称 N N N 为集群通信的并行宽度)时,所有设备之间的通信量是正比增长的,而且在所有设备上造成的冗余显存是正比例增长的。当然,完成一个特定的集群通信所需要的时间基本和并行宽度 N N N 无关。

因此,增加并行宽度 N N N 是一个双刃剑,一方面它让每个设备处理的数据量更小,即 Φ N \frac{\Phi}{N} NΦ ,从而让计算的时间更短,但另一方面,它会需要更多的通信带宽 ( N − 1 ) Φ (N-1)\Phi (N1)Φ ,以及更多的显存空间 ( N − 1 ) Φ (N-1)\Phi (N1)Φ

5 环状算法的最优性

我们在前面提了一个问题:你能不能想出比环状算法更好的集群算法实现?答案是,理论上不可能有更好的算法了。

我们已经分析过了要完成 reduce-scatter 和 all-gather 每个设备至少要向外发送(以及同时接收)的数据量是 ( N − 1 ) Φ N \frac{(N-1)\Phi}{N} N(N1)Φ ,无论使用什么算法,这个数据量是不可能更少了。在这个数据量下,最少需要多少时间呢?出口带宽是 B B B ,因此一张卡向外发送数据至少需要的时间是 ( N − 1 ) Φ N B \frac{(N-1)\Phi}{NB} NB(N1)Φ ,这也正是环状算法需要的时间。

当然,我们这里的通信时间只考虑传输带宽,而没有考虑每次传输都包含的延迟(latency)。当数据量V比较大时,延迟项可以忽略,上文的分析就是成立的。当 Φ \Phi Φ 特别小,或者设备数 N N N 特别大时,带宽 B B B 就变得不重要了,反而是延迟比较关键,这时更好地实现就不是环状算法了,而应该使用树状通信。

这也是为什么英伟达 NCCL 里既实现了ring all-reduce,也实现了 double-tree all-reduce 算法(https://developer.nvidia.com/blog/massively-scale-deep-learning-training-nccl-2-4/)。

高效实现一个集群通信的关键是如何充分利用设备和设备之间的带宽,基于环状(ring)通信实现的集群通信算法就是这一思想的体现。

评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值