Offload：CPU / NVMe 卸载的代价与收益

6.1 Offload 的核心思想

ZeRO 已经把每张卡的占用降到 $16\Psi/N$。如果卡还是不够呢？把更不常用的张量搬到 CPU 内存乃至 NVMe SSD，需要时再搬回 GPU。这就是 ZeRO-Offload / ZeRO-Infinity。

显存等级金字塔：

越往下，容量越大但带宽越低。Offload 的本质就是用更便宜的容量换更小的吞吐损失。

6.2 可以 offload 的对象

6.3 三种 offload 配置组合

组合 A：ZeRO-2 + CPU offload（性价比最高）

仓库自带 examples/deepspeed/ds_z2_offload_config.json（在 LLaMA-Factory 仓库里也有镜像）：

{
  "bf16": { "enabled": "auto" },
  "zero_optimization": {
    "stage": 2,
    "offload_optimizer": {
      "device": "cpu",
      "pin_memory": true
    },
    "overlap_comm": true,
    "contiguous_gradients": true,
    "allgather_partitions": true,
    "reduce_scatter": true,
    "reduce_bucket_size": "auto"
  }
}

显存收益（13B 模型，8 卡 A100 40G）：

$$ M_\text{ZeRO-2} = 2\Psi + \tfrac{14\Psi}{8} \approx 3.75\Psi \quad\Rightarrow\quad M_\text{+ offload\_optim} = 2\Psi + \tfrac{2\Psi}{8} = 2.25\Psi $$

即从 $\approx$ 49 GB 降到 $\approx$ 29 GB，开 offload 后 13B 在 40G 卡上就能跑了。

组合 B：ZeRO-3 + CPU offload（70B 神器）

{
  "bf16": { "enabled": "auto" },
  "zero_optimization": {
    "stage": 3,
    "offload_optimizer": {
      "device": "cpu",
      "pin_memory": true
    },
    "offload_param": {
      "device": "cpu",
      "pin_memory": true
    },
    "overlap_comm": true,
    "contiguous_gradients": true,
    "sub_group_size": 1e9,
    "reduce_bucket_size": "auto",
    "stage3_prefetch_bucket_size": "auto",
    "stage3_param_persistence_threshold": "auto",
    "stage3_max_live_parameters": 1e9,
    "stage3_max_reuse_distance": 1e9,
    "stage3_gather_16bit_weights_on_model_save": true
  }
}

用法：在单机 8 卡 A100 上训 70B 全参，offload_param=cpu + offload_optimizer=cpu 把每张卡占用压到能塞下的水平。代价是训练速度下降 30%-60%，因为每层 forward 都要从 CPU 拷参数到 GPU。

组合 C：ZeRO-Infinity（NVMe offload）

{
  "zero_optimization": {
    "stage": 3,
    "offload_optimizer": {
      "device": "nvme",
      "nvme_path": "/local_nvme/ds_swap",
      "pin_memory": true,
      "buffer_count": 4,
      "fast_init": false
    },
    "offload_param": {
      "device": "nvme",
      "nvme_path": "/local_nvme/ds_swap",
      "buffer_count": 5,
      "buffer_size": 1e8,
      "max_in_cpu": 1e9
    },
    "aio": {
      "block_size": 1048576,
      "queue_depth": 8,
      "thread_count": 1,
      "single_submit": false,
      "overlap_events": true
    }
  }
}

用法：CPU 内存也不够时，把对象进一步卸到 NVMe SSD。前置要求：

1. 装 libaio：apt install libaio-dev；
2. 编 aio op：DS_BUILD_AIO=1 pip install --force-reinstall deepspeed；
3. 给个真的 NVMe 盘路径（不能是 HDD，否则会慢到没法用）。

实测 NVMe offload + ZeRO-3 让一台 4 × V100 32G 上能训 GPT-3 13B 全参 —— 本来连个 7B 都装不下。代价是吞吐降到 1/3 左右。

6.4 速度损失到底是多少

非常粗的"经验值"（A100 80G、PCIe Gen4、NVMe 7 GB/s）：

选择思路：

• 能不开就不开（速度优先）；
• 能 CPU 就不上 NVMe；
• 非要 NVMe 也是本地 NVMe，不要走网络盘；
• 多机不要 offload_param 到 NVMe（每台机得有 NVMe）。

6.5 关键调优字段

6.6 用 Profiler 看 offload 是不是瓶颈

DS 自带的 wall clock breakdown 是判断 offload 是否拖速度的最快方法：

