检查点与容错¶
大模型训练的 checkpoint 不是一个权重文件,而是恢复计算轨迹所需的分布式状态快照。系统需要同时回答:
- 哪些状态决定下一步计算?
- 多个 shard 怎样形成一个不可分割的提交?
- 保存时怎样避免阻塞训练或复制正在变化的 buffer?
- world size 与并行拓扑改变后怎样恢复同一个逻辑 tensor?
只恢复模型权重属于 warm start,不等价于严格 resume。
恢复状态¶
严格 resume 通常需要:
- 模型参数与 buffer;
- optimizer state、参数分组和更新步数;
- learning-rate scheduler;
- gradient scaler、amax history 与低精度 scale;
- global step、已见 token 数、样本与 packing 游标;
- CPU、GPU、dropout 与随机舍入 RNG;
- data loader、shuffle 与动态采样状态;
- DP、TP、PP、CP、EP 的 mesh 与 tensor placement;
- 模型、tokenizer、代码、容器、依赖和数据版本。
状态可按恢复强度分层:
| 模式 | 状态 | 可承诺的语义 |
|---|---|---|
| 权重加载 | 模型参数 | 新优化轨迹的 warm start |
| 训练恢复 | 权重、optimizer、scheduler、scaler | 优化状态连续 |
| 严格恢复 | 再加数据游标、RNG、拓扑和版本 | 在误差契约内延续同一轨迹 |
分布式 reduction 和低精度 kernel 可能使恢复后不具备 bitwise 一致性;应提前定义允许的 loss、梯度与收敛误差,而不是事后降低标准。
逻辑 tensor 与物理 shard¶
checkpoint 应优先描述逻辑 tensor:
每个 shard 记录 global shape、global offset、local shape、dtype 和 placement。旧 topology 恢复到新 topology 时,流程是
若格式只用文件名中的 local rank 表示身份,就难以在 TP / PP / EP size 改变后安全 reshard。global parameter ID、expert ID 和 pipeline layer range 必须独立于当前 rank 编号。
PyTorch Distributed Checkpoint 提供分布式 state dict、planner 与跨 rank 保存 / 读取接口。具体格式仍需由调用方补齐数据游标、版本、外部状态和提交协议。
Manifest¶
一个 snapshot manifest 至少包含:
- 不可变 snapshot ID;
- 创建时间、global step、token 与 data cursor;
- 模型、tokenizer、代码、数据和 schema version;
- 原 topology 与 process mesh;
- 每个逻辑 tensor 的 shape、dtype、placement;
- 每个 shard 的对象名、offset、长度与 checksum;
- optimizer、scheduler、scaler 和 RNG 的存在性;
- parent snapshot 或增量依赖;
- writer 数量和格式版本;
- committed 状态。
manifest 是一致性边界,不只是文件清单。恢复端必须先验证 schema,再验证全部必需 shard 和 checksum,最后才分配或加载模型状态。
Manifest 完整性检查¶
下面用 byte payload 简化表示逻辑 tensor shard。每项输入包含全局 offset、逻辑 length、不可变 payload 与记录的 SHA-256;只有 range 无缺口、无重叠且 checksum 全部匹配时才返回 True。
import hashlib
def make_shard(offset, payload):
return {
"offset": offset,
"length": len(payload),
"payload": payload,
"sha256": hashlib.sha256(payload).hexdigest(),
}
def manifest_complete(shards, global_length):
cursor = 0
for shard in sorted(shards, key=lambda item: item["offset"]):
if shard["offset"] != cursor or shard["length"] != len(shard["payload"]):
return False
if hashlib.sha256(shard["payload"]).hexdigest() != shard["sha256"]:
return False
cursor += shard["length"]
return cursor == global_length
a, b = make_shard(0, b"ab"), make_shard(2, b"cd")
assert manifest_complete([b, a], 4)
assert not manifest_complete([a], 4)
assert not manifest_complete([a, {**b, "sha256": "0" * 64}], 4)
该不变量直接阻止缺 shard、重复 range 与静默损坏进入恢复路径。真实 tensor 的逻辑元素范围和对象字节范围应分别记录;持久化、跨 rank 共识与原子 marker 仍由存储协议保证,只有这个检查成功后 coordinator 才能发布 committed 状态。
原子提交协议¶
对象存储或分布式文件系统通常不能让一组文件天然原子。可采用两阶段协议:
- 为本次保存生成唯一 snapshot ID。
- 各 rank 写入临时命名空间,只写不可变 shard。
- writer 完成后返回 size、checksum 与逻辑 range。
- coordinator 收齐并验证所有必需 shard。
- 写入完整 immutable manifest。
- 最后原子创建很小的 committed marker,或原子更新 latest 指针。
恢复端只枚举 committed snapshot:
marker 必须最后写入。仅依赖目录存在、文件数量或修改时间,会把部分写入误认为可恢复状态。latest 只是可重建索引,不应成为唯一真相。
失败保存产生的临时对象可以延迟回收,但回收器不得删除仍被某个 committed 增量快照引用的 parent 或 shard。
同步与异步保存¶
同步保存使所有 rank 在一致点停下,语义简单但暂停时间长。异步保存通常分两阶段:
- 在一致 step 边界把本地状态复制为不可变 staging snapshot;
- 后台将 staging 数据写入持久存储。
设本 rank checkpoint 大小为 \(M_{\mathrm{local}}\),异步路径至少需要约
的额外 CPU 或其他稳定存储空间,除非采用更细粒度流水。保存中的 buffer 不能被下一训练 step 修改或 allocator 复用。
PyTorch 异步 checkpoint recipe 明确把 state 搬到 CPU staging,再由后台 future 完成写入,并建议限制同时只有一个异步保存。若无界累积多个 future,CPU 内存、PCIe、网络和存储队列都可能失控。
异步并不意味着零开销:
- device-to-host copy 竞争 PCIe / 内存带宽;
- state snapshot 会短暂增加显存或 CPU 内存;
- 后台压缩、checksum 和写入消耗 CPU;
- 共享网络上的 checkpoint 会干扰 collective;
- 下一次保存可能等待上一次完成。
需要在 trace 中分别记录 snapshot pause 与 background drain。
一致保存点¶
checkpoint 边界应明确位于:
若在 gradient accumulation 中间保存,还需保存已累积 gradient、microbatch index 和相应 RNG。否则恢复后无法重建同一更新。
各 rank 的 global step 相同并不足够。还要检查:
- optimizer 是否全部完成;
- scheduler 是否已经推进;
- data cursor 是当前 batch 还是下一 batch;
- scaler overflow 分支是否完成;
- pipeline 是否已 drain;
- 异步 collective 与 kernel 是否完成使用相关 buffer。
一致点可通过 barrier 或协议实现,但 barrier 本身不能修复不同 rank 已处于不同逻辑 step 的错误。
跨拓扑恢复¶
从旧网格
恢复到新网格时,需要分别处理:
- replicated state 是否只读一个副本;
- TP shard 怎样按全局维度重切;
- PP layer range 怎样重新分配;
- optimizer state 是否与参数使用相同 placement;
- expert global ID 怎样映射到新 EP rank;
- data-parallel batch 与数据游标怎样保持 token 语义;
- RNG 是按 global sample 还是按 rank 派生。
跨拓扑恢复不应要求某个 rank 聚合完整模型,否则大模型会在 CPU 或单卡产生不可接受的峰值。更好的 planner 直接计算旧 shard 与新 shard range 的交集并流式传输。ByteCheckpoint 研究了大规模训练中高效、弹性的分布式 checkpoint,可作为 2024 年的系统设计实例,而非通用格式标准。
保存间隔¶
设一次 checkpoint 的不可隐藏耗时为 \(C\),系统平均故障间隔为 \(M\),保存间隔为 \(\tau\)。忽略恢复时间和高阶项时,单位时间浪费近似为
第一项是保存开销,第二项是假设故障在区间内均匀发生时的平均 lost work。求极值得到 Young 一阶近似:
Daly 的高阶分析在 Poisson 单组件故障、长作业等假设下,对“两个 checkpoint 之间的有效计算时间”给出更精细的近似:
恢复时间会增加总 wall time,但在该高阶最优间隔近似中不直接出现。实际选择仍需加入异步保存的隐藏比例、集群规模相关 MTBF、存储限流和业务 RPO;故障相关、重尾分布、检测延迟不可忽略或 \(C\) 接近 \(M\) 时,应使用完整模型或实测优化,不能机械套用 Young / Daly。
增量与分层 checkpoint¶
完整快照简单但写入量大;增量快照只保存变化块,却引入依赖链。恢复成本可写为
其中 \(d\) 是增量链深度。链越长,任一 parent 丢失或损坏的影响越大。
分层策略可以组合:
- 高频本地 / 节点级快照;
- 低频远端持久快照;
- optimizer state 与权重不同周期;
- 训练恢复格式与部署权重导出格式分离。
只有远端持久层能覆盖整机、机架或站点故障。本地快照不应被统计为相同故障域下的完整 RPO 保证。
正确性契约¶
- snapshot ID 全局唯一且提交后不可变。
- 每个逻辑 tensor 的 ranges 不重叠、不缺失,除非 schema 明确允许复制。
- manifest 写入前,全部必需 shard 已持久化并校验。
- committed marker 只在 manifest 完整后出现。
- 恢复端永不读取未提交 snapshot。
- staging buffer 在后台完成前不可修改或释放。
- 同时进行的异步保存数量有界。
- data cursor、global token、scheduler step 与 optimizer step 关系明确。
- RNG、scaler、amax 与低精度 scale 属于恢复状态。
- 格式升级有显式 schema version 和转换器。
- strict resume、approximate resume 与 warm start 在接口和日志中区分。
ZeRO/FSDP shard 与 reshard 的基础语义见集合通信与状态分片。
失效模式¶
- 半写目录被发现:缺少最终 committed marker。
- 不同 rank 保存不同 step:只有局部 step,没有全局一致点。
- 异步 buffer 被覆盖:快照不是不可变副本。
- 只校验文件存在:静默截断或 bit rot 未被发现。
- rank 文件名即身份:拓扑改变后无法定位逻辑 range。
- 只恢复模型与 optimizer:数据、scheduler 或 RNG 漂移。
- 增量 parent 被清理:manifest 仍在但依赖链断裂。
- 保存导致训练降速:PCIe、CPU、网络或存储争用未进入关键路径模型。
- 重试写入同名对象:不同 attempt 互相覆盖,manifest 混合。
对于短实验或可从头重跑的小模型,复杂增量和跨拓扑格式可能得不偿失;一个有 checksum 的同步完整快照通常更可靠。
故障注入¶
至少覆盖:
- 某 rank 在 shard 写入前、写入中、manifest 前和 marker 前退出;
- shard 截断、checksum 错误、重复 range 和缺失 range;
- coordinator 失败并由新 coordinator 重试;
- 对象存储出现超时、重排可见性或重复响应;
- 异步保存尚未完成时触发下一次保存;
- 保存过程中训练进程 OOM 或被抢占;
- 从不同 DP / TP / PP / EP 网格恢复;
- schema、模型、tokenizer 或 optimizer 版本不兼容;
- data cursor 恢复后出现重样本或跳样本;
- 本地快照可用但远端持久层不可用。
每个故障点都应证明:未提交快照不会被读取,最近一次已提交快照仍可恢复,临时对象可安全回收。
恢复验证¶
- 在独立进程和干净环境中读取,而不是仅由 writer 自检。
- 核对 manifest、shard checksum、global shape、dtype 与 range 覆盖。
- 比较恢复前后的模型参数、optimizer step、scheduler 和 scaler。
- 对固定 batch 比较 loss、logits、梯度统计和下一次更新。
- 比较样本 ID、token cursor 与 RNG 派生序列。
- 在至少一个不同 topology 上完成 reshard 与下一 step。
- 记录 snapshot pause、后台 drain、存储吞吐、RTO 和 lost work。
- 长期演练 partial write、worker failure 和存储故障。
恢复排查可结合调试手册,manifest 原子提交与恢复校验的紧凑 reference 见手撕:分布式与容错。通过一次读取测试只证明一个样本可读;可靠 checkpoint 需要持续的故障注入、跨版本迁移和独立恢复证据。