微软更新 DeepSpeed：可用更少的 GPU 训练更多的 AI 模型

今年 2 月，我们发布了 DeepSpeed和 ZeRO（Zero Redundancy Optimizer，零冗余优化器），前者是一个开源的深度学习训练优化库，后者是该库中的一种新型内存优化技术，通过提高规模、速度、成本和可用性，极大地推进了大型模型训练。DeepSpeed 研究人员创建了 Turing-NLG，这是最大的语言模型，拥有 170 亿个参数，发布时具有最先进的正确率。

今年 5 月，我们发布了 ZeRO-2，它支持 2000 亿个参数的模型训练，比现有技术快 10 倍，同时还提供了一系列计算、I/O 和收敛优化，以实现最快的 BERT 训练。从那时起，我们一直在快速持续创新，推动深度学习训练速度和规模的边界。

今天，我们很高兴与大家分享我们的新进展，这些新进展不仅将深度学习训练推向极致，而且还为更多的人实现了深度学习训练的民主化：从在大规模超级计算机上进行训练的数据科学家，到在低端集群甚至在单 GPU 上进行训练的数据科学家。更具体地说，DeepSpeed 增加了四项新的技术，进一步推进了 “AI at Scale” 计划，以在 Microsoft 的人工智能产品和平台上进行创新。它们提供了极高的计算、内存和通信效率，并为模型训练提供了数十亿到数万亿个参数。这些技术还允许在硬件系统上使用极长的输入序列，这些硬件系统包括单 GPU 的硬件系统、具有数千个 GPU 的高端集群或具有非常慢的以太网的低端集群。

• 具有 3D 并行性的万亿参数模型训练：DeepSpeed 支持三种并行方法的灵活组合——ZeRO 支持的数据并行、流水线并行和张量切片模型并行。3D 并行可适应工作负载需求的变化，为具有超过一万亿个参数的超大型模型提供支持，同时实现近乎完美的内存扩展和吞吐量扩展效率。此外，其改进的通信效率允许用户在网络带宽有限的常规集群上将数十亿参数模型的训练速度提高 2~7 倍。

• 使用 ZeRO-Offload，在单 GPU 上进行 10 倍的模型训练。我们扩展了 ZeRO-2，利用 CPU 和 GPU 内存来训练大型模型。使用一台 NVIDIA V100 GPU 机器，我们的用户可以在不耗尽内存的情况下运行多达 130 亿个参数模型，比现有方法大 10 倍，同时获得具有竞争力的吞吐量。这一特性实现了数十亿参数模型训练的民主化，为众多深度学习从业者打开了探索更大更好模型的窗口。

• 通过 DeepSpeed 稀疏注意力机制，提供 10 倍长序列和 6 倍快的执行速度。DeepSpeed 提供了稀疏注意力内核，这是一种支持长序列模型输入的工具性技术，无论是文本、图像还是声音。与经典的密集 Transformer 相比，它提供了一个数量级更长的输入序列，并以相当的正确率获得高达 6 倍的执行速度。它的执行速度也比最先进的稀疏实现快 1.5~3 倍。此外，我们的稀疏内核支持高效执行灵活的稀疏格式，并赋予用户对其定制的稀疏结构进行创新的能力。

• 1 位 Adam，通信量最多可减少 5 倍。Adam 是训练许多大规模深度学习模型的有效（且可能是利用最充分的）优化器。然而，Adam 通常与通信高效的优化算法不兼容。因此，在跨分布式设备扩展时，通信成本可能会成为瓶颈。我们引入了一种新的算法，具有高效实现的 1 位 Adam，在达到与 Adam 相似的收敛效率的同时，通信量降低了 5 倍。在通信受限的场景中，我们观察到，分布式训练的速度最高可提高 3.5 倍，从而可以扩展到不同类型的 GPU 集群和网络。

这篇博文更深入地探讨了这四种技术。我们已经在开源库 DeepSpeed中提供了所有这些令人兴奋的新优化。

3D 并行性：扩展到万亿参数模型

随着现代 GPU 集群上可用计算的快速增长，训练一个具有不可思议能力的、强大的万亿参数模型不再是遥不可及的梦想，而是一个近在眼前的现实。DeepSpeed 结合了三种强大的技术，可以训练万亿规模的模型，并扩展到数千个 GPU：数据并行训练、模型并行训练和流水线并行训练。这种共生关系使深度学习训练的规模远远超出了每种策略单独提供的能力。3D 并行性同时解决了训练万亿参数模型的两个基本挑战：内存效率和计算效率。因此，DeepSpeed 可以在不牺牲速度的情况下，扩展以适应内存中最大规模的模型。

