Enforce cross-rank parameter order in the Muon preflight - #66
Conversation
The grad-presence preflight tolerated rank-divergent parameter-group registration order after its content-sort hardening, but exact and distributed stepping enter gradient collectives and assign momentum owners by insertion index, so divergent order with consistent gradient presence could still misalign those collectives. Reduce each sorted parameter's local insertion position to its mesh-wide max and min and fail with the same error on every rank before any mutation. Also complete the release-gate script's --help output (the header outgrew the hardcoded sed range), resolve --wheel paths against the invocation directory, and reject a missing --wheel argument with a usage error instead of a raw set -u failure.
|
Warning Review limit reachedYou’ve reached a temporary PR review limit under our Fair Usage Limits Policy. Next review available in: 23 minutes Enable usage-based reviews in Billing to review now. Otherwise, wait until the next included review is available. How can I continue?After more reviews become available, a review can be triggered using the To avoid repeated limits, reduce automatic review volume by pausing incremental auto-reviews earlier, using label-based review opt-in, excluding WIP or generated PR titles, or requesting reviews manually when the PR is ready. If your team needs uninterrupted high-volume reviews, an organization admin can enable usage-based reviews. How do review limits work?CodeRabbit enforces per-developer PR review limits for each organization. Most developers receive the normal plan review availability. For paid Pro and Pro+ PR reviews, CodeRabbit uses adaptive limits for sustained high-volume activity. When a developer's recent PR review activity reaches the 95th percentile or higher among CodeRabbit users, additional reviews become available more gradually as earlier reviews age out of the rolling window. Please refer docs for additional details. Review details⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (4)
📝 WalkthroughWalkthroughThe change adds cross-rank parameter-group order validation before distributed Muon collectives, expands regression coverage, and hardens release GPU gate help and wheel-path handling. ChangesDistributed Muon preflight
GPU release gate
Estimated code review effort: 3 (Moderate) | ~20 minutes Sequence Diagram(s)sequenceDiagram
participant Rank0
participant Rank1
participant Preflight as GefenMuon._assert_sharded_grad_presence_consistent
participant ProcessGroup
Rank0->>Preflight: Build local parameter order probe
Rank1->>Preflight: Build local parameter order probe
Preflight->>ProcessGroup: all_reduce order probe
ProcessGroup-->>Preflight: Reduced cross-rank values
Preflight-->>Rank0: Reject divergent order
Preflight-->>Rank1: Reject divergent order
Possibly related PRs
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
src/gefen/gefen_muon.py (1)
1376-1409: 🩺 Stability & Availability | 🔴 Critical | 🏗️ Heavy liftReject bare or non-uniquely named distributed parameters before this order probe.
This aligns ranks by
name, shape, and dtype, but bare parameters receive positional auto-names. If same-shapeda, bare registered as[a, b]on one rank and[b, a]on another, both ranks still produce positions0, 1for the same auto-names, so this check passes while laterfull_tensor()/owner collectives follow divergent orders. Require globally unique explicit names for exact/distributed DTensor groups (or provide another stable cross-rank ID), and add a reversed bare-parameter regression case.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/gefen/gefen_muon.py` around lines 1376 - 1409, Before constructing order_probe in the distributed parameter-group path, reject bare parameters and any non-unique auto-generated names for exact/distributed DTensor groups. Require globally unique explicit parameter names, or another stable cross-rank identifier, so identical-shaped parameters cannot be aligned only by positional auto-names; preserve the existing name/shape/dtype alignment for valid groups and add coverage for reversed bare-parameter registration order.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@scripts/release_gpu_gate.sh`:
- Around line 53-57: Update the --wheel argument handling in release_gpu_gate.sh
to reject a missing value when the next token begins with “-”, including
option-like values such as --help, before calling realpath. Preserve accepting
valid wheel paths and the existing usage-error exit behavior for missing or
invalid --wheel arguments.
---
Outside diff comments:
In `@src/gefen/gefen_muon.py`:
- Around line 1376-1409: Before constructing order_probe in the distributed
parameter-group path, reject bare parameters and any non-unique auto-generated
names for exact/distributed DTensor groups. Require globally unique explicit
parameter names, or another stable cross-rank identifier, so identical-shaped
parameters cannot be aligned only by positional auto-names; preserve the
existing name/shape/dtype alignment for valid groups and add coverage for
reversed bare-parameter registration order.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: 7704b347-bdaf-4604-b640-9727d6228bfa
📒 Files selected for processing (4)
CHANGELOG.mdscripts/release_gpu_gate.shsrc/gefen/gefen_muon.pytests/test_muon_grad_presence.py
![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: ef58f58b37
ℹ️ About Codex in GitHub
Your team has set up Codex to 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 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
Use optimizer-wide insertion positions in the order probe so swaps between parameters on different DeviceMeshes are caught (step() enters per-mesh collectives in global registration order). Fail closed when parameters in one mesh share a name, shape, and dtype: construction accepts duplicate explicit names, and such parameters cannot be identified across ranks by any rank-local property. Same-shaped bare parameters have positionally auto-generated names with the same ambiguity; warn once and recommend named construction instead of breaking the common identically-constructed case. The gate script also rejects an option-like value after --wheel instead of treating it as a path.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 83c29eec4b
ℹ️ About Codex in GitHub
Your team has set up Codex to 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 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
Ride the duplicate-label count on the MAX-reduced order probe so a rank with clean labels learns about duplicates on a peer and raises the same error instead of blocking in a collective the peer abandoned, and defer duplicate/order/presence raises until every local mesh has completed its probe collectives so overlapping but non-identical meshes cannot be stranded mid-preflight by an early exit.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: eed3554403
ℹ️ About Codex in GitHub
Your team has set up Codex to 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 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| if order_mismatches: | ||
| raise RuntimeError( |
There was a problem hiding this comment.
When DeviceMeshes overlap but are not identical, an order mismatch on only one mesh still is not propagated to ranks that only share a later clean mesh. For example, with meshes over ranks [0,1] and [1,2], ranks 0/1 raise here for an order mismatch on [0,1] after all probes complete, but rank 2 never records order_mismatches and proceeds into _maybe_refresh_gefen_codebook()/step() collectives for [1,2] after rank 1 has exited. This preserves the unmatched-collective failure the preflight is meant to avoid; the error flag needs to reach overlapping peers before any post-preflight collectives run.
Useful? React with 👍 / 👎.
| collisions = {} | ||
| for name, p, _ in items: | ||
| if _AUTO_PARAM_NAME.match(str(name)): | ||
| key = (tuple(p.shape), str(p.dtype)) |
There was a problem hiding this comment.
When a DTensor optimizer is built from bare parameters and two ranks register [a, b] vs [b, a], the generated param_0/param_1 labels follow local position rather than tensor identity. This collision map only warns when multiple auto-named params share the same (shape, dtype), so differently shaped or differently typed bare params produce no warning/error; the order probe then compares identical local positions for param_0/param_1 and passes, leaving later full_tensor()/distributed collectives to run for different tensors on each rank. The auto-name path needs to fail closed or otherwise avoid trusting positional labels even when shapes differ.
Useful? React with 👍 / 👎.
1 participant
Summary
Follow-ups from the final pre-merge review sweep of #65 (the sweep found no critical or major issues; these are its three minor findings):
step()still assigns momentum owners and enters gradient collectives by insertion index, so rank-divergent registration order with consistent gradient presence would misalign the step. Each sorted parameter's local insertion position is now reduced to its mesh-wide max/min (one extra small all-reduce per mesh dimension on the already-guarded exact/distributed path only) and any divergence raises the same error on every rank before mutation.scripts/release_gpu_gate.shfixes:--helpprints the full header again (it had outgrown a hardcodedsedrange and silently dropped theGEFEN_GATE_EXTRA_SITE/CUDA_VISIBLE_DEVICESdocs); a relative--wheelpath now resolves against the invocation directory instead of breaking after the script'scd; a missing--wheelargument exits with a usage error instead of a rawset -ufailure.Validation
reversed_consistentmust now raise the order error) fails on the previous code and passes with the fix.bash -nclean; full--helpverified; both--wheelerror paths exercised from a foreign working directory.🤖 Generated with Claude Code
Summary by CodeRabbit
Bug Fixes
Chores