AMP Gradient Clipping: Why Unscale Comes First

黎 浩然/ 17 10 月, 2026/ 大语言模型/LARGELANGUAGEMODEL/LLM, 教程/TUTORIAL, 计算机/COMPUTER/ 0 comments

With GradScaler, the clipping threshold is normally defined for unscaled gradients. Clipping scaled gradients first and letting the later step unscale them also divides the effective threshold by the scale factor. Training can keep running while the applied gradient is smaller than intended.

Contents
  1. Separate scale, threshold and learning rate
  2. The wrong order changes the effective threshold
  3. Where unscaling belongs in PyTorch
  4. Twelve local checks reproduce the difference
  5. Unscaling halfway through accumulation mixes units
  6. Match the clipping set to the logged quantity
  7. Do not use clipping to conceal nonfinite values
  8. References and validation scope

Take gradient (3,4), norm threshold 1 and scale factor 8. Unscale then clip gives (0.6,0.8); reversing those operations gives (0.075,0.1), whose norm is only 0.125. A small counterexample can expose the ordering problem before a full model run. 中文版.

Separate scale, threshold and learning rate

The loss scale S, clipping threshold M and learning rate eta control different operations. Multiplying a loss by a positive S fixed with respect to the parameters multiplies its gradient by S in ideal arithmetic. Unscaling divides that factor back out to recover the original objective’s gradient; it is not a decision to give that objective a different weight.

Norm clipping intentionally changes an oversized gradient. For a finite vector g and positive M, define C_M(g)=g*min(1,M/||g||₂), leaving the zero vector unchanged. It preserves direction and moves a vector exceeding the threshold onto the boundary. This differs from independently clamping every component into an interval.

Learning rate acts later in the optimizer update. Ordinary SGD without momentum gives a change of −eta*C_M(g), but a stateful optimizer such as AdamW does not generally turn a gradient-norm bound into the same bound on parameter displacement. Keep the quantity being bounded explicit when interpreting a training chart.

The wrong order changes the effective threshold

Clipping Sg at the unchanged M and dividing afterward gives C_M(Sg)/S=g*min(1,M/(S||g||₂)). That is clipping g with threshold M/S, whereas the intended order gives C_M(g). They can agree when S=1, the gradient is zero, or neither path clips; the operations are not generally interchangeable.

In the counterexample, ||g||₂=5. The correct path divides (24,32) by 8 and then multiplies by 1/5, producing (0.6,0.8). The wrong path clips a norm of 40 to 1, producing that same intermediate vector, then divides it by 8 again. Identical intermediate numbers can represent different stages.

If S changes during training, incorrect ordering makes M/S change as well. A newly smaller gradient cannot be interpreted from learning rate alone. First identify the measurement stage. This algebra assumes positive S, rather than assuming the scale is necessarily greater than one.

Blue: unscale then clip. Brown: clip then unscale.

S=1

S=8

S=64

Original ideal-arithmetic comparison: original norm 5, threshold 1. Correct order ends at norm 1; wrong order gives 1, 0.125 and 0.015625. Not a training curve.

Where unscaling belongs in PyTorch

The PyTorch 2.14 AMP examples place unscaling before gradient inspection or modification, once per optimizer between steps and after its gradient accumulation completes. After an explicit unscale_, scaler.step does not repeat the unscaling.

The operation therefore belongs after all backward contributions for that update and before clipping, followed by scaler.step and scaler.update. With several optimizers, identify which owns each parameter. Unscaling one optimizer’s gradients does not establish that parameters belonging to another have been processed too.

Do not read AMP as requiring the same scaling configuration for every datatype or training path. This article concerns a path that explicitly uses GradScaler. Check the actual configuration and framework version before applying the order; the mathematical argument explains why the units matter once scaling is present.

1. Scaled backward
2. Finish accumulation
3. Unscale gradients
4. Clip → step → update
Original update flow: keep one scale within an effective batch, then unscale and clip the completed accumulation.

Twelve local checks reproduce the difference

This standard-library code implements the ideal L2 clipping equation for the finite-vector teaching scenario. It is not a bit-for-bit reimplementation of PyTorch, whose execution may include numerical safeguards, device branches and finite-precision differences. No GPU or model-training run was performed.

# Ideal finite-vector model: no framework epsilon, GPU, or optimizer.
from math import hypot, isclose

def mul(v, k):
    return tuple(x*k for x in v)

def clip(v, bound):
    n = hypot(*v)
    return v if n == 0 or n <= bound else mul(v, bound/n)

