[论文笔记]ZeRO: Memory Optimizations Toward Training Trillion Parameter Models

引言

今天带来ZeRO: Memory Optimizations Toward Training Trillion Parameter Models的论文笔记。

大型深度模型提供了显著的准确性提升，但训练数十亿到数万亿个参数是具有挑战性的。现有的解决方案，如数据并行和模型并行，存在基本的局限性，难以将这些模型适配到有限的设备内存中，同时又能达到计算、通信和开发效率。

作者开发了一种新颖的解决方案，零冗余优化器(Zero Redundancy Optimizer, ZeRO)，以优化内存，大幅提升训练速度，同时增加可以高效训练的模型大小。ZeRO在数据和模型并行训练中消除了内存冗余，同时保持低通信量和高计算粒度，可以根据设备数量按比例扩展模型大小，保持高效的可持续性。

1. 总体介绍

深度学习模型正变得越来越大，模型规模的增加带来了显著的准确性提升。在自然语言处理领域，Transformer技术使得大模型得以问世。为了使模型规模从数十亿参数继续增长到数万亿参数，我们面临着训练它们的挑战——这些模型显然无法适应单一设备的内存，例如GPU或TPU，单纯增加设备数量也无法有效扩展训练。

基本的数据并行(data parallelism, DP)并不能减少每个设备的内存，在32GB内存的GPU上，参数超过14亿的模型很容易就会出现内存不足的情况。其他现有的解决方案，如流水线并行(Pipeline Parallelism,PP)、模型并行(Model Parallelism,MP)、**CPU卸载(CPU-Offloading)**等，虽然在功能、可用性以及内存和计算/通信效率之间进行了权衡，但这些都是实现快速和大规模训练的关键因素。

DP： 将模型复制到多个GPU上，每个GPU设备可以并行接受不同的数据批次，计算的梯度结果也不一致，因此需要在模型间交换梯度计算平均值。由于模型在每个GPU上复制了一份，产生了很多GPU内存消耗。

MP： 当模型本身的参数量很大无法在单卡上加载时，无法进行数据并行。模型并行选择将模型按层拆分，每个GPU只加载模型参数的一部分，在前向传播过程中汇总，数据将依次通过各个GPU，带来的时间消耗很大。

PP： 流水线并行是MP的一种变体，将每批输入数据拆分成很多小的微批次，以减少设备的空闲等待时间。只有在整个模型处理完每个微批次后，才会更新模型参数，因此当其他GPU设备仍在处理上一个微批次时，空闲的GPU就可以开始处理下一个微批次。

Megatron-LM提出了1F1B流水线策略，即一个前向通道和一个后向通道。1F1B流水线策略引入了任务调度机制，使得下游设备能够在等待上游计算的同时执行其他可并行的任务，从而提高设备的利用率。

在现存的训练大模型的解决方案中，模型并行或许是最有前景的。然而，MP无法进一步扩展到更大的模型。模型并行是通过将模型垂直切分，将每层的计算和参数分配到多个设备上，从而需要在每一层之间进行大量通信。因此，它们在单个节点内工作良好，因为节点间的GPU通信带宽很高，但在超出单节点时效率迅速下降。

作者在两个DGX-2节点上使用Megatron-LM测试了一个40亿参数的模型，观察到每个V100 GPU的性能约为5T flops(不到硬件峰值的5%)。

那么，如何克服现有解决方案的局限性，更高效地训练大模型呢？为了解答这个问题，作者首先分析了现有系统在模型训练中的内存消耗全景，并将其分类为两个部分：1)对于大模型，大部分内存被模型状态(model states)占用，包括优化器状态(optimizer states, 如Adam中的动量和方差)、梯度(gradients)和参数(parameters)；2）剩余的内存被激活值(activation)、临时缓冲区( temporary buffers)和不可用的碎片内存(unusable fragmented memory)消耗，将其统称为剩余状态(residual states)。

作者开发了ZeRO——零冗余优化器(Zero Redundancy Optimizer)，旨在优化这两部分的内存效率，同时实现高计算和通信效率。由于这两部分面临不同的挑战，分别来讨论其解决方案。

优化模型状态内存

模型状态在训练过程中通常消耗最多的内存，但现有的方法，如数据并行(DP)和模型并行(MP)，并未提供令人满意的解决方案。DP在计算/通信效率方面表现良好，但内存效率较差，而MP则可能导致计算/通信效率不理想。具体而言，DP在所有数据并行进程中复制所有模型状态，导致冗余内存消耗；而MP通过将这些状态进行划分以获得高内存效率，但通常会导致过于细粒度的计算和昂贵的通信，从而降低扩展效率。此外，所有这些方法在整个训练过程中都静态维护所需的所有模型状态，尽管在训练期间并不是所有的模型状态都是随时需要的。

