Training dynamics: warmup, clipping, loss spikes, large batches, mixed precision

Before this lesson: lesson 02 (gradient descent on $\frac{\lambda}{2}w^2$ and the limit $2/\lambda$), lesson 05 (pre-LN and post-LN, and why post-LN needs warmup), lesson 08 (gradient clipping, first look), lesson 09 (attention logits and softmax saturation) and lesson 11 (the transformer block). Floating-point numbers are in linear-algebra lesson 13. Adam, AdamW and schedules are covered in the optimization track; mixed precision in more depth in the ML-systems track. Time: about 80 minutes, plus the notebook.

Why this matters #

A large training run can be healthy for 80,000 steps and then diverge in 20 steps. The cause is rarely a bug in the model code. It is usually a problem of training dynamics: a step too large for the current curvature, a saturated softmax, an out-of-date optimizer statistic, or a number that does not fit in 16 bits. Every large language model (GPT-3, PaLM, LLaMA, OPT) reports these problems and the fixes in its paper, and the default recipe of nanoGPT contains all of them: warmup, cosine decay, clipping at 1.0, $\beta_2 = 0.95$ and bf16. Interviews ask "your loss spiked at step 50k: what do you do?". In ML for science the same problems appear when a weather or Earth-system foundation model is trained in bf16 on many GPUs.

1. Predict first #

2. The picture #

One parameter. Take gradient descent on $L(w) = \tfrac{\lambda}{2}w^2$ with curvature $\lambda > 0$ and learning rate $\eta$ (lesson 02). One step is $w \leftarrow w - \eta\lambda w = (1 - \eta\lambda)\,w$. It shrinks $|w|$ when $|1 - \eta\lambda| < 1$, that is, when $\eta\lambda < 2$. Above this limit every step makes $|w|$ larger and the loss explodes.

In a network, the role of $\lambda$ is played by the sharpness: the largest eigenvalue of the Hessian of the loss. It is not constant. It is often large at initialization and smaller later. So a learning rate with $\eta\lambda < 2$ at step 10,000 can have $\eta\lambda > 2$ at step 10.

The rest of the lesson follows from this picture and from finite number formats:

3. Definitions and main results #

3.1 Learning-rate warmup and decay #

Let $t$ be the step, $W$ the warmup length, $T$ the total number of steps and $\eta_{\max}$ the peak learning rate. A linear warmup is

(nanoGPT uses $(t + 1)/(W + 1)$, so the first step does not have $\eta = 0$.) After the warmup comes a decay to a small $\eta_{\min}$:

• cosine decay: $\eta(t) = \eta_{\min} + \tfrac12(\eta_{\max} - \eta_{\min})\big(1 + \cos(\pi r)\big)$ with $r = \frac{t - W}{T - W}$;
• linear decay: $\eta(t) = \eta_{\max} - (\eta_{\max} - \eta_{\min})\, r$;
• warmup–stable–decay (WSD): $\eta = \eta_{\max}$ until step $T - D$, then a short decay (linear, for example) over the last $D$ steps, often 10–20% of the run.

WSD has a practical advantage: you can start the decay from any saved checkpoint, so you do not need to fix $T$ before the run starts (Hägele et al. 2024). nanoGPT's defaults are a typical cosine recipe: $\eta_{\max} = 6 \cdot 10^{-4}$, $W = 2{,}000$, cosine decay to $\eta_{\min} = 6 \cdot 10^{-5}$ at step 600,000, AdamW with $\beta_2 = 0.95$, weight decay 0.1 and clipping at 1.0.

Why warmup helps. Three mechanisms are supported by the literature:

1. High curvature at the start. The sharpness is often largest at initialization, as in the picture of section 2. A small learning rate first lets training move to flatter regions, where the later large learning rate is below $2/\lambda$ (Kalra and Barkeshli 2024). In post-LN transformers the gradients near the output are large at initialization (lesson 05, Xiong et al. 2020), which makes this worse.
2. Adam's statistics are noisy at the start. In the first step Adam's bias-corrected moments are $\hat m_1 = g_1$ and $\hat v_1 = g_1^2$, so the update $\eta\,\hat m_1/\sqrt{\hat v_1} = \eta\,g_1/|g_1|$ has size $\eta$ in every coordinate, whatever the gradient. In the next few steps $\hat v$ still averages very few squared gradients, so the per-coordinate step size $\eta/\sqrt{\hat v}$ is noisy and can be very large (Liu et al. 2020, RAdam).
3. Large batches. The learning rate of a large batch (section 3.4) is too large while the network changes fast, which happens at the start (Goyal et al. 2017).

