What is gradient accumulation and why is it needed on limited memory?

Gradient accumulation is a technique that allows efficient training of models with a larger batch size than the GPU memory permits, by accumulating gradients over several mini-batches before updating the weights.

When memory is limited, it is not possible to load a large batch entirely. Instead, the model processes several small batches sequentially, accumulating gradients (without updating parameters), and after a specified number of such steps, performs one optimization step.

This allows:

• Emulating a larger batch size, which can improve training stability and quality.
• Bypassing GPU memory limitations.

Example pseudocode in PyTorch:

optimizer.zero_grad()
for i, (inputs, targets) in enumerate(dataloader):
    outputs = model(inputs)
    loss = criterion(outputs, targets) / accumulation_steps
    loss.backward()
    if (i + 1) % accumulation_steps == 0:
        optimizer.step()
        optimizer.zero_grad()

Here, accumulation_steps is the number of mini-batches after which the weights are updated.