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

3、DeepSpeed核心概念总览：ZeRO、Offload、混合精度、梯度累积、通信优化等术语解析

DeepSpeed 的配置项看起来很多，但真正决定训练能否跑起来、跑得快、跑得大的，往往就是几个核心术语：ZeRO、Offload、混合精度、梯度累积和通信优化。理解它们之间的关系，比死记参数更重要。可以把它们看成一条协同链路：显存不够时，ZeRO 负责切分模型状态；切完还紧张，Offload 把部分状态搬到 CPU/NVMe；混合精度降低单份数据体积并加速计算；梯度累积在有限显存下模拟更大 batch；通信优化则保证切分和搬运之后，多卡之间仍然高效同步。

一、ZeRO：把模型状态切开来放

ZeRO（Zero Redundancy Optimizer）解决的是数据并行中“每张卡都存一份完整模型状态”的冗余问题。训练时每张 GPU 通常要保存参数、梯度和优化器状态，其中优化器状态（如 Adam 的动量和方差）往往最占显存。ZeRO 按阶段逐步切分：

ZeRO-1：只切分优化器状态，显存收益明显，通信量增加较少，是最常用的起点。
ZeRO-2：在 ZeRO-1 基础上再切分梯度，进一步降低显存，适合中等规模模型。
ZeRO-3：参数也切分，每张卡只保留一部分参数，需要时再通过通信收集，能支撑超大模型，但通信开销最高。

选择阶段时不要盲目追求 ZeRO-3。模型能放进显存时，ZeRO-1 或 ZeRO-2 往往性价比更高；只有单卡显存实在放不下完整参数时，才需要 ZeRO-3 的参数切分。

二、Offload：把冷数据搬到更便宜的存储

Offload 是 ZeRO 的“外援”。当 GPU 显存仍然不够时，可以把优化器状态、梯度甚至参数卸载到 CPU 内存或 NVMe 磁盘。常见配置是 offload_optimizer 和 offload_param。它的代价是 PCIe 或 NVMe 带宽远低于显存带宽，因此会拖慢训练速度。实践中通常优先 offload 优化器状态，因为它在反向传播后才更新，对计算路径影响相对小；参数 offload 则更激进，适合显存极度受限但能接受速度下降的场景。

判断顺序：先尝试 ZeRO-1/2 切分；仍 OOM 再开启 optimizer offload；最后才考虑 param offload 或 NVMe offload。

三、混合精度：用 FP16/BF16 换显存和速度

混合精度训练让前向和反向计算主要使用 FP16 或 BF16，同时保留 FP32 主权重用于更新。这样既减少显存占用，又利用 Tensor Core 提升吞吐。DeepSpeed 中常用 fp16 或 bf16 配置，并配合动态损失缩放（loss scaling）防止梯度下溢。BF16 动态范围更大，通常比 FP16 更稳定，在支持 BF16 的硬件上优先选用。

{
  "bf16": {
    "enabled": true
  },
  "zero_optimization": {
    "stage": 2,
    "offload_optimizer": {
      "device": "cpu"
    }
  }
}

四、梯度累积与通信优化：小显存跑大 batch

梯度累积是在显存不足以放下大 batch 时，把多个 micro-batch 的梯度累加后再更新一次参数，从而模拟大 batch 的训练效果。它不减少单步显存峰值，但能改善梯度稳定性。通信优化则关注多卡之间的 AllReduce、AllGather、ReduceScatter 等集合通信。DeepSpeed 通过通信桶（bucket）合并小张量、重叠计算与通信、以及针对 ZeRO 阶段的通信调度，减少同步等待。

下面这张图展示了这些概念在训练迭代中的协作关系：

五、术语对照与常见误区

术语	主要作用	典型收益	主要代价
ZeRO-1/2/3	切分优化器状态、梯度、参数	显著降低单卡显存	阶段越高通信越多
Offload	把状态搬到 CPU/NVMe	突破显存上限	PCIe/NVMe 带宽瓶颈
混合精度	FP16/BF16 计算 + FP32 主权重	省显存、提吞吐	需处理数值稳定性
梯度累积	多 micro-batch 累加后更新	模拟大 batch	不降低单步显存峰值
通信优化	合并、重叠、调度集合通信	提升多卡扩展效率	配置与调优复杂度上升

常见误区：把 ZeRO-3 当成万能加速器。实际上 ZeRO-3 会增加 AllGather 通信，若模型本可放入显存，ZeRO-1/2 往往更快。另一个误区是同时开启所有优化，导致通信与卸载互相争抢带宽，反而变慢。建议逐项开启并观察吞吐与显存曲线。

六、场景复盘：一次 OOM 的排查路径

假设训练 7B 模型时出现 OOM。第一步，检查是否已启用 BF16，若仍是 FP32，先切混合精度，通常能省下近一半激活和参数显存。第二步，启用 ZeRO-2，切分优化器状态和梯度。第三步，若仍不足，开启 optimizer offload 到 CPU。第四步，若模型参数本身过大，再考虑 ZeRO-3 或 param offload。第五步，若多卡吞吐不升反降，检查通信桶大小和梯度累积步数是否匹配，避免过小的 micro-batch 导致通信占比过高。整个过程的核心逻辑是：先降精度，再切状态，后卸载，最后调通信。

把这些术语串起来看，DeepSpeed 的配置就不是零散开关，而是一套围绕“显存—计算—通信”三角的权衡体系。理解每个术语解决哪一类瓶颈，才能在具体模型和硬件上做出合理选择。