Repository: ray-project/ray_lightning (language: python)

Issue:
unhandled cuda error, NCCL version 2.7.8
Hi, I'm getting the following error when trying to run with `RayPlugin`:
```
E                         File "python/ray/_raylet.pyx", line 501, in ray._raylet.execute_task
E                         File "python/ray/_raylet.pyx", line 451, in ray._raylet.execute_task.function_executor
E                         File "/home/rizhiy/miniconda3/envs/ntf/lib/python3.8/site-packages/ray/_private/function_manager.py", line 563, in actor_method_executor
E                           return method(__ray_actor, *args, **kwargs)
E                         File "/home/rizhiy/miniconda3/envs/ntf/lib/python3.8/site-packages/ray_lightning/ray_ddp.py", line 41, in execute
E                           return fn(*args, **kwargs)
E                         File "/home/rizhiy/miniconda3/envs/ntf/lib/python3.8/site-packages/ray_lightning/ray_ddp.py", line 265, in execute_remote
E                           super(RayPlugin, self).new_process(
E                         File "/home/rizhiy/miniconda3/envs/ntf/lib/python3.8/site-packages/pytorch_lightning/plugins/training_type/ddp_spawn.py", line 168, in new_process
E                           self.configure_ddp()
E                         File "/home/rizhiy/miniconda3/envs/ntf/lib/python3.8/site-packages/pytorch_lightning/plugins/training_type/ddp_spawn.py", line 215, in configure_ddp
E                           self._model = DistributedDataParallel(
E                         File "/home/rizhiy/miniconda3/envs/ntf/lib/python3.8/site-packages/torch/nn/parallel/distributed.py", line 496, in __init__
E                           dist._verify_model_across_ranks(self.process_group, parameters)
E                       RuntimeError: NCCL error in: /pytorch/torch/lib/c10d/ProcessGroupNCCL.cpp:911, unhandled cuda error, NCCL version 2.7.8
E                       ncclUnhandledCudaError: Call to CUDA function failed.
```


Discussion / hints:
Tried with `gloo` backend and it seems to be working fine.
Can you set this initialization hook and then share the full output from the run:

```
def init_hook():
    os.environ["NCCL_DEBUG"] = "INFO"

RayPlugin(init_hook=init_hook, ...)
```
```
tests/training/test_runner.py (pid=3250757) raven:3250757:3253258 [0] NCCL INFO Channel 00/02 :    0   1
(pid=3250757) raven:3250757:3253258 [0] NCCL INFO Channel 01/02 :    0   1
(pid=3250757) raven:3250757:3253258 [0] NCCL INFO threadThresholds 8/8/64 | 16/8/64 | 8/8/64
(pid=3250757) raven:3250757:3253258 [0] NCCL INFO Trees [0] 1/-1/-1->0->-1|-1->0->1/-1/-1 [1] 1/-1/-1->0->-1|-1->0->1/-1/-1
(pid=3250757) raven:3250757:3253258 [0] NCCL INFO Channel 00 : 0[1000] -> 1[2000] via P2P/IPC
(pid=3250757)
(pid=3250757) raven:3250757:3253258 [0] transport/p2p.cc:238 NCCL WARN failed to open CUDA IPC handle : 217 peer access is not supported between these two devices
(pid=3250757) raven:3250757:3253258 [0] NCCL INFO transport.cc:68 -> 1
(pid=3250757) raven:3250757:3253258 [0] NCCL INFO init.cc:766 -> 1
(pid=3250757) raven:3250757:3253258 [0] NCCL INFO init.cc:840 -> 1
(pid=3250757) raven:3250757:3253258 [0] NCCL INFO group.cc:73 -> 1 [Async thread]
(pid=3250782) raven:3250782:3253259 [0] NCCL INFO threadThresholds 8/8/64 | 16/8/64 | 8/8/64
(pid=3250782) raven:3250782:3253259 [0] NCCL INFO Trees [0] -1/-1/-1->1->0|0->1->-1/-1/-1 [1] -1/-1/-1->1->0|0->1->-1/-1/-1
(pid=3250782) raven:3250782:3253259 [0] NCCL INFO Channel 00 : 1[2000] -> 0[1000] via P2P/IPC
(pid=3250782)
(pid=3250782) raven:3250782:3253259 [0] transport/p2p.cc:238 NCCL WARN failed to open CUDA IPC handle : 217 peer access is not supported between these two devices
(pid=3250782) raven:3250782:3253259 [0] NCCL INFO transport.cc:68 -> 1
(pid=3250782) raven:3250782:3253259 [0] NCCL INFO init.cc:766 -> 1
(pid=3250782) raven:3250782:3253259 [0] NCCL INFO init.cc:840 -> 1
(pid=3250782) raven:3250782:3253259 [0] NCCL INFO group.cc:73 -> 1 [Async thread]
```
I guess it is similar to #24 
Can you try disabling CUDA_VISIBLE_DEVICES?

