Distributed Training Basics: 解释 ZeRO 的三个阶段。

Core Concept: ZeRO 分片优化器状态（阶段 1）、梯度（阶段 2）、参数（阶段 3），显存逐级下降、通信逐级增加。

数学机理：以参数量 Ψ、混合精度（FP16 参数 + FP32 master + Adam 的 m/v）计，单卡显存约为：参数 2Ψ（FP16）+ master 4Ψ（FP32）+ m 4Ψ + v 4Ψ = <strong>14Ψ 字节</strong>（另有激活与临时缓冲）。可见<strong>优化器状态（m+v+master）占 12Ψ、远大于参数的 2Ψ</strong>——这是显存的主要瓶颈。<strong>ZeRO 的分片思路</strong>（Rajbhandari 等 2020）：<strong>ZeRO-1</strong> 把<strong>优化器状态</strong>按 DP 度 N 分片，每卡只存 1/N 的 m/v/master → 优化器显存降 N 倍（12Ψ→12Ψ/N）；<strong>ZeRO-2</strong> 再把<strong>梯度</strong>分片（每卡只存 1/N 梯度）→ 梯度显存降 N 倍；<strong>ZeRO-3</strong> 再把<strong>参数</strong>分片（每卡只存 1/N 参数）→ 参数显存也降 N 倍，总显存降到约 14Ψ/N。<strong>代价</strong>：ZeRO-1 通信量与原 DP 相同（一次 reduce-scatter + all-gather 替代 all-reduce）；ZeRO-2 略增；ZeRO-3 因每层前向需 <strong>all-gather 参数</strong>、反向后丢弃，通信量显著增加（约 1.5 倍）。因此 ZeRO-3 适合'显存极度紧张、带宽尚可'的场景，ZeRO-1/2 适合'显存够用、追吞吐'的场景。

深度剖析与工程权衡：① <strong>与 FSDP 的关系</strong>——PyTorch 的 FSDP（Fully Sharded Data Parallel）是 ZeRO-3 的等价实现（参数、梯度、优化器状态全分片），两者思想一致、API 不同；DeepSpeed 的 ZeRO-3 与 FSDP 可视为同一技术的两个工程实现。② <strong>通信与计算的重叠</strong>——ZeRO-3 的 all-gather 可与计算重叠（预取下一层参数）；DeepSpeed 与 FSDP 都实现了 prefetch 以隐藏通信延迟。③ <strong>offload 的组合</strong>——ZeRO-3 可把参数/优化器状态 offload 到 CPU（ZeRO-Offload/Infinity），进一步降低 GPU 显存但受 PCIe 带宽限制。④ <strong>与 TP/PP 的组合</strong>——ZeRO 是'DP 维度的分片'，与 TP（层内切）、PP（层间切）正交；大模型训练常用 'ZeRO-1 + TP + PP' 或 'FSDP + TP' 的组合。⑤ <strong>显存账本的实用价值</strong>——面试中能算出'14Ψ 字节'并指出'优化器状态是大头'，是理解 ZeRO 动机的关键；这比记住三个阶段名字更重要。⑥ <strong>面试要点</strong>——被问'ZeRO 是什么'，应从'<strong>显存账本 → 优化器状态占大头 → 逐级分片 → 通信代价递增</strong>'这条逻辑链回答；并主动区分 ZeRO-1/2/3 的适用场景。