DL S024 · Gradient Clipping and Accumulation

Series: Deep Learning & LLMs From Scratch — 80 sessions · Session 24 / 80 · Module M04 · ~3 hours

The story · two problems, one session, both about "the gradient is the wrong size"

There are two failure modes that ruin more LLM training runs than everything else combined, and both live in the ten lines of code between .backward() and optimizer.step().

Problem 1: the gradient is too big. You are 47 hours into a $30,000 training run. Loss curve is clean. Val perplexity is beating your baseline. Then at step 47,193, one particular minibatch — maybe it contains an unusually long log-scale outlier token, maybe a code sample where the model is very confident about the wrong token, maybe just a stochastic freak — produces a gradient with L2 norm of 47.2 instead of the usual 0.6. Your careful warmup and cosine schedule and Adam moments cannot save you. That one update lands the model in a bad region of the loss surface. Loss jumps from 3.1 to 8.7 in one step. Adam's variance estimate is now polluted, so the next 200 updates are also huge. By step 48,000, loss is NaN. Forty-seven hours of compute, gone.

This scenario is exactly why every modern LLM training script starts with three lines: forward, backward, then torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0). This is gradient clipping. The idea, from Pascanu, Mikolov, Bengio 2013 "On the difficulty of training recurrent neural networks": compute the total L2 norm of all gradients, and if it exceeds some threshold, uniformly scale the entire gradient down so the norm becomes exactly the threshold. Direction of the gradient is preserved, only magnitude is bounded. Cheap, one line, and it turns "occasionally-catastrophic training" into "training." Every frontier LLM — GPT-3, GPT-4, Llama, Gemma, DeepSeek, Mistral, Qwen — uses gradient clipping at max_norm = 1.0. It is not a stylistic choice. It is why the model exists.

Problem 2: the batch is too small. You are trying to reproduce a paper that trained at batch size 256. Your GPU is a single 24GB RTX 4090 that fits batch size 32. If you just run batch 32, you get results the paper does not report and cannot compare to. If you rent 8× A100s, that is $12/hr you cannot justify. Cornered.

Enter gradient accumulation. Run 8 mini-batches of size 32 sequentially. Backward pass each one — gradients get added into .grad because PyTorch accumulates by default. After all 8, take one optimizer step and zero the grads. Mathematically identical to a single batch of 256 (with one tiny caveat around loss scaling and BatchNorm statistics — §4). You just pretended to have 8× more GPU RAM.

Together, these two techniques are the difference between "you need Google-scale infrastructure to train LLMs" and "a graduate student with a single 3090 can pretrain a 400M-parameter transformer on TinyStories in a weekend." Both are one-line changes to your loop. Both are how you actually train big models on modest hardware. Both are why "you don't have a big GPU" stopped being an excuse around 2020.

Today: what clipping is, when it fires, what max_norm to pick, how to make accumulation math-correct, how the two interact with AMP (Session 018), and the six ways to get any of it wrong.

Prerequisites

• Session 017 — training loop skeleton.
• Session 018 — AMP interacts with both features; get the order right.
• Session 023 — clipping catches what schedules cannot (occasional bad batches).

1 · Gradient clipping — two flavors (one worth using)

Two APIs. Only one you should ever use in practice.

1.1 clip_grad_norm_ — clip by TOTAL norm (recommended)

Compute the L2 norm across ALL parameters (as if they were concatenated into one big vector). If it exceeds max_norm, scale every gradient DOWN by the same factor so the total norm becomes exactly max_norm. Preserves the direction of the gradient — only shrinks its size.

loss.backward()
total_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
optimizer.step()

The returned total_norm is the PRE-clip norm — log it to TensorBoard. Sudden spikes are your early-warning system for pathological batches.

The math:

where c is max_norm and ‖g‖₂ is the total L2 norm across all parameter gradients. If ‖g‖₂ ≤ c, the gradient is untouched. Otherwise, it is scaled to length exactly c, direction unchanged.

1.2 clip_grad_value_ — clip each element to [-v, v] (avoid)

Clamps every gradient element individually. Changes the direction of the gradient (some components clip harder than others). Loses geometric meaning. Do not use this unless you have a very specific reason — it exists mostly for backward compatibility with old code and for one narrow use case (per-neuron sparsity constraints in specific research).

2 · Choosing max_norm

Standard practice by field:

When in doubt: 1.0. Log grad_norm for a few hundred steps first. If typical grad_norm is 0.1-0.3 and clipping to 1.0 rarely fires, you are safe. If it fires every step, you are clipping the signal — either raise max_norm or fix a bigger issue (LR too high, bad init, dying activations).

