DeepSpeed 与 1 位 Adam：通信量减少 5 倍，训练速度提升 3.4 倍

2020 年 9 月 8 日

1. 概述

大型模型（如 BERT 和 GPT-3）的可扩展训练需要基于模型设计、架构和系统功能的仔细优化。从系统的角度来看，通信已成为主要瓶颈，尤其是在标准 TCP 互连（提供有限网络带宽）的通用系统上。通信压缩是在此类系统上减少训练时间的重要技术。压缩通信最有效的方法之一是误差补偿压缩，它即使在 1 位压缩下也能提供稳健的收敛速度。然而，最先进的误差补偿技术仅适用于随机梯度下降 (SGD) 和动量 SGD 等基本优化器，它们与梯度呈线性关系。它们不适用于 Adam 等非线性梯度优化器，Adam 在包括 BERT 类模型训练在内的许多任务中提供最先进的收敛效率和准确性。对于像 ADAM 这样强大的优化器，对梯度（方差项）的非线性依赖性使得开发基于误差补偿的压缩技术极具挑战性，限制了最先进的通信压缩技术的实际价值。

1.1 背景：经典压缩技术

通信压缩的一种方法是 1 位压缩，它可以表示为

$x\to \frac{\|x\|}{\|Sign(x)\|}Sign(x)$

通过这种压缩，我们可以通过使用一位来表示每个数字，从而将内存大小减少 32 倍。问题是使用这种直接的方法会显着降低收敛速度，这使得这种方法不可行。为了解决这个问题，最近的研究表明，通过使用误差补偿压缩，我们可以预期在通信压缩的情况下获得几乎相同的收敛速度。误差补偿的概念可以概括为：1) 进行压缩，2) 记住压缩误差，然后 3) 在下一轮迭代中将压缩误差加回去。对于 SGD，进行误差压缩会导致

$x_t= x_{t-1} - \gamma C(g_t + e_{t-1}), \quad e_t = g_t+e_{t-1}-C(g_t+e_{t-1} )$

其中 C(⋅) 是 1 位压缩运算符。进行这种误差补偿的好处是，历史压缩误差 (e_t 和 e_(t-1)) 最终会自行抵消，这可以通过以下公式看出

$x_t=x_{t-1}-\gamma(g_t+e_{t-1}-e_t )$

这种策略已被证明适用于与梯度呈线性关系的优化算法，如 SGD 和动量 SGD。

1.2 将误差补偿应用于 Adam 的挑战

我们提供了 Adam 算法的概述。更新规则如下。

$m_{t+1}=\beta_1 m_t+(1-\beta_1 ) g_t$

$v_{t+1}=\beta_2 v_t+(1-\beta_2 ) (g_t )^2$

$x_{t+1}=x_t-\gamma \frac{m_{t+1}}{\sqrt{v_{t+1}} +\eta}$

如上述公式所示，方差项 v_t 与梯度 g_t 呈非线性关系。如果我们将基本的误差补偿压缩应用于 Adam，我们会发现 Adam 不会收敛，如图 1 所示。

Inapplicability of Error-compensation Compression for Adam due to non-linear dependence on the gradient

图 1：由于对梯度的非线性依赖，误差补偿压缩不适用于 Adam

2. 使用 1 位 Adam 压缩通信

为了在使用 Adam 优化器时压缩通信，我们开发了 1 位 Adam，它通过预处理解决了梯度中的非线性问题。我们观察到，在训练了几轮后，非线性项方差 (v_t) 上的变化幅度显著降低，之后将 v_t 设置为常数不会改变收敛速度。提出的 1 位 Adam 优化器（如图 2 所示）由两部分组成：预热阶段，本质上是 vanilla Adam 算法；压缩阶段，它使方差项保持恒定，并将剩余的线性项（即动量）压缩成 1 位表示。

