One Train Step

The five-line ritual

The unit of training is one step. It always has the same five lines, in the same order:

opt.zero_grad() — clear gradients from the previous step.
logits = model(xb) — forward pass.
loss = loss_fn(logits, yb) — scalar loss.
loss.backward() — compute gradients via autograd.
opt.step() — apply gradients to parameters.

Optionally insert clip_grad_norm_ between backward and step, and scheduler.step() after step. That's the whole loop body. Everything else is plumbing around it.

Tip: If you can't draw the five-line ritual on a whiteboard from memory, you don't yet have the muscle. Write it twenty times. Eventually you will type opt.zero_grad() before you finish thinking the sentence.

The two famous bugs of step 1

Forgotten zero_grad() — gradients accumulate across steps, so each step sees the sum of every previous gradient. Loss stays high, training collapses. The fix is one line.

Calling opt.step() before loss.backward() — you step on stale gradients (the previous batch's). Loss curve looks vaguely like training but the model isn't actually learning the current batch.

Gradient accumulation (the right way to skip `zero_grad`)

Sometimes you want a larger effective batch size than fits in VRAM. The trick: skip zero_grad() for several batches, divide the loss by the number of accumulation steps, then opt.step() + opt.zero_grad() at the end. This produces gradients equivalent to a single batch accum_steps times bigger.

Principle: The five-line ritual is the smallest unit of training. Memorize it. Then learn one variation at a time (gradient accumulation, mixed precision, gradient clipping). Do not skip the basic version.

Code

One train step, with all the optional pieces·python

from torch.nn.utils import clip_grad_norm_

opt.zero_grad()
logits = model(xb)
loss = loss_fn(logits, yb)
loss.backward()
clip_grad_norm_(model.parameters(), max_norm=1.0)   # optional
opt.step()
scheduler.step()                                    # optional

Gradient accumulation for big effective batch·python

ACCUM = 4
opt.zero_grad()
for micro_idx, (xb, yb) in enumerate(loader):
    logits = model(xb)
    loss = loss_fn(logits, yb) / ACCUM        # rescale for averaging
    loss.backward()                            # accumulates into .grad
    if (micro_idx + 1) % ACCUM == 0:
        clip_grad_norm_(model.parameters(), max_norm=1.0)
        opt.step()
        opt.zero_grad()

The five-line ritual

The two famous bugs of step 1

Gradient accumulation (the right way to skip `zero_grad`)

Code

External links

Exercise

Progress

Comments 0

The five-line ritual

The two famous bugs of step 1

Gradient accumulation (the right way to skip zero_grad)

Code

External links

Exercise

Progress

Gradient accumulation (the right way to skip `zero_grad`)