PaLM's 540-billion-parameter run had a problem its authors could not fully explain. About 20 times during training, despite gradient clipping, the loss suddenly jumped, sometimes by a lot, at irregular intervals and often late in training. Their fix was blunt. Rewind to a checkpoint about 100 steps before the spike, skip roughly 200 to 500 batches of data, and continue. The spike did not come back. And when they replayed exactly those skipped batches starting from a different checkpoint, there was no spike either (Chowdhery et al. 2022). So the data alone was not to blame. The spikes came from the combination of particular batches with a particular state of the model and the optimizer.

OPT-175B's public logbook tells a similar story, with loss divergences handled by lowering the learning rate and restarting from earlier checkpoints, and gradient clipping tightened from 1.0 to 0.3 along the way (Zhang et al. 2022). Every lab that trains large models has a version of this experience. This section collects what is now understood about why large runs become unstable and the fixes that labs actually use, z-loss, QK-norm, logit soft-capping, clipping, embedding scaling, rewinding, and the signals to watch so you see trouble coming. The deep-Transformer mechanisms (residual scaling, attention entropy collapse, warmup) are in §5.7. Here we focus on what goes wrong during a long run, and on the knobs a pretraining recipe turns.

What a loss spike is, and where spikes come from

A loss spike is a sudden rise in training loss, usually with a jump in the gradient norm, that either recovers within tens to thousands of steps or turns into divergence, where the loss climbs and never comes back. Recoverable spikes still cost compute and can leave lasting damage. Divergence ends the run unless you rewind.

Wortsman et al. (2023) made an important methodological point. Many instabilities seen in huge runs can be reproduced in small models trained at high learning rates, so they can be studied cheaply, and the fixes found at small scale transfer. Their study, together with the reports of PaLM, OPT, OLMo, and others, suggests a short list of causes.

The data row deserves a real example. OLMo 2's team traced many of their spikes to documents containing long runs of repeated $n$-grams, such as a token repeated thousands of times, and removed every document with 32 or more consecutive repeats of an $n$-gram of 1 to 13 tokens (Team OLMo 2024). Some spikes are data bugs.

The optimizer row is easy to see in a simulation. The widget below runs Adam on a noisy quadratic and injects a few "bad batches" whose gradients are many times larger than usual. With a long-memory second moment, the denominator $\sqrt{\hat{\mathbf{v}}}$ cannot keep up with the sudden growth, the update becomes several times its usual size, and the loss spikes. The same batches with $\beta_2 = 0.95$, with gradient clipping, or skipped altogether, cause little or no spike.

z-loss, a leash on the output normalizer

The cross-entropy of a softmax over a vocabulary of $V$ tokens, with logits $\mathbf{z} \in \R^V$ and target token $y$, is

Add the same constant $c$ to every logit, and $\log Z$ grows by $c$ while $z_y$ grows by $c$, so the loss is unchanged. The overall level of the logits is a direction the loss cannot see. Its gradient confirms it. With $\mathbf{p} = \operatorname{softmax}(\mathbf{z})$ and one-hot $\mathbf{y}$,

so no gradient ever pushes the mean of the logits one way or the other. Yet the optimizer moves it anyway, because Adam rescales each coordinate separately and its updates do not preserve that zero sum, and because the logits are produced by shared weights that move for other reasons. Over a long run, $\log Z$ drifts. That matters because large logits are imprecise in 16-bit arithmetic (bfloat16 has a spacing of 0.5 near 100, §8.5), and because large, drifting logits were associated with divergence in Wortsman et al.'s experiments.

z-loss, as used in PaLM (Chowdhery et al. 2022), adds a small penalty on the normalizer,

whose gradient follows from $\partial\log Z/\partial z_v = p_v$,

In plain words, z-loss scales the "push down" part of the gradient by $1 + 2\alpha\log Z$. Its component along $\mathbf{1}$ is $\mathbf{1}^\top(\ldots) = 2\alpha\log Z$, the only force on the logits' overall level, and it always points back toward $\log Z = 0$. Since the cross-entropy is indifferent to the level, a model can satisfy z-loss at no cost in prediction quality, by learning to subtract a constant. Wortsman et al. confirmed that z-loss with $\alpha = 10^{-4}$ removes the output-logit divergence in their small-scale reproductions, and z-loss is part of the OLMo 2 recipe. The same idea keeps routers in mixture-of-experts models well behaved (the router z-loss of §13.3).

Without z-loss, $\log Z$ settles wherever the optimizer happened to push it, here around 6.6. Increasing $\alpha$ pulls it toward zero, all the way to about 0.06 at $\alpha = 0.1$, while the cross-entropy barely moves. (In this short run it rises slightly, from 1.654 to 1.713, because the optimizer spends some of its 800 steps on the shift. With a long run and the tiny $\alpha$ of real recipes, the cost is negligible.)

