BytePS源码解析(1)-理论

bad_sheep

已于 2022-01-29 21:18:40 修改

阅读量2.7k

点赞数 4

文章标签：机器学习分布式 BytePS All-reduce

于 2022-01-29 17:38:33 首次发布

本文链接：https://blog.csdn.net/Bad_Sheep/article/details/122284797

版权

最近因为工作需要，研究了一下字节跳动开源的BytePS，虽然其名字上带有"PS"两字，但是研究后发现它其实并不是一个传统的PS架构，而是all-reduce和PS的一个融合，这个在后面会详细说。网络上有一些关于BytePS创新性的一些质疑，这个见仁见智，我个人感觉撇开创新性，BytePS的理论还是比较优雅的，另外根据初步的测试，效果也确实很不错，思路和代码都很有参考价值。因此尝试写个系列文章来记录一下，本篇是第一篇，先从BytePS的论文入手，了解它的一些理论要点。（作者水平有限，如有谬误敬请指正）

BytePS的论文地址在https://www.usenix.org/system/files/osdi20-jiang.pdf，感兴趣的读者可以阅读原文以了解更多细节。

一.并行范式

1.1 All-reduce

1.2 PS

1.2.1 Non-colocate PS

一.并行范式

了解机器学习模型分布式训练的同学可能都知道，在做分布式训练时有两大类方式，一种是模型并行，一种是数据并行。简单理解模型并行就是把一个完整的模型切分成若干部分，然后放到不同的worker上去运行，每个worker运行模型的不同部分，如果切分时能够做到不同部分没有依赖，能够独立运行，则可以并行运行不同部分而达到提高效率的目的；数据并行则是另一个思路，每个worker都运行完整的模型，但是把整个训练数据切分成若干个部分，每个worker都只用一个部分的数据做训练，这样会有一个问题，每个worker使用的数据不一样，训练时同一个模型参数在不同worker上的梯度就会不一样，因此数据并行在训练时会需要一个梯度的聚合机制，不同的梯度聚合机制就会产生不同的训练方法。就目前而言，主流的有两种大的数据并行范式：PS和all-reduce。BytePS也同样属于一种数据并行方法，即它也是通过对数据进行切分来达到分布式的目的。不同在于BytePS融合了all-reduce和PS的一些特性，这个从论文的名字《A Unified Architecture for Accelerating Distributed DNN Training in Heterogeneous GPU/CPU Clusters》也可以看出来，它想做的是对PS和all-reduce进行统一。下面先简单介绍下all-reduce和PS。

1.1 All-reduce

关于all-reduce的一些概念可以参考这篇文章，写的非常好，这里就不再赘述了。以深度学习中最常用的ring all-reduce为例，其通信过程可以划分成两个阶段：ReduceScatter和AllGather，如图2，图3所示：

图1.Ring all-reduce初始状态（图片截自腾讯机智团队分享--AllReduce算法的前世今生 - 知乎）

图2.ReduceScatter （图片截自腾讯机智团队分享--AllReduce算法的前世今生 - 知乎）

图3.AllGather （图片截自腾讯机智团队分享--AllReduce算法的前世今生 - 知乎）

图1中的a，b，c，d表示4个节点（可以理解为参与训练的4个worker），需要同步的数据（例如参数梯度）在每个节点上都被切分成了4份，例如a0，a1，a2，a3是节点a上的4份数据。最终我们需要聚合得到数据如图3所示：最终每个节点上都得到了4份数据各自对应的聚合结果（这里示例是加和）。现在来分析一下从初始状态到最终结果中每个节点的通信量。

在第一个阶段ReduceScatter中，每个节点都会对下一个节点发送4份数据中的某一份，同时接收来自上一个节点发送过来的4份数据中的某一份（因为逻辑上把4个节点构成了一个环）。如图2最上端所示，第一步中a向b发送a0，b向c发送b1，c向d发送c2，d向a发送d3；第二步类似，但是每个节点都是发送不同部分的数据；第三步完成时，每个节点都有一个聚合完成的部分数据，例如a上持有了a1+b1+c1+d1，b持有了a2+b2+c2+d2。推广至一般情况，假设要传输的参数量为M，节点数为n，那完成第一阶段每个节点需要发送/接收n-1次数据，每次的发送/接收的数据量为M/n，因此在ReduceScatter阶段每个节点的发送/接收数据量为(n-1)M/n。

