Training Infrastructure

2026年10月2日约 32 分钟taking-lying-flat

🐊 Training Deep Nets with Sublinear Memory Cost

本文提出一种系统性方法，用于降低深度神经网络训练过程中的内存消耗。主要关注如何减少存储中间结果（特征图）及梯度所需的内存，因为在许多常见的深度网络架构中，参数所占用的存储空间相较于中间特征图而言较小

通过计算图分析，自动执行原地操作和内存共享优化。更为关键的是，提出一种以计算换取内存的新方法。基于这一方法，给出一种实用算法：对于一个包含 $n$ 层的网络，仅需将前向传播的计算开销增加至原来的两倍，即可将特征图的存储开销降至 $O(\sqrt{n})$
证明在极端情况下，训练一个包含 $n$ 层的网络时，特征图的存储开销可以低至 $O(\log n)$

由于存储特征图及其梯度的开销随网络深度线性增长，探索更深模型的能力受到设备内存容量的限制。最优模型的规模与复杂度通常会随着训练数据量的增加而增长，因此，具备较高内存效率的训练算法至关重要

降低内存消耗不仅能够支持更大规模模型的训练，还允许采用更大的批大小，从而提高设备利用率，并改善批归一化等按批次计算的算子的稳定性
对于内存受限的设备，这有助于改善内存访问的局部性，并可能带来更优的内存访问模式。此外，在训练深度卷积神经网络时，降低内存消耗还支持从模型并行切换至数据并行，这在某些情况下能够带来收益

🏗️ Memory Optimization with Computation Graph

计算图由表示运算的节点以及表示运算之间依赖关系的边组成。给定网络配置（即前向计算图）后，可以构建相应的反向路径以计算梯度。具体而言，按照逆拓扑序遍历前向计算图，并像常规反向传播算法一样应用反向算子，即可构建反向路径

在训练深度卷积神经网络或循环神经网络时，通常有很大一部分内存用于存储中间输出及梯度。每个中间结果都对应计算图中的一个节点。合理的内存分配算法能够通过内存共享，尽可能减少为这些节点分配的内存

原地操作（Inplace operation）： 将输出值直接存储到某个输入值所占用的内存中
内存共享（Memory sharing）： 回收不再需要的中间结果所占用的内存，并将其用于其他节点

算法按拓扑序遍历计算图，并使用计数器表示每条记录的活跃状态。当不存在其他尚未执行且依赖于某个输入的运算时，就可以对该输入执行原地操作。当已回收的内存被另一个节点复用时，就实现内存共享。这一方法也可以作为动态运行时算法，在遍历计算图的过程中，利用垃圾回收器回收不再使用的内存。本文将其用作静态内存分配算法，在执行开始之前为各个节点分配内存，以避免运行时垃圾回收带来的开销

数据依赖关系会延长各个输出的生命周期，并增加大型网络的内存消耗。因此，深度学习框架应当：仅声明梯度算子所必需的依赖关系。基于依赖信息进行活跃性分析，并支持内存共享

🌊 Trade Computation for Memory

由于大多数梯度算子依赖于前向传播的中间结果，在训练一个包含 $n$ 层的卷积神经网络，或处理长度为 $n$ 的序列的循环神经网络时，仍然需要 $O(n)$ 的内存来存储中间结果。为进一步降低内存消耗，提出丢弃部分中间结果，并在需要时通过额外的前向计算恢复这些结果

具体而言，在反向传播阶段，可以从最近的已保存结果出发，执行前向计算，重新生成被丢弃的中间结果。为此，将神经网络划分为若干区段，算法仅保存各区段的输出，并丢弃区段内部的全部中间结果。在反向传播过程中，以区段为单位重新计算这些被丢弃的结果。因此，所需内存仅包括存储各区段输出的开销，以及各区段执行反向传播时所需内存的最大值

基于相同的核心思想，引入一种通用的梯度图构建算法

