来源 | 北京智源人工智能研究院

作者 | 骆沁毅

导读：

2020年4月12日上午，北京智源人工智能研究院和北京大学高能效计算与应用中心联合主办了“AI芯片体系架构和软件专题报告会”，五位学者结合在2020年计算机体系结构顶级会议（ASPLOS和HPCA）中发表的最新研究成果。本文整理了其中一位学者美国南加州大学PhD Candidate骆沁毅分享的《Prague: High-Performance Heterogeneity-Aware Asynchronous Decentralized Training》（Prague：高性能异构异步分散训练）。

骆沁毅，美国南加州大学研究助理，在读博士生，研究方向为分布式机器学习和数据分析，本科毕业于清华大学电子系。在ASPLOS、MICRO、PLDI等顶会上发表数篇论文，曾获国家奖学金、WeTech高通奖学金、提名微软博士生奖学金等荣誉。

本次分享中，骆沁毅带来了她和同事们发表在ASPLOS 2020上的最新作品——Prague[1]。这是一种在异构平台上进行高效机器学习分布式训练的方法，其特点是融合了目前的主流分布式训练算法All-Reduce和前沿算法AD-PSGD的优点，既能在同质环境下获得高性能，又具有良好的异质性容忍度。在本次演讲中，骆沁毅条理清晰地介绍了Prague的整个设计过程和它们背后的思考，包括：探索了算法和系统实现之间、统计效率与硬件效率之间的相互作用，以及深入的同步优化等。

一、背景介绍：机器学习分布式训练

近些年来机器学习已经渗透到社会的方方面面，对日常生活产生了广泛影响。为了使得机器学习模型达到足够高的精度，常用的方法是提高模型自身的复杂度以及采用更大规模的训练数据集。然而这使得模型训练成为一个非常耗时的工作。一般情况下，我们会通过分布式训练来加速模型训练过程。

所谓的分布式训练是把训练工作分配到多个机器（或者称之为Worker）上，每个Worker先是独立处理一部分工作，然后各Worker再进行同步，从而协调完成整个训练任务。根据任务分配和协调方式的不同，分布式训练可以分为四类：数据并行，模型并行，混合并行和流水线并行。

骆沁毅演讲中涉及的方法属于数据并行。和其它方式相比，数据并行具有较高的吞吐量，也被应用于多个机器学习框架中。简单来讲，在数据并行方式中，每个Worker都保有一个完整的模型副本。Worker先会在本地处理分配到的训练数据，比如计算梯度，然后再和其他Worker同步模型参数。

Parameter Server [2] (图1.a所示)，或简称为PS， 是一类传统的数据并行的实现方法。PS是一个中心化的结构，每一个Worker会把自己计算得到的梯度上传到中心节点，中心节点对所有Worker的梯度做平均化处理后，再传回各个Worker。作为中心化的结构，PS中存在中心节点通信瓶颈问题。且Worker数量越多，该问题越突出，这制约了PS的扩展能力。

All-Reduce [3] (图1.b所示)，采用了基于流水线式的同步方式，使得节点间的通信变得更加高效。图1.c则展示了基于去中心化式 (Decentralized) 的通信方式。

图1: 三种不同的数据并行同步方法

二、落后者（Straggler）问题

All-Reduce虽然解决了中心节点通信瓶颈问题，但却因其流水线式的同步方式而遇到新的挑战：落后者 (Straggler)问题。Straggler问题也被称之为异质性 (Heterogenity)问题，指的是当一个平台上不同的Worker具有不同的运行性能时，整个系统的性能受限于运行效率最低的Worker。Straggler可能是确定性的，比如某台机器的计算能力或者网络带宽低于其他机器；也可能是随机性的，比如低性能是因为资源共享，后台操作系统任务，缓存，功耗限制等引起的。

解决Straggler问题的一种常用方式是进行异步训练。然而All-Reduce流水线式的通信结构对异步训练带来了很大的挑战。上节中介绍过的去中心化的结构更有利于异步训练的实施。发表在PMLR 2018上的AD-PSGD [4] 就是去中心化异步训练的一个代表算法。

AD-PSGD算法中，每个Worker都保有整个模型，worker间的通信由一个通信图控制。只有通信图中相邻位置的节点 (即由边直接相连的节点)之间才能进行通信。对于某一特定Worker，在一次参数迭代中，会进行如下三步操作：

1. 在本地计算和应用梯度

