ccl: hierarchical cross-node AllGather (intra-node SDMA + inter-node RDMA)#441
Open
inkcherry wants to merge 38 commits into
Open
ccl: hierarchical cross-node AllGather (intra-node SDMA + inter-node RDMA)#441inkcherry wants to merge 38 commits into
inkcherry wants to merge 38 commits into
Conversation
…RDMA)
Add mori.ccl.HierAllGather: an all_gather_into_tensor-compatible collective that
keeps intra-node traffic on the SDMA copy engines (XGMI) and moves inter-node
traffic over RDMA. A fused ring||local-gather kernel runs the inter-node RDMA
ring concurrently with the ring-independent local node-block SDMA gather in one
grid (stream-ordered, direct-to-output, no staging copy).
Bit-exact vs torch.distributed.all_gather_into_tensor for {bf16,fp16,fp32,int32}.
On 2 nodes x 4 GPUs (MI355X), fp32: standalone bandwidth >= RCCL for sizes
>=8MB (1.19-1.35x); under a concurrent GEMM the SDMA path overlaps with compute
and is 16-20% faster than RCCL at 128-512MB (copy engines vs CU contention).
Includes tests (test_hier_allgather*), size-sweep + gemm-overlap benches, a
plot script, and the measured result charts/CSVs under benchmarks/.
…ignature HierAllGather now auto-detects the node-local rank count (LOCAL_WORLD_SIZE, else hostname grouping, else npes) so callers use the same constructor/call signature as the flat AllgatherSdma with no new required argument. ranks_per_node is now optional and keyword-only; added transit_buffer_size for signature parity. Single node still degenerates to the pure intra-node SDMA path.
Resolve conflicts with the upstream param-contiguous SDMA allgather: - oneshot_sdma_kernel.hpp: keep both the hierarchical sub-group/broadcast SDMA kernels and the upstream param-contiguous kernel (additive). - symmetric_memory.cpp: keep deviceHandles_d indexing by global pe (the array is worldSize-sized and all SDMA kernels index by global pe); adopt the upstream non-fatal GPU-metadata teardown.
Add HierAllGather.all_gather(tensor_list, tensor) matching
torch.distributed.all_gather (list output), built on the same hierarchical
intra-node SDMA / inter-node RDMA path as the contiguous all_gather_into_tensor.
Bit-exact vs torch across {bf16,fp16,fp32,int32}; adds test_hier_allgather_list.
…DMA + inter-RDMA)
…A + inter RDMA) Adds a single drop-in FSDP2 AllGather backend, MoriAllGather, used identically for single-node and cross-node runs via the stock FSDPModule.set_custom_all_gather API. It routes intra-node traffic over SDMA copy engines (XGMI) and, when the process group spans multiple nodes, inter-node traffic over RDMA — the same object handles both, so user code is unchanged between one node and many. (MoriHierAllGather kept as a backward-compat alias.) Includes the Qwen-7B FSDP2 step benchmark, a 2-node driver, and the chart script. No mori source change needed; the HierAllGather primitive already exists and handles the single-node case as pure intra-node SDMA.
inkcherry
force-pushed
the
inkcherry/sdma-hier-allgather
branch
from
c8d4eca to
4188668
Compare
… to beat RCCL xnode) Motivation: cross-node FSDP2 lost to RCCL only because HierAllGather had no param-contiguous output, forcing the backend to copy-out rank-major->param on every gather. HierAllGather.enqueue_param_contiguous now PUSHES the gathered result straight into FSDP's [param][rank] output via the existing per-slot direct gather (no new C++ kernel): per (node-block m, param s) it scatters with dst_block_offset=O_s*W+m*G*E_s, dst_slot_stride=E_s so rank r lands at O_s*W+r*E_s. Adapter sets supports_param_contiguous_output + builds dtype-elem splits. 2-node bit-exact vs torch all_gather_into_tensor (bf16/fp16/fp32/int32, 3 reps).
…l N*P launch overhead to beat RCCL xnode) Motivation: zero-copy killed FSDP copy-out but the per-(node-block,param) loop issued N_nodes*N_params SubGroup launches per all-gather; that launch overhead regressed 2-node FSDP to 106 TFLOPS (< RCCL 128). New fused OneShotAllGatherSdmaSubGroupParamContiguousKernel loops all blocks+splits inside one launch. 2-node standalone bit-exact PASS bf16/fp16/fp32/int32.
…the inter-node RDMA ring (recover ring||gather overlap zero-copy lost to serial path; lever to beat RCCL xnode)
…fault under FSDP; standalone bit-exact) The ring||local-scatter overlap zero-copy path is bit-exact in the 2-node standalone test but triggers an HSA memory-exception under FSDP's repeated-call/ buffer-reuse pattern. Keep the proven non-overlap fused scatter as the default zero-copy path; enable overlap with MORI_HIER_PC_OVERLAP=1 to iterate the fault.
…ng path); catches non-bit-exact scatter (num_blocks=1/W=G) behind the +17.6% HSDP result
… register-once regression test (isolates num_blocks=1 scatter bug; num_blocks=N cross-node path re-confirmed bit-exact)
…ntiguous kernel (align to proven flat kernel); add receiver-slot diagnostic to intra test Fused subgroup param-contiguous scatter proven buggy via single-node repro (world=4 G=4 num_nodes=1): concurrent-warp multi-put scrambles/drops sender data despite correct offsets. threadfence_system (SDMA async, unaffected) only shifted the pattern. Per-put gather_kernel_direct loop predecessor was bit-exact.
…per-op cross-node barrier is the FSDP gap (SDMA 102.9->112.1, beats RCCL 105.6 when removed); justifies generation-counter barrier-free ring
…item()/.tolist() drained the pipeline every call, destroying AG<->backward overlap under FSDP); cache u32 split tensors
…intra barriers) — decisive A/B falsifies barrier-skew: all-barrier removal gives 0% recovery (111.98->111.79), gap is not per-op barrier serialization
… was an undersized-output IPC-registration artifact, NOT a concurrent-put race Root cause (proven via fast single-node nproc=4 repro + a size sweep): the intra param-contiguous direct scatter writes to peerPtrs[remotePe]+dstBaseOffset assuming the registered peer pointer == the output buffer base. When the standalone test's output was small (bf16, ~16MB) torch SUB-allocated it inside a larger pool segment, so ShmemSymmetricRegister/hipIpcGetMemHandle resolved the peer pointer to the SEGMENT base (not the buffer) and the scatter SILENTLY CORRUPTED (wrong slots). float32 (2x larger) crossed the own-segment threshold and passed, which made the failure look dtype/num_blocks-specific. Enlarging the output to its own segment makes num_blocks=1 bit-exact for bf16/fp16/fp32/int32 (12 reps). The kernel was never buggy; a race is not dtype-size-deterministic. Bumps _PARAM_SPLITS to force own-segment; adds probe_nb.py.
…contiguous scatters (shared per-groupPos flag slot) The overlapped param-contiguous zero-copy path (MORI_HIER_PC_OVERLAP=1, the historic best 113.9 TFLOPS win candidate) ran the side-stream LOCAL-block scatter concurrently with the main-stream REMOTE-block scatters. Both call gather_kernel_direct_param_contiguous, which shares ONE per-groupPos flag slot + seq token on the intra handle; concurrent use let a receiver observe the other scatter's flag bump -> premature completion / spin-deadlock under FSDP (the Turn 4/14 hang). Move main.wait_stream(side) BEFORE the remote loop so the two scatter phases are serialized. The key overlap (side local scatter || inter-node RDMA ring) is preserved: the ring finish is enqueued on main before the wait, so it still runs concurrent with the side scatter; only the ring-dependent remote scatters wait. Kernel correctness itself is now proven bit-exact (see prior commit).
…ut on the side stream so the caching allocator does not recycle FSDP buffers mid-scatter
…hurn per AG call (each change = a cross-node ShmemSymmetric register/deregister collective that cannot overlap; candidate for the in-FSDP per-AG inflation vs RCCL)
…oss-nondeterminism is an async completion-capture bug — forced stream.synchronize() at op return gives bit-exact loss==native (11.0992556 x2), while on-device ShmemBarrierOnStream does not drain the local SDMA copy-engine DMA before FSDP consumes the output
…ns on the peer-completion flag with a SYSTEM-scope acquire (AtomicLoadSeqCstSystem) + __threadfence_system instead of an AGENT-scope relaxed load; the flag+data are produced by a REMOTE peer GPU (different HSA agent) so AGENT-scope gives no cross-agent happens-before and the copy-OUT could consume not-yet-visible data under rapid FSDP reuse (host sync masked it)
… elementwise copy-out) to test the copy-engine<->CU coherence root cause for FSDP loss-drift; both CU-write-output and CU-read-transit variants leave loss in the same drift band => the stale bytes are in out_ itself (SDMA receiver drain), not the copy-out engine
…iguous zero-copy — reproduces the FSDP AG->backward hazard (AG on comm stream, consumer waits on a recorded event on the main stream, per-rep varying inputs + layout-weighted consumer + compute pressure); proves the zero-copy direct path is bit-exact AND deterministic under overlap at ~2M-elem splits => the FSDP loss drift is NOT in the AG kernel completion at this regime (look to copy-out path / small-layer routing next)
… deployed FSDP perf path) + small/single-split size profiles — cross-node world=8 shows BOTH the copy-out and zero-copy AllGather paths are bit-exact AND deterministic under the FSDP AG->consumer cross-stream overlap at every size band (0/40 wrong); NaN-safe verdict. => the ~0.15% FSDP loss drift is NOT in HierAllGather completion ordering; redirect to the downstream FSDP path (reduce-scatter / bf16 accum) or an internal side-stream event-join not captured by the caller event
… mode; mixing copy-OUT __call__ and zero-copy enqueue_param_contiguous on ONE handle contaminated shared intra flag/seq + output-registration state (copy-OUT registers transit out_, zero-copy registers the USER output), spuriously yielding a stable-NaN for [zerocopy bf16 nsplit=5] right after copy-OUT ops. The pure-mode bit-exact test PASSES bf16 at that exact config, and FSDP uses one mode per run, so per-mode handles are the faithful harness. Now 12/12 configs 0/40-wrong => both AG output paths bit-exact + deterministic under FSDP AG->consumer cross-stream overlap at every size band; the ~0.15%% FSDP loss drift is NOT in the AG
…ther truth ref Large profile (68M elems/rank; gathered int32 ~2.18GB, crosses 2^31 bytes) probes u32 byte-offset overflow at the embed+lm_head band -> 2-node bit-exact PASS (bf16/fp32/int32), ruling out size-dependent AG corruption. Overlap UT now checks the async output against an INDEPENDENT all_gather reference (not the self-golden), so a stable-but-wrong copy-out drain would be caught -> still 0 wrong at all bands, falsifying the in-kernel drain-race theory for the FSDP loss drift.
… receiver flag wait The InterNodeRing receiver spun on the completion flag with core::AtomicLoadRelaxed and no acquire fence, unlike the intra SDMA gather (AtomicLoadSeqCstSystem + __threadfence_system). The flag is bumped by a REMOTE peer's RDMA AMO and the chunk it guards is landed by that peer's RDMA put -- both cross-agent writes; a relaxed load establishes no happens-before, so the received data need not be coherently visible to the forward-put / copy-out. Harden to a system-scope acquire matching the proven intra pattern. NOTE: in-situ AG-output probe (MORI_FSDP_AG_VERIFY) shows this alone does NOT close the FSDP loss race (184/384 AG calls still read stale vs RCCL, 0/384 under host sync) -- the residual race is in the copy-out drain (ring-buffer->output visibility), the next target.
…r corrupts ~48% of FSDP all-gathers (remote half stale); serial path is 2-3%, DEBUG_SYNC 0 In-situ 2-node FSDP AGVERIFY (world=8, copy-out, VOCAB=32000 LAYERS=28, s6/w1) localizes the cross-node loss drift to the FUSED kernel: fuse_local=1 -> 184/384 stale AG calls (remote/RDMA half of the big embed/lm_head band read before the concurrent ring CTA's remote puts are globally visible to finish_ring_stream + the remote-block direct gathers); fuse_local=0 -> 8-11/384; DEBUG_SYNC -> 0/384. Ruled OUT the ring flag/data QP-ordering theory (numQp=1 and numQp=4 both 184). Ship the serial direct path (correctness) until the fused ring-completion visibility is fixed on-device; opt back in via MORI_HIER_FUSE_LOCAL=1.
… send — drain the RDMA send-queue (all QPs) before the completion-flag AMO for a CROSS-NODE neighbour, instead of the SDMA-typed memObj quiet that only drains P2P. Closes the flag-beats-data remote-landing race on the serial ring path: FSDP AGVERIFY 1/384->0/384 and combo loss spread 0.26%->0.03% at +1.7% vs RCCL; same-node P2P keeps the SDMA quiet (crash-free). Residual full-speed loss offset remains (receiver/copy-out drain)
…place the copy-engine hipMemcpyAsync D2D drain of the RDMA-landed ring buffer with a RingFinishCopyKernel_u32 compute-unit copy that reads the same HBM the ring kernel fenced (system-scope acquire + threadfence). Motivation: close the receiver/copy-out remote-visibility residual left after the T14 sender-side RDMA quiet. Result: perf-neutral (combo 124.8 vs 125.1 mean, all reps > RCCL 123.0) and modestly improves mean loss drift (-0.18%->-0.09%) but A/B (MORI_HIER_RING_CU_COPYOUT=0) shows on~=off, so copy-engine coherence is NOT the dominant residual. Env-gated (default ON), no host sync, keeps ring<->gather overlap.
…RRIER) — force a full cross-PE ShmemBarrierOnStream on the first Phase-B reassembly gather even when slice_fuse_ib drops it, so the intra SDMA gather cannot read a peer's ring 'collection' before that peer's ring finish is globally visible. Targets the residual host-sync-recoverable FSDP loss race on the combo path; default OFF (perf-neutral until enabled) for A/B.
…UT_SIGNAL, default OFF) — fuse the single-warp RDMA data WRITE + completion-flag AMO into one ShmemPutMemNbiSignal so the flag WQE rides the SAME QP strictly after the data (RC in-order => flag can never be observed before its data lands remotely, no host sync). The last untried transport-level lever for the residual FSDP loss completion race. TESTED: standalone 2-node bit-exact PASS (small+Qwen-large bands, bf16/fp16/fp32/int32). NEGATIVE at FSDP: the perf combo runs numQp=4 -> the fanOut path (already quiet-before-flag), so this single-warp branch is inert there (loss unchanged -0.2%); forcing numQp=1 engages it but forfeits the +7.6% perf and does not conclusively close the race (-0.09%, 1 rep, within the straddle). Default OFF; safe. Same-window seq2048: native 208.48 / combo+ps mean 224.4 = +7.6%.
…I_HIER_RING_PUT_SIGNAL, default OFF) Each fan-out warp now fuses its data WRITE + a per-QP flag AMO_ADD(1) via ShmemPutMemNbiSignalWarp on qpId=warpId, and the receiver waits for the flag to reach fanActive (the # of active fan-out warps) instead of just nonzero. On RC each QP executes WRITE-then-AMO in order, so ALL QPs' data is globally visible before the receiver proceeds -- no separate quiet, no host sync, ring<->gather overlap kept. The T22 single-warp put-signal was inert in the perf combo (numQp=4 -> fanOut path); this exercises it exactly there. TESTED: standalone 2-node bit-exact PASS (small + Qwen-large, bf16/fp16/fp32/int32, numQp=4). FSDP NEGATIVE: loss still non-deterministic (r1 -0.022%, r2 -0.33%, spread 0.31%) at +7.6-8.0% vs RCCL -> the residual cross-node loss race is NOT per-op flag ordering; it is the physical RDMA remote GPU-visibility timing that only host-sync (SYNC_BIG) closes. Exhausts the last per-op transport-completion avenue. Default OFF; safe.
…Imm) Phase-5 scaffold for the inline-flag ring: RDMA_WRITE_WITH_IMM carries a 32-bit immediate that reaches the receiver's recv-CQ strictly AFTER its payload lands remotely (RC in-order), so a CQ-observed completion proves data visibility - unlike the separate GPU-memory flag AMO that the residual FSDP loss race exploits. Additive/uncalled (send-side only); receiver recv-CQ-poll wrapper + ring wiring next. TESTED: worktree gfx950 rebuild OK; standalone 2-node HierAllGather param-contiguous small-band bit-exact PASS (bf16/fp16/fp32/int32) - header compiles into the JIT'd inter-node ring kernels, working path unchanged.
…RecvCqImm) Phase-5 inline-flag ring, receiver half (pairs with T29 PostWriteImm send side). Polls a recv CQ for the RDMA-WRITE-with-immediate completion: color-bit readiness (as PollCq), op must be RDMA_IMM=3, immediate decoded from recv.imm_data_rkey. The CQE cannot be observed before its payload DMA lands remotely, so consuming it proves global data visibility -- the ordering guarantee the separate GPU-memory flag lacks (the residual FSDP cross-node loss race). Strictly additive + UNCALLED scaffold; working ring/perf path untouched. TESTED: worktree gfx950 rebuild OK, HierAllGather import OK; standalone 2-node param-contiguous test [small] band bit-exact PASS (bf16/fp16/fp32/int32) => the edited transport headers JIT-compile cleanly into the inter-node ring kernels and the working AG path is unchanged.
…ig.dedicatedRecvCq, default OFF) Phase-5 recv-infra step (a): the inline-flag ring rewrite needs RDMA_IMM recv CQEs on a CQ distinct from the ring's send completions (T31 blocker: collective QP shares one CQ for send+recv, so a WRITE_IMM recv CQE would be mis-consumed by the send-CQ drains). Add an opt-in second IonicCqContainer wired as the QP recv_cq, expose its device ring/db_val/mask, and publish a new RdmaEndpoint.recvCqHandle. Default OFF => recv_cq == send cq and recvCqHandle mirrors cqHandle, so the working ring/perf path is byte-identical. Standalone 2-node param-contiguous bit-exact PASS (bf16/fp16/fp32/int32); gfx950 build clean.
inkcherry
force-pushed
the
inkcherry/sdma-hier-allgather
branch
from
e968e15 to
8bfc9b3
Compare
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
1 participant
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.
Motivation
Technical Details
Test Plan
Test Result
Submission Checklist