Add DTensor optimizer contracts and sharded portable resharding

CHANGELOG.md

@@ -2,6 +2,15 @@
 
 All notable changes to this project are documented here. This project adheres to [Semantic Versioning](https://semver.org/).
 
+## Unreleased
+
+- Add stable post-sharding rebinding, local training, and exact same-topology native checkpoint continuation for narrow one-dimensional default-world DTensors, including uneven and empty shards.
+- Freeze GefenMuon's DTensor parameter membership, order, rebinding, mesh route, and `sharded_mode` after the first collective routing validation; later edits fail before optimizer or parameter mutation because initialized state has no mode-migration rule.
+- Extend `GefenMuonHybrid` atomic rebinding and native load staging to exact AdamW backups and mixed DTensor layouts while preserving nested child state and fail-before-mutation semantics.
+- Bind finalized scoped native checkpoints to their exact process-group identity and local layout, including direct distributed Muon; unscoped distributed Muon retains its complete cross-world-size native checkpoint path.
+- Add representation-qualified checkpoint transitions, including the directional `defined_projection_factored_to_block_live_fp32_target_period_one_v1` portable projection; reverse block-to-factored conversion remains unsupported.
+- Replace the dense-everywhere portable DCP save format with a load-compatible sharded version that stores multi-member canonical fields as DCP DTensors and singleton fields as ordinary full tensors, supports plain-Gefen DTensor resharding and Gefen-backed Hybrid child namespaces, and bounds metadata, chunks, and aggregate payloads.
+
 ## [0.4.0] - 2026-07-12
 
 Correctness and compatibility:

COMPATIBILITY.md

@@ -95,7 +95,7 @@ Native single-process optimizer `state_dict()`/`load_state_dict()` and the expli
 
 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`.
 
-Finalized exact period-one plain Gefen, `GefenMuon(sharded_mode="distributed")`, and Gefen-backed `GefenMuonHybrid` additionally expose portable global-state DCP through `save_portable_dcp(...)` and `load_portable_dcp(...)`. This separate synchronous path gathers and verifies complete logical optimizer state, stores a tensor-only bounded canonical-wire envelope through PyTorch DCP, and projects it onto supported replicated, flattened, or whole-owner targets after load. It can change flattened placement and redistribute Muon owners across checkpoint world sizes; factored second moments remain replicated and same-topology. Gefen-backed Hybrid uses a separately versioned, digested composite wrapper around unchanged child v3 documents, exact disjoint FQN routing, and one all-child freshness/commit boundary. Every checkpoint member temporarily holds the complete dense global document and encoded CPU payload. The target import is fail-before-local-mutation after DCP has completed the read, but DCP storage publication itself is not transactionally atomic. AdamW-backed Hybrid, DTensor targets, asynchronous saving, mixed model/optimizer `Stateful` composition, singleton checkpoint scopes inside a larger initialized default world, and multi-member checkpoint groups whose coordinate zero is not global rank zero are not supported by this path.
+Finalized exact period-one plain Gefen, `GefenMuon(sharded_mode="distributed")`, and Gefen-backed `GefenMuonHybrid` additionally expose portable global-state DCP through `save_portable_dcp(...)` and `load_portable_dcp(...)`. This separate synchronous path stores bounded semantic metadata plus DCP-native sharded fp32 fields, so checkpoint ranks retain canonical local chunks instead of all materializing one dense optimizer document. Plain Gefen can reshard exact block state among replicated, flattened, and narrow one-dimensional default-world DTensor layouts; replicated factored state restores exactly to replicated factored state or follows the versioned `defined_projection_factored_to_block_live_fp32_target_period_one_v1` rule to a replicated, flattened, or narrow-DTensor block target using compact factor slices and the saved exact fp32 row denominator. The standard DCP reader can transiently deserialize one saved source storage chunk before narrowing it, so each source chunk must fit the load-time fragment ceiling and a lower-limit reader may reject a checkpoint written with larger chunks. Distributed Muon remains limited to replicated and whole-owner portable layouts and can redistribute owners across checkpoint world sizes. Gefen-backed Hybrid stores each present child in a derived sharded namespace and uses exact disjoint FQN routing plus one all-child freshness/commit boundary. The target import is fail-before-local-mutation after DCP completes the read, but DCP storage publication itself is not transactionally atomic. The earlier dense DCP envelope remains load-compatible and retains its original complete-document memory cost. AdamW-backed Hybrid portable state, Muon DTensor resharding, block-to-factored conversion, asynchronous saving, mixed model/optimizer `Stateful` composition, singleton checkpoint scopes inside a larger initialized default world, and multi-member checkpoint groups whose coordinate zero is not global rank zero are not supported by this path.
 
 ## Transformers Trainer DDP
 

README.md

@@ -604,6 +604,8 @@ Under FSDP2, Muon's orthogonalization needs each full weight matrix, but every G
 - **`"distributed"`** (experimental) — each matrix is assigned by its stable position in the full distributed parameter set to one GPU, which does the work and shares the result. Bit-identical to `"exact"` on homogeneous GPUs and faster as you add them. Momentum ownership is stable when the active gradient set varies, and `state_dict()` collectively gathers owner-local momentum so rank 0 can write a complete checkpoint after all ranks call it.
 - **`"approx"`** — each GPU works on just its slice. Fastest, but results genuinely differ — an accuracy trade.
 
+Choose the DTensor route before its first optimizer collective. After GefenMuon validates a multi-rank DTensor route, its parameter membership and order, live tensor/mesh binding, and `sharded_mode` are frozen; even a unanimous later edit is rejected because initialized local, global, and owner momentum have no implicit migration rule. Construct a new optimizer and use a supported checkpoint path when changing those semantics.
+
 ```python
 from gefen import GefenMuonHybrid
 
@@ -615,7 +617,7 @@ opt = GefenMuonHybrid(
 # only takes effect under FSDP2 (DTensor params); no-op single-GPU
 ```
 
-> **`"distributed"` checkpointing is collective.** `state_dict()` gathers each owner's momentum across ranks, so **every rank must call it** (as in a standard FSDP full-state-dict flow). Calling `state_dict()` on rank 0 only — e.g. a rank-0-only save loop — **deadlocks**. Save and load also transiently materialize the full unsharded momentum on every rank, so peak memory at checkpoint time approaches `"exact"` mode's. The saved checkpoint is complete on every rank and resumes under any world size, including a single-process optimizer; incomplete or inconsistent owner state fails closed before any state is touched.
+> **`"distributed"` checkpointing is collective.** `state_dict()` gathers each owner's momentum across ranks, so **every rank must call it** (as in a standard FSDP full-state-dict flow). Calling `state_dict()` on rank 0 only — e.g. a rank-0-only save loop — **deadlocks**. Save and load also transiently materialize the full unsharded momentum on every rank, so peak memory at checkpoint time approaches `"exact"` mode's. An unscoped `GefenMuon` checkpoint is complete on every rank and resumes under any world size, including a single-process optimizer; incomplete or inconsistent owner state fails closed before mutation. Once direct Muon or `GefenMuonHybrid` is finalized with an explicit codebook scope, its ordinary native checkpoint carries an exact scope, topology, and local-layout guard and is same-topology only; use the separate portable DCP path for supported Gefen-backed topology changes.
 
 ![Gefen-Muon exact / distributed / approx sharded — eval loss](https://raw.githubusercontent.com/thad0ctor/Gefen-X/main/docs/benchmarks/muon_shard_loss.png)
 ![Gefen-Muon exact / distributed / approx sharded — throughput & VRAM](https://raw.githubusercontent.com/thad0ctor/Gefen-X/main/docs/benchmarks/muon_shard_perf.png)
@@ -632,13 +634,13 @@ Measured (Qwen3-0.6B, 2 and 4 GPUs): `"distributed"` matched `"exact"` in the re
 
 `Gefen`, `GefenMuon`, and `GefenMuonHybrid` preserve the parameter groups you pass to the optimizer, so list-indexed layer-wise LR recipes, per-group LR logging, and `state_dict()["param_groups"]` see the same group boundaries as conventional `torch.optim` optimizers. Per-parameter names are stored in optimizer state and mirrored as each group's `param_names` list for integrations that need name-level routing. Checkpoints from the older one-group-per-parameter layout are migrated on load when the total parameter order still matches and the old per-param hyperparameters can be represented by the new group layout.
 
-Finalized exact period-one plain Gefen, distributed-owner GefenMuon, and Gefen-backed `GefenMuonHybrid` also provide a separate portable global optimizer-state path through `save_portable_dcp(...)` and `load_portable_dcp(...)`. It can reshard supported block-state parameters, redistribute Muon owners, and restore both Gefen-backed Hybrid children as one validated composite transaction across checkpoint topologies; it is synchronous, temporarily materializes the complete dense optimizer document on every checkpoint rank, and is distinct from ordinary FSDP2 optimizer checkpoints. See the [optimizer integration contracts](https://github.com/thad0ctor/Gefen-X/blob/main/docs/optimizer_contracts.md#portable-global-state-v3) for the supported layouts, setup, and exclusions.
+Finalized exact period-one plain Gefen, distributed-owner GefenMuon, and Gefen-backed `GefenMuonHybrid` also provide a separate portable global optimizer-state path through `save_portable_dcp(...)` and `load_portable_dcp(...)`. It uses DCP-native sharded fields rather than materializing the complete dense optimizer document on every checkpoint rank, can reshard plain-Gefen replicated/flattened/narrow-DTensor state and Muon replicated/whole-owner state, defines a versioned bounded factored-to-block projection, and restores every present Gefen-backed Hybrid child as one validated composite transaction. Retained payload is target-canonical-local, although the standard DCP reader may transiently deserialize one saved source storage chunk and rejects it when that chunk exceeds the current fragment ceiling. It is synchronous and distinct from ordinary FSDP2 optimizer checkpoints; the earlier dense DCP envelope remains load-compatible. See the [optimizer integration contracts](https://github.com/thad0ctor/Gefen-X/blob/main/docs/optimizer_contracts.md#portable-global-state-v3) for the exact optimizer/layout matrix, setup, and exclusions.
 
 Plain `Gefen` with ordinary replicated CUDA parameters can keep its persistent per-parameter optimizer state on CPU between eager steps with `optimizer.offload_state_("cpu")`. Each step synchronously stages only the parameter currently being updated to its CUDA device, copies the updated state back to CPU, and releases the temporary device state; the small shared codebook remains CUDA-resident. `optimizer.restore_state_()` atomically returns all state to the parameter devices, while `move_state_()` also disables an active offload policy. This path intentionally excludes `GefenMuon`, `GefenMuonHybrid`, sharded or DTensor parameters, multi-member explicit codebook scopes, capturable optimizers, `torch.compile`, and CUDA graph capture.
 
 ## Known limitations
 
-- **Hybrid checkpoint schema.** `GefenMuonHybrid`'s ordinary `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. The separate topology-neutral DCP path above supports only a finalized Gefen-backed Hybrid; AdamW-backed Hybrid remains same-topology through its ordinary nested checkpoint.
+- **Hybrid checkpoint schema.** `GefenMuonHybrid`'s ordinary `state_dict()` uses its own nested `{"muon": ..., "backup": ..., "backup_optimizer": "gefen" | "adamw"}` layout and adds a versioned `finalized_binding` record after `post_sharding`. Resume from a checkpoint the Hybrid itself saved—not one consolidated or converted to the flat torch `{state, param_groups}` layout. Finalized loads require the exact child set, routing, shard identity, and local topology before either child is mutated; finalized and legacy-unfinalized schemas do not cross-load. Cross-backend loads reject, while legacy untagged Hybrid checkpoints are interpreted as Gefen-backed. The separate topology-neutral DCP path above supports only a finalized Gefen-backed Hybrid; AdamW-backed Hybrid remains same-topology through its ordinary nested checkpoint.
 - **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).
 - **CPU state offload is synchronous.** Plain-Gefen state offload reduces persistent CUDA optimizer-state residency by paging one parameter at a time, but it adds blocking CPU↔CUDA transfers to every updated parameter and is not an asynchronous overlap engine.
 - **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.