Gradient clipping in PyTorch: norms, timing, and AMP

In PyTorch, apply gradient clipping after the backward pass and before the optimizer updates the parameters. If you accumulate gradients, finish the accumulation window first. If you use a gradient scaler, unscale the gradients before clipping them.

This guide uses four small CPU examples to make that order visible. You will inspect a norm, compare norm clipping with value clipping, and see how two ordering mistakes change the result. Twenty local checks passed with Python 3.12.14 and PyTorch 2.15.0+cpu. The examples require no dataset or downloaded model.

Start with a gradient you can check by hand

The clip_grad_norm_ API modifies gradients in place and uses an L2 norm by default. Its parameter list defines the group being clipped. In a model with one optimizer covering all parameters, that list will commonly be model.parameters().

Here, a synthetic linear loss produces the gradient [3, 4], whose L2 norm is 5. The threshold of 1 is chosen for easy arithmetic. It is not a recommended setting for every model. Each Python block below can run independently in an environment with PyTorch installed.

import torch

weights = torch.nn.Parameter(torch.tensor([1.0, -1.0], dtype=torch.float64))
optimizer = torch.optim.SGD([weights], lr=0.1)
optimizer.zero_grad(set_to_none=True)

# Synthetic loss chosen to produce the gradient [3, 4].
loss = 3 * weights[0] + 4 * weights[1]
loss.backward()
before = torch.nn.utils.clip_grad_norm_(
    [weights], max_norm=1.0, error_if_nonfinite=True
)
print('Returned norm:', round(before.item(), 6))
print('Norm after:', round(weights.grad.norm().item(), 6))
print('Gradient:', [round(x, 6) for x in weights.grad.tolist()])
optimizer.step()
print('Weights:', [round(x, 6) for x in weights.tolist()])

Observed output:

Returned norm: 5.0
Norm after: 1.0
Gradient: [0.6, 0.8]
Weights: [0.94, -1.08]

The returned norm is 5, measured before clipping. The gradient becomes approximately [0.6, 0.8], with a norm just below 1 that rounds to 1.0. The returned value does not become the new norm; inspect the gradients afterward when you need that measurement.

Clipping alone leaves the parameter values unchanged. The subsequent SGD step subtracts 0.1 times the clipped gradient, moving the weights from [1, -1] to approximately [0.94, -1.08]. This linear loss exists to expose the arithmetic; it is not a useful training objective or evidence of convergence.

Clip one selected group consistently

We also tested the same gradient split across two separate parameters. Clipping them together returned 5 and scaled their entries to approximately 0.6 and 0.8. Calling the function separately on each parameter would express a different rule: each would receive its own limit.

Choose the parameter group deliberately. If an optimizer updates only part of a model, use the intended group rather than assuming every model parameter belongs in that clipping operation. The examples here use one optimizer and dense gradients.

Our checks also covered zero gradients, a gradient below the threshold, and a parameter whose gradient was None. Zero remained zero, [0.3, 0.4] remained unchanged with a limit of 1, and the missing gradient did not affect the calculated norm. A successful clipping call does not prove that every expected parameter received a gradient.

Distinguish norm clipping from value clipping

clip_grad_value_ limits each individual entry to the interval from the negative threshold to the positive threshold. Run it on the original [3, 4] gradient:

import torch

weights = torch.nn.Parameter(torch.zeros(2, dtype=torch.float64))
(3 * weights[0] + 4 * weights[1]).backward()
torch.nn.utils.clip_grad_value_([weights], clip_value=1.0)
print('Value-clipped gradient:', weights.grad.tolist())
print('Its L2 norm:', round(weights.grad.norm().item(), 6))

Observed output:

Value-clipped gradient: [1.0, 1.0]
Its L2 norm: 1.414214

Both entries now satisfy the value limit of 1, but their combined L2 norm is about 1.414214. In the first example, norm clipping applied a shared scale and preserved the ratio between the entries. Value clipping changed the ratio from 3:4 to 1:1. Select the operation that matches the rule you intend to impose.

Finish accumulation before clipping

Consider two equally weighted microbatches whose unnormalized gradients are [6, 0] and [0, 8]. Dividing each synthetic loss by two gives a final accumulated gradient of [3, 4]. Clip that final gradient once for the upcoming optimizer step.

The following comparison deliberately includes the incorrect alternative of clipping after every backward call:

import torch

def accumulated_gradient(clip_after_each):
    weights = torch.nn.Parameter(torch.zeros(2, dtype=torch.float64))
    # Two equally weighted synthetic microbatches.
    for coefficients in ([6.0, 0.0], [0.0, 8.0]):
        loss = (weights * weights.new_tensor(coefficients)).sum() / 2
        loss.backward()
        if clip_after_each:
            torch.nn.utils.clip_grad_norm_(
                [weights], 1.0, error_if_nonfinite=True
            )
    if not clip_after_each:
        torch.nn.utils.clip_grad_norm_(
            [weights], 1.0, error_if_nonfinite=True
        )
    return [round(x, 6) for x in weights.grad.tolist()]