2. 随机的选择一个相邻节点

3. 跟此相邻节点进行原子模型参数平均化操作。所谓的原子操作是指在多个节点同时选中了相同的相邻节点情况下，这些节点的平均化操作需要逐次排他性地进行。

下图展示了基于Tensorflow实现的All-Reduce和AD-PSGD应用在同构和异构平台上时的性能比较。相比于同构平台，All-Reduce应用在存在Straggler问题的异构平台上时，性能会明显下降。而AD-PSGD则可以很好地抵抗Straggler问题，从而在同构和异构平台上表现出相当的性能。但是AD-PSGD也有其明显的短板。当应用在没有Straggler的同构平台上时，其性能明显低于All-Reduce。这是由于AD-PSGD中模型参数平均化操作的原子性所带来的大量开销导致的。

图2: All-Reduce和AD-PSGD在同构和异构平台上的性能比较

接下来我们通过一个具体的例子来进一步了解AD-PSGD通信开销的来源。如图3所示，以Worker0为例，首先它会在本地计算出梯度，然后随机选中Worker3来平均化模型参数。模型参数平均化的操作可以借助图3右半部分的矩阵来完成。然而同一时间可能会有多个Worker选择了同一个目标Worker来做参数平均化。如图4所示，Worker0和Worker4同时选择了Worker3做参数平均化。因为参数平均化操作的原子性，同一时间Worker0和Worker4中只能有一个和Worker3进行参数平均化，另外一个只能等待前者完成。正是这类因冲突而串行化的过程使得AD-PSGD具有较大的通信开销。

图3: AD-PSGD中的同步操作

图4: 两个Worker同步操作出现冲突

三、Partial All-Reduce机制

如图4右半部分所示，出现两个Worker同步操作冲突的情况下，相应的参数平均化操作依然可以通过矩阵来表示。既然如此，为何不通过一次联合同步来代替两次单独的同步来降低同步开销呢？此外，既然两次同步可能以任意的顺序发生，比如可能是Worker0先和Worker3同步，也可能是Worker4先和Worker3同步，因此如图4所示，Worker0和Worker4收到不同参数的做法不符合常理。

鉴于上述分析，骆沁毅在Prague中使用了基于组的Worker同步机制，也就是Partial All-Reduce (简称为P-Reduce) 机制。P-Reduce采用图5中所示的参数平均化矩阵。

图5: Partial All-Reduce中的同步操作

在P-Reduce中，一组Workers 作为同步的基本单元。Prague系统的设计思路如下：

1. 每个Worker在本地训练模型，进行梯度计算及更新模型参数。

2. 从Group Generator（GG）中获得自己的分组信息。

3. 同一组中的Worker统一进行P-Reduce操作。

Group Generator（GG）是一个独立中心化部件，专门用于为Worker集群生成组，需要放在一个比较稳定的Node上面进行操作。当Worker想要同步的时候，只需要联系GG，GG会返回代表当前Worker的组信息（即针对每个Worker接下来将进行的操作）。如图6中所示，当W0选中W3之后，W0先联系Group Generator，GG生成组信息{0,3,4}传给W0，同时将生成的组信息放入组内其他成员W3和W4 的Group Buffers中。当W4联系Group Generator时，如果它对应的Group Buffer不为空，则GG会直接读取Buffer中的第一条组信息返回给W4。只有在W0、W3、W4都联系Group Generator之后，才进行P-Reduce的操作（即所有Group Member都拿到自己的分组信息之后Collective Communication才能开始）。

图6：Worker从Group Generator获取组信息过程

Group Generator设计的关键点在于如何生成组，组的质量和数量会直接决定同步的速度。针对于生成组主要考虑两个问题，第一是组间冲突（Conflict）问题，第二是Straggler问题。

1. 避免组间冲突

组间冲突（Conflict）是指不同组的Worker有重叠，这些组在执行P-Reduce操作时仍然需要满足原子性（即串行操作）。这也是Group Generator采用中心化组织的原因，如果每个Worker自己生成组，产生Conflict的可能性非常大，需要多轮communication才能消除冲突。虽然采用中心化组织，Worker和GG之间也需要通信，但是交互信息非常小，不会导致通信或者规模瓶颈。Group Generator减少冲突的方法有两种：Fixed Schedule 和 Randomized Schedule。

固定调度 (Fixed Schedule)

