A Unified Architecture for Accelerating Distributed DNN Training in Heterogeneous GPU/CPU Clusters 论文阅读

发表于 更新于 分类于 阅读次数： 本文字数： 1.4k 阅读时长 ≈ 5 分钟

#1. 引言

近年来，为了训练更大的模型，深度学习训练集群在不断发展。这样的集群通常包括 GPU，CPU 和高速互联网络，例如高速以太网或者 Infiniband。

两个主要的架构是参数服务器（Parameter Server）和All-Reduce。他们都是基于数据并行（Data Parallel）的方法。

在 All-Reduce 架构中，只涉及 GPU 机器。在参数服务器架构中，CPU 和 GPU 机器都参与训练。

这两种架构在理论和实践中都有很大的差异。对于仅有 GPU 的集群，All-Reduce 是带宽最优的。然而，随着 CPU 和带宽资源的增加，All-Reduce 的最优性不再成立。

从理论上讲，使用额外的 CPU 机器来帮助 GPU 机器，参数服务器架构能够取得更好的性能。然而，现有的实现不够令人满意。

在本文中，作者提出了一种新的架构，称为BytePS，它通过仔细分配流量负载，统一了参数服务器和 All-Reduce 两种架构，能够适用于任意比例的不同 PCIe/NVLink 配置的 CPU/GPU 机器上。

#2. 背景

#2.1 分布式 DNN 训练

#2.2 All-Reduce

#2.3 参数服务器

#3. BytePS 架构

#3.1 动机

在 BytePS 之前，用户大多使用 All-Reduce 架构来训练 DNN 模型。作者观察到几个现象：

  1. 生产 GPU 集群通常有闲置的 CPU 和带宽。
  2. 现有的 All-Reduce 和 PS 实现都没有充分利用 CPU 和带宽。

作者的解决方案是 BytePS，它是一个统一的架构，可以利用空闲的 CPU 和带宽来加速训练。

#3.2 架构概述

BytePS 有两个主要模块：通信服务（Communication Service）和汇总服务（Summation Service）。

汇总服务 SS 运行在所有类型的机器上。

通信服务 CS 只运行在 GPU 机器上，负责在多个本地 GPU 之间同步 Tensor，并且与 SS 进行通信。

#4. BytePS 通信设计

#4.1 机器间通信

在 BytePS 中，所有通信都通过通信服务（CS）进行。假设网络具有全双工带宽，而且 DCQCN 类似的流量控制算法可以在网络中使用。

CPU 机器的网络流量等于其上 SS 的总和。GPU 机器的网络流量等于其上 SS 和 CS 的总和。

为了尽量减少通信时间，BytePS 给每个 CPU 上的 SS 分配 MssCpu 字节的求和负载。给每个 GPU 上的 CS 分配 MssGpu 字节的求和负载。

MssCpu=2(n−1)n2+kn−2kMMssGpu=(n−k)n2+kn−2kM\text{MssCpu} = \frac{2(n-1)}{n^2+kn-2k}M \\ \text{MssGpu} = \frac{(n-k)}{n^2+kn-2k}M

在实践中，DNN 模型的参数大小通常是不固定的。作者将参数分割成不超过 4MB 的小块。然后，所有 CS 将这些小块哈希到[0,n2+kn−2k)[0, n^2+kn-2k)的范围内。根据哈希值，CS 发送和接收对应的数据给 SS。

#4.2 机器内通信

#4.2.1 PCIe 机器

作者生产环境的机器模型是纯 PCIe 拓扑的，每个 CPU 直连 4 块 GPU，两个 CPU 之间通过 QPI 链接。使用现有的框架（例如 TensorFlow PS，MXNet PS，Horovod），在进行 Reduce 或者 Reduce-Scatter 时，会导致跨 PCIe 的内存拷贝，性能很差。

BytePS 考虑到 PCIe 拓扑的特点，允许同一个 PCIe 下的 GPU 先进行一个求和操作，然后再拷贝到 CPU 进行全局的求和操作，称为 CPU-assisted aggregation。完整步骤如下：

  1. Reduce-Scatter
  2. GPU->CPU 拷贝
  3. CPU-Reduce
  4. Networking
  5. CPU->GPU 拷贝
  6. All-Gather

实践结果表明 CPU-assisted aggregation 比单纯的 Ring-AllReduce 要好。

作者数据中心的另一种机器模型是带有 NVLink 连接的 GPU 机器。每台机器有 4 个 PCIe Switch，每个 Switch 连接 2 个 GPU，GPU 之间通过 NVLink 连接。NVLink 的带宽远高于 PCIe。

因为有 NVLink 的存在，GPU 之间的通信不需要占用 PCIe 带宽。然而，由于 NVLink 不是完全对称的，仍然存在 PCIe 带宽的瓶颈。

#4.2.3 讨论

以上分析表明，最佳的方案与机器的拓扑结构密切相关。

尽管存在差异，但作者总结了两个原则：

  1. 当两个 GPU 不在同一个 PCIe 交换机下时，始终避免直接的 GPU 到 GPU 内存复制，因为在实践中它很慢。
  2. 尽量减少在由 GPU 和 NIC 共享的 PCIe 交换机到 CPU 链路上的数据流量。

作者提出了以下最佳实践程序。让 Sn 表示具有 GPU 和 NIC 的 PCIe 交换机的数量，Sg 表示仅具有 GPU 的 PCIe 交换机的数量。

  1. 如果 Sn > 0，Sg > 0，CS 应该使用 Reduce 和 Broadcast，就不会竞争网卡带宽。
  2. 如果 Sn = 0 或者 Sg = 0，CS 应该使用 Reduce-Scatter 和 All-Gather 来平衡 PCIe 交换机上的流量。

多网卡情况同理。

GDR 技术(GPU-direct RDMA)可以减少 PCIe 流量，然而需要 GPU 和 RDMA 网卡在同一个 PCIe 交换机上。因此没有什么好处。

构建一个分析器来自动选择最佳的通信模式是一个有趣的研究方向。

#5. 汇总服务 SS

CPU 的内存带宽在某些时候可能会成为瓶颈。

作者提出了一个新的解决方案，将计算密集型的汇总操作分配给 GPU，并仅在 CPU 上执行 CPU 擅长的汇总操作。

虽然汇总操作提升了训练性能，但是打破了 PS 的异步特性。为此，作者提出了一个新的异步汇总算法。