Checkpointing is the practice of periodically saving a model's complete training state — weights, optimizer state, epoch number, learning rate schedule, and training metrics — during training — enabling crash recovery (resume from hour 70 of a 72-hour training run instead of restarting from scratch), best model selection (save the model with the lowest validation loss, which is often NOT the final epoch), and preemption resilience (cloud spot instances can be killed at any time, and checkpoints are insurance against lost compute).
What Is Checkpointing?
- Definition: The periodic serialization of a model's complete training state to disk, enabling training to be paused, resumed, or rewound to any saved state — including the model weights, optimizer momentum/state, current epoch, learning rate scheduler state, and random number generator seeds.
- Why Just Saving Weights Isn't Enough: If you only save model weights, you can't resume training correctly because the optimizer's momentum buffers (Adam: running mean and variance of gradients) are lost, the learning rate scheduler doesn't know what epoch to resume from, and the data loader doesn't know which batches were already seen.
- The Cost of Not Checkpointing: A 3-day GPU training run costs ~$500-$2000 on cloud. A single power outage, OOM error, or spot instance preemption without checkpointing means starting over completely.
What to Save
| Component | Why It Matters | What Happens Without It |
|---|---|---|
| model.state_dict() | Model weights and biases | Can't use or resume the model at all |
| optimizer.state_dict() | Momentum, adaptive learning rates (Adam) | Optimizer restarts cold → training diverges |
| epoch / step | Current position in training | Don't know where to resume |
| scheduler.state_dict() | Learning rate schedule position | LR schedule restarts → wrong learning rate |
| best_val_metric | Best validation score seen | Can't determine if new checkpoints are improvements |
| RNG states | Random seeds for reproducibility | Non-reproducible training |
PyTorch Implementation
# Save checkpoint
def save_checkpoint(model, optimizer, scheduler, epoch, best_val, path):
torch.save({
'epoch': epoch,
'model_state_dict': model.state_dict(),
'optimizer_state_dict': optimizer.state_dict(),
'scheduler_state_dict': scheduler.state_dict(),
'best_val_loss': best_val,
}, path)
# Load checkpoint
checkpoint = torch.load('checkpoint.pt')
model.load_state_dict(checkpoint['model_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
scheduler.load_state_dict(checkpoint['scheduler_state_dict'])
start_epoch = checkpoint['epoch'] + 1
Checkpointing Strategies
| Strategy | When to Save | Use Case |
|---|---|---|
| Every N epochs | Save every 5 or 10 epochs | Standard training — periodic insurance |
| Best only | Save only when validation metric improves | Long training — disk space efficient |
| Last + Best | Keep most recent + best validation | Resume from latest OR use best |
| Top K | Keep K best checkpoints | Model selection, ensemble from top-K |
| Every step | Save after every batch/step | Very expensive training (LLMs) |
Framework Support
| Framework | Checkpointing API |
|---|---|
| PyTorch | torch.save() / torch.load() (manual) |
| PyTorch Lightning | ModelCheckpoint callback (automatic) |
| Keras | ModelCheckpoint callback (automatic) |
| Hugging Face | Trainer(save_strategy="epoch") (automatic) |
| DeepSpeed | Built-in distributed checkpointing |
Checkpointing is the essential training infrastructure that protects against compute loss — saving the complete training state periodically so that expensive GPU hours are never wasted due to crashes, preemption, or hardware failures, while simultaneously enabling best-model selection by preserving the weights from the optimal validation epoch rather than the final (potentially overfit) epoch.
Explore 500+ Semiconductor & AI Topics
From EUV lithography to CUDA optimization — search the full knowledge base or chat with our AI assistant.