因此作者开发了ZeRO-DP(ZeRO支持的数据并行)，它在实现DP的计算/通信效率的同时，也达成了MP的内存效率。ZeRO-DP通过对模型状态进行划分而不是复制，消除了数据并行进程之间的内存状态冗余，并通过在训练期间使用动态通信调度，保持DP的计算粒度和通信量，从而保持计算/通信效率。

图1：比较采用ZeRO-DP优化的三个阶段下，每个设备的模型状态内存消耗。$\Psi$表示模型大小(参数量），$K$表示优化器状态的内存乘子，$N_d$表示数据并行度。在此示例中，我们假设模型大小为$\Psi =7.5B$，数据并行度$N_d=64$，基于使用Adam优化器的模型并行训练，优化器状态的内存乘子$K=12$，Adam优化器会以FP32保存模型参数备份，一阶动量和二阶动量也都采用FP32格式存储，因此是12。

ZeRO-DP主要有三个优化阶段(图1)，分别对应优化器状态、梯度和参数的划分。当累积启用时：

1. 优化器状态划分($P_{os}$)：内存减少4倍，与DP的通信量相同；
2. 添加梯度划分($P_{os+g}$)：内存减少8倍，与DP的通信量相同；
3. 添加参数划分$(P_{os+g+p}$)：内存减少与数据并行程度$N_d$成线性关系。

例如，在64个GPU上划分($N_d=64$)将实现64倍的内存减少。通信量有适度增加，约50%。

ZeRO-DP消除了内存冗余，使集群的全部聚合内存容量可用。在启用所有三个阶段的情况下，ZeRO能够在仅有1024个NVIDIA GPU上训练一个万亿参数模型。一个采用16位精度的优化器(如Adam)的大规模万亿参数模型大约需要16TB的内存来存储优化器状态、梯度和参数。将16TB除以1024，得到16GB，这在GPU(具有32GB设备内存)可接受的范围内。

优化剩余状态内存

在ZeRO-DP提升了模型状态的内存效率之后，激活、临时缓冲区和不可用内存碎片所消耗的剩余内存可能成为次要的内存瓶颈。作者开发了ZeRO-R，以分别优化这三种因素所消耗的剩余内存。

1. 对于激活(在前向传播中存储以便执行反向传播)，检查点(checkpointing)¹技术有所帮助，但对于大模型而言仍不够。因此，ZeRO-R通过激活分区(activation partitioning)优化激活内存，识别并消除现有模型并行方法中的激活复制，同时在适当时将激活卸载到CPU。

2. ZeRO-R为临时缓冲区定义了适当的大小，以期在内存和计算效率之间达到平衡。

3. 在训练过程中，由于不同张量的生命周期差异，导致内存碎片化。由于碎片化而缺乏连续内存可能导致内存分配失败，尽管可用内存充足。ZeRO-R基于不同张量的生命周期主动管理内存，防止内存碎片化。

ZeRO-DP和ZeRO-R结合在一起，形成了一个强大的深度学习训练内存优化系统，统称为ZeRO。

ZeRO与模型并行

由于ZeRO消除了数据并行中的内存低效性，自然会问：我们是否还需要模型并行，何时需要？ZeRO如何与模型并行协同工作？在使用ZeRO的情况下，模型并行作为单纯为了适应大模型而采用的选项变得不那么吸引人。ZeRO-DP在减少每个设备的内存占用方面至少与模型并行同样有效，有时在模型并行无法均匀划分模型时甚至更有效。它的扩展效率也相当或更好。此外，数据并行性易于使用，广泛适用于不同的工作负载。

尽管如此，仍然存在一些情况下我们希望利用模型并行：1)在与ZeRO-R结合使用时，模型并行可以减少非常大型模型的激活内存占用。2) 对于较小的模型，当仅使用数据并行时聚合批量大小过大而导致收敛效果不佳时，模型并行也可以带来好处。在这些情况下，可以将ZeRO与模型并行结合使用，以适应可接受的聚合批量大小。

ZeRO可以与模型并行结合，从而在每个设备上实现理论上的最大内存减少，减少倍数为$N_d\times N_m$，其中$N_d$为数据并行度，$N_m$为模型并行度。

实现与评估

ZeRO中的完整优化集使我们能够在今天的高端硬件集群上运行万亿参数的模型，然而，硬件计算能力仍然太有限，训练时间可能过长(超过1年)。因此，作者实现的重点是高效支持参数量是SOTA的10倍(约1000亿参数)的模型，同时仍然在当前硬件的计算能力范围内。作者实现并评估了ZeRO的部分优化，称为ZeRO-100B——这是ZeRO-DP的$P_{os+g}$​与ZeRO-R的结合。结果显示：

