Skip to content

← Memory, Copies and In-Place step 6 of 6

Medium Primitives

Gradient accumulation, and the buffer that never moved

Accumulating over chunks, so a batch too large for memory can be processed in pieces:

def accumulate(acc, chunks):
    for c in chunks:
        acc = acc + c        # a new tensor every iteration
    return acc

The returned value is right. The caller’s acc is still zero.

Why this one is worse than it looks

The single-write version of this mistake is easy to spot. In a loop it hides, because the function does return the right answer, so it passes any test that only checks the return value. It fails the moment someone relies on the buffer, which is precisely the pattern accumulation exists for:

optimizer.zero_grad()
for micro_batch in split(batch):
    loss = model(micro_batch).mean()
    loss.backward()          # adds into p.grad, in place, for every p
optimizer.step()

backward accumulates into p.grad rather than replacing it, and the optimiser reads p.grad. If accumulation rebound instead of adding in place, the optimiser would see whatever was there before and the micro-batches would be silently discarded.

This is also why zero_grad() has to exist at all: the accumulation is the default, so somebody has to clear it.

The fix

acc.add_(c)

One buffer for the whole loop, written to len(chunks) times, and the caller sees every one of them. It also allocates nothing per iteration, where the rebinding version allocates a full-size tensor per chunk and throws away all but the last.

Your task

def accumulate(acc: torch.Tensor, chunks: list) -> torch.Tensor

Add every chunk into acc, in place, and return acc. The caller holds the same tensor and expects to see the total in it.