通过 3D 并行性实现内存和计算效率

数据、模型和流水线并行性各自在提高内存和计算效率方面起着特定的作用。图 1 说明了我们的 3D 策略。

内存效率：将模型的各层划分为流水线阶段，并通过模型的并行性进一步划分每个阶段的层。这种 2D 组合同时降低了模型、优化器和激活所消耗的内存。然而，我们并不能无限期地对模型进行划分而不屈服通信开销，从而限制了计算效率。

计算效率：为了在不牺牲计算效率的情况下，允许工作器（worker）的数量超过模型和流水线的并行性，我们使用了 ZeRO 支持的数据并行性（ZeRO-DP）。ZeRO-DP 不仅可以通过优化器状态分区进一步提高内存效率，而且还可以通过拓扑感知映射以最小的通信开销扩展到任意数量的 GPU。

拓扑感知的 3D 映射（图 2）：3D 并行中的每个维都被仔细地映射到工作器上，通过利用两个关键的架构属性来实现最大的计算效率。

1. 优化节点内和节点间的通信带宽：在三种策略中，模型并行的通信开销最大，因此我们优先将模型并行组放置在节点内以利用更大的节点内带宽。在这里，我们应用 NVIDIA Megatron-LM 进行张量分割风格的模型并行。当模型并行不跨越节点中的所有工作器时，则将数据并行组放置在节点内；否则，将数据并行组放置在节点间；流水线并行具有最低的通信量，因此我们可以在不受通信带宽限制的情况下跨节点调度流水线阶段。

2. 在通信中通过并行实现带宽放大：通过流水线并行和模型并行，每个数据并行组通信的梯度大小线性下降，从而使总通信量从纯数据并行减少，并且每个数据并行组在局部工作器子集之间独立并行地执行通信。结果，通过减少通信量、增加局部性和并行性的结合，数据并行通信的有效带宽被放大。

图 1：具有 32 个工作器的示例 3D 并行性。神经网络的层被划分为四个流水线阶段

每个流水线阶段内的层在四个模型并行工作器之间进一步划分。最后，每个流水线跨两个数据并行实例进行复制，ZeRO 将优化器状态跨数据并行副本进行分区。

图 2：图 1 中的工作器映射到系统上的 GPU，该系统具有八个节点，每个节点具有四个 GPU。着色表示同一节点上的 GPU

利用线性效率扩展支持万亿级参数模型训练

DeepSpeed 可以使用至少 800 个 NVIDIA V100 GPU 来训练一个具有一万亿参数的语言模型（图 3）。我们通过扩展模型的规模，并观察到线性增长，无论是在模型的规模还是训练的吞吐量方面，都展示了同步的内存和计算效率。在每个配置中，我们每个 GPU 可以训练大约 14 亿个参数，这是单个 GPU 可以支持的最大模型大小，而不会耗尽内存，这表明了完美的内存扩展。我们还获得了接近完美的线性计算效率扩展和每个 V100 GPU 47 每秒万亿次浮点运算的吞吐量。对于给定的硬件来说，这是令人印象深刻的可扩展性和吞吐量。

图 3：作为 GPU 函数的模型大小（以十亿参数为单位）和训练吞吐量（以万亿次浮点运算为单位）。

DeepSpeed 可以使用 800 个具有 32 GB 内存的 NVIDIA V100 张量核心 GPU 来训练具有 1 万亿个参数的模型。每种配置都使用由NVIDIA Megatron-LM提供的 16 路模型并行，其余 GPU 使用流水线并行进行排列。万亿参数模型有 298 层 Transformer，隐维度为 17408，用序列长度为 2048，批大小为 2048。对于较小的模型，我们按 GPU 数量的比例减少 Transformer 层数和批处理大小。

ZeRO-Offload：使用单 GPU 将模型训练扩大 10 倍

ZeRO-Offload 通过利用 GPU 及其主机 CPU 上的计算和内存资源，推入最大模型大小的边界，从而可以使用最小的 GPU 资源有效地训练。它允许在单个 NVIDIA V100 GPU 上训练多达 130 亿个参数模型，比最先进的 GPU 大 10 倍，同时保持每个 GPU 超过 30 每秒万亿次浮点运算的高训练吞吐量。

