Table of Contents
What Gradient Accumulation Means in PyTorch
In PyTorch, gradients are stored on parameter tensors in the .grad field. When you call loss.backward(), PyTorch computes gradients for all parameters that participated in the forward pass and then adds those gradients into the existing .grad buffers. This additive behavior is called gradient accumulation, and it is the default.
This has an immediate consequence for training loops: if you do not clear gradients before the next backward pass, you will silently accumulate gradients from multiple batches. Sometimes this is a bug, sometimes it is a deliberate technique to simulate a larger batch size.
PyTorch accumulates gradients by default. If you want one update per batch, you must clear gradients once per update step.
Why You Must “Zero Gradients” Before Backprop
A typical per batch training step wants gradients that reflect only the current batch. Because .grad accumulates, you need to clear old values before the next backward pass. Otherwise, your effective gradient becomes a sum across many batches, which changes optimization dynamics and usually makes training unstable.
Conceptually, if your parameter is $\theta$ and you do two backward passes without clearing, you get:
$$g = g_1 + g_2$$
where $g_1$ is the gradient from the first batch and $g_2$ from the second batch.
The Correct Place to Clear Gradients
You clear gradients right before computing new gradients for an optimizer step. The standard pattern is:
- Clear gradients.
- Forward pass.
- Compute loss.
- Backward pass.
- Optimizer step.
In code, it looks like this:
model.train()
for x, y in train_loader:
x, y = x.to(device), y.to(device)
optimizer.zero_grad(set_to_none=True)
preds = model(x)
loss = loss_fn(preds, y)
loss.backward()
optimizer.step()
Clearing gradients after optimizer.step() can also work if you never early exit, but clearing before backward() is clearer and safer when code becomes more complex.
Do not call optimizer.zero_grad() after loss.backward() and before optimizer.step() unless you intentionally want to discard the gradients.
`optimizer.zero_grad()` vs `model.zero_grad()`
Both clear gradients, but optimizer.zero_grad() clears only the parameters that the optimizer knows about, while model.zero_grad() clears gradients for all parameters in the model. In most training scripts, optimizer.zero_grad() is preferred because it matches exactly what will be updated. If you have multiple optimizers or you are freezing parts of a model, being explicit about which parameters are cleared can prevent surprises.
`set_to_none=True` and Why It Is Often Better
optimizer.zero_grad() supports set_to_none=True. This sets .grad to None instead of filling it with zeros. It often reduces memory operations and can be faster.
The difference matters in a few edge cases. A None gradient means “no gradient computed yet,” while a zero tensor means “gradient computed and it happened to be zero.” Most optimizers handle both correctly, but if you write custom logic that checks gradients, you must handle None.
If you use set_to_none=True, do not assume every p.grad is a tensor. It can be None.
Intentional Gradient Accumulation to Simulate Larger Batches
Gradient accumulation is also a useful technique when a batch size that you want does not fit in memory. You can split a large batch into several smaller micro batches, run backward() on each micro batch, and only call optimizer.step() after several micro batches. This approximates training with a larger effective batch size.
If you accumulate for accum_steps micro batches, your effective batch size is:
$$B_{\text{effective}} = B_{\text{micro}} \times \text{accum\_steps}$$
A correct accumulation loop looks like this:
model.train()
accum_steps = 4
optimizer.zero_grad(set_to_none=True)
for step, (x, y) in enumerate(train_loader, start=1):
x, y = x.to(device), y.to(device)
preds = model(x)
loss = loss_fn(preds, y)
loss = loss / accum_steps
loss.backward()
if step % accum_steps == 0:
optimizer.step()
optimizer.zero_grad(set_to_none=True)
Two details are critical here. First, you do not clear gradients every micro batch, you clear them only when you are done accumulating and after you take an optimizer step. Second, you scale the loss by accum_steps.
When accumulating gradients over accum_steps, divide the loss by accum_steps before calling backward(), otherwise your gradients become accum_steps times larger than intended.
What to Do with the Last Partial Accumulation
If the number of batches is not divisible by accum_steps, you can end the epoch with some accumulated gradients that never get applied. A simple fix is to perform a final optimizer step at the end if needed.
model.train()
accum_steps = 4
optimizer.zero_grad(set_to_none=True)
for step, (x, y) in enumerate(train_loader, start=1):
x, y = x.to(device), y.to(device)
loss = loss_fn(model(x), y) / accum_steps
loss.backward()
if step % accum_steps == 0:
optimizer.step()
optimizer.zero_grad(set_to_none=True)
if step % accum_steps != 0:
optimizer.step()
optimizer.zero_grad(set_to_none=True)Interactions With Logging and Loss Reporting
If you divide the loss by accum_steps for correct gradients, that scaled loss is not the same as the original per micro batch loss. For human readable logs, you often want to log the unscaled loss, but backprop with the scaled one.
A common pattern is:
raw_loss = loss_fn(preds, y)
(raw_loss / accum_steps).backward()
Then report raw_loss.item() for monitoring.
Common Mistakes and Their Symptoms
One common mistake is calling optimizer.zero_grad() inside the accumulation loop every micro batch. That defeats accumulation and results in tiny updates if you also divide the loss. Another common mistake is forgetting to divide the loss. That usually shows up as sudden instability, exploding loss, or needing a much smaller learning rate.
If you see training diverge right after introducing gradient accumulation, first check whether you forgot loss = loss / accum_steps.
A Minimal Mental Checklist
Before each optimizer.step(), make sure you know whether gradients currently represent one batch or several. If it is one batch, you should have cleared gradients once since the previous step. If it is several batches, you should not have cleared gradients in between, and you should have scaled the loss so the gradient magnitude matches your intended effective batch size.