Repository: Lightning-AI/pytorch-lightning
Pull request: optimizer clean up
## 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.

If we didn't discuss your PR in Github issues there's a high chance it will not be merged.
-->

This PR is a first attempt at cleaning out optimization by introducing LightningOptimizer.

Here is how SGD will get printed.
```
LightningSGD (
Parameter Group 0
    dampening: 0
    lr: 0.1
    momentum: 0
    nesterov: False
    weight_decay: 0
)
```

Fixes #4663

## 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
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 in short, see following bullet-list:  

 - [ ] Is this pull request ready for review? (if not, please submit in draft mode)
 - [ ] Check that all items from **Before submitting** are resolved
 - [ ] Make sure the title is self explanatory and the description concisely explains the PR
 - [ ] Add labels and milestones (and optionally projects) to the PR so it can be classified; _Bugfixes should be including in bug-fix release milestones (m.f.X) and features should be included in (m.X.b) releases._
 

## Did you have fun?
Make sure you had fun coding 🙃


Problem statement:
Wrap optimizer in LightningOptimizer object for manual optimization
Currently users are overriding optimizer_step in Lightning model, and therefore lose support for tpu, amp, etc...
By wrapping their optimizers, we can make sure that when they call optimizer.step, everything is done properly.

It allows us to remove the need for manual_backward / manual_optimizer_step making API cleaner and reduce possible bugs.

new API for manual_optimization, but it also works for automatic optimization.

        Example::

            def training_step(...):
                (opt_a, opt_b) = self.optimizers()
                loss_a = ...

                # automatically applies scaling, toggle model, accumulated_gradients, tpu, zero_grad, etc...
                opt_a.backward(loss_a)
                opt_a.step()

        Example::

            def training_step(...):
                (opt_a, opt_b) = self.optimizers()
                loss_a = ...

                # automatically applies scaling, etc...
                def closure_a():
                    loss_a = ...
                    opt_a.backward(loss)

                opt_a.step(closure=closure_a)

        Example::

           For a GAN

            def training_step(self, batch, batch_idx):

                # emulate gans training
                opt_gen, opt_dis = self.optimizers()

                # Note: Be careful, don't log on the same key in self.log in both closure
                # as they will be aggregated together on epoch_end

                def gen_closure():
                    ... forward and compute loss for generator
                    loss_gen = ...
                    self.log("loss_gen", loss_gen, on_step=True, on_epoch=True)
                    opt_gen.backward(loss_gen)

                def dis_closure():
                    ... forward and compute loss for discriminator
                    loss_dis = ...
                    self.log("loss_dis", loss_dis, on_step=True, on_epoch=True)
                    opt_dis.backward(loss_dis)

                # this will accumulate gradients for 2 batches and then call opt_gen.step()
                opt_gen.step(closure=gen_closure, make_optimizer_step=batch_idx % 2 == 0)

                # update discriminator every 4 batches
                # therefore, no gradient accumulation for discriminator
                if batch_idx % 4 == 0 :
                    # Note: Set make_optimizer_step to True or it will use by default
                    # Trainer(accumulate_grad_batches=x)
                    opt_dis.step(closure=optimizer_closure, make_optimizer_step=True)
here is how you would do it with manual optimization


            def training_step(...):
                (opt_a, opt_b) = self.optimizers()
                loss = ...
                # automatically applies scaling, etc...

                self.manual_backward(loss, opt_a)

                # This will use accumulate gradients for `accumulate_grad_batches` batches
                # and then run opt_a.step()
                self.manual_optimizer_step(opt_a)
Here is how people would do it before by overriding optimizer_step

                def optimizer_step(self, epoch, batch_idx, optimizer, optimizer_idx,
                                   optimizer_closure, on_tpu, using_native_amp, using_lbfgs):
                    # update generator opt every 2 steps
                    if optimizer_idx == 0:
                        if batch_idx % 2 == 0 :
                            optimizer.step(closure=optimizer_closure)
                            optimizer.zero_grad()

                    # update discriminator opt every 4 steps
                    if optimizer_idx == 1:
                        if batch_idx % 4 == 0 :
                            optimizer.step(closure=optimizer_closure)
                            optimizer.zero_grad()
There is a flag enable_pl_optimizer to switch it off.

It will allow us to move AMP logic from accelerators to LightningOptimizer, cleaning out the internal code.

Finally, we are working on some DDPPlugin and we need the step function to be wrapped.

It is still an idea and we didn't decided if we will integrate it. I am really keen to hear your feedbacks @carmocca @awaelchli @justusschock @williamFalcon :)

Hints from issue discussion:
@awaelchli

Commit message: Update pytorch_lightning/core/optimizer.py

Co-authored-by: Carlos Mocholí <carlossmocholi@gmail.com>

