ZeRO论文精读

ZeRO: Memory Optimizations Toward Training Trillion Parameter Models

将大型深度学习模型载入有限的设备内存是一个挑战，其计算、通信、开发效率存在根本限制，ZeRO，零冗余优化器（Zero Redundancy Optimizer）是一种内存优化方法。它主要完成的工作是：消除数据和模型并行训练中的内存冗余，保留了低通信量和高计算粒度，以持续的高效率按设备数量比例扩展模型。

问题描述

首先，论文阐述了训练大规模模型参数时，模型的参数不能容纳在单个设备的GPU或TPU内存中，因此增加更多设备无助于扩展训练。（此处，为何在分布式训练当中，其参数不能分开来摆放？）

1. 基本数据并行（Basic Data Parallelism）不能减少每台设备的内存，模型参数超过1.4B时，32GB的GPU内存将会耗尽。
   1. 此处简单计算一下，32GB的GPU内部如果模型以全精度即FP32存储，则32GB的内存可以放置大约$\frac{32\times 2^{30}\mathrm{bytes}}{4bytes/parameter}\approx 8B$，此处的计算假设所有的32GB的内存都用于存储模型参数，而实际情况中，还需要考虑其他内存开销，如梯度存储、优化器状态、中间激活层。因此，实际能够支持的模型参数数量会小于8B，而论文中给出的1.4B实际很小了。
   2. 一个直观的解决方案是采用模型并行（MP），在多个GPU上分布模型的参数，从而在有限的显存总量下训练更大的模型。
   3. 另一种解决方案是采用低精度训练，如采用FP16或INT8、INT4等更低精度的训练，但低精度是有代价的
2. 模型并行虽然直观上解决了问题，即模型参数以及中间量的数据将会分布在不同GPU的内存中，但其不具有可扩展性，即在这些模型尺寸之外进一步扩展。模型并行通常对模型进行垂直的划分（即层与层之间进行划分）。模型划分本身上不具有问题，但其划分导致各模型部分之间的通信不再是在GPU内部的通信，而是转化为GPU与GPU的通信。在此认知基础上，如果单节点内部其IO或许不会成为GPU与GPU通信的瓶颈，但集群内部，节点与节点的GPU与GPU的通信，即使是高速互联网络也依旧无法满足大型模型的通信需求。论文描述，在GPU间通信带宽较高的单节点内部，他们能够很好的工作，但超过单个节点之后，效率将会迅速下降。
3. 解决方案如流水线并行，模型并行与CPU-Offloading，都需要在功能性，可用性与内存和计算/通信效率做出权衡，但这几者对于分布式训练而言都很重要，如果能够都要就没有什么好权衡的。

而后，基于上面的分析，为了回答如何克服现有的局限性从而高效训练大型模型，论文分析了现有系统在模型训练中内存消耗的全部情况，并将其分为两种类型的内存消耗：

1. 模型状态：大模型中，大部分内存被模型状态占用，包括Optimizer（优化器）状态（如Adam中的动量下降与方差等等中间值），以及梯度和参数（模型状态，计算中间值）
   因此有： 内存占用公式为$4\varphi +K\varphi$
   1. 文中以Adam为例，Adam需要存储两个优化器状态，时间平均动量与梯度方差计算更新。因此，使用Adam训练，需要有足够的内存来保存动量和梯度方差的副本。如果采用混合精度训练（其中参数和激活状态以 fp16 格式存储，从而能够在这些 GPU 上使用高吞吐量的张量核心单元。在混合精度训练过程中，前向传播和后向传播都使用 fp16 权重和激活值。不过，为了在后向传播结束时有效计算和应用更新，混合精度优化器会保留参数的 fp32 副本以及所有其他优化器状态的 fp32 副本），那么对于一个具有$\varphi$个参数的模型，其参数所需内存大小为$2\varphi$字节，梯度所需内存为$2\varphi$字节，Adam所需的参数、动量、方差的全精度所需内存分别为$4\varphi$字节，$4\varphi$字节，$4\varphi$字节，因此总体上需要$2+2+4+4+4=16\varphi$字节的内存
