PyTorch autograd anomaly detection 怎样定位 NaN 梯度

训练跑到 loss.backward() 才出现 NaN 或 Inf 梯度时，最有效的第一步不是立刻调学习率，而是让 PyTorch 记录“哪个反向算子出错，以及它对应的前向调用位置”。torch.autograd.detect_anomaly(check_nan=True)正好用于这类调试：它能在反向异常时给出前向 traceback，并在反向计算产生 NaN 时抛出错误。

官方文档：https://docs.pytorch.org/docs/stable/autograd.html

下面只围绕一个边界展开：定位 NaN/Inf 梯度的来源。它不替代数据清洗、混合精度策略或模型结构调参，定位完成后应关闭异常检测，再用小样本复查修复结果。

先判断异常是在 forward 还是 backward

先把训练循环拆成三个观测点：输入是否有限、loss 是否有限、反向后的参数梯度是否有限。这样可以避免把 forward 已经产生的 NaN 误判成 autograd 的问题。

import torch

def assert_finite(name, value):
    # 统一检查张量中是否出现 NaN 或 Inf，先定位异常阶段。
    if not torch.isfinite(value).all():
        raise FloatingPointError(f"{name} 包含 NaN 或 Inf")

def inspect_batch(model, batch, target):
    # 只做一次小批量检查，不把调试逻辑扩散到完整训练任务。
    output = model(batch)
    assert_finite("模型输出", output)
    loss = torch.nn.functional.cross_entropy(output, target)
    assert_finite("loss", loss)
    return loss

如果输入或输出在 backward() 之前就不有限，先检查数据归一化、除法分母、对数输入和混合精度转换；只有 forward 与 loss 都有限、反向才报错时，才进入 autograd 异常定位。

用 anomaly detection 找到反向失败算子

确认问题发生在反向后，把异常检测包在产生计算图和 backward() 的同一小段代码外侧。前向也要放在上下文中，因为 PyTorch 需要保存前向信息，才能在 backward 失败时回溯到对应位置。

from torch import autograd

def debug_step(model, batch, target):
    # 异常检测只包住一次小样本反向，避免拖慢整轮训练。
    with autograd.detect_anomaly(check_nan=True):
        output = model(batch)
        loss = torch.nn.functional.cross_entropy(output, target)
        # check_nan=True 会在 backward 生成 NaN 时抛出异常。
        loss.backward()
    return loss.detach()

看到 traceback 后，先找“反向报错”之前的前向调用位置，再回到该算子的输入和参数。这个位置是最有价值的线索，但不是自动证明根因：同一算子可能只是第一个暴露上游异常的地方。

把 NaN 缩小到输入、loss 和梯度

异常检测指出反向链路后，再沿着张量和梯度边界做有限值检查。重点看三处：进入可疑模块的输入、该模块输出参与计算的 loss、反向后参数的 .grad。

def inspect_gradients(model):
    # 只报告出现非有限梯度的参数名，便于和 traceback 对照。
    bad_names = []
    for name, parameter in model.named_parameters():
        if parameter.grad is not None and not torch.isfinite(parameter.grad).all():
            bad_names.append(name)
    if bad_names:
        raise FloatingPointError("非有限梯度: " + ", ".join(bad_names))

def debug_backward(model, batch, target):
    # retain_grad 只用于需要观察的中间张量，避免无边界保存梯度。
    output = model(batch)
    loss = torch.nn.functional.cross_entropy(output, target)
    if not torch.isfinite(loss).all():
        raise FloatingPointError("loss 在 backward 前已经异常")
    loss.backward()
    inspect_gradients(model)

如果 loss 有限而某一组参数梯度先变成 NaN，优先检查 traceback 指向的算子是否存在除零、对数负数、极端指数或溢出；如果梯度全部有限但训练指标异常，再回头看标签、损失定义和梯度累积是否符合预期。

修复数值边界，再做一次小样本反向复查

定位后只修复对应边界，不要把异常检测长期留在生产训练路径。常见修复包括给除法分母加入明确的 epsilon、在进入 log 前限制输入范围、检查标签与 loss 的约定，以及把不稳定的中间值用更合适的 dtype 计算。

def safe_ratio(numerator, denominator, eps=1e-8):
    # epsilon 只保护分母接近零的情况，具体值要结合数据尺度选择。
    safe_denominator = denominator.clamp_min(eps)
    return numerator / safe_denominator

def finite_loss(loss):
    # 修复后先在小批量上确认 loss 可用于反向，再恢复正常训练。
    if not torch.isfinite(loss).all():
        raise FloatingPointError("修复后 loss 仍不是有限值")
    return loss

复查时重复一次“输入、loss、梯度”三点检查，并确认 traceback 不再出现同一个异常算子。复查通过后删除或关闭 detect_anomaly，保留轻量的失败日志与必要的有限值监控。

常见疑问

为什么打开 anomaly detection 后训练明显变慢？

因为它会为调试保存前向上下文并增加异常检查。它适合一次小批量定位，不适合整轮生产训练。

traceback 指向的算子一定就是根因吗？

不一定。它通常是第一个让反向失败的观察点，上游输入、loss 定义或数值范围仍需用 isfinite 检查确认。

只检查参数梯度够不够？

不够。参数梯度只能说明异常已经传到参数边界；要缩短定位路径，还应在可疑模块输入、输出和 loss 处检查有限值。

总结一下：先把 forward 与 backward 分开，再用 detect_anomaly(check_nan=True) 获取前向 traceback，最后用输入、loss 和参数梯度的有限值检查确认根因。定位完成后关闭调试开关，保留小样本回归和必要的数值监控。