模型规模 结合模型并行，ZeRO-100B能高效运行170亿参数的模型，如图2所示。这相比于SOTA，模型规模增加了超过8倍。

速度 改进的内存效率提升了吞吐量和训练速度。如图2所示，ZeRO能在一个400个Nvidia V100 GPU集群上运行100亿参数的模型，每个GPU超过38 TFlops，整体性能超过15 Petaflops。这在相同模型规模下相比于SOTA训练速度提升超过10倍。

可扩展性 在64-400个GPU的范围内，观察到超线性加速，其性能在双倍增加GPU数量时超过两倍。这是ZeRO-DP的特性，在增加数据并行度时减少模型状态的内存占用，使我们能在每个GPU上适配更大的批量大小，从而提高性能。

大型模型训练的民主化 ZeRO-100B使数据科学家能够在没有任何需要重构模型的模型并行或流水线并行的情况下训练最多130亿参数的模型。数据科学家因此可以自由实验大型模型，而不用担心并行性。相比之下，现有系统在1.4亿参数的模型上会出现内存不足的情况。

新一代最先进模型 ZeRO支持了最大的语言模型，参数为170亿，并且取得了创纪录的准确性，即Turing-NLG。

2. 相关工作

2.1 数据、模型和流水线并行

并行化是大规模训练大型模型的关键策略。对于能够适应设备内存进行训练的模型，可以使用数据并行（DP）将训练扩展到多个设备。在数据并行中，模型参数在每个设备上被复制。在每一步中，mini-batch被均匀地分配到所有数据并行进程，使每个进程在不同的数据样本子集上执行前向和后向传播，并使用跨进程的平均梯度在本地更新模型。

当模型无法适应设备内存时，模型并行（MP）和流水线并行（PP）会在进程之间分割模型，分别从垂直和水平的方式进行。

流水线并行水平分割模型，通过在不同设备上运行每个部分，并使用micro-batch处理来隐藏流水线气泡。由于水平分割和micro-batch，某些模型功能（如绑定权重和批归一化）的实现变得困难。

流行的流水线并行实现，如G-pipe，在分割模型参数和总激活值时，需要一个与流水线分区数量成比例的批量大小来隐藏流水线气泡。大的批量大小会影响收敛速度，同时还需要大量内存来存储激活值。

PipeDream的另一种流水线并行实现保留多个过时参数副本，以隐藏流水线气泡，而不会显著增加批量大小，这使得其内存效率较低。此外，该实现与标准的深度学习训练并不等同，并对训练收敛有影响。相比之下，ZeRO在不产生流水线并行的功能、性能和收敛相关限制的情况下，获得了与PP相同或更好的内存效率。

2.2 非并行方法减少内存

2.2.1 减少激活内存

多项研究致力于通过压缩、激活检查点或实时分析(live analysis)来降低激活的内存占用。这些方法是互补的，可以与ZeRO协同工作。实际上，ZeRO-R中的激活内存减少与激活检查点一起工作。

2.2.2 CPU卸载

一些研究利用当前计算节点的异构特性，通过算法设计或虚拟化内存将模型状态卸载到CPU内存中，使得训练时间的多达50%可能用于GPU-CPU-GPU的传输。与这些方法不同，ZeRO显著减少内存消耗，而不是将模型状态存储到由于PCI-E带宽限制而严重受限的CPU内存中。在罕见情况下，ZeRO-R可能会将激活检查点卸载到CPU中，以提升性能。

2.2.3 内存高效优化器

另一些研究关注通过维护粗粒度的模型参数和梯度统计来减少自适应优化方法的内存消耗，这可能对模型收敛性保证产生影响。ZeRO与这些努力是正交的，其优化不会改变模型优化方法或影响模型收敛，而是有效减少每个设备的优化器状态和梯度的内存占用。自适应优化方法对于大模型的有效训练，达到SOTA性能和准确性至关重要。与随机梯度下降相比，由于需要维护每个模型参数和梯度的细粒度一阶和二阶统计信息，这会产生显著的内存占用。ZeRO能够将这些优化器的内存占用降低几个数量级，使得这些复杂的优化方法在内存资源有限的硬件上变得可行。这也使得开发和使用更复杂、内存需求更高的优化器成为可能，可能会带来更好的收敛性。

3 内存都去哪儿了？

