DeepSpeed超大规模模型训练工具

DeepSpeed超大规模模型训练工具

2021年 2 月份发布了 DeepSpeed。这是一个开源深度学习训练优化库，包含的一个新的显存优化技术—— ZeRO（零冗余优化器），通过扩大规模，提升速度，控制成本，提升可用性，极大地推进了大模型训练能力。DeepSpeed 已经帮助研究人员，开发了图灵自然语言生成模型( Turing-NLG)，在发表时，为世界上最大的语言模型（拥有 170 亿参数），有着最佳的精度。在2021年 5 月份发布了 ZeRO-2——支持有着 2000 亿参数的模型训练，与最新技术相比，训练速度可达 10 倍——系列计算、IO 和收敛优化功能，从而助力最快速的 BERT 训练。自那时起，持续高速进行创新，不断突破深度学习模型训练的速度和规模的边界。

这些进展不仅会推动深度学习训练走向极致，同时也让这份技术的使用范围更加广泛——上至数据科学家们在超算上训练，下至在低端集群甚至仅仅一张 GPU 上训练。具体来说，DeepSpeed 加入了 4 项系统性新技术来进一步拓展 AI at Scale 倡议。推动了微软的AI产品与平台的创新。这些技术提供了极为高效的计算、显存和通信的利用效率，并助力训练有着十亿至万亿量级参数的模型。支持超长输入序列，无论在单卡GPU、千卡GPU的高端集群上，还是在慢速以太网的低端集群上，均可以使用。

• 用 3D 并行化实现万亿参数模型训练： DeepSpeed 实现了三种并行方法的灵活组合：ZeRO 支持的数据并行，流水线并行和张量切片模型并行。3D 并行性适应了不同工作负载的需求，支持具有万亿参数的超大型模型，实现了近乎完美的显存扩展性和吞吐量扩展效率。提高的通信效率，使用户可以在网络带宽有限的常规群集上，以 2-7 倍的速度训练有数十亿参数的模型。
• ZeRO-Offload 使 GPU 单卡，能够训练 10 倍大的模型： 为了同时利用 CPU 和 GPU 内存来训练大型模型，扩展了 ZeRO-2。用户在使用带有单张英伟达 V100 GPU 的机器时，可以在不耗尽显存的情况下，运行多达 130 亿个参数的模型，模型规模扩展至现有方法的10倍，保持有竞争力的吞吐量。此功能使数十亿参数的模型训练，更加大众化，，并为许多深度学习从业人员，打开了一扇探索更大更好的模型的窗户。
• 通过 DeepSpeed Sparse Attention，用6倍速度执行10倍长的序列： DeepSpeed提供了稀疏 attention kernel ——一种工具性技术，可支持长序列的模型输入，包括文本输入，图像输入和语音输入。与经典的稠密 Transformer 相比，支持的输入序列长一个数量级，在保持相当的精度下，获得最高 6 倍的执行速度提升。还比最新的稀疏实现快 1.5–3 倍。稀疏 kernel 灵活支持稀疏格式，使用户能够通过自定义稀疏结构进行创新。
• 1 比特 Adam 减少 5 倍通信量： Adam 是一个在大规模深度学习模型训练场景下的，有效的（也许是最广为应用的）优化器。与通信效率优化算法往往不兼容。在跨设备进行分布式扩展时，通信开销可能成为瓶颈。推出了一种 1 比特 Adam 新算法，高效实现。该算法最多可减少 5 倍通信量，实现了与Adam相似的收敛率。在通信受限的场景下，观察到分布式训练速度，提升了 3.5 倍，这使得该算法，可以扩展到不同类型的 GPU 群集和网络环境。

a screenshot of a cell phone

将深入探究这 4 项技术。已经将这些激动人心的优化技术公布在了开源项目 DeepSpeed中。

随着现代 GPU 群集上计算量的快速增长，训练具有惊人的功能的、强大的万亿参数模型，不再是遥不可及的，可能在不久的将来就能实现。DeepSpeed 结合了三项强大的技术，可以训练数万亿规模的模型，扩展到数千个 GPU：数据并行训练，模型并行训练和流水线并行训练。这三者的共生让深度学习训练的规模，远远超出了单独使用每种策略。3D 并行解决了训练万亿参数模型的两个基本挑战：显存效率和计算效率。DeepSpeed 可以扩展至在显存中，放下最巨大的模型，不会牺牲速度。

了解训练巨大模型的显存和计算效率的挑战

显存效率：训练万亿参数模型所需的显存，远远超出了单张 GPU 的显存大小。在使用 Adam 优化器，进行混合精度训练时，存储模型状态量（参数、梯度和优化器状态量）需要约 16TB 的显存。最先进的英伟达 A100 GPU 只有 40 GB 的显存。仅仅为了存储模型状态，需要 400 张这样的 GPU。

