ZeRO & 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. 优化器状态分区（P_os）：会带4倍的内存消耗降低，与数据并行相同的通信量；
2. 添加梯度分区（P_os+g）：会带8倍的内存消耗降低，与数据并行相同的通信量；
3. 添加参数分区（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是优化器特定的常量。在具体的样例中，我们展现了具有7.5B个参数并使用Adam的内存占用，其中在64个GPU的时候，K=12。我们也看到了ZeRO与基线对比的通信量。

下面的视频展示了ZeRO（包括所有三个阶段）如何执行一个训练步骤，包括前向传播（forward pass）、反向传播（backward pass）和参数更新。

Turing-Animation.mp4?_=1未知大小

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（团队领导），我们热衷于大型系统的性能优化。我们最近专注于深度学习系统，优化其训练速度、收敛速度和发展速度!