ZeRO-Infinity: Breaking the GPU Memory Wall for Extreme Scale Deep Learning

Samyam Rajbhandari, Olatunji Ruwase, Jeff Rasley, Shaden Smith, Yuxiong He {samyamr, olruwase, jerasley, shsmit, yuxhe}@microsoft.com

A1 主要贡献

本文旨在解决大规模深度学习模型训练中日益严峻的“GPU内存墙”问题。过去三年，最大的密集模型参数量增长了1000倍以上，而GPU内存仅增长了5倍，这使得训练超大规模模型（如万亿参数模型）需要海量的GPU资源，且对数据科学家提出了重构模型的巨大负担。

本文的核心问题与研究目标如下：
1. 支持模型规模的未来增长：如何支持模型从百亿级别（如GPT-3）增长到百-万亿级别的下一轮1000倍增长？
2. 提升大模型的可及性：如何让没有大规模GPU集群的数据科学家也能访问和微调现有的大模型？
3. 简化大模型训练：能否通过消除模型重构和多种并行方式的需求，使大模型训练更加简便？

为应对上述挑战，本文提出了ZeRO-Infinity，一个新颖的异构系统技术，其主要贡献和创新点如下：
* 内存和性能特征分析：对大规模模型训练的内存需求（第3节）和带宽需求（第4节）进行了详细的特征分析，为系统设计提供了理论依据。
* ZeRO-Infinity系统技术：这是一个包含五项创新技术的深度学习训练系统，旨在满足大规模、易用且高效的训练需求：
1. Infinity Offload Engine：通过同时利用GPU、CPU和NVMe内存以及GPU和CPU计算，充分发挥现代集群的异构架构优势。
2. Memory-Centric Tiling：一种新颖的GPU内存优化技术，用于处理巨大的算子，避免了模型并行的需要。
3. Bandwidth-Centric Partitioning：一种数据分区策略，旨在利用所有并行设备的聚合内存带宽。
4. Overlap-Centric Design：用于重叠计算和通信，以隐藏数据传输开销。
5. Ease-Inspired Implementation：通过自动化通信和数据分区，避免了模型代码重构。
* 全面的评估验证：
* 前所未有的规模：在32个NVIDIA DGX-2节点（512个V100 GPU）上运行了32万亿参数的模型。
* 卓越的训练效率：在相同硬件上实现了超过25 petaflops的吞吐量。
* 超线性可扩展性：展示了万亿参数模型的超线性扩展能力。
* 可及性和易用性：在单个DGX-2节点上微调了高达万亿参数的模型，且无需模型并行或代码重构。
* 对未来硬件系统设计的启示：讨论了ZeRO-Infinity对未来硬件设计的潜在影响。
* 开源实现：在深度学习优化库DeepSpeed中开源了ZeRO-Infinity的实现。

图1：ZeRO-Infinity可以在32个NVIDIA V100 DGX-2节点（512个GPU）上训练一个32万亿参数的模型，比现有的SOTA技术3D并行大50倍。

A3 背景知识/关键Observation/设计原则

2 背景知识与相关工作

数据、模型、流水线和3D并行：并行化是训练大规模模型的重要策略。对于可以装入设备内存的模型，数据并行（DP）可用于扩展训练到多个设备。当模型无法装入设备内存时，模型并行（MP）【索引7，Megatron-lm: Training multi-billion parameter language models using model parallelism，2019】、【索引17，Mesh-tensorflow: Deep learning for supercomputers，2018】、【索引18，Supporting very large models using automatic dataflow graph partitioning，2019】和流水线并行（PP）【索引7，Megatron-lm: Training multi-billion parameter language models using model parallelism，2019】、【索引8，Pipedream: Fast and efficient pipeline parallel DNN training，2018】、【索引9，Gpipe: Efficient training of giant neural networks using pipeline parallelism，2018】可以分别在垂直和水平方向上切分模型。3D并行【索引14，Megatron-LM: Software repository，2021】、【索引15，DeepSpeed: Extreme-scale model training for everyone，2020】结合了数据、模型和流水线并行，以利用各自的优点，使其能够高效地扩展到万亿参数。虽然3D并行效率很高，但它需要：i) 大量的模型代码重构，以将模型拆分为模型和流水线并行组件；ii) 具有复杂依赖图的模型难以表达为负载均衡的流水线阶段；iii) 模型大小受限于总可用GPU内存。读者可以参考Ben-Nun和Hoefler的综述【索引19，Demystifying parallel and distributed deep learning: An in-depth concurrency analysis，2019】，了解深度学习中并行的全面概述。
ZeRO：零冗余优化器：ZeRO【索引11，ZeRO: Memory Optimizations toward Training Trillion Parameter Models，2020】通过在数据并行进程之间划分三种模型状态（即优化器状态、梯度和参数）而不是复制它们，来消除数据并行进程之间的内存冗余。通过这样做，它在保持计算粒度和通信效率的同时，提高了与经典数据并行相比的内存效率。ZeRO有三个阶段，对应三种模型状态：第一阶段（ZeRO-1）仅划分优化器状态，第二阶段（ZeRO-2）划分优化器状态和梯度，最后阶段（ZeRO-3）划分所有三种模型状态。在ZeRO-3中，模型每一层的参数都由一个唯一的数据并行进程拥有。在训练期间，ZeRO-3通过从业主进程发出广播通信集合操作，确保在算子的前向或后向传递执行之前，所需的参数是可用的。在算子执行之后，ZeRO-3还会移除参数，因为在下一次前向或后向传递之前不再需要它们。此外，在训练的参数更新阶段，ZeRO-3确保每个数据并行进程只更新其拥有的参数所对应的优化器状态。因此，ZeRO-3可以在整个训练过程中保持所有模型状态的分区，除了立即计算所需的参数。
异构训练方法：在几种基于异构CPU内存的训练方法中【索引20，Autotm: Automatic tensor movement in heterogeneous memory systems using integer linear programming，2020】、【索引21，Swapadvisor: Pushing deep learning beyond the gpu memory limit via smart swapping，2020】、【索引22，Layer-centric memory reuse and data migration for extremescale deep learning on many-core architectures，2018】、【索引23，Capuchin: Tensor-based gpu memory management for deep learning，2020】、【索引24，Sentinel: Efficient tensor migration and allocation on heterogeneous memory systems for deep learning，2021】、【索引25，vdnn: Virtualized deep neural networks for scalable, memory-efficient neural network design，2016】、【索引26，Superneurons: Dynamic gpu memory management for training deep neural networks，2018】，ZeRO-Offload【索引12，ZeRO-Offload: Democratizing Billion-Scale Model Training，2021】是多GPU上大规模模型训练的最新技术（SOTA）。ZeRO-Offload构建于ZeRO-2之上，将梯度和优化器状态存储在CPU内存中。当没有足够的GPU设备来存储优化器状态和梯度时，ZeRO-Offload利用CPU内存。然而，它仍然需要参数存储在GPU内存中并在所有设备上复制。因此，ZeRO-Offload的模型规模受限于单个GPU设备内存可以容纳的总参数数量。由于次优的数据分区和有限的PCIe带宽，ZeRO-Offload还需要较大的批量大小才能保持高效。我们用ZeRO-Infinity解决了ZeRO-Offload的这些限制。在基于NVMe的方法方面，Zhao等人【索引27，Distributed hierarchical gpu parameter server for massive scale deep learning ads systems，2020】使用分层参数服务器设计将稀疏参数卸载到SSD，以创建一个大规模深度学习广告系统。相比之下，ZeRO-Infinity被设计为用于训练大规模密集模型的通用深度学习系统。
减少激活内存：激活是在前向传播过程中产生的中间结果，需要保留下来以便在后向传播中计算梯度。已经有多种努力致力于通过压缩【索引28，Gist: Efficient data encoding for deep neural network training，2018】、激活检查点【索引29，Training deep nets with sublinear memory cost，2016】、【索引30，Checkmate: Breaking the memory wall with optimal tensor rematerialization，2019】或实时分析【索引31，Superneurons: Dynamic GPU memory management for training deep neural networks，2018】来减少激活所需的内存。ZeRO-Infinity与激活检查点技术协同工作，以减少激活内存。
Adam优化器和混合精度训练：自适应优化方法【索引32，Adaptive subgradient methods for online learning and stochastic optimization，2011】、【索引33，Adam: A method for stochastic optimization，2015】、【索引34，Scaling SGD batch size to 32k for imagenet training，2017】、【索引35，Reducing BERT pre-training time from 3 days to 76 minutes，2019】对于实现大规模模型有效训练的SOTA性能和准确性至关重要。与SGD相比，Adam【索引33，Adam: A method for stochastic optimization，2015】为每个模型参数和梯度维护细粒度的一阶和二阶统计量，但这会带来显著的内存开销，它是大规模模型训练中最常用的优化器。
混合精度训练细节：大规模模型训练通常采用混合精度训练，其中前向和后向传播以FP16进行，参数更新以FP32进行【索引36，Mixed precision training，2017】。这利用了现代GPU上可用的张量核心单元的性能加速【索引37，NVIDIA Tensor Cores，2018】。

