【大模型训练】ZeRO-infinity: breaking the GPU memory wall for extreme scale deep learning

摘要

在过去的三年中，最大的密集深度学习模型已经增长了1000倍以上，达到数千亿个参数，而GPU内存仅增长了5倍（从16 GB到80 GB）。因此，模型规模的增长主要通过系统创新来支持，这些创新使得大型模型能够适应多个GPU的聚合内存。然而，我们正接近GPU内存的极限。仅仅为了适应一个拥有万亿个参数的模型的训练，就需要800个NVIDIA V100 GPU，而这样的集群对于大多数数据科学家来说是不可企及的。此外，以这样的规模训练模型需要复杂的并行技术组合，这给数据科学家重构模型带来了巨大负担。
在本文中，我们提出了一种名为ZeRO-Infinity的新型异构系统技术，利用GPU、CPU和NVMe内存，在有限的资源上实现了前所未有的模型规模，而无需对模型代码进行重构。同时，它实现了出色的训练吞吐量和可扩展性，不受有限的CPU或NVMe带宽的限制。ZeRO-Infinity可以在当前一代GPU集群上适应具有数万亿甚至数百万亿参数的模型进行训练。它可以用于在单个NVIDIA DGX-2节点上微调万亿参数模型，使得大型模型更易于使用。在训练吞吐量和可扩展性方面，它在512个NVIDIA V100 GPU上可以持续达到超过25 petaflops的性能（达到峰值性能的40%），同时还展示了超线性的可扩展性。ZeRO-Infinity的开源实现可通过DeepSpeed 1获得。

引言

在过去的三年中，深度学习中最大的已训练密集模型增长了超过1000倍，从一亿个参数（ELMo [6]）增长到超过一千亿个参数（GPT-3 [4]）。相比之下，单个GPU的内存仅增加了5倍（从16 GB增长到80 GB）。因此，模型规模的增长主要通过系统技术的进步来实现，以训练大型深度学习模型，其中包括模型并行 [7]、流水线并行[8-10]和ZeRO [11, 12]等并行技术，为训练更大、更强大的模型铺平了道路。

目前大型模型训练技术的最新进展是三维并行（3D parallelism [13, 14]），它将模型（张量切片）和流水线并行与数据并行相结合，以高效地将深度学习训练扩展到拥有数万亿个参数的模型。例如，DeepSpeed实现的三维并行ism可以通过充分利用集群的总体GPU内存，在800个NVIDIA V100 GPU上扩展到超过一万亿个参数的规模[15]。

尽管3D并行在大型模型训练方面具有很大的能力，但我们现在已经接近GPU内存的极限[16]。聚合GPU内存简单地无法满足模型规模的增长。即使使用具有80 GB内存的最新NVIDIA A100 GPU，3D并行也需要320个GPU才能容纳一个万亿参数的模型进行训练，而要扩展到未来拥有一百万亿参数的模型，即使假设未来几年内GPU内存增加5倍，也需要超过6,000个GPU。我们无法再继续通过GPU内存来支持模型规模的持续增长。

GPU内存墙也限制了数据科学家甚至是访问当今的大型模型，尤其是对于微调来说如此。大型模型首先在大量通用数据上进行预训练，然后通过微调可以使同一模型适用于各种应用。虽然对一个拥有数千亿参数的模型进行预训练可能需要数百万个GPU计算小时，但进行微调要便宜得多，只需要较少的GPU计算小时，并且可以在单个计算节点上使用少量的GPU完成。虽然许多企业和用户可以获得这样的计算资源，但不幸的是，它们受到这些计算节点可用内存的限制，从而限制了可以进行微调的模型规模。这使得大型模型的微调对于大多数没有大规模GPU集群访问权限的研究人员和公司来说不可行。例如，即使单个DGX-2节点（16个GPU）具有足够的计算能力，在合理的时间内微调GPT-3也需要超过8个DGX-2节点（128个GPU）进行训练，仅仅是为了适应模型的大小。

除了GPU内存墙之外，训练大型模型的最新技术在可用性和灵活性方面也存在限制。如上所述，三维并行需要以复杂的方式结合数据并行、模型并行和流水线并行，以实现数千亿或数万亿个参数。虽然这样的系统可以非常高效，但需要数据科学家进行重大的模型代码重构，将单个GPU操作符替换为张量切片版本，并将模型分割成负载平衡的流水线阶段。这也使得三维并行在支持的模型类型方面缺乏灵活性。具有复杂依赖关系的模型不能轻松地转换为负载平衡的流水线模式。

鉴于大型模型训练所面临的挑战和问题，这篇论文提出了一种名为ZeRO-Infinity的新型系统，旨在解决这些问题并推动模型规模的进一步增长。ZeRO-Infinity旨在解决以下挑战：

• 支持模型规模的未来1000倍增长：ZeRO-Infinity提供了一种解决方案，以支持从GPT-3等拥有1750亿参数的模型到拥有数万亿参数的模型的指数级增长。
• 提高数据科学家的可访问性：ZeRO-Infinity使资源有限的数据科学家能够有效地训练和微调大型模型。
• 简化大型模型训练：ZeRO-Infinity旨在简化大型模型训练的过程，消除对大规模模型重构和多种并行处理形式的需求。它引入了简化的技术，减轻了训练大规模模型所涉及的复杂性。

整体阐述

ZeRO-Infinity是一种数据并行训练的形式，但与标准的数据并行训练不同，标准数据并行训练将模型状态（如参数、梯度和优化器状态）在所有数据并行进程之间进行复制，而ZeRO-Infinity将它们分区以充分利用所有数据并行进程的聚合内存。在训练过程中，ZeRO-Infinity使用通信集合（communication collectives）仅收集当前所需的模型状态。这类似于一种内存高效的数据并行形式，称为ZeRO [11]。然而，与仅在GPU上保留本地分区的ZeRO不同，ZeRO-Infinity将它们卸载到CPU或NVMe内存，并根据需要在GPU、CPU和NVMe之间移动，使ZeRO-Infinity能够充分利用CPU和NVMe内存，而不仅仅是GPU内存。接下来，我们将讨论ZeRO-Infinity的创新之处，使其能够在不影响易用性的情况下实现前所未有的模型规模和出色的训练效率。

前所未有的模型规模

ZeRO-Infinity通过称为无限卸载引擎的异构内存访问创新技术，扩展了ZeRO技术家族[11, 12]。这使得ZeRO-Infinity能够在有限的GPU资源上支持大规模模型，同时利用CPU和NVMe内存。此外，ZeRO-Infinity还引入了一种称为内存中心分块的新颖GPU内存优化技术，以支持极大的单个层，即使逐层加载，也无法适应GPU内存。通过无限卸载引擎和内存中心分块，ZeRO-Infinity不仅支持模型规模的未来1000倍增长，还使数据科学家能够利用有限的GPU资源来访问大型模型。