激活函数额外消耗的显存，随 batch 大小而增加。batch 设置为1的情况下，训练万亿参数模型，产生超过 1 TB 的激活函数用的显存（后文称为激活显存）。用 checkpoint 处理激活显存，用计算换显存，可以将显存减少到大约20 GB，对于训练而言，仍然过高了。

必须在多个 GPU 设备之间，有效地划分模型状态量和激活显存，才能让这种大模型，在不耗尽显存的情况下，开始训练。

计算效率：经估算端到端训练一个万亿参数的模型，大约需要 5000 Zflops（即 5 后面带有 24 个零；这个估算结果基于 OpenAI 的研究 law of scaling）。训练一个模型需要 4000 张 A100，以 50％ 的计算效率运行大约 100 天。

尽管大型超级计算 GPU 集群，可以拥有超过 4000 个 GPU，但是由于 batch 大小的限制，要在这种规模上实现高计算效率，仍然是一项挑战。计算效率随着计算时间，对通信时间的比例的增加而增加。该比例与 batch 大小成正比。训练模型的 batch 大小，有一个上限——超过这个上限，收敛情况会明显变差。

实际上最大的模型之一，GPT-3 的训练 batch 大小约 1500。如果使用大约 4000 张 GPU, 即使可以自由设置 batch 大小为 4000，每张卡上的 batch 大小，也只有 1，这将影响扩展性。

理解数据并行、模型并行和流水线并行之间的权衡

数据并行是深度学习中的一种普遍使用的技术。每批输入的训练数据，都在数据并行的 worker 之间平分。反向传播后，需要通信规约梯度，保证优化器在各个 worker 上进行相同的更新。数据并行性，具有几个明显的优势，包括计算效率高和实现起来工作量小。数据并行的 batch 大小，随 worker 数量提高，往往无法在不影响收敛性的情况下，一直增加 batch 大小。

• 显存效率：数据并行会在所有 worker 之间，进行模型和优化器的复制，显存效率不高。DeepSpeed 开发了 ZeRO ，一系列用于提高数据并行的显存效率的优化器。依赖于 ZeRO 的 1 阶段，在 worker 之间划分优化器状态量，减少冗余。
• 计算效率：随着提高并行度，每个 worker 执行的计算量是恒定的。数据并行，在小规模上实现近乎线性扩展。在 worker 之间规约梯度的通信开销，跟模型大小成正相关，当模型很大或通信带宽很低时，计算效率会受限。。梯度累积是一种均摊通信成本的一种常用策略。进一步增加batch大小，在本地使用 micro-batch，多次进行正向和反向传播积累梯度后，进行梯度规约和优化器更新。

模型并行是包含范围很广的一类技术。在多个 worker 之间，划分模型的各个层。模型并行性的计算和通信因模型结构而异，在实现上有很大的工作量。DeepSpeed 借用了英伟达的 Megatron-LM，基于 Transformer 的语言模型，提供大规模模型并行功能。模型并行根据 worker 数量，成比例地减少显存使用量，也是这三种并行度中显存效率最高的。代价是计算效率最低。

• 显存效率：模型并行会根据 worker 数量，成比例地减少显存使用量。减少单个网络层的激活显存的唯一方法。DeepSpeed 通过在模型并行 worker 之间，划分激活显存，进一步提高显存效率。
• 计算效率：每次前向和反向传播需要额外通信激活值，模型并行的计算效率很低。模型并行需要高通信带宽，不能很好地扩展到通信带宽受限的节点。每个模型并行worker，都会减少每个通信阶段之间执行的计算量，影响计算效率。模型并行性通常与数据并行性结合使用，在内存和计算效率之间进行权衡。

流水线并行训练引擎，包含在了这次发布的DeepSpeed中！流水线并行将模型的各层划分，可以并行处理。当完成一个 micro-batch 的正向传递时，激活内存将被通信至流水线的下一个阶段。当下一阶段完成反向传播时，将通过管道反向通信梯度。必须同时计算多个 micro-batch，确保流水线的各个阶段能并行计算。已经开发出了几种用于权衡内存和计算效率收敛行为的方法，例如 PipeDream。DeepSpeed通过梯度累积来实现并行，保持与传统数据并行和模型并行训练，在相同的总 batch 大小下收敛。

• 显存效率：流水线并行减少的显存与流水线的阶段数成正比，使模型的大小，可以随 worker 的数量线性扩展。流水线并行不会减少每一层的激活函数的显存占用量。每个 worker 必须存储，同时运行的各个 micro-batch激活值。导致流水线第一阶段的激活内存，与单个 mirco batch 的总激活内存，大致相同。一个万亿参数模型，需要为一个 micro batch，提供大约 19 GB 的显存的激活内存，几乎占到新推出的英伟达 A100 GPU 总显存的一半。
• 计算效率：流水线并行具有最低的通信量，通信量只和在各阶段边界的各层的激活值大小成正比。但是不能无限扩展。像模型并行一样，增加流水线大小，会减少每个流水线阶段的计算量，降低计算与通信的比率。如果要实现好的计算效率，流水线并行，要求其每个阶段的计算负载完美均衡。

