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

[pengdurice]

Fixes #8248.

The change

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

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

[chatgpt-codex-connector]

Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits.

[sfc-gh-truwase]

@pengdurice thanks for the PR. I don't fully understand the problem. It appears that this a recent regression. Did you identify regression?

[ebarkhordar]

@sfc-gh-truwase on the regression question: yes, and the boundary is the one the reporter names. I traced this in git on a full clone rather than reproducing it, since I have no CUDA machine here.

The device_id binding entered deepspeed/comm/torch.py in ee286e5 (#7266), and v0.17.2 is the last release without it:

$ git rev-parse --is-shallow-repository
false
$ git log -G 'kwargs\.update\(device_id=' --oneline --follow -- deepspeed/comm/torch.py
ee286e53 set `device_id` in torch's `init_process_group` (#7266)
$ git tag --contains ee286e53 | sort -V | head -3
v0.17.3
v0.17.4
v0.17.5
$ git log --oneline v0.17.2..v0.17.3 -- deepspeed/comm/
ee286e53 set `device_id` in torch's `init_process_group` (#7266)
f485e136 Improvements to Communication Logger (#7404)

#8248 reports 0.17.2 working and later versions failing, so the two agree. The only other comm change in that window is the logger one.

Two things in that history bear on the fix. First, #7266's stated purpose was narrow: it set device_id to silence the ProcessGroupNCCL warning that "devices used by this process are currently unknown" during barrier(). A single-rank job has no rank-to-GPU mapping to get wrong, which is the case this PR stops binding for.

Second, this block has already been narrowed once for the same shape of reason. #7542 (08879a3, first released in v0.17.6) added the get_accelerator().device_name() == 'cuda' condition because binding a device blocked creating CPU process groups afterwards. So get_init_process_group_device_id is the second carve-out for a side effect of binding rather than a new kind of exception, and that seems worth saying in the PR description.

What I did not check: I could not run the NCCL path, so I cannot confirm that leaving the device unbound is sufficient to clear CUDA error 600 on the reporter's setup. The git evidence only establishes that the binding is what changed between the two versions they compared.

[tohtana]

Do we actually know why this happens only with world_size=1? It seems like an issue with PyTorch side.
I hesitate a bit to add a special case without knowing the real cause.

#7266 originally tried to remove this warning and chose the second approach ("call init_process_group() with a device_id). Does anyone think using barrier with device_ids works? (Any thought, @stas00?) without causing the issue in #8248?

[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.

While investigating this, I also found DeepSpeed ignores device_ids passed to barrier. It should be fixed regardless of the approach in this PR. I opened #8312 to address it.

[stas00]

I have never tried barrier with device id, @tohtana

question, why would anybody try to init dist on a single gpu. I suppose to make some code work on 1 or many gpus, but nccl won't really be used on a single gpu, correct? if I remember correctly this is just a no-op on nccl level.

if the culprit is the init of nccl on a single gpu, why not solve this in a much simpler way - just add if world_size>1

    def init_process_group(self, backend, timeout, init_method, rank, world_size):
        if not torch.distributed.is_initialized():
            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
            # 3. do not set device_id unless more than 1 gpu
            if 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")):
                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)

[pengdurice]

#8248 (comment)
Sorry for the delay, I did some preliminary investigations.
Hi @sfc-gh-truwase , @stas00 and @tohtana, thank you for reviewing this PR, I added some investigation to the link above and am still investigating it and waiting on the reporter's reply.
@stas00 to your questions: my understanding is that the communication is nearly no-op (e.g all-reduce), but the init is not. CMIIW. The DS_SET_DEVICE_ID maybe an overkill, removing it, it reduces to what you suggests.
@tohtana , to your question about if this happens only on world-size = 1, the customer reports it works fine on world size = 2. I am waiting for more information and confirmation from them. Thank you!

[stas00]

yeah, I propose to keep things simple, your code is valid but I think a simple world>1 check would have a much smaller footprint and it's a way easier to track/understand.

[tohtana]

Thank you @stas00 for sharing your thoughts. Then we can take the simplest approach.
One corner case we need to handle is having a world size in the init URL:

init_method="tcp://host:port?rank=0&world_size=2"
rank=-1
world_size=-1

In this case, world_size won't reflect the real size. We should skip the nccl init only when we are sure about the world size (!=-1).

@tohtana , to your question about if this happens only on world-size = 1, the customer reports it works fine on world size = 2. I am waiting for more information and confirmation from them. Thank you!

@pengdurice If they don't reply, let's move forward and merge this. I think this PR is already useful.

[stas00]

before merging that get_env_flag shouldn't be there, surely DS already has that somewhere in utils.

no need to pollute the code.