Fixed Schedule即Static GG，手动实现无冲突的分组。通过多轮实验，发现好的Schedule是周期性的。图7中展示了一个周期 Schedule，交替于Inter phase和Intra phase之间。图中有16个Worker，W0-W3、W4-W7、W8-W11、W12-W15分别属于不同的Node。Inter phase 是指会有node之间的Worker communication。Intra phase是指node内部的Worker communication，在node内部communication可以加快执行速度。图中第一个iteration是Inter phase, 所以G5中的Worker 在不同的node上。下一个iteration是Intra phase，所以每个组中的Worker都在同一个node上。

图7：无冲突静态调度策略

静态化的schedule可以避免出现conflict，但是当存在straggler时抵抗能力较弱。

随机调度 (Randomized Schedule)

每个Worker 发送请求给GG时，如果GG总是仅生成一个随机的组，则可能会产生很多冲突。Randomized Schedule提出Group Division想法，即GG每次不只是生成一个组，而是生成多个组，这些组之间没有conflict，并将生成的组信息放入各个Worker对应的Group Buffer中。如图8所示，第一个iteration生成组{0,3,4}同时生成组{1,2}，下一轮生成组{0,2}同时生成组{1,3,4}，依次放入对应的Group Buffer中。

图8：随机调度策略

图9：Inter-Intra Synchronization

为了进一步提高性能，借鉴静态Schedule中的Inter phase和Intra phase设计，在生成Group Division时也在Inter-Intra之间交替。如图9所示，每次GG会同时生成2个Group Divisions，一个是在Inter Phase, 一个是在Intra Phase。下一轮又会生成不同的2个Group Divsions，以此类推。

2. 对抗落后者 (Straggler) 问题

如果恰好选中一个straggler和别的Worker在同一组，其他成员需要等待这个straggler。为了优化这个问题, Prague使用了Threshold Rule方法，定义如图10所示。GG通过实时追踪每个Worker多长时间请求一次组信息，确定哪个Worker是straggler。当给快Worker生成组的时候不把慢Worker放在组中，这样慢Worker永远不会影响快Worker。当慢Worker请求组时， 也给其分配快Worker，提供Catch Up的机会，让其可以继续回到训练的过程当中。

图10：Threshold Rule 定义

四、系统实现

Pargue基于TensorFlow实现，其中：P-Reduce使用NCCL作为后端，通过MPI创建NCCL communicator；GG使用gRPC python 包实现。实现两个版本的Pargue系统，分别是：

1. Static GG：手动设计 Group Schedules。

2. Smart GG：实现前面提到的 Randomized Schedules, 包括randomized group division, inter-intra synchronization 和 threshold rule。

同时，实现两个baseline 算法，分别是：

All-Reduce：Horovod with NCCL backend on TensorFlow,

AD-PSGD：TensorFlow remote variable access。

五、实验与结果

1. 实验过程

在图像和自然语言处理领域都进行了实验，使用的模型和数据集为：

1. Vision：VGG-16 on CIFAR-10；Resnet-18、Resnet-50和Resnet-200 on ImageNet。

2. NLP：Transformer on News-Commentary dataset。

硬件环境：GTX partition in the Maverick2 cluster of TACC Super Computer（4 GPUs/node），共有16个Workers，分布在4个nodes上，每个node有4个Workers, 每个Worker占用一个GPU。同时，针对Resnet-50，在8个GTX nodes（32个Workers）上也进行了实验。

在实验中，通过人工随机添加slowdown，即在每个iteration中，每个Worker以1/#Workers的概率随机slowing down。每个Worker slowing down 的概率是相互独立的。

2. 实验结果

① 吞吐量：分为有slowdown和没有slowdown两种情况，结果如图11所示。

图11：Speedup w.r.t. throughput (normalized to All-Reduce without slowdown)

可以看到在没有slowdown的情况下，无论是Static GG还是Smart GG吞吐量都与All-Reduce相近或者相优。针对Transformer模型可以实现将近4倍的性能提升，主要是因为Transformer相较于上面的CNN模型计算量要少一些，通信量比较大，再加上GG主要是对通信做的优化，所以相比之下GG在Transformer上会比All-Reduce快很多——即使在没有slowdown的情况下，Smart GG也有明显优势。在有slowdown的情况下，Smart GG比All Reduce快很多。以上数据都是以没有slowdown情况下All-Reduce的吞吐量作为单位1。

② 收敛情况：达到某个loss需要的时间，结果如图12所示。

