ZeRO-1/2/3（Zero Redundancy Optimizer 三阶段）

3 秒看懂

ZeRO 是微软 DeepSpeed 团队提出的分布式数据并行内存优化技术（OSDI 2020），核心思想极简：标准数据并行中每个 GPU 都持有完整的模型副本（优化器状态+梯度+参数），ZeRO 把这些状态分片到各 GPU 上，按需 AllGather 拼回。三个阶段逐步加大分片范围，内存从”每 GPU 持有全量”下降到近似”全量/N”。

N = 数据并行度（参与 ZeRO 分片的 GPU 数量）；具体倍数取决于优化器类型与精度配置，此处以混合精度 + Adam 为例的典型估算。

3 分钟产业解释

为什么 ZeRO 存在？

大模型训练的核心矛盾：单卡装不下，模型并行又太复杂。

传统的分布式训练方案主要有两条路：

• 模型并行（Tensor/Pipeline Parallelism）：把模型切开放到不同 GPU。精度无损，但需要大量侵入式代码改造，通信拓扑设计复杂。
• 数据并行（Data Parallelism, DP）：每个 GPU 持有完整模型，各算各的梯度，再 AllReduce 同步。实现简单，但每个 GPU 的内存冗余极大。

以混合精度训练 + Adam 优化器为例，每个参数需要：

此处 16B/param 是混合精度（fp16 参数/梯度 + fp32 优化器状态）+ Adam 的经典拆解，见 ZeRO 原始论文（Rajbhandari et al., OSDI 2020）。如果使用 SGD+Momentum，优化器状态部分减半。

一个 175B 参数的模型（GPT-3 量级）→ 每个 GPU 需要 175B × 16B = 2,800 GB，远超任何单 GPU 显存。即使只考虑 7B 参数，也需要 112 GB，超过单张 80GB A100 的容量。

ZeRO 的回答是：既然 DP 中每个 GPU 都存了完整的 16Φ 字节，何不把它切分成 N 份，每个 GPU 只存 16Φ/N？需要用时再 AllGather 回来。

产业价值：ZeRO 让数据并行也能训练远超单卡显存的模型，且不需要侵入式修改模型代码。它与模型并行（Tensor/Pipeline Parallelism）可以叠加使用，是当前大模型训练基础设施的关键拼图之一。

15 分钟专家深入

三阶段机制详解

设模型参数量为 Φ，数据并行度为 N_d（参与 ZeRO 分片的 GPU 数），混合精度 + Adam 训练。

Stage 1：分片优化器状态（ZeRO-OS）

分片对象：Adam 的一阶动量、二阶动量、fp32 主权重副本（共 12Φ 字节）。

机制：

1. 每个 GPU 只保存 1/N_d 的优化器状态（对应 1/N_d 的参数切片）
2. 前向/反向传播时，每个 GPU 使用本地持有的 fp16 参数（完整参数仍在每张卡上）
3. 反向传播完成后，每张卡得到完整的本地梯度（fp16）
4. ReduceScatter：各卡将梯度按参数切片做规约，每个 GPU 拿到属于自己负责的那 1/N_d 参数切片的聚合梯度
5. 用聚合后的梯度更新本地负责的那部分 fp32 优化器状态和 fp32 主权重
6. 将更新后的 fp32 主权重截断为 fp16，写回本地 fp16 参数对应切片
7. AllGather：各卡对更新后的 fp16 参数切片执行 AllGather，所有 GPU 获得完整的更新后 fp16 参数

内存：

Stage 1 per-GPU memory = 2Φ + 2Φ + 12Φ/N_d = 4Φ + 12Φ/N_d
         (params)  (grads)  (opt states)

通信开销：与标准 DP 的 AllReduce 相同（AllReduce = ReduceScatter + AllGather，总通信量 2Φ/GPU）。Stage 1 没有额外通信。[来源：ZeRO 原始论文]

Stage 2：分片优化器状态 + 梯度（ZeRO-OS+G）

在 Stage 1 基础上进一步分片：梯度也按参数切片分区，每个 GPU 只保存属于自己的 1/N_d 参数切片对应的梯度。

机制：

