连着两个月的训练计划前阵子栽在了一个看起来特别小的环节上凌晨两点服务器因为内存占用过载被系统自动重启我跑了一周多的扩散模型训练进程直接没了。重新把代码拉起来一看好家伙从第一个epoch从头开始。那一周时间和上千元的GPU算力就此清零。痛定思痛之后我把断点续训这件事当成一个正经的工程来做——前四篇讲了数据、模型结构和训练策略这第五篇就专门聊模型训练中断了怎么让它记住上次干到哪。不卖关子了核心答案就四个字checkpoint。但怎么存、存哪些东西、怎么恢复细节比大多数人想的要多得多今天这篇就把我踩过的坑和最终方案一并写出来。1. 断点续训是什么从一次崩溃聊起1.1 训练中断的典型场景在讲方案之前先把“敌人”看清楚。我说的“模型”不限于大语言模型扩散模型、图像分类模型、目标检测模型都一样只要训练耗时超过几个小时断点续训就是刚需。常见的训练中断原因我按出现频率排了个序服务器自动重启或远程连接断开。云服务器经常在凌晨做维护磁盘、驱动一旦出问题进程直接被系统杀掉。显存溢出也就是OOM。训练过程中偶尔出现一个高维batch或者显存泄露越积越多内存尖峰一上来整个程序直接崩掉。网络抖动导致计算集群把当前节点摘除。多机训练时节点之间的心跳断了调度器默认这个节点失联任务被回收。宿主机管理员做资源回收或者共享GPU机器上别人把我的进程挤掉。磁盘空间被日志、历史模型占满日志进程先写不进去随后训练进程报错退出。这五个场景我全部亲身遇到过。一个容易被忽略的事实是真正因为代码bug导致的中断反而占比不高绝大多数中断发生在外部的、你无法控制的环节上。如果没有主动做防护系统不会给你任何后悔的机会甚至中断发生的当下你还在睡觉等第二天醒来才看到空荡荡的终端。1.2 不保存checkpoint的真实代价算一笔细账假设一个模型训练任务需要7天你已经跑到第4天这时候断了。如果没有checkpoint意味着你要重启完整训练。表面上看只是浪费4天但如果把后续影响摊开代价远不止于此。第一是机器费用。按一张GPU每小时几块到十几块算一周连续训练就是一笔不小的开销多卡并行之后费用直接翻倍。训练中断一次原本计划的实验预算可能就超了。第二是时间成本。一个任务反复重启调参窗口被拉长好几周整个项目排期全线漂移。我们常说“早一天出模型早一天迭代”中断很可能让你错过关键交付节点。第三是心智负担。上次训练跑到第几步、loss曲线长什么样、学习率调度到什么阶段这些上下文全部丢失之后你连怎么判断这次改动的效果都很难。我一直认为训练到一半因为外部因素“归零”挫败感远超代码bug。代码bug起码是你自己写出来的可以快速定位而外部中断往往神出鬼没只能靠机制兜底。所以我们不做“事后补救派”要做“事前工程派”断点续训必须在训练脚本的第一版内置进去而不是等到快训完才想着加两行保存代码。2. 到底要把哪些东西塞进checkpoint2.1 模型权重最基础但不完整很多新手第一次做保存时代码长这样torch.save(model.state_dict(), model_final.pth)这个写法本身没有错但只适用于“最终模型落盘之后拿去推理”的场景。如果拿去断点续训马上会碰到两个问题。第一个问题是优化器成了“陌生人”。网络参数确实原样载回来了但Adam等优化器内部维护的动量、自适应学习率等状态全是空的。恢复后的第一个step优化器会认为网络参数是全新初始化的更新幅度完全不在正常轨道上轻则loss波动重则直接训崩。第二个问题是学习率调度器的进度丢了。如果用的是cosine衰减或者带warmup的调度器恢复时调度器以为训练才刚开始会把学习率拉回初始高位已经学到中后段的模型突然被灌了一大口“猛药”表现自然是过山车。所以这里送给大家第一句经验保存训练中期的checkpoint绝对不只是模型权重。整个优化器状态、调度器状态、训练进度一个都不能少。2.2 优化器状态与调度器真正决定“连续性”优化器状态到底是什么我举个生活化的例子假设你在跑马拉松每5公里需要补水。模型权重是你的跑步速度优化器状态是体内水分和糖原的存量。只保存跑步速度、不保存体内存量下一段路你根本跑不动。具体到代码层面优化器里除了param_groups这些配置信息还会为每个参数维护momentum、variance等内部缓冲。尤其是Adam这类自适应优化器这些缓冲直接决定下一步权重往哪儿走、走多远。PyTorch里最省事的做法是连优化器对象一起整包存取torch.save({ model: model.state_dict(), optimizer: optimizer.state_dict(), scheduler: scheduler.state_dict(), epoch: epoch, step: global_step, best_loss: best_loss, rng_state: torch.get_rng_state(), }, fcheckpoints/checkpoint_epoch{epoch}.pth)恢复时也整包恢复。假设你在第20个epoch中断恢复后调度器却认为还在第0个epoch接下来所有学习率都会错位就相当于你前面20个epoch的学习率规划全白做了。2.3 训练状态epoch、step、随机种子与loss记录真正完整的checkpoint除了模型、优化器、调度器还应该包含下面这些训练状态我用一个表格说明它们各自的价值保存字段作用不保存会发生什么model.state_dict()网络参数模型内容对不上optimizer.state_dict()动量/方差等内部状态恢复后第一步更新异常scheduler.state_dict()学习率进度warmup和衰减从头再来epoch / global_step训练进度无法计算剩余轮数rng_state随机数状态数据shuffle顺序被打乱best_loss历史最佳损失无法继续做模型存档与早停平时大家容易忽略的是随机数状态和best_loss这两个字段。随机数状态不只包括PyTorch主进程的RNG状态如果DataLoader开了多进程worker还要考虑每个worker的采样顺序。不恢复随机状态最直接的后果是模型参数确实接上了但数据加载的shuffle顺序从头再来严格意义上不叫“无缝续训”。best_loss则是续跑后做early stopping的重要依据你不存它恢复之后没法判断新epoch到底有没有在进步。我把完整checkpoint比作游戏存档存档不只保存角色血量还得保存当前地图进度、背包物品和系统时间。训练也是如此上下文越全恢复越真实。3. 保存与恢复的完整实操3.1 一份可复用的checkpoint保存模板直接贴我目前项目里在用的模板结构按照“主训练循环每N步调用一次”来设计import os import torch def save_checkpoint(network, optimizer, scheduler, epoch, step, loss, tagmodel, save_dir./checkpoints, keep_last3): os.makedirs(save_dir, exist_okTrue) file_path os.path.join(save_dir, f{tag}_epoch{epoch:03d}_step{step:06d}.pth) tmp_path file_path .tmp state { network: network.state_dict(), optimizer: optimizer.state_dict(), scheduler: scheduler.state_dict() if scheduler else None, epoch: epoch, step: step, best_loss: loss, rng_state: torch.get_rng_state(), cuda_rng_state: torch.cuda.get_rng_state_all() if torch.cuda.is_available() else None, } # 先写临时文件再改名避免保存到一半进程崩溃留下半个文件 torch.save(state, tmp_path) os.replace(tmp_path, file_path) # 保留最近keep_last个防止磁盘被堆满 files sorted([f for f in os.listdir(save_dir) if f.startswith(tag)]) for old_file in files[:-keep_last]: os.remove(os.path.join(save_dir, old_file)) print(f[checkpoint] saved to {file_path})这个函数里有两个细节容易被忽略。第一文件名带上epoch和step方便一眼看出训练进度第二定期清理旧存档。一个checkpoint动辄几个GB不清理的话半个月就能把磁盘塞满而磁盘满了又会训练中断形成恶性循环。另外先写.tmp临时文件再os.replace改名这个习惯救过我一次当时保存过程中进程被系统杀掉普通写法会留下一个写到一半的损坏文件而临时文件转正能保证checkpoint文件总是完整的。3.2 恢复训练从checkpoint说起恢复代码也不复杂但有一个坑很多时候我们不知道checkpoint文件里到底存了哪些键所以恢复时要先load出来再按既有结构取值。下面是我一直在用的模板def load_checkpoint(file_path, network, optimizerNone, schedulerNone): if not os.path.exists(file_path): print([checkpoint] no checkpoint found, start fresh) return 0, None state torch.load(file_path, map_locationcpu) network.load_state_dict(state[network]) if optimizer is not None and optimizer in state: optimizer.load_state_dict(state[optimizer]) if scheduler is not None and scheduler in state: scheduler.load_state_dict(state[scheduler]) epoch state.get(epoch, 0) step state.get(step, 0) best_loss state.get(best_loss, float(inf)) print(f[checkpoint] loaded from epoch {epoch}, step {step}, best_loss{best_loss:.4f}) return epoch, step, best_loss恢复之后主循环要从start_epoch 1开始跑注意别用原epoch直接覆盖控制变量DataLoader需要重新构造但如果是分布式训练数据采样的shuffle状态要和checkpoint里的step对齐。实际操作中我一般把“决定shuffle的种子”也单独作为一个字段保存。经验是想完全复现一个训练过程靠的是“代码版本 数据版本 随机种子 checkpoint”这四件套缺一样都可能复现不出来。3.3 自动化周期保存与备份策略手动保存太依赖记忆我的习惯是写进训练循环里每500或1000步自动保存一次EVAL_INTERVAL 500 SAVE_INTERVAL 1000 start_epoch 0 global_step 0 for epoch in range(start_epoch, total_epochs): for batch_idx, batch in enumerate(train_loader): loss train_one_step(batch) global_step 1 if global_step % SAVE_INTERVAL 0: save_checkpoint(model, optimizer, scheduler, epoch, global_step, loss)除了按step保存还有两个策略值得加进来。第一个是“模型提升时额外保存best版”也就是if loss best_loss: best_loss loss; save(...)这样随时保留一个性能最高的存档供推理使用。第二个是“里程碑节点手动双备份”比如训练到20%、50%、80%这种节点我会把checkpoint复制到另一个磁盘或者远端存储防止本机磁盘损坏后全军覆没。训练后期突然发现磁盘挂了这种倒霉事一旦碰上多一份备份是真的能救命。4. 进阶技巧大模型、多卡与框架的resume玩法4.1 超大模型的checkpoint思路当模型大到单个文件保存都吃力时普通torch.save就不够看了。我见过不少团队踩过这个大文件的坑保存几个GB的state_dict看起来还好但当模型到几十GB甚至上百GB时单文件保存和加载都极度耗时还容易触发内存瓶颈。可以换几个思路。第一是分片保存把state_dict按层名或者按shard索引拆成多个小文件恢复时先加载一个配置索引再按需加载分片。这样可以避开单个大文件的序列化限制也方便并行读取。第二是使用更高效、专门为模型权重设计的序列化格式加载速度更快也更节约磁盘。第三是直接依赖训练框架的resume功能很多框架已经把这些底层逻辑封装好了比如通过resume_from_checkpoint参数指定上一次的checkpoint目录它会自动处理分片和优化器状态的匹配。分片保存还有一个额外好处如果我改动模型结构后只想恢复部分层可以直接按层名单独加载不必整个文件打包排查问题会方便很多。4.2 多卡并行的恢复细节多卡训练比如DataParallel、DistributedDataParallel时checkpoint里要额外注意优化器状态是以哪个设备为基准保存的。常见的坑是保存时用了map_locationcuda:0恢复时机器换卡或者卡号变动加载直接报device mismatch。稳妥做法是恢复时统一先load到CPU再to(device)到当前进程对应的设备上state torch.load(file_path, map_locationcpu) model.load_state_dict(state[network]) optimizer.load_state_dict(state[optimizer]) model model.to(device) optimizer optimizer.to(device)多卡场景还容易忽略rank和world_size这两个概念。DDP训练时每个进程的rank不一样data sampler会依据epoch做分片恢复后一定要把sampler的epoch同步到start_epoch否则数据会被重复取或者漏取。换句话说多卡的断点续训不只是“把权重导回去”还包括让每张卡恢复到同样的数据read指针位置。4.3 主流训练框架自带resume参数现在很多开源训练框架已经把断电续训做成了开箱功能不见得非要自己实现全套。比如常见的transformers类训练器、diffusers训练脚本它们的命令行都带resume_from_checkpoint参数还有一些加速工具会单独提供分布式状态保存和加载函数专门处理多卡下的checkpoint。这些框架级的方案通常比手写方案更可靠因为它们会连带处理优化器、调度器、RNG状态、sampler状态等一系列细节。我的建议是先看懂框架的resume机制再决定要不要自己写。框架支持完善就依赖框架边缘场景多、自定义逻辑重就用前面第3节那份模板自己兜底。很多时候两种可以同时用框架负责模型和优化器部分你自己负责随机种子、日志历史这类框架覆盖不到的字段。5. 常见问题与避坑实录5.1 恢复后loss异常升高的排查“恢复模型后loss比中断前高了一大截是checkpoint坏了吗”这个问题我至少被问过十次自己也踩过同样的坑。排查顺序一般是这样确认模型是否真的加载了。打印几层参数的和或者对比训练前后权重的hash能快速排除“没加载上”的可能。确认优化器是否加载了。光加载model不加载optimizer第一个step的更新幅度会异常大loss肯定跳。确认学习率调度器是否加载了。如果没加载warmup会重来往往表现为前几个step的曲线异常高。确认数据是否从正确的step续上。如果数据shuffle种子被重置哪怕模型参数是接续的曲线也会有一小段波动。确认batch大小、分辨率、混合精度设置有没有变。训练中途切换精度会改变数值稳定性也是隐藏的背锅位。大多数情况不是checkpoint坏了而是我们保存的信息不完整。这也是为什么我一直强调保存模型时要把优化器、调度器、随机种子一起打包而不是只存一份model.state_dict()。5.2 state_dict不匹配的玄学恢复时比较常见的一个报错是“键不匹配”也就是Missing key(s)或者Unexpected key(s)。多数原因是模型结构代码变了比如把某层从3层改成4层或者改了一个层的命名。遇到这种情况别硬来先看两边结构差异能改代码就改代码能改加载方式就改加载方式。一个很有用的兼容加载技巧是只加载权重中的交集部分其他层保持随机初始化model_dict model.state_dict() pretrained_dict {k: v for k, v in state[network].items() if k in model_dict and model_dict[k].shape v.shape} model_dict.update(pretrained_dict) model.load_state_dict(model_dict)这种写法适合在原有checkpoint基础上新增或删除少量层的场景。如果结构差异太大直接保留旧配置重新生成一个中间版本再继续训练别硬凑。我见过有人强行用strictFalse加载结果模型前面几层新结构参数全是随机值后面旧结构参数保留调和训练出来的结果非常诡异。5.3 我的checkpoint目录组织习惯最后分享一个我用了很久的目录组织方式谈不上多先进但省了很多事runs/exp001/ ├── checkpoints/ │ ├── exp001_epoch003_step12000.pth │ ├── exp001_epoch006_step25000.pth │ ├── exp001_best.pth │ └── latest.pth # 永远指向最近一个存档 ├── logs/ │ ├── metrics.jsonl │ └── tensorboard/ └── config.yaml # 保存训练超参恢复时一起读其中latest.pth这个设计很关键它解决了一个实际痛点恢复脚本不需要手动输入一长串带epoch和step的文件名。训练中断后我只要写成load_checkpoint(runs/exp001/checkpoints/latest.pth, ...)就能自动接上。为了不让这个文件在每次保存时都复制一份大文件我通常直接用软链接指向最新实际存档脚本里再判断一下os.path.islink如果存在软链就readlink拿到目标路径。如果你在多台机器间同步训练目录千万别让骨干网络高频同步所有历史checkpoint几千个进程同时读同一份大文件会把网络打满同步latest和best两个存档就足够了。我在实际项目里还发现真正影响断点续训可靠性的往往不是代码逻辑而是你根本不知道什么时候会断。所以我这两年养成的好习惯是除了把保存逻辑写进脚本还坚持隔两小时瞄一眼训练机器的资源占用看到磁盘空间低于阈值就立刻清理每个里程碑存档一定复制一份到远端备份。有人觉得这些行为很琐碎但吃过凌晨三点因服务器重启而丢进度的亏之后每次多按一次备份键心里都踏实。存好档再关电脑才敢让训练任务自己跑一整夜。希望这篇第五部分能让你少进一次凌晨重新起训练的坑。 
长河君 · 2026/10/19 13:26:40 · 阅读 1.2k 深度学习断点续训全指南：checkpoint保存与恢复的工程实践
连着两个月的训练计划前阵子栽在了一个看起来特别小的环节上凌晨两点服务器因为内存占用过载被系统自动重启我跑了一周多的扩散模型训练进程直接没了。重新把代码拉起来一看好家伙从第一个epoch从头开始。那一周时间和上千元的GPU算力就此清零。痛定思痛之后我把断点续训这件事当成一个正经的工程来做——前四篇讲了数据、模型结构和训练策略这第五篇就专门聊模型训练中断了怎么让它记住上次干到哪。不卖关子了核心答案就四个字checkpoint。但怎么存、存哪些东西、怎么恢复细节比大多数人想的要多得多今天这篇就把我踩过的坑和最终方案一并写出来。1. 断点续训是什么从一次崩溃聊起1.1 训练中断的典型场景在讲方案之前先把“敌人”看清楚。我说的“模型”不限于大语言模型扩散模型、图像分类模型、目标检测模型都一样只要训练耗时超过几个小时断点续训就是刚需。常见的训练中断原因我按出现频率排了个序服务器自动重启或远程连接断开。云服务器经常在凌晨做维护磁盘、驱动一旦出问题进程直接被系统杀掉。显存溢出也就是OOM。训练过程中偶尔出现一个高维batch或者显存泄露越积越多内存尖峰一上来整个程序直接崩掉。网络抖动导致计算集群把当前节点摘除。多机训练时节点之间的心跳断了调度器默认这个节点失联任务被回收。宿主机管理员做资源回收或者共享GPU机器上别人把我的进程挤掉。磁盘空间被日志、历史模型占满日志进程先写不进去随后训练进程报错退出。这五个场景我全部亲身遇到过。一个容易被忽略的事实是真正因为代码bug导致的中断反而占比不高绝大多数中断发生在外部的、你无法控制的环节上。如果没有主动做防护系统不会给你任何后悔的机会甚至中断发生的当下你还在睡觉等第二天醒来才看到空荡荡的终端。1.2 不保存checkpoint的真实代价算一笔细账假设一个模型训练任务需要7天你已经跑到第4天这时候断了。如果没有checkpoint意味着你要重启完整训练。表面上看只是浪费4天但如果把后续影响摊开代价远不止于此。第一是机器费用。按一张GPU每小时几块到十几块算一周连续训练就是一笔不小的开销多卡并行之后费用直接翻倍。训练中断一次原本计划的实验预算可能就超了。第二是时间成本。一个任务反复重启调参窗口被拉长好几周整个项目排期全线漂移。我们常说“早一天出模型早一天迭代”中断很可能让你错过关键交付节点。第三是心智负担。上次训练跑到第几步、loss曲线长什么样、学习率调度到什么阶段这些上下文全部丢失之后你连怎么判断这次改动的效果都很难。我一直认为训练到一半因为外部因素“归零”挫败感远超代码bug。代码bug起码是你自己写出来的可以快速定位而外部中断往往神出鬼没只能靠机制兜底。所以我们不做“事后补救派”要做“事前工程派”断点续训必须在训练脚本的第一版内置进去而不是等到快训完才想着加两行保存代码。2. 到底要把哪些东西塞进checkpoint2.1 模型权重最基础但不完整很多新手第一次做保存时代码长这样torch.save(model.state_dict(), model_final.pth)这个写法本身没有错但只适用于“最终模型落盘之后拿去推理”的场景。如果拿去断点续训马上会碰到两个问题。第一个问题是优化器成了“陌生人”。网络参数确实原样载回来了但Adam等优化器内部维护的动量、自适应学习率等状态全是空的。恢复后的第一个step优化器会认为网络参数是全新初始化的更新幅度完全不在正常轨道上轻则loss波动重则直接训崩。第二个问题是学习率调度器的进度丢了。如果用的是cosine衰减或者带warmup的调度器恢复时调度器以为训练才刚开始会把学习率拉回初始高位已经学到中后段的模型突然被灌了一大口“猛药”表现自然是过山车。所以这里送给大家第一句经验保存训练中期的checkpoint绝对不只是模型权重。整个优化器状态、调度器状态、训练进度一个都不能少。2.2 优化器状态与调度器真正决定“连续性”优化器状态到底是什么我举个生活化的例子假设你在跑马拉松每5公里需要补水。模型权重是你的跑步速度优化器状态是体内水分和糖原的存量。只保存跑步速度、不保存体内存量下一段路你根本跑不动。具体到代码层面优化器里除了param_groups这些配置信息还会为每个参数维护momentum、variance等内部缓冲。尤其是Adam这类自适应优化器这些缓冲直接决定下一步权重往哪儿走、走多远。PyTorch里最省事的做法是连优化器对象一起整包存取torch.save({ model: model.state_dict(), optimizer: optimizer.state_dict(), scheduler: scheduler.state_dict(), epoch: epoch, step: global_step, best_loss: best_loss, rng_state: torch.get_rng_state(), }, fcheckpoints/checkpoint_epoch{epoch}.pth)恢复时也整包恢复。假设你在第20个epoch中断恢复后调度器却认为还在第0个epoch接下来所有学习率都会错位就相当于你前面20个epoch的学习率规划全白做了。2.3 训练状态epoch、step、随机种子与loss记录真正完整的checkpoint除了模型、优化器、调度器还应该包含下面这些训练状态我用一个表格说明它们各自的价值保存字段作用不保存会发生什么model.state_dict()网络参数模型内容对不上optimizer.state_dict()动量/方差等内部状态恢复后第一步更新异常scheduler.state_dict()学习率进度warmup和衰减从头再来epoch / global_step训练进度无法计算剩余轮数rng_state随机数状态数据shuffle顺序被打乱best_loss历史最佳损失无法继续做模型存档与早停平时大家容易忽略的是随机数状态和best_loss这两个字段。随机数状态不只包括PyTorch主进程的RNG状态如果DataLoader开了多进程worker还要考虑每个worker的采样顺序。不恢复随机状态最直接的后果是模型参数确实接上了但数据加载的shuffle顺序从头再来严格意义上不叫“无缝续训”。best_loss则是续跑后做early stopping的重要依据你不存它恢复之后没法判断新epoch到底有没有在进步。我把完整checkpoint比作游戏存档存档不只保存角色血量还得保存当前地图进度、背包物品和系统时间。训练也是如此上下文越全恢复越真实。3. 保存与恢复的完整实操3.1 一份可复用的checkpoint保存模板直接贴我目前项目里在用的模板结构按照“主训练循环每N步调用一次”来设计import os import torch def save_checkpoint(network, optimizer, scheduler, epoch, step, loss, tagmodel, save_dir./checkpoints, keep_last3): os.makedirs(save_dir, exist_okTrue) file_path os.path.join(save_dir, f{tag}_epoch{epoch:03d}_step{step:06d}.pth) tmp_path file_path .tmp state { network: network.state_dict(), optimizer: optimizer.state_dict(), scheduler: scheduler.state_dict() if scheduler else None, epoch: epoch, step: step, best_loss: loss, rng_state: torch.get_rng_state(), cuda_rng_state: torch.cuda.get_rng_state_all() if torch.cuda.is_available() else None, } # 先写临时文件再改名避免保存到一半进程崩溃留下半个文件 torch.save(state, tmp_path) os.replace(tmp_path, file_path) # 保留最近keep_last个防止磁盘被堆满 files sorted([f for f in os.listdir(save_dir) if f.startswith(tag)]) for old_file in files[:-keep_last]: os.remove(os.path.join(save_dir, old_file)) print(f[checkpoint] saved to {file_path})这个函数里有两个细节容易被忽略。第一文件名带上epoch和step方便一眼看出训练进度第二定期清理旧存档。一个checkpoint动辄几个GB不清理的话半个月就能把磁盘塞满而磁盘满了又会训练中断形成恶性循环。另外先写.tmp临时文件再os.replace改名这个习惯救过我一次当时保存过程中进程被系统杀掉普通写法会留下一个写到一半的损坏文件而临时文件转正能保证checkpoint文件总是完整的。3.2 恢复训练从checkpoint说起恢复代码也不复杂但有一个坑很多时候我们不知道checkpoint文件里到底存了哪些键所以恢复时要先load出来再按既有结构取值。下面是我一直在用的模板def load_checkpoint(file_path, network, optimizerNone, schedulerNone): if not os.path.exists(file_path): print([checkpoint] no checkpoint found, start fresh) return 0, None state torch.load(file_path, map_locationcpu) network.load_state_dict(state[network]) if optimizer is not None and optimizer in state: optimizer.load_state_dict(state[optimizer]) if scheduler is not None and scheduler in state: scheduler.load_state_dict(state[scheduler]) epoch state.get(epoch, 0) step state.get(step, 0) best_loss state.get(best_loss, float(inf)) print(f[checkpoint] loaded from epoch {epoch}, step {step}, best_loss{best_loss:.4f}) return epoch, step, best_loss恢复之后主循环要从start_epoch 1开始跑注意别用原epoch直接覆盖控制变量DataLoader需要重新构造但如果是分布式训练数据采样的shuffle状态要和checkpoint里的step对齐。实际操作中我一般把“决定shuffle的种子”也单独作为一个字段保存。经验是想完全复现一个训练过程靠的是“代码版本 数据版本 随机种子 checkpoint”这四件套缺一样都可能复现不出来。3.3 自动化周期保存与备份策略手动保存太依赖记忆我的习惯是写进训练循环里每500或1000步自动保存一次EVAL_INTERVAL 500 SAVE_INTERVAL 1000 start_epoch 0 global_step 0 for epoch in range(start_epoch, total_epochs): for batch_idx, batch in enumerate(train_loader): loss train_one_step(batch) global_step 1 if global_step % SAVE_INTERVAL 0: save_checkpoint(model, optimizer, scheduler, epoch, global_step, loss)除了按step保存还有两个策略值得加进来。第一个是“模型提升时额外保存best版”也就是if loss best_loss: best_loss loss; save(...)这样随时保留一个性能最高的存档供推理使用。第二个是“里程碑节点手动双备份”比如训练到20%、50%、80%这种节点我会把checkpoint复制到另一个磁盘或者远端存储防止本机磁盘损坏后全军覆没。训练后期突然发现磁盘挂了这种倒霉事一旦碰上多一份备份是真的能救命。4. 进阶技巧大模型、多卡与框架的resume玩法4.1 超大模型的checkpoint思路当模型大到单个文件保存都吃力时普通torch.save就不够看了。我见过不少团队踩过这个大文件的坑保存几个GB的state_dict看起来还好但当模型到几十GB甚至上百GB时单文件保存和加载都极度耗时还容易触发内存瓶颈。可以换几个思路。第一是分片保存把state_dict按层名或者按shard索引拆成多个小文件恢复时先加载一个配置索引再按需加载分片。这样可以避开单个大文件的序列化限制也方便并行读取。第二是使用更高效、专门为模型权重设计的序列化格式加载速度更快也更节约磁盘。第三是直接依赖训练框架的resume功能很多框架已经把这些底层逻辑封装好了比如通过resume_from_checkpoint参数指定上一次的checkpoint目录它会自动处理分片和优化器状态的匹配。分片保存还有一个额外好处如果我改动模型结构后只想恢复部分层可以直接按层名单独加载不必整个文件打包排查问题会方便很多。4.2 多卡并行的恢复细节多卡训练比如DataParallel、DistributedDataParallel时checkpoint里要额外注意优化器状态是以哪个设备为基准保存的。常见的坑是保存时用了map_locationcuda:0恢复时机器换卡或者卡号变动加载直接报device mismatch。稳妥做法是恢复时统一先load到CPU再to(device)到当前进程对应的设备上state torch.load(file_path, map_locationcpu) model.load_state_dict(state[network]) optimizer.load_state_dict(state[optimizer]) model model.to(device) optimizer optimizer.to(device)多卡场景还容易忽略rank和world_size这两个概念。DDP训练时每个进程的rank不一样data sampler会依据epoch做分片恢复后一定要把sampler的epoch同步到start_epoch否则数据会被重复取或者漏取。换句话说多卡的断点续训不只是“把权重导回去”还包括让每张卡恢复到同样的数据read指针位置。4.3 主流训练框架自带resume参数现在很多开源训练框架已经把断电续训做成了开箱功能不见得非要自己实现全套。比如常见的transformers类训练器、diffusers训练脚本它们的命令行都带resume_from_checkpoint参数还有一些加速工具会单独提供分布式状态保存和加载函数专门处理多卡下的checkpoint。这些框架级的方案通常比手写方案更可靠因为它们会连带处理优化器、调度器、RNG状态、sampler状态等一系列细节。我的建议是先看懂框架的resume机制再决定要不要自己写。框架支持完善就依赖框架边缘场景多、自定义逻辑重就用前面第3节那份模板自己兜底。很多时候两种可以同时用框架负责模型和优化器部分你自己负责随机种子、日志历史这类框架覆盖不到的字段。5. 常见问题与避坑实录5.1 恢复后loss异常升高的排查“恢复模型后loss比中断前高了一大截是checkpoint坏了吗”这个问题我至少被问过十次自己也踩过同样的坑。排查顺序一般是这样确认模型是否真的加载了。打印几层参数的和或者对比训练前后权重的hash能快速排除“没加载上”的可能。确认优化器是否加载了。光加载model不加载optimizer第一个step的更新幅度会异常大loss肯定跳。确认学习率调度器是否加载了。如果没加载warmup会重来往往表现为前几个step的曲线异常高。确认数据是否从正确的step续上。如果数据shuffle种子被重置哪怕模型参数是接续的曲线也会有一小段波动。确认batch大小、分辨率、混合精度设置有没有变。训练中途切换精度会改变数值稳定性也是隐藏的背锅位。大多数情况不是checkpoint坏了而是我们保存的信息不完整。这也是为什么我一直强调保存模型时要把优化器、调度器、随机种子一起打包而不是只存一份model.state_dict()。5.2 state_dict不匹配的玄学恢复时比较常见的一个报错是“键不匹配”也就是Missing key(s)或者Unexpected key(s)。多数原因是模型结构代码变了比如把某层从3层改成4层或者改了一个层的命名。遇到这种情况别硬来先看两边结构差异能改代码就改代码能改加载方式就改加载方式。一个很有用的兼容加载技巧是只加载权重中的交集部分其他层保持随机初始化model_dict model.state_dict() pretrained_dict {k: v for k, v in state[network].items() if k in model_dict and model_dict[k].shape v.shape} model_dict.update(pretrained_dict) model.load_state_dict(model_dict)这种写法适合在原有checkpoint基础上新增或删除少量层的场景。如果结构差异太大直接保留旧配置重新生成一个中间版本再继续训练别硬凑。我见过有人强行用strictFalse加载结果模型前面几层新结构参数全是随机值后面旧结构参数保留调和训练出来的结果非常诡异。5.3 我的checkpoint目录组织习惯最后分享一个我用了很久的目录组织方式谈不上多先进但省了很多事runs/exp001/ ├── checkpoints/ │ ├── exp001_epoch003_step12000.pth │ ├── exp001_epoch006_step25000.pth │ ├── exp001_best.pth │ └── latest.pth # 永远指向最近一个存档 ├── logs/ │ ├── metrics.jsonl │ └── tensorboard/ └── config.yaml # 保存训练超参恢复时一起读其中latest.pth这个设计很关键它解决了一个实际痛点恢复脚本不需要手动输入一长串带epoch和step的文件名。训练中断后我只要写成load_checkpoint(runs/exp001/checkpoints/latest.pth, ...)就能自动接上。为了不让这个文件在每次保存时都复制一份大文件我通常直接用软链接指向最新实际存档脚本里再判断一下os.path.islink如果存在软链就readlink拿到目标路径。如果你在多台机器间同步训练目录千万别让骨干网络高频同步所有历史checkpoint几千个进程同时读同一份大文件会把网络打满同步latest和best两个存档就足够了。我在实际项目里还发现真正影响断点续训可靠性的往往不是代码逻辑而是你根本不知道什么时候会断。所以我这两年养成的好习惯是除了把保存逻辑写进脚本还坚持隔两小时瞄一眼训练机器的资源占用看到磁盘空间低于阈值就立刻清理每个里程碑存档一定复制一份到远端备份。有人觉得这些行为很琐碎但吃过凌晨三点因服务器重启而丢进度的亏之后每次多按一次备份键心里都踏实。存好档再关电脑才敢让训练任务自己跑一整夜。希望这篇第五部分能让你少进一次凌晨重新起训练的坑。