LLM Training Debugging, in depth: expected curves, a symptom triage table, the shrink-the-problem ladder, data and NaN localisation, gradient flow and distributed-only bugs

What a healthy run predicts

You cannot spot a wrong curve without a prediction of the right one, and several predictions follow from first principles.

• Step-zero loss is about ln(V). A freshly initialised language model predicts a near-uniform distribution over its V-token vocabulary, so cross-entropy starts near ln(V): 10.37 for a 32,000-token vocabulary, 10.82 for 50,257 and 11.76 for 128,256. Starting far above that means logits are too large (an initialisation or scaling bug); starting well below it on a model you think is fresh means weights were loaded from somewhere or labels leak.
• Loss falls fast, then slowly. Within a few hundred steps the model learns unigram and bigram statistics and loss drops by several nats. A curve that stays near ln(V) means no learning signal is arriving.
• Too good is a bug. A loss far below comparable published runs, or one that reaches near zero on real text, almost always means the model can see the answer: unshifted labels, a broken causal mask, or validation data in the training set.
• Gradient norm settles. After warmup the global gradient norm (returned by torch.nn.utils.clip_grad_norm_) usually falls and then drifts slowly. Sustained growth, or clipping on nearly every step, signals instability long before loss shows it.

Triage: symptom to suspect

The table encodes a habit worth more than any row: classify the symptom before you change anything, and change one thing at a time.

Shrink the problem: the debugging ladder

The fastest general technique is to make the problem small. Step one is overfitting a single batch: take eight sequences, turn off dropout and weight decay, and train on them repeatedly. A correct model drives loss toward zero within a few hundred steps. If it cannot, the bug is in the model, loss or optimizer, and no amount of cluster debugging will help.

def overfit_one_batch(model, batch, steps=300, lr=1e-3):
    model.train()
    opt = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=0.0)
    n_opt = sum(p.numel() for g in opt.param_groups for p in g["params"])
    n_req = sum(p.numel() for p in model.parameters() if p.requires_grad)
    assert n_opt == n_req, f"optimizer sees {n_opt} params, model trains {n_req}"
    for step in range(steps):
        loss = model(**batch).loss
        opt.zero_grad(set_to_none=True)
        loss.backward()
        opt.step()
        if step % 50 == 0:
            print(step, round(loss.item(), 4))
    assert loss.item() < 0.1, "model cannot memorise one batch: bug before scale"

Then climb the ladder one rung at a time, comparing each run's first few hundred steps with the rung below on the same seed and data. Mixed precision, fused kernels, data parallelism, gradient accumulation, tensor and pipeline parallelism and expert routing are separate axes; adding them together turns one bug into an unsearchable product of configurations.

Look at the data the model sees

Most training bugs are data bugs, and the cheapest check in this article is to decode what the model actually receives. Pull one batch from the real loader, after collation, packing and masking, and print it as text with its labels and masks aligned.

def show_batch(batch, tok, row=0, n=40):
    ids = batch["input_ids"][row].tolist()
    labels = batch["labels"][row].tolist()
    for i in range(min(n, len(ids))):
        lab = "IGN" if labels[i] == -100 else repr(tok.decode([labels[i]]))
        print(f"{i:4d} in={tok.decode([ids[i]])!r:14} label={lab}")

Read that output with four questions. Shift: with Hugging Face causal-LM models the labels should equal the inputs, because the model shifts them internally; a pipeline that also shifts them trains the model to predict two tokens ahead, which still learns, only worse. With a hand-written loss, label i must be input i+1. Ignore index: padding and prompt tokens you do not want to train on must carry -100, the default ignore_index of CrossEntropyLoss; a pad id used as a label teaches the model to emit padding. Packing: when several documents share a sequence, attention must not cross document boundaries unless you intend it, which requires per-document position ids and a variable-length attention kernel or block-diagonal mask. Special tokens: check that BOS and EOS appear once each where the tokenizer that will be used at inference expects them.

Localising NaN and Inf

When a NaN appears, the loss reports the symptom many layers after the cause. Forward hooks find the first module whose output goes non-finite.

def install_nan_hooks(model):
    def hook(mod, inp, out):
        outs = out if isinstance(out, (tuple, list)) else (out,)
        for o in outs:
            if torch.is_tensor(o) and not torch.isfinite(o).all():
                amax = o.detach().abs().nan_to_num(0, 0, 0).max().item()
                raise RuntimeError(f"non-finite output in {mod._dbg_name}; max finite |x| = {amax:.3g}")
    for name, mod in model.named_modules():
        mod._dbg_name = name
        mod.register_forward_hook(hook)

If the forward pass is clean and the NaN appears in gradients, torch.autograd.set_detect_anomaly(True) records forward stack traces and reports the backward operation that produced it. Both tools slow training badly and hooks can break torch.compile graphs, so use them on a replayed step, not in production: save the batch, RNG state and checkpoint at the failing step and rerun it in isolation.

Precision explains many of these. fp16 has a maximum of 65504, so attention logits or activations in a badly scaled layer overflow; that is why fp16 needs loss scaling. bf16 shares fp32's exponent range and rarely overflows, but its 8 significant bits round small updates away and make sums over long sequences inaccurate. Keep softmax, normalisation statistics, the loss and optimizer state in fp32, and suspect any custom kernel that does not.

Is the gradient reaching everything?