2. 残留状态（Residual states）：其指的是并非必须的量，例如激活值（Activation），它无需被存储在GPU中，可以通过重新前向计算（forward）而得出，临时缓冲区（temporary buffers），类似线程间通信所需的消息缓冲区，它存储从其他GPU发送而来的数据等待当前进程使用，此外是无法使用的内存碎片。
   1. 激活值：以序列长度为1K，批量大小为32的1.5B参数GPT-2模型进行训练，将会需要60GB的内存。可按照该计算公式验算：$12\times hidden\_dim\times batch\times seq\_length\times transformer\_layers$
      1. 文中提到了Activation checkpointing，激活值检查点方法是一种重新计算方法，它通过重新进行forward操作，降低了激活值的内存占用量为原先的总激活值的占用量的平方根，但增加了33%的重新计算开销。具体可以看这篇论文：Training deep nets with sublinear memory cost. CoRR, abs/1604.06174, 2016.
   2. 临时缓冲区：梯度all-reduce和正则化计算在执行计算之前将所有的梯度存储在一个缓冲区内部以提高吞吐量，但这样做，如果一个1.5B参数的模型，产生梯度矩阵大小为6B字节，因此该缓冲区实际上需要6GB的内存。
   3. 内存碎片：连续内存不足的问题。

针对这两部分，开发了ZeRO来优化这两部分的内存效率。

ZeRO

针对优化模型状态所需内存

模型状态在训练时占用了大部分计算内存，但现有的方法中，数据并行和模型并行并不能提供良好的效果。数据并行的计算与通信资源消耗较少但是其内存消耗大；模型并行对于计算与通信资源的消耗大但其内存消耗较少。

1. 数据并行在各个节点上使用相同的模型计算不同的数据，这意味着在数据并行处理过程中复制整个模型状态，导致更多的内存消耗；
2. 模型并行将模型状态进行分区，以获得较高的内存效率，这将导致计算过于细分（too finegrained computation）以及过大的通信资源的消耗，使得其可扩展性变低
   1. 计算过于细分：当模型被分割得非常细小时，每个计算设备可能只能执行很小一部分的计算任务。这可能导致计算任务的启动和结束开销（例如，任务调度和上下文切换）相对于实际计算时间来说变得相对较大，从而降低效率。

上述的方法中，在整个训练过程中都需要保存完整的模型状态，即使不是所有的模型状态都需要被使用。
在此基础上开发了ZeRO-DP，通过分区（partitioning）而非复制的方式，消除数据并行进程中的内存冗余，通过在训练过程中使用动态通信调度，保持了计算/通信的效率。ZeRO-DP主要考虑到：

1. DP 比 MP 有更好的扩展效率，因为 MP 在降低计算粒度的同时也增加了通信开销。超过一定程度后，较低的计算粒度会降低每个 GPU 的效率，而增加的通信开销则会阻碍跨 GPU 的可扩展性，尤其是在跨越节点边界时。相反，DP 既有较高的计算粒度，又有较低的通信量，因此效率要高得多。
2. DP 的内存效率低，因为模型状态在所有数据并行进程中都是冗余存储的。相反，MP 对模型状态进行了分区，从而提高了内存效率。
3. DP 和 MP 都会保留整个训练过程中所需的所有模型状态，但并非所有状态都是必需的。

它具有三个主要的优化阶段，对应优化器状态，梯度和参数的分区。

1. 优化器状态分区：提供4倍的内存减少
2. 梯度分区：额外提供2倍的内存减少
3. 参数分区：内存减少与数据并行阶数成线性关系，在64GPU之间进行分割，内存可减少64倍。通信量增加50%
   

ZeRO-DP

$P_{os}$Optimizer State Partitioning

对于一个$N_d$阶的数据并行，我们将优化器状态分割为$N_d$个相等的分区，因此第$i$个的数据并行进程只需要存储，更新对应的第$i$个分区的优化器状态。因此每个数据并行进程只需要更新总优化器状态的$\frac{1}{N_d}$，也只需要更新其中的$\frac{1}{N_d}$的参数。而当在数据并行过程中，通过实现all-gather原语来从所有的数据并行进程获得全部的参数更新。
问题：使用这样的方法提升了数据并行过程中多少的通信量？