{
  ...
  "wall_clock_breakdown": true
}

训练时会按 step 打印每个阶段耗时：

step= 100 fwd= 320ms bwd= 410ms bwd_inner=190 bwd_allreduce=110 bwd_other=110
                                step= 230ms ...
optim_step= 95ms allgather= 32ms reduce_scatter= 21ms ...

如果 optim_step 远大于 fwd+bwd，多半是 CPU offload 的 cpu_adam 慢了——升级 CPU、关 offload，或加 nvme_path。

6.7 Offload 常见错误

6.8 OffloadStateTypeEnum：到底能 offload 哪些"状态"

翻 deepspeed/runtime/zero/offload_config.py 第 109–115 行能看到一个枚举，把"可被 offload 的对象"列得很清楚：

# offload_config.py:109–115
class OffloadStateTypeEnum(Enum):
    optim_states           = "optim_states"           # m, v 等 Adam 状态
    hp_params              = "hp_params"              # high-precision (fp32) master 权重
    lp_params              = "lp_params"              # low-precision 工作副本（fp16/bf16）
    lp_grads               = "lp_grads"               # 低精度梯度
    contiguous_grad_buffer = "contiguous_grad_buffer" # ZeRO 拍平后的连续梯度桶

这 5 种状态的关键区别：

新版 DS 允许通过 offload_states() / reload_states() 接口在运行时选择性 offload/reload 这些状态（RLHF 训完 actor、要释放显存做 reference rollout 时常用）。

6.9 NVMe AIO 缓冲区设计

开 NVMe offload 时容易遇到一个困惑：buffer_count / buffer_size / max_in_cpu 三个字段都跟 buffer 有关，到底谁管谁？翻 offload_config.py:21–50 的 DeepSpeedZeroOffloadParamConfig 能把它们分清楚：

数据流图：

调优经验：buffer_size 调到 5e8 上下，buffer_count 保持 4–8，max_in_cpu 不超过总 RAM 的 70%。aio.queue_depth=8 配合 overlap_events=true 通常打满 NVMe Gen4 的 7 GB/s。

6.10 SuperOffload 异步 CPU 进程模型

ch05.15 已经讲了 SuperOffload 跟 ZeRO-Infinity 在概念上的不同——这里展开它在 offload 字段里要怎么开：

"zero_optimization": {
  "stage": 3,
  "offload_optimizer": {
    "device": "cpu",
    "pin_memory": true,
    "super_offload":      true,    // 启用
    "cpuadam_cores_perc": 0.8,     // CPU adam worker 占用 80% CPU 核
    "ratio":              1.0      // 100% 走 SuperOffload；可调 0.5 等做 hybrid
  }
}

对应字段在 offload_config.py：super_offload（line 96）、cpuadam_cores_perc（99）、ratio（93）。后台跑的就是 superoffload_utils.py:38 的 superoffload_optimizer_worker()——一个独立 python 进程，通过 SimpleQueue 收任务、用 DeepSpeedCPUAdam 算 step、再把结果发回主进程。

注意：

• SuperOffload 只对 ZeRO-3 + offload_optimizer.device=cpu 有效；NVMe 路径不走它；
• worker 进程会吃 RAM，估算公式 ≈ cpuadam_cores_perc × 主进程 worker buffer；机器内存紧时反而拖慢；
• 有 issue 报告在 H100 单机 8 卡 70B 训练上 SuperOffload 比纯 ZeRO-Infinity offload 快 30-50%，主要省的是"GPU 等 CPU"那段时间。

6.11 DeepNVMe / GPU-Direct Storage

repos/deepspeed/csrc/gds/ 目录是 NVIDIA 的 GPU-Direct Storage（GDS）接入。区别于普通 AIO（GPU → CPU → NVMe 的三段拷贝），GDS 让 GPU 通过 NVLink/PCIe 直接读写 NVMe，跳过 CPU 内存：

开启方法：把 "offload_optimizer.device": "nvme" 配合 GDS 编译的 wheel 即可；DS 在 op 装好时自动选 GDS 后端。校验：ds_report 输出里 gds .................. [OKAY]。

6.12 ZenFlow 与 offload 的协议

普通认识里"开 offload = 慢"。ZenFlow 想颠覆的是这一点：

ZenFlow 与 SuperOffload 不冲突，可以叠加：ZenFlow 负责"只更新重要梯度"，SuperOffload 负责"被更新的那部分也异步算"。但叠加目前还在 blog 验证阶段，生产用先开一个跑稳定再考虑组合。