3 内存需求

内存需求分类：本节描述了深度学习训练的内存需求。虽然我们的方法是通用的，但我们将具体分析集中在基于Transformer【索引38，Attention is all you need，2017】的架构上，因为所有超过十亿参数的SOTA模型都遵循该架构。我们的分析假设使用Adam优化器进行混合精度训练，因为这是训练基于Transformer模型的实际标准。训练所需的内存可分为两个部分：i) 模型状态，包括优化器状态、梯度和模型参数；ii) 残差状态，主要指激活内存。为了研究异构资源上的训练，我们还描述了GPU工作内存，即假设模型和残差状态可以成功地从GPU内存中卸载，支持训练必须在GPU上可用的最小内存量。
模型状态内存：模型状态由优化器状态、梯度和参数组成。对于使用Adam优化器的混合精度训练，参数和梯度以FP16存储，而优化器状态由FP32的动量、方差、参数和梯度组成。总共，每个参数需要20字节的内存。基于Transformer的模型中的总参数数量主要取决于隐藏维度（$h_{d}$）和Transformer层的数量（$n_{l}$）。Transformer块中的几乎所有参数都来自块内的四个线性层，大小分别为：($h_{d}$, 3$h_{d}$)，($h_{d}$, $h_{d}$)，($h_{d}$, 4$h_{d}$)和(4$h_{d}$, $h_{d}$)。因此，基于Transformer模型的总参数可以近似为：

公式1

存储模型状态所需的总内存（字节）为：

公式2
模型状态内存需求示例：图2a的第5列显示了通过改变隐藏维度和层数创建的、从1000亿到100万亿参数的类GPT-3 Transformer模型的模型状态内存需求。为了将内存需求置于具体情境中，图2b的第3列显示了单个NVIDIA V100 DGX-2盒子以及DGX2 SuperPOD集群上可用的聚合GPU内存。请注意，仅存储一个1000亿参数模型的模型状态就需要64个GPU。容纳一个万亿参数模型需要超过512个GPU，而一个10万亿参数模型甚至超出了一个拥有1536个GPU的大型集群的能力范围。
残差状态内存：残差状态主要由激活内存组成，它取决于模型架构、批量大小（$b_{s}$）和序列长度（$s_{l}$），并且可能非常大。积极的一面是，激活所需的内存可以通过激活检查点技术【索引29，Training deep nets with sublinear memory cost，2016】显著减少，该技术以0.33倍的额外重计算为代价来换取激活内存。像Turing-NLG 17.2B和GPT-3 175B这样的大型模型都是使用激活检查点进行训练的。存储激活检查点所需的内存估计为：

公式3

(a) (b) 图2：(a) 大规模模型的内存需求。(b) NVIDIA V100 DGX-2集群上的可用内存和可实现带宽（报告的带宽表示所有GPU并行从指定内存读取数据时的每个GPU的带宽）。

图片a

图片b

