From d02874e1ba0bcc043de4ecdb39e62470a3a49149 Mon Sep 17 00:00:00 2001 From: thad0ctor Date: Tue, 14 Jul 2026 22:58:48 -0700 Subject: [PATCH 1/4] Synchronize pre-collective Muon step failures --- src/gefen/gefen.py | 125 +++++-- src/gefen/gefen_muon.py | 246 +++++++++++++- src/gefen/hybrid.py | 87 +++-- tests/test_precollective_failure_sync.py | 402 +++++++++++++++++++++++ 4 files changed, 800 insertions(+), 60 deletions(-) create mode 100644 tests/test_precollective_failure_sync.py diff --git a/src/gefen/gefen.py b/src/gefen/gefen.py index c76d67f..1a77ac3 100644 --- a/src/gefen/gefen.py +++ b/src/gefen/gefen.py @@ -497,22 +497,75 @@ def _resolve_find_period_backend(grad: torch.Tensor) -> str: return "cuda_kernel" if grad_work.device.type == "cuda" else "cpu" +def _step_failure_collective_device(process_group) -> torch.device: + """Choose a backend-compatible device for a process-group control flag.""" + import torch.distributed as dist + + group = process_group if process_group is not None else dist.group.WORLD + bound_device = getattr(group, "bound_device_id", None) + if bound_device is not None: + return torch.device(bound_device) + + backend = str(dist.get_backend(process_group)).lower() + if "nccl" in backend: + return torch.device("cuda", torch.cuda.current_device()) + if "xccl" in backend: + return torch.device("xpu", torch.xpu.current_device()) + # Gloo, MPI, UCC, and PyTorch's multi-backend default process group all + # accept CPU control tensors. Keeping their flag on CPU also avoids an + # unrelated accelerator initialization in spawned Gloo workers. + return torch.device("cpu") + + @torch.no_grad() -def _amp_prepare_optimizer_step(optimizer) -> bool: - """Honor PyTorch's native ``GradScaler`` optimizer-step protocol. +@torch._dynamo.disable +def _synchronize_step_failure(local_failed, process_group) -> bool: + """Return whether any process-group member reported a step failure.""" + failed = bool(local_failed) + if not torch.distributed.is_available() or not torch.distributed.is_initialized(): + return failed - ``GradScaler.step`` gives AMP-aware optimizers two temporary attributes: - ``found_inf`` contains the result of the non-finite scan and ``grad_scale`` - contains the scale that still needs to be removed. When users explicitly - call ``scaler.unscale_(optimizer)`` first, ``grad_scale`` is ``None`` and - the gradients must not be divided a second time. + import torch.distributed as dist - Return ``False`` on overflow before callers learn a codebook or mutate any - optimizer/parameter state. On a finite automatic-unscale step, multiply the - gradients in place by the same fp32 reciprocal used by ``GradScaler``. This - keeps every existing fused/unfused update path operating on ordinary - unscaled gradients and also works for local DTensor shards. - """ + if dist.get_world_size(process_group) < 2: + return failed + flag = torch.tensor( + int(failed), + dtype=torch.int32, + device=_step_failure_collective_device(process_group), + ) + dist.all_reduce(flag, op=dist.ReduceOp.MAX, group=process_group) + return bool(flag.item()) + + +@torch.no_grad() +@torch._dynamo.disable +def _synchronize_step_control_range(local_minimum, local_maximum, process_group): + """Expand control-value bounds across one process group.""" + minimum = tuple(float(item) for item in local_minimum) + maximum = tuple(float(item) for item in local_maximum) + if len(minimum) != len(maximum): + raise ValueError("step-control bounds must have equal lengths") + if not torch.distributed.is_available() or not torch.distributed.is_initialized(): + return minimum, maximum + + import torch.distributed as dist + + if dist.get_world_size(process_group) < 2: + return minimum, maximum + bounds = torch.tensor( + minimum + tuple(-item for item in maximum), + dtype=torch.float64, + device=_step_failure_collective_device(process_group), + ) + dist.all_reduce(bounds, op=dist.ReduceOp.MIN, group=process_group) + values = bounds.cpu().tolist() + width = len(minimum) + return tuple(values[:width]), tuple(-item for item in values[width:]) + + +def _amp_optimizer_step_controls(optimizer): + """Parse GradScaler's temporary controls without mutating gradients.""" found_inf = getattr(optimizer, "found_inf", None) if found_inf is not None: if torch.is_tensor(found_inf): @@ -524,20 +577,54 @@ def _amp_prepare_optimizer_step(optimizer) -> bool: overflow = bool(found_inf.detach().item()) else: overflow = bool(found_inf) - if overflow: - return False + else: + overflow = False grad_scale = getattr(optimizer, "grad_scale", None) - if grad_scale is None: - # Explicit scaler.unscale_(optimizer), or an ordinary non-AMP step. - return True - if torch.is_tensor(grad_scale): if grad_scale.numel() != 1: raise RuntimeError( "GradScaler supplied a non-scalar grad_scale tensor with shape " "{}".format(tuple(grad_scale.shape)) ) + scale_value = float(grad_scale.detach().item()) + elif grad_scale is not None: + scale_value = float(grad_scale) + else: + scale_value = 0.0 + if grad_scale is not None and ( + not math.isfinite(scale_value) or scale_value <= 0.0 + ): + raise RuntimeError( + "GradScaler supplied a non-finite or non-positive grad_scale" + ) + return overflow, grad_scale, scale_value + + +@torch.no_grad() +def _amp_prepare_optimizer_step(optimizer) -> bool: + """Honor PyTorch's native ``GradScaler`` optimizer-step protocol. + + ``GradScaler.step`` gives AMP-aware optimizers two temporary attributes: + ``found_inf`` contains the result of the non-finite scan and ``grad_scale`` + contains the scale that still needs to be removed. When users explicitly + call ``scaler.unscale_(optimizer)`` first, ``grad_scale`` is ``None`` and + the gradients must not be divided a second time. + + Return ``False`` on overflow before callers learn a codebook or mutate any + optimizer/parameter state. On a finite automatic-unscale step, multiply the + gradients in place by the same fp32 reciprocal used by ``GradScaler``. This + keeps every existing fused/unfused update path operating on ordinary + unscaled gradients and also works for local DTensor shards. + """ + overflow, grad_scale, _ = _amp_optimizer_step_controls(optimizer) + if overflow: + return False + if grad_scale is None: + # Explicit scaler.unscale_(optimizer), or an ordinary non-AMP step. + return True + + if torch.is_tensor(grad_scale): scale = grad_scale.detach() # Match torch.amp.GradScaler.unscale_: computing the reciprocal in # fp64 avoids compile-option-dependent fp32 division differences. diff --git a/src/gefen/gefen_muon.py b/src/gefen/gefen_muon.py index 02e435f..6963101 100644 --- a/src/gefen/gefen_muon.py +++ b/src/gefen/gefen_muon.py @@ -9,8 +9,11 @@ from gefen.gefen import ( Gefen, + _amp_optimizer_step_controls, _amp_prepare_optimizer_step, _assert_optimizer_gradients_structurally_valid, + _synchronize_step_control_range, + _synchronize_step_failure, ) EPS = 1e-7 @@ -1302,6 +1305,208 @@ def _dist_available() -> bool: return False return torch.distributed.is_initialized() + @torch._dynamo.disable + def _step_failure_process_groups(self): + """Return the control scope for eager exact/distributed Muon steps.""" + if not self._dist_available(): + return () + + params = [param for group in self.param_groups for param in group["params"]] + # The protocol returns host-readable flags and therefore cannot be + # captured. Captured steps already require an eager warmup with a fixed + # control-flow and gradient-presence pattern. + if any(param.device.type == "cuda" for param in params) and ( + torch.cuda.is_current_stream_capturing() + ): + return () + + import torch.distributed as dist + + meshes = {} + 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: + 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), + ) + meshes.setdefault(key, (mesh, process_groups)) + + if not meshes: + return () + + member_sets = {frozenset(key[2]) for key in meshes} + world_members = frozenset(range(dist.get_world_size())) + participant_union = frozenset().union(*member_sets) + if world_members in member_sets: + return (dist.group.WORLD,) + + enclosing_sets = [ + members for members in member_sets if participant_union <= members + ] + if len(member_sets) > 1 and not enclosing_sets: + # No existing group encloses every upcoming collective participant. + # Entering only a rank-local subset of those groups could create a + # new deadlock before the established mesh/order preflight. Leave + # this uncommon topology on its existing behavior until the caller + # supplies an explicit enclosing control group. + return () + + selected_members = ( + min(enclosing_sets, key=lambda item: (len(item), sorted(item))) + if enclosing_sets + else next(iter(member_sets)) + ) + candidates = [ + (key, entry) + for key, entry in meshes.items() + if frozenset(key[2]) == selected_members + ] + _, (mesh, process_groups) = min(candidates, key=lambda item: item[0]) + if mesh.size() < 2: + return () + if mesh.ndim == 1: + return (process_groups[0],) + return tuple( + process_group + for process_group in process_groups + if dist.get_world_size(process_group) > 1 + ) + + @staticmethod + def _synchronize_sharded_step_flag(local_value, process_groups) -> bool: + synchronized = bool(local_value) + for process_group in process_groups: + synchronized = _synchronize_step_failure( + synchronized, process_group + ) + return synchronized + + def _synchronize_sharded_step_error( + self, error, phase: str, process_groups + ) -> None: + if not process_groups: + if error is not None: + raise error + return + failed = self._synchronize_sharded_step_flag( + error is not None, process_groups + ) + if not failed: + return + + import torch.distributed as dist + + if error is not None: + raise RuntimeError( + "GefenMuon {} failed on local rank {}: {}".format( + phase, dist.get_rank(), error + ) + ) from error + raise RuntimeError( + "GefenMuon {} failed on another process-group member".format( + phase + ) + ) + + @staticmethod + def _synchronize_sharded_step_control_range(local_control, process_groups): + minimum = tuple(float(item) for item in local_control) + maximum = minimum + for process_group in process_groups: + minimum, maximum = _synchronize_step_control_range( + minimum, maximum, process_group + ) + return minimum, maximum + + def _prepare_synchronized_amp_step(self, 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" + ) + if not process_groups: + if not local_present: + return True + return _amp_prepare_optimizer_step(optimizer) + + try: + if local_present: + local_overflow, grad_scale, scale_value = ( + _amp_optimizer_step_controls(optimizer) + ) + local_scale_present = grad_scale is not None + else: + local_overflow = False + local_scale_present = False + scale_value = 0.0 + local_amp_error = None + except Exception as exc: + local_overflow = False + local_scale_present = False + scale_value = 0.0 + local_amp_error = exc + self._synchronize_sharded_step_error( + local_amp_error, "AMP control preflight", process_groups + ) + + minimum, maximum = self._synchronize_sharded_step_control_range( + ( + int(local_present), + int(local_overflow), + int(local_scale_present), + scale_value, + ), + process_groups, + ) + if minimum[0] != maximum[0]: + raise RuntimeError( + "GefenMuon AMP protocol presence differs across process-group " + "members; use the same group-aware gradient scaler on every member" + ) + if not bool(maximum[0]): + return True + if minimum[1] != maximum[1]: + raise RuntimeError( + "GefenMuon AMP found_inf differs across process-group members; " + "use a group-aware gradient scaler" + ) + if minimum[2] != maximum[2]: + raise RuntimeError( + "GefenMuon AMP grad_scale presence differs across process-group " + "members; use the same group-aware gradient scaler on every member" + ) + if bool(maximum[2]) and minimum[3] != maximum[3]: + raise RuntimeError( + "GefenMuon AMP grad_scale differs across process-group members; " + "use a group-aware gradient scaler" + ) + if bool(maximum[1]): + return False + + try: + should_step = _amp_prepare_optimizer_step(optimizer) + local_amp_error = None + except Exception as exc: + should_step = False + local_amp_error = exc + self._synchronize_sharded_step_error( + local_amp_error, "AMP preparation", process_groups + ) + return should_step + @torch._dynamo.disable def _assert_sharded_grad_presence_consistent(self) -> None: """Fail before collectives when mesh ranks disagree on the step inputs. @@ -2728,24 +2933,33 @@ def _load_state_dict_impl(self, state_dict): @torch.no_grad() def step(self, closure=None): - self._assert_capturable_if_capturing() loss = None - if closure is not None: - with torch.enable_grad(): - loss = closure() - - _assert_optimizer_gradients_structurally_valid( - self, require_2d_params=True + hybrid_preflight_complete = bool( + getattr(self, "_gefen_hybrid_precollective_preflight", False) ) - - # Native GradScaler calls AMP-aware optimizers even when its non-finite - # scan found an overflow. Skip before the sharded preflight, first-step - # codebook learning, or any optimizer/parameter mutation; finite scaled - # gradients are unscaled once here for every existing Muon path. - if ( - hasattr(self, "found_inf") or hasattr(self, "grad_scale") - ) and not _amp_prepare_optimizer_step(self): - return loss + if not hybrid_preflight_complete: + process_groups = self._step_failure_process_groups() + try: + self._assert_capturable_if_capturing() + if closure is not None: + with torch.enable_grad(): + loss = closure() + _assert_optimizer_gradients_structurally_valid( + self, require_2d_params=True + ) + local_preflight_error = None + except Exception as exc: + loss = None + local_preflight_error = exc + self._synchronize_sharded_step_error( + local_preflight_error, "step preflight", process_groups + ) + + # Every mesh member enters the AMP control agreement, including a + # member with no local GradScaler attributes. This prevents the + # protocol-presence decision itself from becoming rank-divergent. + if not self._prepare_synchronized_amp_step(self, process_groups): + return loss # Partition the work once so distributed-mode sharded params can take the # stable-owner Parallel-Muon path while every other param keeps the normal diff --git a/src/gefen/hybrid.py b/src/gefen/hybrid.py index 8b59fe4..0e2fc2d 100644 --- a/src/gefen/hybrid.py +++ b/src/gefen/hybrid.py @@ -656,7 +656,11 @@ def _assert_capturable_devices_if_capturing(self) -> None: ) def step(self, closure=None): - self._assert_capturable_devices_if_capturing() + process_groups = ( + self.muon._step_failure_process_groups() + if self.muon is not None + else () + ) # 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 @@ -667,41 +671,74 @@ def step(self, closure=None): # each child sub-optimizer's (wrapped) step; see the class docstring. args = (self, closure) if closure is not None else (self,) kwargs = {} - for pre_hook in self._optimizer_step_pre_hooks.values(): - result = pre_hook(self, args, kwargs) - if result is not None: - if isinstance(result, tuple) and len(result) == 2: - args, kwargs = result - else: - raise RuntimeError( - f"{self.__class__.__name__}.step pre hook must return None " - f"or a tuple of (new_args, new_kwargs), but got {result}." - ) - # Re-read the closure from the (possibly hook-rewritten) call args, as - # torch's wrapper would by calling step(*args, **kwargs). - closure = args[1] if len(args) > 1 else kwargs.get("closure") - loss = None - if closure is not None: - with torch.enable_grad(): - loss = closure() - for child in self._subopts: - _assert_optimizer_gradients_structurally_valid( - child, require_2d_params=child is self.muon + try: + self._assert_capturable_devices_if_capturing() + for pre_hook in self._optimizer_step_pre_hooks.values(): + result = pre_hook(self, args, kwargs) + if result is not None: + if isinstance(result, tuple) and len(result) == 2: + args, kwargs = result + else: + raise RuntimeError( + f"{self.__class__.__name__}.step pre hook must return None " + f"or a tuple of (new_args, new_kwargs), but got {result}." + ) + # Re-read the closure from the (possibly hook-rewritten) call args, + # as torch's wrapper would by calling step(*args, **kwargs). + closure = args[1] if len(args) > 1 else kwargs.get("closure") + if closure is not None: + with torch.enable_grad(): + loss = closure() + for child in self._subopts: + _assert_optimizer_gradients_structurally_valid( + child, require_2d_params=child is self.muon + ) + local_preflight_error = 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 + # 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 ( - hasattr(self, "found_inf") or hasattr(self, "grad_scale") - ) and not _amp_prepare_optimizer_step(self): + 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 + if not should_step: for post_hook in self._optimizer_step_post_hooks.values(): post_hook(self, args, kwargs) return loss + with torch.no_grad(): for o in self._subopts: - o.step() + if o is not self.muon: + o.step() + continue + marker = object() + previous = getattr( + o, "_gefen_hybrid_precollective_preflight", marker + ) + o._gefen_hybrid_precollective_preflight = True + try: + o.step() + finally: + if previous is marker: + del o._gefen_hybrid_precollective_preflight + else: + o._gefen_hybrid_precollective_preflight = previous 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 new file mode 100644 index 0000000..91ff3b7 --- /dev/null +++ b/tests/test_precollective_failure_sync.py @@ -0,0 +1,402 @@ +"""CPU/Gloo coverage for failures before sharded Muon step collectives.""" + +from __future__ import annotations + +import copy +from datetime import timedelta +import os +from pathlib import Path +import queue as queue_module +import tempfile +import time +import traceback + +import pytest +import torch +import torch.distributed as dist +from torch import nn + +from gefen import GefenMuon, GefenMuonHybrid + + +_WORLD_SIZE = 2 +_PROCESS_GROUP_TIMEOUT_SECONDS = 20 +_STEP_DEADLINE_SECONDS = 10.0 +_WORKER_DEADLINE_SECONDS = 60.0 +_CASES = ("muon_exact", "muon_distributed", "hybrid_exact") + + +def _local_clone(value: torch.Tensor) -> torch.Tensor: + value = value.to_local() if hasattr(value, "to_local") else value + value = value.wait() if hasattr(value, "wait") else value + return value.detach().clone() + + +def _clone_tree(value): + if torch.is_tensor(value): + return _local_clone(value) + if isinstance(value, dict): + return {key: _clone_tree(item) for key, item in value.items()} + if isinstance(value, list): + return [_clone_tree(item) for item in value] + if isinstance(value, tuple): + return tuple(_clone_tree(item) for item in value) + return copy.deepcopy(value) + + +def _trees_equal(left, right) -> bool: + if torch.is_tensor(left) or torch.is_tensor(right): + return ( + torch.is_tensor(left) + and torch.is_tensor(right) + and left.dtype == right.dtype + and tuple(left.shape) == tuple(right.shape) + and torch.equal(left, right) + ) + if isinstance(left, dict) or isinstance(right, dict): + return ( + isinstance(left, dict) + and isinstance(right, dict) + and left.keys() == right.keys() + and all(_trees_equal(left[key], right[key]) for key in left) + ) + if isinstance(left, (list, tuple)) or isinstance(right, (list, tuple)): + return ( + type(left) is type(right) + and len(left) == len(right) + and all(_trees_equal(a, b) for a, b in zip(left, right)) + ) + return type(left) is type(right) and left == right + + +def _attribute_snapshot(owner, name: str): + if not hasattr(owner, name): + return (False, None, None) + value = getattr(owner, name) + return (True, id(value), _clone_tree(value)) + + +def _optimizer_snapshot(optimizer, parameters): + children = tuple(getattr(optimizer, "_subopts", (optimizer,))) + child_states = [] + codebooks = [] + global_steps = [] + for child in children: + child_states.append( + ( + id(child.state), + tuple( + (id(parameter), id(state), _clone_tree(state)) + for parameter, state in child.state.items() + ), + ) + ) + codebooks.append( + tuple( + _attribute_snapshot(child, name) + for name in ( + "_gefen_codebook", + "_gefen_codebook_by_device", + "_gefen_codebook_lut_by_device", + ) + ) + ) + global_steps.append( + tuple( + _attribute_snapshot(child, name) + for name in ( + "_gefen_global_step", + "_gefen_global_step_by_device", + ) + ) + ) + return { + "parameters": tuple(_local_clone(parameter) for parameter in parameters), + "gradients": tuple( + None if parameter.grad is None else _local_clone(parameter.grad) + for parameter in parameters + ), + "state": tuple(child_states), + "codebook": tuple(codebooks), + "global_step": tuple(global_steps), + } + + +def _snapshot_comparison(before, after): + return { + name: _trees_equal(before[name], after[name]) + for name in ("parameters", "gradients", "state", "codebook", "global_step") + } + + +def _make_optimizer(mesh, case: str): + from torch.distributed.tensor import Shard, distribute_tensor + + full_weight = torch.linspace(-0.8, 0.9, 64, dtype=torch.float32).reshape(8, 8) + 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": + bias = nn.Parameter(torch.linspace(-0.4, 0.3, 8, dtype=torch.float32)) + optimizer = GefenMuonHybrid( + [("weight", weight)], + [("bias", bias)], + lr=1e-3, + fused=False, + ns_steps=1, + ns_schedule="standard", + sharded_mode="exact", + backup_optimizer="gefen", + normuon=False, + ) + parameters = (weight, bias) + bias.grad = torch.linspace(0.2, -0.3, 8, dtype=torch.float32) + else: + mode = case.removeprefix("muon_") + optimizer = GefenMuon( + [("weight", weight)], + lr=1e-3, + fused=False, + ns_steps=1, + ns_schedule="standard", + sharded_mode=mode, + ) + parameters = (weight,) + + weight.grad = distribute_tensor(full_weight_grad.clone(), mesh, [Shard(0)]) + return optimizer, parameters + + +def _closure_failure_result(rank: int, mesh, case: str): + optimizer, parameters = _make_optimizer(mesh, case) + before = _optimizer_snapshot(optimizer, parameters) + + def closure(): + if rank == 0: + raise RuntimeError("rank-zero closure failure") + return torch.tensor(1.0) + + started = time.monotonic() + try: + optimizer.step(closure) + message = None + error_traceback = None + except BaseException as exc: + message = "{}: {}".format(type(exc).__name__, exc) + error_traceback = traceback.format_exc() + elapsed = time.monotonic() - started + after = _optimizer_snapshot(optimizer, parameters) + return { + "message": message, + "traceback": error_traceback, + "elapsed": elapsed, + "unchanged": _snapshot_comparison(before, after), + } + + +def _amp_control_result(rank: int, mesh, case: str, scenario: str): + optimizer, parameters = _make_optimizer(mesh, case) + if scenario == "divergent_found_inf": + optimizer.found_inf = torch.tensor(float(rank == 0)) + optimizer.grad_scale = torch.tensor(8.0) + elif scenario == "protocol_presence": + if rank == 0: + optimizer.found_inf = torch.tensor(0.0) + optimizer.grad_scale = torch.tensor(8.0) + elif scenario == "scale_disagreement": + optimizer.found_inf = torch.tensor(0.0) + optimizer.grad_scale = torch.tensor(8.0 if rank == 0 else 4.0) + elif scenario == "group_wide_overflow": + optimizer.found_inf = torch.tensor(1.0) + optimizer.grad_scale = torch.tensor(8.0) + else: + raise AssertionError("unknown AMP-control scenario: {}".format(scenario)) + + post_hook_calls = [] + hook_handle = None + if case == "hybrid_exact" and scenario == "group_wide_overflow": + hook_handle = optimizer.register_step_post_hook( + lambda _optimizer, _args, _kwargs: post_hook_calls.append(True) + ) + before = _optimizer_snapshot(optimizer, parameters) + + started = time.monotonic() + try: + optimizer.step() + message = None + error_traceback = None + except BaseException as exc: + message = "{}: {}".format(type(exc).__name__, exc) + error_traceback = traceback.format_exc() + elapsed = time.monotonic() - started + after = _optimizer_snapshot(optimizer, parameters) + if hook_handle is not None: + hook_handle.remove() + return { + "message": message, + "traceback": error_traceback, + "elapsed": elapsed, + "controls_present": hasattr(optimizer, "found_inf") or hasattr(optimizer, "grad_scale"), + "local_found_inf": ( + float(optimizer.found_inf.item()) if hasattr(optimizer, "found_inf") else None + ), + "local_grad_scale": ( + float(optimizer.grad_scale.item()) if hasattr(optimizer, "grad_scale") else None + ), + "post_hook_calls": ( + len(post_hook_calls) + if case == "hybrid_exact" and scenario == "group_wide_overflow" + else None + ), + "unchanged": _snapshot_comparison(before, after), + } + + +def _distributed_worker(rank: int, init_method: str, case: str, result_queue) -> None: + 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() + divergent_found_inf = _amp_control_result(rank, mesh, case, "divergent_found_inf") + dist.barrier() + protocol_presence = _amp_control_result(rank, mesh, case, "protocol_presence") + dist.barrier() + scale_disagreement = _amp_control_result(rank, mesh, case, "scale_disagreement") + dist.barrier() + group_wide_overflow = _amp_control_result(rank, mesh, case, "group_wide_overflow") + dist.barrier() + result_queue.put( + { + "rank": rank, + "closure_failure": closure_failure, + "divergent_found_inf": divergent_found_inf, + "protocol_presence": protocol_presence, + "scale_disagreement": scale_disagreement, + "group_wide_overflow": group_wide_overflow, + } + ) + 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): + context = torch.multiprocessing.get_context("spawn") + result_queue = context.Queue() + descriptor, rendezvous_path = tempfile.mkstemp(prefix="gefen-precollective-sync-") + os.close(descriptor) + os.unlink(rendezvous_path) + init_method = Path(rendezvous_path).resolve().as_uri() + processes = [ + context.Process( + target=_distributed_worker, + args=(rank, init_method, case, result_queue), + ) + for rank in range(_WORLD_SIZE) + ] + results = [] + hung_ranks = [] + exit_codes = [] + try: + for process in processes: + process.start() + deadline = time.monotonic() + _WORKER_DEADLINE_SECONDS + for process in processes: + process.join(timeout=max(0.0, deadline - time.monotonic())) + hung_ranks = [rank for rank, process in enumerate(processes) if process.is_alive()] + exit_codes = [process.exitcode for process in processes] + while len(results) < _WORLD_SIZE: + try: + results.append(result_queue.get(timeout=0.5)) + except queue_module.Empty: + break + finally: + for process in processes: + if process.is_alive(): + process.terminate() + process.join(timeout=5) + result_queue.close() + result_queue.join_thread() + try: + os.unlink(rendezvous_path) + except FileNotFoundError: + pass + + assert not hung_ranks, (case, "workers hung", hung_ranks, exit_codes, results) + assert all(code == 0 for code in exit_codes), (case, "nonzero worker exit", exit_codes, results) + assert len(results) == _WORLD_SIZE, (case, "missing worker result", exit_codes, results) + return sorted(results, key=lambda result: result["rank"]) + + +@pytest.mark.skipif( + not dist.is_available() or not dist.is_gloo_available(), + reason="pre-collective failure synchronization coverage requires Gloo", +) +@pytest.mark.parametrize("case", _CASES) +def test_rank_local_closure_and_amp_controls_are_synchronized(case): + results = _run_distributed_case(case) + 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 + phase = "hybrid step preflight" if case == "hybrid_exact" else "step preflight" + assert "GefenMuon {} failed on local rank 0".format(phase) in closure_results[0]["message"] + assert "rank-zero closure failure" in closure_results[0]["message"] + assert "GefenMuon {} failed on another process-group member".format(phase) in closure_results[1]["message"] + 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 + + overflow_results = [result["divergent_found_inf"] for result in results] + assert [item["local_found_inf"] for item in overflow_results] == [1.0, 0.0] + assert all( + item["message"] is not None + and "AMP found_inf differs across process-group members" in item["message"] + and "group-aware gradient scaler" in item["message"] + for item in overflow_results + ), overflow_results + assert all(item["elapsed"] < _STEP_DEADLINE_SECONDS for item in overflow_results), overflow_results + assert all(all(item["unchanged"].values()) for item in overflow_results), overflow_results + + presence_results = [result["protocol_presence"] for result in results] + assert [item["controls_present"] for item in presence_results] == [True, False] + assert all( + item["message"] is not None + and "AMP protocol presence differs across process-group members" in item["message"] + and "group-aware gradient scaler" in item["message"] + for item in presence_results + ), presence_results + assert all(item["elapsed"] < _STEP_DEADLINE_SECONDS for item in presence_results), presence_results + assert all(all(item["unchanged"].values()) for item in presence_results), presence_results + + scale_results = [result["scale_disagreement"] for result in results] + assert [item["local_found_inf"] for item in scale_results] == [0.0, 0.0] + assert [item["local_grad_scale"] for item in scale_results] == [8.0, 4.0] + assert all( + item["message"] is not None + and "AMP grad_scale differs across process-group members" in item["message"] + and "group-aware gradient scaler" in item["message"] + for item in scale_results + ), scale_results + assert all(item["elapsed"] < _STEP_DEADLINE_SECONDS for item in scale_results), scale_results + assert all(all(item["unchanged"].values()) for item in scale_results), scale_results + + skip_results = [result["group_wide_overflow"] for result in results] + assert all(item["message"] is None for item in skip_results), skip_results + assert all(item["elapsed"] < _STEP_DEADLINE_SECONDS for item in skip_results), skip_results + assert all(all(item["unchanged"].values()) for item in skip_results), skip_results + if case == "hybrid_exact": + 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] From 70ba3a0a28dfb4f1d44342e019956c8fc716e85a Mon Sep 17 00:00:00 2001 From: thad0ctor Date: Wed, 15 Jul 2026 06:46:53 -0700 Subject: [PATCH 2/4] Disable Dynamo tracing on the pre-collective sync wrappers _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). --- src/gefen/gefen_muon.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/src/gefen/gefen_muon.py b/src/gefen/gefen_muon.py index 6963101..b0d9e9a 100644 --- a/src/gefen/gefen_muon.py +++ b/src/gefen/gefen_muon.py @@ -1387,6 +1387,7 @@ def _step_failure_process_groups(self): ) @staticmethod + @torch._dynamo.disable def _synchronize_sharded_step_flag(local_value, process_groups) -> bool: synchronized = bool(local_value) for process_group in process_groups: @@ -1395,6 +1396,7 @@ def _synchronize_sharded_step_flag(local_value, process_groups) -> bool: ) return synchronized + @torch._dynamo.disable def _synchronize_sharded_step_error( self, error, phase: str, process_groups ) -> None: @@ -1423,6 +1425,7 @@ def _synchronize_sharded_step_error( ) @staticmethod + @torch._dynamo.disable def _synchronize_sharded_step_control_range(local_control, process_groups): minimum = tuple(float(item) for item in local_control) maximum = minimum @@ -1432,6 +1435,7 @@ def _synchronize_sharded_step_control_range(local_control, process_groups): ) return minimum, maximum + @torch._dynamo.disable def _prepare_synchronized_amp_step(self, optimizer, process_groups) -> bool: """Agree on AMP controls before unscaling or entering Muon collectives.""" local_present = hasattr(optimizer, "found_inf") or hasattr( From 58ca39d5a88e5bc38a795857718c00d5e9540137 Mon Sep 17 00:00:00 2001 From: thad0ctor Date: Wed, 15 Jul 2026 08:00:01 -0700 Subject: [PATCH 3/4] Sync the pre-collective failure flag over every mesh, on the shard device 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). --- src/gefen/gefen.py | 23 ++++++++--- src/gefen/gefen_muon.py | 87 ++++++++++++++--------------------------- 2 files changed, 47 insertions(+), 63 deletions(-) diff --git a/src/gefen/gefen.py b/src/gefen/gefen.py index 1a77ac3..6e6b5b9 100644 --- a/src/gefen/gefen.py +++ b/src/gefen/gefen.py @@ -497,10 +497,15 @@ def _resolve_find_period_backend(grad: torch.Tensor) -> str: return "cuda_kernel" if grad_work.device.type == "cuda" else "cpu" -def _step_failure_collective_device(process_group) -> torch.device: +def _step_failure_collective_device( + process_group, *, collective_device=None +) -> torch.device: """Choose a backend-compatible device for a process-group control flag.""" import torch.distributed as dist + if collective_device is not None: + return torch.device(collective_device) + group = process_group if process_group is not None else dist.group.WORLD bound_device = getattr(group, "bound_device_id", None) if bound_device is not None: @@ -519,7 +524,9 @@ def _step_failure_collective_device(process_group) -> torch.device: @torch.no_grad() @torch._dynamo.disable -def _synchronize_step_failure(local_failed, process_group) -> bool: +def _synchronize_step_failure( + local_failed, process_group, *, collective_device=None +) -> bool: """Return whether any process-group member reported a step failure.""" failed = bool(local_failed) if not torch.distributed.is_available() or not torch.distributed.is_initialized(): @@ -532,7 +539,9 @@ def _synchronize_step_failure(local_failed, process_group) -> bool: flag = torch.tensor( int(failed), dtype=torch.int32, - device=_step_failure_collective_device(process_group), + device=_step_failure_collective_device( + process_group, collective_device=collective_device + ), ) dist.all_reduce(flag, op=dist.ReduceOp.MAX, group=process_group) return bool(flag.item()) @@ -540,7 +549,9 @@ def _synchronize_step_failure(local_failed, process_group) -> bool: @torch.no_grad() @torch._dynamo.disable -def _synchronize_step_control_range(local_minimum, local_maximum, process_group): +def _synchronize_step_control_range( + local_minimum, local_maximum, process_group, *, collective_device=None +): """Expand control-value bounds across one process group.""" minimum = tuple(float(item) for item in local_minimum) maximum = tuple(float(item) for item in local_maximum) @@ -556,7 +567,9 @@ def _synchronize_step_control_range(local_minimum, local_maximum, process_group) bounds = torch.tensor( minimum + tuple(-item for item in maximum), dtype=torch.float64, - device=_step_failure_collective_device(process_group), + device=_step_failure_collective_device( + process_group, collective_device=collective_device + ), ) dist.all_reduce(bounds, op=dist.ReduceOp.MIN, group=process_group) values = bounds.cpu().tolist() diff --git a/src/gefen/gefen_muon.py b/src/gefen/gefen_muon.py index b0d9e9a..8c3f385 100644 --- a/src/gefen/gefen_muon.py +++ b/src/gefen/gefen_muon.py @@ -1307,7 +1307,19 @@ def _dist_available() -> bool: @torch._dynamo.disable def _step_failure_process_groups(self): - """Return the control scope for eager exact/distributed Muon steps.""" + """Return the control scope for eager exact/distributed Muon steps. + + Every sharded (non-approx) mesh this rank participates in contributes its + collective groups, deduplicated by process-group name and returned in a + globally consistent (sorted) order -- exactly the meshes + ``_assert_sharded_grad_presence_consistent`` enters. Selecting only one + mesh, or abstaining when the meshes share no enclosing group, desyncs the + control collective: with overlapping meshes (e.g. members {0,1,2} and + {2,3}) a rank in both would skip while a rank in only one enters that + mesh's all-reduce, deadlocking. Each group carries the local shard device + so the control flag lands on the same device as the real collectives + rather than the ambient ``current_device``. + """ if not self._dist_available(): return () @@ -1322,7 +1334,7 @@ def _step_failure_process_groups(self): import torch.distributed as dist - meshes = {} + groups = {} for group in self.param_groups: if group["sharded_mode"] == "approx": continue @@ -1330,69 +1342,28 @@ def _step_failure_process_groups(self): if not self._is_sharded(param): continue mesh = param.device_mesh - if mesh.get_coordinate() is None: + if mesh.get_coordinate() is None or mesh.size() < 2: continue + device = self._state_tensor_device(param) 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), + mesh_groups = ( + process_groups[:1] if mesh.ndim == 1 else process_groups ) - meshes.setdefault(key, (mesh, process_groups)) - - if not meshes: - return () - - member_sets = {frozenset(key[2]) for key in meshes} - world_members = frozenset(range(dist.get_world_size())) - participant_union = frozenset().union(*member_sets) - if world_members in member_sets: - return (dist.group.WORLD,) - - enclosing_sets = [ - members for members in member_sets if participant_union <= members - ] - if len(member_sets) > 1 and not enclosing_sets: - # No existing group encloses every upcoming collective participant. - # Entering only a rank-local subset of those groups could create a - # new deadlock before the established mesh/order preflight. Leave - # this uncommon topology on its existing behavior until the caller - # supplies an explicit enclosing control group. - return () - - selected_members = ( - min(enclosing_sets, key=lambda item: (len(item), sorted(item))) - if enclosing_sets - else next(iter(member_sets)) - ) - candidates = [ - (key, entry) - for key, entry in meshes.items() - if frozenset(key[2]) == selected_members - ] - _, (mesh, process_groups) = min(candidates, key=lambda item: item[0]) - if mesh.size() < 2: - return () - if mesh.ndim == 1: - return (process_groups[0],) - return tuple( - process_group - for process_group in process_groups - if dist.get_world_size(process_group) > 1 - ) + for process_group in mesh_groups: + if dist.get_world_size(process_group) > 1: + groups.setdefault( + str(process_group.group_name), + (process_group, device), + ) + return tuple(groups[name] for name in sorted(groups)) @staticmethod @torch._dynamo.disable def _synchronize_sharded_step_flag(local_value, process_groups) -> bool: synchronized = bool(local_value) - for process_group in process_groups: + for process_group, collective_device in process_groups: synchronized = _synchronize_step_failure( - synchronized, process_group + synchronized, process_group, collective_device=collective_device ) return synchronized @@ -1429,9 +1400,9 @@ def _synchronize_sharded_step_error( def _synchronize_sharded_step_control_range(local_control, process_groups): minimum = tuple(float(item) for item in local_control) maximum = minimum - for process_group in process_groups: + for process_group, collective_device in process_groups: minimum, maximum = _synchronize_step_control_range( - minimum, maximum, process_group + minimum, maximum, process_group, collective_device=collective_device ) return minimum, maximum From dbbc001b80f1ccf909b020985e7bc4b346ebd932 Mon Sep 17 00:00:00 2001 From: thad0ctor Date: Wed, 15 Jul 2026 08:42:39 -0700 Subject: [PATCH 4/4] Preserve DeviceMesh dimension order in the pre-collective scope 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. --- src/gefen/gefen_muon.py | 56 +++++++++++++++++++++++++---------------- 1 file changed, 35 insertions(+), 21 deletions(-) diff --git a/src/gefen/gefen_muon.py b/src/gefen/gefen_muon.py index 8c3f385..0d275b1 100644 --- a/src/gefen/gefen_muon.py +++ b/src/gefen/gefen_muon.py @@ -1309,16 +1309,21 @@ def _dist_available() -> bool: def _step_failure_process_groups(self): """Return the control scope for eager exact/distributed Muon steps. - Every sharded (non-approx) mesh this rank participates in contributes its - collective groups, deduplicated by process-group name and returned in a - globally consistent (sorted) order -- exactly the meshes - ``_assert_sharded_grad_presence_consistent`` enters. Selecting only one - mesh, or abstaining when the meshes share no enclosing group, desyncs the - control collective: with overlapping meshes (e.g. members {0,1,2} and - {2,3}) a rank in both would skip while a rank in only one enters that - mesh's all-reduce, deadlocking. Each group carries the local shard device - so the control flag lands on the same device as the real collectives - rather than the ambient ``current_device``. + Enter EXACTLY the collectives ``_assert_sharded_grad_presence_consistent`` + enters, in the same order: every sharded (non-approx) mesh this rank + participates in, keyed and deduplicated by mesh content and iterated in + sorted-key order, and within each mesh its ``get_all_groups()`` dimension + order. Preserving the per-mesh dimension order matters for multi-dim + (HSDP/TP) meshes -- flattening and sorting all groups by name could make + one rank enter a row all-reduce while a peer is already blocked in its + column all-reduce, deadlocking the preflight itself. Each group carries + the local shard device so the control flag lands on the same device as + the real collectives rather than the ambient ``current_device``. + + Like the grad-presence preflight, this is a single pass per group; a + failure only reaches ranks sharing a mesh with the failing rank. Deeply + overlapping non-enclosing meshes therefore keep that preflight's existing + cross-mesh limitation. """ if not self._dist_available(): return () @@ -1334,7 +1339,7 @@ def _step_failure_process_groups(self): import torch.distributed as dist - groups = {} + by_mesh = {} for group in self.param_groups: if group["sharded_mode"] == "approx": continue @@ -1344,18 +1349,27 @@ def _step_failure_process_groups(self): mesh = param.device_mesh if mesh.get_coordinate() is None or mesh.size() < 2: continue - device = self._state_tensor_device(param) process_groups = tuple(mesh.get_all_groups()) - mesh_groups = ( - process_groups[:1] if mesh.ndim == 1 else process_groups + members = tuple( + int(item) + for item in mesh.mesh.detach().cpu().reshape(-1).tolist() ) - for process_group in mesh_groups: - if dist.get_world_size(process_group) > 1: - groups.setdefault( - str(process_group.group_name), - (process_group, device), - ) - return tuple(groups[name] for name in sorted(groups)) + 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) @staticmethod @torch._dynamo.disable