此外，流水线并行性，在每个 batch 的开始和结束时，因为需要重新填充，或排空流水线，产生 bubble overhead。使用流水线阶段数的 4 倍或 8 倍的梯度，累积步骤（以及 batch 大小）进行训练，相较于只有一个流水线阶段，分别达到了 81％ 和 90％ 的扩展性。

数据，模型和流水线并行，在提高内存和计算效率方面，均起到特定的作用。图 1 说明了 3D 策略。

显存效率：先将模型的各层，划分到不同的流水线阶段，把每个阶段的层通过模型并行，进行划分。这种 2D 组合同时减少了模型、优化器和激活函数所消耗的内存。不能在不引入通信开销的情况下，无限划分模型，通信开销会限制计算效率。

计算效率：为了在不牺牲计算效率，将 worker 数量扩展至超出模型和流水线并行规模，使用了 ZeRO 支持的数据并行功能（ZeRO-DP）。ZeRO-DP 不仅可以通过划分优化器状态量，进一步提高显存利用效率，通过利用基于通信拓扑的映射关系，以最小的通信开销，扩展到任意数量的 GPU。

基于通信拓扑的 3D 映射（图2）：通过利用两个关键的架构属性，将 3D 并行中的每个维度，仔细映射到 worker 上，实现最大的计算效率。

1. 优化节点内和节点间的通信带宽：模型并行是这三种策略中通信开销最大的，优先考虑将模型并行 worker 组放置在节点内，利用更大的节点内带宽。基于英伟达 Megatron-LM，进行了张量切分式的模型并行。当模型并行组不占满节点内的所有 worker 时，选择将数据并行组放置在节点内。不然，就跨节点进行数据并行。流水线并行的通信量最低，可以跨节点调度流水线的各个阶段，不受通信带宽的限制。
2. 通过并行通信增大带宽：每个数据并行组，需要通信的梯度量，随着流水线和模型并行的规模线性减小，总通信量少于单纯使用数据并行。此外，每个数据并行组，在局部的一小部分 worker 内部，独立进行通信，组间通信可以相互并行。通过减少通信量和增加局部性与并行性，数据并行通信的有效带宽被增大了。

该图显示了一个有 32 个 worker 进行 3D 并行的例子。神经网络的各层分为四个流水线阶段。每个流水线阶段中的层在四个模型并行 worker 之间进一步划分。最后，每个流水线阶段有两个数据并行实例，且 ZeRO 在这 2 个副本之间划分优化器状态量。

图 1：一个有 32 个 worker 进行 3D 并行的例子。神经网络的各层分为四个流水线阶段。每个流水线阶段中的层在四个模型并行 worker 之间，进一步划分。每个流水线阶段，有两个数据并行实例，ZeRO 在这 2 个副本之间，划分优化器状态量。

彩色块显示图 1 中的 worker 到八个节点（每个节点有四个 GPU）的系统上的 GPU 的映射。同一颜色的 GPU 在同一节点上。

图 2：图 1 中的 worker 到八个节点（每个节点有四个 GPU）的系统上的 GPU 的映射。同一颜色的 GPU 在同一节点上。

了解关于 3D 并行训练万亿参数模型的更多信息

使用 8 路模型并行，64 路流水线并行和 8 路数据并行，在 4096 个英伟达 A100 GPU 上，扩展训练一个万亿参数模型。

通过结合模型并行和流水线并行，3D 并行可实现出色的内存效率和跨多个节点的高效计算效率。模型并行性，提高了节点内的激活内存和模型状态量的存储效率，流水线并行，相较于仅使用模型并行，可以在不牺牲计算效率的情况下，跨节点高效存储模型状态。在 micro-batch 大小为 1 的万亿参数例子中，在使用激活值 checkpoint及上述 3D 并行后，模型状态量会消耗 30 GB 的显存，划分后的激活值消耗 2.5 GB 的内存。总显存占用为 32.5 GB，就能够使用具有 40 GB 内存的英伟达 A100 GPU，容纳和训练这样的模型了。

结合模型并行与流水线并行，可以使流水线并行在非常小的 batch 下，以最小的 bubble overhead，实现高计算效率。在 8 路模型并行下，每个模型使用 micro-batch 为 1 个微批处理，导致每个 GPU 的有效 micro-batch 大小为 1/8。使用 8 倍于管道并行度的梯度累加步骤，让每张 GPU 上的总累计 batch 大小为 1，流水并行处理，可以实现 90％ 的计算效率。与数据并行性结合使用时，让 4096 张 GPU 上的总有效 batch 大小为 4096，仍然可以达到 90％ 的流水线效率。

