【深度学习】DeepSpeed，ZeRO 数据并行的三个阶段是什么？

ZeRO

ZeRO（Zero Redundancy Optimizer）是一种用于分布式训练的大规模深度学习模型的优化技术。它通过分片模型状态（参数、梯度和优化器状态）来消除数据并行训练中的内存冗余，从而显著提高内存效率。ZeRO 分为三个阶段，每个阶段逐步增加内存优化程度。分别对应优化器状态、梯度和参数的划分。

1. 优化器状态分区（P os）——内存减少 4 倍，通信量与数据并行相同
   阶段 1：跨数据并行工作器/GPU 的分片优化器状态

   • 描述：在 ZeRO 的第一个阶段，优化器状态被分片到不同的数据并行工作器或 GPU 上，而不是在每个 GPU 上重复。比如，Adam 优化器的状态包括权重、动量和方差等，这些状态信息会被均匀分布到各个 GPU 上。
   • 内存节省：这种方法可以实现大约 4 倍的内存节省，同时保持与传统数据并行相同的通信量。
   • 应用场景：适用于训练参数在亿级别的模型，比如 GPT-2，能有效降低单个 GPU 的内存使用.

2. 添加梯度分区（P os+g）——内存减少 8 倍，通信量与数据并行相同
   阶段 2：分片优化器状态和梯度

   • 描述