DeepSpeed ZeRO-3 参数详解 结合4张16g显卡微调14b模型

DeepSpeed ZeRO-3 参数详解 结合4张16g显卡微调14b模型

{
  "zero_optimization": {
    "stage": 3,
    "overlap_comm": true,
    "contiguous_gradients": true,
    "sub_group_size": 1e9,
    "reduce_bucket_size": 5e8,                      // 必须改大！
    "stage3_prefetch_bucket_size": 5e8,             // 必须改大！
    "stage3_param_persistence_threshold": 1e5,
    "stage3_max_live_parameters": 1e9,              // 必须改大！
    "stage3_max_reuse_distance": 1e9,               // 必须改大！
    "stage3_gather_16bit_weights_on_model_save": false,
    "offload_optimizer": {
      "device": "cpu",
      "pin_memory": true
    }
    // 确认没有 offload_param！
  }
}

在这里插入代码片### 1. "stage": 3
ZeRO 分三个阶段，逐步切分更多东西：

2. "overlap_comm": true

作用：通信和计算重叠执行。

关闭时（串行）：
计算 ████████          通信          ████████  计算
                      ████
开启时（重叠）：
计算 ████████
                      通信 ████
                              计算继续...

类比：你一边做饭（计算），一边让外卖小哥送食材（通信），不用等食材到了才开火。
必须开，否则每一步都要等通信完了才能计算，纯浪费时间。

3. "contiguous_gradients": true

作用：把分散在内存各处的梯度拷贝，合并成一块连续内存。

关闭时：
梯度碎片 [块1]...[块2]...[块3]...[块4]  → 通信要发4次小包
开启时：
梯度连续 [块1块2块3块4]                   → 通信1次大包搞定

类比：搬家时把零碎东西装进一个箱子，一次搬走，比来回跑4趟快。
必须开，减少通信次数和内存碎片。

4. "sub_group_size": 1e9

作用：ZeRO-3 参数归约时，把参数分成子组，每个子组 1e9 = 1GB 一批处理。

在计算机中：
**1 GB = $10^9$ 字节 = 1,000,000,000 字节**
而 `1e9` 是科学计数法，正好等于 $10^9$，也就是 1,000,000,000。
所以 **1e9 字节 = 1 GB**。

28GB 参数，sub_group_size=1e9：
组1 [1GB] → 归约
组2 [1GB] → 归约
...
组28 [1GB] → 归约
总共 28 轮，每轮处理量适中

如果设太小（比如默认 1e6 = 1MB），28GB 要拆 28000 轮，启动开销巨大。
类比：高速公路收费站，1GB = 一次通过一列车队，太小的车队会导致频繁开关闸。

5. "reduce_bucket_size": 5e8 ⭐ 最关键

作用：梯度归约（All-Reduce）时，一次通信的最大数据量。

5e8 = 500MB：
28GB 梯度 ÷ 500MB = 56 次通信 ✅
2e7 = 20MB（你之前的值）：
28GB 梯度 ÷ 20MB = 1400 次通信 ❌

这就是你 239 秒/步 的元凶！ 每次 all-reduce 有固定启动延迟（哪怕是 SHM 也有 ~0.01 秒）：

6. "stage3_prefetch_bucket_size": 5e8 ⭐ 第二关键

作用：在当前层计算时，提前预取下一层参数的批量大小。

前向传播过程：
Layer 1 计算 ████████
              Layer 2 参数预取 ████████（和计算重叠！）
                       Layer 2 计算 ████████
                                    Layer 3 参数预取 ████████

如果预取太小（2e7 = 20MB），每次只预取一点点，下一层计算时参数还没到，CPU→GPU 的搬运速度跟不上计算速度，计算单元要等。

2e7 时（预取太慢）：
Layer1计算 ████████  等待...  Layer2计算 ████████  等待...
                    ████████                      ████████
                    预取太慢！                    预取太慢！
5e8 时（预取够快）：
Layer1计算 ████████
          Layer2预取 ████████（并行完成）
                   Layer2计算 ████████（无缝衔接！）

类比：炒菜时提前把下道菜的食材从冰箱拿出来放案板上（预取），不用等开火了才去翻冰箱。

7. "stage3_param_persistence_threshold": 1e5

作用：参数体积 ≤ 100KB 的小张量，不分片，每张卡保留完整副本。

为什么？小张量分片反而更慢：
大张量（比如 4096×4096 = 32MB）：
  分片：每卡 8MB，all-gather 一次拿到完整 → 值得
小张量（比如 32×32 = 2KB）：
  分片：每卡 0.5KB，all-gather 通信开销 >> 节省的显存 → 不值得！
  不分片：每卡 2KB，省掉一次通信 → 更快