在该算法中，用户在计算图的节点上定义函数 $m:\mathcal{V}\rightarrow\mathbb{N}$ ，用于指定某个结果可以被重新计算的次数。由于重计算本质上是在复制节点，即对节点进行镜像，因此将 $m$ 称为镜像计数函数。当所有镜像计数均设为 0 时，该算法退化为常规梯度图的构建过程
为指定上述重计算模式，用户只需对各区段内部的节点设置 $m(v)=1$ ，并对各区段的输出节点设置 $m(v)=0$ 。镜像计数也可以大于 1，由此可以得到递归形式的推广

上述通用方法的一种直接应用，是丢弃低开销运算的结果，并保留计算耗时较长的结果。这通常适用于卷积神经网络中的 Conv-BatchNorm-Activation 计算流程。可以始终保留卷积运算的结果，同时丢弃批归一化、激活函数和池化运算的结果。由于批归一化和激活函数的计算开销都较低，因此在实际应用中，这种方法能够以很小的额外计算开销节省内存

假设将一个包含 $n$ 层的网络划分为 $k$ 个区段，则训练该网络所需的内存开销如下

式中的第一部分表示在各区段上执行反向传播所需的内存开销。若各区段均匀划分，则该开销为 $O(n/k)$ 。第二部分表示存储区段之间中间输出的开销。令 $k=\sqrt{n}$ ，即可得到 $O(2\sqrt{n})$ 的内存开销
该算法仅需在每个训练步中额外执行一次前向传播，便能将内存开销降低至次线性水平。由于反向运算的耗时接近前向运算的两倍，因此，这只会使整体计算速度略有下降

也可以从另一种视角理解上述内存优化方案。可以将每个区段视为一个复合算子，将区段内部的全部运算封装在一起。该复合算子通过执行描述其内部计算过程的子图来计算梯度。这一视角使我们能够将一系列运算视为子程序。子图内部的优化不会影响外部计算，因此，可以将所提出的内存优化方案递归地应用于各个子图

递归地保存检查点并重计算子网络，可进一步降低特征图存储开销。在各层特征图大小相近时，以单层特征图占用的内存为单位，令 $g(n)$ 表示一个包含 $n$ 层的网络在前向和反向传播中的特征图存储开销。记 $k$ 为每一级递归保存的中间检查点数量。将当前的 $n$ 层划分为 $k+1$ 组，每组约 $n/(k+1)$ 层，并依次递归处理。各组复用临时内存，因此总存储开销为本级检查点的存储开销，加上处理一组所需的存储开销

当 $k=1$ 时，每一级递归只保存一个中间检查点，并将需要处理的层数减半，因此得到 $g(n)=\log_2 n$ 。对于一个包含 $n$ 层、各层特征图大小相近的网络，特征图的存储开销由 $O(n)$ 降至 $O(\log_2 n)$ 。代价是在反向传播期间反复执行前向重计算

🍄 Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism

随着模型规模不断增大，其内存需求逐渐超出现代处理器的容量限制，因此需要采用激活检查点等额外的内存管理技术

ADAM 等广泛使用的优化算法需要为每个参数分配额外内存，以存储动量及其他优化器状态，从而进一步限制能够有效训练的模型规模
若干模型并行方法通过对模型进行划分来突破这一限制，使模型权重及其对应的优化器状态无需同时驻留于同一处理器的内存中
GPipe 和 Mesh-TensorFlow 提供支持不同类型模型并行的框架。然而，这些方法需要重写模型，并依赖仍处于开发阶段的定制编译器和框架

在本研究中，采用层内模型并行，实现一种简单而高效的模型并行方法。利用基于 Transformer 的语言模型所固有的结构，构建一种简洁的模型并行实现，使其能够在 PyTorch 中高效训练，而无需编写定制的 C++ 代码或使用定制编译器。该方法与 GPipe 等工作所采用的流水线模型并行方法相互正交

