Debugging — nan, inf, and Overfitting a Single Batch

Series: Deep Learning & LLMs From Scratch — 80 sessions · Session 19 / 80 · Module M03 · ~2 hours

0 · Debugging tests claims, not just curves

A falling loss proves that the written objective fell on the supplied data. It does not prove that labels are aligned, the split is independent, the metric is correct or the intended task is being learned. Debugging requires both program logic and numerical inspection. No checklist diagnoses every training bug in ten minutes.

The original attributed a specific count and debugging postmortem to Llama-2 without adequate primary support. Here the useful lesson is retained as a hypothetical incident: a previously finite run reaches NaN. Save enough trusted local state and data identifiers to reproduce the first failing update; do not invent a hardware cause from the curve.

Prerequisites

Use the training-loop state contract and precision controls. Run the CPU blocks in order. No MNIST/model download, package install, paid service or GPU is required.

1 · Rule zero — overfit a single batch

A fixed-batch test removes changing sampling as a variable. Disable augmentation and dropout, remove regularization for the diagnostic, choose a model capable of expressing the tiny task, and train repeatedly on the same examples. Record loss and predictions, gradients and actual parameter differences.

Success says the chosen model/optimizer can fit these supplied examples. It does not establish label meaning, data integrity, generalization or leakage freedom. A model can memorize a fixed random permutation of labels. Failure is not proof of broken wiring: insufficient capacity, contradictory labels on identical inputs, an unsuitable loss floor, label smoothing, regularization or an inappropriate learning rate can prevent near-zero loss. More parameters than constraints is not a general nonlinear optimization theorem.

1.1 A reproducible fit, and a failure the fit cannot catch

This program fits an eight-example identity-feature classification task twice: once with arbitrary fixed labels and once with another label assignment. Both can be memorized. That is the negative control against treating successful fitting as semantic validation.

import torch
from torch import nn
import torch.nn.functional as F
 
torch.set_num_threads(1)
torch.manual_seed(13)
x = torch.eye(8)
y = torch.arange(8) % 3
 
def fit(labels, lr=1., steps=250):
    torch.manual_seed(13)
    model = nn.Linear(8, 3)
    opt = torch.optim.SGD(model.parameters(), lr=lr)
    initial = F.cross_entropy(model(x), labels).item()
    for _ in range(steps):
        opt.zero_grad(set_to_none=True)
        loss = F.cross_entropy(model(x), labels)
        loss.backward(); opt.step()
    final = F.cross_entropy(model(x), labels).item()
    return model, initial, final
 
m, initial, final = fit(y)
assert final < .04 and torch.equal(m(x).argmax(1), y)
wrong_labels = y[torch.tensor([3, 0, 7, 2, 1, 5, 4, 6])]
_, _, shuffled_final = fit(wrong_labels)
assert shuffled_final < .04  # pass does NOT certify label meaning
# Identical inputs with conflicting labels have an irreducible average CE floor.
z = torch.zeros(2, 2, requires_grad=True)
assert abs(F.cross_entropy(z, torch.tensor([0,1])).item()-torch.log(torch.tensor(2.)).item()) < 1e-7
print("tiny fit", initial, final, "fixed-shuffled fit", shuffled_final)

Removing a ReLU collapses two affine maps into one, but that linear model can still fit a linearly separable task. MSE on correctly shaped one-hot targets is a different, often less convenient classification objective, not intrinsically an invalid loss. Deliberate bug experiments need an actual failure assertion rather than assuming every unconventional choice must fail.

2 · Loss is NaN — a decision tree

Step 1 — locate the first failing boundary

Timing narrows the search, not the diagnosis. Step-zero NaN suggests inputs, initialization or the forward/loss; a later failure can be a data outlier, state corruption, scale/precision issue or an invalid operation reached only later. A fixed reproducible step does not prove a hardware fault. Recreate the preceding checkpoint, batch order, RNG and configuration on a tiny eager FP32 reproduction.

Step 2 — check inputs, labels and objective contract

Check shapes, finite values, dtypes, range, masks and denominators before optimization. For class-index cross entropy, logits have a class dimension and integer labels are in range or the configured ignore index. Binary/multilabel logits pair with BCEWithLogitsLoss; probabilities pair with BCELoss. The old instruction saying logits+BCEWithLogits was wrong was reversed.

Count supervised tokens after masking. Zero valid targets can give an undefined mean. Confirm next-token target shifting and that padding is not scored. Incorrect broadcasting can silently turn one prediction per example into a matrix of losses.

Step 3 — anomaly detection has a bounded promise

torch.autograd.detect_anomaly(check_nan=True) retains forward traceback information and errors on NaNs produced in backward. Enter its context around forward and backward to capture the useful forward stack. It is not a complete detector of every forward nonfinite value or every infinite backward value. Use explicit finite checks too. It slows execution; there is no fixed 2–4× factor for every graph.

