1. 论文基本信息

1.1. 标题

ZeRO-Infinity: Breaking the GPU Memory Wall for Extreme Scale Deep Learning (ZeRO-Infinity：打破 GPU 内存墙以实现超大规模深度学习)

1.2. 作者

Samyam Rajbhandari, Olatunji Ruwase, Jeff Rasley, Shaden Smith, Yuxiong He 他们的研究背景和隶属机构是微软 (Microsoft)，研究领域集中于深度学习系统优化和分布式训练。

1.3. 发表期刊/会议

该论文发布于 arXiv 预印本平台，日期为 2021-04-16T02:22:12.000Z。通常，arXiv 上的预印本论文会在经过同行评审后，被正式接受并发表在顶级的学术会议或期刊上，例如高性能计算领域的 SC (Supercomputing)、HPCA (High Performance Computer Architecture) 或 PPoPP (Principles and Practice of Parallel Programming) 等。DeepSpeed 系列的工作通常在这些顶级会议上发表，具有很高的影响力。

1.4. 发表年份

2021年。

1.5. 摘要

在过去三年中，最大的密集深度学习模型规模增长了 1000 倍，达到数千亿参数，而 GPU 内存仅增长了 5 倍（从 16 GB 到 80 GB）。因此，模型规模的增长主要依赖于系统创新，这些创新允许大型模型适应多个 GPU 的聚合内存。然而，我们正接近 GPU 内存墙 (GPU memory wall)。训练一个万亿参数的模型仅需要 800 块 NVIDIA V100 GPU，而这样的集群对于大多数数据科学家来说是遥不可及的。此外，训练如此规模的模型需要复杂的并行技术组合，这给数据科学家带来了重大的模型代码重构负担。

本文提出了 ZeRO-Infinity，一种新颖的 异构系统技术 (heterogeneous system technology)，它利用 GPU、CPU 和 NVMe 内存 (NVMe memory)，在不 重构模型代码 (model code refactoring) 的情况下，以有限的资源实现了前所未有的模型规模。同时，它实现了出色的训练吞吐量和可扩展性，不受限于 CPU 或 NVMe 的带宽瓶颈。ZeRO-Infinity 可以在当前一代 GPU 集群上训练数万亿甚至数百兆亿参数的模型。它可以在单个 NVIDIA DGX-2 节点上 微调 (fine-tune) 万亿参数模型，使得大型模型更易于访问。在训练吞吐量和可扩展性方面，它在 512 块 NVIDIA V100 GPU 上维持超过 25 petaflops（峰值性能的 40%），并展示了 超线性可扩展性 (super linear scalability)。ZeRO-Infinity 的开源实现通过 DeepSpeed 提供，DeepSpeed 是一个使分布式训练变得简单、高效和有效的深度学习优化库。

1.6. 原文链接

https://arxiv.org/abs/2104.07857v1

1.7. PDF 链接

https://arxiv.org/pdf/2104.07857v1.pdf

2. 整体概括

2.1. 研究背景与动机

2.1.1. 核心问题

论文试图解决的核心问题是GPU 内存墙 (GPU memory wall)。随着深度学习模型规模的爆炸式增长（在三年内增长 1000x），单块 GPU 的内存增长（仅 5x）远远滞后。这导致训练和微调 超大规模模型 (extreme scale models) 需要极大量的 GPU 资源，甚至数百个 GPU 才能容纳万亿参数模型，这对于大多数数据科学家来说是不可承受的。

2.1.2. 问题的重要性与挑战

• 资源限制： 现有的 3D 并行 (3D parallelism) 技术虽然能够扩展到万亿参数，但仍需要大量 GPU 才能容纳模型，例如训练一个万亿参数模型需要 320 块 NVIDIA A100 GPU。对于未来的百兆亿参数模型，甚至需要超过 6000 块 GPU。这使得大型模型训练成为只有少数拥有大规模集群的机构才能进行的活动。
• 可访问性挑战： 即使是相对较小的任务，如 微调 (fine-tuning) 数百亿参数模型，也需要多节点 GPU 集群（例如微调 GPT-3 需要超过 8 个 DGX-2 节点），这使得许多研究人员和公司无法进行。
• 复杂性与易用性： 现有的 大规模模型训练 (large model training) 方法（如 3D 并行 (3D parallelism)）需要数据科学家对模型代码进行大量的 重构 (refactoring)，以实现模型并行和流水线并行，这增加了使用的复杂性和门槛，也限制了支持的模型类型。

2.1.3. 论文的切入点与创新思路

论文的切入点在于利用现代计算系统中的异构内存体系。现有方法主要依赖于聚合 GPU 内存，但现代集群同时拥有大量的 CPU 内存和 NVMe 存储，这些内存虽然速度较慢，但容量巨大且成本低廉。论文的创新思路是通过智能地将模型状态和激活值卸载到这些更慢但容量更大的内存介质上，并结合高效的通信和计算重叠机制，来打破 GPU 内存限制，同时保持训练效率和易用性。

2.2. 核心贡献/主要发现

论文的主要贡献和发现可以总结为以下几点：

1. 前所未有的模型规模 (Unprecedented Model Scale)：
   • 提出了 ZeRO-Infinity，一种新颖的 异构系统技术 (heterogeneous system technology)，它能同时利用 GPU、CPU 和 NVMe 内存来存储模型状态。
   • 引入了 内存中心平铺 (memory-centric tiling) 技术，以处理单个无法放入 GPU 内存的超大算子，从而避免了 模型并行 (model parallelism)。
   • 实验证明，ZeRO-Infinity 可以在 32 个 NVIDIA DGX-2 节点（512 块 V100 GPU）上训练高达 32 万亿参数的模型，比现有的 3D 并行 (3D parallelism) 技术高出 50x。
2. 卓越的训练效率 (Excellent Training Efficiency)：
   • 开发了 带宽中心分区 (bandwidth-centric partitioning) 策略，以聚合所有设备的内存带宽，解决了 CPU 和 NVMe 的带宽限制问题。
   • 采用了 重叠中心设计 (overlap-centric design)，将通信和计算（包括 GPU-GPU、NVMe-CPU 和 CPU-GPU 传输）高效重叠，显著提高了效率。
   • 在 512 块 NVIDIA V100 GPU 上，实现了超过 25 petaflops 的吞吐量（峰值性能的 40%），并展示了在训练万亿参数模型时的 超线性可扩展性 (superlinear scalability)。
3. 易用性和可访问性 (Ease of Use and Accessibility)：
   • 通过 内存中心平铺 (memory-centric tiling) 避免了对大型算子进行 模型并行 (model parallelism)，从而消除了复杂模型重构的需求。
   • 引入了 易用性启发实现 (ease-inspired implementation)，自动化了数据移动和模型初始化过程，使得数据科学家无需进行手动模型代码重构，即可训练万亿参数模型。
   • ZeRO-Infinity 可以在单个 NVIDIA DGX-2 节点上微调万亿参数模型，极大地降低了大型模型训练和微调的门槛，使得更多研究人员和企业能够访问这些模型。
4. 开源实现： ZeRO-Infinity 作为 DeepSpeed 优化库的一部分开源，促进了其在深度学习社区中的广泛采用。

3. 预备知识与相关工作

3.1. 基础概念

为了理解 ZeRO-Infinity，初学者需要了解以下核心概念：

• 深度学习 (Deep Learning, DL)：机器学习的一个分支，使用多层神经网络（深度神经网络 (Deep Neural Networks)）从数据中学习复杂的模式。
• GPU (Graphics Processing Unit)：图形处理器，最初用于图形渲染，但因其并行计算能力而广泛用于加速深度学习训练。
• CPU (Central Processing Unit)：中央处理器，计算机的“大脑”，执行通用计算任务，通常具有较低的并行度但更高的单核性能和更大的内存容量。
• NVMe (Non-Volatile Memory Express)：一种高速存储接口协议，用于固态硬盘 (SSD)，提供比传统 SATA 硬盘更快的读写速度，可作为大容量、相对低速的存储介质。
• 模型参数 (Model Parameters)：深度学习模型中的可学习权重和偏置，它们在训练过程中被优化以使模型更好地拟合数据。例如，一个 Transformer 模型可能包含数十亿甚至数万亿个参数。
• 梯度 (Gradients)：在深度学习训练中，梯度 (gradients) 表示损失函数相对于模型参数的变化率。优化器使用梯度来更新参数，从而最小化损失。
• 优化器状态 (Optimizer States)：许多高级优化器（如 Adam）除了存储参数和梯度外，还需要维护额外的状态信息。例如，Adam 优化器会为每个参数维护一阶矩（momentum）和二阶矩（variance）的估计值，这些都是 优化器状态 (optimizer states) 的一部分。这些状态通常比参数和梯度本身占用更多的内存。
• 激活值 (Activations)：神经网络前向传播过程中，每一层输出的中间结果。在 反向传播 (backward propagation) 计算梯度时，这些 激活值 (activations) 需要被保留下来。如果模型很深或 批大小 (batch size) 很大，激活值 (activations) 可能会占用大量 GPU 内存。
• 混合精度训练 (Mixed Precision Training)：一种训练技巧，结合使用 16 位浮点数 (FP16) 和 32 位浮点数 (FP32)。通常，模型参数和激活值以 FP16 存储和计算，而 优化器状态 (optimizer states) 和某些关键计算（如 梯度累积 (gradient accumulation)）使用 FP32，以利用 GPU 的 Tensor Core 加速 FP16 运算，同时保持数值稳定性并减少内存使用。
• Adam 优化器 (Adam Optimizer)：一种流行的 自适应学习率优化器 (adaptive learning rate optimizer)，通过计算每个参数的 自适应学习率 (adaptive learning rate) 来提高训练效率和收敛速度。它需要维护额外的 优化器状态 (optimizer states)，因此内存开销较大。
• Transformer 架构 (Transformer Architecture)：一种流行的神经网络架构，广泛应用于 自然语言处理 (Natural Language Processing, NLP) 领域，也是许多 大型语言模型 (Large Language Models, LLMs) 的基础。它主要由 多头自注意力机制 (multi-head self-attention mechanisms) 和 前馈网络 (feed-forward networks) 组成。
• 微调 (Fine-tuning)：在预训练了一个大型模型（通常在海量通用数据上）之后，将其在特定任务的小型数据集上进一步训练，以适应特定应用的过程。这比从头开始训练（预训练 (pretraining)）要便宜得多。