A longer warmup also makes the final loss less sensitive to the choice of $\eta_{\max}$, more so for larger models (Wortsman et al. 2023).

3.2 Global-norm gradient clipping #

Let $g = (g^{(1)}, \dots, g^{(m)})$ be the gradients of all $m$ parameter tensors, seen as one long vector, with the global norm $\|g\| = \big(\sum_i \|g^{(i)}\|_2^2\big)^{1/2}$. For a threshold $\tau > 0$, global-norm clipping is

Result 1. $\operatorname{clip}_\tau(g)$ has the same direction as $g$, and $\|\operatorname{clip}_\tau(g)\| = \min(\|g\|, \tau)$.

Proof.

1. $\operatorname{clip}_\tau(g) = \alpha g$ with $\alpha = \min(1, \tau/\|g\|) > 0$. A positive multiple keeps the direction.
2. If $\|g\| \le \tau$, then $\alpha = 1$ and the norm is $\|g\|$.
3. If $\|g\| > \tau$, then $\alpha = \tau/\|g\|$ and the norm is $\alpha\|g\| = \tau$. $\blacksquare$

Lesson 08 has a figure that compares clipping by norm with clipping each entry. The details that matter in practice:

• PyTorch (torch.nn.utils.clip_grad_norm_) uses the factor $\tau/(\|g\| + 10^{-6})$, clamped at 1, and returns the norm before clipping. Log it: a rising gradient norm is often the first sign of a coming spike.
• Global, not per tensor or per entry. Clipping each tensor separately, or each entry (clip_grad_value_), changes the direction of the full gradient.
• With SGD the step length is at most $\eta\tau$. With Adam clipping does not cap the step, because Adam divides by $\sqrt{\hat v}$. It limits how much one bad gradient enters $m$ and $v$.
• With fp16 loss scaling (section 3.5) the gradients are multiplied by the scale $S$. Unscale first (scaler.unscale_(optimizer)), then clip.

3.3 Loss spikes and divergence #

A loss spike is a sudden jump of the training loss, after which the loss either comes back or diverges. Each cause below shows up first in one quantity. So log: the gradient norm before clipping, the update size per layer, the largest attention logit, the mean $\log Z$ of the output softmax, and in fp16 the loss scale. The table at the end of this section connects the signals to the causes.

(a) Attention-logit growth

The logit of query $i$ and key $j$ is $x_i^\top W_Q W_K^\top x_j/\sqrt{d_k}$ (lesson 09). It is a product of two weight matrices, so it grows with the square of the weight scale: if $W_Q$ and $W_K$ both grow by a factor $c$, every logit grows by $c^2$. When the weights grow during training, the logits become very large, the attention weights become one-hot (entropy near 0), the softmax saturates and its gradient vanishes, and the loss diverges. This was seen in ViT-22B (Dehghani et al. 2023) and reproduced in small models at high learning rates (Wortsman et al. 2023).

Fix: qk-layernorm. Apply a LayerNorm to each query and each key (per head) before the dot product. With it, small models train stably at learning rates where the plain models diverge (Wortsman et al. 2023). Logit soft-capping is an alternative: $s \leftarrow \kappa \tanh(s/\kappa)$, so $|s| < \kappa$ (Gemma 2 uses $\kappa = 50$ for the attention logits and $\kappa = 30$ for the final logits).

Result 2 (qk-layernorm bounds the logits). Let $\hat q = \operatorname{LN}(q)$ and $\hat k = \operatorname{LN}(k)$ in $\mathbb{R}^{d_k}$, without gain and bias. Then $|\hat q^\top \hat k|/\sqrt{d_k} \le \sqrt{d_k}$, whatever the size of $W_Q$ and $W_K$.

Proof.