出色的训练效率

ZeRO-Infinity引入了一种新颖的数据分区策略，利用所有设备的聚合内存带宽，我们称之为带宽中心分区，并将其与强大的通信重叠设计相结合，以及针对无限卸载引擎中高性能NVMe访问的优化。综合起来，ZeRO-Infinity提供了出色的训练效率，尽管将数据卸载到CPU或NVMe，并不受其有限带宽的限制。

易用性

有了ZeRO-Infinity，数据科学家不再需要将模型适应于3D并行处理等多种并行形式。这是由于上述ZeRO-Infinity中的内存中心分块，旨在降低大型单个层的GPU内存需求，否则需要模型并行处理（张量切片）才能将层适应GPU内存。此外，ZeRO-Infinity通过一种简化的实现，消除了对手动模型代码重构的需求，即使在通过自动化进行通信和数据分区所需的训练任意模型架构的情况下，也能实现对数万亿参数的扩展。

这篇论文的主要贡献如下：

• 大型模型训练的内存和性能特征化：描述了大型模型训练不同组件的内存需求（第3节）以及为了实现高效训练所需的带宽需求（第4节）。
• ZeRO-Infinity（第5和第6节）：一种新颖的深度学习训练系统技术，由五种创新技术组成，以解决内存和带宽需求，提供前所未有的可访问和易于使用的模型规模，并实现出色的训练效率：
  • i）无限卸载引擎，以同时利用现代集群上的异构架构，包括GPU、CPU和NVMe内存，以及GPU和CPU计算；
  • ii）内存中心分块，处理大型操作符而无需模型并行处理；
  • iii）带宽中心分区，利用所有并行设备的聚合内存带宽；
  • iv）重叠中心设计，实现计算和通信的重叠；
  • v）受易用性启发的实现，避免模型代码重构。
• 对ZeRO-Infinity的广泛评估：展示了32个NVIDIA DGX-2节点（512个V100 GPU）上运行320万亿参数的前所未有的规模，以及在相同硬件上实现超过25 petaflops的吞吐量的出色训练效率，以及万亿参数模型的超线性可扩展性，能够在单个DGX-2节点上微调万亿参数模型，无需使用任何模型并行处理或模型代码重构，并讨论了ZeRO-Infinity中不同技术对模型规模和效率的影响（第7节）。
• 关于ZeRO-Infinity及其对未来硬件系统设计的潜在影响的讨论（第8节）。
• ZeRO-Infinity的开源实现，在深度学习社区中得到广泛采用。

2 背景和相关工作

数据、模型、流水线和3D并行性
并行化是在大规模模型训练中的一种重要策略。对于适合设备内存的模型，可以使用数据并行性（DP）将训练扩展到多个设备上。当模型不适合设备内存时，可以使用模型并行性（MP）[7, 17, 18]和流水线并行性（PP）[7-9]来将模型在进程之间进行垂直和水平划分。3D并行性[14, 15]结合了数据、模型和流水线并行性的优点，能够高效地扩展到数万亿个参数。虽然3D并行性可以非常高效，但它需要：i）对模型进行重构，将模型划分为模型并行和流水线并行组件；ii）具有复杂依赖图的模型很难表达为负载平衡的流水线；iii）模型大小受限于总可用的GPU内存。关于深度学习中并行性的详细调查，可以参考Ben-Nun和Hoefler的研究[19]。

ZeRO：零冗余优化器
ZeRO [11]通过将三个模型状态（即优化器状态、梯度和参数）在数据并行进程之间进行划分而不是复制，消除了数据并行进程之间的内存冗余。通过这样做，相比经典的数据并行性，它提高了内存效率，同时保持了计算粒度和通信效率。ZeRO有三个阶段，对应于三个模型状态：第一阶段（ZeRO-1）只划分优化器状态，第二阶段（ZeRO-2）划分优化器状态和梯度，最后一阶段（ZeRO-3）划分所有三个模型状态。在ZeRO-3中，模型的每个层的参数由唯一的数据并行进程拥有。在训练过程中，ZeRO-3确保在运算符执行之前，前向或反向传递所需的参数可供使用，通过从拥有者进程发出广播通信集合。在运算符执行之后，ZeRO-3还会删除参数，因为在下次运算符的前向或反向传递之前不再需要它们。此外，在训练的参数更新阶段，ZeRO-3确保每个数据并行进程仅更新其拥有的参数对应的优化器状态。因此，ZeRO-3可以在整个训练过程中保持所有模型状态的划分，除了立即计算所需的参数。

异构训练方法
在几种基于CPU内存的异构训练方法[20-26]中，ZeRO-Offload [12]是多GPU上大型模型训练的最新技术。ZeRO-Offload建立在ZeRO-2之上，将梯度和优化器状态存储在CPU内存中。在GPU设备不足以存储优化器状态和梯度时，ZeRO-Offload利用CPU内存。然而，它仍然需要将参数存储在GPU内存中，并在所有设备上复制。因此，ZeRO-Offload的模型规模受限于单个GPU设备的内存可容纳的参数总数。由于数据划分不理想和PCIe带宽有限，ZeRO-Offload还需要较大的批量大小以保持高效。我们通过ZeRO-Infinity来解决ZeRO-Offload的这些限制。

减少激活内存
激活是在前向传播过程中产生的中间结果，需要保留以计算反向传播中的梯度。已经有多个工作致力于通过压缩[28]、激活检查点[29, 30]或实时分析[31]来减少激活所需的内存。ZeRO-Infinity与激活检查点一起工作，以减少激活内存的使用。

Adam优化器和混合精度训练
自适应优化方法[32-35]对于实现最先进的性能和准确性，以有效地训练大型模型非常重要。与随机梯度下降（SGD）相比，它们在每个模型参数和梯度上维护了细粒度的一阶和二阶统计信息，但代价是显著的内存占用。Adam [33]是在大型模型训练中最常用的优化器。大型模型训练通常以混合精度进行，即前向和反向传播使用FP16，参数更新使用FP32 [36]。这利用了现代GPU上可用的张量核心单元的性能加速。

3 内存需求

本节描述了深度学习训练的内存需求。虽然我们的方法是通用的，但我们重点关注基于Transformer [38]的架构，因为所有超过十亿参数的SOTA模型都遵循这个架构。我们的分析假设使用Adam优化器进行混合精度训练，因为这是训练Transformer模型的事实标准。训练所需的内存可以分为两个组件：
内存可以从两个方面优化

模型状态：梯度，参数，优化器状态
剩余状态：中间激活值、其取决于模型架构、批量大小（𝑏𝑠𝑧）和序列长度（𝑠𝑒𝑞），可能非常大

模型状态的内存

总显存占用$240l{h}^2$ 字节

为了方便分析，定义 transformer 模型的层数为 l，隐藏层维度为 h，注意力头数为 a。词表大小为 v，训练数据的批次大小为 b，序列长度为 s。transformer模型由 l 个相同的层组成，每个层分为 self-attention 块和 MLP 块

