Synchronize pre-collective step failures across the process group - #71
Conversation
|
Warning Review limit reachedYou’ve reached a temporary PR review limit under our Fair Usage Limits Policy. Next review available in: 16 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 (1)
📝 WalkthroughWalkthroughGefen adds distributed preflight and AMP-control synchronization for Muon and hybrid optimizer steps, including backend-aware reductions, coordinated failure propagation, synchronized GradScaler decisions, and distributed Gloo coverage for closure and AMP-control scenarios. ChangesDistributed AMP step synchronization
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant Caller
participant GefenMuonHybrid
participant GefenMuon
participant ProcessGroup
participant GradScalerControls
Caller->>GefenMuonHybrid: step()
GefenMuonHybrid->>GefenMuon: select process groups and preflight
GefenMuon->>ProcessGroup: synchronize failures and AMP control ranges
ProcessGroup-->>GefenMuon: collective step decision
GefenMuon->>GradScalerControls: parse and unscale controls
GefenMuon-->>GefenMuonHybrid: proceed or skip step
GefenMuonHybrid-->>Caller: return loss or completed step
Possibly related PRs
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 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.
🧹 Nitpick comments (1)
src/gefen/gefen_muon.py (1)
1389-1509: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winMissing
@torch._dynamo.disableon the new host-branching sync wrappers.
_synchronize_sharded_step_flag,_synchronize_sharded_step_error,_synchronize_sharded_step_control_range, and_prepare_synchronized_amp_stepall perform Python-level branching/raising on values derived fromdist.all_reduceresults, yet none carries@torch._dynamo.disable— unlike the sibling_step_failure_process_groups(line 1308) and the pre-existing_assert_sharded_grad_presence_consistent(line 1510), which both use it for the same kind of host-branching-around-collectives pattern. Rawdist.all_reduce(tensor, group=process_group)calls are a known Dynamo pain point (frequently causing graph breaks orUnsupportederrors when traced), so undertorch.compile+capturable=True(a combination this class's docstring explicitly advertises), these undecorated wrappers risk extra graph breaks/recompiles compared to the established pattern in this same file.Consider adding
@torch._dynamo.disableto these four methods for consistency with_step_failure_process_groups/_assert_sharded_grad_presence_consistent.♻️ Proposed fix
+ `@torch._dynamo.disable` `@staticmethod` def _synchronize_sharded_step_flag(local_value, process_groups) -> bool: ... + `@torch._dynamo.disable` def _synchronize_sharded_step_error( self, error, phase: str, process_groups ) -> None: ... + `@torch._dynamo.disable` `@staticmethod` def _synchronize_sharded_step_control_range(local_control, process_groups): ... + `@torch._dynamo.disable` def _prepare_synchronized_amp_step(self, optimizer, process_groups) -> bool: ...🤖 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 1389 - 1509, Add `@torch._dynamo.disable` to _synchronize_sharded_step_flag, _synchronize_sharded_step_error, _synchronize_sharded_step_control_range, and _prepare_synchronized_amp_step, matching the existing decoration pattern used by _step_failure_process_groups and _assert_sharded_grad_presence_consistent. Preserve their current synchronization, branching, and error behavior.
🤖 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.
Nitpick comments:
In `@src/gefen/gefen_muon.py`:
- Around line 1389-1509: Add `@torch._dynamo.disable` to
_synchronize_sharded_step_flag, _synchronize_sharded_step_error,
_synchronize_sharded_step_control_range, and _prepare_synchronized_amp_step,
matching the existing decoration pattern used by _step_failure_process_groups
and _assert_sharded_grad_presence_consistent. Preserve their current
synchronization, branching, and error behavior.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: 606cadce-c148-404a-b3fd-6595a647feb9
📒 Files selected for processing (4)
src/gefen/gefen.pysrc/gefen/gefen_muon.pysrc/gefen/hybrid.pytests/test_precollective_failure_sync.py
_synchronize_sharded_step_flag, _synchronize_sharded_step_error, _synchronize_sharded_step_control_range, and _prepare_synchronized_amp_step all branch/raise on values derived from dist.all_reduce, so they must not be traced -- matching the existing @torch._dynamo.disable on the sibling _step_failure_process_groups and _assert_sharded_grad_presence_consistent. Under torch.compile + capturable (advertised in the class docstring) the raw all_reduce + host branching otherwise forces avoidable graph breaks/recompiles. Tracing-only change; no runtime behavior change (precollective test still 3/3).
Match the mainline fix (#71): the sharded sync wrappers (_synchronize_sharded_step_flag/_error/_control_range, _prepare_synchronized_amp_step) and their codebook-scope analogs (_synchronize_codebook_scope_failure, _prepare_scoped_amp_optimizer_step) all branch/raise on dist.all_reduce results, so they carry @torch._dynamo.disable like the sibling _step_failure_process_groups and _assert_sharded_grad_presence_consistent. Tracing-only; no runtime change (precollective + scoped-agreement tests still 8/8).
…vice
Two distributed-correctness fixes in _step_failure_process_groups (found in
review):
* Deadlock on overlapping meshes: the old logic selected a single mesh, or
abstained (returned ()) only when a rank saw multiple non-enclosing meshes.
With overlapping meshes (e.g. members {0,1,2} and {2,3}) a rank in both
abstained while a rank in only one entered that mesh's all-reduce, so the
shared rank never joined and the control collective hung. Sync over EVERY
sharded mesh the rank participates in, deduplicated by process-group name and
ordered consistently -- exactly the meshes _assert_sharded_grad_presence_-
consistent already enters, so participation is symmetric and cannot desync.
* Wrong flag device: the flag used the ambient torch.cuda.current_device(),
which can differ from the parameter's local shard device and target the wrong
GPU. Derive the device from the shard (as the grad-presence preflight does)
and thread it through as _synchronize_step_failure/_control_range's optional
collective_device.
test_precollective_failure_sync + muon distributed suite stay green (466 passed).
![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: 58ca39d5a8
ℹ️ 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".
The previous rewrite flattened every sharded mesh's groups and sorted them by process-group name. For a multi-dim (HSDP/TP) mesh whose group names do not lexically follow the dimension order, that could make one rank enter a row all-reduce while a peer is blocked in its column all-reduce -- a lock-ordering deadlock in the preflight itself. Mirror _assert_sharded_grad_presence_- consistent exactly instead: dedup meshes by content key, iterate them in sorted-key order, and within each mesh keep the get_all_groups() dimension order. Each group still carries the local shard device. (Single pass per group, like the grad-presence preflight -- deeply overlapping non-enclosing meshes keep that preflight's existing cross-mesh limitation.) Full suite with GPU (2x3090 NCCL): 761 passed.
…ier-2 convention side) (#72) * Preserve legacy vmean loads on the rank-local sharded path (#70) * Preserve legacy vmean loads on the rank-local sharded path _unwrap_rank_local_sharded_checkpoint validated the payload with the strict default (allow_legacy_vmean_counter=False), so a rank-local (DTensor/FSDP) checkpoint carrying vmean without the separate vmean_step counter -- a pre-counter state the step-time resume path backfills from step -- was rejected at load. The native load path already opts into that tolerance (allow_legacy_vmean_counter=True); the rank-local path did not, so the two disagreed and an otherwise valid legacy resume failed only on the sharded path. This regressed against pre-load-atomicity behavior, which had no such check. Pass allow_legacy_vmean_counter=True at the rank-local call site to match the native path. This only relaxes the vmean-without-vmean_step case; a current checkpoint (which carries vmean_step) is unaffected, and the inverse vmean_step-without-vmean corruption check is unchanged. tests/test_cpu_step_checkpoint.py::test_rank_local_validator_tolerates_legacy_vmean_without_step pins the validator tolerance both ways (the strict default rejects the pre-counter payload; the tolerance the rank-local path now opts into accepts it). * Guard the rank-local unwrap path directly, not just the validator Add test_rank_local_unwrap_tolerates_and_backfills_legacy_vmean: a single-rank gloo world builds a rank-local (rank_local_dtensor_v2) checkpoint, drops the separate vmean_step counter to simulate a pre-counter state, loads it through _unwrap_rank_local_sharded_checkpoint (the call site the fix touches), and asserts the first resumed step backfills vmean_step from step. The existing validator test calls _validate_rank_local_states directly, so it stays green if the unwrap call site regresses to strict validation. This test fails without allow_legacy_vmean_counter=True on the rank-local path (verified: reverting the fix raises "block second moment is missing vmean_step" at load). CPU-only, no multiprocessing (~2s). * Synchronize pre-collective Muon step failures * Disable Dynamo tracing on the pre-collective sync wrappers Match the mainline fix (#71): the sharded sync wrappers (_synchronize_sharded_step_flag/_error/_control_range, _prepare_synchronized_amp_step) and their codebook-scope analogs (_synchronize_codebook_scope_failure, _prepare_scoped_amp_optimizer_step) all branch/raise on dist.all_reduce results, so they carry @torch._dynamo.disable like the sibling _step_failure_process_groups and _assert_sharded_grad_presence_consistent. Tracing-only; no runtime change (precollective + scoped-agreement tests still 8/8). * Close pre-collective sync gaps on the convention side Four distributed-correctness fixes (found in review), on top of re-pointing the codebook scope onto the mainline primitive: * Overlapping-mesh deadlock + wrong flag device: same _step_failure_process_- groups rewrite as the mainline PR -- sync over every participated mesh (deduped, ordered) with the shard device, mirroring _assert_sharded_grad_presence_consistent. Threads collective_device through _synchronize_step_control_range too. * Scoped Gefen AMP presence hang: plain Gefen.step gated _prepare_scoped_amp_optimizer_step() on local found_inf/grad_scale, so with a multi-member codebook scope a rank with those attributes entered the presence all-gather while a rank without them skipped it -- deadlock. Run the scoped AMP agreement on EVERY member of a multi-member scope, matching the unconditional Muon/Hybrid preflight. * Dropped post-closure binding recheck (the CI failure): the merged Tier-2 step rewrite lost the convention's _assert_finalized_binding_layout() recheck after closure() on the mesh (non-scoped) branch, so a closure that rebinds param_groups was not caught -- test_finalized_layout_guard_rechecks_after_- closure[muon] failed. Re-assert it inside the preflight try (caught + synchronized) so a closure rebind fails before mutation and on every member. Full convention suite: 1095 passed, 0 failed. * Preserve mesh dimension order + close hybrid preflight gaps Three more review fixes: * DeviceMesh dimension order (same as the mainline PR): mirror _assert_sharded_grad_presence_consistent -- dedup meshes by content key, sorted-key order, get_all_groups() dimension order per mesh -- instead of flattening and sorting all groups by name, which could reorder a 2-D HSDP/TP mesh's row/column groups per rank and lock-order deadlock the preflight. * Hybrid post-closure binding recheck (unscoped path): the binding-is-None branch stepped children without re-asserting _assert_finalized_binding_layout() after the closure/pre-hooks, so a rank-local closure swapping a same-shaped backup parameter passed the gradient scan and stepped under stale routing. The scoped path already rechecked; re-assert in the unscoped preflight try (caught + synchronized). * Hybrid capture readiness: the _gefen_hybrid_precollective_preflight marker suppresses the muon child's own _assert_codebook_capture_ready() guard, and the hybrid preamble only checked devices, so a hybrid capturing before its Muon codebook initialized could run host-driven codebook init during capture. Call the muon capture/codebook readiness guard in the hybrid preflight. Full convention suite with GPU (2x3090 NCCL): 1420 passed.
1 participant
What & why
In sharded/FSDP2 runs,
GefenMuon.step()andGefenMuonHybrid.step()ranclosure()and the AMP overflow-skip before the collective section — the codebookfull_tensor()all-gather in exact/distributed modes, and the Parallel-Muondist.all_reduce/broadcastin the distributed step. Base already synchronized gradient presence, but not:closure()that raises on one rank (e.g. a data-loader or loss error on a single rank), orfound_infthat differs across ranks (each rank scans its own FSDP shard, so overflow is genuinely rank-local).Either one makes a rank exit early while its peers enter the next collective → indefinite hang / NCCL timeout.
The fix
Synchronize failure across the process group before entering any collective, so every rank raises or skips symmetrically, without introducing a new hang and preserving fail-before-mutation.
_synchronize_step_failure(local_failed, process_group)— all-reduces anint32flag withReduceOp.MAXover a supplied process group (mesh /dist.group.WORLD), guarded onis_initialized()/world_size < 2._step_failure_collective_deviceplaces the flag on the right device per backend (nccl→cuda, gloo/default→cpu, honoringbound_device_id)._synchronize_step_control_rangeagrees min/max control bounds in one collective. The primitive takes the PG as an argument so other scopes (e.g. the convention's codebook-scope path) can re-point onto it rather than keep a divergent copy.closure(), the capturable check, gradient-structure validation, and the AMPfound_inf/grad_scaleparse are wrapped and synchronized before work partitioning, codebook learning, or any state mutation._step_failure_process_groupsselects a control group that encloses every upcoming collective participant (or abstains, returning()), so the sync itself cannot deadlock. CUDA-graph capture returns()(host-readable flags can't be captured; captured steps require an eager warmup)._prepare_synchronized_amp_stepagrees on(present, overflow, scale_present, scale_value)across ranks before any unscale, raising symmetric errors on rank-divergent scaler config._gefen_hybrid_precollective_preflightmarker — no double-collective.Tests
tests/test_precollective_failure_sync.py(gloo, world=2,muon_exact/muon_distributed/hybrid_exact): a one-rankclosure()raise and a divergentfound_infboth make all ranks fail fast (per-step liveness deadline + no-hang assertion) with symmetric messages, and leave optimizer/param state byte-for-byte unchanged. Would catch a re-introduced hang (removing the sync blocks rank 1 in the collective past the deadline). 3/3 pass.distributed_checkpoint_safety,grad_presence,capturable_fsdp2,fsdp2_parity,fused_multirankall pass.Scope
Sharded / FSDP2 (
sharded_modeexact or distributed) users. This is a standalone extraction ontomain; the primitive is factored so the convention's codebook-scope failure sync re-points onto it (paired convention PR).Summary by CodeRabbit