Scope hybrid pre-collective failure sync over both children's meshes

src/gefen/gefen_muon.py

@@ -1305,6 +1305,57 @@ def _dist_available() -> bool:
             return False
         return torch.distributed.is_initialized()
 
+    @staticmethod
+    @torch._dynamo.disable
+    def _collect_sharded_failure_groups(param_groups):
+        """Deduped, deterministically-ordered (process_group, device) pairs.
+
+        Shared by the raw ``GefenMuon`` step scope and ``GefenMuonHybrid``'s
+        union scope. Scans every sharded (non-``approx``) mesh represented in
+        ``param_groups``, keyed and deduplicated by mesh content and iterated in
+        sorted-key order, and within each mesh in its ``get_all_groups()``
+        dimension order. Folding several optimizers' groups through ONE dict
+        keeps a single deterministic order across ranks: an identical mesh owned
+        by more than one child (the standard fully_shard case) collapses to a
+        single collective, and a child-only mesh is still folded in -- never a
+        per-child concatenation that could interleave one rank's row all-reduce
+        against a peer's column all-reduce. Groups with no ``sharded_mode`` key
+        (a foreign backup's conventional torch groups) count as non-approx.
+        """
+        import torch.distributed as dist
+
+        by_mesh = {}
+        for group in param_groups:
+            if group.get("sharded_mode") == "approx":
+                continue
+            for param in group["params"]:
+                if not GefenMuon._is_sharded(param):
+                    continue
+                mesh = param.device_mesh
+                if mesh.get_coordinate() is None or mesh.size() < 2:
+                    continue
+                process_groups = tuple(mesh.get_all_groups())
+                members = tuple(
+                    int(item)
+                    for item in mesh.mesh.detach().cpu().reshape(-1).tolist()
+                )
+                key = (
+                    str(mesh.device_type),
+                    tuple(int(item) for item in mesh.shape),
+                    members,
+                    tuple(str(pg.group_name) for pg in process_groups),
+                )
+                by_mesh.setdefault(
+                    key, (GefenMuon._state_tensor_device(param), process_groups)
+                )
+        result = []
+        for key in sorted(by_mesh):
+            device, process_groups = by_mesh[key]
+            for process_group in process_groups:
+                if dist.get_world_size(process_group) > 1:
+                    result.append((process_group, device))
+        return tuple(result)
+
     @torch._dynamo.disable
     def _step_failure_process_groups(self):
         """Return the control scope for eager exact/distributed Muon steps.
@@ -1337,39 +1388,7 @@ def _step_failure_process_groups(self):
         ):
             return ()
 
-        import torch.distributed as dist
-
-        by_mesh = {}
-        for group in self.param_groups:
-            if group["sharded_mode"] == "approx":
-                continue
-            for param in group["params"]:
-                if not self._is_sharded(param):
-                    continue
-                mesh = param.device_mesh
-                if mesh.get_coordinate() is None or mesh.size() < 2:
-                    continue
-                process_groups = tuple(mesh.get_all_groups())
-                members = tuple(
-                    int(item)
-                    for item in mesh.mesh.detach().cpu().reshape(-1).tolist()
-                )
-                key = (
-                    str(mesh.device_type),
-                    tuple(int(item) for item in mesh.shape),
-                    members,
-                    tuple(str(pg.group_name) for pg in process_groups),
-                )
-                by_mesh.setdefault(
-                    key, (self._state_tensor_device(param), process_groups)
-                )
-        result = []
-        for key in sorted(by_mesh):
-            device, process_groups = by_mesh[key]
-            for process_group in process_groups:
-                if dist.get_world_size(process_group) > 1:
-                    result.append((process_group, device))
-        return tuple(result)
+        return self._collect_sharded_failure_groups(self.param_groups)
 
     @staticmethod
     @torch._dynamo.disable

src/gefen/hybrid.py

@@ -655,12 +655,44 @@ def _assert_capturable_devices_if_capturing(self) -> None:
                 )
             )
 
+    @torch._dynamo.disable
+    def _step_failure_process_groups(self):
+        """Failure-sync scope for the composite preflight: the UNION of both
+        children's sharded-mesh (process_group, device) pairs.
+
+        The composite preflight (closure + structural gradient validation + AMP
+        controls) covers BOTH children, so the pre-collective failure sync must
+        span every sharded mesh EITHER child owns -- not only the Muon child's.
+        Deriving the scope from the Muon child alone missed any mesh owned only
+        by the backup child (a sharded backup weight whose Muon half is
+        non-sharded or absent): a one-rank preflight failure on that backup-only
+        mesh then raised on the failing rank while its mesh peers proceeded to
+        step and mutate their shard, diverging cross-rank state. Both children's
+        param_groups are folded through ONE deduped, sorted scan
+        (``GefenMuon._collect_sharded_failure_groups``), so the standard
+        fully_shard case -- both halves sharded on the SAME mesh -- collapses to
+        exactly the Muon-only scope with no extra collective, while a
+        backup-only mesh is included in one deterministic cross-rank order.
+        """
+        if not GefenMuon._dist_available():
+            return ()
+        param_groups = [
+            group
+            for optimizer in self._subopts
+            for group in optimizer.param_groups
+        ]
+        params = [param for group in param_groups for param in group["params"]]
+        # The protocol returns host-readable flags and cannot be captured; a
+        # captured step already requires an eager warmup with fixed control flow
+        # (mirrors GefenMuon._step_failure_process_groups).
+        if any(param.device.type == "cuda" for param in params) and (
+            torch.cuda.is_current_stream_capturing()
+        ):
+            return ()
+        return GefenMuon._collect_sharded_failure_groups(param_groups)
+
     def step(self, closure=None):
-        process_groups = (
-            self.muon._step_failure_process_groups()
-            if self.muon is not None
-            else ()
-        )
+        process_groups = self._step_failure_process_groups()
         # Dispatch the INSTANCE step hooks around the composite step, mirroring
         # torch.optim.Optimizer.profile_hook_step exactly: hooks receive
         # (optimizer, args, kwargs) where args are the raw step() call args

tests/test_precollective_failure_sync.py

@@ -136,7 +136,33 @@ def _make_optimizer(mesh, case: str):
     full_weight_grad = torch.linspace(0.6, -0.7, 64, dtype=torch.float32).reshape(8, 8)
     weight = nn.Parameter(distribute_tensor(full_weight.clone(), mesh, [Shard(0)]))
 
-    if case == "hybrid_exact":
+    if case == "hybrid_backup_sharded":
+        # Invert the ownership of hybrid_exact: the Muon half holds a REPLICATED
+        # (non-sharded) 2D weight while the BACKUP half owns the sharded mesh.
+        # Deriving the failure-sync scope from the Muon child alone would then
+        # miss that backup-only mesh, so a one-rank preflight failure would raise
+        # on the failing rank while its mesh peers stepped and mutated their
+        # backup shard (cross-rank divergence). The hybrid's union scope must
+        # fold the backup mesh in so every rank fails fast instead.
+        muon_weight = nn.Parameter(
+            torch.linspace(-0.5, 0.6, 64, dtype=torch.float32).reshape(8, 8)
+        )
+        optimizer = GefenMuonHybrid(
+            [("muon.weight", muon_weight)],
+            [("backup.weight", weight)],
+            lr=1e-3,
+            fused=False,
+            ns_steps=1,
+            ns_schedule="standard",
+            sharded_mode="exact",
+            backup_optimizer="gefen",
+            normuon=False,
+        )
+        parameters = (muon_weight, weight)
+        muon_weight.grad = torch.linspace(
+            0.6, -0.7, 64, dtype=torch.float32
+        ).reshape(8, 8)
+    elif case == "hybrid_exact":
         bias = nn.Parameter(torch.linspace(-0.4, 0.3, 8, dtype=torch.float32))
         optimizer = GefenMuonHybrid(
             [("weight", weight)],
@@ -292,7 +318,37 @@ def _distributed_worker(rank: int, init_method: str, case: str, result_queue) ->
             dist.destroy_process_group()
 
 
-def _run_distributed_case(case: str):
+def _backup_sharded_worker(rank: int, init_method: str, case: str, result_queue) -> None:
+    """Worker for the backup-only-mesh case: only the closure-failure scenario.
+
+    Constructs the hybrid where the BACKUP half owns the sharded mesh and the
+    Muon half is non-sharded, induces a rank-0 closure failure, and reports
+    whether the failure fanned out to every mesh rank (fail-fast) with no
+    parameter/state mutation and no hang.
+    """
+    try:
+        torch.set_num_threads(1)
+        dist.init_process_group(
+            "gloo",
+            init_method=init_method,
+            rank=rank,
+            world_size=_WORLD_SIZE,
+            timeout=timedelta(seconds=_PROCESS_GROUP_TIMEOUT_SECONDS),
+        )
+        from torch.distributed.tensor import init_device_mesh
+
+        mesh = init_device_mesh("cpu", (_WORLD_SIZE,), mesh_dim_names=("dp",))
+        closure_failure = _closure_failure_result(rank, mesh, case)
+        dist.barrier()
+        result_queue.put({"rank": rank, "closure_failure": closure_failure})
+    except BaseException:
+        result_queue.put({"rank": rank, "fatal_error": traceback.format_exc()})
+    finally:
+        if dist.is_available() and dist.is_initialized():
+            dist.destroy_process_group()
+
+
+def _run_distributed_case(case: str, worker_target=_distributed_worker):
     context = torch.multiprocessing.get_context("spawn")
     result_queue = context.Queue()
     descriptor, rendezvous_path = tempfile.mkstemp(prefix="gefen-precollective-sync-")
@@ -301,7 +357,7 @@ def _run_distributed_case(case: str):
     init_method = Path(rendezvous_path).resolve().as_uri()
     processes = [
         context.Process(
-            target=_distributed_worker,
+            target=worker_target,
             args=(rank, init_method, case, result_queue),
         )
         for rank in range(_WORLD_SIZE)
@@ -400,3 +456,44 @@ def test_rank_local_closure_and_amp_controls_are_synchronized(case):
         assert [item["post_hook_calls"] for item in skip_results] == [1, 1]
     else:
         assert [item["post_hook_calls"] for item in skip_results] == [None, None]
+
+
+@pytest.mark.skipif(
+    not dist.is_available() or not dist.is_gloo_available(),
+    reason="pre-collective failure synchronization coverage requires Gloo",
+)
+def test_backup_only_mesh_preflight_failure_is_synchronized():
+    """A backup-owned sharded mesh the Muon half does not touch is in scope.
+
+    Regression for the composite failure-sync deriving its process-group scope
+    from the Muon child only: when the BACKUP half owns the sharded mesh (and
+    the Muon half is non-sharded), a rank-0 preflight failure must fan out to
+    every mesh rank so ALL ranks fail fast with no mutation -- rather than
+    rank 0 raising while its mesh peer silently steps and mutates its backup
+    shard. Also guards against a hang (bounded worker/step deadlines).
+    """
+    results = _run_distributed_case(
+        "hybrid_backup_sharded", worker_target=_backup_sharded_worker
+    )
+    assert all("fatal_error" not in result for result in results), results
+
+    closure_results = [result["closure_failure"] for result in results]
+    # Both ranks fail fast (the failure was synchronized across the backup mesh).
+    assert all(item["message"] is not None for item in closure_results), closure_results
+    assert (
+        "GefenMuon hybrid step preflight failed on local rank 0"
+        in closure_results[0]["message"]
+    ), closure_results
+    assert "rank-zero closure failure" in closure_results[0]["message"], closure_results
+    assert (
+        "GefenMuon hybrid step preflight failed on another process-group member"
+        in closure_results[1]["message"]
+    ), closure_results
+    # Liveness: neither rank hung, and no parameter/state/gradient moved on
+    # either rank (the failing rank AND its mesh peer both roll back).
+    assert all(
+        item["elapsed"] < _STEP_DEADLINE_SECONDS for item in closure_results
+    ), closure_results
+    assert all(
+        all(item["unchanged"].values()) for item in closure_results
+    ), closure_results