Checkpoint 存储与快速恢复：从本地盘到并行文件系统的两级设计
说明：本篇为综述/无第一手实验数据，采用公开实践与架构解析路线。
实验目标
大模型训练的天数尺度里，一次意外中断（节点宕机、电力抖动、OOM 误杀）可能丢掉 10 小时以上的训练进度。Checkpoint 是唯一的"后悔药"，但它的存储设计常被忽略：存太快拖慢训练（IO 突发挤占带宽），存太慢丢进度（恢复要重训）；都存本地盘会丢（节点没了盘也没了），都存并行文件系统又太贵太慢。本篇目标：给出一个可落地的两级 Checkpoint 存储架构（本地 NVMe 热存 + 并行文件系统冷存），并给出保存频率、双写时机、恢复路径的核算方法。
代码逐段拆解
第一段：训练脚本里的 Checkpoint 频率控制
python
# 训练循环片段：按 token 而非按步数保存
import time
ckpt_interval_tokens = 50_000_000 # 每 5 千万 token 落一次盘
tokens_since_ckpt = 0
last_ckpt_time = time.time()
for step, batch in enumerate(train_loader):
loss = train_step(batch)
tokens_since_ckpt += batch.num_tokens()
# 策略 1：跨步数保存不稳定（步长随数据变化），按 token 更均匀
if tokens_since_ckpt >= ckpt_interval_tokens:
save_checkpoint(step, local_path='/mnt/local-nvme/ckpt-latest',
async_copy_to='/mnt/pfs/ckpt/step-%d' % step)
tokens_since_ckpt = 0
# 记录保存耗时，作为 IO 预算监控
elapsed = time.time() - last_ckpt_time
print(f"[ckpt] step={step} save+copy={elapsed:.1f}s")
last_ckpt_time = time.time()要点：
- 按 token 计数比按 step 稳定：数据并行下每个 step 的 token 数因填充、多卡 batch 对齐而波动，按 token 保存保证"每存一次，损失的上限进度"是恒定值（这里是 5 千万 token ≈ 一定训练时长）。
- 先写本地 NVMe、再异步拷到 PFS：本地 NVMe 顺序写可达 GB/s 级，几十秒完成；PFS 拷备是异步的，不阻塞训练主线程（下一段看怎么实现异步）。
- 打印保存耗时：这是 IO 预算的仪表盘——如果
save+copy超过几分钟，说明 Checkpoint 已经拖慢训练，需要调频率或换存储层。
第二段：两级存储与异步拷贝
bash
# 节点上的一级 Checkpoint 盘（本地 NVMe，训练前挂载）
mkdir -p /mnt/local-nvme/ckpt-latest && chmod 1777 /mnt/local-nvme
# 二级：并行文件系统目录（多节点可见，跨节点恢复用）
mkdir -p /mnt/pfs/ckpt && chmod 1777 /mnt/pfs/ckptpython
# 异步双写：本地先落盘，PFS 拷贝异步进行（简化示意）
import threading, shutil
def async_copy(src, dst):
threading.Thread(target=lambda: (shutil.copytree(src, dst, dirs_exist_ok=True)), daemon=True).start()
save_checkpoint(step, '/mnt/local-nvme/ckpt-latest')
async_copy('/mnt/local-nvme/ckpt-latest', f'/mnt/pfs/ckpt/step-{step}')设计逻辑：
- 一级（本地 NVMe）：保证"保存动作不阻塞训练"——写本地盘几十秒内完成，训练继续。风险是节点故障时盘上数据丢失，所以它只是缓冲。
- 二级（PFS）：保证"跨节点可恢复"——PFS 是全集群共享的，任何节点都能读到最新 Checkpoint。PFS 拷备虽慢（几百 MB/s~GB/s），但它是异步的。
- 两级之间靠异步任务衔接：训练进程不等待 PFS 完成，只保证"下一轮保存前上一轮拷贝已完成"（实践中用队列 + 水位确认）。
原理讲解
为什么 Checkpoint 是 IO 突发型负载
训练的平均 IO 很小（读数据流、写日志），但保存 Checkpoint 的瞬间会产生一个大而集中的写突发：70B 模型 BF16 的权重 + 优化器态（Adam 要存两份动量）+ 数据并行下的梯度统计，往往 300GB-1TB 级别。这带来两个问题：
- 突发挤占共享存储带宽：如果所有节点同时保存（step 对齐），PFS 瞬时被打满，其他任务的读延迟飙升——这就是"训练一保存，全集群卡顿"的原因。
- 写入时间不可忽略：按 300GB 算，PFS 单客户端写 500MB/s 要 10 分钟——如果保存是同步的，训练每 30 分钟就停 10 分钟，MFU 直接掉 25%。
两级架构为什么是对的
| 层级 | 介质 | 延迟/带宽 | 可靠性 | 角色 |
|---|---|---|---|---|
| 一级 | 本地 NVMe | 数 GB/s | 节点故障即丢 | 热存：保存不阻塞训练 |
| 二级 | 并行文件系统 | 百 MB/s-GB/s | 多副本/纠删 | 冷存：跨节点可恢复 |
核心权衡：训练连续性 vs 恢复完整性。一级盘保"训练不停"，二级盘保"能恢复"。实际生产中，二级盘甚至可以是"每 N 个一级 Checkpoint 才拷一个"——丢掉中间一两个进度，换训练全程不被 IO 拖累。
恢复路径优化
恢复的耗时 = 读 Checkpoint + 重新初始化分布式状态。优化三板斧：
- 分级恢复：先恢复权重（几秒内可推理/评估），优化器态异步恢复（训练再等），把"能恢复"和"全恢复"解耦。
- 尽量恢复最近的一级盘副本：同节点还在时直接从本地 NVMe 读，秒级恢复；节点换了才走 PFS。
- 异步保存 + 固定频率：固定 token 间隔、异步双写，恢复损失上限恒定，不会出现"保存太勤拖慢训练、太疏丢太多"的两难。
实验结果
以下为 70B 训练场景的估算（估算口径，供设计参考）：
| 参数 | 估算值 | 说明 |
|---|---|---|
| Checkpoint 体积 | 300-400GB | 权重 140GB + 优化器态 160GB+ |
| 本地 NVMe 写入 | ~2-4 min | 数 GB/s 顺序写 |
| PFS 异步拷贝 | 10-15 min | 500MB/s 级 |
| 恢复（本地盘） | < 5 min | 权重先行，优化器态后到 |
| 恢复（PFS） | 15-30 min | 全量读 + 分布式状态重建 |
工程启示
- Checkpoint 频率按 token 定：先定"可接受的最大损失进度"，再反推保存间隔，而不是拍脑袋每 500 步存一次。
- 本地盘是必需品不是奢侈品：两级架构里一级盘的成本换来训练全程不卡顿，性价比远高于"同步写 PFS"。
- 把 IO 预算变成监控项：保存耗时、PFS 队列水位要进 Grafana（上篇的 DCGM 看板旁边加一块存储面板），突发异常早发现。
- 与调度器联动：Checkpoint 完成时间点要能被调度器感知——SLURM/K8s 在 Checkpoint 后才允许缩容或维护节点，避免"存了白存"。
下一篇预告：本篇解决了"训练中断后怎么快速回来"，推理侧的显存压力则是另一个战场——"KV Cache 优化与显存管理：长上下文推理的算账方法"。