1. $\hat q_i = (q_i - \mu)/\sqrt{\sigma^2 + \varepsilon}$, where $\mu$ and $\sigma^2$ are the mean and the variance of the entries of $q$. So $\frac{1}{d_k}\sum_i \hat q_i^2 = \sigma^2/(\sigma^2 + \varepsilon) \le 1$, that is, $\|\hat q\|^2 \le d_k$. The same holds for $\hat k$.
2. Cauchy–Schwarz: $|\hat q^\top \hat k| \le \|\hat q\|\,\|\hat k\| \le \sqrt{d_k}\sqrt{d_k} = d_k$.
3. Divide by $\sqrt{d_k}$. $\blacksquare$

With learned gains $\gamma_q, \gamma_k$ (no biases), the bound is multiplied by $\max_i|\gamma_{q,i}| \cdot \max_i|\gamma_{k,i}|$. For $d_k = 128$ the bound is $\sqrt{128} \approx 11.3$. The bound is reached only when $\hat q$ and $\hat k$ are parallel; typical logits are much smaller.

(b) Output-logit drift

Without weight decay, the output logits of a language model can drift to very large negative values late in training, and the loss diverges (Wortsman et al. 2023). The reason is a flat direction of the cross-entropy.

Result 3 (cross-entropy has a flat direction). For logits $s \in \mathbb{R}^V$, a label $y$, $Z = \sum_j e^{s_j}$ and $p = \operatorname{softmax}(s)$, the loss is $\ell(s) = \log Z - s_y$. Then (i) $\ell(s + a\mathbf{1}) = \ell(s)$ for every $a \in \mathbb{R}$, and (ii) $\nabla_s \ell = p - e_y$, whose entries sum to 0.

Proof.

1. $\log \sum_j e^{s_j + a} = a + \log Z$, and $s_y + a$ also grows by $a$. The difference is unchanged.
2. $\partial \log Z/\partial s_j = e^{s_j}/Z = p_j$ and $\partial s_y/\partial s_j = \delta_{jy}$, so $\nabla_s \ell = p - e_y$. Its entries sum to $\sum_j p_j - 1 = 0$. $\blacksquare$

So nothing in the loss pulls $\log Z$ back when all logits drift together: the gradient has no component along $\mathbf{1}$. Fix: the z-loss (PaLM, Chowdhery et al. 2022), an extra term $c_z(\log Z)^2$ with $c_z = 10^{-4}$, which keeps $\log Z$ near 0. Its gradient is $2c_z \log Z \cdot p$, because $\partial \log Z/\partial s_j = p_j$. The entries of this gradient sum to $2c_z\log Z$: it pushes all logits down when $\log Z > 0$ and up when $\log Z < 0$. Weight decay also prevents this drift in Wortsman et al.'s experiments.

(c) Adam's second moment is out of date

Adam divides by $\sqrt{v}$, an exponential average of $g^2$ with a memory of about $1/(1 - \beta_2)$ steps: 1,000 steps for $\beta_2 = 0.999$. Suppose some parameters had tiny gradients for a long time, so $v$ is tiny. If they then receive a large gradient, $g/\sqrt{v}$ is huge and the update is far too large. In CLIP training, 28 of 30 loss spikes came 1–8 steps after the RMS of $g/\sqrt{v}$ in the patch-embedding layer jumped far above 1 (Wortsman, Dettmers et al. 2023, "stuck in the past").

Fix: a lower $\beta_2$. With $\beta_2 = 0.95$ (common for LLMs: GPT-3, nanoGPT), $v$ follows the gradients faster. A $\beta_2$ that is too low makes the steps noisy and slows training. PaLM uses $\beta_2 = 1 - k^{-0.8}$ at step $k$. Update clipping, as in AdaFactor, is another fix.

(d) Adam's $\varepsilon$ is too large. When the gradient RMS of a layer falls to the order of $\varepsilon$, the update $m/(\sqrt v + \varepsilon)$ becomes small and the layer stops learning. This is not a spike but a silent stall. Wortsman et al. (2023) found that a much smaller $\varepsilon$ (for example $10^{-15}$ instead of $10^{-8}$) helps large models.

(e) The practical procedure. PaLM 540B had about 20 loss spikes despite gradient clipping. Each time, the authors restarted from a checkpoint about 100 steps before the spike and skipped about 200–500 data batches. The same batches did not cause a spike from a different, earlier checkpoint. So the cause was the combination of specific batches with a specific parameter state, not "bad data" alone. OPT-175B restarted from earlier checkpoints with a lower learning rate and lowered the clipping threshold from 1.0 to 0.3; its divergences came with the fp16 loss scale falling toward 0.