import numpy as np

def ce_zloss(logits, targets, alpha=1e-4):
    raise NotImplementedError

QK-norm, bounding the attention logits

The attention logit $s_{ij} = \mathbf{q}_i\cdot\mathbf{k}_j/\sqrt{d_k}$, for a head of width $d_k$ with query $\mathbf{q}_i = \mathbf{W}_Q^\top\mathbf{x}_i$ and key $\mathbf{k}_j = \mathbf{W}_K^\top\mathbf{x}_j$, is bilinear in the weights, so if $\mathbf{W}_Q$ and $\mathbf{W}_K$ each grow by a factor $a$, the logits grow by $a^2$. Nothing in the loss stops that growth when a head wants to attend sharply, and once the logits are large, the softmax saturates, its Jacobian vanishes, and the head's gradients become erratic (§4.3, §5.7). Dehghani et al. (2023) hit exactly this while scaling vision Transformers. At around 8 billion parameters, the attention logits grew past 50,000 in magnitude within a couple of thousand steps, the attention weights became almost one-hot, and the loss diverged. Normalizing queries and keys removed the divergence, and their 22-billion-parameter ViT was trained with it from the start. Wortsman et al. reproduced the same growth in small language models at high learning rates.

QK-norm normalizes each head's query and key vectors before the dot product, usually with RMSNorm and a learned gain $\mathbf{g}$ (one per head dimension),

and then $s_{ij} = \hat{\mathbf{q}}_i\cdot\hat{\mathbf{k}}_j/\sqrt{d_k}$. A normalized vector has $\lVert\mathbf{q}/\operatorname{RMS}(\mathbf{q})\rVert = \sqrt{d_k}$ (ignoring $\varepsilon$), so by Cauchy-Schwarz, with gains bounded by $g_{\max}$,

