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
@@ -1381,15 +1400,16 @@ def _synchronize_sharded_step_flag(local_value, process_groups) -> bool:
             )
         return synchronized
 
+    @staticmethod
     @torch._dynamo.disable
     def _synchronize_sharded_step_error(
-        self, error, phase: str, process_groups
+        error, phase: str, process_groups
     ) -> None:
         if not process_groups:
             if error is not None:
                 raise error
             return
-        failed = self._synchronize_sharded_step_flag(
+        failed = GefenMuon._synchronize_sharded_step_flag(
             error is not None, process_groups
         )
         if not failed:
@@ -1420,8 +1440,9 @@ def _synchronize_sharded_step_control_range(local_control, process_groups):
             )
         return minimum, maximum
 
+    @staticmethod
     @torch._dynamo.disable
-    def _prepare_synchronized_amp_step(self, optimizer, process_groups) -> bool:
+    def _prepare_synchronized_amp_step(optimizer, process_groups) -> bool:
         """Agree on AMP controls before unscaling or entering Muon collectives."""
         local_present = hasattr(optimizer, "found_inf") or hasattr(
             optimizer, "grad_scale"
@@ -1447,11 +1468,11 @@ def _prepare_synchronized_amp_step(self, optimizer, process_groups) -> bool:
             local_scale_present = False
             scale_value = 0.0
             local_amp_error = exc
-        self._synchronize_sharded_step_error(
+        GefenMuon._synchronize_sharded_step_error(
             local_amp_error, "AMP control preflight", process_groups
         )
 
-        minimum, maximum = self._synchronize_sharded_step_control_range(
+        minimum, maximum = GefenMuon._synchronize_sharded_step_control_range(
             (
                 int(local_present),
                 int(local_overflow),
@@ -1491,7 +1512,7 @@ def _prepare_synchronized_amp_step(self, optimizer, process_groups) -> bool:
         except Exception as exc:
             should_step = False
             local_amp_error = exc
-        self._synchronize_sharded_step_error(
+        GefenMuon._synchronize_sharded_step_error(
             local_amp_error, "AMP preparation", process_groups
         )
         return should_step

src/gefen/hybrid.py

@@ -62,7 +62,6 @@
 from gefen.gefen import (
     Gefen,
     _amp_native_scaling_required,
-    _amp_prepare_optimizer_step,
     _assert_optimizer_gradients_structurally_valid,
 )
 from gefen.gefen_muon import GefenMuon
@@ -655,12 +654,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
@@ -698,25 +729,25 @@ def step(self, closure=None):
         except Exception as exc:
             loss = None
             local_preflight_error = exc
-        if self.muon is not None:
-            self.muon._synchronize_sharded_step_error(
-                local_preflight_error, "hybrid step preflight", process_groups
-            )
-        elif local_preflight_error is not None:
-            raise local_preflight_error
+        # Synchronize the preflight failure across the UNION scope
+        # unconditionally -- including a muon-only-None (backup-only) hybrid
+        # whose backup owns a sharded mesh. When process_groups is empty (no
+        # sharded mesh) the static helper just re-raises any local error, which
+        # is the previous rank-local behavior.
+        GefenMuon._synchronize_sharded_step_error(
+            local_preflight_error, "hybrid step preflight", process_groups
+        )
 
         # A non-finite gradient in either half skips BOTH children before their
         # codebooks, states, counters, or parameters can move. Explicit
         # scaler.unscale_(hybrid) is detected by grad_scale=None and is not
         # repeated; automatic unscale covers every child parameter exactly once.
-        if self.muon is not None:
-            should_step = self.muon._prepare_synchronized_amp_step(
-                self, process_groups
-            )
-        elif hasattr(self, "found_inf") or hasattr(self, "grad_scale"):
-            should_step = _amp_prepare_optimizer_step(self)
-        else:
-            should_step = True
+        # The static AMP agreement subsumes the old muon-present/absent split:
+        # with an empty scope and no local controls it returns True, and with
+        # local controls it falls back to plain _amp_prepare_optimizer_step.
+        should_step = GefenMuon._prepare_synchronized_amp_step(
+            self, process_groups
+        )
         if not should_step:
             for post_hook in self._optimizer_step_post_hooks.values():
                 post_hook(self, args, kwargs)