← 全部文章模型训练

训练 NaN 的失败现场与状态恢复

区分 AMP 跳步、非有限 loss 和日志异常,整理复现训练 NaN 所需的实际 batch、模型、优化器和随机状态。

研究记录: 2026-10-28 — 2026-10-30
本站发布:

训练中断时,最后一条异常可能来自日志系统。一次运行在写优化器审计记录时,严格 JSON 序列化拒绝了非有限浮点数;后续另一段运行则明确记录了非有限 loss,并由保护机制停止。两种退出都涉及 NaN,却不能只凭最终异常定位产生 NaN 的运算。

恢复最近正常权重后,用少量真实图像做 FP32 与 AMP 前反向检查,结果均有限。这说明所测状态没有再次触发问题,但故障批次及其之前的更新没有被重放,数值问题仍未定位。排查的关键由此转向:需要保存什么,才能重建失败时的计算条件。

一、分清异常出现在哪个阶段

训练的一次迭代包含输入处理、前向、loss、反向、优化器更新及日志记录。它们的检查对象不同:

阶段 检查对象 能说明什么
前向与 loss 输入、中间激活、loss 分量 是否已经出现非有限值
反向 梯度,及其是否仍带缩放 反向是否暴露异常
更新 GradScaler 状态、实际更新事件 参数更新有没有执行
更新后 参数、buffer、优化器状态 更新后是否仍然有限
记录 被序列化字段和原始异常 观测系统是否完整保存现场

最后正常 checkpoint 中的参数和 buffer 有限,只能描述保存时的状态。它没有包含此后尚未保存的变化,也没有证明优化器状态、下一批输入和中间激活都正常。

NaN 出现在 loss 中,也不足以直接归因于 AMP 或某种注意力实现。它可能在更早的前向中产生,也可能来自 loss 的某个分支,需要沿计算路径缩小范围。

输入和参数有限,也不保证中间结果有限。PyTorch 的数值精度说明给出了范数计算的例子:FP32 输入仍在表示范围内,中间平方和可能溢出,最终数学结果却本来能够被 FP32 表示。对诊断的启发是,应同时检查运算的中间范围和执行精度;只扫 checkpoint 中有没有 NaN,会漏掉这种运行时溢出。

同时,首个被观测到的异常位置,只能作为继续缩小范围的起点。一个归约处出现非有限值,可能是其输入已经异常,也可能是局部归约才溢出;只有把输入、输出与必要的局部中间量对应起来,才能区分两者。

二、AMP 跳步与非有限 loss 退出

GradScaler 会检查优化器对应的梯度,在发现 Inf 或 NaN 时跳过该优化器的更新,并调整缩放因子。因此,循环计数增加一次,不一定意味着参数更新了一次。记录更新尝试和实际更新,可以解释训练是否真正向前推进。

一次少量的梯度溢出跳步,与训练已经产生非有限 loss、随后触发退出,是不同事件。缩放因子降低说明保护机制作出了响应,不足以证明它制造了异常,更不足以证明下一次就能恢复。

检查或裁剪缩放后的梯度时,还要明确位置。若希望使用未缩放梯度的范数,应先调用 scaler.unscale_(optimizer),再做检查或裁剪。把带缩放的范数与普通训练的范数直接比较,会混入缩放因子的影响。

梯度裁剪发生在前向和 loss 计算之后。若 loss 已经非有限,不能期待单纯调整裁剪阈值恢复此前的计算。

跳过优化器更新也不等于整次迭代回滚。训练前向可能已经修改 BatchNorm 的统计量,也已经消耗随机数;一些训练框架还会推进自己的计数或调度状态。GradScaler 负责的是对应优化器的梯度检查和更新控制,不会自动撤销这些副作用。因此,重放异常时应取失败前的状态,而不是仅恢复一次跳步后的权重,再假定其余状态完全没有变化。

这也解释了为什么“换成 FP32 后正常”更适合用于缩小范围,而不能直接作为 AMP 根因结论。改变精度会改变可表示范围、舍入和可能使用的内核。它说明故障依赖当前执行条件,仍需找到最小的触发计算;若只是扩大了数值范围,原先异常增长的中间量可能仍然存在。

三、保存实际计算使用的 batch

只保存样本编号,重读原图时可能无法还原随机裁剪、颜色变换和其他增强。若目标是重放异常前向,优先保存送入模型的实际张量、处理后的标签,以及 loss 所需的尺寸、掩码等元数据。

保存时应生成独立快照,避免后续原地操作或缓冲区复用改变内容。输入张量之外,还需要记录当时的训练或评估模式、autocast 设置、输入 dtype,以及运行环境。只写“FP16”无法区分全模型 .half() 与 AMP 的逐算子精度选择。

下面几种恢复目标需要不同的状态:

恢复目标 主要材料
重放一次前向与 loss 失败前模型状态、实际 batch、模式与精度设置
重放反向 前向条件,以及必要的梯度累积和缩放状态
重放一次更新 优化器状态、GradScaler 状态、更新前梯度与训练状态
重放此前训练轨迹 更早的完整恢复点,以及数据和随机过程的连续状态