3.2. 前人工作

论文提及了以下与分布式深度学习和内存优化相关的关键前人工作：

• 数据并行 (Data Parallelism, DP)：最常见的并行策略。在每个设备（如 GPU）上复制完整的模型，然后将 批数据 (batch data) 分割成小批次并分配给每个设备。每个设备独立计算 梯度 (gradients)，然后通过 AllReduce 操作将 梯度 (gradients) 聚合 (aggregate) 起来，并更新模型参数。这种方法的优点是易于实现且计算效率高，但要求模型本身能够完全放入单个设备内存。
• 模型并行 (Model Parallelism, MP)：当模型过大无法放入单个设备内存时使用。它将模型的不同层或一个层内的不同部分放置在不同的设备上。张量切片 (tensor-slicing) 是一种常见的 模型并行 (model parallelism) 形式，将大型 线性层 (linear layer) 或 注意力层 (attention layer) 的权重张量分割到多个设备上。这需要复杂的模型重构。
• 流水线并行 (Pipeline Parallelism, PP)：将模型的层序列分割成不同的阶段，每个阶段由一个或一组设备负责计算。数据以流水线方式流经这些阶段，从而减少了内存需求并提高了吞吐量。
• 3D 并行 (3D Parallelism)：结合了 数据并行 (DP)、模型并行 (MP) 和 流水线并行 (PP)。这是当前训练 万亿参数模型 (trillion-parameter models) 的 最先进技术 (state-of-the-art)。例如，DeepSpeed 的 3D 并行 (3D parallelism) 实现可以扩展到万亿参数模型，但它需要大量 GPU 的聚合内存，并且要求数据科学家进行大量的模型代码重构。
• ZeRO: Zero Redundancy Optimizer (零冗余优化器)：
  • 一种内存优化技术，通过在数据并行进程之间 分区 (partitioning) 模型状态而不是 复制 (replicating) 它们来消除内存冗余。
  • ZeRO-1：仅分区 优化器状态 (optimizer states)。
  • ZeRO-2：分区 优化器状态 (optimizer states) 和 梯度 (gradients)。
  • ZeRO-3：分区所有三个模型状态（参数 (parameters)、梯度 (gradients) 和 优化器状态 (optimizer states)）。在训练过程中，ZeRO-3 动态地在需要时 收集 (gather) 参数到 GPU，并在使用后 释放 (release) 它们。每个数据并行进程只拥有模型参数的一部分。
• ZeRO-Offload (ZeRO 卸载)：基于 ZeRO-2，将 梯度 (gradients) 和 优化器状态 (optimizer states) 卸载 (offload) 到 CPU 内存。它解决了 优化器状态 (optimizer states) 的内存限制，但模型参数仍需存储在 GPU 内存中并在所有设备上复制，因此模型规模仍然受限于单个 GPU 的内存容量。ZeRO-Offload 还需要较大的 批大小 (batch size) 才能保持效率，因为 PCIe 带宽 (PCIe bandwidth) 有限。
• 异构训练方法 (Heterogeneous Training Approaches)：除了 ZeRO-Offload，还有其他基于 CPU 内存的训练方法 [20-26] 和基于 NVMe 的方法 [27]，但这些方法通常针对特定场景（如稀疏参数 [27]）或未充分利用异构存储的聚合带宽。
• 减少激活内存 (Reducing Activation Memory)：激活值 (activations) 的内存消耗可以通过 压缩 (compression) [28]、激活检查点 (activation checkpointing) [29, 30] 或 实时分析 (live analysis) [31] 等技术减少。ZeRO-Infinity 与 激活检查点 (activation checkpointing) 协同工作。

3.3. 技术演进

深度学习训练技术从最初的单 GPU 训练，逐渐演进到分布式训练以处理更大的模型和数据集。

1. 单 GPU 瓶颈：最早的模型可以完全放入单块 GPU 内存中。
2. 数据并行 (DP)：当数据集变大时，使用 DP 扩展训练，但模型规模仍受限于单个 GPU 内存。
3. 模型并行 (MP) 和流水线并行 (PP)：为了训练无法放入单个 GPU 的超大模型，MP 和 PP 被引入，将模型分割到多个 GPU 上。但它们通常需要复杂的模型重构。
4. ZeRO 系列优化器：ZeRO 进一步通过 分区 (partitioning) 模型状态来消除内存冗余，尤其是 ZeRO-3，它将所有模型状态分区，显著提升了单卡的模型容量。
5. 异构内存利用：ZeRO-Offload 将 梯度 (gradients) 和 优化器状态 (optimizer states) 卸载到 CPU，是利用异构内存的第一步。然而，参数仍需在 GPU 内存中，且 PCIe 带宽限制效率。
6. 3D 并行 (3D Parallelism)：结合 DP、MP 和 PP，是当前 最先进的 (state-of-the-art) 万亿参数模型 (trillion-parameter model) 训练方法，但仍面临 GPU 内存墙和代码重构的挑战。
7. ZeRO-Infinity：本文的工作，通过全面利用 GPU、CPU 和 NVMe 异构内存，并结合带宽和重叠优化，旨在突破现有的 GPU 内存墙，同时提升易用性，是异构内存利用和分布式优化的最新进展。

3.4. 差异化分析

ZeRO-Infinity 与现有主要方法的核心区别和创新点在于：

• 与 3D 并行 (3D Parallelism) 的区别：

  • 内存利用： 3D 并行 (3D Parallelism) 主要依赖于聚合 GPU 内存来容纳模型，当模型过大时，即使所有 GPU 内存加起来也可能不足。而 ZeRO-Infinity 突破性地利用了 GPU、CPU 和 NVMe 存储的三层异构内存，容量潜力远超 3D 并行 (3D Parallelism)。
  • 模型规模： ZeRO-Infinity 能够训练的模型规模比 3D 并行 (3D Parallelism) 大 50 倍（例如，从 650B 参数到 32T 参数）。
  • 易用性与代码重构： 3D 并行 (3D Parallelism) 需要数据科学家进行大量的模型代码重构，以实现 模型并行 (model parallelism) 和 流水线并行 (pipeline parallelism)。ZeRO-Infinity 通过 内存中心平铺 (memory-centric tiling) 和 自动化数据移动 (automated data movement) 机制，消除了对 模型并行 (model parallelism) 和代码重构的需求，极大地提升了易用性。

• 与 ZeRO-Offload 的区别：

  • 内存层级： ZeRO-Offload 主要将 优化器状态 (optimizer states) 和 梯度 (gradients) 卸载到 CPU 内存，但模型参数 (model parameters) 仍然完全存在于 GPU 内存中。这意味着 ZeRO-Offload 的模型规模上限仍然受限于单个 GPU 的内存容量。

  • 参数卸载： ZeRO-Infinity 不仅可以卸载 优化器状态 (optimizer states) 和 梯度 (gradients)，更重要的是，它能够将模型参数 (model parameters) 卸载到 CPU 甚至 NVMe 内存，从而彻底解除了模型规模对 GPU 内存的依赖。

  • 带宽利用： ZeRO-Offload 受限于单个 PCIe 链路的带宽，需要较大的 批大小 (batch size) 才能保持效率。ZeRO-Infinity 引入了 带宽中心分区 (bandwidth-centric partitioning)，通过 allgather 操作聚合所有 PCIe 链路的带宽，实现了远超 ZeRO-Offload 的有效带宽，即使在小 批大小 (batch size) 下也能保持高效。

    总的来说，ZeRO-Infinity 在模型规模、训练效率和易用性上都超越了前人工作，其核心创新在于对异构内存的全面和高效利用，以及为实现此目标而设计的创新技术。

