DeepSpeed ZeRO分布式训练优化 - 第4节

4、ZeRO-Stage 1：优化器状态切分原理、显存节省计算、通信开销分析、适用场景与配置方法

在大规模深度学习模型训练中，当模型参数量达到数十亿甚至千亿级别时，GPU显存瓶颈往往首先出现在优化器状态（Optimizer States）上，而非模型参数本身。以Adam优化器为例，每个参数需要维护一阶动量（moment）和二阶动量（variance），再加上参数本身和梯度，显存占用呈4倍甚至更高的放大效应。ZeRO-Stage 1正是针对这一痛点设计的第一级优化策略：将优化器状态按数据并行维度进行切分，每个GPU只负责存储和更新自己分片对应的那部分优化器状态，从而在不改变模型逻辑和训练语义的前提下，显著降低单卡显存峰值。

上图清晰地展示了Baseline与ZeRO-Stage 1在显存布局上的本质差异。在Baseline模式下，每个GPU都保存完整的参数、梯度和优化器状态，显存冗余度极高；而Stage 1将Adam的动量m和方差v按数据并行度N切分成N份，每个GPU只持有其中一份。这种设计带来的直接收益是：当N=64时，优化器状态从原来的8ψ字节（Adam）降至0.125ψ字节，显存占用大幅下降。更关键的是，Stage 1不改变通信模式——它仍然只需要在反向传播后进行梯度All-Reduce，通信量与数据并行Baseline完全一致，因此不会引入额外的通信瓶颈。

4.1 核心原理：从“全量冗余”到“分片自治”

4.1.1 优化器状态的内存放大效应

理解Stage 1的前提是掌握优化器状态的内存放大公式。以混合精度训练（FP16参数+FP32主权重）为例，假设模型参数量为ψ，则：

• FP16参数副本：2ψ字节（用于前向/反向计算）
• FP32主权重：4ψ字节（用于参数更新）
• FP32梯度：4ψ字节（在Adam中梯度以FP32累积）
• Adam动量m：4ψ字节（FP32）
• Adam方差v：4ψ字节（FP32）

合计为18ψ字节，其中优化器状态（m+v）占8ψ字节，约占总显存的44%。当ψ=10B（百亿参数）时，仅优化器状态就需要80GB显存，远超单张A100（80GB）的容量。Stage 1的核心思想就是将这8ψ字节的优化器状态按数据并行度N切分，使每卡只需存储8ψ/N字节。

4.1.2 切分与更新流程

Stage 1的训练流程分为四个关键步骤：

1. 前向传播：每卡使用完整的FP16参数副本计算各自数据分片的损失。
2. 反向传播：每卡计算出完整的FP32梯度（梯度本身不切分）。
3. 梯度All-Reduce：跨GPU对梯度进行求和平均，确保每卡获得一致的全局梯度。
4. 局部更新：每卡仅使用自己分片对应的梯度部分，更新自己持有的m、v分片和对应的参数分片。更新完成后，通过Broadcast操作将更新后的参数分片广播给所有GPU，确保下一轮前向传播时每卡都有完整的参数。

4.2 显存节省计算：量化收益与边界

4.2.1 显存公式推导

设模型参数量为ψ，数据并行度为N，使用Adam优化器和混合精度训练，则：

存储项	Baseline（每卡）	ZeRO-Stage 1（每卡）	节省比例
FP16参数	2ψ	2ψ	0%
FP32主权重	4ψ	4ψ	0%
FP32梯度	4ψ	4ψ	0%
Adam动量m	4ψ	4ψ/N	(N-1)/N
Adam方差v	4ψ	4ψ/N	(N-1)/N
总计	18ψ	14ψ + 4ψ/N	随N增大趋近于22%

以 175B模型为例，ψ=175e9，N=64时：Baseline每卡需要18×175=3150GB显存（显然无法单卡运行），而Stage 1每卡需要14×175+4×175/64=2450+10.9=2460.9GB，仍然超出单卡容量。这说明Stage 1单独使用时，对于超大模型仍力不从心，但它为Stage 2和Stage 3奠定了基础。

4.2.2 实际训练中的显存分布

在实际训练中，除了上述静态显存外，还需要考虑激活值（Activation）、临时缓冲区（如通信缓冲区）和CUDA上下文开销。Stage 1的显存节省主要体现在静态部分，对于激活值显存没有直接影响。因此，在判断是否使用Stage 1时，需要先估算激活值显存占比。如果激活值显存已经接近GPU上限，那么Stage 1的收益会被稀释。

