Repository: Lightning-AI/pytorch-lightning
Pull request: [feat] Allow overriding optimizer_zero_grad and/or optimizer_step when using accumulate_grad_batches
## What does this PR do?

Adds the ability to override optimizer_zero_grad when using accumulate_grad_batches as discussed in #6910 

Fixes #6910

## Before submitting
- [x] Was this discussed/approved via a GitHub issue? (not for typos and docs)
- [x] 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 update the [CHANGELOG](https://github.com/PyTorchLightning/pytorch-lightning/blob/master/CHANGELOG.md)? (not for typos, docs, test updates, or internal minor changes/refactorings)

## PR review
Anyone in the community is free to review the PR once the tests have passed.
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:
Allow overriding optimizer_zero_grad and/or optimizer_step when using accumulate_grad_batches
## 🚀 Feature

Currently, PyTorch Lightning refuses to allow custom optimizer_step and optimizer_zero_grad functions when using accumulated gradients, saying:
```
When overriding `LightningModule` optimizer_step or optimizer_zero_grad, `accumulate_grad_batches` in `Trainer` should be 1. It ensures optimizer_step or optimizer_zero_grad are called on every batch.
```

It doesn't seem like there's any reason to avoid allowing custom gradient steps/optimizer reset behaviors with the automated accumulated_grad_batches? Perhaps I'm not as familiar with the underlying procedure, however, and there's something technical blocking this functionality. 

### Motivation

Currently, I would like to make modifications to an optimizer every time zero_grad is called. The perfect time to do that would be in an overriden zero_grad method, however that prevents the use of gradient accumulation. 

### Pitch

Remove the check at lines 80-86 in the configuration_validator.py (Or change it to a warning, instead of an error). I can submit a PR for this, if reasonable.

### Alternatives

We can solve this problem by writing custom optimizers, or by keeping track of the number of iterations of gradient accumulation in the standard steps, however this seems like it runs contrary to the point of using pytorch lightning's APIs.

Hints from issue discussion:
This issue has been automatically marked as stale because it hasn't had any recent activity. This issue will be closed in 7 days if no further activity occurs. Thank you for your contributions, Pytorch Lightning Team!

@PyTorchLightning/core-contributors 
@tchaton Are you aware of the motivation for this?

Is it a hard limitation or just to avoid potential user errors?
This issue has been automatically marked as stale because it hasn't had any recent activity. This issue will be closed in 7 days if no further activity occurs. Thank you for your contributions, Pytorch Lightning Team!

I'm still watching this issue for conversation - I can go ahead and just submit a PR for this, if it would be more likely to generate conversation on the topic @carmocca / @tchaton. 
Dove into git history and found the why:

https://github.com/PyTorchLightning/pytorch-lightning/pull/4984#discussion_r537302856

I feel like we can remove this constraint and explicitly mention it in the `optimizer_step` docs.

Feel free to open that PR 👍 
As a note, the PR should also uncomment this snippet of code:

https://github.com/PyTorchLightning/pytorch-lightning/blob/7f1d3fa0806e680d23728a92ad207f823e377de1/tests/trainer/test_trainer.py#L256-L272

Commit message: Allow for overriding optimizer_step and optimizer_zero_grad when using accumulate_grad_batches

Patch to review:
diff --git a/pytorch_lightning/core/lightning.py b/pytorch_lightning/core/lightning.py
index 02633d3df16fa..9f6039353e15a 100644
--- a/pytorch_lightning/core/lightning.py
+++ b/pytorch_lightning/core/lightning.py
@@ -1447,6 +1447,21 @@ def optimizer_step(
             ``training_step()``, ``optimizer.zero_grad()``, ``backward()`` are called within
             :meth:`~pytorch_lightning.trainer.training_loop.TrainLoop.run_training_batch`.
 
+        Warning:
+            When using the standard trainer and ``accumulate_grad_batches``, both
+            ``optimizer.step()`` and ``optimizer.zero_grad()`` are called only once per
+            *optimization cycle*, and not once per step. This means that code like
+            the following:
+
+            .. code-block:: python
+
+                Trainer(accumulate_grad_batches=2)
+                ...
+                def optimizer_step(...):
+                    batch_idx % 2 == 0: optimizer.step()
+
+            will produce a gradient accumulation of 4, and not the desired 2 as expected.
+
         Args:
             epoch: Current epoch
             batch_idx: Index of current batch
diff --git a/pytorch_lightning/trainer/configuration_validator.py b/pytorch_lightning/trainer/configuration_validator.py
index 3fc7f91348dfa..7c1b02fdd325d 100644
--- a/pytorch_lightning/trainer/configuration_validator.py
+++ b/pytorch_lightning/trainer/configuration_validator.py
@@ -83,10 +83,11 @@ def __verify_train_loop_configuration(self, model: 'pl.LightningModule') -> None
 
         has_overriden_optimization_functions = trainer.overriden_optimizer_step or trainer.overriden_optimizer_zero_grad
         if has_overriden_optimization_functions and going_to_accumulate_grad_batches and automatic_optimization:
-            raise MisconfigurationException(
+            rank_zero_warn(
                 'When overriding `LightningModule` optimizer_step or optimizer_zero_grad,'
-                ' `accumulate_grad_batches` in `Trainer` should be 1.'
-                ' It ensures optimizer_step or optimizer_zero_grad are called on every batch.'
+                ' `accumulate_grad_batches` in `Trainer` is usually 1.'
+                ' If `accumulate_grad_batches` is not 1, optimizer_step and optimizer_zero_grad'
+                ' are not called on every batch (rather, they are called on every optimization step).'
             )
 
     def __verify_eval_loop_configuration(self, model: 'pl.LightningModule', stage: str) -> None:
diff --git a/tests/trainer/test_trainer.py b/tests/trainer/test_trainer.py
index d353c0941d3f6..ae3e9d79259f8 100644
--- a/tests/trainer/test_trainer.py
+++ b/tests/trainer/test_trainer.py
@@ -253,18 +253,18 @@ def backward(self, *args, **kwargs) -> None:
 
             return out
 
-        # def optimizer_step(self, *args, **kwargs):
-        #     pre_opt_step_state_dict = self.state_dict()
-        #     assert self.check(self.start_state_dict, pre_opt_step_state_dict)
+        def optimizer_step(self, *args, **kwargs):
+            pre_opt_step_state_dict = self.state_dict()
+            assert self.check(self.start_state_dict, pre_opt_step_state_dict)
 
-        #     # this calls `backward` and `on_after_backward` inside the closure
-        #     out = super().optimizer_step(*args, **kwargs)
+            # this calls `backward` and `on_after_backward` inside the closure
+            out = super().optimizer_step(*args, **kwargs)
 
-        #     # the state dict changed
-        #     assert self.check(pre_opt_step_state_dict, self.state_dict(), equal=False)
+            # the state dict changed
+            assert self.check(pre_opt_step_state_dict, self.state_dict(), equal=False)
 
-        #     self.opt_step_called = True
-        #     return out
+            self.opt_step_called = True
+            return out
 
         def on_after_backward(self):
             # should override `optimizer_step` instead but can't with `accumulate_grad_batches`


Code review comments:
diff_hunk: @@ -253,18 +253,18 @@ def backward(self, *args, **kwargs) -> None:
 
             return out
 
-        # def optimizer_step(self, *args, **kwargs):
-        #     pre_opt_step_state_dict = self.state_dict()
-        #     assert self.check(self.start_state_dict, pre_opt_step_state_dict)
+        def optimizer_step(self, *args, **kwargs):
+            pre_opt_step_state_dict = self.state_dict()
+            assert self.check(self.start_state_dict, pre_opt_step_state_dict)
 
-        #     # this calls `backward` and `on_after_backward` inside the closure
-        #     out = super().optimizer_step(*args, **kwargs)
+            # this calls `backward` and `on_after_backward` inside the closure
+            out = super().optimizer_step(*args, **kwargs)
 
-        #     # the state dict changed
-        #     assert self.check(pre_opt_step_state_dict, self.state_dict(), equal=False)
+            # the state dict changed
+            assert self.check(pre_opt_step_state_dict, self.state_dict(), equal=False)
 
-        #     self.opt_step_called = True
-        #     return out
+            self.opt_step_called = True
+            return out
 
         def on_after_backward(self):
             # should override `optimizer_step` instead but can't with `accumulate_grad_batches`
path: tests/trainer/test_trainer.py
text: @user1:
This function can be removed now

diff_hunk: @@ -83,10 +83,11 @@ def __verify_train_loop_configuration(self, model: 'pl.LightningModule') -> None
 
         has_overriden_optimization_functions = trainer.overriden_optimizer_step or trainer.overriden_optimizer_zero_grad
         if has_overriden_optimization_functions and going_to_accumulate_grad_batches and automatic_optimization:
-            raise MisconfigurationException(
+            rank_zero_warn(
                 'When overriding `LightningModule` optimizer_step or optimizer_zero_grad,'
-                ' `accumulate_grad_batches` in `Trainer` should be 1.'
-                ' It ensures optimizer_step or optimizer_zero_grad are called on every batch.'
+                ' `accumulate_grad_batches` in `Trainer` is usually 1.'
+                ' If `accumulate_grad_batches` is not 1, optimizer_step and optimizer_zero_grad'
+                ' are not called on every batch (rather, they are called on every optimization step).'
path: pytorch_lightning/trainer/configuration_validator.py
text: @user1:
At this point, we know that `accumulate_grad_batches != 1` so we can reword the warning.

@author:
I reworded it slightly already, but I'm not exactly sure what to say here. Perhaps: "accumulate_grad_batches is not 1, but optimizer_step or optimizer_zero_grad is overridden. Note that optimizer_step and optimizer_zero_grad will not be called on every batch." 

@user1:
```python
When using `Trainer(accumulate_grad_batches != 1)` and overriding `LightningModule.optimizer_{step,zero_grad}`, the hooks will not be called on every batch (rather, they are called on every optimization step)
```