3.4 Large-batch training #

Linear scaling rule (Goyal et al. 2017): when the mini-batch size is multiplied by $k$, multiply the learning rate by $k$ and keep the other hyperparameters.

Result 4 (why the rule works). Let $l(x, w)$ be the loss of example $x$. Compare $k$ SGD steps on mini-batches $B_0, \dots, B_{k-1}$ of size $n$ with learning rate $\eta$, against one step on their union (size $kn$) with learning rate $\hat\eta$.

1. $k$ small steps: $w_{t+k} = w_t - \eta \frac{1}{n}\sum_{j<k}\sum_{x \in B_j} \nabla l(x, w_{t+j})$.
2. One large step: $\hat w_{t+1} = w_t - \hat\eta \frac{1}{kn}\sum_{j<k}\sum_{x \in B_j} \nabla l(x, w_t)$.
3. If $\nabla l(x, w_{t+j}) \approx \nabla l(x, w_t)$ for $j < k$, the two results are equal exactly when $\hat\eta/(kn) = \eta/n$, that is, $\hat\eta = k\eta$. $\blacksquare$

The assumption in step 3 fails in two cases. Early in training the network changes fast, so the gradients at $w_{t+j}$ and $w_t$ differ. The fix is gradual warmup: increase the learning rate linearly from $\eta$ to $k\eta$ over the first epochs (5 epochs in Goyal et al.). With it, ResNet-50 on ImageNet at batch 8,192 reaches 23.74% top-1 validation error, against 23.60% for the batch-256 baseline. Beyond some batch size the accuracy degrades quickly (about 8k for ResNet-50 on ImageNet).

Where the rule stops. Beyond the critical batch size (McCandlish et al. 2018), a larger batch gives little gain per step: the gradient is already accurate, and only the number of steps matters. For Adam a square-root rule $\hat\eta = \sqrt{k}\,\eta$ is a common start (Malladi et al. 2022 derive it).

Other pitfalls (Goyal et al.): normalize the loss by the global batch $kn$, not by the per-worker batch $n$ (an all-reduce sums the gradients of the workers); compute BatchNorm statistics per worker; and weight decay is part of the gradient, so scaling the loss is not the same as scaling the learning rate.

3.5 Mixed precision #

Here machine epsilon $\epsilon$ is the gap between 1 and the next larger number of the format. The number of exponent bits sets the range; the number of fraction bits sets the precision. fp16 also has subnormal numbers down to $2^{-24} \approx 5.96 \cdot 10^{-8}$. bf16 ("brain float") has the range of fp32 but only 8 significant bits; fp16 has more precision but a small range.

Write $\operatorname{ulp}(w)$ ("unit in the last place") for the gap between the format's numbers near $w$.

Result 5 (when an update is lost). Let a format have $p$ fraction bits, and let $w$ be a number of the format with $2^e \le w < 2^{e+1}$. The spacing of the numbers there is $\operatorname{ulp}(w) = 2^{e-p}$. With round-to-nearest, $w + u$ is rounded back to $w$ for every positive update $u < \operatorname{ulp}(w)/2 = 2^{e-p-1}$.

Why. The next number of the format above $w$ is $w + \operatorname{ulp}(w)$. A value between them that is closer to $w$, that is, less than half a spacing above $w$, is rounded to $w$. $\blacksquare$

For $w = 2^e$ and $u > 0$, the update is lost when $u/w < 2^{-(p+1)}$: when the ratio $w/u$ is above 2,048 in fp16 and above 256 in bf16. (Just below a power of two the spacing is half as large, so for a negative update from $w = 2^e$ the threshold is half as large too.)

So a weight stored in 16 bits ignores small updates, and many small updates in a row are all lost. Micikevicius et al. (2018) gave the fixes that define mixed-precision training:

1. fp32 master weights. Keep an fp32 copy of the weights and of the optimizer states. Apply the update in fp32, then cast to 16 bits for the forward and backward passes. With Adam this costs 16 bytes per parameter: 2 (16-bit weights) + 2 (16-bit gradients) + 4 + 4 + 4 (fp32 master weights, $m$, $v$).
2. fp32 accumulation. Dot products multiply in 16 bits but add the products into an fp32 accumulator (GPU tensor cores do this). Large reductions (normalization statistics, softmax sums) run in fp32. A 16-bit accumulator stops growing once it is large compared with the terms, because each new term is below half its spacing (Result 5 again).
3. Loss scaling (fp16 only). Many activation gradients are smaller than fp16's range and become 0. Multiply the loss by a scale $S$ before the backward pass. By the chain rule every gradient is multiplied by $S$, so small values move into the range. Then unscale ($g \leftarrow g/S$, in fp32) before clipping and the update, so no hyperparameter changes. A fixed $S$ must keep $S \cdot \max|g| \le 65{,}504$.

Dynamic loss scaling (torch.amp.GradScaler): start at $S = 2^{16}$. If any gradient is inf or NaN, skip the optimizer step and halve $S$. After 2,000 steps in a row without overflow, double $S$.

bf16 usually needs no loss scaling. It has the 8 exponent bits of fp32, so gradients practically never underflow. Its precision is lower than fp16's, so fp32 master weights and fp32 accumulation matter even more.

Answers to section 1. (1) It stays at $1.0$. Above 1 the fp16 spacing is $2^{-10} \approx 9.8 \cdot 10^{-4}$, and $3 \cdot 10^{-4}$ is less than half of it, so every addition rounds back to 1.0. An fp32 master weight reaches $1.3$. (2) $k = 4096/128 = 32$, so try $\eta = 0.05 \cdot 32 = 1.6$. Not from step 0: ramp up from 0.05 to 1.6 over the first few epochs (gradual warmup), because the assumption of Result 4 fails early in training. (3) No. Result 3: the cross-entropy is unchanged by $s \mapsto s + a\mathbf{1}$. Only $\log Z$ grows by 50; the z-loss sees it.

4. Worked example #

(a) Clipping two tensors. The gradients are $(2, -4)$ for $W$ and $(4)$ for $b$, and $\tau = 1.5$.

1. Global norm: $\sqrt{2^2 + 4^2 + 4^2} = \sqrt{36} = 6 > 1.5$.
2. Factor: $1.5/6 = 0.25$.
3. Clipped gradients: $(0.5, -1)$ and $(1)$. The new global norm is $\sqrt{0.25 + 1 + 1} = 1.5$.

Per-tensor clipping would give $W$: $(2, -4) \cdot 1.5/\sqrt{20} = (0.671, -1.342)$ and $b$: $(1.5)$. That is a different direction, with total norm $\sqrt{1.5^2 + 1.5^2} = 1.5\sqrt 2 \approx 2.12$.

(b) Warmup + cosine with nanoGPT's settings. $\eta_{\max} = 6 \cdot 10^{-4}$, $\eta_{\min} = 6 \cdot 10^{-5}$, $W = 2{,}000$, $T = 600{,}000$, warmup $\eta_{\max}(t+1)/(W+1)$.

1. At $t = 999$ (warmup): $6 \cdot 10^{-4} \cdot 1000/2001 = 3.0 \cdot 10^{-4}$.
2. At $t = 301{,}000$: $r = (301{,}000 - 2{,}000)/(600{,}000 - 2{,}000) = 0.5$ and $\cos(\pi/2) = 0$, so $\eta = \eta_{\min} + \tfrac12(\eta_{\max} - \eta_{\min}) = 6 \cdot 10^{-5} + 2.7 \cdot 10^{-4} = 3.3 \cdot 10^{-4}$.

(c) Large batch. Base: batch 128, $\eta = 0.05$. Target: batch 4,096.

1. $k = 4096/128 = 32$, so $\hat\eta = 32 \cdot 0.05 = 1.6$.
2. ImageNet has 1,281,167 training images, so one epoch is $1{,}281{,}167/4{,}096 \approx 313$ steps, and 5 epochs of warmup are about 1,564 steps.
3. The learning rate grows by $(1.6 - 0.05)/1{,}564 \approx 9.9 \cdot 10^{-4}$ per step.

(d) A lost update. fp16, $w = 1.0$, $u = 3 \cdot 10^{-4}$.