4. 方法论

本部分将详细拆解 ZeRO-Infinity 的技术方案，从内存需求分析到具体的设计和优化。

4.1. 内存需求分析

在深入 ZeRO-Infinity 的设计之前，了解深度学习训练的内存需求至关重要。论文详细分析了 Transformer 模型的内存构成，将其分为 模型状态 (Model states) 和 残差状态 (Residual states)（主要是 激活值 (activations)），并引入了 GPU 工作内存 (GPU working memory) 的概念。分析假设使用 Adam 优化器的 混合精度训练 (mixed precision training)。

4.1.1. 模型状态内存 (Memory for Model States)

模型状态 (Model states) 包括 优化器状态 (optimizer states)、梯度 (gradients) 和 模型参数 (model parameters)。 对于使用 Adam 优化器的 混合精度训练 (mixed precision training)：

• 参数 (parameters) 和 梯度 (gradients) 通常以 FP16 (2 字节) 存储。

• 优化器状态 (optimizer states) 包含 FP32 格式的 momentum、variance、参数副本 (parameter copy) 和 梯度副本 (gradient copy)。具体来说，Adam 优化器通常为每个 FP16 参数维护 3 个 FP32 值（momentum、variance 和 FP32 参数的完整副本），加上 FP16 梯度 (gradients) 的 FP32 副本。这通常意味着每个参数需要 20 字节的内存 ($2(FP16)+2(FP16)+4(FP32)+4(FP32)+4(FP32)+4(FP32)$ -> $2+2+4+4+4+4$ = 20 字节。实际上，通常是 FP16 参数(2字节) + FP16 梯度(2字节) + FP32 优化器状态(12字节，包含 momentum 和 variance，以及参数的 FP32 副本) = 16字节。而论文中提到 20 字节，是更保守的估计或包含额外的 FP32 梯度副本。这里我们严格按照论文的 20 字节来计算)。

  Transformer 模型中的总参数量主要取决于 隐藏维度 (hidden dimension, hd) 和 Transformer 层数 (nl)。几乎所有参数都来自每个 Transformer 块内的四个 线性层 (linear layers)，其大小分别为：$(hd, 3hd)$、(hd, hd)、$(hd, 4hd)$ 和 $(4hd, hd)$。 因此，Transformer 模型中的总参数量可近似为： $12 \times nl \times hd^2$ 存储 模型状态 (model states) 所需的总内存（以字节为单位）为： $240 \times nl \times hd^2$ 符号解释：

• nl: Transformer 模型的层数。

• hd: Transformer 模型的 隐藏维度 (hidden dimension)。

• $12 \times nl \times hd^2$: 近似的总参数量。

• $240 \times nl \times hd^2$: 存储模型状态所需的总内存（字节），假设每个参数需要 20 字节。

4.1.2. 残差状态内存 (Memory for Residual States)

残差状态 (Residual states) 主要指 激活值 (activation memory)。它取决于模型架构、批大小 (batch size, bsz) 和 序列长度 (sequence length, seq)，可能会非常大。通过 激活检查点 (activation checkpointing) [29] 可以显著减少 激活值 (activations) 内存，这以 0.33x 的额外重计算为代价。

存储 激活检查点 (activation checkpoints) 所需的内存估计为： $$2 \times bsz \times seq \times hd \times nl / ci$$ 符号解释：

• bsz: 批大小 (batch size)。
• seq: 序列长度 (sequence length)。
• hd: 隐藏维度 (hidden dimension)。
• nl: Transformer 模型的层数。
• ci: 两个 激活检查点 (activation checkpoints) 之间的 Transformer 块数量。
• $bsz \times seq \times hd$: 每个 Transformer 块输入的尺寸。

4.1.3. GPU 工作内存 (GPU Working Memory)

GPU 工作内存 (GPU working memory) 是在所有 模型状态 (model states) 已 卸载 (offload) 到 CPU 或 NVMe 之后，GPU 上必须可用的最低内存量，以支持模型中最大单个算子的 前向传播 (forward propagation) 或 反向传播 (backward propagation)。

• 模型状态工作内存 (Model State Working Memory, MSWM)：大约等于模型中最大算子的参数和梯度大小。对于 Transformer 模型，最大的算子是将 隐藏状态 (hidden states) 从 hd 转换到 4hd 的 线性层 (linear layer)。该 线性层 (linear layer) 的参数和梯度大小（以字节为单位）为： $4 \times hd \times 4hd$ 这表示 MSWM 可能需要连续的数 GB 内存，在高并发或内存碎片化时可能导致 内存不足 (out of memory)。

• 激活工作内存 (Activation Working Memory, AWM)：是 反向传播 (backward propagation) 中重新计算 激活值 (activations) 所需的内存。它是两个连续 激活检查点 (activation checkpoints) 之间的 激活值 (activations) 总大小。例如，如果每个 Transformer 块创建一个 激活检查点 (activation checkpoint)，则内存由每个 Transformer 块的总 激活值 (activations) 大小给出，大约为： $$bsz \times seq \times ci \times (16 \times hd + 2 \times attn\_heads \times seq)$$ 符号解释：

  • bsz: 批大小 (batch size)。

  • seq: 序列长度 (sequence length)。

  • ci: 两个 激活检查点 (activation checkpoints) 之间的 Transformer 块数量。

  • hd: 隐藏维度 (hidden dimension)。

  • $attn\_heads$: 注意力头 (attention heads) 的数量。

    以下是原文 Figure 2a 和 Figure 2b 展示的内存需求和可用带宽： 以下是原文 Figure 2a 的结果：

以下是原文 Figure 2b 的结果：

4.2. 带宽需求分析

将数据 卸载 (offload) 到 CPU 或 NVMe 内存的一个关键问题是，它们有限的带宽是否会损害训练效率。本节分析了带宽对训练效率的影响。

4.2.1. 效率指标 (Efficiency Metric)

假设工作负载执行时没有计算和通信重叠，效率可以使用峰值计算吞吐量 ($peak_{tp}$)、数据移动带宽 (bw) 及其 算术强度 (arithmetic intensity, ait) 来估计。 算术强度 (ait) 是总计算量与所需数据量之比，描述了每次数据移动的计算量。ait 越高，对数据移动带宽的要求越低。

计算时间 ($compute\_time$)、通信时间 ($communication\_time$) 和效率 (efficiency) 的推导如下： $$compute\_time = \frac{total\_computation}{peak_{tp}}$$ $$ait = \frac{total\_computation}{total\_data\_movement}$$ $$communication\_time = \frac{total\_data\_movement}{bw} = \frac{total\_computation}{ait \times bw}$$ $$efficiency = \frac{compute\_time}{compute\_time + communication\_time}$$ 效率可以表示为 $peak_{tp}$、bw 和 ait 的函数： $$efficiency = \frac{ait \times bw}{ait \times bw + peak_{tp}}$$ 符号解释：

• $compute\_time$: 计算所需的时间。
• $total\_computation$: 总计算量。
• $peak_{tp}$: 峰值计算吞吐量。
• ait: 算术强度 (arithmetic intensity)。
• $total\_data\_movement$: 总数据移动量。
• bw: 数据移动带宽。
• $communication\_time$: 通信所需的时间。
• efficiency: 训练效率。

4.2.2. 量化深度学习训练中的算术强度 (Quantifying AIT in DL training)

为了量化 ait，需要识别每次迭代的总计算量和不同模型状态及 激活检查点 (activation checkpoints) 的数据移动量。

• 每次迭代的总计算量 (Total Computation per Iteration)： 主要由 Transformer 模型的 线性层 (linear layers) 计算决定。 前向传播 (forward propagation) 的计算量可近似为 $2 \times bsz \times seq \times parameters$。 反向传播 (backward propagation) 的计算量大约是 前向传播 (forward propagation) 的两倍。 激活检查点 (activation checkpointing) 需要额外的 前向计算 (forward computation) 作为 反计算 (recomputation) 的一部分。 因此，每次迭代的总计算量为： $$computation\_per\_iter = 2 \times 4 \times bsz \times seq \times parameters = 2 \times 4 \times 12 \times bsz \times seq \times nl \times hd^2$$ 符号解释：

  • bsz: 批大小 (batch size)。
  • seq: 序列长度 (sequence length)。
  • parameters: 模型参数总数。
  • nl: Transformer 模型的层数。
  • hd: 隐藏维度 (hidden dimension)。
  • $2 \times 4$: 考虑了前向传播、反向传播（两倍）和激活检查点的重计算。
  • $12 \times bsz \times seq \times nl \times hd^2$: 将参数量公式代入。