在第二个阶段AllGather中，通信流程与ReduceScatter类似，只是此时每个节点会把从上个节点接收到的数据替换掉本节点中相应位置的数据，如图3最上端所示。最终状态如图3最下端所示，每个节点都有完整的聚合后的数据。同样的假设要传输的数据量为M，节点数为n，AllGather阶段每个节点的发送/接收数据量也为(n-1)M/n。

综合所述，在all-reduce方式下，模型训练的每一个step每个节点需要传输数据量为：

$2M(n-1)/n$

1.2 PS

PS中主要有两种角色，server和worker。server主要用于存储模型参数，并接收来自worker的梯度然后执行优化更新参数，不同的server一般会存储模型的不同部分参数；worker主要执行模型训练的前后向计算，将后向计算的梯度发送给server，并从server拉取最新的参数值。论文中把PS按照部署方式分成了两种方式：Non-colocate PS和Colocate PS。

1.2.1 Non-colocate PS

Non-colocate PS即server和worker是分开部署在不同的机器上的，如图4所示。

图4.Non-colocate PS（图片来自原论文）

下面分析下Non-colocate PS的节点通信量，假设参数量为M，worker节点数为n，server节点数为k。假设参数均匀分配到了每个server上，那么在训练的每一个step中，对于worker来说，需要发送/接收的数据量为M（因为每个worker需要把自己的参数发送给所有的k个server）；对于server来说，需要发送/接收的数据量为n*M/k（因为每个server需要管理所有n个worker的M/k的参数）。一个step中，训练耗时是由M和n*M/k两者的最大值决定的（这里应该是以完全异步的PS为分析前提的），因此可以容易看出，当n=k时能够达到最优，即每个节点（包括server和worker）传输的数据量为：

$M$

1.2.2 Colocate PS

Colocate PS即server和worker是部署在同一个机器上的，每个worker机器上都会有一个server，也就是说server和worker的配比是1:1，如图5所示。

图5.Colocate PS（图片来自原论文）

在Colocate PS方式下，在训练的每一个step中，对于worker来说，需要发送/接收的数据量为(n-1)*M/n（因为每个worker有M/n的数据放在了本机PS上，剩余(n-1)*M/n的参数放在了其他机器PS上）；对于server来说，需要发送/接收的数据量同样也是(n-1)*M/n（因为每个server需要接收来自其他n-1个机器上worker的M/n的参数）。因此总得来看，在Colocate PS方式下，在训练的每一个step中，每个节点（指单个机器）需要传输的数据量为：

$2M(n-1)/n$

从这里可以看出来，Colocate PS与All-reduce的效率是一样的，当然这个结论其实有一个隐含的假设，即本机上的server和worker之间通讯是没有耗时的，实际上本机上的不同进程通信同样是有消耗的，因此理论上讲Colocate PS还是会比All-reduce差一些。

另外我起初在看到这里的时候是感觉到有点违反直觉的，因为Colocate PS相比Non-colocate PS只是把server与worker部署在了一台机器上而已，对于在同一个机器上的server和worker来说，其实是提升了通讯效率的，为什么Colocate PS反而比Non-colocate PS还要差呢？（这个问题留给读者思考下：））

1.3 对比总结

论文中给出了一个对比总结的表，如表1所示。

表1.不同分布式训练方式理论通信耗时总结（表来自原论文）

可以从表中看到，理论上看，PS的通信效率是大于等于all-reduce的，但是在实际使用中，PS的表现却往往不如all-reduce（不考虑大规模稀疏模型），我想可能主要有几点原因吧：一是在上述的PS最优效率是在参数能够在不同server上均匀分配前提下的，这点在实际模型中可能难以做到，模型的参数总是有大有小；二是传统的PS server是需要实现优化器逻辑的且一般是运行在CPU上，这部分计算量在上面的分析中并没有考虑，当模型参数量较大，优化器逻辑又相对复杂（比如adam等带动量的优化器）时，这部分计算耗时是无法忽略的，这一点在论文后面也有提到；三是在PS中，worker和server是n对n的通信，这在节点数较大时，对整个集群的通信带宽也是个挑战。

