Debugging instability and common failure modes

Detecting exploding and vanishing gradients

You want to know early if your Transformer’s gradients are exploding or vanishing, because both will wreck training before loss or metrics make sense.

What exploding / vanishing gradients are and why they matter

• Exploding gradients: gradient magnitudes grow very large (e.g. norms $10^3$–$10^8$).
  • Effect: huge parameter updates, loss jumps up, then NaNs/Infs in params or loss.
• Vanishing gradients: gradient magnitudes shrink towards zero (e.g. norms $10^{-8}$ or smaller).
  • Effect: parameters barely change, loss decreases very slowly or not at all.

In a deep Transformer, gradients are the chained product of many Jacobians (per-layer derivatives). If each layer’s “effective derivative” is > 1 on average, the product can blow up with depth. If it’s < 1, it shrinks to zero.

Residual connections and layer norm help a lot, but you can still get:

• Exploding gradients from:
  • Too-high learning rate
  • Poor init (weights too large)
  • Long sequences + attention instability
  • No or too-large gradient clipping
• Vanishing gradients from:
  • Over-aggressive normalization or regularization
  • Very small learning rate
  • Bad init (weights too small, outputs saturated at near-zero regions)

Monitoring gradient norms

You want to log gradient norms every step or every few steps.

In PyTorch:

pythonimport torch from torch.nn.utils import clip_grad_norm_ def log_grad_norms(model, step, logger=print, max_grad_norm=None): total_norm = 0.0 param_count = 0 for name, p in model.named_parameters(): if p.grad is None: continue param_norm = p.grad.data.norm(2).item() total_norm += param_norm ** 2 param_count += 1 # Optional: log layer-wise norms occasionally # logger(f"step {step} | grad_norm[{name}] = {param_norm:.4e}") total_norm = total_norm ** 0.5 logger(f"step {step} | total_grad_norm = {total_norm:.4e}") if max_grad_norm is not None: # Apply gradient clipping and log clipped norm clip_grad_norm_(model.parameters(), max_grad_norm) logger(f"step {step} | clipped_grad_norm <= {max_grad_norm}")

Usage in a training loop:

pythonfor step, batch in enumerate(train_loader, start=1): optimizer.zero_grad() loss = model(batch).mean() loss.backward() log_grad_norms(model, step, max_grad_norm=1.0) optimizer.step()

How to detect explosions / vanishings in logs

• Exploding:

  • Gradient norms suddenly spike by orders of magnitude between steps (e.g. from 1e1 to 1e5).
  • Loss jumps up, or becomes nan or inf right after a few high-norm steps.
  • Often correlated with a specific training event:
    • LR warmup finishing and LR jumping high
    • Switching schedulers
    • Changing batch size or precision

• Vanishing:

  • Gradient norms are extremely small and stable (e.g. always ~1e-7) even early in training.
  • Loss barely decreases from its initial value for many steps, even though LR is reasonable.
  • You might see that only some layers have near-zero grads (e.g. early or late layers).

Example: detecting an issue

Say your log prints:

step 1 | loss = 8.92 | total_grad_norm = 3.2e+0 step 2 | loss = 8.76 | total_grad_norm = 4.1e+0 step 3 | loss = 8.80 | total_grad_norm = 7.6e+3 step 4 | loss = nan | total_grad_norm = nan

Interpretation:

• Grad norm from ~1e0 to 1e3 in one step → exploding gradients.
• Immediate nan loss → likely overflow in params or activations.
• First places to check:
  • Learning rate too high?
  • Gradient clipping missing or too loose?
  • Any recent code change affecting normalization or scaling?