DL S019 · Debugging — nan, inf, and Overfitting a Single Batch

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

The story — the great Llama-2 training crash of 2023

In July 2023, Meta's Llama-2 pretraining run at the 70B scale crashed 45 times across its 1.7M-GPU-hour training. Some crashes were hardware (a rack died). Most were software — subtle numerical bugs that took the training loop from loss: 2.31 to loss: nan between one step and the next. The postmortem section of the paper describes the workflow: each NaN was traced to a specific op in a specific layer by leaving set_detect_anomaly gates in production code, dumping the offending activation, and manually inspecting it. Meta ended up building a whole internal tooling suite around it. The best-known techniques — z-loss to prevent softmax overflow, query/key clipping to prevent attention scores exploding, careful initialization of the output projection — all came out of the same debugging cycles.

The lesson: debugging isn't the exception in DL, it's the main job. You will spend more time as a working ML engineer figuring out why a loss curve looks weird than you will spend writing model code. The tools in this session are what turn that from an art into a checklist.

Here's a scene that plays out in every ML team: you run the training script, wait ten minutes, come back to loss: NaN on line 5. You look at the model. It looks fine. You look at the data. It looks fine. The loss function looks fine. Everything, individually, is fine. Together, it produces a NaN, and you've got no idea where in the pipeline it started.

Debugging deep learning is not like debugging a webapp. There's no exception with a helpful traceback. The "bug" is a number quietly being 3 instead of 0.03, then propagating through five layers, then exploding on the sixth. By the time you see NaN, the actual bug happened 500 lines ago.

This session gives you a handful of tools that make DL debugging tractable. The most important is a mindset shift: DL bugs are numerical, not logical. You debug them by looking at numbers — activation magnitudes, gradient norms, weight statistics — not by stepping through code. And there's exactly one sanity check that catches 80% of bugs before they cost you a day: overfit a single batch.

Prerequisites

• Session 017 — the training loop that just failed.
• Session 018 — AMP adds its own failure modes.

1 · Rule zero — overfit a single batch

The single most useful sanity check in deep learning:

xb, yb = next(iter(train_loader))
xb, yb = xb.to(device), yb.to(device)
 
for step in range(500):
    optimizer.zero_grad()
    loss = criterion(model(xb), yb)
    loss.backward()
    optimizer.step()
    if step % 50 == 0:
        print(f"step {step}  loss {loss.item():.4f}")

Expected: loss drives to ~0 within a few hundred steps. If it doesn't, your model literally cannot memorize 128 examples. Something is broken and no amount of "more data" or "better hyperparameters" will fix it.

Do this before every serious training run. It takes 30 seconds and saves days.

2 · Loss is NaN — a decision tree

You see nan in the loss print. Now what.

Step 1 — WHEN did it happen?

• Step 0 (before any update): NaN input, NaN weights (uninitialized?), or log(0) in your loss.
• Step 1-5: usually LR too high or a bad batch.
• Step 100+: usually gradient explosion (add clipping — Session 024) or slow-burn numerical issue.

Step 2 — check inputs and labels

print(xb.min(), xb.max(), xb.isnan().any(), xb.isinf().any())
print(yb.min(), yb.max(), yb.isnan().any())

If the input already has NaNs, no model change will help. Fix the data.

Step 3 — turn on anomaly detection

torch.autograd.set_detect_anomaly(True)   # ~2-4× slowdown; turn OFF after

Now when a backward pass produces NaN, PyTorch raises a RuntimeError at the exact op that produced it, with a stack trace pointing at the forward-pass line that constructed that op. This is the single most useful debug flag in PyTorch. Don't leave it on always — it's expensive — but flip it on whenever you're chasing NaN.

Step 4 — common culprits

• log(x) where x reached 0: fix with log(x + 1e-8) or log_softmax instead of log(softmax(...)).
• sqrt(x) at x=0: derivative is inf. Add epsilon.
• softmax on logits with huge positive values: exp overflows. Fix: use F.log_softmax or F.cross_entropy — they subtract the max internally.
• 1 / (var + eps): eps too small for the dtype. See Session 018 §7.
• Division by an integer that is sometimes 0 — happens in custom loss functions.

