diff --git a/src/gefen/gefen.py b/src/gefen/gefen.py index c76d67f..6e6b5b9 100644 --- a/src/gefen/gefen.py +++ b/src/gefen/gefen.py @@ -497,22 +497,88 @@ 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, *, 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: + 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, *, 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(): + 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, collective_device=collective_device + ), + ) + 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, *, 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) + 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, collective_device=collective_device + ), + ) + 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 +590,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..0d275b1 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,197 @@ 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. + + 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 () + + 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 + + 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) + + @staticmethod + @torch._dynamo.disable + def _synchronize_sharded_step_flag(local_value, process_groups) -> bool: + synchronized = bool(local_value) + for process_group, collective_device in process_groups: + synchronized = _synchronize_step_failure( + synchronized, process_group, collective_device=collective_device + ) + return synchronized + + @torch._dynamo.disable + 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 + @torch._dynamo.disable + def _synchronize_sharded_step_control_range(local_control, process_groups): + minimum = tuple(float(item) for item in local_control) + maximum = minimum + for process_group, collective_device in process_groups: + minimum, maximum = _synchronize_step_control_range( + minimum, maximum, process_group, collective_device=collective_device + ) + 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( + 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 +2922,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]