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