Anner-deJong
commented
2 years ago I'd be happy to help fix if agreed this is a bug, but like the opinion + take on a solution from folks more involved with pytorch lightning. Some simple ideas, but none of them great:
- Give the model checkpoint callback a special status compared to others, call it at the end of
training_epoch_loop.on_advance_end- Sounds like this really shouldnt be required + makes codes fragile/ harder to maintain
- have a relevant hook called at the beginning of
training_epoch_loop.on_advance_endthat checkpoints- why would we need yet another hook? And if we move one, I assume any hooks location in code flow was well thought out already
- call
on_train_batch_endafter theincrement_completed- definition based: does _on_train_batchend mean it wont change training module anymore and could be run after the increment? Is there logic in the on_train_batch_end that expects "completed" not to be incremented yet?
krshrimali
lantiga
🐛 Bug
Found a discrepancy between a continued run after checkpointing, and restoring from checkpoint
Observation:
training_batch / val_loop ordering upon checkpoint restoration not the same as original run after checkpoint saving.
There are still the same amount of train steps, but the validation loops are interleaved at a single step later, which can cause the restored run to end up with one less validation loop (see colab)
Assumption / expectation:
Zero difference between a training run after a checkpoint and a run continued from said checkpoint
Investigation so far:
Im new to some of this lightning code, but IIUC:
Key:
TrainingEpochLoop'sself.batch_progress.increment_completed()is called after"on_train_batch_end"hooks, the latter kicking off checkpoint saving.TrainingEpochLoop.batch_progress.current.reset_on_restart()will reset thereadyback tocompletedglobal_step, which refers toTrainingEpochLoop.batch_loop.optimizer_loop.optim_progress.optimizer.step.total.completed, has beenincrement_completed()(called withinTrainingEpochLoop.batch_loop.run) and thus upon restoring,..optimizer.step.total.readyis set to an up to dateoptimizer.step.total.completed, out of sync with the aboveval_check_intervalmode", validation is triggered whenTrainingEpochLoop.batch_progress.current.ready % val_check_interval == 0(throughTrainingEpochLoop.on_advance_end->TrainingEpochLoop._should_check_val_fxbatch_progress.currentready/completedcounter for the continued and restored runs, end up aligned with differentglobal_steps, and hence validation triggers at differentglobal_stepsAnother observation:
The following if statement seems to allow for a zero-difference restart, except that just like 4. above,
_should_check_val_fxwouldnt trigger where in the original run on the checkpointing step it did (although there called inon_advance_end). Not sure if the original intention of this snippet included the current scopePR's relevant to this line:
9681
9563
Potential impact:
Assuming not too worrisome for the more default Lightning use cases:
val_check_interval>> 3 (colab example = 3), or that turned off relying instead oncheck_val_every_n_epochHowever, in theory it can influence all of the following:
To Reproduce
customized google colab
bug_report_model.ipynbwith same observation onBoringModelExpected behavior
Zero difference between a training run continued after a checkpoint and a run continued from said checkpoint
Environment
Note:
v1.6.4.v1.7.4masterbranch last weeks and the relevant code seems unchangedDetails
CUDA:
Lightning:
Packages:
System:
Additional context
cc @awaelchli @ananthsub @ninginthecloud @rohitgr7 @otaj @carmocca @justusschock