Like `os.environ.pop("CUDA_VISIBLE_DEVICES")` in init_hook?

Fix patch:
diff --git a/README.md b/README.md
index 8ff09056..fe1f2e98 100644
--- a/README.md
+++ b/README.md
@@ -34,7 +34,7 @@ Here are the supported PyTorch Lightning versions:
 | Ray Lightning | PyTorch Lightning |
 |---|---|
 | 0.1 | 1.4 |
-| master | 1.5.0-1.5.8 |
+| master | 1.5 |
 
 
 ## PyTorch Distributed Data Parallel Plugin on Ray
diff --git a/ray_lightning/ray_ddp.py b/ray_lightning/ray_ddp.py
index 25affd53..7f4681ce 100644
--- a/ray_lightning/ray_ddp.py
+++ b/ray_lightning/ray_ddp.py
@@ -54,6 +54,9 @@ def get_node_ip(self):
         """Returns the IP address of the node that this Ray actor is on."""
         return ray.util.get_node_ip_address()
 
+    def get_node_and_gpu_ids(self):
+        return ray.get_runtime_context().node_id.hex(), ray.get_gpu_ids()
+
     def execute(self, fn: Callable, *args, **kwargs):
         """Execute the provided function and return the result."""
         return fn(*args, **kwargs)
