现有的分布式 DNN 训练架构无法充分利用异构资源实现高性能训练。近期,来自字节跳动和清华大学的研究人员提出一种新型分布式 DNN 训练架构——BytePS,解决了这一问题,实现了大规模训练性能的显著提升。这项工作已在国际顶级计算机系统会议 OSDI’20 上发表,其开源代码在 GitHub 上获得 2400 stars。
近年来 DNN 为计算机视觉、语音识别与合成、自然语言处理等领域带来了突破性进展。然而,训练这些 DNN 模型通常需要大量算力。大型公司一般会建设具备数千块 GPU 卡的大型训练集群,其核心资源是 GPU 机器,同时这些 GPU 机器也具备高端 CPU 资源。此外,也有一些仅具备 CPU 的机器,用于训练数据预处理和生成等任务。这些 GPU/CPU 机器通过高速网络连接,以加速分布式训练的通信。目前工业界主流的分布式训练是基于数据并行方式实现的,其中具有代表性的两种架构是 All-reduce 和参数服务器(PS)。如下图所示,目前的 All-reduce 和 PS 架构在训练性能上距离最优情况都有较大差距。在异构集群场景下,All-reduce 和 PS 架构对资源的利用情况如下图所示。All-reduce 架构中仅用到 GPU 机器,这是因为其设计假定了每个节点都是同构节点。迭代过程中,GPU 独立计算模型参数的梯度,然后使用 All-reduce 通信聚合梯度。
PS 架构则包含 GPU worker 和 CPU server。迭代过程中,GPU worker 将梯度传输至 CPU server;后者将接收到的不同 workers 的梯度做聚合,然后执行 DNN 优化器(如 RMSProp 或 Adam 等)并将更新后的参数传输回 GPU workers。
图 2:All-reduce 和 PS 架构示意图。可以看出,All-reduce 的同构化设计导致其无法充分利用这种异构资源,即只有 GPU workers 之间通信,而无法利用其他 CPU 和带宽资源。而 PS 虽然能够利用 CPU 机器作为 server,却可能在 CPU server 数量较少的时候产生流量热点(例如形成多对一的情况),从而导致网络拥塞。 如今一台训练机器通常都具备有多张 GPU 卡(例如 4 或 8 卡)。在做机器间的通信前,机器内部的多 GPU 之间需要首先做一次本地通信,该通信过程一般是基于 PCIe 或 NVLink 链路。 我们观察到,目前的网卡如 Mellanox CX5 带宽已达 100Gbps,已经很接近 PCIe 3.0 x16 的 128Gbps 链路带宽。而不幸的是,现在流行的机器内聚合方式(例如 8 卡直接做 all-reduce)会使 PCIe 成为瓶颈,导致网卡无法达到其 100Gbps 带宽上限。即使对含有 NVLink 的拓扑,我们也能发现类似的 PCIe 竞争导致的瓶颈问题。 从前面的问题 1 中可以看出,相比于 All-reduce 而言,PS 架构实际上是存在更大的潜力的,因为它能充分利用异构 GPU/CPU 资源。然而,目前的 PS 甚至比 All-reduce 性能显著低(图 1),似乎与之矛盾。这是因为 PS 架构中还存在另外一种瓶颈限制了其性能 – CPU 瓶颈。 顾名思义,PS (参数服务器)需要将参数存储在 CPU server 上,这就意味着需要将优化器 (如 Adam/RMSProp 等) 放在 Server 上去执行。然而,优化器通常包含复杂的数学运算,将会消耗大量的 CPU 内存带宽。在 100Gbps 的网络输入情况下,CPU 将无法满足将完整优化器放置在其上运行的需求。针对以上三个问题,BytePS 逐一提出了解决方案。该优化涉及到对问题的建模,在给出具体的公式前,先从直观上介绍其思想。(1)由前文可知,PS 仅利用了 GPU 机器与 CPU 机器之间的带宽。在 CPU 机器数量较少时,GPU 机器的带宽 B 无法充分利用。图 3 给出了一个例子,在这种情况下,GPU 机器仅能达到 2/3 的最大带宽,剩余 1/3 带宽未得到利用。(2)All-reduce 仅利用了 GPU 机器之间的带宽。此时,GPU 机器与 CPU 机器之间的带宽未得到利用。(3)BytePS 结合两者之长,同时利用了 GPU 与 GPU 之间、GPU 与 CPU 之间的带宽,使得每台机器的带宽都能被充分利用。这就是 BytePS 机器间通信的思路。图 5: BytePS 从通信层面结合 PS 与 All-reduce。该思路在实现过程中,需要考虑如何分配 GPU 与 GPU 之间(设为 x%)、GPU 与 CPU 之间(设为 y%)的流量比例。经过计算,最优比例如下:(其中 n 表示 GPU 机器的数量,k 表示 CPU 机器的数量,)以上即为最优通信策略,对于不同的 n 与 k,采用该策略可使得机器间通信时间最小。目前市面上主流有两种机器拓扑:PCIe-only 型 8 卡机器和 NVLink-based 型 8 卡机器。BytePS 针对这两类拓扑提出了通用的解决方案和设计原则。如图 6 所示,标记 0-7 的灰框表示 GPU,P0 和 P1 表示 PCIe switch。现实当中, P0-CPU0 以及 P1-CPU1 是带宽最小的链路,因此优化目标是最小化这条链路上传输的数据量。目前主流的做法是对这 8 卡做一次全局 All-reduce,这样 P0-CPU0 需要传输的数据量是 7M/4(根据 All-reduce 的通信量计算得出),其中 M 是每张卡上的梯度大小。注意到该做法并没有利用 CPU 的计算能力。BytePS 的核心思想是利用 CPU 的计算能力减少瓶颈链路的传输数据量。如图 7 所示,首先每个 PCIe switch 下的 4 张卡先进行一次 Local Reduce-scatter。该步骤之后,每张卡上已聚合的梯度大小为 M/4。图 7:PCIe-only 拓扑解决方案步骤(1)。下一步,每张卡将自身已聚合的 M/4 梯度拷贝到主机的内存上,如图 8 所示。注意这个过程使得 P0-CPU0 只传输了 M/4*4=M 的流量。图 8:PCIe-only 拓扑解决方案步骤(2)。此时,两个 NUMA node 的主机内存上各自有一份大小为 M 的梯度。我们再利用 CPU 将这两份梯度做一次聚合。但是这个过程的传输只发生在带宽较大的 QPI 上(>300Gbps),并不会产生瓶颈。于是这一系列步骤不但实现了预期中的梯度聚合效果,还使瓶颈链路的传输量从 7M/4 降低为 M,显著降低了通信时间。这里的核心设计原则是:尽量避免跨 NUMA GPU 的直接通信,而可以利用 CPU 的聚合能力来间接完成。图 9 是 NVLink-based 机型的示意图。对于这种拓扑,GPU 之间可以通过超高带宽的 NVLink 链路进行通信。由于 NVLink 带宽显著大于 PCIe 带宽,PCIe 瓶颈问题显得更加严重。可以看到,图 9 中 P0-CPU0 链路(标红色的线段)会同时被以下两种传输同时竞争:(1)CPU0 的内存往 GPU0/GPU1 拷贝数据;(2)CPU0 的内存往网卡 NIC 发送数据。由于 P0-CPU0 链路的带宽与网卡带宽很接近,这种竞争会导致网卡无法发挥最大带宽。为解决这一竞争问题,BytePS 利用了 NVLink 带宽显著高于 PCIe 链路的事实,利用 Reduce(而非 Reduce-scatter)方式避免 PCIe 竞争。如图 10 中红线所示,所有卡先将其梯度通过 NVLink 传输至 GPU2 上并做 Reduce,接着 GPU2 将聚合后的梯度拷贝到 CPU0 内存,再经由网卡发送出去。由于 NVLink 带宽很高,这种做法不会导致 GPU2 产生流量热点问题,但却能够避免在 P0-CPU0 链路上发生的竞争。图 10:BytePS 对 NVLink-based 拓扑的解决方案。3、CPU 瓶颈优化:Summation Service前文提到,优化器对于 CPU 而言是比较重的任务,这也是 PS 架构的性能缺陷之一。然而,如何高效利用 CPU 的异构计算能力是 BytePS 的核心诉求之一,这就需要克服数据同步过程中的 CPU 瓶颈。经分析,优化器可被拆解为两部分:(1)Sum:将来自其他 GPU workers 的梯度求和并得到一份聚合后的新梯度;(2)Update:利用新梯度对参数进行更新。后者对于 CPU 而言的确是非常消耗内存带宽的操作,但前者却能够在 CPU 上高效实现(例如 AVX 指令集)。如图 11 所示,求和操作在 CPU 上可以达到远超网络带宽的吞吐率,即不会引入 CPU 瓶颈。受到这个发现的启发,BytePS 提出了 Summation Service 概念,对传统 PS 的 CPU 瓶颈问题做了改进。如图 12 所示,不同于 PS 将完整优化器放置在 CPU 上的设计,Summation Service 只将 Sum 操作放置在 CPU 上,而将 Update 操作交由计算能力更强大、内存带宽更充足的 GPU 来执行。这种设计能够避免同步过程中的 CPU 瓶颈。图 12:PS 架构与 Summation Service 的对比。将前述 BytePS 三个设计点结合起来,我们得到 BytePS 完整系统架构,如下图所示。 每台 GPU 机器上部署了一个 Communication Service 模块,负责聚合本地多卡的梯度(即机器内多卡通信),其能充分考虑机器内部复杂的拓扑,避免产生 PCIe 瓶颈。
每台 GPU/CPU 机器上部署了一个 Summation Service 模块,处理来自其他 GPU 机器的梯度,其能够高效运行在 CPU 上。
模块之间通过网络互连,通信策略使用的是前述设计中提到的最优网络通信方案。经证明,该方案不仅有最佳的性能,且能够从通信角度统一 All-reduce 和 PS 两种架构。
BytePS 在实现中充分利用了流水线的思想,尽可能地将计算和 PCIe 和网络传输重叠起来。此外,BytePS 还利用了一些 RDMA 实现上的优化技巧(例如内存页对齐等),解决了实际部署时遇到的 RDMA slow receiver 问题,实现了高性能的网络传输。对于用户侧,BytePS 提供了丰富的用户接口,能够兼容 TensorFlow,PyTorch 和 MXNet 等主流深度学习框架。通常只需要几行至十几行代码的修改,就可将现有基于其他框架(如 Horovod 或 PyTorch DDP 等)的代码迁移至 BytePS 上运行。BytePS 对多种 CV 类(包括 ResNet-50,VGG-16 ,GAN)和 NLP 类(Transformer,BERT-Large,GPT-2)模型都做了分布式性能评测,规模从 8 卡 - 256 卡。所使用的硬件是 V100 GPU 和 100Gbps RDMA 网络。对照组为目前广泛使用的 All-reduce 和原生 PS 实现。图 14 和图 15 分别展示了 CV 和 NLP 模型上的评估结果。总体而言,BytePS 在各类模型上都取得了正向收益,且相比于 All-reduce 和 PS 能够达到的最大提升幅度达 84% 和 245%。在字节跳动内部,BytePS 已经得到了部署并在真实的业务场景中取得显著收益。BytePS 及实验代码都已公开:https://github.com/bytedance/byteps。理论字节跳动清华大学分布式DNN训练架构BytePS