仅对现有的 PyTorch Transformer 实现进行少量有针对性的修改，便实现一种简单而高效的模型并行方法
对所提出的模型并行与数据并行技术进行深入的实证分析，并在使用 512 张 GPU 时实现最高 76% 的扩展效率
证明在类 BERT 模型中，合理设置层归一化的位置，对于在模型规模增大时持续提高准确率至关重要

🪐 Model Parallel Transformers

记 $X$ 为输入， $A$ 为权重矩阵， $Y$ 为中间激活。MLP 模块的第一部分由一次通用矩阵乘法（GEMM）及其后的 GeLU 非线性激活函数组成

对这一 GEMM 进行并行化的一种选择，是沿行维度划分权重矩阵 $A$ ，并沿列维度划分输入 $X$

此时 $Y = \operatorname{GeLU}(X_1A_1 + X_2A_2)$ 。由于 GeLU 是非线性函数， $\operatorname{GeLU}(X_1A_1 + X_2A_2) \neq \operatorname{GeLU}(X_1A_1) + \operatorname{GeLU}(X_2A_2)$ ，因此，这种方法需要在执行 GeLU 函数之前设置一个同步点
另一种选择是沿列维度划分 $A$ ，即 $A = [A_1, A_2]$ 。这种划分方式允许对各个划分后的 GEMM 输出独立应用 GeLU 非线性激活函数

这种划分方式能够消除一个同步点。因此，第一个 GEMM 采用列并行，第二个 GEMM 的权重沿行维度划分，使第二个 GEMM 能直接接收 GeLU 层的分片输出，无需在两个 GEMM 之间通信。随后，在将第二个 GEMM 的输出传递至 Dropout 层之前，通过 all-reduce 汇总各 GPU 的局部结果。MLP 模块在前向和反向传播中各需一次 all-reduce。图中的 $f$ 在前向执行恒等映射，在反向执行 all-reduce； $g$ 则在前向执行 all-reduce，在反向执行恒等映射。两者互为共轭，可通过 PyTorch 的自动求导接口实现

对于自注意力模块，利用多头注意力运算所固有的并行性，以列并行方式划分与键、查询和值相关的 GEMM，使每个注意力头对应的矩阵乘法均在单张 GPU 上本地完成。由此，能够按照注意力头将参数和计算负载分配至不同 GPU，且完成自注意力计算无需立即进行通信

自注意力之后的输出线性层所对应的 GEMM 则沿行维度进行并行化，直接接收并行注意力层的输出，而无需在 GPU 之间进行通信
对于 MLP 层和自注意力层，这种方法均将两个 GEMM 作为一组进行融合，消除二者之间的同步点，从而实现更好的扩展性。因此，仅需在前向传播中执行两次 all-reduce，并在反向传播中执行两次 all-reduce，即可完成一个基本 Transformer 层中的全部 GEMM 运算

Transformer 语言模型具有一个输出嵌入矩阵，其维度为隐藏层维度 $H$ 乘以词表大小 $v$ 。由于现代语言模型的词表通常包含数万个 token，例如 GPT-2 的词表大小为 50,257，因此，对输出嵌入对应的 GEMM 进行并行化是有益的

在 Transformer 语言模型中，输出嵌入层与输入嵌入层共享权重，因此需要同时修改二者。沿词表维度，即列维度，对输入嵌入权重矩阵 $E_{H \times v}$ 进行并行划分，得到 $E = [E_1, E_2]$ 。由于每个分区此时仅包含嵌入表的一部分，因此需要在输入嵌入操作之后执行一次 all-reduce
对于输出嵌入，一种方法是通过并行 GEMM $[Y_1, Y_2] = [XE_1, XE_2]$ 获得 logits，随后执行一次 all-gather，即 $Y = \operatorname{all-gather}([Y_1, Y_2])$ ，再将结果传递给交叉熵损失函数。然而，在这种情况下，all-gather 需要传输 $b \times s \times v$ 个元素，其中 $b$ 为批大小， $s$ 为序列长度。由于词表规模较大，这会产生巨大的通信量
为降低通信量，将并行 GEMM 的输出 $[Y_1, Y_2]$ 与交叉熵损失计算进行融合，使通信数据的维度降至 $b \times s$ 。通过传输标量损失而非 logits，可以大幅减少通信量，从而提高模型并行方法的效率