Patch to review:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index c87c3811a8003..4d2c8f29b918a 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -44,6 +44,8 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 - Added printing of total num of params, trainable and non-trainable params in ModelSummary ([#4521](https://github.com/PyTorchLightning/pytorch-lightning/pull/4521))
 
+- Added `LightningOptimizer` ([#4658](https://github.com/PyTorchLightning/pytorch-lightning/pull/4658))
+
 
 ### Changed
 
diff --git a/docs/source/lightning_module.rst b/docs/source/lightning_module.rst
index 4e88aaa3471d0..b6010409d64a0 100644
--- a/docs/source/lightning_module.rst
+++ b/docs/source/lightning_module.rst
@@ -1009,11 +1009,6 @@ 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 7f1bcc97662b4..6ca72b8069d6d 100644
--- a/docs/source/optimizers.rst
+++ b/docs/source/optimizers.rst
@@ -36,7 +36,7 @@ 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)
-        self.manual_optimizer_step(opt_g)
+        opt_g.step()
 
 
         # do anything you want
@@ -45,7 +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)
-        self.manual_optimizer_step(opt_d)
+        opt_d.step()
 
         # log losses
         self.log('loss_a', loss_a)
diff --git a/pytorch_lightning/accelerators/accelerator_connector.py b/pytorch_lightning/accelerators/accelerator_connector.py
index c40600785a558..ed0379e46dc5b 100644
--- a/pytorch_lightning/accelerators/accelerator_connector.py
+++ b/pytorch_lightning/accelerators/accelerator_connector.py
@@ -11,18 +11,18 @@
 # 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 pytorch_lightning import accelerators
 import os
+
 import torch
 
-from pytorch_lightning.utilities import device_parser, XLA_AVAILABLE
-from pytorch_lightning.utilities import rank_zero_only
-from pytorch_lightning.utilities.distributed import rank_zero_warn, rank_zero_info
-from pytorch_lightning.utilities.exceptions import MisconfigurationException
 from pytorch_lightning import _logger as log
+from pytorch_lightning import accelerators
+from pytorch_lightning.accelerators.accelerator import Accelerator
 from pytorch_lightning.cluster_environments.slurm_environment import SLURMEnvironment
 from pytorch_lightning.cluster_environments.torchelastic_environment import TorchElasticEnvironment
-from pytorch_lightning.accelerators.accelerator import Accelerator
+from pytorch_lightning.utilities import XLA_AVAILABLE, device_parser, rank_zero_only
+from pytorch_lightning.utilities.distributed import rank_zero_info, rank_zero_warn
+from pytorch_lightning.utilities.exceptions import MisconfigurationException
 
 try:
     import horovod.torch as hvd
@@ -397,8 +397,7 @@ def set_nvidia_flags(self, is_slurm_managing_tasks, data_parallel_device_ids):
     def determine_local_rank(self):
         if self.trainer.is_slurm_managing_tasks:
             return int(os.environ['SLURM_LOCALID'])
-        else:
-            return int(os.environ.get('LOCAL_RANK', 0))
+        return int(os.environ.get('LOCAL_RANK', 0))
 
     def determine_ddp_node_rank(self):
         if self.trainer.is_slurm_managing_tasks:
diff --git a/pytorch_lightning/accelerators/cpu_accelerator.py b/pytorch_lightning/accelerators/cpu_accelerator.py
index 66f9e4f0201b2..fe0ab59fb554f 100644
--- a/pytorch_lightning/accelerators/cpu_accelerator.py
+++ b/pytorch_lightning/accelerators/cpu_accelerator.py
@@ -11,7 +11,7 @@
 # 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 typing import Optional, Union, Any
+from typing import Any, Optional, Union
 
 import torch
 
@@ -47,6 +47,8 @@ def setup(self, model):
         # allow for lr schedulers as well
         self.setup_optimizers(model)
 
+        self.trainer.convert_to_lightning_optimizers()
+
         self.trainer.model = model
 
     def train(self):
diff --git a/pytorch_lightning/accelerators/ddp2_accelerator.py b/pytorch_lightning/accelerators/ddp2_accelerator.py
index 8f7f58c74684d..f43866881cabb 100644
--- a/pytorch_lightning/accelerators/ddp2_accelerator.py
+++ b/pytorch_lightning/accelerators/ddp2_accelerator.py
@@ -12,23 +12,23 @@
 # See the License for the specific language governing permissions and
 # limitations under the License
 import os
+from typing import Any, List, Optional, Union
 
 import torch
 import torch.distributed as torch_distrib
+from torch.nn.parallel import DistributedDataParallel
 
+from pytorch_lightning import _logger as log
+from pytorch_lightning.accelerators.accelerator import Accelerator, ReduceOp
 from pytorch_lightning.core.lightning import LightningModule
 from pytorch_lightning.core.step_result import Result
 from pytorch_lightning.distributed.dist import LightningDistributed
-from pytorch_lightning import _logger as log
-from pytorch_lightning.accelerators.accelerator import Accelerator, ReduceOp
-from pytorch_lightning.utilities import AMPType, HYDRA_AVAILABLE
+from pytorch_lightning.utilities import HYDRA_AVAILABLE, AMPType
 from pytorch_lightning.utilities.distributed import rank_zero_only, sync_ddp_if_available
-from torch.nn.parallel import DistributedDataParallel
-from typing import List, Optional, Union, Any
 
 if HYDRA_AVAILABLE:
-    from hydra.utils import to_absolute_path, get_original_cwd
     from hydra.core.hydra_config import HydraConfig
+    from hydra.utils import get_original_cwd, to_absolute_path
 
 
 class DDP2Accelerator(Accelerator):
@@ -170,6 +170,8 @@ def ddp_train(self, process_idx, mp_queue, model):
         # 16-bit
         model = self.trainer.precision_connector.connect(model)
 
+        self.trainer.convert_to_lightning_optimizers()
+
         # device ids change depending on the DDP setup
         device_ids = self.get_device_ids()
 
diff --git a/pytorch_lightning/accelerators/ddp_accelerator.py b/pytorch_lightning/accelerators/ddp_accelerator.py
index a4904d315fc5d..687b5c21874fb 100644
--- a/pytorch_lightning/accelerators/ddp_accelerator.py
+++ b/pytorch_lightning/accelerators/ddp_accelerator.py
@@ -12,32 +12,29 @@
 # See the License for the specific language governing permissions and
 # limitations under the License
 import os
-import torch
-import torch.distributed as torch_distrib
 import subprocess
 import sys
 from os.path import abspath
 from time import sleep
-from typing import Any, Optional, List, Union
+from typing import Any, List, Optional, Union
 
 import numpy as np
+import torch
+import torch.distributed as torch_distrib
+from torch.nn.parallel import DistributedDataParallel
 
 from pytorch_lightning import _logger as log
 from pytorch_lightning.accelerators.accelerator import Accelerator, ReduceOp
 from pytorch_lightning.core.lightning import LightningModule
 from pytorch_lightning.distributed.dist import LightningDistributed
-from pytorch_lightning.utilities import AMPType, HYDRA_AVAILABLE
-from pytorch_lightning.utilities.distributed import find_free_network_port
-from pytorch_lightning.utilities.distributed import rank_zero_only
-from pytorch_lightning.utilities.distributed import sync_ddp_if_available
+from pytorch_lightning.utilities import HYDRA_AVAILABLE, AMPType
+from pytorch_lightning.utilities.distributed import find_free_network_port, rank_zero_only, sync_ddp_if_available
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
 from pytorch_lightning.utilities.seed import seed_everything
-from torch.nn.parallel import DistributedDataParallel
-
 
 if HYDRA_AVAILABLE:
-    from hydra.utils import to_absolute_path, get_original_cwd
     from hydra.core.hydra_config import HydraConfig
+    from hydra.utils import get_original_cwd, to_absolute_path
 
 
 class DDPAccelerator(Accelerator):
@@ -266,6 +263,8 @@ def ddp_train(self, process_idx, model):
         # 16-bit
         model = self.trainer.precision_connector.connect(model)
 
+        self.trainer.convert_to_lightning_optimizers()
+
         # device ids change depending on the DDP setup
         device_ids = self.get_device_ids()
 
diff --git a/pytorch_lightning/accelerators/ddp_cpu_spawn_accelerator.py b/pytorch_lightning/accelerators/ddp_cpu_spawn_accelerator.py
index feab74bd23f6b..982da2f53216b 100644
--- a/pytorch_lightning/accelerators/ddp_cpu_spawn_accelerator.py
+++ b/pytorch_lightning/accelerators/ddp_cpu_spawn_accelerator.py
@@ -23,10 +23,14 @@
 from pytorch_lightning import _logger as log
 from pytorch_lightning.accelerators.accelerator import Accelerator, ReduceOp
 from pytorch_lightning.core.lightning import LightningModule
-from pytorch_lightning.utilities import AMPType, HYDRA_AVAILABLE
-from pytorch_lightning.utilities.distributed import rank_zero_only, rank_zero_warn
-from pytorch_lightning.utilities.distributed import find_free_network_port, sync_ddp_if_available
 from pytorch_lightning.distributed.dist import LightningDistributed
+from pytorch_lightning.utilities import HYDRA_AVAILABLE, AMPType
+from pytorch_lightning.utilities.distributed import (
+    find_free_network_port,
+    rank_zero_only,
+    rank_zero_warn,
+    sync_ddp_if_available,
+)
 
 if HYDRA_AVAILABLE:
     from hydra.core.hydra_config import HydraConfig
@@ -130,6 +134,8 @@ def ddp_train(self, process_idx, mp_queue, model):
         # 16-bit
         model = self.trainer.precision_connector.connect(model)
 
+        self.trainer.convert_to_lightning_optimizers()
+
         # DDP spawn already spawned off each process... no need to do anything
         device_ids = self.get_device_ids()
 
diff --git a/pytorch_lightning/accelerators/ddp_hpc_accelerator.py b/pytorch_lightning/accelerators/ddp_hpc_accelerator.py
index 96f743bb76d43..28817c6845f5b 100644
--- a/pytorch_lightning/accelerators/ddp_hpc_accelerator.py
+++ b/pytorch_lightning/accelerators/ddp_hpc_accelerator.py
@@ -23,13 +23,12 @@
 from pytorch_lightning.accelerators.accelerator import Accelerator, ReduceOp
 from pytorch_lightning.core.lightning import LightningModule
 from pytorch_lightning.distributed.dist import LightningDistributed
-from pytorch_lightning.utilities import AMPType, HYDRA_AVAILABLE
+from pytorch_lightning.utilities import HYDRA_AVAILABLE, AMPType
 from pytorch_lightning.utilities.distributed import rank_zero_only, sync_ddp_if_available
 
-
 if HYDRA_AVAILABLE:
-    from hydra.utils import to_absolute_path, get_original_cwd
     from hydra.core.hydra_config import HydraConfig
+    from hydra.utils import get_original_cwd, to_absolute_path
 
 
 class DDPHPCAccelerator(Accelerator):
@@ -164,6 +163,8 @@ def ddp_train(self, process_idx, model):
         # 16-bit
         model = self.trainer.precision_connector.connect(model)
 
+        self.trainer.convert_to_lightning_optimizers()
+
         # device ids change depending on the DDP setup
         device_ids = self.get_device_ids()
 
diff --git a/pytorch_lightning/accelerators/ddp_spawn_accelerator.py b/pytorch_lightning/accelerators/ddp_spawn_accelerator.py
index c34fe18cac67f..a06d0b82d6d15 100644
--- a/pytorch_lightning/accelerators/ddp_spawn_accelerator.py
+++ b/pytorch_lightning/accelerators/ddp_spawn_accelerator.py
@@ -16,20 +16,25 @@
 from typing import Any, List, Optional, Union
 
 import torch
-import torch.multiprocessing as mp
 import torch.distributed as torch_distrib
 import torch.distributed as dist
+import torch.multiprocessing as mp
 from torch.nn.parallel import DistributedDataParallel
 
 from pytorch_lightning import _logger as log
 from pytorch_lightning.accelerators.accelerator import Accelerator, ReduceOp
 from pytorch_lightning.core.lightning import LightningModule
-from pytorch_lightning.utilities import AMPType, HYDRA_AVAILABLE
-from pytorch_lightning.utilities.cloud_io import atomic_save, load as pl_load
-from pytorch_lightning.utilities.distributed import rank_zero_only, rank_zero_warn, find_free_network_port
-from pytorch_lightning.utilities.distributed import sync_ddp_if_available
+from pytorch_lightning.distributed import LightningDistributed
+from pytorch_lightning.utilities import HYDRA_AVAILABLE, AMPType
+from pytorch_lightning.utilities.cloud_io import atomic_save
+from pytorch_lightning.utilities.cloud_io import load as pl_load
+from pytorch_lightning.utilities.distributed import (
+    find_free_network_port,
+    rank_zero_only,
+    rank_zero_warn,
+    sync_ddp_if_available,
+)
 from pytorch_lightning.utilities.seed import seed_everything
-from pytorch_lightning.distributed.dist import LightningDistributed
 
 if HYDRA_AVAILABLE:
     from hydra.core.hydra_config import HydraConfig
@@ -141,6 +146,8 @@ def ddp_train(self, process_idx, mp_queue, model, is_master=False, proc_offset=0
         # 16-bit
         model = self.trainer.precision_connector.connect(model)
 
+        self.trainer.convert_to_lightning_optimizers()
+
         # device ids change depending on the DDP setup
         device_ids = self.get_device_ids()
 
diff --git a/pytorch_lightning/accelerators/dp_accelerator.py b/pytorch_lightning/accelerators/dp_accelerator.py
index a894afba27b71..4b4e1eac8a66c 100644
--- a/pytorch_lightning/accelerators/dp_accelerator.py
+++ b/pytorch_lightning/accelerators/dp_accelerator.py
@@ -16,10 +16,10 @@
 import torch
 from torch import optim
 
-from pytorch_lightning.core.lightning import LightningModule
 from pytorch_lightning.accelerators.accelerator import Accelerator
-from pytorch_lightning.distributed import LightningDistributed
+from pytorch_lightning.core.lightning import LightningModule
 from pytorch_lightning.core.step_result import Result
+from pytorch_lightning.distributed import LightningDistributed
 from pytorch_lightning.overrides.data_parallel import LightningDataParallel
 from pytorch_lightning.utilities import AMPType
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
@@ -63,6 +63,8 @@ def setup(self, model):
         if self.trainer.amp_backend:
             model = self.__init_half_precision(model)
 
+        self.trainer.convert_to_lightning_optimizers()
+
         self.trainer.model = model
 
     def __init_torch_data_parallel(self, model):
diff --git a/pytorch_lightning/accelerators/gpu_accelerator.py b/pytorch_lightning/accelerators/gpu_accelerator.py
index 1a52c4037c8d3..b12d275c8ac26 100644
--- a/pytorch_lightning/accelerators/gpu_accelerator.py
+++ b/pytorch_lightning/accelerators/gpu_accelerator.py
@@ -11,13 +11,13 @@
 # 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 typing import Union, Optional, Any
+from typing import Any, Optional, Union
 
 import torch
 
 from pytorch_lightning.accelerators.accelerator import Accelerator, ReduceOp
-from pytorch_lightning.utilities import AMPType
 from pytorch_lightning.distributed.dist import LightningDistributed
+from pytorch_lightning.utilities import AMPType
 
 
 class GPUAccelerator(Accelerator):
@@ -52,6 +52,8 @@ def setup(self, model):
         # 16-bit
         model = self.trainer.precision_connector.connect(model)
 
+        self.trainer.convert_to_lightning_optimizers()
+
         self.trainer.model = model
 
     def train(self):
diff --git a/pytorch_lightning/accelerators/horovod_accelerator.py b/pytorch_lightning/accelerators/horovod_accelerator.py
index 3d9191914566d..b2cec906178f9 100644
--- a/pytorch_lightning/accelerators/horovod_accelerator.py
+++ b/pytorch_lightning/accelerators/horovod_accelerator.py
@@ -93,6 +93,8 @@ def _filter_named_parameters(model, optimizer):
         # 16-bit
         model = self.trainer.precision_connector.connect(model)
 
+        self.trainer.convert_to_lightning_optimizers()
+
         # Update logger rank info from Horovod to avoid race conditions from  different ranks
         # creating directories / writing files in the same locations.
         self.trainer.global_rank = hvd.rank()
diff --git a/pytorch_lightning/accelerators/tpu_accelerator.py b/pytorch_lightning/accelerators/tpu_accelerator.py
index 6e8314106112e..30cf6c9dbf169 100644
--- a/pytorch_lightning/accelerators/tpu_accelerator.py
+++ b/pytorch_lightning/accelerators/tpu_accelerator.py
@@ -14,7 +14,7 @@
 import io
 import os
 import re
-from typing import Optional, Union, Any
+from typing import Any, Optional, Union
 
 import torch
 import torch.multiprocessing as mp
@@ -23,7 +23,7 @@
 from pytorch_lightning import _logger as log
 from pytorch_lightning.accelerators.accelerator import Accelerator, ReduceOp
 from pytorch_lightning.core import LightningModule
-from pytorch_lightning.utilities import rank_zero_info, rank_zero_only, rank_zero_warn, TPU_AVAILABLE
+from pytorch_lightning.utilities import TPU_AVAILABLE, rank_zero_info, rank_zero_only, rank_zero_warn
 from pytorch_lightning.utilities.cloud_io import atomic_save
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
 
@@ -230,6 +230,8 @@ def __setup_tpu_training(self, model: LightningModule, trainer):
                  f' global rank: {trainer.tpu_global_core_rank}'
                  f' with XLA_USE_BF16={os.environ.get("XLA_USE_BF16")}')
 
+        self.trainer.convert_to_lightning_optimizers()
+
     def backward(self, closure_loss, optimizer, opt_idx, *args, **kwargs):
         # do backward pass
         if self.trainer.train_loop.automatic_optimization:
diff --git a/pytorch_lightning/core/lightning.py b/pytorch_lightning/core/lightning.py
index e888c6f4a87df..175dc56094b89 100644
--- a/pytorch_lightning/core/lightning.py
+++ b/pytorch_lightning/core/lightning.py
@@ -14,31 +14,34 @@
 
 """nn.Module with additional great features."""
 
-import os
-import tempfile
 import collections
 import copy
 import inspect
+import os
 import re
+import tempfile
 import types
 from abc import ABC
 from argparse import Namespace
-from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple, Union, Mapping
+from typing import Any, Callable, Dict, List, Mapping, Optional, Sequence, Tuple, Union
 
 import torch
+from torch import ScriptModule, Tensor
+from torch.nn import Module
+from torch.optim.optimizer import Optimizer
+
 from pytorch_lightning import _logger as log
+from pytorch_lightning.callbacks import Callback
 from pytorch_lightning.core.grads import GradInformation
 from pytorch_lightning.core.hooks import CheckpointHooks, DataHooks, ModelHooks
 from pytorch_lightning.core.memory import ModelSummary
+from pytorch_lightning.core.optimizer import LightningOptimizer
 from pytorch_lightning.core.saving import ALLOWED_CONFIG_TYPES, PRIMITIVE_TYPES, ModelIO
 from pytorch_lightning.core.step_result import Result
-from pytorch_lightning.utilities import rank_zero_warn, AMPType, TPU_AVAILABLE
+from pytorch_lightning.utilities import TPU_AVAILABLE, AMPType, rank_zero_warn
 from pytorch_lightning.utilities.device_dtype_mixin import DeviceDtypeModuleMixin
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
 from pytorch_lightning.utilities.parsing import AttributeDict, collect_init_args, get_init_args
-from torch import ScriptModule, Tensor
-from torch.nn import Module
-from torch.optim.optimizer import Optimizer
 
 if TPU_AVAILABLE:
     import torch_xla.core.xla_model as xm
@@ -116,10 +119,8 @@ def optimizers(self):
         # single optimizer
         if isinstance(opts, list) and len(opts) == 1 and isinstance(opts[0], Optimizer):
             return opts[0]
-
         # multiple opts
-        else:
-            return opts
+        return opts
 
     @property
     def example_input_array(self) -> Any:
@@ -1090,7 +1091,7 @@ 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)`
+           Trainer(accumulate_grad_batches=x) is set and you use `optimizer.step()`
 
         Example::
 
@@ -1099,7 +1100,7 @@ def training_step(...):
                 loss = ...
                 # automatically applies scaling, etc...
                 self.manual_backward(loss, opt_a)
-                self.manual_optimizer_step(opt_a)
+                opt_a.step()
         """
         # make sure we're using manual opt
         self._verify_is_manual_optimization('manual_backward')
@@ -1109,152 +1110,6 @@ def training_step(...):
         self.trainer.train_loop.backward(loss, optimizer, -1, *args, **kwargs)
         self._running_manual_backward = False
 
-    def manual_optimizer_step(self,
-                              optimizer: Optimizer,
-                              *args,
-                              make_optimizer_step: Optional[bool] = None,
-                              optimizer_closure: Optional[Callable] = None,
-                              ** kwargs) -> 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
-
-            make_optimizer_step: Whether to force an optimizer step. When nothing is provided,
-                we will use `accumulate_grad_batches` for accumulation frequency by default.
-                However, one coud provide True and False based on its own scheduling.
-                c.f example 2 and 3
-
-            optimizer_closure: One could provide its own optimizer_closure. Set to None by default.
-
-            args: Any parameters provided to optimizer.step()
-
-            kwargs: Any parameters provided to optimizer.step()
-
-        Example::
-
-            def training_step(...):
-                (opt_a, opt_b) = self.optimizers()
-                loss = ...
-                # automatically applies scaling, etc...
-
-                self.manual_backward(loss, opt_a)
-
-                # This will use accumulate gradients for `accumulate_grad_batches` batches
-                # and then run opt_a.step()
-                self.manual_optimizer_step(opt_a)
-
-        Example::
-
-            def training_step(self, batch, batch_idx):
-                # using Boring Model
-                opt = self.optimizers() # only 1 optimizer
-
-                def compute_loss():
-                    x = batch[0]
-                    x = F.dropout(x, 0.1)
-                    predictions = self(x)
-                    predictions = F.dropout(predictions, 0.1)
-                    loss = self.loss(None, predictions)
-                    return loss
-
-                def optimizer_closure():
-                    # emulate MC dropout training
-                    num_backward = 1
-                    losses = []
-                    for backward_idx in range(num_backward + 1):
-                        loss = compute_loss()
-                        losses.append(loss)
-                        retain_graph = num_backward!= backward_idx
-                        self.manual_backward(loss, opt, retain_graph=retain_graph)
-                    loss_mean = torch.stack(losses).mean()
-                    loss_std = torch.stack(losses).std()
-                    self.log("train_loss_mean", loss_mean, on_step=True, prog_bar=True, on_epoch=True)
-                    self.log("train_loss_std", loss_std, on_step=True, prog_bar=True, on_epoch=True)
-
-                self.manual_optimizer_step(opt, optimizer_closure=optimizer_closure)
-
-        Example::
-
-            # Scenario for a gan.
-
-            def training_step(self, batch, batch_idx, optimizer_idx):
-
-                # emulate gans training
-                opt_gen, opt_dis = self.optimizers()
-
-                # Note: Be careful, don't log on the same key in self.log in both closure
-                # as they will be aggregated together on epoch_end
-
-                def gen_closure():
-                    ... forward and compute loss for generator
-                    loss_gen = ...
-                    self.log("loss_gen", loss_gen, on_step=True, on_epoch=True)
-                    self.manual_backward(loss_gen, opt_gen)
-
-                def dis_closure():
-                    ... forward and compute loss for discriminator
-                    loss_dis = ...
-                    self.log("loss_dis", loss_dis, on_step=True, on_epoch=True)
-                    self.manual_backward(loss_dis, opt_dis)
-
-                # this will accumulate gradients for 2 batches and then call opt_gen.step()
-                self.manual_optimizer_step(
-                    opt_gen,
-                    optimizer_closure=gen_closure,
-                    make_optimizer_step=batch_idx % 2 == 0)
-
-                # update discriminator every 4 batches
-                # therefore, no gradient accumulation for discriminator
-                if batch_idx % 4 == 0 :
-                    # Note: Set make_optimizer_step to True or it will use by default
-                    # Trainer(accumulate_grad_batches=x)
-                    self.manual_optimizer_step(
-                        opt_dis,
-                        optimizer_closure=dis_closure,
-                        make_optimizer_step=True)
-        """
-        # make sure we're using manual opt
-        self._verify_is_manual_optimization('manual_optimizer_step')
-
-        should_make_optimizer_step = not self.trainer.train_loop.should_accumulate()
-        make_optimizer_step = make_optimizer_step if make_optimizer_step is not None else should_make_optimizer_step
-
-        if make_optimizer_step:
-
-            # mock closure function as the user is responsible to call `manual_backward`
-            def do_nothing_optimizer_closure():
-                return
-
-            is_callable = isinstance(optimizer_closure, types.FunctionType)
-            optimizer_closure = optimizer_closure if is_callable else do_nothing_optimizer_closure
-
-            self.trainer.train_loop.optimizer_step(
-                optimizer,
-                None,
-                self.trainer.batch_idx,
-                optimizer_closure,
-                *args,
-                **kwargs,
-            )
-
-            # update will be called after every optimizer_step call
-            if self.trainer.amp_backend == AMPType.NATIVE:
-                self.trainer.scaler.update()
-
-            # perform zero grad
-            optimizer.zero_grad()
-
-        else:
-            # make sure to call optimizer_closure when accumulating
-            if isinstance(optimizer_closure, types.FunctionType):
-                optimizer_closure()
-
     def backward(self, loss: Tensor, optimizer: Optimizer, optimizer_idx: int, *args, **kwargs) -> None:
         """
         Override backward with your own implementation if you need to.
@@ -1316,7 +1171,7 @@ def optimizer_step(
         By default, Lightning calls ``step()`` and ``zero_grad()`` as shown in the example
         once per optimizer.
 
-        .. tip:: Consider using `manual_optimizer_step` instead of overriding this method as done previously.
+        .. tip:: Consider using `.step()` instead of overriding this method as done previously.
 
         Warning:
             If you are overriding this method, make sure that you pass the ``optimizer_closure`` parameter
@@ -1386,16 +1241,10 @@ def optimizer_step(self, epoch, batch_idx, optimizer, optimizer_idx,
         """
         if on_tpu:
             xm.optimizer_step(optimizer, optimizer_args={'closure': optimizer_closure, **kwargs})
-        elif self.trainer.amp_backend == AMPType.NATIVE:
-            # native amp does not yet support closures.
-            # TODO: pass the closure to the step ASAP
-            optimizer_closure()
-            self.trainer.scaler.step(optimizer)
-        elif self.trainer.amp_backend == AMPType.APEX:
-            # apex amp does not yet support closures.
-            # TODO: pass the closure to the step ASAP
-            optimizer_closure()
-            optimizer.step(*args, **kwargs)
+
+        elif self.trainer.amp_backend is not None:
+            self.trainer.precision_connector.backend.step(self.trainer, optimizer, optimizer_closure)
+
         else:
             optimizer.step(closure=optimizer_closure, *args, **kwargs)
 
@@ -1707,11 +1556,10 @@ def to_onnx(self, file_path: str, input_sample: Optional[Tensor] = None, **kwarg
                 raise ValueError(
                     f"Received `input_sample` of type {type(input_sample)}. Expected type is `Tensor`"
                 )
-            else:
-                raise ValueError(
-                    "Could not export to ONNX since neither `input_sample` nor"
-                    " `model.example_input_array` attribute is set."
-                )
+            raise ValueError(
+                "Could not export to ONNX since neither `input_sample` nor"
+                " `model.example_input_array` attribute is set."
+            )
         input_data = input_data.to(self.device)
         if "example_outputs" not in kwargs:
             self.eval()
@@ -1813,7 +1661,6 @@ def hparams(self, hp: Union[dict, Namespace, Any]):
             self._hparams_initial = copy.deepcopy(self._hparams)
 
     def __get_hparams_assignment_variable(self):
-        """"""
         """
         looks at the code of the class to figure out what the user named self.hparams
         this only happens when the user explicitly sets self.hparams
diff --git a/pytorch_lightning/core/optimizer.py b/pytorch_lightning/core/optimizer.py
new file mode 100644
index 0000000000000..000e925c915c2
--- /dev/null
+++ b/pytorch_lightning/core/optimizer.py
@@ -0,0 +1,245 @@
+# Copyright The PyTorch Lightning team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+import collections
+import copy
+import inspect
+import os
+import re
+import tempfile
+import types
+from abc import ABC
+from argparse import Namespace
+from typing import Any, Callable, Dict, List, Mapping, Optional, Sequence, Tuple, Union
+from weakref import proxy
+
+import torch
+from torch import ScriptModule, Tensor
+from torch.nn import Module
+from torch.optim import SGD
+from torch.optim.optimizer import Optimizer
+
+from pytorch_lightning import _logger as log
+from pytorch_lightning.callbacks import Callback
+from pytorch_lightning.core.grads import GradInformation
+from pytorch_lightning.core.hooks import CheckpointHooks, DataHooks, ModelHooks
+from pytorch_lightning.core.memory import ModelSummary
+from pytorch_lightning.core.saving import ALLOWED_CONFIG_TYPES, PRIMITIVE_TYPES, ModelIO
+from pytorch_lightning.core.step_result import Result
+from pytorch_lightning.utilities import AMPType, rank_zero_warn
+from pytorch_lightning.utilities.device_dtype_mixin import DeviceDtypeModuleMixin
+from pytorch_lightning.utilities.exceptions import MisconfigurationException
+from pytorch_lightning.utilities.parsing import AttributeDict, collect_init_args, get_init_args
+from pytorch_lightning.utilities.xla_device_utils import XLADeviceUtils
+
+TPU_AVAILABLE = XLADeviceUtils.tpu_device_exists()
+
+if TPU_AVAILABLE:
+    import torch_xla.core.xla_model as xm
+
+
+def is_lightning_optimizer(optimizer):
+    return isinstance(optimizer, LightningOptimizer)
+
+
+def do_nothing_closure():
+    return
+
+
+class LightningOptimizer:
+    """
+    This class is used to wrap the user optimizers and handle properly
+    the backward and optimizer_step logic across accelerators, AMP, accumulated_grad_batches
+    """
+    def __init__(self,
+                 optimizer: Optimizer,
+                 accumulate_grad_batches: Optional[int] = None):
+
+        assert accumulate_grad_batches is None or isinstance(accumulate_grad_batches, int)
+        if isinstance(accumulate_grad_batches, int) and accumulate_grad_batches < 1:
+            raise MisconfigurationException(f"accumulate_grad_batches parameters "
+                                            f"{accumulate_grad_batches} should be >= 1")
+
+        self.__dict__ = {k: v for k, v in optimizer.__dict__.items() if k != 'step'}
+
+        # For Horovod
+        if hasattr(optimizer, "skip_synchronize"):
+            self.__class__ = type("Lightning" + optimizer.__class__.__name__, (self.__class__, optimizer.__class__.__bases__[0]), {})
+            self.skip_synchronize = optimizer.skip_synchronize
+            self.synchronize = optimizer.synchronize
+        else:
+            self.__class__ = type("Lightning" + optimizer.__class__.__name__, (self.__class__, optimizer.__class__), {})
+
+        self._trainer = None
+        self._optimizer = optimizer
+        self._accumulate_grad_batches = accumulate_grad_batches
+        self._use_accumulate_grad_batches_from_trainer = accumulate_grad_batches is None
+        # used to find back the step
+        self._lightning_step = self.step
+
+    def _on_trainer_init(self, trainer):
+        self._trainer = proxy(trainer)
+
+    def _accumulated_batches_reached(self):
+        if self._use_accumulate_grad_batches_from_trainer:
+            accumulate_grad_batches = self._trainer.accumulate_grad_batches
+        else:
+            accumulate_grad_batches = self._accumulate_grad_batches
+        return (self._trainer.batch_idx + 1) % accumulate_grad_batches == 0
+
+    @property
+    def _should_accumulate(self):
+        # checks if backward or backward + optimizer step (via closure)
+        accumulation_done = self._accumulated_batches_reached()
+        is_final_batch = self._trainer.train_loop._num_training_batches_reached()
+        return not (accumulation_done or is_final_batch)
+
+    def step(self, *args, closure: Optional[Callable] = None, make_optimizer_step: Optional[bool] = None, **kwargs):
+        """
+        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:
+
+            closure: One could provide its own optimizer_closure. Set to None by default.
+
+            make_optimizer_step: Whether to force an optimizer step. When nothing is provided,
+                we will use `accumulate_grad_batches` for accumulation frequency by default.
+                However, one coud provide True and False based on its own scheduling.
+                Refer to example 2 and 3
+
+            args: Any parameters provided to wrapped optimizer.step()
+
+            kwargs: Any parameters provided to wrapped optimizer.step()
+
+        Example::
+
+            def training_step(...):
+                (opt_a, opt_b) = self.optimizers()
+                loss_a = ...
+                # automatically applies scaling, etc...
+                self.manual_backward(loss_a, opt_a)
+                opt_a.step()
+
+        Example::
+
+            def training_step(self, batch, batch_idx):
+                # using Boring Model
+                opt = self.optimizers() # only 1 optimizer
+
+                def compute_loss():
+                    x = batch[0]
+                    x = F.dropout(x, 0.1)
+                    predictions = self(x)
+                    predictions = F.dropout(predictions, 0.1)
+                    loss = self.loss(None, predictions)
+                    return loss
+
+                def closure():
+                    # emulate MC dropout training
+                    num_backward = 1
+                    losses = []
+                    for backward_idx in range(num_backward + 1):
+                        loss = compute_loss()
+                        losses.append(loss)
+                        retain_graph = num_backward!= backward_idx
+                        self.manual_backward(loss, opt, retain_graph=retain_graph)
+                    loss_mean = torch.stack(losses).mean()
+                    loss_std = torch.stack(losses).std()
+                    self.log("train_loss_mean", loss_mean, on_step=True, prog_bar=True, on_epoch=True)
+                    self.log("train_loss_std", loss_std, on_step=True, prog_bar=True, on_epoch=True)
+
+                opt.step(loss, closure=closure)
+
+        Example::
+
+            # Scenario for a gan.
+
+            def training_step(self, batch, batch_idx, optimizer_idx):
+
+                # emulate gans training
+                opt_gen, opt_dis = self.optimizers()
+
+                # Note: Be careful, don't log on the same key in self.log in both closure
+                # as they will be aggregated together on epoch_end
+
+                def gen_closure():
+                    ... forward and compute loss for generator
+                    loss_gen = ...
+                    self.log("loss_gen", loss_gen, on_step=True, on_epoch=True)
+                    self.manual_backward(loss_gen, opt_gen)
+
+                def dis_closure():
+                    ... forward and compute loss for discriminator
+                    loss_dis = ...
+                    self.log("loss_dis", loss_dis, on_step=True, on_epoch=True)
+                    self.manual_backward(loss_dis, opt_dis)
+
+                # this will accumulate gradients for 2 batches and then call opt_gen.step()
+                opt_gen.step(closure=gen_closure, make_optimizer_step=batch_idx % 2 == 0)
+
+                # update discriminator every 4 batches
+                # therefore, no gradient accumulation for discriminator
+                if batch_idx % 4 == 0 :
+                    # Note: Set make_optimizer_step to True or it will use by default
+                    # Trainer(accumulate_grad_batches=x)
+                    opt_dis.step(closure=optimizer_closure, make_optimizer_step=True)
+        """
+
+        profiler_name = "optimizer_step_and_closure"
+        if closure is None:
+            closure = do_nothing_closure
+            profile_name = "optimizer_step"
+        else:
+            if not isinstance(closure, types.FunctionType):
+                raise MisconfigurationException("When closure is provided, it should be a function")
+
+        if make_optimizer_step is None:
+            make_optimizer_step = not self._should_accumulate
+
+        trainer = self._trainer
+        optimizer = self._optimizer
+
+        if make_optimizer_step:
+            if trainer.on_tpu:
+                with trainer.profiler.profile(profiler_name):
+                    xm.optimizer_step(optimizer, optimizer_args={'closure': closure, **kwargs})
+
+            elif trainer.amp_backend is not None:
+                trainer.precision_connector.backend.step(trainer, optimizer, closure)
+
+            else:
+                with trainer.profiler.profile(profiler_name):
+                    optimizer.step(closure=closure, *args, **kwargs)
+
+            # perform zero grad
+            optimizer.zero_grad()
+        else:
+            # make sure to call optimizer_closure when accumulating
+            with trainer.train_loop.block_ddp_sync_behaviour():
+                with trainer.profiler.profile("closure"):
+                    closure()
+
+    def __repr__(self):
+        groups = [
+            {
+                k: round(v, 12) if isinstance(v, float) else v
+                for k, v in sorted(group.items())
+                if k != "params"
+            }
+            for group in self.param_groups
+        ]
+        return f"{self.__class__.__name__}(groups={groups})"
diff --git a/pytorch_lightning/plugins/apex.py b/pytorch_lightning/plugins/apex.py
index 5cd5afa47a8cb..1b971b89da2d5 100644
--- a/pytorch_lightning/plugins/apex.py
+++ b/pytorch_lightning/plugins/apex.py
@@ -36,10 +36,6 @@ def connect(self, model, optimizers):
         self.trainer.reinit_scheduler_properties(optimizers, self.trainer.lr_schedulers)
         return model, optimizers
 
-    def training_step(self, fx, args):
-        output = fx(args)
-        return output
-
     def backward(self, closure_loss, optimizer, opt_idx, *args, **kwargs):
         closure_loss = amp.scale_loss(closure_loss, optimizer)
 
@@ -131,3 +127,11 @@ def clip_gradients(self, grad_clip_val: Union[int, float], optimizer: Optimizer,
     @property
     def norm_clipping_epsilon(self):
         return 1e-5
+
+    def step(self, trainer, optimizer, closure):
+        # apex amp does not yet support closures.
+        # TODO: pass the closure to the step ASAP
+        with trainer.profiler.profile("closure"):
+            closure()
+        with trainer.profiler.profile("optimizer_step"):
+            optimizer.step()
diff --git a/pytorch_lightning/plugins/native_amp.py b/pytorch_lightning/plugins/native_amp.py
index ff371965b7c83..1fb21338dafe7 100644
--- a/pytorch_lightning/plugins/native_amp.py
+++ b/pytorch_lightning/plugins/native_amp.py
@@ -51,11 +51,6 @@ def backward(self, closure_loss, optimizer, opt_idx, *args, **kwargs):
 
         return closure_loss
 
-    def training_step(self, fx, args):
-        with torch.cuda.amp.autocast():
-            output = fx(*args)
-        return output
-
     def clip_gradients(self, grad_clip_val: Union[int, float], optimizer: Optimizer, norm_type: float):
         model = self.trainer.get_model()
         torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=grad_clip_val, norm_type=norm_type)
@@ -63,3 +58,12 @@ def clip_gradients(self, grad_clip_val: Union[int, float], optimizer: Optimizer,
     @property
     def scaler(self):
         return torch.cuda.amp.GradScaler()
+
+    def step(self, trainer, optimizer, closure):
+        # native amp does not yet support closures.
+        # TODO: pass the closure to the step ASAP
+        with trainer.profiler.profile("closure"):
+            closure()
+        with trainer.profiler.profile("optimizer_step"):
+            trainer.scaler.step(optimizer)
+            trainer.scaler.update()
diff --git a/pytorch_lightning/plugins/sharded_plugin.py b/pytorch_lightning/plugins/sharded_plugin.py
index 84b23c3230059..1aeec7e1db96a 100644
--- a/pytorch_lightning/plugins/sharded_plugin.py
+++ b/pytorch_lightning/plugins/sharded_plugin.py
@@ -14,12 +14,14 @@
 from typing import List, Optional, Union
 
 from pytorch_lightning.core.lightning import LightningModule
+from pytorch_lightning.core.optimizer import is_lightning_optimizer
 from pytorch_lightning.plugins.ddp_plugin import DDPPlugin
-from pytorch_lightning.utilities import rank_zero_only, FAIRSCALE_AVAILABLE
+from pytorch_lightning.utilities import FAIRSCALE_AVAILABLE, rank_zero_only
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
 
 if FAIRSCALE_AVAILABLE:
     from fairscale.optim import OSS
+
     from pytorch_lightning.overrides.fairscale import LightningShardedDataParallel
 
 
@@ -62,6 +64,8 @@ def _wrap_optimizers(self, model):
     def _reinit_with_fairscale_oss(self, trainer):
         optimizers = trainer.optimizers
         for x, optimizer in enumerate(optimizers):
+            if is_lightning_optimizer(optimizer):
+                optimizer = optimizer._optimizer
             if not isinstance(optimizer, OSS):
                 optim_class = type(optimizer)
                 zero_optimizer = OSS(
@@ -71,6 +75,7 @@ def _reinit_with_fairscale_oss(self, trainer):
                 )
                 optimizers[x] = zero_optimizer
                 del optimizer
+        trainer.convert_to_lightning_optimizers()
 
     def get_model_from_plugin(
             self,
diff --git a/pytorch_lightning/trainer/connectors/optimizer_connector.py b/pytorch_lightning/trainer/connectors/optimizer_connector.py
index 3987b1b64ac25..70b18e037eba9 100644
--- a/pytorch_lightning/trainer/connectors/optimizer_connector.py
+++ b/pytorch_lightning/trainer/connectors/optimizer_connector.py
@@ -19,7 +19,8 @@ class OptimizerConnector:
     def __init__(self, trainer):
         self.trainer = trainer
 
-    def on_trainer_init(self):
+    def on_trainer_init(self, enable_pl_optimizer):
+        self.trainer.enable_pl_optimizer = enable_pl_optimizer
         self.trainer.lr_schedulers = []
         self.trainer.optimizers = []
         self.trainer.optimizer_frequencies = []
diff --git a/pytorch_lightning/trainer/connectors/precision_connector.py b/pytorch_lightning/trainer/connectors/precision_connector.py
index 8866607dc678c..c38931e9ccf5d 100644
--- a/pytorch_lightning/trainer/connectors/precision_connector.py
+++ b/pytorch_lightning/trainer/connectors/precision_connector.py
@@ -79,6 +79,8 @@ def _setup_amp_backend(self, amp_type: str, plugins: Optional[list]):
                 log.info('Using APEX 16bit precision.')
                 self.trainer.amp_backend = AMPType.APEX
                 self.backend = ApexPlugin(self.trainer)
+                log.warn("LightningOptimizer doesn't support Apex")
+                self.trainer.enable_pl_optimizer = False
 
         if not self.trainer.amp_backend:
             raise ModuleNotFoundError(
diff --git a/pytorch_lightning/trainer/optimizers.py b/pytorch_lightning/trainer/optimizers.py
index a6b63002dcc40..2e7e85a041ab2 100644
--- a/pytorch_lightning/trainer/optimizers.py
+++ b/pytorch_lightning/trainer/optimizers.py
@@ -20,6 +20,7 @@
 from torch.optim.optimizer import Optimizer
 
 from pytorch_lightning.core.lightning import LightningModule
+from pytorch_lightning.core.optimizer import LightningOptimizer
 from pytorch_lightning.utilities import rank_zero_warn
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
 
@@ -78,6 +79,16 @@ def init_optimizers(self, model: LightningModule) -> Tuple[List, List, List]:
 
         return optimizers, lr_schedulers, optimizer_frequencies
 
+    def convert_to_lightning_optimizers(self):
+        def convert_to_lightning_optimizer(trainer, optimizer):
+            if not isinstance(optimizer, LightningOptimizer):
+                optimizer = LightningOptimizer(optimizer)
+            optimizer._on_trainer_init(trainer)
+            return optimizer
+
+        if self.enable_pl_optimizer:
+            self.optimizers = [convert_to_lightning_optimizer(self, opt) for opt in self.optimizers]
+
     def configure_schedulers(self, schedulers: list, monitor: Optional[str] = None):
         # Convert each scheduler into dict structure with relevant information
         lr_schedulers = []
diff --git a/pytorch_lightning/trainer/properties.py b/pytorch_lightning/trainer/properties.py
index af06b1bbc1352..09022d5b37094 100644
--- a/pytorch_lightning/trainer/properties.py
+++ b/pytorch_lightning/trainer/properties.py
@@ -15,10 +15,14 @@
 import os
 from abc import ABC
 from argparse import ArgumentParser, Namespace
-from typing import List, Optional, Union, Type, TypeVar, cast
+from typing import List, Optional, Type, TypeVar, Union, cast
 
-from pytorch_lightning.callbacks import Callback, ProgressBarBase, ModelCheckpoint
+from pytorch_lightning.accelerators.accelerator import Accelerator
+from pytorch_lightning.callbacks import Callback, ModelCheckpoint, ProgressBarBase
 from pytorch_lightning.core.lightning import LightningModule
+from pytorch_lightning.core.optimizer import is_lightning_optimizer
+from pytorch_lightning.loggers.base import LightningLoggerBase
+from pytorch_lightning.loggers.tensorboard import TensorBoardLogger
 from pytorch_lightning.trainer.connectors.checkpoint_connector import CheckpointConnector
 from pytorch_lightning.trainer.connectors.logger_connector import LoggerConnector
 from pytorch_lightning.trainer.connectors.model_connector import ModelConnector
@@ -26,9 +30,6 @@
 from pytorch_lightning.utilities import argparse_utils
 from pytorch_lightning.utilities.cloud_io import get_filesystem
 from pytorch_lightning.utilities.model_utils import is_overridden
-from pytorch_lightning.accelerators.accelerator import Accelerator
-from pytorch_lightning.loggers.base import LightningLoggerBase
-from pytorch_lightning.loggers.tensorboard import TensorBoardLogger
 
 
 class TrainerProperties(ABC):
@@ -232,6 +233,16 @@ def save_checkpoint(self, filepath, weights_only: bool = False):
     def get_model(self):
         return self.model_connector.get_model()
 
+    def __getstate__(self):
+        # unwrap optimizer
+        self.optimizers = [opt._optimizer if is_lightning_optimizer(opt) else opt for opt in self.optimizers]
+        return self.__dict__
+
+    def __setstate__(self, d):
+        self.__dict__ = d
+        # wrap optimizers in enable_pl_optimzer is True
+        self.convert_to_lightning_optimizers()
+
 
 # Used to represent the concrete type TrainerProperties class methods are called on.
 _T = TypeVar('_T', bound=TrainerProperties)
diff --git a/pytorch_lightning/trainer/trainer.py b/pytorch_lightning/trainer/trainer.py
index eeb8b85fa4fa9..99a9eb24c3448 100644
--- a/pytorch_lightning/trainer/trainer.py
+++ b/pytorch_lightning/trainer/trainer.py
@@ -133,6 +133,7 @@ def __init__(
         distributed_backend: Optional[str] = None,
         automatic_optimization: Optional[bool] = None,
         move_metrics_to_cpu: bool = False,
+        enable_pl_optimizer: bool = True,
     ):
         r"""
         Customize every aspect of training via flags
@@ -275,8 +276,12 @@ def __init__(
                 Can be remote file paths such as `s3://mybucket/path` or 'hdfs://path/'
                 Defaults to `default_root_dir`.
 
-            move_metrics_to_cpu: Whether to force internal logged metrics to be moved to CPU.
-                This can save some GPU memory but can make training slower. Use with attention.
+            move_metrics_to_cpu: Whether to force internal logged metrics to be moved to cpu.
+                This can save some gpu memory, but can make training slower. Use with attention.
+
+            enable_pl_optimizer: If True, each optimizer will be wrapped by
+                `pytorch_lightning.core.optimizer.LightningOptimizer`. It allows Lightning to
+                handle AMP, TPU, accumulated_gradients, etc..
         """
         super().__init__()
 
@@ -322,7 +327,7 @@ def __init__(
         self.on_init_start()
 
         # init optimizer + lr scheduler related flags
-        self.optimizer_connector.on_trainer_init()
+        self.optimizer_connector.on_trainer_init(enable_pl_optimizer)
 
         # init data flags
         self.data_connector.on_trainer_init(
@@ -532,12 +537,11 @@ def train(self):
                     if met_min_epochs and met_min_steps:
                         self.train_loop.on_train_end()
                         return
-                    else:
-                        log.info(
-                            'Trainer was signaled to stop but required minimum epochs'
-                            f' ({self.min_epochs}) or minimum steps ({self.min_steps}) has'
-                            ' not been met. Training will continue...'
-                        )
+                    log.info(
+                        'Trainer was signaled to stop but required minimum epochs'
+                        f' ({self.min_epochs}) or minimum steps ({self.min_steps}) has'
+                        ' not been met. Training will continue...'
+                    )
 
             # hook
             self.train_loop.on_train_end()
diff --git a/pytorch_lightning/trainer/training_loop.py b/pytorch_lightning/trainer/training_loop.py
index 3d261d99a604b..74ce9c4b85625 100644
--- a/pytorch_lightning/trainer/training_loop.py
+++ b/pytorch_lightning/trainer/training_loop.py
@@ -22,6 +22,7 @@
 from pytorch_lightning.callbacks import ModelCheckpoint
 from pytorch_lightning.core.lightning import LightningModule
 from pytorch_lightning.core.memory import ModelSummary
+from pytorch_lightning.core.optimizer import LightningOptimizer
 from pytorch_lightning.core.step_result import EvalResult, Result
 from pytorch_lightning.trainer.states import TrainerState
 from pytorch_lightning.trainer.supporters import Accumulator, TensorRunningAccum
@@ -217,7 +218,8 @@ def check_checkpoint_callback(self, should_save, is_last=False):
             if is_last and any(c.save_last for c in checkpoint_callbacks):
                 rank_zero_info("Saving latest checkpoint...")
             model = self.trainer.get_model()
-            [c.on_validation_end(self.trainer, model) for c in checkpoint_callbacks]
+            for c in checkpoint_callbacks:
+                c.on_validation_end(self.trainer, model)
 
     def on_train_epoch_start(self, epoch):
 
@@ -474,11 +476,14 @@ def _process_result(self, training_step_output, split_batch):
         return training_step_output_for_epoch_end
 
     def optimizer_step(self, optimizer, opt_idx, batch_idx, train_step_and_backward_closure, *args, **kwargs):
-        with self.trainer.profiler.profile("optimizer_step"):
-            # optimizer step lightningModule hook
-            self.trainer.accelerator_backend.optimizer_step(
-                optimizer, batch_idx, opt_idx, train_step_and_backward_closure, *args, **kwargs
-            )
+        # optimizer step lightningModule hook
+        if isinstance(optimizer, LightningOptimizer):
+            optimizer.step(closure=train_step_and_backward_closure)
+        else:
+            with self.trainer.profiler.profile("optimizer_step"):
+                self.trainer.accelerator_backend.optimizer_step(
+                    optimizer, batch_idx, opt_idx, train_step_and_backward_closure, *args, **kwargs
+                )
 
     def on_before_zero_grad(self, optimizer):
         self.trainer.call_hook('on_before_zero_grad', optimizer)
@@ -948,7 +953,6 @@ def zero_grad_handler(self, batch_idx, optimizer, opt_idx):
             self.on_before_zero_grad(optimizer)
             optimizers = enumerate([optimizer])
         else:
-            # should be called handled in `manual_optimizer_step`
             optimizers = []
 
         for idx, optimizer in optimizers:
diff --git a/setup.cfg b/setup.cfg
index e559d9a30971d..2802806f6074d 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -38,6 +38,8 @@ exclude_lines =
     warnings
     pass
     rank_zero_warn
+    raise NotImplementedError
+
 # TODO: figure out how to get codecov to pick up the test results on these backends
 #  The actual coverage for each is 90%+
 # *metrics (94%+) are temporarily removed from testing while tests speed up
diff --git a/tests/checkpointing/test_model_checkpoint.py b/tests/checkpointing/test_model_checkpoint.py
index c8e7643556460..3a1aa54dc5e17 100644
--- a/tests/checkpointing/test_model_checkpoint.py
+++ b/tests/checkpointing/test_model_checkpoint.py
@@ -615,7 +615,8 @@ def validation_epoch_end(self, outputs):
 
 
 @mock.patch.dict(os.environ, {"PL_DEV_DEBUG": "1"})
-def test_checkpoint_repeated_strategy(tmpdir):
+@pytest.mark.parametrize("enable_pl_optimizer", [False, True])
+def test_checkpoint_repeated_strategy(enable_pl_optimizer, tmpdir):
     """
     This test validates that the checkpoint can be called when provided to callacks list
     """
@@ -637,7 +638,8 @@ def validation_step(self, batch, batch_idx):
         limit_train_batches=2,
         limit_val_batches=2,
         limit_test_batches=2,
-        callbacks=[checkpoint_callback]
+        callbacks=[checkpoint_callback],
+        enable_pl_optimizer=enable_pl_optimizer
     )
 
     trainer.fit(model)
@@ -657,7 +659,8 @@ def get_last_checkpoint():
                              limit_train_batches=2,
                              limit_val_batches=2,
                              limit_test_batches=2,
-                             resume_from_checkpoint=chk)
+                             resume_from_checkpoint=chk,
+                             enable_pl_optimizer=enable_pl_optimizer)
         trainer.fit(model)
         trainer.test(model)
 
@@ -665,7 +668,8 @@ def get_last_checkpoint():
 
 
 @mock.patch.dict(os.environ, {"PL_DEV_DEBUG": "1"})
-def test_checkpoint_repeated_strategy_tmpdir(tmpdir):
+@pytest.mark.parametrize("enable_pl_optimizer", [False, True])
+def test_checkpoint_repeated_strategy_tmpdir(enable_pl_optimizer, tmpdir):
     """
     This test validates that the checkpoint can be called when provided to callacks list
     """
@@ -688,7 +692,8 @@ def validation_step(self, batch, batch_idx):
         limit_train_batches=2,
         limit_val_batches=2,
         limit_test_batches=2,
-        callbacks=[checkpoint_callback])
+        callbacks=[checkpoint_callback],
+        enable_pl_optimizer=enable_pl_optimizer)
 
     trainer.fit(model)
     assert sorted(os.listdir(tmpdir)) == sorted(['epoch=00.ckpt', 'lightning_logs'])
@@ -711,7 +716,8 @@ def get_last_checkpoint():
                              limit_train_batches=2,
                              limit_val_batches=2,
                              limit_test_batches=2,
-                             resume_from_checkpoint=chk)
+                             resume_from_checkpoint=chk,
+                             enable_pl_optimizer=enable_pl_optimizer)
 
         trainer.fit(model)
         trainer.test(model)
@@ -720,7 +726,8 @@ def get_last_checkpoint():
 
 
 @mock.patch.dict(os.environ, {"PL_DEV_DEBUG": "1"})
-def test_checkpoint_repeated_strategy_extended(tmpdir):
+@pytest.mark.parametrize("enable_pl_optimizer", [False, True])
+def test_checkpoint_repeated_strategy_extended(enable_pl_optimizer, tmpdir):
     """
     This test validates checkpoint can be called several times without
     increasing internally its global step if nothing run.
@@ -772,6 +779,7 @@ def get_model():
         limit_train_batches=limit_train_batches,
         limit_val_batches=3,
         limit_test_batches=4,
+        enable_pl_optimizer=enable_pl_optimizer
     )
 
     trainer = pl.Trainer(
diff --git a/tests/checkpointing/test_torch_saving.py b/tests/checkpointing/test_torch_saving.py
index f9279976a0069..493aa0dabe126 100644
--- a/tests/checkpointing/test_torch_saving.py
+++ b/tests/checkpointing/test_torch_saving.py
@@ -18,16 +18,19 @@
 import torch
 
 from pytorch_lightning import Trainer
+from pytorch_lightning.core.optimizer import LightningOptimizer
 from tests.base import BoringModel
 
 
-def test_model_torch_save(tmpdir):
+@pytest.mark.parametrize("enable_pl_optimizer", [False, True])
+def test_model_torch_save(tmpdir, enable_pl_optimizer):
     """Test to ensure torch save does not fail for model and trainer."""
     model = BoringModel()
     num_epochs = 1
     trainer = Trainer(
         default_root_dir=tmpdir,
         max_epochs=num_epochs,
+        enable_pl_optimizer=enable_pl_optimizer,
     )
     temp_path = os.path.join(tmpdir, 'temp.pt')
     trainer.fit(model)
@@ -35,6 +38,9 @@ def test_model_torch_save(tmpdir):
     # Ensure these do not fail
     torch.save(trainer.model, temp_path)
     torch.save(trainer, temp_path)
+    trainer = torch.load(temp_path)
+    is_lightning_optimizer = isinstance(trainer.optimizers[0], LightningOptimizer)
+    assert is_lightning_optimizer if enable_pl_optimizer else not is_lightning_optimizer
 
 
 @pytest.mark.skipif(platform.system() == "Windows",
diff --git a/tests/core/test_lightning_optimizer.py b/tests/core/test_lightning_optimizer.py
new file mode 100644
index 0000000000000..9b787848d2100
--- /dev/null
+++ b/tests/core/test_lightning_optimizer.py
@@ -0,0 +1,192 @@
+# Copyright The PyTorch Lightning team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+import os
+from unittest.mock import patch
+
+import pytest
+import torch
+import torch.nn as nn
+from torch.optim import Adam, Optimizer
+
+from pytorch_lightning import LightningModule, Trainer
+from pytorch_lightning.core.optimizer import LightningOptimizer
+from tests.base.boring_model import BoringModel, RandomDictDataset, RandomDictStringDataset
+
+
+def test_lightning_optimizer(tmpdir):
+    """
+    Test that optimizer are correctly wrapped by our LightningOptimizer
+    """
+    class TestModel(BoringModel):
+        def configure_optimizers(self):
+            optimizer = torch.optim.SGD(self.layer.parameters(), lr=0.1)
+            # optimizer = LightningOptimizer(self.trainer, optimizer)
+            lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=1)
+            return [optimizer], [lr_scheduler]
+    model = TestModel()
+    trainer = Trainer(
+        default_root_dir=os.getcwd(),
+        limit_train_batches=1,
+        limit_val_batches=1,
+        max_epochs=1,
+        weights_summary=None,
+        enable_pl_optimizer=True)
+    trainer.fit(model)
+
+    groups = "{'dampening': 0, 'initial_lr': 0.1, 'lr': 0.01, 'momentum': 0, 'nesterov': False, 'weight_decay': 0}"
+    expected = f"LightningSGD(groups=[{groups}])"
+    assert trainer.optimizers[0].__repr__() == expected
+
+
+def test_lightning_optimizer_from_user(tmpdir):
+    """
+    Test that the user can use our LightningOptimizer. Not recommended.
+    """
+
+    class TestModel(BoringModel):
+        def configure_optimizers(self):
+            optimizer = torch.optim.Adam(self.layer.parameters(), lr=0.1)
+            optimizer = LightningOptimizer(optimizer)
+            lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=1)
+            return [optimizer], [lr_scheduler]
+    model = TestModel()
+    trainer = Trainer(
+        default_root_dir=os.getcwd(),
+        limit_train_batches=1,
+        limit_val_batches=1,
+        max_epochs=1,
+        weights_summary=None,
+        enable_pl_optimizer=True)
+    trainer.fit(model)
+
+    groups = "{'amsgrad': False, 'betas': (0.9, 0.999), 'eps': 1e-08, 'initial_lr': 0.1, 'lr': 0.01, 'weight_decay': 0}"
+    expected = f"LightningAdam(groups=[{groups}])"
+    assert trainer.optimizers[0].__repr__() == expected
+
+
+@patch("torch.optim.Adam.step")
+@patch("torch.optim.SGD.step")
+def test_lightning_optimizer_manual_optimization(mock_sgd_step, mock_adam_step, tmpdir):
+    """
+    Test that the user can use our LightningOptimizer. Not recommended for now.
+    """
+    class TestModel(BoringModel):
+
+        def training_step(self, batch, batch_idx, optimizer_idx=None):
+            (opt_1, opt_2) = self.optimizers()
+            assert isinstance(opt_1, LightningOptimizer)
+            assert isinstance(opt_2, LightningOptimizer)
+
+            output = self.layer(batch)
+            loss_1 = self.loss(batch, output)
+            self.manual_backward(loss_1, opt_1)
+            opt_1.step(idx="1")
+
+            def closure():
+                output = self.layer(batch)
+                loss_2 = self.loss(batch, output)
+                self.manual_backward(loss_2, opt_2)
+            opt_2.step(closure=closure, idx="2")
+
+        def configure_optimizers(self):
+            optimizer_1 = torch.optim.SGD(self.layer.parameters(), lr=0.1)
+            optimizer_2 = torch.optim.Adam(self.layer.parameters(), lr=0.1)
+            optimizer_1 = LightningOptimizer(optimizer_1, 4)
+
+            lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer_1, step_size=1)
+            return [optimizer_1, optimizer_2], [lr_scheduler]
+
+    model = TestModel()
+    model.training_step_end = None
+    model.training_epoch_end = None
+    trainer = Trainer(
+        default_root_dir=os.getcwd(),
+        limit_train_batches=8,
+        limit_val_batches=1,
+        max_epochs=1,
+        weights_summary=None,
+        automatic_optimization=False,
+        enable_pl_optimizer=True)
+    trainer.fit(model)
+
+    assert len(mock_sgd_step.mock_calls) == 2
+    assert len(mock_adam_step.mock_calls) == 8
+
+
+@patch("torch.optim.Adam.step")
+@patch("torch.optim.SGD.step")
+def test_lightning_optimizer_manual_optimization_and_accumulated_gradients(mock_sgd_step, mock_adam_step, tmpdir):
+    """
+    Test that the user can use our LightningOptimizer. Not recommended.
+    """
+    class TestModel(BoringModel):
+
+        def training_step(self, batch, batch_idx, optimizer_idx=None):
+            (opt_1, opt_2) = self.optimizers()
+            assert isinstance(opt_1, LightningOptimizer)
+            assert isinstance(opt_2, LightningOptimizer)
+
+            output = self.layer(batch)
+            loss_1 = self.loss(batch, output)
+            self.manual_backward(loss_1, opt_1)
+            opt_1.step(idx="1")
+
+            def closure():
+                output = self.layer(batch)
+                loss_2 = self.loss(batch, output)
+                self.manual_backward(loss_2, opt_2)
+            opt_2.step(closure=closure, idx="2")
+
+        def configure_optimizers(self):
+            optimizer_1 = torch.optim.SGD(self.layer.parameters(), lr=0.1)
+            optimizer_2 = torch.optim.Adam(self.layer.parameters(), lr=0.1)
+            optimizer_1 = LightningOptimizer(optimizer_1, 4)
+
+            lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer_1, step_size=1)
+            return [optimizer_1, optimizer_2], [lr_scheduler]
+
+    model = TestModel()
+    model.training_step_end = None
+    model.training_epoch_end = None
+    trainer = Trainer(
+        default_root_dir=os.getcwd(),
+        limit_train_batches=8,
+        limit_val_batches=1,
+        max_epochs=1,
+        weights_summary=None,
+        automatic_optimization=False,
+        accumulate_grad_batches=2,
+        enable_pl_optimizer=True)
+    trainer.fit(model)
+
+    assert len(mock_sgd_step.mock_calls) == 2
+    assert len(mock_adam_step.mock_calls) == 4
+
+
+def test_state(tmpdir):
+    model = torch.nn.Linear(3, 4)
+    optimizer = torch.optim.Adam(model.parameters())
+    lightning_optimizer = LightningOptimizer(optimizer)
+    assert isinstance(lightning_optimizer, LightningOptimizer)
+    assert isinstance(lightning_optimizer, Adam)
+    assert isinstance(lightning_optimizer, Optimizer)
+    lightning_dict = {}
+    special_attrs = ["_accumulate_grad_batches", "_optimizer", "_optimizer_idx",
+                     "_trainer", "_use_accumulate_grad_batches_from_trainer", "_lightning_step"]
+    for k, v in lightning_optimizer.__dict__.items():
+        if k not in special_attrs:
+            lightning_dict[k] = v
+    assert lightning_dict == optimizer.__dict__
+    assert optimizer.state_dict() == lightning_optimizer.state_dict()
+    assert optimizer.state == lightning_optimizer.state
diff --git a/tests/models/test_amp.py b/tests/models/test_amp.py
index 29277d6fccbd2..b1e623490b50d 100644
--- a/tests/models/test_amp.py
+++ b/tests/models/test_amp.py
@@ -22,9 +22,9 @@
 from pytorch_lightning import Trainer
 from pytorch_lightning.loggers import WandbLogger
 from pytorch_lightning.trainer.states import TrainerState
