AI 原生开放教材

外观

AI 工程从零到一（中文）源码研读·中文原文 ↗原件 ↓仓库 ↗

Checkpoint 保存与恢复 ​

训练中断会杀死整个 run；checkpoint 让它接着跑。把模型、优化器、调度器、loss 历史、step 计数器和 RNG 状态原子地存下来，这样任意时刻被 kill 都能在磁盘上留下一个合法文件。

类型： Build 语言： Python 前置要求： 第19阶段第42-45课 预计时间： ~90 分钟

学习目标 ​

  • 把完整的训练状态打包到一个 payload 中，能在全新进程里加载恢复。
  • 用"先写临时文件再 rename"的方式实现原子保存，crash 不会留下写到一半的文件。
  • 恢复 Python、NumPy 和 PyTorch 的 RNG 状态，使 resume 后的 loss 与未中断的 baseline 一致。
  • 为单文件装不下的大模型构建分片 checkpoint 布局，带 hash 校验的 shard 加上 JSON 索引。

问题背景 ​

你设了一个 18 小时的训练任务。挂钟上限是 4 小时。集群在第 11 小时重启了——因为你上面的人批准了一次内核升级。没有 checkpoint 就从头来过。没有 resume，就连优化器花了前 11 小时才学到的状态也丢了：即使模型权重幸存，AdamW 的矩估计也没了，下一步会朝着训练轨迹早已走过的方向猛冲。

正确的产物是一个包含续跑所需一切的文件：模型参数、优化器状态、调度器状态、画图用的 loss 历史、当前 step/epoch/batch-in-epoch 计数器，以及每种随机源的 RNG 状态。没有 RNG 状态，resume 后的 loss 曲线就是另一条曲线——同一个模型、同一份数据，不同的 shuffle、不同的 dropout mask、仪表盘上不同的数字。

原子保存是合约的另一半。直接写目标文件名意味着写到一半 crash 会留下一个损坏文件；resume 读到的是垃圾。先写同目录下的临时文件再 rename，crash 时之前的好文件依然完好。rename 在 POSIX 文件系统上是原子操作。

核心概念 ​

mermaid
flowchart TD
  ckpt[checkpoint 载荷] --> m[模型 state_dict]
  ckpt --> o[优化器 state_dict]
  ckpt --> s[调度器 state_dict]
  ckpt --> tr[训练状态: step, epoch, batch_in_epoch, losses]
  ckpt --> rng[RNG 状态: python, numpy, torch_cpu, torch_cuda]
  ckpt --> meta[wall_saved_at, schema]
  ckpt --> write[原子写入: 临时文件然后 os.replace]

五类状态 ​

类别为什么重要
Model权重和 buffer；模型"是什么"
OptimizerMomentum 和 adaptive moments；没有这些下一步就是另一个优化问题
Scheduler学习率在曲线上的位置；cosine schedule 尤其在意
Train countersStep、epoch、batch-in-epoch，加上画仪表盘的 loss 历史
RNG statedropout、数据 shuffle 和模型内部采样的确定性保证

原子保存 ​

mermaid
flowchart LR
  payload[载荷] --> tmpf[写入 .ckpt.pt.XXXX.tmp]
  tmpf --> rename[os.replace 到 ckpt.pt]
  rename --> done[ckpt.pt 有效]
  crash1[rename 前崩溃] --> orig[ckpt.pt 不变]
  crash2[rename 后崩溃] --> done

两条规则。第一，临时文件和目标放在同一个目录下，保证 rename 在同一个文件系统内；跨设备 rename 不是原子的。第二，每次临时文件名唯一，避免两个 writer 互踩。

分片 checkpoint ​

模型变大后，单文件 payload 加载太慢、检查太难，网络共享传输中断更痛苦。解决办法是把参数状态拆成 shard，再写一个小索引把它们串起来。

mermaid
flowchart LR
  state[state_dict] --> split[按轮询方式拆分 key 到 N 个分片]
  split --> s0[model.shard-000.pt]
  split --> s1[model.shard-001.pt]
  split --> sN[model.shard-NNN.pt]
  s0 --> idx[index.json]
  s1 --> idx
  sN --> idx
  meta[meta.pt: 优化器 + 调度器 + 训练状态 + RNG] --> idx

索引记录 shard 数量、每个 shard 的 sha256 以及 meta 文件的 sha256。加载器在任何 hash 不匹配时直接报错。shard 可以落在不同物理磁盘上；meta 很小，先读。

Resume 从 epoch 中间继续 ​

如果 resume 只能跳到下一个 epoch 开头，浪费几分钟到一整天不等。解决办法是 (epoch, batch_in_epoch) 加上 RNG 状态。加载后，训练循环把随机数生成器快进到当前 epoch 中已消费的 batch 之后，从 batch_in_epoch 继续。课程代码就是这么做的；断言是 resume 后的 loss 轨迹与未中断 baseline 的差异在 1e-4 以内。