• all-gather: 主要做了广播的操作，单个GPU向外广播的通信量大小为$(N_d-1)\frac{\varphi }{N_d}$，接收的通信量大小为$(N_d-1)\frac{\varphi }{N_d}$。

问题：现在总的内存消耗量为多少？

• $4\varphi +\frac{K}{N_d}\varphi$
  

$P_g$ Gradient Partitioning

由于每个数据并行进程只需更新一部分模型参数，因此其只需要对应的规约后梯度。因此，每一层的梯度在后向传播过程中可用时，只在负责更新相应参数的数据并行进程中规约部分的梯度。具体而言，对于一个$N_d$阶的数据并行，我们将每个进程上的模型也进行分区，那么对于前向传播之后每个模型分区产生的梯度进行梯度规约操作，对于每个数据并行进程对应的需要更新的参数分区对应的规约后梯度，可以通过梯度的reduce-scatter实现。

问题：使用这样的方法提升了数据并行过程中多少的通信量？

• reduce-scatter: 主要做了广播的操作，单个GPU向外单次传输的通信量大小为$\frac{N_d-1}{N_d}\varphi$，接收的通信量大小为$\frac{N_d-1}{N_d}\varphi$.

问题：现在总的内存消耗量为多少？

• $2\varphi +\frac{2+K}{N_d}\varphi$

计算和通信的重叠

由于梯度反向传播是一层一层产出的，因此我们可以将最先完成的分区先进行通信在节点间完成梯度的规约，在这个时候，节点内继续进行反向传播的计算，这样可以实现计算与通信的重叠，从而减少通信的时间。
*注意，上述所有的规约操作都是基于Ring-AllReduce场景下的规约场景，实际的规约操作并不一定是按照Ring-AllReduce的方法进行操作的，也可能采取了不同的规约方法

$P_p$ Parameter Partitioning

每个进程只存储与其分区相对应的参数。当其分区之外的参数需要用于向前或向后传播时，则通过广播从相应的数据并行进程接收。这似乎会产生很大的通信开销，但ZeRO的研究表明，这种方法只将基准 DP 系统的总通信量提高了 1.5 倍，同时还减少了与 Nd 成比例的内存。这很有质疑性，*此处未完待续。
问题：内存消耗为？

• $\frac{4+K}{N_d}\varphi$
  

针对残留状态所需内存

激活值、临时缓冲区和不可用的内存碎片所消耗的剩余内存可能会成为第二个内存瓶颈。

1. 对于激活值（为执行后向传递而存储的前向传递），我们注意到检查点[7]对大型模型有帮助，但还不够。因此，ZeRO-R 通过激活值分区来识别和消除现有 MP 方法中的激活值副本，从而优化激活值所需内存。它还能在适当的时候将激活值卸载到 CPU。
2. ZeRO-R 为临时缓冲区定义了适当的大小，以实现内存和计算效率的平衡。
3. 由于不同张量的生命周期不同，内存碎片将会出现。即使有足够的可用内存，碎片导致的连续内存不足也会导致内存分配失败。ZeRO-R 可根据张量的不同寿命主动管理内存，防止内存碎片。

在此基础上，开发了ZeRO-R，考虑到：

1. 激活值
   1. MP 对模型状态进行分区，但往往需要复制激活内存。例如，如果垂直拆分线性层的参数并在两个 GPU 上并行计算，则每个 GPU 都需要整个激活内存来计算其分区。
   2. 对于 GPT-2 或更大的模型，算术强度（每次迭代的计算量与每次迭代的激活检查点数量之比）非常大（≥ 10K），并随隐藏维度线性增加，因此即使带宽较低，也有可能隐藏激活检查点的数据移动成本。
2. 临时缓冲区
   1. 采用了固定的缓冲区大小
3. 内存碎片
   1. 内存碎片是短时和长时内存交错使用的结果。
      1. 在前向传播过程中，激活检查点是长效的，但重新计算的激活是短效的。
      2. 在反向传播中，激活梯度是短时的，而参数梯度是长时的。
   2. 基于这一认识，ZeRO 通过将激活检查点和梯度移动到预先分配的连续内存缓冲区，执行即时内存碎片整理。这不仅提高了内存可用性，还减少了内存分配器寻找空闲连续内存的时间，从而提高了效率。

