Repository navigation
collectives: binomial-tree reduce and broadcast, tests at any size and root; optional bounded spin; two build fixes - #11
Merged
Conversation
CI ran the collectives at 2 ranks only, where a tree and a linear walk take the same path. test_collectives.py checks, for every root, that reduce (SUM) leaves the sum on the root and every other rank's input unchanged, that broadcast reaches every rank, and that all_reduce gives every rank the sum. CI now also runs it at 3 and 4 ranks.
reduce_locked and broadcast_locked each walked every peer from the root, so an
allreduce cost 2*(P-1) blocking messages serialized on one rank. The source said
so ("recursive-doubling is a future optimization... payloads are tiny, so linear
is fine") -- true of the payload, not of the message count.
Measured on one 128-core node, all shared memory, allreduce of a 1-element f64
through torch.distributed (no application), linear version:
ranks 2 4 8 16 32 64 128
ucx us 20.0 54.4 77.0 144.3 252.7 513.8 992.3
Every doubling of P doubles the cost from P=8 up (1.87, 1.75, 2.03, 1.93x),
fitting 7.6 us/rank over P=8..128. A halo-sized p2p ring over the same ranks
stays flat (17.8 -> 82.7 us), so this is the collective, not the fabric.
The tree sends 2*ceil(log2(P)) messages per rank instead: 14 rather than 254 at
P=128. Virtual ranks (vrank = rank - root mod size) keep any root working, and
the mask walk handles non-power-of-two sizes.
Measured in an application, 128 ranks, 512 cells per rank, one block per rank
(snapy, a dry box, two allreduces per step), on a second cluster: the step went
from 12.47 to 9.95 ms wall (1.25x) and from 7.81 to 6.48 ms CPU on rank 0, and
from 12.31 to 9.94 ms in a repeat on other nodes. Output was bit-identical.
A reduce must leave a non-root rank's input alone: an interior node of the tree
accumulates into a clone, not into t, so after a bare reduce() every rank still
holds its own input (the minimal case of the in-place version is size 4, root 0,
rank 2). allreduce would not show it, because the broadcast overwrites t.
An out-of-range root still fails loudly, as the linear version did through
coll_send; with virtual ranks it would otherwise hang.
MIN and MAX are exact under reordering, so per-step MIN/MAX collectives are
bit-identical. SUM and PRODUCT reassociate, so a reduce-based diagnostic may
differ in the last ulp.
Every completion not already satisfied by ucp_worker_progress went straight to
ucp_worker_arm + poll(), so each cost a transport transition into interrupt mode
plus a scheduler wakeup. Measured on the matched dry box, 128 ranks, per step:
park cpu 7.94 wall 12.76
spin, never parking cpu 9.43 wall 9.48
Idle collapsed 4.82 -> 0.05 ms and reappeared as CPU: wakeup latency, not waiting
on a peer. Never parking is not shippable -- a rank then burns a core through a
genuinely long wait -- so this spins a bounded window and otherwise parks exactly
as before.
OFF by default (COMMUX_SPIN_US=0), matching COMMUX_COALESCE and COMMUX_GROUP in
this file: the consumer opts in, because spinning is not a universal win. Runs
with few large waits (GPU, 3-D) would burn a core for nothing.
The window bounds one CALL, not one wait. Both call sites reach here only after
ucp_worker_progress returned 0, so a busy path never enters the loop; but an
unrelated completion mid-spin returns early and the caller, finding its own
request still pending, re-enters with a fresh window. Total spin across a single
wait is therefore unbounded, and a root fielding many peers is exactly that
traffic. Budget CPU as "up to one window per unrelated completion observed",
not as one window per wait.
Untested: the concurrent-consumer path. worker_mu_ is held across this spin at
the wait_request call site, and the multi-threaded consumer it guards against is
snapy's BlockWorkerPool, built only when a rank owns more than one block
(mesh.cpp:244). Every run gating this used one block per rank. The pre-existing
idle path already held that lock across a poll() of up to COMMUX_WAIT_POLL_MS,
so the change there is 50 us then park rather than park immediately.
Note COMMUX_WAIT_POLL_MS=0 still means 1 ms, not "do not park" -- the clamp is
v > 0 ? v : 1.
tests/test_commux.cpp starts std::threads (the COMMUX_MT test and the deadlock
watchdog), but the target linked only commux and torch. glibc 2.34 folded
libpthread into libc, which hides this. With glibc 2.17 the link fails:
libpthread.so.0: error adding symbols: DSO missing from command line
With Threads::Threads it links there, and test_commux passes at 2 ranks with
COMMUX_MT=1 COMMUX_DEADLOCK=1.
A torch that bundles pybind11 2.x (2.6, for one) builds its extensions with
PYBIND11_BUILD_ABI="_cxxabi1011" and exposes it as torch._C._PYBIND11_BUILD_ABI.
commux._C took the compiler's own tag instead (gcc 11: _cxxabi1016), so its
pybind11 internals were not torch's and it could not see c10d::Store:
TypeError: create_backend(): incompatible function arguments. The following
argument types are supported:
1. (store: c10d::Store, rank: int, size: int, timeout: ...) -> c10d::Backend
Read the tag in the existing torch probe and, when it is set, define it for
commux_pybind, as torch.utils.cpp_extension does. With it, test_torch_dist.py
passes at 2 ranks on torch 2.6. torch 2.10 bundles pybind11 3 and has no such
attribute, so there the probe returns an empty tag and nothing is defined.
Unrelated UCX completions return from worker_wait_timed so the caller can reap, and the next call used to start a new window. One wait now keeps a single deadline. COMMUX_SPIN_US unset or 0 still skips the spin and parks.
Google style, matching pre-commit. No behavior change.
The outer progress loops treated any ucp_worker_progress as a reason to skip the park. A back-to-back stream never read the per-wait deadline, so the CPU spin ran for the whole wait. Once that deadline is past, sleep one poll interval even when ucp_worker_arm returns BUSY. No deadline (COMMUX_SPIN_US unset or 0) still returns immediately on BUSY.
UCzhangxi
marked this pull request as ready for review
September 30, 2026 04:29
UCzhangxi
requested review from
zoeyzyhu
and
a balanced review from Copilot
September 30, 2026 04:29
There was a problem hiding this comment.
Copilot review overview
🟡 Changes recommended
Late completions can incur unnecessary poll latency, and the spin regression test is not enabled in CI.
Review effort: Balanced
Findings: 2
Open (2)
What changed in this PR
Improves collective scalability, wait behavior, testing, and build compatibility.
Changes:
- Replaces linear collectives with binomial-tree reduce/broadcast.
- Adds bounded spinning and regression tests.
- Fixes thread linking and PyBind11 ABI handling.
| File | Description |
|---|---|
src/process_group_ucx.cpp |
Implements tree collectives and bounded spinning. |
tests/test_collectives.py |
Tests collectives across roots and world sizes. |
tests/test_commux.cpp |
Adds spin-budget regression checks. |
CMakeLists.txt |
Fixes threading and PyBind11 ABI configuration. |
.github/workflows/ci.yml |
Runs multi-rank collective tests. |
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Comment on lines
+329
to
+334
| if (ucp_worker_progress(worker_)) { | ||
| // Advanced. A stream of unrelated completions stays on this branch and | ||
| // would never consult the one per-wait deadline below. | ||
| if (spin.expired()) worker_wait_timed(worker_, spin.deadline()); | ||
| continue; | ||
| } |
Comment on lines
+343
to
+344
| if (const char* spin_env = std::getenv("COMMUX_SPIN_US"); | ||
| spin_env != nullptr && std::atoi(spin_env) > 0 && size == 2) { |
zoeyzyhu
approved these changes
Sep 30, 2026
UCzhangxi
pushed a commit
to UCzhangxi/commux
that referenced
this pull request
Sep 30, 2026
ucp_worker_progress can complete the request this wait is blocked on. Parking after that, or sleeping on UCS_ERR_BUSY while idle, waits out the rest of COMMUX_WAIT_POLL_MS without watching the wakeup fd. Reap first. The idle path and wait_request yield on BUSY and progress again. The p2p throttle still sleeps one poll when unrelated completions are queued and this request is still pending, so the zoeyzyhu#11 flood stays off the core.
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
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.
Seven commits on main cf41b46. Each commit message carries its own evidence; commits 6 and 7 are by Sihe Chen.
1. Tests (2bd56fa): reduce, broadcast and all_reduce at any world size and root.
tests/test_collectives.py: for every root, reduce SUM gives the root the sum and leaves the non-root inputs unchanged, broadcast reaches every rank, and all_reduce SUM agrees. Integer inputs keep the sums exact in any order. CI now also runs it at 3 and 4 ranks, since at 2 ranks a tree and a linear walk coincide.2. Collectives (8510922): a binomial tree instead of a linear reduce-to-root.
reduce_lockedandbroadcast_lockedwalked every peer from the root, so an allreduce cost 2(P-1) blocking messages serialized on one rank. At P = 8..128 on one node the linear allreduce doubles in cost with every doubling of P (7.6 us per rank; 992 us at 128 ranks), while a p2p ring over the same ranks stays flat. The tree sends 2 ceil(log2 P) messages per rank: 14 instead of 254 at P = 128. Virtual ranks keep any root working, and the mask walk handles sizes that are not a power of two. A reduce accumulates into a clone on non-root ranks, so their inputs are not mutated. In a 128-rank application step, the wall time went from 12.47 to 9.95 ms (1.25x) with bit-identical output. Allreduce of 4096 float64 on one node, spin off, microseconds per call (main cf41b46 vs this branch's library):3. Wait (0da2698): an optional bounded spin (budgeted per wait by commit 6) before parking on the worker's eventfd.
COMMUX_SPIN_US, default 0 = off, as withCOMMUX_COALESCEandCOMMUX_GROUP. Measured at 128 ranks, per step: park takes 7.94 ms CPU and 12.76 ms wall; spinning takes 9.43 ms CPU and 9.48 ms wall. The message states the CPU budget and the one untested path (a multi-threaded consumer).4. CMake (432b550): link
test_commuxwithThreads::Threads. The test starts std::threads but did not link Threads. glibc >= 2.34 hides this; on glibc 2.17 the link fails ("libpthread.so.0: error adding symbols: DSO missing from command line"). With the fix it links, and test_commux passes at 2 ranks withCOMMUX_MT=1 COMMUX_DEADLOCK=1.5. CMake (c8005b6): build
commux._Cwith torch's pybind11 ABI tag. A torch that bundles pybind11 2.x (2.6, for one) exposestorch._C._PYBIND11_BUILD_ABI. commux._C took the compiler's own tag, so it could not seec10d::Store(TypeError: create_backend(): incompatible function arguments). The existing torch probe now reads the tag and, when it is set, defines it forcommux_pybind, astorch.utils.cpp_extensiondoes. torch 2.10 (the CI pin) has no such attribute, so there nothing changes.6. Spin budget (013f5a1): one spin deadline per wait, not per call. Review found that in commit 3, every unrelated completion restarted the spin window, so one wait could spin without bound, on the blocking path while holding the worker mutex. The wait now carries one deadline across re-entries. With
COMMUX_SPIN_USunset or 0, the path is unchanged. A test intest_commuxruns withCOMMUX_SPIN_US=20000while a peer sends 40 unrelated messages 3 ms apart before the one being waited on: the waiter's thread CPU is 124.5 ms with commit 3's source (it spins through the whole drip) and 22.3 ms with this commit (one window, then it parks).7. Format (357bad7): clang-format of the new test code only.
Checks at 357bad7 (CPU, UCX, spin off unless stated): test_collectives passes at 2, 3, 4, 5 and 8 ranks; test_torch_dist.py passes at 2 ranks; test_commux passes at 2 ranks with
COMMUX_MT=1 COMMUX_DEADLOCK=1, and its spin-budget check passes withCOMMUX_SPIN_US=20000. clang-format 20.1.4 and every pre-commit hook pass.