微软发布 DeepSpeed 开源库，支持 1000 亿个参数模型的训练

微软发布了名为 DeepSpeed 的开源库，它通过改善规模、速度、成本和可用性，能够极大地推进大型模型的训练，支持 1000 亿个参数模型的训练。

本文最初发表于微软的Research网站，由 InfoQ 中文站翻译分享。

在 AI 方面，最新的趋势是更大的自然语言模型会提供更好的准确性，但是，考虑到成本、时间以及代码集成的便捷性，更大的模型通常会难以训练。微软发布了名为 DeepSpeed 的开源库，它通过改善规模、速度、成本和可用性，能够极大地推进大型模型的训练，支持 1000 亿个参数的模型。DeepSpeed 能够与 PyTorch 兼容。这个库包含一个名为 ZeRO 的组成部分，这是一个新的并行优化器，在可训练参数的数量大幅度增长的时候，它能够大幅减少模型和数据并行处理所需的资源。研究人员利用这些突破创建了 Turing 自然语言生成（Turing Natural Language Generation，Turing-NLG）模型，这是已知的最大的模型，有 170 亿个参数，读者可以通过该篇博客文章了解它的更多信息。

ZeRO（Zero Redundancy Optimizer）是一项新型的内存优化技术，适用于大规模的分布式深度学习。ZeRO 能够在当前的 GPU 集群上训练具有 1000 亿个参数的深度学习模型，其吞吐量是当前最佳系统的三到五倍。它还提出了一个清晰的路径，以训练带有数万亿参数的模型，在深度学习系统技术中，这体现出了一个前所未有的飞跃。我们将 ZeRO 发布为 DeepSpeed 的一部分，DeepSpeed 是我们用于加速分布式深度学习训练的高性能库。

训练大型深度学习模型所面临的挑战

大型模型会带来明显的精准度提升，但是数十亿到数万亿个参数通常会遇到基本硬件的限制。为了将这些模型放入到内存中，现有的解决方案不得不在计算能力、通信和开发效率之间做出权衡：

• 数据并行无助于减少每台设备的内存占用。即便是在具有 32GB 内存的 GPU 上，如果一个模型有超过 10 亿个参数，也会导致内存的耗尽；

• 鉴于细粒度计算和昂贵的通信，模型并行无法有效扩展至带个节点之外。模型并行框架通常会需要大量的代码集成，而这些集成工作可能是与特定的模型架构相关的。例如，NVIDIA Megatron-LM 创造了一个 83 亿参数的新模型记录。在单个节点上有多个 GPU 时，该模型能够扩展地非常好，但是如果跨节点扩展的话，它的性能就会下降。例如，当在 NVIDIA DGX-2 节点上运行 400 亿个参数时，我们观察到每个 GPU 大约 5 万亿次浮点运算。

使用 ZeRO 克服数据并行和模型并行的挑战

我们开发 ZeRO 就是为了克服数据并行性和模型并行性的限制，兼得两者的优点。ZeRO 通过将模型状态（参数、梯度和优化器状态）跨数据并行进程（而不是复制它们）分区来消除数据并行进程之间的内存冗余。在训练期间，它使用一个动态通信调度机制来跨分布式设备共享必要的状态，以保持数据并行的计算粒度和通信量。

我们将其称为 ZeRO 支撑的数据并行（ZeRO-powered data parallelism），它让每个设备的内存使用能够随数据并行度线性扩展，并产生与数据并行类似的通信量。只要聚合的设备内存足以共享模型状态，那么 ZeRO 支撑的数据并行性能够适用于任意规模的模型。

ZeRO 的三个阶段及其收益

ZeRO 有三个主要的优化阶段（如图 1 所示），分别对应于优化器状态（optimizer state）、梯度（gradient）和参数分区。在逐步启用它们的时候：

1. 优化器状态分区（Pos）：会带 4 倍的内存消耗降低，与数据并行相同的通信量；

2. 添加梯度分区（Pos+g）：会带 8 倍的内存消耗降低，与数据并行相同的通信量；

3. 添加参数分区（Pos+g+p）：内存减少与数据并行度 Nd 呈线性关系。例如，跨 64 个 GPU (即 Nd = 64)进行拆分将会减少 64 倍的内存。通信量温和性地增长了 50%。

ZeRO 消除了内存冗余，并使集群的全部聚合内存容量均可用。启用了全部三个阶段后，ZeRO 可以在仅仅 1024 个 NVIDIA GPU 上训练一个万亿参数的模型。借助像 Adam 这样的 16 位精度的优化器，一个万亿参数的模型需要大约 16 TB 的内存来保存优化器状态、梯度和参数。16TB 除以 1024 等于 16GB，这对于 GPU 来说是一个合理的界限。

