Lightning-AI / Lightning-AI/pytorch-lightning
ModelCheckpoint: manual-optimization pre-step snapshot is not applied when the save is deferred to on_validation_end
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 31.4k
- Forks
- 3.8k
- Avg merge
- 6d 7h
- Merged PRs (30d)
- 6
Description
### Bug description
`ModelCheckpoint` with manual optimization, `every_n_train_steps`, and a validation-only monitored metric (e.g. `monitor="val_loss"` logged in `validation_step`) saves the checkpoint credited as "best" with the **live, fully-trained, post-optimization weights at the end of the epoch**, not the pre-optimization weights captured for the step that actually produced the best score. This defeats the entire point of the pre-optimization-state-snapshot mechanism added in #21239.
`on_train_batch_end` (`src/lightning/pytorch/callbacks/model_checkpoint.py`, manual-optimization branch) correctly swaps in `pl_module.saved_models[latest_step]` before saving and restores the live weights afterward, when the monitored metric is already available:
```python
# lines 381-397
with torch.no_grad():
original_state = {k: v.detach().clone() for k, v in pl_module.layer.state_dict().items()}
try:
saved_state = saved_models[latest_step]
pl_module.layer.load_state_dict(saved_state)
self._save_topk_checkpoint(trainer, monitor_candidates)
self._save_last_checkpoint(trainer, monitor_candidates)
self._last_time_checked = time.monotonic()
finally:
pl_module.layer.load_state_dict(original_state)
```
But when the monitored metric isn't available yet (the normal case for a validation-only metric during training), the save is deferred instead (line 361-363):
```python
if self.monitor is not None and self.monitor not in monitor_candidates:
self._defer_save_until_validation = True
return
```
And the deferred path, `on_validation_end` (lines 507-511), does the save with **no swap at all**:
```python
if self._defer_save_until_validation:
self._save_topk_checkpoint(trainer, monitor_candidates)
self._save_last_checkpoint(trainer, monitor_candidates)
self._defer_save_until_validation = False
return
```
This calls straight into `_save_checkpoint` -> `trainer.save_checkpoint(...)`, which just dumps whatever the live model state is at that moment, i.e. after every `optimizer.step()` call for the whole epoch (or however many steps ran before validation), not the snapshot corresponding to the step that produced the best monitored value.
`_save_checkpoint`'s own docstring claims: *"For manual optimization, we rely on the fact that the model's training_step method saves the model state before the optimizer step, so we can use that state directly"*, but that guarantee is only actually implemented at the one call site inside `on_train_batch_end`, not in `on_validation_end`.
This is a regression of #20947 ("the best checkpointed model is the model obtained after backpropagation and not the one used for computing the loss"), reintroduced by an incomplete fix in #21239. #21239's own test (`test_model_checkpoint_manual_opt`) only monitors a training-step metric (`monitor="loss"`), which is always present in `on_train_batch_end`'s `monitor_candidates` and therefore never exercises the `_defer_save_until_validation` branch at all, so it never caught that the swap logic doesn't exist in `on_validation_end`.
### How to reproduce the bug
```python
import torch
from torch.utils.data import DataLoader, Dataset
from lightning.pytorch import LightningModule, Trainer
from lightning.pytorch.callbacks import ModelCheckpoint
class FakeDataset(Dataset):
def __init__(self, n=4):
self.data = [torch.randn(3) for _ in range(n)]
self.labels = [torch.randint(0, 2, (1,)) for _ in range(n)]
def __len__(self):
return len(self.data)
def __getitem__(self, idx):
return self.data[idx], self.labels[idx]
class SimpleModule(LightningModule):
def __init__(self):
super().__init__()
self.layer = torch.nn.Linear(3, 1)
self.automatic_optimization = False
self.saved_models = {}
self.fake_val_losses = [1.0, 1.0, 0.0, 1.0] # batch_idx=2 is "best"
def training_step(self, batch, batch_idx):
out = self.layer(batch[0])
loss = torch.nn.functional.binary_cross_entropy_with_logits(out, batch[1].float())
# documented pattern from the library's own warning message: snapshot BEFORE optimizer.step()
self.saved_models[batch_idx] = {k: v.detach().clone() for k, v in self.layer.state_dict().items()}
opt = self.optimizers()
opt.zero_grad()
self.manual_backward(loss)
opt.step()
return loss
def validation_step(self, batch, batch_idx):
self.log("val_loss", torch.tensor(self.fake_val_losses[batch_idx]), on_epoch=True)
def configure_optimizers(self):
return torch.optim.SGD(self.parameters(), lr=0.5)
model = SimpleModule()
trainer = Trainer(
max_epochs=1,
callbacks=[ModelCheckpoint(
monitor="val_loss", mode="min", save_top_k=1, every_n_train_steps=1,
every_n_epochs=0, save_on_train_epoch_end=False, save_weights_only=True, save_last=False,
)],
num_sanity_val_steps=0,
)
trainer.fit(model, DataLoader(FakeDataset(), batch_size=1), DataLoader(FakeDataset(), batch_size=1))
best_ckpt = torch.load(trainer.checkpoint_callback.best_model_path, weights_only=True)["state_dict"]
def matches(saved_state, ckpt):
return all(torch.equal(saved_state[name.removeprefix("layer.")], val) for name, val in ckpt.items())
any_pre_opt_match = any(matches(s, best_ckpt) for s in model.saved_models.values())
live_weights = {k: v.detach().clone() for k, v in model.layer.state_dict().items()}
matches_live = all(torch.equal(live_weights[name.removeprefix("layer.")], val) for name, val in best_ckpt.items())
print("Checkpoint matches ANY pre-optimization snapshot (expected):", any_pre_opt_match)
print("Checkpoint matches the LIVE, fully-trained end-of-epoch weights instead:", matches_live)
```
Actual output:
```
Checkpoint matches ANY pre-optimization snapshot (expected): False
Checkpoint matches the LIVE, fully-trained end-of-epoch weights instead: True
```
The saved "best" checkpoint doesn't match any of the four pre-optimization snapshots captured during the epoch (it should match the one from the step with the best `val_loss`), it matches the fully-trained, post-all-optimizer-steps weights instead.
### What version are you seeing the problem on?
master, v2.6
### How to reproduce the bug
(see script above)
### Environment
- PyTorch Lightning Version: 2.6.5 (pip install lightning), confirmed the same logic is present on current `master` by reading `src/lightning/pytorch/callbacks/model_checkpoint.py` directly
- PyTorch Version: 2.x, CPU only, no GPU needed to reproduce
- Python version: 3.12.10
- OS: Windows 11
- How installed: pip
### More info
Related: #20947 (the original bug), #21239 (the fix that introduced this regression by only covering the non-deferred `on_train_batch_end` path). A fix would need to apply the same pre-optimization-state swap in `on_validation_end`'s `_defer_save_until_validation` branch that already exists in `on_train_batch_end`.
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start in src/lightning/pytorch/callbacks/model_checkpoint.py, comparing the pre-optimization snapshot swap in on_train_batch_end with the deferred _defer_save_until_validation path in on_validation_end. Extend test_model_checkpoint_manual_opt to cover a validation-only monitor and verify that the best checkpoint matches the pre-optimization snapshot for the best validation step.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- machine-learning
- Issue type
- Bug
- Difficulty
- 3/5
- Estimated time
- 1-2 days
- Activity status
- Quiet
- Clarity
- Clearly specified
- Newbie friendliness
- 74/100