Fix post-merge DCP resharding review findings + harden the flaky handoff test

COMPATIBILITY.md

@@ -105,7 +105,14 @@ Native single-process optimizer `state_dict()`/`load_state_dict()` and the expli
 
 The native format is deliberately same-topology only and currently requires one 1-D DeviceMesh spanning the default process-group world. It is bit-exact on resume and is the recommended path when the world size does not change. Plain Gefen additionally exposes `GefenDCPState`, a standalone `torch.distributed.checkpoint` Stateful wrapper whose purpose is **resharding**: a checkpoint saved on N ranks can be loaded on M ranks for one-dimensional default-world `Shard(0)` DTensors. Save dequantizes the quantized momentum against each rank's learned codebook into dense `Shard(0)` DTensors, so DCP can reshard them without gathering the full optimizer state onto one rank. Load validates every field before mutating the live optimizer, then re-blocks the resharded dense momentum back into Gefen's compact per-block state on each new local shard — re-running the block-variance period search, relearning the exact codebook, and re-quantizing — so the restored optimizer keeps Gefen's ~1 byte/param footprint rather than expanding to per-element state.
 
-Because the momentum is routed through a dense reshard and re-blocked against a freshly learned per-shard codebook, resume is a **correct, finite continuation within 256-level quantization noise, not a bit-exact restore** — even at the same world size. Same-topology resumes should therefore use the native `state_dict` path above. `GefenDCPState` requires `factored_v_2d=False` and `capturable=False` (the factored row/column second moment is not shard-addressable, and capturable/compiled counter and seed semantics are not host-serializable); Muon, Hybrid, `sharded_mode`, multidimensional meshes, subgroups, non-`Shard(0)` placements, and the native full-state format remain same-topology and fail closed.
+Resume through `GefenDCPState` is a **correct, finite continuation, not a bit-exact restore** — even at the same world size. It approximates in two independent ways, and only the first is bounded by quantization noise:
+
+- **Momentum re-quantization.** The momentum is routed through a dense reshard and re-quantized against a codebook relearned on the new shard, so it returns within 256-level quantization noise (measured ~2e-4 relative on the covered shards).
+- **Second-moment re-aggregation.** Gefen never stores a per-element second moment: it keeps one `vmean` per block. Save can therefore only write each source block's `vmean` repeated across its elements, and load averages those repeated values into the *target* blocks. When the target blocking refines or matches the source blocking, every target block lies inside one source block and the original `vmean` comes back to fp32 round-off. When a target block instead spans several source blocks, their distinct `vmean`s are averaged together — and because the per-element history they came from never existed, that averaging is **irreversible**. Its size is set by how much `vmean` varies across the merged blocks, which is data-dependent and can be far larger than the momentum quantization noise (a 2x-coarser blocking measured ~29% relative on a synthetic shard with realistic block-to-block spread).
+
+Block geometry is re-derived per shard, so a world-size change that changes a shard's chosen period is what triggers the second effect; a reshard that keeps the period is subject only to the first. Same-topology resumes should use the bit-exact native `state_dict` path above.
+
+`GefenDCPState` requires `factored_v_2d=False` and `capturable=False` (the factored row/column second moment is not shard-addressable, and capturable/compiled counter and seed semantics are not host-serializable); Muon, Hybrid, `sharded_mode`, multidimensional meshes, subgroups, non-`Shard(0)` placements, and the native full-state format remain same-topology and fail closed.
 
 ## Transformers Trainer DDP
 

README.md

@@ -126,7 +126,7 @@ Gefen drops into standard distributed training like any other PyTorch optimizer,
 | DeepSpeed ZeRO 1-3 | Plain Gefen (client optimizer); optimizer CPU-offload verified at ZeRO-2/3, parameter offload at ZeRO-3 (full fine-tune and LoRA); 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.** `GefenDCPState` is for **resharding** plain Gefen: wrap the optimizer with it for `torch.distributed.checkpoint`, and a checkpoint saved on N ranks loads on M ranks for a one-dimensional default-world `Shard(0)` mesh. Load re-blocks the resharded state back to Gefen's compact ~1 byte/param form, so resume is a correct continuation within quantization noise rather than a bit-exact restore; same-topology resumes should use the bit-exact native full-state path (Muon `approx` too). Requires `factored_v_2d=False` and `capturable=False` — [full details](https://github.com/thad0ctor/Gefen-X/blob/main/COMPATIBILITY.md#optimizer-checkpoint-scope).