数据并行怎样影响计算效率呢？难道数据并行不是需要每张 GPU，都有大 batch 才能保持高效吗？

模型并行可以将每张GPU上的有效 batch 大小，减小到小于 1。使流水线并行，即使在小 batch 下，仍可以隐藏流水线 bubble overhead。通过跨节点使用流水线并行性，可以让流水线每个阶段的数据并行节点之间的独立进行通信，与其它流水线并行进行。实际上，在高端 GPU 集群中常见的，完全连接的网络拓扑中，对可用于数据并行训练的有效通信带宽，具有重要意义。流水线阶段中的每个节点，都可以与其对应的数据并行节点，并行通信，有效的通信带宽与流水线阶段数成正比。通过设置64个并行流水线，有效带宽将变为往返单个节点的带宽的 64 倍。流水线并行带给数据并行，如此大的有效带宽，使数据并行在计算与通信比率非常低的小 batch 情况下，实现高效扩展。

DeepSpeed 可以只用 800 张英伟达 V100 GPU，训练具有一个万亿参数的语言模型（图 3）。展示了模型大小和训练吞吐量，可以观察到显存和计算效率，随模型的大小的扩展线性增长。在各种配置中，在每个 GPU 上训练，大约 14 亿个参数，这是单个 GPU 在不耗尽内存的情况下，可以支持的最大模型大小，表明了完美的显存扩展性。获得了接近完美的线性计算效率扩展，每张 V100 GPU 的吞吐量为 47 Tflops。对于上述的硬件，这是令人印象深刻的扩展性和吞吐量。

图3：模型大小（以十亿个参数为单位）和训练吞吐量（以 Pflops 为单位）随 GPU 数量变化趋势的图表。DeepSpeed 可以使用 800 张，具有 32 GB 内存的英伟达 V100 Tensor Core GPU，训练有 1 万亿个参数的模型。每种配置都使用 NVIDIA Megatron-LM 提供的，16路模型并行性，剩余的GPU，负责进行流水线并行。万亿参数模型，具有 298 层 Transformer，隐藏层大小为 17408，训练的序列长度为 2048，batch 大小 2048。对于较小的模型，根据 GPU 数量，按比例减少了 Transformer 层的数量和 batch 大小。

深入研究 3D 并行，如何加速训练 GPT-3 规模的模型

图 4：使用 2D 和 3D 并行，使用 800 个 GPU，训练具有 1800 亿参数的 GPT-3 规模模型的系统性能。模型具有 100 个 Transformer 层，隐藏层尺寸为 12288，有 96 个 attention head。训练使用的 batch 大小为 2048，序列长度为 2048。ZeRO-1可以跟数据并行结合使用。P、M 和 D 分别表示流水线，模型和数据并行维度。

在图 4 中，使用具有超过 1,750 亿个参数的最新 GPT-3 模型架构，作为 3D 并行性的基准：

• 首先评估了 2D 配置（C1-C3）。配置 C1 和 C2 仅使用流水线和模型并行——可以训练模型，由于过度分解模型，导致吞吐量较低，GPU 利用率较低。C3 尝试仅使用流水线和数据并行，不通过 Megatron 的模型并行，减少激活量，无法解决显存不足的问题。
• 3D 配置（C4-C10）依次增加了流水线并行度；中间的平衡了并行性的配置，实现最佳性能，实现了显存，计算和通信效率三高。
• 最佳的 3D 方法，每个GPU可实现 49 Tflops，超过硬件的理论峰值的 40％。

混合并行如何在低带宽集群上，7 倍加速训练 GPT-2

训练了一个 15 亿参数的 GPT-2 模型，在图 5，展示了混合并行的通信优势。为了突出展示训练的通信阶段，训练在节点间带宽较低的四节点的群集上进行：

• 模型并行在这种情况下没有优势，因为模型较小，节点内带宽较低。
• 流水线并行的通信量，比配置数据和模型并行的情况，小一个数量级。在 batch 较小时，训练速度快 7 倍。
• 数据并行使用通过梯度累积增加 batch 大小，均摊通信开销，在更大的 batch 大小下，配置了流水线并行的情况的性能，仍是数据并行的两倍。
• 混合流水线和数据并行配置，通过将数据并行组限制在节点内的 GPU 上，避免了梯度通信瓶颈，因此梯度通信受益于更快的节点内带宽。