1. $w = 2^0$, so $e = 0$; fp16 has $p = 10$. The threshold is $2^{e-p-1} = 2^{-11} \approx 4.88 \cdot 10^{-4}$.
2. $u = 3 \cdot 10^{-4} < 4.88 \cdot 10^{-4}$ (equivalently, $w/u \approx 3{,}333 > 2{,}048$), so $w + u$ rounds to 1.0 at every step.
3. The fp32 master weight reaches $1 + 1000 \cdot 3 \cdot 10^{-4} = 1.3$; cast to fp16 it is 1.2998.

(e) Loss scaling. fp16, $S = 2^{16} = 65{,}536$.

1. A gradient of $2 \cdot 10^{-8}$ is below half of fp16's smallest subnormal ($2^{-25} \approx 2.98 \cdot 10^{-8}$), so it becomes 0.
2. Scaled, it is $2 \cdot 10^{-8} \cdot 65{,}536 = 1.31 \cdot 10^{-3}$, a normal fp16 number. Unscaled in fp32 it is $1.9994 \cdot 10^{-8}$: a relative error of $3 \cdot 10^{-4}$.
3. A gradient of 1.5 becomes $1.5 \cdot 65{,}536 = 98{,}304 > 65{,}504$: inf. The scaler skips this step and sets $S = 32{,}768$.

(f) Logit shift and z-loss. Logits $s = (10, 9, 8)$, label 0.

1. $\log Z = 10 + \log(1 + e^{-1} + e^{-2}) = 10.408$, and $p = (0.665, 0.245, 0.090)$.
2. Loss: $\log Z - s_0 = 10.408 - 10 = 0.408$.
3. Shift by $-10$: $(0, -1, -2)$ gives the same $p$ and the same loss, but $\log Z = 0.408$.
4. z-loss ($c_z = 10^{-4}$): $10^{-4} \cdot 10.408^2 = 0.0108$ before the shift and $10^{-4} \cdot 0.408^2 = 1.7 \cdot 10^{-5}$ after it. Only the z-loss sees the shift.

5. In code #

PyTorch's own functions check the numbers of the worked example:

import torch
import torch.nn.functional as F

torch.manual_seed(0)

# (a) global-norm clipping over two tensors
w, b = torch.nn.Parameter(torch.zeros(2)), torch.nn.Parameter(torch.zeros(1))
w.grad, b.grad = torch.tensor([2.0, -4.0]), torch.tensor([4.0])
total = torch.nn.utils.clip_grad_norm_([w, b], max_norm=1.5)
print(total.item(), [round(v, 4) for v in w.grad.tolist() + b.grad.tolist()])
# 6.0 [0.5, -1.0, 1.0]

# the three formats: machine epsilon, largest value, smallest normal value
for dt in (torch.float16, torch.bfloat16, torch.float32):
    f = torch.finfo(dt)
    print(dt, f"{f.eps:.3g} {f.max:.3g} {f.smallest_normal:.3g}")
# torch.float16 0.000977 6.55e+04 6.1e-05
# torch.bfloat16 0.00781 3.39e+38 1.18e-38
# torch.float32 1.19e-07 3.4e+38 1.18e-38

# (d) lost updates: fp16 weight against an fp32 master weight
w16, w32 = torch.tensor(1.0, dtype=torch.float16), torch.tensor(1.0)
u16 = torch.tensor(3e-4, dtype=torch.float16)
for _ in range(1000):
    w16, w32 = w16 + u16, w32 + 3e-4
print(w16.item(), round(w32.item(), 4))
# 1.0 1.3

# (e) underflow, loss scaling and overflow in fp16
S = 2.0**16
print(torch.tensor(2e-8, dtype=torch.float16).item(),
      f"{torch.tensor(2e-8 * S, dtype=torch.float16).float().item() / S:.4e}",
      torch.tensor(1.5 * S, dtype=torch.float16).item())
# 0.0 1.9994e-08 inf

# GradScaler: an inf gradient skips the step and halves the scale
model = torch.nn.Linear(2, 1)
opt = torch.optim.SGD(model.parameters(), lr=0.1)
scaler = torch.amp.GradScaler("cpu")
before = model.weight.detach().clone()
scaler.scale(model(torch.ones(1, 2)).sum()).backward()
model.weight.grad[0, 0] = float("inf")          # pretend the backward pass overflowed
scaler.step(opt)
scaler.update()
print(scaler.get_scale(), torch.equal(before, model.weight.detach()))
# 32768.0 True

