Checkpoint with optimizer state in PyTorch - Deep Dive

Overview - Checkpoint with optimizer state

What is it?

Checkpointing with optimizer state means saving both the model's learned parameters and the optimizer's internal settings during training. This allows you to pause and later resume training exactly where you left off. Without saving the optimizer state, resuming training might not continue learning properly because the optimizer loses track of its progress.

Why it matters

Saving optimizer state solves the problem of interrupted training sessions, such as power failures or time limits on computers. Without it, you would have to start training from scratch or lose the benefits of previous learning steps. This saves time, computing resources, and helps build better models faster.

Where it fits

Before learning checkpointing with optimizer state, you should understand basic PyTorch model training and saving/loading model weights. After this, you can explore advanced training techniques like learning rate scheduling, mixed precision training, and distributed training that also rely on checkpointing.

Mental Model

Core Idea

Checkpointing with optimizer state saves both the model's parameters and the optimizer's progress so training can resume seamlessly.

Think of it like...

It's like saving a video game where you not only save your character's position but also your inventory and current mission progress, so when you reload, you continue exactly where you left off.

┌─────────────────────────────┐
│        Checkpoint File       │
├─────────────┬───────────────┤
│ Model State │ Optimizer State│
│ (weights)   │ (learning rate,│
│             │ momentum, etc) │
└─────────────┴───────────────┘

Build-Up - 7 Steps

1

Foundation

Saving model parameters only

Concept: Learn how to save just the model's weights during training.

In PyTorch, you save the model's parameters using torch.save(model.state_dict(), 'model.pth'). This stores the learned weights but not the optimizer's state.

Result

A file named 'model.pth' containing the model's weights is created.

Understanding how to save model weights is the first step before adding optimizer state to checkpointing.

2

Foundation

Loading model parameters only

Concept: Learn how to load saved model weights back into a model.

You load weights with model.load_state_dict(torch.load('model.pth')). This restores the model's parameters but does not restore optimizer progress.

Result

The model has the saved weights loaded and is ready for inference or further training.

Loading weights alone is not enough to resume training perfectly because optimizer state is missing.

3

Intermediate

Saving optimizer state with model

🤔Before reading on: do you think saving only model weights is enough to resume training perfectly? Commit to yes or no.

Concept: Learn to save both model weights and optimizer state together.

You create a dictionary with keys 'model_state' and 'optimizer_state' and save it: torch.save({'model_state': model.state_dict(), 'optimizer_state': optimizer.state_dict()}, 'checkpoint.pth'). This saves all necessary info to resume training.

Result

A checkpoint file 'checkpoint.pth' contains both model and optimizer states.

Saving optimizer state captures the optimizer's internal variables like momentum, which are crucial for continuing training smoothly.

4

Intermediate

Loading optimizer state with model

🤔Before reading on: do you think loading optimizer state is as simple as loading model weights? Commit to yes or no.

Concept: Learn to load both model and optimizer states from a checkpoint.

Load checkpoint = torch.load('checkpoint.pth') then model.load_state_dict(checkpoint['model_state']) and optimizer.load_state_dict(checkpoint['optimizer_state']). This restores both model and optimizer to the saved state.

Result

Training can resume exactly from the saved point with optimizer progress intact.

Loading optimizer state ensures the optimizer continues updating weights correctly, avoiding training disruptions.

5

Intermediate

Checkpointing during training loop

Concept: Learn how to save checkpoints periodically during training.

Inside the training loop, after some epochs, save checkpoint with model and optimizer states. Example: torch.save({'model_state': model.state_dict(), 'optimizer_state': optimizer.state_dict(), 'epoch': epoch}, 'checkpoint.pth'). This allows resuming from the last saved epoch.

Result

You have multiple checkpoints to resume training from different points.

Periodic checkpointing protects against data loss and allows flexible training management.

6

Advanced

Handling device compatibility in checkpoints

🤔Before reading on: do you think loading a checkpoint saved on GPU will always work on CPU? Commit to yes or no.

Concept: Learn how to save and load checkpoints that work across CPU and GPU devices.

When loading, use map_location=torch.device('cpu') if loading on CPU. Example: torch.load('checkpoint.pth', map_location=torch.device('cpu')). This avoids errors when device differs between saving and loading.

Result

Checkpoint files become portable across different hardware setups.

Handling device mapping prevents common errors and makes checkpoints flexible for different environments.

7

Expert

Checkpointing with learning rate schedulers

🤔Before reading on: do you think saving optimizer state alone is enough when using learning rate schedulers? Commit to yes or no.

Concept: Learn to save and restore learning rate scheduler state along with model and optimizer.

Extend checkpoint dict to include 'scheduler_state': scheduler.state_dict(). Save and load it similarly. This ensures learning rate adjustments continue correctly after resuming.

Result

Training resumes with correct learning rate schedule, avoiding sudden jumps or drops.

Saving scheduler state prevents subtle training issues that degrade model performance after resuming.

Under the Hood

PyTorch models and optimizers store their internal states as Python dictionaries of tensors and variables. When you call state_dict(), it returns these dictionaries. Saving them with torch.save serializes these dictionaries to disk. Loading restores these dictionaries into the model and optimizer objects, preserving all internal variables like weights, momentum buffers, and learning rates. This allows training to continue exactly where it left off.