+from pytorch_lightning.utilities import APEX_AVAILABLE
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
 from tests.base import EvalModelTemplate
-from pytorch_lightning.utilities import APEX_AVAILABLE
 
 
 @pytest.mark.skip(reason='dp + amp not supported currently')  # TODO
@@ -145,7 +145,8 @@ def test_amp_gpu_ddp_slurm_managed(tmpdir):
     assert trainer.slurm_connector.resolve_root_node_address('abc[23-24, 45-40, 40]') == 'abc23'
 
 
-def test_cpu_model_with_amp(tmpdir):
+@pytest.mark.parametrize("enable_pl_optimizer", [False, True])
+def test_cpu_model_with_amp(enable_pl_optimizer, tmpdir):
     """Make sure model trains on CPU."""
     trainer_options = dict(
         default_root_dir=tmpdir,
@@ -153,7 +154,8 @@ def test_cpu_model_with_amp(tmpdir):
         max_epochs=1,
         limit_train_batches=0.4,
         limit_val_batches=0.4,
-        precision=16
+        precision=16,
+        enable_pl_optimizer=enable_pl_optimizer
     )
 
     model = EvalModelTemplate()
diff --git a/tests/models/test_cpu.py b/tests/models/test_cpu.py
index d71f9243e5a83..e7ab4d0eb53d1 100644
--- a/tests/models/test_cpu.py
+++ b/tests/models/test_cpu.py
@@ -21,13 +21,13 @@
 import tests.base.develop_pipelines as tpipes
 import tests.base.develop_utils as tutils
 from pytorch_lightning import Trainer
