ZeRO

分类: 训练优化

定义

ZeRO (Zero Redundancy Optimizer) 是 DeepSpeed 框架中的核心分布式训练优化策略，通过将模型状态（优化器状态、梯度、参数）分片到多个 GPU 上，消除数据并行中的内存冗余

标准数据并行：每个 GPU 存储完整的 $\{\theta, g, O\}$ （参数、梯度、优化器状态）

ZeRO Stage 1: 分片 $O$ ，内存降至 $\frac{|O|}{N} + |\theta| + |g|$

ZeRO Stage 2: 分片 $O + g$ ，内存降至 $\frac{|O|+|g|}{N} + |\theta|$

ZeRO Stage 3: 分片 $O + g + \theta$ ，内存降至 $\frac{|O|+|g|+|\theta|}{N}$

Stage 1-3 逐步增加分片粒度，Stage 3 可训练万亿参数模型

通过通信换内存：需要时从其他 GPU 收集完整参数/梯度

ZeRO-Offload: 将优化器状态卸载到 CPU 内存，进一步降低 GPU 内存需求

ZeRO-Infinity: 支持 NVMe SSD 卸载，可在单节点训练超大模型

与模型并行正交，可组合使用

ZeRO (2020, Rajbhandari et al.): 提出三阶段分片策略

ZeRO-Offload (2021): CPU 卸载扩展

ZeRO-Infinity (2021): NVMe 卸载扩展