Gradient Accumulation - PyTorch Basics

Gradient Accumulation

What you'll build: gradient_accumulation(micro_batches_grads, accumulation_steps), which combines the gradients of several micro-batches into the gradient of one large batch.

Background

When a batch of 512 does not fit in GPU memory, you can run 8 micro-batches of 64, add up their gradients, and take one optimizer step. If each micro-batch loss is a mean over its samples, the mean of the micro-batch gradients equals the gradient of the full 512-sample batch, so training behaves as if the large batch fit.

In PyTorch this falls out of a design choice: loss.backward() adds into .grad instead of overwriting it (which is also why every loop calls optimizer.zero_grad()):

optimizer.zero_grad()
for i, (xb, yb) in enumerate(micro_batches):
    loss = loss_fn(model(xb), yb) / accumulation_steps   # scale so the sum is a mean
    loss.backward()                                      # adds into p.grad
optimizer.step()

Formula

For gradients g_1, ..., g_K with K = accumulation_steps:

effective_grad = (1 / K) * (g_1 + g_2 + ... + g_K)

Inputs

• micro_batches_grads: list of gradient arrays, all with the same shape (any shape).
• accumulation_steps: integer K, with 1 <= K <= len(micro_batches_grads).

Expected Return Format

An array with the same shape as each gradient: the mean of the first K gradients.

Edge Cases and Conventions

• Gradients after the first K belong to the next optimizer step and are ignored.
• With K = 1 the result is the first gradient itself.

Constraints

• PyTorch track: the reference solution lets backward() do the accumulation; a torch.stack(...).mean(dim=0) one-liner also passes.

Follow-up

• Accumulation matches a large batch exactly for a mean loss. Which layer breaks that equivalence? (Hint: it normalises with batch statistics.)
• How should the learning rate change if you raise the effective batch size 8 times?

Input: gradient_accumulation([np.array([1.0, 2.0]), np.array([3.0, 4.0])], 2)
Output: array([2., 3.])

Input: gradient_accumulation([np.array([1.0, 0.0]), np.array([3.0, 2.0]), np.array([100.0, 100.0])], 2)
Output: array([2., 1.])

import numpy as np

def gradient_accumulation(micro_batches_grads, accumulation_steps):
    """
    Average the gradients of the first accumulation_steps micro-batches.

    Args:
        micro_batches_grads: List of gradient arrays (all the same shape)
        accumulation_steps: Number of micro-batches to accumulate (K)

    Returns:
        Mean of the first K gradients, same shape as each gradient
    """
    # Your code here
    pass