# (f) shift invariance of cross-entropy, and the z-loss
s = torch.tensor([[10.0, 9.0, 8.0]], requires_grad=True)
y = torch.tensor([0])
ce = F.cross_entropy(s, y)
ce.backward()
log_z = torch.logsumexp(s, dim=-1)
print(round(ce.item(), 4), round(F.cross_entropy(s - 10, y).item(), 4),
      round(s.grad.sum().item(), 6), round(log_z.item(), 4), round(1e-4 * log_z.item() ** 2, 4))
# 0.4076 0.4076 0.0 10.4076 0.0108

# attention logits grow with the square of the weight scale; qk-layernorm bounds them
d = 64
x = torch.randn(32, d)
W_q, W_k = torch.randn(d, d) / d**0.5, torch.randn(d, d) / d**0.5
for c in (1.0, 4.0, 16.0):
    q, k = x @ (c * W_q), x @ (c * W_k)
    raw = (q @ k.T / d**0.5).abs().max().item()
    qn, kn = F.layer_norm(q, (d,)), F.layer_norm(k, (d,))
    normed = (qn @ kn.T / d**0.5).abs().max().item()
    print(c, f"{raw:.1f}", f"{normed:.2f}")
# 1.0 4.2 3.51
# 4.0 67.0 3.51
# 16.0 1072.7 3.51

The largest plain logit grows by a factor of 16 each time $c$ grows by 4; the normalized logit stays at 3.51, below the bound $\sqrt{64} = 8$. The notebook has these cells and the exercises.

6. Where it shows up #

• LLM pretraining. AdamW with $\beta_2 = 0.95$, warmup then cosine or WSD decay, global clipping at 1.0, and bf16 with fp32 master weights is the default recipe (nanoGPT, PaLM, LLaMA). The z-loss is used in PaLM; qk-layernorm in ViT-22B and many later models; soft-capping in Gemma 2.
• Data-parallel training on many GPUs. The global batch is (batch per GPU) × (number of GPUs) × (gradient-accumulation steps). When you add GPUs, the linear rule (SGD) or the square-root rule (Adam) gives the start of the learning-rate search. With sharded training (FSDP), the gradient norm for clipping must be reduced over all shards; FSDP has its own clip_grad_norm_ for this.
• Mixture-of-experts models add one more softmax, the router, with its own z-loss (lesson 15).
• ML for science. Weather and Earth-system foundation models are large transformers or graph networks trained in bf16. Physical fields have very different scales (pressure in Pa, humidity near $10^{-3}$), so normalize inputs and targets, or small values lose precision. Fine-tuning on long autoregressive rollouts gives large, rare gradients, where clipping and a lower learning rate help.

7. Common mistakes #

• Clipping before unscaling in fp16. The gradients are still multiplied by $S$, so almost every step is clipped by a huge factor (check yourself 2).
• Per-entry or per-tensor clipping when global-norm clipping is meant. It changes the direction.
• Logging the gradient norm after clipping. It is always at most $\tau$ and hides the problem.
• Scaling the learning rate with the batch but not warming up, or keeping the warmup length in steps fixed when the number of steps per epoch changes.
• Summing the gradients over workers without dividing by the global batch. This multiplies the effective learning rate by the number of workers.
• 16-bit weights with no fp32 master copy. Small updates vanish silently, and the loss stalls instead of diverging.
• No loss scaling with fp16. Small gradients are lost (bf16 does not need it).

8. Check yourself #

1. The global gradient norm is 3 and max_norm is 6. What does clip_grad_norm_ do? What would value clipping at 1 do to the gradient $(1.5, -2, 1.5)$?
   Answer

   Nothing: the factor $6/(3 + 10^{-6})$ is clamped to 1. Value clipping gives $(1, -1, 1)$, which points in a different direction.
2. Why must you call scaler.unscale_(optimizer) before clipping in fp16 training?
   Answer

   The gradients are still multiplied by $S$ (for example $2^{16}$). The clip would compare $S\|g\|$ with $\tau$, clip almost every step by a factor of about $\tau/(S\|g\|)$, and shrink the steps by orders of magnitude.
3. An fp16 weight equals 64. Which positive updates survive the addition?
   Answer

   In $[64, 128)$ the spacing is $64 \cdot 2^{-10} = 1/16$. Only updates larger than half the spacing, $1/32 = 0.03125$, change the stored value. An update of exactly $1/32$ is a tie: round-to-nearest-even keeps 64 (its last significand bit is 0), so it is lost too.
