Gradient Accumulation & Clipping¶

Overview¶

Two small hammers that fix two classic problems:

Gradient accumulation: mimic a large batch from small micro-batches when memory can't fit a big batch (or when you want to tune batch size without touching memory).
Gradient clipping: prevent exploding gradients (RNNs, transformers pre-norm, GANs) by capping the global gradient norm.

💡 Clipping is per–optimizer-step; accumulation changes when the optimizer steps. They compose.

Gradient Accumulation — the Correct Pattern¶

import torch, torch.nn as nn

model = nn.Linear(8, 4)
opt = torch.optim.SGD(model.parameters(), lr=0.1)
accum_steps = 4          # "fake" batch = 4x micro-batch

opt.zero_grad()
for step, (x, y) in enumerate(loader):          # micro-batches
    loss = model(x).abs().mean() / accum_steps  # scale DOWN by steps
    loss.backward()                             # accumulates into .grad

    if (step + 1) % accum_steps == 0:
        opt.step()          # optimizer update only every accum_steps
        opt.zero_grad()     # clear AFTER real step

Why divide by `accum_steps`?¶

Without scaling, accumulated grads are accum_stepsx larger than a real batch → LR must shrink — breaks your intended LR semantics.

The Traps¶

zero_grad timing — zero at the optimizer step, not every micro-batch.
BatchNorm — BN uses micro-batch stats; accumulation does NOT reproduce big-batch BN behavior.
DDP interplay — gradient sync happens per backward unless you defer:

# DDP + accumulation: avoid all-reduce per micro-batch
from torch.nn.parallel import DistributedDataParallel as DDP
ddp = DDP(model, device_ids=[rank])
# use context manager to sync only at the real step:
from torch.distributed.algorithms.ddp_comm_hooks import default_hooks
from contextlib import nullcontext

sync = ddp.no_sync() if (step + 1) % accum_steps != 0 else nullcontext()
with sync:
    loss.backward()

AMP + scaler: scaler.step(opt) only every accum_steps; scaler.update() at the same point.

Gradient Clipping — the Global Norm Way¶

import torch, torch.nn as nn

model = nn.Sequential(nn.Linear(8, 8), nn.GELU(), nn.Linear(8, 1))
opt = torch.optim.Adam(model.parameters(), lr=1e-3)
max_norm = 1.0

loss = model(torch.randn(4, 8)).abs().mean()
loss.backward()

# clip GLOBAL norm (recommended; scales the whole vector)
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm)

# vs per-param value clamp (rarely right):
# torch.nn.utils.clip_grad_value_(model.parameters(), clip_value=0.5)

opt.step()

Why global norm?¶

Exploding is driven by the entire gradient vector growing (e.g., repeated matrix products in RNNs). Per-param clipping destroys the direction; scaling the global norm preserves it.

# you can watch the norm before clipping:
total_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm, error_if_nonfinite=True)
# returns the PRE-clip norm (or raises on NaN/Inf)

Clipping + AMP Scaler (must unscale first)¶

scaler.scale(loss).backward()
scaler.unscale_(opt)                          # expose true grads
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm)
scaler.step(opt)
scaler.update()

⚠️ Clipping the scaled grads clips the wrong thing. Unscale, clip, then step.

Schedules of Clipping¶

When	Value	Note
RNN/unstable early training	0.1-1.0	tight
Transformers pre-norm	0.5-5.0	LoRA often 1.0
Stable vision nets	5-10+	mostly unnecessary
GAN discriminators	1.0	helps mode collapse

Debugging Exploding/Vanishing¶

def log_grad_norm(model, step):
    norm = sum(p.grad.norm().item() ** 2 for p in model.parameters() if p.grad is not None) ** 0.5
    if step % 100 == 0:
        print(f"step {step}: grad norm {norm:.3e}")
    return norm

Key Takeaways¶

Accumulation = divide loss by accum_steps, zero grads only at optimizer step.
BN and DDP need care: BN stats per micro-batch; DDP sync only at real step.
Clip the global norm for stability; unscale before clipping under AMP.
Log the pre-clip norm — it tells you why clipping matters.
These compose with FSDP/AMP/compile — they're orthogonal levers.