摘要

       大型深度学习模型能提供显著的准确率提升，但训练亿至万亿的参数是个挑战。现有解决方案如数据和模型并行存在基本限制，难以将这些模型装入有限的设备内存中，同时还要获得计算、通信和开发效率。我们开发了一种新颖的解决方案，即零冗余优化器（ZeRO），以优化内存，大大提高训练速度，同时增加了可高效训练的模型大小。ZeRO消除了数据和模型并行训练中的内存冗余，同时保持了低通信量和高计算粒度，使我们能够按设备数量的比例扩展模型大小，并保持高效率。我们对内存需求和通信量进行了分析，结果显示：利用现今的硬件，ZeRO有潜力扩展到超过1万亿个参数。

       我们实现并评估了ZeRO：它在400个GPU上以超线性速度加速训练了超过1000亿参数的大型模型，实现了15 Petaflops的吞吐量。这代表了模型大小的8倍增加和性能的10倍提升，相比于当前最先进的技术。在可用性方面，ZeRO能够训练多达130亿参数的大型模型（例如，比Megatron GPT的83亿和T5的110亿更大）而无需模型并行，这对科学家来说更难应用。最后但同样重要的是，研究人员利用ZeRO的系统突破创造了世界上最大的语言模型（170亿参数），并创造了准确率的新纪录。

1 引言

       随着模型的不断壮大，深度学习（DL）模型正在取得显著的准确性提升。在自然语言处理（NLP）领域，变换器（Transformers）的出现为大型模型如Bert-large（0.3B）、GPT-2（1.5B）、Megatron-LM（8.3B）、T5（11B）的发展铺平了道路。为了实现从数百亿到数万亿参数规模的模型大小增长，我们在训练这些模型时遇到了挑战 - 显然，它们无法全部装入单个设备的内存中，比如 GPU 或 TPU，而简单地增加更多设备并不能帮助扩展训练。

       基本的数据并行（DP）无法减少每个设备的内存占用，并且在当前具有32GB内存的GPU上，对于超过1.4B参数的模型会出现内存不足的问题。其它现有的解决方案如管道并行（PP）、模型并行（MP）、CPU-卸载等，在功能、可用性以及内存和计算/通信效率之间做出了权衡，但所有这些因素对于实现速度和规模的训练都至关重要。在训练大型模型的不同现有解决方案中，MP可能是最有前途的。目前文献中最大的模型，11B T5模型和8.3B Megatron-LM模型，都是由模型并行性支持，分别在Mesh-Tensorflow和Megatron-LM中实现。但是，MP在这些模型大小之外无法进一步扩展。MP通过垂直分割模型，在每层之间分割计算和参数，并在多个设备上进行分配，这需要每层之间进行大量的通信。因此，它们在单个节点内效果很好，其中的GPU间通信带宽很高，但是效率在超过一个节点之后会快速下降。

       我们在两个DGX-2节点上使用Megatron-LM测试了一个40B参数的模型，并观察到每个V100 GPU大约有5 T flops（低于硬件峰值的5%）。那么，我们如何克服现有解决方案的局限性，更有效地训练大型模型呢？为了回答这个问题，我们首先分析了现有系统在模型训练中的内存消耗全谱，并将其分类为两部分：1) 对于大型模型，大部分内存被模型状态占据，这包括优化器状态（例如Adam中的动量和方差）、梯度和参数；2) 其余的内存被激活、临时缓冲和不可用的碎片化内存消耗，我们统称为残留状态。我们开发了ZeRO——零冗余优化器——以在保持高计算和通信效率的同时优化内存效率。由于这两部分面临不同的挑战，我们分别为它们开发和讨论了解决方案。

       优化模型状态内存模型状态通常在训练过程中消耗最多的内存，但现有的方法如DP和MP并未提供令人满意的解决方案。DP具有良好的计算/通信效率但内存效率低，而MP可能具有较差的计算/通信效率。更具体地说，DP通过在所有数据并行进程中复制整个模型状态，导致了冗余的内存消耗；而MP通过分割这些状态来获得高内存效率，但通常会导致过细的计算和昂贵的通信，这对于缩放效率不利。此外，所有这些方法在整个训练过程中都静态地维护所有必需的模型状态，尽管在训练过程中并不是所有时间都需要所有模型状态。基于这些观察，我们开发了ZeRO-DP，ZeRO支持的数据并行，它实现了DP的计算/通信效率，同时实现了MP的内存效率。ZeRO-DP通过分割模型状态而不是复制它们来消除数据并行进程中的内存状态冗余，并通过在训练期间使用动态通信计划来保留DP的计算粒度和通信量。

       ZeRO-DP具有三个主要的优化阶段（如图1所示），分别对应于优化器状态、梯度和参数的分区。当按累加方式启用时：

1）优化器状态分区(Pos)：内存减少4倍，通信量与DP相同；

2）添加梯度分区(Pos+g)：内存减少8倍，通信量与DP相同；