@@ -132,6 +135,7 @@ def __init__(self,
             self.num_gpus_per_worker = resources_per_worker.pop("GPU")
         else:
             self.num_gpus_per_worker = int(use_gpu)
+
         self.use_gpu = self.num_gpus_per_worker > 0
         self.additional_resources_per_worker = resources_per_worker
         self.workers = []
@@ -178,12 +182,12 @@ def setup_environment(self) -> None:
             precision_plugin = current_accelerator.precision_plugin
             new_accelerator = DelayedGPUAccelerator(
                 precision_plugin=precision_plugin, training_type_plugin=self)
-            self.lightning_module.trainer.accelerator_connector \
+            self.lightning_module.trainer._accelerator_connector \
                 ._training_type_plugin = \
                 proxy(new_accelerator.training_type_plugin)
-            self.lightning_module.trainer.accelerator_connector \
+            self.lightning_module.trainer._accelerator_connector \
                 ._precision_plugin = proxy(new_accelerator.precision_plugin)
-            self.lightning_module.trainer.accelerator_connector.accelerator \
+            self.lightning_module.trainer._accelerator_connector.accelerator \
                 = new_accelerator
 
     def _setup_env_vars(self):
@@ -201,6 +205,52 @@ def _setup_env_vars(self):
         values = [os.getenv(k) for k in keys]
         ray.get([w.set_env_vars.remote(keys, values) for w in self.workers])
 
+    def _share_cuda_visible_devices(self):
+        """Sets CUDA_VISIBLE_DEVICES on all workers.
+
+        For each worker, CUDA_VISIBLE_DEVICES will be set to the GPU IDs
+        visible to all workers on that worker's node.
+
+        This allows GPU workers on the same node to communicate with one
+        another.
+
+        Example:
+
+            Setup:
+            - Node1:
+                - Worker1: {0, 1}
+                - Worker2: {2, 3}
+            - Node2:
+                - Worker3: {0, 1}
+
+            CUDA_VISIBLE_DEVICES:
+            - Worker1: "0,1,2,3"
+            - Worker2: "0,1,2,3"
+            - Worker2: "0,1"
+
+        """
+        node_ids_and_gpu_ids = ray.get(
+            [w.get_node_and_gpu_ids.remote() for w in self.workers])
+
+        node_id_to_worker_id = defaultdict(set)
+        node_id_to_gpu_ids = defaultdict(set)
+
+        for worker_id, (node_id, gpu_ids) in enumerate(node_ids_and_gpu_ids):
+            node_id_to_worker_id[node_id].add(worker_id)
+            node_id_to_gpu_ids[node_id].update(gpu_ids)
+
+        futures = []
+        for node_id, gpu_ids in node_id_to_gpu_ids.items():
+            all_gpu_ids = ",".join([str(gpu_id) for gpu_id in gpu_ids])
+
+            def set_gpu_ids():
+                os.environ["CUDA_VISIBLE_DEVICES"] = all_gpu_ids
+
+            for worker_id in node_id_to_worker_id[node_id]:
+                futures.append(
+                    self.workers[worker_id].execute.remote(set_gpu_ids))
+        ray.get(futures)
+
     def start_training(self, trainer):
         results = self.execution_loop(tune_enabled=True)
         # reset optimizers, since main process is never used for training and
@@ -257,6 +307,10 @@ def execution_loop(self, tune_enabled: bool = True):
         # This will set the MASTER_ADDR and MASTER_PORT on each Ray actor.
         self._setup_env_vars()
 
+        if self.use_gpu:
+            # Set the CUDA_VISIBLE_DEVICES for all workers.
+            self._share_cuda_visible_devices()
+
         # Get the mapping from global ranks to the respective local ranks.
         self.global_to_local = self.get_local_ranks()
 
@@ -459,9 +513,14 @@ def node_rank(self) -> int:
     @property
     def root_device(self):
         if self.use_gpu and torch.cuda.is_available():
-            # Ray already sets CUDA_VISIBLE_DEVICES for each process.
-            # So the device is the 0th index in CUDA_VISIBLE_DEVICES
-            return torch.device("cuda", 0)
+            if self._is_remote:
+                # Adjust for if there are multiple GPUs per worker.
+                device_id = self.local_rank * self.num_gpus_per_worker
+                return torch.device("cuda", device_id)
+            else:
+                # If the root device is requested on the driver, just return
+                # the 0th device.
+                return torch.device("cuda", 0)
         else:
             return torch.device("cpu")
 
diff --git a/ray_lightning/ray_horovod.py b/ray_lightning/ray_horovod.py
index 8fb5a08e..dc715608 100644
--- a/ray_lightning/ray_horovod.py
+++ b/ray_lightning/ray_horovod.py
@@ -1,5 +1,6 @@
 import torch
 import pytorch_lightning as pl
+from pytorch_lightning.accelerators import CPUAccelerator
 from pytorch_lightning.plugins import HorovodPlugin
 from pytorch_lightning.utilities import rank_zero_only
 
@@ -124,6 +125,24 @@ def setup(self):
             use_gpu=self.use_gpu)
         self.executor.start(executable_cls=get_executable_cls())
 
+    def setup_environment(self) -> None:
+        # Swap out the accelerator if necessary.
+        # This is needed to support CPU head with GPU workers or Ray Client.
+        current_accelerator = self.lightning_module.trainer.accelerator
+        if self.use_gpu and isinstance(current_accelerator, CPUAccelerator):
+            from weakref import proxy
+            from ray_lightning.util import DelayedGPUAccelerator
+            precision_plugin = current_accelerator.precision_plugin
+            new_accelerator = DelayedGPUAccelerator(
+                precision_plugin=precision_plugin, training_type_plugin=self)
+            self.lightning_module.trainer._accelerator_connector \
+                ._training_type_plugin = \
+                proxy(new_accelerator.training_type_plugin)
+            self.lightning_module.trainer._accelerator_connector \
+                ._precision_plugin = proxy(new_accelerator.precision_plugin)
+            self.lightning_module.trainer._accelerator_connector.accelerator \
+                = new_accelerator
+
     def pre_dispatch(self):
         """All pre-dispatch logic should be done in train_remote instead."""
         pass
@@ -212,6 +231,9 @@ def is_distributed(self):
     @property
     def root_device(self):
         if self.use_gpu and torch.cuda.is_available():
-            return torch.device("cuda", hvd.local_rank())
+            if hvd.is_initialized():
+                return torch.device("cuda", hvd.local_rank())
+            else:
+                return torch.device("cuda", 0)
         else:
             return torch.device("cpu")
