Checkpointing and Recovery › Fault Recovery
Language AI HandbookAll chapters Checkpointing and Recovery › Fault RecoveryOne section · Sections in this chapter5 / 9 ↓
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
A complete checkpoint is only half the system. The other half is the recovery logic that detects failures, loads the checkpoint, and resumes training with minimal disruption.
The checkpoint and recovery subsystems must be designed together. A checkpoint saved in a format that is slow to load creates a system where failures are infrequent but recovery is painful. A checkpoint format that loads quickly but is expensive to write creates a system where checkpointing overhead limits how often you can protect your work. Production training frameworks co-design these two halves, often using sharded file formats that write and load in parallel across many storage nodes, and maintaining a lightweight metadata store that makes "find the latest valid checkpoint" a constant-time operation rather than a directory scan.
Failure Detection
Distributed training frameworks use heartbeat mechanisms and timeouts to detect failures. When a process crashes, adjacent processes eventually time out waiting for communication and raise an exception. The NCCL collective communication library, which handles gradient synchronization in most GPU training setups, has its own timeout mechanism that triggers when a collective operation does not complete within a configurable window.
The time from failure to detection varies considerably. NCCL timeouts are typically set to 1,800 seconds (30 minutes) to tolerate temporary network congestion without false positives. A slow network switch, a brief NIC hiccup, or a transient memory allocation stall can all look like a hung collective operation from NCCL's perspective, and triggering a full job restart in response to a 5-second network blip would be extremely disruptive. The conservative default timeout prevents this.
But 30 minutes of hanging is an enormous waste. A job that checkpoints every 10 minutes and then hangs for 30 minutes after each failure is losing more time to failure detection than to failed work. Modern training setups are moving toward faster detection using in-band monitoring: a watchdog process that actively monitors process health, GPU utilization, and collective operation latency, and triggers recovery within seconds rather than waiting for NCCL to time out.
The watchdog approach requires careful calibration. A watchdog that is too sensitive will trigger false restarts during legitimate pauses (for example, during a checkpoint write or a data loading bottleneck). A watchdog that is too conservative is no improvement over NCCL's default timeout. The right setting depends on the specific training workload and the network characteristics of the cluster. Teams running production training often tune their watchdog parameters empirically based on the failure patterns they observe.
When the orchestrator (SLURM, Kubernetes, or a custom job manager) detects the failed job, it terminates all running ranks and begins the recovery sequence. Terminating all ranks is necessary even if only one has failed, because the surviving ranks are hung waiting for the failed rank's contributions to collective operations.
Recovery Workflow
The recovery workflow after a failure typically proceeds in a specific order. Understanding each step and its potential failure modes helps you build a reliable recovery system rather than one that fails silently.
The orchestrator detects the failure and terminates all ranks in the job. It then determines the most recent valid checkpoint by reading the checkpoint metadata (the "latest checkpoint" pointer file). It restarts the job with fresh processes on the available, non-failed hardware. Each new process initializes its model and optimizer to default values, then loads the checkpoint, restoring all state to the values at the checkpoint step. Training resumes from the checkpointed step.
The step "restart on available hardware" deserves elaboration. If a GPU has permanently failed (a common scenario with high-memory-bandwidth GPU memory errors), the replacement hardware may not be immediately available. The job may run with fewer GPUs than the original, which requires adjusting the parallelism configuration. In elastic training setups (discussed in the context of PyTorch Elastic), the restarted job may run with fewer GPUs than the original, adjusting the batch size or gradient accumulation steps accordingly to maintain the effective batch size.
Maintaining the effective batch size is important for training dynamics. If you trained with an effective batch size of 2,048 and you restart with 25% fewer GPUs, you have two choices: reduce the batch size to 1,536 (which may change training dynamics if you are near a threshold where batch size affects convergence) or increase gradient accumulation steps to compensate (which maintains the effective batch size at the cost of slower throughput). Most production systems prefer the gradient accumulation approach.
Loading State
Loading a checkpoint involves the reverse of saving: reading sharded files from storage, distributing the appropriate shards to each rank, and populating the model and optimizer state dictionaries.
The straightforward case is loading into the same parallelism configuration as the original: the same number of ranks, the same tensor parallelism degree, the same FSDP sharding strategy. In this case, each rank simply reads its own shard file and loads it directly into the appropriate tensors. This is fast and requires no redistribution.
The complex case arises when recovering onto different hardware with a different parallelism configuration. If the original run used 64 GPUs in a 4-way tensor parallel, 16-way data parallel configuration and you restart with 48 GPUs (after losing 16 to hardware failure), the shard layout has changed. PyTorch's distributed checkpoint format handles this through resharding: it reads all shards and redistributes the tensors according to the new sharding plan, without requiring the checkpoint to be re-saved in the new format.
Resharding is computationally inexpensive but communication-intensive: each rank must receive tensor slices from multiple source ranks to assemble its new shard. In practice, this takes a few minutes for large models and is acceptable as a one-time cost per recovery event.
The load_state_dict approach for a FSDP model using distributed checkpointing:
import torch.distributed.checkpoint as dcp
from torch.distributed.checkpoint.state_dict import (
StateDictOptions,
set_state_dict,
)
# After model and optimizer are created with new parallelism config:
state_dict = {"model": model.state_dict(), "optimizer": optimizer.state_dict()}
dcp.load(
state_dict=state_dict,
storage_reader=dcp.FileSystemReader("/checkpoints/step_1000"),
)
set_state_dict(
model,
optimizer,
model_state_dict=state_dict["model"],
optim_state_dict=state_dict["optimizer"],
options=StateDictOptions(full_state_dict=False), # keep sharded
)After loading, restore the scheduler, scaler, and data loader state from the metadata in the checkpoint before calling optimizer.step(). The order matters: the scheduler must be restored before the optimizer, because some scheduler implementations modify the optimizer's learning rate in their step() method, and you want to restore to the correct learning rate, not override it with an incorrect one.
Verifying Recovery
After loading a checkpoint, it is good practice to verify that the resumed run matches expectations before continuing for thousands of steps. A corrupted or incomplete checkpoint can cause training to silently diverge, wasting compute and potentially making the model worse. Catching this early is far cheaper than discovering it after 12 hours of additional training.
Common verification steps include:
- Confirm the global step counter matches the checkpoint metadata
- Run one forward pass and compare the loss to the logged loss at that step
- Verify the learning rate matches what the scheduler should produce at that step
- Check that gradient norms in the first resumed step are consistent with the run history
A mismatch in any of these suggests an incomplete or corrupt checkpoint, and it is better to roll back to an earlier checkpoint than to continue with a corrupted state.
Why run a forward pass and compare to logged loss? During a training run, the loss at each step is typically logged to a metrics system. After recovery, if you run the same batch of data (using the restored data loader position and RNG state) and compute the loss, you should get the same value as what was logged at that step. An exact match confirms that the model parameters and all randomness-affecting state have been correctly restored. Even a small discrepancy (beyond floating-point rounding) indicates something was missed in the checkpoint.
This kind of checkpoint integrity check is particularly important in long-running research jobs where silent errors can propagate undetected for hours before manifesting as anomalous evaluation metrics. The forward pass check is fast (a single batch, no backward pass needed), and the cost of skipping it is potentially many hours of corrupt training.
Some teams automate this verification step into the recovery logic itself: the recovery script runs the forward pass check, and if it fails, it automatically tries the previous checkpoint in the retention window. This removes the need for human intervention in most recovery scenarios and reduces mean recovery time from tens of minutes (waiting for a human to notice and act) to seconds.