DL0201 Pretraining Loss Spike Debugging

Written by

During pretraining of a 100B parameter model on a large GPU cluster, you suddenly see loss spikes. How do you debug this, and is it a data issue, a learning-rate problem, or hardware failure?

Answer

All three causes are real at this scale, and they leave different fingerprints, so the answer is a triage order rather than a guess. Run the checks in order of cost: first the numerics and per-rank logs you already have on disk (NaN or Inf, one data-parallel rank far from the mean loss, XID or ECC events, NCCL timeouts), then the optimizer traces (global gradient norm, Adam update ratio, maximum attention logit, output logit magnitude), and only then the data, which is proven guilty only by replaying the exact batch from the pre-spike checkpoint. The shape of the curve already narrows the space: a single sharp spike that recovers within a few thousand steps is usually a batch-times-model-state interaction, a slow creeping divergence preceded by growing logits is a learning-rate or numerics problem, and an instant jump to NaN with a stalled step time is hardware. The industry-standard mitigation for the first case is the one PaLM reported: rewind to a checkpoint roughly 100 to 200 steps before the spike and skip a few hundred batches, which removes the spike even though the same data is harmless from a different model state. That asymmetry is the single most important fact in this debugging problem, because it means “the data was bad” is almost never the correct conclusion on its own.

(1) Read The Shape First: recoverable spike, creeping divergence, and hard fault are three distinct signatures, and each maps to a different subsystem.
(2) Cheapest Check First: numerics and per-rank divergence cost nothing to inspect, optimizer traces cost one dashboard query, and a batch replay costs about one node-hour.
(3) Never Trust The Global Loss: with 512 data-parallel ranks a single sick rank moves the average by a fraction of a percent, so log per-rank loss and gradient-norm maxima, not just the mean.
(4) Silent Data Corruption Has No NaN: a degrading GPU can return wrong arithmetic with no error flag, so keep a bitwise-reproducibility canary that reruns one batch on two hosts and compares.
(5) Data Is Guilty Only On Replay: rerun batch $b_s$ from the checkpoint at step $s-1$ on a fresh replica; if the spike does not reproduce, the batch is not the cause.
(6) Skip-And-Resume Versus Config Change: rollback is right for isolated spikes, but recurring spikes plus growing attention logits mean you must pay for QK-norm, z-loss, a lower peak learning rate, or fp32 optimizer state.

Two-row by three-column chart grid. Top row shows training loss versus step for three cases: a single sharp spike at step 9000 that recovers to trend, a creeping divergence starting near step 11500 that never returns, and a curve that ends abruptly at step 15000 with a NaN marker. Bottom row shows the corresponding global gradient norm on a log scale with a clipping threshold line at 1.0: a single narrow 15x spike, a gradient norm that creeps up for two thousand steps before the loss moves, and a flat trace that terminates in an infinite value.

Figure 1: The same symptom, three diseases. Case A is a recoverable spike whose gradient norm spikes for a handful of steps and then returns, case B is a creeping divergence in which the gradient norm rises for roughly 2,000 steps before the loss visibly moves, and case C is a hard fault with no precursor at all. Note that gradient clipping at 1.0 caps the applied update but does not hide the raw norm, which is exactly why the raw norm must be logged.

The mechanism behind case B is worth being precise about, because it is the one people misdiagnose as bad data. As training proceeds, attention logits $q_i^\top k_j / \sqrt{d_h}$ can grow without bound; the softmax saturates, one head collapses onto a single position, and the gradient through that head becomes both tiny and extremely sharp. In parallel, the output logits drift so that $\log Z$ wanders, which makes the cross-entropy surface stiff in a direction the optimizer cannot see from the loss value alone. Case A has a different mechanism: in bf16 or when optimizer state is stored in low precision, the Adam second moment for a rarely-activated parameter block (embeddings for rare tokens are the classic example) can decay toward the epsilon floor, after which one moderate gradient produces an update whose effective step size is orders of magnitude larger than intended. The debugging pipeline is therefore always the same: detect → localize → replay → mitigate, with instrumentation decided before the run rather than after the first spike.

Mathematical Formulation:
$\ell_t > \mu_{t-W:t} + 3\,\sigma_{t-W:t}$
$\theta_{t+1} = \theta_t - \eta_t \dfrac{\hat m_t}{\sqrt{\hat v_t} + \epsilon}$
$r_t = \dfrac{|\hat m_t|}{\sqrt{\hat v_t} + \epsilon}$
$r_{\max} \approx \dfrac{1-\beta_1}{\sqrt{1-\beta_2}}$
$s_{ij} = q_i^\top k_j / \sqrt{d_h}$
$\mathcal{L} = \mathcal{L}_{\mathrm{CE}} + \alpha (\log Z)^2$

