只保存模型权重只能用于推理或微调初始化,不能严格恢复训练。 一个可恢复训练 checkpoint 必须保存所有会影响后续参数更新的状态: 模型、优化器、学习率调度器、混合精度 scaler、训练步数、随机数状态、sampler 进度、数据版本和配置。 少任何一项,都可能让恢复后的 loss 曲线发生微妙偏移。

单卡时代常用 torch.save 把一个 state dict 写成单文件。 大模型分布式训练不同:参数和 optimizer state 可能被 FSDP/ZeRO 分片,单个 rank 聚合完整权重会慢、占内存、甚至 OOM。 PyTorch Distributed Checkpoint(DCP) 支持多 rank 并行保存和加载,并支持在加载时按当前并行拓扑重新分片。

最小 DCP 状态封装

import torch
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, scheduler, scaler, meta):
        self.model = model
        self.optimizer = optimizer
        self.scheduler = scheduler
        self.scaler = scaler
        self.meta = meta

    def state_dict(self):
        model_sd, optim_sd = get_state_dict(self.model, self.optimizer)
        return {
            "model": model_sd,
            "optim": optim_sd,
            "scheduler": self.scheduler.state_dict(),
            "scaler": self.scaler.state_dict() if self.scaler else {},
            "meta": self.meta,
            "rng_cpu": torch.get_rng_state(),
            "rng_cuda": torch.cuda.get_rng_state_all(),
        }

    def load_state_dict(self, state):
        set_state_dict(
            self.model,
            self.optimizer,
            model_state_dict=state["model"],
            optim_state_dict=state["optim"],
        )
        self.scheduler.load_state_dict(state["scheduler"])
        if self.scaler and state["scaler"]:
            self.scaler.load_state_dict(state["scaler"])
        self.meta.update(state["meta"])
        torch.set_rng_state(state["rng_cpu"])
        torch.cuda.set_rng_state_all(state["rng_cuda"])

def save_checkpoint(path, train_state):
    dcp.save({"app": train_state}, checkpoint_id=path)

def load_checkpoint(path, train_state):
    dcp.load({"app": train_state}, checkpoint_id=path)

checkpoint 保存常常是训练 step 的长尾:GPU 算完了,所有 rank 等着权重从 GPU 拷到 CPU,再写本地盘或对象存储。 异步保存的思路是把 checkpoint 从训练关键路径移走: 当前 step 触发保存后继续训练,后台线程/进程做 staging 和 upload。 代价是 CPU 内存、pinned memory 和并发管理更复杂。

一次只允许一个异步 checkpoint

checkpoint_future = None

for step, batch in enumerate(loader, start=start_step):
    loss = train_one_step(batch)

    if should_save(step):
        # 避免异步保存堆积:上一份还没写完,不要再发起下一份
        if checkpoint_future is not None:
            checkpoint_future.result()

        train_state.meta["step"] = step
        tmp_path = f"ckpt/step_{step:08d}.tmp"
        final_path = f"ckpt/step_{step:08d}"
        checkpoint_future = dcp.async_save({"app": train_state}, checkpoint_id=tmp_path)

        # 完成后再把 tmp 标记为 final。可放后台回调或保存管理线程里做。
        checkpoint_future.result()
        os.rename(tmp_path, final_path)
        write_text_atomic("ckpt/latest", final_path)

if checkpoint_future is not None:
    checkpoint_future.result()

恢复训练不是简单 load checkpoint。 可靠流程是:启动时发现 latest 指针,校验 checkpoint 完整性和配置一致性, 构建当前并行拓扑下的模型/optimizer,加载分布式状态,恢复 step、scheduler、scaler、RNG 和 sampler, 最后做一个 sanity step 检查 loss 和 LR 是否连续。

resume 入口模板

def maybe_resume(train_state, ckpt_root):
    latest_file = os.path.join(ckpt_root, "latest")
    if not os.path.exists(latest_file):
        return 0

    ckpt_path = open(latest_file).read().strip()
    meta = read_json(os.path.join(ckpt_path, "metadata.json"))
    assert meta["model_config_hash"] == current_model_config_hash()
    assert meta["data_fingerprint"] == current_data_fingerprint()

    dcp.load({"app": train_state}, checkpoint_id=ckpt_path)
    start_step = int(train_state.meta["step"]) + 1

    if dist.get_rank() == 0:
        print(f"resumed from {ckpt_path}, next step={start_step}")
    return start_step

torchrun elastic 能在 worker 失败后重启进程组,甚至支持节点数量在范围内变化。 但它只负责重新拉起 worker 和 rendezvous,不负责让你的训练状态神奇回来。 你的脚本必须在启动时加载最近 checkpoint,并能处理 rank/world size 变化带来的数据和分片重建。

# 固定节点数,允许最多 3 次失败重启
torchrun \
  --nnodes=2 \
  --nproc-per-node=8 \
  --max-restarts=3 \
  --rdzv-id=aiinfra-day24 \
  --rdzv-backend=c10d \
  --rdzv-endpoint=node0:29400 \
  train.py --ckpt-root /mnt/checkpoints/job_001

# 弹性节点数:最少 1 台,最多 4 台。节点加入/离开也算 restart budget。
torchrun \
  --nnodes=1:4 \
  --nproc-per-node=8 \
  --max-restarts=3 \
  --rdzv-id=aiinfra-day24 \
  --rdzv-backend=c10d \
  --rdzv-endpoint=node0:29400 \
  train.py --ckpt-root /mnt/checkpoints/job_001

今天必须做一次"故意杀进程"演练。没有故障演练的 checkpoint 都是心理安慰。 目标是证明:训练能从最近完整 checkpoint 恢复,global step 不倒退太多,LR 连续,loss 没有明显跳变, 未完成 checkpoint 不会被 latest 指向。

# 1. 启动一个允许重启的单机多卡任务
torchrun --standalone --nproc-per-node=4 --max-restarts=2 \
  train_day24_ckpt.py --ckpt-root ./ckpt_day24 --save-every 50

# 2. 另开终端,随机杀掉一个 rank 对应进程
ps aux | grep train_day24_ckpt.py
kill -9 PID_OF_ONE_WORKER

# 3. 观察 torchrun 是否重启 worker group,脚本是否从 latest 恢复
tail -f logs/train.log

# 4. 检查 checkpoint 目录
find ckpt_day24 -maxdepth 2 -type f | sort
cat ckpt_day24/latest

1. 列出一个严格恢复训练必须保存的状态清单,说明每项漏掉会造成什么问题。
2. DCP 和 torch.save 的差异是什么?为什么分布式训练更适合 DCP?
3. 解释异步 checkpoint 的 staging、upload、latest 更新三步,以及每一步失败时如何恢复。
4. 为什么 DCP load 需要先构建模型和 optimizer,再把 state dict 加载进去?
5. torchrun --max-restarts 能解决什么问题?不能解决什么问题?
6. 设计一次故障演练:你会杀哪个进程、观察哪些日志、如何判断恢复成功?