Deliverable 4: the slimmed platform-optimizer convention (contracts, portable state, DCP)

.github/workflows/ci.yml

@@ -124,11 +124,12 @@ jobs:
         run: |
           python - <<'PY'
           import gefen
-          # Derive the public optimizers from __all__ so this adapts per branch
-          # (e.g. main has no GefenMuonHybrid). Skip non-class exports: "kernels"
-          # is a submodule and "__version__" is a string, neither has __name__.
-          names = [n for n in gefen.__all__ if n not in ("kernels", "__version__")]
-          loaded = {n: getattr(gefen, n).__name__ for n in names}
+          # Resolve every public export in __all__ so this adapts per branch and
+          # exercises the lazy __getattr__ hooks. Exports are a mix of classes,
+          # a submodule ("kernels"), the version string, and integer schema
+          # constants (e.g. CONTRACT_SCHEMA_VERSION), so report __name__ where it
+          # exists and fall back to repr for value exports.
+          loaded = {n: getattr(getattr(gefen, n), "__name__", repr(getattr(gefen, n))) for n in gefen.__all__}
           assert "Gefen" in loaded, "Gefen must always be importable"
           print("import OK ->", loaded)
           PY

COMPATIBILITY.md

@@ -157,6 +157,8 @@ Async save returns once the state dict has been *staged* — copied to CPU — a
 
 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.
 
+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.
+
 ## Transformers Trainer DDP
 
 The `benchmarks.trainer_resume` gate exercises plain Gefen, GefenMuon+AdamW, and GefenMuon+Gefen through Trainer's internal Accelerate wrapper with tied weights, gradient accumulation, a changing scheduler, native Trainer checkpoint files, BF16 fused updates, and two-rank DDP replica hashes. All three recipes have passed its deterministic fused-BF16 two-rank configuration on homogeneous GPUs, which requires exact model, optimizer, scheduler, LR, and logged-loss agreement between uninterrupted and resumed runs. Run the gate with:

README.md

@@ -184,6 +184,8 @@ Mixed precision works out of the box: BF16 and standard AMP behave exactly as wi
 
 > **DeepSpeed ZeRO config.** Set `"zero_allow_untested_optimizer": true` and leave the config's `optimizer` section unset. With optimizer or parameter CPU-offload, also set `"zero_force_ds_cpu_optimizer": false` — otherwise raw DeepSpeed refuses to initialize, and accelerate-based launchers (axolotl) silently swap in DeepSpeed's own CPU Adam. With those two flags, plain Gefen (`optimizer: gefenx`) steps the CPU-resident fp32 partitions directly and trains normally under ZeRO-2 and ZeRO-3 offload, for both full fine-tuning and LoRA. Activation offloading (`activation_offloading: true`) is optimizer-agnostic and composes with `gefenx`, including alongside ZeRO CPU-offload.
 