print('Clip once at the end:', accumulated_gradient(False))
print('Clip after each backward:', accumulated_gradient(True))

Observed output:

Clip once at the end: [0.6, 0.8]
Clip after each backward: [0.242535, 0.970142]

In the second path, clipping reduces the first contribution before the second is added. The resulting direction differs from the once-at-the-end result even though both final norms are near 1. This is why placing clipping inside a microbatch loop can change the update you intended to compute.

The division by two assumes two equally weighted contributions. Unequal batch sizes, masked targets, or a partial final window need the appropriate loss normalization. Use our gradient accumulation guide for the full training loop and window handling. Keep the normalization decision separate from the clipping threshold.

Unscale before clipping when using GradScaler

The PyTorch mixed-precision examples put scaler.unscale_(optimizer) after all backward calls for that optimizer’s step and before clipping. Call it once per optimizer per step, then use scaler.step(optimizer) and scaler.update(). Keep the scale constant during an accumulation window.

This standalone CPU demonstration tests that ordering with an initial scale of 8. It uses float64 tensors without autocast; it does not test FP16 arithmetic or CUDA execution.

import torch

# CPU demonstration of scaling order, without autocast or FP16.
weights = torch.nn.Parameter(torch.tensor([1.0, -1.0], dtype=torch.float64))
optimizer = torch.optim.SGD([weights], lr=0.1)
scaler = torch.amp.GradScaler('cpu', init_scale=8.0)
optimizer.zero_grad(set_to_none=True)
loss = 3 * weights[0] + 4 * weights[1]
scaler.scale(loss).backward()
print('Scaled norm:', round(weights.grad.norm().item(), 6))

scaler.unscale_(optimizer)  # Once, after all backward calls for this step.
before = torch.nn.utils.clip_grad_norm_(
    [weights], 1.0, error_if_nonfinite=True
)
print('Unscaled norm before clipping:', round(before.item(), 6))
print('Clipped norm:', round(weights.grad.norm().item(), 6))
scaler.step(optimizer)
scaler.update()
print('Weights:', [round(x, 6) for x in weights.tolist()])

Observed output:

Scaled norm: 40.0
Unscaled norm before clipping: 5.0
Clipped norm: 1.0
Weights: [0.94, -1.08]

The scaled norm is 40. Unscaling restores 5, and clipping then produces the same weights as the original unscaled example. In a separate deliberately incorrect test, we clipped the scaled gradient to 1 before unscaling. Its final unscaled norm was about 0.125, one eighth of the intended limit.

The test suite also confirmed that calling unscale_ twice before the step raises an error. For the precision and range questions that motivate mixed precision, see our FP16 guide. The CPU scaling check here establishes operation order, not GPU performance or numerical behavior.

Treat nonfinite gradients as a separate problem

The norm-clipping calls use error_if_nonfinite=True. In our tests, gradients containing NaN, positive infinity, or negative infinity caused an exception before the following optimizer step, and the weights stayed unchanged. That is a fail-fast policy. Clipping did not repair the bad values.

When this happens, inspect the batch, loss calculation, precision settings, and preceding operations. Preserve enough information to reproduce the failure. Simply lowering the clipping threshold does not explain where a nonfinite value came from.

With a gradient scaler, this fail-fast choice also means an exception stops the example before the scaler can perform its usual step-skipping behavior. If your training loop should recover by skipping a nonfinite update, design that control flow explicitly; do not assume the code above already implements recovery.

Choose a threshold from observed training behavior

Start by recording the returned pre-clipping norm, the configured threshold, loss, learning rate, and how often clipping applies. Keep loss reduction and accumulation settings alongside the measurements, because changing their scale changes what a fixed threshold means.

Compare candidate thresholds against the outcome you care about, such as validation loss, while keeping the rest of the experiment controlled. If clipping applies on nearly every step, investigate whether the threshold is intentionally restrictive or whether another part of the training setup needs correction. These small examples establish no universal best threshold.

For plain SGD without momentum or weight decay, the update in our example is the learning rate times the clipped gradient. Do not generalize that relationship to every optimizer: Adam uses moving gradient statistics, while momentum and weight decay can also change the update. A gradient-norm limit is not a universal bound of the same size on parameter movement.

The 20 checks cover dense CPU gradients, the displayed snippets, accumulation order, nonfinite failures, and CPU gradient scaling. They do not validate sparse gradients, distributed or sharded training, CUDA autocast, or model accuracy. The runtime emitted a missing-NumPy initialization warning; these tensor-only examples completed without NumPy.