How does gradient accumulation work in PyTorch and when would you use it?

Gradient accumulation simulates a larger effective batch size than fits in GPU memory by summing gradients over several smaller forward/backward passes before calling optimizer.step(). This is useful when training large models on limited GPU memory.

import torch
import torch.nn as nn

model     = nn.Linear(100, 10)
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)
criterion = nn.CrossEntropyLoss()

# Simulate effective batch_size=128 using micro_batch=32 (4 accumulation steps)
accumulation_steps = 4

optimizer.zero_grad()
for step, (X_micro, y_micro) in enumerate(loader):  # loader yields micro-batches
    logits = model(X_micro)
    loss = criterion(logits, y_micro)

    # CRITICAL: scale loss by 1/accumulation_steps before backward
    # so the accumulated gradient matches what a single large-batch
    # backward pass would have produced
    loss = loss / accumulation_steps
    loss.backward()           # gradients ACCUMULATE (not cleared)

    if (step + 1) % accumulation_steps == 0:
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
        optimizer.step()                          # update only every N micro-batches
        optimizer.zero_grad(set_to_none=True)     # clear for next accumulation cycle

# Effective batch size = micro_batch_size * accumulation_steps
# This trades extra forward/backward compute for lower peak memory usage

Invest now in Acorns!!! 🚀
Join Acorns and get your $5 bonus!

Acorns is a micro-investing app that automatically invests your "spare change" from daily purchases into diversified, expert-built portfolios of ETFs. It is designed for beginners, allowing you to start investing with as little as $5. The service automates saving and investing. Disclosure: I may receive a referral bonus.

Invest now!!! Get Free equity stock (US, UK only)!

Use Robinhood app to invest in stocks. It is safe and secure. Use the Referral link to claim your free stock when you sign up!.

The Robinhood app makes it easy to trade stocks, crypto and more.

Webull! Receive free stock by signing up using the link: Webull signup.