Microsoft DeepSpeed 实现了 BERT 最快的训练时间
好消息！DeepSpeed 创造了 BERT 训练速度的新纪录：在 1024 个 NVIDIA V100 GPU 上仅需 44 分钟。 这比最佳已发布结果提高了 30%，最佳已发布结果是在相同数量和代的 GPU 上进行端到端训练，需要 67 分钟才能达到相同的准确率。这一改进并非以过度使用硬件资源为代价，而是来自软件效率的提高。例如，DeepSpeed 在 NVIDIA V100 GPU 上可以实现惊人的 64 teraflops 单 GPU 性能，超过硬件峰值的 50%。
在这篇博文中，我们将讨论四项技术改进，它们使 DeepSpeed 能够实现这项打破纪录的 BERT 训练时间。
- 高度优化的 Transformer 内核以提高计算效率
- 通过异步预取队列将 I/O 与计算重叠
- 稀疏输出处理以消除浪费的计算
- 层归一化重新排序以提高训练稳定性和更快收敛
这些优化不仅对 BERT 有益；它们也适用于许多其他基于 Transformer 的模型，例如 RoBERTa、XLNet 和 UniLM。此外，除了针对预训练提到的改进之外，DeepSpeed 还能为下游任务（例如 Bing-BERT SQuAD 的微调）提高高达 1.5 倍的速度。
BERT 预训练的性能结果
与 SOTA 相比，DeepSpeed 显着提高了基于 Transformer 的模型（如 BERT）的单 GPU 性能。图 1 显示了通过 DeepSpeed 优化的 BERT-Large 训练的单 GPU 吞吐量，与来自 NVIDIA BERT 和 Hugging Face BERT 的两个知名 PyTorch 实现进行了比较。DeepSpeed 的吞吐量分别达到了 64 和 53 teraflops（对应于每秒 272 和 52 个样本），序列长度分别为 128 和 512，与 NVIDIA BERT 相比，吞吐量提高了 28%，与 HuggingFace BERT 相比，吞吐量提高了 62%。我们还支持高达 1.8 倍的更大批次大小，而不会出现内存不足。
为了实现这种性能，DeepSpeed 实施了一个随机 Transformer，它表现出一定程度的非确定性噪声，而不会影响整体收敛。此外，DeepSpeed 还实现了一个确定性 Transformer 内核，该内核完全可重现，但以平均约 2% 的性能回归为代价。用户可以根据自己的使用场景轻松地在两个版本之间进行选择和切换：随机版本追求最终的训练性能目标，而确定性版本可以通过更好地促进实验和调试来节省开发时间。我们在图 1 中报告了这两个内核的性能指标。所有批次大小和配置的性能指标都是使用 10 的梯度累积步骤收集的，因为在实际场景中使用的总体批次大小平均范围从几百到几千不等。
图 1：BERT-Large 在单个 V100 GPU 上的性能评估，将 DeepSpeed 与 NVIDIA 和 HuggingFace 版本的 BERT 在混合序列长度训练中进行比较。标记的点显示了每个实现的最高吞吐量（以 teraflops（Tflops）为单位）。DeepSpeed 提高了吞吐量，并允许在不出现内存不足的情况下使用更大的批次大小。
查看跨 GPU 的分布式训练，表 1 显示了我们使用 16 到 1024 个 GPU 的端到端 BERT-Large 预训练时间（SQUAD 的 F1 分数为 90.5）。我们使用 1024 个 V100 GPU（64 个 NVIDIA DGX-2 节点）在 44 分钟内完成了 BERT 预训练。相比之下，NVIDIA 的先前 SOTA 使用 1472 个 V100 GPU 需要 47 分钟。DeepSpeed 不仅更快，而且使用资源减少了 30%。使用相同的 1024 个 GPU，NVIDIA BERT 使用相同的 1024 个 GPU 需要 67 分钟 [1] BERT，而 DeepSpeed 需要 44 分钟，将训练时间缩短了 30%。同样，在 256 个 GPU 上，NVIDIA BERT 需要 236 分钟，而 DeepSpeed 需要 144 分钟（快了 39%）。
| 节点数量 | V100 GPU 数量 | 时间 |
|---|---|---|
| 1 个 DGX-2 | 16 | 33 小时 13 分钟 |
| 4 个 DGX-2 | 64 | 8 小时 41 分钟 |
| 16 个 DGX-2 | 256 | 144 分钟 |
| 64 个 DGX-2 | 1024 | 44 分钟 |
表 1：使用 DeepSpeed 在 1 到 64 个 DGX-2 上的 BERT-Large 训练时间。
在最近的 GTC 2020 上，NVIDIA 宣布了下一代硬件 A100，该硬件现在提供了比 V100 GPU 高 2.5 倍的硬件峰值性能。假设 A100 GPU 允许我们获得与我们在 V100 GPU 上获得的相同百分比的硬件峰值性能（50%），我们预计通过将我们的软件优化与新硬件相结合，我们将获得更高的吞吐量。我们预计这将进一步缩短 BERT 的训练时间，在 1024 个 A100 GPU 的集群上少于 25 分钟。
微调任务的性能结果
除了我们针对预训练显示的性能优势之外，我们还评估了我们自定义内核在下游任务微调方面的性能。表 2 和 3 显示了使用 PyTorch 和 DeepSpeed Transformer 内核，在 NVIDIA V100 上使用 16 和 32 GB 内存运行 Bing-BERT SQuAD 时获得的每秒样本数。对于 16 GB V100，我们可以在支持每个 GPU 大两倍的批次大小的同时，实现高达 1.5 倍的速度提升。另一方面，我们可以使用 32 GB 内存支持高达 32 的批次大小（比 Pytorch 大 2.6 倍），同时为端到端微调训练提供 1.3 倍的速度提升。请注意，我们使用最佳每秒样本数来计算 PyTorch 出现内存不足 (OOM) 情况下的速度提升。
| 微批次大小 | PyTorch | DeepSpeed | 加速（x） |
|---|---|---|---|
| 4 | 36.34 | 50.76 | 1.4 |
| 6 | OOM | 54.28 | 1.5 |
| 8 | OOM | 54.16 | 1.5 |
表 2. 在 NVIDIA V100 (16 GB) 上运行 SQuAD 微调时的每秒样本数，使用 PyTorch 和 DeepSpeed Transformer 内核。
| 微批次大小 | PyTorch | DeepSpeed | 加速（x） |
|---|---|---|---|
| 4 | 37.78 | 50.82 | 1.3 |
| 6 | 43.81 | 55.97 | 1.3 |
| 12 | 49.32 | 61.41 | 1.2 |
| 24 | OOM | 60.70 | 1.2 |
| 32 | OOM | 63.01 | 1.3 |
表 3：在 NVIDIA V100 (32 GB) 上运行 SQuAD 微调时的每秒样本数，使用 PyTorch 和 DeepSpeed Transformer 内核。
BERT 高度优化的 Transformer 内核
GPU 的峰值浮点吞吐量非常高，但大多数框架实现中的默认 Transformer 模块远未达到这个峰值。图 2 显示了 Transformer 模块的结构，层归一化放置在两个子层（注意力和前馈）的输入流上。为了接近 GPU 的峰值性能，我们在自己的 Transformer 内核实现中采用了两种优化方法：高级融合和可逆算子。
图 2：具有预层归一化架构的 Transformer 层
(a) 先进的融合内核以减少数据移动
我们观察到，基于 Transformer 的网络会触发许多以生产者-消费者方式运行的 CUDA 内核调用，这会增加将数据传输到全局内存和从全局内存传输数据的成本，以及内核启动的开销。现有的基于编译器的方法执行细粒度融合（例如，元素级运算的融合），导致错过融合机会。相比之下，我们充分利用了针对 Transformer 模块量身定制的细粒度和粗粒度融合。
QKV 和各种融合。 我们将三个查询 (Q)、键 (K) 和值 (V) 权重矩阵合并，以调度更大的 QKV GEMM，以暴露更多并行性并提高 GPU 共享内存和寄存器文件上的数据局部性，如图 3 所示。接下来，我们将 QKV 输出矩阵的数据布局转换与偏差添加结合起来。然后，我们将大型 QKV 矩阵划分为三个转换后的矩阵，用于以下自注意力计算。
如图 3 所示，我们以连续行读取 QKV 矩阵（由红色方框表示），并将它们写入三个转换后的 Q、K 和 V 矩阵中。由于每个矩阵都从不同的偏移量开始，因此我们可能无法对主内存进行合并访问。因此，我们使用共享内存作为中间缓冲区，以重新排列数据，以便我们可以将数据放在内存的连续部分中。即使我们在访问共享内存时产生了非合并模式，我们也减少了对主内存进行非合并访问的成本，以更好地利用内存带宽，从而在端到端训练中实现 3% 到 5% 的性能提升。
图 3：QKV 的 GEMM 和变换内核融合
我们执行额外的融合，例如将注意力输出 GEMM 中的偏差添加与残差连接和 dropout 的添加合并，这允许在寄存器文件和共享内存中进行访问，这些访问比昂贵的写回全局内存快得多。
warp 级通信。 为了减轻并行 GPU 内核之间的同步开销，并进一步提高融合内核的资源利用率，我们使用 warp 级（数据混洗指令）而不是默认的跨 warp 通信。以层归一化和 SoftMax 内核为例，我们在 warp 内执行每个归约操作，同时将不同的归约分布在不同的 warp 中。这样，我们减轻了并行线程之间的同步，并进一步提高了 GPU 资源利用率。
随机内核与确定性内核。 DL 训练通常对一定程度的随机性具有鲁棒性，在某些情况下，诸如 dropout 之类的受控噪声充当正则化器，从而提高泛化能力。在设计我们的 Transformer 内核时，我们接受了一定程度的随机性，以允许在内核中存在有限的数据竞争条件来提高吞吐量：我们利用隐式 warp 同步编程来提高 warp 级合作操作的性能 [3]。缺乏明确的 warp 级同步充当非确定性噪声，而不会影响 Transformer 内核的整体收敛行为，同时提供相当大的吞吐量提升。
此外，DeepSpeed 还实现了一个非随机 Transformer 内核，该内核具有显式 warp 同步，它产生确定性结果，但以轻微的性能回归为代价。用户可以根据自己的使用场景轻松地在两个版本之间进行选择和切换：随机版本追求最终的训练性能目标，而确定性版本可以通过更好地促进实验和调试来节省开发时间。
在我们的实验中，我们对预训练 BERT 使用随机内核，而对微调使用非随机内核以获得完全可重现的结果。我们建议对涉及大量数据的训练任务（如预训练）使用随机内核，而在使用有限数据的训练时（如微调）使用非随机版本，以获得更一致的结果。
经济高效的重计算。 在融合不同操作的内核时，我们观察到某些操作计算起来不费力，但会产生昂贵的数据移动成本，例如偏差添加和 dropout。对于这些操作，我们避免在正向传递中保存它们的结果，而是改为在反向传递期间重新计算它们，这比将它们的结果写入主内存并从主内存中重新加载快得多。
(b) 可逆算子以节省内存并运行大型批次
我们还观察到 Transformer 块中几个运算符的中间激活会导致大量的内存消耗，例如 SoftMax 和 Layer Norm。对于这些运算符，我们利用它们是可逆函数这一事实来丢弃这些层的输入，从而减少激活内存的占用。可逆函数是指其反向传播独立于输入，并且可以仅基于输出进行公式化 [2]。图 4 和图 5 展示了 PyTorch 中 SoftMax 和 Layer-Norm 的原始实现与 DeepSpeed 中可逆 SoftMax 实现的示例。通过此优化，我们能够将运算符的激活内存减少一半，而减少的内存使我们能够以更大的批量大小进行训练，这再次提高了 GPU 效率。
图 4：DeepSpeed 可逆 SoftMax 运算与默认 PyTorch SoftMax 运算
图 5：DeepSpeed 可逆 LayerNorm 运算与默认 PyTorch LayerNorm 运算
通过异步预取队列将 I/O 与计算重叠
除了高度优化的 Transformer 内核之外，BERT 训练还存在其他性能限制因素，例如数据加载。我们开发了自己的异步工作程序，它仅在“安全点”将数据批次预取到队列中，这些“安全点”是 CPU 空闲时的点（例如，异步启动正向传播之后）。通过这种方式，我们确保在 CPU 端进行计算时，不会从 CPU 到 GPU 进行出队和数据复制。这与默认的 PyTorch 数据加载器不同，后者可以在任何点预取数据，并导致性能干扰。通过使用这种方法，我们隐藏了几乎所有 I/O 负担，这占原始训练时间的 4%。
利用 BERT 输出处理的稀疏性
通过识别和利用 BERT 输出处理中的稀疏性，我们将端到端训练时间提高了 5.4%。输出处理涉及两个步骤：i）使用矩阵乘法，将 BERT 从最终 Transformer 层的隐藏输出维度投影到语言词汇表，以及 ii）对掩盖的输出标记进行交叉熵，以获得每个序列的预测误差。第一步的成本与词汇量大小、隐藏输出维度和序列长度成正比，可能与 Transformer 层计算一样昂贵，甚至更昂贵。但是，只有大约 15% 的标记被掩盖，我们只需要掩盖标记的交叉熵。因此，投影可以作为有效的稀疏计算来完成。为此，我们在进行投影之前丢弃与非掩盖标记相对应的最终 Transformer 层的行，将输出处理的计算成本降低了 85%。
预层归一化与后层归一化架构
我们观察到，在较大的批量大小（例如 64K）下，默认的 BERT 预训练会遇到训练不稳定，这可能导致模型发散或收敛到糟糕的/可疑的局部最优。进一步调查表明，默认的 BERT 存在梯度消失问题。为了缓解这个问题，我们更改了 LayerNorm 的放置位置（后 LayerNorm），将其仅放置在 Transformer 块中子层的输入流上（称为前 LayerNorm），这是最近几项神经机器翻译工作的修改。前 LayerNorm 带来了一些有用的特性，例如避免梯度消失、稳定优化和性能提升。它使我们能够以 64K 的聚合批量大小进行训练，并提高学习率和更快的收敛速度。
要试用这些优化和训练配方，请查看我们的 BERT 训练教程 以及 DeepSpeed GitHub 仓库中的源代码 DeepSpeed GitHub 仓库.
参考文献
[1] “NVIDIA 创造了世界上最快的 BERT 训练时间和基于 Transformer 的最大模型，为高级对话式 AI 铺平道路” https://devblogs.nvidia.com/training-bert-with-gpus/.
[2] S. R. Bulo, L. Porzi, 和 P. Kontschieder，“用于 DNN 内存优化训练的原地激活批归一化” 2017。 http://arxiv.org/abs/1712.02616.
[3] Mark Harris 和 Kyrylo Perelygin，“协作组：灵活的 CUDA 线程编程”，https://devblogs.nvidia.com/cooperative-groups/.