Gradient Clipping in LLM Training: Global Norm and AMP Order

Why architecture matters here

Architecture matters here because gradient explosion is not a rare curiosity — it is a structural property of composing many layers. Backpropagation multiplies Jacobians layer by layer, and when the product of those per-layer factors is greater than one, the gradient grows geometrically as it flows backward toward the early layers. In a deep transformer or a recurrent network unrolled over a long sequence, this compounding means the gradient norm can span many orders of magnitude between an ordinary step and a bad one. Without a cap, the optimizer faithfully applies whatever it is handed, so the single worst batch in an epoch dictates whether the whole run survives.

The cost of an unclipped spike is total and unrecoverable. A normal training step nudges the parameters a small distance; an exploding step can move them so far that activations saturate, the softmax produces infinities, and the loss becomes NaN. Once NaN enters the weights it propagates on the next forward pass and every subsequent gradient is NaN — the model is bricked, and unless you were checkpointing frequently you have lost hours or days of compute. Because the trigger is a rare input, the failure is also non-deterministic and maddening to reproduce: the run is fine for ten thousand steps and then dies at step 10,001 on a batch you cannot easily identify.

Clipping converts this catastrophic, run-ending risk into a bounded, benign one. By capping the update length, you guarantee that no single step can move the parameters more than a fixed distance, no matter how large the raw gradient was. The worst case becomes 'this step made little useful progress' instead of 'this step destroyed the model.' That guarantee is what lets practitioners train at higher learning rates and with larger batches than would otherwise be safe, because the tail risk of an occasional huge gradient has been amputated. In effect, clipping decouples the aggressive average-case step size from the fragile worst-case one.

There is an architectural elegance to why global-norm clipping in particular is the right tool: it respects the geometry of the optimization. The gradient is a direction in a very high-dimensional space, and that direction — where downhill lies — is the valuable information; the magnitude is merely how far to trust it. Scaling the whole vector by one shared factor discards only the magnitude when it is untrustworthy while keeping the direction intact. Contrast this with clip-by-value, which caps each coordinate independently and therefore rotates the update vector, corrupting the direction precisely when the model most needs a faithful descent direction. The choice between the two is a choice about whether you are willing to distort where you are heading, and for most training the answer is emphatically no.

The architecture: every piece explained

The core algorithm is clip-by-global-norm and it has four steps. First, after the backward pass has populated a gradient for every parameter, compute the global L2 norm: square every gradient element, sum all those squares across the entire model, and take the square root. This single scalar summarizes the total magnitude of the update the optimizer is about to apply. Second, compare it to the threshold max_norm. Third, if the norm exceeds the threshold, compute the scale factor max_norm / (norm + eps), which is a number less than one. Fourth, multiply every gradient in the model by that same factor. If the norm is already below the threshold, do nothing. The result is that the post-clip global norm is exactly max_norm on spiky steps and unchanged on normal ones.

The key structural detail is global: the norm is taken over all parameters jointly, not per-layer or per-tensor. This matters because it means the relative sizes of gradients across layers are preserved — a layer that legitimately has larger gradients keeps its larger share after clipping. A per-tensor clip would flatten those relationships and change which layers learn fastest. Because a single shared scalar multiplies everything, the direction of the combined update vector in the full parameter space is exactly preserved; only its overall length changes. This is the property that makes the technique safe to apply on every step: on a normal step it is a no-op, and on a spiky step it is a pure length cap.

The alternative primitive is clip-by-value: independently clamp each gradient component to the range [-v, +v]. This is simpler and needs no global reduction, but it changes the update's direction because different components are clipped by different amounts — a component at +100 and one at +0.5 both survive differently, rotating the vector. Clip-by-value is occasionally used in reinforcement learning or when a specific pathology produces a few enormous coordinates, but for language-model and general deep-learning training, clip-by-norm is the standard because direction fidelity dominates. A related idea is adaptive or auto clipping, where the threshold tracks a running percentile of recent norms rather than being a fixed constant, which removes the need to hand-tune the number at the cost of a little state.

Finally, clipping has a precise place in the training step relative to two other mechanisms. When you use gradient accumulation (summing gradients over several micro-batches before a step), the clip must be applied to the accumulated gradient, once, just before the optimizer step — not to each micro-batch, which would clip a partial gradient. And when you use automatic mixed precision with loss scaling, the gradients in memory are inflated by the loss-scale factor, so they must be unscaled back to their true magnitude before the norm is measured, or the comparison against max_norm is meaningless. Both constraints are about measuring the true, final gradient that will actually be applied.

End-to-end flow

Walk one optimizer step of a transformer trained with mixed precision and gradient accumulation. The training loop runs the forward pass on a micro-batch, multiplies the loss by the loss-scale factor S (a large number like 65536, chosen to push tiny half-precision gradients up into the representable range), and calls backward. The gradients now sitting in each parameter are S times their true value. The loop repeats this for, say, four micro-batches, accumulating the scaled gradients so that the total is S times the true accumulated gradient. Only now is a step imminent.

Before anything touches the weights, the framework unscales: it divides every gradient by S, recovering the true accumulated gradient, and in the same pass checks for infinities or NaNs (a sign the loss scale was too high and produced overflow). If an overflow is detected, the step is skipped entirely and the loss scale is reduced — the inf or NaN gradients are never applied to the weights, even though the clipping call itself still runs on them. Assuming the gradients are finite, the framework now computes the global norm over the unscaled gradients. Suppose the norm comes out to 8.0 and the threshold is 1.0; the scale factor is 1.0 / 8.0 = 0.125, and every gradient is multiplied by 0.125, bringing the global norm down to exactly 1.0 while preserving direction.

