ZeRO 三阶段拆解：到底分了什么、什么时候 all-gather

5.1 ZeRO 要解决的问题

普通 DDP 训练一个 $\Psi$ 参数的模型，每张卡都要存：

合计单卡占用约 $16\Psi$（不算激活）。对 13B 模型就是 208 GB —— 一张 H100 80G 都装不下。DDP 里这些每张卡都是同样一份，浪费惊人。

ZeRO 的洞察：把这 $16\Psi$ 里"全卡都存"的部分，按 DP 大小 $N$ 切分。

$$ M_\text{DDP} = 16\Psi \quad\Longrightarrow\quad M_\text{ZeRO-1} = 4\Psi + \tfrac{12\Psi}{N} \quad\Longrightarrow\quad M_\text{ZeRO-2} = 2\Psi + \tfrac{14\Psi}{N} \quad\Longrightarrow\quad M_\text{ZeRO-3} = \tfrac{16\Psi}{N} $$

ZeRO-3 在 $N\to\infty$ 时单卡显存趋近 0（只剩激活）。这就是为什么 1024 卡能训万亿模型。

5.2 三个 stage 切的是什么

5.3 ZeRO-1 / ZeRO-2 时序

Stage-1 只切优化器状态，所以 reduce-scatter 的对象不变，只是 optimizer.step 时只动自己那片 m/v/master。

5.4 ZeRO-3 时序（最复杂）

所以 ZeRO-3 训练中每一层都有"取-用-丢"三步，通信非常密集。overlap_comm: true 让通信和计算并发，能藏掉大部分开销，开 ZeRO-3 一定要开这个。

5.5 一组关键 JSON 字段对照

5.6 怎么选 stage

5.7 stage 3 的"参数生命周期"

读 stage3.py 会看到三个状态机：

PartitionedParamCoordinator 是这个状态机的调度器，前向中按拓扑顺序 prefetch，反向中按相反顺序释放。看懂这部分基本就理解 ZeRO-3 的精髓。

5.8 ZeRO++ 与 SuperOffload（前沿）

README "Latest News" 里能看到几个新名字：

这些算是"进阶之路"，初次入门不用动，先用 ZeRO-2/3 把基本盘吃透。

5.10 ZeRO-1 / 2 优化器源码导览

Stage 1 / 2 的实现都在 deepspeed/runtime/zero/stage_1_and_2.py，类名 DeepSpeedZeroOptimizer（line 126）。它内部用一个叫 IPGBucket 的小数据类管理"凑齐再发出"的桶：

# stage_1_and_2.py:110–124
@dataclass
class IPGBucket:
    buffer:    Tensor             # 拍平后的桶；所有梯度按字节连续放进来
    params:    List[Parameter]    # 桶里有谁
    grads:     List[Tensor]       # 对应梯度视图
    elements:  int                # 已填字节数
    index:     int                # 多桶 round-robin 时桶编号
    has_moe_params: bool          # 桶里是否有 MoE 专家参数（决定要不要走 expert DP group）

读这个文件时几个关键 flag 与设计选择：

如果你看到一段 DS 调试日志说"reduce-scatter to expert group X，bucket size Y"，那就是上面这套 per-param-group 通信路径在发声。

5.11 ZeRO-3 状态机：PartitionedParameterCoordinator

ZeRO-3 真正的"大脑"在 deepspeed/runtime/zero/partitioned_param_coordinator.py。它跟踪每一个被切分参数当前在哪、谁正在 all-gather 它、它什么时候该被释放。理解这个类的状态枚举，是看懂 stage 3 hang / 性能波动的钥匙。

5.11.1 三种 trace 状态

# partitioned_param_coordinator.py:44–50
class ZeRoTraceMode(Enum):
    RECORD   = 1   # 第一个 forward + backward：把模块访问顺序录下来
    COMPLETE = 2   # 之后每步：按录像 prefetch，最大化 overlap
    INVALID  = 3   # 当前 forward 顺序和录像对不上：退化为同步 all-gather

三种状态的迁移：

这是为什么"第一个 step 慢、后面快"，也是为什么"模型里加了 if rank == 0: x = some_module(x) 这种 rank 分支后整训慢 5 倍" —— trace 进了 INVALID 自己不会嚷嚷。打开 __log_trace_cache_warnings 才会看到 warning。

5.11.2 关键内部数据结构

5.11.3 八个 timer 名（调优时认这个）

# partitioned_param_coordinator.py:65–71
FORWARD_FETCH_SUBMIT   = "...forward_fetch.submit"
FORWARD_FETCH_WAIT     = "...forward_fetch.wait"
FORWARD_PREFETCH_SUBMIT= "...forward_prefetch.submit"
BACKWARD_FETCH_SUBMIT  = "...backward_fetch.submit"
BACKWARD_FETCH_WAIT    = "...backward_fetch.wait"
BACKWARD_PREFETCH_SUBMIT="...backward_prefetch.submit"
FORWARD_ALL_GATHER     = "...forward_all_gather"
BACKWARD_ALL_GATHER    = "...backward_all_gather"

