Pretraining Stability Engineering: Loss-Spike Detection and Checkpoint Restart Strategies

Note: This is a review article (no first-hand experimental data), based on public practices and engineering experience.

Experiment Objective

In day-scale LLM pretraining, the biggest enemy is not slowness but instability: loss suddenly spiking and not recovering, one node OOM invalidating a whole checkpoint batch, gradient explosion blowing out the training parameters. The goal: build the engineering closed loop for pretraining stability — real-time monitoring (loss/gradient norm) → automatic spike detection and attribution → auto-save and checkpoint restart → hyperparameter repair. After reading, you can answer operationally: "training crashed, how do I rescue it, should I roll back, and how do I tell bad data from bad hyperparameters?"

Code Walkthrough

Part 1: Stability monitoring in the training loop

import torch, numpy as np

EMA_ALPHA = 0.99
loss_ema = None
grad_norm_ema = None
spike_ratio = 3.0        # trigger: current loss / EMA baseline

for step, batch in enumerate(train_loader):
    loss, grad_norm = train_step(batch)   # train_step returns loss and grad norm

    # Exponential moving average: build the "normal baseline"
    loss_ema = loss if loss_ema is None else EMA_ALPHA * loss_ema + (1 - EMA_ALPHA) * loss
    grad_norm_ema = grad_norm if grad_norm_ema is None else \
        EMA_ALPHA * grad_norm_ema + (1 - EMA_ALPHA) * grad_norm

    # Spike detection: loss jumps above the ratio vs baseline
    if loss > spike_ratio * loss_ema and step > 500:   # warmup excluded
        print(f"[SPIKE] step={step} loss={loss:.3f} ema={loss_ema:.3f} "
              f"grad_norm={grad_norm:.2f} ema={grad_norm_ema:.2f}")
        save_marker(step, loss, grad_norm)   # record spike scene for attribution

    # Gradient explosion guardrail: clip and cut LR beyond threshold
    if grad_norm > 10 * grad_norm_ema:
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
        lr_scheduler.eta *= 0.5              # halving decay to help recovery

Key points:

1. EMA is a "baseline," not a fixed threshold: pretraining loss trends downward overall; fixed thresholds expire. EMA follows the trend, and loss > 3×EMA is "a deviation from trend."
2. Do not judge spikes during warmup: the first few hundred steps drop fast from the initial value with high variance; detection would false-positive.
3. Gradient norm is the second signal: loss spikes are usually preceded by a gradient-norm surge. Dual signals (loss + grad_norm) separate "bad batch" (loss jumps, gradient normal) from "training divergence" (gradient explodes first).
4. Guardrails should "intervene gently," not "cut hard": clip gradients + lower LR (halving) lets training recover by itself, losing less progress than an immediate checkpoint rollback.

Part 2: Auto-save and checkpoint restart

# Dual insurance: periodic checkpoint + anomaly checkpoint
save_every = 1000                       # periodic save (paired with ch04 storage)
marker_hits = 0

if step % save_every == 0:
    save_checkpoint(step, tag=f'periodic-{step}')
if loss > spike_ratio * loss_ema:       # also save the spike scene
    marker_hits += 1
    if marker_hits <= 3:                # keep at most 3 scenes (avoid disk blowout)
        save_checkpoint(step, tag=f'spike-marker-{step}')

# Anomaly restart: roll back if loss does not recover for N consecutive steps
if loss > spike_ratio * loss_ema for last 200 steps:
    restore_checkpoint(nearest_before_spike)   # roll back to nearest pre-spike periodic
    lr_scheduler.eta *= 0.5
    print("[RECOVER] rolled back to pre-spike checkpoint, LR halved")

Design logic:

• Periodic saves handle "normal rollback": one per 1000 steps; post-spike rollback loses at most 1000 steps.
• Spike-scene saves handle "attribution": weight snapshots at the spiking steps feed post-analysis (is some batch toxic? did some layer's parameters fly?).
• Rollback requires "consecutive," not "single": a single spike may recover by itself; roll back only after 200 steps without recovery, avoiding harm to normal fluctuation.

Principles

Five attribution classes for spikes

Class	Signature	Handling
Bad sample/dirty data	loss single-point jump, gradient normal	record sample IDs, filter in the data pipeline
Gradient explosion	grad_norm surges first	clip + lower LR; check whether LR peak is too high
Learning rate too high	periodic small spikes, loss oscillation	lower LR or change warmup schedule
Numerical instability	precursor to NaN/Inf in specific layers	check AMP loss scale, LayerNorm eps
Hardware/comm jitter	node errors/network alerts alongside	back to ch02 DCGM and network troubleshooting

Attribution order: check gradient norm first (exploded → hyperparameter/numerical), then hardware alerts (yes → environment), and suspect data last — this order avoids spending time in the wrong direction.

Why "auto-restart" must pair with checkpoint storage

Pretraining recovery depends on chapter 04's two-tier checkpoint: rollback only happens at points where both local disk and PFS hold copies. If checkpoints are too sparse (every 10000 steps), rolling back after a spike discards 10000 steps; too frequent and training slows. Engineering balance: periodic saves balance recovery granularity and cost (1000-step level), and anomaly scenes are stored separately (at most 3 copies).

Results

(Estimated, not measured)

Stability measure	Recoverable progress	Cost
Periodic checkpoints only	lose ≤ save interval	save IO overhead
+ spike detection alerting	timely human intervention	monitoring overhead negligible
+ auto guardrails (clip/LR halving)	most spikes self-heal	ineffective for extreme divergence
+ auto rollback	recover nearest stable point	rollback loses 1000 steps

Engineering Takeaways

1. Dual-signal monitoring is the floor: watching only loss misses the window where gradients exploded and loss has not yet jumped; dual EMA on loss + grad_norm is standard.
2. Grade the guardrails instead of one hammer: clip → lower LR → rollback, escalating gently; self-heal when possible — rollback is the last resort.
3. Spike scenes must be reproducible: record the sample indices and weight snapshots of the triggering steps, or "the training died and the cause stays a mystery."
4. Stability engineering, storage, and observability are one system: chapter 04 storage, chapter 02 DCGM, and this article's loss monitoring belong in the same monitoring stack, forming a "detect-attribute-recover" loop.

Next article preview: pretraining stabilized and the inference-side engineering chain is nearly complete — next we return to chapter 13 applications, dissecting "Scientific Agent Landing: From Tool Calling to Scientific Workflow Automation."