企业官网建设流程全解析

54 天崩 466 次，训练进度怎么才不白跑？

摘要

大规模训练一定会出故障，区别只在于崩一次要损失多少进度。PyTorch 自带的分布式 Checkpoint（torch.distributed.checkpoint，简称 DCP）让每张卡并行写自己那一份状态，恢复时还能换一种卡数和并行方式重新切分；再配合异步保存和 torchrun 的自动重启，就能把一次崩溃的代价压到“最多重跑几分钟”。

背景与问题

故障不是小概率事件。Llama 3 的技术报告披露，在 1.6 万张 H100 上做预训练，其中 54 天的统计窗口里出现了 466 次作业中断，419 次是意料之外的，约 78% 归因于硬件问题，仅 GPU 本身就占了意外中断的一半以上。Meta 最终把有效训练时间做到了 90% 以上，靠的是更快的启动、更快的 checkpoint 和自动化的故障处理。

卡数越多，整体平均无故障时间越短，所以问题其实有三个：

1. 怎么存：模型加优化器状态动辄上百 GB 到数 TB，传统做法是把所有分片汇总到 rank 0 再用torch.save写一个文件，既占满单卡显存，又让其余几千张卡干等。
2. 存的时候训练得停多久：每隔几百步存一次，每次停一分钟，累积起来的损耗相当可观。
3. 恢复时卡数变了怎么办：坏了一个节点，换成少一台机器继续跑，或者训练和推理用不同的并行配置，旧 checkpoint 还能不能读？

DCP 针对的就是这三件事。

核心思路与优势

• 每个 rank 并行写自己的分片：不再汇总到 rank 0。FileSystemWriter默认每个 rank 写一个文件，还可以用thread_count调写入线程数，把存储带宽吃满。
• 加载时重新切分：DCP 可以在任意数量的 rank 上并行保存，再在不同的集群拓扑上加载。保存时 64 卡 FSDP，恢复时 48 卡甚至换成别的并行方式，都能读。
• 原地加载：dcp.load把数据直接写进你已经分配好的 state_dict，没有“先在 CPU 上拼一份完整模型”的内存峰值。
• 异步保存：dcp.async_save先把状态拷到暂存区（默认是 CPU 内存），再由后台线程落盘，训练循环很快就能接着往下走。PyTorch 2.9 起还有DefaultStager，支持锁页内存和非阻塞拷贝。
• 统一的状态字典：get_state_dict返回、set_state_dict接收的都是以完整限定名（FQN）为键的状态，不依赖你用的是 FSDP、DDP 还是张量并行，这也是跨拓扑恢复能成立的前提。

需要提醒两点：async_save目前仍标注为实验性接口；官方明确说明，不保证不同 PyTorch 版本之间保存的 state_dict 能向后兼容，所以升级框架前要先验证旧 checkpoint 能读。

面向人群

• 用 FSDP、FSDP2 或张量并行训练大模型，已经被“一崩就从头来”折磨过的工程师
• 负责训练平台，需要给任务配置自动重启和断点续训的人
• 想在恢复时更换卡数或并行策略的团队
• 受不了每次保存训练都停顿很久的人

实践步骤

第一步：把需要保存的东西包成 Stateful

DCP 会在保存和加载时自动调用对象的state_dict()和load_state_dict()。除了模型和优化器，训练步数、学习率调度器、数据加载位置这些“续训必需”的东西也应该放进去。

importtorch.distributed.checkpointasdcpfromtorch.distributed.checkpoint.state_dictimportget_state_dict,set_state_dictfromtorch.distributed.checkpoint.statefulimportStatefulclassTrainState(Stateful):def__init__(self,model,optimizer,step=0):self.model=model self.optimizer=optimizer self.step=stepdefstate_dict(self):model_sd,optim_sd=get_state_dict(self.model,self.optimizer)return{"model":model_sd,"optim":optim_sd,"step":self.step}defload_state_dict(self,state_dict):set_state_dict(self.model,self.optimizer,model_state_dict=state_dict["model"],optim_state_dict=state_dict["optim"],)self.step=state_dict["step"]

第二步：同步保存与恢复

CKPT_DIR="/mnt/shared/ckpt"# 保存：每个 rank 并行写自己的分片state={"app":TrainState(model,optimizer,step)}dcp.save(state,checkpoint_id=f"{CKPT_DIR}/step_{step}")# 恢复：先建好模型和优化器，再原地加载state={"app":TrainState(model,optimizer)}dcp.load(state,checkpoint_id=f"{CKPT_DIR}/step_{latest}")start_step=state["app"].step

注意恢复前模型必须已经构建好（FSDP 包装完成），DCP 要先拿到分片信息，才能做加载时的重新切分。全新创建的优化器还没有状态，不用自己先跑一步：get_state_dict发现优化器状态为空时，会用零梯度、学习率临时置 0 的方式调用一次step()把状态初始化出来。另外官方要求set_state_dict/set_optimizer_state_dict只能在backward()之前或step()之后调用，别夹在两者之间加载。