二.BytePS

2.1 动机

论文作者提到主要基于以下几个观察产生了做BytePS的想法：

一般训练集群都是既有CPU机器又有GPU机器的，GPU的利用率一般相对较高，而CPU的利用率一般较低。
集群中往往有一些非分布式的训练任务，对于这些任务来说，其实是没有利用到本机的带宽的。
All-reduce和Colocate PS无法利用额外CPU机器，Non-Colocate PS在n=k时达到理论最优，但是实际中由于参数分片和优化器逻辑往往无法达到理论最优。

作者的上述观察我觉得是在很多集群中是真实存在的，在我们的集群中也确实有类似现象。基于这些观察，提出的BytePS就是考虑如何将额外的CPU和带宽资源整合到训练任务中来提高效率。

2.2 整体架构

BytePS提出的架构如图6所示。

图6.BytePS架构（图片来自原论文）

BytePS主要由3个部分构成：

GPU Computation：运行在每个GPU机器上，执行训练的前后向计算以及优化器的运算。
Summation Service（SS）：运行在参与计算的所有机器上，只执行简单的参数聚合，而并不实现优化器逻辑。（论文中提到SS聚合的是参数梯度，但是从后面的分析可以看出来SS聚合的实际上是参数的变化量，而不是梯度），类似于PS，模型的参数也是分配到不同的SS上进行聚合的。
Communication Service（CS）：运行在每个GPU机器上，负责本机多卡的参数聚合以及与SS间的通信，即BytePS的参数聚合其实是分成了两级，一级是本机内的（CS），一级是机器间的（SS）。CS在做本机内部参数聚合时会根据卡间及机器间的拓扑结构（例如卡数量，是否支持NVLink，RDMA等）进行通信算法的优化（在开源版本中拓扑结构依赖于手工配置，还没有支持拓扑的自动感知）。

BytePS通过这样的架构设计达到了几个目的，一是通过将参数切分到不同SS进行聚合利用了额外的CPU算力及通信带宽（这点与Non-colocate PS类似），二是将PS中的server优化器计算逻辑和参数聚合逻辑分别拆分到了GPU Computation和SS上，这样优化器的计算逻辑就可以在GPU上运行，从而解决了CPU运行大规模优化逻辑的算力不足的问题，三是通过机内和机间两级聚合并针对不同拓扑结构适配优化了通信效率（这个后面会详细介绍）。

2.3 机间通信分析

BytePS中，所有机间通信都发生在CS和SS间。通信时，会先对参数进行切分成不大于4M的part（小于4M的参数不切分），然后将各个part hash到[0,n^2+kn-2k)整数区间，再映射到不同的SS，n^2+kn-2k的由来下面会分析。通过这样的参数划分，一方面可以一定程度上解决由于参数大小不一导致的负载不均衡问题，另一方面也可以实现发送和接收的overlap以提高通信效率。

假设模型参数量为M，CPU机器数为k，GPU机器数为n，机器间带宽为B（双向），CPU上的SS分配的数据量为Msscpu，GPU上的SS分配的数据量为Mssgpu，在每个训练step中，GPU机器的总发送/接收耗时为tg，CPU机器的总发送/接收耗时为tc。

首先看tg，GPU本机的SS存放了Mssgpu的参数，剩下的M-Mssgpu需要通过CS从其他机器的SS来发送/接收，因此CS的通信量为M-Mssgpu。GPU上的SS需要向其他n-1个CS每个发送/接收Mssgpu，因此GPU上SS的通信量为(n-1)Mssgpu。合起来得到：

$t_g=\frac{M+(n-2)*M_{ssgpu}}{B}$

再看tc，CPU机器上的通信就是由SS决定，在每个step中，其需要向所有n个CS每个发送/接收Msscpu，因此得到（注意这里原论文上面有笔误，漏掉了n）：

$t_c=\frac{{}n*M_{sscpu}}{B}$

