Checkpoint

下一次更新还取决于优化器历史、学习率进度、随机状态和下一批数据。Adam 的一阶、二阶矩重置后，即使模型权重相同，下一次更新也会改变。使用动态 loss scaling 时，还要恢复 scaler 状态。

DCP 保存你交给它的状态。模型和优化器可以通过 Stateful 封装；scheduler、训练步数、RNG 和数据进度也要接入保存与加载。模型配置、tokenizer 和代码版本则用于重建运行环境。

最容易管理的边界是一次完整的参数更新结束后：模型、优化器、scheduler 和训练步数都推进到同一步，数据状态指向下一批。保存接口应在这组状态一致时调用。

若在梯度累积中途保存，还要记录已累积的梯度、microbatch 位置及相关状态。仅保存模型和优化器，会丢掉这次尚未完成的更新。

state_dict() 中的张量仍可引用训练正在使用的存储。把这个字典交给后台线程后继续修改参数，后台读到的值就可能跨越不同训练步。

DCP 的 staging 将状态复制到供保存使用的缓冲区。源状态在复制完成前保持不变，复制完成后，训练才能修改这些源张量。

async_save 返回表示调用已交还控制权。采用异步 staging 时，staging_completion 表示本地状态拷贝完成；upload_completion 表示这次 checkpoint 保存完成。同步 staging 路径通常返回一个代表保存完成的 Future，使用时要按所用版本和配置区分。

训练应在下一次修改源状态之前等待 staging 完成。参数通常在 optimizer.step() 更新，但模型 buffer、RNG 或数据状态可能更早变化。退出训练、释放保存缓冲区或复用同一 stager 前，还要等待保存完成并检查异常。

默认分布式保存需要参与进程组中的各 rank 一起调用。每个 rank 提供自己的局部分片，DCP 协调保存计划和写入结果。协调者整理元数据，并不意味着所有模型张量都要先汇集到 rank 0。

已经收集成完整模型的单进程导出是另一种调用路径。它需要承担汇集和完整副本的开销，不能替代分片训练中的集体调用。

默认 Planner 会按键处理重复的普通值。相同键下，各 rank 内容不同的 RNG 或数据状态，不能直接当成同一份共享状态。

一种实现方式是先汇总各 rank 的小型状态，形成所有 rank 一致的映射，再交给 DCP。加载后每个 rank 取回对应项。需要改变并行度时，还要定义状态怎样映射到新 rank；简单按旧 rank 编号取值就不够了。

先按新的并行布局创建模型和优化器，再提供目标 state_dict。DCP 根据保存的全局张量元数据与目标分片，计算需要读取的区间。每个 rank 读取自己需要的部分，不要求它沿用原来的文件编号。

这依赖稳定的参数名、全局形状和受支持的分片表示。优化器状态也需要对应同一组参数。框架内部的扁平 buffer 还可能绑定原布局，例如 Megatron 的 dp_reshardable 与 fully_reshardable 支持范围就不同。

DCP 加载前，需要创建目标对象并准备目标张量存储。模型结构和目标分片来自当前程序。使用教程中的 get_state_dict / set_state_dict，可以统一模型和优化器的参数名映射。

Stateful 封装会接入 state_dict 与 load_state_dict。若直接传普通字典，加载后还要把状态应用回模型和优化器，让各对象的加载逻辑完成。

数据文件出现时，其他分片仍可能在写入。在默认集体保存路径中，DCP 等待各 rank 的写入结果，再调用 FileSystemWriter.finish() 写元数据，最后将临时元数据文件改名为 .metadata。

每次保存使用独立目录，等完成信号成功后再更新 latest。这样，写入中断时仍保留上一份完整 checkpoint。文件同步、重命名和持久性的保障取决于存储后端；关闭 sync_files 会改变故障后的保障。

只有数据顺序能由这些信息完整重建时才够用。流式数据、随机变换和序列拼接还可能有迭代器游标、worker RNG 与尚未用完的 token buffer。checkpoint 要对应训练已经消费的数据，而不是后台预取已经读到的位置。

StatefulDataLoader 提供状态接口，并汇总 worker 状态。它的 iterable dataset 示例要求恢复时保持相同 num_workers。改变 DP 大小或 worker 数，需要数据管线另行支持状态重分配。

按实际保存的张量逐项相加。设有 $P$ 个参数，同时保存 BF16 权重、FP32 主权重和两个 FP32 动量，则这几项共 $14P$ 字节。70B 参数约为 980 GB，采用十进制单位。

这是该状态组合的未压缩张量大小。优化器可能省略或重建部分副本；共享存储、分片副本和文件格式也会改变磁盘占用。更新边界上的 checkpoint 通常不需要梯度，不能直接把训练显存估算当作文件大小。

后台保存仍会使用 CPU、主机内存、PCIe、网络和存储。PyTorch 的吞吐图展示了 staging 结束后，线程间 GIL 竞争仍会影响训练；后台进程与保存计划缓存分别降低这类干扰和重复协调成本。

测量时同时看训练停顿、保存完成延迟和整段训练吞吐。若每隔 $\Delta$ 秒保存 $S$ 字节，存储至少要承受平均 $S/\Delta$ 的写入负载。长期写不完时，增加后台请求只会积压缓冲区。

在固定配置下，从同一步分成连续训练和保存后重启两条路径。先核对模型、优化器、scheduler 和数据位置，再比较下一批样本、loss 与更新后的参数。这个对照能发现漏存状态和加载顺序错误。

验证故障恢复时，再中断一次正在写入的保存，检查程序是否选回上一份完整 checkpoint。换卡数时重做数据进度与分片检查。逐位一致还取决于算子、通信归约顺序和软件版本。

恢复取决于故障后还有哪些副本可读。训练进程退出、整台机器丢失和共享存储不可用，对应不同的恢复路径。保存在原进程内的 CPU 缓冲区会随进程消失；本地 SSD 的数据也需要原节点可访问。

Gemini 把内存副本放到其他机器，并保留持久存储这一层。副本位置决定哪些节点故障组合能恢复，复制流量又会与训练争用网络。

先记录 checkpoint 大小、staging 停顿、写入尾延迟、主机内存和重启耗时，再结合中断频率与可接受的重算量选择。后续用实际训练结果补充。