pytorch - 💡(How to fix) Fix Add torch.utils.nan_guard for forward-pass NaN/Inf debugging [1 participants]

from torch.utils.nan_guard import NaNGuard

with NaNGuard(model):
    out = model(x)
# NaNGuardError: detected NaN/Inf in output of submodule '3.attn'
# (MultiheadAttention):
#   [0]: shape=(8, 128, 512) dtype=torch.float32 device=cuda:0 nan=512 inf=0
Walks every submodule of model, scans its output (handling tuples / dicts /
nested structures via pytree), and raises naming the first offending
submodule, the path inside its output structure, and tensor stats. Optional
check_inputs=True distinguishes the producer of NaN/Inf from downstream
propagators. on_detect="warn" for non-raising mode.

Motivation
PyTorch ships torch.autograd.detect_anomaly for the backward pass, but
there's no lightweight forward-only equivalent. When forward activations go
NaN, the standard workaround is a copy-pasted snippet:


for n, m in model.named_modules():
    m.register_forward_hook(lambda mod, inp, out, n=n: ...)
Every reinvention of this gets at least one of the following wrong:

Doesn't handle non-tensor outputs (dicts, tuples, dataclasses).
Doesn't separate producers from propagators — every downstream module fires.
Leaks hooks on exception.
Doesn't skip integer / empty / sparse tensors.
Doesn't tell you the path inside a structured output.
A correct version belongs in torch.utils.*.

Pitch
Pure-Python utility next to other debug helpers under torch.utils. No C++,
no autograd-engine changes, no new top-level API surface. Composes with
nn.Module's existing forward-hook machinery — no new abstractions. Cost is
opt-in (only inside the with block) and the docstring is explicit that it's
a debugging tool, not for steady-state training.

Complements torch.autograd.detect_anomaly rather than replacing it.



cc @albanD @mruberry @jbschlosser @walterddr @mikaylagawarecki