Harden convention codebook-scope preambles and contract validation

docs/optimizer_contracts.md

@@ -130,7 +130,7 @@ Portable v3 currently excludes non-period-one initialized state, second-moment r
 
 Plain Gefen declares replicated, flattened element-shard, and the narrow DTensor training layouts. Its PyTorch rank-local checkpoint transport is same-topology only. GefenMuon declares replicated and narrow DTensor training, with mode-specific state extents: `approx` state is local, `exact` state is logically global, and `distributed` momentum is held by the parameter owner while non-owners retain metadata only. Native Parallel-Muon checkpoints separately declare world-size owner redistribution, not placement-changing resharding. `GefenMuonHybrid` retains its nested child namespaces and does not flatten AdamW or Gefen child state into a fabricated common schema.
 
-Gefen and GefenMuon native loads, rank-local payload restoration through `load_state_dict`, and distributed-owner payload restoration prepare their complete core restore before changing local live optimizer state. Their transport entries report `atomic_load=True`. This is a per-optimizer-instance fail-before-mutation guarantee at the optimizer load boundary, not a coordinated all-rank commit or a guarantee over work an external checkpoint orchestrator performs before calling the optimizer. Load pre-hooks run before that boundary and load post-hooks run afterward, so arbitrary side effects in user hooks are also outside the guarantee. Hybrid composite loads do not yet provide the same guarantee and report `atomic_load=False`.
+Gefen and GefenMuon native loads, rank-local payload restoration through `load_state_dict`, and distributed-owner payload restoration prepare their complete core restore before changing local live optimizer state. Their transport entries report `atomic_load=True`. This is a per-optimizer-instance fail-before-mutation guarantee at the optimizer load boundary, not a coordinated all-rank commit or a guarantee over work an external checkpoint orchestrator performs before calling the optimizer. Load pre-hooks run before that boundary and load post-hooks run afterward, so arbitrary side effects in user hooks are also outside the guarantee. Ordinary nested `GefenMuonHybrid.load_state_dict()` loads do not yet provide the same guarantee and report `atomic_load=False`; portable composite imports retain the separate fail-before-mutation guarantee described above.
 
 ## Adapter requirements
 

src/gefen/checkpoint.py

@@ -122,11 +122,10 @@ def _validate_collective_device(self, backend: object) -> None:
                 raise ValueError(
                     "checkpoint collective device is incompatible with the runtime backend"
                 )
-        elif "gloo" in backend_name or "mpi" in backend_name:
-            if self.collective_device.type != "cpu":
-                raise ValueError(
-                    "checkpoint collective device is incompatible with the runtime backend"
-                )
+        # Gloo (and MPI) support CUDA tensors in addition to CPU, so no
+        # backend-level device restriction applies here; a CUDA collective
+        # device on those backends still falls through to the availability
+        # check below.
 
         if self.collective_device.type == "cuda":
             index = self.collective_device.index

src/gefen/contracts.py

@@ -798,6 +798,8 @@ def __post_init__(self) -> None:
             raise TypeError("StateField.geometry must be a StateGeometry")
         if not isinstance(self.key_match, StateKeyMatch):
             raise TypeError("StateField.key_match must be a StateKeyMatch")
+        if type(self.checkpointed) is not bool:
+            raise TypeError("StateField.checkpointed must be a bool")
         if type(self.optional) is not bool:
             raise TypeError("StateField.optional must be a bool")
 
@@ -853,6 +855,9 @@ def __post_init__(self) -> None:
             raise TypeError("StateVariant.extent must be a StateExtent")
         if not isinstance(self.role, ParameterStateRole):
             raise TypeError("StateVariant.role must be a ParameterStateRole")