diff --git a/ray_lightning/util.py b/ray_lightning/util.py
index f696e1ef..52718b2b 100644
--- a/ray_lightning/util.py
+++ b/ray_lightning/util.py
@@ -3,7 +3,7 @@
 
 import torch
 from pytorch_lightning.accelerators import GPUAccelerator
-from pytorch_lightning import Trainer, LightningModule
+from pytorch_lightning import Trainer
 
 import ray
 
@@ -15,12 +15,25 @@ class DelayedGPUAccelerator(GPUAccelerator):
     like the laptop) but have training still execute on GPU.
     """
 
-    def setup(self, trainer: Trainer, model: LightningModule) -> None:
-        return super(GPUAccelerator, self).setup(trainer, model)
+    def setup_environment(self) -> None:
+        # Don't do any CUDA setup.
+        # Directly call the setup_environment method of the superclass of
+        # GPUAccelerator.
+        super(GPUAccelerator, self).setup_environment()
+
+    def setup(
+            self,
+            trainer: Trainer,
+    ) -> None:
+        # Don't do any CUDA setup.
+        # Directly call the setup_environment method of the superclass of
+        # GPUAccelerator.
+        return super(GPUAccelerator, self).setup(trainer)
 
     def on_train_start(self) -> None:
         if "cuda" not in str(self.root_device):
             raise RuntimeError("GPUs were requested but are not available.")
+        torch.cuda.set_device(self.root_device)
         super(DelayedGPUAccelerator, self).on_train_start()
 
 
diff --git a/setup.py b/setup.py
index 32017e65..3a70bd2d 100644
--- a/setup.py
+++ b/setup.py
@@ -9,4 +9,4 @@
     long_description="Custom Pytorch Lightning distributed plugins "
     "built on top of distributed computing framework Ray.",
     url="https://github.com/ray-project/ray_lightning_accelerators",
-    install_requires=["pytorch-lightning==1.5.*,<1.5.9", "ray"])
+    install_requires=["pytorch-lightning==1.5.*", "ray"])


Test patch:
diff --git a/ray_lightning/tests/test_ddp_gpu.py b/ray_lightning/tests/test_ddp_gpu.py
index 04708c88..3df8b4ec 100644
--- a/ray_lightning/tests/test_ddp_gpu.py
+++ b/ray_lightning/tests/test_ddp_gpu.py
@@ -20,6 +20,13 @@ def ray_start_2_gpus():
     ray.shutdown()
 
 
+@pytest.fixture
+def ray_start_4_gpus():
+    address_info = ray.init(num_cpus=4, num_gpus=4)
+    yield address_info
+    ray.shutdown()
+
+
 @pytest.fixture
 def seed():
     pl.seed_everything(0)
@@ -32,7 +39,7 @@ def test_train(tmpdir, ray_start_2_gpus, num_workers):
     """Tests if training modifies model weights."""
     model = BoringModel()
     plugin = RayPlugin(num_workers=num_workers, use_gpu=True)
-    trainer = get_trainer(tmpdir, plugins=[plugin], use_gpu=True)
+    trainer = get_trainer(tmpdir, plugins=[plugin])
     train_test(trainer, model)
 
 
@@ -52,11 +59,7 @@ def test_predict(tmpdir, ray_start_2_gpus, seed, num_workers):
         data_dir=tmpdir, num_workers=1, batch_size=config["batch_size"])
     plugin = RayPlugin(num_workers=num_workers, use_gpu=True)
     trainer = get_trainer(
-        tmpdir,
-        limit_train_batches=20,
-        max_epochs=1,
-        plugins=[plugin],
-        use_gpu=True)
+        tmpdir, limit_train_batches=20, max_epochs=1, plugins=[plugin])
     predict_test(trainer, model, dm)
 
 
@@ -72,30 +75,29 @@ def on_epoch_end(self, trainer, pl_module):
 
     plugin = RayPlugin(num_workers=2, use_gpu=True)
     trainer = get_trainer(
-        tmpdir, plugins=[plugin], use_gpu=True, callbacks=[CheckGPUCallback()])
+        tmpdir, plugins=[plugin], callbacks=[CheckGPUCallback()])
     trainer.fit(model)
 
 
 @pytest.mark.skipif(
     torch.cuda.device_count() < 2, reason="test requires multi-GPU machine")
-def test_correct_devices(tmpdir, ray_start_2_gpus):
+@pytest.mark.parametrize("num_gpus_per_worker", [1, 2])
+def test_correct_devices(tmpdir, ray_start_4_gpus, num_gpus_per_worker):
     """Tests if GPU devices are correctly set."""
     model = BoringModel()
 
     class CheckDevicesCallback(Callback):
         def on_epoch_end(self, trainer, pl_module):
-            assert trainer.root_gpu == 0
-            assert int(os.environ["CUDA_VISIBLE_DEVICES"]) == \
-                trainer.local_rank
+            assert trainer.root_gpu == trainer.local_rank * num_gpus_per_worker
             assert trainer.root_gpu == pl_module.device.index
             assert torch.cuda.current_device() == trainer.root_gpu
 
-    plugin = RayPlugin(num_workers=2, use_gpu=True)
-    trainer = get_trainer(
-        tmpdir,
-        plugins=plugin,
+    plugin = RayPlugin(
+        num_workers=2,
         use_gpu=True,
-        callbacks=[CheckDevicesCallback()])
+        resources_per_worker={"GPU": num_gpus_per_worker})
+    trainer = get_trainer(
+        tmpdir, plugins=[plugin], callbacks=[CheckDevicesCallback()])
     trainer.fit(model)
 
 
@@ -109,5 +111,5 @@ def test_multi_node(tmpdir):
     num_gpus = ray.available_resources()["GPU"]
     model = BoringModel()
     plugin = RayPlugin(num_workers=num_gpus, use_gpu=True)
-    trainer = get_trainer(tmpdir, plugins=[plugin], use_gpu=True)
+    trainer = get_trainer(tmpdir, plugins=[plugin])
     train_test(trainer, model)
diff --git a/ray_lightning/tests/test_horovod.py b/ray_lightning/tests/test_horovod.py
index 185d7c83..775ae61a 100644
--- a/ray_lightning/tests/test_horovod.py
+++ b/ray_lightning/tests/test_horovod.py
@@ -7,7 +7,6 @@
 
 try:
     import horovod  # noqa: F401
-    from horovod.common.util import nccl_built
 except ImportError:
     HOROVOD_AVAILABLE = False
 else:
@@ -20,15 +19,6 @@
     train_test, load_test, predict_test, LightningMNISTClassifier
 
 
-def _nccl_available():
-    if not HOROVOD_AVAILABLE:
-        return False
-    try:
-        return nccl_built()
-    except AttributeError:
-        return False
-
-
 @pytest.fixture
 def ray_start_2_cpus():
     address_info = ray.init(num_cpus=2)
@@ -121,8 +111,6 @@ def test_predict_client(tmpdir, start_ray_client_server_2_cpus, seed,
     predict_test(trainer, model, dm)
 
 
-@pytest.mark.skipif(
-    not _nccl_available(), reason="test requires Horovod with NCCL support")
 @pytest.mark.skipif(
     torch.cuda.device_count() < 2, reason="test requires multi-GPU machine")
 @pytest.mark.parametrize("num_workers", [1, 2])
@@ -130,12 +118,10 @@ def test_train_gpu(tmpdir, ray_start_2_gpus, seed, num_workers):
     """Tests if training modifies model weights."""
     model = BoringModel()
     plugin = HorovodRayPlugin(num_workers=num_workers, use_gpu=True)
-    trainer = get_trainer(tmpdir, plugins=[plugin], use_gpu=True)
+    trainer = get_trainer(tmpdir, plugins=[plugin])
     train_test(trainer, model)
 
 
-@pytest.mark.skipif(
-    not _nccl_available(), reason="test requires Horovod with NCCL support")
 @pytest.mark.skipif(
     torch.cuda.device_count() < 2, reason="test requires multi-GPU machine")
 @pytest.mark.parametrize("num_workers", [1, 2])
@@ -143,12 +129,10 @@ def test_load_gpu(tmpdir, ray_start_2_gpus, seed, num_workers):
     """Tests if model checkpoint can be loaded."""
     model = BoringModel()
     plugin = HorovodRayPlugin(num_workers=num_workers, use_gpu=True)
-    trainer = get_trainer(tmpdir, plugins=[plugin], use_gpu=True)
+    trainer = get_trainer(tmpdir, plugins=[plugin])
     load_test(trainer, model)
 
 
-@pytest.mark.skipif(
-    not _nccl_available(), reason="test requires Horovod with NCCL support")
 @pytest.mark.skipif(
     torch.cuda.device_count() < 2, reason="test requires multi-GPU machine")
 @pytest.mark.parametrize("num_workers", [1, 2])
@@ -165,9 +149,5 @@ def test_predict_gpu(tmpdir, ray_start_2_gpus, seed, num_workers):
         data_dir=tmpdir, num_workers=1, batch_size=config["batch_size"])
     plugin = HorovodRayPlugin(num_workers=num_workers, use_gpu=True)
     trainer = get_trainer(
-        tmpdir,
-        limit_train_batches=20,
-        max_epochs=1,
-        plugins=[plugin],
-        use_gpu=True)
+        tmpdir, limit_train_batches=20, max_epochs=1, plugins=[plugin])
     predict_test(trainer, model, dm)
diff --git a/ray_lightning/tests/test_tune.py b/ray_lightning/tests/test_tune.py
index 613e180b..937ad51b 100644
--- a/ray_lightning/tests/test_tune.py
+++ b/ray_lightning/tests/test_tune.py
@@ -18,12 +18,18 @@ def ray_start_4_cpus():
     ray.shutdown()
 
 
-def train_func(dir, plugin, use_gpu=False, callbacks=None):
+@pytest.fixture
+def ray_start_4_cpus_4_gpus():
+    address_info = ray.init(num_cpus=4, num_gpus=4)
+    yield address_info
+    ray.shutdown()
+
+
+def train_func(dir, plugin, callbacks=None):
     def _inner_train(config):
         model = BoringModel()
         trainer = get_trainer(
             dir,
-            use_gpu=use_gpu,
             callbacks=callbacks,
             plugins=[plugin],
             checkpoint_callback=False,
@@ -85,16 +91,16 @@ def test_checkpoint_horovod(tmpdir, ray_start_4_cpus):
 
 
 @pytest.mark.skipif(
-    torch.cuda.device_count() < 2, reason="test requires multi-GPU machine")
-def test_checkpoint_ddp_gpu(tmpdir, ray_start_4_cpus):
+    torch.cuda.device_count() < 4, reason="test requires multi-GPU machine")
+def test_checkpoint_ddp_gpu(tmpdir, ray_start_4_cpus_4_gpus):
     """Tests if Tune checkpointing works with RayAccelerator."""
-    plugin = RayPlugin(num_workers=2, use_gpu=False)
+    plugin = RayPlugin(num_workers=2, use_gpu=True)
     checkpoint_test(tmpdir, plugin)
 
 
 @pytest.mark.skipif(
-    torch.cuda.device_count() < 2, reason="test requires multi-GPU machine")
-def test_checkpoint_horovod_gpu(tmpdir, ray_start_4_cpus):
+    torch.cuda.device_count() < 4, reason="test requires multi-GPU machine")
+def test_checkpoint_horovod_gpu(tmpdir, ray_start_4_cpus_4_gpus):
     """Tests if Tune checkpointing works with HorovodRayAccelerator."""
-    plugin = HorovodRayPlugin(num_hosts=1, num_slots=2, use_gpu=False)
+    plugin = HorovodRayPlugin(num_workers=2, use_gpu=True)
     checkpoint_test(tmpdir, plugin)
diff --git a/ray_lightning/tests/utils.py b/ray_lightning/tests/utils.py
index 0cc2a1c5..c797ce5d 100644
--- a/ray_lightning/tests/utils.py
+++ b/ray_lightning/tests/utils.py
@@ -149,7 +149,6 @@ def validation_epoch_end(self, outputs):
 
 def get_trainer(dir,
                 plugins: List[PLUGIN_INPUT],
-                use_gpu: bool = False,
                 max_epochs: int = 1,
                 limit_train_batches: int = 10,
                 limit_val_batches: int = 10,
@@ -160,7 +159,6 @@ def get_trainer(dir,
     callbacks = [] if not callbacks else callbacks
     trainer = pl.Trainer(
         default_root_dir=dir,
-        gpus=1 if use_gpu else 0,
         callbacks=callbacks,
         plugins=plugins,
         max_epochs=max_epochs,