Numerical stability & convergence

Parts I–V made the run fast. This lesson keeps it from diverging. At 70B / trillions-of-tokens / three weeks, a single loss spike can torch days of compute — so convergence is active defensive engineering, not luck.

The thing you are actually afraid of

Lesson 23 taught you to diagnose a run that is slow. This one teaches you to survive a run that goes divergent. They are different emergencies. A slow run wastes money at a steady rate; a divergent run can throw away every step since the last good checkpoint in a single update. On a 1024-GPU job that has been grinding for a week, "the loss went to NaN at 3am" is not a curiosity — it is the most expensive event in the lifecycle of the run.

Here is the shape of it. You are watching the loss tick down — 2.41, 2.40, 2.40 — and then one step posts 2.40 → 11.7, or 2.40 → NaN. Two outcomes follow. Sometimes the loss recovers over the next few hundred steps and rejoins its old trajectory, as if nothing happened (a survivable spike). Sometimes it never comes back: the loss pins flat-high, or every subsequent step is NaN because a NaN in the weights poisons every matmul forever (a fatal spike). The whole job is now generating garbage gradients on garbage weights, and it will keep doing so, burning $X00/hour, until a human notices.

What causes a spike? It is almost always one of a small set of dynamic-range failures:

• An outlier batch. A handful of pathological sequences — a 100k-token run of a single repeated byte, a corrupt shard, a document that is all whitespace — produces an enormous loss and therefore an enormous gradient.
• A large gradient on its own. Even on clean data, the tail of the gradient-norm distribution is heavy. Occasionally a step lands a gradient 50× the running norm.
• A precision overflow. fp16 has only ~65k of headroom; a big activation or gradient overflows to inf, and inf − inf = NaN spreads from there. fp8 is even tighter. (This is the lesson-17 dynamic-range story, now as a stability problem.)
• An attention-logit blowup. A query·key dot product grows without bound during training; the pre-softmax logits hit the format's ceiling; the softmax saturates and the gradient dies or explodes.
• A drifting softmax normalizer. Over a 128k-vocab head, the log-sum-exp normalizer log Z wanders away from zero and the cross-entropy arithmetic loses precision.
• Dead optimizer state. A near-zero Adam second moment v in the denominator turns a normal gradient into a giant step.

The rest of the lesson is the standard kit of defenses, each aimed at one of these. None of them is exotic. The expensive mistake is not knowing which one you forgot to turn on — and, in a distributed run, getting one of them subtly wrong so that every rank does something slightly different.

Defense 1 · Global gradient-norm clipping

The workhorse. Before the optimizer step, measure the L2 norm of the entire gradient vector and, if it exceeds a threshold c, scale the whole vector down so its norm is exactly c:

with c ≈ 1.0 for most LM pretraining. Two properties make this the right tool. First, it caps the worst case: no single step can move the weights further than η · c in norm, no matter how violent the batch. Second, it preserves direction — it scales the entire gradient by one scalar, so the relative sizes of per-parameter updates are untouched. You are throttling the magnitude of one bad step, not reshaping it.

The word doing all the work is global. ‖g‖₂ is the norm over all parameters at once:

— one number for the whole model, summed across every parameter p and every element i within it.

Written out, the correct distributed clip is a three-line dance:

