DeepSpeed 完全图解：ZeRO / 3D 并行 / Offload 从零讲清楚

要理解 ZeRO，先得把「训练一个模型，显存被谁吃了」算清楚。以 Adam 优化器 + 混合精度（FP16 前向 + FP32 优化器）训练为例：

1.1 16 字节/参数法则

1.1.1 用一个极简神经网络理解 P / G / O / 激活

假设有一个只有 2 个参数的极简全连接网络（输入 x，中间一个权重 w₁，再一个权重 w₂，输出 ŷ），用 MSE loss 训练。这个例子虽然简单，但能把所有概念对应清楚。

现在我们来逐一看，训练这个网络时 GPU 里存了什么：

1.2 激活（Activation）占用

激活是前向传播时产生的中间张量（每层的 attention 矩阵、FFN 输出 …），反向时要用来算梯度，所以默认会全量保留在显存。它的大小跟 batch size × seq len × 模型宽度成正比，跟参数量无关：

2.1 ZeRO-1：只切优化器状态 O

最保守的一档。只切「Adam 的 FP32 副本 + m + v」（共 12 字节/参数）。参数 P 和梯度 G 每张卡仍然完整持有。

2.2 ZeRO-2：再把梯度 G 也切了

在 ZeRO-1 基础上，梯度也分片。每张卡反向后立刻 reduce-scatter，只保留自己负责的那段梯度，不再持有全量 G。

2.3 ZeRO-3：连参数 P 都切（最猛）

P/G/O 全部分片，每张卡只持有 1/N 的模型权重。需要某层时 all-gather 凑回来，算完立刻释放。

2.4 通信量分析：ZeRO 为什么没变慢多少

2.5 三档对比总表

ZeRO-3 还不够？还是 OOM？那就接着甩——把优化器状态、梯度，甚至参数本身，从 GPU HBM 搬到 CPU RAM。这就是 ZeRO-Offload（2021）。

// ds_config.json 关键 Offload 字段
"zero_optimization": {
  "stage": 3,
  "offload_optimizer": {
    "device": "cpu",
    "pin_memory": true    // pin_memory=true 让 CPU 到 GPU 传输更快
  },
  "offload_param": {
    "device": "cpu",      // 只在 OOM 时才加这个！
    "pin_memory": true
  }
}

ZeRO-Infinity（SC'21）把同样的思路推到极致：P/G/O 全部可以 offload 到本地 NVMe SSD，同时设计了 bandwidth-centric partitioning（以带宽为中心的分片）来最大化 NVMe 读写吞吐。

ZeRO 是「数据并行里的优化版」——多张卡处理不同 batch，但模型的每一层都在这 N 张卡里某张完整地算完。当模型大到单层都装不下时，光靠 ZeRO 还不够，需要三种并行同时上。

5.1 数据并行（DP）— ZeRO 所在的维度

5.2 张量并行（TP / Megatron Style）

把同一层的权重矩阵切到多卡上分别计算，再通信合并结果。由 Megatron-LM 提出，节点内（NVLink）效率最高。

5.3 流水线并行（PP）

把不同 Transformer 层放到不同 GPU 上，前向像工厂流水线一样推进。

5.4 三维组合：DP × TP × PP

6.1 浮点数格式基础：指数 vs 尾数

要理解 FP16 和 BF16 的区别，必须先搞清楚浮点数的两个核心字段：

6.2 三种格式的位布局对比

6.3 FP16 vs BF16：核心差距在哪里？

具体问题一：FP16 的溢出 / 下溢

具体问题二：BF16 的精度代价

6.4 为什么需要「混合」精度？直接全用 FP16 不行吗？

不行，原因是 Adam 优化器的更新步长极小：

所以必须在优化器里保留一份 FP32 主权重来做精确更新，这就是「混合」的含义：

6.5 Loss Scaling：专门解决 FP16 梯度下溢的工程手段

6.6 实际选型建议

// ds_config.json 精度配置（A100/H100 用 BF16）
"bf16": { "enabled": true  },
"fp16": { "enabled": false }

// 旧 GPU（V100/3090）用 FP16 + 动态 Loss Scaling
"fp16": {
  "enabled": true,
  "loss_scale": 0,              // 0 = 动态自适应
  "initial_scale_power": 16,   // 初始 S = 2^16 = 65536
  "loss_scale_window": 1000    // 连续 1000 step 无 overflow 则 S 翻倍
}

ZeRO 切的是 P/G/O，但激活（activations）没碰。激活是前向中所有中间张量，反向时要用，默认全量保留。Gradient Checkpointing 的思路很暴力：

# HuggingFace Trainer / LLaMA-Factory 一行开启
model.gradient_checkpointing_enable()
# 或在 TrainingArguments 里
gradient_checkpointing=True

