A training iteration usually has five stages: zero gradients, run forward, compute a scalar loss, run backward, and update the optimizer.

Zero gradients

The iteration starts by zeroing gradient buffers. Autograd accumulates new gradients into each parameter's existing gradient buffer instead of overwriting it. Clearing those buffers prevents gradients from an earlier step from being mixed with gradients computed using the current parameters.

Putting zeroing at the start makes this invariant explicit: every backward pass begins with empty gradient buffers. The buffers are not arbitrary garbage in a normal framework, but they may hold gradients from the previous iteration. Zeroing after the optimizer step can be equivalent if nothing uses the buffers in between, but zeroing first makes the order of each iteration easier to reason about.

Forward pass

The forward pass runs the model and records the intermediate activations required to compute gradients later. These activations can consume substantial memory. Activation checkpointing can discard selected activations during forward and recompute them during backward, trading extra computation for lower activation memory.

Scalar loss

The model output is compared with the target to compute a loss, usually one scalar value. A scalar loss provides the single starting value for backpropagation. A non-scalar output can also be differentiated, but it requires an explicit vector-Jacobian seed.

Backward pass

Calling backward starts at the loss node and traverses the autograd graph in reverse toward its leaf tensors. It applies the chain rule and accumulates gradients for trainable parameters in their gradient buffers.

Optimizer step

After backward, parameter-gradient buffers are ready for the optimizer. The optimizer updates its internal state, such as momentum or Adam moments, then updates the model parameters. The next iteration clears the old gradients and computes new ones using those updated parameters.