一、DeepSpeed介绍
DeepSpeed 是 PyTorch 的开源深度学习优化库。 该库旨在降低计算能力和内存使用，并在现有计算机硬件上训练具有更好并行性的大型分布式模型。
DeepSpeed 针对低延迟、高吞吐量训练进行了优化。它包括零冗余优化器 (ZeRO)，用于训练具有 1 万亿或更多参数的模型。功能包括混合精度训练、单 GPU、多 GPU 和多节点训练以及自定义模型并行性。
DeepSpeed 源代码已根据 MIT 许可证获得许可，并可在 GitHub 上获取。
该团队声称吞吐量提高了 6.2 倍，收敛速度提高了 2.8 倍，通信量减少了 4.6 倍。
DeepSpeed 是一款易于使用的深度学习优化软件套件，可为深度学习训练和推理提供前所未有的规模和速度。
deepspeed.ai
Github 存储库
deepspedd docs
deepspeed huggingface
二、正文
人工智能的最新趋势是更大的自然语言模型提供更好的准确性；然而，由于成本、时间和代码集成的难易程度，较大的模型很难训练。微软正在发布一个名为 DeepSpeed 的开源库，该库通过提高规模、速度、成本和可用性，极大地推进了大型模型训练，释放了训练 1000 亿个参数模型的能力。 DeepSpeed 与 PyTorch 兼容。该库的一个部分称为 ZeRO，是一种新的并行优化器，它大大减少了模型和数据并行所需的资源，同时大量增加了可训练的参数数量。研究人员利用这些突破创建了图灵自然语言生成 (Turing-NLG)，这是拥有 170 亿个参数的最大的公开语言模型，您可以在这篇附带的博客文章中了解更多信息。
零冗余优化器（Zero Redundancy Optimizer，简称ZeRO）是一种针对大规模分布式深度学习的新型内存优化技术。 ZeRO 可以在当前一代 GPU 集群上训练具有 1000 亿个参数的深度学习模型，吞吐量是当前最佳系统的三到五倍。它还为训练具有数万亿参数的模型提供了一条清晰的路径，展示了深度学习系统技术前所未有的飞跃。我们将发布 ZeRO 作为 DeepSpeed 的一部分，DeepSpeed 是我们用于加速分布式深度学习训练的高性能库。
训练大型深度学习模型的挑战
大型模型可显着提高准确性，但训练数十亿至数万亿参数经常会遇到基本的硬件限制。为了将这些模型放入内存中，现有的解决方案在计算、通信和开发效率之间进行了权衡：
- 数据并行性无助于减少每个设备的内存占用：具有超过 10 亿个参数的模型即使在具有 32GB 内存的 GPU 上也会出现内存不足的情况。
- 由于细粒度计算和昂贵的通信，模型并行性无法有效地扩展到单个节点之外。模型并行框架经常需要广泛的代码集成，这些代码集成可能是特定于模型架构的。例如，NVIDIA Megatron-LM 创下了 83 亿个参数的模型大小新纪录。对于适合单个节点的多个 GPU 的模型来说，它的扩展性非常好，但是当跨节点扩展时，其性能会下降。例如，当跨 NVIDIA DGX-2 节点运行 400 亿个参数时，我们观察到大约 5 teraflops/GPU。
使用 ZeRO 克服数据并行性和模型并行性的限制
我们开发ZeRO是为了克服数据并行和模型并行的局限性，同时实现两者的优点。 ZeRO 通过跨数据并行进程划分模型状态（参数、梯度和优化器状态）而不是复制它们，消除了数据并行进程中的内存冗余。它在训练期间使用动态通信调度来跨分布式设备共享必要的状态，以保留数据并行性的计算粒度和通信量。
我们称之为 ZeRO 支持的数据并行性，它允许每个设备的内存使用量随着数据并行性的程度线性扩展，并产生与数据并行性类似的通信量。 ZeRO 支持的数据并行性可以适应任意大小的模型 - 只要聚合的设备内存足够大以共享模型状态。
ZeRO的三个阶段及其好处
ZeRO 具有三个主要优化阶段（如图 1 所示），分别对应于优化器状态、梯度和参数的划分。累积启用时：
- 优化器状态分区 (P os ) – 内存减少 4 倍，通信量与数据并行相同
- 添加梯度分区 (P os+g) – 内存减少 8 倍，通信量与数据并行相同
- 添加参数分区 (P os+g+p ) – 内存减少与数据并行度 N d 呈线性关系。例如，拆分为 64 个 GPU (N d = 64) 将导致内存减少 64 倍。通信量略有增加 50%。
ZeRO 消除了内存冗余，并使集群的全部聚合内存容量可用。启用所有三个阶段后，ZeRO 可以在 1024 个 NVIDIA GPU 上训练万亿参数模型。具有 Adam 等 16 位精度优化器的万亿参数模型需要大约 16 TB 的内存来保存优化器状态、梯度和参数。 16TB 除以 1024 就是 16GB，这对于 GPU 来说完全在合理范围内。
图 1：与标准数据并行基线相比，ZeRO 三个阶段的内存节省和通信量。在内存消耗公式中，Ψ 指模型中的参数数量，K 是优化器特定常数项。作为一个具体示例，我们展示了使用 Adam 优化器（其中 K=12）在 64 个 GPU 上的 7.5B 参数模型的内存消耗。我们还显示了 ZeRO 相对于基线的通信量。
下面的视频展示了 ZeRO（包含所有三个阶段）如何执行训练步骤，包括前向传递、后向传递和参数更新。
DeepSpeed：PyTorch 兼容性和系统性能
我们实施了 ZeRO 第一阶段——优化器状态分区（简称 ZeRO-OS）——它已被证明能够支持 1000 亿个参数模型。该代码与我们的训练优化库 DeepSpeed 一起发布。 DeepSpeed 通过与 PyTorch 兼容的轻量级 API 带来了最先进的训练技术，例如 ZeRO、分布式训练、混合精度和检查点。只需对 PyTorch 模型进行几行代码更改，您就可以利用 DeepSpeed 来解决潜在的性能挑战并提高训练的速度和规模。
DeepSpeed 在四个方面表现出色（如图 2 所示）：
-
规模：OpenAI GPT-2、NVIDIA Megatron-LM 和 Google T5 等最先进的大型模型的参数大小分别为 15 亿、83 亿和 110 亿。 DeepSpeed 中的 ZeRO 第一阶段提供系统支持，可运行多达 1000 亿个参数的模型，参数数量是原来的 10 倍。未来，我们计划增加对 ZeRO 第二阶段和第三阶段的支持，解锁训练具有 2000 亿个参数到数万亿个参数的模型的能力。
-
速度：我们发现各种硬件的吞吐量比现有技术高出五倍。例如，为了在 GPT 系列工作负载上训练大型模型，DeepSpeed 将 ZeRO 支持的数据并行性与 NVIDIA Megatron-LM 模型并行性相结合。在具有低带宽互连（无 NVIDIA NVLink 或 Infiniband）的 NVIDIA GPU 集群上，与单独使用 Megatron-LM 处理具有 15 亿个参数的标准 GPT-2 模型相比，我们实现了 3.75 倍的吞吐量提升。在具有高带宽互连的 NVIDIA DGX-2 集群上，对于 20 至 800 亿个参数的模型，我们的速度提高了三到五倍。这些吞吐量的提高来自 DeepSpeed 更高的内存效率以及使用较低程度的模型并行性和更大的批量大小来适应这些模型的能力。
-
成本：吞吐量的提高可以转化为培训成本的显着降低。例如，要训练具有 200 亿个参数的模型，DeepSpeed 所需的资源减少了三倍。
-
可用性：只需更改几行代码即可使PyTorch 模型使用DeepSpeed 和ZeRO。与当前的模型并行库相比，DeepSpeed 不需要重新设计代码或模型重构。它还不对模型维度（例如注意力头的数量、隐藏大小等）、批量大小或任何其他训练参数施加限制。对于多达 60 亿个参数的模型，您可以方便地使用数据并行（由 ZeRO 提供支持），而无需模型并行，而相比之下，对于超过 13 亿个参数的模型，标准数据并行会耗尽内存。 ZeRO 的第二阶段和第三阶段将进一步增加仅通过数据并行性即可训练的模型大小。此外，DeepSpeed 支持 ZeRO 支持的数据并行性与模型并行性的灵活组合。
图 2：DeepSpeed 在规模、速度、成本和可用性方面表现出色。左下图描绘了 DeepSpeed（将 ZeRO 支持的数据并行性与 Megatron-LM 的模型并行性相结合）相对于单独使用 Megatron-LM 的系统吞吐量改进。右下图比较了单独使用数据并行性（使用和不使用 ZeRO）的可训练模型大小。
Turing-NLG 和 DeepSpeed 支持的大型模型训练
我们利用 DeepSpeed 中的 ZeRO-OS 训练 170 亿参数的 Turing-NLG 模型，其精度和训练效率比当前最先进的方法更高。请参阅此博客，其中展示了该模型建立的新准确率记录及其在自由格式文本生成、摘要和答案合成方面的广泛应用。
ZeRO-OS 与不同类型的模型并行性互补并兼容，对于不适合单个节点的大型模型（大约 200 亿个参数或更多），与其他模型相比，它提供了显着的性能提升、资源节省和模型设计灵活性单独使用模型并行性。
我们在 DeepSpeed 中使用 ZeRO-OS 与 NVIDIA 的 Megatron-LM 结合来训练 Turing-NLG 模型。与单独使用 NVIDIA Megatron-LM 相比，ZeRO-OS 节省的内存允许 Turning-NLG 模型以小 4 倍的模型并行度和大 4 倍的批量大小运行。结果，我们实现了 3 倍的吞吐量增益。此外，我们可以仅使用 256 个 GPU 以 512 的批量大小进行训练，而单独使用 Megatron-LM 则需要 1024 个 GPU。最后，Megatron-LM 无法运行这个精确的模型——模型结构不受支持，因为它的注意力头 (=28) 不能被模型并行度 (=16) 整除。 DeepSpeed 将模型从不可行的运行变为可行且高效的训练！
有关更多详细信息，请参阅 DeepSpeed GitHub 存储库和 ZeRO 论文。我们还与 ONNX 和 ONNX Runtime 社区合作，进一步集成这些技术。
关于 DeepSpeed 团队：我们是一群系统研究人员和工程师 - Samyam Rajbhandari、Jeff Rasley、Olatunji Ruwase、Arash Ashari、Elton Cheng、Jing Chao、Minjia Zhang、Niranjan Uma Naresh、Reza Yazdani Aminabadi、Shaden Smith、Yyuxiong He（团队负责人）——热衷于大型系统的性能优化。我们最近专注于深度学习系统，优化其训练速度、收敛速度和开发速度！