+> **FSDP2 checkpoint scope.** `GefenDCPState` is for **resharding** plain Gefen: wrap the optimizer with it for `torch.distributed.checkpoint`, and a checkpoint saved on N ranks loads on M ranks for a one-dimensional default-world `Shard(0)` mesh. Load re-blocks the resharded state back to Gefen's compact ~1 byte/param form, so resume is an approximate continuation, not a bit-exact restore: the momentum is re-quantized, and where the new shard's block boundaries do not line up with the old one's, the second moment is re-aggregated and some of its history is lost for good. Use the bit-exact native full-state path whenever the topology is unchanged (Muon `approx` too). Requires `factored_v_2d=False` and `capturable=False` — [full details](https://github.com/thad0ctor/Gefen-X/blob/main/COMPATIBILITY.md#optimizer-checkpoint-scope).
 >
 > **FSDP2 CPU offload.** Training-time CPU offload via `CPUOffloadPolicy` (`fully_shard(module, offload_policy=CPUOffloadPolicy())`) is validated for plain Gefen on single and multiple GPUs: each rank steps its CPU-resident local shard directly (the codebook is learned rank-locally, with no cross-rank codebook collective), and the multi-GPU run completes on an NCCL-only process group.
 

src/gefen/dcp.py

@@ -4,20 +4,10 @@
 
 from dataclasses import dataclass
 import math
-import re
 
 import torch
 
 
-# Names Gefen synthesizes positionally when the caller passes no parameter name
-# (``group_<gi>_param_<pi>`` for a multi-param group, ``param_<n>`` for an
-# implicit/single group). Such names encode *registration order*, not stable
-# parameter identity, so two differently ordered reconstructions produce the
-# same identity list and would cross-assign momentum on load. DCP resharding
-# therefore requires caller-provided names and rejects the synthesized pattern.
-_SYNTHESIZED_NAME_RE = re.compile(r"^(group_\d+_param_\d+|param_\d+)$")
-
-
 # Schema history:
 #   v1 -- counters + dense momentum/second-moment slots only.
 #   v2 -- adds per-slot parameter identities (name + group + global shape) and
@@ -68,24 +58,53 @@ def _dtensor_from_local(local: torch.Tensor, parameter):
 
 
 def _dense_momentum(optimizer, parameter, state) -> torch.Tensor:
-    indices = state["m_codebook"].reshape(-1).long()
+    """Dequantize one shard's block state into dense fp32 momentum.
+
+    Uses Gefen's chunked dequantize rather than ``codebook[indices.long()]``.
+    Advanced indexing makes a full-size int64 copy of the uint8 indices (8N
+    bytes) plus a full-size fp32 gather (4N) on top of the 4N result, so the
+    transient peaked at ~16 bytes/param -- ~16 GiB on a 1B-param local shard,
+    enough to OOM an otherwise-viable save in the one optimizer whose premise is
+    memory efficiency. ``gefen_dequantize_unpacked_indices`` gathers bounded
+    element chunks straight into the preallocated output, and the magnitude
+    scaling is applied in place, so only the 4N result plus one bounded chunk is
+    live. The gathered values, the broadcast multiply, and therefore the saved
+    dense momentum are bit-for-bit what advanced indexing produced.
+    """
+    from gefen.gefen import gefen_dequantize_unpacked_indices
+
+    stored = state["m_codebook"].reshape(-1)
     magnitude = state["m_magnitude"].reshape(-1).float()
     codebook = optimizer._gefen_codebook
     if codebook is None:
         raise RuntimeError("Gefen DCP save requires an initialized codebook")