因为在每个step中，上面的tg，tc是同时发生的，因此step耗时就是Max(tg, tc)，这里很容易看出，当tg=tc时，通信耗时达到最优。由上面的tg和tc的公式，结合M=k*Msscpu+n*Mssgpu约束，可以很容易推导出来得到：

$M_{sscpu}=\frac{2(n-1)M}{n^2+kn-2k}$

$M_{ssgpu}=\frac{(n-k)M}{n^2+nk-2k}$

上面就是CPU和GPU上SS的最优参数量分配公式，从公式里面也可以看出来为什么在参数分片hash空间是[0, n^2+nk-2k)。同时从上面公式也可以推导出来最优通信耗时：

$t_{opt}=\frac{2n(n-1)M}{(n^2+nk-2k)B}$

根据topt公式，当没有额外CPU机器时，即k=0时，就与All-reduce及Colocate PS方式下一致，而当k=n时，就与Non-colocate PS方式一致了，因此在理论通信耗时方面，BytePS确实统一了All-reduce和PS。

在一般情况下，即0<k<n时，可以计算一下相对All-reduce以及Non-colocate PS的加速比，即用All-reduce及Non-colocate PS理论通信耗时除以BytePS的理论耗时，分别用 $\gamma_a$ 及 $\gamma_p$ 表示（All-reduce及Non-colocate PS的理论通信耗时可以从表1中知道）：

$\gamma_a=\frac{n^2+kn-2k}{n^2}$

$\gamma_p=\frac{n^2+kn-2k}{2k(n-1)}$

上面的加速比公式可以看出，在0<k<n时且n>=2时， $\gamma_a$ 和 $\gamma_p$ 都是大于等于1的，即理论上BytePS的通信效率要优于All-reduce及Non-colocate PS。这里可以进一步分析一下，对于 $\gamma_a$ 来说，是关于k的递增函数，在k<n时，是k越大越好，那是不是越大越好呢？当k>n时，从公式上看似乎 $\gamma_a$ 更大，但实际上达不到，因此网卡带宽是有限的，前面的分析都是建立在每个机器的网卡出入带宽都是一样的，当k>n时，GPU机器上的网卡会成为瓶颈。对于 $\gamma_p$ 来说，可以对k求导，求导结果是：

$\frac{\partial \gamma_p}{\partial k}=\frac{-n^2}{2k^2(n-1)}$

可以看到导数是小于0的，即加速比是随着k增大而降低的，个人理解这是因为对于Non-colocate PS来说，最优情况是k=n时达到M/B，这个是数据并行情况下的通信效率上限了，当k逐渐接n时，Non-colocate PS逐渐接近理论最优，因此加速比呈递减趋势，当k=n时， $\gamma_p$ =1，此时BytePS和Non-colocate PS都达到了最优。

2.4 机内通信分析

机内通信这块比较复杂一些，跟GPU卡以及网卡的硬件连接方式有关系，BytePS论文分了两种情况来进行讨论（硬件部分不是我所擅长，所以以下主要是复述论文）。

2.4.1 PCIe-Only

PCIe-Only的硬件连接如图7所示：

图7.PCIe-Only连接示例（图片来自原论文）

示例设定中单台GPU机器有一张网卡，两个CPU，8张GPU卡；网卡的带宽是100Gbps，通过PCIe总线与其中一个CPU连接，PCIe总线带宽大约是128Gbps，CPU间通过QPI连接，QPI连接的带宽大于300Gbps，8张GPU卡分成了2组，每组4卡，组内通过PCIe switch连接在一起再通过PCIe总线连接到某一个CPU，PCIe switch内部的GPU-GPU通信带宽大约是105Gbps，跨PCIe switch的GPU-GPU通信带宽大约是80Gbps。可以看到在整个通信链路上不同的带宽是存在不一致的。BytePS就是利用这些带宽不一致的先验信息去优化机内通信流程。

BytePS提出了一个CPU-assisted aggregation的参数聚合算法，总体思路是先在PCIe switch内部先做局部聚合，然后再做CPU间的全局聚合，最后再把全局聚合结果广播回到各卡上。具体分为如下6步：