-from pytorch_lightning.callbacks import EarlyStopping
-from pytorch_lightning.callbacks import ModelCheckpoint
+from pytorch_lightning.callbacks import EarlyStopping, ModelCheckpoint
 from pytorch_lightning.core.step_result import TrainResult
 from tests.base import EvalModelTemplate
 
 
-def test_cpu_slurm_save_load(tmpdir):
+@pytest.mark.parametrize("enable_pl_optimizer", [False, True])
+def test_cpu_slurm_save_load(enable_pl_optimizer, tmpdir):
     """Verify model save/load/checkpoint on CPU."""
     hparams = EvalModelTemplate.get_default_hparams()
     model = EvalModelTemplate(**hparams)
@@ -44,6 +44,7 @@ def test_cpu_slurm_save_load(tmpdir):
         limit_train_batches=0.2,
         limit_val_batches=0.2,
         checkpoint_callback=ModelCheckpoint(dirpath=tmpdir),
+        enable_pl_optimizer=enable_pl_optimizer
     )
     result = trainer.fit(model)
     real_global_step = trainer.global_step
@@ -80,6 +81,7 @@ def test_cpu_slurm_save_load(tmpdir):
         max_epochs=1,
         logger=logger,
         checkpoint_callback=ModelCheckpoint(dirpath=tmpdir),