字节，其中 $k$ 是两个激活检查点之间的Transformer块数，$b_{s} \times s_{l} \times h_{d}$ 是每个Transformer块的输入大小。图2a的第7列显示了在批量大小为32、序列长度为1024，并假设每个Transformer块存储一个激活时存储激活检查点所需的内存。许多现代GPU集群每节点有8-16个GPU，因此我们选择每个GPU的批量大小为2-4，从而得到批量大小为32作为每个节点内激活的保守估计。虽然由此产生的激活检查点比完整的激活集（第6列）小几个数量级，但在超过万亿参数后，对于所考虑的批量大小和序列长度，它们仍然变得太大而无法装入GPU内存。
模型状态工作内存（MSWM）：模型状态工作内存（MSWM）是指在所有模型状态都已卸载到CPU或NVMe后，在模型中最大的单个算子上执行前向或后向传播所需的最小GPU内存量。这大约由该算子的参数和梯度大小决定，因为必须有足够的内存来容纳参数及其梯度以进行后向传播。对于基于Transformer的模型，最大的算子是一个将隐藏状态从$h_{d}$转换为4$h_{d}$的线性层。该线性层的参数和梯度大小（字节）为：

公式4

请注意，MSWM（图2a第8列）在超过1000亿参数后显著增长，需要数GB的连续内存，这可能因缺乏足够的连续内存来满足这些要求而导致训练期间内存不足。像3D并行这样的最新方法通过模型并行来解决这个问题，即将单个算子拆分到多个GPU上。在第5.1.3节中，我们讨论了一种无需模型并行即可解决这些巨大模型状态工作内存的新方法。
激活工作内存（AWM）：激活工作内存（AWM）是在后向传播中，在执行实际的后向传播之前，为重新计算激活所需的内存。这是两个连续激活检查点之间的激活大小。例如，如果我们为每个Transformer块创建一个激活检查点，则内存由每个Transformer块的总激活大小决定。这大约由以下字节数给出：

公式5

图2a的第8列显示，即使$k=1$，AWM在超过10万亿参数后也会变得很大。与仅由单个参数和梯度组成的MSWM不同，AWM由数十个激活组成，只要总AWM可以装入GPU内存，就不会因缺乏连续内存而导致内存问题。

4 带宽需求

带宽对效率的影响：一个关键问题是，卸载到CPU和NVMe内存是否会因其有限的带宽而损害训练效率。本节描述了带宽对训练效率的影响。我们首先定义一个效率指标。假设工作负载执行没有任何计算和通信重叠，我们可以使用峰值计算吞吐量（$P_{peak}$）、数据移动带宽（$B$）及其算术强度（$AI_{T}$）来估计训练效率。工作负载的算术强度（AIT）是总计算量与计算所需数据量之比。它描述了每次数据移动的计算量。更高的AIT意味着对数据移动带宽的要求更低，因为对于加载的每个数据，加速器可以进行更多的计算。效率指标可以推导如下：

公式6

效率可以写成$P_{peak}$、$B$和$AI_{T}$的函数：

公式7

我们将使用这个简单的效率方程来描述训练大规模模型所需的数据移动带宽。但在此之前，我们将首先量化深度学习训练工作负载的$AI_{T}$。

4.1 量化深度学习训练中的AIT

不同状态的AIT：模型状态和激活检查点可以有不同的$AI_{T}$。我们可以通过首先确定深度学习训练每次迭代中的总计算量，然后确定每个模型状态和激活的数据移动量来量化它们。
每次迭代的总计算量：每次迭代的总计算量主要由Transformer的线性层中的计算主导。对于前向传播，这可以近似为参数数量、序列长度和批量大小的函数，由 $2 \times b_{s} \times s_{l} \times \text{Parameters}$ 给出。后向传播的成本大约是前向传播的两倍。此外，激活检查点在后向传播期间需要额外的正向计算作为重计算的一部分。因此，每次迭代的总计算量为：

公式8
参数和梯度的AIT：在前向和后向传播期间，模型参数必须至少两次从源位置加载到GPU寄存器中：i) 在前向传播期间，ii) 在实际的后向传播期间，导致数据移动量为 $2 \times \text{Memory}_{\text{Parameters}}$。在存在激活检查点的情况下，参数可能在后向传播期间为重计算而额外加载一次，增加了 $1 \times \text{Memory}_{\text{Parameters}}$。此外，梯度必须至少一次从GPU寄存器存储到其最终位置，增加了最后的 $1 \times \text{Memory}_{\text{Parameters}}$ 的数据移动。因此，假设参数和梯度存储在相同的最终位置，前向和后向传播期间的总数据移动量将是 $4 \times \text{Memory}_{\text{Parameters}}$，即 $2 \times 4 \times \text{Parameters}$ 字节。每次迭代的总计算量由第4.1节给出。因此，关于参数和梯度的$AI_{T}$是：

公式9
优化器状态的AIT：在优化器步骤中，优化器状态必须至少读取一次，并且优化器状态必须至少写入一次。所以总数据移动量是 $2 \times \text{Memory}_{\text{Optimizer}\_\text{States}}$，大约是 $2 \times 16 \times \text{Parameters}$ 字节。每次迭代的总计算量由第4.1节给出。因此，在完整训练迭代中，关于优化器状态的$AI_{T}$是：

公式10
激活检查点的AIT：在前向传播期间，激活检查点必须保存到其最终位置，并在后向传播期间必须被检索。因此，关于激活检查点的总数据移动量（字节）由 $2 \times \text{Memory}_{\text{Activation}\_\text{Checkpoint}\_\text{in}\_\text{bytes}}$ 给出，根据公式(3)为 $4 \times k/n_{l} \times h_{d} \times b_{s} \times s_{l}$。每次迭代的总计算量由第4.1节给出。所以，关于激活检查点的$AI_{T}$由以下公式给出：

公式11

4.2 带宽需求