下面是一份通用 ZeRO-3 + offload optimizer 配置，适用于 LLaMA-7B/13B 等大模型，配合 HuggingFace Trainer / LLaMA-Factory 直接使用。

{
  // ── Batch 大小（写 "auto" 让 Trainer 自动填入）──
  "train_batch_size": "auto",
  "gradient_accumulation_steps": "auto",
  "gradient_clipping": 1.0,         // 梯度裁剪，防止 NaN

  // ── 精度（A100/H100 用 BF16，V100 用 FP16）──
  "bf16": { "enabled": true  },
  "fp16": { "enabled": false },

  // ── ZeRO 核心配置 ──
  "zero_optimization": {
    "stage": 3,                     // ZeRO 档位：1/2/3
    "overlap_comm": true,           // 通信与计算重叠，提速 10~20%
    "contiguous_gradients": true,   // 梯度内存连续，减碎片
    "reduce_bucket_size": 5e8,      // 每桶梯度 ~500MB，太大通信延迟增大

    // ZeRO-3 参数分片细节
    "stage3_prefetch_bucket_size": 5e8,      // 提前 gather 下一层参数（流水线）
    "stage3_param_persistence_threshold": 1e6, // 小于 1M 参数的层不分片（节省通信）
    "stage3_max_live_parameters": 1e9,
    "stage3_max_reuse_distance": 1e9,
    "stage3_gather_16bit_weights_on_model_save": true, // ⚠️ 必须开！否则 ckpt 只存 1/N

    // ── CPU Offload ──
    "offload_optimizer": {
      "device": "cpu",          // 优化器状态（FP32副本+m+v）移到 CPU
      "pin_memory": true        // pin_memory 提升 PCIe 传输速度
    }
    // "offload_param": { "device": "cpu", "pin_memory": true }
    // ↑ 仅在极端 OOM 时才加，会大幅降低训练速度
  },

  // ── 优化器（"auto" = 让 Trainer 接管）──
  "optimizer": {
    "type": "AdamW",
    "params": { "lr": "auto", "betas": "auto", "eps": "auto", "weight_decay": "auto" }
  },

  // ── 学习率调度 ──
  "scheduler": {
    "type": "WarmupDecayLR",
    "params": {
      "total_num_steps": "auto", "warmup_num_steps": "auto",
      "warmup_min_lr": "auto",   "warmup_max_lr": "auto"
    }
  }
}

启动命令

# 通过 HuggingFace Trainer
deepspeed train.py \
  --deepspeed ds_config.json \
  --model_name_or_path meta-llama/Llama-3-8B \
  --per_device_train_batch_size 2 \
  --learning_rate 2e-5 ...

# 通过 LLaMA-Factory（直接用预设模板）
llamafactory-cli train \
  examples/train_full/llama3_full_sft_ds3.yaml

# torchrun 多节点
torchrun --nproc_per_node=8 --nnodes=4 \
  --master_addr=... --master_port=29500 \
  train.py --deepspeed ds_config.json ...

场景 × 配置速查表

7 个常见坑

一句话总结

Q1：ZeRO 跟模型并行有啥区别？也是一个模型拆成几份放不同机器吗？

虽然都是「拆开放不同机器」，但拆的东西不一样，更关键的是计算方式不一样。

Layer 1,2,3 → GPU 0
Layer 4,5,6 → GPU 1
Layer 7,8,9 → GPU 2

全部 GPU 都跑完整的前向+反向
但 optimizer states / gradients / params
分片存储

Q2：为什么 ZeRO 要分 Stage 1/2/3，而不是直接全部拆了（Stage 3）更好呢？

先看三个 Stage 各自切了什么：

Stage 3 的代价不可忽视：

Q3：梯度 G 为什么要全部存下来，不能算完一层立刻覆盖掉吗？

表面上看"算完一层梯度就用掉然后覆盖"是合理的，但这里有一个关键时序问题。

你说的"边算边发、覆盖内存"是 DDP 的 bucket AllReduce 优化，确实是可行的——反向传播期间，每算完一层就立刻把这层梯度发出去做 AllReduce，同时继续算前面层的梯度。

但 AllReduce 的结果（平均后的梯度）必须写回显存，一直驻留到 Optimizer Step 用完为止：

关键在于 Step 1 和 Step 2 是串行的两个阶段，Step 1 全部完成后 Step 2 才开始。所以在 Step 2 进行期间，所有层的"平均后梯度"必须同时在显存里，等 optimizer 按参数顺序逐一消费完才能释放。

Q4：混合精度里 FP32 副本是干什么的？既然 FP16 精度不够，为什么不直接全用 FP32，还要多存一份 FP16 不是更浪费？

先澄清一个常见误解：FP32 副本不是用来做前向/反向计算的，它只在 Optimizer Step 的最后一步用。

FP32 副本解决的问题：更新步长太小被 FP16 舍入成 0。

那为什么不全用 FP32，反而还要多存一份 FP16？