enable_pl_optimizer causes optimizers to not be restored properly

[PiotrDabkowski]

🐛 Bug

enable_pl_optimizer (default!) causes optimizers to not be restored properly from the checkpoint specified by resume_from_checkpoint.

BoringModel Colab Reproduction

The model is trained for 3 epochs and saved in a checkpoint. The checkpoint is then restored and further trained for 1 epoch (with different values of enable_pl_optimizer), the training loss is printed at each step.
The setup where enable_pl_optimizer=True shows a huge loss spike after the first optimizer step, suggesting that the optimizer is not restored properly.

https://colab.research.google.com/drive/1lHYXm4MpnmXwPZTcPem4D4wwwU5vJhHc?usp=sharing

Expected behavior

PL Optimizers are restored such that there is no huge loss spike after restore, just like when enable_pl_optimizer=False.

Environment

See Colab.

[github-actions]

Hi! thanks for your contribution!, great first issue!

[heng-yuwen]

Hi, I think you can check #4655. I think these problems are related. @PiotrDabkowski

[heng-yuwen]

The loaded optimiser is restored as the initial one.

[PiotrDabkowski]

Possibly, but not sure. Does disabling pl_optimizer fix #4655 as well?

[tchaton]

Hey @PiotrDabkowski,

enable_pl_optimizer has been reset to False by default.

We will look into this bug.

Thanks,
T.C

[heng-yuwen]

I didn't try to disable enable_pl_optimizer, @tchaton is there anything I should take care of when enable_pl_optimizer=False? I use ddp to train on Slurm cluster.

[tchaton]

Hey @Hyw1994,

I checked your notebook ! You are entirely right.
I will look into today as it is high priority to re-enable LightningOptimizer :)

Best regards,
T.C