我们看看当前训练系统的内存消耗。例如，一个有15亿参数的GPT-2模型在16位精度下需要3GB的内存来存储其权重(或参数)，然而，它无法在32GB内存的单个GPU上使用Tensorflow或PyTorch进行训练。这让人不禁要问，所有内存都去哪儿了。在模型训练过程中，大部分内存被模型状态(model states)所占用，即由优化器状态、梯度和参数组成的张量。除了这些模型状态，剩余的内存则被激活、临时缓冲区和我们称之为剩余状态的碎片内存所消耗。

3.1 模型状态：优化器状态、梯度和参数

在训练过程中，设备内存的大部分被模型状态所占用。以Adam为例。Adam需要存储两个优化器状态：1)时间平均动量，2)梯度的方差，以计算更新。因此，使用Adam训练模型时，必须有足够的内存来存储梯度的动量和方差的副本。此外，还需要有足够的内存来存储梯度和权重本身。在这三种类型的参数相关张量中，优化器状态通常占用最多的内存，尤其是在进行混合精度训练时。

混合精度训练 在当前一代NVIDIA GPU上训练大型模型的最先进方法是采用混合精度(fp16/32)训练，其中参数和激活以fp16存储，从而能够使用这些GPU高吞吐量的张量核心单元。在混合精度训练中，前向和反向传播都使用fp16权重和激活进行。然而，为了在反向传播结束时有效计算和应用更新，混合精度优化器还需要保持参数的fp32副本以及所有其他优化器状态的fp32副本。

以Adam为具体例子，使用Adam训练一个具有$\Psi$个参数的模型的混合精度版本，需要足够的内存来存储参数和梯度的fp16副本，需要的内存为$2\Psi$和$2\Psi$字节。此外，还需要存储优化器状态：参数、动量和方差的fp32副本，需要的内存为$4\Psi$、$4\Psi$和$4\Psi$字节。我们用$K$表示优化器状态的内存倍增因子，即存储它们所需的额外内存为$K\Psi$字节。在混合精度Adam中，K = 12。总之，这导致内存需求为$2\Psi +2\Psi +K\Psi =16\Psi$​字节。对于一个像GPT-2这样的模型，具有15亿个参数，这意味着至少需要24GB的内存，这显著高于仅存储fp16参数所需的微薄的3GB内存。

3.2 剩余内存消耗

在训练过程中，激活值会占用大量内存。以具体例子来说，参数为15亿的GPT-2模型在序列长度为1K和批量大小为32时，需要约60 GB的内存。激活检查点(或激活重计算)是一种常见的方法，可以在重计算开销为33%的情况下，将激活内存降低约为总激活的平方根。这将使该模型的激活内存消耗减少到约8 GB。

尽管显著减少了内存，但对于更大的模型，即使使用激活检查点，激活内存仍然可能变得相当庞大。例如，一个具有1000亿参数的类GPT的模型，在批量大小为32时仍需要约60 GB的内存。

用于存储中间结果的临时缓冲区在较大模型中也会消耗相当可观的内存。诸如梯度all-reduce或梯度norm计算等操作通常会将所有梯度融合到一个单一的扁平化缓冲区中，以提高吞吐量。例如，跨设备的全规约带宽会随着大消息大小的增加而提高。虽然梯度本身通常存储为fp16张量，但根据操作不同，融合的缓冲区可能是fp32张量。当模型规模较大时，这些临时缓冲区的大小是相当不容小觑的。例如，对于一个具有15亿参数的模型，扁平化的fp32缓冲区将需要6 GB的内存。

内存碎片化 到目前为止，我们讨论了训练过程中实际的内存消耗。除此之外，即使有大量可用内存，也可能会出现可用内存不足的情况。这种情况可能是由于内存碎片化造成的。如果没有足够的连续内存来满足内存请求，即使总可用内存大于请求的内存，该请求也会失败。在训练非常大的模型时，我们观察到显著的内存碎片化，导致出现内存不足的问题，某些极端情况下，仍然有超过30%的内存可用。

4 ZeRO：洞察与概述

ZeRO有两组优化措施：1) ZeRO-DP旨在减少模型状态的内存占用，2) ZeRO-R则针对减少剩余内存消耗。我们将介绍这些优化的概述和背后的洞察，这使得ZeRO能够在保持高效的同时减少内存占用。这里效率很关键：如果没有这个约束，像将所有参数状态移动到CPU内存，或任意增加MP度数等简单解决方案都可以减少内存占用。

洞察：独一无二的观点

4.1 ZeRO-DP

ZeRO驱动的DP基于三个关键洞察：