3）添加参数分区(Pos+g+p)：内存减少与DP程度Nd线性相关。例如，跨64个GPU分割（Nd = 64）将实现64倍的内存减少。通信量会有适度的50%增加。

       ZeRO-DP消除了内存冗余，并使群集的全部聚合内存容量可用。启用所有三个阶段后，ZeRO可以仅用1024个NVIDIA GPU训练一个万亿参数的模型。具有Adam优化器的万亿参数模型在16位精度下需要大约16TB的内存来保存优化器状态、梯度和参数。16TB除以1024等于16GB，这对于GPU来说是合理的（例如，具有32GB的设备内存）。

       优化残余状态内存(Optimizing Residual State Memory) 在ZeRO-DP提高模型状态的内存效率后，由激活、临时缓冲区和不可用的内存片段消耗的剩余内存可能会成为次要的内存瓶颈。我们开发了ZeRO-R来分别优化这三个因素消耗的残余内存。

       1）对于激活（为了执行反向传递而存储的前向传递），我们注意到检查点技术（checkpointing）有所帮助，但对于大型模型来说不够。因此，ZeRO-R通过识别和删除现有MP方法中的激活复制来优化激活内存，并通过激活分区实现。它还在适当时将激活卸载到CPU。

       2）ZeRO-R为临时缓冲区定义了合适的大小，以实现内存和计算效率之间的平衡。

       3）我们在训练过程中观察到了由于不同张量的生命周期变化而产生的内存碎片。由于碎片化导致连续内存的缺乏可能会导致内存分配失败，即使有足够的空闲内存。ZeRO-R根据不同张量的生命周期主动管理内存，防止内存碎片化。

       ZeRO-DP和ZeRO-R结合在一起构成了我们统称为ZeRO的DL训练内存优化的强大系统。

       ZeRO和MP： 由于ZeRO消除了DP中的内存低效问题，很自然地会问：我们还需要MP吗，什么时候需要？ZeRO如何与MP一起工作？有了ZeRO，MP对于仅适应大型模型而言变得不那么有吸引力。ZeRO-DP至少与MP一样有效，有时在MP无法均匀分割模型时更有效。它还具有相当或更好的缩放效率。此外，数据并行非常容易使用，因此可以广泛应用于不同的工作负载，而MP方法通常需要模型开发人员修改其模型，系统开发人员设计分布式运算符，而像Megatron-LM这样的现有工作仅支持有限的运算符和模型。

       也就是说，仍然有一些情况我们希望利用MP：i) 当与ZeRO-R一起使用时，MP可以减少非常大型模型的激活内存占用。ii) 对于激活内存不是问题的较小模型，当使用DP时累积的批处理大小过大以至于无法获得良好的收敛性时，MP也可能有益。在这些情况下，可以将ZeRO与MP结合，以适应具有可接受的累积批处理大小的模型。

       我们展示了ZeRO可以与MP结合，每个设备上的最大理论内存减少为Nd × Nm倍，其中DP的度为Nd，MP的度为Nm。这可能使我们能够在1024个GPU上适应一个万亿参数模型，每个DGX2节点内具有16路模型并行性，跨节点具有64路数据并行性，并使用适度的批处理大小有效运行！

       实现与评估(Implementation & Evaluation) ZeRO中的完整优化集可以让我们在今天的高端硬件集群上运行具有万亿参数的模型（例如，带有1K V100 GPUs），但是硬件计算能力仍然非常有限，训练时间可能不切实际地长（>1年）。因此，我们此次实现的重点是有效支持具有10倍参数（∼100B参数）的模型，而这些参数相对于最先进的技术（SOTA）仍然在当前硬件的计算能力范围内。我们实现并评估了ZeRO的一个子集，称为ZeRO-100B——ZeRO-DP的Pos+g加上ZeRO-R——以实现这个目标。结果显示：

       模型大小(Model Size) 与MP结合，ZeRO-100B可以有效运行170B参数模型，而像仅使用Megatron的现有系统不能有效地超过40B参数的规模，如图2所示。与SOTA相比，这是模型大小超过8倍的增加。

       速度(Speed) 改进的内存效率提高了吞吐量和训练速度。如图2所示，ZeRO在400个Nvidia V100 GPU集群上运行100B参数模型，每个GPU超过38 TFlops，总性能超过15 Petaflops。对于相同模型大小，这比SOTA的训练速度提高了10倍以上。

       可扩展性(Scalability) 我们在64-400 GPU的范围内观察到超线性加速，其中性能在我们加倍GPU数量时增加了一倍以上。这是ZeRO-DP的一个属性，它随着我们增加DP度而减少模型状态的内存占用，从而允许我们每个GPU适应更大的批处理大小，从而获得更好的性能。我们预期当我们将GPU数量增加到400以上时，这种行为将继续。

       大型模型训练的民主化(Democratization of Large Model Training) ZeRO-100B使数据科学家能够训练多达130亿参数的模型，而无需进行任何MP或PP的模型重构，其中13B的参数比文献中最大的模型（具有11B参数的T5）还要多。因此，数据科学家可以自由地尝试大型模型，而无需担心并行性。相比之下，现有系统（例如，PyTorch Distributed Data Parallel）在14亿参数模型中会耗尽内存。

       新SOTA模型(New SOTA Model) ZeRO支持具有170亿参数和创纪录准确性的最大语言模型，即Turing-NLG。

       我们将ZeRO作为我们名为DeepSpeed2的开源DL训练优化库的一部分共享。我们计划在2020年5月底之前发布本文中描述的所有实现，并进一步扩展以通过启用ZeRO-DP阶段3参数分区（Pos+g+p）来支持1万亿参数。我们计划使DL社区能够完全访问ZeRO，以催化大规模大型模型训练的进化和民主化。

2 相关工作