带宽需求分析：由于AIT的变化，模型状态和激活检查点对实现良好效率有非常不同的带宽要求。前者仅取决于批量大小和序列长度，而后者仅取决于激活检查点的频率和模型的隐藏维度大小。除了AIT，效率的带宽要求还取决于$P_{peak}$，如公式（6）所示。我们使用$P_{peak}$和$AI_{T}$首先展示效率如何随带宽变化而变化，涉及不同的模型和残差状态，然后讨论这些状态对深度学习训练高效所需的带宽。我们的方法是通用的，可应用于理解任何当前或未来集群的带宽要求。在这里，我们以NVIDIA V100 DGX-2 SuperPOD集群为例。
带宽与效率关系：使用第4.1节的$AI_{T}$表达式和基于公式（6）的效率指标，图3显示了效率与可用带宽在参数和梯度、优化器状态以及激活检查点方面的关系。为了生成这些图表，我们根据第4.1节中推导的表达式计算了$AI_{T}$，针对不同的批量大小、序列长度和模型配置。具体来说，我们使用的序列长度为1024，与GPT-2【索引2，Language models are unsupervised multitask learners，2019】、Megatron-LM【索引7，Megatron-lm: Training multi-billion parameter language models using model parallelism，2019】和Turing-NLG【索引39，turing-nlg: A 17-billion-parameter language model by microsoft】使用的序列长度相同。我们变化的批量大小范围从1到16，以分别捕捉大规模GPU和小规模GPU实验。在大量GPU上运行时使用小的每GPU批量大小，而在相对较少的GPU上训练时使用大的每GPU批量大小，以保持合理的有效批量大小。我们的隐藏大小范围从8K-64K，代表了从数百亿参数到数十万亿参数的模型，如图2a所示。为了确定此分析的$P_{peak}$，我们使用了一种经验方法。我们在一个NVIDIA V100 DGX-2盒子上运行了上述配置的模型，并关闭了所有非GPU通信，以模拟一个几乎无限带宽的场景。根据8K-64K的隐藏大小，实现的性能范围为62-78 TFlops/GPU。我们使用平均值70 TFlops/GPU来代表此分析的$P_{peak}$。

图3：假设加速器的单GPU峰值可实现吞吐量为70 TFlops时，带宽对效率的影响

(a) 参数和梯度带宽

(b) 优化器状态带宽

参数和梯度的带宽：图3a显示，当参数和梯度的带宽超过70 GB/s时，即使是最小的批量大小也能实现超过50%的效率。理论上，在此带宽下，数据移动可以与计算完全重叠，从而实现100%的效率。
优化器状态的带宽：图3b显示，与参数和梯度相比，优化器状态需要近4倍的带宽才能达到50%的效率。此外，优化器状态在前向和后向传播结束时更新，无法与计算重叠。因此，它们需要明显更大的带宽来保持整个深度学习工作负载的高效。例如，要在每个GPU批量大小为2的情况下实现90%的效率，需要近1.5 TB/s的有效带宽，这甚至超过了GPU内存带宽。
激活内存的带宽：图3c也显示，启用激活检查点后，仅2 GB/s的微薄带宽就能在隐藏大小为2K的情况下维持超过50%的效率。一旦隐藏大小增长到8K以上，带宽需求会降至不到1 GB/s。

A2 方法细节

5 ZERO-INFINITY 设计概述

ZeRO-Infinity设计概览：本节我们概述了ZeRO-Infinity的设计选择，这些选择使其能够达到前所未有的模型规模，同时提供卓越的训练效率和易用性。ZeRO-Infinity的鸟瞰图如图4所示，并将在下文讨论。

5.1 为前所未有的规模设计

利用异构内存：现代GPU集群在内存存储方面高度异构。除了GPU内存，它们还拥有CPU内存以及巨大的NVMe存储，其容量比GPU内存大50倍以上，比CPU内存大近20倍（见图2b）。我们开发了ZeRO-Infinity，一个用于深度学习训练的并行系统，它可以通过利用现代GPU集群中的这些异构内存系统来超越GPU内存墙。图1比较了3D并行和ZeRO-Infinity实现的最大模型大小。ZeRO-Infinity在每个NVIDIA V100 DGX-2节点上支持一万亿个参数，比3D并行增加了50倍。

5.1.1 用于模型状态的Infinity offload engine

利用Infinity Offload Engine：ZeRO-Infinity构建于ZeRO-3【索引11，ZeRO: Memory Optimizations toward Training Trillion Parameter Models，2020】之上，后者通过划分所有模型状态来消除内存冗余，如第2节所述。与任何现有的ZeRO系列技术不同，ZeRO-Infinity设计了一个强大的卸载机制，称为infinity offload engine，它可以根据内存需求将所有分区的模型状态卸载到CPU或NVMe内存，或将它们保留在GPU上。从图2a和图2b中可以看出，即使是一个100万亿参数模型所需的模型状态也可以装入一个拥有96个节点（1536个GPU）的DGX-2集群的聚合NVMe内存中。因此，infinity offload engine使ZeRO-Infinity能够容纳数百-万亿参数模型的模型状态。更多细节见第6节。

5.1.2 用于激活的CPU卸载

激活的CPU卸载：除了模型状态，ZeRO-Infinity在必要时可以将激活内存卸载到CPU内存。请注意，一个10万亿参数模型所需的激活检查点（0.76 TB）可以轻松装入DGX-2系统上可用的1.5TB CPU内存中，而一个100万亿参数模型所需的3 TB激活检查点也在下一代硬件的CPU内存可及范围内。因此，通过将激活检查点卸载到CPU内存，ZeRO-Infinity可以容纳数百-万亿参数模型的激活检查点。

图4：ZeRO-Infinity在四个数据并行（DP）rank上训练一个两层模型的快照。图中描绘了第一层反向传播的通信过程。分区后的参数从慢速内存移动到GPU，然后被收集起来形成完整的层。梯度计算完毕后，它们被聚合、重新分区，然后卸载到慢速内存。层用下标表示，DP rank用上标表示。例如，$P_0^{(2)}$ 是第0层参数中由DP rank(2)拥有的部分。

5.1.3 用于工作内存的内存中心切片（Memory-centric tiling）

Memory-Centric Tiling技术：为了减少大型模型深度学习训练的工作内存需求，ZeRO-Infinity引入了一种称为内存中心切片（memory-centric tiling）的新技术，该技术利用ZeRO-3的数据获取和释放模式，通过将大型算子分解为可以顺序执行的较小瓦片（tile）来减少工作内存需求。例如，为了减少大型线性算子的工作内存，ZeRO-Infinity将该算子表示为一个数学上等效的、由原始算子参数的瓦片组成的较小线性算子序列，并按顺序执行它们。当与ZeRO-3结合时，每个瓦片的参数和梯度可以一次一个地获取和释放，从而按瓦片数量的比例减少工作内存。因此，ZeRO-Infinity可以支持任意大小的算子，而无需依赖模型并行来将其装入有限的GPU内存。