• self-attention：参数有 q, k, v 的权重矩阵和 wq, wk, wv 的偏置 bq, bk, bv，输出权重矩阵 wo 和偏置 bo，4 个权重矩阵的形状为 [h, h]，4个偏置的形状为 [h]。所以 self-attention 块的参数量为 4h^2 + 4h
• MLP：由 2 个线性层组成，一般第一个线性层是先将维度从 h 映射到 4h ，第二个线性层再将维度从 4h 映射到 h。第一个线性层的权重矩阵 的形状为 [h,4h] ，偏置的形状为 [4h] 。第二个线性层权重矩阵的形状为 [4h,h] ，偏置形状为 [h] 。所以 MLP 块的参数量为 8h^2+5h
• layer normalization：self-attention 块和 MLP 块各有一个 layer normalization，包含可训练的缩放参数和平移参数，形状都是 [h] 。2个layer normalization的参数量为 4h

所以每个 transformer 层的参数量为 12h^2+13h，此外还有

• embdding 层：维度通常等于隐藏层维度 h，参数量为 vh。最后的输出层的权重矩阵通常与 embdding 层是参数共享的
• 位置编码：如果采用可训练式的位置编码，会有一些可训练模型参数，数量比较少。如果采用相对位置编码，例如 RoPE 和 ALiBi，则不包含可训练的模型参数。这里忽略这部分参数
  综上，l 层 transformer 模型的可训练模型参数量为 $l(12h^2+13h)+vh$。当隐藏维度 h 较大时，可以忽略一次项，近似为$12l{h}^2$

目前主流的大模型训练方法是使用 Adam 优化器进行混合精度训练，参数和梯度存储为 FP16，优化器状态包括 FP32 的动量、方差、参数和梯度，因此每个参数需要 $20B（2\ast 2b+4\ast 4b）$ 的显存，总显存占用就是$240l{h}^2$ 字节

图2a第5列显示了存储类似于GPT-3的Transformer模型的模型状态所需的内存，该模型具有1000亿到1万亿个参数，通过调整隐藏维度和层数进行创建。为了将内存需求放入上下文中，图2b第3列显示了在单个NVIDIA V100 DGX-2设备以及DGX2 SuperPOD集群上可用的总GPU内存。请注意，仅为了适应1000亿参数的模型，就需要64个GPU。适应1万亿参数的模型需要超过512个GPU，而10万亿参数的模型甚至超出了一个庞大的1536个GPU集群的范围。

剩余状态内存

内存用于残余状态：残余状态主要包括激活内存，其取决于模型架构、批量大小（b）和序列长度（s），可能非常大。从积极的一面来看，可以通过激活检查点技术[29]显著减少激活内存，但代价是增加了0.33倍的额外计算量。

中间激活值可以理解为前向传递过程中计算得到的，并在后向传递过程中需要用到的所有 tensor

• layer normalization 层：计算梯度时需要用到层的输入 bh字节，输入的均值和方差 b字节，由于 h 一般是千位，所以 layer normalization 显存近似为bh字节
• dropout 层：在训练中存储为 mask 矩阵，每个元素只占 1 字节

每个 transformer 层包含了一个 self-attention 块和 MLP 块，并分别对应一个 layer normalization 连接

• self-attention
  
  • Q, K, V：需要保存共同输入 x，这就是中间激活，x shape 为 [b, s, h]，元素个数为 bsh，占用显存为 2bsh 字节。这里的 shape 中没有 a 是因为在计算时，每个 head 对应的 hidden size 被除以了 head，一乘一除，化简了$Q{K}^T$，需要保存 Q, K，shape 都是 [b, s, h]，占用显存大小为 $4bsh$
  • softmax()：需要保存函数输入 $Q{K}^T$，Q shape [b, a, s, h]，$K^T$ shape [b, a, h, s]，$Q{K}^T$ shape [b, a, s, s]，占用显存为 $2b{s}^{2}a$。
  • dropout：计算完 softmax 后得到 score，需要计算 dropout，保存 mask 矩阵，shape 与$Q{K}^T$ 相同，占用显存为$b{s}^{2}a$
  • 计算 v 上的 attention ，score * v，需要保存 score：$2b{s}^{2}a$，v：$2bsh$，总合 $2b{s}^{2}a+2bsh$
  • 计算输出映射和一个 dropout，输出映射需要保存其输入，2bsh 字节，结合 dropout 3bsh

综上所述，self-attention 块中间激活占用显存为 11bsh + 5bs^2a 字节

• mlp
  
  linear_1：2bsh
  激活函数：8bsh
  linear_2：8bsh
  dropout：bsh

综上所述，MLP 块中间激活占用显存为 19bsh 字节

另外，self-attention 块和 MLP 块共对应两个 layer normalization，其输入合计为 2bsh，中间激活 4bsh

每个 transformer 层中间激活占用显存为 $34bsh+5b{s}^{2}a$ 字节

以 GPT3-175B 模型为例，模型配置如下

layer（l）：96
hidden size（h）：12288
attention head（a）：96
sequence length（s）：2048
假设采用混合精度训练，都采用 bf16 或 fp16 存储，每个元素占2个bytes，则模型参数占用显存为 350 GB

图2a的第7列显示了假设我们在每个Transformer块中存储一个激活时，用于存储激活检查点的内存需求，假设批量大小为32，序列长度为1024。许多现代GPU集群每个节点有8-16个GPU，因此我们选择每个GPU的批量大小为2-4，从而得到每个节点内的激活批量大小为32，作为对每个节点内激活的保守估计。虽然生成的激活检查点比完整的激活集合（第6列）小几个数量级，但在超过一万亿参数的情况下，对于所考虑的批量大小和序列长度，它们仍然变得过大，无法适应GPU内存。

模型状态工作内存 (MSWM)

模型状态工作内存（MSWM）是在将所有模型状态卸载到CPU或NVMe之后，在模型中执行正向或反向传播所需的最小GPU内存量。
这大致等于模型中最大单个运算符的参数和梯度的大小，因为至少需要足够的内存来保存参数及其梯度以供反向传播使用。
对于基于Transformer的模型，最大的运算符是一个线性层，将隐藏状态从ℎ𝑑转换为4ℎ𝑑。该线性层的参数和梯度大小为4 × ℎ𝑑 × 4ℎ𝑑字节。请注意，MSWM（图2a第8列）在超过1000亿个参数时显著增长，需要多个连续的内存千兆字节，这可能导致训练过程中由于缺少足够的连续内存来满足这些要求而导致内存耗尽。3D并行等最先进的方法通过模型并行性来解决这个问题，将单个运算符分割到多个GPU上。在5.1.3节中，我们将讨论一种新颖的方法，可以解决这种大规模模型状态工作内存的问题，而无需使用模型并行性。