2.1 数据、模型和管道并行性

    并行化是大规模训练大型模型的关键策略。对于一个适合在设备内存中训练的模型，数据并行性（DP）被用来将训练扩展到多个设备。在DP中，模型参数在每个设备上都被复制。在每个步骤中，一个小批量被均匀地分布在所有的数据并行过程中，使得每个过程在不同的数据样本子集上执行前向和后向传播，并使用跨过程的平均梯度在本地更新模型。

    当模型不适合在设备内存中时，模型并行性（MP）和管道并行性（PP）分别以垂直和水平的方式将模型分割在过程中。第1节讨论了ZeRO如何与DP和MP相关。我们现在讨论PP以及它如何降低内存消耗。

    PP将模型在层之间水平分割，每个分区在不同的设备上运行，并使用微批处理来隐藏管道泡沫。由于水平分割和微批处理，模型功能如绑定权重和批量归一化很难实现。流行的PP实现，如G-pipe，划分了模型参数和总激活，但需要与管道分区数量成比例的批大小以隐藏管道泡沫。大批量大小可能会影响收敛速率，同时还需要大量的内存来存储激活。PipeDream中的不同PP实现保留了多份过期参数的副本以隐藏管道泡沫，而不会显著增加批量大小，从而使其内存效率较低。此外，该实现不等同于标准的DL训练，并对训练收敛具有影响。相比之下，ZeRO在不产生PP的功能、性能和收敛相关限制的情况下，获得了与PP相同或更好的内存效率。

2.2 基于非并行方法来减少内存

    除了MP和PP，还有多条旨在减少DL训练内存开销的工作线。

2.2.1 减少激活内存

    通过压缩、激活检查点或实时分析，有多种努力专注于减少激活的内存占用。这些努力是互补的，并且可以与ZeRO一起工作。实际上，ZeRO-R中的激活内存减少与激活检查点并行。

2.2.2 CPU卸载

    利用了当今计算节点的异构性，通过算法设计或虚拟内存将模型状态卸载到CPU内存。多达50%的训练时间可以花费在GPU-CPU-GPU传输上。ZeRO的不同之处在于，它在不将模型状态存储到由于PCI-E严重受限的CPU内存的情况下显著减少了内存消耗。在极少数情况下，ZeRO-R可能仅卸载非常大型模型的激活检查点以提高性能。

2.2.3 内存高效优化器

    着重于通过保持模型参数和梯度的粗粒度统计来减少自适应优化方法的内存消耗，可能会影响模型收敛保证。ZeRO与这些努力是正交的，它的优化不会改变模型优化方法或影响模型收敛，但会有效减少每个设备上优化器状态和梯度的内存占用。

2.3 训练优化器

    自适应优化方法对于实现大型模型的有效训练和达到最先进的性能和准确性至关重要。与SGD相比，通过保持每个模型参数和梯度的精细的一阶和二阶统计，以显著的内存占用为代价。ZeRO可以将这些优化器的内存占用减少数个数量级，使得这些复杂的优化方法能够在具有适度设备内存的硬件上训练大型模型成为可能。它还使得开发和使用可能具有更好收敛性的更复杂和更占用内存的优化器成为可能。

3 内存都去哪了？

    让我们退一步来检查当前训练系统的内存消耗。例如，具有15亿参数的GPT-2模型需要3GB内存来存储其以16位精度的权重（或参数），但是，它不能在单个拥有32GB内存的GPU上使用Tensorflow或PyTorch进行训练。人们可能会想知道所有的内存都去哪了。在模型训练期间，大部分内存被模型状态所消耗，即，由优化器状态、梯度和参数组成的张量。除了这些模型状态，其余的内存被激活、临时缓冲区和我们称之为残留状态的碎片化内存所消耗。我们将详细查看这两方面的内存消耗。

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

    在训练期间，设备内存的大部分被模型状态所消耗。例如，考虑Adam，它是DL训练中最受欢迎的优化器之一。Adam需要存储两个优化器状态，即时间平均动量和梯度的方差来计算更新。因此，要用Adam训练模型，就必须有足够的内存来保存梯度的动量和方差的副本。另外，还需要有足够的内存来存储梯度和权重本身。在这三种与参数相关的张量类型中，优化器状态通常消耗最多的内存，特别是当应用混合精度训练时。

    混合精度训练(Mixed-Precision Training) 通过混合精度（fp16/32）训练在当前代NVIDIA GPU上训练大型模型是最先进的方法，其中参数和激活存储为fp16，使得可以使用这些GPU上的高吞吐量张量核单元。在混合精度训练期间，使用fp16权重和激活执行前向和后向传播。然而，为了在反向传播结束时有效计算和应用更新，混合精度优化器保留了参数的fp32副本以及所有其他优化器状态的fp32副本。

    以Adam为具体例子。使用Adam进行模型的混合精度训练需要足够的内存来保存参数和梯度的fp16副本，其内存要求分别为2Ψ和2Ψ字节。另外，它需要保存优化器状态：参数、动量和方差的fp32副本，其内存要求分别为4Ψ、4Ψ和4Ψ字节。我们用K表示优化器状态的内存乘数，即，存储它们所需的额外内存是KΨ字节。混合精度Adam的K值为12。总的来说，这导致2Ψ + 2Ψ + KΨ = 16Ψ字节的内存要求。对于如GPT-2这样具有15亿参数的模型，这导致了至少24GB的内存要求，这比仅保存fp16参数所需的微薄的3GB内存要高得多。