• 关于参数和梯度的算术强度 (AIT w.r.t. Parameters and Gradients)： 在 前向传播 (forward propagation) 和 反向传播 (backward propagation) 期间，模型参数至少需要从源位置加载到 GPU 寄存器两次（一次用于前向，一次用于实际反向）。如果存在 激活检查点 (activation checkpointing)，参数可能需要额外加载一次用于反向传播中的重计算。梯度 (gradients) 必须至少存储一次。 因此，假设参数和梯度存储在同一位置，前向传播 (forward propagation) 和 反向传播 (backward propagation) 期间的总数据移动量为 $4 \times parameters$（即 $2 \times 4 \times parameters$ 字节）。 所以，关于参数和梯度的 ait 为： $seq \times bsz$ 符号解释：

  • seq: 序列长度 (sequence length)。
  • bsz: 批大小 (batch size)。

• 关于优化器状态的算术强度 (AIT w.r.t. Optimizer States)： 在优化器步骤中，优化器状态 (optimizer states) 必须至少读取一次并写入一次。因此，总数据移动量为 $2 \times optimizer\_states$，约为 $2 \times 16 \times parameters$ 字节。 所以，一个完整训练迭代中关于 优化器状态 (optimizer states) 的 ait 为： $seq \times bsz / 4$ 符号解释：

  • seq: 序列长度 (sequence length)。
  • bsz: 批大小 (batch size)。
  • /4: 反映了优化器状态相对于参数有更高的内存占用（如 16 字节 vs 4 字节），导致 ait 降低。

• 关于激活检查点的算术强度 (AIT w.r.t. Activation Checkpoints)： 在 前向传播 (forward propagation) 期间，激活检查点 (activation checkpoints) 必须保存到最终位置；在 反向传播 (backward propagation) 期间，必须检索它们。因此，激活检查点 (activation checkpoints) 的总数据移动量由 $2 \times total\_activation\_checkpoints\_in\_bytes$ 给出。 所以，关于 激活检查点 (activation checkpoints) 的 ait 为： $24 \times hd \times ci$ 符号解释：

  • hd: 隐藏维度 (hidden dimension)。
  • ci: 两个 激活检查点 (activation checkpoints) 之间的 Transformer 块数量。

4.2.3. 带宽需求 (Bandwidth Requirements)

由于 算术强度 (ait) 的变化，不同模型状态和 激活检查点 (activation checkpoints) 需要非常不同的带宽才能实现高效率。

• 参数和梯度 (Parameter and Gradients)：图 2 (原文 Figure 3a) 显示，对于参数和梯度，超过 $70 \mathrm{GB/s}$ 的带宽可以实现超过 50% 的效率，即使对于最小 批大小 (batch size)。理论上，在此带宽下，数据移动可以与计算完全重叠，达到 100% 效率。
• 优化器状态 (Optimizer States)：图 3 (原文 Figure 3b) 显示，优化器状态 (optimizer states) 需要近 4 倍的带宽才能达到 50% 的效率。由于 优化器状态 (optimizer states) 在 前向传播 (forward propagation) 和 反向传播 (backward propagation) 结束后更新，无法与计算重叠，因此需要显著更高的带宽才能保持整体 DL 工作负载高效。例如，在每个 GPU 批大小 (batch size) 为 2 的情况下，要达到 90% 的效率，需要近 $1.5 \mathrm{TB/s}$ 的有效带宽。
• 激活内存 (Activation Memory)：图 4 (原文 Figure 3c) 显示，启用 激活检查点 (activation checkpointing) 后，即使是 $2 \mathrm{GB/s}$ 的微薄带宽，也能在 隐藏维度 (hidden size) 为 2K 时维持超过 50% 的效率。当 隐藏维度 (hidden size) 超过 8K 时，带宽需求甚至降至 $1 \mathrm{GB/s}$ 以下。

4.3. ZeRO-Infinity 设计概览 (ZeRO-Infinity Design Overview)

ZeRO-Infinity 旨在通过利用现代 GPU 集群中的异构内存（GPU、CPU 和 NVMe）来突破 GPU 内存墙，同时提供卓越的训练效率和易用性。

4.3.1. 前所未有的规模设计 (Design for Unprecedented Scale)

• Infinity Offload Engine (无限卸载引擎) 用于模型状态：

  • ZeRO-Infinity 基于 ZeRO-3 构建，它分区所有 模型状态 (model states) 以消除内存冗余。
  • 无限卸载引擎 (infinity offload engine) 能够将所有分区后的 模型状态 (model states) 卸载 (offload) 到 CPU 或 NVMe 内存，或根据内存需求将其保留在 GPU 上。
  • 这使得 ZeRO-Infinity 能够容纳数万亿参数模型的 模型状态 (model states)，因为即使是 100 万亿参数模型的 模型状态 (model states) 也可以放入 96 节点 DGX-2 集群的聚合 NVMe 内存中。

• CPU 卸载 (CPU Offload) 用于激活值：

  • ZeRO-Infinity 还可以将 激活内存 (activation memory) 卸载到 CPU 内存。
  • 例如，一个 10 万亿参数模型所需的 激活检查点 (activation checkpoints) （0.76 TB）可以轻松放入 DGX-2 系统上可用的 1.5TB CPU 内存中。这使得 ZeRO-Infinity 能够容纳数百兆亿参数模型的 激活检查点 (activation checkpoints)。

• 内存中心平铺 (Memory-centric Tiling) 用于工作内存：

  • ZeRO-Infinity 引入了一种新颖的技术，即 内存中心平铺 (memory-centric tiling)，通过将大型算子分解为可以顺序执行的较小 切片 (tiles)，从而减少深度学习训练的 工作内存 (working memory) 需求。

  • 例如，对于大型 线性算子 (linear operator)，ZeRO-Infinity 将其表示为数学上等效的较小 线性算子 (linear operators) 序列。结合 ZeRO-3，每个 切片 (tile) 的参数和梯度可以一次一个地获取和释放，按 切片 (tiles) 数量的比例减少 工作内存 (working memory)。

  • 这使得 ZeRO-Infinity 能够支持任意大小的算子，而无需依赖 模型并行 (model parallelism) 来适应有限的 GPU 内存。

    以下是原文 Figure 4，展示了 ZeRO-Infinity 的设计概览：

    该图像是论文中的示意图，展示了ZeRO-Infinity在四个数据并行（DP）rank上训练两层模型时的参数和梯度在GPU、慢速内存（CPU+NVMe）及网络中的迁移与通信过程。

Figure 4: A snapshot of ZeRO-Infinity training a model with two layers on four data parallel (DP) ranks. Communication for the backward pass of the first layer is depicted. Partitioned parameters are moved from slow memory to GPU and then collected to form the full layer. After gradients are computed, they are aggregated, repartitoned, and then offloaded to slow memory. Layers are denoted with subscripts and DP ranks are denoted with superscripts. For example, p(2) is the portion of layer O's parameters owned by $G P U ^ { ( 2 ) }$ .

4.3.2. 高效训练设计 (Design for Excellent Training Efficiency)

尽管 CPU 和 NVMe 的带宽有限，ZeRO-Infinity 仍能实现高效率。

• 参数和梯度的效率 (Efficiency w.r.t Parameter and Gradients)：

  • 参数和梯度的带宽需要超过 $70 \mathrm{GB/s}$。ZeRO-3 中参数从拥有者 GPU 广播到其他 GPU 的方式可以高效运行，只要通信可以重叠。
  • ZeRO-Infinity 通过两种方式解决了传统异构解决方案（如 ZeRO-Offload）中 PCIe 带宽不足的问题：
    1. 带宽中心分区 (Bandwidth-centric Partitioning)：一种新颖的数据映射和并行数据检索策略，用于 卸载 (offloaded) 的参数和梯度，使得 ZeRO-Infinity 能够实现几乎无限的异构内存带宽。它将单个参数在所有数据并行进程中分区，并使用 allgather 而非 broadcast 来访问参数，从而并行激活所有 PCIe 链路。
    2. 重叠中心设计 (Overlap Centric Design)：不仅将 GPU-GPU 通信与计算重叠，还同时重叠 NVMe-CPU 和 CPU-GPU 通信（通过 PCIe）。

• 优化器状态的效率 (Efficiency w.r.t Optimizer States)：

  • 优化器状态 (optimizer states) 可以并行更新。ZeRO-Infinity 利用聚合的 GPU 和 CPU 内存带宽以及聚合的 CPU 计算能力来进行优化器步骤。
  • 对于 NVMe 卸载 (offload)，数据分块从 NVMe 传输到 CPU 内存，然后进行优化器步骤。这受限于 NVMe-CPU 内存带宽。
  • 无限卸载引擎 (infinity offload engine) 能够实现接近峰值的 NVMe 带宽，并重叠 NVMe 到 CPU 的读取、CPU 到 NVMe 的写入以及优化器步骤的 CPU 计算，以保持效率。它通过仔细重用临时缓冲区来最小化内存碎片。

• 激活值的效率 (Efficiency w.r.t Activations)：

  • 在 DGX-2 节点上，每个 GPU 可以以大约 $3 \mathrm{GB/s}$ 的速度通过 PCIe 并行读写 CPU 内存，从而可以将 激活检查点 (activation checkpoints) 卸载到 CPU 内存，同时对于较大的 隐藏维度 (hidden sizes) 保持超过 80% 的效率。
  • ZeRO-Infinity 可以降低 激活检查点 (activation checkpoints) 的频率，并有效地重叠 激活检查点 (activation checkpoints) 与 前向传播 (forward) 和 反向传播 (backward) 计算之间的通信。

