来源：知乎
作者：微软DeepSpeed
庞大 AI 模型正在改变数字世界。基于庞大语言模型 (LLM) 的 Turing-NLG、ChatGPT 和 GPT-4 等生成语言模型用途广泛，能够执行摘要、代码生成和翻译等任务。同样，DALL・E、Microsoft Designer 和 Bing Image Creator 等庞大多模态生成模型可以生成艺术、建筑、视频和其他数字资产，使内容创作者、建筑师和工程师能够探索全新的创意生产力。
然而，训练这些庞大模型需要在数百甚至数千个 GPU 设备上利用大量内存和计算资源。例如，训练 Megatron-Turing NLG 530B 模型需要利用超过 4,000 个 NVidia A100 GPU。有效地利用这些资源需要一个复杂的优化系统，以将模型合理分配到各个设备的内存中，并有效地并行化这些设备上的计算。同时，为了使深度学习社区能够轻松从事庞大模型训练，这些优化必须易于利用。
DeepSpeed 的 ZeRO 优化系列为这些挑战提供了强大的解决方案，并已广泛用于庞大深度学习模型例如 TNLG-17B、Bloom-176B、MPT-7B、Jurrasic-1 的训练中 。尽管它具有变革性的能力 ，在一些关键场景中，ZeRO 会在 GPU 之间产生大量数据传输开消，这降低了训练效率。这种情况特别发生在以下场景中：a) 全局 batch size 较小，而 GPU 数量多，这导致每一个 GPU 上 batch size 较小，需要频繁通讯；或者 b) 在低端集群上从事训练，其中跨节点网络带宽有限，导致高通讯延迟。在这些情况下，ZeRO 的训练效率会受到限制。
为了解决这些限制，我们发布了 ZeRO++ 。ZeRO++ 相比 ZeRO 将总通讯量减轻了 4 倍，而不会影响模型质量。这有两个关键意义：
1. ZeRO++ 加快庞大模型预训练和微调
每一个 GPU 上 batch size 较小时：无论是在数千个 GPU 上预训练庞大模型，还是在数百个甚至数十个 GPU 上对其从事微调，当每一个 GPU 的 batch size 较小时，ZeRO++ 提供比 ZeRO 高 2.2 倍的吞吐量，直接减轻训练时间和成本。
低带宽计算集群: ZeRO++ 使低带宽集群能够实行与带宽高 4 倍的高端集群类似的吞吐量。因此，ZeRO++ 可以跨更广泛的集群从事高效的庞大模型训练。
2. ZeRO++ 加快 ChatGPT 类的 RLHF 训练
虽然 ZeRO++ 主要是为训练而设计的，但它的优化也自动适用于 ZeRO-Inference，因为通讯开消对于 ZeRO 的训练和推理同样适用。因此，ZeRO++ 可以提高人类反馈强化学习 (RLHF) 等算法的效率，因为 RLHF 结合了训练和推理。
通过与 DeepSpeed-Chat 的集成，与原始 ZeRO 相比，ZeRO++ 可以将 RLHF 训练的生成阶段效率提高多达 2 倍，强化学习训练阶段效率提高多达 1.3 倍。
接下来，我们将更深入地解释 ZeRO 及其通讯开消，并讨论 ZeRO++ 中为解决这些问题而从事的关键优化。然后我们将展示 ZeRO++ 对不同模型巨细、批量巨细和带宽限制的训练吞吐量的影响。我们还将讨论 ZeRO++ 如何应用于 DeepSpeed-Chat，以加快利用 RLHF 的对话模型的训练。
ZeRO++ 详解
图2：ZeRO optimizer 工作流程图（此为部分展示，完整流程请看知乎原文）
ZeRO 是数据并行 (Data Parallelism) 的一种内存高效版本，其中模型状态会被分割储存在所有 GPU 上，而不需要在训练期间利用基于 gather/broadcas 的通讯从事复制和重建。这使 ZeRO 能够有效地利用所有设备的聚合 GPU 内存和计算力，同时提供简单易用的数据并行训练。
假设模型巨细为 M。在前向传播过程中，ZeRO 执行全收集 / 广播 (all-gather/broadcast) 操作以在需要之时为每一个模型层收集参数（总共巨细为 M）。在向后传递中，ZeRO 对每一层的参数采用类似的通讯模式来计算其局部梯度（总巨细为 M）。此外，ZeRO 在对每一个局部梯度计算完毕后会立刻利用 reduce 或 reduce-scatter 通讯从事平均和分割储存（总巨细为 M）。因此，ZeRO 总共有 3M 的通讯量，平均分布在两个全收集 / 广播 (all-gather/broadcast) 和一个减轻分散 / 减轻 (reduce-scatter/reduce) 操作中。
为了减轻这些通讯开消，ZeRO++ 从事了三组通讯优化，分别针对上述三个通讯集合：
图 3:qwZ 的分区量化图例
ZeRO 通讯过程中的权重量化 (qwZ)
首先，为了减轻 all-gather 期间的参数通讯量，我们采用权重量化在通讯前将每一个模型参数从 FP16（两个字节）动态缩小为 INT8（一个字节）数据类型，并在通讯后对权重从事反量化。然而，简单地对权重从事量化会降低模型训练的准确性。为了保持良好的模型训练精度，我们采用分区量化，即对模型参数的每一个子集从事独立量化。目前尚且没有针对分区量化的高性能现有实行。因此，我们自行从头开始实行了一套高度优化的量化 CUDA 内核，与基本量化相比，精度提高 3 倍，速度提高 5 倍。
图 4: 权重的分层分割存储 (hpZ)
ZeRO 模型权重的分层分割存储 (hpZ)
其次，为了减轻向后传递期间全收集 (all-gather) 权重的通讯开消，我们用 GPU 内存从事通讯。更具体地说，我们不像在 ZeRO 中那样将整个模型权重分布在所有机器上，而是在每台机器中维护一个完整的模型副本。以更高的内存开消为代价，这允许我们用机器内的模型权重全收集 / 广播 (all-gather/broadcast) 代替昂贵的跨机器全收集 / 广播 (all-gather/broadcast)，由于机器内通讯带宽更高，这使得通讯速度大幅提升。
图 5: qgZ 端到端的工作流程
ZeRO 通讯过程中梯度量化 (qgZ)
第三，要降低梯度的 reduce-scatter 通讯成本更具挑战性。因为直接应用量化来减轻通讯量是不可行的。即使我们利用分区量化来降低量化误差，梯度 reduce 也会累积并放大量化误差。为了解决这个问题，我们只在通讯之前量化梯度，但在任何 reduce 操作之前将它们反量化到原有精度。为了有效地做到这一点，我们发明了一种名为 qgZ 的基于 all-to-all 的新型量化梯度通讯范式，它在功能上等同于压缩的归约 – 分散 (reduce-scatter) 操作。
qgZ 旨在解决两个挑战：i) 如果我们简单地在 INT4/INT8 中实施 reduce-scatter 会导致显著精度损失，以及 ii) 在传统 tree 或 ring-based reduce-scatter 中利用量化需要一长串量化和反量化步骤，这直接导致误差积累和显著的延迟，即使我们在全精度上从事 reduce。为了解决这两个挑战，qgZ 不利用 tree 或 ring-based reduce-scatter 算法，而是基于一种新颖的分层 all-to-all 方法。
qgZ 中有三个主要步骤：
梯度切片重新排序；
节点内通讯和 reduce；
节点间通讯和 reduce。
首先，在任何通讯发生之前，我们对梯度从事切片并对张量切片重新排序，以保证通讯结束时每一个 GPU 上的最终梯度位置（即图 5 中的绿色块）是正确的。其次，我们量化重新排序的梯度切片，在每一个节点内从事 all-to-all 通讯，从 all-to-all 中对接收到的梯度切片从事反量化，并从事局部 reduce。第三，我们再次量化局部 reduce 后的梯度，从事节点间的 all-to-all 通讯，再次对接收到的梯度从事反量化，并计算最终的高精度梯度 reduce，得到图 5 中绿色块的结果。
这种分层方法的原因是为了减轻跨节点通讯量。更准确地说，给定每一个节点 N 个 GPU、M 的模型巨细和 Z 的量化比率，单跳 all-to-all 将生成 M*N/Z 跨节点流量。相比之下，通过这种分层方法，我们将每一个 GPU 的跨节点流量从 M/Z 减轻到 M/(Z*N)。因此，总通讯量从 M*N/Z 减轻到 M*N/(Z*N) = M/Z。我们通过重叠节点内和节点间通讯以及融合 CUDA 内核来进一步优化 qgZ 的端到端延迟（张量切片重新排序 (Tensor Slice Reordering)+ 节点内量化 (Intra-node quantization)）和（节点内反量化 (Intra-node Dequantization) + 节点内梯度整合 (Intra-node Reduction) + 节点间量化 (inter-node quantization)）。
通讯总量优化
通过结合以上所有三个组件，我们将跨节点通讯量从 3M 减轻到 0.75M。更具体地说，我们利用 qwZ 将模型权重的前向全收集 / 广播从 M 减轻到 0.5M。我们利用 hpZ 消除了反向传播期间的跨节点 all-gather，将通讯从 M 减轻到 0。最后，我们利用 qgZ 将反向传播期间的跨节点 reduce-scatter 通讯从 M 减轻到 0.25M。
ZeRO++ 加快庞大语言模型训练
在这里，我们展示了 ZeRO++ 在 384 个 Nvidia V100 GPU 上的真实 LLM 训练场景的测试结果。
图 6: 在 384 个 V100 GPU 上的各种模型巨细下 ZeRO++ 与 ZeRO 的吞吐量，节点间利用 4 个 Infiniband (IB) 从事互连，每一个以 100 Gbps 运行。
在 GPU 小 batch size 情况下 ZeRO++ 实行更高的训练效率
高带宽集群：如图 6 所示，我们首先展示了 ZeRO++ 相对于 ZeRO 的吞吐量改进，针对不同的模型巨细和微批量 (micro-batch size) 巨细，测试利用 4x Infiniband (IB) 以实行 400Gbps 跨节点互连带宽，每一个以 100Gbps 运行。在 micro-batch size 为每 GPU 1k tokens 时，ZeRO++ 比 ZeRO-3 的吞吐量提高了 28% 到 36%。对于 2k tokens micro-batch size 巨细，ZeRO++ 比 ZeRO-3 实行了 24% 到 29% 的吞吐量增益。
图 7: 在 384 个 V00 GPU 上 100Gbps 跨节点带宽时各种 LLM 的吞吐量
低带宽集群：在 100Gbps 等低带宽网络环境中，ZeRO++ 的性能明显优于 ZeRO-3。如图 7 所示，与 ZeRO-3 相比，ZeRO++ 在端到端吞吐量方面实行了高达 2.2 倍的加快。平均而言，ZeRO++ 比 ZeRO-3 基线实行了大约 2 倍的加快。
图 8: ZeRO++ 以显著降低的带宽实行高带宽集群性能
实行高带宽 ZeRO 和低带宽 ZeRO++ 集群之间的模型训练效率等效
此外，与 ZeRO 在高得多的带宽环境下相比，ZeRO ++ 可以在低带宽集群中实行相当的系统吞吐量。如图 8 所示，对于 18B 和 138B 模型巨细，具有 200Gbps 跨节点带宽的 ZeRO++ 可以达到与 800Gbps 跨节点带宽的 ZeRO-3 相似的 TFLOP。
鉴于 ZeRO++ 出色的可扩展性，我们将 ZeRO++ 视为用于训练庞大 AI 模型的下一代 ZeRO。
DeepSpeed-Chat 与 ZeRO++ 结合用于 RLHF 训练
RLHF 训练简介
ChatGPT 类模型由 LLM 提供支持，并利用 RLHF 从事微调。RLHF 由生成（推理）阶段和训练阶段组成。在生成阶段，演员 (actor) 模型将部分对话作为输入，并利用一系列前向传递生成响应。然后在训练阶段，评论 (critic) 模型根据质量对生成的响应从事排名，为演员模型提供强化信号。利用这些排名对参与者模型从事微调，使其能够在后续迭代中生成更准确和适当的响应。
RLHF 训练带来了巨大的内存压力，因为它利用了四种模型（演员、参考、评论、奖励）。常见的解决方案是采用低秩自适应训练 (LoRA) 来解决 RLHF 的内存压力。LoRA 冻结了预训练模型的权重，并将可训练的秩分解矩阵注入到 Transformer 架构的每一层中，显著减轻了可训练参数的数量。LoRA 通过减轻内存利用来加快 RLHF，允许更大的批处理 (batch) 巨细，从而大大提高吞吐量。
DeepSpeed-Chat with ZeRO++ 用于 RLHF 训练
图 9: ZeRO++ 加快了 RLHF 训练的生成和训练阶段
ZeRO++ 在 RLHF + LoRA 的场景下有着独特的应用，因为大多数模型权重都被冻结了。这意味着 ZeRO++ 可以将这些冻结的权重量化保存到 INT4/8 中，而不是将它们存储在 fp16 中并在每次通讯操作之前对其从事量化。通讯后的反量化仍然是为了让权重为计算做好准备，但反量化后的权重在计算后被简单地丢弃。
以这种方式利用 ZeRO++ 从事 RLHF 训练可以减轻内存利用和通讯量。这意味着通过减轻通讯以及由于减轻内存利用而启用更大的批处理巨细来提高训练吞吐量。在生成阶段，ZeRO++ 利用 hpZ 将所有权重通讯保持在每一个节点内，以利用更高的节点内通讯带宽，减轻通讯量，进一步提高生成吞吐量。
ZeRO++ 已集成到 DeepSpeed-Chat 中，以支持 ChatGPT 类模型的 RLHF 训练。在图 9 中，我们比较了不同巨细的 actor 模型的 RLHF 生成吞吐量。测试配置为 32 个 V100 GPU ，actor 模型巨细为 30B 和 66B 以测试 ZeRO 和 ZeRO++ 性能。结果表明，ZeRO++ 的 RLHF 生成吞吐量比 ZeRO 高出 2.25 倍。我们还展示了在 16 个 V100 GPU 上训练阶段的加快，其中 ZeRO++ 实行了比 ZeRO 高 1.26 倍的吞吐量，这是由于 ZeRO++ 支持的更低通讯量和更大批量巨细。
DeepSpeed ZeRO++ 现已发布！
我们非常高兴能够发布 DeepSpeed ZeRO++ 并让 AI 社区中的每一个人都可以利用它。请访问我们的 GitHub 页面以获取 LLM 训练教程。用于 DeepSpeed-Chat 的 ZeRO++ 将在未来几周内发布。有关 ZeRO++ 的更多技术细节，请查看我们的 arxiv 论文。
DeepSpeed-ZeRO++ 是 DeepSpeed 生态系统的一部分。要了解更多信息，请访问我们的网站，在那里您可以找到详细的博客文章、教程和有用的文档。
您还可以在我们的英文 Twitter、日文 Twitter 和中文知乎 上获取最新的 DeepSpeed 新闻。
DeepSpeed 欢迎您的贡献！我们鼓励您在 DeepSpeed GitHub 页面上报告问题、贡献 PR 并加入讨论。有关更多详细信息，请参阅我们的贡献指南。我们对与大学、研究实验室和公司的合作持开放态度。对于此类请求（以及其他不适合 GitHub 的请求），请直接发送电子邮件至 [email protected]。
贡献者:
DeepSpeed 团队的以下人员的贡献使该项目成为可能：
Guanhua Wang, Heyang Qin, Sam Ade Jacobs, Connor Holmes, Samyam Rajbhandari, Olatunji Ruwase, Ammar Ahmad Awan, Jeff Rasley, Michael Wyatt, Yuxiong He (team lead)
本文转载自微软DeepSpeed组。
官方知乎账号：zhihu.com/people/deepspeed
原文链接：https://zhuanlan.zhihu.com/p/639002087