Zero Infinity原理

如上图，一次加载一个layer的一个weights分片（一层的1/DP的参数量），Broadcast至所有rank，计算各自的梯度，再Reduce至其中一个负责的rank，offload存放至CPU Memory，释放GPU里的weights和梯度；之后再对本layer的下一个weights分片，重复该操作；

Zero-Offload是和Zero-2一起用的；

Zero-Infinity适合Zero-3一起用的；

1. Infinity offload engine

用来在GPU Memory、CPU Memory、NVMe(SSD)之间搬运模型states的组件；

Pinned memory管理器：尽量复用同一块pinned memory，以减少内存碎片；直接在pinned memory里计算（例如optimizer更新），省得搬运的时候又要copy一遍；

-----------

2. activations的CPU offload

将用于backward的activations，offload至CPU memory；

-----------

3. Memory-centric tiling

遇到大的operator，显存放不下整个weights矩阵。则将其分片，每次一片的offload加载、计算、offload结果、释放，顺序执行每个分片；（可以理解为一种更加灵活的model并行）

-----------

4. Bandwidth-Centric partitioning

Zero Offload，每个layer的weights，是整个被一个rank所own的；在要用之前，从这个rank加载该层weights，从CPU memory至GPU memory，只有该rank的带宽是忙碌的，其他rank的带宽空闲；

Zero Infinity，每个layer的weights，划分为DP个分片，每个rank只own一个分片；在要用之前，所有rank同时传输并做AllGather，这样所有rank的PCIe都忙碌起来；（我：这样需要每个rank放得下一整个layer的weights，浪费显存；可以一个分片一个分片的broadcast，结合双buffer去overlap计算和传输，每个rank同时只保留2个分片在显存里）

-----------

5. Overlap-Centric design

提前得知道计算图，才能提前知道operator执行顺序；

Weights预取：GPU忙着计算第i层operator时，后端忙着搬运：GPU->GPU:第i+1层operator，CPU->GPU:第i+2层operator，NVMe->CPU:第i+3层operator; 注意：动态图可能根据数值条件，改变operator执行顺序，造成prefech白整了。

Gradients放回：同理，在传输第i+1、i+2、i+3层operator时，GPU在计算第i层的operator；

-----------

6.PyTorch实现

对PyTorch的哪些module类的forward、backward进行hook(加decorator)，在hook里进行参数、梯度的搬运操作；

一上来并不把整个model加载至所有GPU（可能会撑爆显存）；而是hook了module的__init__函数，将该module分片放至各个rank上；