Use current device id in dist.barrier

Summary:
Pull Request resolved: facebookresearch#3350

`get_local_rank` relies on a global variable set by Detectron2's `launch` utils.
Since other frameworks might use Detectron2's distribute utils but don't launch with Detectron2's `launch` utils. Use `torch.cuda.current_device` to get the current device instead.

Reviewed By: HarounH, ppwwyyxx

Differential Revision: D30233746

fbshipit-source-id: 0b140ed5c1e7cd87ccf05235127f338ffc40a53d

Author: kazhang

detectron2/engine/launch.py

@@ -116,11 +116,11 @@ def _distributed_worker(
         if i == machine_rank:
             comm._LOCAL_PROCESS_GROUP = pg
 
+    assert num_gpus_per_machine <= torch.cuda.device_count()
+    torch.cuda.set_device(local_rank)
+
     # synchronize is needed here to prevent a possible timeout after calling init_process_group
     # See: https://github.com/facebookresearch/maskrcnn-benchmark/issues/172
     comm.synchronize()
 
-    assert num_gpus_per_machine <= torch.cuda.device_count()
-    torch.cuda.set_device(local_rank)
-
     main_func(*args)

detectron2/utils/comm.py

@@ -83,7 +83,7 @@ def synchronize():
     if dist.get_backend() == dist.Backend.NCCL and TORCH_VERSION >= (1, 8):
         # This argument is needed to avoid warnings.
         # It's valid only for NCCL backend.
-        dist.barrier(device_ids=[get_local_rank()])
+        dist.barrier(device_ids=[torch.cuda.current_device()])
     else:
         dist.barrier()