激活值工作内存 (AWM)

是在进行 BWD 之前，在 BWD 中重计算激活所需的内存。这取决于两个连续激活值 checkpoint 之间的激活大小

例如，如果我们在每个Transformer块中创建一个激活checkpoint ，那么内存大小由每个Transformer块的总激活大小确定。这可以用以下公式近似表示（以字节为单位）：$16h+2\ast att{n}_{h}eads\ast s$。图2a第8列显示，AWM在超过1万亿个参数时也会变得很大。与只由单个参数和梯度组成的MSWM不同，AWM 由几十个中间激活层参数组成，只要总的 AWM 占用满足显存，不会遇到连续空间不足的问题

4 带宽需求

假设工作负载的执行没有任何计算和通信重叠，我们可以使用峰值计算吞吐量（peaktp）、数据传输带宽（bw）及其算术强度（ait）来估计训练效率。

工作负载的算术强度（AIT）是总计算量与计算所需的数据之间的比率。它描述了每次数据传输的计算量。较高的算术强度意味着对数据传输带宽的要求较低，因为加速器在每次加载数据时可以执行更多的计算。

效率指标可以如下推导：

通过考虑效率指标，我们可以评估有限带宽（例如，当带到CPU和NVMe内存时）对训练效率的影响。如果数据传输成为瓶颈并限制计算吞吐量，较低的带宽可能会导致效率降低。因此，在带到其他内存时，重要的是在计算吞吐量、数据传输带宽和算术强度之间进行权衡，以优化训练效率。

在DL训练中量化AIT

模型状态和激活检查点可以具有不同的算术强度（AIT）。我们可以通过首先确定DL训练每个迭代中的总计算量，然后确定每个模型状态和激活的数据传输量来对其进行量化。

每次迭代的总计算量
每次迭代的总计算量主要由Transformer的线性层中的计算支配。对于正向传播，可以近似为参数数量、序列长度和批量大小的函数，表示为2bsp。反向传播的成本大约是正向传播的两倍。此外，激活检查点需要在反向传播期间作为重新计算的一部分进行额外的正向计算。因此，每次迭代的总计算量为：8bsp

每次迭代的总计算量为 8 bsp

参数和梯度AIT
在正向传播和反向传播过程中，模型参数必须至少加载两次，即在正向传播期间和实际的反向传播期间，从源位置加载到GPU寄存器中，导致数据传输量为2p。在存在激活检查点的情况下，参数可能会额外加载一次，用于在反向传播过程中的重新计算，增加了1p的数据传输量。此外，梯度必须至少被存储一次，从GPU寄存器存储到最终位置，增加了1p的数据传输量。因此，假设参数和梯度被存储在相同的最终位置上，正向传播和反向传播期间的总数据传输量将为4p，即以字节表示的8p。每次迭代的总计算量由第4.1节给出。因此，相对于参数和梯度，算术强度（AIT）为bs。

优化器状态AIT
在优化器步骤中，优化器状态必须至少读取一次，并且优化器状态必须至少写入一次。因此，总的数据传输量为2optimizer_states，这大约是216*p字节。每次迭代的总计算量由第4.1节给出。因此，在完整的训练迭代中，相对于优化器状态的算术强度（AIT）为$bs/4$。

激活检查点的AIT
在正向传播期间，激活检查点必须保存到其最终位置，并在反向传播期间检索。因此，相对于激活检查点的总数据传输量（以字节为单位）由2*total_activation_checkpoints_in_bytes给出，该值由公式（1）中的$4nl/c\ast hbs$给出。每次迭代的总计算量由第4.1节给出。因此，相对于激活检查点的算术强度（AIT）为24hc。

带宽需求

由于AIT的变化，模型状态和激活检查点对于实现良好的效率具有非常不同的带宽需求。前者仅取决于批量大小和序列长度，而后者仅取决于激活检查点的频率和模型的隐藏维度大小。除了AIT，效率的带宽需求还取决于paektp。

使用paektp和ait，我们首先展示了效率如何随着不同模型和残差状态的带宽变化而变化，然后讨论了这些状态对于DL训练的带宽需求。在这里，我们以NVIDIA V100 DGX-2 SuperPOD集群作为示例平台。使用来自第4.1节的ait表达式和基于公式（2）的效率指标，图3显示了效率与可用带宽之间的关系，相对于参数和梯度、优化器状态以及激活检查点。

为了生成这些图，我们根据第4.1节推导出的表达式计算了ait，针对不同的批量大小、序列长度和模型配置进行变化。具体而言，我们使用序列长度为1024，与GPT-2 [2]、Megatron-LM [7]和Turing-NLG [41]使用的序列长度相同。我们将批量大小范围从1到16进行变化，分别捕捉大型GPU和小型GPU实验。当在大量GPU上运行时，每个GPU使用小批量大小，而当在相对较少的GPU上进行训练时，每个GPU使用大批量大小，以保持合理的有效批量大小进行训练。我们的隐藏大小范围从8K到64K，代表具有数千亿参数的模型到如图2a所示的数万亿参数的模型。为了确定此分析的paektp，我们使用经验方法。我们在单个NVIDIA V100 DGX-2盒上运行具有上述配置的模型，并关闭所有非GPU通信，以模拟零外部通信开销的情况，或者等效地模拟无限的外部带宽情况。根据8K-64K的隐藏大小，实现的性能范围为62-78 TFlops/GPU。我们使用70 TFlops/GPU的平均值来表示此分析的paektp。

a：参数和梯度的带宽超过70 GB/s，即使是最小的 batch size，也可以实现50%以上的效率，在这个带宽下，只要 bz 足够大，数据移动可以与计算完全重叠，以实现100%的效率
b：相比 a，优化器状态需要高出近 4 倍的带宽才能达到 50% 效率。此外，优化器状态在 FWD 和 BWD 结束时更新，并且不能与计算重叠
c：启用激活值 checkpoint 后，即使 h 为 2k，2GB/s 的带宽也能维持 50% 以上的效率。一旦 h 超过 8k，带宽需求将降至 1 GB/s 以下

5 设计概况

ZeRO-Infinity的整体概述如图4所示，并在下面进行讨论。

5.1 Design for Unprecedented Scale

现代GPU集群在存储器方面非常异构化。除了GPU内存之外，它们还具有CPU内存以及超过GPU内存大小近50倍、CPU内存大小近20倍的大型NVMe存储（参见图2b）。我们开发了ZeRO-Infinity，这是一个用于DL训练的并行系统，通过利用现代GPU集群中的这些异构内存系统，可以突破GPU内存的限制。图1比较了3D并行性和ZeRO-Infinity的最大可实现模型规模。ZeRO-Infinity支持每个NVIDIA V100 DGX-2节点一万亿个参数，相比3D并行性增加了50倍。

5.1.1 模型状态的无限卸载引擎