Where:

$\ell_t$ is the training loss at step $t$ , and $\mu$ and $\sigma$ are the mean and standard deviation over a trailing window of $W$ steps (a few hundred works well); this is the spike detector that should fire an alert, not a human eye on a dashboard.
$\theta_t$ are the parameters, $\eta_t$ the scheduled learning rate, and $\hat m_t$ and $\hat v_t$ the bias-corrected first and second Adam moments with decays $\beta_1$ and $\beta_2$ .
$r_t$ is the update ratio, the per-block quantity to log alongside the gradient norm, because it exposes the actual step size the optimizer is taking rather than the size of the gradient.
$r_{\max}$ is the ratio a single outlier gradient can produce, roughly $0.1/\sqrt{0.05} \approx 0.45$ for $\beta_1 = 0.9$ and $\beta_2 = 0.95$ ; this bound holds only while $\sqrt{\hat v_t}$ dominates $\epsilon$ , and it is epsilon underflow that breaks it.
$s_{ij}$ is the pre-softmax attention logit for query $i$ and key $j$ with head dimension $d_h$ ; tracking $\max_{ij} |s_{ij}|$ per layer is the earliest available warning of the creeping-divergence mode.
$Z = \sum_{v=1}^{V} \exp(z_v)$ is the softmax partition function over vocabulary size $V$ , and $\alpha$ is the z-loss coefficient (commonly $10^{-4}$ ) that pins $\log Z$ near zero at negligible quality cost.

The $r_{\max} \approx 0.45$ bound is the reason “one weird document blew up the run” is usually wrong: with healthy second moments, Adam already damps a single 30x outlier gradient to less than a normal step. What actually hurts is a sequence of correlated large gradients, or a block whose $\hat v_t$ has collapsed so that the denominator is dominated by $\epsilon$ . That distinction determines the fix: correlated gradients are a data-ordering and learning-rate problem, while a collapsed second moment is a precision problem solved by keeping optimizer state in fp32 and raising $\epsilon$ .

Left-to-right triage diagram. A trigger box for a detected loss spike fans out to three ordered check boxes: numerics and per-rank divergence, optimizer signals such as gradient norm and maximum attention logit, and a replay of the offending batch from the previous checkpoint. Each check has a yes branch to an evidence box listing concrete symptoms, and each evidence box points to a verdict box: hardware, learning rate and numerics, or data crossed with model state. Vertical clean arrows connect check one to check two and check two to check three.

Figure 2: The triage order that makes this question answerable in an interview. The three branches are ordered by cost per bit of information, not by likelihood, and the replay step is what separates a genuinely corrupt shard from a batch-times-state interaction that a 400-batch skip will silently fix.

Signal	Data and model state	Learning rate and numerics	Hardware and interconnect
Loss shape	Sharp isolated spike, recovers to trend within a few thousand steps	Repeated spikes of growing size, or a divergence that never recovers	Instant NaN or Inf, or a flat line while step time stalls
Gradient norm	One narrow burst of a few steps, clipping engages briefly	Slow upward drift over thousands of steps, clipping engages permanently	Jumps straight to Inf, or becomes exactly zero on one rank
Per-rank pattern	Elevated on the ranks holding the offending microbatches only	Uniform across all ranks, since the cause is the shared update	One rank or one host is an extreme outlier, or a canary mismatches bitwise
Replay test	Reproduces from the same checkpoint, disappears from a different one	Reproduces from any nearby checkpoint with the same schedule	Does not reproduce on a healthy node with identical inputs
Corroborating logs	Sample dump shows repeated n-grams, a truncated shard, or non-text bytes	Max attention logit and update ratio rising, parameter norm inflating	XID or ECC entries, NCCL timeouts, thermal throttling, link flaps
First mitigation	Rewind about 200 steps, skip 200 to 500 batches, then dedup or drop the shard	Lower peak LR, extend warmup, add QK-norm and z-loss, keep fp32 optimizer state	Fence the host, restore the last checkpoint, rerun on a spare node

The operational context matters for how much of this you should automate. Meta reported 419 unexpected interruptions during Llama 3 405B pretraining, with roughly 78% traced to hardware, which means at 100B scale a fault-driven anomaly is not an exotic event but a weekly one. That argues for treating spike response as a control loop: an automatic detector on the trailing-window rule, automatic checkpoint retention dense enough to rewind 200 steps, automatic host fencing on XID events, and a human decision only for the residual class where the numerics themselves are unstable.

DL0201 Pretraining Loss Spike Debugging

Comments

Leave a Reply Cancel reply

More posts

MSD0064 Notification Ranking System

MSD0063 Image Copyright Detection

MSD0062 Ads Ranking System

MSD0061 Nearby Place Recommendation