📚 课程目录第4章 DeepSpeed核心架构

4、DeepSpeed核心架构：ZeRO优化器、ZeRO-Offload、ZeRO-Infinity、DeepSpeed引擎、配置系统

DeepSpeed 的核心价值并不只是“把模型并行起来”，而是围绕显存、通信与计算三者的平衡，构建了一套可组合的训练优化体系。理解它的架构，关键是抓住五块拼图：ZeRO 优化器解决显存冗余，ZeRO-Offload把状态搬到 CPU，ZeRO-Infinity进一步扩展到 NVMe，DeepSpeed 引擎负责调度与执行，配置系统则把这些能力组织成可复用的 JSON 描述。它们不是并列关系，而是层层递进、彼此配合的整体。

一、ZeRO 优化器：消除数据并行中的显存冗余

传统数据并行（DP）中，每张 GPU 都保存一份完整的模型参数、梯度和优化器状态。以 Adam 为例，混合精度下每个参数大约需要 16 字节以上的状态开销，当模型达到数十亿参数时，单卡显存很快被撑爆。ZeRO（Zero Redundancy Optimizer）的核心思想是：把原本每卡重复保存的状态切分到各卡上，需要时再通过通信临时聚合。

Stage 1：只切分优化器状态，显存下降约 4 倍，通信量与标准 DP 接近，是最容易落地的起点。
Stage 2：在 Stage 1 基础上切分梯度，显存进一步下降约 8 倍，反向传播时通过 reduce-scatter 收集梯度分片。
Stage 3：连模型参数也切分，显存下降与并行度成正比，但前向/反向需要 all-gather 参数，通信量显著增加。

选择 Stage 的原则：先用 Stage 1 或 2 解决显存瓶颈；只有当单卡连一层参数都放不下时，才启用 Stage 3，并配合通信优化手段。

二、ZeRO-Offload 与 ZeRO-Infinity：把状态搬出 GPU

当 ZeRO 切分后 GPU 显存仍然不够时，下一步就是把数据下沉到更便宜的存储层级。ZeRO-Offload 将优化器状态和梯度卸载到 CPU 内存，由 CPU 完成参数更新，GPU 只负责前向和反向计算。这样单卡可以训练比自身显存大数倍的模型，代价是 PCIe 传输和 CPU 计算带来的额外延迟。

ZeRO-Infinity 则把层级再向下延伸到 NVMe SSD，允许参数、梯度和优化器状态在 GPU、CPU、NVMe 之间按需流动。它引入了带宽感知的预取与卸载策略，尽量让数据搬运与计算重叠，从而在极大规模模型上维持可接受的吞吐。两者常与 ZeRO Stage 3 组合使用，形成“GPU 计算 + CPU 内存 + NVMe 存储”的三级显存体系。

三、DeepSpeed 引擎与配置系统：把能力组织起来

DeepSpeed 引擎是实际执行训练的入口。它接管模型的前向、反向、梯度规约、参数更新和混合精度管理，并根据配置决定启用哪些 ZeRO 阶段和卸载策略。用户通常只需用 deepspeed.initialize() 包装模型和优化器，即可获得分布式训练能力。

import deepspeed

model_engine, optimizer, _, _ = deepspeed.initialize(
    args=args,
    model=model,
    model_parameters=model.parameters(),
    config="ds_config.json"
)

for step, batch in enumerate(dataloader):
    loss = model_engine(batch)
    model_engine.backward(loss)
    model_engine.step()

配置系统是 DeepSpeed 的“控制面板”，所有优化能力都通过 JSON 字段声明。下面是一个启用 ZeRO Stage 2 与 CPU 卸载的典型配置片段：

{
  "train_batch_size": 64,
  "fp16": { "enabled": true },
  "zero_optimization": {
    "stage": 2,
    "offload_optimizer": {
      "device": "cpu",
      "pin_memory": true
    },
    "allgather_partitions": true,
    "overlap_comm": true
  },
  "optimizer": {
    "type": "Adam",
    "params": { "lr": 1e-4 }
  }
}

四、核心组件对比与选型清单

组件	主要作用	显存收益	主要代价	适用场景
ZeRO Stage 1	切分优化器状态	约 4 倍	通信量基本不变	显存轻度不足
ZeRO Stage 2	切分优化器状态与梯度	约 8 倍	增加 reduce-scatter	中等规模模型
ZeRO Stage 3	切分参数、梯度、状态	与并行度成正比	all-gather 通信明显	超大模型、单层可放下
ZeRO-Offload	状态与梯度卸载到 CPU	释放大量 GPU 显存	PCIe 传输、CPU 计算	单机多卡、显存紧张
ZeRO-Infinity	扩展到 NVMe 存储	支持超大规模	存储带宽与延迟	千亿参数级训练

五、常见误区与复盘建议

误区一：Stage 越高越好。Stage 3 通信开销大，在模型本身能放进显存时反而拖慢训练。应先测量显存占用，再决定阶段。

误区二：开启 Offload 就一定更快。Offload 是用时间换空间，CPU 与 PCIe 带宽不足时可能成为瓶颈，需结合 pin_memory 和通信重叠使用。

误区三：配置文件字段可以随意堆叠。不同 ZeRO 阶段与卸载策略存在兼容性约束，错误组合会导致初始化失败或静默降级，建议先用小规模任务验证配置。

复盘时建议记录三项指标：单卡峰值显存、每步耗时、通信占比。通过对比不同 Stage 与 Offload 组合下的指标变化，才能找到当前硬件与模型规模下的最优配置，而不是盲目追求“最高级别”的优化。