通过在单个 GPU 上实现数十亿参数模型训练，ZeRO-Offload 使大型模型训练民主化，使资源有限的深度学习从业者也可以使用大型模型。

图 4：最大的模型可以在单 GPU 上使用默认的 PyTorch 和 ZeRO-Offload 进行训练

ZeRO-Offload 背后的关键技术是我们在ZeRO-2的基础上，将优化器状态和梯度卸载到 CPU 内存上的新功能。这种方法允许 ZeRO-Offload 最小化 CPU 卸载带来的计算效率损失，同时还可以实现与原始 ZeRO-2 相同的、有时甚至更好的效率。下图显示了 ZeRO-Offload 的架构。

图 5：ZeRO-Offload 概述

DeepSpeed 稀疏注意力机制：为 10 倍长的序列提供支持，并将执行速度提高 6 倍

基于注意力的深度学习模型，比如 Transformer，在捕获输入序列中令牌之间的关系方面非常有效，甚至可以跨越很长的距离。因此，它们与文本、图像和基于声音的输入一起使用，其中序列长度可以用数千个令牌表示。然而，尽管注意力模块能够有效地捕捉长期的依赖关系，但在实践中，它们对长序列输入的应用受到注意力计算的计算量和记忆需求的限制，这些计算量和内存需求随着序列长度 $n$ 的增长而呈二次方增长 $O(n^2)$。

为了解决这一局限性，DeepSpeed 提供了一套稀疏注意力机制，这是一种工具性技术，可以通过块稀疏计算以数量级减少注意力计算的计算和内存需求。该套件不仅缓解了注意力计算的内存瓶颈，而且能有效地进行稀疏计算。它的 API 允许与任何基于 Transformer 的模型进行方便的集成。除了提供广泛的稀疏结构之外，它还具有处理任何用户定义的块稀疏结构的灵活性。

更具体地说，稀疏注意力可以被设计成计算附近令牌之间的局部注意力，或者经由利用局部注意力计算的汇总令牌来计算全局注意力。此外，稀疏注意力还可以允许随机注意力或局部注意力，全局注意力和随机注意力的任意组合，如图 6 所示，分别使用蓝色、橙色和绿色块。结果，稀疏注意力将内存占用减少为 $O(wn)$，其中 $1<w\le n$ 是一个参数，其值取决于注意力结构。

图 6：可变稀疏结构

GPU 的有效实现：虽然稀疏注意力的基本实现可能显示出节省内存的好处，但在计算方面来看，它可能甚至比完全计算更糟糕。这主要是由于稀疏数据给全图增加了分歧和未合并的内存访问。一般来说，开发高效的稀疏内核，特别是在 GPU 上，是一项具有挑战性的工作。DeepSpeed 提供了Triton开发的高效稀疏注意力内核。这些内核是以块稀疏范式结构的，它支持对齐内存访问，缓解线程分歧，并平衡处理器上的工作负载。

系统性能：稀疏注意力支持 10 倍以上的序列，计算速度快达 6.3 倍以上，如图 7 所示。左图显示了三种设置下 BERT-Base 和 BERT-Large 模型中可运行的最长序列长度：密集、带激活检查点的密集和带激活检查点的稀疏。对于 BERT-Base 和 BERT-Large，与密集注意力相比，稀疏注意力算法减少了总计算量，提高了训练速度：随着序列长度的增加，其 boost 值增大，而 BERT-Base 和 BERT-Large 分别提高了 6.3 倍和 5.3 倍。

图 7：BERT 模型的最大可能序列长度（左）；BERT-Base（中）和 BERT-Large（右）在单个 NVIDIA V100 GPU 上的训练时间，具有不同的序列长度

灵活处理任何块状稀疏结构：DeepSpeed 稀疏注意力套件并不针对任何特定的稀疏结构，而是让模型科学家能够在高效的系统支持下探索任何块状稀疏结构。目前，我们已经添加了流行的稀疏结构，如Fixed（来自 OpenAI Sparse Transformer）、BigBird（来自 Google）和 BSLongformer（AI2Longformer的 Block-Sparse 实现）。我们还定义了一个具有 "变量 "结构的模板，如图 6 所示，它可以用来简单地定制任何块-稀疏的随机、局部或全局注意力模式。