Step 5 — if still stuck, halve the LR and try again

Nine times out of ten, LR was too high. Halve it. Try again. If halving 3 times fixes it, you know it's an optimization issue and can then tune more carefully.

3 · Loss won't budge — a different decision tree

Loss is stuck at ~ln(num_classes) (= 2.303 for 10 classes). Model has converged to "predict uniform distribution".

Check 1 — gradients ARE flowing?

loss.backward()
for name, p in model.named_parameters():
    if p.grad is None:
        print(f"{name}: NO GRAD")
    else:
        print(f"{name}: grad_norm={p.grad.norm().item():.3e}")

If any parameter's grad is None, it's not being reached by backprop. Usually you forgot to requires_grad_(True), or it's inside a container that doesn't register (see S015 war stories). If a grad is 1e-20, the layer is effectively frozen — probably a sigmoid/tanh saturating (dying gradient). Fix with better init (S022) or a different activation.

Check 2 — is LR way too low?

Try LR = 1e-1 or 1e-2 (huge, deliberately). If loss now moves but noisily, you were too low. Binary-search from there.

Check 3 — are labels shuffled correctly?

Shuffle labels randomly and retrain. If loss goes down JUST AS FAST as with the real labels, your model isn't learning the label mapping at all — probably it's collapsing to a constant. Real training should be visibly slower with shuffled labels.

Check 4 — is your model literally trivial?

Print sum(p.numel() for p in model.parameters()). If it's 10, well, that's the problem.

4 · Val diverges from train — the overfit/underfit map

• Both low: shipping. Go home.
• Both high: underfitting. Bigger model, longer training, less regularization, higher LR.
• Train low, val high: overfitting. More data, more regularization (dropout, weight decay), smaller model, early stop.
• Train high, val low (rare): usually a bug — train pipeline broken, augmentation too aggressive, LR way too high, or a data leak in the val set that makes it artificially easy.

Rule: always look at both curves together. Train alone can lie in either direction.

5 · Weight and gradient histograms

writer.add_histogram(name, tensor, epoch) in TensorBoard shows the distribution over time. Two things to watch:

• Weight histograms: should stay Gaussian-ish, mean near 0, std stable. If a layer's std explodes or collapses, either init is wrong (S022) or that layer is dying.
• Gradient histograms: should stay symmetric around 0. If they collapse to a spike at 0, the layer is dead (usually saturating activation or too-small init). If they explode, need clipping.

Do this once, per epoch, for every layer:

for name, p in model.named_parameters():
    writer.add_histogram(f"weights/{name}", p, epoch)
    if p.grad is not None:
        writer.add_histogram(f"grads/{name}", p.grad, epoch)

The output looks like a stack of ridges — one per epoch. Bad training is visually obvious.

6 · The "become one with the data" step

Karpathy's #1 rule. Before you touch the model, print raw batches. Look at them with your eyes.

• For images: matplotlib.pyplot.imshow a batch. Are they upside-down? Colors inverted? Labels aligned?
• For text: print(dataset[0]). Are your token IDs correct? Is the label 0 for one class and 1 for the other?
• For tabular: describe() in pandas. Any NaNs? Outliers? Categorical column stored as numeric?

You will find bugs. Every time. I've been doing this ten years and I still find bugs at this step.

7 · War stories

probs = F.softmax(logits, dim=-1)
loss = -probs[torch.arange(len(y)), y].log().mean()

log_probs = F.log_softmax(logits, dim=-1)
loss = -log_probs[torch.arange(len(y)), y].mean()
# or, better, just use F.cross_entropy(logits, y)

8 · Diagram — the debugging decision tree

Print this. Tape it above your desk.

8.5 · The 2024–2025 tool belt — memory viz, torch.compile logs, profiler

Three tools that shipped or matured in the last two years and that every serious PyTorch user should have muscle memory for.

8.5.1 torch.cuda.memory._record_memory_history() — the OOM detective

