Between backward and step

Gradient accumulation is usually introduced as a way around a memory limit, and the mechanical part is trivial: skip zero_grad, let .grad sum across several backward passes, step once at the end. The claim underneath is the interesting half, because you are asserting that the sum you built equals the gradient of a batch you never had room for.

It does, on one condition. mse_loss and cross_entropy average over their own batch by default, so four micro-batches of 16 give you four averages over 16 rows, and summing those is four times the average over 64 rows. Divide each micro-batch loss by four before backward and the identity holds.

Run both and the difference is not subtle. With the divisor, the accumulated gradient and the full-batch gradient agree to 2.4e-07, which is float32 rounding on a sum of 64 terms. Without it, the largest element is off by 4.8, and the gradient norms sit at a ratio of exactly 4.0.

import torch
from torch import nn


def fresh():
    torch.manual_seed(0)
    return nn.Sequential(nn.Linear(8, 32), nn.ReLU(), nn.Linear(32, 1))


def grads(m):
    return [p.grad.detach().clone() for p in m.parameters()]


def gap(a, b):
    return max((u - v).abs().max().item() for u, v in zip(a, b))


torch.manual_seed(1)
x = torch.randn(64, 8)
y = x.sum(dim=1, keepdim=True)

m = fresh()
nn.functional.mse_loss(m(x), y).backward()
full = grads(m)                                  # one batch of 64

m = fresh()
for i in range(4):
    s = slice(i * 16, (i + 1) * 16)
    (nn.functional.mse_loss(m(x[s]), y[s]) / 4).backward()
print(gap(full, grads(m)))                       # 2.384185791015625e-07

m = fresh()
for i in range(4):
    s = slice(i * 16, (i + 1) * 16)
    nn.functional.mse_loss(m(x[s]), y[s]).backward()
print(gap(full, grads(m)))                       # 4.807855606079102

Dividing by the number of micro-batches is a shortcut that happens to be correct when every micro-batch is the same size. Datasets are rarely so obliging, and the last group in an epoch is the one that catches people.

Split 64 rows as 24, 24 and 16 and divide each loss by three, and the result misses the full-batch gradient by 0.202 in its largest element. Weight each micro-batch by its own sample count over the total instead, 24/64 and 24/64 and 16/64, and the gap drops back to 2.4e-07.

The rule that survives every batch shape: the weight on a micro-batch is its share of the samples, not its share of the micro-batches. A loop that hardcodes loss / accum_steps is asserting the two are the same thing.

def accumulate(weights):
    m = fresh()
    off = 0
    for n, w in zip([24, 24, 16], weights):
        s = slice(off, off + n)
        off += n
        (nn.functional.mse_loss(m(x[s]), y[s]) * w).backward()
    return grads(m)


print(gap(full, accumulate([1 / 3, 1 / 3, 1 / 3])))          # 0.20217061042785645
print(gap(full, accumulate([24 / 64, 24 / 64, 16 / 64])))    # 2.384185791015625e-07

The identity covers gradients and stops there. Anything in the model that computes a statistic over the rows in front of it sees micro-batches, because micro-batches are what it was given.

Batch normalization is the clear case. One BatchNorm1d fed 64 rows at once and another fed the same 64 rows in four passes end up with running means that differ by 0.049, and their num_batches_tracked counters read 1 and 4. Nothing about that is a bug in accumulation. The layer normalized over 16 rows because 16 rows is what arrived.

There is a distributed consequence too, and its mechanism is chapter eight's. Bucketed all-reduce fires as soon as a bucket's gradients are ready, so a plain accumulation loop under DDP pays a reduction per micro-batch rather than one per optimizer step. DistributedDataParallel.no_sync is the context manager that suppresses it for every micro-batch but the last.

import torch
from torch import nn

torch.manual_seed(1)
x = torch.randn(64, 8)

one = nn.BatchNorm1d(8)
one(x)                                        # a single batch of 64

four = nn.BatchNorm1d(8)
for i in range(4):
    four(x[i * 16 : (i + 1) * 16])            # the same 64 rows, four at a time

print((one.running_mean - four.running_mean).abs().max().item())
# 0.048792045563459396
print(int(one.num_batches_tracked), int(four.num_batches_tracked))   # 1 4

