Use DistributedSampler when running with custom accelerator

pytorch_lightning/plugins/training_type/ddp.py

@@ -95,6 +95,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) -> torch.device:
         return self.parallel_devices[self.local_rank]

pytorch_lightning/trainer/connectors/accelerator_connector.py

@@ -403,7 +403,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,
                 cluster_environment=self.cluster_environment,

tests/trainer/connectors/test_accelerator_connector.py

@@ -0,0 +1,30 @@
+import torch
+
+from pytorch_lightning import Trainer
+from pytorch_lightning.accelerators import Accelerator
+from pytorch_lightning.plugins import DDPPlugin, PrecisionPlugin, SingleDevicePlugin
+
+
+def test_accelerator_training_type_plugin():
+    """Test that training_type_plugin is pulled from accelearator"""
+
+    # check that this works for different types of plugins to ensure
+    # there are no dependencies on TrainingTypePlugin class refinements
+
+    precision_plugin = PrecisionPlugin()
+    singledev_plugin = SingleDevicePlugin(torch.device('cpu'))
+    accelerator = Accelerator(
+        precision_plugin=precision_plugin,
+        training_type_plugin=singledev_plugin,
+    )
+    trainer = Trainer(accelerator=accelerator)
+    assert trainer.accelerator_connector.training_type_plugin is singledev_plugin
+
+    precision_plugin = PrecisionPlugin()
+    ddp_plugin = DDPPlugin()
+    accelerator = Accelerator(
+        precision_plugin=precision_plugin,
+        training_type_plugin=ddp_plugin,
+    )
+    trainer = Trainer(accelerator=accelerator)
+    assert trainer.accelerator_connector.training_type_plugin is ddp_plugin