第三步：改成异步保存，让训练别停

importtorch.distributedasdist dist.init_process_group(backend="cpu:gloo,cuda:nccl")ckpt_future=Noneforstepinrange(start_step,total_steps):train_one_step()done=step+1# 已完成的步数，恢复后从这里接着跑ifdone%500==0:ifckpt_futureisnotNone:ckpt_future.result()# 上一次还没落盘完就先等它，避免堆积state={"app":TrainState(model,optimizer,done)}ckpt_future=dcp.async_save(state,checkpoint_id=f"{CKPT_DIR}/step_{done}")ifckpt_futureisnotNone:ckpt_future.result()# 退出前等最后一次保存落盘

不传暂存器时，async_save返回的是一个普通的Future，GPU 到 CPU 的拷贝在函数返回前已经同步做完，所以返回后就可以继续训练、修改参数。

异步保存的后台线程要靠 CPU 通信组做协调，所以初始化时给 CPU 配上 gloo。在 2.9 及以上版本，还可以换成更细的暂存器，把 GPU 到 CPU 的拷贝也放到后台：

fromtorch.distributed.checkpoint.stagingimportDefaultStager stager=DefaultStager()# 建一次，循环里复用resp=dcp.async_save(state,checkpoint_id=f"{CKPT_DIR}/step_{done}",async_stager=stager,)train_forward_backward()# 拷贝在后台进行，前向和反向可以先跑resp.staging_completion.result()# 拷贝完成后，才能 optimizer.step() 修改参数...resp.upload_completion.result()# 训练结束前等最后一次落盘stager.close()# 自己创建的暂存器，用完要手动关闭

传入暂存器后，返回值就不再是单个Future，而是一个带两个 future 的对象：staging_completion表示本地拷贝已完成，这之后改动模型是安全的；upload_completion表示 checkpoint 已经完整落盘。DefaultStager默认开启锁页内存和后台拷贝，需要 CUDA 环境。

第四步：用 torchrun 做自动重启

DCP 解决“存下来”，进程怎么自动拉起来，交给 torchrun：

torchrun\--nnodes=$NUM_NODES\--nproc-per-node=8\--max-restarts=3\--rdzv-id=$JOB_ID\--rdzv-backend=c10d\--rdzv-endpoint=$HOST_NODE_ADDR\train.py

官方文档里有几条必须记住的规则：

• 只要有 worker 失败，所有worker 都会被停掉，再整体重启，最多重启--max-restarts次，所以脚本必须能从 checkpoint 续跑
• 重启之后RANK不保证稳定，不要把 rank 编号和数据分片、文件名硬绑定
• 如果用弹性模式（比如--nnodes=1:4），WORLD_SIZE也会变，同样不能写死；这正是 DCP 加载时重新切分的用武之地

torchrun 文档推荐的结构是“加载 checkpoint → 初始化 → 训练”；用 DCP 时模型和优化器必须先建好才能原地加载，所以实际顺序是建模型、加载、训练：

defmain():latest=find_latest_checkpoint(CKPT_DIR)# 自己写：扫描目录，找最新的完整 checkpointstate={"app":TrainState(model,optimizer)}iflatest:dcp.load(state,checkpoint_id=latest)train(start_step=state["app"].step)

第五步：几个容易踩的坑

• 别在半写的目录上恢复：崩溃可能发生在保存途中。保存完成后（upload_completion完成后）再写一个标记文件，恢复时只认带标记的目录。
• 保留多份：只保留最新一份，遇到最后一份恰好损坏就全完了。通常保留最近 2 到 3 份，再定期归档里程碑版本。
• 频率是个权衡：官方的说法是 checkpoint 频率取决于你的作业能容忍丢多少进度。保存越频繁，故障损失越小，但存储和带宽开销越大，异步保存的意义就在于把这个开销压下去。
• 存储要共享：重启后的 worker 可能落在不同的机器上，checkpoint 目录必须是所有节点都能访问的共享存储或对象存储。
• 导出给推理用：DCP 目录不是单个.pt文件，需要转成普通格式时，可以用torch.distributed.checkpoint.format_utils里的dcp_to_torch_save，反方向则用torch_save_to_dcp。

我的看法

一个 checkpoint 方案好不好，看三个指标：存得快不快、训练被打断多久、坏了能不能换个形状恢复。DCP 在这三点上都有对应的设计，而且是 PyTorch 原生的，不用额外引入第三方库。

如果你的任务还在单机几张卡上跑，torch.save够用；一旦上了多机、用了 FSDP，就应该尽早切到 DCP，等到集群规模上去、故障频繁了再改，迁移成本只会更高。