Reduce-Scatter：PICe switch内部的参数聚合，这里与All-reduce中的Reduce-Scatter步骤是一致的。这一步中，每个卡通过PCIe switch传输数据量为(n-1)M/n=3M/4，最终每个卡上会有M/n=M/4的聚合结果数据。
GPU-CPU copy：每个GPU将自己的M/4的聚合数据发送到CPU，单PCIe switch到CPU的传输数据量为M/4*4 = M。
CPU-Reduce：CPU聚合来自PCIe switch的各卡数据，得到本机上的数据聚合结果。
Networking：CS将本机数据聚合结果发送到各个SS并从SS得到最终的所有机器的聚合结果。
CPU-GPU copy：GPU将属于自己那部分的最终聚合数组（每个GPU M/4）从CPU拷贝回来。
All-Gather：每个PCIe switch内部的所有卡做一次all-gather，这个与All-reduce中的All-Gather也是一致的。同样，每个卡通过PCIe switch传输数据量为(n-1)M/n=3M/4，最终每个卡都会得到最终的完整聚合结果数据。

从这个流程来看，其实就相当于在原始的All-reduce算法的Reduce-Scatter和All-Gather两步之间插入了2到4步做了CPU及网络间的聚合工作，体现了“CPU-assisted”。通过这个方式避免直接做跨PCIe switch的通信。

关于CPU-assisted Aggregation的最优性分析如下：

将PCIe-only的硬件拓扑抽象成一个图G=(V, E)，N表示总的GPU卡数，S表示PCIe switch个数，C表示CPU个数。N，S，C即G中的节点，因此有V=N∪S∪C，每一条边b(Vx, Vy)表示节点Vx和Vy之间的带宽。用t(Vx, Vy)表示节点Vx和节点Vy间的数据传输总量，p表示总的PCIe switch数，n表示每个PCIe switch下的卡数。后面的分析做了几个假设：（1）每条边都是全双工且双向带宽是相等的，即b(Vx, Vy)=b(Vy, Vx)。（2）G是对称的，即每个PCIe switch和CPU之间的带宽是一样的，PCIe Switch和GPU之间的带宽也是一样的。（3）QPI带宽远高于其他部分带宽。

所有GPU间的数据聚合，有两种策略，一种是使用CPU-assisted Aggregation（下面简称CAA），一种是使用直接拷贝，即每个GPU把自己的所有数据M直接发送给CPU（下面简称BF），那么一般情况下，是将x比例的数据使用BF聚合，y比例的数据使用CAA，x+y=1。每个PCIe switch和CPU之间的通信量为：

$t(S_j, C_j)=n*xM+\frac{yM}{n}*n=(nx+y)M$

上式第一个等号后的第一项就是BF方式下n个GPU向CPU传输的数据量，第二项是CAA方式下的传输量。

每个GPU和PCIe switch的通信量为：

$t(N_i, S_j)=xM+(\frac{2(n-1)}{n}+\frac{1}{n})yM=(\frac{2n-1}{n}y+x)M$

上式第一个等号后的第一项是BF方式下的单个GPU传输量，第二项括号中的第一项是CAA方式下all-reduce过程（即CAA第一步Reduce-Scatter和第六步All-Gather）的传输量，括号中的第二项是CPU上聚合需要的传输量。那最终一个step的耗时就是：

$J=max(\frac{t(S_j, C_j)}{b(S_j, C_j)}, \frac{t(N_i, S_j)}{b(N_i, S_j)})$

按照上述假设的设置及带宽，可以得到最优情况下，x=0.2，y=0.8。也就是说理论最优是两种策略的混合，但是在最终实现BytePS还是只采用了CAA，而没有用混合策略，是因为一方面CAA相对混合策略的效率差异比较小，另一方面BF策略其实对QPI的带宽要求也会增大，前面分析是建立在QPI带宽不成为瓶颈的前提下，而一旦卡数增加，BF方式下的带宽需求可能会超过QPI带宽，最后一个方面我想也是因为实现上的简洁。因此最终CAA的耗时是（x=0， y=1）：

$J_{ca}=\frac{M}{b(S_j, C_j)}*max(1, \frac{2n-1}{kn}), k=\frac{b(N_i, S_j)}{b(S_j, C_j)}$

如果对于n*p个GPU采用原始的All-Reduce算法聚合数据，那么耗时为：

