Gradient clipping and accumulation

PyTorch · Practical · 10 min

Both are standard in large-model training: clipping prevents loss spikes, accumulation fits big effective batches in small memory.

Gradient clipping rescales gradients when their combined norm exceeds a threshold:

torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

Call it after backward() and before step(). It returns the total norm before clipping, which is worth logging: sudden spikes signal instability. Clipping is routine for transformers and RNNs, where occasional huge gradients would otherwise wreck the weights.

Gradient accumulation uses the fact that gradients add up across backward() calls. Run several micro-batches, calling backward() on each, and only then step() and zero_grad(). With 4 micro-batches of 8, the update matches one batch of 32, at the memory cost of 8.

Two details make it correct:

  • Divide each micro-batch loss by the number of accumulation steps, so the summed gradient equals the gradient of the mean loss.
  • Count learning-rate scheduler steps per optimizer step, not per micro-batch.

Layers that use batch statistics, like BatchNorm, still only see the micro-batch, so accumulation is not perfectly equivalent for them. Under DistributedDataParallel, wrap the non-final micro-batches in model.no_sync() to skip redundant gradient communication.

Example

import torch
from torch import nn

torch.manual_seed(0)
model = nn.Linear(10, 1)
opt = torch.optim.SGD(model.parameters(), lr=0.01)
accum = 4                                         # effective batch = 4 * 8
data = [(torch.randn(8, 10), torch.randn(8, 1)) for _ in range(8)]

opt.zero_grad()
for i, (xb, yb) in enumerate(data):
    loss = nn.functional.mse_loss(model(xb), yb) / accum
    loss.backward()                               # gradients add up
    if (i + 1) % accum == 0:
        norm = nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
        opt.step()
        opt.zero_grad()
        print(f"update after batch {i}, norm before clipping {norm:.2f}")
# prints two updates: after batch 3 and batch 7

Remember

  • Clip after backward(), before step(); log the returned norm
  • Accumulate: scale loss by 1/steps, step every N micro-batches
  • Scheduler steps count optimizer updates, not micro-batches

Watch out for

  • Forgetting to divide the loss by the accumulation steps, which multiplies the effective learning rate.

Check yourself

  1. You accumulate gradients over 4 micro-batches without dividing the loss by 4. The effect is roughly that...

Sign in to answer, see the explanations, and get this lesson back as a quick review later.

Official docs: pytorch.org/docs/stable/generated/torch.nn.utils.clip_grad_norm_.html

Builds on

Unlocks

Nothing builds on this yet.