文章目录
- ZeRO
- 实验
-
-
- 实验设置
- DeepSpeed ZeRO Stage-2 实验
- 性能比较
- 进一步优化
- DeepSpeed ZeRO Stage-3 和 CPU 卸载
- 结论
-
ZeRO
ZeRO(Zero Redundancy Optimizer)是一种用于分布式训练的大规模深度学习模型的优化技术。它通过分片模型状态(参数、梯度和优化器状态)来消除数据并行训练中的内存冗余,从而显著提高内存效率。ZeRO 分为三个阶段,每个阶段逐步增加内存优化程度。分别对应优化器状态、梯度和参数的划分。
-
优化器状态分区(P os)——内存减少 4 倍,通信量与数据并行相同
阶段 1:跨数据并行工作器/GPU 的分片优化器状态- 描述:在 ZeRO 的第一个阶段,优化器状态被分片到不同的数据并行工作器或 GPU 上,而不是在每个 GPU 上重复。比如,Adam 优化器的状态包括权重、动量和方差等,这些状态信息会被均匀分布到各个 GPU 上。
- 内存节省:这种方法可以实现大约 4 倍的内存节省,同时保持与传统数据并行相同的通信量。
- 应用场景:适用于训练参数在亿级别的模型,比如 GPT-2,能有效降低单个 GPU 的内存使用.
-
添加梯度分区(P os+g)——内存减少 8 倍,通信量与数据并行相同
阶段 2:分片优化器状态和梯度- 描述