🎓 A1Learner
🔍

PyTorch Gradient Accumulation

Gradient accumulation is a training pattern that adds gradients from several small batches before one step. One step then carries the effect of one big batch.

Stepping on every slice

Stepping inside the slice loop applies four small updates. Each update reacts to partial data only.

for start in range(0, 32, 8):
    loss = loss_fn(model(x[start:start + 8]), y[start:start + 8])
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

Scaling each micro loss

Each micro loss is a mean over its own 8 rows. Dividing by 4 makes the four scaled losses add to the mean over all 32 rows.

micro_loss = loss_fn(model(x[start:start + 8]), y[start:start + 8])
(micro_loss / 4).backward()

Gradients add up

Two backward() calls with no zero_grad() between them leave one summed gradient. The weight starts at 1.0 in this small check.

import torch

torch.manual_seed(0)
model = torch.nn.Linear(1, 1, bias=False)
with torch.no_grad():
    model.weight.fill_(1.0)
loss_fn = torch.nn.MSELoss()
x1 = torch.tensor([[1.0], [2.0]])
y1 = torch.tensor([[2.0], [4.0]])
x2 = torch.tensor([[3.0]])
y2 = torch.tensor([[6.0]])
model.zero_grad()
loss1 = loss_fn(model(x1), y1)
loss1.backward()
print(f"After first: {model.weight.grad.item():.4f}")
loss2 = loss_fn(model(x2), y2)
loss2.backward()
print(f"After second: {model.weight.grad.item():.4f}")

Example

import torch

torch.manual_seed(0)
x = torch.randn(32, 3)
true_w = torch.tensor([[2.0], [-1.0], [0.5]])
y = x @ true_w + 0.3 + torch.randn(32, 1) * 0.1
loss_fn = torch.nn.MSELoss()


def make_model():
    torch.manual_seed(7)
    return torch.nn.Linear(3, 1)


big_model = make_model()
big_optimizer = torch.optim.SGD(big_model.parameters(), lr=0.1)
big_optimizer.zero_grad()
big_loss = loss_fn(big_model(x), y)
big_loss.backward()
big_optimizer.step()

small_model = make_model()
small_optimizer = torch.optim.SGD(small_model.parameters(), lr=0.1)
small_optimizer.zero_grad()
micro_loss_total = 0.0
for start in range(0, 32, 8):
    micro_loss = loss_fn(small_model(x[start:start + 8]), y[start:start + 8])
    (micro_loss / 4).backward()
    micro_loss_total += micro_loss.item() / 4
small_optimizer.step()

max_diff = 0.0
for big_param, small_param in zip(big_model.parameters(), small_model.parameters()):
    diff = (big_param - small_param).abs().max().item()
    max_diff = max(max_diff, diff)

print(f"Big batch loss: {big_loss.item():.4f}")
print(f"Accumulated loss: {micro_loss_total:.4f}")
print(f"Max parameter difference: {max_diff:.8f}")

Key points

  • Micro-batch: a small slice of a batch, processed one slice at a time
  • Use (loss / micro_count) to scale each micro-batch loss before backward()
  • Equal micro-batch sizes make the scaled losses add to the big-batch loss
  • Gradients add up across backward() calls until zero_grad() clears them
  • Accumulation gives one step the reach of a big batch, without holding a big batch in memory

In this example, two copies of one model start from the same seed. The data holds 32 rows made from the rule 2, -1 and 0.5 times the inputs, plus 0.3. The first model takes one step on all 32 rows at once. The second model runs backward() on four slices of 8 rows. Each slice loss is divided by 4 before backward(). One step() then applies the summed gradients. Both losses print as 4.0989. The max parameter difference prints as 0.00000000 at eight decimal places.

Note

Accumulation saves memory, and it does not make the maths itself faster

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

Exercise

Why is each micro-batch loss divided by the number of micro-batches?