Enforce cross-rank parameter order in the Muon preflight

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 before exact/distributed DTensor Muon collectives, avoiding unmatched-collective deadlocks and cross-rank momentum-owner misalignment 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 — 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.

scripts/release_gpu_gate.sh

@@ -41,15 +41,22 @@
 
 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=""
 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; }
+      # 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 ;;

src/gefen/gefen_muon.py

@@ -1297,15 +1297,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
@@ -1339,12 +1342,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),
@@ -1369,6 +1373,10 @@ def _assert_sharded_grad_presence_consistent(self) -> None:
         for mesh_key in sorted(by_mesh):
             entry = by_mesh[mesh_key]
             mesh = entry["mesh"]
+            local_position = {
+                id(item[1]): position
+                for position, item in enumerate(entry["items"])
+            }
             items = sorted(
                 entry["items"],
                 key=lambda item: (
@@ -1383,15 +1391,51 @@ def _assert_sharded_grad_presence_consistent(self) -> None:
                 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 local insertion position to the mesh-wide max and
+            # min (one MAX all_reduce over [position, -position]) exposes any
+            # rank whose parameter-group order differs.
+            order_probe = torch.tensor(
+                [
+                    signed
+                    for _, p, _ in items
+                    for signed in (
+                        local_position[id(p)],
+                        -local_position[id(p)],
+                    )
+                ],
+                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
+                    )
+
+            order_flat = order_probe.cpu().tolist()
+            order_mismatches = [
+                items[index][0]
+                for index in range(len(items))
+                if order_flat[2 * index] != -order_flat[2 * index + 1]
+            ]
+            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))
+                )
 
             mesh_size = mesh.size()
             inconsistent = torch.nonzero(

tests/test_muon_grad_presence.py

@@ -572,7 +572,12 @@ def _reversed_order_worker(rank, world, port, result_queue):
         mesh = init_device_mesh("cpu", (world,))
         results = []
 
-        for case in ("reversed_mismatch", "reversed_consistent", "aligned_mismatch"):
+        for case in (
+            "reversed_mismatch",
+            "reversed_consistent",
+            "aligned_mismatch",
+            "aligned_consistent",
+        ):
             generator = torch.Generator(device="cpu").manual_seed(
                 6100 + len(results)
             )
@@ -604,10 +609,11 @@ 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()
@@ -630,13 +636,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
 
@@ -685,28 +694,32 @@ def test_sharded_grad_presence_check_is_order_insensitive():
             "reversed_mismatch",
             "reversed_consistent",
             "aligned_mismatch",
+            "aligned_consistent",
         }, (rank, cases)
-        for case in ("reversed_mismatch", "aligned_mismatch"):
+        # 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]
             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 "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"]
+        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"] is None, (
             rank,
-            cases["reversed_consistent"],
+            cases["aligned_consistent"],
         )
 
-    # Sorted item names make the diagnostic identical on every rank.
-    for case in ("reversed_mismatch", "aligned_mismatch"):
+    # 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"):
         assert rank_results[0][case] == rank_results[1][case]