数据截至 (上游 commit 32e301ffaf5a)
04 · Offload：把状态搬到 CPU 与 NVMe
这一章讲什么： 分片（第 1、3 章）把显存除以 N；offload 在这之上再除以「卡外存储的容量」——优化器状态、梯度、参数都可以落到 CPU 内存或 NVMe 盘。本章讲两条 offload 轴（optimizer / param）各自的实现，以及它们共同的底层：pinned 内存 + AIO 异步交换。
1. 它要解决的小问题
ZeRO-3 之后每卡只剩 16Ψ/N，但 N 不能无限加（卡就那么多、通信随 N 变差）。一个 500B 的模型在 64 卡上仍要每卡约 125GB 的模型状态——还是放不下。
卡外还有两级便宜得多的存储：
| 层级 | 量级（单节点典型） | 带宽代价 |
|---|---|---|
| GPU HBM | 80GB/卡 | —— |
| CPU 内存 | 1-2TB/节点 | 走 PCIe，比 HBM 慢一个量级 |
| NVMe SSD | 10TB+/节点 | 再走一级，还吃并发读写模式 |
offload 的问题因此不是「能不能搬」，而是：怎么搬才能不让 GPU 饿着——答案全部是流水线与异步。
配置的两条轴
offload 在配置里是两个独立旋钮（deepspeed/runtime/zero/offload_config.py）：
{
"zero_optimization": {
"stage": 3,
"offload_optimizer": {"device": "cpu", "pin_memory": true},
"offload_param": {"device": "nvme", "nvme_path": "/local_nvme",
"max_in_cpu": 1e9, "buffer_count": 5}
}
}
offload_optimizer（DeepSpeedZeroOffloadOptimizerConfig，:55）：优化器状态 + fp32 主权重 + 梯度分片外移，stage 1/2/3 都可用。注意 docstring 明确写了：无论 device 选 cpu 还是 nvme，优 化器计算（Adam 更新）都在 CPU 上跑。offload_param（DeepSpeedZeroOffloadParamConfig，:21）：fp16 参数分片外移，只在 stage 3 有意义（1/2 参数本来就不分片）。
2. 直觉：计算与搬运并行
offload 版的训练步，本质是三条流水叠在一起：
时间 ───────────────────────────────────────────────────►
GPU: [forward/backward 计算] [收梯度分片] [等 CPU 结果?]
PCIe: [梯度分片 H2D→CPU 异步拷] [更新后参数 D2H 回拷]
CPU: [DeepSpeedCPUAdam 逐 sub-group 更新]
NVMe: [参数分片 swap-in 提前于 all-gather] [优化器状态 swap-out 流水线]
关键设计常量：pin_memory: true（两个配置的默认值都是 true，见 offload_config.py:45、:75）——只有页锁定内存才能做真正的异步 DMA 拷贝，否则「异步」拷贝会先隐式走一遍可分页缓冲，带宽腰斩。
3. 真实实现 I：offload_optimizer（stage 1/2 路径）
stage 1/2 开 offload_optimizer 后（stage_1_and_2.py 里 self.cpu_offload=True，:188-191），变化集中在三处：
① 优化器换成 CPUAdam。 引擎侧强制：zero_use_cpu_optimizer() 为真时，配置造的优化器直接是 DeepSpeedCPUAdam（deepspeed/runtime/engine.py:2180-2197 的 get_optimizer_configuration）；用户自带 GPU 优化器会报错并指引换（engine.py:2043-2048）。
② 梯度在 CPU 上累积。 backward 里梯度分片 reduce 到本 rank 后，异步拷到 CPU pinned 缓冲逐 micro-step 累加（async_accumulate_grad_in_cpu_via_gpu，stage_1_and_2.py:1526）；到累积边界再 _finalize_cpu_offload_gradient_accumulation（:989）搬回 GPU 算范数、查溢出。
③ step 在 CPU 上原地更新。 step 的 cpu_offload 分支（:2322 起）：fp32 分片的梯度就在 CPU，DeepSpeedCPUAdam.step 直接在 CPU 上跑 Adam；更新结果先拷进一个 fp16 的 pinned 暂存 param_buffer_of_bit16_for_cpu_offload_groups，再 non_blocking=True 异步拷回 GPU 分片，最后照常 all-gather。
DeepSpeedCPUAdam（deepspeed/ops/adam/cpu_adam.py:13）本身是编译型 C++ 内核：构造函数 CPUAdamBuilder().load() 编译/加载 ds_opt_adam，step（:107）逐参数调 adam_update（:160）——SIMD 向量化的 Adam/AdamW。step 里有硬断言「参数必须在 CPU 上」（:137-139），这就是 stage 1/2 里那条「fp32 分片的 device 是 cpu」的约定。
4. 真实实现 II：stage 3 的 offload
stage 3 里 offload 的装配点在 _configure_offloading（deepspeed/runtime/zero/stage3.py:776）：
| 配置 | 落点 | 效果 |
|---|---|---|
offload_optimizer.device = cpu | self.offload_optimizer=True | fp32 主权重分片、Adam 状态在 CPU；_optimizer_step 走 CPUAdam |
offload_optimizer.device = nvme | + self.swap_optimizer=True | 再进一步：优化器状态按 sub-group 换到 NVMe |
offload_param.device = cpu | self.offload_param=True | fp16 参数分片常驻 pinned CPU 内存，all-gather 时从 CPU 取 |
offload_param.device = nvme | + self.params_in_nvme_and_cpu=True | 参数分片写盘，CPU 只缓存 max_in_cpu（默认 1e9 元素） |
4.1 优化器状态的 NVMe 交换：按 sub-group 流水
NVMe 优化器交换由 _configure_tensor_swapping（stage3.py:794）建 swapper：offload_optimizer.pipeline 为真用 PipelinedOptimizerSwapper（deepspeed/runtime/swap_tensor/pipelined_optimizer_swapper.py:52），否则用 PartitionedOptimizerSwapper（partitioned_optimizer_swapper.py:26）。
与第 3 章的 step 循环咬合得很紧——_prepare_sub_group 在更新 sub-group i 之前先经 _optimizer_states_and_gradient_swap_in 换入状态（stage3.py:2435，swap_in 调用在 :2440），pipelined 版会同时异步预取 sub-group i+1（pipelined_optimizer_swapper.py:102 的 async_parameter）；更新完 _release_sub_group 把状态 swap_out 回盘。sub-group 因此成为「GPU 正在算的 / 正在换入的 / 正在换出的」三段流水的调度单位——sub_group_size 旋钮控制粒度。
4.2 参数分片的 NVMe 交换：AIO buffer pool
参数侧是 AsyncPartitionedParameterSwapper（deepspeed/runtime/swap_tensor/partitioned_param_swapper.py:37），机制拆四点：
- AIO 引擎二选一：装了 GDS（GPU Direct Storage）走
GDSBuilder，否则AsyncIOBuilder的 Linux AIO（:96-97）；读写各有独立 handle，队列深度、块大小由aio_config控制（:111-120）。 - 固定 buffer pool：构造时一次性分配
buffer_count（默认 5）个 pinned buffer（:124-133），参数分片换入换出都在池里轮转——不做临时分配，I/O 对齐到aligned_bytes（:101）。 - 太小的参数不换：
swappable_tensor（:139）要求numel × element_size ≥ min_aio_bytes，小于最小 I/O 粒度的参数留在内存。 - 换入有「提前一拍」的调度：协调器的
__prefetch_nvme_param_partitions（deepspeed/runtime/zero/partitioned_param_coordinator.py:667）在 all-gather 预取之前，先把更靠后的 NVMe 分片swap_in进 buffer——磁盘这一拍藏在网络那一拍前面。
offload_ratio（配置 ratio，offload_config.py:99）还允许部分 offload：_create_fp32_partitions（stage3.py:1055-1064）按比例把一部分 sub-group 留在 GPU 用 backup_optimizer（torch AdamW）更新，其余走 CPUAdam——GPU 有余量时的折中。
5. 代价与坑
- 吞吐换规模，掉多少看硬件。 offload 后 step 的瓶颈通常变成 PCIe 与 SSD；官方的定位就是「让 1 张卡训百亿级、几台机器训千亿级」，不 是「不掉速」。代码里给不出固定数字，需自己 profile（开
wall_clock_breakdown）。 - pinned 内存是硬通货。 两个配置的
pin_memory默认 true 且注释明确警告：仅在宿主机 memlock 上限紧（ulimit -l）或常驻内存极有限时才关（offload_config.py:45-51的 docstring）。关了等于放弃异步拷贝。 - CPUAdam 的核数争用。 dataloader worker、主进程、CPUAdam 共享 CPU；
cpuadam_cores_perc（默认 0.8，offload_config.py:105）限制 Adam 用核比例，留核给数据管道。 - NVMe 路径要「每台机器本地盘」。 swap 文件按 rank 分目录（
optimizer_utils.py:174的base_folder/optimizer/rank{r}、partitioned_param_swapper.py:152的{param_id}_param.tensor.swp）；nvme_path指向共享网络盘会把 I/O 模型彻底打乱。 - offload 与若干特性的互斥。 Muon + reduce_scatter + optimizer offload 直接
ValueError（stage_1_and_2.py:233）；offload_param只在 stage 3 生效（配到 stage 2 上是静默无效，不会报错）。 - 更新的 offload 引擎在别的目录。 本 commit 已有
deepspeed/runtime/zenflow/（ZenFlow，stall-free offload）与deepspeed/runtime/superoffload/（SuperOffload，面向 GH200/GB200 类 superchip）两条新路线，分别由zenflow配置与offload_optimizer.super_offload开关启用（引擎分流见engine.py:2401与:2485的两处优化器类选择）；本文主线只讲经典 ZeRO-Offload/Infinity 路径，二者是演化方向而非替代。
6. 代码地图
| 主题 | 文件路径 | 符号名 |
|---|---|---|
| offload 配置 | deepspeed/runtime/zero/offload_config.py | DeepSpeedZeroOffloadParamConfig、DeepSpeedZeroOffloadOptimizerConfig、OffloadDeviceEnum |
| stage 3 offload 装配 | deepspeed/runtime/zero/stage3.py:776、:790 | _configure_offloading、_configure_tensor_swapping |
| stage 3 step 的换入换出 | deepspeed/runtime/zero/stage3.py:2427、:2449 | _prepare_sub_group、_release_sub_group |
| CPU Adam | deepspeed/ops/adam/cpu_adam.py:13 | DeepSpeedCPUAdam.step |
| stage 1/2 梯度 CPU 累积 | deepspeed/runtime/zero/stage_1_and_2.py:1526、:989 | async_accumulate_grad_in_cpu_via_gpu、_finalize_cpu_offload_gradient_accumulation |
| 参数 NVMe 交换 | deepspeed/runtime/swap_tensor/partitioned_param_swapper.py:37 | AsyncPartitionedParameterSwapper、swap_in、swappable_tensor |
| 优化器状态 NVMe 交换 | deepspeed/runtime/swap_tensor/pipelined_optimizer_swapper.py:52、partitioned_optimizer_swapper.py:26 | PipelinedOptimizerSwapper、PartitionedOptimizerSwapper |
| NVMe 预取调度 | deepspeed/runtime/zero/partitioned_param_coordinator.py:667 | __prefetch_nvme_param_partitions |
| 新一代 offload | deepspeed/runtime/zenflow/、deepspeed/runtime/superoffload/ | （目录级，自行下钻） |