import warnings
import torch
import torch.nn.functional as F
 
logits = torch.tensor([[1000., -1000.]], requires_grad=True)
target = torch.tensor([1])
naive = -logits.softmax(1)[0,1].log()
stable = F.cross_entropy(logits, target)
assert torch.isinf(naive) and stable.item() == 2000.
stable.backward(); assert torch.isfinite(logits.grad).all()
# A zero-over-zero backward is caught; capture expected diagnostic warnings.
with warnings.catch_warnings(record=True) as notices:
    warnings.simplefilter("always")
    try:
        with torch.autograd.detect_anomaly(check_nan=True):
            a = torch.tensor(0., requires_grad=True)
            (a/a).backward()
    except RuntimeError as error:
        assert "nan" in str(error).lower()
    else:
        raise AssertionError("negative control was not detected")
assert notices
# sqrt(0) is finite forward and infinite backward: test both boundaries.
b = torch.tensor(0., requires_grad=True); b.sqrt().backward()
assert torch.isinf(b.grad)
print("stable-logit and anomaly/nonfinite controls passed")

Step 4 — fix the operation, not only its symptom

log(softmax(logits)) can underflow a probability to zero even if softmax internally subtracts the maximum. Use log_softmax or cross entropy on logits. Raw exp can overflow; subtracting a maximum alone does not make taking a rounded zero's logarithm safe. log(x+eps) and sqrt(x+eps) change the objective and do not fix negative-domain inputs. A reciprocal's epsilon must be meaningful at the active dtype and scale. Avoid evaluating invalid branches and then masking their outputs: autograd may still encounter the invalid derivative.

Step 5 — learning-rate experiments after diagnosis

A smaller LR is a controlled hypothesis, not a nine-in-ten diagnosis. For $L(w)=aw^2/2$ and SGD, $w_{t+1}=(1-\eta a)w_t$; convergence requires $0<\eta a<2$. That gives a known test case, not a neural-network stability bound. If loss is invalid before any optimizer step, lowering LR cannot repair that forward computation.

3 · Loss will not move — gradients, updates and baselines

Uniform class probabilities give natural-log cross entropy $\log C$; for ten classes this is about 2.302585. The converse is false: the same average loss can come from many nonuniform predictions. Random initialization need not produce uniform logits. Test uniform zero logits as a loss-unit oracle instead of treating every initial loss different from 2.30 as a bug.

Inspect the whole chain:

1. Does the intended parameter require gradients and participate in this graph? A legitimate frozen or unused branch can have grad=None.
2. Is the parameter registered and present in the optimizer's groups? An unregistered tensor can receive gradients yet never be stepped by that optimizer.
3. Is the gradient finite/nonzero where expected? Small magnitude depends on units; it is not automatically a dead layer.
4. Does the parameter actually change after stepping? LR zero, wrong optimizer instance, a stale checkpoint reload, underflow or an omitted step can stop changes.
5. Are inputs and labels paired through the same shuffled indices? A fixed wrong mapping can still fit, so audit example identities separately.

import torch
from torch import nn
import torch.nn.functional as F
 
assert abs(F.cross_entropy(torch.zeros(7,10), torch.arange(7)).item()-2.302585093) < 1e-6
p = nn.Parameter(torch.tensor(2.)); q = nn.Parameter(torch.tensor(2.))
opt = torch.optim.SGD([q], lr=.1)  # wrong parameter group
(p*p).backward(); before = p.detach().clone(); opt.step()
assert p.grad.item() == 4. and torch.equal(p, before)
# Loss broadcasting negative control: enforce the intended shape first.
pred = torch.ones(4,1); labels = torch.ones(4)
assert pred.shape != labels.shape
# Excessive SGD rate on a known quadratic diverges; no random bug diagnosis needed.
w = torch.tensor(1.)
for _ in range(8): w = w - 3*w
assert abs(w.item()) == 256.
# Forgetting zero_grad accumulates old gradients even though each graph is new.
r = nn.Parameter(torch.tensor(2.)); (r*r).backward(); (r*r).backward()
assert r.grad.item() == 8.
print("loss units, wrong optimizer, shape, LR and accumulation controls passed")

Counting parameters helps catch an accidentally empty network, but ten parameters are enough for some tasks. Fixed shuffled labels are a useful memorization control; comparing their training speed does not reliably identify a broken pipeline. Evaluate semantics and held-out behavior independently.

4 · Val diverges from train — compare like with like