3.2 残留内存消耗

    激活(Activations) 在训练期间，可能会占用大量内存。举一个具体的例子，具有15亿参数的GPT-2模型，以1K的序列长度和32的批量大小进行训练，需要大约60GB的内存。激活检查点（或激活重计算）是一种常见的方法，可通过大约减少总激活的平方根来减少激活内存，但代价是33%的重新计算开销。这将使该模型的激活内存消耗减少到大约8GB。

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

    临时缓冲区(Temporary buffers) 用于存储中间结果，对于大型模型来说消耗了不少内存。诸如梯度全归约或梯度范数计算等操作倾向于在应用操作以努力提高吞吐量之前，将所有梯度融合到一个单一的扁平缓冲区中。例如，设备间的全归约带宽会随着大消息大小而提高。虽然梯度通常存储为fp16张量，但根据操作的不同，融合缓冲区可能是一个fp32张量。当模型的大小较大时，这些临时缓冲区大小是不可忽视的。例如，对于具有15亿参数的模型，一个扁平的fp32缓冲区将需要6GB的内存。

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

4 ZeRO: 洞察与概述

    ZeRO有两组优化：i) ZeRO-DP旨在减少模型状态的内存占用，ii) ZeRO-R针对减少残留内存消耗。我们提供了优化的概述和背后的洞察，这使得ZeRO能够在保持效率的同时减少内存占用。请注意，效率是这里的关键：如果没有这个约束，像将所有参数状态移动到CPU内存或随意增加MP程度这样的简单解决方案可以减少内存占用。

4.1 洞察和概述：ZeRO-DP

    ZeRO提供的DP基于三个关键洞察：

    a) 与MP相比，DP具有更好的扩展效率，因为MP减少了计算的粒度，同时也增加了通信开销。超过某个点后，较低的计算粒度会降低每个GPU的效率，而增加的通信开销则会阻碍跨GPU的可扩展性，特别是在跨节点边界时。相反，DP具有更高的计算粒度和更低的通信量，从而实现了更高的效率。
    b) DP在内存效率上不高，因为模型状态在所有数据并行进程中都被冗余存储。相反，MP划分模型状态以获得内存效率。
    c) DP和MP在整个训练过程中保留所有需要的模型状态，但并非所有时间都需要。例如，仅在每层的前向传播和反向传播期间才需要相应层的参数。

    根据这些洞察，ZeRO-DP保留了DP的训练效率，同时实现了MP的内存效率。ZeRO-DP通过划分模型状态而不是复制它们，并使用动态通信计划，该计划利用模型状态的固有时态性同时最小化通信量。通过这样做，ZeRO-DP随着增加的DP程度线性减少了模型的每设备内存占用，同时保持通信量接近于默认的DP，从而保留了效率。

4.2 洞察和概述：ZeRO-R

4.2.1 减少激活内存

    两个关键见解：

    a) MP划分模型状态，但通常需要复制激活内存。例如，如果我们垂直分割线性层的参数，并在两个GPU上并行计算它们，每个GPU都需要整个激活来计算其划分。
    b) 对于GPT-2或更大的模型，算术强度（每次迭代的计算量与每次迭代的激活检查点量之比）非常大（≥ 10K），并且随着隐藏维度的增加而线性增加，使得即使带宽很低，也有可能隐藏激活检查点的数据移动成本。

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

4.2.2 管理临时缓冲区

    ZeRO-R使用常量大小的缓冲区，以避免随着模型大小的增加临时缓冲区的爆炸，同时使它们足够大以保持效率。

4.2.3 管理碎片化内存

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

5 对ZeRO-DP的深入研究

    尽管现有的DP方法在每个设备上复制模型状态并引入了显著的内存开销，但ZeRO-DP通过划分它们 - 优化器状态、梯度和参数 - 来消除这种内存冗余，分布在数据并行过程中。图1量化并可视化了有无ZeRO-DP的内存需求。该图显示了在划分（1）优化器状态，（2）梯度和（3）参数冗余后的内存占用。我们将它们称为ZeRO-DP的三个优化阶段：Pos，Pg和Pp，我们将在下面详细说明。

5.1 Pos : 优化器状态划分

    对于DP度数Nd，我们将优化器状态分组成Nd个相等的分区，使得第i个数据并行过程仅更新与第i个分区对应的优化器状态。因此，每个数据并行过程只需要存储和更新1/Nd的总优化器状态，然后只更新1/Nd的参数。在每个训练步骤结束时，我们在数据并行过程中执行all-gather，以获得所有数据并行过程中的完全更新参数。

    内存节省(Memory Savings)：如图1所示，优化状态划分后的内存消耗从4Ψ + KΨ减少到4Ψ + KΨ/Nd。如图1中所示的具体示例，一个7.5 B参数模型使用64路DP（Nd = 64）的Pos需要31.4GB的内存，而标准DP需要120GB。此外，当Nd很大时，模型状态的内存需求从4Ψ + 12Ψ = 16Ψ字节减少到4Ψ + 12Ψ/Nd ≈ 4Ψ字节，从而实现了4倍的减少。

