[SC 2020] ZeRO: Memory Optimizations Toward Training Trillion Parameter Models

Introduction

• Promblems with MP (Model Parallel) and DP (Data Parallel). MP 把模型垂直切分到不同设备上，在前向和反向传播时都需要卡间通信，对通信带宽要求比较高，如果使用多机训练，那么 GPU 利用率会非常低 (因此 MP 一般不会跨机)；DP 通信效率更高，仅需在反向传播时进行梯度聚合，梯度计算和梯度聚合可以 overlap 进而掩盖掉通信开销，但 DP 的问题是并不能降低单卡的显存占用，大模型训练往往单卡内存不够，因此无法使用 DP
• 针对上述问题，作者提出了新的大模型数据并行训练方法 ZeRO (Zero Redundancy Optimizer)，仅需少量通信开销就能极大地降低训练内存开销 (retains the training efficiency of DP while achieving the memory efficiency of MP)，并且能与模型并行结合使用

Method

ZeRO-DP

• 模型混合精度训练时 Model State 的显存占用可以分成 3 个部分：optimizer states (FP32 model parameters $4\Phi$ + FP32 Adam momentums $4\Phi$ + FP32 Adam variances $4\Phi$)、gradients (FP16 gradients $2\Phi$)、parameters (FP16 model parameters $2\Phi$)；其中 $\Phi$ 为模型参数量；ZeRO-DP 主要用于降低 Model State 的显存占用
• ZeRO-1/2/3 通过将这 3 部分存储到不同设备上，有效降低了单卡的内存占用
  

---

Preliminaries – DDP 通信量分析
DDP 训练中需要在反向传播结束时用 Ring All-reduce 汇聚梯度，Ring All-reduce 由 Ring Reduce-scatter 和 Ring All-gather 两个步骤组成，各有 $\Phi$ 的通信量，全部通信量为 $2\Phi$

ZeRO-1 – $P_{os}$ : Optimizer State Partitioning

• ZeRO-1 中，每个设备都保存有完整的 parameters 和 gradients，而 optimizer states 被切分为 $N_d$ 块存放在不同设备上，其中 $N_d$ 为设备数 (DP degree)；training step 结束时每个设备独立进行前向和反向传播计算出参数梯度，然后通过 reduce-scatter 汇聚多卡梯度使得 $i$-th device 得到 $i$-th block 对应的 gradients (梯度聚合的通信过程和梯度计算过程还可以 overlap 来进一步降低通信开销)，可以用其更新 $i$-th block 对应的 optimizer states 并得到更新后的 parameters，等到所有设备都更新完毕后，使用 all-gather 即可使得所有设备上都得到全部更新后的 parameters；注意到在这个过程里，通信量仍然由 reduce-scatter + all-gather 带来，因此ZeRO 的通信量相比 DDP 不变！
• Memory Saving. 单卡内存占用由 $16\Phi$ 降低到 $4\Phi +\frac{12}{N_d}\Phi$

ZeRO-2 – $P_{os+g}$ : Gradient Partitioning

• 可以发现在 ZeRO-1 中，更新 $i$-th block 对应的 optimizer states 其实只需要 $i$-th block 对应的 gradients，其余参数的 gradients 是冗余的，将其丢掉后就成了 ZeRO-2. ZeRO-2 在 ZeRO-1 的基础上继续将 gradients 切分为 $N_d$ 块存放在不同设备上；具体实现时，ZeRO-2 在计算梯度时进行 reduce 操作，假如计算的是 $i$-th block 对应的 gradients，那么就通过 reduce 把梯度汇聚到 $i$-th device 上，其余设备上的该部分 gradients 就可以直接丢掉；其余训练过程和 ZeRO-1 一致；总的来说，ZeRO-2 的通信量相比 DDP 仍然不变！
• Memory Saving. 单卡内存占用由 $16\Phi$ 降低到 $2\Phi +\frac{14}{N_d}\Phi$

ZeRO-3 – $P_{os+g+p}$ : Parameter Partitioning

