How to clip gradients in PyTorch

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.

1. Create a smoke script that clips gradients after the backward pass.

   clip_grad_demo.py

   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.

2. Run the smoke script to verify the clipped gradient norm.

   $ 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.

3. Add norm clipping to the real training step.

   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.

4. Log the returned total norm while tuning the threshold.

   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

5. Unscale mixed-precision gradients before clipping.

   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

6. Remove the smoke script after copying the clipping pattern into the project.

   $ rm clip_grad_demo.py