1. DP的扩展效率优于MP，因为MP降低了计算的粒度，同时增加了通信开销。在某个临界点之后，较低的计算粒度会降低每个GPU的效率，而增加的通信开销则会妨碍跨GPU的可扩展性，尤其是在跨节点边界时。相反，DP具有更高的计算粒度和更低的通信量，从而实现更高的效率。
2. DP在内存使用上效率较低，因为模型状态在所有数据并行进程中冗余存储。相反，MP通过对模型状态进行分区来获得内存效率。
3. DP和MP在整个训练过程中都保持所有所需的模型状态，但并不是所有状态在所有时间都需要。例如，层参数仅在该层的前向传播和反向传播期间需要。

基于这些思考，ZeRO-DP在保持DP的训练效率的同时，实现了MP的内存效率。ZeRO-DP通过分区模型状态而不是复制它们，并使用动态通信调度，利用模型状态的内在时间特性，同时最小化通信量。

通过这样做，ZeRO-DP在增加DP度数的同时线性减少每个设备的模型内存占用，同时保持通信量接近默认DP的水平，从而保持效率。

4.2 ZeRO-R

4.2.1 减少激活内存

两个关键洞察是：

1. MP对模型状态进行分区，但通常需要复制激活内存。例如，如果我们将线性层的参数垂直拆分在两个GPU上并行计算，则每个GPU都需要完整的激活来计算其分区。
2. 对于像GPT-2或更大规模的模型，算术强度(每次迭代的计算量与每次迭代的激活检查点量之比)非常大，并且随着隐藏层维度的增加而线性增加，使得即使带宽较低，也能够隐藏激活检查点的数据移动开销。

ZeRO通过在GPU之间分区激活检查点，减少了MP中的内存冗余，并使用allgather根据需要重构它们。激活内存的占用量根据MP度数成比例减少。对于非常大的模型，ZeRO甚至可以选择将激活部分移动到CPU内存中，同时由于这些模型的算术强度大，仍然能够实现良好的效率。

4.2.2 管理临时缓冲区

ZeRO使用恒定大小的缓冲区，以避免随着模型规模的增加而导致临时缓冲区膨胀，同时确保缓冲区的大小足够以保持效率。

4.2.3 管理内存碎片

内存碎片是短期内存对象与长期内存对象之间交错的结果。在前向传播期间，激活检查点是长期存在的，而重新计算的激活则是短期存在的。类似地，在反向计算中，激活梯度是短期存在的，而参数梯度是长期存在的。基于这一洞察，ZeRO通过将激活检查点和梯度移动到预分配的连续内存缓冲区，进行即时内存碎片整理。这不仅增加了内存的可用性，还通过减少内存分配器寻找空闲连续内存所需的时间，提高了效率。

5. 深入ZeRO-DP

虽然现有的数据并行(DP)方法在每个设备上复制模型状态，并引入显著的内存开销，但 ZeRO-DP 通过在数据并行进程之间对优化器状态、梯度和参数进行分区，消除了这种内存冗余。图1定量并可视化了使用和不使用 ZeRO-DP 的内存需求。该图展示了在分区后(1)优化器状态、(2)梯度和(3)参数冗余的累计内存占用。将它们称为 ZeRO-DP 的三个优化阶段：$P_{os}$、$P_g$ 和 $P_p$，以下将详细说明。

5.1 Pos：优化器状态分区

对于 DP 程度 $N_d$，将优化器状态分为 $N_d$ 个相等的分区，使得第i个数据并行进程只更新与第i个分区对应的优化器状态。因此，每个数据并行进程只需存储和更新总优化器状态的 $\frac{1}{N_d}$部分，同时只更新 $\frac{1}{N_d}$的参数。在每个训练步骤结束时，在所有数据并行进程之间执行一次 all-gather，以获取所有数据并行进程的完全更新参数。

内存节省：如图1所示，经过优化状态分区后的内存消耗从 $4\Psi +K\Psi$ 降低到 $4\Psi +\frac{K\Psi }{N_d}$。以图1中的具体例子为例，使用 64 路数据并行的 $P_{os}$ 处理($N_d=64$)，一个 75 亿参数的模型需要 31.4GB 的内存，而使用标准 DP 则需要 120 GB。此外，当 $N_d$ 较大时，模型状态的内存需求从 $4\Psi +12\Psi =16\Psi$字节降至 $4\Psi +\frac{12\Psi }{N_d}\approx 4\Psi$​字节，得到 4 倍的减少。

5.2 Pg：梯度分区

由于每个数据并行进程只更新其对应的参数分区，因此它只需要对应参数聚合(reduce)的梯度。因此，在反向传播过程中，当每层的梯度可用时，仅在负责更新相应参数的数据并行进程中进行reduce。之后不再需要这些梯度，可以释放其内存。这将需要持有梯度的内存占用从 $2\Psi$ 字节减少到 $\frac{2\Psi }{N_d}$。

