fix optimizer loop with frequencies (#9507)

Author: awaelchli

pytorch_lightning/loops/batch/training_batch_loop.py

@@ -123,8 +123,7 @@ def advance(self, batch, batch_idx):
 
         if self.trainer.lightning_module.automatic_optimization:
             # in automatic optimization, hand over execution to the OptimizerLoop
-            optimizers = [optimizer for _, optimizer in self.get_active_optimizers(batch_idx)]
-            batch_outputs = self.optimizer_loop.run(split_batch, optimizers, batch_idx)
+            batch_outputs = self.optimizer_loop.run(split_batch, self.get_active_optimizers(batch_idx), batch_idx)
             # combine outputs from each optimizer
             for k in range(len(batch_outputs)):
                 self.batch_outputs[k].extend(batch_outputs[k])

pytorch_lightning/loops/optimization/optimizer_loop.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 from dataclasses import dataclass, field
 from functools import partial
-from typing import Any, Callable, Dict, List, Optional
+from typing import Any, Callable, Dict, List, Optional, Tuple
 
 import torch
 from torch import Tensor
@@ -188,36 +188,42 @@ def __init__(self) -> None:
         self._skip_backward: bool = False
         self._batch_idx: int = 0
         self._optimizers: List[Optimizer] = []
+        self._indices: List[int] = []
         self._hiddens: Optional[Any] = None
 
+    @property
+    def optimizer_idx(self) -> int:
+        return self._indices[self.optim_progress.optimizer_position]
+
     @property
     def done(self) -> bool:
         """Returns ``True`` when the last optimizer in the sequence has run."""
-        return self.optim_progress.optimizer_idx >= len(self._optimizers)
+        return self.optim_progress.optimizer_position >= len(self._indices)
 
     def connect(self, **kwargs: "Loop") -> None:
         raise NotImplementedError(f"{self.__class__.__name__} does not connect any child loops.")
 
     def reset(self) -> None:
         if not self.restarting or self.done:
-            self.optim_progress.optimizer_idx = 0
+            self.optim_progress.optimizer_position = 0
         self.outputs = [[] for _ in range(len(self.trainer.optimizers))]
 
-    def on_run_start(self, batch: Any, optimizers: List[Optimizer], batch_idx: int) -> None:  # type: ignore[override]
+    def on_run_start(  # type: ignore[override]
+        self, batch: Any, optimizers: List[Tuple[int, Optimizer]], batch_idx: int
+    ) -> None:
         self._batch_idx = batch_idx
-        self._optimizers = optimizers
+        self._indices, self._optimizers = zip(*optimizers)
 
     def advance(self, batch: Any, *args: Any, **kwargs: Any) -> None:  # type: ignore[override]
         result = self._run_optimization(
             batch,
             self._batch_idx,
-            self._optimizers[self.optim_progress.optimizer_idx],
-            self.optim_progress.optimizer_idx,
+            self._optimizers[self.optim_progress.optimizer_position],
+            self.optimizer_idx,
         )
         if result.loss is not None:
-            self.outputs[self.optim_progress.optimizer_idx].append(result.drop_closure_loss())
-
-        self.optim_progress.optimizer_idx += 1
+            self.outputs[self.optimizer_idx].append(result.drop_closure_loss())
+        self.optim_progress.optimizer_position += 1
 
     def on_run_end(self) -> _OUTPUTS_TYPE:
         outputs, self.outputs = self.outputs, []  # free memory

pytorch_lightning/trainer/progress.py

@@ -209,12 +209,15 @@ class OptimizationProgress(BaseProgress):
 
     Args:
         optimizer: Tracks optimizer progress.
-        optimizer_idx: The index of the current optimizer. Used to know which optimizer we were using when restarting.
+        optimizer_position: The index of the current optimizer amongst the currently active optimizers.
+            Used to know which optimizer we were using when restarting.
+            Since not all optimizers may be active at a given time, this index is different from the ``optimizer_idx``
+            seen in the optimization loops.
     """
 
     # TODO: support for multiple optimizers
     optimizer: OptimizerProgress = field(default_factory=OptimizerProgress)
-    optimizer_idx: int = 0
+    optimizer_position: int = 0
 
     @property
     def optimizer_steps(self) -> int:
@@ -225,4 +228,4 @@ def reset_on_epoch(self) -> None:
 
     def load_state_dict(self, state_dict: dict) -> None:
         self.optimizer.load_state_dict(state_dict["optimizer"])
-        self.optimizer_idx = state_dict["optimizer_idx"]
+        self.optimizer_position = state_dict["optimizer_position"]

tests/loops/optimization/test_optimizer_loop.py

@@ -11,9 +11,16 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+from unittest.mock import Mock
+
+import pytest
 import torch
+from torch.optim import Adam, SGD
 
+from pytorch_lightning import Trainer
+from pytorch_lightning.core.optimizer import LightningOptimizer
 from pytorch_lightning.loops.optimization.optimizer_loop import ClosureResult
+from tests.helpers import BoringModel
 
 
 def test_closure_result_deepcopy():
@@ -37,3 +44,68 @@ def test_closure_result_apply_accumulation():
     closure_loss = torch.tensor(25.0)
     result = ClosureResult.from_training_step_output(closure_loss, 5)
     assert result.loss == 5
+
+
+@pytest.mark.parametrize(
+    "frequencies,expected",
+    [
+        (
+            (3, 1),
+            [
+                (0, "SGD"),
+                (0, "SGD"),
+                (0, "SGD"),
+                (1, "Adam"),
+                (0, "SGD"),
+                (0, "SGD"),
+                (0, "SGD"),
+                (1, "Adam"),
+                (0, "SGD"),
+                (0, "SGD"),
+            ],
+        ),
+        (
+            (1, 2),
+            [
+                (0, "SGD"),
+                (1, "Adam"),
+                (1, "Adam"),
+                (0, "SGD"),
+                (1, "Adam"),
+                (1, "Adam"),
+                (0, "SGD"),
+                (1, "Adam"),
+                (1, "Adam"),
+                (0, "SGD"),
+            ],
+        ),
+    ],
+)
+def test_optimizer_frequencies(tmpdir, frequencies, expected):
+    """Test that the optimizer loop runs optimization for the correct optimizer and optimizer idx when different
+    frequencies are requested."""
+
+    class CurrentModel(BoringModel):
+        def training_step(self, batch, batch_idx, optimizer_idx):
+            return super().training_step(batch, batch_idx)
+
+        def configure_optimizers(self):
+            opt0 = SGD(self.parameters(), lr=0.1)
+            opt1 = Adam(self.parameters(), lr=0.1)
+            return {"optimizer": opt0, "frequency": frequencies[0]}, {"optimizer": opt1, "frequency": frequencies[1]}
+
+    model = CurrentModel()
+    model.optimizer_step = Mock(wraps=model.optimizer_step)
+    trainer = Trainer(
+        default_root_dir=tmpdir,
+        fast_dev_run=10,
+        progress_bar_refresh_rate=0,
+    )
+    trainer.fit(model)
+
+    positional_args = [c[0] for c in model.optimizer_step.call_args_list]
+    pl_optimizer_sequence = [args[2] for args in positional_args]
+    opt_idx_sequence = [args[3] for args in positional_args]
+    assert all(isinstance(opt, LightningOptimizer) for opt in pl_optimizer_sequence)
+    optimizer_sequence = [opt._optimizer.__class__.__name__ for opt in pl_optimizer_sequence]
+    assert list(zip(opt_idx_sequence, optimizer_sequence)) == expected

tests/loops/test_loop_state_dict.py

@@ -67,7 +67,7 @@ def test_loops_state_dict_structure():
                         "current": {"ready": 0, "started": 0, "completed": 0},
                     },
                 },
-                "optimizer_idx": 0,
+                "optimizer_position": 0,
             },
             "epoch_loop.val_loop.state_dict": {},
             "epoch_loop.val_loop.dataloader_progress": {

tests/loops/test_loops.py

@@ -468,7 +468,7 @@ def configure_optimizers_multiple(self):
         "epoch_loop.batch_loop.manual_loop.state_dict": ANY,
         "epoch_loop.batch_loop.optimizer_loop.state_dict": {},
         "epoch_loop.batch_loop.optimizer_loop.optim_progress": {
-            "optimizer_idx": stop_optimizer,
+            "optimizer_position": stop_optimizer,
             "optimizer": {
                 "step": {
                     "total": {
@@ -611,7 +611,7 @@ def configure_optimizers_multiple(self):
         "epoch_loop.batch_loop.manual_loop.state_dict": ANY,
         "epoch_loop.batch_loop.optimizer_loop.state_dict": {},
         "epoch_loop.batch_loop.optimizer_loop.optim_progress": {
-            "optimizer_idx": n_optimizers,
+            "optimizer_position": n_optimizers,
             "optimizer": {
                 "step": {
                     "total": {
@@ -697,12 +697,12 @@ def mid_epoch_reset_assertions():
 
     # resetting from a mid-epoch checkpoint should not change progress counters
     mid_epoch_reset_assertions()
-    assert optimizer_loop.optim_progress.optimizer_idx == 1
+    assert optimizer_loop.optim_progress.optimizer_position == 1
     fit_loop.reset()
     epoch_loop.reset()
     optimizer_loop.reset()
     mid_epoch_reset_assertions()
-    assert optimizer_loop.optim_progress.optimizer_idx == 0
+    assert optimizer_loop.optim_progress.optimizer_position == 0
 
     # reset state loaded from a checkpoint from the end of an epoch
     end_of_epoch_ckpt = torch.load(str(tmpdir / "epoch=0-step=3.ckpt"))
@@ -726,7 +726,7 @@ def mid_epoch_reset_assertions():
     assert epoch_loop.batch_progress.current.ready == 4
     assert epoch_loop.batch_progress.current.completed == 4
 
-    assert optimizer_loop.optim_progress.optimizer_idx == 1
+    assert optimizer_loop.optim_progress.optimizer_position == 1
 
     # resetting from a end-of-epoch checkpoint should reset the current counters to 0
     fit_loop.reset()
@@ -745,4 +745,4 @@ def mid_epoch_reset_assertions():
     assert epoch_loop.batch_progress.current.ready == 0
     assert epoch_loop.batch_progress.current.completed == 0
 
-    assert optimizer_loop.optim_progress.optimizer_idx == 0
+    assert optimizer_loop.optim_progress.optimizer_position == 0