CPU Offloading：把显存墙拆成 GPU / CPU / NVMe 三层

先算清一个参数在混合精度 Adam 下占多少字节，才能判断该卸载什么、不该卸载什么。

ZeRO-Offload §2 给出的账本是：fp16 参数 2 B、fp16 梯度 2 B、fp32 master weight 4 B、Adam 的 momentum 与 variance 各 4 B，合计 16 B/参数。 ZeRO-Infinity §3 把 fp32 梯度也计入 optimizer states，得到 20 B/参数。两篇的记账口径不同，本页的计算统一采用 ZeRO-Offload 的 16 B 口径，并在手算里注明。

按十进制 GB（1 GB = 10⁹ B）计；未计 activation、临时 buffer 与碎片。

这不是“把 70B 权重塞进一张卡”的问题：fp16 权重本身就要 140 GB，单张 80GB GPU 仍放不下；offload 解决的是把 980 GB 的梯度与 optimizer states 移出 GPU，从而让每张卡保留完整的 fp16 参数，只把梯度与 optimizer states 的分片移到 CPU。 真正的前向/反向仍必须在 GPU 上完成，ZeRO-Offload §3.2 的原则是只把计算复杂度低于 O(MB) 的更新类计算放到 CPU。

推理的账本不同：bf16 权重是 2 B/参数，70B 约 140 GB；此外还有随上下文增长的 KV cache。vLLM 把 KV 作为独立的 block 池管理，并在 vllm/config/cache.py:399-408 提供单独的 kv_offloading_size，与权重显存分开核算。

ZeRO-Offload 的脚注明确把 activation offload 排除在范围外；ZeRO-Infinity §3 则用 activation checkpointing 把 residual states 压到可接受范围。也就是说，offload 主要处理 model states，不负责解决所有显存项。

重计算用算力换显存，并行用通信与多卡换显存，offload 用带宽换容量；先分清代价，才能判断何时该用哪条。

FSDP 落在“并行”和“offload”之间：它先把参数、梯度、optimizer states 分片到各 rank，减少数据并行的冗余；再打开 CPUOffload(offload_params=True) 时，才把 sharded 参数与梯度进一步下沉到 CPU。 PyTorch 源码把这点写得很直白：参数不在计算中时 offload 到 CPU，梯度也一起下沉，因此 optimizer step 在 CPU 上跑（torch/distributed/fsdp/api.py:231-242）。FSDP2 的 CPUOffloadPolicy 则明确写到“parameters, gradients, and optimizer states to CPU”，并在 all-gather 前把 shard 拷回设备（torch/distributed/fsdp/_fully_shard/_fsdp_api.py:239-255）。

三条路不是互斥的。ZeRO-Offload 本身就是 ZeRO-2 的分片 + CPU offload；ZeRO-Infinity 又在 ZeRO-3 的分片之上加 CPU/NVMe；FSDP 也可以和 activation checkpointing 叠加。真正要问的是：被搬到慢速层级的那部分，是否能在下一次使用前被取回，且取回时间小于它遮住的 compute。

训练 offload 的对象是 optimizer states / gradients / params；推理 offload 的对象是静态权重或动态 KV。对象不同，数据流和付费频率也不同。

3.1 训练侧：ZeRO-Offload 的 CPU 层，ZeRO-Infinity 的 NVMe 层

ZeRO-Offload §3.5 的策略是：fp16 参数留在 GPU，fp16 梯度、fp32 master weight、momentum、variance 全部放 CPU，参数更新在 CPU 执行。 前向/反向仍在 GPU；反向产生的梯度可以先在 GPU 上暂存，再按参数或小批量拷到 CPU，并与剩余反向计算重叠（§4.1）。每步 CPU↔GPU 的最小通信量是 4M：2M 梯度下行 + 2M 更新后的参数上行。

ZeRO-Offload §3.5/§4.1；M 为参数个数，按每参数字节数直接换算。

ZeRO-Infinity 在这之上再加一层 NVMe。它不再要求参数常驻 GPU：partitioned params、gradients、optimizer states 都可以在 GPU / CPU / NVMe 之间移动，由 infinity offload engine 负责调度；memory-centric tiling 把大算子切成 tile，避免单个 layer 必须靠 tensor parallel 才放得下（§5.1.3/§8.5）。 DeepSpeed 的实现把这些抽象落成配置：offload_optimizer 支持 stage 1/2/3，device 可选 cpu 或 nvme，且 optimizer computation 无论如何都在 CPU（deepspeed/runtime/zero/offload_config.py:55-62）；offload_param 只允许 stage 3（同文件 21-28）。 NVMe 路径不是简单 write()：它走 swap_tensor 的异步 IO 与 buffer pool（deepspeed/runtime/swap_tensor/partitioned_optimizer_swapper.py:26-37），并用 pipeline_read/pipeline_write 把 tile 读写与 CPU optimizer 计算重叠（offload_config.py:84-94）。

