本文作者木木大佬,系美团资深研发专家,发表于网络那些事儿
1 概述
针对分布式训练服务器集群进行架构涉及的主流架构,目前主流的主要分为两种分布式架构:
参数服务器(ParameterServer),简称PS
去中心化架构(Decentralized Network)
分布式训练通常在计算集群上进行,集群中每个节点执行一部分计算。想当然不同节点则会有数据共享和依赖,需要在不同节点上传输数据,这就产生了通信的概念,分布式通讯按照通信方式又可以分为
集合通信方式(Collective commutication):一组节点内进行通信
点对点通信方式(Point to point comunication, P2P):在两个节点之间进行通信
深度学习训练过程中因为要涉及大量的网络模型权重参数和训练过程中产生的临时变量等数据,主要都是采用集合通信,所以我们可以理解主要是采用在PS架构下的集合通讯模式;而在大模型的分布式训练中,因为减少跟单点参数服务器统一更新,更多直接采用纯集合通讯模式。
在深度学习框架中,分布式训练的通讯原语和通讯实现方式对AI框架分布式训练起着非常重要的作用,如果想要训练大模型(Foundation Model)肯定离不开进行通讯。
2 通信原语
集合通讯中包含多个sender和多个receiver,一般的通信原语包括broadcast、gather、all-gather、scatter、reduce、all-reduce、reduce-scatter、all-to-all等通信操作进行数据传输。
假设我们现在有一台集成了4张训练卡的服务器:
2.1 Broadcast
在集合通信中,如果某个节点想把自身的数据发送到集群中的其他节点,那么就可以使用广播Broadcast的操作。
如图所示,圆圈表示分布式系统中的独立节点,一共4个节点,小方块则代表了数据,颜色相同表示数据一样。Broadcast代表广播行为,执行Broadcast时,数据从主节点0广播至其他各个指定的节点(0~3)。
Broadcast是数据的1对多的同步,它将一张XPU卡上的数据同步到其他所有的XPU卡上,其应用场景有:
1)数据并行的参数初始化,确保每张卡上的初始参数是一致的;
2)allReduce里的 broadcast + reduce组合里的broadcast操作;
3)分布式训练parameter server 参数服务器结构里的 master节点 broadcast 数据到worker节点,再从worker节点reduce数据回master节点里的broadcast操作;
2.2 Scatter
同Broadcast一样,Scatter也是一个1对多的通信原语,也是一个数据发送者,多个数据接收者,可以在集群内把一个节点自身的数据发散到其他节点上。与Broadcast不同的是Broadcast把主节点0的数据发送给所有节点,而Scatter则是将数据的进行切片再分发给集群内所有的节点,如下图所示,不相同的颜色的小方块代表不相同的数据,主节点 0 将数据分为四份分发到了节点0-3。
2.3 Reduce
Reduce称为规约运算,是一系列简单运算操作的统称,细分可以包括:SUM、MIN、MAX、PROD、LOR等类型的规约操作。Reduce意为减少/精简,因为其操作在每个节点上获取一个输入元素数组,通过执行操作后,将得到精简的更少的元素。下面以Reduce sum为例子。
在NCCL中的Reduce,从多个sender那里接收数据,最终combine到一个节点上。
2.4 All Reduce
Reduce是一系列简单运算操作的统称,All Reduce则是在所有的节点上都应用同样的Reduce操作。以All Reduce Sum为例。
All Reduce操作可通过单节点上Reduce + Broadcast操作完成。在NCCL中的All Reduce中,则是从多个sender那里接收数据,最终合并和分发到每一个节点上。
2.5 Gather
All Reduce操作可通过单节点上Reduce + Broadcast操作完成。在NCCL中的All Reduce中,则是从多个sender那里接收数据,最终合并和分发到每一个节点上。
2.6 All Gather
很多时候发送多个元素到多个节点也很有用,即在多对多通信模式的场景。这个时候就需要 All Gather操作。
对于分发在所有节点上的一组数据来说,All Gather会收集所有数据到所有节点上。从最基础的角度来看,All Gather相当于一个Gather操作之后跟着一个Broadcast操作。
2.7 Reduce Scatter
Reduce Scatter操作会将个节点的输入先进行求和,然后在第0维度按卡数切分,将数据分发到对应的卡上。例如上图所示,每卡的输入均为4x1的Tensor。Reduce Scatter先对输入求和得到[0, 4, 8, 12]的Tensor,然后进行分发,每卡获得1x1大小的Tensor。例如卡0对应的输出结果为[[0.0]],卡1对应的输出结果为[[4.0]]。
2.8 All to All
All to All作为全交换操作,通过All to All通信,可以让每个节点都获取其他节点的值。
在使用 All to All 时,每一个节点都会向任意一个节点发送消息,每一个节点也都会接收到任意一个节点的消息。每个节点的接收缓冲区和发送缓冲区都是一个分为若干个数据块的数组。All to All 的具体操作是:将节点i的发送缓冲区中的第j块数据发送给节点j,节点j将接收到的来自节点i的数据块放在自身接收缓冲区的第i块位置。
All to All 与 All Gather 相比较,区别在于:All Gather 操作中,不同节点向某一节点收集到的数据是完全相同的,而在 All to All 中,不同的节点向某一节点收集到的数据是不同的。在每个节点的发送缓冲区中,为每个节点都单独准备了一块数据。
领取专属 10元无门槛券
私享最新 技术干货