Large gradient norms in PyTorch can turn one backward pass into an unstable optimizer update. Gradient clipping sets an upper bound on the combined parameter-gradient norm before optimizer.step(), which helps training loops keep occasional spikes from dominating the next weight update.
PyTorch handles norm clipping with torch.nn.utils.clip_grad_norm_(). The function computes the total norm across the supplied parameters, modifies existing gradient tensors in place, and returns the total norm before clipping so the training loop can log how large the update would have been.
Place the clipping call after loss.backward() has filled .grad buffers and before the optimizer uses those buffers. A CPU smoke script is enough to prove the order, and mixed-precision loops should unscale gradients with GradScaler before clipping so the threshold applies to the real gradient values.
Related: How to run a training loop in PyTorch
Related: How to debug NaN loss in PyTorch
Related: How to enable mixed precision in PyTorch
import torch from torch import nn from torch.nn.utils import clip_grad_norm_ def total_grad_norm(model: nn.Module) -> float: norms = [ parameter.grad.detach().norm(2) for parameter in model.parameters() if parameter.grad is not None ] return torch.linalg.vector_norm(torch.stack(norms), 2).item() torch.manual_seed(7) model = nn.Sequential( nn.Linear(4, 8), nn.ReLU(), nn.Linear(8, 1), ) optimizer = torch.optim.SGD(model.parameters(), lr=0.05) loss_fn = nn.MSELoss() inputs = torch.tensor( [ [1.0, 0.5, -1.0, 2.0], [0.0, -1.5, 2.0, 1.0], [2.0, 1.0, 0.5, -0.5], [-1.0, 2.0, 1.5, 0.0], ], dtype=torch.float32, ) target = torch.full((4, 1), 25.0) max_norm = 0.25 optimizer.zero_grad(set_to_none=True) loss = loss_fn(model(inputs), target) loss.backward() grad_norm_before = total_grad_norm(model) returned_norm = clip_grad_norm_( model.parameters(), max_norm=max_norm, error_if_nonfinite=True, ) grad_norm_after = total_grad_norm(model) optimizer.step() optimizer.zero_grad(set_to_none=True) gradients_cleared = all(parameter.grad is None for parameter in model.parameters()) print(f"torch_version={torch.__version__}") print(f"loss={loss.item():.4f}") print(f"max_norm={max_norm:.2f}") print(f"grad_norm_before={grad_norm_before:.4f}") print(f"clip_grad_norm_returned={returned_norm.item():.4f}") print(f"grad_norm_after={grad_norm_after:.4f}") print(f"after_within_limit={grad_norm_after <= max_norm + 1e-6}") print("optimizer_step=complete") print(f"gradients_cleared={gradients_cleared}")
The target values intentionally make the first backward pass produce a norm above the threshold. error_if_nonfinite=True fails fast if the total norm is NaN or infinite instead of silently clipping an unusable gradient set.
$ python clip_grad_demo.py torch_version=2.12.1+cpu loss=613.9304 max_norm=0.25 grad_norm_before=75.0652 clip_grad_norm_returned=75.0652 grad_norm_after=0.2500 after_within_limit=True optimizer_step=complete gradients_cleared=True
clip_grad_norm_returned matches the pre-clip norm, while grad_norm_after is at the configured threshold. optimizer_step=complete confirms the clipped gradients still reached the optimizer.
optimizer.zero_grad(set_to_none=True) prediction = model(inputs) loss = loss_fn(prediction, target) loss.backward() total_norm = torch.nn.utils.clip_grad_norm_( model.parameters(), max_norm=1.0, error_if_nonfinite=True, ) optimizer.step()
Use the same parameter iterable that the optimizer updates. If the optimizer trains only part of a model, clip that same parameter group instead of every parameter object in the module.
if step % 100 == 0: print(f"step={step} grad_norm_before_clip={total_norm.item():.4f}")
A returned norm that is usually far below max_norm means clipping rarely changes the update. A returned norm that is almost always above max_norm can indicate an aggressive learning rate, a bad batch, or a loss that needs separate debugging.
Related: How to debug NaN loss in PyTorch
scaler.scale(loss).backward() scaler.unscale_(optimizer) torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0) scaler.step(optimizer) scaler.update()
Call scaler.unscale_(optimizer) once for that optimizer after all gradients for the upcoming step have been accumulated. Clipping before unscaling applies the threshold to scaled gradients instead of the values the optimizer will use.
Related: How to enable mixed precision in PyTorch
$ rm clip_grad_demo.py