+        enable_pl_optimizer=enable_pl_optimizer
     )
     model = EvalModelTemplate(**hparams)
 
@@ -99,7 +101,8 @@ def assert_pred_same():
     trainer.fit(model)
 
 
-def test_early_stopping_cpu_model(tmpdir):
+@pytest.mark.parametrize("enable_pl_optimizer", [False, True])
+def test_early_stopping_cpu_model(enable_pl_optimizer, tmpdir):
     """Test each of the trainer options."""
     stopping = EarlyStopping(monitor='early_stop_on', min_delta=0.1)
     trainer_options = dict(
@@ -111,6 +114,7 @@ def test_early_stopping_cpu_model(tmpdir):
         track_grad_norm=2,
         limit_train_batches=0.1,
         limit_val_batches=0.1,
+        enable_pl_optimizer=enable_pl_optimizer
     )
 
     model = EvalModelTemplate()
@@ -126,7 +130,8 @@ def test_early_stopping_cpu_model(tmpdir):
 @pytest.mark.skipif((platform.system() == "Darwin" and
                      LooseVersion(torch.__version__) < LooseVersion("1.3.0")),
                     reason="Distributed training is not supported on MacOS before Torch 1.3.0")
-def test_multi_cpu_model_ddp(tmpdir):
+@pytest.mark.parametrize("enable_pl_optimizer", [False, True])
+def test_multi_cpu_model_ddp(enable_pl_optimizer, tmpdir):
     """Make sure DDP works."""
     tutils.set_random_master_port()
 
