Only validate schedulers in automatic optimization (#18092)

(cherry picked from commit ea92c21)

Author: awaelchli

src/lightning/pytorch/CHANGELOG.md

@@ -16,6 +16,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - `LightningCLI` not saving correctly `seed_everything` when `run=True` and `seed_everything=True` ([#18056](https://github.com/Lightning-AI/lightning/pull/18056))
 
 
+- Fixed validation of non-PyTorch LR schedulers in manual optimization mode ([#18092](https://github.com/Lightning-AI/lightning/pull/18092))
+
+
 ## [2.0.5] - 2023-07-07
 
 ### Fixed

src/lightning/pytorch/core/optimizer.py

@@ -325,7 +325,11 @@ def _validate_scheduler_api(lr_scheduler_configs: List[LRSchedulerConfig], model
                 " It should have `state_dict` and `load_state_dict` methods defined."
             )
 
-        if not isinstance(scheduler, LRSchedulerTypeTuple) and not is_overridden("lr_scheduler_step", model):
+        if (
+            not isinstance(scheduler, LRSchedulerTypeTuple)
+            and not is_overridden("lr_scheduler_step", model)
+            and model.automatic_optimization
+        ):
             raise MisconfigurationException(
                 f"The provided lr scheduler `{scheduler.__class__.__name__}` doesn't follow PyTorch's LRScheduler"
                 " API. You should override the `LightningModule.lr_scheduler_step` hook with your own logic if"

tests/tests_pytorch/trainer/optimization/test_manual_optimization.py

@@ -22,6 +22,7 @@
 import torch.distributed as torch_distrib
 import torch.nn.functional as F
 
+from lightning.fabric.utilities.exceptions import MisconfigurationException
 from lightning.fabric.utilities.imports import _TORCH_GREATER_EQUAL_2_0
 from lightning.pytorch import seed_everything, Trainer
 from lightning.pytorch.demos.boring_classes import BoringModel, ManualOptimBoringModel
@@ -886,3 +887,37 @@ def configure_optimizers(self):
 
     assert set(trainer.logged_metrics) == {"loss_d", "loss_g"}
     assert set(trainer.progress_bar_metrics) == {"loss_d", "loss_g"}
+
+
+@pytest.mark.parametrize("automatic_optimization", [True, False])
+def test_manual_optimization_with_non_pytorch_scheduler(automatic_optimization):
+    """In manual optimization, the user can provide a custom scheduler that doesn't follow PyTorch's interface."""
+
+    class IncompatibleScheduler:
+        def __init__(self, optimizer):
+            self.optimizer = optimizer
+
+        def state_dict(self):
+            return {}
+
+        def load_state_dict(self, _):
+            pass
+
+    class Model(BoringModel):
+        def __init__(self):
+            super().__init__()
+            self.automatic_optimization = automatic_optimization
+
+        def configure_optimizers(self):
+            optimizer = torch.optim.SGD(self.layer.parameters(), lr=0.1)
+            scheduler = IncompatibleScheduler(optimizer)
+            return [optimizer], [scheduler]
+
+    model = Model()
+    trainer = Trainer(accelerator="cpu", max_epochs=0)
+    if automatic_optimization:
+        with pytest.raises(MisconfigurationException, match="doesn't follow PyTorch's LRScheduler"):
+            trainer.fit(model)
+    else:
+        # No error for manual optimization
+        trainer.fit(model)