在 GPU算力平台 上进行大模型训练时，保存了模型和优化器并不等于能“原地续跑”。GPU服务器租用实例重启后，损失曲线突然跳变、样本顺序重复，往往是随机状态、学习率调度器或数据游标没有恢复。本文给出一套可验证的断点续训方案。

问题背景

训练状态不只有权重。优化器动量、学习率步数、混合精度缩放器、Python 与 NumPy 随机数、PyTorch CPU/GPU 随机状态、当前 epoch 和批次位置都会影响下一步计算。深度学习任务若只保存 state_dict，恢复后虽然能运行，却可能重复消费数据或改变增强结果。对于长周期训练，这类“静默偏移”比直接报错更难定位，也会影响后续推理部署的模型验收。

环境准备

准备 Python、PyTorch、固定版本的数据集清单和可写检查点目录。多进程训练还需记录 world_size、全局 batch 和梯度累积步数。按项目实际依赖选择实例；本文不假设固定硬件或镜像版本。

编号实操步骤

1. 固化训练配置

把影响数据顺序和有效 batch 的参数写入 JSON，恢复时先比较而不是直接覆盖：

import json
config = {
    "seed": 2026,
    "batch_size": 4,
    "grad_accum": 8,
    "world_size": 2,
    "dataset_manifest": "manifest-v3.jsonl"
}
open("run_config.json", "w").write(
    json.dumps(config, ensure_ascii=False, indent=2)
)

数据清单应包含稳定样本 ID，避免仅依赖随时变化的目录遍历顺序。

2. 统一初始化随机种子

import random, numpy as np, torch

def seed_all(seed):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)

seed_all(2026)

初始化种子只解决起点一致，断点处仍需保存当时的随机状态。

3. 保存完整检查点

state = {
    "model": model.state_dict(),
    "optimizer": optimizer.state_dict(),
    "scheduler": scheduler.state_dict(),
    "scaler": scaler.state_dict(),
    "epoch": epoch,
    "batch_in_epoch": batch_idx + 1,
    "global_step": global_step,
    "python_rng": random.getstate(),
    "numpy_rng": np.random.get_state(),
    "torch_rng": torch.get_rng_state(),
    "cuda_rng": torch.cuda.get_rng_state_all(),
}
torch.save(state, "checkpoint.tmp")

保存完成后再把临时文件原子改名为正式检查点，防止中断留下半文件：

mv checkpoint.tmp checkpoint.pt
sha256sum checkpoint.pt > checkpoint.pt.sha256

4. 按正确顺序恢复

先构建相同结构的模型、优化器、调度器和缩放器，再加载状态：

ckpt = torch.load("checkpoint.pt", map_location="cpu")
model.load_state_dict(ckpt["model"])
optimizer.load_state_dict(ckpt["optimizer"])
scheduler.load_state_dict(ckpt["scheduler"])
scaler.load_state_dict(ckpt["scaler"])

random.setstate(ckpt["python_rng"])
np.random.set_state(ckpt["numpy_rng"])
torch.set_rng_state(ckpt["torch_rng"])
torch.cuda.set_rng_state_all(ckpt["cuda_rng"])
global_step = ckpt["global_step"]

优化器状态加载后，要检查其中张量是否已迁移到正确设备。

5. 恢复数据游标

对可索引数据集，可从保存的 epoch 重新构建采样器并跳过已完成批次：

sampler.set_epoch(ckpt["epoch"])
for i, batch in enumerate(loader):
    if i < ckpt["batch_in_epoch"]:
        continue
    train_step(batch)

这是一种易理解的基线方案，但跳过过程仍会读取数据。超大数据集或流式数据应保存可恢复的样本游标或分片偏移，并以数据组件的官方能力为准。

6. 校验恢复一致性

在首次训练时，额外保存断点后一个批次的样本 ID、学习率、损失和参数摘要。恢复后只跑一步并对比：

print({
    "step": global_step,
    "lr": optimizer.param_groups[0]["lr"],
    "sample_ids": batch["id"][:5],
})

浮点结果可能受硬件和算子实现影响，不应把“逐位一致”设为所有项目的硬性标准；重点是样本位置、调度步数和曲线连续性。

7. 处理并行规模变化

如果恢复时 GPU 数、全局 batch 或梯度累积发生改变，学习率轨迹和数据切分也会变化。AI算力平台 迁移任务前，应先决定是追求严格续跑，还是接受新配置下的再基线测试。GPU算力平台 的恢复记录需保留配置差异、检查点哈希和首步验证结果。

常见问题与解决方案

• 恢复后学习率回到初值：调度器未保存，或加载顺序错误；检查 global_step 与调度器状态。
• 首批数据重复：仅恢复 epoch，未记录批次或样本游标。
• 损失突然增大：核对优化器、混合精度缩放器和有效全局 batch。
• 多卡恢复卡住：确认所有 rank 使用同一检查点版本，并检查并行规模与启动参数。

总结

可靠断点续训应同时恢复计算状态和数据位置，并通过“断点后第一步”验证连续性。润云智算可作为计算环境入口，具体资源以官网实时信息为准。把配置快照、哈希校验和恢复回归加入训练脚本，能减少长任务因重启或迁移产生的隐性偏差。

FAQ

1. 只保存模型权重适合什么情况？

适合推理或重新开始优化，不适合要求曲线连续的严格续训。

2. 每个 step 都保存检查点吗？

通常没有必要。应结合单步耗时、可接受回退量和存储开销确定周期。

3. 更换 GPU 后还能严格复现吗？

不一定。算子、驱动和硬件差异可能引入数值变化，应先做短程回归。

4. 检查点需要保存数据集吗？

不必复制全部数据，但必须锁定数据版本、清单和可恢复游标。