Scope hybrid pre-collective failure sync over both children's meshes - #74
Conversation
GefenMuonHybrid.step derived its pre-collective failure-sync scope from the Muon child alone, while the composite preflight (closure + structural grad validation + AMP controls) validates BOTH children. When the backup half owns a sharded mesh the Muon half does not touch (a sharded backup weight with a non-sharded or absent Muon half), a one-rank preflight failure on that backup-only mesh was never all-reduced across it: the failing rank raised while its mesh peers stepped and mutated their backup shard, diverging cross-rank state. Derive the scope from the UNION of both children's sharded meshes. Extract the dedup/order logic into GefenMuon._collect_sharded_failure_groups and fold both children's param_groups through one deduped, sorted scan, so the standard fully_shard case (both halves on the same mesh) collapses to exactly the Muon-only scope with no extra collective, while a backup-only mesh is folded in with one deterministic cross-rank order. A plain-torch backup (AdamW) with no sharded_mode groups contributes only its DTensor params, if any. Add a Gloo CPU regression test: a hybrid whose backup owns the sharded mesh and whose Muon half is non-sharded; a rank-0 closure failure must fan out so all ranks fail fast with no mutation and no hang. Fails before this change (peer rank steps and mutates), passes after.
|
Warning Review limit reachedYou’ve reached a temporary PR review limit under our Fair Usage Limits Policy. Next review available in: 47 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 (3)
📝 WalkthroughWalkthroughGefenMuon centralizes sharded failure-group discovery. GefenMuonHybrid now collects groups from both Muon and backup optimizers, and tests cover synchronization when only the backup optimizer owns a sharded mesh. ChangesHybrid sharded failure synchronization
Estimated code review effort: 3 (Moderate) | ~20 minutes Possibly related issues
Possibly related PRs
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
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: cbe2681f29
ℹ️ 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".
|
@coderabbitai review |
✅ Action performedReview finished.
|
A backup-only GefenMuonHybrid (self.muon is None) whose backup owns a sharded mesh skipped the pre-collective failure sync entirely: step() gated both the error sync and the AMP agreement on self.muon being present, so a one-rank preflight/closure failure raised only on the failing rank while its mesh peer stepped and mutated its backup shard (cross-rank state divergence). _synchronize_sharded_step_error and _prepare_synchronized_amp_step touch no GefenMuon instance state (they only reach already-static helpers), so make both @staticmethod and call them unconditionally from the hybrid over the union scope. When the scope is empty the static error sync just re-raises any local error and the AMP agreement returns True / falls back to _amp_prepare_optimizer_step -- identical to the previous muon-absent branches, which are removed. Drop the now-unused _amp_prepare_optimizer_step import. Add a Gloo CPU regression test: a hybrid with every param routed to a sharded backup (muon is None); a rank-0 closure failure must fan out so all ranks fail fast with no mutation and no hang. Fails before this change (peer steps and mutates), passes after.
Reconcile the hybrid pre-collective failure sync with #74 (merged to main): gefen_muon.py auto-merges (the sync helpers are now @staticmethod and the mesh scope is collected via _collect_sharded_failure_groups). In GefenMuonHybrid.step, keep the convention's codebook-scope binding branch (_synchronize_prevalidated_codebook_scope_failure) unchanged, and in the no-binding branch adopt #74's static UNION-scope sync (GefenMuon._synchronize_sharded_step_error / _prepare_synchronized_amp_step over self._step_failure_process_groups()), which also covers the muon=None backup-only case. Also picks up the v0.4.1 release bump.
1 participant
Fixes a cross-rank correctness bug in
GefenMuonHybrid's pre-collective failure synchronization, surfaced by a CodeRabbit review on the convention PR (#67) but rooted in the Tier-2 code already onmain.The bug — silent cross-rank state divergence (not a hang)
GefenMuonHybridderived its failure-sync process-group scope from the Muon child only (hybrid.pyself.muon._step_failure_process_groups()), while its composite step preflight validates both children. When the backup child owns a sharded mesh that the Muon child does not participate in, a preflight/closure failure on one rank of that backup-only mesh is not all-reduced across it — so ranks disagree on whether the step failed.The original review framed this as a deadlock. Adversarial verification showed a hang is not reachable: the backup child is only ever
Gefenortorch.optim.AdamW, and neither performs a post-preflight sharded collective that a stranded peer would block on. What actually happens is worse in a quieter way — reproduced on a 2-rank gloo hybrid (non-sharded Muon weight + sharded backup, rank-0 closure failure): rank 0 raises and rolls back while rank 1 sees no error and mutates its backup shard, leaving the two ranks with divergent optimizer state and no error surfaced.The fix
gefen_muon.py: extracted the sharded-mesh dedup/order logic into a reusable_collect_sharded_failure_groups(param_groups)staticmethod (byte-identical to the previous_step_failure_process_groupsbody; keyed ongroup.get("sharded_mode")so a foreign backup's groups are handled)._step_failure_process_groupsnow delegates to it; the Muon path is unchanged.hybrid.py:GefenMuonHybrid._step_failure_process_groupsnow derives the scope from the union of both children's param-groups through a single deduped, deterministically-ordered scan — a shared mesh yields one collective, a backup-only mesh is included, and there is exactly one cross-rank ordering (no per-child concatenation that could break 2D-mesh row/column lock order).@torch._dynamo.disableand the capture guard mirror the Muon path.Under standard
fully_shard(both halves shard on the same mesh) the dedup collapses to the Muon-only scope, so this is a verified no-op in the common topology; it only changes behavior in the asymmetric-mesh case that was previously unsynchronized.Test
tests/test_precollective_failure_sync.py: newtest_backup_only_mesh_preflight_failure_is_synchronized(backup owns the sharded mesh, Muon non-sharded, rank-0 closure fails) asserts both ranks fail fast and neither mutates, with bounded no-hang deadlines. Verified it fails without the fix (rank 1 reports no error and mutates state) and passes with it.Validation (CPU / gloo)
test_precollective_failure_sync.py: 4 passedNo failures.
Known limitation (out of scope)
A degenerate backup-only hybrid (
muon=None) still does not synchronize this path — the sync helpers are GefenMuon-instance-bound, so covering it would require promoting them to module-level functions. That single-child case is outside the composite-hybrid scenario this fixes and is left as a follow-up.Summary by CodeRabbit
Bug Fixes
Tests