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

COMPATIBILITY.md

@@ -105,7 +105,35 @@ 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, and that `vmean` is an EMA of the block's mean squared gradient. Save can therefore only write each source block's `vmean` repeated across its elements, and load re-aggregates those values onto the *target* blocks. Because both the EMA and the block mean are linear, a target block that is an exact union of whole source blocks reproduces what a natively blocked run at the target period would hold, to fp32 round-off (measured ~3e-7) — **coarsening onto source boundaries costs nothing**. The lossy direction is the opposite one: when a target block falls strictly inside a source block, or straddles a source boundary, the single stored `vmean` can only be spread flat across sub-blocks that a native run would have given distinct values. That sub-block detail was never in the checkpoint, so it cannot be recovered (a 4x-finer blocking measured ~16% relative on a synthetic shard with realistic within-block variation). Its size is set by how much the true second moment varies *inside* a source block, which is data-dependent and can be far larger than the momentum quantization noise.
+
+Block geometry is re-derived per shard, so a world-size change that moves a shard's chosen period off the source's block boundaries is what triggers the second effect; a reshard that keeps the period, or that only merges whole blocks, 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, and non-`Shard(0)` placements fail closed. The native full-state format is same-topology with one carve-out: a `GefenMuon(sharded_mode="distributed")` checkpoint consolidates every owner's complete state onto every rank, so it resumes under a different world size — including into a single-process optimizer that never initializes a process group. That portability is replication, not resharding: each rank materializes the full optimizer state instead of addressing its own shard, so it does not buy what `GefenDCPState` buys.
+
+Pass `GefenSavePlanner` to `dcp.save` so the save stays memory-bounded. Resharding forces the dense expansion above — 4 bytes/param of momentum plus 4 bytes/param of second moment, eight times what the optimizer actually holds — and without the planner every slot's dense pair is expanded up front and kept live for the whole write, peaking higher than a training step. The planner expands one slot at a time, at the moment the writer asks for it, so the save peaks at the resident state plus a single slot's dense form:
+
+```python
+from gefen import Gefen, GefenDCPState, GefenSavePlanner
+import torch.distributed.checkpoint as dcp
+
+dcp.save(
+    {"model": model, "optimizer": GefenDCPState(optimizer)},
+    storage_writer=dcp.FileSystemWriter(path),
+    planner=GefenSavePlanner(),          # required for the memory bound
+)
+dcp.load(                                 # load needs no planner
+    {"model": model, "optimizer": GefenDCPState(optimizer)},
+    storage_reader=dcp.FileSystemReader(path),
+)
+```
+
+On a 64M-param local shard the peak over the optimizer's resident state falls from 12.0 to 6.0 bytes/param for one large tensor, and from 8.25 to 0.50 bytes/param for a 32-tensor model; the checkpoint bytes are unchanged. The bound is per *slot*, so the win grows as the state is split across more parameters, and the remaining transient is a single slot's dense form plus up to ~128 MiB of dequantize scratch. That scratch is a ceiling, not a fixed cost: it is bounded by the dequantize gather chunk (`GEFEN_DEQUANT_GATHER_CHUNK`, 8M elements), so only slots larger than the chunk pay the full ~128 MiB, and a smaller slot pays proportionally less — ~24 MiB per 2M-param slot in the 32-tensor case above, which is what makes 0.50 bytes/param reachable there. Omitting the planner is not a correctness bug — the save still writes exactly the same checkpoint — it only costs the old 8 bytes/param. `GefenSavePlanner` subclasses `DefaultSavePlanner` and defers to it for everything that is not Gefen optimizer state, so the model and any other Stateful items in the same state dict are unaffected.
+
+Save through `GefenDCPState` must be the synchronous `dcp.save`. `dcp.async_save` is **not supported**: its CPU staging copies the whole state dict up front, which is the all-slots-live dense expansion the save bound exists to avoid. Where that copy happens the save fails loudly with a `RuntimeError` and writes nothing — that is every save of CUDA-resident state, and every save on the torch versions whose staging builds its CPU copy with `zeros_like`. The one gap is torch 2.5 staging state that is *already* on the CPU: `tensor.to(cpu)` short-circuits to the same object, so nothing is copied, the call is not refused, and it writes a correct checkpoint by doing the whole expansion in the writer thread. Do not rely on that — nothing was staged, so the write is not actually asynchronous with respect to the optimizer, and resuming training while it runs races the writer against your own state. Use `dcp.save`, or the native full-state path for a same-topology resume.
 
 ## Transformers Trainer DDP
 

