参数服务器

一、参数服务器的背景

当我们从单个 GPU 训练模型，升级到多个 GPU，再到多台服务器（可能跨机架、跨交换机）训练时，分布式并行训练的复杂度会大大提升。一方面是不同设备间的连接带宽差异极大：比如 NVLink 能达到 100GB/s 的带宽，PCIe4.0 是 32GB/s，而 100GbE 以太网只有约 10GB/s；另一方面，要求开发者既懂统计学习建模，又精通系统和网络，这不太现实。

参数服务器的想法最早是在分布式隐变量模型里提出的，后来逐渐完善了 Push 和 Pull 的语义，还有对应的系统和开源库，它的作用就是帮我们解决分布式训练里的参数同步、梯度聚合这些复杂的问题，提升计算效率。

二、数据并行训练

在分布式训练里，数据并行是比较简单实用的策略，现在除了图深度学习，大部分场景都推荐用这个。

1. 基本的多 GPU 数据并行

之前我们学过的多 GPU 数据并行，是把所有梯度都聚合到某一个 GPU（比如 GPU 0）上，更新参数后再把新参数广播给所有 GPU。其实聚合的位置不一定是 GPU 0，也可以是 CPU，甚至可以把不同的参数梯度分到不同 GPU 上聚合。

2. 不同同步策略的时间差异

不同的参数同步策略，耗时差别很大。举个例子：假设所有梯度一共 160MB。

• 如果把 3 个 GPU 的梯度都传到第 4 个 GPU 上聚合，再把更新后的参数传回来，总共需要 60 毫秒（每次传 160MB，带宽 16GB/s，每次传输 10 毫秒，3 次传过去，3 次传回来，一共 60 毫秒）。

• 如果把所有数据都传到 CPU 上聚合，需要 80 毫秒，因为 4 个 GPU 都要把数据传到 CPU，耗时更多。

• 如果把梯度分成 4 块，每块 40MB，在不同 GPU 上同时聚合，只需要 15 毫秒，因为 PCIe 交换机可以全带宽工作，传输时间大大缩短。

另外，还有个小技巧：在深度网络里，计算梯度是需要时间的，我们可以在还没算完所有梯度的时候，就先把已经算好的参数梯度同步了，这样能节省时间，Horovod 这个工具就用到了这个思路。

三、环同步（Ring Synchronization）

现在的深度学习硬件，比如 AWS p3.16xlarge 和 NVIDIA DGX-2，用了很多定制的网络连接，比如 NVLink，它的带宽比 PCIe 高很多（每个 NVLink 双向传输速度约 18GB/s，比 PCIe 的 16GB/s 快）。

1. 环同步的思路

研究发现，把这些 GPU 组成环来同步数据是最优的。比如 8 个 GPU，可以分成两个环：一个是 1-2-3-4-5-6-7-8-1，这个环用的是双 NVLink 带宽；另一个是 1-4-6-3-5-8-2-7-1，是常规带宽。

2. 高效的环同步方法

如果直接在环里逐个传递梯度，比如从第一个 GPU 传到第二个，第二个把本地梯度和传来的梯度相加再传到第三个，这样聚合梯度的时间会随着节点数增加而线性增长，效率很低。

但如果我们把梯度分成 n 块（n 是 GPU 数量），每个 GPU 负责同步自己的那一块，比如节点 i 同步块 i，这样总时间就几乎不会随着环的大小增加而变长。比如 4 个 GPU 的情况，每个 GPU 开始向左邻居传部分梯度，最后在右邻居那里拿到聚合后的梯度。

还是用 160MB 的例子，8 个 V100 GPU 用环同步，只需要约 6 毫秒，比用 PCIe 总线快很多，不过实际中因为深度学习框架的限制，可能会比这个时间稍长一点。

这里要注意，环同步其实和树算法本质上差别不大，只是同步的路径更精细一点。

四、多机训练

当训练规模扩大到多台机器时，新的问题就来了：机器之间的连接带宽比较低，可能比机器内部慢一个数量级，而且不同机器的训练速度有点差别，要同步这些机器很麻烦。

1. 多机多 GPU 训练的流程

多机多 GPU 的训练流程是这样的：

1. 每台机器读取不同的批量数据，把数据分到本机的多个 GPU 显存里。

2. 每个 GPU 用自己的数据计算预测和梯度。

3. 把本机所有 GPU 的梯度聚合到一个（或多个）GPU 上。

4. 把本机的梯度传到本机的 CPU 里。

5. 所有机器的 CPU 把梯度传到中央参数服务器，服务器聚合所有梯度。

6. 用聚合后的梯度更新参数，再把新参数广播回各个 CPU。

7. 把新参数传到本机的 GPU 里。

8. 所有 GPU 更新参数，完成一次训练。

2. 参数服务器的瓶颈与解决方法

如果只用一个中央参数服务器，它的带宽会成为瓶颈，因为有 m 个工作节点的话，把所有梯度传到服务器的时间是 O (m)。解决方法是用多个参数服务器，每个服务器只存储一部分参数（比如 1/n 的参数，n 是参数服务器数量），这样更新和优化的总时间就变成 O (m/n)，伸缩性更好。实际中，我们可以让同一台机器既做工作节点，又做参数服务器。

五、键值存储

分布式多 GPU 训练的实现很复杂，所以我们可以用键值存储这个抽象来简化操作。

1. 梯度聚合的特点

梯度的聚合是交换归约，也就是不管先聚合哪个梯度，结果都一样，而且不同的梯度之间是独立的。这就意味着我们不需要精细控制什么时候接收哪个梯度。

2. Push 和 Pull 操作

我们可以定义两个操作：

• Push（key，value）：把工作节点上的梯度值传到公共存储里，在那里把这些值聚合（比如相加）。

• Pull（key，value）：从公共存储里拿到聚合后的梯度值。

这样做的好处是，把统计建模人员和系统工程师的工作分开了：建模人员只需要用简单的 Push 和 Pull 来表达优化，不用管复杂的同步细节；系统工程师负责处理分布式同步的复杂问题。

六、小结

1. 参数同步的效率很依赖网络基础设施和服务器内部的连接，不同的同步策略耗时差别很大。

2. 环同步在 p3 和 DGX-2 服务器上是最优的，但不是所有服务器都适合。

3. 用多个参数服务器可以增加带宽，分层同步策略的效果不错。

（注：文档部分内容可能由 AI 生成）