Use DistributedSampler when running with custom accelerator (#7814)

Co-authored-by: Jirka Borovec <[email protected]> Co-authored-by: Adrian Wälchli <[email protected]>
Lightning-AI · Jun 22, 2021 · 7c58c4d · 7c58c4d
1 parent dc44a53
commit 7c58c4d
Show file tree

Hide file tree

Showing 3 changed files with 30 additions and 2 deletions.
diff --git a/pytorch_lightning/plugins/training_type/ddp.py b/pytorch_lightning/plugins/training_type/ddp.py
@@ -84,6 +84,10 @@ def __init__(
         self._ddp_comm_wrapper = ddp_comm_wrapper
         self.set_world_ranks()
 
+    @property
+    def is_distributed(self) -> bool:
+        return True
+
     @property
     def root_device(self):
         return self.parallel_devices[self.local_rank]

diff --git a/pytorch_lightning/trainer/connectors/accelerator_connector.py b/pytorch_lightning/trainer/connectors/accelerator_connector.py
@@ -382,7 +382,11 @@ def select_precision_plugin(self) -> PrecisionPlugin:
         raise NotImplementedError("We only support precisions 64, 32 and 16!")
 
     def select_training_type_plugin(self) -> TrainingTypePlugin:
-        if self.use_ddp2:
+        if isinstance(
+            self.distributed_backend, Accelerator
+        ) and self.distributed_backend.training_type_plugin is not None:
+            plugin = self.distributed_backend.training_type_plugin
+        elif self.use_ddp2:
             plugin = DDP2Plugin(
                 parallel_devices=self.parallel_devices,
                 num_nodes=self.num_nodes,

diff --git a/tests/accelerators/test_accelerator_connector.py b/tests/accelerators/test_accelerator_connector.py
@@ -383,8 +383,9 @@ class Prec(PrecisionPlugin):
     class TrainTypePlugin(SingleDevicePlugin):
         pass
 
+    ttp = TrainTypePlugin(device=torch.device("cpu"))
     accelerator = Accel(
-        training_type_plugin=TrainTypePlugin(device=torch.device("cpu")),
+        training_type_plugin=ttp,
         precision_plugin=Prec(),
     )
     trainer = Trainer(
@@ -395,6 +396,25 @@ class TrainTypePlugin(SingleDevicePlugin):
     assert isinstance(trainer.accelerator, Accel)
     assert isinstance(trainer.training_type_plugin, TrainTypePlugin)
     assert isinstance(trainer.precision_plugin, Prec)
+    assert trainer.accelerator_connector.training_type_plugin is ttp
+
+    class DistributedPlugin(DDPPlugin):
+        pass
+
+    ttp = DistributedPlugin()
+    accelerator = Accel(
+        training_type_plugin=ttp,
+        precision_plugin=Prec(),
+    )
+    trainer = Trainer(
+        accelerator=accelerator,
+        fast_dev_run=True,
+        num_processes=2,
+    )
+    assert isinstance(trainer.accelerator, Accel)
+    assert isinstance(trainer.training_type_plugin, DistributedPlugin)
+    assert isinstance(trainer.precision_plugin, Prec)
+    assert trainer.accelerator_connector.training_type_plugin is ttp
 
 
 @mock.patch.dict(