Debugging Gradients That Misbehave

When backward goes wrong, you have tools

Training that diverges, NaN losses, gradients that mysteriously become zero or infinity — these are the classic "what's wrong with my model" symptoms. Four PyTorch-specific tools cover most of the diagnostic surface:

Gradient norm logging. Add clip_grad_norm_(model.parameters(), float('inf')) after backward and log the returned value. A spike to 1e6 is your "what just happened" signal.
NaN / Inf detection. torch.isfinite(loss) at every step is cheap insurance — it catches both Inf and NaN. torch.autograd.set_detect_anomaly(True) at the top of training will tag the exact op that produced a NaN backward — slow, but priceless when you're stuck. Watch the scope mismatch: anomaly detection only fires on NaN, not Inf (see "Inf vs NaN" below).
Per-parameter gradient stats. Walk model.named_parameters() after backward and print mean/std/max for each. Vanishing-gradient bugs are obvious in this view.
Gradient hooks for shape mismatches. Register a hook on a layer's output to verify gradient shapes match what you expect — useful when adding skip connections or custom layers.

Inf vs NaN — anomaly detection has a narrow scope

The gap most tutorials skip: set_detect_anomaly(True) raises only on NaN values in backward output, not Inf. The same function can land in either bucket depending on its input.

Inf sources — anomaly detection stays silent; use torch.isfinite() to catch:

sqrt(0).backward() — derivative is 1/(2·sqrt(x)) = ∞ at x=0. Fix: (x + 1e-9).sqrt().
log(0).backward() — derivative is 1/x = ∞ at x=0. Fix: (p + 1e-9).log().

NaN sources — anomaly detection fires and points at the exact op:

sqrt(-1), log(-1) — NaN in forward, propagates to backward.
0 / 0 — usually from a normalize-by-norm where the norm is zero.
0 * log(0) — 0 × (-∞) = NaN.
log(softmax(x)) where softmax collapses to ~0 on an entry — produces NaN backward (use F.log_softmax for numerical stability — see the exercise).
fp16 overflow that turns into 0×∞ — switch to bf16 or add a GradScaler.

Code

set_detect_anomaly — the Inf trap (won't raise)·python

import torch

# Turn on at the start of training when debugging — turn OFF for real runs
torch.autograd.set_detect_anomaly(True)

# Inf trap: anomaly detection WON'T catch this — sqrt(0) backward = inf, not nan
x = torch.tensor(0.0, requires_grad=True)
y = torch.sqrt(x)
y.backward()
print(x.grad, torch.isfinite(x.grad))   # tensor(inf), tensor(False) — no RuntimeError
# Use isfinite() to catch Inf; anomaly detection only fires on NaN.

set_detect_anomaly — NaN backward DOES raise·python

import torch

torch.autograd.set_detect_anomaly(True)

# Same sqrt function, different domain — sqrt(-1) IS caught
x = torch.tensor(-1.0, requires_grad=True)
y = torch.sqrt(x)            # forward: nan, propagates to backward
y.backward()
# RuntimeError: Function 'SqrtBackward0' returned nan values in its 0th output.
# The traceback points at the exact forward call that created the nan-producing op.

Per-parameter gradient stats — the diagnostic walk·python

import torch
import torch.nn as nn

model = nn.Sequential(nn.Linear(100, 50), nn.ReLU(), nn.Linear(50, 10))
loss = nn.functional.cross_entropy(model(torch.randn(8, 100)), torch.randint(0, 10, (8,)))
loss.backward()

print(f"{'name':25s}{'mean':>12s}{'std':>12s}{'max_abs':>12s}{'has_nan':>10s}")
for name, p in model.named_parameters():
    if p.grad is None:
        print(f"{name:25s}  NO GRADIENT")
        continue
    g = p.grad
    print(f"{name:25s}{g.mean():12.2e}{g.std():12.2e}{g.abs().max():12.2e}"
          f"{str(torch.isnan(g).any().item()):>10s}")

Cheap NaN insurance — every training loop·python

import torch
import torch.nn as nn

model = nn.Linear(10, 2)
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)

for step, (x, y) in enumerate(loader):
    optimizer.zero_grad()
    loss = nn.functional.mse_loss(model(x), y)

    # Cheap. Worth it. isfinite() catches both Inf and NaN.
    if not torch.isfinite(loss):
        print(f"Step {step}: non-finite loss = {loss.item()}; halting.")
        # Save a snapshot of the offending batch for later analysis
        torch.save({'x': x, 'y': y, 'state': model.state_dict()},
                   f'nan_snapshot_step_{step}.pt')
        break

    loss.backward()
    optimizer.step()

Debugging Gradients That Misbehave

When backward goes wrong, you have tools

Inf vs NaN — anomaly detection has a narrow scope

Code

External links

Exercise

Progress

Comments 0