Unit 03.03: Optimizer step and zeroing gradients

Three lines, always in the same order. Getting the order wrong produces a model that trains slowly, erratically, or not at all — with no error message.

Why gradients accumulate, and what to do about it

PyTorch adds each new gradient to whatever is already in .grad. That is deliberate: it lets you accumulate across several batches when memory forces small ones. But it means that forgetting to clear gradients makes every step use the sum of all previous steps.

The first half of this code proves the accumulation, then shows the two ways to clear it:

import torch
from torch import nn

torch.manual_seed(0)
model = nn.Linear(4, 1)
optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
x, y = torch.randn(8, 4), torch.randn(8, 1)

# Gradients ACCUMULATE. Run backward twice without clearing and you get the
# sum of both, which is the single most common training bug.
nn.MSELoss()(model(x), y).backward()
first = model.weight.grad.clone()

nn.MSELoss()(model(x), y).backward()
second = model.weight.grad.clone()

print("gradients accumulated:", torch.allclose(second, first * 2))   # True

# The older behaviour: overwrite the gradient with zeros.
optimizer.zero_grad(set_to_none=False)
print("set_to_none=False -> sum:", model.weight.grad.abs().sum().item())   # 0.0

# The modern default: release the tensor entirely, which is cheaper. The
# attribute becomes None, not a tensor of zeros, so code that reaches for
# .grad.abs() straight after zero_grad() raises AttributeError.
optimizer.zero_grad()
print("default zero_grad() ->  :", model.weight.grad)                      # None

# The correct order, every step: zero -> backward -> step
optimizer.zero_grad()
loss = nn.MSELoss()(model(x), y)
loss.backward()
before = model.weight.clone()
optimizer.step()
print("weights changed:", not torch.allclose(before, model.weight))   # True

Since PyTorch 2.0, zero_grad() releases the gradient tensors entirely rather than filling them with zeros — it is cheaper. So .grad becomes None, not a tensor of zeros. Code that reaches straight for .grad.abs() after clearing raises AttributeError. Pass set_to_none=False for the older behaviour.

The order is always: zero, backward, step.

The mistake this prevents

Calling optimizer.step() before loss.backward(). The step applies whatever stale gradients are lying around, so the model updates on the previous batch. Nothing errors and the loss curve just looks noisy.

Takeaway

Zero, backward, step — in that order, every iteration.

Mark complete