Fix LR scheduler behaviour with AMP

src/lightning/fabric/plugins/precision/amp.py

@@ -79,9 +79,11 @@ def optimizer_step(
             return super().optimizer_step(optimizer, **kwargs)
         if isinstance(optimizer, LBFGS):
             raise TypeError("AMP and the LBFGS optimizer are not compatible.")
+        previous_scale = self.scaler.get_scale()
         # note: the scaler will skip the `optimizer.step` if nonfinite gradients are found
         step_output = self.scaler.step(optimizer, **kwargs)
         self.scaler.update()
+        optimizer._skip_next_scheduler_step = self.scaler.get_scale() < previous_scale
         return step_output
 
     def state_dict(self) -> Dict[str, Any]:

src/lightning/pytorch/loops/training_epoch_loop.py

@@ -344,6 +344,9 @@ def _update_learning_rates(self, interval: str, update_plateau_schedulers: bool)
                         )
                         continue
 
+                if getattr(self.trainer.optimizers[config.opt_idx], "_skip_next_scheduler_step", False):
+                    continue
+
                 self.scheduler_progress.increment_ready()
 
                 # update LR

src/lightning/pytorch/plugins/precision/amp.py

@@ -80,8 +80,10 @@ def optimizer_step(  # type: ignore[override]
         # in manual optimization, the closure does not return a value
         if not model.automatic_optimization or not skipped_backward:
             # note: the scaler will skip the `optimizer.step` if nonfinite gradients are found
+            previous_scale = self.scaler.get_scale()
             step_output = self.scaler.step(optimizer, **kwargs)
             self.scaler.update()
+            optimizer._skip_next_scheduler_step = self.scaler.get_scale() < previous_scale
             return step_output
         return closure_result
 

tests/tests_fabric/plugins/precision/test_amp.py

@@ -69,6 +69,7 @@ def test_amp_precision_backward():
 def test_amp_precision_optimizer_step_with_scaler():
     precision = MixedPrecision(precision="16-mixed", device="cuda")
     precision.scaler = Mock()
+    precision.scaler.get_scale = Mock(return_value=1.0)
     optimizer = Mock()
 
     precision.optimizer_step(optimizer, keyword="arg")

tests/tests_pytorch/models/test_hooks.py

@@ -300,11 +300,7 @@ def _auto_train_batch(
                         name="optimizer_step",
                         args=(current_epoch, i, ANY, ANY),
                     ),
-                    *(
-                        [dict(name="lr_scheduler_step", args=(ANY, None))]
-                        if i == (trainer.num_training_batches - 1)
-                        else []
-                    ),
+                    *([dict(name="lr_scheduler_step", args=ANY)] if i == (trainer.num_training_batches - 1) else []),
                     dict(name="Callback.on_train_batch_end", args=(trainer, model, dict(loss=ANY), ANY, i)),
                     dict(name="on_train_batch_end", args=(dict(loss=ANY), ANY, i)),
                 ]