[WIP] ref: decoupled ddp, ddp spawn

benchmarks/test_parity.py

@@ -11,7 +11,7 @@
 
 @pytest.mark.parametrize('cls_model,max_diff', [
     (ParityModuleRNN, 0.05),
-    (ParityModuleMNIST, 0.5)
+    (ParityModuleMNIST, 0.55)
 ])
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="test requires GPU machine")
 def test_pytorch_parity(tmpdir, cls_model, max_diff):

pl_examples/basic_examples/autoencoder.py

@@ -97,7 +97,8 @@ def cli_main():
     # ------------
     # testing
     # ------------
-    trainer.test(test_dataloaders=test_loader)
+    result = trainer.test(test_dataloaders=test_loader)
+    print(result)
 
 
 if __name__ == '__main__':

pytorch_lightning/accelerators/accelerator_connector.py

@@ -81,12 +81,11 @@ def on_trainer_init(
 
         # override with environment flag
         gpus = os.environ.get('PL_TRAINER_GPUS', gpus)
+        self.trainer.gpus = gpus
 
         # for gpus allow int, string and gpu list
         if auto_select_gpus and isinstance(gpus, int):
             self.trainer.gpus = self.trainer.tuner.pick_multiple_gpus(gpus)
-        else:
-            self.trainer.gpus = gpus
 
         self.trainer.data_parallel_device_ids = device_parser.parse_gpu_ids(self.trainer.gpus)
         self.trainer.root_gpu = device_parser.determine_root_gpu_device(self.trainer.data_parallel_device_ids)
@@ -126,6 +125,9 @@ def on_trainer_init(
         self.trainer.replace_sampler_ddp = replace_sampler_ddp
 
     def select_accelerator(self):
+        if self.trainer.accelerator_backend is not None:
+            return self.trainer.accelerator_backend
+
         # SLURM ddp
         use_slurm_ddp = self.trainer.use_ddp and self.trainer.is_slurm_managing_tasks
 
@@ -294,7 +296,8 @@ def set_nvidia_flags(self, is_slurm_managing_tasks, data_parallel_device_ids):
                 os.environ["CUDA_VISIBLE_DEVICES"] = gpu_str
 
         # don't make this debug... this is good UX
-        rank_zero_info(f'CUDA_VISIBLE_DEVICES: [{os.environ["CUDA_VISIBLE_DEVICES"]}]')
+        devices = os.environ["CUDA_VISIBLE_DEVICES"]
+        log.info(f'LOCAL_RANK: {self.trainer.local_rank} - CUDA_VISIBLE_DEVICES: [{devices}]')
 
     def determine_local_rank(self):
         if self.trainer.is_slurm_managing_tasks:

pytorch_lightning/accelerators/base_backend.py

@@ -170,6 +170,9 @@ def _clip_gradients(self, optimizer):
     def on_train_epoch_end(self):
         pass
 
+    def on_train_end(self):
+        pass
+
     def early_stopping_should_stop(self, pl_module):
         return self.trainer.should_stop
 

pytorch_lightning/accelerators/ddp_backend.py

@@ -13,17 +13,22 @@
 # 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 Optional
 
 import numpy as np
-import torch
 
+from pytorch_lightning import _logger as log
 from pytorch_lightning.utilities.distributed import find_free_network_port
-from pytorch_lightning.accelerators.ddp_base_backend import DDPBase
+from pytorch_lightning.accelerators.base_backend import Accelerator
+from pytorch_lightning.utilities.distributed import rank_zero_only
+from pytorch_lightning.utilities import AMPType
+
 
 try:
     from hydra.utils import to_absolute_path, get_original_cwd
@@ -34,13 +39,14 @@
     HYDRA_AVAILABLE = True
 
 
-class DDPBackend(DDPBase):
+class DDPBackend(Accelerator):
 
     def __init__(self, trainer, mode: str = 'ddp'):
         super().__init__(trainer)
         self.task_idx = None
         self._has_spawned_children = False
         self.mode = mode
+        self.interactive_ddp_procs = []
 
     def setup(self, model):
         if self.mode == 'ddp':
@@ -59,6 +65,10 @@ def __torchelastic_setup(self):
         self.task_idx = int(os.environ['LOCAL_RANK'])
 
     def __ddp_script_mode_setup(self):
+        # do nothing when already in a ddp subprocess
+        if os.environ.get('PL_IN_DDP_SUBPROCESS', '0') == '1':
+            return
+
         assert self.trainer.global_rank == 0
         self._check_can_spawn_children()
         self._has_spawned_children = True
@@ -91,21 +101,27 @@ def __ddp_script_mode_setup(self):
         # when the trainer script was called the device has already been scoped by the time
         # code reaches this point. so, to call the scripts, we need to leave cuda visible devices alone
         # but forward the GPUs selected via environment variables
+        # set the flag for ddp scripts
+
+        os.environ['PL_TRAINER_GPUS'] = ','.join([str(i) for i in self.trainer.data_parallel_device_ids])
+        os.environ['PL_IN_DDP_SUBPROCESS'] = '1'
+
+        if self.trainer.logger is not None:
+            os.environ['PL_EXP_VERSION'] = str(self.trainer.logger.version)
+
         gpu_ids = os.environ.get('CUDA_VISIBLE_DEVICES', '')
         if len(gpu_ids) == 1:
             gpu_ids = f'{gpu_ids},'
 
         num_gpus = max(1, len(gpu_ids.split(',')))
 
-        # set the flag for ddp scripts
-        os.environ['PL_TRAINER_GPUS'] = gpu_ids
-
         os.environ['WORLD_SIZE'] = f'{num_gpus * self.trainer.num_nodes}'
 
-        self.trainer.interactive_ddp_procs = []
+        self.interactive_ddp_procs = []
         for local_rank in range(1, self.trainer.num_processes):
             env_copy = os.environ.copy()
             env_copy['LOCAL_RANK'] = f'{local_rank}'
+            env_copy['PL_DDP_PID'] = str(self.trainer.data_parallel_device_ids[local_rank])
 
             # start process
             # if hydra is available and initialized, make sure to set the cwd correctly
@@ -114,7 +130,7 @@ def __ddp_script_mode_setup(self):
                 if HydraConfig.initialized():
                     cwd = get_original_cwd()
             proc = subprocess.Popen(command, env=env_copy, cwd=cwd)
-            self.trainer.interactive_ddp_procs.append(proc)
+            self.interactive_ddp_procs.append(proc)
 
             # starting all processes at once can cause issues
             # with dataloaders delay between 1-10 seconds
@@ -123,14 +139,116 @@ def __ddp_script_mode_setup(self):
 
         self.task_idx = 0
 
+        # wait for all the procs to start
+        sleep(2)
+
     def train(self):
         model = self.trainer.model
         if self.mode == 'ddp':
-            results = self.ddp_train_tmp(process_idx=self.task_idx, mp_queue=None, model=model, is_master=True)
-            del os.environ['WORLD_SIZE']
+            results = self.ddp_train(process_idx=self.task_idx, mp_queue=None, model=model, is_master=True)
+            if 'WORLD_SIZE' in os.environ:
+                del os.environ['WORLD_SIZE']
             return results
         else:
-            self.ddp_train_tmp(process_idx=self.task_idx, mp_queue=None, model=model)
+            return self.ddp_train(process_idx=self.task_idx, mp_queue=None, model=model)
+
+    def ddp_train(self, process_idx, mp_queue, model, is_master=False, proc_offset=0):
+        """
+        Entry point for ddp
+        Args:
+            process_idx:
+            mp_queue: multiprocessing queue
+            model:
+            is_master:
+            proc_offset:
+        Returns:
+        """
+        # offset the process id if requested
+        process_idx = process_idx + proc_offset
+
+        # show progressbar only on progress_rank 0
+        if (self.trainer.node_rank != 0 or process_idx != 0) and self.trainer.progress_bar_callback is not None:
+            self.trainer.progress_bar_callback.disable()
+
+        # determine which process we are and world size
+        self.set_world_ranks(process_idx)
+
+        # set warning rank
+        rank_zero_only.rank = self.trainer.global_rank
+
+        # set up server using proc 0's ip address
+        # try to init for 20 times at max in case ports are taken
+        # where to store ip_table
+        model.trainer = self.trainer
+        model.init_ddp_connection(
+            self.trainer.global_rank,
+            self.trainer.world_size,
+            self.trainer.is_slurm_managing_tasks
+        )
+
+        # call setup after the ddp process has connected
+        self.trainer.call_setup_hook(model)
+
+        # on world_size=0 let everyone know training is starting
+        if self.trainer.is_global_zero and not torch.distributed.is_initialized():
+            log.info('-' * 100)
+            log.info(f'distributed_backend={self.trainer.distributed_backend}')
+            log.info(f'All DDP processes registered. Starting ddp with {self.trainer.world_size} processes')
+            log.info('-' * 100)
+
+            # call sync_bn before .cuda(), configure_apex and configure_ddp
+            if self.trainer.sync_batchnorm:
+                model = model.configure_sync_batchnorm(model)
+
+        # MODEL
+        # copy model to each gpu
+        self.model_to_device(model, process_idx, is_master)
+
+        # CHOOSE OPTIMIZER
+        # allow for lr schedulers as well
+        self.setup_optimizers(model)
+
+        # set model properties before going into wrapper
+        self.trainer.model_connector.copy_trainer_model_properties(model)
+
+        # AMP - run through amp wrapper before going to distributed DP
+        # DDP uses all GPUs on the machine
+        device_ids = self.get_device_ids()
+
+        # allow user to configure ddp
+        model = model.configure_ddp(model, device_ids)
+
+        # set up training routine
+        self.barrier('ddp_setup')
+        self.trainer.train_loop.setup_training(model)
+
+        # train or test
+        results = self.train_or_test()
+
+        # clean up memory
+        torch.cuda.empty_cache()
+
+        return results
+
+    def training_step(self, args):
+        if self.trainer.amp_backend == AMPType.NATIVE:
+            with torch.cuda.amp.autocast():
+                output = self.trainer.model(*args)
+        else:
+            output = self.trainer.model(*args)
+        return output
+
+    def validation_step(self, args):
+        output = self.training_step(args)
+        return output
+
+    def test_step(self, args):
+        output = self.training_step(args)
+        return output
+
+    def barrier(self, name: str = None):
+        if torch_distrib.is_initialized():
+            torch_distrib.barrier()
 
     def _check_can_spawn_children(self):
         if self._has_spawned_children:
@@ -145,15 +263,7 @@ def set_world_ranks(self, process_idx):
         self.trainer.world_size = self.trainer.num_nodes * self.trainer.num_processes
 
     def model_to_device(self, model, process_idx, is_master):
-        gpu_idx = process_idx
-
-        # when using ddp, the master process (proc 0) continues running as the main one
-        # this means that the local rank will always be 0
-        # (even if cuda visible devices has other visible gpus)
-        # this means that the master process needs to pull the 0th visible index as the device number
-        if is_master:
-            available_gpus = os.environ['CUDA_VISIBLE_DEVICES'].split(',')
-            gpu_idx = int(available_gpus[self.trainer.local_rank])
+        gpu_idx = int(os.environ.get('PL_DDP_PID', process_idx))
 
         self.trainer.root_gpu = gpu_idx
         torch.cuda.set_device(self.trainer.root_gpu)
@@ -162,3 +272,6 @@ def model_to_device(self, model, process_idx, is_master):
     def get_device_ids(self):
         device_ids = [self.trainer.root_gpu]
         return device_ids
+
+    def on_train_end(self):
+        pass

pytorch_lightning/accelerators/ddp_base_backend.py

@@ -58,7 +58,8 @@ def test_step(self, args):
         return output
 
     def barrier(self, name: str = None):
-        torch_distrib.barrier()
+        if torch_distrib.is_initialized():
+            torch_distrib.barrier()
 
     def early_stopping_should_stop(self, pl_module):
         stop = torch.tensor(int(self.trainer.should_stop), device=pl_module.device)
@@ -132,7 +133,7 @@ def ddp_train_tmp(self, process_idx, mp_queue, model, is_master=False, proc_offs
         self.trainer.call_setup_hook(model)
 
         # on world_size=0 let everyone know training is starting
-        if self.trainer.is_global_zero:
+        if self.trainer.is_global_zero and not torch.distributed.is_initialized():
             log.info('-' * 100)
             log.info(f'distributed_backend={self.trainer.distributed_backend}')
             log.info(f'All DDP processes registered. Starting ddp with {self.trainer.world_size} processes')

pytorch_lightning/accelerators/ddp_cpu_spawn_backend.py

@@ -95,9 +95,7 @@ def ddp_train(self, process_idx, mp_queue, model):
             self.trainer.progress_bar_callback.disable()
 
         # determine which process we are and world size
-        self.trainer.local_rank = process_idx
-        self.trainer.global_rank = self.trainer.node_rank * self.trainer.num_processes + process_idx
-        self.trainer.world_size = self.trainer.num_nodes * self.trainer.num_processes
+        self.set_world_ranks(process_idx)
 
         # set warning rank
         rank_zero_only.rank = self.trainer.global_rank
@@ -116,7 +114,7 @@ def ddp_train(self, process_idx, mp_queue, model):
         self.trainer.call_setup_hook(model)
 
         # on world_size=0 let everyone know training is starting
-        if self.trainer.is_global_zero:
+        if self.trainer.is_global_zero and not torch.distributed.is_initialized():
             log.info('-' * 100)
             log.info(f'distributed_backend={self.trainer.distributed_backend}')
             log.info(f'All DDP processes registered. Starting ddp with {self.trainer.world_size} processes')
@@ -126,6 +124,9 @@ def ddp_train(self, process_idx, mp_queue, model):
         if self.trainer.sync_batchnorm:
             model = model.configure_sync_batchnorm(model)
 
+        # move the model to the correct device
+        self.model_to_device(model, process_idx)
+
         # CHOOSE OPTIMIZER
         # allow for lr schedulers as well
         self.setup_optimizers(model)
@@ -137,7 +138,7 @@ def ddp_train(self, process_idx, mp_queue, model):
         model = self.trainer.precision_connector.connect(model)
 
         # DDP spawn already spawned off each process... no need to do anything
-        device_ids = None
+        device_ids = self.get_device_ids()
 
         # allow user to configure ddp
         model = model.configure_ddp(model, device_ids)
@@ -174,7 +175,8 @@ def test_step(self, args):
         return output
 
     def barrier(self, name: str = None):
-        torch_distrib.barrier()
+        if torch_distrib.is_initialized():
+            torch_distrib.barrier()
 
     def broadcast(self, obj, src=0):
         return self.dist.broadcast(obj)
@@ -186,6 +188,19 @@ def early_stopping_should_stop(self, pl_module):
         should_stop = stop == self.trainer.world_size
         return should_stop
 
+    def set_world_ranks(self, process_idx):
+        self.trainer.local_rank = process_idx
+        self.trainer.global_rank = self.trainer.node_rank * self.trainer.num_processes + process_idx
+        self.trainer.world_size = self.trainer.num_nodes * self.trainer.num_processes
+
+    def model_to_device(self, model, process_idx):
+        # in ddp cpu we don't actually move models to a device
+        pass
+
+    def get_device_ids(self):
+        device_ids = None
+        return device_ids
+
     def transfer_distrib_spawn_state_on_fit_end(self, model, mp_queue, results):
         # track the best model path
         best_model_path = None