Vol. XI · Optimisation at Scale Plate 5 of 10loss spikes, divergence, rollback, skipping batches, instability, gradient norm, monitoringLE

Loss Spikes.

The loss leaps; the gradient norm leaps first. Roll back, skip the batches, and fix the cause if you can find it.

A loss spike is a sudden rise in training loss — sometimes for a few steps, sometimes followed by a slow recovery, sometimes by divergence. Published logbooks show them in many large runs. PaLM’s 540-billion-parameter run had about 20, despite gradient clipping. Its team found that a spike was not caused by bad data alone: replaying the same batches from an earlier checkpoint did not reproduce it. Spikes come from a particular state of the model meeting particular data.

The standard response is therefore procedural: restart from a checkpoint roughly 100 steps before the spike and skip about 200 to 500 batches — the batches seen just before and during it. The lasting responses are architectural and numerical: the logit controls of plate VII.7, a lower β₂ (plate XI.3), better initialisation, and careful precision (plate XI.6). The workshop shows the recorded runs of plate VII.7 at high learning rates, where the same instabilities appear at small scale.

Below: left, training loss of the selected recorded run against step. Right, its gradient norm on a log scale — jumps in the norm precede and accompany jumps in the loss.

What you will learn
Workshop
Recorded 8-layer runs of plate VII.7, without gradient clipping, 3 000 steps. Left: training loss averaged over each 30-step logging interval. Right: gradient norm at each logged step. A gradient-norm spike is a logged norm more than five times the run’s median.

Signatures and responses

At learning rate 0.03 the plain 8-layer run logs 5 gradient-norm spikes — norms more than five times its median, reaching 25.1 — and at 0.1 it logs 11; the same model with QK-norm at 0.03 logs 0, and its norm never exceeds 0.46. In the loss, averaged over 30-step intervals, the same events appear only as stalls and small bumps: the gradient norm is the earlier and clearer signal. That is why large runs alert on gradient norms and on the largest attention logits, not only on loss (plate XV.6).

PaLM’s response to its roughly 20 spikes was to restart about 100 steps before each one and skip 200 to 500 batches, after which the spike did not recur at that point. Rollback costs the compute since the checkpoint and some data, and it treats the symptom. The runs here show the alternative: with QK-norm the instability that produced the spikes is gone. Kimi K2’s report describes pretraining on 15.5 trillion tokens without a single spike, using a clipping of attention logits built into its optimiser — the same idea as plate VII.7, applied at the scale of a trillion parameters.

Skipping batches is not throwing data away carelessly: the skipped window is a few hundred batches out of hundreds of thousands, and the PaLM team found the same batches harmless when the model met them in a different state.

Terms

Loss spike
A sudden rise in training loss, with or without recovery.
Rollback
Restarting from an earlier checkpoint.
Batch skipping
Leaving out the batches seen just before and during a spike after a rollback.
Gradient-norm alert
A monitor that flags steps whose gradient norm jumps.
Spikes come from a model state meeting particular data: roll back and skip to recover, and remove the instability to prevent them.

Sources

  1. Chowdhery et al., “PaLM: Scaling Language Modeling with Pathways”, JMLR 2023, section 5.1.
  2. Zhang et al., “OPT: Open Pre-trained Transformer Language Models”, arXiv:2205.01068, 2022, and its published logbook.
  3. Wortsman et al., “Small-scale proxies for large-scale Transformer training instabilities”, arXiv:2309.14322, 2023.