4.3 通信开销分析：零额外成本的背后

4.3.1 通信模式对比

Stage 1的通信模式与数据并行Baseline完全一致，都只需要在反向传播结束后进行一次梯度All-Reduce。通信量公式为：

通信量 = 2 × ψ × 4字节 × (N-1)/N  （All-Reduce的通信量近似为2倍数据量）

以ψ=10B、N=32为例，单次通信量约为2×10×4×31/32≈77.5GB。在NVLink（600GB/s）或InfiniBand（200Gb/s）环境下，这会产生数十秒的通信延迟。但关键在于，Stage 1没有增加任何额外的通信步骤——它不像Stage 2那样需要Scatter梯度，也不像Stage 3那样需要Broadcast参数。因此，从通信开销角度看，Stage 1是“零成本”的优化。

4.3.2 通信与计算重叠

在实际工程中，梯度All-Reduce可以与反向传播的后半段重叠执行。PyTorch的DDP（DistributedDataParallel）和DeepSpeed都支持梯度桶（Gradient Bucket）机制，将梯度分成多个桶，每计算完一个桶就立即开始通信，从而隐藏通信延迟。Stage 1完全继承了这一优化能力，不会因为切分优化器状态而破坏梯度桶的连续性。

4.4 适用场景与配置方法

4.4.1 适用场景判断

Stage 1最适合以下场景：

• 模型参数量在1B-10B之间：此时优化器状态是主要显存瓶颈，但参数和梯度仍占较大比例，Stage 1的切分收益明显。
• 数据并行度N适中（8-64）：N太小则节省有限，N太大则通信开销可能超过收益。
• 已有数据并行训练代码：从DDP迁移到Stage 1几乎零成本，只需修改几行配置。
• 对通信延迟敏感：如果网络带宽有限，Stage 1是比Stage 2/3更安全的选择。

4.4.2 DeepSpeed配置示例

在DeepSpeed中启用Stage 1只需在配置文件中设置即可：

{
  "train_batch_size": 1024,
  "gradient_accumulation_steps": 4,
  "optimizer": {
    "type": "Adam",
    "params": {
      "lr": 1e-4,
      "betas": [0.9, 0.999],
      "eps": 1e-8
    }
  },
  "zero_optimization": {
    "stage": 1,
    "allgather_partitions": true,
    "reduce_scatter": false,
    "overlap_comm": true,
    "contiguous_gradients": true
  }
}

关键参数说明：stage: 1启用Stage 1；allgather_partitions控制参数更新后是否进行全量聚合；overlap_comm开启通信与计算重叠；contiguous_gradients将梯度存储在连续内存块中以提高通信效率。

4.4.3 从DDP迁移的实战案例

某团队训练一个6.7B参数的风格模型，使用8×A100（80GB）GPU。初始使用PyTorch DDP，遇到OOM错误。分析显存分布：参数FP16（13.4GB）+ 主权重FP32（26.8GB）+ 梯度FP32（26.8GB）+ Adam状态（53.6GB）= 120.6GB，远超80GB。激活值通过Checkpointing控制在10GB左右。

迁移到Stage 1后，Adam状态降至53.6/8=6.7GB，总显存需求变为13.4+26.8+26.8+6.7+10=83.7GB，仍然略超80GB。团队进一步将gradient_accumulation_steps从1调整为2，使梯度显存减半至13.4GB，最终总需求降至70.3GB，成功运行。训练吞吐量从DDP的无法运行提升至约1500 tokens/s/GPU，通信开销占比约12%，完全在可接受范围内。

4.5 总结复盘：Stage 1的定位与边界

ZeRO-Stage 1是迈向极致显存优化的第一步，它的价值在于以零通信代价换取约22%的显存节省（当N足够大时）。在实际工程中，Stage 1最适合作为从DDP升级的“第一跳”，因为它几乎不改变训练脚本和通信模式，风险极低。然而，它的局限性也很明显：参数和梯度仍然全量冗余，对于超大模型（>20B）力不从心。

从系统设计的角度看，Stage 1揭示了一个重要原则：显存优化的核心在于识别“哪些数据可以安全地分片存储，而不影响计算正确性”。优化器状态之所以可以切分，是因为每个参数的更新只依赖于其自身的梯度，而梯度在All-Reduce后每卡都有一份完整副本。这一洞察为Stage 2（切分梯度）和Stage 3（切分参数）铺平了道路。建议读者在掌握Stage 1后，继续深入理解Stage 2和Stage 3的递进关系，形成完整的ZeRO知识体系。