pass in ddp_kwargs (#28)

Author: amogkam

ray_lightning/ray_ddp.py

@@ -1,4 +1,4 @@
-from typing import Callable, Dict, List
+from typing import Callable, Dict, List, Union, Any
 
 import os
 from collections import defaultdict
@@ -65,6 +65,8 @@ class RayPlugin(DDPSpawnPlugin):
             Trainer to a value > 0.
         init_hook (Callable): A function to run on each worker
             upon instantiation.
+        **ddp_kwargs: Additional arguments to pass into
+            ``DistributedDataParallel`` initialization
 
     Example:
 
@@ -88,10 +90,12 @@ def __init__(self,
                  num_workers: int = 1,
                  num_cpus_per_worker: int = 1,
                  use_gpu: bool = False,
-                 init_hook: Callable = None):
+                 init_hook: Callable = None,
+                 **ddp_kwargs: Union[Any, Dict[str, Any]]):
         if not ray.is_initialized():
             ray.init()
-        super().__init__(sync_batchnorm=None, parallel_devices=[])
+        super().__init__(
+            sync_batchnorm=None, parallel_devices=[], **ddp_kwargs)
         self.nickname = "ddp_ray"
         self.num_workers = num_workers
         self.num_cpus_per_worker = num_cpus_per_worker

ray_lightning/tests/test_ddp.py

@@ -130,3 +130,18 @@ def test_early_stop(tmpdir, ray_start_2_cpus):
     trained_model = BoringModel.load_from_checkpoint(
         trainer.checkpoint_callback.best_model_path)
     assert trained_model.val_epoch == 2, trained_model.val_epoch
+
+
+def test_unused_parameters(tmpdir, ray_start_2_cpus):
+    """Tests if find_unused_parameters is properly passed to model."""
+    model = BoringModel()
+    plugin = RayPlugin(
+        num_workers=2, use_gpu=False, find_unused_parameters=False)
+
+    class UnusedParameterCallback(Callback):
+        def on_train_start(self, trainer, pl_module):
+            assert trainer.model.find_unused_parameters is False
+
+    trainer = get_trainer(
+        tmpdir, plugins=[plugin], callbacks=[UnusedParameterCallback()])
+    trainer.fit(model)