We can't find the internet
Attempting to reconnect
Something went wrong!
Attempting to reconnect
← Memory, Copies and In-Place step 6 of 6
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.
Stuck?
PyTorch reference solution
Sign in to attempt this problem and reveal the reference solution.