5.2 为卓越的训练效率设计

应对效率挑战：将所有模型状态和激活卸载到CPU或NVMe只有在ZeRO-Infinity能够在这种情况下实现高效率时才是可行的。实际上，这是极具挑战性的，因为CPU内存带宽比GPU内存带宽慢一个数量级，而NVMe带宽又比CPU内存带宽慢一个数量级。此外，从GPU读取和写入这些内存甚至更慢（见图2b）。在一个像DGX-2这样的系统上，根据我们在第4节的分析，为了使深度学习训练高效，参数和梯度的带宽必须大于70GB/s，优化器状态的带宽必须大于1.5TB/s，激活检查点的带宽必须大于1-4GB/s。这里我们讨论ZeRO-Infinity如何实现必要的带宽以达到卓越的效率。

5.2.1 参数和梯度的效率

参数和梯度的带宽需求：参数和梯度的数据移动带宽必须大于70GB/s，这接近DGX-2集群上可用的GPU-GPU带宽【索引40，NVIDIA DGX SuperPOD delivers world record supercomputing to any enterprise，2019】。因此，像ZeRO-3【索引11，ZeRO: Memory Optimizations toward Training Trillion Parameter Models，2020】这样的深度学习并行训练解决方案，其中参数在用于前向或后向传播之前从业主GPU广播到其他GPU，只要通信被重叠，就可以高效运行。
异构训练的带宽瓶颈：相反，从单个GPU到CPU内存或NVMe的微薄12 GB/s PCIe带宽（见图2b）根本不足以支持大规模的异构训练。因此，现有的异构解决方案，如ZeRO-Offload，其中参数必须首先从CPU移动到所有者GPU，然后才能广播，需要非常大的每GPU批量大小才能获得在有限带宽下保持高效所需的足够$AI_T$。这带来了两个问题：i) 对于大规模模型，激活内存会变得太大，甚至无法装入CPU内存；ii) 在扩展到数百或数千个GPU时，有效批量大小变得太大，不利于有效收敛。
ZeRO-Infinity的解决方案：ZeRO-Infinity通过两种方式解决这些挑战：i) 带宽中心分区（bandwidth-centric partitioning）：一种新颖的数据映射和并行数据检索策略，用于卸载的参数和梯度，使ZeRO-Infinity能够实现几乎无限的异构内存带宽（详见第6.1节）；ii) 重叠中心设计（overlap centric design），使ZeRO-Infinity不仅能将GPU-GPU通信与计算重叠，还能将NVMe-CPU和CPU-GPU通过PCIe的通信重叠（详见第5.1.3节）。

5.2.2 优化器状态的效率

优化器状态的并行更新：与在前向和后向传播期间顺序消耗和产生的参数和梯度不同，优化器状态可以并行地一次性更新。ZeRO-3和ZeRO-Offload都利用了这一特性，它们分别在所有可用的GPU和CPU上并行存储和更新GPU和CPU内存中的优化器状态。因此，随着GPU或CPU数量的增加，聚合的GPU或CPU内存带宽可以远高于所需的1.5TB/s。
NVMe卸载的挑战与优化：由于ZeRO-Infinity建立在ZeRO-3之上，当将优化器状态卸载到CPU内存时，它也可以利用聚合的GPU和CPU内存带宽以及聚合的CPU计算来进行优化器步骤。然而，对于NVMe卸载，有必要将数据从NVMe分块地带到CPU内存再返回，这些分块必须能装入CPU内存以执行优化器步骤，一次一个分块。因此，优化器步骤受限于NVMe-CPU内存带宽：虽然ZeRO-Infinity可以在多个节点上实现聚合的NVMe带宽，但实现每个节点的近峰值NVMe带宽至关重要，以便在尽可能少的节点和尽可能小的批量大小下支持超过1.5 TB/s的必要带宽。此外，将数据从NVMe移入移出CPU内存，或从CPU内存移入GPU内存的过程，可能会在GPU和CPU中导致CPU内存碎片化，即使还有大量可用内存，也可能导致内存不足。
Infinity Offload Engine的优化：infinity offload engine不仅能实现接近峰值的NVMe带宽，还能让ZeRO-Infinity将NVMe到CPU的读取与CPU到NVMe的写入以及CPU的优化器步骤计算同时重叠，从而使ZeRO-Infinity在少量GPU上以适中的批量大小和在大量GPU上以小的批量大小保持高效。同时，它通过仔细重用用于数据移动的临时缓冲区来最小化内存碎片。我们将在第6节详细讨论infinity offload engine中的优化。

5.2.3 激活的效率

激活卸载的效率：在DGX-2节点上，每个GPU可以通过PCIe以大约3 GB/s的速度并行地向CPU内存读写数据，这使得激活检查点可以被卸载到CPU内存，同时对于大于8K的隐藏尺寸仍能保持超过80%的效率。为了在较小的隐藏尺寸下也能实现高效率，ZeRO-Infinity可以降低激活检查点的频率，并有效地将激活检查点往返CPU内存的通信与GPU上的前向和后向计算重叠。

5.3 为易用性设计

简化模型并行：使用ZeRO-Infinity，数据科学家不再需要像在3D并行中那样，将他们的模型适应多种并行形式。这得益于ZeRO-Infinity中的内存中心切片（memory-centric tiling）技术（在第5.1.3节中讨论），该技术旨在减少大型单个层的GPU内存需求，否则这些层需要模型并行（张量切片）才能装入GPU内存。
消除代码重构：此外，ZeRO-Infinity在PyTorch中的实现方式消除了即使在扩展到万亿参数时也需要手动重构模型代码的需要。这通过一个受易用性启发的实现得以实现，该实现具有两个自动化功能：
- i) 自动化数据移动：在训练期间，在参数需要之前和之后立即收集和划分参数。ZeRO-Infinity通过向PyTorch子模块注入i) 前向/后向预处理钩子（pre forward/backward hooks），触发allgather集合操作以收集其前向/后向传递所需的参数，以及ii) 后向/后向后处理钩子（post forward/backward hooks），触发参数/梯度划分并可选择地将其卸载到CPU或NVMe（详见第7.1节）。
- ii) 初始化期间的自动模型划分：使得无法在单个GPU或CPU内存中容纳的模型仍然可以被初始化，而无需跨数据并行进程手动划分模型。ZeRO-Infinity通过包装所有模块类的构造函数来实现这一点，以便每个子模块的参数在初始化期间创建后立即被划分和卸载。整个模型永远不会在单个数据并行进程上完全实例化（详见第7.2节）。

