训练状态与激活卸载：按下一次使用时刻安排存储

01 · 训练内存为什么比模型文件大得多

训练既要做前向，也要反向并更新参数。除权重外，还可能常驻梯度、FP32 主权重、优化器动量，以及等待反向的激活。一个对象暂时不用，不代表可以丢掉；卸载的思路是保存到设备外，再在需要时恢复。

它主要解决“设备放不下”，速度收益来自允许更合适的批量、避免其他更昂贵的方案等间接效果。单纯把同样训练从设备内存换到 CPU/NVMe，不会让数据访问自动变快。

依据：DeepSpeed · ZeRO-Offload与CPU优化器 ↗ · DeepSpeed · ZeRO ↗

02 · 四类对象，下一次使用时刻不同

对象	典型使用时机	卸载要配套什么
参数	所属模块前向/反向时	提前回迁；分片时可能还需 AllGather；保存正确权重版本。
梯度	反向产生，随后同步/累积/更新	归约完成、累积规则与 CPU 优化器消费顺序。
优化器状态	更新对应参数时	状态可留在 CPU 并由 CPU 优化器计算，不必每层全部搬回设备。
保存激活	某段前向之后，到对应反向之间	保存所需 Tensor，释放其他引用；反向前恢复原值与元数据。

Adam 的一阶/二阶动量跟踪过去梯度，不能用本步激活重建。激活通常只活到本步反向结束。二者生命周期不同，因此不宜将所有卸载对象都画成“前向后搬出、反向前搬回”。

依据：DeepSpeed · ZeRO-Offload与CPU优化器 ↗ · PyTorch 2.14 · save_on_cpu与saved_tensors_hooks ↗

03 · 和 ZeRO、激活重计算是什么关系

方法	改变什么	增加什么
ZeRO/FSDP 分片	不同 rank 各保存一部分状态	使用时汇集/归约；不等于设备外存储。
激活重计算	不保存部分激活，反向时重做前向	额外计算；重建随机与权重语义。
卸载	保存原数据但放在 CPU/NVMe	复制、容量管理，可能还有 CPU 更新或存储 IO。

三者可以组合。例如参数分片并卸载到 CPU、激活一部分重算另一部分卸载。但峰值要加上临时完整参数、回迁激活和通信缓冲；“每项单独省了多少”不能直接相加。

依据：DeepSpeed · ZeRO ↗ · PyTorch · 激活 Checkpoint（main） ↗

04 · 保存一个激活，需要满足哪些条件

教学示意 · 不代表硬件实测

前向生成 A，后续前向消费者先使用 A
提交 A 的 D2H；源 A 与主机目标都保留
复制完成 + 设备其他读者结束 → 可释放设备副本
反向即将用 A → 提交 H2D 到有效设备缓冲
复制完成 → 对应反向读取 A
反向/复制引用全部结束 → 两侧按策略回收

如果后续层还持有 A 的 view，或 Python 对象仍保留同一存储，移动一份保存激活不一定释放原设备存储。卸载收益应按真正能释放的 storage 计算。目标设备、shape、dtype 和内容都要恢复；数值若压缩，还需要单独验证误差。

前向对应哪一版参数，反向就要保持那次计算的语义。不能为了取回方便，使用已经 optimizer.step 更新后的权重替代原前向权重。异步 CPU 读取也必须等待 D2H 完成。

依据：PyTorch 2.14 · save_on_cpu与saved_tensors_hooks ↗ · PyTorch 2.14 · CUDA stream、等待与内存生命周期 ↗ · PyTorch 2.14 · Tensor view、stride与连续布局 ↗

05 · 预取提前多久，才够覆盖回迁

自设单对象例子：一块激活为 256 MiB，有效 H2D 带宽 32 GiB/s，忽略固定开销时复制需 0.25/32 秒 ≈ 7.8125 ms。若反向需要它之前只有 3 ms 独立计算，则仍须等约 4.8125 ms；提前 10 ms 且复制资源空闲时可按时到达。

用更易算的时间线：需要时刻 D=12 ms，复制耗时 H=8 ms。如果提前 L=3 ms，9 ms 开始、17 ms 完成，等待 5 ms；若 L=10 ms，2 ms 开始、10 ms 完成，不等。它只有在 2 ms 前已保存完，且设备有落点时才可这么早预取。

反向要在12 ms使用，回迁赶得上吗？

提前L ms回迁H ms

自设单对象时间线：数据在0 ms已保存，通路空闲、缓冲充足、独立计算原计划在12 ms结束。未计D2H、多对象排队与争用。

