fix(comm): run collectives on their operand's device - #57
Merged
Conversation
lvyufeng
force-pushed
the
fix-comm-device-index
branch
2 times, most recently
from
August 6, 2026 05:27
0a2c22c to
a74a9cc
Compare
ProcessGroupFlagOS delegates to an inner backend (FlagCX or RCCL/HCCL) that
resolves the device it works on from whatever is *current*, not from the tensor.
A caller that binds only with torch.cuda.set_device(rank) -- the natural thing to
do -- had its collectives enqueued on the wrong device.
The cause is in the runtime, not the process group. accelerator/cuda/device.cc
kept a thread_local gCurrentDevice shadow copy of the current device index and
GetDevice returned it instead of asking the driver. flagos aliases the same
physical GPU as torch's CUDA backend, and cudaSetDevice's current device is
already thread-local, so a second copy could only diverge from it:
torch.cuda.set_device() moves the driver's index without passing through our
SetDevice, leaving the shadow stale. Measured on DTK before this change:
torch.cuda.set_device(1) -> cuda=1 flagos=0 diverged
torch.ones(4, device='flagos:1') -> cuda=0 flagos=0 cuda reset to 0
torch.cuda.set_device(0) -> cuda=0 flagos=1 diverged
The middle line is a DeviceGuard reading the stale 0 as "previous" and restoring
the process to a device that was never current. After it, every rank had device 0
current, so rank 1's collective targeted the wrong GPU.
FlagCX turns that into a GPU fault rather than a wrong answer: getStreamByIndex
lazily streamCreate()s one stream on the current device and caches it by index
alone, so the first collective binds the comm to device 0 permanently. Rank 1
then enqueued device-1 buffers onto a device-0 stream and RCCL faulted
(VMFault / "invalid resource handle") on the first all_reduce. A process kill,
not an exception.
Dropping the shadow variable and calling cudaGetDevice directly makes the two
currents identical by construction; all three lines above now agree. No guard is
needed in the process group, which is why this touches 20 delegation sites less
than it looks like it should. It also removes the need for the Python-side
DeviceGuard added to the CallPythonOp_* callers in flagos-ai#54 -- the compute-side twin
of the same accounting bug.
Only accelerator/cuda/device.cc is changed here. ascend, metax and tsingmicro
carry the same shadow variable, but their GetDevice semantics need confirming on
that hardware before being changed the same way.
Also merges two pairs of duplicate _allgather_base / _reduce_scatter_base
definitions in process_group.py, where the second silently shadowed the first.
That is unrelated to the device index; the surviving copy keeps the more
informative comment.
test_flagos_dist_live.py passed with this bug present -- it happened to leave the
two currents agreeing -- so the new tests/manual/test_comm_device_index.py
deliberately calls only torch.cuda.set_device(rank) and never touches
torch_fl.flagos.set_device.
Verified on Hygon DCU (DTK) with the FlagCX inner backend confirmed live
(pg._inner.name() == "flagcx"), with process_group.py reverted to its pre-change
state to prove the runtime fix alone is sufficient:
test_comm_device_index.py 2 and 4 cards, exit 0 (previously VMFault)
test_flagos_dist_live.py 4 cards, 20/20 collectives + DDP, grads
identical across ranks, exit 0
test_compute_device_index.py 15 passed
Local ruff check and ruff format pass across the repo (122 files).
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
lvyufeng
force-pushed
the
fix-comm-device-index
branch
from
August 6, 2026 05:39
a74a9cc to
65e6e01
Compare
zhaoyinglia
approved these changes
Aug 6, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Fixes the device-index bug that causes collectives to enqueue device-N buffers onto a device-0 stream, resulting in VMFault on FlagCX/DCU and silent wrong results on other backends.
Root Cause
csrc/runtime/accelerator/cuda/device.ccmaintained athread_local int gCurrentDeviceshadow variable thatGetDevicereturned instead of asking the driver. Since flagos aliases the same physical GPUs as torch.cuda, andcudaSetDevice's current device is already thread-local, this shadow copy could only diverge:torch.cuda.set_device(1)moves the driver's index without going through ourSetDevice, leaving the shadow at 0DeviceGuardthen reads the stale 0 as "previous" and restores the process to a device that was never currentflagos:1tensor resetstorch.cuda.current_device()back to 0ProcessGroupFlagOSdelegates to an inner backend (FlagCX or RCCL) that resolves its device from whatever is current, so rank 1's device-1 collective gets enqueued on the device-0 stream FlagCX cached. On DCU this triggers VMFault; on other hardware it silently computes wrong results or hangs.Changes
GetDevicecallcudaGetDevicedirectly_allgather_baseand_reduce_scatter_basedefinitions that were silently shadowing each othertests/manual/test_comm_device_index.pydeliberately calls onlytorch.cuda.set_device(rank)to catch this class of bugVerification
test_comm_device_index.py+test_flagos_dist_live.py(20/20 collectives + DDP), all passtest_compute_device_index.py15 passed (this fix also resolves the twin bug that PR fix(flagos): run FlagGems Python ops on their operand's device #54 patched in Python)Impact
This also fixes the device-index hazard that #54 addressed on the compute side (
CallPythonOp_*). Both were symptoms of the same C++ accounting bug.