diff --git a/CHANGELOG.md b/CHANGELOG.md index 5e40920..20b6b0f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -12,7 +12,7 @@ Correctness and compatibility: - 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. +- Reject rank-divergent `grad is None` patterns and rank-divergent parameter-group registration order (including swaps across DeviceMeshes) before exact/distributed DTensor Muon collectives, avoiding unmatched-collective deadlocks and cross-rank momentum-owner misalignment without adding collectives to plain/DDP/`approx` steps. Parameters sharing a name, shape, and dtype in one mesh are refused as cross-rank-unidentifiable, and same-shaped unnamed parameters warn once that order divergence between them is undetectable. - 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. diff --git a/scripts/release_gpu_gate.sh b/scripts/release_gpu_gate.sh index 1f1106b..8fb15fc 100755 --- a/scripts/release_gpu_gate.sh +++ b/scripts/release_gpu_gate.sh @@ -41,7 +41,8 @@ set -euo pipefail -usage() { sed -n '2,32p' "$0" | sed 's/^# \{0,1\}//'; } +# Print the whole leading comment block, however long it grows. +usage() { awk 'NR == 1 { next } /^#/ { sub(/^# ?/, ""); print; next } { exit }' "$0"; } TAG="" WHEEL="" @@ -49,7 +50,16 @@ FRESH=0 TAG_CHECK=1 while [ $# -gt 0 ]; do case "$1" in - --wheel) WHEEL="$2"; shift 2 ;; + --wheel) + [ $# -ge 2 ] || { echo "error: --wheel requires a path" >&2; exit 2; } + case "$2" in + -*) echo "error: --wheel requires a path, got option-like '$2'" >&2; exit 2 ;; + esac + # Resolve now: the script cd's to the repo root later, which would break + # a relative path given from another directory. + WHEEL="$(realpath -m -- "$2")" + shift 2 + ;; --fresh) FRESH=1; shift ;; --no-tag-check) TAG_CHECK=0; shift ;; -h|--help) usage; exit 0 ;; diff --git a/src/gefen/gefen_muon.py b/src/gefen/gefen_muon.py index 4196895..02e435f 100644 --- a/src/gefen/gefen_muon.py +++ b/src/gefen/gefen_muon.py @@ -1,4 +1,5 @@ import math +import re import warnings from collections import OrderedDict from typing import Iterable, Optional, Tuple, Union @@ -27,6 +28,12 @@ def _stable_distributed_owner(stable_index: int, world: int) -> int: raise ValueError("stable_index must be non-negative, got {}".format(stable_index)) return stable_index % world +# Auto-generated names for bare (unnamed) parameters, as assigned during group +# registration in gefen.py: "param_{i}", "group_{g}_param_{i}", plus the +# "_{n}" uniquing suffix. These names follow registration position, so they +# carry no cross-rank identity of their own. +_AUTO_PARAM_NAME = re.compile(r"^(?:param|group_\d+_param)_\d+(?:_\d+)?$") + # Tuned per-iteration Newton-Schulz coefficient schedules. Each entry is a list # of (a, b, c) quintic coefficients applied one per iteration. Unlike the fixed # (DEFAULT_A, DEFAULT_B, DEFAULT_C) quintic -- which uses identical, conservative @@ -1297,15 +1304,18 @@ def _dist_available() -> bool: @torch._dynamo.disable def _assert_sharded_grad_presence_consistent(self) -> None: - """Fail before collectives when mesh ranks disagree on ``grad is None``. + """Fail before collectives when mesh ranks disagree on the step inputs. 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. + collectives entered in parameter-group insertion order, and + ``_step_distributed_pg`` assigns momentum owners by that same insertion + index. Two rank-local properties therefore must agree mesh-wide before + any of those collectives start: which parameters have gradients, and + the order the parameters were registered in. Pack both per DeviceMesh + — an activity bit summed across every mesh dimension, and an insertion + position reduced to its mesh-wide max and min — so every participating + rank sees the same global result 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 @@ -1330,6 +1340,11 @@ def _assert_sharded_grad_presence_consistent(self) -> None: import torch.distributed as dist by_mesh = OrderedDict() + # Optimizer-wide insertion positions: step() walks parameters and + # first-seen process groups in this order, so cross-mesh order + # divergence is just as collective-fatal as divergence inside one + # mesh. Per-mesh enumeration would hide it. + global_position = {} for group in self.param_groups: if group["sharded_mode"] == "approx": continue @@ -1339,12 +1354,13 @@ def _assert_sharded_grad_presence_consistent(self) -> None: mesh = p.device_mesh if mesh.get_coordinate() is None or mesh.size() < 2: continue - # 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. + # Two equivalent DeviceMesh objects can still refer to distinct + # c10d groups, and a mis-built training script can register the + # same parameters in a different order per rank. Key each mesh + # by content instead of object identity, and sort meshes and + # items below, so this preflight's own collectives stay matched + # under both — and rank-divergent registration order is then + # detected positionally rather than corrupting the comparison. process_groups = tuple(mesh.get_all_groups()) key = ( str(mesh.device_type), @@ -1364,8 +1380,12 @@ def _assert_sharded_grad_presence_consistent(self) -> None: } by_mesh[key] = entry entry["items"].append((str(name), p, p.grad is not None)) + global_position[id(p)] = len(global_position) mismatches = [] + order_mismatches = [] + duplicate_labels = [] + duplicates_anywhere = False for mesh_key in sorted(by_mesh): entry = by_mesh[mesh_key] mesh = entry["mesh"] @@ -1377,21 +1397,110 @@ def _assert_sharded_grad_presence_consistent(self) -> None: str(item[1].dtype), ), ) + # Two parameters carrying the same explicit name, shape, and dtype + # are indistinguishable across ranks, so neither this order probe + # nor the presence vector below can align them; construction + # accepts duplicate explicit names, so fail closed on them. The + # flag rides the reduced probe below rather than raising here: + # a rank whose own labels are clean must learn about duplicates on + # its peers and take the same error path, not enter a collective + # its peer already abandoned. + duplicate_keys = sorted( + { + items[index][0] + for index in range(1, len(items)) + if items[index][0] == items[index - 1][0] + and tuple(items[index][1].shape) + == tuple(items[index - 1][1].shape) + and str(items[index][1].dtype) + == str(items[index - 1][1].dtype) + } + ) + # Bare parameters receive positional auto-names, so two of them + # with the same shape and dtype have no rank-invariant identity at + # all: the probes below would align labels that themselves follow + # position, and no rank-local property can do better. Ranks built + # by the same construction code (the normal case) remain correct; + # warn once that divergence between such parameters is invisible + # to this preflight, and recommend named construction. + if not getattr(self, "_gefen_warned_unverifiable_order", False): + collisions = {} + for name, p, _ in items: + if _AUTO_PARAM_NAME.match(str(name)): + key = (tuple(p.shape), str(p.dtype)) + collisions.setdefault(key, []).append(str(name)) + ambiguous = sorted( + name + for names in collisions.values() + if len(names) > 1 + for name in names + ) + if ambiguous: + warnings.warn( + "GefenMuon cannot verify cross-rank parameter order " + "for unnamed parameters that share a shape and dtype " + "({}): their auto-generated names follow registration " + "position, so rank-divergent registration order " + "between them is undetectable. Construct the " + "optimizer from model.named_parameters() to make " + "this check complete.".format(", ".join(ambiguous)), + RuntimeWarning, + stacklevel=3, + ) + self._gefen_warned_unverifiable_order = True 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. + # Each sorted slot is the same parameter identity on every rank, so + # reducing its optimizer-wide insertion position to the mesh-wide + # max and min (one MAX all_reduce over [position, -position]) + # exposes any rank whose registration order differs — including + # order swaps between parameters on different meshes. The final + # slot carries this rank's duplicate-label count so the reduction + # also propagates duplicates to every rank of the mesh. + order_probe = torch.tensor( + [ + signed + for _, p, _ in items + for signed in ( + global_position[id(p)], + -global_position[id(p)], + ) + ] + + [len(duplicate_keys)], + dtype=torch.int32, + device=device, + ) + # Reducing the same vectors once along every Cartesian mesh + # dimension propagates mesh-global results to every rank. This is + # two collectives per mesh dimension, independent of the number of + # optimizer parameters. 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 ) + dist.all_reduce( + order_probe, op=dist.ReduceOp.MAX, group=process_group + ) + + # Do not raise inside this loop: with DTensors on overlapping but + # non-identical meshes, a rank exiting early would strand peers + # that only share a later mesh inside that mesh's all_reduce. + # Accumulate every violation and raise after all local meshes + # have completed their probe collectives, like the presence path. + order_flat = order_probe.cpu().tolist() + if order_flat[-1] > 0: + duplicates_anywhere = True + duplicate_labels.extend(duplicate_keys) + order_mismatches.extend( + items[index][0] + for index in range(len(items)) + if order_flat[2 * index] != -order_flat[2 * index + 1] + ) mesh_size = mesh.size() inconsistent = torch.nonzero( @@ -1408,6 +1517,29 @@ def _assert_sharded_grad_presence_consistent(self) -> None: ) ) + if duplicates_anywhere: + raise RuntimeError( + "GefenMuon cannot verify cross-rank parameter order or " + "gradient presence when parameters in one " + "sharded_mode='exact' or 'distributed' DTensor mesh share " + "a name, shape, and dtype{}: the shared label makes them " + "indistinguishable across ranks. Give every parameter a " + "unique name.".format( + " ({})".format(", ".join(sorted(set(duplicate_labels)))) + if duplicate_labels + else " on at least one mesh rank" + ) + ) + if order_mismatches: + raise RuntimeError( + "GefenMuon requires identical parameter-group order on " + "every rank of a DTensor/FSDP mesh before " + "sharded_mode='exact' or 'distributed' stepping: gradient " + "collectives and momentum-owner assignment follow that " + "order. Parameters at rank-divergent positions: {}. " + "Construct parameter groups in the same order on every " + "rank.".format(", ".join(order_mismatches)) + ) if mismatches: raise RuntimeError( "GefenMuon requires identical gradient presence on every rank " diff --git a/tests/test_muon_grad_presence.py b/tests/test_muon_grad_presence.py index ccb92db..7899f46 100644 --- a/tests/test_muon_grad_presence.py +++ b/tests/test_muon_grad_presence.py @@ -570,23 +570,62 @@ def _reversed_order_worker(rank, world, port, result_queue): timeout=timedelta(seconds=12), ) mesh = init_device_mesh("cpu", (world,)) + second_mesh = init_device_mesh("cpu", (world,)) results = [] - for case in ("reversed_mismatch", "reversed_consistent", "aligned_mismatch"): + import warnings + + for case in ( + "reversed_mismatch", + "reversed_consistent", + "aligned_mismatch", + "aligned_consistent", + "bare_ambiguous", + "duplicate_names", + "duplicate_one_rank", + "cross_mesh_reversed", + ): generator = torch.Generator(device="cpu").manual_seed( 6100 + len(results) ) + # The ambiguity cases need two same-shape parameters: identical + # shape and dtype is exactly what makes positional auto-names or + # shared explicit names ambiguous. + b_shape = ( + (8, 8) + if case + in ("bare_ambiguous", "duplicate_names", "duplicate_one_rank") + else (6, 10) + ) + b_mesh = second_mesh if case == "cross_mesh_reversed" else mesh 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 + full_b = torch.randn(*b_shape, generator=generator) + full_b_grad = torch.randn(*b_shape, 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)) + b = nn.Parameter( + distribute_tensor(full_b.clone(), b_mesh, [Shard(0)]) + ) + if case == "bare_ambiguous": + params = [a, b] + elif case == "duplicate_names": + params = [("w", a), ("w", b)] + elif case == "duplicate_one_rank": + # Only rank 1 carries the duplicate labels; rank 0 must still + # learn about them through the reduced flag and take the same + # error path instead of blocking in a peer-abandoned + # collective. + params = ( + [("w", a), ("w", b)] if rank == 1 else [("a", a), ("b", b)] + ) + else: + params = [("a", a), ("b", b)] + if ( + case.startswith("reversed") or case == "cross_mesh_reversed" + ) and rank == 1: + params = list(reversed(params)) optimizer = GefenMuon( - named, + params, lr=1e-3, fused=False, ns_steps=1, @@ -596,7 +635,7 @@ def _reversed_order_worker(rank, world, port, result_queue): # 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)]) + b.grad = distribute_tensor(full_b_grad.clone(), b_mesh, [Shard(0)]) if case.endswith("mismatch"): if rank == 0: b.grad = None @@ -604,16 +643,25 @@ def _reversed_order_worker(rank, world, port, result_queue): 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. + # order hardening, the reversed cases either falsely passed here + # (complementary grads) or sailed into step()'s misaligned owner + # assignment (consistent grads), 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}) + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + try: + optimizer._assert_sharded_grad_presence_consistent() + except RuntimeError as exc: + message = str(exc) + warned = any( + "cannot verify cross-rank parameter order" in str(item.message) + for item in caught + ) + results.append( + {"case": case, "message": message, "warned": warned} + ) a.grad = None b.grad = None dist.barrier() @@ -630,13 +678,16 @@ def _reversed_order_worker(rank, world, port, result_queue): 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. +def test_sharded_grad_presence_check_rejects_rank_divergent_order(): + """Rank-divergent param-group order must be rejected, not merely survived. 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. + and the check falsely passes right before full_tensor() deadlocks. And + even with consistent gradient presence, exact/distributed stepping enters + collectives and assigns momentum owners by insertion order, so divergent + order itself must raise on every rank before any of that starts. """ import torch.multiprocessing as mp @@ -674,7 +725,7 @@ def test_sharded_grad_presence_check_is_order_insensitive(): 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} + rank: {item["case"]: item for item in payload} for kind, rank, payload in messages if kind == "result" } @@ -685,29 +736,76 @@ def test_sharded_grad_presence_check_is_order_insensitive(): "reversed_mismatch", "reversed_consistent", "aligned_mismatch", + "aligned_consistent", + "bare_ambiguous", + "duplicate_names", + "duplicate_one_rank", + "cross_mesh_reversed", }, (rank, cases) - for case in ("reversed_mismatch", "aligned_mismatch"): - message = cases[case] + # Duplicate explicit labels defeat cross-rank identification, so they + # fail closed even with aligned order and full gradient presence — + # and the rank whose own labels are clean must raise too, via the + # reduced flag, instead of blocking in an abandoned collective. + for case in ("duplicate_names", "duplicate_one_rank"): + message = cases[case]["message"] 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, ( + assert "unique name" in message, (rank, case, message) + # Order swaps between parameters on different meshes are step-fatal + # too; the optimizer-wide position probe must catch them. + message = cases["cross_mesh_reversed"]["message"] + assert message is not None, (rank, "cross_mesh_reversed") + assert "identical parameter-group order" in message, (rank, message) + # Rank-divergent order is rejected up front, with or without a + # gradient-presence mismatch layered on top. + for case in ("reversed_mismatch", "reversed_consistent"): + message = cases[case]["message"] + assert message is not None, (rank, case) + assert "identical parameter-group order" in message, ( rank, case, message, ) - assert cases["reversed_consistent"] is None, ( + assert "a" in message and "b" in message, (rank, case, message) + message = cases["aligned_mismatch"]["message"] + assert message is not None, (rank, "aligned_mismatch") + assert "identical gradient presence" in message, (rank, message) + assert "a (1/2 mesh ranks have gradients)" in message, (rank, message) + assert "b (1/2 mesh ranks have gradients)" in message, (rank, message) + assert cases["aligned_consistent"]["message"] is None, ( rank, - cases["reversed_consistent"], + cases["aligned_consistent"], + ) + # Same-shape bare parameters have no cross-rank identity, so the + # preflight cannot police their order; it must say so once instead + # of failing, and must stay silent for named construction. + assert cases["bare_ambiguous"]["message"] is None, ( + rank, + cases["bare_ambiguous"], + ) + assert cases["bare_ambiguous"]["warned"], (rank, cases["bare_ambiguous"]) + for case in ( + "reversed_mismatch", + "reversed_consistent", + "aligned_mismatch", + "aligned_consistent", + "duplicate_names", + "duplicate_one_rank", + "cross_mesh_reversed", + ): + assert not cases[case]["warned"], (rank, case) + + # Reduced order/activity vectors make the diagnostic identical on every + # rank, which is what lets the job fail synchronously. + for case in ( + "reversed_mismatch", + "reversed_consistent", + "aligned_mismatch", + "duplicate_names", + "cross_mesh_reversed", + ): + assert ( + rank_results[0][case]["message"] == rank_results[1][case]["message"] ) - - # 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):