Compat: DCP resharding for FSDP2 Gefen state

CHANGELOG.md

@@ -4,7 +4,9 @@ All notable changes to this project are documented here. This project adheres to
 
 ## [0.4.1] - 2026-07-15
 
-Distributed checkpoint-load and step-failure hardening. All fixes are backward-compatible; no public API changes.
+Distributed checkpoint-load and step-failure hardening. All changes are backward-compatible; the only public API addition is the opt-in `GefenDCPState` resharding wrapper.
+
+- Add `GefenDCPState`, a standalone DCP Stateful wrapper for **resharding** plain Gefen on one-dimensional default-world FSDP2 `Shard(0)` parameters (`factored_v_2d=False`, `capturable=False`). Save writes shard-addressable dense optimizer tensors so DCP can reshard N ranks to M without a rank-0 full-state gather; load validates before mutation, then re-blocks the resharded dense momentum back into Gefen's compact ~1 byte/param block state (re-running the period search, relearning the codebook, and re-quantizing on each new shard). Resume is a correct continuation within 256-level quantization noise, not bit-exact — use the unchanged native full-state path for same-topology resumes.
 
 Correctness — checkpoint loads are now atomic:
 

COMPATIBILITY.md

@@ -103,7 +103,9 @@ Plain Gefen supports FSDP2 training-time CPU offload via `CPUOffloadPolicy` (`fu
 
 Native single-process optimizer `state_dict()`/`load_state_dict()` and the explicitly documented GefenMuon `distributed` state path remain separate from the rank-local DCP adapter. For plain Gefen and `GefenMuon(sharded_mode="approx")`, each rank can learn different local codebook and block geometry, so `state_dict()` collectively replaces ordinary rank-local tensors with one tagged, rank-indexed CPU payload that PyTorch full-state DCP preserves. Both optimizers' two-GPU FSDP2 `fully_shard` paths are exercised through `get_optimizer_state_dict(..., full_state_dict=True, cpu_offload=True)` and `set_optimizer_state_dict(..., full_state_dict=True, broadcast_from_rank0=True)`, with exact next-step continuation; the flattened optimizer-state form is also covered.
 
-This format is deliberately same-topology only and currently requires one 1-D DeviceMesh spanning the default process-group world. Every rank must participate in both save and restore; each process temporarily holds all serialized rank payloads on CPU, so the leading checkpoint-time CPU cost approaches `world_size ×` its local optimizer-state size plus local serialization scratch. Loading validates world size, parameter order and names, global and local shapes and dtypes, mesh membership and names, structural placements, rank coordinates, global step, deterministic policy, frozen codebook, and sharded mode before mutation. Multidimensional meshes, subgroups, pipeline-local optimizers, world-size/topology changes, and old unsafe untagged full checkpoints fail closed rather than silently applying rank 0's state to every shard. No optimizer-state reshard portability is claimed, model-only DCP is unaffected, and the full-state DCP support described here does not extend beyond plain Gefen and Muon `approx`.
+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.
 
 ## 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 and Muon `approx`; resume needs the same GPU count — scope note below |
+| 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) |
 | 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.** Plain Gefen and Muon `approx` save and resume exactly through PyTorch's standard full-state checkpoint calls, as long as the GPU count and sharding layout are unchanged and every GPU joins the save. Anything outside that scope refuses to load instead of corrupting state — [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 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 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 checkpoints don't reshard.** Plain Gefen and Muon `approx` resume only on the same GPU count and layout; changing either refuses to load. Model weights are unaffected — [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)}, ...)` 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).
 - **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

@@ -11,6 +11,7 @@
     "Gefen",
     "GefenMuon",
     "GefenMuonHybrid",
+    "GefenDCPState",
     "split_params_for_muon",
     "validate_split",
     "kernels",
@@ -31,6 +32,10 @@ def __getattr__(name):
         from .hybrid import GefenMuonHybrid
 
         return GefenMuonHybrid
+    if name == "GefenDCPState":
+        from .dcp import GefenDCPState
+
+        return GefenDCPState
     if name in ("split_params_for_muon", "validate_split"):
         from . import params