g, bound, scale = (3.0, 4.0), 1.0, 8.0
right = clip(mul(mul(g, scale), 1/scale), bound)
wrong = mul(clip(mul(g, scale), bound), 1/scale)
assert hypot(*g) == 5.0
assert all(isclose(x,y) for x,y in zip(right, (0.6,0.8)))
assert all(isclose(x,y) for x,y in zip(wrong, (0.075,0.1)))
assert isclose(hypot(*right), 1.0)
assert isclose(hypot(*wrong), 0.125)
assert clip((0.0,0.0), bound) == (0.0,0.0)
for s in (1.0,8.0,64.0):
    good = clip(mul(mul(g,s),1/s),bound)
    bad = mul(clip(mul(g,s),bound),1/s)
    assert isclose(hypot(*good),1.0)
    assert isclose(hypot(*bad),1/s)
print("12 clipping checks passed")

Twelve checks cover the original norm, both paths’ components and norms, the zero vector, and scales 1, 8 and 64. The comparison chart uses those teaching values directly. Its bars do not describe convergence, model quality or training throughput. The example isolates an ordering error rather than trying to reproduce every aspect of a training framework.

Unscaling halfway through accumulation mixes units

Let two microbatches contribute g1 and g2. A common scale gives S*g1+S*g2, which can be divided by S at the end. Unscale the first contribution early, then add the next scaled one, and the buffer becomes g1+S*g2. Its terms no longer share a common multiplicative factor.

More generally, separately scaled contributions S1*g1 and S2*g2 cannot be recovered for arbitrary vectors by dividing their sum by one guessed average scale. The repair is to restore the accumulation boundary, not to invent a numerical correction from the final buffer. The loss contributions and their scale need to remain interpretable throughout that boundary.

The earlier gradient accumulation and loss-normalization article addresses contribution weights. Here the issue is numerical units. Establish that microbatches represent the intended objective, then accumulate them under a common scale. Clipping cannot compensate for an incorrectly weighted objective; it introduces another transformation rather than reconstructing missing information.

Match the clipping set to the logged quantity

The clip_grad_norm_ reference describes a total norm over the supplied parameter gradients, in-place modification and a returned norm measured before clipping. That return value is not the post-clipping norm. A logged number exceeding the threshold may therefore be expected rather than evidence that clipping failed.

Clipping each layer separately at the same threshold differs from clipping the combined parameter set. Two groups can each have norm 1 after separate clipping while their concatenation has norm sqrt(2). Record the parameter set and norm_type alongside the threshold. Otherwise two configurations with an identical max_norm can impose different constraints.

For diagnosis, distinguish scaled, unscaled, pre-clipping and post-clipping observations. Record S, the update boundary and the selected parameter set. A single series named grad_norm may discard the context required to interpret it. Measurements are most useful when their units and position in the update are inspectable, not merely when the chart looks smooth.

Do not use clipping to conceal nonfinite values

The equations and code assume finite vectors. NaN and Inf are not merely large gradients, and this example does not establish that clipping makes them safe. Investigate the loss, inputs, operations and scaling state rather than first lowering the threshold. The simplified finite case should remain separate from a diagnosis of nonfinite arithmetic.

Real training may also include distributed reduction, sharded parameters or framework-specific clipping APIs. A single-process vector example cannot verify the global-norm behavior of those systems. Check the relevant framework’s update order and which parameters or shards its operation can see before interpreting its result as a model-wide bound.

A practical debugging starting point is a small fixed finite gradient with expected numbers at each observation point. Once units, parameter set and ordering agree, restore the full training conditions. If these counterexamples help untangle AMP clipping order in a training loop, voluntary support through the section below is welcome.

References and validation scope

Checked October 17, 2026 against the PyTorch 2.14 AMP examples and gradient-norm clipping reference. This explains an engineering mechanism rather than announcing a release. Both diagrams, the derivation and finite-vector code are original teaching material. Twelve checks ran locally; real-model and hardware performance were not measured.

Support

If the explanations or examples helped you solve a problem or save time, you’re welcome to support this site with a voluntary tip.

WeChat support QR code; click to enlarge
WeChat
Alipay support QR code; click to enlarge
Alipay
Buy Me a Coffee; support this site
Buy Me a Coffee

Click a QR code to enlarge. More options: support page。

Share this Post

Leave a Comment

您的邮箱地址不会被公开。 必填项已用 * 标注

*
*