1. 前向传播：每张卡持有完整 fp16 参数，正常计算
2. 反向传播：每张卡计算本地梯度（完整 Φ 个参数的梯度）
3. ReduceScatter：各卡梯度按切片规约，每个 GPU 只保留自己负责的 1/N_d 切片的聚合梯度，其余梯度立即释放
4. 用聚合梯度更新本地优化器状态
5. 更新 fp16 参数

内存：

Stage 2 per-GPU memory = 2Φ + (2Φ + 12Φ)/N_d = 2Φ + 14Φ/N_d
         (params)  (grads+opt states)

通信开销：仍与标准 DP AllReduce 等价。ReduceScatter 聚合梯度（Φ）+ 后续参数更新后隐式同步（Φ），总 2Φ/GPU。[来源：ZeRO 原始论文]

关键洞察：Stage 2 的核心优势在于释放了 (N_d-1)/N_d 的梯度内存。对于大模型，梯度内存占比与参数内存相同（都是 2Φ），这一步节省非常显著。

Stage 3：分片一切（ZeRO-OS+G+P）

终极形态：参数也按切片分片，每个 GPU 只保存 1/N_d 的 fp16 参数。

机制：

1. 前向传播前：AllGather 拼回当前层（或整个模型）的完整 fp16 参数 → 计算 → 释放非本切片参数
2. 反向传播前：再次 AllGather 拼回完整参数 → 计算梯度 → 释放非本切片参数
3. ReduceScatter：梯度按切片规约 → 每卡只保留 1/N_d 聚合梯度
4. 更新本地优化器状态和参数

内存：

Stage 3 per-GPU memory = (2Φ + 2Φ + 12Φ)/N_d = 16Φ/N_d
             (所有状态均分片)

通信开销：高于标准 DP。每张卡每步额外需要两次 AllGather（前向 + 反向各一次）来拼回参数：

Stage 3 通信量 = AllGather_fwd(Φ) + AllGather_bwd(Φ) + ReduceScatter(Φ) = 3Φ/GPU
标准 DP 通信量 = AllReduce(Φ) = ReduceScatter(Φ) + AllGather(Φ) = 2Φ/GPU
通信比 = 3Φ / 2Φ = 1.5×

[来源：ZeRO 原始论文，§4.3]

三阶段内存公式汇总

设 N_d 为 ZeRO 并行度，混合精度 + Adam（16B/param）：

N_d 越大，Stage 3 的内存优势越极端。但 Stage 3 也引入了 1.5× 通信量和更频繁的 AllGather，实际吞吐受通信带宽制约。

ZeRO-Offload 与 ZeRO-Infinity

ZeRO 框架的后续扩展（同团队）：

• ZeRO-Offload（2021）：将 fp32 优化器状态和部分计算 offload 到 CPU 内存和 CPU 计算，使单 GPU 也能训练大模型。代价是 CPU-GPU PCIe 带宽瓶颈。
• ZeRO-Infinity（2021）：进一步扩展到 NVMe SSD，实现”无限”模型规模（论文标题字面含义）。利用 GPU 显存 + CPU 内存 + NVMe 的层次化存储。

ZeRO++

微软后续工作，主要优化 Stage 3 的通信效率：

• qwZ（Quantized Weights ZeRO）：对 AllGather 传输的权重做量化（如 int8/int4），减少通信量
• hpZ（Hierarchical Partitioning ZeRO）：在节点内和节点间做两级分片，优先在节点内（NVLink 高带宽）做 AllGather，减少跨节点通信

[注：ZeRO++ 的具体量化精度和性能数据，以 DeepSpeed 官方 GitHub 为准]

技术原理（最深）

数据并行中的内存冗余问题

标准数据并行（DP）的工作流：

┌─────────────────────────────────────────────────┐
│                  标准 Data Parallel              │
│                                                  │
│  GPU-0          GPU-1          ...    GPU-(N-1)  │
│  ┌──────┐      ┌──────┐            ┌──────┐     │
│  │完整Φ │      │完整Φ │            │完整Φ │     │
│  │完整G │      │完整G │            │完整G │     │
│  │完整OS│      │完整OS│            │完整OS│     │
│  └──────┘      └──────┘            └──────┘     │
│       ↓              ↓                  ↓        │
│     AllReduce(g)    AllReduce(g)    AllReduce(g) │
│       ↓              ↓                  ↓        │
│    本地更新         本地更新          本地更新      │
│                                                  │
│  Φ=参数, G=梯度, OS=优化器状态                    │
│  冗余: 每张卡存 16Φ, N张卡共存 16Φ×N             │
└─────────────────────────────────────────────────┘