所采用的大部分模型并行技术，旨在减少通信，并使 GPU 保持计算受限状态。对于 Dropout、层归一化和残差连接等操作，选择在各张 GPU 上重复执行计算，而不是由某张 GPU 完成部分计算后，再将结果广播至其他 GPU

在每张 GPU 上保留层归一化参数的副本，并对模型并行区域的输出张量执行 Dropout 和残差连接操作，然后将结果作为输入传递至下一个模型并行区域。在模型优化过程中，每个模型并行工作进程负责优化各自的参数集合。由于所有数值要么仅保存在某张 GPU 本地，要么在各张 GPU 上均有副本，因此，在这一方案中无需对更新后的参数值进行通信

🦫 ZeRO: Memory Optimizations Toward Training Trillion Parameter Models

基础数据并行（DP）不会降低单个设备的模型状态显存需求。在 ZeRO 论文的 32GB GPU 实验配置下，使用标准数据并行训练时，模型参数量超过约 1.4B 即会出现显存不足

其他现有方案，例如流水线并行、模型并行和 CPU 卸载等，需要在功能支持、易用性、内存效率以及计算与通信效率之间进行权衡，而这些因素对于实现高速、大规模训练都至关重要
在现有的大模型训练方案中，模型并行（MP）或许是最具潜力的方法。参数量为 11B 的 T5 和参数量为 8.3B 的 Megatron-LM，均借助模型并行实现训练，分别采用 Mesh-TensorFlow 和 Megatron-LM 中的实现
然而，MP 难以在这些模型规模的基础上继续大幅扩展。MP 沿纵向切分模型，将每一层的计算与参数划分到多个设备上，因此需要在各层的计算过程中进行大量通信。由此，MP 在 GPU 间通信带宽较高的单个节点内能够取得良好效果，但一旦扩展到多个节点，其效率便会迅速下降

如何克服现有方案的局限，更高效地训练大规模模型？为回答这一问题，作者首先全面分析现有系统在模型训练过程中的内存消耗，并将其划分为两部分

对于大规模模型，大部分内存由模型状态占据，包括优化器状态（例如 Adam 中的一阶矩和二阶矩）、梯度以及参数
其余内存由激活值、临时缓冲区和无法利用的内存碎片占据，将这些统称为剩余状态
提出零冗余优化器（Zero Redundancy Optimizer，ZeRO），在优化这两部分内存效率的同时，保持较高的计算和通信效率

模型状态的内存优化。模型状态通常占据训练过程中最大比例的内存，但 DP 和 MP 等现有方法均未提供令人满意的解决方案

DP 具有较高的计算与通信效率，但内存效率较低；MP 则可能存在计算与通信效率不足的问题。具体而言，DP 在所有数据并行进程中复制完整的模型状态，造成冗余的内存消耗；MP 通过划分这些状态提高内存效率，但通常会使计算粒度过细，并引入高昂的通信开销，从而降低扩展效率
这些方法均以静态方式保留整个训练过程所需的全部模型状态，尽管训练过程中并非每个时刻都需要所有状态
基于这些观察，提出由 ZeRO 支持的数据并行方法 ZeRO-DP，使其在获得 DP 的计算与通信效率的同时，实现 MP 的内存效率
ZeRO-DP 通过对模型状态进行分片，消除数据并行进程之间的冗余状态存储；同时，在训练过程中采用动态通信调度，保留 DP 的计算粒度和通信量，从而维持其计算与通信效率

ZeRO-DP 依次对优化器状态、梯度和参数进行分片。记 $\Psi$ 为模型参数量， $N_d$ 为数据并行度， $K$ 为每个参数的优化器相关状态占用的字节数。对于 FP16/FP32 混合精度 Adam， $K=4+4+4=12$ ，分别对应 FP32 主权重、一阶矩和二阶矩。三个阶段累积启用后