1 位 Adam：通信量减少 5 倍，训练速度提高 3.4 倍

大型模型（如 BERT 和 GPT-3）的可扩展训练需要根据模型设计、架构和系统功能进行仔细的优化。从系统的角度来看，通信已成为主要瓶颈，特别是在具有标准 TCP 互连的商用系统上，这些互连提供了有限的网络带宽。

通信压缩是减少此类系统训练时间的一项重要技术。压缩通信最有效的方法之一是通过误差补偿压缩，它提供了强大的收敛速度，即使在 1 位压缩下也是如此。然而，最先进的错误补偿技术只适用于基本的优化器，如随机梯度下降（SGD）和动量 SGD，它们是线性依赖梯度的。它们不能与 Adam 等基于梯度的非线性优化器一起工作，Adam 为许多任务提供了最先进的收敛效率和正确率，包括训练类 BERT 的模型。

对于像 Adam 这样强大的优化器来说，由于梯度（在方差项中）的非线性依赖性，使得开发基于误差补偿的压缩技术具有挑战性，限制了最先进的通信压缩技术的实用价值。

用 1 位 Adam 压缩通信

为了在使用 Adam 优化器的同时压缩通信，我们开发了 1 位 Adam，它通过预处理来解决梯度中的非线性问题。我们观察到非线性项方差（$v_t$）的变化幅度在经过几个轮数的训练后显著减小，并且在之后设置 $v_t$ 常数不会改变收敛速度。建议的 1 位 Adam 优化器，如图 14 所示，由两部分组成：预热阶段，本质上是普通 Adam 算法；压缩阶段，保持方差项不变，并将剩余的线性项，即动量，压缩为 1 位表示。

该算法的压缩阶段由一个阈值参数控制（如图 8 所示）。当我们检测到“方差”的变化低于某个阈值时，我们就切换到压缩阶段。我们的研究表明，预热阶段只需整个训练步骤的 15%~20%。

图 8：经典 Adam 算法与建议的 1 位压缩 Adam 算法中分布式训练步骤的比较

解决 1 位 Adam 的系统挑战

除了算法挑战之外，在训练系统中应用 1 位 Adam 还有两个系统挑战。首先，我们需要将动量转换为 1 位表示的高效内核。第二，我们需要高效的通信方案来跨不同的 GPU 交换这种压缩的动量。压缩的目标是减少整体训练时间，以便具有带宽受限互联的商品系统可以用来训练大型模型。我们在 DeepSpeed 中解决了这些挑战，并引入了一个完全优化的 1 位 Adam 实现，用于在通信受限的系统上进行训练。

1 位 Adam 在通信受限系统中的优势

1 位 Adam 提供了与 Adam 相同的收敛性，产生的通信量减少了 5 倍，使得 BERT-Large 预训练的吞吐量提高了 3.5 倍，SQuAD 微调的吞吐量提高了 2.7 倍。在压缩阶段观察到 6.6 倍（图 9 左）和 6.2 倍（图 9 右）的速度提升，使得这种端到端的吞吐量提升得以实现。值得一提的是，我们的 1 位 Adam 优化器在 40 千兆以太网系统上的扩展性非常好，其性能与 Adam 在 40 千兆 InfiniBand QDR 系统上的扩展性相当。我们注意到，根据 iPerf 基准，40 千兆以太网上的有效带宽为 4.1 Gbps，而根据 InfiniBand perftest 微基准，InfiniBand 提供了接近峰值的 32 Gbps 带宽。

图 9：1-bit Adam 在 NVIDIA V100 GPU 上进行 BERT-大型预训练（左）和 SQuAD 微调（右）的可扩展性。BERT 预训练和 SQuAD 微调的批大小分别为 16/GPU 和 3/GPU

有关这些新技术的代码、教程和文档，请访问DeepSpeed 网站和GitHub 仓库。我们还将其中一些技术集成到ONNX 运行时中。

作者介绍：

Rangan Majumder，Microsoft 负责搜索和人工智能的副总裁。

Junhua Wang，WebXT 搜索与人工智能平台团队的副总裁、杰出工程师。

原文链接：

https://www.microsoft.com/en-us/research/blog/deepspeed-extreme-scale-model-training-for-everyone/