ZeRO-Infinity是基于ZeRO-3 [11]构建的，它对所有模型状态进行分区，以消除内存冗余，如第2节所讨论的。与现有的ZeRO系列技术不同，ZeRO-Infinity设计了一个强大的离线机制，称为无限卸载引擎，可以根据内存需求将所有分区的模型状态卸载到CPU或NVMe内存，或者保留在GPU上。请注意，根据图2a和图2b，即使是一万亿参数模型所需的模型状态也可以适应DGX-2集群（1536个GPU）的聚合NVMe内存。因此，无限卸载引擎使得ZeRO-Infinity能够适应具有数百万亿参数的模型的模型状态。有关更多详细信息，请参阅第6节。

5.1.2 激活的CPU卸载

除了模型状态之外，当需要时，ZeRO-Infinity还可以将激活内存卸载到CPU内存。请注意，一个一万亿参数模型所需的激活检查点（0.76TB）可以轻松适应DGX-2系统上可用的1.5TB CPU内存，而一个一百万亿参数模型所需的3TB激活检查点则在下一代硬件的CPU内存范围内。因此，通过将激活检查点卸载到CPU内存，ZeRO-Infinity可以适应具有数百万亿参数的模型的激活检查点。

5.1.3 基于内存的工作内存划分

为了减少大型模型DL训练的工作内存需求，ZeRO-Infinity引入了一种称为内存中心划分的新技术，利用ZeRO-3的数据获取和释放模式，通过将大型操作符拆分为可以顺序执行的较小块（划分为小块）来降低工作内存需求。例如，为了减少大型线性操作符的工作内存，ZeRO-Infinity将操作符表示为由原始操作符中的参数块组成的较小线性操作符的数学等效序列，并依次执行它们。当与ZeRO-3结合使用时，可以逐个获取和释放每个块的参数和梯度，将工作内存按照块的数量进行降低。因此，ZeRO-Infinity可以支持任意大小的操作符，而不依赖于模型并行性来适应有限的GPU内存。

5.2 Design for Excellent Training Efficiency

将所有模型状态和激活卸载到CPU或NVMe只有在ZeRO-Infinity能够实现高效率的情况下才是可行的。实际上，这是非常具有挑战性的，因为CPU内存的带宽比GPU内存慢一个数量级，而NVMe的带宽比CPU内存带宽还慢一个数量级。此外，从GPU读取和写入这些内存的速度更慢（参见图2b）。在像DGX-2这样的系统上，根据我们在第4节中的分析，参数和梯度、优化器状态以及激活检查点的数据传输带宽必须大于70GB/s、1.5TB/s和1-4GB/s，才能实现高效的DL训练。在这里，我们讨论ZeRO-Infinity如何实现所需的带宽，以实现出色的效率。

5.2.1 参数和梯度的效率

参数和梯度的数据传输带宽必须大于70GB/s，接近DGX-2集群上可用的GPU-GPU带宽[42]。因此，像ZeRO3 [11]这样的DL并行训练解决方案，在从拥有者GPU广播参数到其他GPU之前，只要通信可以重叠，就可以高效运行，用于前向或反向传播。相反，从单个GPU到CPU内存或NVMe的12GB/s PCIe带宽（参见图2b），或者反过来，根本无法支持大规模的异构训练。

因此，现有的异构解决方案（如ZeRO-Offload）要求首先将参数从CPU移动到拥有者GPU，然后再进行广播，这就需要每个GPU具有显著大的批大小，以在有限的带宽下实现足够的效率。这带来了两个问题：i）对于大型模型，激活内存将变得过大，甚至无法适应CPU内存；ii）当扩展到数百或数千个GPU以实现有效收敛时，有效批大小将变得过大。

ZeRO-Infinity通过两种方式解决这些挑战：
i）基于带宽的分区：一种新颖的数据映射和并行数据检索策略，用于卸载的参数和梯度，使ZeRO-Infinity能够实现几乎无限的异构内存带宽（详见第6.1节）；
ii）一种以重叠为中心的设计，使ZeRO-Infinity能够将GPU-GPU通信与计算重叠，同时还可以通过PCIe实现NVMe-CPU和CPU-GPU通信的重叠（详见第5.1.3节）。

5.2.2 优化器状态的效率

与在前向和反向传播过程中按顺序使用和产生的参数和梯度不同，优化器状态可以同时并行更新。这个特性被ZeRO-3和ZeRO-Offload所利用，在所有可用的GPU和CPU之间并行存储和更新优化器状态，分别存储于GPU内存和CPU内存中。因此，聚合的GPU或CPU内存带宽可以随着GPU或CPU数量的增加而大大提高，超过所需的1.5TB/s。

由于ZeRO-Infinity是基于ZeRO-3构建的，当将优化器状态卸载到CPU内存时，它也可以利用聚合的GPU和CPU内存带宽以及聚合的CPU计算进行优化器步骤。然而，对于NVMe卸载，需要将数据从NVMe带到CPU内存，并以适合CPU内存的块的形式进行优化器步骤，一次处理一个块。因此，优化器步骤受到NVMe-CPU内存带宽的限制：虽然ZeRO-Infinity可以在多个节点上实现聚合的NVMe带宽，但至关重要的是要实现每个节点的接近峰值的NVMe带宽，以支持超过1.5TB/s的必要带宽，节点数量尽可能少，批大小尽可能小。此外，将数据从NVMe带入CPU内存，或者从CPU内存带入GPU内存的过程可能会导致GPU和CPU的内存碎片，即使还有大量内存可用，也可能导致内存不足。Infinity卸载引擎不仅可以实现接近峰值的NVMe带宽，还可以同时允许ZeRO-Infinity将NVMe到CPU的读取与CPU到NVMe的写入以及优化器步骤的CPU计算重叠，以便ZeRO-Infinity在少数GPU上以适度的批大小保持高效，在大量GPU上以较小的批大小保持高效。同时，它通过精心重用临时缓冲区来最小化内存碎片化。我们在第6节中详细讨论了Infinity卸载引擎中的优化。

5.2.3 激活的效率

在DGX-2节点上，每个GPU可以通过PCIe以大约3 GB/s的速度与CPU内存并行读写数据，允许将激活检查点卸载到CPU内存，同时保持大于80%的效率，适用于隐藏大小大于8K或更大的情况。为了在较小的隐藏大小下实现高效率，ZeRO-Infinity可以降低激活检查点的频率，并有效地将激活检查点的通信与GPU上的前向和反向计算重叠，包括向CPU内存传输和从CPU内存传输。

5.3 Design for Ease of Use

有了ZeRO-Infinity，数据科学家再也不需要调整他们的模型以适应多种形式的并行计算，比如3D并行计算。这得益于ZeRO-Infinity中讨论的基于内存的切片技术，旨在减少大型单层所需的GPU内存，否则这些层将需要模型并行计算（张量切片）来将层放入GPU内存中。

