ZeRO & DeepSpeed: New system optimizations enable training models with over 100 billion parameters

ZeRO & DeepSpeed: New system optimizations enable training models with over 100 billion parameters

人工智能领域的最新趋势是，更大的自然语言模型提供了更高的准确性；然而，由于成本、时间和代码集成便捷性等因素，更大的模型训练起来比较困难。微软正在发布一个名为 DeepSpeed（在新标签页中打开）的开源库，该库通过提升规模、速度、成本和可用性，极大地推动了大型模型的训练，解锁了训练拥有 1000 亿参数的模型的能力。DeepSpeed 与 PyTorch（在新标签页中打开）兼容。该库中的一个组件名为 ZeRO，是一种新的并行优化器，能够极大地减少模型和数据并行所需的资源，同时大幅增加可训练参数的数量。研究人员利用这些突破创建了 Turing Natural Language Generation（Turing-NLG，在新标签页中打开），这是目前已知最大的公开语言模型，拥有 170 亿参数，您可以在随附的博客文章中了解更多信息。（在新标签页中打开）
零冗余优化器（简称 ZeRO）是一种针对大规模分布式深度学习的全新内存优化技术。ZeRO 能够在当前一代 GPU 集群上训练拥有 1000 亿参数的深度学习模型，吞吐量是当前最佳系统的三到五倍。它还为训练拥有万亿参数的模型提供了一条清晰的路径，展现了深度学习系统技术前所未有的飞跃。我们将 ZeRO 作为 DeepSpeed 的一部分发布，DeepSpeed 是我们用于加速分布式深度学习训练的高性能库。

评：使用zero优化器能节约显存的占用。

训练大模型的挑战

大型模型在显著提高准确性的同时，其训练过程中的数十亿到数万亿参数常常会受到基本硬件限制的影响。为了将这些模型放入内存中，现有的解决方案在计算、通信和开发效率之间做出了权衡：
• 数据并行性无法帮助减少每个设备的内存占用：一个拥有超过10亿参数的模型即使在具有32GB内存的GPU上也会耗尽内存。
• 模型并行性在单个节点之外无法有效扩展，因为计算粒度细且通信成本高。模型并行性框架通常需要广泛的代码集成，而这些集成可能特定于模型架构。例如，NVIDIA的Megatron-LM（在新标签页中打开）创造了83亿参数的新模型规模记录。它对于单个节点中多个GPU中的此类模型具有很好的扩展性，但在跨节点扩展时，其性能会下降。例如，我们在NVIDIA DGX-2节点上运行400亿参数时，观察到每个GPU的算力约为5 teraflops。

评：模型并行的技术需要大量的修改代码。
使用ZeRO克服数据并行性和模型并行性的限制
我们开发了ZeRO来克服数据并行性和模型并行性的限制，同时保留两者的优点。ZeRO通过跨数据并行进程划分模型状态（参数、梯度和优化器状态）而不是复制它们，来消除数据并行进程之间的内存冗余。在训练过程中，它使用动态通信计划来跨分布式设备共享必要的状态，以保持数据并行性的计算粒度和通信量。
我们称这种由ZeRO驱动的数据并行性为“ZeRO增强型数据并行性”，它允许每个设备的内存使用量随数据并行度的增加而线性扩展，并且产生的通信量与数据并行性相似。只要聚合的设备内存足够大以共享模型状态，ZeRO增强型数据并行性就可以适应任意大小的模型。

评：使用通信的方式来从其他GPU中获取资源

ZeRO的三个阶段及其优势

ZeRO有三个主要的优化阶段（如图1所示），分别对应于优化器状态、梯度和参数的划分。当它们被累积启用时：
优化器状态划分（Pos） - 内存减少4倍，通信量与数据并行性相同
添加梯度划分（Pos+g） - 内存减少8倍，通信量与数据并行性相同
添加参数划分（Pos+g+p） - 内存减少与数据并行度Nd成线性关系。例如，跨64个GPU（Nd = 64）划分将实现64倍内存减少。通信量适度增加50%。
ZeRO消除了内存冗余，并使集群的全部聚合内存容量可用。当启用所有三个阶段时，ZeRO可以在仅1024个NVIDIA GPU上训练一个万亿参数模型。一个使用像Adam（在新标签页中打开）这样的优化器，在16位精度下的万亿参数模型，大约需要16TB（太字节）的内存来存储优化器状态、梯度和参数。16TB除以1024等于16GB，这对于GPU来说是一个合理的范围。
评：将模型的存储分散。相对于数据并行。

DeepSpeed：PyTorch兼容性与系统性能

我们实现了ZeRO的第一阶段——优化器状态划分（简称ZeRO-OS），该阶段已展示出支持1000亿参数模型的能力。此代码将与我们的训练优化库DeepSpeed一同发布。DeepSpeed通过轻量级API提供了最前沿的训练技术，如ZeRO、分布式训练、混合精度和检查点保存，这些API与PyTorch（在新标签页中打开）兼容。只需对您的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集群上，对于标准的15亿参数GPT-2模型，我们使用DeepSpeed实现的吞吐量比单独使用Megatron-LM提高了3.75倍。在具有高带宽互连的NVIDIA DGX-2集群上，对于20到800亿参数的模型，我们的速度是三到五倍。这些吞吐量的提升来自于DeepSpeed更高的内存效率和能够在较低的模型并行度和更大的批量大小下适应这些模型的能力。
评：使用这个并行技术比模型并行技术还要快。但是从原理上来看，模型并行应该更快。但是这两个方法其实是有点问题的。当我使用模型并行的时候，我再使用数据并行，我似乎已经并行不了了。因为模型并行将模型的权重分配到各个GPU上了，你再使用数据并行的技术。还将将模型的一部分来进行切分吗？但实际上是可以兼容的，我之前也认为是可以兼容的。
• 成本：吞吐量的提升可以转化为显著降低的训练成本。例如，训练一个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中结合使用NVIDIA的Megatron-LM和ZeRO-OS来训练Turing-NLG模型。ZeRO-OS的内存节省使得Turing-NLG模型能够在比仅使用NVIDIA Megatron-LM时低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 Zheng、Jing Zhao、Minjia Zhang、Niranjan Uma Naresh、Reza Yazdani Aminabadi、Shaden Smith、Yuxiong He（团队负责人）——我们热衷于大规模系统的性能优化。最近，我们专注于深度学习系统，优化其训练速度、收敛速度和开发速度！
如果您对这种类型的工作感兴趣，DeepSpeed团队正在招聘！请访问我们的职业页面（在新标签页中打开）。