From fc159a70264898f2677d53792df595011346b5a3 Mon Sep 17 00:00:00 2001 From: thad0ctor Date: Sun, 12 Jul 2026 00:25:16 -0700 Subject: [PATCH 01/14] Harden optimizer state and runtime validation --- pyproject.toml | 2 +- src/gefen/gefen.py | 129 ++++++++++++--- src/gefen/gefen_muon.py | 175 ++++++++++++++++++--- tests/test_capturable.py | 16 +- tests/test_cpu_step_checkpoint.py | 79 ++++++++++ tests/test_dispatch_gating_cpu.py | 30 ++++ tests/test_epsilon_muon_param_group_cpu.py | 157 ++++++++++++++++++ 7 files changed, 540 insertions(+), 48 deletions(-) create mode 100644 tests/test_epsilon_muon_param_group_cpu.py diff --git a/pyproject.toml b/pyproject.toml index 8bf9db5..f909c3f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -38,7 +38,7 @@ classifiers = [ dependencies = [ "numpy>=1.24", "torch>=2.5", - "numba" + "numba>=0.65" ] [project.optional-dependencies] diff --git a/src/gefen/gefen.py b/src/gefen/gefen.py index b498d62..e44dd7c 100644 --- a/src/gefen/gefen.py +++ b/src/gefen/gefen.py @@ -724,8 +724,10 @@ def _validate_group_options(lr, betas, eps, weight_decay): ) elif not 0.0 <= weight_decay: raise ValueError("Invalid weight_decay value: {}".format(weight_decay)) - elif not 0.0 <= eps: - raise ValueError("Invalid epsilon value: {}".format(eps)) + elif not math.isfinite(eps) or eps <= 0.0: + raise ValueError( + "Invalid epsilon value: {}; eps must be finite and > 0".format(eps) + ) @staticmethod def _iter_params_with_names(group_params): @@ -1831,20 +1833,16 @@ def _predict_period_from_grad_sq( self, param_name: str, param: torch.Tensor, grad: torch.Tensor ) -> int: backend = _resolve_find_period_backend(grad) - # If the fused CUDA toolchain GENUINELY failed to build (the step-1 probe - # ran and returned False -> _fused_build_ok is False), fall back to the - # pure-torch "gpu" period backend rather than crashing in the separate - # period-variance extension load. Gate on the explicit build-failure - # verdict, NOT on _gefen_fused_toolchain_ok(): the latter also returns - # False for a user-chosen fused=False optimizer, which must keep using - # the cuda_kernel period backend (bit-identical to before this scrub). - # Only override the AUTO resolution -- an explicit FIND_PERIOD_BACKEND is - # the user's deliberate choice. + # ``fused=False`` is an explicit request for a pure-PyTorch optimizer + # step, so it must not cross a different lazy-JIT boundary merely to + # choose the block period. A failed fused-kernel probe has the same + # requirement. Keep honoring an explicit FIND_PERIOD_BACKEND override; + # callers that deliberately select ``cuda_kernel`` accept its JIT. if ( backend == "cuda_kernel" and FIND_PERIOD_BACKEND is None and grad.device.type == "cuda" - and self._fused_build_ok is False + and (not self.fused or self._fused_build_ok is False) ): backend = "gpu" if backend == "cuda_kernel": @@ -1998,14 +1996,14 @@ def _learn_gefen_exact_codebook( force_endpoints=True, verbose=self.verbose, compute_mse_logging=compute_mse_logging, - # Use the fused histogram UNLESS the toolchain probe explicitly - # failed (_fused_build_ok is False). Gating on the build-failure - # verdict rather than self.fused preserves the pre-scrub behavior of - # the histogram (it ran on CUDA tensors regardless of the optimizer's - # fused flag), so a fused=False run stays bit-identical; only a - # genuinely broken toolchain drops to the pure-torch bincount path. + # ``fused=False`` promises a JIT-free pure-PyTorch path. Do not load + # the separate histogram extension in that mode (or after the fused + # toolchain probe failed). An explicit global period-backend choice + # remains independent of this histogram implementation decision. use_fused_histogram=( - FUSE_HISTOGRAM_FOR_EXACT and self._fused_build_ok is not False + FUSE_HISTOGRAM_FOR_EXACT + and self.fused + and self._fused_build_ok is not False ), ) if codebook is None: @@ -3298,6 +3296,20 @@ def _compact(value): # these custom top-level keys; _maybe_refresh_gefen_codebook handles that # fallback by reusing the restored periods.) state_dict["gefen_codebook"] = self._gefen_codebook + # PyTorch Distributed Checkpoint's ``get_optimizer_state_dict`` keeps + # only the conventional ``state`` and ``param_groups`` top-level keys. + # Mirror the optimizer-level values inside every serialized group so a + # DCP/FSDP normalization round-trip cannot silently discard the frozen + # codebook and reinterpret restored uint8 momentum indices with a newly + # learned one. The loader removes this private transport metadata before + # handing groups to torch, so it never leaks into live scheduler groups. + checkpoint_metadata = { + "format_version": 1, + "global_step": self._gefen_global_step, + "codebook": self._gefen_codebook, + } + for group in state_dict["param_groups"]: + group["_gefen_checkpoint_metadata"] = checkpoint_metadata return state_dict @staticmethod @@ -3378,8 +3390,85 @@ def load_state_dict(self, state_dict): # dict (torch.optim.Optimizer.load_state_dict leaves its input intact); # otherwise a second load of the same dict loses the gefen_* keys. state_dict = dict(state_dict) - gefen_global_step = state_dict.pop("gefen_global_step", 0) + # Copy the group dictionaries before removing private transport + # metadata so repeated loads leave the caller's checkpoint untouched. + state_dict["param_groups"] = [ + dict(group) for group in state_dict.get("param_groups", ()) + ] + group_metadata = [] + for group in state_dict["param_groups"]: + metadata = group.pop("_gefen_checkpoint_metadata", None) + if metadata is not None: + group_metadata.append(metadata) + + gefen_global_step = state_dict.pop("gefen_global_step", None) gefen_codebook = state_dict.pop("gefen_codebook", None) + if group_metadata: + if len(group_metadata) != len(state_dict["param_groups"]): + raise ValueError( + "Gefen checkpoint metadata is present on only some parameter groups" + ) + first_metadata = group_metadata[0] + if first_metadata.get("format_version") != 1: + raise ValueError( + "Unsupported Gefen checkpoint metadata format_version: {}".format( + first_metadata.get("format_version") + ) + ) + for metadata in group_metadata[1:]: + same_step = metadata.get("global_step") == first_metadata.get( + "global_step" + ) + left_codebook = metadata.get("codebook") + right_codebook = first_metadata.get("codebook") + same_codebook = ( + left_codebook is right_codebook + or ( + torch.is_tensor(left_codebook) + and torch.is_tensor(right_codebook) + and torch.equal(left_codebook, right_codebook) + ) + ) + if not same_step or not same_codebook: + raise ValueError( + "Gefen checkpoint parameter groups carry inconsistent " + "optimizer metadata" + ) + metadata_step = first_metadata.get("global_step", 0) + metadata_codebook = first_metadata.get("codebook") + if gefen_global_step is None: + gefen_global_step = metadata_step + elif gefen_global_step != metadata_step: + raise ValueError( + "Gefen checkpoint top-level and parameter-group global steps disagree" + ) + if gefen_codebook is None: + gefen_codebook = metadata_codebook + elif not ( + gefen_codebook is metadata_codebook + or ( + torch.is_tensor(gefen_codebook) + and torch.is_tensor(metadata_codebook) + and torch.equal(gefen_codebook, metadata_codebook) + ) + ): + raise ValueError( + "Gefen checkpoint top-level and parameter-group codebooks disagree" + ) + if gefen_global_step is None: + gefen_global_step = 0 + has_quantized_momentum = any( + isinstance(param_state, dict) and "m_codebook" in param_state + for param_state in (state_dict.get("state", {}) or {}).values() + ) + if has_quantized_momentum and gefen_codebook is None: + raise ValueError( + "Gefen checkpoint contains quantized momentum indices but no frozen " + "codebook. Loading it would reinterpret the restored indices with a " + "different codebook and silently corrupt optimizer state. Resume from " + "an unmodified Gefen checkpoint or a DCP checkpoint written by a " + "version that preserves _gefen_checkpoint_metadata." + ) state_dict = self._pack_legacy_param_groups_for_load(state_dict) # torch.optim.Optimizer.load_state_dict casts *every* per-param state diff --git a/src/gefen/gefen_muon.py b/src/gefen/gefen_muon.py index fefd132..9992718 100644 --- a/src/gefen/gefen_muon.py +++ b/src/gefen/gefen_muon.py @@ -554,9 +554,9 @@ def __init__( raise ValueError( "momentum should be in [0, 1) but is: {}".format(momentum) ) - if not eps > 0.0: + if not math.isfinite(eps) or eps <= 0.0: raise ValueError( - "eps should be > 0 but is: {}. The Newton-Schulz input " + "eps should be finite and > 0 but is: {}. The Newton-Schulz input " "normalization divides by the momentum matrix norm clamped at " "eps, so eps=0 turns an all-zero momentum matrix into 0/0 = " "NaN".format(eps) @@ -682,11 +682,36 @@ def __init__( raise ValueError( "normuon_beta2 must be in [0, 1) but is: {}".format(normuon_beta2) ) + if not math.isfinite(normuon_eps) or normuon_eps <= 0.0: + raise ValueError( + "normuon_eps must be finite and > 0 but is: {}".format(normuon_eps) + ) # Cautious masking (opt-in): zero update coordinates whose sign # disagrees with the current gradient, rescaled to preserve magnitude. self._normuon = normuon self._cautious = cautious + # torch.optim.Optimizer calls ``self.add_param_group`` while the base + # constructor is still running. Publish the complete Muon group schema + # first so both constructor groups and groups added later take the same + # validation/default-injection path. + self._muon_group_defaults = { + "momentum": momentum, + "nesterov": nesterov, + "ns_coefficients": ns_coefficients, + "ns_steps": ns_steps, + "adjust_lr_fn": adjust_lr_fn, + "sharded_mode": sharded_mode, + "fp8_ns": fp8_ns, + "fp8_ns_compile": fp8_ns_compile, + "batched_ns": bool(batched_ns), + "batched_ns_workspace_bytes": batched_ns_workspace_bytes, + "normuon": normuon, + "normuon_beta2": normuon_beta2, + "normuon_eps": normuon_eps, + "cautious": cautious, + } + super().__init__( params, lr=lr, @@ -699,28 +724,132 @@ def __init__( verbose=verbose, ) - for group in self.param_groups: - group["momentum"] = momentum - group["nesterov"] = nesterov - group["ns_coefficients"] = ns_coefficients - group["ns_steps"] = ns_steps - group["adjust_lr_fn"] = adjust_lr_fn - group["sharded_mode"] = sharded_mode - group["fp8_ns"] = fp8_ns - group["fp8_ns_compile"] = fp8_ns_compile - group["batched_ns"] = bool(batched_ns) - group["batched_ns_workspace_bytes"] = batched_ns_workspace_bytes - group["normuon"] = normuon - group["normuon_beta2"] = normuon_beta2 - group["normuon_eps"] = normuon_eps - group["cautious"] = cautious - for p in group["params"]: - if p.ndim != 2: + def add_param_group(self, param_group): + """Add a validated 2D Muon parameter group atomically. + + Muon-specific options may be supplied per group; omitted values inherit + the constructor defaults. The group is fully validated before the base + Gefen registration mutates ``param_groups`` or per-parameter state. + """ + if not isinstance(param_group, dict) or "params" not in param_group: + # Preserve Gefen's public error types/messages for malformed group + # containers and missing ``params``. + return super().add_param_group(param_group) + + group = dict(param_group) + raw_params = group["params"] + if isinstance(raw_params, torch.Tensor): + raw_params = [raw_params] + else: + raw_params = list(raw_params) + group["params"] = raw_params + + # Validate tensor type/complex support and Muon's dimensionality before + # registration. Gefen.add_param_group repeats its own checks, but doing + # this first keeps a mixed valid/invalid addition all-or-nothing. + for _, param in self._iter_params_with_names(raw_params): + if param.ndim != 2: + raise ValueError( + "GefenMuon only supports 2D parameters whereas we found a " + "parameter with size: {}".format(param.size()) + ) + + defaults = self._muon_group_defaults + momentum = group.get("momentum", defaults["momentum"]) + if not 0.0 <= momentum < 1.0: + raise ValueError( + "momentum should be in [0, 1) but is: {}".format(momentum) + ) + + adjust_lr_fn = group.get("adjust_lr_fn", defaults["adjust_lr_fn"]) + if adjust_lr_fn is not None and adjust_lr_fn not in ( + "original", + "match_rms_adamw", + ): + raise ValueError( + "Adjust learning rate function {} is not supported".format( + adjust_lr_fn + ) + ) + + sharded_mode = group.get("sharded_mode", defaults["sharded_mode"]) + if sharded_mode not in ("exact", "approx", "distributed"): + raise ValueError( + "sharded_mode must be 'exact', 'approx' or 'distributed' but is: " + "{}".format(sharded_mode) + ) + + ns_coefficients = group.get( + "ns_coefficients", defaults["ns_coefficients"] + ) + ns_steps = group.get("ns_steps", defaults["ns_steps"]) + ns_schedule = group.get("ns_schedule") + if ns_schedule is not None: + if isinstance(ns_schedule, str): + if ns_schedule not in NS_SCHEDULES: raise ValueError( - "GefenMuon only supports 2D parameters whereas we found a parameter with size: {}".format( - p.size(), + "Unknown ns_schedule {!r}; choose from {}".format( + ns_schedule, sorted(NS_SCHEDULES) ) ) + resolved_schedule = NS_SCHEDULES[ns_schedule] + else: + resolved_schedule = ns_schedule + if resolved_schedule is not None: + ns_coefficients = _normalize_ns_schedule( + resolved_schedule, ns_steps + ) + ns_steps = len(ns_coefficients) + # Validate the fixed-coefficient path and direct explicit schedules too. + _normalize_ns_schedule(ns_coefficients, ns_steps) + + workspace_bytes = group.get( + "batched_ns_workspace_bytes", + defaults["batched_ns_workspace_bytes"], + ) + if not isinstance(workspace_bytes, int) or isinstance(workspace_bytes, bool): + raise TypeError("batched_ns_workspace_bytes must be an integer") + if workspace_bytes <= 0: + raise ValueError("batched_ns_workspace_bytes must be positive") + + normuon_beta2 = group.get("normuon_beta2", defaults["normuon_beta2"]) + if not 0.0 <= normuon_beta2 < 1.0: + raise ValueError( + "normuon_beta2 must be in [0, 1) but is: {}".format(normuon_beta2) + ) + normuon_eps = group.get("normuon_eps", defaults["normuon_eps"]) + if not math.isfinite(normuon_eps) or normuon_eps <= 0.0: + raise ValueError( + "normuon_eps must be finite and > 0 but is: {}".format(normuon_eps) + ) + + group.update( + { + "momentum": momentum, + "nesterov": group.get("nesterov", defaults["nesterov"]), + "ns_coefficients": ns_coefficients, + "ns_steps": ns_steps, + "adjust_lr_fn": adjust_lr_fn, + "sharded_mode": sharded_mode, + "fp8_ns": group.get("fp8_ns", defaults["fp8_ns"]), + "fp8_ns_compile": group.get( + "fp8_ns_compile", defaults["fp8_ns_compile"] + ), + "batched_ns": bool( + group.get("batched_ns", defaults["batched_ns"]) + ), + "batched_ns_workspace_bytes": workspace_bytes, + "normuon": group.get("normuon", defaults["normuon"]), + "normuon_beta2": normuon_beta2, + "normuon_eps": normuon_eps, + "cautious": group.get("cautious", defaults["cautious"]), + # Gefen's inherited storage still carries beta1/beta2. Keep it + # consistent with the Muon momentum option instead of retaining + # irrelevant or contradictory caller-supplied Adam betas. + "betas": (momentum, 0.0), + } + ) + return super().add_param_group(group) def _init_gefen_muon_state(self, state, grad_view: torch.Tensor) -> None: self._init_gefen_state(state, grad_view) @@ -742,7 +871,9 @@ def _iter_gefen_grad_periods(self, reuse_existing_periods: bool = False): # approx mode learns the codebook/period from the LOCAL shard # (no all-gather) so periods divide the local numel that the # approximate step operates on; exact mode gathers the full matrix. - if self._sharded_mode == "approx" and hasattr(grad, "to_local"): + if group["sharded_mode"] == "approx" and hasattr( + grad, "to_local" + ): grad = grad.to_local() elif hasattr(grad, "full_tensor"): grad = grad.full_tensor() diff --git a/tests/test_capturable.py b/tests/test_capturable.py index 763a36f..adf8be3 100644 --- a/tests/test_capturable.py +++ b/tests/test_capturable.py @@ -145,10 +145,10 @@ def test_parity_gefen_factored_fused(): # per-step scalars flow through the device buffer). The two sides run the # identical kernel arithmetic: the fp32 scalars the kernel consumes are # asserted BIT-IDENTICAL (device float64 tensor ops vs host python doubles - # cast to the same fp32 values), so the only residual divergence is the - # factored kernel's own documented run-to-run atomicAdd nondeterminism in - # the row/col grad^2 sums (a re-run of the SAME config differs by a couple - # of 1-ulp bf16 flips) -- hold the params to that envelope. + # cast to the same fp32 values). Pin momentum period selection so a near tie + # in the separate variance-search reduction cannot change state geometry; + # this test then isolates capturable scalar plumbing, with only the factored + # row/col reduction's small run-to-run atomicAdd noise remaining. lr, wd, beta1, beta2 = 1e-3, 0.1, 0.9, 0.999 def run(capturable, steps=6, shape=(768, 512), seed=3): @@ -156,7 +156,13 @@ def run(capturable, steps=6, shape=(768, 512), seed=3): p = torch.nn.Parameter( (torch.randn(*shape, device=DEVICE) * 0.02).bfloat16() ) - opt = Gefen([("w", p)], lr=lr, weight_decay=wd, capturable=capturable) + opt = Gefen( + [("w", p)], + lr=lr, + weight_decay=wd, + capturable=capturable, + force_2d_period_one=True, + ) torch.manual_seed(seed + 1) for _ in range(steps): p.grad = torch.randn(*shape, device=DEVICE).bfloat16() * 1e-3 diff --git a/tests/test_cpu_step_checkpoint.py b/tests/test_cpu_step_checkpoint.py index 35d7b94..ae39758 100644 --- a/tests/test_cpu_step_checkpoint.py +++ b/tests/test_cpu_step_checkpoint.py @@ -144,6 +144,49 @@ def test_state_dict_roundtrip_bit_exact(factored): assert torch.equal(pa, pb), "resumed run diverged from continuous run" +@pytest.mark.parametrize("factored", [True, False]) +def test_distributed_checkpoint_optimizer_state_roundtrip_bit_exact(factored): + """PyTorch DCP normalization must retain Gefen's frozen codebook.""" + from torch.distributed.checkpoint.state_dict import ( + get_optimizer_state_dict, + set_optimizer_state_dict, + ) + + grads_all = _synthetic_grads(_small_model(), 4) + model_a = _small_model() + opt_a = Gefen( + list(model_a.named_parameters()), lr=1e-3, fused=False, factored_v_2d=factored + ) + for step_grads in grads_all[:2]: + _apply_grads(model_a, step_grads) + opt_a.step() + opt_a.zero_grad() + + saved_opt = copy.deepcopy(get_optimizer_state_dict(model_a, opt_a)) + assert "gefen_codebook" not in saved_opt + assert all( + "_gefen_checkpoint_metadata" in group for group in saved_opt["param_groups"] + ) + + model_b = _small_model() + model_b.load_state_dict(copy.deepcopy(model_a.state_dict())) + opt_b = Gefen( + list(model_b.named_parameters()), lr=1e-3, fused=False, factored_v_2d=factored + ) + set_optimizer_state_dict(model_b, opt_b, saved_opt) + + for step_grads in grads_all[2:]: + _apply_grads(model_a, step_grads) + _apply_grads(model_b, step_grads) + opt_a.step() + opt_b.step() + opt_a.zero_grad() + opt_b.zero_grad() + + for pa, pb in zip(model_a.parameters(), model_b.parameters()): + assert torch.equal(pa, pb), "DCP-resumed run diverged from continuous run" + + def _zero_style_flat_swap(opt): """Simulate DeepSpeed ZeRO-1/2 wrapping: replace every group's ``params`` with one fresh flat fp32 partition tensor, leaving ``opt.state`` untouched @@ -261,6 +304,42 @@ def test_load_state_dict_does_not_mutate_caller_dict(): opt.load_state_dict(sd) +def test_load_rejects_quantized_momentum_without_codebook(): + model = _small_model() + opt = Gefen(list(model.named_parameters()), lr=1e-3, fused=False) + _apply_grads(model, _synthetic_grads(model, 1)[0]) + opt.step() + sd = copy.deepcopy(opt.state_dict()) + sd.pop("gefen_codebook") + for group in sd["param_groups"]: + group.pop("_gefen_checkpoint_metadata") + + fresh_model = _small_model() + fresh_opt = Gefen(list(fresh_model.named_parameters()), lr=1e-3, fused=False) + with pytest.raises(ValueError, match="quantized momentum indices but no frozen codebook"): + fresh_opt.load_state_dict(sd) + + +def test_load_rejects_partial_or_conflicting_group_metadata(): + _, opt = _run_and_save(factored=False) + sd = copy.deepcopy(opt.state_dict()) + # Create a second group so partial metadata can be represented. + group = sd["param_groups"][0] + first_param = group["params"][0] + second_group = dict(group) + second_group["params"] = [first_param] + group["params"] = group["params"][1:] + second_group.pop("_gefen_checkpoint_metadata") + sd["param_groups"].append(second_group) + with pytest.raises(ValueError, match="only some parameter groups"): + opt.load_state_dict(sd) + + conflicting = copy.deepcopy(opt.state_dict()) + conflicting["param_groups"][0]["_gefen_checkpoint_metadata"]["global_step"] += 1 + with pytest.raises(ValueError, match="global steps disagree"): + opt.load_state_dict(conflicting) + + def test_load_legacy_flattened_param_group_checkpoint(): model_src, opt_src = _run_and_save(factored=False) legacy_sd = _legacy_flattened_param_groups(opt_src.state_dict()) diff --git a/tests/test_dispatch_gating_cpu.py b/tests/test_dispatch_gating_cpu.py index 91089a4..19f7d1e 100644 --- a/tests/test_dispatch_gating_cpu.py +++ b/tests/test_dispatch_gating_cpu.py @@ -125,6 +125,36 @@ def boom(): assert opt._fused_build_ok is None # unprobed +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") +def test_nonfused_cuda_step_never_loads_jit_extensions(monkeypatch): + """The public ``fused=False`` contract is JIT-free even on CUDA. + + Period selection and exact-codebook histogramming have their own optional + extensions. Historically those still built on the first non-fused step, + so selecting the documented pure-PyTorch path could fail on a machine with + an otherwise usable CUDA PyTorch install but no matching nvcc toolchain. + """ + import gefen.kernels.exact_histogram_fused as histogram_mod + import gefen.kernels.period_variance as period_mod + + def unexpected_load(): + raise AssertionError("fused=False must not load a CUDA extension") + + monkeypatch.setattr(gefen_mod, "_ensure_gefen_fused_extension_loaded", unexpected_load) + monkeypatch.setattr(histogram_mod, "_load_extension", unexpected_load) + monkeypatch.setattr(period_mod, "_load_extension", unexpected_load) + + w = nn.Parameter(torch.randn(16, 16, device="cuda")) + opt = Gefen([("w", w)], lr=1e-2, fused=False) + before = w.detach().clone() + (w.square().sum()).backward() + opt.step() + + assert opt._fused_build_ok is None + assert torch.isfinite(w).all() + assert not torch.equal(w.detach(), before) + + # --------------------------------------------------------------------------- # m14 -- empty-string env override means UNSET # --------------------------------------------------------------------------- diff --git a/tests/test_epsilon_muon_param_group_cpu.py b/tests/test_epsilon_muon_param_group_cpu.py new file mode 100644 index 0000000..335e7d6 --- /dev/null +++ b/tests/test_epsilon_muon_param_group_cpu.py @@ -0,0 +1,157 @@ +"""CPU regressions for epsilon safety and GefenMuon group registration.""" + +import pytest +import torch +import torch.nn as nn + +from gefen import Gefen, GefenMuon + + +NON_POSITIVE_OR_NONFINITE = [ + pytest.param(0.0, id="zero"), + pytest.param(-1e-8, id="negative"), + pytest.param(float("nan"), id="nan"), + pytest.param(float("inf"), id="positive-infinity"), + pytest.param(float("-inf"), id="negative-infinity"), +] + + +@pytest.mark.parametrize("eps", NON_POSITIVE_OR_NONFINITE) +def test_gefen_requires_finite_positive_eps(eps): + with pytest.raises(ValueError, match="eps|epsilon"): + Gefen([nn.Parameter(torch.ones(8))], eps=eps, fused=False) + + +@pytest.mark.parametrize("argument", ["eps", "normuon_eps"]) +@pytest.mark.parametrize("value", NON_POSITIVE_OR_NONFINITE) +def test_gefen_muon_requires_finite_positive_epsilon_values(argument, value): + with pytest.raises(ValueError, match=argument): + GefenMuon( + [nn.Parameter(torch.ones(8, 8))], + fused=False, + **{argument: value}, + ) + + +def test_gefen_zero_gradient_step_stays_finite(): + param = nn.Parameter(torch.randn(8)) + before = param.detach().clone() + opt = Gefen([("weight", param)], eps=1e-8, weight_decay=0.0, fused=False) + param.grad = torch.zeros_like(param) + + opt.step() + + assert torch.equal(param, before) + assert torch.isfinite(param).all() + + +def test_gefen_muon_normuon_zero_gradient_step_stays_finite(): + param = nn.Parameter(torch.randn(8, 8)) + before = param.detach().clone() + opt = GefenMuon( + [("weight", param)], + eps=1e-7, + normuon=True, + normuon_eps=1e-8, + weight_decay=0.0, + fused=False, + ) + param.grad = torch.zeros_like(param) + + opt.step() + + assert torch.equal(param, before) + assert torch.isfinite(param).all() + + +def test_gefen_muon_add_param_group_populates_all_options_and_steps(): + torch.manual_seed(0) + first = nn.Parameter(torch.randn(8, 8)) + opt = GefenMuon([("first", first)], weight_decay=0.0, fused=False) + + # Establish the frozen codebook before adding the new group. The new matrix + # must initialize against that existing codebook without missing group keys. + first.grad = torch.randn_like(first) + opt.step() + opt.zero_grad() + + added = nn.Parameter(torch.randn(8, 8)) + opt.add_param_group( + { + "params": [("added", added)], + "lr": 2e-4, + "eps": 2e-7, + "weight_decay": 0.0, + "momentum": 0.8, + "nesterov": False, + "ns_schedule": "tuned3", + "adjust_lr_fn": "match_rms_adamw", + "sharded_mode": "approx", + "fp8_ns": False, + "fp8_ns_compile": False, + "batched_ns": True, + "batched_ns_workspace_bytes": 1 << 20, + "normuon": True, + "normuon_beta2": 0.9, + "normuon_eps": 1e-6, + "cautious": True, + } + ) + + group = opt.param_groups[-1] + assert group["param_names"] == ["added"] + assert group["lr"] == 2e-4 + assert group["eps"] == 2e-7 + assert group["weight_decay"] == 0.0 + assert group["beta1"] == group["momentum"] == 0.8 + assert group["beta2"] == 0.0 + assert group["nesterov"] is False + assert group["ns_steps"] == 3 + assert len(group["ns_coefficients"]) == 3 + assert group["adjust_lr_fn"] == "match_rms_adamw" + assert group["sharded_mode"] == "approx" + assert group["fp8_ns"] is False + assert group["fp8_ns_compile"] is False + assert group["batched_ns"] is True + assert group["batched_ns_workspace_bytes"] == 1 << 20 + assert group["normuon"] is True + assert group["normuon_beta2"] == 0.9 + assert group["normuon_eps"] == 1e-6 + assert group["cautious"] is True + + first.grad = torch.randn_like(first) + added.grad = torch.randn_like(added) + before = added.detach().clone() + opt.step() + + assert "step" in opt.state[added] + assert torch.isfinite(added).all() + assert not torch.equal(added, before) + + +def test_gefen_muon_add_param_group_rejects_non_2d_atomically(): + opt = GefenMuon([nn.Parameter(torch.randn(8, 8))], fused=False) + valid = nn.Parameter(torch.randn(8, 8)) + invalid = nn.Parameter(torch.randn(8)) + group_count = len(opt.param_groups) + + with pytest.raises(ValueError, match="only supports 2D"): + opt.add_param_group({"params": [("valid", valid), ("invalid", invalid)]}) + + assert len(opt.param_groups) == group_count + assert valid not in opt.state + + +@pytest.mark.parametrize("normuon_eps", NON_POSITIVE_OR_NONFINITE) +def test_gefen_muon_add_param_group_rejects_bad_normuon_eps_atomically( + normuon_eps, +): + opt = GefenMuon([nn.Parameter(torch.randn(8, 8))], fused=False) + added = nn.Parameter(torch.randn(8, 8)) + group_count = len(opt.param_groups) + + with pytest.raises(ValueError, match="normuon_eps"): + opt.add_param_group({"params": [added], "normuon_eps": normuon_eps}) + + assert len(opt.param_groups) == group_count + assert added not in opt.state From 83019bdba9446fe8183b305a581489f45f875a11 Mon Sep 17 00:00:00 2001 From: thad0ctor Date: Sun, 12 Jul 2026 01:57:40 -0700 Subject: [PATCH 02/14] Add replica-exact fused optimizer mode --- README.md | 9 + src/gefen/gefen.py | 145 +++++++++-- src/gefen/gefen_muon.py | 12 + src/gefen/hybrid.py | 11 +- tests/test_deterministic_mode.py | 412 +++++++++++++++++++++++++++++++ 5 files changed, 563 insertions(+), 26 deletions(-) create mode 100644 tests/test_deterministic_mode.py diff --git a/README.md b/README.md index 9684f5e..0e90950 100644 --- a/README.md +++ b/README.md @@ -118,6 +118,14 @@ Gefen drops into standard distributed training like any other PyTorch optimizer, > **DeepSpeed ZeRO config.** Set `"zero_allow_untested_optimizer": true` and leave the config's `optimizer` section unset. With optimizer CPU-offload, also set `"zero_force_ds_cpu_optimizer": false` — otherwise raw DeepSpeed refuses to initialize, and accelerate-based launchers (axolotl) silently swap in DeepSpeed's own CPU Adam. ZeRO steps flattened 1-D partitions, so `GefenMuon`/`GefenMuonHybrid` raise a clear error under ZeRO; use FSDP2, DDP, or single-GPU for the Muon family. +## Replica-exact fused updates (`deterministic`) + +Set `deterministic=True` when data-parallel replicas must remain bit-exact on homogeneous GPUs. Automatic block periods are selected with fixed-order GPU reductions instead of the faster atomic CUDA search, block-vmean parameters remain fused through the fixed-order v1 CUDA reduction, and factored-v parameters use the deterministic decomposed factored update instead of the fused stats kernel's unordered floating-point atomics. The default is `False`, so existing performance routing is unchanged. Tagged checkpoints must resume with the same deterministic policy; legacy checkpoints without a tag remain loadable. Plain Gefen does not allow `deterministic=True`, `factored_v_2d=True`, and `stochastic_round=True` together because the deterministic factored fallback uses nearest-codeword quantization. + +```python +optimizer = Gefen(model.named_parameters(), lr=3e-4, fused=True, deterministic=True) +``` + ## CUDA Graphs & torch.compile (`capturable`) All three optimizers accept `capturable=True` (same meaning as `torch.optim`'s argument): `opt.step()` can then be captured in a `torch.cuda.CUDAGraph` or wrapped in `torch.compile(mode="reduce-overhead")` at no step-time cost — and the compiled hybrid step is about 10% faster than eager. Usage, caveats, and measured numbers: [COMPATIBILITY.md](https://github.com/thad0ctor/Gefen-X/blob/main/COMPATIBILITY.md#cuda-graphs--torchcompile-capturable). @@ -457,6 +465,7 @@ This keeps both halves on quantized Gefen state and is the low-memory recommenda | `normuon` | `True` | free per-neuron 2nd moment on the NS output; recovers tuned3 quality in SFT — [details](#quality-lever-per-neuron-2nd-moment-on-the-newton-schulz-output-normuon) | keep on for SFT; disable for classic pretraining | | `backup_2d_period_one` | `False` | per-element 2nd moment on a Gefen-backed embedding/LM head — extra memory — [details](#experimental-lever-per-element-gefen-backup-state-on-embed--lm-head-backup_2d_period_one) | Gefen backup only; AdamW already keeps per-element moments | | `stochastic_round` | `False` | unbiased rounding for the 8-bit momentum (free, loss-neutral) | optional | +| `deterministic` | `False` | replica-exact fused routing on homogeneous GPUs; persists in child checkpoints | enable when distributed replica hashes must match bit-for-bit | | `capturable` | `False` | CUDA-graph-capturable `step()`, like `torch.optim`'s `capturable` — [details](https://github.com/thad0ctor/Gefen-X/blob/main/COMPATIBILITY.md#cuda-graphs--torchcompile-capturable) | turn on to capture `step()` in a `torch.cuda.CUDAGraph` | It supports `step()`, `zero_grad()`, `state_dict()`/`load_state_dict()`, and LR schedulers (e.g. `torch.optim.lr_scheduler.StepLR(optimizer, ...)`) like any optimizer. Because it splits params at construction rather than taking a single iterable, build it yourself and hand it to the Hugging Face `Trainer` via `optimizers=` (not `optimizer_cls_and_kwargs`): diff --git a/src/gefen/gefen.py b/src/gefen/gefen.py index e44dd7c..5f53fde 100644 --- a/src/gefen/gefen.py +++ b/src/gefen/gefen.py @@ -550,6 +550,16 @@ class Gefen(torch.optim.Optimizer): nearest-rounding accumulates in the EMA momentum over a long horizon. Only the fused CUDA automatic-step kernels honor it; a warning fires when the fused path is off. + deterministic (bool, default False): require replica-exact update + arithmetic on identical devices and inputs. Automatic block periods + use fixed-order GPU reductions instead of the atomic CUDA search. + Fused block-vmean parameters stay fused but use the fixed-order v1 + reduction instead of the atomic v2-full path. Fused factored-v + parameters use the deterministic decomposed factored update because + the fast fused stats kernel accumulates column sums with unordered + floating-point atomics. This is an explicit throughput-for- + reproducibility mode; the default keeps the existing performance + routing unchanged. capturable (bool, default False): CUDA-graph capturability (mirrors torch.optim's argument). Everything that varies across steps -- step counters, bias corrections, kernel scalars -- lives in device @@ -579,9 +589,19 @@ def __init__( factored_v_2d: bool = True, codebook_refresh_every: int = 0, stochastic_round: bool = False, + deterministic: bool = False, capturable: bool = False, verbose: bool = False, ): + if not isinstance(deterministic, bool): + raise TypeError("deterministic must be a bool") + if deterministic and factored_v_2d and stochastic_round: + raise ValueError( + "deterministic=True with factored_v_2d=True cannot honor " + "stochastic_round=True: the replica-exact factored path uses " + "deterministic nearest-codeword quantization. Disable stochastic " + "rounding or factored_v_2d." + ) if fused and not torch.cuda.is_available(): warnings.warn( "Gefen optimizer got fused=True, but CUDA is not available. " @@ -609,6 +629,7 @@ def __init__( s.lower() for s in period_one_substrings ) self._factored_v_2d = factored_v_2d + self._deterministic = deterministic if codebook_refresh_every < 0: raise ValueError( "codebook_refresh_every must be >= 0 but is: {}".format( @@ -1833,20 +1854,31 @@ def _predict_period_from_grad_sq( self, param_name: str, param: torch.Tensor, grad: torch.Tensor ) -> int: backend = _resolve_find_period_backend(grad) + # The dedicated CUDA period-search kernel finishes each candidate's + # cross-block variance reduction with a floating-point atomicAdd. Its + # scheduling-dependent last bits normally do not affect the selected + # period, but a near tie can route otherwise-identical replicas to + # different state geometries before their first update. Deterministic + # mode therefore uses PyTorch's fixed-order CUDA reductions, including + # when a process-wide FIND_PERIOD_BACKEND override requested the faster + # atomic kernel. + if self._deterministic and backend == "cuda_kernel": + backend = "gpu" # ``fused=False`` is an explicit request for a pure-PyTorch optimizer # step, so it must not cross a different lazy-JIT boundary merely to # choose the block period. A failed fused-kernel probe has the same - # requirement. Keep honoring an explicit FIND_PERIOD_BACKEND override; - # callers that deliberately select ``cuda_kernel`` accept its JIT. + # requirement. Outside deterministic mode, keep honoring an explicit + # FIND_PERIOD_BACKEND override; callers that deliberately select + # ``cuda_kernel`` accept its JIT. + grad_work = grad.to_local() if hasattr(grad, "to_local") else grad if ( backend == "cuda_kernel" and FIND_PERIOD_BACKEND is None - and grad.device.type == "cuda" + and grad_work.device.type == "cuda" and (not self.fused or self._fused_build_ok is False) ): backend = "gpu" if backend == "cuda_kernel": - grad_work = grad.to_local() if hasattr(grad, "to_local") else grad grad_flat = grad_work.detach().reshape(-1) try: return find_period_by_block_variance( @@ -1878,11 +1910,13 @@ def _predict_period_from_grad_sq( ) if backend == "cpu": - period_input = grad.detach().float().square().reshape(-1).cpu().numpy() + period_input = ( + grad_work.detach().float().square().reshape(-1).cpu().numpy() + ) elif backend == "gpu": - if grad.device.type != "cuda": + if grad_work.device.type != "cuda": raise ValueError("FIND_PERIOD_BACKEND='gpu' requires a CUDA tensor") - period_input = grad.detach().float().square().reshape(-1) + period_input = grad_work.detach().float().square().reshape(-1) else: raise ValueError("Unexpected FIND_PERIOD_BACKEND: {}".format(backend)) return find_period_by_block_variance( @@ -2102,6 +2136,7 @@ def _step_automatic_factored( self._use_fused_gefen_automatic_step() and p.is_cuda and p.dtype != torch.float64 + and not self._deterministic ): # Fused path. Phase 1 of the kernel computes the per-block absmax # AND the raw row/col grad^2 sums in ONE pass over grad + m_sign; @@ -2197,9 +2232,13 @@ def _step_automatic_factored( bias_correction_1 = 1 - beta1 ** state["step"] bias_correction_2 = 1 - beta2 ** state["factored_step"] - # Decomposed fallback (CPU / fused disabled). Row/col mean-square EMAs - # over bounded fp32 row-chunks (one chunk serves both reductions, so - # the transient stays ~a few chunk sizes), then whole-tensor torch ops. + # Decomposed fallback (CPU / fused disabled / deterministic factored-v). + # The fast factored CUDA stats kernel uses unordered floating-point + # atomicAdd operations for row/column partials, so deterministic mode + # deliberately reaches this fixed-order reduction path while block-vmean + # parameters continue using the fused v1 kernel. Row/col mean-square + # EMAs run over bounded fp32 row-chunks (one chunk serves both reductions, + # so the transient stays ~a few chunk sizes), then whole-tensor torch ops. grad2d = grad.detach() row_ms = torch.empty(rows, dtype=torch.float32, device=p.device) col_sq = torch.zeros(cols, dtype=torch.float32, device=p.device) @@ -2222,7 +2261,7 @@ def _step_automatic_factored( state["v_row"].mul_(beta2).add_(row_ms, alpha=1 - beta2) state["v_col"].mul_(beta2).add_(col_sq, alpha=1 - beta2) - # Decomposed fallback (CPU / fused disabled): whole-tensor torch ops. + # Decomposed fallback: whole-tensor torch ops. # The kernel above is the canonical numerics; this matches it within # float tolerance (different op association), not bit-for-bit. v_hat = torch.outer(state["v_row"], state["v_col"]).div_( @@ -2813,7 +2852,13 @@ def _step_automatic( and grad_view.is_cuda and grad_view.dtype != torch.float64 ) - route_v2 = _should_use_v2_full( + # v2-full forms each block's grad^2 sum through floating-point atomics. + # The result is convergence-equivalent but not replica-exact because + # independent launches may observe a different accumulation order. + # Deterministic mode keeps the fused update and selects v1-full, whose + # fixed thread tree produces bit-identical vmean/parameter writes for + # identical inputs on a homogeneous architecture. + route_v2 = False if self._deterministic else _should_use_v2_full( grad_view.shape[0], automatic_period, grad_view.device ) if ( @@ -3197,10 +3242,11 @@ def state_dict(self): Beyond the base ``torch.optim.Optimizer.state_dict`` contents (per-param state keyed by index and caller-preserved ``param_groups``), the dict carries each parameter's stable ``name`` in its state, plus - ``gefen_global_step`` and the frozen exact-DP ``gefen_codebook`` (without - it a resume would re-learn the codebook and desync the restored block - periods). Per-step scratch buffers are stripped, and non-owning state - views are compacted to tight clones so checkpoints stay small. + ``gefen_global_step``, the frozen exact-DP ``gefen_codebook`` (without it + a resume would re-learn the codebook and desync the restored block + periods), and the replica-determinism policy. Per-step scratch buffers + are stripped, and non-owning state views are compacted to tight clones + so checkpoints stay small. Only state entries whose parameter is still reachable from ``param_groups`` are serialized. Wrappers like DeepSpeed ZeRO-1/2 @@ -3287,6 +3333,7 @@ def _compact(value): for pid, pstate in state_dict["state"].items() } state_dict["gefen_global_step"] = self._gefen_global_step + state_dict["gefen_deterministic"] = self._deterministic # The exact-DP codebook is learned once on the first step and then frozen # for the rest of the run. It is not per-param state, so persist it # explicitly; without it resume re-learns the codebook (see @@ -3307,6 +3354,7 @@ def _compact(value): "format_version": 1, "global_step": self._gefen_global_step, "codebook": self._gefen_codebook, + "deterministic": self._deterministic, } for group in state_dict["param_groups"]: group["_gefen_checkpoint_metadata"] = checkpoint_metadata @@ -3378,13 +3426,16 @@ def load_state_dict(self, state_dict): """Restore optimizer state saved by :meth:`state_dict`. On top of the base load this restores ``gefen_global_step`` and the - frozen codebook, and undoes the base class's dtype coercion of aux - state: ``torch.optim`` casts every floating state tensor to the owning - param's dtype, which would corrupt Gefen's fp32 moments and uint8 - momentum indices on bf16 models, so each aux tensor is written back in - its ORIGINAL saved dtype. Step counters are normalized to this - optimizer's ``capturable`` mode (host ints vs 0-dim device tensors), so - checkpoints are portable across the toggle in either direction. + frozen codebook, verifies that an explicitly tagged checkpoint uses the + same replica-determinism policy as this optimizer, and undoes the base + class's dtype coercion of aux state: ``torch.optim`` casts every floating + state tensor to the owning param's dtype, which would corrupt Gefen's + fp32 moments and uint8 momentum indices on bf16 models, so each aux + tensor is written back in its ORIGINAL saved dtype. Step counters are + normalized to this optimizer's ``capturable`` mode (host ints vs 0-dim + device tensors), so checkpoints are portable across that toggle in + either direction. Legacy checkpoints with no deterministic tag remain + loadable and retain the live optimizer's configured policy. """ # Shallow-copy so the pop()s below don't mutate the caller's checkpoint # dict (torch.optim.Optimizer.load_state_dict leaves its input intact); @@ -3403,6 +3454,13 @@ def load_state_dict(self, state_dict): gefen_global_step = state_dict.pop("gefen_global_step", None) gefen_codebook = state_dict.pop("gefen_codebook", None) + has_top_level_deterministic = "gefen_deterministic" in state_dict + gefen_deterministic = state_dict.pop("gefen_deterministic", None) + if has_top_level_deterministic and type(gefen_deterministic) is not bool: + raise ValueError( + "Gefen checkpoint top-level deterministic policy must be a bool, " + "got {!r}".format(gefen_deterministic) + ) if group_metadata: if len(group_metadata) != len(state_dict["param_groups"]): raise ValueError( @@ -3415,6 +3473,15 @@ def load_state_dict(self, state_dict): first_metadata.get("format_version") ) ) + for metadata in group_metadata: + if ( + "deterministic" in metadata + and type(metadata["deterministic"]) is not bool + ): + raise ValueError( + "Gefen checkpoint parameter-group deterministic policy " + "must be a bool, got {!r}".format(metadata["deterministic"]) + ) for metadata in group_metadata[1:]: same_step = metadata.get("global_step") == first_metadata.get( "global_step" @@ -3429,13 +3496,17 @@ def load_state_dict(self, state_dict): and torch.equal(left_codebook, right_codebook) ) ) - if not same_step or not same_codebook: + same_deterministic = metadata.get( + "deterministic" + ) == first_metadata.get("deterministic") + if not same_step or not same_codebook or not same_deterministic: raise ValueError( "Gefen checkpoint parameter groups carry inconsistent " "optimizer metadata" ) metadata_step = first_metadata.get("global_step", 0) metadata_codebook = first_metadata.get("codebook") + metadata_deterministic = first_metadata.get("deterministic") if gefen_global_step is None: gefen_global_step = metadata_step elif gefen_global_step != metadata_step: @@ -3455,6 +3526,32 @@ def load_state_dict(self, state_dict): raise ValueError( "Gefen checkpoint top-level and parameter-group codebooks disagree" ) + if gefen_deterministic is None: + gefen_deterministic = metadata_deterministic + elif ( + metadata_deterministic is not None + and gefen_deterministic != metadata_deterministic + ): + raise ValueError( + "Gefen checkpoint top-level and parameter-group deterministic " + "policies disagree" + ) + if gefen_deterministic is not None: + if type(gefen_deterministic) is not bool: + raise ValueError( + "Gefen checkpoint deterministic policy must be a bool, got " + "{!r}".format(gefen_deterministic) + ) + if gefen_deterministic != self._deterministic: + raise ValueError( + "Gefen checkpoint deterministic={!r}, but this optimizer was " + "constructed with deterministic={!r}. Resuming under a " + "different replica-determinism policy requires an intentional " + "state migration or a fresh optimizer restart; refusing to " + "change the policy silently.".format( + gefen_deterministic, self._deterministic + ) + ) if gefen_global_step is None: gefen_global_step = 0 has_quantized_momentum = any( diff --git a/src/gefen/gefen_muon.py b/src/gefen/gefen_muon.py index 9992718..98f9b8a 100644 --- a/src/gefen/gefen_muon.py +++ b/src/gefen/gefen_muon.py @@ -508,6 +508,11 @@ class GefenMuon(Gefen): opt-in path's approximate numerical result. stochastic_round: stochastically round the 8-bit momentum quantization (debiases it; throughput-neutral opt-in). + deterministic: persist and enforce Gefen's replica-determinism policy + across checkpoints. GefenMuon's fused momentum reduction is already + replica-exact on homogeneous GPUs (it uses order-independent max + reductions); the flag is forwarded so hybrid Muon/backup children + share one explicit policy. normuon: NorMuon-style per-row 2nd-moment normalization of the NS output (default False here; GefenMuonHybrid turns it on). normuon_beta2: EMA coefficient in [0, 1) for the normuon row statistic. @@ -539,6 +544,7 @@ def __init__( batched_ns: bool = False, batched_ns_workspace_bytes: int = BATCHED_NS_DEFAULT_WORKSPACE_BYTES, stochastic_round: bool = False, + deterministic: bool = False, normuon: bool = False, normuon_beta2: float = 0.95, normuon_eps: float = 1e-8, @@ -719,7 +725,13 @@ def __init__( eps=eps, weight_decay=weight_decay, fused=fused, + # GefenMuon owns its second-moment/update pipeline; plain Gefen's + # factored-v routing is unused here. Pinning it off also keeps the + # valid deterministic+stochastic-round Muon combination distinct + # from plain Gefen's deterministic factored fallback. + factored_v_2d=False, stochastic_round=stochastic_round, + deterministic=deterministic, capturable=capturable, verbose=verbose, ) diff --git a/src/gefen/hybrid.py b/src/gefen/hybrid.py index d7fc529..ba315b9 100644 --- a/src/gefen/hybrid.py +++ b/src/gefen/hybrid.py @@ -252,6 +252,7 @@ def __init__( batched_ns=False, batched_ns_workspace_bytes=256 << 20, stochastic_round=False, + deterministic=False, normuon=True, normuon_beta2=0.95, normuon_eps=1e-8, @@ -318,9 +319,13 @@ def __init__( Note ``normuon=True`` and ``ns_schedule="tuned3"`` are the hybrid's defaults, unlike raw GefenMuon. capturable: forwarded to both halves. ``stochastic_round`` and - ``verbose`` are forwarded to Gefen children; with an AdamW - backup they apply only to the Muon half. + ``deterministic`` are forwarded to Gefen children; with an + AdamW backup they apply only to the Muon half. ``verbose`` is + likewise forwarded to Gefen children. """ + if not isinstance(deterministic, bool): + raise TypeError("deterministic must be a bool") + self._deterministic = deterministic if backup_named_params is None: # Single-argument convenience form: the first arg is a model or a # named-param iterable to split internally. @@ -438,6 +443,7 @@ def __init__( batched_ns=batched_ns, batched_ns_workspace_bytes=batched_ns_workspace_bytes, stochastic_round=stochastic_round, + deterministic=deterministic, normuon=normuon, normuon_beta2=normuon_beta2, normuon_eps=normuon_eps, @@ -506,6 +512,7 @@ def _is_no_decay(name): # embedding/head, not the untested factored-v combination. factored_v_2d=False, stochastic_round=stochastic_round, + deterministic=deterministic, capturable=capturable, verbose=verbose, ) diff --git a/tests/test_deterministic_mode.py b/tests/test_deterministic_mode.py new file mode 100644 index 0000000..84c2067 --- /dev/null +++ b/tests/test_deterministic_mode.py @@ -0,0 +1,412 @@ +"""Replica-exact routing and checkpoint coverage for ``deterministic=True``.""" + +import copy +import hashlib +import multiprocessing as mp +import traceback +from unittest import mock + +import pytest +import torch +import torch.nn as nn + +import gefen.gefen as gefen_mod +from gefen import Gefen, GefenMuon, GefenMuonHybrid + + +def _deep_equal(actual, expected): + if type(actual) is not type(expected): + return False + if isinstance(actual, dict): + return set(actual) == set(expected) and all( + _deep_equal(actual[key], expected[key]) for key in actual + ) + if isinstance(actual, (list, tuple)): + return len(actual) == len(expected) and all( + _deep_equal(left, right) for left, right in zip(actual, expected) + ) + if torch.is_tensor(actual): + return actual.dtype == expected.dtype and torch.equal(actual, expected) + return actual == expected + + +def _stepped_cpu_optimizer(*, deterministic): + generator = torch.Generator().manual_seed(11) + param = nn.Parameter(torch.randn(16, 12, generator=generator)) + optimizer = Gefen( + [("weight", param)], + lr=1e-3, + fused=False, + deterministic=deterministic, + ) + param.grad = torch.randn(param.shape, generator=generator) + optimizer.step() + return param, optimizer + + +def test_deterministic_configuration_validation(): + param = nn.Parameter(torch.ones(2, 2)) + with pytest.raises(TypeError, match="deterministic must be a bool"): + Gefen([param], fused=False, deterministic=1) + with pytest.raises(ValueError, match="cannot honor stochastic_round"): + Gefen( + [param], + fused=False, + deterministic=True, + factored_v_2d=True, + stochastic_round=True, + ) + + # Muon's momentum-only fused kernel is replica-exact and supports its + # stateless, step-seeded stochastic rounding under the same policy. + muon = GefenMuon( + [("weight", nn.Parameter(torch.ones(2, 2)))], + fused=False, + deterministic=True, + stochastic_round=True, + ) + assert muon._deterministic is True + assert muon._factored_v_2d is False + + +def test_deterministic_checkpoint_tags_and_matching_roundtrip(): + source_param, source = _stepped_cpu_optimizer(deterministic=True) + checkpoint = copy.deepcopy(source.state_dict()) + assert checkpoint["gefen_deterministic"] is True + assert all( + group["_gefen_checkpoint_metadata"]["deterministic"] is True + for group in checkpoint["param_groups"] + ) + + target_param = nn.Parameter(source_param.detach().clone()) + target = Gefen( + [("weight", target_param)], + lr=1e-3, + fused=False, + deterministic=True, + ) + target.load_state_dict(checkpoint) + target_param.grad = torch.full_like(target_param, 0.125) + source_param.grad = torch.full_like(source_param, 0.125) + target.step() + source.step() + assert torch.equal(target_param, source_param) + + +def test_deterministic_checkpoint_mismatch_rejected_before_mutation(): + _, source = _stepped_cpu_optimizer(deterministic=True) + checkpoint = copy.deepcopy(source.state_dict()) + + target_param = nn.Parameter(torch.zeros(16, 12)) + target = Gefen( + [("weight", target_param)], + lr=1e-3, + fused=False, + deterministic=False, + ) + before = copy.deepcopy(target.state_dict()) + with pytest.raises(ValueError, match="intentional state migration"): + target.load_state_dict(checkpoint) + assert _deep_equal(target.state_dict(), before) + + # DCP keeps conventional top-level keys only. The mirrored group tag must + # preserve the same safety check when the top-level Gefen extras are gone. + dcp_style = { + "state": checkpoint["state"], + "param_groups": checkpoint["param_groups"], + } + with pytest.raises(ValueError, match="intentional state migration"): + target.load_state_dict(dcp_style) + assert _deep_equal(target.state_dict(), before) + + +@pytest.mark.parametrize("invalid_tag", [None, 0, 1, "true"]) +def test_deterministic_checkpoint_top_level_tag_requires_actual_bool(invalid_tag): + _, source = _stepped_cpu_optimizer(deterministic=True) + checkpoint = copy.deepcopy(source.state_dict()) + checkpoint["gefen_deterministic"] = invalid_tag + + target = Gefen( + [("weight", nn.Parameter(torch.zeros(16, 12)))], + lr=1e-3, + fused=False, + deterministic=True, + ) + before = copy.deepcopy(target.state_dict()) + with pytest.raises(ValueError, match="top-level deterministic policy must be a bool"): + target.load_state_dict(checkpoint) + assert _deep_equal(target.state_dict(), before) + + +def test_deterministic_checkpoint_group_tag_requires_actual_bool(): + _, source = _stepped_cpu_optimizer(deterministic=True) + checkpoint = copy.deepcopy(source.state_dict()) + checkpoint["param_groups"][0]["_gefen_checkpoint_metadata"][ + "deterministic" + ] = 1 + + target = Gefen( + [("weight", nn.Parameter(torch.zeros(16, 12)))], + lr=1e-3, + fused=False, + deterministic=True, + ) + before = copy.deepcopy(target.state_dict()) + with pytest.raises( + ValueError, match="parameter-group deterministic policy must be a bool" + ): + target.load_state_dict(checkpoint) + assert _deep_equal(target.state_dict(), before) + + +def test_legacy_untagged_checkpoint_loads_under_live_policy(): + _, source = _stepped_cpu_optimizer(deterministic=False) + legacy = copy.deepcopy(source.state_dict()) + legacy.pop("gefen_deterministic") + for group in legacy["param_groups"]: + group["_gefen_checkpoint_metadata"].pop("deterministic") + + target = Gefen( + [("weight", nn.Parameter(torch.zeros(16, 12)))], + lr=1e-3, + fused=False, + deterministic=True, + ) + target.load_state_dict(legacy) + assert target._deterministic is True + + +def test_hybrid_plumbs_one_deterministic_policy_to_gefen_children(): + model = nn.Sequential(nn.Linear(8, 8), nn.LayerNorm(8)) + optimizer = GefenMuonHybrid.from_model( + model, + lr=1e-3, + backup_optimizer="gefen", + fused=False, + deterministic=True, + ) + assert optimizer._deterministic is True + assert optimizer.muon is not None and optimizer.muon._deterministic is True + assert optimizer.backup is not None and optimizer.backup._deterministic is True + # Keep the established nested hybrid schema; each Gefen child carries the + # policy in its own ordinary state dict and mirrored group metadata. + assert set(optimizer.state_dict()) == {"muon", "backup", "backup_optimizer"} + + +def _fixed_codebook(device): + return torch.linspace(-1.0, 1.0, 256, device=device, dtype=torch.float32) + + +def _tensor_digest(tensor): + tensor = tensor.detach().contiguous().cpu() + payload = tensor.view(torch.uint8).numpy().tobytes() + header = "{}:{}:".format(tuple(tensor.shape), tensor.dtype).encode() + return hashlib.sha256(header + payload).hexdigest() + + +def _result_digest(param, optimizer): + state = optimizer.state[param] + result = {"param": _tensor_digest(param)} + for key in sorted(state): + value = state[key] + if torch.is_tensor(value): + result[key] = _tensor_digest(value) + elif isinstance(value, (bool, int, float, str)): + result[key] = repr(value) + result["codebook"] = _tensor_digest(optimizer._gefen_codebook) + return result + + +def _run_fused_deterministic_cases(device_index): + device = torch.device("cuda", device_index) + + generator = torch.Generator(device="cpu").manual_seed(1234) + factored_init = torch.randn(384, 128, generator=generator) * 0.02 + factored_grad = torch.randn(384, 128, generator=generator) * 1e-3 + factored_param = nn.Parameter(factored_init.to(device)) + factored = Gefen( + [("linear_qkv.weight", factored_param)], + lr=1e-3, + weight_decay=0.01, + fused=True, + factored_v_2d=True, + force_2d_period_one=True, + deterministic=True, + ) + factored._gefen_codebook = _fixed_codebook(device) + factored_param.grad = factored_grad.to(device) + factored.step() + + block_init = torch.randn(384, 128, generator=generator) * 0.02 + block_grad = torch.randn(384, 128, generator=generator) * 1e-3 + block_param = nn.Parameter(block_init.to(device)) + block = Gefen( + [("projection.weight", block_param)], + lr=1e-3, + weight_decay=0.01, + fused=True, + factored_v_2d=False, + deterministic=True, + ) + block._gefen_codebook = _fixed_codebook(device) + # Two very large blocks are normally routed to atomic v2-full. Pre-seeding + # the period also removes first-step period-search variability from this + # update-kernel regression. + block.state[block_param]["automatic_period"] = 24_576 + block_param.grad = block_grad.to(device) + block.step() + + torch.cuda.synchronize(device) + return { + "factored": _result_digest(factored_param, factored), + "block": _result_digest(block_param, block), + } + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_deterministic_fused_repeated_runs_are_bit_exact(): + first = _run_fused_deterministic_cases(0) + second = _run_fused_deterministic_cases(0) + assert first == second + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_deterministic_routes_fused_block_v1_and_factored_fallback(): + device = torch.device("cuda", 0) + generator = torch.Generator(device="cpu").manual_seed(55) + + block_param = nn.Parameter(torch.randn(64, 32, generator=generator).to(device)) + block = Gefen( + [("block", block_param)], + fused=True, + factored_v_2d=False, + deterministic=True, + ) + block._gefen_codebook = _fixed_codebook(device) + block.state[block_param]["automatic_period"] = 1_024 + block_param.grad = torch.randn(64, 32, generator=generator).to(device) + v1_update = gefen_mod._automatic_gefen_fused_full_update_cuda + with mock.patch.object( + gefen_mod, + "_should_use_v2_full", + side_effect=AssertionError("deterministic mode consulted v2 routing"), + ), mock.patch.object( + gefen_mod, + "_automatic_gefen_fused_full_update_cuda", + wraps=v1_update, + ) as v1_spy: + block.step() + assert v1_spy.call_count == 1 + + factored_param = nn.Parameter(torch.randn(64, 32, generator=generator).to(device)) + factored = Gefen( + [("factored", factored_param)], + fused=True, + factored_v_2d=True, + force_2d_period_one=True, + deterministic=True, + ) + factored._gefen_codebook = _fixed_codebook(device) + factored_param.grad = torch.randn(64, 32, generator=generator).to(device) + with mock.patch.object( + gefen_mod, + "_gefen_factored_update_cuda", + side_effect=AssertionError("deterministic mode used atomic factored stats"), + ): + factored.step() + assert "v_row" in factored.state[factored_param] + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_deterministic_automatic_period_avoids_atomic_cuda_backend(): + device = torch.device("cuda", 0) + generator = torch.Generator(device="cpu").manual_seed(88) + param = nn.Parameter(torch.randn(96, 64, generator=generator).to(device)) + optimizer = Gefen( + [("weight", param)], + lr=1e-3, + fused=True, + deterministic=True, + ) + param.grad = torch.randn(param.shape, generator=generator).to(device) + + period_search = gefen_mod.find_period_by_block_variance + with mock.patch.object( + gefen_mod, "FIND_PERIOD_BACKEND", "cuda_kernel" + ), mock.patch.object( + gefen_mod, "find_period_by_block_variance", wraps=period_search + ) as period_spy: + optimizer.step() + + assert period_spy.call_count >= 1 + assert all(call.kwargs["backend"] == "gpu" for call in period_spy.call_args_list) + + +def _run_fused_automatic_first_step(device_index): + device = torch.device("cuda", device_index) + generator = torch.Generator(device="cpu").manual_seed(2027) + initial = torch.randn(192, 128, generator=generator) * 0.02 + grad = torch.randn(192, 128, generator=generator) * 1e-3 + param = nn.Parameter(initial.to(device)) + optimizer = Gefen( + [("weight", param)], + lr=1e-3, + weight_decay=0.01, + fused=True, + deterministic=True, + ) + param.grad = grad.to(device) + optimizer.step() + torch.cuda.synchronize(device) + return _result_digest(param, optimizer) + + +def _replica_worker(rank, queue, case): + try: + torch.cuda.set_device(rank) + if case == "preseeded": + result = _run_fused_deterministic_cases(rank) + elif case == "automatic_first_step": + result = _run_fused_automatic_first_step(rank) + else: + raise ValueError("unknown deterministic replica case: {}".format(case)) + queue.put(("ok", rank, result)) + except Exception: + queue.put(("error", rank, traceback.format_exc())) + raise + + +def _run_two_process_replica_case(case): + capabilities = [torch.cuda.get_device_capability(index) for index in range(2)] + if capabilities[0] != capabilities[1]: + pytest.skip("replica-exact test requires two GPUs with one compute capability") + + context = mp.get_context("spawn") + queue = context.Queue() + processes = [ + context.Process(target=_replica_worker, args=(rank, queue, case)) + for rank in range(2) + ] + for process in processes: + process.start() + messages = [queue.get(timeout=120) for _ in processes] + for process in processes: + process.join(timeout=120) + assert not process.is_alive() + assert process.exitcode == 0 + + errors = [message for message in messages if message[0] == "error"] + assert not errors, errors + return {rank: result for _, rank, result in messages} + + +@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="two CUDA devices required") +def test_deterministic_fused_two_process_replicas_are_bit_exact(): + results = _run_two_process_replica_case("preseeded") + assert results[0] == results[1] + + +@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="two CUDA devices required") +def test_deterministic_fused_first_step_replicas_are_bit_exact(): + results = _run_two_process_replica_case("automatic_first_step") + assert results[0] == results[1] From 077ee7afc8088a8a67420498c030af45dbff635d Mon Sep 17 00:00:00 2001 From: thad0ctor Date: Sun, 12 Jul 2026 02:54:58 -0700 Subject: [PATCH 03/14] Fix CUDA graph checkpoint step tracking --- CHANGELOG.md | 9 ++ COMPATIBILITY.md | 7 +- README.md | 2 +- src/gefen/gefen.py | 132 ++++++++++++++-- src/gefen/hybrid.py | 32 +++- tests/test_capturable.py | 287 ++++++++++++++++++++++++++++++++++- tests/test_validation_cpu.py | 12 ++ 7 files changed, 462 insertions(+), 19 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 8a11f94..69014ea 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,15 @@ All notable changes to this project are documented here. This project adheres to [Semantic Versioning](https://semver.org/). +## [Unreleased] + +Correctness and compatibility: + +- Add `deterministic=True` to `Gefen`, `GefenMuon`, and `GefenMuonHybrid` for replica-exact fused routing on homogeneous GPUs. Automatic periods use fixed-order reductions, block-vmean parameters use the deterministic fused v1 path, factored-v parameters use the decomposed deterministic update, and tagged checkpoints enforce the saved policy. +- Capturable optimizers maintain device-resident global-step counters on every parameter device. CUDA-graph replays now serialize the true global step, including steps with no gradients, so stochastic-rounding checkpoints resume with the correct seed. +- Checkpoint loading preserves compact optimizer-state dtypes for bf16 parameters, validates frozen codebooks and hybrid backend metadata, and keeps legacy untagged checkpoints loadable. +- Reject host-driven gradient-histogram output under `capturable=True`, matching the existing periodic-codebook-refresh guard. + ## [0.3.0] - 2026-07-11 Lands the Muon optimization suite (#62) and validated DeepSpeed ZeRO support (#64): a selectable AdamW backup for the hybrid, an opt-in batched Newton-Schulz experiment, faster fused Muon momentum kernels, and plain `Gefen` as a validated DeepSpeed ZeRO 1-3 client optimizer with bit-exact checkpoint resume. diff --git a/COMPATIBILITY.md b/COMPATIBILITY.md index 36e4b5e..eecb7b0 100644 --- a/COMPATIBILITY.md +++ b/COMPATIBILITY.md @@ -87,7 +87,7 @@ Setup: from scratch — MNIST (paper recipe, 3 seeds), CIFAR-10 ResNet-18; fine- ## CUDA Graphs & torch.compile (`capturable`) -`Gefen`, `GefenMuon`, and `GefenMuonHybrid` accept `capturable=True` (same meaning as `torch.optim`'s argument): step counters, bias corrections, and a tensor `lr` live on the GPU, so `opt.step()` stays correct inside a replayed CUDA graph or a compiled region. The default `capturable=False` is bit-identical to previous behavior, and capturing with it raises instead of silently freezing the step counters. +`Gefen`, `GefenMuon`, and `GefenMuonHybrid` accept `capturable=True` (same meaning as `torch.optim`'s argument): step counters, bias corrections, the optimizer-global checkpoint counter, and a tensor `lr` live on the GPU, so `opt.step()` stays correct inside a replayed CUDA graph or a compiled region. The default `capturable=False` is bit-identical to previous behavior, and capturing with it raises instead of silently freezing the step counters. Measured step times (386M-parameter census, RTX 3090 Ti, tail-100 mean over 500 CUDA-event-timed steps): @@ -119,8 +119,9 @@ compiled_step() Caveats: - `dynamic=False` is required — dynamic shapes trip a dynamo symbolic-shapes bug (torch 2.12) on the per-param optimizer state. -- `capturable=True` rejects the host-driven option `codebook_refresh_every > 0` at construction. +- `capturable=True` rejects host-driven periodic codebook refresh and gradient-histogram output. +- Every parameter in a captured optimizer must live on the current CUDA capture device; multi-process distributed training should construct one optimizer per rank/device. - A float `lr` bakes into the graph (as in `torch.optim`); pass a tensor `lr` and update it in place to drive an LR schedule. -- Checkpoints are portable across the `capturable` toggle in both directions. +- Checkpoints are portable across the `capturable` toggle in both directions; saving after graph replay synchronizes the true replayed global step so stochastic-rounding resumes continue from the correct seed. - FSDP2 / DTensor row-shard capture is validated by `tests/test_capturable_fsdp2.py` on 2 ranks for plain Gefen plus GefenMuon `sharded_mode="approx"`, `"exact"`, and `"distributed"` (including empty-shard coverage); ranks must capture and replay in lockstep. - Capturing the sharded `"exact"`/`"distributed"` modes records NCCL collectives (the in-`step()` all-gather) into the graph, which requires an NCCL/PyTorch build that supports collective graph capture — validated on PyTorch 2.12 / CUDA 13.3; on older stacks capture may raise (use `"approx"` or plain Gefen, which take no collectives, or `capturable=False`). diff --git a/README.md b/README.md index 0e90950..28b6b25 100644 --- a/README.md +++ b/README.md @@ -128,7 +128,7 @@ optimizer = Gefen(model.named_parameters(), lr=3e-4, fused=True, deterministic=T ## CUDA Graphs & torch.compile (`capturable`) -All three optimizers accept `capturable=True` (same meaning as `torch.optim`'s argument): `opt.step()` can then be captured in a `torch.cuda.CUDAGraph` or wrapped in `torch.compile(mode="reduce-overhead")` at no step-time cost — and the compiled hybrid step is about 10% faster than eager. Usage, caveats, and measured numbers: [COMPATIBILITY.md](https://github.com/thad0ctor/Gefen-X/blob/main/COMPATIBILITY.md#cuda-graphs--torchcompile-capturable). +All three optimizers accept `capturable=True` (same meaning as `torch.optim`'s argument): `opt.step()` can then be captured in a `torch.cuda.CUDAGraph` or wrapped in `torch.compile(mode="reduce-overhead")` at no step-time cost — and the compiled hybrid step is about 10% faster than eager. Device-resident global counters advance on every replay, so checkpoints record the true replayed step and stochastic-rounding resumes continue from the correct seed. Usage, caveats, and measured numbers: [COMPATIBILITY.md](https://github.com/thad0ctor/Gefen-X/blob/main/COMPATIBILITY.md#cuda-graphs--torchcompile-capturable).
diff --git a/src/gefen/gefen.py b/src/gefen/gefen.py index 5f53fde..c93402a 100644 --- a/src/gefen/gefen.py +++ b/src/gefen/gefen.py @@ -569,7 +569,9 @@ class Gefen(torch.optim.Optimizer): float ``lr`` is baked at capture time. Codebook learning and the period search are host-driven and run on the FIRST step, so do the standard warmup steps before capturing. Incompatible with - ``codebook_refresh_every > 0``. + ``codebook_refresh_every > 0``. Every captured parameter must live + on the current CUDA capture device; distributed jobs should use + one optimizer per rank/device. verbose (bool, default False): extra diagnostics during codebook learning and period prediction. """ @@ -704,6 +706,11 @@ def __init__( # nothing it would mark has changed. See _static_mark_signature. self._static_mark_sig = None self._gefen_global_step = 0 + # Manual CUDA-graph replay does not execute Python, so the host counter + # above cannot be the checkpoint authority under capturable=True. Keep + # one device counter per parameter device and advance it in the captured + # step tail; state_dict synchronizes the host mirror before serializing. + self._gefen_global_step_by_device = {} defaults = dict( lr=lr, @@ -926,6 +933,7 @@ def add_param_group(self, param_group): for param, param_name in zip(params, param_names): self._param_names[param] = param_name self.state[param]["name"] = param_name + self._ensure_gefen_global_step_devices() def _lr_scalar(self, group) -> float: """Resolve ``group["lr"]`` to a python float, caching a tensor lr's ``.item()``. @@ -1016,16 +1024,29 @@ def _assert_capturable_if_capturing(self) -> None: # graph without capturable=True would silently freeze the host-side # step counters / bias corrections at their capture-time values, so # fail loudly instead. - if ( - not self.capturable - and torch.cuda.is_available() - and torch.cuda.is_current_stream_capturing() - ): + capturing = ( + torch.cuda.is_available() and torch.cuda.is_current_stream_capturing() + ) + if capturing and not self.capturable: raise RuntimeError( "Attempting CUDA graph capture of {}.step() but capturable=False. " "Construct the optimizer with capturable=True to make step() " "graph-safe.".format(type(self).__name__) ) + if capturing: + devices = { + param.device + for group in self.param_groups + for param in group["params"] + } + capture_device = torch.device("cuda", torch.cuda.current_device()) + if devices != {capture_device}: + raise RuntimeError( + "CUDA graph capture requires every Gefen parameter on the " + "current capture device {}; found {}".format( + capture_device, sorted(map(str, devices)) + ) + ) def _static_mark_signature(self): # Cheap identity fingerprint of every object the static-marking pass @@ -1169,13 +1190,94 @@ def _sr_seed_on(self, device: torch.device): return None seed = self._sr_seed_by_device.get(device) if seed is None: + self._ensure_gefen_global_step_devices() + seed = self._sr_seed_by_device.get(device) + if seed is None: + global_counter = self._gefen_global_step_by_device.get(device) + initial_step = ( + self._gefen_global_step + if global_counter is None + else int(global_counter.item()) + ) seed = torch.full( - (), int(self._gefen_global_step), + (), initial_step, dtype=torch.int64, device=device, ) self._sr_seed_by_device[device] = seed return seed + def _gefen_cuda_param_devices(self): + return sorted( + { + param.device + for group in self.param_groups + for param in group["params"] + if param.device.type == "cuda" + }, + key=lambda device: -1 if device.index is None else device.index, + ) + + def _device_gefen_global_step(self): + if not self._gefen_global_step_by_device: + return None + steps = [ + int(counter.item()) + for counter in self._gefen_global_step_by_device.values() + ] + if any(step != steps[0] for step in steps[1:]): + raise RuntimeError( + "Gefen capturable global-step counters disagree across devices: " + "{}".format(steps) + ) + return steps[0] + + def _ensure_gefen_global_step_devices(self) -> None: + """Create capturable counters before capture, including no-grad devices.""" + if not self.capturable: + return + devices = self._gefen_cuda_param_devices() + missing = [ + device + for device in devices + if device not in self._gefen_global_step_by_device + ] + if missing: + device_step = self._device_gefen_global_step() + initial_step = self._gefen_global_step if device_step is None else device_step + self._gefen_global_step = initial_step + for device in missing: + self._gefen_global_step_by_device[device] = torch.full( + (), initial_step, dtype=torch.int64, device=device + ) + if self._stochastic_round: + for device in devices: + if device not in self._sr_seed_by_device: + self._sr_seed_by_device[device] = ( + self._gefen_global_step_by_device[device].detach().clone() + ) + + def _reset_gefen_global_step_devices(self) -> None: + self._gefen_global_step_by_device.clear() + self._ensure_gefen_global_step_devices() + + def _synchronize_gefen_global_step_for_checkpoint(self) -> None: + """Refresh the serialized host mirror after manual graph replays.""" + self._ensure_gefen_global_step_devices() + device_step = self._device_gefen_global_step() + if device_step is not None: + self._gefen_global_step = device_step + if self.capturable and self._stochastic_round: + seed_steps = [ + int(seed.item()) for seed in self._sr_seed_by_device.values() + ] + if any(seed_step != self._gefen_global_step for seed_step in seed_steps): + raise RuntimeError( + "Gefen capturable stochastic-round seeds disagree with the " + "optimizer global step: {} != {}".format( + seed_steps, self._gefen_global_step + ) + ) + @torch._dynamo.disable def _advance_sr_seeds(self) -> None: # Advance every per-device rounding seed by one, ON DEVICE, at the @@ -2350,6 +2452,11 @@ def _maybe_save_gefen_grad_histogram(self) -> None: requested_steps = quantization_module.LIST_STEPS_SAVE_HIST_GRAD if requested_steps is None: return + if self.capturable and requested_steps: + raise ValueError( + "capturable=True is incompatible with LIST_STEPS_SAVE_HIST_GRAD: " + "histogram collection and file output are host-driven" + ) if self._gefen_global_step not in requested_steps: return if not hasattr(quantization_module, "SAVE_CODEBOOK_PREFIX"): @@ -3258,6 +3365,8 @@ def state_dict(self): live ``self.state`` untouched (``_param_name`` still reads them); they are only withheld from the serialized copy. """ + self._synchronize_gefen_global_step_for_checkpoint() + # Withhold state entries orphaned from param_groups around the base # call (see docstring). Tensor hashing is identity-based, so set # membership matches the id()-keyed param_mappings the base class @@ -3590,11 +3699,9 @@ def load_state_dict(self, state_dict): self._gefen_global_step = gefen_global_step # Capturable SR seeds are optimizer-level scratch (a device mirror of # gefen_global_step): drop them so the first post-load SR kernel call - # rebuilds them from the restored counter. (Under manual graph-replay - # training the host counter -- like every host-side value -- reflects - # host step() calls, not replays; checkpoint semantics are unchanged - # from the pre-SR capturable behavior.) + # uses a scalar rebuilt from the restored, replay-synchronized counter. self._sr_seed_by_device.clear() + self._reset_gefen_global_step_devices() # Restoring the frozen codebook keeps _maybe_refresh_gefen_codebook a # no-op on the first resume step, so the restored automatic_period values @@ -3795,4 +3902,7 @@ def _advance_gefen_global_step(self) -> None: # step. @torch._dynamo.disable makes the increment a single opaque # call at the very tail of step() instead (one cheap graph break # after all device work). Eager behavior is unchanged. + self._ensure_gefen_global_step_devices() + for counter in self._gefen_global_step_by_device.values(): + counter.add_(1) self._gefen_global_step += 1 diff --git a/src/gefen/hybrid.py b/src/gefen/hybrid.py index ba315b9..60aa1b3 100644 --- a/src/gefen/hybrid.py +++ b/src/gefen/hybrid.py @@ -321,11 +321,13 @@ def __init__( capturable: forwarded to both halves. ``stochastic_round`` and ``deterministic`` are forwarded to Gefen children; with an AdamW backup they apply only to the Muon half. ``verbose`` is - likewise forwarded to Gefen children. + likewise forwarded to Gefen children. Every captured parameter + must live on the current CUDA capture device. """ if not isinstance(deterministic, bool): raise TypeError("deterministic must be a bool") self._deterministic = deterministic + self.capturable = capturable if backup_named_params is None: # Single-argument convenience form: the first arg is a model or a # named-param iterable to split internally. @@ -611,7 +613,35 @@ def zero_grad(self, set_to_none: bool = True): for o in self._subopts: o.zero_grad(set_to_none=set_to_none) + def _assert_capturable_devices_if_capturing(self) -> None: + capturing = ( + torch.cuda.is_available() and torch.cuda.is_current_stream_capturing() + ) + if not capturing: + return + if not self.capturable: + raise RuntimeError( + "Attempting CUDA graph capture of GefenMuonHybrid.step() but " + "capturable=False. Construct the optimizer with capturable=True " + "to make step() graph-safe." + ) + devices = { + param.device + for optimizer in self._subopts + for group in optimizer.param_groups + for param in group["params"] + } + capture_device = torch.device("cuda", torch.cuda.current_device()) + if devices != {capture_device}: + raise RuntimeError( + "CUDA graph capture requires every GefenMuonHybrid parameter on " + "the current capture device {}; found {}".format( + capture_device, sorted(map(str, devices)) + ) + ) + def step(self, closure=None): + self._assert_capturable_devices_if_capturing() # 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 diff --git a/tests/test_capturable.py b/tests/test_capturable.py index adf8be3..17ad0cb 100644 --- a/tests/test_capturable.py +++ b/tests/test_capturable.py @@ -32,6 +32,8 @@ Run on the current GPU: python tests/test_capturable.py """ +import copy + import torch try: @@ -550,11 +552,13 @@ def run(compiled): else: opt.step() torch.cuda.synchronize() - return [p.detach().clone() for _, p in params] + return [p.detach().clone() for _, p in params], opt - ref = run(compiled=False) - got = run(compiled=True) + ref, ref_opt = run(compiled=False) + got, got_opt = run(compiled=True) _assert_params_close(got, ref, rtol=1e-4, atol=1e-6, what="gefen-compile") + assert ref_opt.state_dict()["gefen_global_step"] == steps + assert got_opt.state_dict()["gefen_global_step"] == steps torch._dynamo.reset() @@ -708,6 +712,282 @@ def build(capturable): assert "step" in seen, seen +@pytest.mark.parametrize( + ("make_opt", "specs"), + [ + ( + lambda params, lr: Gefen( + params, lr=lr, capturable=True, fused=False + ), + GEFEN_SPECS, + ), + ( + lambda params, lr: GefenMuon( + params, lr=lr, capturable=True, fused=False, normuon=True + ), + MUON_SPECS, + ), + ], + ids=("gefen", "muon"), +) +def test_graph_replay_checkpoint_serializes_true_global_step(make_opt, specs): + warmup, replays = 2, 4 + steps = warmup + replays + grads = _grad_sequence(11, specs, steps) + params, opt = _run_captured( + make_opt, + specs, + grads, + [1e-3] * steps, + warmup, + seed=10, + ) + + state_dict = copy.deepcopy(opt.state_dict()) + + assert opt._gefen_global_step == steps + assert state_dict["gefen_global_step"] == steps + assert { + group["_gefen_checkpoint_metadata"]["global_step"] + for group in state_dict["param_groups"] + } == {steps} + assert { + int(counter.item()) + for counter in opt._gefen_global_step_by_device.values() + } == {steps} + + restored_params = [ + (name, torch.nn.Parameter(param.detach().clone())) for name, param in params + ] + restored = make_opt( + restored_params, + torch.tensor(1e-3, device=DEVICE), + ) + restored.load_state_dict(state_dict) + assert restored._gefen_global_step == steps + assert { + int(counter.item()) + for counter in restored._gefen_global_step_by_device.values() + } == {steps} + + +def test_graph_replay_global_step_advances_without_gradients(): + parameter = torch.nn.Parameter(torch.ones(32, device=DEVICE)) + opt = Gefen( + [("unused", parameter)], lr=1e-3, capturable=True, fused=False + ) + warmup, replays = 2, 3 + + for _ in range(warmup): + opt.step() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + opt.step() + for _ in range(replays): + graph.replay() + torch.cuda.synchronize() + + assert opt.state_dict()["gefen_global_step"] == warmup + replays + + +def test_no_grad_replays_seed_first_active_stochastic_step_correctly(): + warmup, replays = 2, 3 + steps = warmup + replays + params = _named_params(30, GEFEN_SPECS) + opt = _sr_gefen(params, True, lr=torch.tensor(1e-3, device=DEVICE)) + + for _ in range(warmup): + opt.step() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + opt.step() + for _ in range(replays): + graph.replay() + torch.cuda.synchronize() + + reference_params = _named_params(30, GEFEN_SPECS) + reference = _sr_gefen( + reference_params, True, lr=torch.tensor(1e-3, device=DEVICE) + ) + for _ in range(steps): + reference.step() + + active_grads = _grad_sequence(31, GEFEN_SPECS, 1)[0] + for (_, captured), (_, eager), grad in zip( + params, reference_params, active_grads + ): + captured.grad = grad.clone() + eager.grad = grad.clone() + opt.step() + reference.step() + torch.cuda.synchronize() + + for (_, captured), (_, eager) in zip(params, reference_params): + assert torch.equal(captured, eager) + for captured, eager in zip(_m_codebooks(opt), _m_codebooks(reference)): + assert torch.equal(captured, eager) + for candidate in (opt, reference): + assert candidate.state_dict()["gefen_global_step"] == steps + 1 + assert { + int(seed.item()) for seed in candidate._sr_seed_by_device.values() + } == {steps + 1} + + +@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="needs two CUDA devices") +def test_capturable_global_step_stays_aligned_across_devices(): + first = torch.nn.Parameter(torch.ones(32, device="cuda:0")) + second = torch.nn.Parameter(torch.ones(32, device="cuda:1")) + opt = Gefen( + [ + {"params": [("first", first)]}, + {"params": [("second", second)]}, + ], + lr=1e-3, + capturable=True, + fused=False, + ) + + for step in range(4): + first.grad = torch.full_like(first, 0.1) if step != 2 else None + second.grad = torch.full_like(second, -0.1) if step != 1 else None + opt.step() + + assert opt.state_dict()["gefen_global_step"] == 4 + assert { + int(counter.item()) + for counter in opt._gefen_global_step_by_device.values() + } == {4} + + +@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="needs two CUDA devices") +def test_add_param_group_uses_replay_authoritative_global_step(): + first = torch.nn.Parameter(torch.ones(32, device="cuda:0")) + opt = Gefen( + [("first", first)], lr=1e-3, capturable=True, fused=False + ) + for _ in range(2): + opt.step() + + # Model the state after five graph replays: only the captured device + # counter advances until Python runs again. + opt._gefen_global_step_by_device[first.device].add_(5) + second = torch.nn.Parameter(torch.ones(32, device="cuda:1")) + opt.add_param_group({"params": [("second", second)]}) + + assert opt._gefen_global_step == 7 + assert { + int(counter.item()) + for counter in opt._gefen_global_step_by_device.values() + } == {7} + + +@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="needs two CUDA devices") +def test_cuda_graph_rejects_multi_device_optimizer(): + first = torch.nn.Parameter(torch.ones(32, device="cuda:0")) + second = torch.nn.Parameter(torch.ones(32, device="cuda:1")) + opt = Gefen( + [ + {"params": [("first", first)]}, + {"params": [("second", second)]}, + ], + lr=1e-3, + capturable=True, + fused=False, + ) + + graph = torch.cuda.CUDAGraph() + with pytest.raises(RuntimeError, match="current capture device"): + with torch.cuda.graph(graph): + opt.step() + + +@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="needs two CUDA devices") +def test_cuda_graph_rejects_parameter_on_noncurrent_device(): + parameter = torch.nn.Parameter(torch.ones(32, device="cuda:1")) + opt = Gefen( + [("other", parameter)], lr=1e-3, capturable=True, fused=False + ) + + with torch.cuda.device(0): + graph = torch.cuda.CUDAGraph() + with pytest.raises(RuntimeError, match="current capture device cuda:0"): + with torch.cuda.graph(graph): + opt.step() + + +def test_cuda_graph_rejects_cpu_parameter(): + parameter = torch.nn.Parameter(torch.ones(32, device="cpu")) + opt = Gefen( + [("cpu", parameter)], lr=1e-3, capturable=True, fused=False + ) + + graph = torch.cuda.CUDAGraph() + with pytest.raises(RuntimeError, match="current capture device"): + with torch.cuda.graph(graph): + opt.step() + + +@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="needs two CUDA devices") +def test_hybrid_cuda_graph_rejects_children_split_across_devices(): + muon = [("hidden", torch.nn.Parameter(torch.ones(32, 32, device="cuda:0")))] + backup = [("bias", torch.nn.Parameter(torch.ones(32, device="cuda:1")))] + opt = GefenMuonHybrid( + muon, + backup, + lr=1e-3, + backup_optimizer="adamw", + fused=False, + capturable=True, + ) + + with torch.cuda.device(0): + graph = torch.cuda.CUDAGraph() + with pytest.raises(RuntimeError, match="every GefenMuonHybrid parameter"): + with torch.cuda.graph(graph): + opt.step() + + +def test_graph_replay_stochastic_round_checkpoint_continues_exactly(): + warmup, replays = 2, 4 + steps = warmup + replays + grads = _grad_sequence(21, GEFEN_SPECS, steps) + params, opt = _run_captured( + lambda named, lr: _sr_gefen(named, True, lr=lr), + GEFEN_SPECS, + grads, + [1e-3] * steps, + warmup, + seed=20, + ) + state_dict = copy.deepcopy(opt.state_dict()) + restored_params = [ + (name, torch.nn.Parameter(param.detach().clone())) for name, param in params + ] + restored = _sr_gefen( + restored_params, + True, + lr=torch.tensor(1e-3, device=DEVICE), + ) + restored.load_state_dict(state_dict) + + next_grads = _grad_sequence(22, GEFEN_SPECS, 1)[0] + for (_, original), (_, resumed), grad in zip( + params, restored_params, next_grads + ): + original.grad = grad.clone() + resumed.grad = grad.clone() + opt.step() + restored.step() + torch.cuda.synchronize() + + for (_, original), (_, resumed) in zip(params, restored_params): + assert torch.equal(original, resumed) + for original, resumed in zip(_m_codebooks(opt), _m_codebooks(restored)): + assert torch.equal(original, resumed) + assert opt.state_dict()["gefen_global_step"] == steps + 1 + assert restored.state_dict()["gefen_global_step"] == steps + 1 + + # --------------------------------------------------------------------------- # 5. stochastic rounding under capturable # --------------------------------------------------------------------------- @@ -926,6 +1206,7 @@ def run(compiled): for o in (ref_opt, got_opt): (seed,) = o._sr_seed_by_device.values() assert seed.item() == steps + assert o.state_dict()["gefen_global_step"] == steps torch._dynamo.reset() diff --git a/tests/test_validation_cpu.py b/tests/test_validation_cpu.py index 6930669..137c029 100644 --- a/tests/test_validation_cpu.py +++ b/tests/test_validation_cpu.py @@ -13,6 +13,7 @@ import torch.nn as nn import gefen +import gefen.quantization as gefen_quantization from gefen import Gefen, GefenMuon, GefenMuonHybrid, split_params_for_muon, validate_split from gefen.gefen_muon import ( NS_SCHEDULES, @@ -69,6 +70,17 @@ def test_gefen_rejects_capturable_with_codebook_refresh(): Gefen(_params_1d(), fused=False, capturable=True, codebook_refresh_every=10) +def test_gefen_rejects_capturable_gradient_histogram(monkeypatch): + monkeypatch.setattr( + gefen_quantization, "LIST_STEPS_SAVE_HIST_GRAD", [0], raising=False + ) + opt = Gefen(_params_1d(), fused=False, capturable=True) + opt.param_groups[0]["params"][0].grad = torch.ones(8) + + with pytest.raises(ValueError, match="LIST_STEPS_SAVE_HIST_GRAD"): + opt.step() + + def test_gefen_rejects_non_tensor_param(): with pytest.raises(TypeError): Gefen([("bad", "not a tensor")], fused=False) From 624608c0d8bce2158db54bd0327356aa9a0b0622 Mon Sep 17 00:00:00 2001 From: thad0ctor Date: Sun, 12 Jul 2026 05:17:13 -0700 Subject: [PATCH 04/14] Harden distributed optimizer integration and release gates --- .github/workflows/ci.yml | 42 +- .github/workflows/release.yml | 343 +++- .gitignore | 1 + CHANGELOG.md | 18 +- COMPATIBILITY.md | 31 +- CONTRIBUTING.md | 4 +- MANIFEST.in | 8 +- README.md | 47 +- benchmarks/README.md | 9 + benchmarks/trainer_resume/README.md | 45 + benchmarks/trainer_resume/__init__.py | 2 + benchmarks/trainer_resume/run.py | 646 ++++++++ pyproject.toml | 15 +- src/gefen/gefen.py | 1456 ++++++++++++++++- src/gefen/gefen_muon.py | 821 +++++++++- src/gefen/hybrid.py | 32 +- tests/test_amp_grad_scaler.py | 608 +++++++ tests/test_cpu_step_checkpoint.py | 67 + tests/test_gefen_fsdp2_checkpoint.py | 738 +++++++++ ...test_muon_distributed_checkpoint_safety.py | 839 ++++++++++ tests/test_muon_grad_presence.py | 554 +++++++ tests/test_step_preflight_atomicity.py | 246 +++ tests/test_training_matrix_harness.py | 10 +- tests/test_transformers_trainer_resume.py | 73 + 24 files changed, 6543 insertions(+), 112 deletions(-) create mode 100644 benchmarks/trainer_resume/README.md create mode 100644 benchmarks/trainer_resume/__init__.py create mode 100644 benchmarks/trainer_resume/run.py create mode 100644 tests/test_amp_grad_scaler.py create mode 100644 tests/test_gefen_fsdp2_checkpoint.py create mode 100644 tests/test_muon_distributed_checkpoint_safety.py create mode 100644 tests/test_muon_grad_presence.py create mode 100644 tests/test_step_preflight_atomicity.py create mode 100644 tests/test_transformers_trainer_resume.py diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index a7f78f6..38e1060 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -31,11 +31,11 @@ jobs: name: Lint (ruff) runs-on: ubuntu-latest steps: - - uses: actions/checkout@v4 + - uses: actions/checkout@34e114876b0b11c390a56381ad16ebd13914f8d5 # v4.3.1 with: persist-credentials: false - - uses: actions/setup-python@v5 + - uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0 with: python-version: "3.12" cache: pip @@ -52,11 +52,11 @@ jobs: name: Build sdist + wheel runs-on: ubuntu-latest steps: - - uses: actions/checkout@v4 + - uses: actions/checkout@34e114876b0b11c390a56381ad16ebd13914f8d5 # v4.3.1 with: persist-credentials: false - - uses: actions/setup-python@v5 + - uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0 with: python-version: "3.12" cache: pip @@ -72,27 +72,43 @@ jobs: run: python -m twine check dist/* cpu: - name: CPU tests (py${{ matrix.python-version }}) + name: CPU tests (py${{ matrix.python-version }}, ${{ matrix.torch-label }}) runs-on: ubuntu-latest strategy: fail-fast: false matrix: - python-version: ["3.10", "3.11", "3.12", "3.13"] + include: + - python-version: "3.10" + torch-spec: "torch==2.5.0" + torch-label: "torch floor 2.5.0" + - python-version: "3.11" + torch-spec: "torch" + torch-label: "latest torch" + - python-version: "3.12" + torch-spec: "torch" + torch-label: "latest torch" + - python-version: "3.13" + torch-spec: "torch" + torch-label: "latest torch" steps: - - uses: actions/checkout@v4 + - uses: actions/checkout@34e114876b0b11c390a56381ad16ebd13914f8d5 # v4.3.1 + with: + persist-credentials: false - - uses: actions/setup-python@v5 + - uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0 with: python-version: ${{ matrix.python-version }} cache: pip - name: Install CPU PyTorch + package + env: + TORCH_SPEC: ${{ matrix.torch-spec }} run: | python -m pip install --upgrade pip # CPU-only torch wheel: the CUDA wheel is ~2GB and pointless here. - pip install torch --index-url https://download.pytorch.org/whl/cpu + pip install "$TORCH_SPEC" --index-url https://download.pytorch.org/whl/cpu # torch is already satisfied, so this only pulls numba/numpy/ninja + gefen. - pip install .[test] + pip install '.[test]' - name: Byte-compile all sources run: python -m compileall -q src/gefen @@ -124,12 +140,14 @@ jobs: if: github.event_name == 'workflow_dispatch' && inputs.run_gpu_tests runs-on: [self-hosted, gpu] steps: - - uses: actions/checkout@v4 + - uses: actions/checkout@34e114876b0b11c390a56381ad16ebd13914f8d5 # v4.3.1 + with: + persist-credentials: false - name: Install package (CUDA torch expected on the runner) run: | python -m pip install --upgrade pip - pip install .[test] + pip install '.[test]' - name: Run full test suite (CUDA) # Multi-GPU / FSDP2 cases skip themselves when the runner has one GPU. diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index ad6170d..6661ef5 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -5,8 +5,10 @@ name: Release # git tag v0.2.1.dev1 && git push origin v0.2.1.dev1 -> build + TestPyPI only # git tag v0.2.1 && git push origin v0.2.1 -> build + TestPyPI + PyPI # -# The SAME artifacts built once in `build` are promoted through both indexes, so -# what lands on PyPI is byte-identical to what you smoke-tested on TestPyPI. +# `ci.yml` intentionally runs on branch pushes and pull requests, not tag pushes. +# This workflow therefore carries installed-wheel CPU and Transformers Trainer +# gates plus a mandatory two-GPU CUDA/JIT/distributed gate. Every gate tests the +# artifact built in `build`, and every publish job downloads that same artifact. # # Auth is OIDC trusted publishing (no API tokens stored). The manual approval # gates are GitHub Environment "required reviewers", configured in @@ -33,6 +35,7 @@ jobs: build: name: Build & verify artifacts runs-on: ubuntu-latest + timeout-minutes: 15 outputs: version: ${{ steps.ver.outputs.version }} prerelease: ${{ steps.ver.outputs.prerelease }} @@ -49,11 +52,36 @@ jobs: # could taint the wheel uploaded to PyPI. Build tooling installs fast. - name: Install build tooling - run: python -m pip install --upgrade pip build twine + run: | + python -m pip install 'pip==26.1.2' + python -m pip install 'build==1.5.1' 'twine==6.2.0' + + - name: Pin artifact timestamps to the tagged commit + run: echo "SOURCE_DATE_EPOCH=$(git log -1 --format=%ct)" >> "$GITHUB_ENV" - name: Build sdist + wheel # Pure-Python build (kernels JIT at runtime), so no CUDA toolchain needed. - run: python -m build + run: | + normalize_sdist() { + local SDIST SCRATCH TOPDIR + SDIST=$(find dist -maxdepth 1 -type f -name '*.tar.gz' -print -quit) + SCRATCH=$(mktemp -d) + tar -xzf "$SDIST" -C "$SCRATCH" + TOPDIR=$(find "$SCRATCH" -mindepth 1 -maxdepth 1 -type d -printf '%f\n') + test -n "$TOPDIR" + tar --sort=name --mtime="@${SOURCE_DATE_EPOCH}" --owner=0 --group=0 --numeric-owner --format=posix --pax-option=delete=atime,delete=ctime -C "$SCRATCH" -cf - "$TOPDIR" | gzip -n > "${SDIST}.normalized" + mv "${SDIST}.normalized" "$SDIST" + rm -rf "$SCRATCH" + } + build_once() { + rm -rf dist + python -m build + normalize_sdist + } + build_once + sha256sum dist/* | sort -k2 > /tmp/gefen-dist.sha256 + build_once + sha256sum -c /tmp/gefen-dist.sha256 - name: twine check run: python -m twine check dist/* @@ -72,22 +100,312 @@ jobs: exit 1 fi echo "version=$PKG_VER" >> "$GITHUB_OUTPUT" - # PEP 440 prerelease markers (.devN / aN / bN / rcN) => TestPyPI only. - if echo "$PKG_VER" | grep -Eq '(\.dev|a|b|rc)[0-9]+$'; then - echo "prerelease=true" >> "$GITHUB_OUTPUT" + # Any PEP 440 prerelease/dev release stops after TestPyPI. + PRERELEASE=$(python -c 'from packaging.version import Version; import sys; print(str(Version(sys.argv[1]).is_prerelease).lower())' "$PKG_VER") + echo "prerelease=$PRERELEASE" >> "$GITHUB_OUTPUT" + + - name: Upload artifacts + uses: actions/upload-artifact@ea165f8d65b6e75b540449e92b4886f43607fa02 # v4.6.2 + with: + name: dist + path: dist/ + + cpu_wheel_tests: + name: Installed-wheel CPU tests (${{ matrix.torch-label }}) + needs: build + runs-on: ubuntu-latest + timeout-minutes: 30 + strategy: + fail-fast: false + matrix: + include: + - python-version: "3.10" + torch-spec: "torch==2.5.0" + torch-label: "torch floor 2.5.0" + - python-version: "3.12" + torch-spec: "torch" + torch-label: "latest torch" + steps: + - uses: actions/checkout@34e114876b0b11c390a56381ad16ebd13914f8d5 # v4.3.1 + with: + persist-credentials: false + + - uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0 + with: + python-version: ${{ matrix.python-version }} + + - uses: actions/download-artifact@d3f86a106a0bac45b974a628896c90dbdf5c8093 # v4.3.0 + with: + name: dist + path: dist/ + + - name: Install the built wheel + env: + TORCH_SPEC: ${{ matrix.torch-spec }} + run: | + python -m pip install --upgrade pip + python -m pip install "$TORCH_SPEC" --index-url https://download.pytorch.org/whl/cpu + WHEEL=$(ls dist/*.whl) + python -m pip install "$WHEEL" pytest + + - name: Verify installed metadata and packaged JIT resources + env: + EXPECTED_VERSION: ${{ needs.build.outputs.version }} + run: | + python - <<'PY' + import importlib.metadata as metadata + import importlib.resources as resources + import os + from pathlib import Path + + import gefen + + version = metadata.version("gefen-x") + assert version == os.environ["EXPECTED_VERSION"], (version, os.environ["EXPECTED_VERSION"]) + module_path = Path(gefen.__file__).resolve() + checkout_src = (Path.cwd() / "src").resolve() + assert checkout_src not in module_path.parents, module_path + for symbol in ("Gefen", "GefenMuon", "GefenMuonHybrid"): + assert hasattr(gefen, symbol), symbol + + kernel_root = resources.files("gefen.kernels") + expected_sources = ( + "automatic_gefen_fused_binding.cpp", + "automatic_gefen_fused_kernel.cu", + "automatic_vmean_binding.cpp", + "automatic_vmean_kernel.cu", + "exact_histogram_fused_binding.cpp", + "exact_histogram_fused_kernel.cu", + "period_variance_binding.cpp", + "period_variance_kernel.cu", + ) + missing = [name for name in expected_sources if not kernel_root.joinpath(name).is_file()] + assert not missing, missing + assert resources.files("gefen").joinpath("py.typed").is_file() + requirements = metadata.requires("gefen-x") or [] + normalized_requirements = [requirement.partition(";")[0].strip().lower() for requirement in requirements] + assert any(requirement == "ninja" for requirement in normalized_requirements) + assert any(requirement.startswith("setuptools>=") for requirement in normalized_requirements) + print("wheel OK:", module_path, version, len(expected_sources), "JIT sources") + PY + ninja --version + + - name: Run CPU test suite against the wheel + env: + CUDA_VISIBLE_DEVICES: "" + run: python -m pytest tests -q -ra + + gpu_release_tests: + name: Two-GPU JIT & distributed release gate + needs: build + runs-on: [self-hosted, gpu] + timeout-minutes: 120 + steps: + - uses: actions/checkout@34e114876b0b11c390a56381ad16ebd13914f8d5 # v4.3.1 + with: + persist-credentials: false + + - uses: actions/download-artifact@d3f86a106a0bac45b974a628896c90dbdf5c8093 # v4.3.0 + with: + name: dist + path: dist/ + + - name: Preflight homogeneous two-GPU CUDA runner + run: | + python - <<'PY' + import re + import shutil + import subprocess + + import torch + + assert torch.version.cuda is not None, "release gate requires a CUDA-enabled PyTorch build" + assert torch.cuda.is_available(), "release gate requires CUDA" + assert torch.cuda.device_count() >= 2, "release gate requires at least two visible GPUs" + assert torch.distributed.is_available(), "torch.distributed is unavailable" + assert torch.distributed.is_nccl_available(), "release gate requires NCCL" + names = [torch.cuda.get_device_name(index) for index in range(2)] + capabilities = [torch.cuda.get_device_capability(index) for index in range(2)] + assert names[0] == names[1], ( + "replica-exact gate requires identical GPU models", + names, + ) + assert capabilities[0] == capabilities[1], ( + "replica-exact gate requires homogeneous GPUs", + capabilities, + ) + nvcc = shutil.which("nvcc") + assert nvcc is not None, "release gate requires nvcc on PATH" + nvcc_result = subprocess.run( + [nvcc, "--version"], check=True, capture_output=True, text=True + ) + match = re.search(r"release\s+([0-9]+)\.", nvcc_result.stdout) + assert match is not None, "could not parse nvcc CUDA version" + assert int(match.group(1)) == int(torch.version.cuda.split(".")[0]), ( + "nvcc and PyTorch CUDA major versions differ", + nvcc_result.stdout, + torch.version.cuda, + ) + print("torch:", torch.__version__, "CUDA:", torch.version.cuda) + print("GPUs:", names) + print("capabilities:", capabilities) + print(nvcc_result.stdout) + PY + if [ -n "${CUDA_VISIBLE_DEVICES:-}" ]; then + RELEASE_GPUS=$(printf '%s' "$CUDA_VISIBLE_DEVICES" | cut -d, -f1,2) else - echo "prerelease=false" >> "$GITHUB_OUTPUT" + RELEASE_GPUS=0,1 fi + echo "CUDA_VISIBLE_DEVICES=$RELEASE_GPUS" >> "$GITHUB_ENV" + echo "Mandatory release tests will use CUDA_VISIBLE_DEVICES=$RELEASE_GPUS" - - name: Upload artifacts + - name: Install the built wheel + env: + EXPECTED_VERSION: ${{ needs.build.outputs.version }} + run: | + python -m pip install --upgrade pip + python -m pip uninstall -y gefen-x || true + WHEEL=$(ls dist/*.whl) + python -m pip install "$WHEEL" pytest 'accelerate==1.14.0' + python - <<'PY' + import importlib.metadata as metadata + import os + from pathlib import Path + + import gefen + + version = metadata.version("gefen-x") + assert version == os.environ["EXPECTED_VERSION"], (version, os.environ["EXPECTED_VERSION"]) + module_path = Path(gefen.__file__).resolve() + assert (Path.cwd() / "src").resolve() not in module_path.parents, module_path + print("installed wheel:", module_path, version) + PY + ninja --version + + - name: Run mandatory JIT and distributed tests + env: + GEFEN_KERNEL_BUILD_ROOT: ${{ runner.temp }}/gefen-release-jit-${{ github.run_id }}-${{ github.run_attempt }} + GEFEN_VERBOSE_BUILD: "1" + run: | + rm -rf "$GEFEN_KERNEL_BUILD_ROOT" + python -m pytest -q -ra --junitxml=release-gpu.xml \ + tests/test_amp_grad_scaler.py \ + tests/test_capturable.py \ + tests/test_capturable_fsdp2.py \ + tests/test_deterministic_mode.py \ + tests/test_factored_v_ema_parity.py \ + tests/test_fused_fsdp2_noncontig.py \ + tests/test_fused_full_update_parity.py \ + tests/test_fused_update_v2_full_parity.py \ + tests/test_gefen_fsdp2_checkpoint.py \ + tests/test_muon_distributed_checkpoint_safety.py \ + tests/test_muon_grad_presence.py \ + tests/test_step_preflight_atomicity.py \ + tests/test_muon_fsdp2_parity.py::test_muon_fsdp2_fused_cuda_parity \ + tests/test_muon_fsdp2_parity.py::test_muon_fsdp2_fused_multirank_parity \ + tests/test_muon_fsdp2_parity.py::test_muon_fsdp2_distributed_parity \ + tests/test_muon_fsdp2_parity.py::test_muon_fsdp2_distributed_fluctuating_grad_set \ + tests/test_muon_fsdp2_parity.py::test_muon_fsdp2_distributed_checkpoint_restore + + - name: Reject skipped release-gate tests + run: | + python - <<'PY' + import xml.etree.ElementTree as ET + + root = ET.parse("release-gpu.xml").getroot() + cases = root.findall(".//testcase") + assert cases, "release gate collected no tests" + skipped = [ + "{}::{}".format(case.attrib.get("classname", ""), case.attrib.get("name", "")) + for case in cases + if case.find("skipped") is not None + ] + assert not skipped, "mandatory release tests skipped:\n{}".format("\n".join(skipped)) + print("mandatory GPU release tests:", len(cases), "passed with zero skips") + PY + + - name: Upload GPU release-gate report + if: always() uses: actions/upload-artifact@ea165f8d65b6e75b540449e92b4886f43607fa02 # v4.6.2 + with: + name: release-gpu-junit-${{ github.run_attempt }} + path: release-gpu.xml + if-no-files-found: ignore + + framework_wheel_tests: + name: Installed-wheel Transformers Trainer resume + needs: build + runs-on: ubuntu-latest + timeout-minutes: 30 + steps: + - uses: actions/checkout@34e114876b0b11c390a56381ad16ebd13914f8d5 # v4.3.1 + with: + persist-credentials: false + + - uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0 + with: + python-version: "3.12" + + - uses: actions/download-artifact@d3f86a106a0bac45b974a628896c90dbdf5c8093 # v4.3.0 with: name: dist path: dist/ + - name: Install wheel and pinned framework surface + env: + EXPECTED_VERSION: ${{ needs.build.outputs.version }} + run: | + python -m pip install --upgrade pip + python -m pip install torch --index-url https://download.pytorch.org/whl/cpu + WHEEL=$(ls dist/*.whl) + python -m pip install "$WHEEL" pytest 'transformers==5.5.0' 'accelerate==1.14.0' + python - <<'PY' + import importlib.metadata as metadata + import os + from pathlib import Path + + import gefen + + version = metadata.version("gefen-x") + assert version == os.environ["EXPECTED_VERSION"], (version, os.environ["EXPECTED_VERSION"]) + module_path = Path(gefen.__file__).resolve() + assert (Path.cwd() / "src").resolve() not in module_path.parents, module_path + print("installed wheel:", module_path, version) + PY + + - name: Run Trainer checkpoint-continuation matrix + env: + CUDA_VISIBLE_DEVICES: "" + run: python -m pytest -q -ra --junitxml=release-framework.xml tests/test_transformers_trainer_resume.py + + - name: Reject skipped framework-gate tests + run: | + python - <<'PY' + import xml.etree.ElementTree as ET + + root = ET.parse("release-framework.xml").getroot() + cases = root.findall(".//testcase") + assert cases, "framework gate collected no tests" + skipped = [ + "{}::{}".format(case.attrib.get("classname", ""), case.attrib.get("name", "")) + for case in cases + if case.find("skipped") is not None + ] + assert not skipped, "mandatory framework tests skipped:\n{}".format("\n".join(skipped)) + print("mandatory framework tests:", len(cases), "passed with zero skips") + PY + + - name: Upload framework release-gate report + if: always() + uses: actions/upload-artifact@ea165f8d65b6e75b540449e92b4886f43607fa02 # v4.6.2 + with: + name: release-framework-junit-${{ github.run_attempt }} + path: release-framework.xml + if-no-files-found: ignore + testpypi: name: Publish to TestPyPI (gate 1) - needs: build + needs: [build, cpu_wheel_tests, framework_wheel_tests, gpu_release_tests] runs-on: ubuntu-latest # Approval gate #1: the `testpypi` environment's required reviewers. environment: @@ -104,13 +422,10 @@ jobs: uses: pypa/gh-action-pypi-publish@cef221092ed1bacb1cc03d23a2d87d1d172e277b # v1.14.0 with: repository-url: https://test.pypi.org/legacy/ - # A real release also passes through here; if that version was already - # tested on TestPyPI, don't hard-fail on the duplicate. - skip-existing: true pypi: name: Publish to PyPI (gate 2) - needs: [build, testpypi] + needs: [build, cpu_wheel_tests, framework_wheel_tests, gpu_release_tests, testpypi] # Prerelease tags stop at TestPyPI; only clean vX.Y.Z tags reach PyPI. if: needs.build.outputs.prerelease == 'false' runs-on: ubuntu-latest diff --git a/.gitignore b/.gitignore index 61cc2a3..02c7457 100644 --- a/.gitignore +++ b/.gitignore @@ -18,6 +18,7 @@ benchmarks/data/ # Generated optimizer-matrix results, logs, and checkpoints benchmarks/training_matrix/out/ benchmarks/training_matrix/RESULTS_*.md +benchmarks/trainer_resume/out/ # Toy fine-tuning example outputs (generated data + saved checkpoints) examples/toy-finetune/data/ diff --git a/CHANGELOG.md b/CHANGELOG.md index 69014ea..c6514c8 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,15 +2,27 @@ All notable changes to this project are documented here. This project adheres to [Semantic Versioning](https://semver.org/). -## [Unreleased] +## [0.4.0] - 2026-07-12 Correctness and compatibility: - Add `deterministic=True` to `Gefen`, `GefenMuon`, and `GefenMuonHybrid` for replica-exact fused routing on homogeneous GPUs. Automatic periods use fixed-order reductions, block-vmean parameters use the deterministic fused v1 path, factored-v parameters use the decomposed deterministic update, and tagged checkpoints enforce the saved policy. - Capturable optimizers maintain device-resident global-step counters on every parameter device. CUDA-graph replays now serialize the true global step, including steps with no gradients, so stochastic-rounding checkpoints resume with the correct seed. -- Checkpoint loading preserves compact optimizer-state dtypes for bf16 parameters, validates frozen codebooks and hybrid backend metadata, and keeps legacy untagged checkpoints loadable. +- Checkpoint loading preserves compact optimizer-state dtypes for bf16 parameters, validates frozen codebooks and hybrid backend metadata, and keeps safe legacy untagged native checkpoints loadable. +- Plain Gefen and `GefenMuon(sharded_mode="approx")` collectively encode rank-local DTensor optimizer state for ordinary and flattened PyTorch full-state DCP. Two-GPU FSDP2 `fully_shard` get/set tests resume both optimizers' next update exactly on the same one-dimensional default-world mesh; all ranks must participate, per-process CPU checkpoint memory approaches `world_size ×` local serialized state plus local scratch, topology changes fail closed, and old unsafe untagged full checkpoints are rejected. +- Add a Transformers Trainer checkpoint-continuation gate for plain Gefen, GefenMuon+AdamW, and GefenMuon+Gefen, covering Trainer's internal Accelerate wrapper, tied weights, gradient accumulation, changing LRs, fused BF16 DDP replica hashes, and exact deterministic resume state. +- Integrate PyTorch GradScaler without changing the ordinary FP32-master/BF16 Accelerate path: true-FP16 gradients use the optimizer-side protocol, overflow skips every GefenMuonHybrid child atomically, DTensor overflow decisions remain synchronized across ranks, and FSDP1 FP16 directs callers to `ShardedGradScaler`. +- Reject rank-divergent `grad is None` patterns before exact/distributed DTensor Muon collectives, avoiding an unmatched-collective deadlock without adding collectives to plain/DDP/`approx` steps. +- Parallel-Muon `distributed` checkpoints now carry a versioned saved-world ownership manifest, validate full owner-state geometry before mutation, preserve released consolidated-v1 and cross-world-size resume, and reject markerless, partial, or internally inconsistent populated state instead of warning and risking divergent momentum. +- Preflight every active gradient before AMP, codebook work, or child dispatch so a later sparse, compressed, MKLDNN, complex, or malformed gradient cannot leave an earlier parameter partially updated; periodic codebook refresh likewise stages every replacement index before committing shared state. - Reject host-driven gradient-histogram output under `capturable=True`, matching the existing periodic-codebook-refresh guard. +Packaging: + +- Require `numba>=0.65` for the compiled exact-DP codebook solver. +- Make `ninja` and `setuptools>=77` core dependencies because PyTorch's runtime CUDA extension loader requires both to build the fused kernels; the former `perf` extra is no longer needed. +- Pin release build tooling, verify byte-reproducible wheel and normalized-sdist rebuilds, and gate the installed artifact on PyTorch 2.5.0 plus latest CPU, Transformers Trainer resume, and fresh-build two-GPU CUDA/distributed tests before either package index can publish it. + ## [0.3.0] - 2026-07-11 Lands the Muon optimization suite (#62) and validated DeepSpeed ZeRO support (#64): a selectable AdamW backup for the hybrid, an opt-in batched Newton-Schulz experiment, faster fused Muon momentum kernels, and plain `Gefen` as a validated DeepSpeed ZeRO 1-3 client optimizer with bit-exact checkpoint resume. @@ -75,7 +87,7 @@ Packaging: - Distribution renamed to `gefen-x`; project URLs point at the fork repository. - Version single-sourced at runtime via `gefen.__version__`. -- Hard runtime dependencies are `torch>=2.5`, `numpy`, and `numba` (the compiled codebook solver; the pure-Python fallback is too slow for real training). `ninja` is optional under the `perf` extra (`pip install gefen-x[perf]`) for faster CUDA JIT builds. `requires-python = ">=3.10"`. +- Hard runtime dependencies are `torch>=2.5`, `numpy`, and `numba` (the compiled codebook solver; the pure-Python fallback is too slow for real training). At this release, `ninja` was exposed under the `perf` extra; 0.4.0 corrects it to a core dependency because PyTorch's runtime extension loader requires it. `requires-python = ">=3.10"`. - SPDX `license = "MIT"` metadata and a shipped `py.typed` (PEP 561) marker. - CI adds lint, an sdist/wheel build check, and Python 3.12 and 3.13 to the test matrix; the CPU job now runs the real pytest suite (66 new CPU-only tests) instead of an import smoke. - Package moved into a `src/gefen/` layout, retiring the flat-layout import shadowing (`import gefen` from a clone; the repo's `kernels/` no longer shadows Hugging Face's `kernels`). diff --git a/COMPATIBILITY.md b/COMPATIBILITY.md index eecb7b0..a831707 100644 --- a/COMPATIBILITY.md +++ b/COMPATIBILITY.md @@ -82,12 +82,35 @@ Setup: from scratch — MNIST (paper recipe, 3 seeds), CIFAR-10 ResNet-18; fine- - **Real-dataset rows** (the Vision and Audio tables) report held-out accuracy, COCO mAP, or — where no accuracy metric exists (TTS) — validation loss, after a full training run: a convergence comparison against AdamW at matched model / data / LR / schedule / epochs. - **Smoke-test rows** (the LLM / VLM / diffusion matrix above) train a fixed batch for N optimizer steps and report the loss as `first → last (Δ%)`, where Δ% is the reduction from the first step's loss to the last (`(first − last) / first`); PASS if it exceeds 30% with no NaN/Inf. Peak VRAM from `torch.cuda.max_memory_allocated`. -- The **Method** column: `full-param` trains every native parameter tensor; `full-param FSDP2` / `device_map` shard that same full-parameter training for models over one card; `LoRA` covers models too large to full-fine-tune even sharded — a weaker claim, since only the adapter matrices are trained. +- The **Method** column: `full-param` trains every native parameter tensor; `full-param FSDP2` / `device_map` shard that same full-parameter training for models over one card; `LoRA` covers models too large to full-fine-tune even sharded — a weaker claim, since only the adapter matrices are trained. The FSDP2 rows validate training and loss reduction; the narrower same-topology full-state optimizer checkpoint contract is described below. - **Versions**: LLM / VLM / diffusion smokes on torch 2.12.0+cu133, transformers 5.10.2 (≥ 5.11 for PaddleOCR-VL), diffusers 0.39.0. Vision / audio / real-data runs on torch 2.13.0+cu133, torchvision 0.28, torchaudio 2.11, transformers 5.12, ultralytics 8.4, rfdetr 1.8. Harness and per-model recipes: [`benchmarks/arch-compat/`](benchmarks/arch-compat/). +## Megatron-LM integration scope + +Megatron-LM integration tests ran tiny mock-data GPT training through the production pretraining entry point for plain Gefen, GefenMuon+AdamW, and GefenMuon+Gefen. Fused deterministic coverage on two homogeneous RTX 3090 Ti GPUs included DP2, TP2, PP2, CP2 with Transformer Engine, and EP2, with replica or tied-weight hashes appropriate to each topology; TP2 and PP2 also covered legacy optimizer checkpoint continuation for all three recipes. Unfused four-rank coverage included EP2×DP2, TP2×DP2, PP2×DP2, and TP2×PP2 for all three recipes, plus EP2×ETP2 and EP2 checkpoint continuation for GefenMuon+Gefen. This scope does not include Megatron's distributed optimizer, FSDP, optimizer CPU offload, fp16, or non-legacy optimizer checkpoint formats. + +## Optimizer checkpoint scope + +Native single-process optimizer `state_dict()`/`load_state_dict()` and the explicitly documented GefenMuon `distributed` state path remain separate from the rank-local DCP adapter. For plain Gefen and `GefenMuon(sharded_mode="approx")`, each rank can learn different local codebook and block geometry, so `state_dict()` collectively replaces ordinary rank-local tensors with one tagged, rank-indexed CPU payload that PyTorch full-state DCP preserves. Both optimizers' two-GPU FSDP2 `fully_shard` paths are exercised through `get_optimizer_state_dict(..., full_state_dict=True, cpu_offload=True)` and `set_optimizer_state_dict(..., full_state_dict=True, broadcast_from_rank0=True)`, with exact next-step continuation; the flattened optimizer-state form is also covered. + +This format is deliberately same-topology only and currently requires one 1-D DeviceMesh spanning the default process-group world. Every rank must participate in both save and restore; each process temporarily holds all serialized rank payloads on CPU, so the leading checkpoint-time CPU cost approaches `world_size ×` its local optimizer-state size plus local serialization scratch. Loading validates world size, parameter order and names, global and local shapes and dtypes, mesh membership and names, structural placements, rank coordinates, global step, deterministic policy, frozen codebook, and sharded mode before mutation. Multidimensional meshes, subgroups, pipeline-local optimizers, world-size/topology changes, and old unsafe untagged full checkpoints fail closed rather than silently applying rank 0's state to every shard. No optimizer-state reshard portability is claimed, model-only DCP is unaffected, and the full-state DCP support described here does not extend beyond plain Gefen and Muon `approx`. + +## Transformers Trainer DDP + +The `benchmarks.trainer_resume` gate exercises plain Gefen, GefenMuon+AdamW, and GefenMuon+Gefen through Trainer's internal Accelerate wrapper with tied weights, gradient accumulation, a changing scheduler, native Trainer checkpoint files, BF16 fused updates, and two-rank DDP replica hashes. All three recipes have passed its deterministic fused-BF16 two-rank configuration on homogeneous GPUs, which requires exact model, optimizer, scheduler, LR, and logged-loss agreement between uninterrupted and resumed runs. Run the gate with: + +```bash +CUDA_VISIBLE_DEVICES=0,1 PYTHONPATH=.:src torchrun --standalone --nproc-per-node=2 \ + -m benchmarks.trainer_resume.run --output-dir benchmarks/trainer_resume/out/ddp \ + --device cuda --dtype bfloat16 --fused --deterministic --steps 3 --split-step 1 \ + --gradient-accumulation-steps 2 +``` + +This covers Trainer's DDP frontend and internal Accelerate optimizer wrapper, not a standalone `Accelerator` loop, FSDP, or `model_init`. + ## CUDA Graphs & torch.compile (`capturable`) -`Gefen`, `GefenMuon`, and `GefenMuonHybrid` accept `capturable=True` (same meaning as `torch.optim`'s argument): step counters, bias corrections, the optimizer-global checkpoint counter, and a tensor `lr` live on the GPU, so `opt.step()` stays correct inside a replayed CUDA graph or a compiled region. The default `capturable=False` is bit-identical to previous behavior, and capturing with it raises instead of silently freezing the step counters. +`Gefen`, `GefenMuon`, and `GefenMuonHybrid` accept `capturable=True` (same meaning as `torch.optim`'s argument): step counters, bias corrections, the optimizer-global checkpoint counter, and a tensor `lr` live on the GPU, so `opt.step()` stays correct inside a replayed CUDA graph or compiled regions. The default `capturable=False` preserves the established eager behavior, and capturing with it raises instead of silently freezing the step counters. Parity is bitwise for fixed-order components and tolerance-bounded where fused atomic reductions can vary at the ULP scale. Measured step times (386M-parameter census, RTX 3090 Ti, tail-100 mean over 500 CUDA-event-timed steps): @@ -98,6 +121,8 @@ Measured step times (386M-parameter census, RTX 3090 Ti, tail-100 mean over 500 | manual `torch.cuda.CUDAGraph` replay | 22.9 ms | 168.9 ms | | `torch.compile(mode="reduce-overhead")` | 23.0 ms | **147.6 ms** | +These measurements are path- and shape-specific, not a zero-overhead guarantee: `capturable=True` and manual replay are nearly flat for plain Gefen but add modest overhead to the measured hybrid, while compilation improves the measured hybrid and is slightly slower for plain Gefen. + **Manual capture** — run a few real warmup steps first (codebook learning happens on the first step), then capture one `opt.step()` and drive training by refilling the static grad buffers (and the `lr` tensor, for schedules) before each `graph.replay()`: ```python @@ -106,7 +131,7 @@ from gefen import GefenMuonHybrid opt = GefenMuonHybrid(model, lr=torch.tensor(3e-5, device="cuda"), capturable=True) ``` -**torch.compile** — the fused kernels are registered as torch custom ops, so dynamo traces the whole step into one graph and CUDA-graphs it end to end: +**torch.compile** — the fused kernels are registered as torch custom ops, so they remain inside compiled regions. The step also contains deliberate `@torch._dynamo.disable` host helpers for static-address registration, batched scalar refresh, and global-step bookkeeping, so Dynamo may partition one optimizer call into multiple compiled regions rather than one end-to-end graph: ```python compiled_step = torch.compile(opt.step, mode="reduce-overhead", dynamic=False) diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 8032400..c526ec3 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -12,7 +12,7 @@ pip install torch --index-url https://download.pytorch.org/whl/cpu pip install -e ".[test]" ``` -Add the `perf` extra (`pip install -e ".[perf,test]"`) for ninja-backed CUDA kernel builds (faster JIT; optional — torch builds without it). +`ninja` and setuptools are core dependencies because PyTorch's runtime extension loader requires them for the fused CUDA kernel build; the editable install above includes both automatically. ## Running the tests @@ -22,7 +22,7 @@ CI runs the CPU suite from the repo root after installing the package. Mirror th python -m pytest tests -q -ra ``` -CUDA-dependent tests skip themselves automatically on CPU. GPU kernel-parity tests require an NVIDIA device plus `nvcc` and are gated behind the manual `workflow_dispatch` GPU job in CI; run them locally with the same command on a CUDA host. +CUDA-dependent tests skip themselves automatically on CPU. GPU kernel-parity tests require an NVIDIA device plus `nvcc`; branch CI exposes them through the manual `workflow_dispatch` job, while release tags must pass the two-GPU JIT and distributed gate in `release.yml`. Run the full suite locally with the same command on a CUDA host. ## Code style diff --git a/MANIFEST.in b/MANIFEST.in index 5546082..efe54dd 100644 --- a/MANIFEST.in +++ b/MANIFEST.in @@ -1,6 +1,2 @@ -# Include CUDA/C++ sources in source distributions so PyTorch JIT can build kernels after install. -recursive-include src/gefen/kernels *.py *.cu *.cpp -recursive-include src/gefen/tools *.md -# PEP 561 typing marker. -include src/gefen/py.typed -prune src/gefen/kernels/.build +prune benchmarks +prune tests diff --git a/README.md b/README.md index 28b6b25..6148330 100644 --- a/README.md +++ b/README.md @@ -19,10 +19,10 @@ - **New: matches AdamW's loss out of the box** at about a quarter of its optimizer memory (default `factored_v_2d`). See [Benchmarks](#benchmarks) and [the factored-v lever](#quality-lever-factored-second-moment-on-2d-params-factored_v_2d). - **Works on modern decoders** (Qwen3, Llama-3, Mistral). Upstream loses its memory advantage on these architectures (~9 B/param — worse than AdamW); this fork keeps the intended ~1 B/param. - **Validated across 2026 architectures** — two dozen modern LLMs, VLMs, and image/video/audio media-gen models full-fine-tuned on 24 GB GPUs, plus 20-30B MoEs via LoRA. See the [compatibility matrix](https://github.com/thad0ctor/Gefen-X/blob/main/COMPATIBILITY.md). - - **~2× faster `opt.step()`** via fused CUDA kernels, with identical results. + - **~2× faster `opt.step()`** via fused CUDA kernels, with validated numerical parity; fixed-order paths have bitwise checks, while atomic reduction paths are tolerance-bounded. - **Whole-model Muon option** (`GefenMuonHybrid`) with a selectable AdamW quality backup or ~1 B/param Gefen low-memory backup, plus task-specific [SFT and pretraining recipes](#which-muon-recipe-should-i-use). - - **Reliable checkpoint save/resume and FSDP2 support** — broken or absent in the shipped release. - - **Hardened against crashes** (device/dtype guards, bounds checks, race fixes) with a bit-exact test suite. + - **Reliable native checkpoint save/resume and FSDP2 training support** — broken or absent in the shipped release. Plain Gefen and rank-local `GefenMuon(sharded_mode="approx")` also support same-topology PyTorch full-state DCP optimizer resume; see [Distributed Training](#distributed-training) for its collective and portability limits. + - **Hardened against crashes** (device/dtype guards, bounds checks, race fixes) with bitwise tests where reduction order is fixed and tolerance-bounded tests for CUDA atomic reductions.
Detailed Fork Improvements (vs upstream) @@ -33,19 +33,19 @@ >| **Loss vs AdamW** | trails AdamW by ~0.06 | **matches AdamW** via the default `factored_v_2d` — [details](#quality-lever-factored-second-moment-on-2d-params-factored_v_2d) | >| **Modern decoders** (Qwen3 / Llama-3 / Mistral — SwiGLU + grouped-query attention) | uses *more* optimizer memory than AdamW on these | keeps the full ~1 B/param optimizer state (about a quarter of bf16 AdamW's) | >| **Learning rate(s)** | no guidance — silently over-steps | documented ~0.6× AdamW, so quality matches AdamW | ->| **Optimizer-step speed** | baseline | ~2× faster `opt.step()` (fused kernels), identical results | +>| **Optimizer-step speed** | baseline | ~2× faster `opt.step()` (fused kernels), with fixed-order bitwise checks and bounded atomic-reduction differences | >| **Peak memory**| large transient spikes | much lower peak — room for bigger models / batches | >| **Sharded multi-GPU training (FSDP2)** | breaks with the fast path | works — for plain Gefen *and* Muon | >| **Whole-model Muon** | 2D weight matrices only | `GefenMuonHybrid` trains the entire model | >| **Muon step efficiency** | generic momentum hack + redundant dequant gather | single-pass bit-exact momentum kernel | ->| **Save / resume checkpoints** | can corrupt state or lose tuning on resume | saves & resumes correctly | +>| **Save / resume checkpoints** | can corrupt state or lose tuning on resume | native optimizer checkpoints save and resume correctly; distributed checkpoint formats have documented limits | >| **Crash safety** | missing device / edge-case guards | guarded against wrong-device, empty-tensor, and race bugs | ->| **Correctness** | no fused-kernel tests | bit-exact + distributed parity test suite | +>| **Correctness** | no fused-kernel tests | bitwise kernel checks, tolerance-bounded atomic-reduction parity, and distributed tests | >| **Documentation** | no Axolotl / fork-install guidance | Axolotl how-to + fair loss/speed/memory benchmarks | >| **Muon Usage** | no whole-model recipe or task split | measured SFT/pretraining recipes pair quantized Muon with an AdamW backup and document the observed quality/throughput tradeoff; a Gefen backup remains available for minimum state — [details](#which-muon-recipe-should-i-use) | >| **Muon (Newton-Schulz) speed** | fixed 5-step schedule | tunable `ns_schedule`; tuned3 is the balanced SFT choice, while quality-first pretraining retains classic NS5 — [details](#experimental-lever-faster-newton-schulz-ns_schedule-fp8_ns) | >| **Low-precision orthogonalization** | bf16 only | opt-in `fp8_ns` for large matrices on newer GPUs, safe fallback elsewhere — [details](#experimental-lever-faster-newton-schulz-ns_schedule-fp8_ns) | ->| **Sharded Muon under FSDP2** | every GPU redundantly repeats the same work | `sharded_mode="distributed"` splits the work across GPUs at identical results — [details](#experimental-lever-sharded-newton-schulz-under-fsdp2-sharded_mode) | +>| **Sharded Muon under FSDP2** | every GPU redundantly repeats the same work | `sharded_mode="distributed"` splits the work across GPUs and matches `"exact"` bitwise on homogeneous GPUs in the parity suite — [details](#experimental-lever-sharded-newton-schulz-under-fsdp2-sharded_mode) |
@@ -81,7 +81,7 @@ CUDA toolkit and host compiler compatible with your PyTorch build (the JIT compi | PyTorch | 2.5+ (verified on 2.12 / cu133) | | Platform | Linux with a CUDA GPU | -The Hugging Face Trainer flow, the `find_lr` tool, and the `examples/` all use `transformers` and `datasets`, which are not core dependencies — install them alongside (`pip install transformers datasets`). Optional benchmark baselines: `bitsandbytes` (AdamW-8bit), `torchao` (AdamW-4bit). +The Hugging Face Trainer flow uses `transformers` and Accelerate, while the data-backed examples also use `datasets`; these are not core dependencies — install them alongside (`pip install transformers 'accelerate>=1.1' datasets`). Optional benchmark baselines: `bitsandbytes` (AdamW-8bit), `torchao` (AdamW-4bit). ## Compatibility @@ -114,7 +114,11 @@ One knock-on effect: weight decay in AdamW-style optimizers is applied as `lr × ## Distributed Training -Gefen drops into standard distributed training like any other PyTorch optimizer, with either `fused=True` or `fused=False`. Validated setups: single-GPU, PyTorch DDP, FSDP2 (`fully_shard` / DTensor), and DeepSpeed ZeRO 1-3 (plain `Gefen` as the client optimizer, direct or via axolotl `gefenx`; bit-exact ZeRO-2 checkpoint resume). +Gefen drops into standard distributed training like any other PyTorch optimizer, with either `fused=True` or `fused=False`. Validated training setups include single-GPU, PyTorch DDP, FSDP2 (`fully_shard` / DTensor), and DeepSpeed ZeRO 1-3 (plain `Gefen` as the client optimizer, direct or via axolotl `gefenx`; bit-exact ZeRO-2 checkpoint resume). FSDP2 optimizer checkpoint support is mode- and topology-specific as described below. + +> **FSDP2 optimizer checkpoint scope.** Plain Gefen and `GefenMuon(sharded_mode="approx")` collectively encode every rank's local DTensor optimizer state into PyTorch DCP `StateDictOptions(full_state_dict=True)` output, including the flattened optimizer-state form. `get_optimizer_state_dict()` and `set_optimizer_state_dict()` resume the next update exactly when world size, mesh, placements, rank coordinates, parameter ordering, shapes, names, and sharded mode are unchanged; the actual two-GPU `fully_shard` get/set test covers both optimizers. The adapter currently requires one 1-D DeviceMesh spanning the default world; multidimensional meshes, subgroups, and pipeline-local optimizers fail before its collectives. Save and restore are collective, so every rank must participate. Each process temporarily stages all serialized rank payloads on CPU, making the leading checkpoint-time CPU cost about `world_size ×` that rank's local optimizer state plus local serialization scratch. World-size or topology changes fail before mutation, and older unsafe untagged full checkpoints fail closed. This is not a reshardable optimizer-state format. + +`torch.amp.GradScaler` keeps PyTorch's ordinary externally skipped step for FP32-master and BF16 training, including Trainer/Accelerate gradient clipping and scheduler behavior. Actual FP16 gradient storage opts into PyTorch's native optimizer-side scaling protocol so finite gradients are unscaled once and overflow returns before either Hybrid child, codebook, state, parameter, or counter changes. DTensor/FSDP2 non-finite flags are reduced across the mesh; FSDP1 FlatParameters must use `torch.distributed.fsdp.ShardedGradScaler`. > **DeepSpeed ZeRO config.** Set `"zero_allow_untested_optimizer": true` and leave the config's `optimizer` section unset. With optimizer CPU-offload, also set `"zero_force_ds_cpu_optimizer": false` — otherwise raw DeepSpeed refuses to initialize, and accelerate-based launchers (axolotl) silently swap in DeepSpeed's own CPU Adam. ZeRO steps flattened 1-D partitions, so `GefenMuon`/`GefenMuonHybrid` raise a clear error under ZeRO; use FSDP2, DDP, or single-GPU for the Muon family. @@ -122,13 +126,15 @@ Gefen drops into standard distributed training like any other PyTorch optimizer, Set `deterministic=True` when data-parallel replicas must remain bit-exact on homogeneous GPUs. Automatic block periods are selected with fixed-order GPU reductions instead of the faster atomic CUDA search, block-vmean parameters remain fused through the fixed-order v1 CUDA reduction, and factored-v parameters use the deterministic decomposed factored update instead of the fused stats kernel's unordered floating-point atomics. The default is `False`, so existing performance routing is unchanged. Tagged checkpoints must resume with the same deterministic policy; legacy checkpoints without a tag remain loadable. Plain Gefen does not allow `deterministic=True`, `factored_v_2d=True`, and `stochastic_round=True` together because the deterministic factored fallback uses nearest-codeword quantization. +**Megatron-LM validation scope.** The Megatron integration was exercised through its GPT pretraining entry point with plain Gefen, GefenMuon+AdamW, and GefenMuon+Gefen. Fused deterministic runs covered DP2, TP2, PP2, CP2 with Transformer Engine, and EP2 on two homogeneous RTX 3090 Ti GPUs, with replica or tied-weight hashes appropriate to each topology; legacy optimizer checkpoint continuation was also exercised under TP2 and PP2 for all three recipes. Unfused four-rank coverage additionally exercised EP2×DP2, TP2×DP2, PP2×DP2, and TP2×PP2 for all three recipes, plus EP2×ETP2 and EP2 checkpoint continuation for GefenMuon+Gefen. These are tiny mock-data integration gates, not large-scale convergence claims, and they do not cover Megatron's distributed optimizer, FSDP, optimizer CPU offload, fp16, or non-legacy optimizer checkpoint formats. + ```python optimizer = Gefen(model.named_parameters(), lr=3e-4, fused=True, deterministic=True) ``` ## CUDA Graphs & torch.compile (`capturable`) -All three optimizers accept `capturable=True` (same meaning as `torch.optim`'s argument): `opt.step()` can then be captured in a `torch.cuda.CUDAGraph` or wrapped in `torch.compile(mode="reduce-overhead")` at no step-time cost — and the compiled hybrid step is about 10% faster than eager. Device-resident global counters advance on every replay, so checkpoints record the true replayed step and stochastic-rounding resumes continue from the correct seed. Usage, caveats, and measured numbers: [COMPATIBILITY.md](https://github.com/thad0ctor/Gefen-X/blob/main/COMPATIBILITY.md#cuda-graphs--torchcompile-capturable). +All three optimizers accept `capturable=True` (same meaning as `torch.optim`'s argument): `opt.step()` can then be captured in a `torch.cuda.CUDAGraph` or wrapped in `torch.compile(mode="reduce-overhead")`. Performance is path-dependent rather than guaranteed cost-free: the retained measurements are nearly flat for plain Gefen, show modest eager/manual-capture overhead for the hybrid, and make the compiled hybrid step about 10% faster than default eager. Device-resident global counters advance on every replay, so checkpoints record the true replayed step and stochastic-rounding resumes continue from the correct seed. Usage, graph-partition caveats, and measured numbers: [COMPATIBILITY.md](https://github.com/thad0ctor/Gefen-X/blob/main/COMPATIBILITY.md#cuda-graphs--torchcompile-capturable).
@@ -229,7 +235,7 @@ optimizer: gefenx learning_rate: 6.0e-6 # ≈0.6× your AdamW LR — Gefen takes larger effective steps bf16: true optim_args: - fused: true # ≈2× faster opt.step, bit-exact (default) + fused: true # ≈2× faster opt.step; validated numerical parity (default) factored_v_2d: true # Adafactor-style 2D 2nd moment — matches AdamW loss (default) ``` @@ -483,6 +489,17 @@ trainer = Trainer(model=model, args=training_args, train_dataset=train_dataset, optimizers=(optimizer, None)) # (optimizer, lr_scheduler) ``` +For an executable save/resume gate covering plain Gefen and both Muon backup variants, including gradient accumulation, changing LRs, tied weights, optional BF16, and DDP replica hashes, use [`benchmarks/trainer_resume/`](https://github.com/thad0ctor/Gefen-X/tree/main/benchmarks/trainer_resume). Its Transformers 5 optimizer factory constructs the optimizer from Trainer's model and validates Trainer's internal Accelerate wrapping. All three recipes have passed the deterministic fused-BF16 two-rank gate on homogeneous GPUs, which requires exact uninterrupted-versus-resumed model, optimizer, scheduler, loss-history, and cross-replica hashes: + +```bash +CUDA_VISIBLE_DEVICES=0,1 PYTHONPATH=.:src torchrun --standalone --nproc-per-node=2 \ + -m benchmarks.trainer_resume.run --output-dir benchmarks/trainer_resume/out/ddp \ + --device cuda --dtype bfloat16 --fused --deterministic --steps 3 --split-step 1 \ + --gradient-accumulation-steps 2 +``` + +The harness tests Trainer's DDP frontend and internal Accelerate optimizer wrapper; it does not test a standalone `Accelerator` loop, FSDP, or `model_init`. + ## Quality Lever: factored second moment on 2D params (`factored_v_2d`) > [!NOTE] @@ -509,7 +526,7 @@ opt = Gefen(model.named_parameters(), lr=0.6 * ADAMW_LR, fused=True) # factored Details: validation, checkpoint migration, and limits - Validated at both scales in the fair-LR regime, with a second-seed replication at 0.6B and a learning-rate sweep at 1.7B (best LR `3e-5` ~ 0.6× AdamW's, matching Gefen's documented heuristic). The fused kernel computes the per-element step size in registers (no extra temporaries; step transients measured 0 MiB) and is covered by `tests/test_gefen_factored_v.py`. -- Checkpoints migrate automatically in both directions (old checkpoints work with the new default and vice versa; the second-moment statistics re-warm briefly and harmlessly). +- Native Gefen checkpoints migrate automatically in both directions (old checkpoints work with the new default and vice versa; the second-moment statistics re-warm briefly and harmlessly). Rank-local DTensor full-state DCP uses a separate tagged, collective format and requires the same world size and topology. - With `backup_optimizer="gefen"`, `GefenMuonHybrid` pins the backup half to Gefen's legacy block-vmean path (the factored-v combination has not been benchmarked). With `backup_optimizer="adamw"`, the backup uses conventional per-element AdamW state. Under FSDP2, sharded 2D Gefen params also fall back to the legacy path. - Two sibling experiments from the same investigation ship off by default because they measured **no effect**: `period_one_substrings` (per-element state on name-matched tensors) and `codebook_refresh_every` (periodic codebook refit). @@ -580,7 +597,7 @@ A related opt-in, `stochastic_round=True`, switches the 8-bit momentum to unbias Under FSDP2, Muon's orthogonalization needs each full weight matrix, but every GPU only holds a slice. Three strategies: -- **`"exact"`** (default) — every GPU rebuilds every matrix and does the same work. Identical to single-GPU results, but redundant. +- **`"exact"`** (default) — every GPU rebuilds every matrix and does the same work. It matches the single-GPU oracle within the tested BF16 tolerance, but the cross-rank GEMM path is not promised bitwise identical; the work is redundant. - **`"distributed"`** (experimental) — each matrix is assigned by its stable position in the full distributed parameter set to one GPU, which does the work and shares the result. Bit-identical to `"exact"` on homogeneous GPUs and faster as you add them. Momentum ownership is stable when the active gradient set varies, and `state_dict()` collectively gathers owner-local momentum so rank 0 can write a complete checkpoint after all ranks call it. - **`"approx"`** — each GPU works on just its slice. Fastest, but results genuinely differ — an accuracy trade. @@ -600,7 +617,7 @@ opt = GefenMuonHybrid( ![Gefen-Muon exact / distributed / approx sharded — eval loss](https://raw.githubusercontent.com/thad0ctor/Gefen-X/main/docs/benchmarks/muon_shard_loss.png) ![Gefen-Muon exact / distributed / approx sharded — throughput & VRAM](https://raw.githubusercontent.com/thad0ctor/Gefen-X/main/docs/benchmarks/muon_shard_perf.png) -Measured (Qwen3-0.6B, 2 and 4 GPUs): `"distributed"` is a free speedup at identical results (1.06× / 1.12×); `"approx"` is faster still (1.25× / 1.39×) but visibly costs training quality, and the cost grows with GPU count. The `"distributed"` win grows with model size. +Measured (Qwen3-0.6B, 2 and 4 GPUs): `"distributed"` matched `"exact"` in the retained convergence runs while improving throughput by 1.06× / 1.12×; `"approx"` is faster still (1.25× / 1.39×) but visibly costs training quality, and the cost grows with GPU count. The `"distributed"` win grows with model size.
@@ -615,6 +632,8 @@ Measured (Qwen3-0.6B, 2 and 4 GPUs): `"distributed"` is a free speedup at identi ## Known limitations - **Hybrid checkpoint schema.** `GefenMuonHybrid`'s `state_dict()` uses its own nested `{"muon": ..., "backup": ..., "backup_optimizer": "gefen" | "adamw"}` layout. Resume from a checkpoint the hybrid itself saved—not one consolidated or converted to the flat torch `{state, param_groups}` layout. Cross-backend loads are rejected before either child is mutated; legacy untagged hybrid checkpoints are interpreted as Gefen-backed. +- **FSDP2 full-state optimizer DCP is same-topology only.** Plain Gefen and `GefenMuon(sharded_mode="approx")` preserve rank-local state in a tagged collective payload and resume exactly with the same world size, one-dimensional default-world mesh, placements, rank coordinates, parameter layout, and mode. Every rank must enter save and restore, and each process can transiently use about `world_size ×` its local optimizer-state size in CPU memory plus local serialization scratch. Multidimensional meshes, subgroups, pipeline-local optimizers, and resharding are intentionally rejected; old untagged full checkpoints that could have reused rank 0's codebook on every rank are also rejected. This limitation does not apply to model-only DCP. +- **Accelerate cannot observe native true-FP16 overflow skips.** PyTorch's native AMP optimizer protocol calls `optimizer.step()` even when the optimizer returns before mutation, so Accelerate's `step_was_skipped` flag remains false and a scheduler driven solely by that flag can advance. Normal FP32-master autocast and BF16 use the ordinary GradScaler path and are unaffected; prefer those modes in Trainer/Accelerate, or gate a true-FP16 scheduler from the scaler's scale change. ## Troubleshooting diff --git a/benchmarks/README.md b/benchmarks/README.md index 443600d..77710fc 100644 --- a/benchmarks/README.md +++ b/benchmarks/README.md @@ -8,6 +8,7 @@ These suites measure optimizer quality, throughput, memory, and Muon-specific ex |---|---|---| | [`optimizer-sweep/`](optimizer-sweep/README.md) | Compare AdamW, Gefen, and Gefen-Muon on full-model SFT: validation loss, throughput, peak VRAM, and optimizer-state bytes per parameter | `bash benchmarks/optimizer-sweep/run.sh` | | [`training_matrix/`](training_matrix/README.md) | Compare controlled AdamW/Muon recipes across full-model HF SFT, small-model pretraining, and checkpoint handoff | `PYTHONPATH=.:src python -m benchmarks.training_matrix.run_matrix` | +| [`trainer_resume/`](trainer_resume/README.md) | Validate Trainer's internal Accelerate wrapping, accumulation, scheduling, tied weights, DDP replicas, and exact native Trainer checkpoint continuation | `PYTHONPATH=.:src python -m benchmarks.trainer_resume.run` | | [`sharding-sweep/`](sharding-sweep/README.md) | Compare Gefen-Muon `sharded_mode` choices under FSDP2 across world sizes | `bash benchmarks/sharding-sweep/run.sh` | | [`microbench/`](microbench/) | Measure Newton–Schulz schedules, fp8 and batched kernels, capturable steps, and distributed Muon internals | `PYTHONPATH=.:src python benchmarks/microbench/bench_ns_schedule.py --help` | @@ -44,6 +45,14 @@ PYTHONPATH=.:src python -m benchmarks.training_matrix.run_matrix --cells adamw - --device cpu --dtype float32 --no-fused ``` +Transformers Trainer checkpoint continuation: + +```bash +PYTHONPATH=.:src python -m benchmarks.trainer_resume.run \ + --output-dir benchmarks/trainer_resume/out/cpu \ + --device cpu --dtype float32 --no-fused +``` + FSDP2 sharding comparison: ```bash diff --git a/benchmarks/trainer_resume/README.md b/benchmarks/trainer_resume/README.md new file mode 100644 index 0000000..12c8588 --- /dev/null +++ b/benchmarks/trainer_resume/README.md @@ -0,0 +1,45 @@ +# Transformers Trainer resume harness + +This harness validates plain Gefen, Gefen-Muon with an AdamW backup, and Gefen-Muon with a Gefen backup through the Hugging Face `Trainer` lifecycle. It compares an uninterrupted deterministic run against a native Trainer checkpoint continuation and fails unless model, optimizer, scheduler, learning-rate, and logged-loss state agree exactly. + +The run also verifies that Trainer creates the optimizer from the model instance passed to it, Trainer's internal Accelerate layer wraps the optimizer, gradient accumulation produces one optimizer step per logical update, tied parameters are routed once, checkpoint hooks reach the underlying optimizer, and DDP replicas retain identical state hashes. This is not a standalone Accelerate harness: it does not exercise a direct `Accelerator` loop or `Accelerator.save_state()`/`load_state()`. It also does not exercise FSDP or a `model_init` callback. + +Install the optional framework dependencies and Gefen-X from the repository root: + +```bash +python -m pip install -e . +python -m pip install 'transformers>=5,<6' 'accelerate>=1.1' +``` + +Run the fast CPU matrix: + +```bash +PYTHONPATH=.:src python -m benchmarks.trainer_resume.run \ + --output-dir benchmarks/trainer_resume/out/cpu \ + --device cpu --dtype float32 --no-fused +``` + +Run the fused BF16 matrix against a local model: + +```bash +CUDA_VISIBLE_DEVICES=0 PYTHONPATH=.:src python -m benchmarks.trainer_resume.run \ + --output-dir benchmarks/trainer_resume/out/qwen3-06b \ + --model /path/to/Qwen3-0.6B --attn-implementation eager \ + --device cuda --dtype bfloat16 --fused --gradient-checkpointing \ + --steps 6 --split-step 3 --batch-size 1 --gradient-accumulation-steps 2 --seq-len 64 +``` + +Launch the same validation through Trainer's DDP frontend with homogeneous GPUs: + +```bash +CUDA_VISIBLE_DEVICES=0,1 PYTHONPATH=.:src torchrun --standalone --nproc-per-node=2 \ + -m benchmarks.trainer_resume.run \ + --output-dir benchmarks/trainer_resume/out/qwen3-06b-ddp \ + --model /path/to/Qwen3-0.6B --attn-implementation eager \ + --device cuda --dtype bfloat16 --fused --gradient-checkpointing \ + --steps 6 --split-step 3 --batch-size 1 --gradient-accumulation-steps 2 --seq-len 64 +``` + +Each recipe retains the staged Trainer checkpoint and writes a machine-readable `summary.json` under the selected output directory. Use a new output directory for each run; the harness refuses to mix results into a non-empty phase directory. + +The tiny-model CPU command is the fast lifecycle gate. The local-model BF16 and DDP commands exercise the additional options shown; their exact-continuation assertion is a test condition for the selected model, attention implementation, hardware, and framework versions, not a claim that every Transformers architecture is bitwise reproducible across restart. The harness uses native Trainer checkpoints and does not validate FSDP/DCP optimizer checkpoint conversion. diff --git a/benchmarks/trainer_resume/__init__.py b/benchmarks/trainer_resume/__init__.py new file mode 100644 index 0000000..61c1103 --- /dev/null +++ b/benchmarks/trainer_resume/__init__.py @@ -0,0 +1,2 @@ +"""Hugging Face Trainer checkpoint-continuation validation.""" + diff --git a/benchmarks/trainer_resume/run.py b/benchmarks/trainer_resume/run.py new file mode 100644 index 0000000..dc5aa7f --- /dev/null +++ b/benchmarks/trainer_resume/run.py @@ -0,0 +1,646 @@ +#!/usr/bin/env python3 +"""Validate Gefen optimizer recipes through Hugging Face Trainer save/resume. + +The harness compares an uninterrupted run with a stopped-and-resumed run. It +uses Trainer's model-aware optimizer-factory lifecycle, Accelerate wrapping, +gradient accumulation, a changing learning-rate schedule, and native Trainer +checkpoints. The default tiny tied-weight GPT-2 model is dependency-light; +``--model`` accepts a local causal-LM checkpoint for a release gate. +""" + +from __future__ import annotations + +import argparse +import gc +import hashlib +import json +import os +import platform +import struct +import subprocess +from dataclasses import asdict, dataclass +from pathlib import Path +from typing import Any, Iterable + +import torch +from torch.utils.data import Dataset + +from gefen import Gefen, GefenMuonHybrid, __version__ as gefen_version + + +RECIPES = ("gefen", "gefen_muon_adamw", "gefen_muon_gefen") + + +def _positive_int(value: str) -> int: + parsed = int(value) + if parsed <= 0: + raise argparse.ArgumentTypeError("must be > 0") + return parsed + + +def _nonnegative_int(value: str) -> int: + parsed = int(value) + if parsed < 0: + raise argparse.ArgumentTypeError("must be >= 0") + return parsed + + +def _recipe_list(value: str) -> tuple[str, ...]: + recipes = tuple(part.strip() for part in value.split(",") if part.strip()) + unknown = sorted(set(recipes) - set(RECIPES)) + if not recipes or unknown: + choices = ", ".join(RECIPES) + detail = "no recipes supplied" if not recipes else f"unknown recipes: {', '.join(unknown)}" + raise argparse.ArgumentTypeError(f"{detail}; choose from {choices}") + if len(recipes) != len(set(recipes)): + raise argparse.ArgumentTypeError("recipes must not contain duplicates") + return recipes + + +def parse_args(argv: list[str] | None = None) -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__, allow_abbrev=False) + parser.add_argument("--output-dir", type=Path, required=True) + parser.add_argument("--model", help="local causal-LM checkpoint; omit for a tiny random GPT-2") + parser.add_argument("--trust-remote-code", action="store_true") + parser.add_argument("--attn-implementation", choices=("auto", "eager", "sdpa"), default="auto") + parser.add_argument("--recipes", type=_recipe_list, default=RECIPES) + parser.add_argument("--steps", type=_positive_int, default=6) + parser.add_argument("--split-step", type=_positive_int, default=3) + parser.add_argument("--warmup-steps", type=_nonnegative_int, default=1) + parser.add_argument("--batch-size", type=_positive_int, default=2) + parser.add_argument("--gradient-accumulation-steps", type=_positive_int, default=2) + parser.add_argument("--seq-len", type=_positive_int, default=64) + parser.add_argument("--seed", type=int, default=17) + parser.add_argument("--lr", type=float, default=1e-3) + parser.add_argument("--backup-lr", type=float, help="Gefen-backup LR; defaults to half --lr") + parser.add_argument("--weight-decay", type=float, default=0.0) + parser.add_argument("--max-grad-norm", type=float, default=1.0) + parser.add_argument("--ns-steps", type=_positive_int, default=3) + parser.add_argument("--device", choices=("cpu", "cuda"), default="cuda" if torch.cuda.is_available() else "cpu") + parser.add_argument("--dtype", choices=("auto", "float32", "bfloat16"), default="auto") + parser.add_argument("--fused", action=argparse.BooleanOptionalAction, default=None) + parser.add_argument("--deterministic", action=argparse.BooleanOptionalAction, default=True) + parser.add_argument("--gradient-checkpointing", action=argparse.BooleanOptionalAction, default=False) + return parser.parse_args(argv) + + +class FixedTokenDataset(Dataset): + """Deterministic causal-LM rows with no online tokenization or transforms.""" + + def __init__(self, *, samples: int, seq_len: int, vocab_size: int, seed: int): + generator = torch.Generator().manual_seed(seed) + self.input_ids = torch.randint(0, vocab_size, (samples, seq_len), generator=generator) + + def __len__(self) -> int: + return self.input_ids.shape[0] + + def __getitem__(self, index: int) -> dict[str, torch.Tensor]: + input_ids = self.input_ids[index] + return { + "input_ids": input_ids, + "attention_mask": torch.ones_like(input_ids), + "labels": input_ids.clone(), + } + + +@dataclass +class HookCounts: + steps: int = 0 + state_dicts: int = 0 + loads: int = 0 + + +def _increment(tracker: HookCounts, field: str) -> None: + setattr(tracker, field, getattr(tracker, field) + 1) + + +class GefenTrainerOptimizerFactory: + """Transformers 5 model-aware optimizer factory used by ``Trainer``. + + Trainer recognizes non-Optimizer classes as factories, constructs one with + no arguments, then calls it with the model it owns. This preserves names, + tied-weight/module-aware Muon routing, and wrapper lifecycle ordering. + """ + + def __call__( + self, + model: torch.nn.Module, + *, + recipe: str, + lr: float, + backup_lr: float, + weight_decay: float, + fused: bool, + deterministic: bool, + ns_steps: int, + hook_counts: HookCounts, + ) -> torch.optim.Optimizer: + if recipe == "gefen": + optimizer = Gefen( + model.named_parameters(), + lr=lr, + weight_decay=weight_decay, + fused=fused, + deterministic=deterministic, + ) + elif recipe in ("gefen_muon_adamw", "gefen_muon_gefen"): + backup_optimizer = recipe.removeprefix("gefen_muon_") + optimizer = GefenMuonHybrid.from_model( + model, + lr=lr, + muon_lr=lr, + backup_lr=lr if backup_optimizer == "adamw" else backup_lr, + backup_optimizer=backup_optimizer, + weight_decay=weight_decay, + fused=fused, + deterministic=deterministic, + ns_steps=ns_steps, + ns_schedule="tuned3" if ns_steps == 3 else "standard", + normuon=True, + backup_1d_period_one=backup_optimizer == "gefen", + ) + else: # pragma: no cover - guarded by argparse and direct-call validation + raise ValueError(f"unknown recipe {recipe!r}") + + optimizer.register_step_post_hook(lambda *_args, **_kwargs: _increment(hook_counts, "steps")) + optimizer.register_state_dict_post_hook( + lambda _optimizer, state_dict: (_increment(hook_counts, "state_dicts"), state_dict)[1] + ) + optimizer.register_load_state_dict_post_hook(lambda _optimizer: _increment(hook_counts, "loads")) + return optimizer + + +class StopAfterStep: + """Callback implementation is completed lazily after Transformers imports.""" + + def __new__(cls, stop_step: int): + from transformers import TrainerCallback + + class _StopAfterStep(TrainerCallback): + def on_step_end(self, args, state, control, **kwargs): + if state.global_step >= stop_step: + control.should_training_stop = True + return control + + return _StopAfterStep() + + +def _feed_digest(digest: Any, value: Any) -> None: + if torch.is_tensor(value): + tensor = value.detach().contiguous().cpu() + digest.update(b"tensor\0") + digest.update(str(tensor.dtype).encode()) + digest.update(b"\0") + digest.update(json.dumps(list(tensor.shape), separators=(",", ":")).encode()) + digest.update(b"\0") + digest.update(tensor.reshape(-1).view(torch.uint8).numpy().tobytes()) + elif isinstance(value, dict): + digest.update(b"dict\0") + keys = sorted(value, key=lambda item: (type(item).__qualname__, repr(item))) + for key in keys: + _feed_digest(digest, key) + _feed_digest(digest, value[key]) + elif isinstance(value, (list, tuple)): + digest.update(b"list\0" if isinstance(value, list) else b"tuple\0") + for item in value: + _feed_digest(digest, item) + elif isinstance(value, bool): + digest.update(b"bool\1" if value else b"bool\0") + elif isinstance(value, int): + encoded = str(value).encode() + digest.update(b"int\0" + encoded + b"\0") + elif isinstance(value, float): + digest.update(b"float\0" + struct.pack("!d", value)) + elif isinstance(value, str): + encoded = value.encode("utf-8", errors="surrogateescape") + digest.update(b"str\0" + str(len(encoded)).encode() + b"\0" + encoded) + elif value is None: + digest.update(b"none\0") + else: + raise TypeError(f"unsupported digest value {type(value).__qualname__}") + + +def state_digest(value: Any) -> str: + digest = hashlib.sha256() + _feed_digest(digest, value) + return digest.hexdigest() + + +def _git_metadata() -> dict[str, Any]: + def _git(*args: str) -> str | None: + try: + return subprocess.run( + ("git", *args), check=True, text=True, stdout=subprocess.PIPE, stderr=subprocess.DEVNULL + ).stdout.strip() + except (OSError, subprocess.CalledProcessError): + return None + + status = _git("status", "--porcelain") + return { + "commit": _git("rev-parse", "HEAD"), + "branch": _git("branch", "--show-current"), + "dirty": None if status is None else bool(status), + } + + +def _distributed_values(value: Any) -> list[Any]: + if not torch.distributed.is_available() or not torch.distributed.is_initialized(): + return [value] + values: list[Any] = [None] * torch.distributed.get_world_size() + torch.distributed.all_gather_object(values, value) + return values + + +def _assert_replicated(label: str, value: Any) -> None: + values = _distributed_values(value) + if any(item != values[0] for item in values[1:]): + raise AssertionError(f"{label} differs across ranks: {values}") + + +def _world_size() -> int: + if torch.distributed.is_available() and torch.distributed.is_initialized(): + return torch.distributed.get_world_size() + return int(os.environ.get("WORLD_SIZE", "1")) + + +def _rank() -> int: + if torch.distributed.is_available() and torch.distributed.is_initialized(): + return torch.distributed.get_rank() + return int(os.environ.get("RANK", "0")) + + +def _wait() -> None: + if torch.distributed.is_available() and torch.distributed.is_initialized(): + kwargs = {} + if "nccl" in str(torch.distributed.get_backend()).lower(): + kwargs["device_ids"] = [torch.cuda.current_device()] + torch.distributed.barrier(**kwargs) + + +def _device(args: argparse.Namespace) -> torch.device: + if args.device == "cpu": + return torch.device("cpu") + if not torch.cuda.is_available(): + raise RuntimeError("--device cuda requested but CUDA is unavailable") + local_rank = int(os.environ.get("LOCAL_RANK", "0")) + torch.cuda.set_device(local_rank) + return torch.device("cuda", local_rank) + + +def _dtype(args: argparse.Namespace) -> torch.dtype: + name = args.dtype + if name == "auto": + name = "bfloat16" if args.device == "cuda" else "float32" + return {"float32": torch.float32, "bfloat16": torch.bfloat16}[name] + + +def _load_model(args: argparse.Namespace, dtype: torch.dtype, device: torch.device): + from transformers import AutoModelForCausalLM, GPT2Config, GPT2LMHeadModel, set_seed + + set_seed(args.seed) + if args.model: + kwargs: dict[str, Any] = { + "trust_remote_code": args.trust_remote_code, + "local_files_only": Path(args.model).exists(), + } + if args.attn_implementation != "auto": + kwargs["attn_implementation"] = args.attn_implementation + try: + model = AutoModelForCausalLM.from_pretrained(args.model, dtype=dtype, **kwargs) + except TypeError: # transformers<4.56 used torch_dtype + model = AutoModelForCausalLM.from_pretrained(args.model, torch_dtype=dtype, **kwargs) + else: + config = GPT2Config( + vocab_size=128, + n_positions=max(args.seq_len, 16), + n_ctx=max(args.seq_len, 16), + n_embd=32, + n_layer=1, + n_head=4, + n_inner=64, + resid_pdrop=0.0, + embd_pdrop=0.0, + attn_pdrop=0.0, + use_cache=False, + tie_word_embeddings=True, + ) + model = GPT2LMHeadModel(config).to(dtype=dtype) + model.config.use_cache = False + model.to(device) + return model + + +def _routing_census(model: torch.nn.Module, optimizer: torch.optim.Optimizer) -> dict[str, Any]: + trainable = [parameter for parameter in model.parameters() if parameter.requires_grad] + routed = [parameter for group in optimizer.param_groups for parameter in group["params"]] + if len(routed) != len({id(parameter) for parameter in routed}): + raise AssertionError("optimizer routes a trainable parameter more than once") + if {id(parameter) for parameter in routed} != {id(parameter) for parameter in trainable}: + raise AssertionError("optimizer routing does not cover every trainable parameter exactly once") + + aliases: dict[int, list[str]] = {} + for name, parameter in model.named_parameters(remove_duplicate=False): + if parameter.requires_grad: + aliases.setdefault(id(parameter), []).append(name) + tied = [names for names in aliases.values() if len(names) > 1] + census: dict[str, Any] = { + "parameter_tensors": len(trainable), + "parameters": sum(parameter.numel() for parameter in trainable), + "optimizer_groups": len(optimizer.param_groups), + "tied_parameter_aliases": tied, + } + if isinstance(optimizer, GefenMuonHybrid): + census.update( + { + "muon_tensors": sum(len(group["params"]) for group in optimizer.muon.param_groups) + if optimizer.muon is not None + else 0, + "backup_tensors": sum(len(group["params"]) for group in optimizer.backup.param_groups) + if optimizer.backup is not None + else 0, + "backup_optimizer": optimizer.backup_optimizer, + } + ) + return census + + +def _unwrap_optimizer(optimizer: Any) -> torch.optim.Optimizer: + current = optimizer + while type(current).__module__.startswith("accelerate.") and hasattr(current, "optimizer"): + current = current.optimizer + if not isinstance(current, torch.optim.Optimizer): + raise TypeError(f"Trainer produced an unexpected optimizer wrapper {type(current)!r}") + return current + + +def _loss_history(history: Iterable[dict[str, Any]]) -> list[dict[str, Any]]: + keys = ("step", "loss", "grad_norm", "learning_rate", "epoch") + return [{key: row[key] for key in keys if key in row} for row in history if "loss" in row] + + +def _checkpoint_files(path: Path) -> list[str]: + files = sorted(item.name for item in path.iterdir() if item.is_file()) + required = {"optimizer.pt", "scheduler.pt", "trainer_state.json"} + missing = sorted(required - set(files)) + if missing: + raise AssertionError(f"Trainer checkpoint is missing {missing}: {path}") + if not any(name.startswith("rng_state") and name.endswith(".pth") for name in files): + raise AssertionError(f"Trainer checkpoint has no RNG state: {path}") + if not any(name in files for name in ("model.safetensors", "pytorch_model.bin")): + raise AssertionError(f"Trainer checkpoint has no model weights: {path}") + return files + + +def _phase_dir(output_dir: Path, recipe: str, phase: str) -> Path: + path = output_dir / recipe / phase + if path.exists() and any(path.iterdir()): + raise FileExistsError(f"refusing to mix with non-empty phase directory: {path}") + path.mkdir(parents=True, exist_ok=True) + return path + + +def _run_phase( + args: argparse.Namespace, + *, + recipe: str, + phase: str, + dataset: Dataset, + dtype: torch.dtype, + device: torch.device, + resume_from: Path | None = None, + stop_step: int | None = None, +) -> dict[str, Any]: + from transformers import Trainer, TrainingArguments + + model = _load_model(args, dtype, device) + tracker = HookCounts() + callbacks = [StopAfterStep(stop_step)] if stop_step is not None else [] + phase_dir = _phase_dir(args.output_dir, recipe, phase) + training_args = TrainingArguments( + output_dir=str(phase_dir), + per_device_train_batch_size=args.batch_size, + gradient_accumulation_steps=args.gradient_accumulation_steps, + max_steps=args.steps, + learning_rate=args.lr, + lr_scheduler_type="linear", + warmup_steps=args.warmup_steps, + weight_decay=args.weight_decay, + max_grad_norm=args.max_grad_norm, + bf16=dtype == torch.bfloat16, + gradient_checkpointing=args.gradient_checkpointing, + save_strategy="steps" if stop_step is not None else "no", + save_steps=args.split_step, + save_total_limit=1, + logging_strategy="steps", + logging_steps=1, + logging_first_step=True, + report_to="none", + disable_tqdm=True, + full_determinism=args.deterministic, + seed=args.seed, + data_seed=args.seed, + use_cpu=args.device == "cpu", + ddp_find_unused_parameters=False, + dataloader_num_workers=0, + dataloader_pin_memory=args.device == "cuda", + ) + optimizer_kwargs = { + "recipe": recipe, + "lr": args.lr, + "backup_lr": args.backup_lr, + "weight_decay": args.weight_decay, + "fused": args.fused, + "deterministic": args.deterministic, + "ns_steps": args.ns_steps, + "hook_counts": tracker, + } + trainer = Trainer( + model=model, + args=training_args, + train_dataset=dataset, + callbacks=callbacks, + optimizer_cls_and_kwargs=(GefenTrainerOptimizerFactory, optimizer_kwargs), + ) + trainer.train(resume_from_checkpoint=str(resume_from) if resume_from is not None else None) + wrapped_optimizer = trainer.optimizer + optimizer = _unwrap_optimizer(wrapped_optimizer) + if not type(wrapped_optimizer).__module__.startswith("accelerate."): + raise AssertionError(f"Trainer optimizer was not wrapped by Accelerate: {type(wrapped_optimizer)!r}") + expected_steps = stop_step if stop_step is not None else args.steps - (args.split_step if resume_from else 0) + if tracker.steps != expected_steps: + raise AssertionError( + f"{recipe}/{phase}: optimizer step hook fired {tracker.steps} times, expected {expected_steps}" + ) + if resume_from is not None and tracker.loads < 1: + raise AssertionError(f"{recipe}/{phase}: Trainer did not load the underlying optimizer state") + + framework_hook_counts = asdict(tracker) + routing = _routing_census(model, optimizer) + model_hash = state_digest(model.state_dict()) + optimizer_hash = state_digest(optimizer.state_dict()) + scheduler_hash = state_digest(trainer.lr_scheduler.state_dict()) + _assert_replicated(f"{recipe}/{phase} model hash", model_hash) + _assert_replicated(f"{recipe}/{phase} optimizer hash", optimizer_hash) + _assert_replicated(f"{recipe}/{phase} scheduler hash", scheduler_hash) + result = { + "global_step": trainer.state.global_step, + "model_sha256": model_hash, + "optimizer_sha256": optimizer_hash, + "scheduler_sha256": scheduler_hash, + "loss_history": _loss_history(trainer.state.log_history), + "hook_counts": framework_hook_counts, + "routing": routing, + "optimizer_wrapper": f"{type(wrapped_optimizer).__module__}.{type(wrapped_optimizer).__qualname__}", + "learning_rates": [float(group["lr"]) for group in optimizer.param_groups], + } + del trainer, optimizer, wrapped_optimizer, model + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + _wait() + return result + + +def _validate_args(args: argparse.Namespace) -> None: + if not 0 < args.split_step < args.steps: + raise ValueError(f"--split-step must satisfy 0 < split < steps, got {args.split_step}/{args.steps}") + if args.warmup_steps >= args.steps: + raise ValueError("--warmup-steps must be smaller than --steps") + if args.lr <= 0: + raise ValueError("--lr must be > 0") + if args.weight_decay < 0: + raise ValueError("--weight-decay must be >= 0") + if args.max_grad_norm <= 0: + raise ValueError("--max-grad-norm must be > 0") + if args.fused is None: + args.fused = args.device == "cuda" + if args.device == "cpu" and args.fused: + raise ValueError("--fused requires --device cuda; pass --no-fused for a CPU run") + if args.backup_lr is None: + args.backup_lr = 0.5 * args.lr + if args.backup_lr <= 0: + raise ValueError("--backup-lr must be > 0") + + +def run(args: argparse.Namespace) -> dict[str, Any]: + try: + import accelerate + import transformers + except ImportError as exc: # pragma: no cover - exercised in minimal installs + raise RuntimeError( + "this harness requires transformers and accelerate; install them with " + "`python -m pip install transformers 'accelerate>=1.1'`" + ) from exc + + _validate_args(args) + device = _device(args) + dtype = _dtype(args) + args.output_dir.mkdir(parents=True, exist_ok=True) + + probe_model = _load_model(args, dtype, device) + vocab_size = int(probe_model.config.vocab_size) + del probe_model + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + samples = max( + 16, + args.steps * args.batch_size * args.gradient_accumulation_steps * _world_size() * 2, + ) + dataset = FixedTokenDataset(samples=samples, seq_len=args.seq_len, vocab_size=vocab_size, seed=args.seed + 1) + + results: dict[str, Any] = {} + for recipe in args.recipes: + baseline = _run_phase( + args, recipe=recipe, phase="baseline", dataset=dataset, dtype=dtype, device=device + ) + staged = _run_phase( + args, + recipe=recipe, + phase="staged", + dataset=dataset, + dtype=dtype, + device=device, + stop_step=args.split_step, + ) + checkpoint = args.output_dir / recipe / "staged" / f"checkpoint-{args.split_step}" + _wait() + checkpoint_files = _checkpoint_files(checkpoint) + resumed = _run_phase( + args, + recipe=recipe, + phase="resumed", + dataset=dataset, + dtype=dtype, + device=device, + resume_from=checkpoint, + ) + + if staged["global_step"] != args.split_step: + raise AssertionError(f"{recipe}: staged Trainer stopped at {staged['global_step']}") + if baseline["global_step"] != args.steps or resumed["global_step"] != args.steps: + raise AssertionError(f"{recipe}: baseline/resumed Trainer did not reach step {args.steps}") + for key in ("model_sha256", "optimizer_sha256", "scheduler_sha256", "learning_rates"): + if baseline[key] != resumed[key]: + raise AssertionError(f"{recipe}: baseline/resumed {key} mismatch") + if baseline["loss_history"][: args.split_step] != staged["loss_history"]: + raise AssertionError(f"{recipe}: pre-checkpoint loss/schedule history mismatch") + if baseline["loss_history"] != resumed["loss_history"]: + raise AssertionError(f"{recipe}: post-resume loss/schedule history mismatch") + if staged["hook_counts"]["state_dicts"] < 1: + raise AssertionError(f"{recipe}: Trainer checkpoint did not call optimizer.state_dict()") + if baseline["routing"] != resumed["routing"]: + raise AssertionError(f"{recipe}: tied-weight/parameter routing changed after resume") + + results[recipe] = { + "passed": True, + "baseline": baseline, + "staged": staged, + "resumed": resumed, + "checkpoint": str(checkpoint), + "checkpoint_files": checkpoint_files, + } + + cuda_devices = [] + if args.device == "cuda": + cuda_devices = _distributed_values(torch.cuda.get_device_name(device)) + summary = { + "schema_version": 1, + "passed": True, + "config": { + **vars(args), + "output_dir": str(args.output_dir), + "recipes": list(args.recipes), + "dtype": str(dtype).removeprefix("torch."), + }, + "runtime": { + "python": platform.python_version(), + "torch": torch.__version__, + "cuda": torch.version.cuda, + "transformers": transformers.__version__, + "accelerate": accelerate.__version__, + "gefen_x": gefen_version, + "world_size": _world_size(), + "cuda_devices": cuda_devices, + "git": _git_metadata(), + }, + "recipes": results, + } + _wait() + if _rank() == 0: + summary_path = args.output_dir / "summary.json" + summary_path.write_text(json.dumps(summary, indent=2, sort_keys=True) + "\n") + print(json.dumps(summary, indent=2, sort_keys=True)) + return summary + + +def main(argv: list[str] | None = None) -> int: + run(parse_args(argv)) + if torch.distributed.is_available() and torch.distributed.is_initialized(): + _wait() + torch.distributed.destroy_process_group() + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/pyproject.toml b/pyproject.toml index f909c3f..7b70710 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,13 +1,13 @@ # Python packaging metadata used by pip/build to install Gefen as an importable package. # Distribution name is `gefen-x` (this fork); the import name stays `gefen`. [build-system] -requires = ["setuptools>=77", "wheel"] +requires = ["setuptools==80.10.2", "wheel==0.47.0"] build-backend = "setuptools.build_meta" [project] name = "gefen-x" # Single-sourced at runtime via importlib.metadata (see __init__.py __version__). -version = "0.3.0" +version = "0.4.0" description = "Gefen optimizer for memory-efficient PyTorch training" readme = "README.md" requires-python = ">=3.10" @@ -33,17 +33,18 @@ classifiers = [ # Hard runtime deps. The code needs FSDP2 / public DTensor namespace # (torch.distributed.tensor), so torch>=2.5. numba is required: it compiles the # exact-DP codebook solver, and the pure-Python fallback in quantization.py is -# too slow (O(bins^2 * codebooks)) for real training. ninja is optional (CUDA -# JIT builds are faster with it, but torch can build without) -> the `perf` extra. +# too slow (O(bins^2 * codebooks)) for real training. ninja and setuptools are +# required by torch.utils.cpp_extension.load, which builds the fused CUDA kernels +# at runtime. dependencies = [ "numpy>=1.24", "torch>=2.5", - "numba>=0.65" + "numba>=0.65", + "ninja", + "setuptools>=77" ] [project.optional-dependencies] -# ninja speeds up the CUDA kernel JIT build (optional; torch builds without it). -perf = ["ninja"] test = ["pytest"] [project.urls] diff --git a/src/gefen/gefen.py b/src/gefen/gefen.py index c93402a..ee226c7 100644 --- a/src/gefen/gefen.py +++ b/src/gefen/gefen.py @@ -10,10 +10,12 @@ helpers it steps through; ``gefen_muon``/``hybrid`` build on it. """ +import io import logging import math import os import warnings +from collections import OrderedDict from itertools import chain from typing import Iterable, Optional, Tuple, Union @@ -49,6 +51,15 @@ logger = logging.getLogger(__name__) +_RANK_LOCAL_PAYLOAD_KEY_PREFIX = "_gefen_rank_local_payload_" +_RANK_LOCAL_MEMBER_KEY = "_gefen_rank_local_member" +_RANK_LOCAL_FORMAT = "rank_local_dtensor_v2" +_RANK_LOCAL_METADATA_VERSION = 3 + + +def _rank_local_payload_key(global_rank: int) -> str: + return "{}{}".format(_RANK_LOCAL_PAYLOAD_KEY_PREFIX, int(global_rank)) + # Optional explicit override for the period-search backend ("cuda_kernel" / "cpu" # / "gpu"). None (default) means resolve per call from each tensor's own device # -- see _resolve_find_period_backend. NOT used as an auto-populated cache. @@ -486,6 +497,308 @@ def _resolve_find_period_backend(grad: torch.Tensor) -> str: return "cuda_kernel" if grad_work.device.type == "cuda" else "cpu" +@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. + """ + found_inf = getattr(optimizer, "found_inf", None) + if found_inf is not None: + if torch.is_tensor(found_inf): + if found_inf.numel() != 1: + raise RuntimeError( + "GradScaler supplied a non-scalar found_inf tensor with " + "shape {}".format(tuple(found_inf.shape)) + ) + overflow = bool(found_inf.detach().item()) + else: + overflow = bool(found_inf) + if overflow: + return 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 = grad_scale.detach() + # Match torch.amp.GradScaler.unscale_: computing the reciprocal in + # fp64 avoids compile-option-dependent fp32 division differences. + inv_scale = scale.double().reciprocal().float() + else: + inv_scale = torch.tensor(float(grad_scale), dtype=torch.float64) + inv_scale = inv_scale.reciprocal().float() + + grads_by_device = OrderedDict() + for group in optimizer.param_groups: + for param in group["params"]: + grad = param.grad + if grad is None: + continue + # A DTensor's local shard is a writable view of its local storage; + # scaling that view avoids introducing a DTensor collective. + if hasattr(grad, "placements") and hasattr(grad, "to_local"): + grad = grad.to_local() + if grad.is_sparse: + grad = grad._values() + grads_by_device.setdefault(grad.device, []).append(grad) + + inv_scale_by_device = {} + for device, grads in grads_by_device.items(): + device_inv_scale = inv_scale_by_device.get(device) + if device_inv_scale is None: + device_inv_scale = inv_scale.to(device=device, non_blocking=True) + inv_scale_by_device[device] = device_inv_scale + # One foreach launch per parameter device. Mixed floating dtypes are + # supported; each tensor rounds the fp32 reciprocal in its own dtype's + # elementwise multiply, just as GradScaler's regular unscale path does. + torch._foreach_mul_(grads, device_inv_scale) + return True + + +@torch._dynamo.disable +def _amp_dtensor_protocol_preflight(optimizer) -> bool: + """Synchronize FP16-DTensor AMP protocol selection before GradScaler scans. + + ``GradScaler.step`` reads ``_step_supports_amp_scaling`` before it checks + gradients for non-finite values. DTensor dispatch may make that check + collective, so choosing the protocol from rank-local active gradients can + leave one rank in the ordinary path while another enters a DTensor + collective. When an optimizer contains any multi-rank DTensor, non-FSDP1 + FP16 parameter storage makes the native protocol a topology property + instead. Before returning it, compare a fixed gradient-presence vector for + every DTensor parameter on the same mesh so every rank either proceeds with + the same protocol or raises before GradScaler touches any gradient. + + Return ``True`` when a distributed optimizer containing DTensors also has + any non-FSDP1 FP16 parameter storage, even if that parameter is local or + inactive on this step. Every DTensor in the optimizer is included because + GradScaler scans the complete active set after selecting the native + protocol. Single-rank and DTensor-free optimizers retain the active-gradient + protocol and do not pay a collective here. + """ + if not torch.distributed.is_available() or not torch.distributed.is_initialized(): + return False + if torch.cuda.is_available() and torch.cuda.is_current_stream_capturing(): + return False + + import torch.distributed as dist + + params = [param for group in optimizer.param_groups for param in group["params"]] + has_fp16_storage = any( + param.dtype == torch.float16 + and not getattr(param, "_is_flat_param", False) + for param in params + ) + if not has_fp16_storage: + return False + + by_mesh = OrderedDict() + for group in optimizer.param_groups: + names = group.get("param_names") + for index, param in enumerate(group["params"]): + if not ( + hasattr(param, "to_local") + and hasattr(param, "placements") + and hasattr(param, "device_mesh") + ): + continue + mesh = param.device_mesh + if mesh.get_coordinate() is None or mesh.size() < 2: + continue + key = ( + str(mesh.device_type), + tuple(int(item) for item in mesh.shape), + tuple( + int(item) + for item in mesh.mesh.detach().cpu().reshape(-1).tolist() + ), + ) + entry = by_mesh.get(key) + if entry is None: + entry = {"mesh": mesh, "items": []} + by_mesh[key] = entry + name = ( + names[index] + if isinstance(names, (list, tuple)) and index < len(names) + else getattr(optimizer, "_param_name", lambda _: "parameter")( + param + ) + ) + entry["items"].append((str(name), param, param.grad is not None)) + + if not by_mesh: + return False + + mismatches = [] + for entry in by_mesh.values(): + mesh = entry["mesh"] + items = sorted( + entry["items"], + key=lambda item: ( + item[0], + tuple(item[1].shape), + str(item[1].dtype), + ), + ) + local = items[0][1].to_local() + if hasattr(local, "wait"): + local = local.wait() + active_counts = torch.tensor( + [int(active) for _, _, active in items], + dtype=torch.int32, + device=local.device, + ) + for process_group in mesh.get_all_groups(): + if dist.get_world_size(process_group) > 1: + dist.all_reduce( + active_counts, op=dist.ReduceOp.SUM, group=process_group + ) + + mesh_size = mesh.size() + inconsistent = torch.nonzero( + (active_counts != 0) & (active_counts != mesh_size), + as_tuple=False, + ).flatten() + if inconsistent.numel() == 0: + continue + counts_cpu = active_counts.cpu() + for index in inconsistent.cpu().tolist(): + mismatches.append( + "{} ({}/{} mesh ranks have gradients)".format( + items[index][0], int(counts_cpu[index]), mesh_size + ) + ) + + if mismatches: + raise RuntimeError( + "Gefen GradScaler integration requires identical DTensor gradient " + "presence on every mesh rank before its distributed non-finite " + "scan. Mismatched parameters: {}. Ensure conditional or unused " + "parameters produce the same `.grad is None` pattern on every " + "mesh rank.".format(", ".join(mismatches)) + ) + return True + + +def _amp_native_scaling_required(optimizer) -> bool: + """Use native scaling only when generic GradScaler cannot unscale safely. + + PyTorch's ordinary GradScaler path has the best integration semantics: it + does not call ``optimizer.step()`` on overflow, so Accelerate/Trainer can + observe the skip and keep schedulers aligned. Its one unsupported case is a + true FP16 gradient tensor. Opt into the optimizer-side protocol only for + that case. + + FSDP1 FlatParameter FP16 needs ``ShardedGradScaler``'s scaler-owned, + cross-rank non-finite reduction, so native handling is disabled for that + representation. DTensor/FSDP2 is different: torch's DTensor handler for the + AMP non-finite op MAX-reduces the scaler-owned flag over every mesh + dimension even in the native check path, so all ranks skip and update their + scales identically. + """ + has_fp16_grad = _amp_dtensor_protocol_preflight(optimizer) + has_fsdp1_flat_param = False + for group in optimizer.param_groups: + for param in group["params"]: + if getattr(param, "_is_flat_param", False): + has_fsdp1_flat_param = True + grad = param.grad + if grad is not None and grad.dtype == torch.float16: + has_fp16_grad = True + + if has_fp16_grad and has_fsdp1_flat_param: + if not getattr(optimizer, "_gefen_warned_sharded_fp16_amp", False): + warnings.warn( + "Gefen detected FP16 gradients on an FSDP1 FlatParameter and " + "disabled optimizer-side AMP unscaling. Use " + "torch.distributed.fsdp.ShardedGradScaler (or explicitly call " + "its unscale_ before step) so non-finite detection is reduced " + "across ranks. Base GradScaler will fail fast while unscaling " + "FP16 gradients instead of risking rank-divergent optimizer " + "steps.", + RuntimeWarning, + stacklevel=3, + ) + optimizer._gefen_warned_sharded_fp16_amp = True + return False + return has_fp16_grad + + +def _assert_optimizer_gradients_structurally_valid( + optimizer, *, require_2d_params: bool = False +) -> None: + """Validate every active gradient before any optimizer mutation. + + Optimizer implementations normally discover an invalid later gradient only + after earlier parameters have already stepped. Gefen also learns/refreshes a + shared codebook before its per-parameter loop, and native AMP may unscale + gradients in place. Scan the complete parameter set first so sparse, + compressed-sparse, MKLDNN, complex, or malformed gradients fail atomically. + """ + for group in optimizer.param_groups: + names = group.get("param_names") + for index, param in enumerate(group["params"]): + name = ( + names[index] + if isinstance(names, (list, tuple)) and index < len(names) + else getattr(optimizer, "_param_name", lambda _: "parameter")(param) + ) + if require_2d_params and param.ndim != 2: + error_factory = getattr( + optimizer, "_step_non_2d_parameter_error", None + ) + if callable(error_factory): + raise ValueError(error_factory(param)) + raise ValueError( + "GefenMuon requires every parameter to remain a 2D matrix " + "before step, but {!r} has {} dimensions.".format( + str(name), param.ndim + ) + ) + grad = param.grad + if grad is None: + continue + layout = getattr(grad, "layout", None) + if getattr(grad, "is_sparse", False) or layout != torch.strided: + raise RuntimeError( + "Gefen does not support sparse gradients or other " + "non-strided layouts; parameter {!r} has gradient layout " + "{}.".format( + str(name), layout + ) + ) + if torch.is_complex(grad): + raise RuntimeError( + "Gefen optimizers do not support complex gradients, but " + "parameter {!r} has dtype {}.".format(str(name), grad.dtype) + ) + if tuple(grad.shape) != tuple(param.shape): + raise RuntimeError( + "Gefen gradient shape {} for parameter {!r} does not match " + "parameter shape {}.".format( + tuple(grad.shape), str(name), tuple(param.shape) + ) + ) + + class Gefen(torch.optim.Optimizer): """Adam-family optimizer with 8-bit quantized momentum and block-shared or factored second moments, at roughly 1 byte of optimizer state per parameter. @@ -720,7 +1033,22 @@ def __init__( weight_decay=weight_decay, ) self._param_names = {} + # ``set_optimizer_state_dict(flatten_optimizer_state_dict=True)`` uses + # the *live* optimizer state/group keys as its unflattening schema before + # it calls our loader. Publish the private rank-local transport keys only + # after the base constructor has finished registering every group. + self._gefen_checkpoint_schema_ready = False super().__init__(self._normalize_param_groups(params), defaults) + self._gefen_checkpoint_schema_ready = True + self._install_rank_local_checkpoint_schema() + + @property + def _step_supports_amp_scaling(self) -> bool: + # Keep standard GradScaler/Accelerate skip semantics for ordinary + # FP32-master AMP. Native handling is needed for active local FP16 + # gradients, or statically for distributed optimizers that combine any + # true-FP16 storage with DTensors after a collective presence preflight. + return _amp_native_scaling_required(self) @staticmethod def _normalize_param_groups(params): @@ -934,6 +1262,8 @@ def add_param_group(self, param_group): self._param_names[param] = param_name self.state[param]["name"] = param_name self._ensure_gefen_global_step_devices() + if getattr(self, "_gefen_checkpoint_schema_ready", False): + self._install_rank_local_checkpoint_schema() def _lr_scalar(self, group) -> float: """Resolve ``group["lr"]`` to a python float, caching a tensor lr's ``.item()``. @@ -2397,19 +2727,29 @@ def _refresh_codebook_with_requant(self) -> None: ) if new_codebook is None: return + staged_indices = [] for pgroup in self.param_groups: for p in pgroup["params"]: pstate = self.state.get(p) if not pstate or "m_codebook" not in pstate: continue stored_device = pstate["m_codebook"].device - old_codebook_local = self._gefen_codebook_on(stored_device) + # Do not populate/mutate the per-device cache until the whole + # refresh transaction succeeds. A failure while staging a later + # parameter must leave both indices and cache identity untouched. + old_codebook_local = old_codebook.to(stored_device) new_codebook_local = new_codebook.to(stored_device) coeffs = gefen_dequantize_unpacked_indices( old_codebook_local, pstate["m_codebook"], old_codebook_local ) indices = gefen_nearest_codebook_indices(new_codebook_local, coeffs) - gefen_set_unpacked_indices(pstate["m_codebook"], indices) + staged_indices.append((pstate["m_codebook"], indices)) + + # Commit only after every allocation/dequantization/search succeeded. + # The staged tensors already have the destination shape/device, so this + # tail is a deterministic sequence of infallible in-place copies. + for stored_indices, indices in staged_indices: + gefen_set_unpacked_indices(stored_indices, indices) self._gefen_codebook = new_codebook self._gefen_codebook_by_device.clear() self._gefen_codebook_lut_by_device.clear() @@ -3344,6 +3684,549 @@ def _step_automatic_merged(self, items) -> None: ) def state_dict(self): + """Run optimizer state-dict hooks around Gefen's complete schema.""" + + for pre_hook in self._optimizer_state_dict_pre_hooks.values(): + pre_hook(self) + state_dict = self._state_dict_impl() + for post_hook in self._optimizer_state_dict_post_hooks.values(): + hook_result = post_hook(self, state_dict) + if hook_result is not None: + state_dict = hook_result + return state_dict + + def _base_state_dict_without_hooks(self): + """Call ``Optimizer.state_dict`` without double-firing public hooks.""" + + pre_hooks = self._optimizer_state_dict_pre_hooks + post_hooks = self._optimizer_state_dict_post_hooks + self._optimizer_state_dict_pre_hooks = OrderedDict() + self._optimizer_state_dict_post_hooks = OrderedDict() + try: + return super().state_dict() + finally: + self._optimizer_state_dict_pre_hooks = pre_hooks + self._optimizer_state_dict_post_hooks = post_hooks + + def _uses_rank_local_sharded_state(self) -> bool: + """Whether this topology stores optimizer state in rank-local geometry.""" + + for group in self.param_groups: + mode = group.get("sharded_mode") + if mode not in (None, "approx"): + continue + for p in group["params"]: + if ( + hasattr(p, "to_local") + and hasattr(p, "placements") + and hasattr(p, "device_mesh") + ): + return True + return False + + @staticmethod + def _is_dtensor_parameter(param) -> bool: + return ( + hasattr(param, "to_local") + and hasattr(param, "placements") + and hasattr(param, "device_mesh") + ) + + @staticmethod + def _placement_checkpoint_signature(placement): + """Encode a DTensor placement structurally rather than through repr().""" + + encoded = { + "type": type(placement).__name__, + "module": type(placement).__module__, + } + for attribute in ("dim", "split_factor"): + value = getattr(placement, attribute, None) + if value is not None: + encoded[attribute] = int(value) + reduce_op = getattr(placement, "reduce_op", None) + if reduce_op is not None: + encoded["reduce_op"] = str(reduce_op) + return encoded + + @staticmethod + def _device_mesh_checkpoint_signature(mesh): + mesh_ranks = mesh.mesh.detach().cpu().reshape(-1).tolist() + return { + "device_type": str(mesh.device_type), + "shape": list(mesh.shape), + "dim_names": list(mesh.mesh_dim_names or ()), + "ranks": [int(item) for item in mesh_ranks], + } + + def _rank_local_checkpoint_context(self, *, fail: bool = True): + """Return the deliberately narrow process-group contract for the adapter. + + Rank-local full-state encoding currently supports exactly one 1-D + DeviceMesh spanning the default process group. Subgroups, pipeline-local + optimizers, and multidimensional meshes need a native sharded DCP format; + silently routing them through the global collectives below can deadlock. + """ + + if not self._uses_rank_local_sharded_state(): + return None + + def reject(message): + if fail: + raise RuntimeError(message) + return None + + if not torch.distributed.is_available() or not torch.distributed.is_initialized(): + return reject( + "Gefen rank-local DTensor checkpointing requires an initialized " + "default distributed process group" + ) + + import torch.distributed as dist + + dtensor_params = [ + p + for group in self.param_groups + for p in group["params"] + if self._is_dtensor_parameter(p) + ] + if not dtensor_params: + return None + + first_mesh = dtensor_params[0].device_mesh + first_mesh_signature = self._device_mesh_checkpoint_signature(first_mesh) + if len(first_mesh_signature["shape"]) != 1: + return reject( + "Gefen rank-local optimizer checkpoints support only one 1-D " + "DeviceMesh spanning the default process group" + ) + for param in dtensor_params[1:]: + mesh_signature = self._device_mesh_checkpoint_signature(param.device_mesh) + if mesh_signature != first_mesh_signature: + return reject( + "Gefen rank-local optimizer checkpoints require every DTensor " + "parameter to use the same 1-D DeviceMesh" + ) + + world = dist.get_world_size() + world_ranks = list(range(world)) + if ( + first_mesh_signature["shape"] != [world] + or sorted(first_mesh_signature["ranks"]) != world_ranks + ): + return reject( + "Gefen rank-local optimizer checkpoints do not support DeviceMesh " + "subgroups; the one-dimensional mesh must span the default world" + ) + + try: + mesh_group = first_mesh.get_group() + mesh_group_ranks = [ + dist.get_global_rank(mesh_group, group_rank) + for group_rank in range(dist.get_world_size(mesh_group)) + ] + except Exception as exc: + if fail: + raise RuntimeError( + "Gefen could not resolve the rank-local DeviceMesh process group" + ) from exc + return None + if sorted(mesh_group_ranks) != world_ranks: + return reject( + "Gefen rank-local optimizer checkpoints require the DeviceMesh " + "process group to contain exactly the default-world ranks" + ) + + coordinate = first_mesh.get_coordinate() + if coordinate is None: + return reject( + "The current rank is not a member of Gefen's checkpoint DeviceMesh" + ) + return { + "world_size": world, + "global_rank": dist.get_rank(), + "group_rank": dist.get_group_rank(mesh_group, dist.get_rank()), + "world_ranks": world_ranks, + "mesh": first_mesh_signature, + "coordinate": [int(item) for item in coordinate], + } + + def _install_rank_local_checkpoint_schema(self) -> None: + """Expose private live keys required by PyTorch's flat OSD unflattener.""" + + for group in self.param_groups: + group.pop("_gefen_checkpoint_metadata", None) + for param in group["params"]: + pstate = self.state.get(param) + if pstate is not None: + for key in tuple(pstate): + if key.startswith(_RANK_LOCAL_PAYLOAD_KEY_PREFIX): + pstate.pop(key) + pstate.pop(_RANK_LOCAL_MEMBER_KEY, None) + + context = self._rank_local_checkpoint_context(fail=False) + if context is None: + return + params = [param for group in self.param_groups for param in group["params"]] + if not params: + return + for global_rank in context["world_ranks"]: + self.state[params[0]][_rank_local_payload_key(global_rank)] = torch.zeros( + 1, dtype=torch.uint8 + ) + for param in params[1:]: + self.state[param][_RANK_LOCAL_MEMBER_KEY] = True + placeholder = { + "format_version": _RANK_LOCAL_METADATA_VERSION, + "global_step": self._gefen_global_step, + "codebook": None, + "deterministic": self._deterministic, + "device_anchor": self._checkpoint_device_anchor(), + "rank_local_sharded_state": {"format": _RANK_LOCAL_FORMAT}, + } + for group in self.param_groups: + group["_gefen_checkpoint_metadata"] = placeholder + + def _rank_local_sharded_signature(self, context=None): + if context is None: + context = self._rank_local_checkpoint_context() + signature = [] + for group_index, group in enumerate(self.param_groups): + for param_index, (name, p) in enumerate( + self._iter_group_params_with_names(group) + ): + identifier = "group_{}_param_{}_{}".format( + group_index, param_index, str(name).lower() + ) + if not self._is_dtensor_parameter(p): + signature.append( + { + "identifier": identifier, + "group": group_index, + "param": param_index, + "name": str(name).lower(), + "sharded": False, + "shape": list(p.shape), + "local_shape": list(p.shape), + "dtype": str(p.dtype), + "requires_grad": bool(p.requires_grad), + "sharded_mode": group.get("sharded_mode"), + } + ) + continue + local = p.to_local() + if hasattr(local, "wait"): + local = local.wait() + mesh = p.device_mesh + coordinate = mesh.get_coordinate() + signature.append( + { + "identifier": identifier, + "group": group_index, + "param": param_index, + "name": str(name).lower(), + "sharded": True, + "shape": list(p.shape), + "local_shape": list(local.shape), + "dtype": str(p.dtype), + "local_dtype": str(local.dtype), + "requires_grad": bool(p.requires_grad), + "placements": [ + self._placement_checkpoint_signature(item) + for item in p.placements + ], + "mesh": self._device_mesh_checkpoint_signature(mesh), + "coordinate": None + if coordinate is None + else list(coordinate), + "sharded_mode": group.get("sharded_mode"), + } + ) + return signature + + @staticmethod + def _rank_local_parameter_manifest(signature): + keys = ( + "identifier", + "group", + "param", + "name", + "sharded", + "shape", + "dtype", + "requires_grad", + "sharded_mode", + ) + return [{key: item.get(key) for key in keys} for item in signature] + + def _checkpoint_device_anchor(self) -> torch.Tensor: + """One byte that lets DCP locate a fresh lazy optimizer's device.""" + + if self._uses_rank_local_sharded_state(): + # The collective rank-indexed adapter deliberately stages every + # payload on CPU to avoid multiplying live VRAM by world size at + # checkpoint time. DCP requires every tensor in the fresh local + # schema to name one device, so its anchor must be CPU as well. + return torch.zeros(1, dtype=torch.uint8) + for group in self.param_groups: + for p in group["params"]: + local = p.to_local() if hasattr(p, "to_local") else p + if hasattr(local, "wait"): + local = local.wait() + return torch.zeros(1, dtype=torch.uint8, device=local.device) + return torch.zeros(1, dtype=torch.uint8) + + @classmethod + def _pack_checkpoint_payload(cls, value, tensors): + if torch.is_tensor(value): + tensor = value.to_local() if hasattr(value, "to_local") else value + if hasattr(tensor, "wait"): + tensor = tensor.wait() + tensor = tensor.detach().contiguous() + index = len(tensors) + tensors.append(tensor) + return ("tensor", index) + if isinstance(value, dict): + return ( + "dict", + [ + (key, cls._pack_checkpoint_payload(item, tensors)) + for key, item in value.items() + ], + ) + if isinstance(value, list): + return ( + "list", + [cls._pack_checkpoint_payload(item, tensors) for item in value], + ) + if isinstance(value, tuple): + return ( + "tuple", + [cls._pack_checkpoint_payload(item, tensors) for item in value], + ) + return ("object", value) + + @classmethod + def _unpack_checkpoint_payload(cls, packed, tensors): + kind, value = packed + if kind == "tensor": + return tensors[value] + if kind == "dict": + return { + key: cls._unpack_checkpoint_payload(item, tensors) + for key, item in value + } + if kind == "list": + return [cls._unpack_checkpoint_payload(item, tensors) for item in value] + if kind == "tuple": + return tuple( + cls._unpack_checkpoint_payload(item, tensors) for item in value + ) + if kind == "object": + return value + raise ValueError("Unknown Gefen checkpoint payload kind: {}".format(kind)) + + @staticmethod + def _serialize_rank_local_payload(payload) -> torch.Tensor: + buffer = io.BytesIO() + torch.save(payload, buffer) + return torch.frombuffer( + bytearray(buffer.getvalue()), dtype=torch.uint8 + ).clone() + + @staticmethod + def _deserialize_rank_local_payload(payload: torch.Tensor): + if not torch.is_tensor(payload) or payload.dtype != torch.uint8 or payload.dim() != 1: + raise ValueError( + "Gefen rank-local checkpoint payload must be a 1-D uint8 tensor" + ) + buffer = io.BytesIO(payload.detach().cpu().contiguous().numpy().tobytes()) + try: + return torch.load(buffer, map_location="cpu", weights_only=True) + except Exception as exc: + raise ValueError("Gefen rank-local checkpoint payload is invalid") from exc + + def _broadcast_checkpoint_payload(self, payload, src, group=None): + import torch.distributed as dist + + rank = dist.get_rank(group) + tensors = [] + if rank == src: + packed = self._pack_checkpoint_payload(payload, tensors) + specs = [(tuple(tensor.shape), tensor.dtype) for tensor in tensors] + header = [(packed, specs)] + else: + header = [None] + dist.broadcast_object_list(header, src=src, group=group) + packed, specs = header[0] + + backend = str(dist.get_backend(group)).lower() + comm_device = ( + torch.device("cuda", torch.cuda.current_device()) + if "nccl" in backend + else torch.device("cpu") + ) + received = [] + for index, (shape, dtype) in enumerate(specs): + if rank == src: + tensor = tensors[index].to(device=comm_device).contiguous() + else: + tensor = torch.empty(shape, dtype=dtype, device=comm_device) + dist.broadcast(tensor, src=src, group=group) + received.append(tensor.detach().cpu().clone()) + return self._unpack_checkpoint_payload(packed, received) + + def _consolidate_rank_local_sharded_state( + self, state_dict, checkpoint_metadata + ) -> None: + """Collect rank-local DTensor state for exact same-topology resume. + + Plain Gefen and ``GefenMuon(sharded_mode='approx')`` learn rank-local + codebooks and block geometry. PyTorch DCP treats their ordinary local + tensors as replicated and otherwise keeps rank 0 silently. Make every + rank return the same rank-indexed payload so generic DCP preserves all + shards; loading selects the payload for the current rank and rejects a + topology/world-size change. + """ + + if not self._uses_rank_local_sharded_state(): + return + + import torch.distributed as dist + + context = self._rank_local_checkpoint_context() + world = context["world_size"] + rank = context["global_rank"] + saved_ids = list(state_dict["state"]) + local_signature = self._rank_local_sharded_signature(context) + local_manifest = self._rank_local_parameter_manifest(local_signature) + local_control = { + "saved_ids": saved_ids, + "manifest": local_manifest, + "signature": local_signature, + "global_step": self._gefen_global_step, + "deterministic": self._deterministic, + } + controls = [None] * world + dist.all_gather_object(controls, local_control) + if any(control["saved_ids"] != saved_ids for control in controls): + raise RuntimeError( + "Gefen DTensor checkpoint requires identical optimizer parameter " + "ordering on every rank; got state keys {}".format( + [control["saved_ids"] for control in controls] + ) + ) + if any(control["manifest"] != local_manifest for control in controls): + raise RuntimeError( + "Gefen DTensor checkpoint parameter identifiers, names, dtypes, " + "or global shapes differ across ranks" + ) + global_steps = [control["global_step"] for control in controls] + if any(type(step) is not int for step in global_steps) or len( + set(global_steps) + ) != 1: + raise RuntimeError( + "Gefen rank-local checkpoint global_step differs across ranks: " + "{}".format(global_steps) + ) + deterministic_values = [control["deterministic"] for control in controls] + if any(type(value) is not bool for value in deterministic_values) or len( + set(deterministic_values) + ) != 1: + raise RuntimeError( + "Gefen rank-local checkpoint deterministic policy differs across " + "ranks: {}".format(deterministic_values) + ) + + signatures = { + str(global_rank): controls[global_rank]["signature"] + for global_rank in context["world_ranks"] + } + local_payload = { + "format": _RANK_LOCAL_FORMAT, + "global_rank": rank, + "group_rank": context["group_rank"], + "world_size": world, + "world_ranks": context["world_ranks"], + "mesh": context["mesh"], + "signature": local_signature, + "parameter_manifest": local_manifest, + "global_step": global_steps[0], + "deterministic": deterministic_values[0], + "states": [state_dict["state"][saved_id] for saved_id in saved_ids], + "codebook": self._gefen_codebook, + } + # Serialize locally before transfer. The prior tensor-by-tensor gather + # retained every rank's decoded state and then a second serialized copy + # on every process. Exchanging one opaque byte tensor per rank keeps CPU + # peak near one global serialized optimizer state plus this rank's live + # local state. Share serialization status first so a deterministic local + # schema/serialization failure makes every rank raise before broadcasts. + try: + local_serialized = self._serialize_rank_local_payload(local_payload) + serialization_status = {"ok": True, "error": None} + except Exception as exc: + local_serialized = None + serialization_status = { + "ok": False, + "error": "{}: {}".format(type(exc).__name__, exc), + } + serialization_statuses = [None] * world + dist.all_gather_object(serialization_statuses, serialization_status) + if any(not item.get("ok", False) for item in serialization_statuses): + raise RuntimeError( + "Gefen rank-local checkpoint payload serialization failed across " + "ranks: {}".format(serialization_statuses) + ) + serialized_payloads = { + str(src): self._broadcast_checkpoint_payload( + local_serialized if rank == src else None, src + ) + for src in context["world_ranks"] + } + state_dict["state"] = {} + carrier_state = { + # PyTorch 2.5's flat optimizer-state loader requires every key + # present in the fresh live state to have a direct tensor/int/float + # leaf in the flattened checkpoint. The real stable name remains + # inside the validated payload; this integer is only an outer + # transport-schema sentinel. + "name": 0, + } + carrier_state.update( + { + _rank_local_payload_key(global_rank): serialized_payloads[ + str(global_rank) + ] + for global_rank in context["world_ranks"] + } + ) + for index, saved_id in enumerate(saved_ids): + state_dict["state"][saved_id] = ( + dict(carrier_state) + if index == 0 + else {"name": 0, _RANK_LOCAL_MEMBER_KEY: True} + ) + marker = { + "format": _RANK_LOCAL_FORMAT, + "world_size": world, + "world_ranks": context["world_ranks"], + "mesh": context["mesh"], + "signatures": signatures, + "parameter_manifest": local_manifest, + "global_step": global_steps[0], + "deterministic": deterministic_values[0], + } + checkpoint_metadata["format_version"] = _RANK_LOCAL_METADATA_VERSION + checkpoint_metadata["global_step"] = global_steps[0] + checkpoint_metadata["deterministic"] = deterministic_values[0] + checkpoint_metadata["codebook"] = None + checkpoint_metadata["rank_local_sharded_state"] = marker + # DCP drops custom top-level keys, and a raw state_dict must not carry + # rank 0's codebook as if it applied to every payload. + state_dict["gefen_codebook"] = None + + def _state_dict_impl(self, *, consolidate_rank_local: bool = True): """Serialize optimizer state, plus Gefen's run-level extras. Beyond the base ``torch.optim.Optimizer.state_dict`` contents (per-param @@ -3380,7 +4263,7 @@ def state_dict(self): original_order = list(self.state) withheld = [(param, self.state.pop(param)) for param in orphaned] try: - state_dict = super().state_dict() + state_dict = self._base_state_dict_without_hooks() finally: # Re-add only the withheld orphans, without clobbering # anything a registered state_dict pre-hook wrote to live @@ -3396,7 +4279,7 @@ def state_dict(self): for key in original_order: self.state[key] = current[key] else: - state_dict = super().state_dict() + state_dict = self._base_state_dict_without_hooks() # stepsize/_h_buf are per-step scratch buffers (recomputed from vmean every # step); they live in self.state only to be reused across steps, so strip # them from the serialized dict instead of bloating every checkpoint. Build @@ -3415,8 +4298,18 @@ def state_dict(self): # falsely mark restored params as batch-covered. "_capt_stack", "_capt_row", + # Live-only schema hints consumed by PyTorch's flattened optimizer + # state unflattener. The collective adapter rebuilds their serialized + # values after ordinary state has been compacted. + _RANK_LOCAL_MEMBER_KEY, ) + def _is_scratch_key(key): + return key in scratch_keys or ( + isinstance(key, str) + and key.startswith(_RANK_LOCAL_PAYLOAD_KEY_PREFIX) + ) + def _compact(value): # K4 rebinds m_magnitude onto a [2, num_blocks, 1] scratch buffer via # set_() (it shares storage with the transient sumsq/stepsize row). @@ -3435,7 +4328,11 @@ def _compact(value): state_dict["state"] = { pid: ( - {k: _compact(v) for k, v in pstate.items() if k not in scratch_keys} + { + k: _compact(v) + for k, v in pstate.items() + if not _is_scratch_key(k) + } if isinstance(pstate, dict) else pstate ) @@ -3464,7 +4361,18 @@ def _compact(value): "global_step": self._gefen_global_step, "codebook": self._gefen_codebook, "deterministic": self._deterministic, + # PyTorch DCP's full-state loader first inspects a freshly-created + # optimizer's local state to choose a broadcast/distribution device. + # Gefen initializes momentum lazily, so before step 1 it otherwise + # contains only parameter names and DCP raises "Expected device to + # be set". This one-byte transport value is removed with the rest + # of the private group metadata during load. + "device_anchor": self._checkpoint_device_anchor(), } + if consolidate_rank_local: + self._consolidate_rank_local_sharded_state( + state_dict, checkpoint_metadata + ) for group in state_dict["param_groups"]: group["_gefen_checkpoint_metadata"] = checkpoint_metadata return state_dict @@ -3532,6 +4440,523 @@ def _pack_legacy_param_groups_for_load(self, state_dict): return migrated def load_state_dict(self, state_dict): + """Run load hooks around Gefen's complete restore transaction.""" + + state_dict = state_dict.copy() + for pre_hook in self._optimizer_load_state_dict_pre_hooks.values(): + hook_result = pre_hook(self, state_dict) + if hook_result is not None: + state_dict = hook_result + self._load_state_dict_impl(state_dict) + for post_hook in self._optimizer_load_state_dict_post_hooks.values(): + post_hook(self) + + def _base_load_state_dict_without_hooks(self, state_dict): + pre_hooks = self._optimizer_load_state_dict_pre_hooks + post_hooks = self._optimizer_load_state_dict_post_hooks + self._optimizer_load_state_dict_pre_hooks = OrderedDict() + self._optimizer_load_state_dict_post_hooks = OrderedDict() + try: + return super().load_state_dict(state_dict) + finally: + self._optimizer_load_state_dict_pre_hooks = pre_hooks + self._optimizer_load_state_dict_post_hooks = post_hooks + + @staticmethod + def _validate_rank_local_codebook(codebook, *, required: bool) -> None: + if codebook is None: + if required: + raise ValueError( + "Gefen rank-local checkpoint has quantized momentum but no codebook" + ) + return + if ( + not torch.is_tensor(codebook) + or codebook.dtype != torch.float32 + or codebook.dim() != 1 + or codebook.numel() != 256 + ): + raise ValueError( + "Gefen rank-local checkpoint codebook must be a 256-element fp32 tensor" + ) + codebook_cpu = codebook.detach().cpu() + if not bool(torch.isfinite(codebook_cpu).all()): + raise ValueError("Gefen rank-local checkpoint codebook must be finite") + if not bool(torch.all(codebook_cpu[1:] >= codebook_cpu[:-1])): + raise ValueError("Gefen rank-local checkpoint codebook must be sorted") + if codebook_cpu[0].item() != -1.0 or codebook_cpu[-1].item() != 1.0: + raise ValueError( + "Gefen rank-local checkpoint codebook must retain endpoints -1 and 1" + ) + + @staticmethod + def _validate_rank_local_counter(name, value) -> float: + if torch.is_tensor(value): + if ( + value.numel() != 1 + or value.dtype == torch.bool + or not bool(torch.isfinite(value.detach()).all()) + ): + raise ValueError( + "Gefen rank-local checkpoint counter {} must be a finite scalar".format( + name + ) + ) + scalar = float(value.detach().cpu().item()) + elif type(value) is int: + scalar = float(value) + else: + raise ValueError( + "Gefen rank-local checkpoint counter {} has invalid type {}".format( + name, type(value).__name__ + ) + ) + if scalar < 0 or not scalar.is_integer(): + raise ValueError( + "Gefen rank-local checkpoint counter {} must be a nonnegative integer".format( + name + ) + ) + return scalar + + def _validate_rank_local_states(self, states, signature, codebook) -> None: + if not isinstance(states, list) or len(states) != len(signature): + raise ValueError( + "Gefen rank-local checkpoint payload has the wrong parameter count" + ) + has_quantized_momentum = False + counter_keys = ("step", "vmean_step", "factored_step", "normuon_step") + for pstate, param_signature in zip(states, signature): + if not isinstance(pstate, dict): + raise ValueError( + "Gefen rank-local checkpoint parameter state must be a dict" + ) + expected_name = param_signature["name"] + if pstate.get("name") != expected_name: + raise ValueError( + "Gefen rank-local checkpoint state name/order differs: " + "checkpoint={!r} expected={!r}".format( + pstate.get("name"), expected_name + ) + ) + counters = {} + for key in counter_keys: + if key in pstate: + counters[key] = self._validate_rank_local_counter( + key, pstate[key] + ) + + local_shape = param_signature["local_shape"] + state_shape = ( + param_signature["shape"] + if param_signature.get("sharded_mode") in ("exact", "distributed") + else local_shape + ) + state_numel = math.prod(state_shape) + period = pstate.get("automatic_period") + if period is not None: + if type(period) is not int or period <= 0: + raise ValueError( + "Gefen rank-local checkpoint automatic_period must be a positive int" + ) + if state_numel == 0 or state_numel % period != 0: + raise ValueError( + "Gefen rank-local checkpoint automatic_period does not divide " + "the parameter state geometry" + ) + + momentum_keys = ("m_codebook", "m_magnitude") + carries_momentum = any(key in pstate for key in momentum_keys) + if carries_momentum: + has_quantized_momentum = True + if not all(key in pstate for key in momentum_keys) or period is None: + raise ValueError( + "Gefen rank-local checkpoint momentum state is incomplete" + ) + if "step" not in pstate: + raise ValueError( + "Gefen rank-local checkpoint momentum state is missing step" + ) + if counters["step"] < 1: + raise ValueError( + "Gefen rank-local checkpoint initialized momentum requires step >= 1" + ) + blocks = state_numel // period + indices = pstate["m_codebook"] + magnitude = pstate["m_magnitude"] + if ( + not torch.is_tensor(indices) + or indices.dtype != torch.uint8 + or tuple(indices.shape) != (blocks, period) + ): + raise ValueError( + "Gefen rank-local checkpoint m_codebook geometry/dtype is invalid" + ) + if ( + not torch.is_tensor(magnitude) + or magnitude.dtype != torch.float32 + or tuple(magnitude.shape) != (blocks, 1) + or not bool(torch.isfinite(magnitude).all()) + or not bool((magnitude >= 0).all()) + ): + raise ValueError( + "Gefen rank-local checkpoint m_magnitude geometry/dtype/values " + "are invalid" + ) + vmean = pstate.get("vmean") + if vmean is not None and ( + not torch.is_tensor(vmean) + or vmean.dtype != torch.float32 + or tuple(vmean.shape) != (blocks, 1) + or not bool(torch.isfinite(vmean).all()) + or not bool((vmean >= 0).all()) + ): + raise ValueError( + "Gefen rank-local checkpoint vmean geometry/dtype/values " + "are invalid" + ) + + # GefenMuon groups carry a sharded_mode and intentionally use + # quantized momentum without Adam's second moment. Plain Gefen + # must carry one complete second-moment representation. + if param_signature.get("sharded_mode") is None: + carries_factored = "v_row" in pstate or "v_col" in pstate + if carries_factored: + if "factored_step" not in pstate: + raise ValueError( + "Gefen rank-local checkpoint factored state is missing " + "factored_step" + ) + if counters["factored_step"] < 1: + raise ValueError( + "Gefen rank-local checkpoint initialized factored " + "state requires factored_step >= 1" + ) + elif "vmean" not in pstate or "vmean_step" not in pstate: + raise ValueError( + "Gefen rank-local checkpoint plain momentum is missing " + "vmean/vmean_step" + ) + elif counters["vmean_step"] < 1: + raise ValueError( + "Gefen rank-local checkpoint initialized vmean requires " + "vmean_step >= 1" + ) + + factored = (pstate.get("v_row"), pstate.get("v_col")) + if (factored[0] is None) != (factored[1] is None): + raise ValueError( + "Gefen rank-local checkpoint factored second moment is incomplete" + ) + if factored[0] is not None: + global_shape = param_signature["shape"] + if len(global_shape) != 2: + raise ValueError( + "Gefen rank-local checkpoint factored state requires a 2-D parameter" + ) + for key, tensor, expected in ( + ("v_row", factored[0], global_shape[0]), + ("v_col", factored[1], global_shape[1]), + ): + if ( + not torch.is_tensor(tensor) + or tensor.dtype != torch.float32 + or tuple(tensor.shape) != (expected,) + or not bool(torch.isfinite(tensor).all()) + or not bool((tensor >= 0).all()) + ): + raise ValueError( + "Gefen rank-local checkpoint {} geometry/dtype/values " + "are invalid".format(key) + ) + + has_normuon_v = "normuon_v" in pstate + has_normuon_step = "normuon_step" in pstate + if has_normuon_v != has_normuon_step: + raise ValueError( + "Gefen rank-local checkpoint NorMuon state is incomplete" + ) + if has_normuon_v: + if param_signature.get("sharded_mode") is None: + raise ValueError( + "Gefen rank-local checkpoint NorMuon state is invalid for " + "a plain Gefen parameter" + ) + if not carries_momentum: + raise ValueError( + "Gefen rank-local checkpoint NorMuon state is missing " + "initialized momentum" + ) + if len(state_shape) != 2: + raise ValueError( + "Gefen rank-local checkpoint NorMuon state requires a 2-D " + "parameter" + ) + normuon_v = pstate["normuon_v"] + if ( + not torch.is_tensor(normuon_v) + or normuon_v.dtype != torch.float32 + or tuple(normuon_v.shape) != (state_shape[0], 1) + or not bool(torch.isfinite(normuon_v).all()) + or not bool((normuon_v >= 0).all()) + ): + raise ValueError( + "Gefen rank-local checkpoint normuon_v geometry/dtype/values " + "are invalid" + ) + if counters["normuon_step"] < 1: + raise ValueError( + "Gefen rank-local checkpoint initialized NorMuon state " + "requires normuon_step >= 1" + ) + + self._validate_rank_local_codebook( + codebook, required=has_quantized_momentum + ) + + def _unwrap_rank_local_sharded_checkpoint(self, state_dict) -> None: + groups = state_dict.get("param_groups", ()) + if not isinstance(groups, (list, tuple)): + raise ValueError("Gefen checkpoint param_groups must be a sequence") + metadata = [] + for group in groups: + if not isinstance(group, dict): + raise ValueError("Gefen checkpoint parameter groups must be dicts") + metadata.append(group.get("_gefen_checkpoint_metadata")) + markers = [ + item.get("rank_local_sharded_state") + for item in metadata + if isinstance(item, dict) + and item.get("rank_local_sharded_state") is not None + ] + wrapped_state = any( + isinstance(pstate, dict) + and ( + any( + isinstance(key, str) + and key.startswith(_RANK_LOCAL_PAYLOAD_KEY_PREFIX) + for key in pstate + ) + or _RANK_LOCAL_MEMBER_KEY in pstate + ) + for pstate in (state_dict.get("state", {}) or {}).values() + ) + if not markers: + if wrapped_state: + raise ValueError( + "Gefen checkpoint carries rank-local DTensor payloads but " + "is missing their topology metadata" + ) + has_quantized_momentum = any( + isinstance(pstate, dict) and "m_codebook" in pstate + for pstate in (state_dict.get("state", {}) or {}).values() + ) + if self._uses_rank_local_sharded_state() and has_quantized_momentum: + raise ValueError( + "Cannot safely load an untagged/full Gefen checkpoint into " + "rank-local DTensor optimizer state. Older generic FSDP/DCP " + "full optimizer checkpoints kept rank 0's codebook and block " + "state for every rank. Re-save from the original run with a " + "Gefen version that writes {} payloads.".format( + _RANK_LOCAL_FORMAT + ) + ) + return + if len(markers) != len(groups) or any(item is None for item in metadata): + raise ValueError( + "Gefen rank-local DTensor checkpoint metadata is present on only " + "some parameter groups" + ) + marker = markers[0] + if not isinstance(marker, dict): + raise ValueError("Gefen rank-local checkpoint marker must be a dict") + if any(not isinstance(item, dict) or item != marker for item in markers[1:]): + raise ValueError( + "Gefen rank-local checkpoint parameter groups carry inconsistent markers" + ) + if marker.get("format") != _RANK_LOCAL_FORMAT: + raise ValueError( + "Unsupported Gefen rank-local checkpoint format: {!r}".format( + marker.get("format") + ) + ) + if not self._uses_rank_local_sharded_state(): + raise ValueError( + "A rank-local DTensor Gefen checkpoint can only load into the " + "same sharded optimizer topology" + ) + try: + context = self._rank_local_checkpoint_context() + except RuntimeError as exc: + raise ValueError(str(exc)) from exc + world = context["world_size"] + rank = context["global_rank"] + if marker.get("world_size") != world: + raise ValueError( + "Gefen rank-local DTensor checkpoints require the same world " + "size; checkpoint={} current={}".format( + marker.get("world_size"), world + ) + ) + if marker.get("world_ranks") != context["world_ranks"]: + raise ValueError( + "Gefen rank-local checkpoint default-world rank membership differs" + ) + if marker.get("mesh") != context["mesh"]: + raise ValueError("Gefen rank-local checkpoint DeviceMesh differs") + signatures = marker.get("signatures") + expected_rank_keys = {str(item) for item in context["world_ranks"]} + if not isinstance(signatures, dict) or set(signatures) != expected_rank_keys: + raise ValueError("Gefen rank-local checkpoint has invalid signatures") + current_signature = self._rank_local_sharded_signature(context) + if signatures[str(rank)] != current_signature: + raise ValueError( + "Gefen rank-local DTensor checkpoint topology differs on rank {}: " + "checkpoint={!r} current={!r}".format( + rank, signatures[str(rank)], current_signature + ) + ) + current_manifest = self._rank_local_parameter_manifest(current_signature) + if marker.get("parameter_manifest") != current_manifest: + raise ValueError( + "Gefen rank-local checkpoint ordered parameter manifest differs" + ) + marker_step = marker.get("global_step") + if type(marker_step) is not int or marker_step < 0: + raise ValueError("Gefen rank-local checkpoint has invalid global_step") + marker_deterministic = marker.get("deterministic") + if type(marker_deterministic) is not bool: + raise ValueError( + "Gefen rank-local checkpoint has invalid deterministic policy" + ) + saved_ids = list( + chain.from_iterable(group.get("params", ()) for group in groups) + ) + try: + unique_saved_ids = set(saved_ids) + except TypeError as exc: + raise ValueError( + "Gefen rank-local checkpoint parameter IDs must be hashable" + ) from exc + if len(unique_saved_ids) != len(saved_ids): + raise ValueError( + "Gefen rank-local checkpoint contains duplicate outer parameter IDs" + ) + expected_carrier_keys = { + _rank_local_payload_key(item) for item in context["world_ranks"] + } + serialized_payloads = None + for param_id in saved_ids: + pstate = (state_dict.get("state", {}) or {}).get(param_id) + if not isinstance(pstate, dict): + raise ValueError( + "Gefen rank-local checkpoint is missing parameter state {!r}".format( + param_id + ) + ) + carrier_keys = { + key + for key in pstate + if isinstance(key, str) + and key.startswith(_RANK_LOCAL_PAYLOAD_KEY_PREFIX) + } + extras = set(pstate) - carrier_keys - { + _RANK_LOCAL_MEMBER_KEY, + "name", + } + if extras: + raise ValueError( + "Gefen rank-local checkpoint parameter state has an invalid " + "schema: {!r}".format(pstate) + ) + if pstate.get("name") != 0: + raise ValueError( + "Gefen rank-local checkpoint outer name sentinel is invalid" + ) + has_carrier = bool(carrier_keys) + has_member = _RANK_LOCAL_MEMBER_KEY in pstate + if has_carrier and has_member: + raise ValueError( + "Gefen rank-local checkpoint state cannot be both carrier and member" + ) + if has_carrier: + if carrier_keys != expected_carrier_keys: + raise ValueError( + "Gefen rank-local checkpoint carrier has missing or " + "unexpected per-rank payload keys" + ) + if serialized_payloads is not None: + raise ValueError( + "Gefen rank-local checkpoint repeats its payload transport" + ) + serialized_payloads = { + str(global_rank): pstate[ + _rank_local_payload_key(global_rank) + ] + for global_rank in context["world_ranks"] + } + elif pstate.get(_RANK_LOCAL_MEMBER_KEY) is not True: + raise ValueError( + "Gefen rank-local checkpoint has an invalid member marker" + ) + if ( + not isinstance(serialized_payloads, dict) + or set(serialized_payloads) != expected_rank_keys + or any(not torch.is_tensor(item) for item in serialized_payloads.values()) + ): + raise ValueError("Gefen rank-local checkpoint has invalid rank payloads") + selected_payload = self._deserialize_rank_local_payload( + serialized_payloads[str(rank)] + ) + expected_payload_keys = { + "format", + "global_rank", + "group_rank", + "world_size", + "world_ranks", + "mesh", + "signature", + "parameter_manifest", + "global_step", + "deterministic", + "states", + "codebook", + } + if not isinstance(selected_payload, dict) or set(selected_payload) != expected_payload_keys: + raise ValueError("Gefen rank-local checkpoint payload has an invalid schema") + payload_identity = ( + selected_payload["format"] == _RANK_LOCAL_FORMAT + and selected_payload["global_rank"] == rank + and selected_payload["group_rank"] == context["group_rank"] + and selected_payload["world_size"] == world + and selected_payload["world_ranks"] == context["world_ranks"] + and selected_payload["mesh"] == context["mesh"] + and selected_payload["signature"] == current_signature + and selected_payload["parameter_manifest"] == current_manifest + and selected_payload["global_step"] == marker_step + and selected_payload["deterministic"] == marker_deterministic + ) + if not payload_identity: + raise ValueError( + "Gefen rank-local checkpoint payload rank/topology/parameter binding differs" + ) + selected_states = selected_payload["states"] + selected_codebook = selected_payload["codebook"] + self._validate_rank_local_states( + selected_states, current_signature, selected_codebook + ) + state_dict["state"] = dict(zip(saved_ids, selected_states)) + state_dict["gefen_global_step"] = marker_step + state_dict["gefen_codebook"] = selected_codebook + state_dict["gefen_deterministic"] = marker_deterministic + for group in groups: + group_metadata = dict(group["_gefen_checkpoint_metadata"]) + group_metadata["global_step"] = marker_step + group_metadata["codebook"] = selected_codebook + group_metadata["deterministic"] = marker_deterministic + group["_gefen_checkpoint_metadata"] = group_metadata + + def _load_state_dict_impl(self, state_dict): """Restore optimizer state saved by :meth:`state_dict`. On top of the base load this restores ``gefen_global_step`` and the @@ -3555,6 +4980,7 @@ def load_state_dict(self, state_dict): state_dict["param_groups"] = [ dict(group) for group in state_dict.get("param_groups", ()) ] + self._unwrap_rank_local_sharded_checkpoint(state_dict) group_metadata = [] for group in state_dict["param_groups"]: metadata = group.pop("_gefen_checkpoint_metadata", None) @@ -3576,7 +5002,11 @@ def load_state_dict(self, state_dict): "Gefen checkpoint metadata is present on only some parameter groups" ) first_metadata = group_metadata[0] - if first_metadata.get("format_version") != 1: + if first_metadata.get("format_version") not in ( + 1, + 2, + _RANK_LOCAL_METADATA_VERSION, + ): raise ValueError( "Unsupported Gefen checkpoint metadata format_version: {}".format( first_metadata.get("format_version") @@ -3695,7 +5125,7 @@ def load_state_dict(self, state_dict): # and re-aliases lazily). self._capt_invalidate() - super().load_state_dict(state_dict) + self._base_load_state_dict_without_hooks(state_dict) self._gefen_global_step = gefen_global_step # Capturable SR seeds are optimizer-level scratch (a device mirror of # gefen_global_step): drop them so the first post-load SR kernel call @@ -3776,6 +5206,7 @@ def load_state_dict(self, state_dict): elif torch.is_tensor(value): pstate[key] = int(value.item()) self._sync_param_names_to_state() + self._install_rank_local_checkpoint_schema() @torch.no_grad() def step(self, closure=None): @@ -3796,6 +5227,17 @@ def step(self, closure=None): with torch.enable_grad(): loss = closure() + _assert_optimizer_gradients_structurally_valid(self) + + # GradScaler invokes native-AMP optimizers even on overflow. Decide + # before codebook learning, periodic refresh, capturable counters, or + # parameter/state mutation. The attribute gate compiles away on the + # ordinary BF16/FP32 path where GradScaler attaches nothing. + if ( + hasattr(self, "found_inf") or hasattr(self, "grad_scale") + ) and not _amp_prepare_optimizer_step(self): + return loss + self._maybe_refresh_gefen_codebook() self._maybe_save_gefen_grad_histogram() # Periodic codebook re-learn (opt-in): every N steps, refit the exact-DP diff --git a/src/gefen/gefen_muon.py b/src/gefen/gefen_muon.py index 98f9b8a..8bb3b26 100644 --- a/src/gefen/gefen_muon.py +++ b/src/gefen/gefen_muon.py @@ -6,7 +6,11 @@ import torch import torch.nn as nn -from gefen.gefen import Gefen +from gefen.gefen import ( + Gefen, + _amp_prepare_optimizer_step, + _assert_optimizer_gradients_structurally_valid, +) EPS = 1e-7 DEFAULT_A = 3.4445 @@ -85,6 +89,7 @@ def _stable_distributed_owner(stable_index: int, world: int) -> int: BATCHED_NS_MAX_MIN_DIM = 512 BATCHED_NS_MAX_NUMEL = 1 << 20 BATCHED_NS_DEFAULT_WORKSPACE_BYTES = 256 << 20 +_DISTRIBUTED_CHECKPOINT_METADATA_KEY = "muon_distributed_state" def _batched_ns_shape_eligible(rows: int, cols: int) -> bool: @@ -524,6 +529,10 @@ class GefenMuon(Gefen): verbose: print codebook/quantization diagnostics. """ + @staticmethod + def _step_non_2d_parameter_error(param: torch.Tensor) -> str: + return _swapped_param_groups_error(param) + def __init__( self, params: Iterable[Union[nn.Parameter, Tuple[str, nn.Parameter]]], @@ -1286,6 +1295,96 @@ def _dist_available() -> bool: return False return torch.distributed.is_initialized() + @torch._dynamo.disable + def _assert_sharded_grad_presence_consistent(self) -> None: + """Fail before collectives when mesh ranks disagree on ``grad is None``. + + Exact and distributed Muon reconstruct full DTensor gradients with + collectives. If one mesh rank skips a parameter while another enters its + ``full_tensor()``/broadcast path, the latter waits forever. Pack every + relevant parameter's activity bit by DeviceMesh and sum it across each + mesh dimension in sequence; the final count is global to that mesh, so + every participating rank sees and raises on the same disagreement before + codebook learning or optimizer state/parameter mutation begins. + + Plain tensors/DDP and local-shard ``approx`` mode take no Muon gradient + collectives and therefore pay no preflight collective. Manual CUDA graph + capture also skips this host-branching check: capturable Gefen requires + eager warmup before capture, and those warmup steps establish that the + graph's fixed gradient-presence pattern is rank-consistent. + """ + if not self._dist_available(): + return + if torch.cuda.is_available() and torch.cuda.is_current_stream_capturing(): + return + + import torch.distributed as dist + + by_mesh = OrderedDict() + for group in self.param_groups: + if group["sharded_mode"] == "approx": + continue + for name, p in self._iter_group_params_with_names(group): + if not self._is_sharded(p): + continue + mesh = p.device_mesh + if mesh.get_coordinate() is None or mesh.size() < 2: + continue + # All parameters created from one FSDP2/DTensor mesh normally + # share the same DeviceMesh object. Key by identity so exotic + # optimizers carrying distinct meshes run one matched vector + # collective per mesh without conflating their parameter order. + key = id(mesh) + entry = by_mesh.get(key) + if entry is None: + entry = {"mesh": mesh, "items": []} + by_mesh[key] = entry + entry["items"].append((str(name), p, p.grad is not None)) + + mismatches = [] + for entry in by_mesh.values(): + mesh = entry["mesh"] + items = entry["items"] + device = self._state_tensor_device(items[0][1]) + active_counts = torch.tensor( + [int(active) for _, _, active in items], + dtype=torch.int32, + device=device, + ) + # Reducing the same vector once along every Cartesian mesh + # dimension propagates a global activity count to every mesh rank. + # This is one collective per mesh dimension, independent of the + # number of optimizer parameters. + for process_group in mesh.get_all_groups(): + if dist.get_world_size(process_group) > 1: + dist.all_reduce( + active_counts, op=dist.ReduceOp.SUM, group=process_group + ) + + mesh_size = mesh.size() + inconsistent = torch.nonzero( + (active_counts != 0) & (active_counts != mesh_size), + as_tuple=False, + ).flatten() + if inconsistent.numel() == 0: + continue + counts_cpu = active_counts.cpu() + for index in inconsistent.cpu().tolist(): + mismatches.append( + "{} ({}/{} mesh ranks have gradients)".format( + items[index][0], int(counts_cpu[index]), mesh_size + ) + ) + + if mismatches: + raise RuntimeError( + "GefenMuon requires identical gradient presence on every rank " + "of a DTensor/FSDP mesh before sharded_mode='exact' or " + "'distributed' stepping. Mismatched parameters: {}. Ensure " + "conditional or unused parameters produce the same `.grad is " + "None` pattern on every mesh rank.".format(", ".join(mismatches)) + ) + def _distributed_process_group(self, p: torch.Tensor): if not self._dist_available() or not self._is_sharded(p): return None @@ -1645,6 +1744,7 @@ def _consolidate_distributed_state_dict(self, state_dict): if not by_pg: return state_dict + marker_groups = [] for pg, pg_items in by_pg.items(): world = dist.get_world_size(pg) my_coord = dist.get_group_rank(pg, dist.get_rank()) @@ -1708,10 +1808,40 @@ def _consolidate_distributed_state_dict(self, state_dict): state_dict["state"][saved_id] = owner_state + marker_groups.append( + { + "world_size": world, + "params": [ + { + "saved_id": saved_id, + "name": str(name), + "shape": tuple(p.shape), + "owner": owner_by_saved_id[saved_id], + "state_keys": tuple( + sorted( + str(key) + for key in state_dict["state"] + .get(saved_id, {}) + .keys() + ) + ), + "initialized": any( + str(key) != "name" + for key in state_dict["state"] + .get(saved_id, {}) + .keys() + ), + } + for name, p, saved_id in pg_items + ], + } + ) + state_dict["gefen_muon_distributed"] = { - "version": 1, + "version": 2, "ownership": "stable_full_param_index_v1", "consolidated": True, + "groups": marker_groups, } return state_dict @@ -1741,53 +1871,644 @@ def _drop_non_owned_distributed_state(self) -> None: pstate.clear() pstate["name"] = str(name).lower() - def state_dict(self): - state_dict = super().state_dict() - return self._consolidate_distributed_state_dict(state_dict) + def _state_dict_impl(self): + # Parallel-Muon owner state must be made complete before Gefen's + # rank-local DTensor adapter serializes the per-rank payload. This + # ordering matters for mixed optimizers carrying both ``approx`` and + # ``distributed`` groups: wrapping first replaces every real state entry + # with opaque transport tensors, which are not owner momentum. + state_dict = super()._state_dict_impl(consolidate_rank_local=False) + state_dict = self._consolidate_distributed_state_dict(state_dict) + + checkpoint_metadata = None + groups = state_dict.get("param_groups", ()) + if groups: + checkpoint_metadata = groups[0].get("_gefen_checkpoint_metadata") + if self._uses_rank_local_sharded_state(): + if not isinstance(checkpoint_metadata, dict): + raise RuntimeError( + "GefenMuon rank-local checkpoint metadata was not initialized" + ) + self._consolidate_rank_local_sharded_state( + state_dict, checkpoint_metadata + ) + + marker = state_dict.get("gefen_muon_distributed") + if marker is not None: + # Generic DCP keeps only conventional state/param_groups top-level + # keys. Mirror the owner proof into every group's private transport + # metadata so ordinary and flattened DCP round-trips retain it. + for group in groups: + metadata = dict(group.get("_gefen_checkpoint_metadata", {})) + metadata[_DISTRIBUTED_CHECKPOINT_METADATA_KEY] = marker + group["_gefen_checkpoint_metadata"] = metadata + return state_dict + + @staticmethod + def _distributed_checkpoint_error(detail: str) -> ValueError: + return ValueError( + "GefenMuon.load_state_dict refused populated " + "sharded_mode='distributed' optimizer state before mutation: {}. " + "Parallel-Muon momentum must be consolidated by calling " + "GefenMuon.state_dict() on every save rank and must carry either a " + "canonical released version-1 consolidation proof or a valid " + "version-2 'gefen_muon_distributed' saved-world/owner manifest. " + "Re-save from the original distributed job; do not load a rank-local " + "or manually stripped optimizer state.".format(detail) + ) + + def _distributed_checkpoint_marker(self, state_dict): + """Recover one consistent owner proof from top-level/group transport.""" + + marker = state_dict.pop("gefen_muon_distributed", None) + groups = state_dict.get("param_groups", ()) + metadata_markers = [] + metadata_presence = [] + for group in groups if isinstance(groups, (list, tuple)) else (): + metadata = ( + group.get("_gefen_checkpoint_metadata") + if isinstance(group, dict) + else None + ) + present = ( + isinstance(metadata, dict) + and _DISTRIBUTED_CHECKPOINT_METADATA_KEY in metadata + ) + metadata_presence.append(present) + if present: + metadata_markers.append( + metadata[_DISTRIBUTED_CHECKPOINT_METADATA_KEY] + ) + + if any(metadata_presence) and not all(metadata_presence): + raise self._distributed_checkpoint_error( + "the parameter-group copies of the owner manifest are incomplete" + ) + if metadata_markers: + metadata_marker = metadata_markers[0] + try: + copies_agree = all( + item == metadata_marker for item in metadata_markers[1:] + ) + except Exception: + copies_agree = False + if not copies_agree: + raise self._distributed_checkpoint_error( + "the parameter-group copies of the owner manifest disagree" + ) + if marker is None: + marker = metadata_marker + else: + try: + top_level_agrees = marker == metadata_marker + except Exception: + top_level_agrees = False + if not top_level_agrees: + raise self._distributed_checkpoint_error( + "the top-level and parameter-group owner manifests disagree" + ) + return marker + + def _distributed_load_state_items(self, state_dict): + """Bind saved distributed groups to live params without zip truncation.""" + saved_groups = state_dict.get("param_groups") + if not isinstance(saved_groups, (list, tuple)): + raise self._distributed_checkpoint_error( + "the checkpoint param_groups payload is not a sequence" + ) + if not any( + isinstance(group, dict) and group.get("sharded_mode") == "distributed" + for group in saved_groups + ): + return [], OrderedDict() + if len(saved_groups) != len(self.param_groups): + raise self._distributed_checkpoint_error( + "the checkpoint has {} parameter groups but the live optimizer " + "has {}".format(len(saved_groups), len(self.param_groups)) + ) + + saved_state = state_dict.get("state", {}) + expected_items = [] + seen_saved_ids = set() + by_pg = OrderedDict() + for group_index, (saved_group, live_group) in enumerate( + zip(saved_groups, self.param_groups) + ): + if not isinstance(saved_group, dict): + raise self._distributed_checkpoint_error( + "checkpoint parameter group {} is not a mapping".format( + group_index + ) + ) + saved_ids = saved_group.get("params") + if not isinstance(saved_ids, (list, tuple)): + raise self._distributed_checkpoint_error( + "checkpoint parameter group {} has no parameter-id list".format( + group_index + ) + ) + live_items = list(self._iter_group_params_with_names(live_group)) + if len(saved_ids) != len(live_items): + raise self._distributed_checkpoint_error( + "checkpoint parameter group {} has {} parameters but the live " + "group has {}".format( + group_index, len(saved_ids), len(live_items) + ) + ) + if saved_group.get("sharded_mode") != "distributed": + continue + + saved_names = saved_group.get("param_names") + if saved_names is not None and ( + not isinstance(saved_names, (list, tuple)) + or len(saved_names) != len(saved_ids) + ): + raise self._distributed_checkpoint_error( + "checkpoint distributed group {} has an invalid param_names " + "manifest".format(group_index) + ) + for param_index, ((live_name, live_param), saved_id) in enumerate( + zip(live_items, saved_ids) + ): + if saved_id in seen_saved_ids: + raise self._distributed_checkpoint_error( + "checkpoint distributed parameter id {!r} appears more " + "than once".format(saved_id) + ) + seen_saved_ids.add(saved_id) + state_name = None + if isinstance(saved_state, dict): + pstate = saved_state.get(saved_id) + if isinstance(pstate, dict): + state_name = pstate.get("name") + saved_name = ( + saved_names[param_index] + if saved_names is not None + else state_name + ) + if saved_name is not None and str(saved_name).lower() != str( + live_name + ).lower(): + raise self._distributed_checkpoint_error( + "checkpoint distributed group {} parameter {} name {!r} " + "does not match live name {!r}".format( + group_index, param_index, saved_name, str(live_name) + ) + ) + item = (str(live_name).lower(), live_param, saved_id) + expected_items.append(item) + pg = self._distributed_process_group(live_param) + if pg is not None: + by_pg.setdefault(pg, []).append(item) + return expected_items, by_pg + + @staticmethod + def _distributed_checkpoint_is_pristine(state_dict, expected_items) -> bool: + steps = [] + if "gefen_global_step" in state_dict: + steps.append(state_dict.get("gefen_global_step")) + codebooks = [state_dict.get("gefen_codebook")] + for group in state_dict.get("param_groups", ()): + if not isinstance(group, dict): + return False + metadata = group.get("_gefen_checkpoint_metadata") + if isinstance(metadata, dict): + if "global_step" in metadata: + steps.append(metadata.get("global_step")) + codebooks.append(metadata.get("codebook")) + if not steps or any(type(step) is not int or step != 0 for step in steps): + return False + if any(codebook is not None for codebook in codebooks): + return False + + saved_state = state_dict.get("state", {}) + if not isinstance(saved_state, dict): + return False + for expected_name, _, saved_id in expected_items: + pstate = saved_state.get(saved_id) + if ( + not isinstance(pstate, dict) + or set(pstate) != {"name"} + or str(pstate.get("name")).lower() != expected_name + ): + return False + return True + + def _distributed_checkpoint_codebook(self, state_dict): + candidates = [] + if "gefen_codebook" in state_dict: + candidates.append(state_dict.get("gefen_codebook")) + for group in state_dict.get("param_groups", ()): + if not isinstance(group, dict): + continue + metadata = group.get("_gefen_checkpoint_metadata") + if isinstance(metadata, dict) and "codebook" in metadata: + candidates.append(metadata.get("codebook")) + non_null = [value for value in candidates if value is not None] + if not non_null: + return None + codebook = non_null[0] + for other in non_null[1:]: + if not ( + torch.is_tensor(codebook) + and torch.is_tensor(other) + and torch.equal(codebook, other) + ): + raise self._distributed_checkpoint_error( + "checkpoint copies of the frozen codebook disagree" + ) + return codebook + + def _validate_distributed_codebook(self, state_dict, *, required: bool) -> None: + try: + self._validate_rank_local_codebook( + self._distributed_checkpoint_codebook(state_dict), + required=required, + ) + except ValueError as exc: + raise self._distributed_checkpoint_error(str(exc)) from exc + + @staticmethod + def _distributed_counter_is_valid(value, *, minimum: int = 1) -> bool: + if torch.is_tensor(value): + if ( + value.dim() != 0 + or value.dtype == torch.bool + or not bool(torch.isfinite(value).all()) + ): + return False + scalar = float(value.detach().cpu().item()) + return scalar >= minimum and scalar.is_integer() + return type(value) is int and value >= minimum + + def _validate_distributed_param_state( + self, pstate, expected_name: str, param: torch.Tensor + ) -> bool: + """Validate owner momentum independently of the consolidation marker.""" + if not isinstance(pstate, dict): + raise self._distributed_checkpoint_error( + "state for parameter {!r} is missing or is not a mapping".format( + expected_name + ) + ) + if not all(isinstance(key, str) for key in pstate): + raise self._distributed_checkpoint_error( + "state for parameter {!r} has non-string keys".format(expected_name) + ) + if str(pstate.get("name")).lower() != expected_name: + raise self._distributed_checkpoint_error( + "state name {!r} does not match parameter {!r}".format( + pstate.get("name"), expected_name + ) + ) + if set(pstate) == {"name"}: + return False + + core = {"name", "automatic_period", "step", "m_codebook", "m_magnitude"} + missing = core - set(pstate) + if missing: + raise self._distributed_checkpoint_error( + "state for parameter {!r} is initialized but missing core keys " + "{}".format(expected_name, sorted(missing)) + ) + period = pstate["automatic_period"] + if type(period) is not int or period <= 0 or param.numel() % period != 0: + raise self._distributed_checkpoint_error( + "parameter {!r} has invalid automatic_period {!r} for {} " + "elements".format(expected_name, period, param.numel()) + ) + blocks = param.numel() // period + indices = pstate["m_codebook"] + magnitude = pstate["m_magnitude"] + if ( + not torch.is_tensor(indices) + or indices.dtype != torch.uint8 + or tuple(indices.shape) != (blocks, period) + ): + raise self._distributed_checkpoint_error( + "parameter {!r} has invalid m_codebook dtype/geometry".format( + expected_name + ) + ) + if ( + not torch.is_tensor(magnitude) + or magnitude.dtype != torch.float32 + or tuple(magnitude.shape) != (blocks, 1) + or not bool(torch.isfinite(magnitude).all()) + or not bool((magnitude >= 0).all()) + ): + raise self._distributed_checkpoint_error( + "parameter {!r} has invalid m_magnitude dtype/geometry/values".format( + expected_name + ) + ) + if not self._distributed_counter_is_valid(pstate["step"]): + raise self._distributed_checkpoint_error( + "parameter {!r} has invalid step counter".format(expected_name) + ) + normuon_v = pstate.get("normuon_v") + normuon_step = pstate.get("normuon_step") + if (normuon_v is None) != (normuon_step is None): + raise self._distributed_checkpoint_error( + "parameter {!r} has incomplete NorMuon state".format(expected_name) + ) + if normuon_v is not None and ( + not torch.is_tensor(normuon_v) + or normuon_v.dtype != torch.float32 + or tuple(normuon_v.shape) != (param.shape[0], 1) + or not bool(torch.isfinite(normuon_v).all()) + or not bool((normuon_v >= 0).all()) + or not self._distributed_counter_is_valid(normuon_step) + ): + raise self._distributed_checkpoint_error( + "parameter {!r} has invalid NorMuon dtype/geometry/counter".format( + expected_name + ) + ) + return True - def _warn_if_unconsolidated_distributed_load(self, state_dict, marker) -> None: + def _validate_distributed_checkpoint_load(self, state_dict, marker) -> None: # In distributed mode the persistent momentum lives only on each matrix's # stable owner rank; state_dict() runs a collective that broadcasts every - # owner's state to all ranks and stamps the consolidation marker. A - # checkpoint that skipped that consolidation (an old/foreign save, or a - # rank-0-only dump) leaves non-owner ranks without the owner's momentum, so - # a resumed run silently diverges. Warn (do NOT raise) so loading such a - # checkpoint degrades rather than crashes. - if not self._dist_available(): - return - consolidated = isinstance(marker, dict) and marker.get("consolidated") is True - if consolidated: - return + # owner's state to all ranks and stamps a manifest describing the saved + # process-group world, stable owner, parameter identity/order, and state + # keys. A rank-local dump can otherwise look loadable on every rank while + # silently resetting the non-owner momentum. Validate before delegating to + # the base loader so every rejection is mutation-free. try: - by_pg = self._distributed_state_items(state_dict) - except (KeyError, TypeError): - return - if not by_pg: + expected_items, by_pg = self._distributed_load_state_items(state_dict) + except (KeyError, TypeError, ValueError) as exc: + if isinstance(exc, ValueError) and str(exc).startswith( + "GefenMuon.load_state_dict refused" + ): + raise + raise self._distributed_checkpoint_error( + "the checkpoint parameter-group layout cannot be matched to the " + "live distributed DTensor parameters ({})".format(exc) + ) from exc + if not expected_items: return + saved_state = state_dict.get("state", {}) - carries_state = any( - saved_id in saved_state - for pg_items in by_pg.values() - for (_, _, saved_id) in pg_items + if not isinstance(saved_state, dict): + raise self._distributed_checkpoint_error( + "the checkpoint 'state' payload is not a mapping" + ) + pristine = self._distributed_checkpoint_is_pristine( + state_dict, expected_items ) - if not carries_state: + if pristine and marker is None: return - warnings.warn( - "GefenMuon.load_state_dict: loading a sharded_mode='distributed' " - "checkpoint that is not marked consolidated (missing/incomplete " - "'gefen_muon_distributed' marker written by GefenMuon.state_dict()). " - "Non-owner ranks may start with incomplete momentum and the resumed " - "run can diverge. Re-save with GefenMuon.state_dict() (called on every " - "rank) to produce a consolidated checkpoint.", - RuntimeWarning, - stacklevel=2, - ) - def load_state_dict(self, state_dict): + if not isinstance(marker, dict): + # A world-1 or unsupported-mesh "distributed" group takes Muon's + # replicated exact fallback and never had owner-local state to + # consolidate. Preserve such markerless legacy loads only when every + # expected entry is independently proven complete. + if not by_pg: + initialized = [ + self._validate_distributed_param_state( + saved_state.get(saved_id), expected_name, param + ) + for expected_name, param, saved_id in expected_items + ] + if initialized and all(initialized): + self._validate_distributed_codebook( + state_dict, required=True + ) + return + raise self._distributed_checkpoint_error( + "the consolidation marker is missing or is not a mapping" + ) + version = marker.get("version") + if version not in (1, 2): + raise self._distributed_checkpoint_error( + "the consolidation marker version is {!r}, expected 1 or 2".format( + version + ) + ) + if marker.get("ownership") != "stable_full_param_index_v1": + raise self._distributed_checkpoint_error( + "the ownership scheme is {!r}, expected " + "'stable_full_param_index_v1'".format(marker.get("ownership")) + ) + if marker.get("consolidated") is not True: + raise self._distributed_checkpoint_error( + "the consolidation marker does not assert consolidated=True" + ) + + if version == 1: + if pristine: + return + initialized = [ + self._validate_distributed_param_state( + saved_state.get(saved_id), expected_name, param + ) + for expected_name, param, saved_id in expected_items + ] + if not initialized or not all(initialized): + raise self._distributed_checkpoint_error( + "released version-1 state mixes initialized and empty owner " + "entries, so completeness cannot be proven" + ) + self._validate_distributed_codebook(state_dict, required=True) + return + + marker_groups = marker.get("groups") + if not isinstance(marker_groups, (list, tuple)): + raise self._distributed_checkpoint_error( + "the saved-world/owner group manifest is missing or is not a list" + ) + + # A sharded_mode='distributed' optimizer group may contain a mixture of + # Parallel-Muon-eligible parameters and replicated fallbacks (plain + # tensors, multi-dimensional meshes, or world-one meshes). Save-side + # consolidation creates one owner manifest per eligible process group + # and deliberately leaves fallback state to the ordinary Gefen loader. + # Bind the proof to those same ordered process-group partitions instead + # of requiring it to cover every parameter in the optimizer group. + expected_pg_groups = list(by_pg.values()) + if len(marker_groups) != len(expected_pg_groups): + raise self._distributed_checkpoint_error( + "the owner manifest has {} process-group entries but the live " + "optimizer has {} eligible distributed process groups".format( + len(marker_groups), len(expected_pg_groups) + ) + ) + + eligible_ids = { + saved_id + for pg_items in expected_pg_groups + for _, _, saved_id in pg_items + } + fallback_items = [ + item for item in expected_items if item[2] not in eligible_ids + ] + initialized_any = False + for expected_name, param, saved_id in fallback_items: + initialized_any = ( + self._validate_distributed_param_state( + saved_state.get(saved_id), expected_name, param + ) + or initialized_any + ) + + for group_index, (saved_group, expected_pg_items) in enumerate( + zip(marker_groups, expected_pg_groups) + ): + if not isinstance(saved_group, dict): + raise self._distributed_checkpoint_error( + "owner manifest group {} is not a mapping".format(group_index) + ) + saved_world = saved_group.get("world_size") + if ( + not isinstance(saved_world, int) + or isinstance(saved_world, bool) + or saved_world <= 0 + ): + raise self._distributed_checkpoint_error( + "owner manifest group {} has invalid saved world_size {!r}".format( + group_index, saved_world + ) + ) + saved_params = saved_group.get("params") + if not isinstance(saved_params, (list, tuple)): + raise self._distributed_checkpoint_error( + "owner manifest group {} has no ordered parameter list".format( + group_index + ) + ) + expected_ids = [saved_id for _, _, saved_id in expected_pg_items] + manifest_ids = [] + for param_index, saved_param in enumerate(saved_params): + if not isinstance(saved_param, dict): + raise self._distributed_checkpoint_error( + "owner manifest group {} parameter {} is not a mapping".format( + group_index, param_index + ) + ) + expected_owner = _stable_distributed_owner(param_index, saved_world) + saved_owner = saved_param.get("owner") + if ( + not isinstance(saved_owner, int) + or isinstance(saved_owner, bool) + or saved_owner != expected_owner + ): + raise self._distributed_checkpoint_error( + "owner manifest group {} parameter {} has owner {!r}, " + "expected {} for saved world {}".format( + group_index, + param_index, + saved_owner, + expected_owner, + saved_world, + ) + ) + manifest_ids.append(saved_param.get("saved_id")) + + try: + manifest_id_set = set(manifest_ids) + except TypeError as exc: + raise self._distributed_checkpoint_error( + "the owner manifest contains an unhashable saved_id" + ) from exc + if ( + len(manifest_ids) != len(expected_ids) + or len(manifest_id_set) != len(manifest_ids) + or manifest_ids != expected_ids + ): + raise self._distributed_checkpoint_error( + "owner manifest group {} parameter ids {!r} do not exactly " + "match the ordered eligible checkpoint/live ids {!r}".format( + group_index, manifest_ids, expected_ids + ) + ) + + for param_index, (saved_param, live_item) in enumerate( + zip(saved_params, expected_pg_items) + ): + live_name, live_param, saved_id = live_item + if saved_param.get("name") != str(live_name): + raise self._distributed_checkpoint_error( + "owner manifest group {} parameter {} name {!r} does not " + "match live name {!r}".format( + group_index, + param_index, + saved_param.get("name"), + str(live_name), + ) + ) + saved_shape = saved_param.get("shape") + if not isinstance(saved_shape, (list, tuple)) or tuple( + saved_shape + ) != tuple(live_param.shape): + raise self._distributed_checkpoint_error( + "owner manifest group {} parameter {} shape {!r} does not " + "match live shape {!r}".format( + group_index, + param_index, + saved_shape, + tuple(live_param.shape), + ) + ) + saved_keys = saved_param.get("state_keys") + if ( + not isinstance(saved_keys, (list, tuple)) + or not all(isinstance(key, str) for key in saved_keys) + or len(set(saved_keys)) != len(saved_keys) + ): + raise self._distributed_checkpoint_error( + "owner manifest group {} parameter {} has no state_keys " + "manifest".format(group_index, param_index) + ) + actual_param_state = saved_state.get(saved_id) + initialized = self._validate_distributed_param_state( + actual_param_state, str(live_name).lower(), live_param + ) + initialized_any = initialized_any or initialized + actual_keys = tuple(sorted(str(key) for key in actual_param_state)) + if tuple(sorted(saved_keys)) != actual_keys: + raise self._distributed_checkpoint_error( + "owner manifest group {} parameter {} state keys {!r} do " + "not match checkpoint state keys {!r}".format( + group_index, + param_index, + tuple(saved_keys), + actual_keys, + ) + ) + marker_initialized = saved_param.get("initialized") + if ( + type(marker_initialized) is not bool + or marker_initialized != initialized + ): + raise self._distributed_checkpoint_error( + "owner manifest group {} parameter {} initialized={!r} " + "does not match checkpoint state initialized={}".format( + group_index, + param_index, + marker_initialized, + initialized, + ) + ) + self._validate_distributed_codebook(state_dict, required=initialized_any) + + def _load_state_dict_impl(self, state_dict): state_dict = dict(state_dict) - marker = state_dict.pop("gefen_muon_distributed", None) - self._warn_if_unconsolidated_distributed_load(state_dict, marker) - super().load_state_dict(state_dict) + state_dict = self._pack_legacy_param_groups_for_load(state_dict) + marker = self._distributed_checkpoint_marker(state_dict) + + # Unwrap only a copied transaction for Parallel-Muon validation. The + # base loader unwraps and loads the original exactly once after every + # owner/fallback invariant has passed, keeping all rejection paths + # mutation-free while composing with rank-local approx state. + validation_state = dict(state_dict) + validation_state["param_groups"] = [ + dict(group) for group in state_dict.get("param_groups", ()) + ] + self._unwrap_rank_local_sharded_checkpoint(validation_state) + self._validate_distributed_checkpoint_load(validation_state, marker) + super()._load_state_dict_impl(state_dict) # Old checkpoints predate shape-batched NS. Keep their historical, # bit-identical serial behavior even when this optimizer was constructed # with the new opt-in; current checkpoints carry and restore both keys. @@ -1800,6 +2521,7 @@ def load_state_dict(self, state_dict): BATCHED_NS_DEFAULT_WORKSPACE_BYTES, ) self._drop_non_owned_distributed_state() + self._install_rank_local_checkpoint_schema() @torch.no_grad() def step(self, closure=None): @@ -1809,6 +2531,19 @@ def step(self, closure=None): with torch.enable_grad(): loss = closure() + _assert_optimizer_gradients_structurally_valid( + self, require_2d_params=True + ) + + # 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 + # Partition the work once so distributed-mode sharded params can take the # stable-owner Parallel-Muon path while every other param keeps the normal # per-param path. @@ -1834,6 +2569,14 @@ def step(self, closure=None): elif grad is not None: regular_items.append((group, name, p, grad)) + # This must precede _maybe_refresh_gefen_codebook(): the first-step + # Muon codebook iterator itself calls full_tensor() in exact/distributed + # modes, before either update dispatcher gets a chance to validate the + # active set. The preflight is mutation-free and gives every mesh rank + # the same clear error instead of leaving active ranks in an unmatched + # collective. + self._assert_sharded_grad_presence_consistent() + self._maybe_refresh_gefen_codebook() self._maybe_save_gefen_grad_histogram() diff --git a/src/gefen/hybrid.py b/src/gefen/hybrid.py index 60aa1b3..8837e18 100644 --- a/src/gefen/hybrid.py +++ b/src/gefen/hybrid.py @@ -60,7 +60,12 @@ import torch import torch.nn as nn -from gefen.gefen import Gefen +from gefen.gefen import ( + Gefen, + _amp_native_scaling_required, + _amp_prepare_optimizer_step, + _assert_optimizer_gradients_structurally_valid, +) from gefen.gefen_muon import GefenMuon from gefen.params import ( DEFAULT_BACKUP_SUBSTRINGS, @@ -580,6 +585,17 @@ def param_groups(self): groups.extend(o.param_groups) return groups + @property + def _step_supports_amp_scaling(self) -> bool: + # Inspect the union of both halves so GradScaler chooses one protocol + # and one overflow decision for the whole composite. FP32-master AMP + # stays on the ordinary externally-skipped path; true FP16 uses the + # native path because generic unscale rejects FP16 tensors. A hybrid + # combining any true-FP16 storage with multi-rank DTensors selects that + # protocol statically after a union-wide collective presence preflight + # (FSDP1 FlatParameters require ShardedGradScaler). + return _amp_native_scaling_required(self) + @param_groups.setter def param_groups(self, value): # A composite cannot accept a wholesale param_groups assignment: the @@ -670,6 +686,20 @@ def step(self, closure=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 + ) + # 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): + 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() diff --git a/tests/test_amp_grad_scaler.py b/tests/test_amp_grad_scaler.py new file mode 100644 index 0000000..e86ff4e --- /dev/null +++ b/tests/test_amp_grad_scaler.py @@ -0,0 +1,608 @@ +"""Native and ordinary GradScaler protocols for Gefen optimizers.""" + +import copy +from datetime import timedelta +import os +import queue +import socket +import subprocess +import sys +import traceback + +import pytest +import torch +import torch.nn as nn + +from gefen import Gefen, GefenMuon, GefenMuonHybrid + + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available(), reason="GradScaler integration requires CUDA" +) + + +class _Fp16ManualGradScaler(torch.amp.GradScaler): + """Test-only equivalent of ShardedGradScaler's FP16-capable unscale.""" + + def _unscale_grads_(self, optimizer, inv_scale, found_inf, allow_fp16): + return super()._unscale_grads_(optimizer, inv_scale, found_inf, True) + + +def _initial_values(dtype): + generator = torch.Generator(device="cpu").manual_seed(20260712) + matrix = (torch.randn(32, 32, generator=generator) * 0.02).to( + device="cuda", dtype=dtype + ) + bias = (torch.randn(32, generator=generator) * 0.02).to( + device="cuda", dtype=dtype + ) + return matrix, bias + + +def _grad_values(dtype): + generator = torch.Generator(device="cpu").manual_seed(20260713) + matrix = (torch.randn(32, 32, generator=generator) * 0.002).to( + device="cuda", dtype=dtype + ) + bias = (torch.randn(32, generator=generator) * 0.002).to( + device="cuda", dtype=dtype + ) + return matrix, bias + + +def _make_optimizer(kind, dtype, fused, *, backup_optimizer="gefen"): + matrix_init, bias_init = _initial_values(dtype) + matrix = nn.Parameter(matrix_init.clone()) + if kind == "gefen": + optimizer = Gefen( + [("layer.weight", matrix)], + lr=2e-3, + fused=fused, + factored_v_2d=False, + ) + return optimizer, [matrix] + if kind == "muon": + optimizer = GefenMuon( + [("layer.weight", matrix)], + lr=2e-3, + fused=fused, + ns_steps=2, + ns_schedule="standard", + adjust_lr_fn="original", + ) + return optimizer, [matrix] + if kind == "hybrid": + bias = nn.Parameter(bias_init.clone()) + optimizer = GefenMuonHybrid( + [("layer.weight", matrix)], + [("layer.bias", bias)], + lr=2e-3, + fused=fused, + ns_steps=2, + ns_schedule="standard", + normuon=False, + backup_optimizer=backup_optimizer, + ) + return optimizer, [matrix, bias] + raise AssertionError("unknown optimizer kind: {}".format(kind)) + + +def _new_scaler(scale=128.0): + scaler = torch.amp.GradScaler("cuda", init_scale=scale) + # GradScaler initializes its device scale lazily on scale(loss). + scaler.scale(torch.ones((), device="cuda")) + return scaler + + +def _new_fp16_manual_scaler(scale=128.0): + scaler = _Fp16ManualGradScaler("cuda", init_scale=scale) + scaler.scale(torch.ones((), device="cuda")) + return scaler + + +def _free_port(): + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.bind(("127.0.0.1", 0)) + return str(sock.getsockname()[1]) + + +def _assign_grads(params, grads): + for param, grad in zip(params, grads): + param.grad = grad.clone() + + +def _assign_scaled_grads(params, grads, scale): + for param, grad in zip(params, grads): + param.grad = (grad * scale).to(dtype=param.dtype) + + +def _clone_nested(value): + if torch.is_tensor(value): + return value.detach().clone() + if isinstance(value, dict): + return {key: _clone_nested(item) for key, item in value.items()} + if isinstance(value, list): + return [_clone_nested(item) for item in value] + if isinstance(value, tuple): + return tuple(_clone_nested(item) for item in value) + return copy.deepcopy(value) + + +def _assert_nested_equal(actual, expected): + if torch.is_tensor(expected): + assert torch.is_tensor(actual) + assert actual.dtype == expected.dtype + assert actual.device == expected.device + assert torch.equal(actual, expected) + return + if isinstance(expected, dict): + assert isinstance(actual, dict) + assert set(actual) == set(expected) + for key in expected: + _assert_nested_equal(actual[key], expected[key]) + return + if isinstance(expected, (list, tuple)): + assert isinstance(actual, type(expected)) + assert len(actual) == len(expected) + for actual_item, expected_item in zip(actual, expected): + _assert_nested_equal(actual_item, expected_item) + return + assert actual == expected + + +def _snapshot(optimizer, params): + return { + "params": [param.detach().clone() for param in params], + "state_dict": _clone_nested(optimizer.state_dict()), + } + + +def _assert_snapshot_equal(optimizer, params, snapshot): + for param, expected in zip(params, snapshot["params"]): + assert torch.equal(param.detach(), expected) + _assert_nested_equal(optimizer.state_dict(), snapshot["state_dict"]) + + +def _clone_local_state(state): + cloned = {} + for key, value in state.items(): + if torch.is_tensor(value): + if hasattr(value, "to_local"): + value = value.to_local() + cloned[key] = value.detach().clone() + else: + cloned[key] = copy.deepcopy(value) + return cloned + + +def _local_state_equal(actual, expected): + if set(actual) != set(expected): + return False + for key, expected_value in expected.items(): + actual_value = actual[key] + if torch.is_tensor(actual_value) and hasattr(actual_value, "to_local"): + actual_value = actual_value.to_local() + if torch.is_tensor(expected_value): + if not torch.is_tensor(actual_value) or not torch.equal( + actual_value, expected_value + ): + return False + elif actual_value != expected_value: + return False + return True + + +def _case_grads(kind, dtype): + matrix_grad, bias_grad = _grad_values(dtype) + return [matrix_grad, bias_grad] if kind == "hybrid" else [matrix_grad] + + +def _dtensor_amp_overflow_worker(rank, world, port, result_queue): + import torch.distributed as dist + from torch.distributed.tensor import Shard, distribute_tensor, init_device_mesh + + try: + os.environ["MASTER_ADDR"] = "127.0.0.1" + os.environ["MASTER_PORT"] = port + os.environ["RANK"] = str(rank) + os.environ["WORLD_SIZE"] = str(world) + torch.cuda.set_device(rank) + dist.init_process_group( + "nccl", + rank=rank, + world_size=world, + timeout=timedelta(seconds=15), + ) + mesh = init_device_mesh("cuda", (world,)) + results = [] + + for mode in ("exact", "distributed"): + for phase in ("first_step", "initialized"): + generator = torch.Generator(device="cpu").manual_seed( + 5000 + len(results) + ) + full_init = (torch.randn(16, 16, generator=generator) * 0.02).to( + rank, torch.float16 + ) + full_grad = ( + torch.randn(16, 16, generator=generator) * 0.002 + ).to(rank, torch.float16) + param = nn.Parameter( + distribute_tensor(full_init.clone(), mesh, [Shard(0)]) + ) + optimizer = GefenMuon( + [("weight", param)], + lr=2e-3, + fused=False, + ns_steps=2, + ns_schedule="standard", + sharded_mode=mode, + ) + scaler = _new_scaler(scale=8.0) + + if phase == "initialized": + finite_grad = distribute_tensor( + (full_grad * scaler.get_scale()).clone(), mesh, [Shard(0)] + ) + param.grad = finite_grad + scaler.step(optimizer) + scaler.update() + optimizer.zero_grad(set_to_none=True) + + param_before = param.detach().to_local().clone() + state_before = _clone_local_state(optimizer.state[param]) + codebook_before = ( + None + if optimizer._gefen_codebook is None + else optimizer._gefen_codebook.detach().clone() + ) + global_step_before = optimizer._gefen_global_step + scale_before = scaler.get_scale() + + overflow_grad = distribute_tensor( + (full_grad * scale_before).clone(), mesh, [Shard(0)] + ) + if rank == 0: + overflow_grad.to_local().view(-1)[0] = float("inf") + param.grad = overflow_grad + native_protocol = optimizer._step_supports_amp_scaling + scaler.step(optimizer) + scaler.update() + + codebook_after = optimizer._gefen_codebook + codebook_unchanged = ( + codebook_before is None and codebook_after is None + ) or ( + codebook_before is not None + and codebook_after is not None + and torch.equal(codebook_before, codebook_after) + ) + results.append( + { + "mode": mode, + "phase": phase, + "native_protocol": native_protocol, + "param_unchanged": torch.equal( + param.detach().to_local(), param_before + ), + "state_unchanged": _local_state_equal( + optimizer.state[param], state_before + ), + "codebook_unchanged": codebook_unchanged, + "global_step_unchanged": optimizer._gefen_global_step + == global_step_before, + "scale_before": scale_before, + "scale_after": scaler.get_scale(), + } + ) + param.grad = None + dist.barrier(device_ids=[rank]) + + result_queue.put(("result", rank, results)) + except Exception: + result_queue.put(("error", rank, traceback.format_exc())) + finally: + if dist.is_initialized(): + dist.destroy_process_group() + + +@pytest.mark.parametrize("fused", [False, True], ids=["unfused", "fused"]) +@pytest.mark.parametrize("kind", ["gefen", "muon", "hybrid"]) +def test_finite_true_fp16_scaler_step_matches_explicit_unscaled(kind, fused): + reference, reference_params = _make_optimizer(kind, torch.float16, fused) + scaled, scaled_params = _make_optimizer(kind, torch.float16, fused) + grads = _case_grads(kind, torch.float16) + + _assign_grads(reference_params, grads) + reference.step() + + scaler = _new_scaler() + _assign_scaled_grads(scaled_params, grads, scaler.get_scale()) + assert scaled._step_supports_amp_scaling + scaler.step(scaled) + scaler.update() + + for scaled_grad, unscaled_grad in zip( + (param.grad for param in scaled_params), grads + ): + assert torch.equal(scaled_grad, unscaled_grad) + for actual, expected in zip(scaled_params, reference_params): + assert torch.equal(actual.detach(), expected.detach()) + _assert_nested_equal(scaled.state_dict(), reference.state_dict()) + + +@pytest.mark.parametrize("fused", [False, True], ids=["unfused", "fused"]) +@pytest.mark.parametrize("kind", ["gefen", "muon", "hybrid"]) +@pytest.mark.parametrize("initialized", [False, True], ids=["first_step", "initialized"]) +def test_true_fp16_overflow_is_a_complete_noop(kind, fused, initialized): + optimizer, params = _make_optimizer(kind, torch.float16, fused) + scaler = _new_scaler() + grads = _case_grads(kind, torch.float16) + + if initialized: + _assign_scaled_grads(params, grads, scaler.get_scale()) + scaler.step(optimizer) + scaler.update() + optimizer.zero_grad(set_to_none=True) + + snapshot = _snapshot(optimizer, params) + _assign_scaled_grads(params, grads, scaler.get_scale()) + params[0].grad.view(-1)[0] = float("inf") + scale_before = scaler.get_scale() + result = scaler.step(optimizer) + scaler.update() + + assert result is None + assert scaler.get_scale() < scale_before + _assert_snapshot_equal(optimizer, params, snapshot) + + +@pytest.mark.parametrize("fused", [False, True], ids=["unfused", "fused"]) +@pytest.mark.parametrize("kind", ["gefen", "muon", "hybrid"]) +def test_manual_unscale_is_not_applied_twice(kind, fused): + # Explicit GradScaler.unscale_ supports FP32 master gradients and is the + # path used by gradient clipping in Accelerate/Trainer. True FP16 tensors + # require native automatic unscale because base GradScaler rejects them. + reference, reference_params = _make_optimizer(kind, torch.float32, fused) + scaled, scaled_params = _make_optimizer(kind, torch.float32, fused) + grads = _case_grads(kind, torch.float32) + clipped_grads = [grad * 0.25 for grad in grads] + + _assign_grads(reference_params, clipped_grads) + reference.step() + + scaler = _new_scaler() + assert not scaled._step_supports_amp_scaling + _assign_scaled_grads(scaled_params, grads, scaler.get_scale()) + scaler.unscale_(scaled) + for param, grad in zip(scaled_params, grads): + assert torch.equal(param.grad, grad) + param.grad.mul_(0.25) + scaler.step(scaled) + scaler.update() + + for actual, expected in zip(scaled_params, reference_params): + assert torch.equal(actual.detach(), expected.detach()) + _assert_nested_equal(scaled.state_dict(), reference.state_dict()) + + +@pytest.mark.parametrize("fused", [False, True], ids=["unfused", "fused"]) +@pytest.mark.parametrize("kind", ["gefen", "muon", "hybrid"]) +def test_native_fp16_manual_unscale_stage_is_not_divided_twice(kind, fused): + reference, reference_params = _make_optimizer(kind, torch.float16, fused) + scaled, scaled_params = _make_optimizer(kind, torch.float16, fused) + grads = _case_grads(kind, torch.float16) + clipped_grads = [(grad * 0.5).to(torch.float16) for grad in grads] + + _assign_grads(reference_params, clipped_grads) + reference.step() + + scaler = _new_fp16_manual_scaler() + _assign_scaled_grads(scaled_params, grads, scaler.get_scale()) + scaler.unscale_(scaled) + assert scaled._step_supports_amp_scaling + for param, grad in zip(scaled_params, grads): + assert torch.equal(param.grad, grad) + param.grad.mul_(0.5) + scaler.step(scaled) + scaler.update() + + for actual, expected in zip(scaled_params, reference_params): + assert torch.equal(actual.detach(), expected.detach()) + _assert_nested_equal(scaled.state_dict(), reference.state_dict()) + + +@pytest.mark.parametrize("fused", [False, True], ids=["unfused", "fused"]) +@pytest.mark.parametrize("initialized", [False, True], ids=["first_step", "initialized"]) +@pytest.mark.parametrize("overflow_half", ["muon", "backup"]) +def test_hybrid_overflow_in_one_half_skips_both_halves( + fused, initialized, overflow_half +): + optimizer, params = _make_optimizer( + "hybrid", torch.float16, fused, backup_optimizer="adamw" + ) + scaler = _new_scaler() + grads = _case_grads("hybrid", torch.float16) + + if initialized: + _assign_scaled_grads(params, grads, scaler.get_scale()) + scaler.step(optimizer) + scaler.update() + optimizer.zero_grad(set_to_none=True) + + snapshot = _snapshot(optimizer, params) + _assign_scaled_grads(params, grads, scaler.get_scale()) + overflow_index = 0 if overflow_half == "muon" else 1 + params[overflow_index].grad.view(-1)[0] = float("inf") + scaler.step(optimizer) + scaler.update() + + _assert_snapshot_equal(optimizer, params, snapshot) + + +@pytest.mark.parametrize("kind", ["gefen", "muon", "hybrid"]) +def test_bf16_stays_on_the_ordinary_optimizer_path(kind): + if not torch.cuda.is_bf16_supported(): + pytest.skip("GPU does not support BF16") + optimizer, params = _make_optimizer(kind, torch.bfloat16, fused=False) + grads = _case_grads(kind, torch.bfloat16) + assert not optimizer._step_supports_amp_scaling + before = [param.detach().clone() for param in params] + _assign_grads(params, grads) + optimizer.step() + assert any( + not torch.equal(param.detach(), initial) + for param, initial in zip(params, before) + ) + + +def test_accelerate_fp32_master_overflow_reports_step_skipped(): + pytest.importorskip("accelerate") + # Isolate Accelerator's process-global state so this regression composes + # with Trainer/Accelerate tests that choose a different mixed-precision mode. + script = r""" +import torch +from accelerate import Accelerator +from gefen import Gefen + +accelerator = Accelerator(mixed_precision="fp16") +param = torch.nn.Parameter(torch.randn(16, 16, device="cuda")) +optimizer = Gefen([("weight", param)], fused=False) +wrapped = accelerator.prepare_optimizer(optimizer) +wrapped.scaler.scale(torch.ones((), device="cuda")) +param.grad = torch.full_like(param, float("inf")) +before = param.detach().clone() +wrapped.step() +assert wrapped.step_was_skipped +assert torch.equal(param.detach(), before) +assert optimizer._gefen_global_step == 0 +assert optimizer._gefen_codebook is None +""" + result = subprocess.run( + [sys.executable, "-c", script], + cwd=os.getcwd(), + text=True, + capture_output=True, + timeout=30, + ) + assert result.returncode == 0, result.stdout + result.stderr + + +def test_accelerate_bf16_keeps_the_unscaled_ordinary_step(): + pytest.importorskip("accelerate") + script = r""" +import torch +from accelerate import Accelerator +from gefen import Gefen + +accelerator = Accelerator(mixed_precision="bf16") +param = torch.nn.Parameter(torch.randn(16, 16, device="cuda")) +optimizer = Gefen([("weight", param)], fused=False) +wrapped = accelerator.prepare_optimizer(optimizer) +assert wrapped.scaler is None +inputs = torch.randn(8, 16, device="cuda") +with accelerator.autocast(): + loss = (inputs @ param).square().mean() +accelerator.backward(loss) +before = param.detach().clone() +wrapped.step() +assert not wrapped.step_was_skipped +assert not torch.equal(param.detach(), before) +assert optimizer._gefen_global_step == 1 +""" + result = subprocess.run( + [sys.executable, "-c", script], + cwd=os.getcwd(), + text=True, + capture_output=True, + timeout=30, + ) + assert result.returncode == 0, result.stdout + result.stderr + + +def test_fsdp1_flat_fp16_disables_native_protocol_and_fails_before_step(): + optimizer, params = _make_optimizer("muon", torch.float16, fused=False) + param = params[0] + # FSDP1 exposes FlatParameter as a regular local tensor; unlike DTensor it + # has no dispatcher that globally reduces scaler-owned found_inf flags. + param._is_flat_param = True + grad = _case_grads("muon", torch.float16)[0] + scaler = _new_scaler() + _assign_scaled_grads(params, [grad], scaler.get_scale()) + snapshot = _snapshot(optimizer, params) + + with pytest.warns(RuntimeWarning, match="ShardedGradScaler"): + assert not optimizer._step_supports_amp_scaling + with pytest.raises(ValueError, match="Attempting to unscale FP16 gradients"): + scaler.step(optimizer) + _assert_snapshot_equal(optimizer, params, snapshot) + + +@pytest.mark.skipif( + torch.cuda.device_count() < 2 + or not torch.distributed.is_available() + or not torch.distributed.is_nccl_available(), + reason="DTensor AMP overflow regression needs two CUDA GPUs and NCCL", +) +def test_dtensor_rank_local_fp16_overflow_skips_every_rank(): + import torch.multiprocessing as mp + + world = 2 + context = mp.get_context("spawn") + result_queue = context.Queue() + port = _free_port() + processes = [ + context.Process( + target=_dtensor_amp_overflow_worker, + args=(rank, world, port, result_queue), + ) + for rank in range(world) + ] + for process in processes: + process.start() + + messages = [] + try: + for _ in range(world): + messages.append(result_queue.get(timeout=45)) + except queue.Empty: + pass + finally: + for process in processes: + process.join(timeout=5) + for process in processes: + if process.is_alive(): + process.terminate() + process.join(timeout=5) + + assert all(process.exitcode == 0 for process in processes), [ + process.exitcode for process in processes + ] + errors = [payload for kind, _, payload in messages if kind == "error"] + assert not errors, "\n".join(errors) + rank_results = { + rank: payload for kind, rank, payload in messages if kind == "result" + } + assert set(rank_results) == {0, 1}, messages + + expected_cases = { + (mode, phase) + for mode in ("exact", "distributed") + for phase in ("first_step", "initialized") + } + for rank, results in rank_results.items(): + assert {(item["mode"], item["phase"]) for item in results} == expected_cases + for item in results: + assert item["native_protocol"], (rank, item) + assert item["param_unchanged"], (rank, item) + assert item["state_unchanged"], (rank, item) + assert item["codebook_unchanged"], (rank, item) + assert item["global_step_unchanged"], (rank, item) + assert item["scale_after"] == item["scale_before"] / 2, (rank, item) + + # DTensor's scaler-owned MAX reduction must keep every rank's scale in lockstep. + for case_index in range(len(rank_results[0])): + assert ( + rank_results[0][case_index]["scale_after"] + == rank_results[1][case_index]["scale_after"] + ) diff --git a/tests/test_cpu_step_checkpoint.py b/tests/test_cpu_step_checkpoint.py index ae39758..03fdbc7 100644 --- a/tests/test_cpu_step_checkpoint.py +++ b/tests/test_cpu_step_checkpoint.py @@ -841,3 +841,70 @@ def stamp(optimizer): assert opt.state[flat_params[0]]["hook_stamp"] == 123 # Key order is still restored when the hook leaves the key set unchanged. assert list(opt.state.keys()) == live_keys_before + + +def test_gefen_state_dict_post_hooks_see_and_can_replace_final_schema(): + model = _small_model() + opt = Gefen(list(model.named_parameters()), lr=1e-3, fused=False) + _apply_grads(model, _synthetic_grads(model, 1)[0]) + opt.step() + + observed = [] + + def inspect_final(_optimizer, state_dict): + observed.append(copy.deepcopy(state_dict)) + + opt.register_state_dict_post_hook(inspect_final) + opt.register_state_dict_post_hook( + lambda _optimizer, _state_dict: {"custom": "replacement"} + ) + assert opt.state_dict() == {"custom": "replacement"} + assert len(observed) == 1 + assert { + "state", + "param_groups", + "gefen_global_step", + "gefen_codebook", + "gefen_deterministic", + }.issubset(observed[0]) + assert all( + "stepsize" not in pstate and "_h_buf" not in pstate + for pstate in observed[0]["state"].values() + ) + + +def test_gefen_load_hooks_wrap_complete_restore(): + source_model = _small_model() + source = Gefen(list(source_model.named_parameters()), lr=1e-3, fused=False) + for step_grads in _synthetic_grads(source_model, 2): + _apply_grads(source_model, step_grads) + source.step() + source.zero_grad() + saved = copy.deepcopy(source.state_dict()) + + target_model = _small_model() + target = Gefen(list(target_model.named_parameters()), lr=9e-3, fused=False) + seen_pre = [] + seen_post = [] + + def pre_hook(_optimizer, state_dict): + seen_pre.append(set(state_dict)) + + def post_hook(optimizer): + seen_post.append( + ( + optimizer._gefen_global_step, + optimizer._gefen_codebook.detach().clone(), + [state.get("step") for state in optimizer.state.values()], + ) + ) + + target.register_load_state_dict_pre_hook(pre_hook) + target.register_load_state_dict_post_hook(post_hook) + target.load_state_dict(saved) + + assert seen_pre == [set(saved)] + assert len(seen_post) == 1 + assert seen_post[0][0] == 2 + assert torch.equal(seen_post[0][1], saved["gefen_codebook"]) + assert all(not torch.is_tensor(step) for step in seen_post[0][2] if step is not None) diff --git a/tests/test_gefen_fsdp2_checkpoint.py b/tests/test_gefen_fsdp2_checkpoint.py new file mode 100644 index 0000000..2fcbbb0 --- /dev/null +++ b/tests/test_gefen_fsdp2_checkpoint.py @@ -0,0 +1,738 @@ +"""Same-topology FSDP2/DTensor optimizer checkpoint continuation.""" + +from __future__ import annotations + +import copy +import os +import queue +import socket +import traceback + +import pytest +import torch + + +_CARRIER_PREFIX = "_gefen_rank_local_payload_" +_MEMBER = "_gefen_rank_local_member" +_FORMAT = "rank_local_dtensor_v2" + + +def _carrier_key(rank: int) -> str: + return f"{_CARRIER_PREFIX}{rank}" + + +def _free_port() -> str: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.bind(("127.0.0.1", 0)) + return str(sock.getsockname()[1]) + + +def _local_value(value): + if hasattr(value, "to_local"): + value = value.to_local() + if hasattr(value, "wait"): + value = value.wait() + return value + + +def _clone_value(value): + if torch.is_tensor(value): + return _local_value(value).detach().cpu().clone() + if isinstance(value, dict): + return {key: _clone_value(item) for key, item in value.items()} + if isinstance(value, list): + return [_clone_value(item) for item in value] + if isinstance(value, tuple): + return tuple(_clone_value(item) for item in value) + return copy.deepcopy(value) + + +def _values_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 torch.equal(left, right) + if type(left) is not type(right): + return False + if isinstance(left, dict): + return set(left) == set(right) and all( + _values_equal(left[key], right[key]) for key in left + ) + if isinstance(left, (list, tuple)): + return len(left) == len(right) and all( + _values_equal(a, b) for a, b in zip(left, right) + ) + return left == right + + +def _optimizer_snapshot(optimizer): + params = [param for group in optimizer.param_groups for param in group["params"]] + return { + "state": [_clone_value(optimizer.state[param]) for param in params], + "groups": [ + _clone_value( + { + key: value + for key, value in group.items() + if key not in ("params", "_gefen_checkpoint_metadata") + } + ) + for group in optimizer.param_groups + ], + "params": [_clone_value(param) for param in params], + "global_step": optimizer._gefen_global_step, + "codebook": _clone_value(optimizer._gefen_codebook), + "codebook_cache": _clone_value(optimizer._gefen_codebook_by_device), + "seed_cache": _clone_value(optimizer._sr_seed_by_device), + } + + +def _persistent_optimizer_snapshot(optimizer): + scratch = { + "stepsize", + "_h_buf", + "_capt_scalars", + "_capt_consts", + "_capt_consts_key", + "_capt_stack", + "_capt_row", + _MEMBER, + } + params = [param for group in optimizer.param_groups for param in group["params"]] + return { + "state": [ + _clone_value( + { + key: value + for key, value in optimizer.state[param].items() + if key not in scratch + and not ( + isinstance(key, str) and key.startswith(_CARRIER_PREFIX) + ) + } + ) + for param in params + ], + "groups": [ + _clone_value( + { + key: value + for key, value in group.items() + if key not in ("params", "_gefen_checkpoint_metadata") + } + ) + for group in optimizer.param_groups + ], + "params": [_clone_value(param) for param in params], + "global_step": optimizer._gefen_global_step, + "codebook": _clone_value(optimizer._gefen_codebook), + } + + +def _worker( + rank: int, world: int, port: str, optimizer_kind: str, result_queue +) -> None: + import torch.distributed as dist + import torch.nn as nn + from torch.distributed.checkpoint.state_dict import StateDictOptions, get_optimizer_state_dict, set_optimizer_state_dict + from torch.distributed.tensor import Shard, distribute_tensor, init_device_mesh + + from gefen import Gefen, GefenMuon + + os.environ["MASTER_ADDR"] = "127.0.0.1" + os.environ["MASTER_PORT"] = port + os.environ["RANK"] = str(rank) + os.environ["WORLD_SIZE"] = str(world) + dist.init_process_group("gloo", rank=rank, world_size=world) + try: + mesh = init_device_mesh("cpu", (world,), mesh_dim_names=("dp",)) + + def make_model(full_values): + model = nn.Module() + model.register_parameter( + "first", + nn.Parameter(distribute_tensor(full_values[0].clone(), mesh, [Shard(0)])), + ) + model.register_parameter( + "second", + nn.Parameter(distribute_tensor(full_values[1].clone(), mesh, [Shard(0)])), + ) + return model + + def make_optimizer(model): + if optimizer_kind == "muon_normuon": + return GefenMuon( + [ + {"params": [("first", model.first)]}, + {"params": [("second", model.second)]}, + ], + lr=1e-3, + fused=False, + sharded_mode="approx", + normuon=True, + ) + return Gefen( + [ + {"params": [("first", model.first)]}, + {"params": [("second", model.second)]}, + ], + lr=1e-3, + fused=False, + factored_v_2d=False, + ) + + def assign_grads(model, grads): + model.first.grad = distribute_tensor(grads[0].clone(), mesh, [Shard(0)]) + model.second.grad = distribute_tensor(grads[1].clone(), mesh, [Shard(0)]) + + initial = [ + torch.linspace(-1, 1, 64).reshape(8, 8), + torch.linspace(1, -1, 64).reshape(8, 8), + ] + first_grads = [ + torch.cat((torch.ones(4, 8), torch.arange(32).reshape(4, 8).sin())), + torch.arange(64).reshape(8, 8).cos() * 0.75, + ] + second_grads = [ + torch.arange(64).reshape(8, 8).cos(), + torch.arange(64).reshape(8, 8).sin() * 1.25, + ] + model = make_model(initial) + optimizer = make_optimizer(model) + + pristine_osd = optimizer.state_dict() + pristine_target_model = make_model(initial) + pristine_target = make_optimizer(pristine_target_model) + pristine_target.load_state_dict(pristine_osd) + pristine_restored = _values_equal( + _persistent_optimizer_snapshot(optimizer), + _persistent_optimizer_snapshot(pristine_target), + ) + + assign_grads(model, first_grads) + optimizer.step() + + current_params = [ + model.first.detach().full_tensor().clone(), + model.second.detach().full_tensor().clone(), + ] + full_osd = get_optimizer_state_dict( + model, + optimizer, + options=StateDictOptions(full_state_dict=True, cpu_offload=True), + ) + handoff = [full_osd if rank == 0 else None] + dist.broadcast_object_list(handoff, src=0) + full_osd = handoff[0] + + saved_ids = [ + item + for group in full_osd["param_groups"] + for item in group["params"] + ] + carrier_state = full_osd["state"][saved_ids[0]] + member_state = full_osd["state"][saved_ids[1]] + expected_carrier_keys = {_carrier_key(item) for item in range(world)} + markers = [ + group["_gefen_checkpoint_metadata"]["rank_local_sharded_state"] + for group in full_osd["param_groups"] + ] + marker = markers[0] + payload = Gefen._deserialize_rank_local_payload( + carrier_state[_carrier_key(rank)] + ) + schema_ok = ( + len(saved_ids) == 2 + and set(carrier_state) == expected_carrier_keys | {"name"} + and carrier_state["name"] == 0 + and member_state == {"name": 0, _MEMBER: True} + and all( + torch.is_tensor(carrier_state[key]) + for key in expected_carrier_keys + ) + and marker["format"] == _FORMAT + and marker["world_size"] == world + and marker["world_ranks"] == list(range(world)) + and marker["mesh"]["shape"] == [world] + and marker["mesh"]["ranks"] == list(range(world)) + and all(item == marker for item in markers[1:]) + and payload is not None + and payload["format"] == _FORMAT + and payload["global_rank"] == rank + and payload["group_rank"] == rank + and payload["world_ranks"] == list(range(world)) + and payload["signature"] == marker["signatures"][str(rank)] + and payload["parameter_manifest"] == marker["parameter_manifest"] + and [item["name"] for item in payload["parameter_manifest"]] + == ["first", "second"] + and payload["global_step"] == optimizer._gefen_global_step + ) + + resumed_model = make_model(current_params) + resumed = make_optimizer(resumed_model) + resumed.load_state_dict(full_osd) + restored = _values_equal( + _persistent_optimizer_snapshot(optimizer), + _persistent_optimizer_snapshot(resumed), + ) + + flat_osd = get_optimizer_state_dict( + model, + optimizer, + options=StateDictOptions( + full_state_dict=True, + cpu_offload=True, + flatten_optimizer_state_dict=True, + ), + ) + flat_handoff = [flat_osd if rank == 0 else None] + dist.broadcast_object_list(flat_handoff, src=0) + flat_osd = flat_handoff[0] + flat_model = make_model(current_params) + flat_resumed = make_optimizer(flat_model) + set_optimizer_state_dict( + flat_model, + flat_resumed, + flat_osd, + options=StateDictOptions( + flatten_optimizer_state_dict=True, + ), + ) + flatten_restored = _values_equal( + _persistent_optimizer_snapshot(optimizer), + _persistent_optimizer_snapshot(flat_resumed), + ) + + assign_grads(model, second_grads) + assign_grads(resumed_model, second_grads) + assign_grads(flat_model, second_grads) + optimizer.step() + resumed.step() + flat_resumed.step() + exact = all( + torch.equal(source.detach().full_tensor(), target.detach().full_tensor()) + for source, target in zip(model.parameters(), resumed_model.parameters()) + ) + flatten_exact = all( + torch.equal(source.detach().full_tensor(), target.detach().full_tensor()) + for source, target in zip(model.parameters(), flat_model.parameters()) + ) + + original_step = optimizer._gefen_global_step + optimizer._gefen_global_step = original_step + rank + before_divergent_save = _optimizer_snapshot(optimizer) + divergent_rejected = False + try: + optimizer.state_dict() + except RuntimeError as exc: + divergent_rejected = "global_step differs across ranks" in str(exc) + divergent_unchanged = _values_equal( + before_divergent_save, _optimizer_snapshot(optimizer) + ) + optimizer._gefen_global_step = original_step + + def rewrite_rank_payload(checkpoint, mutate): + local_ids = [ + item + for group in checkpoint["param_groups"] + for item in group["params"] + ] + local_carrier = checkpoint["state"][local_ids[0]] + payload_key = _carrier_key(rank) + local_payload = Gefen._deserialize_rank_local_payload( + local_carrier[payload_key] + ) + mutate(local_payload) + local_carrier[payload_key] = Gefen._serialize_rank_local_payload( + local_payload + ) + + corruptions = {} + + bad = copy.deepcopy(full_osd) + bad["state"][saved_ids[0]][_carrier_key(rank)] = torch.zeros( + 8, dtype=torch.uint8 + ) + corruptions["corrupt_bytes"] = bad + + bad = copy.deepcopy(full_osd) + for carrier_key in expected_carrier_keys: + bad["state"][saved_ids[0]].pop(carrier_key) + corruptions["missing_carrier"] = bad + + bad = copy.deepcopy(full_osd) + bad["state"][saved_ids[0]].pop(_carrier_key(rank)) + corruptions["missing_rank_payload"] = bad + + bad = copy.deepcopy(full_osd) + bad["state"][saved_ids[1]] = copy.deepcopy( + bad["state"][saved_ids[0]] + ) + corruptions["duplicate_carrier"] = bad + + bad = copy.deepcopy(full_osd) + rewrite_rank_payload( + bad, + lambda item: item.__setitem__( + "codebook", torch.full((256,), float("nan"), dtype=torch.float32) + ), + ) + corruptions["bad_codebook"] = bad + + bad = copy.deepcopy(full_osd) + bad["param_groups"][-1]["params"] = [saved_ids[0]] + corruptions["duplicate_outer_parameter_id"] = bad + + bad = copy.deepcopy(full_osd) + for group in bad["param_groups"]: + group["_gefen_checkpoint_metadata"]["rank_local_sharded_state"][ + "signatures" + ][str(rank)][0]["local_shape"] = [999] + corruptions["bad_topology"] = bad + + bad = copy.deepcopy(full_osd) + rewrite_rank_payload( + bad, + lambda item: item.__setitem__("global_rank", (rank + 1) % world), + ) + corruptions["bad_rank_binding"] = bad + + bad = copy.deepcopy(full_osd) + rewrite_rank_payload( + bad, + lambda item: item["states"].reverse(), + ) + corruptions["bad_parameter_order"] = bad + + bad = copy.deepcopy(full_osd) + rewrite_rank_payload( + bad, + lambda item: item["states"][0].pop("step"), + ) + corruptions["missing_step"] = bad + + bad = copy.deepcopy(full_osd) + rewrite_rank_payload( + bad, + lambda item: item["states"][0]["m_magnitude"].view(-1).__setitem__( + 0, -1.0 + ), + ) + corruptions["negative_m_magnitude"] = bad + + if optimizer_kind == "gefen": + bad = copy.deepcopy(full_osd) + rewrite_rank_payload( + bad, + lambda item: item["states"][0].pop("vmean"), + ) + corruptions["missing_vmean"] = bad + + bad = copy.deepcopy(full_osd) + rewrite_rank_payload( + bad, + lambda item: item["states"][0]["vmean"].view(-1).__setitem__( + 0, -1.0 + ), + ) + corruptions["negative_vmean"] = bad + + def install_bad_factored_state(item, key): + pstate = item["states"][0] + rows, cols = item["signature"][0]["shape"] + pstate["v_row"] = torch.zeros(rows, dtype=torch.float32) + pstate["v_col"] = torch.zeros(cols, dtype=torch.float32) + pstate["factored_step"] = _clone_value(pstate["step"]) + pstate[key][0] = -1.0 + + for key in ("v_row", "v_col"): + bad = copy.deepcopy(full_osd) + rewrite_rank_payload( + bad, + lambda item, key=key: install_bad_factored_state(item, key), + ) + corruptions[f"negative_{key}"] = bad + else: + bad = copy.deepcopy(full_osd) + rewrite_rank_payload( + bad, + lambda item: item["states"][0].pop("normuon_step"), + ) + corruptions["incomplete_normuon_pair"] = bad + + bad = copy.deepcopy(full_osd) + rewrite_rank_payload( + bad, + lambda item: item["states"][0]["normuon_v"].view(-1).__setitem__( + 0, -1.0 + ), + ) + corruptions["negative_normuon_v"] = bad + + bad = copy.deepcopy(full_osd) + rewrite_rank_payload( + bad, + lambda item: item["states"][0].__setitem__( + "normuon_v", + torch.zeros( + item["states"][0]["normuon_v"].shape[0] + 1, + 1, + dtype=torch.float32, + ), + ), + ) + corruptions["wrong_normuon_shape"] = bad + + bad = copy.deepcopy(full_osd) + rewrite_rank_payload( + bad, + lambda item: item["states"][0].__setitem__( + "normuon_v", item["states"][0]["normuon_v"].to(torch.float64) + ), + ) + corruptions["wrong_normuon_dtype"] = bad + + bad = copy.deepcopy(full_osd) + rewrite_rank_payload( + bad, + lambda item: item["states"][0].__setitem__("normuon_step", 0), + ) + corruptions["zero_normuon_step"] = bad + + bad = copy.deepcopy(full_osd) + second_metadata = copy.deepcopy( + bad["param_groups"][1]["_gefen_checkpoint_metadata"] + ) + second_metadata["rank_local_sharded_state"]["global_step"] += 1 + bad["param_groups"][1]["_gefen_checkpoint_metadata"] = second_metadata + corruptions["inconsistent_group_marker"] = bad + + rejection_checks = {} + for name, bad_checkpoint in corruptions.items(): + checkpoint_before = _clone_value(bad_checkpoint) + target_model = make_model(current_params) + target = make_optimizer(target_model) + assign_grads(target_model, first_grads) + target.step() + before = _optimizer_snapshot(target) + rejected = False + try: + target.load_state_dict(bad_checkpoint) + except (ValueError, RuntimeError): + rejected = True + rejection_checks[name] = ( + rejected + and _values_equal(before, _optimizer_snapshot(target)) + and _values_equal(checkpoint_before, bad_checkpoint) + ) + + rank_checks = [None] * world + dist.all_gather_object( + rank_checks, + { + "schema": schema_ok, + "pristine_restored": pristine_restored, + "restored": restored, + "flatten_restored": flatten_restored, + "exact": exact, + "flatten_exact": flatten_exact, + "divergent_rejected": divergent_rejected, + "divergent_unchanged": divergent_unchanged, + **rejection_checks, + }, + ) + if rank == 0: + result_queue.put(rank_checks) + except BaseException: + result_queue.put({"rank": rank, "traceback": traceback.format_exc()}) + raise + finally: + dist.destroy_process_group() + + +@pytest.mark.skipif( + not torch.distributed.is_available(), reason="torch.distributed is unavailable" +) +@pytest.mark.parametrize("optimizer_kind", ["gefen", "muon_normuon"]) +def test_full_dcp_handoff_is_exact_on_same_dtensor_topology(optimizer_kind): + import torch.multiprocessing as mp + + context = mp.get_context("spawn") + result_queue = context.Queue() + port = _free_port() + processes = [ + context.Process( + target=_worker, + args=(rank, 2, port, optimizer_kind, result_queue), + ) + for rank in range(2) + ] + for process in processes: + process.start() + try: + rank_checks = result_queue.get(timeout=180) + except queue.Empty: + rank_checks = None + for process in processes: + process.join(timeout=180) + for process in processes: + if process.is_alive(): + process.terminate() + process.join(timeout=10) + assert rank_checks is not None, "distributed checkpoint workers timed out" + assert all(process.exitcode == 0 for process in processes) + assert isinstance(rank_checks, list), rank_checks + failures = [ + {name: value for name, value in checks.items() if not value} + for checks in rank_checks + ] + assert all(not rank_failures for rank_failures in failures), failures + + +def _fully_shard_worker( + rank: int, world: int, port: str, optimizer_kind: str, result_queue +) -> None: + import torch.distributed as dist + import torch.nn as nn + from torch.distributed.checkpoint.state_dict import ( + StateDictOptions, + get_optimizer_state_dict, + set_optimizer_state_dict, + ) + from torch.distributed.device_mesh import init_device_mesh + from torch.distributed.fsdp import fully_shard + from torch.distributed.tensor import Shard, distribute_tensor + + from gefen import Gefen, GefenMuon + + os.environ["MASTER_ADDR"] = "127.0.0.1" + os.environ["MASTER_PORT"] = port + os.environ["RANK"] = str(rank) + os.environ["WORLD_SIZE"] = str(world) + torch.cuda.set_device(rank) + dist.init_process_group("nccl", rank=rank, world_size=world) + try: + mesh = init_device_mesh("cuda", (world,), mesh_dim_names=("dp",)) + full = torch.linspace(-1, 1, 64, device="cuda").reshape(8, 8) + first_grad = torch.cat( + ( + torch.ones(4, 8, device="cuda"), + torch.arange(32, device="cuda").reshape(4, 8).sin(), + ), + dim=0, + ) + model = nn.Linear(8, 8, bias=False, device="cuda") + with torch.no_grad(): + model.weight.copy_(full) + fully_shard(model, mesh=mesh) + if optimizer_kind == "gefen": + optimizer = Gefen( + model.named_parameters(), + lr=1e-3, + fused=False, + factored_v_2d=False, + ) + else: + optimizer = GefenMuon( + model.named_parameters(), + lr=1e-3, + fused=False, + sharded_mode="approx", + ) + model.weight.grad = distribute_tensor(first_grad, mesh, [Shard(0)]) + optimizer.step() + current = model.weight.detach().full_tensor().clone() + original_codebook = optimizer._gefen_codebook.detach().cpu().clone() + original_indices = optimizer.state[model.weight]["m_codebook"].detach().cpu().clone() + full_osd = get_optimizer_state_dict( + model, + optimizer, + options=StateDictOptions(full_state_dict=True, cpu_offload=True), + ) + + resumed_model = nn.Linear(8, 8, bias=False, device="cuda") + with torch.no_grad(): + resumed_model.weight.copy_(current) + fully_shard(resumed_model, mesh=mesh) + if optimizer_kind == "gefen": + resumed = Gefen( + resumed_model.named_parameters(), + lr=1e-3, + fused=False, + factored_v_2d=False, + ) + else: + resumed = GefenMuon( + resumed_model.named_parameters(), + lr=1e-3, + fused=False, + sharded_mode="approx", + ) + set_optimizer_state_dict( + resumed_model, + resumed, + full_osd if rank == 0 else {}, + options=StateDictOptions( + full_state_dict=True, + broadcast_from_rank0=True, + ), + ) + restored = ( + torch.equal(resumed._gefen_codebook.detach().cpu(), original_codebook) + and torch.equal( + resumed.state[resumed_model.weight]["m_codebook"].detach().cpu(), + original_indices, + ) + ) + + second_grad = torch.arange(64, device="cuda").reshape(8, 8).cos() + model.weight.grad = distribute_tensor(second_grad, mesh, [Shard(0)]) + resumed_model.weight.grad = distribute_tensor(second_grad, mesh, [Shard(0)]) + optimizer.step() + resumed.step() + uninterrupted = model.weight.detach().full_tensor() + after_resume = resumed_model.weight.detach().full_tensor() + exact = torch.equal(uninterrupted, after_resume) + checks = [None] * world + dist.all_gather_object(checks, (restored, exact)) + if rank == 0: + result_queue.put(checks) + finally: + dist.destroy_process_group() + + +@pytest.mark.skipif( + not torch.cuda.is_available() + or torch.cuda.device_count() < 2 + or not torch.distributed.is_nccl_available(), + reason="fully_shard checkpoint continuation requires two CUDA GPUs and NCCL", +) +@pytest.mark.parametrize("optimizer_kind", ["gefen", "muon_approx"]) +def test_rank_local_full_dcp_set_optimizer_state_is_exact_under_fully_shard( + optimizer_kind, +): + import torch.multiprocessing as mp + + context = mp.get_context("spawn") + result_queue = context.Queue() + port = _free_port() + processes = [ + context.Process( + target=_fully_shard_worker, + args=(rank, 2, port, optimizer_kind, result_queue), + ) + for rank in range(2) + ] + for process in processes: + process.start() + try: + checks = result_queue.get(timeout=180) + except queue.Empty: + checks = None + for process in processes: + process.join(timeout=180) + for process in processes: + if process.is_alive(): + process.terminate() + process.join(timeout=10) + assert checks is not None, "fully_shard checkpoint workers timed out" + assert all(process.exitcode == 0 for process in processes) + assert all(all(rank_check) for rank_check in checks), checks diff --git a/tests/test_muon_distributed_checkpoint_safety.py b/tests/test_muon_distributed_checkpoint_safety.py new file mode 100644 index 0000000..7f53b9b --- /dev/null +++ b/tests/test_muon_distributed_checkpoint_safety.py @@ -0,0 +1,839 @@ +"""Fail-closed checkpoint validation for Parallel-Muon owner state.""" + +import copy +from datetime import timedelta +import os +import queue +import socket +import traceback + +import pytest +import torch +import torch.nn as nn + +from gefen import GefenMuon + +_DISTRIBUTED_METADATA = "muon_distributed_state" + + +def _cpu_optimizer(*, mode="distributed"): + generator = torch.Generator().manual_seed(771) + params = [ + ( + name, + nn.Parameter(torch.randn(*shape, generator=generator) * 0.02), + ) + for name, shape in (("left", (8, 8)), ("right", (6, 10))) + ] + optimizer = GefenMuon( + params, + lr=2e-3, + fused=False, + ns_steps=1, + sharded_mode=mode, + ) + return optimizer, params + + +def _step_cpu(optimizer, params, seed=772): + generator = torch.Generator().manual_seed(seed) + for _, param in params: + param.grad = torch.randn(param.shape, generator=generator) * 0.002 + optimizer.step() + optimizer.zero_grad(set_to_none=True) + + +def _marker_v2(optimizer, state_dict, world=2): + saved_ids = [ + saved_id + for group in state_dict["param_groups"] + for saved_id in group["params"] + ] + live = [ + (name, param) + for group in optimizer.param_groups + for name, param in optimizer._iter_group_params_with_names(group) + ] + return { + "version": 2, + "ownership": "stable_full_param_index_v1", + "consolidated": True, + "groups": [ + { + "world_size": world, + "params": [ + { + "saved_id": saved_id, + "name": str(name), + "shape": tuple(param.shape), + "owner": index % world, + "state_keys": tuple( + sorted(str(key) for key in state_dict["state"][saved_id]) + ), + "initialized": any( + str(key) != "name" + for key in state_dict["state"][saved_id] + ), + } + for index, ((name, param), saved_id) in enumerate( + zip(live, saved_ids) + ) + ], + } + ], + } + + +def _clone(value): + if torch.is_tensor(value): + return value.detach().clone() + if isinstance(value, dict): + return {key: _clone(item) for key, item in value.items()} + if isinstance(value, list): + return [_clone(item) for item in value] + if isinstance(value, tuple): + return tuple(_clone(item) for item in value) + return copy.deepcopy(value) + + +def _assert_equal(actual, expected): + if torch.is_tensor(expected): + assert torch.is_tensor(actual) + assert actual.dtype == expected.dtype + assert torch.equal(actual, expected) + return + if isinstance(expected, dict): + assert isinstance(actual, dict) + assert set(actual) == set(expected) + for key in expected: + _assert_equal(actual[key], expected[key]) + return + if isinstance(expected, (list, tuple)): + assert isinstance(actual, type(expected)) + assert len(actual) == len(expected) + for actual_item, expected_item in zip(actual, expected): + _assert_equal(actual_item, expected_item) + return + assert actual == expected + + +def _simulate_live_dtensor(monkeypatch, optimizer): + process_group = object() + monkeypatch.setattr(optimizer, "_dist_available", lambda: True) + monkeypatch.setattr( + optimizer, "_distributed_process_group", lambda param: process_group + ) + monkeypatch.setattr(optimizer, "_drop_non_owned_distributed_state", lambda: None) + + +def _stateful_checkpoint(): + source, source_params = _cpu_optimizer() + _step_cpu(source, source_params) + checkpoint = _clone(source.state_dict()) + checkpoint["gefen_muon_distributed"] = _marker_v2(source, checkpoint) + return checkpoint + + +def _corrupt(checkpoint, case): + corrupted = _clone(checkpoint) + marker = corrupted["gefen_muon_distributed"] + entry = marker["groups"][0]["params"][0] + saved_id = corrupted["param_groups"][0]["params"][0] + if case == "missing_marker": + corrupted.pop("gefen_muon_distributed") + elif case == "markerless_deleted_state": + corrupted.pop("gefen_muon_distributed") + for group in corrupted["param_groups"]: + for index, param_id in enumerate(group["params"]): + corrupted["state"][param_id] = { + "name": group["param_names"][index] + } + elif case == "incomplete_marker": + corrupted["gefen_muon_distributed"] = {"consolidated": True} + elif case == "wrong_version": + marker["version"] = 99 + elif case == "wrong_ownership": + marker["ownership"] = "rank_zero_v0" + elif case == "not_consolidated": + marker["consolidated"] = False + elif case == "missing_groups": + marker.pop("groups") + elif case == "bad_world": + marker["groups"][0]["world_size"] = 0 + elif case == "bad_owner": + entry["owner"] = 1 + elif case == "bad_saved_id": + entry["saved_id"] = 999 + elif case == "duplicate_saved_id": + marker["groups"][0]["params"][1]["saved_id"] = entry["saved_id"] + elif case == "bad_name": + entry["name"] = "other" + elif case == "bad_shape": + entry["shape"] = (1, 64) + elif case == "missing_state": + corrupted["state"].pop(saved_id) + elif case == "partial_core": + corrupted["state"][saved_id].pop("m_magnitude") + entry["state_keys"] = tuple(sorted(corrupted["state"][saved_id])) + elif case == "fractional_step": + corrupted["state"][saved_id]["step"] = 1.5 + elif case == "negative_magnitude": + corrupted["state"][saved_id]["m_magnitude"].fill_(-1) + elif case == "bad_codebook": + bad = torch.zeros(12, dtype=torch.float32) + corrupted["gefen_codebook"] = bad + for group in corrupted["param_groups"]: + group["_gefen_checkpoint_metadata"]["codebook"] = bad.clone() + elif case == "nonstring_manifest_key": + entry["state_keys"] = (*entry["state_keys"], 7) + elif case == "bad_initialized": + entry["initialized"] = False + else: + raise AssertionError(case) + return corrupted + + +@pytest.mark.parametrize( + "case", + [ + "missing_marker", + "markerless_deleted_state", + "incomplete_marker", + "wrong_version", + "wrong_ownership", + "not_consolidated", + "missing_groups", + "bad_world", + "bad_owner", + "bad_saved_id", + "duplicate_saved_id", + "bad_name", + "bad_shape", + "missing_state", + "partial_core", + "fractional_step", + "negative_magnitude", + "bad_codebook", + "nonstring_manifest_key", + "bad_initialized", + ], +) +def test_populated_distributed_schema_rejects_before_mutation(monkeypatch, case): + checkpoint = _corrupt(_stateful_checkpoint(), case) + target, target_params = _cpu_optimizer() + before_state = _clone(target.state) + before_groups = _clone(target.param_groups) + before_params = [param.detach().clone() for _, param in target_params] + before_codebook = target._gefen_codebook + before_step = target._gefen_global_step + _simulate_live_dtensor(monkeypatch, target) + + with pytest.raises(ValueError, match="refused populated"): + target.load_state_dict(checkpoint) + + _assert_equal(target.state, before_state) + _assert_equal(target.param_groups, before_groups) + for (_, param), expected in zip(target_params, before_params): + assert torch.equal(param.detach(), expected) + assert target._gefen_codebook is before_codebook + assert target._gefen_global_step == before_step + + +def test_valid_v2_and_released_v1_schemas_load(monkeypatch): + checkpoint = _stateful_checkpoint() + target, _ = _cpu_optimizer() + _simulate_live_dtensor(monkeypatch, target) + target.load_state_dict(_clone(checkpoint)) + assert target._gefen_global_step == 1 + + released_v1 = _clone(checkpoint) + released_v1["gefen_muon_distributed"] = { + "version": 1, + "ownership": "stable_full_param_index_v1", + "consolidated": True, + } + target_v1, _ = _cpu_optimizer() + _simulate_live_dtensor(monkeypatch, target_v1) + target_v1.load_state_dict(released_v1) + assert target_v1._gefen_global_step == 1 + + +def test_markerless_pristine_and_non_distributed_legacy_remain_safe(monkeypatch): + fresh_source, _ = _cpu_optimizer() + fresh = _clone(fresh_source.state_dict()) + fresh_target, _ = _cpu_optimizer() + _simulate_live_dtensor(monkeypatch, fresh_target) + fresh_target.load_state_dict(fresh) + assert fresh_target._gefen_global_step == 0 + + exact_source, exact_params = _cpu_optimizer(mode="exact") + _step_cpu(exact_source, exact_params) + exact_target, _ = _cpu_optimizer(mode="exact") + exact_target.load_state_dict(_clone(exact_source.state_dict())) + assert exact_target._gefen_global_step == 1 + + +def test_released_v1_mixed_owner_state_fails_closed(monkeypatch): + checkpoint = _stateful_checkpoint() + checkpoint["gefen_muon_distributed"] = { + "version": 1, + "ownership": "stable_full_param_index_v1", + "consolidated": True, + } + second_id = checkpoint["param_groups"][0]["params"][1] + checkpoint["state"][second_id] = {"name": "right"} + target, _ = _cpu_optimizer() + _simulate_live_dtensor(monkeypatch, target) + with pytest.raises(ValueError, match="mixes initialized and empty"): + target.load_state_dict(checkpoint) + + +def _free_port(): + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.bind(("127.0.0.1", 0)) + return str(sock.getsockname()[1]) + + +def _distributed_grads(step, shapes, device): + generator = torch.Generator(device="cpu").manual_seed(8800 + step) + return [ + (torch.randn(*shape, generator=generator) * 0.002).to( + device, torch.bfloat16 + ) + for shape in shapes + ] + + +def _clone_local_state(state): + return { + key: ( + value.to_local().detach().clone() + if torch.is_tensor(value) and hasattr(value, "to_local") + else value.detach().clone() + if torch.is_tensor(value) + else copy.deepcopy(value) + ) + for key, value in state.items() + } + + +def _local_state_equal(actual, expected): + if set(actual) != set(expected): + return False + for key, expected_value in expected.items(): + actual_value = actual[key] + if torch.is_tensor(actual_value) and hasattr(actual_value, "to_local"): + actual_value = actual_value.to_local() + if torch.is_tensor(expected_value): + if not torch.is_tensor(actual_value) or not torch.equal( + actual_value, expected_value + ): + return False + elif actual_value != expected_value: + return False + return True + + +def _mixed_cpu_checkpoint_worker(rank, world, port, result_queue): + import torch.distributed as dist + from torch.distributed.tensor import Shard, distribute_tensor, init_device_mesh + + try: + os.environ["MASTER_ADDR"] = "127.0.0.1" + os.environ["MASTER_PORT"] = port + os.environ["RANK"] = str(rank) + os.environ["WORLD_SIZE"] = str(world) + dist.init_process_group( + "gloo", + rank=rank, + world_size=world, + timeout=timedelta(seconds=30), + ) + mesh = init_device_mesh("cpu", (world,)) + + def full(seed): + return torch.randn(8, 8, generator=torch.Generator().manual_seed(seed)) + + def sharded(value): + return nn.Parameter(distribute_tensor(value.clone(), mesh, [Shard(0)])) + + def assign(param, value): + param.grad = distribute_tensor(value.clone(), mesh, [Shard(0)]) + + def live_snapshot(optimizer, params): + return ( + optimizer._gefen_global_step, + None + if optimizer._gefen_codebook is None + else optimizer._gefen_codebook.detach().clone(), + [_clone_local_state(optimizer.state[param]) for param in params], + [ + param.detach().to_local().clone() + if hasattr(param, "to_local") + else param.detach().clone() + for param in params + ], + ) + + def snapshot_equal(optimizer, params, snapshot): + step, codebook, states, values = snapshot + current_codebook = optimizer._gefen_codebook + return ( + optimizer._gefen_global_step == step + and ( + (current_codebook is None and codebook is None) + or ( + torch.is_tensor(current_codebook) + and torch.is_tensor(codebook) + and torch.equal(current_codebook, codebook) + ) + ) + and all( + _local_state_equal(optimizer.state[param], expected) + for param, expected in zip(params, states) + ) + and all( + torch.equal( + param.detach().to_local() + if hasattr(param, "to_local") + else param.detach(), + expected, + ) + for param, expected in zip(params, values) + ) + ) + + # One distributed optimizer group may mix a Parallel-Muon DTensor with + # a replicated plain-tensor fallback. Only the former belongs in the + # saved-world owner proof. + initial_parallel, initial_fallback = full(920), full(921) + + def make_fallback_pair(parallel_value, fallback_value): + parallel = sharded(parallel_value) + fallback = nn.Parameter(fallback_value.clone()) + optimizer = GefenMuon( + [ + { + "params": [ + ("parallel", parallel), + ("fallback", fallback), + ], + "sharded_mode": "distributed", + } + ], + lr=2e-3, + fused=False, + ns_steps=1, + ) + return optimizer, [parallel, fallback] + + source, source_params = make_fallback_pair( + initial_parallel, initial_fallback + ) + assign(source_params[0], full(922) * 0.01) + source_params[1].grad = full(923) * 0.01 + source.step() + checkpoint = _clone(source.state_dict()) + saved_ids = checkpoint["param_groups"][0]["params"] + marker = checkpoint["gefen_muon_distributed"] + manifest_ids = [ + item["saved_id"] + for group in marker["groups"] + for item in group["params"] + ] + fallback_schema_ok = manifest_ids == [saved_ids[0]] + + current_parallel = source_params[0].detach().full_tensor().clone() + current_fallback = source_params[1].detach().clone() + resumed, resumed_params = make_fallback_pair( + current_parallel, current_fallback + ) + resumed.load_state_dict(_clone(checkpoint)) + assign(source_params[0], full(924) * 0.01) + source_params[1].grad = full(925) * 0.01 + assign(resumed_params[0], full(924) * 0.01) + resumed_params[1].grad = full(925) * 0.01 + source.step() + resumed.step() + fallback_continuation = torch.equal( + source_params[0].detach().to_local(), + resumed_params[0].detach().to_local(), + ) and torch.equal(source_params[1], resumed_params[1]) + + # Adding the replicated fallback to every copy of the owner proof must + # fail before touching even a pre-existing target transaction. + corrupted = _clone(checkpoint) + marker_copies = [corrupted["gefen_muon_distributed"]] + marker_copies.extend( + group["_gefen_checkpoint_metadata"][ + _DISTRIBUTED_METADATA + ] + for group in corrupted["param_groups"] + ) + for marker_copy in marker_copies: + extra = _clone(marker_copy["groups"][0]["params"][0]) + extra.update( + { + "saved_id": saved_ids[1], + "name": "fallback", + "shape": tuple(initial_fallback.shape), + "owner": 1, + } + ) + marker_copy["groups"][0]["params"].append(extra) + rejected, rejected_params = make_fallback_pair( + initial_parallel, initial_fallback + ) + before = live_snapshot(rejected, rejected_params) + rejection_message = None + try: + rejected.load_state_dict(corrupted) + except ValueError as exc: + rejection_message = str(exc) + fallback_atomic = ( + rejection_message is not None + and "ordered eligible" in rejection_message + and snapshot_equal(rejected, rejected_params, before) + ) + + # Rank-local approx state and stable-owner distributed state compose in + # one optimizer: consolidate owners first, then wrap all real state into + # rank-indexed payloads, and reverse that order on load validation. + initial_approx, initial_distributed = full(930), full(931) + + def make_composed_pair(approx_value, distributed_value): + approx = sharded(approx_value) + distributed = sharded(distributed_value) + optimizer = GefenMuon( + [ + { + "params": [("approx", approx)], + "sharded_mode": "approx", + }, + { + "params": [("distributed", distributed)], + "sharded_mode": "distributed", + }, + ], + lr=2e-3, + fused=False, + ns_steps=1, + ) + return optimizer, [approx, distributed] + + composed, composed_params = make_composed_pair( + initial_approx, initial_distributed + ) + assign(composed_params[0], full(932) * 0.01) + assign(composed_params[1], full(933) * 0.01) + composed.step() + composed_checkpoint = _clone(composed.state_dict()) + current_composed = [ + param.detach().full_tensor().clone() for param in composed_params + ] + composed_resumed, composed_resumed_params = make_composed_pair( + *current_composed + ) + composed_resumed.load_state_dict(_clone(composed_checkpoint)) + metadata_only_checkpoint = _clone(composed_checkpoint) + metadata_only_checkpoint.pop("gefen_muon_distributed") + metadata_resumed, metadata_resumed_params = make_composed_pair( + *current_composed + ) + metadata_resumed.load_state_dict(metadata_only_checkpoint) + metadata_transport_ok = ( + metadata_resumed._gefen_global_step + == composed_resumed._gefen_global_step + and all( + _local_state_equal( + metadata_resumed.state[left], + composed_resumed.state[right], + ) + for left, right in zip( + metadata_resumed_params, composed_resumed_params + ) + ) + ) + assign(composed_params[0], full(934) * 0.01) + assign(composed_params[1], full(935) * 0.01) + assign(composed_resumed_params[0], full(934) * 0.01) + assign(composed_resumed_params[1], full(935) * 0.01) + composed.step() + composed_resumed.step() + composed_continuation = all( + torch.equal(left.detach().to_local(), right.detach().to_local()) + for left, right in zip(composed_params, composed_resumed_params) + ) + + result_queue.put( + ( + "result", + rank, + fallback_schema_ok, + fallback_continuation, + fallback_atomic, + composed_continuation, + metadata_transport_ok, + rejection_message, + ) + ) + except Exception: + result_queue.put(("error", rank, traceback.format_exc())) + finally: + if dist.is_initialized(): + dist.destroy_process_group() + + +def test_mixed_parallel_fallback_and_rank_local_composition_cpu(): + import torch.multiprocessing as mp + + world = 2 + context = mp.get_context("spawn") + result_queue = context.Queue() + port = _free_port() + processes = [ + context.Process( + target=_mixed_cpu_checkpoint_worker, + args=(rank, world, port, result_queue), + ) + for rank in range(world) + ] + for process in processes: + process.start() + + messages = [] + try: + for _ in range(world): + messages.append(result_queue.get(timeout=60)) + except queue.Empty: + pass + finally: + for process in processes: + process.join(timeout=5) + for process in processes: + if process.is_alive(): + process.terminate() + process.join(timeout=5) + + assert all(process.exitcode == 0 for process in processes), [ + process.exitcode for process in processes + ] + errors = [item[2] for item in messages if item[0] == "error"] + assert not errors, "\n".join(errors) + results = {item[1]: item[2:] for item in messages if item[0] == "result"} + assert set(results) == {0, 1}, messages + for rank, result in results.items(): + ( + schema, + fallback_resume, + atomic, + composed_resume, + metadata_transport, + message, + ) = result + assert schema, rank + assert fallback_resume, rank + assert atomic, (rank, message) + assert composed_resume, rank + assert metadata_transport, rank + + +def _checkpoint_worker(rank, world, port, result_queue): + import torch.distributed as dist + from torch.distributed.tensor import Shard, distribute_tensor, init_device_mesh + + try: + os.environ["MASTER_ADDR"] = "127.0.0.1" + os.environ["MASTER_PORT"] = port + os.environ["RANK"] = str(rank) + os.environ["WORLD_SIZE"] = str(world) + torch.cuda.set_device(rank) + device = torch.device("cuda", rank) + dist.init_process_group( + "nccl", + rank=rank, + world_size=world, + timeout=timedelta(seconds=15), + ) + mesh = init_device_mesh("cuda", (world,)) + specs = (("left", (16, 16)), ("right", (12, 20))) + generator = torch.Generator(device="cpu").manual_seed(8799) + initial = [ + (torch.randn(*shape, generator=generator) * 0.02).to(torch.bfloat16) + for _, shape in specs + ] + + def make_params(full_values): + return [ + ( + name, + nn.Parameter( + distribute_tensor( + value.to(device).clone(), mesh, [Shard(0)] + ) + ), + ) + for (name, _), value in zip(specs, full_values) + ] + + def make_optimizer(params): + return GefenMuon( + params, + lr=2e-3, + fused=False, + ns_steps=1, + sharded_mode="distributed", + ) + + source_params = make_params(initial) + source = make_optimizer(source_params) + shapes = [shape for _, shape in specs] + for step in range(2): + for (_, param), grad in zip( + source_params, _distributed_grads(step, shapes, device) + ): + param.grad = distribute_tensor(grad, mesh, [Shard(0)]) + source.step() + source.zero_grad(set_to_none=True) + + checkpoint = _clone(source.state_dict()) + checkpoint_params = [ + param.detach().full_tensor().clone() for _, param in source_params + ] + + invalid_params = make_params(initial) + invalid_target = make_optimizer(invalid_params) + for (_, param), grad in zip( + invalid_params, _distributed_grads(9, shapes, device) + ): + param.grad = distribute_tensor(grad, mesh, [Shard(0)]) + invalid_target.step() + invalid_target.zero_grad(set_to_none=True) + params_before = [ + param.detach().to_local().clone() for _, param in invalid_params + ] + states_before = [ + _clone_local_state(invalid_target.state[param]) + for _, param in invalid_params + ] + codebook_before = invalid_target._gefen_codebook.detach().clone() + step_before = invalid_target._gefen_global_step + invalid_checkpoint = _clone(checkpoint) + invalid_checkpoint.pop("gefen_muon_distributed") + for group in invalid_checkpoint["param_groups"]: + group["_gefen_checkpoint_metadata"].pop( + _DISTRIBUTED_METADATA, None + ) + error = None + try: + invalid_target.load_state_dict(invalid_checkpoint) + except ValueError as exc: + error = str(exc) + rejection_ok = ( + error is not None + and "refused populated" in error + and invalid_target._gefen_global_step == step_before + and torch.equal(invalid_target._gefen_codebook, codebook_before) + and all( + torch.equal(param.detach().to_local(), expected) + for (_, param), expected in zip(invalid_params, params_before) + ) + and all( + _local_state_equal(invalid_target.state[param], expected) + for (_, param), expected in zip(invalid_params, states_before) + ) + ) + + for (_, param), grad in zip( + source_params, _distributed_grads(2, shapes, device) + ): + param.grad = distribute_tensor(grad, mesh, [Shard(0)]) + source.step() + + resumed_params = make_params(checkpoint_params) + resumed = make_optimizer(resumed_params) + resumed.load_state_dict(_clone(checkpoint)) + for (_, param), grad in zip( + resumed_params, _distributed_grads(2, shapes, device) + ): + param.grad = distribute_tensor(grad, mesh, [Shard(0)]) + resumed.step() + continuation_ok = all( + torch.equal( + resumed_param.detach().to_local(), source_param.detach().to_local() + ) + for (_, resumed_param), (_, source_param) in zip( + resumed_params, source_params + ) + ) + marker = checkpoint.get("gefen_muon_distributed", {}) + result_queue.put( + ( + "result", + rank, + rejection_ok, + continuation_ok, + marker.get("version"), + error, + ) + ) + except Exception: + result_queue.put(("error", rank, traceback.format_exc())) + finally: + if dist.is_initialized(): + dist.destroy_process_group() + + +@pytest.mark.skipif( + not torch.cuda.is_available() + or torch.cuda.device_count() < 2 + or not torch.distributed.is_available() + or not torch.distributed.is_nccl_available(), + reason="Parallel-Muon checkpoint safety requires two CUDA GPUs and NCCL", +) +def test_parallel_muon_rejection_is_atomic_and_valid_v2_continues(): + import torch.multiprocessing as mp + + world = 2 + context = mp.get_context("spawn") + result_queue = context.Queue() + port = _free_port() + processes = [ + context.Process( + target=_checkpoint_worker, + args=(rank, world, port, result_queue), + ) + for rank in range(world) + ] + for process in processes: + process.start() + + messages = [] + try: + for _ in range(world): + messages.append(result_queue.get(timeout=45)) + except queue.Empty: + pass + finally: + for process in processes: + process.join(timeout=5) + for process in processes: + if process.is_alive(): + process.terminate() + process.join(timeout=5) + + assert all(process.exitcode == 0 for process in processes), [ + process.exitcode for process in processes + ] + errors = [item[2] for item in messages if item[0] == "error"] + assert not errors, "\n".join(errors) + results = { + rank: (rejection, continuation, version, message) + for kind, rank, rejection, continuation, version, message in messages + if kind == "result" + } + assert set(results) == {0, 1}, messages + for rank, (rejection, continuation, version, message) in results.items(): + assert rejection, (rank, message) + assert continuation, rank + assert version == 2, rank diff --git a/tests/test_muon_grad_presence.py b/tests/test_muon_grad_presence.py new file mode 100644 index 0000000..abc2ef3 --- /dev/null +++ b/tests/test_muon_grad_presence.py @@ -0,0 +1,554 @@ +"""Rank-consistent gradient-presence guards for sharded GefenMuon.""" + +import copy +from datetime import timedelta +import os +import queue +import socket +import traceback + +import pytest +import torch +import torch.nn as nn + +from gefen import Gefen, GefenMuon, GefenMuonHybrid + + +def _free_port(): + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.bind(("127.0.0.1", 0)) + return str(sock.getsockname()[1]) + + +def _clone_state(state): + return { + key: value.detach().clone() if torch.is_tensor(value) else copy.deepcopy(value) + for key, value in state.items() + } + + +def _state_equal(actual, expected): + if set(actual) != set(expected): + return False + for key, expected_value in expected.items(): + actual_value = actual[key] + if torch.is_tensor(expected_value): + if not torch.is_tensor(actual_value) or not torch.equal( + actual_value, expected_value + ): + return False + elif actual_value != expected_value: + return False + return True + + +def _mismatch_worker(rank, world, port, result_queue): + import torch.distributed as dist + from torch.distributed.tensor import Shard, distribute_tensor, init_device_mesh + + try: + os.environ["MASTER_ADDR"] = "127.0.0.1" + os.environ["MASTER_PORT"] = port + os.environ["RANK"] = str(rank) + os.environ["WORLD_SIZE"] = str(world) + torch.cuda.set_device(rank) + dist.init_process_group( + "nccl", + rank=rank, + world_size=world, + timeout=timedelta(seconds=12), + ) + mesh = init_device_mesh("cuda", (world,)) + results = [] + + for mode in ("exact", "distributed"): + for phase in ("first_step", "initialized"): + generator = torch.Generator(device="cpu").manual_seed( + 1000 + len(results) + ) + full_init = torch.randn(8, 8, generator=generator).to(rank) + first_grad = (torch.randn(8, 8, generator=generator) * 0.01).to( + rank + ) + mismatch_grad = ( + torch.randn(8, 8, generator=generator) * 0.01 + ).to(rank) + param = nn.Parameter( + distribute_tensor(full_init.clone(), mesh, [Shard(0)]) + ) + optimizer = GefenMuon( + [("weight", param)], + lr=1e-3, + fused=False, + sharded_mode=mode, + ) + + first_grad_dt = distribute_tensor(first_grad, mesh, [Shard(0)]) + mismatch_grad_dt = distribute_tensor( + mismatch_grad, mesh, [Shard(0)] + ) + if phase == "initialized": + param.grad = first_grad_dt + optimizer.step() + optimizer.zero_grad(set_to_none=True) + + # Build the DTensor on every rank so the test itself contributes + # no unmatched collective; only assignment differs by rank. + param.grad = None if rank == 0 else mismatch_grad_dt + param_before = param.detach().to_local().clone() + state_before = _clone_state(optimizer.state[param]) + codebook_before = ( + None + if optimizer._gefen_codebook is None + else optimizer._gefen_codebook.detach().clone() + ) + global_step_before = optimizer._gefen_global_step + + message = None + try: + optimizer.step() + except RuntimeError as exc: + message = str(exc) + + codebook_after = optimizer._gefen_codebook + codebook_unchanged = ( + codebook_before is None and codebook_after is None + ) or ( + codebook_before is not None + and codebook_after is not None + and torch.equal(codebook_before, codebook_after) + ) + results.append( + { + "mode": mode, + "phase": phase, + "message": message, + "clear_error": message is not None + and "identical gradient presence" in message + and "weight (1/2 mesh ranks have gradients)" in message, + "param_unchanged": torch.equal( + param.detach().to_local(), param_before + ), + "state_unchanged": _state_equal( + optimizer.state[param], state_before + ), + "codebook_unchanged": codebook_unchanged, + "global_step_unchanged": optimizer._gefen_global_step + == global_step_before, + } + ) + param.grad = None + dist.barrier(device_ids=[rank]) + + result_queue.put(("result", rank, results)) + except Exception: + result_queue.put(("error", rank, traceback.format_exc())) + finally: + if dist.is_initialized(): + dist.destroy_process_group() + + +def _distributed_amp_mismatch_worker(rank, world, port, result_queue): + import torch.distributed as dist + from torch.distributed.tensor import Shard, distribute_tensor, init_device_mesh + + try: + os.environ["MASTER_ADDR"] = "127.0.0.1" + os.environ["MASTER_PORT"] = port + os.environ["RANK"] = str(rank) + os.environ["WORLD_SIZE"] = str(world) + dist.init_process_group( + "gloo", + rank=rank, + world_size=world, + timeout=timedelta(seconds=12), + ) + mesh = init_device_mesh("cpu", (world,)) + equivalent_mesh = init_device_mesh("cpu", (world,)) + results = [] + + for kind in ( + "gefen", + "gefen_distinct_mesh", + "muon_exact", + "muon_distributed", + "muon_approx", + "hybrid", + "hybrid_mixed", + ): + generator = torch.Generator(device="cpu").manual_seed( + 4000 + len(results) + ) + full_weight = torch.randn(8, 8, generator=generator).to(torch.float16) + full_weight_grad = ( + torch.randn(8, 8, generator=generator) * 0.01 + ).to(torch.float16) + weight = nn.Parameter( + distribute_tensor(full_weight.clone(), mesh, [Shard(0)]) + ) + weight_grad = distribute_tensor( + (full_weight_grad * 8.0).clone(), mesh, [Shard(0)] + ) + params = [weight] + + if kind == "gefen": + optimizer = Gefen( + [("weight", weight)], + lr=1e-3, + fused=False, + factored_v_2d=False, + ) + elif kind == "gefen_distinct_mesh": + full_other = torch.randn(8, generator=generator) + full_other_grad = torch.randn(8, generator=generator) * 0.01 + other = nn.Parameter( + distribute_tensor( + full_other.clone(), equivalent_mesh, [Shard(0)] + ) + ) + other_grad = distribute_tensor( + (full_other_grad * 8.0).clone(), equivalent_mesh, [Shard(0)] + ) + named_params = [("weight", weight), ("other", other)] + if rank == 1: + named_params.reverse() + optimizer = Gefen( + named_params, + lr=1e-3, + fused=False, + factored_v_2d=False, + ) + weight.grad = weight_grad + other.grad = None if rank == 0 else other_grad + params.append(other) + elif kind.startswith("muon_"): + optimizer = GefenMuon( + [("weight", weight)], + lr=1e-3, + fused=False, + sharded_mode=kind.removeprefix("muon_"), + ) + elif kind == "hybrid": + full_bias = torch.randn(8, generator=generator) + full_bias_grad = torch.randn(8, generator=generator) * 0.01 + bias = nn.Parameter( + distribute_tensor(full_bias.clone(), mesh, [Shard(0)]) + ) + bias_grad = distribute_tensor( + (full_bias_grad * 8.0).clone(), mesh, [Shard(0)] + ) + optimizer = GefenMuonHybrid( + [("weight", weight)], + [("bias", bias)], + lr=1e-3, + fused=False, + sharded_mode="exact", + backup_optimizer="adamw", + normuon=False, + ) + bias.grad = bias_grad + params.append(bias) + else: + # Native AMP is selected by this ordinary local FP16 matrix, + # while the only DTensor is a BF16 backup parameter. GradScaler + # still scans that DTensor, so its rank-divergent presence must + # be caught by the same union-wide preflight. + weight = nn.Parameter(full_weight.clone()) + full_bias = torch.randn(8, generator=generator).to(torch.bfloat16) + full_bias_grad = ( + torch.randn(8, generator=generator) * 0.01 + ).to(torch.bfloat16) + bias = nn.Parameter( + distribute_tensor(full_bias.clone(), mesh, [Shard(0)]) + ) + bias_grad = distribute_tensor( + (full_bias_grad * 8.0).clone(), mesh, [Shard(0)] + ) + optimizer = GefenMuonHybrid( + [("weight", weight)], + [("bias", bias)], + lr=1e-3, + fused=False, + sharded_mode="exact", + backup_optimizer="adamw", + normuon=False, + ) + weight.grad = (full_weight_grad * 8.0).clone() + bias.grad = None if rank == 0 else bias_grad + params = [weight, bias] + + # Construct the same DTensor gradient on both ranks so the test + # itself is collective-safe, then deliberately assign it on only + # one rank. Hybrid also keeps a shared FP32 backup gradient active, + # proving the inactive FP16 rank cannot select the ordinary path. + if kind not in ("gefen_distinct_mesh", "hybrid_mixed"): + weight.grad = None if rank == 0 else weight_grad + scaler = torch.amp.GradScaler("cpu", init_scale=8.0) + scaler.scale(torch.ones(())) + + def local_clone(tensor): + tensor = tensor.to_local() if hasattr(tensor, "to_local") else tensor + if hasattr(tensor, "wait"): + tensor = tensor.wait() + return tensor.detach().clone() + + param_before = [local_clone(param) for param in params] + grad_before = [ + None if param.grad is None else local_clone(param.grad) + for param in params + ] + children = getattr(optimizer, "_subopts", [optimizer]) + state_before = { + id(param): _clone_state(child.state.get(param, {})) + for child in children + for group in child.param_groups + for param in group["params"] + } + codebook_before = { + id(child): ( + None + if getattr(child, "_gefen_codebook", None) is None + else child._gefen_codebook.detach().clone() + ) + for child in children + } + step_before = { + id(child): getattr(child, "_gefen_global_step", None) + for child in children + } + scale_before = scaler.get_scale() + + message = None + try: + scaler.step(optimizer) + except RuntimeError as exc: + message = str(exc) + + params_unchanged = all( + torch.equal(local_clone(param), expected) + for param, expected in zip(params, param_before) + ) + grads_unchanged = all( + (param.grad is None and expected is None) + or ( + param.grad is not None + and expected is not None + and torch.equal(local_clone(param.grad), expected) + ) + for param, expected in zip(params, grad_before) + ) + states_unchanged = all( + _state_equal(child.state.get(param, {}), state_before[id(param)]) + for child in children + for group in child.param_groups + for param in group["params"] + ) + codebooks_unchanged = all( + ( + codebook_before[id(child)] is None + and getattr(child, "_gefen_codebook", None) is None + ) + or ( + codebook_before[id(child)] is not None + and getattr(child, "_gefen_codebook", None) is not None + and torch.equal( + child._gefen_codebook, codebook_before[id(child)] + ) + ) + for child in children + ) + steps_unchanged = all( + getattr(child, "_gefen_global_step", None) == step_before[id(child)] + for child in children + ) + results.append( + { + "kind": kind, + "message": message, + "clear_error": message is not None + and "GradScaler integration requires identical DTensor gradient presence" + in message + and "{} (1/2 mesh ranks have gradients)".format( + ( + "bias" + if kind == "hybrid_mixed" + else "other" + if kind == "gefen_distinct_mesh" + else "weight" + ) + ) + in message, + "params_unchanged": params_unchanged, + "grads_unchanged": grads_unchanged, + "states_unchanged": states_unchanged, + "codebooks_unchanged": codebooks_unchanged, + "steps_unchanged": steps_unchanged, + "scale_unchanged": scaler.get_scale() == scale_before, + } + ) + for param in params: + param.grad = None + dist.barrier() + + result_queue.put(("result", rank, results)) + except Exception: + result_queue.put(("error", rank, traceback.format_exc())) + finally: + if dist.is_initialized(): + dist.destroy_process_group() + + +def test_nonsharded_step_never_collects(monkeypatch): + param = nn.Parameter(torch.randn(4, 4)) + optimizer = GefenMuon([("weight", param)], fused=False) + monkeypatch.setattr(optimizer, "_dist_available", lambda: True) + + def unexpected_collective(*args, **kwargs): + raise AssertionError("plain parameters entered the DTensor preflight") + + monkeypatch.setattr(torch.distributed, "all_reduce", unexpected_collective) + param.grad = torch.randn_like(param) + optimizer.step() + + +@pytest.mark.skipif( + not torch.distributed.is_available() or not torch.distributed.is_gloo_available(), + reason="DTensor GradScaler protocol regression needs Gloo", +) +def test_dtensor_fp16_grad_presence_mismatch_fails_before_grad_scaler_scan(): + import torch.multiprocessing as mp + + world = 2 + context = mp.get_context("spawn") + result_queue = context.Queue() + port = _free_port() + processes = [ + context.Process( + target=_distributed_amp_mismatch_worker, + args=(rank, world, port, result_queue), + ) + for rank in range(world) + ] + for process in processes: + process.start() + + messages = [] + try: + for _ in range(world): + messages.append(result_queue.get(timeout=45)) + except queue.Empty: + pass + finally: + for process in processes: + process.join(timeout=5) + for process in processes: + if process.is_alive(): + process.terminate() + process.join(timeout=5) + + assert all(process.exitcode == 0 for process in processes), [ + process.exitcode for process in processes + ] + errors = [payload for kind, _, payload in messages if kind == "error"] + assert not errors, "\n".join(errors) + rank_results = { + rank: payload for kind, rank, payload in messages if kind == "result" + } + assert set(rank_results) == {0, 1}, messages + + expected_kinds = { + "gefen", + "gefen_distinct_mesh", + "muon_exact", + "muon_distributed", + "muon_approx", + "hybrid", + "hybrid_mixed", + } + for rank, results in rank_results.items(): + assert {item["kind"] for item in results} == expected_kinds + for item in results: + assert item["clear_error"], (rank, item) + assert item["params_unchanged"], (rank, item) + assert item["grads_unchanged"], (rank, item) + assert item["states_unchanged"], (rank, item) + assert item["codebooks_unchanged"], (rank, item) + assert item["steps_unchanged"], (rank, item) + assert item["scale_unchanged"], (rank, item) + + for case_index in range(len(rank_results[0])): + assert ( + rank_results[0][case_index]["message"] + == rank_results[1][case_index]["message"] + ) + + +@pytest.mark.skipif( + not torch.cuda.is_available() + or torch.cuda.device_count() < 2 + or not torch.distributed.is_available() + or not torch.distributed.is_nccl_available(), + reason="gradient-presence mismatch regression needs two CUDA GPUs and NCCL", +) +def test_sharded_grad_presence_mismatch_fails_on_every_rank(): + import torch.multiprocessing as mp + + world = 2 + context = mp.get_context("spawn") + result_queue = context.Queue() + port = _free_port() + processes = [ + context.Process( + target=_mismatch_worker, + args=(rank, world, port, result_queue), + ) + for rank in range(world) + ] + for process in processes: + process.start() + + messages = [] + try: + for _ in range(world): + messages.append(result_queue.get(timeout=30)) + except queue.Empty: + pass + finally: + for process in processes: + process.join(timeout=5) + for process in processes: + if process.is_alive(): + process.terminate() + process.join(timeout=5) + + assert all(process.exitcode == 0 for process in processes), [ + process.exitcode for process in processes + ] + errors = [payload for kind, _, payload in messages if kind == "error"] + assert not errors, "\n".join(errors) + rank_results = { + rank: payload for kind, rank, payload in messages if kind == "result" + } + assert set(rank_results) == {0, 1}, messages + + expected_cases = { + (mode, phase) + for mode in ("exact", "distributed") + for phase in ("first_step", "initialized") + } + for rank, results in rank_results.items(): + assert {(item["mode"], item["phase"]) for item in results} == expected_cases + for item in results: + assert item["clear_error"], (rank, item) + assert item["param_unchanged"], (rank, item) + assert item["state_unchanged"], (rank, item) + assert item["codebook_unchanged"], (rank, item) + assert item["global_step_unchanged"], (rank, item) + + # The all-reduced counts and stable names make the diagnostic identical on + # every mesh rank, which is what lets the job fail synchronously. + for case_index in range(len(rank_results[0])): + assert ( + rank_results[0][case_index]["message"] + == rank_results[1][case_index]["message"] + ) diff --git a/tests/test_step_preflight_atomicity.py b/tests/test_step_preflight_atomicity.py new file mode 100644 index 0000000..3027522 --- /dev/null +++ b/tests/test_step_preflight_atomicity.py @@ -0,0 +1,246 @@ +"""Step preflight and periodic-codebook refresh transaction atomicity.""" + +import copy +import warnings + +import pytest +import torch +import torch.nn as nn + +from gefen import Gefen, GefenMuon, GefenMuonHybrid +import gefen.gefen as gefen_module + + +def _clone(value): + if torch.is_tensor(value): + return value.detach().clone() + if isinstance(value, dict): + return {key: _clone(item) for key, item in value.items()} + if isinstance(value, list): + return [_clone(item) for item in value] + if isinstance(value, tuple): + return tuple(_clone(item) for item in value) + return copy.deepcopy(value) + + +def _assert_equal(actual, expected): + if torch.is_tensor(expected): + assert torch.is_tensor(actual) + assert actual.dtype == expected.dtype + assert actual.layout == expected.layout + if actual.layout == torch.strided: + assert torch.equal(actual, expected) + else: + assert torch.equal(actual.to_dense(), expected.to_dense()) + return + if isinstance(expected, dict): + assert isinstance(actual, dict) + assert set(actual) == set(expected) + for key in expected: + _assert_equal(actual[key], expected[key]) + return + if isinstance(expected, (list, tuple)): + assert isinstance(actual, type(expected)) + assert len(actual) == len(expected) + for actual_item, expected_item in zip(actual, expected): + _assert_equal(actual_item, expected_item) + return + assert actual == expected + + +def _gefen_children(optimizer): + if isinstance(optimizer, GefenMuonHybrid): + return [ + child + for child in (optimizer.muon, optimizer.backup) + if isinstance(child, Gefen) + ] + return [optimizer] if isinstance(optimizer, Gefen) else [] + + +def _snapshot(optimizer, params): + runtime = [] + for child in _gefen_children(optimizer): + runtime.append( + { + "global_step": child._gefen_global_step, + "codebook": _clone(child._gefen_codebook), + "codebook_by_device": _clone(child._gefen_codebook_by_device), + "codebook_lut_by_device": _clone( + child._gefen_codebook_lut_by_device + ), + "sr_seed_by_device": _clone(child._sr_seed_by_device), + } + ) + return { + "params": [ + { + "layout": param.layout, + "dense": param.detach().to_dense().clone() + if param.layout != torch.strided + else param.detach().clone(), + } + for param in params + ], + "state_dict": _clone(optimizer.state_dict()), + "runtime": runtime, + } + + +def _assert_snapshot(optimizer, params, expected): + for param, saved in zip(params, expected["params"]): + assert param.layout == saved["layout"] + actual = param.detach().to_dense() if param.layout != torch.strided else param + assert torch.equal(actual, saved["dense"]) + _assert_equal(optimizer.state_dict(), expected["state_dict"]) + runtime = [] + for child in _gefen_children(optimizer): + runtime.append( + { + "global_step": child._gefen_global_step, + "codebook": _clone(child._gefen_codebook), + "codebook_by_device": _clone(child._gefen_codebook_by_device), + "codebook_lut_by_device": _clone( + child._gefen_codebook_lut_by_device + ), + "sr_seed_by_device": _clone(child._sr_seed_by_device), + } + ) + _assert_equal(runtime, expected["runtime"]) + + +def _invalid_tensor(layout, shape=(8, 8), device="cpu"): + dense = torch.eye(*shape, device=device) + if layout == "coo": + return dense.to_sparse() + if layout == "csr": + with warnings.catch_warnings(): + warnings.simplefilter("ignore", UserWarning) + return dense.to_sparse_csr() + if layout == "mkldnn": + return dense.to_mkldnn() + raise AssertionError(layout) + + +def _make_case(kind, layout, *, device="cpu", fused=False): + generator = torch.Generator().manual_seed(991) + first = nn.Parameter( + (torch.randn(8, 8, generator=generator) * 0.02).to(device) + ) + invalid = nn.Parameter(_invalid_tensor(layout, device=device)) + if kind == "gefen": + optimizer = Gefen( + [("first", first), ("invalid", invalid)], + lr=2e-3, + fused=fused, + factored_v_2d=False, + ) + elif kind == "muon": + optimizer = GefenMuon( + [("first", first), ("invalid", invalid)], + lr=2e-3, + fused=fused, + ns_steps=1, + ) + elif kind in ("hybrid_gefen", "hybrid_adamw"): + optimizer = GefenMuonHybrid( + [("first", first)], + [("invalid", invalid)], + lr=2e-3, + fused=fused, + ns_steps=1, + ns_schedule="standard", + normuon=False, + backup_optimizer=kind.removeprefix("hybrid_"), + ) + else: + raise AssertionError(kind) + return optimizer, [first, invalid] + + +@pytest.mark.parametrize("phase", ["first_step", "initialized"]) +@pytest.mark.parametrize("layout", ["coo", "csr", "mkldnn"]) +@pytest.mark.parametrize( + "kind", ["gefen", "muon", "hybrid_gefen", "hybrid_adamw"] +) +def test_later_invalid_gradient_is_rejected_before_any_mutation( + kind, layout, phase +): + optimizer, params = _make_case(kind, layout) + first, invalid = params + if phase == "initialized": + first.grad = torch.full_like(first, 0.002) + invalid.grad = None + optimizer.step() + optimizer.zero_grad(set_to_none=True) + + first.grad = torch.full_like(first, -0.003) + invalid.grad = _invalid_tensor(layout) + before = _snapshot(optimizer, params) + with pytest.raises(RuntimeError, match="sparse|strided|layout"): + optimizer.step() + _assert_snapshot(optimizer, params, before) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is unavailable") +@pytest.mark.parametrize("fused", [False, True], ids=["unfused", "fused"]) +@pytest.mark.parametrize( + "kind", ["gefen", "muon", "hybrid_gefen", "hybrid_adamw"] +) +def test_cuda_initialized_preflight_is_atomic_for_fused_and_unfused(kind, fused): + optimizer, params = _make_case( + kind, "coo", device="cuda", fused=fused + ) + first, invalid = params + first.grad = torch.full_like(first, 0.002) + optimizer.step() + optimizer.zero_grad(set_to_none=True) + + first.grad = torch.full_like(first, -0.003) + invalid.grad = _invalid_tensor("coo", device="cuda") + before = _snapshot(optimizer, params) + with pytest.raises(RuntimeError, match="sparse|strided|layout"): + optimizer.step() + _assert_snapshot(optimizer, params, before) + + +def test_periodic_codebook_refresh_stages_every_index_before_commit(monkeypatch): + generator = torch.Generator().manual_seed(997) + params = [ + nn.Parameter(torch.randn(8, 8, generator=generator) * 0.02) + for _ in range(2) + ] + optimizer = Gefen( + [("first", params[0]), ("second", params[1])], + lr=2e-3, + fused=False, + factored_v_2d=False, + codebook_refresh_every=1, + ) + for index, param in enumerate(params): + param.grad = torch.full_like(param, 0.001 * (index + 1)) + optimizer.step() + optimizer.zero_grad(set_to_none=True) + + first_indices = optimizer.state[params[0]]["m_codebook"].detach().clone() + calls = 0 + + def fail_on_second(codebook, coefficients): + nonlocal calls + calls += 1 + if calls == 1: + return (255 - first_indices.to(torch.int16)).to(torch.uint8) + raise RuntimeError("injected second-parameter requantization failure") + + monkeypatch.setattr( + gefen_module, "gefen_nearest_codebook_indices", fail_on_second + ) + for index, param in enumerate(params): + param.grad = torch.full_like(param, -0.003 * (index + 1)) + before = _snapshot(optimizer, params) + + with pytest.raises(RuntimeError, match="second-parameter requantization"): + optimizer.step() + + assert calls == 2 + _assert_snapshot(optimizer, params, before) diff --git a/tests/test_training_matrix_harness.py b/tests/test_training_matrix_harness.py index 6d6b63f..acd402b 100644 --- a/tests/test_training_matrix_harness.py +++ b/tests/test_training_matrix_harness.py @@ -399,6 +399,8 @@ def test_cell_parsers_reject_flag_abbreviations(parse, extra): @pytest.mark.parametrize("cell", ("adamw", "torch_muon_adamw")) def test_stock_cells_record_batched_ns_request_as_unsupported(cell): + if cell == "torch_muon_adamw" and not hasattr(torch.optim, "Muon"): + pytest.skip("torch.optim.Muon requires PyTorch 2.9+") optimizer, resolved = build_optimizer( _model(), cell, @@ -813,8 +815,8 @@ def _hf_datasets_cache_dir() -> Path: @pytest.mark.skipif( - TINY_SHAKESPEARE_REVISION_CACHE is None, - reason="the pinned Tiny Shakespeare revision is an optional local integration asset", + TINY_SHAKESPEARE_REVISION_CACHE is None or importlib.util.find_spec("datasets") is None, + reason="the pinned Tiny Shakespeare revision and datasets package are optional local integration assets", ) def test_cached_tiny_shakespeare_hash_groups_are_disjoint(): bundle = build_dataset( @@ -872,8 +874,8 @@ def test_single_text_document_uses_exact_contiguous_byte_ranges(tmp_path): @pytest.mark.skipif( - not WIKITEXT_DOCUMENT_CACHE.exists(), - reason="pinned document-level WikiText is an optional local integration asset", + not WIKITEXT_DOCUMENT_CACHE.exists() or importlib.util.find_spec("datasets") is None, + reason="pinned document-level WikiText and the datasets package are optional local integration assets", ) def test_pinned_wikitext_official_split_provenance(): bundle = build_dataset( diff --git a/tests/test_transformers_trainer_resume.py b/tests/test_transformers_trainer_resume.py new file mode 100644 index 0000000..8d4ae3f --- /dev/null +++ b/tests/test_transformers_trainer_resume.py @@ -0,0 +1,73 @@ +from __future__ import annotations + +import argparse +import copy + +import pytest +import torch + +from benchmarks.trainer_resume import run as harness + + +def test_state_digest_is_structural_and_tensor_exact(): + left = {"b": [torch.tensor([1.0, 2.0]), 3], "a": (True, None)} + reordered = {"a": (True, None), "b": [torch.tensor([1.0, 2.0]), 3]} + changed = copy.deepcopy(left) + changed["b"][0][1] = 2.0001 + + assert harness.state_digest(left) == harness.state_digest(reordered) + assert harness.state_digest(left) != harness.state_digest(changed) + + +def test_recipe_parser_rejects_unknown_and_duplicate_cells(): + assert harness._recipe_list("gefen,gefen_muon_adamw") == ("gefen", "gefen_muon_adamw") + with pytest.raises(argparse.ArgumentTypeError, match="unknown recipes"): + harness._recipe_list("gefen,adamw") + with pytest.raises(argparse.ArgumentTypeError, match="duplicates"): + harness._recipe_list("gefen,gefen") + + +def test_fixed_token_dataset_is_reproducible(): + first = harness.FixedTokenDataset(samples=4, seq_len=8, vocab_size=32, seed=9) + second = harness.FixedTokenDataset(samples=4, seq_len=8, vocab_size=32, seed=9) + assert torch.equal(first.input_ids, second.input_ids) + row = first[0] + assert torch.equal(row["input_ids"], row["labels"]) + assert torch.equal(row["attention_mask"], torch.ones(8, dtype=torch.long)) + + +def test_trainer_all_recipes_exact_resume_cpu(tmp_path): + pytest.importorskip("transformers") + pytest.importorskip("accelerate") + args = harness.parse_args( + [ + "--output-dir", + str(tmp_path), + "--device", + "cpu", + "--dtype", + "float32", + "--no-fused", + "--steps", + "3", + "--split-step", + "1", + "--warmup-steps", + "1", + "--batch-size", + "1", + "--gradient-accumulation-steps", + "2", + "--seq-len", + "16", + "--ns-steps", + "1", + ] + ) + summary = harness.run(args) + assert summary["passed"] is True + assert set(summary["recipes"]) == set(harness.RECIPES) + for result in summary["recipes"].values(): + assert result["baseline"]["model_sha256"] == result["resumed"]["model_sha256"] + assert result["baseline"]["optimizer_sha256"] == result["resumed"]["optimizer_sha256"] + assert result["baseline"]["scheduler_sha256"] == result["resumed"]["scheduler_sha256"] From eaf57d0c6c2e7e4c9319a5164eac2cc74ecaf80f Mon Sep 17 00:00:00 2001 From: thad0ctor Date: Sun, 12 Jul 2026 05:35:02 -0700 Subject: [PATCH 05/14] Order DTensor AMP preflight collectives --- src/gefen/gefen.py | 28 +++++++++++++++++++++++----- 1 file changed, 23 insertions(+), 5 deletions(-) diff --git a/src/gefen/gefen.py b/src/gefen/gefen.py index ee226c7..25cd23b 100644 --- a/src/gefen/gefen.py +++ b/src/gefen/gefen.py @@ -596,8 +596,6 @@ def _amp_dtensor_protocol_preflight(optimizer) -> bool: """ if not torch.distributed.is_available() or not torch.distributed.is_initialized(): return False - if torch.cuda.is_available() and torch.cuda.is_current_stream_capturing(): - return False import torch.distributed as dist @@ -623,6 +621,7 @@ def _amp_dtensor_protocol_preflight(optimizer) -> bool: mesh = param.device_mesh if mesh.get_coordinate() is None or mesh.size() < 2: continue + process_groups = tuple(mesh.get_all_groups()) key = ( str(mesh.device_type), tuple(int(item) for item in mesh.shape), @@ -630,10 +629,15 @@ def _amp_dtensor_protocol_preflight(optimizer) -> bool: int(item) for item in mesh.mesh.detach().cpu().reshape(-1).tolist() ), + tuple(str(group.group_name) for group in process_groups), ) entry = by_mesh.get(key) if entry is None: - entry = {"mesh": mesh, "items": []} + entry = { + "mesh": mesh, + "process_groups": process_groups, + "items": [], + } by_mesh[key] = entry name = ( names[index] @@ -646,9 +650,23 @@ def _amp_dtensor_protocol_preflight(optimizer) -> bool: if not by_mesh: return False + # Do not initialize or query CUDA for a CPU-only DTensor optimizer. Spawned + # Gloo workers may run alongside a CUDA-heavy parent release gate, and an + # unrelated current-stream query can otherwise block before the CPU + # collective preflight. CUDA-backed optimizers retain the capture guard. + if any(param.device.type == "cuda" for param in params) and ( + torch.cuda.is_current_stream_capturing() + ): + return False mismatches = [] - for entry in by_mesh.values(): + # Parameter-group order may legitimately differ between ranks while two + # equivalent DeviceMesh objects still refer to distinct c10d groups. Sort + # by each mesh's process-group names so every rank enters those collectives + # in creation order instead of whichever parameter happened to appear + # first locally. + for mesh_key in sorted(by_mesh): + entry = by_mesh[mesh_key] mesh = entry["mesh"] items = sorted( entry["items"], @@ -666,7 +684,7 @@ def _amp_dtensor_protocol_preflight(optimizer) -> bool: dtype=torch.int32, device=local.device, ) - for process_group in mesh.get_all_groups(): + for process_group in entry["process_groups"]: if dist.get_world_size(process_group) > 1: dist.all_reduce( active_counts, op=dist.ReduceOp.SUM, group=process_group From 73ff12b0f0b054576bb8e50ff701d973bd9aae56 Mon Sep 17 00:00:00 2001 From: thad0ctor Date: Sun, 12 Jul 2026 14:25:25 -0700 Subject: [PATCH 06/14] Fix v2 single-process resume and order-sensitive presence preflight Consolidated version-2 Parallel-Muon checkpoints refused to load into an optimizer with no eligible distributed process group (single-process resume/eval), even though released-v1 and markerless loads of the same state succeed; rebind the saved owner manifest to live parameters by saved id and run the same per-parameter proof. Harden the Muon grad-presence preflight the way the AMP DTensor preflight already is: content-keyed meshes, sorted mesh and item order so rank-divergent parameter-group order cannot falsely pass the positional activity vector or interleave per-mesh collectives, and no CUDA query for CPU-only mesh optimizers. Refresh the stale codebook-fallback comment and disclose the fused=False CUDA backend change, the eps>0 constructor requirement, and the quantized-momentum codebook refusal in the changelog. Both new gloo regression tests fail on the previous code. --- CHANGELOG.md | 6 +- README.md | 11 +- src/gefen/gefen.py | 12 +- src/gefen/gefen_muon.py | 91 +++++- tests/test_deterministic_mode.py | 29 ++ ...test_muon_distributed_checkpoint_safety.py | 173 +++++++++++ tests/test_muon_grad_presence.py | 286 ++++++++++++++++++ 7 files changed, 589 insertions(+), 19 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index c6514c8..ce0814e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,12 +8,14 @@ Correctness and compatibility: - Add `deterministic=True` to `Gefen`, `GefenMuon`, and `GefenMuonHybrid` for replica-exact fused routing on homogeneous GPUs. Automatic periods use fixed-order reductions, block-vmean parameters use the deterministic fused v1 path, factored-v parameters use the decomposed deterministic update, and tagged checkpoints enforce the saved policy. - Capturable optimizers maintain device-resident global-step counters on every parameter device. CUDA-graph replays now serialize the true global step, including steps with no gradients, so stochastic-rounding checkpoints resume with the correct seed. -- Checkpoint loading preserves compact optimizer-state dtypes for bf16 parameters, validates frozen codebooks and hybrid backend metadata, and keeps safe legacy untagged native checkpoints loadable. +- Checkpoint loading preserves compact optimizer-state dtypes for bf16 parameters, validates frozen codebooks and hybrid backend metadata, and keeps safe legacy untagged native checkpoints loadable. A checkpoint whose quantized momentum lost its frozen codebook — for example a pre-0.4.0 FSDP/DCP optim-state consolidation round-trip that stripped Gefen's custom top-level keys — is now refused with a clear error instead of silently relearning a codebook against the restored uint8 indices; resume such runs from the unmodified native checkpoint, or re-save with 0.4.0, whose per-group metadata mirror survives those round-trips. - Plain Gefen and `GefenMuon(sharded_mode="approx")` collectively encode rank-local DTensor optimizer state for ordinary and flattened PyTorch full-state DCP. Two-GPU FSDP2 `fully_shard` get/set tests resume both optimizers' next update exactly on the same one-dimensional default-world mesh; all ranks must participate, per-process CPU checkpoint memory approaches `world_size ×` local serialized state plus local scratch, topology changes fail closed, and old unsafe untagged full checkpoints are rejected. - Add a Transformers Trainer checkpoint-continuation gate for plain Gefen, GefenMuon+AdamW, and GefenMuon+Gefen, covering Trainer's internal Accelerate wrapper, tied weights, gradient accumulation, changing LRs, fused BF16 DDP replica hashes, and exact deterministic resume state. - Integrate PyTorch GradScaler without changing the ordinary FP32-master/BF16 Accelerate path: true-FP16 gradients use the optimizer-side protocol, overflow skips every GefenMuonHybrid child atomically, DTensor overflow decisions remain synchronized across ranks, and FSDP1 FP16 directs callers to `ShardedGradScaler`. - Reject rank-divergent `grad is None` patterns before exact/distributed DTensor Muon collectives, avoiding an unmatched-collective deadlock without adding collectives to plain/DDP/`approx` steps. -- Parallel-Muon `distributed` checkpoints now carry a versioned saved-world ownership manifest, validate full owner-state geometry before mutation, preserve released consolidated-v1 and cross-world-size resume, and reject markerless, partial, or internally inconsistent populated state instead of warning and risking divergent momentum. +- Parallel-Muon `distributed` checkpoints now carry a versioned saved-world ownership manifest, validate full owner-state geometry before mutation, preserve released consolidated-v1 and cross-world-size resume — including resume into a single-process or otherwise non-distributed optimizer — and reject markerless, partial, or internally inconsistent populated state instead of warning and risking divergent momentum. +- On CUDA, `fused=False` now keeps period search and codebook-histogram learning on the pure-PyTorch backends instead of crossing the fused kernels' lazy-JIT boundary. Fresh `fused=False` runs may resolve near-tie block periods or learned codebooks differently than 0.3.x; resumed checkpoints keep their restored periods, and an explicit `FIND_PERIOD_BACKEND` override is still honored outside deterministic mode. +- `eps` must now be finite and strictly positive at construction; `eps=0` was previously accepted. Checkpoints whose parameter groups carry `eps=0` still load unchanged. - Preflight every active gradient before AMP, codebook work, or child dispatch so a later sparse, compressed, MKLDNN, complex, or malformed gradient cannot leave an earlier parameter partially updated; periodic codebook refresh likewise stages every replacement index before committing shared state. - Reject host-driven gradient-histogram output under `capturable=True`, matching the existing periodic-codebook-refresh guard. diff --git a/README.md b/README.md index 6148330..17abd6c 100644 --- a/README.md +++ b/README.md @@ -116,6 +116,15 @@ One knock-on effect: weight decay in AdamW-style optimizers is applied as `lr × Gefen drops into standard distributed training like any other PyTorch optimizer, with either `fused=True` or `fused=False`. Validated training setups include single-GPU, PyTorch DDP, FSDP2 (`fully_shard` / DTensor), and DeepSpeed ZeRO 1-3 (plain `Gefen` as the client optimizer, direct or via axolotl `gefenx`; bit-exact ZeRO-2 checkpoint resume). FSDP2 optimizer checkpoint support is mode- and topology-specific as described below. +| System | Current status | +|---|---| +| DDP | Supported and tested, including fused BF16 resume | +| FSDP2 / DTensor training | Tested for plain Gefen and Muon `approx` / `exact` / `distributed` | +| FSDP2 full-state checkpoints (DCP) | Plain Gefen and Muon `approx` only; same 1-D topology and world size — scope note below | +| Muon `distributed` native checkpoints | Versioned owner manifest; collective save resumes across world sizes, including into a single-process optimizer — [details](#experimental-lever-sharded-newton-schulz-under-fsdp2-sharded_mode) | +| DeepSpeed ZeRO 1-3 | Plain Gefen supported and tested, direct and via axolotl `gefenx`; Muon/Hybrid fail fast with a clear error — config note below | +| Megatron DP/TP/PP/CP/EP/ETP | Tested through the Megatron GPT pretraining entry point with legacy optimizer checkpoints — [scope](https://github.com/thad0ctor/Gefen-X/blob/main/COMPATIBILITY.md) | + > **FSDP2 optimizer checkpoint scope.** Plain Gefen and `GefenMuon(sharded_mode="approx")` collectively encode every rank's local DTensor optimizer state into PyTorch DCP `StateDictOptions(full_state_dict=True)` output, including the flattened optimizer-state form. `get_optimizer_state_dict()` and `set_optimizer_state_dict()` resume the next update exactly when world size, mesh, placements, rank coordinates, parameter ordering, shapes, names, and sharded mode are unchanged; the actual two-GPU `fully_shard` get/set test covers both optimizers. The adapter currently requires one 1-D DeviceMesh spanning the default world; multidimensional meshes, subgroups, and pipeline-local optimizers fail before its collectives. Save and restore are collective, so every rank must participate. Each process temporarily stages all serialized rank payloads on CPU, making the leading checkpoint-time CPU cost about `world_size ×` that rank's local optimizer state plus local serialization scratch. World-size or topology changes fail before mutation, and older unsafe untagged full checkpoints fail closed. This is not a reshardable optimizer-state format. `torch.amp.GradScaler` keeps PyTorch's ordinary externally skipped step for FP32-master and BF16 training, including Trainer/Accelerate gradient clipping and scheduler behavior. Actual FP16 gradient storage opts into PyTorch's native optimizer-side scaling protocol so finite gradients are unscaled once and overflow returns before either Hybrid child, codebook, state, parameter, or counter changes. DTensor/FSDP2 non-finite flags are reduced across the mesh; FSDP1 FlatParameters must use `torch.distributed.fsdp.ShardedGradScaler`. @@ -612,7 +621,7 @@ opt = GefenMuonHybrid( # only takes effect under FSDP2 (DTensor params); no-op single-GPU ``` -> **`"distributed"` checkpointing is collective.** `state_dict()` gathers each owner's momentum across ranks, so **every rank must call it** (as in a standard FSDP full-state-dict flow). Calling `state_dict()` on rank 0 only — e.g. a rank-0-only save loop — **deadlocks**. Save and load also transiently materialize the full unsharded momentum on every rank, so peak memory at checkpoint time approaches `"exact"` mode's. +> **`"distributed"` checkpointing is collective.** `state_dict()` gathers each owner's momentum across ranks, so **every rank must call it** (as in a standard FSDP full-state-dict flow). Calling `state_dict()` on rank 0 only — e.g. a rank-0-only save loop — **deadlocks**. Save and load also transiently materialize the full unsharded momentum on every rank, so peak memory at checkpoint time approaches `"exact"` mode's. The saved checkpoint carries a versioned owner manifest and is complete on every rank, so it resumes under a different world size or in a single-process optimizer; populated state whose manifest is missing, partial, or inconsistent is rejected before any mutation. ![Gefen-Muon exact / distributed / approx sharded — eval loss](https://raw.githubusercontent.com/thad0ctor/Gefen-X/main/docs/benchmarks/muon_shard_loss.png) ![Gefen-Muon exact / distributed / approx sharded — throughput & VRAM](https://raw.githubusercontent.com/thad0ctor/Gefen-X/main/docs/benchmarks/muon_shard_perf.png) diff --git a/src/gefen/gefen.py b/src/gefen/gefen.py index 25cd23b..6cfd3e5 100644 --- a/src/gefen/gefen.py +++ b/src/gefen/gefen.py @@ -4360,12 +4360,12 @@ def _compact(value): state_dict["gefen_deterministic"] = self._deterministic # The exact-DP codebook is learned once on the first step and then frozen # for the rest of the run. It is not per-param state, so persist it - # explicitly; without it resume re-learns the codebook (see - # _maybe_refresh_gefen_codebook) which re-predicts and overwrites every - # restored automatic_period, desyncing them from the saved - # vmean/m_codebook. (Note: under FSDP optim-state consolidation strips - # these custom top-level keys; _maybe_refresh_gefen_codebook handles that - # fallback by reusing the restored periods.) + # explicitly; without it resume would reinterpret the restored uint8 + # m_codebook indices against a freshly learned codebook. (Note: FSDP/DCP + # optim-state consolidation strips custom top-level keys; the per-group + # _gefen_checkpoint_metadata mirror below carries the codebook through + # that round-trip, and load_state_dict refuses quantized-momentum + # checkpoints that lost both copies rather than relearning.) state_dict["gefen_codebook"] = self._gefen_codebook # PyTorch Distributed Checkpoint's ``get_optimizer_state_dict`` keeps # only the conventional ``state`` and ``param_groups`` top-level keys. diff --git a/src/gefen/gefen_muon.py b/src/gefen/gefen_muon.py index 8bb3b26..4196895 100644 --- a/src/gefen/gefen_muon.py +++ b/src/gefen/gefen_muon.py @@ -1315,7 +1315,16 @@ def _assert_sharded_grad_presence_consistent(self) -> None: """ if not self._dist_available(): return - if torch.cuda.is_available() and torch.cuda.is_current_stream_capturing(): + # Do not initialize or query CUDA for a CPU-only mesh optimizer. + # Spawned Gloo workers may run alongside a CUDA-heavy parent process, + # and an unrelated current-stream query can otherwise block before the + # CPU collective preflight. CUDA-backed optimizers retain the capture + # guard. + if any( + param.device.type == "cuda" + for group in self.param_groups + for param in group["params"] + ) and torch.cuda.is_current_stream_capturing(): return import torch.distributed as dist @@ -1330,21 +1339,44 @@ def _assert_sharded_grad_presence_consistent(self) -> None: mesh = p.device_mesh if mesh.get_coordinate() is None or mesh.size() < 2: continue - # All parameters created from one FSDP2/DTensor mesh normally - # share the same DeviceMesh object. Key by identity so exotic - # optimizers carrying distinct meshes run one matched vector - # collective per mesh without conflating their parameter order. - key = id(mesh) + # Parameter-group order may legitimately differ between ranks + # while two equivalent DeviceMesh objects still refer to + # distinct c10d groups. Key each mesh by content instead of + # object identity, and sort meshes and items below, so every + # rank compares the same positional activity vector and enters + # the per-mesh collectives in the same order. + process_groups = tuple(mesh.get_all_groups()) + key = ( + str(mesh.device_type), + tuple(int(item) for item in mesh.shape), + tuple( + int(item) + for item in mesh.mesh.detach().cpu().reshape(-1).tolist() + ), + tuple(str(pg.group_name) for pg in process_groups), + ) entry = by_mesh.get(key) if entry is None: - entry = {"mesh": mesh, "items": []} + entry = { + "mesh": mesh, + "process_groups": process_groups, + "items": [], + } by_mesh[key] = entry entry["items"].append((str(name), p, p.grad is not None)) mismatches = [] - for entry in by_mesh.values(): + for mesh_key in sorted(by_mesh): + entry = by_mesh[mesh_key] mesh = entry["mesh"] - items = entry["items"] + items = sorted( + entry["items"], + key=lambda item: ( + item[0], + tuple(item[1].shape), + str(item[1].dtype), + ), + ) device = self._state_tensor_device(items[0][1]) active_counts = torch.tensor( [int(active) for _, _, active in items], @@ -1355,7 +1387,7 @@ def _assert_sharded_grad_presence_consistent(self) -> None: # dimension propagates a global activity count to every mesh rank. # This is one collective per mesh dimension, independent of the # number of optimizer parameters. - for process_group in mesh.get_all_groups(): + for process_group in entry["process_groups"]: if dist.get_world_size(process_group) > 1: dist.all_reduce( active_counts, op=dist.ReduceOp.SUM, group=process_group @@ -2328,6 +2360,45 @@ def _validate_distributed_checkpoint_load(self, state_dict, marker) -> None: # Bind the proof to those same ordered process-group partitions instead # of requiring it to cover every parameter in the optimizer group. expected_pg_groups = list(by_pg.values()) + if not expected_pg_groups and marker_groups: + # A consolidated version-2 checkpoint carries the complete owner + # state on every rank, so it must stay loadable when the live + # optimizer has no Parallel-Muon-eligible process group at all + # (single-process resume/eval, uninitialized torch.distributed, or + # an all-fallback topology) — released-v1 and markerless loads of + # the same state already support exactly that. The manifest cannot + # bind to live process groups then; rebind each saved partition to + # the live parameters by saved id and run the identical + # per-parameter owner-state proof below on the rebound partitions. + items_by_saved_id = {item[2]: item for item in expected_items} + rebound_ids = set() + rebound_groups = [] + for group_index, saved_group in enumerate(marker_groups): + group_items = [] + saved_params = ( + saved_group.get("params") + if isinstance(saved_group, dict) + else None + ) + if isinstance(saved_params, (list, tuple)): + for saved_param in saved_params: + if not isinstance(saved_param, dict): + continue + saved_id = saved_param.get("saved_id") + try: + item = items_by_saved_id.get(saved_id) + except TypeError: + item = None + if item is None or saved_id in rebound_ids: + raise self._distributed_checkpoint_error( + "owner manifest group {} saved id {!r} does not " + "map to exactly one live distributed " + "parameter".format(group_index, saved_id) + ) + rebound_ids.add(saved_id) + group_items.append(item) + rebound_groups.append(group_items) + expected_pg_groups = rebound_groups if len(marker_groups) != len(expected_pg_groups): raise self._distributed_checkpoint_error( "the owner manifest has {} process-group entries but the live " diff --git a/tests/test_deterministic_mode.py b/tests/test_deterministic_mode.py index 84c2067..cec296d 100644 --- a/tests/test_deterministic_mode.py +++ b/tests/test_deterministic_mode.py @@ -120,6 +120,35 @@ def test_deterministic_checkpoint_mismatch_rejected_before_mutation(): assert _deep_equal(target.state_dict(), before) +def test_nondeterministic_checkpoint_into_deterministic_optimizer_rejected(): + # The reverse policy mismatch of the test above: a deterministic=False + # tagged checkpoint must not silently adopt a deterministic=True live + # policy either, in either the native or the DCP-normalized shape. + _, source = _stepped_cpu_optimizer(deterministic=False) + checkpoint = copy.deepcopy(source.state_dict()) + assert checkpoint["gefen_deterministic"] is False + + target_param = nn.Parameter(torch.zeros(16, 12)) + target = Gefen( + [("weight", target_param)], + lr=1e-3, + fused=False, + deterministic=True, + ) + before = copy.deepcopy(target.state_dict()) + with pytest.raises(ValueError, match="intentional state migration"): + target.load_state_dict(checkpoint) + assert _deep_equal(target.state_dict(), before) + + dcp_style = { + "state": checkpoint["state"], + "param_groups": checkpoint["param_groups"], + } + with pytest.raises(ValueError, match="intentional state migration"): + target.load_state_dict(dcp_style) + assert _deep_equal(target.state_dict(), before) + + @pytest.mark.parametrize("invalid_tag", [None, 0, 1, "true"]) def test_deterministic_checkpoint_top_level_tag_requires_actual_bool(invalid_tag): _, source = _stepped_cpu_optimizer(deterministic=True) diff --git a/tests/test_muon_distributed_checkpoint_safety.py b/tests/test_muon_distributed_checkpoint_safety.py index 7f53b9b..7e2ca0b 100644 --- a/tests/test_muon_distributed_checkpoint_safety.py +++ b/tests/test_muon_distributed_checkpoint_safety.py @@ -2,6 +2,7 @@ import copy from datetime import timedelta +import io import os import queue import socket @@ -837,3 +838,175 @@ def test_parallel_muon_rejection_is_atomic_and_valid_v2_continues(): assert rejection, (rank, message) assert continuation, rank assert version == 2, rank + + +def _single_process_resume_worker(rank, world, port, result_queue): + import torch.distributed as dist + from torch.distributed.tensor import Shard, distribute_tensor, init_device_mesh + + try: + os.environ["MASTER_ADDR"] = "127.0.0.1" + os.environ["MASTER_PORT"] = port + os.environ["RANK"] = str(rank) + os.environ["WORLD_SIZE"] = str(world) + dist.init_process_group( + "gloo", + rank=rank, + world_size=world, + timeout=timedelta(seconds=60), + ) + mesh = init_device_mesh("cpu", (world,)) + specs = (("left", (8, 8)), ("right", (6, 10))) + generator = torch.Generator().manual_seed(9314) + params = [ + ( + name, + nn.Parameter( + distribute_tensor( + (torch.randn(*shape, generator=generator) * 0.02), + mesh, + [Shard(0)], + ) + ), + ) + for name, shape in specs + ] + optimizer = GefenMuon( + params, + lr=2e-3, + fused=False, + ns_steps=1, + sharded_mode="distributed", + ) + grad_generator = torch.Generator().manual_seed(9315) + for _, param in params: + param.grad = distribute_tensor( + torch.randn(param.shape, generator=grad_generator) * 0.002, + mesh, + [Shard(0)], + ) + optimizer.step() + optimizer.zero_grad(set_to_none=True) + checkpoint = _clone(optimizer.state_dict()) + full_values = [ + param.detach().full_tensor().clone() for _, param in params + ] + if rank == 0: + # Serialize instead of queueing live tensors: tensor transport + # over mp.Queue shares memory via file descriptors that die with + # this worker process, racing the parent's read + # (ConnectionResetError on slow CI hosts). Bytes pickle normally. + buffer = io.BytesIO() + torch.save( + {"checkpoint": checkpoint, "full_values": full_values}, buffer + ) + result_queue.put(("checkpoint", rank, buffer.getvalue())) + else: + result_queue.put(("done", rank, None)) + except Exception: + result_queue.put(("error", rank, traceback.format_exc())) + finally: + if dist.is_initialized(): + dist.destroy_process_group() + + +@pytest.mark.skipif( + not torch.distributed.is_available() + or not torch.distributed.is_gloo_available(), + reason="consolidated single-process resume regression needs Gloo", +) +def test_consolidated_v2_checkpoint_loads_into_single_process_optimizer(): + """A consolidated distributed save must resume without torch.distributed. + + Consolidation exists so every rank's ``state_dict()`` carries the complete + owner state; the natural consumer is a single-process resume or eval run + that never initializes a process group. Released-v1 and markerless loads + already support that, so the version-2 manifest must not be stricter. + """ + import torch.multiprocessing as mp + + world = 2 + context = mp.get_context("spawn") + result_queue = context.Queue() + port = _free_port() + processes = [ + context.Process( + target=_single_process_resume_worker, + args=(rank, world, port, result_queue), + ) + for rank in range(world) + ] + for process in processes: + process.start() + + messages = [] + try: + for _ in range(world): + messages.append(result_queue.get(timeout=90)) + except queue.Empty: + pass + finally: + for process in processes: + process.join(timeout=10) + for process in processes: + if process.is_alive(): + process.terminate() + process.join(timeout=5) + + assert all(process.exitcode == 0 for process in processes), [ + process.exitcode for process in processes + ] + errors = [payload for kind, _, payload in messages if kind == "error"] + assert not errors, "\n".join(errors) + payloads = [payload for kind, _, payload in messages if kind == "checkpoint"] + assert len(payloads) == 1, messages + saved = torch.load(io.BytesIO(payloads[0]), weights_only=False) + checkpoint, full_values = saved["checkpoint"], saved["full_values"] + + marker = checkpoint.get("gefen_muon_distributed") + assert isinstance(marker, dict) and marker.get("version") == 2, marker + assert not torch.distributed.is_initialized() + + specs = (("left", (8, 8)), ("right", (6, 10))) + params = [ + (name, nn.Parameter(value.clone())) + for (name, _), value in zip(specs, full_values) + ] + optimizer = GefenMuon( + params, + lr=2e-3, + fused=False, + ns_steps=1, + sharded_mode="distributed", + ) + optimizer.load_state_dict(_clone(checkpoint)) + + assert optimizer._gefen_global_step == 1 + saved_ids = [ + saved_id + for group in checkpoint["param_groups"] + for saved_id in group["params"] + ] + for (name, param), saved_id in zip(params, saved_ids): + live_state = optimizer.state[param] + saved_state = checkpoint["state"][saved_id] + for key in ("automatic_period", "step", "m_codebook", "m_magnitude"): + assert key in live_state, (name, key) + expected = saved_state[key] + actual = live_state[key] + if torch.is_tensor(expected): + assert torch.is_tensor(actual) and torch.equal( + actual, expected + ), (name, key) + else: + assert actual == expected, (name, key) + + grad_generator = torch.Generator().manual_seed(9316) + before = [param.detach().clone() for _, param in params] + for _, param in params: + param.grad = torch.randn(param.shape, generator=grad_generator) * 0.002 + optimizer.step() + assert all( + not torch.equal(param.detach(), previous) + for (_, param), previous in zip(params, before) + ) diff --git a/tests/test_muon_grad_presence.py b/tests/test_muon_grad_presence.py index abc2ef3..218b2a0 100644 --- a/tests/test_muon_grad_presence.py +++ b/tests/test_muon_grad_presence.py @@ -552,3 +552,289 @@ def test_sharded_grad_presence_mismatch_fails_on_every_rank(): rank_results[0][case_index]["message"] == rank_results[1][case_index]["message"] ) + + +def _reversed_order_worker(rank, world, port, result_queue): + import torch.distributed as dist + from torch.distributed.tensor import Shard, distribute_tensor, init_device_mesh + + try: + os.environ["MASTER_ADDR"] = "127.0.0.1" + os.environ["MASTER_PORT"] = port + os.environ["RANK"] = str(rank) + os.environ["WORLD_SIZE"] = str(world) + dist.init_process_group( + "gloo", + rank=rank, + world_size=world, + timeout=timedelta(seconds=12), + ) + mesh = init_device_mesh("cpu", (world,)) + results = [] + + for case in ("reversed_mismatch", "reversed_consistent", "aligned_mismatch"): + generator = torch.Generator(device="cpu").manual_seed( + 6100 + len(results) + ) + full_a = torch.randn(8, 8, generator=generator) + full_a_grad = torch.randn(8, 8, generator=generator) * 0.01 + full_b = torch.randn(6, 10, generator=generator) + full_b_grad = torch.randn(6, 10, generator=generator) * 0.01 + a = nn.Parameter(distribute_tensor(full_a.clone(), mesh, [Shard(0)])) + b = nn.Parameter(distribute_tensor(full_b.clone(), mesh, [Shard(0)])) + named = [("a", a), ("b", b)] + if case.startswith("reversed") and rank == 1: + named = list(reversed(named)) + optimizer = GefenMuon( + named, + lr=1e-3, + fused=False, + ns_steps=1, + sharded_mode="exact", + ) + # distribute_tensor is itself collective, so materialize both + # gradients on every rank first, then drop one per rank to build + # the divergent presence pattern. + a.grad = distribute_tensor(full_a_grad.clone(), mesh, [Shard(0)]) + b.grad = distribute_tensor(full_b_grad.clone(), mesh, [Shard(0)]) + if case.endswith("mismatch"): + if rank == 0: + b.grad = None + else: + a.grad = None + + # Call the preflight directly instead of step(): before the + # order-insensitivity fix, the reversed_mismatch case falsely + # passed here and the corresponding step() would deadlock inside + # full_tensor(), so the isolated call is what keeps a regression + # loud instead of hung. + message = None + try: + optimizer._assert_sharded_grad_presence_consistent() + except RuntimeError as exc: + message = str(exc) + results.append({"case": case, "message": message}) + a.grad = None + b.grad = None + dist.barrier() + + result_queue.put(("result", rank, results)) + except Exception: + result_queue.put(("error", rank, traceback.format_exc())) + finally: + if dist.is_initialized(): + dist.destroy_process_group() + + +@pytest.mark.skipif( + not torch.distributed.is_available() or not torch.distributed.is_gloo_available(), + reason="order-insensitive grad-presence regression needs Gloo", +) +def test_sharded_grad_presence_check_is_order_insensitive(): + """Rank-divergent param-group order must not defeat the presence check. + + The activity vector is compared positionally across ranks, so without + content-keyed meshes and sorted items a reversed group order with + complementary gradient presence sums to a full count on every position + and the check falsely passes right before full_tensor() deadlocks. + """ + import torch.multiprocessing as mp + + world = 2 + context = mp.get_context("spawn") + result_queue = context.Queue() + port = _free_port() + processes = [ + context.Process( + target=_reversed_order_worker, + args=(rank, world, port, result_queue), + ) + for rank in range(world) + ] + for process in processes: + process.start() + + messages = [] + try: + for _ in range(world): + messages.append(result_queue.get(timeout=45)) + except queue.Empty: + pass + finally: + for process in processes: + process.join(timeout=5) + for process in processes: + if process.is_alive(): + process.terminate() + process.join(timeout=5) + + assert all(process.exitcode == 0 for process in processes), [ + process.exitcode for process in processes + ] + errors = [payload for kind, _, payload in messages if kind == "error"] + assert not errors, "\n".join(errors) + rank_results = { + rank: {item["case"]: item["message"] for item in payload} + for kind, rank, payload in messages + if kind == "result" + } + assert set(rank_results) == {0, 1}, messages + + for rank, cases in rank_results.items(): + assert set(cases) == { + "reversed_mismatch", + "reversed_consistent", + "aligned_mismatch", + }, (rank, cases) + for case in ("reversed_mismatch", "aligned_mismatch"): + message = cases[case] + assert message is not None, (rank, case) + assert "identical gradient presence" in message, (rank, case, message) + assert "a (1/2 mesh ranks have gradients)" in message, ( + rank, + case, + message, + ) + assert "b (1/2 mesh ranks have gradients)" in message, ( + rank, + case, + message, + ) + assert cases["reversed_consistent"] is None, ( + rank, + cases["reversed_consistent"], + ) + + # Sorted item names make the diagnostic identical on every rank. + for case in ("reversed_mismatch", "aligned_mismatch"): + assert rank_results[0][case] == rank_results[1][case] + + +def _cpu_mesh_no_cuda_query_worker(rank, world, port, result_queue): + import torch.distributed as dist + from torch.distributed.tensor import Shard, distribute_tensor, init_device_mesh + + from gefen.gefen import _amp_dtensor_protocol_preflight + + try: + os.environ["MASTER_ADDR"] = "127.0.0.1" + os.environ["MASTER_PORT"] = port + os.environ["RANK"] = str(rank) + os.environ["WORLD_SIZE"] = str(world) + dist.init_process_group( + "gloo", + rank=rank, + world_size=world, + timeout=timedelta(seconds=12), + ) + mesh = init_device_mesh("cpu", (world,)) + generator = torch.Generator(device="cpu").manual_seed(7300 + rank * 0) + + full_weight = torch.randn(8, 8, generator=generator) + weight = nn.Parameter(distribute_tensor(full_weight.clone(), mesh, [Shard(0)])) + muon = GefenMuon( + [("weight", weight)], + lr=1e-3, + fused=False, + ns_steps=1, + sharded_mode="exact", + ) + weight.grad = distribute_tensor( + (torch.randn(8, 8, generator=generator) * 0.01), mesh, [Shard(0)] + ) + + full_fp16 = torch.randn(8, 8, generator=generator).to(torch.float16) + fp16_weight = nn.Parameter( + distribute_tensor(full_fp16.clone(), mesh, [Shard(0)]) + ) + gefen = Gefen( + [("weight", fp16_weight)], + lr=1e-3, + fused=False, + factored_v_2d=False, + ) + fp16_weight.grad = distribute_tensor( + (torch.randn(8, 8, generator=generator) * 0.01).to(torch.float16), + mesh, + [Shard(0)], + ) + + # Spawned Gloo workers can share a machine with a CUDA-heavy parent, so + # a CPU-only mesh optimizer must never reach the CUDA capture query in + # either collective preflight (gefen.py gates it on a CUDA param; the + # Muon presence check mirrors that gate). Force is_available() to True + # so the regression also bites on CPU-only CI, where the old + # availability-based guard would short-circuit and hide it. + def forbidden(): + raise AssertionError( + "torch.cuda.is_current_stream_capturing queried for a " + "CPU-only mesh optimizer" + ) + + original_capturing = torch.cuda.is_current_stream_capturing + original_available = torch.cuda.is_available + torch.cuda.is_current_stream_capturing = forbidden + torch.cuda.is_available = lambda: True + try: + muon._assert_sharded_grad_presence_consistent() + native_amp = _amp_dtensor_protocol_preflight(gefen) + finally: + torch.cuda.is_current_stream_capturing = original_capturing + torch.cuda.is_available = original_available + + result_queue.put(("result", rank, {"native_amp": bool(native_amp)})) + except Exception: + result_queue.put(("error", rank, traceback.format_exc())) + finally: + if dist.is_initialized(): + dist.destroy_process_group() + + +@pytest.mark.skipif( + not torch.distributed.is_available() or not torch.distributed.is_gloo_available(), + reason="CPU-mesh CUDA-query regression needs Gloo", +) +def test_cpu_mesh_preflights_never_query_cuda_capture_state(): + import torch.multiprocessing as mp + + world = 2 + context = mp.get_context("spawn") + result_queue = context.Queue() + port = _free_port() + processes = [ + context.Process( + target=_cpu_mesh_no_cuda_query_worker, + args=(rank, world, port, result_queue), + ) + for rank in range(world) + ] + for process in processes: + process.start() + + messages = [] + try: + for _ in range(world): + messages.append(result_queue.get(timeout=45)) + except queue.Empty: + pass + finally: + for process in processes: + process.join(timeout=5) + for process in processes: + if process.is_alive(): + process.terminate() + process.join(timeout=5) + + assert all(process.exitcode == 0 for process in processes), [ + process.exitcode for process in processes + ] + errors = [payload for kind, _, payload in messages if kind == "error"] + assert not errors, "\n".join(errors) + rank_results = { + rank: payload for kind, rank, payload in messages if kind == "result" + } + assert set(rank_results) == {0, 1}, messages + # The fp16 multi-rank DTensor optimizer must still select the native AMP + # protocol; skipping the CUDA query must not change the decision. + for rank, payload in rank_results.items(): + assert payload["native_amp"] is True, (rank, payload) From db1cb06a63c0158ccc4e939a7e16a0c5d1ff9b20 Mon Sep 17 00:00:00 2001 From: thad0ctor Date: Sun, 12 Jul 2026 14:37:52 -0700 Subject: [PATCH 07/14] Move the two-GPU release gate to a local attested script Tag pushes no longer require a registered self-hosted GPU runner: the gpu_release_tests job is replaced by scripts/release_gpu_gate.sh, which downloads the release run's dist artifact via gh, installs the wheel into a cached --system-site-packages venv instead of mutating the host Python, and runs the identical preflight, mandatory GPU test list, and zero-skip enforcement. Approving the testpypi environment is the release manager's attestation that the local gate passed; the hosted CPU and Transformers Trainer wheel gates still block both publish jobs mechanically. The workflow filename, tag trigger, both environment gates, prerelease-skips-PyPI logic, and OIDC permission scoping are unchanged. --- .github/workflows/release.yml | 154 ++-------------------- CHANGELOG.md | 2 +- CONTRIBUTING.md | 2 +- scripts/release_gpu_gate.sh | 235 ++++++++++++++++++++++++++++++++++ 4 files changed, 249 insertions(+), 144 deletions(-) create mode 100755 scripts/release_gpu_gate.sh diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index 6661ef5..4f037d4 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -7,8 +7,14 @@ name: Release # # `ci.yml` intentionally runs on branch pushes and pull requests, not tag pushes. # This workflow therefore carries installed-wheel CPU and Transformers Trainer -# gates plus a mandatory two-GPU CUDA/JIT/distributed gate. Every gate tests the -# artifact built in `build`, and every publish job downloads that same artifact. +# gates. Every gate tests the artifact built in `build`, and every publish job +# downloads that same artifact. +# +# The two-GPU CUDA/JIT/distributed gate runs locally, not on a hosted runner: +# `scripts/release_gpu_gate.sh ` downloads this run's `dist` artifact and +# runs the mandatory GPU test list against the installed wheel with zero skips +# allowed. Approving the `testpypi` environment below is the release manager's +# attestation that the local GPU gate passed for this exact artifact. # # Auth is OIDC trusted publishing (no API tokens stored). The manual approval # gates are GitHub Environment "required reviewers", configured in @@ -194,144 +200,6 @@ jobs: env: CUDA_VISIBLE_DEVICES: "" run: python -m pytest tests -q -ra - - gpu_release_tests: - name: Two-GPU JIT & distributed release gate - needs: build - runs-on: [self-hosted, gpu] - timeout-minutes: 120 - steps: - - uses: actions/checkout@34e114876b0b11c390a56381ad16ebd13914f8d5 # v4.3.1 - with: - persist-credentials: false - - - uses: actions/download-artifact@d3f86a106a0bac45b974a628896c90dbdf5c8093 # v4.3.0 - with: - name: dist - path: dist/ - - - name: Preflight homogeneous two-GPU CUDA runner - run: | - python - <<'PY' - import re - import shutil - import subprocess - - import torch - - assert torch.version.cuda is not None, "release gate requires a CUDA-enabled PyTorch build" - assert torch.cuda.is_available(), "release gate requires CUDA" - assert torch.cuda.device_count() >= 2, "release gate requires at least two visible GPUs" - assert torch.distributed.is_available(), "torch.distributed is unavailable" - assert torch.distributed.is_nccl_available(), "release gate requires NCCL" - names = [torch.cuda.get_device_name(index) for index in range(2)] - capabilities = [torch.cuda.get_device_capability(index) for index in range(2)] - assert names[0] == names[1], ( - "replica-exact gate requires identical GPU models", - names, - ) - assert capabilities[0] == capabilities[1], ( - "replica-exact gate requires homogeneous GPUs", - capabilities, - ) - nvcc = shutil.which("nvcc") - assert nvcc is not None, "release gate requires nvcc on PATH" - nvcc_result = subprocess.run( - [nvcc, "--version"], check=True, capture_output=True, text=True - ) - match = re.search(r"release\s+([0-9]+)\.", nvcc_result.stdout) - assert match is not None, "could not parse nvcc CUDA version" - assert int(match.group(1)) == int(torch.version.cuda.split(".")[0]), ( - "nvcc and PyTorch CUDA major versions differ", - nvcc_result.stdout, - torch.version.cuda, - ) - print("torch:", torch.__version__, "CUDA:", torch.version.cuda) - print("GPUs:", names) - print("capabilities:", capabilities) - print(nvcc_result.stdout) - PY - if [ -n "${CUDA_VISIBLE_DEVICES:-}" ]; then - RELEASE_GPUS=$(printf '%s' "$CUDA_VISIBLE_DEVICES" | cut -d, -f1,2) - else - RELEASE_GPUS=0,1 - fi - echo "CUDA_VISIBLE_DEVICES=$RELEASE_GPUS" >> "$GITHUB_ENV" - echo "Mandatory release tests will use CUDA_VISIBLE_DEVICES=$RELEASE_GPUS" - - - name: Install the built wheel - env: - EXPECTED_VERSION: ${{ needs.build.outputs.version }} - run: | - python -m pip install --upgrade pip - python -m pip uninstall -y gefen-x || true - WHEEL=$(ls dist/*.whl) - python -m pip install "$WHEEL" pytest 'accelerate==1.14.0' - python - <<'PY' - import importlib.metadata as metadata - import os - from pathlib import Path - - import gefen - - version = metadata.version("gefen-x") - assert version == os.environ["EXPECTED_VERSION"], (version, os.environ["EXPECTED_VERSION"]) - module_path = Path(gefen.__file__).resolve() - assert (Path.cwd() / "src").resolve() not in module_path.parents, module_path - print("installed wheel:", module_path, version) - PY - ninja --version - - - name: Run mandatory JIT and distributed tests - env: - GEFEN_KERNEL_BUILD_ROOT: ${{ runner.temp }}/gefen-release-jit-${{ github.run_id }}-${{ github.run_attempt }} - GEFEN_VERBOSE_BUILD: "1" - run: | - rm -rf "$GEFEN_KERNEL_BUILD_ROOT" - python -m pytest -q -ra --junitxml=release-gpu.xml \ - tests/test_amp_grad_scaler.py \ - tests/test_capturable.py \ - tests/test_capturable_fsdp2.py \ - tests/test_deterministic_mode.py \ - tests/test_factored_v_ema_parity.py \ - tests/test_fused_fsdp2_noncontig.py \ - tests/test_fused_full_update_parity.py \ - tests/test_fused_update_v2_full_parity.py \ - tests/test_gefen_fsdp2_checkpoint.py \ - tests/test_muon_distributed_checkpoint_safety.py \ - tests/test_muon_grad_presence.py \ - tests/test_step_preflight_atomicity.py \ - tests/test_muon_fsdp2_parity.py::test_muon_fsdp2_fused_cuda_parity \ - tests/test_muon_fsdp2_parity.py::test_muon_fsdp2_fused_multirank_parity \ - tests/test_muon_fsdp2_parity.py::test_muon_fsdp2_distributed_parity \ - tests/test_muon_fsdp2_parity.py::test_muon_fsdp2_distributed_fluctuating_grad_set \ - tests/test_muon_fsdp2_parity.py::test_muon_fsdp2_distributed_checkpoint_restore - - - name: Reject skipped release-gate tests - run: | - python - <<'PY' - import xml.etree.ElementTree as ET - - root = ET.parse("release-gpu.xml").getroot() - cases = root.findall(".//testcase") - assert cases, "release gate collected no tests" - skipped = [ - "{}::{}".format(case.attrib.get("classname", ""), case.attrib.get("name", "")) - for case in cases - if case.find("skipped") is not None - ] - assert not skipped, "mandatory release tests skipped:\n{}".format("\n".join(skipped)) - print("mandatory GPU release tests:", len(cases), "passed with zero skips") - PY - - - name: Upload GPU release-gate report - if: always() - uses: actions/upload-artifact@ea165f8d65b6e75b540449e92b4886f43607fa02 # v4.6.2 - with: - name: release-gpu-junit-${{ github.run_attempt }} - path: release-gpu.xml - if-no-files-found: ignore - framework_wheel_tests: name: Installed-wheel Transformers Trainer resume needs: build @@ -405,9 +273,11 @@ jobs: testpypi: name: Publish to TestPyPI (gate 1) - needs: [build, cpu_wheel_tests, framework_wheel_tests, gpu_release_tests] + needs: [build, cpu_wheel_tests, framework_wheel_tests] runs-on: ubuntu-latest # Approval gate #1: the `testpypi` environment's required reviewers. + # Approving attests that `scripts/release_gpu_gate.sh ` passed locally + # against this run's built wheel (see the header comment). environment: name: testpypi url: https://test.pypi.org/project/gefen-x/${{ needs.build.outputs.version }}/ @@ -425,7 +295,7 @@ jobs: pypi: name: Publish to PyPI (gate 2) - needs: [build, cpu_wheel_tests, framework_wheel_tests, gpu_release_tests, testpypi] + needs: [build, cpu_wheel_tests, framework_wheel_tests, testpypi] # Prerelease tags stop at TestPyPI; only clean vX.Y.Z tags reach PyPI. if: needs.build.outputs.prerelease == 'false' runs-on: ubuntu-latest diff --git a/CHANGELOG.md b/CHANGELOG.md index ce0814e..5e40920 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -23,7 +23,7 @@ Packaging: - Require `numba>=0.65` for the compiled exact-DP codebook solver. - Make `ninja` and `setuptools>=77` core dependencies because PyTorch's runtime CUDA extension loader requires both to build the fused kernels; the former `perf` extra is no longer needed. -- Pin release build tooling, verify byte-reproducible wheel and normalized-sdist rebuilds, and gate the installed artifact on PyTorch 2.5.0 plus latest CPU, Transformers Trainer resume, and fresh-build two-GPU CUDA/distributed tests before either package index can publish it. +- Pin release build tooling, verify byte-reproducible wheel and normalized-sdist rebuilds, and gate the installed artifact on PyTorch 2.5.0 plus latest CPU and Transformers Trainer resume tests before either package index can publish it. The fresh-build two-GPU CUDA/distributed gate runs locally via `scripts/release_gpu_gate.sh` against the release run's built wheel, with zero skips allowed, and is attested by the TestPyPI environment approval. ## [0.3.0] - 2026-07-11 diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index c526ec3..5300f93 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -22,7 +22,7 @@ CI runs the CPU suite from the repo root after installing the package. Mirror th python -m pytest tests -q -ra ``` -CUDA-dependent tests skip themselves automatically on CPU. GPU kernel-parity tests require an NVIDIA device plus `nvcc`; branch CI exposes them through the manual `workflow_dispatch` job, while release tags must pass the two-GPU JIT and distributed gate in `release.yml`. Run the full suite locally with the same command on a CUDA host. +CUDA-dependent tests skip themselves automatically on CPU. GPU kernel-parity tests require an NVIDIA device plus `nvcc`; branch CI exposes them through the manual `workflow_dispatch` job, while release tags must pass the two-GPU JIT and distributed gate by running `scripts/release_gpu_gate.sh ` against the release run's built wheel before the TestPyPI environment is approved. Run the full suite locally with the same command on a CUDA host. ## Code style diff --git a/scripts/release_gpu_gate.sh b/scripts/release_gpu_gate.sh new file mode 100755 index 0000000..1f1106b --- /dev/null +++ b/scripts/release_gpu_gate.sh @@ -0,0 +1,235 @@ +#!/usr/bin/env bash +# Local two-GPU CUDA/JIT/distributed release gate for gefen-x. +# +# release.yml runs the CPU and Transformers Trainer wheel gates on hosted +# runners; this script is the GPU half of the release gate and runs on a local +# CUDA machine. Approving the `testpypi` environment in release.yml is the +# release manager's attestation that this script passed against that run's +# exact `dist` artifact. It mirrors the former gpu_release_tests workflow job: +# same runner preflight, same wheel install, same mandatory test list, and +# zero skipped tests allowed. +# +# Usage: +# scripts/release_gpu_gate.sh v0.4.0 gate the tag's release-run +# artifact (downloads `dist` +# via `gh run download`) +# scripts/release_gpu_gate.sh v0.4.0 --wheel PATH gate a local wheel instead +# +# Options: +# --fresh rebuild the cached gate venv from scratch +# --no-tag-check skip verifying that HEAD is the tagged commit (the test +# suite must otherwise come from the tag being gated) +# +# Environment: +# GEFEN_GATE_PYTHON base interpreter for the venv (default: python3). It +# must already provide a CUDA-enabled torch build; the +# venv is created with --system-site-packages so torch +# is inherited while the wheel, pytest, and accelerate +# installs stay inside the venv. +# GEFEN_GATE_CACHE venv cache directory +# (default: ~/.cache/gefen-x/release-gate) +# GEFEN_GATE_REPO GitHub repo whose release run holds the artifact +# (default: thad0ctor/Gefen-X) +# GEFEN_GATE_EXTRA_SITE optional site-packages directory grafted into the +# gate venv via a .pth file, searched after the venv's +# own packages. Use it when the CUDA torch build that +# matches the local nvcc lives in another virtualenv +# rather than in the base interpreter (a venv base +# python cannot be inherited with +# --system-site-packages). Same-minor Python required. +# CUDA_VISIBLE_DEVICES choose GPUs; the gate uses the first two listed. + +set -euo pipefail + +usage() { sed -n '2,32p' "$0" | sed 's/^# \{0,1\}//'; } + +TAG="" +WHEEL="" +FRESH=0 +TAG_CHECK=1 +while [ $# -gt 0 ]; do + case "$1" in + --wheel) WHEEL="$2"; shift 2 ;; + --fresh) FRESH=1; shift ;; + --no-tag-check) TAG_CHECK=0; shift ;; + -h|--help) usage; exit 0 ;; + v*) TAG="$1"; shift ;; + *) echo "unknown argument: $1" >&2; usage >&2; exit 2 ;; + esac +done +if [ -z "$TAG" ]; then + echo "error: a release tag (vX.Y.Z[...]) is required" >&2 + usage >&2 + exit 2 +fi +EXPECTED_VERSION="${TAG#v}" + +REPO_ROOT="$(cd "$(dirname "$0")/.." && pwd)" +cd "$REPO_ROOT" + +if [ "$TAG_CHECK" -eq 1 ]; then + TAG_COMMIT="$(git rev-parse --verify --quiet "refs/tags/${TAG}^{commit}" || true)" + if [ -z "$TAG_COMMIT" ]; then + echo "error: tag $TAG not found locally; fetch it or pass --no-tag-check" >&2 + exit 1 + fi + if [ "$(git rev-parse HEAD)" != "$TAG_COMMIT" ]; then + echo "error: HEAD is not $TAG — the gate must run the tagged test suite." >&2 + echo " git checkout $TAG (or pass --no-tag-check if you know better)" >&2 + exit 1 + fi +fi + +WORK="$(mktemp -d)" +JIT_ROOT="$(mktemp -d)" +trap 'rm -rf "$WORK" "$JIT_ROOT"' EXIT + +if [ -z "$WHEEL" ]; then + command -v gh >/dev/null || { echo "error: gh CLI required to download the run artifact" >&2; exit 1; } + # Tags publish from the fork, so don't let gh resolve `origin` (upstream). + # The default matches the trusted-publisher registration in release.yml. + GATE_REPO="${GEFEN_GATE_REPO:-thad0ctor/Gefen-X}" + RUN_ID="$(gh run list -R "$GATE_REPO" --workflow release.yml --branch "$TAG" \ + --limit 1 --json databaseId --jq '.[0].databaseId')" + if [ -z "$RUN_ID" ] || [ "$RUN_ID" = "null" ]; then + echo "error: no release.yml run found for tag $TAG on $GATE_REPO" >&2 + exit 1 + fi + echo "downloading dist artifact from $GATE_REPO release run $RUN_ID" + gh run download "$RUN_ID" -R "$GATE_REPO" --name dist --dir "$WORK/dist" + WHEEL="$(ls "$WORK"/dist/*.whl)" +fi +[ -f "$WHEEL" ] || { echo "error: wheel not found: $WHEEL" >&2; exit 1; } +echo "gating wheel: $WHEEL" +sha256sum "$WHEEL" + +BASE_PYTHON="${GEFEN_GATE_PYTHON:-python3}" +CACHE_DIR="${GEFEN_GATE_CACHE:-$HOME/.cache/gefen-x/release-gate}" +VENV="$CACHE_DIR/venv" +if [ "$FRESH" -eq 1 ]; then + rm -rf "$VENV" +fi +if [ ! -x "$VENV/bin/python" ]; then + echo "creating gate venv at $VENV (base: $BASE_PYTHON)" + mkdir -p "$CACHE_DIR" + "$BASE_PYTHON" -m venv --system-site-packages "$VENV" +fi +PYTHON="$VENV/bin/python" +SITE_DIR="$("$PYTHON" -c 'import sysconfig; print(sysconfig.get_paths()["purelib"])')" +if [ -n "${GEFEN_GATE_EXTRA_SITE:-}" ]; then + [ -d "$GEFEN_GATE_EXTRA_SITE" ] || { echo "error: GEFEN_GATE_EXTRA_SITE is not a directory: $GEFEN_GATE_EXTRA_SITE" >&2; exit 1; } + printf '%s\n' "$GEFEN_GATE_EXTRA_SITE" > "$SITE_DIR/gefen_gate_extra.pth" + echo "grafted extra site dir: $GEFEN_GATE_EXTRA_SITE" +else + rm -f "$SITE_DIR/gefen_gate_extra.pth" +fi + +"$PYTHON" - <<'PY' +import re +import shutil +import subprocess + +import torch + +assert torch.version.cuda is not None, "release gate requires a CUDA-enabled PyTorch build" +assert torch.cuda.is_available(), "release gate requires CUDA" +assert torch.cuda.device_count() >= 2, "release gate requires at least two visible GPUs" +assert torch.distributed.is_available(), "torch.distributed is unavailable" +assert torch.distributed.is_nccl_available(), "release gate requires NCCL" +names = [torch.cuda.get_device_name(index) for index in range(2)] +capabilities = [torch.cuda.get_device_capability(index) for index in range(2)] +assert names[0] == names[1], ( + "replica-exact gate requires identical GPU models", + names, +) +assert capabilities[0] == capabilities[1], ( + "replica-exact gate requires homogeneous GPUs", + capabilities, +) +nvcc = shutil.which("nvcc") +assert nvcc is not None, "release gate requires nvcc on PATH" +nvcc_result = subprocess.run( + [nvcc, "--version"], check=True, capture_output=True, text=True +) +match = re.search(r"release\s+([0-9]+)\.", nvcc_result.stdout) +assert match is not None, "could not parse nvcc CUDA version" +assert int(match.group(1)) == int(torch.version.cuda.split(".")[0]), ( + "nvcc and PyTorch CUDA major versions differ", + nvcc_result.stdout, + torch.version.cuda, +) +print("torch:", torch.__version__, "CUDA:", torch.version.cuda) +print("GPUs:", names) +print("capabilities:", capabilities) +print(nvcc_result.stdout) +PY + +if [ -n "${CUDA_VISIBLE_DEVICES:-}" ]; then + RELEASE_GPUS="$(printf '%s' "$CUDA_VISIBLE_DEVICES" | cut -d, -f1,2)" +else + RELEASE_GPUS=0,1 +fi +echo "mandatory release tests will use CUDA_VISIBLE_DEVICES=$RELEASE_GPUS" + +"$PYTHON" -m pip install --upgrade pip +"$PYTHON" -m pip uninstall -y gefen-x || true +"$PYTHON" -m pip install "$WHEEL" pytest 'accelerate==1.14.0' +EXPECTED_VERSION="$EXPECTED_VERSION" REPO_ROOT="$REPO_ROOT" "$PYTHON" - <<'PY' +import importlib.metadata as metadata +import os +from pathlib import Path + +import gefen + +version = metadata.version("gefen-x") +assert version == os.environ["EXPECTED_VERSION"], (version, os.environ["EXPECTED_VERSION"]) +module_path = Path(gefen.__file__).resolve() +repo_src = (Path(os.environ["REPO_ROOT"]) / "src").resolve() +assert repo_src not in module_path.parents, module_path +print("installed wheel:", module_path, version) +PY +"$VENV/bin/ninja" --version 2>/dev/null || ninja --version + +JUNIT="$WORK/release-gpu.xml" +CUDA_VISIBLE_DEVICES="$RELEASE_GPUS" \ +GEFEN_KERNEL_BUILD_ROOT="$JIT_ROOT" \ +GEFEN_VERBOSE_BUILD=1 \ +"$PYTHON" -m pytest -q -ra --junitxml="$JUNIT" \ + tests/test_amp_grad_scaler.py \ + tests/test_capturable.py \ + tests/test_capturable_fsdp2.py \ + tests/test_deterministic_mode.py \ + tests/test_factored_v_ema_parity.py \ + tests/test_fused_fsdp2_noncontig.py \ + tests/test_fused_full_update_parity.py \ + tests/test_fused_update_v2_full_parity.py \ + tests/test_gefen_fsdp2_checkpoint.py \ + tests/test_muon_distributed_checkpoint_safety.py \ + tests/test_muon_grad_presence.py \ + tests/test_step_preflight_atomicity.py \ + tests/test_muon_fsdp2_parity.py::test_muon_fsdp2_fused_cuda_parity \ + tests/test_muon_fsdp2_parity.py::test_muon_fsdp2_fused_multirank_parity \ + tests/test_muon_fsdp2_parity.py::test_muon_fsdp2_distributed_parity \ + tests/test_muon_fsdp2_parity.py::test_muon_fsdp2_distributed_fluctuating_grad_set \ + tests/test_muon_fsdp2_parity.py::test_muon_fsdp2_distributed_checkpoint_restore + +JUNIT="$JUNIT" "$PYTHON" - <<'PY' +import os +import xml.etree.ElementTree as ET + +root = ET.parse(os.environ["JUNIT"]).getroot() +cases = root.findall(".//testcase") +assert cases, "release gate collected no tests" +skipped = [ + "{}::{}".format(case.attrib.get("classname", ""), case.attrib.get("name", "")) + for case in cases + if case.find("skipped") is not None +] +assert not skipped, "mandatory release tests skipped:\n{}".format("\n".join(skipped)) +print("mandatory GPU release tests:", len(cases), "passed with zero skips") +PY + +echo +echo "GPU release gate PASSED for $TAG" +echo "wheel: $(basename "$WHEEL") sha256: $(sha256sum "$WHEEL" | cut -d' ' -f1)" +echo "You may now approve the testpypi environment for this run." From dbaea4e3c1c5cbc53c0a87f1fec74397cb72476f Mon Sep 17 00:00:00 2001 From: thad0ctor Date: Sun, 12 Jul 2026 15:30:06 -0700 Subject: [PATCH 08/14] Write the distributed compatibility table as capabilities, not test status --- README.md | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/README.md b/README.md index 17abd6c..68a0ab0 100644 --- a/README.md +++ b/README.md @@ -116,14 +116,14 @@ One knock-on effect: weight decay in AdamW-style optimizers is applied as `lr × Gefen drops into standard distributed training like any other PyTorch optimizer, with either `fused=True` or `fused=False`. Validated training setups include single-GPU, PyTorch DDP, FSDP2 (`fully_shard` / DTensor), and DeepSpeed ZeRO 1-3 (plain `Gefen` as the client optimizer, direct or via axolotl `gefenx`; bit-exact ZeRO-2 checkpoint resume). FSDP2 optimizer checkpoint support is mode- and topology-specific as described below. -| System | Current status | +| System | Works with | |---|---| -| DDP | Supported and tested, including fused BF16 resume | -| FSDP2 / DTensor training | Tested for plain Gefen and Muon `approx` / `exact` / `distributed` | -| FSDP2 full-state checkpoints (DCP) | Plain Gefen and Muon `approx` only; same 1-D topology and world size — scope note below | -| Muon `distributed` native checkpoints | Versioned owner manifest; collective save resumes across world sizes, including into a single-process optimizer — [details](#experimental-lever-sharded-newton-schulz-under-fsdp2-sharded_mode) | -| DeepSpeed ZeRO 1-3 | Plain Gefen supported and tested, direct and via axolotl `gefenx`; Muon/Hybrid fail fast with a clear error — config note below | -| Megatron DP/TP/PP/CP/EP/ETP | Tested through the Megatron GPT pretraining entry point with legacy optimizer checkpoints — [scope](https://github.com/thad0ctor/Gefen-X/blob/main/COMPATIBILITY.md) | +| DDP | All optimizers, including fused BF16 resume | +| FSDP2 / DTensor training | Plain Gefen and Muon `approx` / `exact` / `distributed` | +| FSDP2 full-state checkpoints (DCP) | Plain Gefen and Muon `approx`; same 1-D topology and world size — scope note below | +| Muon `distributed` native checkpoints | Resume under any world size, including a single-process optimizer — [details](#experimental-lever-sharded-newton-schulz-under-fsdp2-sharded_mode) | +| DeepSpeed ZeRO 1-3 | Plain Gefen, direct or via axolotl `gefenx`; the Muon family needs FSDP2, DDP, or single-GPU — config note below | +| Megatron DP/TP/PP/CP/EP/ETP | Plain Gefen, GefenMuon+AdamW, and GefenMuon+Gefen with legacy optimizer checkpoints — [scope](https://github.com/thad0ctor/Gefen-X/blob/main/COMPATIBILITY.md) | > **FSDP2 optimizer checkpoint scope.** Plain Gefen and `GefenMuon(sharded_mode="approx")` collectively encode every rank's local DTensor optimizer state into PyTorch DCP `StateDictOptions(full_state_dict=True)` output, including the flattened optimizer-state form. `get_optimizer_state_dict()` and `set_optimizer_state_dict()` resume the next update exactly when world size, mesh, placements, rank coordinates, parameter ordering, shapes, names, and sharded mode are unchanged; the actual two-GPU `fully_shard` get/set test covers both optimizers. The adapter currently requires one 1-D DeviceMesh spanning the default world; multidimensional meshes, subgroups, and pipeline-local optimizers fail before its collectives. Save and restore are collective, so every rank must participate. Each process temporarily stages all serialized rank payloads on CPU, making the leading checkpoint-time CPU cost about `world_size ×` that rank's local optimizer state plus local serialization scratch. World-size or topology changes fail before mutation, and older unsafe untagged full checkpoints fail closed. This is not a reshardable optimizer-state format. @@ -621,7 +621,7 @@ opt = GefenMuonHybrid( # only takes effect under FSDP2 (DTensor params); no-op single-GPU ``` -> **`"distributed"` checkpointing is collective.** `state_dict()` gathers each owner's momentum across ranks, so **every rank must call it** (as in a standard FSDP full-state-dict flow). Calling `state_dict()` on rank 0 only — e.g. a rank-0-only save loop — **deadlocks**. Save and load also transiently materialize the full unsharded momentum on every rank, so peak memory at checkpoint time approaches `"exact"` mode's. The saved checkpoint carries a versioned owner manifest and is complete on every rank, so it resumes under a different world size or in a single-process optimizer; populated state whose manifest is missing, partial, or inconsistent is rejected before any mutation. +> **`"distributed"` checkpointing is collective.** `state_dict()` gathers each owner's momentum across ranks, so **every rank must call it** (as in a standard FSDP full-state-dict flow). Calling `state_dict()` on rank 0 only — e.g. a rank-0-only save loop — **deadlocks**. Save and load also transiently materialize the full unsharded momentum on every rank, so peak memory at checkpoint time approaches `"exact"` mode's. The saved checkpoint is complete on every rank and resumes under any world size, including a single-process optimizer; incomplete or inconsistent owner state fails closed before any state is touched. ![Gefen-Muon exact / distributed / approx sharded — eval loss](https://raw.githubusercontent.com/thad0ctor/Gefen-X/main/docs/benchmarks/muon_shard_loss.png) ![Gefen-Muon exact / distributed / approx sharded — throughput & VRAM](https://raw.githubusercontent.com/thad0ctor/Gefen-X/main/docs/benchmarks/muon_shard_perf.png) From f3d6c16910a63d141c8379cf01a6cb37659a770b Mon Sep 17 00:00:00 2001 From: thad0ctor Date: Sun, 12 Jul 2026 15:31:13 -0700 Subject: [PATCH 09/14] Simplify the distributed compatibility table --- README.md | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/README.md b/README.md index 68a0ab0..7f43e06 100644 --- a/README.md +++ b/README.md @@ -118,12 +118,12 @@ Gefen drops into standard distributed training like any other PyTorch optimizer, | System | Works with | |---|---| -| DDP | All optimizers, including fused BF16 resume | -| FSDP2 / DTensor training | Plain Gefen and Muon `approx` / `exact` / `distributed` | -| FSDP2 full-state checkpoints (DCP) | Plain Gefen and Muon `approx`; same 1-D topology and world size — scope note below | -| Muon `distributed` native checkpoints | Resume under any world size, including a single-process optimizer — [details](#experimental-lever-sharded-newton-schulz-under-fsdp2-sharded_mode) | -| DeepSpeed ZeRO 1-3 | Plain Gefen, direct or via axolotl `gefenx`; the Muon family needs FSDP2, DDP, or single-GPU — config note below | -| Megatron DP/TP/PP/CP/EP/ETP | Plain Gefen, GefenMuon+AdamW, and GefenMuon+Gefen with legacy optimizer checkpoints — [scope](https://github.com/thad0ctor/Gefen-X/blob/main/COMPATIBILITY.md) | +| DDP | All optimizers | +| FSDP2 | All optimizers | +| FSDP2 checkpoints | Plain Gefen and Muon `approx`; resume needs the same GPU count — scope note below | +| Muon `distributed` checkpoints | Resume on any GPU count, even a single GPU — [details](#experimental-lever-sharded-newton-schulz-under-fsdp2-sharded_mode) | +| DeepSpeed ZeRO 1-3 | Plain Gefen; use FSDP2 or DDP for the Muon family — config note below | +| Megatron-LM | All optimizers, including checkpoint resume — [scope](https://github.com/thad0ctor/Gefen-X/blob/main/COMPATIBILITY.md) | > **FSDP2 optimizer checkpoint scope.** Plain Gefen and `GefenMuon(sharded_mode="approx")` collectively encode every rank's local DTensor optimizer state into PyTorch DCP `StateDictOptions(full_state_dict=True)` output, including the flattened optimizer-state form. `get_optimizer_state_dict()` and `set_optimizer_state_dict()` resume the next update exactly when world size, mesh, placements, rank coordinates, parameter ordering, shapes, names, and sharded mode are unchanged; the actual two-GPU `fully_shard` get/set test covers both optimizers. The adapter currently requires one 1-D DeviceMesh spanning the default world; multidimensional meshes, subgroups, and pipeline-local optimizers fail before its collectives. Save and restore are collective, so every rank must participate. Each process temporarily stages all serialized rank payloads on CPU, making the leading checkpoint-time CPU cost about `world_size ×` that rank's local optimizer state plus local serialization scratch. World-size or topology changes fail before mutation, and older unsafe untagged full checkpoints fail closed. This is not a reshardable optimizer-state format. From a9677da7cd583c2a10274cc8cb836b72ada8081e Mon Sep 17 00:00:00 2001 From: thad0ctor Date: Sun, 12 Jul 2026 15:37:20 -0700 Subject: [PATCH 10/14] Trim README technical walls to plain summaries with doc links --- README.md | 49 +++++++++++++++++++------------------------------ 1 file changed, 19 insertions(+), 30 deletions(-) diff --git a/README.md b/README.md index 7f43e06..3877e13 100644 --- a/README.md +++ b/README.md @@ -17,12 +17,12 @@ - **Fine-tuning and pre-training** — upstream Gefen was validated for pre-training; this fork makes it a drop-in optimizer for fine-tuning: full fine-tune (FFT), LoRA, and QLoRA, with a [VRAM sizing guide](https://github.com/thad0ctor/Gefen-X/blob/main/docs/hardware.md). - **New: matches AdamW's loss out of the box** at about a quarter of its optimizer memory (default `factored_v_2d`). See [Benchmarks](#benchmarks) and [the factored-v lever](#quality-lever-factored-second-moment-on-2d-params-factored_v_2d). - - **Works on modern decoders** (Qwen3, Llama-3, Mistral). Upstream loses its memory advantage on these architectures (~9 B/param — worse than AdamW); this fork keeps the intended ~1 B/param. + - **Works on modern decoders** (Qwen3, Llama-3, Mistral). Upstream loses its memory advantage on these architectures (≈9 B/param — worse than AdamW); this fork keeps the intended ≈1 B/param. - **Validated across 2026 architectures** — two dozen modern LLMs, VLMs, and image/video/audio media-gen models full-fine-tuned on 24 GB GPUs, plus 20-30B MoEs via LoRA. See the [compatibility matrix](https://github.com/thad0ctor/Gefen-X/blob/main/COMPATIBILITY.md). - - **~2× faster `opt.step()`** via fused CUDA kernels, with validated numerical parity; fixed-order paths have bitwise checks, while atomic reduction paths are tolerance-bounded. + - **≈2× faster `opt.step()`** via fused CUDA kernels with validated numerical parity. - **Whole-model Muon option** (`GefenMuonHybrid`) with a selectable AdamW quality backup or ~1 B/param Gefen low-memory backup, plus task-specific [SFT and pretraining recipes](#which-muon-recipe-should-i-use). - - **Reliable native checkpoint save/resume and FSDP2 training support** — broken or absent in the shipped release. Plain Gefen and rank-local `GefenMuon(sharded_mode="approx")` also support same-topology PyTorch full-state DCP optimizer resume; see [Distributed Training](#distributed-training) for its collective and portability limits. - - **Hardened against crashes** (device/dtype guards, bounds checks, race fixes) with bitwise tests where reduction order is fixed and tolerance-bounded tests for CUDA atomic reductions. + - **Reliable native checkpoint save/resume and FSDP2 training support** — broken or absent in the shipped release; see [Distributed Training](#distributed-training) for the checkpoint compatibility table. + - **Hardened against crashes** (device/dtype guards, bounds checks, race fixes).
Detailed Fork Improvements (vs upstream) @@ -33,19 +33,19 @@ >| **Loss vs AdamW** | trails AdamW by ~0.06 | **matches AdamW** via the default `factored_v_2d` — [details](#quality-lever-factored-second-moment-on-2d-params-factored_v_2d) | >| **Modern decoders** (Qwen3 / Llama-3 / Mistral — SwiGLU + grouped-query attention) | uses *more* optimizer memory than AdamW on these | keeps the full ~1 B/param optimizer state (about a quarter of bf16 AdamW's) | >| **Learning rate(s)** | no guidance — silently over-steps | documented ~0.6× AdamW, so quality matches AdamW | ->| **Optimizer-step speed** | baseline | ~2× faster `opt.step()` (fused kernels), with fixed-order bitwise checks and bounded atomic-reduction differences | +>| **Optimizer-step speed** | baseline | ≈2× faster `opt.step()` with fused kernels | >| **Peak memory**| large transient spikes | much lower peak — room for bigger models / batches | >| **Sharded multi-GPU training (FSDP2)** | breaks with the fast path | works — for plain Gefen *and* Muon | >| **Whole-model Muon** | 2D weight matrices only | `GefenMuonHybrid` trains the entire model | >| **Muon step efficiency** | generic momentum hack + redundant dequant gather | single-pass bit-exact momentum kernel | >| **Save / resume checkpoints** | can corrupt state or lose tuning on resume | native optimizer checkpoints save and resume correctly; distributed checkpoint formats have documented limits | >| **Crash safety** | missing device / edge-case guards | guarded against wrong-device, empty-tensor, and race bugs | ->| **Correctness** | no fused-kernel tests | bitwise kernel checks, tolerance-bounded atomic-reduction parity, and distributed tests | +>| **Correctness** | no fused-kernel tests | kernel parity and distributed tests | >| **Documentation** | no Axolotl / fork-install guidance | Axolotl how-to + fair loss/speed/memory benchmarks | >| **Muon Usage** | no whole-model recipe or task split | measured SFT/pretraining recipes pair quantized Muon with an AdamW backup and document the observed quality/throughput tradeoff; a Gefen backup remains available for minimum state — [details](#which-muon-recipe-should-i-use) | >| **Muon (Newton-Schulz) speed** | fixed 5-step schedule | tunable `ns_schedule`; tuned3 is the balanced SFT choice, while quality-first pretraining retains classic NS5 — [details](#experimental-lever-faster-newton-schulz-ns_schedule-fp8_ns) | >| **Low-precision orthogonalization** | bf16 only | opt-in `fp8_ns` for large matrices on newer GPUs, safe fallback elsewhere — [details](#experimental-lever-faster-newton-schulz-ns_schedule-fp8_ns) | ->| **Sharded Muon under FSDP2** | every GPU redundantly repeats the same work | `sharded_mode="distributed"` splits the work across GPUs and matches `"exact"` bitwise on homogeneous GPUs in the parity suite — [details](#experimental-lever-sharded-newton-schulz-under-fsdp2-sharded_mode) | +>| **Sharded Muon under FSDP2** | every GPU redundantly repeats the same work | `sharded_mode="distributed"` splits the work across GPUs with identical results on matching GPUs — [details](#experimental-lever-sharded-newton-schulz-under-fsdp2-sharded_mode) |
@@ -114,7 +114,7 @@ One knock-on effect: weight decay in AdamW-style optimizers is applied as `lr × ## Distributed Training -Gefen drops into standard distributed training like any other PyTorch optimizer, with either `fused=True` or `fused=False`. Validated training setups include single-GPU, PyTorch DDP, FSDP2 (`fully_shard` / DTensor), and DeepSpeed ZeRO 1-3 (plain `Gefen` as the client optimizer, direct or via axolotl `gefenx`; bit-exact ZeRO-2 checkpoint resume). FSDP2 optimizer checkpoint support is mode- and topology-specific as described below. +Gefen drops into standard distributed training like any other PyTorch optimizer, with either `fused=True` or `fused=False`. | System | Works with | |---|---| @@ -123,19 +123,17 @@ Gefen drops into standard distributed training like any other PyTorch optimizer, | FSDP2 checkpoints | Plain Gefen and Muon `approx`; resume needs the same GPU count — scope note below | | Muon `distributed` checkpoints | Resume on any GPU count, even a single GPU — [details](#experimental-lever-sharded-newton-schulz-under-fsdp2-sharded_mode) | | DeepSpeed ZeRO 1-3 | Plain Gefen; use FSDP2 or DDP for the Muon family — config note below | -| Megatron-LM | All optimizers, including checkpoint resume — [scope](https://github.com/thad0ctor/Gefen-X/blob/main/COMPATIBILITY.md) | +| Megatron-LM | All optimizers, including checkpoint resume — [scope](https://github.com/thad0ctor/Gefen-X/blob/main/COMPATIBILITY.md#megatron-lm-integration-scope) | -> **FSDP2 optimizer checkpoint scope.** Plain Gefen and `GefenMuon(sharded_mode="approx")` collectively encode every rank's local DTensor optimizer state into PyTorch DCP `StateDictOptions(full_state_dict=True)` output, including the flattened optimizer-state form. `get_optimizer_state_dict()` and `set_optimizer_state_dict()` resume the next update exactly when world size, mesh, placements, rank coordinates, parameter ordering, shapes, names, and sharded mode are unchanged; the actual two-GPU `fully_shard` get/set test covers both optimizers. The adapter currently requires one 1-D DeviceMesh spanning the default world; multidimensional meshes, subgroups, and pipeline-local optimizers fail before its collectives. Save and restore are collective, so every rank must participate. Each process temporarily stages all serialized rank payloads on CPU, making the leading checkpoint-time CPU cost about `world_size ×` that rank's local optimizer state plus local serialization scratch. World-size or topology changes fail before mutation, and older unsafe untagged full checkpoints fail closed. This is not a reshardable optimizer-state format. +> **FSDP2 checkpoint scope.** Plain Gefen and Muon `approx` save and resume exactly through PyTorch's standard full-state checkpoint calls, as long as the GPU count and sharding layout are unchanged and every GPU joins the save. Anything outside that scope fails with a clear error instead of corrupting state — [full details](https://github.com/thad0ctor/Gefen-X/blob/main/COMPATIBILITY.md#optimizer-checkpoint-scope). -`torch.amp.GradScaler` keeps PyTorch's ordinary externally skipped step for FP32-master and BF16 training, including Trainer/Accelerate gradient clipping and scheduler behavior. Actual FP16 gradient storage opts into PyTorch's native optimizer-side scaling protocol so finite gradients are unscaled once and overflow returns before either Hybrid child, codebook, state, parameter, or counter changes. DTensor/FSDP2 non-finite flags are reduced across the mesh; FSDP1 FlatParameters must use `torch.distributed.fsdp.ShardedGradScaler`. +Mixed precision works out of the box: BF16 and standard AMP behave exactly as with any PyTorch optimizer, and true-FP16 `GradScaler` training is handled safely — an overflow step changes nothing. FSDP1 FP16 needs `torch.distributed.fsdp.ShardedGradScaler`. -> **DeepSpeed ZeRO config.** Set `"zero_allow_untested_optimizer": true` and leave the config's `optimizer` section unset. With optimizer CPU-offload, also set `"zero_force_ds_cpu_optimizer": false` — otherwise raw DeepSpeed refuses to initialize, and accelerate-based launchers (axolotl) silently swap in DeepSpeed's own CPU Adam. ZeRO steps flattened 1-D partitions, so `GefenMuon`/`GefenMuonHybrid` raise a clear error under ZeRO; use FSDP2, DDP, or single-GPU for the Muon family. +> **DeepSpeed ZeRO config.** Set `"zero_allow_untested_optimizer": true` and leave the config's `optimizer` section unset. With optimizer CPU-offload, also set `"zero_force_ds_cpu_optimizer": false` — otherwise raw DeepSpeed refuses to initialize, and accelerate-based launchers (axolotl) silently swap in DeepSpeed's own CPU Adam. ## Replica-exact fused updates (`deterministic`) -Set `deterministic=True` when data-parallel replicas must remain bit-exact on homogeneous GPUs. Automatic block periods are selected with fixed-order GPU reductions instead of the faster atomic CUDA search, block-vmean parameters remain fused through the fixed-order v1 CUDA reduction, and factored-v parameters use the deterministic decomposed factored update instead of the fused stats kernel's unordered floating-point atomics. The default is `False`, so existing performance routing is unchanged. Tagged checkpoints must resume with the same deterministic policy; legacy checkpoints without a tag remain loadable. Plain Gefen does not allow `deterministic=True`, `factored_v_2d=True`, and `stochastic_round=True` together because the deterministic factored fallback uses nearest-codeword quantization. - -**Megatron-LM validation scope.** The Megatron integration was exercised through its GPT pretraining entry point with plain Gefen, GefenMuon+AdamW, and GefenMuon+Gefen. Fused deterministic runs covered DP2, TP2, PP2, CP2 with Transformer Engine, and EP2 on two homogeneous RTX 3090 Ti GPUs, with replica or tied-weight hashes appropriate to each topology; legacy optimizer checkpoint continuation was also exercised under TP2 and PP2 for all three recipes. Unfused four-rank coverage additionally exercised EP2×DP2, TP2×DP2, PP2×DP2, and TP2×PP2 for all three recipes, plus EP2×ETP2 and EP2 checkpoint continuation for GefenMuon+Gefen. These are tiny mock-data integration gates, not large-scale convergence claims, and they do not cover Megatron's distributed optimizer, FSDP, optimizer CPU offload, fp16, or non-legacy optimizer checkpoint formats. +Set `deterministic=True` when every GPU replica must produce bit-identical results on matching GPUs. It is off by default (the fastest routing), checkpoints remember the setting, and older checkpoints without it still load. One combination is rejected: plain Gefen with `deterministic=True`, `factored_v_2d=True`, and `stochastic_round=True` together. ```python optimizer = Gefen(model.named_parameters(), lr=3e-4, fused=True, deterministic=True) @@ -143,7 +141,7 @@ optimizer = Gefen(model.named_parameters(), lr=3e-4, fused=True, deterministic=T ## CUDA Graphs & torch.compile (`capturable`) -All three optimizers accept `capturable=True` (same meaning as `torch.optim`'s argument): `opt.step()` can then be captured in a `torch.cuda.CUDAGraph` or wrapped in `torch.compile(mode="reduce-overhead")`. Performance is path-dependent rather than guaranteed cost-free: the retained measurements are nearly flat for plain Gefen, show modest eager/manual-capture overhead for the hybrid, and make the compiled hybrid step about 10% faster than default eager. Device-resident global counters advance on every replay, so checkpoints record the true replayed step and stochastic-rounding resumes continue from the correct seed. Usage, graph-partition caveats, and measured numbers: [COMPATIBILITY.md](https://github.com/thad0ctor/Gefen-X/blob/main/COMPATIBILITY.md#cuda-graphs--torchcompile-capturable). +All three optimizers accept `capturable=True` (same meaning as `torch.optim`'s argument): `opt.step()` can then be captured in a `torch.cuda.CUDAGraph` or wrapped in `torch.compile(mode="reduce-overhead")` — the compiled hybrid step is about 10% faster than eager. Checkpoints stay correct across graph replays. Usage, caveats, and measured numbers: [COMPATIBILITY.md](https://github.com/thad0ctor/Gefen-X/blob/main/COMPATIBILITY.md#cuda-graphs--torchcompile-capturable).
@@ -480,7 +478,7 @@ This keeps both halves on quantized Gefen state and is the low-memory recommenda | `normuon` | `True` | free per-neuron 2nd moment on the NS output; recovers tuned3 quality in SFT — [details](#quality-lever-per-neuron-2nd-moment-on-the-newton-schulz-output-normuon) | keep on for SFT; disable for classic pretraining | | `backup_2d_period_one` | `False` | per-element 2nd moment on a Gefen-backed embedding/LM head — extra memory — [details](#experimental-lever-per-element-gefen-backup-state-on-embed--lm-head-backup_2d_period_one) | Gefen backup only; AdamW already keeps per-element moments | | `stochastic_round` | `False` | unbiased rounding for the 8-bit momentum (free, loss-neutral) | optional | -| `deterministic` | `False` | replica-exact fused routing on homogeneous GPUs; persists in child checkpoints | enable when distributed replica hashes must match bit-for-bit | +| `deterministic` | `False` | bit-identical updates on matching GPUs; remembered by checkpoints | enable when replicas must match exactly | | `capturable` | `False` | CUDA-graph-capturable `step()`, like `torch.optim`'s `capturable` — [details](https://github.com/thad0ctor/Gefen-X/blob/main/COMPATIBILITY.md#cuda-graphs--torchcompile-capturable) | turn on to capture `step()` in a `torch.cuda.CUDAGraph` | It supports `step()`, `zero_grad()`, `state_dict()`/`load_state_dict()`, and LR schedulers (e.g. `torch.optim.lr_scheduler.StepLR(optimizer, ...)`) like any optimizer. Because it splits params at construction rather than taking a single iterable, build it yourself and hand it to the Hugging Face `Trainer` via `optimizers=` (not `optimizer_cls_and_kwargs`): @@ -498,16 +496,7 @@ trainer = Trainer(model=model, args=training_args, train_dataset=train_dataset, optimizers=(optimizer, None)) # (optimizer, lr_scheduler) ``` -For an executable save/resume gate covering plain Gefen and both Muon backup variants, including gradient accumulation, changing LRs, tied weights, optional BF16, and DDP replica hashes, use [`benchmarks/trainer_resume/`](https://github.com/thad0ctor/Gefen-X/tree/main/benchmarks/trainer_resume). Its Transformers 5 optimizer factory constructs the optimizer from Trainer's model and validates Trainer's internal Accelerate wrapping. All three recipes have passed the deterministic fused-BF16 two-rank gate on homogeneous GPUs, which requires exact uninterrupted-versus-resumed model, optimizer, scheduler, loss-history, and cross-replica hashes: - -```bash -CUDA_VISIBLE_DEVICES=0,1 PYTHONPATH=.:src torchrun --standalone --nproc-per-node=2 \ - -m benchmarks.trainer_resume.run --output-dir benchmarks/trainer_resume/out/ddp \ - --device cuda --dtype bfloat16 --fused --deterministic --steps 3 --split-step 1 \ - --gradient-accumulation-steps 2 -``` - -The harness tests Trainer's DDP frontend and internal Accelerate optimizer wrapper; it does not test a standalone `Accelerator` loop, FSDP, or `model_init`. +To prove save/resume in a real Trainer run, use [`benchmarks/trainer_resume/`](https://github.com/thad0ctor/Gefen-X/tree/main/benchmarks/trainer_resume): it compares an uninterrupted run against a save-and-resume run and requires them to match exactly. The run command and covered scope are in [COMPATIBILITY.md](https://github.com/thad0ctor/Gefen-X/blob/main/COMPATIBILITY.md#transformers-trainer-ddp). ## Quality Lever: factored second moment on 2D params (`factored_v_2d`) @@ -535,7 +524,7 @@ opt = Gefen(model.named_parameters(), lr=0.6 * ADAMW_LR, fused=True) # factored Details: validation, checkpoint migration, and limits - Validated at both scales in the fair-LR regime, with a second-seed replication at 0.6B and a learning-rate sweep at 1.7B (best LR `3e-5` ~ 0.6× AdamW's, matching Gefen's documented heuristic). The fused kernel computes the per-element step size in registers (no extra temporaries; step transients measured 0 MiB) and is covered by `tests/test_gefen_factored_v.py`. -- Native Gefen checkpoints migrate automatically in both directions (old checkpoints work with the new default and vice versa; the second-moment statistics re-warm briefly and harmlessly). Rank-local DTensor full-state DCP uses a separate tagged, collective format and requires the same world size and topology. +- Native Gefen checkpoints migrate automatically in both directions (old checkpoints work with the new default and vice versa; the second-moment statistics re-warm briefly and harmlessly). Distributed checkpoints follow the [compatibility table](#distributed-training). - With `backup_optimizer="gefen"`, `GefenMuonHybrid` pins the backup half to Gefen's legacy block-vmean path (the factored-v combination has not been benchmarked). With `backup_optimizer="adamw"`, the backup uses conventional per-element AdamW state. Under FSDP2, sharded 2D Gefen params also fall back to the legacy path. - Two sibling experiments from the same investigation ship off by default because they measured **no effect**: `period_one_substrings` (per-element state on name-matched tensors) and `codebook_refresh_every` (periodic codebook refit). @@ -641,8 +630,8 @@ Measured (Qwen3-0.6B, 2 and 4 GPUs): `"distributed"` matched `"exact"` in the re ## Known limitations - **Hybrid checkpoint schema.** `GefenMuonHybrid`'s `state_dict()` uses its own nested `{"muon": ..., "backup": ..., "backup_optimizer": "gefen" | "adamw"}` layout. Resume from a checkpoint the hybrid itself saved—not one consolidated or converted to the flat torch `{state, param_groups}` layout. Cross-backend loads are rejected before either child is mutated; legacy untagged hybrid checkpoints are interpreted as Gefen-backed. -- **FSDP2 full-state optimizer DCP is same-topology only.** Plain Gefen and `GefenMuon(sharded_mode="approx")` preserve rank-local state in a tagged collective payload and resume exactly with the same world size, one-dimensional default-world mesh, placements, rank coordinates, parameter layout, and mode. Every rank must enter save and restore, and each process can transiently use about `world_size ×` its local optimizer-state size in CPU memory plus local serialization scratch. Multidimensional meshes, subgroups, pipeline-local optimizers, and resharding are intentionally rejected; old untagged full checkpoints that could have reused rank 0's codebook on every rank are also rejected. This limitation does not apply to model-only DCP. -- **Accelerate cannot observe native true-FP16 overflow skips.** PyTorch's native AMP optimizer protocol calls `optimizer.step()` even when the optimizer returns before mutation, so Accelerate's `step_was_skipped` flag remains false and a scheduler driven solely by that flag can advance. Normal FP32-master autocast and BF16 use the ordinary GradScaler path and are unaffected; prefer those modes in Trainer/Accelerate, or gate a true-FP16 scheduler from the scaler's scale change. +- **FSDP2 optimizer checkpoints don't reshard.** Plain Gefen and Muon `approx` resume only on the same GPU count and layout; changing either fails with a clear error. Model weights are unaffected — [details](https://github.com/thad0ctor/Gefen-X/blob/main/COMPATIBILITY.md#optimizer-checkpoint-scope). +- **True-FP16 overflow skips are invisible to Accelerate's `step_was_skipped` flag.** BF16 and standard AMP are unaffected and are the recommended modes in Trainer/Accelerate. ## Troubleshooting From 24ec6ec5a666ec0fc8bcaf5de3c460e0001fd117 Mon Sep 17 00:00:00 2001 From: thad0ctor Date: Sun, 12 Jul 2026 15:41:56 -0700 Subject: [PATCH 11/14] Say refuses to load, not clear error, for out-of-scope FSDP2 resume --- README.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index 3877e13..571ade0 100644 --- a/README.md +++ b/README.md @@ -125,7 +125,7 @@ Gefen drops into standard distributed training like any other PyTorch optimizer, | DeepSpeed ZeRO 1-3 | Plain Gefen; use FSDP2 or DDP for the Muon family — config note below | | Megatron-LM | All optimizers, including checkpoint resume — [scope](https://github.com/thad0ctor/Gefen-X/blob/main/COMPATIBILITY.md#megatron-lm-integration-scope) | -> **FSDP2 checkpoint scope.** Plain Gefen and Muon `approx` save and resume exactly through PyTorch's standard full-state checkpoint calls, as long as the GPU count and sharding layout are unchanged and every GPU joins the save. Anything outside that scope fails with a clear error instead of corrupting state — [full details](https://github.com/thad0ctor/Gefen-X/blob/main/COMPATIBILITY.md#optimizer-checkpoint-scope). +> **FSDP2 checkpoint scope.** Plain Gefen and Muon `approx` save and resume exactly through PyTorch's standard full-state checkpoint calls, as long as the GPU count and sharding layout are unchanged and every GPU joins the save. Anything outside that scope refuses to load instead of corrupting state — [full details](https://github.com/thad0ctor/Gefen-X/blob/main/COMPATIBILITY.md#optimizer-checkpoint-scope). Mixed precision works out of the box: BF16 and standard AMP behave exactly as with any PyTorch optimizer, and true-FP16 `GradScaler` training is handled safely — an overflow step changes nothing. FSDP1 FP16 needs `torch.distributed.fsdp.ShardedGradScaler`. @@ -630,7 +630,7 @@ Measured (Qwen3-0.6B, 2 and 4 GPUs): `"distributed"` matched `"exact"` in the re ## Known limitations - **Hybrid checkpoint schema.** `GefenMuonHybrid`'s `state_dict()` uses its own nested `{"muon": ..., "backup": ..., "backup_optimizer": "gefen" | "adamw"}` layout. Resume from a checkpoint the hybrid itself saved—not one consolidated or converted to the flat torch `{state, param_groups}` layout. Cross-backend loads are rejected before either child is mutated; legacy untagged hybrid checkpoints are interpreted as Gefen-backed. -- **FSDP2 optimizer checkpoints don't reshard.** Plain Gefen and Muon `approx` resume only on the same GPU count and layout; changing either fails with a clear error. Model weights are unaffected — [details](https://github.com/thad0ctor/Gefen-X/blob/main/COMPATIBILITY.md#optimizer-checkpoint-scope). +- **FSDP2 optimizer checkpoints don't reshard.** Plain Gefen and Muon `approx` resume only on the same GPU count and layout; changing either refuses to load. Model weights are unaffected — [details](https://github.com/thad0ctor/Gefen-X/blob/main/COMPATIBILITY.md#optimizer-checkpoint-scope). - **True-FP16 overflow skips are invisible to Accelerate's `step_was_skipped` flag.** BF16 and standard AMP are unaffected and are the recommended modes in Trainer/Accelerate. ## Troubleshooting From b1fa6235d9444d91a1e2f8395e834e696db24a29 Mon Sep 17 00:00:00 2001 From: thad0ctor Date: Sun, 12 Jul 2026 15:47:03 -0700 Subject: [PATCH 12/14] Address PR review: gloo skipif guard and explicit shared seed --- tests/test_muon_distributed_checkpoint_safety.py | 5 +++++ tests/test_muon_grad_presence.py | 4 +++- 2 files changed, 8 insertions(+), 1 deletion(-) diff --git a/tests/test_muon_distributed_checkpoint_safety.py b/tests/test_muon_distributed_checkpoint_safety.py index 7e2ca0b..8688999 100644 --- a/tests/test_muon_distributed_checkpoint_safety.py +++ b/tests/test_muon_distributed_checkpoint_safety.py @@ -585,6 +585,11 @@ def make_composed_pair(approx_value, distributed_value): dist.destroy_process_group() +@pytest.mark.skipif( + not torch.distributed.is_available() + or not torch.distributed.is_gloo_available(), + reason="mixed parallel/fallback CPU composition needs Gloo", +) def test_mixed_parallel_fallback_and_rank_local_composition_cpu(): import torch.multiprocessing as mp diff --git a/tests/test_muon_grad_presence.py b/tests/test_muon_grad_presence.py index 218b2a0..ccb92db 100644 --- a/tests/test_muon_grad_presence.py +++ b/tests/test_muon_grad_presence.py @@ -728,7 +728,9 @@ def _cpu_mesh_no_cuda_query_worker(rank, world, port, result_queue): timeout=timedelta(seconds=12), ) mesh = init_device_mesh("cpu", (world,)) - generator = torch.Generator(device="cpu").manual_seed(7300 + rank * 0) + # Identical seed on every rank: distribute_tensor needs the same full + # tensor across ranks to stay collective-safe. + generator = torch.Generator(device="cpu").manual_seed(7300) full_weight = torch.randn(8, 8, generator=generator) weight = nn.Parameter(distribute_tensor(full_weight.clone(), mesh, [Shard(0)])) From 3ebec7eaa606b16264dadbf3abfc0eac94323fc6 Mon Sep 17 00:00:00 2001 From: thad0ctor Date: Sun, 12 Jul 2026 15:52:00 -0700 Subject: [PATCH 13/14] Document deterministic=True usage for Trainer and axolotl --- README.md | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/README.md b/README.md index 571ade0..8f42ee2 100644 --- a/README.md +++ b/README.md @@ -139,6 +139,8 @@ Set `deterministic=True` when every GPU replica must produce bit-identical resul optimizer = Gefen(model.named_parameters(), lr=3e-4, fused=True, deterministic=True) ``` +It is an ordinary constructor argument, so it works the same everywhere you build the optimizer: pass it in the [Hugging Face Trainer](#hugging-face-trainer) construction, or in [axolotl](#using-gefen-with-axolotl) add `optim_args: { deterministic: true }`. + ## CUDA Graphs & torch.compile (`capturable`) All three optimizers accept `capturable=True` (same meaning as `torch.optim`'s argument): `opt.step()` can then be captured in a `torch.cuda.CUDAGraph` or wrapped in `torch.compile(mode="reduce-overhead")` — the compiled hybrid step is about 10% faster than eager. Checkpoints stay correct across graph replays. Usage, caveats, and measured numbers: [COMPATIBILITY.md](https://github.com/thad0ctor/Gefen-X/blob/main/COMPATIBILITY.md#cuda-graphs--torchcompile-capturable). @@ -275,7 +277,7 @@ optim_args: For quality-first pretraining, keep the AdamW backup at full LR and set `ns_schedule: standard` plus `normuon: false`. In every recipe, `adjust_lr_fn="match_rms_adamw"` keeps Muon updates on the AdamW LR scale. -**How config maps:** `learning_rate`, `weight_decay`, `adam_beta1`/`adam_beta2`, and `adam_epsilon` forward to the selected constructor. Under `optim_args`, `fused` applies to both optimizer names; `factored_v_2d` is `gefenx`-only; and `backup_optimizer`, `backup_lr`, `ns_schedule`, `normuon`, and `sharded_mode` are `gefenx_muon`-only. String values are coerced to type. For plain Gefen, measure the LR instead of relying only on the `~0.6×` heuristic: `python -m gefen.tools.find_lr --model --optimizer gefen --method sweep` — see [`tools/README.md`](https://github.com/thad0ctor/Gefen-X/blob/main/src/gefen/tools/README.md). +**How config maps:** `learning_rate`, `weight_decay`, `adam_beta1`/`adam_beta2`, and `adam_epsilon` forward to the selected constructor. Under `optim_args`, `fused` and `deterministic` apply to both optimizer names; `factored_v_2d` is `gefenx`-only; and `backup_optimizer`, `backup_lr`, `ns_schedule`, `normuon`, and `sharded_mode` are `gefenx_muon`-only. String values are coerced to type. For plain Gefen, measure the LR instead of relying only on the `~0.6×` heuristic: `python -m gefen.tools.find_lr --model --optimizer gefen --method sweep` — see [`tools/README.md`](https://github.com/thad0ctor/Gefen-X/blob/main/src/gefen/tools/README.md). | Lever | axolotl config | Notes | |---|---|---| @@ -285,6 +287,7 @@ For quality-first pretraining, keep the AdamW backup at full LR and set `ns_sche | Factored 2D 2nd moment | `optim_args: { factored_v_2d: true }` | matches AdamW loss; default on | | Muon backup LR | `optim_args: { backup_lr: }` | LR for the selected Gefen/AdamW backup; the Axolotl factory defaults to `0.5 × learning_rate`, while balanced SFT/pretraining use full LR | | Sharded Newton-Schulz | `optim_args: { sharded_mode: exact }` | `exact` (default) or `distributed` (splits NS across GPUs under FSDP2) | +| Replica-exact updates | `optim_args: { deterministic: true }` | bit-identical GPU replicas on matching GPUs; both optimizer names — [details](#replica-exact-fused-updates-deterministic) | | Learning rate | `learning_rate: ` | the ~0.6× heuristic applies to plain Gefen; tune Muon by task. Background: [Learning rate](#learning-rate-when-porting-an-adamw-config) | | `period==1` memory fallback | *on by default in this fork* | restores ~1 B/param on modern decoders; module flag `MEMORY_SAFE_FALLBACK`, not a YAML key | From e8d86775a42e410733e24c8a19d6e6de621ec42a Mon Sep 17 00:00:00 2001 From: thad0ctor Date: Sun, 12 Jul 2026 15:53:45 -0700 Subject: [PATCH 14/14] Rename deterministic section heading to Determinism --- README.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index 8f42ee2..92ca5f8 100644 --- a/README.md +++ b/README.md @@ -131,7 +131,7 @@ Mixed precision works out of the box: BF16 and standard AMP behave exactly as wi > **DeepSpeed ZeRO config.** Set `"zero_allow_untested_optimizer": true` and leave the config's `optimizer` section unset. With optimizer CPU-offload, also set `"zero_force_ds_cpu_optimizer": false` — otherwise raw DeepSpeed refuses to initialize, and accelerate-based launchers (axolotl) silently swap in DeepSpeed's own CPU Adam. -## Replica-exact fused updates (`deterministic`) +## Determinism (`deterministic`) Set `deterministic=True` when every GPU replica must produce bit-identical results on matching GPUs. It is off by default (the fastest routing), checkpoints remember the setting, and older checkpoints without it still load. One combination is rejected: plain Gefen with `deterministic=True`, `factored_v_2d=True`, and `stochastic_round=True` together. @@ -287,7 +287,7 @@ For quality-first pretraining, keep the AdamW backup at full LR and set `ns_sche | Factored 2D 2nd moment | `optim_args: { factored_v_2d: true }` | matches AdamW loss; default on | | Muon backup LR | `optim_args: { backup_lr: }` | LR for the selected Gefen/AdamW backup; the Axolotl factory defaults to `0.5 × learning_rate`, while balanced SFT/pretraining use full LR | | Sharded Newton-Schulz | `optim_args: { sharded_mode: exact }` | `exact` (default) or `distributed` (splits NS across GPUs under FSDP2) | -| Replica-exact updates | `optim_args: { deterministic: true }` | bit-identical GPU replicas on matching GPUs; both optimizer names — [details](#replica-exact-fused-updates-deterministic) | +| Determinism | `optim_args: { deterministic: true }` | bit-identical GPU replicas on matching GPUs; both optimizer names — [details](#determinism-deterministic) | | Learning rate | `learning_rate: ` | the ~0.6× heuristic applies to plain Gefen; tune Muon by task. Background: [Learning rate](#learning-rate-when-porting-an-adamw-config) | | `period==1` memory fallback | *on by default in this fork* | restores ~1 B/param on modern decoders; module flag `MEMORY_SAFE_FALLBACK`, not a YAML key |