Repository: Lightning-AI/pytorch-lightning

Issue:
Optimizer initialization with DDP
## ❓ Questions and Help
#### What is your question?
I would have expected optimizers to always be initialized after parameters have been moved to their destination device.
However, some ddp backends such as
[ddp_backend](https://github.com/PyTorchLightning/pytorch-lightning/blob/master/pytorch_lightning/accelerators/ddp_backend.py#L174),             [ddp_spawn_backend](https://github.com/PyTorchLightning/pytorch-lightning/blob/master/pytorch_lightning/accelerators/ddp_spawn_backend.py#L115),            [ddp2_backend](https://github.com/PyTorchLightning/pytorch-lightning/blob/master/pytorch_lightning/accelerators/ddp2_backend.py#L111)
initialize the optimizer with the CPU parameters before moving the model to the GPU while others such as [gpu_backend](https://github.com/PyTorchLightning/pytorch-lightning/blob/master/pytorch_lightning/accelerators/gpu_backend.py#L39) pass the GPU parameters.

I'm currently trying to understand two things: 
(1) Where does the linking from the CPU to GPU parameters happen?
(2) Is it actually necessary to initialize the optimizer before moving to the specific device or could it be done afterwards? (Most tutorials initialize the optimizer *after* placing the parameters on the corresponding device, e.g.  [this one](https://pytorch.org/tutorials/intermediate/ddp_tutorial.html#basic-use-case))

The reason I'm asking is that I did some parameter/gradient bending to be views into other tensors. This does not work with the current implementation as the optimizer keeps the reference to the CPU parameters with these tweaks but works fine when adapting the pytorch_lightning code by moving the optimizer creation after the model has been moved to the correct device.


Discussion:
Hi, could you share what you changed to defer the initialization? 
For ddp_backend, I've moved [lines 174-177](https://github.com/PyTorchLightning/pytorch-lightning/blob/master/pytorch_lightning/accelerators/ddp_backend.py#L174-L177) ( which is mainly `optimizers, lr_schedulers, optimizer_frequencies = self.trainer.init_optimizers(model)`) to line 199, i.e., after `model.cuda(self.trainer.root_gpu)`. Note that so far I have only verified that the parameters receive gradients again, so I don't know if this breaks something I am unaware of.
Looks good, I don't see a problem with this. Would you like to send a draft PR, then we will see if the tests pass? Then we should also check the other backends to make it consistent.

Resolving patch:
diff --git a/pytorch_lightning/accelerators/ddp2_backend.py b/pytorch_lightning/accelerators/ddp2_backend.py
--- a/pytorch_lightning/accelerators/ddp2_backend.py
+++ b/pytorch_lightning/accelerators/ddp2_backend.py
@@ -106,13 +106,6 @@ def ddp_train(self, process_idx, mp_queue, model, is_master=False, proc_offset=0
             log.info(f'All DDP processes registered. Starting ddp with {self.trainer.world_size} processes')
             log.info('-' * 100)
 
-        # CHOOSE OPTIMIZER
-        # allow for lr schedulers as well
-        optimizers, lr_schedulers, optimizer_frequencies = self.trainer.init_optimizers(model)
-        self.trainer.optimizers = optimizers
-        self.trainer.lr_schedulers = lr_schedulers
-        self.trainer.optimizer_frequencies = optimizer_frequencies
-
         # MODEL
         # copy model to each gpu
         if self.trainer.on_gpu:
@@ -130,6 +123,13 @@ def ddp_train(self, process_idx, mp_queue, model, is_master=False, proc_offset=0
             torch.cuda.set_device(self.trainer.root_gpu)
             model.cuda(self.trainer.root_gpu)
 
+        # CHOOSE OPTIMIZER
+        # allow for lr schedulers as well
+        optimizers, lr_schedulers, optimizer_frequencies = self.trainer.init_optimizers(model)
+        self.trainer.optimizers = optimizers
+        self.trainer.lr_schedulers = lr_schedulers
+        self.trainer.optimizer_frequencies = optimizer_frequencies
+
         # set model properties before going into wrapper
         self.trainer.copy_trainer_model_properties(model)
 
diff --git a/pytorch_lightning/accelerators/ddp_backend.py b/pytorch_lightning/accelerators/ddp_backend.py
--- a/pytorch_lightning/accelerators/ddp_backend.py
+++ b/pytorch_lightning/accelerators/ddp_backend.py
@@ -169,13 +169,6 @@ def ddp_train(self, process_idx, mp_queue, model, is_master=False, proc_offset=0
             log.info(f'All DDP processes registered. Starting ddp with {self.trainer.world_size} processes')
             log.info('-' * 100)
 
-        # CHOOSE OPTIMIZER
-        # allow for lr schedulers as well
-        optimizers, lr_schedulers, optimizer_frequencies = self.trainer.init_optimizers(model)
-        self.trainer.optimizers = optimizers
-        self.trainer.lr_schedulers = lr_schedulers
-        self.trainer.optimizer_frequencies = optimizer_frequencies
-
         # call sync_bn before .cuda(), configure_apex and configure_ddp
         if self.trainer.sync_batchnorm:
             model = model.configure_sync_batchnorm(model)
@@ -197,6 +190,13 @@ def ddp_train(self, process_idx, mp_queue, model, is_master=False, proc_offset=0
             torch.cuda.set_device(self.trainer.root_gpu)
             model.cuda(self.trainer.root_gpu)
 
+        # CHOOSE OPTIMIZER
+        # allow for lr schedulers as well
+        optimizers, lr_schedulers, optimizer_frequencies = self.trainer.init_optimizers(model)
+        self.trainer.optimizers = optimizers
+        self.trainer.lr_schedulers = lr_schedulers
+        self.trainer.optimizer_frequencies = optimizer_frequencies
+
         # set model properties before going into wrapper
         self.trainer.copy_trainer_model_properties(model)
 
diff --git a/pytorch_lightning/accelerators/ddp_spawn_backend.py b/pytorch_lightning/accelerators/ddp_spawn_backend.py
--- a/pytorch_lightning/accelerators/ddp_spawn_backend.py
+++ b/pytorch_lightning/accelerators/ddp_spawn_backend.py
@@ -110,13 +110,6 @@ def ddp_train(self, process_idx, mp_queue, model):
             log.info(f'All DDP processes registered. Starting ddp with {self.trainer.world_size} processes')
             log.info('-' * 100)
 
-        # CHOOSE OPTIMIZER
-        # allow for lr schedulers as well
-        optimizers, lr_schedulers, optimizer_frequencies = self.trainer.init_optimizers(model)
-        self.trainer.optimizers = optimizers
-        self.trainer.lr_schedulers = lr_schedulers
-        self.trainer.optimizer_frequencies = optimizer_frequencies
-
         # call sync_bn before .cuda(), configure_apex and configure_ddp
         if self.trainer.sync_batchnorm:
             model = model.configure_sync_batchnorm(model)
@@ -129,6 +122,13 @@ def ddp_train(self, process_idx, mp_queue, model):
             torch.cuda.set_device(self.trainer.root_gpu)
             model.cuda(self.trainer.root_gpu)
 
+        # CHOOSE OPTIMIZER
+        # allow for lr schedulers as well
+        optimizers, lr_schedulers, optimizer_frequencies = self.trainer.init_optimizers(model)
+        self.trainer.optimizers = optimizers
+        self.trainer.lr_schedulers = lr_schedulers
+        self.trainer.optimizer_frequencies = optimizer_frequencies
+
         # set model properties before going into wrapper
         self.trainer.copy_trainer_model_properties(model)
 


Test patch: