[关闭]
@levinzhang 2020-02-15T19:19:42.000000Z 字数 3987 阅读 1182

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是我们用于加速分布式深度学习训练的高性能库。

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

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

使用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与基线对比的通信量。

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

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

DeepSpeed：PyTorch兼容性和系统性能

我们实现了ZeRO的第一个阶段，即优化区状态分区（简称为ZeRO-OS），它能够支撑1000亿个参数的模型。代码与我们的训练优化库DeepSpeed一起发布。DeepSpeed通过与PyTorch兼容的轻量级API，提供了最先进的训练技术，如ZeRO、分布式训练、混合精度和检查点。只需对PyTorch模型进行几行代码的更改，就可以利用DeepSpeed来解决底层的性能挑战，并提高训练的速度和规模。

DeepSpeed在四个方面都有很好的表现（如图2所示）：

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

添加新批注
在作者公开此批注前，只有你和作者可见。
回复批注