A model can be numerically clean and still not learn because gradients never reach some parameters. Log per-parameter gradient norms and the update-to-weight ratio for the first steps.

@torch.no_grad()
def grad_report(model, lr):
    rows = []
    for name, p in model.named_parameters():
        if not p.requires_grad:
            continue
        if p.grad is None:
            rows.append((name, "NO GRAD", ""))
            continue
        g, w = p.grad.norm().item(), p.norm().item()
        rows.append((name, f"{g:.2e}", f"{lr * g / (w + 1e-12):.1e}"))
    for r in rows:
        print(*r, sep="  ")

A parameter with no gradient is detached from the loss: a branch built with .detach() or under no_grad, a module replaced after the optimizer was created, or an unused output head. Gradient norms that shrink by orders of magnitude from the last layer to the first point to missing residual connections or a misplaced normalisation. The update ratio (learning rate times gradient norm over weight norm, a rough proxy for Adam's step size) is a heuristic; values around 1e-3 are commonly cited as healthy, and values far above or below suggest the learning rate is wrong for that layer. Check also that weight decay is not applied to norms and biases if your recipe excludes them.

Bugs that only exist at scale

Some bugs exist only at scale. The first is rank divergence: data-parallel replicas should hold identical weights after every step, and a nondeterministic op, a rank-local RNG used in the model, or a skipped all-reduce makes them drift apart silently. Compare checksums.

import torch.distributed as dist

@torch.no_grad()
def assert_replicas_agree(model, group=None):
    s = torch.zeros(1, device="cuda", dtype=torch.float64)
    for p in model.parameters():
        s += p.detach().double().sum()
    lo, hi = s.clone(), s.clone()
    dist.all_reduce(lo, op=dist.ReduceOp.MIN, group=group)
    dist.all_reduce(hi, op=dist.ReduceOp.MAX, group=group)
    assert lo.item() == hi.item(), f"replicas diverged: {lo.item()} vs {hi.item()}"

The checksum applies to plain data parallelism; with sharded optimizers or tensor parallelism, compare matching shards within each replica group instead. The second bug is duplicate data: every rank reading the same shard means an N-GPU run sees one GPU's worth of unique data. Log sample ids per rank for a few steps and check they are disjoint.

The third is loss normalisation. Suppose gradient accumulation runs two micro-batches holding 100 and 900 target tokens with mean losses 2.0 and 1.0. Averaging the two means gives 1.5; the true per-token loss is (200 + 900) / 1000 = 1.1. Averaging means overweights short micro-batches, so the result depends on how sequences happen to be grouped and differs from a run with no accumulation. Sum token losses across micro-batches and ranks, then divide by the global token count. Hugging Face Transformers fixed a version of this bug in its Trainer in late 2024.

Compare against a reference, and a worked example

When the curve is merely worse than expected, compare against a trusted implementation of the same architecture, often the Hugging Face reference. Load identical weights into both, run one batch in fp32 with dropout off, and compare outputs layer by layer with forward hooks. The first layer whose maximum absolute difference jumps from about 1e-6 to something large contains the discrepancy: a rotary embedding using a different frequency base, a norm epsilon of 1e-5 against 1e-6, or a fused kernel with a different mask convention. In bf16, expect differences around 1e-2 and look for jumps rather than absolute values.

Worked example. A team training a 1B model sees loss stuck at 10.4 after 2,000 steps with a 32,000-token vocabulary, exactly ln(V). Gradient norms are logged and nonzero, so they suspect the learning rate, which reads 3e-4. The overfit test fails too, which rules out data. The parameter count assertion fires: the optimizer holds 0.9 billion fewer parameters than the model trains. The optimizer was built before the model was wrapped for parallelism and before a weight-tied model rebuilt its layers, so it was stepping stale tensors. Building the optimizer after wrapping fixes it, and the overfit test passes in 150 steps. Twenty minutes of the ladder replaced two days of learning-rate sweeps.

Failure modes of debugging itself

• Changing several things at once. The curve improves and nobody knows why; the original bug is still there. One change per run, on a fixed seed.
• Heisenbugs. Anomaly mode, hooks and synchronising prints change timing and kernel choice, so race-dependent bugs vanish. Replay with the same kernels and add instrumentation outside the step.
• Comparing across noise. Two healthy runs with different seeds differ; decide from repeated seeds, as described in LLM training reproducibility.
• Blaming the model for the machine. Silent data corruption and a flaky link produce spikes that look like optimisation problems; check the node history in GPU hardware faults.
• Debugging memory as numerics. Out-of-memory workarounds such as reduced batch or extra checkpointing change the curve; attribute memory first with LLM memory profiling.

Trade-offs

What to do next

1. Write down the expected step-zero loss, ln(V), and the loss at step 1,000 from a comparable run before launching.
2. Make the overfit-one-batch test, with the parameter-count assertion, a required check for every new model or recipe.
3. Add a debug command that decodes one real batch with labels and masks aligned, and read it once per data-pipeline change.
4. Keep the NaN hooks and batch-replay tooling ready, so a failing step can be rerun in isolation within minutes.
5. Normalise loss by global token count across micro-batches and ranks, and test that accumulation does not change the curve.
6. Run replica checksums periodically and log per-rank sample ids for the first steps of every launch.
7. Climb the ladder one axis at a time whenever a new parallelism or precision feature is enabled.