Zero 1、2、3的原理

较完整且简明的ZeRO系列介绍：大模型并行训练技术（一）—— ZeRO系列 - 知乎 (zhihu.com)

动图演示：【大规模训练】Optimizer state sharding (ZeRO) - 掘金 (juejin.cn)

朴素的Zero-1、2、3，是指参数们全部在GPU显存上，不利用任何CPU-Memory-Offloading；

Baseline(上图第1行）：Mixed Precision训练；FP16的weight和gradient；FP32的optimizer状态(weight、2组参数）

Zero-1(上图第2行)：Gradient还是完整的AllReduce；更新完自己这部分weight后，新增了AllGather让所有rank都拿到最新的weight；

Zero-2(上图第3行)：Gradient分成N份，每个rank只保存自己那份；原来对完整梯度的AllReduce，变成了N次对1/N梯度的Reduce（也就是Reduce-Scatter)；其余同上；

Zero-3(上图第4行): Weight分成N份，每个rank只保存自己那份；每次forward、backward之前，都要把该层的weight进行AllGather，使得每个rank都有该层完整的weight，才能计算。

实际上，Gradient在每层完成weight更新之后，就可以删除释放显存了，同一时刻最多只保留一层即可。我不知道为什么Zero-1需要gradient占用跟weight同样大小的空间。对于Zero-2，我认为可以每次Reduce完1/N，不负责这1/N的其余N-1个rank，可以释放掉该1/N梯度；

通信开销(按照数据move量来统计）：

0. 原始版，每轮只有1次所有gradient的AllReduce；AllReduce的state-of-art实现：先Reduce-Scatter，再AllGather; 两者都是pipeline实现的

Reduce-Scatter: 两种实现方式：

        1. 每轮发送skip加1(1-bit quantization所用）：所有节点先向id-1发送(3, 0, 1, 2)号数据包；再向id-2发送(2, 3, 0, 1)号数据包；再向id-3发送(1, 2, 3, 0)号数据包，至此，0号节点包含所有0号包，1号节点包含所有1号包...加和即可；设S为每个节点上的所有数据，N为节点数，则，每个节点发送数据为(N-1)*S/N约等于S，每个节点接收数据为(N-1)*S/N约等于S

        2. Ring-AllReduce的前半部分，每轮发送skip都是1；

All-Gather: 也是两种方式，类似上面；每个节点发送、接收数据量，都为S；

综上，原始版的两者加起来，对每个节点，发送数据量和接收数据量都是S+S=2S;

2. Zero-2: 

        1. 先对Gradient做Reduce-Scatter; 同上，每个节点，发送、接收数据量，都是S；

        2. 每个节点用从所有节点接收到的“属于自己”的partial梯度，加和之后，更新“属于自己”的optimizer状态和weight；通信量为0；

        3. 每个节点，对“属于自己”的partial weight，进行AllGather；同上，每个节点，发送、接收数据量，都是S；

综上，Zero-2的1和3加起来，每节点的单向通信数据量是S+S=2S;

3. Zero-3:

        1. Forward时，每一层，每个节点要把所有weight都凑齐，所以是个AllGather通信；所有层，每个节点，总通信量为S；（实际中，每层的通信可以和上一层的计算，overlap起来）；

        2. Backward时，同上，weight总通信量为S；

        3. Backward时，同Zero-2，gradient的通信操作是Reduce-Scatter，总通信量为S；

综上, Zero-3的1、2、3加起来，每节点的单向通信数据量是S+S+S=3S;

其他memory开销：

1. Activation

- 可使用activation checkpointing(又叫gradient checkpointing); 增加33%的计算时间（1个Forward矩阵乘，对应2个Backward矩阵乘；这招增加1遍Forward，因此增加1/(1+2)=33%)

- GPU->CPU->GPU数据搬运，如果能被计算所overlap住，则可使用offload到主存的方式，不占用GPU显存；

- 如果使用模型并行，因为backward时，每个rank需要一份完整的activation；可以在forward时，每个rank只保存1/N的activation；backward时，AllGather拼成完整的；

- 我想到的更激进：可以每个rank只保存1/N的activation；用到一个1/N就Broadcast接收到一个1/N，用完即释放；每次用1/N的activation和1/N的weight做矩阵乘法，加和到上一个1/N*1/N的计算结果矩阵上，最终得到output activation的1/N结果(N是GPU数目)

以上所有方式，可混合使用；例如checkpointing、1/N保存、offload一起用，只保存lgK个层的activation，每个rank每层保存1/N，不在显存而是在主存保存；

2. Temporary buffer

all-reduce的buffer；gradient norm之前的buffer；这些操作要求数据放在连续内存里；

- 使用固定大小的大块memory；

3. 显存碎片

有时显存还剩30%，但无大连续块儿了，会因无可用memory而报错；

- on-the-fly显存数据搬运（把松散的搬运到一起）；

- 我想到的更激进的: 提前规划好，哪些tensor是long-live的，哪些是short-live的，各自放到各自的区域去；

官网详细配置参数：

ZeRO — DeepSpeed 0.14.3 documentation