读完这篇，你能做什么
能区分 ZeRO 各阶段分摊哪些状态，解释 FSDP 怎样临时聚合参数，并验证激活 checkpoint 前后输出与梯度一致。
学习理由、准备要求与知识点
为什么学它
不同省内存方法减少的是不同部分，也会增加通信或计算。组合使用前，要先找到内存峰值来源，并检查重复执行会不会改变行为。
先读这些内容
知识点清单
- 模型状态与激活分类
ZeRO三阶段的累计分片- FSDP参数聚合与梯度分片
- 稳态和临时峰值
- 激活重计算
- 非重入checkpoint
- 随机数与副作用
- 预算下候选选择
运行前先确认Java17、PyTorch2.9.1 CPU。真实运行非重入activation checkpoint；FSDP/ZeRO以论文和官方文档解释，未运行多GPU分片。Java背包模型只针对相互独立的候选，不是框架自动分区器。
01先弄懂它是什么
训练内存不够时，先查空间花在哪里。权重、梯度、优化器动量属于模型状态；为了反向传播而保留的前向中间值叫激活。分片会把部分模型状态分摊给 W 个 rank，减少每个进程都存完整副本的浪费；激活检查点则少保存中间值，反向需要时再算一遍。两种方法可以一起用，后者是在用额外计算换存储空间，这里的 checkpoint 也不是保存训练断点文件。
ZeRO 分三个逐步增加的阶段：先分摊优化器状态，再加上梯度，最后连参数也一起分摊。设参数、梯度、优化器分别占 P、G、O 字节，忽略通信缓冲时，每个 rank 的理想主体大小依次为 P+G+O/W、P+(G+O)/W、(P+G+O)/W。实际峰值还要加上激活、临时聚齐的完整参数、未均匀分片的部分和运行时缓冲，因此这些公式不是内存上限保证。
FSDP 可以先让不同进程各存一部分参数；某个模块要计算时，再把它需要的完整参数临时收集起来，用完后按策略重新分片。这样减少的是“每个进程一直拿着多少参数”，也会增加通信；具体何时收集、何时分回去，要看版本和配置。
FSDP1 的 FullyShardedDataParallel 与 FSDP2 的 fully_shard 不是换个名字就能互换的接口。本文讲共同思路，可执行实验只在单个 CPU 上检查激活重计算，没有运行分布式 FSDP，也不能用这次实验宣称测过多机分片效果。
文中这几个词，先认一认
- AllGather
- 把各个 rank 上的参数片段收集起来，让当前计算拿到所需的完整参数。
- ReduceScatter
- 先把各 rank 的梯度按规定合并，再把合并结果分开，每个 rank 只保留自己负责的那一段。
- 激活检查点
- 前向时少存一些中间值，反向需要时重新计算对应片段，用计算时间换内存。它不同于把训练进度保存到磁盘。
02看看它怎么工作
假设一个计算块使用输入 x 和权重 w。普通 autograd 会保存反向要用的中间值；checkpoint 改为保留必要的边界信息，求梯度时再执行相关操作。本篇 PyTorch2.9 基线显式使用 use_reentrant=False。这个非重入实现拿到所需中间值后可以提前停止，因此开始重算一次，不代表一定跑完整个函数。放在函数尾部的写入、计数等副作用，不能据此假定会固定执行几次。
重新计算必须得到与原来兼容的中间结果，否则梯度会变。例如 dropout 使用随机数，重算时需要对应的随机状态；PyTorch 默认会保留相关状态，但换设备或使用未覆盖的外部随机源仍要另查。重算还可能再次读取已变化的全局数据、增加日志计数、更新 BatchNorm 运行统计。本实验的计数只观察是否进入函数，不参与数值计算，避免计数本身改变结果。
CPU实验用固定输入与权重分别执行普通和检查点路径，比较输出、输入梯度与权重梯度；加入dropout并重置种子验证默认随机状态处理。Java算法把相互独立、可选择保存的激活候选抽象成背包：每个候选占用整数容量、带来避免重算的收益，在预算内最大化总收益。真实计算图有依赖与生命期，这个简化模型只训练预算思考，不能直接优化任意Transformer检查点布局。
先跑通一个最小例子
运行 experiment.py。真实调用torch.utils.checkpoint，比较CPU输出与梯度；计数只证明重算发生，未测量GPU节省比例。
import torch
import torch.nn.functional as F
from torch.utils.checkpoint import checkpoint
torch.set_num_threads(1)
def run(use_checkpoint, dropout):
torch.manual_seed(29)
x = torch.linspace(-1, 1, 8, dtype=torch.float64).requires_grad_()
w = torch.tensor(0.7, dtype=torch.float64, requires_grad=True)
calls = [0]
def block(value):
# 观察入口次数，计数不能进入模型数值
calls[0] += 1
hidden = torch.sin(value * w)
if dropout:
hidden = F.dropout(hidden, p=0.25, training=True)
return hidden.square()
# 显式选择非重入版本，保留默认随机状态机制
y = checkpoint(block, x, use_reentrant=False) if use_checkpoint else block(x)
y.sum().backward()
return y.detach(), x.grad.clone(), w.grad.clone(), calls[0]
plain, checked = run(False, False), run(True, False)
assert torch.allclose(plain[0], checked[0])
assert torch.allclose(plain[1], checked[1])
assert torch.allclose(plain[2], checked[2])
assert plain[3] == 1 and checked[3] == 2
# 使用相同种子重新开始，核对dropout重计算的梯度
random_plain, random_checked = run(False, True), run(True, True)
assert torch.allclose(random_plain[0], random_checked[0])
assert torch.allclose(random_plain[1], random_checked[1])
assert torch.allclose(random_plain[2], random_checked[2])
print("7 activation checkpoint equivalence checks passed")
核对运行结果普通路径一次入口，checkpoint路径两次入口；两种情况下输出和梯度均一致，包括dropout。
正在加载流程图…
查看流程图源码（Mermaid）
flowchart TD
A["识别内存主体"] --> B["权重梯度优化器"]
A --> C["前向激活"]
B --> D["跨rank分片"]
D --> E["计算前聚合与梯度归约"]
C --> F["保存边界并丢弃中间值"]
F --> G["反向按需重算"]
E --> H["测量峰值与耗时"]
G --> H按顺序看，每一步做了什么
- 在相同训练配置下定位状态或激活主导的内存。
- 列出分片主体和计算时暂时聚合的参数规模。
- 对纯计算块添加显式非重入checkpoint。
- 比较输出、梯度、随机行为和耗时，再扩大模型。
03对照实现代码
PyTorch v2.9.1 · checkpoint；_checkpoint_without_reentrant_generator；_CheckpointFrame
下例为使用真实checkpoint的原创包装；它不是内部保存张量钩子的复刻，也不实现FSDP。
from torch.utils.checkpoint import checkpoint
def memory_tradeoff(block, x, enabled):
# 保留普通路径，作为数值正确性对照
if not enabled:
return block(x)
# 非重入模式在反向需要时重建中间张量
# block应避免改变后续计算依赖的外部状态
return checkpoint(block, x, use_reentrant=False, preserve_rng_state=True)
checkpoint入口选择重入或非重入实现；显式指定参数可以避免依赖默认变化。
非重入路径通过生成器与保存张量钩子组织重算，调用方不需要手工重新执行backward。
early_stop允许只重建必要中间结果，所以函数执行次数和执行到末尾的次数不是一回事。
随机状态保存有成本与覆盖边界；它不让任意外部副作用变得幂等，也不自动保存训练恢复文件。
代码先调用谁，再调用谁
- checkpoint选择非重入生成器路径
- 首次前向建立边界与保存张量钩子
- 反向请求中间量时触发重计算
04动手写一个例子
独立激活候选的预算背包
每个候选有正整数容量cost和非负重算节省value，只能保存一次。求预算budget内最大节省。容量单位由调用方统一，例如MiB。
- 输入 1
cost=[3,2,2],value=[5,3,3],budget=4- 输出 1
6；保存后两个候选
令dp[b]表示已经看过的候选在容量不超过b时的最大收益，初始全部为零代表不保存。
加入一个候选时比较不保存与保存后的收益；容量从大到小扫描，避免当前候选在同一轮被重复使用。
这个动态规划精确解决独立零一选择模型。真实激活生命周期互相重叠、重算依赖相互影响时，收益未必可加，需要更强的图调度模型。
为什么这样写是对的处理前i个候选后，dp[b]等于只从这i个候选各取至多一次且容量不超过b的最大总收益。
n是候选数，B是整数容量预算；这是伪多项式复杂度，输入预算采用二进制表示时不能称为一般多项式算法。
保存为 Main.java，运行 javac --release 17 -encoding UTF-8 Main.java，再运行 java Main。程序包含固定输入自测，断言失败时会抛出错误。
public class Main {
static long save(int[] cost,long[] value,int budget){
if(cost.length!=value.length || budget<0 || budget>1_000_000)
throw new IllegalArgumentException("shape or teaching capacity bound");
long[] dp=new long[budget+1];
for(int i=0;i<cost.length;i++){
if(cost[i]<=0 || value[i]<0)throw new IllegalArgumentException("item");
// 逆向更新保证当前候选最多使用一次
for(int b=budget;b>=cost[i];b--){
// 超出long范围显式失败，不让收益回绕
long take=Math.addExact(dp[b-cost[i]],value[i]);
dp[b]=Math.max(dp[b],take);
}
}
// 容量不要求填满，零预算对应零收益
return dp[budget];
}
static void check(boolean ok){if(!ok)throw new AssertionError();}
public static void main(String[] args){
check(save(new int[]{3,2,2},new long[]{5,3,3},4)==6);
check(save(new int[]{2},new long[]{3},4)==3);
check(save(new int[]{1},new long[]{2},0)==0);
check(save(new int[]{},new long[]{},3)==0);
check(save(new int[]{2,3},new long[]{4,5},5)==9);
try{save(new int[]{0},new long[]{1},2);throw new AssertionError();}
catch(IllegalArgumentException expected){}
try{save(new int[]{1,1},new long[]{Long.MAX_VALUE,1},2);throw new AssertionError();}
catch(ArithmeticException expected){}
System.out.println("7 independent activation budget checks passed");
}
}
这些特殊情况也要试一试
- 容量必须统一单位，过细粒度会使动态规划数组过大。
- 零容量候选被拒绝，负收益也不符合本模型定义。
- 收益累加超出
long范围显式报错。
进阶选读 · 含生命期的图调度
场景保存一个中间量会减少多个下游重算，但它又与其他激活同时存活。
做法先分析计算图依赖和张量生命期，使用真实运行profile校准代价；比较手工分段、选择性检查点和编译器策略，逐个校验梯度及峰值。
复杂度一般调度含组合选择，独立背包的O(nB)不适用于任意图；搜索成本取决于约束和算法。
好处与代价精细策略可能增加编译与维护成本，粗分段虽非最优但更容易调试和迁移。
05项目里怎么用
分片并不是越细越好。一次聚合太大的模块，可能临时占用大量内存；拆得太碎，又可能把时间耗在大量小通信上。估算峰值时，除了参数分到每台设备后有多少，还要加上同时预取的模块、通信缓冲、激活和优化器更新中的临时数据，再结合网络性能选择模块边界。
激活检查点通常优先放在占用较大且重计算相对可接受的计算块。观察BatchNorm运行统计、缓存写入和随机生成是否重复发生；对于训练随机性，先用确定性版本做输出与梯度对照，再恢复随机层。默认保存随机状态不能覆盖所有外部库与设备迁移情形。
想在崩溃后继续训练，需要把模型、优化器、随机状态和采样进度保存下来。分片训练还要决定保存完整状态还是分片状态，并验证恢复时是否允许改变 world size。激活 checkpoint 只决定一次前向、反向过程中怎样保存或重算中间值，不能替代这类持久化断点。最后仍应在目标硬件上一起测实际内存峰值和有效 Token 吞吐，确认省内存的代价。
容易踩的坑
- 把激活checkpoint当成磁盘训练断点。
- 用总状态除以world size作为FSDP峰值保证。
- 在重算块里更新影响数值的全局计数或缓存。
- 不区分FSDP版本就照搬封装和状态保存API。
几种做法，怎么选
表格有四列，窄屏可左右滑动查看代价和适用条件。
| 方案 | 优势 | 代价与限制 | 适用条件 |
|---|---|---|---|
| 状态分片 | 降低模型状态复制 | 增加通信和状态管理 | 参数与优化器主导的内存瓶颈 |
| 激活检查点 | 减少保存的中间张量 | 反向重算与随机状态成本 | 长序列或深网络激活主导场景 |
| CPU卸载 | 用主机容量缓解设备限制 | 传输带宽与延迟可能主导 | 容量优先且接受额外耗时的配置 |
06检验是否读懂了
先自己想、自己写，再展开答案核对。场景题是假设的练习情境；遇到还没学过的内容，可以做完基础后再回来。
先用自己的话解释
01ZeRO第二阶段相对第一阶段多分什么？
展开参考答案与推导
在优化器状态之外进一步分片梯度；参数仍按该阶段基本模型复制。
02激活检查点为什么不改变理论目标？
展开参考答案与推导
在重计算与首次前向数值等价时，仅改变保存中间值的方法；随机与副作用破坏等价条件时会出错。
03为什么背包算法是伪多项式？
展开参考答案与推导
运行次数依赖预算数值B，而不是只依赖表示B所需的log B位数。
换个情况，还会处理吗
01分片后仍在某一层前向OOM怎么办？
展开参考答案与推导
检查该聚合单元完整参数、预取缓冲和激活峰值，调整分组或预取策略，并测量而非只看稳态分片字节。
02加checkpoint后梯度不同如何排查？
展开参考答案与推导
固定种子与输入，逐块对照，检查随机状态、设备迁移、全局变量以及训练模式副作用。
03要从中断处继续训练，是否只需checkpoint函数？
展开参考答案与推导
不需要这个函数来完成持久化；应保存并恢复模型、优化器、日程、随机和数据游标等状态。
写代码试一试
01推导第三阶段理想主体。
- 输入 / 条件
P=4,G=4,O=8，W=4，单位GiB- 输出 / 预期
4GiB每rank
提示 三类主体全部分摊。
展开参考答案与推导
(4+4+8)/4=4，但激活和暂时聚合另计。
02为何背包容量要逆序？
- 输入 / 条件
一个候选容量2收益3，预算4- 输出 / 预期
正确收益3
提示 正序会复用当前轮dp。
展开参考答案与推导
正序可能先写dp[2]=3再用它写dp[4]=6，错误重复同一个候选；逆序避免复用。
03设计checkpoint正确性测试。
- 输入 / 条件
同输入、权重、种子，两条执行路径- 输出 / 预期
输出及输入和权重梯度接近
提示 不能只比较loss。
展开参考答案与推导
逐张量比较前向输出和各梯度，并加入随机层、零输入与重算次数观察。
读完之后，独立完成
课后作业：自己做一遍
本篇共 8 道作业，先完成 6 道基础与综合任务，再按前置知识挑战进阶题。其中包含 1 道思维拓展，试着换个条件，看看原来的办法是否还管用。不附参考答案。
能区分 ZeRO 各阶段分摊哪些状态，解释 FSDP 怎样临时聚合参数，并验证激活 checkpoint 前后输出与梯度一致。
本篇重点加练
非重入重计算仍依赖执行条件稳定，重点检验闭包状态和随机源是否被完整保留。
- 非重入checkpoint
- 随机数与副作用
重点内容会多练几次。先说清楚怎么做，再试试空数据、出错或条件改变时会发生什么。
先自己做，卡住时记下问题，再回来查。问答题写清“为什么、什么时候适用”；动手题留下输入、输出和运行步骤，方便下次重做。
基础必做
- 作业 01原理问答
建立权重、梯度、优化器状态和激活的归属表，分别描述
ZeRO各阶段与一种 FSDP 策略；标注哪些信息需要临时聚合。- 提交内容
- 状态分布图
- 怎样算完成
- 分片对象不遗漏
- 激活不混入模型状态
练到的知识点模型状态与激活分类 · ZeRO三阶段的累计分片 · FSDP参数聚合与梯度分片
- 作业 02代码推演
模型含三个大小不同的模块，在 world size=4 的理想化分片假设下，画前向和反向的聚合时序，估算各时刻主体存储。
- 提交内容
- 时序与预算表
- 怎样算完成
- 聚合窗口明确
- 理想估计不声称真实显存
练到的知识点FSDP参数聚合与梯度分片 · 稳态和临时峰值
综合应用
- 作业 03编程实践
在 CPU 含 Dropout 的小网络中比较普通执行与非重入 checkpoint，核对输出及梯度；记录随机数保存设置。
- 提交内容
- CPU 实验与等价检查
- 怎样算完成
- 随机条件一致
- 反向实际执行被验证
练到的知识点激活重计算 · 非重入checkpoint · 随机数与副作用
- 作业 04排错实验
给被 checkpoint 的函数加入可观察计数或列表写入，记录普通与重计算版本的事件次数，并审阅该副作用是否符合训练契约。
- 提交内容
- 副作用复现与契约审查
- 怎样算完成
- 调用次数有证据
- 不把重计算视作无副作用操作
练到的知识点激活重计算 · 随机数与副作用
- 作业 05面试场景
有四个重计算候选，各自节省内存与增加时间不同。自建整数预算表并实现选择器，报告预算紧张与不可行两种情况。
- 提交内容
- CPU 选择器及选择记录
- 怎样算完成
- 预算单位明确
- 选择可与穷举核对
练到的知识点稳态和临时峰值 · 预算下候选选择
- 作业 06排错实验
在 CPU 非重入 checkpoint 函数中读取一个可修改的闭包配置，前向后、反向前改变该配置；记录输出与梯度并核对重计算契约，不进行任何分布式性能推断。
- 提交内容
- 故障实验与状态依赖表
- 怎样算完成
- 配置变更时间明确
- 普通执行提供同条件对照
练到的知识点非重入checkpoint · 随机数与副作用
进阶挑战 · 掌握前置后选做
- 作业 07设计评审
选做真实多 GPU 分片实验：预先列出设备、模型、通信后端和峰值采集协议；没有设备时交付待执行方案及不能验证的结论。
- 提交内容
- 实测报告或待执行协议
- 怎样算完成
- 真实通信证据明确
- 不以单机模拟代替多卡验证
练到的知识点ZeRO三阶段的累计分片 · FSDP参数聚合与梯度分片 · 稳态和临时峰值
思维拓展
先写出你的猜想，再用例子或实验验证。答案可以不同，但要说清依据、适用情况，以及还有什么没弄懂。
- 作业 08思维拓展 · 条件变化
被 checkpoint 的函数同时使用
PyTorch与自建 Python 随机生成器，自动保存随机状态是否足以保证重计算一致？设计限定 CPU 的验证并定义需额外管理的状态范围。- 提交内容
- 随机源矩阵与小实验
- 怎样算完成
- 每个随机源身份公开
- 验证结果限定实际版本
练到的知识点非重入checkpoint · 随机数与副作用
完成后，再检验一次
给每题标记“独立完成 / 查阅后完成 / 尚未完成”。查过资料或没做出的题，记下卡在哪里。隔一天换一组输入再试，能解释结果和出错原因，再把这个知识点标为掌握。
继续阅读
先看能帮你完成当前练习的部分。例子跑通后，再读更深入的解释。
- 必读参考文档英文
PyTorch2.9.1 checkpoint 源码 （新标签页）↗
核对非重入路径、随机状态与提前停止。
- 深入参考文档英文
PyTorch2.9 FSDP1 文档 （新标签页）↗
阅读参数聚合、分片策略和状态保存边界，避免与FSDP2接口混淆。
- 深入参考文档英文
ZeRO 原论文 （新标签页）↗
理解三阶段累计分片及残余内存分类。
- 深入参考文档英文
PyTorch2.9 checkpoint 文档 （新标签页）↗
核对
use_reentrant和副作用的使用限制。