实际上，这是一种 Reduce-Scatter 操作，其中对应于不同参数的梯度被reduce到不同的进程。为了在实践中提高效率，使用了一种分桶策略，将所有与特定分区对应的梯度归入一个桶，并对整个桶进行一次性reduce。执行在分区边界的reduce，而不是all-reduce，以减少内存占用并重叠计算和通信。

内存节省：通过消除梯度和优化器状态冗余，将内存占用进一步降低到 $2\Psi +\frac{14\Psi }{N_d}\approx 2\Psi$。如图1中的例子所示，使用 $P_{os+g}$ 处理的 75 亿参数模型，在 64 路数据并行情况下仅需 16.6 GB 的内存，而使用标准 DP 则需要 120 GB。当 $N_d$较大时，模型状态的内存需求从 $2\Psi +14\Psi =16\Psi$字节降低到 $2\Psi +\frac{14\Psi }{N_d}\approx 2\Psi$​字节，得到8 倍的减少。

5.3 Pp：参数分区

与优化器状态和梯度一样，每个进程仅存储其对应分区的(模型)参数。当需要进行前向和反向传播时，超出其分区的参数将通过广播从适当的数据并行进程接收。尽管乍一看这似乎会产生显著的通信开销，但这种方法仅将基线数据并行系统的总通信量增加到 1.5 倍，同时使内存减少与数据并行度 $N_d$成正比。

内存节省：通过参数分区，将一个 $\Psi$参数模型的内存消耗从 $16\Psi$降低到 $\frac{16\Psi }{N_d}$。如图1所示，使用 $P_{os+p+g}$ 和 64路数据并行$N_d=64$​的情况下，7.5 亿参数模型需要 1.9 GB 的模型状态内存，而使用标准数据并行则需要 120 GB。ZeRO 可以使数据并行适应任意大小的模型，只要有足够的设备共享模型状态。

5.4 对模型大小的影响

对 $P_{os}$、$P_{os+g}$ 和 $P_{os+g+p}$三个阶段的分区分别将每个数据并行进程在模型状态的内存消耗减少至最多 4 倍、8 倍和 $N_d$ 倍。表1分析了在 ZeRO-DP 优化的 3 个阶段下，几个示例模型在不同数据并行度下的模型状态内存消耗。如果没有 ZeRO，内存消耗等于表中的第一行，无论数据并行度如何。

当 $N_d=64$时，分别使用$P_{os}$、$P_{os+g}$ 和 $P_{os+g+p}$，ZeRO 可以训练最多 7.5B、14B 和 128B 参数的模型；当 $N_d=1024$时，启用所有 ZeRO 优化$P_{os+g+p}$​能够训练具有 1 万亿参数的模型！或任何规模的模型！而没有 ZeRO，最大的数据并行模型只能运行少于 15 亿参数的模型。

6. 深入了解 ZeRO-R

6.1 Pa：分区激活检查点

如 4.2 节所述，模型并行设计上需要复制激活，从而导致在模型并行 GPU 之间激活的冗余副本。ZeRO 通过对激活进行分区来消除这种冗余，并且仅在计算过程中使用激活之前，以复制形式逐层生成激活，具体来说，一旦计算完成模型某一层的前向传播，输入激活便会在所有模型并行进程中进行分区，直到在反向传播中再次需要此激活为止。在这一点上，ZeRO 使用all-gather操作重新生成激活的复制副本。将这种优化称为 $P_a$。它与激活检查点技术一起工作，只存储分区的激活检查点，而不是复制的副本。此外，对于非常大的模型和极为有限的设备内存，这些分区的激活检查点也可以卸载到 CPU，以将激活的内存开销减少到几乎为零，尽管会产生额外的通信成本，将其称为 $P_{a+cpu}$​。

内存节省：通过分区激活检查点，ZeRO 将激活的内存占用降低至与 MP 程度成正比的因子。考虑训练一个 100B 模型，批量大小为 32，序列长度为 1024，MP 度为 16。如果对每个transformer层进行一个激活检查点，那么每个 GPU 所需的内存将约为 33 GB，仅用于存储激活检查点。但使用 ZeRO 的 $P_a$​ 后，这个内存需求可以降低至每个 GPU 大约 2 GB。此外，这 2 GB 内存可以卸载到 CPU，进一步将激活的内存占用减少到几乎为零。

6.2 CB：恒定大小缓冲区