优化器状态分片 $P_{os}$ ：可实现 4 倍的内存压缩，通信量与 DP 相同
进一步进行梯度分片 $P_{os+g}$ ：可实现 8 倍的内存压缩，通信量与 DP 相同
进一步进行参数分片 $P_{os+g+p}$ ：内存压缩倍数与数据并行度 $N_d$ 呈线性关系，同时通信量增加 50%

剩余状态的内存优化：在 ZeRO-DP 提高模型状态的内存效率之后，激活值、临时缓冲区和无法利用的内存碎片所占据的其余内存，可能成为下一个内存瓶颈。为此，提出 ZeRO-R，分别优化这三类因素造成的剩余内存消耗

对于激活值，即前向传播期间保存、供反向传播使用的中间结果，激活检查点技术虽然有所帮助，但对于大规模模型仍然不够。因此，ZeRO-R 通过激活值分片，识别并消除现有 MP 方法中的激活值冗余副本，从而优化激活值的内存占用。在适当的情况下，它还会将激活值卸载到 CPU
ZeRO-R 为临时缓冲区设置适当的大小，以平衡内存效率与计算效率
不同张量的生命周期存在差异，因此训练过程中会产生内存碎片。内存碎片导致连续内存空间不足，即使仍有足够的空闲内存，也可能发生内存分配失败。ZeRO-R 根据张量生命周期的差异主动管理内存，以防止内存碎片化

ZeRO-DP 与 ZeRO-R 共同构成一套面向深度学习训练的高效内存优化系统，将其统称为 ZeRO

🍀 Where Did All the Memory Go?

内存究竟消耗在哪里？在模型训练过程中，大部分内存由模型状态占据，即用于存储优化器状态、梯度和参数的张量。除模型状态之外，其余内存由激活值、临时缓冲区和内存碎片占据，将其称为剩余状态

在训练过程中，模型状态占据大部分设备内存。以深度学习训练中最常用的优化器之一 Adam 为例，Adam 需要存储两类优化器状态来计算更新

梯度的一阶矩，即梯度的指数移动平均；梯度的二阶矩，即梯度平方的指数移动平均。此外，训练还需保存参数和梯度。在混合精度 Adam 中，优化器相关状态还包括 FP32 主权重，通常占据模型状态显存的主要部分

混合精度训练同时保存 FP16 计算参数与 FP32 优化器状态。这里采用 FP16/FP32 混合精度训练，使用 FP16 权重与激活值进行前向和反向计算，以利用 Tensor Core 的计算吞吐量。优化器保留 FP32 主权重、一阶矩和二阶矩，用于计算参数更新，三者合计占 $K\Psi=12\Psi$ 字节。FP16 参数和梯度各占 $2\Psi$ 字节，因此，未分片的模型状态总占用为 $(2+2+K)\Psi=16\Psi$ 字节

激活值存储也是训练内存开销的重要组成部分

例如，对于参数量为 1.5B 的 GPT-2 模型，当序列长度为 1K、批大小为 32 时，激活值约需占用 60GB 内存。激活检查点技术，也称为激活重计算，是一种常用的优化方法：它以增加 33% 的重计算开销为代价，将激活值的内存需求近似降低至原总量的平方根量级
尽管这种方法能够显著减少内存需求，但对于规模更大的模型，即使采用激活检查点技术，激活值的内存占用仍可能十分庞大

临时缓冲区：对于大规模模型，用于存储中间结果的临时缓冲区也会占据不可忽略的内存。梯度 all-reduce 或梯度范数计算等操作，为提高吞吐量，通常会先将所有梯度合并到一个展平的缓冲区中，再执行相应操作

