DeepSpeed ZeRO-3

定义

DeepSpeed ZeRO-3是零冗余优化器（Zero Redundancy Optimizer，ZeRO）三阶段方案中的最后一个阶段，核心手段是把优化器状态、梯度与模型参数这三类模型状态全部按数据并行进程切分，任何一张卡上都不再保留完整的模型副本^[1]。按 ZeRO 的分级设计，第 1 阶段只切分优化器状态，第 2 阶段再加上梯度，第 3 阶段才把参数本身也纳入切分范围^[1]。

与张量并行、流水线并行这类会改变算子结构的模型并行方式不同，ZeRO-3 仍属于数据并行路线的优化，模型代码基本无需重构，因此能以较低改造成本换取显存占用下降^[1]。DeepSpeed 的 PyTorch 实现把它与 CPU、固态硬盘卸载组合成 ZeRO-3 Offload 等配置形式，通过配置文件开关即可启用^[1]。

原理

在混合精度加 Adam 优化的标准数据并行中，每个进程都要保存低精度参数、低精度梯度，以及包括 fp32 主权重、一阶矩和二阶矩在内的优化器状态，合计约每参数 16 字节^[2]。设模型参数量为 $\Psi$、数据并行度为 $N_d$，则未优化时单卡模型状态约为 $16\Psi$ 字节，而 ZeRO-3 把它降到 $(2+2+K)\Psi/N_d$ 字节，其中 $K$ 在 Adam 下取 12^[2]。

参数被切分后无法常驻显存，运行期改为按需重建：某个子模块执行前向计算时，先把分散在各进程上的参数分片聚合成完整张量，算完立即释放非本地分片；反向传播时再聚合一次以计算梯度，随后用归约散射把梯度分片送回各自负责的进程。聚合与释放通常以子模块的 forward 方法为粒度自动完成，因此第 3 阶段的通信被直接嵌入前向与反向路径，通信强度高于前两个阶段^[3]。

flowchart LR

A[参数分片存于各 GPU] --> B[前向: AllGather 重建本层参数]

B --> C[计算并释放非本地分片]

C --> D[反向: 再次 AllGather 参数]

D --> E[ReduceScatter 梯度分片]

E --> F[各进程用本地分片更新优化器状态]

F --> A

按 DeepSpeed 给出的通信量分析，第 3 阶段每个训练步除梯度归约外还要额外做参数广播，总通信量约为基线数据并行的 1.5 倍，而显存下降幅度与并行度成正比^[4]。为压住这部分开销，实现中引入参数预取、通信与计算重叠、按重要性决定参数常驻等机制，把额外的数据搬运尽量藏在计算之后^[5]。

发展历程

ZeRO 由微软研究院的 Samyam Rajbhandari、Jeff Rasley、Olatunji Ruwase 与 Yuxiong He 提出，论文《ZeRO: Memory Optimizations Toward Training Trillion Parameter Models》于 2019 年 10 月提交至 arXiv，后在 2020 年的国际高性能计算、网络、存储与分析会议（SC 2020）上发表^[6][7]。论文给出了三阶段的设计与显存、通信量分析，并报告在 400 块 GPU 上训练超过 1000 亿参数模型的结果^[6]。

工程落地由微软开源库 DeepSpeed 承担。该库在 2020 年 9 月提供 ZeRO-2 Offload，把第 2 阶段的优化器状态与梯度卸载到 CPU 内存；2021 年 3 月又发布 ZeRO-3 Offload，补齐参数分片与参数卸载能力，使 ZeRO 三个阶段的特性集中在同一套配置中^[1]。

此后围绕第 3 阶段的开销问题陆续出现改进工作，例如把 GPU、CPU 内存和本地固态硬盘统一编址的 ZeRO-Infinity，以及以减少通信量为目标的 ZeRO++，均以 ZeRO-3 的分片结构为基础扩展^[8]。在生态层面，PyTorch 官方的全分片数据并行（FSDP）采用了与 ZeRO-3 相近的全分片思想，两套实现可在同一模型规模上互相替代

应用

第 3 阶段主要用于单卡或单阶段方案装不下的大模型训练。按官方公布的数据，ZeRO-3 Offload 可在单块 GPU 上微调超过 400 亿参数的模型，在 512 块 V100 上训练超过 2 万亿参数的模型；在 1 万亿参数规模、512 块 V100 的条件下可持续 25 PFlops 的算力，折合每卡 49 TFlops，单卡吞吐相比 ZeRO-2 Offload 最高提升 2 倍^[1]。

在实际训练框架中，它通常以配置项形式被调用。Hugging Face 的 Transformers 生态通过 Accelerate 和训练参数指定 DeepSpeed 配置文件来切换阶段，模型工具链也提供 ZeRO-3 与 ZeRO-3 Offload 的现成样例，后者会把优化器状态、梯度、参数甚至激活值一并搬到 CPU^[9]。对大语言模型的全参数微调而言，第 3 阶段常被列为模型权重本身无法放入单卡时的选择^[10]。

分片也带来检查点保存方式的变化：直接用框架默认接口导出权重，各进程只能得到互不重叠的碎片，因此需要先聚合为完整权重或保存分片权重后再用配套脚本合并^[10]。

局限

第 3 阶段节省的是模型状态占用的显存，激活值、临时聚合缓冲区、通信工作区与显存碎片并不在其优化范围内，因此显存瓶颈若来自长序列带来的激活膨胀，单靠切换阶段无法解决^[11]。

额外通信是其固有代价。每个训练步要多做参数聚合，总通信量约达基线数据并行的 1.5 倍，在跨节点带宽有限时更容易成为瓶颈；有实测报告显示，在带宽受限的集群上 ZeRO-3 的算力利用率明显低于张量并行与流水线并行组合的方案^[12]。配置不当还可能出现实际性能下降 10% 至 20% 的情况^[10]。

实现上还有若干前提条件：子模块的前向与反向必须各自能装进设备显存，否则要借助分块线性层等内存中心化的切分手段；模块参数若在自身构造与 forward 之外被访问，属于外部参数，需要显式收拢或注册才能正确参与切分^[3]。此外，卸载模式与 PyTorch 原生的梯度检查点机制并不直接兼容，需要改用框架自身的激活检查点配置^[9]。当较低阶段已能避免显存耗尽时，采用较高阶段会额外付出通信成本，因而并非优化等级越高越合适^[11]。

参见

• DeepSpeed —— 提供 ZeRO 系列实现的微软开源深度学习优化库。

• 零冗余优化器 —— ZeRO-3 所属的三阶段显存优化技术体系。

• 全分片数据并行 —— PyTorch 中与 ZeRO-3 思路相近的原生实现。

• 张量并行 —— 与 ZeRO-3 常组合使用、在算子内部切分计算的并行方式。

• 数据并行 —— ZeRO-3 所优化的基础并行范式。

• 混合精度训练 —— 决定模型状态显存占用的训练方式。