diff --git a/src/gefen/gefen_muon.py b/src/gefen/gefen_muon.py index 0d275b1..45d4e9c 100644 --- a/src/gefen/gefen_muon.py +++ b/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 diff --git a/src/gefen/hybrid.py b/src/gefen/hybrid.py index 0e2fc2d..eb0d837 100644 --- a/src/gefen/hybrid.py +++ b/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) diff --git a/tests/test_precollective_failure_sync.py b/tests/test_precollective_failure_sync.py index 91ff3b7..a8e0989 100644 --- a/tests/test_precollective_failure_sync.py +++ b/tests/test_precollective_failure_sync.py @@ -136,7 +136,52 @@ 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_only_none": + # Degenerate composite: EVERY param routes to the backup, so + # ``optimizer.muon is None``. The backup still owns the sharded mesh, so + # the hybrid's failure sync must run even with no Muon child -- otherwise + # a one-rank preflight failure raises only on the failing rank while its + # mesh peer steps and mutates its backup shard. + optimizer = GefenMuonHybrid( + [], + [("backup.weight", weight)], + lr=1e-3, + fused=False, + ns_steps=1, + ns_schedule="standard", + sharded_mode="exact", + backup_optimizer="gefen", + normuon=False, + ) + assert optimizer.muon is None + parameters = (weight,) + elif 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 +337,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 +376,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 +475,82 @@ 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 + + +@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_hybrid_none_muon_preflight_failure_is_synchronized(): + """A backup-only hybrid (``muon is None``) still synchronizes preflight. + + Regression for the composite skipping the failure sync entirely when there + is no Muon child: when every param routes to a SHARDED backup half, a + rank-0 preflight failure must still fan out across the backup mesh so ALL + ranks fail fast with no mutation -- rather than rank 0 raising while its mesh + peer silently steps and mutates its backup shard. Bounded deadlines guard + against a hang. + """ + results = _run_distributed_case( + "hybrid_backup_only_none", 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] + 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 + 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