我们在看AIGC大模型、算力集群相关文献的时候,经常会看到“All-Reduce”、“All-to-All”这样的词组。
大家知道它们是什么意思吗?
故事还是要从英伟达说起。
2014年,为了取代传统的PCIe协议,英伟达推出了全新的NVLINK技术,用于GPU和GPU之间的点对点高速互联。
后来,NVLINK技术不断迭代。2018年,为了实现8颗GPU之间的互连,英伟达又发布了NVSwitch 1.0。
在捣鼓GPU互连技术的过程中,英伟达还发明了一个名叫NCCL(NVIDIA Collective Communications Library)的集合通信库。
NCCL支持多种编程语言和网络,在算法层面提供了丰富的外部API,可以大幅提升通信网络性能,也可以让GPU之间的通信设计变得更简单。
在NCCL中,定义了两种通信模式:
1、点对点通信(Point to point communication,P2P)
2、集合通信(Collective communication,CC)
点对点大家一看就懂,就是两个点之间进行通信。一个是Sender,一个是Receiver。
什么是集合通信呢?是指一组(多个)节点内进行通信。在我们传统通信里,就是点到多点,多点到多点,涉及到组网(网状、星状、环状、mesh等)那种。
NCCL还定义了一些计算节点之间数据交换的基本操作模式,并将其命名为——“通信原语(也有写作“通信元语”)。
这些通信原语包括:Broadcast、Scatter、Gather、All-Gather、Reduce、All-Reduce、Reduce-Scatter、All-to-All等。
没错,All-Reduce和All-to-All,就是其中之二。
这些通信原语是构建复杂通信行为的“原子操作”。现在所有复杂的AI算力集群,内部通信都是基于这些通信原语。它们极大地提升了并行计算的效率和便利性。
接下来,我们就逐个解释一下,这些通信原语的意思。
▉ Broadcast(1对多的广播)
这个最简单。当主节点执行Broadcast操作时,数据会从主节点发送至其他所有节点。
Broadcast是一个典型的分发、散播行为。在分布式机器学习中,Broadcast常用于网络参数的初始化。
▉ Scatter(1对多的发散)
Scatter也是一种分发、散播行为。它也是将主节点的数据发送至其他所有节点。只不过,Broadcast发送的是完整数据,而Scatter是将数据进行切割后,再分发,就像分生日蛋糕。
▉ Gather(多对1的收集)
Gather,是将多个sender(发送节点)上的数据收集到单个节点上,可以理解为反向的Scatter。
▉ All-Gather(多对多的收集)
Gather是多个到一个,All-Gather是多个到多个。
All-Gather是将多个sender(发送节点)上的数据收集到多个节点上。它相当于多个Gather操作。或者说,是一个Gather操作之后,跟着一个Broadcast操作。
▉ Reduce(多对1的规约)
Reduce的英文意思是“减少、降低”。在集合通信里,它表示“规约”运算,是一系列简单运算操作(包括:SUM、MIN、MAX、PROD、LOR等)的统称。
经常用Excel表格的童鞋,对这些简单运算应该不陌生。例如SUM,就是求和。MIN,就是找出最小值。
其实说白了,Reduce就是:输入多个数,执行操作后,得到更少的数(例如1个数)。
下面这个,就是以ReduceSum(求和规约)为例:
▉ All-Reduce(多对多的规约)
All-Reduce,这个是我们在文章开头提到的,AI领域非常常见的一个词组。
在大模型训练中,经常会用到数据并行(DP)这个并行方式。里面就有AIl Reduce这个关键操作。
我们以All Reduce Sum(求和)为例:
首先,对所有节点进行数据收集。然后,对数据进行求和。再然后,把结果重新发回给所有节点。
在大模型训练中,Server GPU节点收集的数据,就是各个Worker GPU节点计算得出的“梯度”。求和之后再发回的过程,是“更新梯度”。看不懂没关系,以后小枣君会再介绍。
▉ Reduce-Scatter(组合的规约与发散)
Reduce-Scatter稍微有点复杂、烧脑。
它是先归约(Reduce),再分散(Scatter)。具体来说:
首先,在所有参与计算的GPU节点上,对位于相同位置或索引的数据块执行指定的规约运算(例如求和SUM)。
接着,将规约后的完整结果按维度切分,并将不同的数据块分发给各个节点。最终,每个节点只得到整个规约结果的一部分,而不是全部。
简单来说,它先对所有数据进行“汇总计算”,然后再将计算好的结果“分散下发”。
▉ All-to-All(多对多的全互连)
AIl-to-AII也是AI领域出现频率很高的一个词组。它是全交换操作,可以让每个节点都获取其他节点的值。
在使用All-to-All时,每一个节点都会向任意一个节点发送消息,每一个节点也都会接收到任意一个节点的消息。每个节点的接收缓冲区和发送缓冲区都是一个分为若干个数据块的数组。
All-to-All的具体操作是:将节点i的发送缓冲区中的第j块数据发送给节点j。节点j将接收到的来自节点i的数据块,放在自身接收缓冲区的第i块位置。
All-to-All与All-Gather相比较,区别在于:All-Gather操作中,不同节点向某一节点收集到的数据是完全相同的。而在All-to-All中,不同的节点向某一节点收集到的数据是不同的。在每个节点的发送缓冲区中,为每个节点都单独准备了一块数据。
上面这个图,大家如果学过工程数学的话,就会发现,它就是一个矩阵倒置。或者说,是Excel里的行列倒转。
All-to-All的核心目标是重分布。它不进行聚合运算,而是专注于在不同节点间重新分布数据块。
以后小枣君会给大家介绍,All-to-All操作在大模型训练中的混合并行策略里至关重要。例如,当需要从数据并行组切换到模型并行组时,All-to-All可以高效地重组数据。
▉ Ring-base collective(基于环的集合)
最后还要提一个有趣的结构——环(Ring)。
Ring-base collective是将所有的通信节点通过首位相连形成一个单向环,数据在环上依次传输。
传输方式有两种,一种是一次性传输全部,还有一种,是对数据进行切割,然后分别发送。
All-Reduce里有一种Ring All-Reduce(环形全规约)算法。它是通过组合Reduce-Scatter和All-Gather两个操作来实现的。
Ring All-Reduce算法分为两个阶段:
第一阶段,将N个worker分布在一个环上,并且把每个worker的数据分成N份。
对于第k个worker,这个worker会把第k份数据发给下一个worker,同时从前一个worker收到第k-1份数据。
然后,第k个worker会把收到的第k-1份数据和自己的第k-1份数据整合,再将整合的数据发送给下一个worker。
以此循环N次之后,每一个worker都会包含最终整合结果的一份。
第二阶段,每个worker将整合好的部分发送给下一个worker。worker在收到数据之后,更新自身数据对应的部分即可。
很显然,这种环形算法可以解决传统All-Reduce中Server节点的能力瓶颈问题。
▉ 最后的话
好啦,以上就是常见通信原语的具体工作原理。
AI大模型训练推理任务,是由海量的GPU共同完成的。而这些GPU之间的通信,就是基于上面这些通信原语模型。
下一期,小枣君会详细介绍一下大模型训推任务中的并行计算方式,以及这些通信原语究竟是如何运用于不同的并行计算方式中。
敬请关注!