如上图,一次加载一个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的组件;
2. activations的CPU offload
将用于backward的activations,offload至CPU memory;
3. Memory-centric tiling
遇到大的operator,显存放不下整个weights矩阵。则将其分片,每次一片的offload加载、计算、offload结果、释放,顺序执行每个分片;