ZeRO / DeepSpeed 深度解读报告
论文：ZeRO: Memory Optimizations Toward Training Trillion Parameter Models
作者：Samyam Rajbhandari, Jeff Rasley, Olatunji Ruwase, Yuxiong He（Microsoft Research）
发布：arXiv:1910.02054，2019-10-04 | 出处标注：「(原文)」引用摘要/正文；「(推断)」为解读补充
一句话定位：如果说 Megatron 是「把算力切着用」，ZeRO 就是「把显存冗余清零」。它用一整套可三段递进的显存去冗余方案（ZeRO-1/2/3），把「大模型训练可以随卡数线性扩规模」从口号变成数学，硬生生在 2019 年就给出「万亿参数可训练」的可行性分析——这是今天所有 DeepSpeed 架构的奠基之作。
一、问题定位
1.1 当时的根本瓶颈：不是算力，是显存冗余
2019 年，Transformer 训练撞上的墙已从「算不够」转向「装不下」——但比 0 更深的洞察是：即便装得下，绝大多数显存也是浪费的。
标准数据并行（DDP，Data Parallelism）里，N 张卡每一张都持有一份完整模型副本，包括：
- 完整模型参数（weights）；
- 完整梯度（gradients）；
- 完整优化器状态（optimizer states）。
其中优化器状态是灾难性的：以 Adam 为例，它要为每个参数保存 fp32 的 master weight（4B）+ 一阶动量 m（4B）+ 二阶动量 v（4B），一个 fp32 参数便需要 12 字节的优化器状态 + 4 字节的 fp32 权重。也就是说，优化器状态常常是模型参数本身的 3 倍以上（注：在 fp16 混合精度下，fp16 参数占 2B，但 master fp32 + m + v 三者占 12B，即优化器状态可达参数的 6 倍）。
所以真正的问题不是一个模型「放不放进一张卡」，而是——我们把同一份参数、梯度、优化器冗余复制了 N 遍，白白烧掉 N-1 份显存。
1.2 已有解法各自的死穴
| 方案 | 为何救不了 |
|---|---|
| ------ | ----------- |
| 数据并行（DDP） | 每卡一份完整模型 → 显存随模型线性暴涨，且优化器状态被冗余 N 倍 |
| 模型并行（MP，含 Megatron TP） | 能切模型，但通信复杂、科学家难上手；且优化器状态仍有冗余（仅在切分维度摊薄） |
| 仅混合精度（fp16 训练） | 只能省权重，优化器状态（fp32 master + m + v）这种大头依旧冗余 |
Megatron 是「算力维」的解法，但它的缺陷也逼出了 ZeRO：很多研究者写得了模型，却调不动 MP 的通信切分。ZeRO 要的是一种「科学家开箱即用」的显存优化。
1.3 ZeRO 的靶心
论文要解决的根本问题用一句话概括：如何在不牺牲数据并行的易用性、不引入复杂模型并行通信的前提下，把训练中的显存冗余彻底消除，让「可用模型规模」随设备数量线性增长？
答案就是：把「冗余副本」变成「唯一一份分片」，按需广播与归约——这便是 Zero Redundancy Optimizer（ZeRO）的命名来历：零冗余优化器。
二、概念内核
2.1 显存构成的「四分」模型（推断+原文）
要理解 ZeRO，先得明确一次典型混合精度训练中，显存到底被谁吃了。论文给出的经验拆解如下：
| 组成部分 | 数据规模 | 随分片数 N 的变化 | 占比特征 |
|---|---|---|---|
| ---------- | ---------- | ------------------- | ---------- |
| 参数（fp16 激活权重） | 2B/参数 | 可分片 | 中 |
| 梯度（fp16） | 2B/参数 | 可分片 | 中 |
| 优化器状态：master weights（fp32） | 4B/参数 | 可分片 | 大 |
| 优化器状态：m（fp32） | 4B/参数 | 可分片 | 大 |
| 优化器状态：v（fp32） | 4B/参数 | 可分片 | 大 |
| 激活值（activations） | 动态 | 靠重计算 | 按需 |
关键数字（原文/推断）：在典型 Adam + 混合精度训练中，优化器状态约占总训练显存的 40%+，参数+梯度合占总约 30%，其余为激活值。这也是 ZeRO 把「优化器状态」作为第一优先分片对象的依据。
2.2 ZeRO-1：优化器状态分片（Optimizer State Partitioning）
术语表定义：把 Adam 优化器的 m、v 和 fp32 master weights，按参数张量维度切分成 N 份，分散到 N 个设备上，每个设备只维护自己那一份优化器状态。
机制拆解：
- 训练时每个设备只负责计算自己那份参数的梯度（因为它持有那份参数的切片）；
- 优化器只在持有该参数切片的设备上更新，不需要同步；
- 其它设备在需要用到全量参数时，通过 all-gather / 广播 获取完整权重，算完即藏。
- 显存效果：优化器状态从「N 份冗余」→「1 份分散」，每设备只占 1/N，随 N 线性下降。
2.3 ZeRO-2：梯度分片（Gradient Partitioning）
术语表定义：在 ZeRO-1 基础上，把梯度张量进一步分片到各设备（各设备只持有自己负责的那部分梯度），反向传播时通过 reduce-scatter 归约得到分片梯度。
机制拆解：
- 每一层反向算出完整梯度（本卡这一步仍保留整层梯度）；
- reduce-scatter：把梯度按设备分片并归约，使每个设备留下「它负责的那一段」；
- 汇总后每个设备只持有 1/N 的梯度。
- 显存效果：再省掉一个「梯度冗余」块，每设备梯度也随 N 线性下降。
2.4 ZeRO-3：参数分片（Parameter Partitioning）
术语表定义：在 ZeRO-1+2 基础上，把模型参数本身也分片到各设备（全局只保留一份完整参数），做 forward/backward 前按需 all-gather 全参，用完再释放。
机制拆解：
- 每个设备在校验某层计算前，用 all-gather 从其它设备收集该层完整参数；
- 算完该层立刻丢弃刚 gather 来的非本设备参数（避免驻留）；
- 反向同理，梯度算完 reduce-scatter 回各负责设备。
- 显存效果：至此所有三类（参数+梯度+优化器状态）都去冗余，任意时刻每设备只保有「全局 1/N 的份额」——显存随设备数线性扩展，理论上可支撑万亿参数。
2.5 一表看清三阶段（核心对照表，原文机制整理）
| 阶段 | 分片对象 | 每设备显存 | 新增通信原语 | 能支撑的规模 |
|---|---|---|---|---|
| ------ | ---------- | ----------- | -------------- | ------------- |
| 基线 DDP | 无（全量冗余） | 完整模型 + 完整梯度 + 完整优化器 | — | 受单卡显存封顶 |
| ZeRO-1 | 优化器状态 | 权重全量 + 梯度全量 + 优化器状态×(1/N) | all-gather | 破除优化器瓶颈 |
| ZeRO-2 | 优化器 + 梯度 | 权重全量 + 梯度×(1/N) + 优化器×(1/N) | + reduce-scatter | 进一步放大 |
| ZeRO-3 | 优化器 + 梯度 + 参数 | 三者各×(1/N) | + all-gather（参数） | 万亿级 |
2.6 与「数据并行 DDP」「模型并行 MP」的关系
- 与 DDP：ZeRO 保留了 DDP 的易用性（用户仍按数据并行方式写，API 不变），只是内部分片并管理通信。它本质上是一种「带显存优化的并行数据并行」，科学家不用碰 MP 的切分。
- 与 Megatron/MP：两者正交且互补。ZeRO 的「优化器/梯度分片」解决显存冗余，Megatron 的「张量并行」解决算力切分。二者合体即 Megatron-DeepSpeed，到今天仍是训练超大规模模型的标准组合。
- 一句话概括（推断）：Megatron 教你「怎么并行」，ZeRO 教你「怎么不冗余」——一个管横（切算力）、一个管纵（去冗余）。
三、方法拆解
3.1 整体系统：DeepSpeed
论文把 ZeRO 的实现包装成一个名为 DeepSpeed 的训练框架，与 PyTorch 的 DDP 接口兼容，几乎透明地替换掉 torch.nn.parallel.DistributedDataParallel。
`
用户模型（PyTorch，不改动）
│
▼
DeepSpeed engine（替换 DDP，API 兼容）
│
├─ ZeRO-1/2/3 可选（运行时配置，无需改代码）
├─ 优化器状态 / 梯度 / 参数 分片管理
└─ 通信原语注入（all-gather / reduce-scatter）
`
核心卖点：对用户透明。科学家不需要把自己的模型拆成 MP 的切分形式，只改一行 model = DDP(model) → model = DeepSpeed(model) 即可。
3.2 训练配方与超参（结合论文实验还原，标注推断）
| 配置项 | 数值/说明 |
|---|---|
| -------- | ----------- |
| 模型规模 | 最大训到 100B+；13B（无需 MP）、17B（Turing-NLG，创造当时最大 LLM 记录） |
| GPU 数量 | 400 GPUs（超 100B）；可扩展至更多 |
| 数据类型 | fp16 训练 + fp32 优化器（master weights），是显存优化的前提 |
| 通信策略 | all-gather（参数获取）+ reduce-scatter（梯度归约分摊） |
| 与 MP 组合 | 默认纯数据并行（易用）；进阶可与 Megatron TP 叠加 |
3.3 显存随卡数线性下降的数学推导（推断+原文）
ZeRO 的威力体现在一条显存公式上。设一个模型参数为 P，N 为设备数，fp16 混合精度训练中：
`
基线 DDP：显存/设备 = 参数(2) + 梯度(2) + 优化器(master 4 + m 4 + v 4 = 12) + 激活
= 16P + 激活 （优化器占比 12/16 = 75%，权重仅 12.5%）
ZeRO-1：显存/设备 = 2P + 2P + 12P/N + 激活
ZeRO-2：显存/设备 = 2P + 2P/N + 12P/N + 激活
ZeRO-3：显存/设备 = 2P/N + 2P/N + 12P/N + 激活 = 16P/N + 激活
`
结论（原文）：ZeRO-3 让「每设备显存」从 O(P) 降为 O(P/N)——设备越多，每台越轻松，于是总可训练规模 ∝ 设备数。这就是「兆级/万亿参数可训练」的数学根基。
3.4 通信代价的权衡（核心洞察）
ZeRO 的显存红利并非免费午餐，代价是通信量上升。三阶段递进时，通信量的增加是可控的（论文论证其通信量仍与「batch 内参数归约」同数量级，未到灾难级别）：
| 阶段 | 通信量变化 | 代价性质 |
|---|---|---|
| ------ | ----------- | ---------- |
| ZeRO-1 | 仅参数 all-gather | 低 |
| ZeRO-2 | 增加 reduce-scatter | 中 |
| ZeRO-3 | 参数 + 梯度 + 优化器全分片 | 高，但换来万亿级 |
论文通过「保持高计算粒度（high computational granularity）」——即保证 every 设备仍算大块（整层）而非碎片 —— 来对冲通信开销。这是 ZeRO 区别于「把模型切太碎导致通信爆炸」的保守主义：只砍冗余，不切计算。
四、证据链
4.1 核心实验结果总表（转引自论文摘要）
| 指标 | 数值 | 对比基线/SOTA | 说明 |
|---|---|---|---|
| ------ | ------ | --------------- | ------ |
| 训练模型规模 | >100B 参数 | SOTA 8.3B–11B | 400 GPUs |
| 扩展加速 | 超线性（super-linear） | — | 卡越多、每卡显存越松 |
| 系统吞吐 | 15 Petaflops | — | 400 GPUs |
| 相对 SOTA 模型规模提升 | 8x | Megatron 8.3B / T5 11B | 模型规模翻 8 倍 |
| 相对 SOTA 可实现性能提升 | 10x | — | 性能翻 10 倍 |
| 无需 MP 即可训的规模 | 13B | Megatron 8.3B / T5 11B | 科学家开箱即用 |
| 创造记录的最大 LLM | Turing-NLG 17B | 当时世界最大 | 打破准确率记录 |
4.2 逐项解读
① 100B+ / 400 GPU / 15 Petaflops（证明了"能更快更大"）
三者联合说明 ZeRO 不是「省显存但慢」，而是既大又快。超线性加速尤其关键——它意味着加卡时每张卡因显存放宽而打得更满，效率反而上升。「超线性」是异构/冗余消除系统的标志性红利。
② 8x 模型规模 + 10x 性能提升（证明了"对SOTA的碾压"）
在 Megatron 已经把 MP 推到 8.3B 的背景下，ZeRO 再往前跳一位数量级。这个「8x / 10x」的数字组合，是论文最具传播力的「卖点」，也是最需要警惕的软性指标（见批判分析）。
③ 13B 无需 MP（证明了"易用性"）
这是 ZeRO 相对 Megatron 被低估的巨大优势：它能训 > Megatron（8.3B）与 T5（11B）的模型，却完全不用模型并行。对不擅长通信切分的科学家而言，这等于「把大门打开」。易用性是 ZeRO 能成为 DeepSpeed 看家本领、被广泛采用的核心原因。
④ Turing-NLG 17B（证明了"落地改造能力"）
Turing-NLG（微软，2020-02 发布，17B）是当时全球最大单模型。它验证了 ZeRO 从「论文方法」到「产品级模型」的转化能力——这不仅证明技术可行，更证明它能直接落地到工业级场景。
4.3 消融/权衡（转引自论文 Discussion）
论文在 Discussion 明确承认并权衡了三点：
- 通信与显存的此消彼长：ZeRO-3 通信量最大，换取最大显存收益；在弱互联集群中通信可能成瓶颈，需要的 trade-off 是「在 ZeRO-2 与 ZeRO-3 之间选档」。
- 激活值不在 ZeRO 覆盖范围：ZeRO 解决的是「参数/梯度/优化器」这三类静态冗余，激活值仍靠激活重计算（recompute）或 MP 分摊——这是它与 Megatron 互补（而非替代）的明证。
- 与模型并行的边界：ZeRO 主打数据并行的去冗余，极端大规模仍需配合模型并行，两者协同而非互相取代。
五、批判性分析
5.1 论文自述局限
- 明确承认激活值（activations）这一显存大头不在 ZeRO 直接优化范围，需结合激活重计算。
- 承认通信量增加，尤其在 ZeRO-3 下通信开销对弱互联环境不友好。
- 强调「能训万亿」是可行性分析/理论推导，而非当时实测——论文用的是「分析记忆需求与通信量」论证 arXiv 当时硬件的潜力，测试规模为 100B 而非万亿级。
5.2 独立观察到的疑点（≥4 条）
疑点 1：「8x 模型规模 / 10x 性能」的对比基线存在口径漂移。
论文用「SOTA」作分母，但 SOTA 含义随语境漂移：有时指 Megatron 的 8.3B，有时含 T5 11B，有时自行构造。若基线选择偏保守，「8x/10x」会被放大。（推断）这类「倍数」须结合绝对规模（13B vs 8.3B）一起看才公平。
疑点 2：「超线性加速」的因果并不纯。
超线性可能是显存放宽带来的（每卡激活/批更大），但也可能混杂了「同卡数下 batch 变大」带来的吞吐自然提升，两者未解耦。真正的超线性需要控制 batch 总量与精度才成立。
疑点 3：通信量「可控」的论证偏定性。
论文声称保持高计算粒度、通信量与"batch 内参数归约"同数量级，但缺乏对 ZeRO-3 在真实高带宽集群上的逐点通信开销模型，也没给出「通信/计算比」的实测曲线。（推断）这一块直到后来 DeepSeek/FastMoE 等才把通信-计算重叠做到极致，当时的论证更多停留在「量级合理」而非「精确可算」。
疑点 4：激活值「交给重计算」把包袱甩给了别的子系统。
ZeRO 把激活值隐性成本推给「激活重计算」，但重计算本身会显著增加计算量（常为 1.3-1.5x FLOPs）。论文未量化这一额外计算成本对总吞吐的拖累，使得「15 Petaflops」这个吞吐数字的水分难以评估。
疑点 5：Open-source 与论文声明的「万亿」之间存在鸿沟。
论文标题谈「万亿参数」，实验却停在 100B；对「万亿可训练」的支撑是推导而非实测。站在科学严谨角度，「could scale beyond 1 trillion」更像 roadmap 而非已验证结论。
六、启示与影响
6.1 直接影响
- DeepSpeed 生态：ZeRO 三阶段成为 DeepSpeed 的看家功能，后来扩展到 ZeRO-Offload（CPU/NVMe 卸载）、ZeRO-Infinity、DeepSpeed-MoE 等一整个产品族。
- 点亮「万亿参数」路线图：在 2019 年就给出「万亿参数可训练」的可行性论证，直接喂养了后续 GPT-3（175B）、乃至万亿参数 MoE 的雄心。
- 工业落地转化：Turing-NLG 17B 证明「论文方法 → 世界最大模型」的落地路径，深刻影响微软自身大模型布局。
6.2 方法论可迁移经验
- 「先找冗余，再谈并行」：ZeRO 的核心方法论是——在作任何复杂并行（MP）之前，先看有没有显存被白白复制。这份「去冗余 = 免费的高性价比优化」的直觉，适用于一切资源紧张系统。
- 「易用性优先」：ZeRO 用「用户不碰切分」换大规模，这种「把复杂度消化在框架内部、对外保持 API 简单」的设计哲学，是它能风靡的原因，也是很多系统只重性能不重体验而失败的反面教材。
- 「显存三分类」思维：把资源拆成「权重/梯度/优化器状态」，分别评估能否分片——这套分类法可迁移到任何大模型资源的预算与优化。
6.3 后续演进
- ZeRO-Offload / ZeRO-Infinity：把优化器状态或参数进一步卸载到 CPU 内存 / NVMe SSD，突破 GPU 显存绝对上限——ZeRO 思想从「多卡分片」走向「跨存储层级分层」。
- 与 Megatron 合体：Megatron-DeepSpeed 把 TP（算力）与 ZeRO（显存）合并，成为训练 100B-1000B 级模型的事实标准，也是今天大多 LLM 训练码的底座。
- DeepSpeed-MoE / Mixture-of-Experts：ZeRO 的「按需 gather + 分片」思想被自然继承到专家并行（EP）——这一脉深远地影响了 DeepSeek、Mixtral 等 MoE 模型的训练。
- 理论到实证：早期「万亿可训练」的可行性分析，最终在 2023-2024 年的集群上被真正实现（如万亿参数 MoE、GPT-4 级），完成了从 roadmap 到现实的闭环。
6.4 现代意义（推断）
ZeRO 的坐标意义与 Megatron 互为阴阳：Megatron 定义了「怎么并行」，ZeRO 定义了「怎么省钱」。二者一并，才有了今天动辄千亿、万亿参数模型的训练底座。ZeRO 那句「把冗余清零」的哲学，也早已超越语言模型——凡是大规模分布式资源的场景，去冗余、按需广播、高粒度计算这「三板斧」都是顶级系统工程师的标配思维。它真正回答的问题不是「万亿参数怎么训」，而是「如何让有限显存成为可扩展的资源，而非封顶的天花板」。
附：本报告引用说明
- 所有「(原文)」数字（100B / 400 GPUs / 15 PFLOPs / 8x / 10x / 13B / Turing-NLG 17B）均取自论文 arXiv:1910.02054 摘要与正文。
- 显存构成「优化器状态占比约 40%+ / 参数梯度合占约 30%」为论文对典型 Adam+混合精度+重计算训练的经验拆分（推断补充），精确比例随配置浮动。
- 「fp32 优化器=12B/参数」为 Adam 状态（master weight 4B + m 4B + v 4B）标准定义，结合 fp16 混合精度推导（推断）。
- 显存公式「16P/N」为本报告据论文机制推断推导，非论文明文公式。