[PR #767] fix: reuse PyTorch NCCL comm only when all ranks own one (#726) #768

New issue

Open

opened 2026-10-27 21:06:10 +00:00 by bsakaguchi · 0 comments

bsakaguchi commented

2026-10-27 21:06:10 +00:00

Owner

📋 Pull Request Information

Original PR: https://github.com/deepseek-ai/DeepEP/pull/767
Author: @Dashener2
Created: 9/29/2026
Status: 🔄 Open

Base: main ← Head: fix/comm-handle-null-comm-reuse

📝 Commits (1)

f581515 fix: reuse PyTorch NCCL comm only when all ranks own one (#726)

📊 Changes

2 files changed (+164 additions, -2 deletions)

View changed files

📝 deep_ep/comm/handle.py (+27 -2)
➕ tests/comm/test_comm_reuse.py (+137 -0)

📄 Description

Fixes #726.

Problem

get_nccl_comm_handle() cached whatever PyTorch's _comm_ptr() returned, including the
null pointer it yields before the group's NCCL communicator is lazily materialized
(eager creation only happens when init_process_group received device_id=...). The
cached null handle is later dereferenced in C++ — e.g. ncclTeamWorld reads
comm->nRanks — crashing the rank during buffer construction. This is the
vllm serve 16/16-rank failure reported in the issue.

Root cause

deep_ep/comm/handle.py:

backend = group._get_backend(torch.device('cuda'))
if not force_new_comm and hasattr(backend, '_comm_ptr') and int(os.getenv('EP_REUSE_NCCL_COMM', '1')):
    _storage[group] = NCCLCommHandle(backend._comm_ptr(), False)
    return _storage[group]

The pointer is used without checking that it is non-null. A per-rank != 0 check would
not be enough either: ranks can disagree on _comm_ptr() (a peer may have materialized
its communicator via an earlier point-to-point op, or be on a different current device),
and then some ranks would return early while others enter the group-wide unique-id
exchange below — a hang rather than an error.

Fix

Decide reuse with a group-wide all_gather_object so every rank takes the same branch,
and fall back to a DeepEP-managed comm when any rank has not materialized its
communicator yet:

comm_ptr = backend._comm_ptr()
have_comm = [None] * group.size()
dist.all_gather_object(have_comm, comm_ptr != 0, group)
if all(have_comm):
    _storage[group] = NCCLCommHandle(comm_ptr, False)
    return _storage[group]

The collective is reached unconditionally on the reuse path, so ranks cannot split
between "return PyTorch's comm" and "run the unique-id exchange".

Relationship to #727

#727 addresses the same bug but targets deep_ep/utils/comm.py, which has since moved
to deep_ep/comm/handle.py on main, so it is stale/conflicting. This ports the fix to
the current path.

Testing

tests/comm/test_comm_reuse.py adds 5 mocked cases:

all ranks hold a live comm → PyTorch's comm is reused, exactly one collective
local _comm_ptr() is null → fall through to a managed comm
rank skew (this rank live, a peer null) → every rank falls through
EP_REUSE_NCCL_COMM=0 → forces a managed comm regardless of pointer state
cache hit on the same group → no collective at all

The collectives are monkeypatched, so the test needs no GPU. Importing
deep_ep.comm.handle still requires the host extension to be built, like the rest of the
suite.

python tests/comm/test_comm_reuse.py
# All comm-reuse tests passed

_{🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.}

## 📋 Pull Request Information **Original PR:** https://github.com/deepseek-ai/DeepEP/pull/767 **Author:** [@Dashener2](https://github.com/Dashener2) **Created:** 9/29/2026 **Status:** 🔄 Open **Base:** `main` ← **Head:** `fix/comm-handle-null-comm-reuse` --- ### 📝 Commits (1) - [`f581515`](https://github.com/deepseek-ai/DeepEP/commit/f5815151bab3118a892cf85868539ed6e8833cff) fix: reuse PyTorch NCCL comm only when all ranks own one (#726) ### 📊 Changes **2 files changed** (+164 additions, -2 deletions) <details> <summary>View changed files</summary> 📝 `deep_ep/comm/handle.py` (+27 -2) ➕ `tests/comm/test_comm_reuse.py` (+137 -0) </details> ### 📄 Description Fixes #726. ## Problem `get_nccl_comm_handle()` cached whatever PyTorch's `_comm_ptr()` returned, including the null pointer it yields before the group's NCCL communicator is lazily materialized (eager creation only happens when `init_process_group` received `device_id=...`). The cached null handle is later dereferenced in C++ — e.g. `ncclTeamWorld` reads `comm->nRanks` — crashing the rank during buffer construction. This is the `vllm serve` 16/16-rank failure reported in the issue. ## Root cause `deep_ep/comm/handle.py`: ```python backend = group._get_backend(torch.device('cuda')) if not force_new_comm and hasattr(backend, '_comm_ptr') and int(os.getenv('EP_REUSE_NCCL_COMM', '1')): _storage[group] = NCCLCommHandle(backend._comm_ptr(), False) return _storage[group] ``` The pointer is used without checking that it is non-null. A per-rank `!= 0` check would not be enough either: ranks can disagree on `_comm_ptr()` (a peer may have materialized its communicator via an earlier point-to-point op, or be on a different current device), and then some ranks would return early while others enter the group-wide unique-id exchange below — a hang rather than an error. ## Fix Decide reuse with a group-wide `all_gather_object` so every rank takes the same branch, and fall back to a DeepEP-managed comm when any rank has not materialized its communicator yet: ```python comm_ptr = backend._comm_ptr() have_comm = [None] * group.size() dist.all_gather_object(have_comm, comm_ptr != 0, group) if all(have_comm): _storage[group] = NCCLCommHandle(comm_ptr, False) return _storage[group] ``` The collective is reached unconditionally on the reuse path, so ranks cannot split between "return PyTorch's comm" and "run the unique-id exchange". ## Relationship to #727 #727 addresses the same bug but targets `deep_ep/utils/comm.py`, which has since moved to `deep_ep/comm/handle.py` on `main`, so it is stale/conflicting. This ports the fix to the current path. ## Testing `tests/comm/test_comm_reuse.py` adds 5 mocked cases: - all ranks hold a live comm → PyTorch's comm is reused, exactly one collective - local `_comm_ptr()` is null → fall through to a managed comm - rank skew (this rank live, a peer null) → every rank falls through - `EP_REUSE_NCCL_COMM=0` → forces a managed comm regardless of pointer state - cache hit on the same group → no collective at all The collectives are monkeypatched, so the test needs no GPU. Importing `deep_ep.comm.handle` still requires the host extension to be built, like the rest of the suite. ```bash python tests/comm/test_comm_reuse.py # All comm-reuse tests passed ``` --- <sub>🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.</sub>

bsakaguchi added the

pull-request

label

2026-10-27 21:06:10 +00:00

bsakaguchi referenced this issue

2026-10-27 21:06:11 +00:00

[PR #768] test: add GPU-free unit tests for host-side utility modules #770