5.2 Pg: 梯度划分

    由于每个数据并行过程只更新其对应的参数分区，它只需要对应参数的减少梯度。因此，当每层的每个梯度在反向传播期间变得可用时，我们只在负责更新相应参数的数据并行过程中减少它们。减少后，我们不再需要梯度，它们的内存可以被释放。这将保持梯度所需的内存占用从2Ψ字节减少到2Ψ/Nd。

    实际上，这是一个Reduce-Scatter操作，其中不同参数的梯度减少到不同的过程。为了使这在实践中更有效，我们使用一个桶化策略，我们将所有与特定分区对应的梯度放入一个桶中，并一次对整个桶执行减少。这与NVIDIA的AMP优化器如何桶化所有减少梯度计算以重叠通信和计算非常相似。在我们的情况下，我们在分区边界上执行减少而不是all-reduce，以减少内存占用并重叠计算和通信。

    内存节省：通过消除梯度和优化器状态冗余，我们将内存占用进一步减少到2Ψ + 14Ψ/Nd ≈ 2Ψ。如图1中的示例，一个7.5 B参数模型使用64路DP（Nd = 64）的Pos+g仅需要16.6 GB的内存，而标准DP需要120 GB。当Nd很大时，模型状态的内存需求从2Ψ + 14Ψ = 16Ψ字节减少到2Ψ + 14Ψ/Nd ≈ 2Ψ字节，从而实现了8倍的减少。

5.3 Pp: 参数划分

    就像优化器状态和梯度一样，每个过程只存储对应于其分区的参数。当需要其分区外的参数进行前向和反向传播时，它们将通过广播从适当的数据并行过程中接收。虽然这一开始可能看起来会产生很大的通信开销，但我们表明，这种方法只会将基线DP系统的总通信量增加到1.5倍，同时使内存减少与Nd成比例。

    内存节省：通过参数划分，我们将一个Ψ参数模型的内存消耗从16Ψ减少到16Ψ/Nd。如图1中的示例，一个7.5 B参数模型使用64路DP（Nd = 64）的Pos+p+g仅需要1.9 GB的模型状态内存，而标准DP需要120 GB。这具有深远的意义：ZeRO使DP能够适应任意大小的模型，只要有足够数量的设备共享模型状态。

5.4 对模型大小的影响

    划分的三个阶段Pos，Pos+g和Pos+g+p将每个数据并行过程在模型状态上的内存消耗减少了最多4倍、8倍和Nd。表1分析了在3个阶段的ZeRO-DP优化下，不同DP度数的几个示例模型的模型状态内存消耗。没有ZeRO，内存消耗等于表中的第一行，而不考虑DP度数。注意，当Nd = 64时，ZeRO可以使用Pos，Pos+g和Pos+g+p分别训练最多7.5B、14B和128B参数的模型。当Nd = 1024时，启用所有优化（Pos+g+p）的ZeRO可以训练具有1万亿参数的模型！或者可能，具有任意大小的模型！没有ZeRO，DP单独可以运行的最大模型参数少于15亿。

6 ZeRO-R深入剖析

6.1 Pa: 分区激活检查点

       如4.2节所述，MP（模型并行）的设计要求复制激活值，导致模型并行GPU间存在冗余的激活值副本。ZeRO通过对激活值进行分区来消除这种冗余，并且只在计算中使用激活值时才将它们复制到一个激活层，以此为时机。更具体地说，一旦计算出模型的一个层的前向传播，输入激活值就会在所有模型并行进程中分区，直到在反向传播过程中再次需要它们。此时，ZeRO使用一个all-gather操作来重新生成激活值的复制副本。我们将这种优化称为Pa。它与激活检查点一起工作，只存储分区激活检查点而非复制副本。此外，在模型非常大且设备内存非常有限的情况下，这些分区激活检查点也可以被卸载到CPU，以减少激活内存开销，几乎降至零，但通信成本会增加，我们将在7节中讨论这一点。我们将其称为Pa+cpu。

       内存节省(Memory Saving) 通过分区激活检查点，ZeRO按照与MP程度成比例的因子减少了激活占用空间。考虑训练一个100B的模型，如表4所示，其批次大小为32，序列长度为1024，MP度为16。如果我们为每个变压器层检查一个激活值，那么每个GPU只需存储激活检查点就需要约33 GB的内存。但是使用ZeRO中的Pa，它可以减少到每个GPU约2 GB。此外，这2GB可以被卸载到CPU，将激活的内存占用几乎降至零。

6.2 CB: 固定大小缓冲区

       ZeRO仔细选择临时数据缓冲区的大小，以平衡内存和计算效率。在训练过程中，某些操作的计算效率可能高度依赖于输入大小，较大的输入能够实现更高的效率。例如，大型all-reduce操作的带宽远高于小型操作。因此，为了获得更好的效率，诸如NVIDIA Apex或Megatron等高性能库会在应用这些操作前将所有参数融合到一个缓冲区中。然而，融合缓冲区的内存开销与模型大小成正比，可能会成为阻碍。例如，对于一个30亿参数的模型，一个32位的融合缓冲区将需要12 GB的内存。为了解决这个问题，当模型变得太大时，我们简单地使用性能高效的固定大小融合缓冲区。通过这样做，缓冲区大小不依赖于模型大小，而且通过保持缓冲区大小足够大，我们仍然可以实现良好的效率。