Both low is not permission to ship. Dropout-active training loss can naturally exceed evaluation loss; that is not automatically leakage. Compare the same loss components (data loss versus data+regularizer), normalization, class weighting and denominators. Split by deployment unit—person, document family, time—not simply by shuffled rows. Fit preprocessing on training data only.

5 · Weight and gradient histograms

Track per-layer mean, standard deviation, finite fraction, zero fraction, norms and relative updates. Histograms need not be Gaussian or symmetric; embeddings, biases, sparsity and learned structure often are not. A spike at zero might be an inactive branch, convergence or a mask rather than proof of a broken graph. Compare a layer with its own earlier distribution and the architecture's expected scale.

For TensorBoard, add_histogram(f"weights/{name}", p.detach(), update) and, when present, the detached gradient expose distributions. The histogram after backward but before zeroing describes that batch, not an epoch average. Hook-based finite checks and .item() on GPU can synchronize and add material overhead; 'manual instrumentation is free' is false. Sample strategically, remove hooks in finally, and retain the failing step identifier.

6 · Become one with the data

Inspect an example before and after each transform with its stable ID and target. For images, check orientation, channels, value range and label-preserving augmentations. For text, decode token IDs, inspect BOS/EOS, target shift and masks. For tabular inputs, check missing values, units, categorical encodings and train-only normalization. Related samples in both train and validation can create a beautiful but misleading curve.

Printing a few batches cannot prove all data correct, but it can falsify assumptions cheaply. A negative example for every data rule—misaligned label, NaN, out-of-range class, all-masked target, wrong shape—makes the audit repeatable. Do not upload private failing examples to a tracker simply to make debugging convenient.

7 · Failure scenarios, corrected

Mismatched labels: shuffling image and label arrays separately corrupts pairs. A fixed batch may still be memorized, so the single-batch test does not reliably catch this. Shuffle joint indices and compare stable IDs to expected labels.

Unstable log-probability: the CPU test's logits (1000,-1000) produce finite CE 2000 while naive log-softmax composition is infinite. Replace the mathematical implementation, not the label or epsilon blindly.

Dropout left active: repeated validation outputs become stochastic. Call model.eval() and use no-grad/inference mode independently. Functional dropout must pass training=self.training; .eval() cannot disable an unrelated call using its default training=True.

8.5 · Tool belt — memory, compiler, profiler and distributed failures

8.5.1 CUDA allocation history — unexecuted reference

PyTorch 2.9 exposes private _record_memory_history and _dump_snapshot APIs. They record allocations managed by the PyTorch CUDA allocator, not every allocation made by every library; NCCL and direct CUDA allocations may be invisible. Inspect allocated versus reserved memory and reference lifetimes. losses.append(loss) can retain graphs, whereas loss.item() stores a number. .detach().clone() makes a copy, not a free alias. Ordinary FP32 Adam moments are two parameter-sized buffers, not three (weights make the common three-tensor subtotal).

# GPU reference — unexecuted
from pathlib import Path
import torch
from torch import nn
 
if not torch.cuda.is_available(): raise RuntimeError("CUDA required")
model = nn.Sequential(nn.Linear(64,64), nn.ReLU(), nn.Linear(64,8)).cuda()
x = torch.randn(16,64,device="cuda")
output = Path("batch6-dl-training-cuda-memory.pickle")
torch.cuda.memory._record_memory_history(max_entries=100_000)
try:
    model(x).square().mean().backward()
finally:
    try: torch.cuda.memory._dump_snapshot(str(output))
    finally: torch.cuda.memory._record_memory_history(enabled=None)
print(output)

Use the official memory visualizer according to its current privacy handling; snapshots contain stack paths and allocation metadata. Avoid uploading confidential artifacts to an unapproved service. This reference has not produced a CUDA snapshot here.

8.5.2 Compiler logs and graph breaks

Set TORCH_LOGS=recompiles,graph_breaks before starting Python, or use torch._logging.set_logs(recompiles=True, graph_breaks=True) after import when TORCH_LOGS was not set: the environment variable takes precedence. This logging API is marked prototype/version-sensitive. Setting an environment variable late is not a reliable reconfiguration API. torch._dynamo.explain(model)(sample_input) reports capture/guard issues. For graph export, use the versioned torch.export.export(module, (sample_input,)) contract for supported models instead of assuming every private Dynamo export signature remains stable.

Recompilation can arise from shapes, dtypes, Python values or guards; 'ten recompiles means varying shapes' is not a theorem. torch.where is not a universally safe replacement for control flow: both expressions may be evaluated, including a mathematically invalid branch. Reduce the eager failure first; then compare compiled outputs/gradients. No compilation run or compiler throughput is claimed in this CPU lesson.

8.5.3 Profiler — a complete CPU trace