图 1 与标准的数据并行性对比，ZeRO 的内存节省和通信量。在内存消耗公式中，Ψ指的是模型中的参数数量，K 是优化器特定的常量。在具体的样例中，我们展现了具有 7.5B 个参数并使用 Adam 的内存占用，其中在 64 个 GPU 的时候，K=12。我们也看到了 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 亿个参数至数万亿参数。

• 速度：我们观察到，在不同的硬件上，吞吐量能够比现在提升高达 5 倍。例如，为了在 GPT 家族的工作负载上训练大型模型，DeepSpeed 将 ZeRO 支撑的数据并行与 NVIDIA Megatron-LM 模型并行结合在了一起。在低带宽互连的 NVIDIA GPU 集群上（没有 NVIDIA NVLink 或 Infiniband），对于标准的 15 亿参数的 GPT-2 模型，相对于单独运行 Megatron-LM，我们实现了 3.75 倍的吞吐量提升。在具有高带宽互连的 NVIDIA DGX-2 集群上，对于 200 到 800 亿个参数的模型，我们要快 3 到 5 倍。这些吞吐量提升源于 DeepSpeed 更高的内存效率，以及能够使用更低的模型并行度和更大的批处理大小来适应这些模型的能力。

• 成本：吞吐量的提升可以显著降低训练成本。例如，要训练一个具有 200 亿个参数的模型，DeepSpeed 所需的资源仅是原来的三分之一。

• 可用性：只需几行代码更改，PyTorch 模型就可以使用 DeepSpeed 和 ZeRO。与当前的模型并行性库相比，DeepSpeed 不需要重新设计代码或重构模型。它也没有对模型维度、批量大小或任何其他训练参数进行限制。对于多达 60 亿个参数的模型，我们可以方便地使用数据并行（由 ZeRO 提供支持），而不必使用模型并行，相反，对于参数超过 13 亿个的模型，标准的数据并行就将耗尽内存。ZeRO 的第二阶段和第三阶段将会进一步增加可训练模型的大小。此外，DeepSpeed 还支持将 ZeRO 支撑的数据并行与模型并行灵活地结合起来。

Turing-NLG 和基于 DeepSpeed 的大型模型训练

我们利用 DeepSpeed 中的 ZeRO-OS 来训练具有 170 亿个参数的 Turing-NLG 模型，与目前最先进的方法相比，这种方式具有更高的精确度和更高的训练效率。请参考该博客文章，它展示了该模型所建立的新的准确性记录及其在自由格式文本生成、摘要和答案合成方面的广泛应用。

ZeRO-OS 与不同类型的模型并行是互补和兼容的，对于不适合单个节点的大型模型（约 200 亿个参数或更多），与单独使用模型并行相比，它提供了显著的性能收益、资源节省和模型设计的灵活性。

在 DeepSpeed 中，我们使用 ZeRO-OS 和 NVIDIA 的 Megatron-LM 组合训练 Turning-NLG 模型。与单独使用 NVIDIA Megatron-LM 相比，ZeRO-OS 节省的内存使 Turning-NLG 模型的并行度降低了 4 倍，而且批处理大小增加了 4 倍。因此，我们实现了 3 倍的吞吐量增益。此外，我们可以只使用 256 个 GPU 就可以实现批处理大小（batch size）为 512 的训练，而单独使用 Megatron-LM 需要 1024 个 GPU。最后，Megatron-LM 不能运行该精确模型，它不支持这种模型结构，因为它的 attention head（=28）不能被模型并行度（=16）整除。DeepSpeed 使该模型从不可行变为可行，实现了高效的训练!

有关详细信息，请参见 DeepSpeed GitHub 仓库和 ZeRO 论文。我们还与 ONNX 和 ONNX Runtime 社区合作，进一步集成这些技术。

关于 DeepSpeed 团队：我们是一群系统研究人员和工程师，包括 Samyam Rajbhandari、Jeff Rasley、Olatunji Ruwase、Arash Ashari、Elton Zheng、Jing Zhao、Minjia Zhang、Niranjan Uma Naresh、Reza Yazdani Aminabadi、Shaden Smith、Yuxiong He（团队领导），我们热衷于大型系统的性能优化。我们最近专注于深度学习系统，优化其训练速度、收敛速度和发展速度！