Only after unscaling and clipping does the optimizer step run. Adam takes the clipped gradient, updates its first- and second-moment running averages, and applies the parameter update, which is now guaranteed bounded because the gradient feeding it has a norm of at most 1.0. The learning-rate scheduler advances, the gradients are zeroed, and the loop starts the next accumulation cycle. On the vast majority of steps the measured norm is below the threshold — say 0.3 — and the clip is a no-op: the gradient passes through untouched and the model trains at full effective step size. The clip earns its keep only on the rare steps where the norm spikes.

It is worth watching what happens across a spike to see why the run survives. At step 10,000 an unusual batch produces a raw global norm of 400. Without clipping, plain SGD would take a step 400× the size of a normal one. Adam bounds each coordinate’s step by dividing by its second-moment estimate, but the spike still corrupts both moment estimates for many steps, and that instability can push the weights into a region where activations overflow and the next forward pass yields NaN — run over. With clipping at 1.0, that same batch's gradient is scaled by 1/400 down to a norm of exactly 1.0; the optimizer takes an ordinary-sized, correctly-directed step; the loss ticks up slightly because this batch was genuinely hard, and by step 10,010 training has moved on as if nothing happened. The spike was absorbed, not amplified, and the total cost was one slightly-wasted step rather than the entire run.

Gradient clipping in LLM training: a worked example in code

Gradient clipping in LLM training almost always means clipping by global norm: compute one L2 norm over every parameter gradient, and if it exceeds max_norm, multiply all gradients by max_norm / norm. Published LLM recipes commonly use max_norm = 1.0; the GPT-3 paper states it clips the global gradient norm at 1.0, and the original LLaMA paper reports gradient clipping of 1.0. In PyTorch this is one call, torch.nn.utils.clip_grad_norm_, placed after backward() and before optimizer.step().

The formula, with g the concatenation of all parameter gradients:

total_norm = sqrt(sum over all params of sum(g_i ** 2))
scale      = min(1, max_norm / (total_norm + eps))
g_i       <- g_i * scale        # same scalar for every tensor

Runnable demo: norm before and after, direction preserved

This NumPy script builds two gradient tensors, injects one spiky coordinate as a bad batch would, then compares global-norm clipping with per-value clipping. The cosine similarity between the original and clipped gradient shows whether the update direction survived.

import numpy as np

def clip_by_global_norm(grads, max_norm, eps=1e-6):
    total = np.sqrt(sum(float((g ** 2).sum()) for g in grads))
    scale = min(1.0, max_norm / (total + eps))
    return [g * scale for g in grads], total

def flat(gs):
    return np.concatenate([g.ravel() for g in gs])

rng = np.random.default_rng(0)
grads = [rng.normal(0, 1, (4, 4)), rng.normal(0, 1, 4)]
grads[0][0, 0] = 40.0          # one spiky coordinate

clipped, pre = clip_by_global_norm(grads, max_norm=1.0)
post = np.linalg.norm(flat(clipped))
cos = flat(grads) @ flat(clipped) / (np.linalg.norm(flat(grads)) * post)
print(f"global norm before: {pre:.4f}  after: {post:.4f}  cosine: {cos:.6f}")

by_value = [np.clip(g, -1.0, 1.0) for g in grads]
bv = flat(by_value)
cos_v = flat(grads) @ bv / (np.linalg.norm(flat(grads)) * np.linalg.norm(bv))
print(f"clip-by-value norm: {np.linalg.norm(bv):.4f}  cosine: {cos_v:.6f}")

Output (NumPy 2.4):

global norm before: 40.1887  after: 1.0000  cosine: 1.000000
clip-by-value norm: 3.1287  cosine: 0.405382

Global-norm clipping shrinks the norm from 40.19 to 1.0 with cosine 1.0, so the direction is unchanged. Clipping each value to [-1, 1] leaves a norm of 3.13 and a cosine of 0.41: the update now points somewhere quite different, because the spiky coordinate was cut by a factor of 40 while the small ones were untouched.

Where it goes in a PyTorch training step

scaler.scale(loss).backward()          # AMP: grads are multiplied by the loss scale
scaler.unscale_(optimizer)             # restore true magnitudes first
pre_clip = torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
scaler.step(optimizer)                 # skips the step if grads hold inf/NaN
scaler.update()
optimizer.zero_grad(set_to_none=True)
log({"grad_norm": pre_clip.item(), "loss": loss.item()})

Without AMP, drop the scaler lines and call clip_grad_norm_ directly after loss.backward(). With gradient accumulation, call it once per optimizer step, after the last micro-batch.

Log the pre-clip norm

clip_grad_norm_ returns the total norm measured before clipping, which is the most useful number to chart next to the loss. A healthy run shows a fairly steady grad-norm curve. Isolated spikes that line up with loss spikes point to bad batches or a learning rate that is too high; a norm that sits above max_norm on most steps means the threshold is rescaling almost every update, so it is no longer a spike guard; with SGD that acts like a lower learning rate, while Adam largely normalizes the scale away, so check the learning rate and data rather than relying on the clip. On steps where the AMP scaler found an overflow, the returned norm is inf or nan and the step is skipped, so filter those values before treating them as spikes.