4.3.3. 易用性设计 (Design for Ease of Use)

ZeRO-Infinity 旨在无需任何模型代码重构即可使用，类似于 PyTorch 中的标准 数据并行 (data-parallel) 训练。

• 内存中心平铺 (Memory-centric Tiling)：如前所述，它减少了大型单个层的 GPU 内存需求，消除了对 模型并行 (model parallelism) 的依赖，简化了模型训练。

• 易用性启发实现 (Ease Inspired Implementation)：通过两个自动化特性实现：

  1. 自动化数据移动 (Automated Data Movement)：在 PyTorch 子模块中注入 前向/后向钩子 (pre/post forward/backward hooks)。这些钩子在需要参数时触发 allgather 集合操作来收集参数，并在使用后 分区 (partition) 参数并可选地 卸载 (offload) 到 CPU 或 NVMe。
  2. 自动化模型分区初始化 (Automated Model Partitioning during Initialization)：通过包装 torch.nn.Module 的 __init__ 方法，使得每个子模块的参数在初始化后立即在数据并行进程组中 分区 (partition) 和 卸载 (offload)。这样，整个模型永远不会在单个数据并行进程上完全实例化，从而解决了超大模型初始化时内存不足的问题。

4.4. 效率优化 (Efficiency Optimizations)

本节深入探讨了 ZeRO-Infinity 提高效率的优化细节。

4.4.1. 带宽中心分区 (Bandwidth-Centric Partitioning)

• 核心思想： ZeRO-Infinity 采用一种新颖的数据映射和检索策略，以解决 NVMe 和 CPU 内存带宽限制。
• 与 ZeRO/ZeRO-Offload 的区别： 传统的 ZeRO 和 ZeRO-Offload 中，每个层的参数由单个数据并行进程拥有，并在需要时 广播 (broadcast) 给其他进程。这意味着从慢速内存到 GPU 的传输只能通过单个 PCIe 链路进行。
• ZeRO-Infinity 的方法： ZeRO-Infinity 将单个参数在所有数据并行进程中分区。当需要访问参数时，它使用 allgather 操作而不是 broadcast。
  • 这样，所有的 PCIe 链路都可以并行活动，每个链路传输参数的 $1/dp$ 部分（其中 dp 是数据并行度）。
  • 结果是，从 NVMe 或 CPU 到 GPU 的有效通信带宽会随着 dp 的增加而线性增加。
  • 例如，在 DGX-2 上，allgather 方法可以将有效带宽从单 PCIe 的 $12 \mathrm{GB/s}$ 提高到多 GPU 并行下的 $48/25 \mathrm{GB/s}$ (分别为 CPU/NVMe)，并且随着节点数量的增加而线性增长，最终提供“几乎无限”的异构内存带宽。

4.4.2. 重叠中心设计 (Overlap Centric Design)

为了解决单 GPU 或单节点设置中的带宽瓶颈，以及复杂的 NVMe 访问过程（NVMe 到 CPU，CPU 到 GPU，然后 allgather），ZeRO-Infinity 采用了一个 重叠引擎 (overlap engine)。

• 动态预取器 (Dynamic Prefetcher)：
  • 在运行时跟踪 前向传播 (forward) 和 反向传播 (backward) 计算，构建操作符序列的内部映射。
  • 在每次迭代中，预取器会跟踪当前操作符序列的位置，并预取 (prefetches) 未来操作符所需的参数。
  • 预取器了解三步通信过程，因此可以重叠不同参数的 NVMe-CPU 传输（nc-transfer）、CPU-GPU 传输（cg-transfer）和 GPU allgather 传输（gg-transfer）。例如，在执行第 $i$ 个操作符之前，它可以调用第 $i+3$ 个操作符的 nc-transfer，第 $i+2$ 个操作符的 cg-transfer，以及第 $i+1$ 个操作符的 gg-transfer。所有这些数据移动都可以与第 $i$ 个操作符的执行并行进行。
  • 预取器可以更新操作符序列映射以适应动态工作流。
• 通信与卸载重叠机制 (Communication and Offload Overlapping Mechanism)：
  • 在 反向传播 (backward pass) 中，ZeRO-Infinity 可以将第 $i+1$ 个操作符参数的 梯度 (gradients) reduce-scatter 与第 $i$ 个操作符的计算重叠。
  • 同时，它还可以将第 $i+2$ 个操作符 梯度 (gradients) 的 reduce-scatter 后的分区 梯度 (gradients) 传输到 CPU 或 NVMe。 这种强大的 重叠中心设计 (overlap centric design) 即使在少量 GPU 和小 批大小 (batch size) 的情况下也能隐藏大量数据移动开销。

4.4.3. 无限卸载引擎 (Infinity Offload Engine)

无限卸载引擎 (infinity offload engine) 由两个主要组件构成：

• DeepNVMe：
  • 一个功能强大的 C++ NVMe 读写库。
  • 支持 批量读写请求 (bulk read/write requests) 用于 异步完成 (asynchronous completion)，以及 显式同步请求 (explicit synchronization requests) 以刷新正在进行的读写操作。
  • 异步 (asynchrony) 支持使得 ZeRO-Infinity 可以将这些请求与 GPU/GPU 或 GPU/CPU 通信或计算重叠。
  • 通过一系列优化（包括 I/O 请求的并行化、智能工作调度、避免数据复制和内存固定），DeepNVMe 能够实现接近峰值的 NVMe 顺序读写带宽。
• 固定内存管理层 (Pinned Memory Management Layer)：
  • 为了确保从 NVMe/CPU 存储到（或从）GPU 的高性能张量读写，源（或目标）张量必须驻留在 固定内存缓冲区 (pinned memory buffers) 中。
  • 固定内存缓冲区 (pinned memory buffers) 是稀缺的系统资源，过度使用可能导致系统性能下降或不稳定。
  • 此层通过重用少量（数十 GB）的 固定内存 (pinned memory) 来管理整个模型状态（高达数十 TB）的 卸载 (offload)。
  • 内存缓冲区的重用防止了 CPU 和 GPU 内存碎片化。
  • 它还为 PyTorch 张量提供了 固定内存数据 (pinned memory data)，允许 张量 (tensors) 的 原地计算 (inplace computation)，这样它们就可以直接写入 NVMe，无需额外的复制以提高带宽。

4.5. 易用性启发实现 (Ease Inspired Implementation)

ZeRO-Infinity 在 PyTorch 上实现，旨在无需任何模型代码重构即可使用。

4.5.1. 自动化数据移动 (Automating Data Movement)

ZeRO-Infinity 必须协调模型参数、梯度和优化器状态的移动。当张量不活跃时，它会保持在各 工作节点 (workers) 之间 分区 (partitioned)，并可能 卸载 (offloaded) 到 CPU 或 NVMe 内存。系统必须确保张量在需要时及时驻留在 GPU 内存中，并在之后重新 分区 (partition)。

• 钩子注入 (Hooks Injection)：ZeRO-Infinity 递归地将 钩子 (hooks) 注入模型的子模块中，以自动化所需的数据移动。
  • 在子模块 前向传播 (forward pass) 开始时，这些 钩子 (hooks) 确保子模块的参数可用，否则会执行相应的 allgather 集合操作并阻塞直到参数可用。
  • 在子模块 前向传播 (forward pass) 结束时，再次 分区 (partition) 参数并可选地 卸载 (offload) 它们。反向传播 (backward pass) 以类似的方式处理。
• 外部参数的自动注册 (Auto Registration of External Parameters)：
  • 理想情况下，子模块的参数和梯度只在其自己的 前向传播 (forward) 和 反向传播 (backward) 中访问。但有些模型（如 GPT）会在不同子模块中共享参数（如 嵌入层 (embedding layer)）。这些被称为 外部参数 (external parameters)。
  • 手动注册 (Manual Registration)：提供 API 允许用户手动注册 外部参数 (external parameters)。
  • 自动注册 (Automatic Registration)：为了提升用户体验，提供两种机制自动检测和注册 外部参数 (external parameters)：
    1. 拦截分区参数访问 (Intercepting partitioned parameter accesses)：在初始化时，替换 PyTorch 模块中存储张量参数的哈希表，使用一个重写了张量访问的子类。当访问分区参数时，会进行一个阻塞的 allgather 操作，将其注册为 外部参数 (external parameter)，然后返回 收集 (gathered) 到的参数。
    2. 激活内省 (Activation introspection)：检查每个子模块 前向传播 (forward pass) 返回的 激活输出 (activation outputs) 中是否有分区参数。如果发现，则将其 收集 (collect) 并注册为 外部参数 (external parameter)。

4.5.2. 初始化时的自动模型分区 (Automatic Model Partitioning during Initialization)

如果模型过大，传统的 数据并行 (data parallel) 方法在初始化时，需要先在每个数据并行进程上复制完整的模型，这可能导致内存不足。