In plain words, the logits can only grow as fast as the learned gains do, which is slowly and deliberately, instead of quadratically with every weight in $\mathbf{W}_Q$ and $\mathbf{W}_K$. With unit gains and $d_k = 128$, no logit can exceed about 11.3. QK-norm is now standard in many recent models, including OLMo 2, Qwen3, and Gemma 3 (which replaced Gemma 2's soft-capping with it). It is applied before rotary position embeddings, since RoPE rotates vectors without changing their norm (§5.5).

QK-norm is not always available. Multi-head latent attention, which compresses keys into a small shared latent to shrink the KV cache, never materializes each head's full key at inference time, so there is no key vector to normalize (§12.5 explains the design). Moonshot hit this with Kimi K2, a 1-trillion-parameter mixture of experts with MLA, trained with Muon (§8.6), and in their experiments its maximum attention logits kept growing until training became unstable. Their fix, QK-Clip, acts on the weights instead of the activations (Kimi Team 2025). After each optimizer step, for every head $h$ compute the largest attention logit $S^h_{\max}$ seen in the batch, a number the forward pass already has. If it exceeds a threshold $\tau$ (they used 100), multiply that head's query and key projections by $\sqrt{\gamma_h}$ each, with $\gamma_h = \tau/S^h_{\max}$, so the head's logits shrink by exactly $\gamma_h$ and its largest logit returns to $\tau$. Heads below the threshold are untouched, and so is the training step itself, since only the weights change after it. K2 was pretrained on 15.5 trillion tokens with no loss spikes, and the cap was active mainly early in training, after which the logits settled below it on their own.

import numpy as np

def qk_norm_scores(Q, K, g_q, g_k, eps=1e-6):
    raise NotImplementedError

Logit soft-capping

A blunter tool bounds the logits directly with a smooth saturating function. Gemma 2 replaced every attention logit and every final output logit $z$ by

with $c = 50$ for attention and $c = 30$ for the final layer (Gemma Team 2024). In plain words, small logits pass through almost unchanged (since $\tanh x \approx x$ near zero), and large ones are squashed smoothly toward $\pm c$, so no logit can ever exceed $c$. The derivative shows the price. For $|z| \gg c$ it goes to zero, so a capped logit stops receiving gradient, which is the point (it cannot keep growing) and also a risk (a head stuck at the cap learns slowly). Soft-capping needs care in fused attention kernels, since it changes the scores between the matrix product and the softmax, which is one reason later models moved to QK-norm.

import numpy as np

def softcap(z, c):
    return np.asarray(z, dtype=float)         # placeholder: no cap

def softcap_ce_grad(z, y, c):
    raise NotImplementedError

Gradient clipping at scale

Nearly every large run clips the global gradient norm at 1.0, rescaling the whole gradient whenever $\lVert\mathbf{g}\rVert > 1$ (§0.6, §6.3). At scale, three facts about clipping are worth knowing.

It protects Adam's memory more than the step. §0.6 showed why an optimizer that normalizes its updates still needs clipping. A single huge gradient barely changes Adam's step on the spike itself, but it sits in $\hat{\mathbf{v}}$ for many steps afterward, shrinking every later update of those coordinates, and with a long $\beta_2$ the first moment can also carry it past a lagging denominator, which is the mechanism of the simulator above. Clipping keeps the spike out of both moments.

If it is active on most steps, it is a hidden learning-rate change. A threshold that is always exceeded rescales every gradient, and with SGD that is just a smaller learning rate. Watch the fraction of steps that are clipped. Near the start of training it is often high, and later it should be rare.

Its threshold is not scale-free. The gradient norm depends on model size and on the loss scale, so 1.0 means different things for different models. OPT-175B tightened its clipping from 1.0 to 0.3 mid-run to tame instability, while DeepSeek-V3 and OLMo 2 kept 1.0.

Embedding scaling and initialization

Initialization choices that are invisible in small models matter at scale. Takase et al. (2023) analyzed spikes in pre-LN Transformers and identified two conditions that keep gradients well behaved, small sublayers (small initialization of each residual branch) and a large shortcut, meaning the residual stream itself, which starts as the embedding, should have entries of standard deviation near 1. With the common initialization of embeddings at standard deviation 0.02, the stream starts tiny, and the first layer norm divides by a tiny RMS, which amplifies gradients. Two fixes satisfy the condition. Multiply the embeddings by $\sqrt{d}$ (Vaswani et al. 2017 did this, and Gemma's code does too), or put a LayerNorm right after the embedding. Both prevented spikes in their experiments, which the authors present as verification of the analysis rather than as new methods.

Recipes differ in the details. OLMo 2 initializes every parameter from a normal distribution with standard deviation 0.02, DeepSeek-V3 uses 0.006, and many models scale the initialization of each residual branch's output projection by $1/\sqrt{2L}$ (§5.7). The general lesson holds across them. Pick the initialization so that activation and gradient scales are controlled at step zero, and check it by logging per-layer RMS for the first few hundred steps. Width-aware initialization rules that keep these scales fixed as models grow are the subject of §9.5.

Rewinding from a checkpoint

When a spike does not recover, the standard procedure is the one PaLM used.

1. Rewind. Restore the model, optimizer state, scheduler step, and data loader position from a checkpoint safely before the spike, typically 100 or more steps earlier, so the state that led to the spike is gone.
2. Skip data. Skip the batches around the spike, a few hundred in PaLM's case. This changes the combination of data and state that triggered it.
3. Change something if it recurs. Lower the learning rate, tighten clipping, lower $\epsilon$, or reset Adam's second moment for the affected layers (Molybog et al. 2023).

This only works if the run can be rewound exactly. The data loader must be deterministic and resumable at any batch, checkpoints must include the optimizer state and random number generator states, and they must be frequent enough that rewinding does not waste much compute. LLaMA 3's 405B run saw 466 job interruptions in a 54-day period, 419 of them unexpected and most caused by hardware, so frequent, fast checkpoints are needed for reasons beyond loss spikes (Dubey et al. 2024, and §10.7).

import numpy as np

def spike_steps(x, window=1000, k=7.0):
    raise NotImplementedError

def spike_score(x, window=1000, k=7.0):
    raise NotImplementedError

def rewind_plan(spike_step, checkpoints, back=100, skip=300):
    raise NotImplementedError

What to monitor

Most instabilities announce themselves before the loss moves, if you are watching the right numbers. A good dashboard for a pretraining run includes these, logged at least every few steps.

In PyTorch

Here are z-loss and QK-norm as they appear in a real model, an attention module with RMS-normalized queries and keys and a language-model loss with a z-loss term, checked against our NumPy functions.

Read more

• Wortsman et al. (2023). Small-scale proxies for large-scale Transformer training instabilities. Attention logit growth, output logit divergence, and their fixes, all reproduced at small scale.
• Chowdhery et al. (2022). PaLM. Section 5 documents the loss spikes, the rewind-and-skip fix, and z-loss.
• Dehghani et al. (2023). Scaling Vision Transformers to 22 Billion Parameters, with Kimi Team (2025), Kimi K2: Open Agentic Intelligence. The divergence that made QK-norm necessary, and QK-Clip for attention that QK-norm cannot reach.
• Gemma Team (2024). Gemma 2. Logit soft-capping in attention and output layers.
• Takase et al. (2023). Spike No More. Why small embeddings cause spikes and how scaling them fixes it.
• Team OLMo (2024). 2 OLMo 2 Furious. A full account of stability changes (QK-norm, reordered norms, z-loss, data filters) and the spike score.
• Molybog et al. (2023). A Theory on Adam Instability in Large-Scale Machine Learning. How a tiny second moment and correlated gradients produce spikes.
• Zhang et al. (2022). OPT and its training logbook. What a large run looks like from the inside, restarts included.

Watch next

▶

▶