Core ML

Debugging Training Runs

Run one ordered protocol instead of guessing: overfit a single batch first, check the loss at initialization against chance level, then shapes and masks, then the learning rate, then precision — each step's result decides the next.

Five checks in order, cheapest and most decisive first: overfit one batch, check loss at initialization against chance level, verify shapes and masks, tune the learning rate, then check numerical precision

Why Does This Exist?

A training run's loss sits flat for 500 steps, or explodes to nan at step 3, or slowly falls but the model is useless at the end — and the instinct, every time, is to reach for the most familiar lever: a different learning rate, more data, a different optimizer. Sometimes one of those happens to fix it, and the underlying cause never gets identified, waiting to resurface on the next architecture change. Sometimes none of them work, because the real bug was somewhere the learning rate can't reach — a label misaligned with its input, a mask applied backward, a loss function computing something other than intended.

The fix isn't a longer list of things that might be wrong. It's an order: a small number of checks, cheapest and most diagnostic first, where each result tells you either the specific next thing to look at, or that this layer of the stack is fine and the bug is elsewhere. Spending an afternoon tuning a learning rate before confirming the model can even memorize eight examples is debugging the wrong layer — and it's the single most common way a training bug eats a full day instead of ten minutes.

Think of It Like This

A mechanic's diagnostic order, not a parts-swapping spree

A car won't start, and there's a wrong way to diagnose it: replace the battery, then the starter, then the fuel pump, whatever order parts happen to sit on the shelf, hoping one swap is the fix. Slow, expensive, and it teaches nothing even when a swap works.

A mechanic runs a sequence instead: check whether the battery has voltage — thirty seconds, a meter, decisive either way. If it does, check whether the starter engages. If it does, check fuel delivery. Each check is cheap, conclusive about that one thing, and tells you exactly where to look next rather than which part to guess at buying.

Debugging a training run works the same way: overfitting one batch is the "does the battery have voltage" check — cheap, fast, and if it fails, everything downstream is irrelevant until it's fixed.

How It Actually Works

Step 1 — overfit a single batch

Take eight examples, disable any regularization, and train on nothing else for a few hundred steps. A correctly-wired model should drive the loss on those eight examples to nearly zero — no generalization required, just memorizing a tiny, fixed set. If it can't, the bug is structural: a broken gradient path, a frozen layer, a loss that doesn't actually connect to the parameters being updated. No amount of learning-rate tuning fixes a broken gradient path, which is why this check comes first — cheap to run, and it rules out or confirms a whole category of bug in under a minute.

Step 2 — check the loss at initialization against chance

Before training moves the weights at all, the loss on the very first batch should sit close to what pure random guessing would produce. For $C$ -way classification with cross-entropy, chance level is $-\ln(1/C)$ — about 2.30 for 10 classes, 4.61 for 100. A first-batch loss wildly different from that number, in either direction, points at something concrete: a label indexing error, a broken loss function, or weights initialized in a way that's already accidentally informative. This needs zero training steps and catches a class of bugs step 1 alone wouldn't isolate as clearly.

Step 3 — shapes and masks

Print every tensor's shape at each stage of the forward pass, and check any mask — padding, causal, attention — against what it should select, by eye, on a small concrete example. Shape bugs are usually silent: broadcasting in NumPy or PyTorch frequently produces a plausible-looking wrong answer rather than an error, so a shape mismatch that broadcasts instead of failing runs to completion and just computes something other than intended.

Step 4 — the learning rate

Only after the first three checks pass does the learning rate become the likely suspect. A loss that diverges or oscillates wildly usually means the step size is too large for the local curvature; one that falls implausibly slowly, everything else confirmed correct, usually means it's too small. Gradient clipping and warmup are the standard companions here, not a substitute for finding the right base rate.

Step 5 — numerical precision

Last, because it's the least likely cause and the most annoying to isolate: mixed-precision training can silently underflow or overflow in ways full precision never would, and a loss fine in float32 but nan under bfloat16 is a precision bug, not an optimization one. Checking this before the cheaper four checks wastes far more time than it saves — precision bugs are real but comparatively rare next to the other four.

Show Me the Code

Step 1's overfit-a-single-batch check, and step 2's chance-level comparison, both computed directly on a tiny linearly-separable batch.

1import numpy as np2
3rng = np.random.default_rng(0)4X = rng.normal(size=(8, 4))5true_w = rng.normal(size=(4, 3))6y = np.argmax(X @ true_w, axis=1)  # 8 examples, 3 classes, separable by construction7
8
9def train_step(w: np.ndarray, lr: float = 1.0) -> tuple:10    logits = X @ w11    probs = np.exp(logits) / np.exp(logits).sum(axis=1, keepdims=True)12    grad = X.T @ (probs - np.eye(3)[y]) / len(y)13    loss = -np.mean(np.log(probs[np.arange(len(y)), y] + 1e-9))14    return w - lr * grad, loss15
16
17w = rng.normal(scale=0.1, size=(4, 3))18for step in range(300):19    w, loss = train_step(w)20    if step in (0, 100, 299):21        print(f"step {step:3d}: loss = {loss:.4f}")22print(f"chance level (3 classes): {-np.log(1 / 3):.4f}")23# -> step   0: loss = 1.215624# -> step 100: loss = 0.032625# -> step 299: loss = 0.012326# -> chance level (3 classes): 1.0986

Step 0's loss of 1.22 sits close to the 1.10 chance level — step 2 passes. By step 299 the loss on these same eight examples has collapsed to 0.012 — step 1 passes too. Either number landing far off would have stopped the investigation here, before the learning rate entered the picture at all.

Watch Out For

Jumping straight to the learning rate on any training failure

The learning rate is the most familiar knob, so it's reached for almost by reflex — smaller, with warmup, a different schedule. When the actual bug is a shape mismatch or a mislabeled batch, no amount of learning-rate tuning fixes it, and the time spent sweeping rates isn't spent on the checks that would have found the real problem in minutes. Run steps 1 through 3 first, every time, before touching the optimizer.

Trusting that a shape mismatch would have raised an error

A common assumption is that wrong shapes would crash the code. Broadcasting in NumPy and PyTorch frequently produces a plausible, non-erroring result instead — a tensor silently broadcast along the wrong axis computes something that runs to completion and looks like real output. The only reliable check is printing and verifying shapes explicitly at each stage, not assuming an absent exception means the shapes were correct.

The Quick Version

Debug training failures as an ordered protocol, not a list of causes tried in arbitrary order.
Overfit a single batch first — cheap, fast, and it rules out or confirms a whole category of structural bugs immediately.
Check the loss at initialization against chance level before touching anything else; a mismatch points at labels or the loss function.
Shapes and masks come next, checked explicitly — broadcasting hides shape bugs rather than erroring on them.
The learning rate and numerical precision come last, because they're the most familiar levers but the least likely actual cause once the first three checks pass.