dp + manual_optimization is not working on PL 1.2.1
🐛 Bug
dp + manual optimization is not working on PL 1.2.1
I am setting automatic optimization = False in my model and giving Trainer() accelerator='dp'
Expected behavior
Traceback (most recent call last):
  File "/usr/local/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 644, in run_train
    self.train_loop.run_training_epoch()
  File "/usr/local/lib/python3.8/site-packages/pytorch_lightning/trainer/training_loop.py", line 492, in run_training_epoch
    batch_output = self.run_training_batch(batch, batch_idx, dataloader_idx)
  File "/usr/local/lib/python3.8/site-packages/pytorch_lightning/trainer/training_loop.py", line 653, in run_training_batch
    self._curr_step_result = self.training_step(
  File "/usr/local/lib/python3.8/site-packages/pytorch_lightning/trainer/training_loop.py", line 293, in training_step
    training_step_output = self.trainer.accelerator.training_step(args)
  File "/usr/local/lib/python3.8/site-packages/pytorch_lightning/accelerators/accelerator.py", line 157, in training_step
    return self.training_type_plugin.training_step(*args)
  File "/usr/local/lib/python3.8/site-packages/pytorch_lightning/plugins/training_type/dp.py", line 61, in training_step
    return self.model(*args, **kwargs)
  File "/usr/local/lib/python3.8/site-packages/torch/nn/modules/module.py", line 727, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/usr/local/lib/python3.8/site-packages/torch/nn/parallel/data_parallel.py", line 161, in forward
    outputs = self.parallel_apply(replicas, inputs, kwargs)
  File "/usr/local/lib/python3.8/site-packages/torch/nn/parallel/data_parallel.py", line 171, in parallel_apply
    return parallel_apply(replicas, inputs, kwargs, self.device_ids[:len(replicas)])
  File "/usr/local/lib/python3.8/site-packages/torch/nn/parallel/parallel_apply.py", line 86, in parallel_apply
    output.reraise()
  File "/usr/local/lib/python3.8/site-packages/torch/_utils.py", line 428, in reraise
    raise self.exc_type(msg)
KeyError: Caught KeyError in replica 1 on device 1.
Original Traceback (most recent call last):
  File "/usr/local/lib/python3.8/site-packages/torch/nn/parallel/parallel_apply.py", line 61, in _worker
    output = module(*input, **kwargs)
  File "/usr/local/lib/python3.8/site-packages/torch/nn/modules/module.py", line 727, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/usr/local/lib/python3.8/site-packages/pytorch_lightning/overrides/data_parallel.py", line 74, in forward
    output = super().forward(*inputs, **kwargs)
  File "/usr/local/lib/python3.8/site-packages/pytorch_lightning/overrides/base.py", line 48, in forward
    output = self.module.training_step(*inputs, **kwargs)
  File "/workdir/flow-based-test/StableGlowModel.py", line 131, in training_step
    opt.step()
  File "/usr/local/lib/python3.8/site-packages/pytorch_lightning/core/optimizer.py", line 219, in step
    self.__optimizer_step(*args, closure=closure, profiler_name=profiler_name, **kwargs)
  File "/usr/local/lib/python3.8/site-packages/pytorch_lightning/core/optimizer.py", line 135, in __optimizer_step
    trainer.accelerator.optimizer_step(optimizer, self._optimizer_idx, lambda_closure=closure, **kwargs)
  File "/usr/local/lib/python3.8/site-packages/pytorch_lightning/accelerators/accelerator.py", line 278, in optimizer_step
    self.run_optimizer_step(optimizer, opt_idx, lambda_closure, **kwargs)
  File "/usr/local/lib/python3.8/site-packages/pytorch_lightning/accelerators/accelerator.py", line 283, in run_optimizer_step
    self.training_type_plugin.optimizer_step(optimizer, lambda_closure=lambda_closure, **kwargs)
  File "/usr/local/lib/python3.8/site-packages/pytorch_lightning/plugins/training_type/training_type_plugin.py", line 160, in optimizer_step
    optimizer.step(closure=lambda_closure, **kwargs)
  File "/usr/local/lib/python3.8/site-packages/torch/autograd/grad_mode.py", line 26, in decorate_context
    return func(*args, **kwargs)
  File "/usr/local/lib/python3.8/site-packages/torch/optim/adamax.py", line 67, in step
    exp_avg, exp_inf = state['exp_avg'], state['exp_inf']
KeyError: 'exp_avg'

Environment

PyTorch Version (e.g., 1.0): 1.7.1
OS (e.g., Linux): Linux
How you installed PyTorch (conda, pip, source): pip
Build command you used (if compiling from source):
Python version: 3.8.7
CUDA/cuDNN version: 10.2
GPU models and configuration:
Any other relevant information: