DeepSpeed breaks learning rate scheduler checkpointing

🐛 Bug

Activating DeepSpeed prevents learning rate schedulers configured in configure_optimizers from being loaded properly from checkpoints. I've confirmed using my own setup that activating/deactivating DeepSpeed reintroduces/fixes the issue, respectively:

I'm using a simple learning rate scheduler that warms up the learning rate for one epoch and then flatlines. I ran the model for a couple of epochs (red) and then saved a checkpoint. I loaded the checkpoint with DeepSpeed enabled (green) and the learning rate scheduler was improperly reset. I then disabled DeepSpeed and continued training from the same checkpoint, keeping all other settings the same (blue). The learning rate scheduler resumed as expected.

Environment

• CUDA:
  - GPU:
  - available: False
  - version: 11.3
• Packages:
  - numpy: 1.21.1
  - pyTorch_debug: False
  - pyTorch_version: 1.10.1+cu113
  - pytorch-lightning: 1.5.10
  - tqdm: 4.62.2
• System:
  - OS: Linux
  - architecture:
  - 64bit
  -
  - processor: x86_64
  - python: 3.7.0
  - version: #1 SMP Wed Dec 1 05:10:33 EST 2021

Additional context

I tried testing this using version 1.6, but that crashed during the import. In any case, I can't find an issue addressing this bug, so I assume that it's still present.

cc @SeanNaren @awaelchli @rohitgr7

Start at the configure_optimizers entry point and reproduce checkpoint save/load with DeepSpeed enabled and disabled, comparing the learning rate scheduler state after resuming. Done means schedulers configured there resume at the checkpointed state when DeepSpeed is active, matching non-DeepSpeed behavior.