Loss Spike Debugging: Locating Training Divergence

Abstract: When loss suddenly explodes or jumps to NaN mid-training, it is the most expensive failure mode in large-scale training — every divergence can waste hours of compute. This article provides a four-layer investigation framework (data / numeric / communication / hardware) with clear criteria and actions for each layer.

1. Classify the Symptom First: Different Spikes, Different Causes

Symptom	Signature	Most likely direction
Instant NaN/Inf	Loss becomes NaN in one step, everything NaN after	Numeric layer: gradient overflow, bad sample, optimizer state corruption
Single spike then recovery	One step explodes, then back to normal	Data layer: anomalous batch; communication layer: transient retransmission
Spike without recovery	Stays elevated or keeps climbing	LR / data distribution shift, corrupted checkpoint resume
Creeping increase	No spike; loss trend up	Data order change (RNG not restored), precision issue, LR policy

The first step is always recording: training logs must visualize step / loss / lr / grad_norm / GPU temp / network retransmits — without that, any investigation is guesswork.

2. Four-Layer Framework (Execute in Order)

2.1 Layer 1: Data — Is the Input Broken?

Actions:
1. Locate the global step of the spike → find that step's batch (data_loader state)
2. Inspect the batch: corrupted samples (broken images/overlong text/empty samples)
3. Check shuffle consistency after checkpoint resume (RNG not restored)
4. Check data loading misalignment (per-rank batch inconsistency in distributed mode)

• Distributed data misalignment is a silent killer: one rank reads a corrupted file, its gradients become garbage, and AllReduce contaminates the whole cluster.
• Verify: grad_norm spike + abnormal per-rank memory fluctuation → add torch.distributed.barrier() and sampling-seed checks in the dataloader.

2.2 Layer 2: Numeric — Did Gradients Explode?

Actions:
1. Enable grad clip first (protect): max_grad_norm 1.0 to start
2. Check grad_norm curve: does it explode *before* the spike?
3. torch.autograd.detect_anomaly() to locate the operator producing NaN
4. Check LR too large? (after spike, test 0.3× LR to see recovery)
5. Check FP16 Loss Scaling: scale collapsed to tiny → gradient underflow

# Locate NaN source (debug only; slow; run a few steps)
with torch.autograd.detect_anomaly():
    loss.backward()

# Production: record grad_norm and NaN flags
grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
if not torch.isfinite(grad_norm):
    logger.error(f"NaN grad at step {step}")

• BF16 normally cannot underflow; if BF16 training produces NaN, suspect data or operator implementation before precision itself.
• Inf inputs: special-token overflow, abnormal embedding initialization can also produce Inf.

2.3 Layer 3: Communication — Did Sync Break?

Actions:
1. Re-run with NCCL_DEBUG=INFO; check for retransmits/timeouts (TME errors) at the spike time
2. Switch side: ethtool -S for pause frames/drops rising at the same time
3. Confirm communication ops match the parallelism strategy (PP send/recv misalignment)
4. Check whether different NCCL versions are mixed across images

• Communication-layer failures move everyone at once: all ranks' loss jumps together, with network counters rising.
• Distinguish from data layer: data misalignment is "per-rank abnormal"; communication is "all-ranks simultaneous".

2.4 Layer 4: Hardware — Did a Machine Die?

Actions:
1. nvidia-smi -q -d ECC: ECC counts growing around the spike time?
2. dmesg / nvidia-smi -q -d XID: Xid 79/45/63 etc.
3. Temperature/power curves: any node overheating and throttling (throughput drops but loss shouldn't jump)
4. IB links: ibstatus rate negotiation, SM lid changes (link flapping)

• Hardware faults are node-correlated: the spike only appears when training runs on a specific node (re-map ranks off that node and it stops).
• Action: isolate the node first — never keep training while investigating; a faulty node keeps polluting global gradients.

3. Rapid Response Matrix

Symptom	Protect	Root cause confirmation	Recovery
Instant NaN	Stop training, preserve state	detect_anomaly + grad_norm curve	Fix root cause; resume from nearest good checkpoint
Single spike	Record, observe 50 more steps	Recurrence + network counters	No recurrence: continue; recurrence: data layer
Spike without recovery	Test 0.3× LR	Does it improve?	Improved = LR issue; no = data/checkpoint
Creeping after resume	Verify RNG restoration	Compare with uninterrupted loss curve	Fix dataloader/RNG resume logic

4. Production Case: The Real Culprit Behind Repeated NaN

Symptom: 1000-GPU LLM training produced global NaN every 2–3 days; restarts fixed it for about a day, then it returned.

Investigation:

Step	Action	Finding
1	grad_norm before NaN	Creeping up over 10 steps, then exploded
2	detect_anomaly	Pointed to an attention-layer matmul outputting Inf
3	Data layer check	The batch contained an overlong text sample (abnormal token count)
4	Reproduction	Feeding that sample alone on a single GPU reproduces NaN
5	Fix	Data pipeline: length filter + max_length truncation + sampling dedup

Root cause: the upstream cleaning pipeline let a corrupted overlong sample through; attention compute grows quadratically with sequence length and exceeded the numeric range. Data-layer problems masquerading as numeric-layer faults is the most common misdiagnosis — which is why the four-layer framework starts with data.

5. Prevention Checklist

• [ ] Training logs include step/loss/lr/grad_norm/temp/network counters, visualized
• [ ] Data pipeline has length filters, empty-sample filters, dedup validation
• [ ] Distributed sampling RNG seed verified; data order consistent after resume
• [ ] Grad clip enabled (1.0 to start), grad_norm monitoring with alerts
• [ ] ECC/Xid monitoring into alerting; faulty nodes auto-isolated
• [ ] All resumes start from the "nearest good checkpoint" (with checksums)
• [ ] Periodic short "healthy runs" (pure compute, no communication/data change) as control

Further reading: Mixed Precision Training & Checkpointing (stability and recovery foundations), NCCL Tuning in Practice (deep communication-layer troubleshooting).