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

COMPATIBILITY.md

@@ -105,7 +105,57 @@ 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's GPU memory 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 from a synchronous `dcp.save` is not a correctness bug — the save still writes exactly the same checkpoint — it only costs the old 8 bytes/param. That is a statement about `dcp.save` alone: `dcp.async_save` additionally requires `GefenFileSystemWriter` (below) and fails without it, because staging, unlike writing, cannot resolve a stand-in through the planner at all. `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.
+
+The bound holds at any `FileSystemWriter(thread_count=N)`. N writer threads ask for N write items at once, so the planner serializes slot expansion and hands each thread its shard already on the CPU, keeping one slot's dense form on the device rather than one per thread. Expanding slots therefore does not run in parallel across writer threads; their file writes still do, and a threaded save stays faster than a single-threaded one.
+
+That bound is on GPU memory. Host memory during a save belongs to the writer, and torch's `FileSystemWriter` keeps every tensor it has written until it closes the file, so the host holds the full dense state dict by the end of a save — the cost of writing a resharded checkpoint at all, with or without the planner. What the planner bounds is the page-locked share of it: slot copies pass through one reusable page-locked buffer and are handed to the writer in pageable memory, so a save page-locks a single slot's dense form no matter how many threads the writer runs. The shards the writer accumulates are ordinary pageable memory the kernel may reclaim or swap, rather than page-locked memory it never can — but they still count against host memory, so a save needs room for the full dense state dict on the host regardless.
+
+`dcp.async_save` writes a `GefenDCPState` when given `storage_writer=GefenFileSystemWriter(path)`:
+
+```python
+from gefen import GefenFileSystemWriter
+
+response = dcp.async_save(
+    {"model": model, "optimizer": GefenDCPState(optimizer)},
+    storage_writer=GefenFileSystemWriter(path),
+    planner=GefenSavePlanner(),
+)
+# ... training continues ...
+# torch 2.5 returns a bare Future; newer torch wraps it in an AsyncSaveResponse.
+# Waiting is what surfaces a background write failure.
+getattr(response, "upload_completion", response).result()
+```
+
+Async save returns once the state dict has been *staged* — copied to CPU — and writes it from a background thread, which is what lets training resume against a checkpoint that still reflects the moment of the call. `GefenFileSystemWriter` stages Gefen's slots one at a time: it expands a slot's dense form, copies it to CPU, and releases it before expanding the next, so the device holds one slot's dense form during staging rather than every slot's at once. On the 32-tensor model above that is a **0.50 bytes/param** GPU peak over the optimizer's resident state — the same bound `GefenSavePlanner` gives the synchronous save. CPU RAM holds the full dense snapshot for the duration of the write, which is what async staging costs for any optimizer. Everything that is not Gefen optimizer state is staged by `FileSystemWriter` unchanged.
+
+The writer is required because staging has to copy the state dict, and `GefenDCPState` hands DCP zero-storage stand-ins that have nothing to copy. Given the default writer, `async_save` raises a `RuntimeError` and writes nothing, on every supported torch version and on both CPU- and CUDA-resident state. The async checkpoint's file *layout* differs from the synchronous save's — staging sets the writer's `per_thread_copy_ahead` to 0, which orders the write items by size — and it restores identically.
 
 ## Transformers Trainer DDP
 

README.md

@@ -121,12 +121,58 @@ 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).
+### Resharding an FSDP2 checkpoint (N ranks → M)
+
+Opt in by wrapping plain Gefen in `GefenDCPState`. Save on 2 ranks, load on 4 — or on 1.
+
+FSDP2 resharding needs torch **2.6+**, where `fully_shard` is public (on 2.5 it lives at `torch.distributed._composable.fsdp`).
+
+```python
+from gefen import Gefen, GefenDCPState, GefenFileSystemWriter, GefenSavePlanner
+from torch.distributed.fsdp import fully_shard
+import torch.distributed.checkpoint as dcp
+
+fully_shard(model)                            # params must be Shard(0) DTensors
+optimizer = Gefen(model.named_parameters(), lr=3e-5, factored_v_2d=False)
+state = {"model": model, "optimizer": GefenDCPState(optimizer)}
+
+dcp.save(
+    state,
+    storage_writer=dcp.FileSystemWriter(path),
+    planner=GefenSavePlanner(),   # optional — see below. Drop it and the save
+)                                 # still writes the same checkpoint, using more memory
+
+dcp.load(state, storage_reader=dcp.FileSystemReader(path))   # never needs a planner
+```
+
+Requires `factored_v_2d=False` and `capturable=False`.
+
+**`dcp.async_save` needs `GefenFileSystemWriter`.** Async save returns as soon as the state is staged and writes from a background thread, so it has to snapshot the state dict to CPU first. `GefenFileSystemWriter` is what stages Gefen's slots — expanding each one's dense form, copying it to CPU, and releasing it before the next — so training can resume immediately and the checkpoint still reflects the moment of the call.
+
+```python
+response = dcp.async_save(
+    state,
+    storage_writer=GefenFileSystemWriter(path),
+    planner=GefenSavePlanner(),
+)
+
+# ... keep training; the checkpoint holds the state as of the call above ...
+
+# torch 2.5 returns a bare Future; newer torch wraps it in an AsyncSaveResponse.
+# Waiting is what surfaces a background write failure.
+getattr(response, "upload_completion", response).result()
+```
+
+**`GefenSavePlanner` is optional and recommended.** Resharding has to expand Gefen's compact state to dense fp32 to be portable across world sizes. Without the planner, every parameter's dense form is held for the whole write; with it, one is expanded at a time. On a 32-tensor model that is the difference between **8.25 and 0.50 bytes/param** of peak memory over the optimizer's resident state — enough that a model which trains fine could otherwise OOM while saving.
+
+> **Scope.** For *changing* topology only: resume is a correct continuation, not a bit-exact restore. Same-topology resumes should use the native full-state path (Muon `approx` too), which is bit-exact — [full details](https://github.com/thad0ctor/Gefen-X/blob/main/COMPATIBILITY.md#optimizer-checkpoint-scope).
+>
+> **This is a library-level API — you call `dcp.save` yourself.** There is no config key for it in Axolotl: Axolotl and Accelerate do their own checkpointing and pass DCP's default planner, so `GefenDCPState` never enters that path.
 >
 > **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 +682,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 and approximate.** Resharding plain Gefen across world sizes needs `GefenDCPState` + `GefenSavePlanner` ([usage](#resharding-an-fsdp2-checkpoint-n-ranks--m)), and resume is a correct continuation rather than a bit-exact restore. Muon, Hybrid, and other placements fail closed; `dcp.async_save` needs `GefenFileSystemWriter` — [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,8 @@
     "GefenMuon",
     "GefenMuonHybrid",
     "GefenDCPState",
+    "GefenSavePlanner",
+    "GefenFileSystemWriter",
     "split_params_for_muon",
     "validate_split",
     "kernels",
@@ -36,6 +38,14 @@ def __getattr__(name):
         from .dcp import GefenDCPState
 
         return GefenDCPState
+    if name == "GefenSavePlanner":
+        from .dcp import GefenSavePlanner
+
+        return GefenSavePlanner
+    if name == "GefenFileSystemWriter":
+        from .dcp import GefenFileSystemWriter
+
+        return GefenFileSystemWriter
     if name in ("split_params_for_muon", "validate_split"):
         from . import params