ccl: hierarchical cross-node AllGather (intra-node SDMA + inter-node RDMA)#441
Open
inkcherry wants to merge 123 commits into
Open
ccl: hierarchical cross-node AllGather (intra-node SDMA + inter-node RDMA)#441inkcherry wants to merge 123 commits into
inkcherry wants to merge 123 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
July 3, 2026 12:20
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.
…licates inline (A#14)
…ctor mutation (C6)
…N.md refs & M2 milestone (C9,C10)
…d missing MIT header (C7)
…nts — batch 1 (C7)
…nts — batch 2 (C7)
…nts — batch 3 (C7)
…nts — batch 4 (C7)
…tor narration (C7)
Rename ambiguous `O`->`off` (E741) and hoist `import time as _time` to the top import block (E402). Behavior-neutral lint fixes.
Delete dead _use_local_fsdp() (only call site was commented out) and its now-unused importlib/sys imports; remove the /root/wuyl private-path comment; drop the import-time [verify-fsdp] diagnostic prints and their imports. Fixes the E402 import errors. Behavior-neutral.
Superseded by the per-UT launchers (run_ut_ag_perf.sh / run_ut_overlap.sh / run_e2e.sh) that the README's Reproduce section actually drives. Proven no external reference via git grep. Also drops leaked internal hostnames/IP/container names it carried.
…ts, docs (A#1) Strip the developer's private paths and infra identifiers out of the committed benchmark artifacts: - /root/wuyl/pytorch-main -> <torch> in raw run logs - internal hostnames/IP/container name -> placeholders in logs, launcher scripts (run_e2e/run_ut_ag_perf/run_ut_overlap), and mi355x README - fix mi355x README's stale 'Files:' log names to the actual filenames Logs kept (not deleted) because the plot scripts parse them; sanitized in place. Behavior-neutral.
…A#6/A#7) Soften the 'every figure is regenerated' claim to state that only the standalone AllGather figure has a plot script; the overlap/E2E figures are captured artifacts. Add a 'Second platform' section linking the previously orphaned mi355x_ainic result tree.
…(C7) Remove the 'Turn N / Director <ts>Z / Team A / T47/T48/...' development-run narration from shipped comments, rewriting the technically useful ones in timeless terms (the invariant/why) and dropping the process references. No code changes; comments only.
…py (C5) The MORI_FSDP_ENABLE_SDMA branch imports MoriSdmaAllGather from a torch-internal module (_mori_sdma_allgather) that ships in a local PyTorch FSDP overlay, not upstream torch. Add a comment stating the requirement and pointing at the overlay-free MORI_FSDP_ENABLE_HIER path. Comment only.
…reen (finding D#0 E702/E701/F821)
…uard stub (finding D#1)
…nesses (finding D#3)
…is covered (finding D#7)
… files Running the pinned pre-commit (clang-format v20.1.8 / black 25.1.0) over the full PR diff reformats ~13 C++ headers/sources and 2 Python files that were committed un-formatted earlier in the branch. Pure formatting, no behavior change; brings the whole PR diff to a green pre-commit gate.
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