图 5：在训练序列长度为 1024 的 GPT-2（1.5B 参数）时，吞吐量与 batch 大小的关系。使用四个节点，每个节点配备四个具有 16 GB 内存的 V100 GPU 训练。GPU 之间用每秒 50 Gbps 的节点内带宽和 4 Gbps 的节点间带宽连接。DP 表示启用 ZeRO-1 的数据并行性。所有方法都通过增加梯度累积的步数，扩展批量大小。

ZeRO-Offload 通过同时利用GPU和宿主机 CPU 的计算和存储资源，提升了较少的 GPU 资源下，可以高效训练的最大模型规模。在单张 V100 上，进行最高至 1300 亿参数的模型训练，10 倍于当前最高水平，保持每 GPU 30Tflop 的高训练吞吐量。

通过使单 GPU，具备训练数十亿参数的模型的能力，ZeRO-Offload 让大模型训练变得亲民，让硬件资源有限的深度学习，从业者也能参与其中。

在单 GPU 上使用默认的 PyTorch 和 ZeRO-Offload，能训练的最大模型规模的柱状图。

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

ZeRO-Offload 背后的核心技术，在 ZeRO-2 的基础上，将优化器状态和梯度，卸至 CPU 内存。让 ZeRO-Offload 能最大程度，降低拷贝至 CPU 导致的计算效率损失，达到和 ZeRO-2 相同，甚至有时超过的效率。下图展示了 Zero-OffLoad 的架构：

图7: ZeRO-Offload 概述。

了解 ZeRO-Offload 是如何在单GPU上，训练数十亿参数模型的

训练 GPT 和 T5，这样有数十亿参数的模型，需要多个 GPU 来存储模型和状态量。大模型训练大多通过跨 GPU 的模型并行，解决显存限制问题。最近，发布了 ZeRO，一个高效利用显存的优化器，将模型状态量（优化器状态量、梯度和模型参数），分布在多个并行 GPU 上，让数十亿参数模型，在不使用模型并行的情况下，进行训练。ZeRO 还是需要大量数据并行的 GPU，保存划分后的模型状态量，只有少数人，有条件进行这种模型训练。

ZeRO-Offload 让单 GPU 可以进行大模型训练，使这种训练变得平民化。为了在不使用多个 GPU 的情况下，训练数十亿个参数的模型，ZeRO-Offload 继承了 ZeRO-2 的划分优化器状态量和梯度的方法。与ZeRO-2 不同，ZeRO-Offload没有在每个 GPU 上保存一部分优化器状态量和梯度，把两者都移到了本机内存上。Optimizer 状态，在整体训练过程中，都保存在内存中。梯度则是在反向计算过程中，在 GPU 上进行计算，通过 reduce-scatter 进行平均，每个数据并行进程，把平均后的梯度，卸到 CPU 上（图7中的 g offload），弃掉不属于自己负责的部分。

一旦梯度到了 CPU 上，划分后的优化状态量，就会并行地在 CPU 上，进行更新(图7中的 p update)。在更新进行完后，划分后的参数就被移回GPU并用 all gather 操作进行更新 (图7中的 g swap）。Zero-Offload 也通过使用不同 CUDA stream，重叠通信（如 g offload 和 g swap）和计算（如反向传播和 p update)，提高训练效率。

从模型规模，训练速度和扩展性看 ZeRO-Offload 的优势

10 倍模型扩展：在单张 32GB V100 GPU 上，图 6 显示 PyTorch 能最多训练有 13 亿个参数的模型，ZeRO-Offload 能训练 130 亿个参数的模型， PyTorch 的 10 倍。ZeRO-Offload 在整个训练过程中，将消耗了大部分 GPU 显存的优化器状态，保留在本机内存中，在反向传播过程中，将计算出来的梯度移至 CPU。节省的 GPU 显存，可用于训练更大的模型。

高效的训练吞吐量：如图 8 所示，在训练 100 亿参数模型时，即使仅使用单个 GPU 进行训练，使用 ZeRO-Offload，可让每个 GPU 有超过 30 Tflops 的吞吐量，吞吐量随 GPU 数量增长，呈近完美的线性增长。

ZeRO-Offload 是 ZeRO-2 的完美补充，支持在少量 GPU 上，高效训练大型模型。通过利用 CPU 内存，减少了模型所需的 GPU 显存，ZeRO-Offload 让在 1 到 16 个 GPU，训练大模型变得可行。在 32 个 GPU 上，ZeRO-Offload 的性能，略高于 ZeRO-2; 性能提升来源于 ZeRO-Offload 节省的 GPU 显存，可以在更大 batch 下，训练了模型，尽管存在拷贝至 CPU 的开销，GPU 计算效率仍然可以提高。在有更多的 GPU（例如 64 和 128）的情况下，ZeRO-2 的性能，优于 ZeRO-Offload，两者都可以运行类似大小的batch，ZeRO-2 没有将数据移至 CPU 的开销，GPU 上进行优化器更新，比 CPU 上快得多。总之，ZeRO-Offload 是 ZeRO-2 的补充，扩展了 ZeRO 家族的优化范围，从单台设备到数千台设备，有大型模型训练的优化方案。