• ZeRO-3 在 ZeRO-2 的基础上继续将 parameters 切分为 $N_d$ 块存放在不同设备上；这样在模型前向和反向传播时，各要做一次 all-gather 操作得到全部权重 (权重的 all-gather 通信过程和前向/反向传播的计算过程同样可以 overlap 来降低通信开销)，反向传播完之后再用 reduce-scatter 做梯度聚合，最后更新 optimizer states 和 parameters. 值得注意的是，由于各个设备只保留部分权重，因此就没必要再来一次 all-gather 做权重聚合了，这样 ZeRO-3 的通信量仅为 DDP 的 1.5 倍！
• Memory Saving. 单卡内存占用由 $16\Phi$ 降低到 $\frac{16}{N_d}\Phi$

---

• 下图对比了 DDP 和 ZeRO-1/2/3 的通信量
  

ZeRO-R

• 模型训练时 Residual State 的显存占用可以分成 3 个部分：activations、temporary buffers、un-usable memory fragments；ZeRO-R 主要用于降低 Residual State 的显存占用
  

$P_a$: Partitioned Activation Checkpointing

• ZeRO 使用 checkpointing 技术进一步降低激活值占用的显存，另外 Megatron-LM 做模型并行的时候需要在每个设备上都保存完整的 FFN / Attn 输入激活值，对此 ZeRO 做了进一步优化，把输入激活值 (checkpointed activation) 也切分到了不同设备上，在反向传播前需要先用 all-gather 得到完整的 checkpointed activation 进行 recomputation (Megatron-LM 的 Transformer layer 在前向/recomputation/反向计算时各需要两次 all-reduce，而对 checkpointed activation 分块仅需要多增加一次 all-gather，新增的通信量仅为原来的 $1/12$)；在极端情况下，还可以通过把 partitioned activation checkpoints 都 offload 到 CPU 来消除 activation memory overhead (i.e., $P_{a+cpu}$)
• 总的来说，通过 $P_a$，ZeRO 可以极大地增加训练时的 batch size，而数据并行在训练过程中总的通信量又是和 batch size 成反比，因此相当于是变相地降低了总的训练通信量
• Memory Saving. 通过将输入激活值分块，ZeRO 能够使得 activation checkpointing 占用内存减小到原来的 $\frac{1}{MP}$，其中 $MP$ 为模型并行的设备数

$C_B$: Constant Size Buffers

• 训练时存储中间结果的临时缓冲区也会占用大量内存，例如 gradient all-reduce 就会在缓冲区里存储 FP16 的聚合梯度，之后还会在缓冲区里转为 FP32 的梯度信息用于更新 optimizer states，并且由于很多操作都是 input size 越大效率就越高，因此出于性能的考虑，NVIDIA Apex or Megatron 等高性能库会尝试把所有参数都装在一个缓冲区里 (fused buffer) 执行上述操作以提升带宽占用或计算效率，这样会导致临时缓冲区大小与模型大小成正比，占用大量内存 (这种实现方式应该是内存占用 $20\Phi$ 的原因？)
• ZeRO 主要的目的是省显存，因此使用 performance-effcient constant-size fused buffer，这样可以避免模型太大时临时缓冲区占用过多显存

$M_D$: Memory Defragmentation

• 模型训练时前向和反向传播时不同张量的生命周期不同，这会导致模型训练时产生 fragmented memory. 具体来说，前向传播时采用 activation checkpointing，除了这些激活值以外的其他激活值所占用的内存在前向传播结束后都会被释放，而 pytorch 是按照激活值产生的时间顺序分配内存的，这会导致内存空间里 short lived memory (discarded activations) 和 long lived memory (checkpointed activation) 交替存储，而 short lived memory 释放后就会产生 memory fragmentation；同理，反向传播时参数梯度属于 long lived，而 activation gradients and any other buffers 都属于 short lived，这也会导致 memory fragmentation
• memory fragmentation 可能会使得大模型训练时即使有足够空间也无法分配出一段连续内存，也会降低 memory allocator 效率，使其必须花费更多时间来搜寻一段连续内存
• 为了解决上述问题，ZeRO 提前给 long lived memory (activation checkpoints and parameter gradients) 分配一段连续内存

References

• Rajbhandari, Samyam, et al. “Zero: Memory optimizations toward training trillion parameter models.” SC20: International Conference for High Performance Computing, Networking, Storage and Analysis. IEEE, 2020.
• code: https://github.com/microsoft/DeepSpeed
• doc: https://www.deepspeed.ai
• Zero 论文精读【论文精读】
• ZeRO: Zero Redundancy Optimizer，一篇就够了