论文阅读: ZeRO: Memory Optimizations Toward Training Trillion Parameter Models

2023-09-30

懵逼的深度学习

摘要

ZeRO 是一种用于大规模深度学习模型的内存优化解决方案，通过消除数据和模型并行训练中的内存冗余，同时保持了低通信量和高计算粒度。

paper: https://arxiv.org/abs/1910.02054

code: https://github.com/microsoft/DeepSpeed

背景

模型训练的挑战：

传统的 DP 方法不能减少每个设备的内存使用量，导致在当时 32G 显存的 GPU 上无法训练超过 14 亿个参数的模型

现有的方法：流水线并行，模型并行， CPU offloading，在功能性，内存使用，计算通信效率之间做出了取舍

作者提到模型并行最有潜力，需要将模型在垂直方向上进行切分，并将每个层的计算和参数在多个设备之间进行划分，当时的 11B 参数的 T5，8.3B 的 Megatron-LM 都使用了模型并行。但是这需要大量的层间通信，虽然在单个节点内表现良好，作者在两个 DGX-2 节点上使用 Megatron-LM 测试了 40B 的模型，每个 V100 GPU 的计算性能仅为 5 TFLOPS

作者发现内存消耗主要集中在

模型状态：优化器，梯度，参数
residual states：激活值，buffer，内存碎片

针对这两个部分提出了 ZeRO（Zero Redundancy Optimizer）

引论

优化模型状态内存

数据并行（DP）不足：不需要频繁地通信，但需要在每个进程中复制整个模型状态，浪费内存

模型并行（MP）不足：对模型状态分区以高效率使用内存，但会导致过细的计算粒度，需要频繁通信

整体来看，上述方法需要在整个训练过程中静态地维护所有模型状态，但在训练时不总是需要所有的模型状态

改进：

提出了 ZeRO-DP，ZeRO-powered DP ，通过分区模型状态而不是复制来消除 DP 时的内存冗余，并使用动态通信调度策略优化计算通信效率，提出了三个优化阶段

优化器分区（Pos）：内存减少4倍，与DP具有相同的通信量
梯度分区（Pos+g）：内存减少8倍，与DP具有相同的通信量
模型参数分区（Pos+g+p）：内存减少量与 DP 的份数成线性关系。如在 64 个 GPU 上拆分将减少64倍内存。通信仅增加 50%

图中 Ψ 表示模型参数量，K 表示优化器状态的内存倍数，Nd 表示 DP 份数。在该示例中，假设基于 Adam 优化器的混合精度训练，Ψ=7.5B，Nd=64，K=12，其中的详细计算方法后面会介绍

作者在这里计算了通过上述优化在 1024 张 GPU 上训练万亿参数模型，使用 fp16，需要 16TB 空间存储优化器状态，梯度和参数，每张卡占用 16GB 显存

优化 residual states 内存

为了优化激活值，buffer，内存碎片的占用，提出了 ZeRO-R，有以下几个优化点

使用对激活值的 checkpointing 来节省内存，还对激活值切片，根据需要将激活数据转移至 CPU 来优化激活值的内存占用
ZeRO-R 定义了适当的临时缓冲区大小，使内存和计算效率平衡
根据不同 tensor 的生命周期，主动管理内存，预防内存碎片化

综上，ZeRO 主要是由 ZeRO-DP 和 ZeRO-R 两种优化结合

ZeRO 搭配模型并行（MP）

虽然使用 ZeRO 的策略后，MP 的策略变得不那么重要了，MP 在使用时还需要修改模型，相比 DP，有诸多限制

但在激活内存占用非常大时，这时 ZeRO-R 的策略也不能满足训练优化，可以搭配 MP 减少激活内存占用。并且，在结合 ZeRO 和 MP 时，理论上可以优化到 Nd * Nm（MP 份数）倍的内存占用

在小模型的情况下，单独使用 DP 会导致 batch size 过大可能无法收敛，使用 MP 可以在加速的同时减小 batch size 到合适的值，帮助模型收敛

Implementation & Evaluation

作者进行了一些上述工作的实验，得出如下结论

模型大小：将 Megatron 与 MP 相结合，ZeRO-100B 可以高效地运行 1700 亿参数的模型，而单独使用 Megatron 等现有系统在 40 亿参数以上的规模上无法高效扩展。相比于SOTA，模型大小增加了8倍以上
训练速度：改进的内存效率提高了吞吐量和训练速度。通过在 400 台 Nvidia V100 GPU 上运行，ZeRO 可以以每个 GPU 38 TFlops 的速度训练 1000 亿参数的模型，总性能超过 15 Petaflops。与 SOTA 相比，对于相同模型大小，训练速度提高了 10 倍以上
可扩展性：当使用64-400个GPU时，性能呈现超线性加速，即当GPU数量加倍时，性能增加了一倍以上。这是 ZeRO-DP 的特性，它随着DP度数的增加减少了模型状态的内存占用，使得每个GPU能够承载更大的 batch_size ，从而提高性能
大规模训练可行性：ZeRO-100B 使得 130亿参数的模型只需重构即可训练。相比之下，现有系统（如PyTorch Distributed Data Parallel）在 14 亿参数的规模上就会出现内存不足的情况
SOTA：ZeRO 支持拥有 170 亿参数的 Turing-NLG 模型，并取得了 SOTA 的成绩

