续训之后 Adam 从头开始
主线位置：训练主线的 optimizer step，跨过一次存盘和读回。优化器要在新进程里认出每个参数，才能接着用它的矩。
版本：PyTorch 2.15.0（CPU）、burn 0.22.1（Flex 设备，autodiff）
flowchart TB tr["训练 5 步"] --> ms["model.safetensors"] tr --> os["optim.bpk 按 ParamId"] ni["新进程 init"] -->|"新 ParamId"| ld["load_from"] ms --> ld ld -->|"文件里没有 id"| m["模型 新 id"] os --> ol["optim.load 成功"] m --> st["step 查不到"] ol --> st st --> rs["矩和步数从零开始"]
实验代码：PyTorch 侧 resume.py，burn 侧 crate pitb-opt 的 bin resume（feature std、flex、autodiff、optim、store、safetensors，release）。模型、权重、数据同 A1（见一步），Adam 学习率 0.01，两边 ε 都是 1e-8（burn 默认 1e-5，见优化器状态）。两边都分两次运行：save 训练 5 步、存盘、在同一进程里接着训练 5 步当参照；resume 在新进程里读回再训练 5 步。平台 M2 Max、macOS 27、rustc 1.99.0。
PyTorch 的写法
存一个字典，里面是模型和优化器的 state_dict；续训时先按代码建出模型和优化器，再各自 load_state_dict：
if sys.argv[1] == "save":
model = a1_model()
opt = torch.optim.Adam(model.parameters(), lr=0.01)
steps(model, opt, 5)
torch.save({"model": model.state_dict(), "optim": opt.state_dict()}, "out/ckpt/ckpt.pt")
show("continuous, steps 6-10", steps(model, opt, 5))
else:
ck = torch.load("out/ckpt/ckpt.pt")
model = a1_model()
model.load_state_dict(ck["model"])
opt = torch.optim.Adam(model.parameters(), lr=0.01)
opt.load_state_dict(ck["optim"])
show("resumed", steps(model, opt, 5))
第 6 到 10 步更新前的 loss：
continuous, steps 6-10 0.97519988 0.94985896 0.92572045 0.90281999 0.88104737
resumed 0.97519988 0.94985896 0.92572045 0.90281999 0.88104737
new Adam, no state 0.97519988 0.95010257 0.92367375 0.89935386 0.87664640
续训和连续训练逐位相同。最后一行是只读模型、不读优化器的对照：第 6 步的 loss 在更新前算，三行相同，从第 7 步起不同。
搬到 burn 的现象
burn 侧照同样的结构写。模型用 burn-store 存成 safetensors（和 PyTorch 生态互通的格式），优化器用 optim.save：
let mut model = a1_model(&device);
let mut optim = adam();
steps(&mut model, &mut optim, 5, &device);
let mut st = SafetensorsStore::from_file(format!("{DIR}/model.safetensors")).overwrite(true);
if let Err(e) = model.save_into(&mut st) {
println!("save_into: {e}");
}
if let Err(e) = optim.save(format!("{DIR}/optim.bpk")) {
println!("optim.save: {e:?}");
}
新进程里先按代码建出模型，读参数，再读优化器（src 依次是 "bpk" 和 "safetensors"，bpk 的结果放在「改法」里）：
let mut model = a1_model(&device);
let fresh_id = model.fc1.weight.id;
let loaded = match src {
"bpk" => model.try_load_file(format!("{DIR}/model.bpk")).map_err(|e| e.to_string()),
_ => {
let mut st = SafetensorsStore::from_file(format!("{DIR}/model.safetensors"));
model.load_from(&mut st).map(|_| model).map_err(|e| e.to_string())
}
};
// ... 出错时打印并跳过
println!("{src}: fc1.weight id {} (fresh init had {fresh_id})", model.fc1.weight.id);
let mut optim = match adam().load(format!("{DIR}/optim.bpk")) {
// ... 同上
};
let before = optim.to_record().len();
let mut losses = steps(&mut model, &mut optim, 1, &device);
println!("{src}: record tensors after load {before}, after first step {}", optim.to_record().len());
losses.extend(steps(&mut model, &mut optim, 4, &device));
两次加载都返回 Ok：
continuous, steps 6-10 0.97520036 0.94985944 0.92572105 0.90282059 0.88104808
...
safetensors: fc1.weight id qt86fsd8fkfmi (fresh init had qt86fsd8fkfmi)
safetensors: record tensors after load 8, after first step 16
resumed from safetensors 0.97520036 0.95010293 0.92367423 0.89935452 0.87664700
new Adam, no state 0.97520036 0.95010293 0.92367423 0.89935452 0.87664700
「续训」那一行和「新 Adam、不读状态」那一行逐位相同：读进来的 Adam 状态一点都没用上。没有报错，没有警告。优化器记录里的张量读进来是 8 个（4 个参数各 2 个矩），走了一步变成 16 个：旧的 8 个原封不动，又为新参数建了 8 个。
（burn 和 PyTorch 两份「连续」的 loss 差在 7.2e-7 以内，来历见优化器状态。）
原因
两边的优化器状态用的键不同。PyTorch 的 state_dict() 按 model.parameters() 的顺序编号 0、1、2、3，新模型参数顺序不变，就能对上。burn 的 ModuleOptimizer::to_record 按 ParamId 给每项命名（{ParamId}.momentum.moment_1 这样，见优化器状态的「状态存在哪」），load_record 照原样按这些 id 放回去。它手里没有模型，没法检查这些 id 在模型里还在不在，所以不会因为 id 对不上而失败。
ParamId 在新进程里是新的：a1_model 里每个参数是 Param::from_tensor 装的，LinearConfig::init 也一样，每次都生成随机的新 id（参数与遍历）。读权重时，burn-store 的加载器在 Applier::map_float 里决定用哪个 id：
Some((tensor, source_id)) => {
// Prefer the source's persisted ParamId so optimizer state (keyed by ParamId)
// survives save/load cycles, but keep the target's id when the source format
// carries no identity of its own.
param.transform_for_load(tensor, source_id.unwrap_or(param_id))
}
safetensors 文件里没有 ParamId，source_id 是 None，参数保留新建模型的 id（上面输出里 qt86fsd8fkfmi 前后不变就是这个）。读 PyTorch 的 .pt 也一样，from_pytorch 的文档写着「PyTorch carries no parameter identity, so param_id stays None」。
接下来的 step 在 map_float 里按参数的 id 从状态表里 remove_entry，查不到，就当成第一次见到这个参数：矩从零开始，步数回到 1，偏差修正也从头算。旧 id 的状态留在表里，不会再被用到，也不会被清掉。
同样的事发生在任何换了 ParamId 的地方：训练中途用 Param::from_tensor 重新装参数（优化器状态的「动脑筋」逐位复现了这种情况），或者重新 init 一份模型再套进旧的数值。
改法
模型存成 burn 自己的 burnpack（save_file），它连 ParamId 一起存，加载时 Applier 用文件里的 id：
// save 那次
if let Err(e) = model.clone().save_file(format!("{DIR}/model.bpk")) {
println!("save_file: {e}");
}
// resume 那次：上面 match 的 "bpk" 分支
"bpk" => model.try_load_file(format!("{DIR}/model.bpk")).map_err(|e| e.to_string()),
saved fc1.weight id 3svgcbl5ee9o8
continuous, steps 6-10 0.97520036 0.94985944 0.92572105 0.90282059 0.88104808
bpk: fc1.weight id 3svgcbl5ee9o8 (fresh init had 2d720q6ijejvs)
bpk: record tensors after load 8, after first step 8
resumed from bpk 0.97520036 0.94985944 0.92572105 0.90282059 0.88104808
加载后 fc1.weight 的 id 换成了存盘时的那个，续训和连续训练逐位相同，第一步之后记录里仍是 8 个张量。save_file 不需要额外的 feature，格式见保存与加载。Learner 每个 epoch 存的模型 checkpoint 也是 burnpack 文件（训练循环）。
safetensors 留着给别的工具读，checkpoint 另存一份 burnpack。只有 safetensors（比如从 PyTorch 那边转过来的权重）时，没有能对上的优化器状态，就别读旧的优化器记录，直接建一个新的；效果一样，但不会在记录里留下一半用不上的状态。
怎么检查
- 续训后第一步更新之前和之后各打一次
optim.to_record().len()，像上面那样。状态接上了，两个数相同（A1 上 8 和 8）；变大（8 到 16）说明有参数换了 id，它们的状态重新开始了。 - 存盘时记下一个参数的 id（
model.fc1.weight.id，打印成 base32），续训加载后再打印一次，两者应相同。 - 在同一个进程里做一次「存、读、接着训练」，和不中断的训练逐位比较 loss，像本页这样。差一点（第 7 步差 2.4e-4）不是数值误差，是状态没接上。
- 训练中途要改参数的值，用
param.map(..)保住 id，不要Param::from_tensor（参数与遍历）。
小结
- PyTorch 的优化器
state_dict按参数序号存，新建模型后load_state_dict就能接上；burn 按ParamId存，新进程里参数的 id 是新的。 - 模型从 safetensors 或
.pt读回时保留新建模型的 id，optim.load照样成功，Adam 静默从第 1 步重来；A1 上续训的 loss 和「新 Adam」逐位相同，记录从 8 个张量变成 16 个。 save_file存的 burnpack 带ParamId，读回后续训和连续训练逐位相同。- 用
optim.to_record().len()在续训第一步前后对比，能发现状态没接上。