算法的压缩阶段由阈值参数控制（如图 2 所示）。当我们检测到“方差”的变化低于某个阈值时，我们就切换到压缩阶段。我们的研究表明，预热阶段只需要整个训练步骤的 15%-20%。

Comparison of distributed training steps in classic Adam and the proposed 1-bit compressed Adam algorithm

图 2：经典 Adam 和提出的 1 位压缩 Adam 算法的分布式训练步骤对比

2.1 1 位 Adam 的工作原理

1 位 Adam 的权重更新规则由以下公式控制。

对于第 i 个工作节点，在压缩阶段

$m_{t+1}^{(i)}=\beta_1 m_t+(1-\beta_1 ) g_t^{(i)}$

$\widehat{m}_{t+1}^{(i)}=C(m_{t+1}^{(i)}+e_t^{(i)}), \quad e_{t+1}^{(i)}=(m_{t+1}^{(i)}+e_t^{(i)} )-\hat{m}_{t+1}^{(i)}$

$m_{t+1}^{(ave)} = \frac{1}{n}\sum_{i=1}^n \hat{m}_{t+1}^{(i)}$

$\hat{m}_{t+1}^{(ave)}=C(m_{t+1}^{(ave)}+e_t^{(ave)} ),\quad e_{t+1}^{(ave)}=(\hat{m}_{t+1}^{(ave)}+e_t^{(ave)} )-\hat{m}_{t+1}^{(ave)}$

$m_{t+1}=\hat{m}_{t+1}^{(ave)}$

$x_{t+1}=x_t-\gamma \frac{m_{t+1}}{\sqrt{v_{warmup}}+\eta}$

其中 x_t 是迭代 (t-1) 后的模型，m_t^(i)，e_t^(i) 是迭代 (t-1) 后工作节点 i 上的动量和压缩误差，v_warmup 是预热阶段后的方差项。

2.2 解决 1 位 Adam 的系统挑战

除了算法挑战之外，在训练系统中应用 1 位 Adam 还存在两个系统挑战。首先，我们需要高效的内核来将动量转换为 1 位表示。其次，我们需要高效的通信方案来跨不同 GPU 交换这种压缩后的动量。压缩的目的是减少总的训练时间，以便可以使用带宽受限互连的通用系统来训练大型模型。我们在 DeepSpeed 中解决了这些挑战，并为在通信受限系统上进行训练引入了完全优化的 1 位 Adam 实现。

3. 1 位 Adam 在通信受限系统上的优势

1 位 Adam 提供与 Adam 相同的收敛速度，通信量减少了 5 倍，使 BERT-Large 预训练的吞吐量提高了 3.5 倍，SQuAD 微调的吞吐量提高了 2.7 倍。这种端到端吞吐量的提升得益于在压缩阶段观察到的 6.6 倍（图 3）和 6.2 倍（图 4）的加速。值得一提的是，我们的 1 位 Adam 优化器在 40 吉比特以太网系统上具有如此好的可扩展性，以至于它的性能可以与 Adam 在 40 吉比特 InfiniBand QDR 系统上的可扩展性相媲美。我们注意到，根据 iperf 基准测试，40 吉比特以太网的有效带宽为 4.1 Gbps，而 InfiniBand 根据 InfiniBand perftest 微基准测试提供接近峰值的 32 Gbps 带宽。

BERT-Large Pretraining

图 3：1 位 Adam 在 V100 GPU 上进行 BERT-Large 预训练的可扩展性，批大小为 16/GPU。

SQuAD Finetuning

图 4：1 位 Adam 在 V100 GPU 上进行 SQuAD 微调的可扩展性，批大小为 3/GPU。

4. 深入探讨 1 位 Adam 的评估结果

与 Adam 相同的收敛速度

使用 1 位 Adam 的一个主要问题是收敛速度，我们发现 1 位 Adam 可以使用与 Adam 相同数量的训练样本实现相同的收敛速度和可比的测试性能，如图 5 所示。

1-bit Adam convergence

