Zero原理

DeepSpeed: Extreme-scale model training for everyone - Microsoft Research

Why model-parallelism is most communication cost?

模型参数在CPU上有一份（只保留自己这个GPU的partition即可）；update的时候直接把offload到CPU上的gradients以及常驻CPU内存的optimizer-states, 更新到同样在CPU内存里的模型参数，然后把模型参数swap到GPU上，最后要调GPU的allgather来让所有GPU上的weight-partition都gather到所有GPU上；

Forward&Backward, gradient offload(GPU->CPU), update model, swap(CPU->GPU), 这几个环节之间可以pipeline/overlapped执行！

• Gradients: Backward的时候，每算出来一部分(partition)，就reduce到某一个GPU里去；该GPU再offload到CPU内存里去；update model时直接在CPU内存里计算；
• Optimizer states: 每个进程存自己的partition; 常驻CPU内存；update model时既做输入，也做输出；
• Model wegiths: 整个模型常驻GPU显存（我认为，可以每次swap进来几层即可（ZeRO-Infinity?））；每个进程把自己那份partition常驻在CPU内存；

GPU上的是FP16的model、activation、gradient；

CPU上的是FP32的model、gradient、optimizer-states;

CPUßàGPU的数据传输，全部是FP16的，省带宽速度快！

Activation checkpoint: Backward时需要用到的activation，可以checkpoint到CPU内存保存一会儿，Backward用的时候再swap到GPU显存；

以上Zero Offload就是Zero-3

分层来取，是Zero Infinity

DeepSpeed 本质上是一种“节省显存”的数据并行

DeepSpeed 假设了单层参数量可以在单张显卡上放得下，如果不满足这个假设，那么仍然需要使用模型并行，而且 DeepSpeed 的模型并行是通过调用 Megatron 来实现的。
根据 NVIDIA 最新的那篇 论文 （也是下面本文重点要介绍的），Megatron 在大规模训练的效率是超过 DeepSpeed 不少的。
而 DeepSpeed 的论文一直强调：我可以用更少机器训练更大的模型，但没有突出过在效率上的优势。
DeepSpeed 后来又出了一篇论文： ZeRO-Infinity ，当单层参数量在单张显卡上放不下的时候，它通过对这一层算子切片，一片一片来执行，使得单卡也能跑得起来一个巨大的层，可以理解成一种 “时间”轴上展开的模型并行。