6.3 MD： 内存碎片整理

       在模型训练中的内存碎片产生是由于激活检查点和梯度计算的结果。在带有激活检查点的前向传播过程中，只有选定的激活值被存储用于反向传播，而大多数激活值被丢弃，因为它们可以在反向传播过程中再次计算。这产生了短期内存（丢弃的激活值）和长期内存（检查点激活值）的交错，导致内存碎片化。类似地，在反向传播过程中，参数梯度是长期存在的，而激活梯度和计算参数梯度所需的任何其他缓冲区是短期存在的。再次，短期和长期内存的交错导致了内存碎片化。

       当有大量内存可用时，有限的内存碎片通常不是问题，但对于在内存有限的情况下运行的大型模型训练，内存碎片会导致两个问题，i) 即使有足够的可用内存，也会由于缺乏连续内存而出现OOM（内存溢出），ii) 由于内存分配器花费大量时间寻找连续的内存片以满足内存请求，从而导致效率低下。

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

7 ZeRO-DP的通信分析

       随着ZeRO通过消除内存冗余来增加模型大小，自然会问我们是否在交换通信量以获得内存效率。换句话说，与基线DP方法相比，ZeRO支持的DP方法的通信量是多少？答案分为两部分：i) ZeRO-DP在使用Pos和Pg时不会产生额外的通信，同时能够实现多达8倍的内存减少，ii) 当除了Pos和Pg之外还使用Pp时，ZeRO-DP最多会产生1.5倍的通信，同时将内存占用进一步减少Nd倍。我们在这一节中展示了分析。首先，我们简单介绍了标准DP的通信量。

7.1 数据并行通信量

       在数据并行训练期间，所有数据并行进程在计算下一步的更新之前，都会在反向传播结束时对梯度进行平均处理。平均操作是通过一个全归约通信集合来完成的。对于大型模型，全归约通信完全受通信带宽的限制，因此我们将分析限制为每个数据并行进程发送和接收的总通信量。

       最先进的全归约实现采用了两步方法，其中第一步是一个减少-分散操作，它在不同的进程上减少了数据的不同部分。下一步是一个全收集操作，每个进程收集所有进程上的减少数据。这两步的结果是一个全归约。减少-分散和全收集都是通过流水线方法实现的，导致每个操作的总数据移动为Ψ元素（对于具有Ψ元素的数据），因此，标准DP在每个训练步骤中产生2Ψ的数据移动。

7.2 ZeRO-DP通信量

7.2.1 使用Pos+g的通信量

       通过梯度划分，每个进程只存储更新其对应参数划分所需的梯度部分。因此，ZeRO只需要在梯度上执行一个散列-减少操作，产生Ψ的通信量。每个进程更新其负责的参数划分后，执行一个全收集操作以收集所有数据并行进程的所有更新参数。这也产生了Ψ的通信量。因此，每个训练步骤的总通信量为Ψ + Ψ = 2Ψ，与基线DP完全相同。

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

       在参数划分后，每个数据并行进程只存储它更新的参数。因此，在前向传播期间，它需要接收所有其他划分的参数。然而，这可以通过流水线化来避免内存开销。在计算与特定划分相对应的模型部分的前向传播之前，负责该划分的数据并行进程可以将权重广播到所有数据并行进程。该划分的前向传播完成后，参数可以被丢弃。总通信量为Ψ×Nd/Nd = Ψ。换句话说，我们通过将参数全收集操作分布在整个前向传播过程中，并丢弃一旦使用过的参数来重新安排全收集操作。但请注意，对于反向传播，这个全收集需要以相反的顺序再次发生。

       因此，总通信量是这些全收集产生的通信量加上梯度的减少-分散产生的通信量。总体来说，总体积为3Ψ，与基线相比增加了1.5倍。梯度和参数划分都利用了这样一个见解，即不是所有时间都需要梯度和参数的所有状态，以通过审慎地通信状态来优化内存。

8 ZeRO-R的通信分析

       我们比较了分区激活检查点（Pa）在ZeRO-R中的通信量与基线MP，并显示Pa的通信量增加通常不到基线MP的十分之一。此外，我们分析了Pa与DP通信量之间的通信开销，以确定Pa何时通过允许更大批量大小并减少DP通信来提高效率。我们利用这种分析来决定是否以及何时应用Pa以及Pa+cpu。

       通信量的分区激活检查点的权衡取决于模型大小、检查点策略和MP策略。为了分享具体见解，我们在使用SOTA MP方法实现的基于变压器的模型背景下执行了分析，例如Megatron-LM。

       在Megatron-LM中使用激活检查点时，每个变压器块在前向传播中执行两个全归约操作，大小为批量×序列长度×隐藏维度，前向重新计算和反向传播中再次执行两个全归约操作。每个块的总通信量是12×序列长度×隐藏维度，因为全归约的通信量是2×消息大小。

       当ZeRO-R分区激活检查点时，在每个激活检查点上的反向传播前向重新计算之前，需要执行额外的全收集操作。通常，我们为每个变压器块检查点输入激活，每个变压器块需要一个全收集。因此，Pa的通信开销是序列长度*隐藏维度，因为全收集的通信量是消息大小。因此，Pa的总通信开销不到模型并行性原始通信量的10%。

       当与DP一起使用MP时，可以使用Pa减少数据并行通信量的数量级，但以模型并行通信量增加10%为代价，并在数据并行通信成为性能瓶颈时显著提高效率。请注意，Pa通过MP程度减少了激活内存消耗，从而成比例地增加了批量大小。对于大型模型，MP可以达到16（DGX-2节点上的#GPU），允许最多增加16倍的批量大小。数据并行训练的通信量与批量大小成反比。因此，由于Pa，批量大小的数量级增加可能导致数据并行通信量的数量级减少。

       最后，如果应用Pa+cpu，分区激活检查点将被卸载到CPU，将激活内存需求减少到接近零，但与Pa相比，增加了2倍的数据移动到和从CPU内存的开销。在极端情况下，由于小批量大小，即使有Pa，DP通信量也是主要瓶颈，Pa+cpu可以通过增加批量大小来提高效率，只要CPU数据传输开销小于DP通信量开销，通常对于小批量大小是正确的。

       根据模型和硬件特性，我们利用上述分析来决定是否以及何时应用Pa和Pa+cpu。

