#3325

pytorch_lightning/strategies/ddp.py

@@ -27,6 +27,7 @@
 import numpy as np
 import torch
 import torch.distributed
+from torch.distributed.algorithms.join import Join
 from torch.nn import Module
 from torch.nn.parallel.distributed import DistributedDataParallel
 
@@ -88,6 +89,7 @@ def __init__(
         ddp_comm_hook: Optional[callable] = None,
         ddp_comm_wrapper: Optional[callable] = None,
         model_averaging_period: Optional[int] = None,
+        uneven_inputs_support: bool = False,
         **kwargs: Union[Any, Dict[str, Any]],
     ) -> None:
         super().__init__(
@@ -106,6 +108,7 @@ def __init__(
         self._ddp_comm_hook = ddp_comm_hook
         self._ddp_comm_wrapper = ddp_comm_wrapper
         self._model_averaging_period = model_averaging_period
+        self._uneven_inputs_support = uneven_inputs_support
         self._pids: Optional[List[int]] = None
         self._sync_dir: Optional[str] = None
         self._rank_0_has_called_call_children_scripts: bool = False
@@ -142,6 +145,14 @@ def distributed_sampler_kwargs(self):
     def _is_single_process_single_device(self) -> bool:
         return True
 
+    @property
+    def uneven_inputs_support(self) -> bool:
+        return self._uneven_inputs_support
+
+    @uneven_inputs_support.setter
+    def uneven_inputs_support(self, uneven_inputs_support: bool) -> None:
+        self._uneven_inputs_support = uneven_inputs_support
+
     def setup_environment(self) -> None:
         # start the other scripts
         if not self.cluster_environment.creates_processes_externally:
@@ -397,6 +408,10 @@ def reduce(self, tensor, group: Optional[Any] = None, reduce_op: Union[ReduceOp,
 
     def training_step(self, *args, **kwargs) -> STEP_OUTPUT:
         with self.precision_plugin.train_step_context():
+            # TODO: Currently a placeholder, implement Joinable and custom join hooks
+            if self.uneven_inputs_support:
+                with Join([self.model]):
+                    return self.model(*args, **kwargs)
             return self.model(*args, **kwargs)
 
     def validation_step(self, *args, **kwargs) -> Optional[STEP_OUTPUT]:
@@ -428,6 +443,12 @@ def register_strategies(cls, strategy_registry: Dict) -> None:
             description="DDP Strategy with `find_unused_parameters` as False",
             find_unused_parameters=False,
         )
+        strategy_registry.register(
+            "ddp_uneven_inputs_support",
+            cls,
+            description="DDP Strategy with `uneven_inputs_support` as True",
+            uneven_inputs_support=True,
+        )
 
     def _should_run_deadlock_detection(self) -> bool:
         """Determines whether the plugin will perform process reconciliation in case of errors.

pytorch_lightning/strategies/ddp_spawn.py

@@ -21,6 +21,7 @@
 import torch
 import torch.distributed
 import torch.multiprocessing as mp
+from torch.distributed.algorithms.join import Join
 from torch.nn import Module
 from torch.nn.parallel.distributed import DistributedDataParallel
 
@@ -65,6 +66,7 @@ def __init__(
         ddp_comm_state: Optional[object] = None,
         ddp_comm_hook: Optional[callable] = None,
         ddp_comm_wrapper: Optional[callable] = None,
+        uneven_inputs_support: bool = False,
         **kwargs: Any,
     ):
         super().__init__(
@@ -80,6 +82,7 @@ def __init__(
         self._ddp_comm_state = ddp_comm_state
         self._ddp_comm_hook = ddp_comm_hook
         self._ddp_comm_wrapper = ddp_comm_wrapper
+        self._uneven_inputs_support = uneven_inputs_support
         self._local_rank = 0
         self.set_world_ranks()
 
@@ -114,6 +117,14 @@ def distributed_sampler_kwargs(self):
     def _is_single_process_single_device(self):
         return True
 
+    @property
+    def uneven_inputs_support(self) -> bool:
+        return self._uneven_inputs_support
+
+    @uneven_inputs_support.setter
+    def uneven_inputs_support(self, uneven_inputs_support: bool) -> None:
+        self._uneven_inputs_support = uneven_inputs_support
+
     def setup(self, trainer: "pl.Trainer") -> None:
         os.environ["MASTER_PORT"] = str(self.cluster_environment.main_port)
         super().setup(trainer)
@@ -311,6 +322,10 @@ def reduce(self, tensor, group: Optional[Any] = None, reduce_op: Union[ReduceOp,
 
     def training_step(self, *args, **kwargs) -> STEP_OUTPUT:
         with self.precision_plugin.train_step_context():
+            # TODO: Currently a placeholder, implement Joinable and custom join hooks
+            if self.uneven_inputs_support:
+                with Join([self.model]):
+                    return self.model(*args, **kwargs)
             return self.model(*args, **kwargs)
 
     def validation_step(self, *args, **kwargs) -> Optional[STEP_OUTPUT]:
@@ -367,6 +382,12 @@ def register_strategies(cls, strategy_registry: Dict) -> None:
             description="DDPSpawn Strategy with `find_unused_parameters` as False",
             find_unused_parameters=False,
         )
+        strategy_registry.register(
+            "ddp_spawn_uneven_inputs_support",
+            cls,
+            description="DDP Spawn Strategy with `uneven_inputs_support` as True",
+            uneven_inputs_support=True,
+        )
 
     def teardown(self) -> None:
         super().teardown()

pytorch_lightning/trainer/connectors/data_connector.py

@@ -409,8 +409,8 @@ def _resolve_sampler(self, dataloader: DataLoader, shuffle: bool, mode: Optional
 
         return dataloader.sampler
 
-    @staticmethod
     def _get_distributed_sampler(
+        self,
         dataloader: DataLoader,
         shuffle: bool,
         overfit_batches: Union[int, float],
@@ -420,6 +420,15 @@ def _get_distributed_sampler(
         """This function is used to created the distributed sampler injected within the user DataLoader."""
         kwargs["shuffle"] = shuffle and not overfit_batches
         kwargs.setdefault("seed", int(os.getenv("PL_GLOBAL_SEED", 0)))
+        if getattr(self.trainer.strategy, "uneven_inputs_support", False) and mode == RunningStage.TRAINING:
+            if len(dataloader.dataset) % kwargs["num_replicas"] != 0:
+                return UnrepeatedDistributedSampler(dataloader.dataset, **kwargs)
+            rank_zero_warn(
+                f"You have passed `uneven_inputs_support=True` for the {self.trainer.strategy.name} strategy. "
+                "But as the dataset length is evenly divisible by number of replicas, then there "
+                "is no need to support uneven inputs, since the dataset will be split equally."
+            )
+            self.trainer.strategy.uneven_inputs_support = False
         cls = UnrepeatedDistributedSampler if mode == RunningStage.PREDICTING else DistributedSampler
         sampler = cls(dataloader.dataset, **kwargs)
         return sampler