这个系列写到第五部分我想聊聊一个很多人都不太愿意面对的时刻模型训练中断。如果你自己跑过那种一跑就是几天的训练任务一定体会过“根本不敢离开”的心情。我早年的一次经历印象特别深某次训练到第三天机房通知要维护我在心里安慰自己“应该赶得完”结果维护提前了两小时实验直接报废。后面又遇到过云主机被回收、同事误杀进程每一回都是几个日夜的心血说没就没。直到我把断点续训老老实实做进训练流程才真正把“中断”从一场灾难降级成一次普通等待。模型中断了怎么让它记住上次干到哪听起来像个小问题但真要做对牵扯到模型权重、优化器动量、学习率调度器、随机数序列、数据迭代位置甚至分布式训练里的进程编号。这篇不打算讲太高深的内容而是想把这套东西拆开揉碎聊聊恢复到底在恢复什么、checkpoint文件里应该存哪些家当、写代码时从哪里接入、以及恢复后曲线不对劲时该从哪几个方向排查。适合正在自己跑实验的读者也适合要在团队里搭建训练流程的朋友拿去做对照。1. 先搞清楚“续训”到底续的是什么三层现场我见过太多人理解的断点续训就是“把模型权重存一下下次load回来接着跑”这是表面上对但实战里几乎一定会出问题。因为训练过程不是只有一组权重数字在变化它是一个持续运转的系统。你要恢复的不是某个文件而是把整个训练现场恢复到中断前的那个瞬间。这里我把“训练现场”拆成三层来看。1.1 进度现场你干到了哪一步第一层是最直观的进度信息当前跑了多少个epoch、多少步、这个batch算到哪个位置。训练日志里能看到的都是这一层。这一层的作用主要是“人对得上”——恢复之后你知道是从哪里接续的下游做报告、做对照、算剩余时间都要用到它。但严格来说模型本身并不知道“今天是第几个epoch”它只知道当前权重长什么样。进度现场只是一份给人类看的地图不是给模型看的记忆。1.2 数值现场模型和优化器此刻的状态第二层是数值状态也就是模型权重、优化器的动量项、学习率调度器当前给出的学习率。这一层比多数人以为的“进度”重要得多。关键在于优化器。SGD的动量、Adam的一阶二阶矩估计这些状态直接决定了下一步怎么走。如果你只加载模型权重不加载优化器状态那相当于你带着一个已经走了很远的背包客忽然把他脑中的地图和体力记忆全清掉让他重新学怎么走路。训练照样能跑但会退回到起步状态甚至因为学习率曲线已经降低模型还可能出现不收敛。1.3 环境现场随机数、混合精度、数据顺序这些“看不见”的状态第三层是运行环境的状态。这一层最容易被忽略却往往是恢复后训练曲线“不对劲”的元凶。包括CPU和GPU上随机数生成器的内部状态、混合精度训练里梯度缩放器的内部值、DataLoader在当前epoch内的具体迭代位置、分布式训练中各进程的上下文。说句通俗的训练里每一个“随机操作”——数据增强、Dropout、参数初始化时的抖动——都依赖随机数序列。随机数序列像一本流水账你中断时账记到哪一页就得从哪一页继续往下记。只把模型权重存下来等于把账本核心数字抄走了但记到哪一行这个信息丢了。一个更生活化的类比是看书。第300页的书签只告诉你“看到哪一页”但让你能继续流畅读下去的是你对前299页的记忆和理解。恢复训练也一样光有权重数字是不够的还得把前面那些“记忆”一起捞回来。明白这三层之后你在设计保存方案时就有了一个检查清单每次落盘时进度、数值、环境三样都要有缺一样以后恢复大概率要花更多时间排查。2. checkpoint文件里的七件家当与两种存档风格明白了三层现场接下来就具体到“一个checkpoint文件里应该装什么”。我把个人实践中反复确认过的东西整理成一张清单你可以直接对着抄。要保存的项大致内容不保存会怎样model.state_dict()模型全部可学习参数等于没存从头开始optimizer.state_dict()动量、一阶二阶矩、学习率相关内部量优化器状态归零训练回退scheduler.state_dict()当前学习率、内部epoch计数相关变量学习率回到初始可能重新预热epoch 与 global_step当前循环到哪无法准确记录训练进度best_metric当前最优验证分数及对应步数无法选择最佳模型CPU/GPU RNG statetorch随机数生成器状态数据增强、Dropout行为改变附加状态EMA影子权重、DataLoader seed、分布式sampler epoch特定场景下恢复不稳定第七件“附加状态”在普通实验里可以不存但只要你的训练脚本里用到了EMA、分布式采样器或者自定义数据迭代逻辑它就会从“可选”变成“必需”。我下面细说两种存档风格顺便讲讲保存文件本身的坑。2.1 两种存档风格全量单文件与目录式分存存这么多东西实际的组织方式一般有两种。第一种是单文件全量型直接用torch.save打包一个dict里面塞model、optimizer、scheduler、rng_state等所有状态。优点是一次加载、一个文件、逻辑简单个人实验和中小型模型完全够用缺点是每次保存都要序列化全部内容一个几百兆的模型加上优化器状态可能占掉几个GB保存耗时几秒甚至更久。第二种是目录式分存。每个checkpoint是一个目录内部结构类似checkpoints/run_123/ ├── meta.json # epoch, step, best_score, 训练超参数 ├── model.pt # 只存模型权重 ├── optimizer.pt # 优化器状态 ├── scheduler.pt # 调度器状态 └── trainer_state.pt # rng状态, dataloader状态等这种我主要在分布式和长期任务里用。原因很直接模型文件可以保存得频繁一些因为单文件小、成本低优化器和运行环境状态虽然也重要但变化频率没那么高保存间隔可以放宽。另外如果模型文件几百MB分目录后同步到共享存储在带宽紧张时也更可控。2.2 存档过程中最容易踩的几个坑存文件这件事本身也有几个坑。第一个是“只存最新一份”。我有一个朋友吃过这个亏训练时每5分钟覆盖式保存一次看日志和代码确实都在存但磁盘损坏导致唯一的checkpoint文件坏了中断恢复彻底无从谈起。我现在不管任务多小至少保留“best”和“latest”两份空间宽裕的话再留一份“latest_old”做滚动。第二个是写入过程本身可能不完整。保存一个3GB大文件时如果进程在写一半时被杀留在磁盘上的就是一个截断文件。解决办法是先写临时文件写完再用原子重命名替换也就是先save到一个.tmp再通过os.replace移动到正式路径。这样读取方永远看不到半个文件。3. 最小可用续训代码三处接入点与两个函数有了设计思路就该写代码了。这里给一套基于PyTorch的最小实现不一定覆盖你所有的训练框架但接入点足够通用改成你的框架也不难。3.1 保存函数与原子写入import os import torch def save_checkpoint(state, path): # 先写临时文件防止写一半出现半截文件 tmp_path path .tmp torch.save(state, tmp_path) os.replace(tmp_path, path)这里state本身是一个dict由调用方组装。我习惯在训练脚本里单独写一个build_checkpoint_state函数把该存的都打包进去def build_checkpoint_state(model, optimizer, scheduler, epoch, step, score, ema_modelNone, scalerNone): state { model: model.state_dict(), optimizer: optimizer.state_dict(), scheduler: scheduler.state_dict(), epoch: epoch, step: step, best_score: score, cpu_rng: torch.get_rng_state(), } if torch.cuda.is_available(): state[cuda_rng] torch.cuda.get_rng_state_all() if ema_model is not None: state[ema_model] ema_model.state_dict() if scaler is not None: # AMP的GradScaler state[scaler] scaler.state_dict() return state有人会问为什么不直接torch.save(model.state_dict())一个模型文件就完了答案是只存模型权重断掉的优化器动量和学习率会全部丢失训练进度会明显倒退。这条路上走过的人都懂。3.2 加载函数与设备映射def load_checkpoint(model, optimizerNone, schedulerNone, pathcheckpoint.pt, devicecuda, scalerNone): ckpt torch.load(path, map_locationdevice) model.load_state_dict(ckpt[model]) if optimizer is not None and optimizer in ckpt: optimizer.load_state_dict(ckpt[optimizer]) if scheduler is not None and scheduler in ckpt: scheduler.load_state_dict(ckpt[scheduler]) if scaler is not None and scaler in ckpt: scaler.load_state_dict(ckpt[scaler]) if cpu_rng in ckpt: torch.set_rng_state(ckpt[cpu_rng]) if torch.cuda.is_available() and cuda_rng in ckpt: torch.cuda.set_rng_state_all(ckpt[cuda_rng]) return ckpt注意torch.load加map_locationdevice很关键否则在CPU加载GPU权重会报张量存储设备不匹配。如果用的是比较新的PyTorch版本并且checkpoint里保存了非tensor的普通dict遇到weights_only报错时需要显式传weights_onlyFalse。这里提醒一下不要为了省事把所有torch.load都无脑加这个参数加载来源不可信的权重文件时还是要保持默认严格模式。3.3 训练循环里的三处改动第一处每个epoch结束或固定步数后保存。第二处训练启动前判断是否恢复。第三处恢复时要把epoch和step赋值回循环变量。下面这段是训练循环的骨架start_epoch 0 global_step 0 if resume_path: ckpt load_checkpoint(model, optimizer, scheduler, resume_path) start_epoch ckpt[epoch] global_step ckpt[step] for epoch in range(start_epoch, total_epochs): for batch_idx, batch in enumerate(train_loader): global_step 1 # 前向、反向、更新参数 if global_step % save_interval 0: state build_checkpoint_state( model, optimizer, scheduler, epoch, global_step, best_score, ) save_checkpoint(state, checkpoints/latest.pt) state build_checkpoint_state( model, optimizer, scheduler, epoch 1, global_step, best_score, ) save_checkpoint(state, fcheckpoints/epoch-{epoch 1}.pt)这里有三处容易错的地方。一是epoch循环的起点。如果你保存时是完成第3个epoch之后存的那么checkpoint里存的epoch应该是4或者至少能让你从第4个epoch正确开始取决于你是“已完成”还是“正在执行”的语义建议整个项目统一一种约定并在meta信息里写明。二是train_loader的起点。上面这个骨架默认每个epoch会重新遍历一遍完整的数据集断点恢复后batch从当前epoch的第0个batch开始而不是中断时的第几个batch。如果你的业务要求严格接续到中断时精确的batch位置就得在构建DataLoader时手动seek到对应位置或者使用支持断点迭代的采样器。绝大多数任务其实允许“当前epoch重新开始”关键是别把这个选择和“严格复原”混为一谈以免恢复后发现数据流和预想不一致。三是分布式训练。如果是DDP每张卡都要加载同一个checkpoint但恢复后各rank的DistributedSampler需要重新set_epoch才能保证数据切分顺序一致另外多卡训练里存checkpoint通常只在rank0上做加载时其他rank用同样的文件加载。4. 恢复后曲线不对劲先查这三样种子、AMP、数据取样按理说模型状态、优化器状态都加载回去了恢复后就应该和没中断时保持一致。但实际上很多人恢复训练后会发现loss曲线在恢复点附近出现异常波动甚至越来越差。排查下来高频原因基本就三个随机数状态没恢复、混合精度缩放器没保存、数据取样顺序乱了。4.1 随机数状态丢失数据增强、Dropout全换了“剧本”第一个元凶最常见症状是恢复后的第一个epoch验证指标和中断前最后一个epoch对不上或者train loss在恢复点上下跳变。原因是训练过程里的随机操作都依赖一个伪随机数序列你不保存RNG状态下一次取随机数时会从一个新的序列重新开始。数据增强打乱了图片顺序也好dropout mask重排了也好都会让模型看到的内容和原来预期的“时间线”发生偏移。我给个直观例子假设你训练到第10个epoch即将结束时中断恢复后没加载RNG state那么新epoch的数据增强对同一张图的裁剪、翻转结果会和原本第11个epoch应该看到的完全不同。单看一个epoch可能只是指标波动但如果你训练的模型对数据分布很敏感连续几个epoch下来恢复前精心调整的收敛轨迹就被破坏了。所以我的建议是保存checkpoint时把torch.get_rng_state()、每个GPU的cuda_rng_state一起存恢复时原样写回。如果你还用到了第三方库的随机数生成器比如NumPy也要一并存一下np.random.get_state()千万别只盯着torch。4.2 AMP的GradScaler状态没保存混合精度下的“隐性回退”第二个常见问题是用了混合精度训练却没有把GradScaler的状态保存下来。torch.cuda.amp里的GradScaler会在训练过程中根据梯度大小动态调整loss scale这个值在训练早期变化很大。恢复时不恢复它缩放器会从默认初始值重新开始梯度太小或太大时的处理策略可能在接下来若干个epoch都不稳定。最棘手的是这种异常不会立刻表现为“崩了”而是表现为训练变得“进展缓慢”或间歇性出现inf/nan警告。排查时很容易被当成是学习率或模型结构的问题。所以只要代码里出现过GradScaler保存checkpoint时就顺手把scaler.state_dict()存进去代价几乎为零。这个坑我至少在三个项目里见人踩过每次排查都要绕一大圈。4.3 DataLoader的取样顺序变了不是模型坏了是“喂饭的顺序”变了第三个问题在常规单机训练里比较隐蔽在分布式训练里则非常明显。DistributedSampler在每次进行数据切分之前需要调用set_epoch(epoch)这个调用的目的是让不同epoch的数据分到不同rank时随机顺序不一样。如果你恢复训练后忘了把sampler的epoch重新设置各卡拿到的数据顺序会和原训练计划错位分布式场景下各rank的batch对不上loss曲线自然乱。即使不用DistributedSampler只要你的训练流程在epoch维度上做过shuffle也要确认恢复后“下一个epoch的shuffle行为”和你中断前预期的行为一致。比如你在每个epoch之前调用了random.shuffle(indices)那就得把indices的当前状态一起存下来或者在恢复后重新执行一次与原来同种子的shuffle。这三类问题有个共同点排查时都容易被误判成“模型权重坏了”。所以如果恢复后看到曲线异常先别急着删checkpoint从头训按这个顺序查数据是不是原来该批数据、随机数序列是否一致、loss scale是否连续。大部分情况都在这里。5. 三种中断场景的真实应对定时值守、手动暂停、突发故障理论聊完了说点实战里最常见的三种中断场景以及对应的恢复策略。5.1 定时存档 best双轨我最常用的基础方案我默认的存档节奏是三层临时存档、周期存档、最优存档。临时存档是每500步左右存一个很小的当前状态用于“万一只丢了最近几十分钟工作”周期存档是每个epoch结束存一个带时间戳的完整checkpoint最优存档是每次验证分数刷新时单独把best模型拎出来存一份。这样设计有个实际好处恢复时你可以选择最近的周期存档点也可以选择训练效果最好的存档点。比如你中断前最后一个epoch验证分下滑了那你完全没必要非从最新点恢复直接从之前分数更高的存档点重来就行。5.2 手动暂停用信号在循环里优雅退出训练到一半你想调个超参数或者就是想停下来这种事很常见。打断训练时应该让程序在退出前完成一次保存而不是直接kill进程。在PyTorch训练循环里最直接的方法就是在主循环捕获KeyboardInterrupttry: for epoch in range(start_epoch, total_epochs): for batch in train_loader: # 训练逻辑 ... except KeyboardInterrupt: print(手动中断保存当前状态) save_checkpoint(state, checkpoints/manual_stop.pt)如果是集群环境不方便用CtrlC也可以监听SIGTERM信号在信号处理函数里只把“请求退出”标记置位主循环检测到标记后再优雅保存退出。这里有个细节不要在信号处理函数里直接执行torch.save这类耗时操作容易与主线程发生资源竞争应该只记录标记由主循环去处理保存。5.3 突发故障掉电、kill -9、宿主机回收意外断电、作业被杀、云主机被回收这类场景留给你的回调时间基本等于零唯一能依赖的就是前面那个周期存档。这时候最需要的是一个快速的“恢复决策”打开日志看最近一次保存的step和最近一次打印的step相差多少。如果差距大说明checkpoint落盘不够频繁下次把保存周期调小如果差距小直接加载即可。还要提一个比较隐蔽的坑磁盘写坏。掉电不一定会损坏模型文件但在写文件过程中掉电文件损坏的概率很高。所以除了原子写入之外最好永远留有上一个版本的存档。我现在训练时latest.pt背后永远还有一份latest.old.pt脚本里用滚动重命名维护它。别看这个细节小真遇到一次坏文件能省下重训一个礼拜的时间。5.4 什么时候不该恢复而是直接重跑断点续训不是万能的至少有两种情况我强烈建议直接重跑。一是训练进度太浅刚跑到5%左右就崩了。这时候恢复本身带来的工程复杂度对齐环境、验证随机状态、确认记录比重跑还高直接重跑往往更省心。二是代码或数据变了。模型结构和训练逻辑跟之前不一致时强行加载旧checkpoint反而会产生莫名其妙的错误。拿不准的时候先确认“这个checkpoint是怎么产生的”和“我现在要恢复的代码是什么样”这两个信息能不能对上对不上就不要硬恢复。6. 从个人脚本到团队基础设施checkpoint的工程化如果你已经把断点续训做进了自己的训练脚本那接下来自然要考虑一个问题怎么让整个团队、整个集群都稳定地依赖这套机制6.1 目录设计让恢复入口一眼可见先给出一套我目前在用的目录规范project/outputs/run_id/ ├── events/ # 训练日志、tensorboard事件 ├── checkpoints/ │ ├── meta.json # 每个存档对应的关键信息 │ ├── step-500-latest.pt │ ├── step-1000-best.pt │ └── ... └── resume.cmd # 根据最新进度自动生成的恢复命令meta.json里通常只放轻量字段run_id、epoch、global_step、best_score、保存时间、训练超参摘要。这样即使模型文件有几百MB你也能在不加载checkpoint的情况下快速知道“最近干到哪了”。resume.cmd是脚本自动生成的内容形如python train.py --resume checkpoints/step-500-latest.pt给队友用非常省心。6.2 启动参数与自动恢复把续训变成默认行为工程里有三种典型的启动方式传参指定恢复、自动寻找最新存档、团队配置中心下发恢复点。传参最直接适合人工介入自动恢复适合定时任务比如“容器重启后进入原工作目录找到最新的checkpoint继续跑”配置中心适合多人共用环境。三个我都见过没有绝对优劣核心是命令要可重复。一个容易忽略的点是恢复行为必须写在日志里。启动脚本恢复训练时打一行醒目的日志“从checkpoint step 500恢复训练”不要让队友事后翻半天才知道是从哪个点开始续的。6.3 跨机存储与共享存储的取舍checkpoint存放位置直接影响可用性。存本地磁盘最稳速度也快但宿主机没了就全没了存共享文件系统可以跨机恢复但写入瓶颈和高并发写容易踩坑。我现在的做法是训练过程中都写本地盘每次保存完成后由后台上传线程异步同步到对象存储或共享存储恢复时优先读本地盘没有就从远端拉回。这样既保证保存速度快又保证节点故障后还能恢复。对象存储的同步要讲究幂等同一个checkpoint可以在不同时间重复上传但最终覆盖路径要保持一致避免两个进程并发写同一个路径打起来。每次上传前给文件名带上run_id和global_step这样天然避免冲突。6.4 多机多卡恢复的额外检查项多卡训练恢复比单机多几个检查项每张卡加载相同的模型和优化器状态DistributedSampler的epoch要统一重设保存时按rank切分还是rank0汇总要提前定好。现在PyTorch有专门的分片式checkpoint接口如果团队用的是标准DDP可以直接研究一下如果还是自己管理务必在自测阶段模拟一次多卡中断不要等到正式训练时才暴露问题。多卡训练的checkpoint一旦出问题往往不止恢复不了还会出现不同步导致的隐性错误比单机难排查得多。说实话断点续训这套东西不会让模型效果变好它省下的是时间挣来的是底气。我现在习惯了这种心态任何训练任务启动前先确认checkpoint路径存在、保存周期合理再开始跑。万一真遇到中断第一反应也不是懊恼而是打开日志看最近存档点在哪然后用一条命令续上。希望这篇也能让你在下次面对中断时少一点手忙脚乱多一点从容。