此外，ZeRO-Infinity在PyTorch中的实现方式消除了手动重构模型代码的需求，即使在扩展到数万亿个参数的规模时也是如此。这得益于三个自动化功能的简化实现：
i）在训练过程中，自动化数据移动以在参数在训练过程中被使用之前收集它们。ZeRO-Infinity通过将预先的前向/后向钩子注入PyTorch子模块中来实现这一点，这些钩子触发allgather集合操作来收集所需的参数。
ii）在子模块的前向/后向传递结束时，自动化参数分区，因为子模块所需的参数已不再需要。再次，ZeRO-Infinity通过将后向/前向钩子注入子模块中来分区参数，并可选择将其卸载到CPU或NVMe。
iii）在初始化过程中自动化模型分区，即使模型无法在单个GPU或CPU内存中容纳，仍然可以初始化模型，而无需手动将模型分区到数据并行处理中。ZeRO-Infinity通过包装所有模块类的构造函数来实现这一点，以便在初始化过程中创建每个子模块的参数后立即对其进行分区和卸载。整个模型不会完全实例化在单个数据并行处理中。

6 效率优化

6.1 基于带宽的分区

ZeRO-Infinity采用了一种新颖的数据映射和检索策略，以解决NVMe和CPU内存带宽的限制。与ZeRO [11]和ZeRO-Offload [12]不同，其中每个层的参数由单个数据并行进程拥有，并在需要时向其他进程进行广播，ZeRO-Infinity将单个参数分区到所有数据并行进程，并在需要访问参数时使用全局收集（allgather）而不是广播。

需要注意的是，当数据位于GPU上时，广播和全局收集通信集合对于数据移动量来说具有相同的通信成本。因此，对于仅基于GPU的训练而言，这没有任何区别。然而，当数据位于NVMe或CPU上时，情况就不同了。

在基于广播的方法中，由于每个参数完全由一个数据并行进程拥有，参数必须首先通过PCIe从其源位置（CPU或NVMe）传输到GPU内存，然后才能进行广播[8]。需要注意的是，此过程只能激活一个PCIe，而连接到所有其他GPU的所有PCIe链路都是空闲的。

相反，在ZeRO-Infinity中，采用了分区参数和基于全局收集的方法，所有PCIe链路都可以并行激活，每个链路带来参数的$1/d{p}^{th}$部分，其中𝑑𝑝是数据并行度。因此，从NVMe或CPU到GPU之间的有效通信带宽随𝑑𝑝的增加而线性增加。例如，基于广播的方法在DGX-2盒子上进行16路数据并行时，CPU/NVMe到GPU的带宽保持在约12 GB/s（使用PCIe Gen 3），但采用基于全局收集的方法时，可达到约48/25 GB/s（每个GPU分别为3.0/1.6 GB/s）的有效可实现带宽（见图2b），仅受到最大聚合PCIe带宽和DGX-2节点的最大NVMe带宽的限制。

从这里开始，带宽随着更多节点的增加而线性增长。当在大规模训练中训练大规模模型时，ZeRO-Infinity因此可以提供比训练所需的更多异构内存带宽（实际上是无限的）。例如，在64个DGX-2节点上，ZeRO-Infinity可以获得超过3TB/s的CPU内存带宽和超过1.5TB/s的NVMe带宽。

6.2 Overlap Centric Design

尽管ZeRO-Infinity可以在多节点设置中利用足够的异构内存带宽，但在单个GPU或单个节点设置中，带宽仍然可能成为瓶颈。即使是GPU-GPU的全局收集通信，在使用小批量大小时也会对效率产生重大影响（见图3）。此外，访问NVMe内存需要一个三步过程：i）从NVMe读取数据到CPU内存（nc-transfer），ii）将数据从CPU内存复制到GPU内存（cg-transfer），iii）执行全局收集以在所有GPU上构建完整的参数（gg-transfer）。这些数据移动的顺序性意味着，如果进行简单的处理，总通信时间将是这三个数据移动成本的总和，即使每个阶段的数据移动带宽在各自的情况下是足够的，也会导致效率低下。

为了解决这些问题，ZeRO-Infinity具有一个重叠引擎，不仅可以将GPU-GPU通信与GPU计算重叠，还可以同时重叠NVMe到CPU和CPU到GPU的通信。重叠引擎包括两个组件：i）用于重叠数据移动的动态预取器，以在前向或后向传递中使用参数之前重构参数所需的数据移动，ii）用于在后向计算中与梯度一起并行执行所需的数据移动的通信和卸载重叠机制。

ZeRO-Infinity中的动态预取器实时跟踪前向和后向计算，构建每个迭代的操作符序列的内部映射。在每个迭代中，预取器跟踪操作符序列的位置，并预取将来操作符所需的参数。预取器知道三步通信过程，因此可以将一个参数的nc-transfer与其他参数的cg-transfer和gg-transfer同时重叠。例如，在执行第ith个操作符之前，预取器可以分别为第i+3、i+2和i+1 1个操作符所需的参数调用nc、cg和gg-transfer。需要注意的是，所有这些数据移动都可以与执行第ith个操作符同时进行。此外，对于动态工作流程，ZeRO-Infinity可以更新操作符序列映射，从而在迭代过程中进行适当的预取，即使前向和后向传播在迭代之间发生变化。

类似地，在后向传递中，ZeRO-Infinity可以将（i+1）𝑡ℎ操作符的梯度的reduce-scatter与ith操作符的计算重叠，同时将（i+2）𝑡ℎ操作符的分区梯度从梯度的reduce-scatter传输到CPU或NVMe。

通过这种强大的重叠中心设计，即使在使用少量GPU和每个GPU的小批量大小进行训练时，ZeRO-Infinity可以隐藏大部分数据移动。

6.3 Infinity Offload Engine

无限卸载引擎由两个主要组件组成：DeepNVMe和固定内存管理层。

DeepNVMe是一个强大的C++ NVMe读/写库，它位于无限卸载引擎中。它支持批量读/写请求，并且可以异步完成，这意味着这些请求可以独立地发出并运行，而不会阻塞主要的执行流程。DeepNVMe还提供了显式的同步请求，用于刷新正在进行的读/写操作，以确保数据的一致性。通过利用DeepNVMe的异步特性，ZeRO-Infinity可以将这些NVMe请求与GPU/GPU或GPU/CPU的通信或计算重叠，提高整体效率。

DeepNVMe通过采用多种优化技术实现高性能。它对I/O请求进行积极的并行化处理，允许同时处理多个请求，无论这些请求来自单个用户线程还是多个用户线程。智能工作调度技术用于优化请求的顺序和时机，以最大化吞吐量。DeepNVMe避免了不必要的数据拷贝，最大程度地减少了在不同内存位置之间移动数据的开销。此外，它利用内存固定技术，将数据保留在固定的内存位置，进一步提高性能。这些优化共同使得DeepNVMe能够在NVMe存储设备上接近峰值的顺序读写带宽。