FSDP 的 CPU offload 与 ZeRO-Offload 不是同一个切法。ZeRO-Offload 保留完整 fp16 参数副本在每张 GPU 上，只切梯度与 optimizer states；FSDP 先把参数也按 rank 切开，再把 shard 放到 CPU。 FSDP1 在初始化时要求开启 offload 的模块留在 CPU（torch/distributed/fsdp/_init_utils.py:672-673），local shard 与 _cpu_grad 预分配在 pinned CPU（_flat_param.py:1258-1265）；前向 all-gather 前用 flat_param_to(device, non_blocking=True) 把 shard 拷回设备（同文件 1335-1338）。梯度在 post-backward 阶段以 non_blocking 拷到 CPU（_runtime_utils.py:969-991），并在 CPU optimizer step 前显式同步（同文件 1107-1111）。 FSDP2 的 CPUOffloadPolicy(pin_memory=True) 默认用 pinned memory，让 H2D/D2H 拷贝可以与 compute 重叠（_fsdp_api.py:248-252）；shard 在 all-gather 前以 non_blocking=True 拷回设备（_fsdp_param.py:1203-1206）。

3.2 推理侧：weight offload 与 KV cache offload 是两笔账

vLLM 的 cpu_offload_gb 是“每张 GPU 卸载多少 GiB 权重”。源码把它描述成一种虚拟扩大显存的方式：24GB 卡设 10，就相当于 34GB 卡；代价是每次 forward 都要从 CPU 取一部分权重（vllm/config/offload.py:23-32）。 UVA 后端把参数放在 pinned CPU 内存，再用 CUDA view 让 GPU 直接读，省掉显式 copy，但仍然消耗 PCIe 带宽（vllm/model_executor/offloader/uva.py:23-28）；prefetch 后端则按层分组，用独立的 copy stream 异步 H2D，把传输藏到前面的层后面（vllm/config/offload.py:54-68、offloader/prefetch.py:171-172）。

KV cache offload 搬的是另一块内存。vLLM 用 kv_offloading_size 指定 CPU 侧的 KV 缓冲大小，默认不开启；kv_offloading_backend 支持 native 与 lmcache（vllm/config/cache.py:399-408）。 SimpleCPUOffloadConnector 的 CPU 容量默认是 8 GiB（.../simple_cpu_offload_connector.py:55），按 block 做 LRU；scheduler 只在 prefix cache 命中时请求把连续的 CPU block 取回（vllm/v1/simple_kv_offload/manager.py:371-374）。 传输走低优先级 stream，让 KV I/O 给 compute 让路（worker.py:170-173）；CPU block 是 pinned 的，并且只在 GPU→CPU store 完成后才允许被 load（同文件 241-251）。

两者可以同时开，但代价相加：weight offload 让每个 token 多读权重，KV offload 让 block 换入换出多走一次 PCIe；如果模型已经能装下，或者 decode 对每 token 延迟敏感，先不要动 weight offload。

ZeRO-Offload 把每步 CPU↔GPU 流量压到 4M；是否值得，取决于传输时间与计算时间的比值，以及这段传输能否与反向/下一轮前向重叠。

按 §3.5/§4.1，ZeRO-Offload 每步只搬 4M 字节：2M 梯度从 GPU 到 CPU，2M 更新后的参数从 CPU 回 GPU。这个量比 L2L 的 28M 小 7×（28M / 4M，ZeRO-Offload §6.2.2）。通信量小是它能用 CPU optimizer 还不掉太多吞吐的前提。

全双工对称链路下，下行梯度与上行参数可同时进行，墙钟时间取 2M/BW；串行则取 4M/BW。计算量按 ZeRO-Infinity §4 Eq.7 的 8ND 口径，T 为每步 token 数，P 为有效算力。

ZeRO-Infinity §4 进一步给出了效率公式：efficiency = AIT × BW / (AIT × BW + peak)，其中 AIT 是 arithmetic intensity（总计算 / 总数据移动）。带宽越低，AIT × BW 越小，通信占比越高。 §4.2 还指出一个关键差别：optimizer states 在 fwd/bwd 结束后才更新，不能像梯度那样与反向计算重叠，所以它对带宽的要求更高；NVMe 路径必须把 tile 读、CPU compute、tile 写三者流水起来（§5.2.2 补充：optimizer states 可以跨设备并行更新，但更新时点仍在 fwd/bwd 之后）。

这张表的含义是：在 V100 世代，每设备 3 GB/s 的慢速内存带宽就够；算力涨 10× 时，需要 30 GB/s。PCIe 4.0 ×16 的理论单向带宽是 31.5 GB/s，刚好压线；按 80% 有效系数算是 25.2 GB/s，低于 30 GB/s。 ZeRO-Infinity §9 给出的对照是 Summit 用 NVLink 把 V100 连到 CPU 内存，达到 40 GB/s/GPU；这解释了为什么论文强调“慢速内存”不等于“慢速链路”，NVLink 级别的 CPU 链路可以满足 10× V100 的要求。

