Repository: Lightning-AI/pytorch-lightning
Pull request: Fix `LightningModule.{un,}toggle_model` when only 1 optimizer is used
## What does this PR do?
Fixes #12086.

### Does your PR introduce any breaking changes? If yes, please list them.

<!-- FILL IN or None -->

## Before submitting

- [x] Was this **discussed/approved** via a GitHub issue? (not for typos and docs)
- [ ] Did you read the [contributor guideline](https://github.com/PyTorchLightning/pytorch-lightning/blob/master/.github/CONTRIBUTING.md), **Pull Request** section?
- [x] Did you make sure your **PR does only one thing**, instead of bundling different changes together?
- [x] Did you make sure to **update the documentation** with your changes? (if necessary)
- [x] Did you write any **new necessary tests**? (not for typos and docs)
- [x] Did you verify new and **existing tests pass** locally with your changes?
- [x] Did you list all the **breaking changes** introduced by this pull request?
- [x] Did you **update the [CHANGELOG](https://github.com/PyTorchLightning/pytorch-lightning/blob/master/CHANGELOG.md)**? (not for typos, docs, test updates, or internal minor changes/refactorings)

<!-- In the CHANGELOG, separate each item in the unreleased section by a blank line to reduce collisions -->

## PR review

Anyone in the community is welcome to review the PR.
Before you start reviewing make sure you have read [Review guidelines](https://github.com/PyTorchLightning/pytorch-lightning/wiki/Review-guidelines). In short, see the following bullet-list:

- [x] Is this pull request ready for review? (if not, please submit in draft mode)
- [x] Check that all items from **Before submitting** are resolved
- [x] Make sure the title is self-explanatory and the description concisely explains the PR
- [x] Add labels and milestones (and optionally projects) to the PR so it can be classified

## Did you have fun?

Make sure you had fun coding 🙃


Problem statement:
"object is not iterable" in lightning.toggle_optimizer when using single optimizer
## 🐛 Bug
Error thrown when calling optimizer.toggle_model(...) where only one optimizer is used in manual optimization mode.
This is caused by LightningModule.optimizers(...) trying to return the optimizer but not a list in this scenario.

### To Reproduce
```
import torch
import torch.nn.functional as F
from pytorch_lightning import seed_everything, Trainer
from tests.helpers.boring_model import BoringModel
class TesManualOptimizationModel(BoringModel):
    def __init__(self):
        super().__init__()
        self.automatic_optimization=False

    def training_step(self, batch, batch_idx):
        opt=self.optimizers()

        def compute_loss():
            x=batch[ 0 ]
            x=F.dropout(x, 0.1)
            predictions=self(x)
            predictions=F.dropout(predictions, 0.1)
            loss=torch.nn.functional.mse_loss(predictions, torch.ones_like(predictions))
            return loss

        with opt.toggle_model():
            loss=compute_loss()
            self.manual_backward(loss)
            opt.step()
            opt.zero_grad()

    def configure_optimizers(self):
        optimizer=torch.optim.SGD(self.layer.parameters(), lr=0.1)
        return optimizer


def train_manual_optimization(tmpdir):
    seed_everything(42)
    model=TesManualOptimizationModel()
    model.val_dataloader=None
    model.training_epoch_end=None

    limit_train_batches=8
    trainer=Trainer(
        default_root_dir=tmpdir,
        limit_train_batches=limit_train_batches,
        limit_val_batches=2,
        max_epochs=1,
        log_every_n_steps=1,
        gpus=0,
        accelerator=None,
    )
    trainer.fit(model)
if __name__ == '__main__':
    train_manual_optimization('E:/tmp/abcd')
```

### Output
  File "E:/Toolkits/github/pytorch-lightning\pytorch_lightning\trainer\trainer.py", line 717, in _call_and_handle_interrupt
    return trainer_fn(*args, **kwargs)
  File "E:/Toolkits/github/pytorch-lightning\pytorch_lightning\trainer\trainer.py", line 805, in _fit_impl
    results = self._run(model, ckpt_path=self.ckpt_path)
  File "E:/Toolkits/github/pytorch-lightning\pytorch_lightning\trainer\trainer.py", line 1226, in _run
    results = self._run_stage()
  File "E:/Toolkits/github/pytorch-lightning\pytorch_lightning\trainer\trainer.py", line 1311, in _run_stage
    return self._run_train()
  File "E:/Toolkits/github/pytorch-lightning\pytorch_lightning\trainer\trainer.py", line 1340, in _run_train
    self.fit_loop.run()
  File "E:/Toolkits/github/pytorch-lightning\pytorch_lightning\loops\base.py", line 204, in run
    self.advance(*args, **kwargs)
  File "E:/Toolkits/github/pytorch-lightning\pytorch_lightning\loops\fit_loop.py", line 274, in advance
    self._outputs = self.epoch_loop.run(self._data_fetcher)
  File "E:/Toolkits/github/pytorch-lightning\pytorch_lightning\loops\base.py", line 204, in run
    self.advance(*args, **kwargs)
  File "E:/Toolkits/github/pytorch-lightning\pytorch_lightning\loops\epoch\training_epoch_loop.py", line 199, in advance
    batch_output = self.batch_loop.run(batch, batch_idx)
  File "E:/Toolkits/github/pytorch-lightning\pytorch_lightning\loops\base.py", line 204, in run
    self.advance(*args, **kwargs)
  File "E:/Toolkits/github/pytorch-lightning\pytorch_lightning\loops\batch\training_batch_loop.py", line 90, in advance
    outputs = self.manual_loop.run(split_batch, batch_idx)
  File "E:/Toolkits/github/pytorch-lightning\pytorch_lightning\loops\base.py", line 204, in run
    self.advance(*args, **kwargs)
  File "E:/Toolkits/github/pytorch-lightning\pytorch_lightning\loops\optimization\manual_loop.py", line 115, in advance
    training_step_output = self.trainer._call_strategy_hook("training_step", *step_kwargs.values())
  File "E:/Toolkits/github/pytorch-lightning\pytorch_lightning\trainer\trainer.py", line 1728, in _call_strategy_hook
    output = fn(*args, **kwargs)
  File "E:/Toolkits/github/pytorch-lightning\pytorch_lightning\strategies\strategy.py", line 334, in training_step
    return self.model.training_step(*args, **kwargs)
  File "E:/Toolkits/github/kvp_lightning/sandbox.py", line 60, in training_step
    with opt.toggle_model():
  File "C:\Users\86136\.conda\envs\torch110\lib\contextlib.py", line 112, in __enter__
    return next(self.gen)
  File "E:/Toolkits/github/pytorch-lightning\pytorch_lightning\core\optimizer.py", line 96, in toggle_model
    lightning_module.toggle_optimizer(self, self._optimizer_idx)
  File "E:/Toolkits/github/pytorch-lightning\pytorch_lightning\core\lightning.py", line 1379, in toggle_optimizer
    for opt in self.optimizers(use_pl_optimizer=False):
TypeError: 'SGD' object is not iterable



### Expected behavior

No running exception.

### Environment
- PyTorch Lightning Version (1.5.9)

Hints from issue discussion:

Commit message: Make optimizers iterable in toggle_optimizer and untoggle_optimizer

Patch to review:
diff --git a/pytorch_lightning/core/lightning.py b/pytorch_lightning/core/lightning.py
index e9c129c09b57c..e4065ccaf69f4 100644
--- a/pytorch_lightning/core/lightning.py
+++ b/pytorch_lightning/core/lightning.py
@@ -1376,7 +1376,10 @@ def toggle_optimizer(self, optimizer: Union[Optimizer, LightningOptimizer], opti
         # Iterate over all optimizer parameters to preserve their `requires_grad` information
         # in case these are pre-defined during `configure_optimizers`
         param_requires_grad_state = {}
-        for opt in self.optimizers(use_pl_optimizer=False):
+        opts = self.optimizers(use_pl_optimizer=False)
+        if isinstance(opts, (Optimizer, LightningOptimizer)):
+            opts = [opts]
+        for opt in opts:
             for group in opt.param_groups:
                 for param in group["params"]:
                     # If a param already appear in param_requires_grad_state, continue
@@ -1400,7 +1403,10 @@ def untoggle_optimizer(self, optimizer_idx: int) -> None:
         Args:
             optimizer_idx: The index of the optimizer to untoggle.
         """
-        for opt_idx, opt in enumerate(self.optimizers(use_pl_optimizer=False)):
+        opts = self.optimizers(use_pl_optimizer=False)
+        if isinstance(opts, (Optimizer, LightningOptimizer)):
+            opts = [opts]
+        for opt_idx, opt in enumerate(opts):
             if optimizer_idx != opt_idx:
                 for group in opt.param_groups:
                     for param in group["params"]:


Code review comments:
diff_hunk: @@ -1400,7 +1403,10 @@ def untoggle_optimizer(self, optimizer_idx: int) -> None:
         Args:
             optimizer_idx: The index of the optimizer to untoggle.
         """
-        for opt_idx, opt in enumerate(self.optimizers(use_pl_optimizer=False)):
+        opts = self.optimizers(use_pl_optimizer=False)
+        if isinstance(opts, (Optimizer, LightningOptimizer)):
+            opts = [opts]
+        for opt_idx, opt in enumerate(opts):
path: pytorch_lightning/core/lightning.py
text: @user1:
```suggestion
        for opt_idx, opt in enumerate(self.trainer.optimizers):
```

@author:
Nice suggestion, thanks

diff_hunk: @@ -1376,7 +1376,10 @@ def toggle_optimizer(self, optimizer: Union[Optimizer, LightningOptimizer], opti
         # Iterate over all optimizer parameters to preserve their `requires_grad` information
         # in case these are pre-defined during `configure_optimizers`
         param_requires_grad_state = {}
-        for opt in self.optimizers(use_pl_optimizer=False):
+        opts = self.optimizers(use_pl_optimizer=False)
+        if isinstance(opts, (Optimizer, LightningOptimizer)):
+            opts = [opts]
+        for opt in opts:
path: pytorch_lightning/core/lightning.py
text: @user1:
We can simply do:

```suggestion
        for opt in self.trainer.optimizers:
```