优化内存占用

前面提到内存消耗主要集中在

模型状态：优化器，梯度，参数
residual states：激活值，buffer，内存碎片

下面展开讨论为何会这样

模型状态：优化器状态，梯度与参数

混合精度训练

在使用adam优化器训练时，Adam使用指数移动平均来计算梯度，这需要保存梯度的拷贝，以稳定更新参数。Adam还使用了自适应学习率机制，会为每个参数自动调整学习率。学习率的自适应性依赖于每个参数的梯度方差。为了计算梯度方差，就需要保存梯度的平方的移动平均值，以便在参数更新时更好地适应局部梯度的特性

在使用混合精度训练时，将参数和梯度存储为 fp16，并在前反向传播时都使用 fp16 更新，但是为了反向传播结束后保证计算的精确，需要保存参数和优化器状态的 fp32 副本，以Adam优化器为例，使用Adam对具有Ψ个参数的模型进行混合精度训练需要足够的内存来存储参数和梯度的fp16副本，内存需求分别为 2Ψ 和 2Ψ 字节。此外，还需要存储优化器状态：参数、动量和方差的 fp32 副本，内存需求分别为 4Ψ、4Ψ 和 4Ψ 字节。文中使用 K 来表示优化器状态的内存乘数，即存储它们所需的额外内存为 KΨ 字节。混合精度 Adam 的 K 值为 12。这导致了 2+2+12=16 Ψ 字节的内存需求。对于像 GPT-2 这样有15亿参数的模型，至少需要24GB的内存，远远高于 3GB 内存来存储 fp16 参数的需求

Residual 内存消耗

在训练过程中，使用的激活函数会占用大量内存。以 GPT-2 模型为例，当序列长度为 1K，batch_size 为 32时，1.5B 参数的模型需要大约 60GB 的内存

计算公式：激活值内存 = Transformer 层数 × hidden_dim × batch_size × seq_len × Transformer 层数

使用激活值 checkpoint 方法可以减少激活函数内存的消耗，会增加 33% 的 recompute 开销，但可以将激活函数内存消耗降低到约 8GB 但对于更大的模型，激活函数的内存消耗仍然可能非常大。如一个拥有 1000 亿参数的 GPT-like 模型，在 batch_size 为32 时，即使使用了激活值 checkpoint 仍需60GB的内存。此外，对于大型模型，用于存储中间结果的临时 buffer 也会占用相当大的内存。例如，对梯度 all_reduce 或梯度计算时会将所有梯度融合到一个 flattened buffer 中，尽管梯度可以以 fp16 存储，但 buffer 可能还是 fp32。对于一个具有 15 亿参数的模型，一个 flattened fp32 buffer 要占用6GB的内存

此外，内存碎片的问题也要注意，在极端情况下，内存碎片可浪费 30% 的内存

ZeRO

ZeRO 提出了两组优化：

ZeRO-DP：优化模型状态内存消耗
ZeRO-R：优化 Residual 内存消耗

ZeRO-DP

DP：优点：计算粒度高，通信低；不足： DP 进程之间冗余存储

MP：优点：通过分区模型利用内存；不足：计算粒度降低

ZeRO-DP 通过分区模型状态并使用动态的通信调度，同时有 DP 和 MP 的优点

Optimizer State Partitioning

在DP中，通过将优化器的状态分成N个分区，使得每个DP进程只更新对应的分区的优化器状态，也就是1/N的总优化器状态参数量

Gradient Partitioning

梯度的计算被分为不同的分区，每个 DP 进程只处理和更新对应参数分区的梯度。文中还采用了一种 ucketization
策略，将同一参数分区的梯度进行分组，并一次性对整个组进行归约操作。类似于NVIDIA的 AMP 优化器中将全局梯度计算进行 bucketization 以重叠计算通信。通过在最后一个分区进行 all-reduce，以减少内存占用，实现计算通信重叠

Parameter Partitioning

参数分区是在 DP 训练中减少内存消耗的一种方式。每个进程只存储与其分区相对应的参数，当需要使用到其他分区的参数进行前向和反向传播时，通过 broadcast 从相应的 DP 进程接收这些参数

Implication on Model Size