使用 ZeRO-Offload 和 ZeRO-2 在 128 张 GPU 上，训练有 100 亿参数的 GPT-2 模型的的吞吐量的柱状图。

图 8：使用 128 张 GPU，训练 100 亿参数 GPT-2 模型的 ZeRO-Offload 和 ZeRO-2 的训练吞吐量比较。

基于注意力机制的深度学习模型（例如，Transformers），在捕获输入序列中的 token 之间的关系（即使是两者之间距离很长）方面，非常有效。与文本，图像和语音相关的输入配合使用。这些输入的序列长度可至数千 token。注意力模块有效地捕获了长序列内的依赖关系，对长序列输入的支持，计算量和显存的限制。计算量和显存需求，关于序列长度\\(n\\)呈二次方级增长。

为了解决此限制，DeepSpeed 提供了一套稀疏注意力 kernel——一种工具性技术，通过块状稀疏计算，将注意力计算的计算和显存需求，降低几个数量级。这套工具不仅缓解了注意力计算的内存瓶颈，稀疏计算非常高效。API 可以方便地集成进，任何基于 Transformer 的模型。除了提供各种稀疏结构外，可以灵活处理任何用户，自定义的块状稀疏结构。

稀疏注意力（SA）可以设计计算，靠近的 token 之间的局部注意力，或通过使用局部注意力计算，得到 summary token，进而得到全局注意力。SA 既支持随机注意力，也支持局部、全局和随机注意力的任意组合，如图 10 中的蓝色，橙色和绿色块。使SA将内存占用减小到\\(O(wn)\\)，其中1\\(<w≤n \\)是一个参数，取决于注意力结构。

彩色小方块显示可变的稀疏度结构

图 10：可变稀疏结构

在 GPU 上的高效实现：尽管稀疏注意力的基本实现，节省显存，但在计算上，可能会比稠密计算要差。稀疏数据导致了内存访问的分散性。开发高效的稀疏内核，通常是颇具挑战性的，尤其是在 GPU 上。DeepSpeed 提供了在 Triton 中开发的高效的稀疏注意力 kernel。这些 kernel 呈块状稀疏范式结构，实现对齐的内存访问，减少GPU线程分支平衡处理器上的工作负载。

系统性能：如图11所示，SA 支持 10 倍长的序列和最高 6.3 倍的计算提速。左图显示了可在 BERT-Base 和 BERT-Large 中，运行的最长序列长度。实验有以下三种设置：稠密模式，具有激活 checkpoint 的稠密模式和具有激活 checkpoint 的稀疏（SA）模式。与 BERT-Base 和 BERT-Large 的稠密模式相比，SA 的序列分别长 10 倍和 16 倍。与稠密模式相比，SA 减少了总计算量，提高了训练速度：提高的效率随着序列长度的增加而提高，对于 BERT-Base 而言，提升速度高达 6.3 倍，对于 BERT-Large，高达 5.3 倍。

图11：BERT 模型的可支持的最大序列长度（左）；在单英伟达 V100 GPU 上，训练具有不同序列长度的 BERT-Base（中）和 BERT-Large（右）的时间。

了解 SA 如何使其准确率与全稠密注意力相当，甚至比它更高

涉及稀疏注意力的相关工作（Sparse Transformer，Longformer，BigBird），均显示出比全注意力更高的准确性，与经验一致。除了降低内存开销和加快计算速度外，在生产模型中观察到 SA，有更高准确性，更快收敛的情况。下图说明了训练基于 BERT 的长文本，理解（序列长度 2048）生产模型的准确性。实验在以下三种设置中进行：从头开始进行稠密训练，从头开进行 SA 训练，从使用序列长度为 512 的密集型 checkpoint，继续进行 SA 训练。对于从头开始进行预训练，SA较于稠密设置收敛的速度更高，精度更好。时间和准确性，从用 SA 继续训练，预先训练好的 checkpoint 的效果，甚至更好。

图12：长文本理解应用的准确性

了解 SA 与最新的 LongFormer 的比较情况

将 SA 与 Longformer（一种最新的稀疏结构及其实现）进行了比较。在实验中，SA 使用“Fixed”稀疏性。两种实现的准确性相当。在系统性能方面，SA在训练和推断方面均优于Longformer：

• 运行 Wikitext103 上的预训练MLM的速度，提高了 1.5 倍
• BERT-Base 的推理速度，提高3倍（batch 大小 1，序列长度 2,048）