1e5 = 100KB 是合理的阈值，小于这个的参数（如 LayerNorm 的 bias）就别分片了。
类比：大件家具拆开搬（分片），小零碎直接整件拿（持久化），不值得拆。

8. "stage3_max_live_parameters": 1e9 ⭐

作用：在 GPU 显存中同时驻留的参数最大总量。

1e9 = 1GB：
前向传播时，当前层 + 预取层 + 刚用完还没回收的层，总共不超过 1GB
2e7 = 20MB（你之前的值）：
同一时刻只有 20MB 参数在 GPU 上！

2e7 就是灾难！ 意味着 GPU 几乎是空的，参数随用随丢，用完立刻回收，下一步又要重新 all-gather。这导致：

• 预取缓冲区太小，计算和通信无法有效重叠
• 反复 all-gather 同样的参数（反向传播时还要再来一遍）
• GPU 计算单元大量空闲等参数
  1e9 = 1GB 意味着 GPU 可以同时持有几十层的参数，刚算完的层可以暂时留着给反向传播用，不用重新通信获取。
  类比：
• 2e7：办公桌只有巴掌大，一次只放一页纸，看完了收起来，下次又要去档案室取
• 1e9：办公桌够大，一次铺开几十页，翻来覆去看都方便

9. "stage3_max_reuse_distance": 1e9 ⭐

作用：一个参数被使用后，多久之后还会被用到的判断阈值。如果距离 ≤ 1e9，就先不回收，留在 GPU 上。

前向传播：Layer1 → Layer2 → ... → Layer32
反向传播：Layer32 → Layer31 → ... → Layer1
Layer1 的参数：
  前向用了 → 等了 31 层 → 反向还要用！
  
2e7 时：距离太远，立刻回收 → 反向时重新 all-gather → 又一次通信
1e9 时：判断距离在范围内，先留着 → 反向时直接用 → 零通信

2e7 又是灾难！ 前向传播用完的参数立刻被回收，反向传播时又要重新 all-gather，等于通信量翻倍。
1e9 让前向传播的参数大概率能保留到反向传播复用，避免重复通信。
类比：

• 2e7：刚看完的书立刻还图书馆，明天又要借
• 1e9：看完先放书架上，知道过两天还要翻

10. "stage3_gather_16bit_weights_on_model_save": false

作用：保存 checkpoint 时，不收集完整 16 位权重。

true：保存时 all-gather 完整模型 → 占大量显存 → 可能 OOM
false：保存时只存分片的参数 → 省显存，但保存的文件不能直接用

因为你用了 LoRA + save_only_model: true，保存的只是 LoRA 权重，不需要收集完整模型。
类比：存钱只存零钱（LoRA 权重），不用把全家现金都搬到银行（完整模型）。

11. "offload_optimizer": {"device": "cpu", "pin_memory": true}

作用：把优化器状态（Adam 的 m、v）放到 CPU 内存。

14B 模型的 Adam 优化器状态（fp32）：
  m (一阶矩): 14B × 4 bytes = 56 GB
  v (二阶矩): 14B × 4 bytes = 56 GB
  总计: 112 GB！
放 GPU：每卡 28 GB → 爆了 ❌
放 CPU：112 GB 分散到 4 卡的 CPU 内存，每卡 ~28 GB → CPU 内存够 → ✅

但你是 LoRA，优化器状态只有 LoRA 参数的（约 184M params × 12 bytes ≈ 2.2GB），其实不大。保留 offload_optimizer 作为保险，万一没 offload 时 OOM。
pin_memory: true：使用页锁定内存，CPU→GPU 拷贝更快（走 DMA 而不是普通拷贝）。
类比：优化器状态像仓库里的存货，不常用，放在大仓库（CPU），用专门的快速通道（pin_memory）取。

12. 为什么删掉 offload_param？

offload_param 开启时：
每一步：CPU → GPU 搬运 28GB 参数 → 计算 → GPU → CPU 回收 28GB
搬运速度：CPU-DDR4 带宽 ~50 GB/s
单次搬运：28GB / 50 = 0.56 秒
每步要搬运 2 次（前向+反向）：1.12 秒
加上通信延迟：实际 5-15 秒纯搬运开销
offload_param 关闭时（参数在 GPU）：
参数已经在 GPU 上（ZeRO-3 分片），只需要 all-gather
SHM all-gather 速度：~100 GB/s
几乎无搬运开销

只有 16GB 卡装不下完整参数时才需要 offload_param，ZeRO-3 已经帮你切分了，每卡只有 7GB，装得下！

总结对比表