6 效率优化

效率优化概述：在本节中，我们深入探讨了第5节中介绍的优化措施，这些措施使ZeRO-Infinity能够实现卓越的效率。

6.1 带宽中心分区（Bandwidth-Centric Partitioning）

分区与数据检索策略：ZeRO-Infinity实现了一种新颖的数据映射和检索策略，以解决NVMe和CPU内存带宽的限制。与ZeRO【索引11，ZeRO: Memory Optimizations toward Training Trillion Parameter Models，2020】和ZeRO-Offload【索引12，ZeRO-Offload: Democratizing Billion-Scale Model Training，2021】不同，在这些方法中，每层的参数由单个数据并行进程拥有，并在需要时广播给其余进程；而ZeRO-Infinity将单个参数划分到所有数据并行进程中，并在需要访问参数时使用allgather而不是broadcast。请注意，如果数据位于GPU上，broadcast和allgather通信集合在数据移动量方面的通信成本是相同的。因此，这对于纯GPU训练没有区别。然而，当数据位于NVMe或CPU中时，这是一个颠覆性的改变。
提升异构带宽：在基于广播的方法中，由于每个参数完全由一个数据并行进程拥有，因此在广播发生之前，参数必须首先通过PCIe从其源位置传送到GPU内存。请注意，此过程只能激活一个PCIe，而连接到所有其他GPU的所有PCIe链路都处于空闲状态。相反，通过ZeRO-Infinity中分区参数和基于allgather的方法，所有PCIe链路都并行激活，每个链路带入参数的1/N部分，其中N是数据并行度。因此，NVMe或CPU到GPU的有效通信带宽随着N度的增加而线性增加。
带宽增长示例：例如，使用基于广播的方法，即使在DGX-2盒子上进行16路数据并行，CPU/NVMe到GPU的带宽也保持在PCIe Gen 3的大约12 GB/s。然而，使用基于all-gather的方法，有效可实现带宽分别增加到大约48/25 GB/s（每个GPU 3.0/1.6 GB/s）（见图2b），仅受限于每个DGX-2节点的最大聚合PCIe带宽和最大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的allgather通信，在以小批量运行时也会对效率产生重大影响（图3）。此外，访问NVMe内存需要三个步骤：i) 从NVMe读取数据到CPU内存（nc-transfer），ii) 将数据从CPU内存复制到GPU内存（cg-transfer），iii) 执行allgather以在所有GPU上构建完整的参数（gg-transfer）。这些数据移动的顺序性意味着，如果天真地执行，总通信时间将是这三个数据移动成本的总和，即使每个阶段的数据移动带宽本身是足够的，也会导致效率低下。
重叠引擎设计：为了解决这些问题，ZeRO-Infinity有一个重叠引擎，它不仅将GPU-GPU通信与GPU计算重叠，而且还同时重叠NVMe到CPU和CPU到GPU的通信。该重叠引擎有两个组件：i) 一个动态预取器，用于重叠在参数被前向或后向传递消耗之前重建参数所需的数据移动；ii) 一个通信和卸载重叠机制，用于与后向计算并行执行梯度所需的数据移动。
动态预取器工作原理：ZeRO-Infinity中的动态预取器实时跟踪前向和后向计算，为每次迭代构建一个内部的算子序列映射。在每次迭代期间，预取器跟踪它在算子序列中的位置，并预取未来算子所需的参数。预取器了解三步通信过程，因此可以为一个参数的nc-transfer与另一个参数的cg-transfer和gg-transfer重叠。例如，在执行第i个算子之前，预取器可以分别为第i+3、i+2和i+1个算子所需的参数调用nc、cg和gg-transfer。请注意，所有这些数据移动都可以与第i个算子的执行并行进行。此外，ZeRO-Infinity可以在动态工作流的情况下更新算子序列映射，即使前向和后向传播在迭代之间发生变化，也能进行适当的预取。
梯度通信重叠：同样，在反向传播过程中，ZeRO-Infinity可以将第(i+1)个算子参数梯度的reduce-scatter操作与第i个算子的计算重叠，同时将第(i+2)个算子梯度reduce-scatter得到的分区梯度传输到CPU或NVMe。通过这种强大的以重叠为中心的设计，ZeRO-Infinity即使在使用少量GPU和每个GPU小批量大小进行训练时，也能隐藏大部分数据移动开销。

6.3 Infinity Offload Engine

引擎组成：infinity offload engine由两个主要组件组成：
DeepNVMe：这是infinity offload engine中一个强大的C++ NVMe读/写库，支持批量读/写请求的异步完成，以及用于刷新正在进行的读/写的显式同步请求。对异步的支持使ZeRO-Infinity能够将这些请求与GPU/GPU或GPU/CPU的通信或计算重叠。最重要的是，DeepNVMe能够在NVMe存储设备上实现接近峰值的顺序读写带宽。它通过多项优化实现了这种高性能，包括积极并行化I/O请求（无论是来自单个用户线程还是跨多个用户线程）、智能工作调度、避免数据复制和内存固定（memory pinning）。
固定内存管理层：为确保从（到）NVMe/CPU存储中高性能地读取（或写入）张量，源（或目标）张量必须位于固定内存缓冲区中。然而，固定内存缓冲区是稀缺的系统资源，单个进程对其过度订阅可能会降低整体系统性能或导致系统不稳定。该层通过重用少量（数十GB）的固定内存来管理有限的供应，以卸载整个模型状态（高达数十TB）到CPU或NVMe。内存缓冲区的重用可防止CPU和GPU内存中的内存碎片。该层还为PyTorch张量提供固定内存数据，允许对张量进行原地计算，以便它们可以被写入NVMe而无需任何进一步的复制，从而提高带宽。

