[Do not merge] Experiment: iteration-flag combine gate (EP_COMBINE_GATE) - #10
Closed
KeitaW wants to merge 6 commits into
Closed
[Do not merge] Experiment: iteration-flag combine gate (EP_COMBINE_GATE)#10KeitaW wants to merge 6 commits into
KeitaW wants to merge 6 commits into
Conversation
get_rdma_gbs() only knew how to ask ibstat for a CA named EP_NIC_NAME (default mlx5_0). EFA devices expose no umad CA, so on EFA hosts the probe returned 0 and get_theoretical_num_sms() divided by it -- any multi-node run without an explicit --num-sms crashed with ZeroDivisionError. Read /sys/class/infiniband/<nic>/ports/*/rate first, which works for every verbs provider, and fall back to ibstat for setups whose rate only shows there. When EP_NIC_NAME is unset and the default device is absent, pick the fastest device under /sys/class/infiniband instead of failing; an explicitly named NIC still fails loudly rather than guessing.
Add --pressure-iterations to bound the pressure-test loop (upstream's --do-pressure-test runs int(1e9) seeds, i.e. until killed; 0 keeps that behavior), with argument validation. Add a barrier before dist.destroy_process_group() so a fast rank cannot tear down the TCPStore while slower ranks are still in destroy. Signed-off-by: Xuan Jiang <xuanj@amazon.com>
…nting barrier
On EFA the shipped default refuses to initialize past 22 NVLink domains. Measured on
p6-b200.48xlarge: 22 domains complete, 23 refuse at `nccl.cu` with
gin_config.gin_indexed_signals_cnt >= (num_rdma_ranks - 1) and
"GIN indexed-signal budget cannot give each peer rail team a dedicated signal..."
The assert is correct and load-bearing, not redundant: `gin_barrier_wo_local_sync` really did
need one indexed-signal slot per peer, so the rail barrier's demand grew with the team while
the per-context budget did not. A GIN signal id is not free on EFA -- one id is one
`gdaki_sc_endpoint`, a complete QP and CQ -- so the budget is
`gin_indexed_signals_for(c) = (kTotalQPBudget - 2c)/c`, which is 21 at the shipped 11 contexts.
21 slots is 22 domains, and the 23rd is refused. Raising the budget is not available: it is a
NIC resource, independently measured at 512 completion counters per NIC (256 QPs at two
counters each), exactly what `kTotalQPBudget = 256` already assumes.
The fix removes the team-size term instead. A barrier is a counting predicate and does not need
to distinguish senders, so every peer now adds 1 to the SAME signal id and the waiter advances
its shadow by `kNumRanks - 1`. That costs ONE slot whatever the team size. The pattern is the
one the unordered data path already relies on (`hybrid_combine_unordered.cuh`'s
`num_expected_arrivals` wait), and `SignalAdd{.., 1}` matches that precedent.
This is scoped to the RAIL instantiation, deliberately, because the two teams have opposite
requirements:
* Only rail has a ceiling. The unordered-hybrid arm of `NCCLSymmetricMemoryContext` asks for
`gin_indexed_signals_cnt`, the per-context budget above. The direct / ordered arm asks for
`num_ranks + 2 * 2`, commented "Customized RDMA barrier needs extra signals" -- the world
barrier's per-peer slots are already budgeted there and scale with the team.
* Only world is used as a release barrier. Every rail call site passes `kFlushStores = false`
(`hybrid_{dispatch,combine}{,_unordered}.cuh`, `barrier.cuh`), so the rail barrier never even
issues the QP flush; it is a pure synchronisation point. The hybrid kernels' two "ensure data
arrival" barriers pass `do_scaleout = false` and run over NVLink. By contrast
`dispatch.cuh` and `combine.cuh` use the world path with `kFlushStores = true` and then read
what peers wrote.
* A counting barrier cannot carry release. Its counter is anonymous, so a peer one round ahead
can supply an increment standing in for a delayed current-round arrival: the count reaches
its target without every distinct peer having arrived. "Everyone arrived" survives; "every
peer's prior writes are visible to me" does not. Signal strength does not repair this --
strong signals order a sender's own prior puts, they never say which sender incremented.
Identity is the missing half and only per-peer slots have it, so this would hold on
InfiniBand too.
So the world body is left unchanged apart from the four-space re-indent that the new
`if constexpr` block forces around it, and can be reviewed as untouched. To check that
mechanically, strip comments and whitespace from the block in both revisions and compare --
746 characters of code, identical.
A second, live defect had to be fixed for the counting barrier to be safe at all. The barrier
hardcoded signal id 0 and nothing stopped the data path from producing that same id. In
`cached_mode` the unordered hybrid kernels run with zero notify warps, which makes
`kQPStartIdx = 0` and puts data channels on context 0 -- the barrier's own context --
where 0-based `channel_to_signal_id` yields id 0 for channel 0 part 0. NCCL addresses a shadow
by (context, signal) alone, so that was one shared counter with each side inflating the other's
arrival count. `kNumReservedBarrierSignals` now takes one id off the bottom of every context's
id space and `data_signal_id()` is the single place the offset is applied, so both derivations
in `comm.cuh` shift together. Note the offset is added AFTER the per-part multiply; going
through `get_qp_signal_id` would scale it by `kNumParts` and burn ids.
Two invariants were repaired alongside. `all_gin_context_counts_cover_warps()` compared a
cross-context TOTAL against a warp count, which says nothing about the most crowded context;
it is replaced by a serviceability check over every legal context count, verified non-vacuous
(forcing the reservation to 200 fails the build). `compute_part_allocation` is split into a pure
`_raw` and a diagnosing wrapper so the invariant can be a `static_assert` -- a `printf` or a
throwing assert reached during constant evaluation makes the expression non-constant.
`gpu_barrier` gains a static assert, mirrored by a host assert in `NCCLSymmetricMemoryContext`,
that a GIN scale-up and a GIN scale-out barrier are never live concurrently: their id ranges
overlap, since world's per-peer slots start at 0 and rail's counting slot IS 0. No instantiation
can violate it today; the host assert is what fails early, because these kernels are
NVRTC-generated and a device static assert would surface as a JIT exception in production.
Cost, measured with `compute_part_allocation` at the shipped context count: the reservation
costs one part at 16-17 SMs (3 -> 2) and one channel per SM at 51-52 SMs (4 -> 3, with the host
warning). No second id is needed, so the usable per-context budget stays at 20.
Verified on 36 x p6-b200.48xlarge (8x B200, 8 EFA/node, efa kmod 3.3.0g), treatment and control
built from the same base and each pod printing the md5 of the JIT headers it compiled:
23 domains, stock -> refuses at init, verbatim assert above
23 domains, this change -> completes, correctness checks pass
32 domains, this change -> completes (256 ranks)
22 domains, this change -> completes
2 domains, this change -> completes
3 domains, direct mode -> completes; exercises the world instantiation
3 domains, direct, stock -> completes; control for the row above
Barrier cost at 22 domains is unchanged to within run-to-run spread, and the counting barrier is
never slower than the per-peer one: 8.2% faster at 2 domains, 4.9% at 8, converging to no
difference by 22. The single-counter hotspot does not materialise through 32 domains -- marginal
cost 3.38 us/domain over 22->32 against the stock barrier's 3.25 us/domain over 16->22.
End-to-end dispatch and combine throughput is unchanged.
Remaining limitations and untested paths are tracked separately rather than in this change.
EP_COMBINE_GATE=flag replaces the unordered combine's counting-signal return gate: the sender drops the per-put SignalAdd, flushes its data puts at ring completion, then writes the combine iteration number via putValue into a per-(channel, sender) workspace slot on every rail peer; the receiver polls that slot with a wrap-safe compare. signal (the default) keeps the existing path; its SASS is unchanged from baseline except assert line numbers.
PTX ISA 9.0 added support for a 32-bit st.bulk size operand, but ptxas 13.0.88 rejects the register form even when nvcc emits PTX .version 9.0. This breaks the affected JIT kernels on sm_100-family GPUs with CUDA 13.0. Pass kNumBytes through a 64-bit "l" constraint instead. The value is unchanged, and the 64-bit form works with both CUDA 13.0 and CUDA 13.1. A CUDA 13.1.80 compile probe emits PTX .version 9.1 and accepts both the 32-bit and 64-bit register forms. Observed on 2 x p6-b300.48xlarge instances (B300, sm_103): tests/elastic/test_ep.py fails at the first dispatch on every rank with CUDA 13.0. With this fix, the 2-node internode run of tests/elastic/test_ep.py --test-first-only passes end-to-end over EFA-GDA, with RUN_EXIT=0 on both nodes and all correctness checks green. Signed-off-by: Keita Watanabe <keitaw09@gmail.com>
- Post flags via a non-blocking drain test (new try_flush(), the timeout flush overload with a 1-cycle budget) so the SM100 sweeper keeps sweeping other rings while a completed ring drains, instead of stalling in a blocking flush per ring. - Document the two invariants the flag gate rides on (EFA SRD ACK = remote settlement; the drain covers the previous iteration's flag write) at the flush sites. Both are stronger than the portable flush() contract; FlushAsync/Wait are stubs on the EFA GDA backend, so the drain-test route is the only non-blocking option. - Make combine_iteration uint32_t end-to-end; the signed int host counter hit undefined behavior at 2^31 increments. Signal arm SASS is bit-identical to 74ec698.
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
3 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.
Do not merge. This is an experiment branch for the discussion in #5 (see #5 (comment)): gate the unordered combine's return on a putValue iteration flag instead of counting signals. It stacks on #5; the flag-gate work is the last three commits.
EP_COMBINE_GATE=flagdrops the per-put SignalAdd, waits for the sender's data puts to drain at the NIC, then writes the combine iteration number via putValue into a per-(channel, sender) workspace slot on every rail peer. The receiver polls that slot with a wrap-safe compare.signal(the default) keeps the existing path, and its SASS is bit-identical to the base branch.How we measured
tests/elastic/test_ep.pyfrom this branch; the combine and dispatch times below come from its kineto-based bench.--num-tokens={4096|128} --hidden=7168 --num-topk=8 --num-experts=256 --num-sms=24 --num-allocated-qps=11(the channel auto-tuner settled at 4 channels per SM, so 96 channels over 11 QPs).EP_COMBINE_GATEenv var.Results
The absolute penalty in round 2 is ~570-620 us and does not change with payload, so the residual is not sender implementation cost. The receiver waits for the sender's whole channel batch to drain plus one flag RTT (96 channels co-drain over 11 shared QPs), while the signal gate credits arrivals put by put.
Notes for reviewers:
Our read: keep the signal gate for latency. The proposal's resource win would need a different mechanism than the combine return gate.