9 向1万亿参数迈进

        目前发布的最大模型在100亿参数的范围内，这已经很难训练了。到达1万亿参数，大3个数量级，将不可避免地发生，但道路将充满障碍、惊喜和创新。虽然我们不声称知道或解决所有这些问题，但ZeRO解决了最基本的系统挑战：使得当前硬件能够适应这种规模的模型，同时允许以良好的系统可扩展性进行训练。

       从最先进的技术跃升(A Leap from State-of-Art) 最先进的框架Megatron可以在DGX-2系统中以可接受的吞吐量训练16 - 20B参数模型。通过在多个DGX节点之间实现模型并行性来进一步扩展，会因为节点间带宽有限而导致效率显著下降。

       ZeRO大大增加了可以有效运行的模型大小。它使当前硬件能够运行显著更大的模型，而无需将精细的模型并行性扩展到节点边界之外。如表1所示，ZeRO，在打开所有优化（Pos+g+p）的情况下，可以在仅使用DP的情况下在1024个GPU上容纳超过1万亿参数。另外，当与模型并行性结合时（如表2所示），ZeRO可以在1024个GPU上容纳超过1万亿参数，每个DGX2节点内有16路模型并行性，并且节点之间有64路数据并行性。有效运行一个拥有1万亿参数的模型不再是不可能的！

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

       为了理解资源需求，我们与Bert-Large进行了简短的比较。Bert-Large可以在1024 GPU DGX-2H集群上在67分钟内进行训练。1万亿参数模型可以轻松包含比Bert-Large模型更多的3000倍（1万亿/3.3亿）计算量。即使我们假设相同的序列长度和训练模型所需的样本总数相同，训练1T模型将需要140天，假设硬件相同且计算效率相似。实际上，随着模型大小的增加，数据样本和序列长度可能会增加，需要超过一年的时间来训练。在合理的时间内训练1T参数模型可能需要一个exa-flop系统。但是，当这样的计算能力变得可用时，我们希望ZeRO将提供系统技术来有效运行1T模型。

 

10 实现和评估

我们将实现重点放在支持高效训练约100B参数的模型上，这些模型比今天发布的最大模型（例如，T5-11B）大一个数量级，同时可以在当前硬件（例如，1K V100 GPUs）上在合理的时间范围内进行训练。我们实现并评估了ZeRO中的一部分优化—— ZeRO-DP中的Pos+g以及ZeRO-R —— 这使我们能够实现这个目标。我们将此实现称为ZeRO-100B。我们的结果显示，ZeRO-100B可以在400个GPU上高效运行最多170B参数的模型，比SOTA大8倍，速度快10倍，并且可用性得到改善。ZeRO-100B为Turing-NLG提供动力，这是目前世界上最大的模型，具有新的SOTA准确度。

10.1 实现和方法

       实现我们在PyTorch中实现了ZeRO-100B，包括Pos+g和ZeRO-R中的完整优化集。其接口与作为torch.nn.module实现的任何模型兼容。用户可以简单地使用此接口包装他们的模型，并利用ZeRO增强的DP，就像他们使用经典的DP一样。用户不需要修改他们的模型。ZeRO增强的DP可以与任何形式的MP结合使用，包括Megatron-LM。 硬件我们在400 V100 GPUs（25个DGX-2节点）的集群上进行了实验，其节点间通信带宽为800 Gbps。 基线对于没有MP的实验，我们使用torch的分布式数据并行（DDP）作为基线。对于具有MP的实验，我们使用Megatron-LM，因为据我们所知，它是最先进的。我们使用了NVIDIA 4的开源版本的Megatron-LM，截止日期为2019年9月。最近的Megatron-LM结果报告了使用32个DGX-2节点（共512个32GB V100 GPU）扩展到16B参数模型的能力。 ZeRO实验没有MP，使用ZeRO-100B中的ZeRO增强的DP实现。具有MP的实验将ZeRO增强的DP与Megatron-LM的MP结合使用。 模型配置本节介绍的模型是基于GPT-2的变压器模型。我们改变了隐藏维度和层数以获得不同数量参数的模型。表4显示了我们实验中使用的配置参数，附录中提供了更多细节。

10.2 速度和模型大小

       ZeRO-100B可以在400个GPU上高效运行最多170B参数的模型，比Megatron-LM大8倍以上。图2显示了使用ZeRO-100B与MP相比使用Megatron MP单独的不同模型大小的每GPU吞吐量。ZeRO-100B对于8B至100B参数的模型平均实现了15 PetaFlops的持续吞吐量（超过峰值的30％）。相比之下，基线MP性能随着模型大小的增加而迅速下降：MP在GPU之间产生高通信量，并且为了适应更大的模型而超出单个节点会导致通信带宽从每个链接的300GB / sec（NVSwitch）下降到12.5 GB / sec（Infiniband EDR），从而导致显著的性能下降。ZeRO-100B实现了高达10倍的速度提升，显著优于大型模型。 对于ZeRO-100B，超过100B的性能稍有下降是由于缺乏足够的内存来运行更大的批量大小。随着我们增加GPU数量，我们预计性能会提高，因为ZeRO-100B随着我们讨论的下一个GPU数量的增加而显示出超线性的速度提升。

