[Fix] DDP init with device id

references/detection/ddp_utils.py

@@ -84,15 +84,21 @@ def sync_val_metric(
 def barrier_download(rank: int, distributed: bool):
     """Context manager that lets rank 0 populate the docTR cache before the others read it"""
 
+    def _barrier():
+        if torch.cuda.is_available() and dist.get_backend() == "nccl":
+            dist.barrier(device_ids=[torch.cuda.current_device()])
+        else:
+            dist.barrier()
+
     class _BarrierDownload:
         def __enter__(self):
             if distributed and rank != 0:
-                dist.barrier()  # wait for rank 0 to finish downloading
+                _barrier()  # wait for rank 0 to finish downloading
             return self
 
         def __exit__(self, *exc_info):
             if distributed and rank == 0:
-                dist.barrier()  # release the other ranks
+                _barrier()  # release the other ranks
             return False
 
     return _BarrierDownload()

references/detection/train.py

@@ -230,9 +230,9 @@ def main(args):
     # GPU setup
     if distributed:
         rank = int(os.environ.get("LOCAL_RANK", 0))
-        dist.init_process_group(backend=args.backend)
         device = torch.device("cuda", rank)
         torch.cuda.set_device(device)
+        dist.init_process_group(backend=args.backend, device_id=device)
 
     else:
         # single process

references/layout/ddp_utils.py

@@ -84,15 +84,21 @@ def sync_val_metric(
 def barrier_download(rank: int, distributed: bool):
     """Context manager that lets rank 0 populate the docTR cache before the others read it"""
 
+    def _barrier():
+        if torch.cuda.is_available() and dist.get_backend() == "nccl":
+            dist.barrier(device_ids=[torch.cuda.current_device()])
+        else:
+            dist.barrier()
+
     class _BarrierDownload:
         def __enter__(self):
             if distributed and rank != 0:
-                dist.barrier()  # wait for rank 0 to finish downloading
+                _barrier()  # wait for rank 0 to finish downloading
             return self
 
         def __exit__(self, *exc_info):
             if distributed and rank == 0:
-                dist.barrier()  # release the other ranks
+                _barrier()  # release the other ranks
             return False
 
     return _BarrierDownload()

references/layout/train.py

@@ -233,9 +233,9 @@ def main(args):
     # GPU setup
     if distributed:
         rank = int(os.environ.get("LOCAL_RANK", 0))
-        dist.init_process_group(backend=args.backend)
         device = torch.device("cuda", rank)
         torch.cuda.set_device(device)
+        dist.init_process_group(backend=args.backend, device_id=device)
     else:
         # single process
         rank = 0

references/recognition/ddp_utils.py

@@ -84,15 +84,21 @@ def sync_val_metric(
 def barrier_download(rank: int, distributed: bool):
     """Context manager that lets rank 0 populate the docTR cache before the others read it"""
 
+    def _barrier():
+        if torch.cuda.is_available() and dist.get_backend() == "nccl":
+            dist.barrier(device_ids=[torch.cuda.current_device()])
+        else:
+            dist.barrier()
+
     class _BarrierDownload:
         def __enter__(self):
             if distributed and rank != 0:
-                dist.barrier()  # wait for rank 0 to finish downloading
+                _barrier()  # wait for rank 0 to finish downloading
             return self
 
         def __exit__(self, *exc_info):
             if distributed and rank == 0:
-                dist.barrier()  # release the other ranks
+                _barrier()  # release the other ranks
             return False
 
     return _BarrierDownload()

references/recognition/train.py

@@ -210,10 +210,9 @@ def main(args):
     # GPU setup
     if distributed:
         rank = int(os.environ.get("LOCAL_RANK", 0))
-        dist.init_process_group(backend=args.backend)
         device = torch.device("cuda", rank)
         torch.cuda.set_device(device)
-
+        dist.init_process_group(backend=args.backend, device_id=device)
     else:
         # single process
         rank = 0

references/table/ddp_utils.py

@@ -84,15 +84,21 @@ def sync_val_metric(
 def barrier_download(rank: int, distributed: bool):
     """Context manager that lets rank 0 populate the docTR cache before the others read it"""
 
+    def _barrier():
+        if torch.cuda.is_available() and dist.get_backend() == "nccl":
+            dist.barrier(device_ids=[torch.cuda.current_device()])
+        else:
+            dist.barrier()
+
     class _BarrierDownload:
         def __enter__(self):
             if distributed and rank != 0:
-                dist.barrier()  # wait for rank 0 to finish downloading
+                _barrier()  # wait for rank 0 to finish downloading
             return self
 
         def __exit__(self, *exc_info):
             if distributed and rank == 0:
-                dist.barrier()  # release the other ranks
+                _barrier()  # release the other ranks
             return False
 
     return _BarrierDownload()

references/table/train.py

@@ -197,9 +197,9 @@ def main(args):
 
     if distributed:
         rank = int(os.environ.get("LOCAL_RANK", 0))
-        dist.init_process_group(backend=args.backend)
         device = torch.device("cuda", rank)
         torch.cuda.set_device(device)
+        dist.init_process_group(backend=args.backend, device_id=device)
     else:
         rank = 0
         if isinstance(args.device, int):