ZeRO 精心选择临时数据缓冲区的大小，以平衡内存和计算效率。在训练过程中，一些操作的计算效率可能高度依赖于输入大小，较大的输入能够实现更高的效率。例如，大规模的all-reduce操作比小规模操作实现更高的带宽。因此，为了获得更好的效率，高性能库如 NVIDIA Apex 或 Megatron 会在应用这些操作之前，将所有参数融合到一个单一的缓冲区中。然而，融合缓冲区的内存开销与模型大小成正比，可能会造成限制。例如，对于一个 3B 参数的模型，32 位的融合缓冲区将需要 12 GB 的内存。为了解决这个问题，在模型过大时简单使用高性能、高效的恒定大小的融合缓冲区。这样，缓冲区大小不依赖于模型大小，并且通过保持缓冲区大小足够大，仍然可以实现良好的效率。

6.3 MD：内存碎片整理

在模型训练中，由于激活检查点和梯度计算而导致的内存碎片化现象。在进行激活检查点的前向传播过程中，仅存储选定的激活用于反向传播，而大部分激活被丢弃，因为它们可以在反向传播期间重新计算。这导致了短期内存(被丢弃的激活)和长期内存(检查点激活)之间的交错，从而造成内存碎片。同样，在反向传播过程中，参数梯度是长期存在的，而激活梯度和计算参数梯度所需的其他缓冲区则是短期存在的。这种短期内存和长期内存的交错再次导致内存碎片。

有限的内存碎片通常不是问题，尤其是在有充足内存的情况下，但在有限内存下进行大模型训练时，内存碎片会导致两个问题：1)即使有足够的可用内存，也会因缺乏连续内存而出现内存溢出；2)由于内存分配器花费大量时间寻找满足内存请求的连续内存块，导致效率低下。

ZeRO 通过为激活检查点和梯度预分配连续内存块，并在生成时将它们复制到预分配的内存中，实时进行内存碎片整理。$M_D$不仅使 ZeRO 能够以更大的批量大小训练更大的模型，还在有限内存下提高了训练效率。

7. ZeRO-DP 的通信分析

由于 ZeRO 通过消除内存冗余来提升模型大小，ZeRO 驱动的 DP 方法的通信量与基线 DP 方法相比如何？答案分为两部分：1) 使用 $P_{os}$ 和 $P_g$ 时，ZeRO-DP 不会产生额外的通信，同时实现最多 8 倍的内存减少；2)在使用 $P_p$ 而不是$P_{os}$ 和 $P_g$ 时，ZeRO-DP 的通信量最多为 1.5 倍，同时进一步减少内存占用 $N_d$ 倍。

7.1 数据并行通信量

在数据并行训练过程中，所有数据并行进程的梯度在反向传播结束时进行平均，然后再计算下一步的更新。平均操作使用all-reduce集体通信。对于大模型，all-reduce通信完全受限于通信带宽，因此将分析限制在发送和接收每个数据并行进程的总通信量上。最先进的all-reduce实现采用两步法，第一步是reduce-scatter操作，它在不同的进程上减少数据的不同部分。下一步是all-reduce操作，每个进程在所有进程上收集减少后的数据。这两个步骤的结果就是all-reduce。reduce-scatter和all-reduce都是通过管道化方法实现的，这导致每个步骤的数据移动总量为 $\Psi$ 元素(对于一个包含 $\Psi$ 元素的数据)。因此，标准 DP 在每个训练步骤中会产生 $2\Psi$ 的数据移动。

7.2 ZeRO-DP通信量

7.2.1 使用Pos+g的通信量

通过梯度分区，每个进程仅存储更新其对应参数分区所需的梯度部分。因此，ZeRO仅需对梯度执行scatter-reduce操作，而不需要all-reduce，这产生的通信量为$\Psi$。在每个进程更新其负责的参数分区后，会执行all-gather操作，以从所有数据并行进程收集所有更新后的参数。这也产生了$\Psi$的通信量。因此，每个训练步骤的总通信量为$\Psi +\Psi =2\Psi$，正好与基线DP相同。

7.2.2 使用Pos+g+p的通信量

在参数分区后，每个数据并行进程仅存储其更新的参数。因此，在前向传播期间，它需要接收所有其他分区的参数。然而，这可以进行管道化以避免内存开销。在计算与特定分区对应的模型部分的前向传播之前，负责该分区的数据并行进程可以将权重广播到所有数据并行进程。一旦该分区的前向传播完成，参数可以被丢弃。因此，总的通信量为$\frac{\Psi \times N_d}{N_d}=\Psi$。