A useful diagnostic: percentile of grad_norm relative to your max_norm. If p50 is 0.3 and p99 is 0.9 and max_norm is 1.0, you are in the sweet spot. If p50 is 0.9 and p99 is hitting 1.0 constantly, your clip is doing too much work — probably lower LR.

3 · Clipping under AMP — the order matters (a lot)

Repeat from Session 018 §6 because this is easy to get wrong:

Why this order matters: AMP's loss-scaling multiplies the loss (and therefore all gradients) by a large factor (typically 65,536 = 2^16) to prevent underflow in fp16. If you clip before unscaling, you are clipping gradients that are 65,536× too big against a threshold of 1.0 — clipping essentially everything to near zero. Model refuses to learn; loss curve looks like it is frozen.

Always: scale → backward → unscale → clip → step → update.

For bf16 (which does not need loss scaling), you can skip the GradScaler entirely and clip normally. Most 2024-2025 pretraining uses bf16 for this reason; fp16 is now mostly for consumer-GPU inference.

4 · Gradient accumulation — the math

You want batch size B_effective = 256. Your GPU fits B_micro = 32. Run 8 micro-batches, backprop each, sum gradients in .grad, then step ONCE:

Two things worth staring at.

4.1 Why divide by accum_steps?

Standard classification losses (nn.CrossEntropyLoss, nn.MSELoss) default to reduction="mean" — they average over the batch. If you sum 8 backward passes of size-32 mean losses, you get 8 × mean_of_32. But you WANT mean_of_256, which equals sum_of_256 / 256 = (8 × sum_of_32) / 256 = sum_of_32 / 32 = mean_of_32. So the gradient accumulated across 8 micro-batches of mean-loss is 8× larger than the gradient of one true batch of 256. Dividing by accum_steps=8 cancels this exactly.

Concretely:

where a = accum_steps, B_m = micro-batch size, N_e = a · B_m, and bar_ℓ_j is the mean loss on micro-batch j. The factor 1/a on the right is exactly what loss = loss / accum_steps provides.

For reduction="sum" losses, no division is needed — the sums add correctly.

Rule: if your loss is a mean (it is, unless you set reduction="sum"), divide by accum_steps before .backward().

4.2 The BatchNorm gotcha

BatchNorm computes stats over the micro-batch (of 32), not the effective batch (of 256). Its running-mean/var updates therefore accumulate noisier estimates than they would with a real batch of 256. In many cases this is fine — the EMA smoothing over hundreds of micro-batches converges to the same distribution regardless of micro-batch size. But if you are doing very small micro-batches (say 1-4), BN becomes badly noisy and accuracy suffers.

Fix options: (1) use LayerNorm or GroupNorm which are per-example — no BN issue at all; (2) use SyncBatchNorm in a distributed setting to combine batch stats across GPUs; (3) simply increase micro-batch size if you have the memory. Transformers with LN are unaffected — nice.

4.3 The DDP gotcha — no_sync() on all-but-last micro-batch

In DistributedDataParallel training, every .backward() triggers an all-reduce of gradients across GPUs. If you naively accumulate for 8 micro-batches, you do 8 all-reduces per optimizer step, when you only need one (after the accumulated grad is final).

Use model.no_sync() on all-but-last micro-batches:

for i, (xb, yb) in enumerate(train_loader):
    is_last_micro = (i + 1) % accum_steps == 0
    context = model.no_sync() if not is_last_micro else contextlib.nullcontext()
    with context:
        loss = criterion(model(xb), yb) / accum_steps
        loss.backward()
    if is_last_micro:
        optimizer.step()
        optimizer.zero_grad()

At scale this is huge — a 7× reduction in all-reduce traffic when accum_steps=8. HuggingFace's Accelerate and Meta's torchtitan handle this for you automatically.

5 · Accumulation + AMP — the full recipe

accum_steps = 8
scaler = torch.cuda.amp.GradScaler()          # for fp16; skip for bf16
 
optimizer.zero_grad()
for i, (xb, yb) in enumerate(train_loader):
    xb = xb.to(device, non_blocking=True)
    yb = yb.to(device, non_blocking=True)
 
    with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
        loss = criterion(model(xb), yb) / accum_steps
 
    scaler.scale(loss).backward()
 
    if (i + 1) % accum_steps == 0:
        scaler.unscale_(optimizer)
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        scaler.step(optimizer)
        scaler.update()
        optimizer.zero_grad(set_to_none=True)

Rules to internalize:

1. optimizer.zero_grad() at the START of each accumulation window (not every micro-batch).
2. loss / accum_steps inside the autocast, before backward.
3. scaler.scale(loss).backward() on EVERY micro-batch.
4. Unscale, clip, step, update, zero — only after the last micro-batch of a window.
5. Scheduler step (Session 023) also happens only after the last micro-batch — one scheduler step per optimizer step, not per micro-batch.