wall_clock_breakdown: true 时这 8 个 timer 会逐步打印（详见 ch11.10）。看 timer 名你就能定位是"提交慢"还是"等慢"，是 forward 还是 backward 在拖。

5.12 ZeROOrderedDict：参数自动 all-gather 的"魔法"

ZeRO-3 用户经常困惑：参数明明被切到别张卡了，为什么 forward 时直接写 self.weight 还能用？答案在 deepspeed/runtime/zero/partition_parameters.py 第 40–72 行的 ZeROOrderedDict：

# partition_parameters.py:57–68（简化版）
class ZeROOrderedDict(OrderedDict):
    def __getitem__(self, key):
        param = super().__getitem__(key)
        if param is not None and param.ds_status == ZeroParamStatus.NOT_AVAILABLE:
            param.all_gather()   # ★ 触发同步 all-gather
        return param

parameter_offload.py 的 _inject_parameters()（line 158–161）会递归遍历所有 submodule，把它们的 _parameters dict 替换成 ZeROOrderedDict。于是每一次访问 module._parameters["weight"]（也就是常规 PyTorch 写法的 self.weight 背后）都会过这个 __getitem__，状态是 NOT_AVAILABLE 就触发一次 all-gather。

这是为什么：

• 你的 forward 完全不用知道自己在 ZeRO-3 下；
• 但写"外部参数"（一个 module 里访问另一个 module 的参数）时必须调 register_external_parameter(consumer_module, param)——否则 Coordinator 不知道该 prefetch 它，第一次访问会同步 all-gather，触发 INVALID trace；
• HF Trainer 在保存 ZeRO-3 模型时之所以要 stage3_gather_16bit_weights_on_model_save，是因为它要绕过 __getitem__ 直接拿完整参数。

5.13 NCCL 4-byte alignment 与 padding

NCCL 的 collective 要求每张卡上的张量首地址按 4 字节对齐。BF16 / FP16 单元素 2 字节，所以"4 字节 = 2 个 element"。bf16_optimizer.py 第 87 行硬编码了：

self.nccl_start_alignment_factor = 2   # 4-byte / sizeof(bf16)

切分参数时如果某一片不是 2 的整数倍，DS 会在末尾 padding 几个元素，存到 self.group_paddings。这就是为什么 ZeRO 切分后单卡 partition 看起来"比 ⌈Ψ/N⌉ 多几个元素"——不是 bug，是 NCCL 要的对齐。

5.14 进阶：ZenFlow（选择性梯度更新）

ZenFlow 是 2025-10 引入的 stall-free offload 方案。核心 idea：CPU offload 真正的瓶颈不是带宽，而是GPU 等 CPU optimizer step 完才能走下一步。ZenFlow 用 top-k 选择把"重要梯度"留在 GPU 立刻更新，剩下的丢给 CPU 异步累积。代码在 deepspeed/runtime/zenflow/。

关键配置（zenflow_config.py 第 12 行附近的 ZenFlowConfig）：

两种工厂在 zenflow_stage_1_and_2.py:82 的 ZenFlowZeroOptimizer.create() 里分发：Sequential 简单稳定，Parallel 更快但要求 PyTorch ≥ 2.1。前置：必须开 cpu_offload=true，ZenFlow 才有意义。

5.15 进阶：SuperOffload vs ZeRO-Infinity 的本质区别

初看 SuperOffload 像是"ZeRO-Infinity 的升级版"，实际定位完全不同：

源码在 deepspeed/runtime/superoffload/。superoffload_stage3.py:27 的 SuperOffloadOptimizer_Stage3 继承自 DeepSpeedZeroOptimizer_Stage3，只重写了 _optimizer_step() —— 不直接调 cpu_adam，而是把任务发到 worker 进程。superoffload_utils.py:19–36 定义 IPC 协议的 TaskKeys / ResultKeys / EventTypes。预分配 pinned buffer 通过 max_grad_numel 估算（避免运行时再 alloc）。

启用前置：

• 必须 ZeRO-3 + offload_optimizer.device=cpu；
• config 段加 "offload_optimizer": { ..., "super_offload": true, "cpuadam_cores_perc": 0.8 }（offload_config.py:96/99）；
• CPU 核心数要够——worker 进程会吃掉一部分 thread。

5.16 调 stage / 调字段的经验法则

1. 能 ZeRO-1 不 ZeRO-2，能 ZeRO-2 不 ZeRO-3 —— stage 越高通信越重；
2. ZeRO-3 必须开 overlap_comm: true；
3. 通信慢的网络（千兆以下）尽量不用 ZeRO-3；
4. 显存边界一旦突破，立即降一档 batch / 升一档 stage / 加 offload，不要恋战；
5. 能用 auto 字段就用 "auto"，让 HF Trainer 或框架去填，减少自己出错。