Support distributed init outside runtime init

deepspeed/__init__.py

@@ -8,6 +8,7 @@
 from deepspeed.pt.log_utils import logger
 from deepspeed.pt.deepspeed_cuda import DeepSpeedTransformerLayer, DeepSpeedTransformerConfig
 from deepspeed.pt.deepspeed_config import DeepSpeedConfig
+from deepspeed.pt.deepspeed_distributed import init_distributed, mpi_discovery
 
 import deepspeed.pt.deepspeed_checkpointing as checkpointing
 
@@ -62,8 +63,7 @@ def initialize(args,
         mpu: Optional: A model parallelism unit object that implements
             get_{model,data}_parallel_{rank,group,world_size}()
 
-        dist_init_required: Optional: None will auto-initialize torch.distributed if needed,
-            otherwise the user can force it to be initialized or not via boolean.
+        dist_init_required: deprecated argument, torch.distributed will be auto-initialized if needed
 
         collate_fn: Optional: Merges a list of samples to form a
             mini-batch of Tensor(s).  Used when using batched loading from a
@@ -90,14 +90,18 @@ def initialize(args,
             __git_branch__),
     )
 
+    if dist_init_required is not None:
+        logger.warning(
+            "deepspeed.initialize argument of 'dist_init_required' is deprecated and not used, torch.distributed will be auto-initialized if needed."
+        )
+
     engine = DeepSpeedLight(args=args,
                             model=model,
                             optimizer=optimizer,
                             model_parameters=model_parameters,
                             training_data=training_data,
                             lr_scheduler=lr_scheduler,
                             mpu=mpu,
-                            dist_init_required=dist_init_required,
                             collate_fn=collate_fn,
                             config_params=config_params)
 

deepspeed/pt/deepspeed_distributed.py

@@ -0,0 +1,63 @@
+import os
+import torch
+from deepspeed.pt.log_utils import logger
+from deepspeed.pt.deepspeed_constants import TORCH_DISTRIBUTED_DEFAULT_PORT
+
+
+def init_distributed(dist_backend="nccl", auto_mpi_discovery=True):
+    """
+    Initialize torch.distributed backend, potentially performing MPI discovery if needed
+
+    Arguments:
+        dist_backend: torch distributed backend, e.g., nccl, mpi, gloo
+        auto_mpi_discovery: if distributed environment variables are not set, attempt to discover them from MPI
+    """
+    required_env = ["RANK", "WORLD_SIZE", "MASTER_ADDR", "MASTER_PORT", "LOCAL_RANK"]
+    if auto_mpi_discovery and not all(map(lambda v: v in os.environ, required_env)):
+        logger.info(
+            "Not using the DeepSpeed or torch.distributed launchers, attempting to detect MPI environment..."
+        )
+        mpi_discovery()
+
+    if not torch.distributed.is_initialized():
+        logger.info("Initializing torch distributed with backend: {}".format(
+            self.dist_backend))
+        torch.distributed.init_process_group(backend=self.dist_backend)
+
+
+def mpi_discovery():
+    """
+    Discovery MPI environment and map to relevant torch.distributed state
+    """
+    from mpi4py import MPI
+    import subprocess
+
+    comm = MPI.COMM_WORLD
+    rank = comm.Get_rank()
+    world_size = comm.Get_size()
+
+    master_addr = None
+    if rank == 0:
+        hostname_cmd = ["hostname -I"]
+        result = subprocess.check_output(hostname_cmd, shell=True)
+        master_addr = result.decode('utf-8').split()[0]
+    master_addr = comm.bcast(master_addr, root=0)
+
+    # Determine local rank by assuming hostnames are unique
+    proc_name = MPI.Get_processor_name()
+    all_procs = comm.allgather(proc_name)
+    local_rank = sum([i == proc_name for i in all_procs[:rank]])
+
+    os.environ['RANK'] = str(rank)
+    os.environ['WORLD_SIZE'] = str(world_size)
+    os.environ['LOCAL_RANK'] = str(local_rank)
+    os.environ['MASTER_ADDR'] = master_addr
+    os.environ['MASTER_PORT'] = TORCH_DISTRIBUTED_DEFAULT_PORT
+
+    logger.info(
+        "Discovered MPI settings of world_rank={}, local_rank={}, world_size={}, master_addr={}, master_port={}"
+        .format(os.environ['RANK'],
+                os.environ['LOCAL_RANK'],
+                os.environ['WORLD_SIZE'],
+                os.environ['MASTER_ADDR'],
+                os.environ['MASTER_PORT']))