含随机层的前向还依赖运算发生前的 RNG 状态。异常捕获时保存的 RNG 通常已经被本次执行消耗过,不能直接当作该次执行的起点;多 worker 预取和增强的随机流也未必包含在主进程 RNG 中。

因此,故障 batch 加最近 checkpoint 可以用于探索,但完整复现还取决于两者之间发生过什么。恢复权重后换几张未增强图像没有出错,并不矛盾。

PyTorch 的可复现性说明将随机源控制与确定性算子选择分开讨论。相同 seed 只规定随机流的起点;恢复到训练中途,还要知道随机流已经推进到哪里。即使随机状态相同,不同版本、设备或执行算法也未必逐位一致。诊断报告因此需要说明恢复的是一次具体计算,还是从更早位置继续执行整段训练。

cuDNN 的 benchmark 与确定性设置也有区别。前者影响卷积算法搜索与选择,关闭它并不自动保证所有算子确定。为了更容易复现而改变这些设置,应记录为诊断条件的变化;若故障因此消失,仍需回到原设置检查,而不能把运行环境变化当成模型修复。

四、让日志保留异常本身

严格 JSON 写入常使用 allow_nan=False,这能够防止生成不符合标准的非有限数值,却也可能在训练最需要日志时抛出新的异常。

诊断记录可以把非有限值表示成明确的状态,例如 { "finite": false, "value": "nan" },同时保留字段名、迭代位置及原始异常。不能把 NaN 替换成零再继续使用该值参与统计,否则日志会把失败伪装成正常样本。

记录失败也不应覆盖最初的训练异常。异常处理可以先保留原始错误,再尽力保存张量与结构化摘要,最后报告记录过程本身是否发生了第二个错误。

这次最早一段记录无法确定究竟是哪个被序列化字段首先变成非有限,不能把它归为单纯的日志故障。后续保护逻辑能够记录非有限 loss 并终止当前运行,改善的是可观测性,模型数值根因仍需单独查找。

大张量也不必全部写进日志。先记录 dtype、shape、有限值数量,以及有限元素上的最小最大值,再单独保存少量关键张量,更容易保留诊断信息。直接对含 NaN 的张量求最大值或均值,摘要自身也可能只剩一个 NaN,丢失其他元素的范围;没有有限元素时,应明确记录为空,而不是计算无意义的极值。

五、编译后非有限值要单独重现

另一次执行检查中,同一组权重在 eager FP16 下输出有限,经过 torch.compile 后出现非有限值。它与长时间训练后出现 NaN 的问题,触发条件不同,应保留各自的复现样例。

比较前先对齐输入、模式、缓存及参数。派生缓存如果没有跟随精度变化,可能先引入状态差异,参见PyTorch 派生缓存的精度转换与失效处理。状态对齐后,再比较 eager 与编译输出,先检查有限性,再计算误差,避免只留下一个无法解释的误差断言。

进一步定位可以分别使用 backend="eager"、"aot_eager" 和 "inductor",逐步区分图捕获、自动微分图处理与代码生成路径。这里是后续诊断路线;现有记录尚未定位到具体编译阶段,改变输入形状后通过也只能说明触发条件发生变化。

公开的数值精度经验也提醒我们区分普通误差与非有限值。浮点运算不满足结合律,融合、归约顺序和中间精度变化可以带来有限的小偏差;出现 Inf 或 NaN 则需要单独检查表示范围和非法运算。不能用“编译后不保证逐位相同”笼统解释所有差异,也不能仅凭编译路径失败就认定是编译器缺陷。

遇到 NaN 输出时,直接计算相对 L2 往往只会得到 NaN,无法说明异常覆盖范围。更有信息的比较顺序是:先检查双方的非有限值数量和位置,再比较共同有限位置上的绝对误差与相对误差,最后观察异常从哪一段输出开始出现。分母接近零时也要同时查看绝对误差,避免把很小的绝对差放大成惊人的相对数值。

定位首个异常时,模块 hook 可以帮助缩小范围,但模块内部的函数式运算不会自动成为独立模块;增加检查、强制同步或拆图还可能改变执行路径。探针发现线索后,仍需回到原先失败的执行方式验证。

torch.autograd.detect_anomaly(check_nan=True) 是另一种公开的反向诊断工具:在其上下文内执行前向与反向,可以在反向失败时提供关联的前向栈,并检查反向产生的 NaN。它有运行开销,也不能替代前向中对 Inf、NaN 和数值范围的检查。使用时应围绕可复现的局部样例,而不是将开启异常检测后的吞吐与正常训练直接比较。

六、参考资料

来源与说明

本文保留原始技术记录中的实验条件与限制,不代表当前所有产品均支持文中模型、设备或性能。

阅读原文(中文) · 查看固定版本 · MIT License

© 2026 moonnryan。本站提供英文翻译,英文版由 AI 辅助翻译并对照中文原文核对。

Read this article in English ↗