换句话说，通过将其分布在整个前向传播的参数all-gather的过程下重新安排，并在使用后丢弃参数。这个all-gather在反向传播中还需要按照反向顺序再次进行。因此，总通信量是由这些all-gather操作产生的通信量加上梯度的scatter-reduce所产生的通信量之和。因此，总量为$3\Psi$​​，与基线相比为1.5倍。梯度和参数分区利用了这样一个见解——并非所有状态的梯度和参数都需要随时存在——通过明智地通信这些状态来优化内存。

8. ZeRO-R的通信分析

将分区激活检查点($P_a$)在ZeRO-R中的通信量与基线模型并行进行比较，发现$P_a$引入的通信量增加通常小于基线MP的十分之一。分区激活检查点的通信量权衡取决于模型大小、检查点策略和MP策略。

在使用激活检查点的Megatron-LM中，每个transformer块在前向传播中执行两次大小为batch × seq_length × hidden_dim维度的all-reduce操作，在前向重新计算中又执行两次all-reduce，在反向传播中再执行两次。每个块的总通信量为12 × seq_length × hidden_dim，因为all-reduce的通信量为2 × message_size。

当ZeRO-R对激活检查点进行分区时，需要在每个激活检查点的反向传播前进行一次额外的all-reduce操作。一般来说，对每个transformer块检查输入激活，因此每个transformer块需要一次all-reduce。因此，$P_a$的通信开销为seq_length × hidden_dim，因为all-reduce的通信量为message_size。因此，$P_a$的总通信开销小于原始模型并行通信量的10%。

当MP与DP结合使用时，$P_a$可以在增加10%的模型并行通信量的情况下，将数据并行通信量减少一个数量级，并在数据并行通信成为性能瓶颈时显著提升效率。$P_a$通过模型并行度减少激活内存消耗，从而允许批量大小成比例增加。对于大型模型，MP可以达到16，这使得批量大小增加最多可达16倍。

数据并行训练的通信量与批量大小成反比。因此，由于$P_a$引起的批量大小增加一个数量级可能导致数据并行通信量减少一个数量级。

最后，如果应用$P_{a+cpu}$，分区激活检查点将被卸载到CPU，激活内存需求几乎降低到零，但相对于$P_a$，往返于CPU内存的额外数据移动增加了2倍。在极端情况下，如果由于小批量大小，即使使用$P_a$，DP通信量仍是主要瓶颈，那么$P_{a+cpu}$可以通过增加批量大小来提高效率，只要CPU数据传输开销少于DP通信量开销，这通常对小批量大小成立。

9. 迈向万亿参数

目前发布的最大模型在100亿参数的范围内，训练起来已经相当具有挑战性。达到万亿参数——大三个数量级的规模，必然会发生。ZeRO并不声称能够知道或解决所有问题，但ZeRO解决了从系统角度来看最根本的挑战之一：在当前硬件上适配如此规模的模型，同时确保良好的系统可扩展性。

从SOTA的跨越 目前，SOTA框架Megatron能够以可接受的吞吐量训练的最大模型为160亿至200亿参数的模型，运行在DGX-2系统中。而通过在多个DGX节点之间实现模型并行，进一步扩展会由于节点间带宽的限制而导致显著的效率下降。

ZeRO大幅提升了可高效运行的模型规模。它使得当前一代硬件可以运行显著更大的模型，而不需要细粒度的模型并行来跨越节点边界。如表1所示，启用所有优化$P_{os+g+p}$​的ZeRO能够仅使用数据并行在1024个GPU上容纳超过1万亿参数的模型。同时，在结合模型并行时(如表2所示)，ZeRO能够在1024个GPU上以16倍模型并行和64倍数据并行跨节点，容纳超过1万亿参数的模型。

计算能力的差距 然而，在可接受的时间范围内端到端地训练一个万亿参数的模型，仍然可能需要大量计算能力，而这在今天的AI集群中是缺乏的。

为了理解资源需求，与Bert-Large进行了简要比较。Bert-Large可以在1024个GPU的DGX-2H集群上以67分钟的时间训练完成。一个一万亿参数的模型，针对每个数据样本，其计算量可能是Bert-Large模型的3000倍。即使我们假设序列长度和训练模型所需的样本总数相同，训练一个一万亿参数的模型也需140天，假设相同硬件和类似计算效率。实际上，数据样本和序列长度很可能随模型规模的增加而增加，训练可能需要超过一年。为了在合理的时间内训练一万亿参数的模型，需要一个exa-flop级的系统。

总结

⭐ 主要介绍了DeepSpeed框架用到的零冗余优化器(ZeRO)技术，使得训练万亿级别参数的大模型成为了可能。通过ZeRO-DP优化模型状态，通过ZeRO-R优化剩余状态。

引用

---

1. Training Deep Nets with Sublinear Memory Cost ↩︎