The profiler can answer a different question from anomaly detection: where time and allocations go. CPU traces below are executed. GPU profiling would add CUDA activity and synchronization-aware interpretation. Idle GPU intervals suggest inspecting host work, transfers and dependencies; a long kernel is not automatically optimal, and many short kernels are only a candidate for fusion.

import tempfile
from pathlib import Path
import torch
from torch import nn
from torch.profiler import profile, ProfilerActivity, schedule
 
net = nn.Linear(16,4); opt = torch.optim.SGD(net.parameters(), lr=.01)
batch = torch.randn(8,16)
with tempfile.TemporaryDirectory(prefix="batch6-dl-training-profile-") as tmp:
    trace = Path(tmp)/"trace.json"
    with profile(activities=[ProfilerActivity.CPU],
                 schedule=schedule(wait=1,warmup=1,active=3,repeat=1),
                 on_trace_ready=lambda p:p.export_chrome_trace(str(trace))) as prof:
        for _ in range(5):
            opt.zero_grad(); net(batch).square().mean().backward(); opt.step()
            prof.step()
    assert trace.exists() and 'traceEvents' in trace.read_text()
    assert any("linear" in event.key for event in prof.key_averages())
print("CPU scheduled profiler and temporary Chrome trace passed")

A Chrome trace can be inspected in a compatible local viewer. TensorBoard's profiler UI depends on additional tooling; exporting a trace is not proof that a particular browser plugin is installed.

8.5.4 Distributed debugging preview — unexecuted

Record the failing rank, collective sequence, tensor shape/dtype and process-group timeout. NCCL_DEBUG=INFO gives NCCL logs, while TORCH_DISTRIBUTED_DEBUG=DETAIL adds checks with overhead. In the PyTorch ProcessGroupNCCL contract, TORCH_NCCL_ASYNC_ERROR_HANDLING=1 aborts the communicator and tears down the process on detected watchdog errors; it is not a guarantee that every hang is instantly prevented. Flight-recorder dumps require their trace-buffer and timeout-dump settings, not merely one debug flag. Preserve all ranks' logs and coordinate shutdown. No multi-rank job or NCCL runtime has been launched.

9 · Retention scaffold — worked answers

Q1 — What does one-batch overfitting prove? It demonstrates that this configuration can fit the supplied fixed examples. Wrong fixed labels and leaked target features survive; contradictory labels can make zero loss unattainable. The executed shuffled-label control is an explicit counterexample to semantic certification.

Q2 — What does anomaly detection do? It reports failing backward operations with captured forward provenance and raises on backward NaNs when enabled. It is debugging instrumentation with variable overhead, not a universal nonfinite detector. Pair it with finite checks on inputs, logits, loss, gradients and optimizer state.

Q3 — What does loss near log(K) mean? Uniform predictions yield that baseline, but the same mean can arise otherwise. First verify loss/labels/reduction, then graph/optimizer membership and actual updates, then a controlled LR/capacity experiment. Inspect logits rather than infer a distribution from one scalar.

Q4 — Why avoid log(softmax)? A tiny softmax probability may round to zero before logarithm. Use F.cross_entropy(logits, labels) or F.log_softmax(logits, dim=-1) and an appropriate reduction. The exact test gives loss 2000 instead of infinity.

Q5 — What starts CUDA memory recording? torch.cuda.memory._record_memory_history(...); dump with _dump_snapshot, disable in finally, inspect with the official memory visualizer. It sees allocator-managed activity, not every device allocation. API is private/version-sensitive and unexecuted here.

Q6 — Five deliberately broken configurations? (1) Inject NaN input: input finite assertion fails before model; (2) optimize the wrong parameter: finite gradient but unchanged weight; (3) use huge SGD rate on a quadratic: the recurrence diverges; (4) forget zeroing: duplicate backward doubles gradient; (5) pair wrong fixed labels: tiny fit succeeds but ID/label audit fails. Additional GPU-only OOM/compile failures use the respective tools, not anomaly mode for everything. Each diagnosis names the observable and a falsifying control rather than a promised loss-curve shape.

Sources and execution boundary

PyTorch autograd mechanics 2.9, anomaly implementation v2.9.0, CrossEntropyLoss, BCEWithLogitsLoss and CUDA memory source v2.9.0 establish the bounded APIs. Synthetic CPU fitting, negative controls and profiler are executed. No named-model training history, GPU memory trace, distributed incident or application UI is reproduced.

Optional pre-read resources

These inherited resource links are optional, not the primary evidence for this review. Video identity, duration and current availability are not certified.

• A Recipe for Training Neural Networks (Karpathy)
• Common PyTorch Mistakes
• A Recipe for Training Neural Networks
• autograd anomaly detection

Previous: ← DL S018 · GPU + Mixed Precision · Next: DL S020 · Dropout →