-    codebook = codebook.detach().to(device=indices.device, dtype=torch.float32)
+    codebook = codebook.detach().to(device=stored.device, dtype=torch.float32)
     period = _counter(state["automatic_period"], "automatic_period")
-    if period < 1 or indices.numel() != magnitude.numel() * period:
+    if period < 1 or stored.numel() != magnitude.numel() * period:
         raise ValueError("Gefen momentum block geometry is invalid")
+    # An empty fp32 exemplar: the helper takes the output dtype/device from it
+    # (and stays a plain tensor, so no DTensor is reconstructed here).
+    like = torch.empty(0, dtype=torch.float32, device=stored.device)
+    dense = gefen_dequantize_unpacked_indices(codebook, stored, like)
     return (
-        codebook[indices]
-        .reshape(-1, period)
-        .mul(magnitude.reshape(-1, 1))
+        dense.reshape(-1, period)
+        .mul_(magnitude.reshape(-1, 1))
         .reshape(_local(parameter).shape)
     )
 
 
 def _dense_second_moment(parameter, state) -> torch.Tensor:
+    """Expand one shard's per-block ``vmean`` into a dense per-element tensor.
+
+    This is a *repeat*, not a reconstruction: Gefen keeps one second-moment value
+    per block, so the per-element history this dense form implies never existed.
+    Load averages these repeated values into the target blocks, which is exact
+    (to fp32 round-off) only while each target block stays inside one source
+    block -- i.e. when the target blocking matches or refines the source's. A
+    coarser target block averages several source blocks' distinct vmeans
+    together, and that loss is irreversible. See the resume-error note in
+    COMPATIBILITY.md.
+    """
     period = _counter(state["automatic_period"], "automatic_period")
     vmean = state["vmean"].reshape(-1).float()
     local_numel = _local(parameter).numel()
@@ -270,9 +289,9 @@ def _validate_layout(self):
             names = group.get("param_names")
             # Index-addressed slots are only safe to reshard if each carries a
             # caller-stable, unique identity. A missing/short param_names list or
-            # a positionally synthesized name (Gefen's default for unnamed
-            # parameters) encodes registration order, not identity, so refuse it
-            # rather than silently cross-assigning momentum on a reordered load.
+            # a name Gefen synthesized for an unnamed parameter encodes
+            # registration order, not identity, so refuse it rather than silently
+            # cross-assigning momentum on a reordered load.
             if names is None or len(names) != len(group["params"]):
                 raise RuntimeError(
                     "GefenDCPState requires caller-provided stable param_names on "
@@ -285,13 +304,19 @@ def _validate_layout(self):
                     )
                 )
             for name, parameter in zip(names, group["params"]):
-                if _SYNTHESIZED_NAME_RE.match(str(name)):
+                # Ask whether Gefen *generated* this name, rather than matching
+                # its spelling: the synthesized forms are ordinary identifiers a
+                # model may legitimately use for a real parameter, and rejecting
+                # a caller's own "param_0" locked such an optimizer out of
+                # GefenDCPState entirely. Provenance is recorded at registration.
+                if self.optimizer._param_name_is_synthesized(parameter):
                     raise RuntimeError(
                         "GefenDCPState requires caller-provided stable parameter "
-                        "names, but slot {} carries the synthesized positional "
-                        "name {!r}. Construct the optimizer with explicit "
-                        "parameter names so resharded state is addressed by "
-                        "identity, not registration order.".format(index, str(name))
+                        "names, but slot {} carries the positional name {!r} that "
+                        "Gefen synthesized for an unnamed parameter. Construct the "
+                        "optimizer with explicit parameter names so resharded "
+                        "state is addressed by identity, not registration "
+                        "order.".format(index, str(name))
                     )
                 if not _is_dtensor(parameter):
                     raise RuntimeError(
@@ -356,6 +381,16 @@ def _key(slot: _Slot, field: str) -> str:
         "vmean",
     )
 