例如，跨设备 all-reduce 的有效带宽会随着消息大小的增加而提高。尽管梯度本身通常以 fp16 张量存储，但根据具体操作的要求，合并后的缓冲区可能采用 fp32 张量。当模型规模较大时，这些临时缓冲区的大小不可忽视

内存碎片：此外，即使仍有大量空闲内存，也可能出现可用内存不足的情况，内存碎片化便会导致这一问题。如果没有足够大的连续内存空间来满足分配请求，即使空闲内存总量大于请求大小，内存分配仍然会失败

🌟 ZeRO: Insights and Overview

ZeRO-DP 在保留 DP 训练效率的同时，实现 MP 的内存效率。ZeRO-DP 对模型状态进行分片，而不是复制完整状态；同时采用动态通信调度，利用模型状态需求随时间变化的特性，并尽量降低通信量。由此，随着数据并行度提高，ZeRO-DP 能够按比例降低模型的单设备内存占用，同时使通信量保持接近标准 DP 的水平，从而维持训练效率

数据并行保留较大的计算粒度，通常具有更高的扩展效率。MP 在减小计算粒度的同时，也增加通信开销。当划分程度超过一定限度时，过细的计算粒度会降低单块 GPU 的执行效率，而增加的通信开销则会限制跨 GPU 的可扩展性，尤其是在跨节点通信时。DP 具有更大的计算粒度和更低的通信量，因此能够实现高得多的效率
数据并行在各进程重复保存模型状态，增加内存占用。模型状态被冗余地存储在所有数据并行进程中。MP 通过对模型状态进行分片来提高内存效率
DP 和 MP 都会保留整个训练过程所需的全部模型状态，但并非所有状态在每个时刻都需要使用：例如，每一层的参数仅在该层执行前向传播和反向传播时才会被使用

ZeRO 将激活检查点分片存储在不同 GPU 上，并在需要时通过 all-gather 重建完整的激活检查点，从而消除 MP 中的内存冗余。激活值的内存压缩倍数与模型并行度成正比。对于超大规模模型，ZeRO 甚至可以选择将激活值分片迁移到 CPU 内存中；由于这些模型具有较高的算术强度，这种方式仍然能够保持良好的效率

MP 虽然对模型状态进行分片，但通常仍需要复制激活值。例如，如果将一个线性层的参数沿纵向切分，并在两块 GPU 上并行计算，那么每块 GPU 都需要完整的激活值，才能计算其负责的参数分片
对于 GPT-2 或规模更大的模型，其算术强度非常高。这里的算术强度是指每次迭代的计算量与该次迭代中激活检查点数据量的比值。该比值很大，并且随隐藏维度呈线性增长。因此，即使通信带宽较低，也能够隐藏激活检查点的数据传输开销

ZeRO-R 使用固定大小的缓冲区，避免临时缓冲区随模型规模增长而急剧膨胀，同时确保缓冲区足够大，以维持执行效率

内存碎片化源于短生命周期与长生命周期内存对象的交错分配。在前向传播过程中，激活检查点属于长生命周期对象，而重计算得到的激活值属于短生命周期对象。类似地，在反向传播过程中，激活值梯度的生命周期较短，而参数梯度的生命周期较长。基于这一观察，ZeRO 将激活检查点和梯度迁移到预先分配的连续内存缓冲区中，在运行过程中动态整理内存碎片。这不仅提高内存的可用性，还减少内存分配器查找空闲连续内存所需的时间，从而提高执行效率

🌊 Deep Dive into ZeRO-DP

优化器状态分片：对于数据并行度为 $N_d$ 的情况，将优化器状态划分为 $N_d$ 个大小相等的分区，使第 $i$ 个数据并行进程仅更新第 $i$ 个分区对应的优化器状态。因此，每个数据并行进程只需存储和更新全部优化器状态的 $\frac{1}{N_d}$ ，并且仅更新全部参数的 $\frac{1}{N_d}$

在每个训练步骤结束时，在数据并行进程之间执行一次 all-gather，使所有数据并行进程都获得完整的更新后参数
采用优化器状态分片后，每个设备的模型状态显存占用从 $4\Psi+K\Psi$ 降低至 $4\Psi+\frac{K\Psi}{N_d}$