ZeRO 的核心洞察：AllReduce 之后每张卡都有相同的聚合梯度，用于更新相同的参数和优化器状态 → N-1 份是完全冗余的。

ZeRO Stage 3 的 AllGather 与释放机制

以一个简化的两层模型（Layer1, Layer2）为例，Stage 3 中每张卡只存 1/N 的参数分片：

Forward Pass (以 GPU-0 的视角):

  ┌─ AllGather Layer1 params ──→ 得到完整 L1 参数 ─┐
  │                                                  │
  │  Compute Layer1 forward                          │
  │                                                  │
  ├─ 释放非本切片的 L1 参数分片                        │
  │                                                  │
  ├─ AllGather Layer2 params ──→ 得到完整 L2 参数 ─┐ │
  │                                                │ │
  │  Compute Layer2 forward                        │ │
  │                                                │ │
  └─ 释放非本切片的 L2 参数分片                      ┘ ┘

Backward Pass (类似):

  AllGather Layer2 params → Compute L2 backward → 释放
  AllGather Layer1 params → Compute L1 backward → 释放
  ReduceScatter gradients → 保留 1/N_d 聚合梯度

注意：在实际实现（DeepSpeed）中，AllGather 可以按层粒度或按模型整体粒度执行。按层粒度可以降低峰值内存（不需要同时拼回所有层的参数），但增加通信轮次。[具体粒度实现细节见 DeepSpeed 源码]

AllReduce vs ReduceScatter + AllGather 的通信等价性

标准 AllReduce = ReduceScatter + AllGather

ReduceScatter:  每个 GPU 贡献自己的数据，收到一个切片的聚合结果
                通信量: Φ/GPU（在环形拓扑下）

AllGather:      每个 GPU 贡献一个切片，收到完整的聚合数据
                通信量: Φ/GPU（在环形拓扑下）

总计: 2Φ/GPU

ZeRO-1/2 的 ReduceScatter 不是"额外"的——它替代了 AllReduce 的前半段，
而后半段 AllGather 在参数更新后以不同形式发生。
总通信量不变，仍为 2Φ/GPU。

通信量对比（Stage 3 vs 基线 DP）

标准 DP:
  Forward:  无通信
  Backward: AllReduce = ReduceScatter(Φ) + AllGather(Φ) = 2Φ/GPU
  总计: 2Φ/GPU

ZeRO Stage 3:
  Forward:  AllGather(Φ)         ← 拼回完整参数
  Backward: AllGather(Φ)         ← 再次拼回（因前向后已释放）
            + ReduceScatter(Φ)   ← 聚合梯度
  总计: 3Φ/GPU

通信比: 3Φ / 2Φ = 1.5×

技术演进史

Megatron-DeepSpeed 被用于训练 BLOOM-176B（BigScience 项目），这是 ZeRO + 模型并行组合在公开大模型训练中的标志性案例。

技术路线对比（量化表）

ZeRO 三阶段横向对比

ZeRO vs 其他并行策略

注意：上表中 ZeRO 和 Tensor/Pipeline Parallel 不是互斥的，实际大规模训练中通常组合使用。常见组合：Tensor Parallel（节点内 NVLink）+ Pipeline Parallel（跨节点）+ ZeRO（数据并行维度）。

上下游

上游（ZeRO 依赖什么）

下游（ZeRO 支撑什么）

关键指标

内存指标

通信指标

吞吐指标

供需与市场数据

需求侧

供给侧

市场规模参考

大模型训练基础设施市场规模尚无权威公开数据。根据公开信息粗略估算：2024 年全球 AI 训练 GPU 市场规模超过 300 亿美元 [估算，来源：综合 NVIDIA 财报及行业分析]，ZeRO/FSDP 等内存优化技术是该市场的效率乘数——它们不直接产生收入，但决定了同样的硬件预算能训练多大的模型。

代表公司与资本映射

核心贡献方

间接受益方

投资逻辑

核心观点

1. ZeRO 是基础设施层的”隐性标准”：不是独立产品，而是嵌入 DeepSpeed/PyTorch FSDP 等框架的核心技术。投资机会在使用 ZeRO 的公司（训练大模型），而非 ZeRO 本身。