7 受易用性启发的实现

实现目标：ZeRO-Infinity是基于PyTorch实现的，旨在无需任何模型代码重构即可使用，类似于PyTorch中的标准数据并行训练。本节详细介绍了实现这样一个系统所面临的一些挑战。

7.1 自动化数据移动

数据移动协调：ZeRO-Infinity必须协调构成模型参数、梯度和优化器状态的张量的移动。当一个张量不处于活动使用状态时，它会在工作节点之间保持分区状态，并可能被卸载到CPU或NVMe内存中。系统必须确保张量在需要使用时及时驻留在GPU内存中，之后再重新分区。
通过钩子实现自动化：PyTorch模型表示为代表神经网络层的模块层次结构。例如，Transformer架构【索引38，Attention is all you need，2017】包含自注意力和前馈网络等子模块。自注意力子模块又由线性变换和其他子模块组成。ZeRO-Infinity递归地向模型的子模块中注入钩子，以自动化所需的数据移动。在子模块的前向传递开始时，这些钩子确保子模块的参数可用于计算，否则它将执行适当的allgather集合操作并阻塞直到参数可用。第6.2节中详述的以重叠为中心的设计对于最小化因参数通信造成的停顿至关重要。在子模块的前向传递结束时，我们再次对参数进行分区，并可选择地将其卸载。后向传递以类似的方式处理。

7.1.1 外部参数的自动注册

外部参数问题：在理想情况下，子模块的参数和梯度仅在其自身的前向和后向传递中被访问，这使得识别和自动化数据移动变得直接，如上节所述。然而，一些模型架构是例外，其中在一个子模块中定义和分配的参数在另一个子模块的前向和后向传播中使用。例如，像GPT【索引41，Improving language understanding by generative pre-training，2018】这样的语言模型在网络的开始和结束处共享嵌入层的权重，以将单词映射到向量，反之亦然。我们将跨模块边界使用的参数称为外部参数。在存在外部参数的情况下，很难知道在一个子模块的前向和后向传递开始时应该收集哪些参数。
手动注册方案：解决这个问题的一种方法是向ZeRO-Infinity注册外部参数，以便为访问它们的子模块的前向和后向传递收集这些参数。注册后，外部参数将被视为与其他参数一样，并被包含在第6.2节描述的预取系统中。我们提供了用于手动注册外部参数的API。
自动注册方案：为了改善用户体验，我们还提供了检测这些情况并自动注册外部参数的机制，这样用户就不必进行任何代码更改：
- 拦截分区参数访问：PyTorch模块将其张量参数存储在哈希表中。在初始化时，我们用一个覆盖张量访问的子类类型替换该哈希表。当访问一个分区参数时，我们对该参数进行阻塞式allgather，将其注册为外部参数，然后返回收集到的参数。
- 激活内省：一个子模块可能会从其前向传递中返回一个参数，供另一个子模块的前向和后向传递使用。例如，Megatron-LM从线性层的前向传递中返回偏置向量，它们被父Transformer层模块消耗。我们检查每个子模块前向传递返回的激活输出中是否有分区参数。如果发现一个，我们收集并将其注册为外部参数。

7.2 初始化期间的自动模型划分

大规模模型初始化挑战：如果模型很大，那么用传统的数据并行方法可能无法完全初始化模型，即在为ZeRO-Infinity分区之前，在每个数据并行进程上复制它。例如，一个5000亿参数的模型在半精度下将占用1TB的内存，因此一个每节点8个GPU的系统仅初始数据并行分配步骤就需要8TB的聚合CPU或GPU内存。这超出了一个节点上可用的GPU或CPU内存。
解决方案：上下文管理器：为了解决这个限制，模型每一层对应的参数必须在初始化时就进行分区，而不是在整个模型初始化之后。为此，我们提供了一个Python ZeRO-Infinity上下文，它装饰了torch.nn.Module的__init__方法，以便在每个模块/子模块下分配的参数在其初始化后立即在数据并行进程组中进行分区。因此，只有单个子模块在被分区之前是完全初始化的，而完整的模型永远不会在所有数据并行进程上被复制。在上面的例子中，5000亿参数的模型因此可以在其初始化期间完全分区，仅需要1TB的聚合CPU内存，而与数据并行进程的总数无关。

A4 实验环境

硬件配置：实验在一个由最多512个V100 SXM3 32GB GPU（32个DGX-2节点）组成的集群上进行，节点间通信带宽为800 Gbps。
软件与基线：
- 对于不使用模型并行的实验，基线为PyTorch的分布式数据并行（DDP）【索引42，Pytorch distributed: Experiences on accelerating data parallel training，2020】。
- 对于使用模型并行的实验，基线为Megatron-LM【索引7，Megatron-lm: Training multi-billion parameter language models using model parallelism，2019】。
- 根据具体实验，将与当前最先进的方法进行比较，包括3D并行【索引13，DeepSpeed: Extreme-scale model training for everyone，2020】、ZeRO【索引11，ZeRO: Memory Optimizations toward Training Trillion Parameter Models，2020】或ZeRO-Offload【索引12，ZeRO-Offload: Democratizing Billion-Scale Model Training，2021】。
模型与数据集：
- 使用类GPT的基于Transformer的模型。
- 固定序列长度为1024，通过改变隐藏维度和层数来获得不同参数数量的模型。
- 具体的模型配置见表1及附录A。

A4 实验结果

模型规模与速度

模型规模：ZeRO-Infinity能够训练超过32万亿参数的模型，而当前最先进的3D并行技术只能支持约6500亿参数，实现了50倍的模型规模提升（见图1）。
模型速度：在512个GPU上，ZeRO-Infinity在训练高达20万亿参数的模型时表现出色。对于5000亿参数的模型（接近3D并行在该资源下的极限），ZeRO-Infinity的吞吐量与3D并行几乎相同，表明其训练效率与SOTA相当。当模型规模进一步增大时，3D并行因内存不足而失败，而ZeRO-Infinity能以高达49 TFlops/GPU的优异吞吐量训练20万亿参数的模型（规模大40倍）。在极端规模下（10T和20T模型），性能有所下降，但这并非由NVMe带宽引起，而是因为CPU内存有限导致激活检查点存储受限，使得每个GPU的批量大小极小。此问题可通过增加CPU内存或未来将激活检查点卸载到NVMe来改善（见图5a）。