• ZeRO-Infinity 的方法：提供一个 Python ZeRO-Infinity 上下文，它装饰 torch.nn.Module 的 __init__ 方法。
  • 这样，每个模块/子模块中分配的参数在其初始化后立即在数据并行进程组中 分区 (partitioned)。
  • 因此，只有单个子模块在被分区之前完全初始化，而完整的模型永远不会在所有数据并行进程上复制。
  • 例如，一个 5000 亿参数的模型在初始化时只需要 1 TB 的聚合 CPU 内存，而不管数据并行进程的总数。

5. 实验设置

5.1. 数据集

实验使用了 GPT 类的 Transformer 模型。通过调整 隐藏维度 (hidden dimension) 和层数 (number of layers) 来获得不同参数量的模型。序列长度 (sequence length) 固定为 1024。

以下是原文 Table 1 展示的实验配置：

符号解释：

• # nodes: 使用的节点数量。

• # params: 模型参数总量（B 代表 billions，T 代表 trillions，K 代表 1024）。

• hidden dim: 隐藏维度 (hidden dimension) 大小。

• # layers: 模型层数。

• batch/GPU: 每个 GPU 的 批大小 (batch size)。

• mp: 模型并行 (model parallelism) 度。1 表示没有 模型并行 (model parallelism)。

• fp16 param: FP16 参数存储的位置（GPU, CPU, NVMe）。

• Opt State: 优化器状态 (optimizer states) 存储的位置（GPU, CPU, NVMe）。

  这些配置涵盖了从数十亿到数十万亿参数的模型，以及单节点到多节点集群的训练场景，旨在全面评估 ZeRO-Infinity 在不同规模下的性能。

5.2. 评估指标

论文主要使用以下指标来评估 ZeRO-Infinity 的性能：

1. 模型规模 (Model Scale)：能够成功训练的最大模型参数量。这是 ZeRO-Infinity 旨在突破的核心限制。
2. 吞吐量 (Throughput)：通常以 TFlops/GPU（每 GPU 的 Tera 浮点运算每秒）衡量。
   • 概念定义 (Conceptual Definition)：TFlops/GPU 量化了每个 GPU 在单位时间内执行的浮点运算次数。它直接反映了模型训练的速度和计算资源的利用效率。更高的 TFlops/GPU 意味着更快的训练速度。
   • 数学公式 (Mathematical Formula)： TFlops/GPU = $\frac{\text{Total Floating Point Operations}}{\text{Training Time} \times \text{Number of GPUs}} \times 10^{-12}$
   • 符号解释 (Symbol Explanation)：
     • $\text{Total Floating Point Operations}$: 在训练一个 批次 (batch) 数据时，模型执行的总浮点运算次数。
     • $\text{Training Time}$: 训练一个 批次 (batch) 所需的时间。
     • $\text{Number of GPUs}$: 参与训练的 GPU 总数。
     • $10^{-12}$: 将 GFlops 转换为 TFlops（$1 \text{ Tera} = 10^{12}$）。
3. 效率 (Efficiency)：通常以峰值性能的百分比衡量，或者使用第 4.2.1 节中定义的 ait、bw 和 peak_tp 计算。
   • 概念定义 (Conceptual Definition)：效率 (efficiency) 衡量了实际观察到的计算性能与理论峰值计算性能之间的比率。它反映了系统在特定工作负载下，计算资源（如 GPU）被有效利用的程度，以及通信开销对计算速度的影响。
   • 数学公式 (Mathematical Formula)： $$efficiency = \frac{ait \times bw}{ait \times bw + peak_{tp}}$$
   • 符号解释 (Symbol Explanation)：
     • ait: 算术强度 (arithmetic intensity)，表示每次数据移动的计算量。
     • bw: 数据移动带宽。
     • $peak_{tp}$: 峰值计算吞吐量。
4. 可扩展性 (Scalability)：评估随着 GPU 数量的增加，吞吐量如何增长。论文特别提到了 超线性可扩展性 (superlinear scalability)，这意味着吞吐量的增长速度甚至超过了 GPU 数量的线性增长。
5. 内存占用 (Memory Footprint)：虽然没有直接作为评估指标列表列出，但 ZeRO-Infinity 的核心优势在于显著降低了 GPU 内存占用，使得模型可以适应更小的资源。通过比较不同配置下最大可训练模型规模来间接体现。
6. 易用性 (Ease-of-Use)：通过验证是否无需模型代码重构即可训练大规模模型，以及是否能在单节点上进行 微调 (fine-tuning) 来体现。

5.3. 对比基线

论文将 ZeRO-Infinity 与以下 最先进的 (state-of-the-art) 方法进行了比较：

• torch's distributed data parallel (DDP)：对于不涉及 模型并行 (model parallelism) 的实验，DDP 是 PyTorch 中标准的 数据并行 (data parallel) 实现。它作为最基础的分布式训练基线。
• Megatron-LM：一个由 NVIDIA 开发的，用于训练数十亿参数 大型语言模型 (large language models) 的框架，主要通过 模型并行 (model parallelism)（张量切片 (tensor-slicing)）实现。在涉及 模型并行 (model parallelism) 的实验中，Megatron-LM 是一个重要的基线。
• 3D Parallelism：结合了 数据并行 (data parallelism)、模型并行 (model parallelism) 和 流水线并行 (pipeline parallelism) 的 最先进的 (state-of-the-art) 大规模模型训练技术。它是 ZeRO-Infinity 主要超越的对象，特别是在模型规模方面。
• ZeRO (Zero Redundancy Optimizer)：特别是 ZeRO-3，它通过 分区 (partitioning) 模型所有状态来消除内存冗余。ZeRO-Infinity 是在其基础上进一步扩展的。
• ZeRO-Offload：ZeRO-Infinity 的直接前身，它将 梯度 (gradients) 和 优化器状态 (optimizer states) 卸载到 CPU 内存。ZeRO-Infinity 通过引入参数卸载和更高效的带宽利用来改进 ZeRO-Offload。

6. 实验结果与分析

本节将详细分析 ZeRO-Infinity 的实验结果，展示其在模型规模、速度和可访问性方面的优势，并探讨不同系统特性对其性能的影响。

6.1. 模型规模与速度 (Model Size and Speed)

6.1.1. 模型规模

ZeRO-Infinity 实现了前所未有的模型规模。

• 与 3D 并行 (3D parallelism) 的对比：ZeRO-Infinity 可以在 32 个 NVIDIA DGX-2 节点（512 块 GPU）上训练超过 32 万亿参数的模型，而 3D 并行 (3D parallelism) 的 最先进技术 (state-of-the-art) 只能达到约 6500 亿参数。这代表了模型规模上 50 倍的飞跃（如 Figure 1 所示）。

  以下是原文 Figure 1 的结果：

  该图像是一个柱状图，展示了图1中ZeRO-Infinity在不同数量的NVIDIA V100 DGX-2节点（GPU）上支持的模型参数规模（单位：万亿）。图中用3D并行性与ZeRO-Infinity的实际测量和预测参数规模进行了对比，最大显示128T参数的预测值，显著超过3D并行性。

Figure 1: ZeRO-Infinity can train a model with 32 trillion parameters on 32 NVIDIA V100 DGX-2 nodes (512 GPUs), $\mathbf { 50 x }$ larger than 3D parallelism, the existing state-of-the-art.

6.1.2. 模型速度

ZeRO-Infinity 在支持巨大模型规模的同时，也保持了出色的训练吞吐量。

• 与 3D 并行 (3D parallelism) 的对比：对于 5000 亿参数模型（这是 3D 并行 (3D parallelism) 在相同资源下能运行的最大模型规模），ZeRO-Infinity 和 3D 并行 (3D parallelism) 实现了几乎相同的吞吐量，这表明 ZeRO-Infinity 在效率上与 最先进的 (state-of-the-art) 技术相当。

• 超大规模模型的性能：当模型规模进一步增大时，3D 并行 (3D parallelism) 因内存不足而无法运行，而 ZeRO-Infinity 能够训练高达 20 万亿参数的模型，并实现高达 49 TFlops/GPU 的卓越吞吐量。

• 极端规模下的吞吐量下降：在 10 万亿（43 TFlops/GPU）和 20 万亿（34 TFlops/GPU）参数的极端规模下，观察到性能下降。这并非由于 NVMe 带宽，而是因为 20 万亿参数模型需要极小的每个 GPU 批大小 (batch size)（如 Table 1 所示），这限制了 算术强度 (arithmetic intensity)。另一个原因是 CPU 内存有限，无法存储大量的 激活检查点 (activation checkpoints)，导致需要更小的 批大小 (batch size)。论文指出，未来可以通过增加 CPU 内存或将 激活检查点 (activation checkpoints) 卸载到 NVMe 来改善。

  以下是原文 Figure 5a 的结果：

  该图像是图表，展示了ZeRO-Infinity在训练超大规模模型上的效率与扩展性，分别对比了3D并行、ZeRO-Offload与ZeRO-Infinity在不同模型规模及GPU数量下的吞吐率表现，体现了其超线性扩展能力和在单节点上训练1万亿参数模型的优势。