ZeRO-R

$P_a$ Partitioned Activation Checkpointing

模型并行在设计上需要对激活值进行复制，从而导致在模型并行 GPU 上出现激活值的冗余副本。ZeRO 通过对激活值进行分区消除了这种冗余，每次只在计算中使用激活值之前以复制的形式将其具体化到一个激活层。更具体地说，一旦模型某层的前向传播计算完成，输入激活值就会被分割到所有模型并行进程中，直到反向推演过程中再次需要时才会被使用。此时，ZeRO 会使用全收集操作来重新生成激活值的副本。文章将这种优化称为$P_a$。它与激活检查点（activation checkpointing）结合使用，只存储分区的激活值检查点，而不是激活值检查点的完整副本。此外，在超大模型和设备内存非常有限的情况下，这些分区激活值检查点也可以卸载到 CPU，从而将激活值的内存开销降至几乎为零，但需要额外的通信成本

$C_B$ Constant Size Buffers

ZeRO 通过选择临时数据缓冲区的大小，以平衡内存和计算效率。在训练过程中，某些操作的计算效率与输入大小有很大关系，输入越大，效率越高。例如，大型All-Reduce操作的带宽远高于小型All-Reduce操作。因此，为了获得更高的效率，英伟达 Apex 或 Megatron 等高性能库会在应用这些操作前将所有参数融合到一个缓冲区中。不过，融合缓冲区的内存开销与模型大小成正比，可能会产生负面的作用。例如，对于 3B 参数模型，32 位融合缓冲区需要 12 GB 内存。为了解决这个问题，文章在模型过大时使用性能高效的恒定大小融合缓冲区。这样，缓冲区的大小就不取决于模型的大小，而且通过保持足够大的缓冲区，仍然可以实现良好的效率。

$M_D$ Memory Defragmentation

激活检查点和梯度计算会造成模型训练中的内存碎片。在带有激活检查点的反向传播过程中，只有被选中的激活点会被存储起来用于反向传播，而大部分激活点会被丢弃，因为它们可以在反向传播过程中被重新计算。这就形成了短时记忆（丢弃的激活）和长时记忆（检查点激活）的交错，导致记忆碎片化。同样，在后向传播过程中，参数梯度是长效的，而计算参数梯度所需的激活梯度和其他缓冲区是短效的。这种长短期内存的交错再次导致内存碎片化。
在内存充足的情况下，有限的内存碎片通常不是问题，但对于使用有限内存运行的大型模型训练，内存碎片会导致两个问题：

• 即使有足够的可用内存，也会因缺乏连续内存而出现 OOM；
• 内存分配器需要花费大量时间来搜索连续内存以满足内存请求，从而导致效率低下。

ZeRO 通过为激活检查点和梯度预分配连续的内存块，并在产生这些内存块时将其复制到预分配的内存中，从而即时进行内存碎片整理。$M_D$不仅能让 ZeRO 用更大的批次规模训练更大的模型，还能在内存有限的情况下提高训练效率。

ZeRO的结论

以DeepSpeed主流的大模型训练框架，为了高效使用集群内的资源，通过提升通信量来高度利用显存计算资源。其主要综合考虑，通信量、计算资源、环境成本和运行成本，实现大模型的高效训练。对于集群训练的大模型而言，通信资源并非其训练的瓶颈，因此为解决显存不足的问题，经常会牺牲一部分通信资源。随后出现的ZeRO++用于解决ZeRO通信量过大的问题，但其并不能做到既能够节省显存资源同时节省通信资源。ZeRO++中提出的方法中，hpZ（模型权重的分层分割存储）反而又抛弃了ZeRO的低显存，通过在每台机器上维护一个完整模型副本，也就是将ZeRO的权重分区方法不再使用，节省了ZeRO的跨机全收集/广播。ZeRO++最终的成果是将跨节点ZeRO产生的通信量降低了4倍。