在 1024 DP 的情况下，搭配 Pos+g+p，可以实现 1.5 万亿参数的训练，如果只使用 DP 训练，仅能训练 1.5 Billion 参数量

ZeRO-R

将使用的内存分为两类：

长期存在：前向传播时的激活值 checkpoint，反向传播时的参数梯度
短期存在：前向传播时的 recompute，反向传播时的激活值梯度

ZeRO-R 通过将激活值 checkpoint 和梯度移动到预先分配的连续 buffer 中，进行实时内存碎片整理，还减少了查找空闲连续内存的时间

Partitioned Activation Checkpointing

通过激活值 checkpoint 分区，消除了 MP 中的内存冗余，只有在计算中需要使用激活时，才会将激活值复制

在模型中一层完成前向完成时，会将激活值分区到所有并行进程上，如果在反向传播过程中被使用时，则使用 all-gather 重新创建激活值的 copy

Constant Size Buffers

通过保持足够大的常量 buffer ，在计算之前将所有参数融合到这个单独的 buffer 中，可以加速内存读写效率

Memory Defragmentation

前向传播过程中通过激活值 checkpoint 只保留了部分激活值，其余的激活值会被丢弃，因为它们在后向传播时可以重新计算。同样，在后向传播过程中，参数梯度是长期存在的，而激活梯度和其他用于计算参数梯度的缓冲区是短期存在的。这种长期和短期内存的交织导致了内存碎片化的问题

文中提出了 In-place Activation Reuse 的方法，在反向传播过程中可以重复使用激活值的内存，而无需每次都重新分配内存。将不再需要的激活值内存标记为可重用，并在下一次需要相同大小内存的地方重用它们。减少了内存分配和释放的次数，减少了碎片化

ZeRO-DP 通信分析

DP 通信量

在每次反向后执行 all-reduce 平均，这种 all-reduce 会在大模型上会完全依赖通信带宽，现在的 all-reduce 通过流水线执行 reduce-scatter 再 all-gather 达到 all-reduce 的效果，会有两倍的数据通信量

ZeRO-DP 通信量

Pos+g

因为每个进程只保存其分区的梯度，对梯度进程进行 scatter-reduce，再执行 all-gather，通信量与 DP 相同，为 2Ψ

Pos+g+p

通过参数分区，每个进程只存储更新的参数。在计算前向传播结果前，每个分区的进程将权重 braodcast 到所有进程，在前向传播时，通过流水线执行 all-gather 接受其他分区的参数，以减少内存占用，前向传播后，丢弃权重。在反向传播时需要再次 all-gather，因此，总通信量为 (Ψ * N) / N = Ψ

综上，总通信量为 3Ψ，为 DP 的 1.5 倍

ZeRO-R 通信分析

下面通过分析分区激活值 checkpoint（Pa）与MP，DP通信量决定使用Pa还是Pa+cpu

Pa 的通信量权衡取决于模型大小、checkpoint 策略和 MP 策略。论文使用 Megatron-LM 实现的模型背景下进行分析。在带有激活值 checkpoint 的 Megatron-LM 中，每个transformer在正向传播中执行两个大小为 batch × seq_length × hidden_dim 的 all-reduce 操作用于正向传播时的重计算，另外两个 all-reduce 操作用于反向传播。每个块的总通信量为 12 × seq length × hidden dim，因为 all-reduce 的通信量为 2 × message_size。当 ZeRO-R 对激活值 checkpoint 进行分区时，在每个激活值 checkpoint 上的反向传播的正向重新计算之前需要额外的 all-gather 操作。会检查每个transformer块的输入激活，需要一个 all-gather，因此，Pa 的通信量为 seq_length ∗ hidden_dim。因为 all-gather 的通信量为 message_size，计算 Pa 的总通信量小于 MP 原始通信量的 10%。当MP与DP结合使用时，Pa可用于将Pa通信量减少一个数量级，而模型并行通信量增加10%，并在 DP 通信成为性能瓶颈时显著提高效率。另外，Pa将激活内存消耗降低了 MP 并行度，从而允许按比例增加 batch_size。由于 Pa 导致 batch_size 增加一个数量级可能导致 DP 通信量减少一个数量级。如果采用 Pa+cpu，分区激活值 checkpoint 将卸载到cpu，就不再需要激活内存了，与Pa相比，cpu内存之间增加了2倍的数据移动。在极端情况下，DP 通信量是主要瓶颈，因为即使使用Pa，batch_size 也很小，在小batch_size的情况下，只要 cpu 数据传输开销小于DP通信量开销，Pa+cpu 就可以通过增加 batch_size 来提高效率。在给定模型和硬件特性的情况下，可以利用上述分析来决定是否以及何时使用Pa还是Pa+cpu

万亿参数

