Checkpointing and Recovery › What Goes Into a Checkpoint

A checkpoint must contain every piece of state that, if absent, would cause the resumed run to diverge from the original. This is a stricter requirement than it first appears.

The obvious items are model parameters and optimizer state. The less obvious items are the random number generator state, the data loader position, and the learning rate schedule step count. Miss any of these and your resumed run may converge differently, or produce non-reproducible results across restarts. Each component plays a specific role, and understanding that role helps you reason about what happens if it is omitted.

Model parameters are the weights and biases of every layer. For a 70B parameter model in BF16, this is roughly 140 GB of raw tensor data. In a data-parallel setup, every rank holds a complete copy of the model, so in principle only one rank needs to write the checkpoint. In a model-parallel or FSDP setup, parameters are sharded across ranks, and the checkpoint may either consolidate shards into a single full-model file or store per-rank shard files.

Storing shards is faster: each rank writes only its slice independently, achieving bandwidth proportional to the number of ranks. If you have 64 ranks and 10 GB/s storage bandwidth per rank, the effective aggregate write bandwidth can approach 640 GB/s, cutting the checkpoint time from minutes to seconds. Consolidating to a single file is slower but simpler to load, especially when resuming with a different parallelism configuration. Most production systems use a middle ground: they write sharded checkpoints for fast recovery in place, and periodically consolidate a full-model checkpoint for portability and model release.

The consolidation step is usually done off the critical path, either on a scheduled basis or when a permanent checkpoint milestone is reached. A script reads all the per-rank shards and writes a single merged file. This merged file is what gets published, fine-tuned, or used for inference. The sharded checkpoints serve only the operational recovery use case.

Optimizer state is often larger than the model parameters. For Adam, each parameter requires two additional tensors: the first-moment estimate $m_t$ and the second-moment estimate $v_t$. Both tensors have the same shape as the parameter tensor they correspond to, and both are typically stored in FP32 (four bytes each) for numerical stability even when the model parameters are stored in BF16 (two bytes each). This means Adam optimizer state alone requires eight bytes per parameter, compared to two bytes for the BF16 parameter itself. The combined storage cost is ten bytes per parameter, roughly five times the parameter-only storage.

The optimizer state must be checkpointed alongside the parameters. Resuming with the parameters at step $N$ but with optimizer state reset to zero would cause the optimizer to behave as though it were just starting from scratch, leading to a transient spike in the effective learning rate and a loss curve that dips before recovering. The first-moment estimate $m_t$ captures the exponentially weighted average of past gradients, giving a smoothed gradient signal that prevents oscillation. The second-moment estimate $v_t$ captures the variance of past gradients, allowing adaptive step sizes for each parameter. Both take many steps to warm up from zero to useful values. Discarding them throws away potentially thousands of steps of accumulated curvature information, and the optimizer must rebuild that information from scratch before training can progress efficiently.

This is particularly damaging in the middle of a long run where the optimizer has accumulated information about the loss landscape. A fresh Adam optimizer at step 50,000 is a fundamentally different optimizer than one that has been running since step 1, even if both are given the same parameters and the same learning rate.

When training with mixed precision (covered in the Mixed Precision Training chapter), a gradient scaler maintains a dynamic loss scale factor $s$ that prevents underflow in FP16 gradients. The scaler multiplies the loss by $s$ before backpropagation, producing gradients that are large enough to be representable in FP16. After backpropagation, it divides the gradients by $s$ before the optimizer step, restoring them to the correct scale.

The value of $s$ is adapted during training. If gradients overflow (producing NaN or Inf), the scaler reduces $s$ by a factor and skips the optimizer step. If gradients remain finite for many consecutive steps, the scaler increases $s$ to use more of the FP16 dynamic range. The scaler also tracks how many consecutive steps have been free of overflow, using this to decide when to attempt a scale increase.

If you do not checkpoint the scaler state, the resumed run starts with the default initial scale rather than the converged scale, causing the scaler to go through its warmup phase again and potentially skipping optimizer updates due to overflow checks. The scaler state is tiny, typically just a single float and a few integers, but omitting it causes visible artifacts in the loss curve immediately after resuming. The loss will be flat for a few hundred steps while the scaler finds its footing, then resume its original trajectory.

The learning rate scheduler tracks the current step count to implement schedules like cosine annealing or warmup followed by decay. The scheduler does not automatically infer the current step from the model weights; it maintains its own counter. The counter is what maps from step number to learning rate value.

