训练过程 loss 突然变成 NaN，直接调小学习率属于盲猜，最靠谱的办法就是注册 forward hook 逐层监控每一层输出，精准抓到最先出现 inf、NaN 的网络层和训练步数，从源头定位溢出点，而不是等 loss 直接崩了之后再复盘，这个方法也是 AI 工程师调试模型发散最常用的稳定手段。之所以 loss 容易出现 NaN，是因为 NaN 有传导特性，只要前向计算里任意一步算出 inf，后续所有张量、梯度、loss 都会跟着变成 NaN，常见诱因有学习率太大梯度爆炸、自定义损失函数 log (0)、除零、FP16 混合精度溢出、输入数据自带异常值，单纯看 loss 日志根本分不清问题出在哪一层，很难一次性修复。用 hook 逐层定位的操作流程很清晰，第一步先写好 NaN 检测 hook 函数，用来判断每层输出张量是否存在 inf、NaN；第二步遍历模型所有子模块，批量注册 forward hook；第三步启动训练，一旦某一层输出出现异常，脚本立刻打印层名称、步数、张量极值；第四步拿到异常层之后，单独抽这一层的输入做离线复现，判断是数值计算问题还是超参问题；第五步修复完成后，记得注销 hook，避免 hook 持续占用算力拖慢训练速度。总的来说，hook 定位是定位 loss 变 NaN 的精准手段，它可以锁定第一个溢出节点，搭配 autograd anomaly 检测、梯度裁剪、加 epsilon 稳定计算，就能解决大部分训练 NaN 问题，不过 hook 会增加一定开销，只适合调试阶段，正式训练要关掉。
一、loss 变 NaN 常见根因清单
- 超参问题：学习率过大，梯度爆炸，权重更新直接溢出
- 损失函数：log (x) 没有加 epsilon，出现 log (0)；分母为 0
- 精度问题：FP16 混合精度训练 loss scale 异常、softmax 数值溢出
- 数据问题：训练集包含 NaN、无穷大、极端异常数值
- 归一化：LayerNorm、RMSNorm 方差接近 0，开方除零
- 算子问题：attention 分数过大，exp 计算溢出产生 inf
二、两种调试方案对比表
| 对比项 | hook 逐层监控 | torch.autograd.detect_anomaly |
|---|---|---|
| 定位目标 | 前向传播每一层输出，找到最先爆 inf/NaN 的层 | 反向传播报错，定位产生异常梯度的算子 |
| 优点 | 直接锁定前向溢出层，方便复现，日志清晰 | 框架原生，不用自己写 hook 代码 |
| 缺点 | 需要编写注册代码，会增加训练开销 | 反向报错，不一定能定位源头，速度更慢 |
| 适用场景 | 排查 loss 突然 NaN、前向溢出 | 梯度计算异常、反向传播产生 NaN |
三、hook 逐层定位完整操作步骤
- 编写 nan 检测 hook：读取 module 输入输出，判断张量是否存在 isnan、isinf
- 遍历模型模块注册 hook：model.modules () 循环，调用 register_forward_hook
- 启动训练脚本，正常跑训练循环
- 触发异常时自动打印：层名、当前 step、张量最大最小值、NaN 数量
- 离线复现：拿到该层输入，单独跑前向，确认触发条件
- 修复验证：调整学习率、增加 epsilon、梯度裁剪、修正损失写法
- 调试完成：移除 hook，正式训练
四、修复手段清单
- 梯度裁剪：clip_grad_norm 限制梯度范数，防止梯度爆炸
- 数值稳定改造：log (x+1e-8)、分母加极小 epsilon，优先使用 log_softmax
- 降低学习率、使用 warmup，避免前期权重剧烈跳动
- 清洗数据集，过滤掉输入中的 NaN、inf 异常样本
- 混合精度场景，调整 loss scale，关闭不稳定算子 FP16 计算
- attention 计算做分数 clamp，防止 exp 溢出
【常见问题】
问题 1：loss 变 NaN 用 hook 逐层定位，hook 到底监控的是什么？
回答 1：loss 变 NaN 用 hook 逐层定位，forward hook 是监控每一层前向传播的输出张量，实时检测张量里有没有 inf 和 NaN，抓到最先出现异常的网络层，用来定位 loss 变 NaN 的源头。
问题 2：loss 变 NaN 用 hook 逐层定位，和 autograd anomaly 检测有什么区别？
回答 2：loss 变 NaN 用 hook 逐层定位是抓前向传播的异常层，autograd anomaly 是抓反向梯度异常算子，优先用 hook 逐层定位，适合大部分训练 loss 突然变 NaN 的场景。
问题 3：loss 变 NaN 用 hook 逐层定位，正式训练的时候要不要保留 hook？
回答 3：loss 变 NaN 用 hook 逐层定位只适合调试阶段，hook 会增加计算开销，定位完成、修复 NaN 问题之后，必须注销或者删掉 hook，不要放在正式训练代码里。
问题 4：loss 变 NaN 用 hook 逐层定位，抓到异常层之后怎么进一步排查？
回答 4：loss 变 NaN 用 hook 逐层定位抓到异常层之后，把该层输入单独保存下来离线跑前向，复现溢出，再检查损失、归一化、exp、除法这些容易出数值不稳定的计算。
【参考资料】
[2026-11-11]（手把手教你：loss 变 NaN 用 hook 逐层定位｜AI 工程师踩坑实录・第 4 篇）博客园
[2026-11-11]（手把手教你：loss 变 NaN 用 hook 逐层定位）阿里云开发者社区
[2026-11-11]（手把手教你：loss 变 NaN 用 hook 逐层定位）腾讯云开发者社区
[2026-11-10]（PyTorch hook 定位 NaN、inf 溢出层调试方案）PyTorch 官方论坛
【免责声明】
本文仅为 loss 变 NaN 使用 hook 逐层定位的调试技术科普，不代表万能方案；不同框架、混合精度、分布式训练场景下 NaN 成因差异较大，定位脚本、修复策略需要结合自己模型实测验证。