Figure 5: Efficiency and scalability of ZeRO-Infinity for training multi-trillion parameter models.

6.2. 超线性可扩展性 (Superlinear Scalability)

ZeRO-Infinity 在训练万亿参数模型时表现出 超线性可扩展性 (super-linear scalability)。

• 实验设置：在从 4 个节点（64 块 GPU）扩展到 32 个节点（512 块 GPU）训练 1 万亿参数模型时，ZeRO-Infinity 展示了 超线性可扩展性 (super-linear scalability)。这是一种 弱扩展性 (weak scaling) 结果，其中每个节点的 批大小 (batch size) 保持不变，总 批大小 (batch size) 随节点数量增加而增加。

• 原因分析：

  • 聚合带宽的有效利用：随着节点数量的增加，聚合的 PCIe 和 NVMe 带宽线性增加。ZeRO-Infinity 有效利用了这些额外的带宽来加速参数和优化器状态的 卸载 (offloading)。
  • 聚合 CPU 计算：额外的节点提供了更多的 CPU 计算资源，用于参数更新和优化器步骤。

• 效率表现：即使在只有 4 个节点时，ZeRO-Infinity 也达到了超过 2.8 petaflops（44 TFlops/GPU）的性能，这表明聚合的 NVMe 带宽足以在适度规模下实现良好的效率。

  以下是原文 Figure 5b 的结果：

  该图像是图表，展示了ZeRO-Infinity在训练超大规模模型上的效率与扩展性，分别对比了3D并行、ZeRO-Offload与ZeRO-Infinity在不同模型规模及GPU数量下的吞吐率表现，体现了其超线性扩展能力和在单节点上训练1万亿参数模型的优势。

Figure 5: Efficiency and scalability of ZeRO-Infinity for training multi-trillion parameter models.

6.3. 普及大型模型训练 (Democratizing Large Model Training)

ZeRO-Infinity 显著降低了大型模型训练和 微调 (fine-tuning) 的门槛。

• 单节点能力：在单个 DGX-2 节点（16 块 GPU）上，ZeRO-Infinity 能够训练从 100 亿到 1 万亿参数的模型，而无需任何 模型并行 (model parallelism)。

• 性能：对于高达 1000 亿参数的模型，ZeRO-Infinity 实现了超过 40 TFlops/GPU 的卓越性能。这意味着像 GPT-3 这样的模型可以在单个 DGX-2 盒子中进行 微调 (fine-tuning)，这对于大多数数据科学家来说是可行的。

• 易用性：用户无需结合 模型并行 (model parallelism) 或 流水线并行 (pipeline parallelism)，也无需进行模型代码重构，大大简化了大规模模型的训练过程。这与 3D 并行 (3D parallelism) 形成鲜明对比，后者在 200 亿参数以上就无法扩展到单个 DGX-2 节点。

  以下是原文 Figure 5c 的结果：

  该图像是图表，展示了ZeRO-Infinity在训练超大规模模型上的效率与扩展性，分别对比了3D并行、ZeRO-Offload与ZeRO-Infinity在不同模型规模及GPU数量下的吞吐率表现，体现了其超线性扩展能力和在单节点上训练1万亿参数模型的优势。

Figure 5: Efficiency and scalability of ZeRO-Infinity for training multi-trillion parameter models.

6.4. 系统特性对模型规模的影响 (Impact of System Features on Model Scale)

本节评估了不同设备放置策略和 内存中心平铺 (memory-centric tiling) 对模型规模的影响。

6.4.1. 最大模型规模 (Maximum Model Size)

不同的设备放置和分区策略对最大模型规模有显著影响。

以下是原文 Table 2 展示的设备放置选项和分区策略：

符号解释：

• Optimizer + Grad(devices/partitioned): 优化器状态 (optimizer states) 和 梯度 (gradients) 存储的设备以及是否分区（✓ 表示分区，X 表示未分区）。

• $Parameters(devices/partitioned)$: 参数存储的设备以及是否分区。

• [GPU]: 存储在 GPU 内存中。

• [CPU,GPU]: 存储在 CPU 和 GPU 内存中（通常是部分在 CPU）。

• [NVMe,CPU,GPU]: 存储在 NVMe、CPU 和 GPU 内存中（通常是部分在 NVMe）。

• 数据并行 (Data Parallel)：仅使用 数据并行 (data parallel) 时，模型规模受限于 1.4B 参数，因为 GPU 内存有限且模型状态冗余。

• ZeRO-2 和 ZeRO-Offload：通过 优化器 (optimizer) 和 梯度分区 (gradient partitioning) 并 卸载 (offloading) 到 CPU（ZeRO-2 和 ZeRO-Offload），模型规模可扩展到 9 倍，达到 13B 参数。

• ZeRO-Infinity-CPU：ZeRO-Infinity 将参数状态 分区 (partition) 并 卸载 (offload) 到 CPU，使得模型规模接近 100B 参数。

• ZeRO-Infinity-NVMe：最大的规模飞跃来自将模型状态 卸载 (offload) 到 NVMe，最终达到 1T 参数，相对于纯 数据并行 (data parallel) 增加了 700 倍。

  以下是原文 Figure 6a 的结果：

  该图像是论文中图6，展示了系统特性对模型规模和性能的影响。图(a)比较了不同ZeRO策略的最大模型规模，图(b)展示了不同切片因子下最大隐藏维度，图(c)比较了ZeRO-Infinity与ZeRO Offload的反向时间，图(d)显示了通信重叠带来的加速比，图(e)反映了激活检查点转移至CPU带来的开销。

Figure 6: Impact of system features on model scale and performance.

6.4.2. 最大隐藏维度 (Maximum Hidden Size)

内存中心平铺 (memory-centric tiling) 在支持大 隐藏维度 (hidden sizes) 方面发挥了关键作用。

• 实验设置：训练单个 Transformer 层模型，使用不同 隐藏维度 (hidden sizes) 和 平铺因子 (tiling factors)。为确保内存碎片化一致，预先将总 GPU 内存碎片化为 2 GB 的连续块，使得所有大于 2GB 的内存分配请求都会失败。
• 结果分析：
  • 没有 内存中心平铺 (memory-centric tiling) 时，最大可训练的 隐藏维度 (hidden size) 为 8K。
  • 使用 内存中心平铺 (memory-centric tiling)（平铺因子 (tiling factor) 为 16）时，可以训练高达 64K 的 隐藏维度 (hidden size)。 这表明 内存中心平铺 (memory-centric tiling) 大大简化了 DL 系统栈，避免了对 模型并行 (model parallelism) 的需求，使得数据科学家能够轻松训练具有大 隐藏维度 (hidden sizes) 的模型。

以下是原文 Figure 6b 的结果：

该图像是论文中图6，展示了系统特性对模型规模和性能的影响。图(a)比较了不同ZeRO策略的最大模型规模，图(b)展示了不同切片因子下最大隐藏维度，图(c)比较了ZeRO-Infinity与ZeRO Offload的反向时间，图(d)显示了通信重叠带来的加速比，图(e)反映了激活检查点转移至CPU带来的开销。

Figure 6: Impact of system features on model scale and performance.

6.5. 系统特性对性能的影响 (Impact of System Features on Performance)

本节评估了 无限卸载引擎 (infinity offload engine)、带宽中心分区 (bandwidth-centric partitioning)、重叠中心设计 (overlap-centric design) 和 激活检查点卸载 (activation checkpoint offloading) 对训练速度的影响。

6.5.1. ZeRO-Infinity 与 ZeRO-Offload 的对比 (ZeRO-Infinity vs ZeRO-Offload)

• 对比目标：比较 ZeRO-Infinity 和 ZeRO-Offload 在 反向传播 (back propagation) 时间上 梯度 (gradients) 卸载到 CPU 内存的影响，模型为 80 亿参数。

• 结果：在 64 块 GPU 时，ZeRO-Infinity 比 ZeRO-Offload 提速近 2 倍。

• 原因分析：ZeRO-Infinity 利用了 GPU 之间的聚合 PCIe 带宽来 卸载 (offload) 梯度 (gradients)，而 ZeRO-Offload 受限于单个 PCIe 带宽。

  以下是原文 Figure 6c 的结果：

  该图像是论文中图6，展示了系统特性对模型规模和性能的影响。图(a)比较了不同ZeRO策略的最大模型规模，图(b)展示了不同切片因子下最大隐藏维度，图(c)比较了ZeRO-Infinity与ZeRO Offload的反向时间，图(d)显示了通信重叠带来的加速比，图(e)反映了激活检查点转移至CPU带来的开销。

Figure 6: Impact of system features on model scale and performance.

6.5.2. 预取和重叠 (Prefetching and Overlapping)

• 对比目标：比较在 80 亿参数模型和 64 块 GPU 上，启用和关闭通信重叠与预取时的相对吞吐量差异。

