📚 课程目录第3章 DeepSpeed核心概念

3、DeepSpeed核心概念：ZeRO优化器、混合精度训练、梯度累积、通信优化

DeepSpeed 之所以能在单卡显存不变的前提下把大模型训练规模推高一个数量级，靠的不是某一个“银弹”，而是四类核心机制协同工作：ZeRO 优化器解决显存冗余，混合精度训练降低计算与存储开销，梯度累积在有限显存下模拟大 batch，通信优化保证多卡扩展时不被通信拖垮。理解这四者的分工与配合，是后续做 DeepSpeed 选型的基础。

一、ZeRO 优化器：把冗余状态切分掉

数据并行（DP）下，每张卡都保存一份完整的模型参数、梯度和优化器状态。以 Adam 为例，混合精度训练时每个参数大约需要 16 字节（fp16 参数 2B + fp16 梯度 2B + fp32 主权重 4B + fp32 动量 4B + fp32 方差 4B），一个 10B 参数模型仅状态就要 160GB，远超单卡容量。ZeRO 的核心思想是：这些状态在数据并行组内本来就是重复的，完全可以切分到各卡上，需要时再通过通信收集。

ZeRO-1（Pos）：只切分优化器状态，显存约降为 1/4，通信量与 DP 基本持平，性价比最高。
ZeRO-2（Pos+g）：在 Stage 1 基础上再切分梯度，显存约降为 1/8，适合中等规模模型。
ZeRO-3（Pos+g+p）：参数也切分，显存随卡数线性下降，可训练超大模型，但通信量显著增加，需要配合通信重叠。

选型口诀：显存够用优先 ZeRO-1；单卡放不下完整模型时上 ZeRO-2；模型大到单卡连一层都放不下，才用 ZeRO-3。

二、混合精度训练：fp16/bf16 与 fp32 主权重

混合精度指前向与反向用 fp16 或 bf16 计算，同时保留一份 fp32 主权重用于参数更新。这样既利用了 Tensor Core 的高吞吐，又避免了小梯度在低精度下被舍入为零。DeepSpeed 通过 fp16 或 bf16 配置块开启，并支持动态损失缩放（dynamic loss scaling）来防止梯度下溢。

{
  "bf16": {
    "enabled": true
  },
  "optimizer": {
    "type": "AdamW",
    "params": { "lr": 3e-4, "betas": [0.9, 0.95] }
  },
  "zero_optimization": {
    "stage": 2,
    "offload_optimizer": { "device": "cpu" }
  }
}

bf16 动态范围与 fp32 相同，通常不需要 loss scaling，稳定性优于 fp16，是当前大模型训练的首选；fp16 精度更高但需要缩放，适合对数值精度敏感的场景。

三、梯度累积：小显存模拟大 batch

梯度累积把一个大 batch 拆成若干 micro-batch 依次前向反向，梯度累加后再统一更新一次参数。它不改变优化器看到的等效 batch size，却把峰值显存压到单个 micro-batch 的水平。DeepSpeed 通过 train_batch_size 与 gradient_accumulation_steps 自动推导 micro-batch：

micro_batch = train_batch_size / (gradient_accumulation_steps × world_size)

例如 train_batch_size=1024、gradient_accumulation_steps=8、world_size=8，则每卡 micro-batch 为 16。注意：梯度累积会拉长单步时间，且与 ZeRO 的梯度切分配合时，累积发生在本地，通信发生在累积之后，顺序不能颠倒。

四、通信优化：让扩展效率不掉队

ZeRO 切分状态后，每次前向/反向都需要 all-gather 参数、reduce-scatter 梯度，通信量随 stage 升高而增大。DeepSpeed 提供多种手段把通信藏到计算背后：

通信与计算重叠：在反向传播计算某一层梯度的同时，异步通信上一层已就绪的梯度。
梯度压缩：通过 compression 配置对梯度做量化或稀疏化，降低带宽占用。
通信桶（bucket）调优：合理设置 reduce_bucket_size 与 allgather_bucket_size，平衡延迟与吞吐。
拓扑感知：优先在 NVLink/NVSwitch 内通信，跨节点走 InfiniBand，减少跨机流量。

五、四类机制协同关系

下表总结四者在训练流程中的位置与主要作用，便于选型时快速对照。

机制	主要解决的问题	关键配置	典型代价
ZeRO	参数/梯度/优化器状态显存冗余	`zero_optimization.stage`	通信量随 stage 上升
混合精度	计算吞吐与显存占用	`fp16` / `bf16`	数值稳定性需调优
梯度累积	小显存模拟大 batch	`gradient_accumulation_steps`	单步时间变长
通信优化	多卡扩展效率	`overlap_comm`、`compression`	配置复杂度上升

六、案例与常见误区

某团队用 8×A100 80GB 训练 7B 模型，最初直接开 ZeRO-3，发现吞吐只有 ZeRO-2 的 60%。复盘发现模型单卡完全放得下，ZeRO-3 带来的额外 all-gather 纯属浪费。改为 ZeRO-2 + bf16 + 梯度累积 4 后，吞吐提升约 40%，显存仍有余量。

常见误区：① 认为 stage 越高越好，忽略通信开销；② 把 train_batch_size 设成小于 world_size × gradient_accumulation_steps，导致 micro-batch 为 0 报错；③ 开了 bf16 又手动加 loss scaling，反而引入不稳定；④ 只调 ZeRO 不调 bucket size，通信效率上不去。

实操建议：先用 ZeRO-1 或 ZeRO-2 跑通，观察显存与吞吐；显存不足再升 stage；始终开启 bf16；梯度累积步数按等效 batch 反推；通信瓶颈时优先调 bucket 与 overlap，再考虑压缩。

把这四类机制看成一套组合拳：ZeRO 决定显存上限，混合精度决定算力利用率，梯度累积决定 batch 灵活性，通信优化决定多卡扩展效率。选型时先明确瓶颈在显存、算力还是通信，再针对性开启对应能力，才能避免“全开反而更慢”的常见陷阱。