$J_{ar}=\frac{2(np-1)M}{np*b_{bottleneck}}$

在论文作者的示例中，k=1（即CPU和PCIe switch之间的带宽和GPU与PCIe switch间的带宽，这个不太清楚是不是一般情况）。又p>=1，bbottleneck <= b(Sj, Cj)，所以CAA在PCIe-only拓扑下理论上也是优于All-reduce的。

2.4.2 NVLink-based

带有NVLink的拓扑结构如图8所示：

图8.NVlink-based连接示例（图片来自原论文）

相比于图7，NVLink在单机的不同GPU卡之间增加了一些卡间的硬件P2P连接，PCIe switch从2个变成了4个，每个PCIe switch下连接了2个GPU卡，此外网卡也连接到某一个PCIe switch上。NVLink的带宽远高于PCIe带宽，达到1.2Tbps，在这个拓扑下GPU间通信不再是瓶颈，因此不再需要使用CAA方法。这里的关键是结构的不对称性，即网卡与GPU0,1共享了PCIe switch P0，因为多机训练中会有网络通信，如果还是使用CAA，网卡与GPU0,1就会争用P0的PCIe带宽，使这里的带宽成为整个链路中的瓶颈。为了避免瓶颈，在这个拓扑下，BytePS采用的方法分为如下几步：

Reduce：所有GPU卡都先通过NVLink把数据聚合到GPU2上
GPU-CPU copy：GPU2将数据通过PCIe switch P1发送到CPU0
Networking：CPU0将本机数据聚合结果发送到各个SS并从SS得到最终的所有机器的聚合结果。
CPU-GPU copy：CPU0将最终聚合数据发送到GPU2
Broadcast：GPU2将最终聚合数据通过NVLink广播给所有GPU

通过这个方法，BytePS把P0上的PCIe带宽完全留给了网卡，使得网卡带宽能够得到充分利用。

2.4.3 异步收敛性分析

从上面BytePS的流程分析里面可以看出，BytePS的训练过程是一个同步的过程，即每个GPU上都是聚合了所有梯度之后才会更新参数，但是传统PS是支持异步模式的，即PS的server可以不用等待所有的worker的梯度都收到后才更新参数，而是收到任何一个worker的梯度都可以执行参数更新，因此各个worker拉取到的参数其实是不对齐的，这个在理论上是有问题的，但是实践上看异步模式对最终模型指标其实影响有限，且可以大幅提升训练效率，所以PS很多时候都是采用异步模式。在BytePS里面很难实现PS这种异步模式，因为PS可以使用异步是因为有server来集中存放最新参数，但是在BytePS里面，参数的存储和更新都在不同卡上，SS只是做了简单加和而已，而不存放参数。

BytePS采用了另一种方式来模拟PS这种异步训练，即各GPU上计算出梯度且用优化器计算出本地参数变化量，再把这些变化量通过CS及SS聚合成全局变化量，最后应用全局变化量更新参数。如图9中(b)所示：

图9.PS及BytePS参数更新示例（图片来自原论文）

PS的参数更新公式如下：

$w_{ps}=w_0+\sum_{t=1}^{T}\sum_{i=0}^{n-1}f(g_{i,t})$

其中t表示迭代步数，i表示worker编号，gi,t表示在迭代t中worker i上的参数梯度，f(gi,t)表示在梯度gi,t上应用优化器得到参数变化量。这个公式其实对应的是完全异步的PS。

BytePS的参数更新公式如下：

$w_{byteps}=w_0+\sum_{t=1}^{T}\sum_{i=0}^{n-1}\Delta w_{i,t}$

从更新公式上其实就可以看出来，两者在收敛性上是等价的，但是BytePS的这个“异步”流程只能是与完全异步的PS收敛性等价，理论上应该收敛性应该不如带有局部聚合的PS。另外BytePS的这个所谓“异步”也只是收敛性意义上的，训练流程其实还是同步的。因此这里我也比较疑惑，这种“异步”流程的意义到底是什么呢？

至此，BytePS论文的相关理论要点就介绍这么多，一些实验数据以及其他细节有兴趣的读者可以阅读原论文。下一篇开始介绍BytePS的源码。