本文来自 SysML 2018，由星云 Cluster 编译并授权 InfoQ 发布，原文链接：http://www.sysml.cc/doc/151.pdf
大规模分布式机器学习应用中，GPU 间参数聚合的速度对整体训练速度起到至关重要的作用，尤其当今 GPU 算力越来越强，参数聚合速度的重要性也日益显著。伯克利与微软在 SysML 2018 上推出了一个基于 NVLink 构建的高性能的参数聚合通讯库 Blink 并发表了相关论文。文中重点介绍了 Blink 的设计，并通过实验来证明其有效性。
1. 引言
大型深度学习模型进行训练时，需要花费不少的时间，如 ImageNet 1K 等模型通常需要数天甚至数周才能在单个 GPU 上进行训练，所以就需要从单 GPU 扩展到多 GPU 进行训练。减少 DNN 训练时间最广泛使用的方法是使用数据并行随机梯度下降（SGD）来并行化训练。在数据并行训练中，每个 GPU 具有模型参数的完整副本，并且在输入数据的子集上独立地训练， GPU 经常需要与参与训练的其他 GPU 交换参数。在大规模训练时，跨 GPU 同步参数会带来显著的开销 - 这个问题会因为 GPU 计算越来越快，模型规模越来越大而变得更加严重，从而使得通信成本上的问题变得愈发突出。
模型参数交换通常使用 collective 通信原语实现，例如 All-Reduce [2]。NVIDIA collective 通信库（NCCL）[7] 是一种先进的实现方式，可通过 PCIe 或更新的互连方式（如 NVLink）提供 GPU 间 collective 通信原语。通过将 NCCL 纳入 Tensorfow [1]，Uber 的研究人员表明，端到端的训练过程可以加快 60％。
但是，对于某个特定拓扑结构，NCCL 并不总是有效地使用所有可用链路。这是因为 NCCL 使用的是基于环的模式进行数据传输，并在给定拓扑中创建尽可能多的环。考虑图 1 中的拓扑，是现代 NVIDIA DGX-1 的拓扑实现，这里我们从 GPU A 使用 Broadcast 操作。由于每个链路都是双向的，我们可以构建两个环，如图 1 所示（a）。从一个环开始是 A-> B-> D-> C-> A，另一个反方向是 A-> C-> D-> B-> A。要做广播 A 可以将数据分成两部分，并在每个环上发送一个部分。因此，如果数据大小为 n 并且链路带宽大小为 b，则所花费的时间将是 n 除以 2b。请注意，A < - > D 和 B < - > C（虚线表示）的两个交叉链接没被有效利用起来。
为实现更高的链路利用率，从而加快速度传输，我们提出 Blink，这是一系列使用基于广播的数据传输方案的协议，能够充分利用相连的 GPU 组来实现更高效率。Blink 还旨在克服拓扑异构性的困难，这种困难是由于使用的 GPU 数量不同、硬件层的异构性，多机器训练而引起的。例如，在图 1（b）中，我们在 Blink 中展示了一个广播协议。在这种情况下，从 GPU A 发送到 GPU B 的数据，然后再被广播到 GPU C 和 GPU D。我们可以构造三个这样的转发树，因此链路利用率提高，并且所花费的总时间变为 n 除以 3b。
通常，Blink 的协议遵循如下分层方案：对于给定网络拓扑，首先将网络划分为组内所有节点完全连接的组。在第一阶段，执行内部广播，其中使用组内的每个全连接节点进行通信交换数据。在第二阶段，执行跨组转发，跨组进行通信，并转发相应组内的跨组数据。我们在 Blink 中为四个通信原语（Broadcast，Gather，All-Gather，All-Reduce）设计了对应的通信协议。在 NVIDIA DGX-1 机器上使用多达 8 个 GPU 的实验表明，与最先进的库相比，Blink 可以实现高达 2 倍的加速。
2. 背景
GPU 互连拓扑：在这项工作中使用的主要测试平台是 NVIDIA DGX-1，它是一个配备 8 个 P100 或 V100 GPU 的架构。GPU 不仅可以通过 PCIe 连接，还可以通过新设计的被称为 NVLink [8] 的互连技术连接。NVLink 是一种高带宽且节能的互联技术，可实现 20-25 GB / s 的吞吐量。如图 2（a）所示，基于 P100 的 DGX-1 具有 NVLink 拓扑。该拓扑结构由两个方形 GPU 组连接组成，此外还有 2 条跨 GPU 组的长连接。
广播与环：基于环的 NCCL 数据传输浪费了无法形成新环的链路。基于网络拓扑结构，NCCL 将拓扑划分为不相交的环。在大多数情况下，我们有剩余的、不能形成新的环的链接，导致这些链接闲置或浪费。在我们基于 P100 的 DGX-1 机器的 NCCL 基准测试中，如果我们使用 4 到 7 个 GPU，它只能形成一个双向环（如图 2（a）所示），超过一半的链路空闲。例如，如果我们在 DGX-1 机器内的 8 个 GPU 中使用 6 个 GPU，则它仅使用 16 个 NVLink 中的 6 个。
3．BLINK
接下来，我们来说明 Blink 中基于广播的协议，将如何提高链路利用率，并处理可变数量的 GPU。通常在调度数据传输之前，我们首先将节点划分为完全连接的组，在每个组内执行 collective 通信，然后跨组进行同步。虽然可以在任何网络拓扑中实现此节点分组，但为了简单起见，我们将讨论基于 DGX-1 拓扑的 AllReduce 协议。
3.1 Blink 中的 AllReduce
基于 DGX-1 拓扑，我们的协议以分层方式进行。首先，我们将此拓扑划分为两组，每组包含 4 个完全连接的 GPU。我们使用内部广播在每个组内执行 reduce，然后跨组通过跨组转发进行通信。
内部广播：在内部广播阶段，我们使用 reduce-scatter 协议。如果，每个组中有 4 个 GPU，并将每个 GPU 上的数据分成 4 个块。然后，我们让 GPU1，GPU2，GPU3 将他们的第一块数据（即图 2（b）中的 B1，C1，D1）传输到 GPU0，这样在这个 reduce-scatter 步骤之后，GPU0 将拥有最终的第一块组内的结果（即 ABCD1）。同时，我们将所有第二块数据传输到 GPU1，以便 GPU1 具有第二块的最终结果，依此类推。在此阶段结束时，每个 GPU 将具有来自组内的最终结果的一个块。注意，即使当不使用组中的所有节点时（例如，当跨 6 个 GPU 执行 all-reduce 时），该方案也适用。由于来自完全连接组的节点的子集也完全连接，所以，我们可以复用上述相同的内部广播协议。
跨组转发：在跨组转发阶段，我们使用跨组链接来聚合其他组的内部结果，并在组内转发最终结果。在我们的示例中，GPU0 从另一个组中接收 GPU4 的结果（即图 2（c）中的 EFGH1），并和自身的结果（即 ABCD1）进行聚合，然后在组内转发最终结果的第一个数据块。整个过程用图 2（c）中带箭头的线表示。注意，可以同时执行该聚合和转发。
3.2 Benchmark
接下来，介绍上述 AllReduce 方案的性能基准。
我们使用 NVIDIA DGX-1 机器并在每个 GPU 上初始化 1GB 数据。随着 GPU 的数量从 2 变为 8，我们比较了 NCCL 2 和 Blink 实现的吞吐量。如图 3 所示，与 NCCL 2 相比，Blink 可以在 4 到 7 个 GPU 情况下实现 2 倍的吞吐量。吞吐量差异的主要原因是 Blink 可以使用我们的两步数据传输方案来利用更多 NVLink 链路。
4．未来工作与结论
本文介绍了 Blink，一个包含了一系列基于广播的 collective 通信原语的通讯库。与 NCCL 中基于环的数据传输相比，此方案可以实现更好的链路利用率、更高的吞吐量并在不同数量的 GPU 场景下依然有效。在当下硬件异构，以及跨机器 Collective 通信的环境中，我们将继续推广和拓展 Blink。