+    # Fields whose presence means the slot was materialized (or partially so).
+    # ``vmean_step`` is deliberately NOT *required*: a checkpoint written before
+    # the counter existed carries vmean without it, and native Gefen backfills it
+    # from ``step`` at step time (the save below mirrors that backfill). But its
+    # presence still marks a materialize, so a slot carrying only a name and an
+    # orphaned vmean_step is partial -- not fresh. Classifying it as fresh wrote
+    # an "uninitialized" slot with a nonzero counter, which every later load
+    # rejects as incoherent: a silently dead checkpoint.
+    _MATERIALIZED_STATE_FIELDS = _REQUIRED_STATE_FIELDS + ("vmean_step",)
+
     def _identities(self):
         """Stable per-slot identity list (replicated, in slot order).
 
@@ -468,6 +503,43 @@ def _validate_hyper_entry(group_index, entry):
             )
         return parsed
 
+    @staticmethod
+    def _validate_hyper_destinations(group_index, group, entry):
+        """Reject a live tensor hyperparameter that cannot represent its value.
+
+        The commit below fills a tensor hyperparameter in place to preserve its
+        identity/device for any fused kernel holding a reference. ``fill_`` casts
+        to the destination dtype, so restoring a fractional checkpoint lr (1e-3)
+        into a one-element *integral* lr tensor truncates it to 0 and silently
+        freezes every subsequent update. Checking representability here -- during
+        staging, before anything is committed -- keeps that failure fail-atomic
+        instead of discovering it after the state was already replaced.
+        Floating-point destinations are accepted: rounding an fp64 python float
+        into an fp32/bf16 lr tensor is ordinary, pre-existing behavior that
+        matches the native path.
+        """
+        for key in _GROUP_HYPER_KEYS:
+            current = group.get(key)
+            if not torch.is_tensor(current) or current.is_floating_point():
+                continue
+            value = entry[key]
+            if float(value) != float(int(value)):
+                raise ValueError(
+                    "Gefen DCP checkpoint group {} restores {}={!r}, but this "
+                    "optimizer holds it in a {} tensor that cannot represent it "
+                    "(the in-place fill would truncate to {}). Rebuild the "
+                    "optimizer with a floating-point {} tensor (or a plain "
+                    "float) before resuming; refusing to commit a silently "
+                    "truncated value.".format(
+                        group_index,
+                        key,
+                        value,
+                        current.dtype,
+                        int(value),
+                        key,
+                    )
+                )
+
     def _slot_initialized(self, slot: _Slot) -> bool:
         """Classify a slot's live state; reject a partial/malformed materialize.
 
@@ -479,13 +551,11 @@ def _slot_initialized(self, slot: _Slot) -> bool:
         state = self.optimizer.state.get(slot.parameter)
         if not state:
             return False
-        present = [key for key in self._REQUIRED_STATE_FIELDS if key in state]
+        present = [key for key in self._MATERIALIZED_STATE_FIELDS if key in state]
         if not present:
             return False
-        if len(present) != len(self._REQUIRED_STATE_FIELDS):
-            missing = [
-                key for key in self._REQUIRED_STATE_FIELDS if key not in state
-            ]
+        missing = [key for key in self._REQUIRED_STATE_FIELDS if key not in state]
+        if missing:
             raise ValueError(
                 "Gefen DCP slot {} ({}) has partial optimizer state; missing "
                 "{}. Refusing to save it as uninitialized (which would zero the "
@@ -603,6 +673,12 @@ def load_state_dict(self, state_dict):
             self._validate_hyper_entry(group_index, entry)
             for group_index, entry in enumerate(saved_hypers)
         ]
+        # Also verify every live destination can hold the value it is about to be
+        # given, while the commit is still ahead of us (see the helper).
+        for group_index, (group, entry) in enumerate(
+            zip(self.optimizer.param_groups, staged_hypers)
+        ):
+            self._validate_hyper_destinations(group_index, group, entry)
 
         # Reject a deterministic-policy mismatch instead of silently overwriting
         # it, matching native Gefen.load_state_dict: the flag changes fused-routing