clip_grad_norm_ does not clip tensors. It concatenates every gradient you hand it into one conceptual vector, takes that vector's norm, and if the norm exceeds your threshold it multiplies every gradient by the same scalar. Direction is preserved exactly; only magnitude moves.

Two details of the interface are easy to miss. The value it returns is the norm before clipping, not after, which makes it the cheapest gradient-norm logger available and a confusing one if you assume otherwise. And it returns that value whether or not it clipped anything.

The per-tensor norms afterward are the tell that this is a global operation. Clip a model whose total norm is 206.35 down to 1.0 and the four parameters come back at 0.380, 0.087, 0.921 and 0.010: nowhere near 1.0 individually, exactly 1.0 together.

import torch
from torch import nn


def fresh():
    torch.manual_seed(0)
    return nn.Sequential(nn.Linear(8, 32), nn.ReLU(), nn.Linear(32, 1))


def norm(m):
    return torch.cat([p.grad.flatten() for p in m.parameters()]).norm().item()


torch.manual_seed(1)
x = torch.randn(64, 8)
y = x.sum(dim=1, keepdim=True) * 50.0            # far enough away to make big grads

m = fresh()
nn.functional.mse_loss(m(x), y).backward()
print(norm(m))                                                        # 206.35205078125
print(nn.utils.clip_grad_norm_(m.parameters(), max_norm=1.0).item())  # 206.35205078125
print(norm(m))                                                        # 1.0
print([round(p.grad.norm().item(), 6) for p in m.parameters()])
# [0.380417, 0.086949, 0.920664, 0.010048]

Put those two sections together and the placement question answers itself. Clipping rescales a vector, and the vector you meant to clip is the accumulated gradient. Clipping each partial sum rescales four different vectors, and their sum is not the clipped total.

The measurement makes the trap visible. Clip every micro-batch and the final norm reads 1.0. Clip once at the end and it also reads 1.0. The two gradient sets differ by 0.245 in their largest element, so the norm you were watching cannot tell you which loop you wrote.

Mixed precision adds one more ordering constraint on top. Chapter nine covers the scaler itself; what it imposes here is that gradients have to be unscaled before they are clipped, because a norm measured on scaled gradients is not the quantity your threshold was chosen against.

def accumulate(clip_each):
    m = fresh()
    for i in range(4):
        s = slice(i * 16, (i + 1) * 16)
        (nn.functional.mse_loss(m(x[s]), y[s]) / 4).backward()
        if clip_each:
            nn.utils.clip_grad_norm_(m.parameters(), max_norm=1.0)
    if not clip_each:
        nn.utils.clip_grad_norm_(m.parameters(), max_norm=1.0)
    return [p.grad.clone() for p in m.parameters()]


per_micro = accumulate(True)
once = accumulate(False)
print(torch.cat([g.flatten() for g in per_micro]).norm().item())   # 1.0
print(torch.cat([g.flatten() for g in once]).norm().item())        # 1.0
print(max((a - b).abs().max().item() for a, b in zip(per_micro, once)))
# 0.24483658373355865

One inf in one gradient entry, out of 321 in this small model, and the whole clip goes quiet in a specific way. The total norm comes back inf, the coefficient is computed as max_norm / (total_norm + 1e-6), which is 0.0, and torch.clamp(clip_coef, max=1.0) leaves it there.

Multiplying by zero then produces 320 zeros and one nan, because inf * 0 is nan. Your step runs on a model whose entire gradient is zero except for one poisoned entry, and the only signal is the return value you probably discarded.

error_if_nonfinite=True turns this into a RuntimeError naming the norm order and telling you the flag exists. It defaults to False on torch 2.2.2, and the current reference still documents the default as False, so the loud version is the one you have to ask for.

import torch
from torch import nn

torch.manual_seed(0)
m = nn.Sequential(nn.Linear(8, 32), nn.ReLU(), nn.Linear(32, 1))
torch.manual_seed(1)
x = torch.randn(64, 8)
nn.functional.mse_loss(m(x), x.sum(dim=1, keepdim=True) * 50.0).backward()

m[0].weight.grad[0, 0] = float("inf")            # one entry out of 321
print(nn.utils.clip_grad_norm_(m.parameters(), max_norm=1.0).item())   # inf

flat = torch.cat([p.grad.flatten() for p in m.parameters()])
print(flat.numel(), int((flat == 0).sum()), int(flat.isnan().sum()))   # 321 320 1