fix: Only bind device id when needed, Fixes #8248 (#8269)

Fixes #8248.

### The change

`init_process_group` now asks a small resolver instead of inlining the
condition:

```python
device_id = get_init_process_group_device_id(world_size)
if device_id is not None:
    kwargs.update(device_id=device_id)
```

`DEEPSPEED_SET_DEVICE_ID` overrides the decision, parsed by a new
`get_env_flag` helper:

| Value | Effect |
|---|---|
| `1` / `true` / `yes` / `on` | always bind |
| `0` / `false` / `no` / `off` | never bind |
| unset | the default above: bind iff `world_size > 1` |
| anything else | warn and ignore, so the default still decides |

Multi-rank behaviour is unchanged.

### Tests

`tests/unit/comm/test_dist.py`

---------

Signed-off-by: pengdurice <pengduhit@gmail.com>

Author: pengdurice

deepspeed/comm/torch.py

@@ -28,6 +28,22 @@ def disable_compiler_collective(func):
     return compiler.disable(func)
 
 
+def known_world_size(world_size):
+    """The world size when it can be determined, otherwise None.
+
+    ``init_distributed`` defaults its ``world_size`` argument to -1, and the value can also be
+    carried in the ``init_method`` URL (``tcp://host:port?world_size=2``), where neither the
+    argument nor ``WORLD_SIZE`` reflects it. Return None rather than a guess so that callers do
+    not act on an assumed size.
+    """
+    if world_size is not None and world_size > 0:
+        return world_size
+    env_world_size = os.environ.get('WORLD_SIZE', '')
+    if env_world_size.isdigit() and int(env_world_size) > 0:
+        return int(env_world_size)
+    return None
+
+
 def build_shm_op():
     builder = get_accelerator().create_op_builder("ShareMemCommBuilder")
     if builder is None or not deepspeed.ops.__compatible_ops__.get(builder.NAME, False):
@@ -251,7 +267,11 @@ def init_process_group(self, backend, timeout, init_method, rank, 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
             # 3. device_id works and is needed for `cuda`, other accelerators may have issues at the moment. Therefore only do it for the `cuda` accelerator.
-            if ('device_id' in inspect.signature(torch.distributed.init_process_group).parameters
+            # 4. binding a device also makes torch build every later new_group() with ncclCommSplit;
+            #    a single-rank job has no peer to reach, so skip it there (#8248). Only skip when the
+            #    world size is actually known, never on an assumed one.
+            if (known_world_size(world_size) != 1
+                    and '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"))
                     and get_accelerator().device_name() == 'cuda'):
                 local_rank = int(os.environ.get('LOCAL_RANK', 0))

tests/unit/comm/test_dist.py

@@ -3,6 +3,7 @@
 
 # DeepSpeed Team
 
+import importlib
 import os
 import torch
 import deepspeed.comm as dist
@@ -396,3 +397,58 @@ def test_no_init(self, dist_init_required):
                                                   config=config_dict,
                                                   model_parameters=model.parameters(),
                                                   dist_init_required=dist_init_required)
+
+
+# `deepspeed.comm.torch` is shadowed by the real torch module in the `deepspeed.comm` namespace.
+ds_comm_torch = importlib.import_module("deepspeed.comm.torch")
+
+
+@pytest.mark.parametrize("world_size,env,expected", [
+    (2, None, 2),
+    (1, None, 1),
+    (-1, "4", 4),
+    (-1, None, None),
+    (-1, "", None),
+])
+def test_known_world_size(monkeypatch, world_size, env, expected):
+    # -1 with nothing in the environment means the size is genuinely unknown: it may be carried
+    # in the init_method URL. Returning None keeps the caller from skipping the device binding
+    # on an assumed size.
+    monkeypatch.delenv("WORLD_SIZE", raising=False)
+    if env is not None:
+        monkeypatch.setenv("WORLD_SIZE", env)
+    assert ds_comm_torch.known_world_size(world_size) == expected
+
+
+def assert_device_binding(expect_bound):
+    """A device is bound for multi-rank jobs, and not for a single-rank one (#8248)."""
+    if get_accelerator().communication_backend_name() != 'nccl':
+        pytest.skip("device_id is only bound for the nccl backend")
+
+    deepspeed.init_distributed(dist_backend='nccl', auto_mpi_discovery=False)
+    default_pg = torch.distributed.distributed_c10d._get_default_group()
+    assert (default_pg.bound_device_id is not None) == expect_bound
+
+    # The eager split new_group() makes when a device is bound is the call that fails in #8248.
+    device = get_accelerator().device(int(os.environ["LOCAL_RANK"]))
+    default_backend = default_pg._get_backend(device)
+    if hasattr(default_backend, "comm_split_count"):
+        splits_before = default_backend.comm_split_count()
+        torch.distributed.new_group(ranks=list(range(dist.get_world_size())))
+        assert (default_backend.comm_split_count() > splits_before) == expect_bound
+
+
+class TestSingleRankDeviceId(DistributedTest):
+    world_size = 1
+    init_distributed = False
+
+    def test(self):
+        assert_device_binding(expect_bound=False)
+
+
+class TestMultiRankDeviceId(DistributedTest):
+    world_size = 2
+    init_distributed = False
+
+    def test(self):
+        assert_device_binding(expect_bound=True)