fix NCCL error with non-consecutive trainer gpus (#8165)

* device ids in barrier

x

x

s

same fix for spawn

fix non-nccl

x

* add changelog

* get nccl backend

* get backend

Co-authored-by: Kaushik B <[email protected]>

Author: 2 people

pytorch_lightning/plugins/training_type/ddp.py

@@ -295,8 +295,12 @@ def pre_dispatch(self):
     def post_dispatch(self) -> None:
         self.cluster_environment.teardown()
 
-    def barrier(self, *args, **kwargs):
-        if torch_distrib.is_available() and torch_distrib.is_initialized():
+    def barrier(self, *args, **kwargs) -> None:
+        if not torch_distrib.is_initialized():
+            return
+        if _TORCH_GREATER_EQUAL_1_8 and torch.distributed.get_backend() == "nccl":
+            torch_distrib.barrier(device_ids=self.determine_ddp_device_ids())
+        else:
             torch_distrib.barrier()
 
     def broadcast(self, obj: object, src: int = 0) -> object:

pytorch_lightning/plugins/training_type/ddp_spawn.py

@@ -271,8 +271,12 @@ def __recover_child_process_weights(self, best_path, last_path):
             ckpt = pl_load(last_path, map_location=lambda storage, loc: storage)
             self.lightning_module.load_state_dict(ckpt)
 
-    def barrier(self, *args, **kwargs):
-        if torch_distrib.is_initialized():
+    def barrier(self, *args, **kwargs) -> None:
+        if not torch_distrib.is_initialized():
+            return
+        if _TORCH_GREATER_EQUAL_1_8 and torch.distributed.get_backend() == "nccl":
+            torch_distrib.barrier(device_ids=self.determine_ddp_device_ids())
+        else:
             torch_distrib.barrier()
 
     def broadcast(self, obj: object, src: int = 0) -> object: