A framework's training call hides most of what one step touches. That is fine until you checkpoint, restart, or spread the work across many ranks, and state you never looked at has to stay consistent.
What one step actually does
A loader produces token IDs, masks, labels and optional metadata. The forward pass computes activations and the loss. Automatic differentiation walks the graph backward to produce gradients, and across ranks a gradient synchronization (an all-reduce or a reduce-scatter) makes the ranks agree on the gradients each one needs. The optimizer turns those gradients and its own internal state into parameter updates, while a scheduler changes the learning rate. Gradient accumulation may repeat forward and backward several times before one optimizer step.
Mixed precision adds more state: autocast, loss scaling, overflow detection, unscaling and master-precision copies, depending on the format and the optimizer. FP16 commonly uses loss scaling; BF16 usually does not need it. Dropout and data augmentation consume random numbers that depend on the seed and the rank.
What a checkpoint has to capture
The course lists nine things to capture:
| State | What it holds |
|---|---|
| Parameters | the weights |
| Optimizer state | the optimizer's internal state, which shapes every update |
| Scheduler | how far the learning-rate schedule has run |
| Scaler | the loss-scaling state, when the precision format uses one |
| Random generators | the streams that dropout and augmentation draw from |
| Data cursor | where the loader has reached |
| Sampler epoch | which epoch's order the sampler is on |
| Global step | the number of optimizer steps taken so far |
| Configuration | the settings the run started with |
Omit one and the run can still resume, but, in the course's words, it can create a plausible but different continuation.
Engineering Insight
Time the phases separately: data wait, host-to-device copy, forward, backward, gradient synchronization, optimizer, checkpoint and idle barrier. Accelerator kernels run asynchronously, so synchronize the timing correctly. Then validate the loss and gradients on a tiny deterministic batch. If that baseline is not trustworthy, distributed equivalence cannot be established.
