Move init_ddp_connection to distributed utilities (#9044)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: kaushikb11

pytorch_lightning/plugins/training_type/ddp.py

@@ -45,7 +45,7 @@
 )
 from pytorch_lightning.utilities.distributed import (
     distributed_available,
-    rank_zero_info,
+    init_ddp_connection,
     rank_zero_only,
     ReduceOp,
     sync_ddp_if_available,
@@ -253,7 +253,7 @@ def setup_distributed(self):
         # 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
-        self.init_ddp_connection()
+        init_ddp_connection(self.cluster_environment, self.torch_distributed_backend)
 
         # set the ranks and devices
         self.dist.rank = self.global_rank
@@ -316,25 +316,6 @@ def determine_ddp_device_ids(self):
             return None
         return [self.root_device.index]
 
-    def init_ddp_connection(self, global_rank: Optional[int] = None, world_size: Optional[int] = None) -> None:
-        global_rank = global_rank if global_rank is not None else self.cluster_environment.global_rank()
-        world_size = world_size if world_size is not None else self.cluster_environment.world_size()
-        os.environ["MASTER_ADDR"] = self.cluster_environment.master_address()
-        os.environ["MASTER_PORT"] = str(self.cluster_environment.master_port())
-        if torch.distributed.is_available() and not torch.distributed.is_initialized():
-            log.info(f"initializing ddp: GLOBAL_RANK: {global_rank}, MEMBER: {global_rank + 1}/{world_size}")
-            torch.distributed.init_process_group(
-                self.torch_distributed_backend, rank=global_rank, world_size=world_size
-            )
-
-            # on rank=0 let everyone know training is starting
-            rank_zero_info(
-                f"{'-' * 100}\n"
-                f"distributed_backend={self.torch_distributed_backend}\n"
-                f"All DDP processes registered. Starting ddp with {self.world_size} processes\n"
-                f"{'-' * 100}\n"
-            )
-
     def pre_dispatch(self):
         # move the model to the correct device
         self.model_to_device()

pytorch_lightning/plugins/training_type/ddp_spawn.py

@@ -40,7 +40,7 @@
 from pytorch_lightning.utilities.cloud_io import load as pl_load
 from pytorch_lightning.utilities.distributed import (
     distributed_available,
-    rank_zero_info,
+    init_ddp_connection,
     rank_zero_only,
     ReduceOp,
     sync_ddp_if_available,
@@ -185,7 +185,7 @@ def new_process(self, process_idx: int, trainer: "pl.Trainer", mp_queue: SimpleQ
         # 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
-        self.init_ddp_connection(self.global_rank, self.world_size)
+        init_ddp_connection(self.cluster_environment, self.torch_distributed_backend, self.global_rank, self.world_size)
 
         # TODO: we moved it to the trainer.fit after calling pre_dispatch
         #   ... need to double check that it is the correct place
@@ -261,27 +261,6 @@ def configure_ddp(self):
         )
         self._register_ddp_hooks()
 
-    def init_ddp_connection(self, global_rank: Optional[int], world_size: Optional[int]) -> None:
-        # TODO: this code is duplicated in DDP and DDPSpawn, make this a function
-        global_rank = global_rank if global_rank is not None else self.cluster_environment.global_rank()
-        world_size = world_size if world_size is not None else self.cluster_environment.world_size()
-        os.environ["MASTER_ADDR"] = self.cluster_environment.master_address()
-        os.environ["MASTER_PORT"] = str(self.cluster_environment.master_port())
-
-        if not torch.distributed.is_initialized():
-            log.info(f"initializing ddp: GLOBAL_RANK: {global_rank}, MEMBER: {global_rank + 1}/{world_size}")
-            torch.distributed.init_process_group(
-                self.torch_distributed_backend, rank=global_rank, world_size=world_size
-            )
-
-            # on rank=0 let everyone know training is starting
-            rank_zero_info(
-                f"{'-' * 100}\n"
-                f"distributed_backend={self.torch_distributed_backend}\n"
-                f"All DDP processes registered. Starting ddp with {self.world_size} processes\n"
-                f"{'-' * 100}\n"
-            )
-
     def determine_ddp_device_ids(self):
         if self.root_device.type == "cpu":
             return None

pytorch_lightning/utilities/distributed.py

@@ -21,6 +21,7 @@
 import torch
 from torch.nn.parallel.distributed import DistributedDataParallel
 
+import pytorch_lightning as pl
 from pytorch_lightning.utilities.imports import _TORCH_GREATER_EQUAL_1_8, _TORCH_GREATER_EQUAL_1_9, _TPU_AVAILABLE
 
 if _TPU_AVAILABLE:
@@ -345,3 +346,40 @@ def register_ddp_comm_hook(
 
 def tpu_distributed() -> bool:
     return _TPU_AVAILABLE and xm.xrt_world_size() > 1
+
+
+def init_ddp_connection(
+    cluster_environment: "pl.plugins.environments.ClusterEnvironment",
+    torch_distributed_backend: str,
+    global_rank: Optional[int] = None,
+    world_size: Optional[int] = None,
+    **kwargs,
+) -> None:
+    """
+    Utility function to initialize DDP connection by setting env variables
+    and initiliazing the distributed process group.
+
+    Args:
+        cluster_environment: ``ClusterEnvironment`` instance
+        torch_distributed_backend: backend to use (includes `nccl` and `gloo`)
+        global_rank: rank of the current process
+        world_size: number of processes in the group
+        kwargs: kwargs for ``init_process_group``
+    """
+    global_rank = global_rank if global_rank is not None else cluster_environment.global_rank()
+    world_size = world_size if world_size is not None else cluster_environment.world_size()
+    os.environ["MASTER_ADDR"] = cluster_environment.master_address()
+    os.environ["MASTER_PORT"] = str(cluster_environment.master_port())
+    if torch.distributed.is_available() and not torch.distributed.is_initialized():
+        log.info(f"initializing ddp: GLOBAL_RANK: {global_rank}, MEMBER: {global_rank + 1}/{world_size}")
+        torch.distributed.init_process_group(
+            torch_distributed_backend, rank=global_rank, world_size=world_size, **kwargs
+        )
+
+        # on rank=0 let everyone know training is starting
+        rank_zero_info(
+            f"{'-' * 100}\n"
+            f"distributed_backend={torch_distributed_backend}\n"
+            f"All DDP processes registered. Starting ddp with {world_size} processes\n"
+            f"{'-' * 100}\n"
+        )