+        for name in ("initialized", "migration_only"):
+            if type(getattr(self, name)) is not bool:
+                raise TypeError("StateVariant.{} must be a bool".format(name))
         if not set(self.inactive_fields).issubset(self.fields):
             raise ValueError("StateVariant.inactive_fields must be present in fields")
         if self.parameter_ranks is not None and set(self.parameter_ranks) & set(
@@ -991,6 +996,9 @@ def __post_init__(self) -> None:
             raise TypeError(
                 "CheckpointSupport.process_group_scope must be a ProcessGroupScope"
             )
+        for name in ("requires_collective", "atomic_load"):
+            if type(getattr(self, name)) is not bool:
+                raise TypeError("CheckpointSupport.{} must be a bool".format(name))
         if bool(self.topology_changing) != bool(self.topology_change_kinds):
             raise ValueError(
                 "topology-changing layouts and change kinds must be declared together"
@@ -1094,7 +1102,10 @@ def __post_init__(self) -> None:
         object.__setattr__(self, "children", _tuple(self.children))
         if not self.implementation:
             raise ValueError("OptimizerContract.implementation must be non-empty")
-        if self.schema_version != CONTRACT_SCHEMA_VERSION:
+        if (
+            type(self.schema_version) is not int
+            or self.schema_version != CONTRACT_SCHEMA_VERSION
+        ):
             raise ValueError(
                 "unsupported optimizer contract schema version: {}".format(
                     self.schema_version

src/gefen/gefen.py

@@ -5166,9 +5166,25 @@ def _maybe_refresh_gefen_codebook(self) -> None:
     def initialize_codebook(self) -> bool:
         """Collectively initialize the learned codebook without taking a step."""
 
-        self._assert_finalized_binding_layout(full=True)
-        self._assert_runtime_codebook_process_group()
-        self._assert_codebook_capture_ready()
+        # Capture a usable failure-vote binding before inspecting live layout so
+        # a one-sided preamble failure -- for example capture-readiness raising
+        # only on the gradient-owning member while non-owners proceed -- is
+        # reported through the scope instead of stranding peers inside the
+        # scoped operation-header collective below.
+        scope_binding = self._capture_codebook_scope_binding_for_step()
+        try:
+            self._assert_finalized_binding_layout(full=True)
+            self._assert_runtime_codebook_process_group()
+            self._assert_codebook_capture_ready()
+            local_preamble_error = None
+        except Exception as exc:
+            local_preamble_error = exc
+        if scope_binding is not None:
+            self._synchronize_prevalidated_codebook_scope_failure(
+                local_preamble_error, "initialize preamble", scope_binding
+            )
+        elif local_preamble_error is not None:
+            raise local_preamble_error
         self._validate_codebook_scope_operation_header("initialize")
         try:
             _assert_optimizer_gradients_structurally_valid(
@@ -8557,9 +8573,17 @@ def step(self, closure=None):
             raise local_preamble_error
 
         # The closure can replace a finalized parameter or otherwise invalidate
-        # the runtime binding. Recheck before the operation header and synchronize
-        # structural failures before any peer enters a scoped codebook collective.
+        # the runtime binding. It can also rebind (or clear) the runtime
+        # process-group between capture and the operation header; a rank that
+        # silently swapped to None or a different binding would enter a
+        # different header collective than its peers and deadlock. Recheck
+        # against the captured binding and synchronize structural failures
+        # before any peer enters a scoped codebook collective.
         try:
+            if self._gefen_codebook_process_group is not scope_binding:
+                raise RuntimeError(
+                    "Gefen codebook process-group binding changed during the step preamble"
+                )
             self._assert_finalized_binding_layout()
             self._assert_runtime_codebook_process_group()
             _assert_optimizer_gradients_structurally_valid(self)

src/gefen/hybrid.py

@@ -997,6 +997,14 @@ def _hybrid_layout_forensics_fast_tokens(self):
             self._subopts,
             len(self._subopts),
         ]
+        # Capture the owner registry contents by identity (keys, then the
+        # (parameter, child) values) so an in-place entry replacement that
+        # preserves the dict object and its length -- which the identity/len
+        # tokens above cannot see -- flips a token and forces a full rebuild
+        # instead of reusing the stale verdict. Mirrors the _param_names
+        # snapshot in the base fast-token path.
+        live.extend(self._state_param_owner.keys())
+        live.extend(self._state_param_owner.values())
         for child in self._subopts:
             live.append(child._finalized_binding_layout_matches())
             GefenMuonHybrid._hybrid_child_param_group_tokens(child, live)

tests/_state_snapshot.py

@@ -34,6 +34,21 @@
 _CHILD_REGISTRY_ATTRS = ("_param_names", "_gefen_shard_bindings")
 
 
+# Per-device cache dicts each Gefen/GefenMuon child keeps live across steps.
+# Their tensors are also reachable from ``__dict__``, but the tensor-value
+# snapshot keeps a strong reference to each tensor, so an in-place ``clear()`` or
+# entry removal leaves every retained tensor still matching its clone while the
+# cache membership silently vanishes. Recording the dict identity plus its keyed
+# contents makes that membership change fail the assertion. Absent on optimizer
+# types that never build these caches (skipped below).
+_CHILD_CACHE_ATTRS = (
+    "_gefen_codebook_by_device",
+    "_gefen_codebook_lut_by_device",
+    "_sr_seed_by_device",
+    "_gefen_global_step_by_device",
+)
+
+
 def _cloned(value):
     if torch.is_tensor(value):
         return value.detach().clone()
@@ -133,10 +148,29 @@ def _registry_snapshot(optimizer):
     return registries
 
 
+def _cache_snapshot(optimizer):
+    """Live dict identity plus cloned contents of each per-device cache.
+
+    Returns ``name -> (live_dict, cloned_contents)`` for every cache present as a
+    ``dict`` on ``optimizer``. The live dict is kept by reference so the assertion
+    can confirm it was not replaced, and its contents are cloned so a later
+    ``clear()`` or entry removal (or an in-place value edit) produces a mismatch.
+    """
+
+    caches = {}
+    for name in _CHILD_CACHE_ATTRS:
+        cache = getattr(optimizer, name, None)
+        if type(cache) is not dict:
+            continue
+        caches[name] = (cache, _cloned(cache))
+    return caches
+
+
 def deep_state_snapshot(optimizer):
     return {
         "attributes": optimizer.__dict__.copy(),
         "registries": _registry_snapshot(optimizer),
+        "caches": _cache_snapshot(optimizer),
         "state": optimizer.state,
         "state_items": tuple(
             (parameter, state, _cloned(dict(state)))
@@ -194,3 +228,7 @@ def assert_deep_state_snapshot(optimizer, snapshot):
         ):
             assert live_key is expected_key
             _nested_equal(live_value, expected_value)
+    for name, (expected_ref, expected_contents) in snapshot["caches"].items():
+        live = getattr(optimizer, name, None)
+        assert live is expected_ref
+        _nested_equal(live, expected_contents)

tests/test_checkpoint_binding.py

@@ -46,6 +46,30 @@ def test_checkpoint_binding_validates_descriptor_types_membership_and_device():
         CheckpointProcessGroupBinding(identity, "worker:0", None, torch.device("meta"))
 
 
+def test_validate_collective_device_allows_gloo_cuda_but_requires_available_device():
+    identity = ProcessGroupIdentity("local", ("worker:0",))
+    unavailable_index = (
+        torch.cuda.device_count() if torch.cuda.is_available() else 0
+    )
+    cuda_binding = CheckpointProcessGroupBinding(
+        identity, "worker:0", None, torch.device("cuda", unavailable_index)
+    )
+    # Gloo supports CUDA tensors, so a CUDA collective device is not rejected on
+    # backend grounds; it must instead fail the later availability check when the
+    # configured device does not exist.
+    with pytest.raises(ValueError, match="CUDA device is unavailable"):
+        cuda_binding._validate_collective_device("gloo")
+
+    cpu_binding = CheckpointProcessGroupBinding(
+        identity, "worker:0", None, torch.device("cpu")
+    )
+    # NCCL still requires a CUDA collective device.
+    with pytest.raises(ValueError, match="runtime backend"):
+        cpu_binding._validate_collective_device("nccl")
+    # Gloo with a CPU device remains valid.
+    cpu_binding._validate_collective_device("gloo")
+
+
 def test_checkpoint_binding_requires_explicit_multi_member_handle_and_local_singleton():
     singleton = ProcessGroupIdentity("local", ("worker:0",))
     multiple = ProcessGroupIdentity("data", ("worker:0", "worker:1"))
@@ -111,16 +135,23 @@ def _distributed_binding_worker(rank, world_size, init_file, queue):
             "world size",
         )
 
-        wrong_device_binding = CheckpointProcessGroupBinding(
+        # Gloo accepts CUDA tensors, so a CUDA collective device is not rejected
+        # on backend grounds; it must instead fail the later availability check
+        # when the configured device does not exist. An index at or beyond the
+        # visible device count (0 when no CUDA is present) is always unavailable.
+        unavailable_index = (
+            torch.cuda.device_count() if torch.cuda.is_available() else 0
+        )
+        unavailable_device_binding = CheckpointProcessGroupBinding(
             world_identity,
             world_members[rank],
             dist.group.WORLD,
-            torch.device("cuda:0"),
+            torch.device("cuda", unavailable_index),
         )
-        device_mismatch_rejected = _expect_rejection(
-            wrong_device_binding.validate_runtime,
+        cuda_unavailable_rejected = _expect_rejection(
+            unavailable_device_binding.validate_runtime,
             ValueError,
-            "runtime backend",
+            "CUDA device is unavailable",
         )
 
         subgroup_global_ranks = (0, 2)
@@ -162,7 +193,7 @@ def _distributed_binding_worker(rank, world_size, init_file, queue):
                 "world_validated": world_validated,
                 "order_mismatch_rejected": order_mismatch_rejected,
                 "size_mismatch_rejected": size_mismatch_rejected,
-                "device_mismatch_rejected": device_mismatch_rejected,
+                "cuda_unavailable_rejected": cuda_unavailable_rejected,
                 "subgroup_validated": subgroup_validated,
                 "nonmember_rejected": nonmember_rejected,
                 "uninitialized_rejected": uninitialized_rejected,
@@ -221,7 +252,7 @@ def test_checkpoint_binding_validates_real_world_and_subgroup_membership():
     assert all(item["world_validated"] for item in results)
     assert all(item["order_mismatch_rejected"] for item in results)
     assert all(item["size_mismatch_rejected"] for item in results)
-    assert all(item["device_mismatch_rejected"] for item in results)
+    assert all(item["cuda_unavailable_rejected"] for item in results)
     assert all(item["subgroup_validated"] for item in results if item["rank"] in (0, 2))
     assert results[1]["nonmember_rejected"]
     assert all(item["uninitialized_rejected"] for item in results)

tests/test_codebook_scope_cpu.py

@@ -135,6 +135,23 @@ def _replace_native_guard(checkpoint, guard):
         group["_gefen_checkpoint_metadata"]["native_local_shards"] = copy.deepcopy(guard)
 
 
+def test_deep_snapshot_detects_per_device_cache_membership_clear():
+    # A per-device cache clear preserves the dict attribute identity, and the
+    # retained cache tensors still match their tensor-value clones, so only a
+    # membership-aware snapshot catches the removal.
+    parameter = torch.nn.Parameter(torch.randn(4, 4))
+    optimizer = Gefen([("Layer.Weight", parameter)], fused=False)
+    device = torch.device("cpu")
+    optimizer._gefen_codebook_by_device[device] = torch.randn(8)
+
+    snapshot = deep_state_snapshot(optimizer)
+    assert_deep_state_snapshot(optimizer, snapshot)  # unchanged membership passes
+
+    optimizer._gefen_codebook_by_device.clear()
+    with pytest.raises(AssertionError):
+        assert_deep_state_snapshot(optimizer, snapshot)
+
+
 def test_codebook_process_group_binding_is_public_frozen_and_ordered():
     group = ProcessGroupIdentity("replica", ("worker:b", "worker:a"))
     binding = CodebookProcessGroupBinding(group, "worker:b", object(), torch.device("cpu"))