Lightning-AI / Lightning-AI/pytorch-lightning

ModelCheckpoint: manual-optimization pre-step snapshot is not applied when the save is deferred to on_validation_end

Open
#21,842 3 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

callback: model checkpoint ver: 2.6.x ver: 2.7.x
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

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.