使用 DeepSpeed 训练大型模型

概述

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

DeepSpeed 提供了一系列系统技术，使得训练这些规模的模型成为可能。训练大型模型的最佳技术取决于多种因素，例如模型架构、批处理大小、互连带宽等。鉴于可用的选择数量众多，这可能令人困惑甚至望而却步。本页面旨在为您提供一个入门指南，帮助您踏上训练大型模型的旅程。

训练大型模型的可能方法

从宏观层面看，训练大型模型主要有两条途径

  • 基于 ZeRO (零冗余优化器) 的技术
  • 基于 3D 并行的技术

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

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

基于 3D 并行的技术：3D 并行是指三种不同形式的并行技术组合，即张量切片（tensor-slicing）、流水线并行（pipeline-parallelism）和数据并行（或由 ZeRO 提供支持的数据并行）。结合这三种形式可以发挥每种技术的优势，而避免其缺点。在仅依赖 ZeRO 技术可能不足的场景中，3D 并行使 DeepSpeed 能够实现出色的训练吞吐量效率。然而，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 高效运行。这在数千个 GPU 上训练或网络带宽有限时尤其如此。例如，在 4K 个 GPU 上训练 GPT-3 模型，且批处理大小限制为 2K，将导致每个 GPU 的批处理大小为 0.5，这取决于序列长度和网络带宽，可能不足以仅使用 ZeRO 维持良好的性能。

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

更新时间：