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.

Memory during training

Training memory is shared among several consumers: model weights, parameter gradients, optimizer states, saved activations, communication buffers, and memory that the allocator has reserved but cannot currently reuse because of fragmentation. Knowing which consumer is largest matters before optimizing. By Amdahl's law, reducing a small part cannot reduce total memory by much.

During forward, memory usually rises as autograd saves activations for backward. Saving them avoids recomputing each operation while calculating gradients, which trades memory for faster backward computation. The forward peak often appears near the loss, after most activations have been created.

Backward adds parameter gradients, but it also consumes and releases saved activations once their gradients have been computed. Activation memory is often larger than gradient memory, so total live memory can fall through backward even as gradients accumulate. The precise curve depends on the model, autograd schedule, and allocator.

The optimizer step can create a temporary peak. An unfused Adam implementation may materialize intermediate tensors for its update formula before releasing them. Fused optimizers reduce or avoid these buffers. After the step, the gradients are no longer needed; the next zero-grad call clears them, and the framework may retain their storage in its allocator cache for a later iteration rather than returning it to the GPU driver.

For practical memory traces and allocation inspection, see Understanding GPU Memory 1, PyTorch memory visualizations, Model training anatomy, and TRL memory reduction.

Memory flame graphs visualize memory stack traces, showing which call paths allocate the most memory. See also Flame Graphs and the FlameGraph tools.