Learning Rate Warmup › Diagnosing Warmup Failures

When training fails or degrades, understanding whether the problem is warmup-related is an important diagnostic skill. Warmup failures leave characteristic signatures in the loss curve and optimizer statistics that you can learn to recognize.

The most common symptom of insufficient warmup is a loss spike in the first few hundred steps. The loss begins decreasing normally for a handful of steps (while the learning rate is very small), then suddenly jumps to a value higher than the initial loss. If warmup is too short, this spike happens at the point where the learning rate crosses the instability threshold. After the spike, training may recover and continue to converge (if the spike was mild), or the loss may plateau at a high value, or diverge entirely.

You can distinguish a warmup-related spike from other causes by checking the timing. A warmup spike will occur at approximately $t \approx T_{\text{warmup}}$ or shortly after, when the learning rate first reaches values high enough to cause instability. A spike at a later point in training (say, at step 50,000 in a 100,000 step run) is more likely due to a bad batch, a numerical overflow in the model, or gradient clipping interaction.

A useful diagnostic is to monitor the gradient norm (the L2 norm of all gradient tensors concatenated) throughout training. Under healthy training, the gradient norm should start relatively high, decline during the warmup phase as the model begins to learn structure, and then stabilize. If the gradient norm spikes sharply during early training, this indicates that warmup was not sufficient to prevent destabilizing updates.

Many training frameworks log gradient norms automatically, and plotting the gradient norm against the learning rate schedule gives you a direct view of the interaction between warmup dynamics and gradient behavior.

A more targeted diagnostic is to track the second-moment estimate $v_t$ for a representative set of parameters. The second-moment estimate should start near zero, grow rapidly in the first few hundred steps, and then stabilize as it converges to its long-run estimate. If you observe $v_t$ growing erratically or oscillating in early training, this is a sign that the learning rate was too high before the second-moment estimate had stabilized, confirming that warmup was too short.

In practice, you can log optimizer.state[param]['exp_avg_sq'] for a few parameters in PyTorch to inspect the second-moment estimates directly. Watching these values across the warmup phase gives you direct evidence of whether the optimizer statistics have stabilized.

If a training run has already failed due to insufficient warmup, the options depend on how far training has progressed. If you catch the failure early (within the first 1-2% of training), the safest option is to restart from scratch with a longer warmup. If the model has a checkpoint from before the spike, you can restart from that checkpoint with a lower peak learning rate or longer warmup.

If the failure occurs after significant training, a common practice is to load the checkpoint from before the spike and continue training with a short warmup of a few hundred steps. This "mini-warmup" lets the optimizer re-stabilize after the disruption. This approach works because the model's weights are much better initialized at this point than they were at step 0, so the optimizer only needs a short period to rebuild reliable moment estimates.