背景:检查点存在,不等于训练能够恢复
单机微调时,torch.save(model.state_dict()) 往往已经足够完成模型导出;进入 FSDP、ZeRO、Tensor Parallel 或 Pipeline Parallel 后,训练状态被拆散到多个 Rank,检查点也从”一个文件”变成一组相互依赖的分片。
生产事故通常不是完全没有检查点,而是出现以下情况:
- 某些 Rank 已写完,另一些 Rank 在节点故障前只写了一半;
latest已指向新目录,但新检查点尚未完整提交;- 模型权重恢复成功,优化器动量、学习率或 GradScaler 没有恢复;
- 数据集从 Epoch 开头重新读取,导致样本重复;
- 原训练使用 64 张 GPU,恢复时只有 48 张,旧分片无法按新拓扑加载;
- 异步保存积压,占满主机内存,最终把训练进程一起拖垮。
因此,生产级断点恢复的目标不是”能找到一批文件”,而是确保一个检查点同时满足四个条件:状态完整、提交可判定、拓扑可兼容、恢复可验证。
PyTorch Distributed Checkpoint(DCP)支持多个 Rank 并行保存和加载,并能依据加载端的分片信息,在支持的场景中跨不同 World Size 重新分片。它生成多个文件,而不是传统的单一 torch.save 文件,加载前还需要先由目标模型分配状态存储并提供分片信息。
核心原理:把恢复点当成一笔分布式事务
1. 明确训练状态边界
一个可续训检查点至少应覆盖:
- 模型状态:基础模型参数、LoRA 或其他 Adapter 参数、可训练头部;
- 优化器状态:Adam 的一阶和二阶动量、参数组设置、ZeRO/FSDP 分片;
- 调度状态:当前学习率、Warmup 进度、Scheduler 内部计数;
- 数值状态:GradScaler、混合精度配置、梯度累积所在 Micro Step;
- 随机状态:Python、NumPy、PyTorch CPU、各 CUDA 设备和数据采样器 RNG;
- 进度状态:Global Step、Epoch、已消费样本数、数据游标;
- 配置指纹:模型结构、Tokenizer、Chat Template、数据集版本、训练参数和代码版本;
- 并行拓扑:World Size、DP/TP/PP 维度、ZeRO Stage、FSDP Sharding Strategy。
Hugging Face Accelerate 的检查点接口会保存模型、优化器、RNG 和 GradScaler,也允许注册学习率调度器等自定义 Stateful 对象;若在 Epoch 中途恢复,还需要恢复或跳过已经消费的 DataLoader Batch。这个细节经常被遗漏:权重相同并不代表下一步看到的样本相同。
2. 用分片格式保存,不在训练主路径强制聚合
大模型训练中,将所有分片先汇聚为一个完整 FP32 State Dict,可能带来明显的通信、CPU 内存和保存停顿。更合理的主检查点是 Sharded Checkpoint:每个 Rank 并行写自己的分片,同时保存足够的元数据,供加载端重新组织。
完整权重导出和故障恢复是两个不同目标:
| 目标 | 训练恢复检查点 | 发布制品 |
|---|---|---|
| 首要追求 | 写入快、状态全、可重分片 | 格式通用、文件稳定、便于推理加载 |
| 写入方式 | 各 Rank 并行写分片 | 聚合为完整权重 |
| 生命周期 | 训练期间频繁产生,恢复后即可清理 | 长期保留,供下游消费 |
不要为了方便发布,强迫每次训练检查点都执行全量权重聚合。DeepSpeed 提供从 ZeRO 检查点提取 FP32 权重的工具,但官方文档也提醒,这一步可能受到 CPU 内存约束,而且将聚合权重加载回模型后通常需要重新初始化 DeepSpeed Engine。
3. 用提交标记区分”文件存在”和”恢复点有效”
建议把一次保存拆成两个阶段:
阶段 A:写入临时前缀
checkpoints/
step-00012000.inprogress/
rank-00000.distcp
rank-00001.distcp
...
metadata.part.json
每个 Rank 完成写入后,记录文件大小、校验和和状态摘要。协调 Rank 只有在确认全部预期分片都存在、校验通过后,才进入提交阶段。
阶段 B:发布提交证据
对象存储通常不应依赖目录 Rename 来模拟原子操作。更稳妥的工程方案是使用不可变前缀,再写入一个体积很小的 COMMITTED 或 Manifest,最后更新 latest.json 指针。恢复端只承认具备有效提交证据的版本。
{
"checkpoint_id": "step-00012000",
"status": "COMMITTED",
"global_step": 12000,
"saved_world_size": 64,
"parallelism": { "dp": 8, "tp": 8, "pp": 1 },
"training_state": {
"epoch": 2,
"micro_step": 3,
"consumed_samples": 1572864
},
"fingerprints": {
"model_config": "sha256:...",
"tokenizer": "sha256:...",
"dataset_manifest": "sha256:...",
"training_config": "sha256:..."
},
"shards": [
{ "rank": 0, "path": "rank-00000.distcp", "size": 4831838208, "sha256": "..." }
]
}
latest.json 只是索引,不是可信事实。即使指针损坏,系统也应能扫描最近的已提交 Manifest 回退,而不是选择目录名最大的版本。
工程落地
保存流程:先冻结一致状态,再异步落盘
PyTorch DCP 的 async_save 会先把训练状态复制到内部 CPU Buffer,再异步写入存储,从而缩短 GPU 训练主路径中的阻塞时间。但这不是无成本操作:官方文档指出,CPU 内存压力大致与每 Rank 检查点大小和 Rank 数量相关;使用 Pinned Memory 还能加速暂存复制,却会让峰值页锁定内存长期驻留。
一个简化的保存控制器:
from concurrent.futures import Future
import torch.distributed as dist
import torch.distributed.checkpoint as dcp
checkpoint_future: Future | None = None
def save_checkpoint(state: dict, checkpoint_id: str) -> Future:
global checkpoint_future
# 避免同时堆积多个异步保存任务
if checkpoint_future is not None:
checkpoint_future.result()
checkpoint_future = dcp.async_save(
state_dict=state,
checkpoint_id=f"{checkpoint_id}.inprogress",
)
return checkpoint_future
def commit_checkpoint(checkpoint_id: str) -> None:
assert checkpoint_future is not None
checkpoint_future.result()
dist.barrier()
if dist.get_rank() == 0:
verify_all_shards(checkpoint_id)
write_manifest(checkpoint_id)
publish_latest_pointer(checkpoint_id)
dist.barrier()
这里的 verify_all_shards、write_manifest 和 publish_latest_pointer 是业务侧提交协议,不是 DCP 自动提供的完整事务系统。关键约束是:未完成的 Future 不能被当成成功检查点,未发布 Manifest 的目录不能用于恢复。
恢复流程:兼容性判断必须早于加载
恢复任务启动后,建议按以下顺序执行:
- 查找最近的已提交 Manifest;
- 校验 Manifest、分片数量、大小和校验和;
- 比较模型结构、Tokenizer、数据集和训练配置指纹;
- 根据目标 GPU 数量创建新的 Device Mesh 或并行拓扑;
- 构造并预分配目标模型与优化器 State Dict;
- 加载并重新分片模型和优化器状态;
- 恢复 Scheduler、GradScaler、RNG、Global Step 和数据游标;
- 执行恢复后验证,再允许任务进入正式训练。
state = build_stateful_training_state(
model=model,
optimizer=optimizer,
scheduler=scheduler,
scaler=scaler,
)
dcp.load(
state_dict=state,
checkpoint_id=selected_checkpoint,
)
restore_rng_state(manifest["rng_state"])
restore_data_cursor(manifest["training_state"])
validate_resume_state(model, optimizer, manifest)
DCP 的加载是原地写入,依赖目标模型已经分配好的 State Dict 和分片信息。因此,“先读取检查点,再决定新拓扑”的顺序是错误的;应先决定恢复资源和分片布局,再执行加载。
World Size 变化:分三类处理
第一类:仅数据并行度变化。 在模型结构和受支持的分片方式不变时,DCP 可以根据目标端分片信息重新分片。这类变更最适合作为弹性恢复的第一阶段目标。
第二类:TP、PP、DP 或 ZeRO 组合发生变化。 这通常需要通用检查点格式或框架特定转换。DeepSpeed 文档将 Universal Checkpoints 标记为仍在开发中的能力,用于支持 3D Parallelism 下 GPU 数量变化;Universal Checkpointing 论文的核心思路,则是将保存时的分布式表示与加载时的统一参数表示、映射元数据解耦。
生产上不能因为功能名叫”Universal”就假设任意组合都能恢复。必须建立明确的兼容矩阵:
resume_compatibility:
- from: {dp: 8, tp: 8, pp: 1, zero: 1}
to: {dp: 6, tp: 8, pp: 1, zero: 1}
status: verified
- from: {dp: 8, tp: 8, pp: 1, zero: 1}
to: {dp: 8, tp: 4, pp: 2, zero: 1}
status: blocked_until_drill_passes
第三类:模型结构或状态语义变化。 例如新增可训练层、改变 Optimizer、调整词表、切换 ZeRO Stage,不能简单视为 World Size 变化。应进入显式迁移流程;没有迁移器时必须 Fail Closed,而不是跳过缺失 Key 后继续训练。
恢复后如何证明”真的续上了”
加载成功只证明反序列化没有报错。生产验收还需要短程对照:
- Global Step、Learning Rate、GradScaler 与保存点一致;
- 优化器参数组和动量张量数量一致;
- 数据游标没有回到 Epoch 开头;
- 恢复后的第一个 Batch ID 符合预期;
- 下一步 Loss、Gradient Norm 和参数更新幅度没有异常跳变;
- 连续运行若干步后,与不间断对照任务的指标差异处于预先定义的容差范围内。
分布式 GPU 训练未必能在所有 Kernel、拓扑和并行配置下实现逐 Bit 一致,因此恢复验收不宜只比较最终权重哈希。更实际的目标是验证 状态连续性、数据不重不漏和训练轨迹没有不可解释的突变。
故障演练:不要等第一次事故才验证恢复
建议在预发布环境周期性执行以下演练:
- 在部分 Rank 写入期间直接终止进程,验证
.inprogress不会被选中; - 删除一个 Shard,验证恢复端拒绝加载并回退到前一提交点;
- 篡改文件内容,验证校验和门禁生效;
- 让
latest.json指向不存在的版本,验证 Manifest 扫描回退; - 使用不同 World Size 恢复,验证重分片兼容矩阵;
- 在梯度累积中间保存和恢复,验证 Micro Step 不漂移;
- 模拟对象存储超时,验证异步保存 Future 不会无限积压;
- 恢复后运行固定样本短程对照,验证 Loss 和数据游标。
没有故障注入的检查点系统,实际只是一个未经验证的备份脚本。
监控指标与告警
至少记录以下指标:
| 指标 | 含义 |
|---|---|
checkpoint_save_duration_seconds | 单次保存耗时 |
checkpoint_async_staging_seconds | 异步暂存阶段耗时 |
checkpoint_bytes_total | 检查点总大小 |
checkpoint_commit_lag_seconds | 从保存完成到提交完成的延迟 |
checkpoint_inprogress_count | 当前未完成检查点数量 |
checkpoint_last_committed_step | 最近已提交的 Global Step |
checkpoint_restore_duration_seconds | 恢复耗时 |
checkpoint_restore_failures_total{reason} | 按原因分类的恢复失败次数 |
checkpoint_checksum_failures_total | 校验和失败次数 |
checkpoint_reprocessed_samples_total | 恢复后重复处理的样本数 |
resume_loss_delta | 恢复前后 Loss 差异 |
resume_world_size_change_total | World Size 变更恢复次数 |
告警不应只看”最近是否生成目录”,而应看最近一次已提交且演练验证通过的恢复点距离当前训练进度有多远。这才是训练任务真正的恢复点目标(RPO)。
适用场景
这套方案特别适合:
- 跨多节点、运行时间较长的 SFT、DPO、持续预训练任务;
- 使用抢占式或弹性 GPU 资源的训练集群;
- 采用 FSDP、ZeRO 或混合并行,无法用单文件完整保存的模型;
- 训练期间可能调整 GPU 数量,但仍要求保留优化器轨迹的任务;
- 对样本不重不漏、训练审计和成本控制有明确要求的企业环境。
对于数小时内即可完成的小型 LoRA 任务,也不必照搬复杂系统,但仍建议保存 Optimizer、Scheduler、RNG、数据游标和配置指纹,而不是只保留 Adapter 权重。
常见误区
误区一:模型能加载,就代表可以续训。 只恢复模型参数会丢失优化器动量、学习率进度和数值缩放状态,结果更接近”从旧权重重新开始训练”,而不是严格续训。
误区二:目录名最大的就是最新检查点。 写入一半的目录同样可能拥有最大的 Step。必须以 Manifest 和提交标记判断有效性。
误区三:所有 Rank 写同一个文件更简单。 这容易造成覆盖、锁竞争和损坏。分布式训练应使用框架支持的分片保存协议。
误区四:异步保存可以无限并发。 每个任务都需要 CPU 暂存和 I/O。PyTorch 官方建议多数用户一次只保留一个异步请求,避免额外内存压力。
误区五:World Size 可变等于所有并行维度可变。 DP 数量变化、TP/PP 变化和模型结构变化是不同级别的迁移,必须分别验证。
误区六:恢复时重新 Shuffle 数据没关系。 样本重复或跳过会改变训练分布,尤其会影响短数据集、偏好数据和按阶段混合的数据配方。
上线检查清单
- 已定义模型、优化器、Scheduler、Scaler、RNG 和数据游标的状态边界;
- 每个检查点都有不可变 ID、Manifest、校验和与提交标记;
- 未提交目录不会被 latest 或自动恢复逻辑选中;
- 对象存储发布不依赖目录 Rename;
- 异步保存同时最多运行一个任务,且有 CPU/Pinned Memory 预算;
- 建立原 World Size 与目标 World Size 的兼容矩阵;
- 恢复前校验模型、Tokenizer、数据集、代码和训练配置指纹;
- 恢复后验证 LR、Scaler、Global Step、数据游标、Loss 和 Gradient Norm;
- 已演练分片缺失、分片损坏、Rank 中断、指针损坏和存储超时;
- 监控的是最近已提交恢复点,而不是最近创建目录;
- 保留至少一个前序健康检查点,避免最新版本损坏后无处回退;
- 发布制品导出与训练恢复检查点采用不同生命周期。
参考资料
- PyTorch — Getting Started with Distributed Checkpoint
- PyTorch — Asynchronous Saving with Distributed Checkpoint
- Hugging Face Accelerate — Checkpointing
- DeepSpeed — Model Checkpointing
- Universal Checkpointing: A Flexible and Efficient Distributed Checkpointing System for Large-Scale DNN Training with Reconfigurable Parallelism