使用 DeepSpeed 训练大型模型

概述

DeepSpeed 已被用于训练或正在训练一些现有的最大密集模型。这些包括但不限于

DeepSpeed 提供了一系列系统技术，使这些规模的模型训练成为可能。训练大型模型的最佳技术取决于各种因素，例如模型架构、批大小、互连带宽等。鉴于可用选项的数量，这可能令人困惑且令人望而生畏。此页面旨在作为入门指南，帮助您在训练大型模型的旅程中找到方向。

训练大型模型的可能方法

从广义上讲，训练大型模型主要有两条路径

  • 基于 ZeRO（零冗余优化器）的技术
  • 基于 3D 并行性的技术

基于 ZeRO 的技术：简单来说，ZeRO 是一种内存高效的数据并行形式，它让您能够访问所有可用 GPU 设备的聚合 GPU 内存，而不会因数据并行中的数据复制导致效率低下。此外，DeepSpeed 还提供基于 ZeRO 的异构内存技术，例如 ZeRO-Offload 和 ZeRO-Infinity，这些技术允许您在目标系统上可用时有效地利用 CPU 和 NVMe 内存。

由于 ZeRO 替代了数据并行，因此它提供了无缝集成，无需对现有数据并行模型进行模型代码重构。在大多数情况下，基于 ZeRO 的技术提供了模型可扩展性、训练吞吐量效率，而不会影响易用性。

基于 3D 并行性的技术：3D 并行性指的是三种不同形式的并行技术的组合，即张量切片、流水线并行和数据并行（或 ZeRO 支持的数据并行）。结合这三种形式可以发挥每种技术的优势，而不会有任何缺点。3D 并行使 DeepSpeed 能够在仅依靠基于 ZeRO 的技术可能不足的情况下实现出色的训练吞吐量效率。但是，3D 并行需要非平凡的模型代码重构，因此需要仔细考虑以识别 3D 并行可以带来非平凡吞吐量优势的情况。

决定使用哪种技术

适用于 GPT-2/GPT-3 类模型的 3D 并行性：如果您尝试训练的模型架构与 GPT-2 或 GPT-3 非常相似，那么我们已经完成了将 3D 并行移植到基于 GPT-2/GPT-3 架构的模型的艰苦工作，并创建了一个您可以用来高效训练具有数百亿甚至数万亿参数的模型的训练管道。Megatron-Turing NLG 530B 和 Big Science 都使用此代码库的变体来扩展模型训练。您可以在 DeepSpeed-Megatron GPT-3 存储库中找到入门代码和教程。有关 3D 并行性的更多信息，请查看以下资源

3D 并行教程 有关如何将您的模型移植以使用 DeepSpeed 3D 并行的通用教程

3D 并行深入探讨 一篇微软研究院博客文章，深入探讨了 DeepSpeed 中 3D 并行的实现。

基于 ZeRO 的技术：对于大多数训练场景，ZeRO 提供的训练效率与 3D 并行相当，而无需进行模型代码重构。因此，如果您尚未将代码移植以使用 3D 并行，我们建议您首先尝试 ZeRO 系列技术，看看它是否满足您的需求。使用 DeepSpeed 将 ZeRO 添加到您的训练管道非常简单，无需您对模型进行更改。鉴于使用 DeepSpeed 试用 ZeRO 的成本微不足道，它是评估和决定是否应进一步投资将模型移植以使用 3D 并行的最快方法。启用 DeepSpeed 中的 ZeRO 还使您可以访问 ZeRO-Offload 和 ZeRO-Infinity，这可以帮助您在有限的 GPU 资源上微调大型模型。要开始，请查看我们的 ZeRO 教程。

有关 ZeRO 系列技术的更深入信息，请查看我们的论文

ZeRO (SC20)、ZeRO Offload (ATC21) 和 ZeRO-Infinity (SC21)，

以及博客文章

ZeRO & DeepSpeed、ZeRO-2 & DeepSpeed、ZeRO-Offload 和 ZeRO-Infinity & DeepSpeed

了解 ZeRO 和 3D 并行之间的性能权衡

当每个 GPU 的批大小不是极小时，ZeRO 和 3D 并行的性能通常相当。ZeRO 是一种更内存高效的数据并行形式，并且 ZeRO 的通信成本与数据并行本身非常相似。因此，在数据并行运行良好的所有场景中，ZeRO 也会运行良好。事实上，与数据并行相比，ZeRO 能够为大型模型拟合更大的批大小，因为它具有内存效率，从而实现比数据并行更好的吞吐量效率。

但是，在某些情况下，批大小可能不够大以使 ZeRO 效率更高。在使用数千个 GPU 或网络带宽有限的情况下，尤其如此。例如，在 4K 个 GPU 上训练 GPT-3 模型，并且批大小限制为 2K 将导致每个 GPU 的批大小为 0.5，这取决于序列长度和网络带宽，可能不足以维持使用 ZeRO 独自获得良好的性能。

在这种情况下，应该考虑是否可以增加批大小以获得更好的效率。但是，如果由于与收敛相关的考虑而无法增加批大小，那么 3D 并行中的流水线并行可以使有效网络带宽与流水线阶段数成正比地增加，从而使 3D 并行能够实现比 ZeRO 更好的吞吐量。

更新：