@@ -139,6 +144,7 @@ def test_multi_cpu_model_ddp(tmpdir):
         gpus=None,
         num_processes=2,
         distributed_backend='ddp_cpu',
+        enable_pl_optimizer=enable_pl_optimizer
     )
 
     model = EvalModelTemplate()
@@ -278,7 +284,8 @@ def test_cpu_model(tmpdir):
     tpipes.run_model_test(trainer_options, model, on_gpu=False)
 
 
-def test_all_features_cpu_model(tmpdir):
+@pytest.mark.parametrize("enable_pl_optimizer", [False, True])
+def test_all_features_cpu_model(enable_pl_optimizer, tmpdir):
     """Test each of the trainer options."""
     trainer_options = dict(
         default_root_dir=tmpdir,
@@ -289,7 +296,8 @@ def test_all_features_cpu_model(tmpdir):
         accumulate_grad_batches=2,
         max_epochs=1,
         limit_train_batches=0.4,
-        limit_val_batches=0.4
+        limit_val_batches=0.4,
+        enable_pl_optimizer=enable_pl_optimizer
     )
 
     model = EvalModelTemplate()
diff --git a/tests/models/test_horovod.py b/tests/models/test_horovod.py
index ace1577514b7d..0fc68a226eae6 100644
--- a/tests/models/test_horovod.py
+++ b/tests/models/test_horovod.py
@@ -21,14 +21,13 @@
 import numpy as np
 import pytest
 import torch