The AllReduce in the middle is the whole point. Each rank contributes a single scalar (its local sum of squares), gets back the global sum, takes the square root, and now every rank has bit-identical global_norm and applies bit-identical scaling to its own shard. Skip the collective and clip on the local norm instead, and every rank computes a different scale factor. The "model" the optimizer then steps is no longer a single coherent model — each shard has been throttled by a different amount. Nothing crashes. The loss just quietly stops converging, or converges to something worse, and you spend two days bisecting your config before you find it. This is one of the most common silent bugs in a from-scratch trainer; the framework helpers (torch.nn.utils.clip_grad_norm_ under FSDP, Megatron's clip) exist precisely to get the collective right for you.

Cost-wise the clip is cheap: one AllReduce of a single scalar per step. Latency-bound (lesson 02's α term), negligible bytes. You will not see it in the MFU budget. You will absolutely see its absence in the loss curve.

Watching one spike happen, in numbers

Make it concrete. Say the model has been training with a healthy gradient norm hovering around ‖g‖ ≈ 0.8 and a learning rate η = 3 × 10⁻⁴. Each step therefore nudges the weights by roughly η · ‖g‖ ≈ 2.4 × 10⁻⁴ in norm — small, steady, exactly what convergence wants. Now an outlier batch arrives and produces ‖g‖ = 40 (a 50× spike — well inside the tail you will actually see over millions of steps).

Without clipping, that one step moves the weights by η · 40 ≈ 1.2 × 10⁻² in norm — fifty times the usual jump, in a direction set by a single pathological batch. That is enough to fling the parameters out of the basin the optimizer had been carefully descending. The next forward pass runs on those displaced weights, produces an even larger loss, and therefore an even larger gradient; the step after that is larger still. This is the positive-feedback loop that turns one bad batch into a runaway: big step → worse weights → bigger gradient → bigger step. A few iterations of that and an activation overflows to inf, the inf contaminates a matmul, and the loss is NaN for the rest of the run.

With clipping at c = 1.0, the same spike is scaled by min(1, 1.0/40) = 0.025 before the step. The update norm becomes η · 1.0 = 3 × 10⁻⁴ — back in the normal range. The bad batch still nudges the weights in its (bad) direction, but only by a normal-sized amount, which the next few good batches easily absorb. The spike shows up as a one-step blip in the loss and then vanishes. The clip converted a fatal event into a non-event by capping a single scalar. That is the entire value proposition, and it is why c is one of the few hyperparameters you should set conservatively (≈ 1.0) and rarely touch.

Note what clipping does not do: it does not change the direction of the update, only its length. All parameters are scaled by the same 0.025, so the relative shape of the step is identical to the unclipped step — you are stepping the same way, just less far. That direction-preservation is why global-norm clipping is gentle enough to leave on for the entire run without hurting convergence, unlike per-element gradient value clipping (clamping each g_i to [−c, c] independently), which distorts the direction and is rarely used for large LMs.

Defense 2 · The learning-rate schedule

A constant learning rate is a stability liability at both ends of training, for two different reasons. The schedule — warmup, then decay — fixes both.

Warmup: why the first thousand steps are dangerous

Early in training two things are simultaneously true. The model is at a random initialization, far from any good basin, so the gradient direction is noisy and frequently bad. And Adam's second-moment estimate v — the running average of squared gradients that sits in the denominator of the update — has only seen a few batches, so it is a poor estimate of the true gradient scale. A full learning rate multiplied by a confidently-wrong direction with a mis-calibrated adaptive scale is exactly the recipe for a step that launches the weights out of the sensible region and never recovers.

Warmup ramps the LR linearly from (near) zero up to its peak over the first few hundred to few thousand steps:

This buys time. The early steps are tiny, so even bad directions barely move the weights; by the time the LR reaches its peak, the second moments have stabilized and the model sits somewhere reasonable. Typical warmup is 1–3% of total steps (a 500B-token run might warm up over ~2000 steps). Skipping warmup is one of the classic ways to make a large model diverge in the first 100 steps — it is not a luxury, it is load-bearing.

Decay: cosine vs WSD

Late in training you want the opposite: a small LR so the optimizer settles into a minimum instead of bouncing around it. Two schedules dominate.

• Cosine. After warmup, anneal the LR along a half-cosine from η_peak down to a floor (often 0.1·η_peak or 0) over the whole run: η(t) = η_min + ½(η_peak−η_min)(1+cos(π·t/T)). Smooth, well-understood, and the long-time default. Its flaw is that the shape is welded to a fixed total length T — decide to train longer and the curve you already ran is the wrong curve.
• WSD (warmup–stable–decay). Warm up, then hold the LR constant for the bulk of training, then decay sharply over the final ~10–20% of steps. The constant middle is the operationally pleasant part: you can stop, checkpoint, branch, or extend the run from any point in the stable phase without having committed to a horizon, because the LR there does not depend on T. The sharp final decay recovers most of the loss that cosine spends its whole tail collecting. WSD (and its cousins) is increasingly preferred for exactly this checkpoint-/continuation-friendliness at frontier scale.

Defense 3 · Bounding the logits

The softmax is the single most overflow-prone operation in a transformer, and it appears in two load-bearing places: the attention scores and the final vocabulary projection. Both involve an exponential of a dot product, and an exponential of an unbounded input is an overflow waiting to happen.

z-loss: pin the softmax normalizer

Cross-entropy over a vocabulary computes log Z = log Σ_v e^x_v, the log of the softmax normalizer. Mathematically the loss is invariant to adding a constant to all logits, so nothing in the objective stops log Z from drifting — the logits can all creep upward together over training while the predicted distribution stays the same. That drift is harmless in exact arithmetic and dangerous in low precision: a large Z summed over 128k+ vocabulary entries in bf16 (with its 8-bit exponent but only 7 mantissa bits) loses precision in the sum, and in the worst case e^x_v overflows outright.

z-loss adds a tiny auxiliary penalty that gives the model a reason to keep log Z near zero:

with α small (≈ 10⁻⁴). It barely touches the loss value, but it actively pulls the normalizer back to a sane magnitude every step, which keeps the exponentials in range. This is standard in large-vocab LM heads and is essential for the MoE router (lesson 09): the router is a tiny softmax over experts whose logits, if left unconstrained, drift until routing collapses or overflows. A router z-loss is one of the first things you add when MoE training gets unstable.

Attention soft-capping and QK-norm

For the attention logits, two related stabilizers cap the dot products before the softmax sees them:

• Logit soft-capping (Gemma-style): pass the pre-softmax scores through cap · tanh(scores / cap). The tanh squashes any score into (−cap, +cap), so no logit can run away, but small scores pass through almost linearly so you barely perturb normal attention.
• QK-norm: apply a normalization (e.g. RMSNorm) to the query and key vectors before the dot product. This bounds the magnitude of each vector, which bounds their dot product, which keeps the logits in a controlled range without a hard cap. Increasingly the preferred fix because it has no cap hyperparameter to tune and composes cleanly with the rest of the block.

All three — z-loss, soft-capping, QK-norm — are the same idea applied in different spots: do not let a quantity that gets exponentiated grow without bound.

Defense 4 · The mixed-precision interaction

Lesson 17 was about the formats. Here is how each one behaves under the stability lens, because precision choice and spike-proneness are the same conversation.

fp16 → dynamic loss scaling. Multiply the loss by a scale S (start ~2¹⁵) before backward so gradients land in fp16's healthy range instead of underflowing; unscale by S before the optimizer step. The "dynamic" part is itself a stability mechanism: if any gradient overflows to inf, GradScaler detects it, skips the entire step, and halves S; after many clean steps it doubles S back up. That "skip the step on overflow" reflex is the same recovery instinct as Defense 5 — fp16 just bakes it into the scaler.

bf16 → mostly nothing. Its 8-bit exponent matches fp32, so the underflow problem that motivates the scaler disappears. This is why bf16 became the default: you delete the entire loss-scaler apparatus. You still keep the reduction-heavy ops in fp32.

fp8 → the careful one. Smallest range, so it is the most likely to overflow and the hardest to keep stable. You keep per-tensor or per-block scale factors (lesson 17), and you deliberately leave the most numerically sensitive components — the normalization layers, the MoE router, the embedding and LM-head — in bf16 or fp32. fp8 is for the bulk matmuls, not for the parts where a small error becomes a spike.

And underneath all three, the fp32 gradient accumulation rule from lesson 11: when you accumulate gradients over many micro-batches, accumulate in fp32. The accumulator is summed over potentially hundreds of micro-batches; doing that sum in bf16 loses the small contributions exactly as a bf16 vocabulary sum does. The unifying principle across this whole section is one sentence:

Defense 5 · Detection and recovery

Even with every static defense in place, the heavy tail eventually produces a spike. So the final layer is operational, and it ties directly to lesson 26 (fault tolerance): watch every step, and react automatically.

Monitor two scalars on every optimizer step — they cost nothing because you already computed them:

• the training loss, and
• the global gradient norm (you computed it for the clip in Defense 1; log it).

The grad-norm is the better early-warning signal because it spikes before the loss does — a bad batch shows up as an anomalous norm one step before that step's update corrupts the loss. The standard trigger compares the current norm to a running statistic:

with k ≈ 3–10. A median (or other robust statistic) is used rather than a mean precisely because the mean is contaminated by the spikes you are trying to detect. On a trigger, two responses, in increasing order of severity:

1. Skip the step. Throw away this batch's gradient entirely — do not call optimizer.step() — and continue from the same weights with the next batch. The offending data is gone, the weights are untouched, and you have lost one batch of compute. This handles the common case (a single outlier batch) almost for free.
2. Roll back. If the loss has already gone NaN (a step slipped through, or several bad steps compounded), the in-memory weights are poisoned and skipping won't help. Reload the last good checkpoint, and — crucially — skip past the data shard that caused it so you don't immediately re-trigger on the same poison batch. This is exactly lesson 26's checkpoint-and-resume machinery, repurposed from "a GPU died" to "the math died."

Widget · the loss-spike sandbox

A simulated 70B training run. The blue curve is a smooth loss decay; every so often an outlier batch injects a spike. Toggle the three defenses and move the sliders to feel how they interact. With all defenses off and a healthy LR, a single big spike sends the run to NaN and it flatlines — that flat red line is days of compute gone. Turn on clipping (caps the step), warmup (survives the fragile early phase), and z-loss (tames the spike magnitude), and watch the same spikes become survivable wobbles. Try raising the LR with defenses off, then turning them on.

Widget · global-norm clipping as a collective

Per-step gradient norms drawn as bars over a stretch of training. Most steps are well-behaved; a few are outliers (sampled deterministically from the step index, so the picture is stable as you drag). The dashed line is the clip threshold c. Bars above it are clipped down to c (the darker cap shows how much was shaved); bars below pass through untouched. The annotation is the load-bearing systems fact: that single threshold comparison requires every rank to first agree on ‖g‖ via an AllReduce.

Initialization & residual scaling — divergence you can prevent before step 1

Every defense so far reacts to a spike during the run. Initialization is the one lever that decides whether the run is stable before step 0. Init sets the variance of activations and gradients at the very first forward/backward pass; get it wrong and the model diverges before warmup (Defense 2) ever has a chance to save it. The baseline rule is variance-preserving: draw weights with std ∝ 1/√(fan_in) (Xavier/Kaiming). That keeps the activation variance ≈ 1 as a signal passes through a layer, so neither activations nor gradients shrink toward zero or blow up toward overflow as they propagate.

The depth problem: a growing residual stream

One layer being well-conditioned is not enough. In a residual network the residual stream variance grows with depth L, because every block adds its output back into the stream — x ← x + f(x), repeated L times. Variances of independent contributions add, so after L blocks the stream variance scales like L. The fix is to scale each residual-branch output down by 1/√(2L) (the GPT-2 trick) or use DeepNorm, so the accumulated stream variance stays bounded as L grows.

μP: hyperparameters that transfer across width

μP (maximal update parametrization) goes one step further. It parametrizes the init and the per-layer learning rates so that the optimal hyperparameters are invariant to width. You tune LR and init on a small, narrow proxy model, then transfer those exact values to the full-width model without re-tuning. This is how large runs avoid an LR sweep at full scale — sweeping a 70B model is the expense you cannot afford, so you sweep a cheap proxy and let μP carry the result across.

These are the cheapest defenses of all: they cost zero runtime and are decided once, in the config, before the GPUs (lesson 01) ever spin up. A run that is initialized right rarely needs the reactive defenses to fire at all.

Putting the defenses in order

If you are bringing up a new large-model run, this is the order to add the defenses and the failure each one removes:

Notice how cheap the column on the right is. Almost none of this shows up in your MFU. Convergence engineering is not a throughput tax — it is insurance with a near-zero premium and a payout measured in saved GPU-days.