3、DeepSpeed核心概念：ZeRO优化器、混合精度训练、梯度累积、检查点机制

DeepSpeed 之所以能在单卡到千卡规模上显著降低显存占用并提升吞吐，核心并不只是“把模型切到多卡”，而是围绕训练状态管理建立了一整套机制。理解 ZeRO、混合精度、梯度累积和检查点机制，是后续配置 JSON、调参与排障的基础。本节从训练状态出发，把这四个概念串成一条可操作的认知链。

一、训练状态与 ZeRO 优化器

在标准数据并行中，每张 GPU 都保存一份完整模型副本，训练状态主要包括三部分：优化器状态（如 Adam 的动量与方差）、梯度、模型参数。以 Adam 为例，混合精度下每个参数通常需要 2 字节参数、2 字节梯度、4 字节动量、4 字节方差，再加 4 字节 fp32 主权重，显存开销远大于参数本身。ZeRO（Zero Redundancy Optimizer）的核心思想是：把这些状态按 GPU 数量切分，消除冗余，需要时再通过通信收集。

  • ZeRO-1：只切分优化器状态，显存下降明显，通信量最小，适合显存瓶颈初现的场景。
  • ZeRO-2：切分优化器状态与梯度，进一步降低显存，通信量适中，是很多训练任务的默认选择。
  • ZeRO-3：参数也切分，显存占用最低，可训练超大模型，但通信与实现复杂度最高，常配合 offload 使用。

ZeRO 的收益并非“免费”：切分越彻底，前向、反向和更新阶段需要的 all-gather、reduce-scatter 通信越多。因此选择 stage 时要在显存、吞吐和通信带宽之间权衡。下面这张图展示了 ZeRO 从 1 到 3 的切分范围与通信关系。

ZeRO 三阶段：切分范围与通信代价 ZeRO-1 仅切分优化器状态 优化器状态：切分 梯度：完整副本 参数：完整副本 通信：低 显存节省：中 ZeRO-2 切分优化器状态 + 梯度 优化器状态：切分 梯度：切分 参数：完整副本 通信：中 显存节省：较高 ZeRO-3 优化器状态 + 梯度 + 参数 优化器状态：切分 梯度：切分 参数：切分 通信：高 显存节省：最高

二、混合精度训练：fp16/bf16 与 loss scaling

混合精度训练让前向和反向的大部分计算使用 fp16 或 bf16，从而减少显存带宽压力并利用 Tensor Core 加速，同时保留 fp32 主权重用于参数更新。DeepSpeed 通过 fp16 或 bf16 配置开启。fp16 动态范围窄，梯度容易下溢，因此需要 loss scaling：先放大 loss，反向得到放大后的梯度，更新前再缩放回去。bf16 动态范围接近 fp32，通常不需要 loss scaling，但精度略低。

{
  "fp16": {
    "enabled": true,
    "loss_scale": 0,
    "initial_scale_power": 16,
    "loss_scale_window": 1000,
    "hysteresis": 2,
    "min_loss_scale": 1
  },
  "bf16": {
    "enabled": false
  }
}

其中 loss_scale: 0 表示使用动态 loss scaling，DeepSpeed 会根据是否出现 inf/nan 自动调整缩放系数。若训练中出现 loss 突然变成 nan，优先检查 loss scaling 是否过小或过大，以及是否误用了 bf16 与 fp16 同时开启。

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

梯度累积通过多次前向反向、累积梯度后再更新一次参数，使有效 batch size 等于 micro_batch × gradient_accumulation_steps × data_parallel_size。它不降低单步显存峰值，但能提升训练稳定性与吞吐利用率。DeepSpeed 中通过 train_batch_size 与 gradient_accumulation_steps 配合控制。

配置示例：若单卡 micro batch 为 4，数据并行度为 8，希望全局 batch 为 256，则 gradient_accumulation_steps = 256 / (4 × 8) = 8。DeepSpeed 会在内部按此拆分，用户只需保证 train_batch_size 与 micro_batch_per_gpu 一致。

四、检查点机制：断点续训与状态一致性

DeepSpeed 的检查点不仅保存模型权重，还保存优化器状态、学习率调度器、随机数状态和 ZeRO 切分信息。使用 model_engine.save_checkpoint(path) 保存，load_checkpoint(path) 恢复。ZeRO-3 下参数被切分，必须由 DeepSpeed 统一管理加载，不能直接用 torch.load 读取单个分片。

机制 主要作用 关键配置 常见误区
ZeRO 切分优化器状态、梯度、参数，降低单卡显存 zero_optimization.stage 盲目上 ZeRO-3 导致通信瓶颈
混合精度 加速计算、降低显存带宽 fp16.enabled / bf16.enabled fp16 与 bf16 同时开启
梯度累积 小显存模拟大 batch gradient_accumulation_steps 与 train_batch_size 不匹配
检查点 断点续训、状态恢复 save_checkpoint / load_checkpoint 手动 torch.load 读取 ZeRO 分片

五、案例与复盘：一次显存不足的排查

某团队在 8 卡 A100 上训练 7B 模型，初始配置为 ZeRO-2 + fp16，单卡 micro batch 为 2，仍出现 OOM。排查后发现：优化器状态未 offload，且梯度累积步数设为 1，导致有效 batch 过小、梯度噪声大。调整方案为：开启 offload_optimizer，将 gradient_accumulation_steps 提升到 4，并保持 ZeRO-2 不变。结果显存下降约 30%，吞吐仅下降 8%，训练稳定收敛。

实践建议：优先从 ZeRO-2 + bf16 起步；显存仍不足再考虑 ZeRO-3 或 offload；梯度累积步数要与全局 batch 目标对齐；每次修改配置后，先用小规模数据跑通保存与恢复流程，再启动长训练。

注意：ZeRO-3 下模型参数在每层前向时动态收集，若在训练循环中直接访问 model.module 或手动修改参数，可能破坏切分一致性。检查点恢复时务必使用 DeepSpeed 提供的接口，并确保 stage 与保存时一致。

把 ZeRO、混合精度、梯度累积和检查点放在一起看，它们分别对应显存、计算、batch 规模和训练连续性四个维度。实际调参时，先确定显存上限，再选择精度与 ZeRO stage，最后用梯度累积补齐 batch，并用检查点保证可恢复。这四个概念共同构成了 DeepSpeed 训练配置的骨架。

← 上一章 📖 返回目录 下一章 →