+Platform adapters can query the immutable [`optimizer_contract()` capability and state-layout descriptors](https://github.com/thad0ctor/Gefen-X/blob/main/docs/optimizer_contracts.md) instead of depending on Gefen's private optimizer attributes.
+
 ## Determinism (`deterministic`)
 
 Set `deterministic=True` when every GPU replica must produce bit-identical results on matching GPUs. It is off by default (the fastest routing), checkpoints remember the setting, and older checkpoints without it still load. One combination is rejected: plain Gefen with `deterministic=True`, `factored_v_2d=True`, and `stochastic_round=True` together.
@@ -683,9 +685,11 @@ 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.
+
 ## 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.
+- **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.
 - **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.
 

benchmarks/microbench/bench_layout_guard.py

@@ -0,0 +1,198 @@
+"""Benchmark the per-step layout-forensics guard cost on a synthetic manifest.
+
+Builds a finalized plain Gefen over a large synthetic flattened-shard
+``ShardingManifest`` (defaults: 512 process-group members x 300 parameters =
+153,600 global ``ShardIdentity`` records) with no real process group, then
+compares:
+
+  * old per-step guard cost — the pre-fix step() sequence re-ran the complete
+    O(params x world) forensic rebuild on every guard call (2 passes per
+    unscoped step, up to 7 under an explicit multi-member codebook scope) and
+    recomputed the manifest sha256 fingerprint inside every scoped operation
+    header (2 per step);
+  * new per-step guard cost — the exact warm step() guard sequence, which
+    reuses one cached forensic verdict through O(local params) identity
+    tokens and the manifest digest computed once at post_sharding
+    finalization.
+
+The headline is the scoped worst case (7 forensic passes + 2 digest
+computes); the unscoped floor (2 passes, no digests) is printed alongside.
+Exits nonzero if the scoped per-step reduction is below the required 100x.
+
+Run from the repo root (CPU only, no distributed init needed):
+
+    PYTHONPATH=src python benchmarks/microbench/bench_layout_guard.py
+"""
+
+from __future__ import annotations
+
+import argparse
+import sys
+import time
+
+import torch
+
+from gefen import (
+    Gefen,
+    LogicalSlice,
+    ParameterIdentity,
+    ParameterLayout,
+    ParameterRebinding,
+    PlacementKind,
+    ProcessGroupIdentity,
+    ShardIdentity,
+    ShardPlacement,
+    ShardingManifest,
+)
+
+
+def build_finalized_optimizer(members_count: int, params_count: int, local_length: int):
+    members = tuple("member{:04d}".format(index) for index in range(members_count))
+    group = ProcessGroupIdentity("data_parallel", members)
+    local_member = members[0]
+    manifest_shards = []
+    local_shards = []
+    for index in range(params_count):
+        identity = ParameterIdentity(
+            "Model.Block{}.Weight".format(index), (members_count * local_length,)
+        )
+        offset = 0
+        for coordinate, member in enumerate(members):
+            shard = ShardIdentity(
+                identity,
+                ParameterLayout.FLATTENED_ELEMENT_SHARD,
+                LogicalSlice(offset, local_length),
+                placements=(
+                    ShardPlacement(
+                        "data_parallel",
+                        PlacementKind.FLAT_SHARD,
+                        coordinate,
+                        members_count,
+                    ),
+                ),
+                process_group=group,
+                local_member=member,
+            )
+            manifest_shards.append(shard)
+            if member == local_member:
+                local_shards.append(shard)
+            offset += local_length
+
+    start = time.perf_counter()
+    manifest = ShardingManifest(tuple(manifest_shards))
+    manifest_seconds = time.perf_counter() - start
+
+    parameters = [
+        torch.nn.Parameter(torch.randn(local_length)) for _ in range(params_count)
+    ]
+    optimizer = Gefen(
+        [
+            ("model.block{}.weight".format(index), parameters[index])
+            for index in range(params_count)
+        ],
+        fused=False,
+        factored_v_2d=False,
+    )
+    start = time.perf_counter()
+    optimizer.post_sharding(
+        tuple(
+            ParameterRebinding(parameters[index], parameters[index], local_shards[index])
+            for index in range(params_count)
+        ),
+        manifest=manifest,
+    )
+    finalize_seconds = time.perf_counter() - start
+    return optimizer, parameters, manifest, manifest_seconds, finalize_seconds
+
+
+def timed(callable_, repeats: int) -> float:
+    start = time.perf_counter()
+    for _ in range(repeats):
+        callable_()
+    return (time.perf_counter() - start) / repeats
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--members", type=int, default=512)
+    parser.add_argument("--params", type=int, default=300)
+    parser.add_argument("--local-length", type=int, default=4)
+    parser.add_argument("--full-repeats", type=int, default=3)
+    parser.add_argument("--fast-repeats", type=int, default=200)
+    args = parser.parse_args()
+
+    optimizer, parameters, manifest, manifest_seconds, finalize_seconds = (
+        build_finalized_optimizer(args.members, args.params, args.local_length)
+    )
+    print(
+        "manifest: {} shards ({} members x {} params), built in {:.2f}s; "
+        "post_sharding finalized in {:.2f}s".format(
+            len(manifest.shards),
+            args.members,
+            args.params,
+            manifest_seconds,
+            finalize_seconds,
+        )
+    )
+
+    # Warm the runtime and the forensic verdict exactly the way training does.
+    for parameter in parameters:
+        parameter.grad = torch.full_like(parameter, 0.5)
+    optimizer.step()
+
+    full_pass = timed(
+        lambda: optimizer._finalized_binding_layout_matches(full=True),
+        args.full_repeats,
+    )
+    digest_compute = timed(
+        lambda: optimizer._compute_codebook_manifest_fingerprint(manifest),
+        args.full_repeats,
+    )
+    cached_digest = timed(
+        optimizer._codebook_manifest_fingerprint, args.fast_repeats
+    )
+
+    def warm_step_guards():
+        # The finalized-layout/process-group step guard sequence (both the
+        # pre-closure and post-closure blocks), on a warm verdict.
+        optimizer._assert_finalized_binding_layout()
+        optimizer._assert_runtime_codebook_process_group()
+        optimizer._assert_finalized_binding_layout()
+        optimizer._assert_runtime_codebook_process_group()
+
+    warm_guards = timed(warm_step_guards, args.fast_repeats)
+
+    # Pre-fix per-step guard cost. Unscoped step(): 2 complete forensic
+    # passes. Scoped step(): up to 7 complete passes (step entry/re-entry,
+    # scope asserts, operation headers, failure synchronization, scope
+    # agreement) plus 2 manifest fingerprint recomputes in the exchanged
+    # "step" and "periodic_step" headers.
+    old_unscoped = 2 * full_pass
+    old_scoped = 7 * full_pass + 2 * digest_compute
+    new_unscoped = warm_guards
+    new_scoped = warm_guards + 2 * cached_digest
+
+    print("one full forensic pass:        {:>12.6f}s".format(full_pass))
+    print("one manifest digest compute:   {:>12.6f}s".format(digest_compute))
+    print("one cached digest fetch:       {:>12.6f}s".format(cached_digest))
+    print("warm step guard sequence:      {:>12.6f}s".format(warm_guards))
+    print(
+        "old per-step guards (unscoped): {:>11.6f}s -> new: {:.6f}s ({:.0f}x)".format(
+            old_unscoped, new_unscoped, old_unscoped / new_unscoped
+        )
+    )
+    scoped_ratio = old_scoped / new_scoped
+    print(
+        "old per-step guards (scoped):   {:>11.6f}s -> new: {:.6f}s ({:.0f}x)".format(
+            old_scoped, new_scoped, scoped_ratio
+        )
+    )
+    if scoped_ratio < 100.0:
+        print("FAIL: scoped per-step guard reduction is below 100x")
+        return 1
+    print("PASS: scoped per-step guard reduction is >= 100x")
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())