图12：Speedup w.r.t. time to reach the same loss (VGG-16)

图中展示了四个算法在不同的异构性平台上训练模型的效率。其中，蓝色代表没有slowdown，橙色代表有2倍slowdown， 灰色代表有5倍slowdown。纵坐标是Speedup，数值越高表示性能越好。从图中可以看出当存在slowdown时，All-Reduce影响比较明显。AD-PSGD对slowdown不敏感，基本无影响。但是AD-PSGD 本身存在Communication Overhead的问题，在没有slowdown的情况下，效果并没有All-Reduce那么好。Static GG在有slowdown的情况下，也受到很明显的影响。Smart GG在没有slowdown的情况下达到与All-Reduce相近的效果，同时又在有slowdown的情况下相较于All-Reduce有更好的对抗Straggler的能力。

③ 模型准确率：最终得到的模型的准确率，结果如图13所示。

图13：Accuracy of the model

图中实验是在没有slowdown的情况下进行。对于Resnet-50，Prague与All-Reduce的准确率相近。而对于Transformer（5小时定时实验），与All-Reduce相比Prague在相同时间内取得了更高的BLEU score。

Q&A

Q1：你的工作是不是针对于数据中心的？你对物联网终端的分布式训练有什么看法？

A1：可能是想问关于 Federated Learning的情况，确实现在相当火，同样也有Straggler问题，并且由于手机终端性能高低不一且受网络带宽的影响Straggler问题比数据中心更加严重。我也看到一些基于 Federated Learning Decentralized Training 的工作。其他有一些工作比如最近看到一篇发表在MLSys 2020上的论文，该论文把Federated Learning中System Heterogeneity看作Data Heterogeneity，比如如果没有完成训练，可以看成是训练数据不够多。于是该论文从算法的角度通过调整Loss Function来对抗Heterogeneity问题。我也看到了一些Decentralize的工作，一些关注怎样选择Worker的工作，这些都是很好的研究方向，也有很多机遇和挑战。

Q2：你说这个技术适用于Data Parallelism，对于Model Parallelism呢？

A2：没有考虑用Model Parallelism 原因是Model Parallelism的吞吐量没有Data Parallelism好，同时Model Parallelism在对抗Straggler问题上有一定难度。最近有一些Pipeline Parallelism的论文，是基于Model Parallelism的情况下，减少Worker的idle time。由于是pipeline的原因，对于解决Heterogeneity也是一个很好的工作方向。我看到已经有相关论文发表，思路是可以在一个Pipeline的一节或者在Model Parallelism中的一个Worker出问题的情况下，把它的工作传给另外一个Worker去完成，但是这样可能会有一些开销存在。再考虑基于Pipeline的想法，看到Pipeline Parallelism工作中，有些Pipeline Stage中有Data Parallelism，在这种情况下是不是可以考虑用backup worker ？如果有一个Worker出问题了，其他Worker还能正常工作，我们是不是可以把这个batch扔掉？所以在这个方向上考虑Straggler问题还是很有趣的。

Q3：Slowdown Node只测了一个吗？有没有对Prague性能的敏感度分析？

A3：与其说只测了一个，其实是说我们只测了一个Slowdown的Probability。我们的做法是每一个Worker在每一个Iteration都有一定的概率会Slowdown。即同时可能会有多个Worker Slowdown ,甚至可能没有Worker Slowdown。这都是一个随机概率出现的问题。我们认为这是一个比较好的模拟方式。我们确实只测了一个Probability，但是我们相信在更多的Probability下结果应该是满足一致性的。

参考文献

[1] Luo Q, He J, Zhuo Y, et al. Prague: High-Performance Heterogeneity-Aware Asynchronous Decentralized Training[C]//Proceedings of the Twenty-Fifth International Conference on Architectural Support for Programming Languages and Operating Systems. 2020: 401-416.

[2] Li M, Andersen D G, Park J W, et al. Scaling Distributed Machine Learning with the Parameter Server[C]//11th {USENIX} Symposium on Operating Systems Design and Implementation ({OSDI} 14). 2014: 583-598.

[3] Sergeev A, Del Balso M. Horovod: Fast and Easy Distributed Deep Learning in TensorFlow[J]. arXiv Preprint arXiv:1802.05799, 2018.

[4] Lian X, Zhang W, Zhang C, et al. Asynchronous Decentralized Parallel Stochastic Gradient Descent[J]. arXiv Preprint arXiv:1710.06952, 2017.