-
 from sklearn.metrics import accuracy_score
 
 import tests.base.develop_pipelines as tpipes
 import tests.base.develop_utils as tutils
 from pytorch_lightning import Trainer
 from pytorch_lightning.accelerators.horovod_accelerator import HorovodAccelerator
-from pytorch_lightning.core.step_result import Result, TrainResult, EvalResult
+from pytorch_lightning.core.step_result import EvalResult, Result, TrainResult
 from pytorch_lightning.metrics.classification.accuracy import Accuracy
 from pytorch_lightning.utilities import APEX_AVAILABLE, NATIVE_AMP_AVAILABLE
 from tests.base import EvalModelTemplate
@@ -78,7 +77,8 @@ def _run_horovod(trainer_options, on_gpu=False):
 
 
 @pytest.mark.skipif(platform.system() == "Windows", reason="Horovod is not supported on Windows")
-def test_horovod_cpu(tmpdir):
+@pytest.mark.parametrize("enable_pl_optimizer", [False, True])
+def test_horovod_cpu(enable_pl_optimizer, tmpdir):
     """Test Horovod running multi-process on CPU."""
     trainer_options = dict(
         default_root_dir=str(tmpdir),
@@ -90,12 +90,14 @@ def test_horovod_cpu(tmpdir):
         limit_val_batches=0.2,
         distributed_backend='horovod',
         deterministic=True,
+        enable_pl_optimizer=enable_pl_optimizer,
     )
     _run_horovod(trainer_options)
 
 
 @pytest.mark.skipif(platform.system() == "Windows", reason="Horovod is not supported on Windows")
-def test_horovod_cpu_implicit(tmpdir):
+@pytest.mark.parametrize("enable_pl_optimizer", [False, True])
+def test_horovod_cpu_implicit(enable_pl_optimizer, tmpdir):
     """Test Horovod without specifying a backend, inferring from env set by `horovodrun`."""
     trainer_options = dict(
         default_root_dir=str(tmpdir),
@@ -106,6 +108,7 @@ def test_horovod_cpu_implicit(tmpdir):
         limit_train_batches=0.4,
         limit_val_batches=0.2,
         deterministic=True,
+        enable_pl_optimizer=enable_pl_optimizer,
     )
     _run_horovod(trainer_options)
 
@@ -211,7 +214,8 @@ def validation_step(self, batch, *args, **kwargs):
 
 
 @pytest.mark.skipif(platform.system() == "Windows", reason="Horovod is not supported on Windows")
-def test_horovod_multi_optimizer(tmpdir):
+@pytest.mark.parametrize("enable_pl_optimizer", [False, True])
+def test_horovod_multi_optimizer(enable_pl_optimizer, tmpdir):
     model = BasicGAN(**EvalModelTemplate.get_default_hparams())
 
     # fit model
@@ -223,6 +227,7 @@ def test_horovod_multi_optimizer(tmpdir):
         limit_val_batches=0.2,
         deterministic=True,
         distributed_backend='horovod',
+        enable_pl_optimizer=enable_pl_optimizer,
     )
     result = trainer.fit(model)
     assert result == 1, 'model failed to complete'
@@ -244,7 +249,8 @@ def get_optimizer_params(optimizer):
 
 @pytest.mark.skipif(not HOROVOD_AVAILABLE, reason="Horovod is unavailable")
 @pytest.mark.skipif(platform.system() == "Windows", reason="Horovod is not supported on Windows")
-def test_result_reduce_horovod(tmpdir):
+@pytest.mark.parametrize("enable_pl_optimizer", [False, True])
+def test_result_reduce_horovod(enable_pl_optimizer, tmpdir):
     """Make sure result logging works with Horovod.
 
     This test mirrors tests/core/test_results.py::_ddp_test_fn
@@ -257,10 +263,10 @@ def hvd_test_fn():
         path_root = os.path.abspath(os.path.join(path_here, '..', '..'))
         sys.path.insert(0, os.path.abspath(path_root))
 
-        from tests.base.boring_model import BoringModel
-
         import horovod.torch as hvd
 
+        from tests.base.boring_model import BoringModel
+
         class TestModel(BoringModel):
             def training_step(self, batch, batch_idx):
                 self.training_step_called = True
@@ -288,6 +294,7 @@ def training_epoch_end(self, outputs) -> None:
             max_epochs=1,
             log_every_n_steps=1,
             weights_summary=None,
+            enable_pl_optimizer=enable_pl_optimizer,
         )
 
         trainer.fit(model)
diff --git a/tests/models/test_restore.py b/tests/models/test_restore.py
index 1977c5545b726..26b6a36d2314b 100644
--- a/tests/models/test_restore.py
+++ b/tests/models/test_restore.py
@@ -25,7 +25,7 @@
 
 import tests.base.develop_pipelines as tpipes
 import tests.base.develop_utils as tutils
-from pytorch_lightning import Trainer, LightningModule, Callback, seed_everything
+from pytorch_lightning import Callback, LightningModule, Trainer, seed_everything
 from pytorch_lightning.callbacks import ModelCheckpoint
 from tests.base import EvalModelTemplate, GenericEvalModelTemplate, TrialMNIST
 
@@ -52,7 +52,8 @@ def on_train_end(self, trainer, pl_module):
         self._check_properties(trainer, pl_module)
 
 
-def test_model_properties_resume_from_checkpoint(tmpdir):
+@pytest.mark.parametrize("enable_pl_optimizer", [False, True])
+def test_model_properties_resume_from_checkpoint(enable_pl_optimizer, tmpdir):
     """ Test that properties like `current_epoch` and `global_step`
     in model and trainer are always the same. """
     model = EvalModelTemplate()
@@ -61,6 +62,7 @@ def test_model_properties_resume_from_checkpoint(tmpdir):
         default_root_dir=tmpdir,
         max_epochs=1,
         logger=False,
+        enable_pl_optimizer=enable_pl_optimizer,
         callbacks=[checkpoint_callback, ModelTrainerPropertyParity()]  # this performs the assertions
     )
     trainer = Trainer(**trainer_args)
@@ -78,7 +80,8 @@ def on_train_start(self, trainer, pl_module):
         self.callbacks = deepcopy(trainer.callbacks)
 
 
-def test_callbacks_state_resume_from_checkpoint(tmpdir):
+@pytest.mark.parametrize("enable_pl_optimizer", [False, True])
+def test_callbacks_state_resume_from_checkpoint(enable_pl_optimizer, tmpdir):
     """ Test that resuming from a checkpoint restores callbacks that persist state. """
     model = EvalModelTemplate()
     callback_capture = CaptureCallbacksBeforeTraining()
@@ -89,6 +92,7 @@ def get_trainer_args():
             default_root_dir=tmpdir,
             max_steps=1,
             logger=False,
+            enable_pl_optimizer=enable_pl_optimizer,
             callbacks=[
                 checkpoint,
                 callback_capture,
@@ -115,10 +119,11 @@ def get_trainer_args():
             assert before.best_model_score == after.best_model_score
 
 
-def test_callbacks_references_resume_from_checkpoint(tmpdir):
+@pytest.mark.parametrize("enable_pl_optimizer", [False, True])
+def test_callbacks_references_resume_from_checkpoint(enable_pl_optimizer, tmpdir):
     """ Test that resuming from a checkpoint sets references as expected. """
     model = EvalModelTemplate()
-    args = {'default_root_dir': tmpdir, 'max_steps': 1, 'logger': False}
+    args = {'default_root_dir': tmpdir, 'max_steps': 1, 'logger': False, "enable_pl_optimizer": enable_pl_optimizer}
 
     # initial training
     checkpoint = ModelCheckpoint(dirpath=tmpdir, monitor="early_stop_on", save_last=True)
diff --git a/tests/trainer/optimization/test_manual_optimization.py b/tests/trainer/optimization/test_manual_optimization.py
index a58ebfc85383c..1a904daebefbc 100644
--- a/tests/trainer/optimization/test_manual_optimization.py
+++ b/tests/trainer/optimization/test_manual_optimization.py
@@ -14,11 +14,11 @@
 import collections
 import os
 from unittest import mock
+from unittest.mock import ANY, call, patch
 
 import pytest
 import torch
 import torch.nn.functional as F
-from unittest.mock import patch, call, ANY
 
 from pytorch_lightning import Trainer, seed_everything
 from pytorch_lightning.utilities import APEX_AVAILABLE
@@ -383,7 +383,7 @@ def training_step(self, batch, batch_idx):
         if self.should_update:
 
             self.manual_backward(loss, opt)
-            self.manual_optimizer_step(opt)
+            opt.step()
 
         return loss.detach() if self.detach else loss
 
@@ -434,6 +434,7 @@ def test_manual_optimization_and_return_tensor(tmpdir):
         amp_backend='native',
         accelerator="ddp_spawn",
         gpus=2,
+        enable_pl_optimizer=True
     )
     trainer.fit(model)
 
@@ -473,7 +474,7 @@ def test_manual_optimization_and_return_detached_tensor(tmpdir):
 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.
+    step is being called only when we shouldn't accumulate.
     """
     seed_everything(234)
 