10.3 超线性可扩展性

       ZeRO-100B展示了非常大型模型大小的超线性可扩展性。图3显示了从64到400个GPU的60B参数模型的可扩展性结果，我们预计这种趋势将在有更多GPU的情况下继续下去。Pos+g减少了ZeRO-100B的每GPU内存消耗随DP程度的增加，使ZeRO-100B能够在每GPU上适应更大的批量大小，从而提高吞吐量，因为算术强度在增加。

10.4 大模型训练的民主化

       使用MP和PP对许多数据科学家来说具有挑战性，这是训练大型模型的众所周知的障碍。ZeRO不需要对模型本身进行任何更改，可以像基线DP一样简单地使用，同时显著提高了模型大小和速度。图4显示，ZeRO-100B可以在128个GPU上训练最多13B参数的模型，而不使用MP，平均每GPU的吞吐量超过40 TFlops。相比之下，没有ZeRO，仅使用DP的最大可训练模型具有1.4B参数，每GPU的吞吐量不到20 TFlops。此外，在没有MP的通信开销的情况下，可以在没有非常快的节点间互连（如NVLINK或NVSwitch）的低端计算节点上训练这些模型，而这是实现MP效率的必要条件。

10.5 内存和性能分析

       我们探讨了不同优化对最大模型大小、内存消耗和性能的益处和影响。表3中将这些优化称为配置1到5（C1-C5）。

       最大模型大小(Maximum Model Size) 图6显示了通过启用不同的ZeRO优化，对于固定批量大小和MP为16的情况下，可以训练的最大模型。当使用C1与C2进行训练时，模型大小从40B增加到60B，这是因为使用Pa使激活内存减少了16倍（MP程度），而使用C4跳到140B是因为启用Pos+g，与C2中的Pos相比，它将模型状态的内存需求减少了一半。使用C5将其增加到150B完全是因为通过将分区激活检查点卸载到CPU内存来进一步减少激活内存。

       最大缓存内存(Max Cached Memory) 图7显示了在每次训练迭代期间PyTorch缓存的最大内存，针对40B和100B参数模型。缓存内存大小的减少如从C1到C2所预期的那样。C2和C3之间的内存消耗差异取决于模型状态的大小与激活内存的大小比较，当激活内存较大时可能会增加，或者当模型状态较大时可能会减少。值得注意的是，对于40B，缓存的内存从C4到C5并没有减少，但对于100B确实减少了。这仅仅是因为100B的激活内存要大得多，以至于减少是明显的。当我们处理非常大的模型时，这使Pa+cpu成为适应更大批量大小的有价值的工具。在图8中，需要Pa+cpu才能使170B模型在不耗尽内存的情况下执行。

       最大可实现性能(Max Achievable Performance) 图8显示了不同优化集的最佳可实现性能。请注意，性能改进对应于优化之间的内存消耗减少。如前所述，较低的内存消耗允许更大的批量大小，从而提高性能。唯一的警告是60B参数模型在C4和C5之间的性能下降。尽管内存消耗较低，但C5会导致激活在CPU之间的移动，这在大多数情况下会导致性能较差，除非少数情况，模型非常大，以至于没有C5就无法运行，或者没有C5可以运行的批量大小非常小（例如图8中的170B参数模型）。在训练过程中，仅在Pa+cpu有益时才会打开它。

10.6 Turing-NLG，具有17B参数的SOTA语言模型

       截至2020年5月12日，Turing-NLG是世界上最大的模型，拥有超过17B参数。它使用Webtext-103困惑度10.21实现了新的SOTA语言模型。使用ZeRO-100B从头到尾训练了Turing-NLG，图5显示了与先前的SOTA模型Megatron-LM 8.3B参数模型相比，300K迭代过程中的验证困惑度。ZeRO-100B为此模型实现了41.4 TFlops/GPU的持续吞吐量。

11 结束语

       从高性能计算和系统的角度来看，我们相信ZeRO代表了大型模型训练领域的革命性转变。虽然我们的实现，ZeRO-100B，使模型大小增加了8倍，吞吐量提高了10倍以上，实现了现代GPU集群上的超线性加速，并训练了世界上最大的模型，但这只是冰山一角。ZeRO的全部内容有潜力将模型大小再增加一个数量级，使未来的万亿参数模型的训练成为可能。

       也许，我们对ZeRO最乐观的地方是它对数据科学家没有任何障碍。与现有的诸如MP和PP之类的方法不同，不需要对模型进行重构，使用起来就像标准的DP一样简单，使ZeRO成为未来大型模型训练研究的首选候选者。通过开源和社区反馈，我们计划使DL社区完全可以访问ZeRO，以催化大规模大型模型训练的演变和民主化。

致谢

      我们感谢Junhua Wang的宝贵支持和建议。我们感谢Minjia Zhang、Elton Zheng、Shaden Smith、Reza Yazdani Aminabadi、Arash Ashari和Niranjan Uma Naresh对评估工作的良好反馈和帮助。我们感谢Brandon Norick、Corby Rossett、Gopi Kumar、Jack Zhang、Jing Zhao、Payal Bajaj、Rangan Majumder、Saksham Singhal、Saurabh Tiwary和Xia Song的许多有益的讨论和建议。