固定内存管理层负责管理有限的固定内存缓冲区。固定内存缓冲区对于在NVMe和CPU存储之间实现高性能的张量读取和写入是至关重要的。然而，固定内存是一种稀缺的系统资源，单个进程过度使用固定内存可能会对系统性能或稳定性产生负面影响。固定内存管理层通过重用少量固定内存（通常为数十GB）将整个模型状态（高达数十TB）卸载到CPU或NVMe中，以解决这个问题。内存缓冲区的重用策略有助于防止CPU和GPU内存的内存碎片化，优化内存利用率。此外，固定内存管理层为PyTorch张量提供了固定内存数据，使得可以对张量进行原地计算。这意味着可以直接对张量进行操作，无需进行额外的内存拷贝，从而在将其写入NVMe时提高带宽和性能。

总体而言，DeepNVMe和固定内存管理层的组合使得ZeRO-Infinity能够在训练过程中充分利用GPU/GPU和GPU/CPU之间的并行性，并实现高性能的NVMe读写操作。

实现细节

ZeRO Infinity 基于 PyTorch 代码实现，并且不用重构任何模型代码

Automating Data Movement
ZeRO-Infinity 需要协调模型参数、梯度和优化器状态数据移动，必须确保这些数据在被使用前移动到 GPU 内存中，并在使用之后重新分配位置

PyTorch 模型以层级模块的形式表示，代表着神经网络的各个层次。例如，Transformer 架构中包含了诸如自注意力和前馈网络等子模块。自注意力子模块又包含了线性层和其他子模块

ZeRO-Infinity 会递归地向模型的子模块中插入 hooks ，以自动化所需的数据移动。在子模块的 FWD 开始时，这些 hooks 会确保子模块的参数可用于计算，否则它们将执行 allgather 操作，并阻塞到参数可用

类似上文介绍的基于重叠的设计，在子模块的 FWD/BWD 结束时，再次对参数进行分区，并可选择将其转移，减少参数通信时的阻塞

Auto Registration of External Parameters
在理想情况下，一个子模块的参数和梯度只在自己的 FWD/BWD 中使用，这样可以很容易地识别并自动化数据移动，但某些模型架构例外，其中在一个子模块中定义和分配的参数在不同子模块的 FWD/BWD 中都被使用

例如，GPT 等模型在网络的开头和结尾都共享 embedding 层的权重

将上面这种跨模块边界使用的参数称为外部参数，这很难知道在一个子模块的 FWD/BWD 的开始时应该收集哪些参数

于是提出了将外部参数手动注册到 ZeRO-Infinity 中，以便在访问它们的子模块的 FWD/BWD 中进行 gather，注册后，外部参数将像其他参数一样，将被包含在 dynamic prefetcher 中

此外还提供了检测这些场景并自动注册外部参数机制，不用更改任何代码

截获分区参数访问：PyTorch 模块将其数据参数存储在哈希表中。在初始化时，使用子类类型替换哈希表，覆盖数据的访问方法。当访问一个分区参数时，对该参数进行阻塞的 allgather，将其注册为外部参数，然后返回收集到的参数
激活信息检测：子模块的 FWD 可能会返回一个参数，供另一个子模块的 FWD/BWD 使用。例如，Megatron-LM 中线性层 FWD 后返回 bias，在父 Transformer 层模块中使用，再检查每个子模块 FWD 返回的激活值输出中是否包含分区参数，如果包含，则对其进行收集和注册，作为外部参数

Automatic Model Partitioning during Initialization
需要在初始化时划分模型的每个层对应的参数，而不是在整个模型初始化之后再进行划分，以节省峰值内存

提供了一个 Python ZeRO-Infinity 上下文，用于修饰 torch.nn.Module 的 init 方法，在每个模块初始化之后，立即将其分配的参数划分在所在进程组中

只有单独的子模块在完全初始化之后才会被划分，整个模型不会在所有并行进程上复制。一个拥有 5 千亿个参数的模型只需要 1TB 的聚合 CPU 内存就可以在初始化期间完全划分

7 评估

这一部分评估了ZeRO-Infinity，并展示了它在具有数万亿参数的模型中实现了出色的训练效率和可扩展性。我们还展示了ZeRO-Infinity中各种技术对模型规模和性能的影响。

7.1 方法

硬件: 我们在一个由最多512个V100 SXM3 32GB GPU（32个DGX-2节点）组成的集群上进行了实验，节点之间的通信带宽为800 Gbps。

基准: 对于没有模型并行性的实验，我们使用torch的分布式数据并行（DDP [43]）作为基准。对于具有模型并行性的实验，我们使用Megatron-LM [7]。每个实验的基准是3D并行性[13]、ZeRO [11]或ZeRO-Offload [12]中的相关最先进方法。

模型配置: 我们使用基于GPT的Transformer模型。我们将序列长度固定为1024，并根据需要调整隐藏维度和层数，以获得不同参数数量的模型。表1提供了我们在整个评估过程中使用的具体模型配置，请参阅附录获取其他配置。

ZeRO-Infinity卸载策略: 仅当聚合GPU内存不足时，我们才进行卸载操作。首先，我们将优化器状态和梯度卸载到具有足够容量的最快内存中，因为这样可以在最小的通信开销下实现最大的内存节省。接下来，在参数和激活检查点之间，如果只需要将其中一个卸载到CPU内存，我们经验性地选择卸载那个能提供更好性能的部分。当两者都需要卸载时，激活检查点被卸载到CPU，而参数被卸载到具有足够容量的最快内存中。

7.2 模型规模和速度

模型规模: ZeRO-Infinity相比于3D并行性（6500亿参数）可以训练具有32000亿参数的模型，模型规模提升了50倍（图1）。

模型速度: 图5a展示了ZeRO-Infinity在512个GPU上训练多达20万亿参数模型的性能。对于5000亿参数的模型（接近3D并行性在这些资源上可以运行的最大模型），ZeRO-Infinity和3D并行性实现了几乎相同的吞吐量，表明ZeRO-Infinity的训练效率与最先进方法相当。当进一步增加模型规模时，3D并行性会出现内存不足的情况，而ZeRO-Infinity可以训练多达20万亿参数的模型（大约是原来的40倍），吞吐量高达49 TFlops/GPU。这约占该集群的理论峰值性能的40％，但超过了可实现峰值性能（𝑝𝑒𝑎𝑘𝑡𝑝）的70 TFlops的70％。在极限规模下，图5a显示了从10万亿参数（43 TFlops/GPU）到20万亿参数（34 TFlops/GPU）的性能下降。这个下降不是由于NVMe带宽，因为这两个模型都使用了NVMe卸载，而是由于每个GPU的极小批处理大小（表1），这是由于有限的CPU内存来存储激活检查点所导致的。在未来的实现中，可以通过增加CPU内存或将激活检查点卸载到NVMe来改善这一点。

