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 7Remember
- Clip after
backward(), beforestep(); 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.