Once you build muscle memory for this, it is boilerplate. Copy-paste with confidence.

6 · The effective-batch math and hyperparameter scaling

Effective batch size = micro_batch × accum_steps × world_size (world_size is number of GPUs in distributed training, S050+).

The reason people care: hyperparameters (LR, weight decay, warmup) transfer with the effective batch, not the micro-batch. GPT-3's paper says "batch size 3.2M tokens" — that is effective, achieved through some combination of micro-batching, accumulation, and thousands of GPUs. When you reproduce, you match the effective size, not the physical GPU count.

LR scaling rule (rough — Goyal et al. 2017 "Accurate, Large Minibatch SGD" for SGD; adjusted for Adam): when you 2× effective batch, you can 2× LR (linear scaling) up to some point (~2K batch size for SGD, ~1M tokens for Adam on LLMs), then it starts to hurt. For LLM pretraining, the 2020-2024 consensus was ~4M tokens per batch as the sweet spot; Llama 3 at 405B used 16M tokens per batch, so the sweet spot is drifting upward as models grow. See McCandlish et al. 2018 "An Empirical Model of Large-Batch Training" for the theory.

Weight decay scaling: mostly does not depend on batch size, but does depend on the effective number of updates — halve your training length, halve your weight decay (approximately).

7 · Gradient checkpointing — the third leg of the tripod

Accumulation gets you effective batch → memory savings by running mini-batches sequentially. Gradient checkpointing (Chen et al. 2016, torch.utils.checkpoint.checkpoint) gets memory savings by recomputing activations on backward instead of storing them. Trades ~30% more compute for often 40-70% less activation memory.

The mental picture: normally, forward pass stores every intermediate activation so backward can compute local gradients from them. With checkpointing, the forward pass stores only every k-th activation, and the backward pass re-runs the forward within each checkpoint segment to recompute the intermediates it needs. You pay one extra forward per segment.

Combined:

• Micro-batch = 4 with grad checkpointing (fits huge model layers).
• Accumulate 64 times (effective batch = 256).
• Add AMP with bf16 for another 50% memory savings.
• On a single 24 GB consumer GPU (RTX 4090), you can pretrain a small transformer (300-500M params) that would otherwise need an A100.

The 2024 torchtitan reference implementation uses all three plus FSDP and tensor parallelism to train Llama-3-like models at scale. For hobby-scale training, the trio (AMP + accum + checkpointing) gets you 90% of the way.

We build gradient checkpointing in Session 040 (nanoGPT). For now: know it exists, know it is the third leg of the "train big models on small GPU" tripod (with accumulation and AMP), and know that adding it costs ~30% wall time but often unlocks 2-4× larger models.

8 · War stories

9 · Diagram — clipping and accumulation in one picture

10 · Try it yourself

11 · Recap

Gradient clipping bounds the L2 norm of the total gradient to some threshold, uniformly scaling all parameter gradients while preserving direction. Standard for LLMs at max_norm=1.0. Unscale before clip under AMP. Log grad_norm to TensorBoard. Gradient accumulation runs k micro-batches, adds their gradients into .grad, and steps the optimizer once — mathematically equivalent to one big batch of k · B_micro examples, as long as you divide the mean loss by k. Combine with AMP (bf16 is easier than fp16) and gradient checkpointing to train big models on modest GPUs. Effective batch is what matters for hyperparameter scaling. In DDP, use no_sync() on all-but-last micro-batches to avoid wasted all-reduces. That is the M04 optimizer/regularization toolbox complete.

🧠 Retention scaffold

One-line summary (write it in your own words): _______________________________

Spaced review: re-read §1 (clip math) and §4 (accumulation math) in 24 hours. Revisit the full session on day 7 with focus on §5 (full recipe) and the war stories.

Next session (S025): you have finished M04. You know how to init, normalize, regularize, schedule, clip, accumulate. You can train any feed-forward model reliably. Module 5 — Convolutions & CNNs — is where architecture starts to matter for images. Session 025 rebuilds convolution from scratch as sliding-window matrix multiplies, no nn.Conv2d. After tomorrow, every CNN paper you will ever read is transparent.

Sticky note (keep on your desk):

• clip_grad_norm_(params, 1.0) is standard for LLMs. Log grad_norm to TB.
• Accumulation math: loss = loss / accum_steps if using mean loss.
• AMP order: scaler.scale → backward → unscale_ → clip → scaler.step → scaler.update.
• DDP + accum: wrap all-but-last micro-batch in model.no_sync().
• The tripod: AMP + accumulation + gradient checkpointing → 4× larger model on the same GPU.

Previous: ← DL S023 · Schedulers — Warmup, Cosine, and One-Cycle · Next: DL S025 · Convolution →