@@ -511,7 +512,7 @@ def training_step(self, batch, batch_idx):
             if self.should_update:
 
                 self.manual_backward(loss, opt)
-                self.manual_optimizer_step(opt)
+                opt.step()
 
             return loss.detach() if self.detach else loss
 
@@ -550,15 +551,16 @@ def on_train_end(self):
         amp_backend='native',
         accumulate_grad_batches=4,
         gpus=1,
+        enable_pl_optimizer=True,
     )
     trainer.fit(model)
 
 
 @mock.patch.dict(os.environ, {"PL_DEV_DEBUG": "1"})
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="test requires GPU machine")
-def test_multiple_optimizers_manual_optimizer_step(tmpdir):
+def test_multiple_optimizers_step(tmpdir):
     """
-    Tests that `manual_optimizer_step` works with several optimizers
+    Tests that `step` works with several optimizers
     """
     class TestModel(BoringModel):
         def training_step(self, batch, batch_idx, optimizer_idx):
@@ -574,7 +576,7 @@ def training_step(self, batch, batch_idx, optimizer_idx):
                 assert torch.all(self.layer.weight.grad == 0)
 
             self.manual_backward(loss_1, opt_a)
-            self.manual_optimizer_step(opt_a)
+            opt_a.step()
 
             # fake discriminator
             loss_2 = self(x)
@@ -586,7 +588,7 @@ def training_step(self, batch, batch_idx, optimizer_idx):
             self.manual_backward(loss_2, opt_a, retain_graph=True)
 
             assert self.layer.weight.grad is not None
-            self.manual_optimizer_step(opt_b)
+            opt_b.step()
 
         def training_epoch_end(self, outputs) -> None:
             # outputs should be an array with an entry per optimizer
@@ -611,7 +613,8 @@ def configure_optimizers(self):
         weights_summary=None,
         precision=16,
         amp_backend='native',
-        gpus=1
+        gpus=1,
+        enable_pl_optimizer=True,
     )
 
     trainer.fit(model)
@@ -620,9 +623,9 @@ def configure_optimizers(self):
     assert trainer.dev_debugger.count_events('backward_call') == limit_train_batches * num_manual_backward_calls
 
 
-def test_manual_optimizer_step_with_optimizer_closure(tmpdir):
+def test_step_with_optimizer_closure(tmpdir):
     """
-    Tests that `manual_optimizer_step` works with optimizer_closure
+    Tests that `step` works with optimizer_closure
     """
     os.environ['PL_DEV_DEBUG'] = '1'
 
@@ -664,7 +667,7 @@ def optimizer_closure():
 
             weight_before = self.layer.weight.clone()
 
-            self.manual_optimizer_step(opt, optimizer_closure=optimizer_closure)
+            opt.step(closure=optimizer_closure)
 
             weight_after = self.layer.weight.clone()
             assert not torch.equal(weight_before, weight_after)
@@ -689,6 +692,7 @@ def configure_optimizers(self):
         limit_val_batches=2,
         max_epochs=1,
         log_every_n_steps=1,
+        enable_pl_optimizer=True,
     )
 
     trainer.fit(model)
@@ -697,9 +701,9 @@ def configure_optimizers(self):
     assert trainer.logger_connector.progress_bar_metrics["train_loss_epoch"] == torch.stack(model._losses).mean()
 
 
-def test_manual_optimizer_step_with_optimizer_closure_and_accumulated_grad(tmpdir):
+def test_step_with_optimizer_closure_and_accumulated_grad(tmpdir):
     """
-    Tests that `manual_optimizer_step` works with optimizer_closure and accumulated_grad
+    Tests that `step` works with optimizer_closure and accumulated_grad
     """
     os.environ['PL_DEV_DEBUG'] = '1'
 
@@ -721,7 +725,7 @@ def optimizer_closure():
 
             weight_before = self.layer.weight.clone()
 
-            self.manual_optimizer_step(opt, optimizer_closure=optimizer_closure)
+            opt.step(closure=optimizer_closure)
 
             weight_after = self.layer.weight.clone()
             if not self.trainer.train_loop.should_accumulate():
@@ -750,6 +754,7 @@ def configure_optimizers(self):
         max_epochs=1,
         log_every_n_steps=1,
         accumulate_grad_batches=2,
+        enable_pl_optimizer=True,
     )
 
     trainer.fit(model)
@@ -757,9 +762,9 @@ def configure_optimizers(self):
 
 
 @patch("torch.optim.SGD.step")
-def test_manual_optimizer_step_with_optimizer_closure_and_extra_arguments(step_mock, tmpdir):
+def test_step_with_optimizer_closure_and_extra_arguments(step_mock, tmpdir):
     """
-    Tests that `manual_optimizer_step` works with optimizer_closure and extra arguments
+    Tests that `step` works with optimizer_closure and extra arguments
     """
     os.environ['PL_DEV_DEBUG'] = '1'
 
@@ -779,7 +784,7 @@ def optimizer_closure():
                     retain_graph = num_backward != backward_idx # noqa E225
                     self.manual_backward(loss_1, opt, retain_graph=retain_graph)
 
-            self.manual_optimizer_step(opt, 1, optimizer_closure=optimizer_closure, something="new")
+            opt.step(1, closure=optimizer_closure, something="new")
 
         def training_epoch_end(self, outputs) -> None:
             # outputs should be an array with an entry per optimizer
@@ -802,6 +807,7 @@ def configure_optimizers(self):
         max_epochs=1,
         log_every_n_steps=1,
         accumulate_grad_batches=2,
+        enable_pl_optimizer=True,
     )
 
     trainer.fit(model)
@@ -811,9 +817,9 @@ def configure_optimizers(self):
 
 @patch("torch.optim.Adam.step")
 @patch("torch.optim.SGD.step")
-def test_manual_optimizer_step_with_optimizer_closure_with_different_frequencies(mock_sgd_step, mock_adam_step, tmpdir):
+def test_step_with_optimizer_closure_with_different_frequencies(mock_sgd_step, mock_adam_step, tmpdir):
     """
-    Tests that `manual_optimizer_step` works with optimizer_closure and different accumulated_gradient frequency
+    Tests that `step` works with optimizer_closure and different accumulated_gradient frequency
     """
     os.environ['PL_DEV_DEBUG'] = '1'
 
@@ -845,22 +851,14 @@ def dis_closure():
                 self.manual_backward(loss_dis, opt_dis)
 
             # this will accumulate gradients for 2 batches and then call opt_gen.step()
-            self.manual_optimizer_step(
-                opt_gen,
-                optimizer_closure=gen_closure,
-                make_optimizer_step=batch_idx % 2 == 0,
-                optim='sgd')
+            opt_gen.step(closure=gen_closure, make_optimizer_step=batch_idx % 2 == 0, optim='sgd')
 
             # update discriminator every 4 baches
             # therefore, no gradient accumulation for discriminator
             if batch_idx % 4 == 0 :
                 # Note: Set make_optimizer_step to True or it will use by default
                 # Trainer(accumulate_grad_batches=x)
-                self.manual_optimizer_step(
-                    opt_dis,
-                    optimizer_closure=dis_closure,
-                    make_optimizer_step=True,
-                    optim='adam')
+                opt_dis.step(closure=dis_closure, make_optimizer_step=True, optim='adam')
 
         def training_epoch_end(self, outputs) -> None:
             # outputs should be an array with an entry per optimizer
@@ -884,6 +882,7 @@ def configure_optimizers(self):
         max_epochs=1,
         log_every_n_steps=1,
         accumulate_grad_batches=2,
+        enable_pl_optimizer=True,
     )
 
     trainer.fit(model)
diff --git a/tests/trainer/test_optimizers.py b/tests/trainer/test_optimizers.py
index c28e626f2eec0..2e76192836740 100644
--- a/tests/trainer/test_optimizers.py
+++ b/tests/trainer/test_optimizers.py
@@ -14,7 +14,8 @@
 import pytest
 import torch
 
-from pytorch_lightning import Trainer, Callback
+from pytorch_lightning import Callback, Trainer
+from pytorch_lightning.core.optimizer import LightningOptimizer
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
 from tests.base import EvalModelTemplate
 from tests.base.boring_model import BoringModel
@@ -179,8 +180,9 @@ def test_reducelronplateau_scheduling(tmpdir):
     ), 'lr scheduler was not correctly converted to dict'
 
 
-def test_optimizer_return_options():
-    trainer = Trainer()
+@pytest.mark.parametrize("enable_pl_optimizer", [False, True])
+def test_optimizer_return_options(enable_pl_optimizer):
+    trainer = Trainer(enable_pl_optimizer=enable_pl_optimizer)
     model = EvalModelTemplate()
 
     # single optimizer
diff --git a/tests/trainer/test_trainer.py b/tests/trainer/test_trainer.py
index 75c3466bee391..085d361952844 100644
--- a/tests/trainer/test_trainer.py
+++ b/tests/trainer/test_trainer.py
@@ -20,7 +20,7 @@
 from argparse import Namespace
 from copy import deepcopy
 from pathlib import Path
-from unittest.mock import patch, call, ANY
+from unittest.mock import ANY, call, patch
 
 import cloudpickle
 import pytest
@@ -34,10 +34,10 @@
 from pytorch_lightning.loggers import TensorBoardLogger
 from pytorch_lightning.profiler.profilers import AdvancedProfiler, PassThroughProfiler, SimpleProfiler
 from pytorch_lightning.trainer.logging import TrainerLoggingMixin
+from pytorch_lightning.utilities import NATIVE_AMP_AVAILABLE
 from pytorch_lightning.utilities.cloud_io import load as pl_load
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
-from pytorch_lightning.utilities import NATIVE_AMP_AVAILABLE
-from tests.base import EvalModelTemplate, BoringModel
+from tests.base import BoringModel, EvalModelTemplate
 
 
 @pytest.mark.parametrize("url_ckpt", [True, False])
@@ -500,8 +500,9 @@ def test_model_freeze_unfreeze():
     model.unfreeze()
 
 
+@pytest.mark.parametrize("enable_pl_optimizer", [False, True])
 @pytest.mark.parametrize("url_ckpt", [True, False])
-def test_resume_from_checkpoint_epoch_restored(monkeypatch, tmpdir, tmpdir_server, url_ckpt):
+def test_resume_from_checkpoint_epoch_restored(monkeypatch, tmpdir, tmpdir_server, url_ckpt, enable_pl_optimizer):
     """Verify resuming from checkpoint runs the right number of epochs"""
     # set $TORCH_HOME, which determines torch hub's cache path, to tmpdir
     monkeypatch.setenv("TORCH_HOME", tmpdir)
@@ -541,6 +542,7 @@ def increment_on_load_checkpoint(self, _):
         checkpoint_callback=ModelCheckpoint(dirpath=tmpdir, monitor='early_stop_on', save_top_k=-1),
         default_root_dir=tmpdir,
         val_check_interval=1.0,
+        enable_pl_optimizer=enable_pl_optimizer,
     )
 
     trainer = Trainer(**trainer_options)


Code review comments:
diff_hunk: @@ -1316,7 +1171,7 @@ def optimizer_step(
         By default, Lightning calls ``step()`` and ``zero_grad()`` as shown in the example
         once per optimizer.
 
-        .. tip:: Consider using `manual_optimizer_step` instead of overriding this method as done previously.
+        .. tip:: Consider using `.step()` instead of overriding this method as done previously.
path: pytorch_lightning/core/lightning.py
text: @author:
```suggestion
        .. tip:: With `Trainer(enable_pl_optimizer=True)`, you can user `optimizer.step()` directly and it will handle zero_grad, accumulated gradients, AMP, TPU and more automatically for you.
```