Global step reset when restoring checkpoints with trainer.validate

[nicolas-dufour]

Bug description

When restoring checkpoints with trainer.validate, global_step and epoch are overwritten with 0.

It should keep the same global_step and epoch otherwise, it messes with the loggers.

This issue prevents to correctly validate checkpoints of a model as a postprocessing.

How to reproduce the bug

run a model with

trainer.validate(model, datamodule, cpt_path=ckpt_path

and log a metric, the result will be logged at step 0.

Environment

Current environment

#- Lightning Component (e.g. Trainer, LightningModule, LightningApp, LightningWork, LightningFlow):
#- PyTorch Lightning Version (e.g., 1.5.0): 1.9
#- Lightning App Version (e.g., 0.5.2):
#- PyTorch Version (e.g., 2.0): 1.13
#- Python version (e.g., 3.9): 3.10
#- OS (e.g., Linux): Linux
#- CUDA/cuDNN version: 11.7
#- GPU models and configuration:
#- How you installed Lightning(`conda`, `pip`, source):
#- Running environment of LightningApp (e.g. local, cloud):

More info

No response

cc @Borda @awaelchli @carmocca @justusschock

[carmocca]

The global step and epoch counts are part of the fit loop state.

When you load a checkpoint, we only load the state relevant to the current trainer method called https://github.com/Lightning-AI/lightning/blob/1156eb74aef23fa1a43069d20e652ba5a9dfd77e/src/lightning/pytorch/trainer/connectors/checkpoint_connector.py#L331-L334

This explains why trainer.validate() with a new Trainer instance won't have these attributes.

One workaround for this would be to reuse the same trainer instance you used for fitting previously.

We could also load all these states regardless of the trainer method running. Currently, the loops don't define any expensive-to-load states, so this wouldn't be very inefficient.

Another option would be to treat these two attributes differently and load them separately. Irrespective of the trainer's method running.

Any thoughts? cc @awaelchli @justusschock

[nicolas-dufour]

Humm, i get that the training state don't need to be shared with the test state, but it's important to be able to recover which training step we are validating on.

For now, i restore the global_step and epoch from the ckpt and hardcode them to my validate call but i feel this should be part of the native behaviour of validate

[basveeling]

To chime in: I have the same issue. I would like to perform a fit command using the cli, followed by a validate command in a different job, ideally adding new metrics to the original tensorboard log created by the fit command. The following monkey patch fixes the issue for me:

https://github.com/Lightning-AI/lightning/blob/1156eb74aef23fa1a43069d20e652ba5a9dfd77e/src/lightning/pytorch/trainer/connectors/checkpoint_connector.py#L331-L334

To

    if state_dict is not None:
        self.trainer.fit_loop.load_state_dict(state_dict["fit_loop"])
        self.trainer.validate_loop.load_state_dict(state_dict["validate_loop"])
        self.trainer.test_loop.load_state_dict(state_dict["test_loop"])
        self.trainer.predict_loop.load_state_dict(state_dict["

[bkmi]

This is also true with trainer.test btw.

Will attempt to proceed using this:

For now, i restore the global_step and epoch from the ckpt and hardcode them to my validate call but i feel this should be part of the native behaviour of validate