文档中心

Appearance

✨费用中心控制台

编写可恢复的训练代码 ​

平台可以重启任务或把任务重新调度到健康节点，但不能恢复 Python 进程内存中的模型、优化器、dataloader 和随机数状态。训练代码必须自己保存 checkpoint，并在重启后加载最近 checkpoint 继续训练。

在启动时加载最近 checkpoint ​

训练脚本启动时先检查 checkpoint 目录。如果存在最近 checkpoint，就加载模型、优化器、scheduler、epoch、global step 和随机数状态；如果不存在，再从头开始。

language-python
def main():
    state = load_latest_checkpoint(args.checkpoint_dir)
    model, optimizer, scheduler = build_training_objects()

    if state is not None:
        model.load_state_dict(state["model"])
        optimizer.load_state_dict(state["optimizer"])
        scheduler.load_state_dict(state["scheduler"])
        start_step = state["step"] + 1
    else:
        start_step = 0

    train(start_step, model, optimizer, scheduler)

不要只保存模型权重。只保存权重通常可以恢复推理或继续微调，但无法准确恢复长时间训练的学习率、优化器动量和进度。

把 checkpoint 写入共享高性能存储 ​

checkpoint 目录必须位于共享高性能存储中，不能位于 Worker 系统盘。

language-bash
export CHECKPOINT_DIR=/mnt/public/outputs/llama-sft/run-001/checkpoints

如果任务使用多个 Worker，所有 Worker 应能访问同一个 checkpoint 目录。

在 DDP 中只让一个进程保存 ​

PyTorch DDP 训练中，通常只让 rank 0 保存 checkpoint，然后让所有进程在启动时加载同一份 checkpoint。这样可以减少重复写入，也能避免多个进程同时写同一个文件。

language-python
if dist.get_rank() == 0:
    save_checkpoint(checkpoint_path, state)

dist.barrier()

保存完成后使用 barrier，避免其他进程在 checkpoint 尚未写完时继续执行。

让启动脚本可重复执行 ​

任务被重启后会重新执行 启动命令。启动脚本应允许重复运行：

  • 创建目录使用 mkdir -p。
  • 已存在的数据、代码和 checkpoint 不要重复下载或覆盖。
  • 输出路径按 run id、任务名称或环境变量区分，避免重跑覆盖旧结果。
  • 失败时返回非 0 退出码，让平台能触发容错或标记失败。

启动命令 写法详见编写可靠启动命令。

记录恢复证据 ​

每次从 checkpoint 恢复时，在日志中打印 checkpoint 路径、恢复的 step、epoch 和关键训练状态。后续排查 自动重启、任务 重跑 或 Spot 资源中断时，可以直接从日志确认训练是否真的续上。

完成代码改造后，继续验证 checkpoint 保存与恢复。