训练跑到一半挂了。这种事情做深度学习的人应该都遇到过——模型跑到第七个 epochloss 曲线好不容易降下来结果集群被调度抢占、服务器重启、或者手滑把训练进程杀了重启之后只能从第 0 个 epoch 重新来。以前短训练没什么感觉但模型一上规模、数据一多这种“从头再来”的代价是灾难性的。断点续训resume training要解决的就是这个问题让模型中断之后能接着上次的进度继续跑而不是从零开始。它看起来只是“保存模型、再加载模型”这么简单但实际做起来门道不少。我见过很多人做 resume 只保存了模型权重恢复之后虽然能跑但 loss 断层、学习率错乱、随机数状态对不上训练质量受到明显影响。这篇文章就把断点续训里到底该保存什么、恢复时要注意什么、踩过的坑怎么排查完整展开聊聊。适合所有训练周期长、经常被中断打断的工程型选手参考。1. 断点续训到底在解决什么问题1.1 训练中断的三个高频场景先说清楚一个前提不是所有中断都需要续训。如果你的训练脚本跑十分钟就完事那中断了重跑也没关系。真正需要断点续训的是那些训练周期以天甚至周为单位的场景这一类场景里中断几乎是必然发生的。最常见的三类中断资源被抢占。很多团队用共享 GPU 集群跑训练排队排了半天好不容易上机开始训结果来了更高优先级的任务你的进程直接被杀掉。这种中断完全没有预兆kill 信号下来你连准备时间都没有。硬件或系统异常。电源波动、主机重启、GPU 驱动挂掉这一类属于“能不能恢复全靠运气”的情况。运气好的时候 checkpoint 文件还在运气不好连系统都进不去。手动中断。比如你发现 loss 不降了想改学习率策略或者数据有问题要停掉处理再或者你想换个 batch size 重新试——但不想丢掉之前训练的进度。无论哪种中断恢复的关键点都一样尽量精确地回到中断点让训练过程像没有中断过一样继续。1.2 checkpoint 的定位不是存档是“恢复现场”很多人一听到“断点续训”第一反应是“save 模型、load 模型”。模型权重确实是核心但只保存模型权重相当于看书看到一半只记住了“我读到了第 137 页”却把所有上下文丢掉了。训练过程里影响训练走向的不只是参数还有优化器内部状态、学习率当前位置、随机数发生器状态、数据加载进度等等。我习惯把这个整体状态叫“训练现场”。checkpoint 文件本质上不是存档而是一次完整的“现场恢复包”。它要回答的问题是当训练重新跑起来那一刻从外部看它跟没中断过的状态是否一致。这个“一致性”越是接近训练效果就越稳定。如果你只恢复了权重优化器里的动量和方差全部归零Adam 从第一个 step 重新积累统计量——要知道 Adam 的方差项exp_avg_sq是在训练早期就快速建立起来的丢掉之后恢复训练前几百个 step 的更新步长会明显异常直观表现就是 loss 先飙一下再慢慢回来白白浪费一段算力。所以完整考虑 checkpoint 应该包含哪些内容是一切设计的基础。2. 想恢复现场先搞懂这三样东西2.1 模型参数续训的地基模型参数是 checkpoint 里最直观的部分。在 PyTorch 里一个nn.Module实例调用state_dict()得到一个有序字典里面是每个可学习参数的张量和注册的 buffer比如 BatchNorm 的 running_mean、running_var。这个字典是个纯数据容器不包含模型结构代码但它能完整描述“当前权重是什么”。保存时要注意一个细节如果模型被torch.nn.DataParallel或torch.nn.parallel.DistributedDataParallel包过一层model.state_dict()的 key 名会带module.前缀。恢复的时候要判断模型是否包着这层 wrapper否则会出现 key 对不上的问题。我的做法是保存前统一取model.module.state_dict()如果存在恢复时再直接加载回原模型这样 key 名始终干净一致。另外多说一句不要只保存model.state_dict()顺便把模型结构相关的配置也存进去。比如你用了多深的层、多少维度的 embedding、vocab size 是多少。因为网络结构代码本身经常会有小改动有了配置信息至少能帮你判断这份 checkpoint 到底对应哪个版本的结构。2.2 优化器状态比参数更容易被忽略优化器状态是断点续训里最重要的部分也是最容易被忽略的部分。以最常用的 Adam 为例它的state_dict()里保存了每个参数对应的 exp_avg一阶动量和 exp_avg_sq二阶原始动量。这两个是梯度均值和梯度平方均值的历史累积它们决定了每个参数在每个 step 的更新幅度和方向。SGD 加上 momentum 之后optimizer 内部也会保存一个momentum_buffer同样属于历史状态。如果这些丢了优化器就“失忆”了哪怕模型权重一模一样继续训练的轨迹也不会一样。PyTorch 的optimizer对象也有state_dict()同时它还会包含param_groups里面是学习率、权重衰减、beta 这些超参数。这里给了我们一个额外好处如果你在恢复 checkpoint 时想把学习率改小直接改加载之后的 optimizer.param_groups 就可以不需要动其它代码。类似地学习率调度器也要恢复。ReduceLROnPlateau这类调度器内部有 best 值记录CosineAnnealingLR依赖当前 epoch 数计算下一阶段的学习率OneCycleLR有 step 计数。只恢复模型和 optimizer、忘记恢复 scheduler训练能跑但学习率从初始值重新走一遍效果大概率是一团糟。2.3 随机数状态和训练进度元信息第三类是随机数状态这个最容易被忽略但对“严格恢复到中断点”非常重要。现代训练流程里到处都是随机性数据加载时的 shuffle、数据增强随机裁剪、翻转、颜色扰动、dropout 的 mask、初始化噪声、多进程 worker 的采样……这些随机性由 Python 的 random 模块、NumPy 的随机数生成器和 PyTorch 的 torch random 共同管理。每产生一组随机值全局状态就会前进一步。如果不保存这些状态恢复训练时会重新生成一遍随机序列相当于某一段训练用了两批不同的随机增强和 dropout mask。那是不是说必须恢复随机数状态这里要分情况如果数据加载用了 DataLoader 的多 worker且你关心“断点之后的数据顺序必须和不断点时完全一致”那你需要保存每个 worker 的随机状态否则做不到严格一致。如果数据增强强度不太大、batch size 比较大或者你每一轮epoch都重新 shuffle 一遍那恢复随机状态的意义就小很多因为下一轮的随机序列无论怎么走都行。我个人的习惯是在 checkpoint 里保存随机数状态但把是否恢复做成开关。默认开着如果发现因为 DataLoader 多进程导致恢复困难再关掉。保存方式也很简单四个模块的随机状态分别取出来存进 dict 就行import random import numpy as np import torch def capture_rng_state(): return { python: random.getstate(), numpy: np.random.get_state(), torch: torch.random.get_rng_state(), cuda: torch.cuda.get_rng_state_all() if torch.cuda.is_available() else None, }元信息更直接当前训练到第几个 epoch、第几个 step、当前 best metric、best checkpoint 的路径、以及你自定义的任何训练状态。这些信息在恢复时用来决定从哪个 epoch 继续、调度器从哪里接上、以及判断模型性能是否已经刷新到最优。说到恢复随机状态有一个实践上的难处DataLoader 的 worker seed 是在每个 worker 进程启动时重新初始化的恢复全部随机状态不一定能保证每个 worker 内部恢复一致。对于这类严格一致性的要求我后面在第 4 部分专门讲到一个替代方案这里先不展开。3. 实操过程从保存到恢复的完整链路3.1 保存 checkpoint 的推荐做法先给一个我自己在用的保存函数模板。它把前面讨论的内容全部落到了一个字典里def save_checkpoint( model, optimizer, scheduler, epoch, global_step, best_metric, config, rng_state, filepath, ): checkpoint { model_state_dict: model.state_dict(), optimizer_state_dict: optimizer.state_dict(), scheduler_state_dict: scheduler.state_dict() if scheduler else None, epoch: epoch, global_step: global_step, best_metric: best_metric, config: config, rng_state: rng_state, } torch.save(checkpoint, filepath)注意几个细节第一torch.save默认使用 Python 的 pickle 序列化。如果你只关心权重不关心完整的训练状态可以加_use_new_zipfile_serializationTrue参数它能生成 zip 格式的 checkpoint加载更快、文件展开更友好。但前提是你的代码环境兼容新格式一般现代 PyTorch 版本都没问题。第二保存前先检查一遍状态有没有落在正确的 device 上。很多人训练用 GPU保存时发现 state_dict 里的张量在 CUDA 上于是 checkpoint 文件巨大而且加载时还得小心翼翼。实际上 PyTorch 的torch.save保存的是张量的数据副本不要求先把模型挪到 CPU 上——它会正常序列化。但在加载时用map_locationcpu再统一搬到目标设备是最稳妥的套路。第三保存频率要有策略。每多少个 step 存一次我常用的策略是“每 N 个 epoch 存一个完整 checkpoint 每个 epoch 内若干个 checkpoint”具体频率看训练时长。如果一次训练要跑一周每 10 分钟一个 checkpoint 都不过分如果一次训练只要两小时每半小时存一次就够了。这里没有标准答案核心理念是checkpoint 文件丢失的代价 vs 频繁保存的磁盘和 IO 开销取一个平衡点。第四最好用“临时文件 原子替换”的方式写盘。正确做法是torch.save(checkpoint, filepath .tmp) os.replace(filepath .tmp, filepath)为什么这么做如果你的训练正好在torch.save写文件写到一半时被 kill —— 无论是 SIGKILL 还是断电 —— 主文件就会是一个截断的损坏文件下次加载直接报错。先写临时文件再原子替换能保证磁盘上的 checkpoint 文件要么是完整的旧版本要么是完整的新版本不会出现半截文件。这个小技巧我吃过亏之后就成了习惯。3.2 恢复训练的完整流程恢复训练的代码比保存稍微复杂一点因为要考虑 device 映射和加载容错。def load_checkpoint(filepath, model, optimizer, scheduler, device): checkpoint torch.load(filepath, map_locationcpu) model.load_state_dict(checkpoint[model_state_dict]) model.to(device) optimizer.load_state_dict(checkpoint[optimizer_state_dict]) # 加载 optimizer 后把参数组里的学习率等超参数同步到当前训练配置 # 例如optimizer.param_groups[0][lr] your_new_lr if scheduler is not None and checkpoint.get(scheduler_state_dict) is not None: scheduler.load_state_dict(checkpoint[scheduler_state_dict]) start_epoch checkpoint[epoch] 1 global_step checkpoint.get(global_step, 0) best_metric checkpoint.get(best_metric, float(inf)) # 可选恢复随机数状态 if checkpoint.get(rng_state) is not None: restore_rng_state(checkpoint[rng_state]) return start_epoch, global_step, best_metricmap_locationcpu是我强推的做法。哪怕你训练时用的是 GPU加载时也先全部放到 CPU然后统一model.to(device)。原因有两个如果你直接加载到原来的设备但当前进程的 GPU 编号和保存时不同比如保存时用的是 cuda:0现在机器上只有 cuda:1 可用会报设备不匹配错误。大模型如果直接加载进显存临时占用很大万一显存本来就很紧张加载那一下也可能 OOM。放 CPU 再转移到 GPU操作更安全。恢复之后训练主循环要跳过已经跑过的数据步数。epoch 从start_epoch开始循环内部用global_step恢复全局步数计数如果用的是 logging 和 checkpoint 保存的步数索引。还有一个细节如果用了分布式训练需要把distributed_sampler.set_epoch(start_epoch)设置正确否则每个 epoch 的数据 shuffle 会从第 0 个 epoch 的顺序重新开始这和恢复的“现场”就不一致了。随机数状态恢复的函数写起来也不复杂def restore_rng_state(state): random.setstate(state[python]) np.random.set_state(state[numpy]) torch.random.set_rng_state(state[torch]) if state[cuda] is not None and torch.cuda.is_available(): torch.cuda.set_rng_state_all(state[cuda])3.3 用“杀进程”实测验证恢复逻辑写完保存和恢复的代码你怎么知道自己做得对不对我强烈建议你一上来就做一个“杀进程中断实验”这种验证可以一次性把恢复逻辑的问题全都试出来。具体操作如下写一个训练脚本跑 20 个 step。每个 step 结束时检查当前 step 是否是 10如果是就保存一个 checkpoint然后手动 raise 一个中断。再写一个恢复脚本或者同一个脚本加 resume 参数从刚才保存的 checkpoint 加载继续跑完剩余 10 个 step。对比两种跑法——一种是完全不中断、从头跑 20 个 step另一种是中断后从 checkpoint 继续跑剩下的 10 个 step——两者第二次 10 个 step 的 loss 曲线、学习率变化以及最终模型权重是否一致。这里要说清楚一个预期如果随机数状态完整恢复且设备一致两次训练的结果理论上应该完全一致但由于浮点运算的并行顺序可能有微小差异实际可能有非常细微的差别。如果 loss 曲线走势明显不一样很大概率是你没有恢复随机数状态、或者 optimizer / scheduler 状态丢了。我的验证经验里有个经典翻车现场只恢复了模型和 optimizer没恢复 scheduler 和随机状态第一次跑 20 步和恢复后半段的 loss 曲线完全对不上。后来补上了 random state 和 scheduler 的恢复曲线基本吻合才确认了问题所在。3.4 文件命名、轮转与防写坏持续训练会产生大量 checkpoint。如果不做管理磁盘很快会被塞满而磁盘被塞满又会诱发新的中断——这就形成了一个恶性循环。我常用的文件管理方案是checkpoints/ ├── last.pt # 最近一次保存的完整 checkpoint ├── best.pt # 当前在验证集上指标最好的 checkpoint ├── epoch_0005_step_2000.pt ├── epoch_0010_step_4000.pt └── ...保存逻辑上每次保存都把last.pt覆盖更新。同时保留一个带 epoch/step 信息的版本文件便于回溯。如果新一轮验证指标超过历史 best再更新best.pt。用一个简单的保留策略最多保留最近 K 个版本文件K 取决于磁盘空间。比如只保留最近 5 个带 epoch 的 checkpoint超出就删除最老的。def prune_checkpoints(checkpoint_dir, keep_last5): files sorted(glob.glob(os.path.join(checkpoint_dir, epoch_*.pt))) for f in files[:-keep_last]: os.remove(f)这套方案的成本很低但价值巨大任何时候崩溃都有一个last.pt保证“最多丢失最后保存间隔内的进度”也有一个best.pt保证“只要保存过最好的结果后续就算训练跑费了也不亏”。这个习惯我一直保留尤其在做长周期训练并且没有专人盯着训练进程的时候它可以救命。4. 常见问题与排查技巧实录4.1 恢复后 loss 跳变是恢复出问题了吗很多人恢复训练后第一眼看到 loss 比中断前高了一点立刻怀疑恢复逻辑有 bug。其实这里有多个原因。最常见的原因是数据处理顺序改变。你在中断时保存了 epoch 和 step但 resume 之后脚本重新创建了 DataLoadershuffle 的顺序由新的随机种子或者新初始化的随机数状态决定如果随机状态没有恢复新一轮前面几步的 batch 就和原来那些 “本应看到” 的 batch 不一样了。这个差异本身有时会带来 loss 波动但通常幅度不大会稳定下来。另一个常见原因是恢复模型后没有把模型设置为训练模式。如果你在恢复流程里先做了验证或者其他 eval 操作模型带着 eval 状态回到训练循环BatchNorm 和 Dropout 行为完全不同loss 当然会异常。检查一下model.train()是否有被调用。如果 loss 跳变幅度很大而且持续不回落那才需要怀疑优化器或 scheduler 没恢复。排查方法很简单加载 checkpoint 后打印 optimizer 的 param_groups 里 lr 值以及 scheduler 当前的 last_epoch / best 值对比中断前保存时日志里的值。两边对不上说明恢复链路掉了一环。4.2 加载报 Missing/Unexpected key 怎么办这个报错几乎每个做迁移和 resume 的人都会遇到。Missing key(s)的意思是 checkpoint 里有某些 key 在当前模型中不存在或者当前模型有某些参数在 checkpoint 中不存在。出现这个错常见原因是模型结构代码改动过比如加了一层、删了一层、改了层名。checkpoint 是从带module.前缀的并行模型保存的加载到一个普通模型上或者反过来。你保存的时候是 part1 part2 的混合模型加载到新模型结构时名字对不上。排查的第一步不是去看网络代码而是先打印两边的 key 列表来做差异state torch.load(last.pt, map_locationcpu) ckpt_keys set(state[model_state_dict].keys()) model_keys set(model.state_dict().keys()) print(missing:, model_keys - ckpt_keys) print(unexpected:, ckpt_keys - model_keys)看清楚差异再动手。如果只是module.前缀问题写一个 remap 就好def remove_module_prefix(state_dict): from collections import OrderedDict new_state OrderedDict() for k, v in state_dict.items(): new_k k.replace(module., ) if k.startswith(module.) else k new_state[new_k] v return new_state如果确实是结构改了那要看差异 keys 是否对应你想剔除或新增的层。比如你想从 checkpoint 加载一个主干网络的权重到新的下游任务模型直接用strictFalse是可行的model.load_state_dict(checkpoint[model_state_dict], strictFalse)但强烈建议你在 strictFalse 之后手动验证那些被忽略的层是哪些。我遇到过一次很坑的情况用 strictFalse 加载完所有权重都提示 missing 了但因为某个命名空间错误实际加载进去的只有一层卷积其余全是随机初始化。训练出来效果极差排查了半天才靠核对 key 数量发现问题。所以记住strictFalse 只是方便迁移不代表权重一定加载正确。4.3 checkpoint 文件太大磁盘快爆了大模型的 checkpoint 动辄几个 GB。如果你每十分钟存一次磁盘压力非常真实。办法有这么几个只保存最近的 N 个版本用前面提到的轮转策略清理。判断是否真的需要每次都保存整个 optimizer 状态。如果你用的是 Adamoptimizer state 和模型参数差不多一样大exp_avg 和 exp_avg_sq 各一份checkpoint 体积直接翻三倍。折中方案是平时每 N 步保存一个“轻量 checkpoint”只含模型和元信息每 N 个 epoch 才保存一个“完整 checkpoint”。万一轻量 checkpoint 恢复的时候 optimizer 状态丢了训练还能继续只是需要一段预热时间。如果保存完整 checkpoint 的需求无法压缩考虑用torch.save的新 zip 格式减少体积同时不要用 float64 保存不需要的中间状态统一用 float32。文件大还有一个隐蔽问题有的文件系统对大文件单次写入不稳定。你需要确保写文件的进程有足够磁盘空间余量尤其临时文件 原子替换方案中临时文件占两份空间磁盘设计时要留出这个余量。4.4 多卡训练里断点续训的额外坑多卡训练尤其是分布式训练的 resume 比单卡复杂一档。基本模式是只有 rank 0 保存 checkpoint所有 rank 加载同一份 checkpoint。保存时要注意直接取模型的 state_dict 而不是带 DDP wrapper 的。如果使用了DistributedDataParallel加载时最快的方式是model DistributedDataParallel(model, device_ids[local_rank]) model.module.load_state_dict(checkpoint[model_state_dict])加载时每个进程都要做但要注意torch.load(map_locationcpu)然后每个进程各自to(device)。这样每个 rank 拿到的权重一样不会出现 device mismatch。DistributedSampler也要额外设置。它的默认行为是每个 epoch 用self.epoch作为随机种子。如果你不把sampler.set_epoch(start_epoch)调成恢复的 epoch那么恢复后第一个 epoch 的 shuffle 顺序其实还是第 0 个 epoch 的顺序数据采样完全错乱。还有一个多卡特有的问题如果中断后的恢复是在不同数量的卡上进行的比如原来 8 卡、现在只有 4 卡这时 batch size 变了。严格来说梯度累积、学习率策略都可能需要相应调整。你可以选择直接继续跑多数情况下也能收敛但曲线会有一点波动也可以选择重启一个小的热身阶段让优化器适应新的分布式状态。4.5 数据加载顺序变了怎么尽量保持一致这个问题前面提过现在展开说一下。如果你用的 DataLoader 开了多进程 workernum_workers 0并且每个 epoch 内也依赖确定的数据顺序那仅仅恢复random、numpy、torch的全局随机状态是不够的。每个 worker 进程在启动时会用自己的一套随机种子通常是主进程 seed worker id 的组合来初始化而主进程在for batch in dataloader的迭代过程中是通过一个迭代器不断从 worker 取数据的中断后重建 DataLoaderworker 内部的数据采样状态早已丢失。我在实践中面对这个问题的解法是这样的如果确认数据增强对最终结果影响不大比如简单的分类任务用了固定 resiz 和 normalize没带随机裁剪就直接不追求数据顺序一致。每次 resume 后从当前 epoch 继续但 batch 顺序可能变化这完全没关系模型很快就能适应。如果是有强随机增强的图像任务、或者是需要在同一 epoch 内精确恢复的对比学习类任务我会改用“索引式恢复”定义一个 DataLoader 包装器自己控制每次迭代的样本索引把当前 index 也存进 checkpoint。恢复时从保存的 index 继续而不是从 epoch 开头。console 上做这种控制更麻烦但是可控性最好。有人可能会问既然随机状态恢复这么麻烦那为什么不干脆每次中断都从 epoch 开头重新跑我的看法是当你想做严格的消融实验、想对比不同随机种子下的结果时数据顺序的一致性真的很重要但大多数生产场景里模型收敛才是硬指标batch 顺序有一点变化完全可以接受。所以这个坑反而可以不用过度纠结把精力花在更值得的地方。5. 几个我长期在用的续训习惯5.1 把“可恢复性”设计到训练脚本最前面不要等到训练跑起来了才想断点续训的事。我现在的训练脚本统一在命令行参数里加一个--resume参数默认从last.pt恢复。脚本启动的时候做三件事先解析命令行、再检查 checkpoint 是否存在、存在就加载并覆盖训练配置。没有一个模块是后加的全部是训练脚本的组成部分。这样训练任务的启动命令就非常简单无论是手动跑还是配合调度系统都能拿到一样的语义python train.py --config config.yaml --resume checkpoints/last.pt5.2 每次保存都顺手打印状态摘要训练中断后恢复最怕的就是不知道当前 checkpoint 对应训练到哪个阶段。所以我每次 save_checkpoint 之后都会往训练日志里打印一行[checkpoint] saved to checkpoints/last.pt | epoch 7 | step 2860 | loss 1.832 | lr 3.2e-4虽然看起来是个不起眼的习惯但它有实际价值恢复脚本启动时我会把 checkpoint 里的 epoch 和 step 打印出来和日志里的这行对比能快速发现“到底恢复了哪个版本”的问题。有一次我朋友在调试一个项目发现恢复后训练挂了一整天才意识到加载的是半个月前的旧 checkpoint就是因为日志里没有这种摘要信息排查非常痛苦。最后再分享一个小习惯checkpoint 保存目录和日志目录尽量分开并且定期检查磁盘剩余空间。我做训练时会把 checkpoint 挂到存储较大的磁盘分区上因为训练中断导致进度丢失很恼火但更恼火的是磁盘满了导致 checkpoint 保存不了结果连最后的“避难所”都没了。先把现场完整保存下来后面出错都好说。这套经验基本能覆盖绝大多数断点续训的开发和使用场景希望你能在下次训练中断时少一点手忙脚乱。 
长河君 · 2026/10/19 13:26:40 · 阅读 1.2k 深度学习断点续训实战：保存恢复checkpoint的完整指南
训练跑到一半挂了。这种事情做深度学习的人应该都遇到过——模型跑到第七个 epochloss 曲线好不容易降下来结果集群被调度抢占、服务器重启、或者手滑把训练进程杀了重启之后只能从第 0 个 epoch 重新来。以前短训练没什么感觉但模型一上规模、数据一多这种“从头再来”的代价是灾难性的。断点续训resume training要解决的就是这个问题让模型中断之后能接着上次的进度继续跑而不是从零开始。它看起来只是“保存模型、再加载模型”这么简单但实际做起来门道不少。我见过很多人做 resume 只保存了模型权重恢复之后虽然能跑但 loss 断层、学习率错乱、随机数状态对不上训练质量受到明显影响。这篇文章就把断点续训里到底该保存什么、恢复时要注意什么、踩过的坑怎么排查完整展开聊聊。适合所有训练周期长、经常被中断打断的工程型选手参考。1. 断点续训到底在解决什么问题1.1 训练中断的三个高频场景先说清楚一个前提不是所有中断都需要续训。如果你的训练脚本跑十分钟就完事那中断了重跑也没关系。真正需要断点续训的是那些训练周期以天甚至周为单位的场景这一类场景里中断几乎是必然发生的。最常见的三类中断资源被抢占。很多团队用共享 GPU 集群跑训练排队排了半天好不容易上机开始训结果来了更高优先级的任务你的进程直接被杀掉。这种中断完全没有预兆kill 信号下来你连准备时间都没有。硬件或系统异常。电源波动、主机重启、GPU 驱动挂掉这一类属于“能不能恢复全靠运气”的情况。运气好的时候 checkpoint 文件还在运气不好连系统都进不去。手动中断。比如你发现 loss 不降了想改学习率策略或者数据有问题要停掉处理再或者你想换个 batch size 重新试——但不想丢掉之前训练的进度。无论哪种中断恢复的关键点都一样尽量精确地回到中断点让训练过程像没有中断过一样继续。1.2 checkpoint 的定位不是存档是“恢复现场”很多人一听到“断点续训”第一反应是“save 模型、load 模型”。模型权重确实是核心但只保存模型权重相当于看书看到一半只记住了“我读到了第 137 页”却把所有上下文丢掉了。训练过程里影响训练走向的不只是参数还有优化器内部状态、学习率当前位置、随机数发生器状态、数据加载进度等等。我习惯把这个整体状态叫“训练现场”。checkpoint 文件本质上不是存档而是一次完整的“现场恢复包”。它要回答的问题是当训练重新跑起来那一刻从外部看它跟没中断过的状态是否一致。这个“一致性”越是接近训练效果就越稳定。如果你只恢复了权重优化器里的动量和方差全部归零Adam 从第一个 step 重新积累统计量——要知道 Adam 的方差项exp_avg_sq是在训练早期就快速建立起来的丢掉之后恢复训练前几百个 step 的更新步长会明显异常直观表现就是 loss 先飙一下再慢慢回来白白浪费一段算力。所以完整考虑 checkpoint 应该包含哪些内容是一切设计的基础。2. 想恢复现场先搞懂这三样东西2.1 模型参数续训的地基模型参数是 checkpoint 里最直观的部分。在 PyTorch 里一个nn.Module实例调用state_dict()得到一个有序字典里面是每个可学习参数的张量和注册的 buffer比如 BatchNorm 的 running_mean、running_var。这个字典是个纯数据容器不包含模型结构代码但它能完整描述“当前权重是什么”。保存时要注意一个细节如果模型被torch.nn.DataParallel或torch.nn.parallel.DistributedDataParallel包过一层model.state_dict()的 key 名会带module.前缀。恢复的时候要判断模型是否包着这层 wrapper否则会出现 key 对不上的问题。我的做法是保存前统一取model.module.state_dict()如果存在恢复时再直接加载回原模型这样 key 名始终干净一致。另外多说一句不要只保存model.state_dict()顺便把模型结构相关的配置也存进去。比如你用了多深的层、多少维度的 embedding、vocab size 是多少。因为网络结构代码本身经常会有小改动有了配置信息至少能帮你判断这份 checkpoint 到底对应哪个版本的结构。2.2 优化器状态比参数更容易被忽略优化器状态是断点续训里最重要的部分也是最容易被忽略的部分。以最常用的 Adam 为例它的state_dict()里保存了每个参数对应的 exp_avg一阶动量和 exp_avg_sq二阶原始动量。这两个是梯度均值和梯度平方均值的历史累积它们决定了每个参数在每个 step 的更新幅度和方向。SGD 加上 momentum 之后optimizer 内部也会保存一个momentum_buffer同样属于历史状态。如果这些丢了优化器就“失忆”了哪怕模型权重一模一样继续训练的轨迹也不会一样。PyTorch 的optimizer对象也有state_dict()同时它还会包含param_groups里面是学习率、权重衰减、beta 这些超参数。这里给了我们一个额外好处如果你在恢复 checkpoint 时想把学习率改小直接改加载之后的 optimizer.param_groups 就可以不需要动其它代码。类似地学习率调度器也要恢复。ReduceLROnPlateau这类调度器内部有 best 值记录CosineAnnealingLR依赖当前 epoch 数计算下一阶段的学习率OneCycleLR有 step 计数。只恢复模型和 optimizer、忘记恢复 scheduler训练能跑但学习率从初始值重新走一遍效果大概率是一团糟。2.3 随机数状态和训练进度元信息第三类是随机数状态这个最容易被忽略但对“严格恢复到中断点”非常重要。现代训练流程里到处都是随机性数据加载时的 shuffle、数据增强随机裁剪、翻转、颜色扰动、dropout 的 mask、初始化噪声、多进程 worker 的采样……这些随机性由 Python 的 random 模块、NumPy 的随机数生成器和 PyTorch 的 torch random 共同管理。每产生一组随机值全局状态就会前进一步。如果不保存这些状态恢复训练时会重新生成一遍随机序列相当于某一段训练用了两批不同的随机增强和 dropout mask。那是不是说必须恢复随机数状态这里要分情况如果数据加载用了 DataLoader 的多 worker且你关心“断点之后的数据顺序必须和不断点时完全一致”那你需要保存每个 worker 的随机状态否则做不到严格一致。如果数据增强强度不太大、batch size 比较大或者你每一轮epoch都重新 shuffle 一遍那恢复随机状态的意义就小很多因为下一轮的随机序列无论怎么走都行。我个人的习惯是在 checkpoint 里保存随机数状态但把是否恢复做成开关。默认开着如果发现因为 DataLoader 多进程导致恢复困难再关掉。保存方式也很简单四个模块的随机状态分别取出来存进 dict 就行import random import numpy as np import torch def capture_rng_state(): return { python: random.getstate(), numpy: np.random.get_state(), torch: torch.random.get_rng_state(), cuda: torch.cuda.get_rng_state_all() if torch.cuda.is_available() else None, }元信息更直接当前训练到第几个 epoch、第几个 step、当前 best metric、best checkpoint 的路径、以及你自定义的任何训练状态。这些信息在恢复时用来决定从哪个 epoch 继续、调度器从哪里接上、以及判断模型性能是否已经刷新到最优。说到恢复随机状态有一个实践上的难处DataLoader 的 worker seed 是在每个 worker 进程启动时重新初始化的恢复全部随机状态不一定能保证每个 worker 内部恢复一致。对于这类严格一致性的要求我后面在第 4 部分专门讲到一个替代方案这里先不展开。3. 实操过程从保存到恢复的完整链路3.1 保存 checkpoint 的推荐做法先给一个我自己在用的保存函数模板。它把前面讨论的内容全部落到了一个字典里def save_checkpoint( model, optimizer, scheduler, epoch, global_step, best_metric, config, rng_state, filepath, ): checkpoint { model_state_dict: model.state_dict(), optimizer_state_dict: optimizer.state_dict(), scheduler_state_dict: scheduler.state_dict() if scheduler else None, epoch: epoch, global_step: global_step, best_metric: best_metric, config: config, rng_state: rng_state, } torch.save(checkpoint, filepath)注意几个细节第一torch.save默认使用 Python 的 pickle 序列化。如果你只关心权重不关心完整的训练状态可以加_use_new_zipfile_serializationTrue参数它能生成 zip 格式的 checkpoint加载更快、文件展开更友好。但前提是你的代码环境兼容新格式一般现代 PyTorch 版本都没问题。第二保存前先检查一遍状态有没有落在正确的 device 上。很多人训练用 GPU保存时发现 state_dict 里的张量在 CUDA 上于是 checkpoint 文件巨大而且加载时还得小心翼翼。实际上 PyTorch 的torch.save保存的是张量的数据副本不要求先把模型挪到 CPU 上——它会正常序列化。但在加载时用map_locationcpu再统一搬到目标设备是最稳妥的套路。第三保存频率要有策略。每多少个 step 存一次我常用的策略是“每 N 个 epoch 存一个完整 checkpoint 每个 epoch 内若干个 checkpoint”具体频率看训练时长。如果一次训练要跑一周每 10 分钟一个 checkpoint 都不过分如果一次训练只要两小时每半小时存一次就够了。这里没有标准答案核心理念是checkpoint 文件丢失的代价 vs 频繁保存的磁盘和 IO 开销取一个平衡点。第四最好用“临时文件 原子替换”的方式写盘。正确做法是torch.save(checkpoint, filepath .tmp) os.replace(filepath .tmp, filepath)为什么这么做如果你的训练正好在torch.save写文件写到一半时被 kill —— 无论是 SIGKILL 还是断电 —— 主文件就会是一个截断的损坏文件下次加载直接报错。先写临时文件再原子替换能保证磁盘上的 checkpoint 文件要么是完整的旧版本要么是完整的新版本不会出现半截文件。这个小技巧我吃过亏之后就成了习惯。3.2 恢复训练的完整流程恢复训练的代码比保存稍微复杂一点因为要考虑 device 映射和加载容错。def load_checkpoint(filepath, model, optimizer, scheduler, device): checkpoint torch.load(filepath, map_locationcpu) model.load_state_dict(checkpoint[model_state_dict]) model.to(device) optimizer.load_state_dict(checkpoint[optimizer_state_dict]) # 加载 optimizer 后把参数组里的学习率等超参数同步到当前训练配置 # 例如optimizer.param_groups[0][lr] your_new_lr if scheduler is not None and checkpoint.get(scheduler_state_dict) is not None: scheduler.load_state_dict(checkpoint[scheduler_state_dict]) start_epoch checkpoint[epoch] 1 global_step checkpoint.get(global_step, 0) best_metric checkpoint.get(best_metric, float(inf)) # 可选恢复随机数状态 if checkpoint.get(rng_state) is not None: restore_rng_state(checkpoint[rng_state]) return start_epoch, global_step, best_metricmap_locationcpu是我强推的做法。哪怕你训练时用的是 GPU加载时也先全部放到 CPU然后统一model.to(device)。原因有两个如果你直接加载到原来的设备但当前进程的 GPU 编号和保存时不同比如保存时用的是 cuda:0现在机器上只有 cuda:1 可用会报设备不匹配错误。大模型如果直接加载进显存临时占用很大万一显存本来就很紧张加载那一下也可能 OOM。放 CPU 再转移到 GPU操作更安全。恢复之后训练主循环要跳过已经跑过的数据步数。epoch 从start_epoch开始循环内部用global_step恢复全局步数计数如果用的是 logging 和 checkpoint 保存的步数索引。还有一个细节如果用了分布式训练需要把distributed_sampler.set_epoch(start_epoch)设置正确否则每个 epoch 的数据 shuffle 会从第 0 个 epoch 的顺序重新开始这和恢复的“现场”就不一致了。随机数状态恢复的函数写起来也不复杂def restore_rng_state(state): random.setstate(state[python]) np.random.set_state(state[numpy]) torch.random.set_rng_state(state[torch]) if state[cuda] is not None and torch.cuda.is_available(): torch.cuda.set_rng_state_all(state[cuda])3.3 用“杀进程”实测验证恢复逻辑写完保存和恢复的代码你怎么知道自己做得对不对我强烈建议你一上来就做一个“杀进程中断实验”这种验证可以一次性把恢复逻辑的问题全都试出来。具体操作如下写一个训练脚本跑 20 个 step。每个 step 结束时检查当前 step 是否是 10如果是就保存一个 checkpoint然后手动 raise 一个中断。再写一个恢复脚本或者同一个脚本加 resume 参数从刚才保存的 checkpoint 加载继续跑完剩余 10 个 step。对比两种跑法——一种是完全不中断、从头跑 20 个 step另一种是中断后从 checkpoint 继续跑剩下的 10 个 step——两者第二次 10 个 step 的 loss 曲线、学习率变化以及最终模型权重是否一致。这里要说清楚一个预期如果随机数状态完整恢复且设备一致两次训练的结果理论上应该完全一致但由于浮点运算的并行顺序可能有微小差异实际可能有非常细微的差别。如果 loss 曲线走势明显不一样很大概率是你没有恢复随机数状态、或者 optimizer / scheduler 状态丢了。我的验证经验里有个经典翻车现场只恢复了模型和 optimizer没恢复 scheduler 和随机状态第一次跑 20 步和恢复后半段的 loss 曲线完全对不上。后来补上了 random state 和 scheduler 的恢复曲线基本吻合才确认了问题所在。3.4 文件命名、轮转与防写坏持续训练会产生大量 checkpoint。如果不做管理磁盘很快会被塞满而磁盘被塞满又会诱发新的中断——这就形成了一个恶性循环。我常用的文件管理方案是checkpoints/ ├── last.pt # 最近一次保存的完整 checkpoint ├── best.pt # 当前在验证集上指标最好的 checkpoint ├── epoch_0005_step_2000.pt ├── epoch_0010_step_4000.pt └── ...保存逻辑上每次保存都把last.pt覆盖更新。同时保留一个带 epoch/step 信息的版本文件便于回溯。如果新一轮验证指标超过历史 best再更新best.pt。用一个简单的保留策略最多保留最近 K 个版本文件K 取决于磁盘空间。比如只保留最近 5 个带 epoch 的 checkpoint超出就删除最老的。def prune_checkpoints(checkpoint_dir, keep_last5): files sorted(glob.glob(os.path.join(checkpoint_dir, epoch_*.pt))) for f in files[:-keep_last]: os.remove(f)这套方案的成本很低但价值巨大任何时候崩溃都有一个last.pt保证“最多丢失最后保存间隔内的进度”也有一个best.pt保证“只要保存过最好的结果后续就算训练跑费了也不亏”。这个习惯我一直保留尤其在做长周期训练并且没有专人盯着训练进程的时候它可以救命。4. 常见问题与排查技巧实录4.1 恢复后 loss 跳变是恢复出问题了吗很多人恢复训练后第一眼看到 loss 比中断前高了一点立刻怀疑恢复逻辑有 bug。其实这里有多个原因。最常见的原因是数据处理顺序改变。你在中断时保存了 epoch 和 step但 resume 之后脚本重新创建了 DataLoadershuffle 的顺序由新的随机种子或者新初始化的随机数状态决定如果随机状态没有恢复新一轮前面几步的 batch 就和原来那些 “本应看到” 的 batch 不一样了。这个差异本身有时会带来 loss 波动但通常幅度不大会稳定下来。另一个常见原因是恢复模型后没有把模型设置为训练模式。如果你在恢复流程里先做了验证或者其他 eval 操作模型带着 eval 状态回到训练循环BatchNorm 和 Dropout 行为完全不同loss 当然会异常。检查一下model.train()是否有被调用。如果 loss 跳变幅度很大而且持续不回落那才需要怀疑优化器或 scheduler 没恢复。排查方法很简单加载 checkpoint 后打印 optimizer 的 param_groups 里 lr 值以及 scheduler 当前的 last_epoch / best 值对比中断前保存时日志里的值。两边对不上说明恢复链路掉了一环。4.2 加载报 Missing/Unexpected key 怎么办这个报错几乎每个做迁移和 resume 的人都会遇到。Missing key(s)的意思是 checkpoint 里有某些 key 在当前模型中不存在或者当前模型有某些参数在 checkpoint 中不存在。出现这个错常见原因是模型结构代码改动过比如加了一层、删了一层、改了层名。checkpoint 是从带module.前缀的并行模型保存的加载到一个普通模型上或者反过来。你保存的时候是 part1 part2 的混合模型加载到新模型结构时名字对不上。排查的第一步不是去看网络代码而是先打印两边的 key 列表来做差异state torch.load(last.pt, map_locationcpu) ckpt_keys set(state[model_state_dict].keys()) model_keys set(model.state_dict().keys()) print(missing:, model_keys - ckpt_keys) print(unexpected:, ckpt_keys - model_keys)看清楚差异再动手。如果只是module.前缀问题写一个 remap 就好def remove_module_prefix(state_dict): from collections import OrderedDict new_state OrderedDict() for k, v in state_dict.items(): new_k k.replace(module., ) if k.startswith(module.) else k new_state[new_k] v return new_state如果确实是结构改了那要看差异 keys 是否对应你想剔除或新增的层。比如你想从 checkpoint 加载一个主干网络的权重到新的下游任务模型直接用strictFalse是可行的model.load_state_dict(checkpoint[model_state_dict], strictFalse)但强烈建议你在 strictFalse 之后手动验证那些被忽略的层是哪些。我遇到过一次很坑的情况用 strictFalse 加载完所有权重都提示 missing 了但因为某个命名空间错误实际加载进去的只有一层卷积其余全是随机初始化。训练出来效果极差排查了半天才靠核对 key 数量发现问题。所以记住strictFalse 只是方便迁移不代表权重一定加载正确。4.3 checkpoint 文件太大磁盘快爆了大模型的 checkpoint 动辄几个 GB。如果你每十分钟存一次磁盘压力非常真实。办法有这么几个只保存最近的 N 个版本用前面提到的轮转策略清理。判断是否真的需要每次都保存整个 optimizer 状态。如果你用的是 Adamoptimizer state 和模型参数差不多一样大exp_avg 和 exp_avg_sq 各一份checkpoint 体积直接翻三倍。折中方案是平时每 N 步保存一个“轻量 checkpoint”只含模型和元信息每 N 个 epoch 才保存一个“完整 checkpoint”。万一轻量 checkpoint 恢复的时候 optimizer 状态丢了训练还能继续只是需要一段预热时间。如果保存完整 checkpoint 的需求无法压缩考虑用torch.save的新 zip 格式减少体积同时不要用 float64 保存不需要的中间状态统一用 float32。文件大还有一个隐蔽问题有的文件系统对大文件单次写入不稳定。你需要确保写文件的进程有足够磁盘空间余量尤其临时文件 原子替换方案中临时文件占两份空间磁盘设计时要留出这个余量。4.4 多卡训练里断点续训的额外坑多卡训练尤其是分布式训练的 resume 比单卡复杂一档。基本模式是只有 rank 0 保存 checkpoint所有 rank 加载同一份 checkpoint。保存时要注意直接取模型的 state_dict 而不是带 DDP wrapper 的。如果使用了DistributedDataParallel加载时最快的方式是model DistributedDataParallel(model, device_ids[local_rank]) model.module.load_state_dict(checkpoint[model_state_dict])加载时每个进程都要做但要注意torch.load(map_locationcpu)然后每个进程各自to(device)。这样每个 rank 拿到的权重一样不会出现 device mismatch。DistributedSampler也要额外设置。它的默认行为是每个 epoch 用self.epoch作为随机种子。如果你不把sampler.set_epoch(start_epoch)调成恢复的 epoch那么恢复后第一个 epoch 的 shuffle 顺序其实还是第 0 个 epoch 的顺序数据采样完全错乱。还有一个多卡特有的问题如果中断后的恢复是在不同数量的卡上进行的比如原来 8 卡、现在只有 4 卡这时 batch size 变了。严格来说梯度累积、学习率策略都可能需要相应调整。你可以选择直接继续跑多数情况下也能收敛但曲线会有一点波动也可以选择重启一个小的热身阶段让优化器适应新的分布式状态。4.5 数据加载顺序变了怎么尽量保持一致这个问题前面提过现在展开说一下。如果你用的 DataLoader 开了多进程 workernum_workers 0并且每个 epoch 内也依赖确定的数据顺序那仅仅恢复random、numpy、torch的全局随机状态是不够的。每个 worker 进程在启动时会用自己的一套随机种子通常是主进程 seed worker id 的组合来初始化而主进程在for batch in dataloader的迭代过程中是通过一个迭代器不断从 worker 取数据的中断后重建 DataLoaderworker 内部的数据采样状态早已丢失。我在实践中面对这个问题的解法是这样的如果确认数据增强对最终结果影响不大比如简单的分类任务用了固定 resiz 和 normalize没带随机裁剪就直接不追求数据顺序一致。每次 resume 后从当前 epoch 继续但 batch 顺序可能变化这完全没关系模型很快就能适应。如果是有强随机增强的图像任务、或者是需要在同一 epoch 内精确恢复的对比学习类任务我会改用“索引式恢复”定义一个 DataLoader 包装器自己控制每次迭代的样本索引把当前 index 也存进 checkpoint。恢复时从保存的 index 继续而不是从 epoch 开头。console 上做这种控制更麻烦但是可控性最好。有人可能会问既然随机状态恢复这么麻烦那为什么不干脆每次中断都从 epoch 开头重新跑我的看法是当你想做严格的消融实验、想对比不同随机种子下的结果时数据顺序的一致性真的很重要但大多数生产场景里模型收敛才是硬指标batch 顺序有一点变化完全可以接受。所以这个坑反而可以不用过度纠结把精力花在更值得的地方。5. 几个我长期在用的续训习惯5.1 把“可恢复性”设计到训练脚本最前面不要等到训练跑起来了才想断点续训的事。我现在的训练脚本统一在命令行参数里加一个--resume参数默认从last.pt恢复。脚本启动的时候做三件事先解析命令行、再检查 checkpoint 是否存在、存在就加载并覆盖训练配置。没有一个模块是后加的全部是训练脚本的组成部分。这样训练任务的启动命令就非常简单无论是手动跑还是配合调度系统都能拿到一样的语义python train.py --config config.yaml --resume checkpoints/last.pt5.2 每次保存都顺手打印状态摘要训练中断后恢复最怕的就是不知道当前 checkpoint 对应训练到哪个阶段。所以我每次 save_checkpoint 之后都会往训练日志里打印一行[checkpoint] saved to checkpoints/last.pt | epoch 7 | step 2860 | loss 1.832 | lr 3.2e-4虽然看起来是个不起眼的习惯但它有实际价值恢复脚本启动时我会把 checkpoint 里的 epoch 和 step 打印出来和日志里的这行对比能快速发现“到底恢复了哪个版本”的问题。有一次我朋友在调试一个项目发现恢复后训练挂了一整天才意识到加载的是半个月前的旧 checkpoint就是因为日志里没有这种摘要信息排查非常痛苦。最后再分享一个小习惯checkpoint 保存目录和日志目录尽量分开并且定期检查磁盘剩余空间。我做训练时会把 checkpoint 挂到存储较大的磁盘分区上因为训练中断导致进度丢失很恼火但更恼火的是磁盘满了导致 checkpoint 保存不了结果连最后的“避难所”都没了。先把现场完整保存下来后面出错都好说。这套经验基本能覆盖绝大多数断点续训的开发和使用场景希望你能在下次训练中断时少一点手忙脚乱。