现有的分布式 DNN 训练架构无法充分利用异构资源实现高性能训练。近期，来自字节跳动和清华大学的研究人员提出一种新型分布式 DNN 训练架构——BytePS，解决了这一问题，实现了大规模训练性能的显著提升。这项工作已在国际顶级计算机系统会议 OSDI’20 上发表，其开源代码在 GitHub 上获得 2400 stars。
论文：https://www.usenix.org/conference/
osdi20/presentation/jiang
代码：https://github.com/bytedance/byteps
背景
近年来 DNN 为计算机视觉、语音识别与合成、自然语言处理等领域带来了突破性进展。然而，训练这些 DNN 模型通常需要大量算力。大型公司一般会建设具备数千块 GPU 卡的大型训练集群，其核心资源是 GPU 机器，同时这些 GPU 机器也具备高端 CPU 资源。此外，也有一些仅具备 CPU 的机器，用于训练数据预处理和生成等任务。这些 GPU/CPU 机器通过高速网络连接，以加速分布式训练的通信。目前工业界主流的分布式训练是基于数据并行方式实现的，其中具有代表性的两种架构是 All-reduce 和参数服务器（PS）。
研究动机：现有架构的设计缺陷
如下图所示，目前的 All-reduce 和 PS 架构在训练性能上距离最优情况都有较大差距。
图 1：VGG16 训练性能。
究其原因，有如下三个方面的问题。
1、机器间网络通信问题
在异构集群场景下，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 数量较少的时候产生流量热点（例如形成多对一的情况），从而导致网络拥塞。
2、机器内多卡 PCIe 带宽竞争问题
如今一台训练机器通常都具备有多张 GPU 卡（例如 4 或 8 卡）。在做机器间的通信前，机器内部的多 GPU 之间需要首先做一次本地通信，该通信过程一般是基于 PCIe 或 NVLink 链路。 我们观察到，目前的网卡如 Mellanox CX5 带宽已达 100Gbps，已经很接近 PCIe 3.0 x16 的 128Gbps 链路带宽。而不幸的是，现在流行的机器内聚合方式（例如 8 卡直接做 all-reduce）会使 PCIe 成为瓶颈，导致网卡无法达到其 100Gbps 带宽上限。即使对含有 NVLink 的拓扑，我们也能发现类似的 PCIe 竞争导致的瓶颈问题。
3、通信链路中的 CPU 瓶颈问题
从前面的问题 1 中可以看出，相比于 All-reduce 而言，PS 架构实际上是存在更大的潜力的，因为它能充分利用异构 GPU/CPU 资源。然而，目前的 PS 甚至比 All-reduce 性能显著低（图 1），似乎与之矛盾。这是因为 PS 架构中还存在另外一种瓶颈限制了其性能 – CPU 瓶颈。 顾名思义，PS (参数服务器)需要将参数存储在 CPU server 上，这就意味着需要将优化器 (如 Adam/RMSProp 等) 放在 Server 上去执行。然而，优化器通常包含复杂的数学运算，将会消耗大量的 CPU 内存带宽。在 100Gbps 的网络输入情况下，CPU 将无法满足将完整优化器放置在其上运行的需求。
BytePS 主要设计
针对以上三个问题，BytePS 逐一提出了解决方案。
1、机器间网络通信优化
该优化涉及到对问题的建模，在给出具体的公式前，先从直观上介绍其思想。
（1）由前文可知，PS 仅利用了 GPU 机器与 CPU 机器之间的带宽。在 CPU 机器数量较少时，GPU 机器的带宽 B 无法充分利用。图 3 给出了一个例子，在这种情况下，GPU 机器仅能达到 2/3 的最大带宽，剩余 1/3 带宽未得到利用。
图 3: PS 带宽利用示意图。
（2）All-reduce 仅利用了 GPU 机器之间的带宽。此时，GPU 机器与 CPU 机器之间的带宽未得到利用。
图 4：All-reduce 带宽利用示意图。
（3）BytePS 结合两者之长，同时利用了 GPU 与 GPU 之间、GPU 与 CPU 之间的带宽，使得每台机器的带宽都能被充分利用。这就是 BytePS 机器间通信的思路。
图 5: BytePS 从通信层面结合 PS 与 All-reduce。
该思路在实现过程中，需要考虑如何分配 GPU 与 GPU 之间（设为 x%）、GPU 与 CPU 之间（设为 y%）的流量比例。经过计算，最优比例如下：
（其中 n 表示 GPU 机器的数量，k 表示 CPU 机器的数量，）
以上即为最优通信策略，对于不同的 n 与 k，采用该策略可使得机器间通信时间最小。
2、机器内多卡通信优化
目前市面上主流有两种机器拓扑：PCIe-only 型 8 卡机器和 NVLink-based 型 8 卡机器。BytePS 针对这两类拓扑提出了通用的解决方案和设计原则。
（1）PCIe-only 型拓扑
如图 6 所示，标记 0-7 的灰框表示 GPU，P0 和 P1 表示 PCIe switch。现实当中， P0-CPU0 以及 P1-CPU1 是带宽最小的链路，因此优化目标是最小化这条链路上传输的数据量。目前主流的做法是对这 8 卡做一次全局 All-reduce，这样 P0-CPU0 需要传输的数据量是 7M/4（根据 All-reduce 的通信量计算得出），其中 M 是每张卡上的梯度大小。注意到该做法并没有利用 CPU 的计算能力。
图 6：PCIe-only 型拓扑。
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 的聚合能力来间接完成。
（2）NVLink-based 拓扑
图 9 是 NVLink-based 机型的示意图。对于这种拓扑，GPU 之间可以通过超高带宽的 NVLink 链路进行通信。由于 NVLink 带宽显著大于 PCIe 带宽，PCIe 瓶颈问题显得更加严重。可以看到，图 9 中 P0-CPU0 链路（标红色的线段）会同时被以下两种传输同时竞争：（1）CPU0 的内存往 GPU0/GPU1 拷贝数据；（2）CPU0 的内存往网卡 NIC 发送数据。由于 P0-CPU0 链路的带宽与网卡带宽很接近，这种竞争会导致网卡无法发挥最大带宽。
图 9：NVLink-based 拓扑示意图。
为解决这一竞争问题，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 瓶颈。
图 11：求和操作在 CPU 上的吞吐率。
受到这个发现的启发，BytePS 提出了 Summation Service 概念，对传统 PS 的 CPU 瓶颈问题做了改进。如图 12 所示，不同于 PS 将完整优化器放置在 CPU 上的设计，Summation Service 只将 Sum 操作放置在 CPU 上，而将 Update 操作交由计算能力更强大、内存带宽更充足的 GPU 来执行。这种设计能够避免同步过程中的 CPU 瓶颈。
图 12：PS 架构与 Summation Service 的对比。
总体系统架构
将前述 BytePS 三个设计点结合起来，我们得到 BytePS 完整系统架构，如下图所示。
图 13：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%。
图 14：CV 类模型性能评估。
图 15：NLP 类模型性能评估。
在字节跳动内部，BytePS 已经得到了部署并在真实的业务场景中取得显著收益。BytePS 及实验代码都已公开：https://github.com/bytedance/byteps。
你点的每个“在看”，我都认真当成了喜欢