仅使用 DP，ZeRO能够在1024个GPU上容纳超过1万亿参数的模型。此外，如下表所示，当与 MP 结合使用时，每个DGX2节点内使用16路MP，跨节点使用64路DP，ZeRO能够在1024个GPU上运行超过1万亿参数的模型，但训练时长会超过一年，期待未来算力提升

实验评估

实施：基于 PyTorch 的 ZeRO-100B，包括Pos+g 和 ZeRO-R 中的全部优化点

硬件：由 400 个 V100 GPU（25个DGX-2节点）组成的集群，节点间通信带宽为800 Gbps

Baseline：没有MP的实验使用了 torch 的 DDP，MP 的实验使用 Megatron-LM 的 MP

ZeRO：没有MP的实验使用了 ZeRO-100B 中基于 ZeRO 的 DP 实现。MP 的实验将 ZeRO-100B 中的 ZeRO-powered DP 与 Megatron-LM 的 MP 相结合

模型：基于GPT-2的 transformer 模型，下表是参数配置

Speed and Model size

Baseline 的 Megatron MP 在模型规模增大时性能会快速下降，因为 MP 在 GPU 之间产生了高额通信量，而在超过单个节点以适应更大的模型时，每条链路（NVSwitch）的通信带宽从300GB/秒下降为12.5GB/秒，导致性能严重下降。对比之下，ZeRO-100B 会有 10 倍的训练速度提升

Super-Linear Scalability

如下图所示，使用ZeRO-100B可以实现超线性可扩展性，并通过增加 DP 并行度来提高每个GPU的吞吐量，从而适应更大的 batch_size

Democratizing Large Model Training

下图使用128个GPU，ZeRO-100B 可以训练13B参数的模型，平均每个GPU吞吐量超过40 TFlops。相比之下，没有使用ZeRO，仅使用DP的最大可训练模型仅有 1.4B 参数，每个 GPU 的吞吐量不到 20 TFlops。此外，由于没有 MP 带来的通信开销，这些模型可以在具有较低端计算节点上进行训练，无需 NVLINK 或 NVSwitch 这种高速互联方式

Memory and Performance Analysis

作者讨论了不同优化方法对最大模型大小、内存消耗和性能的影响。作者将这些优化方法分为配置1到5（C1-C5），如下表所示，通过使用固定 batch_size 和 MP 为16，观察启用不同ZeRO优化的可训练模型的最大尺寸

在最大模型尺寸方面，如 Figure 6 所示，通过使用 Pa 优化，模型大小从40B增加到了60B。而通过使用Pos+g优化，在C2的基础上，模型大小增加到了140B，这是因为与C2相比，该优化使模型状态的内存需求减半。使用C5进一步减少了激活内存，将分区激活值 checkpoint 转移到CPU内存，使模型大小增加到150B

对于每个训练迭代中 PyTorch 缓存的最大内存，如 Figure 7 所示，作者观察了40B和100B模型的情况。从C1到C2，缓存的内存大小如期减少。C2到C3的内存消耗差异取决于模型状态与激活内存的大小关系，当激活内存较大时，差异可能增加，当模型状态较大时，差异可能减小。值得注意的是，在40B模型中，从C4到C5时，缓存的内存大小没有减少，但在100B模型中有减少。这是因为100B的激活内存较大，减少不明显。作者指出，当我们处理非常大的模型时，Pa+cpu优化可用于适应更大的 batch_size

对于不同优化设置的最佳性能，如 Figure 8 所示，性能提升与内存消耗的减少相对应。较低的内存消耗可以实现更大的 batch_size，从而提高性能。唯一的例外是60B参数模型在C4和C5之间的性能下降。尽管内存消耗较低，但C5会导致激活在CPU之间移动，这通常会导致性能下降，除非模型非常大以至于无法在没有C5的情况下运行，或者可以在没有C5的情况下运行的 batch_size 很小（例如在 Figure 8 中具有170B参数的模型）。在训练过程中，Pa+cpu优化只在有益时才启用

Turing-NLG, the SOTA language model with 17B parameters

下图展示了在300K次迭代中与之前最先进的Megatron-LM 8.3B参数模型的验证 Perplexity 对比，使用 ZeRO 训练得到的模型指标优于 Megatron-LM，此外 ZeRO100B 还实现了持续的 41.4 TFlops/GPU的吞吐量

结论

ZeRO 牛逼

摘要

背景

引论

优化模型状态内存

优化 residual states 内存

ZeRO 搭配模型并行（MP）

Implementation & Evaluation

相关工作

DP, TP and PP

DP（数据并行）

TP（模型并行）

PP（流水线并行）

非并行化方法节省内存

减少激活内存

压缩内存

激活内存 checkpoint

实时内存管理

CPU Offload

优化执行算法

虚拟内存 Offload

Memory Efficient Optimizer

自适应学习率

自适应优化器

训练优化器