训练跑到一半崩了，用DCP从断点接着跑

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

摘要

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

背景与问题

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

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

怎么存：模型加优化器状态动辄上百 GB 到数 TB，传统做法是把所有分片汇总到 rank 0 再用 torch.save 写一个文件，既占满单卡显存，又让其余几千张卡干等。
存的时候训练得停多久：每隔几百步存一次，每次停一分钟，累积起来的损耗相当可观。
恢复时卡数变了怎么办：坏了一个节点，换成少一台机器继续跑，或者训练和推理用不同的并行配置，旧 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()。除了模型和优化器，训练步数、学习率调度器、数据加载位置这些“续训必需”的东西也应该放进去。

import torch.distributed.checkpoint as dcp
from torch.distributed.checkpoint.state_dict import get_state_dict, set_state_dict
from torch.distributed.checkpoint.stateful import Stateful


class TrainState(Stateful):
    def __init__(self, model, optimizer, step=0):
        self.model = model
        self.optimizer = optimizer
        self.step = step

    def state_dict(self):
        model_sd, optim_sd = get_state_dict(self.model, self.optimizer)
        return {"model": model_sd, "optim": optim_sd, "step": self.step}

    def load_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() 之后调用，别夹在两者之间加载。

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

import torch.distributed as dist

dist.init_process_group(backend="cpu:gloo,cuda:nccl")

ckpt_future = None
for step in range(start_step, total_steps):
    train_one_step()

    done = step + 1  # 已完成的步数，恢复后从这里接着跑
    if done % 500 == 0:
        if ckpt_future is not None:
            ckpt_future.result()  # 上一次还没落盘完就先等它，避免堆积
        state = {"app": TrainState(model, optimizer, done)}
        ckpt_future = dcp.async_save(state, checkpoint_id=f"{CKPT_DIR}/step_{done}")

if ckpt_future is not None:
    ckpt_future.result()  # 退出前等最后一次保存落盘

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

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

from torch.distributed.checkpoint.staging import DefaultStager

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 时模型和优化器必须先建好才能原地加载，所以实际顺序是建模型、加载、训练：

def main():
    latest = find_latest_checkpoint(CKPT_DIR)   # 自己写：扫描目录，找最新的完整 checkpoint
    state = {"app": TrainState(model, optimizer)}
    if latest:
        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，等到集群规模上去、故障频繁了再改，迁移成本只会更高。

posted @ 2026-10-13 09:36 RobinDevNotes 阅读(8) 评论(0) 收藏举报

刷新页面返回顶部