Added in PyTorch 2.1, hugely improved in 2.5. Records every allocation, deallocation, and reference to a tensor with a Python stack trace, dumps a .pickle you drop into pytorch.org/memory_viz for an interactive flame graph. This is how you find OOM causes without brute-force bisection.

import torch
torch.cuda.memory._record_memory_history(max_entries=100_000)
 
# ... run one iteration that OOMs or is peaky ...
 
torch.cuda.memory._dump_snapshot("mem.pickle")
torch.cuda.memory._record_memory_history(enabled=None)

Open mem.pickle in the viz. You'll see each tensor's lifetime as a colored bar, the exact stack frame that allocated it, and the peak-memory frame at the top. This tool replaced roughly six previous debugging techniques.

Common OOM causes it uncovers:

• Activations of the whole model kept alive because you accumulated loss tensors instead of .item() (S014 war story).
• KV cache growing unboundedly during long-context inference.
• A stray .detach().clone() that copies rather than aliases.
• Optimizer state (Adam m, v — 3× param count) being much larger than expected.

8.5.2 torch.compile debugging — _dynamo.explain and TORCH_LOGS

When torch.compile(model) is slower than eager, or throws an error, or silently gives different outputs, three tools help:

# 1. Explain what got compiled and where breaks happened
from torch._dynamo import explain
print(explain(model)(sample_input))
 
# 2. See every recompilation as it happens
import os
os.environ["TORCH_LOGS"] = "recompiles,graph_breaks"
# then run your loop; recompiles print to stderr
 
# 3. Dump the actual FX graph
from torch._dynamo import export
fx_graph, guards = export(model, sample_input)
print(fx_graph.code)

Rule of thumb: if you see ≥10 recompiles/epoch, you have a shape-varying input. Either pad to fixed shapes or pass dynamic=True to torch.compile. If you see graph breaks on if x > 0: — that's a data-dependent branch, TorchDynamo has to fall back to eager. Rewrite as torch.where when perf matters.

8.5.3 torch.profiler — the flame graph you need eventually

from torch.profiler import profile, ProfilerActivity, schedule
 
with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
             schedule=schedule(wait=1, warmup=1, active=3),
             on_trace_ready=torch.profiler.tensorboard_trace_handler("./log")) as prof:
    for step, batch in enumerate(loader):
        train_step(batch)
        prof.step()
        if step >= 5: break

Open in TensorBoard's Profile tab (or Perfetto — ui.perfetto.dev) and you'll see every kernel launch, every CUDA sync, every CPU→GPU copy timeline-aligned. The pattern to look for: big idle bars on the GPU timeline = data starvation. Many short kernels back to back = kernel launch overhead, candidate for torch.compile fusion. A single fat kernel = well-optimized inner loop, leave it alone.

8.5.4 Distributed debugging preview

When we hit multi-GPU in M09, three more envvars become essential and I'll mention them here for muscle memory: NCCL_DEBUG=INFO (verbose NCCL logs), TORCH_NCCL_ASYNC_ERROR_HANDLING=1 (crash the process on collective failure instead of hanging), TORCH_DISTRIBUTED_DEBUG=DETAIL (per-rank pretty errors). The 2024 rewrite of ProcessGroupNCCL added flight recorders that make hangs debuggable — see PyTorch Distributed Flight Recorder docs.

9 · 🧠 Retention scaffold

One-line summary (write it in your own words): _______________________________

Spaced review: re-read §1 (overfit one batch) and §8 (the decision tree) in 24 hours. Full session on day 7. Print §8's mermaid tree and pin it above your desk.

Next session (S020, kicks off Module M04): dropout — why randomly zeroing 30% of activations is a good idea (co-adaptation prevention + implicit ensemble), the 3-line implementation, the 1/(1-p) scaling detail that trips up everyone, and why model.eval() is what makes it correct at inference.

Sticky note (keep on your desk): Overfit ONE batch first. NaN → set_detect_anomaly + halve LR. Stuck loss → grad-norms per layer. OOM → _record_memory_history + pytorch.org/memory_viz.

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