2. 内存优化技术持续迭代：随着模型规模继续增长，ZeRO 的后续版本（++、Infinity 等）和其他新技术（如梯度检查点、选择性激活重计算、CPU offload）将持续演进。这是一个持续投入、持续回报的技术方向。

3. 与硬件协同演进：HBM 容量增长（HBM3e 36GB/stack → 未来更大）+ ZeRO 内存优化 = 更大的可训练模型。两者是乘法关系，而非替代关系。

4. 竞争格局：DeepSpeed（Microsoft）和 PyTorch FSDP（Meta）是两强格局。NVIDIA 在 Megatron 侧有自己的并行策略但与 ZeRO 互补。Google 的 JAX 走 GSPMD 自动分片路线。

风险提示

• 如果未来硬件显存足够大（如 HBM 容量 × 数量远超模型需求），ZeRO 的价值边际递减
• 新的训练范式（如状态空间模型、稀疏训练）可能改变内存瓶颈的位置
• DeepSpeed 的 Microsoft 属性可能影响其他公司的采纳意愿（FSDP 的中立性优势）

常见误读纠偏

❌ 误读 1：ZeRO Stage 3 与张量并行（Tensor Parallelism）是同一回事

纠偏：两者本质不同。

关键区别：ZeRO 切的是”冗余”——每个 DP 进程的完整参数副本；Tensor Parallel 切的是”必须”——单层内权重矩阵的计算维度。

❌ 误读 2：ZeRO Stage 3 的通信量是 Stage 1/2 的 1.5 倍，所以性能差很多

纠偏：

1. 通信量 ≠ 通信时间。实际通信时间取决于通信/计算重叠程度。在大模型中，计算时间远大于通信时间（计算密集），额外的 AllGather 可以被计算掩盖（overlap）。
2. 内存节省 → 更大 batch size → 更高 MFU。Stage 3 节省的内存可以用于增大 batch size，提升计算效率，部分甚至完全抵消额外通信开销。
3. 通信量是相对于模型参数 Φ 而言的，而非相对于总数据量。对于大模型（Φ 很大），计算量（约 6Φ × sequence_length × batch_size）远大于通信量（3Φ），通信/计算比很低。

❌ 误读 3：ZeRO 只适用于预训练

纠偏：ZeRO 适用于任何需要梯度更新的分布式训练场景：

• 全参数微调（Full Fine-tuning）：与预训练相同，ZeRO 直接适用
• LoRA / QLoRA：这些方法本身减少了可训练参数，但 ZeRO 仍然可以优化基础模型参数的内存占用（即使 frozen 参数不更新，仍需显存存储）
• RLHF：PPO 训练中同时存在 policy model、reference model、reward model、value model，内存压力极大，ZeRO 可以优化每个模型的状态内存
• 推理（有限适用）：ZeRO-Infinity 的 offload 思路可用于推理时的模型加载

❌ 误读 4：FSDP 和 ZeRO-3 完全等价

纠偏：概念上等效（都是分片参数/梯度/优化器状态），但实现和功能有差异：

• FSDP 是 PyTorch 原生 API，与 PyTorch 生态（编译、自动混合精度等）深度集成
• ZeRO-3 (DeepSpeed) 有更多配置选项（offload、infinity、++等扩展），灵活性更高
• 两者在具体通信调度、内存碎片管理、activation checkpointing 策略等方面有工程差异
• 选择取决于团队技术栈偏好和具体需求

学习路径

入门（1-2 小时）

1. 阅读本文
2. 观看 Microsoft DeepSpeed 官方 YouTube 讲解（搜索 “DeepSpeed ZeRO”）
3. 理解混合精度训练的内存构成

进阶（1-2 天）

1. 精读原论文：Samyam Rajbhandari et al., “ZeRO: Memory Optimizations Toward Training Trillion Parameter Models”, OSDI 2020
2. 动手实验：用 DeepSpeed 在 2-4 张 GPU 上跑 ZeRO Stage 1/2/3 的示例，观察显存占用变化
3. 对比阅读：PyTorch FSDP 官方教程

专家（持续）

1. DeepSpeed 源码：deepspeed/runtime/zero/ 目录下的实现
2. ZeRO-Offload/Infinity 论文：理解层次化存储的工程挑战
3. Megatron-DeepSpeed 融合方案与万亿参数训练实践