Why designed this way?

The state_dict design separates model parameters and optimizer states into simple dictionaries, making saving and loading flexible and transparent. This design avoids saving entire objects, which can cause compatibility issues. It also allows users to customize what to save and supports partial loading. Alternatives like saving entire objects were less flexible and more error-prone.

┌───────────────┐       ┌───────────────┐
│ Model Object  │       │ Optimizer Obj │
├───────────────┤       ├───────────────┤
│ state_dict()  │       │ state_dict()  │
│ ────────────▶│       │ ────────────▶│
│ Dict of Tensors│       │ Dict of Vars  │
└───────────────┘       └───────────────┘
         │                       │
         │                       │
         └──────────────┬────────┘
                        │
                torch.save(dict) 
                        │
                Serialized file
                        │
                torch.load(file)
                        │
         ┌──────────────┴────────┐
         │                       │
 model.load_state_dict()   optimizer.load_state_dict()

Myth Busters - 4 Common Misconceptions

Quick: If you save only the model weights, can you resume training with the same optimizer progress? Commit to yes or no.

Common Belief: Saving just the model weights is enough to resume training perfectly.

Reality: You must save and load the optimizer state too, or training will not continue correctly because optimizer variables like momentum are lost.

Why it matters: Without optimizer state, training can become unstable or slower after resuming, wasting time and resources.

Quick: Do you think loading a checkpoint saved on GPU will always work on a CPU-only machine? Commit to yes or no.

Common Belief: Checkpoints are device-independent and can be loaded anywhere without extra steps.

Reality: You must specify device mapping when loading if devices differ, or loading will fail with errors.

Why it matters: Ignoring device mapping causes crashes and confusion, blocking training continuation.

Quick: Is saving the optimizer state enough when using learning rate schedulers? Commit to yes or no.

Common Belief: Optimizer state includes everything needed, so scheduler state does not need saving.

Reality: Schedulers have their own state that must be saved and restored separately to keep learning rates consistent.

Why it matters: Not saving scheduler state leads to unexpected learning rate changes, harming model convergence.

Quick: Do you think checkpoint files always contain the entire training history? Commit to yes or no.

Common Belief: Checkpoint files store all past training data and history.

Reality: Checkpoints only save current states, not the full training history or logs.

Why it matters: Expecting full history in checkpoints can cause confusion when logs or metrics are missing after resuming.

Expert Zone

1

Optimizer state can be large and complex, especially for adaptive optimizers like Adam, so checkpoint size can grow significantly.

2

When using distributed training, checkpointing requires careful synchronization to save consistent states across devices.

3

Partial loading of checkpoints is possible by manipulating state_dicts, allowing fine control over which parts to restore.

When NOT to use

Checkpointing with optimizer state is not needed if you only want to use the model for inference or evaluation. In such cases, saving just the model weights is sufficient and more lightweight.

Production Patterns

In production, checkpointing is often combined with early stopping and best model saving based on validation metrics. Checkpoints are saved periodically and after improvements, enabling robust training pipelines that can recover from failures.

Connections

Version Control Systems

Both checkpointing and version control save states to allow resuming or reverting work.

Understanding checkpointing like version control helps appreciate the importance of saving progress and being able to return to exact points in complex workflows.

Database Transactions

Checkpointing is similar to committing a transaction that saves a consistent state to avoid data loss.

Knowing how databases ensure consistency through transactions helps understand why saving optimizer state is critical for consistent training continuation.

Human Learning and Memory

Checkpointing with optimizer state is like taking notes on both what you learned and how you plan to learn next.

This connection shows that saving both knowledge and learning strategy is essential for effective progress, just like in machine learning.

Common Pitfalls

#1 Saving only model weights and ignoring optimizer state.

Wrong approach:

torch.save(model.state_dict(), 'model.pth')

Correct approach:

torch.save({'model_state': model.state_dict(), 'optimizer_state': optimizer.state_dict()}, 'checkpoint.pth')

Root cause:

Misunderstanding that optimizer state is needed to continue training properly.

#2 Loading checkpoint without device mapping when devices differ.

Wrong approach:

checkpoint = torch.load('checkpoint.pth') # fails if saved on GPU, loaded on CPU

Correct approach:

checkpoint = torch.load('checkpoint.pth', map_location=torch.device('cpu'))

Root cause:

Not accounting for hardware differences between saving and loading environments.

#3 Not saving learning rate scheduler state when using schedulers.

Wrong approach:

torch.save({'model_state': model.state_dict(), 'optimizer_state': optimizer.state_dict()}, 'checkpoint.pth')

Correct approach:

torch.save({'model_state': model.state_dict(), 'optimizer_state': optimizer.state_dict(), 'scheduler_state': scheduler.state_dict()}, 'checkpoint.pth')

Root cause:

Overlooking that schedulers maintain their own internal state separate from optimizer.

Key Takeaways

Checkpointing with optimizer state saves both model parameters and optimizer progress to allow seamless training resumption.

Saving only model weights is insufficient for continuing training because optimizer variables like momentum are lost.

Loading checkpoints requires careful device mapping to avoid errors when hardware differs between saving and loading.

Including learning rate scheduler state in checkpoints ensures consistent training behavior after resuming.

Periodic checkpointing protects training progress from interruptions and supports flexible training workflows.