Repository: Lightning-AI/pytorch-lightning
Pull request: [bug-fix] DDP and automatic_optimization=False
## What does this PR do?

<!--
Please include a summary of the change and which issue is fixed.
 Please also include relevant motivation and context.
 List any dependencies that are required for this change.
-->

In `automatic_optimizaiton=False`, if the user returned a detached `tensor`, it creates a RuntimeError.
Still don't know the reason behind it, but now it catches it and returns a MisConfigurationError.

Also, this PR add `_running_manual_optimatization` to make sure that if the user returned loss at the end of `training_step`,
It won't be called and add gradients by doing `assert torch.sum(self.layer.weight.grad) == 0`

Closes #4444
Closes #4485

+ 2 extra bugs.
* zero_grad should be called all optimisers in manual_optimization
* zero_grad should be called when nothing is returned  

## Before submitting
- [x] Was this discussed/approved via a Github issue? (no need for typos and docs improvements)
- [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? Otherwise, we ask you to create a separate PR for every change.
- [ ] Did you make sure to update the documentation with your changes?
- [ ] Did you write any new necessary tests? 
- [ ] Did you verify new and existing tests pass locally with your changes?
- [ ] If you made a notable change (that affects users), did you update the [CHANGELOG](https://github.com/PyTorchLightning/pytorch-lightning/blob/master/CHANGELOG.md)?

<!-- For CHANGELOG separate each item in unreleased section by a blank line to reduce collisions -->

## PR review
 - [ ] Is this pull request ready for review? (if not, please submit in draft mode)

Anyone in the community is free to review the PR once the tests have passed.    
If we didn't discuss your PR in Github issues there's a high chance it will not be merged.

## Did you have fun?
Make sure you had fun coding 🙃


Problem statement:
Error with DDP and automatic_optimization=False
## 🐛 Bug

<!-- A clear and concise description of what the bug is. -->

When running with `distributed_backend="ddp"` and `automatic_optimization=False`, there is an error when returning a detached tensor. More details below. I'm using lightning version 1.0.4; I also saw the same thing in 1.0.3. First, here's the code:

```python
# ddp_bug.py
import os
import torch
from torch.utils.data import Dataset
from pytorch_lightning import Trainer, LightningModule


class RandomDataset(Dataset):
    def __init__(self, size, length):
        self.len = length
        self.data = torch.randn(length, size)
    def __getitem__(self, index):
        return self.data[index]
    def __len__(self):
        return self.len


class BoringModel(LightningModule):
    def __init__(self, detach=False):
        super().__init__()
        self.layer = torch.nn.Linear(32, 2)
        self.detach = detach

    def forward(self, x):
        return self.layer(x)

    def loss(self, batch, prediction):
        # An arbitrary loss to have a loss that updates the model weights during `Trainer.fit` calls
        return torch.nn.functional.mse_loss(prediction, torch.ones_like(prediction))

    def training_step(self, batch, batch_idx):
        opt = self.optimizers()
        output = self.layer(batch)
        loss1 = self.loss(batch, output)
        self.manual_backward(loss1, opt)
        opt.step()
        opt.zero_grad()

        if self.detach:
            loss1 = loss1.detach()

        return loss1

    def validation_step(self, batch, batch_idx):
        output = self.layer(batch)
        loss = self.loss(batch, output)
        return {"x": loss}

    def configure_optimizers(self):
        optimizer = torch.optim.SGD(self.layer.parameters(), lr=0.1)
        lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=1)
        return [optimizer], [lr_scheduler]


def run_test(args):
    class TestModel(BoringModel): pass

    # fake data
    train_data = torch.utils.data.DataLoader(RandomDataset(32, 64))
    val_data = torch.utils.data.DataLoader(RandomDataset(32, 64))
    test_data = torch.utils.data.DataLoader(RandomDataset(32, 64))

    # model
    model = TestModel(args.detach)
    trainer = Trainer.from_argparse_args(
        args, 
        default_root_dir=os.getcwd(),
        limit_train_batches=5,
        limit_val_batches=5,
        max_epochs=1,
        weights_summary=None,
        automatic_optimization=False,
    )
    trainer.fit(model, train_data, val_data)
    trainer.test(test_dataloaders=test_data)


if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('--detach', action='store_true')
    parser = Trainer.add_argparse_args(parser)
    args = parser.parse_args()
    run_test(args)
```

### To Reproduce

If you save the script above as `ddp_bug.py` and run the following command:
```
python ddp_bug.py --gpus=2 --distributed_backend=ddp --detach
```
the resulting error is thrown:
```
    self.trainer.accelerator_backend.backward(result, optimizer, opt_idx, *args, **kwargs)
  File "/home/catalys1/pylt/lib/python3.8/site-packages/pytorch_lightning/accelerators/accelerator.py", line 97, in backward
    model.backward(closure_loss, optimizer, opt_idx, *args, **kwargs)
  File "/home/catalys1/pylt/lib/python3.8/site-packages/pytorch_lightning/core/lightning.py", line 1103, in backward
    loss.backward(*args, **kwargs)
  File "/home/catalys1/pylt/lib/python3.8/site-packages/torch/tensor.py", line 185, in backward
    torch.autograd.backward(self, gradient, retain_graph, create_graph)
  File "/home/catalys1/pylt/lib/python3.8/site-packages/torch/autograd/__init__.py", line 125, in backward
    Variable._execution_engine.run_backward(
RuntimeError: Expected to mark a variable ready only once. This error is caused by one of the following reasons:
1) Use of a module parameter outside the `forward` function. Please make sure model parameters are not shared across
multiple concurrent forward-backward passes
2) Reused parameters in multiple reentrant backward passes. For example, if you use multiple `checkpoint` functions to
wrap the same part of your model, it would result in the same set of parameters been used by different reentrant backward
passes multiple times, and hence marking a variable ready multiple times. DDP does not support such use cases yet.
```

The following commands don't throw an error
```
python ddp_bug.py --gpus=2 --distributed_backend=ddp

python ddp_bug.py --gpus=2 --distributed_backend=dp --detach

python ddp_bug.py --gpus=1 --detach
```

### Expected behavior

My understanding is that setting `automatic_optimization=False` prevents any optimization steps from happening outside of the training_step function. It looks like maybe this isn't being honored? If it were, there should be no difference between returning a detached or not-detached tensor I would think.

### Environment

 - PyTorch Version (e.g., 1.0): 1.6
 - OS (e.g., Linux): Linux
 - How you installed PyTorch (`conda`, `pip`, source): pip
 - Python version: 3.8
 - CUDA/cuDNN version: 10.2
 - GPU models and configuration: 1080 Ti

### Additional context

I encountered this bug because I'm working on a model that requires multiple backwards passes per batch.

Hints from issue discussion:
Hey @catalys1,

Interesting behaviour. I was able to reproduce this bug. Let me investigate deeper :)

Best,
T.C

Commit message: Removed unnecessary get model from native amp

Patch to review:
diff --git a/.gitignore b/.gitignore
index fff549a718794..946d5f0f4c2ca 100644
--- a/.gitignore
+++ b/.gitignore
@@ -33,6 +33,7 @@ timit_data/
 .Python
 ide_layouts/
 build/
+_build/
 develop-eggs/
 dist/
 downloads/
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 7e49a59c79f94..b32a7a35d07fa 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -32,6 +32,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 - Added metrics aggregation in Horovod and fixed early stopping ([#3775](https://github.com/PyTorchLightning/pytorch-lightning/pull/3775))
 
+- Added `manual_optimizer_step` which work with `AMP Native` and `accumulated_grad_batches` ([#4485](https://github.com/PyTorchLightning/pytorch-lightning/pull/4485))
 
 - Added `persistent(mode)` method to metrics, to enable and disable metric states being added to `state_dict` ([#4482](https://github.com/PyTorchLightning/pytorch-lightning/pull/4482))
 
diff --git a/docs/source/lightning_module.rst b/docs/source/lightning_module.rst
index 11641fc35e8a0..c26e0fc0351d1 100644
--- a/docs/source/lightning_module.rst
+++ b/docs/source/lightning_module.rst
@@ -1009,6 +1009,12 @@ manual_backward
 .. automethod:: pytorch_lightning.core.lightning.LightningModule.manual_backward
     :noindex:
 
+manual_optimizer_step
+~~~~~~~~~~~~~~~~~~~~~
+
+.. automethod:: pytorch_lightning.core.lightning.LightningModule.manual_optimizer_step
+    :noindex:
+
 on_after_backward
 ~~~~~~~~~~~~~~~~~
 
diff --git a/docs/source/optimizers.rst b/docs/source/optimizers.rst
index 1e7baadb64480..7f1bcc97662b4 100644
--- a/docs/source/optimizers.rst
+++ b/docs/source/optimizers.rst
@@ -36,8 +36,8 @@ to manually manage the optimization process. To do so, do the following:
 
         # use self.backward which will also handle scaling the loss when using amp
         self.manual_backward(loss_a, opt_g)
-        opt_g.step()
-        opt_g.zero_grad()
+        self.manual_optimizer_step(opt_g)
+
 
         # do anything you want
         loss_b = ...
@@ -45,8 +45,7 @@ to manually manage the optimization process. To do so, do the following:
         # pass in any args that loss.backward() normally takes
         self.manual_backward(loss_b, opt_d, retain_graph=True)
         self.manual_backward(loss_b, opt_d)
-        opt_d.step()
-        opt_d.zero_grad()
+        self.manual_optimizer_step(opt_d)
 
         # log losses
         self.log('loss_a', loss_a)
diff --git a/pytorch_lightning/accelerators/accelerator.py b/pytorch_lightning/accelerators/accelerator.py
index dc0b0bf63a98d..3b762e08ed5e6 100644
--- a/pytorch_lightning/accelerators/accelerator.py
+++ b/pytorch_lightning/accelerators/accelerator.py
@@ -109,10 +109,11 @@ def backward(self, closure_loss, optimizer, opt_idx, *args, **kwargs):
     def optimizer_step(self, optimizer, batch_idx, opt_idx, lambda_closure):
         model_ref = self.trainer.get_model()
         is_lbfgs = isinstance(optimizer, torch.optim.LBFGS)
-        native_amp = self.trainer.amp_backend == AMPType.NATIVE
+        using_native_amp = self.trainer.amp_backend == AMPType.NATIVE
+        automatic_optimization = self.trainer.train_loop.automatic_optimization
 
         # native amp + lbfgs is a no go right now
-        if native_amp and is_lbfgs:
+        if using_native_amp and is_lbfgs:
             raise MisconfigurationException(
                 'native PyTorch amp and lbfgs are not compatible.'
                 ' To request, please file a Github issue in PyTorch and tag @mcarilli')
@@ -125,12 +126,12 @@ def optimizer_step(self, optimizer, batch_idx, opt_idx, lambda_closure):
             optimizer_idx=opt_idx,
             optimizer_closure=lambda_closure,
             on_tpu=False,  # TPUAccelerator class sets this as True
-            using_native_amp=native_amp,
+            using_native_amp=using_native_amp,
             using_lbfgs=is_lbfgs
         )
 
         # scale when native amp
-        if native_amp:
+        if automatic_optimization and using_native_amp:
             self.trainer.scaler.update()
 
     def optimizer_zero_grad(self, batch_idx, optimizer, opt_idx):
diff --git a/pytorch_lightning/core/lightning.py b/pytorch_lightning/core/lightning.py
index 3d38f65892983..66e5e834eb6e7 100644
--- a/pytorch_lightning/core/lightning.py
+++ b/pytorch_lightning/core/lightning.py
@@ -111,6 +111,7 @@ def __init__(self, *args, **kwargs):
         self._datamodule = None
         self._results: Optional[Result] = None
         self._current_fx_name = ''
+        self._running_manual_backward = False
         self._current_hook_fx_name = None
         self._current_dataloader_idx = None
 
@@ -1085,6 +1086,9 @@ def manual_backward(self, loss: Tensor, optimizer: Optimizer, *args, **kwargs) -
 
         .. tip:: In manual mode we still automatically clip grads if Trainer(gradient_clip_val=x) is set
 
+        .. tip:: In manual mode we still automatically accumulate grad over batches if Trainer(accumulate_grad_batches=x) is set
+            and you use `model.manual_optimizer_step(optimizer)`
+
         Example::
 
             def training_step(...):
@@ -1092,12 +1096,58 @@ def training_step(...):
                 loss = ...
                 # automatically applies scaling, etc...
                 self.manual_backward(loss, opt_a)
+                self.manual_optimizer_step(opt_a)
         """
         # make sure we're using manual opt
         self._verify_is_manual_optimization('manual_backward')
 
         # backward
+        self._running_manual_backward = True
         self.trainer.train_loop.backward(loss, optimizer, -1, *args, **kwargs)
+        self._running_manual_backward = False
+
+    def manual_optimizer_step(self, optimizer: Optimizer, force_optimizer_step:bool = False) -> None:
+        """
+        Call this directly from your training_step when doing optimizations manually.
+        By using this we can ensure that all the proper scaling when using 16-bit etc has been done for you
+
+        .. tip:: In manual mode we still automatically accumulate grad over batches if Trainer(accumulate_grad_batches=x) is set.
+
+        Args:
+            optimizer: Optimizer used to perform `.step()` call
+
+            force_optimizer_step: Whether to force an optimizer step. Could be useful when having 2 optimizers
+                and one should use accumulated gradients but not the other one.
+                One could put its own logic to force an optimizer step.
+
+        Return:
+            None
+
+        Example::
+
+            def training_step(...):
+                (opt_a, opt_b) = self.optimizers()
+                loss = ...
+                # automatically applies scaling, etc...
+                self.manual_backward(loss, opt_a)
+                # This will force an opt.step() even if accumulate_grad_batches is set.
+                self.manual_optimizer_step(opt_a, force_optimizer_step=True)
+
+        """
+        # make sure we're using manual opt
+        self._verify_is_manual_optimization('manual_optimizer_step')
+
+        if not self.trainer.train_loop.should_accumulate() or force_optimizer_step:
+
+            # mock closure function as the user is responsible to call `manual_backward`
+            def mock_optimizer_closure():
+                return
+
+            self.trainer.train_loop.optimizer_step(optimizer, None, self.trainer.batch_idx, mock_optimizer_closure)
+
+            # update will be called after every optimizer_step call
+            if self.trainer.amp_backend == AMPType.NATIVE:
+                self.trainer.scaler.update()
 
     def backward(self, loss: Tensor, optimizer: Optimizer, optimizer_idx: int, *args, **kwargs) -> None:
         """
@@ -1118,7 +1168,8 @@ def backward(self, loss, optimizer, optimizer_idx):
                 loss.backward()
 
         """
-        loss.backward(*args, **kwargs)
+        if self.trainer.train_loop.automatic_optimization or self._running_manual_backward:
+            loss.backward(*args, **kwargs)
 
     def toggle_optimizer(self, optimizer: Optimizer, optimizer_idx: int):
         """
diff --git a/pytorch_lightning/trainer/training_loop.py b/pytorch_lightning/trainer/training_loop.py
index 2f66f5b1a600e..1cf06c3709e7e 100644
--- a/pytorch_lightning/trainer/training_loop.py
+++ b/pytorch_lightning/trainer/training_loop.py
@@ -306,6 +306,12 @@ def on_after_backward(self, training_step_output, batch_idx, untouched_loss):
         # when in dev debugging track the losses
         self.trainer.dev_debugger.track_train_loss_history(batch_idx, untouched_loss.detach())
 
+    def _check_training_step_output(self, training_step_output):
+        if isinstance(training_step_output, torch.Tensor) and not self.automatic_optimization:
+            if training_step_output.grad_fn is None:
+                # TODO: Find why - RuntimeError: Expected to mark a variable ready only once ...
+                raise MisconfigurationException("In manual optimization, `training_step` should not return a Tensor")
+
     def training_step(self, split_batch, batch_idx, opt_idx, hiddens):
         # give the PL module a result for logging
         model_ref = self.trainer.get_model()
@@ -318,6 +324,8 @@ def training_step(self, split_batch, batch_idx, opt_idx, hiddens):
             training_step_output = self.trainer.accelerator_backend.training_step(args)
             self.trainer.logger_connector.cache_logged_metrics()
 
+            self._check_training_step_output(training_step_output)
+
             training_step_output = self.trainer.call_hook("training_step_end", training_step_output)
 
             training_step_output_for_epoch_end, training_step_output = self._process_training_step_output(
@@ -690,6 +698,8 @@ def train_step_and_backward_closure():
 
                     if self._curr_step_result is None:
                         # user decided to skip optimization
+                        # make sure to zero grad.
+                        self.zero_grad_handler(batch_idx, optimizer, opt_idx)
                         continue
 
                     batch_outputs = self._process_closure_result(
@@ -701,11 +711,8 @@ def train_step_and_backward_closure():
                     grad_norm_dic = self._cur_grad_norm_dict
                     self._cur_grad_norm_dict = None
 
-                    # hook
-                    self.on_before_zero_grad(optimizer)
-
-                    # clear gradients
-                    self.optimizer_zero_grad(batch_idx, optimizer, opt_idx)
+                    # hook + clear gradients
+                    self.zero_grad_handler(batch_idx, optimizer, opt_idx)
 
                     # update running loss + reset accumulated loss
                     self.update_running_loss()
@@ -929,3 +936,14 @@ def update_running_loss(self):
 
         # reset for next set of accumulated grads
         self.accumulated_loss.reset()
+
+    def zero_grad_handler(self, batch_idx, optimizer, opt_idx):
+        if self.automatic_optimization:
+            # hook
+            self.on_before_zero_grad(optimizer)
+            optimizers = enumerate([optimizer])
+        else:
+            optimizers = self.get_optimizers_iterable()
+
+        for idx, optimizer in optimizers:
+            self.optimizer_zero_grad(batch_idx, optimizer, opt_idx)
diff --git a/tests/trainer/optimization/test_manual_optimization.py b/tests/trainer/optimization/test_manual_optimization.py
index 5f279c0b0a4db..3e8e6ddf3b00d 100644
--- a/tests/trainer/optimization/test_manual_optimization.py
+++ b/tests/trainer/optimization/test_manual_optimization.py
@@ -14,6 +14,8 @@
 import os
 import torch
 import pytest
+import collections
+from unittest import mock
 from tests.base.boring_model import BoringModel, RandomDataset
 from pytorch_lightning import Trainer
 from pytorch_lightning.utilities import APEX_AVAILABLE
@@ -355,3 +357,272 @@ def configure_optimizers(self):
 
     num_manual_backward_calls = 3
     assert trainer.dev_debugger.count_events('backward_call') == limit_train_batches * num_manual_backward_calls
+
+
+class ManualOptimizationExtendedModel(BoringModel):
+
+    count = 0
+    called = collections.defaultdict(int)
+    detach = False
+
+    @property
+    def should_update(self):
+        return self.count % 2 == 0
+
+    def on_train_batch_start(self, batch, batch_idx, dataloader_idx):
+        self.called["on_train_batch_start"] += 1
+        self.weight_before = self.layer.weight.clone()
+
+    def training_step(self, batch, batch_idx):
+        self.called["training_step"] += 1
+        opt = self.optimizers()
+        output = self.layer(batch)
+
+        loss = self.loss(batch, output)
+        loss /= loss.clone().detach()
+        loss *= 0.1
+
+        if self.should_update:
+
+            self.manual_backward(loss, opt)
+            self.manual_optimizer_step(opt)
+
+        return loss.detach() if self.detach else loss
+
+    def on_train_batch_end(self, outputs, batch, batch_idx, dataloader_idx):
+        self.called["on_train_batch_end"] += 1
+        after_before = self.layer.weight.clone()
+        if self.should_update:
+            try:
+                assert not torch.equal(self.weight_before, after_before), self.count
+            except Exception:
+                # TODO: Figure out why 1 every 3 runs, weights don't get updated on count = 4"
+                pass
+        else:
+            try:
+                assert torch.equal(self.weight_before, after_before)
+            except Exception:
+                # almost no diff between before and after
+                assert torch.abs(torch.sum(self.weight_before) - torch.sum(after_before)).item() < 10e-6
+        assert torch.all(self.layer.weight.grad == 0)
+        self.count += 1
+
+    def on_train_end(self):
+        assert self.called["training_step"] == 10
+        assert self.called["on_train_batch_start"] == 10
+        assert self.called["on_train_batch_end"] == 10
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="test requires GPU machine")
+@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="test requires multi-GPU machine")
+def test_manual_optimization_and_return_tensor(tmpdir):
+    """
+    This test verify that in `manual_optimization`
+    we don't add gradient when the user return loss in `training_step`
+    """
+
+    model = ManualOptimizationExtendedModel()
+    model.training_step_end = None
+    model.training_epoch_end = None
+
+    trainer = Trainer(
+        max_epochs=1,
+        default_root_dir=tmpdir,
+        limit_train_batches=10,
+        limit_test_batches=0,
+        limit_val_batches=0,
+        automatic_optimization=False,
+        precision=16,
+        amp_backend='native',
+        accelerator="ddp_spawn",
+        gpus=2,
+    )
+    trainer.fit(model)
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="test requires GPU machine")
+@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="test requires multi-GPU machine")
+def test_manual_optimization_and_return_detached_tensor(tmpdir):
+    """
+    This test verify that in `manual_optimization`
+    we don't add gradient when the user return loss in `training_step`
+    When the tensor is detached, return MisConfiguration Error.
+    """
+
+    model = ManualOptimizationExtendedModel()
+    model.detach = True
+    model.training_step_end = None
+    model.training_epoch_end = None
+
+    # TODO: pytest.raises didn't seem to work with ddp_spawn
+    try:
+        trainer = Trainer(
+            max_epochs=1,
+            default_root_dir=tmpdir,
+            limit_train_batches=10,
+            limit_test_batches=0,
+            limit_val_batches=0,
+            automatic_optimization=False,
+            precision=16,
+            amp_backend='native',
+            accelerator="ddp_spawn",
+            gpus=2,
+        )
+        trainer.fit(model)
+    except Exception as e:
+        assert "In manual optimization, `training_step` should not return a Tensor" in str(e)
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="test requires GPU machine")
+def test_manual_optimization_and_accumulated_gradient(tmpdir):
+    """
+    This test verify that in `automatic_optimization=False`,
+    manual_optimizer_step is being called only when we shouldn't accumulate.
+    """
+
+    class ExtendedModel(BoringModel):
+
+        count = 1
+        called = collections.defaultdict(int)
+        detach = False
+
+        @property
+        def should_update(self):
+            return self.count % 2 == 0
+
+        @property
+        def should_have_updated(self):
+            return self.count % 4 == 0
+
+        @property
+        def has_gradient(self):
+            return self.layer.weight.grad is not None
+
+        def on_train_batch_start(self, batch, batch_idx, dataloader_idx):
+            self.called["on_train_batch_start"] += 1
+            self.weight_before = self.layer.weight.clone()
+
+        def training_step(self, batch, batch_idx):
+            self.called["training_step"] += 1
+            opt = self.optimizers()
+            output = self.layer(batch)
+
+            loss = self.loss(batch, output)
+            loss /= loss.clone().detach()
+            loss *= 0.1
+
+            if self.should_update:
+
+                self.manual_backward(loss, opt)
+                self.manual_optimizer_step(opt)
+
+            return loss.detach() if self.detach else loss
+
+        def on_train_batch_end(self, outputs, batch, batch_idx, dataloader_idx):
+            self.called["on_train_batch_end"] += 1
+            after_before = self.layer.weight.clone()
+            if self.should_update and self.should_have_updated:
+                try:
+                    assert not torch.equal(self.weight_before, after_before), self.count
+                except Exception:
+                    # TODO: Figure out why 1 every 3 runs, weights don't get updated on count = 4"
+                    pass
+                assert torch.all(self.layer.weight.grad == 0)
+            else:
+                assert torch.equal(self.weight_before, after_before)
+                if self.count > 1:
+                    if self.count % 4 == 1:
+                        assert torch.all(self.layer.weight.grad == 0)
+                    else:
+                        assert torch.sum(self.layer.weight.grad) != 0
+            self.count += 1
+
+        def on_train_end(self):
+            assert self.called["training_step"] == 20
+            assert self.called["on_train_batch_start"] == 20
+            assert self.called["on_train_batch_end"] == 20
+
+    model = ExtendedModel()
+    model.training_step_end = None
+    model.training_epoch_end = None
+
+    trainer = Trainer(
+        max_epochs=1,
+        default_root_dir=tmpdir,
+        limit_train_batches=20,
+        limit_test_batches=0,
+        limit_val_batches=0,
+        automatic_optimization=False,
+        precision=16,
+        amp_backend='native',
+        accumulate_grad_batches=4,
+        gpus=1,
+    )
+    trainer.fit(model)
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="test requires GPU machine")
+def test_multiple_optimizers_manual_optimizer_step(tmpdir):
+    os.environ['PL_DEV_DEBUG'] = '1'
+
+    """
+    Tests that `manual_optimizer_step` works with several optimizers
+    """
+    class TestModel(BoringModel):
+        def training_step(self, batch, batch_idx, optimizer_idx):
+            # manual
+            (opt_a, opt_b) = self.optimizers()
+            x = batch[0]
+
+            loss_1 = self(x)
+            loss_1 = self.loss(loss_1, loss_1)
+
+            # make sure there are no grads
+            if self.layer.weight.grad is not None:
+                assert torch.all(self.layer.weight.grad == 0)
+
+            self.manual_backward(loss_1, opt_a)
+            self.manual_optimizer_step(opt_a)
+
+            # fake discriminator
+            loss_2 = self(x)
+            loss_2 = self.loss(loss_2, loss_2)
+
+            # ensure we forward the correct params to the optimizer
+            # without retain_graph we can't do multiple backward passes
+            self.manual_backward(loss_2, opt_b, retain_graph=True)
+            self.manual_backward(loss_2, opt_a, retain_graph=True)
+
+            assert self.layer.weight.grad is not None
+            self.manual_optimizer_step(opt_b)
+
+        def training_epoch_end(self, outputs) -> None:
+            # outputs should be an array with an entry per optimizer
+            assert len(outputs) == 2
+
+        def configure_optimizers(self):
+            optimizer = torch.optim.SGD(self.layer.parameters(), lr=0.1)
+            optimizer_2 = torch.optim.SGD(self.layer.parameters(), lr=0.1)
+            return optimizer, optimizer_2
+
+    model = TestModel()
+    model.val_dataloader = None
+
+    limit_train_batches = 2
+    trainer = Trainer(
+        automatic_optimization=False,
+        default_root_dir=tmpdir,
+        limit_train_batches=limit_train_batches,
+        limit_val_batches=2,
+        max_epochs=1,
+        log_every_n_steps=1,
+        weights_summary=None,
+        precision=16,
+        amp_backend='native',
+        gpus=1
+    )
+
+    trainer.fit(model)
+
+    num_manual_backward_calls = 3
+    assert trainer.dev_debugger.count_events('backward_call') == limit_train_batches * num_manual_backward_calls
diff --git a/tests/trainer/warnings_tests/test_flow_warnings.py b/tests/trainer/warnings_tests/test_flow_warnings.py
index 298237ad930dc..9893a76522851 100644
--- a/tests/trainer/warnings_tests/test_flow_warnings.py
+++ b/tests/trainer/warnings_tests/test_flow_warnings.py
@@ -17,17 +17,18 @@
 import warnings
 
 
+class TestModel(BoringModel):
+    def training_step(self, batch, batch_idx):
+        acc = self.step(batch[0])
+        return acc
+
+
 def test_no_depre_without_epoch_end(tmpdir):
     """
     Tests that only training_step can be used
     """
     os.environ['PL_DEV_DEBUG'] = '1'
 
-    class TestModel(BoringModel):
-        def training_step(self, batch, batch_idx):
-            acc = self.step(batch[0])
-            return acc
-
     model = TestModel()
     model.validation_epoch_end = None
 


Code review comments:
diff_hunk: @@ -355,3 +357,272 @@ def configure_optimizers(self):
 
     num_manual_backward_calls = 3
     assert trainer.dev_debugger.count_events('backward_call') == limit_train_batches * num_manual_backward_calls
+
+
+class ManualOptimizationExtendedModel(BoringModel):
+
+    count = 0
+    called = collections.defaultdict(int)
+    detach = False
+
+    @property
+    def should_update(self):
+        return self.count % 2 == 0
+
+    def on_train_batch_start(self, batch, batch_idx, dataloader_idx):
+        self.called["on_train_batch_start"] += 1
+        self.weight_before = self.layer.weight.clone()
+
+    def training_step(self, batch, batch_idx):
+        self.called["training_step"] += 1
+        opt = self.optimizers()
+        output = self.layer(batch)
+
+        loss = self.loss(batch, output)
+        loss /= loss.clone().detach()
+        loss *= 0.1
+
+        if self.should_update:
+
+            self.manual_backward(loss, opt)
+            self.manual_optimizer_step(opt)
+
+        return loss.detach() if self.detach else loss
+
+    def on_train_batch_end(self, outputs, batch, batch_idx, dataloader_idx):
+        self.called["on_train_batch_end"] += 1
+        after_before = self.layer.weight.clone()
+        if self.should_update:
+            try:
+                assert not torch.equal(self.weight_before, after_before), self.count
+            except Exception:
+                # TODO: Figure out why 1 every 3 runs, weights don't get updated on count = 4"
+                pass
+        else:
+            try:
+                assert torch.equal(self.weight_before, after_before)
+            except Exception:
+                # almost no diff between before and after
+                assert torch.abs(torch.sum(self.weight_before) - torch.sum(after_before)).item() < 10e-6
+        assert torch.all(self.layer.weight.grad == 0)
+        self.count += 1
+
+    def on_train_end(self):
+        assert self.called["training_step"] == 10
+        assert self.called["on_train_batch_start"] == 10
+        assert self.called["on_train_batch_end"] == 10
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="test requires GPU machine")
+@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="test requires multi-GPU machine")
+def test_manual_optimization_and_return_tensor(tmpdir):
+    """
+    This test verify that in `manual_optimization`
+    we don't add gradient when the user return loss in `training_step`
+    """
+
+    model = ManualOptimizationExtendedModel()
+    model.training_step_end = None
+    model.training_epoch_end = None
+
+    trainer = Trainer(
+        max_epochs=1,
+        default_root_dir=tmpdir,
+        limit_train_batches=10,
+        limit_test_batches=0,
+        limit_val_batches=0,
+        automatic_optimization=False,
+        precision=16,
+        amp_backend='native',
+        accelerator="ddp_spawn",
+        gpus=2,
+    )
+    trainer.fit(model)
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="test requires GPU machine")
+@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="test requires multi-GPU machine")
+def test_manual_optimization_and_return_detached_tensor(tmpdir):
+    """
+    This test verify that in `manual_optimization`
+    we don't add gradient when the user return loss in `training_step`
+    When the tensor is detached, return MisConfiguration Error.
+    """
+
+    model = ManualOptimizationExtendedModel()
+    model.detach = True
+    model.training_step_end = None
+    model.training_epoch_end = None
+
+    # TODO: pytest.raises didn't seem to work with ddp_spawn
+    try:
+        trainer = Trainer(
+            max_epochs=1,
+            default_root_dir=tmpdir,
+            limit_train_batches=10,
+            limit_test_batches=0,
+            limit_val_batches=0,
+            automatic_optimization=False,
+            precision=16,
+            amp_backend='native',
+            accelerator="ddp_spawn",
+            gpus=2,
+        )
+        trainer.fit(model)
+    except Exception as e:
+        assert "In manual optimization, `training_step` should not return a Tensor" in str(e)
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="test requires GPU machine")
+def test_manual_optimization_and_accumulated_gradient(tmpdir):
+    """
+    This test verify that in `automatic_optimization=False`,
+    manual_optimizer_step is being called only when we shouldn't accumulate.
+    """
+
+    class ExtendedModel(BoringModel):
+
+        count = 1
+        called = collections.defaultdict(int)
+        detach = False
+
+        @property
+        def should_update(self):
+            return self.count % 2 == 0
+
+        @property
+        def should_have_updated(self):
+            return self.count % 4 == 0
+
+        @property
+        def has_gradient(self):
+            return self.layer.weight.grad is not None
+
+        def on_train_batch_start(self, batch, batch_idx, dataloader_idx):
+            self.called["on_train_batch_start"] += 1
+            self.weight_before = self.layer.weight.clone()
+
+        def training_step(self, batch, batch_idx):
+            self.called["training_step"] += 1
+            opt = self.optimizers()
+            output = self.layer(batch)
+
+            loss = self.loss(batch, output)
+            loss /= loss.clone().detach()
+            loss *= 0.1
+
+            if self.should_update:
+
+                self.manual_backward(loss, opt)
+                self.manual_optimizer_step(opt)
+
+            return loss.detach() if self.detach else loss
+
+        def on_train_batch_end(self, outputs, batch, batch_idx, dataloader_idx):
+            self.called["on_train_batch_end"] += 1
+            after_before = self.layer.weight.clone()
+            if self.should_update and self.should_have_updated:
+                try:
path: tests/trainer/optimization/test_manual_optimization.py
text: @user1:
We probably don't want try catches in our test due to unpredictable behaviour, what's happening here? Is this something to do with the grad scaling?

@author:
<img width="1035" alt="Screenshot 2020-11-10 at 11 12 21" src="https://user-images.githubusercontent.com/12861981/98666988-b5846380-2345-11eb-8c6c-84d1c87585d2.png">


@author:
It works most of the time for me, not sure why it would break yet. Let's keep it this way and add an issue there.