从梯度检查点面试题延伸:训练恢复为什么不能只保存权重

图中把恢复所需状态拆成参数、优化过程、随机与数据位置,文件存在只是第一步。

训练进程退出后,把模型权重加载回来,为什么接下来的损失曲线仍可能和原来不同?因为训练不是一个只依赖当前参数的函数。优化器记住了过去的更新,调度器记录了进度,数据加载和随机数也有自己的位置。恢复文件能被读取,与训练真正接着原轨迹继续,是两种不同承诺。

先说你要恢复到什么程度

最宽松的目标是拿权重重新开始一轮训练,允许优化器、顺序和随机状态重置。更严格的目标是延续原来的训练过程。再严格一些,希望在固定环境下得到相同的后续数值。目标不同,必须保存的内容和验证方式也不同。

如果只需要推理,参数与必要的模型配置可能就足够。但若目标是续训,不能拿一个能做推理的权重文件,直接宣称拥有完整恢复能力。文件格式解决“怎样装载数据”,状态契约解决“装载后意味着什么”。

参数之外,优化器也有记忆

以带动量的更新为例,即使当前位置相同,只要累积的动量不同,下一步参数变化就可能不同。某些优化器还有一阶、二阶统计量及步数;把它们清零,相当于改变优化过程,而不是单纯继续原来的下一步。

学习率调度器也要与步数对应。如果恢复到较晚参数,却让调度器从预热起点重新开始,实际更新幅度会变化。混合精度训练的缩放状态、梯度累积进度等,也应根据具体实现纳入恢复契约,不能只记一个 epoch 标签就认为状态已经齐全。

保存时最好选择清楚的训练边界,例如一次完整优化器更新之后。若在若干微批次之间保存,尚未应用的累积梯度是否需要恢复、计数从哪继续,都必须明确,否则会漏算或重复计算样本贡献。

能加载权重,不等于接着训练

数据位置和随机状态要对上

同一份数据集,不同打乱顺序就会产生不同的后续批次。只记“第几个 epoch”,通常不足以精确恢复到该轮中间位置。需要结合采样器状态、已消费位置、数据版本以及分布式切分方式,定义下一批到底是什么。

随机状态也不仅是最开始设过的种子。运行到中途后,生成器已经消耗了一段随机序列;重新设置初始种子会回到开头,不是回到中断处。不同库、设备和工作进程还可能拥有各自的随机源,应按实际流程逐项盘点。

即使这些状态齐备,也不能保证跨软件版本、跨硬件和不同并行配置逐位一致。某些算子具有非确定性,计算次序变化也可能改变舍入。能恢复有效训练,与每一位完全相同仍应分开描述。

一组状态必须来自同一个时间点

如果参数保存于更新之后,优化器却保存于更新之前,两个文件各自都能读取,组合起来仍然是不一致快照。分布式或异步保存时,尤其需要明确版本、提交边界及哪些分片属于同一次快照。

可以先写入临时位置,完成必要的完整性校验后再发布一份清单,让恢复端只选择完整版本。保留旧的可用版本,也比“先覆盖旧文件,写一半再发现失败”更有恢复余地。这是快照发布原则,不代表任何简单重命名都能自动保证跨存储系统的事务原子性。

配置和代码版本同样值得记录。分词器、数据清洗规则或模型结构变化后,张量形状相同也不一定意味着语义相同。对不上时应显式报错或声明迁移,不要静默忽略关键字段。

做一次真正的恢复演练

验证方法可以先在小任务上连续执行几步,保存中间状态,然后从该状态重新启动,比较后续输入、更新计数、学习率和参数。再故意漏掉优化器或数据位置,看比较在哪一步出现差异。

这个实验应在独立恢复进程中进行,避免原进程残留的全局变量替缺失字段“兜底”。只在同一个对象上保存再立即加载,可能掩盖根本没有写入文件的状态。

恢复能力不是“磁盘上有文件”,而是能说明快照边界、恢复目标和未保证的范围。你会怎样证明下一批数据没有被重复消费,也没有从恢复点之后悄悄跳过?