fix deepspeed lr scheduler state dumping

x
x
fix confirmed

pytorch_lightning/strategies/deepspeed.py

@@ -741,7 +741,7 @@ def save_checkpoint(self, checkpoint: Dict, filepath: _PATH) -> None:
             )
         # Use deepspeed's internal checkpointing function to handle partitioned weights across processes
         # dump states as a checkpoint dictionary object
-        _exclude_keys = ["state_dict", "optimizer_states", "lr_schedulers"]
+        _exclude_keys = ["state_dict", "optimizer_states"]
         checkpoint = {k: v for k, v in checkpoint.items() if k not in _exclude_keys}
         self.deepspeed_engine.save_checkpoint(filepath, client_state=checkpoint)
 

tests/plugins/test_deepspeed_plugin.py

@@ -558,6 +558,8 @@ def on_load_checkpoint(self, checkpoint: Dict[str, Any]) -> None:
         if not hasattr(self, "model"):
             self.configure_sharded_model()
 
+        assert checkpoint["lr_schedulers"] is not None
+
 
 class ManualModelParallelClassificationModel(ModelParallelClassificationModel):
     @property