📚 课程目录第4章模型状态分片详解

4. 模型状态分片详解

在 ZeRO（Zero Redundancy Optimizer）体系中，模型状态被拆成三块：优化器状态、梯度和参数。DeepSpeed 的 Stage 1/2/3 分别对应这三块的分片。理解每一块“切什么、切在哪、什么时候通信”，是判断显存收益与通信代价的关键。本节围绕分片对象、切分维度、通信时机与开销展开，并给出可落地的配置与排障思路。

一、核心原理：三块状态，三种切法

数据并行（DP）下，每张卡都保存完整模型状态，显存冗余度等于 DP 规模。ZeRO 的本质是：把原本每卡一份的状态，按数据并行组（data parallel group）切分，每卡只保留 1/N，需要时再通过集合通信临时聚合。

优化器状态分片（Stage 1）：Adam 的 exp_avg、exp_avg_sq 以及 FP32 master weights 按 DP 组切分。每卡只更新自己负责的那一片参数，显存收益约 4P 字节/参数（FP32 动量+方差+主权重）。
梯度分片（Stage 2）：反向传播产生的梯度按 DP 组切分，每卡只保留 1/N 梯度。收益约 2P 字节/参数（FP16 梯度）。
参数分片（Stage 3）：模型参数本身也按 DP 组切分，前向/反向时按需 all_gather 临时拼回。收益约 2P 字节/参数（FP16 参数），但通信量显著上升。

  记忆口诀：Stage 1 切“优化器状态”，Stage 2 再切“梯度”，Stage 3 连“参数”也切。切得越深，显存越省，通信越贵。

二、方法流程：一次训练迭代中的分片与通信

以 Stage 2 为例，一次迭代的典型流程如下：

前向：每卡用完整参数计算，得到本地输出与 loss。
反向：梯度逐层产生。每层梯度就绪后，立即对该层梯度执行 reduce_scatter，把归约结果按 DP 组切分，每卡只保留自己负责的梯度分片。
优化器步：每卡用本地梯度分片更新本地优化器状态分片，并更新对应的 FP32 主权重分片。
参数同步：通过 all_gather 把各卡更新后的参数分片拼回完整参数，供下一轮前向使用。

Stage 3 的区别在于：前向之前需要 all_gather 把参数拼回，反向时再按需释放；反向中同样需要 all_gather 获取参数来计算梯度。因此 Stage 3 的通信次数与模型层数强相关。

{
  "zero_optimization": {
    "stage": 2,
    "allgather_partitions": true,
    "allgather_bucket_size": 5e8,
    "overlap_comm": true,
    "reduce_scatter": true,
    "reduce_bucket_size": 5e8,
    "contiguous_gradients": true
  },
  "train_batch_size": 64,
  "fp16": { "enabled": true }
}

三、通信开销分析：什么时候贵，贵多少

通信开销取决于通信量与通信次数。设模型参数量为 P，DP 规模为 N，单次迭代的通信量大致如下：

阶段	分片对象	单次迭代通信量（近似）	通信次数	显存收益
Stage 1	优化器状态	2P（reduce_scatter + all_gather）	1 次/迭代	约 4P 字节/参数
Stage 2	优化器状态 + 梯度	2P（梯度 reduce_scatter + 参数 all_gather）	1 次/迭代	约 6P 字节/参数
Stage 3	优化器状态 + 梯度 + 参数	约 3P～4P（前向 all_gather + 反向 all_gather + 梯度 reduce_scatter）	与层数相关，多次	约 8P 字节/参数

当 overlap_comm 开启时，DeepSpeed 会把梯度 reduce_scatter 与反向计算重叠，把参数 all_gather 与前向计算重叠，从而隐藏部分通信延迟。但通信量本身不会减少，带宽仍是瓶颈。

四、案例与场景：如何选择 Stage

假设单卡 32GB，模型 7B 参数，FP16 训练，DP=8：

纯 DP：每卡需保存参数 14GB + 梯度 14GB + 优化器状态 56GB，远超 32GB，无法启动。
Stage 1：优化器状态切 8 份，每卡约 7GB，加上参数与梯度共约 35GB，仍偏紧。
Stage 2：梯度再切 8 份，每卡约 1.75GB，总计约 30GB，可运行。
Stage 3：参数也切 8 份，每卡约 1.75GB，总计约 12GB，显存充裕，但通信量最大。

因此，优先用 Stage 2：它在显存与通信之间取得较好平衡。只有当 Stage 2 仍 OOM，或需要训练更大模型时，才升级到 Stage 3，并配合 offload 进一步降低显存。

五、误区与复盘

误区一：Stage 3 一定比 Stage 2 快。 实际上 Stage 3 通信次数多，若带宽不足或模型层数多，速度可能明显下降。应先用 Stage 2 压测，再决定是否升级。

误区二：开了 ZeRO 就不需要调 batch size。 分片后每卡计算量不变，但通信占比上升，过小的 batch 会让通信无法被计算掩盖，反而降低吞吐。

误区三：忽略 bucket size。 allgather_bucket_size 与 reduce_bucket_size 过小会导致通信次数增多，过大则增加显存峰值。通常从 5e8 起步，按实际显存与带宽微调。

复盘时建议记录三个指标：每卡显存峰值、单步通信耗时占比、吞吐（tokens/s）。若通信占比超过 30%，应优先检查 bucket size、overlap 配置与网络带宽，而不是盲目升级 Stage。

# 查看通信与显存日志
deepspeed --num_gpus=8 train.py \
  --deepspeed ds_config.json \
  --log_level INFO 2>&1 | grep -E "comm|memory|throughput"

总结：Stage 1 切优化器状态，Stage 2 加切梯度，Stage 3 再切参数。分片越深，显存越省，通信越贵。实战中应以 Stage 2 为默认起点，结合 bucket size、overlap 与 offload 做精细调优，并用日志验证通信占比是否在可接受范围内。