Training Stability: Loss Spikes, Divergence & Recovery

Table of Contents

1. Learning Objectives
2. Notation
3. Core Intuition
4. Common Training Instabilities
5. Loss Spikes: Causes and Solutions
6. Gradient Clipping
7. Normalization for Stability
8. muP: Maximal Update Parameterization
9. Recovery Strategies
10. Monitoring & Early Detection
11. Common Pitfalls
12. Summary
13. Exercises

Learning Objectives

1. Diagnose the cause of loss spikes from training metrics.
2. Apply gradient clipping with the correct norm.
3. Explain QK-LayerNorm and Z-loss for attention stability.
4. Describe muP and why it enables hyperparameter transfer.
5. Design recovery procedures for diverged training runs.

Notation

• $\|g\|$ — gradient norm
• $c$ — gradient clip threshold
• $\mu P$ — maximal update parameterization

Core Intuition

Large-scale training is a high-wire act: billions of parameters, trillions of tokens, weeks of compute. A single instability event (loss spike, divergence) can waste days of GPU time. Understanding WHY instabilities happen (data quality, learning rate, attention entropy collapse) and having robust PREVENTION (gradient clipping, normalization, muP) and RECOVERY (rollback, LR reduction) strategies is essential for successful large-scale training.

Common Training Instabilities

1. Loss spikes: Sudden increase in loss, followed by recovery or divergence.

• Cause: Batch with outlier data, gradient explosion, numerical overflow.

2. Slow divergence: Loss gradually increases over many steps.

• Cause: Learning rate too high, weight decay too low, training too long.

3. Attention entropy collapse: Attention becomes one-hot (all mass on one token).

• Cause: QK dot products grow unboundedly during training.

4. Embedding instability: Embedding norms grow without bound.

• Cause: Lack of normalization on embeddings; large learning rate.

5. NaN/Inf: Numerical overflow in FP16/BF16.

• Cause: Large activations × large weights exceeding format range.

Loss Spikes: Causes and Solutions

Data-related:

• Corrupt/adversarial examples in training data.
• Sudden domain shift (data shuffling artifact).
• Fix: Better data cleaning; skip batches with unusual loss.

Optimization-related:

• Gradient explosion from unstable dynamics.
• Fix: Gradient clipping (global norm).

Numerical:

• FP16 overflow (values exceed 65504).
• Fix: Use BF16, or dynamic loss scaling.

Attention-related:

• Large QK products → softmax saturation → zero gradients for most tokens.
• Fix: QK-LayerNorm, scaled initialization.

Gradient Clipping

Global norm clipping (standard):

Typical threshold: $c = 1.0$ for LLMs.

Per-parameter clipping: Clip each parameter's gradient independently. Less common; can distort relative gradient magnitudes.

Why global norm: Preserves the DIRECTION of the gradient (only reduces magnitude). Relative relationships between parameter gradients are maintained.

Monitoring: Track gradient norm over time. Healthy training: stable norm. Spikes in norm predict loss spikes by a few steps.

Normalization for Stability

QK-LayerNorm (Dehghani et al., 2023): Apply LayerNorm to queries and keys before dot product:

Prevents QK dot products from growing unboundedly → stable attention throughout training.

Z-loss: Add a small penalty on the log-partition function:

Prevents logits from growing too large (which causes softmax saturation).

RMSNorm: Lighter than LayerNorm (no mean subtraction):

muP: Maximal Update Parameterization

Yang et al. (2022): Standard parameterization makes hyperparameters width-dependent. muP makes them width-INDEPENDENT:

Standard: Optimal LR changes with model width → must retune for every scale.

muP rules:

• Initialize weights: $W \sim \mathcal{N}(0, 1/d_{\text{in}})$ (standard).
• Learning rate for hidden layers: $\eta_{\text{hidden}} = \eta_{\text{base}} / d$ (scales down with width).
• Output layer LR: $\eta_{\text{out}} = \eta_{\text{base}} / d$.

Result: Hyperparameters tuned on a SMALL model (e.g., 100M) transfer directly to large models (70B) without retuning.

Practical impact: Saves millions of dollars in hyperparameter search. Tune LR, batch size, etc. on a 100M proxy, then apply to the production 70B model.

Recovery Strategies

When training diverges:

1. Rollback + LR reduction:

• Load checkpoint from before the instability.
• Reduce LR by 2-5x.
• Resume training.

2. Skip bad data:

• Identify the batch that caused the spike.
• Skip it and continue.
• Only works for data-caused spikes.

3. Gradient norm monitoring + auto-skip:

• If gradient norm exceeds 10x the running average: skip the step.
• Log the skipped step for post-hoc data analysis.

4. Increased weight decay:

• If instability is from growing norms: increase $\lambda$ temporarily.

Monitoring & Early Detection

Key metrics to track:

• Loss (per-step and smoothed).
• Gradient norm (per-layer and global).
• Weight norms (should grow slowly, not explode).
• Attention entropy (should remain moderate, not collapse to 0).
• Learning rate (verify schedule is correct).
• Activation magnitudes (detect overflow risk).

Alert thresholds:

• Loss spike greater than 2x running average → alert.
• Gradient norm greater than 10x baseline → alert.
• Any NaN/Inf in any tensor → STOP.

Common Pitfalls

Pitfall 1. Not using gradient clipping for LLM pre-training. Without it, a single bad batch can send gradients to infinity and corrupt the entire model.

Pitfall 2. Setting gradient clip too low (e.g., 0.1). This over-constrains optimization, making training extremely slow. Use 1.0 as default.

Pitfall 3. Recovering from divergence by continuing training (without rollback). Once parameters enter a bad region, continuing usually makes it worse. Always rollback.

Summary

• Loss spikes: Data, optimization, or numerical causes; prevent with clipping + normalization.
• Gradient clipping (c=1.0): Universal prevention for gradient explosion.
• QK-LayerNorm + Z-loss: Prevent attention instability.
• muP: Width-independent hyperparameters; tune small, transfer to large.
• Recovery: Rollback + LR reduction is the standard procedure.
• Monitor: Gradient norms, attention entropy, weight norms — detect problems early.

Exercises

Exercise 1. A training run shows gradient norm jumping from 0.5 to 50 at step 10,000. List the possible causes and debugging steps.

Exercise 2. Derive why QK dot products grow with training depth (without QK-norm): compute the expected magnitude after $L$ layers of residual attention.

Exercise 3. For muP: if the optimal LR for a 125M model is $6 \times 10^{-4}$, predict the optimal LR for a 1.3B model (10x wider).

Exercise 4. Design a monitoring dashboard for a 2-month pre-training run. Specify all metrics, alert thresholds, and automated recovery actions.

Exercise 5. Compare the cost of: (a) losing 3 days to divergence + recovery, vs (b) using muP to tune hyperparameters on a small proxy model for 1 day.