由于每个数据并行进程只更新其对应的参数分区，它只需要这些参数归约后的梯度。在反向传播过程中，当各层的梯度计算完成后，仅将这些梯度归约到负责更新相应参数的数据并行进程。归约完成后，原有梯度便不再需要，其占用的内存可以释放。存储梯度所需的内存便从 $2\Psi$ 字节降低至 $\frac{2\Psi}{N_d}$ 字节

这一过程实质上是 Reduce-Scatter 操作，即将不同参数对应的梯度归约到不同进程。为提高实际执行效率，采用分桶策略，将属于某个特定分区的全部梯度放入同一个桶中，并对整个桶一次性执行归约。这与 NVIDIA AMP 优化器对梯度 all-reduce 计算进行分桶，以实现通信与计算重叠的思路相似
针对各个分区执行归约，而不是 all-reduce，从而减少内存占用，并实现计算与通信的重叠
同时对梯度和优化器状态分片后，每个设备的模型状态显存占用进一步降低至 $2\Psi+\frac{14\Psi}{N_d}$

与优化器状态和梯度的处理方式相同，每个进程仅存储其所属分区对应的参数。当执行前向传播或反向传播需要使用本分区之外的参数时，便通过广播从相应的数据并行进程获取这些参数

这种方法仅将标准 DP 系统的总通信量增加至原来的 1.5 倍，同时实现与 $N_d$ 成正比的内存压缩
采用参数分片后，对于一个包含 $\Psi$ 个参数的模型，其内存消耗从 $16\Psi$ 降低至 $\frac{16\Psi}{N_d}$

⚓ Deep Dive into ZeRO-R

MP 的设计需要复制激活值，因此在参与模型并行的 GPU 之间产生冗余的激活值副本。ZeRO 通过对激活值进行分片来消除这种冗余，并且仅在激活值即将用于计算之前，逐层恢复其完整副本，每次只恢复一层的激活值

具体而言，当模型某一层的前向传播计算完成后，该层的输入激活值被分片存储在所有模型并行进程中，直到反向传播再次需要使用它们。此时，ZeRO 通过一次 all-gather 操作，在各进程上重新恢复完整的激活值副本。将这一优化记为 $P_a$ 。它与激活检查点技术配合使用，仅存储分片后的激活检查点，而不保留完整的冗余副本
当模型规模极大且设备内存十分有限时，还可以将这些分片后的激活检查点卸载到 CPU，使设备端的激活值内存开销降至接近零，但代价是额外的通信开销。并将这种优化记为 $P_{a+cpu}$
通过分片式激活检查点，ZeRO 能够实现与模型并行度成正比的激活值内存压缩

ZeRO 通过审慎选择临时数据缓冲区的大小，平衡内存效率与计算效率。在训练过程中，某些操作的计算效率高度依赖输入规模，较大的输入通常能够获得更高的效率。例如，大规模 all-reduce 操作能够获得远高于小规模 all-reduce 的有效带宽

为提高执行效率，NVIDIA Apex 或 Megatron 等高性能库会先将所有参数融合到一个缓冲区中，再执行相应操作。然而，融合缓冲区的内存开销与模型规模成正比，可能成为限制因素。例如，对于一个参数量为 3B 的模型，一个采用 32 位精度的融合缓冲区就需要 12GB 内存
为解决这一问题，当模型规模过大时，采用能够保持较高执行效率的固定大小融合缓冲区。这样，缓冲区大小便不再依赖模型规模；同时，只要缓冲区足够大，仍然能够保持良好的执行效率

模型训练过程中的内存碎片化源于激活检查点机制和梯度计算。在采用激活检查点的前向传播过程中，仅有选定的激活值会被保留供反向传播使用，而大部分激活值会被丢弃，因为它们可以在反向传播时重新计算

