ZeRO：一种去除冗余的数据并行方案
目前训练超大规模语言模型主要有两条技术路线：
- TPU + XLA + TensorFlow/JAX
- GPU + Pytorch + Megatron + DeepSpeed
前者由Google主导，由于TPU和自家云平台GCP深度绑定，对于非Googler来说并不友好
后者背后则有NVIDIA、Meta、MS等大厂加持，社区氛围活跃，也更受群众欢迎
另，上面提到的DeepSpeed的核心就是ZeRO(Zero Redundancy Optimizer)，它是一种显存优化的数据并行(data parallelism，DP)方案
ZeRO：论文链接：https://arxiv.org/abs/1910.02054
背景
如今训练大模型离不开各种分布式并行策略，常用的并行策略包括：
- 数据并行(Data Parallelism，DP)
假设有N张卡，每张卡都要保存一个模型，每次迭代(iteration/step)都将batch数据分隔成N个大小的micro-batch，每张卡根据拿到的micro-batch数据独立计算梯度，然后调用AllReduce计算梯度均值，每张卡在独立进行参数更新。
PS：
模型每张卡都存，数据切分，由每张卡单独计算
# https://huggingface.co/docs/transformers/parallelism#model-parallelism
# 假设模型有三层：L0, L1, L2
# 每层有两个神经元
# 两张卡
GPU0:
L0 | L1 | L2
---|----|---
a0 | b0 | c0
a1 | b1 | c1
GPU1:
L0 | L1 | L2
---|----|---
a0 | b0 | c0
a1 | b1 | c1
- 模型并行(Model Parallelism/Tensor Parallelism,MP/TP)
有的tensor/layer很大，一张卡放不下，将tensor分割成多块，一张卡存一块
如果模型的规模比较大，单个 GPU 的内存承载不下时，我们可以将模型网络结构进行拆分，将模型的单层分解成若干份，把每一份分配到不同的 GPU 中，从而在训练时实现模型并行。
训练过程中，正向和反向传播计算出的数据通过使用All gather或者All reduce的方法完成整合。这样的特性使得模型并行成为处理模型中大 layer 的理想方案之一。然而，深度神经网络层与层之间的依赖，使得通信成本和模型并行通信群组中的计算节点 (GPU) 数量正相关。其他条件不变的情况下，模型规模的增加能够提供更好的计算通信比。
# https://huggingface.co/docs/transformers/parallelism#model-parallelism
# 假设模型有三层：L0, L1, L2
# 每层有两个神经元
# 两张卡
GPU0:
L0 | L1 | L2
---|----|---
a0 | b0 | c0
GPU1:
L0 | L1 | L2
---|----|---
a1 | b1 | c1
- 流水线并行(Pipline parallelism,PP)
将网络按层切割，划分成多组，一张卡存一组。
流水线并行，可以理解为层与层之间的重叠计算，也可以理解为按照模型的结构和深度，将不同Layer分配给指定GPU进行计算。
相较于数据并行需要GPU之间的通信，流水线并行只需其之间点对点通讯部分activations，这样的特性可以使流水并行对通讯带宽的需求降到更低。然而，流水并行需要相对稳定的通讯频率来确保效率，这导致在应用时需要手动进行网络分段，并插入繁琐的通信原语。同时，流水线并行的并行效率也依赖各卡负载的手动调优。这些操作都对应用该技术的研究员提出了更高的要求。
# https://huggingface.co/docs/transformers/parallelism#model-parallelism
# 假设模型有8层
# 两张卡
====================== =====================
| L0 | L1 | L2 | L3 | | L4 | L5 | L6 | L7 |
====================== =====================
GPU0 GPU1
# 设想一下，当GPU0在进行（前向/后向）计算时，GPU1在干嘛？闲着
# 当GPU1在进行（前向/后向)计算时，GPU0在干嘛？闲着
# 为了防止”一卡工作，众卡围观“，实践中PP也会把batch数据分割成
# 多个micro-batch，流水线执行
流水线并行：
为什么需要ZeRO
上述三种并行方式中，数据并行因其易用性，得到了最为广泛的应用。然而，数据并行会产生大量冗余Model State的空间占用。
ZeRO的本质，是在数据并行的基础上，对冗余空间进行深度优化。
PS：
大模型训练中的显存占用可以分为Model State和Activation两部分
1、Model State
- 优化器状态(Optimizer States)：是Optimizer 在进行梯度更新时所需要用到数据。一些优化器(如Adam)需要存储额外的状态信息，如梯度的移动平均值和平方梯度的移动平均值。例如SGD中的Momentum亦即使用混合精度训练时的Float32 Master Parameters
- 模型参数(Model Parameters)：存储在显存中的模型权重和偏置项
- 梯度(Gradients): 在反向传播过程中计算得到的梯度，用于更新模型参数。其决定了参数的更新方向
它们三个简称OPG，其中优化器状态会占据大约2倍参数量的显存空间，这取决于选择的优化器，也是整个训练中占据最大空间的部分。
2、 Activation - 中间激活值(Intermediate Activations)： 在前向传播过程中，神经网络的每一层会产生中间激活值，这些激活值需要再反向传播过程中用来计算梯度
- 输入数据(Input Data): 批处理中输入数据也占用显存，尤其是当批处理较大时。
在传统数据并行下，每个进程都使用同样参数进行训练。每个进程也会持有对 Optimizer States 的完整拷贝，同样使用了大量显存。在混合精度场景下，以参数量为\(\Psi\) 的模型和Adam Optimizer为例，Adam需要保存：
- Float16 的 参数 和 梯度 的备份，这两项分别消耗了为\(2\Psi\) 和 \(2\Psi\)内存。(1 Float16 = 2Bytes)
- Float32 的 参数， Momentum， Variance备份，对应到3份 \(4 \Psi\) 的内存空间。(1 Float32 = 4Bytes)
终需要\(2\Psi + 2\Psi + K\Psi = 16\Psi\)Bytes 的显存。
一个7.5B参数量的模型，就需要至少120GB的显存空间才能装下这些Model Stats。当数据并行时，这些重复的Model State会在N个GPU上复制N份
ZeRO则在数据并行的基础上，引入了对冗余Model States的优化。使用ZeRO后，各个进程之后只保存完整状态的1/GPUs，互不重叠，不再存在冗余。在本文中，就以7.5B参数量的模型为例，量化各个级的ZeRO对于内存的优化表现。
ZeRO的三个级别
相比传统数据并行的简单复制，ZeRO通过将模型的 参数、梯度 和 Optimizer State划分到不同的进程来消除冗余的内存占用
ZeRO 有三个不同级别，分别对应Model States 不同程度的分割(Partition)：
- ZeRO-1： 分割Optimizer States
- ZeRO-2： 分割Optimizer States 与 Gradients
- ZeRO-3： 分割Optimizer States、Gradients 与 Parameters；
ZeRO-1
模型训练中，正向传播和反向传播并不会用到优化器状态，只有在梯度更新的时候才会使用梯度和优化器状态计算新参数。因此每个进程单独使用一段优化器状态，对各自进程的参数更新完之后，再把各个进程的模型参数合并形成完整的模型。
假设我们有