4. $\log Z$ is near 0, and then all output logits drift up by 50. What happens to the cross-entropy, to the z-loss and to the precision of the logits in bf16?
   Answer

   The cross-entropy does not change (Result 3). $\log Z$ grows by 50, so the z-loss grows to about $10^{-4} \cdot 50^2 = 0.25$, and its gradient pushes the logits back down. In bf16, numbers in $[32, 64)$ have spacing $32 \cdot 2^{-7} = 0.25$, so the logit differences, which decide $p$, lose precision.
5. Why does Goyal et al.'s gradual warmup work better than a constant low learning rate for the first 5 epochs?
   Answer

   With the constant warmup, the learning rate jumps from $\eta$ to $k\eta$ after 5 epochs, and the training error spikes. At batch 8k the top-1 validation error is 25.88% with constant warmup against 23.74% with gradual warmup. A gradual ramp never makes a sudden jump.
6. An fp16 run shows the loss scale halving again and again, down to values below 1, and then NaN losses. Is this underflow? What do you do?
   Answer

   No. The scale halves only on overflow, so some gradients are inf at every step: the values are truly huge, which is a divergence. OPT-175B saw the same. Restart from a checkpoint where the scale was healthy, lower the learning rate (and possibly the clipping threshold), and look for logit growth. bf16 removes the overflow, not the divergence.

9. Interview angle #

• "Your loss spiked. What do you do?" — Read the logs of section 3.3 first. Restart from an earlier checkpoint; skip the batches around the spike or lower the learning rate; add qk-layernorm, a z-loss or a lower $\beta_2$ if a signal points to them; change one thing at a time.
• "Write gradient clipping." — One global norm over all tensors; factor $\min(1, \tau/\|g\|)$; keeps the direction; with fp16, unscale first.
• "Batch size ×8. What changes?" — Learning rate ×8 for SGD (about ×$\sqrt 8$ for Adam) with a gradual warmup; loss normalized by the global batch; returns diminish beyond the critical batch size.
• "fp16 or bf16? Why master weights?" — bf16 has fp32's range, so no loss scaling, but fewer significant bits; master weights keep small updates (ratio above 2,048 in fp16, 256 in bf16) from being rounded away.

10. Practice #

The notebook for this lesson runs in Colab. It has five exercises, each with a check cell and a folded solution:

1. Write warmup_cosine (nanoGPT's convention) and wsd, and check them against the values of the worked example.
2. Write global-norm clipping over a list of tensors and check it against clip_grad_norm_.
3. Write half_ulp(w, p) and test it against real fp16 and bf16 additions.
4. Write a dynamic loss scaler and check it step by step against torch.amp.GradScaler.
5. Compute attention logits with and without qk-layernorm and check the quadratic growth and the bound $\sqrt{d_k}$.

11. Watch and read #

• Wortsman et al., Small-scale proxies for large-scale Transformer training instabilities (2023): attention-logit growth and qk-layernorm, output-logit divergence and the z-loss, warmup and the AdamW $\varepsilon$.
• Goyal et al., Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour (2017): the linear scaling rule, gradual warmup and the pitfalls of distributed SGD.
• Micikevicius et al., Mixed Precision Training (2018): fp32 master weights, loss scaling and fp32 accumulation.
• Chowdhery et al., PaLM: Scaling Language Modeling with Pathways (2022), section 5: the z-loss, the $\beta_2$ schedule and how loss spikes were handled.
• Wortsman, Dettmers et al., Stable and low-precision training for large-scale vision-language models (2023): loss spikes from an out-of-date second moment.
• Kalra and Barkeshli, Why Warmup the Learning Rate? Underlying Mechanisms and Improvements (2024); Cohen et al., Gradient Descent on Neural Networks Typically Occurs at the Edge of Stability (2021).
• Hägele et al., Scaling Laws and Compute-Optimal Training Beyond Fixed Training Durations (2024): warmup–stable–decay schedules.
• McCandlish et al., An Empirical Model of Large-Batch Training (2018): the gradient noise scale and the critical batch size.
• Karpathy, nanoGPT train.py: get_lr (warmup + cosine) and the training loop (scaler.unscale_, then clip_grad_norm_, then scaler.step).