这使短生命周期内存对象，即被丢弃的激活值，与长生命周期内存对象，即作为检查点保留的激活值，交错分布，从而产生内存碎片。类似地，在反向传播过程中，参数梯度属于长生命周期对象，而激活值梯度以及计算参数梯度所需的其他缓冲区属于短生命周期对象。这种短生命周期与长生命周期内存对象的交错分布，同样会导致内存碎片化
在空闲内存充足的情况下，少量内存碎片通常不会造成问题；但在内存受限的大模型训练中，内存碎片化会导致两个问题：即使空闲内存总量充足，也可能因缺少足够大的连续内存空间而发生内存不足；内存分配器需要花费大量时间寻找能够满足分配请求的连续内存空间，导致执行效率下降
ZeRO 为激活检查点和梯度预先分配连续内存块，并在这些数据生成后将其复制到预分配的内存中，从而在运行过程中动态整理内存碎片

⛵ Communication Analysis of ZeRO

ZeRO 通过消除内存冗余来扩大可训练模型的规模，一个自然的问题是：是否以增加通信量为代价来提高内存效率？换言之，与标准 DP 方法相比，采用 ZeRO 的 DP 方法需要多少通信量

采用 $P_{os}$ 和 $P_g$ 时，ZeRO-DP 不会引入额外通信，同时能够实现最高 8 倍的内存压缩
在 $P_{os}$ 和 $P_g$ 的基础上进一步采用 $P_p$ 时，ZeRO-DP 的通信量最高为标准 DP 的 1.5 倍，同时使内存压缩倍数达到 $N_d$

在数据并行训练中，反向传播结束之后、计算下一步的参数更新之前，需要对所有数据并行进程中的梯度取平均。这一操作通过 all-reduce 集合通信完成。对于大规模模型，all-reduce 通信完全受通信带宽限制

先进的 all-reduce 实现通常采用两阶段方法。第一阶段执行 reduce-scatter，将数据的不同部分归约到不同进程；第二阶段执行 all-gather，使每个进程收集所有进程上已经归约的数据。这两个阶段共同完成一次 all-reduce
reduce-scatter 和 all-gather 均采用流水线方式实现，对于包含 $\Psi$ 个元素的数据，两者各自需要传输总计 $\Psi$ 个元素。因此，标准 DP 在每个训练步骤中的数据传输量为 $2\Psi$

采用梯度分片后，每个进程仅存储更新其对应参数分区所需的那部分梯度。因此，ZeRO 只需要对梯度执行 scatter-reduce 操作，而无须执行 all-reduce，由此产生的通信量为 $\Psi$ 。当每个进程完成其负责的参数分区更新后，再执行一次 all-gather，收集所有数据并行进程中已经更新的参数。这同样产生 $\Psi$ 的通信量。因此，每个训练步骤的总通信量为 $\Psi+\Psi=2\Psi$ ，与标准 DP 完全相同

采用参数分片后，每个数据并行进程仅存储由自身负责更新的参数。因此，在前向传播过程中，它需要接收其他所有分区的参数。不过，这一过程可以采用流水线方式执行，以避免额外的内存开销

在计算某个特定分区对应模型部分的前向传播之前，负责该分区的数据并行进程可以将权重广播到所有数据并行进程。该分区的前向传播完成后，便可以丢弃这些参数。因此，总通信量为 $\frac{\Psi\times N_d}{N_d}=\Psi$
重新调度参数的 all-gather，将其分散到整个前向传播过程中，并在参数使用完毕后将其丢弃。不过，在反向传播过程中，还需要以相反的顺序再次执行这一 all-gather。总通信量由这些 all-gather 产生的通信量，以及梯度 reduce-scatter 产生的通信量共同构成，总计为 $3\Psi$ ，即标准 DP 的 1.5 倍。梯度分片和参数分片都利用同一个核心洞见：并非所有梯度和参数状态在每个时刻都需要使用。因此，可以通过合理安排这些状态的通信来优化内存占用