7.3 超线性可扩展性

图5b显示，ZeRO-Infinity在训练1万亿参数模型时从4个节点（64个GPU）到32个节点（512个GPU）实现了超线性可扩展性。这是一个弱可扩展性结果，其中我们保持每个节点的批处理大小不变，并随着节点数量的增加而增加总批处理大小。ZeRO-Infinity通过有效利用总PCIe和NVMe带宽的线性增加来加速参数和优化器状态的卸载，并利用额外节点的CPU计算来进行参数更新，超过了完美的线性可扩展性。此外，ZeRO-Infinity仅使用4个节点就已经实现了超过2.8 petaflops（44 Tflops/GPU）的性能，这表明聚合NVMe带宽即使在适度规模下也足以实现良好的效率。

7.4 大规模模型训练的民主化

图5c展示了在单个节点（16个GPU）上使用ZeRO-Infinity训练10亿到1万亿参数模型的性能，没有任何模型并行性。对于高达1000亿参数的模型，ZeRO-Infinity实现了超过40 TFlops/GPU的出色性能，使得只使用单个DGX-2机箱就能够对GPT-3等模型进行微调成为可能。相比之下，3D并行性无法扩展到超过200亿参数的模型。这些结果展示了ZeRO-Infinity的两个方面：i）可以在单个NVIDIA DGX-2节点上高效地微调具有万亿参数的大型模型，使得没有大型GPU集群访问权限的用户也能够进行训练。ii）易于使用：使用ZeRO-Infinity可以训练这一规模的模型，无需结合模型并行性或流水线并行性，也无需对模型代码进行重构，使得数据科学家能够轻松扩大他们的模型规模。

7.5 系统特性对模型规模的影响

我们展示了不同设备放置策略对模型规模的影响，以及内存中心分块（第5.1.3节）对最大隐藏大小的影响，使用单个DGX-2系统（16个GPU）。

最大模型规模：图6a展示了不同设备放置和分块策略（见表2）对最大模型规模的影响。仅使用数据并行性，由于有限的GPU内存和显著的模型状态冗余，我们仅限于使用14亿参数。当我们引入优化器/梯度分块和使用ZeRO-2和ZeRO-Offload将其卸载到CPU时，我们能够将单个节点的模型规模扩展9倍至130亿参数。在ZeRO-Infinity中，通过将参数状态分块和卸载到CPU，我们几乎可以达到1000亿参数的规模。然而，规模的最终大幅增长来自于将模型状态卸载到NVMe，这使我们的模型规模达到1万亿参数，相对于仅使用数据并行性，模型规模增加了700倍。

最大隐藏大小：我们评估了内存中心分块在存在内存碎片化情况下实现大型隐藏大小的影响。我们使用不同的隐藏大小和分块因子训练了一个单层Transformer模型，以确定可以使用和不使用分块训练的最大隐藏大小。为了保持所有实验中的内存碎片化一致，我们将总GPU内存预先分割为2GB连续块，以便所有大于2GB的内存分配请求将失败。图6b显示，不使用内存中心分块时，可以训练的最大隐藏大小为8K，而使用分块因子为16的内存中心分块，甚至可以训练具有64K的巨大隐藏大小。通过内存中心分块，ZeRO-Infinity通过避免了对模型并行性的需求，极大地简化了DL系统堆栈，使数据科学家能够轻松地训练具有大型隐藏大小的模型。

7.6 系统特性对性能的影响

我们评估了无限卸载引擎（第5节）、带宽中心分块（第6.1节）、重叠设计（第6.2节）和激活检查点卸载（第4.1节）对训练速度的影响。

ZeRO-Infinity vs ZeRO-Offload：图6c展示了使用ZeRO-Infinity和ZeRO-Offload将梯度卸载到CPU内存对8B参数模型的反向传播时间的影响。ZeRO-Infinity利用GPU之间的聚合PCIe带宽将梯度卸载，相对于受单个PCIe带宽限制的ZeRO-Offload，在64个GPU上实现了近2倍的加速。

预取和重叠：图6d展示了对具有64个GPU的8B参数模型开启和关闭通信重叠和预取时的相对吞吐量差异。该图表显示，预取和重叠对于在每个GPU上使用小批量大小时实现良好性能至关重要，而在大批量大小下，其影响减弱。

激活检查点卸载：图6e显示了ZeRO-Infinity中将激活检查点卸载到CPU内存的训练吞吐量降低了最多1.2倍的情况，对于32K和64K的隐藏大小，影响较小，表明可以在不影响大型隐藏大小的情况下将激活检查点卸载到CPU内存。

7.7 测量与预测的性能

本节提供了通过第4节中的带宽分析将我们的测量性能结果与预测结果进行深入比较，并确定ZeRO-Infinity中剩余的性能瓶颈。

图5a显示了在批量大小为3、2、1.25的情况下运行的5T、10T和20T参数模型分别达到了44、43和34 TFlops/GPU的性能。请注意，这三个配置都将参数、梯度和优化器状态卸载到NVMe。基于图2b中显示的每个GPU带宽为1.6GB/s，512个GPU的聚合NVMe带宽为819GB/s。测量得到的跨GPU带宽约为70GB/s（100GB/s峰值的70%），GPU到CPU的测量带宽为每个GPU约3GB/s，如图2b所示。

现在让我们使用这些测量带宽和第4节中的公式来预测这些模型的预期效率。首先，我们定义效率降低为1.0减去效率，并使用公式（2）计算给定批量大小的三个模型的效率降低值。具体而言，它们相对于参数和梯度的效率降低分别为0.25、0.33和0.44；相对于优化器状态的效率降低分别为0.10、0.14和0.21；最后，相对于激活检查点，它们分别约为0.020、0.015和0.01。将效率降低相加，我们可以得到预期的效率分别为0.63、0.515和0.34。请记住，从第4节得知我们的 𝑝𝑒𝑎𝑘𝑡𝑝 范围为62-78 TFlops（平均70 TFlops），而我们的通信重叠可以提高性能高达1.3倍，如图6d所示。基于这些和上述计算出的效率，表4显示了与5T、10T和20T模型的实现性能进行比较的预测性能范围。表4中的比较结果表明，在所有情况下，实现的性能都在预测范围内，验证了第4节中的带宽分析。此外，对于较小的模型，性能更接近较低的范围，而对于较大的模型，性能更接近较高的范围。这是因为随着隐藏大小从62增加到78 TFlops，𝑝𝑒𝑎𝑘𝑡𝑝也增加。最后，虽然很明显，效率随着批量大小的减小而下降，但请注意，在512个GPU规模下，相对于参数和梯度，计算得到的效率降低更大，这是由于受限的GPU-GPU带宽（70 GB/s）是ZeRO-Infinity在批量大小较小时性能瓶颈的主要来源。