deepspeed/pt/deepspeed_launch.py

@@ -96,7 +96,7 @@ def main():
         current_env["CUDA_VISIBLE_DEVICES"]))
     exclusion_counts_per_node = None
 
-    # set PyTorch distributed related environmental variables
+    # Set torch distributed related environmental variables
     current_env["MASTER_ADDR"] = args.master_addr
     current_env["MASTER_PORT"] = str(args.master_port)
     current_env["WORLD_SIZE"] = str(dist_world_size)
@@ -106,6 +106,7 @@ def main():
         # each process's rank
         dist_rank = global_rank_mapping[local_node][local_rank]
         current_env["RANK"] = str(dist_rank)
+        current_env["LOCAL_RANK"] = str(local_rank)
 
         # spawn the processes
         cmd = [

deepspeed/pt/deepspeed_light.py

@@ -26,12 +26,13 @@
 from deepspeed.pt.deepspeed_dataloader import DeepSpeedDataLoader
 from deepspeed.pt.deepspeed_constants import \
     ROUTE_TRAIN, ROUTE_PREDICT, ROUTE_EVAL, \
-    TORCH_DISTRIBUTED_DEFAULT_PORT, \
     ZERO_OPTIMIZATION_OPTIMIZER_STATES, ZERO_OPTIMIZATION_GRADIENTS
 
 import deepspeed.pt.deepspeed_lr_schedules as lr_schedules
 from deepspeed.pt.deepspeed_csr_tensor import CSRTensor
 
+from deepspeed.pt.deepspeed_distributed import init_distributed
+
 MEMORY_OPT_ALLREDUCE_SIZE = 500000000
 SUMMARY_WRITER_DIR_NAME = "JobId"
 
@@ -102,7 +103,6 @@ def __init__(self,
                  training_data=None,
                  lr_scheduler=None,
                  mpu=None,
-                 dist_init_required=None,
                  collate_fn=None,
                  config_params=None):
         super(DeepSpeedLight, self).__init__()
@@ -121,21 +121,9 @@ def __init__(self,
         self.warn_unscaled_loss = True
         self.config_params = config_params
 
-        if dist_init_required is None:
-            dist_init_required = not dist.is_initialized()
-
-        self._mpi_check(args, dist_init_required)
-
-        self.dist_backend = "nccl"
-        if dist_init_required:
-            if not dist.is_initialized():
-                logger.info("Initializing torch distributed with backend: {}".format(
-                    self.dist_backend))
-                dist.init_process_group(backend=self.dist_backend)
-            else:
-                logger.warning(
-                    "Was given dist_init_required=True but detected that torch"
-                    "distributed was already initialized, cannot initialize twice.")
+        # Initialize torch distributed backend
+        init_distributed()
+        self.local_rank = int(os.environ["LOCAL_RANK"])
 
         self._do_args_sanity_check(args)
         self._configure_with_arguments(args, mpu)
@@ -145,8 +133,6 @@ def __init__(self,
         if self.tensorboard_enabled():
             self.summary_writer = self.get_summary_writer()
 
-        self._init_distributed(dist_init_required)
-
         # Configure distributed model
         self._configure_distributed_model(model)
 
@@ -181,51 +167,13 @@ def __init__(self,
 
         self.save_non_zero_checkpoint = False
         self.save_zero_checkpoint = False
-        self._configure_checkpointing(dist_init_required)
+        self._configure_checkpointing()
 
         if self.global_rank == 0:
             self._config.print('DeepSpeedLight configuration')
             if self.dump_state():
                 print_configuration(self, 'DeepSpeedLight')
 
-    def _mpi_check(self, args, dist_init_required):
-        if hasattr(args, 'deepspeed_mpi') and args.deepspeed_mpi:
-            from mpi4py import MPI
-            import subprocess
-            comm = MPI.COMM_WORLD
-            rank = comm.Get_rank()
-            world_size = comm.Get_size()
-
-            master_addr = None
-            if rank == 0:
-                hostname_cmd = ["hostname -I"]
-                result = subprocess.check_output(hostname_cmd, shell=True)
-                master_addr = result.decode('utf-8').split()[0]
-            master_addr = comm.bcast(master_addr, root=0)
-
-            # Determine local rank by assuming hostnames are unique
-            proc_name = MPI.Get_processor_name()
-            all_procs = comm.allgather(proc_name)
-            local_rank = sum([i == proc_name for i in all_procs[:rank]])
-
-            os.environ['RANK'] = str(rank)
-            os.environ['WORLD_SIZE'] = str(world_size)
-            args.local_rank = local_rank
-            os.environ['MASTER_ADDR'] = master_addr
-            os.environ['MASTER_PORT'] = TORCH_DISTRIBUTED_DEFAULT_PORT
-
-            logger.info(
-                "Discovered MPI settings of world_rank={}, local_rank={}, world_size={}, master_addr={}, master_port={}"
-                .format(os.environ['RANK'],
-                        args.local_rank,
-                        os.environ['WORLD_SIZE'],
-                        os.environ['MASTER_ADDR'],
-                        os.environ['MASTER_PORT']))
-
-            if not dist_init_required and dist.is_initialized():
-                assert dist.get_rank() == rank, "MPI rank {} does not match torch rank {}".format(rank, dist.get_rank())
-                assert dist.get_world_size() == world_size, "MPI world size {} does not match torch world size {}".format(world_size, dist.get_world_size())
-
     def tensorboard_enabled(self):
         return self._config.tensorboard_enabled
 
@@ -360,7 +308,7 @@ def _configure_lr_scheduler(self, client_lr_scheduler):
             self.lr_scheduler = client_lr_scheduler
         logger.info(f'DeepSpeed LR Scheduler = {self.lr_scheduler}')
 
-    def _configure_checkpointing(self, dist_init_required):
+    def _configure_checkpointing(self):
 
         dp_rank = self.global_rank
         if self.mpu:
@@ -393,19 +341,6 @@ def _scheduler_from_config(self, optimizer):
         else:
             return None
 
-    def _init_distributed(self, dist_init_required):
-        if self.local_rank >= 0:
-            torch.cuda.set_device(self.local_rank)
-            self.device = torch.device("cuda", self.local_rank)
-            self.world_size = dist.get_world_size()
-            self.global_rank = dist.get_rank()
-            logger.info("Set device to local rank {} within node.".format(
-                self.local_rank))
-        else:
-            self.world_size = 1
-            self.global_rank = 0
-            self.device = torch.device("cuda")
-
     # Configure based on command line arguments
     def _configure_with_arguments(self, args, mpu):
         self.local_rank = args.local_rank if hasattr(args, 'local_rank') else 0
@@ -450,10 +385,23 @@ def _do_sanity_check(self):
                 'DeepSpeed {} optimizer requires dynamic loss scaling'.format(self.optimizer_name())
 
     def _configure_distributed_model(self, model):
+        if self.local_rank >= 0:
+            torch.cuda.set_device(self.local_rank)
+            self.device = torch.device("cuda", self.local_rank)
+            self.world_size = dist.get_world_size()
+            self.global_rank = dist.get_rank()
+            logger.info("Set device to local rank {} within node.".format(
+                self.local_rank))
+        else:
+            self.world_size = 1
+            self.global_rank = 0
+            self.device = torch.device("cuda")
+
         self.module = model
         if self.fp16_enabled():
             self.module.half()
         self.module.to(self.device)
+
         if self.mpu is None:
             self.data_parallel_group = _initialize_parameter_parallel_groups()
             self.dp_world_size = dist.get_world_size()
@@ -467,11 +415,6 @@ def _configure_distributed_model(self, model):
             if torch.is_tensor(p):
                 dist.broadcast(p, src_rank, group=self.data_parallel_group)
 
-        # TODO: support new AMP optimizer
-        # self.module.half()
-        # self.module.to(self.local_rank)
-        #self.module, self.optimizer = amp.initialize(self.module, self.optimizer, opt_level="O2")
-
     # Configure optimizer
     def _configure_optimizer(self, client_optimizer, model_parameters):
         if client_optimizer is not None:

tests/unit/common.py

@@ -33,6 +33,10 @@ def dist_init(local_rank, num_procs, *func_args, **func_kwargs):
             """Initialize torch.distributed and execute the user function. """
             os.environ['MASTER_ADDR'] = '127.0.0.1'
             os.environ['MASTER_PORT'] = '29500'
+            os.environ['LOCAL_RANK'] = str(local_rank)
+            # multi-node tests are not supported, local_rank == rank
+            os.environ['RANK'] = str(local_rank)
+            os.environ['WORLD_SIZE'] = str(num_procs)
             dist.init_process_group(backend=backend,
                                     init_method='env://',
                                     rank=local_rank,