Lightning-AI / Lightning-AI/pytorch-lightning

Error learning rate when load ckpt for cotinue training if check_val_every_n_epoch > 1

Open
#20,495 3 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

bug callback: lr monitor checkpointing ver: 2.4.x
Dominant language
Python
Stars
31.4k
Forks
3.8k
Avg merge
6d 7h
Merged PRs (30d)
6

Description

### Bug description

I run with LightningCLI. Set check_val_every_n_epoch > 1 (e.g. 2) to run an experiment with 20 max_epoches, the model ckpt is save by lightning.pytorch.callbacks.ModelCheckpoint. The learning rate schedular is torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(optimizer, T_0=self.trainer.max_epochs, T_mult=1, eta_min=self.eta_min) and update each epoch. When I load a ckpt (e.g. saved at epoch 3) to continue training, the learning rate will update 1 epoch quicker than expected.
![image](https://github.com/user-attachments/assets/45ce4ea3-9c86-46ce-8316-133b992336eb)
Here the red is original learning rate curve and the yellow is the continued. The lr is logged by lightning.pytorch.callbacks.LearningRateMonitor.

### What version are you seeing the problem on?

v2.4

### How to reproduce the bug

```yaml
the trainer setting
seed_everything: 0
trainer:
accelerator: gpu
strategy: auto
devices:
- 0
num_nodes: 1
precision: 32-true
logger:
class_path: lightning.pytorch.loggers.TensorBoardLogger
init_args:
save_dir: runs
name: name
version: null
log_graph: false
default_hp_metric: true
prefix: ''
sub_dir: test
callbacks:
- class_path: utils.log_manager.LogManager
- class_path: lightning.pytorch.callbacks.LearningRateMonitor
init_args:
logging_interval: epoch
- class_path: lightning.pytorch.callbacks.ModelCheckpoint
init_args:
dirpath: null
filename: 'epoch={epoch}-psnr={hp_metric:.4f}'
monitor: hp_metric
verbose: false
save_last: true
save_top_k: 5
save_weights_only: false
mode: max
auto_insert_metric_name: false
every_n_train_steps: null
train_time_interval: null
every_n_epochs: 1
save_on_train_epoch_end: null
fast_dev_run: false
max_epochs: 20
min_epochs: null
max_steps: -1
min_steps: null
max_time: null
limit_train_batches: null
limit_val_batches: null
limit_test_batches: null
limit_predict_batches: null
overfit_batches: 0.0
val_check_interval: null
check_val_every_n_epoch: 2
num_sanity_val_steps: null
log_every_n_steps: 20
enable_checkpointing: null
enable_progress_bar: null
enable_model_summary: null
accumulate_grad_batches: 1
gradient_clip_val: null
gradient_clip_algorithm: null
deterministic: null
benchmark: null
inference_mode: true
use_distributed_sampler: true
profiler: null
detect_anomaly: false
barebones: false
plugins: null
sync_batchnorm: false
reload_dataloaders_every_n_epochs: 0
default_root_dir: null
```
```py
def configure_optimizers(self):
params = []
params.append({'params': self.network.parameters(), 'lr': self.lr, 'weight_decay': self.weight_decay, 'name': 'network'})
optimizer = torch.optim.AdamW(params)
lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(optimizer, T_0=self.trainer.max_epochs, T_mult=1, eta_min=self.eta_min)
lr_scheduler_config = {
"scheduler": lr_scheduler,
"interval": 'epoch',
"name": 'AdamW'
}
return {"optimizer": optimizer, "lr_scheduler": lr_scheduler_config}
```

### Error messages and logs

```
# Error messages and logs here please
```

### Environment

Current environment

```
#- PyTorch Lightning Version (e.g., 2.4.0):
#- PyTorch Version (e.g., 2.4):
#- Python version (e.g., 3.12):
#- OS (e.g., Linux):
#- CUDA/cuDNN version:
#- GPU models and configuration:
#- How you installed Lightning(`conda`, `pip`, source):
```

### More info

_No response_

cc @lantiga

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 with the provided configure_optimizers scheduler configuration and the Trainer resume flow using ModelCheckpoint; reproduce with check_val_every_n_epoch: 2 and LearningRateMonitor. Compare scheduler state and epoch progression before and after loading the epoch-3 checkpoint, with the fix validated when continued learning-rate curves align with uninterrupted training.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
machine-learning
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
45/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.