• 结果：预取和重叠对于在每个 GPU 小 批大小 (batch size) 下实现良好性能至关重要。

• 影响变化：当 批大小 (batch size) 增大时，其影响会减小，因为此时 算术强度 (arithmetic intensity) 足够高，通信开销相对不那么显著。

  以下是原文 Figure 6d 的结果：

  该图像是论文中图6，展示了系统特性对模型规模和性能的影响。图(a)比较了不同ZeRO策略的最大模型规模，图(b)展示了不同切片因子下最大隐藏维度，图(c)比较了ZeRO-Infinity与ZeRO Offload的反向时间，图(d)显示了通信重叠带来的加速比，图(e)反映了激活检查点转移至CPU带来的开销。

Figure 6: Impact of system features on model scale and performance.

6.5.3. 激活检查点卸载 (Activation Checkpoint Offload)

• 对比目标：评估 ZeRO-Infinity 中将 激活检查点 (activation checkpoints) 卸载到 CPU 对训练吞吐量的影响。

• 结果：对于小的 隐藏维度 (hidden sizes)，CPU 卸载 激活检查点 (activation checkpoints) 会使训练吞吐量降低高达 1.2x。但对于 32K 和 64K 等较大的 隐藏维度 (hidden sizes)，影响微乎其微。

• 原因分析：这表明对于大型模型，由于其更高的 算术强度 (arithmetic intensity)，可以将 激活检查点 (activation checkpoints) 卸载到 CPU 内存而不会显著影响效率。

  以下是原文 Figure 6e 的结果：

  该图像是论文中图6，展示了系统特性对模型规模和性能的影响。图(a)比较了不同ZeRO策略的最大模型规模，图(b)展示了不同切片因子下最大隐藏维度，图(c)比较了ZeRO-Infinity与ZeRO Offload的反向时间，图(d)显示了通信重叠带来的加速比，图(e)反映了激活检查点转移至CPU带来的开销。

Figure 6: Impact of system features on model scale and performance.

7. 总结与思考

7.1. 结论总结

本文提出了 ZeRO-Infinity，一种新颖的 异构系统技术 (heterogeneous system technology)，旨在解决 GPU 内存墙 (GPU memory wall) 的问题，并实现超大规模深度学习模型的训练。ZeRO-Infinity 的核心创新在于全面利用了 GPU、CPU 和 NVMe 的异构内存，结合 带宽中心分区 (bandwidth-centric partitioning)、重叠中心设计 (overlap-centric design) 和 内存中心平铺 (memory-centric tiling) 等优化技术。

其主要贡献和优势体现在：

• 前所未有的模型规模：能够训练高达 32 万亿参数的模型，比 最先进的 (state-of-the-art) 3D 并行 (3D parallelism) 提高了 50 倍。

• 卓越的训练效率：在 512 块 NVIDIA V100 GPU 上实现了超过 25 petaflops 的吞吐量，并展示了在训练万亿参数模型时的 超线性可扩展性 (superlinear scalability)。

• 高可访问性与易用性：允许在单个 NVIDIA DGX-2 节点上 微调 (fine-tune) 万亿参数模型，无需 模型并行 (model parallelism) 或复杂的代码重构，极大地降低了大规模深度学习的门槛。

• 系统创新：通过 无限卸载引擎 (infinity offload engine)、带宽中心分区 (bandwidth-centric partitioning) 和 重叠中心设计 (overlap-centric design) 等技术，高效地管理和利用异构内存带宽，将计算与通信重叠。

  ZeRO-Infinity 改变了我们对大规模模型训练内存使用的看法，证明即使是廉价、缓慢但容量巨大的 CPU 或 NVMe 内存，通过并行化和智能管理，也能提供高效训练所需的聚合带宽。

7.2. 局限性与未来工作

论文指出，尽管 ZeRO-Infinity 消除了设备内存对模型规模和训练效率的限制，但训练数万亿甚至数百兆亿参数模型在合理时间内完成，仍需要计算能力的大幅飞跃。

以下是原文 Table 3 展示的带宽要求：

符号解释：

• Total devices: 设备总数。

• Achievable peak (pflops/device): 每个设备可达到的峰值 pflops（petaflops）。

• Slow memory bw requirement(GB/s per device): 每个设备与慢速内存之间所需的带宽。

• Slow memory aggregate bw (TB/s): 慢速内存的聚合带宽。

• GPU-to-GPU bw (GB/s): GPU-to-GPU 的带宽。

  Table 3 显示，即使加速器计算能力比 NVIDIA V100 提高 10 倍，在一个拥有 512 个加速器的集群中，ZeRO-Infinity 也只需要每个加速器与慢速内存之间有 $30 \mathrm{GB/s}$ 的带宽即可保持高效。这在当今技术（如通过 NVLink 连接加速器到 CPU 内存）是可实现的。然而，要高效运行在未来的设备上，仍需要设备间带宽的相应飞跃。

具体的未来工作方向可能包括：

• 激活检查点到 NVMe 的卸载 (Offloading Activation Checkpoints to NVMe)：论文提到，当前在 20 万亿参数规模下性能下降的原因之一是 CPU 内存有限，无法存储 激活检查点 (activation checkpoints)。未来的实现可以探索将 激活检查点 (activation checkpoints) 也卸载到 NVMe 存储。
• 动态内存管理 (Dynamic Memory Management)：虽然论文提到了 固定内存管理层 (Pinned Memory Management Layer)，但更动态、更智能的内存管理策略可能进一步优化异构内存的使用。
• 计算与带宽的平衡：随着计算能力的提升，如何持续优化通信和计算的重叠，以及如何更好地平衡不同的带宽需求，将是持续的挑战。

7.3. 个人启发与批判

7.3.1. 个人启发

ZeRO-Infinity 的工作提供了几个重要的启发：

1. 异构内存的巨大潜力：它明确指出，不应将深度学习训练的内存限制在昂贵且有限的 GPU 内存（HBM2）上。CPU DRAM 和 NVMe SSD 等更慢但容量更大的存储介质，通过智能的系统设计，可以成为打破模型规模瓶颈的关键。这种思想对于数据中心和超级计算集群的设计具有深远影响，意味着未来硬件设计可能需要更强调 CPU、GPU 和 NVMe 之间的高带宽互联，而非仅仅追求 GPU 内存的绝对容量。
2. “软件定义硬件”的成功案例：ZeRO-Infinity 通过软件层面的创新，有效地解决了硬件（GPU 内存）带来的限制，这体现了软件在推动计算能力边界方面的重要性。
3. 易用性是关键：论文强调了消除模型代码重构的必要性。即使是最强大的技术，如果使用门槛过高，也难以普及。ZeRO-Infinity 在实现高性能和大规模的同时，注重易用性，这是其能够被广泛采用的重要原因。
4. 超线性可扩展性：通过聚合慢速内存的带宽，实现 超线性可扩展性 (superlinear scalability) 是一个非常有趣的现象，它表明在某些场景下，分布式系统的整体性能甚至可以超越其组件的简单加和，这是对传统线性扩展思维的挑战。

7.3.2. 批判

尽管 ZeRO-Infinity 是一个里程碑式的工作，但仍存在一些可以批判或思考的方面：

1. 对带宽的假设：论文的效率模型和许多优化都基于对带宽的有效利用。虽然 带宽中心分区 (bandwidth-centric partitioning) 能够聚合带宽，但实际网络拓扑、拥塞以及不同设备的 I/O 特性可能会使理想情况复杂化。例如，当节点间通信成为瓶颈时，聚合带宽的优势可能会打折扣。

2. NVMe 的寿命与可靠性：频繁地将数万亿参数 卸载 (offload) 和加载到 NVMe SSD，可能会加速其磨损，影响其寿命和数据可靠性。虽然这对于短期训练可能不是问题，但对于长期运行或关键任务，需要更深入考虑。

3. 对动态工作负载的适应性：动态预取器 (dynamic prefetcher) 旨在适应动态工作流。但在高度动态的模型或数据访问模式下，预取器的准确性和效率可能会受到挑战，导致更多的停顿。

4. 能耗考量：将数据在 GPU、CPU 和 NVMe 之间频繁移动，尤其是大规模 NVMe 读写，可能会带来额外的能耗。在追求极端规模和效率的同时，能耗效率也是一个日益重要的指标。论文没有详细讨论这方面的影响。

5. 模型结构的普适性：论文主要关注 Transformer 模型。虽然 Transformer 具有代表性，但对于具有非常规结构或高度不规则访问模式的模型，ZeRO-Infinity 的某些优化（如 内存中心平铺 (memory-centric tiling)）的有效性可能需要进一步验证。

6. 超参数调整的复杂性：尽管模型代码重构减少了，但引入了更多的系统级超参数（例如，卸载策略、平铺因子、预取深度等）。如何为不同的模型和硬件配置自动或半自动地调整这些参数，以实现最佳性能，仍是一个挑战。

   总的来说，ZeRO-Infinity 为深度学习的未来发展开辟了新的道路，但其在更广阔的场景下，仍有许多值得深究和完善之处。