figure
cc-atomic-checkpoint

动手构建 ​

code/main.py 提供四个原语和一个 demo 驱动。

第一步：捕获和恢复 RNG 状态 ​

capture_rng_state 返回一个 dict，包含 Python 的 random.getstate、NumPy 的 np.random.get_state，以及 PyTorch CPU 和 CUDA 的 RNG 字节。restore_rng_state 做反向操作。CPU tensor 是一个 uint8 字节 buffer，PyTorch 的 RNG 知道怎么消费它。

第二步：原子保存 ​

atomic_save 把 payload 写到目标目录的临时文件，再用 os.replace 换到最终文件名。atomic_write_json 对分片索引做同样的事。

第三步：完整 checkpoint 往返 ​

save_checkpoint 把 model、optimizer、scheduler、train state 和 RNG 打包成一个 dict。load_checkpoint 做反向操作并返回 TrainState。schema 字段是升级钩子：未来格式变更只需 bump 版本字符串，loader 按版本分发。

第四步：分片变体 ​

save_sharded_checkpoint 用 round-robin 把参数 key 分配到 N 个 shard，对每个 shard 做原子保存，写一个包含 optimizer、scheduler 和 train state 的 meta 文件，再写一个带 shard sha256 的 JSON 索引。load_sharded_checkpoint 在合并前校验每个 shard。

第五步：resume demo ​

run_resume_demo 训练一个小模型跑 total_steps 步，在 interrupt_at 处保存 checkpoint，然后继续。第二个进程恢复 checkpoint 并跑完剩余步骤。函数返回中断点之后两条 loss 轨迹的 max absolute difference。RNG 恢复后，差值为零或浮点噪音。

运行：

bash
python3 code/main.py

单文件和分片 demo 都断言 max-diff < 1e-4。摘要写入 outputs/resume-demo.json。

实际应用 ​

生产训练栈把 checkpointing 作为 trainer 的一部分。形状一样：model + optimizer + scheduler + counters + RNG，原子写入，以 step 命名方便找最新的。分片布局让大模型可以并行读取；index.json 是让这一切运转的关键。

三个需要坚持的模式：

  • Schema 是 payload 里的一个字符串。 迁移逻辑按它分发。没有它你就没法演进格式而不破坏旧 run。
  • 每个 shard 都做 sha256。 静默截断的下载是最坏的 bug 类型；loader 要么早报错，要么就是太晚才发现。
  • 保持 checkpoint 频率诚实。 每 N 步和每隔一定挂钟分钟数保存一次，取较短者。否则那个 crash 的超长 step 会浪费一整个窗口的工作量。

交付产物 ​

outputs/skill-checkpoint-save-resume.md 是任何新训练脚本的 recipe：payload 结构、原子写入、RNG 捕获、分片索引。把 skill 丢进 repo，在周期性保存点接上 save_checkpoint，在启动时接上 load_checkpoint，run 就能扛住 kill。

练习 ​

  1. 把 round-robin 分片替换为按参数组分片（.weight 结尾的 vs .bias 结尾的）。什么时候哪种布局更好？
  2. 扩展保存逻辑，保留最近 K 个 checkpoint 并清理更旧的。磁盘小的时候 K 取多少合适？
  3. 加一个 --ckpt-every-seconds 标志，按挂钟间隔触发保存，而非仅按 step 数。
  4. 加一条 checksum 校验路径，在启动时扫描目录下的所有 checkpoint，报告哪些已损坏。
  5. 实现一个 migrate_v1_to_v2 函数，给 payload 加一个新字段并 bump schema 字符串。让 load 兼容两个版本。

关键术语 ​

术语日常说法实际含义
Atomic save"写完祈祷"先写到同目录的临时文件，再用 os.replace 替换到目标文件名
State dict"权重"模型参数和 buffer，以参数名为 key
Sharded checkpoint"大模型文件"多个文件，每个 shard 一个，外加一个 meta 文件和一个带 sha256 的 JSON 索引
RNG state"随机种子"捕获的 Python random、NumPy、torch CPU、torch CUDA 的状态；不只是 seed
Mid-epoch resume"重启"快进 RNG，从同一个 epoch 的下一个 batch 继续

延伸阅读 ​

  • POSIX rename 语义，os.replace 原子性声明的基础。
  • PyTorch torch.save 和 torch.load 文档，包括跨设备恢复的 map_location。
  • 第19阶段第46课覆盖了本课的 checkpoint payload 需要跨越的梯度累积。
  • 第19阶段第48课覆盖了分布式 wrapper，本方案需要兼容其 state dict 格式。
  • Linux 内核 fsync 文档，原子 rename 背后的持久性保证。