README.md

@@ -121,12 +121,12 @@ Gefen drops into standard distributed training like any other PyTorch optimizer,
 |---|---|
 | DDP | All optimizers |
 | FSDP2 | All optimizers; training-time CPU offload (`CPUOffloadPolicy`) validated for plain Gefen, single and multi-GPU |
-| FSDP2 checkpoints | Plain Gefen can use `GefenDCPState` to reshard N ranks to M; same-topology resumes use the bit-exact native full-state path (Muon `approx` too) |
+| FSDP2 checkpoints | Plain Gefen can use `GefenDCPState` + `GefenSavePlanner` to reshard N ranks to M; same-topology resumes use the bit-exact native full-state path (Muon `approx` too) |
 | 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 (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 pass `planner=GefenSavePlanner()` to the synchronous `dcp.save` (which keeps the save memory-bounded; `dcp.load` needs no planner, and `dcp.async_save` is not supported), 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 blocks are finer than the old one's (or straddle its boundaries), the second moment is re-aggregated and detail the checkpoint never stored cannot be recovered. 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.
 
@@ -636,7 +636,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 resharding is explicit.** Use `dcp.save({"optimizer": GefenDCPState(optimizer)}, ...)` and the matching `dcp.load` call to reshard plain Gefen across world sizes on a one-dimensional default-world `Shard(0)` mesh (`factored_v_2d=False`, `capturable=False`). Load re-blocks the state to Gefen's compact form and continues within quantization noise; for a same-topology resume use the bit-exact native path instead. Muon, Hybrid, other placements, and the native full-state format remain same-topology and fail closed — [details](https://github.com/thad0ctor/Gefen-X/blob/main/COMPATIBILITY.md#optimizer-checkpoint-scope).
+- **FSDP2 optimizer resharding is explicit.** Use `dcp.save({"optimizer": GefenDCPState(optimizer)}, storage_writer=..., planner=GefenSavePlanner())` and the matching `dcp.load` call (no planner needed) to reshard plain Gefen across world sizes on a one-dimensional default-world `Shard(0)` mesh (`factored_v_2d=False`, `capturable=False`). Resharding forces a dense expansion of the optimizer state, and `GefenSavePlanner` is what keeps it bounded — it expands one parameter at a time as the writer asks for it, instead of holding every parameter's dense form for the whole write. The save must be the synchronous `dcp.save`; `dcp.async_save` is not supported, because its CPU staging would copy the whole state dict up front. Load re-blocks the state to Gefen's compact form and continues correctly but not bit-exactly, even at the same world size: momentum returns within quantization noise, and the second moment is re-aggregated onto the target blocks, which costs nothing when they merge whole source blocks but cannot recover detail when they split or straddle one. For a same-topology resume use the bit-exact native path instead. Muon, Hybrid, and other placements fail closed; the native full-state format is same-topology except a `GefenMuon(sharded_mode="distributed")` checkpoint, which carries every owner's complete state on every rank and so resumes under a different world size — [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

src/gefen/__init__.py

@@ -12,6 +12,7 @@
     "GefenMuon",
     "GefenMuonHybrid",
     "GefenDCPState",
+    "GefenSavePlanner",
     "split_params_for_muon",
     "validate_split",
     "kernels",
@@ -36,6 +37,10 @@ def __getattr__(name):
         from .dcp import GefenDCPState
 
         return GefenDCPState
+    if name == "GefenSavePlanner":
+        from .dcp import GefenSavePlanner
+
+        return GefenSavePlanner
     if name in ("split_params_for_muon", "validate_split"):
         from . import params