请介绍 DeepSpeed 中 ZeRO Stage-2 在前向和反向传播中的计算与通信逻辑。
考察说明
考察对 ZeRO 分片策略在模型并行训练中具体执行流程的理解
回答思路
- 准确描述 ZeRO Stage-2 对优化器状态和梯度进行分片，而非参数
- 清晰说明前向传播中如何收集（all-gather）完整的模型参数并计算
- 清晰说明反向传播中如何利用分片梯度进行局部规约（reduce-scatter）并更新优化器状态
- 能够说明通信量与计算量的权衡，以及相比 Stage-1 和 Stage-3 的差异