set device_id in torch's init_process_group (#7266)

This PR overcomes this issue when using any `torch.distributed` calls w/
deepspeed:
```
[W404 00:15:21.693690333 ProcessGroupNCCL.cpp:4561] [PG ID 0 PG GUID 0 Rank 0]  using GPU 0 
to perform barrier as devices used by this process are currently unknown. This can
 potentially cause a hang if this rank to GPU mapping is incorrect. Specify device_ids in
 barrier() to force use of a particular device, or call init_process_group() with a device_id.
```
by setting `device_id` to the correct device corresponding to
`LOCAL_RANK` env var.

-------------------

Update: discovered `torch.dist` deadlocks with `torch=>2.7.0` when using
`device_id` arg - switching to draft for now as we can't commit this
until we know how to work around this.

---------

Signed-off-by: Stas Bekman <stas@stason.org>
Signed-off-by: Stas Bekman <stas.bekman@snowflake.com>
Co-authored-by: Olatunji Ruwase <tunji.ruwase@snowflake.com>
Co-authored-by: Logan Adams <114770087+loadams@users.noreply.github.com>
Co-authored-by: Stas Bekman <stas.bekman@snowflake.com>

Author: 4 people

accelerator/cuda_accelerator.py

@@ -64,7 +64,7 @@ def communication_backend_version(self):
         return torch.cuda.nccl.version()
 
     def device(self, device_index=None):
-        return torch.cuda.device(device_index)
+        return torch.device('cuda', device_index)
 
     def set_device(self, device_index):
         torch.cuda.set_device(device_index)

deepspeed/comm/torch.py

@@ -5,6 +5,8 @@
 
 import deepspeed
 from deepspeed import utils
+from packaging import version
+import inspect
 
 from .utils import *
 from .backend import *
@@ -145,11 +147,21 @@ def has_reduce_scatter_tensor(self):
 
     def init_process_group(self, backend, timeout, init_method, rank, world_size):
         if not torch.distributed.is_initialized():
-            torch.distributed.init_process_group(backend,
-                                                 timeout=timeout,
-                                                 init_method=init_method,
-                                                 rank=rank,
-                                                 world_size=world_size)
+            kwargs = dict(
+                timeout=timeout,
+                init_method=init_method,
+                rank=rank,
+                world_size=world_size,
+            )
+
+            # 1. device_id arg was added in torch==2.3
+            # 2. setting device_id leads to hanging in 2.6.0<torch<2.7.1 https://github.com/pytorch/pytorch/issues/153960
+            if 'device_id' in inspect.signature(torch.distributed.init_process_group).parameters and not (
+                    version.parse("2.6.0") < version.parse(torch.__version__) < version.parse("2.7.1")):
+                local_rank = int(os.environ.get('LOCAL_RANK', 0))
+                kwargs.update(device_id=get_accelerator().device(local_rank))
+            torch.distributed.init_process_group(backend, **kwargs)
+
         self.using_mpi = torch.distributed.get_backend() == 'mpi'
 
     @disable_compiler_collective