Forward pass, backward pass, optimizer step, checkpointing, logging. What each phase costs and where engineers spend their debugging time.
Forward pass: the cheap part
Input goes through the model, activations are computed layer by layer, loss comes out at the end.
- Relatively fast compared to backward pass - roughly 1/3 of the compute
- Memory cost is storing activations for every layer (needed for backprop)
- Activation checkpointing trades compute for memory - recompute activations instead of storing them
Backward pass: where the time goes
Gradients are computed layer by layer in reverse. This is the expensive part of each training step.
- Roughly 2x the compute of the forward pass
- Memory peaks here - you need both the stored activations and the gradients
- Gradient accumulation lets you simulate larger batch sizes without the memory cost
- Most numerical instability (NaN gradients, loss spikes) surfaces here first
Optimizer step: updating the weights
Gradients are computed. Now you update the model parameters.
- Adam and its variants (AdamW, LAMB) maintain running averages per parameter - doubles or triples the memory footprint of the model itself
- Learning rate scheduling happens here - warmup, cosine decay, linear decay
- Gradient clipping happens before the optimizer step to prevent instability
- In distributed training, gradients must be synchronized across all workers before this step
Checkpointing: the insurance you forget to test
Periodically save the model state so you can resume if something fails.
- A large model checkpoint can be tens or hundreds of gigabytes
- Writing to disk or cloud storage takes time - synchronous checkpointing stalls training
- Async checkpointing helps but adds complexity
- The checkpoint you didn't test restoring from is the one you'll need at 3am
Logging and monitoring: the stuff nobody budgets time for
Loss curves, learning rate, gradient norms, GPU memory, throughput.
- Without gradient norm tracking, you won't see instability coming until the loss explodes
- Without throughput monitoring, you won't know your data loader became the bottleneck
- Weights & Biases, TensorBoard, or even just printing to stdout - anything beats nothing
- Most debugging time is spent here, looking at logs after something went wrong
Data loading: the phase that isn't a phase
Data loading isn't a step in the training loop - it runs in parallel. But when it falls behind, everything else waits.
- num_workers, prefetch_factor, pin_memory - the PyTorch DataLoader knobs that matter
- Slow storage, complex preprocessing, or insufficient CPU cores all cause the GPU to idle
- Easy to overlook because it doesn't show up as a named phase in profiling - it shows up as gaps