1.1.4 通信范式
分布式机器学习的通信范式规定了数据并行模式下计算节点应以何种通信拓扑及协议同步模型参数(或梯度)。基于消息传输接口(Message Passing Interface,MPI)的AllReduce通信和参数服务器(Parameter Server,PS)通信是两种主流通信范式,此外,流言(Gossip)异步去中心化通信范式也有一些成功案例。本节将对上述三种通信范式的典型实现做简要介绍。
1.1.4.1 基于MPI的AllReduce通信范式
1.MPI集群通信原语概览
MPI[61]是一个性能优异的分布式通信库,支持各种网络设备和网络拓扑结构,在多种分布式通信技术中得到广泛使用,典型通信原语包括点对点(P2P)、广播(Broadcast)、汇集(Gather)、全汇集(AllGather)、聚合(Reduce)、全聚合(AllReduce)、分发(Scatter),如图1-13所示。
图1-13 MPI的七种常见通信原语示意图
集合通信(Collective Communication)[62],顾名思义,即一个集合内多个进程间的通信。这些原语使我们能在多进程之间更简便地交换数据。图1-13中除点对点(P2P)通信外,其他原语都属于集合通信。不同的通信原语功能不同,适用于不同的通信任务。首先是多对一和一对多原语:广播(Broadcast)原语可以从一个进程向所有进程发送数据;汇集(Gather)原语可以将所有进程的数据收集到一个进程;聚合(Reduce)原语与汇集原语相似,但会将收集的数据通过求和、最大等方式聚合;分发(Scatter)原语可以将一个进程的不同数据切片分发给不同的进程。全汇集(AllGather)和全聚合(AllReduce)是两个多对多原语,即所有进程都能获得Gather或Reduce的结果。
2.AllReduce原语及其典型实现
分布式机器学习的模型同步操作需要聚合所有计算节点的模型数据,并同步给所有计算节点,这个过程可以通过调用AllReduce原语实现。AllReduce原语的实现是多种多样的,最简单直接的一类方案是Reduce+Broadcast,如二叉树型同步(Binary Tree)[63],Reduce原语先将计算节点的数据沿二叉树聚合到根节点,Broadcast原语再将聚合结果原路广播回计算节点。另一类方案是ReduceScatter+AllGather,ReduceScatter原语在聚合所有进程数据的同时,不同进程只获得聚合结果的一个切片,随后利用AllGather原语交换切片得到完整的聚合数据。这类方案的典型算法包括蝶式同步(Recursive Halving and Doubling,或称Butterfly)、二进制块同步(Binary Blocks)和环形同步(Ring)[63]。研究表明,在上述算法中,仅Binary Blocks和Ring适用于较大数据量,其中Ring适用于小规模集群(少于32节点),Binary Blocks适用于较大规模集群(32~256节点)[63]。由于AllReduce原语的具体实现算法繁多,下面仅简要介绍近年来在工业实践中表现卓越的Ring AllReduce、Three-Phase Ring AllReduce、2D-Torus Ring AllReduce三种算法,其他算法实现请参阅文献[25]、[26]、[33]、[63]~[66]、[68]。
三种算法示意图如图1-14所示。百度首先将Ring AllReduce算法[67]引入分布式机器学习。在这种算法中,所有节点首尾相连形成通信环,每个节点有且仅有一个前继节点和一个后继节点,且只能从前继节点接收数据,向后继节点发送数据。Ring AllReduce算法整体分为ReduceScatter和AllGather两个阶段。以图1-15所示个计算节点为例,每个节点的数据被平均分割为个数据切片。在ReduceScatter阶段,第次通信时,每个节点向其后继节点发送第个数据切片,同时从前继节点接收数据切片,并与本地第个数据切片聚合。重复上述过程次,此时每个节点都拥有一个全局聚合的数据切片,ReduceScatter阶段完成。随后,AllGather阶段交换各节点的全局聚合数据切片,以获得完整的全局聚合数据。AllGather的过程与ReduceScatter相似,只是AllGather不执行聚合操作,而是简单的覆盖。再经过次通信后,所有节点都获得完整的全局聚合数据,至此AllReduce阶段完成。在上述算法中,每个节点总计发送个数据,为模型参数量,可见算法的通信开销恒定,与集群规模无关。在实际部署中,将大带宽互联的两个GPU作为邻居,可以最大限度地发挥Ring AllReduce算法的性能。然而,在较大规模集群中,模型数据会被过度切片并导致低带宽利用率,所以Ring AllReduce算法只适用于小规模集群。
图1-14 三种算法示意图
图1-15 Ring算法工作流程
Three-Phase Ring AllReduce算法由腾讯机智提出[30],如图1-14(b)所示。该算法从本质上讲就是分层的Ring AllReduce算法,由内环同步、外环同步和内环广播三个阶段组成。算法将所有节点分为多个组,组内节点通过PCIe/NVLink形成内部通信环,组内代表节点与其他组的代表节点通过GPUDirect RDMA形成外部通信环。在第一阶段,组内节点调用Ring AllReduce(或其他Reduce算法)算法将组内数据聚合到代表节点。在第二阶段,代表节点之间再次调用Ring AllReduce算法完成全局数据聚合,每个代表节点都获得完整聚合数据。在第三阶段,代表节点广播全局聚合数据给组内其他节点。该算法执行步骤少,适用于小模型、大集群的训练任务,其在1024个节点的大集群上表现优越,并取得2018年度4 min训练ImageNet+AlexNet的世界纪录。但是,在传输数据量大时,Ring AllReduce算法表现更佳。
2D-Torus Ring AllReduce算法[32]由索尼提出,如图1-14(c)所示,集群节点可视为二维网格排列。该算法由三步组成:第一步,横向节点执行ReduceScatter;第二步,纵向节点执行AllReduce;第三步,横向节点执行AllGather。该算法具有较小的通信开销,可用于更大规模的机器学习集群。索尼将该算法应用于更大规模的(如3456个节点)集群,在ImageNet+ResNet-50训练任务上取得2min的世界纪录,证明了该算法的优越性能。类似的层次化环形算法还有谷歌的2D-Mesh[33]和IBM的BlueConnect[68],它们均可以有效提高Ring AllReduce算法在大规模集群中的训练性能。
3.MPI实现库及其支持的AllReduce算法
目前业界广泛使用的MPI实现库主要为OpenMPI和MPICH,它们都采用MPI标准,提供了一系列集合通信原语的高性能实现。OpenMPI[69]的适用面广泛,可在多种互联模式的大规模集群中运行,如InfiniBand、Myrinet、Quadrics、TCP/IP等,其支持的AllReduce算法也最多,适用于生产系统。MPICH[64]实现了每一版的MPI标准,内置调试器,适用于开发系统,MPICH及其衍生产品MVAPICH和Intel MPI都提供了广泛的网络支持。另外,Gloo、NCCL、Horovod也提供了MPI集合通信算法支持。Gloo[70]是META开发的机器学习集合通信库,现已被深度学习开源框架PyTorch应用于集合通信后端,支持CPU、GPU以及GPUDirect RDMA通信。NVIDIA集合通信库(NCCL)[71]针对NVIDIA系列GPU通信进行了特别优化,在PCIe、NVLink和InfiniBand互联集群中传输带宽得到明显提升。目前,NCCL已被多数深度学习开源框架支持,包括PyTorch、TensorFlow、MXNET、Caffe等。Horovod[72]原是Uber为TensorFlow开发的分布式通信库,现也支持PyTorch和MXNET。Horovod对MPI、Gloo、NCCL进行了封装,用户可以灵活使用不同的底层通信库,实现多样化的分布式通信能力。五种MPI通信库支持的AllReduce算法见表1-2,深度学习框架支持的MPI通信库见表1-3。
表1-2 五种MPI通信库支持的AllReduce算法
注:表中数据源于2022年5月各通信库的文档和代码,支持的算法后续仍可能更新。
表1-3 深度学习框架支持的MPI通信库
注:表中数据源于2022年5月各深度学习框架的文档和代码,支持的通信库后续仍可能更新。
这些MPI通信库针对不同的用例而特别设计,所以各有优劣。除了支持的算法有所不同,OpenMPI在CPU通信和小张量上性能表现最佳。Gloo适用于CPU上的半精度(Float16)张量通信。OpenMPI和Gloo都支持CPU和GPU通信,而NCCL仅支持GPU通信。尽管如此,NCCL在NVIDIA系列GPU上拥有绝对性能优势[73]。
总的来说,MPI的优势和劣势都比较明显。优势在于提供了高性能的集合通信支持,能够有效加速分布式机器学习的数据并行。另外,MPI的几种高效实现(如Butterfly、Binary Blocks和Ring系列)采取去中心化的通信拓扑,能够缓解有中心通信拓扑的通信瓶颈,实现负载均衡。劣势在于只支持同步并行,不支持容灾恢复与动态扩展,这意味着运行时不仅不能动态增删节点,而且任意节点故障或响应变慢都将阻塞整个集群,甚至导致任务失败。因此,MPI主要用于具有同构计算与通信资源和高可靠保证的理想集群。
1.1.4.2 参数服务器通信范式
参数服务器的最早抽象来源于分布式键值存储,采用了分布式内存对象缓存系统Memcached存放共享参数,实现在分布式系统的不同计算节点之间同步采样器状态[74]。但是,第一代参数服务器缺少灵活性且性能不佳。后来,谷歌大脑提出第二代参数服务器DistBelief[75],这是一种有中心的分片参数服务器,使用多个参数服务器节点分别存储大模型的不同分片,模型分片之间相互独立且异步运行。计算机内部以多卡模型并行模式训练本地模型副本,再经由对应的参数服务器分片异步更新全局参数。DistBelief针对大规模深度模型做出了许多改进,如异步运行的Downpour SGD能够容忍不同计算节点运行速度的差异,以及在Sandblaster L-BFGS中引入调度器,协调计算节点的工作量以平衡处理时间。目前使用最多的是第三代参数服务器,由卡内基梅隆大学(CMU)的李沐等人提出[76-77]。第三代参数服务器的整体结构与第二代的类似,但提供了更为通用的设计,具有高效通信、灵活的一致性模型、弹性可扩展性、系统容错和耐用性、易于使用五大优势。
1.架构与节点功能概述
后文默认参数服务器为第三代参数服务器,其架构简图如图1-16所示。请注意,图1-16抽象于PS-LITE的工程实现[78],这是参数服务器架构的一个高效且轻量的实现库,提供了灵活且高性能的通信接口(PUSH/PULL)和服务器端用户可编程能力。参数服务器架构有三类功能节点:参数服务器、计算节点和调度器。参数服务器是核心的功能节点,用于存储和共享全局模型参数,且总是维护着最新版本的参数。由于单个参数服务器节点无法存储大模型的完整参数,所以实践中往往部署多个参数服务器节点构成参数服务器组,每个参数服务器节点分担不同的模型分片。为保障系统可靠性,应对参数服务器节点掉线的问题,参数服务器节点之间可以复制和迁移参数以实现冗余备份。
图1-16 参数服务器架构
计算节点负责利用本地训练数据计算模型参数的更新量(典型如模型梯度),并上报参数服务器以更新全局参数,所以计算节点可以仅与参数服务器通信,计算节点之间相互不建立连接。计算节点的训练数据可以全局打乱并均匀分配,也可以从分布式文件系统或分布式数据库中读取得到。对于大模型在计算节点上的存储问题,计算节点不一定需要存储完整的参数。例如,在广告点击预测等逻辑回归或线性回归应用中,输入数据特征可能是稀疏的,于是计算节点只需要从参数服务器下载与输入数据的关键特征相关的参数。另一方面,我们可以将大模型切片放置到多机或多卡上,通过模型并行解决大模型存储和计算的问题,这些模型切片可以再次被细分,由不同的参数服务器节点管理,进而实现参数服务器节点之间通信负载的均衡。
调度器主要用于协调计算节点和参数服务器节点之间通信框架的建立、销毁以及运行时的节点状态管理。在集群启动时,调度器优先启动,其他节点与调度器建立通信连接。调度器注册新节点,告知新节点其他节点的网络地址,引导参数服务器节点和计算节点之间建立通信连接,从而初始化参数服务器通信框架。在目前的PS-LITE实现中,调度器在运行时只需负责节点的活性检测和掉线恢复。每隔一定时间间隔,其他节点就会向调度器发送心跳消息报告自己的存活状态。另外,当掉线节点重连时,调度器帮助掉线节点恢复掉线前的状态以继续参与训练。在后续版本中,PS-LITE可能会支持更丰富的运行时功能。当集群训练完成时,调度器要协调各节点销毁通信套接字,随后停止各节点的进程。
2.同步、异步工作流程与流量模型
计算节点与参数服务器之间依靠PUSH和PULL两个基本操作原语实现数据通信。计算节点首先调用PULL原语从参数服务器拉取最新模型参数,随后调用PUSH原语将本地计算的模型梯度推送给参数服务器。这两个原语的实现采用经典的请求/响应协议,计算节点主动发起PUSH/PULL请求,参数服务器被动响应。注意,模型梯度和参数由多个张量构成,从用户视角来看,每个网络层的可训练参数都是一个张量,如卷积层的卷积核是一个张量,全连接层的参数矩阵也是一个张量。这些张量有各自的唯一索引值KEY,计算节点在调用PUSH/PULL时,需要指定待传输张量的KEY值,这意味着一次模型同步过程需要多次连续调用PUSH(KEY)或PULL(KEY),且相应的请求/响应相互独立、互不影响。另外,由于不同张量大小不一,在底层实际传输时,大张量还会被继续切分并均匀分发给多个参数服务器节点,以实现存储和流量的负载均衡,该问题将在后面讨论。
参数服务器架构相比MPI AllReduce的明显优势是支持同步、异步两种步调。图1-17(a)所示为整体同步并行(Bulk Synchronous Parallel,BSP)下参数服务器的流量模型。①在一次同步通信开始时,计算节点发起PULL(KEY)请求,请求从参数服务器节点下载索引值为KEY的最新参数张量。②参数服务器节点从本地键值数据库KVStore读取对应参数张量,将其附加在PULL(KEY)响应消息中发回计算节点。③计算节点用本地训练数据和下载的模型副本计算模型梯度。④计算节点发起PUSH(KEY)请求,将索引值为KEY的梯度张量发送给参数服务器节点。⑤在BSP下,参数服务器节点等待收齐所有计算节点的索引值为KEY的梯度张量,然后用聚合梯度更新全局参数。⑥参数服务器节点返回PUSH(KEY)响应消息给计算节点,表示上传的梯度张量已收到,可启动下一轮次的参数拉取。重复上述步骤直至达到指定轮数。上述步骤中多个张量的PUSH/PULL流同时存在,它们在步骤①②④⑤⑥中互不依赖,可以同时执行。
完全异步并行(Total Asynchronous Parallel,TAP)的流量模型如图1-17(b)所示,其与BSP的主要区别在步骤⑤。不同于BSP需要等待所有计算节点到位,TAP下参数服务器节点在收到某一计算节点上传的梯度张量后,即刻用于更新全局参数,而不等待其他计算节点。在此种模式下,计算节点之间完全独立并行,解决了BSP中的掉队者问题,可以有效提高计算吞吐率,加速训练进程。但是,TAP不能保证计算节点间模型参数的一致性,理论上破坏了梯度优化准则,引发延迟梯度问题,劣化算法的收敛表现。本书将在3.3节详细讨论异步更新中的延迟梯度问题及其解决方案。
图1-17 整体同步并行和完全异步并行模式下的流量模型示意图
BSP和TAP两种模式的时间线如图1-18(a)和图1-18(b)所示。BSP存在明确的同步屏障,优点是所有计算节点的步调保持一致,分布式算法的收敛性能得以保障;缺点是快节点容易被阻塞以等待慢节点(掉队者问题),会加剧资源空闲,降低系统吞吐率,最终拉低训练效率。TAP与BSP完全相反,TAP没有同步过程,计算节点即刻更新全局参数且相互不等待,因而具有不同的步调。在TAP下,模型计算和传输紧密排列,没有阻塞延迟,全局参数更新更加频繁,但代价是延迟梯度引发的算法收敛性损伤。延迟同步并行(Stale Synchronous Parallel,SSP)[79]是BSP和TAP的折中,如图1-18(c)所示,SSP允许计算节点以异步方式运行,但最快和最慢的节点之间的差距不能超过步,否则快节点将被阻塞。由于深度学习训练过程中的小偏差不一定会损害模型的准确性,有界步差可以控制较小偏差,能够对收敛性进行数学分析和证明,同时也在一定程度上缓解了掉队者问题,提高系统吞吐率和训练效率。但是,SSP仅适用于动态异构性的集群环境,即快节点后续可能成为慢节点,慢节点后续也可能变快。否则,若快节点总是最快,它最终将被持续阻塞,使SSP退化为带TAP过时梯度的BSP,取二者之劣而无所增益,得不偿失。
图1-18 BSP、TAP、SSP三种模式的时间线示意图
3.多参数服务器负载均衡
上文提到,一个完整模型由多个网络层参数张量构成,所以一种朴素的多参数服务器分配方式是轮询调度,即以轮询的方式将多个网络层参数张量分配给不同的参数服务器节点。但是,不同网络层的参数张量大小不一,轮询调度无法保证负载均衡。PS-LITE实现了一套更为细致的参数分配规则。以四层卷积和三层全连接网络为例,卷积核的参数量通常小于全连接参数矩阵。多参数服务器负载均衡示意图如图1-19所示,PS-LITE设定了一个参数量阈值,参数量小于该阈值的小张量均匀分配给各个参数服务器节点,反之,大张量被进一步平均切分成多个更小的分片,并分配给各个参数服务器节点。由于大张量的参数量占总体的大比例,这种参数分配方案总体趋于平均分配,有良好的负载均衡效果,同时避免了对小张量的过度切分。
图1-19 多参数服务器负载均衡示意图
总体来说,参数服务器的优势在于同时支持同步和异步两种更新模式,可以灵活使用BSP、TAP、SSP等多种同步算法,也能够提供容灾恢复和动态扩展等功能,这些功能使得参数服务器更适合不理想的计算机集群,如可用资源异构分布且动态变化的商业云环境,或者多任务竞争有限资源的低可靠数据中心。参数服务器架构已得到TensorFlow、PyTorch、MXNET、PaddlePaddle等主流深度学习框架的支持,可供读者简单部署使用。
1.1.4.3 Gossip异步去中心化通信范式
为了避免中心化的通信瓶颈和单点故障问题,近些年来涌现了很多去中心化通信范式。在这些通信范式下,各个计算节点在一个连通图中,通过点对点通信来交互训练信息。随着训练的推进,计算节点的模型信息会通过图的边逐步扩散到其他计算节点。这种通信范式可以将计算集群的数据流量均匀分布到各个计算节点之间的通信链路上,从而避免中心化的通信瓶颈。
根据各个计算节点的训练步调是否保持一致,去中心化通信范式也可以分为同步和异步两类。上文介绍的MPI AllReduce就是典型的同步去中心化通信范式,各个计算节点的训练步调保持完全一致,每个计算节点都需要等待其他计算节点完成通信才能进入下一轮迭代。显而易见,同步去中心化通信范式面临掉队者问题。为了在去中心化通信范式下解决掉队者问题,异步去中心化通信范式被相继提出,如Gossip范式[80-82]。
Gossip范式下各节点运行视角的示意图如图1-20所示。在Gossip范式下,各个计算节点的训练进程相互独立,计算节点在每次迭代中从一个或多个邻居节点拉取最新模型参数,然后使用本地计算的梯度和拉取的参数更新本地模型副本,随后立刻进入下一轮迭代。可以看出,这种异步去中心化通信范式允许计算节点通过完全去中心化的对等网络来协同训练机器学习模型,计算节点既用本地数据学习本地个性模型,也与其网络邻居交互模型知识。文献[80]、[83]显示,不管是在16~32个节点的小集群中,还是在112个节点的较大集群中,Gossip都比有中心算法表现得更好。
图1-20 Gossip范式下各节点运行视角的示意图(灰色节点表示该节点的视角,虚线表示有邻居关系但无数据通信,实线箭头表示有向数据通信,节点5处于异常离线状态)