处理任何块状稀疏结构的灵活性： DeepSpeed 稀疏注意力套件，不针对任何特定的稀疏结构，能有效支持模型研究人员探索任何块状稀疏结构。添加了流行的稀疏结构，例如 Fixed（来自OpenAI稀疏Transformer），[BigBird](https://arxiv.org/pdf/2007.14062 .pdf)（来自Google）和BSLongformer（AI2 Longformer的块稀疏实现）。定义了一个具有“可变”结构的模板，如图 10 所示，模板可用于简单地自定义任何随机，局部或全局注意力模式的块状稀疏结构。

大型模型（如 BERT 和 GPT-3）的扩展训练，需要基于模型设计，体系结构和系统功能的细致优化。从系统的角度来看，通信效率已成为主要的瓶颈，在使用标准 TCP 且网络带宽有限的商用系统上。

通信压缩是减少在此类系统上的训练时间的重要技术。压缩通信的最有效方法之一，误差补偿压缩，即使在1比特压缩下，可以提供稳定的收敛速度。但是，最新的误差补偿技术，仅适用于一些和梯度线性相关的简单优化器，随机梯度下降（SGD）和 Momentum SGD。这些技术无法和 Adam 之类的非线性优化器整合，在许多任务（包括训练类似 BERT 的模型）中，带来了最好的收敛率和精度。

对于像 Adam 之类的强大优化器，依赖于梯度的非线性特征（在方差项上），开发基于误差补偿的压缩技术，颇具挑战性的工作，限制了先进的通信压缩技术的实用价值。

理解经典压缩技术的背景

通信压缩的一种方法是1比特压缩，可以被表示为：

在这种压缩中，用 1 比特表示每个数字，将内存需求减少 32 倍。这种直接的方法会大大降低收敛速度，没什么实用价值。研究表明，通过使用误差补偿压缩，有望在通信压缩下，保证几乎相同的收敛率。

误差补偿的思想，可以概括为：1）进行压缩，2）记忆压缩误差，3）在下一次迭代中，把压缩误差加回来。对于 SGD，误差压缩相当于：

其中\\(C(?)\\)是1比特压缩算子。这种误差压缩的优点在于压缩误差的历史值\\(e_t\\)和\\(e_t-1\\)最终会相互抵消, 这使得：

该策略已经被证明，适用于所有线性依赖于梯度的优化算法，例如 SGD 和 Momentum SGD。

了解将误差补偿应用于 Adam 的挑战

在下面提供了 Adam 算法的概述。更新规则如下：

如上图的公式所示，方差项 \\(v_t\\) 和梯度 \\(g_t\\) 呈非线程关系。如果对 Adam 进行普通的误差补偿，发现（见图 13）Adam 将无法收敛。

图13：由于对梯度的非线性依赖，误差补偿压缩不适用于 Adam

为了在使用 Adam 优化器时压缩通信，开发了 1 比特 Adam，通过预处理解决了梯度中的非线性依赖问题。观察到非线性项方差（\\(v_t\\)）的变化幅度，在几个训练周期后显著降低，将 \\(v_t\\) 设置为常数不会改变收敛速度。1 位 Adam 优化器，由两部分组成（如图 14 所示）：预热阶段，本质上就是原始的 Adam 算法。压缩阶段，使方差项保持恒定，将剩余的线性项（即动量）压缩为 1 位表示形式。

算法的压缩阶段，由阈值参数控制（如图 14 所示）。当检测到“方差”的变化，降至某个阈值以下时，切换到压缩阶段。研究表明，热身阶段，需要全部训练步骤的 15-20％。

进一步了解 1 比特 Adam 的底层机制

1 比特 Adam 的权重按以下公式进行更新。对于第 i 个 worker，在压缩阶段：

a screenshot of text

a screenshot of a cell phone

图 14：使用经典 Adam 算法和使用 1 比特压缩 Adam 算法，进行分布式训练的流程对比

除了算法上的挑战外，在训练系统中应用 1 比特 Adam，有两个系统挑战。首先，需要具备将动量转换为 1 比特表示形式的功能的高效 kernel。其次，需要高效的通信方案，在不同的 GPU 之间，传输压缩后的动量。压缩的目的是减少总体训练时间，使带宽受限的商品系统，用来训练大型模型。在 DeepSpeed 中，解决了这些具有挑战性的问题，针对在通信效率受限的系统上，进行训练的场景，对 1 比特 Adam 实现，进行了全面的优化。

1 比特 Adam 提供了和 Adam 相同的收敛能力，最多可以减少 5 倍的通信量，进行 BERT-Large 预训练任务时，可达最高 3.5 倍的吞吐量，用于 SQuAD fine-tuning 任务时，可达 2.7 倍的高吞吐量。端到端吞吐量的提高，来源于在压缩阶段观察到的 6.6 倍（图 15 左）和 6.2 倍（图 15 右）速度提升。1 位 Adam 优化器，在 40 Gb 以太网系统上的扩展性非常好，性能可与 Adam 在 40 Gb InfiniBand QDR 系统上的扩展性相媲美。基于 iPerf 基准，40 Gb 以太网上的有效带宽为 4.1 Gbps，基于 InfiniBand perftest 微基准，InfiniBand 提供了 32 Gbps 的近峰带宽。

图 15：NVIDIA V100 GPU 上的 BERT-Large 预训练（左）和 SQuAD fine-tuning（右）的 1 比特 Adam 扩展性。BERT 预训练的 batch 大小为 16/GPU，SQuAD fine-tuning 为 3/GPU。

深入研究 1 比特 Adam 的评测结果

与 Adam 相同的收敛性：使用 1 比特 Adam 的一个主要问题是收敛速度。发现在使用相同数量的训练样本时，1 比特 Adam，可以达到相同的收敛速度和相当的性能，见图 16。

图 16：使用相同数量的训练样本，1 比特 Adam，可以像 Adam 一样收敛。

表 1 显示了 BERT-Base 和 BERT-Large 的详细结果。看到，对于未压缩和压缩情况，1 比特 Adam 的性能，均与原始模型相当，有些则优于原始模型。

表 1：在各种测试任务上验证 1 比特 Adam 的正确性

最多可减少 5 倍的通信量： 1 比特 Adam，提供了与 Adam 相同的收敛能力，在压缩阶段（对于 16 位（FP16）训练），将通信量减少了 16 倍。 对于 BERT 预训练模型，观察到预热阶段，仅为端到端训练时间的 15％，总体通信减少了 5 倍。

原始 Adam 和 1 比特 Adam 的通信量之比的公式如下：

1 / (warmup + (1 – warmup)/16)

1 比特 Adam 使训练 BERT-Large 的速度快 3.5 倍： 提供了在两个具有有限带宽限制的系统上，训练 BERT-Large 的结果：1）40 Gbps 以太网（图 17 左）和 2）40 Gbps InfiniBand QDR（图 17 右）。在压缩阶段，发现使用以太网的系统吞吐量，提高了 6.6 倍，使用 InfiniBand的系统吞吐量，提高了 2 倍，端到端的速度（包括预热和压缩阶段），分别提高了 3.5 倍和 2.7 倍。1 比特 Adam 主要得益于通信量的减少（因为对动量通信的压缩实现），及自定义的 allreduce 操作，操作通过高效的 1 比特，无阻塞 gather 和一个 allgather 操作实现。

使用 LAMB 而不是 Adam 优化器，进行 BERT 预训练，通过增加总 batch 大小，减少通信量。1 比特的 Adam，避免了这种要求严格的超参数调参。根据经验，大 batch 下进行调参，通常会更加困难。此外，1 比特 Adam，对于临界批处理量较小（无法在大 batch 下良好收敛，例如许多 fine-tuning 任务）的工作也非常适用。

图 17：在压缩阶段，使用 1 比特 Adam 在 40 Gbps 以太网（左）和 InfiniBand（右）上，进行 BERT-Large 训练时的性能

1 比特 Adam 使 SQuAD fine-tuning 任务加速 2.7 倍： 1 比特 Adam，不仅在大规模训练任务上提供扩展性，而且在 SQuAD 微调之类的任务上，也有效果。如图 18 所示，1 比特 Adam，在基于以太网和基于 InfiniBand 的系统上很好地扩展，在基于以太网的系统上，提供高达 6.2 倍的高吞吐量（在压缩阶段），带来端到端的 2.7 倍提速（预热阶段占 25％，压缩阶段占 75％）。对于 SQuAD fine-tuning，观察到总 batch 大小为 96 时，F1 得分最高。 batch 大小大于此值降低收敛率，需要额外的超参数调整。为了扩展到 32 个 GPU，在每个 GPU 上运行值为 3-4 的小 batch。使得 fine-tuning 任务的通信强度大，难以扩展。1 比特 Adam，很好地解决了扩展性的难题，在不增大 batch 的情况下，减少了 3.4 倍的通信量，实现了 2.7 倍的端到端加速。

图 18：在 40 Gbps 以太网（左）和 InfiniBand（右）上的 SQuAD fine-tuning 任务中，使用 1 比特 Adam 时，压缩阶段的性能。

一群热衷于大规模系统性能优化的研究员和工程师——Samyam Rajbhandari, Jeff Rasley, Olatunji Ruwase, Reza Yazdani Aminabadi, Elton Zheng, Arash Ashari, Jing Zhao, Minjia Zhang, Niranjan Uma Naresh, Shaden Smith, Ammar Ahmad Awan, Conglong Li, Yuxiong He (team lead)。

参考链接：

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

https://zhuanlan.zhihu.com/p/343570325