图 5：1 位 Adam 使用与 Adam 相同数量的训练样本实现类似的收敛速度。

表 1 显示了 BERT-Base 和 BERT-Large 的详细结果。我们可以看到，无论是未压缩情况还是压缩情况，分数都与原始模型相当或更好。

1-bit Adam convergence table

表 1：在各种测试任务上验证 1 位 Adam 的正确性

通信量减少了 5 倍：1 位 Adam 提供与 Adam 相同的收敛速度，并在 16 位 (FP16) 训练的压缩阶段将通信量减少了 16 倍。对于 BERT 预训练，这会导致总的通信量减少 5 倍，因为我们观察到预热阶段只占端到端训练时间的 15%。

计算原始 Adam 与 1 位 Adam 的通信量比率的公式如下

1 / (warmup + (1 – warmup)/16)

如果预热阶段占 15%，那么原始 Adam 的通信量是 1 位 Adam 的 5 倍。

1 位 Adam 训练 BERT-Large 的速度提升 3.5 倍

我们展示了在具有两种不同带宽受限互连的系统上训练 BERT-Large 的两个主要结果：1) 40 吉比特以太网（图 5）和 2) 40 Gbps InfiniBand QDR（图 6）。在压缩阶段，我们在以太网系统上观察到高达 6.6 倍的吞吐量提升，在 InfiniBand 系统上观察到高达 2 倍的吞吐量提升，这导致端到端加速（包括预热和压缩阶段）分别达到 3.5 倍和 2.7 倍。1 位 Adam 的主要优势来自于通信量的减少（得益于我们的压缩动量交换）以及我们自定义的 allreduce 操作，它使用非阻塞 gather 操作（然后是 allgather 操作）实现了高效的 1 位通信。

需要注意的是，也可以使用 LAMB 等优化器（而不是 Adam）来增加总批大小，从而减少 BERT 预训练的通信量。但是，1 位 Adam 避免了对超参数进行严格的调整，根据我们的经验，对于大型批次来说，这通常更困难。此外，1 位 Adam 也非常适合批大小限制较小（不能很好地收敛于大批大小）的工作负载，比如许多微调任务。

Performance of 1-bit Adam for BERT-Large training on 40 gbps Ethernet

图 5：1 位 Adam 在 40 Gbps 以太网互连上进行 BERT-Large 训练的压缩阶段性能。

Performance of 1-bit Adam for BERT-Large training on 40 gbps InfiniBand

图 6：1 位 Adam 在 40 Gbps InfiniBand 互连上进行 BERT-Large 训练的压缩阶段性能。

1 位 Adam 对 SQuAD 微调的速度提升 2.7 倍

1 位 Adam 不仅在大型训练任务上具有可扩展性，而且在 SQuAD 微调等任务上也具有可扩展性。如图 7 和图 8 所示，1 位 Adam 在以太网和 InfiniBand 系统上均具有良好的可扩展性，并在以太网系统上提供了高达 6.2 倍的吞吐量提升（在压缩阶段），从而导致端到端加速达到 2.7 倍（25% 预热加上 75% 压缩阶段）。对于 SQuAD 微调，我们观察到总批大小为 96 时 F1 分数最高。大于此值的批大小会降低收敛速度，并需要额外的超参数调整。因此，为了扩展到 32 个 GPU，我们只能应用每个 GPU 3-4 个的小批大小。这使得微调任务通信量很大，难以扩展。1 位 Adam 很好地解决了可扩展性问题，在不扩大批大小的情况下实现了 3.4 倍的通信量减少，并导致端到端加速达到 2.7 倍。

1-bit Adam convergence

图 7：1 位 Adam 在 40 Gbps 以太网上进行 SQuAD 微调的压缩阶段性能。

1-bit Adam convergence

图 8：1 位 Adam 在 40 Gbps InfiniBand 互连上进行 SQuAD 微调的压缩阶段性能。

Twitter Facebook LinkedIn