ZeRO 显存优化

定义

ZeRO(Zero Redundancy Optimizer)是 DeepSpeed 提出的显存优化技术。传统数据并行让每块 GPU 保存完整的模型参数、梯度和优化器状态,存在大量冗余;ZeRO 把这三类状态切分到所有并行进程,按需通信聚合。

三个阶段

Stage 1 分片优化器状态(Adam 的动量与方差在混合精度下占显存大头);Stage 2 再分片梯度;Stage 3 连模型参数也分片,前向反向时临时 all-gather。理论上 N 卡可把单卡状态显存降至约 1/N,代价是通信量增加。

扩展与影响

ZeRO-Offload 将状态卸载到 CPU 内存,ZeRO-Infinity 进一步用上 NVMe,号称打破 GPU 显存墙。PyTorch 原生的 FSDP 即是 ZeRO Stage 3 思想的官方实现。ZeRO 大幅降低了大模型训练门槛,是数据并行路线上最重要的显存技术。