早点预取增加设备驻留时间，多个提前取回的对象会叠加峰值。选择要同时满足计算依赖、容量预算和共享复制队列，而不是把所有参数与激活一起提前取回。

06 · 把容量和带宽一起算

自设 1 亿参数的 FP32 Adam 两个动量状态共 1亿×2×4B=800 MB（十进制）。若这两项驻留 CPU，可从设备状态账本移走这 800 MB，但仍要计参数、梯度、主权重、激活、工作区和复制缓冲。CPU 更新还需要计算时间及梯度/更新结果的传输。

另一例：10 个可卸载激活各 256 MiB，共 2.5 GiB。完整 D2H+H2D 字节为 5 GiB。假设共享通路有效总吞吐 32 GiB/s且串行处理这些字节，纯复制时间下界约 156.25 ms；实际能隐藏多少要看时间线。真实双向通路与争用不同，需要分别测量。

07 · 先看最小入口，再研究完整调度

PyTorch CUDA 保存激活入口 · 非完整启动脚本，未作硬件性能测试

import torch
# model与inputs已在CUDA设备；optimizer已创建
optimizer.zero_grad(set_to_none=True)
with torch.autograd.graph.save_on_cpu(pin_memory=True):
    prediction = model(inputs)
    loss = loss_fn(prediction, labels)
loss.backward()
optimizer.step()

save_on_cpu 作用于上下文内算子为反向保存的 Tensor，不是把模型所有对象都迁到 CPU，也不保证提前到达的自定义预取计划。用 profiler 看恢复发生在哪里；需要更复杂方案时再研究 saved_tensors_hooks 或框架卸载调度。锁页缓冲不能无上限占用主机内存。

DeepSpeed ZeRO-3 CPU 卸载配置片段 · 其余训练配置、初始化和版本约束需配套

{
  "zero_optimization": {
    "stage": 3,
    "offload_optimizer": {"device": "cpu"},
    "offload_param": {"device": "cpu"}
  }
}

offload_optimizer 和 offload_param 是不同选项；本片段没有激活卸载。CPU 优化器路径可在主机更新状态，然后把设备需要的参数准备好。不能根据这个配置推断每层会把两个 Adam 动量一起回迁。

依据：PyTorch 2.14 · save_on_cpu与saved_tensors_hooks ↗ · DeepSpeed · ZeRO ↗ · DeepSpeed · ZeRO-Offload与CPU优化器 ↗

08 · 教学脚本验证数值，硬件实验验证收益

CPU 脚本用 h=tanh(w×x)、y=v×h、Loss=½(y−target)² 做前向，将保存值复制到模拟主机记录再释放设备记录。恢复后算梯度，与常驻解析梯度和有限差分对照；故意用错版本的 v 时会暴露梯度不一致。

固定有效 batch、数据与精度，比较 Loss、梯度、更新后参数及恢复训练。
测设备和主机峰值、锁页内存、复制字节、CPU 优化器时间与存储吞吐。
测稳态 step time、有效 tokens/s，以及同质量训练总时间。
在时间线区分反向本身、回迁等待、参数 AllGather 等通信和 CPU 更新。

性能比较要写明基线。如果基线 OOM，只能说明可训练规模变化，不能报告与“无法执行的每步耗时”相比的加速倍数。

09 · NVIDIA 与昇腾怎样对照

NVIDIA 可沿 PyTorch 保存激活 API 和 DeepSpeed 卸载路径阅读。昇腾要分别核对框架的参数/优化器/激活卸载支持与 CANN 搬运约束；不将 CUDA 接入代码当作 NPU 的已验证方案。两侧对比应列出同样对象、驻留层级、精度、并行分片与回迁计划。

依据：CANN Community 9.1.0 · aclrtMemcpyAsync与Host内存约束 ↗ · MindSpeed Core · 特性清单 ↗

10 · 自查与扩展实验

解释为什么“CPU Adam”不要求把所有动量搬回设备再更新。
把交互中的回迁耗时改成 16 ms；最多提前 12 ms 时还剩多少等待？
如果一块激活的 view 仍存活，卸载该激活为什么可能没有释放预期设备容量？

流程位置与关联方法

这个方法可能出现在多个位置；计算线包含模型计算与算子实现，执行线包含运行时、内存和通信。

T-B3 · 调度前向与 Loss T-B4 · 管理激活与调度反向 T-B6 · 调度参数更新与预取

资料与核查范围

原理依据以原始论文和官方文档为主；算例与时间线是教学推导。文档核查日期：2026-10-10。引用 main / latest 文档时，实际接口以安装版本为准。

代码为接入片段或教学伪代码；本页没有声称完成 NVIDIA / 昇腾硬件实测。