Move implementation to scales comparison

Author: milesial

src/lightning_fabric/plugins/precision/native_amp.py

@@ -73,8 +73,10 @@ def optimizer_step(
             return super().optimizer_step(optimizer, **kwargs)
         if isinstance(optimizer, LBFGS):
             raise TypeError("Native 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)
+        model._skip_next_scheduler_step = self.scaler.get_scale() < previous_scale
         self.scaler.update()
         return step_output
 

src/pytorch_lightning/core/module.py

@@ -1649,10 +1649,6 @@ def lr_scheduler_step(self, scheduler, optimizer_idx, metric):
                 scheduler.step(epoch=self.current_epoch)
 
         """
-        optimizer = self.trainer.optimizers[optimizer_idx]
-        if hasattr(optimizer, "_step_count") and optimizer._step_count <= 0:
-            return
-
         if metric is None:
             scheduler.step()  # type: ignore[call-arg]
         else:

src/pytorch_lightning/loops/epoch/training_epoch_loop.py

@@ -216,8 +216,9 @@ def advance(self, data_fetcher: AbstractDataFetcher) -> None:
 
         # update non-plateau LR schedulers
         # update epoch-interval ones only when we are at the end of training epoch
-        self.update_lr_schedulers("step", update_plateau_schedulers=False)
-        if self._num_ready_batches_reached():
+        if not getattr(self.trainer.lightning_module, "_skip_next_scheduler_step", False):
+            self.update_lr_schedulers("step", update_plateau_schedulers=False)
+        elif self._num_ready_batches_reached():
             self.update_lr_schedulers("epoch", update_plateau_schedulers=False)
 
         batch_end_outputs = self._prepare_outputs_training_batch_end(