超线性可扩展性

实验内容与结果：在训练1万亿参数模型时，ZeRO-Infinity在从4个节点（64 GPU）扩展到32个节点（512 GPU）的过程中表现出超线性扩展性。这是一个弱扩展实验，即保持每节点的批量大小不变，随节点数增加而增加总批量大小。
结论分析：ZeRO-Infinity之所以能超越完美的线性扩展，是因为它有效利用了聚合PCIe和NVMe带宽的线性增长来加速参数和优化器状态的卸载，并利用新增节点的CPU计算能力进行参数更新。此外，仅用4个节点就实现了超过2.8 petaflops（44 Tflops/GPU）的性能，证明了即使在适度规模下，聚合的NVMe带宽也足以实现高效率（见图5b）。

普及大规模模型训练

实验内容与结果：在单个节点（16 GPU）上，ZeRO-Infinity无需任何模型并行即可训练100亿到1万亿参数的模型。对于高达1000亿参数的模型，ZeRO-Infinity实现了超过40 TFlops/GPU的优异性能，使得在单个DGX-2上微调GPT-3等模型成为可能。相比之下，3D并行无法扩展到超过200亿参数的模型（见图5c）。
结论分析：这些结果证明了ZeRO-Infinity的两个方面：
1. 可及性：使得没有大型GPU集群的用户也能在单个NVIDIA DGX-2节点上微调高达万亿参数的大模型。
2. 易用性：如此规模的模型可以使用ZeRO-Infinity进行训练，而无需结合模型或流水线并行，也无需重构模型代码，极大地简化了数据科学家扩展模型的过程。

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

不同设备放置策略对模型规模的影响：实验展示了在单个DGX-2系统（16 GPU）上，不同设备放置和分区策略对最大模型规模的影响。仅使用数据并行，模型大小受限于14亿参数。通过ZeRO-2和ZeRO-Offload将优化器/梯度分区并卸载到CPU，规模扩大了9倍至130亿参数。进一步使用ZeRO-Infinity将参数状态也分区并卸载到CPU，模型规模接近1000亿。最终，将模型状态卸载到NVMe使得模型规模达到1万亿，相比单纯数据并行提升了700倍（见图6a和表2）。
内存中心切片对隐藏层大小的影响：为了评估内存中心切片（memory-centric tiling）在存在内存碎片时对支持大隐藏维度的影响，实验在一个单层Transformer模型上进行，并通过预先将GPU内存碎片化为2GB的连续块来模拟内存碎片。结果显示，不使用切片技术时，可训练的最大隐藏维度为8K。而使用16倍的切片因子，可以训练高达64K的巨大隐藏维度。该技术避免了模型并行的需要，简化了DL系统栈（见图6b）。

系统特性对性能的影响

ZeRO-Infinity vs. ZeRO-Offload：在80亿参数模型的反向传播时间对比中，ZeRO-Infinity利用聚合的PCIe带宽卸载梯度，在64个GPU上相比受限于单个PCIe带宽的ZeRO-Offload实现了近2倍的加速（见图6c）。
预取和重叠的影响：对于80亿参数模型在64个GPU上的实验，开启通信重叠和预取功能后，吞吐量有显著提升。结果表明，预取和重叠对于在小批量/GPU下实现高性能至关重要，其影响随批量增大而减小（见图6d）。
激活检查点卸载的影响：将激活检查点卸载到CPU会使小隐藏尺寸模型的训练吞吐量降低最多1.2倍。但对于32K和64K的大隐藏尺寸，影响微乎其微，证明了在不牺牲效率的情况下将大隐藏尺寸模型的激活检查点卸载到CPU是可行的（见图6e）。

A5 结论

本文提出了ZeRO-Infinity，一种新颖的异构系统技术，它通过利用GPU、CPU和NVMe内存，实现了前所未有的模型规模，同时保证了训练的可及性、易用性和卓越效率。这项工作改变了我们对大规模模型训练中内存的看法，证明了不再需要将所有训练数据都放在昂贵且有限的HBM2等高速内存中。ZeRO-Infinity展示了通过并行利用多个设备上廉价、低速但容量巨大的CPU或NVMe内存，可以获得高效训练所需的聚合带宽，从而突破GPU内存墙。

展望未来，随着GPU和其他加速器的计算能力变得更强，高效训练所需的聚合带宽也将增加。如表3所示，即使加速器的计算能力比NVIDIA V100 GPU强10倍，在一个512个加速器的集群上，ZeRO-Infinity也仅需每个加速器与慢速内存之间有30 GB/s的带宽即可保持高效。这在当前技术下，例如通过NVLink连接加速器和慢速内存，是完全可以实现的。

显然，有了ZeRO-Infinity，加速器设备内存不再是模型规模或训练效率的瓶颈。然而，在合理时间内训练数十万亿或数百万亿参数的模型仍需要计算能力的巨大飞跃，并且在未来设备上高效运行也需要设备间带宽的相应提升。我们希望，随着设备内存不再是限制因素，ZeRO-Infinity将激励未来在超强计算能力加速器和超级计算集群方面进行更多以计算和设备间带宽为中心的创新，以支持模型规模的下一个1000倍增长及其带来的进步。

表3：在一个拥有512个加速器设备（其可实现计算能力分别为NVIDIA V100 GPU的10倍和100倍）的集群上，ZeRO-Infinity保持高效所需的带宽（bw）要求。

A6 附录

图6(a)的模型配置

表4：图6(a)的模型配置
图6(b)的模型配置

表5：图6(b)的模型配置
图6(c)的模型配置

表6：图6(c)的模型配置
图6(d)的模型配置

表7：图6(d)的模型配置
图6(e)的模型配置

表8：图6(e)的模型配置