For a cosine annealing schedule with a warmup period, the learning rate at step $t$ is:

$$\eta_t = \begin{cases} \eta_{\max} \cdot \dfrac{t}{t_{\text{warmup}}} & \text{if } t < t_{\text{warmup}} \\[6pt] \eta_{\min} + \dfrac{1}{2}(\eta_{\max} - \eta_{\min})\left(1 + \cos\!\left(\pi \cdot \dfrac{t - t_{\text{warmup}}}{T - t_{\text{warmup}}}\right)\right) & \text{otherwise} \end{cases}$$

where:

• $\eta_{\max}$: peak learning rate, reached at the end of the warmup period
• $\eta_{\min}$: minimum learning rate, approached at the end of the cosine decay
• $t_{\text{warmup}}$: number of warmup steps
• $T$: total training steps
• $t$: current step count, stored in the scheduler state

If you restore the model and optimizer but not the scheduler, the step counter resets to zero. For a cosine schedule with a long decay, this means the resumed run may suddenly jump from a very low learning rate (late in the cosine decay) back to the initial warmup value. A sudden jump from $10^{-5}$ to $3 \times 10^{-4}$ midway through training corrupts the training dynamics. The model may never fully recover to the trajectory it would have followed if uninterrupted, or it may recover only after several thousand steps of wasted compute.

This is the subtlest item. Every GPU maintains its own RNG state, which is consumed by operations like dropout, weight initialization, and stochastic data augmentation. If you do not restore the RNG state at checkpoint time, the sequence of random numbers after resumption will differ from what it would have been in the uninterrupted run.

For most practical purposes this does not affect final accuracy, because different random seeds rarely change it. The loss curves of two training runs that start from the same checkpoint but use different random seeds will be nearly identical in expectation. But it does affect reproducibility. If you need to reproduce a training run exactly (for debugging a divergence, for example, where you want to determine whether the divergence is caused by a specific gradient update or by random variation), you must save and restore the RNG state on every device.

In PyTorch, each GPU has its own CUDA RNG state, separate from the CPU RNG state. A complete checkpoint saves both:

Restoring both of these at load time ensures that every dropout mask, every random data augmentation, and every stochastic operation unfolds in exactly the same order as it would have in the uninterrupted run.

The data loader position determines which examples have been seen and which have not. Without it, resumption may either repeat examples already processed (wasting compute and potentially biasing the model toward examples that appear twice in the same epoch) or skip examples that have not yet been seen (creating an epoch that is effectively shorter than intended).

For large-scale training, datasets are typically pre-shuffled and streamed in order. Restoring the data loader state means recording the global step and reconstructing which shard files and which positions within those shards were being read. For a simple case where each step consumes exactly $B$ tokens from a sequential stream, you can reconstruct the position from the step number: after step $N$ with batch size $B$, you have consumed $N \times B$ tokens. You resume by skipping the first $N \times B$ tokens of the stream.

This approach works well for deterministic, sequential data loading. It becomes more complex with streaming datasets or on-the-fly shuffling, which may require replaying the shuffle RNG state as well. Some frameworks sidestep this complexity by using a deterministic mapping from step number to data sample: a fixed global shuffle computed before training begins. Given the step number, you can always compute exactly which samples belong to that step without tracking any additional state beyond the step counter.

The cost of getting this wrong is not always obvious immediately. If your data loader resumes from the wrong position, you may train on a subtly different distribution of examples than intended. For large enough datasets, this is usually harmless. For smaller datasets where seeing the exact right examples matters, or for datasets with known challenging examples that you want to ensure are seen at the right training stage, it matters more.

A complete checkpoint contains:

• Model parameters: weights and biases for all layers
• Optimizer state: first and second moment estimates for each parameter (for Adam)
• Gradient scaler state: loss scale and overflow counters (for mixed precision training)
• Scheduler state: current step count and any internal counters
• RNG state: per-device random number generator state (for full reproducibility)
• Data loader state: global step index and dataset position
• Training metadata: step number, epoch, wall-clock time, configuration hash

The training metadata is not strictly necessary for exact resumption, but it is invaluable for bookkeeping and debugging. Knowing the wall-clock time helps you estimate how long the run took; the configuration hash lets you verify you are resuming with the same hyperparameters and not accidentally loading a checkpoint from a different experiment. The step number lets you verify that the recovery logic loaded the right checkpoint and that the scheduler counter is consistent with it.