NVLink 本身主要是 GPU↔GPU 链路；CPU↔GPU 是否走 NVLink，取决于平台（Summit 是论文给出的例子）。在普通 PCIe 服务器上，offload 的重叠能力就取决于 pinned memory、copy stream 和 kernel 是否能同时跑。ZeRO-Offload 用 DPU（delayed parameter update）把 optimizer update 与下一轮 forward 重叠，代价是 1-step staleness（§6.2.4）；ZeRO-Infinity 用 tile pipeline 把 NVMe 读写、CPU Adam 和 GPU compute 串起来；FSDP 用 non_blocking 拷贝；vLLM 用独立 copy stream 与低优先级 KV stream。

把两条数据流画在同一张分层图上，才能看出谁在重复付费：训练搬梯度与参数，weight offload 搬权重，KV offload 搬 block。

训练一步的时序（ZeRO-Offload）

1. GPU 用常驻的 fp16 参数做 forward，不需要访问 CPU。
2. backward 产生梯度；梯度可以在计算过程中按参数或小批量拷到 CPU，与剩余反向图重叠（§4.1）。
3. 多卡时先做 reduce-scatter，每张卡只把自己负责的梯度分片写到 CPU（§4.2）。
4. CPU 上用 CPU-Adam 更新 fp32 master weight、momentum、variance，并得到更新后的参数分片。
5. 参数分片从 CPU 回 GPU，再 all-gather 成完整的 fp16 参数，供下一轮 forward 使用。
6. DPU 打开时，optimizer update 可以与下一轮 forward 重叠，代价是参数晚一步生效（1-step staleness，§6.2.4）。

推理的两种时序

• Weight offload：第 k 层 forward 前，UVA 直接读 CPU 上的权重，或 prefetch 线程把第 k+1 层权重拷到 GPU；每个 token 都要为被卸载的权重付一次 PCIe 带宽。KV 仍在 GPU，不因 weight offload 变小。
• KV cache offload：GPU KV 池按 LRU 把不活跃 block 写到 CPU；新请求命中 prefix 时，把连续 block 取回 GPU。当前 decode token 的 KV 留在 GPU，因此不会在每个 token 上重复搬整段 context。它降低的是 KV 常驻占用，不降低权重占用。

这些数字都带硬件与 batch 条件；换链路、换模型规模或换并行策略，结论会变。

三个反直觉的点

• 卡多了，offload 可能反而更慢。ZeRO-Offload §6.2.3 报告：到 64/128 张 GPU 时 ZeRO-2 开始超过 ZeRO-Offload，因为两边都能跑相近的 batch，而 ZeRO-Offload 多了一层 CPU↔GPU 通信。offload 是容量工具，不是无限扩展的加速器。
• 能装下不等于能扩展。L2L 在单卡上能训练 17B，比 ZeRO-Offload 的 13B 还大，但多卡训练时模型规模不再随卡数增长（§6.2.1/§6.2.2）。offload 策略决定了容量能否随卡数一起走。
• superlinear 来自聚合带宽与 CPU 算力。ZeRO-Infinity §8.3 在 1T 模型上从 64 卡扩到 512 卡，超过线性；原因是每加一个节点，PCIe/NVMe 聚合带宽和 CPU optimizer 算力也一起增加。反过来说，如果链路或 CPU 不随节点增长，这个收益会消失。

ZeRO-Infinity §8.2 还给出一个边界信号：20T 时吞吐从 10T 的 43 TFLOPS/GPU 降到 34 TFLOPS/GPU，原因不是 NVMe 带宽，而是每 GPU 的 batch 太小、CPU 内存放不下 activation checkpoint。offload 把容量问题转移到 CPU/NVMe 后，新的瓶颈会出现在主机内存与 batch 上。

offload 是把容量问题换成带宽问题；带宽不够、无法重叠或 CPU 侧成为瓶颈时，收益为负。

一张决策表

不要把训练 offload 的数字外推到推理。训练每步搬一次梯度与参数，推理 weight offload 每个 token 搬一次被卸载的权重，KV offload 每个 block 搬一次 KV；三者的频率、体积、可重叠性都不同。ZeRO-Offload 的 4M/步和 vLLM 的 cpu_offload_gb 不是同一个量。

也不要把 CPU offload 的数字外推到 NVMe。ZeRO-Infinity 的 NVMe 层需要 tile 读写、CPU 内存中转与 fragmentation 控制；§5.2.2 明确提醒，NVMe↔CPU↔GPU 的往返会造成 CPU/GPU 内存碎片，容量够也可能 OOM。NVMe 适合极端规模，不适合作为默认的第一层。