feat(distributed): support ragged (non-divisible) lengths in HOST ring allreduce - #2603
Conversation
|
Important Review skippedAuto incremental reviews are disabled on this repository. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Team Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Team Run ID: 📒 Files selected for processing (5)
Included review availability: Your plan provides up to 2 included reviews per hour; 1 remains after this review. 📝 WalkthroughWalkthroughChangesRagged ring allreduce
Estimated code review effort: 3 (Moderate) | ~20 minutes Merge Risk: 🟡 Moderate · up to The PR enables ragged and dynamic tensor lengths, but inconsistent runtime metadata across ranks could cause one rank to exit while peers wait, potentially hanging the collective and leaving communication state unsuitable for reuse. Merge should require explicit owner acceptance or a coordinated failure and reset plan. Sequence Diagram(s)sequenceDiagram
participant HostProgram
participant LoweringPass
participant RingKernel
participant RankOutputs
HostProgram->>LoweringPass: lower non-divisible allreduce
LoweringPass->>RingKernel: emit builtin.tensor.allreduce_ring
RingKernel->>RankOutputs: process balanced ragged chunks
RankOutputs-->>HostProgram: return summed tensors
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
Full details: Docstring CoverageExplanation Docstring coverage is 42.86% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 21 functions across 4 files. (1 skipped: 1 unsupported.) Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: bb1fb3fa6f
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
daafcb9 to
0cbbd47
Compare
YunjiQin
left a comment
YunjiQin
left a comment
There was a problem hiding this comment.
Two review notes on the compile-time-check removals — the runtime ragged geometry itself looks right to me (the floor(numel*r/NR) boundaries match the InCore composite ring's FP32 convention in lower_composite_ops_pass.cpp:1413, numel is already accumulated at runtime from data_tensor->shapes[], empty chunks at numel < NR are handled, and the divisible path is numerically identical to before). Both notes are about scope and dead code, not the algorithm.
…g allreduce (hw-native-sys#2242) Relax the ring allreduce numel % NR == 0 constraint so the host builtin ring allreduce handles arbitrary/ragged lengths (issue hw-native-sys#2242). - lower_host_tensor_collectives_pass.cpp: delete the static-shape + divisibility checks and the now-redundant CheckRingChunkConstraints wrapper (its remaining DistributedTensorType assert is duplicated by MakeBuiltinAllReduce and the ring builtin deducer); relocate the ragged-partition comment to the ring-builtin selection in MakeBuiltinAllReduce. The kernel partitions src into NR balanced, potentially ragged chunks at runtime ([floor(numel*r/NR), floor(numel*(r+1)/NR))). - Ring builtin deducer (collective.cpp): drop the compile-time static-shape + numel % nr validation block (the kMaxSupportedRanks check stays). - allreduce_ring kernel template: RingChunkStart() helper, soft guard no longer rejects non-divisible numel, max_chunk_elems = ceil(numel/nranks), tile_cols capped at kTileCount, and both RS/AG loops iterate ragged send spans. - ST: ragged host ring test {1,17,4097,65537} x P=2/P=4 (Sum), stage-in/out chunked with 32-byte-aligned sub-tiles. - UT: reject_nondivisible -> accept_nondivisible (SIZE=27, P=4); ring dynamic- extent deducer test (dynamic src + [7,4] INT32 signal type-deduces). - docs: update the ring allreduce contract in 43-lower_host_tensor_collectives.md (EN+ZH) and 01-faq.md (EN+ZH) — drop the obsolete static + numel%NR==0 restriction; document the runtime ragged partitioning. Verified: UT 51/51; a2a3sim 12/12; NPU P=2 (d4-5) 6/6 and P=4 (d4-7) 12/12; pre-commit green for the changed files.
0cbbd47 to
7c0ca78
Compare
…ring site `builtin.tensor.allreduce_ring` is `set_internal_only(true)`, so its type deducer is only ever reached through `LowerHostTensorCollectives` -> `MakeBuiltinAllReduce` -> `OpRegistry::CreateInternal`. A UT that builds the user-level `pld.tensor.allreduce(mode="ring")` runs `DeduceTensorAllReduceType` instead, which never inspected `shape_` and never read `mode` — so it cannot observe the static-shape requirement this branch removed. Add the lowering-side companion: a host_orch whose ring `src` window is `[pld.world_size(), SIZE]`, i.e. a `pld.system.world_size` Call rather than a `ConstInt` dim, which is exactly what the deleted check rejected. Verified non-vacuous by temporarily restoring that check: this test fails with "requires a statically-known src shape", while the existing test_tensor_allreduce_ring_accepts_dynamic_shape still passes.
2 participants
Summary
Closes the HOST-ring ragged / arbitrary-length leftover of issue #2242.
The host builtin ring allreduce previously required a statically-known src shape with
numel % NR == 0. This PR relaxes both enforcement points and makes the kernel partition each rank's buffer into NR balanced, potentially ragged chunks at runtime — chunkrspans[floor(numel·r/NR), floor(numel·(r+1)/NR))— mirroring the InCore composite ring (#2161).Changes
CheckRingChunkConstraints(src/ir/transforms/lower_host_tensor_collectives_pass.cpp): dropped the static-shape +src_numel % NR == 0CHECK_SPAN. Only the distributed-type internal check remains, plus a comment documenting the runtime ragged partition.src/ir/op/distributed/collective.cpp): dropped the compile-time static-shape +numel % nrvalidation block (thekMaxSupportedRankscheck stays).allreduce_ringkernel template (python/pypto/runtime/builtins/collectives/allreduce_ring/templates/kernel.cpp.in): addedRingChunkStart(); the soft guard no longer rejects non-divisiblenumel;max_chunk_elems = ceil(numel/nranks);tile_cols = min(max_chunk_elems, kTileCount); both RS/AG loops iterate ragged send spans, with each TPUT transfer narrowed to the chunk's exact extent via the staging tile's valid shape (ColMaskInternal).tests/st/distributed/test_l3_host_tensor_allreduce_ring.py): new ragged test{1, 17, 4097, 65537}× P=2/P=4 (all non-divisible by both P=2 and P=4; 65537 > UB); stage-in/out chunked (STAGE_CHUNK=8192, 32-byte[8,1]tile forSIZE=1, stage cols rounded to 32-byte row alignment).tests/ut/ir/transforms/test_lower_host_tensor_collectives.py):test_host_allreduce_ring_rejects_nondivisible_numel→test_host_allreduce_ring_accepts_nondivisible_numel(SIZE=27, P=4).Verification
test_lower_host_tensor_collectives.pytest_l3_host_tensor_allreduce_ring.py(uniform + ragged + signal-reuse)The NPU runs are the ordering gate: TPUT push +
NeighborBarrier+pipe_barrier(PIPE_ALL)/dsbordering is not modelled in sim, and is proven on silicon for every ragged size × rank count.Scope notes
entry.cpp.inenum ReduceOp { kSum = 0 };TPUT<AtomicType::AtomicAdd>cannot Prod). Prod on the HOST ring requires the atomic-add → non-atomic-push algorithm change, tracked separately. The ragged ST mirrors the InCore ring ragged ST, which is also Sum-only.numel % NR == 0the floor boundaries coincide exactly with the oldchunk_elemsmultiples, sosend_base/send_end/send_len/max_chunk_elems/tile_colsare identical and the generated kernel is unchanged (the only removed code is thenumel % nranks != 0guard branch, which never fired on the supported path).pre-commit run --all-filesreports 4 pyright errors inpython/pypto/runtime/distributed_runner.py(DataType/AccessMode/BackendKindunknown import symbol) that reproduce identically on a pristineorigin/mainworktree @f1a2f35f— an environment artifact of the isolated hook env (missing_task_interfaceC-extension the installedsimplerre-exports from), not caused by this PR. pyright passes on all files changed here.