diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 38e1060..155776c 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -118,11 +118,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 diff --git a/COMPATIBILITY.md b/COMPATIBILITY.md index a831707..ecb244f 100644 --- a/COMPATIBILITY.md +++ b/COMPATIBILITY.md @@ -95,6 +95,8 @@ 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. + ## 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: diff --git a/README.md b/README.md index 92ca5f8..dc1efd8 100644 --- a/README.md +++ b/README.md @@ -131,6 +131,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 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. +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. @@ -630,9 +632,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 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). - **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. diff --git a/benchmarks/microbench/bench_layout_guard.py b/benchmarks/microbench/bench_layout_guard.py new file mode 100644 index 0000000..33bd8b7 --- /dev/null +++ b/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()) diff --git a/docs/optimizer_contracts.md b/docs/optimizer_contracts.md new file mode 100644 index 0000000..62f9403 --- /dev/null +++ b/docs/optimizer_contracts.md @@ -0,0 +1,139 @@ +# Optimizer integration contracts + +Gefen exposes immutable, versioned optimizer contracts for platform adapters that need to inspect state ownership and validated distributed capabilities without depending on private attributes. Calling `optimizer.optimizer_contract()` is read-only and does not change parameters, optimizer state, checkpoint schemas, or step behavior. + +```python +from gefen import CheckpointTransport, Gefen, ParameterLayout + +optimizer = Gefen(model.named_parameters(), lr=3e-5) +contract = optimizer.optimizer_contract() + +assert contract.schema_version == 1 +rank_local_dcp = next( + support + for support in contract.capabilities.checkpoints + if support.transport is CheckpointTransport.PYTORCH_RANK_LOCAL +) +assert ParameterLayout.DTENSOR_1D_DEFAULT_WORLD in rank_local_dcp.same_topology +``` + +## Contract boundaries + +- `OptimizerStateLayout` separates optimizer-common authoritative state, per-parameter authoritative state, derived caches, checkpoint transport fields, and composite child namespaces. +- `StateVariant` identifies valid lazy, initialized, local-shard, global-parameter, owner, non-owner, and migrated state combinations using structured layout, mode, rank, extent, ownership, and inactive-field declarations. +- `TrainingSupport` qualifies each validated parameter layout by process-group source, mesh dimensionality, sharded mode, and whether the update needs complete parameter storage or a transient complete logical matrix. +- `CheckpointSupport` reports same-topology, topology-changing, and fail-before-mutation load support separately for native, PyTorch rank-local, canonical local, canonical global, and composite checkpoint transports. +- Precision, canonical parameter identity, stable shard identity, explicit process-group-scoped codebooks, shard rebinding, post-sharding, canonical state I/O, and the retained `atomic_state_movement` and `state_offload` fields are independent capability declarations. A false field is an explicit unsupported contract, not an invitation for an adapter to infer support from internal state. + +The current DTensor declaration is deliberately narrow: `DTENSOR_1D_DEFAULT_WORLD` means one shared one-dimensional mesh spanning the default world. Multidimensional meshes, subgroups, and placement-changing loads are not implied by that declaration. + +## Canonical parameter and shard identity + +`ParameterIdentity` records an exact, case-preserving model FQN and global logical shape independently of any live tensor object. `ProcessGroupIdentity` records an adapter-defined semantic group name and authoritative ordered member IDs without importing a framework process-group type. `ShardIdentity` combines those values with a contiguous row-major `LogicalSlice` for replicated, flattened, and whole-owner layouts or an axis-aligned `LogicalRegion` for the narrow one-dimensional DTensor layout, plus structured placements, the local member, and an optional whole-parameter owner. `ShardingManifest` validates and deterministically orders the complete identity set; flattened slices and dimension-sharded regions must cover each logical parameter exactly once without gaps or overlaps, replicated manifests carry one complete identity per declared member, and whole-parameter manifests identify one complete owner while retaining empty non-owner records. DTensor regions currently describe one default-world mesh axis with either replication or one parameter-dimension shard, including uneven and empty shards. This identity vocabulary does not by itself claim DTensor post-sharding rebinding or portable checkpoint support; those remain negative until the optimizer data plane consumes the regions. + +These descriptors do not treat legacy `param_names`, generated names, Python tensor identity, rank-local parameter IDs, devices, or dtypes as canonical identity. They also do not contain runtime collective handles. An adapter remains responsible for mapping a stable `ProcessGroupIdentity` to its framework process group and for canonicalizing tied aliases to one primary FQN and one optimizer slot; alias-rich identity is not part of schema version 1. Declaring identity metadata alone does not enable rebinding, canonical checkpoint I/O, topology-changing load, or codebook scoping; those capabilities remain separate. + +## Atomic post-sharding rebinding + +`Gefen.post_sharding(rebindings, manifest=...)` finalizes the complete local optimizer layout as one pre-initialization transaction. Each `ParameterRebinding` maps an existing optimizer slot to its local live tensor, or to `None` for an explicit whole-parameter non-owner. The global `ShardingManifest` remains descriptive: the core compares its FQN set with every local optimizer slot and validates each supplied local shard, but it does not infer live tensors, current members, owners, or runtime collective handles from manifest order. `rebind_parameter` and `rebind_shard` are one-slot conveniences and therefore apply only when they describe the optimizer's complete plan. + +Rebinding is allowed only while the entire optimizer is pristine: global step zero, no learned codebook, no gradients, no authoritative parameter state, no active capture stacks, and no nonzero device counters. The core stages every group, compatibility name, constructor-only state removal, canonical binding, cache invalidation, device counter, and checkpoint-schema update before publishing the result. A failed batch leaves the exact live optimizer objects unchanged. A successful batch preserves group order, group options, and released lowercase compatibility names while storing exact FQNs separately; it seals the layout against later incremental groups or rebindings. Targets must have no internal storage overlap and distinct targets may not overlap one another. Schema version 1 conservatively rejects multidimensional strided layouts whose element disjointness cannot be proven from dense stride spans, as well as distinct noncontiguous targets that share one storage even when their logical elements are disjoint. Tied aliases must already be collapsed to one optimizer slot. + +Plain Gefen currently rebinds complete replicated parameters and contiguous physical 1-D flattened element shards. A flattened logical matrix requires `factored_v_2d=False` because canonical row/column factored-state projection is not implemented. GefenMuon rebinds complete replicated matrices and can finalize/prune whole-parameter owner manifests. A whole-owner instance may step only when the same transaction installs the explicit codebook process group described below; the owner performs the update and the machine-readable training declaration requires the adapter to synchronize the updated matrix afterward. A Gefen-backed `GefenMuonHybrid` atomically partitions one complete manifest and rebinding plan by its frozen exact FQN routing, stages both children, validates cross-child storage disjointness, rebuilds composite state routing, and publishes only after every child succeeds. AdamW-backed Hybrid and DTensor composite rebinding remain unsupported. The portable global-state path described below can reshard supported finalized layouts. + +After finalization, every entry point re-validates the published layout, and this has two costs. Steps and identity queries take an O(local params) fast path: the first complete forensic rebuild caches a verdict keyed by cheap identity tokens — the finalized registries by object identity, every live group container, parameter, and compatibility name, and a version counter that every mutating API bumps — and the full rebuild re-runs only when one of those tokens changes. A fixed set of boundaries always runs the complete rebuild regardless of the cache: checkpoint save and load (`state_dict` / `load_state_dict`), canonical export, import prepare, and import commit; `post_sharding` rebinding; collective codebook initialize and refresh; codebook-scope re-validation; and external contract-readiness queries. `post_sharding` additionally computes the manifest shard set and its sha256 digest once per finalized manifest, for reuse by the scoped operation headers. The practical consequence for an integrator is a clean split: any layout corruption reachable through the public containers — the group `params` / `param_names` slots, per-parameter state names, or the compatibility-name cache — still fails the step guard before any state is mutated, including corruption a closure introduces between the pre- and post-closure guards. Only corruption that leaves every fast-path token intact — in-place value replacement inside the private finalized registries, or an `object.__setattr__` on a frozen identity record — waits until the next boundary above to be caught rather than being caught at the next step. + +## Explicit learned-codebook process groups + +`CodebookProcessGroupBinding` maps one stable `ProcessGroupIdentity` and local semantic member to an opaque PyTorch process-group handle plus an explicit collective device. It is accepted only through the complete `post_sharding(..., codebook_process_group=...)` transaction: every manifest shard must use that one semantic group, each local shard must name the binding's local member, the runtime group size and coordinate must match `ordered_members`, and the backend must support the supplied device. A one-member scope uses `process_group=None`; multi-member scopes must pass a real handle, including `dist.group.WORLD` when the default world is intentionally the semantic scope. Gefen never treats `None` as an implicit default-world selection. + +The `explicit_process_group_codebook_scope` capability reports that Gefen or GefenMuon implements this API; `codebook_process_group_binding()` separately reports whether a particular finalized instance has an active binding. A Gefen-backed `GefenMuonHybrid` implements one atomic composite transaction that installs the same exact binding object on every present child while retaining independent child codebooks. AdamW-backed Hybrid remains negative because AdamW has no corresponding stable rebinding and staged portable-state contract. + +The optimizer owns one learned codebook and therefore accepts one scope. Histogram accounting represents each logical parameter once: the first ordered member contributes a replicated parameter after all members agree on gradient presence and its automatic period, every flattened shard contributes its local logical slice after all nonempty slices agree on gradient presence, and only the declared whole-parameter owner contributes an owned matrix. Local inputs are visited in canonical manifest order. Members first stage periods and an integer histogram without touching live optimizer state, exchange operation, step, scope, manifest, active-slice, period, policy, and old-codebook controls, sum the fixed-size `int64` histogram in the supplied group, solve the same exact-DP problem, verify codebook agreement, and only then publish periods and the codebook. Refresh additionally stages every replacement momentum-index tensor in canonical order and exchanges readiness before replacing indices or invalidating derived codebook/LUT caches. A local preparation, solve, or requantization failure is reported by every participant before optimizer-state commit; this is not rollback after process death or a collective-backend failure during the final commit window. + +`initialize_codebook()` and `refresh_codebook()` expose these operations for adapters that enter their optimizers in a deterministic order; `binding.sort_key` supplies the stable process-group portion of that schedule. Normal `step()` still initializes automatically and plain Gefen still honors `codebook_refresh_every`. Every scoped step exchanges a common operation header before any rank-dependent branch, and the first step after binding or native load additionally verifies codebook bytes and the complete manifest. Scoped native AMP requires every member to select the same protocol and present identical `found_inf` and `grad_scale` values; a mismatch raises collectively and requires a group-aware gradient scaler rather than changing external scaler state behind its back. Multi-member explicit scopes reject `capturable=True` because their host validation and process-group collectives are not CUDA-graph-safe. A one-member local scope may initialize during eager warmup and then use ordinary capturable stepping, but manual codebook replacement remains rejected. Ordinary unscoped behavior remains collective-free. Explicit scope does not replace DTensor mesh collectives, AMP mesh preflights, Parallel-Muon ownership collectives, or checkpoint transport groups. + +Native checkpoints store a primitive rank-neutral scope record—format version, semantic name, ordered members, and refresh schedule—at the top level and in the existing transport mirror. Scoped group metadata uses a new outer format version so an older loader rejects rather than silently discarding the scope. Runtime process-group handles and collective devices are live adapter configuration and are never serialized or reconstructed. Flattened and whole-owner local payloads additionally carry a primitive rank-local shard-identity guard. Guard format v2 records every original logical slot in group and slot order, including its lowercase compatibility name and portable shard identity, so pruned whole-owner nonowners remain bound to their original positions and equal-shaped parameters cannot be reinterpreted positionally. New checkpoints emit v2, while the loader continues to accept v1 guards by comparing their legacy live-slot and separately sorted pruned-shard projection exactly. A scoped checkpoint requires the target to complete `post_sharding` with the same stable group identity, refresh policy, and native local layout before load. A missing or different record rejects inside the optimizer's atomic shadow load. Native flattened-shard checkpoint support is declared only on a finalized scoped instance with this guard; unscoped optimizers no longer claim it. Legacy unscoped checkpoints remain loadable by unscoped optimizers. Native whole-owner completeness, scoped DTensor rank-local transport, and scope migration remain separate from the portable global transport and are not claimed by the native path. Hybrid's ordinary nested native loader also remains distinct from its atomic composite portable path. + +## Exact-binding canonical local state + +`export_canonical_state()` returns the versioned `gefen.bound_state` local-fragment format and `import_canonical_state()` restores it by exact FQN rather than native parameter IDs or group/slot position. `prepare_canonical_state_import()` performs complete schema, implementation, policy, manifest, shard, group-option, common-state, and tensor-geometry validation on a shadow optimizer without mutating the live instance; an adapter may synchronize preparation status in its own process group and then call `commit_canonical_state_import()`. Prepared imports are optimizer-specific, single-use, and rejected if live optimizer state changed after preparation. The convenience import is the local prepare-plus-commit transaction. Canonical methods bypass ordinary native state-dict hooks so arbitrary hook side effects cannot weaken the core fail-before-mutation boundary. + +The v1 document contains only primitive containers and detached tight finite CPU leaves of the ordinary dense `torch.Tensor` type, round-trips through `torch.load(weights_only=True)`, records the complete stable manifest, and maps each locally present parameter's authoritative state and algorithm-shaping group options by exact case-preserving FQN. After finalization the exact FQN is also the name used by period-routing policy, so compatibility names, devices, runtime process-group handles, derived lookup tables, capturable buffers, checkpoint carriers, and native positional IDs are not optimizer meaning in this format; the target retains its own lowercase compatibility names after import. Codebook scope identity is recorded, but its runtime handle and collective device must already be installed through `post_sharding`. + +This is an exact-binding transport-neutral local fragment, distinct from the dense global logical state in portable v3 below. Its dynamic `CANONICAL_LOCAL` checkpoint entry covers finalized plain-Gefen replicated and flattened local shards and finalized replicated GefenMuon, performs no collectives, reports atomic local import, and has an empty topology-changing set. A different member, slice, manifest, algorithm policy, group option, or declared state variant rejects. Export, preparation, and commit are quiescent checkpoint-boundary operations; prepared imports use content-bearing freshness tokens, including device counters, to reject intervening mutation. Export remains available after capturable warmup, but a capturable import target must still be fresh, before authoritative device state or a CUDA graph exists; importing replaces state identities, so an already captured graph cannot safely remain attached. Configurations with `stochastic_round=True` do not claim canonical v1 because the decomposed path intentionally lacks the fused stochastic quantizer, so changing effective fused availability would change the algorithm. DTensor, whole-owner completeness, Hybrid composition, rank-fragment gathering, resharding, world-size change, dense momentum decoding, and target-topology recompression remain unclaimed by `CANONICAL_LOCAL`; supported collective assembly and resharding use `CANONICAL_GLOBAL`. + +## Portable global-state v3 + +The transport-neutral `gefen.portable_state` version-3 format is a complete `global_logical_optimizer` artifact for the supported exact period-one configurations. Its schema carries the implementation and algorithm policy, optimizer-common state, an exact FQN-keyed parameter catalog with global identities, algorithm options, state variants, dense authoritative state, projection hints, source provenance, and a completion marker whose deterministic SHA-256 covers every preceding value including tensor dtype, shape, and canonical little-endian bytes. Optimizer export validates local native state, decodes quantized momentum and block second moments to dense logical fp32 fields, collectively assembles complete parameters by stable shard identity, verifies replicated and owner consensus bit-for-bit, and returns the same normalized complete document to every participant. Import validates the complete document and digest before projecting dense fields to the target shard, recompressing momentum at period one, staging a native shadow load, exchanging unanimous readiness and freshness, and publishing locally only after the final vote. Runtime global ranks are absent from the durable identity. + +`export_portable_state(checkpoint_process_group=..., transaction_id=..., limits=...)` and `import_portable_state(state, checkpoint_process_group=..., transaction_id=..., limits=...)` implement this collective path on exact `Gefen`, `GefenMuon`, and supported Gefen-backed `GefenMuonHybrid` instances. The exact `CheckpointProcessGroupBinding` must match the optimizer's installed `CodebookProcessGroupBinding` in stable identity, local member, live process-group handle, and collective device; all collectives use that optimizer-owned transport. Every member must enter operations in the same deterministic order with the same trimmed transaction ID, limits, target logical slot schema, and, for import, complete document. `PortableStateLimits` bounds per-member and aggregate tensor bytes, metadata, tree shape, strings, tensor count and rank, member count, wire chunk size, and diagnostics; structural limits are checked before dense materialization. `chunk_bytes` controls wire cloning and collective transfer, runtime value validation and freshness hashing use independently fixed bounded chunks, and dense decode/projection allocations remain bounded by the declared fragment and collective tensor-byte ceilings. + +```python +from gefen import PortableStateLimits + +limits = PortableStateLimits( + max_fragment_tensor_bytes=2 << 30, + max_collective_tensor_bytes=16 << 30, + max_collective_metadata_bytes=256 << 20, +) + +document = optimizer.export_portable_state( + checkpoint_process_group=checkpoint_binding, + transaction_id="optimizer-save-0001", + limits=limits, +) + +target_optimizer.import_portable_state( + document, + checkpoint_process_group=target_checkpoint_binding, + transaction_id="optimizer-load-0001", + limits=limits, +) +``` + +`save_portable_dcp(...)` and `load_portable_dcp(...)` provide synchronous PyTorch Distributed Checkpoint storage for the same document. The adapter persists only tensors: one bounded canonical-wire metadata tensor plus numbered dense payload tensors. Its load planner validates the exact namespace, key set, ordinary full-tensor metadata, dtype, rank, chunk coverage, tensor count, and aggregate bytes against `PortableStateLimits` before allocating CPU destinations; canonical-wire and portable-document digests are verified before collective import. The exact storage class and its bounded string or path-like `checkpoint_id` are included in the collective preflight, so every member must address the same checkpoint; a custom storage plugin without that identity fails closed. A plain DCP `Stateful` wrapper is intentionally not used because DCP asks the target for preallocated tensors before loading while portable state variants determine their own key set and shapes. Tensor-only optimizer data avoids DCP's opaque-object payload path, but the framework's own checkpoint metadata remains a trusted-storage boundary and is read before the custom planner runs. Every member temporarily holds the complete dense global document and its encoded CPU payload, so this synchronous path prioritizes portability and validation rather than rank-sharded checkpoint memory. The portable document and canonical-wire envelope are versioned by Gefen; the surrounding on-disk DCP format retains PyTorch's own cross-release compatibility policy. + +A finalized Gefen-backed `GefenMuonHybrid` uses a separate `gefen.portable_composite_state` version-1 wrapper rather than pretending its children share one portable-v3 policy or codebook. The wrapper embeds unchanged Muon and Gefen v3 documents, records the exact disjoint FQN-to-role routing and backup policy, requires both child global steps and deterministic settings to agree, and covers the complete wrapper with its own digest. Export agrees on child presence and routing before entering child collectives in fixed role order. Import stages every child, exchanges one composite target and freshness vote, and then performs only the prevalidated nonthrowing child publications; failure before that vote leaves both children unchanged. The DCP tensor envelope is tree-agnostic and stores this wrapper without an opaque payload. Muon-only and Gefen-backup-only Hybrids are supported, while AdamW-backed Hybrid remains negative because AdamW lacks stable FQN rebinding, a topology-neutral moment codec, and a staged commit primitive. + +```python +import torch.distributed.checkpoint as dcp + +from gefen import load_portable_dcp, save_portable_dcp + +save_portable_dcp( + optimizer, + checkpoint_process_group=checkpoint_binding, + storage_writer=dcp.FileSystemWriter("checkpoint/optimizer"), + transaction_id="optimizer-save-0001", + limits=limits, +) + +load_portable_dcp( + target_optimizer, + checkpoint_process_group=target_checkpoint_binding, + storage_reader=dcp.FileSystemReader("checkpoint/optimizer"), + transaction_id="optimizer-load-0001", + limits=limits, +) +``` + +The dynamic `CANONICAL_GLOBAL` checkpoint declaration appears only while the live finalized optimizer passes the exact runtime readiness checks: explicit process-group scope, stable logical slots and manifest, ordinary built-in containers, supported CPU/CUDA tensor storage, no active compilation or CUDA capture, `capturable=False`, `stochastic_round=False`, a complete declared native state variant, and period one for selected or initialized state. Plain Gefen supports replicated and contiguous flattened element shards. Block-second-moment state can reshard between replicated and flattened targets; a logical matrix using factored second moments remains replicated and same-topology because factored-to-block representation migration is not implemented. GefenMuon supports replicated matrices and whole-parameter ownership when every participating group uses `sharded_mode="distributed"`; the transport can change placement and redistribute owners across world sizes, including NorMuon row state. A ready Hybrid declaration is the union of its heterogeneous child layouts and change kinds; an adapter must inspect each role's child contract together with the finalized immutable `optimizer.parameter_routing()` result rather than apply that union indiscriminately to every parameter. Pristine and period-selected states are supported under the same policy rules, and zero-element parameters remain pristine. + +The collective protocol exchanges fixed-size preparation headers before payload movement, visits member fragments in stable semantic order, bounds metadata and tensor chunks, propagates asymmetric local failures to every participant, and performs no semantic checks after the final freshness vote. Import preserves the target's parameter groups, defaults, parameters, compatibility names, and runtime process-group configuration while restoring portable common state, including the source deterministic setting. The atomic claim is fail-before-local-mutation for live, quiescent optimizer instances; it is not rollback after process death, backend failure, or concurrent mutation after the final vote. Ordinary state-dict hooks are bypassed. Adapters must quiesce training, avoid retaining state-container identities across a successful import, and persist the returned weights-only-safe CPU document with their checkpoint system. + +Portable v3 currently excludes non-period-one initialized state, second-moment representation conversion, DTensor layouts, stochastic rounding, capturable/device-authoritative state, retained or migrated legacy parameter-state variants, tensor-valued or callable algorithm options, GefenMuon modes other than `distributed` for whole-owner transport, AdamW-backed `GefenMuonHybrid`, tied-alias expansion, asynchronous DCP, and mixed model/optimizer `Stateful` composition. The dedicated DCP helpers require every checkpoint member to enter synchronously and use the exact optimizer-owned process group. A singleton binding is rejected inside an initialized default world larger than one because the common PyTorch 2.5–2.12 `dcp.save/load` API interprets `process_group=None` as that world. A multi-member checkpoint group must have global rank zero at group coordinate zero; PyTorch 2.5's DCP coordinator path can otherwise address group coordinate zero as global rank zero and hang, so the adapter applies this compatibility restriction on every supported version. DCP storage publication is not transactionally atomic; the optimizer load remains fail-before-local-mutation after a complete successful read and verification. A CPU portable document remains a detached checkpoint artifact and does not change the live optimizer's declared capabilities. + +## Optimizer-state placement capability fields + +`atomic_state_movement` and `state_offload` remain in schema version 1 so adapters receive explicit negative declarations instead of inferring support from device placement, checkpoint contents, or private implementation details. They are false for live optimizers and do not expose movement or offload operations through this contract. + +Plain Gefen declares replicated, flattened element-shard, and the narrow DTensor training layouts. Its PyTorch rank-local checkpoint transport is same-topology only. GefenMuon declares replicated and narrow DTensor training, with mode-specific state extents: `approx` state is local, `exact` state is logically global, and `distributed` momentum is held by the parameter owner while non-owners retain metadata only. Native Parallel-Muon checkpoints separately declare world-size owner redistribution, not placement-changing resharding. `GefenMuonHybrid` retains its nested child namespaces and does not flatten AdamW or Gefen child state into a fabricated common schema. + +Gefen and GefenMuon native loads, rank-local payload restoration through `load_state_dict`, and distributed-owner payload restoration prepare their complete core restore before changing local live optimizer state. Their transport entries report `atomic_load=True`. This is a per-optimizer-instance fail-before-mutation guarantee at the optimizer load boundary, not a coordinated all-rank commit or a guarantee over work an external checkpoint orchestrator performs before calling the optimizer. Load pre-hooks run before that boundary and load post-hooks run afterward, so arbitrary side effects in user hooks are also outside the guarantee. Hybrid composite loads do not yet provide the same guarantee and report `atomic_load=False`. + +## Adapter requirements + +An adapter should match the exact `TrainingSupport` or `CheckpointSupport` entry it intends to use, including transport, layout, process-group scope, mesh dimensions, and sharded mode. It should not treat successful training as checkpoint support, same-topology checkpointing as resharding support, or accepted caller names as canonical fully qualified parameter identity. + +The contract types do not import a distributed platform and do not perform collectives. Platform adapters remain responsible for topology discovery, deterministic scheduling, and lifecycle orchestration; future core mutation APIs can consume these declarations without changing their meaning. diff --git a/src/gefen/__init__.py b/src/gefen/__init__.py index d2eb309..80075ac 100644 --- a/src/gefen/__init__.py +++ b/src/gefen/__init__.py @@ -11,6 +11,49 @@ "Gefen", "GefenMuon", "GefenMuonHybrid", + "CONTRACT_SCHEMA_VERSION", + "CANONICAL_STATE_FORMAT_VERSION", + "PORTABLE_STATE_FORMAT_VERSION", + "IDENTITY_SCHEMA_VERSION", + "CheckpointSupport", + "CheckpointTransport", + "CanonicalStateProvider", + "CheckpointProcessGroupBinding", + "CodebookProcessGroupBinding", + "PreparedCanonicalStateImport", + "PortableStateLimits", + "PortableStateProvider", + "OptimizerCapabilities", + "OptimizerChildContract", + "OptimizerContract", + "OptimizerContractProvider", + "OptimizerStateLayout", + "LogicalRegion", + "LogicalSlice", + "ParameterLayout", + "ParameterIdentity", + "ParameterStateRole", + "ParameterRebinding", + "PlacementKind", + "Precision", + "ProcessGroupScope", + "ProcessGroupIdentity", + "ShardIdentity", + "ShardPlacement", + "ShardingManifest", + "StateExtent", + "StateField", + "StateGeometry", + "StateKeyMatch", + "StateScope", + "StateVariant", + "TopologyChange", + "TrainingSupport", + "build_portable_state_document", + "normalize_portable_state_document", + "portable_state_digest", + "load_portable_dcp", + "save_portable_dcp", "split_params_for_muon", "validate_split", "kernels", @@ -35,6 +78,78 @@ def __getattr__(name): from . import params return getattr(params, name) + if name == "ParameterRebinding": + from .rebinding import ParameterRebinding + + return ParameterRebinding + if name == "CodebookProcessGroupBinding": + from .codebook import CodebookProcessGroupBinding + + return CodebookProcessGroupBinding + if name == "CheckpointProcessGroupBinding": + from .checkpoint import CheckpointProcessGroupBinding + + return CheckpointProcessGroupBinding + if name in ( + "CANONICAL_STATE_FORMAT_VERSION", + "PreparedCanonicalStateImport", + ): + from . import canonical + + return getattr(canonical, name) + if name in ( + "PORTABLE_STATE_FORMAT_VERSION", + "build_portable_state_document", + "normalize_portable_state_document", + "portable_state_digest", + ): + from . import portable_schema + + return getattr(portable_schema, name) + if name == "PortableStateLimits": + from .portable_state import PortableStateLimits + + return PortableStateLimits + if name in ("load_portable_dcp", "save_portable_dcp"): + from . import portable_dcp + + return getattr(portable_dcp, name) + if name in ( + "CONTRACT_SCHEMA_VERSION", + "IDENTITY_SCHEMA_VERSION", + "CheckpointSupport", + "CheckpointTransport", + "CanonicalStateProvider", + "OptimizerCapabilities", + "OptimizerChildContract", + "OptimizerContract", + "OptimizerContractProvider", + "OptimizerStateLayout", + "LogicalRegion", + "LogicalSlice", + "ParameterLayout", + "ParameterIdentity", + "ParameterStateRole", + "PlacementKind", + "PortableStateProvider", + "Precision", + "ProcessGroupScope", + "ProcessGroupIdentity", + "ShardIdentity", + "ShardPlacement", + "ShardingManifest", + "StateExtent", + "StateField", + "StateGeometry", + "StateKeyMatch", + "StateScope", + "StateVariant", + "TopologyChange", + "TrainingSupport", + ): + from . import contracts + + return getattr(contracts, name) if name == "kernels": # NOT `from . import kernels`: its fromlist handling re-enters this # __getattr__ before the submodule import runs, recursing forever. diff --git a/src/gefen/canonical.py b/src/gefen/canonical.py new file mode 100644 index 0000000..2f3f45c --- /dev/null +++ b/src/gefen/canonical.py @@ -0,0 +1,158 @@ +"""Primitive, device-neutral helpers for canonical optimizer state.""" + +import math + +import torch + + +CANONICAL_STATE_FORMAT_VERSION = 1 +_IMPORT_PLAN_TOKEN = object() + + +class PreparedCanonicalStateImport: + """Opaque, single-use result of canonical import preparation.""" + + __slots__ = ("_optimizer", "_live_token", "_staged", "_consumed") + + def __init__(self, optimizer, live_token, staged, *, _token=None): + if _token is not _IMPORT_PLAN_TOKEN: + raise TypeError("PreparedCanonicalStateImport values are created by an optimizer") + self._optimizer = optimizer + self._live_token = live_token + self._staged = staged + self._consumed = False + + +def make_prepared_canonical_state_import(optimizer, live_token, staged): + return PreparedCanonicalStateImport( + optimizer, + live_token, + staged, + _token=_IMPORT_PLAN_TOKEN, + ) + + +def canonical_value_supported(value, *, finite_tensors=False) -> bool: + """Return whether ``value`` has a deterministic weights-only wire form.""" + + if type(value) is torch.Tensor: + supported = ( + value.layout is torch.strided + and not value.is_meta + and not value.is_nested + and not value.is_quantized + ) + if ( + supported + and finite_tensors + and (value.is_floating_point() or value.is_complex()) + ): + try: + supported = bool(torch.isfinite(value.detach()).all()) + except (NotImplementedError, RuntimeError, TypeError): + supported = False + return supported + if torch.is_tensor(value): + return False + if type(value) is float: + return math.isfinite(value) + if value is None or type(value) in {bool, int, str}: + return True + if type(value) in {list, tuple}: + return all( + canonical_value_supported( + item, finite_tensors=finite_tensors + ) + for item in value + ) + if type(value) is dict: + return all( + type(key) is str + and canonical_value_supported( + item, finite_tensors=finite_tensors + ) + for key, item in value.items() + ) + return False + + +def clone_canonical_value(value, *, path="value"): + """Clone one supported value into a CPU, weights-only-safe wire value.""" + + if torch.is_tensor(value): + if ( + type(value) is not torch.Tensor + or value.layout is not torch.strided + or value.is_meta + or value.is_nested + or value.is_quantized + ): + raise TypeError( + "{} must be a plain materialized strided tensor for canonical state".format( + path + ) + ) + cloned = ( + value.detach() + .to(device="cpu") + .resolve_conj() + .resolve_neg() + .contiguous() + .clone() + ) + if cloned.is_floating_point() or cloned.is_complex(): + try: + finite = bool(torch.isfinite(cloned).all()) + except (NotImplementedError, RuntimeError, TypeError) as exc: + raise ValueError( + "{} tensor dtype does not support finite canonical state".format( + path + ) + ) from exc + if not finite: + raise ValueError("{} tensor must be finite".format(path)) + return cloned + if type(value) is float: + if not math.isfinite(value): + raise ValueError("{} must be finite".format(path)) + return value + if value is None or type(value) in {bool, int, str}: + return value + if type(value) is list: + return [clone_canonical_value(item, path="{}[{}]".format(path, index)) for index, item in enumerate(value)] + if type(value) is tuple: + return tuple(clone_canonical_value(item, path="{}[{}]".format(path, index)) for index, item in enumerate(value)) + if type(value) is dict: + if any(type(key) is not str for key in value): + raise TypeError("{} dictionary keys must be strings".format(path)) + return {key: clone_canonical_value(value[key], path="{}.{}".format(path, key)) for key in sorted(value)} + raise TypeError("{} has unsupported canonical-state type {}".format(path, type(value).__name__)) + + +def canonical_values_equal(left, right) -> bool: + """Compare canonical values exactly while ignoring tensor device.""" + + if torch.is_tensor(left) or torch.is_tensor(right): + return ( + torch.is_tensor(left) + and torch.is_tensor(right) + and left.dtype == right.dtype + and tuple(left.shape) == tuple(right.shape) + and torch.equal( + left.detach().cpu().contiguous(), + right.detach().cpu().contiguous(), + ) + ) + if type(left) is not type(right): + return False + if type(left) is dict: + return set(left) == set(right) and all(canonical_values_equal(left[key], right[key]) for key in left) + if type(left) in {list, tuple}: + return len(left) == len(right) and all(canonical_values_equal(a, b) for a, b in zip(left, right)) + return left == right + + +__all__ = [ + "CANONICAL_STATE_FORMAT_VERSION", + "PreparedCanonicalStateImport", +] diff --git a/src/gefen/checkpoint.py b/src/gefen/checkpoint.py new file mode 100644 index 0000000..945d8f3 --- /dev/null +++ b/src/gefen/checkpoint.py @@ -0,0 +1,142 @@ +"""Runtime process-group binding for checkpoint adapters. + +The stable identity contains adapter-defined semantic member IDs. Runtime +global ranks are inspected only while validating the live PyTorch process +group; they are deliberately not part of this binding's durable identity. +""" + +from dataclasses import dataclass +from typing import Optional + +import torch + +from gefen.contracts import ProcessGroupIdentity + + +def _is_process_group(value: object) -> bool: + if not torch.distributed.is_available(): + return False + process_group_type = getattr(torch.distributed, "ProcessGroup", None) + return process_group_type is not None and isinstance(value, process_group_type) + + +@dataclass(frozen=True, eq=False) +class CheckpointProcessGroupBinding: + """Bind a stable checkpoint scope to one live PyTorch process group. + + Multi-member scopes always carry an explicit runtime handle, including + ``torch.distributed.group.WORLD`` when the default world is the intended + scope. A one-member scope is local and therefore uses ``None``. The runtime + handle and runtime global ranks are not canonical checkpoint metadata. + """ + + identity: ProcessGroupIdentity + local_member: str + process_group: Optional[object] + collective_device: torch.device + + def __post_init__(self) -> None: + if not isinstance(self.identity, ProcessGroupIdentity): + raise TypeError("CheckpointProcessGroupBinding.identity must be a ProcessGroupIdentity") + if not isinstance(self.local_member, str): + raise TypeError("CheckpointProcessGroupBinding.local_member must be a string") + if self.local_member not in self.identity.ordered_members: + raise ValueError("CheckpointProcessGroupBinding.local_member must belong to the identity") + if not isinstance(self.collective_device, torch.device): + raise TypeError("CheckpointProcessGroupBinding.collective_device must be a torch.device") + if self.collective_device.type not in {"cpu", "cuda"}: + raise ValueError("checkpoint collective device must be CPU or CUDA") + + member_count = len(self.identity.ordered_members) + if member_count == 1: + if self.process_group is not None: + raise ValueError("a one-member checkpoint scope must use process_group=None") + else: + if self.process_group is None: + raise ValueError( + "a multi-member checkpoint scope requires an explicit runtime process group" + ) + if not _is_process_group(self.process_group): + raise TypeError( + "CheckpointProcessGroupBinding.process_group must be a torch.distributed.ProcessGroup" + ) + + @property + def sort_key(self): + """Return the rank-invariant key used to order checkpoint scopes.""" + + return (self.identity.semantic_name, self.identity.ordered_members) + + def validate_runtime(self) -> None: + """Fail unless this process belongs at the declared semantic coordinate. + + This performs only local process-group metadata queries. It does not + execute a collective or retain the runtime group's global-rank list. + """ + + members = self.identity.ordered_members + if len(members) == 1: + return + if not torch.distributed.is_available() or not torch.distributed.is_initialized(): + raise RuntimeError( + "a multi-member checkpoint scope requires initialized torch.distributed" + ) + + import torch.distributed as dist + + try: + world_size = dist.get_world_size(self.process_group) + runtime_global_members = tuple(dist.get_process_group_ranks(self.process_group)) + current_global_rank = dist.get_rank() + group_rank = dist.get_group_rank(self.process_group, current_global_rank) + roundtrip_global_rank = dist.get_global_rank(self.process_group, group_rank) + backend = dist.get_backend(self.process_group) + except Exception as exc: + raise ValueError( + "the current rank must belong to the explicit checkpoint process group" + ) from exc + + if world_size != len(members) or len(runtime_global_members) != world_size: + raise ValueError( + "checkpoint process-group world size does not match its stable identity" + ) + if ( + group_rank < 0 + or group_rank >= world_size + or runtime_global_members[group_rank] != current_global_rank + or roundtrip_global_rank != current_global_rank + ): + raise ValueError( + "the current global rank is not a consistent member of the checkpoint process group" + ) + if members[group_rank] != self.local_member: + raise ValueError( + "runtime checkpoint group order does not match ordered semantic members" + ) + self._validate_collective_device(backend) + + def _validate_collective_device(self, backend: object) -> None: + backend_name = str(backend).lower() + if "nccl" in backend_name: + if self.collective_device.type != "cuda": + raise ValueError( + "checkpoint collective device is incompatible with the runtime backend" + ) + elif "gloo" in backend_name or "mpi" in backend_name: + if self.collective_device.type != "cpu": + raise ValueError( + "checkpoint collective device is incompatible with the runtime backend" + ) + + if self.collective_device.type == "cuda": + index = self.collective_device.index + if ( + not torch.cuda.is_available() + or index is None + or index < 0 + or index >= torch.cuda.device_count() + ): + raise ValueError("checkpoint collective CUDA device is unavailable") + + +__all__ = ["CheckpointProcessGroupBinding"] diff --git a/src/gefen/codebook.py b/src/gefen/codebook.py new file mode 100644 index 0000000..7dced8c --- /dev/null +++ b/src/gefen/codebook.py @@ -0,0 +1,46 @@ +"""Runtime binding for one optimizer-wide learned-codebook process group.""" + +from dataclasses import dataclass +from typing import Optional + +import torch + +from gefen.contracts import ProcessGroupIdentity + + +@dataclass(frozen=True, eq=False) +class CodebookProcessGroupBinding: + """Bind a stable semantic group to one framework runtime group handle. + + ``process_group`` is deliberately opaque to the descriptor. Gefen validates + and consumes it through PyTorch's public distributed APIs when the binding + is installed by ``post_sharding``. A one-member scope uses ``None`` rather + than implicitly selecting the default world. + """ + + identity: ProcessGroupIdentity + local_member: str + process_group: Optional[object] + collective_device: torch.device + + def __post_init__(self) -> None: + if not isinstance(self.identity, ProcessGroupIdentity): + raise TypeError("CodebookProcessGroupBinding.identity must be a ProcessGroupIdentity") + if self.local_member not in self.identity.ordered_members: + raise ValueError("CodebookProcessGroupBinding.local_member must belong to the identity") + try: + device = torch.device(self.collective_device) + except (TypeError, RuntimeError) as exc: + raise TypeError("CodebookProcessGroupBinding.collective_device must be a torch device") from exc + if device.type == "meta": + raise ValueError("codebook collectives require a materialized device") + object.__setattr__(self, "collective_device", device) + + @property + def sort_key(self): + """Return the stable adapter scheduling key for this collective scope.""" + + return (self.identity.semantic_name, self.identity.ordered_members) + + +__all__ = ["CodebookProcessGroupBinding"] diff --git a/src/gefen/contracts.py b/src/gefen/contracts.py new file mode 100644 index 0000000..7b5159f --- /dev/null +++ b/src/gefen/contracts.py @@ -0,0 +1,2069 @@ +"""Platform-agnostic optimizer capability and state-layout contracts. + +The descriptors in this module are read-only declarations. They describe the +optimizer state Gefen owns and the integration surfaces implemented today +without importing DCP, DTensor, Megatron, DeepSpeed, or another adapter. +""" + +from dataclasses import dataclass +from enum import Enum +import math +from typing import ( + AbstractSet, + FrozenSet, + Optional, + Protocol, + Sequence, + Tuple, + Union, + runtime_checkable, +) + + +CONTRACT_SCHEMA_VERSION = 1 +IDENTITY_SCHEMA_VERSION = 1 + + +class StateScope(str, Enum): + """Ownership category for one optimizer-state field.""" + + OPTIMIZER_COMMON = "optimizer_common" + PARAMETER = "parameter" + DERIVED = "derived" + + +class StateGeometry(str, Enum): + """Logical geometry kind, independent of local/global extent.""" + + SCALAR = "scalar" + CODEBOOK = "codebook" + PARAMETER = "parameter" + BLOCK = "block" + ROW = "row" + COLUMN = "column" + OPAQUE = "opaque" + + +class StateKeyMatch(str, Enum): + """How a field declaration matches live or serialized dictionary keys.""" + + EXACT = "exact" + PREFIX = "prefix" + + +class StateExtent(str, Enum): + """Logical extent represented by a per-parameter state variant.""" + + METADATA_ONLY = "metadata_only" + LOCAL_STORAGE = "local_storage" + GLOBAL_PARAMETER = "global_parameter" + OWNER_PARAMETER = "owner_parameter" + + +class ParameterStateRole(str, Enum): + """Ownership role on which a per-parameter variant is present.""" + + ANY = "any" + OWNER = "owner" + NON_OWNER = "non_owner" + + +class ParameterLayout(str, Enum): + """Concrete parameter layouts exposed by the current implementation.""" + + REPLICATED = "replicated" + FLATTENED_ELEMENT_SHARD = "flattened_element_shard" + WHOLE_PARAMETER_OWNER = "whole_parameter_owner" + DTENSOR_1D_DEFAULT_WORLD = "dtensor_1d_default_world" + + +class PlacementKind(str, Enum): + """Framework-neutral placement carried by a stable shard identity.""" + + REPLICATE = "replicate" + FLAT_SHARD = "flat_shard" + DIMENSION_SHARD = "dimension_shard" + WHOLE_PARAMETER_OWNER = "whole_parameter_owner" + + +class ProcessGroupScope(str, Enum): + """How an implemented path obtains its collective process group.""" + + NONE = "none" + DEFAULT_WORLD = "default_world" + INFERRED_DEVICE_MESH = "inferred_device_mesh" + ADAPTER_DEFINED = "adapter_defined" + + +class CheckpointTransport(str, Enum): + """Checkpoint transport with independently scoped topology support.""" + + NATIVE_OPTIMIZER = "native_optimizer" + PYTORCH_RANK_LOCAL = "pytorch_rank_local" + COMPOSITE_NATIVE = "composite_native" + CANONICAL_LOCAL = "canonical_local" + CANONICAL_GLOBAL = "canonical_global" + + +class TopologyChange(str, Enum): + """Specific topology mutation supported by a checkpoint transport.""" + + WORLD_SIZE_OWNER_REDISTRIBUTION = "world_size_owner_redistribution" + PLACEMENT_RESHARD = "placement_reshard" + + +class Precision(str, Enum): + """Parameter/gradient storage precision accepted by validated core paths.""" + + FLOAT32 = "float32" + BFLOAT16 = "bfloat16" + FLOAT16 = "float16" + FLOAT64 = "float64" + + +def _tuple(value): + return tuple(value) + + +def _frozenset(value): + return frozenset(value) + + +def _validate_dimensions(name, values, *, positive): + if len(set(values)) != len(values): + raise ValueError("{} must be unique".format(name)) + minimum = 1 if positive else 0 + if any(type(value) is not int or value < minimum for value in values): + relation = "positive" if positive else "nonnegative" + raise ValueError("{} must contain {} integers".format(name, relation)) + + +def _validate_identity_name(name, value): + if not isinstance(value, str) or not value or value != value.strip(): + raise ValueError( + "{} must be a non-empty string without outer whitespace".format(name) + ) + if "\x00" in value: + raise ValueError("{} must not contain NUL".format(name)) + + +def _validate_identity_schema_version(name, value): + if type(value) is not int or value != IDENTITY_SCHEMA_VERSION: + raise ValueError("unsupported {} schema version".format(name)) + + +@dataclass(frozen=True) +class ParameterIdentity: + """Canonical logical parameter identity, independent of tensor storage.""" + + fqn: str + global_shape: Sequence[int] + schema_version: int = IDENTITY_SCHEMA_VERSION + + def __post_init__(self) -> None: + _validate_identity_name("ParameterIdentity.fqn", self.fqn) + if self.fqn.startswith(".") or self.fqn.endswith(".") or ".." in self.fqn: + raise ValueError( + "ParameterIdentity.fqn must contain non-empty dot-separated components" + ) + if isinstance(self.global_shape, (str, bytes, bytearray)): + raise TypeError( + "ParameterIdentity.global_shape must be a sequence of dimensions" + ) + object.__setattr__(self, "global_shape", _tuple(self.global_shape)) + if any(type(dim) is not int or dim < 0 for dim in self.global_shape): + raise ValueError( + "ParameterIdentity.global_shape must contain nonnegative integers" + ) + _validate_identity_schema_version("parameter identity", self.schema_version) + + @property + def numel(self) -> int: + """Return the canonical logical element count.""" + + return math.prod(self.global_shape) + + +@dataclass(frozen=True) +class ProcessGroupIdentity: + """Stable semantic process-group identity supplied by an adapter.""" + + semantic_name: str + ordered_members: Sequence[str] + schema_version: int = IDENTITY_SCHEMA_VERSION + + def __post_init__(self) -> None: + _validate_identity_name( + "ProcessGroupIdentity.semantic_name", self.semantic_name + ) + if isinstance(self.ordered_members, (str, bytes)): + raise TypeError( + "ProcessGroupIdentity.ordered_members must be a sequence of member IDs" + ) + object.__setattr__(self, "ordered_members", _tuple(self.ordered_members)) + if not self.ordered_members: + raise ValueError("ProcessGroupIdentity.ordered_members must be non-empty") + for member in self.ordered_members: + _validate_identity_name("process-group member", member) + if len(set(self.ordered_members)) != len(self.ordered_members): + raise ValueError("ProcessGroupIdentity.ordered_members must be unique") + _validate_identity_schema_version("process-group identity", self.schema_version) + + +@dataclass(frozen=True) +class ShardPlacement: + """One explicit mesh-axis placement for a logical parameter shard.""" + + mesh_axis: str + kind: PlacementKind + coordinate: int + parts: int + parameter_dimension: Optional[int] = None + + def __post_init__(self) -> None: + _validate_identity_name("ShardPlacement.mesh_axis", self.mesh_axis) + if not isinstance(self.kind, PlacementKind): + raise TypeError("ShardPlacement.kind must be a PlacementKind") + if type(self.parts) is not int or self.parts <= 0: + raise ValueError("ShardPlacement.parts must be a positive integer") + if ( + type(self.coordinate) is not int + or self.coordinate < 0 + or self.coordinate >= self.parts + ): + raise ValueError("ShardPlacement.coordinate must be within parts") + if self.kind is PlacementKind.DIMENSION_SHARD: + if ( + type(self.parameter_dimension) is not int + or self.parameter_dimension < 0 + ): + raise ValueError( + "dimension-shard placement requires a nonnegative parameter dimension" + ) + elif self.parameter_dimension is not None: + raise ValueError( + "only a dimension-shard placement may name a parameter dimension" + ) + + +@dataclass(frozen=True) +class LogicalSlice: + """Contiguous range in canonical row-major flattened parameter order.""" + + flat_offset: int + length: int + + def __post_init__(self) -> None: + if type(self.flat_offset) is not int or self.flat_offset < 0: + raise ValueError("LogicalSlice.flat_offset must be a nonnegative integer") + if type(self.length) is not int or self.length < 0: + raise ValueError("LogicalSlice.length must be a nonnegative integer") + + @classmethod + def full(cls, parameter: ParameterIdentity) -> "LogicalSlice": + """Return the complete logical range for ``parameter``.""" + + if not isinstance(parameter, ParameterIdentity): + raise TypeError("parameter must be a ParameterIdentity") + return cls(0, parameter.numel) + + +@dataclass(frozen=True) +class LogicalRegion: + """One axis-aligned region in canonical logical parameter coordinates.""" + + offsets: Sequence[int] + lengths: Sequence[int] + + def __post_init__(self) -> None: + for name, values in ( + ("LogicalRegion.offsets", self.offsets), + ("LogicalRegion.lengths", self.lengths), + ): + if isinstance(values, (str, bytes, bytearray)): + raise TypeError("{} must be a sequence of dimensions".format(name)) + try: + normalized = tuple(values) + except TypeError as exc: + raise TypeError( + "{} must be a sequence of dimensions".format(name) + ) from exc + if any(type(value) is not int or value < 0 for value in normalized): + raise ValueError("{} must contain nonnegative integers".format(name)) + object.__setattr__(self, name.rsplit(".", 1)[1], normalized) + if len(self.offsets) != len(self.lengths): + raise ValueError( + "LogicalRegion offsets and lengths must have the same rank" + ) + + @property + def rank(self) -> int: + """Return the logical tensor rank described by this region.""" + + return len(self.offsets) + + @property + def numel(self) -> int: + """Return the number of logical elements in this region.""" + + return math.prod(self.lengths) + + @classmethod + def full(cls, parameter: ParameterIdentity) -> "LogicalRegion": + """Return the complete axis-aligned region for ``parameter``.""" + + if not isinstance(parameter, ParameterIdentity): + raise TypeError("parameter must be a ParameterIdentity") + return cls((0,) * len(parameter.global_shape), parameter.global_shape) + + def validate_bounds(self, parameter: ParameterIdentity) -> None: + """Raise when this region is not contained by ``parameter``.""" + + if not isinstance(parameter, ParameterIdentity): + raise TypeError("parameter must be a ParameterIdentity") + if self.rank != len(parameter.global_shape): + raise ValueError("LogicalRegion rank must match the global parameter rank") + if any( + offset + length > dimension + for offset, length, dimension in zip( + self.offsets, self.lengths, parameter.global_shape + ) + ): + raise ValueError("LogicalRegion exceeds the global parameter") + + def intersection(self, other: "LogicalRegion") -> "LogicalRegion": + """Return the axis-aligned intersection with another same-rank region.""" + + if not isinstance(other, LogicalRegion): + raise TypeError("other must be a LogicalRegion") + if self.rank != other.rank: + raise ValueError("LogicalRegion intersection requires equal ranks") + offsets = tuple( + max(left, right) for left, right in zip(self.offsets, other.offsets) + ) + ends = tuple( + min(left_offset + left_length, right_offset + right_length) + for left_offset, left_length, right_offset, right_length in zip( + self.offsets, + self.lengths, + other.offsets, + other.lengths, + ) + ) + return LogicalRegion( + offsets, + tuple(max(0, end - offset) for offset, end in zip(offsets, ends)), + ) + + def overlaps(self, other: "LogicalRegion") -> bool: + """Return whether two regions share at least one logical element.""" + + return self.intersection(other).numel > 0 + + @staticmethod + def validate_exact_coverage( + parameter: ParameterIdentity, regions: Sequence["LogicalRegion"] + ) -> None: + """Validate that bounded regions cover a parameter exactly once.""" + + if not isinstance(parameter, ParameterIdentity): + raise TypeError("parameter must be a ParameterIdentity") + if isinstance(regions, (str, bytes, bytearray)): + raise TypeError("regions must be a sequence of LogicalRegion values") + try: + regions = tuple(regions) + except TypeError as exc: + raise TypeError( + "regions must be a sequence of LogicalRegion values" + ) from exc + if any(not isinstance(region, LogicalRegion) for region in regions): + raise TypeError("regions must contain LogicalRegion values") + for region in regions: + region.validate_bounds(parameter) + for index, region in enumerate(regions): + if any(region.overlaps(other) for other in regions[index + 1 :]): + raise ValueError("LogicalRegions must not overlap") + if sum(region.numel for region in regions) != parameter.numel: + raise ValueError("LogicalRegions must exactly cover the global parameter") + + +@dataclass(frozen=True) +class ShardIdentity: + """Stable identity of one process-group member's logical parameter shard.""" + + parameter: ParameterIdentity + layout: ParameterLayout + logical_slice: Union[LogicalSlice, LogicalRegion] + placements: Sequence[ShardPlacement] = () + process_group: Optional[ProcessGroupIdentity] = None + local_member: Optional[str] = None + owner: Optional[str] = None + schema_version: int = IDENTITY_SCHEMA_VERSION + + def __post_init__(self) -> None: + if not isinstance(self.parameter, ParameterIdentity): + raise TypeError("ShardIdentity.parameter must be a ParameterIdentity") + if not isinstance(self.layout, ParameterLayout): + raise TypeError("ShardIdentity.layout must be a ParameterLayout") + uses_logical_region = isinstance(self.logical_slice, LogicalRegion) + if self.layout is ParameterLayout.DTENSOR_1D_DEFAULT_WORLD: + if not uses_logical_region: + raise ValueError( + "stable DTensor identity requires a logical-region descriptor (LogicalRegion)" + ) + self.logical_slice.validate_bounds(self.parameter) + elif not isinstance(self.logical_slice, LogicalSlice): + raise TypeError( + "non-DTensor ShardIdentity.logical_slice must be a LogicalSlice" + ) + placements = _tuple(self.placements) + if any(not isinstance(item, ShardPlacement) for item in placements): + raise TypeError( + "ShardIdentity.placements must contain ShardPlacement values" + ) + axes = tuple(item.mesh_axis for item in placements) + if len(set(axes)) != len(axes): + raise ValueError("ShardIdentity placement mesh axes must be unique") + object.__setattr__( + self, + "placements", + tuple(sorted(placements, key=lambda item: item.mesh_axis)), + ) + if ( + not uses_logical_region + and self.logical_slice.flat_offset + self.logical_slice.length + > self.parameter.numel + ): + raise ValueError("ShardIdentity.logical_slice exceeds the global parameter") + if self.process_group is None: + if self.local_member is not None or self.owner is not None: + raise ValueError( + "ShardIdentity members and owners require a process-group identity" + ) + else: + if not isinstance(self.process_group, ProcessGroupIdentity): + raise TypeError( + "ShardIdentity.process_group must be a ProcessGroupIdentity" + ) + if self.local_member not in self.process_group.ordered_members: + raise ValueError( + "ShardIdentity.local_member must belong to the process group" + ) + if ( + self.owner is not None + and self.owner not in self.process_group.ordered_members + ): + raise ValueError("ShardIdentity.owner must belong to the process group") + if ( + self.layout is not ParameterLayout.WHOLE_PARAMETER_OWNER + and self.owner is not None + ): + raise ValueError( + "ShardIdentity.owner is valid only for whole-parameter ownership" + ) + + full = ( + self.logical_slice == LogicalRegion.full(self.parameter) + if uses_logical_region + else self.logical_slice == LogicalSlice.full(self.parameter) + ) + kinds = tuple(item.kind for item in self.placements) + if self.process_group is not None: + member_index = self.process_group.ordered_members.index(self.local_member) + for placement in self.placements: + if ( + placement.parts != len(self.process_group.ordered_members) + or placement.coordinate != member_index + ): + raise ValueError( + "ShardIdentity placement coordinates must match the ordered process-group members" + ) + for placement in self.placements: + if ( + placement.parameter_dimension is not None + and placement.parameter_dimension >= len(self.parameter.global_shape) + ): + raise ValueError( + "ShardIdentity placement dimension exceeds parameter rank" + ) + if self.layout is ParameterLayout.REPLICATED: + if not full or self.owner is not None: + raise ValueError("replicated identity must cover the full parameter") + if any(kind is not PlacementKind.REPLICATE for kind in kinds): + raise ValueError("replicated identity has a non-replicated placement") + if self.process_group is None and self.placements: + raise ValueError("an ungrouped replicated identity has no placements") + if self.process_group is not None and len(self.placements) != 1: + raise ValueError( + "a process-group replicated identity requires one placement" + ) + elif self.layout is ParameterLayout.FLATTENED_ELEMENT_SHARD: + if self.process_group is None or self.owner is not None: + raise ValueError( + "flattened element shards require a process group and no owner" + ) + if len(kinds) != 1 or kinds[0] is not PlacementKind.FLAT_SHARD: + raise ValueError( + "flattened element shards require one flat-shard placement" + ) + elif self.layout is ParameterLayout.WHOLE_PARAMETER_OWNER: + if self.process_group is None or self.owner is None: + raise ValueError( + "whole-parameter ownership requires a process group and owner" + ) + owns_parameter = self.local_member == self.owner + if owns_parameter and not full: + raise ValueError( + "the owner must carry the full whole-parameter logical slice" + ) + if not owns_parameter and self.logical_slice != LogicalSlice(0, 0): + raise ValueError("a non-owner whole-parameter slice must be empty") + if len(kinds) != 1 or kinds[0] is not PlacementKind.WHOLE_PARAMETER_OWNER: + raise ValueError( + "whole-parameter ownership requires one owner placement" + ) + elif self.layout is ParameterLayout.DTENSOR_1D_DEFAULT_WORLD: + if self.process_group is None or self.owner is not None: + raise ValueError( + "DTensor identities require a process group and no owner" + ) + if len(self.placements) != 1 or kinds[0] not in { + PlacementKind.DIMENSION_SHARD, + PlacementKind.REPLICATE, + }: + raise ValueError( + "one-dimensional DTensor identities require one dimension-shard or replicate placement" + ) + placement = self.placements[0] + if placement.kind is PlacementKind.REPLICATE: + if not full: + raise ValueError( + "replicated DTensor identities must cover the full parameter" + ) + else: + shard_dimension = placement.parameter_dimension + for dimension, (offset, length, global_length) in enumerate( + zip( + self.logical_slice.offsets, + self.logical_slice.lengths, + self.parameter.global_shape, + ) + ): + if dimension != shard_dimension and ( + offset != 0 or length != global_length + ): + raise ValueError( + "a dimension-sharded DTensor region must cover every unsharded parameter dimension" + ) + _validate_identity_schema_version("shard identity", self.schema_version) + + @property + def logical_region(self) -> Optional[LogicalRegion]: + """Return the axis-aligned logical region, when this identity has one.""" + + if isinstance(self.logical_slice, LogicalRegion): + return self.logical_slice + return None + + @property + def sort_key(self): + """Return a deterministic structural ordering key.""" + + member_index = -1 + group_name = "" + if self.process_group is not None: + group_name = self.process_group.semantic_name + member_index = self.process_group.ordered_members.index(self.local_member) + owner_index = -1 + if self.process_group is not None and self.owner is not None: + owner_index = self.process_group.ordered_members.index(self.owner) + placement_key = tuple( + ( + item.mesh_axis, + item.kind.value, + item.coordinate, + item.parts, + -1 if item.parameter_dimension is None else item.parameter_dimension, + ) + for item in self.placements + ) + if isinstance(self.logical_slice, LogicalSlice): + # Preserve the released contiguous-slice key exactly. Besides + # retaining its public structural shape, this leaves all existing + # manifest ordering and freshness tokens byte-for-byte stable. + return ( + self.parameter.fqn, + self.logical_slice.flat_offset, + self.logical_slice.length, + self.layout.value, + group_name, + member_index, + owner_index, + placement_key, + ) + + flat_offset = 0 + stride = 1 + for offset, dimension in reversed( + tuple(zip(self.logical_slice.offsets, self.parameter.global_shape)) + ): + flat_offset += offset * stride + stride *= dimension + return ( + self.parameter.fqn, + flat_offset, + self.logical_slice.numel, + self.layout.value, + group_name, + member_index, + owner_index, + placement_key, + self.logical_slice.offsets, + self.logical_slice.lengths, + ) + + +@dataclass(frozen=True) +class ShardingManifest: + """Complete deterministic shard identity set supplied by one adapter.""" + + shards: Sequence[ShardIdentity] + schema_version: int = IDENTITY_SCHEMA_VERSION + + def __post_init__(self) -> None: + shards = _tuple(self.shards) + if not shards: + raise ValueError("ShardingManifest.shards must be non-empty") + if any(not isinstance(item, ShardIdentity) for item in shards): + raise TypeError("ShardingManifest.shards must contain ShardIdentity values") + ordered = tuple(sorted(shards, key=lambda item: item.sort_key)) + if len(set(ordered)) != len(ordered): + raise ValueError("ShardingManifest.shards must be unique") + object.__setattr__(self, "shards", ordered) + _validate_identity_schema_version("sharding manifest", self.schema_version) + + by_fqn = {} + for shard in ordered: + by_fqn.setdefault(shard.parameter.fqn, []).append(shard) + for fqn, parameter_shards in by_fqn.items(): + parameter = parameter_shards[0].parameter + if any(item.parameter != parameter for item in parameter_shards[1:]): + raise ValueError( + "manifest shards for {!r} disagree on parameter identity".format( + fqn + ) + ) + layouts = {item.layout for item in parameter_shards} + groups = {item.process_group for item in parameter_shards} + if len(layouts) != 1 or len(groups) != 1: + raise ValueError( + "manifest shards for {!r} disagree on layout or process group".format( + fqn + ) + ) + layout = parameter_shards[0].layout + group = parameter_shards[0].process_group + members = tuple(item.local_member for item in parameter_shards) + if group is None: + if len(parameter_shards) != 1: + raise ValueError( + "an ungrouped parameter must have exactly one manifest shard" + ) + elif set(members) != set(group.ordered_members) or len(members) != len( + group.ordered_members + ): + raise ValueError( + "manifest shards must contain each process-group member exactly once" + ) + if group is not None: + placement_shapes = { + tuple( + ( + placement.mesh_axis, + placement.kind, + placement.parts, + placement.parameter_dimension, + ) + for placement in item.placements + ) + for item in parameter_shards + } + if len(placement_shapes) != 1: + raise ValueError( + "manifest member placements must agree on one topology" + ) + + if layout is ParameterLayout.FLATTENED_ELEMENT_SHARD: + cursor = 0 + boundaries = {0} + empty_offsets = [] + for item in sorted( + parameter_shards, + key=lambda shard: ( + shard.logical_slice.flat_offset, + group.ordered_members.index(shard.local_member), + ), + ): + if item.logical_slice.length == 0: + empty_offsets.append(item.logical_slice.flat_offset) + continue + if item.logical_slice.flat_offset != cursor: + raise ValueError( + "flattened manifest slices must be gapless and non-overlapping" + ) + cursor += item.logical_slice.length + boundaries.add(cursor) + if cursor != parameter.numel: + raise ValueError( + "flattened manifest slices must cover the global parameter" + ) + if any(offset not in boundaries for offset in empty_offsets): + raise ValueError( + "empty flattened manifest slices must use a partition boundary" + ) + elif layout is ParameterLayout.REPLICATED: + if any( + item.logical_slice != LogicalSlice.full(parameter) + for item in parameter_shards + ): + raise ValueError("replicated manifest shards must all be complete") + elif layout is ParameterLayout.WHOLE_PARAMETER_OWNER: + owners = {item.owner for item in parameter_shards} + if len(owners) != 1: + raise ValueError( + "whole-parameter manifest shards must agree on one owner" + ) + elif layout is ParameterLayout.DTENSOR_1D_DEFAULT_WORLD: + placements = tuple(item.placements[0] for item in parameter_shards) + placement_kind = placements[0].kind + if placement_kind is PlacementKind.REPLICATE: + if any( + item.logical_region != LogicalRegion.full(parameter) + for item in parameter_shards + ): + raise ValueError( + "replicated DTensor manifest regions must all be complete" + ) + continue + + shard_dimension = placements[0].parameter_dimension + by_coordinate = sorted( + parameter_shards, + key=lambda item: item.placements[0].coordinate, + ) + cursor = 0 + for item in by_coordinate: + region = item.logical_region + if region.offsets[shard_dimension] != cursor: + raise ValueError( + "dimension-sharded DTensor manifest regions must be " + "gapless and non-overlapping in coordinate order" + ) + cursor += region.lengths[shard_dimension] + if cursor != parameter.global_shape[shard_dimension]: + raise ValueError( + "dimension-sharded DTensor manifest regions must cover the global parameter" + ) + + def for_parameter(self, fqn: str) -> Tuple[ShardIdentity, ...]: + """Return one canonical parameter's shards in deterministic order.""" + + return tuple(item for item in self.shards if item.parameter.fqn == fqn) + + +@dataclass(frozen=True) +class StateField: + """One named authoritative field, runtime cache, or transport field.""" + + name: str + scope: StateScope + geometry: StateGeometry + checkpointed: bool + key_match: StateKeyMatch = StateKeyMatch.EXACT + applicable_sharded_modes: AbstractSet[str] = frozenset() + description: str = "" + optional: bool = False + + def __post_init__(self) -> None: + object.__setattr__( + self, + "applicable_sharded_modes", + _frozenset(self.applicable_sharded_modes), + ) + if not self.name: + raise ValueError("StateField.name must be non-empty") + if not isinstance(self.scope, StateScope): + raise TypeError("StateField.scope must be a StateScope") + if not isinstance(self.geometry, StateGeometry): + raise TypeError("StateField.geometry must be a StateGeometry") + if not isinstance(self.key_match, StateKeyMatch): + raise TypeError("StateField.key_match must be a StateKeyMatch") + if type(self.optional) is not bool: + raise TypeError("StateField.optional must be a bool") + + @property + def authoritative(self) -> bool: + """Whether the field carries optimizer meaning rather than a cache.""" + + return self.scope is not StateScope.DERIVED + + def matches(self, key: str) -> bool: + """Return whether a live or serialized key matches this declaration.""" + + if self.key_match is StateKeyMatch.PREFIX: + return isinstance(key, str) and key.startswith(self.name) + return key == self.name + + +@dataclass(frozen=True) +class StateVariant: + """A machine-selectable valid per-parameter state combination.""" + + name: str + fields: Sequence[str] + layouts: AbstractSet[ParameterLayout] + extent: StateExtent + role: ParameterStateRole = ParameterStateRole.ANY + initialized: bool = True + parameter_ranks: Optional[Sequence[int]] = None + excluded_parameter_ranks: Sequence[int] = () + sharded_mode: Optional[str] = None + inactive_fields: Sequence[str] = () + migration_only: bool = False + description: str = "" + + def __post_init__(self) -> None: + object.__setattr__(self, "fields", _tuple(self.fields)) + object.__setattr__(self, "layouts", _frozenset(self.layouts)) + object.__setattr__(self, "inactive_fields", _tuple(self.inactive_fields)) + object.__setattr__( + self, "excluded_parameter_ranks", _tuple(self.excluded_parameter_ranks) + ) + if self.parameter_ranks is not None: + object.__setattr__(self, "parameter_ranks", _tuple(self.parameter_ranks)) + if not self.name: + raise ValueError("StateVariant.name must be non-empty") + if not self.fields: + raise ValueError("StateVariant.fields must be non-empty") + if len(set(self.fields)) != len(self.fields): + raise ValueError("StateVariant.fields must not contain duplicates") + if not self.layouts: + raise ValueError("StateVariant.layouts must be non-empty") + if not isinstance(self.extent, StateExtent): + raise TypeError("StateVariant.extent must be a StateExtent") + if not isinstance(self.role, ParameterStateRole): + raise TypeError("StateVariant.role must be a ParameterStateRole") + if not set(self.inactive_fields).issubset(self.fields): + raise ValueError("StateVariant.inactive_fields must be present in fields") + if self.parameter_ranks is not None and set(self.parameter_ranks) & set( + self.excluded_parameter_ranks + ): + raise ValueError("included and excluded parameter ranks must be disjoint") + if self.parameter_ranks is not None: + _validate_dimensions( + "parameter_ranks", self.parameter_ranks, positive=False + ) + _validate_dimensions( + "excluded_parameter_ranks", + self.excluded_parameter_ranks, + positive=False, + ) + + +@dataclass(frozen=True) +class OptimizerStateLayout: + """Declared fields, per-parameter variants, and composite namespaces.""" + + fields: Sequence[StateField] + parameter_variants: Sequence[StateVariant] + composite_namespaces: Sequence[str] = () + + def __post_init__(self) -> None: + object.__setattr__(self, "fields", _tuple(self.fields)) + object.__setattr__(self, "parameter_variants", _tuple(self.parameter_variants)) + object.__setattr__( + self, "composite_namespaces", _tuple(self.composite_namespaces) + ) + names = tuple(field.name for field in self.fields) + if len(set(names)) != len(names): + raise ValueError("OptimizerStateLayout fields must have unique names") + if len(set(self.composite_namespaces)) != len(self.composite_namespaces): + raise ValueError("composite namespaces must be unique") + variant_names = tuple(variant.name for variant in self.parameter_variants) + if len(set(variant_names)) != len(variant_names): + raise ValueError("StateVariant names must be unique") + parameter_fields = { + field.name for field in self.fields if field.scope is StateScope.PARAMETER + } + for variant in self.parameter_variants: + unknown = set(variant.fields) - parameter_fields + if unknown: + raise ValueError( + "StateVariant {!r} references undeclared parameter fields: {}".format( + variant.name, sorted(unknown) + ) + ) + + def fields_for_scope(self, scope: StateScope) -> Tuple[StateField, ...]: + """Return fields in declaration order for one ownership scope.""" + + return tuple(field for field in self.fields if field.scope is scope) + + def field(self, name: str) -> StateField: + """Return one declared field by name.""" + + for field in self.fields: + if field.matches(name): + return field + raise KeyError(name) + + +@dataclass(frozen=True) +class TrainingSupport: + """One qualified training layout capability.""" + + layout: ParameterLayout + process_group_scope: ProcessGroupScope + mesh_dimensions: Optional[Sequence[int]] = None + sharded_mode: Optional[str] = None + requires_complete_parameter_storage: bool = False + requires_complete_logical_matrix: bool = False + requires_post_step_parameter_sync: bool = False + + def __post_init__(self) -> None: + if self.mesh_dimensions is not None: + object.__setattr__(self, "mesh_dimensions", _tuple(self.mesh_dimensions)) + _validate_dimensions("mesh_dimensions", self.mesh_dimensions, positive=True) + if not isinstance(self.layout, ParameterLayout): + raise TypeError("TrainingSupport.layout must be a ParameterLayout") + if not isinstance(self.process_group_scope, ProcessGroupScope): + raise TypeError( + "TrainingSupport.process_group_scope must be a ProcessGroupScope" + ) + for name in ( + "requires_complete_parameter_storage", + "requires_complete_logical_matrix", + "requires_post_step_parameter_sync", + ): + if type(getattr(self, name)) is not bool: + raise TypeError("TrainingSupport.{} must be a bool".format(name)) + + +@dataclass(frozen=True) +class CheckpointSupport: + """One transport's separately qualified checkpoint capability. + + ``atomic_load`` means each participating optimizer instance validates and + prepares its core restore before local mutation. It does not claim a + coordinated all-rank commit after failures outside the declared process + group or arbitrary user hook side effects. + """ + + transport: CheckpointTransport + same_topology: AbstractSet[ParameterLayout] + topology_changing: AbstractSet[ParameterLayout] + process_group_scope: ProcessGroupScope + topology_change_kinds: AbstractSet[TopologyChange] = frozenset() + mesh_dimensions: Optional[Sequence[int]] = None + required_sharded_modes: AbstractSet[str] = frozenset() + requires_collective: bool = False + atomic_load: bool = False + + def __post_init__(self) -> None: + object.__setattr__(self, "same_topology", _frozenset(self.same_topology)) + object.__setattr__( + self, "topology_changing", _frozenset(self.topology_changing) + ) + object.__setattr__( + self, "topology_change_kinds", _frozenset(self.topology_change_kinds) + ) + object.__setattr__( + self, + "required_sharded_modes", + _frozenset(self.required_sharded_modes), + ) + if self.mesh_dimensions is not None: + object.__setattr__(self, "mesh_dimensions", _tuple(self.mesh_dimensions)) + _validate_dimensions("mesh_dimensions", self.mesh_dimensions, positive=True) + if not isinstance(self.transport, CheckpointTransport): + raise TypeError("CheckpointSupport.transport must be a CheckpointTransport") + if not isinstance(self.process_group_scope, ProcessGroupScope): + raise TypeError( + "CheckpointSupport.process_group_scope must be a ProcessGroupScope" + ) + if bool(self.topology_changing) != bool(self.topology_change_kinds): + raise ValueError( + "topology-changing layouts and change kinds must be declared together" + ) + + +@dataclass(frozen=True) +class OptimizerCapabilities: + """Implemented integration capabilities, including explicit negative claims.""" + + training: Sequence[TrainingSupport] + checkpoints: Sequence[CheckpointSupport] + precisions: AbstractSet[Precision] + supported_parameter_ranks: Optional[Sequence[int]] + accepts_semantic_parameter_names: bool + canonical_parameter_fqns: bool + stable_shard_identity: bool + explicit_process_group_codebook_scope: bool + shard_rebinding: bool + post_sharding: bool + canonical_state_io: bool + atomic_state_movement: bool + state_offload: bool + + def __post_init__(self) -> None: + object.__setattr__(self, "training", _tuple(self.training)) + object.__setattr__(self, "checkpoints", _tuple(self.checkpoints)) + object.__setattr__(self, "precisions", _frozenset(self.precisions)) + if any(not isinstance(item, TrainingSupport) for item in self.training): + raise TypeError( + "OptimizerCapabilities.training must contain TrainingSupport values" + ) + if any(not isinstance(item, CheckpointSupport) for item in self.checkpoints): + raise TypeError( + "OptimizerCapabilities.checkpoints must contain CheckpointSupport values" + ) + if any(not isinstance(item, Precision) for item in self.precisions): + raise TypeError( + "OptimizerCapabilities.precisions must contain Precision values" + ) + for name in ( + "accepts_semantic_parameter_names", + "canonical_parameter_fqns", + "stable_shard_identity", + "explicit_process_group_codebook_scope", + "shard_rebinding", + "post_sharding", + "canonical_state_io", + "atomic_state_movement", + "state_offload", + ): + if type(getattr(self, name)) is not bool: + raise TypeError( + "OptimizerCapabilities.{} must be a bool".format(name) + ) + if self.supported_parameter_ranks is not None: + object.__setattr__( + self, + "supported_parameter_ranks", + _tuple(self.supported_parameter_ranks), + ) + _validate_dimensions( + "supported_parameter_ranks", + self.supported_parameter_ranks, + positive=False, + ) + + +@dataclass(frozen=True) +class OptimizerChildContract: + """One named child of a composite optimizer contract.""" + + role: str + implementation: str + contract: Optional["OptimizerContract"] + + def __post_init__(self) -> None: + if not self.role: + raise ValueError("OptimizerChildContract.role must be non-empty") + if not self.implementation: + raise ValueError("OptimizerChildContract.implementation must be non-empty") + if self.contract is not None and not isinstance( + self.contract, OptimizerContract + ): + raise TypeError( + "OptimizerChildContract.contract must be an OptimizerContract or None" + ) + + +@dataclass(frozen=True) +class OptimizerContract: + """Complete read-only state and capability contract for one optimizer.""" + + implementation: str + state_layout: OptimizerStateLayout + capabilities: OptimizerCapabilities + children: Sequence[OptimizerChildContract] = () + schema_version: int = CONTRACT_SCHEMA_VERSION + + def __post_init__(self) -> None: + object.__setattr__(self, "children", _tuple(self.children)) + if not self.implementation: + raise ValueError("OptimizerContract.implementation must be non-empty") + if self.schema_version != CONTRACT_SCHEMA_VERSION: + raise ValueError( + "unsupported optimizer contract schema version: {}".format( + self.schema_version + ) + ) + roles = tuple(child.role for child in self.children) + if len(set(roles)) != len(roles): + raise ValueError("OptimizerContract child roles must be unique") + + +@runtime_checkable +class OptimizerContractProvider(Protocol): + """Structural protocol implemented by optimizers that expose a contract.""" + + def optimizer_contract(self) -> OptimizerContract: + """Return the optimizer's immutable state and capability declaration.""" + + +@runtime_checkable +class CanonicalStateProvider(Protocol): + """Structural protocol for exact-binding canonical local state.""" + + def export_canonical_state(self): + """Return a versioned local canonical state fragment.""" + + def prepare_canonical_state_import(self, state): + """Validate and stage a canonical fragment without live mutation.""" + + def commit_canonical_state_import(self, prepared) -> None: + """Commit one still-current prepared import.""" + + def import_canonical_state(self, state) -> None: + """Prepare and commit a canonical fragment atomically.""" + + +@runtime_checkable +class PortableStateProvider(Protocol): + """Structural protocol for collective topology-neutral state I/O.""" + + def export_portable_state( + self, *, checkpoint_process_group, transaction_id, limits + ): + """Collectively return one complete portable global-state document.""" + + def import_portable_state( + self, + state, + *, + checkpoint_process_group, + transaction_id, + limits, + ) -> None: + """Collectively stage and atomically publish portable global state.""" + + +_ALL_PRECISIONS = frozenset( + {Precision.FLOAT32, Precision.BFLOAT16, Precision.FLOAT16, Precision.FLOAT64} +) +_DTENSOR_LAYOUT = ParameterLayout.DTENSOR_1D_DEFAULT_WORLD +_MUON_MODES = frozenset({"exact", "approx", "distributed"}) + + +def _common_fields() -> Tuple[StateField, ...]: + return ( + StateField( + "gefen_global_step", + StateScope.OPTIMIZER_COMMON, + StateGeometry.SCALAR, + True, + description="Canonical optimizer-wide update counter.", + ), + StateField( + "gefen_codebook", + StateScope.OPTIMIZER_COMMON, + StateGeometry.CODEBOOK, + True, + description="Canonical learned 256-entry momentum codebook.", + ), + StateField( + "gefen_deterministic", + StateScope.OPTIMIZER_COMMON, + StateGeometry.SCALAR, + True, + description="Checkpoint-bound deterministic execution policy.", + ), + StateField( + "gefen_codebook_scope", + StateScope.OPTIMIZER_COMMON, + StateGeometry.OPAQUE, + True, + description="Stable adapter-defined learned-codebook process-group scope.", + optional=True, + ), + ) + + +def _base_parameter_fields() -> Tuple[StateField, ...]: + return ( + StateField("name", StateScope.PARAMETER, StateGeometry.OPAQUE, True), + StateField( + "automatic_period", StateScope.PARAMETER, StateGeometry.SCALAR, True + ), + StateField("step", StateScope.PARAMETER, StateGeometry.SCALAR, True), + StateField("m_codebook", StateScope.PARAMETER, StateGeometry.PARAMETER, True), + StateField("m_magnitude", StateScope.PARAMETER, StateGeometry.BLOCK, True), + ) + + +def _derived_fields() -> Tuple[StateField, ...]: + return ( + StateField("stepsize", StateScope.DERIVED, StateGeometry.BLOCK, False), + StateField("_h_buf", StateScope.DERIVED, StateGeometry.BLOCK, False), + StateField("_capt_scalars", StateScope.DERIVED, StateGeometry.OPAQUE, False), + StateField("_capt_consts", StateScope.DERIVED, StateGeometry.OPAQUE, False), + StateField("_capt_consts_key", StateScope.DERIVED, StateGeometry.OPAQUE, False), + StateField("_capt_stack", StateScope.DERIVED, StateGeometry.OPAQUE, False), + StateField("_capt_row", StateScope.DERIVED, StateGeometry.OPAQUE, False), + StateField("_param_names", StateScope.DERIVED, StateGeometry.OPAQUE, False), + StateField("_lr_scalar_cache", StateScope.DERIVED, StateGeometry.SCALAR, False), + StateField("_fused_build_ok", StateScope.DERIVED, StateGeometry.SCALAR, False), + StateField("_static_mark_sig", StateScope.DERIVED, StateGeometry.OPAQUE, False), + StateField( + "m_codebook_shape", + StateScope.DERIVED, + StateGeometry.OPAQUE, + True, + description="Inert legacy checkpoint metadata retained for compatibility.", + ), + StateField( + "_gefen_codebook_by_device", + StateScope.DERIVED, + StateGeometry.CODEBOOK, + False, + ), + StateField( + "_gefen_codebook_lut_by_device", + StateScope.DERIVED, + StateGeometry.OPAQUE, + False, + ), + StateField( + "_gefen_codebook_process_group", + StateScope.DERIVED, + StateGeometry.OPAQUE, + False, + description="Live runtime handle for the adapter-defined codebook scope.", + ), + StateField( + "_gefen_codebook_scope_validated", + StateScope.DERIVED, + StateGeometry.SCALAR, + False, + description="Rebuildable cross-member scope-agreement cache.", + ), + StateField( + "_sr_seed_by_device", StateScope.DERIVED, StateGeometry.SCALAR, False + ), + StateField( + "_gefen_global_step_by_device", + StateScope.DERIVED, + StateGeometry.SCALAR, + False, + ), + StateField("_capt_stacks", StateScope.DERIVED, StateGeometry.OPAQUE, False), + StateField( + "_gefen_rank_local_payload_", + StateScope.DERIVED, + StateGeometry.OPAQUE, + True, + key_match=StateKeyMatch.PREFIX, + description="PyTorch rank-local checkpoint transport.", + ), + StateField( + "_gefen_rank_local_member", + StateScope.DERIVED, + StateGeometry.OPAQUE, + True, + description="PyTorch rank-local checkpoint transport marker.", + ), + StateField( + "_gefen_checkpoint_metadata", + StateScope.DERIVED, + StateGeometry.OPAQUE, + True, + description="PyTorch transport mirror of optimizer-common state.", + ), + StateField( + "gefen_native_local_shards", + StateScope.DERIVED, + StateGeometry.OPAQUE, + True, + description="Native rank-local identity guard for scoped physical shards.", + ), + ) + + +def _base_training() -> Tuple[TrainingSupport, ...]: + return ( + TrainingSupport(ParameterLayout.REPLICATED, ProcessGroupScope.NONE), + TrainingSupport( + _DTENSOR_LAYOUT, + ProcessGroupScope.INFERRED_DEVICE_MESH, + mesh_dimensions=(1,), + ), + ) + + +def _negative_capabilities( + *, + training: Tuple[TrainingSupport, ...], + checkpoints: Tuple[CheckpointSupport, ...], + supported_parameter_ranks: Optional[Tuple[int, ...]], + canonical_parameter_fqns: bool = False, + stable_shard_identity: bool = False, + explicit_process_group_codebook_scope: bool = False, + shard_rebinding: bool = False, + post_sharding: bool = False, + canonical_state_io: bool = False, +) -> OptimizerCapabilities: + return OptimizerCapabilities( + training=training, + checkpoints=checkpoints, + precisions=_ALL_PRECISIONS, + supported_parameter_ranks=supported_parameter_ranks, + accepts_semantic_parameter_names=True, + canonical_parameter_fqns=canonical_parameter_fqns, + stable_shard_identity=stable_shard_identity, + explicit_process_group_codebook_scope=explicit_process_group_codebook_scope, + shard_rebinding=shard_rebinding, + post_sharding=post_sharding, + canonical_state_io=canonical_state_io, + atomic_state_movement=False, + state_offload=False, + ) + + +def _gefen_contract( + *, + factored_v_2d: bool, + canonical_parameter_fqns: bool = False, + stable_shard_identity: bool = False, + explicit_process_group_codebook_scope: bool = False, + native_flattened_checkpoint: bool = False, + canonical_state_layouts: AbstractSet[ParameterLayout] = frozenset(), + canonical_global_same_topology: AbstractSet[ParameterLayout] = frozenset(), + canonical_global_topology_changing: AbstractSet[ParameterLayout] = frozenset(), + canonical_global_topology_change_kinds: AbstractSet[TopologyChange] = frozenset(), +) -> OptimizerContract: + canonical_state_layouts = _frozenset(canonical_state_layouts) + canonical_global_same_topology = _frozenset(canonical_global_same_topology) + canonical_global_topology_changing = _frozenset(canonical_global_topology_changing) + canonical_global_topology_change_kinds = _frozenset( + canonical_global_topology_change_kinds + ) + block_fields = ( + StateField("vmean", StateScope.PARAMETER, StateGeometry.BLOCK, True), + StateField("vmean_step", StateScope.PARAMETER, StateGeometry.SCALAR, True), + ) + factored_fields = ( + StateField("v_row", StateScope.PARAMETER, StateGeometry.ROW, True), + StateField("v_col", StateScope.PARAMETER, StateGeometry.COLUMN, True), + StateField("factored_step", StateScope.PARAMETER, StateGeometry.SCALAR, True), + ) + layouts = frozenset( + { + ParameterLayout.REPLICATED, + ParameterLayout.FLATTENED_ELEMENT_SHARD, + _DTENSOR_LAYOUT, + } + ) + base_names = tuple(field.name for field in _base_parameter_fields()) + variants = [ + StateVariant( + "name_only", + ("name",), + layouts, + StateExtent.METADATA_ONLY, + initialized=False, + ), + StateVariant( + "period_selected", + ("name", "automatic_period"), + layouts, + StateExtent.METADATA_ONLY, + initialized=False, + ), + ] + factored_names = tuple(field.name for field in factored_fields) + block_names = tuple(field.name for field in block_fields) + legacy_block_names = ("vmean",) + if factored_v_2d: + variants.extend( + ( + StateVariant( + "block_second_moment_non_2d", + base_names + block_names, + frozenset({ParameterLayout.REPLICATED}), + StateExtent.LOCAL_STORAGE, + excluded_parameter_ranks=(2,), + ), + StateVariant( + "block_second_moment_sharded", + base_names + block_names, + frozenset( + { + ParameterLayout.FLATTENED_ELEMENT_SHARD, + _DTENSOR_LAYOUT, + } + ), + StateExtent.LOCAL_STORAGE, + ), + StateVariant( + "factored_second_moment", + base_names + factored_names, + frozenset({ParameterLayout.REPLICATED}), + StateExtent.GLOBAL_PARAMETER, + parameter_ranks=(2,), + ), + StateVariant( + "block_state_pending_factored_initialization", + base_names + block_names, + frozenset({ParameterLayout.REPLICATED}), + StateExtent.GLOBAL_PARAMETER, + parameter_ranks=(2,), + migration_only=True, + ), + StateVariant( + "legacy_block_state_pending_factored_initialization", + base_names + legacy_block_names, + frozenset({ParameterLayout.REPLICATED}), + StateExtent.GLOBAL_PARAMETER, + parameter_ranks=(2,), + migration_only=True, + ), + StateVariant( + "factored_with_retained_block_state", + base_names + factored_names + block_names, + frozenset({ParameterLayout.REPLICATED}), + StateExtent.GLOBAL_PARAMETER, + parameter_ranks=(2,), + inactive_fields=block_names, + migration_only=True, + description="Factored-v state after loading a block-v checkpoint.", + ), + StateVariant( + "factored_with_retained_legacy_block_state", + base_names + factored_names + legacy_block_names, + frozenset({ParameterLayout.REPLICATED}), + StateExtent.GLOBAL_PARAMETER, + parameter_ranks=(2,), + inactive_fields=legacy_block_names, + migration_only=True, + description="Factored-v state after loading a pre-vmean-counter checkpoint.", + ), + ) + ) + else: + variants.extend( + ( + StateVariant( + "block_second_moment", + base_names + block_names, + layouts, + StateExtent.LOCAL_STORAGE, + ), + StateVariant( + "factored_state_pending_block_initialization", + base_names + factored_names, + frozenset({ParameterLayout.REPLICATED}), + StateExtent.GLOBAL_PARAMETER, + parameter_ranks=(2,), + migration_only=True, + ), + StateVariant( + "legacy_block_without_vmean_counter", + base_names + legacy_block_names, + layouts, + StateExtent.LOCAL_STORAGE, + migration_only=True, + ), + StateVariant( + "block_with_retained_factored_state", + base_names + block_names + factored_names, + frozenset({ParameterLayout.REPLICATED}), + StateExtent.LOCAL_STORAGE, + parameter_ranks=(2,), + inactive_fields=factored_names, + migration_only=True, + description="Block-v state after loading a factored-v checkpoint.", + ), + ) + ) + fields = ( + _common_fields() + _base_parameter_fields() + block_fields + factored_fields + ) + fields += _derived_fields() + training = _base_training() + ( + TrainingSupport( + ParameterLayout.FLATTENED_ELEMENT_SHARD, + ProcessGroupScope.NONE, + ), + ) + checkpoints = [ + CheckpointSupport( + CheckpointTransport.NATIVE_OPTIMIZER, + frozenset({ParameterLayout.REPLICATED}) + | ( + frozenset({ParameterLayout.FLATTENED_ELEMENT_SHARD}) + if native_flattened_checkpoint + else frozenset() + ), + frozenset(), + ProcessGroupScope.NONE, + atomic_load=True, + ), + CheckpointSupport( + CheckpointTransport.PYTORCH_RANK_LOCAL, + frozenset({_DTENSOR_LAYOUT}), + frozenset(), + ProcessGroupScope.DEFAULT_WORLD, + mesh_dimensions=(1,), + requires_collective=True, + atomic_load=True, + ), + ] + if canonical_state_layouts: + checkpoints.append( + CheckpointSupport( + CheckpointTransport.CANONICAL_LOCAL, + canonical_state_layouts, + frozenset(), + ProcessGroupScope.NONE, + atomic_load=True, + ) + ) + if canonical_global_same_topology or canonical_global_topology_changing: + checkpoints.append( + CheckpointSupport( + CheckpointTransport.CANONICAL_GLOBAL, + canonical_global_same_topology, + canonical_global_topology_changing, + ProcessGroupScope.ADAPTER_DEFINED, + topology_change_kinds=canonical_global_topology_change_kinds, + requires_collective=True, + atomic_load=True, + ) + ) + return OptimizerContract( + implementation="gefen.Gefen", + state_layout=OptimizerStateLayout(fields, tuple(variants)), + capabilities=_negative_capabilities( + training=training, + checkpoints=tuple(checkpoints), + supported_parameter_ranks=None, + canonical_parameter_fqns=canonical_parameter_fqns, + stable_shard_identity=stable_shard_identity, + explicit_process_group_codebook_scope=explicit_process_group_codebook_scope, + shard_rebinding=True, + post_sharding=True, + canonical_state_io=bool( + canonical_state_layouts + or canonical_global_same_topology + or canonical_global_topology_changing + ), + ), + ) + + +def _muon_initialized_variant( + *, + name: str, + field_names: Tuple[str, ...], + layout: ParameterLayout, + extent: StateExtent, + mode: str, + role: ParameterStateRole = ParameterStateRole.ANY, +) -> StateVariant: + return StateVariant( + name, + field_names, + frozenset({layout}), + extent, + role=role, + parameter_ranks=(2,), + sharded_mode=mode, + ) + + +def _gefen_muon_contract( + *, + sharded_modes: FrozenSet[str], + normuon_modes: FrozenSet[str], + non_normuon_modes: FrozenSet[str], + canonical_parameter_fqns: bool = False, + stable_shard_identity: bool = False, + explicit_process_group_codebook_scope: bool = False, + whole_parameter_owner: bool = False, + canonical_state_layouts: AbstractSet[ParameterLayout] = frozenset(), + canonical_global_same_topology: AbstractSet[ParameterLayout] = frozenset(), + canonical_global_topology_changing: AbstractSet[ParameterLayout] = frozenset(), + canonical_global_topology_change_kinds: AbstractSet[TopologyChange] = frozenset(), +) -> OptimizerContract: + canonical_state_layouts = _frozenset(canonical_state_layouts) + canonical_global_same_topology = _frozenset(canonical_global_same_topology) + canonical_global_topology_changing = _frozenset(canonical_global_topology_changing) + canonical_global_topology_change_kinds = _frozenset( + canonical_global_topology_change_kinds + ) + sharded_modes = _frozenset(sharded_modes) + normuon_modes = _frozenset(normuon_modes) + non_normuon_modes = _frozenset(non_normuon_modes) + normuon_fields = ( + StateField("normuon_v", StateScope.PARAMETER, StateGeometry.ROW, True), + StateField("normuon_step", StateScope.PARAMETER, StateGeometry.SCALAR, True), + ) + fields = _common_fields() + _base_parameter_fields() + if "distributed" in sharded_modes: + fields += ( + StateField( + "gefen_muon_distributed", + StateScope.DERIVED, + StateGeometry.OPAQUE, + True, + applicable_sharded_modes=frozenset({"distributed"}), + description="Native Parallel-Muon ownership manifest.", + ), + ) + if normuon_modes: + fields += normuon_fields + fields += _derived_fields() + + base_names = tuple(field.name for field in _base_parameter_fields()) + normuon_names = base_names + tuple(field.name for field in normuon_fields) + variants = [] + for mode in sorted(sharded_modes): + variants.append( + StateVariant( + "name_only_replicated_" + mode, + ("name",), + frozenset({ParameterLayout.REPLICATED}), + StateExtent.METADATA_ONLY, + initialized=False, + parameter_ranks=(2,), + sharded_mode=mode, + ) + ) + variants.append( + StateVariant( + "period_selected_replicated_" + mode, + ("name", "automatic_period"), + frozenset({ParameterLayout.REPLICATED}), + StateExtent.METADATA_ONLY, + initialized=False, + parameter_ranks=(2,), + sharded_mode=mode, + ) + ) + if mode == "distributed": + variants.append( + StateVariant( + "distributed_owner_name_only", + ("name",), + frozenset({_DTENSOR_LAYOUT}), + StateExtent.METADATA_ONLY, + role=ParameterStateRole.OWNER, + initialized=False, + parameter_ranks=(2,), + sharded_mode=mode, + ) + ) + variants.append( + StateVariant( + "distributed_owner_period_selected", + ("name", "automatic_period"), + frozenset({_DTENSOR_LAYOUT}), + StateExtent.METADATA_ONLY, + role=ParameterStateRole.OWNER, + initialized=False, + parameter_ranks=(2,), + sharded_mode=mode, + ) + ) + else: + variants.append( + StateVariant( + "name_only_dtensor_" + mode, + ("name",), + frozenset({_DTENSOR_LAYOUT}), + StateExtent.METADATA_ONLY, + initialized=False, + parameter_ranks=(2,), + sharded_mode=mode, + ) + ) + variants.append( + StateVariant( + "period_selected_dtensor_" + mode, + ("name", "automatic_period"), + frozenset({_DTENSOR_LAYOUT}), + StateExtent.METADATA_ONLY, + initialized=False, + parameter_ranks=(2,), + sharded_mode=mode, + ) + ) + if whole_parameter_owner: + variants.append( + StateVariant( + "name_only_whole_owner_" + mode, + ("name",), + frozenset({ParameterLayout.WHOLE_PARAMETER_OWNER}), + StateExtent.METADATA_ONLY, + role=ParameterStateRole.OWNER, + initialized=False, + parameter_ranks=(2,), + sharded_mode=mode, + ) + ) + variants.append( + StateVariant( + "period_selected_whole_owner_" + mode, + ("name", "automatic_period"), + frozenset({ParameterLayout.WHOLE_PARAMETER_OWNER}), + StateExtent.METADATA_ONLY, + role=ParameterStateRole.OWNER, + initialized=False, + parameter_ranks=(2,), + sharded_mode=mode, + ) + ) + for mode_set, field_names, prefix in ( + (non_normuon_modes, base_names, "quantized_muon"), + (normuon_modes, normuon_names, "quantized_normuon"), + ): + if not mode_set: + continue + for mode in sorted(mode_set): + variants.append( + _muon_initialized_variant( + name=prefix + "_replicated_" + mode, + field_names=field_names, + layout=ParameterLayout.REPLICATED, + extent=StateExtent.GLOBAL_PARAMETER, + mode=mode, + ) + ) + if whole_parameter_owner: + variants.append( + _muon_initialized_variant( + name=prefix + "_whole_owner_" + mode, + field_names=field_names, + layout=ParameterLayout.WHOLE_PARAMETER_OWNER, + extent=StateExtent.OWNER_PARAMETER, + mode=mode, + role=ParameterStateRole.OWNER, + ) + ) + if "approx" in mode_set: + variants.append( + _muon_initialized_variant( + name=prefix + "_local", + field_names=field_names, + layout=_DTENSOR_LAYOUT, + extent=StateExtent.LOCAL_STORAGE, + mode="approx", + ) + ) + if "exact" in mode_set: + variants.append( + _muon_initialized_variant( + name=prefix + "_global", + field_names=field_names, + layout=_DTENSOR_LAYOUT, + extent=StateExtent.GLOBAL_PARAMETER, + mode="exact", + ) + ) + if "distributed" in mode_set: + variants.append( + _muon_initialized_variant( + name=prefix + "_owner", + field_names=field_names, + layout=_DTENSOR_LAYOUT, + extent=StateExtent.OWNER_PARAMETER, + mode="distributed", + role=ParameterStateRole.OWNER, + ) + ) + if "distributed" in sharded_modes: + variants.extend( + ( + StateVariant( + "distributed_non_owner_name_only", + ("name",), + frozenset({_DTENSOR_LAYOUT}), + StateExtent.METADATA_ONLY, + role=ParameterStateRole.NON_OWNER, + initialized=False, + parameter_ranks=(2,), + sharded_mode="distributed", + ), + StateVariant( + "distributed_non_owner", + ("name", "automatic_period"), + frozenset({_DTENSOR_LAYOUT}), + StateExtent.METADATA_ONLY, + role=ParameterStateRole.NON_OWNER, + initialized=False, + parameter_ranks=(2,), + sharded_mode="distributed", + ), + ) + ) + + training = tuple( + support + for mode in sorted(sharded_modes) + for support in ( + TrainingSupport( + ParameterLayout.REPLICATED, + ProcessGroupScope.NONE, + sharded_mode=mode, + requires_complete_logical_matrix=True, + ), + TrainingSupport( + _DTENSOR_LAYOUT, + ProcessGroupScope.INFERRED_DEVICE_MESH, + mesh_dimensions=(1,), + sharded_mode=mode, + requires_complete_logical_matrix=mode in {"exact", "distributed"}, + ), + ) + ) + if whole_parameter_owner: + training += tuple( + TrainingSupport( + ParameterLayout.WHOLE_PARAMETER_OWNER, + ProcessGroupScope.ADAPTER_DEFINED, + sharded_mode=mode, + requires_complete_parameter_storage=True, + requires_complete_logical_matrix=True, + requires_post_step_parameter_sync=True, + ) + for mode in sorted(sharded_modes) + ) + checkpoints = [ + CheckpointSupport( + CheckpointTransport.NATIVE_OPTIMIZER, + frozenset({ParameterLayout.REPLICATED}), + frozenset(), + ProcessGroupScope.NONE, + required_sharded_modes=sharded_modes, + atomic_load=True, + ) + ] + if "approx" in sharded_modes and sharded_modes.issubset({"approx", "distributed"}): + checkpoints.append( + CheckpointSupport( + CheckpointTransport.PYTORCH_RANK_LOCAL, + frozenset({_DTENSOR_LAYOUT}), + frozenset(), + ProcessGroupScope.DEFAULT_WORLD, + mesh_dimensions=(1,), + required_sharded_modes=sharded_modes, + requires_collective=True, + atomic_load=True, + ) + ) + if sharded_modes == frozenset({"distributed"}): + checkpoints.append( + CheckpointSupport( + CheckpointTransport.NATIVE_OPTIMIZER, + frozenset({_DTENSOR_LAYOUT}), + frozenset({_DTENSOR_LAYOUT}), + ProcessGroupScope.INFERRED_DEVICE_MESH, + topology_change_kinds=frozenset( + {TopologyChange.WORLD_SIZE_OWNER_REDISTRIBUTION} + ), + mesh_dimensions=(1,), + required_sharded_modes=frozenset({"distributed"}), + requires_collective=True, + atomic_load=True, + ) + ) + if canonical_state_layouts: + checkpoints.append( + CheckpointSupport( + CheckpointTransport.CANONICAL_LOCAL, + canonical_state_layouts, + frozenset(), + ProcessGroupScope.NONE, + required_sharded_modes=sharded_modes, + atomic_load=True, + ) + ) + if canonical_global_same_topology or canonical_global_topology_changing: + checkpoints.append( + CheckpointSupport( + CheckpointTransport.CANONICAL_GLOBAL, + canonical_global_same_topology, + canonical_global_topology_changing, + ProcessGroupScope.ADAPTER_DEFINED, + topology_change_kinds=canonical_global_topology_change_kinds, + required_sharded_modes=sharded_modes, + requires_collective=True, + atomic_load=True, + ) + ) + return OptimizerContract( + implementation="gefen.GefenMuon", + state_layout=OptimizerStateLayout(fields, tuple(variants)), + capabilities=_negative_capabilities( + training=training, + checkpoints=tuple(checkpoints), + supported_parameter_ranks=(2,), + canonical_parameter_fqns=canonical_parameter_fqns, + stable_shard_identity=stable_shard_identity, + explicit_process_group_codebook_scope=explicit_process_group_codebook_scope, + shard_rebinding=True, + post_sharding=True, + canonical_state_io=bool( + canonical_state_layouts + or canonical_global_same_topology + or canonical_global_topology_changing + ), + ), + ) + + +def _hybrid_training( + muon: Optional[OptimizerContract], + backup: Optional[OptimizerContract], + backup_present: bool, +) -> Tuple[TrainingSupport, ...]: + """Compose hybrid training claims from the routed children's own claims. + + The hybrid routes each parameter to exactly one child, so the composite + declares the ordered union of the present children's validated claims; an + adapter must read each role's child contract together with the frozen + parameter routing rather than apply one entry to every parameter. A backup + child without a contract of its own (AdamW) contributes only the plain + replicated layout the composite actually exercises for it, because no code + validates any other layout for that child. + """ + + child_claims = [] + if muon is not None: + child_claims.append(muon.capabilities.training) + if backup is not None: + child_claims.append(backup.capabilities.training) + elif backup_present: + child_claims.append( + (TrainingSupport(ParameterLayout.REPLICATED, ProcessGroupScope.NONE),) + ) + training = [] + for claims in child_claims: + for support in claims: + if support not in training: + training.append(support) + return tuple(training) + + +def _hybrid_contract( + *, + muon: Optional[OptimizerContract], + backup: Optional[OptimizerContract], + backup_implementation: str, + canonical_parameter_fqns: bool = False, + stable_shard_identity: bool = False, + explicit_process_group_codebook_scope: bool = False, + shard_rebinding: bool = False, + post_sharding: bool = False, + canonical_global_same_topology: AbstractSet[ParameterLayout] = frozenset(), + canonical_global_topology_changing: AbstractSet[ParameterLayout] = frozenset(), + canonical_global_topology_change_kinds: AbstractSet[TopologyChange] = frozenset(), +) -> OptimizerContract: + fields = ( + StateField( + "backup_optimizer", + StateScope.OPTIMIZER_COMMON, + StateGeometry.OPAQUE, + True, + ), + ) + children = [] + if muon is not None: + children.append(OptimizerChildContract("muon", muon.implementation, muon)) + if backup_implementation: + children.append(OptimizerChildContract("backup", backup_implementation, backup)) + training = _hybrid_training(muon, backup, bool(backup_implementation)) + canonical_global_same_topology = _frozenset(canonical_global_same_topology) + canonical_global_topology_changing = _frozenset(canonical_global_topology_changing) + canonical_global_topology_change_kinds = _frozenset( + canonical_global_topology_change_kinds + ) + checkpoints = [ + CheckpointSupport( + CheckpointTransport.COMPOSITE_NATIVE, + frozenset({ParameterLayout.REPLICATED}), + frozenset(), + ProcessGroupScope.NONE, + ), + ] + if canonical_global_same_topology or canonical_global_topology_changing: + checkpoints.append( + CheckpointSupport( + CheckpointTransport.CANONICAL_GLOBAL, + canonical_global_same_topology, + canonical_global_topology_changing, + ProcessGroupScope.ADAPTER_DEFINED, + topology_change_kinds=canonical_global_topology_change_kinds, + requires_collective=True, + atomic_load=True, + ) + ) + return OptimizerContract( + implementation="gefen.GefenMuonHybrid", + state_layout=OptimizerStateLayout( + fields, + (), + composite_namespaces=("muon", "backup"), + ), + capabilities=_negative_capabilities( + training=training, + checkpoints=tuple(checkpoints), + supported_parameter_ranks=None, + canonical_parameter_fqns=canonical_parameter_fqns, + stable_shard_identity=stable_shard_identity, + explicit_process_group_codebook_scope=explicit_process_group_codebook_scope, + shard_rebinding=shard_rebinding, + post_sharding=post_sharding, + canonical_state_io=bool( + canonical_global_same_topology or canonical_global_topology_changing + ), + ), + children=tuple(children), + ) + + +__all__ = [ + "CONTRACT_SCHEMA_VERSION", + "IDENTITY_SCHEMA_VERSION", + "CheckpointSupport", + "CheckpointTransport", + "CanonicalStateProvider", + "OptimizerCapabilities", + "OptimizerChildContract", + "OptimizerContract", + "OptimizerContractProvider", + "OptimizerStateLayout", + "LogicalRegion", + "LogicalSlice", + "ParameterLayout", + "ParameterIdentity", + "ParameterStateRole", + "PortableStateProvider", + "PlacementKind", + "Precision", + "ProcessGroupScope", + "ProcessGroupIdentity", + "ShardIdentity", + "ShardPlacement", + "ShardingManifest", + "StateExtent", + "StateField", + "StateGeometry", + "StateKeyMatch", + "StateScope", + "StateVariant", + "TrainingSupport", + "TopologyChange", +] diff --git a/src/gefen/gefen.py b/src/gefen/gefen.py index 6e6b5b9..82de0ab 100644 --- a/src/gefen/gefen.py +++ b/src/gefen/gefen.py @@ -10,19 +10,48 @@ helpers it steps through; ``gefen_muon``/``hybrid`` build on it. """ +import hashlib import io import logging import math import os import warnings -from collections import OrderedDict +from collections import defaultdict, OrderedDict from itertools import chain from typing import Iterable, Optional, Tuple, Union import torch import torch.nn as nn +from gefen.canonical import ( + CANONICAL_STATE_FORMAT_VERSION, + PreparedCanonicalStateImport, + canonical_value_supported, + canonical_values_equal, + clone_canonical_value, + make_prepared_canonical_state_import, +) +from gefen.codebook import CodebookProcessGroupBinding +from gefen.contracts import ( + LogicalSlice, + OptimizerContract, + ParameterIdentity, + ParameterLayout, + ParameterStateRole, + PlacementKind, + ProcessGroupIdentity, + ShardIdentity, + ShardPlacement, + ShardingManifest, + TopologyChange, + _gefen_contract, +) from gefen.partitioning import find_period_by_block_variance +from gefen.portable_identity import ( + _parse_shard_identity, + _serialize_shard_identity, +) +from gefen.rebinding import LogicalSlotBinding, ParameterRebinding import gefen.quantization as quantization_module from gefen.kernels.automatic_vmean import ( automatic_vmean_update_cuda as _automatic_vmean_update_cuda, @@ -55,11 +84,42 @@ _RANK_LOCAL_MEMBER_KEY = "_gefen_rank_local_member" _RANK_LOCAL_FORMAT = "rank_local_dtensor_v2" _RANK_LOCAL_METADATA_VERSION = 3 - - +_CODEBOOK_SCOPE_FORMAT_VERSION = 1 +_SCOPED_NATIVE_METADATA_VERSION = 4 +_NATIVE_LOCAL_SHARDS_LEGACY_FORMAT_VERSION = 1 +_NATIVE_LOCAL_SHARDS_FORMAT_VERSION = 2 +_CANONICAL_PARAMETER_STATE_KEYS = frozenset( + { + "name", + "automatic_period", + "step", + "m_codebook", + "m_magnitude", + "vmean", + "vmean_step", + "v_row", + "v_col", + "factored_step", + "normuon_v", + "normuon_step", + } +) +_CANONICAL_DERIVED_PARAMETER_STATE_KEYS = frozenset( + { + "stepsize", + "_h_buf", + "_capt_scalars", + "_capt_consts", + "_capt_consts_key", + "_capt_stack", + "_capt_row", + "m_codebook_shape", + } +) def _rank_local_payload_key(global_rank: int) -> str: return "{}{}".format(_RANK_LOCAL_PAYLOAD_KEY_PREFIX, int(global_rank)) + # Optional explicit override for the period-search backend ("cuda_kernel" / "cpu" # / "gpu"). None (default) means resolve per call from each tensor's own device # -- see _resolve_find_period_backend. NOT used as an auto-populated cache. @@ -380,24 +440,15 @@ def gefen_automatic_fused_update( ) -def learn_gefen_exact_codebook_from_grad_periods( - *, - grad_periods, - codebook_device: torch.device, - num_codebooks: int, - force_endpoints: bool, - verbose: bool, - compute_mse_logging: bool, - use_fused_histogram: bool = FUSE_HISTOGRAM_FOR_EXACT, -) -> Optional[torch.Tensor]: +def _gefen_exact_histogram_from_grad_periods( + *, grad_periods, num_codebooks: int, use_fused_histogram: bool +) -> torch.Tensor: histogram_bins = num_codebooks * 16 bin_width = 2.0 / float(histogram_bins) # Accumulate counts as int64 and cast to float32 once at assembly: repeated # fp32 adds round for bins past 2^24 counts, so chunked accumulation could # otherwise drift from the whole-tensor form's single cast. bin_counts_cpu = torch.zeros(histogram_bins, dtype=torch.int64, device="cpu") - total_numel = 0 - prev_mode = torch.get_deterministic_debug_mode() torch.set_deterministic_debug_mode(0) try: @@ -416,7 +467,6 @@ def learn_gefen_exact_codebook_from_grad_periods( ) gefen_exact_histogram_cuda(flat_float, period, local_counts_cuda) bin_counts_cpu.add_(local_counts_cuda.cpu()) - total_numel += flat_float.numel() else: blocks = automatic_partition_view(flat_float, period) absmax = blocks.abs().amax(dim=1, keepdim=True) @@ -428,13 +478,35 @@ def learn_gefen_exact_codebook_from_grad_periods( torch.long ) bin_indices = bin_indices.clamp(0, histogram_bins - 1) - local_counts = torch.bincount( - bin_indices, minlength=histogram_bins - ) + local_counts = torch.bincount(bin_indices, minlength=histogram_bins) bin_counts_cpu.add_(local_counts.cpu()) - total_numel += normalized_flat.numel() finally: torch.set_deterministic_debug_mode(prev_mode) + return bin_counts_cpu + + +def _learn_gefen_exact_codebook_from_histogram( + *, + bin_counts_cpu: torch.Tensor, + codebook_device: torch.device, + num_codebooks: int, + force_endpoints: bool, + verbose: bool, + compute_mse_logging: bool, +) -> Optional[torch.Tensor]: + histogram_bins = num_codebooks * 16 + if ( + not torch.is_tensor(bin_counts_cpu) + or bin_counts_cpu.device.type != "cpu" + or bin_counts_cpu.dtype != torch.int64 + or tuple(bin_counts_cpu.shape) != (histogram_bins,) + ): + raise ValueError( + "Gefen exact-codebook histogram must be a CPU int64 vector with {} bins".format( + histogram_bins + ) + ) + total_numel = int(bin_counts_cpu.sum().item()) if total_numel == 0: return None @@ -483,6 +555,33 @@ def learn_gefen_exact_codebook_from_grad_periods( return codebook +def learn_gefen_exact_codebook_from_grad_periods( + *, + grad_periods, + codebook_device: torch.device, + num_codebooks: int, + force_endpoints: bool, + verbose: bool, + compute_mse_logging: bool, + use_fused_histogram: bool = FUSE_HISTOGRAM_FOR_EXACT, +) -> Optional[torch.Tensor]: + """Learn one exact-DP codebook from an unscoped local gradient stream.""" + + histogram = _gefen_exact_histogram_from_grad_periods( + grad_periods=grad_periods, + num_codebooks=num_codebooks, + use_fused_histogram=use_fused_histogram, + ) + return _learn_gefen_exact_codebook_from_histogram( + bin_counts_cpu=histogram, + codebook_device=codebook_device, + num_codebooks=num_codebooks, + force_endpoints=force_endpoints, + verbose=verbose, + compute_mse_logging=compute_mse_logging, + ) + + def _resolve_find_period_backend(grad: torch.Tensor) -> str: # An explicit FIND_PERIOD_BACKEND (set by config/tests) is an override and # wins. Otherwise resolve PER CALL from the tensor's own device. Do NOT cache @@ -701,8 +800,7 @@ def _amp_dtensor_protocol_preflight(optimizer) -> bool: params = [param for group in optimizer.param_groups for param in group["params"]] has_fp16_storage = any( - param.dtype == torch.float16 - and not getattr(param, "_is_flat_param", False) + param.dtype == torch.float16 and not getattr(param, "_is_flat_param", False) for param in params ) if not has_fp16_storage: @@ -726,8 +824,7 @@ def _amp_dtensor_protocol_preflight(optimizer) -> bool: str(mesh.device_type), tuple(int(item) for item in mesh.shape), tuple( - int(item) - for item in mesh.mesh.detach().cpu().reshape(-1).tolist() + int(item) for item in mesh.mesh.detach().cpu().reshape(-1).tolist() ), tuple(str(group.group_name) for group in process_groups), ) @@ -742,9 +839,7 @@ def _amp_dtensor_protocol_preflight(optimizer) -> bool: name = ( names[index] if isinstance(names, (list, tuple)) and index < len(names) - else getattr(optimizer, "_param_name", lambda _: "parameter")( - param - ) + else getattr(optimizer, "_param_name", lambda _: "parameter")(param) ) entry["items"].append((str(name), param, param.grad is not None)) @@ -880,9 +975,7 @@ def _assert_optimizer_gradients_structurally_valid( else getattr(optimizer, "_param_name", lambda _: "parameter")(param) ) if require_2d_params and param.ndim != 2: - error_factory = getattr( - optimizer, "_step_non_2d_parameter_error", None - ) + error_factory = getattr(optimizer, "_step_non_2d_parameter_error", None) if callable(error_factory): raise ValueError(error_factory(param)) raise ValueError( @@ -899,19 +992,17 @@ def _assert_optimizer_gradients_structurally_valid( raise RuntimeError( "Gefen does not support sparse gradients or other " "non-strided layouts; parameter {!r} has gradient layout " - "{}.".format( - str(name), layout - ) + "{}.".format(str(name), layout) ) if torch.is_complex(grad): raise RuntimeError( - "Gefen optimizers do not support complex gradients, but " - "parameter {!r} has dtype {}.".format(str(name), grad.dtype) + "Gefen optimizers do not support complex gradients, but parameter {!r} has dtype {}.".format( + str(name), grad.dtype + ) ) if tuple(grad.shape) != tuple(param.shape): raise RuntimeError( - "Gefen gradient shape {} for parameter {!r} does not match " - "parameter shape {}.".format( + "Gefen gradient shape {} for parameter {!r} does not match parameter shape {}.".format( tuple(grad.shape), str(name), tuple(param.shape) ) ) @@ -1007,6 +1098,18 @@ class Gefen(torch.optim.Optimizer): learning and period prediction. """ + # Pure rebuildable memoization for the layout-forensics guards. These live + # in ``__slots__`` rather than the instance ``__dict__`` so that staged + # ``__dict__`` copies, staged commits, and fail-before-mutation attribute + # snapshots treat them as what they are — caches, never optimizer state. + # Every accessor reads them through ``getattr`` with a cold default, so an + # object whose slots were never populated simply revalidates fully. + __slots__ = ( + "_gefen_layout_version", + "_gefen_layout_forensics_verdict", + "_gefen_manifest_forensics_cache", + ) + def __init__( self, params: Iterable[Union[nn.Parameter, Tuple[str, nn.Parameter]]], @@ -1037,8 +1140,7 @@ def __init__( ) if fused and not torch.cuda.is_available(): warnings.warn( - "Gefen optimizer got fused=True, but CUDA is not available. " - "Changing fused to False.", + "Gefen optimizer got fused=True, but CUDA is not available. Changing fused to False.", UserWarning, stacklevel=2, ) @@ -1058,11 +1160,11 @@ def __init__( # partitioning.memory_safe_fallback_period. self._force_1d_period_one = force_1d_period_one self._force_2d_period_one = force_2d_period_one - self._period_one_substrings = tuple( - s.lower() for s in period_one_substrings - ) + self._period_one_substrings = tuple(s.lower() for s in period_one_substrings) self._factored_v_2d = factored_v_2d self._deterministic = deterministic + if type(codebook_refresh_every) is not int: + raise TypeError("codebook_refresh_every must be an integer") if codebook_refresh_every < 0: raise ValueError( "codebook_refresh_every must be >= 0 but is: {}".format( @@ -1114,6 +1216,12 @@ def __init__( # cudagraph-trees recording rejects whenever the codebook input gets # copied into the graph pool). Cleared with _gefen_codebook_by_device. self._gefen_codebook_lut_by_device = {} + # An explicit learned-codebook scope is installed only as part of the + # same full post_sharding transaction that publishes canonical shard + # identities. It is live runtime configuration, never inferred from a + # default process group or from tensor storage. + self._gefen_codebook_process_group = None + self._gefen_codebook_scope_validated = False # Capturable stochastic rounding: ONE 0-dim int64 device tensor per # device holding the per-step rounding seed (a device-side mirror of # _gefen_global_step). Created lazily by _sr_seed_on and advanced on @@ -1142,7 +1250,6 @@ def __init__( # one device counter per parameter device and advance it in the captured # step tail; state_dict synchronizes the host mirror before serializing. self._gefen_global_step_by_device = {} - defaults = dict( lr=lr, beta1=betas[0], @@ -1151,6 +1258,28 @@ def __init__( weight_decay=weight_decay, ) self._param_names = {} + # Exact canonical identity remains separate from compatibility + # ``param_names``. A full post_sharding transaction publishes these + # registries only after validating every optimizer slot and manifest. + self._gefen_shard_bindings = {} + self._gefen_local_shard_bindings = () + self._gefen_sharding_manifest = None + self._gefen_post_sharding_finalized = False + self._gefen_finalized_slots = () + self._gefen_logical_slots = () + # Layout-forensics fast path. After one full structural pass succeeds, + # the exact container identities it validated are remembered together + # with a version counter that every legitimate mutating API bumps + # (post_sharding and staged checkpoint load commits). Per-step guards + # accept only that unchanged token set; boundary operations (checkpoint + # prepare/commit, rebinding, collective codebook initialize/refresh, + # and contract readiness) always rerun the complete forensic rebuild. + self._gefen_layout_version = 0 + self._gefen_layout_forensics_verdict = None + # (manifest, frozenset(manifest.shards), sha256 digest) computed once + # per finalized manifest object instead of rehashing every identity on + # every scoped step header. + self._gefen_manifest_forensics_cache = None # ``set_optimizer_state_dict(flatten_optimizer_state_dict=True)`` uses # the *live* optimizer state/group keys as its unflattening schema before # it calls our loader. Publish the private rank-local transport keys only @@ -1166,177 +1295,1818 @@ def _step_supports_amp_scaling(self) -> bool: # FP32-master AMP. Native handling is needed for active local FP16 # gradients, or statically for distributed optimizers that combine any # true-FP16 storage with DTensors after a collective presence preflight. + # An explicit codebook scope is also a topology property: whole owners + # and empty nonowners must enter the same GradScaler protocol before + # Gefen can compare found_inf/grad_scale collectively inside step(). + if self._gefen_codebook_process_group is not None: + return True return _amp_native_scaling_required(self) - @staticmethod - def _normalize_param_groups(params): - if isinstance(params, torch.Tensor): - param_groups = [params] - else: - param_groups = list(params) - if len(param_groups) == 0: - raise ValueError("optimizer got an empty parameter list") - if isinstance(param_groups[0], dict): - return param_groups - return [{"params": param_groups, "_gefen_implicit_group": True}] + def optimizer_contract(self) -> OptimizerContract: + """Return the immutable state-layout and integration capability contract.""" - @staticmethod - def _validate_group_options(lr, betas, eps, weight_decay): - if torch.is_tensor(lr) and lr.numel() != 1: - raise ValueError( - "lr as a Tensor must be 1-element, got numel={}".format(lr.numel()) - ) - if not 0.0 <= lr: - raise ValueError("Invalid learning rate: {}".format(lr)) - elif not 0.0 <= betas[0] < 1.0: - raise ValueError( - "Invalid beta parameter at index 0: {}".format(betas[0]) + identity_ready = self._canonical_identity_ready() + canonical_state_layouts = self._canonical_state_layouts() + try: + from gefen.portable_runtime import _portable_runtime_layouts + + canonical_global_same_topology = _portable_runtime_layouts(self) + except Exception: + canonical_global_same_topology = frozenset() + has_factored_matrix = ( + canonical_global_same_topology + and self._factored_v_2d + and any( + len(slot.shard.parameter.global_shape) == 2 + for slot in self._gefen_logical_slots ) - elif not 0.0 <= betas[1] < 1.0: - raise ValueError( - "Invalid beta parameter at index 1: {}".format(betas[1]) + ) + if canonical_global_same_topology and not has_factored_matrix: + canonical_global_topology_changing = frozenset( + { + ParameterLayout.REPLICATED, + ParameterLayout.FLATTENED_ELEMENT_SHARD, + } ) - elif not 0.0 <= weight_decay: - raise ValueError("Invalid weight_decay value: {}".format(weight_decay)) - elif not math.isfinite(eps) or eps <= 0.0: - raise ValueError( - "Invalid epsilon value: {}; eps must be finite and > 0".format(eps) + canonical_global_topology_change_kinds = frozenset( + {TopologyChange.PLACEMENT_RESHARD} ) - - @staticmethod - def _iter_params_with_names(group_params): - if isinstance(group_params, torch.Tensor): - group_params = (group_params,) - for param_or_named_param in group_params: - if isinstance(param_or_named_param, tuple): - param_name, param = param_or_named_param - else: - param_name = None - param = param_or_named_param - if not isinstance(param, torch.Tensor): - raise TypeError( - "Gefen can only optimize Tensors, but one of the params is {}".format( - type(param).__name__, - ) - ) - if torch.is_complex(param): - raise ValueError( - "Gefen does not support complex parameters, but got a " - "parameter with dtype {}".format(param.dtype) + else: + canonical_global_topology_changing = frozenset() + canonical_global_topology_change_kinds = frozenset() + return _gefen_contract( + factored_v_2d=self._factored_v_2d, + canonical_parameter_fqns=identity_ready, + stable_shard_identity=identity_ready, + explicit_process_group_codebook_scope=True, + canonical_state_layouts=canonical_state_layouts, + canonical_global_same_topology=canonical_global_same_topology, + canonical_global_topology_changing=canonical_global_topology_changing, + canonical_global_topology_change_kinds=canonical_global_topology_change_kinds, + native_flattened_checkpoint=( + self._codebook_scope_ready() + and any( + shard.layout is ParameterLayout.FLATTENED_ELEMENT_SHARD + for _, shard in self._gefen_local_shard_bindings ) - yield param_name, param + ), + ) - def _param_name(self, param) -> str: - state = self.state.get(param) - if state is not None and "name" in state: - return state["name"] - return self._param_names.get(param, "none") + def _canonical_state_variant_layout(self): + return _gefen_contract( + factored_v_2d=self._factored_v_2d, + ).state_layout - def _iter_group_params_with_names(self, group): - names = group.get("param_names") - for idx, param in enumerate(group["params"]): - if names is not None and idx < len(names): - yield names[idx], param - else: - yield self._param_name(param), param + def _canonical_state_layout_supported(self, layout) -> bool: + return layout in { + ParameterLayout.REPLICATED, + ParameterLayout.FLATTENED_ELEMENT_SHARD, + } @staticmethod - def _unique_name(base: str, existing_names) -> str: - name = str(base).lower() - if name not in existing_names: - return name - suffix = 1 - while "{}_{}".format(name, suffix) in existing_names: - suffix += 1 - return "{}_{}".format(name, suffix) + def _canonical_group_options_value(group): + return { + key: value + for key, value in group.items() + if key + not in { + "params", + "param_names", + "name", + "_gefen_checkpoint_metadata", + } + } - def _sync_param_names_to_state(self) -> None: - self._param_names = {} + def _canonical_state_layouts(self): + if not self._canonical_identity_ready(): + return frozenset() + if self._stochastic_round: + return frozenset() + if any( + parameter is None + or not self._canonical_state_layout_supported(shard.layout) + for parameter, shard in self._gefen_local_shard_bindings + ): + return frozenset() + if any( + not canonical_value_supported( + self._canonical_group_options_value(group), + finite_tensors=True, + ) + for group in self.param_groups + ): + return frozenset() + if not canonical_value_supported( + self._canonical_policy(), + finite_tensors=True, + ): + return frozenset() + if not canonical_value_supported( + self._gefen_codebook, + finite_tensors=True, + ): + return frozenset() + try: + global_step = self._canonical_common_global_step() + if type(global_step) is not int or global_step < 0: + return frozenset() + self._validate_loaded_native_state() + except Exception: + return frozenset() + state_layout = self._canonical_state_variant_layout() for group in self.param_groups: - params = group["params"] - names = list(group.get("param_names", ())) - if len(names) != len(params): - names = [self._param_name(p) for p in params] - normalized_names = [] - for param, name in zip(params, names): - name = str(name).lower() - normalized_names.append(name) - self._param_names[param] = name - self.state[param]["name"] = name - group["param_names"] = normalized_names + group_options = self._canonical_group_options_value(group) + for parameter in group["params"]: + state = self.state.get(parameter, {}) + if any( + key not in _CANONICAL_PARAMETER_STATE_KEYS + and key not in _CANONICAL_DERIVED_PARAMETER_STATE_KEYS + for key in state + ): + return frozenset() + if any( + not canonical_value_supported( + value, + finite_tensors=True, + ) + for key, value in state.items() + if key in _CANONICAL_PARAMETER_STATE_KEYS + ): + return frozenset() + canonical_state = { + key: value + for key, value in state.items() + if key in _CANONICAL_PARAMETER_STATE_KEYS and key != "name" + } + shard = self._gefen_shard_bindings[parameter] + try: + self._validate_canonical_parameter_semantics( + parameter, + shard, + group_options, + canonical_state, + state_layout, + global_step, + ) + except (TypeError, ValueError, RuntimeError): + return frozenset() + return frozenset(shard.layout for _, shard in self._gefen_local_shard_bindings) - def add_param_group(self, param_group): - """Add a param group while preserving the caller-visible group boundary. + def _codebook_scope_ready(self) -> bool: + return ( + self._gefen_codebook_process_group is not None + and self._canonical_identity_ready() + ) - Accepts the same group dict shape as the constructor (``params`` as - tensors or ``(name, tensor)`` tuples; optional ``lr`` / ``betas`` (or - ``beta1``/``beta2``) / ``eps`` / ``weight_decay``, defaulting to the - constructor values). The public group is preserved; each parameter's - stable lowercase name is stored in its per-param state and in the - group's ``param_names`` list for introspection. - """ - if not isinstance(param_group, dict): - raise TypeError( - "param_group must be a dict, got {}".format( - type(param_group).__name__ - ) - ) - if "params" not in param_group: - raise ValueError("Gefen parameter group is missing the 'params' key.") - defaults = self.defaults - group_lr = param_group.get("lr", defaults["lr"]) - if "betas" in param_group: - group_betas = param_group["betas"] - else: - group_betas = ( - param_group.get("beta1", defaults["beta1"]), - param_group.get("beta2", defaults["beta2"]), + def _canonical_identity_ready(self) -> bool: + if ( + not self._gefen_post_sharding_finalized + or self._gefen_sharding_manifest is None + ): + return False + # Contract readiness is an honest external claim, so it never trusts + # the per-step fast-path verdict. + return self._finalized_binding_layout_matches(full=True) + + def _invalidate_layout_forensics_caches(self) -> None: + self._gefen_layout_version = getattr(self, "_gefen_layout_version", 0) + 1 + self._gefen_layout_forensics_verdict = None + + @staticmethod + def _forensics_tokens_match(cached, live) -> bool: + # Containers, tensors, and strings compare by object identity — the + # cached tuple holds strong references, so an unchanged attribute is + # the same object. Plain ints (version counter, lengths) compare by + # value because CPython interns only small integers. + return len(cached) == len(live) and all( + cached_item is live_item + or ( + type(cached_item) is int + and type(live_item) is int + and cached_item == live_item ) - group_eps = param_group.get("eps", defaults["eps"]) - group_weight_decay = param_group.get("weight_decay", defaults["weight_decay"]) - self._validate_group_options( - group_lr, group_betas, group_eps, group_weight_decay - ) - named_params = list(self._iter_params_with_names(param_group["params"])) - if len(named_params) == 0: - raise ValueError("optimizer got an empty parameter list") - # Keys beyond the ones Gefen manages (e.g. a subclass's extras) are - # carried onto the preserved group, mirroring torch's behavior of - # preserving unknown group keys. - consumed = ( - "params", - "name", - "lr", - "betas", - "beta1", - "beta2", - "eps", - "weight_decay", - "param_names", - "_gefen_implicit_group", + for cached_item, live_item in zip(cached, live) ) - extra = {k: v for k, v in param_group.items() if k not in consumed} - existing_names = set(self._param_names.values()) - group_index = len(self.param_groups) - # Build (and fully validate) the preserved group up front so a failure - # partway through -- a complex/non-tensor param from - # _iter_params_with_names, or a duplicate param -- leaves - # self.param_groups untouched. add_param_group is all-or-nothing: only - # once the whole group is materialized do we register any of it. - existing_params = set() - for g in self.param_groups: - existing_params.update(g["params"]) - batch_params = set() - group_name = param_group.get("name") - implicit_group = bool(param_group.get("_gefen_implicit_group", False)) - params = [] - param_names = [] - for param_index, (param_name, param) in enumerate(named_params): - # Reject duplicates before registering any group, so the base class's - # own (post-registration) duplicate check can't leave us partial. - if param in existing_params or param in batch_params: - raise ValueError( - "some parameters appear in more than one parameter group" + + def _layout_forensics_fast_tokens(self): + # O(local params) identity snapshot of everything user code can reach + # between two guard calls: the finalized registries by object identity + # (they are replaced, never mutated, by legitimate APIs) plus every + # live group container, parameter, and compatibility name element, so + # in-place mutation of the public groups — e.g. a closure swapping one + # ``group["params"]`` slot — still falls back to the full forensic + # pass before any state mutation. + live = [ + getattr(self, "_gefen_layout_version", 0), + self._gefen_post_sharding_finalized, + self._gefen_sharding_manifest, + self._gefen_logical_slots, + self._gefen_finalized_slots, + self._gefen_local_shard_bindings, + self._gefen_shard_bindings, + len(self._gefen_shard_bindings), + self._param_names, + len(self._param_names), + self.param_groups, + len(self.param_groups), + self.state, + ] + live.extend(self._param_names.keys()) + live.extend(self._param_names.values()) + for group in self.param_groups: + params = group.get("params") if type(group) is dict else None + names = group.get("param_names") if type(group) is dict else None + live.append(group) + live.append(params) + live.append(names) + if isinstance(params, (list, tuple)): + live.append(len(params)) + live.extend(params) + for parameter in params: + parameter_state = self.state.get(parameter) + live.append(type(parameter_state)) + live.append( + parameter_state.get("name") + if type(parameter_state) is dict + else None + ) + if isinstance(names, (list, tuple)): + live.append(len(names)) + live.extend(names) + return tuple(live) + + def _manifest_layout_forensics(self, *, refresh: bool = False): + """Return the cached ``(shard set, digest)`` for the live manifest. + + Both values are pure functions of one immutable ``ShardingManifest``, + so they are computed once per manifest object (at post_sharding + finalization) instead of rehashing every ``ShardIdentity`` — including + its world-size ``ordered_members`` tuple — on every step. Full + forensic passes recompute with ``refresh=True`` so boundary checks + never trust a stale cache. + """ + + manifest = self._gefen_sharding_manifest + cache = getattr(self, "_gefen_manifest_forensics_cache", None) + if not refresh and cache is not None and cache[0] is manifest: + return cache[1], cache[2] + shards = frozenset(manifest.shards) + digest = self._compute_codebook_manifest_fingerprint(manifest) + self._gefen_manifest_forensics_cache = (manifest, shards, digest) + return shards, digest + + def _finalized_binding_layout_matches(self, *, full: bool = False) -> bool: + try: + verdict = getattr(self, "_gefen_layout_forensics_verdict", None) + if not full and verdict is not None: + if self._forensics_tokens_match( + verdict, self._layout_forensics_fast_tokens() + ): + return True + self._gefen_layout_forensics_verdict = None + matches = self._finalized_binding_layout_matches_full() + if matches: + self._gefen_layout_forensics_verdict = ( + self._layout_forensics_fast_tokens() + ) + return matches + except AttributeError: + return False + + def _finalized_binding_layout_matches_full(self) -> bool: + try: + if ( + type(self._gefen_logical_slots) is not tuple + or not self._gefen_logical_slots + or type(self._gefen_finalized_slots) is not tuple + or type(self._gefen_local_shard_bindings) is not tuple + or type(self._gefen_shard_bindings) is not dict + or type(self._param_names) is not dict + or not isinstance(self._gefen_sharding_manifest, ShardingManifest) + or len(self.param_groups) != len(self._gefen_finalized_slots) + ): + return False + + logical_groups = [[] for _ in self.param_groups] + logical_fqns = set() + previous_position = None + for logical_slot in self._gefen_logical_slots: + if ( + type(logical_slot) is not LogicalSlotBinding + or type(logical_slot.group_index) is not int + or type(logical_slot.original_slot_index) is not int + or type(logical_slot.compatibility_name) is not str + or logical_slot.compatibility_name + != logical_slot.compatibility_name.lower() + or not isinstance(logical_slot.shard, ShardIdentity) + ): + return False + position = ( + logical_slot.group_index, + logical_slot.original_slot_index, + ) + if ( + logical_slot.group_index < 0 + or logical_slot.group_index >= len(logical_groups) + or logical_slot.original_slot_index + != len(logical_groups[logical_slot.group_index]) + ): + return False + if previous_position is None: + if position != (0, 0): + return False + elif logical_slot.group_index == previous_position[0]: + if logical_slot.original_slot_index != previous_position[1] + 1: + return False + elif position != (previous_position[0] + 1, 0): + return False + fqn = logical_slot.shard.parameter.fqn + if fqn in logical_fqns: + return False + logical_fqns.add(fqn) + logical_groups[logical_slot.group_index].append(logical_slot) + previous_position = position + + if any(not group for group in logical_groups): + return False + + manifest_shards = self._manifest_layout_forensics(refresh=True)[0] + if { + shard.parameter.fqn for shard in manifest_shards + } != logical_fqns or any( + logical_slot.shard not in manifest_shards + for logical_slot in self._gefen_logical_slots + ): + return False + + local_bindings = {} + previous_sort_key = None + for item in self._gefen_local_shard_bindings: + if type(item) is not tuple or len(item) != 2: + return False + parameter, shard = item + if parameter is not None and not isinstance(parameter, torch.Tensor): + return False + if not isinstance(shard, ShardIdentity): + return False + if previous_sort_key is not None and shard.sort_key < previous_sort_key: + return False + previous_sort_key = shard.sort_key + fqn = shard.parameter.fqn + if fqn in local_bindings: + return False + local_bindings[fqn] = (parameter, shard) + if set(local_bindings) != logical_fqns: + return False + + expected_live_groups = [] + expected_name_groups = [] + expected_live_bindings = [] + for logical_group in logical_groups: + live_group = [] + name_group = [] + for logical_slot in logical_group: + parameter, shard = local_bindings[logical_slot.shard.parameter.fqn] + if shard != logical_slot.shard: + return False + pruned = ( + shard.layout is ParameterLayout.WHOLE_PARAMETER_OWNER + and shard.local_member != shard.owner + ) + if pruned: + if parameter is not None: + return False + continue + if parameter is None: + return False + live_group.append(parameter) + name_group.append(logical_slot.compatibility_name) + expected_live_bindings.append((parameter, shard)) + expected_live_groups.append(tuple(live_group)) + expected_name_groups.append(tuple(name_group)) + + if len(self._gefen_shard_bindings) != len(expected_live_bindings): + return False + expected_live_ids = { + id(parameter) for parameter, _ in expected_live_bindings + } + if len(expected_live_ids) != len(expected_live_bindings): + return False + if { + id(parameter) for parameter in self._gefen_shard_bindings + } != expected_live_ids: + return False + for parameter, shard in expected_live_bindings: + if self._gefen_shard_bindings.get(parameter) != shard: + return False + + if len(self._param_names) != len(expected_live_bindings): + return False + if {id(parameter) for parameter in self._param_names} != expected_live_ids: + return False + for group_index, (group, expected_params, expected_names) in enumerate( + zip(self.param_groups, expected_live_groups, expected_name_groups) + ): + params = group.get("params") + names = group.get("param_names") + finalized = self._gefen_finalized_slots[group_index] + if ( + not isinstance(params, (list, tuple)) + or not isinstance(names, (list, tuple)) + or type(finalized) is not tuple + or len(params) != len(expected_params) + or len(names) != len(expected_names) + or len(finalized) != len(expected_params) + ): + return False + if any( + live is not expected + for live, expected in zip(params, expected_params) + ) or any( + bound is not expected + for bound, expected in zip(finalized, expected_params) + ): + return False + if tuple(names) != expected_names: + return False + for parameter, expected_name in zip(expected_params, expected_names): + if self._param_names.get(parameter) != expected_name: + return False + parameter_state = self.state.get(parameter) + if ( + type(parameter_state) is not dict + or parameter_state.get("name") != expected_name + ): + return False + return True + except ( + AttributeError, + IndexError, + KeyError, + RuntimeError, + TypeError, + ValueError, + ): + return False + + def _assert_finalized_binding_layout(self, *, full: bool = False) -> None: + if ( + self._gefen_post_sharding_finalized + and not self._finalized_binding_layout_matches(full=full) + ): + raise RuntimeError( + "Gefen finalized parameter layout changed outside post_sharding" + ) + + def parameter_identity(self, parameter) -> ParameterIdentity: + """Return the exact canonical identity bound to one live parameter.""" + + return self.shard_identity(parameter).parameter + + def shard_identity(self, parameter) -> ShardIdentity: + """Return the stable shard identity bound to one live parameter.""" + + self._assert_finalized_binding_layout() + try: + return self._gefen_shard_bindings[parameter] + except KeyError: + raise KeyError("parameter has no finalized Gefen shard identity") from None + + def shard_bindings(self): + """Return local tensor/identity pairs in canonical structural order.""" + + self._assert_finalized_binding_layout() + return self._gefen_local_shard_bindings + + def sharding_manifest(self): + """Return the finalized global identity manifest, or ``None``.""" + + self._assert_finalized_binding_layout() + return self._gefen_sharding_manifest + + def codebook_process_group_binding(self): + """Return the finalized explicit learned-codebook binding, or ``None``.""" + + self._assert_finalized_binding_layout() + return self._gefen_codebook_process_group + + def _serialized_codebook_scope(self): + binding = self._gefen_codebook_process_group + if binding is None: + return None + return { + "format_version": _CODEBOOK_SCOPE_FORMAT_VERSION, + "semantic_name": binding.identity.semantic_name, + "ordered_members": list(binding.identity.ordered_members), + "refresh_every": self._codebook_refresh_every, + } + + @staticmethod + def _normalize_serialized_codebook_scope(value): + if value is None: + return None + if not isinstance(value, dict) or set(value) != { + "format_version", + "semantic_name", + "ordered_members", + "refresh_every", + }: + raise ValueError("Gefen checkpoint codebook scope has an invalid schema") + if ( + type(value["format_version"]) is not int + or value["format_version"] != _CODEBOOK_SCOPE_FORMAT_VERSION + ): + raise ValueError( + "Gefen checkpoint codebook scope has an unsupported format version" + ) + if not isinstance(value["ordered_members"], list): + raise ValueError( + "Gefen checkpoint codebook scope ordered_members must be a list" + ) + from gefen.contracts import ProcessGroupIdentity + + try: + identity = ProcessGroupIdentity( + value["semantic_name"], tuple(value["ordered_members"]) + ) + except (TypeError, ValueError) as exc: + raise ValueError( + "Gefen checkpoint codebook scope identity is invalid" + ) from exc + refresh_every = value["refresh_every"] + if type(refresh_every) is not int or refresh_every < 0: + raise ValueError( + "Gefen checkpoint codebook scope refresh_every must be a nonnegative integer" + ) + return { + "format_version": _CODEBOOK_SCOPE_FORMAT_VERSION, + "semantic_name": identity.semantic_name, + "ordered_members": list(identity.ordered_members), + "refresh_every": refresh_every, + } + + @staticmethod + def _serialized_native_local_shard(shard): + return { + "fqn": shard.parameter.fqn, + "global_shape": list(shard.parameter.global_shape), + "layout": shard.layout.value, + "flat_offset": shard.logical_slice.flat_offset, + "length": shard.logical_slice.length, + "process_group": { + "semantic_name": shard.process_group.semantic_name, + "ordered_members": list(shard.process_group.ordered_members), + }, + "local_member": shard.local_member, + "owner": shard.owner, + "placements": [ + { + "mesh_axis": placement.mesh_axis, + "kind": placement.kind.value, + "coordinate": placement.coordinate, + "parts": placement.parts, + "parameter_dimension": placement.parameter_dimension, + } + for placement in shard.placements + ], + } + + @staticmethod + def _serialized_canonical_shard(shard): + process_group = None + if shard.process_group is not None: + process_group = { + "schema_version": shard.process_group.schema_version, + "semantic_name": shard.process_group.semantic_name, + "ordered_members": list(shard.process_group.ordered_members), + } + return { + "schema_version": shard.schema_version, + "parameter": { + "schema_version": shard.parameter.schema_version, + "fqn": shard.parameter.fqn, + "global_shape": list(shard.parameter.global_shape), + }, + "layout": shard.layout.value, + "logical_slice": { + "flat_offset": shard.logical_slice.flat_offset, + "length": shard.logical_slice.length, + }, + "process_group": process_group, + "local_member": shard.local_member, + "owner": shard.owner, + "placements": [ + { + "mesh_axis": placement.mesh_axis, + "kind": placement.kind.value, + "coordinate": placement.coordinate, + "parts": placement.parts, + "parameter_dimension": placement.parameter_dimension, + } + for placement in shard.placements + ], + } + + @classmethod + def _normalize_serialized_canonical_shard(cls, record): + if not isinstance(record, dict) or set(record) != { + "schema_version", + "parameter", + "layout", + "logical_slice", + "process_group", + "local_member", + "owner", + "placements", + }: + raise ValueError("Gefen canonical shard has an invalid schema") + parameter_record = record["parameter"] + if not isinstance(parameter_record, dict) or set(parameter_record) != { + "schema_version", + "fqn", + "global_shape", + }: + raise ValueError("Gefen canonical parameter identity is invalid") + if not isinstance(parameter_record["global_shape"], list): + raise ValueError("Gefen canonical global_shape must be a list") + slice_record = record["logical_slice"] + if not isinstance(slice_record, dict) or set(slice_record) != { + "flat_offset", + "length", + }: + raise ValueError("Gefen canonical logical slice is invalid") + group_record = record["process_group"] + if group_record is not None and ( + not isinstance(group_record, dict) + or set(group_record) + != {"schema_version", "semantic_name", "ordered_members"} + or not isinstance(group_record["ordered_members"], list) + ): + raise ValueError("Gefen canonical process-group identity is invalid") + if not isinstance(record["placements"], list): + raise ValueError("Gefen canonical placements must be a list") + try: + parameter = ParameterIdentity( + parameter_record["fqn"], + tuple(parameter_record["global_shape"]), + schema_version=parameter_record["schema_version"], + ) + process_group = None + if group_record is not None: + process_group = ProcessGroupIdentity( + group_record["semantic_name"], + tuple(group_record["ordered_members"]), + schema_version=group_record["schema_version"], + ) + placements = [] + for placement_record in record["placements"]: + if not isinstance(placement_record, dict) or set(placement_record) != { + "mesh_axis", + "kind", + "coordinate", + "parts", + "parameter_dimension", + }: + raise ValueError("invalid placement schema") + placements.append( + ShardPlacement( + placement_record["mesh_axis"], + PlacementKind(placement_record["kind"]), + placement_record["coordinate"], + placement_record["parts"], + placement_record["parameter_dimension"], + ) + ) + shard = ShardIdentity( + parameter, + ParameterLayout(record["layout"]), + LogicalSlice(slice_record["flat_offset"], slice_record["length"]), + process_group=process_group, + local_member=record["local_member"], + owner=record["owner"], + placements=tuple(placements), + schema_version=record["schema_version"], + ) + except (TypeError, ValueError) as exc: + raise ValueError("Gefen canonical shard identity is invalid") from exc + return cls._serialized_canonical_shard(shard) + + def _serialized_native_local_shards_v1(self): + if self._gefen_codebook_process_group is None: + return None + if not any( + shard.layout is not ParameterLayout.REPLICATED + for _, shard in self._gefen_local_shard_bindings + ): + return None + + # Native Optimizer.load_state_dict maps per-parameter state by + # parameter-group and slot position, not by canonical FQN. Bind every + # serialized live slot to its shard identity (including replicated + # slots in a mixed optimizer), and retain pruned whole-owner records + # separately. A canonical shard set alone would not distinguish A/B + # from B/A when equal-shaped parameters exchange positions. + param_groups = [] + for group in self.param_groups: + param_groups.append( + [ + self._serialized_native_local_shard( + self._gefen_shard_bindings[parameter] + ) + for parameter in group["params"] + ] + ) + pruned_shards = [ + self._serialized_native_local_shard(shard) + for parameter, shard in self._gefen_local_shard_bindings + if parameter is None + ] + return { + "format_version": _NATIVE_LOCAL_SHARDS_LEGACY_FORMAT_VERSION, + "param_groups": param_groups, + "pruned_shards": pruned_shards, + } + + def _serialized_native_local_shards(self): + if self._gefen_codebook_process_group is None: + return None + if not any( + shard.layout is not ParameterLayout.REPLICATED + for _, shard in self._gefen_local_shard_bindings + ): + return None + return { + "format_version": _NATIVE_LOCAL_SHARDS_FORMAT_VERSION, + "logical_slots": [ + { + "group_index": slot.group_index, + "original_slot_index": slot.original_slot_index, + "compatibility_name": slot.compatibility_name, + "shard": _serialize_shard_identity(slot.shard), + } + for slot in self._gefen_logical_slots + ], + } + + @classmethod + def _normalize_serialized_native_local_shard(cls, record): + expected = { + "fqn", + "global_shape", + "layout", + "flat_offset", + "length", + "process_group", + "local_member", + "owner", + "placements", + } + if not isinstance(record, dict) or set(record) != expected: + raise ValueError("Gefen native local-shard metadata has an invalid schema") + group_record = record["process_group"] + if not isinstance(group_record, dict) or set(group_record) != { + "semantic_name", + "ordered_members", + }: + raise ValueError( + "Gefen native local-shard process-group metadata is invalid" + ) + if not isinstance(group_record["ordered_members"], list): + raise ValueError("Gefen native local-shard ordered_members must be a list") + if not isinstance(record["global_shape"], list) or not isinstance( + record["placements"], list + ): + raise ValueError( + "Gefen native local-shard shapes and placements must be lists" + ) + try: + parameter = ParameterIdentity(record["fqn"], tuple(record["global_shape"])) + group = ProcessGroupIdentity( + group_record["semantic_name"], + tuple(group_record["ordered_members"]), + ) + placements = [] + for placement_record in record["placements"]: + if not isinstance(placement_record, dict) or set(placement_record) != { + "mesh_axis", + "kind", + "coordinate", + "parts", + "parameter_dimension", + }: + raise ValueError("invalid placement schema") + placements.append( + ShardPlacement( + placement_record["mesh_axis"], + PlacementKind(placement_record["kind"]), + placement_record["coordinate"], + placement_record["parts"], + placement_record["parameter_dimension"], + ) + ) + shard = ShardIdentity( + parameter, + ParameterLayout(record["layout"]), + LogicalSlice(record["flat_offset"], record["length"]), + process_group=group, + local_member=record["local_member"], + owner=record["owner"], + placements=tuple(placements), + ) + except (TypeError, ValueError) as exc: + raise ValueError("Gefen native local-shard metadata is invalid") from exc + if shard.layout not in { + ParameterLayout.REPLICATED, + ParameterLayout.FLATTENED_ELEMENT_SHARD, + ParameterLayout.WHOLE_PARAMETER_OWNER, + }: + raise ValueError( + "Gefen native local-shard metadata has an unsupported layout" + ) + return cls._serialized_native_local_shard(shard) + + @classmethod + def _normalize_serialized_native_local_shards_v1(cls, value): + if not isinstance(value, dict) or set(value) != { + "format_version", + "param_groups", + "pruned_shards", + }: + raise ValueError("Gefen native local-shard metadata has an invalid schema") + format_version = value["format_version"] + if ( + type(format_version) is not int + or format_version != _NATIVE_LOCAL_SHARDS_LEGACY_FORMAT_VERSION + ): + raise ValueError( + "Unsupported Gefen native local-shard format_version: {}".format( + format_version + ) + ) + param_groups = value["param_groups"] + pruned_shards = value["pruned_shards"] + if not isinstance(param_groups, list) or not isinstance(pruned_shards, list): + raise ValueError( + "Gefen native local-shard parameter groups and pruned shards must be lists" + ) + normalized_groups = [] + for group in param_groups: + if not isinstance(group, list): + raise ValueError( + "Gefen native local-shard parameter groups must contain lists" + ) + normalized_groups.append( + [ + cls._normalize_serialized_native_local_shard(record) + for record in group + ] + ) + normalized_pruned = [ + cls._normalize_serialized_native_local_shard(record) + for record in pruned_shards + ] + if any( + record["layout"] != ParameterLayout.WHOLE_PARAMETER_OWNER.value + for record in normalized_pruned + ): + raise ValueError( + "Gefen native pruned-shard metadata must describe whole-parameter nonowners" + ) + return { + "format_version": _NATIVE_LOCAL_SHARDS_LEGACY_FORMAT_VERSION, + "param_groups": normalized_groups, + "pruned_shards": normalized_pruned, + } + + @classmethod + def _normalize_serialized_native_local_shards_v2(cls, value): + if type(value) is not dict or set(value) != { + "format_version", + "logical_slots", + }: + raise ValueError("Gefen native local-shard metadata has an invalid schema") + format_version = value["format_version"] + if ( + type(format_version) is not int + or format_version != _NATIVE_LOCAL_SHARDS_FORMAT_VERSION + ): + raise ValueError( + "Unsupported Gefen native local-shard format_version: {}".format( + format_version + ) + ) + records = value["logical_slots"] + if type(records) is not list or not records: + raise ValueError( + "Gefen native local-shard logical_slots must be a nonempty list" + ) + + normalized_records = [] + seen_fqns = set() + previous_position = None + supported_layouts = { + ParameterLayout.REPLICATED, + ParameterLayout.FLATTENED_ELEMENT_SHARD, + ParameterLayout.WHOLE_PARAMETER_OWNER, + } + for record in records: + if type(record) is not dict or set(record) != { + "group_index", + "original_slot_index", + "compatibility_name", + "shard", + }: + raise ValueError( + "Gefen native local-shard logical slot has an invalid schema" + ) + if ( + type(record["shard"]) is dict + and "process_group" in record["shard"] + and record["shard"]["process_group"] is None + ): + raise ValueError( + "Gefen native local-shard metadata requires a process-group identity" + ) + try: + shard = _parse_shard_identity(record["shard"]) + slot = LogicalSlotBinding( + record["group_index"], + record["original_slot_index"], + record["compatibility_name"], + shard, + ) + except (TypeError, ValueError) as exc: + raise ValueError( + "Gefen native local-shard logical slot is invalid" + ) from exc + position = (slot.group_index, slot.original_slot_index) + if previous_position is None: + valid_position = position == (0, 0) + elif position[0] == previous_position[0]: + valid_position = position[1] == previous_position[1] + 1 + else: + valid_position = position == (previous_position[0] + 1, 0) + if not valid_position: + raise ValueError( + "Gefen native local-shard logical slot positions must be contiguous" + ) + previous_position = position + if slot.shard.layout not in supported_layouts: + raise ValueError( + "Gefen native local-shard metadata has an unsupported layout" + ) + if slot.shard.process_group is None: + raise ValueError( + "Gefen native local-shard metadata requires a process-group identity" + ) + fqn = slot.shard.parameter.fqn + if fqn in seen_fqns: + raise ValueError( + "Gefen native local-shard logical slot FQNs must be unique" + ) + seen_fqns.add(fqn) + normalized_records.append( + { + "group_index": slot.group_index, + "original_slot_index": slot.original_slot_index, + "compatibility_name": slot.compatibility_name, + "shard": _serialize_shard_identity(slot.shard), + } + ) + return { + "format_version": _NATIVE_LOCAL_SHARDS_FORMAT_VERSION, + "logical_slots": normalized_records, + } + + @classmethod + def _normalize_serialized_native_local_shards(cls, value): + if value is None: + return None + if type(value) is not dict or "format_version" not in value: + raise ValueError("Gefen native local-shard metadata has an invalid schema") + format_version = value["format_version"] + if type(format_version) is not int: + raise ValueError( + "Unsupported Gefen native local-shard format_version: {}".format( + format_version + ) + ) + if format_version == _NATIVE_LOCAL_SHARDS_LEGACY_FORMAT_VERSION: + return cls._normalize_serialized_native_local_shards_v1(value) + if format_version == _NATIVE_LOCAL_SHARDS_FORMAT_VERSION: + return cls._normalize_serialized_native_local_shards_v2(value) + raise ValueError( + "Unsupported Gefen native local-shard format_version: {}".format( + format_version + ) + ) + + def _native_local_shards_matches_live(self, value): + if value is None: + return self._serialized_native_local_shards() is None + if value["format_version"] == _NATIVE_LOCAL_SHARDS_LEGACY_FORMAT_VERSION: + return value == self._serialized_native_local_shards_v1() + if value["format_version"] == _NATIVE_LOCAL_SHARDS_FORMAT_VERSION: + return value == self._serialized_native_local_shards() + return False + + @staticmethod + def _parameter_in(parameters, candidate) -> bool: + return any(item is candidate for item in parameters) + + @staticmethod + def _assert_rebound_storage_disjoint(parameters) -> None: + storage_ranges = [] + for parameter in parameters: + if parameter.numel() == 0: + continue + storage = parameter.untyped_storage() + storage_id = (str(parameter.device), storage.data_ptr()) + if not parameter.is_contiguous(): + for other_id, _, _ in storage_ranges: + if other_id == storage_id: + raise ValueError( + "rebound target tensors must not share noncontiguous storage" + ) + storage_ranges.append((storage_id, None, None)) + continue + start = parameter.storage_offset() * parameter.element_size() + end = start + parameter.numel() * parameter.element_size() + for other_id, other_start, other_end in storage_ranges: + if other_id != storage_id: + continue + if other_start is None or max(start, other_start) < min(end, other_end): + raise ValueError( + "rebound target tensor storage ranges must not overlap" + ) + storage_ranges.append((storage_id, start, end)) + + @staticmethod + def _target_may_have_internal_storage_overlap(parameter) -> bool: + """Conservatively prove disjoint positive-stride tensor elements.""" + + required_span = 1 + dimensions = sorted( + (stride, size) + for size, stride in zip(parameter.shape, parameter.stride()) + if size > 1 + ) + for stride, size in dimensions: + if stride < required_span: + return True + required_span += (size - 1) * stride + return False + + def _assert_rebinding_pristine(self, rebindings) -> None: + if self._gefen_post_sharding_finalized: + raise RuntimeError("Gefen post-sharding identity is already finalized") + if ( + self._gefen_shard_bindings + or self._gefen_local_shard_bindings + or self._gefen_sharding_manifest is not None + or self._gefen_finalized_slots + or self._gefen_logical_slots + ): + raise RuntimeError( + "Gefen parameter rebinding found an incomplete prior identity plan" + ) + global_step = self._validate_rank_local_counter( + "gefen_global_step", self._gefen_global_step + ) + if global_step != 0 or self._gefen_codebook is not None: + raise RuntimeError( + "Gefen parameter rebinding is allowed only before optimizer mutation" + ) + if self._capt_stacks is not None: + raise RuntimeError( + "Gefen parameter rebinding cannot discard active capturable stacks" + ) + if self._gefen_codebook_by_device or self._gefen_codebook_lut_by_device: + raise RuntimeError( + "Gefen parameter rebinding cannot discard initialized codebook caches" + ) + for cache_name in ( + "_gefen_global_step_by_device", + "_sr_seed_by_device", + ): + cache = getattr(self, cache_name) + for value in cache.values(): + if ( + not torch.is_tensor(value) + or value.numel() != 1 + or value.dtype == torch.bool + or not bool(torch.isfinite(value.detach()).all()) + or float(value.detach().cpu().item()) != 0.0 + ): + raise RuntimeError( + "Gefen parameter rebinding found non-pristine device counters" + ) + + live = [param for group in self.param_groups for param in group["params"]] + sources = [item.old_parameter for item in rebindings] + allowed_state_keys = tuple(live) + tuple(sources) + for parameter in self._param_names: + if not self._parameter_in(allowed_state_keys, parameter): + raise RuntimeError( + "Gefen parameter rebinding found orphan parameter-name state" + ) + compatibility_name = self._param_names[parameter] + if ( + not isinstance(compatibility_name, str) + or compatibility_name != compatibility_name.lower() + ): + raise RuntimeError( + "Gefen parameter rebinding found an invalid compatibility name" + ) + allowed_names = {"name", _RANK_LOCAL_MEMBER_KEY} + for parameter, pstate in self.state.items(): + if not self._parameter_in(allowed_state_keys, parameter): + raise RuntimeError( + "Gefen parameter rebinding found orphan optimizer state" + ) + if not isinstance(pstate, dict): + raise RuntimeError( + "Gefen parameter rebinding found invalid constructor state" + ) + for key in pstate: + if key in allowed_names: + continue + if isinstance(key, str) and key.startswith( + _RANK_LOCAL_PAYLOAD_KEY_PREFIX + ): + continue + raise RuntimeError( + "Gefen parameter rebinding found initialized optimizer state" + ) + state_name = pstate.get("name") + if state_name is not None and ( + not isinstance(state_name, str) or state_name != state_name.lower() + ): + raise RuntimeError( + "Gefen parameter rebinding found an invalid state name" + ) + + def _validate_rebinding_layout(self, rebinding: ParameterRebinding) -> None: + shard = rebinding.shard + target = rebinding.new_parameter + if shard.layout is ParameterLayout.REPLICATED: + if target is None or tuple(target.shape) != shard.parameter.global_shape: + raise ValueError( + "replicated Gefen rebinding requires complete parameter storage" + ) + elif shard.layout is ParameterLayout.FLATTENED_ELEMENT_SHARD: + if target is None or target.ndim != 1: + raise ValueError( + "flattened Gefen rebinding requires a live 1-D tensor shard" + ) + if not target.is_contiguous(): + raise ValueError( + "flattened Gefen rebinding requires contiguous physical storage" + ) + if self._factored_v_2d and len(shard.parameter.global_shape) == 2: + raise ValueError( + "flattened logical matrices require factored_v_2d=False until " + "matrix-aware factored-state projection is implemented" + ) + else: + raise ValueError( + "plain Gefen rebinding supports replicated or flattened element shards only" + ) + if target.numel() != shard.logical_slice.length: + raise ValueError( + "Gefen rebound tensor numel does not match its logical slice" + ) + + @staticmethod + def _codebook_scope_backend_device_is_valid(backend: str, device) -> bool: + backend = str(backend).lower() + if "nccl" in backend: + return device.type == "cuda" and device.index is not None + if "gloo" in backend or "mpi" in backend: + return device.type == "cpu" + return device.type in {"cpu", "cuda"} + + def _validate_codebook_process_group_binding( + self, binding, rebindings, manifest + ) -> None: + if not isinstance(binding, CodebookProcessGroupBinding): + raise TypeError( + "codebook_process_group must be a CodebookProcessGroupBinding" + ) + identity = binding.identity + if any(shard.process_group != identity for shard in manifest.shards): + raise ValueError( + "every scoped manifest shard must use the codebook process-group identity" + ) + for rebinding in rebindings: + shard = rebinding.shard + if shard.local_member != binding.local_member: + raise ValueError( + "local shard members must match the codebook binding member" + ) + if shard.layout not in { + ParameterLayout.REPLICATED, + ParameterLayout.FLATTENED_ELEMENT_SHARD, + ParameterLayout.WHOLE_PARAMETER_OWNER, + }: + raise ValueError( + "explicit codebook scope supports replicated, flattened, or whole-parameter owner identities" + ) + + self._validate_codebook_runtime_binding(binding) + + def _validate_codebook_runtime_binding(self, binding) -> None: + members = binding.identity.ordered_members + if len(members) == 1: + if binding.process_group is not None: + raise ValueError( + "a one-member codebook scope must use process_group=None" + ) + return + + if self.capturable: + raise ValueError( + "capturable=True does not support a multi-member explicit codebook scope" + ) + + if binding.process_group is None: + raise ValueError( + "a multi-member codebook scope requires an explicit runtime process group" + ) + if ( + not torch.distributed.is_available() + or not torch.distributed.is_initialized() + ): + raise RuntimeError( + "a multi-member codebook scope requires initialized torch.distributed" + ) + import torch.distributed as dist + + try: + world = dist.get_world_size(binding.process_group) + group_rank = dist.get_group_rank(binding.process_group, dist.get_rank()) + backend = dist.get_backend(binding.process_group) + except Exception as exc: + raise ValueError( + "the current rank must belong to the explicit codebook process group" + ) from exc + if world != len(members): + raise ValueError( + "codebook process-group world size does not match its stable identity" + ) + if group_rank < 0 or members[group_rank] != binding.local_member: + raise ValueError( + "runtime codebook group order does not match ordered semantic members" + ) + if not self._codebook_scope_backend_device_is_valid( + backend, binding.collective_device + ): + raise ValueError( + "codebook collective device is incompatible with the runtime backend" + ) + if binding.collective_device.type == "cuda": + index = binding.collective_device.index + if ( + not torch.cuda.is_available() + or index is None + or index < 0 + or index >= torch.cuda.device_count() + ): + raise ValueError("codebook collective CUDA device is unavailable") + + @torch._dynamo.disable + def _assert_runtime_codebook_process_group(self, *, full: bool = False) -> None: + binding = self._gefen_codebook_process_group + if binding is None: + return + self._assert_finalized_binding_layout(full=full) + self._validate_codebook_runtime_binding(binding) + + def _capture_codebook_scope_binding_for_step(self): + """Return a usable failure-vote binding before inspecting live layout.""" + + binding = self._gefen_codebook_process_group + if binding is None: + return None + if not isinstance(binding, CodebookProcessGroupBinding): + raise TypeError("the runtime codebook process-group binding is invalid") + self._validate_codebook_runtime_binding(binding) + return binding + + def _codebook_parameter_contributes(self, parameter) -> bool: + binding = self._gefen_codebook_process_group + if binding is None: + return True + shard = self._gefen_shard_bindings.get(parameter) + if shard is None: + raise RuntimeError( + "scoped codebook parameter has no finalized shard identity" + ) + if shard.layout is ParameterLayout.REPLICATED: + return shard.local_member == binding.identity.ordered_members[0] + return True + + def _stage_post_sharding(self, rebindings, manifest, codebook_process_group=None): + self._assert_rebinding_pristine(rebindings) + live_slots = [] + for group_index, group in enumerate(self.param_groups): + params = list(group["params"]) + names = list(group.get("param_names", ())) + if len(names) != len(params): + names = [self._param_name(param) for param in params] + for parameter_index, (parameter, name) in enumerate(zip(params, names)): + live_slots.append((group_index, parameter_index, parameter, str(name))) + if len(rebindings) != len(live_slots): + raise ValueError( + "post_sharding requires exactly one rebinding for every optimizer slot" + ) + + manifest_fqns = {shard.parameter.fqn for shard in manifest.shards} + local_fqns = [item.shard.parameter.fqn for item in rebindings] + if len(set(local_fqns)) != len(local_fqns): + raise ValueError("local canonical parameter FQNs must be unique") + if set(local_fqns) != manifest_fqns: + raise ValueError( + "post_sharding manifest FQNs must exactly match optimizer slots" + ) + + assigned_positions = {} + binding_names = {} + final_targets = [] + seen_sources = [] + local_member_by_group = {} + for binding_index, rebinding in enumerate(rebindings): + if not isinstance(rebinding.old_parameter, torch.Tensor): + raise TypeError("rebinding source must be a Tensor") + if self._parameter_in(seen_sources, rebinding.old_parameter): + raise ValueError("rebinding source tensors must be unique") + seen_sources.append(rebinding.old_parameter) + process_group = rebinding.shard.process_group + if process_group is not None: + previous_member = local_member_by_group.setdefault( + process_group, rebinding.shard.local_member + ) + if previous_member != rebinding.shard.local_member: + raise ValueError( + "local shard bindings must use one member per process group" + ) + target = rebinding.new_parameter + if target is not None: + if not isinstance(target, torch.Tensor): + raise TypeError("rebinding target must be a Tensor or None") + if self._is_dtensor_parameter(target): + raise ValueError( + "stable DTensor rebinding requires the deferred logical-region identity schema" + ) + if torch.is_complex(target): + raise ValueError( + "Gefen does not support complex rebound parameters" + ) + if target.layout is not torch.strided or target.dtype not in ( + torch.float16, + torch.bfloat16, + torch.float32, + torch.float64, + ): + raise ValueError( + "rebound parameters require strided floating-point storage" + ) + if target.is_meta: + raise ValueError("rebound parameters require materialized storage") + if ( + rebinding.shard.layout is ParameterLayout.FLATTENED_ELEMENT_SHARD + and not target.is_contiguous() + ): + raise ValueError( + "flattened Gefen rebinding requires contiguous physical storage" + ) + if self._target_may_have_internal_storage_overlap(target): + raise ValueError( + "rebound target storage must be provably free of internal storage overlap" + ) + if not target.is_leaf and not target.retains_grad: + raise ValueError("can't optimize a non-leaf rebound Tensor") + if rebinding.old_parameter.grad is not None or ( + target is not None and target.grad is not None + ): + raise RuntimeError( + "post_sharding must run before source or target gradients exist" + ) + if rebinding.shard not in manifest.shards: + raise ValueError("local shard identity is absent from the manifest") + whole_owner = ( + rebinding.shard.layout is ParameterLayout.WHOLE_PARAMETER_OWNER + ) + local_owns = rebinding.shard.local_member == rebinding.shard.owner + if whole_owner and local_owns != (target is not None): + raise ValueError( + "whole-parameter owner bindings require storage only on the owner" + ) + if not whole_owner and target is None: + raise ValueError("only a whole-parameter non-owner may bind None") + + source_positions = [ + index + for index, slot in enumerate(live_slots) + if slot[2] is rebinding.old_parameter + ] + target_positions = [] + if target is not None: + target_positions = [ + index for index, slot in enumerate(live_slots) if slot[2] is target + ] + if len(source_positions) == 1: + position = source_positions[0] + if target_positions and target_positions != [position]: + raise ValueError( + "rebinding target already occupies another optimizer slot" + ) + elif len(source_positions) == 0 and len(target_positions) == 1: + position = target_positions[0] + source_known = ( + rebinding.old_parameter in self.state + or rebinding.old_parameter in self._param_names + ) + if not source_known: + raise ValueError("rebinding source is not registered with Gefen") + else: + raise ValueError( + "rebinding source must identify exactly one optimizer slot" + ) + if position in assigned_positions: + raise ValueError("multiple rebindings target one optimizer slot") + assigned_positions[position] = binding_index + compatibility_name = self._param_names.get(rebinding.old_parameter) + source_state = self.state.get(rebinding.old_parameter) + if compatibility_name is None and isinstance(source_state, dict): + compatibility_name = source_state.get("name") + if compatibility_name is None: + compatibility_name = live_slots[position][3] + if ( + not isinstance(compatibility_name, str) + or compatibility_name != compatibility_name.lower() + ): + raise ValueError( + "rebinding source compatibility name must be a lowercase string" + ) + binding_names[binding_index] = compatibility_name + if target is not None: + if self._parameter_in(final_targets, target): + raise ValueError("rebound target tensors must be unique") + final_targets.append(target) + self._validate_rebinding_layout(rebinding) + + if len(assigned_positions) != len(live_slots): + raise ValueError("post_sharding did not bind every optimizer slot") + self._assert_rebound_storage_disjoint(final_targets) + if codebook_process_group is not None: + self._validate_codebook_process_group_binding( + codebook_process_group, rebindings, manifest + ) + + staged = object.__new__(type(self)) + staged.__dict__ = self.__dict__.copy() + staged.param_groups = [] + staged.state = defaultdict(dict) + staged._param_names = {} + staged._gefen_shard_bindings = {} + local_bindings = [] + logical_slots = [] + slot_cursor = 0 + for group_index, group in enumerate(self.param_groups): + staged_group = dict(group) + staged_group.pop("_gefen_checkpoint_metadata", None) + staged_params = [] + staged_names = [] + names = list(group.get("param_names", ())) + if len(names) != len(group["params"]): + names = [self._param_name(param) for param in group["params"]] + for original_slot_index, _ in enumerate(names): + binding_index = assigned_positions[slot_cursor] + rebinding = rebindings[binding_index] + compatibility_name = binding_names[binding_index] + slot_cursor += 1 + target = rebinding.new_parameter + local_bindings.append((target, rebinding.shard)) + logical_slots.append( + LogicalSlotBinding( + group_index, + original_slot_index, + compatibility_name, + rebinding.shard, + ) + ) + if target is None: + continue + staged_params.append(target) + staged_names.append(compatibility_name) + staged._param_names[target] = compatibility_name + staged.state[target]["name"] = compatibility_name + staged._gefen_shard_bindings[target] = rebinding.shard + staged_group["params"] = staged_params + staged_group["param_names"] = staged_names + staged.param_groups.append(staged_group) + + staged._gefen_local_shard_bindings = tuple( + sorted(local_bindings, key=lambda item: item[1].sort_key) + ) + staged._gefen_logical_slots = tuple(logical_slots) + staged._gefen_sharding_manifest = manifest + staged._gefen_post_sharding_finalized = True + staged._gefen_codebook_process_group = codebook_process_group + staged._gefen_codebook_scope_validated = False + staged._gefen_codebook_by_device = {} + staged._gefen_codebook_lut_by_device = {} + staged._sr_seed_by_device = {} + staged._gefen_global_step_by_device = {} + staged._capt_stacks = None + staged._static_mark_sig = None + staged._lr_scalar_cache = None + staged._ensure_gefen_global_step_devices() + staged._install_rank_local_checkpoint_schema() + staged._gefen_finalized_slots = tuple( + tuple(group["params"]) for group in staged.param_groups + ) + return staged + + def post_sharding( + self, + rebindings, + *, + manifest: ShardingManifest, + codebook_process_group=None, + ) -> None: + """Atomically finalize every local optimizer slot after sharding.""" + + if not isinstance(manifest, ShardingManifest): + raise TypeError("manifest must be a ShardingManifest") + if isinstance(rebindings, (str, bytes)): + raise TypeError("rebindings must be a sequence") + rebindings = tuple(rebindings) + if not rebindings or any( + not isinstance(item, ParameterRebinding) for item in rebindings + ): + raise TypeError( + "rebindings must contain one or more ParameterRebinding values" + ) + sources = [] + for rebinding in rebindings: + if self._parameter_in(sources, rebinding.old_parameter): + raise ValueError("rebinding source tensors must be unique") + sources.append(rebinding.old_parameter) + staged = self._stage_post_sharding(rebindings, manifest, codebook_process_group) + self.__dict__.update(staged.__dict__) + self._invalidate_layout_forensics_caches() + # Compute the manifest shard set and digest exactly once at + # finalization; every later step guard and scoped operation header + # reuses this cache instead of rehashing every shard identity. + self._manifest_layout_forensics(refresh=True) + + def rebind_shard( + self, + old_parameter, + new_parameter, + *, + shard: ShardIdentity, + manifest: ShardingManifest, + ) -> None: + """Apply a complete one-slot shard plan through ``post_sharding``.""" + + self.post_sharding( + (ParameterRebinding(old_parameter, new_parameter, shard),), + manifest=manifest, + ) + + def rebind_parameter( + self, + old_parameter, + new_parameter, + *, + identity: ParameterIdentity, + ) -> None: + """Bind one complete replicated parameter with a canonical identity.""" + + shard = ShardIdentity( + identity, + ParameterLayout.REPLICATED, + LogicalSlice.full(identity), + ) + self.rebind_shard( + old_parameter, + new_parameter, + shard=shard, + manifest=ShardingManifest((shard,)), + ) + + @staticmethod + def _normalize_param_groups(params): + if isinstance(params, torch.Tensor): + param_groups = [params] + else: + param_groups = list(params) + if len(param_groups) == 0: + raise ValueError("optimizer got an empty parameter list") + if isinstance(param_groups[0], dict): + return param_groups + return [{"params": param_groups, "_gefen_implicit_group": True}] + + @staticmethod + def _validate_group_options(lr, betas, eps, weight_decay): + if torch.is_tensor(lr) and lr.numel() != 1: + raise ValueError( + "lr as a Tensor must be 1-element, got numel={}".format(lr.numel()) + ) + if not 0.0 <= lr: + raise ValueError("Invalid learning rate: {}".format(lr)) + elif not 0.0 <= betas[0] < 1.0: + raise ValueError("Invalid beta parameter at index 0: {}".format(betas[0])) + elif not 0.0 <= betas[1] < 1.0: + raise ValueError("Invalid beta parameter at index 1: {}".format(betas[1])) + elif not 0.0 <= weight_decay: + raise ValueError("Invalid weight_decay value: {}".format(weight_decay)) + elif not math.isfinite(eps) or eps <= 0.0: + raise ValueError( + "Invalid epsilon value: {}; eps must be finite and > 0".format(eps) + ) + + @staticmethod + def _iter_params_with_names(group_params): + if isinstance(group_params, torch.Tensor): + group_params = (group_params,) + for param_or_named_param in group_params: + if isinstance(param_or_named_param, tuple): + param_name, param = param_or_named_param + else: + param_name = None + param = param_or_named_param + if not isinstance(param, torch.Tensor): + raise TypeError( + "Gefen can only optimize Tensors, but one of the params is {}".format( + type(param).__name__, + ) + ) + if torch.is_complex(param): + raise ValueError( + "Gefen does not support complex parameters, but got a parameter with dtype {}".format( + param.dtype + ) + ) + yield param_name, param + + def _param_name(self, param) -> str: + state = self.state.get(param) + if state is not None and "name" in state: + return state["name"] + return self._param_names.get(param, "none") + + def _iter_group_params_with_names(self, group): + names = group.get("param_names") + for idx, param in enumerate(group["params"]): + if names is not None and idx < len(names): + yield names[idx], param + else: + yield self._param_name(param), param + + def _iter_codebook_params_with_names(self): + """Yield codebook inputs in canonical order when a scope is bound.""" + + if self._gefen_codebook_process_group is None: + for group in self.param_groups: + for name, parameter in self._iter_group_params_with_names(group): + yield group, name, parameter + return + by_identity = {} + for group in self.param_groups: + for name, parameter in self._iter_group_params_with_names(group): + by_identity[id(parameter)] = (group, name, parameter) + for parameter, _ in self._gefen_local_shard_bindings: + if parameter is not None: + yield by_identity[id(parameter)] + + def _has_local_codebook_gradients(self) -> bool: + return any( + parameter.grad is not None + for _, _, parameter in self._iter_codebook_params_with_names() + ) + + @staticmethod + def _unique_name(base: str, existing_names) -> str: + name = str(base).lower() + if name not in existing_names: + return name + suffix = 1 + while "{}_{}".format(name, suffix) in existing_names: + suffix += 1 + return "{}_{}".format(name, suffix) + + def _sync_param_names_to_state(self) -> None: + self._param_names = {} + logical_groups = None + if self._gefen_post_sharding_finalized: + logical_groups = [[] for _ in self.param_groups] + for logical_slot in self._gefen_logical_slots: + if ( + logical_slot.shard.layout is ParameterLayout.WHOLE_PARAMETER_OWNER + and logical_slot.shard.local_member != logical_slot.shard.owner + ): + continue + logical_groups[logical_slot.group_index].append( + logical_slot.compatibility_name + ) + for group_index, group in enumerate(self.param_groups): + params = group["params"] + if logical_groups is None: + names = list(group.get("param_names", ())) + if len(names) != len(params): + names = [self._param_name(p) for p in params] + else: + names = logical_groups[group_index] + if len(names) != len(params): + raise RuntimeError( + "Gefen finalized logical-slot registry does not match the loaded layout" + ) + normalized_names = [] + for param, name in zip(params, names): + name = str(name).lower() + normalized_names.append(name) + self._param_names[param] = name + self.state[param]["name"] = name + group["param_names"] = normalized_names + + def add_param_group(self, param_group): + """Add a param group while preserving the caller-visible group boundary. + + Accepts the same group dict shape as the constructor (``params`` as + tensors or ``(name, tensor)`` tuples; optional ``lr`` / ``betas`` (or + ``beta1``/``beta2``) / ``eps`` / ``weight_decay``, defaulting to the + constructor values). The public group is preserved; each parameter's + stable lowercase name is stored in its per-param state and in the + group's ``param_names`` list for introspection. + """ + if getattr(self, "_gefen_post_sharding_finalized", False): + raise RuntimeError( + "Gefen cannot add parameter groups after post_sharding finalization" + ) + if not isinstance(param_group, dict): + raise TypeError( + "param_group must be a dict, got {}".format(type(param_group).__name__) + ) + if "params" not in param_group: + raise ValueError("Gefen parameter group is missing the 'params' key.") + defaults = self.defaults + group_lr = param_group.get("lr", defaults["lr"]) + if "betas" in param_group: + group_betas = param_group["betas"] + else: + group_betas = ( + param_group.get("beta1", defaults["beta1"]), + param_group.get("beta2", defaults["beta2"]), + ) + group_eps = param_group.get("eps", defaults["eps"]) + group_weight_decay = param_group.get("weight_decay", defaults["weight_decay"]) + self._validate_group_options( + group_lr, group_betas, group_eps, group_weight_decay + ) + named_params = list(self._iter_params_with_names(param_group["params"])) + if len(named_params) == 0: + raise ValueError("optimizer got an empty parameter list") + # Keys beyond the ones Gefen manages (e.g. a subclass's extras) are + # carried onto the preserved group, mirroring torch's behavior of + # preserving unknown group keys. + consumed = ( + "params", + "name", + "lr", + "betas", + "beta1", + "beta2", + "eps", + "weight_decay", + "param_names", + "_gefen_implicit_group", + ) + extra = {k: v for k, v in param_group.items() if k not in consumed} + existing_names = set(self._param_names.values()) + group_index = len(self.param_groups) + # Build (and fully validate) the preserved group up front so a failure + # partway through -- a complex/non-tensor param from + # _iter_params_with_names, or a duplicate param -- leaves + # self.param_groups untouched. add_param_group is all-or-nothing: only + # once the whole group is materialized do we register any of it. + existing_params = set() + for g in self.param_groups: + existing_params.update(g["params"]) + batch_params = set() + group_name = param_group.get("name") + implicit_group = bool(param_group.get("_gefen_implicit_group", False)) + params = [] + param_names = [] + for param_index, (param_name, param) in enumerate(named_params): + # Reject duplicates before registering any group, so the base class's + # own (post-registration) duplicate check can't leave us partial. + if param in existing_params or param in batch_params: + raise ValueError( + "some parameters appear in more than one parameter group" ) batch_params.add(param) params.append(param) @@ -1483,9 +3253,7 @@ def _assert_capturable_if_capturing(self) -> None: ) if capturing: devices = { - param.device - for group in self.param_groups - for param in group["params"] + param.device for group in self.param_groups for param in group["params"] } capture_device = torch.device("cuda", torch.cuda.current_device()) if devices != {capture_device}: @@ -1518,9 +3286,7 @@ def _static_mark_signature(self): acc += sum(map(id, params)) acc ^= sum(id(p.grad) for p in params) state = self.state - acc += sum( - map(id, chain.from_iterable(map(dict.values, state.values()))) - ) + acc += sum(map(id, chain.from_iterable(map(dict.values, state.values())))) stacks = self._capt_stacks return ( acc, @@ -1573,9 +3339,7 @@ def _mark_state_static_for_compile(self) -> None: mark(self._gefen_codebook) lut = _codebook_search_lut(self._gefen_codebook) mark(lut) - self._gefen_codebook_lut_by_device[ - self._gefen_codebook.device - ] = lut + self._gefen_codebook_lut_by_device[self._gefen_codebook.device] = lut for device, cached in self._gefen_codebook_by_device.items(): mark(cached) lut = _codebook_search_lut(cached) @@ -1648,8 +3412,10 @@ def _sr_seed_on(self, device: torch.device): else int(global_counter.item()) ) seed = torch.full( - (), initial_step, - dtype=torch.int64, device=device, + (), + initial_step, + dtype=torch.int64, + device=device, ) self._sr_seed_by_device[device] = seed return seed @@ -1674,8 +3440,9 @@ def _device_gefen_global_step(self): ] if any(step != steps[0] for step in steps[1:]): raise RuntimeError( - "Gefen capturable global-step counters disagree across devices: " - "{}".format(steps) + "Gefen capturable global-step counters disagree across devices: {}".format( + steps + ) ) return steps[0] @@ -1691,7 +3458,9 @@ def _ensure_gefen_global_step_devices(self) -> None: ] if missing: device_step = self._device_gefen_global_step() - initial_step = self._gefen_global_step if device_step is None else device_step + initial_step = ( + self._gefen_global_step if device_step is None else device_step + ) self._gefen_global_step = initial_step for device in missing: self._gefen_global_step_by_device[device] = torch.full( @@ -1715,13 +3484,10 @@ def _synchronize_gefen_global_step_for_checkpoint(self) -> None: if device_step is not None: self._gefen_global_step = device_step if self.capturable and self._stochastic_round: - seed_steps = [ - int(seed.item()) for seed in self._sr_seed_by_device.values() - ] + seed_steps = [int(seed.item()) for seed in self._sr_seed_by_device.values()] if any(seed_step != self._gefen_global_step for seed_step in seed_steps): raise RuntimeError( - "Gefen capturable stochastic-round seeds disagree with the " - "optimizer global step: {} != {}".format( + "Gefen capturable stochastic-round seeds disagree with the optimizer global step: {} != {}".format( seed_steps, self._gefen_global_step ) ) @@ -1976,10 +3742,7 @@ def _capt_build_stacks(self) -> bool: steps2d = torch.empty((2, n), dtype=torch.float32, device=device) steps2d[0].copy_( torch.stack( - [ - self.state[p][key].detach() - for p, key in zip(rows, d["keys"]) - ] + [self.state[p][key].detach() for p, key in zip(rows, d["keys"])] ) ) steps2d[1].copy_( @@ -2105,9 +3868,7 @@ def _capt_batched_ready(self) -> bool: # per-row alias/routing guard, but validate group-owned hypers once per # contiguous device run instead of once per parameter. Heterogeneous # registries take the fully general path below. - if stacks and all( - len(device_stacks) == 1 for device_stacks in stacks.values() - ): + if stacks and all(len(device_stacks) == 1 for device_stacks in stacks.values()): if self._capt_single_cohort_ready(stacks): return True if capturing: @@ -2148,8 +3909,7 @@ def _capt_batched_ready(self) -> bool: or state["step"] is not stack["step_views"][i] or state.get(info[0]) is not stack["bc2_views"][i] or state.get("_capt_scalars") is not stack["scalar_views"][i] - or stack["betas"][i] - != (group["beta1"], group["beta2"]) + or stack["betas"][i] != (group["beta1"], group["beta2"]) or group["weight_decay"] != stack["weight_decay0"] ): ok = False @@ -2313,6 +4073,11 @@ def _gefen_codebook_device(self) -> torch.device: if p.grad is not None: return p.grad.device return p.device + # A whole-parameter non-owner deliberately has no local parameter + # storage but must still materialize the group-agreed codebook so its + # optimizer-common checkpoint state remains complete. + if self._gefen_codebook_process_group is not None: + return torch.device("cpu") raise ValueError( "Expected at least one parameter when choosing the Gefen codebook device." ) @@ -2333,9 +4098,7 @@ def _gefen_codebook_on(self, device: torch.device) -> Optional[torch.Tensor]: self._gefen_codebook_by_device[device] = cached return cached - def _gefen_codebook_lut_on( - self, device: torch.device - ) -> Optional[torch.Tensor]: + def _gefen_codebook_lut_on(self, device: torch.device) -> Optional[torch.Tensor]: # Resolve once per device and retain the tensor on the optimizer. The # eager Muon momentum path calls this per matrix; returning a stable # cached tensor avoids rebuilding the data_ptr/version cache key in the @@ -2390,15 +4153,19 @@ def _resolve_automatic_period( # Single decision point for a parameter's automatic period so the # codebook-learning pass and the step agree. force_1d_period_one short- # circuits 1D params to per-element (period==1) before the block search. + routing_name = param_name + shard = self._gefen_shard_bindings.get(param) + if shard is not None: + routing_name = shard.parameter.fqn if self._force_1d_period_one and param.ndim == 1: return 1 if self._force_2d_period_one and param.ndim == 2: return 1 if self._period_one_substrings: - lname = str(param_name).lower() + lname = str(routing_name).lower() if any(sub in lname for sub in self._period_one_substrings): return 1 - return self._predict_period_from_grad_sq(param_name, param, grad) + return self._predict_period_from_grad_sq(routing_name, param, grad) def _predict_period_from_grad_sq( self, param_name: str, param: torch.Tensor, grad: torch.Tensor @@ -2460,9 +4227,7 @@ def _predict_period_from_grad_sq( ) if backend == "cpu": - period_input = ( - grad_work.detach().float().square().reshape(-1).cpu().numpy() - ) + period_input = grad_work.detach().float().square().reshape(-1).cpu().numpy() elif backend == "gpu": if grad_work.device.type != "cuda": raise ValueError("FIND_PERIOD_BACKEND='gpu' requires a CUDA tensor") @@ -2509,104 +4274,546 @@ def _print_v_period( ) ) - def _iter_gefen_grad_periods(self, reuse_existing_periods: bool = False): + def _iter_gefen_grad_periods( + self, reuse_existing_periods: bool = False, staged_periods=None + ): + for _, param_name, p in self._iter_codebook_params_with_names(): + if p.grad is None: + continue + grad = p.grad + # Codebook learning runs before the step loop's own guard, so + # reject sparse grads here too with the same clear error. + if getattr(grad, "is_sparse", False): + raise RuntimeError("Gefen does not support sparse gradients") + if torch.is_complex(grad): + raise RuntimeError("Gefen does not support complex gradients") + if hasattr(grad, "to_local"): + grad = grad.to_local() + if hasattr(grad, "wait"): + grad = grad.wait() + grad = grad.detach() + flat = grad.reshape(-1) + if flat.numel() == 0: + continue - for group in self.param_groups: - for param_name, p in self._iter_group_params_with_names(group): - if p.grad is None: - continue - grad = p.grad - # Codebook learning runs before the step loop's own guard, so - # reject sparse grads here too with the same clear error. - if getattr(grad, "is_sparse", False): - raise RuntimeError("Gefen does not support sparse gradients") - if torch.is_complex(grad): - raise RuntimeError("Gefen does not support complex gradients") - if hasattr(grad, "to_local"): - grad = grad.to_local() - if hasattr(grad, "wait"): - grad = grad.wait() - grad = grad.detach() - flat = grad.reshape(-1) - if flat.numel() == 0: - continue + if reuse_existing_periods: + state = self.state[p] + if "automatic_period" not in state: + raise ValueError( + "Expected automatic_period to exist for {} before refreshing Gefen codebook at optimizer step {}".format( + param_name, + self._gefen_global_step, + ) + ) + period = state["automatic_period"] + elif flat.numel() == 1: + period = 1 + else: + period = self._resolve_automatic_period(param_name, p, grad) + + if flat.numel() % period != 0: + raise ValueError( + "Automatic partition period {} does not divide parameter {} with numel {} while learning Gefen codebook".format( + period, + param_name, + flat.numel(), + ) + ) + + if staged_periods is None: + self.state[p]["automatic_period"] = period + elif not reuse_existing_periods: + staged_periods.append((p, period)) + if not self._codebook_parameter_contributes(p): + continue + + yield param_name, flat, period, grad + + @torch._dynamo.disable + def _synchronize_codebook_scope_failure(self, error, phase: str) -> None: + binding = self._gefen_codebook_process_group + if binding is None or len(binding.identity.ordered_members) == 1: + if error is not None: + raise error + return + self._assert_runtime_codebook_process_group() + self._synchronize_prevalidated_codebook_scope_failure(error, phase, binding) + + @staticmethod + @torch._dynamo.disable + def _synchronize_prevalidated_codebook_scope_failure( + error, phase: str, binding + ) -> None: + """Synchronize through a binding validated before rank-local user code.""" + + if binding is None or len(binding.identity.ordered_members) == 1: + if error is not None: + raise error + return + failed = _synchronize_step_failure( + error is not None, + binding.process_group, + collective_device=binding.collective_device, + ) + if not failed: + return + if error is not None: + raise RuntimeError( + "scoped Gefen codebook {} failed on local member {}: {}".format( + phase, binding.local_member, error + ) + ) from error + raise RuntimeError( + "scoped Gefen codebook {} failed on another process-group member".format( + phase + ) + ) + + @torch._dynamo.disable + def _prepare_scoped_amp_optimizer_step(self) -> bool: + binding = self._gefen_codebook_process_group + if binding is None: + return _amp_prepare_optimizer_step(self) + found_inf = getattr(self, "found_inf", None) + local_present = hasattr(self, "found_inf") or hasattr( + self, "grad_scale" + ) + try: + local_overflow, grad_scale, local_scale = ( + _amp_optimizer_step_controls(self) + ) + if ( + found_inf is not None + and not torch.is_tensor(found_inf) + and len(binding.identity.ordered_members) > 1 + ): + raise RuntimeError( + "a multi-member scoped optimizer requires a group-aware " + "gradient scaler to provide tensor found_inf on every member" + ) + scale_present = grad_scale is not None + local_error = None + except Exception as exc: + local_overflow = False + scale_present = False + local_scale = 0.0 + local_error = exc + self._synchronize_codebook_scope_failure(local_error, "AMP overflow preflight") + if len(binding.identity.ordered_members) > 1: + import torch.distributed as dist + + amp_control = torch.tensor( + [ + int(local_present), + int(local_overflow), + int(scale_present), + local_scale, + ], + dtype=torch.float64, + device=binding.collective_device, + ) + controls = [ + torch.empty_like(amp_control) for _ in binding.identity.ordered_members + ] + dist.all_gather(controls, amp_control, group=binding.process_group) + if any(not torch.equal(item, controls[0]) for item in controls[1:]): + raise RuntimeError( + "scoped Gefen AMP requires identical protocol presence, " + "found_inf, and grad_scale on every process-group member; " + "use a group-aware gradient scaler" + ) + if local_overflow: + return False + elif local_overflow: + return False + return _amp_prepare_optimizer_step(self) + + def _reduce_codebook_scope_histogram(self, histogram: torch.Tensor) -> torch.Tensor: + binding = self._gefen_codebook_process_group + if binding is None or len(binding.identity.ordered_members) == 1: + return histogram + self._assert_runtime_codebook_process_group() + import torch.distributed as dist + + reduced = histogram.to(binding.collective_device).clone() + dist.all_reduce(reduced, op=dist.ReduceOp.SUM, group=binding.process_group) + return reduced.cpu() + + def _codebook_manifest_fingerprint(self): + return self._manifest_layout_forensics()[1] + + def _compute_codebook_manifest_fingerprint(self, manifest): + # Scoped manifests always carry contiguous slices and one process-group + # identity, so their payload (and digest) is unchanged; the ungrouped + # and logical-region forms only appear for unscoped finalized layouts, + # whose digest is cached but never exchanged in a scope header. + payload = tuple( + ( + shard.parameter.fqn, + shard.parameter.global_shape, + shard.layout.value, + shard.logical_slice.flat_offset + if isinstance(shard.logical_slice, LogicalSlice) + else tuple(shard.logical_slice.offsets), + shard.logical_slice.length + if isinstance(shard.logical_slice, LogicalSlice) + else tuple(shard.logical_slice.lengths), + None + if shard.process_group is None + else shard.process_group.semantic_name, + None + if shard.process_group is None + else shard.process_group.ordered_members, + shard.local_member, + shard.owner, + tuple( + ( + placement.mesh_axis, + placement.kind.value, + placement.coordinate, + placement.parts, + placement.parameter_dimension, + ) + for placement in shard.placements + ), + ) + for shard in manifest.shards + ) + return self._sha256_int64(repr(payload).encode("utf-8")) + + @staticmethod + def _sha256_int64(payload): + digest = hashlib.sha256(payload).digest() + return tuple( + int.from_bytes(digest[index : index + 8], "big", signed=True) + for index in range(0, len(digest), 8) + ) + + def _codebook_scope_fingerprint(self): + return self._sha256_int64( + repr(self._serialized_codebook_scope()).encode("utf-8") + ) + + def _codebook_value_fingerprint(self): + if self._gefen_codebook is None: + return (0, 0, 0, 0) + raw = bytes( + self._gefen_codebook.detach().cpu().contiguous().view(torch.uint8).tolist() + ) + return self._sha256_int64(raw) + + def _validate_codebook_scope_operation_header(self, operation: str) -> None: + binding = self._gefen_codebook_process_group + if binding is None or len(binding.identity.ordered_members) == 1: + return + operation_codes = { + "initialize": 1, + "refresh": 2, + "periodic_step": 3, + "step": 4, + } + if operation not in operation_codes: + raise ValueError("unknown scoped codebook operation") + self._assert_runtime_codebook_process_group() + import torch.distributed as dist + + header = torch.tensor( + [ + operation_codes[operation], + int(self._gefen_global_step), + int(self._gefen_codebook is not None), + int(self._deterministic), + int(self._codebook_refresh_every), + int(self.capturable), + int(hasattr(self, "found_inf") or hasattr(self, "grad_scale")), + *self._codebook_scope_fingerprint(), + *self._codebook_manifest_fingerprint(), + *self._codebook_value_fingerprint(), + # Trailing decision bit, deliberately excluded from the equality + # check below: collective-free rank-local operations + # (post-sharding, staged/native checkpoint loads) + # legitimately reset _gefen_codebook_scope_validated on a subset + # of members without changing any fingerprinted value. The group + # resolves "does any member need re-validation" here so + # _ensure_codebook_scope_agreement re-validates on every member + # together instead of diverging in front of its collectives. + int(self._gefen_codebook_scope_validated), + ], + dtype=torch.int64, + device=binding.collective_device, + ) + headers = [torch.empty_like(header) for _ in binding.identity.ordered_members] + dist.all_gather(headers, header, group=binding.process_group) + if any(not torch.equal(item[:-1], headers[0][:-1]) for item in headers[1:]): + raise RuntimeError( + "scoped Gefen codebook operation, step, scope, manifest, policy, " + "or old codebook differs across process-group members" + ) + if any(int(item[-1].item()) == 0 for item in headers): + self._gefen_codebook_scope_validated = False + if operation != "step" and self._gefen_codebook is not None: + self._verify_codebook_scope_agreement(self._gefen_codebook) + + def _validate_codebook_scope_contribution_controls( + self, staged_periods, *, reuse_existing_periods: bool + ) -> None: + binding = self._gefen_codebook_process_group + if binding is None or len(binding.identity.ordered_members) == 1: + return + self._assert_runtime_codebook_process_group() + import torch.distributed as dist + staged_by_parameter = { + id(parameter): int(period) for parameter, period in staged_periods + } + local_records = [] + for parameter, shard in self._gefen_local_shard_bindings: + active = int( + parameter is not None + and parameter.grad is not None + and parameter.numel() > 0 + ) + period = -1 + if active: if reuse_existing_periods: - state = self.state[p] - if "automatic_period" not in state: - raise ValueError( - "Expected automatic_period to exist for {} before refreshing Gefen codebook at optimizer step {}".format( - param_name, - self._gefen_global_step, - ) - ) - period = state["automatic_period"] - elif flat.numel() == 1: - period = 1 + period = int(self.state[parameter]["automatic_period"]) else: - period = self._resolve_automatic_period(param_name, p, grad) + period = staged_by_parameter[id(parameter)] + layout_code = { + ParameterLayout.REPLICATED: 1, + ParameterLayout.FLATTENED_ELEMENT_SHARD: 2, + ParameterLayout.WHOLE_PARAMETER_OWNER: 3, + }[shard.layout] + local_records.append( + (layout_code, active, period, shard.logical_slice.length) + ) - self.state[p]["automatic_period"] = period + header = torch.tensor( + list(self._codebook_manifest_fingerprint()) + [len(local_records)], + dtype=torch.int64, + device=binding.collective_device, + ) + headers = [torch.empty_like(header) for _ in binding.identity.ordered_members] + dist.all_gather(headers, header, group=binding.process_group) + if any(not torch.equal(item, headers[0]) for item in headers[1:]): + raise RuntimeError( + "scoped Gefen codebook manifests differ across process-group members" + ) - if flat.numel() % period != 0: - raise ValueError( - "Automatic partition period {} does not divide parameter {} with numel {} while learning Gefen codebook".format( - period, - param_name, - flat.numel(), + flattened = [] + for record in local_records: + flattened.extend(record) + local = torch.tensor( + flattened, dtype=torch.int64, device=binding.collective_device + ) + gathered = [torch.empty_like(local) for _ in binding.identity.ordered_members] + dist.all_gather(gathered, local, group=binding.process_group) + for record_index, (layout_code, _, _, _) in enumerate(local_records): + offset = record_index * 4 + if any(int(item[offset].item()) != layout_code for item in gathered): + raise RuntimeError( + "scoped Gefen codebook local layouts differ across members" + ) + if layout_code != 1: + if layout_code == 2: + nonempty_activity = { + int(item[offset + 1].item()) + for item in gathered + if int(item[offset + 3].item()) > 0 + } + if len(nonempty_activity) > 1: + raise RuntimeError( + "scoped flattened parameters require every nonempty shard to agree on gradient presence" ) - ) + continue + replicated_controls = { + (int(item[offset + 1].item()), int(item[offset + 2].item())) + for item in gathered + } + if len(replicated_controls) != 1: + raise RuntimeError( + "scoped replicated parameters require identical gradient " + "presence and automatic periods on every member" + ) + + def _verify_codebook_scope_agreement(self, codebook: torch.Tensor) -> None: + binding = self._gefen_codebook_process_group + if binding is None or len(binding.identity.ordered_members) == 1: + return + self._assert_runtime_codebook_process_group() + import torch.distributed as dist + + local = ( + codebook.detach() + .to(device=binding.collective_device, dtype=torch.float32) + .contiguous() + ) + gathered = [torch.empty_like(local) for _ in binding.identity.ordered_members] + dist.all_gather(gathered, local, group=binding.process_group) + if any(not torch.equal(item, gathered[0]) for item in gathered[1:]): + raise RuntimeError( + "scoped Gefen codebook values differ across process-group members" + ) + + def _ensure_codebook_scope_agreement(self) -> None: + binding = self._gefen_codebook_process_group + if binding is None or self._gefen_codebook_scope_validated: + return + # Scope re-validation runs only after a mutating API reset the flag, + # so its collectives already amortize one full forensic rebuild. + self._assert_runtime_codebook_process_group(full=True) + if len(binding.identity.ordered_members) == 1: + self._gefen_codebook_scope_validated = self._gefen_codebook is not None + return + import torch.distributed as dist + + control = torch.tensor( + [ + int(self._gefen_global_step), + int(self._gefen_codebook is not None), + int(self._deterministic), + int(self._codebook_refresh_every), + int(hasattr(self, "found_inf") or hasattr(self, "grad_scale")), + int(self.capturable), + *self._codebook_scope_fingerprint(), + *self._codebook_manifest_fingerprint(), + ], + dtype=torch.int64, + device=binding.collective_device, + ) + controls = [torch.empty_like(control) for _ in binding.identity.ordered_members] + dist.all_gather(controls, control, group=binding.process_group) + if any(not torch.equal(item, controls[0]) for item in controls[1:]): + raise RuntimeError( + "scoped Gefen codebook step, presence, policy, scope, manifest, " + "or AMP protocol differs across process-group members" + ) + if self._gefen_codebook is not None: + self._verify_codebook_scope_agreement(self._gefen_codebook) + self._gefen_codebook_scope_validated = True - yield param_name, flat, period, grad + def _codebook_requires_2d_parameters(self) -> bool: + return False + + def _assert_codebook_capture_ready(self) -> None: + if ( + torch.cuda.is_available() + and torch.cuda.is_current_stream_capturing() + and self._gefen_codebook is None + and self._has_local_codebook_gradients() + ): + raise RuntimeError( + "Gefen codebook initialization is host-driven and must complete " + "during eager warmup before CUDA graph capture" + ) def _learn_gefen_exact_codebook( self, reuse_existing_periods: bool = False, compute_mse_logging: bool = True, - ) -> Optional[torch.Tensor]: + ): + fused_build_ok = self._fused_build_ok + try: + return self._prepare_gefen_exact_codebook( + reuse_existing_periods=reuse_existing_periods, + compute_mse_logging=compute_mse_logging, + ) + except Exception: + self._fused_build_ok = fused_build_ok + raise + + def _prepare_gefen_exact_codebook( + self, + reuse_existing_periods: bool = False, + compute_mse_logging: bool = True, + ): + if torch.cuda.is_available() and torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "Gefen codebook initialization is host-driven and cannot run " + "inside CUDA graph capture; initialize during eager warmup" + ) # Fire the fused-toolchain probe ONCE here, at the first fused CUDA # extension boundary (step-1 codebook learning), so a broken toolchain # is discovered and downgraded before the histogram / period-finding / # update paths run. On a healthy box this just builds the extension a # little earlier; for fused=False the probe is a no-op (never builds). - if self.fused: - self._gefen_fused_toolchain_ok() - codebook = learn_gefen_exact_codebook_from_grad_periods( - grad_periods=self._iter_gefen_grad_periods( - reuse_existing_periods=reuse_existing_periods - ), - codebook_device=self._gefen_codebook_device(), - num_codebooks=256, - force_endpoints=True, - verbose=self.verbose, - compute_mse_logging=compute_mse_logging, - # ``fused=False`` promises a JIT-free pure-PyTorch path. Do not load - # the separate histogram extension in that mode (or after the fused - # toolchain probe failed). An explicit global period-backend choice - # remains independent of this histogram implementation decision. - use_fused_histogram=( - FUSE_HISTOGRAM_FOR_EXACT - and self.fused - and self._fused_build_ok is not False - ), + if not reuse_existing_periods: + self._validate_codebook_scope_operation_header("initialize") + staged_periods = [] + try: + if self.fused: + self._gefen_fused_toolchain_ok() + histogram = _gefen_exact_histogram_from_grad_periods( + grad_periods=self._iter_gefen_grad_periods( + reuse_existing_periods=reuse_existing_periods, + staged_periods=staged_periods, + ), + num_codebooks=256, + # ``fused=False`` promises a JIT-free pure-PyTorch path. Do not + # load the separate histogram extension in that mode (or after + # the fused toolchain probe failed). + use_fused_histogram=( + FUSE_HISTOGRAM_FOR_EXACT + and self.fused + and self._fused_build_ok is not False + ), + ) + local_error = None + except Exception as exc: + histogram = None + local_error = exc + self._synchronize_codebook_scope_failure( + local_error, "local histogram preparation" ) - if codebook is None: - return None + self._validate_codebook_scope_contribution_controls( + staged_periods, + reuse_existing_periods=reuse_existing_periods, + ) + histogram = self._reduce_codebook_scope_histogram(histogram) - return codebook + try: + codebook = _learn_gefen_exact_codebook_from_histogram( + bin_counts_cpu=histogram, + codebook_device=self._gefen_codebook_device(), + num_codebooks=256, + force_endpoints=True, + verbose=self.verbose, + compute_mse_logging=compute_mse_logging, + ) + local_error = None + except Exception as exc: + codebook = None + local_error = exc + self._synchronize_codebook_scope_failure( + local_error, "exact-DP solve and placement" + ) + if codebook is not None: + self._verify_codebook_scope_agreement(codebook) + return codebook, staged_periods def _ensure_gefen_codebook(self, reuse_existing_periods: bool = False) -> None: - codebook = self._learn_gefen_exact_codebook( + codebook, staged_periods = self._learn_gefen_exact_codebook( reuse_existing_periods=reuse_existing_periods, compute_mse_logging=False, ) if codebook is not None: + for parameter, period in staged_periods: + self.state[parameter]["automatic_period"] = period self._gefen_codebook = codebook self._gefen_codebook_by_device.clear() self._gefen_codebook_lut_by_device.clear() + self._gefen_codebook_scope_validated = ( + self._gefen_codebook_process_group is not None + ) def _step_automatic_factored( - self, group, param_name: str, p: torch.Tensor, grad: torch.Tensor + self, + group, + param_name: str, + p: torch.Tensor, + grad: torch.Tensor, + *, + state=None, ) -> None: # Adafactor-style factored second moment for a 2D param (opt-in via # factored_v_2d). The quantized-momentum machinery is byte-identical to @@ -2621,7 +4828,7 @@ def _step_automatic_factored( grad = grad.to_local() if hasattr(grad, "wait"): grad = grad.wait() - state = self.state[p] + state = self.state[p] if state is None else state beta1 = group["beta1"] beta2 = group["beta2"] lr = group["lr"] @@ -2637,9 +4844,7 @@ def _step_automatic_factored( elif flat_grad.numel() == 1: automatic_period = 1 else: - automatic_period = self._resolve_automatic_period( - param_name, p, grad - ) + automatic_period = self._resolve_automatic_period(param_name, p, grad) if flat_grad.numel() % automatic_period != 0: raise ValueError( "Automatic partition period {} does not divide parameter {} with numel {}".format( @@ -2697,8 +4902,7 @@ def _step_automatic_factored( codebook = self._gefen_codebook_on(p.device) if codebook is None: raise ValueError( - "Expected Gefen codebook to be initialized before the " - "factored update." + "Expected Gefen codebook to be initialized before the factored update." ) # The raw kernel flat-indexes a contiguous matrix. A strided leaf # Parameter is uncommon but valid (and can arise from model surgery @@ -2837,17 +5041,32 @@ def _refresh_codebook_with_requant(self) -> None: # the codebook that WROTE them, so the momentum survives the swap with # only nearest-codeword rounding error; m_magnitude is unchanged # (coefficients stay normalized to [-1, 1]). + if self.capturable: + raise RuntimeError( + "capturable=True does not support replacing the learned codebook" + ) + self._validate_codebook_scope_operation_header("refresh") old_codebook = self._gefen_codebook if old_codebook is None: return - new_codebook = self._learn_gefen_exact_codebook( + new_codebook, staged_periods = self._learn_gefen_exact_codebook( reuse_existing_periods=True, compute_mse_logging=False ) + if staged_periods: + raise RuntimeError( + "codebook refresh unexpectedly attempted to replace block periods" + ) if new_codebook is None: return staged_indices = [] - for pgroup in self.param_groups: - for p in pgroup["params"]: + if self._gefen_codebook_process_group is None: + parameters = [p for pgroup in self.param_groups for p in pgroup["params"]] + else: + parameters = [ + p for p, _ in self._gefen_local_shard_bindings if p is not None + ] + try: + for p in parameters: pstate = self.state.get(p) if not pstate or "m_codebook" not in pstate: continue @@ -2862,8 +5081,14 @@ def _refresh_codebook_with_requant(self) -> None: ) indices = gefen_nearest_codebook_indices(new_codebook_local, coeffs) staged_indices.append((pstate["m_codebook"], indices)) + local_error = None + except Exception as exc: + local_error = exc + self._synchronize_codebook_scope_failure( + local_error, "momentum requantization staging" + ) - # Commit only after every allocation/dequantization/search succeeded. + # Commit only after every member prepared every local replacement. # The staged tensors already have the destination shape/device, so this # tail is a deterministic sequence of infallible in-place copies. for stored_indices, indices in staged_indices: @@ -2881,6 +5106,28 @@ def _resuming_from_checkpoint(self) -> bool: for p in group["params"] ) + def _scope_agreed_resuming_from_checkpoint(self) -> bool: + # Per-parameter state presence is legitimately rank-asymmetric under an + # explicit scope (empty non-owner slots and zero-length shards never + # create state), so the rank-local _resuming_from_checkpoint() predicate + # must not gate collectives directly: resolve one group-wide decision + # before any member branches on it. Any member that restored periods + # makes the whole scope reuse them. + resuming = self._resuming_from_checkpoint() + binding = self._gefen_codebook_process_group + if binding is None or len(binding.identity.ordered_members) == 1: + return resuming + self._assert_runtime_codebook_process_group() + import torch.distributed as dist + + control = torch.tensor( + int(resuming), + dtype=torch.int32, + device=binding.collective_device, + ) + dist.all_reduce(control, op=dist.ReduceOp.MAX, group=binding.process_group) + return bool(int(control.item())) + def _maybe_refresh_gefen_codebook(self) -> None: if self._gefen_codebook is not None: # A codebook restored from a checkpoint may land on CPU (depending on @@ -2892,6 +5139,17 @@ def _maybe_refresh_gefen_codebook(self) -> None: self._gefen_codebook_lut_by_device.clear() return + binding = self._gefen_codebook_process_group + if ( + binding is None or len(binding.identity.ordered_members) == 1 + ) and not self._has_local_codebook_gradients(): + # No local histogram exists, so learning would return None. This + # fast no-op is also what keeps a no-gradient capturable step free + # of the host-driven exact-DP preparation path. Multi-member scopes + # still enter collectively because another member (for example a + # whole-parameter owner) may hold the only active gradient. + return + # No codebook yet. On a fresh run this learns it and predicts periods. # On resume the persisted codebook is normally restored in # load_state_dict, but FSDP optim-state consolidation strips custom @@ -2901,8 +5159,68 @@ def _maybe_refresh_gefen_codebook(self) -> None: # the refreshed periods desync from the restored vmean/m_codebook block # geometry and the vmean kernel aborts on a block-count mismatch. self._ensure_gefen_codebook( - reuse_existing_periods=self._resuming_from_checkpoint() + reuse_existing_periods=self._scope_agreed_resuming_from_checkpoint() + ) + + @torch.no_grad() + def initialize_codebook(self) -> bool: + """Collectively initialize the learned codebook without taking a step.""" + + self._assert_finalized_binding_layout(full=True) + self._assert_runtime_codebook_process_group() + self._assert_codebook_capture_ready() + self._validate_codebook_scope_operation_header("initialize") + try: + _assert_optimizer_gradients_structurally_valid( + self, require_2d_params=self._codebook_requires_2d_parameters() + ) + local_error = None + except Exception as exc: + local_error = exc + if self._gefen_codebook_process_group is not None: + self._synchronize_codebook_scope_failure(local_error, "gradient preflight") + elif local_error is not None: + raise local_error + self._ensure_codebook_scope_agreement() + if self._gefen_codebook is not None: + return False + self._ensure_gefen_codebook( + reuse_existing_periods=self._scope_agreed_resuming_from_checkpoint() ) + return self._gefen_codebook is not None + + @torch.no_grad() + def refresh_codebook(self) -> bool: + """Collectively relearn and atomically requantize the current codebook.""" + + self._assert_finalized_binding_layout(full=True) + self._assert_runtime_codebook_process_group() + if torch.cuda.is_available() and torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "capturable=True does not support replacing the learned codebook" + ) + self._validate_codebook_scope_operation_header("refresh") + if self.capturable: + raise RuntimeError( + "capturable=True does not support replacing the learned codebook" + ) + if self._gefen_codebook is None: + raise RuntimeError("Gefen codebook must be initialized before refresh") + try: + _assert_optimizer_gradients_structurally_valid( + self, require_2d_params=self._codebook_requires_2d_parameters() + ) + local_error = None + except Exception as exc: + local_error = exc + if self._gefen_codebook_process_group is not None: + self._synchronize_codebook_scope_failure(local_error, "gradient preflight") + elif local_error is not None: + raise local_error + self._ensure_codebook_scope_agreement() + before = self._gefen_codebook + self._refresh_codebook_with_requant() + return self._gefen_codebook is not before def _maybe_save_gefen_grad_histogram(self) -> None: if not hasattr(quantization_module, "LIST_STEPS_SAVE_HIST_GRAD"): @@ -3212,9 +5530,7 @@ def _automatic_momentum_update( * m_magnitude[r0:r1] ) # lerp() needs matching dtypes, so promote the (bf16) grad rows to fp32. - updated_m = current_m.lerp( - grad_view[r0:r1].to(current_m.dtype), 1 - beta1 - ) + updated_m = current_m.lerp(grad_view[r0:r1].to(current_m.dtype), 1 - beta1) # New per-block magnitude = max |updated_m| over each row, written back # into the persistent fp32 m_magnitude state in place. @@ -3298,7 +5614,13 @@ def _automatic_momentum_update_merged( return quantized_m * state["m_magnitude"] def _step_automatic( - self, group, param_name: str, p: torch.Tensor, grad: torch.Tensor + self, + group, + param_name: str, + p: torch.Tensor, + grad: torch.Tensor, + *, + state=None, ) -> None: # Unwrap the gradient to its local shard for both the fused and # non-fused paths. Under FSDP2 the gradient arrives as a sharded DTensor; @@ -3310,7 +5632,7 @@ def _step_automatic( if hasattr(grad, "wait"): grad = grad.wait() - state = self.state[p] + state = self.state[p] if state is None else state beta1 = group["beta1"] beta2 = group["beta2"] lr = group["lr"] @@ -3341,9 +5663,7 @@ def _step_automatic( elif local_numel == 1: automatic_period = 1 elif local_numel > 1: - automatic_period = self._resolve_automatic_period( - param_name, p, grad - ) + automatic_period = self._resolve_automatic_period(param_name, p, grad) else: raise ValueError( "Automatic partition received an empty parameter {}".format( @@ -3393,9 +5713,7 @@ def _step_automatic( # A capturable (tensor) step must be CLONED -- assigning the same # 0-dim tensor would alias the two counters and double-increment. step = state["step"] - state["vmean_step"] = ( - step.clone() if torch.is_tensor(step) else step - ) + state["vmean_step"] = step.clone() if torch.is_tensor(step) else step # Tier-1 uses a separately calibrated predicate for the current full # v1/v2 kernels. Both fold the vmean EMA, per-block stepsize math, and @@ -3423,8 +5741,12 @@ def _step_automatic( # Deterministic mode keeps the fused update and selects v1-full, whose # fixed thread tree produces bit-identical vmean/parameter writes for # identical inputs on a homogeneous architecture. - route_v2 = False if self._deterministic else _should_use_v2_full( - grad_view.shape[0], automatic_period, grad_view.device + route_v2 = ( + False + if self._deterministic + else _should_use_v2_full( + grad_view.shape[0], automatic_period, grad_view.device + ) ) if ( route_v2 @@ -3456,9 +5778,7 @@ def _step_automatic( # local shard's storage. DTensor.to_local() returns a view of that # storage, so in-place ops propagate back to the param. p_local = ( - p.to_local() - if (hasattr(p, "to_local") and hasattr(p, "placements")) - else p + p.to_local() if (hasattr(p, "to_local") and hasattr(p, "placements")) else p ) if self.capturable and (use_full_fused or use_v2_full): @@ -3496,8 +5816,7 @@ def _step_automatic( # contiguous; a non-contiguous shard is updated on a contiguous # copy that is copied back to preserve the in-place semantics. weight_decay_factor = ( - 1.0 if step_scalars is not None - else 1.0 - lr * group["weight_decay"] + 1.0 if step_scalars is not None else 1.0 - lr * group["weight_decay"] ) if p_local.is_contiguous(): self._automatic_gefen_fused_full_update( @@ -3537,8 +5856,7 @@ def _step_automatic( # (see _automatic_gefen_fused_update_v2_full). Same contiguity # handling as the v1-full branch. weight_decay_factor = ( - 1.0 if step_scalars is not None - else 1.0 - lr * group["weight_decay"] + 1.0 if step_scalars is not None else 1.0 - lr * group["weight_decay"] ) if p_local.is_contiguous(): self._automatic_gefen_fused_update_v2_full( @@ -3684,128 +6002,773 @@ def _nonfused_batch_key(self, group, p: torch.Tensor): group["weight_decay"], ) - def _nonfused_state_ready(self, p: torch.Tensor) -> bool: - state = self.state.get(p) - if state is None: - return False - return ( - "step" in state - and "vmean" in state - and "m_magnitude" in state - and "m_codebook" in state - and "automatic_period" in state + def _nonfused_state_ready(self, p: torch.Tensor) -> bool: + state = self.state.get(p) + if state is None: + return False + return ( + "step" in state + and "vmean" in state + and "m_magnitude" in state + and "m_codebook" in state + and "automatic_period" in state + ) + + def _dispatch_nonfused_batch(self, items) -> None: + # items: list of (group, name, p, grad) sharing block geometry + hypers. + # Sub-batch by an element budget so the merged transient stays bounded: + # K wide for tiny params (big launch win), narrow for larger params. + local_numel = items[0][2].grad.reshape(-1).numel() + budget = GEFEN_NONFUSED_FOREACH_BUDGET + k = max(1, budget // max(local_numel, 1)) + for start in range(0, len(items), k): + chunk = items[start : start + k] + if len(chunk) == 1: + group, name, p, grad = chunk[0] + self._step_automatic(group, name, p, grad) + else: + self._step_automatic_merged(chunk) + + def _step_automatic_merged(self, items) -> None: + # All items share (local numel, automatic_period, beta1, beta2, lr, eps, + # weight_decay). Stack their per-block tensors into one merged + # (K*nblocks, period) problem and run the existing non-fused arithmetic on + # it. Each op in that arithmetic is per-block (vmean/magnitude reductions, + # stepsize, weight decay) or per-element (lerp, codebook search, dequant, + # the final update), so the merged result is bit-for-bit identical to + # looping over the params one at a time. The stepsize op sequence here is + # the algebraic equal of the per-param cached-buffer form in + # _step_automatic (scalar/tensor == reciprocal(tensor)*scalar), verified + # bit-identical by the K-series scratch-buffer change. + group0 = items[0][0] + beta1 = group0["beta1"] + beta2 = group0["beta2"] + lr = group0["lr"] + eps = group0["eps"] + weight_decay = group0["weight_decay"] + + first_state = self.state[items[0][2]] + period = first_state["automatic_period"] + k = len(items) + + grad_views = [] + vmeans = [] + mags = [] + codebooks = [] + params = [] + for _, _, p, grad in items: + flat = grad.detach().reshape(-1) + grad_views.append(flat.view(-1, period)) + state = self.state[p] + vmeans.append(state["vmean"]) + mags.append(state["m_magnitude"]) + codebooks.append(state["m_codebook"]) + params.append(p) + nblocks = grad_views[0].shape[0] + + # Merged views. torch.stack copies into fresh contiguous buffers, so the + # per-param state tensors are untouched until we scatter the results back. + merged_grad = torch.stack(grad_views).reshape(k * nblocks, period) + merged_vmean = torch.stack(vmeans).reshape(k * nblocks, 1) + merged_mag = torch.stack(mags).reshape(k * nblocks, 1) + merged_codebook = torch.stack(codebooks).reshape(k * nblocks, period) + + # vmean update (same non-fused arithmetic, on the merged block grid). + automatic_vmean_update(merged_vmean, merged_grad, beta2, use_fused=False) + + for _, _, p, _ in items: + pstate = self.state[p] + pstate["step"] += 1 + # Backfill for pre-counter checkpoints (vmean as old as step), + # then advance the vmean age; the batch key groups by it, so one + # state's correction is valid for the whole batch. + if "vmean_step" not in pstate: + pstate["vmean_step"] = pstate["step"] - 1 + pstate["vmean_step"] += 1 + step_count = first_state["step"] + + if weight_decay > 0.0: + torch._foreach_mul_(params, 1 - lr * weight_decay) + + bias_correction_1 = 1 - beta1**step_count + bias_correction_2 = 1 - beta2 ** first_state["vmean_step"] + h = (merged_vmean.sqrt() / math.sqrt(bias_correction_2)).add_(eps) + stepsize = (1 / bias_correction_1) / h + + merged_state = { + "automatic_period": period, + "m_magnitude": merged_mag, + "m_codebook": merged_codebook, + } + shared_m = self._automatic_momentum_update_merged( + merged_state, merged_grad, beta1 + ) + update = shared_m * stepsize + update.mul_(lr) + + # Scatter results back. The momentum update wrote the new magnitude and + # codebook indices into the merged buffers in place; copy them, the new + # vmean, and the parameter delta into each param/state with one foreach + # call apiece. + update_views = update.reshape(k, nblocks, period) + param_updates = [update_views[i].reshape(params[i].shape) for i in range(k)] + torch._foreach_add_(params, param_updates, alpha=-1.0) + torch._foreach_copy_( + vmeans, list(merged_vmean.reshape(k, nblocks, 1).unbind(0)) + ) + torch._foreach_copy_(mags, list(merged_mag.reshape(k, nblocks, 1).unbind(0))) + torch._foreach_copy_( + codebooks, list(merged_codebook.reshape(k, nblocks, period).unbind(0)) + ) + + def _canonical_policy(self): + return { + "factored_v_2d": self._factored_v_2d, + "force_1d_period_one": self._force_1d_period_one, + "force_2d_period_one": self._force_2d_period_one, + "period_one_substrings": list(self._period_one_substrings), + "stochastic_round": self._stochastic_round, + "codebook_refresh_every": self._codebook_refresh_every, + } + + def _canonical_common_global_step(self) -> int: + device_step = self._device_gefen_global_step() + step = self._gefen_global_step if device_step is None else device_step + if self.capturable and self._stochastic_round: + seed_steps = [int(seed.item()) for seed in self._sr_seed_by_device.values()] + if any(seed_step != step for seed_step in seed_steps): + raise RuntimeError( + "Gefen capturable stochastic-round seeds disagree with the canonical optimizer global step" + ) + return step + + def _serialized_sharding_manifest(self): + return [ + self._serialized_canonical_shard(shard) + for shard in self._gefen_sharding_manifest.shards + ] + + def _canonical_live_entries(self): + entries = {} + for group in self.param_groups: + options = clone_canonical_value( + self._canonical_group_options_value(group), + path="parameter group options", + ) + for compatibility_name, parameter in self._iter_group_params_with_names( + group + ): + shard = self._gefen_shard_bindings[parameter] + entries[shard.parameter.fqn] = ( + parameter, + shard, + str(compatibility_name), + options, + ) + return entries + + @staticmethod + def _canonical_value_token(value): + if torch.is_tensor(value): + raw = ( + value.detach() + .to(device="cpu") + .resolve_conj() + .resolve_neg() + .contiguous() + .reshape(-1) + .view(torch.uint8) + .numpy() + .tobytes() + ) + try: + version = value._version + except RuntimeError: + version = None + return ( + "tensor", + id(value), + version, + hashlib.sha256(raw).digest(), + str(value.device), + str(value.dtype), + tuple(value.shape), + ) + if type(value) is dict: + return ( + "dict", + tuple( + (key, Gefen._canonical_value_token(value[key])) + for key in sorted(value, key=repr) + ), + ) + if type(value) in {list, tuple}: + return ( + type(value).__name__, + tuple(Gefen._canonical_value_token(item) for item in value), + ) + try: + hash(value) + token = value + except TypeError: + token = (id(value), repr(value)) + return (type(value).__name__, token) + + def _canonical_import_live_token(self): + from gefen.portable_runtime import _parameter_storage_token + + groups = tuple( + ( + id(group), + tuple(id(parameter) for parameter in group["params"]), + tuple( + str(name) for name, _ in self._iter_group_params_with_names(group) + ), + self._canonical_value_token(self._canonical_group_options_value(group)), + ) + for group in self.param_groups + ) + states = tuple( + ( + id(parameter), + id(self.state.get(parameter)), + self._canonical_value_token(self.state.get(parameter, {})), + ) + for group in self.param_groups + for parameter in group["params"] + ) + return ( + id(self.param_groups), + id(self.state), + id(self.defaults), + self._canonical_value_token(self.defaults), + groups, + states, + self._gefen_global_step, + self._device_gefen_global_step(), + self._canonical_value_token(self._gefen_codebook), + self._canonical_value_token(self._gefen_global_step_by_device), + self._canonical_value_token(self._sr_seed_by_device), + self._canonical_value_token(self._canonical_policy()), + self._deterministic, + self.capturable, + self.fused, + self.verbose, + self._fused_build_ok, + id(self._gefen_codebook_process_group), + self._canonical_value_token(self._serialized_codebook_scope()), + id(self._gefen_sharding_manifest), + id(self._gefen_logical_slots), + tuple( + ( + id(logical_slot), + logical_slot.group_index, + logical_slot.original_slot_index, + logical_slot.compatibility_name, + self._canonical_value_token( + self._serialized_canonical_shard(logical_slot.shard) + ), + ) + for logical_slot in self._gefen_logical_slots + ), + tuple( + ( + id(parameter), + shard.sort_key, + # Mirror the portable live token: a prepared import must go + # stale when a locally bound parameter's storage is + # retargeted or mutated between prepare and commit. + None + if parameter is None + else _parameter_storage_token(parameter), + ) + for parameter, shard in self._gefen_local_shard_bindings + ), + ) + + def _canonical_parameter_state_matches_variant( + self, shard, group_options, parameter_state, state_layout + ) -> bool: + carries_momentum = any( + key in parameter_state for key in ("m_codebook", "m_magnitude") + ) + carries_normuon = any( + key in parameter_state for key in ("normuon_v", "normuon_step") + ) + if bool(group_options.get("normuon", False)): + if carries_momentum and not carries_normuon: + return False + elif carries_normuon: + return False + fields = frozenset({"name", *parameter_state}) + parameter_rank = len(shard.parameter.global_shape) + sharded_mode = group_options.get("sharded_mode") + for variant in state_layout.parameter_variants: + if frozenset(variant.fields) != fields: + continue + if shard.layout not in variant.layouts: + continue + if variant.role is not ParameterStateRole.ANY: + continue + if variant.sharded_mode != sharded_mode: + continue + if ( + variant.parameter_ranks is not None + and parameter_rank not in variant.parameter_ranks + ): + continue + if parameter_rank in variant.excluded_parameter_ranks: + continue + return True + return False + + def _canonical_period_must_be_one(self, parameter, shard) -> bool: + if self._force_1d_period_one and parameter.ndim == 1: + return True + if self._force_2d_period_one and parameter.ndim == 2: + return True + fqn = shard.parameter.fqn.lower() + return any(substring in fqn for substring in self._period_one_substrings) + + def _validate_canonical_parameter_semantics( + self, + parameter, + shard, + group_options, + parameter_state, + state_layout, + global_step, + ) -> None: + carries_momentum = any( + key in parameter_state for key in ("m_codebook", "m_magnitude") + ) + carries_normuon = any( + key in parameter_state for key in ("normuon_v", "normuon_step") + ) + counters = {} + for counter_key in ( + "step", + "vmean_step", + "factored_step", + "normuon_step", + ): + if counter_key not in parameter_state: + continue + counter = self._validate_rank_local_counter( + counter_key, parameter_state[counter_key] + ) + counters[counter_key] = counter + if counter > global_step: + raise ValueError( + "Gefen canonical parameter counter {} exceeds the optimizer global step".format( + counter_key + ) + ) + if "step" in counters and any( + counters[counter_key] > counters["step"] + for counter_key in ( + "vmean_step", + "factored_step", + "normuon_step", + ) + if counter_key in counters + ): + raise ValueError( + "Gefen canonical secondary counter exceeds the parameter step" + ) + period = parameter_state.get("automatic_period") + if ( + period is not None + and self._canonical_period_must_be_one(parameter, shard) + and (type(period) is not int or period != 1) + ): + raise ValueError( + "Gefen canonical automatic period violates period-one policy" + ) + if bool(group_options.get("normuon", False)): + if carries_momentum and not carries_normuon: + raise ValueError( + "Gefen canonical initialized NorMuon state is incomplete" + ) + elif carries_normuon: + raise ValueError( + "Gefen canonical NorMuon state is invalid for the target policy" + ) + if not self._canonical_parameter_state_matches_variant( + shard, group_options, parameter_state, state_layout + ): + raise ValueError( + "Gefen canonical parameter state does not match a declared state variant" + ) + + @staticmethod + def _assert_canonical_state_outside_cuda_capture(operation) -> None: + if torch.cuda.is_available() and torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "canonical state {} cannot run during CUDA capture".format(operation) + ) + + def _assert_canonical_import_target_safe(self) -> None: + if not self.capturable: + return + device_step = self._device_gefen_global_step() + initialized_state = any( + any(key != "name" for key in self.state.get(parameter, {})) + for group in self.param_groups + for parameter in group["params"] + ) + if ( + self._capt_stacks is not None + or self._gefen_codebook is not None + or self._gefen_global_step != 0 + or (device_step is not None and device_step != 0) + or initialized_state + ): + raise RuntimeError( + "canonical import into a capturable optimizer requires a fresh " + "target before device state or CUDA graphs are initialized" + ) + + def export_canonical_state(self): + """Export an exact-binding, device-neutral local state fragment.""" + + self._assert_finalized_binding_layout(full=True) + self._assert_canonical_state_outside_cuda_capture("export") + if not self._canonical_state_layouts(): + raise RuntimeError( + "canonical state export requires a supported finalized identity layout and primitive algorithm policy" + ) + entries = self._canonical_live_entries() + parameters = {} + for fqn in sorted(entries): + parameter, shard, compatibility_name, options = entries[fqn] + state = {} + for key, value in self.state.get(parameter, {}).items(): + if key in _CANONICAL_DERIVED_PARAMETER_STATE_KEYS: + continue + if key == "name": + continue + if key not in _CANONICAL_PARAMETER_STATE_KEYS: + raise RuntimeError( + "canonical state export found undeclared parameter state key {!r}".format( + key + ) + ) + state[key] = clone_canonical_value( + value, path="parameters.{}.state.{}".format(fqn, key) + ) + parameters[fqn] = { + "compatibility_name": compatibility_name, + "shard": self._serialized_canonical_shard(shard), + "group_options": clone_canonical_value( + options, path="parameters.{}.group_options".format(fqn) + ), + "state": state, + } + return { + "format": "gefen.bound_state", + "format_version": CANONICAL_STATE_FORMAT_VERSION, + "coverage": "local_optimizer_fragment", + "implementation": self.optimizer_contract().implementation, + "policy": clone_canonical_value(self._canonical_policy(), path="policy"), + "common": { + "gefen_global_step": self._canonical_common_global_step(), + "gefen_codebook": clone_canonical_value( + self._gefen_codebook, path="common.gefen_codebook" + ), + "gefen_deterministic": self._deterministic, + "gefen_codebook_scope": clone_canonical_value( + self._serialized_codebook_scope(), + path="common.gefen_codebook_scope", + ), + }, + "manifest": clone_canonical_value( + self._serialized_sharding_manifest(), path="manifest" + ), + "parameters": parameters, + } + + def _normalize_canonical_state_import(self, state): + state = clone_canonical_value(state, path="canonical state") + expected_top = { + "format", + "format_version", + "coverage", + "implementation", + "policy", + "common", + "manifest", + "parameters", + } + if type(state) is not dict or set(state) != expected_top: + raise ValueError("Gefen canonical state has an invalid top-level schema") + if state["format"] != "gefen.bound_state": + raise ValueError("Unsupported Gefen canonical state format") + if ( + type(state["format_version"]) is not int + or state["format_version"] != CANONICAL_STATE_FORMAT_VERSION + ): + raise ValueError( + "Unsupported Gefen canonical state format_version: {}".format( + state["format_version"] + ) + ) + if state["coverage"] != "local_optimizer_fragment": + raise ValueError("Unsupported Gefen canonical state coverage") + implementation = self.optimizer_contract().implementation + if state["implementation"] != implementation: + raise ValueError( + "Gefen canonical state implementation does not match the target" + ) + live_policy = clone_canonical_value( + self._canonical_policy(), path="live policy" + ) + if not canonical_values_equal(state["policy"], live_policy): + raise ValueError( + "Gefen canonical state algorithm policy does not match the target" + ) + + common = state["common"] + if type(common) is not dict or set(common) != { + "gefen_global_step", + "gefen_codebook", + "gefen_deterministic", + "gefen_codebook_scope", + }: + raise ValueError("Gefen canonical common state has an invalid schema") + if ( + type(common["gefen_global_step"]) is not int + or common["gefen_global_step"] < 0 + ): + raise ValueError( + "Gefen canonical global step must be a nonnegative integer" + ) + if type(common["gefen_deterministic"]) is not bool: + raise ValueError("Gefen canonical deterministic policy must be a bool") + normalized_scope = self._normalize_serialized_codebook_scope( + common["gefen_codebook_scope"] ) + if normalized_scope != self._serialized_codebook_scope(): + raise ValueError( + "Gefen canonical codebook scope does not match the live binding" + ) + common["gefen_codebook_scope"] = normalized_scope - def _dispatch_nonfused_batch(self, items) -> None: - # items: list of (group, name, p, grad) sharing block geometry + hypers. - # Sub-batch by an element budget so the merged transient stays bounded: - # K wide for tiny params (big launch win), narrow for larger params. - local_numel = items[0][2].grad.reshape(-1).numel() - budget = GEFEN_NONFUSED_FOREACH_BUDGET - k = max(1, budget // max(local_numel, 1)) - for start in range(0, len(items), k): - chunk = items[start : start + k] - if len(chunk) == 1: - group, name, p, grad = chunk[0] - self._step_automatic(group, name, p, grad) - else: - self._step_automatic_merged(chunk) + manifest = state["manifest"] + if type(manifest) is not list: + raise ValueError("Gefen canonical manifest must be a list") + normalized_manifest = [ + self._normalize_serialized_canonical_shard(record) for record in manifest + ] + if normalized_manifest != self._serialized_sharding_manifest(): + raise ValueError( + "Gefen canonical manifest does not match the finalized target" + ) + state["manifest"] = normalized_manifest + + parameters = state["parameters"] + if type(parameters) is not dict: + raise ValueError("Gefen canonical parameters must be an FQN mapping") + live_entries = self._canonical_live_entries() + state_layout = self._canonical_state_variant_layout() + if set(parameters) != set(live_entries): + raise ValueError( + "Gefen canonical parameter FQNs do not match the finalized target" + ) + for fqn, record in parameters.items(): + if type(record) is not dict or set(record) != { + "compatibility_name", + "shard", + "group_options", + "state", + }: + raise ValueError( + "Gefen canonical parameter {!r} has an invalid schema".format(fqn) + ) + if type(record["compatibility_name"]) is not str: + raise ValueError("Gefen canonical compatibility names must be strings") + parameter, shard, _, live_options = live_entries[fqn] + normalized_shard = self._normalize_serialized_canonical_shard( + record["shard"] + ) + if normalized_shard != self._serialized_canonical_shard(shard): + raise ValueError( + "Gefen canonical parameter shard does not match the live binding" + ) + record["shard"] = normalized_shard + if not canonical_values_equal(record["group_options"], live_options): + raise ValueError( + "Gefen canonical parameter-group options do not match the target" + ) + parameter_state = record["state"] + if type(parameter_state) is not dict or any( + key == "name" or key not in _CANONICAL_PARAMETER_STATE_KEYS + for key in parameter_state + ): + raise ValueError( + "Gefen canonical parameter state has undeclared fields" + ) + self._validate_canonical_parameter_semantics( + parameter, + shard, + live_options, + parameter_state, + state_layout, + common["gefen_global_step"], + ) + record["state"] = clone_canonical_value( + parameter_state, path="parameters.{}.state".format(fqn) + ) + del parameter + return state + + def _canonical_native_state_dict(self, canonical_state): + native = self._base_state_dict_without_hooks() + native["state"] = {} + live_entries = self._canonical_live_entries() + for saved_group, live_group in zip(native["param_groups"], self.param_groups): + for parameter_id, parameter in zip( + saved_group["params"], live_group["params"] + ): + shard = self._gefen_shard_bindings[parameter] + record = canonical_state["parameters"][shard.parameter.fqn] + parameter_state = { + "name": self._param_name(parameter), + **clone_canonical_value( + record["state"], + path="parameters.{}.state".format(shard.parameter.fqn), + ), + } + native["state"][parameter_id] = parameter_state + + common = canonical_state["common"] + native["gefen_global_step"] = common["gefen_global_step"] + native["gefen_codebook"] = common["gefen_codebook"] + native["gefen_deterministic"] = common["gefen_deterministic"] + scope = common["gefen_codebook_scope"] + if scope is not None: + native["gefen_codebook_scope"] = scope + else: + native.pop("gefen_codebook_scope", None) + local_shards = self._serialized_native_local_shards() + if local_shards is not None: + native["gefen_native_local_shards"] = local_shards + else: + native.pop("gefen_native_local_shards", None) + metadata = { + "format_version": ( + _SCOPED_NATIVE_METADATA_VERSION if scope is not None else 1 + ), + "global_step": common["gefen_global_step"], + "codebook": common["gefen_codebook"], + "deterministic": common["gefen_deterministic"], + "device_anchor": self._checkpoint_device_anchor(), + } + if scope is not None: + metadata["codebook_scope"] = scope + if local_shards is not None: + metadata["native_local_shards"] = local_shards + for group in native["param_groups"]: + group["_gefen_checkpoint_metadata"] = metadata + del live_entries + return native + + def _preserve_canonical_target_configuration(self, staged) -> None: + staged.defaults = self.defaults.copy() + staged.param_groups = self.param_groups - def _step_automatic_merged(self, items) -> None: - # All items share (local numel, automatic_period, beta1, beta2, lr, eps, - # weight_decay). Stack their per-block tensors into one merged - # (K*nblocks, period) problem and run the existing non-fused arithmetic on - # it. Each op in that arithmetic is per-block (vmean/magnitude reductions, - # stepsize, weight decay) or per-element (lerp, codebook search, dequant, - # the final update), so the merged result is bit-for-bit identical to - # looping over the params one at a time. The stepsize op sequence here is - # the algebraic equal of the per-param cached-buffer form in - # _step_automatic (scalar/tensor == reciprocal(tensor)*scalar), verified - # bit-identical by the K-series scratch-buffer change. - group0 = items[0][0] - beta1 = group0["beta1"] - beta2 = group0["beta2"] - lr = group0["lr"] - eps = group0["eps"] - weight_decay = group0["weight_decay"] + def prepare_canonical_state_import(self, state): + """Validate and stage a canonical local import without live mutation.""" - first_state = self.state[items[0][2]] - period = first_state["automatic_period"] - k = len(items) + self._assert_finalized_binding_layout(full=True) + self._assert_canonical_state_outside_cuda_capture("import preparation") + self._assert_canonical_import_target_safe() + if not self._canonical_state_layouts(): + raise RuntimeError( + "canonical state import requires a supported finalized identity layout" + ) + normalized = self._normalize_canonical_state_import(state) + live_token = self._canonical_import_live_token() + staged = self._stage_load_state_dict( + self._canonical_native_state_dict(normalized) + ) + self._preserve_canonical_target_configuration(staged) + return make_prepared_canonical_state_import(self, live_token, staged) - grad_views = [] - vmeans = [] - mags = [] - codebooks = [] - params = [] - for _, _, p, grad in items: - flat = grad.detach().reshape(-1) - grad_views.append(flat.view(-1, period)) - state = self.state[p] - vmeans.append(state["vmean"]) - mags.append(state["m_magnitude"]) - codebooks.append(state["m_codebook"]) - params.append(p) - nblocks = grad_views[0].shape[0] + def commit_canonical_state_import(self, prepared) -> None: + """Commit one still-current prepared canonical import exactly once.""" - # Merged views. torch.stack copies into fresh contiguous buffers, so the - # per-param state tensors are untouched until we scatter the results back. - merged_grad = torch.stack(grad_views).reshape(k * nblocks, period) - merged_vmean = torch.stack(vmeans).reshape(k * nblocks, 1) - merged_mag = torch.stack(mags).reshape(k * nblocks, 1) - merged_codebook = torch.stack(codebooks).reshape(k * nblocks, period) + if not isinstance(prepared, PreparedCanonicalStateImport): + raise TypeError( + "prepared must be a PreparedCanonicalStateImport from this optimizer" + ) + if prepared._optimizer is not self: + raise ValueError("prepared canonical state belongs to another optimizer") + if prepared._consumed: + raise RuntimeError("prepared canonical state import was already consumed") + self._assert_finalized_binding_layout(full=True) + self._assert_canonical_state_outside_cuda_capture("import commit") + self._assert_canonical_import_target_safe() + if prepared._live_token != self._canonical_import_live_token(): + raise RuntimeError( + "live optimizer state changed after canonical import preparation" + ) + prepared._consumed = True + self._commit_staged_load_state_dict(prepared._staged) - # vmean update (same non-fused arithmetic, on the merged block grid). - automatic_vmean_update(merged_vmean, merged_grad, beta2, use_fused=False) + def import_canonical_state(self, state) -> None: + """Atomically prepare and commit an exact-binding canonical fragment.""" - for _, _, p, _ in items: - pstate = self.state[p] - pstate["step"] += 1 - # Backfill for pre-counter checkpoints (vmean as old as step), - # then advance the vmean age; the batch key groups by it, so one - # state's correction is valid for the whole batch. - if "vmean_step" not in pstate: - pstate["vmean_step"] = pstate["step"] - 1 - pstate["vmean_step"] += 1 - step_count = first_state["step"] + self.commit_canonical_state_import(self.prepare_canonical_state_import(state)) - if weight_decay > 0.0: - torch._foreach_mul_(params, 1 - lr * weight_decay) + def export_portable_state( + self, + *, + checkpoint_process_group, + transaction_id, + limits, + ): + """Collectively export exact period-one state by logical identity.""" - bias_correction_1 = 1 - beta1 ** step_count - bias_correction_2 = 1 - beta2 ** first_state["vmean_step"] - h = (merged_vmean.sqrt() / math.sqrt(bias_correction_2)).add_(eps) - stepsize = (1 / bias_correction_1) / h + from gefen.portable_runtime import _export_portable_state - merged_state = { - "automatic_period": period, - "m_magnitude": merged_mag, - "m_codebook": merged_codebook, - } - shared_m = self._automatic_momentum_update_merged( - merged_state, merged_grad, beta1 + return _export_portable_state( + self, + checkpoint_process_group=checkpoint_process_group, + transaction_id=transaction_id, + limits=limits, ) - update = shared_m * stepsize - update.mul_(lr) - # Scatter results back. The momentum update wrote the new magnitude and - # codebook indices into the merged buffers in place; copy them, the new - # vmean, and the parameter delta into each param/state with one foreach - # call apiece. - update_views = update.reshape(k, nblocks, period) - param_updates = [update_views[i].reshape(params[i].shape) for i in range(k)] - torch._foreach_add_(params, param_updates, alpha=-1.0) - torch._foreach_copy_(vmeans, list(merged_vmean.reshape(k, nblocks, 1).unbind(0))) - torch._foreach_copy_(mags, list(merged_mag.reshape(k, nblocks, 1).unbind(0))) - torch._foreach_copy_( - codebooks, list(merged_codebook.reshape(k, nblocks, period).unbind(0)) + def import_portable_state( + self, + state, + *, + checkpoint_process_group, + transaction_id, + limits, + ) -> None: + """Collectively stage and atomically publish exact portable state.""" + + from gefen.portable_runtime import _import_portable_state + + _import_portable_state( + self, + state, + checkpoint_process_group=checkpoint_process_group, + transaction_id=transaction_id, + limits=limits, ) + canonical_state_dict = export_canonical_state + load_canonical_state_dict = import_canonical_state + def state_dict(self): """Run optimizer state-dict hooks around Gefen's complete schema.""" + self._assert_finalized_binding_layout(full=True) for pre_hook in self._optimizer_state_dict_pre_hooks.values(): pre_hook(self) + self._assert_finalized_binding_layout(full=True) state_dict = self._state_dict_impl() for post_hook in self._optimizer_state_dict_post_hooks.values(): hook_result = post_hook(self, state_dict) @@ -3894,10 +6857,12 @@ def reject(message): raise RuntimeError(message) return None - if not torch.distributed.is_available() or not torch.distributed.is_initialized(): + if ( + not torch.distributed.is_available() + or not torch.distributed.is_initialized() + ): return reject( - "Gefen rank-local DTensor checkpointing requires an initialized " - "default distributed process group" + "Gefen rank-local DTensor checkpointing requires an initialized default distributed process group" ) import torch.distributed as dist @@ -4054,9 +7019,7 @@ def _rank_local_sharded_signature(self, context=None): for item in p.placements ], "mesh": self._device_mesh_checkpoint_signature(mesh), - "coordinate": None - if coordinate is None - else list(coordinate), + "coordinate": None if coordinate is None else list(coordinate), "sharded_mode": group.get("sharded_mode"), } ) @@ -4148,13 +7111,15 @@ def _unpack_checkpoint_payload(cls, packed, tensors): def _serialize_rank_local_payload(payload) -> torch.Tensor: buffer = io.BytesIO() torch.save(payload, buffer) - return torch.frombuffer( - bytearray(buffer.getvalue()), dtype=torch.uint8 - ).clone() + return torch.frombuffer(bytearray(buffer.getvalue()), dtype=torch.uint8).clone() @staticmethod def _deserialize_rank_local_payload(payload: torch.Tensor): - if not torch.is_tensor(payload) or payload.dtype != torch.uint8 or payload.dim() != 1: + if ( + not torch.is_tensor(payload) + or payload.dtype != torch.uint8 + or payload.dim() != 1 + ): raise ValueError( "Gefen rank-local checkpoint payload must be a 1-D uint8 tensor" ) @@ -4236,24 +7201,27 @@ def _consolidate_rank_local_sharded_state( ) if any(control["manifest"] != local_manifest for control in controls): raise RuntimeError( - "Gefen DTensor checkpoint parameter identifiers, names, dtypes, " - "or global shapes differ across ranks" + "Gefen DTensor checkpoint parameter identifiers, names, dtypes, or global shapes differ across ranks" ) global_steps = [control["global_step"] for control in controls] - if any(type(step) is not int for step in global_steps) or len( - set(global_steps) - ) != 1: + if ( + any(type(step) is not int for step in global_steps) + or len(set(global_steps)) != 1 + ): raise RuntimeError( - "Gefen rank-local checkpoint global_step differs across ranks: " - "{}".format(global_steps) + "Gefen rank-local checkpoint global_step differs across ranks: {}".format( + global_steps + ) ) deterministic_values = [control["deterministic"] for control in controls] - if any(type(value) is not bool for value in deterministic_values) or len( - set(deterministic_values) - ) != 1: + if ( + any(type(value) is not bool for value in deterministic_values) + or len(set(deterministic_values)) != 1 + ): raise RuntimeError( - "Gefen rank-local checkpoint deterministic policy differs across " - "ranks: {}".format(deterministic_values) + "Gefen rank-local checkpoint deterministic policy differs across ranks: {}".format( + deterministic_values + ) ) signatures = { @@ -4293,8 +7261,9 @@ def _consolidate_rank_local_sharded_state( dist.all_gather_object(serialization_statuses, serialization_status) if any(not item.get("ok", False) for item in serialization_statuses): raise RuntimeError( - "Gefen rank-local checkpoint payload serialization failed across " - "ranks: {}".format(serialization_statuses) + "Gefen rank-local checkpoint payload serialization failed across ranks: {}".format( + serialization_statuses + ) ) serialized_payloads = { str(src): self._broadcast_checkpoint_payload( @@ -4373,9 +7342,7 @@ def _state_dict_impl(self, *, consolidate_rank_local: bool = True): # membership matches the id()-keyed param_mappings the base class # builds. The no-orphan path (every non-wrapped run) takes the plain # super() call, bit-for-bit unchanged. - reachable = { - param for group in self.param_groups for param in group["params"] - } + reachable = {param for group in self.param_groups for param in group["params"]} orphaned = [param for param in self.state if param not in reachable] if orphaned: original_order = list(self.state) @@ -4424,8 +7391,7 @@ def _state_dict_impl(self, *, consolidate_rank_local: bool = True): def _is_scratch_key(key): return key in scratch_keys or ( - isinstance(key, str) - and key.startswith(_RANK_LOCAL_PAYLOAD_KEY_PREFIX) + isinstance(key, str) and key.startswith(_RANK_LOCAL_PAYLOAD_KEY_PREFIX) ) def _compact(value): @@ -4446,11 +7412,7 @@ def _compact(value): state_dict["state"] = { pid: ( - { - k: _compact(v) - for k, v in pstate.items() - if not _is_scratch_key(k) - } + {k: _compact(v) for k, v in pstate.items() if not _is_scratch_key(k)} if isinstance(pstate, dict) else pstate ) @@ -4458,6 +7420,12 @@ def _compact(value): } state_dict["gefen_global_step"] = self._gefen_global_step state_dict["gefen_deterministic"] = self._deterministic + codebook_scope = self._serialized_codebook_scope() + if codebook_scope is not None: + state_dict["gefen_codebook_scope"] = codebook_scope + native_local_shards = self._serialized_native_local_shards() + if native_local_shards is not None: + state_dict["gefen_native_local_shards"] = native_local_shards # The exact-DP codebook is learned once on the first step and then frozen # for the rest of the run. It is not per-param state, so persist it # explicitly; without it resume would reinterpret the restored uint8 @@ -4475,7 +7443,9 @@ def _compact(value): # learned one. The loader removes this private transport metadata before # handing groups to torch, so it never leaks into live scheduler groups. checkpoint_metadata = { - "format_version": 1, + "format_version": ( + _SCOPED_NATIVE_METADATA_VERSION if codebook_scope is not None else 1 + ), "global_step": self._gefen_global_step, "codebook": self._gefen_codebook, "deterministic": self._deterministic, @@ -4487,10 +7457,12 @@ def _compact(value): # of the private group metadata during load. "device_anchor": self._checkpoint_device_anchor(), } + if codebook_scope is not None: + checkpoint_metadata["codebook_scope"] = self._serialized_codebook_scope() + if native_local_shards is not None: + checkpoint_metadata["native_local_shards"] = native_local_shards if consolidate_rank_local: - self._consolidate_rank_local_sharded_state( - state_dict, checkpoint_metadata - ) + self._consolidate_rank_local_sharded_state(state_dict, checkpoint_metadata) for group in state_dict["param_groups"]: group["_gefen_checkpoint_metadata"] = checkpoint_metadata return state_dict @@ -4574,11 +7546,13 @@ def _prepare_load_state_dict(self, state_dict): ``_publish_load_state_dict`` for the standalone atomic load. """ + self._assert_finalized_binding_layout(full=True) state_dict = state_dict.copy() for pre_hook in self._optimizer_load_state_dict_pre_hooks.values(): hook_result = pre_hook(self, state_dict) if hook_result is not None: state_dict = hook_result + self._assert_finalized_binding_layout(full=True) return self._stage_load_state_dict(state_dict) def _publish_load_state_dict(self, staged): @@ -4636,16 +7610,15 @@ def _commit_staged_load_state_dict(self, staged) -> None: # preserve that public object identity while publishing the staged # value. Both mappings are ordinary built-in dicts created by Optimizer. live_defaults = self.defaults - live_defaults.update(staged.defaults) + dict.update(live_defaults, staged.defaults) staged.defaults = live_defaults - self.__dict__.update(staged.__dict__) + dict.update(self.__dict__, staged.__dict__) + self._invalidate_layout_forensics_caches() def _validate_loaded_native_state(self) -> None: """Validate the complete prepared native state before publication.""" - self._validate_rank_local_counter( - "gefen_global_step", self._gefen_global_step - ) + self._validate_rank_local_counter("gefen_global_step", self._gefen_global_step) signature = self._rank_local_sharded_signature(context={}) states = [ self.state[param] @@ -4657,6 +7630,9 @@ def _validate_loaded_native_state(self) -> None: signature, self._gefen_codebook, allow_legacy_vmean_counter=True, + allow_preinitialized_periods=( + self._gefen_global_step == 0 and self._gefen_codebook is not None + ), ) for group in self.param_groups: self._validate_group_options( @@ -4741,6 +7717,7 @@ def _validate_rank_local_states( codebook, *, allow_legacy_vmean_counter: bool = False, + allow_preinitialized_periods: bool = False, ) -> None: if not isinstance(states, list) or len(states) != len(signature): raise ValueError( @@ -4756,17 +7733,14 @@ def _validate_rank_local_states( expected_name = param_signature["name"] if pstate.get("name") != expected_name: raise ValueError( - "Gefen rank-local checkpoint state name/order differs: " - "checkpoint={!r} expected={!r}".format( + "Gefen rank-local checkpoint state name/order differs: checkpoint={!r} expected={!r}".format( pstate.get("name"), expected_name ) ) counters = {} for key in counter_keys: if key in pstate: - counters[key] = self._validate_rank_local_counter( - key, pstate[key] - ) + counters[key] = self._validate_rank_local_counter(key, pstate[key]) local_shape = param_signature["local_shape"] state_shape = ( @@ -4783,8 +7757,7 @@ def _validate_rank_local_states( ) if state_numel == 0 or state_numel % period != 0: raise ValueError( - "Gefen rank-local checkpoint automatic_period does not divide " - "the parameter state geometry" + "Gefen rank-local checkpoint automatic_period does not divide the parameter state geometry" ) momentum_keys = ("m_codebook", "m_magnitude") @@ -4801,8 +7774,7 @@ def _validate_rank_local_states( ) if any(key in pstate for key in initialized_keys) and not carries_momentum: raise ValueError( - "Gefen rank-local checkpoint initialized state is missing " - "quantized momentum" + "Gefen rank-local checkpoint initialized state is missing quantized momentum" ) if ( period is not None @@ -4811,10 +7783,13 @@ def _validate_rank_local_states( param_signature.get("sharded") and param_signature.get("sharded_mode") == "distributed" ) + and not ( + allow_preinitialized_periods + and not any(key in pstate for key in initialized_keys) + ) ): raise ValueError( - "Gefen rank-local checkpoint automatic_period is invalid without " - "initialized momentum" + "Gefen rank-local checkpoint automatic_period is invalid without initialized momentum" ) if carries_momentum: has_quantized_momentum = True @@ -4849,8 +7824,7 @@ def _validate_rank_local_states( or not bool((magnitude >= 0).all()) ): raise ValueError( - "Gefen rank-local checkpoint m_magnitude geometry/dtype/values " - "are invalid" + "Gefen rank-local checkpoint m_magnitude geometry/dtype/values are invalid" ) vmean = pstate.get("vmean") if vmean is not None and ( @@ -4861,8 +7835,7 @@ def _validate_rank_local_states( or not bool((vmean >= 0).all()) ): raise ValueError( - "Gefen rank-local checkpoint vmean geometry/dtype/values " - "are invalid" + "Gefen rank-local checkpoint vmean geometry/dtype/values are invalid" ) has_vmean = "vmean" in pstate @@ -4871,14 +7844,9 @@ def _validate_rank_local_states( raise ValueError( "Gefen rank-local checkpoint block second moment is incomplete" ) - if ( - has_vmean - and not has_vmean_step - and not allow_legacy_vmean_counter - ): + if has_vmean and not has_vmean_step and not allow_legacy_vmean_counter: raise ValueError( - "Gefen rank-local checkpoint block second moment is missing " - "vmean_step" + "Gefen rank-local checkpoint block second moment is missing vmean_step" ) # GefenMuon groups carry a sharded_mode and intentionally use @@ -4889,26 +7857,19 @@ def _validate_rank_local_states( if carries_factored: if "factored_step" not in pstate: raise ValueError( - "Gefen rank-local checkpoint factored state is missing " - "factored_step" + "Gefen rank-local checkpoint factored state is missing factored_step" ) if counters["factored_step"] < 1: raise ValueError( - "Gefen rank-local checkpoint initialized factored " - "state requires factored_step >= 1" + "Gefen rank-local checkpoint initialized factored state requires factored_step >= 1" ) elif "vmean" not in pstate: raise ValueError( - "Gefen rank-local checkpoint plain momentum is missing a " - "second moment" + "Gefen rank-local checkpoint plain momentum is missing a second moment" ) - elif ( - "vmean_step" in counters - and counters["vmean_step"] < 1 - ): + elif "vmean_step" in counters and counters["vmean_step"] < 1: raise ValueError( - "Gefen rank-local checkpoint initialized vmean requires " - "vmean_step >= 1" + "Gefen rank-local checkpoint initialized vmean requires vmean_step >= 1" ) factored = (pstate.get("v_row"), pstate.get("v_col")) @@ -4939,8 +7900,9 @@ def _validate_rank_local_states( or not bool((tensor >= 0).all()) ): raise ValueError( - "Gefen rank-local checkpoint {} geometry/dtype/values " - "are invalid".format(key) + "Gefen rank-local checkpoint {} geometry/dtype/values are invalid".format( + key + ) ) has_normuon_v = "normuon_v" in pstate @@ -4952,18 +7914,15 @@ def _validate_rank_local_states( if has_normuon_v: if param_signature.get("sharded_mode") is None: raise ValueError( - "Gefen rank-local checkpoint NorMuon state is invalid for " - "a plain Gefen parameter" + "Gefen rank-local checkpoint NorMuon state is invalid for a plain Gefen parameter" ) if not carries_momentum: raise ValueError( - "Gefen rank-local checkpoint NorMuon state is missing " - "initialized momentum" + "Gefen rank-local checkpoint NorMuon state is missing initialized momentum" ) if len(state_shape) != 2: raise ValueError( - "Gefen rank-local checkpoint NorMuon state requires a 2-D " - "parameter" + "Gefen rank-local checkpoint NorMuon state requires a 2-D parameter" ) normuon_v = pstate["normuon_v"] if ( @@ -4974,18 +7933,14 @@ def _validate_rank_local_states( or not bool((normuon_v >= 0).all()) ): raise ValueError( - "Gefen rank-local checkpoint normuon_v geometry/dtype/values " - "are invalid" + "Gefen rank-local checkpoint normuon_v geometry/dtype/values are invalid" ) if counters["normuon_step"] < 1: raise ValueError( - "Gefen rank-local checkpoint initialized NorMuon state " - "requires normuon_step >= 1" + "Gefen rank-local checkpoint initialized NorMuon state requires normuon_step >= 1" ) - self._validate_rank_local_codebook( - codebook, required=has_quantized_momentum - ) + self._validate_rank_local_codebook(codebook, required=has_quantized_momentum) def _unwrap_rank_local_sharded_checkpoint(self, state_dict) -> None: groups = state_dict.get("param_groups", ()) @@ -5017,8 +7972,7 @@ def _unwrap_rank_local_sharded_checkpoint(self, state_dict) -> None: if not markers: if wrapped_state: raise ValueError( - "Gefen checkpoint carries rank-local DTensor payloads but " - "is missing their topology metadata" + "Gefen checkpoint carries rank-local DTensor payloads but is missing their topology metadata" ) has_quantized_momentum = any( isinstance(pstate, dict) and "m_codebook" in pstate @@ -5030,15 +7984,12 @@ def _unwrap_rank_local_sharded_checkpoint(self, state_dict) -> None: "rank-local DTensor optimizer state. Older generic FSDP/DCP " "full optimizer checkpoints kept rank 0's codebook and block " "state for every rank. Re-save from the original run with a " - "Gefen version that writes {} payloads.".format( - _RANK_LOCAL_FORMAT - ) + "Gefen version that writes {} payloads.".format(_RANK_LOCAL_FORMAT) ) return if len(markers) != len(groups) or any(item is None for item in metadata): raise ValueError( - "Gefen rank-local DTensor checkpoint metadata is present on only " - "some parameter groups" + "Gefen rank-local DTensor checkpoint metadata is present on only some parameter groups" ) marker = markers[0] if not isinstance(marker, dict): @@ -5055,8 +8006,7 @@ def _unwrap_rank_local_sharded_checkpoint(self, state_dict) -> None: ) if not self._uses_rank_local_sharded_state(): raise ValueError( - "A rank-local DTensor Gefen checkpoint can only load into the " - "same sharded optimizer topology" + "A rank-local DTensor Gefen checkpoint can only load into the same sharded optimizer topology" ) try: context = self._rank_local_checkpoint_context() @@ -5066,8 +8016,7 @@ def _unwrap_rank_local_sharded_checkpoint(self, state_dict) -> None: rank = context["global_rank"] if marker.get("world_size") != world: raise ValueError( - "Gefen rank-local DTensor checkpoints require the same world " - "size; checkpoint={} current={}".format( + "Gefen rank-local DTensor checkpoints require the same world size; checkpoint={} current={}".format( marker.get("world_size"), world ) ) @@ -5084,8 +8033,7 @@ def _unwrap_rank_local_sharded_checkpoint(self, state_dict) -> None: current_signature = self._rank_local_sharded_signature(context) if signatures[str(rank)] != current_signature: raise ValueError( - "Gefen rank-local DTensor checkpoint topology differs on rank {}: " - "checkpoint={!r} current={!r}".format( + "Gefen rank-local DTensor checkpoint topology differs on rank {}: checkpoint={!r} current={!r}".format( rank, signatures[str(rank)], current_signature ) ) @@ -5133,14 +8081,19 @@ def _unwrap_rank_local_sharded_checkpoint(self, state_dict) -> None: if isinstance(key, str) and key.startswith(_RANK_LOCAL_PAYLOAD_KEY_PREFIX) } - extras = set(pstate) - carrier_keys - { - _RANK_LOCAL_MEMBER_KEY, - "name", - } + extras = ( + set(pstate) + - carrier_keys + - { + _RANK_LOCAL_MEMBER_KEY, + "name", + } + ) if extras: raise ValueError( - "Gefen rank-local checkpoint parameter state has an invalid " - "schema: {!r}".format(pstate) + "Gefen rank-local checkpoint parameter state has an invalid schema: {!r}".format( + pstate + ) ) if pstate.get("name") != 0: raise ValueError( @@ -5155,17 +8108,14 @@ def _unwrap_rank_local_sharded_checkpoint(self, state_dict) -> None: if has_carrier: if carrier_keys != expected_carrier_keys: raise ValueError( - "Gefen rank-local checkpoint carrier has missing or " - "unexpected per-rank payload keys" + "Gefen rank-local checkpoint carrier has missing or unexpected per-rank payload keys" ) if serialized_payloads is not None: raise ValueError( "Gefen rank-local checkpoint repeats its payload transport" ) serialized_payloads = { - str(global_rank): pstate[ - _rank_local_payload_key(global_rank) - ] + str(global_rank): pstate[_rank_local_payload_key(global_rank)] for global_rank in context["world_ranks"] } elif pstate.get(_RANK_LOCAL_MEMBER_KEY) is not True: @@ -5195,8 +8145,13 @@ def _unwrap_rank_local_sharded_checkpoint(self, state_dict) -> None: "states", "codebook", } - if not isinstance(selected_payload, dict) or set(selected_payload) != expected_payload_keys: - raise ValueError("Gefen rank-local checkpoint payload has an invalid schema") + if ( + not isinstance(selected_payload, dict) + or set(selected_payload) != expected_payload_keys + ): + raise ValueError( + "Gefen rank-local checkpoint payload has an invalid schema" + ) payload_identity = ( selected_payload["format"] == _RANK_LOCAL_FORMAT and selected_payload["global_rank"] == rank @@ -5271,12 +8226,25 @@ def _load_state_dict_impl(self, state_dict): gefen_global_step = state_dict.pop("gefen_global_step", None) gefen_codebook = state_dict.pop("gefen_codebook", None) + has_top_level_codebook_scope = "gefen_codebook_scope" in state_dict + gefen_codebook_scope = state_dict.pop("gefen_codebook_scope", None) + if has_top_level_codebook_scope: + gefen_codebook_scope = self._normalize_serialized_codebook_scope( + gefen_codebook_scope + ) + has_top_level_native_local_shards = "gefen_native_local_shards" in state_dict + native_local_shards = state_dict.pop("gefen_native_local_shards", None) + if has_top_level_native_local_shards: + native_local_shards = self._normalize_serialized_native_local_shards( + native_local_shards + ) has_top_level_deterministic = "gefen_deterministic" in state_dict gefen_deterministic = state_dict.pop("gefen_deterministic", None) if has_top_level_deterministic and type(gefen_deterministic) is not bool: raise ValueError( - "Gefen checkpoint top-level deterministic policy must be a bool, " - "got {!r}".format(gefen_deterministic) + "Gefen checkpoint top-level deterministic policy must be a bool, got {!r}".format( + gefen_deterministic + ) ) if group_metadata: if len(group_metadata) != len(state_dict["param_groups"]): @@ -5284,50 +8252,84 @@ def _load_state_dict_impl(self, state_dict): "Gefen checkpoint metadata is present on only some parameter groups" ) first_metadata = group_metadata[0] - if first_metadata.get("format_version") not in ( - 1, - 2, - _RANK_LOCAL_METADATA_VERSION, - ): - raise ValueError( - "Unsupported Gefen checkpoint metadata format_version: {}".format( - first_metadata.get("format_version") - ) - ) + normalized_metadata_native_local_shards = [] for metadata in group_metadata: + metadata_version = metadata.get("format_version") + if metadata_version not in ( + 1, + 2, + _RANK_LOCAL_METADATA_VERSION, + _SCOPED_NATIVE_METADATA_VERSION, + ): + raise ValueError( + "Unsupported Gefen checkpoint metadata format_version: {}".format( + metadata_version + ) + ) + if ("codebook_scope" in metadata) != ( + metadata_version == _SCOPED_NATIVE_METADATA_VERSION + ): + raise ValueError( + "Gefen scoped checkpoint metadata must use format_version {}".format( + _SCOPED_NATIVE_METADATA_VERSION + ) + ) if ( "deterministic" in metadata and type(metadata["deterministic"]) is not bool ): raise ValueError( - "Gefen checkpoint parameter-group deterministic policy " - "must be a bool, got {!r}".format(metadata["deterministic"]) + "Gefen checkpoint parameter-group deterministic policy must be a bool, got {!r}".format( + metadata["deterministic"] + ) ) - for metadata in group_metadata[1:]: + normalized_metadata_native_local_shards.append( + self._normalize_serialized_native_local_shards( + metadata.get("native_local_shards") + ) + ) + for metadata_index, metadata in enumerate(group_metadata[1:], start=1): + same_version = metadata.get("format_version") == first_metadata.get( + "format_version" + ) same_step = metadata.get("global_step") == first_metadata.get( "global_step" ) left_codebook = metadata.get("codebook") right_codebook = first_metadata.get("codebook") - same_codebook = ( - left_codebook is right_codebook - or ( - torch.is_tensor(left_codebook) - and torch.is_tensor(right_codebook) - and torch.equal(left_codebook, right_codebook) - ) + same_codebook = left_codebook is right_codebook or ( + torch.is_tensor(left_codebook) + and torch.is_tensor(right_codebook) + and torch.equal(left_codebook, right_codebook) ) same_deterministic = metadata.get( "deterministic" ) == first_metadata.get("deterministic") - if not same_step or not same_codebook or not same_deterministic: + same_codebook_scope = metadata.get( + "codebook_scope" + ) == first_metadata.get("codebook_scope") + same_native_local_shards = ( + normalized_metadata_native_local_shards[metadata_index] + == normalized_metadata_native_local_shards[0] + ) + if ( + not same_version + or not same_step + or not same_codebook + or not same_deterministic + or not same_codebook_scope + or not same_native_local_shards + ): raise ValueError( - "Gefen checkpoint parameter groups carry inconsistent " - "optimizer metadata" + "Gefen checkpoint parameter groups carry inconsistent optimizer metadata" ) metadata_step = first_metadata.get("global_step", 0) metadata_codebook = first_metadata.get("codebook") metadata_deterministic = first_metadata.get("deterministic") + metadata_codebook_scope = self._normalize_serialized_codebook_scope( + first_metadata.get("codebook_scope") + ) + metadata_native_local_shards = normalized_metadata_native_local_shards[0] if gefen_global_step is None: gefen_global_step = metadata_step elif gefen_global_step != metadata_step: @@ -5354,14 +8356,35 @@ def _load_state_dict_impl(self, state_dict): and gefen_deterministic != metadata_deterministic ): raise ValueError( - "Gefen checkpoint top-level and parameter-group deterministic " - "policies disagree" + "Gefen checkpoint top-level and parameter-group deterministic policies disagree" ) + if not has_top_level_codebook_scope: + gefen_codebook_scope = metadata_codebook_scope + elif gefen_codebook_scope != metadata_codebook_scope: + raise ValueError( + "Gefen checkpoint top-level and parameter-group codebook scopes disagree" + ) + if not has_top_level_native_local_shards: + native_local_shards = metadata_native_local_shards + elif native_local_shards != metadata_native_local_shards: + raise ValueError( + "Gefen checkpoint top-level and parameter-group local shards disagree" + ) + live_codebook_scope = self._serialized_codebook_scope() + if gefen_codebook_scope != live_codebook_scope: + raise ValueError( + "Gefen checkpoint codebook scope does not match the live explicit binding" + ) + if not self._native_local_shards_matches_live(native_local_shards): + raise ValueError( + "Gefen checkpoint native local-shard identity does not match the live binding" + ) if gefen_deterministic is not None: if type(gefen_deterministic) is not bool: raise ValueError( - "Gefen checkpoint deterministic policy must be a bool, got " - "{!r}".format(gefen_deterministic) + "Gefen checkpoint deterministic policy must be a bool, got {!r}".format( + gefen_deterministic + ) ) if gefen_deterministic != self._deterministic: raise ValueError( @@ -5425,6 +8448,7 @@ def _load_state_dict_impl(self, state_dict): self._gefen_codebook = gefen_codebook self._gefen_codebook_by_device.clear() self._gefen_codebook_lut_by_device.clear() + self._gefen_codebook_scope_validated = False # The identity fingerprint would catch the wholesale state swap anyway; # reset it explicitly so the first post-load capturable step re-marks. self._static_mark_sig = None @@ -5450,7 +8474,10 @@ def _load_state_dict_impl(self, state_dict): if not torch.is_tensor(saved_value): continue live_value = live_state.get(key) - if torch.is_tensor(live_value) and live_value.dtype == saved_value.dtype: + if ( + torch.is_tensor(live_value) + and live_value.dtype == saved_value.dtype + ): continue device = ( live_value.device if torch.is_tensor(live_value) else param.device @@ -5503,24 +8530,76 @@ def step(self, closure=None): closure feed the first step's codebook learning correctly. The returned loss is passed through. """ - self._assert_capturable_if_capturing() + # Capture and validate the control scope independently of public + # parameter containers. A rank-local mutation made between steps can + # then be reported through the last known-good binding instead of + # stranding peers after a one-sided entry-guard failure. + scope_binding = self._capture_codebook_scope_binding_for_step() loss = None - if closure is not None: - with torch.enable_grad(): - loss = closure() + try: + self._assert_finalized_binding_layout() + self._assert_runtime_codebook_process_group() + self._assert_capturable_if_capturing() + self._assert_codebook_capture_ready() + if closure is not None: + with torch.enable_grad(): + loss = closure() + local_preamble_error = None + except Exception as exc: + loss = None + local_preamble_error = exc + if scope_binding is not None: + self._synchronize_prevalidated_codebook_scope_failure( + local_preamble_error, "step preamble", scope_binding + ) + elif local_preamble_error is not None: + raise local_preamble_error - _assert_optimizer_gradients_structurally_valid(self) + # The closure can replace a finalized parameter or otherwise invalidate + # the runtime binding. Recheck before the operation header and synchronize + # structural failures before any peer enters a scoped codebook collective. + try: + self._assert_finalized_binding_layout() + self._assert_runtime_codebook_process_group() + _assert_optimizer_gradients_structurally_valid(self) + local_preflight_error = None + except Exception as exc: + local_preflight_error = exc + if scope_binding is not None: + self._synchronize_prevalidated_codebook_scope_failure( + local_preflight_error, "gradient preflight", scope_binding + ) + elif local_preflight_error is not None: + raise local_preflight_error + self._validate_codebook_scope_operation_header("step") + self._ensure_codebook_scope_agreement() # GradScaler invokes native-AMP optimizers even on overflow. Decide # before codebook learning, periodic refresh, capturable counters, or - # parameter/state mutation. The attribute gate compiles away on the + # parameter/state mutation. A multi-member codebook scope must run the + # AMP protocol agreement on EVERY member -- including one with no local + # GradScaler attributes -- otherwise the presence gather inside + # _prepare_scoped_amp_optimizer_step is entered by only some members and + # deadlocks (mirrors the unconditional Muon/Hybrid AMP preflight). Off a + # multi-member scope the attribute gate still compiles away on the # ordinary BF16/FP32 path where GradScaler attaches nothing. - if ( - hasattr(self, "found_inf") or hasattr(self, "grad_scale") - ) and not _amp_prepare_optimizer_step(self): + scope = self._gefen_codebook_process_group + if scope is not None and len(scope.identity.ordered_members) > 1: + should_step = self._prepare_scoped_amp_optimizer_step() + elif hasattr(self, "found_inf") or hasattr(self, "grad_scale"): + should_step = self._prepare_scoped_amp_optimizer_step() + else: + should_step = True + if not should_step: return loss + if ( + self._gefen_codebook_process_group is not None + and self._codebook_refresh_every + ): + self._validate_codebook_scope_operation_header("periodic_step") + self._maybe_refresh_gefen_codebook() self._maybe_save_gefen_grad_histogram() # Periodic codebook re-learn (opt-in): every N steps, refit the exact-DP @@ -5580,11 +8659,7 @@ def step(self, closure=None): # path) instead of any vmean path. Sharded (DTensor) params # fall through to the standard path: local-shard row/col # statistics would be wrong under sharding. - if ( - self._factored_v_2d - and p.ndim == 2 - and not hasattr(p, "placements") - ): + if self._factored_v_2d and p.ndim == 2 and not hasattr(p, "placements"): self._step_automatic_factored(group, name, p, grad) continue if ( diff --git a/src/gefen/gefen_muon.py b/src/gefen/gefen_muon.py index 0d275b1..8330b3d 100644 --- a/src/gefen/gefen_muon.py +++ b/src/gefen/gefen_muon.py @@ -7,6 +7,12 @@ import torch import torch.nn as nn +from gefen.contracts import ( + OptimizerContract, + ParameterLayout, + TopologyChange, + _gefen_muon_contract, +) from gefen.gefen import ( Gefen, _amp_optimizer_step_controls, @@ -755,6 +761,128 @@ def __init__( verbose=verbose, ) + def optimizer_contract(self) -> OptimizerContract: + """Return the immutable Muon state and capability contract.""" + + sharded_modes = frozenset( + group["sharded_mode"] for group in self.param_groups + ) + normuon_modes = frozenset( + group["sharded_mode"] + for group in self.param_groups + if group.get("normuon", False) + ) + non_normuon_modes = frozenset( + group["sharded_mode"] + for group in self.param_groups + if not group.get("normuon", False) + ) + canonical_state_layouts = self._canonical_state_layouts() + try: + from gefen.portable_runtime import _portable_runtime_layouts + + canonical_global_same_topology = _portable_runtime_layouts(self) + except Exception: + canonical_global_same_topology = frozenset() + if canonical_global_same_topology and sharded_modes == frozenset({"distributed"}): + canonical_global_topology_changing = frozenset( + { + ParameterLayout.REPLICATED, + ParameterLayout.WHOLE_PARAMETER_OWNER, + } + ) + canonical_global_topology_change_kinds = frozenset( + { + TopologyChange.PLACEMENT_RESHARD, + TopologyChange.WORLD_SIZE_OWNER_REDISTRIBUTION, + } + ) + else: + canonical_global_topology_changing = frozenset() + canonical_global_topology_change_kinds = frozenset() + return _gefen_muon_contract( + sharded_modes=sharded_modes, + normuon_modes=normuon_modes, + non_normuon_modes=non_normuon_modes, + canonical_parameter_fqns=self._canonical_identity_ready(), + stable_shard_identity=self._canonical_identity_ready(), + explicit_process_group_codebook_scope=True, + canonical_state_layouts=canonical_state_layouts, + canonical_global_same_topology=canonical_global_same_topology, + canonical_global_topology_changing=canonical_global_topology_changing, + canonical_global_topology_change_kinds=canonical_global_topology_change_kinds, + whole_parameter_owner=( + self._codebook_scope_ready() + and any( + shard.layout is ParameterLayout.WHOLE_PARAMETER_OWNER + for _, shard in self._gefen_local_shard_bindings + ) + ), + ) + + def _canonical_state_layout_supported(self, layout) -> bool: + return layout is ParameterLayout.REPLICATED + + def _canonical_state_variant_layout(self): + sharded_modes = frozenset( + group["sharded_mode"] for group in self.param_groups + ) + normuon_modes = frozenset( + group["sharded_mode"] + for group in self.param_groups + if group.get("normuon", False) + ) + non_normuon_modes = frozenset( + group["sharded_mode"] + for group in self.param_groups + if not group.get("normuon", False) + ) + return _gefen_muon_contract( + sharded_modes=sharded_modes, + normuon_modes=normuon_modes, + non_normuon_modes=non_normuon_modes, + ).state_layout + + def _validate_rebinding_layout(self, rebinding) -> None: + shard = rebinding.shard + target = rebinding.new_parameter + if len(shard.parameter.global_shape) != 2: + raise ValueError("GefenMuon canonical parameters must be logical matrices") + if shard.layout is ParameterLayout.REPLICATED: + if target is None or tuple(target.shape) != shard.parameter.global_shape: + raise ValueError( + "replicated GefenMuon rebinding requires one complete 2-D matrix" + ) + elif shard.layout is ParameterLayout.WHOLE_PARAMETER_OWNER: + local_owns = shard.local_member == shard.owner + if local_owns and ( + target is None + or target.ndim != 2 + or tuple(target.shape) != shard.parameter.global_shape + ): + raise ValueError( + "a GefenMuon whole-parameter owner requires one complete 2-D matrix" + ) + if not local_owns and target is not None: + raise ValueError( + "a GefenMuon whole-parameter non-owner must not retain storage" + ) + else: + raise ValueError( + "GefenMuon rebinding supports replicated complete matrices or " + "whole-parameter ownership only" + ) + if target is not None and target.numel() != shard.logical_slice.length: + raise ValueError( + "GefenMuon rebound tensor numel does not match its logical slice" + ) + + def _has_unscoped_whole_owner_bindings(self) -> bool: + return self._gefen_codebook_process_group is None and any( + shard.layout is ParameterLayout.WHOLE_PARAMETER_OWNER + for _, shard in self._gefen_local_shard_bindings + ) + def add_param_group(self, param_group): """Add a validated 2D Muon parameter group atomically. @@ -885,7 +1013,12 @@ def add_param_group(self, param_group): def _init_gefen_muon_state(self, state, grad_view: torch.Tensor) -> None: self._init_gefen_state(state, grad_view) - def _iter_gefen_grad_periods(self, reuse_existing_periods: bool = False): + def _codebook_requires_2d_parameters(self) -> bool: + return True + + def _iter_gefen_grad_periods( + self, reuse_existing_periods: bool = False, staged_periods=None + ): # Same as Gefen._iter_gefen_grad_periods, but for sharded (DTensor) # gradients gather the FULL matrix (full_tensor) instead of taking the # local shard. The exact-DP codebook and the per-param block period are @@ -894,56 +1027,58 @@ def _iter_gefen_grad_periods(self, reuse_existing_periods: bool = False): # matches). With full grads, flat.numel() is the global numel and is # never 0, so every rank iterates every parameter in the same order and # the full_tensor() collective is matched across ranks. - for group in self.param_groups: - for param_name, p in self._iter_group_params_with_names(group): - if p.grad is None: - continue - grad = p.grad - # approx mode learns the codebook/period from the LOCAL shard - # (no all-gather) so periods divide the local numel that the - # approximate step operates on; exact mode gathers the full matrix. - if group["sharded_mode"] == "approx" and hasattr( - grad, "to_local" - ): - grad = grad.to_local() - elif hasattr(grad, "full_tensor"): - grad = grad.full_tensor() - elif hasattr(grad, "to_local"): - grad = grad.to_local() - if hasattr(grad, "wait"): - grad = grad.wait() - grad = grad.detach() - flat = grad.reshape(-1) - if flat.numel() == 0: - continue - - if reuse_existing_periods: - state = self.state[p] - if "automatic_period" not in state: - raise ValueError( - "Expected automatic_period to exist for {} before refreshing Gefen codebook at optimizer step {}".format( - param_name, - self._gefen_global_step, - ) - ) - period = state["automatic_period"] - elif flat.numel() == 1: - period = 1 - else: - period = self._predict_period_from_grad_sq(param_name, p, grad) - - self.state[p]["automatic_period"] = period + for group, param_name, p in self._iter_codebook_params_with_names(): + if p.grad is None: + continue + grad = p.grad + # approx mode learns the codebook/period from the LOCAL shard + # (no all-gather) so periods divide the local numel that the + # approximate step operates on; exact mode gathers the full matrix. + if group["sharded_mode"] == "approx" and hasattr(grad, "to_local"): + grad = grad.to_local() + elif hasattr(grad, "full_tensor"): + grad = grad.full_tensor() + elif hasattr(grad, "to_local"): + grad = grad.to_local() + if hasattr(grad, "wait"): + grad = grad.wait() + grad = grad.detach() + flat = grad.reshape(-1) + if flat.numel() == 0: + continue - if flat.numel() % period != 0: + if reuse_existing_periods: + state = self.state[p] + if "automatic_period" not in state: raise ValueError( - "Automatic partition period {} does not divide parameter {} with numel {} while learning Gefen codebook".format( - period, + "Expected automatic_period to exist for {} before refreshing Gefen codebook at optimizer step {}".format( param_name, - flat.numel(), + self._gefen_global_step, ) ) + period = state["automatic_period"] + elif flat.numel() == 1: + period = 1 + else: + period = self._predict_period_from_grad_sq(param_name, p, grad) + + if flat.numel() % period != 0: + raise ValueError( + "Automatic partition period {} does not divide parameter {} with numel {} while learning Gefen codebook".format( + period, + param_name, + flat.numel(), + ) + ) + + if staged_periods is None: + self.state[p]["automatic_period"] = period + elif not reuse_existing_periods: + staged_periods.append((p, period)) + if not self._codebook_parameter_contributes(p): + continue - yield param_name, flat, period, grad + yield param_name, flat, period, grad def _quantize_momentum_(self, state, momentum_view: torch.Tensor) -> None: period = state["automatic_period"] @@ -2922,17 +3057,75 @@ def _load_state_dict_impl(self, state_dict): @torch.no_grad() def step(self, closure=None): - loss = None + scope_binding = self._capture_codebook_scope_binding_for_step() hybrid_preflight_complete = bool( getattr(self, "_gefen_hybrid_precollective_preflight", False) ) - if not hybrid_preflight_complete: + if hybrid_preflight_complete or scope_binding is None: + self._assert_finalized_binding_layout() + self._assert_runtime_codebook_process_group() + if self._has_unscoped_whole_owner_bindings(): + raise RuntimeError( + "GefenMuon whole-parameter owner stepping requires the separate " + "explicit process-group codebook scope" + ) + loss = None + if not hybrid_preflight_complete and scope_binding is not None: + # The explicit convention binding is the exclusive control scope: + # synchronize on it before any operation header or mesh collective, + # never in addition to the mainline DTensor-derived scope. + try: + self._assert_finalized_binding_layout() + self._assert_runtime_codebook_process_group() + if self._has_unscoped_whole_owner_bindings(): + raise RuntimeError( + "GefenMuon whole-parameter owner stepping requires the separate " + "explicit process-group codebook scope" + ) + self._assert_capturable_if_capturing() + self._assert_codebook_capture_ready() + if closure is not None: + with torch.enable_grad(): + loss = closure() + local_preamble_error = None + except Exception as exc: + loss = None + local_preamble_error = exc + self._synchronize_prevalidated_codebook_scope_failure( + local_preamble_error, "step preamble", scope_binding + ) + + try: + self._assert_finalized_binding_layout() + self._assert_runtime_codebook_process_group() + _assert_optimizer_gradients_structurally_valid( + self, require_2d_params=True + ) + local_preflight_error = None + except Exception as exc: + local_preflight_error = exc + self._synchronize_prevalidated_codebook_scope_failure( + local_preflight_error, "gradient preflight", scope_binding + ) + self._validate_codebook_scope_operation_header("step") + self._ensure_codebook_scope_agreement() + + if not self._prepare_scoped_amp_optimizer_step(): + return loss + elif not hybrid_preflight_complete: process_groups = self._step_failure_process_groups() try: self._assert_capturable_if_capturing() + self._assert_codebook_capture_ready() if closure is not None: with torch.enable_grad(): loss = closure() + # A closure can rebind param_groups (e.g. swap in a rogue + # parameter), so re-assert the finalized binding layout after it + # runs -- before any state mutation -- matching the scoped path + # and the pre-Tier-2 step. The raise is caught below and + # synchronized so every mesh member fails together. + self._assert_finalized_binding_layout() _assert_optimizer_gradients_structurally_valid( self, require_2d_params=True ) diff --git a/src/gefen/hybrid.py b/src/gefen/hybrid.py index 0e2fc2d..04f7db2 100644 --- a/src/gefen/hybrid.py +++ b/src/gefen/hybrid.py @@ -53,16 +53,26 @@ period-one off, stochastic rounding off), with ``ns_schedule="tuned3"`` and ``normuon=True`` as the hybrid-specific defaults. """ + import logging from collections import OrderedDict, defaultdict import torch import torch.nn as nn +from gefen.codebook import CodebookProcessGroupBinding +from gefen.contracts import ( + LogicalSlice, + OptimizerContract, + ParameterIdentity, + ParameterLayout, + ShardIdentity, + ShardingManifest, + _hybrid_contract, +) from gefen.gefen import ( Gefen, _amp_native_scaling_required, - _amp_prepare_optimizer_step, _assert_optimizer_gradients_structurally_valid, ) from gefen.gefen_muon import GefenMuon @@ -72,6 +82,7 @@ split_params_for_muon, validate_split, ) +from gefen.rebinding import ParameterRebinding logger = logging.getLogger(__name__) @@ -335,9 +346,7 @@ def __init__( if backup_named_params is None: # Single-argument convenience form: the first arg is a model or a # named-param iterable to split internally. - muon_named_params, backup_named_params = self._auto_split( - muon_named_params, backup_substrings - ) + muon_named_params, backup_named_params = self._auto_split(muon_named_params, backup_substrings) elif backup_substrings is not None: raise TypeError( "backup_substrings is only valid with the single-argument " @@ -357,9 +366,7 @@ def __init__( # 1D-only model (all params with real storage) still constructs a # backup-only hybrid exactly as before. if not muon_named_params and backup_named_params: - zero_numel = [ - name for name, p in backup_named_params if p.numel() == 0 - ] + zero_numel = [name for name, p in backup_named_params if p.numel() == 0] # Only the ALL-placeholder set is the ZeRO-3 signature; a mixed # set (e.g. an intentionally empty slot among real 1-D weights) # still constructs a backup-only hybrid as before. @@ -381,11 +388,7 @@ def __init__( ) ) if backup_optimizer not in ("gefen", "adamw"): - raise ValueError( - "backup_optimizer must be 'gefen' or 'adamw' but is: {!r}".format( - backup_optimizer - ) - ) + raise ValueError("backup_optimizer must be 'gefen' or 'adamw' but is: {!r}".format(backup_optimizer)) self.backup_optimizer = backup_optimizer # Catch the silent footguns: a param routed to both halves (stepped # twice) or a duplicate name (codebook cache key collision). Completeness @@ -406,9 +409,7 @@ def __init__( "GefenMuonHybrid: adjust_lr_fn={!r} uses Muon-native LR scaling, so " "an AdamW-scale lr will mis-scale the 2D Muon matrices relative to " "the backup half. Use adjust_lr_fn='match_rms_adamw' (the default) " - "to share one AdamW-scale lr, or set muon_lr explicitly.".format( - adjust_lr_fn - ), + "to share one AdamW-scale lr, or set muon_lr explicitly.".format(adjust_lr_fn), stacklevel=2, ) @@ -428,9 +429,7 @@ def __init__( # matrices while leaving the backup (norms/biases/embeddings/head) on a # different schedule. Both default to the shared weight_decay. muon_weight_decay = weight_decay if muon_weight_decay is None else muon_weight_decay - backup_weight_decay = ( - weight_decay if backup_weight_decay is None else backup_weight_decay - ) + backup_weight_decay = weight_decay if backup_weight_decay is None else backup_weight_decay self.muon = ( GefenMuon( @@ -564,6 +563,26 @@ def _is_no_decay(name): for p in group["params"]: self._state_param_owner[id(p)] = (p, o) + # Composite stable identity is installed only after every present Gefen + # child has staged the same complete post-sharding transaction. These + # fields remain separate from either child's rank-local metadata so a + # failed second-child preparation cannot leave the Hybrid half-bound. + self._hybrid_post_sharding_finalized = False + self._hybrid_sharding_manifest = None + self._hybrid_local_shard_bindings = () + self._hybrid_fqn_roles = () + self._hybrid_codebook_process_group = None + self._hybrid_finalized_slots = () + + # Composite finalized-layout forensics cache, mirroring the O(local + # params) scheme Gefen/GefenMuon use for their own step guards. The + # version counter is bumped by the only hybrid API that reassigns any + # of the finalized composite fields (post_sharding, which every rebind + # helper routes through); the cached verdict is one identity-token + # snapshot of everything the composite forensic rebuild reads. + self._hybrid_layout_version = 0 + self._hybrid_layout_forensics_verdict = None + # Deliberately do NOT call super().__init__(): we expose each # sub-optimizer's real param_groups/state via properties (shared dict # refs), so the LR scheduler's in-place ``group["lr"] = ...`` updates @@ -622,16 +641,605 @@ def state(self): # unknown key (DeepSpeed ZeRO's flattened 1-D fp32 partitions) fails # fast with a self-explanatory KeyError instead of silently landing in # a rebuilt-per-access throwaway dict. + self._assert_finalized_binding_layout() return _HybridMergedState(self._subopts, self._state_param_owner) + def _gefen_rebinding_children(self): + children = [] + if self.muon is not None: + if type(self.muon) is not GefenMuon: + raise TypeError("GefenMuonHybrid post_sharding requires an exact GefenMuon child") + children.append(("muon", self.muon)) + if self.backup is not None: + if self.backup_optimizer != "gefen" or type(self.backup) is not Gefen: + raise NotImplementedError( + "GefenMuonHybrid post_sharding does not yet support an AdamW " + "backup; use backup_optimizer='gefen' until AdamW has stable " + "rebinding identity and atomic staged state I/O" + ) + children.append(("backup", self.backup)) + if not children: + raise RuntimeError("GefenMuonHybrid has no child optimizer to rebind") + return tuple(children) + + @staticmethod + def _same_local_binding(left, right) -> bool: + return left[0] is right[0] and left[1] == right[1] + + @staticmethod + def _reject_rebinding_method_shadows(value) -> None: + if type(value.__dict__) is not dict: + raise TypeError("GefenMuonHybrid rebinding requires exact attribute dictionaries") + for name in value.__dict__: + descriptor = None + for owner in type(value).__mro__: + if name in owner.__dict__: + descriptor = owner.__dict__[name] + break + if isinstance(descriptor, (staticmethod, classmethod)): + descriptor = descriptor.__func__ + if callable(descriptor): + raise TypeError("GefenMuonHybrid rebinding rejects instance-level method shadows") + + def _hybrid_identity_metadata_empty(self) -> bool: + return ( + self._hybrid_sharding_manifest is None + and self._hybrid_local_shard_bindings == () + and self._hybrid_fqn_roles == () + and self._hybrid_codebook_process_group is None + and self._hybrid_finalized_slots == () + ) + + def _assert_composite_rebinding_pristine(self, children) -> None: + if self._hybrid_post_sharding_finalized: + raise RuntimeError("GefenMuonHybrid post-sharding identity is already finalized") + if not self._hybrid_identity_metadata_empty(): + raise RuntimeError("GefenMuonHybrid parameter rebinding found an incomplete prior identity plan") + if type(self._subopts) is not list or len(self._subopts) != len(children): + raise RuntimeError("GefenMuonHybrid child optimizer order changed before post_sharding") + if any(live is not expected for live, (_role, expected) in zip(self._subopts, children)): + raise RuntimeError("GefenMuonHybrid child optimizer order changed before post_sharding") + if type(self._state_param_owner) is not dict: + raise TypeError("GefenMuonHybrid parameter ownership must use an exact dictionary") + child_set = {child for _role, child in children} + owner_counts = {child: 0 for child in child_set} + for parameter_id, live in self._state_param_owner.items(): + if ( + type(live) is not tuple + or len(live) != 2 + or not isinstance(live[0], torch.Tensor) + or parameter_id != id(live[0]) + or live[1] not in child_set + ): + raise RuntimeError("GefenMuonHybrid parameter ownership changed before post_sharding") + owner_counts[live[1]] += 1 + for _role, child in children: + live_slot_count = sum(len(group["params"]) for group in child.param_groups) + if owner_counts[child] != live_slot_count: + raise RuntimeError("GefenMuonHybrid child slot counts changed before post_sharding") + + def _stage_post_sharding( + self, + rebindings, + manifest: ShardingManifest, + codebook_process_group=None, + ): + GefenMuonHybrid._reject_rebinding_method_shadows(self) + children = GefenMuonHybrid._gefen_rebinding_children(self) + for _role, child in children: + GefenMuonHybrid._reject_rebinding_method_shadows(child) + GefenMuonHybrid._assert_composite_rebinding_pristine(self, children) + if type(manifest) is not ShardingManifest: + raise TypeError("manifest must be an exact ShardingManifest") + validated_manifest = ShardingManifest( + manifest.shards, + schema_version=manifest.schema_version, + ) + if validated_manifest != manifest: + raise ValueError("GefenMuonHybrid requires a canonical manifest") + if type(rebindings) is not tuple or not rebindings: + raise TypeError("rebindings must be a non-empty tuple of ParameterRebinding values") + if any(type(item) is not ParameterRebinding for item in rebindings): + raise TypeError("rebindings must contain exact ParameterRebinding values") + if codebook_process_group is not None and type(codebook_process_group) is not CodebookProcessGroupBinding: + raise TypeError("codebook_process_group must be an exact CodebookProcessGroupBinding") + + source_entries = {} + source_roles = {} + child_by_role = dict(children) + for parameter_id, (parameter, child) in self._state_param_owner.items(): + if parameter_id != id(parameter): + raise RuntimeError("GefenMuonHybrid parameter ownership contains a stale key") + role = "muon" if child is self.muon else "backup" + if role not in child_by_role or child_by_role[role] is not child: + raise RuntimeError("GefenMuonHybrid parameter ownership names a foreign child") + source_entries[parameter_id] = parameter + source_roles[parameter_id] = role + if len(rebindings) != len(source_entries): + raise ValueError( + "GefenMuonHybrid post_sharding requires exactly one rebinding for every original child slot" + ) + + seen_sources = set() + seen_targets = set() + seen_fqns = set() + targets = [] + by_role = {role: [] for role, _child in children} + fqn_roles = {} + for rebinding in rebindings: + source = rebinding.old_parameter + source_id = id(source) + if source_id not in source_entries or source_entries[source_id] is not source: + raise ValueError("GefenMuonHybrid rebinding source is not an original child slot") + if source_id in seen_sources: + raise ValueError("GefenMuonHybrid rebinding source tensors must be unique") + seen_sources.add(source_id) + target = rebinding.new_parameter + if target is not None: + if not isinstance(target, torch.Tensor): + raise TypeError("GefenMuonHybrid rebinding target must be a Tensor or None") + target_id = id(target) + if target_id in seen_targets: + raise ValueError("GefenMuonHybrid rebound target tensors must be unique") + seen_targets.add(target_id) + original_target = source_entries.get(target_id) + if original_target is not None and original_target is not source: + raise ValueError("GefenMuonHybrid rebound targets cannot steal another original child slot") + targets.append(target) + fqn = rebinding.shard.parameter.fqn + if fqn in seen_fqns: + raise ValueError("GefenMuonHybrid local canonical parameter FQNs must be unique") + seen_fqns.add(fqn) + role = source_roles[source_id] + fqn_roles[fqn] = role + by_role[role].append(rebinding) + + if seen_sources != set(source_entries): + raise ValueError("GefenMuonHybrid post_sharding did not bind every original child slot") + Gefen._assert_rebound_storage_disjoint(targets) + + manifest_fqns = {shard.parameter.fqn for shard in manifest.shards} + if seen_fqns != manifest_fqns: + raise ValueError("GefenMuonHybrid manifest FQNs must exactly match all child slots") + if codebook_process_group is not None: + if any(shard.process_group != codebook_process_group.identity for shard in manifest.shards): + raise ValueError("every Hybrid manifest shard must use the shared codebook process-group identity") + if any(item.shard.local_member != codebook_process_group.local_member for item in rebindings): + raise ValueError("every Hybrid local shard must match the shared codebook member") + + staged_children = [] + for role, child in children: + child_rebindings = tuple(by_role[role]) + if not child_rebindings: + raise ValueError("GefenMuonHybrid child {!r} has no routed rebinding".format(role)) + child_fqns = {item.shard.parameter.fqn for item in child_rebindings} + child_manifest = ShardingManifest( + tuple(shard for shard in manifest.shards if shard.parameter.fqn in child_fqns), + schema_version=manifest.schema_version, + ) + staged = child._stage_post_sharding( + child_rebindings, + child_manifest, + codebook_process_group, + ) + if ( + type(child.__dict__) is not dict + or type(staged.__dict__) is not dict + or staged._gefen_codebook_process_group is not codebook_process_group + or not staged._finalized_binding_layout_matches() + ): + raise TypeError("GefenMuonHybrid child staging produced an unsafe finalized optimizer") + staged_children.append((role, child, staged, child_manifest)) + + local_bindings = [] + local_fqns = set() + new_state_param_owner = {} + finalized_slots = [] + for role, child, staged, _child_manifest in staged_children: + finalized_slots.append( + ( + role, + tuple(tuple(group["params"]) for group in staged.param_groups), + ) + ) + for parameter, shard in staged._gefen_local_shard_bindings: + fqn = shard.parameter.fqn + if fqn in local_fqns or fqn_roles.get(fqn) != role: + raise ValueError("GefenMuonHybrid staged child identities overlap or changed routing") + local_fqns.add(fqn) + local_bindings.append((parameter, shard)) + for group in staged.param_groups: + for parameter in group["params"]: + parameter_id = id(parameter) + if parameter_id in new_state_param_owner: + raise ValueError("GefenMuonHybrid staged children share a live parameter") + new_state_param_owner[parameter_id] = (parameter, child) + if local_fqns != seen_fqns: + raise ValueError("GefenMuonHybrid staged children do not cover the full manifest") + local_bindings.sort(key=lambda item: item[1].sort_key) + if type(self.__dict__) is not dict: + raise TypeError("GefenMuonHybrid publication requires an exact attribute dictionary") + return { + "children": tuple(staged_children), + "state_param_owner": new_state_param_owner, + "manifest": manifest, + "local_bindings": tuple(local_bindings), + "fqn_roles": tuple(sorted(fqn_roles.items())), + "codebook_process_group": codebook_process_group, + "finalized_slots": tuple(finalized_slots), + } + + def post_sharding( + self, + rebindings, + *, + manifest: ShardingManifest, + codebook_process_group=None, + ) -> None: + """Atomically finalize every present Gefen child after sharding.""" + + if isinstance(rebindings, (str, bytes)): + raise TypeError("rebindings must be a sequence") + try: + rebindings = tuple(rebindings) + except TypeError as exc: + raise TypeError("rebindings must be a sequence") from exc + staged = GefenMuonHybrid._stage_post_sharding( + self, + rebindings, + manifest, + codebook_process_group, + ) + for _role, child, staged_child, _child_manifest in staged["children"]: + dict.update(child.__dict__, staged_child.__dict__) + dict.update( + self.__dict__, + { + "_state_param_owner": staged["state_param_owner"], + "_hybrid_post_sharding_finalized": True, + "_hybrid_sharding_manifest": staged["manifest"], + "_hybrid_local_shard_bindings": staged["local_bindings"], + "_hybrid_fqn_roles": staged["fqn_roles"], + "_hybrid_codebook_process_group": staged["codebook_process_group"], + "_hybrid_finalized_slots": staged["finalized_slots"], + # Reassigning the composite fields invalidates any warm verdict + # (there is none from a pristine hybrid, but bumping keeps the + # counter honest and forces the next guard through a full + # rebuild before it caches a fresh verdict). + "_hybrid_layout_version": getattr(self, "_hybrid_layout_version", 0) + + 1, + "_hybrid_layout_forensics_verdict": None, + }, + ) + + def rebind_shard( + self, + old_parameter, + new_parameter, + *, + shard: ShardIdentity, + manifest: ShardingManifest, + ) -> None: + """Apply a complete one-slot Hybrid shard plan.""" + + self.post_sharding( + (ParameterRebinding(old_parameter, new_parameter, shard),), + manifest=manifest, + ) + + def rebind_parameter( + self, + old_parameter, + new_parameter, + *, + identity: ParameterIdentity, + ) -> None: + """Bind one complete replicated parameter on a one-slot Hybrid.""" + + if type(identity) is not ParameterIdentity: + raise TypeError("identity must be an exact ParameterIdentity") + shard = ShardIdentity( + identity, + ParameterLayout.REPLICATED, + LogicalSlice.full(identity), + ) + self.rebind_shard( + old_parameter, + new_parameter, + shard=shard, + manifest=ShardingManifest((shard,)), + ) + + @staticmethod + def _hybrid_child_param_group_tokens(child, tokens): + # Append the child's live group containers, their ``params`` list + # objects, and every parameter in them so an in-place slot swap in a + # child's public param_groups is visible to the composite fast path. + tokens.append(child) + tokens.append(getattr(child, "_gefen_sharding_manifest", None)) + tokens.append(getattr(child, "_gefen_codebook_process_group", None)) + tokens.append(getattr(child, "_gefen_local_shard_bindings", None)) + tokens.append(getattr(child, "defaults", None)) + groups = child.param_groups + tokens.append(groups) + tokens.append(len(groups)) + for group in groups: + params = group.get("params") if type(group) is dict else None + tokens.append(group) + tokens.append(params) + if isinstance(params, (list, tuple)): + tokens.append(len(params)) + tokens.extend(params) + + def _hybrid_layout_forensics_fast_tokens(self): + # O(local params) identity snapshot of everything the composite + # forensic rebuild reads by identity, plus each child's own O(local) + # cached fast-path verdict. Calling the children's fast path means any + # child-level change they can detect (an in-place ``group['params']`` + # slot swap, a replaced private child registry, a mutated name cache) + # flips a bool in this token and forces the composite through a full + # rebuild; the composite-owned fields are captured directly so a + # replaced manifest/roles/local-bindings/slots/owner/binding container + # is caught even if no child noticed. Legitimate mutating APIs replace + # these containers (never mutate them in place) and bump the version + # counter, so an unchanged attribute is the same object. + live = [ + getattr(self, "_hybrid_layout_version", 0), + self._hybrid_post_sharding_finalized, + self._hybrid_sharding_manifest, + self._hybrid_local_shard_bindings, + self._hybrid_fqn_roles, + self._hybrid_finalized_slots, + self._hybrid_codebook_process_group, + self._state_param_owner, + len(self._state_param_owner), + self.defaults, + self._subopts, + len(self._subopts), + ] + for child in self._subopts: + live.append(child._finalized_binding_layout_matches()) + GefenMuonHybrid._hybrid_child_param_group_tokens(child, live) + return tuple(live) + + def _finalized_binding_layout_matches(self, *, full: bool = False) -> bool: + # Steady-state guard: reuse one cached verdict validated by an + # O(local params) identity-token check. Any real finalized-layout + # change either bumps the hybrid version, replaces a composite + # container, or flips a child's own fast-path verdict -- all captured + # by the token -- so a stale hit is impossible; a full boundary check + # (``full=True`` from base contract-readiness call sites) always + # rebuilds. Detection stays before any mutation. + verdict = getattr(self, "_hybrid_layout_forensics_verdict", None) + if not full and verdict is not None: + try: + fast_tokens = self._hybrid_layout_forensics_fast_tokens() + except (AttributeError, KeyError, TypeError, ValueError, RuntimeError): + fast_tokens = None + if fast_tokens is not None and Gefen._forensics_tokens_match( + verdict, fast_tokens + ): + return True + self._hybrid_layout_forensics_verdict = None + matches = self._finalized_binding_layout_matches_full() + if matches: + try: + self._hybrid_layout_forensics_verdict = ( + self._hybrid_layout_forensics_fast_tokens() + ) + except (AttributeError, KeyError, TypeError, ValueError, RuntimeError): + self._hybrid_layout_forensics_verdict = None + return matches + + def _finalized_binding_layout_matches_full(self) -> bool: + try: + if ( + not self._hybrid_post_sharding_finalized + or type(self._hybrid_sharding_manifest) is not ShardingManifest + or type(self._hybrid_local_shard_bindings) is not tuple + or type(self._hybrid_fqn_roles) is not tuple + or type(self._hybrid_finalized_slots) is not tuple + or type(self._state_param_owner) is not dict + ): + return False + children = self._gefen_rebinding_children() + if ( + type(self._subopts) is not list + or len(self._subopts) != len(children) + or any(live is not child for live, (_role, child) in zip(self._subopts, children)) + or self.defaults is not children[0][1].defaults + ): + return False + role_by_fqn = dict(self._hybrid_fqn_roles) + if ( + len(role_by_fqn) != len(self._hybrid_fqn_roles) + or tuple(sorted(role_by_fqn.items())) != self._hybrid_fqn_roles + or any(role not in {"muon", "backup"} for role in role_by_fqn.values()) + ): + return False + manifest_fqns = {shard.parameter.fqn for shard in self._hybrid_sharding_manifest.shards} + if manifest_fqns != set(role_by_fqn): + return False + binding = self._hybrid_codebook_process_group + if binding is not None: + if type(binding) is not CodebookProcessGroupBinding: + return False + if any(shard.process_group != binding.identity for shard in self._hybrid_sharding_manifest.shards): + return False + + expected_local = [] + expected_owner = {} + expected_slots = [] + for role, child in children: + if ( + not child._finalized_binding_layout_matches() + or child._gefen_codebook_process_group is not self._hybrid_codebook_process_group + ): + return False + expected_child_shards = tuple( + shard + for shard in self._hybrid_sharding_manifest.shards + if role_by_fqn.get(shard.parameter.fqn) == role + ) + if child._gefen_sharding_manifest.shards != expected_child_shards: + return False + expected_slots.append( + ( + role, + tuple(tuple(group["params"]) for group in child.param_groups), + ) + ) + for parameter, shard in child._gefen_local_shard_bindings: + if role_by_fqn.get(shard.parameter.fqn) != role: + return False + if binding is not None and ( + shard.process_group != binding.identity or shard.local_member != binding.local_member + ): + return False + expected_local.append((parameter, shard)) + for group in child.param_groups: + for parameter in group["params"]: + parameter_id = id(parameter) + if parameter_id in expected_owner: + return False + expected_owner[parameter_id] = (parameter, child) + expected_local.sort(key=lambda item: item[1].sort_key) + if len(expected_local) != len(self._hybrid_local_shard_bindings): + return False + if any( + not self._same_local_binding(live, expected) + for live, expected in zip( + self._hybrid_local_shard_bindings, + expected_local, + ) + ): + return False + if tuple(expected_slots) != self._hybrid_finalized_slots: + return False + if set(expected_owner) != set(self._state_param_owner): + return False + for parameter_id, (parameter, child) in expected_owner.items(): + live = self._state_param_owner[parameter_id] + if type(live) is not tuple or len(live) != 2 or live[0] is not parameter or live[1] is not child: + return False + return True + except ( + AttributeError, + KeyError, + NotImplementedError, + RuntimeError, + TypeError, + ValueError, + ): + return False + + def _canonical_identity_ready(self) -> bool: + # Contract readiness is an honest external claim, so it never trusts the + # per-step fast-path verdict and always runs the full forensic rebuild. + return self._hybrid_post_sharding_finalized and self._finalized_binding_layout_matches( + full=True + ) + + def _codebook_scope_ready(self) -> bool: + return self._hybrid_codebook_process_group is not None and self._canonical_identity_ready() + + def _assert_finalized_binding_layout(self, *, full: bool = False) -> None: + if self._hybrid_post_sharding_finalized: + if not self._finalized_binding_layout_matches(full=full): + raise RuntimeError("GefenMuonHybrid finalized parameter layout changed outside post_sharding") + elif not self._hybrid_identity_metadata_empty(): + raise RuntimeError("GefenMuonHybrid found an incomplete post_sharding identity plan") + + def parameter_identity(self, parameter) -> ParameterIdentity: + """Return the canonical identity bound to one live Hybrid parameter.""" + + return self.shard_identity(parameter).parameter + + def shard_identity(self, parameter) -> ShardIdentity: + """Return the stable shard identity bound to one live parameter.""" + + self._assert_finalized_binding_layout() + entry = self._state_param_owner.get(id(parameter)) + if entry is None or entry[0] is not parameter: + raise KeyError("parameter has no finalized GefenMuonHybrid shard identity") + return entry[1].shard_identity(parameter) + + def shard_bindings(self): + """Return all local tensor/identity pairs in canonical order.""" + + self._assert_finalized_binding_layout() + return self._hybrid_local_shard_bindings + + def parameter_routing(self): + """Return the immutable canonical FQN-to-child-role routing.""" + + self._assert_finalized_binding_layout() + if not self._canonical_identity_ready(): + raise RuntimeError("GefenMuonHybrid parameter routing is not finalized") + return self._hybrid_fqn_roles + + def sharding_manifest(self): + """Return the complete composite manifest, or ``None``.""" + + self._assert_finalized_binding_layout() + return self._hybrid_sharding_manifest + + def codebook_process_group_binding(self): + """Return the one binding shared by every present Gefen child.""" + + self._assert_finalized_binding_layout() + return self._hybrid_codebook_process_group + + @property + def _gefen_codebook_process_group(self): + # Read-only alias: the scoped-step protocol borrowed from Gefen + # (_prepare_scoped_amp_optimizer_step and the failure synchronization + # it drives) looks the binding up under the child attribute name. The + # composite's one shared binding lives in + # _hybrid_codebook_process_group; nothing may store under this name. + return self._hybrid_codebook_process_group + + def _synchronize_codebook_scope_failure(self, error, phase: str) -> None: + # Composite-level preflight/AMP failures synchronize on the one + # binding shared by every child, so every scope member raises + # together instead of stranding peers inside a child's scoped step + # collectives. The finalized layout guarantees the first child holds + # the identical binding; without a binding, the local error is raised + # unchanged, exactly as the children do. + binding = self._hybrid_codebook_process_group + if binding is None: + if error is not None: + raise error + return + self._subopts[0]._synchronize_codebook_scope_failure(error, phase) + + def _synchronize_prevalidated_codebook_scope_failure( + self, error, phase: str, binding + ) -> None: + """Synchronize through the shared binding validated at step entry.""" + + self._subopts[0]._synchronize_prevalidated_codebook_scope_failure( + error, phase, binding + ) + + def _prepare_scoped_amp_optimizer_step(self) -> bool: + # GradScaler attaches found_inf/grad_scale to the composite, never to + # the children, so the children's own scoped AMP gates cannot see an + # overflow. Run the base scoped protocol once with the composite as + # the optimizer: under a multi-member codebook binding it validates + # found_inf/grad_scale agreement collectively and makes an overflow + # skip a group-wide decision, and on a finite step it unscales the + # union of both children's gradients exactly once. Without a binding + # this is exactly _amp_prepare_optimizer_step(self). + return Gefen._prepare_scoped_amp_optimizer_step(self) + def zero_grad(self, set_to_none: bool = True): + self._assert_finalized_binding_layout() for o in self._subopts: o.zero_grad(set_to_none=set_to_none) def _assert_capturable_devices_if_capturing(self) -> None: - capturing = ( - torch.cuda.is_available() and torch.cuda.is_current_stream_capturing() - ) + capturing = torch.cuda.is_available() and torch.cuda.is_current_stream_capturing() if not capturing: return if not self.capturable: @@ -641,24 +1249,26 @@ def _assert_capturable_devices_if_capturing(self) -> None: "to make step() graph-safe." ) devices = { - param.device - for optimizer in self._subopts - for group in optimizer.param_groups - for param in group["params"] + param.device for optimizer in self._subopts for group in optimizer.param_groups for param in group["params"] } capture_device = torch.device("cuda", torch.cuda.current_device()) if devices != {capture_device}: raise RuntimeError( "CUDA graph capture requires every GefenMuonHybrid parameter on " - "the current capture device {}; found {}".format( - capture_device, sorted(map(str, devices)) - ) + "the current capture device {}; found {}".format(capture_device, sorted(map(str, devices))) ) def step(self, closure=None): + binding = self._hybrid_codebook_process_group + if binding is None: + self._assert_finalized_binding_layout() + else: + if not isinstance(binding, CodebookProcessGroupBinding): + raise TypeError("the runtime codebook process-group binding is invalid") + self._subopts[0]._validate_codebook_runtime_binding(binding) process_groups = ( self.muon._step_failure_process_groups() - if self.muon is not None + if self.muon is not None and binding is None else () ) # Dispatch the INSTANCE step hooks around the composite step, mirroring @@ -669,10 +1279,20 @@ def step(self, closure=None): # _patch_step_function); without this, registered hooks would silently # never fire. GLOBAL step hooks are NOT dispatched here -- they fire on # each child sub-optimizer's (wrapped) step; see the class docstring. + # + # The capturable-device check, the user step pre-hooks, and the closure + # can all fail on only a subset of scope members. Capture such a failure + # and synchronize it on the one shared codebook binding BEFORE any member + # enters the composite preflight synchronize or a child's scoped step + # collectives, so it raises on every member together instead of stranding + # peers. A present runtime binding is validated independently first, so + # the finalized-layout entry guard can join this synchronized phase. args = (self, closure) if closure is not None else (self,) kwargs = {} loss = None try: + if binding is not None: + self._assert_finalized_binding_layout() self._assert_capturable_devices_if_capturing() for pre_hook in self._optimizer_step_pre_hooks.values(): result = pre_hook(self, args, kwargs) @@ -684,39 +1304,69 @@ def step(self, closure=None): f"{self.__class__.__name__}.step pre hook must return None " f"or a tuple of (new_args, new_kwargs), but got {result}." ) - # Re-read the closure from the (possibly hook-rewritten) call args, - # as torch's wrapper would by calling step(*args, **kwargs). + # Re-read the closure from the (possibly hook-rewritten) call args, as + # torch's wrapper would by calling step(*args, **kwargs). closure = args[1] if len(args) > 1 else kwargs.get("closure") if closure is not None: with torch.enable_grad(): loss = closure() - for child in self._subopts: - _assert_optimizer_gradients_structurally_valid( - child, require_2d_params=child is self.muon - ) - local_preflight_error = None + local_preamble_error = None except Exception as exc: loss = None - local_preflight_error = exc - if self.muon is not None: - self.muon._synchronize_sharded_step_error( - local_preflight_error, "hybrid step preflight", process_groups - ) - elif local_preflight_error is not None: - raise local_preflight_error + local_preamble_error = exc - # A non-finite gradient in either half skips BOTH children before their - # codebooks, states, counters, or parameters can move. Explicit - # scaler.unscale_(hybrid) is detected by grad_scale=None and is not - # repeated; automatic unscale covers every child parameter exactly once. - if self.muon is not None: - should_step = self.muon._prepare_synchronized_amp_step( - self, process_groups + if binding is not None: + self._synchronize_prevalidated_codebook_scope_failure( + local_preamble_error, "step preamble", binding ) - elif hasattr(self, "found_inf") or hasattr(self, "grad_scale"): - should_step = _amp_prepare_optimizer_step(self) + try: + self._assert_finalized_binding_layout() + for child in self._subopts: + _assert_optimizer_gradients_structurally_valid( + child, require_2d_params=child is self.muon + ) + local_preflight_error = None + except Exception as exc: + local_preflight_error = exc + self._synchronize_prevalidated_codebook_scope_failure( + local_preflight_error, "gradient preflight", binding + ) + should_step = self._prepare_scoped_amp_optimizer_step() else: - should_step = True + if local_preamble_error is None: + try: + # A closure/pre-hook can rebind param_groups after the entry + # guard (e.g. swap a same-shaped backup parameter), so + # re-assert the finalized layout before stepping children; + # and the muon child's own capture/codebook readiness guard + # is suppressed below by the preflight marker, so run it here. + # Both raises are caught and synchronized so every mesh member + # fails together. + self._assert_finalized_binding_layout() + if self.muon is not None: + self.muon._assert_codebook_capture_ready() + for child in self._subopts: + _assert_optimizer_gradients_structurally_valid( + child, require_2d_params=child is self.muon + ) + local_preflight_error = None + except Exception as exc: + local_preflight_error = exc + else: + local_preflight_error = local_preamble_error + if self.muon is not None: + self.muon._synchronize_sharded_step_error( + local_preflight_error, + "hybrid step preflight", + process_groups, + ) + should_step = self.muon._prepare_synchronized_amp_step( + self, process_groups + ) + else: + if local_preflight_error is not None: + raise local_preflight_error + should_step = self._prepare_scoped_amp_optimizer_step() if not should_step: for post_hook in self._optimizer_step_post_hooks.values(): post_hook(self, args, kwargs) @@ -727,6 +1377,9 @@ def step(self, closure=None): if o is not self.muon: o.step() continue + if binding is not None: + o.step() + continue marker = object() previous = getattr( o, "_gefen_hybrid_precollective_preflight", marker @@ -745,11 +1398,13 @@ def step(self, closure=None): return loss def state_dict(self): + self._assert_finalized_binding_layout() # Instance state-dict pre/post hooks, mirroring Optimizer.state_dict: # pre-hooks take (optimizer) and return nothing; a post-hook may return # a replacement state_dict. for pre_hook in self._optimizer_state_dict_pre_hooks.values(): pre_hook(self) + self._assert_finalized_binding_layout() state_dict = { "muon": self.muon.state_dict() if self.muon is not None else None, "backup": self.backup.state_dict() if self.backup is not None else None, @@ -761,7 +1416,95 @@ def state_dict(self): state_dict = hook_result return state_dict + def optimizer_contract(self) -> OptimizerContract: + """Return the composite contract without flattening either child schema.""" + + try: + GefenMuonHybrid._reject_rebinding_method_shadows(self) + rebinding_ready = bool(GefenMuonHybrid._gefen_rebinding_children(self)) + identity_ready = GefenMuonHybrid._canonical_identity_ready(self) + except Exception: + rebinding_ready = False + identity_ready = False + try: + from gefen.portable_hybrid import _hybrid_portable_contract_support + + ( + canonical_global_same_topology, + canonical_global_topology_changing, + canonical_global_topology_change_kinds, + ) = _hybrid_portable_contract_support(self) + except Exception: + canonical_global_same_topology = frozenset() + canonical_global_topology_changing = frozenset() + canonical_global_topology_change_kinds = frozenset() + muon_contract = self.muon.optimizer_contract() if self.muon is not None else None + backup_contract = ( + self.backup.optimizer_contract() + if self.backup is not None and hasattr(self.backup, "optimizer_contract") + else None + ) + backup_implementation = "" + if self.backup is not None: + backup_implementation = ( + backup_contract.implementation + if backup_contract is not None + else type(self.backup).__module__ + "." + type(self.backup).__qualname__ + ) + return _hybrid_contract( + muon=muon_contract, + backup=backup_contract, + backup_implementation=backup_implementation, + canonical_parameter_fqns=identity_ready, + stable_shard_identity=identity_ready, + explicit_process_group_codebook_scope=rebinding_ready, + shard_rebinding=rebinding_ready, + post_sharding=rebinding_ready, + canonical_global_same_topology=canonical_global_same_topology, + canonical_global_topology_changing=canonical_global_topology_changing, + canonical_global_topology_change_kinds=canonical_global_topology_change_kinds, + ) + + def export_portable_state( + self, + *, + checkpoint_process_group, + transaction_id, + limits, + ): + """Collectively export one complete Gefen-backed composite document.""" + + from gefen.portable_hybrid import _export_hybrid_portable_state + + return _export_hybrid_portable_state( + self, + checkpoint_process_group=checkpoint_process_group, + transaction_id=transaction_id, + limits=limits, + ) + + def import_portable_state( + self, + state, + *, + checkpoint_process_group, + transaction_id, + limits, + ) -> None: + """Collectively stage and atomically publish composite portable state.""" + + from gefen.portable_hybrid import _import_hybrid_portable_state + + _import_hybrid_portable_state( + self, + state, + checkpoint_process_group=checkpoint_process_group, + transaction_id=transaction_id, + limits=limits, + ) + def load_state_dict(self, state_dict): + self._assert_finalized_binding_layout() # Instance load pre-hooks first (a pre-hook may return a replacement # dict -- e.g. one that converts a foreign schema), mirroring # Optimizer.load_state_dict's shallow copy + hook pass. @@ -770,6 +1513,7 @@ def load_state_dict(self, state_dict): hook_result = pre_hook(self, state_dict) if hook_result is not None: state_dict = hook_result + self._assert_finalized_binding_layout() # Schema guard: this used to silently skip loading whenever the keys # were absent, so resuming from a standard {"state", "param_groups"} @@ -810,9 +1554,7 @@ def load_state_dict(self, state_dict): # optimizer loader otherwise accepts the other backend's group/state # layout and fails only on a later step (or, worse, misinterprets state). if self.backup is not None: - checkpoint_backup_optimizer = state_dict.get( - "backup_optimizer", "gefen" - ) + checkpoint_backup_optimizer = state_dict.get("backup_optimizer", "gefen") if checkpoint_backup_optimizer != self.backup_optimizer: raise ValueError( "GefenMuonHybrid.load_state_dict: checkpoint backup_optimizer " @@ -821,25 +1563,6 @@ def load_state_dict(self, state_dict): checkpoint_backup_optimizer, self.backup_optimizer ) ) - # Load the two children with two-phase composite semantics: validate and - # stage BOTH halves before publishing EITHER, so a rejection on either - # half (a different param count, a foreign layout, a corrupted/truncated - # child state) leaves both live children byte-for-byte untouched. The - # previous code committed the muon child first and only then loaded the - # backup, recovering via a second full ``muon.load_state_dict(snapshot)`` - # reload -- a rollback that could itself raise (e.g. CUDA OOM re-staging - # every muon tensor), masking the real backup error and leaving a - # half-loaded hybrid (new muon, old backup). - # - # Each Gefen child already loads atomically via a stage-then-swap - # primitive (``_prepare_load_state_dict`` returns an isolated shadow; - # ``_publish_load_state_dict`` publishes it through non-throwing dict - # swaps). Staging is process-group-safe: ``_stage_load_state_dict`` - # shallow-copies the child's ``__dict__`` and swaps in fresh state - # containers -- it never deepcopies the live optimizer, so a codebook - # process-group handle (or any live ProcessGroup) is shared by reference, - # never duplicated. We deliberately avoid ``copy.deepcopy`` of a child or - # its process groups for exactly that reason. # Phase one -- stage every child (validate, mutating nothing live). Phase # two -- commit every child's raw state through non-throwing swaps. Phase # three -- only then dispatch any child's load post-hooks. Separating the @@ -943,15 +1666,9 @@ def _auto_split(params_or_model, backup_substrings): T5 ``shared``, would slip to Muon, so a model is preferred). Bare tensors raise, since they carry neither names nor module type to route on. """ - subs = ( - DEFAULT_BACKUP_SUBSTRINGS - if backup_substrings is None - else tuple(backup_substrings) - ) + subs = DEFAULT_BACKUP_SUBSTRINGS if backup_substrings is None else tuple(backup_substrings) if isinstance(params_or_model, nn.Module): - muon_named, backup_named = split_params_for_muon( - params_or_model, backup_substrings=subs - ) + muon_named, backup_named = split_params_for_muon(params_or_model, backup_substrings=subs) validate_split(muon_named, backup_named, model=params_or_model) return muon_named, backup_named items = list(params_or_model) @@ -967,9 +1684,7 @@ def _auto_split(params_or_model, backup_substrings): "an iterable of (name, param) pairs (e.g. model.named_parameters()); " "got a bare {}. Bare tensors can't be routed (embeddings/heads " "would go to Muon) -- pass the model or use " - "GefenMuonHybrid.from_model(model, ...).".format( - type(item).__name__ - ) + "GefenMuonHybrid.from_model(model, ...).".format(type(item).__name__) ) muon_named, backup_named = [], [] for name, param in items: @@ -1001,22 +1716,9 @@ def from_model(cls, model, *, backup_substrings=None, **kwargs): return cls(muon_named, backup_named, **kwargs) def add_param_group(self, param_group): - raise NotImplementedError( - "GefenMuonHybrid splits params at construction; add_param_group is unsupported" - ) + raise NotImplementedError("GefenMuonHybrid splits params at construction; add_param_group is unsupported") def __repr__(self): - nm = ( - sum(len(group["params"]) for group in self.muon.param_groups) - if self.muon is not None - else 0 - ) - nb = ( - sum(len(group["params"]) for group in self.backup.param_groups) - if self.backup is not None - else 0 - ) - return ( - f"GefenMuonHybrid(muon_params={nm}, backup_params={nb}, " - f"backup_optimizer={self.backup_optimizer!r})" - ) + nm = sum(len(group["params"]) for group in self.muon.param_groups) if self.muon is not None else 0 + nb = sum(len(group["params"]) for group in self.backup.param_groups) if self.backup is not None else 0 + return f"GefenMuonHybrid(muon_params={nm}, backup_params={nb}, backup_optimizer={self.backup_optimizer!r})" diff --git a/src/gefen/portable.py b/src/gefen/portable.py new file mode 100644 index 0000000..d759a94 --- /dev/null +++ b/src/gefen/portable.py @@ -0,0 +1,518 @@ +"""Strict pure math for topology-independent Gefen optimizer state.""" + +import math + +import torch + + +_PORTABLE_STATE_CHUNK_ELEMENTS = 1 << 20 + + +def _chunk_element_budget() -> int: + budget = _PORTABLE_STATE_CHUNK_ELEMENTS + if type(budget) is not int or budget <= 0: + raise RuntimeError("portable state chunk budget must be a positive int") + return budget + + +def _element_chunks(numel: int): + budget = _chunk_element_budget() + for start in range(0, numel, budget): + yield start, min(start + budget, numel) + + +def _whole_row_chunks(rows: int, width: int): + if rows == 0: + return + rows_per_chunk = max(1, _chunk_element_budget() // max(1, width)) + for start in range(0, rows, rows_per_chunk): + yield start, min(start + rows_per_chunk, rows) + + +def _read_flat_chunk(value: torch.Tensor, start: int, stop: int) -> torch.Tensor: + """Read one logical row-major flat slice without flattening the whole input.""" + + detached = value.detach() + if detached.is_contiguous(): + return detached.reshape(-1)[start:stop] + linear = torch.arange(start, stop, dtype=torch.int64, device=detached.device) + remainder = linear + reversed_coordinates = [] + for dimension in reversed(detached.shape): + reversed_coordinates.append(torch.remainder(remainder, dimension)) + remainder = torch.div(remainder, dimension, rounding_mode="floor") + coordinates = tuple(reversed(reversed_coordinates)) + return detached[coordinates] + + +def _tensor_values_valid(value: torch.Tensor, *, nonnegative: bool) -> tuple[bool, bool]: + finite = True + nonnegative_values = True + for start, stop in _element_chunks(value.numel()): + chunk = _read_flat_chunk(value, start, stop) + try: + if not bool(torch.isfinite(chunk).all()): + finite = False + break + except (NotImplementedError, RuntimeError, TypeError) as exc: + raise TypeError("tensor dtype does not support finite-state validation") from exc + if nonnegative and not bool((chunk >= 0).all()): + nonnegative_values = False + break + return finite, nonnegative_values + + +def _validate_state_counter(value, *, name: str, minimum: int = 0) -> int: + """Return one strict host counter after validating its lower bound.""" + + if type(name) is not str: + raise TypeError("counter name must be a string") + if not name: + raise ValueError("counter name must not be empty") + if type(minimum) is not int: + raise TypeError("counter minimum must be an int") + if minimum < 0: + raise ValueError("counter minimum must be nonnegative") + if type(value) is not int: + raise TypeError("{} must be a host int".format(name)) + if value < minimum: + raise ValueError("{} must be at least {}".format(name, minimum)) + return value + + +def _validate_plain_tensor( + value, + *, + name: str, + dtype: torch.dtype, + ndim: int | None = None, + nonnegative: bool = False, +) -> torch.Tensor: + if type(value) is not torch.Tensor: + raise TypeError("{} must be a plain torch.Tensor".format(name)) + if ( + value.layout is not torch.strided + or value.is_meta + or value.is_nested + or value.is_quantized + ): + raise TypeError("{} must be a materialized strided tensor".format(name)) + if value.dtype != dtype: + raise TypeError("{} must have dtype {}".format(name, dtype)) + if ndim is not None and value.ndim != ndim: + raise ValueError("{} must be {}-D".format(name, ndim)) + try: + finite, nonnegative_values = _tensor_values_valid( + value, + nonnegative=nonnegative, + ) + except TypeError as exc: + raise TypeError("{} dtype does not support finite-state validation".format(name)) from exc + if not finite: + raise ValueError("{} must be finite".format(name)) + if nonnegative and not nonnegative_values: + raise ValueError("{} must be nonnegative".format(name)) + return value + + +def _validate_logical_shape(logical_shape) -> tuple[int, ...]: + if type(logical_shape) not in {tuple, torch.Size}: + raise TypeError("logical_shape must be a tuple or torch.Size") + shape = tuple(logical_shape) + for dimension in shape: + if type(dimension) is not int: + raise TypeError("logical_shape dimensions must be ints") + if dimension < 0: + raise ValueError("logical_shape dimensions must be nonnegative") + return shape + + +def _validate_period(period, *, numel: int) -> int: + if type(period) is not int: + raise TypeError("period must be an int") + if period <= 0: + raise ValueError("period must be positive") + if numel % period != 0: + raise ValueError( + "logical element count {} is not divisible by period {}".format( + numel, + period, + ) + ) + return period + + +def _validate_codebook(codebook) -> torch.Tensor: + codebook = _validate_plain_tensor( + codebook, + name="codebook", + dtype=torch.float32, + ndim=1, + ) + if not 1 <= codebook.numel() <= 256: + raise ValueError("codebook must contain between 1 and 256 entries") + if codebook.numel() > 1 and not bool( + torch.all(codebook.detach()[1:] >= codebook.detach()[:-1]) + ): + raise ValueError("codebook must be sorted in nondecreasing order") + if not bool(((codebook.detach() >= -1.0) & (codebook.detach() <= 1.0)).all()): + raise ValueError("codebook entries must lie in [-1, 1]") + # ``searchsorted`` warns and performs its own hidden copy for a strided + # boundary tensor. A Gefen codebook has at most 256 elements, so normalize + # this bounded input once and reuse it across every logical-state chunk. + return codebook.detach().contiguous() + + +def _new_output(shape, *, dtype: torch.dtype, device: torch.device) -> torch.Tensor: + return torch.empty(shape, dtype=dtype, device=device) + + +def _finish_output(result: torch.Tensor, *, name: str) -> torch.Tensor: + if result.is_floating_point() and not _tensor_values_valid( + result, + nonnegative=False, + )[0]: + raise ValueError("{} cannot be represented as finite {} state".format(name, result.dtype)) + if ( + result.requires_grad + or not result.is_contiguous() + or result.storage_offset() != 0 + or result.untyped_storage().nbytes() != result.numel() * result.element_size() + ): + raise RuntimeError("{} projection did not produce tight detached state".format(name)) + return result + + +def _nearest_codebook_indices( + codebook: torch.Tensor, + normalized_values: torch.Tensor, +) -> torch.Tensor: + """Return nearest codewords, resolving equal distances to the lower index.""" + + if codebook.device != normalized_values.device: + raise ValueError("codebook and normalized values must be on the same device") + codebook = codebook.detach().contiguous() + output = _new_output( + normalized_values.shape, + dtype=torch.uint8, + device=normalized_values.device, + ) + if codebook.numel() == 1: + output.zero_() + return _finish_output(output, name="momentum indices") + + output_flat = output.reshape(-1) + for start, stop in _element_chunks(normalized_values.numel()): + segment = _read_flat_chunk(normalized_values, start, stop).float() + insertion = torch.searchsorted(codebook, segment) + left = (insertion - 1).clamp_(0, codebook.numel() - 1) + right = insertion.clamp_(0, codebook.numel() - 1) + indices = torch.where( + (segment - codebook[left]).abs() + <= (segment - codebook[right]).abs(), + left, + right, + ) + output_flat[start:stop].copy_(indices.to(torch.uint8)) + return _finish_output(output, name="momentum indices") + + +def _decode_quantized_momentum( + codebook, + indices, + magnitudes, + *, + logical_shape, + period, + step, +) -> torch.Tensor: + """Validate and decode block-quantized momentum to dense logical fp32 state.""" + + _validate_state_counter(step, name="step", minimum=1) + shape = _validate_logical_shape(logical_shape) + numel = math.prod(shape) + period = _validate_period(period, numel=numel) + codebook = _validate_codebook(codebook) + indices = _validate_plain_tensor( + indices, + name="momentum indices", + dtype=torch.uint8, + ndim=2, + ) + magnitudes = _validate_plain_tensor( + magnitudes, + name="momentum magnitudes", + dtype=torch.float32, + ndim=2, + nonnegative=True, + ) + if indices.device != codebook.device or magnitudes.device != codebook.device: + raise ValueError("codebook, momentum indices, and magnitudes must share a device") + + blocks = numel // period + if tuple(indices.shape) != (blocks, period): + raise ValueError( + "momentum indices must have shape ({}, {})".format(blocks, period) + ) + if tuple(magnitudes.shape) != (blocks, 1): + raise ValueError("momentum magnitudes must have shape ({}, 1)".format(blocks)) + for start, stop in _element_chunks(indices.numel()): + if int(_read_flat_chunk(indices, start, stop).max().item()) >= codebook.numel(): + raise ValueError("momentum indices contain an out-of-range codebook entry") + + dense = _new_output(shape, dtype=torch.float32, device=codebook.device) + dense_flat = dense.reshape(-1) + for row_start, row_stop in _whole_row_chunks(blocks, period): + flat_start = row_start * period + flat_stop = row_stop * period + index_chunk = _read_flat_chunk(indices, flat_start, flat_stop).long() + coefficients = codebook.index_select(0, index_chunk) + magnitude_chunk = _read_flat_chunk( + magnitudes, + row_start, + row_stop, + ).reshape(-1, 1) + decoded = coefficients.reshape(-1, period) * magnitude_chunk + dense_flat[flat_start:flat_stop].copy_(decoded.reshape(-1)) + return _finish_output(dense, name="dense momentum") + + +def _recompress_dense_momentum( + momentum, + codebook, + *, + period, + step, +) -> tuple[torch.Tensor, torch.Tensor]: + """Compress dense fp32 momentum into one magnitude and byte indices per block.""" + + _validate_state_counter(step, name="step", minimum=1) + momentum = _validate_plain_tensor( + momentum, + name="dense momentum", + dtype=torch.float32, + ) + codebook = _validate_codebook(codebook) + if momentum.device != codebook.device: + raise ValueError("dense momentum and codebook must share a device") + period = _validate_period(period, numel=momentum.numel()) + preserve_zero_sign = ( + period == 1 + and bool(codebook[0] == -1.0) + and bool(codebook[-1] == 1.0) + ) + + blocks = momentum.numel() // period + indices = _new_output( + (blocks, period), + dtype=torch.uint8, + device=momentum.device, + ) + magnitudes = _new_output( + (blocks, 1), + dtype=torch.float32, + device=momentum.device, + ) + for row_start, row_stop in _whole_row_chunks(blocks, period): + flat_start = row_start * period + flat_stop = row_stop * period + block_values = _read_flat_chunk( + momentum, + flat_start, + flat_stop, + ).reshape(-1, period) + magnitude_chunk = block_values.abs().amax(dim=1, keepdim=True) + normalized = block_values.clone() + nonzero = magnitude_chunk > 0 + normalized.div_(magnitude_chunk) + normalized.masked_fill_(~nonzero, 0.0) + index_chunk = _nearest_codebook_indices(codebook, normalized) + if preserve_zero_sign: + zero_values = block_values == 0 + zero_indices = torch.where( + torch.signbit(block_values), + torch.zeros_like(index_chunk), + torch.full_like(index_chunk, codebook.numel() - 1), + ) + index_chunk = torch.where(zero_values, zero_indices, index_chunk) + indices[row_start:row_stop].copy_(index_chunk) + magnitudes[row_start:row_stop].copy_(magnitude_chunk) + return ( + _finish_output(indices, name="momentum indices"), + _finish_output(magnitudes, name="momentum magnitudes"), + ) + + +def _expand_block_second_moment( + block_values, + *, + logical_shape, + period, + step, +) -> torch.Tensor: + """Expand one nonnegative fp32 scalar per block to dense logical state.""" + + _validate_state_counter(step, name="vmean_step", minimum=1) + shape = _validate_logical_shape(logical_shape) + numel = math.prod(shape) + period = _validate_period(period, numel=numel) + block_values = _validate_plain_tensor( + block_values, + name="block second moment", + dtype=torch.float32, + ndim=2, + nonnegative=True, + ) + blocks = numel // period + if tuple(block_values.shape) != (blocks, 1): + raise ValueError("block second moment must have shape ({}, 1)".format(blocks)) + + dense = _new_output(shape, dtype=torch.float32, device=block_values.device) + dense_flat = dense.reshape(-1) + for row_start, row_stop in _whole_row_chunks(blocks, period): + flat_start = row_start * period + flat_stop = row_stop * period + block_chunk = _read_flat_chunk( + block_values, + row_start, + row_stop, + ).reshape(-1, 1) + dense_flat[flat_start:flat_stop].reshape(-1, period).copy_( + block_chunk.expand(-1, period) + ) + return _finish_output(dense, name="dense second moment") + + +def _reduce_block_second_moment( + dense, + *, + period, + step, +) -> torch.Tensor: + """Reduce dense nonnegative fp32 state to target-block arithmetic means.""" + + _validate_state_counter(step, name="vmean_step", minimum=1) + dense = _validate_plain_tensor( + dense, + name="dense second moment", + dtype=torch.float32, + nonnegative=True, + ) + period = _validate_period(period, numel=dense.numel()) + blocks = dense.numel() // period + reduced = _new_output( + (blocks, 1), + dtype=torch.float32, + device=dense.device, + ) + if period == 1: + reduced_flat = reduced.reshape(-1) + for start, stop in _element_chunks(dense.numel()): + reduced_flat[start:stop].copy_( + _read_flat_chunk(dense, start, stop) + ) + return _finish_output(reduced, name="block second moment") + for row_start, row_stop in _whole_row_chunks(blocks, period): + flat_start = row_start * period + flat_stop = row_stop * period + values64 = _read_flat_chunk(dense, flat_start, flat_stop).reshape( + -1, + period, + ).to(torch.float64) + reduced[row_start:row_stop].copy_( + values64.mean(dim=1, keepdim=True).to(torch.float32) + ) + return _finish_output(reduced, name="block second moment") + + +def _expand_factored_second_moment( + row, + column, + *, + logical_shape, + step, +) -> torch.Tensor: + """Expand Adafactor row/column state as ``outer(row, column) / mean(row)``.""" + + _validate_state_counter(step, name="factored_step", minimum=1) + shape = _validate_logical_shape(logical_shape) + if len(shape) != 2: + raise ValueError("factored second moment requires a 2-D logical shape") + row = _validate_plain_tensor( + row, + name="row second moment", + dtype=torch.float32, + ndim=1, + nonnegative=True, + ) + column = _validate_plain_tensor( + column, + name="column second moment", + dtype=torch.float32, + ndim=1, + nonnegative=True, + ) + if row.device != column.device: + raise ValueError("row and column second moments must share a device") + if tuple(row.shape) != (shape[0],) or tuple(column.shape) != (shape[1],): + raise ValueError("factored second-moment vectors do not match logical_shape") + dense = _new_output(shape, dtype=torch.float32, device=row.device) + if math.prod(shape) == 0: + return _finish_output(dense, name="dense second moment") + + row_mean = row.detach().mean(dtype=torch.float64) + if bool(row_mean == 0): + has_nonzero_column = any( + bool((_read_flat_chunk(column, start, stop) != 0).any()) + for start, stop in _element_chunks(column.numel()) + ) + if has_nonzero_column: + raise ValueError( + "zero row mean is inconsistent with a nonzero column second moment" + ) + dense.zero_() + else: + column64 = column.detach().to(torch.float64) + for row_start, row_stop in _whole_row_chunks(shape[0], shape[1]): + row64 = _read_flat_chunk(row, row_start, row_stop).to(torch.float64) + expanded = torch.outer(row64, column64).div_(row_mean).to(torch.float32) + dense[row_start:row_stop].copy_(expanded) + return _finish_output(dense, name="dense second moment") + + +def _project_factored_second_moment( + dense, + *, + step, +) -> tuple[torch.Tensor, torch.Tensor]: + """Project dense 2-D second moment to its row and column arithmetic means.""" + + _validate_state_counter(step, name="factored_step", minimum=1) + dense = _validate_plain_tensor( + dense, + name="dense second moment", + dtype=torch.float32, + ndim=2, + nonnegative=True, + ) + rows, columns = dense.shape + row = _new_output((rows,), dtype=torch.float32, device=dense.device) + column = _new_output((columns,), dtype=torch.float32, device=dense.device) + if dense.numel() == 0: + row.zero_() + column.zero_() + else: + for matrix, output in ((dense.detach(), row), (dense.detach().t(), column)): + matrix_rows, width = matrix.shape + for row_start, row_stop in _whole_row_chunks(matrix_rows, width): + values64 = matrix[row_start:row_stop].to(torch.float64) + output[row_start:row_stop].copy_( + values64.mean(dim=1).to(torch.float32) + ) + return ( + _finish_output(row, name="row second moment"), + _finish_output(column, name="column second moment"), + ) + + +__all__ = [] diff --git a/src/gefen/portable_collective.py b/src/gefen/portable_collective.py new file mode 100644 index 0000000..246934a --- /dev/null +++ b/src/gefen/portable_collective.py @@ -0,0 +1,826 @@ +"""Symmetric tensor-only collectives for bounded canonical-state fragments. + +This module deliberately keeps the transport private. Callers stage values in +callbacks and publish them only after the complete operation returns. +""" + +from __future__ import annotations + +from dataclasses import fields +import hashlib +import hmac +import struct +import sys + +import torch + +from gefen.checkpoint import CheckpointProcessGroupBinding +from gefen.portable_wire import ( + _CanonicalWireLimits, + _CanonicalWirePlan, + _PreparedCanonicalWireValue, + _parse_canonical_wire_metadata, + _prepare_canonical_wire_value, + _reconstruct_canonical_wire_value, +) + + +_I64_MAX = (1 << 63) - 1 +_HEADER_MAGIC = 0x47464E434F4C4C31 # ``GFNCOLL1`` +_VOTE_MAGIC = 0x47464E564F544531 # ``GFNVOTE1`` +_PROTOCOL_VERSION = 1 +_DIGEST_FIELDS = 4 +_DIAGNOSTIC_BYTES = 2048 +_LABEL_BYTES = 1024 +_LIMIT_NAMES = tuple(field.name for field in fields(_CanonicalWireLimits)) + +_H_MAGIC = 0 +_H_VERSION = 1 +_H_STATUS = 2 +_H_COORDINATE = 3 +_H_MEMBERS = 4 +_H_METADATA_BYTES = 5 +_H_TENSOR_COUNT = 6 +_H_TENSOR_BYTES = 7 +_H_TRANSPORT = 8 +_H_OPERATION_BYTES = 9 +_H_TRANSACTION_BYTES = 10 +_H_CALLBACK_FLAGS = 11 +_H_OPERATION_DIGEST = 12 +_H_TRANSACTION_DIGEST = _H_OPERATION_DIGEST + _DIGEST_FIELDS +_H_CONTEXT_DIGEST = _H_TRANSACTION_DIGEST + _DIGEST_FIELDS +_H_IDENTITY_DIGEST = _H_CONTEXT_DIGEST + _DIGEST_FIELDS +_H_METADATA_DIGEST = _H_IDENTITY_DIGEST + _DIGEST_FIELDS +_H_FRAGMENT_DIGEST = _H_METADATA_DIGEST + _DIGEST_FIELDS +_H_LIMITS = _H_FRAGMENT_DIGEST + _DIGEST_FIELDS +_HEADER_FIELDS = _H_LIMITS + len(_LIMIT_NAMES) + +_V_MAGIC = 0 +_V_VERSION = 1 +_V_PHASE = 2 +_V_SOURCE = 3 +_V_STATUS = 4 +_V_COORDINATE = 5 +_V_RESERVED_0 = 6 +_V_RESERVED_1 = 7 +_VOTE_FIELDS = 8 + +_PHASE_METADATA_BUFFER = 1 +_PHASE_METADATA_PARSE = 2 +_PHASE_PAYLOAD_ALLOCATION = 3 +_PHASE_CHUNK_PREPARE = 4 +_PHASE_CHUNK_STAGE = 5 +_PHASE_RECONSTRUCT = 6 +_PHASE_CONSUME = 7 + + +class _CanonicalCollectiveError(RuntimeError): + """A deterministic, live-member-wide canonical transport rejection.""" + + +def _require_signed_header_int(name: str, value: int) -> None: + if type(value) is not int: + raise TypeError("{} must be an int".format(name)) + if value < 0 or value > _I64_MAX: + raise ValueError("{} must fit a nonnegative signed int64 header field".format(name)) + + +def _label_bytes(name: str, value: str) -> bytes: + if type(value) is not str: + raise TypeError("{} must be a string".format(name)) + if not value or value != value.strip() or "\x00" in value: + raise ValueError("{} must be nonempty, trimmed, and contain no NUL".format(name)) + encoded = value.encode("utf-8") + if len(encoded) > _LABEL_BYTES: + raise ValueError("{} exceeds {} UTF-8 bytes".format(name, _LABEL_BYTES)) + return encoded + + +def _digest_fields(value: bytes) -> tuple[int, int, int, int]: + if type(value) is not bytes or len(value) != hashlib.sha256().digest_size: + raise ValueError("digest header fields require exactly 32 bytes") + return tuple( + int.from_bytes(value[start : start + 8], "big", signed=True) + for start in range(0, len(value), 8) + ) + + +def _fields_digest(values: tuple[int, ...] | list[int]) -> bytes: + if len(values) != _DIGEST_FIELDS: + raise ValueError("a digest requires four signed int64 fields") + return b"".join(int(value).to_bytes(8, "big", signed=True) for value in values) + + +def _identity_digest(binding: CheckpointProcessGroupBinding) -> bytes: + identity = binding.identity + hasher = hashlib.sha256() + hasher.update(b"gefen.canonical_collective.identity.v1\0") + hasher.update(struct.pack(">Q", identity.schema_version)) + semantic_name = identity.semantic_name.encode("utf-8") + hasher.update(struct.pack(">Q", len(semantic_name))) + hasher.update(semantic_name) + hasher.update(struct.pack(">Q", len(identity.ordered_members))) + for member in identity.ordered_members: + encoded = member.encode("utf-8") + hasher.update(struct.pack(">Q", len(encoded))) + hasher.update(encoded) + return hasher.digest() + + +def _transport_code(binding: CheckpointProcessGroupBinding) -> int: + if sys.byteorder == "little": + byteorder = 0 + elif sys.byteorder == "big": + byteorder = 1 + else: + raise RuntimeError("unsupported native byte order") + return (0 if binding.collective_device.type == "cpu" else 2) + byteorder + + +def _diagnostic(exc: Exception | None, *, limit: int = _DIAGNOSTIC_BYTES) -> bytes: + if exc is None: + return bytes(_DIAGNOSTIC_BYTES) + usable = max(1, min(limit, _DIAGNOSTIC_BYTES)) + message = "{}: {}".format(type(exc).__name__, str(exc)).replace("\x00", "\\0") + encoded = message.encode("utf-8", errors="replace")[: usable - 1] + return encoded + bytes(_DIAGNOSTIC_BYTES - len(encoded)) + + +def _decode_diagnostic(value: bytes) -> str: + return value.split(b"\0", 1)[0].decode("utf-8", errors="replace") or "unspecified failure" + + +def _runtime_size(binding: CheckpointProcessGroupBinding) -> int: + if len(binding.identity.ordered_members) == 1: + return 1 + import torch.distributed as dist + + if not dist.is_available() or not dist.is_initialized(): + raise RuntimeError("a multi-member canonical collective requires initialized torch.distributed") + return dist.get_world_size(binding.process_group) + + +def _all_gather_fixed( + binding: CheckpointProcessGroupBinding, + values: torch.Tensor, + *, + member_count: int, +) -> tuple[torch.Tensor, ...]: + if member_count == 1: + return (values.detach().cpu(),) + import torch.distributed as dist + + collective = values.to(binding.collective_device) + gathered = [torch.empty_like(collective) for _ in range(member_count)] + dist.all_gather(gathered, collective, group=binding.process_group) + return tuple(value.cpu() for value in gathered) + + +def _exchange_header( + binding: CheckpointProcessGroupBinding, + header: tuple[int, ...], + diagnostic: bytes, + *, + member_count: int, +) -> tuple[tuple[tuple[int, ...], ...], tuple[bytes, ...]]: + header_tensor = torch.tensor(header, dtype=torch.int64) + diagnostic_tensor = torch.frombuffer(bytearray(diagnostic), dtype=torch.uint8) + gathered_headers = _all_gather_fixed(binding, header_tensor, member_count=member_count) + gathered_diagnostics = _all_gather_fixed(binding, diagnostic_tensor, member_count=member_count) + return ( + tuple(tuple(int(item) for item in value.tolist()) for value in gathered_headers), + tuple(bytes(value.numpy().tobytes()) for value in gathered_diagnostics), + ) + + +def _error_from_diagnostics( + phase: str, + diagnostics: tuple[bytes, ...], + failed: tuple[int, ...], +) -> _CanonicalCollectiveError: + details = "; ".join( + "semantic-coordinate[{}]: {}".format(coordinate, _decode_diagnostic(diagnostics[coordinate])) + for coordinate in failed + ) + return _CanonicalCollectiveError("canonical collective {} failed: {}".format(phase, details)) + + +def _phase_vote( + binding: CheckpointProcessGroupBinding, + *, + member_count: int, + coordinate: int, + source: int, + phase: int, + phase_name: str, + error: Exception | None, + diagnostic_limit: int, +) -> None: + vote = [0] * _VOTE_FIELDS + vote[_V_MAGIC] = _VOTE_MAGIC + vote[_V_VERSION] = _PROTOCOL_VERSION + vote[_V_PHASE] = phase + vote[_V_SOURCE] = source + vote[_V_STATUS] = 1 if error is not None else 0 + vote[_V_COORDINATE] = coordinate + headers, diagnostics = _exchange_header( + binding, + tuple(vote), + _diagnostic(error, limit=diagnostic_limit), + member_count=member_count, + ) + for expected_coordinate, header in enumerate(headers): + if len(header) != _VOTE_FIELDS or header[_V_MAGIC] != _VOTE_MAGIC or header[_V_VERSION] != _PROTOCOL_VERSION: + raise _CanonicalCollectiveError("canonical collective received an invalid fixed-size phase vote") + if header[_V_PHASE] != phase or header[_V_SOURCE] != source: + raise _CanonicalCollectiveError("canonical collective phase/source order diverged") + if header[_V_COORDINATE] != expected_coordinate or header[_V_RESERVED_0] != 0 or header[_V_RESERVED_1] != 0: + raise _CanonicalCollectiveError("canonical collective phase vote has an invalid semantic coordinate") + if header[_V_STATUS] not in {0, 1}: + raise _CanonicalCollectiveError("canonical collective phase vote has an invalid status") + failed = tuple(index for index, header in enumerate(headers) if header[_V_STATUS]) + if failed: + raise _error_from_diagnostics( + phase_name, + diagnostics, + failed, + ) + + +def _failure_header(coordinate: int, member_count: int) -> tuple[int, ...]: + header = [0] * _HEADER_FIELDS + header[_H_MAGIC] = _HEADER_MAGIC + header[_H_VERSION] = _PROTOCOL_VERSION + header[_H_STATUS] = 1 + header[_H_COORDINATE] = coordinate + header[_H_MEMBERS] = member_count + return tuple(header) + + +def _success_header( + binding: CheckpointProcessGroupBinding, + plan: _CanonicalWirePlan | None, + limits: _CanonicalWireLimits, + *, + coordinate: int, + member_count: int, + operation: bytes, + transaction: bytes, + context_digest: bytes, + callback_flags: int, +) -> tuple[int, ...]: + if plan is not None and type(plan) is not _CanonicalWirePlan: + raise TypeError("plan must be a _CanonicalWirePlan or None") + metadata_bytes = len(plan.metadata) if plan is not None else 0 + tensor_count = len(plan.tensor_specs) if plan is not None else 0 + tensor_bytes = plan.total_tensor_bytes if plan is not None else 0 + for name in _LIMIT_NAMES: + _require_signed_header_int("limits.{}".format(name), getattr(limits, name)) + if limits.diagnostic_bytes > _DIAGNOSTIC_BYTES: + raise ValueError("limits.diagnostic_bytes exceeds the collective diagnostic capacity") + for name, value in ( + ("member_count", member_count), + ("coordinate", coordinate), + ("metadata bytes", metadata_bytes), + ("tensor count", tensor_count), + ("tensor bytes", tensor_bytes), + ("operation bytes", len(operation)), + ("transaction bytes", len(transaction)), + ("callback flags", callback_flags), + ): + _require_signed_header_int(name, value) + if member_count > limits.max_members: + raise ValueError("checkpoint member count exceeds limits.max_members") + header = [0] * _HEADER_FIELDS + header[_H_MAGIC] = _HEADER_MAGIC + header[_H_VERSION] = _PROTOCOL_VERSION + header[_H_STATUS] = 0 + header[_H_COORDINATE] = coordinate + header[_H_MEMBERS] = member_count + header[_H_METADATA_BYTES] = metadata_bytes + header[_H_TENSOR_COUNT] = tensor_count + header[_H_TENSOR_BYTES] = tensor_bytes + header[_H_TRANSPORT] = _transport_code(binding) + header[_H_OPERATION_BYTES] = len(operation) + header[_H_TRANSACTION_BYTES] = len(transaction) + if callback_flags > 3: + raise ValueError("callback flags contain unknown bits") + header[_H_CALLBACK_FLAGS] = callback_flags + for offset, digest in ( + (_H_OPERATION_DIGEST, hashlib.sha256(b"operation\0" + operation).digest()), + (_H_TRANSACTION_DIGEST, hashlib.sha256(b"transaction\0" + transaction).digest()), + (_H_CONTEXT_DIGEST, context_digest), + (_H_IDENTITY_DIGEST, _identity_digest(binding)), + (_H_METADATA_DIGEST, plan.metadata_digest if plan is not None else bytes(32)), + (_H_FRAGMENT_DIGEST, plan.fragment_digest if plan is not None else bytes(32)), + ): + header[offset : offset + _DIGEST_FIELDS] = _digest_fields(digest) + for index, name in enumerate(_LIMIT_NAMES): + header[_H_LIMITS + index] = getattr(limits, name) + return tuple(header) + + +def _validate_initial_exchange( + binding: CheckpointProcessGroupBinding, + headers: tuple[tuple[int, ...], ...], + diagnostics: tuple[bytes, ...], + *, + member_count: int, +) -> None: + for expected_coordinate, header in enumerate(headers): + if len(header) != _HEADER_FIELDS or header[_H_MAGIC] != _HEADER_MAGIC or header[_H_VERSION] != _PROTOCOL_VERSION: + raise _CanonicalCollectiveError("canonical collective received an invalid fixed-size header") + if header[_H_STATUS] not in {0, 1}: + raise _CanonicalCollectiveError("canonical collective header has an invalid status") + if header[_H_COORDINATE] != expected_coordinate or header[_H_MEMBERS] != member_count: + raise _CanonicalCollectiveError("canonical collective header has an invalid semantic coordinate/order") + if header[_H_TRANSPORT] not in {0, 1, 2, 3} or header[_H_CALLBACK_FLAGS] not in {0, 1, 2, 3}: + raise _CanonicalCollectiveError("canonical collective header has invalid transport/callback flags") + failed = tuple(index for index, header in enumerate(headers) if header[_H_STATUS]) + if failed: + raise _error_from_diagnostics( + "preparation", + diagnostics, + failed, + ) + consensus_ranges = ( + (_H_OPERATION_BYTES, _H_CONTEXT_DIGEST), + (_H_CONTEXT_DIGEST, _H_METADATA_DIGEST), + (_H_LIMITS, _HEADER_FIELDS), + ) + reference = headers[0] + if any( + header[start:stop] != reference[start:stop] + for header in headers[1:] + for start, stop in consensus_ranges + ): + raise _CanonicalCollectiveError( + "canonical collective operation, transaction, context, identity, callback configuration, or limits diverged" + ) + limits = reference[_H_LIMITS:_HEADER_FIELDS] + limit_by_name = dict(zip(_LIMIT_NAMES, limits)) + total_metadata = 0 + total_tensors = 0 + for header in headers: + metadata_bytes = header[_H_METADATA_BYTES] + tensor_count = header[_H_TENSOR_COUNT] + tensor_bytes = header[_H_TENSOR_BYTES] + for name, value in ( + ("metadata bytes", metadata_bytes), + ("tensor count", tensor_count), + ("tensor bytes", tensor_bytes), + ): + _require_signed_header_int(name, value) + if metadata_bytes > limit_by_name["max_metadata_bytes"]: + raise _CanonicalCollectiveError("canonical collective fragment metadata exceeds max_metadata_bytes") + if tensor_count > limit_by_name["max_tensors"]: + raise _CanonicalCollectiveError("canonical collective fragment exceeds max_tensors") + if tensor_bytes > limit_by_name["max_fragment_tensor_bytes"]: + raise _CanonicalCollectiveError("canonical collective fragment exceeds max_fragment_tensor_bytes") + if total_metadata > limit_by_name["max_collective_metadata_bytes"] - metadata_bytes: + raise _CanonicalCollectiveError("canonical collective metadata exceeds max_collective_metadata_bytes") + if total_tensors > limit_by_name["max_collective_tensor_bytes"] - tensor_bytes: + raise _CanonicalCollectiveError("canonical collective tensors exceed max_collective_tensor_bytes") + total_metadata += metadata_bytes + total_tensors += tensor_bytes + + +def _broadcast_tensor( + binding: CheckpointProcessGroupBinding, + value: torch.Tensor, + *, + source: int, + member_count: int, +) -> None: + if member_count == 1: + return + import torch.distributed as dist + + global_source = dist.get_global_rank(binding.process_group, source) + dist.broadcast(value, src=global_source, group=binding.process_group) + + +def _metadata_from_tensor(value: torch.Tensor) -> bytes: + return bytes(value.detach().cpu().numpy().tobytes()) + + +def _canonical_byte_indices(start: int, stop: int, component_bytes: int) -> torch.Tensor: + canonical = torch.arange(start, stop, dtype=torch.int64) + base = torch.div(canonical, component_bytes, rounding_mode="floor") * component_bytes + return base + (component_bytes - 1 - torch.remainder(canonical, component_bytes)) + + +def _flat_native_bytes(value: torch.Tensor) -> torch.Tensor: + return value.reshape(-1).view(torch.uint8).reshape(-1) + + +def _canonical_cpu_chunk(value: torch.Tensor, start: int, stop: int) -> torch.Tensor: + raw = _flat_native_bytes(value) + if sys.byteorder == "little": + return raw[start:stop].clone() + if sys.byteorder != "big": + raise RuntimeError("unsupported native byte order") + component_bytes = value.element_size() // 2 if value.is_complex() else value.element_size() + if component_bytes == 1: + return raw[start:stop].clone() + return raw.index_select(0, _canonical_byte_indices(start, stop, component_bytes)) + + +def _store_canonical_cpu_chunk( + destination: torch.Tensor, + start: int, + stop: int, + canonical: torch.Tensor, +) -> None: + raw = _flat_native_bytes(destination) + if sys.byteorder == "little": + raw[start:stop].copy_(canonical) + return + if sys.byteorder != "big": + raise RuntimeError("unsupported native byte order") + component_bytes = destination.element_size() // 2 if destination.is_complex() else destination.element_size() + if component_bytes == 1: + raw[start:stop].copy_(canonical) + return + raw.index_copy_(0, _canonical_byte_indices(start, stop, component_bytes), canonical) + + +def _header_digest(header: tuple[int, ...], offset: int) -> bytes: + return _fields_digest(header[offset : offset + _DIGEST_FIELDS]) + + +def _receive_metadata( + binding: CheckpointProcessGroupBinding, + local_plan: _CanonicalWirePlan, + headers: tuple[tuple[int, ...], ...], + limits: _CanonicalWireLimits, + *, + coordinate: int, + source: int, + member_count: int, + validate_plan, +) -> _PreparedCanonicalWireValue: + header = headers[source] + metadata_bytes = header[_H_METADATA_BYTES] + metadata_buffer = None + error = None + try: + if coordinate == source: + metadata_buffer = torch.frombuffer(bytearray(local_plan.metadata), dtype=torch.uint8).to( + binding.collective_device + ) + else: + metadata_buffer = torch.empty( + metadata_bytes, + dtype=torch.uint8, + device=binding.collective_device, + ) + if metadata_buffer.numel() != metadata_bytes: + raise ValueError("canonical collective metadata buffer has an invalid size") + except Exception as exc: + error = exc + _phase_vote( + binding, + member_count=member_count, + coordinate=coordinate, + source=source, + phase=_PHASE_METADATA_BUFFER, + phase_name="metadata-buffer allocation", + error=error, + diagnostic_limit=limits.diagnostic_bytes, + ) + assert metadata_buffer is not None + _broadcast_tensor(binding, metadata_buffer, source=source, member_count=member_count) + + prepared = None + error = None + try: + metadata = _metadata_from_tensor(metadata_buffer) + if not hmac.compare_digest(hashlib.sha256(metadata).digest(), _header_digest(header, _H_METADATA_DIGEST)): + raise ValueError("canonical collective metadata digest does not match its source header") + prepared = _parse_canonical_wire_metadata(metadata, limits=limits) + if len(prepared.tensor_specs) != header[_H_TENSOR_COUNT]: + raise ValueError("canonical collective tensor count does not match its source header") + if prepared.total_tensor_bytes != header[_H_TENSOR_BYTES]: + raise ValueError("canonical collective tensor bytes do not match its source header") + if not hmac.compare_digest(prepared.fragment_digest, _header_digest(header, _H_FRAGMENT_DIGEST)): + raise ValueError("canonical collective fragment digest does not match its source header") + if validate_plan is not None: + validate_plan(binding.identity.ordered_members[source], prepared) + except Exception as exc: + error = exc + _phase_vote( + binding, + member_count=member_count, + coordinate=coordinate, + source=source, + phase=_PHASE_METADATA_PARSE, + phase_name="metadata validation", + error=error, + diagnostic_limit=limits.diagnostic_bytes, + ) + assert prepared is not None + return prepared + + +def _allocate_payloads( + binding: CheckpointProcessGroupBinding, + prepared: _PreparedCanonicalWireValue, + limits: _CanonicalWireLimits, + *, + coordinate: int, + source: int, + member_count: int, +) -> tuple[torch.Tensor, ...]: + payloads = None + error = None + try: + payloads = tuple( + torch.empty(spec.shape, dtype=spec.dtype, device="cpu") + for spec in prepared.tensor_specs + ) + except Exception as exc: + error = exc + _phase_vote( + binding, + member_count=member_count, + coordinate=coordinate, + source=source, + phase=_PHASE_PAYLOAD_ALLOCATION, + phase_name="payload allocation", + error=error, + diagnostic_limit=limits.diagnostic_bytes, + ) + assert payloads is not None + return payloads + + +def _transfer_payloads( + binding: CheckpointProcessGroupBinding, + local_plan: _CanonicalWirePlan, + prepared: _PreparedCanonicalWireValue, + payloads: tuple[torch.Tensor, ...], + limits: _CanonicalWireLimits, + *, + coordinate: int, + source: int, + member_count: int, + direct_cpu: bool, +) -> None: + for tensor_index, (spec, destination) in enumerate(zip(prepared.tensor_specs, payloads)): + for start in range(0, spec.nbytes, limits.chunk_bytes): + stop = min(start + limits.chunk_bytes, spec.nbytes) + wire_buffer = None + error = None + try: + if direct_cpu: + wire_buffer = _flat_native_bytes(destination)[start:stop] + if coordinate == source: + source_bytes = _flat_native_bytes(local_plan.tensors[tensor_index]) + wire_buffer.copy_(source_bytes[start:stop]) + else: + wire_buffer = torch.empty( + stop - start, + dtype=torch.uint8, + device=binding.collective_device, + ) + if coordinate == source: + canonical = _canonical_cpu_chunk(local_plan.tensors[tensor_index], start, stop) + wire_buffer.copy_(canonical) + except Exception as exc: + error = exc + _phase_vote( + binding, + member_count=member_count, + coordinate=coordinate, + source=source, + phase=_PHASE_CHUNK_PREPARE, + phase_name="payload chunk {}:{} preparation".format(tensor_index, start), + error=error, + diagnostic_limit=limits.diagnostic_bytes, + ) + assert wire_buffer is not None + _broadcast_tensor(binding, wire_buffer, source=source, member_count=member_count) + if not direct_cpu: + error = None + try: + canonical = wire_buffer.detach().cpu() + _store_canonical_cpu_chunk(destination, start, stop, canonical) + except Exception as exc: + error = exc + _phase_vote( + binding, + member_count=member_count, + coordinate=coordinate, + source=source, + phase=_PHASE_CHUNK_STAGE, + phase_name="payload chunk {}:{} staging".format(tensor_index, start), + error=error, + diagnostic_limit=limits.diagnostic_bytes, + ) + + +def _collective_unanimous_status( + binding: CheckpointProcessGroupBinding, + local_error: Exception | None, + *, + operation: str, + transaction_id: str, + context_digest: bytes, + limits: _CanonicalWireLimits, +) -> None: + """Exchange one fixed-size pre-publication status across all live members. + + This is the lightweight companion to the fragment visitor for import + readiness and freshness gates. It transfers only fixed-size status/header + tensors and diagnostics; it never broadcasts canonical metadata or payloads. + """ + + if not isinstance(binding, CheckpointProcessGroupBinding): + raise TypeError("binding must be a CheckpointProcessGroupBinding") + member_count = _runtime_size(binding) + coordinate = binding.identity.ordered_members.index(binding.local_member) + header = _failure_header(coordinate, member_count) + error = None + diagnostic_limit = _DIAGNOSTIC_BYTES + try: + binding.validate_runtime() + if type(limits) is not _CanonicalWireLimits: + raise TypeError("limits must be _CanonicalWireLimits") + diagnostic_limit = limits.diagnostic_bytes + if local_error is not None and not isinstance(local_error, Exception): + raise TypeError("local_error must be an Exception or None") + operation_bytes = _label_bytes("operation", operation) + transaction_bytes = _label_bytes("transaction_id", transaction_id) + if type(context_digest) is not bytes or len(context_digest) != hashlib.sha256().digest_size: + raise ValueError("context_digest must be exactly 32 bytes") + header = _success_header( + binding, + None, + limits, + coordinate=coordinate, + member_count=member_count, + operation=operation_bytes, + transaction=transaction_bytes, + context_digest=context_digest, + callback_flags=0, + ) + if local_error is not None: + error = local_error + header = _failure_header(coordinate, member_count) + except Exception as exc: + error = exc + headers, diagnostics = _exchange_header( + binding, + header, + _diagnostic(error, limit=diagnostic_limit), + member_count=member_count, + ) + _validate_initial_exchange( + binding, + headers, + diagnostics, + member_count=member_count, + ) + + +def _collective_visit_canonical_fragments( + binding: CheckpointProcessGroupBinding, + local_fragment, + *, + operation: str, + transaction_id: str, + context_digest: bytes, + limits: _CanonicalWireLimits, + validate_plan=None, + consume=None, +) -> None: + """Visit one staged canonical fragment from every semantic member in order. + + ``validate_plan(member, prepared)`` runs after bounded metadata validation + and before payload allocation. ``consume(member, value)`` receives an owned, + digest-verified value. Callbacks must stage only reversible local work; this + transport does not publish or roll back callback side effects. + """ + + if not isinstance(binding, CheckpointProcessGroupBinding): + raise TypeError("binding must be a CheckpointProcessGroupBinding") + member_count = _runtime_size(binding) + members = binding.identity.ordered_members + coordinate = members.index(binding.local_member) + plan = None + header = _failure_header(coordinate, member_count) + error = None + diagnostic_limit = _DIAGNOSTIC_BYTES + try: + binding.validate_runtime() + if type(limits) is not _CanonicalWireLimits: + raise TypeError("limits must be _CanonicalWireLimits") + diagnostic_limit = limits.diagnostic_bytes + if validate_plan is not None and not callable(validate_plan): + raise TypeError("validate_plan must be callable or None") + if consume is not None and not callable(consume): + raise TypeError("consume must be callable or None") + operation_bytes = _label_bytes("operation", operation) + transaction_bytes = _label_bytes("transaction_id", transaction_id) + if type(context_digest) is not bytes or len(context_digest) != hashlib.sha256().digest_size: + raise ValueError("context_digest must be exactly 32 bytes") + plan = _prepare_canonical_wire_value(local_fragment, limits) + header = _success_header( + binding, + plan, + limits, + coordinate=coordinate, + member_count=member_count, + operation=operation_bytes, + transaction=transaction_bytes, + context_digest=context_digest, + callback_flags=(1 if validate_plan is not None else 0) | (2 if consume is not None else 0), + ) + except Exception as exc: + error = exc + headers, diagnostics = _exchange_header( + binding, + header, + _diagnostic(error, limit=diagnostic_limit), + member_count=member_count, + ) + _validate_initial_exchange( + binding, + headers, + diagnostics, + member_count=member_count, + ) + assert plan is not None + assert type(limits) is _CanonicalWireLimits + direct_cpu = all(header[_H_TRANSPORT] == 0 for header in headers) + + for source in range(member_count): + prepared = _receive_metadata( + binding, + plan, + headers, + limits, + coordinate=coordinate, + source=source, + member_count=member_count, + validate_plan=validate_plan, + ) + payloads = _allocate_payloads( + binding, + prepared, + limits, + coordinate=coordinate, + source=source, + member_count=member_count, + ) + _transfer_payloads( + binding, + plan, + prepared, + payloads, + limits, + coordinate=coordinate, + source=source, + member_count=member_count, + direct_cpu=direct_cpu, + ) + value = None + error = None + try: + value = _reconstruct_canonical_wire_value( + prepared, + payloads, + expected_fragment_digest=_header_digest(headers[source], _H_FRAGMENT_DIGEST), + ) + except Exception as exc: + error = exc + _phase_vote( + binding, + member_count=member_count, + coordinate=coordinate, + source=source, + phase=_PHASE_RECONSTRUCT, + phase_name="fragment reconstruction", + error=error, + diagnostic_limit=limits.diagnostic_bytes, + ) + error = None + if consume is not None: + try: + consume(members[source], value) + except Exception as exc: + error = exc + _phase_vote( + binding, + member_count=member_count, + coordinate=coordinate, + source=source, + phase=_PHASE_CONSUME, + phase_name="fragment consumer", + error=error, + diagnostic_limit=limits.diagnostic_bytes, + ) + + +__all__ = [] diff --git a/src/gefen/portable_dcp.py b/src/gefen/portable_dcp.py new file mode 100644 index 0000000..c698bcb --- /dev/null +++ b/src/gefen/portable_dcp.py @@ -0,0 +1,563 @@ +"""Synchronous PyTorch DCP persistence for collective portable Gefen state.""" + +from __future__ import annotations + +import math +import os + +import torch + +from gefen.checkpoint import CheckpointProcessGroupBinding +from gefen.portable_state import PortableStateLimits +from gefen.portable_wire import ( + _DTYPE_BY_VALUE, + _CanonicalWireLimits, + _CanonicalWirePlan, + _parse_canonical_wire_metadata, + _prepare_canonical_wire_value, + _reconstruct_canonical_wire_value, +) + + +_DCP_METADATA_KEY = "__gefen_portable_metadata_v1__" +_DCP_PAYLOAD_PREFIX = "__gefen_portable_tensor_" +_DCP_PAYLOAD_SUFFIX = "__" +_DCP_PAYLOAD_DIGITS = 16 +_DCP_SAVE_PREFLIGHT_TRANSACTION = "gefen-portable-dcp-save-preflight-v1" +_DCP_LOAD_PREFLIGHT_TRANSACTION = "gefen-portable-dcp-load-preflight-v1" + + +def _require_namespace(namespace, limits: PortableStateLimits) -> str: + if type(namespace) is not str or not namespace or namespace != namespace.strip(): + raise ValueError("namespace must be a non-empty trimmed string") + if "." in namespace or "\x00" in namespace: + raise ValueError("namespace must not contain dots or NUL") + if any(not (character.isascii() and (character.isalnum() or character in {"_", "-"})) for character in namespace): + raise ValueError("namespace must contain only ASCII letters, digits, underscores, and hyphens") + from gefen.portable_runtime import _bounded_utf8_length + + _bounded_utf8_length( + namespace, + limit=limits.max_string_bytes, + name="DCP namespace", + ) + return namespace + + +def _payload_key(index: int) -> str: + if type(index) is not int or index < 0 or index >= 10**_DCP_PAYLOAD_DIGITS: + raise ValueError("portable DCP payload index is out of range") + return "{}{:0{}d}{}".format( + _DCP_PAYLOAD_PREFIX, + index, + _DCP_PAYLOAD_DIGITS, + _DCP_PAYLOAD_SUFFIX, + ) + + +def _flat_key(namespace: str, inner_key: str) -> str: + return "{}.{}".format(namespace, inner_key) + + +def _validate_dcp_runtime(binding: CheckpointProcessGroupBinding) -> None: + import torch.distributed as dist + + members = binding.identity.ordered_members + if len(members) == 1: + if not dist.is_available() or not dist.is_initialized(): + return + if dist.get_world_size() > 1: + raise RuntimeError("a singleton portable DCP binding cannot run inside a larger initialized default world") + backend = str(dist.get_backend()).lower() + else: + if dist.get_global_rank(binding.process_group, 0) != 0: + raise RuntimeError("portable DCP requires checkpoint group coordinate zero to be global rank zero") + backend = str(dist.get_backend(binding.process_group)).lower() + if "nccl" in backend: + if binding.collective_device.type != "cuda": + raise ValueError("portable DCP requires a CUDA checkpoint device with NCCL") + expected = binding.collective_device.index + if expected is None or torch.cuda.current_device() != expected: + raise ValueError("portable DCP requires the current CUDA device to match the checkpoint binding") + elif "gloo" in backend or "mpi" in backend: + if binding.collective_device.type != "cpu": + raise ValueError("portable DCP requires a CPU checkpoint device with this backend") + + +def _storage_identity(storage, limits: PortableStateLimits): + storage_class = type(storage) + module = storage_class.__module__ + qualname = storage_class.__qualname__ + if type(module) is not str or type(qualname) is not str: + raise TypeError("DCP storage classes require exact string identities") + checkpoint_id = getattr(storage, "checkpoint_id", None) + try: + checkpoint_id = os.fspath(checkpoint_id) + except TypeError as exc: + raise TypeError("DCP storage must expose a string or path-like checkpoint_id") from exc + if type(checkpoint_id) is not str or not checkpoint_id or "\x00" in checkpoint_id: + raise ValueError("DCP storage checkpoint_id must be a non-empty string without NUL") + from torch.distributed.checkpoint import FileSystemReader, FileSystemWriter + + if isinstance(storage, (FileSystemReader, FileSystemWriter)): + checkpoint_id = os.path.abspath(checkpoint_id) + from gefen.portable_runtime import _bounded_utf8_length + + for name, value in ( + ("DCP storage class module", module), + ("DCP storage class qualname", qualname), + ("DCP storage checkpoint_id", checkpoint_id), + ): + _bounded_utf8_length( + value, + limit=limits.max_string_bytes, + name=name, + ) + return { + "module": module, + "qualname": qualname, + "checkpoint_id": checkpoint_id, + } + + +def _preflight_dcp_operation( + optimizer, + *, + checkpoint_process_group, + transaction_id, + limits, + namespace, + storage, + storage_type, + operation: str, +): + from gefen import portable_runtime as runtime + from gefen.portable_collective import _collective_unanimous_status + from gefen.hybrid import GefenMuonHybrid + + if type(optimizer) is GefenMuonHybrid: + from gefen.portable_hybrid import _hybrid_transport_binding + + transport = _hybrid_transport_binding( + optimizer, + checkpoint_process_group, + ) + else: + transport = runtime._preflight_transport_binding( + optimizer, + checkpoint_process_group, + ) + binding = None + normalized_limits = None + normalized_transaction = None + normalized_namespace = None + implementation = None + context_digest = bytes(32) + wire_limits = runtime._STATUS_FALLBACK_LIMITS + error = None + try: + if type(optimizer) is GefenMuonHybrid: + from gefen.portable_hybrid import ( + HYBRID_PORTABLE_STATE_IMPLEMENTATION, + _preflight_hybrid_portable_local, + ) + + ( + binding, + normalized_transaction, + normalized_limits, + _children, + _routing, + base_context, + _base_digest, + _live_token, + ) = _preflight_hybrid_portable_local( + optimizer, + checkpoint_process_group=checkpoint_process_group, + transaction_id=transaction_id, + limits=limits, + ) + implementation = HYBRID_PORTABLE_STATE_IMPLEMENTATION + else: + runtime._validate_supplied_binding(checkpoint_process_group, transport) + binding = transport + normalized_limits = runtime._require_limits(limits) + normalized_transaction = runtime._require_transaction_id(transaction_id) + implementation = runtime._optimizer_implementation(optimizer) + runtime._validate_context_identity_bounds(binding, normalized_limits) + prepared = runtime._prepare_local_structure( + optimizer, + implementation, + binding, + normalized_limits, + include_payload=False, + ) + runtime._validate_prepared_local_state( + optimizer, + implementation, + prepared, + ) + base_context = runtime._base_context(binding, implementation) + wire_limits = normalized_limits._wire_limits() + normalized_namespace = _require_namespace(namespace, normalized_limits) + if not isinstance(storage, storage_type): + raise TypeError("storage must be a {}".format(storage_type.__name__)) + storage_identity = _storage_identity(storage, normalized_limits) + _validate_dcp_runtime(binding) + context = { + **base_context, + "dcp_envelope_version": 1, + "dcp_namespace": normalized_namespace, + "dcp_storage": storage_identity, + "transaction_id": normalized_transaction, + } + runtime._preflight_portable_value(context, normalized_limits) + context_digest = runtime._context_digest(context) + except Exception as exc: + error = exc + _collective_unanimous_status( + transport, + error, + operation="portable_dcp_{}_preflight".format(operation), + transaction_id=(_DCP_SAVE_PREFLIGHT_TRANSACTION if operation == "save" else _DCP_LOAD_PREFLIGHT_TRANSACTION), + context_digest=context_digest, + limits=wire_limits, + ) + assert ( + binding is not None + and normalized_limits is not None + and normalized_transaction is not None + and normalized_namespace is not None + and implementation is not None + ) + return ( + binding, + normalized_transaction, + normalized_limits, + normalized_namespace, + ) + + +def _metadata_tensor(metadata: bytes) -> torch.Tensor: + if type(metadata) is not bytes: + raise TypeError("portable DCP metadata must be bytes") + return torch.frombuffer(bytearray(metadata), dtype=torch.uint8) + + +def _state_from_plan( + namespace: str, + plan: _CanonicalWirePlan, + limits: _CanonicalWireLimits, +): + if type(plan) is not _CanonicalWirePlan: + raise TypeError("plan must be a canonical wire plan") + if type(limits) is not _CanonicalWireLimits: + raise TypeError("limits must be canonical wire limits") + if len(plan.payload_tensors) + 1 > limits.max_container_items: + raise ValueError("portable DCP envelope exceeds max_container_items") + envelope = {_DCP_METADATA_KEY: _metadata_tensor(plan.metadata)} + envelope.update({_payload_key(index): tensor for index, tensor in enumerate(plan.payload_tensors)}) + return {namespace: envelope} + + +def _validate_tensor_metadata(value, *, name: str, limits: _CanonicalWireLimits): + from torch.distributed.checkpoint.metadata import ( + ChunkStorageMetadata, + TensorProperties, + TensorStorageMetadata, + ) + + if type(value) is not TensorStorageMetadata: + raise TypeError("{} must be full-tensor DCP metadata".format(name)) + properties = value.properties + if ( + type(properties) is not TensorProperties + or type(value.size) is not torch.Size + or properties.dtype not in _DTYPE_BY_VALUE + or properties.layout is not torch.strided + or properties.requires_grad is not False + or properties.memory_format is not torch.contiguous_format + or properties.pin_memory is not False + ): + raise TypeError("{} has unsupported tensor properties".format(name)) + shape = tuple(value.size) + if len(shape) > limits.max_tensor_rank or any( + type(dimension) is not int or dimension < 0 or dimension > (1 << 63) - 1 for dimension in shape + ): + raise ValueError("{} has invalid tensor geometry".format(name)) + if type(value.chunks) is not list or len(value.chunks) != 1: + raise ValueError("{} must contain exactly one full DCP chunk".format(name)) + chunk = value.chunks[0] + if ( + type(chunk) is not ChunkStorageMetadata + or type(chunk.offsets) is not torch.Size + or type(chunk.sizes) is not torch.Size + or tuple(chunk.offsets) != (0,) * len(shape) + or tuple(chunk.sizes) != shape + ): + raise ValueError("{} must contain one complete unsharded tensor".format(name)) + numel = math.prod(shape) + element_size = _DTYPE_BY_VALUE[properties.dtype][1] + nbytes = numel * element_size + return properties.dtype, shape, nbytes + + +def _allocate_dcp_state( + metadata, + *, + namespace: str, + limits: _CanonicalWireLimits, +): + from torch.distributed.checkpoint.metadata import Metadata + + if type(metadata) is not Metadata: + raise TypeError("portable DCP requires exact checkpoint Metadata") + entries = metadata.state_dict_metadata + if type(entries) is not dict: + raise TypeError("portable DCP metadata entries must be a dict") + if len(entries) < 1 or len(entries) > limits.max_tensors + 1 or len(entries) > limits.max_container_items: + raise ValueError("portable DCP tensor count exceeds limits") + metadata_key = _flat_key(namespace, _DCP_METADATA_KEY) + if metadata_key not in entries: + raise ValueError("portable DCP metadata tensor is missing") + payload_count = len(entries) - 1 + expected_payload_keys = tuple(_flat_key(namespace, _payload_key(index)) for index in range(payload_count)) + expected_keys = {metadata_key, *expected_payload_keys} + if set(entries) != expected_keys or any(type(key) is not str for key in entries): + raise ValueError("portable DCP checkpoint has unexpected tensor keys") + planner_data = metadata.planner_data + expected_planner_data = {key: tuple(key.split(".", 1)) for key in expected_keys} + if ( + type(planner_data) is not dict + or planner_data != expected_planner_data + or any( + type(key) is not str + or type(path) is not tuple + or len(path) != 2 + or any(type(component) is not str for component in path) + for key, path in planner_data.items() + ) + ): + raise ValueError("portable DCP checkpoint has incompatible planner paths") + + metadata_dtype, metadata_shape, metadata_nbytes = _validate_tensor_metadata( + entries[metadata_key], + name="portable DCP wire metadata", + limits=limits, + ) + if ( + metadata_dtype is not torch.uint8 + or len(metadata_shape) != 1 + or metadata_nbytes == 0 + or metadata_nbytes > limits.max_metadata_bytes + ): + raise ValueError("portable DCP wire metadata exceeds its byte limit") + + payload_specs = [] + total_payload_bytes = 0 + for index, key in enumerate(expected_payload_keys): + dtype, shape, nbytes = _validate_tensor_metadata( + entries[key], + name="portable DCP payload {}".format(index), + limits=limits, + ) + if total_payload_bytes > limits.max_fragment_tensor_bytes - nbytes: + raise ValueError("portable DCP payload tensors exceed their byte limit") + total_payload_bytes += nbytes + payload_specs.append((_payload_key(index), dtype, shape)) + + envelope = { + _DCP_METADATA_KEY: torch.empty( + metadata_shape, + dtype=metadata_dtype, + device="cpu", + ) + } + for key, dtype, shape in payload_specs: + envelope[key] = torch.empty(shape, dtype=dtype, device="cpu") + return {namespace: envelope} + + +def _load_planner(namespace: str, limits: _CanonicalWireLimits): + import torch.distributed.checkpoint as dcp + + class _PortableDCPDynamicLoadPlanner(dcp.DefaultLoadPlanner): + def set_up_planner( + self, + state_dict, + metadata=None, + is_coordinator=False, + ) -> None: + if ( + type(state_dict) is not dict + or set(state_dict) != {namespace} + or type(state_dict[namespace]) is not dict + or state_dict[namespace] + ): + raise ValueError("portable DCP load planner requires one empty exact namespace") + allocated = _allocate_dcp_state( + metadata, + namespace=namespace, + limits=limits, + ) + dict.update(state_dict[namespace], allocated[namespace]) + super().set_up_planner(state_dict, metadata, is_coordinator) + + return _PortableDCPDynamicLoadPlanner() + + +def _decode_dcp_state( + state, + *, + namespace: str, + limits: _CanonicalWireLimits, +): + if type(state) is not dict or set(state) != {namespace}: + raise ValueError("portable DCP root state has an invalid schema") + envelope = state[namespace] + if type(envelope) is not dict or _DCP_METADATA_KEY not in envelope: + raise ValueError("portable DCP envelope has an invalid schema") + metadata_tensor = envelope[_DCP_METADATA_KEY] + if ( + type(metadata_tensor) is not torch.Tensor + or metadata_tensor.dtype is not torch.uint8 + or metadata_tensor.device.type != "cpu" + or metadata_tensor.ndim != 1 + or not metadata_tensor.is_contiguous() + or metadata_tensor.numel() > limits.max_metadata_bytes + ): + raise ValueError("portable DCP wire metadata tensor is invalid") + metadata = bytes(memoryview(metadata_tensor.numpy())) + prepared = _parse_canonical_wire_metadata(metadata, limits=limits) + expected_keys = { + _DCP_METADATA_KEY, + *(_payload_key(index) for index in range(len(prepared.tensor_specs))), + } + if set(envelope) != expected_keys: + raise ValueError("portable DCP payload key count disagrees with wire metadata") + payloads = tuple(envelope[_payload_key(index)] for index in range(len(prepared.tensor_specs))) + document = _reconstruct_canonical_wire_value(prepared, payloads) + return document, prepared.fragment_digest + + +def save_portable_dcp( + optimizer, + *, + checkpoint_process_group, + storage_writer, + transaction_id, + limits, + namespace="optimizer", +): + """Collectively save one tensor-only portable optimizer document through DCP.""" + + import torch.distributed.checkpoint as dcp + from torch.distributed.checkpoint.storage import StorageWriter + + binding, transaction_id, limits, namespace = _preflight_dcp_operation( + optimizer, + checkpoint_process_group=checkpoint_process_group, + transaction_id=transaction_id, + limits=limits, + namespace=namespace, + storage=storage_writer, + storage_type=StorageWriter, + operation="save", + ) + document = optimizer.export_portable_state( + checkpoint_process_group=binding, + transaction_id=transaction_id, + limits=limits, + ) + wire_limits = limits._wire_limits(collective=True) + plan = None + state = None + digest = bytes(32) + error = None + try: + plan = _prepare_canonical_wire_value(document, wire_limits) + state = _state_from_plan(namespace, plan, wire_limits) + digest = plan.fragment_digest + except Exception as exc: + error = exc + from gefen.portable_collective import _collective_unanimous_status + + _collective_unanimous_status( + binding, + error, + operation="portable_dcp_save_encode", + transaction_id=transaction_id, + context_digest=digest, + limits=wire_limits, + ) + assert plan is not None and state is not None + return dcp.save( + state, + storage_writer=storage_writer, + process_group=binding.process_group, + ) + + +def load_portable_dcp( + optimizer, + *, + checkpoint_process_group, + storage_reader, + transaction_id, + limits, + namespace="optimizer", +) -> None: + """Collectively load, verify, and atomically import portable optimizer state.""" + + import torch.distributed.checkpoint as dcp + from torch.distributed.checkpoint.storage import StorageReader + + binding, transaction_id, limits, namespace = _preflight_dcp_operation( + optimizer, + checkpoint_process_group=checkpoint_process_group, + transaction_id=transaction_id, + limits=limits, + namespace=namespace, + storage=storage_reader, + storage_type=StorageReader, + operation="load", + ) + wire_limits = limits._wire_limits(collective=True) + state = {namespace: {}} + dcp.load( + state, + storage_reader=storage_reader, + planner=_load_planner(namespace, wire_limits), + process_group=binding.process_group, + ) + document = None + digest = bytes(32) + error = None + try: + document, digest = _decode_dcp_state( + state, + namespace=namespace, + limits=wire_limits, + ) + except Exception as exc: + error = exc + from gefen.portable_collective import _collective_unanimous_status + + _collective_unanimous_status( + binding, + error, + operation="portable_dcp_load_decode", + transaction_id=transaction_id, + context_digest=digest, + limits=wire_limits, + ) + assert document is not None + optimizer.import_portable_state( + document, + checkpoint_process_group=binding, + transaction_id=transaction_id, + limits=limits, + ) + + +__all__ = ["load_portable_dcp", "save_portable_dcp"] diff --git a/src/gefen/portable_fields.py b/src/gefen/portable_fields.py new file mode 100644 index 0000000..8a5441c --- /dev/null +++ b/src/gefen/portable_fields.py @@ -0,0 +1,353 @@ +"""Pure assembly and projection for dense logical optimizer-state fields.""" + +from collections.abc import Sequence + +import torch + +from gefen.contracts import ( + LogicalRegion, + ParameterIdentity, + ParameterLayout, + PlacementKind, + ShardIdentity, + ShardingManifest, +) +from gefen.portable import _element_chunks, _read_flat_chunk, _validate_plain_tensor + + +def _validate_parameter(parameter) -> ParameterIdentity: + if not isinstance(parameter, ParameterIdentity): + raise TypeError("parameter must be a ParameterIdentity") + # Reconstruct the frozen descriptor so an object modified through low-level + # attribute access cannot bypass its invariants at a checkpoint boundary. + return ParameterIdentity( + parameter.fqn, + parameter.global_shape, + schema_version=parameter.schema_version, + ) + + +def _validate_manifest(manifest) -> ShardingManifest: + if not isinstance(manifest, ShardingManifest): + raise TypeError("manifest must be a ShardingManifest") + return ShardingManifest( + manifest.shards, + schema_version=manifest.schema_version, + ) + + +def _validate_dense_tensor(value, *, name: str, shape) -> torch.Tensor: + value = _validate_plain_tensor( + value, + name=name, + dtype=torch.float32, + ) + if value.device.type != "cpu": + raise ValueError("{} must be on CPU".format(name)) + if tuple(value.shape) != tuple(shape): + raise ValueError("{} must have shape {}".format(name, tuple(shape))) + return value + + +def _finish_tight(value: torch.Tensor, *, name: str) -> torch.Tensor: + if ( + value.device.type != "cpu" + or value.dtype != torch.float32 + or value.requires_grad + or not value.is_contiguous() + or value.storage_offset() != 0 + or value.untyped_storage().nbytes() != value.numel() * value.element_size() + ): + raise RuntimeError("{} did not produce tight detached CPU fp32 state".format(name)) + return value + + +def _tight_clone(value: torch.Tensor, *, name: str) -> torch.Tensor: + result = torch.empty(tuple(value.shape), dtype=torch.float32, device="cpu") + flat = result.reshape(-1) + for start, stop in _element_chunks(value.numel()): + flat[start:stop].copy_(_read_flat_chunk(value, start, stop).resolve_conj().resolve_neg()) + return _finish_tight(result, name=name) + + +def _tensor_bits_equal(left: torch.Tensor, right: torch.Tensor) -> bool: + for start, stop in _element_chunks(left.numel()): + left_bytes = ( + _read_flat_chunk(left, start, stop) + .resolve_conj() + .resolve_neg() + .contiguous() + .view(torch.uint8) + ) + right_bytes = ( + _read_flat_chunk(right, start, stop) + .resolve_conj() + .resolve_neg() + .contiguous() + .view(torch.uint8) + ) + if not torch.equal(left_bytes, right_bytes): + return False + return True + + +def _tight_flat_slice( + value: torch.Tensor, + *, + flat_offset: int, + length: int, + name: str, +) -> torch.Tensor: + result = torch.empty((length,), dtype=torch.float32, device="cpu") + for start, stop in _element_chunks(length): + result[start:stop].copy_( + _read_flat_chunk( + value, + flat_offset + start, + flat_offset + stop, + ) + .resolve_conj() + .resolve_neg() + ) + return _finish_tight(result, name=name) + + +def _full_output(parameter: ParameterIdentity) -> torch.Tensor: + return torch.empty( + parameter.global_shape, + dtype=torch.float32, + device="cpu", + ) + + +def _region_index(region: LogicalRegion): + return tuple(slice(offset, offset + length) for offset, length in zip(region.offsets, region.lengths)) + + +def _normalize_fragments(fragments): + if isinstance(fragments, (str, bytes, bytearray)): + raise TypeError("fragments must be a sequence of shard/payload pairs") + try: + fragments = tuple(fragments) + except TypeError as exc: + raise TypeError("fragments must be a sequence of shard/payload pairs") from exc + + normalized = [] + for item in fragments: + if not isinstance(item, Sequence) or isinstance(item, (str, bytes, bytearray)) or len(item) != 2: + raise TypeError("each fragment must be a shard/payload pair") + shard, payload = item + if not isinstance(shard, ShardIdentity): + raise TypeError("each fragment identity must be a ShardIdentity") + normalized.append((shard, payload)) + return tuple(normalized) + + +def _validate_fragment_set( + manifest: ShardingManifest, + parameter: ParameterIdentity, + fragments, +): + expected = manifest.for_parameter(parameter.fqn) + if not expected: + raise ValueError("manifest does not contain the requested parameter") + if any(shard.parameter != parameter for shard in expected): + raise ValueError("manifest parameter identity does not match parameter") + + fragments = _normalize_fragments(fragments) + identities = tuple(shard for shard, _ in fragments) + if len(set(identities)) != len(identities): + raise ValueError("fragments must not contain duplicate shard identities") + if set(identities) != set(expected): + raise ValueError("fragments must exactly cover the requested parameter manifest") + return expected, {shard: payload for shard, payload in fragments} + + +def _assemble_replicated( + parameter: ParameterIdentity, + shards, + payload_by_shard, + *, + name: str, +) -> torch.Tensor: + reference = None + for shard in shards: + payload = payload_by_shard[shard] + if payload is None: + raise ValueError("{} replicas require a payload on every shard".format(name)) + payload = _validate_dense_tensor( + payload, + name="{} fragment".format(name), + shape=parameter.global_shape, + ) + if reference is None: + reference = payload + elif not _tensor_bits_equal(reference, payload): + raise ValueError("{} replicas disagree".format(name)) + if reference is None: # ShardingManifest forbids this, but keep failure local. + raise ValueError("{} requires at least one replica".format(name)) + return _tight_clone(reference, name="assembled {}".format(name)) + + +def _assemble_flattened( + parameter: ParameterIdentity, + shards, + payload_by_shard, +) -> torch.Tensor: + result = _full_output(parameter) + flat = result.reshape(-1) + for shard in shards: + logical_slice = shard.logical_slice + payload = payload_by_shard[shard] + if logical_slice.length == 0: + if payload is not None: + raise ValueError("an empty flattened fragment must have payload None") + continue + if payload is None: + raise ValueError("a nonempty flattened fragment requires a payload") + payload = _validate_dense_tensor( + payload, + name="flattened fragment", + shape=(logical_slice.length,), + ) + start = logical_slice.flat_offset + flat[start : start + logical_slice.length].copy_(payload.detach()) + return _finish_tight(result, name="assembled flattened field") + + +def _assemble_owner( + parameter: ParameterIdentity, + shards, + payload_by_shard, +) -> torch.Tensor: + owner_payload = None + for shard in shards: + payload = payload_by_shard[shard] + if shard.local_member == shard.owner: + if payload is None: + raise ValueError("the whole-parameter owner requires a payload") + owner_payload = _validate_dense_tensor( + payload, + name="whole-parameter owner fragment", + shape=parameter.global_shape, + ) + elif payload is not None: + raise ValueError("whole-parameter non-owners must have payload None") + if owner_payload is None: + raise ValueError("whole-parameter fragments do not contain the owner payload") + return _tight_clone(owner_payload, name="assembled whole-parameter field") + + +def _assemble_dtensor_dimension_shards( + parameter: ParameterIdentity, + shards, + payload_by_shard, +) -> torch.Tensor: + result = _full_output(parameter) + for shard in shards: + region = shard.logical_region + payload = payload_by_shard[shard] + if region.numel == 0: + if payload is not None: + raise ValueError("an empty DTensor fragment must have payload None") + continue + if payload is None: + raise ValueError("a nonempty DTensor fragment requires a payload") + payload = _validate_dense_tensor( + payload, + name="DTensor fragment", + shape=region.lengths, + ) + result[_region_index(region)].copy_(payload.detach()) + return _finish_tight(result, name="assembled DTensor field") + + +def _assemble_dense_logical_field( + manifest, + parameter, + fragments, +) -> torch.Tensor: + """Assemble one complete topology-neutral fp32 CPU logical field.""" + + manifest = _validate_manifest(manifest) + parameter = _validate_parameter(parameter) + shards, payload_by_shard = _validate_fragment_set( + manifest, + parameter, + fragments, + ) + layout = shards[0].layout + if layout is ParameterLayout.REPLICATED: + return _assemble_replicated( + parameter, + shards, + payload_by_shard, + name="replicated", + ) + if layout is ParameterLayout.FLATTENED_ELEMENT_SHARD: + return _assemble_flattened(parameter, shards, payload_by_shard) + if layout is ParameterLayout.WHOLE_PARAMETER_OWNER: + return _assemble_owner(parameter, shards, payload_by_shard) + if layout is ParameterLayout.DTENSOR_1D_DEFAULT_WORLD: + placement_kind = shards[0].placements[0].kind + if placement_kind is PlacementKind.REPLICATE: + return _assemble_replicated( + parameter, + shards, + payload_by_shard, + name="replicated DTensor", + ) + if placement_kind is PlacementKind.DIMENSION_SHARD: + return _assemble_dtensor_dimension_shards( + parameter, + shards, + payload_by_shard, + ) + raise ValueError("unsupported portable field layout") + + +def _project_dense_logical_field( + parameter, + tensor, + shard, +): + """Project one complete logical fp32 CPU field to a target shard identity.""" + + parameter = _validate_parameter(parameter) + if not isinstance(shard, ShardIdentity): + raise TypeError("shard must be a ShardIdentity") + if shard.parameter != parameter: + raise ValueError("target shard parameter identity does not match parameter") + tensor = _validate_dense_tensor( + tensor, + name="global logical field", + shape=parameter.global_shape, + ) + + if shard.layout is ParameterLayout.REPLICATED: + return _tight_clone(tensor, name="projected replicated field") + if shard.layout is ParameterLayout.FLATTENED_ELEMENT_SHARD: + logical_slice = shard.logical_slice + if logical_slice.length == 0: + return None + return _tight_flat_slice( + tensor, + flat_offset=logical_slice.flat_offset, + length=logical_slice.length, + name="projected flattened field", + ) + if shard.layout is ParameterLayout.WHOLE_PARAMETER_OWNER: + if shard.local_member != shard.owner: + return None + return _tight_clone(tensor, name="projected whole-parameter field") + if shard.layout is ParameterLayout.DTENSOR_1D_DEFAULT_WORLD: + placement = shard.placements[0] + if placement.kind is PlacementKind.REPLICATE: + return _tight_clone(tensor, name="projected replicated DTensor field") + if placement.kind is PlacementKind.DIMENSION_SHARD: + region = shard.logical_region + if region.numel == 0: + return None + values = tensor[_region_index(region)] + return _tight_clone(values, name="projected DTensor field") + raise ValueError("unsupported portable field layout") diff --git a/src/gefen/portable_hybrid.py b/src/gefen/portable_hybrid.py new file mode 100644 index 0000000..bf6b117 --- /dev/null +++ b/src/gefen/portable_hybrid.py @@ -0,0 +1,691 @@ +"""Strict composite envelope for Gefen-backed Hybrid portable state.""" + +import hashlib +import hmac + +from gefen.portable_schema import ( + PORTABLE_STATE_DIGEST_ALGORITHM, + normalize_portable_state_document, + portable_state_digest, +) + + +HYBRID_PORTABLE_STATE_FORMAT = "gefen.portable_composite_state" +HYBRID_PORTABLE_STATE_FORMAT_VERSION = 1 +HYBRID_PORTABLE_STATE_COVERAGE = "global_logical_composite_optimizer" +HYBRID_PORTABLE_STATE_IMPLEMENTATION = "gefen.GefenMuonHybrid" + +_HYBRID_PORTABLE_ROLES = ("muon", "backup") +_HYBRID_PORTABLE_CHILD_IMPLEMENTATIONS = { + "muon": "gefen.GefenMuon", + "backup": "gefen.Gefen", +} +_HYBRID_PORTABLE_TOP_LEVEL_KEYS = frozenset( + { + "format", + "format_version", + "coverage", + "implementation", + "backup_optimizer", + "routing", + "children", + "completion", + } +) +_HYBRID_PORTABLE_COMPLETION_KEYS = frozenset({"status", "digest_algorithm", "digest"}) +_HYBRID_EXPORT_PREFLIGHT_TRANSACTION = "gefen-hybrid-portable-export-preflight-v1" +_HYBRID_IMPORT_PREFLIGHT_TRANSACTION = "gefen-hybrid-portable-import-preflight-v1" + + +def _hybrid_portable_payload(document): + return {key: document[key] for key in sorted(_HYBRID_PORTABLE_TOP_LEVEL_KEYS - {"completion"})} + + +def _normalize_routing(routing): + if type(routing) is not dict: + raise ValueError("portable Hybrid routing must be an FQN dictionary") + if any(type(fqn) is not str for fqn in routing): + raise ValueError("portable Hybrid routing keys must be strings") + normalized = {} + for fqn in sorted(routing): + role = routing[fqn] + if not fqn or fqn != fqn.strip(): + raise ValueError("portable Hybrid routing keys must be non-empty trimmed FQNs") + if type(role) is not str or role not in _HYBRID_PORTABLE_ROLES: + raise ValueError("portable Hybrid routing values must name exact child roles") + normalized[fqn] = role + return normalized + + +def _normalize_children(children): + if type(children) is not dict or set(children) != set(_HYBRID_PORTABLE_ROLES): + raise ValueError("portable Hybrid children must contain exact muon and backup roles") + normalized = {} + for role in _HYBRID_PORTABLE_ROLES: + child = children[role] + normalized[role] = ( + None + if child is None + else normalize_portable_state_document( + child, + expected_implementation=_HYBRID_PORTABLE_CHILD_IMPLEMENTATIONS[role], + ) + ) + if all(child is None for child in normalized.values()): + raise ValueError("portable Hybrid state must contain at least one child") + return normalized + + +def _validate_child_routing(children, routing) -> None: + expected_routing = {} + common_step = None + deterministic = None + for role in _HYBRID_PORTABLE_ROLES: + child = children[role] + if child is None: + continue + for fqn in child["parameters"]: + if fqn in expected_routing: + raise ValueError("portable Hybrid child parameter FQNs must be disjoint") + expected_routing[fqn] = role + child_common = child["common"] + if type(child_common) is not dict: + raise ValueError("portable Hybrid child common state must be a dictionary") + child_step = child_common.get("gefen_global_step") + child_deterministic = child_common.get("gefen_deterministic") + if type(child_step) is not int or child_step < 0: + raise ValueError("portable Hybrid children require nonnegative exact global steps") + if type(child_deterministic) is not bool: + raise ValueError("portable Hybrid children require exact deterministic policies") + if common_step is None: + common_step = child_step + deterministic = child_deterministic + elif child_step != common_step: + raise ValueError("portable Hybrid child global steps must agree") + elif child_deterministic is not deterministic: + raise ValueError("portable Hybrid child deterministic policies must agree") + if routing != {fqn: expected_routing[fqn] for fqn in sorted(expected_routing)}: + raise ValueError("portable Hybrid routing does not exactly match child parameters") + + +def _normalize_hybrid_portable_payload(state): + if state["format"] != HYBRID_PORTABLE_STATE_FORMAT: + raise ValueError("unsupported portable Hybrid state format") + if type(state["format_version"]) is not int or state["format_version"] != HYBRID_PORTABLE_STATE_FORMAT_VERSION: + raise ValueError("unsupported portable Hybrid state format_version: {}".format(state["format_version"])) + if state["coverage"] != HYBRID_PORTABLE_STATE_COVERAGE: + raise ValueError("unsupported portable Hybrid state coverage") + if state["implementation"] != HYBRID_PORTABLE_STATE_IMPLEMENTATION: + raise ValueError("portable Hybrid state implementation does not match the target") + if state["backup_optimizer"] != "gefen": + raise ValueError("portable Hybrid state requires a Gefen backup policy") + routing = _normalize_routing(state["routing"]) + children = _normalize_children(state["children"]) + _validate_child_routing(children, routing) + return { + "format": HYBRID_PORTABLE_STATE_FORMAT, + "format_version": HYBRID_PORTABLE_STATE_FORMAT_VERSION, + "coverage": HYBRID_PORTABLE_STATE_COVERAGE, + "implementation": HYBRID_PORTABLE_STATE_IMPLEMENTATION, + "backup_optimizer": "gefen", + "routing": routing, + "children": children, + } + + +def _validate_hybrid_portable_completion(state, payload) -> None: + completion = state["completion"] + if type(completion) is not dict or set(completion) != _HYBRID_PORTABLE_COMPLETION_KEYS: + raise ValueError("portable Hybrid completion marker has an invalid schema") + if completion["status"] != "complete": + raise ValueError("portable Hybrid state is not marked complete") + if completion["digest_algorithm"] != PORTABLE_STATE_DIGEST_ALGORITHM: + raise ValueError("unsupported portable Hybrid state digest algorithm") + digest = completion["digest"] + if ( + type(digest) is not str + or len(digest) != hashlib.sha256().digest_size * 2 + or any(character not in "0123456789abcdef" for character in digest) + ): + raise ValueError("portable Hybrid completion digest is invalid") + if not hmac.compare_digest(digest, portable_state_digest(payload)): + raise ValueError("portable Hybrid completion digest does not match its payload") + + +def build_hybrid_portable_state_document(*, backup_optimizer, routing, children): + """Build one complete composite wrapper around unchanged portable v3 children.""" + + state = { + "format": HYBRID_PORTABLE_STATE_FORMAT, + "format_version": HYBRID_PORTABLE_STATE_FORMAT_VERSION, + "coverage": HYBRID_PORTABLE_STATE_COVERAGE, + "implementation": HYBRID_PORTABLE_STATE_IMPLEMENTATION, + "backup_optimizer": backup_optimizer, + "routing": routing, + "children": children, + } + payload = _normalize_hybrid_portable_payload(state) + payload["completion"] = { + "status": "complete", + "digest_algorithm": PORTABLE_STATE_DIGEST_ALGORITHM, + "digest": portable_state_digest(payload), + } + return payload + + +def normalize_hybrid_portable_state_document(state): + """Clone and strictly validate one complete portable Hybrid wrapper.""" + + if type(state) is not dict or set(state) != _HYBRID_PORTABLE_TOP_LEVEL_KEYS: + raise ValueError("portable Hybrid state has an invalid top-level schema") + payload = _normalize_hybrid_portable_payload(state) + _validate_hybrid_portable_completion(state, payload) + payload["completion"] = dict(state["completion"]) + return payload + + +def _validate_hybrid_portable_limits(state, limits) -> None: + from gefen.portable_wire import _prepare_canonical_wire_value + + _prepare_canonical_wire_value(state, limits._wire_limits(collective=True)) + + +def _hybrid_children(optimizer): + from gefen.gefen import Gefen + from gefen.gefen_muon import GefenMuon + from gefen.hybrid import GefenMuonHybrid + + if type(optimizer) is not GefenMuonHybrid: + raise TypeError("portable Hybrid state requires an exact GefenMuonHybrid") + if optimizer.backup_optimizer != "gefen": + raise NotImplementedError("portable Hybrid state does not support an AdamW backup") + children = [] + if optimizer.muon is not None: + if type(optimizer.muon) is not GefenMuon: + raise TypeError("portable Hybrid state requires an exact GefenMuon child") + children.append(("muon", optimizer.muon)) + if optimizer.backup is not None: + if type(optimizer.backup) is not Gefen: + raise TypeError("portable Hybrid state requires an exact Gefen backup child") + children.append(("backup", optimizer.backup)) + if not children: + raise RuntimeError("portable Hybrid state requires at least one child") + return tuple(children) + + +def _hybrid_transport_binding(optimizer, supplied): + from gefen import portable_runtime as runtime + + # The caller-provided exact binding is the only safe first-status + # transport: child presence/type may differ across faulty ranks and must be + # voted before any rank tries to enter a child collective. + del optimizer + return runtime._require_binding(supplied) + + +def _hybrid_portable_live_token(optimizer): + from gefen import portable_runtime as runtime + + children = _hybrid_children(optimizer) + return ( + optimizer.backup_optimizer, + optimizer._deterministic, + optimizer._hybrid_post_sharding_finalized, + id(optimizer._hybrid_sharding_manifest), + id(optimizer._hybrid_codebook_process_group), + id(optimizer._state_param_owner), + runtime._portable_value_token(optimizer._hybrid_fqn_roles), + runtime._portable_value_token(optimizer._hybrid_finalized_slots), + tuple( + ( + None if parameter is None else id(parameter), + shard.sort_key, + ) + for parameter, shard in optimizer._hybrid_local_shard_bindings + ), + tuple((role, id(child), runtime._portable_live_token(child)) for role, child in children), + ) + + +def _validate_hybrid_portable_readiness(optimizer, binding): + from gefen import portable_runtime as runtime + from gefen.checkpoint import CheckpointProcessGroupBinding + from gefen.hybrid import GefenMuonHybrid + + if type(binding) is not CheckpointProcessGroupBinding: + raise TypeError("portable Hybrid readiness requires an exact checkpoint binding") + GefenMuonHybrid._reject_rebinding_method_shadows(optimizer) + children = _hybrid_children(optimizer) + if not optimizer._canonical_identity_ready(): + raise RuntimeError("portable Hybrid state requires a finalized composite binding") + if optimizer._hybrid_codebook_process_group is None: + raise RuntimeError("portable Hybrid state requires one shared codebook scope") + routing = dict(optimizer._hybrid_fqn_roles) + if ( + type(optimizer._hybrid_fqn_roles) is not tuple + or len(routing) != len(optimizer._hybrid_fqn_roles) + or tuple(sorted(routing.items())) != optimizer._hybrid_fqn_roles + ): + raise RuntimeError("portable Hybrid routing metadata is not canonical") + expected_routing = {} + layouts = {} + common_step = None + deterministic = None + for role, child in children: + if child._gefen_codebook_process_group is not optimizer._hybrid_codebook_process_group: + raise RuntimeError("portable Hybrid children must share one exact codebook binding") + child_binding = runtime._preflight_transport_binding(child, binding) + runtime._validate_supplied_binding(binding, child_binding) + layouts[role] = runtime._validate_live_readiness( + child, + runtime._optimizer_implementation(child), + binding, + ) + for slot in child._gefen_logical_slots: + fqn = slot.shard.parameter.fqn + if fqn in expected_routing: + raise RuntimeError("portable Hybrid child FQNs must be disjoint") + expected_routing[fqn] = role + child_step = child._canonical_common_global_step() + child_deterministic = child._deterministic + if common_step is None: + common_step = child_step + deterministic = child_deterministic + elif child_step != common_step: + raise RuntimeError("portable Hybrid child global steps must agree") + elif child_deterministic is not deterministic: + raise RuntimeError("portable Hybrid child deterministic policies must agree") + if routing != {fqn: expected_routing[fqn] for fqn in sorted(expected_routing)}: + raise RuntimeError("portable Hybrid routing does not match its child slots") + if optimizer._deterministic is not deterministic: + raise RuntimeError("portable Hybrid deterministic policy disagrees with its children") + return children, routing, layouts + + +def _hybrid_child_transaction_id(parent: str, role: str, operation: str) -> str: + if role not in _HYBRID_PORTABLE_ROLES or operation not in {"export", "import"}: + raise ValueError("invalid portable Hybrid child transaction domain") + digest = hashlib.sha256( + b"gefen.portable_hybrid.child.v1\0" + + operation.encode("ascii") + + b"\0" + + role.encode("ascii") + + b"\0" + + parent.encode("utf-8") + ).hexdigest() + return "hybrid-{}-{}-{}".format(operation, role, digest) + + +def _preflight_hybrid_portable_local( + optimizer, + *, + checkpoint_process_group, + transaction_id, + limits, +): + from gefen import portable_runtime as runtime + + binding = _hybrid_transport_binding(optimizer, checkpoint_process_group) + runtime._validate_supplied_binding(checkpoint_process_group, binding) + normalized_limits = runtime._require_limits(limits) + normalized_transaction = runtime._require_transaction_id(transaction_id) + runtime._validate_context_identity_bounds(binding, normalized_limits) + children, routing, layouts = _validate_hybrid_portable_readiness( + optimizer, + binding, + ) + context = { + **runtime._base_context(binding, HYBRID_PORTABLE_STATE_IMPLEMENTATION), + "composite_format": HYBRID_PORTABLE_STATE_FORMAT, + "composite_format_version": HYBRID_PORTABLE_STATE_FORMAT_VERSION, + "backup_optimizer": optimizer.backup_optimizer, + "roles": tuple(role for role, _child in children), + "routing": routing, + "layouts": { + role: tuple(sorted(layout.value for layout in role_layouts)) for role, role_layouts in layouts.items() + }, + "transaction_id": normalized_transaction, + } + runtime._preflight_portable_value(context, normalized_limits) + return ( + binding, + normalized_transaction, + normalized_limits, + children, + routing, + context, + runtime._context_digest(context), + _hybrid_portable_live_token(optimizer), + ) + + +def _hybrid_portable_contract_support(optimizer): + """Return dynamic composite CANONICAL_GLOBAL support without collectives.""" + + try: + from gefen.checkpoint import CheckpointProcessGroupBinding + from gefen.codebook import CodebookProcessGroupBinding + from gefen.contracts import CheckpointTransport + + scope = optimizer._hybrid_codebook_process_group + if type(scope) is not CodebookProcessGroupBinding: + return frozenset(), frozenset(), frozenset() + binding = CheckpointProcessGroupBinding( + scope.identity, + scope.local_member, + scope.process_group, + scope.collective_device, + ) + children, _routing, _layouts = _validate_hybrid_portable_readiness( + optimizer, + binding, + ) + # A composite canonical-global save/load processes EVERY child, so the + # Hybrid may only advertise a checkpoint guarantee that every child + # supports: intersect (not union) the children's guarantee sets. (This + # is the opposite of the per-routed-parameter TRAINING claims, which are + # legitimately unioned.) None marks "no child seen yet" so the first + # child seeds each set and later children narrow it. + same_topology = None + topology_changing = None + topology_change_kinds = None + for _role, child in children: + supports = tuple( + support + for support in child.optimizer_contract().capabilities.checkpoints + if support.transport is CheckpointTransport.CANONICAL_GLOBAL + ) + if len(supports) != 1: + return frozenset(), frozenset(), frozenset() + support = supports[0] + child_same = set(support.same_topology) + child_changing = set(support.topology_changing) + child_kinds = set(support.topology_change_kinds) + same_topology = ( + child_same if same_topology is None else same_topology & child_same + ) + topology_changing = ( + child_changing + if topology_changing is None + else topology_changing & child_changing + ) + topology_change_kinds = ( + child_kinds + if topology_change_kinds is None + else topology_change_kinds & child_kinds + ) + return ( + frozenset(same_topology or ()), + frozenset(topology_changing or ()), + frozenset(topology_change_kinds or ()), + ) + except Exception: + return frozenset(), frozenset(), frozenset() + + +def _export_hybrid_portable_state( + optimizer, + *, + checkpoint_process_group, + transaction_id, + limits, +): + from gefen import portable_runtime as runtime + from gefen.portable_collective import _collective_unanimous_status + + transport = _hybrid_transport_binding(optimizer, checkpoint_process_group) + wire_limits = runtime._STATUS_FALLBACK_LIMITS + binding = None + normalized_transaction = None + normalized_limits = None + children = None + routing = None + context_digest = bytes(32) + live_token = None + error = None + try: + ( + binding, + normalized_transaction, + normalized_limits, + children, + routing, + _context, + context_digest, + live_token, + ) = _preflight_hybrid_portable_local( + optimizer, + checkpoint_process_group=checkpoint_process_group, + transaction_id=transaction_id, + limits=limits, + ) + wire_limits = normalized_limits._wire_limits() + except Exception as exc: + error = exc + _collective_unanimous_status( + transport, + error, + operation="hybrid_portable_export_preflight", + transaction_id=_HYBRID_EXPORT_PREFLIGHT_TRANSACTION, + context_digest=context_digest, + limits=wire_limits, + ) + assert ( + binding is not None + and normalized_transaction is not None + and normalized_limits is not None + and children is not None + and routing is not None + and live_token is not None + ) + + child_documents = {"muon": None, "backup": None} + for role, child in children: + child_documents[role] = child.export_portable_state( + checkpoint_process_group=binding, + transaction_id=_hybrid_child_transaction_id( + normalized_transaction, + role, + "export", + ), + limits=normalized_limits, + ) + + document = None + error = None + try: + document = build_hybrid_portable_state_document( + backup_optimizer=optimizer.backup_optimizer, + routing=routing, + children=child_documents, + ) + _validate_hybrid_portable_limits(document, normalized_limits) + _validate_hybrid_portable_readiness(optimizer, binding) + if live_token != _hybrid_portable_live_token(optimizer): + raise RuntimeError("live Hybrid state changed during portable export") + context_digest = bytes.fromhex(document["completion"]["digest"]) + except Exception as exc: + error = exc + _collective_unanimous_status( + binding, + error, + operation="hybrid_portable_export_finalize", + transaction_id=normalized_transaction, + context_digest=context_digest, + limits=wire_limits, + ) + assert document is not None + return document + + +def _import_hybrid_portable_state( + optimizer, + state, + *, + checkpoint_process_group, + transaction_id, + limits, +) -> None: + from gefen import portable_runtime as runtime + from gefen.portable_collective import _collective_unanimous_status + + transport = _hybrid_transport_binding(optimizer, checkpoint_process_group) + wire_limits = runtime._STATUS_FALLBACK_LIMITS + binding = None + normalized_transaction = None + normalized_limits = None + children = None + routing = None + base_context = None + context_digest = bytes(32) + composite_live_token = None + document = None + error = None + try: + ( + binding, + normalized_transaction, + normalized_limits, + children, + routing, + base_context, + context_digest, + composite_live_token, + ) = _preflight_hybrid_portable_local( + optimizer, + checkpoint_process_group=checkpoint_process_group, + transaction_id=transaction_id, + limits=limits, + ) + wire_limits = normalized_limits._wire_limits() + from gefen.portable_state import ( + _bounded_clone, + _normalize_gefen_portable_state_document, + ) + + document = normalize_hybrid_portable_state_document(_bounded_clone(state, normalized_limits, collective=True)) + for role in _HYBRID_PORTABLE_ROLES: + child_document = document["children"][role] + if child_document is not None: + document["children"][role] = _normalize_gefen_portable_state_document( + child_document, + limits=normalized_limits, + expected_implementation=_HYBRID_PORTABLE_CHILD_IMPLEMENTATIONS[role], + ) + if document["routing"] != routing: + raise ValueError("portable Hybrid document routing does not match the target") + expected_presence = { + role: child is not None for role, child in (("muon", optimizer.muon), ("backup", optimizer.backup)) + } + if any( + (document["children"][role] is not None) is not expected_presence[role] for role in _HYBRID_PORTABLE_ROLES + ): + raise ValueError("portable Hybrid child presence does not match the target") + context_digest = bytes.fromhex(document["completion"]["digest"]) + except Exception as exc: + error = exc + _collective_unanimous_status( + transport, + error, + operation="hybrid_portable_import_document", + transaction_id=_HYBRID_IMPORT_PREFLIGHT_TRANSACTION, + context_digest=context_digest, + limits=wire_limits, + ) + assert ( + binding is not None + and normalized_transaction is not None + and normalized_limits is not None + and children is not None + and routing is not None + and base_context is not None + and composite_live_token is not None + and document is not None + ) + + staged_children = [] + target_children = {} + error = None + try: + for role, child in children: + implementation = runtime._optimizer_implementation(child) + staged, live_token, target_fragment = runtime._stage_portable_import( + child, + implementation, + binding, + normalized_limits, + document["children"][role], + ) + staged_children.append((role, child, implementation, staged, live_token)) + target_children[role] = runtime._target_context( + runtime._base_context(binding, implementation), + target_fragment, + ) + target_context = { + **base_context, + "document_digest": document["completion"]["digest"], + "target_children": target_children, + } + runtime._preflight_portable_value(target_context, normalized_limits) + context_digest = runtime._context_digest(target_context) + except Exception as exc: + error = exc + context_digest = runtime._context_digest( + { + **base_context, + "document_digest": document["completion"]["digest"], + } + ) + _collective_unanimous_status( + binding, + error, + operation="hybrid_portable_import_prepare", + transaction_id=normalized_transaction, + context_digest=context_digest, + limits=wire_limits, + ) + assert len(staged_children) == len(children) + + target_deterministic = next( + document["children"][role]["common"]["gefen_deterministic"] for role, _child in children + ) + from gefen.gefen import Gefen + + commit_staged = Gefen._commit_staged_load_state_dict + + error = None + try: + _validate_hybrid_portable_readiness(optimizer, binding) + if composite_live_token != _hybrid_portable_live_token(optimizer): + raise RuntimeError("live Hybrid state changed after portable import preparation") + for _role, child, _implementation, _staged, live_token in staged_children: + if live_token != runtime._portable_live_token(child): + raise RuntimeError("live Hybrid child changed after portable import preparation") + except Exception as exc: + error = exc + _collective_unanimous_status( + binding, + error, + operation="hybrid_portable_import_freshness", + transaction_id=normalized_transaction, + context_digest=context_digest, + limits=wire_limits, + ) + + for _role, child, _implementation, staged, _live_token in staged_children: + commit_staged(child, staged) + dict.__setitem__( + optimizer.__dict__, + "_deterministic", + target_deterministic, + ) + + +__all__ = [ + "HYBRID_PORTABLE_STATE_COVERAGE", + "HYBRID_PORTABLE_STATE_FORMAT", + "HYBRID_PORTABLE_STATE_FORMAT_VERSION", + "HYBRID_PORTABLE_STATE_IMPLEMENTATION", + "build_hybrid_portable_state_document", + "normalize_hybrid_portable_state_document", +] diff --git a/src/gefen/portable_identity.py b/src/gefen/portable_identity.py new file mode 100644 index 0000000..1a1872d --- /dev/null +++ b/src/gefen/portable_identity.py @@ -0,0 +1,289 @@ +"""Strict primitive wire codecs for portable optimizer identities.""" + +from gefen.contracts import ( + LogicalRegion, + LogicalSlice, + ParameterIdentity, + ParameterLayout, + PlacementKind, + ProcessGroupIdentity, + ShardIdentity, + ShardPlacement, + ShardingManifest, +) + + +_PARAMETER_IDENTITY_KEYS = frozenset({"schema_version", "fqn", "global_shape"}) +_PROCESS_GROUP_IDENTITY_KEYS = frozenset({"schema_version", "semantic_name", "ordered_members"}) +_SHARD_IDENTITY_KEYS = frozenset( + { + "schema_version", + "parameter", + "layout", + "logical_extent", + "placements", + "process_group", + "local_member", + "owner", + } +) +_LOGICAL_SLICE_KEYS = frozenset({"kind", "flat_offset", "length"}) +_LOGICAL_REGION_KEYS = frozenset({"kind", "offsets", "lengths"}) +_SHARD_PLACEMENT_KEYS = frozenset( + { + "mesh_axis", + "kind", + "coordinate", + "parts", + "parameter_dimension", + } +) +_SHARDING_MANIFEST_KEYS = frozenset({"schema_version", "shards"}) + + +def _require_exact_record(value, keys, *, name): + if type(value) is not dict or set(value) != keys: + raise ValueError("{} has an invalid schema".format(name)) + return value + + +def _require_exact_list(value, *, name): + if type(value) is not list: + raise ValueError("{} must be a list".format(name)) + return value + + +def _require_optional_string(value, *, name): + if value is not None and type(value) is not str: + raise ValueError("{} must be a string or None".format(name)) + return value + + +def _serialize_parameter_identity(identity): + if not isinstance(identity, ParameterIdentity): + raise TypeError("identity must be a ParameterIdentity") + return { + "schema_version": identity.schema_version, + "fqn": identity.fqn, + "global_shape": list(identity.global_shape), + } + + +def _parse_parameter_identity(record): + record = _require_exact_record(record, _PARAMETER_IDENTITY_KEYS, name="parameter identity") + shape = _require_exact_list(record["global_shape"], name="parameter identity global_shape") + if type(record["schema_version"]) is not int: + raise ValueError("parameter identity schema_version must be an int") + if type(record["fqn"]) is not str: + raise ValueError("parameter identity fqn must be a string") + if any(type(dimension) is not int for dimension in shape): + raise ValueError("parameter identity global_shape must contain only integers") + try: + return ParameterIdentity( + fqn=record["fqn"], + global_shape=tuple(shape), + schema_version=record["schema_version"], + ) + except (TypeError, ValueError) as exc: + raise ValueError("parameter identity is invalid") from exc + + +def _normalize_parameter_identity(record): + return _serialize_parameter_identity(_parse_parameter_identity(record)) + + +def _serialize_process_group_identity(identity): + if not isinstance(identity, ProcessGroupIdentity): + raise TypeError("identity must be a ProcessGroupIdentity") + return { + "schema_version": identity.schema_version, + "semantic_name": identity.semantic_name, + "ordered_members": list(identity.ordered_members), + } + + +def _parse_process_group_identity(record): + record = _require_exact_record( + record, + _PROCESS_GROUP_IDENTITY_KEYS, + name="process-group identity", + ) + members = _require_exact_list( + record["ordered_members"], + name="process-group identity ordered_members", + ) + if type(record["schema_version"]) is not int: + raise ValueError("process-group identity schema_version must be an int") + if type(record["semantic_name"]) is not str: + raise ValueError("process-group identity semantic_name must be a string") + if any(type(member) is not str for member in members): + raise ValueError("process-group identity ordered_members must contain only strings") + try: + return ProcessGroupIdentity( + semantic_name=record["semantic_name"], + ordered_members=tuple(members), + schema_version=record["schema_version"], + ) + except (TypeError, ValueError) as exc: + raise ValueError("process-group identity is invalid") from exc + + +def _normalize_process_group_identity(record): + return _serialize_process_group_identity(_parse_process_group_identity(record)) + + +def _serialize_logical_extent(logical_extent): + if isinstance(logical_extent, LogicalSlice): + return { + "kind": "logical_slice", + "flat_offset": logical_extent.flat_offset, + "length": logical_extent.length, + } + if isinstance(logical_extent, LogicalRegion): + return { + "kind": "logical_region", + "offsets": list(logical_extent.offsets), + "lengths": list(logical_extent.lengths), + } + raise TypeError("logical extent must be a LogicalSlice or LogicalRegion") + + +def _parse_logical_extent(record): + if type(record) is not dict or type(record.get("kind")) is not str: + raise ValueError("shard logical_extent has an invalid schema") + kind = record["kind"] + if kind == "logical_slice": + _require_exact_record(record, _LOGICAL_SLICE_KEYS, name="shard logical_slice") + if type(record["flat_offset"]) is not int or type(record["length"]) is not int: + raise ValueError("shard logical_slice offsets and lengths must be integers") + try: + return LogicalSlice(record["flat_offset"], record["length"]) + except (TypeError, ValueError) as exc: + raise ValueError("shard logical_slice is invalid") from exc + if kind == "logical_region": + _require_exact_record(record, _LOGICAL_REGION_KEYS, name="shard logical_region") + offsets = _require_exact_list(record["offsets"], name="shard logical_region offsets") + lengths = _require_exact_list(record["lengths"], name="shard logical_region lengths") + if any(type(value) is not int for value in offsets + lengths): + raise ValueError("shard logical_region offsets and lengths must contain only integers") + try: + return LogicalRegion(tuple(offsets), tuple(lengths)) + except (TypeError, ValueError) as exc: + raise ValueError("shard logical_region is invalid") from exc + raise ValueError("unsupported shard logical_extent kind") + + +def _serialize_shard_placement(placement): + if not isinstance(placement, ShardPlacement): + raise TypeError("placement must be a ShardPlacement") + return { + "mesh_axis": placement.mesh_axis, + "kind": placement.kind.value, + "coordinate": placement.coordinate, + "parts": placement.parts, + "parameter_dimension": placement.parameter_dimension, + } + + +def _parse_shard_placement(record): + record = _require_exact_record(record, _SHARD_PLACEMENT_KEYS, name="shard placement") + if type(record["mesh_axis"]) is not str: + raise ValueError("shard placement mesh_axis must be a string") + if type(record["kind"]) is not str: + raise ValueError("shard placement kind must be a string") + if type(record["coordinate"]) is not int: + raise ValueError("shard placement coordinate must be an int") + if type(record["parts"]) is not int: + raise ValueError("shard placement parts must be an int") + if record["parameter_dimension"] is not None and type(record["parameter_dimension"]) is not int: + raise ValueError("shard placement parameter_dimension must be an int or None") + try: + kind = PlacementKind(record["kind"]) + return ShardPlacement( + mesh_axis=record["mesh_axis"], + kind=kind, + coordinate=record["coordinate"], + parts=record["parts"], + parameter_dimension=record["parameter_dimension"], + ) + except (TypeError, ValueError) as exc: + raise ValueError("shard placement is invalid") from exc + + +def _serialize_shard_identity(identity): + if not isinstance(identity, ShardIdentity): + raise TypeError("identity must be a ShardIdentity") + return { + "schema_version": identity.schema_version, + "parameter": _serialize_parameter_identity(identity.parameter), + "layout": identity.layout.value, + "logical_extent": _serialize_logical_extent(identity.logical_slice), + "placements": [_serialize_shard_placement(placement) for placement in identity.placements], + "process_group": ( + None if identity.process_group is None else _serialize_process_group_identity(identity.process_group) + ), + "local_member": identity.local_member, + "owner": identity.owner, + } + + +def _parse_shard_identity(record): + record = _require_exact_record(record, _SHARD_IDENTITY_KEYS, name="shard identity") + if type(record["schema_version"]) is not int: + raise ValueError("shard identity schema_version must be an int") + if type(record["layout"]) is not str: + raise ValueError("shard identity layout must be a string") + placements = _require_exact_list(record["placements"], name="shard identity placements") + _require_optional_string(record["local_member"], name="shard identity local_member") + _require_optional_string(record["owner"], name="shard identity owner") + try: + parameter = _parse_parameter_identity(record["parameter"]) + logical_extent = _parse_logical_extent(record["logical_extent"]) + parsed_placements = tuple(_parse_shard_placement(placement) for placement in placements) + process_group = ( + None if record["process_group"] is None else _parse_process_group_identity(record["process_group"]) + ) + layout = ParameterLayout(record["layout"]) + return ShardIdentity( + parameter=parameter, + layout=layout, + logical_slice=logical_extent, + placements=parsed_placements, + process_group=process_group, + local_member=record["local_member"], + owner=record["owner"], + schema_version=record["schema_version"], + ) + except (TypeError, ValueError) as exc: + raise ValueError("shard identity is invalid") from exc + + +def _normalize_shard_identity(record): + return _serialize_shard_identity(_parse_shard_identity(record)) + + +def _serialize_sharding_manifest(manifest): + if not isinstance(manifest, ShardingManifest): + raise TypeError("manifest must be a ShardingManifest") + return { + "schema_version": manifest.schema_version, + "shards": [_serialize_shard_identity(shard) for shard in manifest.shards], + } + + +def _parse_sharding_manifest(record): + record = _require_exact_record(record, _SHARDING_MANIFEST_KEYS, name="sharding manifest") + if type(record["schema_version"]) is not int: + raise ValueError("sharding manifest schema_version must be an int") + shards = _require_exact_list(record["shards"], name="sharding manifest shards") + try: + return ShardingManifest( + shards=tuple(_parse_shard_identity(shard) for shard in shards), + schema_version=record["schema_version"], + ) + except (TypeError, ValueError) as exc: + raise ValueError("sharding manifest is invalid") from exc + + +def _normalize_sharding_manifest(record): + return _serialize_sharding_manifest(_parse_sharding_manifest(record)) diff --git a/src/gefen/portable_runtime.py b/src/gefen/portable_runtime.py new file mode 100644 index 0000000..353fdc0 --- /dev/null +++ b/src/gefen/portable_runtime.py @@ -0,0 +1,1908 @@ +"""Optimizer-facing collective runtime for exact portable Gefen state.""" + +from __future__ import annotations + +from collections import defaultdict +import hashlib +import math + +import torch +from torch import nn + +from gefen.canonical import canonical_value_supported +from gefen.checkpoint import CheckpointProcessGroupBinding +from gefen.codebook import CodebookProcessGroupBinding +from gefen.contracts import ( + IDENTITY_SCHEMA_VERSION, + LogicalSlice, + ParameterIdentity, + ParameterLayout, + PlacementKind, + ProcessGroupIdentity, + ShardIdentity, + ShardPlacement, + ShardingManifest, +) +from gefen.portable import ( + _decode_quantized_momentum, + _expand_block_second_moment, + _read_flat_chunk, +) +from gefen.portable_collective import ( + _collective_unanimous_status, + _collective_visit_canonical_fragments, +) +from gefen.portable_identity import ( + _serialize_parameter_identity, + _serialize_process_group_identity, + _serialize_shard_identity, + _serialize_sharding_manifest, +) +from gefen.portable_schema import portable_state_digest +from gefen.portable_state import ( + PortableStateLimits, + _MOMENTUM_PROJECTION, + _SECOND_MOMENT_PROJECTION, + _assemble_portable_state_fragments, + _build_portable_state_fragment, + _derived_role, + _local_dense_shape, + _normalize_common, + _normalize_gefen_portable_state_document, + _normalize_options, + _normalize_policy, + _normalize_portable_state_fragment, + _project_portable_parameter_state, + _values_equal, +) +from gefen.rebinding import LogicalSlotBinding + + +_PLAIN_IMPLEMENTATION = "gefen.Gefen" +_MUON_IMPLEMENTATION = "gefen.GefenMuon" +_MAX_TRANSACTION_BYTES = 1024 +_EXPORT_PREFLIGHT_TRANSACTION = "gefen-portable-export-preflight-v1" +_IMPORT_PREFLIGHT_TRANSACTION = "gefen-portable-import-preflight-v1" +_STATUS_FALLBACK_LIMITS = PortableStateLimits( + max_fragment_tensor_bytes=1, + max_collective_tensor_bytes=1, + max_collective_metadata_bytes=1, + max_metadata_bytes=1, +)._wire_limits() + +_PLAIN_GROUP_REQUIRED = frozenset({"params", "param_names", "lr", "beta1", "beta2", "eps", "weight_decay"}) +_PLAIN_GROUP_ALLOWED = _PLAIN_GROUP_REQUIRED | {"name"} +_MUON_GROUP_REQUIRED = _PLAIN_GROUP_REQUIRED | frozenset( + { + "momentum", + "nesterov", + "ns_coefficients", + "ns_steps", + "adjust_lr_fn", + "sharded_mode", + "fp8_ns", + "fp8_ns_compile", + "batched_ns", + "batched_ns_workspace_bytes", + "normuon", + "normuon_beta2", + "normuon_eps", + "cautious", + } +) +_MUON_GROUP_ALLOWED = _MUON_GROUP_REQUIRED | {"name", "ns_schedule"} + +_AUTHORITATIVE_STATE_KEYS = frozenset( + { + "name", + "automatic_period", + "step", + "m_codebook", + "m_magnitude", + "vmean", + "vmean_step", + "v_row", + "v_col", + "factored_step", + "normuon_v", + "normuon_step", + } +) +_IGNORED_DERIVED_STATE_KEYS = frozenset({"stepsize", "_h_buf", "m_codebook_shape"}) +_PLAIN_BLOCK_KEYS = frozenset( + { + "name", + "automatic_period", + "step", + "m_codebook", + "m_magnitude", + "vmean", + "vmean_step", + } +) +_PLAIN_FACTORED_KEYS = frozenset( + { + "name", + "automatic_period", + "step", + "m_codebook", + "m_magnitude", + "v_row", + "v_col", + "factored_step", + } +) +_MUON_KEYS = frozenset({"name", "automatic_period", "step", "m_codebook", "m_magnitude"}) +_NORMUON_KEYS = _MUON_KEYS | {"normuon_v", "normuon_step"} + + +def _bounded_utf8_length(value: str, *, limit: int, name: str) -> int: + total = 0 + for start in range(0, len(value), 4096): + chunk = value[start : start + 4096].encode("utf-8") + if total > limit - len(chunk): + raise ValueError("{} exceeds its UTF-8 byte limit".format(name)) + total += len(chunk) + return total + + +class _PreflightTensor: + __slots__ = ("shape",) + + def __init__(self, shape): + self.shape = tuple(shape) + + +def _preflight_portable_value(value, limits: PortableStateLimits) -> None: + """Apply exact wire structural limits without allocating tensor payloads.""" + + wire = limits._wire_limits() + nodes = 0 + tensor_bytes = 0 + tensor_shapes = [] + metadata_bytes = 16 + + def add_metadata(count): + nonlocal metadata_bytes + metadata_bytes += count + if metadata_bytes > wire.max_metadata_bytes: + raise ValueError("portable fragment exceeds max_metadata_bytes before materialization") + + def visit(item, *, depth): + nonlocal nodes, tensor_bytes + if depth > wire.max_tree_depth: + raise ValueError("portable fragment exceeds max_tree_depth before materialization") + nodes += 1 + if nodes > wire.max_tree_nodes: + raise ValueError("portable fragment exceeds max_tree_nodes before materialization") + item_type = type(item) + if item is None or item_type is bool: + add_metadata(1) + return + if item_type is int: + magnitude_bytes = (abs(item).bit_length() + 7) // 8 + if magnitude_bytes > wire.max_integer_bytes: + raise ValueError("portable fragment integer exceeds max_integer_bytes") + add_metadata(10 + magnitude_bytes) + return + if item_type is float: + if not math.isfinite(item): + raise ValueError("portable fragment contains a nonfinite float") + add_metadata(9) + return + if item_type is str: + encoded_bytes = _bounded_utf8_length( + item, + limit=wire.max_string_bytes, + name="portable fragment string", + ) + add_metadata(9 + encoded_bytes) + return + if item_type is _PreflightTensor: + shape = item.shape + if any(type(dimension) is not int or dimension < 0 for dimension in shape): + raise ValueError("portable fragment tensor has invalid geometry") + if len(shape) > wire.max_tensor_rank: + raise ValueError("portable fragment tensor exceeds max_tensor_rank") + if len(tensor_shapes) >= wire.max_tensors: + raise ValueError("portable fragment exceeds max_tensors") + numel = math.prod(shape) + nbytes = numel * 4 + if tensor_bytes > wire.max_fragment_tensor_bytes - nbytes: + raise ValueError("portable fragment exceeds max_fragment_tensor_bytes before materialization") + tensor_bytes += nbytes + tensor_shapes.append(shape) + add_metadata(9) + return + if item_type in {list, tuple}: + if len(item) > wire.max_container_items: + raise ValueError("portable fragment container exceeds max_container_items") + add_metadata(9) + for child in item: + visit(child, depth=depth + 1) + return + if item_type is dict: + if len(item) > wire.max_container_items: + raise ValueError("portable fragment container exceeds max_container_items") + if any(type(key) is not str for key in item): + raise TypeError("portable fragment dictionary keys must be strings") + add_metadata(9) + for key in sorted(item): + visit(key, depth=depth + 1) + visit(item[key], depth=depth + 1) + return + raise TypeError("portable preflight encountered an unsupported value") + + visit(value, depth=0) + add_metadata(8) + for shape in tensor_shapes: + add_metadata(56 + 8 * len(shape)) + + +def _optimizer_implementation(optimizer) -> str: + # Keep imports lazy so Gefen may delegate to this module without creating a + # module-import cycle. + from gefen.gefen import Gefen + from gefen.gefen_muon import GefenMuon + + if type(optimizer) is Gefen: + return _PLAIN_IMPLEMENTATION + if type(optimizer) is GefenMuon: + return _MUON_IMPLEMENTATION + raise TypeError("portable state supports exact Gefen and GefenMuon instances only") + + +def _reject_method_shadows(optimizer) -> None: + if type(optimizer.__dict__) is not dict: + raise TypeError("portable state requires an exact optimizer attribute mapping") + for name in optimizer.__dict__: + descriptor = None + for owner in type(optimizer).__mro__: + if name in owner.__dict__: + descriptor = owner.__dict__[name] + break + if isinstance(descriptor, (staticmethod, classmethod)): + descriptor = descriptor.__func__ + if callable(descriptor): + raise TypeError("portable state rejects instance-level method shadows") + + +def _require_limits(limits) -> PortableStateLimits: + if type(limits) is not PortableStateLimits: + raise TypeError("limits must be a PortableStateLimits") + return limits + + +def _require_binding(binding) -> CheckpointProcessGroupBinding: + if type(binding) is not CheckpointProcessGroupBinding: + raise TypeError("checkpoint_process_group must be a CheckpointProcessGroupBinding") + return binding + + +def _validate_exact_process_group_identity(identity) -> ProcessGroupIdentity: + if ( + type(identity) is not ProcessGroupIdentity + or type(identity.semantic_name) is not str + or type(identity.ordered_members) is not tuple + or any(type(member) is not str for member in identity.ordered_members) + or type(identity.schema_version) is not int + ): + raise TypeError("portable state requires an exact ProcessGroupIdentity") + validated = ProcessGroupIdentity( + identity.semantic_name, + identity.ordered_members, + schema_version=identity.schema_version, + ) + if validated != identity: + raise ValueError("portable process-group identities must be canonical") + return validated + + +def _validate_exact_shard_identity(shard) -> ShardIdentity: + if type(shard) is not ShardIdentity: + raise TypeError("portable state requires exact ShardIdentity values") + if type(shard.parameter) is not ParameterIdentity: + raise TypeError("portable state requires exact ParameterIdentity values") + if ( + type(shard.parameter.fqn) is not str + or type(shard.parameter.global_shape) is not tuple + or any(type(dimension) is not int for dimension in shard.parameter.global_shape) + or type(shard.parameter.schema_version) is not int + ): + raise TypeError("portable parameter identities require exact primitive fields") + parameter = ParameterIdentity( + shard.parameter.fqn, + shard.parameter.global_shape, + schema_version=shard.parameter.schema_version, + ) + if type(shard.logical_slice) is not LogicalSlice: + raise TypeError("portable state requires exact LogicalSlice values") + if ( + type(shard.logical_slice.flat_offset) is not int + or type(shard.logical_slice.length) is not int + or type(shard.local_member) is not str + or (shard.owner is not None and type(shard.owner) is not str) + or type(shard.schema_version) is not int + ): + raise TypeError("portable shard identities require exact primitive fields") + logical_slice = LogicalSlice( + shard.logical_slice.flat_offset, + shard.logical_slice.length, + ) + process_group = _validate_exact_process_group_identity(shard.process_group) + if type(shard.placements) is not tuple or any( + type(placement) is not ShardPlacement for placement in shard.placements + ): + raise TypeError("portable state requires exact ShardPlacement values") + if any( + type(placement.mesh_axis) is not str + or type(placement.kind) is not PlacementKind + or type(placement.coordinate) is not int + or type(placement.parts) is not int + or (placement.parameter_dimension is not None and type(placement.parameter_dimension) is not int) + for placement in shard.placements + ): + raise TypeError("portable shard placements require exact primitive fields") + placements = tuple( + ShardPlacement( + placement.mesh_axis, + placement.kind, + placement.coordinate, + placement.parts, + placement.parameter_dimension, + ) + for placement in shard.placements + ) + if type(shard.layout) is not ParameterLayout: + raise TypeError("portable shard layouts require exact ParameterLayout values") + validated = ShardIdentity( + parameter, + shard.layout, + logical_slice, + placements=placements, + process_group=process_group, + local_member=shard.local_member, + owner=shard.owner, + schema_version=shard.schema_version, + ) + if validated != shard: + raise ValueError("portable shard identities must be canonical") + return validated + + +def _validate_exact_manifest(manifest) -> None: + if ( + type(manifest) is not ShardingManifest + or type(manifest.shards) is not tuple + or type(manifest.schema_version) is not int + ): + raise TypeError("portable state requires an exact ShardingManifest") + validated = ShardingManifest( + tuple(_validate_exact_shard_identity(shard) for shard in manifest.shards), + schema_version=manifest.schema_version, + ) + if validated != manifest: + raise ValueError("portable sharding manifests must be canonical") + + +def _require_transaction_id(transaction_id) -> str: + if ( + type(transaction_id) is not str + or not transaction_id + or transaction_id != transaction_id.strip() + or "\x00" in transaction_id + ): + raise ValueError("transaction_id must be non-empty, trimmed, and without NUL") + _bounded_utf8_length( + transaction_id, + limit=_MAX_TRANSACTION_BYTES, + name="transaction_id", + ) + return transaction_id + + +def _preflight_transport_binding(optimizer, supplied): + scope = getattr(optimizer, "_gefen_codebook_process_group", None) + if type(scope) is CodebookProcessGroupBinding: + identity = ProcessGroupIdentity( + str(scope.identity.semantic_name), + tuple(str(member) for member in scope.identity.ordered_members), + schema_version=IDENTITY_SCHEMA_VERSION, + ) + return CheckpointProcessGroupBinding( + identity, + str(scope.local_member), + scope.process_group, + scope.collective_device, + ) + if isinstance(supplied, CheckpointProcessGroupBinding): + return supplied + raise TypeError("portable preflight requires a transport-usable process group") + + +def _validate_supplied_binding(supplied, transport) -> None: + supplied = _require_binding(supplied) + _validate_exact_process_group_identity(supplied.identity) + _validate_exact_process_group_identity(transport.identity) + if ( + type(supplied.local_member) is not str + or type(supplied.collective_device) is not torch.device + or type(transport.local_member) is not str + or type(transport.collective_device) is not torch.device + ): + raise TypeError("portable checkpoint bindings require exact primitive fields") + if supplied.identity != transport.identity or supplied.local_member != transport.local_member: + raise ValueError("checkpoint and codebook process-group identities must match") + if ( + supplied.process_group is not transport.process_group + or supplied.collective_device != transport.collective_device + ): + raise ValueError("checkpoint_process_group must exactly match the optimizer-owned runtime transport") + + +def _validate_context_identity_bounds( + binding: CheckpointProcessGroupBinding, + limits: PortableStateLimits, +) -> None: + identity = binding.identity + member_count = len(identity.ordered_members) + if member_count > limits.max_members or member_count > limits.max_container_items: + raise ValueError("checkpoint process-group identity exceeds member limits") + _bounded_utf8_length( + identity.semantic_name, + limit=limits.max_string_bytes, + name="process-group semantic name", + ) + for member in identity.ordered_members: + _bounded_utf8_length( + member, + limit=limits.max_string_bytes, + name="process-group member", + ) + + +def _context_digest(value) -> bytes: + return bytes.fromhex(portable_state_digest(value)) + + +def _base_context(binding: CheckpointProcessGroupBinding, implementation: str): + return { + "format": "gefen.portable_runtime_context", + "format_version": 1, + "implementation": implementation, + "checkpoint_process_group": _serialize_process_group_identity(binding.identity), + } + + +def _plain_tensor( + value, + *, + name: str, + dtype=None, + shape=None, + nonnegative=False, + validate_values=True, +): + if ( + type(value) is not torch.Tensor + or value.layout is not torch.strided + or value.is_meta + or value.is_nested + or value.is_quantized + or value.requires_grad + or value.is_conj() + or value.is_neg() + ): + raise TypeError("{} must be a plain materialized strided tensor".format(name)) + if dtype is not None and value.dtype != dtype: + raise TypeError("{} has an invalid dtype".format(name)) + if shape is not None and tuple(value.shape) != tuple(shape): + raise ValueError("{} has invalid geometry".format(name)) + if value.is_floating_point() and validate_values: + for start in range(0, value.numel(), 1 << 20): + chunk = _read_flat_chunk( + value, + start, + min(start + (1 << 20), value.numel()), + ) + if not bool(torch.isfinite(chunk).all()): + raise ValueError("{} must be finite".format(name)) + if nonnegative and not bool((chunk >= 0).all()): + raise ValueError("{} must be nonnegative".format(name)) + return value + + +def _tight_cpu_fp32(value, *, name: str, shape=None, nonnegative=False): + value = _plain_tensor( + value, + name=name, + dtype=torch.float32, + shape=shape, + nonnegative=nonnegative, + ) + result = torch.empty(tuple(value.shape), dtype=torch.float32, device="cpu") + result.copy_(value.detach()) + return result + + +def _strict_counter(value, *, name: str, minimum=0): + if type(value) is not int or value < minimum or value > (1 << 53) - 1: + raise ValueError("{} must be an exact bounded host int".format(name)) + return value + + +def _live_policy(optimizer, implementation: str): + raw = optimizer._canonical_policy() + if type(raw) is not dict: + raise ValueError("live portable policy must be a plain dictionary") + policy = { + "schema_version": 1, + "factored_v_2d": raw.get("factored_v_2d"), + "force_1d_period_one": raw.get("force_1d_period_one"), + "force_2d_period_one": raw.get("force_2d_period_one"), + "period_one_substrings": raw.get("period_one_substrings"), + "codebook_refresh_every": raw.get("codebook_refresh_every"), + "stochastic_round": raw.get("stochastic_round"), + "momentum_projection": _MOMENTUM_PROJECTION, + "second_moment_projection": _SECOND_MOMENT_PROJECTION, + } + return _normalize_policy(policy, implementation) + + +def _normalized_ns_schedule(group): + from gefen.gefen_muon import _normalize_ns_schedule + + schedule = _normalize_ns_schedule(group["ns_coefficients"], group["ns_steps"]) + return [[float(a), float(b), float(c)] for a, b, c in schedule] + + +def _group_options(group, implementation: str, *, second_moment_policy=None): + if type(group) is not dict: + raise TypeError("portable parameter groups must be plain dictionaries") + allowed = _PLAIN_GROUP_ALLOWED if implementation == _PLAIN_IMPLEMENTATION else _MUON_GROUP_ALLOWED + required = _PLAIN_GROUP_REQUIRED if implementation == _PLAIN_IMPLEMENTATION else _MUON_GROUP_REQUIRED + if not required.issubset(group) or not set(group).issubset(allowed): + raise ValueError("portable parameter group contains missing or unknown keys") + if type(group["params"]) is not list or type(group["param_names"]) is not list: + raise TypeError("portable parameter group params and names must be plain lists") + + if implementation == _PLAIN_IMPLEMENTATION: + return _normalize_options( + { + "lr": group["lr"], + "beta1": group["beta1"], + "beta2": group["beta2"], + "eps": group["eps"], + "weight_decay": group["weight_decay"], + "second_moment_policy": second_moment_policy, + }, + implementation, + ) + + if any(type(group[key]) is not float for key in ("beta1", "beta2", "momentum")): + raise ValueError("Muon inherited betas and momentum must be exact floats") + if group["beta1"] != group["momentum"] or group["beta2"] != 0.0: + raise ValueError("Muon inherited beta values disagree with its momentum policy") + if "ns_schedule" in group and not canonical_value_supported(group["ns_schedule"], finite_tensors=True): + raise ValueError("Muon raw ns_schedule must remain weights-only-safe metadata") + return _normalize_options( + { + "lr": group["lr"], + "weight_decay": group["weight_decay"], + "momentum": group["momentum"], + "nesterov": group["nesterov"], + "ns_schedule": _normalized_ns_schedule(group), + "ns_eps": group["eps"], + "adjust_lr_fn": group["adjust_lr_fn"], + "sharded_mode": group["sharded_mode"], + "fp8_ns": group["fp8_ns"], + "fp8_ns_compile": group["fp8_ns_compile"], + "batched_ns": group["batched_ns"], + "batched_ns_workspace_bytes": group["batched_ns_workspace_bytes"], + "normuon": group["normuon"], + "normuon_beta2": group["normuon_beta2"], + "normuon_eps": group["normuon_eps"], + "cautious": group["cautious"], + }, + implementation, + ) + + +def _parameter_supported(parameter, shard: ShardIdentity) -> bool: + if type(parameter) not in {torch.Tensor, nn.Parameter} or not ( + parameter.layout is torch.strided + and parameter.device.type in {"cpu", "cuda"} + and parameter.dtype in {torch.float16, torch.bfloat16, torch.float32, torch.float64} + and not torch.is_complex(parameter) + and not parameter.is_meta + and not parameter.is_nested + and not parameter.is_quantized + and getattr(parameter, "fake_mode", None) is None + ): + return False + if shard.layout in { + ParameterLayout.REPLICATED, + ParameterLayout.WHOLE_PARAMETER_OWNER, + }: + return ( + tuple(parameter.shape) == shard.parameter.global_shape and parameter.numel() == shard.logical_slice.length + ) + if shard.layout is ParameterLayout.FLATTENED_ELEMENT_SHARD: + return parameter.ndim == 1 and parameter.is_contiguous() and parameter.numel() == shard.logical_slice.length + return False + + +def _validate_optimizer_shell( + optimizer, + implementation: str, + binding: CheckpointProcessGroupBinding, +) -> None: + _reject_method_shadows(optimizer) + _validate_exact_manifest(optimizer._gefen_sharding_manifest) + optimizer._assert_finalized_binding_layout() + if type(optimizer.defaults) is not dict: + raise TypeError("portable state requires exact built-in optimizer defaults") + if type(optimizer.param_groups) is not list: + raise TypeError("portable state requires an exact parameter-group list") + if optimizer.capturable is not False: + raise RuntimeError("portable state does not support capturable optimizers") + if optimizer._stochastic_round is not False: + raise RuntimeError("portable state does not support stochastic rounding") + if optimizer._capt_stacks is not None: + raise RuntimeError("portable state requires inactive capturable stacks") + if optimizer._gefen_global_step_by_device or optimizer._sr_seed_by_device: + raise RuntimeError("portable state does not support device-authoritative counters") + if torch.compiler.is_compiling(): + raise RuntimeError("portable state cannot run while compiling") + if torch.cuda.is_available() and torch.cuda.is_current_stream_capturing(): + raise RuntimeError("portable state cannot run during CUDA capture") + if type(optimizer._gefen_logical_slots) is not tuple or not optimizer._gefen_logical_slots: + raise ValueError("portable state requires immutable logical slots") + if type(optimizer._gefen_local_shard_bindings) is not tuple: + raise TypeError("portable state requires immutable local shard bindings") + + scope = optimizer._gefen_codebook_process_group + if type(scope) is not CodebookProcessGroupBinding: + raise RuntimeError("portable state requires an explicit codebook process group") + _validate_exact_process_group_identity(scope.identity) + if type(scope.local_member) is not str or type(scope.collective_device) is not torch.device: + raise TypeError("portable codebook bindings require exact primitive fields") + if scope.identity != binding.identity or scope.local_member != binding.local_member: + raise ValueError("checkpoint and codebook process-group identities must match") + optimizer._validate_codebook_runtime_binding(scope) + for shard in optimizer._gefen_sharding_manifest.shards: + if shard.process_group != binding.identity: + raise ValueError("portable manifest process groups must match the checkpoint binding") + if implementation == _PLAIN_IMPLEMENTATION: + if shard.layout not in { + ParameterLayout.REPLICATED, + ParameterLayout.FLATTENED_ELEMENT_SHARD, + }: + raise ValueError("plain portable state has an unsupported layout") + elif shard.layout not in { + ParameterLayout.REPLICATED, + ParameterLayout.WHOLE_PARAMETER_OWNER, + }: + raise ValueError("Muon portable state has an unsupported layout") + for slot in optimizer._gefen_logical_slots: + if type(slot) is not LogicalSlotBinding: + raise TypeError("portable logical slots must be exact LogicalSlotBinding values") + _validate_exact_shard_identity(slot.shard) + if slot.shard.local_member != binding.local_member: + raise ValueError("local logical slots must match the checkpoint member") + + +def _live_maps(optimizer, binding: CheckpointProcessGroupBinding): + local_by_fqn = {} + for item in optimizer._gefen_local_shard_bindings: + if type(item) is not tuple or len(item) != 2: + raise TypeError("portable local bindings must be parameter/shard tuples") + parameter, shard = item + _validate_exact_shard_identity(shard) + if shard.process_group != binding.identity: + raise ValueError("portable local shard process groups must exactly match the checkpoint binding") + if shard.local_member != binding.local_member: + raise ValueError("portable local shards must match the checkpoint member") + if shard.parameter.fqn in local_by_fqn: + raise ValueError("portable local bindings contain duplicate FQNs") + if parameter is not None and not _parameter_supported(parameter, shard): + raise TypeError("portable state requires supported live shard geometry and dtype") + local_by_fqn[shard.parameter.fqn] = (parameter, shard) + + expected_live = { + parameter + for parameter, shard in local_by_fqn.values() + if not (shard.layout is ParameterLayout.WHOLE_PARAMETER_OWNER and shard.local_member != shard.owner) + } + if None in expected_live: + raise TypeError("portable state requires plain materialized local parameters") + state_type = type(optimizer.state) + if not (state_type is dict or (state_type is defaultdict and optimizer.state.default_factory is dict)): + raise TypeError("portable optimizer state must use a standard mapping") + if set(optimizer.state) != expected_live: + raise ValueError("portable optimizer state contains missing or foreign parameter keys") + if type(optimizer._param_names) is not dict or set(optimizer._param_names) != expected_live: + raise ValueError("portable parameter-name state does not match live parameters") + return local_by_fqn + + +def _parameter_storage_token(parameter): + try: + version = parameter._version + except RuntimeError: + version = None + return ( + id(parameter), + version, + str(parameter.device), + str(parameter.dtype), + str(parameter.layout), + tuple(parameter.shape), + tuple(parameter.stride()), + parameter.storage_offset(), + parameter.numel(), + parameter.is_contiguous(), + parameter.untyped_storage().data_ptr(), + parameter.untyped_storage().nbytes(), + ) + + +def _portable_value_token(value): + if type(value) is torch.Tensor: + hasher = hashlib.sha256() + element_size = value.element_size() + chunk_elements = max(1, (1 << 20) // max(1, element_size)) + for start in range(0, value.numel(), chunk_elements): + stop = min(start + chunk_elements, value.numel()) + chunk = _read_flat_chunk(value, start, stop).contiguous().cpu() + hasher.update(memoryview(chunk.view(torch.uint8).numpy())) + try: + version = value._version + except RuntimeError: + version = None + return ( + "tensor", + id(value), + version, + hasher.digest(), + str(value.device), + str(value.dtype), + str(value.layout), + tuple(value.shape), + tuple(value.stride()), + value.storage_offset(), + value.requires_grad, + value.is_conj(), + value.is_neg(), + ) + if type(value) is dict: + return ( + "dict", + tuple((key, _portable_value_token(value[key])) for key in sorted(value, key=repr)), + ) + if type(value) in {list, tuple}: + return ( + type(value).__name__, + tuple(_portable_value_token(item) for item in value), + ) + try: + hash(value) + token = value + except TypeError: + token = (id(value), repr(value)) + return (type(value).__name__, token) + + +def _portable_live_token(optimizer): + _reject_method_shadows(optimizer) + + def live_group_options(group): + options = optimizer._canonical_group_options_value(group) + if type(options) is dict: + options = dict(options) + options.pop("ns_schedule", None) + return options + + groups = tuple( + ( + id(group), + tuple(id(parameter) for parameter in group["params"]), + tuple(str(name) for name, _parameter in optimizer._iter_group_params_with_names(group)), + _portable_value_token(live_group_options(group)), + ) + for group in optimizer.param_groups + ) + state_entries = tuple( + sorted( + ( + id(parameter), + type(parameter), + id(state), + _portable_value_token(state), + ) + for parameter, state in optimizer.state.items() + ) + ) + local_storage = tuple( + None if parameter is None else _parameter_storage_token(parameter) + for parameter, _shard in optimizer._gefen_local_shard_bindings + ) + return ( + id(optimizer.param_groups), + id(optimizer.state), + id(optimizer.defaults), + type(optimizer.defaults), + type(optimizer.param_groups), + type(optimizer.state), + _portable_value_token(optimizer.defaults), + groups, + state_entries, + _portable_value_token(optimizer._param_names), + local_storage, + optimizer._gefen_global_step, + _portable_value_token(optimizer._gefen_codebook), + _portable_value_token(optimizer._gefen_global_step_by_device), + _portable_value_token(optimizer._sr_seed_by_device), + _portable_value_token(optimizer._canonical_policy()), + optimizer._deterministic, + optimizer.capturable, + optimizer.fused, + optimizer.verbose, + optimizer._fused_build_ok, + id(optimizer._gefen_codebook_process_group), + _portable_value_token(optimizer._serialized_codebook_scope()), + id(optimizer._gefen_sharding_manifest), + _portable_value_token(_serialize_sharding_manifest(optimizer._gefen_sharding_manifest)), + id(optimizer._gefen_logical_slots), + tuple( + ( + id(slot), + slot.group_index, + slot.original_slot_index, + slot.compatibility_name, + _portable_value_token(_serialize_shard_identity(slot.shard)), + ) + for slot in optimizer._gefen_logical_slots + ), + tuple( + ( + None if parameter is None else id(parameter), + shard.sort_key, + ) + for parameter, shard in optimizer._gefen_local_shard_bindings + ), + ) + + +def _catalog_and_options(optimizer, implementation: str, policy): + group_options = [] + for group in optimizer.param_groups: + # Plain options depend on each logical parameter rank, so validate the + # group shell here and construct per-slot values below. + if implementation == _PLAIN_IMPLEMENTATION: + _group_options(group, implementation, second_moment_policy="block") + group_options.append(group) + else: + group_options.append(_group_options(group, implementation)) + + catalog = {} + options_by_fqn = {} + for slot in optimizer._gefen_logical_slots: + identity = slot.shard.parameter + if slot.group_index >= len(group_options): + raise ValueError("portable logical slot group index is out of range") + if implementation == _PLAIN_IMPLEMENTATION: + second = "factored" if policy["factored_v_2d"] and len(identity.global_shape) == 2 else "block" + options = _group_options( + group_options[slot.group_index], + implementation, + second_moment_policy=second, + ) + if second == "factored" and slot.shard.layout is not ParameterLayout.REPLICATED: + raise ValueError("factored portable state supports replicated matrices only") + else: + options = group_options[slot.group_index] + if slot.shard.layout is ParameterLayout.WHOLE_PARAMETER_OWNER and options["sharded_mode"] != "distributed": + raise ValueError("Muon whole-owner portable state requires distributed mode") + fqn = identity.fqn + if fqn in catalog: + raise ValueError("portable logical slots contain duplicate FQNs") + options_by_fqn[fqn] = options + catalog[fqn] = { + "identity": _serialize_parameter_identity(identity), + "algorithm_options": options, + } + return catalog, options_by_fqn + + +def _common_state(optimizer): + global_step = _strict_counter(optimizer._canonical_common_global_step(), name="gefen_global_step") + codebook = optimizer._gefen_codebook + codebook_cpu = None + if codebook is not None: + codebook_cpu = _tight_cpu_fp32( + codebook, + name="gefen_codebook", + shape=(256,), + ) + return _normalize_common( + { + "gefen_global_step": global_step, + "gefen_codebook": codebook_cpu, + "gefen_deterministic": optimizer._deterministic, + } + ) + + +def _preflight_common_state(optimizer, *, validate_values=True): + global_step = _strict_counter( + optimizer._canonical_common_global_step(), + name="gefen_global_step", + ) + codebook = optimizer._gefen_codebook + codebook_preview = None + if codebook is not None: + codebook = _plain_tensor( + codebook, + name="gefen_codebook", + dtype=torch.float32, + shape=(256,), + validate_values=validate_values, + ) + if validate_values: + if float(codebook[0].item()) != -1.0 or float(codebook[-1].item()) != 1.0: + raise ValueError("portable codebook must retain exact endpoints") + if not bool(torch.all(codebook.detach()[1:] >= codebook.detach()[:-1])): + raise ValueError("portable codebook must be sorted") + codebook_preview = _PreflightTensor((256,)) + if type(optimizer._deterministic) is not bool: + raise TypeError("gefen_deterministic must be a bool") + return { + "gefen_global_step": global_step, + "gefen_codebook": codebook_preview, + "gefen_deterministic": optimizer._deterministic, + } + + +def _parameter_state_core(optimizer, parameter, compatibility_name: str): + state = optimizer.state[parameter] + if type(state) is not dict: + raise TypeError("portable per-parameter state must be a plain dictionary") + if any(key not in _AUTHORITATIVE_STATE_KEYS and key not in _IGNORED_DERIVED_STATE_KEYS for key in state): + raise ValueError("portable parameter state contains an unknown key") + core = {key: value for key, value in state.items() if key not in _IGNORED_DERIVED_STATE_KEYS} + if core.get("name") != compatibility_name: + raise ValueError("portable parameter state name does not match its logical slot") + return core + + +def _decode_local_momentum(core, shard: ShardIdentity, codebook): + local_shape = _local_dense_shape(shard) + if local_shape is None: + raise ValueError("a non-payload shard cannot carry initialized momentum") + indices = _plain_tensor(core["m_codebook"], name="m_codebook", dtype=torch.uint8) + magnitudes = _plain_tensor( + core["m_magnitude"], + name="m_magnitude", + dtype=torch.float32, + nonnegative=True, + ) + if indices.device != magnitudes.device: + raise ValueError("momentum indices and magnitudes must share a device") + local_codebook = codebook.detach().to(device=indices.device).contiguous() + dense = _decode_quantized_momentum( + local_codebook, + indices, + magnitudes, + logical_shape=local_shape, + period=core["automatic_period"], + step=core["step"], + ) + return _tight_cpu_fp32(dense, name="dense momentum", shape=local_shape) + + +def _slot_record( + optimizer, + implementation: str, + slot: LogicalSlotBinding, + parameter, + shard: ShardIdentity, + options, + common, +): + role = _derived_role(shard) + record = { + "group_index": slot.group_index, + "original_slot_index": slot.original_slot_index, + "compatibility_name": slot.compatibility_name, + "shard": _serialize_shard_identity(shard), + "algorithm_options": options, + "role": role, + "source_period": None, + "source_second_moment": None, + "state_variant": "pristine", + "state": {}, + } + payload_role = role in {"live", "whole_owner"} and shard.parameter.numel > 0 + if not payload_role: + if ( + parameter is not None + and shard.layout is ParameterLayout.WHOLE_PARAMETER_OWNER + and shard.local_member != shard.owner + ): + raise ValueError("whole-owner nonowners must not retain parameter storage") + if parameter is not None: + core = _parameter_state_core(optimizer, parameter, slot.compatibility_name) + if frozenset(core) != {"name"}: + raise ValueError("empty portable shards must remain pristine") + return record + if parameter is None: + raise ValueError("a portable payload shard requires local parameter storage") + + core = _parameter_state_core(optimizer, parameter, slot.compatibility_name) + keys = frozenset(core) + if keys == {"name"}: + return record + if keys == {"name", "automatic_period"}: + if type(core["automatic_period"]) is not int or core["automatic_period"] != 1: + raise ValueError("exact portable period-selected state requires period one") + record["source_period"] = 1 + record["state_variant"] = "period_selected" + return record + if common["gefen_codebook"] is None: + raise ValueError("initialized portable state requires a codebook") + if type(core.get("automatic_period")) is not int or core["automatic_period"] != 1: + raise ValueError("exact portable initialized state requires period one") + step = _strict_counter(core.get("step"), name="step", minimum=1) + if step > common["gefen_global_step"]: + raise ValueError("parameter step exceeds optimizer global step") + dense_momentum = _decode_local_momentum( + core, + shard, + optimizer._gefen_codebook, + ) + record["source_period"] = 1 + + if implementation == _PLAIN_IMPLEMENTATION: + if options["second_moment_policy"] == "factored": + if keys != _PLAIN_FACTORED_KEYS: + raise ValueError("plain factored state is incomplete or contains inactive fields") + if shard.layout is not ParameterLayout.REPLICATED: + raise ValueError("factored state requires a replicated shard") + factored_step = _strict_counter(core["factored_step"], name="factored_step", minimum=1) + if factored_step > step: + raise ValueError("factored_step exceeds parameter step") + record.update( + { + "source_second_moment": "factored", + "state_variant": "initialized_factored", + "state": { + "step": step, + "momentum": dense_momentum, + "v_row": _tight_cpu_fp32( + core["v_row"], + name="v_row", + shape=(shard.parameter.global_shape[0],), + nonnegative=True, + ), + "v_col": _tight_cpu_fp32( + core["v_col"], + name="v_col", + shape=(shard.parameter.global_shape[1],), + nonnegative=True, + ), + "factored_step": factored_step, + }, + } + ) + return record + + if keys != _PLAIN_BLOCK_KEYS: + raise ValueError("plain block state is incomplete or contains inactive fields") + vmean_step = _strict_counter(core["vmean_step"], name="vmean_step", minimum=1) + if vmean_step > step: + raise ValueError("vmean_step exceeds parameter step") + local_shape = _local_dense_shape(shard) + vmean = _plain_tensor( + core["vmean"], + name="vmean", + dtype=torch.float32, + nonnegative=True, + ) + second = _expand_block_second_moment( + vmean, + logical_shape=local_shape, + period=1, + step=vmean_step, + ) + record.update( + { + "source_second_moment": "block", + "state_variant": "initialized_dense", + "state": { + "step": step, + "momentum": dense_momentum, + "second_moment": _tight_cpu_fp32( + second, + name="dense second moment", + shape=local_shape, + nonnegative=True, + ), + "second_moment_step": vmean_step, + }, + } + ) + return record + + expected = _NORMUON_KEYS if options["normuon"] else _MUON_KEYS + if keys != expected: + raise ValueError("Muon state is incomplete or conflicts with its NorMuon policy") + state = {"step": step, "momentum": dense_momentum} + variant = "initialized_dense" + if options["normuon"]: + normuon_step = _strict_counter(core["normuon_step"], name="normuon_step", minimum=1) + if normuon_step > step: + raise ValueError("normuon_step exceeds parameter step") + state.update( + { + "normuon_v": _tight_cpu_fp32( + core["normuon_v"], + name="normuon_v", + shape=(shard.parameter.global_shape[0], 1), + nonnegative=True, + ), + "normuon_step": normuon_step, + } + ) + variant = "initialized_dense_normuon" + record["state_variant"] = variant + record["state"] = state + return record + + +def _prepare_local_structure( + optimizer, + implementation: str, + binding: CheckpointProcessGroupBinding, + limits: PortableStateLimits, + *, + include_payload: bool, +): + for name, value in ( + ("logical slots", optimizer._gefen_logical_slots), + ("manifest shards", optimizer._gefen_sharding_manifest.shards), + ("parameter groups", optimizer.param_groups), + ): + if type(value) not in {list, tuple}: + continue + if len(value) > limits.max_container_items or len(value) > limits.max_tree_nodes: + raise ValueError("portable {} exceed structural limits".format(name)) + for group in optimizer.param_groups: + if type(group) is not dict: + continue + for key in ("params", "param_names"): + value = group.get(key) + if type(value) is list and (len(value) > limits.max_container_items or len(value) > limits.max_tree_nodes): + raise ValueError("portable group slots exceed structural limits") + _validate_optimizer_shell(optimizer, implementation, binding) + local_by_fqn = _live_maps(optimizer, binding) + policy = _live_policy(optimizer, implementation) + catalog, options_by_fqn = _catalog_and_options(optimizer, implementation, policy) + common_preview = _preflight_common_state( + optimizer, + validate_values=False, + ) + preview_slots = [] + for logical_slot in optimizer._gefen_logical_slots: + fqn = logical_slot.shard.parameter.fqn + parameter, shard = local_by_fqn[fqn] + if shard != logical_slot.shard: + raise ValueError("logical-slot and local shard identities disagree") + preview_slots.append( + _preflight_slot_record( + optimizer, + implementation, + logical_slot, + parameter, + shard, + options_by_fqn[fqn], + common_preview, + ) + ) + manifest = _serialize_sharding_manifest(optimizer._gefen_sharding_manifest) + target_slots = [ + { + "fqn": slot["shard"]["parameter"]["fqn"], + "group_index": slot["group_index"], + "original_slot_index": slot["original_slot_index"], + "compatibility_name": slot["compatibility_name"], + } + for slot in preview_slots + ] + if include_payload: + bounded_value = { + "format": "gefen.portable_state_fragment", + "format_version": 1, + "coverage": "local_logical_optimizer_fragment", + "implementation": implementation, + "member": binding.local_member, + "policy": policy, + "common": common_preview, + "manifest": manifest, + "catalog": catalog, + "logical_slots": preview_slots, + } + else: + bounded_value = { + "policy": policy, + "manifest": manifest, + "catalog": catalog, + "logical_slots": target_slots, + } + _preflight_portable_value(bounded_value, limits) + return { + "local_by_fqn": local_by_fqn, + "policy": policy, + "catalog": catalog, + "options_by_fqn": options_by_fqn, + "target_descriptor": { + "policy": policy, + "manifest": manifest, + "catalog": catalog, + "logical_slots": target_slots, + }, + } + + +def _validate_prepared_local_state(optimizer, implementation: str, prepared) -> None: + common = _preflight_common_state(optimizer) + local_by_fqn = prepared["local_by_fqn"] + options_by_fqn = prepared["options_by_fqn"] + for slot in optimizer._gefen_logical_slots: + fqn = slot.shard.parameter.fqn + parameter, shard = local_by_fqn[fqn] + if shard != slot.shard: + raise ValueError("logical-slot and local shard identities disagree") + _validate_readiness_slot( + optimizer, + implementation, + slot, + parameter, + shard, + options_by_fqn[fqn], + common, + ) + + +def _materialize_local_fragment( + optimizer, + implementation: str, + binding: CheckpointProcessGroupBinding, + limits: PortableStateLimits, + prepared, +): + local_by_fqn = prepared["local_by_fqn"] + policy = prepared["policy"] + catalog = prepared["catalog"] + options_by_fqn = prepared["options_by_fqn"] + common = _common_state(optimizer) + slots = [] + for logical_slot in optimizer._gefen_logical_slots: + fqn = logical_slot.shard.parameter.fqn + parameter, shard = local_by_fqn[fqn] + if shard != logical_slot.shard: + raise ValueError("logical-slot and local shard identities disagree") + slots.append( + _slot_record( + optimizer, + implementation, + logical_slot, + parameter, + shard, + options_by_fqn[fqn], + common, + ) + ) + return _build_portable_state_fragment( + implementation=implementation, + member=binding.local_member, + policy=policy, + common=common, + manifest=optimizer._gefen_sharding_manifest, + catalog=catalog, + logical_slots=slots, + limits=limits, + ) + + +def _target_context(base, fragment): + return { + **base, + "target_policy": fragment["policy"], + "target_manifest": fragment["manifest"], + "target_catalog": fragment["catalog"], + "target_logical_slots": fragment["logical_slots"], + } + + +def _export_portable_state( + optimizer, + *, + checkpoint_process_group, + transaction_id, + limits, +): + """Collectively export one complete exact portable v3 optimizer document.""" + + transport_binding = _preflight_transport_binding( + optimizer, + checkpoint_process_group, + ) + binding = None + normalized_limits = None + normalized_transaction = None + implementation = None + context = bytes(32) + wire_limits = _STATUS_FALLBACK_LIMITS + local_fragment = None + live_token = None + error = None + try: + _validate_supplied_binding(checkpoint_process_group, transport_binding) + binding = transport_binding + normalized_limits = _require_limits(limits) + wire_limits = normalized_limits._wire_limits() + normalized_transaction = _require_transaction_id(transaction_id) + implementation = _optimizer_implementation(optimizer) + _validate_context_identity_bounds(binding, normalized_limits) + base_context = _base_context(binding, implementation) + _preflight_portable_value(base_context, normalized_limits) + context = _context_digest(base_context) + _prepare_local_structure( + optimizer, + implementation, + binding, + normalized_limits, + include_payload=True, + ) + live_token = _portable_live_token(optimizer) + prepared_local = _prepare_local_structure( + optimizer, + implementation, + binding, + normalized_limits, + include_payload=True, + ) + _validate_prepared_local_state( + optimizer, + implementation, + prepared_local, + ) + local_fragment = _materialize_local_fragment( + optimizer, + implementation, + binding, + normalized_limits, + prepared_local, + ) + if live_token != _portable_live_token(optimizer): + raise RuntimeError("live optimizer state changed during portable export preparation") + except Exception as exc: + error = exc + _collective_unanimous_status( + transport_binding, + error, + operation="portable_export_prepare", + transaction_id=_EXPORT_PREFLIGHT_TRANSACTION, + context_digest=context, + limits=wire_limits, + ) + assert ( + normalized_limits is not None + and normalized_transaction is not None + and implementation is not None + and binding is not None + and local_fragment is not None + ) + + fragments = [] + _collective_visit_canonical_fragments( + binding, + local_fragment, + operation="portable_export_fragments", + transaction_id=normalized_transaction, + context_digest=context, + limits=wire_limits, + consume=lambda _member, value: fragments.append( + _normalize_portable_state_fragment( + value, + limits=normalized_limits, + ) + ), + ) + + document = None + error = None + try: + document = _assemble_portable_state_fragments( + fragments, + process_group_identity=binding.identity, + limits=normalized_limits, + ) + _validate_live_readiness(optimizer, implementation, binding) + if live_token != _portable_live_token(optimizer): + raise RuntimeError("live optimizer state changed during portable export") + except Exception as exc: + error = exc + _collective_unanimous_status( + binding, + error, + operation="portable_export_finalize", + transaction_id=normalized_transaction, + context_digest=context, + limits=wire_limits, + ) + assert document is not None + return document + + +def _stage_portable_import( + optimizer, + implementation: str, + binding: CheckpointProcessGroupBinding, + limits: PortableStateLimits, + document, +): + _prepare_local_structure( + optimizer, + implementation, + binding, + limits, + include_payload=False, + ) + live_token = _portable_live_token(optimizer) + prepared_local = _prepare_local_structure( + optimizer, + implementation, + binding, + limits, + include_payload=False, + ) + _validate_prepared_local_state( + optimizer, + implementation, + prepared_local, + ) + target_fragment = prepared_local["target_descriptor"] + if not _values_equal(document["policy"], target_fragment["policy"]): + raise ValueError("portable document policy does not match the target") + common = document["common"] + if set(document["parameters"]) != set(target_fragment["catalog"]): + raise ValueError("portable parameter catalog does not match the target") + + local_by_fqn = _live_maps(optimizer, binding) + canonical_parameters = {} + for fqn, (parameter, shard) in local_by_fqn.items(): + if parameter is None: + continue + target_options = target_fragment["catalog"][fqn]["algorithm_options"] + target_second = target_options["second_moment_policy"] if implementation == _PLAIN_IMPLEMENTATION else None + projected = _project_portable_parameter_state( + document["parameters"][fqn], + shard, + implementation=implementation, + global_step=common["gefen_global_step"], + codebook=common["gefen_codebook"], + target_algorithm_options=target_options, + target_second_moment=target_second, + ) + canonical_parameters[fqn] = {"state": projected} + + canonical_state = { + "common": { + "gefen_global_step": common["gefen_global_step"], + "gefen_codebook": common["gefen_codebook"], + "gefen_deterministic": common["gefen_deterministic"], + "gefen_codebook_scope": optimizer._serialized_codebook_scope(), + }, + "parameters": canonical_parameters, + } + native = optimizer._canonical_native_state_dict(canonical_state) + # Portable state intentionally carries the source replica-determinism + # setting. Native loading ordinarily rejects a policy change, so perform + # the same complete staging path through an isolated owner configured with + # the document value. The live optimizer remains untouched until commit. + staging_owner = object.__new__(type(optimizer)) + staging_owner.__dict__ = optimizer.__dict__.copy() + staging_owner._deterministic = common["gefen_deterministic"] + staged = staging_owner._stage_load_state_dict(native) + optimizer._preserve_canonical_target_configuration(staged) + if ( + type(staged.__dict__) is not dict + or type(staged.defaults) is not dict + or set(staged.defaults) != set(optimizer.defaults) + or staged.param_groups is not optimizer.param_groups + ): + raise TypeError("portable import staging produced unsafe publication containers") + if live_token != _portable_live_token(optimizer): + raise RuntimeError("live optimizer state changed during portable import preparation") + return staged, live_token, target_fragment + + +def _import_portable_state( + optimizer, + state, + *, + checkpoint_process_group, + transaction_id, + limits, +) -> None: + """Collectively stage, vote, and atomically publish portable v3 state.""" + + transport_binding = _preflight_transport_binding( + optimizer, + checkpoint_process_group, + ) + binding = None + normalized_limits = None + normalized_transaction = None + implementation = None + base = None + base_digest = bytes(32) + wire_limits = _STATUS_FALLBACK_LIMITS + document = None + error = None + try: + _validate_supplied_binding(checkpoint_process_group, transport_binding) + binding = transport_binding + normalized_limits = _require_limits(limits) + wire_limits = normalized_limits._wire_limits() + normalized_transaction = _require_transaction_id(transaction_id) + implementation = _optimizer_implementation(optimizer) + _validate_context_identity_bounds(binding, normalized_limits) + base = _base_context(binding, implementation) + _preflight_portable_value(base, normalized_limits) + base_digest = _context_digest(base) + document = _normalize_gefen_portable_state_document( + state, + limits=normalized_limits, + expected_implementation=implementation, + ) + except Exception as exc: + error = exc + _collective_unanimous_status( + transport_binding, + error, + operation="portable_import_document", + transaction_id=_IMPORT_PREFLIGHT_TRANSACTION, + context_digest=base_digest, + limits=wire_limits, + ) + assert ( + normalized_limits is not None + and normalized_transaction is not None + and implementation is not None + and binding is not None + and base is not None + and document is not None + ) + + staged = None + live_token = None + target_fragment = None + error = None + try: + staged, live_token, target_fragment = _stage_portable_import( + optimizer, + implementation, + binding, + normalized_limits, + document, + ) + target_context = _target_context(base, target_fragment) + target_context["document_digest"] = document["completion"]["digest"] + _preflight_portable_value(target_context, normalized_limits) + import_digest = _context_digest(target_context) + except Exception as exc: + error = exc + import_digest = _context_digest({**base, "document_digest": document["completion"]["digest"]}) + _collective_unanimous_status( + binding, + error, + operation="portable_import_prepare", + transaction_id=normalized_transaction, + context_digest=import_digest, + limits=wire_limits, + ) + assert staged is not None and live_token is not None and target_fragment is not None + + error = None + try: + _validate_live_readiness(optimizer, implementation, binding) + if live_token != _portable_live_token(optimizer): + raise RuntimeError("live optimizer state changed after portable import preparation") + except Exception as exc: + error = exc + _collective_unanimous_status( + binding, + error, + operation="portable_import_freshness", + transaction_id=normalized_transaction, + context_digest=import_digest, + limits=wire_limits, + ) + + # This is the established non-throwing publication primitive. No collective + # or semantic validation is allowed after this point. + from gefen.gefen import Gefen + + Gefen._commit_staged_load_state_dict(optimizer, staged) + + +def _validate_readiness_slot( + optimizer, + implementation: str, + slot: LogicalSlotBinding, + parameter, + shard: ShardIdentity, + options, + common, + *, + validate_values=True, +) -> None: + role = _derived_role(shard) + payload_role = role in {"live", "whole_owner"} and shard.parameter.numel > 0 + if not payload_role: + if parameter is None: + if shard.layout is not ParameterLayout.WHOLE_PARAMETER_OWNER: + raise ValueError("only whole-parameter nonowners may omit storage") + return + core = _parameter_state_core(optimizer, parameter, slot.compatibility_name) + if frozenset(core) != {"name"}: + raise ValueError("empty portable shards must remain pristine") + return + if parameter is None: + raise ValueError("a portable payload shard requires local storage") + + core = _parameter_state_core(optimizer, parameter, slot.compatibility_name) + keys = frozenset(core) + if keys == {"name"}: + return + if keys == {"name", "automatic_period"}: + if type(core["automatic_period"]) is not int or core["automatic_period"] != 1: + raise ValueError("exact portable period-selected state requires period one") + if common["gefen_codebook"] is None: + raise ValueError("period-selected portable state requires a codebook") + return + if common["gefen_codebook"] is None: + raise ValueError("initialized portable state requires a codebook") + if type(core.get("automatic_period")) is not int or core["automatic_period"] != 1: + raise ValueError("exact portable initialized state requires period one") + step = _strict_counter(core.get("step"), name="step", minimum=1) + if step > common["gefen_global_step"]: + raise ValueError("parameter step exceeds optimizer global step") + local_shape = _local_dense_shape(shard) + if local_shape is None: + raise ValueError("initialized portable state requires a local payload") + blocks = math.prod(local_shape) + indices = _plain_tensor( + core.get("m_codebook"), + name="m_codebook", + dtype=torch.uint8, + shape=(blocks, 1), + validate_values=validate_values, + ) + magnitudes = _plain_tensor( + core.get("m_magnitude"), + name="m_magnitude", + dtype=torch.float32, + shape=(blocks, 1), + nonnegative=True, + validate_values=validate_values, + ) + if indices.device != magnitudes.device: + raise ValueError("momentum indices and magnitudes must share a device") + if indices.device != parameter.device: + raise ValueError("portable parameter state must share its parameter device") + + if implementation == _PLAIN_IMPLEMENTATION: + if options["second_moment_policy"] == "factored": + if keys != _PLAIN_FACTORED_KEYS: + raise ValueError("plain factored state is incomplete or contains inactive fields") + factored_step = _strict_counter(core["factored_step"], name="factored_step", minimum=1) + if factored_step > step: + raise ValueError("factored_step exceeds parameter step") + rows, columns = shard.parameter.global_shape + v_row = _plain_tensor( + core["v_row"], + name="v_row", + dtype=torch.float32, + shape=(rows,), + nonnegative=True, + validate_values=validate_values, + ) + v_col = _plain_tensor( + core["v_col"], + name="v_col", + dtype=torch.float32, + shape=(columns,), + nonnegative=True, + validate_values=validate_values, + ) + if v_row.device != parameter.device or v_col.device != parameter.device: + raise ValueError("factored portable state must share its parameter device") + return + if keys != _PLAIN_BLOCK_KEYS: + raise ValueError("plain block state is incomplete or contains inactive fields") + vmean_step = _strict_counter(core["vmean_step"], name="vmean_step", minimum=1) + if vmean_step > step: + raise ValueError("vmean_step exceeds parameter step") + vmean = _plain_tensor( + core["vmean"], + name="vmean", + dtype=torch.float32, + shape=(blocks, 1), + nonnegative=True, + validate_values=validate_values, + ) + if vmean.device != parameter.device: + raise ValueError("block portable state must share its parameter device") + return + + expected = _NORMUON_KEYS if options["normuon"] else _MUON_KEYS + if keys != expected: + raise ValueError("Muon state is incomplete or conflicts with its NorMuon policy") + if options["normuon"]: + normuon_step = _strict_counter(core["normuon_step"], name="normuon_step", minimum=1) + if normuon_step > step: + raise ValueError("normuon_step exceeds parameter step") + normuon_v = _plain_tensor( + core["normuon_v"], + name="normuon_v", + dtype=torch.float32, + shape=(shard.parameter.global_shape[0], 1), + nonnegative=True, + validate_values=validate_values, + ) + if normuon_v.device != parameter.device: + raise ValueError("NorMuon portable state must share its parameter device") + + +def _preflight_slot_record( + optimizer, + implementation: str, + slot: LogicalSlotBinding, + parameter, + shard: ShardIdentity, + options, + common, +): + _validate_readiness_slot( + optimizer, + implementation, + slot, + parameter, + shard, + options, + common, + validate_values=False, + ) + role = _derived_role(shard) + record = { + "group_index": slot.group_index, + "original_slot_index": slot.original_slot_index, + "compatibility_name": slot.compatibility_name, + "shard": _serialize_shard_identity(shard), + "algorithm_options": options, + "role": role, + "source_period": None, + "source_second_moment": None, + "state_variant": "pristine", + "state": {}, + } + payload_role = role in {"live", "whole_owner"} and shard.parameter.numel > 0 + if not payload_role: + return record + core = _parameter_state_core(optimizer, parameter, slot.compatibility_name) + keys = frozenset(core) + if keys == {"name"}: + return record + record["source_period"] = 1 + if keys == {"name", "automatic_period"}: + record["state_variant"] = "period_selected" + return record + + local_shape = _local_dense_shape(shard) + state = { + "step": core["step"], + "momentum": _PreflightTensor(local_shape), + } + if implementation == _PLAIN_IMPLEMENTATION: + if options["second_moment_policy"] == "factored": + rows, columns = shard.parameter.global_shape + record.update( + { + "source_second_moment": "factored", + "state_variant": "initialized_factored", + "state": { + **state, + "v_row": _PreflightTensor((rows,)), + "v_col": _PreflightTensor((columns,)), + "factored_step": core["factored_step"], + }, + } + ) + return record + record.update( + { + "source_second_moment": "block", + "state_variant": "initialized_dense", + "state": { + **state, + "second_moment": _PreflightTensor(local_shape), + "second_moment_step": core["vmean_step"], + }, + } + ) + return record + + if options["normuon"]: + state.update( + { + "normuon_v": _PreflightTensor((shard.parameter.global_shape[0], 1)), + "normuon_step": core["normuon_step"], + } + ) + record["state_variant"] = "initialized_dense_normuon" + else: + record["state_variant"] = "initialized_dense" + record["state"] = state + return record + + +def _validate_live_readiness( + optimizer, + implementation: str, + binding: CheckpointProcessGroupBinding, +): + _validate_optimizer_shell(optimizer, implementation, binding) + local_by_fqn = _live_maps(optimizer, binding) + policy = _live_policy(optimizer, implementation) + _, options_by_fqn = _catalog_and_options(optimizer, implementation, policy) + common = _preflight_common_state(optimizer) + for slot in optimizer._gefen_logical_slots: + parameter, shard = local_by_fqn[slot.shard.parameter.fqn] + if shard != slot.shard: + raise ValueError("logical-slot and local shard identities disagree") + _validate_readiness_slot( + optimizer, + implementation, + slot, + parameter, + shard, + options_by_fqn[slot.shard.parameter.fqn], + common, + ) + return frozenset(slot.shard.layout for slot in optimizer._gefen_logical_slots) + + +def _portable_runtime_layouts(optimizer): + """Return dynamically eligible layouts without executing a collective.""" + + try: + implementation = _optimizer_implementation(optimizer) + scope = optimizer._gefen_codebook_process_group + if type(scope) is not CodebookProcessGroupBinding: + return frozenset() + binding = CheckpointProcessGroupBinding( + scope.identity, + scope.local_member, + scope.process_group, + scope.collective_device, + ) + return _validate_live_readiness(optimizer, implementation, binding) + except Exception: + return frozenset() + + +__all__ = [] diff --git a/src/gefen/portable_schema.py b/src/gefen/portable_schema.py new file mode 100644 index 0000000..1ffa415 --- /dev/null +++ b/src/gefen/portable_schema.py @@ -0,0 +1,403 @@ +"""Strict source-topology-neutral envelope for portable optimizer state.""" + +import hashlib +import hmac +import math +import struct +import sys + +import torch + +from gefen.contracts import ParameterIdentity + + +PORTABLE_STATE_FORMAT = "gefen.portable_state" +PORTABLE_STATE_FORMAT_VERSION = 3 +PORTABLE_STATE_COVERAGE = "global_logical_optimizer" +PORTABLE_STATE_DIGEST_ALGORITHM = "sha256" +_PORTABLE_DIGEST_CHUNK_BYTES = 1 << 23 +_PORTABLE_CLONE_CHUNK_BYTES = 1 << 23 +_PORTABLE_NATIVE_BYTEORDER = sys.byteorder + +_PORTABLE_STATE_TOP_LEVEL_KEYS = frozenset( + { + "format", + "format_version", + "coverage", + "implementation", + "policy", + "common", + "parameters", + "provenance", + "completion", + } +) +_PORTABLE_PARAMETER_KEYS = frozenset( + { + "identity", + "algorithm_options", + "state_variant", + "state", + "projection_hints", + } +) +_PORTABLE_IDENTITY_KEYS = frozenset( + {"schema_version", "fqn", "global_shape"} +) +_PORTABLE_COMPLETION_KEYS = frozenset( + {"status", "digest_algorithm", "digest"} +) + + +def _portable_tensor_chunk_elements(value: torch.Tensor) -> int: + budget = _PORTABLE_CLONE_CHUNK_BYTES + if type(budget) is not int or budget <= 0: + raise RuntimeError("portable clone chunk budget must be a positive int") + scratch_bytes_per_element = value.element_size() + if not value.is_contiguous(): + # The non-contiguous read path (_read_portable_tensor_chunk) materializes + # one int64 coordinate vector per dimension plus the int64 linear index + # and its running quotient. Size the chunk for that scratch too, so a + # high-rank strided tensor stays within the fixed clone budget instead of + # oversubscribing it. Mirrors portable_wire._clone_tensor. + scratch_bytes_per_element += 8 * (value.ndim + 2) + return max(1, budget // scratch_bytes_per_element) + + +def _read_portable_tensor_chunk( + value: torch.Tensor, start: int, stop: int +) -> torch.Tensor: + detached = value.detach() + if detached.is_contiguous(): + return detached.reshape(-1)[start:stop] + linear = torch.arange(start, stop, dtype=torch.int64, device=detached.device) + remainder = linear + reversed_coordinates = [] + for dimension in reversed(detached.shape): + reversed_coordinates.append(torch.remainder(remainder, dimension)) + remainder = torch.div(remainder, dimension, rounding_mode="floor") + return detached[tuple(reversed(reversed_coordinates))] + + +def _clone_portable_value(value, *, path: str): + if torch.is_tensor(value): + if ( + type(value) is not torch.Tensor + or value.layout is not torch.strided + or value.is_meta + or value.is_nested + or value.is_quantized + ): + raise TypeError( + "{} must be a plain materialized strided tensor for portable state".format( + path + ) + ) + cloned = torch.empty(tuple(value.shape), dtype=value.dtype, device="cpu") + cloned_flat = cloned.reshape(-1) + chunk_elements = _portable_tensor_chunk_elements(value) + for start in range(0, value.numel(), chunk_elements): + stop = min(start + chunk_elements, value.numel()) + source = ( + _read_portable_tensor_chunk(value, start, stop) + .resolve_conj() + .resolve_neg() + .reshape(-1) + ) + destination = cloned_flat[start:stop] + destination.copy_(source) + if cloned.is_floating_point() or cloned.is_complex(): + try: + finite = bool(torch.isfinite(destination).all()) + except (NotImplementedError, RuntimeError, TypeError) as exc: + raise ValueError( + "{} tensor dtype does not support finite portable state".format( + path + ) + ) from exc + if not finite: + raise ValueError("{} tensor must be finite".format(path)) + return cloned + if type(value) is float: + if not math.isfinite(value): + raise ValueError("{} must be finite".format(path)) + return value + if value is None or type(value) in {bool, int, str}: + return value + if type(value) is list: + return [ + _clone_portable_value(item, path="{}[{}]".format(path, index)) + for index, item in enumerate(value) + ] + if type(value) is tuple: + return tuple( + _clone_portable_value(item, path="{}[{}]".format(path, index)) + for index, item in enumerate(value) + ) + if type(value) is dict: + if any(type(key) is not str for key in value): + raise TypeError("{} dictionary keys must be strings".format(path)) + return { + key: _clone_portable_value( + value[key], path="{}.{}".format(path, key) + ) + for key in sorted(value) + } + raise TypeError( + "{} has unsupported portable-state type {}".format( + path, type(value).__name__ + ) + ) + + +def _digest_bytes(hasher, value: bytes) -> None: + hasher.update(struct.pack(">Q", len(value))) + hasher.update(value) + + +def _update_portable_digest(hasher, value) -> None: + value_type = type(value) + if value is None: + hasher.update(b"N") + return + if value_type is bool: + hasher.update(b"B1" if value else b"B0") + return + if value_type is int: + hasher.update(b"I") + _digest_bytes(hasher, str(value).encode("ascii")) + return + if value_type is float: + hasher.update(b"F") + hasher.update(struct.pack(">d", value)) + return + if value_type is str: + hasher.update(b"S") + _digest_bytes(hasher, value.encode("utf-8")) + return + if value_type is list: + hasher.update(b"L") + hasher.update(struct.pack(">Q", len(value))) + for item in value: + _update_portable_digest(hasher, item) + return + if value_type is tuple: + hasher.update(b"T") + hasher.update(struct.pack(">Q", len(value))) + for item in value: + _update_portable_digest(hasher, item) + return + if value_type is dict: + hasher.update(b"D") + hasher.update(struct.pack(">Q", len(value))) + for key in sorted(value): + if type(key) is not str: + raise TypeError("portable state dictionary keys must be strings") + _update_portable_digest(hasher, key) + _update_portable_digest(hasher, value[key]) + return + if value_type is torch.Tensor: + if value.device.type != "cpu" or not value.is_contiguous(): + raise ValueError("portable digest tensors must be contiguous CPU tensors") + hasher.update(b"R") + _digest_bytes(hasher, str(value.dtype).encode("ascii")) + _update_portable_digest(hasher, list(value.shape)) + element_size = value.element_size() + hasher.update(struct.pack(">Q", value.numel() * element_size)) + elements_per_chunk = max( + 1, _PORTABLE_DIGEST_CHUNK_BYTES // element_size + ) + flat = value.reshape(-1) + for start in range(0, value.numel(), elements_per_chunk): + stop = min(start + elements_per_chunk, value.numel()) + raw = flat[start:stop].view(torch.uint8).reshape(-1) + component_size = element_size // 2 if value.is_complex() else element_size + if _PORTABLE_NATIVE_BYTEORDER not in {"little", "big"}: + raise RuntimeError("unsupported native byte order") + if _PORTABLE_NATIVE_BYTEORDER == "big" and component_size > 1: + raw = ( + raw.reshape(-1, component_size) + .flip(1) + .contiguous() + .reshape(-1) + ) + hasher.update(memoryview(raw.numpy())) + return + raise TypeError( + "portable state digest does not support {}".format(value_type.__name__) + ) + + +def _canonical_portable_state_digest(canonical) -> str: + hasher = hashlib.sha256() + _update_portable_digest(hasher, canonical) + return hasher.hexdigest() + + +def portable_state_digest(payload) -> str: + """Return the deterministic SHA-256 digest of one canonical wire value.""" + + canonical = _clone_portable_value(payload, path="portable digest payload") + return _canonical_portable_state_digest(canonical) + + +def _normalize_portable_parameter(fqn, record): + if type(record) is not dict or set(record) != _PORTABLE_PARAMETER_KEYS: + raise ValueError( + "portable parameter {!r} has an invalid schema".format(fqn) + ) + identity_record = record["identity"] + if ( + type(identity_record) is not dict + or set(identity_record) != _PORTABLE_IDENTITY_KEYS + or type(identity_record["global_shape"]) is not list + ): + raise ValueError( + "portable parameter {!r} has an invalid identity".format(fqn) + ) + try: + identity = ParameterIdentity( + identity_record["fqn"], + tuple(identity_record["global_shape"]), + schema_version=identity_record["schema_version"], + ) + except (TypeError, ValueError) as exc: + raise ValueError( + "portable parameter {!r} has an invalid identity".format(fqn) + ) from exc + if identity.fqn != fqn: + raise ValueError("portable parameter identity does not match its FQN key") + for key in ("algorithm_options", "state", "projection_hints"): + if type(record[key]) is not dict: + raise ValueError( + "portable parameter {!r} {} must be a dictionary".format( + fqn, key + ) + ) + if type(record["state_variant"]) is not str or not record["state_variant"]: + raise ValueError( + "portable parameter {!r} state_variant must be a non-empty string".format( + fqn + ) + ) + return record + + +def _portable_payload(document): + return { + key: document[key] + for key in sorted(_PORTABLE_STATE_TOP_LEVEL_KEYS - {"completion"}) + } + + +def _validate_portable_payload(document) -> None: + if document["format"] != PORTABLE_STATE_FORMAT: + raise ValueError("unsupported portable state format") + if ( + type(document["format_version"]) is not int + or document["format_version"] != PORTABLE_STATE_FORMAT_VERSION + ): + raise ValueError( + "unsupported portable state format_version: {}".format( + document["format_version"] + ) + ) + if document["coverage"] != PORTABLE_STATE_COVERAGE: + raise ValueError("unsupported portable state coverage") + implementation = document["implementation"] + if type(implementation) is not str or not implementation: + raise ValueError("portable state implementation must be a non-empty string") + if type(document["policy"]) is not dict: + raise ValueError("portable state policy must be a dictionary") + if type(document["common"]) is not dict: + raise ValueError("portable common state must be a dictionary") + if document["provenance"] is not None and type(document["provenance"]) is not dict: + raise ValueError("portable state provenance must be a dictionary or None") + parameters = document["parameters"] + if type(parameters) is not dict: + raise ValueError("portable parameters must be an FQN mapping") + for fqn, record in parameters.items(): + if type(fqn) is not str: + raise ValueError("portable parameter FQN keys must be strings") + _normalize_portable_parameter(fqn, record) + + +def _validate_portable_completion(document) -> None: + completion = document["completion"] + if type(completion) is not dict or set(completion) != _PORTABLE_COMPLETION_KEYS: + raise ValueError("portable state completion marker has an invalid schema") + if completion["status"] != "complete": + raise ValueError("portable state is not marked complete") + if completion["digest_algorithm"] != PORTABLE_STATE_DIGEST_ALGORITHM: + raise ValueError("unsupported portable state digest algorithm") + digest = completion["digest"] + if ( + type(digest) is not str + or len(digest) != hashlib.sha256().digest_size * 2 + or any(character not in "0123456789abcdef" for character in digest) + ): + raise ValueError("portable state completion digest is invalid") + expected_digest = _canonical_portable_state_digest(_portable_payload(document)) + if not hmac.compare_digest(digest, expected_digest): + raise ValueError("portable state completion digest does not match its payload") + + +def normalize_portable_state_document(state, *, expected_implementation=None): + """Clone and validate a complete portable v3 state document.""" + + document = _clone_portable_value(state, path="portable state") + if type(document) is not dict or set(document) != _PORTABLE_STATE_TOP_LEVEL_KEYS: + raise ValueError("portable state has an invalid top-level schema") + _validate_portable_payload(document) + _validate_portable_completion(document) + if ( + expected_implementation is not None + and document["implementation"] != expected_implementation + ): + raise ValueError("portable state implementation does not match the target") + return document + + +def build_portable_state_document( + *, + implementation, + policy, + common, + parameters, + provenance=None, +): + """Build and validate a complete source-topology-neutral v3 document.""" + + document = { + "format": PORTABLE_STATE_FORMAT, + "format_version": PORTABLE_STATE_FORMAT_VERSION, + "coverage": PORTABLE_STATE_COVERAGE, + "implementation": implementation, + "policy": policy, + "common": common, + "parameters": parameters, + "provenance": provenance, + } + canonical = _clone_portable_value(document, path="portable state") + if type(canonical) is not dict or set(canonical) != ( + _PORTABLE_STATE_TOP_LEVEL_KEYS - {"completion"} + ): + raise ValueError("portable state has an invalid top-level schema") + _validate_portable_payload(canonical) + canonical["completion"] = { + "status": "complete", + "digest_algorithm": PORTABLE_STATE_DIGEST_ALGORITHM, + "digest": _canonical_portable_state_digest(canonical), + } + return canonical + + +__all__ = [ + "PORTABLE_STATE_COVERAGE", + "PORTABLE_STATE_DIGEST_ALGORITHM", + "PORTABLE_STATE_FORMAT", + "PORTABLE_STATE_FORMAT_VERSION", + "build_portable_state_document", + "normalize_portable_state_document", + "portable_state_digest", +] diff --git a/src/gefen/portable_state.py b/src/gefen/portable_state.py new file mode 100644 index 0000000..e28f722 --- /dev/null +++ b/src/gefen/portable_state.py @@ -0,0 +1,1171 @@ +"""Exact semantic normalization, assembly, and projection for portable Gefen state.""" + +from __future__ import annotations + +from dataclasses import dataclass +import math +import struct + +import torch + +from gefen.contracts import ( + ParameterLayout, + ProcessGroupIdentity, + ShardIdentity, + ShardingManifest, +) +from gefen.portable import _recompress_dense_momentum, _reduce_block_second_moment, _validate_codebook +from gefen.portable_fields import _assemble_dense_logical_field, _project_dense_logical_field, _tensor_bits_equal +from gefen.portable_identity import ( + _normalize_parameter_identity, + _normalize_shard_identity, + _normalize_sharding_manifest, + _parse_parameter_identity, + _parse_shard_identity, + _parse_sharding_manifest, + _serialize_parameter_identity, + _serialize_sharding_manifest, +) +from gefen.portable_schema import build_portable_state_document, normalize_portable_state_document +from gefen.portable_wire import ( + _CanonicalWireLimits, + _parse_canonical_wire_metadata, + _prepare_canonical_wire_value, + _reconstruct_canonical_wire_value, +) + + +_IMPLEMENTATIONS = frozenset({"gefen.Gefen", "gefen.GefenMuon"}) +_FRAGMENT_FORMAT = "gefen.portable_state_fragment" +_FRAGMENT_FORMAT_VERSION = 1 +_FRAGMENT_COVERAGE = "local_logical_optimizer_fragment" +_MOMENTUM_PROJECTION = "dense_fp32_target_period_one_v1" +_SECOND_MOMENT_PROJECTION = "exact_representation_target_period_one_v1" +_MAX_EXACT_COUNTER = (1 << 53) - 1 + +_FRAGMENT_KEYS = frozenset( + { + "format", + "format_version", + "coverage", + "implementation", + "member", + "policy", + "common", + "manifest", + "catalog", + "logical_slots", + } +) +_CATALOG_KEYS = frozenset({"identity", "algorithm_options"}) +_SLOT_KEYS = frozenset( + { + "group_index", + "original_slot_index", + "compatibility_name", + "shard", + "algorithm_options", + "role", + "source_period", + "source_second_moment", + "state_variant", + "state", + } +) +_COMMON_KEYS = frozenset({"gefen_global_step", "gefen_codebook", "gefen_deterministic"}) +_POLICY_KEYS = frozenset( + { + "schema_version", + "factored_v_2d", + "force_1d_period_one", + "force_2d_period_one", + "period_one_substrings", + "codebook_refresh_every", + "stochastic_round", + "momentum_projection", + "second_moment_projection", + } +) +_PLAIN_OPTION_KEYS = frozenset({"lr", "beta1", "beta2", "eps", "weight_decay", "second_moment_policy"}) +_MUON_OPTION_KEYS = frozenset( + { + "lr", + "weight_decay", + "momentum", + "nesterov", + "ns_schedule", + "ns_eps", + "adjust_lr_fn", + "sharded_mode", + "fp8_ns", + "fp8_ns_compile", + "batched_ns", + "batched_ns_workspace_bytes", + "normuon", + "normuon_beta2", + "normuon_eps", + "cautious", + } +) +_PLAIN_HINT_KEYS = frozenset({"source_periods", "source_second_moment", "target_period"}) +_MUON_HINT_KEYS = frozenset({"source_periods", "target_period"}) +_PARAMETER_KEYS = frozenset({"identity", "algorithm_options", "state_variant", "state", "projection_hints"}) + + +@dataclass(frozen=True, slots=True) +class PortableStateLimits: + """Public resource limits shared by portable semantic and collective I/O.""" + + max_fragment_tensor_bytes: int + max_collective_tensor_bytes: int + max_collective_metadata_bytes: int + chunk_bytes: int = 8 << 20 + max_members: int = 4096 + max_metadata_bytes: int = 64 << 20 + max_tree_nodes: int = 1_000_000 + max_tree_depth: int = 64 + max_container_items: int = 1_000_000 + max_string_bytes: int = 1 << 20 + max_integer_bytes: int = 4096 + max_tensors: int = 262_144 + max_tensor_rank: int = 64 + diagnostic_bytes: int = 2048 + + def __post_init__(self) -> None: + self._wire_limits() + + def _wire_limits(self, *, collective: bool = False) -> _CanonicalWireLimits: + return _CanonicalWireLimits( + max_fragment_tensor_bytes=( + self.max_collective_tensor_bytes if collective else self.max_fragment_tensor_bytes + ), + max_collective_tensor_bytes=self.max_collective_tensor_bytes, + max_collective_metadata_bytes=self.max_collective_metadata_bytes, + chunk_bytes=self.chunk_bytes, + max_members=self.max_members, + max_metadata_bytes=(self.max_collective_metadata_bytes if collective else self.max_metadata_bytes), + max_tree_nodes=self.max_tree_nodes, + max_tree_depth=self.max_tree_depth, + max_container_items=self.max_container_items, + max_string_bytes=self.max_string_bytes, + max_integer_bytes=self.max_integer_bytes, + max_tensors=self.max_tensors, + max_tensor_rank=self.max_tensor_rank, + diagnostic_bytes=self.diagnostic_bytes, + ) + + +def _require_limits(limits) -> PortableStateLimits: + if type(limits) is not PortableStateLimits: + raise TypeError("limits must be a PortableStateLimits") + return limits + + +def _bounded_clone(value, limits: PortableStateLimits, *, collective: bool = False): + wire_limits = limits._wire_limits(collective=collective) + plan = _prepare_canonical_wire_value(value, wire_limits) + prepared = _parse_canonical_wire_metadata(plan.metadata, limits=wire_limits) + return _reconstruct_canonical_wire_value( + prepared, + plan.payload_tensors, + expected_fragment_digest=plan.fragment_digest, + ) + + +def _exact_record(value, keys, *, name: str): + if type(value) is not dict or set(value) != keys: + raise ValueError("{} has an invalid schema".format(name)) + return value + + +def _strict_int(value, *, name: str, minimum: int = 0, maximum=None) -> int: + if type(value) is not int: + raise ValueError("{} must be an int".format(name)) + if value < minimum: + raise ValueError("{} must be at least {}".format(name, minimum)) + if maximum is not None and value > maximum: + raise ValueError("{} must be at most {}".format(name, maximum)) + return value + + +def _strict_float(value, *, name: str, minimum=None, maximum=None, maximum_open=False) -> float: + if type(value) is not float or not math.isfinite(value): + raise ValueError("{} must be a finite float".format(name)) + if minimum is not None and value < minimum: + raise ValueError("{} is below its minimum".format(name)) + if maximum is not None and (value > maximum or (maximum_open and value == maximum)): + raise ValueError("{} is above its maximum".format(name)) + return value + + +def _strict_name(value, *, name: str) -> str: + if type(value) is not str or not value or value != value.strip() or "\x00" in value: + raise ValueError("{} must be a non-empty canonical string".format(name)) + return value + + +def _float_bits_equal(left: float, right: float) -> bool: + return struct.pack(">d", left) == struct.pack(">d", right) + + +def _values_equal(left, right) -> bool: + if type(left) is not type(right): + return False + if type(left) is torch.Tensor: + return left.dtype == right.dtype and tuple(left.shape) == tuple(right.shape) and _tensor_bits_equal(left, right) + if type(left) is float: + return _float_bits_equal(left, right) + if type(left) is dict: + return set(left) == set(right) and all(_values_equal(left[key], right[key]) for key in left) + if type(left) in {list, tuple}: + return len(left) == len(right) and all(_values_equal(a, b) for a, b in zip(left, right)) + return left == right + + +def _tight_fp32(value, *, name: str, shape=None, nonnegative: bool = False) -> torch.Tensor: + if ( + type(value) is not torch.Tensor + or value.layout is not torch.strided + or value.device.type != "cpu" + or value.dtype != torch.float32 + or value.is_meta + or value.is_nested + or value.is_quantized + or value.requires_grad + or not value.is_contiguous() + or value.storage_offset() != 0 + or value.untyped_storage().nbytes() != value.numel() * value.element_size() + ): + raise ValueError("{} must be a tight detached CPU fp32 tensor".format(name)) + if shape is not None and tuple(value.shape) != tuple(shape): + raise ValueError("{} has invalid shape".format(name)) + if not bool(torch.isfinite(value).all()): + raise ValueError("{} must be finite".format(name)) + if nonnegative and not bool((value >= 0).all()): + raise ValueError("{} must be nonnegative".format(name)) + return value + + +def _tight_clone(value: torch.Tensor) -> torch.Tensor: + result = torch.empty(tuple(value.shape), dtype=torch.float32, device="cpu") + result.copy_(value) + return result + + +def _normalize_codebook(value, *, global_step: int): + if value is None: + return None + value = _tight_fp32(value, name="common.gefen_codebook", shape=(256,)) + _validate_codebook(value) + if float(value[0].item()) != -1.0 or float(value[-1].item()) != 1.0: + raise ValueError("portable codebook must retain exact -1 and +1 endpoints") + return value + + +def _normalize_common(value): + value = _exact_record(value, _COMMON_KEYS, name="portable common state") + global_step = _strict_int(value["gefen_global_step"], name="gefen_global_step", maximum=_MAX_EXACT_COUNTER) + if type(value["gefen_deterministic"]) is not bool: + raise ValueError("gefen_deterministic must be a bool") + return { + "gefen_global_step": global_step, + "gefen_codebook": _normalize_codebook(value["gefen_codebook"], global_step=global_step), + "gefen_deterministic": value["gefen_deterministic"], + } + + +def _normalize_policy(value, implementation: str): + value = _exact_record(value, _POLICY_KEYS, name="portable policy") + if value["schema_version"] != 1 or type(value["schema_version"]) is not int: + raise ValueError("unsupported portable semantic policy schema_version") + _strict_int(value["codebook_refresh_every"], name="codebook_refresh_every") + if value["stochastic_round"] is not False: + raise ValueError("portable semantic state requires stochastic_round=False") + if value["momentum_projection"] != _MOMENTUM_PROJECTION: + raise ValueError("unsupported portable momentum projection") + if value["second_moment_projection"] != _SECOND_MOMENT_PROJECTION: + raise ValueError("unsupported portable second-moment projection") + for key in ("factored_v_2d", "force_1d_period_one", "force_2d_period_one"): + if type(value[key]) is not bool: + raise ValueError("{} must be a bool".format(key)) + if implementation == "gefen.GefenMuon" and value["factored_v_2d"]: + raise ValueError("Muon portable policy requires factored_v_2d=False") + substrings = value["period_one_substrings"] + if type(substrings) is not list or any(type(item) is not str or item != item.lower() for item in substrings): + raise ValueError("period_one_substrings must be a canonical lowercase string list") + return {**value, "period_one_substrings": list(substrings)} + + +def _normalize_ns_schedule(value): + if type(value) is not list or not value or len(value) >= 100: + raise ValueError("ns_schedule must contain between 1 and 99 entries") + if any(type(item) is not list or len(item) != 3 for item in value): + raise ValueError("ns_schedule must contain three-float lists") + return [[_strict_float(component, name="ns_schedule") for component in item] for item in value] + + +def _normalize_options(value, implementation: str): + keys = _PLAIN_OPTION_KEYS if implementation == "gefen.Gefen" else _MUON_OPTION_KEYS + value = _exact_record(value, keys, name="portable algorithm options") + result = dict(value) + result["lr"] = _strict_float(value["lr"], name="lr", minimum=0.0) + result["weight_decay"] = _strict_float(value["weight_decay"], name="weight_decay", minimum=0.0) + if implementation == "gefen.Gefen": + result["beta1"] = _strict_float(value["beta1"], name="beta1", minimum=0.0, maximum=1.0, maximum_open=True) + result["beta2"] = _strict_float(value["beta2"], name="beta2", minimum=0.0, maximum=1.0, maximum_open=True) + result["eps"] = _strict_float(value["eps"], name="eps", minimum=0.0) + if result["eps"] == 0.0: + raise ValueError("eps must be positive") + if value["second_moment_policy"] not in {"block", "factored"}: + raise ValueError("second_moment_policy must be 'block' or 'factored'") + return result + + result["momentum"] = _strict_float(value["momentum"], name="momentum", minimum=0.0, maximum=1.0, maximum_open=True) + for key in ("nesterov", "fp8_ns", "fp8_ns_compile", "batched_ns", "normuon", "cautious"): + if type(value[key]) is not bool: + raise ValueError("{} must be a bool".format(key)) + result["ns_schedule"] = _normalize_ns_schedule(value["ns_schedule"]) + result["ns_eps"] = _strict_float(value["ns_eps"], name="ns_eps", minimum=0.0) + if result["ns_eps"] == 0.0: + raise ValueError("ns_eps must be positive") + if value["adjust_lr_fn"] not in {None, "original", "match_rms_adamw"}: + raise ValueError("adjust_lr_fn is unsupported") + if value["sharded_mode"] not in {"exact", "approx", "distributed"}: + raise ValueError("sharded_mode is unsupported") + result["batched_ns_workspace_bytes"] = _strict_int( + value["batched_ns_workspace_bytes"], name="batched_ns_workspace_bytes", minimum=1 + ) + result["normuon_beta2"] = _strict_float( + value["normuon_beta2"], name="normuon_beta2", minimum=0.0, maximum=1.0, maximum_open=True + ) + result["normuon_eps"] = _strict_float(value["normuon_eps"], name="normuon_eps", minimum=0.0) + if result["normuon_eps"] == 0.0: + raise ValueError("normuon_eps must be positive") + return result + + +def _normalize_periods(value, *, name: str): + if type(value) is not list or any(type(period) is not int or period <= 0 for period in value): + raise ValueError("{} must be a list of positive ints".format(name)) + if value != sorted(set(value)): + raise ValueError("{} must be sorted and unique".format(name)) + if any(period != 1 for period in value): + raise ValueError("exact portable v3 state supports only source period one") + return list(value) + + +def _normalize_complete_parameter_record(fqn, value, implementation: str, global_step: int, policy): + value = _exact_record(value, _PARAMETER_KEYS, name="portable parameter {!r}".format(fqn)) + identity_record = _normalize_parameter_identity(value["identity"]) + identity = _parse_parameter_identity(identity_record) + if identity.fqn != fqn: + raise ValueError("portable parameter identity does not match its FQN key") + options = _normalize_options(value["algorithm_options"], implementation) + variant = value["state_variant"] + if type(variant) is not str: + raise ValueError("portable state_variant must be a string") + state = value["state"] + hints = value["projection_hints"] + if type(state) is not dict or type(hints) is not dict: + raise ValueError("portable parameter state and projection_hints must be dictionaries") + shape = identity.global_shape + if identity.numel == 0 and variant != "pristine": + raise ValueError("empty logical parameters must remain pristine") + + if implementation == "gefen.Gefen": + _exact_record(hints, _PLAIN_HINT_KEYS, name="plain projection_hints") + periods = _normalize_periods(hints["source_periods"], name="source_periods") + if hints["target_period"] != 1 or type(hints["target_period"]) is not int: + raise ValueError("portable target_period must be exactly one") + if hints["source_second_moment"] not in {None, "block", "factored"}: + raise ValueError("portable source_second_moment is invalid") + expected_representation = "factored" if policy["factored_v_2d"] and len(shape) == 2 else "block" + if options["second_moment_policy"] != expected_representation: + raise ValueError("parameter second_moment_policy conflicts with the optimizer policy") + if variant == "pristine": + expected_keys = frozenset() + if periods or hints["source_second_moment"] is not None: + raise ValueError("pristine portable state has invalid projection hints") + elif variant == "period_selected": + expected_keys = frozenset() + if not periods or hints["source_second_moment"] is not None: + raise ValueError("period-selected portable state has invalid projection hints") + elif variant == "initialized_dense": + expected_keys = frozenset({"step", "momentum", "second_moment", "second_moment_step"}) + if not periods or hints["source_second_moment"] != "block" or expected_representation != "block": + raise ValueError("dense block state conflicts with its policy or projection hints") + elif variant == "initialized_factored": + expected_keys = frozenset({"step", "momentum", "v_row", "v_col", "factored_step"}) + if not periods or hints["source_second_moment"] != "factored" or expected_representation != "factored": + raise ValueError("factored state conflicts with its policy or projection hints") + else: + raise ValueError("unsupported plain portable state_variant") + _exact_record(state, expected_keys, name="plain portable parameter state") + normalized_state = dict(state) + if variant.startswith("initialized_"): + if identity.numel == 0: + raise ValueError("empty logical parameters cannot carry initialized state") + step = _strict_int(state["step"], name="step", minimum=1, maximum=_MAX_EXACT_COUNTER) + if step > global_step: + raise ValueError("parameter step exceeds optimizer global step") + normalized_state["momentum"] = _tight_fp32(state["momentum"], name="momentum", shape=shape) + if variant == "initialized_dense": + second_step = _strict_int( + state["second_moment_step"], name="second_moment_step", minimum=1, maximum=_MAX_EXACT_COUNTER + ) + normalized_state["second_moment"] = _tight_fp32( + state["second_moment"], name="second_moment", shape=shape, nonnegative=True + ) + else: + if len(shape) != 2: + raise ValueError("factored portable state requires a logical matrix") + second_step = _strict_int( + state["factored_step"], name="factored_step", minimum=1, maximum=_MAX_EXACT_COUNTER + ) + normalized_state["v_row"] = _tight_fp32( + state["v_row"], name="v_row", shape=(shape[0],), nonnegative=True + ) + normalized_state["v_col"] = _tight_fp32( + state["v_col"], name="v_col", shape=(shape[1],), nonnegative=True + ) + if second_step > step: + raise ValueError("secondary parameter counter exceeds step") + return { + "identity": identity_record, + "algorithm_options": options, + "state_variant": variant, + "state": normalized_state, + "projection_hints": { + "source_periods": periods, + "source_second_moment": hints["source_second_moment"], + "target_period": 1, + }, + } + + _exact_record(hints, _MUON_HINT_KEYS, name="Muon projection_hints") + periods = _normalize_periods(hints["source_periods"], name="source_periods") + if hints["target_period"] != 1 or type(hints["target_period"]) is not int: + raise ValueError("portable target_period must be exactly one") + if len(shape) != 2: + raise ValueError("portable Muon parameters must be logical matrices") + if variant == "pristine": + expected_keys = frozenset() + if periods: + raise ValueError("pristine Muon state must not contain source periods") + elif variant == "period_selected": + expected_keys = frozenset() + if not periods: + raise ValueError("period-selected Muon state requires source periods") + elif variant == "initialized_dense": + expected_keys = frozenset({"step", "momentum"}) + if options["normuon"] or not periods: + raise ValueError("initialized Muon state conflicts with its options or periods") + elif variant == "initialized_dense_normuon": + expected_keys = frozenset({"step", "momentum", "normuon_v", "normuon_step"}) + if not options["normuon"] or not periods: + raise ValueError("initialized NorMuon state conflicts with its options or periods") + else: + raise ValueError("unsupported Muon portable state_variant") + _exact_record(state, expected_keys, name="Muon portable parameter state") + normalized_state = dict(state) + if variant.startswith("initialized_"): + if identity.numel == 0: + raise ValueError("empty logical parameters cannot carry initialized state") + step = _strict_int(state["step"], name="step", minimum=1, maximum=_MAX_EXACT_COUNTER) + if step > global_step: + raise ValueError("parameter step exceeds optimizer global step") + normalized_state["momentum"] = _tight_fp32(state["momentum"], name="momentum", shape=shape) + if variant == "initialized_dense_normuon": + normuon_step = _strict_int( + state["normuon_step"], name="normuon_step", minimum=1, maximum=_MAX_EXACT_COUNTER + ) + if normuon_step > step: + raise ValueError("normuon_step exceeds parameter step") + normalized_state["normuon_v"] = _tight_fp32( + state["normuon_v"], name="normuon_v", shape=(shape[0], 1), nonnegative=True + ) + return { + "identity": identity_record, + "algorithm_options": options, + "state_variant": variant, + "state": normalized_state, + "projection_hints": {"source_periods": periods, "target_period": 1}, + } + + +def _normalize_gefen_portable_state_document(state, *, limits, expected_implementation=None): + """Bound, clone, and semantically validate one complete portable v3 document.""" + + limits = _require_limits(limits) + if expected_implementation is not None and expected_implementation not in _IMPLEMENTATIONS: + raise ValueError("unsupported expected portable implementation") + bounded = _bounded_clone(state, limits, collective=True) + document = normalize_portable_state_document(bounded, expected_implementation=expected_implementation) + implementation = document["implementation"] + if implementation not in _IMPLEMENTATIONS: + raise ValueError("unsupported portable Gefen implementation") + if document["provenance"] is not None: + raise ValueError("Gefen portable v3 provenance must be None") + policy = _normalize_policy(document["policy"], implementation) + common = _normalize_common(document["common"]) + parameters = document["parameters"] + if type(parameters) is not dict or not parameters: + raise ValueError("portable parameters must be a non-empty FQN mapping") + normalized_parameters = {} + for fqn in sorted(parameters): + _strict_name(fqn, name="parameter FQN") + normalized_parameters[fqn] = _normalize_complete_parameter_record( + fqn, + parameters[fqn], + implementation, + common["gefen_global_step"], + policy, + ) + if common["gefen_codebook"] is None and any( + record["state_variant"] != "pristine" for record in normalized_parameters.values() + ): + raise ValueError("non-pristine portable parameter state requires a codebook") + normalized = build_portable_state_document( + implementation=implementation, + policy=policy, + common=common, + parameters=normalized_parameters, + provenance=None, + ) + if not _values_equal(normalized, document): + raise ValueError("portable state is not in canonical semantic form") + return normalized + + +def _derived_role(shard: ShardIdentity) -> str: + if shard.layout is ParameterLayout.REPLICATED: + return "live" if shard.parameter.numel else "empty_replicated" + if shard.layout is ParameterLayout.FLATTENED_ELEMENT_SHARD: + return "live" if shard.logical_slice.length else "empty_flat" + if shard.layout is ParameterLayout.WHOLE_PARAMETER_OWNER: + return "whole_owner" if shard.local_member == shard.owner else "whole_nonowner" + raise ValueError("unsupported portable fragment layout") + + +def _local_dense_shape(shard: ShardIdentity): + role = _derived_role(shard) + if role not in {"live", "whole_owner"} or shard.parameter.numel == 0: + return None + if shard.layout in {ParameterLayout.REPLICATED, ParameterLayout.WHOLE_PARAMETER_OWNER}: + return shard.parameter.global_shape + if shard.layout is ParameterLayout.FLATTENED_ELEMENT_SHARD: + return (shard.logical_slice.length,) + raise ValueError("unsupported portable fragment layout") + + +def _normalize_catalog(value, implementation: str, manifest: ShardingManifest, policy): + if type(value) is not dict: + raise ValueError("portable fragment catalog must be an FQN mapping") + manifest_fqns = {shard.parameter.fqn for shard in manifest.shards} + if set(value) != manifest_fqns: + raise ValueError("portable fragment catalog does not match the manifest") + result = {} + for fqn in sorted(value): + record = _exact_record(value[fqn], _CATALOG_KEYS, name="portable catalog entry") + identity_record = _normalize_parameter_identity(record["identity"]) + identity = _parse_parameter_identity(identity_record) + if identity.fqn != fqn or any(shard.parameter != identity for shard in manifest.for_parameter(fqn)): + raise ValueError("portable catalog identity does not match the manifest") + options = _normalize_options(record["algorithm_options"], implementation) + if implementation == "gefen.Gefen": + expected = "factored" if policy["factored_v_2d"] and len(identity.global_shape) == 2 else "block" + if options["second_moment_policy"] != expected: + raise ValueError("catalog second_moment_policy conflicts with the optimizer policy") + if expected == "factored" and any( + shard.layout is not ParameterLayout.REPLICATED for shard in manifest.for_parameter(fqn) + ): + raise ValueError("factored logical matrices require replicated portable shards") + elif len(identity.global_shape) != 2: + raise ValueError("portable Muon catalog parameters must be logical matrices") + result[fqn] = { + "identity": identity_record, + "algorithm_options": options, + } + return result + + +def _normalize_fragment_slot(value, implementation: str, common, catalog, member: str): + value = _exact_record(value, _SLOT_KEYS, name="portable logical slot") + group_index = _strict_int(value["group_index"], name="group_index") + slot_index = _strict_int(value["original_slot_index"], name="original_slot_index") + compatibility_name = _strict_name(value["compatibility_name"], name="compatibility_name") + if compatibility_name != compatibility_name.lower(): + raise ValueError("compatibility_name must be lowercase") + shard_record = _normalize_shard_identity(value["shard"]) + shard = _parse_shard_identity(shard_record) + allowed_layouts = ( + {ParameterLayout.REPLICATED, ParameterLayout.FLATTENED_ELEMENT_SHARD} + if implementation == "gefen.Gefen" + else {ParameterLayout.REPLICATED, ParameterLayout.WHOLE_PARAMETER_OWNER} + ) + if shard.layout not in allowed_layouts: + raise ValueError("portable fragment layout is unsupported for its implementation") + if shard.process_group is None or shard.local_member != member: + raise ValueError("portable fragment member does not match its shard") + role = _derived_role(shard) + if value["role"] != role: + raise ValueError("portable fragment role does not match its shard") + options = _normalize_options(value["algorithm_options"], implementation) + catalog_entry = catalog.get(shard.parameter.fqn) + if catalog_entry is None or not _values_equal(options, catalog_entry["algorithm_options"]): + raise ValueError("portable slot options do not match the catalog") + if ( + implementation == "gefen.GefenMuon" + and shard.layout is ParameterLayout.WHOLE_PARAMETER_OWNER + and options["sharded_mode"] != "distributed" + ): + raise ValueError("whole-parameter Muon state requires sharded_mode='distributed'") + variant = value["state_variant"] + if variant not in { + "pristine", + "period_selected", + "initialized_dense", + "initialized_factored", + "initialized_dense_normuon", + }: + raise ValueError("portable fragment state_variant is unsupported") + period = value["source_period"] + if period is not None: + _strict_int(period, name="source_period", minimum=1) + if period != 1: + raise ValueError("exact portable v3 fragments require source_period=1") + local_shape = _local_dense_shape(shard) + if local_shape is None or math.prod(local_shape) % period != 0: + raise ValueError("source_period does not divide the local logical payload") + source_second = value["source_second_moment"] + if source_second not in {None, "block", "factored"}: + raise ValueError("portable source_second_moment is invalid") + state = value["state"] + if type(state) is not dict: + raise ValueError("portable fragment slot state must be a dictionary") + payload_shape = _local_dense_shape(shard) + payload_role = role in {"live", "whole_owner"} and shard.parameter.numel > 0 + if not payload_role and (variant != "pristine" or period is not None or state or source_second is not None): + raise ValueError("empty and whole-nonowner slots must remain pristine") + if variant == "pristine": + _exact_record(state, frozenset(), name="pristine fragment state") + if period is not None or source_second is not None: + raise ValueError("pristine fragment state has invalid metadata") + elif variant == "period_selected": + _exact_record(state, frozenset(), name="period-selected fragment state") + if period is None or source_second is not None: + raise ValueError("period-selected fragment state has invalid metadata") + elif implementation == "gefen.Gefen" and variant == "initialized_dense": + _exact_record( + state, frozenset({"step", "momentum", "second_moment", "second_moment_step"}), name="dense fragment state" + ) + if period is None or source_second != "block": + raise ValueError("initialized block fragment has invalid metadata") + if options["second_moment_policy"] != "block": + raise ValueError("initialized block fragment conflicts with its algorithm options") + elif implementation == "gefen.Gefen" and variant == "initialized_factored": + _exact_record( + state, frozenset({"step", "momentum", "v_row", "v_col", "factored_step"}), name="factored fragment state" + ) + if period is None or source_second != "factored" or shard.layout is not ParameterLayout.REPLICATED: + raise ValueError("factored fragments require replicated storage and factored metadata") + if options["second_moment_policy"] != "factored": + raise ValueError("initialized factored fragment conflicts with its algorithm options") + elif implementation == "gefen.GefenMuon" and variant in {"initialized_dense", "initialized_dense_normuon"}: + keys = {"step", "momentum"} + if variant == "initialized_dense_normuon": + keys.update({"normuon_v", "normuon_step"}) + _exact_record(state, frozenset(keys), name="Muon fragment state") + if period is None or source_second is not None: + raise ValueError("initialized Muon fragment has invalid metadata") + if (variant == "initialized_dense_normuon") != options["normuon"]: + raise ValueError("initialized Muon fragment conflicts with its NorMuon option") + else: + raise ValueError("fragment state_variant does not match its implementation") + normalized_state = dict(state) + if variant.startswith("initialized_"): + step = _strict_int(state["step"], name="step", minimum=1, maximum=_MAX_EXACT_COUNTER) + if step > common["gefen_global_step"]: + raise ValueError("fragment parameter step exceeds optimizer global step") + normalized_state["momentum"] = _tight_fp32(state["momentum"], name="local momentum", shape=payload_shape) + if variant == "initialized_dense" and implementation == "gefen.Gefen": + second_step = _strict_int( + state["second_moment_step"], name="second_moment_step", minimum=1, maximum=_MAX_EXACT_COUNTER + ) + normalized_state["second_moment"] = _tight_fp32( + state["second_moment"], name="local second_moment", shape=payload_shape, nonnegative=True + ) + if second_step > step: + raise ValueError("fragment second_moment_step exceeds step") + elif variant == "initialized_factored": + rows, columns = shard.parameter.global_shape + factored_step = _strict_int( + state["factored_step"], name="factored_step", minimum=1, maximum=_MAX_EXACT_COUNTER + ) + normalized_state["v_row"] = _tight_fp32(state["v_row"], name="v_row", shape=(rows,), nonnegative=True) + normalized_state["v_col"] = _tight_fp32(state["v_col"], name="v_col", shape=(columns,), nonnegative=True) + if factored_step > step: + raise ValueError("fragment factored_step exceeds step") + elif variant == "initialized_dense_normuon": + normuon_step = _strict_int( + state["normuon_step"], name="normuon_step", minimum=1, maximum=_MAX_EXACT_COUNTER + ) + normalized_state["normuon_v"] = _tight_fp32( + state["normuon_v"], name="normuon_v", shape=(shard.parameter.global_shape[0], 1), nonnegative=True + ) + if normuon_step > step: + raise ValueError("fragment normuon_step exceeds step") + return { + "group_index": group_index, + "original_slot_index": slot_index, + "compatibility_name": compatibility_name, + "shard": shard_record, + "algorithm_options": options, + "role": role, + "source_period": period, + "source_second_moment": source_second, + "state_variant": variant, + "state": normalized_state, + } + + +def _normalize_portable_state_fragment(fragment, *, limits): + """Bound, clone, and validate one reversible local portable-state fragment.""" + + limits = _require_limits(limits) + fragment = _bounded_clone(fragment, limits) + _exact_record(fragment, _FRAGMENT_KEYS, name="portable state fragment") + if fragment["format"] != _FRAGMENT_FORMAT or fragment["format_version"] != _FRAGMENT_FORMAT_VERSION: + raise ValueError("unsupported portable state fragment format") + if type(fragment["format_version"]) is not int or fragment["coverage"] != _FRAGMENT_COVERAGE: + raise ValueError("unsupported portable state fragment coverage") + implementation = fragment["implementation"] + if implementation not in _IMPLEMENTATIONS: + raise ValueError("unsupported portable fragment implementation") + member = _strict_name(fragment["member"], name="portable fragment member") + policy = _normalize_policy(fragment["policy"], implementation) + common = _normalize_common(fragment["common"]) + manifest_record = _normalize_sharding_manifest(fragment["manifest"]) + manifest = _parse_sharding_manifest(manifest_record) + allowed_layouts = ( + {ParameterLayout.REPLICATED, ParameterLayout.FLATTENED_ELEMENT_SHARD} + if implementation == "gefen.Gefen" + else {ParameterLayout.REPLICATED, ParameterLayout.WHOLE_PARAMETER_OWNER} + ) + if any(shard.layout not in allowed_layouts or shard.process_group is None for shard in manifest.shards): + raise ValueError("portable fragment manifest contains an unsupported or ungrouped shard") + catalog = _normalize_catalog(fragment["catalog"], implementation, manifest, policy) + slots_value = fragment["logical_slots"] + if type(slots_value) is not list or not slots_value: + raise ValueError("portable fragment logical_slots must be a non-empty list") + slots = [_normalize_fragment_slot(slot, implementation, common, catalog, member) for slot in slots_value] + if common["gefen_codebook"] is None and any(slot["state_variant"] != "pristine" for slot in slots): + raise ValueError("non-pristine portable fragment state requires a codebook") + positions = [(slot["group_index"], slot["original_slot_index"]) for slot in slots] + if positions[0] != (0, 0) or positions != sorted(positions) or len(set(positions)) != len(positions): + raise ValueError("portable logical slot positions must start at (0, 0) and be strictly ordered") + expected_group = 0 + expected_slot = 0 + for group_index, slot_index in positions: + if group_index == expected_group and slot_index == expected_slot: + expected_slot += 1 + elif group_index == expected_group + 1 and slot_index == 0: + expected_group += 1 + expected_slot = 1 + else: + raise ValueError("portable logical slot positions must be contiguous") + identities = [_parse_shard_identity(slot["shard"]) for slot in slots] + if len(set(identities)) != len(identities): + raise ValueError("portable fragment contains duplicate shard identities") + if any(shard not in manifest.shards for shard in identities): + raise ValueError("portable fragment slot is absent from the manifest") + return { + "format": _FRAGMENT_FORMAT, + "format_version": _FRAGMENT_FORMAT_VERSION, + "coverage": _FRAGMENT_COVERAGE, + "implementation": implementation, + "member": member, + "policy": policy, + "common": common, + "manifest": manifest_record, + "catalog": catalog, + "logical_slots": slots, + } + + +def _build_portable_state_fragment(*, implementation, member, policy, common, manifest, catalog, logical_slots, limits): + """Build one strict local fragment from serialized primitives and local dense fields.""" + + if isinstance(manifest, ShardingManifest): + manifest = _serialize_sharding_manifest(manifest) + fragment = { + "format": _FRAGMENT_FORMAT, + "format_version": _FRAGMENT_FORMAT_VERSION, + "coverage": _FRAGMENT_COVERAGE, + "implementation": implementation, + "member": member, + "policy": policy, + "common": common, + "manifest": manifest, + "catalog": catalog, + "logical_slots": logical_slots, + } + return _normalize_portable_state_fragment(fragment, limits=limits) + + +def _required_payload_slots(slots): + return [ + slot + for slot in slots + if slot["role"] in {"live", "whole_owner"} and _parse_shard_identity(slot["shard"]).parameter.numel > 0 + ] + + +def _consensus(values, *, name: str): + if not values: + raise ValueError("{} requires at least one value".format(name)) + reference = values[0] + if any(not _values_equal(reference, value) for value in values[1:]): + raise ValueError("{} disagree across portable fragments".format(name)) + return reference + + +def _assemble_special_field(slots, *, key: str, shape, name: str) -> torch.Tensor: + values = [] + for slot in slots: + shard = _parse_shard_identity(slot["shard"]) + payload = slot["state"].get(key) + if shard.layout is ParameterLayout.REPLICATED: + values.append(_tight_fp32(payload, name=name, shape=shape, nonnegative=True)) + elif shard.layout is ParameterLayout.WHOLE_PARAMETER_OWNER: + if shard.local_member == shard.owner: + values.append(_tight_fp32(payload, name=name, shape=shape, nonnegative=True)) + elif payload is not None: + raise ValueError("whole-parameter nonowner carries {}".format(name)) + else: + raise ValueError("{} supports only exact replicas or a sole whole owner".format(name)) + return _tight_clone(_consensus(values, name=name)) + + +def _assemble_parameter(fqn, slots, manifest: ShardingManifest, catalog_entry, implementation, policy, global_step): + parameter = _parse_parameter_identity(catalog_entry["identity"]) + required = _required_payload_slots(slots) + variants = {slot["state_variant"] for slot in required} + if not required: + variant = "pristine" + elif variants == {"pristine"}: + variant = "pristine" + elif variants == {"period_selected"}: + variant = "period_selected" + elif len(variants) == 1 and next(iter(variants)).startswith("initialized_"): + variant = next(iter(variants)) + else: + raise ValueError("portable parameter shards disagree on initialization state") + if any(slot["state_variant"] != "pristine" for slot in slots if slot not in required): + raise ValueError("non-payload portable slots must remain pristine") + periods = sorted({slot["source_period"] for slot in required if slot["source_period"] is not None}) + if variant == "pristine" and periods: + raise ValueError("pristine portable parameter has source periods") + if variant != "pristine" and not periods: + raise ValueError("non-pristine portable parameter requires source periods") + state = {} + source_second = None + if variant.startswith("initialized_"): + step = _consensus([slot["state"]["step"] for slot in required], name="parameter step") + if step > global_step: + raise ValueError("parameter step exceeds optimizer global step") + momentum = _assemble_dense_logical_field( + manifest, + parameter, + [(_parse_shard_identity(slot["shard"]), slot["state"].get("momentum")) for slot in slots], + ) + state = {"step": step, "momentum": momentum} + if implementation == "gefen.Gefen" and variant == "initialized_dense": + source_second = "block" + second_step = _consensus( + [slot["state"]["second_moment_step"] for slot in required], name="second_moment_step" + ) + state.update( + { + "second_moment": _assemble_dense_logical_field( + manifest, + parameter, + [(_parse_shard_identity(slot["shard"]), slot["state"].get("second_moment")) for slot in slots], + ), + "second_moment_step": second_step, + } + ) + elif implementation == "gefen.Gefen" and variant == "initialized_factored": + source_second = "factored" + if any(_parse_shard_identity(slot["shard"]).layout is not ParameterLayout.REPLICATED for slot in slots): + raise ValueError("factored portable state must be assembled from exact replicas") + factored_step = _consensus([slot["state"]["factored_step"] for slot in required], name="factored_step") + state.update( + { + "v_row": _assemble_special_field( + slots, key="v_row", shape=(parameter.global_shape[0],), name="v_row" + ), + "v_col": _assemble_special_field( + slots, key="v_col", shape=(parameter.global_shape[1],), name="v_col" + ), + "factored_step": factored_step, + } + ) + elif implementation == "gefen.GefenMuon" and variant == "initialized_dense_normuon": + state.update( + { + "normuon_v": _assemble_special_field( + slots, + key="normuon_v", + shape=(parameter.global_shape[0], 1), + name="normuon_v", + ), + "normuon_step": _consensus( + [slot["state"]["normuon_step"] for slot in required], name="normuon_step" + ), + } + ) + if implementation == "gefen.Gefen": + expected = "factored" if policy["factored_v_2d"] and len(parameter.global_shape) == 2 else "block" + if variant == "initialized_dense" and expected != "block": + raise ValueError("block portable state conflicts with factored_v_2d policy") + if variant == "initialized_factored" and expected != "factored": + raise ValueError("factored portable state conflicts with factored_v_2d policy") + hints = {"source_periods": periods, "source_second_moment": source_second, "target_period": 1} + else: + hints = {"source_periods": periods, "target_period": 1} + return { + "identity": _serialize_parameter_identity(parameter), + "algorithm_options": catalog_entry["algorithm_options"], + "state_variant": variant, + "state": state, + "projection_hints": hints, + } + + +def _assemble_portable_state_fragments(fragments, *, process_group_identity, limits): + """Assemble member-ordered fragments into one complete topology-neutral document.""" + + limits = _require_limits(limits) + if not isinstance(process_group_identity, ProcessGroupIdentity): + raise TypeError("process_group_identity must be a ProcessGroupIdentity") + process_group_identity = ProcessGroupIdentity( + process_group_identity.semantic_name, + process_group_identity.ordered_members, + schema_version=process_group_identity.schema_version, + ) + members = tuple(_strict_name(member, name="ordered member") for member in process_group_identity.ordered_members) + if len(set(members)) != len(members) or len(members) > limits.max_members: + raise ValueError("process-group members must be unique and within max_members") + if isinstance(fragments, (str, bytes, bytearray)): + raise TypeError("fragments must be a sequence") + try: + iterator = iter(fragments) + except TypeError as exc: + raise TypeError("fragments must be a sequence") from exc + bounded_fragments = [] + for fragment in iterator: + bounded_fragments.append(fragment) + if len(bounded_fragments) > limits.max_members: + raise ValueError("portable fragments exceed max_members") + fragments = tuple(bounded_fragments) + if len(fragments) != len(members): + raise ValueError("portable fragments must contain every process-group member exactly once") + + normalized = [] + total_tensor_bytes = 0 + total_metadata_bytes = 0 + for fragment in fragments: + normalized_fragment = _normalize_portable_state_fragment(fragment, limits=limits) + plan = _prepare_canonical_wire_value(normalized_fragment, limits._wire_limits()) + if total_tensor_bytes > limits.max_collective_tensor_bytes - plan.total_tensor_bytes: + raise ValueError("portable fragments exceed max_collective_tensor_bytes") + if total_metadata_bytes > limits.max_collective_metadata_bytes - len(plan.metadata): + raise ValueError("portable fragments exceed max_collective_metadata_bytes") + total_tensor_bytes += plan.total_tensor_bytes + total_metadata_bytes += len(plan.metadata) + normalized.append(normalized_fragment) + if tuple(fragment["member"] for fragment in normalized) != members: + raise ValueError("portable fragments are not in process-group member order") + + implementation = _consensus([fragment["implementation"] for fragment in normalized], name="implementation") + policy = _consensus([fragment["policy"] for fragment in normalized], name="policy") + common = _consensus([fragment["common"] for fragment in normalized], name="common") + manifest_record = _consensus([fragment["manifest"] for fragment in normalized], name="manifest") + catalog = _consensus([fragment["catalog"] for fragment in normalized], name="catalog") + manifest = _parse_sharding_manifest(manifest_record) + if any(shard.process_group != process_group_identity for shard in manifest.shards): + raise ValueError("portable manifest process groups must exactly match process_group_identity") + all_slots = [slot for fragment in normalized for slot in fragment["logical_slots"]] + shard_slots = {_parse_shard_identity(slot["shard"]): slot for slot in all_slots} + if len(shard_slots) != len(all_slots) or set(shard_slots) != set(manifest.shards): + raise ValueError("portable fragments must exactly cover the complete manifest") + + parameters = {} + parameter_positions = {} + for fqn in sorted(catalog): + slots = [shard_slots[shard] for shard in manifest.for_parameter(fqn)] + parameter_positions[fqn] = _consensus( + [(slot["group_index"], slot["original_slot_index"]) for slot in slots], + name="logical slot position", + ) + _consensus([slot["compatibility_name"] for slot in slots], name="compatibility_name") + parameters[fqn] = _assemble_parameter( + fqn, + slots, + manifest, + catalog[fqn], + implementation, + policy, + common["gefen_global_step"], + ) + if len(set(parameter_positions.values())) != len(parameter_positions): + raise ValueError("portable parameters must have unique logical slot positions") + document = build_portable_state_document( + implementation=implementation, + policy=policy, + common=common, + parameters=parameters, + provenance=None, + ) + return _normalize_gefen_portable_state_document( + document, + limits=limits, + expected_implementation=implementation, + ) + + +def _project_portable_parameter_state( + document_record, + target_shard, + *, + implementation, + global_step, + codebook, + target_algorithm_options, + target_second_moment=None, +): + """Project one normalized global record to a target shard's native state fields.""" + + if implementation not in _IMPLEMENTATIONS: + raise ValueError("unsupported portable projection implementation") + global_step = _strict_int(global_step, name="gefen_global_step", maximum=_MAX_EXACT_COUNTER) + common_codebook = _normalize_codebook(codebook, global_step=global_step) + if not isinstance(target_shard, ShardIdentity): + raise TypeError("target_shard must be a ShardIdentity") + target_layouts = ( + {ParameterLayout.REPLICATED, ParameterLayout.FLATTENED_ELEMENT_SHARD} + if implementation == "gefen.Gefen" + else {ParameterLayout.REPLICATED, ParameterLayout.WHOLE_PARAMETER_OWNER} + ) + if target_shard.layout not in target_layouts: + raise ValueError("target shard layout is unsupported for the portable implementation") + if target_shard.process_group is None: + raise ValueError("portable projection requires a grouped target shard") + source_factored = ( + implementation == "gefen.Gefen" + and type(document_record) is dict + and type(document_record.get("algorithm_options")) is dict + and document_record["algorithm_options"].get("second_moment_policy") == "factored" + ) + policy = { + "schema_version": 1, + "factored_v_2d": source_factored, + "force_1d_period_one": False, + "force_2d_period_one": False, + "period_one_substrings": [], + "codebook_refresh_every": 0, + "stochastic_round": False, + "momentum_projection": _MOMENTUM_PROJECTION, + "second_moment_projection": _SECOND_MOMENT_PROJECTION, + } + fqn = target_shard.parameter.fqn + record = _normalize_complete_parameter_record(fqn, document_record, implementation, global_step, policy) + if _parse_parameter_identity(record["identity"]) != target_shard.parameter: + raise ValueError("portable parameter identity does not match the target shard") + target_options = _normalize_options(target_algorithm_options, implementation) + if not _values_equal(record["algorithm_options"], target_options): + raise ValueError("portable algorithm options do not match the target") + if implementation == "gefen.Gefen": + if target_second_moment not in {"block", "factored"}: + raise ValueError("plain Gefen projection requires target_second_moment") + if target_second_moment != target_options["second_moment_policy"]: + raise ValueError("target_second_moment conflicts with target algorithm options") + if ( + target_second_moment == "factored" + and len(target_shard.parameter.global_shape) == 2 + and target_shard.layout is not ParameterLayout.REPLICATED + ): + raise ValueError("factored logical matrices require replicated target shards") + source_second = record["projection_hints"]["source_second_moment"] + if source_second is not None and source_second != target_second_moment: + raise ValueError("portable second-moment representation migration is unsupported") + else: + if target_second_moment is not None: + raise ValueError("Muon projection does not accept target_second_moment") + if ( + target_shard.layout is ParameterLayout.WHOLE_PARAMETER_OWNER + and target_options["sharded_mode"] != "distributed" + ): + raise ValueError("whole-parameter Muon targets require sharded_mode='distributed'") + + variant = record["state_variant"] + if variant == "pristine": + return {} + if common_codebook is None: + raise ValueError("non-pristine portable projection requires a codebook") + if variant == "period_selected": + return {"automatic_period": 1} if _local_dense_shape(target_shard) is not None else {} + local_momentum = _project_dense_logical_field( + target_shard.parameter, + record["state"]["momentum"], + target_shard, + ) + if local_momentum is None: + return {} + step = record["state"]["step"] + indices, magnitudes = _recompress_dense_momentum(local_momentum, common_codebook, period=1, step=step) + result = { + "automatic_period": 1, + "step": step, + "m_codebook": indices, + "m_magnitude": magnitudes, + } + if implementation == "gefen.Gefen" and variant == "initialized_dense": + local_second = _project_dense_logical_field( + target_shard.parameter, + record["state"]["second_moment"], + target_shard, + ) + if local_second is None: + return {} + second_step = record["state"]["second_moment_step"] + result.update( + { + "vmean": _reduce_block_second_moment(local_second, period=1, step=second_step), + "vmean_step": second_step, + } + ) + elif implementation == "gefen.Gefen" and variant == "initialized_factored": + if target_shard.layout is not ParameterLayout.REPLICATED: + raise ValueError("factored portable state projects only to replicated targets") + result.update( + { + "v_row": _tight_clone(record["state"]["v_row"]), + "v_col": _tight_clone(record["state"]["v_col"]), + "factored_step": record["state"]["factored_step"], + } + ) + elif implementation == "gefen.GefenMuon" and variant == "initialized_dense_normuon": + if target_shard.layout not in {ParameterLayout.REPLICATED, ParameterLayout.WHOLE_PARAMETER_OWNER}: + raise ValueError("NorMuon state projects only to replicas or a whole-parameter owner") + result.update( + { + "normuon_v": _tight_clone(record["state"]["normuon_v"]), + "normuon_step": record["state"]["normuon_step"], + } + ) + return result + + +__all__ = ["PortableStateLimits"] diff --git a/src/gefen/portable_wire.py b/src/gefen/portable_wire.py new file mode 100644 index 0000000..9c4a02e --- /dev/null +++ b/src/gefen/portable_wire.py @@ -0,0 +1,747 @@ +"""Bounded binary wire codec for canonical optimizer-state fragments.""" + +from __future__ import annotations + +from dataclasses import dataclass +import hashlib +import hmac +import math +import struct +import sys + +import torch + + +_CANONICAL_WIRE_PREAMBLE = b"GFNCV1\0\0" +_CANONICAL_WIRE_CODEC = 1 +_CANONICAL_WIRE_TENSOR_DOMAIN = b"gefen.canonical_wire.tensor.v1\0" +_CANONICAL_WIRE_FRAGMENT_DOMAIN = b"gefen.canonical_wire.fragment.v1\0" +_SHA256_BYTES = hashlib.sha256().digest_size +_U16_MAX = (1 << 16) - 1 +_I64_MAX = (1 << 63) - 1 +_U64_MAX = (1 << 64) - 1 +_MAX_SAFE_TREE_DEPTH = 256 + +_TAG_NONE = 0x00 +_TAG_FALSE = 0x01 +_TAG_TRUE = 0x02 +_TAG_INTEGER = 0x03 +_TAG_FLOAT = 0x04 +_TAG_STRING = 0x05 +_TAG_LIST = 0x06 +_TAG_TUPLE = 0x07 +_TAG_DICTIONARY = 0x08 +_TAG_TENSOR = 0x09 + +_DTYPE_RECORDS = ( + (1, torch.bool, 1), + (2, torch.uint8, 1), + (3, torch.int8, 1), + (4, torch.int16, 2), + (5, torch.int32, 4), + (6, torch.int64, 8), + (7, torch.float16, 2), + (8, torch.bfloat16, 2), + (9, torch.float32, 4), + (10, torch.float64, 8), + (11, torch.complex64, 8), + (12, torch.complex128, 16), +) +_DTYPE_BY_CODE = {code: (dtype, size) for code, dtype, size in _DTYPE_RECORDS} +_DTYPE_BY_VALUE = {dtype: (code, size) for code, dtype, size in _DTYPE_RECORDS} + + +def _require_positive_limit(name: str, value: int, *, maximum: int = _U64_MAX) -> None: + if type(value) is not int: + raise TypeError("{} must be an int".format(name)) + if value <= 0 or value > maximum: + raise ValueError("{} must be in [1, {}]".format(name, maximum)) + + +@dataclass(frozen=True, slots=True) +class _CanonicalWireLimits: + max_fragment_tensor_bytes: int + max_collective_tensor_bytes: int + max_collective_metadata_bytes: int + chunk_bytes: int = 8 << 20 + max_members: int = 4096 + max_metadata_bytes: int = 64 << 20 + max_tree_nodes: int = 1_000_000 + max_tree_depth: int = 64 + max_container_items: int = 1_000_000 + max_string_bytes: int = 1 << 20 + max_integer_bytes: int = 4096 + max_tensors: int = 262_144 + max_tensor_rank: int = 64 + diagnostic_bytes: int = 2048 + + def __post_init__(self) -> None: + for name in ( + "max_fragment_tensor_bytes", + "max_collective_tensor_bytes", + "max_collective_metadata_bytes", + "chunk_bytes", + "max_members", + "max_metadata_bytes", + "max_tree_nodes", + "max_tree_depth", + "max_container_items", + "max_string_bytes", + "max_integer_bytes", + "max_tensors", + "diagnostic_bytes", + ): + _require_positive_limit(name, getattr(self, name)) + _require_positive_limit("max_tensor_rank", self.max_tensor_rank, maximum=_U16_MAX) + if self.max_tree_depth > _MAX_SAFE_TREE_DEPTH: + raise ValueError( + "max_tree_depth cannot exceed {} for recursive parsing".format( + _MAX_SAFE_TREE_DEPTH + ) + ) + if self.max_fragment_tensor_bytes > self.max_collective_tensor_bytes: + raise ValueError("max_fragment_tensor_bytes cannot exceed max_collective_tensor_bytes") + if self.max_metadata_bytes > self.max_collective_metadata_bytes: + raise ValueError("max_metadata_bytes cannot exceed max_collective_metadata_bytes") + + +@dataclass(frozen=True, slots=True) +class _CanonicalWireTensorSpec: + index: int + dtype_code: int + dtype: torch.dtype + shape: tuple[int, ...] + numel: int + nbytes: int + digest: bytes + + def __post_init__(self) -> None: + if type(self.index) is not int or self.index < 0 or self.index > _U64_MAX: + raise ValueError("tensor index must be a nonnegative int") + if type(self.dtype_code) is not int: + raise TypeError("tensor dtype code must be an int") + record = _DTYPE_BY_CODE.get(self.dtype_code) + if record is None or record[0] != self.dtype: + raise ValueError("tensor dtype code and dtype do not agree") + if ( + type(self.shape) is not tuple + or len(self.shape) > _U16_MAX + or any( + type(dimension) is not int + or dimension < 0 + or dimension > _I64_MAX + for dimension in self.shape + ) + ): + raise ValueError("tensor shape must be a tuple of nonnegative ints") + if type(self.numel) is not int or self.numel < 0 or self.numel > _U64_MAX: + raise ValueError("tensor numel must be a nonnegative int") + if type(self.nbytes) is not int or self.nbytes < 0 or self.nbytes > _U64_MAX: + raise ValueError("tensor nbytes must be a nonnegative int") + if type(self.digest) is not bytes or len(self.digest) != _SHA256_BYTES: + raise ValueError("tensor digest must be a 32-byte value") + if _shape_numel(self.shape) != self.numel or self.numel * record[1] != self.nbytes: + raise ValueError("tensor shape, numel, and nbytes do not agree") + + +@dataclass(frozen=True, slots=True) +class _CanonicalWirePlan: + metadata: bytes + tensors: tuple[torch.Tensor, ...] + tensor_specs: tuple[_CanonicalWireTensorSpec, ...] + metadata_digest: bytes + fragment_digest: bytes + total_tensor_bytes: int + + def __post_init__(self) -> None: + if type(self.metadata) is not bytes: + raise TypeError("canonical wire metadata must be bytes") + if type(self.tensors) is not tuple or any(type(value) is not torch.Tensor for value in self.tensors): + raise TypeError("canonical wire tensors must be a tuple of plain tensors") + if type(self.tensor_specs) is not tuple or any(type(value) is not _CanonicalWireTensorSpec for value in self.tensor_specs): + raise TypeError("canonical wire tensor specs must be a tuple of tensor specs") + if len(self.tensors) != len(self.tensor_specs): + raise ValueError("canonical wire tensors and specs must have equal lengths") + if any(spec.index != index for index, spec in enumerate(self.tensor_specs)): + raise ValueError("canonical wire tensor specs must have consecutive indices") + if any(not _tight_cpu_tensor(tensor, spec) for tensor, spec in zip(self.tensors, self.tensor_specs)): + raise ValueError("canonical wire tensors must be tight CPU tensors matching their specs") + _validate_digest_field("metadata_digest", self.metadata_digest) + _validate_digest_field("fragment_digest", self.fragment_digest) + if type(self.total_tensor_bytes) is not int or self.total_tensor_bytes < 0: + raise ValueError("total_tensor_bytes must be a nonnegative int") + if sum(spec.nbytes for spec in self.tensor_specs) != self.total_tensor_bytes: + raise ValueError("total_tensor_bytes does not match the tensor specs") + if not hmac.compare_digest(hashlib.sha256(self.metadata).digest(), self.metadata_digest): + raise ValueError("metadata_digest does not match canonical wire metadata") + if not hmac.compare_digest(_canonical_fragment_digest(self.metadata_digest, self.tensor_specs), self.fragment_digest): + raise ValueError("fragment_digest does not match canonical wire metadata and tensor specs") + + @property + def payload_tensors(self) -> tuple[torch.Tensor, ...]: + return self.tensors + + +@dataclass(frozen=True, slots=True) +class _TensorReference: + index: int + + +@dataclass(frozen=True, slots=True) +class _ListValue: + values: tuple[object, ...] + + +@dataclass(frozen=True, slots=True) +class _TupleValue: + values: tuple[object, ...] + + +@dataclass(frozen=True, slots=True) +class _DictionaryValue: + values: tuple[tuple[str, object], ...] + + +@dataclass(frozen=True, slots=True) +class _PreparedCanonicalWireValue: + tensor_specs: tuple[_CanonicalWireTensorSpec, ...] + metadata_digest: bytes + fragment_digest: bytes + total_tensor_bytes: int + _value: object + _limits: _CanonicalWireLimits + + def __post_init__(self) -> None: + if type(self.tensor_specs) is not tuple or any(type(value) is not _CanonicalWireTensorSpec for value in self.tensor_specs): + raise TypeError("canonical wire tensor specs must be a tuple of tensor specs") + if any(spec.index != index for index, spec in enumerate(self.tensor_specs)): + raise ValueError("canonical wire tensor specs must have consecutive indices") + _validate_digest_field("metadata_digest", self.metadata_digest) + _validate_digest_field("fragment_digest", self.fragment_digest) + if type(self.total_tensor_bytes) is not int or self.total_tensor_bytes < 0: + raise ValueError("total_tensor_bytes must be a nonnegative int") + if sum(spec.nbytes for spec in self.tensor_specs) != self.total_tensor_bytes: + raise ValueError("total_tensor_bytes does not match the tensor specs") + if not hmac.compare_digest(_canonical_fragment_digest(self.metadata_digest, self.tensor_specs), self.fragment_digest): + raise ValueError("fragment_digest does not match canonical wire metadata and tensor specs") + if type(self._limits) is not _CanonicalWireLimits: + raise TypeError("prepared canonical wire values require canonical wire limits") + + +def _validate_digest_field(name: str, value: bytes) -> None: + if type(value) is not bytes: + raise TypeError("{} must be bytes".format(name)) + if len(value) != _SHA256_BYTES: + raise ValueError("{} must be a 32-byte value".format(name)) + + +def _shape_numel(shape: tuple[int, ...]) -> int: + result = 1 + for dimension in shape: + if dimension != 0 and result > _U64_MAX // dimension: + raise ValueError("tensor shape product overflows u64") + result *= dimension + return result + + +class _MetadataWriter: + __slots__ = ("_data", "_limit") + + def __init__(self, limit: int): + self._data = bytearray() + self._limit = limit + + def write(self, value: bytes) -> None: + if len(self._data) + len(value) > self._limit: + raise ValueError("canonical wire metadata exceeds max_metadata_bytes") + self._data.extend(value) + + def finish(self) -> bytes: + return bytes(self._data) + + +class _PreparationState: + __slots__ = ("limits", "writer", "nodes", "tensors", "specs", "tensor_bytes") + + def __init__(self, limits: _CanonicalWireLimits): + self.limits = limits + self.writer = _MetadataWriter(limits.max_metadata_bytes) + self.nodes = 0 + self.tensors: list[torch.Tensor] = [] + self.specs: list[_CanonicalWireTensorSpec] = [] + self.tensor_bytes = 0 + + def visit(self, value, *, depth: int, path: str) -> None: + if depth > self.limits.max_tree_depth: + raise ValueError("canonical wire value exceeds max_tree_depth at {}".format(path)) + self.nodes += 1 + if self.nodes > self.limits.max_tree_nodes: + raise ValueError("canonical wire value exceeds max_tree_nodes") + value_type = type(value) + if value is None: + self.writer.write(bytes((_TAG_NONE,))) + return + if value_type is bool: + self.writer.write(bytes((_TAG_TRUE if value else _TAG_FALSE,))) + return + if value_type is int: + magnitude = abs(value) + magnitude_size = (magnitude.bit_length() + 7) // 8 + if magnitude_size > self.limits.max_integer_bytes: + raise ValueError("integer exceeds max_integer_bytes at {}".format(path)) + magnitude_bytes = magnitude.to_bytes(magnitude_size, "big") + self.writer.write(bytes((_TAG_INTEGER, 1 if value < 0 else 0))) + self.writer.write(struct.pack(">Q", len(magnitude_bytes))) + self.writer.write(magnitude_bytes) + return + if value_type is float: + if not math.isfinite(value): + raise ValueError("float must be finite at {}".format(path)) + self.writer.write(bytes((_TAG_FLOAT,))) + self.writer.write(struct.pack(">d", value)) + return + if value_type is str: + self._write_string(value, path=path) + return + if torch.is_tensor(value): + self._write_tensor(value, path=path) + return + if value_type is list or value_type is tuple: + if len(value) > self.limits.max_container_items: + raise ValueError("container exceeds max_container_items at {}".format(path)) + self.writer.write(bytes((_TAG_LIST if value_type is list else _TAG_TUPLE,))) + self.writer.write(struct.pack(">Q", len(value))) + for index, item in enumerate(value): + self.visit(item, depth=depth + 1, path="{}[{}]".format(path, index)) + return + if value_type is dict: + if len(value) > self.limits.max_container_items: + raise ValueError("container exceeds max_container_items at {}".format(path)) + if any(type(key) is not str for key in value): + raise TypeError("dictionary keys must be strings at {}".format(path)) + self.writer.write(bytes((_TAG_DICTIONARY,))) + self.writer.write(struct.pack(">Q", len(value))) + for key in sorted(value): + self.visit(key, depth=depth + 1, path="{}.".format(path)) + self.visit(value[key], depth=depth + 1, path="{}.{}".format(path, key)) + return + raise TypeError("{} has unsupported canonical wire type {}".format(path, value_type.__name__)) + + def _write_string(self, value: str, *, path: str) -> None: + encoded = _bounded_utf8(value, limit=self.limits.max_string_bytes, path=path) + self.writer.write(bytes((_TAG_STRING,))) + self.writer.write(struct.pack(">Q", len(encoded))) + self.writer.write(encoded) + + def _write_tensor(self, value: torch.Tensor, *, path: str) -> None: + if type(value) is not torch.Tensor: + raise TypeError("tensor must be a plain torch.Tensor at {}".format(path)) + if value.layout is not torch.strided or value.is_meta or value.is_nested or value.is_quantized: + raise TypeError("tensor must be a materialized strided tensor at {}".format(path)) + dtype_record = _DTYPE_BY_VALUE.get(value.dtype) + if dtype_record is None: + raise TypeError("tensor dtype {} is not supported at {}".format(value.dtype, path)) + shape = tuple(value.shape) + if len(shape) > self.limits.max_tensor_rank: + raise ValueError("tensor rank exceeds max_tensor_rank at {}".format(path)) + if len(self.tensors) >= self.limits.max_tensors: + raise ValueError("canonical wire value exceeds max_tensors") + dtype_code, element_size = dtype_record + numel = value.numel() + nbytes = numel * element_size + if nbytes > _U64_MAX or self.tensor_bytes + nbytes > self.limits.max_fragment_tensor_bytes: + raise ValueError("canonical wire tensors exceed max_fragment_tensor_bytes") + index = len(self.tensors) + cloned = _clone_tensor(value, chunk_bytes=self.limits.chunk_bytes, path=path) + digest = _canonical_tensor_digest( + cloned, + dtype_code=dtype_code, + shape=shape, + nbytes=nbytes, + chunk_bytes=self.limits.chunk_bytes, + ) + spec = _CanonicalWireTensorSpec(index, dtype_code, value.dtype, shape, numel, nbytes, digest) + self.tensors.append(cloned) + self.specs.append(spec) + self.tensor_bytes += nbytes + self.writer.write(bytes((_TAG_TENSOR,))) + self.writer.write(struct.pack(">Q", index)) + + +def _bounded_utf8(value: str, *, limit: int, path: str) -> bytes: + encoded = bytearray() + characters_per_chunk = max(1, min(1 << 14, limit // 4)) + for start in range(0, len(value), characters_per_chunk): + chunk = value[start : start + characters_per_chunk].encode("utf-8") + if len(encoded) + len(chunk) > limit: + raise ValueError("string exceeds max_string_bytes at {}".format(path)) + encoded.extend(chunk) + return bytes(encoded) + + +def _read_tensor_chunk(value: torch.Tensor, start: int, stop: int) -> torch.Tensor: + detached = value.detach() + if detached.is_contiguous(): + return detached.reshape(-1)[start:stop] + linear = torch.arange(start, stop, dtype=torch.int64, device=detached.device) + remainder = linear + reversed_coordinates = [] + for dimension in reversed(detached.shape): + reversed_coordinates.append(torch.remainder(remainder, dimension)) + remainder = torch.div(remainder, dimension, rounding_mode="floor") + return detached[tuple(reversed(reversed_coordinates))] + + +def _clone_tensor(value: torch.Tensor, *, chunk_bytes: int, path: str) -> torch.Tensor: + cloned = torch.empty(tuple(value.shape), dtype=value.dtype, device="cpu") + flat = cloned.reshape(-1) + scratch_bytes_per_element = value.element_size() + if not value.is_contiguous(): + # Logical indexing retains one int64 coordinate vector per dimension, + # plus the linear index and current quotient. Keep that scratch within + # the same caller-supplied chunk budget as the tensor copy. + scratch_bytes_per_element += 8 * (value.ndim + 2) + chunk_elements = max(1, chunk_bytes // scratch_bytes_per_element) + for start in range(0, value.numel(), chunk_elements): + stop = min(start + chunk_elements, value.numel()) + source = _read_tensor_chunk(value, start, stop).resolve_conj().resolve_neg().reshape(-1) + destination = flat[start:stop] + destination.copy_(source) + if cloned.is_floating_point() or cloned.is_complex(): + try: + finite = bool(torch.isfinite(destination).all()) + except (NotImplementedError, RuntimeError, TypeError) as exc: + raise ValueError("tensor dtype does not support finite values at {}".format(path)) from exc + if not finite: + raise ValueError("tensor must be finite at {}".format(path)) + return cloned + + +def _canonical_tensor_digest( + value: torch.Tensor, + *, + dtype_code: int, + shape: tuple[int, ...], + nbytes: int, + chunk_bytes: int, +) -> bytes: + hasher = hashlib.sha256() + hasher.update(_CANONICAL_WIRE_TENSOR_DOMAIN) + hasher.update(struct.pack(">HH", dtype_code, len(shape))) + for dimension in shape: + hasher.update(struct.pack(">Q", dimension)) + hasher.update(struct.pack(">Q", nbytes)) + element_size = value.element_size() + component_size = element_size // 2 if value.is_complex() else element_size + chunk_elements = max(1, chunk_bytes // element_size) + flat = value.reshape(-1) + for start in range(0, value.numel(), chunk_elements): + stop = min(start + chunk_elements, value.numel()) + raw = flat[start:stop].view(torch.uint8).reshape(-1) + if sys.byteorder == "big" and component_size > 1: + raw = raw.reshape(-1, component_size).flip(1).contiguous().reshape(-1) + elif sys.byteorder != "little": + raise RuntimeError("unsupported native byte order") + hasher.update(memoryview(raw.numpy())) + return hasher.digest() + + +def _canonical_fragment_digest(metadata_digest: bytes, specs: tuple[_CanonicalWireTensorSpec, ...]) -> bytes: + hasher = hashlib.sha256() + hasher.update(_CANONICAL_WIRE_FRAGMENT_DOMAIN) + hasher.update(metadata_digest) + hasher.update(struct.pack(">Q", len(specs))) + for spec in specs: + hasher.update(struct.pack(">Q", spec.nbytes)) + hasher.update(spec.digest) + return hasher.digest() + + +def _prepare_canonical_wire_value(value, limits: _CanonicalWireLimits) -> _CanonicalWirePlan: + """Clone and encode one canonical value into metadata and tensor payloads.""" + + if type(limits) is not _CanonicalWireLimits: + raise TypeError("limits must be _CanonicalWireLimits") + state = _PreparationState(limits) + state.writer.write(_CANONICAL_WIRE_PREAMBLE) + state.writer.write(struct.pack(">HHI", _CANONICAL_WIRE_CODEC, 0, 0)) + state.visit(value, depth=0, path="value") + state.writer.write(struct.pack(">Q", len(state.specs))) + for spec in state.specs: + state.writer.write(struct.pack(">HHI", spec.dtype_code, len(spec.shape), 0)) + for dimension in spec.shape: + state.writer.write(struct.pack(">Q", dimension)) + state.writer.write(struct.pack(">QQ", spec.numel, spec.nbytes)) + state.writer.write(spec.digest) + metadata = state.writer.finish() + specs = tuple(state.specs) + metadata_digest = hashlib.sha256(metadata).digest() + fragment_digest = _canonical_fragment_digest(metadata_digest, specs) + return _CanonicalWirePlan( + metadata=metadata, + tensors=tuple(state.tensors), + tensor_specs=specs, + metadata_digest=metadata_digest, + fragment_digest=fragment_digest, + total_tensor_bytes=state.tensor_bytes, + ) + + +class _MetadataReader: + __slots__ = ("metadata", "limits", "position", "nodes", "next_tensor") + + def __init__(self, metadata: bytes, limits: _CanonicalWireLimits): + self.metadata = metadata + self.limits = limits + self.position = 0 + self.nodes = 0 + self.next_tensor = 0 + + def read(self, size: int) -> bytes: + if size < 0 or size > len(self.metadata) - self.position: + raise ValueError("canonical wire metadata is truncated") + start = self.position + self.position += size + return self.metadata[start : self.position] + + def u8(self) -> int: + return self.read(1)[0] + + def u16(self) -> int: + return struct.unpack(">H", self.read(2))[0] + + def u32(self) -> int: + return struct.unpack(">I", self.read(4))[0] + + def u64(self) -> int: + return struct.unpack(">Q", self.read(8))[0] + + def value(self, *, depth: int, dictionary_key: bool = False): + if depth > self.limits.max_tree_depth: + raise ValueError("canonical wire metadata exceeds max_tree_depth") + self.nodes += 1 + if self.nodes > self.limits.max_tree_nodes: + raise ValueError("canonical wire metadata exceeds max_tree_nodes") + tag = self.u8() + if tag == _TAG_NONE: + result = None + elif tag == _TAG_FALSE: + result = False + elif tag == _TAG_TRUE: + result = True + elif tag == _TAG_INTEGER: + sign = self.u8() + if sign not in {0, 1}: + raise ValueError("canonical wire integer has an invalid sign") + size = self.u64() + if size > self.limits.max_integer_bytes: + raise ValueError("canonical wire integer exceeds max_integer_bytes") + encoded = self.read(size) + if encoded[:1] == b"\0": + raise ValueError("canonical wire integer magnitude is not canonical") + magnitude = int.from_bytes(encoded, "big") + if sign == 1 and magnitude == 0: + raise ValueError("canonical wire integer encodes negative zero") + result = -magnitude if sign else magnitude + elif tag == _TAG_FLOAT: + result = struct.unpack(">d", self.read(8))[0] + if not math.isfinite(result): + raise ValueError("canonical wire float must be finite") + elif tag == _TAG_STRING: + result = self.string_after_tag() + elif tag == _TAG_LIST or tag == _TAG_TUPLE: + count = self.u64() + if count > self.limits.max_container_items: + raise ValueError("canonical wire container exceeds max_container_items") + values = tuple(self.value(depth=depth + 1) for _ in range(count)) + result = _ListValue(values) if tag == _TAG_LIST else _TupleValue(values) + elif tag == _TAG_DICTIONARY: + count = self.u64() + if count > self.limits.max_container_items: + raise ValueError("canonical wire dictionary exceeds max_container_items") + values = [] + previous = None + for _ in range(count): + key = self.value(depth=depth + 1, dictionary_key=True) + if type(key) is not str: + raise ValueError("canonical wire dictionary key must be a string") + if previous is not None and key <= previous: + raise ValueError("canonical wire dictionary keys are not strictly increasing") + previous = key + values.append((key, self.value(depth=depth + 1))) + result = _DictionaryValue(tuple(values)) + elif tag == _TAG_TENSOR: + index = self.u64() + if index != self.next_tensor: + raise ValueError("canonical wire tensor references are not ordered") + if index >= self.limits.max_tensors: + raise ValueError("canonical wire metadata exceeds max_tensors") + self.next_tensor += 1 + result = _TensorReference(index) + else: + raise ValueError("canonical wire metadata has an unknown value tag") + if dictionary_key and type(result) is not str: + raise ValueError("canonical wire dictionary key must be a string") + return result + + def string_after_tag(self) -> str: + size = self.u64() + if size > self.limits.max_string_bytes: + raise ValueError("canonical wire string exceeds max_string_bytes") + encoded = self.read(size) + try: + result = encoded.decode("utf-8") + except UnicodeDecodeError as exc: + raise ValueError("canonical wire string is not valid UTF-8") from exc + if result.encode("utf-8") != encoded: + raise ValueError("canonical wire string is not canonically encoded") + return result + + +def _parse_canonical_wire_metadata( + metadata: bytes, + *, + limits: _CanonicalWireLimits, +) -> _PreparedCanonicalWireValue: + """Validate bounded metadata without allocating any tensor payloads.""" + + if type(metadata) is not bytes: + raise TypeError("metadata must be bytes") + if type(limits) is not _CanonicalWireLimits: + raise TypeError("limits must be _CanonicalWireLimits") + if len(metadata) > limits.max_metadata_bytes: + raise ValueError("canonical wire metadata exceeds max_metadata_bytes") + reader = _MetadataReader(metadata, limits) + if reader.read(len(_CANONICAL_WIRE_PREAMBLE)) != _CANONICAL_WIRE_PREAMBLE: + raise ValueError("canonical wire metadata has an invalid preamble") + if reader.u16() != _CANONICAL_WIRE_CODEC: + raise ValueError("canonical wire metadata uses an unsupported codec") + if reader.u16() != 0 or reader.u32() != 0: + raise ValueError("canonical wire metadata has nonzero reserved header fields") + value = reader.value(depth=0) + descriptor_count = reader.u64() + if descriptor_count != reader.next_tensor: + raise ValueError("canonical wire descriptor count does not match tensor references") + if descriptor_count > limits.max_tensors: + raise ValueError("canonical wire metadata exceeds max_tensors") + specs = [] + total_tensor_bytes = 0 + for index in range(descriptor_count): + dtype_code = reader.u16() + dtype_record = _DTYPE_BY_CODE.get(dtype_code) + if dtype_record is None: + raise ValueError("canonical wire tensor descriptor has an unknown dtype") + rank = reader.u16() + if rank > limits.max_tensor_rank: + raise ValueError("canonical wire tensor rank exceeds max_tensor_rank") + if reader.u32() != 0: + raise ValueError("canonical wire tensor descriptor has nonzero flags") + shape = tuple(reader.u64() for _ in range(rank)) + if any(dimension > _I64_MAX for dimension in shape): + raise ValueError("canonical wire tensor dimension exceeds signed int64") + numel = reader.u64() + nbytes = reader.u64() + digest = reader.read(_SHA256_BYTES) + expected_numel = 1 + for dimension in shape: + if dimension != 0 and expected_numel > _U64_MAX // dimension: + raise ValueError("canonical wire tensor shape product overflows u64") + expected_numel *= dimension + if expected_numel != numel: + raise ValueError("canonical wire tensor shape does not match numel") + dtype, element_size = dtype_record + if numel != 0 and numel > _U64_MAX // element_size: + raise ValueError("canonical wire tensor byte count overflows u64") + if numel * element_size != nbytes: + raise ValueError("canonical wire tensor numel does not match nbytes") + if total_tensor_bytes + nbytes > limits.max_fragment_tensor_bytes: + raise ValueError("canonical wire tensors exceed max_fragment_tensor_bytes") + specs.append(_CanonicalWireTensorSpec(index, dtype_code, dtype, shape, numel, nbytes, digest)) + total_tensor_bytes += nbytes + if reader.position != len(metadata): + raise ValueError("canonical wire metadata has trailing bytes") + specs_tuple = tuple(specs) + metadata_digest = hashlib.sha256(metadata).digest() + fragment_digest = _canonical_fragment_digest(metadata_digest, specs_tuple) + return _PreparedCanonicalWireValue( + tensor_specs=specs_tuple, + metadata_digest=metadata_digest, + fragment_digest=fragment_digest, + total_tensor_bytes=total_tensor_bytes, + _value=value, + _limits=limits, + ) + + +def _tight_cpu_tensor(value: torch.Tensor, spec: _CanonicalWireTensorSpec) -> bool: + if ( + type(value) is not torch.Tensor + or value.layout is not torch.strided + or value.device.type != "cpu" + or value.is_meta + or value.is_nested + or value.is_quantized + or value.dtype != spec.dtype + or tuple(value.shape) != spec.shape + or value.numel() != spec.numel + or value.element_size() * value.numel() != spec.nbytes + or not value.is_contiguous() + or value.is_conj() + or value.is_neg() + or value.storage_offset() != 0 + ): + return False + try: + return value.untyped_storage().nbytes() == spec.nbytes + except (AttributeError, RuntimeError): + return False + + +def _materialize_value(value, tensors: tuple[torch.Tensor, ...]): + if type(value) is _TensorReference: + return tensors[value.index] + if type(value) is _ListValue: + return [_materialize_value(item, tensors) for item in value.values] + if type(value) is _TupleValue: + return tuple(_materialize_value(item, tensors) for item in value.values) + if type(value) is _DictionaryValue: + return {key: _materialize_value(item, tensors) for key, item in value.values} + return value + + +def _reconstruct_canonical_wire_value( + prepared: _PreparedCanonicalWireValue, + payload_tensors, + *, + expected_fragment_digest: bytes | None = None, +): + """Verify separate tensor payloads and reconstruct one owned canonical value.""" + + if type(prepared) is not _PreparedCanonicalWireValue: + raise TypeError("prepared must be _PreparedCanonicalWireValue") + if expected_fragment_digest is not None: + _validate_digest_field("expected_fragment_digest", expected_fragment_digest) + if not hmac.compare_digest(prepared.fragment_digest, expected_fragment_digest): + raise ValueError("canonical wire fragment digest does not match") + if type(payload_tensors) not in {tuple, list}: + raise TypeError("payload_tensors must be a tuple or list") + if len(payload_tensors) != len(prepared.tensor_specs): + raise ValueError("canonical wire payload count does not match tensor descriptors") + payloads = tuple(payload_tensors) + for index, (payload, spec) in enumerate(zip(payloads, prepared.tensor_specs)): + if not torch.is_tensor(payload) or not _tight_cpu_tensor(payload, spec): + raise ValueError("canonical wire tensor payload {} does not match its descriptor".format(index)) + cloned = [] + for index, (payload, spec) in enumerate(zip(payloads, prepared.tensor_specs)): + owned = _clone_tensor(payload, chunk_bytes=prepared._limits.chunk_bytes, path="payload_tensors[{}]".format(index)) + digest = _canonical_tensor_digest( + owned, + dtype_code=spec.dtype_code, + shape=spec.shape, + nbytes=spec.nbytes, + chunk_bytes=prepared._limits.chunk_bytes, + ) + if not hmac.compare_digest(digest, spec.digest): + raise ValueError("canonical wire tensor payload {} has an invalid digest".format(index)) + cloned.append(owned) + return _materialize_value(prepared._value, tuple(cloned)) diff --git a/src/gefen/rebinding.py b/src/gefen/rebinding.py new file mode 100644 index 0000000..804e147 --- /dev/null +++ b/src/gefen/rebinding.py @@ -0,0 +1,54 @@ +"""Runtime requests for atomic parameter/shard rebinding.""" + +from dataclasses import dataclass +from typing import Optional + +from gefen.contracts import ShardIdentity + + +@dataclass(frozen=True, slots=True) +class LogicalSlotBinding: + """Tensor-free identity for one original optimizer parameter slot.""" + + group_index: int + original_slot_index: int + compatibility_name: str + shard: ShardIdentity + + def __post_init__(self) -> None: + if type(self.group_index) is not int: + raise TypeError("LogicalSlotBinding.group_index must be an integer") + if self.group_index < 0: + raise ValueError("LogicalSlotBinding.group_index must be nonnegative") + if type(self.original_slot_index) is not int: + raise TypeError( + "LogicalSlotBinding.original_slot_index must be an integer" + ) + if self.original_slot_index < 0: + raise ValueError( + "LogicalSlotBinding.original_slot_index must be nonnegative" + ) + if type(self.compatibility_name) is not str: + raise TypeError("LogicalSlotBinding.compatibility_name must be a string") + if self.compatibility_name != self.compatibility_name.lower(): + raise ValueError( + "LogicalSlotBinding.compatibility_name must be lowercase" + ) + if not isinstance(self.shard, ShardIdentity): + raise TypeError("LogicalSlotBinding.shard must be a ShardIdentity") + + +@dataclass(frozen=True, eq=False) +class ParameterRebinding: + """Bind one optimizer slot to a local tensor or prune it as a non-owner.""" + + old_parameter: object + new_parameter: Optional[object] + shard: ShardIdentity + + def __post_init__(self) -> None: + if not isinstance(self.shard, ShardIdentity): + raise TypeError("ParameterRebinding.shard must be a ShardIdentity") + + +__all__ = ["ParameterRebinding"] diff --git a/tests/_state_snapshot.py b/tests/_state_snapshot.py new file mode 100644 index 0000000..7bb0115 --- /dev/null +++ b/tests/_state_snapshot.py @@ -0,0 +1,196 @@ +"""Shared deep fail-before-mutation snapshot helpers for optimizer tests. + +``deep_state_snapshot`` captures both the container identities and bitwise +clones of every tensor reachable from an optimizer (per-parameter state, +codebook tensors and per-device caches, counters, param-group contents). A +regression that publishes staged data by mutating live containers in place — +``state[param].update(...)``, ``copy_()`` into an existing state tensor, an +in-place ``group['params']`` element swap, or an in-place option edit — +preserves every top-level object identity and is only caught by comparing the +live values against these clones. ``assert_deep_state_snapshot`` therefore +checks identity AND bitwise value equality together. + +The staged post_sharding transaction also publishes two per-parameter registry +dicts on each child — ``_param_names`` (parameter -> compatibility FQN string) +and ``_gefen_shard_bindings`` (parameter -> ``ShardIdentity``) — by rebuilding +them on a staged shadow and swapping the shadow in wholesale. Those dicts keep +their live object identity across a transaction, so a regression that wrote +``self._param_names[target] = ...`` on the LIVE child before a later child +raised would leave a half-populated registry that top-level identity checks +miss. ``deep_state_snapshot`` therefore also records the key set (by parameter +identity) and cloned values of these registries, and +``assert_deep_state_snapshot`` value-compares them after the transaction. +""" + +import copy + +import torch + + +# Per-parameter registry dicts each Gefen/GefenMuon child rebuilds on a staged +# shadow during a post_sharding transaction. Keys are live parameters (compared +# by identity); values are FQN strings / ``ShardIdentity`` objects (compared by +# value). Absent on optimizer types that do not stage these (handled below). +_CHILD_REGISTRY_ATTRS = ("_param_names", "_gefen_shard_bindings") + + +def _cloned(value): + if torch.is_tensor(value): + return value.detach().clone() + if type(value) is dict: + return {key: _cloned(item) for key, item in value.items()} + if type(value) in (list, tuple): + return type(value)(_cloned(item) for item in value) + return copy.deepcopy(value) + + +def _tensors_bitwise_equal(live, expected): + # ``torch.equal`` reports +0.0/-0.0 as equal and every NaN as unequal, so it + # can miss a failed transaction that flips a sign bit and can spuriously fail + # on unchanged NaN payloads. A fail-before-mutation snapshot needs exact byte + # preservation, so compare dtype/layout/shape and then the raw bytes. + if ( + live.dtype != expected.dtype + or live.layout != expected.layout + or tuple(live.shape) != tuple(expected.shape) + ): + return False + # ``view(torch.uint8)`` rejects a 0-dim tensor (a scalar cannot hold the + # several uint8 elements one value reinterprets to), which capturable + # device-resident scalar counters hit. Flatten to 1-D first; shapes were + # already checked equal above, so both sides flatten identically. + live_bytes = live.detach().contiguous().reshape(-1).view(torch.uint8) + expected_bytes = expected.detach().contiguous().reshape(-1).view(torch.uint8) + return bool(torch.equal(live_bytes, expected_bytes)) + + +def _nested_equal(live, expected): + if torch.is_tensor(expected): + assert torch.is_tensor(live) + assert _tensors_bitwise_equal(live, expected) + return + assert type(live) is type(expected) + if isinstance(expected, dict): + assert set(live) == set(expected) + for key in expected: + _nested_equal(live[key], expected[key]) + elif isinstance(expected, (list, tuple)): + assert len(live) == len(expected) + for live_item, expected_item in zip(live, expected): + _nested_equal(live_item, expected_item) + else: + assert live == expected + + +def _tensor_value_pairs(optimizer): + """Every tensor reachable from ``optimizer.__dict__`` with a clone of it.""" + + pairs = [] + visited = set() + + def visit(value): + if torch.is_tensor(value): + if id(value) not in visited: + visited.add(id(value)) + pairs.append((value, value.detach().clone())) + elif isinstance(value, dict): + if id(value) in visited: + return + visited.add(id(value)) + for key, item in value.items(): + visit(key) + visit(item) + elif isinstance(value, (list, tuple, set, frozenset)): + if id(value) in visited: + return + visited.add(id(value)) + for item in value: + visit(item) + + for value in optimizer.__dict__.values(): + visit(value) + return tuple(pairs) + + +def _registry_snapshot(optimizer): + """Cloned contents of each per-parameter registry dict that staging swaps. + + Returns ``name -> ((param, cloned_value), ...)`` for every registry that is + present as a ``dict`` on ``optimizer``; registries absent on this optimizer + type are skipped so the helper stays usable across optimizer kinds. Keys are + kept by reference for identity comparison; values are cloned so a later + in-place value mutation cannot alias the snapshot. + """ + + registries = {} + for name in _CHILD_REGISTRY_ATTRS: + registry = getattr(optimizer, name, None) + if type(registry) is not dict: + continue + registries[name] = tuple( + (key, _cloned(value)) for key, value in registry.items() + ) + return registries + + +def deep_state_snapshot(optimizer): + return { + "attributes": optimizer.__dict__.copy(), + "registries": _registry_snapshot(optimizer), + "state": optimizer.state, + "state_items": tuple( + (parameter, state, _cloned(dict(state))) + for parameter, state in optimizer.state.items() + ), + "param_groups": optimizer.param_groups, + "groups": tuple( + ( + group, + group["params"], + tuple(group["params"]), + _cloned({key: value for key, value in group.items() if key != "params"}), + ) + for group in optimizer.param_groups + ), + "tensors": _tensor_value_pairs(optimizer), + } + + +def assert_deep_state_snapshot(optimizer, snapshot): + assert optimizer.__dict__.keys() == snapshot["attributes"].keys() + for name, value in snapshot["attributes"].items(): + assert optimizer.__dict__[name] is value + assert optimizer.state is snapshot["state"] + assert optimizer.param_groups is snapshot["param_groups"] + assert len(optimizer.state) == len(snapshot["state_items"]) + for live_parameter, (parameter, state, expected) in zip( + optimizer.state, snapshot["state_items"] + ): + assert live_parameter is parameter + assert optimizer.state[parameter] is state + _nested_equal(dict(state), expected) + assert len(optimizer.param_groups) == len(snapshot["groups"]) + for live_group, (group, params, expected_params, expected_options) in zip( + optimizer.param_groups, snapshot["groups"] + ): + assert live_group is group + assert group["params"] is params + assert len(params) == len(expected_params) + for live_param, expected_param in zip(params, expected_params): + assert live_param is expected_param + _nested_equal( + {key: value for key, value in group.items() if key != "params"}, + expected_options, + ) + for tensor, expected in snapshot["tensors"]: + assert _tensors_bitwise_equal(tensor, expected) + for name, expected_entries in snapshot["registries"].items(): + live = getattr(optimizer, name, None) + assert type(live) is dict + live_entries = tuple(live.items()) + assert len(live_entries) == len(expected_entries) + for (live_key, live_value), (expected_key, expected_value) in zip( + live_entries, expected_entries + ): + assert live_key is expected_key + _nested_equal(live_value, expected_value) diff --git a/tests/test_canonical_state_cpu.py b/tests/test_canonical_state_cpu.py new file mode 100644 index 0000000..6e98c0c --- /dev/null +++ b/tests/test_canonical_state_cpu.py @@ -0,0 +1,877 @@ +import copy +import io + +import pytest +import torch + +from gefen import ( + CANONICAL_STATE_FORMAT_VERSION, + CanonicalStateProvider, + CheckpointTransport, + Gefen, + GefenMuon, + LogicalSlice, + ParameterIdentity, + ParameterLayout, + ParameterRebinding, + PlacementKind, + PreparedCanonicalStateImport, + ProcessGroupIdentity, + ProcessGroupScope, + ShardIdentity, + ShardPlacement, + ShardingManifest, +) +from gefen.rebinding import LogicalSlotBinding + +from _state_snapshot import assert_deep_state_snapshot, deep_state_snapshot + + +def _replicated(identity, group=None, member=None): + placements = () + if group is not None: + coordinate = group.ordered_members.index(member) + placements = ( + ShardPlacement( + "dp", + PlacementKind.REPLICATE, + coordinate, + len(group.ordered_members), + ), + ) + return ShardIdentity( + identity, + ParameterLayout.REPLICATED, + LogicalSlice.full(identity), + process_group=group, + local_member=member, + placements=placements, + ) + + +def _flat(identity, group, member, offset, length): + coordinate = group.ordered_members.index(member) + return ShardIdentity( + identity, + ParameterLayout.FLATTENED_ELEMENT_SHARD, + LogicalSlice(offset, length), + process_group=group, + local_member=member, + placements=( + ShardPlacement( + "dp", + PlacementKind.FLAT_SHARD, + coordinate, + len(group.ordered_members), + ), + ), + ) + + +def _finalize(optimizer, bindings, manifest): + optimizer.post_sharding( + tuple(ParameterRebinding(parameter, parameter, shard) for parameter, shard in bindings), + manifest=manifest, + ) + + +def _snapshot(optimizer): + # Deep snapshot: container identities plus bitwise clones of every state + # tensor, codebook tensor, counter, and param-group entry, so a rejected + # import that mutated live containers in place cannot pass unnoticed. + return { + "deep": deep_state_snapshot(optimizer), + "codebook": optimizer._gefen_codebook, + "global_step": optimizer._gefen_global_step, + } + + +def _assert_snapshot_identity(optimizer, snapshot): + assert optimizer._gefen_codebook is snapshot["codebook"] + assert optimizer._gefen_global_step is snapshot["global_step"] + assert_deep_state_snapshot(optimizer, snapshot["deep"]) + + +def _two_parameter_source(): + first = torch.nn.Parameter(torch.arange(1, 5, dtype=torch.float32)) + second = torch.nn.Parameter(torch.arange(11, 15, dtype=torch.float32)) + optimizer = Gefen( + [("first", first), ("second", second)], + fused=False, + factored_v_2d=False, + ) + first_identity = ParameterIdentity("Model.First", (4,)) + second_identity = ParameterIdentity("Model.Second", (4,)) + first_shard = _replicated(first_identity) + second_shard = _replicated(second_identity) + manifest = ShardingManifest((first_shard, second_shard)) + _finalize( + optimizer, + ((first, first_shard), (second, second_shard)), + manifest, + ) + optimizer._resolve_automatic_period = lambda *args: 4 + first.grad = torch.tensor([1.0, 2.0, 3.0, 4.0]) + second.grad = torch.tensor([-8.0, 2.0, 1.0, 3.0]) + optimizer.step() + return optimizer, first, second, first_shard, second_shard, manifest + + +def _reordered_target(first_value, second_value, first_shard, second_shard, manifest): + second = torch.nn.Parameter(second_value.detach().clone()) + first = torch.nn.Parameter(first_value.detach().clone()) + optimizer = Gefen( + [ + {"params": [("second_target", second)]}, + {"params": [("first_target", first)]}, + ], + fused=False, + factored_v_2d=False, + ) + _finalize( + optimizer, + ((second, second_shard), (first, first_shard)), + manifest, + ) + return optimizer, first, second + + +def test_canonical_local_capability_is_dynamic_and_exact_binding_only(): + parameter = torch.nn.Parameter(torch.ones(4)) + optimizer = Gefen([("p", parameter)], fused=False, factored_v_2d=False) + assert isinstance(optimizer, CanonicalStateProvider) + assert not optimizer.optimizer_contract().capabilities.canonical_state_io + assert all( + support.transport is not CheckpointTransport.CANONICAL_LOCAL + for support in optimizer.optimizer_contract().capabilities.checkpoints + ) + + identity = ParameterIdentity("Model.P", (4,)) + shard = _replicated(identity) + _finalize(optimizer, ((parameter, shard),), ShardingManifest((shard,))) + contract = optimizer.optimizer_contract() + support = next( + item for item in contract.capabilities.checkpoints if item.transport is CheckpointTransport.CANONICAL_LOCAL + ) + assert contract.capabilities.canonical_state_io + assert support.same_topology == frozenset({ParameterLayout.REPLICATED}) + assert not support.topology_changing + assert support.process_group_scope is ProcessGroupScope.NONE + assert support.atomic_load + assert not support.requires_collective + + +def test_pristine_export_is_primitive_device_neutral_and_excludes_derived_state(): + parameter = torch.nn.Parameter(torch.ones(8)) + optimizer = Gefen([("p", parameter)], fused=False, factored_v_2d=False) + identity = ParameterIdentity("Model.P", (8,)) + shard = _replicated(identity) + _finalize(optimizer, ((parameter, shard),), ShardingManifest((shard,))) + optimizer.state[parameter]["stepsize"] = torch.ones(1) + + exported = optimizer.export_canonical_state() + + assert exported["format"] == "gefen.bound_state" + assert exported["format_version"] == CANONICAL_STATE_FORMAT_VERSION + assert exported["coverage"] == "local_optimizer_fragment" + assert exported["manifest"][0]["process_group"] is None + assert set(exported["parameters"]) == {"Model.P"} + record = exported["parameters"]["Model.P"] + assert record["state"] == {} + assert "stepsize" not in record["state"] + assert "name" not in record["state"] + assert record["compatibility_name"] == "p" + + buffer = io.BytesIO() + torch.save(exported, buffer) + buffer.seek(0) + loaded = torch.load(buffer, weights_only=True) + assert loaded["format_version"] == CANONICAL_STATE_FORMAT_VERSION + assert loaded["parameters"]["Model.P"]["state"] == {} + + +def test_pristine_continuation_routes_period_policy_by_canonical_fqn(): + source_parameter = torch.nn.Parameter(torch.arange(1, 9, dtype=torch.float32)) + source = Gefen( + [("special_weight", source_parameter)], + fused=False, + factored_v_2d=False, + period_one_substrings=("special",), + ) + identity = ParameterIdentity("Model.SpecialWeight", (8,)) + shard = _replicated(identity) + manifest = ShardingManifest((shard,)) + _finalize(source, ((source_parameter, shard),), manifest) + + target_parameter = torch.nn.Parameter(source_parameter.detach().clone()) + target = Gefen( + [("unrelated_target_name", target_parameter)], + fused=False, + factored_v_2d=False, + period_one_substrings=("special",), + ) + _finalize(target, ((target_parameter, shard),), manifest) + target.import_canonical_state(source.export_canonical_state()) + + gradient = torch.arange(1, 9, dtype=torch.float32) + source_parameter.grad = gradient.clone() + target_parameter.grad = gradient.clone() + source.step() + target.step() + assert source.state[source_parameter]["automatic_period"] == 1 + assert target.state[target_parameter]["automatic_period"] == 1 + assert torch.equal(target_parameter, source_parameter) + assert torch.equal(target._gefen_codebook, source._gefen_codebook) + + +def test_informational_empty_compatibility_name_round_trips_but_stochastic_policy_is_unclaimed(): + parameter = torch.nn.Parameter(torch.ones(4)) + optimizer = Gefen([("", parameter)], fused=False, factored_v_2d=False) + identity = ParameterIdentity("Model.P", (4,)) + shard = _replicated(identity) + manifest = ShardingManifest((shard,)) + _finalize(optimizer, ((parameter, shard),), manifest) + exported = optimizer.export_canonical_state() + assert exported["parameters"]["Model.P"]["compatibility_name"] == "" + optimizer.import_canonical_state(exported) + assert optimizer.state[parameter]["name"] == "" + + stochastic_parameter = torch.nn.Parameter(torch.ones(4)) + with pytest.warns(RuntimeWarning, match="stochastic_round=True"): + stochastic = Gefen( + [("p", stochastic_parameter)], + fused=False, + factored_v_2d=False, + stochastic_round=True, + ) + _finalize(stochastic, ((stochastic_parameter, shard),), manifest) + assert not stochastic.optimizer_contract().capabilities.canonical_state_io + with pytest.raises(RuntimeError, match="supported finalized"): + stochastic.export_canonical_state() + + +def test_initialized_import_maps_by_fqn_across_order_and_group_boundaries(): + source, source_first, source_second, first_shard, second_shard, manifest = _two_parameter_source() + exported = source.export_canonical_state() + exported_magnitude = exported["parameters"]["Model.First"]["state"][ + "m_magnitude" + ] + assert exported_magnitude is not source.state[source_first]["m_magnitude"] + assert ( + exported_magnitude.untyped_storage().nbytes() + == exported_magnitude.numel() * exported_magnitude.element_size() + ) + target, target_first, target_second = _reordered_target( + source_first, + source_second, + first_shard, + second_shard, + manifest, + ) + load_hooks = [] + target.register_load_state_dict_pre_hook(lambda *args: load_hooks.append("pre")) + target.register_load_state_dict_post_hook(lambda *args: load_hooks.append("post")) + input_codebook = exported["common"]["gefen_codebook"] + input_magnitude = exported["parameters"]["Model.First"]["state"]["m_magnitude"] + input_codebook_value = input_codebook.clone() + input_magnitude_value = input_magnitude.clone() + + prepared = target.prepare_canonical_state_import(exported) + logical_slots = target._gefen_logical_slots + assert isinstance(prepared, PreparedCanonicalStateImport) + assert target._gefen_codebook is None + assert exported["common"]["gefen_codebook"] is input_codebook + assert exported["parameters"]["Model.First"]["state"]["m_magnitude"] is input_magnitude + assert torch.equal(input_codebook, input_codebook_value) + assert torch.equal(input_magnitude, input_magnitude_value) + target.commit_canonical_state_import(prepared) + + assert not load_hooks + assert target._gefen_logical_slots is logical_slots + assert torch.equal( + target.state[target_first]["m_magnitude"], + source.state[source_first]["m_magnitude"], + ) + assert torch.equal( + target.state[target_second]["m_magnitude"], + source.state[source_second]["m_magnitude"], + ) + assert target.state[target_first]["name"] == "first_target" + assert target.state[target_second]["name"] == "second_target" + + first_grad = torch.tensor([4.0, -3.0, 2.0, -1.0]) + second_grad = torch.tensor([1.0, 3.0, -5.0, 7.0]) + source_first.grad = first_grad.clone() + target_first.grad = first_grad.clone() + source_second.grad = second_grad.clone() + target_second.grad = second_grad.clone() + source.step() + target.step() + assert torch.equal(target_first, source_first) + assert torch.equal(target_second, source_second) + assert torch.equal(target._gefen_codebook, source._gefen_codebook) + + with pytest.raises(RuntimeError, match="already consumed"): + target.commit_canonical_state_import(prepared) + + +def test_prepared_import_rejects_wrong_optimizer_and_stale_live_state(): + source, source_first, source_second, first_shard, second_shard, manifest = _two_parameter_source() + exported = source.export_canonical_state() + target, _, _ = _reordered_target( + source_first, + source_second, + first_shard, + second_shard, + manifest, + ) + other, _, _ = _reordered_target( + source_first, + source_second, + first_shard, + second_shard, + manifest, + ) + prepared = target.prepare_canonical_state_import(exported) + + other_before = _snapshot(other) + with pytest.raises(ValueError, match="another optimizer"): + other.commit_canonical_state_import(prepared) + _assert_snapshot_identity(other, other_before) + + target._gefen_global_step += 1 + target_before = _snapshot(target) + with pytest.raises(RuntimeError, match="changed after"): + target.commit_canonical_state_import(prepared) + _assert_snapshot_identity(target, target_before) + + +@pytest.mark.parametrize("mutation", ["registry_identity", "registry_structure"]) +def test_prepared_import_freshness_covers_logical_slot_registry(mutation): + parameter = torch.nn.Parameter(torch.ones(4)) + optimizer = Gefen( + [("p", parameter)], + fused=False, + factored_v_2d=False, + ) + identity = ParameterIdentity("Model.P", (4,)) + shard = _replicated(identity) + _finalize(optimizer, ((parameter, shard),), ShardingManifest((shard,))) + prepared = optimizer.prepare_canonical_state_import( + optimizer.export_canonical_state() + ) + slot = optimizer._gefen_logical_slots[0] + if mutation == "registry_identity": + optimizer._gefen_logical_slots = ( + LogicalSlotBinding( + slot.group_index, + slot.original_slot_index, + slot.compatibility_name, + slot.shard, + ), + ) + else: + object.__setattr__( + slot.shard, + "parameter", + ParameterIdentity("Model.Changed", (4,)), + ) + assert optimizer._finalized_binding_layout_matches() + before = _snapshot(optimizer) + + with pytest.raises(RuntimeError, match="changed after"): + optimizer.commit_canonical_state_import(prepared) + _assert_snapshot_identity(optimizer, before) + + +@pytest.mark.parametrize( + "corrupt", + [ + "missing_parameter", + "wrong_manifest", + "wrong_shard", + "wrong_policy", + "wrong_group_options", + "unknown_state", + "bad_state_geometry", + "counter_ahead", + "secondary_counter_ahead", + "bool_version", + ], +) +def test_canonical_corruption_rejects_before_live_mutation(corrupt): + source, source_first, source_second, first_shard, second_shard, manifest = _two_parameter_source() + exported = source.export_canonical_state() + target, _, _ = _reordered_target( + source_first, + source_second, + first_shard, + second_shard, + manifest, + ) + damaged = copy.deepcopy(exported) + if corrupt == "missing_parameter": + damaged["parameters"].pop("Model.Second") + elif corrupt == "wrong_manifest": + damaged["manifest"][0]["parameter"]["fqn"] = "Wrong.First" + elif corrupt == "wrong_shard": + damaged["parameters"]["Model.First"]["shard"] = copy.deepcopy(damaged["parameters"]["Model.Second"]["shard"]) + elif corrupt == "wrong_policy": + damaged["policy"]["stochastic_round"] = True + elif corrupt == "wrong_group_options": + damaged["parameters"]["Model.First"]["group_options"]["lr"] = 9.0 + elif corrupt == "unknown_state": + damaged["parameters"]["Model.First"]["state"]["exp_avg"] = torch.ones(4) + elif corrupt == "bad_state_geometry": + damaged["parameters"]["Model.First"]["state"]["m_magnitude"] = torch.ones(2, 1) + elif corrupt == "counter_ahead": + damaged["parameters"]["Model.First"]["state"]["step"] = ( + damaged["common"]["gefen_global_step"] + 1 + ) + elif corrupt == "secondary_counter_ahead": + damaged["common"]["gefen_global_step"] = 100 + damaged["parameters"]["Model.First"]["state"]["step"] = 1 + damaged["parameters"]["Model.First"]["state"]["vmean_step"] = 2 + else: + damaged["format_version"] = True + before = _snapshot(target) + + with pytest.raises((TypeError, ValueError, RuntimeError)): + target.import_canonical_state(damaged) + + _assert_snapshot_identity(target, before) + + +def test_import_rejects_initialized_period_that_violates_forced_policy(): + source_parameter = torch.nn.Parameter(torch.arange(8, dtype=torch.float32)) + source = Gefen( + [("p", source_parameter)], + fused=False, + factored_v_2d=False, + force_1d_period_one=True, + ) + identity = ParameterIdentity("Model.P", (8,)) + shard = _replicated(identity) + manifest = ShardingManifest((shard,)) + _finalize(source, ((source_parameter, shard),), manifest) + source_parameter.grad = torch.arange(1, 9, dtype=torch.float32) + source.step() + damaged = source.export_canonical_state() + parameter_state = damaged["parameters"]["Model.P"]["state"] + parameter_state["automatic_period"] = 2 + parameter_state["m_codebook"] = torch.zeros(4, 2, dtype=torch.uint8) + parameter_state["m_magnitude"] = torch.ones(4, 1, dtype=torch.float32) + parameter_state["vmean"] = torch.ones(4, 1, dtype=torch.float32) + + target_parameter = torch.nn.Parameter(source_parameter.detach().clone()) + target = Gefen( + [("target", target_parameter)], + fused=False, + factored_v_2d=False, + force_1d_period_one=True, + ) + _finalize(target, ((target_parameter, shard),), manifest) + before = _snapshot(target) + with pytest.raises(ValueError, match="violates period-one policy"): + target.import_canonical_state(damaged) + _assert_snapshot_identity(target, before) + + +def test_codebook_initialized_before_first_step_round_trips_canonically(): + parameter = torch.nn.Parameter(torch.arange(8, dtype=torch.float32)) + source = Gefen([("p", parameter)], fused=False, factored_v_2d=False) + identity = ParameterIdentity("Model.P", (8,)) + shard = _replicated(identity) + manifest = ShardingManifest((shard,)) + _finalize(source, ((parameter, shard),), manifest) + source._resolve_automatic_period = lambda *args: 4 + parameter.grad = torch.arange(1, 9, dtype=torch.float32) + assert source.initialize_codebook() + assert set(source.state[parameter]) == {"name", "automatic_period"} + + target_parameter = torch.nn.Parameter(parameter.detach().clone()) + target = Gefen([("target", target_parameter)], fused=False, factored_v_2d=False) + _finalize(target, ((target_parameter, shard),), manifest) + target.import_canonical_state(source.export_canonical_state()) + + assert target._gefen_global_step == 0 + assert torch.equal(target._gefen_codebook, source._gefen_codebook) + assert target.state[target_parameter] == { + "name": "target", + "automatic_period": 4, + } + + +def test_flat_fragment_is_fqn_keyed_but_rejects_another_member_slice(): + group = ProcessGroupIdentity("flat", ("rank:0", "rank:1")) + identity = ParameterIdentity("Model.Flat", (8,)) + shards = ( + _flat(identity, group, "rank:0", 0, 4), + _flat(identity, group, "rank:1", 4, 4), + ) + manifest = ShardingManifest(shards) + source_parameter = torch.nn.Parameter(torch.arange(4, dtype=torch.float32)) + source = Gefen([("flat", source_parameter)], fused=False, factored_v_2d=False) + _finalize(source, ((source_parameter, shards[0]),), manifest) + source._resolve_automatic_period = lambda *args: 4 + source_parameter.grad = torch.tensor([1.0, 2.0, 3.0, 4.0]) + source.step() + exported = source.export_canonical_state() + support = next( + item + for item in source.optimizer_contract().capabilities.checkpoints + if item.transport is CheckpointTransport.CANONICAL_LOCAL + ) + assert support.same_topology == frozenset({ParameterLayout.FLATTENED_ELEMENT_SHARD}) + + target_parameter = torch.nn.Parameter(torch.arange(4, dtype=torch.float32)) + target = Gefen([("flat", target_parameter)], fused=False, factored_v_2d=False) + _finalize(target, ((target_parameter, shards[1]),), manifest) + before = _snapshot(target) + with pytest.raises(ValueError, match="shard"): + target.import_canonical_state(exported) + _assert_snapshot_identity(target, before) + + +def test_muon_replicated_supports_canonical_state_but_whole_owner_does_not(): + parameter = torch.nn.Parameter(torch.ones(2, 2)) + optimizer = GefenMuon([("matrix", parameter)], fused=False) + identity = ParameterIdentity("Model.Matrix", (2, 2)) + shard = _replicated(identity) + _finalize(optimizer, ((parameter, shard),), ShardingManifest((shard,))) + assert optimizer.optimizer_contract().capabilities.canonical_state_io + + group = ProcessGroupIdentity("owner", ("rank:0",)) + owner_shard = ShardIdentity( + identity, + ParameterLayout.WHOLE_PARAMETER_OWNER, + LogicalSlice.full(identity), + process_group=group, + local_member="rank:0", + owner="rank:0", + placements=( + ShardPlacement( + "dp", + PlacementKind.WHOLE_PARAMETER_OWNER, + 0, + 1, + ), + ), + ) + owner_parameter = torch.nn.Parameter(torch.ones(2, 2)) + owner = GefenMuon([("matrix", owner_parameter)], fused=False) + _finalize(owner, ((owner_parameter, owner_shard),), ShardingManifest((owner_shard,))) + assert not owner.optimizer_contract().capabilities.canonical_state_io + with pytest.raises(RuntimeError, match="supported finalized"): + owner.export_canonical_state() + + +def test_initialized_normuon_requires_its_authoritative_pair_and_continues_exactly(): + source_parameter = torch.nn.Parameter( + torch.tensor([[1.0, -2.0], [3.0, -4.0]]) + ) + source = GefenMuon( + [("matrix", source_parameter)], + fused=False, + normuon=True, + ) + identity = ParameterIdentity("Model.Matrix", (2, 2)) + shard = _replicated(identity) + manifest = ShardingManifest((shard,)) + _finalize(source, ((source_parameter, shard),), manifest) + source._resolve_automatic_period = lambda *args: 4 + source_parameter.grad = torch.tensor([[0.5, -1.0], [1.5, -2.0]]) + source.step() + exported = source.export_canonical_state() + assert {"normuon_v", "normuon_step"}.issubset( + exported["parameters"]["Model.Matrix"]["state"] + ) + + target_parameter = torch.nn.Parameter(source_parameter.detach().clone()) + target = GefenMuon( + [("target", target_parameter)], + fused=False, + normuon=True, + ) + _finalize(target, ((target_parameter, shard),), manifest) + damaged = copy.deepcopy(exported) + damaged_state = damaged["parameters"]["Model.Matrix"]["state"] + damaged_state.pop("normuon_v") + damaged_state.pop("normuon_step") + before = _snapshot(target) + with pytest.raises(ValueError, match="NorMuon state is incomplete"): + target.import_canonical_state(damaged) + _assert_snapshot_identity(target, before) + + target.import_canonical_state(exported) + next_grad = torch.tensor([[-2.0, 0.25], [0.75, -1.25]]) + source_parameter.grad = next_grad.clone() + target_parameter.grad = next_grad.clone() + source.step() + target.step() + assert torch.equal(target_parameter, source_parameter) + assert torch.equal( + target.state[target_parameter]["normuon_v"], + source.state[source_parameter]["normuon_v"], + ) + + +@pytest.mark.parametrize("foreign_state", ["block", "factored"]) +def test_muon_rejects_plain_gefen_authoritative_state_variants(foreign_state): + source_parameter = torch.nn.Parameter(torch.ones(2, 2)) + source = GefenMuon([("matrix", source_parameter)], fused=False) + identity = ParameterIdentity("Model.Matrix", (2, 2)) + shard = _replicated(identity) + manifest = ShardingManifest((shard,)) + _finalize(source, ((source_parameter, shard),), manifest) + source._resolve_automatic_period = lambda *args: 4 + source_parameter.grad = torch.tensor([[1.0, -2.0], [3.0, -4.0]]) + source.step() + damaged = source.export_canonical_state() + parameter_state = damaged["parameters"]["Model.Matrix"]["state"] + if foreign_state == "block": + parameter_state["vmean"] = torch.zeros_like( + parameter_state["m_magnitude"] + ) + parameter_state["vmean_step"] = 1 + else: + parameter_state["v_row"] = torch.zeros(2, dtype=torch.float32) + parameter_state["v_col"] = torch.zeros(2, dtype=torch.float32) + parameter_state["factored_step"] = 1 + + target_parameter = torch.nn.Parameter(source_parameter.detach().clone()) + target = GefenMuon([("target", target_parameter)], fused=False) + _finalize(target, ((target_parameter, shard),), manifest) + before = _snapshot(target) + with pytest.raises(ValueError, match="declared state variant"): + target.import_canonical_state(damaged) + _assert_snapshot_identity(target, before) + + +def test_unsupported_custom_state_or_callable_policy_disables_canonical_claim(): + parameter = torch.nn.Parameter(torch.ones(2, 2)) + optimizer = GefenMuon( + [("matrix", parameter)], + fused=False, + ) + identity = ParameterIdentity("Model.Matrix", (2, 2)) + shard = _replicated(identity) + _finalize(optimizer, ((parameter, shard),), ShardingManifest((shard,))) + optimizer.param_groups[0]["custom_callable"] = lambda: None + assert not optimizer.optimizer_contract().capabilities.canonical_state_io + + plain_parameter = torch.nn.Parameter(torch.ones(4)) + plain = Gefen([("p", plain_parameter)], fused=False) + plain_identity = ParameterIdentity("Model.P", (4,)) + plain_shard = _replicated(plain_identity) + _finalize( + plain, + ((plain_parameter, plain_shard),), + ShardingManifest((plain_shard,)), + ) + plain.state[plain_parameter]["custom_tensor"] = torch.ones(1) + assert not plain.optimizer_contract().capabilities.canonical_state_io + with pytest.raises(RuntimeError, match="primitive algorithm policy"): + plain.export_canonical_state() + + class TensorSubclass(torch.Tensor): + pass + + subclass_parameter = torch.nn.Parameter(torch.ones(4)) + subclassed = Gefen([("p", subclass_parameter)], fused=False) + subclass_identity = ParameterIdentity("Model.Subclass", (4,)) + subclass_shard = _replicated(subclass_identity) + _finalize( + subclassed, + ((subclass_parameter, subclass_shard),), + ShardingManifest((subclass_shard,)), + ) + subclassed.param_groups[0]["custom_tensor"] = torch.ones(1).as_subclass( + TensorSubclass + ) + assert not subclassed.optimizer_contract().capabilities.canonical_state_io + with pytest.raises(RuntimeError, match="primitive algorithm policy"): + subclassed.export_canonical_state() + + subclassed.param_groups[0].pop("custom_tensor") + with torch.no_grad(): + nested = torch.nested.nested_tensor( + [torch.ones(2), torch.ones(3)] + ) + subclassed.param_groups[0]["custom_nested"] = nested + assert not subclassed.optimizer_contract().capabilities.canonical_state_io + with pytest.raises(RuntimeError, match="primitive algorithm policy"): + subclassed.export_canonical_state() + + subclassed.param_groups[0].pop("custom_nested") + subclassed.param_groups[0]["custom_nan"] = torch.tensor([float("nan")]) + assert not subclassed.optimizer_contract().capabilities.canonical_state_io + with pytest.raises(RuntimeError, match="primitive algorithm policy"): + subclassed.export_canonical_state() + + subclassed.param_groups[0].pop("custom_nan") + if hasattr(torch, "float8_e4m3fn"): + subclassed.param_groups[0]["custom_float8"] = torch.ones( + 2, dtype=torch.float8_e4m3fn + ) + assert not subclassed.optimizer_contract().capabilities.canonical_state_io + with pytest.raises(RuntimeError, match="primitive algorithm policy"): + subclassed.export_canonical_state() + subclassed.param_groups[0].pop("custom_float8") + + subclassed.param_groups[0]["custom_conjugate"] = torch.tensor( + [1.0 + 2.0j] + ).conj() + assert subclassed.optimizer_contract().capabilities.canonical_state_io + prepared = subclassed.prepare_canonical_state_import( + subclassed.export_canonical_state() + ) + subclassed.commit_canonical_state_import(prepared) + + +def test_inference_mode_tensor_lr_has_a_stable_canonical_freshness_token(): + with torch.inference_mode(): + learning_rate = torch.tensor(1e-3) + parameter = torch.nn.Parameter(torch.ones(4)) + optimizer = Gefen( + [("p", parameter)], + lr=learning_rate, + fused=False, + factored_v_2d=False, + ) + identity = ParameterIdentity("Model.P", (4,)) + shard = _replicated(identity) + _finalize(optimizer, ((parameter, shard),), ShardingManifest((shard,))) + exported = optimizer.export_canonical_state() + groups = optimizer.param_groups + group = optimizer.param_groups[0] + + prepared = optimizer.prepare_canonical_state_import(exported) + optimizer.commit_canonical_state_import(prepared) + assert optimizer.param_groups is groups + assert optimizer.param_groups[0] is group + assert optimizer.param_groups[0]["lr"] is learning_rate + assert optimizer.defaults["lr"] is learning_rate + + +def test_prepared_import_detects_storage_level_live_state_mutation(): + optimizer, _, _, _, _, _ = _two_parameter_source() + prepared = optimizer.prepare_canonical_state_import( + optimizer.export_canonical_state() + ) + magnitude = next( + state["m_magnitude"] + for state in optimizer.state.values() + if "m_magnitude" in state + ) + version = magnitude._version + magnitude.numpy()[0, 0] += 7.0 + assert magnitude._version == version + before = _snapshot(optimizer) + + with pytest.raises(RuntimeError, match="changed after"): + optimizer.commit_canonical_state_import(prepared) + _assert_snapshot_identity(optimizer, before) + + +def test_nonfinite_authoritative_state_disables_the_export_capability(): + optimizer, first, _, _, _, _ = _two_parameter_source() + optimizer.state[first]["vmean"].fill_(float("inf")) + + assert not optimizer.optimizer_contract().capabilities.canonical_state_io + with pytest.raises(RuntimeError, match="primitive algorithm policy"): + optimizer.export_canonical_state() + + +def test_invalid_live_counter_relationship_disables_the_export_capability(): + optimizer, first, _, _, _, _ = _two_parameter_source() + optimizer.state[first]["step"] = optimizer._gefen_global_step + 1 + + assert not optimizer.optimizer_contract().capabilities.canonical_state_io + with pytest.raises(RuntimeError, match="primitive algorithm policy"): + optimizer.export_canonical_state() + + +def test_canonical_policy_is_constructor_validated_and_capability_checked(): + parameter = torch.nn.Parameter(torch.ones(4)) + with pytest.raises(TypeError, match="codebook_refresh_every must be an integer"): + Gefen( + [("p", parameter)], + fused=False, + codebook_refresh_every=float("nan"), + ) + + optimizer = Gefen([("p", parameter)], fused=False) + identity = ParameterIdentity("Model.P", (4,)) + shard = _replicated(identity) + _finalize(optimizer, ((parameter, shard),), ShardingManifest((shard,))) + optimizer._codebook_refresh_every = float("nan") + assert not optimizer.optimizer_contract().capabilities.canonical_state_io + with pytest.raises(RuntimeError, match="primitive algorithm policy"): + optimizer.export_canonical_state() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +def test_capturable_import_requires_a_fresh_target_before_graph_state_exists(): + parameter = torch.nn.Parameter(torch.arange(8, dtype=torch.float32, device="cuda")) + optimizer = Gefen( + [("p", parameter)], + fused=False, + factored_v_2d=False, + capturable=True, + ) + identity = ParameterIdentity("Model.P", (8,)) + shard = _replicated(identity) + _finalize(optimizer, ((parameter, shard),), ShardingManifest((shard,))) + optimizer._resolve_automatic_period = lambda *args: 4 + parameter.grad = torch.arange(1, 9, dtype=torch.float32, device="cuda") + optimizer.step() + + graph = torch.cuda.CUDAGraph() + torch.cuda.synchronize() + with torch.cuda.graph(graph): + optimizer.step() + exported = optimizer.export_canonical_state() + before = _snapshot(optimizer) + with pytest.raises(RuntimeError, match="requires a fresh target"): + optimizer.prepare_canonical_state_import(exported) + _assert_snapshot_identity(optimizer, before) + + fresh_parameter = torch.nn.Parameter(parameter.detach().clone()) + fresh = Gefen( + [("p", fresh_parameter)], + fused=False, + factored_v_2d=False, + capturable=True, + ) + _finalize(fresh, ((fresh_parameter, shard),), ShardingManifest((shard,))) + fresh.import_canonical_state(exported) + assert fresh._device_gefen_global_step() == exported["common"][ + "gefen_global_step" + ] + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +def test_canonical_cpu_fragment_imports_state_to_cuda_parameter_devices(): + source_parameter = torch.nn.Parameter(torch.arange(8, dtype=torch.float32)) + source = Gefen([("p", source_parameter)], fused=False, factored_v_2d=False) + identity = ParameterIdentity("Model.P", (8,)) + shard = _replicated(identity) + manifest = ShardingManifest((shard,)) + _finalize(source, ((source_parameter, shard),), manifest) + source._resolve_automatic_period = lambda *args: 4 + source_parameter.grad = torch.arange(1, 9, dtype=torch.float32) + source.step() + exported = source.export_canonical_state() + assert all( + not torch.is_tensor(value) or value.device.type == "cpu" + for value in exported["parameters"]["Model.P"]["state"].values() + ) + + target_parameter = torch.nn.Parameter(source_parameter.detach().cuda()) + target = Gefen([("p", target_parameter)], fused=False, factored_v_2d=False) + _finalize(target, ((target_parameter, shard),), manifest) + target.import_canonical_state(exported) + assert all( + not torch.is_tensor(value) or value.device.type == "cuda" for value in target.state[target_parameter].values() + ) diff --git a/tests/test_checkpoint_binding.py b/tests/test_checkpoint_binding.py new file mode 100644 index 0000000..e1c0e49 --- /dev/null +++ b/tests/test_checkpoint_binding.py @@ -0,0 +1,227 @@ +from dataclasses import FrozenInstanceError, fields +from datetime import timedelta +import multiprocessing as mp +import os +import tempfile + +import pytest +import torch +import torch.distributed as dist + +from gefen.checkpoint import CheckpointProcessGroupBinding +from gefen.contracts import ProcessGroupIdentity + + +def test_checkpoint_binding_is_frozen_rank_neutral_and_deterministically_ordered(): + identity = ProcessGroupIdentity("pipeline:1/data", ("worker:b",)) + binding = CheckpointProcessGroupBinding(identity, "worker:b", None, torch.device("cpu")) + + assert binding.identity is identity + assert binding.local_member == "worker:b" + assert binding.process_group is None + assert binding.collective_device == torch.device("cpu") + assert binding.sort_key == ("pipeline:1/data", ("worker:b",)) + assert tuple(field.name for field in fields(binding)) == ( + "identity", + "local_member", + "process_group", + "collective_device", + ) + with pytest.raises(FrozenInstanceError): + binding.local_member = "worker:a" + + +def test_checkpoint_binding_validates_descriptor_types_membership_and_device(): + identity = ProcessGroupIdentity("local", ("worker:0",)) + + with pytest.raises(TypeError, match="ProcessGroupIdentity"): + CheckpointProcessGroupBinding(object(), "worker:0", None, torch.device("cpu")) + with pytest.raises(TypeError, match="local_member"): + CheckpointProcessGroupBinding(identity, 0, None, torch.device("cpu")) + with pytest.raises(ValueError, match="belong"): + CheckpointProcessGroupBinding(identity, "worker:1", None, torch.device("cpu")) + with pytest.raises(TypeError, match="torch.device"): + CheckpointProcessGroupBinding(identity, "worker:0", None, "cpu") + with pytest.raises(ValueError, match="CPU or CUDA"): + CheckpointProcessGroupBinding(identity, "worker:0", None, torch.device("meta")) + + +def test_checkpoint_binding_requires_explicit_multi_member_handle_and_local_singleton(): + singleton = ProcessGroupIdentity("local", ("worker:0",)) + multiple = ProcessGroupIdentity("data", ("worker:0", "worker:1")) + + with pytest.raises(ValueError, match="one-member"): + CheckpointProcessGroupBinding(singleton, "worker:0", object(), torch.device("cpu")) + with pytest.raises(ValueError, match="explicit"): + CheckpointProcessGroupBinding(multiple, "worker:0", None, torch.device("cpu")) + with pytest.raises(TypeError, match="ProcessGroup"): + CheckpointProcessGroupBinding(multiple, "worker:0", object(), torch.device("cpu")) + + +def _expect_rejection(callable_object, exception_type, message): + try: + callable_object() + except exception_type as exc: + return message in str(exc) + return False + + +def _distributed_binding_worker(rank, world_size, init_file, queue): + try: + dist.init_process_group( + "gloo", + init_method="file://{}".format(init_file), + rank=rank, + world_size=world_size, + timeout=timedelta(seconds=45), + ) + world_members = tuple("world:{}".format(index) for index in range(world_size)) + world_identity = ProcessGroupIdentity("world", world_members) + world_binding = CheckpointProcessGroupBinding( + world_identity, + world_members[rank], + dist.group.WORLD, + torch.device("cpu"), + ) + world_binding.validate_runtime() + world_validated = world_binding.process_group is dist.group.WORLD + + wrong_member_binding = CheckpointProcessGroupBinding( + world_identity, + world_members[(rank + 1) % world_size], + dist.group.WORLD, + torch.device("cpu"), + ) + order_mismatch_rejected = _expect_rejection( + wrong_member_binding.validate_runtime, + ValueError, + "ordered semantic members", + ) + + short_identity = ProcessGroupIdentity("short", ("short:0", "short:1")) + short_binding = CheckpointProcessGroupBinding( + short_identity, + short_identity.ordered_members[rank % 2], + dist.group.WORLD, + torch.device("cpu"), + ) + size_mismatch_rejected = _expect_rejection( + short_binding.validate_runtime, + ValueError, + "world size", + ) + + wrong_device_binding = CheckpointProcessGroupBinding( + world_identity, + world_members[rank], + dist.group.WORLD, + torch.device("cuda:0"), + ) + device_mismatch_rejected = _expect_rejection( + wrong_device_binding.validate_runtime, + ValueError, + "runtime backend", + ) + + subgroup_global_ranks = (0, 2) + subgroup = dist.new_group(ranks=list(subgroup_global_ranks), backend="gloo") + subgroup_members = ("subgroup:left", "subgroup:right") + subgroup_validated = False + nonmember_rejected = False + if rank in subgroup_global_ranks: + coordinate = subgroup_global_ranks.index(rank) + subgroup_binding = CheckpointProcessGroupBinding( + ProcessGroupIdentity("subgroup", subgroup_members), + subgroup_members[coordinate], + subgroup, + torch.device("cpu"), + ) + subgroup_binding.validate_runtime() + subgroup_validated = True + else: + nonmember_rejected = _expect_rejection( + lambda: CheckpointProcessGroupBinding( + ProcessGroupIdentity("subgroup", subgroup_members), + subgroup_members[0], + subgroup, + torch.device("cpu"), + ), + TypeError, + "ProcessGroup", + ) + + dist.destroy_process_group() + uninitialized_rejected = _expect_rejection( + world_binding.validate_runtime, + RuntimeError, + "initialized torch.distributed", + ) + queue.put( + { + "rank": rank, + "world_validated": world_validated, + "order_mismatch_rejected": order_mismatch_rejected, + "size_mismatch_rejected": size_mismatch_rejected, + "device_mismatch_rejected": device_mismatch_rejected, + "subgroup_validated": subgroup_validated, + "nonmember_rejected": nonmember_rejected, + "uninitialized_rejected": uninitialized_rejected, + } + ) + except Exception as exc: + queue.put({"rank": rank, "error": repr(exc)}) + finally: + if dist.is_available() and dist.is_initialized(): + dist.destroy_process_group() + + +def _run_distributed_binding_workers(): + context = mp.get_context("spawn") + queue = context.Queue() + fd, init_file = tempfile.mkstemp(prefix="gefen-checkpoint-binding-") + os.close(fd) + os.unlink(init_file) + processes = [ + context.Process( + target=_distributed_binding_worker, + args=(rank, 3, init_file, queue), + ) + for rank in range(3) + ] + try: + for process in processes: + process.start() + results = [queue.get(timeout=60) for _ in processes] + for process in processes: + process.join(timeout=10) + for process in processes: + if process.is_alive(): + process.terminate() + process.join(timeout=5) + pytest.fail("checkpoint-binding worker hung") + assert process.exitcode == 0 + return sorted(results, key=lambda item: item["rank"]) + finally: + for process in processes: + if process.is_alive(): + process.terminate() + process.join(timeout=5) + if os.path.exists(init_file): + os.unlink(init_file) + + +@pytest.mark.skipif( + not dist.is_available() or not dist.is_gloo_available(), + reason="Gloo is required for process-group binding validation", +) +def test_checkpoint_binding_validates_real_world_and_subgroup_membership(): + results = _run_distributed_binding_workers() + + assert all("error" not in item for item in results), results + assert all(item["world_validated"] for item in results) + assert all(item["order_mismatch_rejected"] for item in results) + assert all(item["size_mismatch_rejected"] for item in results) + assert all(item["device_mismatch_rejected"] for item in results) + assert all(item["subgroup_validated"] for item in results if item["rank"] in (0, 2)) + assert results[1]["nonmember_rejected"] + assert all(item["uninitialized_rejected"] for item in results) diff --git a/tests/test_codebook_scope_cpu.py b/tests/test_codebook_scope_cpu.py new file mode 100644 index 0000000..b427d4a --- /dev/null +++ b/tests/test_codebook_scope_cpu.py @@ -0,0 +1,782 @@ +import copy +from dataclasses import FrozenInstanceError +import io + +import pytest +import torch + +from gefen import ( + CodebookProcessGroupBinding, + Gefen, + GefenMuon, + LogicalSlice, + ParameterIdentity, + ParameterLayout, + ParameterRebinding, + ParameterStateRole, + PlacementKind, + ProcessGroupIdentity, + ProcessGroupScope, + ShardIdentity, + ShardPlacement, + ShardingManifest, + StateExtent, +) +import gefen.gefen as gefen_module + +from _state_snapshot import assert_deep_state_snapshot, deep_state_snapshot + + +def _replicated_shard(parameter, group, member): + coordinate = group.ordered_members.index(member) + return ShardIdentity( + parameter, + ParameterLayout.REPLICATED, + LogicalSlice.full(parameter), + process_group=group, + local_member=member, + placements=( + ShardPlacement( + "data_parallel", + PlacementKind.REPLICATE, + coordinate, + len(group.ordered_members), + ), + ), + ) + + +def _flat_shard(parameter, group, member): + coordinate = group.ordered_members.index(member) + return ShardIdentity( + parameter, + ParameterLayout.FLATTENED_ELEMENT_SHARD, + LogicalSlice.full(parameter), + process_group=group, + local_member=member, + placements=( + ShardPlacement( + "data_parallel", + PlacementKind.FLAT_SHARD, + coordinate, + len(group.ordered_members), + ), + ), + ) + + +def _whole_owner_shard(parameter, group, member, owner): + coordinate = group.ordered_members.index(member) + return ShardIdentity( + parameter, + ParameterLayout.WHOLE_PARAMETER_OWNER, + LogicalSlice.full(parameter) if member == owner else LogicalSlice(0, 0), + process_group=group, + local_member=member, + owner=owner, + placements=( + ShardPlacement( + "data_parallel", + PlacementKind.WHOLE_PARAMETER_OWNER, + coordinate, + len(group.ordered_members), + ), + ), + ) + + +def _single_member_binding(group, member="rank:0"): + return CodebookProcessGroupBinding(group, member, None, torch.device("cpu")) + + +def _finalize_replicated(optimizer, parameter, *, fqn="Layer.Weight"): + group = ProcessGroupIdentity("data_parallel", ("rank:0",)) + identity = ParameterIdentity(fqn, tuple(parameter.shape)) + shard = _replicated_shard(identity, group, "rank:0") + binding = _single_member_binding(group) + optimizer.post_sharding( + (ParameterRebinding(parameter, parameter, shard),), + manifest=ShardingManifest((shard,)), + codebook_process_group=binding, + ) + return binding + + +def _snapshot(optimizer): + return { + "dict": optimizer.__dict__.copy(), + "groups": optimizer.param_groups, + "state": optimizer.state, + "codebook": optimizer._gefen_codebook, + "binding": optimizer._gefen_codebook_process_group, + "validated": optimizer._gefen_codebook_scope_validated, + # Bitwise clones of every reachable state tensor, param-group option, and + # mutable registry so a rejected transaction that mutates contents in + # place (without swapping the top-level container) is still caught. + "deep": deep_state_snapshot(optimizer), + } + + +def _assert_snapshot_identity(optimizer, snapshot): + assert optimizer.param_groups is snapshot["groups"] + assert optimizer.state is snapshot["state"] + assert optimizer._gefen_codebook is snapshot["codebook"] + assert optimizer._gefen_codebook_process_group is snapshot["binding"] + assert optimizer._gefen_codebook_scope_validated is snapshot["validated"] + assert optimizer.__dict__.keys() == snapshot["dict"].keys() + for key, value in snapshot["dict"].items(): + assert optimizer.__dict__[key] is value + assert_deep_state_snapshot(optimizer, snapshot["deep"]) + + +def _replace_native_guard(checkpoint, guard): + checkpoint["gefen_native_local_shards"] = copy.deepcopy(guard) + for group in checkpoint["param_groups"]: + group["_gefen_checkpoint_metadata"]["native_local_shards"] = copy.deepcopy(guard) + + +def test_codebook_process_group_binding_is_public_frozen_and_ordered(): + group = ProcessGroupIdentity("replica", ("worker:b", "worker:a")) + binding = CodebookProcessGroupBinding(group, "worker:b", object(), torch.device("cpu")) + + assert binding.identity is group + assert binding.local_member == "worker:b" + assert binding.sort_key == ("replica", ("worker:b", "worker:a")) + with pytest.raises(FrozenInstanceError): + binding.local_member = "worker:a" + with pytest.raises(TypeError, match="ProcessGroupIdentity"): + CodebookProcessGroupBinding(object(), "worker:b", None, "cpu") + with pytest.raises(ValueError, match="local_member"): + CodebookProcessGroupBinding(group, "missing", None, "cpu") + with pytest.raises(ValueError, match="materialized"): + CodebookProcessGroupBinding(group, "worker:b", None, "meta") + + +def test_scoped_one_member_initialization_matches_unscoped_first_step(): + initial = torch.arange(1, 17, dtype=torch.float32).reshape(4, 4) + scoped_param = torch.nn.Parameter(initial.clone()) + reference_param = torch.nn.Parameter(initial.clone()) + scoped = Gefen([("Layer.Weight", scoped_param)], fused=False, factored_v_2d=False) + reference = Gefen([("Layer.Weight", reference_param)], fused=False, factored_v_2d=False) + binding = _finalize_replicated(scoped, scoped_param) + gradient = torch.linspace(-2, 3, initial.numel()).reshape_as(initial) + scoped_param.grad = gradient.clone() + reference_param.grad = gradient.clone() + + assert scoped.initialize_codebook() + assert not scoped.initialize_codebook() + scoped.step() + reference.step() + + assert scoped.codebook_process_group_binding() is binding + assert scoped.optimizer_contract().capabilities.explicit_process_group_codebook_scope + assert torch.equal(scoped._gefen_codebook, reference._gefen_codebook) + assert torch.equal(scoped_param, reference_param) + for key in scoped.state[scoped_param]: + left = scoped.state[scoped_param][key] + right = reference.state[reference_param][key] + if torch.is_tensor(left): + assert torch.equal(left, right) + else: + assert left == right + frozen_codebook = scoped._gefen_codebook + frozen_indices = scoped.state[scoped_param]["m_codebook"].detach().clone() + scoped_param.grad = None + assert not scoped.refresh_codebook() + assert scoped._gefen_codebook is frozen_codebook + assert torch.equal(scoped.state[scoped_param]["m_codebook"], frozen_indices) + + +def test_scope_binding_validation_is_part_of_atomic_post_sharding(): + parameter = torch.nn.Parameter(torch.ones(8)) + optimizer = Gefen([("p", parameter)], fused=False, factored_v_2d=False) + group = ProcessGroupIdentity("dp", ("rank:0",)) + wrong_group = ProcessGroupIdentity("other", ("rank:0",)) + identity = ParameterIdentity("P", (8,)) + shard = _replicated_shard(identity, group, "rank:0") + manifest = ShardingManifest((shard,)) + snapshot = _snapshot(optimizer) + + with pytest.raises(ValueError, match="process-group identity"): + optimizer.post_sharding( + (ParameterRebinding(parameter, parameter, shard),), + manifest=manifest, + codebook_process_group=_single_member_binding(wrong_group), + ) + + _assert_snapshot_identity(optimizer, snapshot) + assert not optimizer._gefen_post_sharding_finalized + + +def test_scope_binding_rejects_implicit_default_world_for_one_member(): + parameter = torch.nn.Parameter(torch.ones(4)) + optimizer = Gefen([parameter], fused=False, factored_v_2d=False) + group = ProcessGroupIdentity("dp", ("rank:0",)) + identity = ParameterIdentity("P", (4,)) + shard = _replicated_shard(identity, group, "rank:0") + + with pytest.raises(ValueError, match="process_group=None"): + optimizer.post_sharding( + (ParameterRebinding(parameter, parameter, shard),), + manifest=ShardingManifest((shard,)), + codebook_process_group=CodebookProcessGroupBinding(group, "rank:0", object(), "cpu"), + ) + + assert not optimizer._gefen_post_sharding_finalized + + +def test_initialization_failure_does_not_publish_periods_or_codebook(monkeypatch): + parameter = torch.nn.Parameter(torch.arange(8, dtype=torch.float32)) + optimizer = Gefen([("p", parameter)], fused=False, factored_v_2d=False) + _finalize_replicated(optimizer, parameter, fqn="P") + parameter.grad = torch.arange(1, 9, dtype=torch.float32) + before_state = copy.deepcopy(optimizer.state[parameter]) + before_cache = optimizer._gefen_codebook_by_device + + def fail_exact_dp(*args, **kwargs): + raise RuntimeError("injected exact-DP failure") + + monkeypatch.setattr(gefen_module.quantization_module, "exact_dp", fail_exact_dp) + with pytest.raises(RuntimeError, match="injected exact-DP failure"): + optimizer.initialize_codebook() + + assert optimizer.state[parameter] == before_state + assert optimizer._gefen_codebook is None + assert optimizer._gefen_codebook_by_device is before_cache + assert not optimizer._gefen_codebook_scope_validated + assert optimizer._gefen_global_step == 0 + + +def test_scoped_native_checkpoint_uses_primitive_scope_and_requires_exact_binding(): + source_param = torch.nn.Parameter(torch.arange(8, dtype=torch.float32)) + source = Gefen([("p", source_param)], fused=False, factored_v_2d=False) + _finalize_replicated(source, source_param, fqn="P") + source_param.grad = torch.arange(1, 9, dtype=torch.float32) + source.step() + checkpoint = source.state_dict() + + scope = checkpoint["gefen_codebook_scope"] + assert scope == { + "format_version": 1, + "semantic_name": "data_parallel", + "ordered_members": ["rank:0"], + "refresh_every": 0, + } + assert all(group["_gefen_checkpoint_metadata"]["codebook_scope"] == scope for group in checkpoint["param_groups"]) + assert all(group["_gefen_checkpoint_metadata"]["format_version"] == 4 for group in checkpoint["param_groups"]) + buffer = io.BytesIO() + torch.save(checkpoint, buffer) + buffer.seek(0) + loaded = torch.load(buffer, weights_only=True) + assert loaded["gefen_codebook_scope"] == scope + + target_param = torch.nn.Parameter(source_param.detach().clone()) + target = Gefen([("p", target_param)], fused=False, factored_v_2d=False) + _finalize_replicated(target, target_param, fqn="P") + live_binding = target.codebook_process_group_binding() + target.load_state_dict(copy.deepcopy(checkpoint)) + assert target.codebook_process_group_binding() is live_binding + assert not target._gefen_codebook_scope_validated + assert torch.equal(target._gefen_codebook, source._gefen_codebook) + continuation_grad = torch.linspace(-1, 1, 8) + source_param.grad = continuation_grad.clone() + target_param.grad = continuation_grad.clone() + source.step() + target.step() + assert torch.equal(target_param, source_param) + assert torch.equal(target._gefen_codebook, source._gefen_codebook) + + conflicting_param = torch.nn.Parameter(torch.arange(8, dtype=torch.float32)) + conflicting_target = Gefen([("p", conflicting_param)], fused=False, factored_v_2d=False) + _finalize_replicated(conflicting_target, conflicting_param, fqn="P") + conflicting = copy.deepcopy(checkpoint) + conflicting["gefen_codebook_scope"] = dict(conflicting["gefen_codebook_scope"]) + conflicting["gefen_codebook_scope"]["semantic_name"] = "other" + conflicting_before = _snapshot(conflicting_target) + with pytest.raises(ValueError, match="scopes disagree"): + conflicting_target.load_state_dict(conflicting) + _assert_snapshot_identity(conflicting_target, conflicting_before) + + unsafe_version = copy.deepcopy(checkpoint) + unsafe_version["param_groups"][0]["_gefen_checkpoint_metadata"]["format_version"] = 1 + with pytest.raises(ValueError, match="format_version 4"): + conflicting_target.load_state_dict(unsafe_version) + _assert_snapshot_identity(conflicting_target, conflicting_before) + + unbound_param = torch.nn.Parameter(torch.arange(8, dtype=torch.float32)) + unbound = Gefen([("p", unbound_param)], fused=False, factored_v_2d=False) + before = _snapshot(unbound) + with pytest.raises(ValueError, match="does not match"): + unbound.load_state_dict(checkpoint) + _assert_snapshot_identity(unbound, before) + + schedule_param = torch.nn.Parameter(torch.arange(8, dtype=torch.float32)) + schedule_target = Gefen( + [("p", schedule_param)], + fused=False, + factored_v_2d=False, + codebook_refresh_every=1, + ) + _finalize_replicated(schedule_target, schedule_param, fqn="P") + schedule_before = _snapshot(schedule_target) + with pytest.raises(ValueError, match="does not match"): + schedule_target.load_state_dict(checkpoint) + _assert_snapshot_identity(schedule_target, schedule_before) + + +@pytest.mark.parametrize( + "second_layout", + [ + ParameterLayout.FLATTENED_ELEMENT_SHARD, + ParameterLayout.REPLICATED, + ], +) +def test_native_flat_checkpoint_guard_rejects_reordered_equal_shape_slots_atomically( + second_layout, +): + group = ProcessGroupIdentity("data_parallel", ("rank:0",)) + first_identity = ParameterIdentity("Model.First", (4,)) + second_identity = ParameterIdentity("Model.Second", (4,)) + first_shard = _flat_shard(first_identity, group, "rank:0") + second_shard = ( + _flat_shard(second_identity, group, "rank:0") + if second_layout is ParameterLayout.FLATTENED_ELEMENT_SHARD + else _replicated_shard(second_identity, group, "rank:0") + ) + manifest = ShardingManifest((first_shard, second_shard)) + + source_first = torch.nn.Parameter(torch.zeros(4)) + source_second = torch.nn.Parameter(torch.zeros(4)) + source = Gefen( + [("first", source_first), ("second", source_second)], + fused=False, + factored_v_2d=False, + ) + source.post_sharding( + ( + ParameterRebinding(source_first, source_first, first_shard), + ParameterRebinding(source_second, source_second, second_shard), + ), + manifest=manifest, + codebook_process_group=_single_member_binding(group), + ) + source._resolve_automatic_period = lambda *args: 4 + source_first.grad = torch.tensor([1.0, 2.0, 3.0, 4.0]) + source_second.grad = torch.tensor([9.0, -1.0, -2.0, -3.0]) + source.step() + checkpoint = source.state_dict() + source_guard = checkpoint["gefen_native_local_shards"] + source_slots = source_guard["logical_slots"] + assert source_guard["format_version"] == 2 + assert ( + checkpoint["param_groups"][0]["_gefen_checkpoint_metadata"]["native_local_shards"] + is source_guard + ) + buffer = io.BytesIO() + torch.save(checkpoint, buffer) + buffer.seek(0) + assert torch.load(buffer, weights_only=True)["gefen_native_local_shards"] == checkpoint["gefen_native_local_shards"] + malformed_version = copy.deepcopy(checkpoint) + malformed_version["gefen_native_local_shards"]["format_version"] = True + for group_record in malformed_version["param_groups"]: + group_record["_gefen_checkpoint_metadata"]["native_local_shards"]["format_version"] = True + source_before = _snapshot(source) + with pytest.raises(ValueError, match="format_version"): + source.load_state_dict(malformed_version) + _assert_snapshot_identity(source, source_before) + assert [ + ( + record["group_index"], + record["original_slot_index"], + record["compatibility_name"], + record["shard"]["parameter"]["fqn"], + ) + for record in source_slots + ] == [ + (0, 0, "first", "Model.First"), + (0, 1, "second", "Model.Second"), + ] + if second_layout is ParameterLayout.REPLICATED: + assert source_slots[1]["shard"]["layout"] == ParameterLayout.REPLICATED.value + + legacy_checkpoint = copy.deepcopy(checkpoint) + legacy_guard = source._serialized_native_local_shards_v1() + legacy_checkpoint["gefen_native_local_shards"] = legacy_guard + for group_record in legacy_checkpoint["param_groups"]: + group_record["_gefen_checkpoint_metadata"]["native_local_shards"] = copy.deepcopy(legacy_guard) + legacy_first = torch.nn.Parameter(torch.zeros(4)) + legacy_second = torch.nn.Parameter(torch.zeros(4)) + legacy_target = Gefen( + [("target_first", legacy_first), ("target_second", legacy_second)], + fused=False, + factored_v_2d=False, + ) + legacy_target.post_sharding( + ( + ParameterRebinding(legacy_first, legacy_first, first_shard), + ParameterRebinding(legacy_second, legacy_second, second_shard), + ), + manifest=manifest, + codebook_process_group=_single_member_binding(group), + ) + legacy_target.load_state_dict(copy.deepcopy(legacy_checkpoint)) + assert torch.equal(legacy_target._gefen_codebook, source._gefen_codebook) + assert legacy_target.param_groups[0]["param_names"] == [ + "target_first", + "target_second", + ] + assert legacy_target.state[legacy_first]["name"] == "target_first" + assert legacy_target.state[legacy_second]["name"] == "target_second" + + mirror_only_checkpoint = copy.deepcopy(legacy_checkpoint) + mirror_only_checkpoint.pop("gefen_native_local_shards") + legacy_target.load_state_dict(mirror_only_checkpoint) + assert torch.equal(legacy_target._gefen_codebook, source._gefen_codebook) + + top_only_checkpoint = copy.deepcopy(legacy_checkpoint) + for group_record in top_only_checkpoint["param_groups"]: + group_record.pop("_gefen_checkpoint_metadata") + legacy_target.load_state_dict(top_only_checkpoint) + assert torch.equal(legacy_target._gefen_codebook, source._gefen_codebook) + + mixed_version_checkpoint = copy.deepcopy(checkpoint) + for group_record in mixed_version_checkpoint["param_groups"]: + group_record["_gefen_checkpoint_metadata"]["native_local_shards"] = ( + copy.deepcopy(legacy_guard) + ) + source_before = _snapshot(source) + with pytest.raises(ValueError, match="local shards disagree"): + source.load_state_dict(mixed_version_checkpoint) + _assert_snapshot_identity(source, source_before) + + target_second = torch.nn.Parameter(torch.zeros(4)) + target_first = torch.nn.Parameter(torch.zeros(4)) + target = Gefen( + [("second", target_second), ("first", target_first)], + fused=False, + factored_v_2d=False, + ) + target.post_sharding( + ( + ParameterRebinding(target_second, target_second, second_shard), + ParameterRebinding(target_first, target_first, first_shard), + ), + manifest=manifest, + codebook_process_group=_single_member_binding(group), + ) + target_before = _snapshot(target) + + with pytest.raises(ValueError, match="local-shard identity"): + target.load_state_dict(copy.deepcopy(checkpoint)) + + _assert_snapshot_identity(target, target_before) + + +def test_native_v2_guard_retains_pruned_logical_positions_and_rejects_corruption_atomically( + monkeypatch, +): + monkeypatch.setattr( + Gefen, + "_validate_codebook_runtime_binding", + lambda self, binding: None, + ) + group = ProcessGroupIdentity("data_parallel", ("rank:0", "rank:1")) + + def build(): + owned = torch.nn.Parameter(torch.ones(2, 2)) + remote = torch.nn.Parameter(torch.full((2, 2), 2.0)) + replicated = torch.nn.Parameter(torch.full((2, 2), 3.0)) + optimizer = GefenMuon( + [ + {"params": [("owned", owned), ("remote", remote)]}, + {"params": [("replicated", replicated)]}, + ], + fused=False, + ) + owned_identity = ParameterIdentity("Model.Owned", (2, 2)) + remote_identity = ParameterIdentity("Model.Remote", (2, 2)) + replicated_identity = ParameterIdentity("Model.Replicated", (2, 2)) + owned_records = tuple( + _whole_owner_shard(owned_identity, group, member, "rank:0") + for member in group.ordered_members + ) + remote_records = tuple( + _whole_owner_shard(remote_identity, group, member, "rank:1") + for member in group.ordered_members + ) + replicated_records = tuple( + _replicated_shard(replicated_identity, group, member) + for member in group.ordered_members + ) + local_owned = owned_records[0] + local_remote = remote_records[0] + local_replicated = replicated_records[0] + optimizer.post_sharding( + ( + ParameterRebinding(remote, None, local_remote), + ParameterRebinding(replicated, replicated, local_replicated), + ParameterRebinding(owned, owned, local_owned), + ), + manifest=ShardingManifest( + owned_records + remote_records + replicated_records + ), + codebook_process_group=CodebookProcessGroupBinding( + group, + "rank:0", + object(), + torch.device("cpu"), + ), + ) + return optimizer + + source = build() + checkpoint = source.state_dict() + guard = checkpoint["gefen_native_local_shards"] + assert guard["format_version"] == 2 + assert [ + ( + record["group_index"], + record["original_slot_index"], + record["compatibility_name"], + record["shard"]["parameter"]["fqn"], + ) + for record in guard["logical_slots"] + ] == [ + (0, 0, "owned", "Model.Owned"), + (0, 1, "remote", "Model.Remote"), + (1, 0, "replicated", "Model.Replicated"), + ] + assert source.param_groups[0]["param_names"] == ["owned"] + + matching = build() + matching.load_state_dict(copy.deepcopy(checkpoint)) + assert matching._gefen_logical_slots == source._gefen_logical_slots + + malformed_second_mirror = copy.deepcopy(checkpoint) + second_metadata = copy.deepcopy( + malformed_second_mirror["param_groups"][1][ + "_gefen_checkpoint_metadata" + ] + ) + second_metadata["native_local_shards"]["logical_slots"] = tuple( + second_metadata["native_local_shards"]["logical_slots"] + ) + malformed_second_mirror["param_groups"][1][ + "_gefen_checkpoint_metadata" + ] = second_metadata + malformed_target = build() + malformed_before = _snapshot(malformed_target) + with pytest.raises(ValueError, match="logical_slots must be a nonempty list"): + malformed_target.load_state_dict(malformed_second_mirror) + _assert_snapshot_identity(malformed_target, malformed_before) + + corruptions = [] + + class GuardDict(dict): + pass + + corruptions.append((GuardDict(copy.deepcopy(guard)), "invalid schema")) + + invalid_schema = copy.deepcopy(guard) + invalid_schema["logical_slots"][0]["extra"] = None + corruptions.append((invalid_schema, "invalid schema")) + + bool_position = copy.deepcopy(guard) + bool_position["logical_slots"][0]["group_index"] = False + corruptions.append((bool_position, "logical slot is invalid")) + + skipped_position = copy.deepcopy(guard) + skipped_position["logical_slots"][2]["group_index"] = 2 + corruptions.append((skipped_position, "positions must be contiguous")) + + uppercase_name = copy.deepcopy(guard) + uppercase_name["logical_slots"][1]["compatibility_name"] = "Remote" + corruptions.append((uppercase_name, "logical slot is invalid")) + + duplicate_fqn = copy.deepcopy(guard) + duplicate_fqn["logical_slots"][1]["shard"]["parameter"] = copy.deepcopy( + duplicate_fqn["logical_slots"][0]["shard"]["parameter"] + ) + corruptions.append((duplicate_fqn, "FQNs must be unique")) + + missing_group = copy.deepcopy(guard) + missing_group["logical_slots"][0]["shard"]["process_group"] = None + corruptions.append((missing_group, "requires a process-group identity")) + + changed_layout = copy.deepcopy(guard) + changed_layout["logical_slots"][0]["shard"]["layout"] = ( + ParameterLayout.REPLICATED.value + ) + changed_layout["logical_slots"][0]["shard"]["owner"] = None + changed_layout["logical_slots"][0]["shard"]["placements"][0]["kind"] = ( + PlacementKind.REPLICATE.value + ) + corruptions.append((changed_layout, "does not match")) + + changed_fqn = copy.deepcopy(guard) + changed_fqn["logical_slots"][1]["shard"]["parameter"]["fqn"] = ( + "Model.OtherRemote" + ) + corruptions.append((changed_fqn, "does not match")) + + changed_name = copy.deepcopy(guard) + changed_name["logical_slots"][1]["compatibility_name"] = "other" + corruptions.append((changed_name, "does not match")) + + for corrupted_guard, match in corruptions: + target = build() + before = _snapshot(target) + corrupted = copy.deepcopy(checkpoint) + _replace_native_guard(corrupted, corrupted_guard) + with pytest.raises(ValueError, match=match): + target.load_state_dict(corrupted) + _assert_snapshot_identity(target, before) + + def build_all_nonowner(): + remote = torch.nn.Parameter(torch.full((2, 2), 2.0)) + optimizer = GefenMuon([("remote", remote)], fused=False) + identity = ParameterIdentity("Model.Remote", (2, 2)) + records = tuple( + _whole_owner_shard(identity, group, member, "rank:1") + for member in group.ordered_members + ) + optimizer.post_sharding( + (ParameterRebinding(remote, None, records[0]),), + manifest=ShardingManifest(records), + codebook_process_group=CodebookProcessGroupBinding( + group, + "rank:0", + object(), + torch.device("cpu"), + ), + ) + return optimizer + + all_nonowner = build_all_nonowner() + all_nonowner_checkpoint = all_nonowner.state_dict() + all_nonowner_guard = all_nonowner_checkpoint["gefen_native_local_shards"] + assert all_nonowner.param_groups[0]["params"] == [] + assert [ + ( + record["group_index"], + record["original_slot_index"], + record["compatibility_name"], + record["shard"]["parameter"]["fqn"], + ) + for record in all_nonowner_guard["logical_slots"] + ] == [(0, 0, "remote", "Model.Remote")] + all_nonowner_target = build_all_nonowner() + all_nonowner_target.load_state_dict(copy.deepcopy(all_nonowner_checkpoint)) + assert all_nonowner_target._gefen_logical_slots == all_nonowner._gefen_logical_slots + + +def test_unscoped_native_checkpoint_does_not_serialize_a_none_scope(): + parameter = torch.nn.Parameter(torch.arange(4, dtype=torch.float32)) + optimizer = Gefen([parameter], fused=False) + checkpoint = optimizer.state_dict() + + assert "gefen_codebook_scope" not in checkpoint + assert all("codebook_scope" not in group["_gefen_checkpoint_metadata"] for group in checkpoint["param_groups"]) + + +def test_capturable_optimizer_rejects_manual_codebook_replacement(): + parameter = torch.nn.Parameter(torch.arange(8, dtype=torch.float32)) + optimizer = Gefen( + [("p", parameter)], + fused=False, + factored_v_2d=False, + capturable=True, + ) + _finalize_replicated(optimizer, parameter, fqn="P") + parameter.grad = torch.arange(1, 9, dtype=torch.float32) + assert optimizer.initialize_codebook() + + with pytest.raises(RuntimeError, match="capturable=True"): + optimizer.refresh_codebook() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +def test_capturable_scope_rejects_first_active_step_inside_graph_capture(): + parameter = torch.nn.Parameter(torch.arange(8, dtype=torch.float32, device="cuda")) + optimizer = Gefen( + [("p", parameter)], + fused=False, + factored_v_2d=False, + capturable=True, + ) + _finalize_replicated(optimizer, parameter, fqn="P") + parameter.grad = torch.arange(1, 9, dtype=torch.float32, device="cuda") + graph = torch.cuda.CUDAGraph() + + with pytest.raises(RuntimeError, match="eager warmup"): + with torch.cuda.graph(graph): + optimizer.step() + + assert optimizer._gefen_codebook is None + assert optimizer._gefen_global_step == 0 + + +def test_capturable_multi_member_scope_binding_is_atomic_rejection(): + parameter = torch.nn.Parameter(torch.arange(8, dtype=torch.float32)) + optimizer = Gefen( + [("p", parameter)], + fused=False, + factored_v_2d=False, + capturable=True, + ) + group = ProcessGroupIdentity("dp", ("rank:0", "rank:1")) + identity = ParameterIdentity("P", (8,)) + records = tuple(_replicated_shard(identity, group, member) for member in group.ordered_members) + before = _snapshot(optimizer) + + with pytest.raises(ValueError, match="capturable=True"): + optimizer.post_sharding( + (ParameterRebinding(parameter, parameter, records[0]),), + manifest=ShardingManifest(records), + codebook_process_group=CodebookProcessGroupBinding(group, "rank:0", object(), "cpu"), + ) + + _assert_snapshot_identity(optimizer, before) + + +def test_muon_whole_owner_scope_enables_owner_update_with_sync_requirement(): + parameter = torch.nn.Parameter(torch.ones(4, 4)) + optimizer = GefenMuon([("matrix", parameter)], fused=False) + group = ProcessGroupIdentity("owner_group", ("rank:0",)) + identity = ParameterIdentity("Matrix", (4, 4)) + shard = _whole_owner_shard(identity, group, "rank:0", "rank:0") + binding = _single_member_binding(group) + optimizer.post_sharding( + (ParameterRebinding(parameter, parameter, shard),), + manifest=ShardingManifest((shard,)), + codebook_process_group=binding, + ) + before = parameter.detach().clone() + parameter.grad = torch.arange(1, 17, dtype=torch.float32).reshape(4, 4) + + optimizer.step() + + support = next( + item + for item in optimizer.optimizer_contract().capabilities.training + if item.layout is ParameterLayout.WHOLE_PARAMETER_OWNER + ) + assert support.process_group_scope is ProcessGroupScope.ADAPTER_DEFINED + assert support.requires_complete_parameter_storage + assert support.requires_post_step_parameter_sync + assert all( + ParameterLayout.WHOLE_PARAMETER_OWNER not in checkpoint.same_topology + for checkpoint in optimizer.optimizer_contract().capabilities.checkpoints + ) + owner_variants = [ + variant + for variant in optimizer.optimizer_contract().state_layout.parameter_variants + if ParameterLayout.WHOLE_PARAMETER_OWNER in variant.layouts and variant.initialized + ] + assert owner_variants + assert all( + variant.extent is StateExtent.OWNER_PARAMETER and variant.role is ParameterStateRole.OWNER + for variant in owner_variants + ) + assert not torch.equal(parameter, before) + assert optimizer._gefen_global_step == 1 diff --git a/tests/test_codebook_scope_distributed.py b/tests/test_codebook_scope_distributed.py new file mode 100644 index 0000000..b55c458 --- /dev/null +++ b/tests/test_codebook_scope_distributed.py @@ -0,0 +1,1421 @@ +from datetime import timedelta +import copy +import multiprocessing as mp +import os +import tempfile + +import pytest +import torch +import torch.distributed as dist + +from gefen import ( + CodebookProcessGroupBinding, + Gefen, + GefenMuon, + LogicalSlice, + ParameterIdentity, + ParameterLayout, + ParameterRebinding, + PlacementKind, + ProcessGroupIdentity, + ShardIdentity, + ShardPlacement, + ShardingManifest, +) +from gefen.gefen import learn_gefen_exact_codebook_from_grad_periods +import gefen.gefen as gefen_module + + +def _replicated(parameter, group, member): + coordinate = group.ordered_members.index(member) + return ShardIdentity( + parameter, + ParameterLayout.REPLICATED, + LogicalSlice.full(parameter), + process_group=group, + local_member=member, + placements=( + ShardPlacement( + "dp", + PlacementKind.REPLICATE, + coordinate, + len(group.ordered_members), + ), + ), + ) + + +def _flat(parameter, group, member, offset, length): + coordinate = group.ordered_members.index(member) + return ShardIdentity( + parameter, + ParameterLayout.FLATTENED_ELEMENT_SHARD, + LogicalSlice(offset, length), + process_group=group, + local_member=member, + placements=( + ShardPlacement( + "dp", + PlacementKind.FLAT_SHARD, + coordinate, + len(group.ordered_members), + ), + ), + ) + + +def _whole_owner(parameter, group, member, owner): + coordinate = group.ordered_members.index(member) + return ShardIdentity( + parameter, + ParameterLayout.WHOLE_PARAMETER_OWNER, + LogicalSlice.full(parameter) if member == owner else LogicalSlice(0, 0), + process_group=group, + local_member=member, + owner=owner, + placements=( + ShardPlacement( + "dp", + PlacementKind.WHOLE_PARAMETER_OWNER, + coordinate, + len(group.ordered_members), + ), + ), + ) + + +def _binding(group, rank, process_group): + return CodebookProcessGroupBinding(group, "rank:{}".format(rank), process_group, torch.device("cpu")) + + +def _finalize(optimizer, parameter, local_shard, manifest, binding): + optimizer.post_sharding( + (ParameterRebinding(parameter, parameter, local_shard),), + manifest=manifest, + codebook_process_group=binding, + ) + + +def _distributed_worker(rank, world, init_file, queue): + try: + dist.init_process_group( + "gloo", + init_method="file://{}".format(init_file), + rank=rank, + world_size=world, + timeout=timedelta(seconds=45), + ) + members = tuple("rank:{}".format(index) for index in range(world)) + group = ProcessGroupIdentity("data_parallel", members) + runtime_group = dist.group.WORLD + + # Runtime coordinate/member validation must fail before post_sharding + # publication, then leave the same optimizer retryable with a valid map. + invalid_param = torch.nn.Parameter(torch.ones(4)) + invalid_optimizer = Gefen([("invalid", invalid_param)], fused=False, factored_v_2d=False) + invalid_identity = ParameterIdentity("Invalid", (4,)) + reversed_group = ProcessGroupIdentity("data_parallel", tuple(reversed(members))) + invalid_records = tuple( + _replicated(invalid_identity, reversed_group, member) for member in reversed_group.ordered_members + ) + invalid_local = next(item for item in invalid_records if item.local_member == "rank:{}".format(rank)) + bad_binding = CodebookProcessGroupBinding( + reversed_group, + "rank:{}".format(rank), + runtime_group, + torch.device("cpu"), + ) + try: + invalid_optimizer.post_sharding( + (ParameterRebinding(invalid_param, invalid_param, invalid_local),), + manifest=ShardingManifest(invalid_records), + codebook_process_group=bad_binding, + ) + invalid_rejected = False + except ValueError: + invalid_rejected = True + invalid_atomic = ( + not invalid_optimizer._gefen_post_sharding_finalized + and invalid_optimizer._gefen_codebook_process_group is None + and invalid_optimizer.param_groups[0]["params"] == [invalid_param] + ) + + # Replicated logical state contributes exactly once from ordered member + # zero. Rank 1 intentionally has different values but the same period; + # including both replicas would produce a different histogram oracle. + replicated_param = torch.nn.Parameter(torch.zeros(4)) + replicated_optimizer = Gefen( + [("replicated", replicated_param)], + fused=False, + factored_v_2d=False, + ) + replicated_identity = ParameterIdentity("Replicated", (4,)) + replicated_records = tuple(_replicated(replicated_identity, group, member) for member in members) + _finalize( + replicated_optimizer, + replicated_param, + replicated_records[rank], + ShardingManifest(replicated_records), + _binding(group, rank, runtime_group), + ) + replicated_optimizer._resolve_automatic_period = lambda *args: 4 + canonical_grad = torch.tensor([-4.0, -1.0, 2.0, 8.0]) + other_grad = torch.tensor([-8.0, 3.0, 4.0, 4.0]) + replicated_param.grad = canonical_grad.clone() if rank == 0 else other_grad + replicated_initialized = replicated_optimizer.initialize_codebook() + replicated_oracle = learn_gefen_exact_codebook_from_grad_periods( + grad_periods=(("Replicated", canonical_grad, 4, canonical_grad),), + codebook_device=torch.device("cpu"), + num_codebooks=256, + force_endpoints=True, + verbose=False, + compute_mse_logging=False, + use_fused_histogram=False, + ) + replicated_logical_once = torch.equal(replicated_optimizer._gefen_codebook, replicated_oracle) + + mismatch_param = torch.nn.Parameter(torch.zeros(4)) + mismatch_optimizer = Gefen([("mismatch", mismatch_param)], fused=False, factored_v_2d=False) + mismatch_identity = ParameterIdentity("Mismatch", (4,)) + mismatch_records = tuple(_replicated(mismatch_identity, group, member) for member in members) + _finalize( + mismatch_optimizer, + mismatch_param, + mismatch_records[rank], + ShardingManifest(mismatch_records), + _binding(group, rank, runtime_group), + ) + mismatch_optimizer._resolve_automatic_period = lambda *args: 4 + mismatch_param.grad = canonical_grad.clone() if rank == 0 else None + try: + mismatch_optimizer.initialize_codebook() + mismatch_rejected = False + except RuntimeError as exc: + mismatch_rejected = "identical gradient presence" in str(exc) + mismatch_atomic = ( + mismatch_optimizer._gefen_codebook is None + and mismatch_optimizer._gefen_global_step == 0 + and mismatch_optimizer.state[mismatch_param] == {"name": "mismatch"} + ) + + amp_param = torch.nn.Parameter(torch.zeros(4)) + amp_optimizer = Gefen([("amp", amp_param)], fused=False, factored_v_2d=False) + amp_identity = ParameterIdentity("Amp", (8,)) + amp_records = tuple(_flat(amp_identity, group, member, index * 4, 4) for index, member in enumerate(members)) + _finalize( + amp_optimizer, + amp_param, + amp_records[rank], + ShardingManifest(amp_records), + _binding(group, rank, runtime_group), + ) + amp_optimizer._resolve_automatic_period = lambda *args: 4 + amp_param.grad = canonical_grad.clone() if rank == 0 else other_grad.clone() + amp_grad_before = amp_param.grad.detach().clone() + amp_optimizer.found_inf = torch.tensor(float(rank == 0)) + amp_optimizer.grad_scale = torch.tensor(8.0) + try: + amp_optimizer.step() + amp_mismatch_rejected = False + except RuntimeError as exc: + amp_mismatch_rejected = "group-aware gradient scaler" in str(exc) + amp_optimizer.found_inf = torch.tensor(1.0) + amp_optimizer.step() + amp_overflow_atomic = ( + amp_optimizer._gefen_global_step == 0 + and amp_optimizer._gefen_codebook is None + and amp_optimizer.state[amp_param] == {"name": "amp"} + and torch.equal(amp_param.grad, amp_grad_before) + and torch.equal(amp_param, torch.zeros_like(amp_param)) + ) + if rank == 0: + amp_optimizer.found_inf = torch.tensor(0.0) + else: + del amp_optimizer.found_inf + del amp_optimizer.grad_scale + try: + amp_optimizer.step() + amp_protocol_rejected = False + except RuntimeError as exc: + amp_protocol_rejected = "policy" in str(exc) + + # Flattened shards each contribute once. The global oracle uses two + # period-4 blocks, exactly matching the two physical local shards. + flat_param = torch.nn.Parameter(torch.zeros(4)) + flat_optimizer = Gefen([("flat", flat_param)], fused=False, factored_v_2d=False) + flat_identity = ParameterIdentity("Flat", (8,)) + flat_records = tuple(_flat(flat_identity, group, member, index * 4, 4) for index, member in enumerate(members)) + _finalize( + flat_optimizer, + flat_param, + flat_records[rank], + ShardingManifest(flat_records), + _binding(group, rank, runtime_group), + ) + flat_optimizer._resolve_automatic_period = lambda *args: 4 + flat_grads = ( + torch.tensor([-3.0, -2.0, 1.0, 7.0]), + torch.tensor([-9.0, 2.0, 5.0, 6.0]), + ) + flat_param.grad = flat_grads[rank].clone() + flat_initialized = flat_optimizer.initialize_codebook() + flat_oracle = learn_gefen_exact_codebook_from_grad_periods( + grad_periods=tuple(("Flat", gradient, 4, gradient) for gradient in flat_grads), + codebook_device=torch.device("cpu"), + num_codebooks=256, + force_endpoints=True, + verbose=False, + compute_mse_logging=False, + use_fused_histogram=False, + ) + flat_global = torch.equal(flat_optimizer._gefen_codebook, flat_oracle) + flat_native_claim = any( + ParameterLayout.FLATTENED_ELEMENT_SHARD in item.same_topology + for item in flat_optimizer.optimizer_contract().capabilities.checkpoints + ) + + flat_mismatch_param = torch.nn.Parameter(torch.zeros(4)) + flat_mismatch_optimizer = Gefen( + [("flat_mismatch", flat_mismatch_param)], + fused=False, + factored_v_2d=False, + ) + flat_mismatch_identity = ParameterIdentity("FlatMismatch", (8,)) + flat_mismatch_records = tuple( + _flat(flat_mismatch_identity, group, member, index * 4, 4) for index, member in enumerate(members) + ) + _finalize( + flat_mismatch_optimizer, + flat_mismatch_param, + flat_mismatch_records[rank], + ShardingManifest(flat_mismatch_records), + _binding(group, rank, runtime_group), + ) + flat_mismatch_optimizer._resolve_automatic_period = lambda *args: 4 + flat_mismatch_param.grad = flat_grads[rank].clone() if rank == 0 else None + try: + flat_mismatch_optimizer.initialize_codebook() + flat_mismatch_rejected = False + except RuntimeError as exc: + flat_mismatch_rejected = "every nonempty shard" in str(exc) + flat_mismatch_atomic = flat_mismatch_optimizer._gefen_codebook is None and flat_mismatch_optimizer.state[ + flat_mismatch_param + ] == {"name": "flat_mismatch"} + + gathered_codebooks = [torch.empty_like(flat_optimizer._gefen_codebook) for _ in range(world)] + dist.all_gather(gathered_codebooks, flat_optimizer._gefen_codebook) + flat_agreement = all(torch.equal(item, gathered_codebooks[0]) for item in gathered_codebooks[1:]) + flat_optimizer.step() + flat_checkpoint = copy.deepcopy(flat_optimizer.state_dict()) + flat_canonical = flat_optimizer.export_canonical_state() + checkpoint_scopes = [None] * world + dist.all_gather_object( + checkpoint_scopes, + flat_checkpoint["gefen_codebook_scope"], + group=runtime_group, + ) + rank_neutral_checkpoint_scope = all(item == checkpoint_scopes[0] for item in checkpoint_scopes[1:]) + local_shard_records = [None] * world + dist.all_gather_object( + local_shard_records, + flat_checkpoint["gefen_native_local_shards"], + group=runtime_group, + ) + rank_local_checkpoint_identity = local_shard_records[0] != local_shard_records[1] + flat_v2_guard_valid = ( + flat_checkpoint["gefen_native_local_shards"]["format_version"] == 2 + and len(flat_checkpoint["gefen_native_local_shards"]["logical_slots"]) + == 1 + and flat_checkpoint["gefen_native_local_shards"]["logical_slots"][0][ + "group_index" + ] + == 0 + and flat_checkpoint["gefen_native_local_shards"]["logical_slots"][0][ + "original_slot_index" + ] + == 0 + and flat_checkpoint["gefen_native_local_shards"]["logical_slots"][0][ + "compatibility_name" + ] + == "flat" + and flat_checkpoint["gefen_native_local_shards"]["logical_slots"][0][ + "shard" + ]["local_member"] + == "rank:{}".format(rank) + ) + rank_zero_checkpoint = [flat_checkpoint if rank == 0 else None] + dist.broadcast_object_list(rank_zero_checkpoint, src=0, group=runtime_group) + cross_param = torch.nn.Parameter(flat_param.detach().clone()) + cross_target = Gefen([("flat", cross_param)], fused=False, factored_v_2d=False) + _finalize( + cross_target, + cross_param, + flat_records[rank], + ShardingManifest(flat_records), + _binding(group, rank, runtime_group), + ) + try: + cross_target.load_state_dict(rank_zero_checkpoint[0]) + cross_member_guard = rank == 0 + except ValueError as exc: + cross_member_guard = rank == 1 and "local-shard identity" in str(exc) + canonical_shards = [None] * world + dist.all_gather_object( + canonical_shards, + flat_canonical["parameters"]["Flat"]["shard"], + group=runtime_group, + ) + canonical_rank_local_identity = canonical_shards[0] != canonical_shards[1] + rank_zero_canonical = [flat_canonical if rank == 0 else None] + dist.broadcast_object_list(rank_zero_canonical, src=0, group=runtime_group) + cross_canonical_param = torch.nn.Parameter(flat_param.detach().clone()) + cross_canonical = Gefen( + [("flat", cross_canonical_param)], + fused=False, + factored_v_2d=False, + ) + _finalize( + cross_canonical, + cross_canonical_param, + flat_records[rank], + ShardingManifest(flat_records), + _binding(group, rank, runtime_group), + ) + try: + cross_canonical.import_canonical_state(rank_zero_canonical[0]) + canonical_cross_member_guard = rank == 0 + except ValueError as exc: + canonical_cross_member_guard = rank == 1 and "shard" in str(exc) + resumed_param = torch.nn.Parameter(flat_param.detach().clone()) + resumed = Gefen([("flat", resumed_param)], fused=False, factored_v_2d=False) + resumed_binding = _binding(group, rank, runtime_group) + _finalize( + resumed, + resumed_param, + flat_records[rank], + ShardingManifest(flat_records), + resumed_binding, + ) + resumed.load_state_dict(flat_checkpoint) + legacy_checkpoint = copy.deepcopy(flat_checkpoint) + legacy_guard = flat_optimizer._serialized_native_local_shards_v1() + legacy_checkpoint["gefen_native_local_shards"] = legacy_guard + for legacy_group in legacy_checkpoint["param_groups"]: + legacy_group["_gefen_checkpoint_metadata"][ + "native_local_shards" + ] = copy.deepcopy(legacy_guard) + legacy_param = torch.nn.Parameter(flat_param.detach().clone()) + legacy_resumed = Gefen( + [("flat", legacy_param)], + fused=False, + factored_v_2d=False, + ) + legacy_binding = _binding(group, rank, runtime_group) + _finalize( + legacy_resumed, + legacy_param, + flat_records[rank], + ShardingManifest(flat_records), + legacy_binding, + ) + legacy_resumed.load_state_dict(legacy_checkpoint) + legacy_v1_guard_accepted = ( + legacy_resumed.codebook_process_group_binding() is legacy_binding + and torch.equal( + legacy_resumed._gefen_codebook, + flat_optimizer._gefen_codebook, + ) + ) + canonical_param = torch.nn.Parameter(flat_param.detach().clone()) + canonical_resumed = Gefen( + [("flat", canonical_param)], + fused=False, + factored_v_2d=False, + ) + canonical_binding = _binding(group, rank, runtime_group) + _finalize( + canonical_resumed, + canonical_param, + flat_records[rank], + ShardingManifest(flat_records), + canonical_binding, + ) + canonical_resumed.import_canonical_state(flat_canonical) + continuation_grad = flat_grads[rank].flip(0).clone() + resumed_param.grad = continuation_grad.clone() + canonical_param.grad = continuation_grad.clone() + flat_param.grad = continuation_grad.clone() + canonical_resumed.step() + resumed.step() + flat_optimizer.step() + flat_checkpoint_continuation = ( + torch.equal(resumed_param, flat_param) + and resumed.codebook_process_group_binding() is resumed_binding + and torch.equal(resumed._gefen_codebook, flat_optimizer._gefen_codebook) + ) + flat_canonical_continuation = ( + torch.equal(canonical_param, flat_param) + and canonical_resumed.codebook_process_group_binding() is canonical_binding + and torch.equal( + canonical_resumed._gefen_codebook, + flat_optimizer._gefen_codebook, + ) + ) + refresh_succeeded = flat_optimizer.refresh_codebook() + continuation_grads = tuple(gradient.flip(0) for gradient in flat_grads) + refresh_oracle = learn_gefen_exact_codebook_from_grad_periods( + grad_periods=tuple(("Flat", gradient, 4, gradient) for gradient in continuation_grads), + codebook_device=torch.device("cpu"), + num_codebooks=256, + force_endpoints=True, + verbose=False, + compute_mse_logging=False, + use_fused_histogram=False, + ) + refresh_matches_oracle = torch.equal(flat_optimizer._gefen_codebook, refresh_oracle) + refreshed_codebooks = [torch.empty_like(flat_optimizer._gefen_codebook) for _ in range(world)] + dist.all_gather(refreshed_codebooks, flat_optimizer._gefen_codebook) + refresh_agreement = all(torch.equal(item, refreshed_codebooks[0]) for item in refreshed_codebooks[1:]) + refresh_failure_grad = flat_grads[rank].roll(1).clone() + flat_param.grad = refresh_failure_grad + old_refresh_codebook_object = flat_optimizer._gefen_codebook + old_refresh_codebook = flat_optimizer._gefen_codebook.detach().clone() + old_refresh_indices = flat_optimizer.state[flat_param]["m_codebook"].detach().clone() + old_refresh_magnitude = flat_optimizer.state[flat_param]["m_magnitude"].detach().clone() + old_refresh_period = flat_optimizer.state[flat_param]["automatic_period"] + old_refresh_step = flat_optimizer._gefen_global_step + codebook_cache = flat_optimizer._gefen_codebook_by_device + lut_cache = flat_optimizer._gefen_codebook_lut_by_device + cache_marker = torch.tensor([17.0]) + lut_marker = torch.tensor([19.0]) + codebook_cache[torch.device("cpu")] = cache_marker + lut_cache[torch.device("cpu")] = lut_marker + original_nearest = gefen_module.gefen_nearest_codebook_indices + if rank == 1: + + def fail_nearest(*args, **kwargs): + raise RuntimeError("rank-local requantization failure") + + gefen_module.gefen_nearest_codebook_indices = fail_nearest + try: + flat_optimizer.refresh_codebook() + refresh_failure_seen = False + except RuntimeError as exc: + refresh_failure_seen = "requantization" in str(exc) + finally: + gefen_module.gefen_nearest_codebook_indices = original_nearest + refresh_failure_atomic = ( + flat_optimizer._gefen_codebook is old_refresh_codebook_object + and torch.equal(flat_optimizer._gefen_codebook, old_refresh_codebook) + and torch.equal( + flat_optimizer.state[flat_param]["m_codebook"], + old_refresh_indices, + ) + and torch.equal( + flat_optimizer.state[flat_param]["m_magnitude"], + old_refresh_magnitude, + ) + and flat_optimizer.state[flat_param]["automatic_period"] == old_refresh_period + and flat_optimizer._gefen_global_step == old_refresh_step + and flat_optimizer._gefen_codebook_by_device is codebook_cache + and flat_optimizer._gefen_codebook_lut_by_device is lut_cache + and codebook_cache[torch.device("cpu")] is cache_marker + and lut_cache[torch.device("cpu")] is lut_marker + ) + refresh_retry_succeeded = flat_optimizer.refresh_codebook() + + periodic_param = torch.nn.Parameter(torch.zeros(4)) + periodic_optimizer = Gefen( + [("periodic", periodic_param)], + fused=False, + factored_v_2d=False, + codebook_refresh_every=2, + ) + periodic_identity = ParameterIdentity("Periodic", (8,)) + periodic_records = tuple( + _flat(periodic_identity, group, member, index * 4, 4) for index, member in enumerate(members) + ) + _finalize( + periodic_optimizer, + periodic_param, + periodic_records[rank], + ShardingManifest(periodic_records), + _binding(group, rank, runtime_group), + ) + periodic_optimizer._resolve_automatic_period = lambda *args: 4 + periodic_param.grad = flat_grads[rank].clone() + periodic_optimizer.step() + first_periodic_codebook = periodic_optimizer._gefen_codebook + + if rank == 0: + periodic_param.grad = torch.sparse_coo_tensor( + torch.tensor([[0]]), + torch.tensor([1.0]), + size=(4,), + ) + else: + periodic_param.grad = flat_grads[rank].roll(1).clone() + try: + periodic_optimizer.step() + periodic_nondue_failure_rejected = False + except RuntimeError as exc: + periodic_nondue_failure_rejected = "gradient preflight" in str(exc) + periodic_nondue_failure_rejected = ( + periodic_nondue_failure_rejected + and periodic_optimizer._gefen_global_step == 1 + and periodic_optimizer._gefen_codebook is first_periodic_codebook + ) + + periodic_param.grad = flat_grads[rank].roll(1).clone() + periodic_optimizer.step() + periodic_nondue_failure_rejected = ( + periodic_nondue_failure_rejected + and periodic_optimizer._gefen_global_step == 2 + and periodic_optimizer._gefen_codebook is first_periodic_codebook + ) + periodic_param.grad = flat_grads[rank].flip(0).clone() + periodic_optimizer.step() + periodic_codebooks = [torch.empty_like(periodic_optimizer._gefen_codebook) for _ in range(world)] + dist.all_gather(periodic_codebooks, periodic_optimizer._gefen_codebook) + periodic_refresh_valid = ( + periodic_optimizer._gefen_global_step == 3 + and periodic_optimizer._gefen_codebook is not first_periodic_codebook + and all(torch.equal(item, periodic_codebooks[0]) for item in periodic_codebooks[1:]) + and periodic_optimizer.state_dict()["gefen_codebook_scope"]["refresh_every"] == 2 + ) + + # Whole-matrix owner mode keeps no fake tensor/state on nonowners. Every + # member still joins codebook collectives and receives common state; + # the adapter performs the separately declared post-step matrix sync. + owner_source = torch.nn.Parameter(torch.ones(2, 2)) + owner_optimizer = GefenMuon([("matrix", owner_source)], fused=False) + owner_identity = ParameterIdentity("Matrix", (2, 2)) + owner_records = tuple(_whole_owner(owner_identity, group, member, "rank:0") for member in members) + owner_optimizer.post_sharding( + ( + ParameterRebinding( + owner_source, + owner_source if rank == 0 else None, + owner_records[rank], + ), + ), + manifest=ShardingManifest(owner_records), + codebook_process_group=_binding(group, rank, runtime_group), + ) + if rank == 0: + owner_source.grad = torch.tensor([[1.0, -2.0], [3.0, -4.0]]) + owner_initialized = owner_optimizer.initialize_codebook() + owner_codebooks = [torch.empty_like(owner_optimizer._gefen_codebook) for _ in range(world)] + dist.all_gather(owner_codebooks, owner_optimizer._gefen_codebook) + owner_agreement = all(torch.equal(item, owner_codebooks[0]) for item in owner_codebooks[1:]) + owner_optimizer.step() + synchronized_matrix = owner_source.detach().clone() if rank == 0 else torch.empty(2, 2, dtype=torch.float32) + dist.broadcast(synchronized_matrix, src=0, group=runtime_group) + owner_step_valid = ( + owner_optimizer._gefen_global_step == 1 + and not torch.equal(synchronized_matrix, torch.ones_like(synchronized_matrix)) + and (rank != 0 or torch.equal(owner_source, synchronized_matrix)) + and (rank == 0 or (not owner_optimizer.param_groups[0]["params"] and not owner_optimizer.state)) + ) + owner_support = next( + item + for item in owner_optimizer.optimizer_contract().capabilities.training + if item.layout is ParameterLayout.WHOLE_PARAMETER_OWNER + ) + owner_contract_valid = owner_support.requires_post_step_parameter_sync + owner_amp_protocol = owner_optimizer._step_supports_amp_scaling + + # A failure after the shared histogram reduction must be observed by + # every member before staged periods or the codebook are committed. + failure_param = torch.nn.Parameter(torch.zeros(4)) + failure_optimizer = Gefen([("failure", failure_param)], fused=False, factored_v_2d=False) + failure_identity = ParameterIdentity("Failure", (8,)) + failure_records = tuple( + _flat(failure_identity, group, member, index * 4, 4) for index, member in enumerate(members) + ) + _finalize( + failure_optimizer, + failure_param, + failure_records[rank], + ShardingManifest(failure_records), + _binding(group, rank, runtime_group), + ) + failure_optimizer._resolve_automatic_period = lambda *args: 4 + failure_param.grad = flat_grads[rank].clone() + original_exact_dp = gefen_module.quantization_module.exact_dp + if rank == 1: + + def fail_exact_dp(*args, **kwargs): + raise RuntimeError("rank-local exact-DP failure") + + gefen_module.quantization_module.exact_dp = fail_exact_dp + try: + failure_optimizer.initialize_codebook() + failure_seen = False + except RuntimeError as exc: + failure_seen = "exact-DP" in str(exc) + finally: + gefen_module.quantization_module.exact_dp = original_exact_dp + failure_atomic = ( + failure_optimizer._gefen_codebook is None + and failure_optimizer._gefen_global_step == 0 + and failure_optimizer.state[failure_param] == {"name": "failure"} + ) + retry_succeeded = failure_optimizer.initialize_codebook() + restored_codebook = failure_optimizer._gefen_codebook + if rank == 0: + failure_optimizer._gefen_codebook = None + try: + failure_optimizer.refresh_codebook() + presence_mismatch_rejected = False + except RuntimeError as exc: + presence_mismatch_rejected = "old codebook differs" in str(exc) + if rank == 0: + failure_optimizer._gefen_codebook = restored_codebook + original_manifest = failure_optimizer._gefen_sharding_manifest + failure_optimizer._gefen_codebook_scope_validated = False + manifest_guard_codebook = failure_optimizer._gefen_codebook + manifest_guard_step = failure_optimizer._gefen_global_step + if rank == 0: + alternate_identity = ParameterIdentity("Alternate", (8,)) + alternate_records = tuple( + _flat(alternate_identity, group, member, index * 4, 4) for index, member in enumerate(members) + ) + failure_optimizer._gefen_sharding_manifest = ShardingManifest(alternate_records) + try: + failure_optimizer.initialize_codebook() + manifest_mismatch_rejected = False + except RuntimeError: + manifest_mismatch_rejected = True + manifest_mismatch_rejected = ( + manifest_mismatch_rejected + and failure_optimizer._gefen_codebook is manifest_guard_codebook + and failure_optimizer._gefen_global_step == manifest_guard_step + ) + failure_optimizer._gefen_sharding_manifest = original_manifest + + queue.put( + { + "rank": rank, + "invalid_rejected": invalid_rejected, + "invalid_atomic": invalid_atomic, + "replicated_initialized": replicated_initialized, + "replicated_logical_once": replicated_logical_once, + "mismatch_rejected": mismatch_rejected, + "mismatch_atomic": mismatch_atomic, + "amp_overflow_atomic": amp_overflow_atomic, + "amp_mismatch_rejected": amp_mismatch_rejected, + "amp_protocol_rejected": amp_protocol_rejected, + "flat_initialized": flat_initialized, + "flat_global": flat_global, + "flat_native_claim": flat_native_claim, + "flat_mismatch_rejected": flat_mismatch_rejected, + "flat_mismatch_atomic": flat_mismatch_atomic, + "flat_agreement": flat_agreement, + "flat_checkpoint_continuation": flat_checkpoint_continuation, + "rank_neutral_checkpoint_scope": rank_neutral_checkpoint_scope, + "rank_local_checkpoint_identity": rank_local_checkpoint_identity, + "flat_v2_guard_valid": flat_v2_guard_valid, + "legacy_v1_guard_accepted": legacy_v1_guard_accepted, + "cross_member_guard": cross_member_guard, + "canonical_rank_local_identity": canonical_rank_local_identity, + "canonical_cross_member_guard": canonical_cross_member_guard, + "flat_canonical_continuation": flat_canonical_continuation, + "refresh_succeeded": refresh_succeeded, + "refresh_matches_oracle": refresh_matches_oracle, + "refresh_agreement": refresh_agreement, + "refresh_failure_seen": refresh_failure_seen, + "refresh_failure_atomic": refresh_failure_atomic, + "refresh_retry_succeeded": refresh_retry_succeeded, + "periodic_nondue_failure_rejected": periodic_nondue_failure_rejected, + "periodic_refresh_valid": periodic_refresh_valid, + "owner_initialized": owner_initialized, + "owner_agreement": owner_agreement, + "owner_step_valid": owner_step_valid, + "owner_contract_valid": owner_contract_valid, + "owner_amp_protocol": owner_amp_protocol, + "failure_seen": failure_seen, + "failure_atomic": failure_atomic, + "retry_succeeded": retry_succeeded, + "presence_mismatch_rejected": presence_mismatch_rejected, + "manifest_mismatch_rejected": manifest_mismatch_rejected, + } + ) + except Exception as exc: + queue.put({"rank": rank, "error": repr(exc)}) + finally: + if dist.is_available() and dist.is_initialized(): + dist.destroy_process_group() + + +def _run_workers(world=2): + context = mp.get_context("spawn") + queue = context.Queue() + fd, init_file = tempfile.mkstemp(prefix="gefen-codebook-scope-") + os.close(fd) + os.unlink(init_file) + processes = [ + context.Process( + target=_distributed_worker, + args=(rank, world, init_file, queue), + ) + for rank in range(world) + ] + try: + for process in processes: + process.start() + results = [queue.get(timeout=60) for _ in processes] + for process in processes: + process.join(timeout=10) + for process in processes: + if process.is_alive(): + process.terminate() + process.join(timeout=5) + pytest.fail("scoped codebook worker hung") + assert process.exitcode == 0 + return sorted(results, key=lambda item: item["rank"]) + finally: + for process in processes: + if process.is_alive(): + process.terminate() + process.join(timeout=5) + if os.path.exists(init_file): + os.unlink(init_file) + + +def test_explicit_gloo_scope_aggregates_logical_state_and_fails_atomically(): + results = _run_workers() + + assert all("error" not in item for item in results), results + for item in results: + assert item["invalid_rejected"], item + assert item["invalid_atomic"], item + assert item["replicated_initialized"], item + assert item["replicated_logical_once"], item + assert item["mismatch_rejected"], item + assert item["mismatch_atomic"], item + assert item["amp_overflow_atomic"], item + assert item["amp_mismatch_rejected"], item + assert item["amp_protocol_rejected"], item + assert item["flat_initialized"], item + assert item["flat_global"], item + assert item["flat_native_claim"], item + assert item["flat_mismatch_rejected"], item + assert item["flat_mismatch_atomic"], item + assert item["flat_agreement"], item + assert item["flat_checkpoint_continuation"], item + assert item["rank_neutral_checkpoint_scope"], item + assert item["rank_local_checkpoint_identity"], item + assert item["flat_v2_guard_valid"], item + assert item["legacy_v1_guard_accepted"], item + assert item["cross_member_guard"], item + assert item["canonical_rank_local_identity"], item + assert item["canonical_cross_member_guard"], item + assert item["flat_canonical_continuation"], item + assert item["refresh_succeeded"], item + assert item["refresh_matches_oracle"], item + assert item["refresh_agreement"], item + assert item["refresh_failure_seen"], item + assert item["refresh_failure_atomic"], item + assert item["refresh_retry_succeeded"], item + assert item["periodic_nondue_failure_rejected"], item + assert item["periodic_refresh_valid"], item + assert item["owner_initialized"], item + assert item["owner_agreement"], item + assert item["owner_step_valid"], item + assert item["owner_contract_valid"], item + assert item["owner_amp_protocol"], item + assert item["failure_seen"], item + assert item["failure_atomic"], item + assert item["retry_succeeded"], item + assert item["presence_mismatch_rejected"], item + assert item["manifest_mismatch_rejected"], item + + +def _subgroup_worker(rank, world, init_file, queue): + try: + dist.init_process_group( + "gloo", + init_method="file://{}".format(init_file), + rank=rank, + world_size=world, + timeout=timedelta(seconds=45), + ) + first_ranks = (0, 2) + second_ranks = (1, 3) + first_group = dist.new_group(ranks=list(first_ranks)) + second_group = dist.new_group(ranks=list(second_ranks)) + if rank in first_ranks: + group_index = 0 + group_ranks = first_ranks + runtime_group = first_group + canonical_grad = torch.tensor([-1.0, -0.5, 0.25, 1.0]) + else: + group_index = 1 + group_ranks = second_ranks + runtime_group = second_group + canonical_grad = torch.tensor([-1.0, -0.9, 0.8, 1.0]) + coordinate = group_ranks.index(rank) + members = ("slot:0", "slot:1") + identity = ProcessGroupIdentity("replica_subgroup:{}".format(group_index), members) + parameter_identity = ParameterIdentity("Subgroup{}.Weight".format(group_index), (4,)) + records = tuple(_replicated(parameter_identity, identity, member) for member in members) + parameter = torch.nn.Parameter(torch.zeros(4)) + optimizer = Gefen([("weight", parameter)], fused=False, factored_v_2d=False) + optimizer.post_sharding( + (ParameterRebinding(parameter, parameter, records[coordinate]),), + manifest=ShardingManifest(records), + codebook_process_group=CodebookProcessGroupBinding( + identity, + members[coordinate], + runtime_group, + torch.device("cpu"), + ), + ) + optimizer._resolve_automatic_period = lambda *args: 4 + parameter.grad = canonical_grad.clone() if coordinate == 0 else torch.tensor([-1.0, -0.2, 0.1, 1.0]) + optimizer.initialize_codebook() + queue.put( + { + "rank": rank, + "group": group_index, + "codebook": optimizer._gefen_codebook.tolist(), + } + ) + except Exception as exc: + queue.put({"rank": rank, "error": repr(exc)}) + finally: + if dist.is_available() and dist.is_initialized(): + dist.destroy_process_group() + + +def _run_subgroup_workers(): + context = mp.get_context("spawn") + queue = context.Queue() + fd, init_file = tempfile.mkstemp(prefix="gefen-codebook-subgroups-") + os.close(fd) + os.unlink(init_file) + processes = [ + context.Process( + target=_subgroup_worker, + args=(rank, 4, init_file, queue), + ) + for rank in range(4) + ] + try: + for process in processes: + process.start() + results = [queue.get(timeout=60) for _ in processes] + for process in processes: + process.join(timeout=10) + for process in processes: + if process.is_alive(): + process.terminate() + process.join(timeout=5) + pytest.fail("scoped codebook subgroup worker hung") + assert process.exitcode == 0 + return sorted(results, key=lambda item: item["rank"]) + finally: + for process in processes: + if process.is_alive(): + process.terminate() + process.join(timeout=5) + if os.path.exists(init_file): + os.unlink(init_file) + + +def test_explicit_gloo_subgroups_are_isolated_from_default_world(): + results = _run_subgroup_workers() + + assert all("error" not in item for item in results), results + first = [item["codebook"] for item in results if item["group"] == 0] + second = [item["codebook"] for item in results if item["group"] == 1] + assert first[0] == first[1] + assert second[0] == second[1] + assert first[0] != second[0] + + +def _zero_length_flat_worker(rank, world, init_file, queue): + try: + dist.init_process_group( + "gloo", + init_method="file://{}".format(init_file), + rank=rank, + world_size=world, + timeout=timedelta(seconds=45), + ) + members = tuple("rank:{}".format(index) for index in range(world)) + group = ProcessGroupIdentity("data_parallel", members) + runtime_group = dist.group.WORLD + + # Uneven sharding of a small parameter: member one holds a legal + # zero-length flattened slice bound to a live numel-0 tensor. The + # empty member creates no per-parameter state and its gradient + # presence is excluded from the nonempty-shard consensus, but it must + # still join every scoped collective symmetrically. + identity = ParameterIdentity("ZeroFlat", (4,)) + records = ( + _flat(identity, group, "rank:0", 0, 4), + _flat(identity, group, "rank:1", 4, 0), + ) + parameter = torch.nn.Parameter(torch.zeros(4 if rank == 0 else 0)) + optimizer = Gefen([("zero_flat", parameter)], fused=False, factored_v_2d=False) + _finalize( + optimizer, + parameter, + records[rank], + ShardingManifest(records), + _binding(group, rank, runtime_group), + ) + optimizer._resolve_automatic_period = lambda *args: 4 + nonempty_grad = torch.tensor([-4.0, -1.0, 2.0, 8.0]) + if rank == 0: + parameter.grad = nonempty_grad.clone() + initialized = optimizer.initialize_codebook() + init_oracle = learn_gefen_exact_codebook_from_grad_periods( + grad_periods=(("ZeroFlat", nonempty_grad, 4, nonempty_grad),), + codebook_device=torch.device("cpu"), + num_codebooks=256, + force_endpoints=True, + verbose=False, + compute_mse_logging=False, + use_fused_histogram=False, + ) + init_matches_oracle = torch.equal(optimizer._gefen_codebook, init_oracle) + init_codebooks = [torch.empty_like(optimizer._gefen_codebook) for _ in range(world)] + dist.all_gather(init_codebooks, optimizer._gefen_codebook) + init_agreement = all(torch.equal(item, init_codebooks[0]) for item in init_codebooks[1:]) + + optimizer.step() + step_valid = optimizer._gefen_global_step == 1 and ( + rank == 0 or optimizer.state[parameter] == {"name": "zero_flat"} + ) + + continuation_grad = nonempty_grad.flip(0) + if rank == 0: + parameter.grad = continuation_grad.clone() + refreshed = optimizer.refresh_codebook() + refresh_oracle = learn_gefen_exact_codebook_from_grad_periods( + grad_periods=(("ZeroFlat", continuation_grad, 4, continuation_grad),), + codebook_device=torch.device("cpu"), + num_codebooks=256, + force_endpoints=True, + verbose=False, + compute_mse_logging=False, + use_fused_histogram=False, + ) + refresh_matches_oracle = torch.equal(optimizer._gefen_codebook, refresh_oracle) + refreshed_codebooks = [torch.empty_like(optimizer._gefen_codebook) for _ in range(world)] + dist.all_gather(refreshed_codebooks, optimizer._gefen_codebook) + refresh_agreement = all(torch.equal(item, refreshed_codebooks[0]) for item in refreshed_codebooks[1:]) + empty_state_inert = rank == 0 or optimizer.state[parameter] == {"name": "zero_flat"} + + # A failure on the nonempty member must be observed by the empty + # member through the failure-sync collective, leave both members + # unchanged, and keep the pair retryable. + failure_param = torch.nn.Parameter(torch.zeros(4 if rank == 0 else 0)) + failure_optimizer = Gefen( + [("zero_failure", failure_param)], + fused=False, + factored_v_2d=False, + ) + failure_identity = ParameterIdentity("ZeroFailure", (4,)) + failure_records = ( + _flat(failure_identity, group, "rank:0", 0, 4), + _flat(failure_identity, group, "rank:1", 4, 0), + ) + _finalize( + failure_optimizer, + failure_param, + failure_records[rank], + ShardingManifest(failure_records), + _binding(group, rank, runtime_group), + ) + failure_optimizer._resolve_automatic_period = lambda *args: 4 + if rank == 0: + failure_param.grad = nonempty_grad.clone() + original_exact_dp = gefen_module.quantization_module.exact_dp + if rank == 0: + + def fail_exact_dp(*args, **kwargs): + raise RuntimeError("rank-local exact-DP failure") + + gefen_module.quantization_module.exact_dp = fail_exact_dp + try: + failure_optimizer.initialize_codebook() + failure_seen = False + except RuntimeError as exc: + failure_seen = "exact-DP" in str(exc) + finally: + gefen_module.quantization_module.exact_dp = original_exact_dp + failure_atomic = ( + failure_optimizer._gefen_codebook is None + and failure_optimizer._gefen_global_step == 0 + and failure_optimizer.state[failure_param] == {"name": "zero_failure"} + ) + retry_succeeded = failure_optimizer.initialize_codebook() + retry_codebooks = [torch.empty_like(failure_optimizer._gefen_codebook) for _ in range(world)] + dist.all_gather(retry_codebooks, failure_optimizer._gefen_codebook) + retry_agreement = all(torch.equal(item, retry_codebooks[0]) for item in retry_codebooks[1:]) + + queue.put( + { + "rank": rank, + "initialized": initialized, + "init_matches_oracle": init_matches_oracle, + "init_agreement": init_agreement, + "step_valid": step_valid, + "refreshed": refreshed, + "refresh_matches_oracle": refresh_matches_oracle, + "refresh_agreement": refresh_agreement, + "empty_state_inert": empty_state_inert, + "failure_seen": failure_seen, + "failure_atomic": failure_atomic, + "retry_succeeded": retry_succeeded, + "retry_agreement": retry_agreement, + } + ) + except Exception as exc: + queue.put({"rank": rank, "error": repr(exc)}) + finally: + if dist.is_available() and dist.is_initialized(): + dist.destroy_process_group() + + +def _run_zero_length_flat_workers(world=2): + context = mp.get_context("spawn") + queue = context.Queue() + fd, init_file = tempfile.mkstemp(prefix="gefen-codebook-zero-flat-") + os.close(fd) + os.unlink(init_file) + processes = [ + context.Process( + target=_zero_length_flat_worker, + args=(rank, world, init_file, queue), + ) + for rank in range(world) + ] + try: + for process in processes: + process.start() + results = [queue.get(timeout=60) for _ in processes] + for process in processes: + process.join(timeout=10) + for process in processes: + if process.is_alive(): + process.terminate() + process.join(timeout=5) + pytest.fail("zero-length flattened shard worker hung") + assert process.exitcode == 0 + return sorted(results, key=lambda item: item["rank"]) + finally: + for process in processes: + if process.is_alive(): + process.terminate() + process.join(timeout=5) + if os.path.exists(init_file): + os.unlink(init_file) + + +def test_zero_length_flattened_shard_member_joins_every_scoped_collective(): + results = _run_zero_length_flat_workers() + + assert all("error" not in item for item in results), results + for item in results: + assert item["initialized"], item + assert item["init_matches_oracle"], item + assert item["init_agreement"], item + assert item["step_valid"], item + assert item["refreshed"], item + assert item["refresh_matches_oracle"], item + assert item["refresh_agreement"], item + assert item["empty_state_inert"], item + assert item["failure_seen"], item + assert item["failure_atomic"], item + assert item["retry_succeeded"], item + assert item["retry_agreement"], item + + +def _nccl_empty_owner_worker(rank, world, init_file, queue): + try: + torch.cuda.set_device(rank) + dist.init_process_group( + "nccl", + init_method="file://{}".format(init_file), + rank=rank, + world_size=world, + timeout=timedelta(seconds=60), + ) + members = tuple("rank:{}".format(index) for index in range(world)) + group = ProcessGroupIdentity("nccl_owner", members) + identity = ParameterIdentity("Matrix", (2, 2)) + records = tuple(_whole_owner(identity, group, member, "rank:0") for member in members) + source = torch.nn.Parameter(torch.ones(2, 2, device=torch.device("cuda", rank))) + optimizer = GefenMuon([("matrix", source)], fused=False) + optimizer.post_sharding( + ( + ParameterRebinding( + source, + source if rank == 0 else None, + records[rank], + ), + ), + manifest=ShardingManifest(records), + codebook_process_group=CodebookProcessGroupBinding( + group, + members[rank], + dist.group.WORLD, + torch.device("cuda", rank), + ), + ) + if rank == 0: + source.grad = torch.tensor([[1.0, -2.0], [3.0, -4.0]], device=source.device) + initialized = optimizer.initialize_codebook() + local = optimizer._gefen_codebook.to(torch.device("cuda", rank)) + gathered = [torch.empty_like(local) for _ in range(world)] + dist.all_gather(gathered, local) + agreement = all(torch.equal(item, gathered[0]) for item in gathered[1:]) + # The codebook binding, not ambient CUDA state, owns the device used by + # every scoped control collective. Deliberately make current_device + # disagree before the step preamble exercises its shared failure flag. + torch.cuda.set_device(1 - rank) + optimizer.step() + queue.put( + { + "rank": rank, + "initialized": initialized, + "agreement": agreement, + "empty_nonowner": rank == 0 + or ( + not optimizer.param_groups[0]["params"] + and not optimizer.state + and optimizer._gefen_codebook.device.type == "cpu" + ), + "step": optimizer._gefen_global_step, + "ambient_device_differed": torch.cuda.current_device() != rank, + } + ) + except Exception as exc: + queue.put({"rank": rank, "error": repr(exc)}) + finally: + if dist.is_available() and dist.is_initialized(): + dist.destroy_process_group() + + +@pytest.mark.skipif( + not dist.is_available() or not dist.is_nccl_available() or torch.cuda.device_count() < 2, + reason="requires NCCL and two CUDA devices", +) +def test_nccl_scope_uses_explicit_collective_device_with_empty_nonowner(): + context = mp.get_context("spawn") + queue = context.Queue() + fd, init_file = tempfile.mkstemp(prefix="gefen-codebook-nccl-") + os.close(fd) + os.unlink(init_file) + processes = [ + context.Process( + target=_nccl_empty_owner_worker, + args=(rank, 2, init_file, queue), + ) + for rank in range(2) + ] + try: + for process in processes: + process.start() + results = [queue.get(timeout=90) for _ in processes] + for process in processes: + process.join(timeout=15) + for process in processes: + if process.is_alive(): + process.terminate() + process.join(timeout=5) + pytest.fail("NCCL scoped codebook worker hung") + assert process.exitcode == 0 + results.sort(key=lambda item: item["rank"]) + assert all("error" not in item for item in results), results + assert all(item["initialized"] for item in results) + assert all(item["agreement"] for item in results) + assert all(item["empty_nonowner"] for item in results) + assert all(item["step"] == 1 for item in results) + assert all(item["ambient_device_differed"] for item in results) + finally: + for process in processes: + if process.is_alive(): + process.terminate() + process.join(timeout=5) + if os.path.exists(init_file): + os.unlink(init_file) + + +def _closure_preamble_worker( + rank, world, init_file, queue, optimizer_kind, failure_mode +): + # Optimizer.step runs the closure before any later scope collectives. A rank-local closure failure must raise on every scope member together instead of leaving the failing rank to exit step() while the peer enters the scoped operation-header / synchronization collectives and hangs. + try: + dist.init_process_group( + "gloo", + init_method="file://{}".format(init_file), + rank=rank, + world_size=world, + timeout=timedelta(seconds=45), + ) + members = tuple("rank:{}".format(index) for index in range(world)) + group = ProcessGroupIdentity("data_parallel", members) + runtime_group = dist.group.WORLD + + matrix = torch.nn.Parameter(torch.zeros(2, 2)) + if optimizer_kind == "muon": + optimizer = GefenMuon([("matrix", matrix)], fused=False) + elif optimizer_kind == "gefen": + optimizer = Gefen( + [("matrix", matrix)], fused=False, factored_v_2d=False + ) + else: + raise AssertionError("unknown closure-preamble optimizer kind") + identity = ParameterIdentity("Matrix", (2, 2)) + records = tuple(_replicated(identity, group, member) for member in members) + _finalize( + optimizer, + matrix, + records[rank], + ShardingManifest(records), + _binding(group, rank, runtime_group), + ) + matrix.grad = torch.tensor([[1.0, -2.0], [3.0, -4.0]]) + rogue = torch.nn.Parameter(torch.ones(2, 2)) + rogue_before = rogue.detach().clone() + if rank == 0 and failure_mode == "replace_before": + optimizer.param_groups[0]["params"][0] = rogue + rogue.grad = torch.ones_like(rogue) + + def closure(): + if rank == 0 and failure_mode == "raise": + raise RuntimeError("closure boom on rank:0") + if rank == 0 and failure_mode == "replace": + optimizer.param_groups[0]["params"][0] = rogue + rogue.grad = torch.ones_like(rogue) + return torch.tensor(1.0) + + try: + optimizer.step(closure) + message = None + except RuntimeError as exc: + message = str(exc) + # The synchronized failure must leave the step fully un-run on both ranks. + untouched = ( + optimizer._gefen_global_step == 0 + and optimizer._gefen_codebook is None + and optimizer.state[matrix] == {"name": "matrix"} + and torch.equal(matrix, torch.zeros(2, 2)) + and torch.equal(rogue, rogue_before) + ) + queue.put({"rank": rank, "message": message, "untouched": untouched}) + except Exception as exc: + queue.put({"rank": rank, "error": repr(exc)}) + finally: + if dist.is_available() and dist.is_initialized(): + dist.destroy_process_group() + + +def _run_closure_preamble_workers( + optimizer_kind="muon", failure_mode="raise", world=2 +): + context = mp.get_context("spawn") + queue = context.Queue() + fd, init_file = tempfile.mkstemp(prefix="gefen-codebook-closure-preamble-") + os.close(fd) + os.unlink(init_file) + processes = [ + context.Process( + target=_closure_preamble_worker, + args=(rank, world, init_file, queue, optimizer_kind, failure_mode), + ) + for rank in range(world) + ] + try: + for process in processes: + process.start() + results = [] + try: + for _ in processes: + results.append(queue.get(timeout=120)) + except Exception: + pass + for process in processes: + process.join(timeout=10) + for process in processes: + if process.is_alive(): + process.terminate() + process.join(timeout=5) + pytest.fail("closure-preamble worker hung") + assert process.exitcode == 0 + return sorted(results, key=lambda item: item["rank"]) + finally: + for process in processes: + if process.is_alive(): + process.terminate() + process.join(timeout=5) + if os.path.exists(init_file): + os.unlink(init_file) + + +@pytest.mark.skipif( + not dist.is_available() or not dist.is_gloo_available(), + reason="scoped closure-preamble coverage requires Gloo", +) +def test_scoped_step_closure_failure_raises_symmetrically_across_the_scope(): + results = _run_closure_preamble_workers() + assert len(results) == 2, results + assert all("error" not in item for item in results), results + assert results[0]["message"] is not None and results[1]["message"] is not None, results + assert "closure boom on rank:0" in results[0]["message"] + assert "step preamble failed on another process-group member" in results[1]["message"] + assert all(item["untouched"] for item in results), results + + +@pytest.mark.skipif( + not dist.is_available() or not dist.is_gloo_available(), + reason="plain Gefen scoped closure-preamble coverage requires Gloo", +) +def test_plain_gefen_scoped_step_closure_failure_raises_symmetrically_across_the_scope(): + results = _run_closure_preamble_workers("gefen") + assert len(results) == 2, results + assert all("error" not in item for item in results), results + assert results[0]["message"] is not None and results[1]["message"] is not None, results + assert "closure boom on rank:0" in results[0]["message"] + assert "step preamble failed on another process-group member" in results[1]["message"] + assert all(item["untouched"] for item in results), results + + +@pytest.mark.skipif( + not dist.is_available() or not dist.is_gloo_available(), + reason="scoped closure layout-mutation coverage requires Gloo", +) +@pytest.mark.parametrize("optimizer_kind", ["gefen", "muon"]) +def test_scoped_step_closure_layout_mutation_raises_symmetrically_across_the_scope( + optimizer_kind, +): + results = _run_closure_preamble_workers(optimizer_kind, "replace") + assert len(results) == 2, results + assert all("error" not in item for item in results), results + assert results[0]["message"] is not None and results[1]["message"] is not None, results + assert "changed outside post_sharding" in results[0]["message"] + assert "gradient preflight failed on another process-group member" in results[1]["message"] + assert all(item["untouched"] for item in results), results + + +@pytest.mark.skipif( + not dist.is_available() or not dist.is_gloo_available(), + reason="scoped step-entry layout-mutation coverage requires Gloo", +) +@pytest.mark.parametrize("optimizer_kind", ["gefen", "muon"]) +def test_scoped_step_entry_layout_mutation_raises_symmetrically_across_the_scope( + optimizer_kind, +): + results = _run_closure_preamble_workers(optimizer_kind, "replace_before") + assert len(results) == 2, results + assert all("error" not in item for item in results), results + assert results[0]["message"] is not None and results[1]["message"] is not None, results + assert "changed outside post_sharding" in results[0]["message"] + assert "step preamble failed on another process-group member" in results[1]["message"] + assert all(item["untouched"] for item in results), results diff --git a/tests/test_hybrid_layout_cache.py b/tests/test_hybrid_layout_cache.py new file mode 100644 index 0000000..f6bbf2a --- /dev/null +++ b/tests/test_hybrid_layout_cache.py @@ -0,0 +1,162 @@ +"""Per-step composite layout-guard cost for a finalized GefenMuonHybrid. + +GefenMuonHybrid.step() calls its finalized-layout guard twice per step, and +each call used to run the full O(params) composite forensic rebuild (routing, +ownership, per-child manifest partitioning, local-binding sort). These tests +pin the contract that a steady-state finalized composite reuses one cached +verdict (an O(local params) identity-token check that folds in each child's +own cached fast-path verdict) instead of rebuilding, while every real +finalized-layout change -- an in-place child ``group['params']`` slot swap, a +replaced composite registry, a closure that tampers between the pre- and +post-closure guards -- is still detected before any state mutation. +""" + +import pytest +import torch + +from gefen import GefenMuonHybrid +from gefen.contracts import ( + LogicalSlice, + ParameterIdentity, + ParameterLayout, + ShardIdentity, + ShardingManifest, +) +from gefen.rebinding import ParameterRebinding + + +def _replicated_shard(fqn, shape): + identity = ParameterIdentity(fqn, shape) + return ShardIdentity( + identity, + ParameterLayout.REPLICATED, + LogicalSlice.full(identity), + ) + + +def _finalized_hybrid(): + matrix = torch.nn.Parameter(torch.randn(2, 2)) + bias = torch.nn.Parameter(torch.randn(4)) + optimizer = GefenMuonHybrid( + [("layer.weight", matrix)], + [("layer.bias", bias)], + lr=1e-3, + fused=False, + sharded_mode="distributed", + ) + matrix_shard = _replicated_shard("Model.Layer.Weight", (2, 2)) + bias_shard = _replicated_shard("Model.Layer.Bias", (4,)) + optimizer.post_sharding( + ( + ParameterRebinding(matrix, matrix, matrix_shard), + ParameterRebinding(bias, bias, bias_shard), + ), + manifest=ShardingManifest((matrix_shard, bias_shard)), + ) + return optimizer, matrix, bias + + +def _count_full_composite_rebuilds(monkeypatch): + # The full composite rebuild is the only step-time caller of + # _gefen_rebinding_children; the cached fast path never touches it. This + # distinguishes a real rebuild in both the pre- and post-cache code. + calls = {"count": 0} + original = GefenMuonHybrid._gefen_rebinding_children + + def counted(self): + calls["count"] += 1 + return original(self) + + monkeypatch.setattr(GefenMuonHybrid, "_gefen_rebinding_children", counted) + return calls + + +def _step_with_grads(optimizer, params): + for parameter in params: + parameter.grad = torch.full_like(parameter, 0.5) + optimizer.step() + + +def test_steady_state_finalized_hybrid_step_does_not_rebuild(monkeypatch): + optimizer, matrix, bias = _finalized_hybrid() + _step_with_grads(optimizer, (matrix, bias)) # warm the cached verdict + + calls = _count_full_composite_rebuilds(monkeypatch) + for _ in range(3): + _step_with_grads(optimizer, (matrix, bias)) + # Pre-cache each step's two guard calls each ran a full rebuild (6 total); + # with the verdict cache a steady-state step rebuilds zero times. + assert calls["count"] == 0 + + +def test_first_finalized_hybrid_step_rebuilds_then_caches(monkeypatch): + optimizer, matrix, bias = _finalized_hybrid() + calls = _count_full_composite_rebuilds(monkeypatch) + _step_with_grads(optimizer, (matrix, bias)) + # The first guard rebuilds once and caches; the second guard in the same + # step reuses the warm verdict. + assert calls["count"] == 1 + assert optimizer._hybrid_layout_forensics_verdict is not None + + +def test_finalized_hybrid_detects_child_param_slot_swap_with_warm_verdict(): + optimizer, matrix, bias = _finalized_hybrid() + _step_with_grads(optimizer, (matrix, bias)) # warm the cached verdict + assert optimizer._hybrid_layout_forensics_verdict is not None + + rogue = torch.nn.Parameter(torch.full((2, 2), 7.0)) + rogue_before = rogue.detach().clone() + optimizer.muon.param_groups[0]["params"][0] = rogue + + with pytest.raises(RuntimeError, match="finalized parameter layout changed"): + _step_with_grads(optimizer, (bias,)) + assert torch.equal(rogue, rogue_before) + + +def test_finalized_hybrid_detects_swapped_child_binding_with_warm_verdict(): + optimizer, matrix, bias = _finalized_hybrid() + _step_with_grads(optimizer, (matrix, bias)) # warm the cached verdict + + # Swap a child's private finalized manifest for a foreign one after + # finalization: the composite fast token captures the child manifest by + # identity, so the swap forces a full rebuild that rejects the layout. + optimizer.muon._gefen_sharding_manifest = optimizer.backup._gefen_sharding_manifest + + with pytest.raises(RuntimeError, match="finalized parameter layout changed"): + _step_with_grads(optimizer, (matrix, bias)) + + +def test_finalized_hybrid_rechecks_layout_after_closure(): + optimizer, matrix, bias = _finalized_hybrid() + _step_with_grads(optimizer, (matrix, bias)) # warm the cached verdict + + rogue = torch.nn.Parameter(torch.full((2, 2), 3.0)) + rogue_before = rogue.detach().clone() + matrix_before = matrix.detach().clone() + + def mutate_layout(): + # Tamper with a public child container between the pre- and + # post-closure guards; the post-closure guard must catch it before any + # child steps. + optimizer.muon.param_groups[0]["params"][0] = rogue + rogue.grad = torch.ones_like(rogue) + bias.grad = torch.ones_like(bias) + return torch.tensor(1.0) + + with pytest.raises(RuntimeError, match="finalized parameter layout changed"): + optimizer.step(mutate_layout) + assert torch.equal(rogue, rogue_before) + assert torch.equal(matrix, matrix_before) + + +def test_composite_registry_replacement_is_detected_with_warm_verdict(): + optimizer, matrix, bias = _finalized_hybrid() + _step_with_grads(optimizer, (matrix, bias)) # warm the cached verdict + + # Replacing the composite ownership registry preserves no container + # identity the fast token recorded, so the next step guard rebuilds and + # rejects. + optimizer._state_param_owner = {} + + with pytest.raises(RuntimeError, match="finalized parameter layout changed"): + _step_with_grads(optimizer, (matrix, bias)) diff --git a/tests/test_hybrid_rebinding.py b/tests/test_hybrid_rebinding.py new file mode 100644 index 0000000..1d91932 --- /dev/null +++ b/tests/test_hybrid_rebinding.py @@ -0,0 +1,549 @@ +"""Focused CPU coverage for atomic Gefen-backed Hybrid rebinding.""" + +import pytest +import torch + +from gefen import GefenMuonHybrid +from gefen.codebook import CodebookProcessGroupBinding +from gefen.contracts import ( + LogicalSlice, + ParameterIdentity, + ParameterLayout, + PlacementKind, + ProcessGroupIdentity, + ShardIdentity, + ShardPlacement, + ShardingManifest, +) +from gefen.rebinding import ParameterRebinding + +from _state_snapshot import assert_deep_state_snapshot, deep_state_snapshot + + +_MEMBER = "rank:0" + + +def _optimizer(*, backup_optimizer="gefen", sharded_mode="distributed"): + matrix = torch.nn.Parameter(torch.arange(4, dtype=torch.float32).reshape(2, 2)) + bias = torch.nn.Parameter(torch.arange(4, dtype=torch.float32)) + optimizer = GefenMuonHybrid( + [("layer.weight", matrix)], + [("layer.bias", bias)], + lr=1e-3, + fused=False, + sharded_mode=sharded_mode, + backup_optimizer=backup_optimizer, + ) + return optimizer, matrix, bias + + +def _ungrouped_replicated(fqn, shape): + identity = ParameterIdentity(fqn, shape) + return ShardIdentity( + identity, + ParameterLayout.REPLICATED, + LogicalSlice.full(identity), + ) + + +def _grouped_replicated(fqn, shape, group, member): + identity = ParameterIdentity(fqn, shape) + coordinate = group.ordered_members.index(member) + return ShardIdentity( + identity, + ParameterLayout.REPLICATED, + LogicalSlice.full(identity), + placements=( + ShardPlacement( + "checkpoint", + PlacementKind.REPLICATE, + coordinate, + len(group.ordered_members), + ), + ), + process_group=group, + local_member=member, + ) + + +def _flat_shards(fqn, shape, group, lengths): + identity = ParameterIdentity(fqn, shape) + offset = 0 + shards = [] + for coordinate, (member, length) in enumerate(zip(group.ordered_members, lengths)): + shards.append( + ShardIdentity( + identity, + ParameterLayout.FLATTENED_ELEMENT_SHARD, + LogicalSlice(offset, length), + placements=( + ShardPlacement( + "checkpoint", + PlacementKind.FLAT_SHARD, + coordinate, + len(group.ordered_members), + ), + ), + process_group=group, + local_member=member, + ) + ) + offset += length + return tuple(shards) + + +def _owner_shards(fqn, shape, group, owner): + identity = ParameterIdentity(fqn, shape) + shards = [] + for coordinate, member in enumerate(group.ordered_members): + shards.append( + ShardIdentity( + identity, + ParameterLayout.WHOLE_PARAMETER_OWNER, + (LogicalSlice.full(identity) if member == owner else LogicalSlice(0, 0)), + placements=( + ShardPlacement( + "checkpoint", + PlacementKind.WHOLE_PARAMETER_OWNER, + coordinate, + len(group.ordered_members), + ), + ), + process_group=group, + local_member=member, + owner=owner, + ) + ) + return tuple(shards) + + +def _snapshot(optimizer): + # Each child is snapshotted deeply (container identities plus bitwise + # clones of per-parameter state, group['params'] entries, and options) + # so a failed composite transaction that partially rebound a child in + # place cannot pass on top-level identity alone. + return { + "children": tuple((child, deep_state_snapshot(child)) for child in optimizer._subopts), + "owner": optimizer._state_param_owner, + # id(param) -> (param, owning child); copy the mapping so an in-place + # remap that keeps the registry object identity is still caught. + "owner_contents": dict(optimizer._state_param_owner), + "finalized": optimizer._hybrid_post_sharding_finalized, + "manifest": optimizer._hybrid_sharding_manifest, + "local": optimizer._hybrid_local_shard_bindings, + "roles": optimizer._hybrid_fqn_roles, + "binding": optimizer._hybrid_codebook_process_group, + "slots": optimizer._hybrid_finalized_slots, + } + + +def _assert_snapshot(optimizer, snapshot): + assert optimizer._state_param_owner is snapshot["owner"] + expected_owner = snapshot["owner_contents"] + assert optimizer._state_param_owner.keys() == expected_owner.keys() + for key, (expected_param, expected_child) in expected_owner.items(): + live_param, live_child = optimizer._state_param_owner[key] + assert live_param is expected_param + assert live_child is expected_child + assert optimizer._hybrid_post_sharding_finalized is snapshot["finalized"] + assert optimizer._hybrid_sharding_manifest is snapshot["manifest"] + assert optimizer._hybrid_local_shard_bindings is snapshot["local"] + assert optimizer._hybrid_fqn_roles is snapshot["roles"] + assert optimizer._hybrid_codebook_process_group is snapshot["binding"] + assert optimizer._hybrid_finalized_slots is snapshot["slots"] + assert len(optimizer._subopts) == len(snapshot["children"]) + for live, (expected_child, expected_snapshot) in zip(optimizer._subopts, snapshot["children"]): + assert live is expected_child + assert_deep_state_snapshot(live, expected_snapshot) + + +@pytest.mark.parametrize("registry", ["_param_names", "_gefen_shard_bindings"]) +def test_deep_snapshot_catches_inplace_child_registry_addition(registry): + # The staged post_sharding transaction rebuilds each child's per-parameter + # registries (``_param_names`` / ``_gefen_shard_bindings``) on a shadow and + # swaps it in, so a regression that instead wrote onto the LIVE registry + # before a later child raised would leave a half-populated dict. The dict + # keeps its object identity, so only value comparison of its contents + # catches the partial publication. + optimizer, matrix, bias = _optimizer() + for child in optimizer._subopts: + live = getattr(child, registry) + snapshot = deep_state_snapshot(child) + # The live registry object identity is preserved by an in-place add, so + # the top-level identity checks stay green; only content comparison sees + # the injected key. + assert getattr(child, registry) is live + injected = torch.nn.Parameter(torch.zeros(1)) + if registry == "_param_names": + live[injected] = "injected" + else: + live[injected] = _ungrouped_replicated("Injected.Param", (1,)) + with pytest.raises(AssertionError): + assert_deep_state_snapshot(child, snapshot) + del live[injected] + # With the injection undone the child is byte-for-byte its snapshot. + assert_deep_state_snapshot(child, snapshot) + + +def test_deep_snapshot_catches_inplace_child_registry_value_change(): + # A value edit that keeps the key set intact (e.g. rebinding a child's shard + # in place onto the live registry) must also be caught by content + # comparison, not just added/removed keys. + optimizer, matrix, bias = _optimizer() + child = optimizer._subopts[0] + (param,) = child._param_names + child._gefen_shard_bindings[param] = _ungrouped_replicated("Seed.Param", (2, 2)) + snapshot = deep_state_snapshot(child) + + child._param_names[param] = "mutated-name" + with pytest.raises(AssertionError): + assert_deep_state_snapshot(child, snapshot) + child._param_names[param] = snapshot["registries"]["_param_names"][0][1] + assert_deep_state_snapshot(child, snapshot) + + child._gefen_shard_bindings[param] = _ungrouped_replicated("Other.Param", (4,)) + with pytest.raises(AssertionError): + assert_deep_state_snapshot(child, snapshot) + + +def test_composite_post_sharding_publishes_children_and_rebuilds_routing(): + optimizer, old_matrix, old_bias = _optimizer() + matrix = torch.nn.Parameter(torch.full((2, 2), 7.0)) + bias = torch.nn.Parameter(torch.full((4,), 11.0)) + group = ProcessGroupIdentity("checkpoint", (_MEMBER,)) + binding = CodebookProcessGroupBinding( + group, + _MEMBER, + None, + torch.device("cpu"), + ) + matrix_shard = _grouped_replicated( + "Model.Layer.Weight", + (2, 2), + group, + _MEMBER, + ) + bias_shard = _grouped_replicated( + "Model.Layer.Bias", + (4,), + group, + _MEMBER, + ) + manifest = ShardingManifest((matrix_shard, bias_shard)) + + optimizer.post_sharding( + ( + ParameterRebinding(old_bias, bias, bias_shard), + ParameterRebinding(old_matrix, matrix, matrix_shard), + ), + manifest=manifest, + codebook_process_group=binding, + ) + + assert optimizer._canonical_identity_ready() + assert optimizer._codebook_scope_ready() + assert optimizer.sharding_manifest() is manifest + assert optimizer.codebook_process_group_binding() is binding + assert optimizer.muon.codebook_process_group_binding() is binding + assert optimizer.backup.codebook_process_group_binding() is binding + assert optimizer._hybrid_fqn_roles == ( + ("Model.Layer.Bias", "backup"), + ("Model.Layer.Weight", "muon"), + ) + assert optimizer.parameter_routing() == optimizer._hybrid_fqn_roles + assert optimizer.muon.sharding_manifest().shards == (matrix_shard,) + assert optimizer.backup.sharding_manifest().shards == (bias_shard,) + assert tuple(shard.parameter.fqn for _, shard in optimizer.shard_bindings()) == ( + "Model.Layer.Bias", + "Model.Layer.Weight", + ) + assert optimizer.parameter_identity(matrix) == matrix_shard.parameter + assert optimizer.shard_identity(bias) == bias_shard + with pytest.raises(KeyError, match="no finalized"): + optimizer.shard_identity(old_matrix) + + optimizer.state[bias]["sentinel"] = 13 + optimizer.state[matrix]["sentinel"] = 17 + assert optimizer.backup.state[bias]["sentinel"] == 13 + assert optimizer.muon.state[matrix]["sentinel"] == 17 + assert set(optimizer._state_param_owner) == {id(matrix), id(bias)} + + +def test_composite_supports_muon_nonowner_and_flat_backup_without_binding(): + optimizer, old_matrix, old_bias = _optimizer() + group = ProcessGroupIdentity("data_parallel", ("rank:0", "rank:1")) + matrix_shards = _owner_shards( + "Model.Layer.Weight", + (2, 2), + group, + owner="rank:1", + ) + bias_shards = _flat_shards( + "Model.Layer.Bias", + (4,), + group, + lengths=(2, 2), + ) + local_matrix = matrix_shards[0] + local_bias = bias_shards[0] + bias = torch.nn.Parameter(torch.full((2,), 5.0)) + manifest = ShardingManifest(matrix_shards + bias_shards) + + optimizer.post_sharding( + ( + ParameterRebinding(old_matrix, None, local_matrix), + ParameterRebinding(old_bias, bias, local_bias), + ), + manifest=manifest, + ) + + assert optimizer._canonical_identity_ready() + assert not optimizer._codebook_scope_ready() + assert optimizer.muon.param_groups[0]["params"] == [] + assert optimizer.muon.state == {} + assert optimizer.backup.param_groups[0]["params"] == [bias] + assert optimizer.codebook_process_group_binding() is None + assert optimizer.shard_bindings() == ( + (bias, local_bias), + (None, local_matrix), + ) + assert set(optimizer._state_param_owner) == {id(bias)} + + +def test_second_child_staging_failure_leaves_both_children_unchanged(): + optimizer, old_matrix, old_bias = _optimizer() + matrix = torch.nn.Parameter(torch.full((2, 2), 3.0)) + invalid_bias = torch.nn.Parameter(torch.full((3,), 4.0)) + matrix_shard = _ungrouped_replicated("Model.Layer.Weight", (2, 2)) + bias_shard = _ungrouped_replicated("Model.Layer.Bias", (4,)) + manifest = ShardingManifest((matrix_shard, bias_shard)) + snapshot = _snapshot(optimizer) + matrix_before = matrix.detach().clone() + bias_before = invalid_bias.detach().clone() + + with pytest.raises(ValueError, match="complete parameter storage"): + optimizer.post_sharding( + ( + ParameterRebinding(old_matrix, matrix, matrix_shard), + ParameterRebinding(old_bias, invalid_bias, bias_shard), + ), + manifest=manifest, + ) + + _assert_snapshot(optimizer, snapshot) + assert torch.equal(matrix, matrix_before) + assert torch.equal(invalid_bias, bias_before) + assert not optimizer._canonical_identity_ready() + + +def test_external_child_parameter_replacement_can_be_finalized_atomically(): + optimizer, old_matrix, old_bias = _optimizer() + matrix = torch.nn.Parameter(torch.full((2, 2), 3.0)) + bias = torch.nn.Parameter(torch.full((4,), 4.0)) + optimizer.muon.param_groups[0]["params"][0] = matrix + optimizer.backup.param_groups[0]["params"][0] = bias + matrix_shard = _ungrouped_replicated("Model.Layer.Weight", (2, 2)) + bias_shard = _ungrouped_replicated("Model.Layer.Bias", (4,)) + + optimizer.post_sharding( + ( + ParameterRebinding(old_matrix, matrix, matrix_shard), + ParameterRebinding(old_bias, bias, bias_shard), + ), + manifest=ShardingManifest((matrix_shard, bias_shard)), + ) + + assert optimizer._canonical_identity_ready() + assert optimizer.muon.param_groups[0]["params"] == [matrix] + assert optimizer.backup.param_groups[0]["params"] == [bias] + assert set(optimizer._state_param_owner) == {id(matrix), id(bias)} + + +def test_cross_child_target_storage_overlap_is_rejected_atomically(): + optimizer, old_matrix, old_bias = _optimizer() + storage = torch.arange(8, dtype=torch.float32) + matrix = torch.nn.Parameter(storage[:4].view(2, 2)) + bias = torch.nn.Parameter(storage[2:6]) + assert matrix.untyped_storage().data_ptr() == bias.untyped_storage().data_ptr() + matrix_shard = _ungrouped_replicated("Model.Layer.Weight", (2, 2)) + bias_shard = _ungrouped_replicated("Model.Layer.Bias", (4,)) + snapshot = _snapshot(optimizer) + + with pytest.raises(ValueError, match="storage ranges must not overlap"): + optimizer.post_sharding( + ( + ParameterRebinding(old_matrix, matrix, matrix_shard), + ParameterRebinding(old_bias, bias, bias_shard), + ), + manifest=ShardingManifest((matrix_shard, bias_shard)), + ) + + _assert_snapshot(optimizer, snapshot) + + +def test_global_source_target_and_fqn_guards_run_before_publication(): + optimizer, old_matrix, old_bias = _optimizer() + matrix = torch.nn.Parameter(torch.ones(2, 2)) + bias = torch.nn.Parameter(torch.ones(4)) + matrix_shard = _ungrouped_replicated("Model.Shared", (2, 2)) + duplicate_fqn = _ungrouped_replicated("Model.Shared", (4,)) + snapshot = _snapshot(optimizer) + + with pytest.raises(ValueError, match="FQNs must be unique"): + optimizer.post_sharding( + ( + ParameterRebinding(old_matrix, matrix, matrix_shard), + ParameterRebinding(old_bias, bias, duplicate_fqn), + ), + manifest=ShardingManifest((matrix_shard,)), + ) + _assert_snapshot(optimizer, snapshot) + + matrix_shard = _ungrouped_replicated("Model.Layer.Weight", (2, 2)) + bias_shard = _ungrouped_replicated("Model.Layer.Bias", (4,)) + with pytest.raises(ValueError, match="source tensors must be unique"): + optimizer.post_sharding( + ( + ParameterRebinding(old_matrix, matrix, matrix_shard), + ParameterRebinding(old_matrix, bias, bias_shard), + ), + manifest=ShardingManifest((matrix_shard, bias_shard)), + ) + _assert_snapshot(optimizer, snapshot) + + with pytest.raises(ValueError, match="cannot steal"): + optimizer.post_sharding( + ( + ParameterRebinding(old_matrix, old_bias, matrix_shard), + ParameterRebinding(old_bias, bias, bias_shard), + ), + manifest=ShardingManifest((matrix_shard, bias_shard)), + ) + _assert_snapshot(optimizer, snapshot) + + +def test_shared_binding_must_match_every_child_shard_before_staging(): + optimizer, old_matrix, old_bias = _optimizer() + shard_group = ProcessGroupIdentity("shards", (_MEMBER,)) + binding_group = ProcessGroupIdentity("checkpoint", (_MEMBER,)) + binding = CodebookProcessGroupBinding( + binding_group, + _MEMBER, + None, + torch.device("cpu"), + ) + matrix_shard = _grouped_replicated( + "Model.Layer.Weight", + (2, 2), + shard_group, + _MEMBER, + ) + bias_shard = _grouped_replicated( + "Model.Layer.Bias", + (4,), + shard_group, + _MEMBER, + ) + snapshot = _snapshot(optimizer) + + with pytest.raises(ValueError, match="shared codebook"): + optimizer.post_sharding( + ( + ParameterRebinding( + old_matrix, + torch.nn.Parameter(torch.ones(2, 2)), + matrix_shard, + ), + ParameterRebinding( + old_bias, + torch.nn.Parameter(torch.ones(4)), + bias_shard, + ), + ), + manifest=ShardingManifest((matrix_shard, bias_shard)), + codebook_process_group=binding, + ) + + _assert_snapshot(optimizer, snapshot) + + +@pytest.mark.parametrize("role", ["muon", "backup"]) +def test_one_child_hybrids_support_convenience_rebinding(role): + if role == "muon": + old = torch.nn.Parameter(torch.ones(2, 2)) + new = torch.nn.Parameter(torch.full((2, 2), 2.0)) + optimizer = GefenMuonHybrid( + [("layer.weight", old)], + [], + lr=1e-3, + fused=False, + ) + identity = ParameterIdentity("Model.Layer.Weight", (2, 2)) + else: + old = torch.nn.Parameter(torch.ones(4)) + new = torch.nn.Parameter(torch.full((4,), 2.0)) + optimizer = GefenMuonHybrid( + [], + [("layer.bias", old)], + lr=1e-3, + fused=False, + ) + identity = ParameterIdentity("Model.Layer.Bias", (4,)) + + optimizer.rebind_parameter(old, new, identity=identity) + + assert optimizer._canonical_identity_ready() + assert optimizer.parameter_identity(new) == identity + assert optimizer.shard_bindings() == ((new, optimizer.shard_identity(new)),) + assert set(optimizer._state_param_owner) == {id(new)} + + +def test_adamw_backup_rejects_before_any_child_or_hybrid_mutation(): + optimizer, old_matrix, old_bias = _optimizer(backup_optimizer="adamw") + matrix_shard = _ungrouped_replicated("Model.Layer.Weight", (2, 2)) + bias_shard = _ungrouped_replicated("Model.Layer.Bias", (4,)) + snapshot = _snapshot(optimizer) + + with pytest.raises(NotImplementedError, match="AdamW"): + optimizer.post_sharding( + ( + ParameterRebinding( + old_matrix, + torch.nn.Parameter(torch.ones(2, 2)), + matrix_shard, + ), + ParameterRebinding( + old_bias, + torch.nn.Parameter(torch.ones(4)), + bias_shard, + ), + ), + manifest=ShardingManifest((matrix_shard, bias_shard)), + ) + + _assert_snapshot(optimizer, snapshot) + assert not optimizer._canonical_identity_ready() + + +def test_composite_guard_fails_closed_after_owner_metadata_corruption(): + optimizer, old_matrix, old_bias = _optimizer() + matrix = torch.nn.Parameter(torch.ones(2, 2)) + bias = torch.nn.Parameter(torch.ones(4)) + matrix_shard = _ungrouped_replicated("Model.Layer.Weight", (2, 2)) + bias_shard = _ungrouped_replicated("Model.Layer.Bias", (4,)) + optimizer.post_sharding( + ( + ParameterRebinding(old_matrix, matrix, matrix_shard), + ParameterRebinding(old_bias, bias, bias_shard), + ), + manifest=ShardingManifest((matrix_shard, bias_shard)), + ) + assert optimizer._canonical_identity_ready() + + optimizer._state_param_owner = {} + + assert not optimizer._canonical_identity_ready() + with pytest.raises(RuntimeError, match="changed outside post_sharding"): + optimizer.sharding_manifest() + with pytest.raises(RuntimeError, match="changed outside post_sharding"): + _ = optimizer.state diff --git a/tests/test_hybrid_scoped_failure_protocol.py b/tests/test_hybrid_scoped_failure_protocol.py new file mode 100644 index 0000000..68e59a6 --- /dev/null +++ b/tests/test_hybrid_scoped_failure_protocol.py @@ -0,0 +1,727 @@ +"""Scoped failure-protocol coverage for GefenMuonHybrid.step. + +GradScaler attaches found_inf/grad_scale to the composite (never to the +children), and the hybrid's structural gradient preflight runs before either +child's scoped step collectives. Under the one multi-member codebook binding +shared by both children, both decisions must therefore be group decisions: + +* a rank-local AMP overflow must not skip the children on one member while the + peers enter the children's scoped step collectives -- found_inf/grad_scale + agreement is validated collectively and an overflow skip is entered/exited + symmetrically on every member; +* a rank-local structural preflight failure must raise on every scope member + together via the children's synchronized failure protocol instead of + stranding peers inside a child's scope header all_gather. + +Single-process behavior (no binding, and the collective-free one-member +binding) keeps the local semantics. +""" + +from datetime import timedelta +import multiprocessing as mp +import os +import queue as queue_module +import tempfile +import traceback + +import pytest +import torch +import torch.distributed as dist + +from gefen import GefenMuonHybrid +from gefen.codebook import CodebookProcessGroupBinding +from gefen.contracts import ( + LogicalSlice, + ParameterIdentity, + ParameterLayout, + PlacementKind, + ProcessGroupIdentity, + ShardIdentity, + ShardPlacement, + ShardingManifest, +) +from gefen.rebinding import ParameterRebinding + + +_WORLD = 2 +_MUON_FQN = "model.matrix" +_BACKUP_FQN = "model.vector" +_MUON_OWNER = "rank:1" +_BACKUP_LENGTHS = (3, 5) + + +def _members(): + return tuple("rank:{}".format(rank) for rank in range(_WORLD)) + + +def _owner_shards(identity, group, owner): + return tuple( + ShardIdentity( + identity, + ParameterLayout.WHOLE_PARAMETER_OWNER, + (LogicalSlice.full(identity) if member == owner else LogicalSlice(0, 0)), + placements=( + ShardPlacement( + "dp", + PlacementKind.WHOLE_PARAMETER_OWNER, + coordinate, + len(group.ordered_members), + ), + ), + process_group=group, + local_member=member, + owner=owner, + ) + for coordinate, member in enumerate(group.ordered_members) + ) + + +def _flat_shards(identity, group, lengths): + offset = 0 + shards = [] + for coordinate, (member, length) in enumerate(zip(group.ordered_members, lengths)): + shards.append( + ShardIdentity( + identity, + ParameterLayout.FLATTENED_ELEMENT_SHARD, + LogicalSlice(offset, length), + placements=( + ShardPlacement( + "dp", + PlacementKind.FLAT_SHARD, + coordinate, + len(group.ordered_members), + ), + ), + process_group=group, + local_member=member, + ) + ) + offset += length + return tuple(shards) + + +def _muon_initial(): + return torch.linspace(-0.3, 0.2, 6, dtype=torch.float32).reshape(3, 2) + + +def _backup_initial(): + return torch.linspace(-0.4, 0.3, 8, dtype=torch.float32) + + +def _muon_gradient(): + return torch.tensor( + [[0.4, -0.7], [1.1, -1.3], [1.7, -1.9]], + dtype=torch.float32, + ) + + +def _backup_gradient(): + return torch.tensor( + [0.25, -0.5, 0.75, -1.0, 1.25, -1.5, 1.75, -2.0], + dtype=torch.float32, + ) + + +def _codebook(): + return torch.linspace(-1.0, 1.0, 256, dtype=torch.float32) + + +def _quantized_period_one(momentum): + flat = momentum.reshape(-1) + indices = torch.where( + torch.signbit(flat), + torch.zeros(flat.numel(), dtype=torch.uint8), + torch.full((flat.numel(),), 255, dtype=torch.uint8), + ) + return indices.reshape(-1, 1), flat.abs().reshape(-1, 1).clone() + + +def _hybrid_kwargs(): + return dict( + lr=2.5e-3, + muon_lr=3.0e-3, + backup_lr=2.5e-3, + weight_decay=0.03, + muon_weight_decay=0.02, + backup_weight_decay=0.03, + backup_optimizer="gefen", + backup_1d_period_one=True, + betas=(0.8, 0.97), + eps=2.0e-8, + fused=False, + momentum=0.85, + nesterov=False, + ns_steps=2, + deterministic=True, + normuon=True, + normuon_beta2=0.9, + normuon_eps=3.0e-8, + ) + + +def _make_scoped_hybrid(rank, group): + muon_identity = ParameterIdentity(_MUON_FQN, (3, 2)) + backup_identity = ParameterIdentity(_BACKUP_FQN, (8,)) + old_muon = torch.nn.Parameter(_muon_initial().clone()) + old_backup = torch.nn.Parameter(_backup_initial().clone()) + optimizer = GefenMuonHybrid( + [("matrix", old_muon)], + [("vector", old_backup)], + sharded_mode="distributed", + **_hybrid_kwargs(), + ) + muon_shards = _owner_shards(muon_identity, group, _MUON_OWNER) + backup_shards = _flat_shards(backup_identity, group, _BACKUP_LENGTHS) + muon_local = old_muon if _members()[rank] == _MUON_OWNER else None + backup_shard = backup_shards[rank] + start = backup_shard.logical_slice.flat_offset + stop = start + backup_shard.logical_slice.length + backup_local = torch.nn.Parameter(_backup_initial()[start:stop].clone()) + binding = CodebookProcessGroupBinding( + group, + _members()[rank], + dist.group.WORLD, + torch.device("cpu"), + ) + optimizer.post_sharding( + ( + ParameterRebinding(old_muon, muon_local, muon_shards[rank]), + ParameterRebinding(old_backup, backup_local, backup_shard), + ), + manifest=ShardingManifest(muon_shards + backup_shards), + codebook_process_group=binding, + ) + return optimizer, muon_local, backup_local, backup_shard + + +def _seed_hybrid_state(optimizer, muon_parameter, backup_parameter, backup_shard): + for child in (optimizer.muon, optimizer.backup): + child._gefen_global_step = 13 + child._gefen_codebook = _codebook() + if muon_parameter is not None: + momentum = torch.tensor( + [[-0.75, 0.5], [1.25, -1.5], [2.0, -2.5]], + dtype=torch.float32, + ) + indices, magnitudes = _quantized_period_one(momentum) + optimizer.muon.state[muon_parameter].update( + { + "automatic_period": 1, + "step": 11, + "m_codebook": indices, + "m_magnitude": magnitudes, + "normuon_v": torch.tensor([[0.5], [1.5], [2.5]], dtype=torch.float32), + "normuon_step": 10, + } + ) + start = backup_shard.logical_slice.flat_offset + stop = start + backup_shard.logical_slice.length + momentum = torch.tensor( + [-0.25, 0.5, -0.75, 1.0, -1.25, 1.5, -1.75, 2.0], + dtype=torch.float32, + )[start:stop] + indices, magnitudes = _quantized_period_one(momentum) + optimizer.backup.state[backup_parameter].update( + { + "automatic_period": 1, + "step": 11, + "m_codebook": indices, + "m_magnitude": magnitudes, + "vmean": torch.linspace(0.2, 0.9, 8, dtype=torch.float32)[start:stop].reshape(-1, 1).clone(), + "vmean_step": 10, + } + ) + + +def _bits_equal(left, right): + return ( + torch.is_tensor(left) + and torch.is_tensor(right) + and left.dtype == right.dtype + and tuple(left.shape) == tuple(right.shape) + and torch.equal( + left.detach().contiguous().view(torch.uint8), + right.detach().contiguous().view(torch.uint8), + ) + ) + + +def _set_local_gradients(muon_parameter, backup_parameter, backup_shard, scale=1.0): + if muon_parameter is not None: + muon_parameter.grad = _muon_gradient() * scale + start = backup_shard.logical_slice.flat_offset + stop = start + backup_shard.logical_slice.length + backup_parameter.grad = _backup_gradient()[start:stop] * scale + + +def _untouched(optimizer, muon_parameter, backup_parameter): + return ( + optimizer.muon._gefen_global_step == 0 + and optimizer.backup._gefen_global_step == 0 + and optimizer.muon._gefen_codebook is None + and optimizer.backup._gefen_codebook is None + and (muon_parameter is None or _bits_equal(muon_parameter, _muon_initial())) + and ( + muon_parameter is None + or optimizer.muon.state[muon_parameter] == {"name": "matrix"} + ) + and optimizer.backup.state[backup_parameter] == {"name": "vector"} + ) + + +def _amp_divergent_overflow_result(rank, group): + optimizer, muon_parameter, backup_parameter, backup_shard = _make_scoped_hybrid(rank, group) + _set_local_gradients(muon_parameter, backup_parameter, backup_shard) + grad_before = backup_parameter.grad.detach().clone() + optimizer.found_inf = torch.tensor(float(rank == 0)) + optimizer.grad_scale = torch.tensor(8.0) + try: + optimizer.step() + message = None + except RuntimeError as exc: + message = str(exc) + return { + "message": message, + "untouched": _untouched(optimizer, muon_parameter, backup_parameter), + "grads_untouched": _bits_equal(backup_parameter.grad, grad_before), + } + + +def _amp_group_wide_overflow_result(rank, group): + optimizer, muon_parameter, backup_parameter, backup_shard = _make_scoped_hybrid(rank, group) + _set_local_gradients(muon_parameter, backup_parameter, backup_shard) + grad_before = backup_parameter.grad.detach().clone() + optimizer.found_inf = torch.tensor(1.0) + optimizer.grad_scale = torch.tensor(8.0) + post_hook_calls = [] + optimizer.register_step_post_hook( + lambda _optimizer, _args, _kwargs: post_hook_calls.append(True) + ) + try: + optimizer.step() + skipped = True + except RuntimeError: + skipped = False + return { + "skipped": skipped, + "post_hook_fired": len(post_hook_calls) == 1, + "untouched": _untouched(optimizer, muon_parameter, backup_parameter), + "grads_untouched": _bits_equal(backup_parameter.grad, grad_before), + } + + +def _amp_agreeing_finite_result(rank, group): + target, target_muon, target_backup, target_shard = _make_scoped_hybrid(rank, group) + reference, reference_muon, reference_backup, reference_shard = _make_scoped_hybrid(rank, group) + _seed_hybrid_state(target, target_muon, target_backup, target_shard) + _seed_hybrid_state(reference, reference_muon, reference_backup, reference_shard) + _set_local_gradients(target_muon, target_backup, target_shard, scale=8.0) + _set_local_gradients(reference_muon, reference_backup, reference_shard) + target.found_inf = torch.tensor(0.0) + target.grad_scale = torch.tensor(8.0) + target.step() + reference.step() + return { + "stepped": ( + target.muon._gefen_global_step == 14 + and target.backup._gefen_global_step == 14 + ), + "unscaled_once_exactly": ( + (target_muon is None or _bits_equal(target_muon, reference_muon)) + and _bits_equal(target_backup, reference_backup) + ), + } + + +def _preflight_divergent_result(rank, group): + optimizer, muon_parameter, backup_parameter, backup_shard = _make_scoped_hybrid(rank, group) + _set_local_gradients(muon_parameter, backup_parameter, backup_shard) + if rank == 0: + backup_parameter.grad = torch.sparse_coo_tensor( + torch.tensor([[0]]), + torch.tensor([1.0]), + size=(backup_shard.logical_slice.length,), + ) + try: + optimizer.step() + message = None + except RuntimeError as exc: + message = str(exc) + return { + "message": message, + "untouched": _untouched(optimizer, muon_parameter, backup_parameter), + } + + +def _closure_divergent_result(rank, group): + # The closure runs BEFORE any scope synchronization. A rank-local closure + # failure must raise on every scope member together instead of leaving the + # failing rank to exit step() while the peer enters the composite preflight + # synchronize / a child's scoped step collectives and hangs. + optimizer, muon_parameter, backup_parameter, backup_shard = _make_scoped_hybrid(rank, group) + _set_local_gradients(muon_parameter, backup_parameter, backup_shard) + + def closure(): + if rank == 0: + raise RuntimeError("closure boom on rank:0") + return torch.tensor(1.0) + + try: + optimizer.step(closure) + message = None + except RuntimeError as exc: + message = str(exc) + return { + "message": message, + "untouched": _untouched(optimizer, muon_parameter, backup_parameter), + } + + +def _closure_layout_divergent_result(rank, group): + optimizer, muon_parameter, backup_parameter, backup_shard = _make_scoped_hybrid( + rank, group + ) + _set_local_gradients(muon_parameter, backup_parameter, backup_shard) + rogue = torch.nn.Parameter(torch.ones_like(backup_parameter)) + rogue_before = rogue.detach().clone() + + def closure(): + if rank == 0: + optimizer.backup.param_groups[0]["params"][0] = rogue + rogue.grad = torch.ones_like(rogue) + return torch.tensor(1.0) + + try: + optimizer.step(closure) + message = None + except RuntimeError as exc: + message = str(exc) + return { + "message": message, + "untouched": ( + _untouched(optimizer, muon_parameter, backup_parameter) + and torch.equal(rogue, rogue_before) + ), + } + + +def _step_entry_layout_divergent_result(rank, group): + optimizer, muon_parameter, backup_parameter, backup_shard = _make_scoped_hybrid( + rank, group + ) + _set_local_gradients(muon_parameter, backup_parameter, backup_shard) + rogue = torch.nn.Parameter(torch.ones_like(backup_parameter)) + rogue_before = rogue.detach().clone() + if rank == 0: + optimizer.backup.param_groups[0]["params"][0] = rogue + rogue.grad = torch.ones_like(rogue) + try: + optimizer.step() + message = None + except RuntimeError as exc: + message = str(exc) + return { + "message": message, + "untouched": ( + _untouched(optimizer, muon_parameter, backup_parameter) + and torch.equal(rogue, rogue_before) + ), + } + + +def _distributed_worker(rank, init_file, result_queue): + try: + dist.init_process_group( + "gloo", + init_method="file://{}".format(init_file), + rank=rank, + world_size=_WORLD, + timeout=timedelta(seconds=45), + ) + group = ProcessGroupIdentity("hybrid_scoped_failure", _members()) + amp_divergent = _amp_divergent_overflow_result(rank, group) + dist.barrier() + amp_group_wide = _amp_group_wide_overflow_result(rank, group) + dist.barrier() + amp_finite = _amp_agreeing_finite_result(rank, group) + dist.barrier() + preflight = _preflight_divergent_result(rank, group) + dist.barrier() + closure_divergent = _closure_divergent_result(rank, group) + dist.barrier() + closure_layout_divergent = _closure_layout_divergent_result(rank, group) + dist.barrier() + step_entry_layout_divergent = _step_entry_layout_divergent_result(rank, group) + dist.barrier() + result_queue.put( + { + "rank": rank, + "amp_divergent": amp_divergent, + "amp_group_wide": amp_group_wide, + "amp_finite": amp_finite, + "preflight": preflight, + "closure_divergent": closure_divergent, + "closure_layout_divergent": closure_layout_divergent, + "step_entry_layout_divergent": step_entry_layout_divergent, + } + ) + except BaseException: + result_queue.put({"rank": rank, "fatal_error": traceback.format_exc()}) + finally: + if dist.is_available() and dist.is_initialized(): + dist.destroy_process_group() + + +def _run_distributed_workers(): + context = mp.get_context("spawn") + result_queue = context.Queue() + descriptor, init_file = tempfile.mkstemp(prefix="gefen-hybrid-scoped-failure-") + os.close(descriptor) + os.unlink(init_file) + processes = [ + context.Process( + target=_distributed_worker, + args=(rank, init_file, result_queue), + ) + for rank in range(_WORLD) + ] + results = [] + try: + for process in processes: + process.start() + try: + for _ in processes: + results.append(result_queue.get(timeout=240)) + except queue_module.Empty: + pass + for process in processes: + process.join(timeout=10) + # A worker that queued its result but then hangs in + # destroy_process_group() (or exits nonzero) means the release-critical + # synchronization did not shut down cleanly; fail rather than silently + # terminate it below. + hung = [index for index, process in enumerate(processes) if process.is_alive()] + exit_codes = [process.exitcode for process in processes] + finally: + for process in processes: + if process.is_alive(): + process.terminate() + process.join(timeout=5) + result_queue.close() + result_queue.join_thread() + if os.path.exists(init_file): + os.unlink(init_file) + assert not hung, ("hybrid scoped-failure workers hung", hung, exit_codes) + assert all(code == 0 for code in exit_codes), exit_codes + assert len(results) == _WORLD, (results, exit_codes) + return sorted(results, key=lambda item: item["rank"]) + + +@pytest.mark.skipif( + not dist.is_available() or not dist.is_gloo_available(), + reason="hybrid scoped failure-protocol coverage requires Gloo", +) +def test_hybrid_step_synchronizes_amp_and_preflight_across_the_scope(): + results = _run_distributed_workers() + assert all("fatal_error" not in result for result in results), results + + # Rank-divergent found_inf must raise the children's collective AMP + # agreement error on EVERY member (never a one-sided silent skip that + # strands the peer inside a child's scope header collective), leaving + # both children and the gradients untouched on both ranks. + divergent = [result["amp_divergent"] for result in results] + assert all( + item["message"] is not None and "group-aware gradient scaler" in item["message"] + for item in divergent + ), divergent + assert all(item["untouched"] and item["grads_untouched"] for item in divergent), divergent + + # A group-wide overflow skips BOTH children symmetrically on every + # member: step returns through the post-hook path with no mutation and + # no gradient unscale, and every member exits together (the barriers + # after this scenario would otherwise desynchronize). + group_wide = [result["amp_group_wide"] for result in results] + assert all( + item["skipped"] and item["post_hook_fired"] and item["untouched"] and item["grads_untouched"] + for item in group_wide + ), group_wide + + # An agreeing finite scale steps both children on every member and + # unscales the union of both children's gradients exactly once: the AMP + # step over 8x-scaled gradients is bit-identical to the plain step. + finite = [result["amp_finite"] for result in results] + assert all(item["stepped"] and item["unscaled_once_exactly"] for item in finite), finite + + # A rank-local structural preflight failure raises on every scope member + # together through the children's synchronized failure protocol, keeping + # the atomic both-children-skip semantics on both ranks. + preflight = [result["preflight"] for result in results] + assert preflight[0]["message"] is not None and preflight[1]["message"] is not None, preflight + assert "gradient preflight failed on local member rank:0" in preflight[0]["message"] + assert "sparse gradients" in preflight[0]["message"] + assert "gradient preflight failed on another process-group member" in preflight[1]["message"] + assert all(item["untouched"] for item in preflight), preflight + + # A rank-local closure failure (before any scope synchronization) raises on + # every scope member together through the composite preamble synchronization + # instead of stranding the peer inside a later scoped collective. Both ranks + # exit with an error and leave both children untouched. + closure_divergent = [result["closure_divergent"] for result in results] + assert closure_divergent[0]["message"] is not None and closure_divergent[1]["message"] is not None, closure_divergent + assert "closure boom on rank:0" in closure_divergent[0]["message"] + assert "step preamble failed on another process-group member" in closure_divergent[1]["message"] + assert all(item["untouched"] for item in closure_divergent), closure_divergent + + # A closure can silently replace one finalized child parameter. Synchronize + # that post-closure layout failure through the binding captured at step entry, + # before either member revalidates the now rank-divergent live binding. + closure_layout_divergent = [ + result["closure_layout_divergent"] for result in results + ] + assert closure_layout_divergent[0]["message"] is not None and closure_layout_divergent[1]["message"] is not None, closure_layout_divergent + assert "changed outside post_sharding" in closure_layout_divergent[0]["message"] + assert "gradient preflight failed on another process-group member" in closure_layout_divergent[1]["message"] + assert all(item["untouched"] for item in closure_layout_divergent), closure_layout_divergent + + # A mutation made between steps fails the first finalized-layout guard. The + # explicit binding is captured independently so that guard joins the same + # preamble vote instead of failing on only one member. + step_entry_layout_divergent = [ + result["step_entry_layout_divergent"] for result in results + ] + assert step_entry_layout_divergent[0]["message"] is not None and step_entry_layout_divergent[1]["message"] is not None, step_entry_layout_divergent + assert "changed outside post_sharding" in step_entry_layout_divergent[0]["message"] + assert "step preamble failed on another process-group member" in step_entry_layout_divergent[1]["message"] + assert all(item["untouched"] for item in step_entry_layout_divergent), step_entry_layout_divergent + + +def _make_plain_hybrid(): + matrix = torch.nn.Parameter(_muon_initial().clone()) + vector = torch.nn.Parameter(_backup_initial().clone()) + optimizer = GefenMuonHybrid( + [("matrix", matrix)], + [("vector", vector)], + **_hybrid_kwargs(), + ) + return optimizer, matrix, vector + + +def test_unscoped_amp_overflow_and_finite_step_behavior_is_preserved(): + optimizer, matrix, vector = _make_plain_hybrid() + matrix.grad = _muon_gradient().clone() + vector.grad = _backup_gradient().clone() + optimizer.found_inf = torch.tensor(1.0) + optimizer.grad_scale = torch.tensor(4.0) + assert optimizer.step() is None + assert optimizer.muon._gefen_global_step == 0 + assert optimizer.backup._gefen_global_step == 0 + assert _bits_equal(matrix, _muon_initial()) + assert _bits_equal(vector, _backup_initial()) + assert _bits_equal(vector.grad, _backup_gradient()) + + matrix.grad = _muon_gradient() * 4.0 + vector.grad = _backup_gradient() * 4.0 + optimizer.found_inf = torch.tensor(0.0) + optimizer.grad_scale = torch.tensor(4.0) + optimizer.step() + assert optimizer.muon._gefen_global_step == 1 + assert optimizer.backup._gefen_global_step == 1 + assert not _bits_equal(matrix, _muon_initial()) + assert not _bits_equal(vector, _backup_initial()) + assert _bits_equal(vector.grad, _backup_gradient()) + + +def test_unscoped_preflight_failure_raises_locally_and_skips_both_children(): + optimizer, matrix, vector = _make_plain_hybrid() + matrix.grad = _muon_gradient().clone() + vector.grad = torch.sparse_coo_tensor( + torch.tensor([[0]]), + torch.tensor([1.0]), + size=(8,), + ) + with pytest.raises(RuntimeError) as excinfo: + optimizer.step() + assert "sparse gradients" in str(excinfo.value) + assert "scoped Gefen codebook" not in str(excinfo.value) + assert optimizer.muon._gefen_global_step == 0 + assert optimizer.backup._gefen_global_step == 0 + assert optimizer.muon._gefen_codebook is None + assert optimizer.backup._gefen_codebook is None + assert _bits_equal(matrix, _muon_initial()) + assert _bits_equal(vector, _backup_initial()) + + +def _make_one_member_scoped_hybrid(): + group = ProcessGroupIdentity("solo", ("rank:0",)) + muon_identity = ParameterIdentity(_MUON_FQN, (3, 2)) + backup_identity = ParameterIdentity(_BACKUP_FQN, (8,)) + matrix = torch.nn.Parameter(_muon_initial().clone()) + vector = torch.nn.Parameter(_backup_initial().clone()) + optimizer = GefenMuonHybrid( + [("matrix", matrix)], + [("vector", vector)], + **_hybrid_kwargs(), + ) + + def _replicated(identity): + return ShardIdentity( + identity, + ParameterLayout.REPLICATED, + LogicalSlice.full(identity), + placements=(ShardPlacement("dp", PlacementKind.REPLICATE, 0, 1),), + process_group=group, + local_member="rank:0", + ) + + muon_shard = _replicated(muon_identity) + backup_shard = _replicated(backup_identity) + binding = CodebookProcessGroupBinding(group, "rank:0", None, torch.device("cpu")) + optimizer.post_sharding( + ( + ParameterRebinding(matrix, matrix, muon_shard), + ParameterRebinding(vector, vector, backup_shard), + ), + manifest=ShardingManifest((muon_shard, backup_shard)), + codebook_process_group=binding, + ) + return optimizer, matrix, vector + + +def test_one_member_scope_amp_skip_and_preflight_raise_stay_collective_free(): + optimizer, matrix, vector = _make_one_member_scoped_hybrid() + matrix.grad = _muon_gradient().clone() + vector.grad = _backup_gradient().clone() + optimizer.found_inf = torch.tensor(1.0) + optimizer.grad_scale = torch.tensor(2.0) + assert optimizer.step() is None + assert optimizer.muon._gefen_global_step == 0 + assert optimizer.backup._gefen_global_step == 0 + assert _bits_equal(matrix, _muon_initial()) + assert _bits_equal(vector, _backup_initial()) + + vector.grad = torch.sparse_coo_tensor( + torch.tensor([[0]]), + torch.tensor([1.0]), + size=(8,), + ) + del optimizer.found_inf + del optimizer.grad_scale + # A one-member binding synchronizes without collectives and re-raises the + # local structural error unchanged, exactly like the children. + with pytest.raises(RuntimeError) as excinfo: + optimizer.step() + assert "sparse gradients" in str(excinfo.value) + assert optimizer.muon._gefen_global_step == 0 + assert optimizer.backup._gefen_global_step == 0 + + matrix.grad = _muon_gradient() * 2.0 + vector.grad = _backup_gradient() * 2.0 + optimizer.found_inf = torch.tensor(0.0) + optimizer.grad_scale = torch.tensor(2.0) + optimizer.step() + assert optimizer.muon._gefen_global_step == 1 + assert optimizer.backup._gefen_global_step == 1 + assert not _bits_equal(matrix, _muon_initial()) + assert not _bits_equal(vector, _backup_initial()) + assert _bits_equal(vector.grad, _backup_gradient()) diff --git a/tests/test_layout_guard_cost.py b/tests/test_layout_guard_cost.py new file mode 100644 index 0000000..a2b17c8 --- /dev/null +++ b/tests/test_layout_guard_cost.py @@ -0,0 +1,314 @@ +"""Per-step layout guard cost: fast-path tokens, dedupe, and boundaries. + +The finalized-layout guards run on every ``step()``; these tests pin the +contract that steady-state steps reuse one cached forensic verdict (an +O(local params) identity token check) while every legitimate mutating API and +every boundary operation (checkpoint prepare/commit, rebinding, and contract +readiness) still runs the complete O(params x world) forensic rebuild. +Detection-before-mutation is preserved: anything a +closure or adapter can corrupt through the public optimizer containers still +raises at the step guard itself. +""" + +import time + +import pytest +import torch + +from gefen import ( + Gefen, + LogicalSlice, + ParameterIdentity, + ParameterLayout, + ParameterRebinding, + PlacementKind, + ProcessGroupIdentity, + ShardIdentity, + ShardPlacement, + ShardingManifest, +) + + +def _replicated_shard(fqn, shape): + identity = ParameterIdentity(fqn, shape) + return ShardIdentity( + identity, + ParameterLayout.REPLICATED, + LogicalSlice.full(identity), + ) + + +def _finalized_replicated_optimizer(count=3): + parameters = [ + torch.nn.Parameter(torch.full((4,), float(index + 1))) + for index in range(count) + ] + optimizer = Gefen( + [("weight{}".format(index), parameter) for index, parameter in enumerate(parameters)], + fused=False, + factored_v_2d=False, + ) + shards = tuple( + _replicated_shard("Model.Weight{}".format(index), (4,)) + for index in range(count) + ) + optimizer.post_sharding( + tuple( + ParameterRebinding(parameter, parameter, shard) + for parameter, shard in zip(parameters, shards) + ), + manifest=ShardingManifest(shards), + ) + return optimizer, parameters + + +def _flat_sharded_optimizer(members_count, params_count, local_length=4): + members = tuple("m{: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 + manifest = ShardingManifest(tuple(manifest_shards)) + 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, + ) + optimizer.post_sharding( + tuple( + ParameterRebinding(parameters[index], parameters[index], local_shards[index]) + for index in range(params_count) + ), + manifest=manifest, + ) + return optimizer, parameters + + +def _count_full_layout_passes(monkeypatch): + calls = {"count": 0} + original = Gefen._finalized_binding_layout_matches_full + + def counted(self): + calls["count"] += 1 + return original(self) + + monkeypatch.setattr(Gefen, "_finalized_binding_layout_matches_full", counted) + return calls + + +def _count_manifest_digest_computes(monkeypatch): + calls = {"count": 0} + original = Gefen._compute_codebook_manifest_fingerprint + + def counted(self, manifest): + calls["count"] += 1 + return original(self, manifest) + + monkeypatch.setattr( + Gefen, "_compute_codebook_manifest_fingerprint", counted + ) + return calls + + +def _step_with_grads(optimizer, parameters): + for parameter in parameters: + parameter.grad = torch.full_like(parameter, 0.5) + optimizer.step() + + +def test_steady_state_step_runs_no_full_layout_forensics(monkeypatch): + optimizer, parameters = _finalized_replicated_optimizer() + _step_with_grads(optimizer, parameters) + + calls = _count_full_layout_passes(monkeypatch) + for _ in range(3): + _step_with_grads(optimizer, parameters) + assert calls["count"] == 0 + + +def test_first_step_after_finalization_validates_fully_once(monkeypatch): + optimizer, parameters = _finalized_replicated_optimizer() + calls = _count_full_layout_passes(monkeypatch) + _step_with_grads(optimizer, parameters) + assert calls["count"] == 1 + + +def test_manifest_digest_is_computed_once_at_finalization(monkeypatch): + calls = _count_manifest_digest_computes(monkeypatch) + optimizer, parameters = _finalized_replicated_optimizer() + finalize_computes = calls["count"] + assert finalize_computes >= 1 + + first = optimizer._codebook_manifest_fingerprint() + second = optimizer._codebook_manifest_fingerprint() + assert calls["count"] == finalize_computes + assert first == second == optimizer._compute_codebook_manifest_fingerprint( + optimizer._gefen_sharding_manifest + ) + # The deliberate recompute above must not mask a per-step regression, so + # anchor the count separately from it. + after_deliberate_compute = calls["count"] + assert after_deliberate_compute > finalize_computes + + # The first step after finalization validates the layout fully exactly once + # (see test_first_step_after_finalization_validates_fully_once), which + # recomputes the digest a single time. Warm past it, then require that every + # subsequent steady-state step recomputes nothing -- a per-step recompute + # regression would keep growing the count here. + _step_with_grads(optimizer, parameters) + warm = calls["count"] + assert warm == after_deliberate_compute + 1 + _step_with_grads(optimizer, parameters) + _step_with_grads(optimizer, parameters) + assert calls["count"] == warm + optimizer._codebook_manifest_fingerprint() + assert calls["count"] == warm + + +def test_closure_layout_mutation_is_detected_with_a_warm_verdict(): + optimizer, parameters = _finalized_replicated_optimizer() + _step_with_grads(optimizer, parameters) # warm the cached verdict + + rogue = torch.nn.Parameter(torch.ones(4) * 7) + rogue_before = rogue.detach().clone() + + def mutate_layout(): + optimizer.param_groups[0]["params"][0] = rogue + rogue.grad = torch.ones_like(rogue) + return torch.tensor(1.0) + + with pytest.raises(RuntimeError, match="changed outside post_sharding"): + optimizer.step(mutate_layout) + assert torch.equal(rogue, rogue_before) + + +@pytest.mark.parametrize( + "corruption", + ["local_bindings", "logical_slots", "group_params", "name_cache", "state_name"], +) +def test_step_detects_public_and_replaced_layout_tampering_after_warmup(corruption): + optimizer, parameters = _finalized_replicated_optimizer() + _step_with_grads(optimizer, parameters) # warm the cached verdict + + if corruption == "local_bindings": + optimizer._gefen_local_shard_bindings = tuple( + reversed(optimizer._gefen_local_shard_bindings) + ) + elif corruption == "logical_slots": + optimizer._gefen_logical_slots = tuple( + reversed(optimizer._gefen_logical_slots) + ) + elif corruption == "group_params": + optimizer.param_groups[0]["params"].reverse() + elif corruption == "name_cache": + optimizer._param_names[parameters[0]] = "changed" + else: + optimizer.state[parameters[0]]["name"] = "changed" + + with pytest.raises(RuntimeError, match="changed outside post_sharding"): + _step_with_grads(optimizer, parameters) + + +def test_private_registry_inplace_tamper_is_detected_at_the_next_boundary(): + # In-place value swaps inside the private finalized registries preserve + # every container identity the per-step tokens can see; per the documented + # contract they are detected by the next full forensic boundary + # (checkpoint prepare here, and contract readiness below) instead of the + # next step. + optimizer, parameters = _finalized_replicated_optimizer() + _step_with_grads(optimizer, parameters) + + first, second = parameters[0], parameters[1] + bindings = optimizer._gefen_shard_bindings + bindings[first], bindings[second] = bindings[second], bindings[first] + + _step_with_grads(optimizer, parameters) # fast tokens cannot see this + assert not optimizer.optimizer_contract().capabilities.stable_shard_identity + with pytest.raises(RuntimeError, match="changed outside post_sharding"): + optimizer.state_dict() + + +def test_load_state_dict_bumps_the_layout_version_and_forces_revalidation(monkeypatch): + optimizer, parameters = _finalized_replicated_optimizer() + _step_with_grads(optimizer, parameters) + checkpoint = optimizer.state_dict() + + calls = _count_full_layout_passes(monkeypatch) + + version = optimizer._gefen_layout_version + optimizer.load_state_dict(checkpoint) + assert optimizer._gefen_layout_version > version + assert optimizer._gefen_layout_forensics_verdict is None + + after_load = calls["count"] + _step_with_grads(optimizer, parameters) + assert calls["count"] > after_load + after_revalidation = calls["count"] + _step_with_grads(optimizer, parameters) + assert calls["count"] == after_revalidation # steady again + + +def test_forensics_caches_stay_out_of_the_public_attribute_namespace(): + optimizer, parameters = _finalized_replicated_optimizer() + _step_with_grads(optimizer, parameters) + assert optimizer._gefen_layout_forensics_verdict is not None + for name in Gefen.__slots__: + assert name not in optimizer.__dict__ + + +def test_warm_step_guards_are_far_cheaper_than_one_forensic_pass(): + # Relative timing with a deliberately generous margin: the old behavior + # ran (at least) two full forensic passes inside every step, so the warm + # per-step guard sequence must be much cheaper than even one full pass on + # a moderately sized manifest. Identity-token checks are microseconds + # while the full pass rehashes 96-member identities for every shard, so + # the 10x assertion holds with orders of magnitude to spare. + optimizer, parameters = _flat_sharded_optimizer(96, 60) + _step_with_grads(optimizer, parameters) # warm caches + + iterations = 50 + start = time.perf_counter() + for _ in range(iterations): + optimizer._assert_finalized_binding_layout() + optimizer._assert_runtime_codebook_process_group() + optimizer._assert_finalized_binding_layout() + optimizer._assert_runtime_codebook_process_group() + warm_guard = (time.perf_counter() - start) / iterations + + start = time.perf_counter() + for _ in range(iterations): + assert optimizer._finalized_binding_layout_matches(full=True) + full_pass = (time.perf_counter() - start) / iterations + + assert warm_guard * 10 < full_pass diff --git a/tests/test_muon_distributed_checkpoint_safety.py b/tests/test_muon_distributed_checkpoint_safety.py index 8688999..66fa37f 100644 --- a/tests/test_muon_distributed_checkpoint_safety.py +++ b/tests/test_muon_distributed_checkpoint_safety.py @@ -339,6 +339,8 @@ def _mixed_cpu_checkpoint_worker(rank, world, port, result_queue): import torch.distributed as dist from torch.distributed.tensor import Shard, distribute_tensor, init_device_mesh + from gefen import ParameterLayout, ParameterStateRole, StateExtent, StateScope + try: os.environ["MASTER_ADDR"] = "127.0.0.1" os.environ["MASTER_PORT"] = port @@ -434,6 +436,38 @@ def make_fallback_pair(parallel_value, fallback_value): assign(source_params[0], full(922) * 0.01) source_params[1].grad = full(923) * 0.01 source.step() + contract = source.optimizer_contract() + if rank == 0: + parallel_variant = next( + item + for item in contract.state_layout.parameter_variants + if item.name == "quantized_muon_owner" + ) + assert "m_codebook" in source.state[source_params[0]] + assert ( + source.state[source_params[0]]["m_codebook"].numel() + == initial_parallel.numel() + ) + assert parallel_variant.role is ParameterStateRole.OWNER + assert parallel_variant.extent is StateExtent.OWNER_PARAMETER + else: + parallel_variant = next( + item + for item in contract.state_layout.parameter_variants + if item.name == "distributed_non_owner" + ) + authoritative_keys = { + key + for key in source.state[source_params[0]] + if contract.state_layout.field(key).scope is StateScope.PARAMETER + } + assert authoritative_keys == {"name", "automatic_period"} + assert parallel_variant.role is ParameterStateRole.NON_OWNER + assert parallel_variant.extent is StateExtent.METADATA_ONLY + assert parallel_variant.layouts == frozenset( + {ParameterLayout.DTENSOR_1D_DEFAULT_WORLD} + ) + assert parallel_variant.sharded_mode == "distributed" checkpoint = _clone(source.state_dict()) saved_ids = checkpoint["param_groups"][0]["params"] marker = checkpoint["gefen_muon_distributed"] diff --git a/tests/test_muon_fsdp2_approx.py b/tests/test_muon_fsdp2_approx.py index 52a0b68..af378b7 100644 --- a/tests/test_muon_fsdp2_approx.py +++ b/tests/test_muon_fsdp2_approx.py @@ -43,6 +43,7 @@ def _worker(rank, world, case, port, q): import torch.distributed as dist from torch.distributed.tensor import Shard, distribute_tensor, init_device_mesh + from gefen import ParameterLayout, StateExtent from gefen.gefen_muon import GefenMuon os.environ["MASTER_ADDR"] = "127.0.0.1" @@ -64,6 +65,24 @@ def _worker(rank, world, case, port, q): p.grad = distribute_tensor(full_grad.clone(), mesh, [Shard(0)]) opt.step() local_rows = p.to_local().shape[0] + contract = opt.optimizer_contract() + if p.to_local().numel() == 0: + variant_name = "name_only_dtensor_approx" + expected_extent = StateExtent.METADATA_ONLY + else: + variant_name = "quantized_muon_local" + expected_extent = StateExtent.LOCAL_STORAGE + assert opt.state[p]["m_codebook"].numel() == p.to_local().numel() + variant = next( + item + for item in contract.state_layout.parameter_variants + if item.name == variant_name + ) + assert variant.layouts == frozenset( + {ParameterLayout.DTENSOR_1D_DEFAULT_WORLD} + ) + assert variant.sharded_mode == "approx" + assert variant.extent is expected_extent gathered = p.detach().full_tensor() # collective; all ranks must call if rank == 0: diff --git a/tests/test_muon_fsdp2_parity.py b/tests/test_muon_fsdp2_parity.py index 9ddf216..1313800 100644 --- a/tests/test_muon_fsdp2_parity.py +++ b/tests/test_muon_fsdp2_parity.py @@ -82,6 +82,7 @@ def _worker(rank, world, case, port, q): import torch.distributed as dist from torch.distributed.tensor import Shard, distribute_tensor, init_device_mesh + from gefen import ParameterLayout, StateExtent from gefen.gefen_muon import GefenMuon os.environ["MASTER_ADDR"] = "127.0.0.1" @@ -109,6 +110,24 @@ def _worker(rank, world, case, port, q): ) opt.step() + contract = opt.optimizer_contract() + variant_name = ( + "quantized_normuon_global" + if case.startswith("normuon-") + else "quantized_muon_global" + ) + variant = next( + item + for item in contract.state_layout.parameter_variants + if item.name == variant_name + ) + assert variant.layouts == frozenset( + {ParameterLayout.DTENSOR_1D_DEFAULT_WORLD} + ) + assert variant.sharded_mode == "exact" + assert variant.extent is StateExtent.GLOBAL_PARAMETER + assert opt.state[p]["m_codebook"].numel() == full_init.numel() + gathered = p.detach().full_tensor() # collective; all ranks must call if rank == 0: diff --git a/tests/test_optimizer_contracts.py b/tests/test_optimizer_contracts.py new file mode 100644 index 0000000..4412acc --- /dev/null +++ b/tests/test_optimizer_contracts.py @@ -0,0 +1,760 @@ +"""CPU coverage for public optimizer capability and state-layout contracts.""" + +import copy +from dataclasses import FrozenInstanceError, replace + +import pytest +import torch + +import gefen +from gefen import ( + CONTRACT_SCHEMA_VERSION, + CheckpointTransport, + Gefen, + GefenMuon, + GefenMuonHybrid, + OptimizerChildContract, + OptimizerContractProvider, + OptimizerStateLayout, + ParameterLayout, + ParameterStateRole, + Precision, + ProcessGroupScope, + StateExtent, + StateField, + StateGeometry, + StateKeyMatch, + StateScope, + StateVariant, + TopologyChange, + TrainingSupport, +) + + +_DTENSOR = ParameterLayout.DTENSOR_1D_DEFAULT_WORLD + + +def _values_equal(left, right): + if torch.is_tensor(left) or torch.is_tensor(right): + return ( + torch.is_tensor(left) + and torch.is_tensor(right) + and torch.equal(left, right) + ) + if type(left) is not type(right): + return False + if isinstance(left, dict): + return set(left) == set(right) and all( + _values_equal(left[key], right[key]) for key in left + ) + if isinstance(left, (list, tuple)): + return len(left) == len(right) and all( + _values_equal(a, b) for a, b in zip(left, right) + ) + return left == right + + +def _declares_state_key(contract, key): + try: + contract.state_layout.field(key) + except KeyError: + return False + return True + + +def _step_with_fixed_grad(optimizer, params): + for index, param in enumerate(params): + values = torch.arange(param.numel(), dtype=param.dtype).reshape_as(param) + param.grad = (values + index + 1) / max(1, param.numel()) + optimizer.step() + + +def _persistent_parameter_keys(contract, state): + return { + key + for key in state + if _declares_state_key(contract, key) + and contract.state_layout.field(key).scope is StateScope.PARAMETER + } + + +def _matching_variants( + contract, + state, + *, + layout=ParameterLayout.REPLICATED, + parameter_rank=2, + sharded_mode=None, + role=ParameterStateRole.ANY, +): + keys = _persistent_parameter_keys(contract, state) + return [ + variant + for variant in contract.state_layout.parameter_variants + if set(variant.fields) == keys + and layout in variant.layouts + and ( + variant.parameter_ranks is None or parameter_rank in variant.parameter_ranks + ) + and parameter_rank not in variant.excluded_parameter_ranks + and variant.sharded_mode == sharded_mode + and (variant.role is ParameterStateRole.ANY or variant.role is role) + ] + + +def _training_support(contract, layout, mode=None): + matches = [ + item + for item in contract.capabilities.training + if item.layout is layout and (mode is None or item.sharded_mode == mode) + ] + assert len(matches) == 1 + return matches[0] + + +def _checkpoint_support(contract, transport, layout): + return [ + item + for item in contract.capabilities.checkpoints + if item.transport is transport + and (layout in item.same_topology or layout in item.topology_changing) + ] + + +@pytest.mark.parametrize("factored_v_2d", [False, True]) +def test_plain_contract_matches_live_persistent_state(factored_v_2d): + param = torch.nn.Parameter(torch.arange(16, dtype=torch.float32).reshape(4, 4)) + optimizer = Gefen( + [("layer.weight", param)], + fused=False, + factored_v_2d=factored_v_2d, + ) + assert isinstance(optimizer, OptimizerContractProvider) + contract = optimizer.optimizer_contract() + + assert contract.schema_version == CONTRACT_SCHEMA_VERSION + assert contract.implementation == "gefen.Gefen" + assert contract.capabilities.supported_parameter_ranks is None + assert {item.layout for item in contract.capabilities.training} == { + ParameterLayout.REPLICATED, + ParameterLayout.FLATTENED_ELEMENT_SHARD, + _DTENSOR, + } + dtensor_training = _training_support(contract, _DTENSOR) + assert dtensor_training.mesh_dimensions == (1,) + assert ( + dtensor_training.process_group_scope is ProcessGroupScope.INFERRED_DEVICE_MESH + ) + dcp = _checkpoint_support( + contract, CheckpointTransport.PYTORCH_RANK_LOCAL, _DTENSOR + ) + assert len(dcp) == 1 + assert dcp[0].process_group_scope is ProcessGroupScope.DEFAULT_WORLD + assert dcp[0].mesh_dimensions == (1,) + assert dcp[0].requires_collective + assert dcp[0].atomic_load + assert not dcp[0].topology_changing + native = next( + item + for item in contract.capabilities.checkpoints + if item.transport is CheckpointTransport.NATIVE_OPTIMIZER + ) + assert native.atomic_load + assert ParameterLayout.FLATTENED_ELEMENT_SHARD not in native.same_topology + assert contract.capabilities.accepts_semantic_parameter_names + assert not contract.capabilities.canonical_parameter_fqns + assert not contract.capabilities.stable_shard_identity + assert contract.capabilities.explicit_process_group_codebook_scope + assert contract.capabilities.shard_rebinding + assert contract.capabilities.post_sharding + assert not contract.capabilities.canonical_state_io + assert not contract.capabilities.atomic_state_movement + assert not contract.capabilities.state_offload + assert Precision.FLOAT64 in contract.capabilities.precisions + flattened = _training_support(contract, ParameterLayout.FLATTENED_ELEMENT_SHARD) + assert flattened.process_group_scope is ProcessGroupScope.NONE + + assert optimizer.state[param] == {"name": "layer.weight"} + name_only = next( + item + for item in contract.state_layout.parameter_variants + if item.name == "name_only" + ) + assert name_only.fields == ("name",) + assert name_only.extent is StateExtent.METADATA_ONLY + assert not name_only.initialized + + _step_with_fixed_grad(optimizer, [param]) + state = optimizer.state[param] + assert all(_declares_state_key(contract, key) for key in state) + expected_variant = ( + "factored_second_moment" if factored_v_2d else "block_second_moment" + ) + variant = next( + item + for item in contract.state_layout.parameter_variants + if item.name == expected_variant + ) + persistent_keys = { + key + for key in state + if contract.state_layout.field(key).scope is StateScope.PARAMETER + } + assert persistent_keys == set(variant.fields) + + state_dict = optimizer.state_dict() + common = { + field.name + for field in contract.state_layout.fields_for_scope(StateScope.OPTIMIZER_COMMON) + } + assert common == { + "gefen_global_step", + "gefen_codebook", + "gefen_deterministic", + "gefen_codebook_scope", + } + required_common = { + field.name + for field in contract.state_layout.fields_for_scope(StateScope.OPTIMIZER_COMMON) + if not field.optional + } + assert required_common.issubset(state_dict) + assert contract.state_layout.field("gefen_codebook_scope").optional + assert "gefen_codebook_scope" not in state_dict + + +@pytest.mark.parametrize( + ("source_factored", "target_factored", "migrated_variant"), + [ + (True, False, "block_with_retained_factored_state"), + (False, True, "factored_with_retained_block_state"), + ], +) +def test_plain_contract_declares_factored_v_checkpoint_migration_states( + source_factored, target_factored, migrated_variant +): + initial = torch.arange(16, dtype=torch.float32).reshape(4, 4) + source_param = torch.nn.Parameter(initial.clone()) + source = Gefen( + [("layer.weight", source_param)], + fused=False, + factored_v_2d=source_factored, + ) + _step_with_fixed_grad(source, [source_param]) + checkpoint = copy.deepcopy(source.state_dict()) + + target_param = torch.nn.Parameter(initial.clone()) + target = Gefen( + [("layer.weight", target_param)], + fused=False, + factored_v_2d=target_factored, + ) + target.load_state_dict(checkpoint) + contract = target.optimizer_contract() + pending = _matching_variants(contract, target.state[target_param]) + expected_pending = ( + "factored_state_pending_block_initialization" + if source_factored + else "block_state_pending_factored_initialization" + ) + assert [variant.name for variant in pending] == [expected_pending] + assert pending[0].migration_only + + _step_with_fixed_grad(target, [target_param]) + matches = _matching_variants(contract, target.state[target_param]) + assert [variant.name for variant in matches] == [migrated_variant] + assert matches[0].inactive_fields + + +@pytest.mark.parametrize("target_factored", [False, True]) +def test_plain_contract_declares_legacy_vmean_counter_transition(target_factored): + initial = torch.arange(16, dtype=torch.float32).reshape(4, 4) + source_param = torch.nn.Parameter(initial.clone()) + source = Gefen( + [("layer.weight", source_param)], + fused=False, + factored_v_2d=False, + ) + _step_with_fixed_grad(source, [source_param]) + checkpoint = copy.deepcopy(source.state_dict()) + checkpoint["state"][0].pop("vmean_step") + + target_param = torch.nn.Parameter(initial.clone()) + target = Gefen( + [("layer.weight", target_param)], + fused=False, + factored_v_2d=target_factored, + ) + target.load_state_dict(checkpoint) + contract = target.optimizer_contract() + pending = _matching_variants(contract, target.state[target_param]) + expected_pending = ( + "legacy_block_state_pending_factored_initialization" + if target_factored + else "legacy_block_without_vmean_counter" + ) + assert [variant.name for variant in pending] == [expected_pending] + + _step_with_fixed_grad(target, [target_param]) + matches = _matching_variants(contract, target.state[target_param]) + expected_after = ( + "factored_with_retained_legacy_block_state" + if target_factored + else "block_second_moment" + ) + assert [variant.name for variant in matches] == [expected_after] + + +def test_plain_contract_classifies_legacy_m_codebook_shape_as_inert(): + initial = torch.arange(16, dtype=torch.float32).reshape(4, 4) + source_param = torch.nn.Parameter(initial.clone()) + source = Gefen([("layer.weight", source_param)], fused=False) + _step_with_fixed_grad(source, [source_param]) + checkpoint = copy.deepcopy(source.state_dict()) + checkpoint["state"][0]["m_codebook_shape"] = tuple( + checkpoint["state"][0]["m_codebook"].shape + ) + + target_param = torch.nn.Parameter(initial.clone()) + target = Gefen([("layer.weight", target_param)], fused=False) + target.load_state_dict(checkpoint) + contract = target.optimizer_contract() + field = contract.state_layout.field("m_codebook_shape") + assert field.scope is StateScope.DERIVED + assert not field.authoritative + assert "m_codebook_shape" in target.state[target_param] + assert _matching_variants(contract, target.state[target_param]) + + +@pytest.mark.parametrize( + ("sharded_mode", "normuon", "same_topology_dtensor", "reshard_dtensor"), + [ + ("exact", False, False, False), + ("approx", False, True, False), + ("distributed", True, True, True), + ], +) +def test_muon_contract_separates_mode_topology_and_state_extent( + sharded_mode, normuon, same_topology_dtensor, reshard_dtensor +): + param = torch.nn.Parameter(torch.arange(16, dtype=torch.float32).reshape(4, 4)) + optimizer = GefenMuon( + [("layer.weight", param)], + fused=False, + sharded_mode=sharded_mode, + normuon=normuon, + ) + contract = optimizer.optimizer_contract() + + assert contract.implementation == "gefen.GefenMuon" + assert contract.capabilities.supported_parameter_ranks == (2,) + assert contract.capabilities.explicit_process_group_codebook_scope + assert not contract.capabilities.atomic_state_movement + assert not contract.capabilities.state_offload + native = next( + item + for item in contract.capabilities.checkpoints + if item.transport is CheckpointTransport.NATIVE_OPTIMIZER + and ParameterLayout.REPLICATED in item.same_topology + ) + assert native.atomic_load + replicated = _training_support(contract, ParameterLayout.REPLICATED) + assert replicated.requires_complete_logical_matrix + assert not replicated.requires_complete_parameter_storage + dtensor = _training_support(contract, _DTENSOR, sharded_mode) + assert dtensor.mesh_dimensions == (1,) + assert dtensor.requires_complete_logical_matrix == ( + sharded_mode in ("exact", "distributed") + ) + assert not dtensor.requires_complete_parameter_storage + + same_topology = any( + _DTENSOR in support.same_topology + for support in contract.capabilities.checkpoints + ) + topology_changing = any( + _DTENSOR in support.topology_changing + for support in contract.capabilities.checkpoints + ) + assert same_topology is same_topology_dtensor + assert topology_changing is reshard_dtensor + if same_topology_dtensor: + support = next( + item + for item in contract.capabilities.checkpoints + if _DTENSOR in item.same_topology + ) + assert support.mesh_dimensions == (1,) + assert support.requires_collective + assert support.atomic_load + if reshard_dtensor: + support = next( + item + for item in contract.capabilities.checkpoints + if _DTENSOR in item.topology_changing + ) + assert support.topology_change_kinds == frozenset( + {TopologyChange.WORLD_SIZE_OWNER_REDISTRIBUTION} + ) + assert TopologyChange.PLACEMENT_RESHARD not in support.topology_change_kinds + + common_names = { + field.name + for field in contract.state_layout.fields_for_scope(StateScope.OPTIMIZER_COMMON) + } + assert "gefen_muon_distributed" not in common_names + if sharded_mode == "distributed": + manifest = contract.state_layout.field("gefen_muon_distributed") + assert manifest.scope is StateScope.DERIVED + assert not manifest.authoritative + owner = next( + item + for item in contract.state_layout.parameter_variants + if item.name == "quantized_normuon_owner" + ) + assert owner.extent is StateExtent.OWNER_PARAMETER + assert owner.role is ParameterStateRole.OWNER + non_owner = next( + item + for item in contract.state_layout.parameter_variants + if item.name == "distributed_non_owner" + ) + assert non_owner.fields == ("name", "automatic_period") + assert non_owner.role is ParameterStateRole.NON_OWNER + owner_lazy = _matching_variants( + contract, + {"name": "layer.weight"}, + layout=_DTENSOR, + sharded_mode="distributed", + role=ParameterStateRole.OWNER, + ) + non_owner_lazy = _matching_variants( + contract, + {"name": "layer.weight"}, + layout=_DTENSOR, + sharded_mode="distributed", + role=ParameterStateRole.NON_OWNER, + ) + assert [item.name for item in owner_lazy] == ["distributed_owner_name_only"] + assert [item.name for item in non_owner_lazy] == [ + "distributed_non_owner_name_only" + ] + + assert optimizer.state[param] == {"name": "layer.weight"} + _step_with_fixed_grad(optimizer, [param]) + state = optimizer.state[param] + assert all(_declares_state_key(contract, key) for key in state) + expected_variant = ( + "quantized_normuon_replicated_" + sharded_mode + if normuon + else "quantized_muon_replicated_" + sharded_mode + ) + variant = next( + item + for item in contract.state_layout.parameter_variants + if item.name == expected_variant + ) + persistent_keys = _persistent_parameter_keys(contract, state) + assert persistent_keys == set(variant.fields) + + +@pytest.mark.parametrize("backup_optimizer", ["gefen", "adamw"]) +def test_hybrid_contract_preserves_child_namespaces(backup_optimizer): + matrix = torch.nn.Parameter(torch.arange(16, dtype=torch.float32).reshape(4, 4)) + bias = torch.nn.Parameter(torch.arange(4, dtype=torch.float32)) + optimizer = GefenMuonHybrid( + [("layer.weight", matrix)], + [("layer.bias", bias)], + lr=1e-3, + fused=False, + backup_optimizer=backup_optimizer, + ) + contract = optimizer.optimizer_contract() + + assert contract.implementation == "gefen.GefenMuonHybrid" + assert tuple(child.role for child in contract.children) == ("muon", "backup") + assert contract.children[0].contract.implementation == "gefen.GefenMuon" + if backup_optimizer == "gefen": + assert contract.children[1].contract.implementation == "gefen.Gefen" + else: + assert contract.children[1].implementation == "torch.optim.adamw.AdamW" + assert contract.children[1].contract is None + assert contract.state_layout.composite_namespaces == ("muon", "backup") + gefen_backed = backup_optimizer == "gefen" + assert contract.capabilities.explicit_process_group_codebook_scope is gefen_backed + assert contract.capabilities.shard_rebinding is gefen_backed + assert contract.capabilities.post_sharding is gefen_backed + assert not contract.capabilities.canonical_state_io + assert not contract.capabilities.atomic_state_movement + assert not contract.capabilities.state_offload + assert contract.children[ + 0 + ].contract.capabilities.explicit_process_group_codebook_scope + if backup_optimizer == "gefen": + assert contract.children[ + 1 + ].contract.capabilities.explicit_process_group_codebook_scope + assert {field.name for field in contract.state_layout.fields} == { + "backup_optimizer" + } + checkpoint = contract.capabilities.checkpoints + assert len(checkpoint) == 1 + assert checkpoint[0].transport is CheckpointTransport.COMPOSITE_NATIVE + assert checkpoint[0].same_topology == frozenset({ParameterLayout.REPLICATED}) + assert not checkpoint[0].topology_changing + assert not checkpoint[0].atomic_load + + +def test_muon_contract_keeps_mixed_normuon_variants_in_one_mode(): + plain = torch.nn.Parameter(torch.ones(4, 4)) + normuon = torch.nn.Parameter(torch.ones(4, 4)) + optimizer = GefenMuon( + [ + {"params": [("plain", plain)], "normuon": False}, + {"params": [("normuon", normuon)], "normuon": True}, + ], + fused=False, + sharded_mode="exact", + ) + variants = { + item.name + for item in optimizer.optimizer_contract().state_layout.parameter_variants + } + assert "quantized_muon_replicated_exact" in variants + assert "quantized_normuon_replicated_exact" in variants + + +def test_muon_mixed_approx_distributed_checkpoint_is_same_topology_only(): + first = torch.nn.Parameter(torch.ones(4, 4)) + second = torch.nn.Parameter(torch.ones(4, 4)) + optimizer = GefenMuon( + [ + {"params": [("first", first)], "sharded_mode": "approx"}, + {"params": [("second", second)], "sharded_mode": "distributed"}, + ], + fused=False, + ) + support = _checkpoint_support( + optimizer.optimizer_contract(), + CheckpointTransport.PYTORCH_RANK_LOCAL, + _DTENSOR, + ) + assert len(support) == 1 + assert support[0].required_sharded_modes == frozenset({"approx", "distributed"}) + assert not support[0].topology_changing + + +@pytest.mark.parametrize("backup_optimizer", ["gefen", "adamw"]) +def test_hybrid_approx_does_not_require_complete_dtensor_matrix(backup_optimizer): + matrix = torch.nn.Parameter(torch.ones(4, 4)) + bias = torch.nn.Parameter(torch.ones(4)) + optimizer = GefenMuonHybrid( + [("layer.weight", matrix)], + [("layer.bias", bias)], + lr=1e-3, + fused=False, + sharded_mode="approx", + backup_optimizer=backup_optimizer, + ) + contract = optimizer.optimizer_contract() + support = _training_support(contract, _DTENSOR, "approx") + assert not support.requires_complete_parameter_storage + assert not support.requires_complete_logical_matrix + + +def test_bare_parameters_do_not_claim_canonical_fqns(): + param = torch.nn.Parameter(torch.ones(4, 4)) + optimizer = Gefen([param], fused=False) + contract = optimizer.optimizer_contract() + assert optimizer.state[param]["name"] == "param_0" + assert contract.capabilities.accepts_semantic_parameter_names + assert not contract.capabilities.canonical_parameter_fqns + assert not contract.capabilities.stable_shard_identity + + +def test_contract_is_deeply_immutable_and_query_is_behavior_neutral(): + initial = torch.linspace(-1, 1, 16).reshape(4, 4) + param = torch.nn.Parameter(initial.clone()) + reference_param = torch.nn.Parameter(initial.clone()) + optimizer = Gefen([("layer.weight", param)], fused=False, factored_v_2d=False) + reference = Gefen( + [("layer.weight", reference_param)], fused=False, factored_v_2d=False + ) + + before = copy.deepcopy(optimizer.state_dict()) + contract = optimizer.optimizer_contract() + assert contract == optimizer.optimizer_contract() + assert _values_equal(before, optimizer.state_dict()) + with pytest.raises(FrozenInstanceError): + contract.implementation = "changed" + + source_fields = [ + StateField("name", StateScope.PARAMETER, StateGeometry.OPAQUE, True) + ] + source_variant_fields = ["name"] + source_layouts = {ParameterLayout.REPLICATED} + source_variants = [ + StateVariant( + "name_only", + source_variant_fields, + source_layouts, + StateExtent.METADATA_ONLY, + initialized=False, + ) + ] + layout = OptimizerStateLayout(source_fields, source_variants) + source_fields.clear() + source_variant_fields.clear() + source_layouts.clear() + source_variants.clear() + assert tuple(field.name for field in layout.fields) == ("name",) + assert layout.parameter_variants[0].fields == ("name",) + assert layout.parameter_variants[0].layouts == frozenset( + {ParameterLayout.REPLICATED} + ) + + grad = torch.arange(16, dtype=torch.float32).reshape(4, 4) / 16 + param.grad = grad.clone() + reference_param.grad = grad.clone() + optimizer.step() + reference.step() + assert torch.equal(param, reference_param) + assert _values_equal(optimizer.state_dict(), reference.state_dict()) + + +def test_contract_rejects_duplicate_variant_names(): + field = StateField("name", StateScope.PARAMETER, StateGeometry.OPAQUE, True) + variant = StateVariant( + "name_only", + ("name",), + frozenset({ParameterLayout.REPLICATED}), + StateExtent.METADATA_ONLY, + initialized=False, + ) + with pytest.raises(ValueError, match="names must be unique"): + OptimizerStateLayout((field,), (variant, variant)) + + +def test_derived_fields_are_explicitly_non_authoritative(): + optimizer = Gefen( + [("layer.weight", torch.nn.Parameter(torch.ones(4, 4)))], + fused=False, + ) + derived = optimizer.optimizer_contract().state_layout.fields_for_scope( + StateScope.DERIVED + ) + assert derived + assert all(not field.authoritative for field in derived) + assert {field.name for field in derived if field.checkpointed} == { + "m_codebook_shape", + "_gefen_rank_local_payload_", + "_gefen_rank_local_member", + "_gefen_checkpoint_metadata", + "gefen_native_local_shards", + } + payload_field = optimizer.optimizer_contract().state_layout.field( + "_gefen_rank_local_payload_3" + ) + assert payload_field.key_match is StateKeyMatch.PREFIX + + +def test_all_public_contract_exports_resolve(): + from gefen import contracts + + assert all(getattr(gefen, name) is not None for name in contracts.__all__) + + +def test_portable_global_transport_is_defined_but_not_claimed_before_integration(): + parameter = torch.nn.Parameter(torch.ones(4)) + optimizer = Gefen([("parameter", parameter)], fused=False) + + assert all( + support.transport is not CheckpointTransport.CANONICAL_GLOBAL + for support in optimizer.optimizer_contract().capabilities.checkpoints + ) + + +@pytest.mark.parametrize("backup_optimizer", ["gefen", "adamw"]) +def test_backup_only_hybrid_training_claims_come_from_backup_child(backup_optimizer): + bias = torch.nn.Parameter(torch.arange(4, dtype=torch.float32)) + optimizer = GefenMuonHybrid( + [], + [("layer.bias", bias)], + lr=1e-3, + fused=False, + backup_optimizer=backup_optimizer, + ) + assert optimizer.muon is None + contract = optimizer.optimizer_contract() + assert tuple(child.role for child in contract.children) == ("backup",) + if backup_optimizer == "gefen": + backup_training = contract.children[0].contract.capabilities.training + assert contract.capabilities.training == backup_training + assert any( + item.layout is ParameterLayout.FLATTENED_ELEMENT_SHARD + for item in contract.capabilities.training + ) + else: + # AdamW publishes no contract, so the composite keeps only the plain + # replicated layout the hybrid actually exercises for that child; a + # DTensor claim no code validated would be an over-claim. + assert contract.children[0].contract is None + assert contract.capabilities.training == ( + TrainingSupport(ParameterLayout.REPLICATED, ProcessGroupScope.NONE), + ) + + +@pytest.mark.parametrize("backup_optimizer", ["gefen", "adamw"]) +def test_hybrid_training_claims_are_the_union_of_routed_children(backup_optimizer): + matrix = torch.nn.Parameter(torch.ones(4, 4)) + bias = torch.nn.Parameter(torch.ones(4)) + optimizer = GefenMuonHybrid( + [("layer.weight", matrix)], + [("layer.bias", bias)], + lr=1e-3, + fused=False, + backup_optimizer=backup_optimizer, + ) + contract = optimizer.optimizer_contract() + training = contract.capabilities.training + muon_training = contract.children[0].contract.capabilities.training + if backup_optimizer == "gefen": + backup_training = contract.children[1].contract.capabilities.training + else: + backup_training = ( + TrainingSupport(ParameterLayout.REPLICATED, ProcessGroupScope.NONE), + ) + assert set(training) == set(muon_training) | set(backup_training) + assert len(set(training)) == len(training) + + +def test_capabilities_reject_untyped_entries_and_flags(): + parameter = torch.nn.Parameter(torch.ones(4)) + optimizer = Gefen([("parameter", parameter)], fused=False) + capabilities = optimizer.optimizer_contract().capabilities + with pytest.raises(TypeError, match="training must contain TrainingSupport"): + replace(capabilities, training=[{"junk": 1}]) + with pytest.raises(TypeError, match="checkpoints must contain CheckpointSupport"): + replace(capabilities, checkpoints=["nonsense"]) + with pytest.raises(TypeError, match="precisions must contain Precision"): + replace(capabilities, precisions={"float32"}) + for name in ( + "accepts_semantic_parameter_names", + "canonical_parameter_fqns", + "stable_shard_identity", + "explicit_process_group_codebook_scope", + "shard_rebinding", + "post_sharding", + "canonical_state_io", + "atomic_state_movement", + "state_offload", + ): + with pytest.raises(TypeError, match="{} must be a bool".format(name)): + replace(capabilities, **{name: "yes"}) + + +def test_child_contract_rejects_untyped_contract_payload(): + with pytest.raises( + TypeError, match="contract must be an OptimizerContract or None" + ): + OptimizerChildContract("backup", "torch.optim.adamw.AdamW", {"junk": 1}) diff --git a/tests/test_portable_collective.py b/tests/test_portable_collective.py new file mode 100644 index 0000000..cc0ed92 --- /dev/null +++ b/tests/test_portable_collective.py @@ -0,0 +1,575 @@ +from datetime import timedelta +import hashlib +import multiprocessing as mp +import os +import tempfile + +import pytest +import torch +import torch.distributed as dist + +from gefen.checkpoint import CheckpointProcessGroupBinding +from gefen.contracts import ProcessGroupIdentity +import gefen.portable_collective as portable_collective +from gefen.portable_collective import ( + _CanonicalCollectiveError, + _collective_unanimous_status, + _collective_visit_canonical_fragments, +) +from gefen.portable_wire import _CanonicalWireLimits, _PreparedCanonicalWireValue + + +def _limits(**overrides): + values = { + "max_fragment_tensor_bytes": 1 << 20, + "max_collective_tensor_bytes": 4 << 20, + "max_collective_metadata_bytes": 4 << 20, + "chunk_bytes": 5, + "max_members": 8, + "max_metadata_bytes": 1 << 20, + "max_tree_nodes": 1000, + "max_tree_depth": 16, + "max_container_items": 1000, + "max_string_bytes": 4096, + "max_integer_bytes": 128, + "max_tensors": 128, + "max_tensor_rank": 8, + "diagnostic_bytes": 512, + } + values.update(overrides) + return _CanonicalWireLimits(**values) + + +def _singleton_binding(): + identity = ProcessGroupIdentity("checkpoint", ("worker:solo",)) + return CheckpointProcessGroupBinding( + identity, + "worker:solo", + None, + torch.device("cpu"), + ) + + +def _context(value=b"context"): + return hashlib.sha256(value).digest() + + +def test_singleton_visits_owned_digest_verified_value_without_distributed_calls(monkeypatch): + def forbidden(*args, **kwargs): + raise AssertionError("object collectives and distributed calls are forbidden") + + monkeypatch.setattr(dist, "all_gather_object", forbidden) + monkeypatch.setattr(dist, "broadcast_object_list", forbidden) + monkeypatch.setattr(dist, "all_gather", forbidden) + monkeypatch.setattr(dist, "broadcast", forbidden) + original = torch.arange(6, dtype=torch.float32).reshape(2, 3) + plans = [] + consumed = [] + + def validate(member, prepared): + assert member == "worker:solo" + assert type(prepared) is _PreparedCanonicalWireValue + plans.append((member, prepared.total_tensor_bytes, len(prepared.tensor_specs))) + + def consume(member, value): + consumed.append((member, value)) + + _collective_visit_canonical_fragments( + _singleton_binding(), + {"flag": True, "payload": [original]}, + operation="export", + transaction_id="transaction-1", + context_digest=_context(), + limits=_limits(), + validate_plan=validate, + consume=consume, + ) + + assert plans == [("worker:solo", original.numel() * original.element_size(), 1)] + assert [member for member, _ in consumed] == ["worker:solo"] + result = consumed[0][1] + torch.testing.assert_close(result["payload"][0], original, rtol=0, atol=0) + assert result["payload"][0] is not original + assert result["payload"][0].device.type == "cpu" + assert result["payload"][0].is_contiguous() + original.add_(100) + torch.testing.assert_close( + result["payload"][0], + torch.arange(6, dtype=torch.float32).reshape(2, 3), + rtol=0, + atol=0, + ) + + +def test_singleton_transfers_scalar_empty_complex_and_small_chunk_payloads(): + value = { + "bool_scalar": torch.tensor(True), + "complex": torch.tensor([1 + 2j, -3 + 4j], dtype=torch.complex64), + "empty": torch.empty((2, 0, 3), dtype=torch.uint8), + "integer_scalar": torch.tensor(-7, dtype=torch.int64), + } + consumed = [] + _collective_visit_canonical_fragments( + _singleton_binding(), + value, + operation="export", + transaction_id="transaction-dtypes", + context_digest=_context(), + limits=_limits(chunk_bytes=3), + consume=lambda member, result: consumed.append(result), + ) + + assert len(consumed) == 1 + for name, expected in value.items(): + torch.testing.assert_close(consumed[0][name], expected, rtol=0, atol=0) + + +@pytest.mark.parametrize( + ("fragment", "validate_plan", "consume", "message"), + [ + (object(), None, None, "unsupported canonical wire type"), + (None, lambda member, plan: (_ for _ in ()).throw(ValueError("invalid plan")), None, "invalid plan"), + (None, None, lambda member, value: (_ for _ in ()).throw(RuntimeError("cannot stage")), "cannot stage"), + ], +) +def test_singleton_converts_local_preparation_and_callback_failures_to_collective_rejections( + fragment, + validate_plan, + consume, + message, +): + with pytest.raises(_CanonicalCollectiveError, match=message) as caught: + _collective_visit_canonical_fragments( + _singleton_binding(), + fragment, + operation="import", + transaction_id="transaction-2", + context_digest=_context(), + limits=_limits(), + validate_plan=validate_plan, + consume=consume, + ) + assert "semantic-coordinate[0]" in str(caught.value) + + +def test_singleton_does_not_invoke_consumer_after_metadata_corruption(monkeypatch): + original = portable_collective._metadata_from_tensor + consumed = [] + + def corrupt(value): + metadata = bytearray(original(value)) + metadata[0] ^= 1 + return bytes(metadata) + + monkeypatch.setattr(portable_collective, "_metadata_from_tensor", corrupt) + with pytest.raises(_CanonicalCollectiveError, match="metadata digest"): + _collective_visit_canonical_fragments( + _singleton_binding(), + torch.arange(4, dtype=torch.float32), + operation="export", + transaction_id="transaction-corrupt", + context_digest=_context(), + limits=_limits(), + consume=lambda member, value: consumed.append(value), + ) + assert consumed == [] + + +@pytest.mark.parametrize( + ("changes", "message"), + [ + ({"operation": ""}, "operation"), + ({"transaction_id": " bad"}, "transaction_id"), + ({"context_digest": b"short"}, "context_digest"), + ({"validate_plan": object()}, "validate_plan"), + ({"consume": object()}, "consume"), + ], +) +def test_singleton_strict_invocation_fields_fail_through_fixed_status_exchange(changes, message): + arguments = { + "operation": "export", + "transaction_id": "transaction-strict", + "context_digest": _context(), + "limits": _limits(), + } + arguments.update(changes) + with pytest.raises(_CanonicalCollectiveError, match=message): + _collective_visit_canonical_fragments( + _singleton_binding(), + None, + **arguments, + ) + + +def test_singleton_rejects_limits_that_do_not_fit_signed_header_fields(): + too_large = (1 << 63) + limits = _limits( + max_fragment_tensor_bytes=too_large, + max_collective_tensor_bytes=too_large, + ) + with pytest.raises(_CanonicalCollectiveError, match="signed int64"): + _collective_visit_canonical_fragments( + _singleton_binding(), + None, + operation="export", + transaction_id="transaction-overflow", + context_digest=_context(), + limits=limits, + ) + + +def test_unanimous_status_uses_only_fixed_exchanges_and_reports_local_error(monkeypatch): + def forbidden(*args, **kwargs): + raise AssertionError("no distributed collective is allowed for a singleton") + + monkeypatch.setattr(dist, "all_gather", forbidden) + monkeypatch.setattr(dist, "broadcast", forbidden) + monkeypatch.setattr(portable_collective, "_prepare_canonical_wire_value", forbidden) + _collective_unanimous_status( + _singleton_binding(), + None, + operation="import-ready", + transaction_id="transaction-ready", + context_digest=_context(), + limits=_limits(), + ) + with pytest.raises(_CanonicalCollectiveError, match="stale target"): + _collective_unanimous_status( + _singleton_binding(), + ValueError("stale target"), + operation="import-fresh", + transaction_id="transaction-ready", + context_digest=_context(), + limits=_limits(), + ) + + +def _capture_collective_error(callable_object): + try: + callable_object() + except _CanonicalCollectiveError as exc: + return str(exc) + return None + + +def _distributed_worker(rank, init_file, queue): + try: + dist.init_process_group( + "gloo", + init_method="file://{}".format(init_file), + rank=rank, + world_size=2, + timeout=timedelta(seconds=60), + ) + members = ("worker:left", "worker:right") + binding = CheckpointProcessGroupBinding( + ProcessGroupIdentity("checkpoint", members), + members[rank], + dist.group.WORLD, + torch.device("cpu"), + ) + original_all_gather_object = dist.all_gather_object + original_broadcast_object_list = dist.broadcast_object_list + + def forbidden(*args, **kwargs): + raise AssertionError("object collectives are forbidden") + + dist.all_gather_object = forbidden + dist.broadcast_object_list = forbidden + successes = [] + validation_order = [] + + _collective_visit_canonical_fragments( + binding, + { + "coordinate": rank, + "payload": torch.arange(4, dtype=torch.float32) + rank * 10, + }, + operation="export", + transaction_id="distributed-success", + context_digest=_context(), + limits=_limits(chunk_bytes=7), + validate_plan=lambda member, plan: validation_order.append((member, len(plan.tensor_specs))), + consume=lambda member, value: successes.append( + (member, value["coordinate"], value["payload"].tolist()) + ), + ) + dist.barrier() + + context_error = _capture_collective_error( + lambda: _collective_visit_canonical_fragments( + binding, + None, + operation="export", + transaction_id="divergent-context", + context_digest=_context(b"left" if rank == 0 else b"right"), + limits=_limits(), + ) + ) + dist.barrier() + + operation_error = _capture_collective_error( + lambda: _collective_visit_canonical_fragments( + binding, + None, + operation="export" if rank == 0 else "import", + transaction_id="divergent-operation", + context_digest=_context(), + limits=_limits(), + ) + ) + dist.barrier() + + transaction_error = _capture_collective_error( + lambda: _collective_visit_canonical_fragments( + binding, + None, + operation="export", + transaction_id="transaction-left" if rank == 0 else "transaction-right", + context_digest=_context(), + limits=_limits(), + ) + ) + dist.barrier() + + limits_error = _capture_collective_error( + lambda: _collective_visit_canonical_fragments( + binding, + None, + operation="export", + transaction_id="divergent-limits", + context_digest=_context(), + limits=_limits(chunk_bytes=5 + rank), + ) + ) + dist.barrier() + + wrong_binding = CheckpointProcessGroupBinding( + ProcessGroupIdentity("checkpoint", members), + members[1] if rank == 0 else members[rank], + dist.group.WORLD, + torch.device("cpu"), + ) + coordinate_error = _capture_collective_error( + lambda: _collective_visit_canonical_fragments( + wrong_binding, + None, + operation="export", + transaction_id="coordinate-mismatch", + context_digest=_context(), + limits=_limits(), + ) + ) + dist.barrier() + + callback_presence_error = _capture_collective_error( + lambda: _collective_visit_canonical_fragments( + binding, + None, + operation="import", + transaction_id="callback-presence", + context_digest=_context(), + limits=_limits(), + consume=None if rank == 0 else (lambda member, value: None), + ) + ) + dist.barrier() + + prepare_error = _capture_collective_error( + lambda: _collective_visit_canonical_fragments( + binding, + object() if rank == 0 else None, + operation="export", + transaction_id="prepare-failure", + context_digest=_context(), + limits=_limits(), + ) + ) + dist.barrier() + + def validate(member, plan): + if rank == 1 and member == members[0]: + raise ValueError("receiver rejected plan") + + validate_error = _capture_collective_error( + lambda: _collective_visit_canonical_fragments( + binding, + torch.arange(2, dtype=torch.float32), + operation="import", + transaction_id="validate-failure", + context_digest=_context(), + limits=_limits(), + validate_plan=validate, + ) + ) + dist.barrier() + + def consume(member, value): + if rank == 0 and member == members[1]: + raise RuntimeError("receiver could not stage value") + + consume_error = _capture_collective_error( + lambda: _collective_visit_canonical_fragments( + binding, + torch.arange(2, dtype=torch.float32), + operation="import", + transaction_id="consume-failure", + context_digest=_context(), + limits=_limits(), + consume=consume, + ) + ) + dist.barrier() + + budget_error = _capture_collective_error( + lambda: _collective_visit_canonical_fragments( + binding, + torch.arange(4, dtype=torch.float32), + operation="export", + transaction_id="aggregate-budget", + context_digest=_context(), + limits=_limits( + max_fragment_tensor_bytes=16, + max_collective_tensor_bytes=24, + ), + ) + ) + dist.barrier() + + original_broadcast = portable_collective._broadcast_tensor + broadcast_count = 0 + + def corrupt_broadcast(bound_binding, value, *, source, member_count): + nonlocal broadcast_count + broadcast_count += 1 + if rank == 0 and source == 0 and broadcast_count == 2: + value[0] ^= 1 + return original_broadcast( + bound_binding, + value, + source=source, + member_count=member_count, + ) + + portable_collective._broadcast_tensor = corrupt_broadcast + try: + corruption_error = _capture_collective_error( + lambda: _collective_visit_canonical_fragments( + binding, + torch.arange(2, dtype=torch.float32), + operation="export", + transaction_id="payload-corruption", + context_digest=_context(), + limits=_limits(), + ) + ) + finally: + portable_collective._broadcast_tensor = original_broadcast + dist.barrier() + + status_error = _capture_collective_error( + lambda: _collective_unanimous_status( + binding, + ValueError("target changed") if rank == 1 else None, + operation="import-fresh", + transaction_id="status-helper", + context_digest=_context(), + limits=_limits(), + ) + ) + dist.barrier() + + dist.all_gather_object = original_all_gather_object + dist.broadcast_object_list = original_broadcast_object_list + queue.put( + { + "rank": rank, + "successes": successes, + "validation_order": validation_order, + "context_error": context_error, + "operation_error": operation_error, + "transaction_error": transaction_error, + "limits_error": limits_error, + "coordinate_error": coordinate_error, + "callback_presence_error": callback_presence_error, + "prepare_error": prepare_error, + "validate_error": validate_error, + "consume_error": consume_error, + "budget_error": budget_error, + "corruption_error": corruption_error, + "status_error": status_error, + } + ) + except Exception as exc: + queue.put({"rank": rank, "fatal_error": repr(exc)}) + finally: + if dist.is_available() and dist.is_initialized(): + dist.destroy_process_group() + + +def _run_distributed_workers(): + context = mp.get_context("spawn") + queue = context.Queue() + fd, init_file = tempfile.mkstemp(prefix="gefen-portable-collective-") + os.close(fd) + os.unlink(init_file) + processes = [ + context.Process(target=_distributed_worker, args=(rank, init_file, queue)) + for rank in range(2) + ] + try: + for process in processes: + process.start() + results = [queue.get(timeout=120) for _ in processes] + for process in processes: + process.join(timeout=15) + for process in processes: + if process.is_alive(): + process.terminate() + process.join(timeout=5) + pytest.fail("portable-collective worker hung") + assert process.exitcode == 0 + return sorted(results, key=lambda item: item["rank"]) + finally: + for process in processes: + if process.is_alive(): + process.terminate() + process.join(timeout=5) + if os.path.exists(init_file): + os.unlink(init_file) + + +@pytest.mark.skipif( + not dist.is_available() or not dist.is_gloo_available(), + reason="Gloo is required for the canonical collective transport test", +) +def test_two_process_gloo_transport_is_symmetric_bounded_and_tensor_only(): + results = _run_distributed_workers() + + assert all("fatal_error" not in result for result in results), results + expected_successes = [ + ("worker:left", 0, [0.0, 1.0, 2.0, 3.0]), + ("worker:right", 1, [10.0, 11.0, 12.0, 13.0]), + ] + assert all(result["successes"] == expected_successes for result in results) + assert all( + result["validation_order"] == [("worker:left", 1), ("worker:right", 1)] + for result in results + ) + for key, expected in ( + ("context_error", "context"), + ("operation_error", "operation"), + ("transaction_error", "transaction"), + ("limits_error", "limits"), + ("coordinate_error", "semantic coordinate/order"), + ("callback_presence_error", "callback configuration"), + ("prepare_error", "semantic-coordinate[0]"), + ("validate_error", "receiver rejected plan"), + ("consume_error", "receiver could not stage value"), + ("budget_error", "max_collective_tensor_bytes"), + ("corruption_error", "invalid digest"), + ("status_error", "target changed"), + ): + messages = [result[key] for result in results] + assert messages[0] == messages[1] + assert expected in messages[0] diff --git a/tests/test_portable_dcp.py b/tests/test_portable_dcp.py new file mode 100644 index 0000000..ac0fd8a --- /dev/null +++ b/tests/test_portable_dcp.py @@ -0,0 +1,637 @@ +"""Warning-strict CPU coverage for tensor-only portable DCP persistence.""" + +from __future__ import annotations + +from types import SimpleNamespace +import warnings + +import pytest +import torch +import torch.distributed.checkpoint as dcp +from torch.distributed.checkpoint.api import CheckpointException +from torch.distributed.checkpoint.metadata import ( + BytesStorageMetadata, + ChunkStorageMetadata, + Metadata, + TensorProperties, + TensorStorageMetadata, +) + +import gefen +import gefen.portable_dcp as portable_dcp +from gefen import load_portable_dcp, save_portable_dcp +from gefen.checkpoint import CheckpointProcessGroupBinding +from gefen.codebook import CodebookProcessGroupBinding +from gefen.contracts import ( + LogicalSlice, + ParameterIdentity, + ParameterLayout, + PlacementKind, + ProcessGroupIdentity, + ShardIdentity, + ShardPlacement, + ShardingManifest, +) +from gefen.gefen import Gefen +from gefen.gefen_muon import GefenMuon +from gefen.portable import _decode_quantized_momentum +from gefen.portable_state import PortableStateLimits +from gefen.rebinding import ParameterRebinding + + +_MEMBER = "rank:0" +_SINGLE_PROCESS_DCP_WARNING = ( + r"^torch\.distributed is (?:disabled, )?unavailable or uninitialized, " + r"assuming the intent is to (?:save|load) in a single process\.$" +) +_TYPED_STORAGE_DCP_WARNING = ( + r"^TypedStorage is deprecated\. It will be removed in the future and UntypedStorage will be the only storage class\. " + r"This should only matter to you if you are using storages directly\. To access UntypedStorage directly, use " + r"tensor\.untyped_storage\(\) instead of tensor\.storage\(\)$" +) + + +class _TensorPropertiesSubclass(TensorProperties): + pass + + +@pytest.fixture(autouse=True) +def _warning_strict(): + with warnings.catch_warnings(): + warnings.simplefilter("error") + warnings.filterwarnings( + "ignore", + message=_SINGLE_PROCESS_DCP_WARNING, + category=UserWarning, + ) + warnings.filterwarnings( + "ignore", + message=_TYPED_STORAGE_DCP_WARNING, + category=UserWarning, + ) + yield + + +def _limits( + *, + max_fragment_tensor_bytes=4 << 20, + max_metadata_bytes=64 << 20, +): + return PortableStateLimits( + max_fragment_tensor_bytes=max_fragment_tensor_bytes, + max_collective_tensor_bytes=16 << 20, + max_collective_metadata_bytes=64 << 20, + max_metadata_bytes=max_metadata_bytes, + ) + + +def _bindings(group): + return ( + CodebookProcessGroupBinding( + group, + _MEMBER, + None, + torch.device("cpu"), + ), + CheckpointProcessGroupBinding( + group, + _MEMBER, + None, + torch.device("cpu"), + ), + ) + + +def _finalize(optimizer, parameter, *, layout): + group = ProcessGroupIdentity("checkpoint", (_MEMBER,)) + identity = ParameterIdentity("layer.weight", tuple(parameter.shape)) + if layout is ParameterLayout.REPLICATED: + kind = PlacementKind.REPLICATE + owner = None + else: + kind = PlacementKind.WHOLE_PARAMETER_OWNER + owner = _MEMBER + shard = ShardIdentity( + identity, + layout, + LogicalSlice.full(identity), + placements=(ShardPlacement("checkpoint", kind, 0, 1),), + process_group=group, + local_member=_MEMBER, + owner=owner, + ) + codebook_binding, checkpoint_binding = _bindings(group) + optimizer.post_sharding( + (ParameterRebinding(parameter, parameter, shard),), + manifest=ShardingManifest((shard,)), + codebook_process_group=codebook_binding, + ) + return checkpoint_binding + + +def _plain(*, factored, deterministic): + parameter = torch.nn.Parameter(torch.arange(1, 7, dtype=torch.float32).reshape(2, 3)) + optimizer = Gefen( + [("weight", parameter)], + fused=False, + factored_v_2d=factored, + force_2d_period_one=True, + deterministic=deterministic, + ) + binding = _finalize( + optimizer, + parameter, + layout=ParameterLayout.REPLICATED, + ) + return optimizer, parameter, binding + + +def _muon(*, deterministic): + parameter = torch.nn.Parameter(torch.arange(1, 7, dtype=torch.float32).reshape(2, 3)) + optimizer = GefenMuon( + [("weight", parameter)], + fused=False, + sharded_mode="distributed", + normuon=True, + deterministic=deterministic, + ) + binding = _finalize( + optimizer, + parameter, + layout=ParameterLayout.WHOLE_PARAMETER_OWNER, + ) + return optimizer, parameter, binding + + +def _optimizer_pair(variant): + if variant == "muon-normuon": + source, source_parameter, source_binding = _muon(deterministic=True) + target, target_parameter, target_binding = _muon(deterministic=False) + else: + factored = variant == "plain-factored" + source, source_parameter, source_binding = _plain( + factored=factored, + deterministic=True, + ) + target, target_parameter, target_binding = _plain( + factored=factored, + deterministic=False, + ) + return ( + source, + source_parameter, + source_binding, + target, + target_parameter, + target_binding, + ) + + +def _initialize(optimizer, parameter, *, variant): + optimizer._gefen_global_step = 4 + optimizer._gefen_codebook = torch.linspace( + -1.0, + 1.0, + 256, + dtype=torch.float32, + ) + state = optimizer.state[parameter] + state.update( + { + "automatic_period": 1, + "step": 4, + "m_codebook": torch.tensor( + [[0], [255], [128], [64], [192], [0]], + dtype=torch.uint8, + ), + "m_magnitude": torch.tensor( + [[1.0], [2.0], [3.0], [4.0], [5.0], [6.0]], + dtype=torch.float32, + ), + } + ) + if variant == "plain-block": + state.update( + { + "vmean": torch.arange(1, 7, dtype=torch.float32).reshape(6, 1), + "vmean_step": 3, + } + ) + elif variant == "plain-factored": + state.update( + { + "v_row": torch.tensor([2.0, 3.0], dtype=torch.float32), + "v_col": torch.tensor([5.0, 7.0, 11.0], dtype=torch.float32), + "factored_step": 3, + } + ) + elif variant == "muon-normuon": + state.update( + { + "normuon_v": torch.tensor([[2.0], [3.0]], dtype=torch.float32), + "normuon_step": 2, + } + ) + else: + raise AssertionError("unknown test variant") + + +def _assert_loaded_state(target, target_parameter, expected, *, variant): + assert target._gefen_global_step == 4 + assert target._deterministic is True + assert torch.equal( + target._gefen_codebook, + expected["common"]["gefen_codebook"], + ) + target_state = target.state[target_parameter] + target_momentum = _decode_quantized_momentum( + target._gefen_codebook, + target_state["m_codebook"], + target_state["m_magnitude"], + logical_shape=tuple(target_parameter.shape), + period=1, + step=target_state["step"], + ) + record = expected["parameters"]["layer.weight"] + assert torch.equal( + target_momentum.view(torch.int32), + record["state"]["momentum"].view(torch.int32), + ) + if variant == "plain-block": + assert torch.equal( + target_state["vmean"].reshape(2, 3), + record["state"]["second_moment"], + ) + assert target_state["vmean_step"] == 3 + elif variant == "plain-factored": + assert torch.equal(target_state["v_row"], record["state"]["v_row"]) + assert torch.equal(target_state["v_col"], record["state"]["v_col"]) + assert target_state["factored_step"] == 3 + else: + assert torch.equal( + target_state["normuon_v"], + record["state"]["normuon_v"], + ) + assert target_state["normuon_step"] == 2 + + +@pytest.mark.parametrize( + "variant", + ["plain-block", "plain-factored", "muon-normuon"], +) +def test_filesystem_dcp_round_trip_is_tensor_only(tmp_path, variant): + ( + source, + source_parameter, + source_binding, + target, + target_parameter, + target_binding, + ) = _optimizer_pair(variant) + _initialize(source, source_parameter, variant=variant) + limits = _limits() + expected = source.export_portable_state( + checkpoint_process_group=source_binding, + transaction_id="dcp-expected-{}".format(variant), + limits=limits, + ) + checkpoint = tmp_path / variant + + metadata = save_portable_dcp( + source, + checkpoint_process_group=source_binding, + storage_writer=dcp.FileSystemWriter(checkpoint), + transaction_id="dcp-save-{}".format(variant), + limits=limits, + ) + + assert type(metadata) is Metadata + assert metadata.state_dict_metadata + assert all(type(value) is TensorStorageMetadata for value in metadata.state_dict_metadata.values()) + assert not any(type(value) is BytesStorageMetadata for value in metadata.state_dict_metadata.values()) + on_disk_metadata = dcp.FileSystemReader(checkpoint).read_metadata() + assert set(on_disk_metadata.state_dict_metadata) == set(metadata.state_dict_metadata) + assert all(type(value) is TensorStorageMetadata for value in on_disk_metadata.state_dict_metadata.values()) + + load_portable_dcp( + target, + checkpoint_process_group=target_binding, + storage_reader=dcp.FileSystemReader(checkpoint), + transaction_id="dcp-load-{}".format(variant), + limits=limits, + ) + + _assert_loaded_state( + target, + target_parameter, + expected, + variant=variant, + ) + + +@pytest.mark.parametrize( + "namespace", + ["", " optimizer", "optimizer.", "optim\x00izer", "optim/izer", "optimé"], +) +def test_namespace_validation_precedes_dcp_io(tmp_path, namespace): + optimizer, parameter, binding = _plain( + factored=False, + deterministic=True, + ) + _initialize(optimizer, parameter, variant="plain-block") + checkpoint = tmp_path / "invalid-namespace" + + with pytest.raises(RuntimeError, match="namespace"): + save_portable_dcp( + optimizer, + checkpoint_process_group=binding, + storage_writer=dcp.FileSystemWriter(checkpoint), + transaction_id="dcp-invalid-namespace", + limits=_limits(), + namespace=namespace, + ) + + assert not (checkpoint / ".metadata").exists() + + +@pytest.mark.parametrize( + ("members", "world_size", "group_zero", "match"), + [ + (("rank:0",), 2, 0, "singleton"), + (("rank:1", "rank:2"), 3, 1, "global rank zero"), + ], +) +def test_dcp_runtime_rejects_unsafe_default_world_or_subgroup( + monkeypatch, + members, + world_size, + group_zero, + match, +): + binding = SimpleNamespace( + identity=SimpleNamespace(ordered_members=members), + process_group=object(), + collective_device=torch.device("cpu"), + ) + monkeypatch.setattr(torch.distributed, "is_available", lambda: True) + monkeypatch.setattr(torch.distributed, "is_initialized", lambda: True) + monkeypatch.setattr(torch.distributed, "get_world_size", lambda: world_size) + monkeypatch.setattr( + torch.distributed, + "get_global_rank", + lambda _group, _rank: group_zero, + ) + + with pytest.raises(RuntimeError, match=match): + portable_dcp._validate_dcp_runtime(binding) + + +def test_load_rejects_namespace_mismatch_without_target_mutation(tmp_path): + ( + source, + source_parameter, + source_binding, + target, + _, + target_binding, + ) = _optimizer_pair("plain-block") + _initialize(source, source_parameter, variant="plain-block") + checkpoint = tmp_path / "namespace-mismatch" + save_portable_dcp( + source, + checkpoint_process_group=source_binding, + storage_writer=dcp.FileSystemWriter(checkpoint), + transaction_id="dcp-save-namespace-mismatch", + limits=_limits(), + namespace="optimizer", + ) + before = target._canonical_import_live_token() + + with pytest.raises(CheckpointException): + load_portable_dcp( + target, + checkpoint_process_group=target_binding, + storage_reader=dcp.FileSystemReader(checkpoint), + transaction_id="dcp-load-namespace-mismatch", + limits=_limits(), + namespace="other", + ) + + assert target._canonical_import_live_token() == before + + +def _tensor_metadata(dtype, shape, *, offsets=None, sizes=None): + shape = torch.Size(shape) + if offsets is None: + offsets = (0,) * len(shape) + if sizes is None: + sizes = shape + return TensorStorageMetadata( + properties=TensorProperties(dtype=dtype), + size=shape, + chunks=[ + ChunkStorageMetadata( + offsets=torch.Size(offsets), + sizes=torch.Size(sizes), + ) + ], + ) + + +def _planner_metadata(case, *, namespace, limits): + metadata_key = portable_dcp._flat_key( + namespace, + portable_dcp._DCP_METADATA_KEY, + ) + payload_key = portable_dcp._flat_key( + namespace, + portable_dcp._payload_key(0), + ) + entries = {metadata_key: _tensor_metadata(torch.uint8, (1,))} + if case == "oversized-wire-metadata": + entries[metadata_key] = _tensor_metadata( + torch.uint8, + (limits.max_metadata_bytes + 1,), + ) + elif case == "oversized-payload": + entries[payload_key] = _tensor_metadata( + torch.uint8, + (limits.max_fragment_tensor_bytes + 1,), + ) + elif case == "extra-key": + entries[portable_dcp._flat_key(namespace, "extra")] = _tensor_metadata( + torch.uint8, + (1,), + ) + elif case == "byte-io": + entries[metadata_key] = BytesStorageMetadata() + elif case == "partial-tensor": + entries[metadata_key] = _tensor_metadata( + torch.uint8, + (8,), + offsets=(0,), + sizes=(4,), + ) + elif case == "partial-planner-data": + entries[payload_key] = _tensor_metadata(torch.float32, (1,)) + elif case == "properties-subclass": + entry = _tensor_metadata(torch.uint8, (1,)) + entry.properties = _TensorPropertiesSubclass(dtype=torch.uint8) + entries[metadata_key] = entry + elif case == "list-size": + entry = _tensor_metadata(torch.uint8, (1,)) + entry.size = [1] + entries[metadata_key] = entry + elif case == "list-chunk-offsets": + entry = _tensor_metadata(torch.uint8, (1,)) + entry.chunks[0].offsets = [0] + entries[metadata_key] = entry + else: + raise AssertionError("unknown planner metadata case") + planner_data = {key: tuple(key.split(".", 1)) for key in entries} + if case == "partial-planner-data": + planner_data.pop(payload_key) + return Metadata( + state_dict_metadata=entries, + planner_data=planner_data, + ) + + +@pytest.mark.parametrize( + ("case", "error_type", "match"), + [ + ("oversized-wire-metadata", ValueError, "metadata exceeds"), + ("oversized-payload", ValueError, "payload tensors exceed"), + ("extra-key", ValueError, "unexpected tensor keys"), + ("byte-io", TypeError, "full-tensor DCP metadata"), + ("partial-tensor", ValueError, "complete unsharded tensor"), + ("partial-planner-data", ValueError, "incompatible planner paths"), + ("properties-subclass", TypeError, "unsupported tensor properties"), + ("list-size", TypeError, "unsupported tensor properties"), + ("list-chunk-offsets", ValueError, "complete unsharded tensor"), + ], +) +def test_load_planner_rejects_untrusted_metadata_before_allocation( + monkeypatch, + case, + error_type, + match, +): + namespace = "optimizer" + limits = _limits( + max_fragment_tensor_bytes=64, + max_metadata_bytes=64, + )._wire_limits(collective=True) + metadata = _planner_metadata( + case, + namespace=namespace, + limits=limits, + ) + planner = portable_dcp._load_planner(namespace, limits) + allocations = [] + + def fail_if_allocated(*args, **kwargs): + allocations.append((args, kwargs)) + raise AssertionError("untrusted DCP metadata reached tensor allocation") + + monkeypatch.setattr(torch, "empty", fail_if_allocated) + + with pytest.raises(error_type, match=match): + planner.set_up_planner( + {namespace: {}}, + metadata, + is_coordinator=True, + ) + + assert allocations == [] + + +def test_dcp_envelope_container_limit_is_symmetric_before_save(): + limits = PortableStateLimits( + max_fragment_tensor_bytes=64, + max_collective_tensor_bytes=64, + max_collective_metadata_bytes=4096, + max_metadata_bytes=4096, + max_container_items=3, + max_tensors=8, + )._wire_limits(collective=True) + plan = portable_dcp._prepare_canonical_wire_value( + ( + torch.ones(1), + torch.ones(1), + torch.ones(1), + ), + limits, + ) + + with pytest.raises(ValueError, match="max_container_items"): + portable_dcp._state_from_plan("optimizer", plan, limits) + + +def test_corrupted_tensor_payload_leaves_target_unchanged(tmp_path): + ( + source, + source_parameter, + source_binding, + target, + target_parameter, + target_binding, + ) = _optimizer_pair("plain-block") + _initialize(source, source_parameter, variant="plain-block") + _initialize(target, target_parameter, variant="plain-block") + target._deterministic = False + target._gefen_global_step = 9 + limits = _limits() + original = tmp_path / "original" + corrupted = tmp_path / "corrupted" + save_portable_dcp( + source, + checkpoint_process_group=source_binding, + storage_writer=dcp.FileSystemWriter(original), + transaction_id="dcp-save-corruption-source", + limits=limits, + ) + + wire_limits = limits._wire_limits(collective=True) + state = {"optimizer": {}} + dcp.load( + state, + storage_reader=dcp.FileSystemReader(original), + planner=portable_dcp._load_planner("optimizer", wire_limits), + process_group=None, + ) + payload_keys = sorted(key for key in state["optimizer"] if key.startswith(portable_dcp._DCP_PAYLOAD_PREFIX)) + assert payload_keys + payload = next(state["optimizer"][key] for key in payload_keys if state["optimizer"][key].numel()) + raw = payload.view(torch.uint8).reshape(-1) + raw[0] = int(raw[0]) ^ 1 + dcp.save( + state, + storage_writer=dcp.FileSystemWriter(corrupted), + process_group=None, + ) + before = target._canonical_import_live_token() + + with pytest.raises(RuntimeError, match="invalid digest"): + load_portable_dcp( + target, + checkpoint_process_group=target_binding, + storage_reader=dcp.FileSystemReader(corrupted), + transaction_id="dcp-load-corrupted-payload", + limits=limits, + ) + + assert target._canonical_import_live_token() == before + + +def test_portable_dcp_helpers_are_public_exports(): + assert gefen.save_portable_dcp is portable_dcp.save_portable_dcp + assert gefen.load_portable_dcp is portable_dcp.load_portable_dcp + assert save_portable_dcp is portable_dcp.save_portable_dcp + assert load_portable_dcp is portable_dcp.load_portable_dcp + assert "save_portable_dcp" in gefen.__all__ + assert "load_portable_dcp" in gefen.__all__ + assert portable_dcp.__all__ == [ + "load_portable_dcp", + "save_portable_dcp", + ] diff --git a/tests/test_portable_dcp_distributed.py b/tests/test_portable_dcp_distributed.py new file mode 100644 index 0000000..8c44e84 --- /dev/null +++ b/tests/test_portable_dcp_distributed.py @@ -0,0 +1,530 @@ +"""CPU-only distributed coverage for portable optimizer DCP persistence.""" + +from datetime import timedelta +import multiprocessing as mp +import os +import queue as queue_module +import tempfile +import traceback +import warnings + +import pytest +import torch +import torch.distributed as dist +from torch.distributed.checkpoint import FileSystemReader, FileSystemWriter + +from gefen import load_portable_dcp, save_portable_dcp +from gefen.checkpoint import CheckpointProcessGroupBinding +from gefen.codebook import CodebookProcessGroupBinding +from gefen.contracts import ( + LogicalSlice, + ParameterIdentity, + ParameterLayout, + PlacementKind, + ProcessGroupIdentity, + ShardIdentity, + ShardPlacement, + ShardingManifest, +) +from gefen.gefen import Gefen +from gefen.portable import _decode_quantized_momentum +from gefen.portable_state import PortableStateLimits +from gefen.rebinding import ParameterRebinding + + +_SAVE_WORLD = 2 +_FQN = "model.weight" + + +def _limits(): + return PortableStateLimits( + max_fragment_tensor_bytes=1 << 20, + max_collective_tensor_bytes=4 << 20, + max_collective_metadata_bytes=4 << 20, + chunk_bytes=11, + max_members=4, + max_metadata_bytes=1 << 20, + max_tree_nodes=10_000, + max_tree_depth=32, + max_container_items=10_000, + max_string_bytes=16 << 10, + max_integer_bytes=128, + max_tensors=256, + max_tensor_rank=8, + diagnostic_bytes=1024, + ) + + +def _save_members(): + return tuple("save:{}".format(rank) for rank in range(_SAVE_WORLD)) + + +def _optimizer(parameter, *, deterministic): + return Gefen( + [("weight", parameter)], + lr=2.5e-3, + betas=(0.8, 0.97), + eps=2.0e-8, + weight_decay=0.03, + fused=False, + force_2d_period_one=True, + factored_v_2d=False, + deterministic=deterministic, + ) + + +def _bindings(identity, local_member, process_group): + return ( + CodebookProcessGroupBinding( + identity, + local_member, + process_group, + torch.device("cpu"), + ), + CheckpointProcessGroupBinding( + identity, + local_member, + process_group, + torch.device("cpu"), + ), + ) + + +def _flat_manifest(identity, group, lengths): + if len(lengths) != len(group.ordered_members) or sum(lengths) != identity.numel: + raise AssertionError("invalid flat test partition") + shards = [] + offset = 0 + for coordinate, (member, length) in enumerate(zip(group.ordered_members, lengths)): + shards.append( + ShardIdentity( + identity, + ParameterLayout.FLATTENED_ELEMENT_SHARD, + LogicalSlice(offset, length), + placements=( + ShardPlacement( + "checkpoint", + PlacementKind.FLAT_SHARD, + coordinate, + len(group.ordered_members), + ), + ), + process_group=group, + local_member=member, + ) + ) + offset += length + return ShardingManifest(tuple(shards)), tuple(shards) + + +def _replicated_manifest(identity, group): + shard = ShardIdentity( + identity, + ParameterLayout.REPLICATED, + LogicalSlice.full(identity), + placements=( + ShardPlacement( + "checkpoint", + PlacementKind.REPLICATE, + 0, + 1, + ), + ), + process_group=group, + local_member=group.ordered_members[0], + ) + return ShardingManifest((shard,)), shard + + +def _make_flat_source(rank, group): + identity = ParameterIdentity(_FQN, (2, 4)) + lengths = (3, 5) + original = torch.nn.Parameter(torch.zeros(identity.global_shape, dtype=torch.float32)) + local = torch.nn.Parameter(torch.zeros(lengths[rank], dtype=torch.float32)) + optimizer = _optimizer(original, deterministic=True) + manifest, shards = _flat_manifest(identity, group, lengths) + codebook_binding, checkpoint_binding = _bindings(group, _save_members()[rank], dist.group.WORLD) + optimizer.post_sharding( + (ParameterRebinding(original, local, shards[rank]),), + manifest=manifest, + codebook_process_group=codebook_binding, + ) + return optimizer, local, shards[rank], checkpoint_binding + + +def _make_replicated_target(initial_parameter, *, deterministic): + identity = ParameterIdentity(_FQN, (2, 4)) + group = ProcessGroupIdentity("portable_dcp_singleton", ("load:0",)) + parameter = torch.nn.Parameter(initial_parameter.clone()) + optimizer = _optimizer(parameter, deterministic=deterministic) + manifest, shard = _replicated_manifest(identity, group) + codebook_binding, checkpoint_binding = _bindings(group, "load:0", None) + optimizer.post_sharding( + (ParameterRebinding(parameter, parameter, shard),), + manifest=manifest, + codebook_process_group=codebook_binding, + ) + return optimizer, parameter, checkpoint_binding + + +def _codebook(): + return torch.linspace(-1.0, 1.0, 256, dtype=torch.float32) + + +def _momentum_values(): + return ( + torch.tensor( + [ + -2147483648, + 0, + 1, + -2147483647, + 1056964608, + -1090519040, + 1078984704, + -1058013184, + ], + dtype=torch.int32, + ) + .view(torch.float32) + .reshape(2, 4) + .clone() + ) + + +def _second_moment_values(): + return ( + torch.tensor( + [ + -2147483648, + 0, + 1, + 8388608, + 1048576000, + 1065353216, + 1073741824, + 2139095039, + ], + dtype=torch.int32, + ) + .view(torch.float32) + .reshape(2, 4) + .clone() + ) + + +def _quantized_period_one(momentum): + flat = momentum.reshape(-1) + indices = torch.where( + torch.signbit(flat), + torch.zeros(flat.numel(), dtype=torch.uint8), + torch.full((flat.numel(),), 255, dtype=torch.uint8), + ) + return indices.reshape(-1, 1), flat.abs().reshape(-1, 1).clone() + + +def _seed_state(optimizer, parameter, momentum, second_moment): + indices, magnitudes = _quantized_period_one(momentum) + optimizer._gefen_global_step = 9 + optimizer._gefen_codebook = _codebook() + optimizer.state[parameter].update( + { + "automatic_period": 1, + "step": 7, + "m_codebook": indices, + "m_magnitude": magnitudes, + "vmean": second_moment.reshape(-1, 1).clone(), + "vmean_step": 6, + } + ) + + +def _bits_equal(left, right): + return ( + type(left) is torch.Tensor + and type(right) is torch.Tensor + and left.dtype == right.dtype + and tuple(left.shape) == tuple(right.shape) + and torch.equal(left.contiguous().view(torch.uint8), right.contiguous().view(torch.uint8)) + ) + + +def _capture_runtime_error(operation): + try: + operation() + except RuntimeError as exc: + return str(exc) + raise AssertionError("portable DCP operation unexpectedly succeeded") + + +def _strict_worker_warnings(): + warnings.simplefilter("error") + warnings.filterwarnings( + "ignore", + message="TypedStorage is deprecated.*", + category=UserWarning, + ) + + +def _save_worker( + rank, + init_file, + checkpoint_dir, + mismatch_dir, + divergent_dir, + invalid_dir, + result_queue, +): + _strict_worker_warnings() + try: + dist.init_process_group( + "gloo", + init_method="file://{}".format(init_file), + rank=rank, + world_size=_SAVE_WORLD, + timeout=timedelta(seconds=60), + ) + group = ProcessGroupIdentity("portable_dcp_save", _save_members()) + optimizer, parameter, shard, checkpoint_binding = _make_flat_source(rank, group) + momentum = _momentum_values().reshape(-1) + second_moment = _second_moment_values().reshape(-1) + start = shard.logical_slice.flat_offset + stop = start + shard.logical_slice.length + _seed_state( + optimizer, + parameter, + momentum[start:stop].clone(), + second_moment[start:stop].clone(), + ) + + save_portable_dcp( + optimizer, + checkpoint_process_group=checkpoint_binding, + storage_writer=FileSystemWriter(checkpoint_dir), + transaction_id="portable-dcp-save-2-to-1", + limits=_limits(), + namespace="optimizer", + ) + dist.barrier() + + token = optimizer._canonical_import_live_token() + namespace_message = _capture_runtime_error( + lambda: save_portable_dcp( + optimizer, + checkpoint_process_group=checkpoint_binding, + storage_writer=FileSystemWriter(mismatch_dir), + transaction_id="portable-dcp-mismatched-namespace", + limits=_limits(), + namespace="optimizer-a" if rank == 0 else "optimizer-b", + ) + ) + namespace_unchanged = optimizer._canonical_import_live_token() == token + dist.barrier() + + token = optimizer._canonical_import_live_token() + storage_message = _capture_runtime_error( + lambda: save_portable_dcp( + optimizer, + checkpoint_process_group=checkpoint_binding, + storage_writer=FileSystemWriter("{}-{}".format(divergent_dir, rank)), + transaction_id="portable-dcp-mismatched-storage", + limits=_limits(), + namespace="optimizer", + ) + ) + storage_unchanged = optimizer._canonical_import_live_token() == token + dist.barrier() + + token = optimizer._canonical_import_live_token() + storage = object() if rank == 0 else FileSystemWriter(invalid_dir) + invalid_message = _capture_runtime_error( + lambda: save_portable_dcp( + optimizer, + checkpoint_process_group=checkpoint_binding, + storage_writer=storage, + transaction_id="portable-dcp-invalid-writer", + limits=_limits(), + namespace="optimizer", + ) + ) + invalid_unchanged = optimizer._canonical_import_live_token() == token + dist.barrier() + + result_queue.put( + { + "rank": rank, + "saved": True, + "namespace_message": namespace_message, + "namespace_unchanged": namespace_unchanged, + "storage_message": storage_message, + "storage_unchanged": storage_unchanged, + "invalid_message": invalid_message, + "invalid_unchanged": invalid_unchanged, + } + ) + except BaseException: + result_queue.put({"rank": rank, "fatal_error": traceback.format_exc()}) + finally: + if dist.is_available() and dist.is_initialized(): + dist.destroy_process_group() + + +def _loaded_state_is_exact(optimizer, parameter): + state = optimizer.state[parameter] + decoded = _decode_quantized_momentum( + optimizer._gefen_codebook, + state["m_codebook"], + state["m_magnitude"], + logical_shape=(2, 4), + period=1, + step=state["step"], + ) + return ( + _bits_equal(decoded, _momentum_values()) + and _bits_equal(state["vmean"].reshape(2, 4), _second_moment_values()) + and _bits_equal(optimizer._gefen_codebook, _codebook()) + and optimizer._gefen_global_step == 9 + and state["automatic_period"] == 1 + and state["step"] == 7 + and state["vmean_step"] == 6 + and optimizer._deterministic is True + ) + + +def _next_step_is_exact(optimizer, parameter, initial_parameter): + oracle, oracle_parameter, _ = _make_replicated_target(initial_parameter, deterministic=True) + _seed_state(oracle, oracle_parameter, _momentum_values(), _second_moment_values()) + gradient = torch.tensor( + [0.25, -0.5, 0.75, -1.0, 1.25, -1.5, 1.75, -2.0], + dtype=torch.float32, + ).reshape(2, 4) + parameter.grad = gradient.clone() + oracle_parameter.grad = gradient.clone() + optimizer.step() + oracle.step() + state = optimizer.state[parameter] + oracle_state = oracle.state[oracle_parameter] + return ( + _bits_equal(parameter.detach(), oracle_parameter.detach()) + and optimizer._gefen_global_step == oracle._gefen_global_step == 10 + and state["step"] == oracle_state["step"] == 8 + and state["vmean_step"] == oracle_state["vmean_step"] == 7 + and all(_bits_equal(state[key], oracle_state[key]) for key in ("m_codebook", "m_magnitude", "vmean")) + ) + + +def _load_worker(rank, init_file, checkpoint_dir, result_queue): + _strict_worker_warnings() + try: + dist.init_process_group( + "gloo", + init_method="file://{}".format(init_file), + rank=rank, + world_size=1, + timeout=timedelta(seconds=60), + ) + initial_parameter = torch.linspace(-0.4, 0.3, 8, dtype=torch.float32).reshape(2, 4) + optimizer, parameter, checkpoint_binding = _make_replicated_target( + initial_parameter, + deterministic=False, + ) + load_portable_dcp( + optimizer, + checkpoint_process_group=checkpoint_binding, + storage_reader=FileSystemReader(checkpoint_dir), + transaction_id="portable-dcp-load-2-to-1", + limits=_limits(), + namespace="optimizer", + ) + restored_exact = _loaded_state_is_exact(optimizer, parameter) + next_step_exact = _next_step_is_exact(optimizer, parameter, initial_parameter) + result_queue.put( + { + "rank": rank, + "restored_exact": restored_exact, + "next_step_exact": next_step_exact, + } + ) + except BaseException: + result_queue.put({"rank": rank, "fatal_error": traceback.format_exc()}) + finally: + if dist.is_available() and dist.is_initialized(): + dist.destroy_process_group() + + +def _run_phase(worker, world_size, *worker_args): + context = mp.get_context("spawn") + result_queue = context.Queue() + descriptor, init_file = tempfile.mkstemp(prefix="gefen-portable-dcp-") + os.close(descriptor) + os.unlink(init_file) + processes = [ + context.Process( + target=worker, + args=(rank, init_file, *worker_args, result_queue), + ) + for rank in range(world_size) + ] + results = [] + try: + for process in processes: + process.start() + try: + for _ in processes: + results.append(result_queue.get(timeout=120)) + except queue_module.Empty: + pass + for process in processes: + process.join(timeout=10) + for process in processes: + if process.is_alive(): + process.terminate() + process.join(timeout=5) + finally: + for process in processes: + if process.is_alive(): + process.terminate() + process.join(timeout=5) + result_queue.close() + result_queue.join_thread() + if os.path.exists(init_file): + os.unlink(init_file) + assert len(results) == world_size, (results, [process.exitcode for process in processes]) + assert all(process.exitcode == 0 for process in processes), [process.exitcode for process in processes] + return sorted(results, key=lambda item: item["rank"]) + + +@pytest.mark.skipif( + not dist.is_available() or not dist.is_gloo_available(), + reason="portable DCP topology-change coverage requires Gloo", +) +def test_portable_dcp_saves_flattened_world_and_loads_replicated_singleton(tmp_path): + checkpoint_dir = str(tmp_path / "checkpoint") + save_results = _run_phase( + _save_worker, + _SAVE_WORLD, + checkpoint_dir, + str(tmp_path / "mismatched-namespace"), + str(tmp_path / "mismatched-storage"), + str(tmp_path / "invalid-writer"), + ) + + assert all("fatal_error" not in result for result in save_results), save_results + assert all(result["saved"] for result in save_results) + namespace_messages = [result["namespace_message"] for result in save_results] + assert namespace_messages[0] == namespace_messages[1] + assert "context" in namespace_messages[0] + storage_messages = [result["storage_message"] for result in save_results] + assert storage_messages[0] == storage_messages[1] + assert "context" in storage_messages[0] + invalid_messages = [result["invalid_message"] for result in save_results] + assert invalid_messages[0] == invalid_messages[1] + assert "StorageWriter" in invalid_messages[0] + assert all( + result["namespace_unchanged"] and result["storage_unchanged"] and result["invalid_unchanged"] + for result in save_results + ) + + load_results = _run_phase(_load_worker, 1, checkpoint_dir) + assert all("fatal_error" not in result for result in load_results), load_results + assert load_results == [{"rank": 0, "restored_exact": True, "next_step_exact": True}] diff --git a/tests/test_portable_dcp_nccl.py b/tests/test_portable_dcp_nccl.py new file mode 100644 index 0000000..e543e4e --- /dev/null +++ b/tests/test_portable_dcp_nccl.py @@ -0,0 +1,364 @@ +"""Warning-strict CUDA/NCCL coverage for portable DCP persistence.""" + +from datetime import timedelta +import multiprocessing as mp +import os +import queue as queue_module +import tempfile +import traceback +import warnings + +import pytest +import torch +import torch.distributed as dist +from torch.distributed.checkpoint import FileSystemReader, FileSystemWriter + +from gefen import load_portable_dcp, save_portable_dcp +from gefen.checkpoint import CheckpointProcessGroupBinding +from gefen.codebook import CodebookProcessGroupBinding +from gefen.contracts import ( + LogicalSlice, + ParameterIdentity, + ParameterLayout, + PlacementKind, + ProcessGroupIdentity, + ShardIdentity, + ShardPlacement, + ShardingManifest, +) +from gefen.gefen import Gefen +from gefen.portable_state import PortableStateLimits +from gefen.rebinding import ParameterRebinding + + +_WORLD = 2 +_FQN = "model.weight" + + +def _limits(): + return PortableStateLimits( + max_fragment_tensor_bytes=1 << 20, + max_collective_tensor_bytes=4 << 20, + max_collective_metadata_bytes=4 << 20, + chunk_bytes=13, + max_members=4, + max_metadata_bytes=1 << 20, + max_tree_nodes=10_000, + max_tree_depth=32, + max_container_items=10_000, + max_string_bytes=16 << 10, + max_integer_bytes=128, + max_tensors=256, + max_tensor_rank=8, + diagnostic_bytes=1024, + ) + + +def _members(): + return tuple("rank:{}".format(rank) for rank in range(_WORLD)) + + +def _replicated_manifest(identity, group): + shards = tuple( + ShardIdentity( + identity, + ParameterLayout.REPLICATED, + LogicalSlice.full(identity), + placements=( + ShardPlacement( + "checkpoint", + PlacementKind.REPLICATE, + coordinate, + len(group.ordered_members), + ), + ), + process_group=group, + local_member=member, + ) + for coordinate, member in enumerate(group.ordered_members) + ) + return ShardingManifest(shards), shards + + +def _optimizer(rank, group, *, deterministic): + device = torch.device("cuda", rank) + parameter = torch.nn.Parameter(torch.linspace(-0.4, 0.3, 8, dtype=torch.float32, device=device).reshape(2, 4)) + optimizer = Gefen( + [("weight", parameter)], + lr=2.5e-3, + betas=(0.8, 0.97), + eps=2.0e-8, + weight_decay=0.03, + fused=False, + force_2d_period_one=True, + factored_v_2d=False, + deterministic=deterministic, + ) + manifest, shards = _replicated_manifest(ParameterIdentity(_FQN, (2, 4)), group) + member = _members()[rank] + codebook_binding = CodebookProcessGroupBinding( + group, + member, + dist.group.WORLD, + device, + ) + checkpoint_binding = CheckpointProcessGroupBinding( + group, + member, + dist.group.WORLD, + device, + ) + optimizer.post_sharding( + (ParameterRebinding(parameter, parameter, shards[rank]),), + manifest=manifest, + codebook_process_group=codebook_binding, + ) + return optimizer, parameter, checkpoint_binding + + +def _codebook(device): + return torch.linspace(-1.0, 1.0, 256, dtype=torch.float32, device=device) + + +def _momentum(device): + return ( + torch.tensor( + [ + -2147483648, + 0, + 1, + -2147483647, + 1056964608, + -1090519040, + 1078984704, + -1058013184, + ], + dtype=torch.int32, + device=device, + ) + .view(torch.float32) + .reshape(2, 4) + .clone() + ) + + +def _second_moment(device): + return ( + torch.tensor( + [ + -2147483648, + 0, + 1, + 8388608, + 1048576000, + 1065353216, + 1073741824, + 2139095039, + ], + dtype=torch.int32, + device=device, + ) + .view(torch.float32) + .reshape(2, 4) + .clone() + ) + + +def _seed_state(optimizer, parameter): + momentum = _momentum(parameter.device) + flat = momentum.reshape(-1) + indices = torch.where( + torch.signbit(flat), + torch.zeros(flat.numel(), dtype=torch.uint8, device=parameter.device), + torch.full((flat.numel(),), 255, dtype=torch.uint8, device=parameter.device), + ) + optimizer._gefen_global_step = 9 + optimizer._gefen_codebook = _codebook(parameter.device) + optimizer.state[parameter].update( + { + "automatic_period": 1, + "step": 7, + "m_codebook": indices.reshape(-1, 1), + "m_magnitude": flat.abs().reshape(-1, 1).clone(), + "vmean": _second_moment(parameter.device).reshape(-1, 1), + "vmean_step": 6, + } + ) + + +def _tensor_bits_equal(left, right): + return ( + torch.is_tensor(left) + and torch.is_tensor(right) + and left.dtype == right.dtype + and tuple(left.shape) == tuple(right.shape) + and torch.equal( + left.detach().cpu().contiguous().view(torch.uint8), + right.detach().cpu().contiguous().view(torch.uint8), + ) + ) + + +def _state_equal(left, right): + if set(left) != set(right): + return False + for key, expected in right.items(): + actual = left[key] + if torch.is_tensor(expected): + if not _tensor_bits_equal(actual, expected): + return False + elif actual != expected: + return False + return True + + +def _optimizer_state_equal(left, left_parameter, right, right_parameter): + return ( + left._gefen_global_step == right._gefen_global_step + and left._deterministic is right._deterministic + and _tensor_bits_equal(left._gefen_codebook, right._gefen_codebook) + and _state_equal(left.state[left_parameter], right.state[right_parameter]) + ) + + +def _strict_worker_warnings(): + warnings.simplefilter("error") + warnings.filterwarnings( + "ignore", + message="TypedStorage is deprecated.*", + category=UserWarning, + ) + + +def _worker(rank, init_file, checkpoint_dir, result_queue): + _strict_worker_warnings() + try: + torch.cuda.set_device(rank) + dist.init_process_group( + "nccl", + init_method="file://{}".format(init_file), + rank=rank, + world_size=_WORLD, + timeout=timedelta(seconds=90), + ) + group = ProcessGroupIdentity("portable_dcp_nccl", _members()) + source, source_parameter, source_binding = _optimizer( + rank, + group, + deterministic=True, + ) + _seed_state(source, source_parameter) + + save_portable_dcp( + source, + checkpoint_process_group=source_binding, + storage_writer=FileSystemWriter(checkpoint_dir), + transaction_id="portable-dcp-nccl-save-v1", + limits=_limits(), + ) + dist.barrier(device_ids=[rank]) + + target, target_parameter, target_binding = _optimizer( + rank, + group, + deterministic=False, + ) + load_portable_dcp( + target, + checkpoint_process_group=target_binding, + storage_reader=FileSystemReader(checkpoint_dir), + transaction_id="portable-dcp-nccl-load-v1", + limits=_limits(), + ) + restored_exact = _tensor_bits_equal(target_parameter, source_parameter) and _optimizer_state_equal( + target, + target_parameter, + source, + source_parameter, + ) + + gradient = torch.tensor( + [0.25, -0.5, 0.75, -1.0, 1.25, -1.5, 1.75, -2.0], + dtype=torch.float32, + device=source_parameter.device, + ).reshape(2, 4) + source_parameter.grad = gradient.clone() + target_parameter.grad = gradient.clone() + source.step() + target.step() + continuation = { + "parameter": _tensor_bits_equal(target_parameter, source_parameter), + "global_step": target._gefen_global_step == source._gefen_global_step, + "deterministic": target._deterministic is source._deterministic, + "codebook": _tensor_bits_equal(target._gefen_codebook, source._gefen_codebook), + "state": _state_equal(target.state[target_parameter], source.state[source_parameter]), + } + next_step_exact = all(continuation.values()) + result_queue.put( + { + "rank": rank, + "restored_exact": restored_exact, + "next_step_exact": next_step_exact, + "continuation": continuation, + } + ) + except BaseException: + result_queue.put({"rank": rank, "fatal_error": traceback.format_exc()}) + finally: + if dist.is_available() and dist.is_initialized(): + dist.destroy_process_group() + + +def _run_workers(checkpoint_dir): + context = mp.get_context("spawn") + result_queue = context.Queue() + descriptor, init_file = tempfile.mkstemp(prefix="gefen-portable-dcp-nccl-") + os.close(descriptor) + os.unlink(init_file) + processes = [ + context.Process( + target=_worker, + args=(rank, init_file, checkpoint_dir, result_queue), + ) + for rank in range(_WORLD) + ] + results = [] + try: + for process in processes: + process.start() + try: + for _ in processes: + results.append(result_queue.get(timeout=180)) + except queue_module.Empty: + pass + for process in processes: + process.join(timeout=10) + for process in processes: + if process.is_alive(): + process.terminate() + process.join(timeout=5) + finally: + for process in processes: + if process.is_alive(): + process.terminate() + process.join(timeout=5) + result_queue.close() + result_queue.join_thread() + if os.path.exists(init_file): + os.unlink(init_file) + assert len(results) == _WORLD, (results, [process.exitcode for process in processes]) + assert all(process.exitcode == 0 for process in processes), [process.exitcode for process in processes] + return sorted(results, key=lambda item: item["rank"]) + + +@pytest.mark.skipif( + not dist.is_available() or not dist.is_nccl_available() or torch.cuda.device_count() < _WORLD, + reason="portable DCP NCCL coverage requires NCCL and two CUDA devices", +) +def test_portable_dcp_round_trip_and_next_step_are_exact_on_nccl(tmp_path): + results = _run_workers(str(tmp_path / "checkpoint")) + + assert all("fatal_error" not in result for result in results), results + assert [result["rank"] for result in results] == [0, 1] + assert all(result["restored_exact"] and result["next_step_exact"] for result in results), results + assert all(all(result["continuation"].values()) for result in results), results diff --git a/tests/test_portable_dcp_topologies.py b/tests/test_portable_dcp_topologies.py new file mode 100644 index 0000000..9dc6a44 --- /dev/null +++ b/tests/test_portable_dcp_topologies.py @@ -0,0 +1,920 @@ +"""Warning-strict CPU/Gloo topology coverage for portable DCP.""" + +from datetime import timedelta +import multiprocessing as mp +import os +import queue as queue_module +import tempfile +import traceback +import warnings + +import pytest +import torch +import torch.distributed as dist +from torch.distributed.checkpoint import FileSystemReader, FileSystemWriter + +from gefen import load_portable_dcp, save_portable_dcp +from gefen.checkpoint import CheckpointProcessGroupBinding +from gefen.codebook import CodebookProcessGroupBinding +from gefen.contracts import ( + LogicalSlice, + ParameterIdentity, + ParameterLayout, + PlacementKind, + ProcessGroupIdentity, + ShardIdentity, + ShardPlacement, + ShardingManifest, +) +from gefen.gefen import Gefen +from gefen.gefen_muon import GefenMuon +from gefen.portable import _decode_quantized_momentum +from gefen.portable_state import PortableStateLimits +from gefen.rebinding import ParameterRebinding + + +_PLAIN_FQN = "model.weight" +_MUON_FQN = "model.matrix" + + +def _limits(): + return PortableStateLimits( + max_fragment_tensor_bytes=1 << 20, + max_collective_tensor_bytes=4 << 20, + max_collective_metadata_bytes=4 << 20, + chunk_bytes=13, + max_members=4, + max_metadata_bytes=1 << 20, + max_tree_nodes=10_000, + max_tree_depth=32, + max_container_items=10_000, + max_string_bytes=16 << 10, + max_integer_bytes=128, + max_tensors=256, + max_tensor_rank=8, + diagnostic_bytes=1024, + ) + + +def _strict_worker_warnings(): + warnings.simplefilter("error") + warnings.filterwarnings( + "ignore", + message="TypedStorage is deprecated.*", + category=UserWarning, + ) + + +def _bindings(group, local_member, process_group): + return ( + CodebookProcessGroupBinding( + group, + local_member, + process_group, + torch.device("cpu"), + ), + CheckpointProcessGroupBinding( + group, + local_member, + process_group, + torch.device("cpu"), + ), + ) + + +def _replicated_manifest(identity, group): + shards = tuple( + ShardIdentity( + identity, + ParameterLayout.REPLICATED, + LogicalSlice.full(identity), + placements=( + ShardPlacement( + "checkpoint", + PlacementKind.REPLICATE, + coordinate, + len(group.ordered_members), + ), + ), + process_group=group, + local_member=member, + ) + for coordinate, member in enumerate(group.ordered_members) + ) + return ShardingManifest(shards), shards + + +def _flat_manifest(identity, group, lengths): + if len(lengths) != len(group.ordered_members) or sum(lengths) != identity.numel: + raise AssertionError("invalid flattened test partition") + offset = 0 + shards = [] + for coordinate, (member, length) in enumerate(zip(group.ordered_members, lengths)): + shards.append( + ShardIdentity( + identity, + ParameterLayout.FLATTENED_ELEMENT_SHARD, + LogicalSlice(offset, length), + placements=( + ShardPlacement( + "checkpoint", + PlacementKind.FLAT_SHARD, + coordinate, + len(group.ordered_members), + ), + ), + process_group=group, + local_member=member, + ) + ) + offset += length + shards = tuple(shards) + return ShardingManifest(shards), shards + + +def _owner_manifest(identity, group, owner): + shards = tuple( + ShardIdentity( + identity, + ParameterLayout.WHOLE_PARAMETER_OWNER, + (LogicalSlice.full(identity) if member == owner else LogicalSlice(0, 0)), + placements=( + ShardPlacement( + "checkpoint", + PlacementKind.WHOLE_PARAMETER_OWNER, + coordinate, + len(group.ordered_members), + ), + ), + process_group=group, + local_member=member, + owner=owner, + ) + for coordinate, member in enumerate(group.ordered_members) + ) + return ShardingManifest(shards), shards + + +def _plain_optimizer(parameter, *, deterministic): + return Gefen( + [("weight", parameter)], + lr=2.5e-3, + betas=(0.8, 0.97), + eps=2.0e-8, + weight_decay=0.03, + fused=False, + force_2d_period_one=True, + factored_v_2d=False, + deterministic=deterministic, + ) + + +def _muon_optimizer(parameter, *, deterministic, normuon): + return GefenMuon( + [("matrix", parameter)], + lr=3.0e-3, + weight_decay=0.02, + momentum=0.85, + nesterov=False, + ns_steps=2, + fused=False, + sharded_mode="distributed", + deterministic=deterministic, + normuon=normuon, + normuon_beta2=0.9, + normuon_eps=3.0e-8, + ) + + +def _codebook(): + return torch.linspace(-1.0, 1.0, 256, dtype=torch.float32) + + +def _plain_initial_parameter(): + return torch.linspace(-0.4, 0.3, 8, dtype=torch.float32).reshape(2, 4) + + +def _plain_momentum(): + return ( + torch.tensor( + [ + -2147483648, + 0, + 1, + -2147483647, + 1056964608, + -1090519040, + 1078984704, + -1058013184, + ], + dtype=torch.int32, + ) + .view(torch.float32) + .reshape(2, 4) + .clone() + ) + + +def _plain_second_moment(): + return ( + torch.tensor( + [ + -2147483648, + 0, + 1, + 8388608, + 1048576000, + 1065353216, + 1073741824, + 2139095039, + ], + dtype=torch.int32, + ) + .view(torch.float32) + .reshape(2, 4) + .clone() + ) + + +def _plain_gradient(): + return torch.tensor( + [0.25, -0.5, 0.75, -1.0, 1.25, -1.5, 1.75, -2.0], + dtype=torch.float32, + ).reshape(2, 4) + + +def _muon_initial_parameter(): + return torch.linspace(-0.3, 0.2, 6, dtype=torch.float32).reshape(3, 2) + + +def _muon_momentum(): + return torch.tensor( + [[-0.75, 0.5], [1.25, -1.5], [2.0, -2.5]], + dtype=torch.float32, + ) + + +def _muon_normuon_v(): + return torch.tensor([[0.5], [1.5], [2.5]], dtype=torch.float32) + + +def _muon_gradient(): + return torch.tensor( + [[0.4, -0.7], [1.1, -1.3], [1.7, -1.9]], + dtype=torch.float32, + ) + + +def _quantized_period_one(momentum): + flat = momentum.reshape(-1) + indices = torch.where( + torch.signbit(flat), + torch.zeros(flat.numel(), dtype=torch.uint8), + torch.full((flat.numel(),), 255, dtype=torch.uint8), + ) + return indices.reshape(-1, 1), flat.abs().reshape(-1, 1).clone() + + +def _seed_plain_state(optimizer, parameter, momentum, second_moment): + indices, magnitudes = _quantized_period_one(momentum) + optimizer._gefen_global_step = 9 + optimizer._gefen_codebook = _codebook() + optimizer.state[parameter].update( + { + "automatic_period": 1, + "step": 7, + "m_codebook": indices, + "m_magnitude": magnitudes, + "vmean": second_moment.reshape(-1, 1).clone(), + "vmean_step": 6, + } + ) + + +def _seed_muon_state(optimizer, parameter, *, normuon): + indices, magnitudes = _quantized_period_one(_muon_momentum()) + optimizer._gefen_global_step = 13 + optimizer._gefen_codebook = _codebook() + state = optimizer.state[parameter] + state.update( + { + "automatic_period": 1, + "step": 11, + "m_codebook": indices, + "m_magnitude": magnitudes, + } + ) + if normuon: + state.update( + { + "normuon_v": _muon_normuon_v(), + "normuon_step": 10, + } + ) + + +def _bits_equal(left, right): + return ( + torch.is_tensor(left) + and torch.is_tensor(right) + and left.dtype == right.dtype + and tuple(left.shape) == tuple(right.shape) + and torch.equal( + left.detach().contiguous().view(torch.uint8), + right.detach().contiguous().view(torch.uint8), + ) + ) + + +def _state_equal(left, right): + if set(left) != set(right): + return False + for key in left: + left_value = left[key] + right_value = right[key] + if torch.is_tensor(left_value) or torch.is_tensor(right_value): + if not _bits_equal(left_value, right_value): + return False + elif type(left_value) is not type(right_value) or left_value != right_value: + return False + return True + + +def _plain_reference(local_parameter, momentum, second_moment): + parameter = torch.nn.Parameter(local_parameter.clone()) + optimizer = _plain_optimizer(parameter, deterministic=True) + _seed_plain_state(optimizer, parameter, momentum, second_moment) + return optimizer, parameter + + +def _muon_reference(*, normuon): + parameter = torch.nn.Parameter(_muon_initial_parameter().clone()) + optimizer = _muon_optimizer( + parameter, + deterministic=True, + normuon=normuon, + ) + _seed_muon_state(optimizer, parameter, normuon=normuon) + return optimizer, parameter + + +def _run_phase(worker, world_size, *worker_args, timeout=180): + context = mp.get_context("spawn") + result_queue = context.Queue() + descriptor, init_file = tempfile.mkstemp(prefix="gefen-portable-dcp-topology-") + os.close(descriptor) + os.unlink(init_file) + processes = [ + context.Process( + target=worker, + args=(rank, world_size, init_file, *worker_args, result_queue), + ) + for rank in range(world_size) + ] + results = [] + try: + for process in processes: + process.start() + try: + for _ in processes: + results.append(result_queue.get(timeout=timeout)) + except queue_module.Empty: + pass + for process in processes: + process.join(timeout=10) + for process in processes: + if process.is_alive(): + process.terminate() + process.join(timeout=5) + finally: + for process in processes: + if process.is_alive(): + process.terminate() + process.join(timeout=5) + result_queue.close() + result_queue.join_thread() + if os.path.exists(init_file): + os.unlink(init_file) + assert len(results) == world_size, (results, [process.exitcode for process in processes]) + assert all(process.exitcode == 0 for process in processes), [process.exitcode for process in processes] + return sorted(results, key=lambda item: item["rank"]) + + +def _plain_singleton_save_worker( + rank, + world_size, + init_file, + checkpoint_dir, + result_queue, +): + _strict_worker_warnings() + try: + dist.init_process_group( + "gloo", + init_method="file://{}".format(init_file), + rank=rank, + world_size=world_size, + timeout=timedelta(seconds=60), + ) + member = "save:0" + group = ProcessGroupIdentity("plain_singleton_save", (member,)) + identity = ParameterIdentity(_PLAIN_FQN, (2, 4)) + parameter = torch.nn.Parameter(_plain_initial_parameter().clone()) + optimizer = _plain_optimizer(parameter, deterministic=True) + manifest, shards = _replicated_manifest(identity, group) + codebook_binding, checkpoint_binding = _bindings(group, member, None) + optimizer.post_sharding( + (ParameterRebinding(parameter, parameter, shards[0]),), + manifest=manifest, + codebook_process_group=codebook_binding, + ) + _seed_plain_state( + optimizer, + parameter, + _plain_momentum(), + _plain_second_moment(), + ) + save_portable_dcp( + optimizer, + checkpoint_process_group=checkpoint_binding, + storage_writer=FileSystemWriter(checkpoint_dir), + transaction_id="plain-singleton-to-flat-save", + limits=_limits(), + ) + result_queue.put({"rank": rank, "saved": True}) + except BaseException: + result_queue.put({"rank": rank, "fatal_error": traceback.format_exc()}) + finally: + if dist.is_available() and dist.is_initialized(): + dist.destroy_process_group() + + +def _plain_flat_load_worker( + rank, + world_size, + init_file, + checkpoint_dir, + result_queue, +): + _strict_worker_warnings() + try: + dist.init_process_group( + "gloo", + init_method="file://{}".format(init_file), + rank=rank, + world_size=world_size, + timeout=timedelta(seconds=60), + ) + members = tuple("load:{}".format(index) for index in range(world_size)) + group = ProcessGroupIdentity("plain_flat_load", members) + identity = ParameterIdentity(_PLAIN_FQN, (2, 4)) + lengths = (3, 5) + manifest, shards = _flat_manifest(identity, group, lengths) + shard = shards[rank] + start = shard.logical_slice.flat_offset + stop = start + shard.logical_slice.length + original = torch.nn.Parameter(_plain_initial_parameter().clone()) + local_initial = _plain_initial_parameter().reshape(-1)[start:stop].clone() + local = torch.nn.Parameter(local_initial.clone()) + optimizer = _plain_optimizer(original, deterministic=False) + codebook_binding, checkpoint_binding = _bindings( + group, + members[rank], + dist.group.WORLD, + ) + optimizer.post_sharding( + (ParameterRebinding(original, local, shard),), + manifest=manifest, + codebook_process_group=codebook_binding, + ) + load_portable_dcp( + optimizer, + checkpoint_process_group=checkpoint_binding, + storage_reader=FileSystemReader(checkpoint_dir), + transaction_id="plain-singleton-to-flat-load", + limits=_limits(), + ) + + expected_momentum = _plain_momentum().reshape(-1)[start:stop].clone() + expected_second = _plain_second_moment().reshape(-1)[start:stop].clone() + state = optimizer.state[local] + decoded = _decode_quantized_momentum( + optimizer._gefen_codebook, + state["m_codebook"], + state["m_magnitude"], + logical_shape=(lengths[rank],), + period=1, + step=state["step"], + ) + restored_exact = ( + _bits_equal(decoded, expected_momentum) + and _bits_equal(state["vmean"].reshape(-1), expected_second) + and state["vmean_step"] == 6 + and state["step"] == 7 + and optimizer._gefen_global_step == 9 + and optimizer._deterministic is True + ) + + reference, reference_parameter = _plain_reference( + local_initial, + expected_momentum, + expected_second, + ) + gradient = _plain_gradient().reshape(-1)[start:stop].clone() + local.grad = gradient.clone() + reference_parameter.grad = gradient.clone() + optimizer.step() + reference.step() + continuation = { + "parameter": _bits_equal(local, reference_parameter), + "state": _state_equal( + optimizer.state[local], + reference.state[reference_parameter], + ), + "global_step": optimizer._gefen_global_step == reference._gefen_global_step, + "codebook": _bits_equal( + optimizer._gefen_codebook, + reference._gefen_codebook, + ), + } + next_step_exact = all(continuation.values()) + result_queue.put( + { + "rank": rank, + "restored_exact": restored_exact, + "next_step_exact": next_step_exact, + "continuation": continuation, + "length": lengths[rank], + } + ) + except BaseException: + result_queue.put({"rank": rank, "fatal_error": traceback.format_exc()}) + finally: + if dist.is_available() and dist.is_initialized(): + dist.destroy_process_group() + + +def _muon_singleton_save_worker( + rank, + world_size, + init_file, + checkpoint_dir, + normuon, + result_queue, +): + _strict_worker_warnings() + try: + dist.init_process_group( + "gloo", + init_method="file://{}".format(init_file), + rank=rank, + world_size=world_size, + timeout=timedelta(seconds=60), + ) + member = "save:0" + group = ProcessGroupIdentity("muon_singleton_save", (member,)) + identity = ParameterIdentity(_MUON_FQN, (3, 2)) + parameter = torch.nn.Parameter(_muon_initial_parameter().clone()) + optimizer = _muon_optimizer( + parameter, + deterministic=True, + normuon=normuon, + ) + manifest, shards = _replicated_manifest(identity, group) + codebook_binding, checkpoint_binding = _bindings(group, member, None) + optimizer.post_sharding( + (ParameterRebinding(parameter, parameter, shards[0]),), + manifest=manifest, + codebook_process_group=codebook_binding, + ) + _seed_muon_state(optimizer, parameter, normuon=normuon) + save_portable_dcp( + optimizer, + checkpoint_process_group=checkpoint_binding, + storage_writer=FileSystemWriter(checkpoint_dir), + transaction_id="muon-singleton-to-owner-save-{}".format(int(normuon)), + limits=_limits(), + ) + result_queue.put({"rank": rank, "saved": True}) + except BaseException: + result_queue.put({"rank": rank, "fatal_error": traceback.format_exc()}) + finally: + if dist.is_available() and dist.is_initialized(): + dist.destroy_process_group() + + +def _muon_owner_load_worker( + rank, + world_size, + init_file, + checkpoint_dir, + normuon, + result_queue, +): + _strict_worker_warnings() + try: + dist.init_process_group( + "gloo", + init_method="file://{}".format(init_file), + rank=rank, + world_size=world_size, + timeout=timedelta(seconds=60), + ) + members = tuple("load:{}".format(index) for index in range(world_size)) + owner = members[-1] + group = ProcessGroupIdentity("muon_owner_load", members) + identity = ParameterIdentity(_MUON_FQN, (3, 2)) + manifest, shards = _owner_manifest(identity, group, owner) + original = torch.nn.Parameter(_muon_initial_parameter().clone()) + optimizer = _muon_optimizer( + original, + deterministic=False, + normuon=normuon, + ) + local = original if members[rank] == owner else None + codebook_binding, checkpoint_binding = _bindings( + group, + members[rank], + dist.group.WORLD, + ) + optimizer.post_sharding( + (ParameterRebinding(original, local, shards[rank]),), + manifest=manifest, + codebook_process_group=codebook_binding, + ) + load_portable_dcp( + optimizer, + checkpoint_process_group=checkpoint_binding, + storage_reader=FileSystemReader(checkpoint_dir), + transaction_id="muon-singleton-to-owner-load-{}".format(int(normuon)), + limits=_limits(), + ) + + if local is None: + restored_exact = ( + not optimizer.param_groups[0]["params"] + and not optimizer.state + and optimizer._gefen_global_step == 13 + and optimizer._deterministic is True + and _bits_equal(optimizer._gefen_codebook, _codebook()) + ) + else: + state = optimizer.state[local] + decoded = _decode_quantized_momentum( + optimizer._gefen_codebook, + state["m_codebook"], + state["m_magnitude"], + logical_shape=(3, 2), + period=1, + step=state["step"], + ) + restored_exact = ( + _bits_equal(decoded, _muon_momentum()) + and state["step"] == 11 + and optimizer._gefen_global_step == 13 + and optimizer._deterministic is True + and ( + not normuon or (_bits_equal(state["normuon_v"], _muon_normuon_v()) and state["normuon_step"] == 10) + ) + ) + + reference = None + reference_parameter = None + if local is not None: + reference, reference_parameter = _muon_reference(normuon=normuon) + local.grad = _muon_gradient().clone() + reference_parameter.grad = _muon_gradient().clone() + optimizer.step() + if reference is not None: + reference.step() + continuation = { + "parameter": _bits_equal(local, reference_parameter), + "state": _state_equal( + optimizer.state[local], + reference.state[reference_parameter], + ), + "global_step": optimizer._gefen_global_step == reference._gefen_global_step, + "codebook": _bits_equal( + optimizer._gefen_codebook, + reference._gefen_codebook, + ), + } + next_step_exact = all(continuation.values()) + else: + continuation = { + "empty_nonowner": not optimizer.param_groups[0]["params"] and not optimizer.state, + "global_step": optimizer._gefen_global_step == 14, + } + next_step_exact = all(continuation.values()) + result_queue.put( + { + "rank": rank, + "owner": members[rank] == owner, + "restored_exact": restored_exact, + "next_step_exact": next_step_exact, + "continuation": continuation, + } + ) + except BaseException: + result_queue.put({"rank": rank, "fatal_error": traceback.format_exc()}) + finally: + if dist.is_available() and dist.is_initialized(): + dist.destroy_process_group() + + +def _subgroup_worker( + rank, + world_size, + init_file, + checkpoint_dir, + result_queue, +): + _strict_worker_warnings() + try: + dist.init_process_group( + "gloo", + init_method="file://{}".format(init_file), + rank=rank, + world_size=world_size, + timeout=timedelta(seconds=90), + ) + subgroup_ranks = (0, 2) + subgroup = dist.new_group( + ranks=list(subgroup_ranks), + backend="gloo", + timeout=timedelta(seconds=60), + ) + if rank not in subgroup_ranks: + dist.barrier() + result_queue.put({"rank": rank, "excluded": True}) + return + + coordinate = subgroup_ranks.index(rank) + members = tuple("subgroup:{}".format(item) for item in subgroup_ranks) + group = ProcessGroupIdentity("portable_dcp_subgroup", members) + identity = ParameterIdentity(_PLAIN_FQN, (2, 4)) + manifest, shards = _replicated_manifest(identity, group) + source_parameter = torch.nn.Parameter(_plain_initial_parameter().clone()) + source = _plain_optimizer(source_parameter, deterministic=True) + source_codebook, source_checkpoint = _bindings( + group, + members[coordinate], + subgroup, + ) + source.post_sharding( + ( + ParameterRebinding( + source_parameter, + source_parameter, + shards[coordinate], + ), + ), + manifest=manifest, + codebook_process_group=source_codebook, + ) + _seed_plain_state( + source, + source_parameter, + _plain_momentum(), + _plain_second_moment(), + ) + save_portable_dcp( + source, + checkpoint_process_group=source_checkpoint, + storage_writer=FileSystemWriter(checkpoint_dir), + transaction_id="supported-subgroup-save", + limits=_limits(), + ) + + target_parameter = torch.nn.Parameter(_plain_initial_parameter().clone()) + target = _plain_optimizer(target_parameter, deterministic=False) + target_codebook, target_checkpoint = _bindings( + group, + members[coordinate], + subgroup, + ) + target.post_sharding( + ( + ParameterRebinding( + target_parameter, + target_parameter, + shards[coordinate], + ), + ), + manifest=manifest, + codebook_process_group=target_codebook, + ) + load_portable_dcp( + target, + checkpoint_process_group=target_checkpoint, + storage_reader=FileSystemReader(checkpoint_dir), + transaction_id="supported-subgroup-load", + limits=_limits(), + ) + restored_exact = ( + _state_equal(target.state[target_parameter], source.state[source_parameter]) + and target._gefen_global_step == source._gefen_global_step + and target._deterministic is True + and _bits_equal(target._gefen_codebook, source._gefen_codebook) + ) + + reference, reference_parameter = _plain_reference( + _plain_initial_parameter(), + _plain_momentum(), + _plain_second_moment(), + ) + target_parameter.grad = _plain_gradient().clone() + reference_parameter.grad = _plain_gradient().clone() + target.step() + reference.step() + continuation = { + "parameter": _bits_equal(target_parameter, reference_parameter), + "state": _state_equal( + target.state[target_parameter], + reference.state[reference_parameter], + ), + "global_step": target._gefen_global_step == reference._gefen_global_step, + } + next_step_exact = all(continuation.values()) + dist.barrier() + result_queue.put( + { + "rank": rank, + "excluded": False, + "restored_exact": restored_exact, + "next_step_exact": next_step_exact, + "continuation": continuation, + } + ) + except BaseException: + result_queue.put({"rank": rank, "fatal_error": traceback.format_exc()}) + finally: + if dist.is_available() and dist.is_initialized(): + dist.destroy_process_group() + + +@pytest.mark.skipif( + not dist.is_available() or not dist.is_gloo_available(), + reason="portable DCP topology coverage requires Gloo", +) +def test_portable_dcp_replicated_singleton_to_uneven_flattened_world(tmp_path): + checkpoint_dir = str(tmp_path / "plain-singleton-to-flat") + save_results = _run_phase( + _plain_singleton_save_worker, + 1, + checkpoint_dir, + ) + assert save_results == [{"rank": 0, "saved": True}] + + load_results = _run_phase( + _plain_flat_load_worker, + 2, + checkpoint_dir, + ) + assert all("fatal_error" not in result for result in load_results), load_results + assert [result["length"] for result in load_results] == [3, 5] + assert all(result["restored_exact"] and result["next_step_exact"] for result in load_results), load_results + + +@pytest.mark.parametrize("normuon", [False, True]) +@pytest.mark.skipif( + not dist.is_available() or not dist.is_gloo_available(), + reason="portable DCP topology coverage requires Gloo", +) +def test_portable_dcp_muon_replicated_singleton_to_world_owner(tmp_path, normuon): + checkpoint_dir = str(tmp_path / "muon-singleton-to-owner-{}".format(int(normuon))) + save_results = _run_phase( + _muon_singleton_save_worker, + 1, + checkpoint_dir, + normuon, + ) + assert save_results == [{"rank": 0, "saved": True}] + + load_results = _run_phase( + _muon_owner_load_worker, + 2, + checkpoint_dir, + normuon, + ) + assert all("fatal_error" not in result for result in load_results), load_results + assert [result["owner"] for result in load_results] == [False, True] + assert all(result["restored_exact"] and result["next_step_exact"] for result in load_results), load_results + + +@pytest.mark.skipif( + not dist.is_available() or not dist.is_gloo_available(), + reason="portable DCP subgroup coverage requires Gloo", +) +def test_portable_dcp_supported_adapter_defined_subgroup(tmp_path): + results = _run_phase( + _subgroup_worker, + 3, + str(tmp_path / "supported-subgroup"), + timeout=240, + ) + assert all("fatal_error" not in result for result in results), results + assert results[1] == {"rank": 1, "excluded": True} + participants = (results[0], results[2]) + assert all(not result["excluded"] for result in participants) + assert all(result["restored_exact"] and result["next_step_exact"] for result in participants), results diff --git a/tests/test_portable_fields.py b/tests/test_portable_fields.py new file mode 100644 index 0000000..abc15d8 --- /dev/null +++ b/tests/test_portable_fields.py @@ -0,0 +1,538 @@ +"""CPU coverage for pure dense logical-field assembly and projection.""" + +import pytest +import torch + +from gefen.contracts import ( + LogicalRegion, + LogicalSlice, + ParameterIdentity, + ParameterLayout, + PlacementKind, + ProcessGroupIdentity, + ShardIdentity, + ShardPlacement, + ShardingManifest, +) +from gefen.portable_fields import ( + _assemble_dense_logical_field, + _project_dense_logical_field, +) + + +def _group(members=("worker:c", "worker:a", "worker:b")): + return ProcessGroupIdentity("data_parallel", members) + + +def _placement(group, member, kind, *, dimension=None): + return ShardPlacement( + "data_parallel", + kind, + group.ordered_members.index(member), + len(group.ordered_members), + dimension, + ) + + +def _local_replicated(parameter): + return ( + ShardIdentity( + parameter, + ParameterLayout.REPLICATED, + LogicalSlice.full(parameter), + ), + ) + + +def _grouped_replicated(parameter, group): + return tuple( + ShardIdentity( + parameter, + ParameterLayout.REPLICATED, + LogicalSlice.full(parameter), + placements=(_placement(group, member, PlacementKind.REPLICATE),), + process_group=group, + local_member=member, + ) + for member in group.ordered_members + ) + + +def _flat_shards(parameter, group, lengths): + offset = 0 + result = [] + for member, length in zip(group.ordered_members, lengths): + result.append( + ShardIdentity( + parameter, + ParameterLayout.FLATTENED_ELEMENT_SHARD, + LogicalSlice(offset, length), + placements=(_placement(group, member, PlacementKind.FLAT_SHARD),), + process_group=group, + local_member=member, + ) + ) + offset += length + return tuple(result) + + +def _owner_shards(parameter, group, owner): + return tuple( + ShardIdentity( + parameter, + ParameterLayout.WHOLE_PARAMETER_OWNER, + (LogicalSlice.full(parameter) if member == owner else LogicalSlice(0, 0)), + placements=( + _placement( + group, + member, + PlacementKind.WHOLE_PARAMETER_OWNER, + ), + ), + process_group=group, + local_member=member, + owner=owner, + ) + for member in group.ordered_members + ) + + +def _dtensor_shards(parameter, group, lengths, *, dimension): + offset = 0 + result = [] + for member, length in zip(group.ordered_members, lengths): + offsets = [0] * len(parameter.global_shape) + region_lengths = list(parameter.global_shape) + offsets[dimension] = offset + region_lengths[dimension] = length + result.append( + ShardIdentity( + parameter, + ParameterLayout.DTENSOR_1D_DEFAULT_WORLD, + LogicalRegion(offsets, region_lengths), + placements=( + _placement( + group, + member, + PlacementKind.DIMENSION_SHARD, + dimension=dimension, + ), + ), + process_group=group, + local_member=member, + ) + ) + offset += length + return tuple(result) + + +def _dtensor_replicated(parameter, group): + return tuple( + ShardIdentity( + parameter, + ParameterLayout.DTENSOR_1D_DEFAULT_WORLD, + LogicalRegion.full(parameter), + placements=(_placement(group, member, PlacementKind.REPLICATE),), + process_group=group, + local_member=member, + ) + for member in group.ordered_members + ) + + +def _assert_tight_cpu_fp32(value, shape): + assert type(value) is torch.Tensor + assert value.device.type == "cpu" + assert value.dtype == torch.float32 + assert tuple(value.shape) == tuple(shape) + assert not value.requires_grad + assert value.is_contiguous() + assert value.storage_offset() == 0 + assert value.untyped_storage().nbytes() == value.numel() * value.element_size() + assert bool(torch.isfinite(value).all()) + + +def _sparse_tensor(shape): + with torch.sparse.check_sparse_tensor_invariants(enable=True): + return torch.sparse_coo_tensor( + torch.tensor([[0], [1]]), + torch.tensor([1.0]), + size=shape, + ) + + +def _round_trip(parameter, shards, logical): + manifest = ShardingManifest(tuple(reversed(shards))) + fragments = tuple( + ( + shard, + _project_dense_logical_field(parameter, logical, shard), + ) + for shard in reversed(shards) + ) + assembled = _assemble_dense_logical_field( + manifest, + parameter, + fragments, + ) + _assert_tight_cpu_fp32(assembled, parameter.global_shape) + assert torch.equal(assembled, logical) + return assembled, fragments + + +def test_local_and_grouped_replicated_fields_round_trip_without_aliasing(): + parameter = ParameterIdentity("layer.weight", (2, 3)) + logical = torch.arange(12, dtype=torch.float32).reshape(2, 6)[:, ::2] + assert not logical.is_contiguous() + + for shards in ( + _local_replicated(parameter), + _grouped_replicated(parameter, _group()), + ): + assembled, fragments = _round_trip(parameter, shards, logical) + snapshots = tuple(payload.clone() for _, payload in fragments if payload is not None) + assembled.add_(1000) + assert torch.equal(logical, torch.tensor([[0.0, 2.0, 4.0], [6.0, 8.0, 10.0]])) + assert all( + torch.equal(payload, snapshot) + for (_, payload), snapshot in zip( + ((shard, payload) for shard, payload in fragments if payload is not None), + snapshots, + ) + ) + + +def test_flattened_uneven_and_empty_shards_preserve_row_major_order(): + parameter = ParameterIdentity("layer.weight", (2, 3)) + group = _group(("rank:0", "rank:1", "rank:2", "rank:3")) + shards = _flat_shards(parameter, group, (2, 0, 3, 1)) + logical = torch.tensor([[10.0, 11.0, 12.0], [20.0, 21.0, 22.0]]) + + assembled, fragments = _round_trip(parameter, shards, logical) + + assert torch.equal(assembled, logical) + by_member = {shard.local_member: payload for shard, payload in fragments} + assert torch.equal(by_member["rank:0"], torch.tensor([10.0, 11.0])) + assert by_member["rank:1"] is None + assert torch.equal(by_member["rank:2"], torch.tensor([12.0, 20.0, 21.0])) + assert torch.equal(by_member["rank:3"], torch.tensor([22.0])) + + +def test_whole_parameter_owner_projects_only_to_owner_and_round_trips(): + parameter = ParameterIdentity("layer.weight", (2, 3)) + group = _group() + shards = _owner_shards(parameter, group, "worker:a") + logical = torch.arange(6, dtype=torch.float32).reshape(2, 3).requires_grad_() + + assembled, fragments = _round_trip(parameter, shards, logical) + + assert torch.equal(assembled, logical.detach()) + by_member = {shard.local_member: payload for shard, payload in fragments} + assert by_member["worker:c"] is None + assert by_member["worker:b"] is None + _assert_tight_cpu_fp32(by_member["worker:a"], parameter.global_shape) + + +@pytest.mark.parametrize( + "shape,lengths,dimension", + ( + ((5, 4), (2, 0, 3), 0), + ((3, 5), (2, 2, 1), 1), + ), +) +def test_dtensor_row_and_column_regions_round_trip_uneven_and_empty_shards( + shape, + lengths, + dimension, +): + parameter = ParameterIdentity("layer.weight", shape) + group = _group() + shards = _dtensor_shards( + parameter, + group, + lengths, + dimension=dimension, + ) + base = torch.arange(shape[0] * (shape[1] * 2), dtype=torch.float32).reshape( + shape[0], + shape[1] * 2, + ) + logical = base[:, ::2] + assert not logical.is_contiguous() + + assembled, fragments = _round_trip(parameter, shards, logical) + + assert torch.equal(assembled, logical) + for shard, payload in fragments: + region = shard.logical_region + if region.numel == 0: + assert payload is None + else: + _assert_tight_cpu_fp32(payload, region.lengths) + + +def test_replicated_dtensor_fields_require_equal_full_replicas(): + parameter = ParameterIdentity("layer.weight", (2, 3)) + shards = _dtensor_replicated(parameter, _group()) + logical = torch.arange(6, dtype=torch.float32).reshape(2, 3) + + assembled, fragments = _round_trip(parameter, shards, logical) + + assert torch.equal(assembled, logical) + assert all(torch.equal(payload, logical) for _, payload in fragments) + + +def test_manifest_may_contain_other_parameters_but_fragments_must_not(): + parameter = ParameterIdentity("wanted.weight", (4,)) + shards = _flat_shards(parameter, _group(), (1, 1, 2)) + other_parameter = ParameterIdentity("other.weight", (1,)) + other_shard = _local_replicated(other_parameter)[0] + manifest = ShardingManifest(shards + (other_shard,)) + logical = torch.arange(4, dtype=torch.float32) + fragments = tuple((shard, _project_dense_logical_field(parameter, logical, shard)) for shard in shards) + + assert torch.equal( + _assemble_dense_logical_field(manifest, parameter, fragments), + logical, + ) + with pytest.raises(ValueError, match="exactly cover"): + _assemble_dense_logical_field( + manifest, + parameter, + fragments + ((other_shard, torch.ones(1)),), + ) + + +def test_fragments_reject_missing_extra_duplicate_and_malformed_identities(): + parameter = ParameterIdentity("layer.weight", (3,)) + shards = _flat_shards(parameter, _group(), (1, 1, 1)) + manifest = ShardingManifest(shards) + logical = torch.arange(3, dtype=torch.float32) + fragments = tuple((shard, _project_dense_logical_field(parameter, logical, shard)) for shard in shards) + + with pytest.raises(ValueError, match="exactly cover"): + _assemble_dense_logical_field(manifest, parameter, fragments[:-1]) + with pytest.raises(ValueError, match="duplicate"): + _assemble_dense_logical_field(manifest, parameter, fragments + (fragments[0],)) + different = _local_replicated(ParameterIdentity("different.weight", (3,)))[0] + with pytest.raises(ValueError, match="exactly cover"): + _assemble_dense_logical_field( + manifest, + parameter, + fragments[:-1] + ((different, torch.arange(3, dtype=torch.float32)),), + ) + with pytest.raises(TypeError, match="pair"): + _assemble_dense_logical_field(manifest, parameter, ((shards[0],),)) + with pytest.raises(TypeError, match="ShardIdentity"): + _assemble_dense_logical_field(manifest, parameter, (("not-a-shard", None),)) + + +def test_requested_parameter_must_match_the_manifest_identity_exactly(): + parameter = ParameterIdentity("layer.weight", (3,)) + shard = _local_replicated(parameter)[0] + manifest = ShardingManifest((shard,)) + + with pytest.raises(ValueError, match="does not contain"): + _assemble_dense_logical_field( + manifest, + ParameterIdentity("other.weight", (3,)), + ((shard, torch.arange(3, dtype=torch.float32)),), + ) + with pytest.raises(ValueError, match="does not match"): + _assemble_dense_logical_field( + manifest, + ParameterIdentity("layer.weight", (4,)), + ((shard, torch.arange(3, dtype=torch.float32)),), + ) + + +def test_replicas_reject_absent_and_disagreeing_payloads(): + parameter = ParameterIdentity("layer.weight", (2, 2)) + shards = _grouped_replicated(parameter, _group()) + manifest = ShardingManifest(shards) + logical = torch.arange(4, dtype=torch.float32).reshape(2, 2) + fragments = [(shard, logical.clone()) for shard in shards] + + fragments[1] = (shards[1], None) + with pytest.raises(ValueError, match="every shard"): + _assemble_dense_logical_field(manifest, parameter, fragments) + fragments[1] = (shards[1], logical.clone()) + fragments[2][1][1, 1] += 1 + with pytest.raises(ValueError, match="disagree"): + _assemble_dense_logical_field(manifest, parameter, fragments) + + +def test_replicas_require_bit_exact_signed_zero_payloads(): + parameter = ParameterIdentity("layer.weight", (1,)) + shards = _grouped_replicated(parameter, _group()) + manifest = ShardingManifest(shards) + fragments = [(shard, torch.tensor([0.0])) for shard in shards] + fragments[1] = (shards[1], torch.tensor([-0.0])) + + with pytest.raises(ValueError, match="disagree"): + _assemble_dense_logical_field(manifest, parameter, fragments) + + +def test_flattened_and_dtensor_empty_payloads_are_explicit_none(): + flat_parameter = ParameterIdentity("flat.weight", (2,)) + flat_group = _group() + flat_shards = _flat_shards(flat_parameter, flat_group, (1, 0, 1)) + flat_manifest = ShardingManifest(flat_shards) + flat_fragments = ( + (flat_shards[0], torch.tensor([1.0])), + (flat_shards[1], torch.empty(0)), + (flat_shards[2], torch.tensor([2.0])), + ) + with pytest.raises(ValueError, match="empty flattened"): + _assemble_dense_logical_field( + flat_manifest, + flat_parameter, + flat_fragments, + ) + missing_flat = list(flat_fragments) + missing_flat[0] = (flat_shards[0], None) + missing_flat[1] = (flat_shards[1], None) + with pytest.raises(ValueError, match="nonempty flattened"): + _assemble_dense_logical_field( + flat_manifest, + flat_parameter, + missing_flat, + ) + + dtensor_parameter = ParameterIdentity("dtensor.weight", (2, 2)) + dtensor_shards = _dtensor_shards( + dtensor_parameter, + _group(), + (1, 1, 0), + dimension=1, + ) + dtensor_manifest = ShardingManifest(dtensor_shards) + logical = torch.arange(4, dtype=torch.float32).reshape(2, 2) + dtensor_fragments = [ + ( + shard, + _project_dense_logical_field(dtensor_parameter, logical, shard), + ) + for shard in dtensor_shards + ] + dtensor_fragments[2] = (dtensor_shards[2], torch.empty((2, 0))) + with pytest.raises(ValueError, match="empty DTensor"): + _assemble_dense_logical_field( + dtensor_manifest, + dtensor_parameter, + dtensor_fragments, + ) + dtensor_fragments[2] = (dtensor_shards[2], None) + dtensor_fragments[0] = (dtensor_shards[0], None) + with pytest.raises(ValueError, match="nonempty DTensor"): + _assemble_dense_logical_field( + dtensor_manifest, + dtensor_parameter, + dtensor_fragments, + ) + + +def test_whole_parameter_payload_is_present_exactly_on_the_owner(): + parameter = ParameterIdentity("layer.weight", (2, 2)) + shards = _owner_shards(parameter, _group(), "worker:a") + manifest = ShardingManifest(shards) + logical = torch.arange(4, dtype=torch.float32).reshape(2, 2) + fragments = [ + ( + shard, + logical.clone() if shard.local_member == shard.owner else None, + ) + for shard in shards + ] + owner_index = next(index for index, shard in enumerate(shards) if shard.local_member == shard.owner) + + missing = list(fragments) + missing[owner_index] = (shards[owner_index], None) + with pytest.raises(ValueError, match="owner requires"): + _assemble_dense_logical_field(manifest, parameter, missing) + nonowner_index = next(index for index, shard in enumerate(shards) if shard.local_member != shard.owner) + extra = list(fragments) + extra[nonowner_index] = (shards[nonowner_index], torch.empty(0)) + with pytest.raises(ValueError, match="non-owners"): + _assemble_dense_logical_field(manifest, parameter, extra) + + +@pytest.mark.parametrize( + "bad_value,error", + ( + (torch.ones((2, 3), dtype=torch.float64), TypeError), + (torch.ones(6, dtype=torch.float32), ValueError), + (torch.tensor([[1.0, 2.0, 3.0], [4.0, float("inf"), 6.0]]), ValueError), + (torch.nn.Parameter(torch.ones((2, 3), dtype=torch.float32)), TypeError), + (_sparse_tensor((2, 3)), TypeError), + (torch.empty((2, 3), device="meta"), TypeError), + ), +) +def test_assembly_rejects_wrong_dtype_shape_nonfinite_subclass_and_sparse( + bad_value, + error, +): + parameter = ParameterIdentity("layer.weight", (2, 3)) + shard = _local_replicated(parameter)[0] + with pytest.raises(error): + _assemble_dense_logical_field( + ShardingManifest((shard,)), + parameter, + ((shard, bad_value),), + ) + + +@pytest.mark.parametrize( + "bad_value,error", + ( + (torch.ones((2, 3), dtype=torch.float64), TypeError), + (torch.ones(6, dtype=torch.float32), ValueError), + (torch.tensor([[1.0, 2.0, 3.0], [4.0, float("nan"), 6.0]]), ValueError), + (torch.nn.Parameter(torch.ones((2, 3), dtype=torch.float32)), TypeError), + (_sparse_tensor((2, 3)), TypeError), + ), +) +def test_projection_rejects_wrong_dtype_shape_nonfinite_subclass_and_sparse( + bad_value, + error, +): + parameter = ParameterIdentity("layer.weight", (2, 3)) + shard = _local_replicated(parameter)[0] + with pytest.raises(error): + _project_dense_logical_field(parameter, bad_value, shard) + + +def test_projection_requires_an_exact_target_parameter_identity(): + parameter = ParameterIdentity("layer.weight", (2, 3)) + logical = torch.ones((2, 3), dtype=torch.float32) + other_shard = _local_replicated(ParameterIdentity("other.weight", (2, 3)))[0] + + with pytest.raises(ValueError, match="does not match"): + _project_dense_logical_field(parameter, logical, other_shard) + with pytest.raises(TypeError, match="ShardIdentity"): + _project_dense_logical_field(parameter, logical, "not-a-shard") + with pytest.raises(TypeError, match="ParameterIdentity"): + _project_dense_logical_field("layer.weight", logical, other_shard) + + +def test_api_requires_valid_manifest_parameter_and_fragment_containers(): + parameter = ParameterIdentity("layer.weight", (1,)) + shard = _local_replicated(parameter)[0] + manifest = ShardingManifest((shard,)) + + with pytest.raises(TypeError, match="ShardingManifest"): + _assemble_dense_logical_field( + "not-a-manifest", + parameter, + ((shard, torch.ones(1)),), + ) + with pytest.raises(TypeError, match="ParameterIdentity"): + _assemble_dense_logical_field( + manifest, + "layer.weight", + ((shard, torch.ones(1)),), + ) + with pytest.raises(TypeError, match="sequence"): + _assemble_dense_logical_field(manifest, parameter, None) + with pytest.raises(TypeError, match="sequence"): + _assemble_dense_logical_field(manifest, parameter, "fragment") diff --git a/tests/test_portable_hybrid.py b/tests/test_portable_hybrid.py new file mode 100644 index 0000000..44466c3 --- /dev/null +++ b/tests/test_portable_hybrid.py @@ -0,0 +1,197 @@ +"""Strict composite-envelope coverage for Gefen-backed Hybrid portable state.""" + +import copy +import io + +import pytest +import torch + +from gefen.portable_hybrid import ( + HYBRID_PORTABLE_STATE_COVERAGE, + HYBRID_PORTABLE_STATE_FORMAT, + HYBRID_PORTABLE_STATE_FORMAT_VERSION, + HYBRID_PORTABLE_STATE_IMPLEMENTATION, + build_hybrid_portable_state_document, + normalize_hybrid_portable_state_document, +) +from gefen.portable_schema import build_portable_state_document + + +def _parameter_record(fqn, tensor): + return { + "identity": { + "schema_version": 1, + "fqn": fqn, + "global_shape": list(tensor.shape), + }, + "algorithm_options": {"period": 1}, + "state_variant": "period_selected", + "state": {"step": 3, "momentum": tensor}, + "projection_hints": {}, + } + + +def _child(role, *, step=7, deterministic=True, fqn=None, tensor=None): + if fqn is None: + fqn = "Model.Weight" if role == "muon" else "Model.Bias" + if tensor is None: + tensor = torch.arange(6, dtype=torch.float32).reshape(2, 3) + implementation = "gefen.GefenMuon" if role == "muon" else "gefen.Gefen" + return build_portable_state_document( + implementation=implementation, + policy={"role": role}, + common={ + "gefen_global_step": step, + "gefen_deterministic": deterministic, + }, + parameters={fqn: _parameter_record(fqn, tensor)}, + provenance={"source_layouts": ["replicated"]}, + ) + + +def _document(*, muon=None, backup=None): + if muon is None: + muon = _child("muon") + if backup is None: + backup = _child("backup") + return build_hybrid_portable_state_document( + backup_optimizer="gefen", + routing={ + "Model.Weight": "muon", + "Model.Bias": "backup", + }, + children={"muon": muon, "backup": backup}, + ) + + +def test_hybrid_portable_document_is_complete_owned_and_weights_only_safe(): + source = torch.arange(30, dtype=torch.float32)[3:15:2].reshape(2, 3) + muon = _child("muon", tensor=source) + document = _document(muon=muon) + + assert document["format"] == HYBRID_PORTABLE_STATE_FORMAT + assert document["format_version"] == HYBRID_PORTABLE_STATE_FORMAT_VERSION + assert document["coverage"] == HYBRID_PORTABLE_STATE_COVERAGE + assert document["implementation"] == HYBRID_PORTABLE_STATE_IMPLEMENTATION + assert document["completion"]["status"] == "complete" + nested = document["children"]["muon"]["parameters"]["Model.Weight"]["state"]["momentum"] + assert nested.device.type == "cpu" + assert nested.is_contiguous() + assert nested.storage_offset() == 0 + assert nested.untyped_storage().nbytes() == nested.numel() * nested.element_size() + assert torch.equal(nested, source) + assert nested is not source + + buffer = io.BytesIO() + torch.save(document, buffer) + buffer.seek(0) + normalized = normalize_hybrid_portable_state_document(torch.load(buffer, weights_only=True)) + assert normalized["completion"] == document["completion"] + + +@pytest.mark.parametrize( + ("key", "value", "match"), + [ + ("format", "other", "format"), + ("format_version", 2, "format_version"), + ("coverage", "partial", "coverage"), + ("implementation", "gefen.Gefen", "implementation"), + ("backup_optimizer", "adamw", "Gefen backup"), + ], +) +def test_hybrid_portable_document_rejects_wrong_envelope_identity(key, value, match): + document = _document() + document[key] = value + + with pytest.raises(ValueError, match=match): + normalize_hybrid_portable_state_document(document) + + +def test_hybrid_portable_document_rejects_outer_and_nested_digest_corruption(): + document = _document() + outer = copy.deepcopy(document) + outer["routing"]["Model.Bias"] = "muon" + with pytest.raises(ValueError, match="routing"): + normalize_hybrid_portable_state_document(outer) + + nested = copy.deepcopy(document) + nested["children"]["muon"]["parameters"]["Model.Weight"]["state"]["step"] += 1 + with pytest.raises(ValueError, match="digest"): + normalize_hybrid_portable_state_document(nested) + + completion = copy.deepcopy(document) + completion["completion"]["digest"] = "0" * 64 + with pytest.raises(ValueError, match="digest"): + normalize_hybrid_portable_state_document(completion) + + +def test_hybrid_portable_document_rejects_schema_and_routing_mismatches(): + document = _document() + extra = copy.deepcopy(document) + extra["extra"] = None + with pytest.raises(ValueError, match="top-level"): + normalize_hybrid_portable_state_document(extra) + + missing_role = copy.deepcopy(document) + del missing_role["children"]["backup"] + with pytest.raises(ValueError, match="children"): + normalize_hybrid_portable_state_document(missing_role) + + missing_fqn = copy.deepcopy(document) + del missing_fqn["routing"]["Model.Bias"] + with pytest.raises(ValueError, match="routing"): + normalize_hybrid_portable_state_document(missing_fqn) + + wrong_child = copy.deepcopy(document) + wrong_child["children"]["muon"] = _child("backup", fqn="Model.Weight") + with pytest.raises(ValueError, match="implementation"): + normalize_hybrid_portable_state_document(wrong_child) + + +@pytest.mark.parametrize( + ("backup", "match"), + [ + (_child("backup", step=8), "global steps"), + (_child("backup", deterministic=False), "deterministic"), + (_child("backup", fqn="Model.Weight"), "disjoint"), + ], +) +def test_hybrid_portable_document_rejects_incompatible_children(backup, match): + with pytest.raises(ValueError, match=match): + build_hybrid_portable_state_document( + backup_optimizer="gefen", + routing={"Model.Weight": "muon", next(iter(backup["parameters"])): "backup"}, + children={"muon": _child("muon"), "backup": backup}, + ) + + +@pytest.mark.parametrize("role", ["muon", "backup"]) +def test_hybrid_portable_document_supports_one_present_child(role): + child = _child(role) + fqn = next(iter(child["parameters"])) + children = {"muon": None, "backup": None} + children[role] = child + + document = build_hybrid_portable_state_document( + backup_optimizer="gefen", + routing={fqn: role}, + children=children, + ) + + assert document["children"][role] is not None + assert document["children"]["backup" if role == "muon" else "muon"] is None + + +def test_hybrid_portable_document_rejects_no_children_and_noncanonical_routing(): + with pytest.raises(ValueError, match="at least one child"): + build_hybrid_portable_state_document( + backup_optimizer="gefen", + routing={}, + children={"muon": None, "backup": None}, + ) + with pytest.raises(ValueError, match="trimmed FQNs"): + build_hybrid_portable_state_document( + backup_optimizer="gefen", + routing={" Model.Weight": "muon"}, + children={"muon": _child("muon"), "backup": None}, + ) diff --git a/tests/test_portable_hybrid_distributed.py b/tests/test_portable_hybrid_distributed.py new file mode 100644 index 0000000..a239593 --- /dev/null +++ b/tests/test_portable_hybrid_distributed.py @@ -0,0 +1,679 @@ +"""Warning-strict distributed topology coverage for portable Hybrid state.""" + +from datetime import timedelta +import multiprocessing as mp +import os +import queue as queue_module +import tempfile +import traceback +import warnings + +import pytest +import torch +import torch.distributed as dist +from torch.distributed.checkpoint import FileSystemReader, FileSystemWriter + +from gefen import GefenMuonHybrid, load_portable_dcp, save_portable_dcp +from gefen.checkpoint import CheckpointProcessGroupBinding +from gefen.codebook import CodebookProcessGroupBinding +from gefen.contracts import ( + LogicalSlice, + ParameterIdentity, + ParameterLayout, + PlacementKind, + ProcessGroupIdentity, + ShardIdentity, + ShardPlacement, + ShardingManifest, +) +from gefen.portable_hybrid import _hybrid_portable_live_token +from gefen.portable_state import PortableStateLimits +from gefen.rebinding import ParameterRebinding + + +_WORLD = 2 +_MUON_FQN = "model.matrix" +_BACKUP_FQN = "model.vector" +_SOURCE_OWNER = "rank:0" +_TARGET_OWNER = "rank:1" +_SOURCE_LENGTHS = (3, 5) +_TARGET_LENGTHS = (5, 3) + + +def _limits(): + return PortableStateLimits( + max_fragment_tensor_bytes=1 << 20, + max_collective_tensor_bytes=4 << 20, + max_collective_metadata_bytes=4 << 20, + chunk_bytes=11, + max_members=4, + max_metadata_bytes=1 << 20, + max_tree_nodes=10_000, + max_tree_depth=32, + max_container_items=10_000, + max_string_bytes=16 << 10, + max_integer_bytes=128, + max_tensors=256, + max_tensor_rank=8, + diagnostic_bytes=1024, + ) + + +def _members(): + return tuple("rank:{}".format(rank) for rank in range(_WORLD)) + + +def _bindings(group, rank): + member = _members()[rank] + return ( + CodebookProcessGroupBinding( + group, + member, + dist.group.WORLD, + torch.device("cpu"), + ), + CheckpointProcessGroupBinding( + group, + member, + dist.group.WORLD, + torch.device("cpu"), + ), + ) + + +def _owner_shards(identity, group, owner): + return tuple( + ShardIdentity( + identity, + ParameterLayout.WHOLE_PARAMETER_OWNER, + (LogicalSlice.full(identity) if member == owner else LogicalSlice(0, 0)), + placements=( + ShardPlacement( + "checkpoint", + PlacementKind.WHOLE_PARAMETER_OWNER, + coordinate, + len(group.ordered_members), + ), + ), + process_group=group, + local_member=member, + owner=owner, + ) + for coordinate, member in enumerate(group.ordered_members) + ) + + +def _flat_shards(identity, group, lengths): + if len(lengths) != len(group.ordered_members) or sum(lengths) != identity.numel: + raise AssertionError("invalid flattened test partition") + offset = 0 + shards = [] + for coordinate, (member, length) in enumerate(zip(group.ordered_members, lengths)): + shards.append( + ShardIdentity( + identity, + ParameterLayout.FLATTENED_ELEMENT_SHARD, + LogicalSlice(offset, length), + placements=( + ShardPlacement( + "checkpoint", + PlacementKind.FLAT_SHARD, + coordinate, + len(group.ordered_members), + ), + ), + process_group=group, + local_member=member, + ) + ) + offset += length + return tuple(shards) + + +def _muon_initial(): + return torch.linspace(-0.3, 0.2, 6, dtype=torch.float32).reshape(3, 2) + + +def _backup_initial(): + return torch.linspace(-0.4, 0.3, 8, dtype=torch.float32) + + +def _muon_momentum(): + return torch.tensor( + [[-0.75, 0.5], [1.25, -1.5], [2.0, -2.5]], + dtype=torch.float32, + ) + + +def _muon_normuon_v(): + return torch.tensor([[0.5], [1.5], [2.5]], dtype=torch.float32) + + +def _backup_momentum(): + return torch.tensor( + [-0.25, 0.5, -0.75, 1.0, -1.25, 1.5, -1.75, 2.0], + dtype=torch.float32, + ) + + +def _backup_second_moment(): + return torch.linspace(0.2, 0.9, 8, dtype=torch.float32) + + +def _muon_gradient(): + return torch.tensor( + [[0.4, -0.7], [1.1, -1.3], [1.7, -1.9]], + dtype=torch.float32, + ) + + +def _backup_gradient(): + return torch.tensor( + [0.25, -0.5, 0.75, -1.0, 1.25, -1.5, 1.75, -2.0], + dtype=torch.float32, + ) + + +def _codebook(): + return torch.linspace(-1.0, 1.0, 256, dtype=torch.float32) + + +def _quantized_period_one(momentum): + flat = momentum.reshape(-1) + indices = torch.where( + torch.signbit(flat), + torch.zeros(flat.numel(), dtype=torch.uint8), + torch.full((flat.numel(),), 255, dtype=torch.uint8), + ) + return indices.reshape(-1, 1), flat.abs().reshape(-1, 1).clone() + + +def _make_hybrid(rank, group, *, owner, lengths, deterministic): + muon_identity = ParameterIdentity(_MUON_FQN, (3, 2)) + backup_identity = ParameterIdentity(_BACKUP_FQN, (8,)) + old_muon = torch.nn.Parameter(_muon_initial().clone()) + old_backup = torch.nn.Parameter(_backup_initial().clone()) + optimizer = GefenMuonHybrid( + [("matrix", old_muon)], + [("vector", old_backup)], + lr=2.5e-3, + muon_lr=3.0e-3, + backup_lr=2.5e-3, + weight_decay=0.03, + muon_weight_decay=0.02, + backup_weight_decay=0.03, + backup_optimizer="gefen", + backup_1d_period_one=True, + betas=(0.8, 0.97), + eps=2.0e-8, + fused=False, + momentum=0.85, + nesterov=False, + ns_steps=2, + sharded_mode="distributed", + deterministic=deterministic, + normuon=True, + normuon_beta2=0.9, + normuon_eps=3.0e-8, + ) + muon_shards = _owner_shards(muon_identity, group, owner) + backup_shards = _flat_shards(backup_identity, group, lengths) + muon_local = old_muon if _members()[rank] == owner else None + backup_shard = backup_shards[rank] + start = backup_shard.logical_slice.flat_offset + stop = start + backup_shard.logical_slice.length + backup_local = torch.nn.Parameter(_backup_initial()[start:stop].clone()) + codebook_binding, checkpoint_binding = _bindings(group, rank) + optimizer.post_sharding( + ( + ParameterRebinding(old_muon, muon_local, muon_shards[rank]), + ParameterRebinding(old_backup, backup_local, backup_shard), + ), + manifest=ShardingManifest(muon_shards + backup_shards), + codebook_process_group=codebook_binding, + ) + return ( + optimizer, + muon_local, + backup_local, + backup_shard, + checkpoint_binding, + ) + + +def _seed_hybrid_state(optimizer, muon_parameter, backup_parameter, backup_shard): + for child in (optimizer.muon, optimizer.backup): + child._gefen_global_step = 13 + child._gefen_codebook = _codebook() + if muon_parameter is not None: + indices, magnitudes = _quantized_period_one(_muon_momentum()) + optimizer.muon.state[muon_parameter].update( + { + "automatic_period": 1, + "step": 11, + "m_codebook": indices, + "m_magnitude": magnitudes, + "normuon_v": _muon_normuon_v(), + "normuon_step": 10, + } + ) + start = backup_shard.logical_slice.flat_offset + stop = start + backup_shard.logical_slice.length + momentum = _backup_momentum()[start:stop] + indices, magnitudes = _quantized_period_one(momentum) + optimizer.backup.state[backup_parameter].update( + { + "automatic_period": 1, + "step": 11, + "m_codebook": indices, + "m_magnitude": magnitudes, + "vmean": _backup_second_moment()[start:stop].reshape(-1, 1).clone(), + "vmean_step": 10, + } + ) + + +def _bits_equal(left, right): + return ( + torch.is_tensor(left) + and torch.is_tensor(right) + and left.dtype == right.dtype + and tuple(left.shape) == tuple(right.shape) + and torch.equal( + left.detach().contiguous().view(torch.uint8), + right.detach().contiguous().view(torch.uint8), + ) + ) + + +def _state_equal(left, right): + if set(left) != set(right): + return False + for key in left: + left_value = left[key] + right_value = right[key] + if torch.is_tensor(left_value) or torch.is_tensor(right_value): + if not _bits_equal(left_value, right_value): + return False + elif type(left_value) is not type(right_value) or left_value != right_value: + return False + return True + + +def _local_child_exact(left, left_parameter, right, right_parameter): + if left_parameter is None or right_parameter is None: + return ( + left_parameter is right_parameter is None + and not left.param_groups[0]["params"] + and not right.param_groups[0]["params"] + and not left.state + and not right.state + ) + return _state_equal( + left.state[left_parameter], + right.state[right_parameter], + ) + + +def _hybrids_exact( + left, + left_muon, + left_backup, + right, + right_muon, + right_backup, +): + return ( + left._canonical_identity_ready() + and right._canonical_identity_ready() + and left._hybrid_fqn_roles == right._hybrid_fqn_roles == ((_MUON_FQN, "muon"), (_BACKUP_FQN, "backup")) + and left._deterministic is right._deterministic is True + and left.muon._deterministic is right.muon._deterministic is True + and left.backup._deterministic is right.backup._deterministic is True + and left.muon._gefen_global_step == right.muon._gefen_global_step + and left.backup._gefen_global_step == right.backup._gefen_global_step + and _bits_equal( + left.muon._gefen_codebook, + right.muon._gefen_codebook, + ) + and _bits_equal( + left.backup._gefen_codebook, + right.backup._gefen_codebook, + ) + and _local_child_exact( + left.muon, + left_muon, + right.muon, + right_muon, + ) + and _local_child_exact( + left.backup, + left_backup, + right.backup, + right_backup, + ) + ) + + +def _step_against_reference( + target, + target_muon, + target_backup, + target_backup_shard, + reference, + reference_muon, + reference_backup, +): + if target_muon is not None: + target_muon.grad = _muon_gradient().clone() + reference_muon.grad = _muon_gradient().clone() + start = target_backup_shard.logical_slice.flat_offset + stop = start + target_backup_shard.logical_slice.length + gradient = _backup_gradient()[start:stop] + target_backup.grad = gradient.clone() + reference_backup.grad = gradient.clone() + target.step() + reference.step() + parameters_exact = ( + (target_muon is reference_muon is None) or _bits_equal(target_muon, reference_muon) + ) and _bits_equal(target_backup, reference_backup) + return parameters_exact and _hybrids_exact( + target, + target_muon, + target_backup, + reference, + reference_muon, + reference_backup, + ) + + +def _fresh_target_and_reference(rank, group): + target = _make_hybrid( + rank, + group, + owner=_TARGET_OWNER, + lengths=_TARGET_LENGTHS, + deterministic=False, + ) + reference = _make_hybrid( + rank, + group, + owner=_TARGET_OWNER, + lengths=_TARGET_LENGTHS, + deterministic=True, + ) + _seed_hybrid_state( + reference[0], + reference[1], + reference[2], + reference[3], + ) + return target, reference + + +def _successful_import_result(rank, group, document): + target, reference = _fresh_target_and_reference(rank, group) + target[0].import_portable_state( + document, + checkpoint_process_group=target[4], + transaction_id="hybrid-direct-import-v1", + limits=_limits(), + ) + restored_exact = _hybrids_exact( + target[0], + target[1], + target[2], + reference[0], + reference[1], + reference[2], + ) + next_step_exact = _step_against_reference( + target[0], + target[1], + target[2], + target[3], + reference[0], + reference[1], + reference[2], + ) + return { + "restored_exact": restored_exact, + "next_step_exact": next_step_exact, + "target_owns_muon": target[1] is not None, + "target_backup_length": target[3].logical_slice.length, + } + + +def _asymmetric_failure_result(rank, group, document): + target = _make_hybrid( + rank, + group, + owner=_TARGET_OWNER, + lengths=_TARGET_LENGTHS, + deterministic=False, + ) + optimizer, muon_parameter, backup_parameter, _shard, binding = target + if rank == 0: + optimizer.backup.param_groups[0]["rank_local_extension"] = "reject-me" + token_before = _hybrid_portable_live_token(optimizer) + child_state_objects = (optimizer.muon.state, optimizer.backup.state) + parameter_values = ( + None if muon_parameter is None else muon_parameter.detach().clone(), + backup_parameter.detach().clone(), + ) + try: + optimizer.import_portable_state( + document, + checkpoint_process_group=binding, + transaction_id="hybrid-asymmetric-import-failure-v1", + limits=_limits(), + ) + except RuntimeError as exc: + message = str(exc) + else: + message = None + parameters_unchanged = ( + (muon_parameter is None and parameter_values[0] is None) or _bits_equal(muon_parameter, parameter_values[0]) + ) and _bits_equal(backup_parameter, parameter_values[1]) + return { + "message": message, + "token_unchanged": _hybrid_portable_live_token(optimizer) == token_before, + "state_objects_unchanged": ( + optimizer.muon.state is child_state_objects[0] and optimizer.backup.state is child_state_objects[1] + ), + "parameters_unchanged": parameters_unchanged, + "common_unchanged": ( + optimizer._deterministic is False + and optimizer.muon._gefen_global_step == 0 + and optimizer.backup._gefen_global_step == 0 + and optimizer.muon._gefen_codebook is None + and optimizer.backup._gefen_codebook is None + ), + } + + +def _dcp_result(rank, group, checkpoint_dir): + target, reference = _fresh_target_and_reference(rank, group) + load_portable_dcp( + target[0], + checkpoint_process_group=target[4], + storage_reader=FileSystemReader(checkpoint_dir), + transaction_id="hybrid-dcp-load-v1", + limits=_limits(), + ) + restored_exact = _hybrids_exact( + target[0], + target[1], + target[2], + reference[0], + reference[1], + reference[2], + ) + next_step_exact = _step_against_reference( + target[0], + target[1], + target[2], + target[3], + reference[0], + reference[1], + reference[2], + ) + return { + "restored_exact": restored_exact, + "next_step_exact": next_step_exact, + } + + +def _strict_worker_warnings(): + warnings.simplefilter("error") + warnings.filterwarnings( + "ignore", + message="TypedStorage is deprecated.*", + category=UserWarning, + ) + + +def _distributed_worker(rank, init_file, checkpoint_dir, result_queue): + _strict_worker_warnings() + try: + dist.init_process_group( + "gloo", + init_method="file://{}".format(init_file), + rank=rank, + world_size=_WORLD, + timeout=timedelta(seconds=120), + ) + group = ProcessGroupIdentity("hybrid_portable_checkpoint", _members()) + source = _make_hybrid( + rank, + group, + owner=_SOURCE_OWNER, + lengths=_SOURCE_LENGTHS, + deterministic=True, + ) + _seed_hybrid_state(source[0], source[1], source[2], source[3]) + document = source[0].export_portable_state( + checkpoint_process_group=source[4], + transaction_id="hybrid-direct-export-v1", + limits=_limits(), + ) + direct = _successful_import_result(rank, group, document) + dist.barrier() + asymmetric = _asymmetric_failure_result(rank, group, document) + dist.barrier() + save_portable_dcp( + source[0], + checkpoint_process_group=source[4], + storage_writer=FileSystemWriter(checkpoint_dir), + transaction_id="hybrid-dcp-save-v1", + limits=_limits(), + ) + dist.barrier() + dcp = _dcp_result(rank, group, checkpoint_dir) + dist.barrier() + result_queue.put( + { + "rank": rank, + "digest": document["completion"]["digest"], + "document_exact": ( + document["routing"] == {_MUON_FQN: "muon", _BACKUP_FQN: "backup"} + and all( + document["children"][role]["common"]["gefen_global_step"] == 13 + and document["children"][role]["common"]["gefen_deterministic"] is True + and _bits_equal( + document["children"][role]["common"]["gefen_codebook"], + _codebook(), + ) + for role in ("muon", "backup") + ) + ), + "direct": direct, + "asymmetric": asymmetric, + "dcp": dcp, + } + ) + except BaseException: + result_queue.put({"rank": rank, "fatal_error": traceback.format_exc()}) + finally: + if dist.is_available() and dist.is_initialized(): + dist.destroy_process_group() + + +def _run_distributed_workers(checkpoint_dir): + context = mp.get_context("spawn") + result_queue = context.Queue() + descriptor, init_file = tempfile.mkstemp(prefix="gefen-portable-hybrid-") + os.close(descriptor) + os.unlink(init_file) + processes = [ + context.Process( + target=_distributed_worker, + args=(rank, init_file, checkpoint_dir, result_queue), + ) + for rank in range(_WORLD) + ] + results = [] + try: + for process in processes: + process.start() + try: + for _ in processes: + results.append(result_queue.get(timeout=240)) + except queue_module.Empty: + pass + for process in processes: + process.join(timeout=10) + for process in processes: + if process.is_alive(): + process.terminate() + process.join(timeout=5) + finally: + for process in processes: + if process.is_alive(): + process.terminate() + process.join(timeout=5) + result_queue.close() + result_queue.join_thread() + if os.path.exists(init_file): + os.unlink(init_file) + assert len(results) == _WORLD, (results, [process.exitcode for process in processes]) + assert all(process.exitcode == 0 for process in processes), [process.exitcode for process in processes] + return sorted(results, key=lambda item: item["rank"]) + + +@pytest.mark.skipif( + not dist.is_available() or not dist.is_gloo_available(), + reason="portable Hybrid topology coverage requires Gloo", +) +def test_two_process_portable_hybrid_topology_change_and_dcp(tmp_path): + results = _run_distributed_workers(str(tmp_path / "hybrid-portable-dcp")) + + assert all("fatal_error" not in result for result in results), results + assert len({result["digest"] for result in results}) == 1 + assert all(result["document_exact"] for result in results) + assert [result["direct"]["target_owns_muon"] for result in results] == [ + False, + True, + ] + assert [result["direct"]["target_backup_length"] for result in results] == [ + 5, + 3, + ] + assert all( + result["direct"]["restored_exact"] + and result["direct"]["next_step_exact"] + and result["dcp"]["restored_exact"] + and result["dcp"]["next_step_exact"] + for result in results + ), results + + messages = [result["asymmetric"]["message"] for result in results] + assert messages[0] == messages[1] + assert messages[0] is not None and "unknown keys" in messages[0] + assert all( + result["asymmetric"]["token_unchanged"] + and result["asymmetric"]["state_objects_unchanged"] + and result["asymmetric"]["parameters_unchanged"] + and result["asymmetric"]["common_unchanged"] + for result in results + ), results diff --git a/tests/test_portable_hybrid_runtime.py b/tests/test_portable_hybrid_runtime.py new file mode 100644 index 0000000..0d43303 --- /dev/null +++ b/tests/test_portable_hybrid_runtime.py @@ -0,0 +1,456 @@ +"""Warning-strict runtime coverage for Gefen-backed Hybrid portable state.""" + +import copy +import warnings + +import pytest +import torch +import torch.distributed.checkpoint as dcp + +from gefen import ( + CheckpointProcessGroupBinding, + CheckpointTransport, + GefenMuonHybrid, + PortableStateLimits, + PortableStateProvider, + load_portable_dcp, + save_portable_dcp, +) +from gefen.codebook import CodebookProcessGroupBinding +from gefen.contracts import ( + LogicalSlice, + ParameterIdentity, + ParameterLayout, + PlacementKind, + ProcessGroupIdentity, + ShardIdentity, + ShardPlacement, + ShardingManifest, +) +from gefen.rebinding import ParameterRebinding +from gefen.portable_schema import portable_state_digest + + +_MEMBER = "rank:0" +_SINGLE_PROCESS_DCP_WARNING = ( + r"^torch\.distributed is (?:disabled, )?unavailable or uninitialized, " + r"assuming the intent is to (?:save|load) in a single process\.$" +) +_TYPED_STORAGE_DCP_WARNING = ( + r"^TypedStorage is deprecated\. It will be removed in the future and UntypedStorage will be the only storage class\. " + r"This should only matter to you if you are using storages directly\. To access UntypedStorage directly, use " + r"tensor\.untyped_storage\(\) instead of tensor\.storage\(\)$" +) + + +@pytest.fixture(autouse=True) +def _warning_strict(): + with warnings.catch_warnings(): + warnings.simplefilter("error") + warnings.filterwarnings( + "ignore", + message=_SINGLE_PROCESS_DCP_WARNING, + category=UserWarning, + ) + warnings.filterwarnings( + "ignore", + message=_TYPED_STORAGE_DCP_WARNING, + category=UserWarning, + ) + yield + + +def _limits(): + return PortableStateLimits( + max_fragment_tensor_bytes=4 << 20, + max_collective_tensor_bytes=16 << 20, + max_collective_metadata_bytes=64 << 20, + ) + + +def _grouped_replicated(fqn, shape, group): + identity = ParameterIdentity(fqn, shape) + return ShardIdentity( + identity, + ParameterLayout.REPLICATED, + LogicalSlice.full(identity), + placements=( + ShardPlacement( + "checkpoint", + PlacementKind.REPLICATE, + 0, + 1, + ), + ), + process_group=group, + local_member=_MEMBER, + ) + + +def _make_hybrid(matrix, bias, *, deterministic): + optimizer = GefenMuonHybrid( + [("layer.weight", matrix)], + [("layer.bias", bias)], + lr=1e-3, + fused=False, + sharded_mode="distributed", + backup_optimizer="gefen", + backup_1d_period_one=True, + normuon=True, + deterministic=deterministic, + ) + group = ProcessGroupIdentity("hybrid_checkpoint", (_MEMBER,)) + matrix_shard = _grouped_replicated( + "Model.Layer.Weight", + tuple(matrix.shape), + group, + ) + bias_shard = _grouped_replicated( + "Model.Layer.Bias", + tuple(bias.shape), + group, + ) + manifest = ShardingManifest((matrix_shard, bias_shard)) + codebook_binding = CodebookProcessGroupBinding( + group, + _MEMBER, + None, + torch.device("cpu"), + ) + checkpoint_binding = CheckpointProcessGroupBinding( + group, + _MEMBER, + None, + torch.device("cpu"), + ) + optimizer.post_sharding( + ( + ParameterRebinding(matrix, matrix, matrix_shard), + ParameterRebinding(bias, bias, bias_shard), + ), + manifest=manifest, + codebook_process_group=codebook_binding, + ) + return optimizer, checkpoint_binding + + +def _set_grads(matrix, bias, offset): + matrix.grad = ( + torch.tensor( + [[0.25, -0.5], [0.75, -1.0]], + dtype=matrix.dtype, + ) + + offset + ) + bias.grad = torch.tensor([0.2, -0.4, 0.6, -0.8], dtype=bias.dtype) - offset + + +def _equal(left, right): + if torch.is_tensor(left) or torch.is_tensor(right): + return ( + torch.is_tensor(left) + and torch.is_tensor(right) + and left.dtype == right.dtype + and tuple(left.shape) == tuple(right.shape) + and torch.equal(left.cpu(), right.cpu()) + ) + if type(left) is not type(right): + return False + if type(left) is dict: + return set(left) == set(right) and all(_equal(left[key], right[key]) for key in left) + if type(left) in {list, tuple}: + return len(left) == len(right) and all(_equal(a, b) for a, b in zip(left, right)) + return left == right + + +def _redigest(document): + payload = {key: document[key] for key in document if key != "completion"} + document["completion"]["digest"] = portable_state_digest(payload) + + +def _assert_same_optimizer_state(left, right): + assert _equal(left.muon.state_dict(), right.muon.state_dict()) + assert _equal(left.backup.state_dict(), right.backup.state_dict()) + assert left._deterministic is right._deterministic + + +def _initialized_source(): + matrix = torch.nn.Parameter(torch.tensor([[1.0, -2.0], [3.0, -4.0]])) + bias = torch.nn.Parameter(torch.tensor([0.5, -1.5, 2.5, -3.5])) + optimizer, binding = _make_hybrid(matrix, bias, deterministic=True) + _set_grads(matrix, bias, 0.0) + optimizer.step() + return optimizer, matrix, bias, binding + + +def _target_from_source(matrix, bias, *, deterministic=False): + target_matrix = torch.nn.Parameter(matrix.detach().clone()) + target_bias = torch.nn.Parameter(bias.detach().clone()) + optimizer, binding = _make_hybrid( + target_matrix, + target_bias, + deterministic=deterministic, + ) + return optimizer, target_matrix, target_bias, binding + + +def test_hybrid_portable_round_trip_and_exact_next_step(): + source, source_matrix, source_bias, source_binding = _initialized_source() + document = source.export_portable_state( + checkpoint_process_group=source_binding, + transaction_id="hybrid-export-v1", + limits=_limits(), + ) + target, target_matrix, target_bias, target_binding = _target_from_source( + source_matrix, + source_bias, + ) + + target.import_portable_state( + document, + checkpoint_process_group=target_binding, + transaction_id="hybrid-import-v1", + limits=_limits(), + ) + + assert isinstance(target, PortableStateProvider) + assert document["routing"] == { + "Model.Layer.Bias": "backup", + "Model.Layer.Weight": "muon", + } + assert document["children"]["muon"]["implementation"] == "gefen.GefenMuon" + assert document["children"]["backup"]["implementation"] == "gefen.Gefen" + assert target._deterministic is True + _assert_same_optimizer_state(source, target) + support = tuple( + item + for item in target.optimizer_contract().capabilities.checkpoints + if item.transport is CheckpointTransport.CANONICAL_GLOBAL + ) + assert len(support) == 1 + assert support[0].atomic_load + + _set_grads(source_matrix, source_bias, 0.125) + _set_grads(target_matrix, target_bias, 0.125) + source.step() + target.step() + + assert torch.equal(source_matrix, target_matrix) + assert torch.equal(source_bias, target_bias) + _assert_same_optimizer_state(source, target) + + +def test_hybrid_portable_import_stages_every_child_before_mutation(): + source, source_matrix, source_bias, source_binding = _initialized_source() + document = source.export_portable_state( + checkpoint_process_group=source_binding, + transaction_id="hybrid-atomic-export-v1", + limits=_limits(), + ) + target, _target_matrix, _target_bias, target_binding = _target_from_source( + source_matrix, + source_bias, + ) + target.backup.param_groups[0]["eps"] = 1e-7 + muon_before = copy.deepcopy(target.muon.state_dict()) + backup_before = copy.deepcopy(target.backup.state_dict()) + muon_state_object = target.muon.state + backup_state_object = target.backup.state + + with pytest.raises((RuntimeError, ValueError), match="policy|option|portable"): + target.import_portable_state( + document, + checkpoint_process_group=target_binding, + transaction_id="hybrid-atomic-import-v1", + limits=_limits(), + ) + + assert target.muon.state is muon_state_object + assert target.backup.state is backup_state_object + assert _equal(target.muon.state_dict(), muon_before) + assert _equal(target.backup.state_dict(), backup_before) + + +def test_hybrid_portable_import_rejects_semantically_noncanonical_child(): + source, source_matrix, source_bias, source_binding = _initialized_source() + document = source.export_portable_state( + checkpoint_process_group=source_binding, + transaction_id="hybrid-semantic-export-v1", + limits=_limits(), + ) + document["children"]["backup"]["provenance"] = {} + _redigest(document["children"]["backup"]) + _redigest(document) + target, _target_matrix, _target_bias, target_binding = _target_from_source( + source_matrix, + source_bias, + ) + muon_before = copy.deepcopy(target.muon.state_dict()) + backup_before = copy.deepcopy(target.backup.state_dict()) + + with pytest.raises(RuntimeError, match="provenance"): + target.import_portable_state( + document, + checkpoint_process_group=target_binding, + transaction_id="hybrid-semantic-import-v1", + limits=_limits(), + ) + + assert _equal(target.muon.state_dict(), muon_before) + assert _equal(target.backup.state_dict(), backup_before) + + +def test_hybrid_portable_composite_uses_collective_tensor_limit(): + source, source_matrix, source_bias, source_binding = _initialized_source() + limits = PortableStateLimits( + max_fragment_tensor_bytes=1536, + max_collective_tensor_bytes=4096, + max_collective_metadata_bytes=64 << 20, + ) + + document = source.export_portable_state( + checkpoint_process_group=source_binding, + transaction_id="hybrid-aggregate-export-v1", + limits=limits, + ) + target, _target_matrix, _target_bias, target_binding = _target_from_source( + source_matrix, + source_bias, + ) + target.import_portable_state( + document, + checkpoint_process_group=target_binding, + transaction_id="hybrid-aggregate-import-v1", + limits=limits, + ) + + _assert_same_optimizer_state(source, target) + + +def test_hybrid_portable_dcp_round_trip(tmp_path): + source, source_matrix, source_bias, source_binding = _initialized_source() + checkpoint = tmp_path / "hybrid-portable" + save_portable_dcp( + source, + checkpoint_process_group=source_binding, + storage_writer=dcp.FileSystemWriter(checkpoint), + transaction_id="hybrid-dcp-save-v1", + limits=_limits(), + ) + target, target_matrix, target_bias, target_binding = _target_from_source( + source_matrix, + source_bias, + ) + + load_portable_dcp( + target, + checkpoint_process_group=target_binding, + storage_reader=dcp.FileSystemReader(checkpoint), + transaction_id="hybrid-dcp-load-v1", + limits=_limits(), + ) + + _assert_same_optimizer_state(source, target) + _set_grads(source_matrix, source_bias, -0.0625) + _set_grads(target_matrix, target_bias, -0.0625) + source.step() + target.step() + assert torch.equal(source_matrix, target_matrix) + assert torch.equal(source_bias, target_bias) + _assert_same_optimizer_state(source, target) + + +def test_adamw_backed_hybrid_remains_explicitly_nonportable(): + matrix = torch.nn.Parameter(torch.ones(2, 2)) + bias = torch.nn.Parameter(torch.ones(4)) + optimizer = GefenMuonHybrid( + [("layer.weight", matrix)], + [("layer.bias", bias)], + lr=1e-3, + fused=False, + sharded_mode="distributed", + backup_optimizer="adamw", + ) + assert not optimizer.optimizer_contract().capabilities.canonical_state_io + assert all( + item.transport is not CheckpointTransport.CANONICAL_GLOBAL + for item in optimizer.optimizer_contract().capabilities.checkpoints + ) + group = ProcessGroupIdentity("hybrid_checkpoint", (_MEMBER,)) + binding = CheckpointProcessGroupBinding( + group, + _MEMBER, + None, + torch.device("cpu"), + ) + with pytest.raises(RuntimeError, match="AdamW"): + optimizer.export_portable_state( + checkpoint_process_group=binding, + transaction_id="adamw-hybrid-reject-v1", + limits=_limits(), + ) + + +def test_hybrid_canonical_global_checkpoint_guarantees_are_intersected(monkeypatch): + # A composite canonical-global save/load processes EVERY child, so the + # Hybrid may advertise a checkpoint guarantee only when both children + # support it. The children's guarantee sets must intersect, not union. + from types import SimpleNamespace + + from gefen.contracts import ( + CheckpointSupport, + ProcessGroupScope, + TopologyChange, + ) + from gefen.portable_hybrid import _hybrid_portable_contract_support + + source, _matrix, _bias, _binding = _initialized_source() + + def _support(same_topology, kinds): + return CheckpointSupport( + transport=CheckpointTransport.CANONICAL_GLOBAL, + same_topology=same_topology, + topology_changing=frozenset({ParameterLayout.REPLICATED}), + process_group_scope=ProcessGroupScope.DEFAULT_WORLD, + topology_change_kinds=kinds, + atomic_load=True, + ) + + muon_support = _support( + frozenset( + {ParameterLayout.REPLICATED, ParameterLayout.FLATTENED_ELEMENT_SHARD} + ), + frozenset( + { + TopologyChange.WORLD_SIZE_OWNER_REDISTRIBUTION, + TopologyChange.PLACEMENT_RESHARD, + } + ), + ) + backup_support = _support( + frozenset({ParameterLayout.REPLICATED}), + frozenset({TopologyChange.WORLD_SIZE_OWNER_REDISTRIBUTION}), + ) + + def _stub(support): + # Patch at the class level: portable readiness rejects instance-level + # method shadows. The muon and backup children are distinct classes. + return lambda _self: SimpleNamespace( + capabilities=SimpleNamespace(checkpoints=(support,)) + ) + + assert type(source.muon) is not type(source.backup) + monkeypatch.setattr(type(source.muon), "optimizer_contract", _stub(muon_support)) + monkeypatch.setattr( + type(source.backup), "optimizer_contract", _stub(backup_support) + ) + + same_topology, topology_changing, topology_change_kinds = ( + _hybrid_portable_contract_support(source) + ) + + assert same_topology == frozenset({ParameterLayout.REPLICATED}) + assert topology_changing == frozenset({ParameterLayout.REPLICATED}) + assert topology_change_kinds == frozenset( + {TopologyChange.WORLD_SIZE_OWNER_REDISTRIBUTION} + ) diff --git a/tests/test_portable_identity.py b/tests/test_portable_identity.py new file mode 100644 index 0000000..0a6e98d --- /dev/null +++ b/tests/test_portable_identity.py @@ -0,0 +1,517 @@ +"""Strict primitive wire coverage for portable optimizer identities.""" + +import copy +import io + +import pytest +import torch + +from gefen.contracts import ( + LogicalRegion, + LogicalSlice, + ParameterIdentity, + ParameterLayout, + PlacementKind, + ProcessGroupIdentity, + ShardIdentity, + ShardPlacement, + ShardingManifest, +) +from gefen.portable_identity import ( + _normalize_parameter_identity, + _normalize_process_group_identity, + _normalize_shard_identity, + _normalize_sharding_manifest, + _serialize_parameter_identity, + _serialize_process_group_identity, + _serialize_shard_identity, + _serialize_sharding_manifest, +) + + +def _group(): + return ProcessGroupIdentity("pipeline:1/data_parallel", ("worker:c", "worker:a", "worker:b")) + + +def _placement(group, member, kind, *, dimension=None): + return ShardPlacement( + "data_parallel", + kind, + group.ordered_members.index(member), + len(group.ordered_members), + dimension, + ) + + +def _replicated_shards(parameter, group): + return tuple( + ShardIdentity( + parameter, + ParameterLayout.REPLICATED, + LogicalSlice.full(parameter), + placements=(_placement(group, member, PlacementKind.REPLICATE),), + process_group=group, + local_member=member, + ) + for member in group.ordered_members + ) + + +def _flat_shards(parameter, group): + boundaries = ((0, 2), (2, 2), (4, 1)) + return tuple( + ShardIdentity( + parameter, + ParameterLayout.FLATTENED_ELEMENT_SHARD, + LogicalSlice(offset, length), + placements=(_placement(group, member, PlacementKind.FLAT_SHARD),), + process_group=group, + local_member=member, + ) + for member, (offset, length) in zip(group.ordered_members, boundaries) + ) + + +def _owner_shards(parameter, group, owner): + return tuple( + ShardIdentity( + parameter, + ParameterLayout.WHOLE_PARAMETER_OWNER, + (LogicalSlice.full(parameter) if member == owner else LogicalSlice(0, 0)), + placements=(_placement(group, member, PlacementKind.WHOLE_PARAMETER_OWNER),), + process_group=group, + local_member=member, + owner=owner, + ) + for member in group.ordered_members + ) + + +def _dtensor_shards(parameter, group, lengths, *, dimension): + offset = 0 + shards = [] + for member, length in zip(group.ordered_members, lengths): + offsets = [0] * len(parameter.global_shape) + region_lengths = list(parameter.global_shape) + offsets[dimension] = offset + region_lengths[dimension] = length + shards.append( + ShardIdentity( + parameter, + ParameterLayout.DTENSOR_1D_DEFAULT_WORLD, + LogicalRegion(offsets, region_lengths), + placements=( + _placement( + group, + member, + PlacementKind.DIMENSION_SHARD, + dimension=dimension, + ), + ), + process_group=group, + local_member=member, + ) + ) + offset += length + return tuple(shards) + + +def _dtensor_replicated_shards(parameter, group): + return tuple( + ShardIdentity( + parameter, + ParameterLayout.DTENSOR_1D_DEFAULT_WORLD, + LogicalRegion.full(parameter), + placements=(_placement(group, member, PlacementKind.REPLICATE),), + process_group=group, + local_member=member, + ) + for member in group.ordered_members + ) + + +def _complete_manifest(): + group = _group() + shards = ( + ShardIdentity( + ParameterIdentity("local.weight", (2, 3)), + ParameterLayout.REPLICATED, + LogicalSlice(0, 6), + ), + ) + shards += _replicated_shards(ParameterIdentity("replicated.weight", (2, 3)), group) + shards += _flat_shards(ParameterIdentity("flat.weight", (5,)), group) + shards += _owner_shards(ParameterIdentity("owner.weight", (2, 3)), group, "worker:a") + shards += _dtensor_shards( + ParameterIdentity("row.weight", (5, 4)), + group, + (2, 2, 1), + dimension=0, + ) + shards += _dtensor_shards( + ParameterIdentity("column_with_empty.weight", (2, 2)), + group, + (1, 1, 0), + dimension=1, + ) + shards += _dtensor_replicated_shards(ParameterIdentity("dtensor_replicated.weight", (2, 3)), group) + return ShardingManifest(tuple(reversed(shards))) + + +def _assert_primitive_tree(value): + assert type(value) in {dict, list, str, int, type(None)} + if type(value) is dict: + assert all(type(key) is str for key in value) + for item in value.values(): + _assert_primitive_tree(item) + elif type(value) is list: + for item in value: + _assert_primitive_tree(item) + + +def test_parameter_and_process_group_codecs_have_exact_canonical_wire_shapes(): + parameter = ParameterIdentity("Encoder.Block.Weight", (0, 4, 7)) + group = _group() + + parameter_wire = _serialize_parameter_identity(parameter) + group_wire = _serialize_process_group_identity(group) + + assert parameter_wire == { + "schema_version": 1, + "fqn": "Encoder.Block.Weight", + "global_shape": [0, 4, 7], + } + assert group_wire == { + "schema_version": 1, + "semantic_name": "pipeline:1/data_parallel", + "ordered_members": ["worker:c", "worker:a", "worker:b"], + } + assert _normalize_parameter_identity(parameter_wire) == parameter_wire + assert _normalize_process_group_identity(group_wire) == group_wire + _assert_primitive_tree(parameter_wire) + _assert_primitive_tree(group_wire) + + +@pytest.mark.parametrize( + "identity,extent_kind", + ( + ( + ShardIdentity( + ParameterIdentity("local.weight", (2, 3)), + ParameterLayout.REPLICATED, + LogicalSlice(0, 6), + ), + "logical_slice", + ), + (_flat_shards(ParameterIdentity("flat.weight", (5,)), _group())[1], "logical_slice"), + ( + _owner_shards( + ParameterIdentity("owner.weight", (2, 3)), + _group(), + "worker:a", + )[0], + "logical_slice", + ), + ( + _dtensor_shards( + ParameterIdentity("row.weight", (5, 4)), + _group(), + (2, 2, 1), + dimension=0, + )[2], + "logical_region", + ), + ( + _dtensor_shards( + ParameterIdentity("empty.weight", (2, 2)), + _group(), + (1, 1, 0), + dimension=1, + )[2], + "logical_region", + ), + ( + _dtensor_replicated_shards( + ParameterIdentity("dtensor_replicated.weight", (2, 3)), + _group(), + )[1], + "logical_region", + ), + ), +) +def test_shard_codec_losslessly_round_trips_every_supported_identity(identity, extent_kind): + wire = _serialize_shard_identity(identity) + + assert wire["logical_extent"]["kind"] == extent_kind + assert _normalize_shard_identity(wire) == wire + _assert_primitive_tree(wire) + + +def test_shard_wire_explicitly_discriminates_slice_and_region_fields(): + legacy = _serialize_shard_identity(_flat_shards(ParameterIdentity("flat.weight", (5,)), _group())[0]) + dtensor = _serialize_shard_identity( + _dtensor_shards( + ParameterIdentity("row.weight", (5, 4)), + _group(), + (2, 2, 1), + dimension=0, + )[0] + ) + + assert legacy["logical_extent"] == { + "kind": "logical_slice", + "flat_offset": 0, + "length": 2, + } + assert dtensor["logical_extent"] == { + "kind": "logical_region", + "offsets": [0, 0], + "lengths": [2, 4], + } + + +def test_manifest_codec_is_deterministic_complete_and_weights_only_safe(): + manifest = _complete_manifest() + wire = _serialize_sharding_manifest(manifest) + reversed_wire = copy.deepcopy(wire) + reversed_wire["shards"].reverse() + + assert [record["parameter"]["fqn"] for record in wire["shards"]] == [ + shard.parameter.fqn for shard in manifest.shards + ] + assert _normalize_sharding_manifest(wire) == wire + assert _normalize_sharding_manifest(reversed_wire) == wire + assert any( + record["logical_extent"] + == { + "kind": "logical_region", + "offsets": [0, 2], + "lengths": [2, 0], + } + for record in wire["shards"] + ) + _assert_primitive_tree(wire) + + buffer = io.BytesIO() + torch.save(wire, buffer) + buffer.seek(0) + loaded = torch.load(buffer, weights_only=True) + assert _normalize_sharding_manifest(loaded) == wire + + +def test_normalizers_rebuild_fresh_canonical_container_trees(): + source = _serialize_sharding_manifest(_complete_manifest()) + normalized = _normalize_sharding_manifest(source) + + assert normalized == source + assert normalized is not source + assert normalized["shards"] is not source["shards"] + assert normalized["shards"][0] is not source["shards"][0] + + source["shards"][0]["parameter"]["global_shape"].append(99) + assert normalized != source + + +@pytest.mark.parametrize( + "corruption", + ( + "not_dict", + "missing_key", + "extra_key", + "bool_version", + "tuple_shape", + "bool_dimension", + "invalid_fqn", + ), +) +def test_parameter_identity_schema_corruption_is_rejected(corruption): + wire = _serialize_parameter_identity(ParameterIdentity("layer.weight", (2, 3))) + if corruption == "not_dict": + wire = [] + elif corruption == "missing_key": + wire.pop("fqn") + elif corruption == "extra_key": + wire["global_rank"] = 0 + elif corruption == "bool_version": + wire["schema_version"] = True + elif corruption == "tuple_shape": + wire["global_shape"] = (2, 3) + elif corruption == "bool_dimension": + wire["global_shape"][0] = True + else: + wire["fqn"] = ".layer" + + with pytest.raises((TypeError, ValueError)): + _normalize_parameter_identity(wire) + + +@pytest.mark.parametrize( + "corruption", + ( + "missing_key", + "extra_handle", + "bool_version", + "tuple_members", + "non_string_member", + "duplicate_member", + "invalid_name", + ), +) +def test_process_group_identity_schema_corruption_is_rejected(corruption): + wire = _serialize_process_group_identity(_group()) + if corruption == "missing_key": + wire.pop("semantic_name") + elif corruption == "extra_handle": + wire["process_group_handle"] = 123 + elif corruption == "bool_version": + wire["schema_version"] = True + elif corruption == "tuple_members": + wire["ordered_members"] = tuple(wire["ordered_members"]) + elif corruption == "non_string_member": + wire["ordered_members"][0] = 0 + elif corruption == "duplicate_member": + wire["ordered_members"][1] = wire["ordered_members"][0] + else: + wire["semantic_name"] = " dp" + + with pytest.raises((TypeError, ValueError)): + _normalize_process_group_identity(wire) + + +@pytest.mark.parametrize( + "corruption", + ( + "missing_key", + "runtime_rank", + "bool_version", + "enum_layout", + "tuple_placements", + "bad_process_group_type", + "non_string_member", + "non_string_owner", + "unknown_extent_kind", + "mixed_slice_region_fields", + "tuple_region_offsets", + "bool_region_length", + "slice_for_dtensor", + "region_for_flat", + "placement_missing_key", + "enum_placement_kind", + "bool_coordinate", + "bool_parts", + "bool_parameter_dimension", + "coordinate_mismatch", + ), +) +def test_shard_identity_schema_and_semantic_corruption_is_rejected(corruption): + identity = _dtensor_shards( + ParameterIdentity("row.weight", (5, 4)), + _group(), + (2, 2, 1), + dimension=0, + )[0] + wire = _serialize_shard_identity(identity) + if corruption == "missing_key": + wire.pop("owner") + elif corruption == "runtime_rank": + wire["global_rank"] = 7 + elif corruption == "bool_version": + wire["schema_version"] = True + elif corruption == "enum_layout": + wire["layout"] = ParameterLayout.DTENSOR_1D_DEFAULT_WORLD + elif corruption == "tuple_placements": + wire["placements"] = tuple(wire["placements"]) + elif corruption == "bad_process_group_type": + wire["process_group"] = [] + elif corruption == "non_string_member": + wire["local_member"] = 0 + elif corruption == "non_string_owner": + wire["owner"] = 0 + elif corruption == "unknown_extent_kind": + wire["logical_extent"]["kind"] = "flat" + elif corruption == "mixed_slice_region_fields": + wire["logical_extent"]["flat_offset"] = 0 + elif corruption == "tuple_region_offsets": + wire["logical_extent"]["offsets"] = (0, 0) + elif corruption == "bool_region_length": + wire["logical_extent"]["lengths"][0] = True + elif corruption == "slice_for_dtensor": + wire["logical_extent"] = { + "kind": "logical_slice", + "flat_offset": 0, + "length": 8, + } + elif corruption == "region_for_flat": + wire["layout"] = "flattened_element_shard" + elif corruption == "placement_missing_key": + wire["placements"][0].pop("parts") + elif corruption == "enum_placement_kind": + wire["placements"][0]["kind"] = PlacementKind.DIMENSION_SHARD + elif corruption == "bool_coordinate": + wire["placements"][0]["coordinate"] = True + elif corruption == "bool_parts": + wire["placements"][0]["parts"] = True + elif corruption == "bool_parameter_dimension": + wire["placements"][0]["parameter_dimension"] = True + else: + wire["placements"][0]["coordinate"] = 1 + + with pytest.raises((TypeError, ValueError)): + _normalize_shard_identity(wire) + + +@pytest.mark.parametrize( + "corruption", + ( + "not_dict", + "missing_key", + "extra_key", + "bool_version", + "tuple_shards", + "empty_shards", + "duplicate_shard", + "incomplete_group", + ), +) +def test_sharding_manifest_schema_and_completeness_corruption_is_rejected( + corruption, +): + wire = _serialize_sharding_manifest(_complete_manifest()) + if corruption == "not_dict": + wire = [] + elif corruption == "missing_key": + wire.pop("schema_version") + elif corruption == "extra_key": + wire["default_process_group"] = "world" + elif corruption == "bool_version": + wire["schema_version"] = True + elif corruption == "tuple_shards": + wire["shards"] = tuple(wire["shards"]) + elif corruption == "empty_shards": + wire["shards"] = [] + elif corruption == "duplicate_shard": + wire["shards"].append(copy.deepcopy(wire["shards"][0])) + else: + fqn = "flat.weight" + wire["shards"] = [ + shard + for shard in wire["shards"] + if not (shard["parameter"]["fqn"] == fqn and shard["local_member"] == "worker:b") + ] + + with pytest.raises((TypeError, ValueError)): + _normalize_sharding_manifest(wire) + + +@pytest.mark.parametrize( + "serializer,value", + ( + (_serialize_parameter_identity, object()), + (_serialize_process_group_identity, object()), + (_serialize_shard_identity, object()), + (_serialize_sharding_manifest, object()), + ), +) +def test_serializers_require_contract_dataclasses(serializer, value): + with pytest.raises(TypeError): + serializer(value) diff --git a/tests/test_portable_runtime.py b/tests/test_portable_runtime.py new file mode 100644 index 0000000..df657bf --- /dev/null +++ b/tests/test_portable_runtime.py @@ -0,0 +1,400 @@ +"""Warning-strict CPU coverage for the optimizer-facing portable runtime.""" + +import copy + +import pytest +import torch + +from gefen.checkpoint import CheckpointProcessGroupBinding +from gefen.codebook import CodebookProcessGroupBinding +from gefen.contracts import ( + CheckpointTransport, + LogicalSlice, + ParameterIdentity, + ParameterLayout, + PlacementKind, + PortableStateProvider, + ProcessGroupScope, + ProcessGroupIdentity, + ShardIdentity, + ShardPlacement, + ShardingManifest, + TopologyChange, +) +from gefen.gefen import Gefen +from gefen.gefen_muon import GefenMuon +from gefen.portable import _decode_quantized_momentum +from gefen.portable_runtime import ( + _export_portable_state, + _import_portable_state, + _optimizer_implementation, + _portable_runtime_layouts, +) +from gefen.portable_state import PortableStateLimits +from gefen.rebinding import ParameterRebinding + + +_MEMBER = "rank:0" + + +def _limits(): + return PortableStateLimits( + max_fragment_tensor_bytes=4 << 20, + max_collective_tensor_bytes=16 << 20, + max_collective_metadata_bytes=64 << 20, + ) + + +def _bindings(group): + return ( + CodebookProcessGroupBinding(group, _MEMBER, None, torch.device("cpu")), + CheckpointProcessGroupBinding(group, _MEMBER, None, torch.device("cpu")), + ) + + +def _finalize(optimizer, parameter, *, layout): + group = ProcessGroupIdentity("checkpoint", (_MEMBER,)) + identity = ParameterIdentity("layer.weight", tuple(parameter.shape)) + if layout is ParameterLayout.REPLICATED: + kind = PlacementKind.REPLICATE + owner = None + else: + kind = PlacementKind.WHOLE_PARAMETER_OWNER + owner = _MEMBER + shard = ShardIdentity( + identity, + layout, + LogicalSlice.full(identity), + placements=(ShardPlacement("checkpoint", kind, 0, 1),), + process_group=group, + local_member=_MEMBER, + owner=owner, + ) + codebook_binding, checkpoint_binding = _bindings(group) + optimizer.post_sharding( + (ParameterRebinding(parameter, parameter, shard),), + manifest=ShardingManifest((shard,)), + codebook_process_group=codebook_binding, + ) + return checkpoint_binding + + +def _plain(*, factored, deterministic): + parameter = torch.nn.Parameter(torch.arange(1, 7, dtype=torch.float32).reshape(2, 3)) + optimizer = Gefen( + [("weight", parameter)], + fused=False, + factored_v_2d=factored, + force_2d_period_one=True, + deterministic=deterministic, + ) + binding = _finalize(optimizer, parameter, layout=ParameterLayout.REPLICATED) + return optimizer, parameter, binding + + +def _muon(*, deterministic): + parameter = torch.nn.Parameter(torch.arange(1, 7, dtype=torch.float32).reshape(2, 3)) + optimizer = GefenMuon( + [("weight", parameter)], + fused=False, + sharded_mode="distributed", + normuon=True, + deterministic=deterministic, + ) + binding = _finalize( + optimizer, + parameter, + layout=ParameterLayout.WHOLE_PARAMETER_OWNER, + ) + return optimizer, parameter, binding + + +def _initialize(optimizer, parameter, *, variant): + optimizer._gefen_global_step = 4 + optimizer._gefen_codebook = torch.linspace(-1.0, 1.0, 256, dtype=torch.float32) + state = optimizer.state[parameter] + state.update( + { + "automatic_period": 1, + "step": 4, + "m_codebook": torch.tensor([[0], [255], [128], [64], [192], [0]], dtype=torch.uint8), + "m_magnitude": torch.tensor([[1.0], [2.0], [3.0], [4.0], [5.0], [6.0]], dtype=torch.float32), + } + ) + if variant == "block": + state.update( + { + "vmean": torch.arange(1, 7, dtype=torch.float32).reshape(6, 1), + "vmean_step": 3, + } + ) + elif variant == "factored": + state.update( + { + "v_row": torch.tensor([2.0, 3.0], dtype=torch.float32), + "v_col": torch.tensor([5.0, 7.0, 11.0], dtype=torch.float32), + "factored_step": 3, + } + ) + elif variant == "normuon": + state.update( + { + "normuon_v": torch.tensor([[2.0], [3.0]], dtype=torch.float32), + "normuon_step": 2, + } + ) + else: + raise AssertionError("unknown test variant") + + +@pytest.mark.parametrize("variant", ["block", "factored", "normuon"]) +def test_initialized_singleton_export_import_round_trip(variant): + if variant == "normuon": + source, source_parameter, source_binding = _muon(deterministic=True) + target, target_parameter, target_binding = _muon(deterministic=False) + else: + source, source_parameter, source_binding = _plain( + factored=variant == "factored", + deterministic=True, + ) + target, target_parameter, target_binding = _plain( + factored=variant == "factored", + deterministic=False, + ) + _initialize(source, source_parameter, variant=variant) + assert isinstance(source, PortableStateProvider) + support = next( + item + for item in source.optimizer_contract().capabilities.checkpoints + if item.transport is CheckpointTransport.CANONICAL_GLOBAL + ) + assert support.requires_collective + assert support.atomic_load + assert support.process_group_scope is ProcessGroupScope.ADAPTER_DEFINED + expected_layout = ParameterLayout.WHOLE_PARAMETER_OWNER if variant == "normuon" else ParameterLayout.REPLICATED + assert expected_layout in support.same_topology + if variant == "factored": + assert not support.topology_changing + elif variant == "block": + assert support.topology_changing == frozenset( + { + ParameterLayout.REPLICATED, + ParameterLayout.FLATTENED_ELEMENT_SHARD, + } + ) + assert support.topology_change_kinds == frozenset({TopologyChange.PLACEMENT_RESHARD}) + else: + assert support.topology_changing == frozenset( + { + ParameterLayout.REPLICATED, + ParameterLayout.WHOLE_PARAMETER_OWNER, + } + ) + assert support.topology_change_kinds == frozenset( + { + TopologyChange.PLACEMENT_RESHARD, + TopologyChange.WORLD_SIZE_OWNER_REDISTRIBUTION, + } + ) + + document = _export_portable_state( + source, + checkpoint_process_group=source_binding, + transaction_id="export-{}".format(variant), + limits=_limits(), + ) + _import_portable_state( + target, + document, + checkpoint_process_group=target_binding, + transaction_id="import-{}".format(variant), + limits=_limits(), + ) + + assert target._gefen_global_step == 4 + assert target._deterministic is True + assert torch.equal(target._gefen_codebook, document["common"]["gefen_codebook"]) + target_state = target.state[target_parameter] + target_momentum = _decode_quantized_momentum( + target._gefen_codebook, + target_state["m_codebook"], + target_state["m_magnitude"], + logical_shape=tuple(target_parameter.shape), + period=1, + step=target_state["step"], + ) + record = document["parameters"]["layer.weight"] + assert torch.equal(target_momentum.view(torch.int32), record["state"]["momentum"].view(torch.int32)) + if variant == "block": + assert torch.equal(target_state["vmean"].reshape(2, 3), record["state"]["second_moment"]) + assert target_state["vmean_step"] == 3 + elif variant == "factored": + assert torch.equal(target_state["v_row"], record["state"]["v_row"]) + assert torch.equal(target_state["v_col"], record["state"]["v_col"]) + assert target_state["factored_step"] == 3 + else: + assert torch.equal(target_state["normuon_v"], record["state"]["normuon_v"]) + assert target_state["normuon_step"] == 2 + + +def test_import_rejects_corrupt_document_without_live_mutation(): + source, source_parameter, source_binding = _plain(factored=False, deterministic=True) + target, _, target_binding = _plain(factored=False, deterministic=False) + _initialize(source, source_parameter, variant="block") + document = _export_portable_state( + source, + checkpoint_process_group=source_binding, + transaction_id="export-corruption-source", + limits=_limits(), + ) + corrupt = copy.deepcopy(document) + corrupt["completion"]["digest"] = "0" * 64 + before = target._canonical_import_live_token() + + with pytest.raises(RuntimeError, match="digest"): + _import_portable_state( + target, + corrupt, + checkpoint_process_group=target_binding, + transaction_id="import-corruption", + limits=_limits(), + ) + + assert target._canonical_import_live_token() == before + + +def test_export_rejects_non_period_one_and_unknown_group_state(): + optimizer, parameter, binding = _plain(factored=False, deterministic=False) + optimizer._gefen_global_step = 2 + optimizer._gefen_codebook = torch.linspace(-1.0, 1.0, 256, dtype=torch.float32) + optimizer.state[parameter].update( + { + "automatic_period": 2, + "step": 2, + "m_codebook": torch.zeros((3, 2), dtype=torch.uint8), + "m_magnitude": torch.ones((3, 1), dtype=torch.float32), + "vmean": torch.ones((3, 1), dtype=torch.float32), + "vmean_step": 2, + } + ) + with pytest.raises(RuntimeError, match="period one"): + _export_portable_state( + optimizer, + checkpoint_process_group=binding, + transaction_id="export-period-two", + limits=_limits(), + ) + + optimizer.state[parameter] = {"name": "weight"} + optimizer.param_groups[0]["extension"] = 1 + with pytest.raises(RuntimeError, match="unknown keys"): + _export_portable_state( + optimizer, + checkpoint_process_group=binding, + transaction_id="export-unknown-group-option", + limits=_limits(), + ) + + +def test_exact_type_and_full_process_group_identity_are_required(): + class GefenSubclass(Gefen): + pass + + parameter = torch.nn.Parameter(torch.ones(2, dtype=torch.float32)) + optimizer = GefenSubclass([("weight", parameter)], fused=False) + with pytest.raises(TypeError, match="exact Gefen"): + _optimizer_implementation(optimizer) + + exact, _, _ = _plain(factored=False, deterministic=False) + wrong_group = ProcessGroupIdentity("different-checkpoint", (_MEMBER,)) + _, wrong_binding = _bindings(wrong_group) + with pytest.raises(RuntimeError, match="identities must match"): + _export_portable_state( + exact, + checkpoint_process_group=wrong_binding, + transaction_id="export-wrong-process-group", + limits=_limits(), + ) + + +def test_runtime_readiness_fails_closed_on_state_group_and_subclass_mutation(): + optimizer, parameter, binding = _plain(factored=False, deterministic=False) + assert _portable_runtime_layouts(optimizer) == frozenset({ParameterLayout.REPLICATED}) + + optimizer.state[parameter]["automatic_period"] = 2 + assert not _portable_runtime_layouts(optimizer) + assert all( + support.transport is not CheckpointTransport.CANONICAL_GLOBAL + for support in optimizer.optimizer_contract().capabilities.checkpoints + ) + optimizer.state[parameter] = {"name": "weight"} + optimizer.param_groups[0]["extension"] = 1 + assert not _portable_runtime_layouts(optimizer) + optimizer.param_groups[0].pop("extension") + + _initialize(optimizer, parameter, variant="block") + optimizer.state[parameter]["m_magnitude"] = optimizer.state[parameter]["m_magnitude"].requires_grad_() + assert not _portable_runtime_layouts(optimizer) + with pytest.raises(RuntimeError, match="plain materialized"): + _export_portable_state( + optimizer, + checkpoint_process_group=binding, + transaction_id="export-grad-state", + limits=_limits(), + ) + + class GefenSubclass(Gefen): + pass + + subclass_parameter = torch.nn.Parameter(torch.ones(2)) + subclassed = GefenSubclass([("weight", subclass_parameter)], fused=False) + assert not _portable_runtime_layouts(subclassed) + + +def test_muon_readiness_rejects_nonexact_beta_and_unsafe_raw_schedule(): + optimizer, _, binding = _muon(deterministic=False) + optimizer.param_groups[0]["beta2"] = torch.tensor(0.0) + assert not _portable_runtime_layouts(optimizer) + with pytest.raises(RuntimeError, match="exact floats"): + _export_portable_state( + optimizer, + checkpoint_process_group=binding, + transaction_id="export-tensor-beta", + limits=_limits(), + ) + + optimizer.param_groups[0]["beta2"] = 0.0 + optimizer.param_groups[0]["ns_schedule"] = lambda: None + assert not _portable_runtime_layouts(optimizer) + with pytest.raises(RuntimeError, match="weights-only-safe"): + _export_portable_state( + optimizer, + checkpoint_process_group=binding, + transaction_id="export-callable-schedule", + limits=_limits(), + ) + + +def test_invalid_collective_arguments_fail_through_the_status_gate(): + optimizer, _, binding = _plain(factored=False, deterministic=False) + with pytest.raises(RuntimeError, match="PortableStateLimits"): + _export_portable_state( + optimizer, + checkpoint_process_group=binding, + transaction_id="invalid-limits", + limits=object(), + ) + with pytest.raises(RuntimeError, match="transaction_id"): + _export_portable_state( + optimizer, + checkpoint_process_group=binding, + transaction_id="", + limits=_limits(), + ) + with pytest.raises(RuntimeError, match="trimmed"): + _export_portable_state( + optimizer, + checkpoint_process_group=binding, + transaction_id=" invalid-transaction ", + limits=_limits(), + ) diff --git a/tests/test_portable_runtime_consensus.py b/tests/test_portable_runtime_consensus.py new file mode 100644 index 0000000..6af5ec0 --- /dev/null +++ b/tests/test_portable_runtime_consensus.py @@ -0,0 +1,477 @@ +"""Warning-strict Gloo regressions for portable runtime consensus gates.""" + +from datetime import timedelta +import multiprocessing as mp +import os +import queue as queue_module +import tempfile +import traceback +import warnings + +import pytest +import torch +import torch.distributed as dist + +from gefen.checkpoint import CheckpointProcessGroupBinding +from gefen.codebook import CodebookProcessGroupBinding +from gefen.contracts import ( + LogicalSlice, + ParameterIdentity, + ParameterLayout, + PlacementKind, + ProcessGroupIdentity, + ShardIdentity, + ShardPlacement, + ShardingManifest, +) +from gefen.gefen import Gefen +from gefen.portable_state import PortableStateLimits +from gefen.rebinding import ParameterRebinding + + +_WORLD = 2 +_SINGLE_FQN = "model.weight" + + +class _CheckpointBindingSubclass(CheckpointProcessGroupBinding): + pass + + +def _members(): + return tuple("rank:{}".format(rank) for rank in range(_WORLD)) + + +def _limits(): + return PortableStateLimits( + max_fragment_tensor_bytes=1 << 20, + max_collective_tensor_bytes=4 << 20, + max_collective_metadata_bytes=4 << 20, + chunk_bytes=17, + max_members=4, + max_metadata_bytes=1 << 20, + max_tree_nodes=10_000, + max_tree_depth=32, + max_container_items=10_000, + max_string_bytes=16 << 10, + max_integer_bytes=128, + max_tensors=256, + max_tensor_rank=8, + diagnostic_bytes=1024, + ) + + +def _bindings(group, rank): + member = _members()[rank] + codebook = CodebookProcessGroupBinding( + group, + member, + dist.group.WORLD, + torch.device("cpu"), + ) + checkpoint = CheckpointProcessGroupBinding( + group, + member, + dist.group.WORLD, + torch.device("cpu"), + ) + return codebook, checkpoint + + +def _replicated_shards(identity, group): + return tuple( + ShardIdentity( + identity, + ParameterLayout.REPLICATED, + LogicalSlice.full(identity), + placements=( + ShardPlacement( + "checkpoint", + PlacementKind.REPLICATE, + coordinate, + len(group.ordered_members), + ), + ), + process_group=group, + local_member=member, + ) + for coordinate, member in enumerate(group.ordered_members) + ) + + +def _new_optimizer(entries, *, deterministic): + return Gefen( + entries, + lr=2.5e-3, + betas=(0.8, 0.97), + eps=2.0e-8, + weight_decay=0.03, + fused=False, + force_2d_period_one=True, + factored_v_2d=False, + deterministic=deterministic, + ) + + +def _make_single(rank, group, *, compatibility_name, deterministic): + identity = ParameterIdentity(_SINGLE_FQN, (2, 3)) + parameter = torch.nn.Parameter(torch.zeros(identity.global_shape, dtype=torch.float32)) + optimizer = _new_optimizer( + [(compatibility_name, parameter)], + deterministic=deterministic, + ) + shards = _replicated_shards(identity, group) + codebook_binding, checkpoint_binding = _bindings(group, rank) + optimizer.post_sharding( + (ParameterRebinding(parameter, parameter, shards[rank]),), + manifest=ShardingManifest(shards), + codebook_process_group=codebook_binding, + ) + return optimizer, checkpoint_binding + + +def _make_two(rank, group, *, reverse, deterministic): + identities = ( + ParameterIdentity("model.alpha", (2, 2)), + ParameterIdentity("model.beta", (2, 2)), + ) + alpha = torch.nn.Parameter(torch.zeros(identities[0].global_shape, dtype=torch.float32)) + beta = torch.nn.Parameter(torch.zeros(identities[1].global_shape, dtype=torch.float32)) + entries = [("alpha", alpha), ("beta", beta)] + if reverse: + entries.reverse() + optimizer = _new_optimizer(entries, deterministic=deterministic) + shards_by_identity = tuple(_replicated_shards(identity, group) for identity in identities) + manifest = ShardingManifest(tuple(shard for shards in shards_by_identity for shard in shards)) + codebook_binding, checkpoint_binding = _bindings(group, rank) + optimizer.post_sharding( + ( + ParameterRebinding(alpha, alpha, shards_by_identity[0][rank]), + ParameterRebinding(beta, beta, shards_by_identity[1][rank]), + ), + manifest=manifest, + codebook_process_group=codebook_binding, + ) + return optimizer, checkpoint_binding + + +def _subclass_binding(binding): + return _CheckpointBindingSubclass( + binding.identity, + binding.local_member, + binding.process_group, + binding.collective_device, + ) + + +def _binding_with_process_group(binding, process_group): + return CheckpointProcessGroupBinding( + binding.identity, + binding.local_member, + process_group, + binding.collective_device, + ) + + +def _attempt_failure(optimizer, operation): + state_mapping = optimizer.state + param_groups = optimizer.param_groups + defaults = optimizer.defaults + parameters = tuple(parameter for group in optimizer.param_groups for parameter in group["params"]) + state_objects = tuple(optimizer.state[parameter] for parameter in parameters) + token = optimizer._canonical_import_live_token() + global_step = optimizer._gefen_global_step + codebook = optimizer._gefen_codebook + deterministic = optimizer._deterministic + try: + operation() + except RuntimeError as exc: + message = str(exc) + else: + message = None + return { + "message": message, + "token_unchanged": optimizer._canonical_import_live_token() == token, + "containers_unchanged": ( + optimizer.state is state_mapping + and optimizer.param_groups is param_groups + and optimizer.defaults is defaults + ), + "state_objects_unchanged": all( + optimizer.state[parameter] is state for parameter, state in zip(parameters, state_objects) + ), + "states_pristine": all( + optimizer.state[parameter] == {"name": optimizer._param_names[parameter]} for parameter in parameters + ), + "common_unchanged": ( + optimizer._gefen_global_step == global_step + and optimizer._gefen_codebook is codebook + and optimizer._deterministic is deterministic + ), + } + + +def _worker(rank, init_file, result_queue): + warnings.simplefilter("error") + alternate_process_group = None + try: + dist.init_process_group( + "gloo", + init_method="file://{}".format(init_file), + rank=rank, + world_size=_WORLD, + timeout=timedelta(seconds=20), + ) + group = ProcessGroupIdentity("portable_consensus", _members()) + alternate_process_group = dist.new_group( + ranks=list(range(_WORLD)), + backend="gloo", + timeout=timedelta(seconds=20), + ) + + source, exact_binding = _make_single( + rank, + group, + compatibility_name="weight", + deterministic=True, + ) + document = source.export_portable_state( + checkpoint_process_group=exact_binding, + transaction_id="consensus-healthy-single-export", + limits=_limits(), + ) + asymmetric_binding = _subclass_binding(exact_binding) if rank == 0 else exact_binding + binding_export = _attempt_failure( + source, + lambda: source.export_portable_state( + checkpoint_process_group=asymmetric_binding, + transaction_id="consensus-subclass-export", + limits=_limits(), + ), + ) + dist.barrier() + object_binding = object() if rank == 0 else exact_binding + object_export = _attempt_failure( + source, + lambda: source.export_portable_state( + checkpoint_process_group=object_binding, + transaction_id="consensus-object-export", + limits=_limits(), + ), + ) + dist.barrier() + alternate_binding = ( + _binding_with_process_group(exact_binding, alternate_process_group) if rank == 0 else exact_binding + ) + alternate_export = _attempt_failure( + source, + lambda: source.export_portable_state( + checkpoint_process_group=alternate_binding, + transaction_id="consensus-alternate-handle-export", + limits=_limits(), + ), + ) + dist.barrier() + + binding_target, target_binding = _make_single( + rank, + group, + compatibility_name="weight", + deterministic=False, + ) + asymmetric_target_binding = _subclass_binding(target_binding) if rank == 0 else target_binding + binding_import = _attempt_failure( + binding_target, + lambda: binding_target.import_portable_state( + document, + checkpoint_process_group=asymmetric_target_binding, + transaction_id="consensus-subclass-import", + limits=_limits(), + ), + ) + dist.barrier() + object_target, object_target_binding = _make_single( + rank, + group, + compatibility_name="weight", + deterministic=False, + ) + asymmetric_object_binding = object() if rank == 0 else object_target_binding + object_import = _attempt_failure( + object_target, + lambda: object_target.import_portable_state( + document, + checkpoint_process_group=asymmetric_object_binding, + transaction_id="consensus-object-import", + limits=_limits(), + ), + ) + dist.barrier() + alternate_target, alternate_target_binding = _make_single( + rank, + group, + compatibility_name="weight", + deterministic=False, + ) + asymmetric_alternate_binding = ( + _binding_with_process_group( + alternate_target_binding, + alternate_process_group, + ) + if rank == 0 + else alternate_target_binding + ) + alternate_import = _attempt_failure( + alternate_target, + lambda: alternate_target.import_portable_state( + document, + checkpoint_process_group=asymmetric_alternate_binding, + transaction_id="consensus-alternate-handle-import", + limits=_limits(), + ), + ) + dist.barrier() + + name_target, name_binding = _make_single( + rank, + group, + compatibility_name="weight" if rank == 0 else "other_weight", + deterministic=False, + ) + compatibility_import = _attempt_failure( + name_target, + lambda: name_target.import_portable_state( + document, + checkpoint_process_group=name_binding, + transaction_id="consensus-compatibility-import", + limits=_limits(), + ), + ) + dist.barrier() + + two_source, two_source_binding = _make_two( + rank, + group, + reverse=False, + deterministic=True, + ) + two_document = two_source.export_portable_state( + checkpoint_process_group=two_source_binding, + transaction_id="consensus-healthy-two-export", + limits=_limits(), + ) + position_target, position_binding = _make_two( + rank, + group, + reverse=rank == 1, + deterministic=False, + ) + position_import = _attempt_failure( + position_target, + lambda: position_target.import_portable_state( + two_document, + checkpoint_process_group=position_binding, + transaction_id="consensus-position-import", + limits=_limits(), + ), + ) + dist.barrier() + + result_queue.put( + { + "rank": rank, + "binding_export": binding_export, + "binding_import": binding_import, + "object_export": object_export, + "object_import": object_import, + "alternate_export": alternate_export, + "alternate_import": alternate_import, + "compatibility_import": compatibility_import, + "position_import": position_import, + } + ) + except Exception: + result_queue.put({"rank": rank, "fatal_error": traceback.format_exc()}) + finally: + if dist.is_available() and dist.is_initialized(): + if alternate_process_group is not None: + dist.destroy_process_group(alternate_process_group) + dist.destroy_process_group() + + +def _run_workers(): + context = mp.get_context("spawn") + result_queue = context.Queue() + descriptor, init_file = tempfile.mkstemp(prefix="gefen-portable-consensus-") + os.close(descriptor) + os.unlink(init_file) + processes = [context.Process(target=_worker, args=(rank, init_file, result_queue)) for rank in range(_WORLD)] + results = [] + try: + for process in processes: + process.start() + try: + for _ in processes: + results.append(result_queue.get(timeout=60)) + except queue_module.Empty: + pass + for process in processes: + process.join(timeout=5) + for process in processes: + if process.is_alive(): + process.terminate() + process.join(timeout=5) + finally: + for process in processes: + if process.is_alive(): + process.terminate() + process.join(timeout=5) + result_queue.close() + result_queue.join_thread() + if os.path.exists(init_file): + os.unlink(init_file) + assert len(results) == _WORLD, (results, [process.exitcode for process in processes]) + assert all(process.exitcode == 0 for process in processes), [process.exitcode for process in processes] + return sorted(results, key=lambda item: item["rank"]) + + +@pytest.mark.skipif( + not dist.is_available() or not dist.is_gloo_available(), + reason="portable runtime consensus coverage requires Gloo", +) +def test_two_process_portable_runtime_rejects_asymmetric_consensus_inputs(): + results = _run_workers() + assert all("fatal_error" not in result for result in results), results + + expected_messages = { + "binding_export": "CheckpointProcessGroupBinding", + "binding_import": "CheckpointProcessGroupBinding", + "object_export": "CheckpointProcessGroupBinding", + "object_import": "CheckpointProcessGroupBinding", + "alternate_export": "optimizer-owned runtime transport", + "alternate_import": "optimizer-owned runtime transport", + "compatibility_import": "context", + "position_import": "context", + } + failures = [] + for case, expected_message in expected_messages.items(): + messages = [result[case]["message"] for result in results] + if messages[0] != messages[1]: + failures.append((case, "messages diverged", messages)) + if messages[0] is None or expected_message not in messages[0]: + failures.append((case, "missing expected failure", messages)) + mutation_free = all( + all( + result[case][key] + for key in ( + "token_unchanged", + "containers_unchanged", + "state_objects_unchanged", + "states_pristine", + "common_unchanged", + ) + ) + for result in results + ) + if not mutation_free: + failures.append((case, "live state changed", [result[case] for result in results])) + assert not failures, (failures, results) diff --git a/tests/test_portable_runtime_distributed.py b/tests/test_portable_runtime_distributed.py new file mode 100644 index 0000000..1996c37 --- /dev/null +++ b/tests/test_portable_runtime_distributed.py @@ -0,0 +1,567 @@ +"""Warning-strict distributed coverage for topology-neutral optimizer state.""" + +from datetime import timedelta +import multiprocessing as mp +import os +import queue as queue_module +import tempfile +import traceback +import warnings + +import pytest +import torch +import torch.distributed as dist + +from gefen.checkpoint import CheckpointProcessGroupBinding +from gefen.codebook import CodebookProcessGroupBinding +from gefen.contracts import ( + LogicalSlice, + ParameterIdentity, + ParameterLayout, + PlacementKind, + ProcessGroupIdentity, + ShardIdentity, + ShardPlacement, + ShardingManifest, +) +from gefen.gefen import Gefen +from gefen.gefen_muon import GefenMuon +from gefen.portable import _decode_quantized_momentum +from gefen.portable_state import PortableStateLimits +from gefen.rebinding import ParameterRebinding + + +_WORLD = 2 +_PLAIN_FQN = "model.weight" +_MUON_FQN = "model.matrix" + + +def _limits(): + return PortableStateLimits( + max_fragment_tensor_bytes=1 << 20, + max_collective_tensor_bytes=4 << 20, + max_collective_metadata_bytes=4 << 20, + chunk_bytes=7, + max_members=4, + max_metadata_bytes=1 << 20, + max_tree_nodes=10_000, + max_tree_depth=32, + max_container_items=10_000, + max_string_bytes=16 << 10, + max_integer_bytes=128, + max_tensors=256, + max_tensor_rank=8, + diagnostic_bytes=1024, + ) + + +def _members(): + return tuple("rank:{}".format(rank) for rank in range(_WORLD)) + + +def _bindings(group, rank): + member = _members()[rank] + return ( + CodebookProcessGroupBinding(group, member, dist.group.WORLD, torch.device("cpu")), + CheckpointProcessGroupBinding(group, member, dist.group.WORLD, torch.device("cpu")), + ) + + +def _replicated_manifest(identity, group): + members = group.ordered_members + shards = tuple( + ShardIdentity( + identity, + ParameterLayout.REPLICATED, + LogicalSlice.full(identity), + placements=( + ShardPlacement( + "checkpoint", + PlacementKind.REPLICATE, + coordinate, + len(members), + ), + ), + process_group=group, + local_member=member, + ) + for coordinate, member in enumerate(members) + ) + return ShardingManifest(shards), shards + + +def _flat_manifest(identity, group, lengths): + members = group.ordered_members + if len(lengths) != len(members) or sum(lengths) != identity.numel: + raise AssertionError("invalid test partition") + offset = 0 + shards = [] + for coordinate, (member, length) in enumerate(zip(members, lengths)): + shards.append( + ShardIdentity( + identity, + ParameterLayout.FLATTENED_ELEMENT_SHARD, + LogicalSlice(offset, length), + placements=( + ShardPlacement( + "checkpoint", + PlacementKind.FLAT_SHARD, + coordinate, + len(members), + ), + ), + process_group=group, + local_member=member, + ) + ) + offset += length + return ShardingManifest(tuple(shards)), tuple(shards) + + +def _owner_manifest(identity, group, owner): + members = group.ordered_members + shards = tuple( + ShardIdentity( + identity, + ParameterLayout.WHOLE_PARAMETER_OWNER, + LogicalSlice.full(identity) if member == owner else LogicalSlice(0, 0), + placements=( + ShardPlacement( + "checkpoint", + PlacementKind.WHOLE_PARAMETER_OWNER, + coordinate, + len(members), + ), + ), + process_group=group, + local_member=member, + owner=owner, + ) + for coordinate, member in enumerate(members) + ) + return ShardingManifest(shards), shards + + +def _codebook(): + return torch.linspace(-1.0, 1.0, 256, dtype=torch.float32) + + +def _momentum_values(shape): + bits = torch.tensor( + [ + -2147483648, + 0, + 1, + -2147483647, + 1056964608, + -1090519040, + 1078984704, + -1058013184, + ], + dtype=torch.int32, + ) + return bits.view(torch.float32).reshape(shape).clone() + + +def _second_moment_values(shape): + bits = torch.tensor( + [ + -2147483648, + 0, + 1, + 8388608, + 1048576000, + 1065353216, + 1073741824, + 2139095039, + ], + dtype=torch.int32, + ) + return bits.view(torch.float32).reshape(shape).clone() + + +def _quantized_period_one(momentum): + indices = torch.where( + torch.signbit(momentum.reshape(-1)), + torch.zeros(momentum.numel(), dtype=torch.uint8), + torch.full((momentum.numel(),), 255, dtype=torch.uint8), + ) + return indices.reshape(-1, 1), momentum.abs().reshape(-1, 1).clone() + + +def _bits_equal(left, right): + return ( + left.dtype == right.dtype == torch.float32 + and tuple(left.shape) == tuple(right.shape) + and torch.equal(left.contiguous().view(torch.int32), right.contiguous().view(torch.int32)) + ) + + +def _make_plain_replicated(rank, group, *, deterministic): + identity = ParameterIdentity(_PLAIN_FQN, (2, 4)) + parameter = torch.nn.Parameter(torch.zeros(identity.global_shape, dtype=torch.float32)) + optimizer = Gefen( + [("weight", parameter)], + lr=2.5e-3, + betas=(0.8, 0.97), + eps=2.0e-8, + weight_decay=0.03, + fused=False, + force_2d_period_one=True, + factored_v_2d=False, + deterministic=deterministic, + ) + manifest, shards = _replicated_manifest(identity, group) + codebook_binding, checkpoint_binding = _bindings(group, rank) + optimizer.post_sharding( + (ParameterRebinding(parameter, parameter, shards[rank]),), + manifest=manifest, + codebook_process_group=codebook_binding, + ) + return optimizer, parameter, checkpoint_binding + + +def _make_plain_flat(rank, group, lengths, *, deterministic): + identity = ParameterIdentity(_PLAIN_FQN, (2, 4)) + original = torch.nn.Parameter(torch.zeros(identity.global_shape, dtype=torch.float32)) + local = torch.nn.Parameter(torch.zeros(lengths[rank], dtype=torch.float32)) + optimizer = Gefen( + [("weight", original)], + lr=2.5e-3, + betas=(0.8, 0.97), + eps=2.0e-8, + weight_decay=0.03, + fused=False, + force_2d_period_one=True, + factored_v_2d=False, + deterministic=deterministic, + ) + manifest, shards = _flat_manifest(identity, group, lengths) + codebook_binding, checkpoint_binding = _bindings(group, rank) + optimizer.post_sharding( + (ParameterRebinding(original, local, shards[rank]),), + manifest=manifest, + codebook_process_group=codebook_binding, + ) + return optimizer, local, shards[rank], checkpoint_binding + + +def _seed_plain_state(optimizer, parameter): + momentum = _momentum_values((2, 4)) + second_moment = _second_moment_values((2, 4)) + indices, magnitudes = _quantized_period_one(momentum) + optimizer._gefen_global_step = 9 + optimizer._gefen_codebook = _codebook() + optimizer.state[parameter].update( + { + "automatic_period": 1, + "step": 7, + "m_codebook": indices, + "m_magnitude": magnitudes, + "vmean": second_moment.reshape(-1, 1).clone(), + "vmean_step": 6, + } + ) + return momentum, second_moment + + +def _make_muon_owner(rank, group, *, owner, deterministic): + identity = ParameterIdentity(_MUON_FQN, (3, 2)) + original = torch.nn.Parameter(torch.zeros(identity.global_shape, dtype=torch.float32)) + optimizer = GefenMuon( + [("matrix", original)], + lr=3.0e-3, + weight_decay=0.02, + momentum=0.85, + nesterov=False, + ns_steps=2, + fused=False, + sharded_mode="distributed", + deterministic=deterministic, + normuon=True, + normuon_beta2=0.9, + normuon_eps=3.0e-8, + ) + manifest, shards = _owner_manifest(identity, group, owner) + local = original if _members()[rank] == owner else None + codebook_binding, checkpoint_binding = _bindings(group, rank) + optimizer.post_sharding( + (ParameterRebinding(original, local, shards[rank]),), + manifest=manifest, + codebook_process_group=codebook_binding, + ) + return optimizer, local, checkpoint_binding + + +def _seed_muon_state(optimizer, parameter): + optimizer._gefen_global_step = 13 + optimizer._gefen_codebook = _codebook() + if parameter is None: + return + momentum = _momentum_values((2, 4)).reshape(-1)[:6].reshape(3, 2).clone() + indices, magnitudes = _quantized_period_one(momentum) + normuon_bits = torch.tensor([-2147483648, 1, 1082130432], dtype=torch.int32) + optimizer.state[parameter].update( + { + "automatic_period": 1, + "step": 11, + "m_codebook": indices, + "m_magnitude": magnitudes, + "normuon_v": normuon_bits.view(torch.float32).reshape(3, 1).clone(), + "normuon_step": 10, + } + ) + + +def _plain_result(rank, group): + source, source_parameter, source_binding = _make_plain_replicated(rank, group, deterministic=True) + expected_momentum, expected_second = _seed_plain_state(source, source_parameter) + document = source.export_portable_state( + checkpoint_process_group=source_binding, + transaction_id="plain-replicated-export-v1", + limits=_limits(), + ) + + lengths = (3, 5) + target, target_parameter, target_shard, target_binding = _make_plain_flat( + rank, + group, + lengths, + deterministic=False, + ) + target.import_portable_state( + document, + checkpoint_process_group=target_binding, + transaction_id="plain-flat-import-v1", + limits=_limits(), + ) + state = target.state[target_parameter] + decoded = _decode_quantized_momentum( + target._gefen_codebook, + state["m_codebook"], + state["m_magnitude"], + logical_shape=(lengths[rank],), + period=1, + step=state["step"], + ) + start = target_shard.logical_slice.flat_offset + stop = start + target_shard.logical_slice.length + expected_local_momentum = expected_momentum.reshape(-1)[start:stop].clone() + expected_local_second = expected_second.reshape(-1)[start:stop].clone() + parameter_record = document["parameters"][_PLAIN_FQN]["state"] + return { + "digest": document["completion"]["digest"], + "document_momentum_exact": _bits_equal(parameter_record["momentum"], expected_momentum), + "document_second_exact": _bits_equal(parameter_record["second_moment"], expected_second), + "local_momentum_exact": _bits_equal(decoded, expected_local_momentum), + "local_second_exact": _bits_equal(state["vmean"].reshape(-1), expected_local_second), + "counters_exact": target._gefen_global_step == 9 and state["step"] == 7 and state["vmean_step"] == 6, + "deterministic_imported": target._deterministic is True, + "document": document, + } + + +def _asymmetric_failure_result(rank, group, document): + target, target_parameter, _, binding = _make_plain_flat(rank, group, (5, 3), deterministic=False) + if rank == 0: + target.param_groups[0]["rank_local_extension"] = "reject-me" + token_before = target._canonical_import_live_token() + state_object_before = target.state[target_parameter] + try: + target.import_portable_state( + document, + checkpoint_process_group=binding, + transaction_id="plain-asymmetric-import-failure-v1", + limits=_limits(), + ) + except RuntimeError as exc: + message = str(exc) + else: + message = None + return { + "message": message, + "token_unchanged": target._canonical_import_live_token() == token_before, + "state_object_unchanged": target.state[target_parameter] is state_object_before, + "still_pristine": target.state[target_parameter] == {"name": "weight"}, + "common_unchanged": target._gefen_global_step == 0 and target._gefen_codebook is None, + } + + +def _muon_result(rank, group): + source_owner = _members()[0] + target_owner = _members()[1] + source, source_parameter, source_binding = _make_muon_owner( + rank, + group, + owner=source_owner, + deterministic=True, + ) + _seed_muon_state(source, source_parameter) + document = source.export_portable_state( + checkpoint_process_group=source_binding, + transaction_id="muon-owner-export-v1", + limits=_limits(), + ) + + target, target_parameter, target_binding = _make_muon_owner( + rank, + group, + owner=target_owner, + deterministic=False, + ) + target.import_portable_state( + document, + checkpoint_process_group=target_binding, + transaction_id="muon-owner-import-v1", + limits=_limits(), + ) + expected_momentum = _momentum_values((2, 4)).reshape(-1)[:6].reshape(3, 2).clone() + expected_normuon = torch.tensor([-2147483648, 1, 1082130432], dtype=torch.int32).view(torch.float32).reshape(3, 1) + record = document["parameters"][_MUON_FQN]["state"] + if target_parameter is None: + local_exact = len(target.state) == 0 and len(target.param_groups[0]["params"]) == 0 + counters_exact = target._gefen_global_step == 13 + else: + state = target.state[target_parameter] + decoded = _decode_quantized_momentum( + target._gefen_codebook, + state["m_codebook"], + state["m_magnitude"], + logical_shape=(3, 2), + period=1, + step=state["step"], + ) + local_exact = _bits_equal(decoded, expected_momentum) and _bits_equal(state["normuon_v"], expected_normuon) + counters_exact = target._gefen_global_step == 13 and state["step"] == 11 and state["normuon_step"] == 10 + return { + "digest": document["completion"]["digest"], + "document_momentum_exact": _bits_equal(record["momentum"], expected_momentum), + "document_normuon_exact": _bits_equal(record["normuon_v"], expected_normuon), + "local_exact": local_exact, + "counters_exact": counters_exact, + "deterministic_imported": target._deterministic is True, + "target_has_storage": target_parameter is not None, + } + + +def _distributed_worker(rank, init_file, result_queue): + warnings.simplefilter("error") + try: + dist.init_process_group( + "gloo", + init_method="file://{}".format(init_file), + rank=rank, + world_size=_WORLD, + timeout=timedelta(seconds=90), + ) + group = ProcessGroupIdentity("portable_checkpoint", _members()) + plain = _plain_result(rank, group) + dist.barrier() + asymmetric = _asymmetric_failure_result(rank, group, plain.pop("document")) + dist.barrier() + muon = _muon_result(rank, group) + dist.barrier() + result_queue.put( + { + "rank": rank, + "plain": plain, + "asymmetric": asymmetric, + "muon": muon, + } + ) + except Exception: + result_queue.put({"rank": rank, "fatal_error": traceback.format_exc()}) + finally: + if dist.is_available() and dist.is_initialized(): + dist.destroy_process_group() + + +def _run_distributed_workers(): + context = mp.get_context("spawn") + result_queue = context.Queue() + descriptor, init_file = tempfile.mkstemp(prefix="gefen-portable-runtime-") + os.close(descriptor) + os.unlink(init_file) + processes = [ + context.Process(target=_distributed_worker, args=(rank, init_file, result_queue)) for rank in range(_WORLD) + ] + results = [] + try: + for process in processes: + process.start() + try: + for _ in processes: + results.append(result_queue.get(timeout=180)) + except queue_module.Empty: + pass + for process in processes: + process.join(timeout=10) + for process in processes: + if process.is_alive(): + process.terminate() + process.join(timeout=5) + finally: + for process in processes: + if process.is_alive(): + process.terminate() + process.join(timeout=5) + result_queue.close() + result_queue.join_thread() + if os.path.exists(init_file): + os.unlink(init_file) + assert len(results) == _WORLD, (results, [process.exitcode for process in processes]) + assert all(process.exitcode == 0 for process in processes), [process.exitcode for process in processes] + return sorted(results, key=lambda item: item["rank"]) + + +@pytest.mark.skipif( + not dist.is_available() or not dist.is_gloo_available(), + reason="portable runtime topology-change coverage requires Gloo", +) +def test_two_process_portable_runtime_reshards_relocates_and_fails_atomically(): + results = _run_distributed_workers() + + assert all("fatal_error" not in result for result in results), results + assert len({result["plain"]["digest"] for result in results}) == 1 + assert len({result["muon"]["digest"] for result in results}) == 1 + for result in results: + plain = result["plain"] + assert all( + plain[key] + for key in ( + "document_momentum_exact", + "document_second_exact", + "local_momentum_exact", + "local_second_exact", + "counters_exact", + "deterministic_imported", + ) + ), result + muon = result["muon"] + assert all( + muon[key] + for key in ( + "document_momentum_exact", + "document_normuon_exact", + "local_exact", + "counters_exact", + "deterministic_imported", + ) + ), result + assert muon["target_has_storage"] is (result["rank"] == 1) + + messages = [result["asymmetric"]["message"] for result in results] + assert messages[0] == messages[1] + assert messages[0] is not None and "unknown keys" in messages[0] + assert all( + all( + result["asymmetric"][key] + for key in ( + "token_unchanged", + "state_object_unchanged", + "still_pristine", + "common_unchanged", + ) + ) + for result in results + ), results diff --git a/tests/test_portable_runtime_hardening.py b/tests/test_portable_runtime_hardening.py new file mode 100644 index 0000000..9d1ccbf --- /dev/null +++ b/tests/test_portable_runtime_hardening.py @@ -0,0 +1,457 @@ +"""Warning-strict CPU regressions for portable runtime publication guards.""" + +import pytest +import torch + +import gefen.portable_runtime as portable_runtime +from gefen.checkpoint import CheckpointProcessGroupBinding +from gefen.codebook import CodebookProcessGroupBinding +from gefen.contracts import ( + CheckpointTransport, + LogicalSlice, + ParameterIdentity, + ParameterLayout, + PlacementKind, + ProcessGroupIdentity, + ShardIdentity, + ShardPlacement, + ShardingManifest, +) +from gefen.gefen import Gefen +from gefen.portable_state import PortableStateLimits +from gefen.rebinding import ParameterRebinding + + +_MEMBER = "rank:0" + + +class _ParameterIdentitySubclass(ParameterIdentity): + pass + + +class _ProcessGroupIdentitySubclass(ProcessGroupIdentity): + pass + + +class _ShardPlacementSubclass(ShardPlacement): + pass + + +class _StringSubclass(str): + pass + + +def _limits(*, max_fragment_tensor_bytes=4 << 20): + return PortableStateLimits( + max_fragment_tensor_bytes=max_fragment_tensor_bytes, + max_collective_tensor_bytes=16 << 20, + max_collective_metadata_bytes=64 << 20, + ) + + +@pytest.mark.parametrize("subclass_kind", ["parameter", "process_group", "placement"]) +def test_nested_identity_subclasses_are_rejected(subclass_kind): + parameter_type = _ParameterIdentitySubclass if subclass_kind == "parameter" else ParameterIdentity + group_type = _ProcessGroupIdentitySubclass if subclass_kind == "process_group" else ProcessGroupIdentity + placement_type = _ShardPlacementSubclass if subclass_kind == "placement" else ShardPlacement + identity = parameter_type("layer.weight", (2, 3)) + group = group_type("checkpoint", (_MEMBER,)) + shard = ShardIdentity( + identity, + ParameterLayout.REPLICATED, + LogicalSlice.full(identity), + placements=( + placement_type( + "checkpoint", + PlacementKind.REPLICATE, + 0, + 1, + ), + ), + process_group=group, + local_member=_MEMBER, + ) + + with pytest.raises(TypeError, match="exact"): + portable_runtime._validate_exact_shard_identity(shard) + + +def _optimizer(*, layout, deterministic): + identity = ParameterIdentity("layer.weight", (2, 3)) + if layout is ParameterLayout.REPLICATED: + parameter = torch.nn.Parameter(torch.arange(1, 7, dtype=torch.float32).reshape(identity.global_shape)) + placement_kind = PlacementKind.REPLICATE + elif layout is ParameterLayout.FLATTENED_ELEMENT_SHARD: + parameter = torch.nn.Parameter(torch.arange(1, 7, dtype=torch.float32)) + placement_kind = PlacementKind.FLAT_SHARD + else: + raise AssertionError("unsupported test layout") + + optimizer = Gefen( + [("weight", parameter)], + fused=False, + factored_v_2d=False, + force_2d_period_one=True, + deterministic=deterministic, + ) + group = ProcessGroupIdentity("checkpoint", (_MEMBER,)) + shard = ShardIdentity( + identity, + layout, + LogicalSlice.full(identity), + placements=(ShardPlacement("checkpoint", placement_kind, 0, 1),), + process_group=group, + local_member=_MEMBER, + ) + optimizer.post_sharding( + (ParameterRebinding(parameter, parameter, shard),), + manifest=ShardingManifest((shard,)), + codebook_process_group=CodebookProcessGroupBinding( + group, + _MEMBER, + None, + torch.device("cpu"), + ), + ) + binding = CheckpointProcessGroupBinding( + group, + _MEMBER, + None, + torch.device("cpu"), + ) + return optimizer, parameter, binding + + +def _initialize(optimizer, parameter): + optimizer._gefen_global_step = 4 + optimizer._gefen_codebook = torch.linspace(-1.0, 1.0, 256, dtype=torch.float32) + optimizer.state[parameter].update( + { + "automatic_period": 1, + "step": 4, + "m_codebook": torch.tensor( + [[0], [255], [128], [64], [192], [0]], + dtype=torch.uint8, + ), + "m_magnitude": torch.tensor( + [[1.0], [2.0], [3.0], [4.0], [5.0], [6.0]], + dtype=torch.float32, + ), + "vmean": torch.arange(1, 7, dtype=torch.float32).reshape(6, 1), + "vmean_step": 3, + } + ) + + +def _has_canonical_global(optimizer): + return any( + support.transport is CheckpointTransport.CANONICAL_GLOBAL + for support in optimizer.optimizer_contract().capabilities.checkpoints + ) + + +@pytest.mark.parametrize( + ("descriptor", "field", "value"), + [ + ("manifest", "schema_version", 2), + ("shard", "schema_version", 2), + ("parameter", "schema_version", 2), + ("process_group", "schema_version", 2), + ("process_group", "semantic_name", ""), + ("process_group", "semantic_name", " checkpoint"), + ("process_group", "semantic_name", "check\x00point"), + ("parameter", "fqn", ""), + ("parameter", "fqn", " layer.weight"), + ("parameter", "fqn", "layer..weight"), + ("logical_slice", "flat_offset", -1), + ("placement", "mesh_axis", " checkpoint"), + ("placement", "coordinate", -1), + ("placement", "parts", 0), + ], +) +def test_mutated_descriptor_invariants_disable_portable_capability( + descriptor, + field, + value, +): + optimizer, _, _ = _optimizer( + layout=ParameterLayout.REPLICATED, + deterministic=True, + ) + manifest = optimizer._gefen_sharding_manifest + shard = manifest.shards[0] + target = { + "manifest": manifest, + "shard": shard, + "parameter": shard.parameter, + "process_group": shard.process_group, + "logical_slice": shard.logical_slice, + "placement": shard.placements[0], + }[descriptor] + object.__setattr__(target, field, value) + + with pytest.raises((TypeError, ValueError)): + portable_runtime._validate_exact_manifest(manifest) + assert not _has_canonical_global(optimizer) + + +def test_codebook_binding_local_member_subclass_disables_portable_capability(): + optimizer, _, _ = _optimizer( + layout=ParameterLayout.REPLICATED, + deterministic=True, + ) + scope = optimizer._gefen_codebook_process_group + object.__setattr__(scope, "local_member", _StringSubclass(scope.local_member)) + + assert not _has_canonical_global(optimizer) + + +def _drift_storage(parameter, drift): + if drift == "replicated_shape": + parameter.data = parameter.detach().reshape(3, 2) + elif drift == "flat_shape": + parameter.data = parameter.detach().reshape(2, 3) + elif drift == "flat_noncontiguous": + parameter.data = torch.arange(12, dtype=torch.float32)[::2] + elif drift == "complex_dtype": + parameter.data = parameter.detach().to(dtype=torch.complex64) + else: + raise AssertionError("unknown storage drift") + + +@pytest.mark.parametrize( + ("layout", "drift"), + [ + (ParameterLayout.REPLICATED, "replicated_shape"), + (ParameterLayout.FLATTENED_ELEMENT_SHARD, "flat_shape"), + (ParameterLayout.FLATTENED_ELEMENT_SHARD, "flat_noncontiguous"), + (ParameterLayout.REPLICATED, "complex_dtype"), + ], +) +def test_finalized_parameter_storage_drift_disables_and_rejects_portable_io( + layout, + drift, +): + source, source_parameter, source_binding = _optimizer( + layout=layout, + deterministic=True, + ) + _initialize(source, source_parameter) + document = source.export_portable_state( + checkpoint_process_group=source_binding, + transaction_id="hardening-source-export-{}".format(drift), + limits=_limits(), + ) + + _drift_storage(source_parameter, drift) + assert not _has_canonical_global(source) + with pytest.raises(RuntimeError, match="geometry|dtype"): + source.export_portable_state( + checkpoint_process_group=source_binding, + transaction_id="hardening-drifted-export-{}".format(drift), + limits=_limits(), + ) + + target, target_parameter, target_binding = _optimizer( + layout=layout, + deterministic=False, + ) + _drift_storage(target_parameter, drift) + assert not _has_canonical_global(target) + state_mapping = target.state + parameter_state = target.state[target_parameter] + defaults = target.defaults + with pytest.raises(RuntimeError, match="geometry|dtype"): + target.import_portable_state( + document, + checkpoint_process_group=target_binding, + transaction_id="hardening-drifted-import-{}".format(drift), + limits=_limits(), + ) + assert target.state is state_mapping + assert target.state[target_parameter] is parameter_state + assert parameter_state == {"name": "weight"} + assert target.defaults is defaults + assert target._gefen_global_step == 0 + assert target._gefen_codebook is None + assert target._deterministic is False + + +class _UpdateBombDict(dict): + update_calls = 0 + + def update(self, *args, **kwargs): + self.update_calls += 1 + self["publication_started"] = True + raise AssertionError("defaults publication must not start") + + +def test_defaults_dict_subclass_is_rejected_before_publication(): + source, source_parameter, source_binding = _optimizer( + layout=ParameterLayout.REPLICATED, + deterministic=True, + ) + _initialize(source, source_parameter) + document = source.export_portable_state( + checkpoint_process_group=source_binding, + transaction_id="hardening-defaults-source", + limits=_limits(), + ) + target, target_parameter, target_binding = _optimizer( + layout=ParameterLayout.REPLICATED, + deterministic=False, + ) + hostile_defaults = dict.__new__(_UpdateBombDict) + dict.update(hostile_defaults, target.defaults) + hostile_defaults.update_calls = 0 + target.defaults = hostile_defaults + before_defaults = dict(hostile_defaults) + state_mapping = target.state + parameter_state = target.state[target_parameter] + + assert not _has_canonical_global(target) + with pytest.raises(RuntimeError, match="exact built-in optimizer defaults"): + target.import_portable_state( + document, + checkpoint_process_group=target_binding, + transaction_id="hardening-defaults-import", + limits=_limits(), + ) + + assert target.defaults is hostile_defaults + assert dict(hostile_defaults) == before_defaults + assert hostile_defaults.update_calls == 0 + assert "publication_started" not in hostile_defaults + assert target.state is state_mapping + assert target.state[target_parameter] is parameter_state + assert parameter_state == {"name": "weight"} + assert target._gefen_global_step == 0 + assert target._gefen_codebook is None + assert target._deterministic is False + + +def test_foreign_state_inserted_after_prepare_fails_freshness_without_publication( + monkeypatch, +): + source, source_parameter, source_binding = _optimizer( + layout=ParameterLayout.REPLICATED, + deterministic=True, + ) + _initialize(source, source_parameter) + document = source.export_portable_state( + checkpoint_process_group=source_binding, + transaction_id="hardening-freshness-source", + limits=_limits(), + ) + target, target_parameter, target_binding = _optimizer( + layout=ParameterLayout.REPLICATED, + deterministic=False, + ) + assert _has_canonical_global(target) + state_mapping = target.state + parameter_state = target.state[target_parameter] + foreign_parameter = torch.nn.Parameter(torch.zeros(1, dtype=torch.float32)) + foreign_state = {"name": "foreign"} + injected = False + original_status = portable_runtime._collective_unanimous_status + + def status_with_concurrent_insertion(*args, **kwargs): + nonlocal injected + result = original_status(*args, **kwargs) + if kwargs["operation"] == "portable_import_prepare" and not injected: + target.state[foreign_parameter] = foreign_state + injected = True + return result + + monkeypatch.setattr( + portable_runtime, + "_collective_unanimous_status", + status_with_concurrent_insertion, + ) + with pytest.raises(RuntimeError, match="foreign parameter keys|changed after"): + target.import_portable_state( + document, + checkpoint_process_group=target_binding, + transaction_id="hardening-freshness-import", + limits=_limits(), + ) + + assert injected + assert target.state is state_mapping + assert target.state[target_parameter] is parameter_state + assert parameter_state == {"name": "weight"} + assert target.state[foreign_parameter] is foreign_state + assert target._gefen_global_step == 0 + assert target._gefen_codebook is None + assert target._deterministic is False + + +def test_fragment_limit_rejects_before_dense_momentum_materialization(monkeypatch): + optimizer, parameter, binding = _optimizer( + layout=ParameterLayout.REPLICATED, + deterministic=True, + ) + _initialize(optimizer, parameter) + reached_decode = False + reached_live_token = False + + def fail_if_decoded(*args, **kwargs): + nonlocal reached_decode + reached_decode = True + raise AssertionError("dense momentum decode was reached") + + def fail_if_tokenized(*args, **kwargs): + nonlocal reached_live_token + reached_live_token = True + raise AssertionError("content-bearing live token was reached") + + monkeypatch.setattr( + portable_runtime, + "_decode_local_momentum", + fail_if_decoded, + ) + monkeypatch.setattr( + portable_runtime, + "_portable_live_token", + fail_if_tokenized, + ) + with pytest.raises(RuntimeError, match="max_fragment_tensor_bytes"): + optimizer.export_portable_state( + checkpoint_process_group=binding, + transaction_id="hardening-fragment-budget", + limits=_limits(max_fragment_tensor_bytes=1), + ) + assert not reached_decode + assert not reached_live_token + + +def test_pristine_import_does_not_bound_or_materialize_obsolete_target_payload(): + source, _source_parameter, source_binding = _optimizer( + layout=ParameterLayout.REPLICATED, + deterministic=True, + ) + tiny_limits = _limits(max_fragment_tensor_bytes=1) + document = source.export_portable_state( + checkpoint_process_group=source_binding, + transaction_id="hardening-pristine-source", + limits=tiny_limits, + ) + + target, target_parameter, target_binding = _optimizer( + layout=ParameterLayout.REPLICATED, + deterministic=False, + ) + _initialize(target, target_parameter) + old_state = target.state[target_parameter] + target.import_portable_state( + document, + checkpoint_process_group=target_binding, + transaction_id="hardening-pristine-target", + limits=tiny_limits, + ) + + assert target.state[target_parameter] is not old_state + assert target.state[target_parameter] == {"name": "weight"} + assert target._gefen_global_step == 0 + assert target._gefen_codebook is None + assert target._deterministic is True diff --git a/tests/test_portable_schema.py b/tests/test_portable_schema.py new file mode 100644 index 0000000..67c02b5 --- /dev/null +++ b/tests/test_portable_schema.py @@ -0,0 +1,337 @@ +"""Strict portable-state v3 envelope coverage.""" + +import copy +import io + +import pytest +import torch + +import gefen +import gefen.contracts as contracts_module +import gefen.portable_schema as portable_schema_module +from gefen.portable_schema import ( + PORTABLE_STATE_COVERAGE, + PORTABLE_STATE_DIGEST_ALGORITHM, + PORTABLE_STATE_FORMAT, + PORTABLE_STATE_FORMAT_VERSION, + build_portable_state_document, + normalize_portable_state_document, + portable_state_digest, +) + + +def _parameter_record(*, tensor=None): + if tensor is None: + tensor = torch.arange(6, dtype=torch.float32).reshape(2, 3) + return { + "identity": { + "schema_version": 1, + "fqn": "Model.Weight", + "global_shape": [2, 3], + }, + "algorithm_options": {"lr": 1.0e-3, "betas": (0.9, 0.999)}, + "state_variant": "quantized_momentum", + "state": {"step": 3, "momentum": tensor}, + "projection_hints": {"source_periods": [3]}, + } + + +def _document(*, tensor=None): + return build_portable_state_document( + implementation="gefen.Gefen", + policy={"factored_v_2d": False}, + common={ + "gefen_global_step": 3, + "gefen_deterministic": False, + }, + parameters={"Model.Weight": _parameter_record(tensor=tensor)}, + provenance={"source_layouts": ["replicated"]}, + ) + + +def test_portable_document_is_complete_device_neutral_and_weights_only_safe(): + backing = torch.arange(30, dtype=torch.float32) + view = backing[5:17:2].reshape(2, 3) + document = _document(tensor=view) + + assert document["format"] == PORTABLE_STATE_FORMAT + assert document["format_version"] == PORTABLE_STATE_FORMAT_VERSION + assert document["coverage"] == PORTABLE_STATE_COVERAGE + assert document["completion"]["status"] == "complete" + assert document["completion"]["digest_algorithm"] == PORTABLE_STATE_DIGEST_ALGORITHM + momentum = document["parameters"]["Model.Weight"]["state"]["momentum"] + assert momentum.device.type == "cpu" + assert momentum.is_contiguous() + assert momentum.storage_offset() == 0 + assert ( + momentum.untyped_storage().nbytes() + == momentum.numel() * momentum.element_size() + ) + assert torch.equal(momentum, view) + assert momentum is not view + + buffer = io.BytesIO() + torch.save(document, buffer) + buffer.seek(0) + loaded = torch.load(buffer, weights_only=True) + normalized = normalize_portable_state_document(loaded) + assert normalized["completion"] == document["completion"] + + +def test_portable_digest_is_deterministic_and_type_shape_dtype_sensitive(): + left = {"b": [1, 2], "a": torch.tensor([1.0, 2.0])} + right = {"a": torch.tensor([1.0, 2.0]), "b": [1, 2]} + baseline = portable_state_digest(left) + + assert baseline == portable_state_digest(right) + assert baseline != portable_state_digest( + {"a": torch.tensor([[1.0, 2.0]]), "b": [1, 2]} + ) + assert baseline != portable_state_digest( + {"a": torch.tensor([1.0, 2.0], dtype=torch.float64), "b": [1, 2]} + ) + assert baseline != portable_state_digest( + {"a": torch.tensor([1.0, 3.0]), "b": [1, 2]} + ) + assert baseline != portable_state_digest( + {"a": torch.tensor([1.0, 2.0]), "b": (1, 2)} + ) + + +def test_portable_digest_streams_tensor_bytes_without_changing_the_digest(monkeypatch): + payload = {"tensor": torch.arange(257, dtype=torch.float32)} + baseline = portable_state_digest(payload) + + monkeypatch.setattr(portable_schema_module, "_PORTABLE_DIGEST_CHUNK_BYTES", 17) + + assert portable_state_digest(payload) == baseline + + +def test_portable_digest_v3_grammar_has_a_cross_version_golden_vector(): + payload = { + "a": None, + "b": True, + "c": -12345678901234567890, + "d": -0.0, + "e": "x\u2603", + "f": [1, (2, 3)], + "g": torch.tensor([[1, -2], [3, -4]], dtype=torch.int16), + "h": torch.tensor([1.5, -2.25], dtype=torch.bfloat16), + "i": torch.tensor(3 + 4j, dtype=torch.complex64), + "j": torch.empty((0, 2), dtype=torch.float64), + } + + assert portable_state_digest(payload) == ( + "348e09a77f1b3eae286adda1573722df938be6d8100631b272665f8c22c6e23a" + ) + + +def test_portable_clone_streams_noncontiguous_values_and_finite_checks(monkeypatch): + source = torch.arange(514, dtype=torch.float32)[1::2] + assert not source.is_contiguous() + calls = [] + original = portable_schema_module._read_portable_tensor_chunk + + def tracked(value, start, stop): + calls.append((start, stop)) + return original(value, start, stop) + + monkeypatch.setattr(portable_schema_module, "_PORTABLE_CLONE_CHUNK_BYTES", 17) + monkeypatch.setattr(portable_schema_module, "_read_portable_tensor_chunk", tracked) + + cloned = portable_schema_module._clone_portable_value(source, path="tensor") + + assert torch.equal(cloned, source) + assert cloned.is_contiguous() + assert cloned.storage_offset() == 0 + assert cloned.untyped_storage().nbytes() == cloned.numel() * cloned.element_size() + assert len(calls) > 1 + assert max(stop - start for start, stop in calls) <= 4 + + +def test_noncontiguous_chunk_elements_stay_within_the_clone_budget(): + budget = portable_schema_module._PORTABLE_CLONE_CHUNK_BYTES + # A high-rank non-contiguous view: the strided read path materializes an + # int64 linear index, its running quotient, and one int64 coordinate tensor + # per dimension. The chunk size must keep that scratch within the fixed clone + # budget rather than sizing on the element size alone (which would + # oversubscribe it and risk OOM on high-rank tensors). + base = torch.arange(2 * 3 * 4 * 5 * 6 * 7, dtype=torch.float32).reshape( + 2, 3, 4, 5, 6, 7 + ) + view = base.permute(5, 4, 3, 2, 1, 0) + assert not view.is_contiguous() + + chunk_elements = portable_schema_module._portable_tensor_chunk_elements(view) + scratch_per_element = view.element_size() + 8 * (view.ndim + 2) + assert chunk_elements >= 1 + assert chunk_elements * scratch_per_element <= budget + # Element-size-only sizing (the prior behavior) would blow past the budget. + assert (budget // view.element_size()) * scratch_per_element > budget + + contiguous = base.reshape(-1) + assert portable_schema_module._portable_tensor_chunk_elements(contiguous) == max( + 1, budget // contiguous.element_size() + ) + + # End-to-end the clone still reproduces a high-rank strided tensor exactly. + cloned = portable_schema_module._clone_portable_value(view, path="tensor") + assert torch.equal(cloned, view) + assert cloned.is_contiguous() + + +def test_builder_and_normalizer_each_hash_the_tensor_tree_once(monkeypatch): + calls = [] + original = portable_schema_module._canonical_portable_state_digest + + def tracked(value): + calls.append(value) + return original(value) + + monkeypatch.setattr( + portable_schema_module, "_canonical_portable_state_digest", tracked + ) + + document = _document() + assert len(calls) == 1 + normalize_portable_state_document(document) + assert len(calls) == 2 + + +def test_multibyte_tensor_digest_is_canonical_little_endian(monkeypatch): + little_endian_values = torch.tensor([0x0102, 0x0304], dtype=torch.int16) + simulated_big_endian_storage = torch.tensor([0x0201, 0x0403], dtype=torch.int16) + baseline = portable_state_digest(little_endian_values) + + monkeypatch.setattr(portable_schema_module, "_PORTABLE_NATIVE_BYTEORDER", "big") + + assert portable_state_digest(simulated_big_endian_storage) == baseline + + +@pytest.mark.parametrize( + "corruption", + ( + "payload", + "digest", + "status", + "algorithm", + "missing_top", + "bool_version", + "coverage", + "implementation", + "fqn", + "shape_type", + "state_variant", + "state_type", + "provenance", + ), +) +def test_portable_schema_and_completion_corruption_are_rejected(corruption): + document = _document() + damaged = copy.deepcopy(document) + if corruption == "payload": + damaged["common"]["gefen_global_step"] = 4 + elif corruption == "digest": + damaged["completion"]["digest"] = "0" * 64 + elif corruption == "status": + damaged["completion"]["status"] = "preparing" + elif corruption == "algorithm": + damaged["completion"]["digest_algorithm"] = "md5" + elif corruption == "missing_top": + damaged.pop("policy") + elif corruption == "bool_version": + damaged["format_version"] = True + elif corruption == "coverage": + damaged["coverage"] = "local_optimizer_fragment" + elif corruption == "implementation": + damaged["implementation"] = "" + elif corruption == "fqn": + damaged["parameters"]["Model.Weight"]["identity"]["fqn"] = "Other.Weight" + elif corruption == "shape_type": + damaged["parameters"]["Model.Weight"]["identity"]["global_shape"] = (2, 3) + elif corruption == "state_variant": + damaged["parameters"]["Model.Weight"]["state_variant"] = "" + elif corruption == "state_type": + damaged["parameters"]["Model.Weight"]["state"] = [] + else: + damaged["provenance"] = [] + + with pytest.raises((TypeError, ValueError)): + normalize_portable_state_document(damaged) + + +def test_expected_implementation_is_checked_after_digest_validation(): + document = _document() + + with pytest.raises(ValueError, match="does not match the target"): + normalize_portable_state_document( + document, expected_implementation="gefen.GefenMuon" + ) + + document["completion"]["digest"] = "0" * 64 + with pytest.raises(ValueError, match="digest does not match"): + normalize_portable_state_document( + document, expected_implementation="gefen.GefenMuon" + ) + + +@pytest.mark.parametrize( + "bad_value", + ( + torch.tensor([float("nan")]), + torch.tensor([float("inf")]), + lambda: None, + ), +) +def test_builder_rejects_nonportable_values(bad_value): + record = _parameter_record() + record["state"]["bad"] = bad_value + + with pytest.raises((TypeError, ValueError)): + build_portable_state_document( + implementation="gefen.Gefen", + policy={}, + common={}, + parameters={"Model.Weight": record}, + ) + + +def test_builder_and_normalizer_do_not_alias_or_mutate_inputs(): + record = _parameter_record() + parameters = {"Model.Weight": record} + source_momentum = record["state"]["momentum"] + + document = build_portable_state_document( + implementation="gefen.Gefen", + policy={}, + common={}, + parameters=parameters, + ) + normalized = normalize_portable_state_document(document) + + assert record["state"]["momentum"] is source_momentum + assert document is not normalized + assert document["parameters"] is not parameters + assert document["parameters"]["Model.Weight"] is not record + assert ( + document["parameters"]["Model.Weight"]["state"]["momentum"] + is not source_momentum + ) + assert ( + normalized["parameters"]["Model.Weight"]["state"]["momentum"] + is not document["parameters"]["Model.Weight"]["state"]["momentum"] + ) + + +def test_portable_schema_and_checkpoint_binding_exports_are_public(): + assert gefen.PORTABLE_STATE_FORMAT_VERSION == PORTABLE_STATE_FORMAT_VERSION + assert gefen.build_portable_state_document is build_portable_state_document + assert gefen.normalize_portable_state_document is normalize_portable_state_document + assert gefen.portable_state_digest is portable_state_digest + assert gefen.CheckpointProcessGroupBinding.__module__ == "gefen.checkpoint" + assert gefen.PortableStateLimits.__module__ == "gefen.portable_state" + assert gefen.PortableStateProvider.__module__ == "gefen.contracts" + assert "PortableStateProvider" in contracts_module.__all__ + assert gefen.load_portable_dcp.__module__ == "gefen.portable_dcp" + assert gefen.save_portable_dcp.__module__ == "gefen.portable_dcp" + assert gefen.LogicalRegion.__module__ == "gefen.contracts" diff --git a/tests/test_portable_state.py b/tests/test_portable_state.py new file mode 100644 index 0000000..1894cd5 --- /dev/null +++ b/tests/test_portable_state.py @@ -0,0 +1,874 @@ +"""Warning-strict CPU coverage for portable Gefen optimizer-state semantics.""" + +import copy +import dataclasses +import math + +import pytest +import torch + +from gefen.contracts import ( + LogicalSlice, + ParameterIdentity, + ParameterLayout, + PlacementKind, + ProcessGroupIdentity, + ShardIdentity, + ShardPlacement, + ShardingManifest, +) +from gefen.portable import _decode_quantized_momentum +from gefen.portable_fields import _project_dense_logical_field +from gefen.portable_identity import ( + _serialize_parameter_identity, + _serialize_shard_identity, +) +from gefen.portable_schema import build_portable_state_document +from gefen.portable_state import ( + PortableStateLimits, + _assemble_portable_state_fragments, + _build_portable_state_fragment, + _normalize_gefen_portable_state_document, + _normalize_portable_state_fragment, + _project_portable_parameter_state, +) + + +def _limits(**changes): + limits = PortableStateLimits( + max_fragment_tensor_bytes=4 << 20, + max_collective_tensor_bytes=16 << 20, + max_collective_metadata_bytes=64 << 20, + ) + return dataclasses.replace(limits, **changes) + + +def _codebook(): + return torch.linspace(-1.0, 1.0, 256, dtype=torch.float32) + + +def _policy(*, factored=False): + return { + "schema_version": 1, + "factored_v_2d": factored, + "force_1d_period_one": False, + "force_2d_period_one": False, + "period_one_substrings": [], + "stochastic_round": False, + "codebook_refresh_every": 0, + "momentum_projection": "dense_fp32_target_period_one_v1", + "second_moment_projection": "exact_representation_target_period_one_v1", + } + + +def _plain_options(*, second="block"): + return { + "lr": 0.001, + "beta1": 0.9, + "beta2": 0.999, + "eps": 1e-8, + "weight_decay": 0.01, + "second_moment_policy": second, + } + + +def _muon_options(*, normuon=False, mode="distributed"): + return { + "lr": 0.01, + "weight_decay": 0.1, + "momentum": 0.95, + "nesterov": True, + "ns_schedule": [[3.4445, -4.775, 2.0315]], + "ns_eps": 1e-7, + "adjust_lr_fn": "match_rms_adamw", + "sharded_mode": mode, + "fp8_ns": False, + "fp8_ns_compile": True, + "batched_ns": False, + "batched_ns_workspace_bytes": 1 << 20, + "normuon": normuon, + "normuon_beta2": 0.95, + "normuon_eps": 1e-8, + "cautious": False, + } + + +def _group(members=("rank:0", "rank:1")): + return ProcessGroupIdentity("checkpoint", members) + + +def _placement(group, member, kind): + return ShardPlacement( + "checkpoint", + kind, + group.ordered_members.index(member), + len(group.ordered_members), + ) + + +def _replicas(parameter, group): + return tuple( + ShardIdentity( + parameter, + ParameterLayout.REPLICATED, + LogicalSlice.full(parameter), + placements=(_placement(group, member, PlacementKind.REPLICATE),), + process_group=group, + local_member=member, + ) + for member in group.ordered_members + ) + + +def _flat(parameter, group, lengths): + offset = 0 + shards = [] + for member, length in zip(group.ordered_members, lengths): + shards.append( + ShardIdentity( + parameter, + ParameterLayout.FLATTENED_ELEMENT_SHARD, + LogicalSlice(offset, length), + placements=(_placement(group, member, PlacementKind.FLAT_SHARD),), + process_group=group, + local_member=member, + ) + ) + offset += length + return tuple(shards) + + +def _owners(parameter, group, owner): + return tuple( + ShardIdentity( + parameter, + ParameterLayout.WHOLE_PARAMETER_OWNER, + LogicalSlice.full(parameter) if member == owner else LogicalSlice(0, 0), + placements=(_placement(group, member, PlacementKind.WHOLE_PARAMETER_OWNER),), + process_group=group, + local_member=member, + owner=owner, + ) + for member in group.ordered_members + ) + + +def _role(shard): + if shard.layout is ParameterLayout.WHOLE_PARAMETER_OWNER: + return "whole_owner" if shard.local_member == shard.owner else "whole_nonowner" + if shard.layout is ParameterLayout.FLATTENED_ELEMENT_SHARD and shard.logical_slice.length == 0: + return "empty_flat" + if shard.parameter.numel == 0: + return "empty_replicated" + return "live" + + +def _fragments( + implementation, + parameter, + shards, + *, + options, + policy, + variant, + momentum=None, + second=None, + v_row=None, + v_col=None, + normuon_v=None, +): + manifest = ShardingManifest(shards) + common = { + "gefen_global_step": 4 if variant != "pristine" else 0, + "gefen_codebook": _codebook() if variant != "pristine" else None, + "gefen_deterministic": True, + } + catalog = { + parameter.fqn: { + "identity": _serialize_parameter_identity(parameter), + "algorithm_options": options, + } + } + result = [] + for shard in shards: + payload = _role(shard) in {"live", "whole_owner"} and parameter.numel > 0 + local_variant = variant if payload else "pristine" + state = {} + source_second = None + period = None + if local_variant == "period_selected": + period = 1 + elif local_variant.startswith("initialized_"): + period = 1 + state = { + "step": 4, + "momentum": _project_dense_logical_field(parameter, momentum, shard), + } + if implementation == "gefen.Gefen" and local_variant == "initialized_dense": + source_second = "block" + state.update( + { + "second_moment": _project_dense_logical_field(parameter, second, shard), + "second_moment_step": 3, + } + ) + elif local_variant == "initialized_factored": + source_second = "factored" + state.update({"v_row": v_row.clone(), "v_col": v_col.clone(), "factored_step": 3}) + elif local_variant == "initialized_dense_normuon": + state.update({"normuon_v": normuon_v.clone(), "normuon_step": 2}) + slot = { + "group_index": 0, + "original_slot_index": 0, + "compatibility_name": "weight", + "shard": _serialize_shard_identity(shard), + "algorithm_options": options, + "role": _role(shard), + "source_period": period, + "source_second_moment": source_second, + "state_variant": local_variant, + "state": state, + } + result.append( + _build_portable_state_fragment( + implementation=implementation, + member=shard.local_member, + policy=policy, + common=common, + manifest=manifest, + catalog=catalog, + logical_slots=[slot], + limits=_limits(), + ) + ) + return result + + +@pytest.mark.parametrize("variant", ["pristine", "period_selected"]) +def test_plain_pristine_and_period_selected_round_trip(variant): + parameter = ParameterIdentity("layer.weight", (2, 3)) + group = _group() + fragments = _fragments( + "gefen.Gefen", + parameter, + _replicas(parameter, group), + options=_plain_options(), + policy=_policy(), + variant=variant, + ) + document = _assemble_portable_state_fragments(fragments, process_group_identity=group, limits=_limits()) + record = document["parameters"][parameter.fqn] + assert record["state_variant"] == variant + assert record["state"] == {} + assert record["projection_hints"]["source_periods"] == ([] if variant == "pristine" else [1]) + normalized = _normalize_gefen_portable_state_document(document, limits=_limits()) + assert normalized["completion"] == document["completion"] + + +@pytest.mark.parametrize("layout", ["replicated", "flat"]) +def test_plain_initialized_block_assembles_and_projects_exact_period_one(layout): + parameter = ParameterIdentity("layer.weight", (2, 3)) + group = _group() + shards = _replicas(parameter, group) if layout == "replicated" else _flat(parameter, group, (2, 4)) + momentum = torch.tensor([[-0.0, 0.0, float.fromhex("0x1p-149")], [1.25, -3.5, torch.finfo(torch.float32).max]]) + second = torch.arange(1, 7, dtype=torch.float32).reshape(2, 3) + fragments = _fragments( + "gefen.Gefen", + parameter, + shards, + options=_plain_options(), + policy=_policy(), + variant="initialized_dense", + momentum=momentum, + second=second, + ) + document = _assemble_portable_state_fragments(fragments, process_group_identity=group, limits=_limits()) + record = document["parameters"][parameter.fqn] + assert torch.equal(record["state"]["second_moment"], second) + assert torch.equal(record["state"]["momentum"].view(torch.int32), momentum.view(torch.int32)) + + for target in _flat(parameter, group, (3, 3)): + projected = _project_portable_parameter_state( + record, + target, + implementation="gefen.Gefen", + global_step=4, + codebook=document["common"]["gefen_codebook"], + target_algorithm_options=_plain_options(), + target_second_moment="block", + ) + local_momentum = _decode_quantized_momentum( + document["common"]["gefen_codebook"], + projected["m_codebook"], + projected["m_magnitude"], + logical_shape=(target.logical_slice.length,), + period=1, + step=4, + ) + expected = _project_dense_logical_field(parameter, momentum, target) + assert torch.equal(local_momentum.view(torch.int32), expected.view(torch.int32)) + assert projected["automatic_period"] == 1 + assert torch.equal(projected["vmean"].reshape(-1), _project_dense_logical_field(parameter, second, target)) + + +def test_empty_flat_fragment_is_pristine_and_ownership_is_enforced(): + parameter = ParameterIdentity("layer.weight", (2,)) + group = _group(("rank:0", "rank:1", "rank:2")) + shards = _flat(parameter, group, (1, 0, 1)) + fragments = _fragments( + "gefen.Gefen", + parameter, + shards, + options=_plain_options(), + policy=_policy(), + variant="initialized_dense", + momentum=torch.tensor([-0.0, 2.0]), + second=torch.tensor([1.0, 2.0]), + ) + assert fragments[1]["logical_slots"][0]["role"] == "empty_flat" + assert fragments[1]["logical_slots"][0]["state_variant"] == "pristine" + corrupted = copy.deepcopy(fragments[1]) + corrupted["logical_slots"][0]["role"] = "live" + with pytest.raises(ValueError, match="role"): + _normalize_portable_state_fragment(corrupted, limits=_limits()) + + +def test_factored_replicas_preserve_authoritative_factor_bytes_and_reject_migration(): + parameter = ParameterIdentity("matrix.weight", (2, 3)) + group = _group() + shards = _replicas(parameter, group) + momentum = torch.arange(6, dtype=torch.float32).reshape(2, 3) + v_row = torch.tensor([-0.0, 4.0]) + v_col = torch.tensor([1.0, 2.0, 3.0]) + fragments = _fragments( + "gefen.Gefen", + parameter, + shards, + options=_plain_options(second="factored"), + policy=_policy(factored=True), + variant="initialized_factored", + momentum=momentum, + v_row=v_row, + v_col=v_col, + ) + document = _assemble_portable_state_fragments(fragments, process_group_identity=group, limits=_limits()) + record = document["parameters"][parameter.fqn] + assert torch.equal(record["state"]["v_row"].view(torch.int32), v_row.view(torch.int32)) + projected = _project_portable_parameter_state( + record, + shards[0], + implementation="gefen.Gefen", + global_step=4, + codebook=_codebook(), + target_algorithm_options=_plain_options(second="factored"), + target_second_moment="factored", + ) + assert torch.equal(projected["v_row"].view(torch.int32), v_row.view(torch.int32)) + with pytest.raises(ValueError, match="target"): + _project_portable_parameter_state( + record, + shards[0], + implementation="gefen.Gefen", + global_step=4, + codebook=_codebook(), + target_algorithm_options=_plain_options(second="factored"), + target_second_moment="block", + ) + fragments[1]["logical_slots"][0]["state"]["v_row"][0] = 0.0 + with pytest.raises(ValueError, match="v_row disagree"): + _assemble_portable_state_fragments(fragments, process_group_identity=group, limits=_limits()) + + +def test_period_selected_projection_binds_target_second_moment_policy(): + parameter = ParameterIdentity("matrix.weight", (2, 3)) + group = _group() + shards = _replicas(parameter, group) + document = _assemble_portable_state_fragments( + _fragments( + "gefen.Gefen", + parameter, + shards, + options=_plain_options(second="factored"), + policy=_policy(factored=True), + variant="period_selected", + ), + process_group_identity=group, + limits=_limits(), + ) + with pytest.raises(ValueError, match="conflicts with target algorithm options"): + _project_portable_parameter_state( + document["parameters"][parameter.fqn], + shards[0], + implementation="gefen.Gefen", + global_step=0, + codebook=document["common"]["gefen_codebook"], + target_algorithm_options=_plain_options(second="factored"), + target_second_moment="block", + ) + + +def test_factored_matrix_rejects_flat_fragments_before_initialization(): + parameter = ParameterIdentity("matrix.weight", (2, 3)) + group = _group() + with pytest.raises(ValueError, match="factored logical matrices require replicated"): + _fragments( + "gefen.Gefen", + parameter, + _flat(parameter, group, (2, 4)), + options=_plain_options(second="factored"), + policy=_policy(factored=True), + variant="period_selected", + ) + document = _assemble_portable_state_fragments( + _fragments( + "gefen.Gefen", + parameter, + _replicas(parameter, group), + options=_plain_options(second="factored"), + policy=_policy(factored=True), + variant="period_selected", + ), + process_group_identity=group, + limits=_limits(), + ) + with pytest.raises(ValueError, match="replicated target shards"): + _project_portable_parameter_state( + document["parameters"][parameter.fqn], + _flat(parameter, group, (2, 4))[0], + implementation="gefen.Gefen", + global_step=document["common"]["gefen_global_step"], + codebook=document["common"]["gefen_codebook"], + target_algorithm_options=_plain_options(second="factored"), + target_second_moment="factored", + ) + + +@pytest.mark.parametrize("normuon", [False, True]) +def test_muon_whole_owner_assembles_and_projects_after_owner_move(normuon): + parameter = ParameterIdentity("muon.weight", (2, 3)) + group = _group() + source = _owners(parameter, group, "rank:0") + momentum = torch.tensor([[-1.0, 0.0, 1.0], [2.0, 3.0, 4.0]]) + normuon_v = torch.tensor([[1.0], [2.0]]) + variant = "initialized_dense_normuon" if normuon else "initialized_dense" + fragments = _fragments( + "gefen.GefenMuon", + parameter, + source, + options=_muon_options(normuon=normuon), + policy=_policy(), + variant=variant, + momentum=momentum, + normuon_v=normuon_v, + ) + document = _assemble_portable_state_fragments(fragments, process_group_identity=group, limits=_limits()) + target = _owners(parameter, group, "rank:1") + nonowner = _project_portable_parameter_state( + document["parameters"][parameter.fqn], + target[0], + implementation="gefen.GefenMuon", + global_step=4, + codebook=_codebook(), + target_algorithm_options=_muon_options(normuon=normuon), + ) + owner = _project_portable_parameter_state( + document["parameters"][parameter.fqn], + target[1], + implementation="gefen.GefenMuon", + global_step=4, + codebook=_codebook(), + target_algorithm_options=_muon_options(normuon=normuon), + ) + assert nonowner == {} + assert owner["automatic_period"] == 1 + if normuon: + assert torch.equal(owner["normuon_v"], normuon_v) + + +def test_muon_whole_owner_projection_requires_distributed_mode(): + parameter = ParameterIdentity("muon.weight", (2, 3)) + group = _group() + document = _assemble_portable_state_fragments( + _fragments( + "gefen.GefenMuon", + parameter, + _replicas(parameter, group), + options=_muon_options(mode="exact"), + policy=_policy(), + variant="period_selected", + ), + process_group_identity=group, + limits=_limits(), + ) + target = _owners(parameter, group, "rank:0")[0] + with pytest.raises(ValueError, match="sharded_mode='distributed'"): + _project_portable_parameter_state( + document["parameters"][parameter.fqn], + target, + implementation="gefen.GefenMuon", + global_step=document["common"]["gefen_global_step"], + codebook=document["common"]["gefen_codebook"], + target_algorithm_options=_muon_options(mode="exact"), + ) + + +def test_strict_document_schema_completion_digest_counters_and_period_gate(): + parameter = ParameterIdentity("layer.weight", (2,)) + record = { + "identity": _serialize_parameter_identity(parameter), + "algorithm_options": _plain_options(), + "state_variant": "period_selected", + "state": {}, + "projection_hints": { + "source_periods": [1], + "source_second_moment": None, + "target_period": 1, + }, + } + document = build_portable_state_document( + implementation="gefen.Gefen", + policy=_policy(), + common={"gefen_global_step": 1, "gefen_codebook": _codebook(), "gefen_deterministic": False}, + parameters={parameter.fqn: record}, + provenance=None, + ) + _normalize_gefen_portable_state_document(document, limits=_limits()) + corrupt = copy.deepcopy(document) + corrupt["completion"]["digest"] = "0" * 64 + with pytest.raises(ValueError, match="digest"): + _normalize_gefen_portable_state_document(corrupt, limits=_limits()) + bad_period = copy.deepcopy(record) + bad_period["projection_hints"]["source_periods"] = [2] + with pytest.raises(ValueError, match="period one"): + _normalize_gefen_portable_state_document( + build_portable_state_document( + implementation="gefen.Gefen", + policy=_policy(), + common=document["common"], + parameters={parameter.fqn: bad_period}, + provenance=None, + ), + limits=_limits(), + ) + too_large = copy.deepcopy(record) + too_large["state_variant"] = "pristine" + too_large["projection_hints"]["source_periods"] = [] + with pytest.raises(ValueError, match="at most"): + _normalize_gefen_portable_state_document( + build_portable_state_document( + implementation="gefen.Gefen", + policy=_policy(), + common={ + "gefen_global_step": (1 << 53), + "gefen_codebook": None, + "gefen_deterministic": False, + }, + parameters={parameter.fqn: too_large}, + provenance=None, + ), + limits=_limits(), + ) + + +def test_global_step_zero_allows_period_selection_only_with_a_valid_codebook(): + parameter = ParameterIdentity("layer.weight", (2,)) + record = { + "identity": _serialize_parameter_identity(parameter), + "algorithm_options": _plain_options(), + "state_variant": "period_selected", + "state": {}, + "projection_hints": { + "source_periods": [1], + "source_second_moment": None, + "target_period": 1, + }, + } + + def document(codebook): + return build_portable_state_document( + implementation="gefen.Gefen", + policy=_policy(), + common={ + "gefen_global_step": 0, + "gefen_codebook": codebook, + "gefen_deterministic": False, + }, + parameters={parameter.fqn: record}, + provenance=None, + ) + + normalized = _normalize_gefen_portable_state_document(document(_codebook()), limits=_limits()) + assert normalized["common"]["gefen_global_step"] == 0 + with pytest.raises(ValueError, match="requires a codebook"): + _normalize_gefen_portable_state_document(document(None), limits=_limits()) + + +def test_fragment_rejects_ungrouped_shards_and_empty_period_projection_stays_empty(): + parameter = ParameterIdentity("layer.weight", (2,)) + ungrouped = ShardIdentity( + parameter, + ParameterLayout.REPLICATED, + LogicalSlice.full(parameter), + ) + catalog = { + parameter.fqn: { + "identity": _serialize_parameter_identity(parameter), + "algorithm_options": _plain_options(), + } + } + with pytest.raises(ValueError, match="ungrouped"): + _build_portable_state_fragment( + implementation="gefen.Gefen", + member="rank:0", + policy=_policy(), + common={ + "gefen_global_step": 0, + "gefen_codebook": None, + "gefen_deterministic": False, + }, + manifest=ShardingManifest((ungrouped,)), + catalog=catalog, + logical_slots=[ + { + "group_index": 0, + "original_slot_index": 0, + "compatibility_name": "weight", + "shard": _serialize_shard_identity(ungrouped), + "algorithm_options": _plain_options(), + "role": "live", + "source_period": None, + "source_second_moment": None, + "state_variant": "pristine", + "state": {}, + } + ], + limits=_limits(), + ) + + group = _group(("rank:0", "rank:1", "rank:2")) + source = _replicas(parameter, group) + document = _assemble_portable_state_fragments( + _fragments( + "gefen.Gefen", + parameter, + source, + options=_plain_options(), + policy=_policy(), + variant="period_selected", + ), + process_group_identity=group, + limits=_limits(), + ) + empty_target = _flat(parameter, group, (1, 0, 1))[1] + assert ( + _project_portable_parameter_state( + document["parameters"][parameter.fqn], + empty_target, + implementation="gefen.Gefen", + global_step=4, + codebook=document["common"]["gefen_codebook"], + target_algorithm_options=_plain_options(), + target_second_moment="block", + ) + == {} + ) + + +@pytest.mark.parametrize("implementation", ["gefen.Gefen", "gefen.GefenMuon"]) +def test_complete_zero_numel_parameter_must_remain_pristine(implementation): + parameter = ParameterIdentity("empty.weight", (0, 3)) + options = _plain_options(second="factored") if implementation == "gefen.Gefen" else _muon_options() + policy = _policy(factored=implementation == "gefen.Gefen") + hints = {"source_periods": [1], "target_period": 1} + if implementation == "gefen.Gefen": + hints["source_second_moment"] = None + record = { + "identity": _serialize_parameter_identity(parameter), + "algorithm_options": options, + "state_variant": "period_selected", + "state": {}, + "projection_hints": hints, + } + document = build_portable_state_document( + implementation=implementation, + policy=policy, + common={ + "gefen_global_step": 0, + "gefen_codebook": _codebook(), + "gefen_deterministic": False, + }, + parameters={parameter.fqn: record}, + provenance=None, + ) + with pytest.raises(ValueError, match="empty logical parameters must remain pristine"): + _normalize_gefen_portable_state_document(document, limits=_limits()) + + +def test_assembly_authenticates_full_process_group_identity(): + parameter = ParameterIdentity("layer.weight", (2,)) + source_group = _group() + fragments = _fragments( + "gefen.Gefen", + parameter, + _replicas(parameter, source_group), + options=_plain_options(), + policy=_policy(), + variant="pristine", + ) + wrong_scope = ProcessGroupIdentity("different-checkpoint", source_group.ordered_members) + with pytest.raises(ValueError, match="exactly match process_group_identity"): + _assemble_portable_state_fragments( + fragments, + process_group_identity=wrong_scope, + limits=_limits(), + ) + + +def test_muon_schedule_respects_native_length_limit(): + parameter = ParameterIdentity("muon.weight", (2, 2)) + options = _muon_options() + options["ns_schedule"] = [[1.0, 2.0, 3.0] for _ in range(100)] + record = { + "identity": _serialize_parameter_identity(parameter), + "algorithm_options": options, + "state_variant": "pristine", + "state": {}, + "projection_hints": {"source_periods": [], "target_period": 1}, + } + document = build_portable_state_document( + implementation="gefen.GefenMuon", + policy=_policy(), + common={ + "gefen_global_step": 0, + "gefen_codebook": None, + "gefen_deterministic": False, + }, + parameters={parameter.fqn: record}, + provenance=None, + ) + with pytest.raises(ValueError, match="between 1 and 99"): + _normalize_gefen_portable_state_document(document, limits=_limits()) + + +def test_limits_are_strict_and_bound_tensor_and_member_aggregation(): + with pytest.raises(TypeError): + PortableStateLimits(True, 2, 64 << 20) + with pytest.raises(ValueError): + PortableStateLimits(2, 1, 64 << 20) + parameter = ParameterIdentity("layer.weight", (2,)) + group = _group() + fragments = _fragments( + "gefen.Gefen", + parameter, + _replicas(parameter, group), + options=_plain_options(), + policy=_policy(), + variant="initialized_dense", + momentum=torch.ones(2), + second=torch.ones(2), + ) + with pytest.raises(ValueError, match="max_fragment_tensor_bytes"): + _normalize_portable_state_fragment( + fragments[0], + limits=PortableStateLimits(1, 64, 64 << 20), + ) + with pytest.raises(ValueError, match="max_members"): + _assemble_portable_state_fragments( + fragments, + process_group_identity=group, + limits=dataclasses.replace(_limits(), max_members=1), + ) + with pytest.raises(ValueError, match="max_collective_tensor_bytes"): + _assemble_portable_state_fragments( + fragments, + process_group_identity=group, + limits=PortableStateLimits( + max_fragment_tensor_bytes=2048, + max_collective_tensor_bytes=2048, + max_collective_metadata_bytes=64 << 20, + ), + ) + + +def test_signed_zero_is_bitwise_consensus_not_numeric_consensus(): + parameter = ParameterIdentity("layer.weight", (1,)) + group = _group() + fragments = _fragments( + "gefen.Gefen", + parameter, + _replicas(parameter, group), + options=_plain_options(), + policy=_policy(), + variant="initialized_dense", + momentum=torch.tensor([-0.0]), + second=torch.tensor([1.0]), + ) + fragments[1]["logical_slots"][0]["state"]["momentum"][0] = 0.0 + with pytest.raises(ValueError, match="replicas disagree"): + _assemble_portable_state_fragments(fragments, process_group_identity=group, limits=_limits()) + + +def test_logical_slot_positions_must_agree_across_members(): + first = ParameterIdentity("first.weight", (1,)) + second = ParameterIdentity("second.weight", (1,)) + group = _group() + shards = _replicas(first, group) + _replicas(second, group) + manifest = ShardingManifest(shards) + common = { + "gefen_global_step": 0, + "gefen_codebook": None, + "gefen_deterministic": False, + } + catalog = { + parameter.fqn: { + "identity": _serialize_parameter_identity(parameter), + "algorithm_options": _plain_options(), + } + for parameter in (first, second) + } + fragments = [] + for member_index, member in enumerate(group.ordered_members): + member_shards = [shard for shard in shards if shard.local_member == member] + if member_index: + member_shards.reverse() + logical_slots = [] + for slot_index, shard in enumerate(member_shards): + logical_slots.append( + { + "group_index": 0, + "original_slot_index": slot_index, + "compatibility_name": shard.parameter.fqn, + "shard": _serialize_shard_identity(shard), + "algorithm_options": _plain_options(), + "role": "live", + "source_period": None, + "source_second_moment": None, + "state_variant": "pristine", + "state": {}, + } + ) + fragments.append( + _build_portable_state_fragment( + implementation="gefen.Gefen", + member=member, + policy=_policy(), + common=common, + manifest=manifest, + catalog=catalog, + logical_slots=logical_slots, + limits=_limits(), + ) + ) + with pytest.raises(ValueError, match="logical slot position disagree"): + _assemble_portable_state_fragments( + fragments, + process_group_identity=group, + limits=_limits(), + ) + + +def test_public_limits_are_frozen_and_only_public_export(): + limits = _limits() + with pytest.raises(dataclasses.FrozenInstanceError): + limits.chunk_bytes = 1 + assert math.isfinite(float(limits.max_collective_tensor_bytes)) diff --git a/tests/test_portable_state_math.py b/tests/test_portable_state_math.py new file mode 100644 index 0000000..b5bfdea --- /dev/null +++ b/tests/test_portable_state_math.py @@ -0,0 +1,751 @@ +import warnings + +import pytest +import torch + +import gefen.portable as portable_math +from gefen.portable import ( + _decode_quantized_momentum, + _expand_block_second_moment, + _expand_factored_second_moment, + _project_factored_second_moment, + _recompress_dense_momentum, + _reduce_block_second_moment, + _validate_state_counter, +) + + +def _assert_tight(tensor, *, shape, dtype=torch.float32): + assert type(tensor) is torch.Tensor + assert tuple(tensor.shape) == tuple(shape) + assert tensor.dtype == dtype + assert tensor.layout is torch.strided + assert not tensor.requires_grad + assert tensor.grad_fn is None + assert tensor.is_contiguous() + assert tensor.storage_offset() == 0 + assert tensor.untyped_storage().nbytes() == tensor.numel() * tensor.element_size() + assert bool(torch.isfinite(tensor).all()) + + +def _codebook(): + return torch.tensor([-1.0, -0.5, 0.5, 1.0], dtype=torch.float32) + + +def _sparse_tensor(): + with warnings.catch_warnings(): + warnings.simplefilter("ignore", UserWarning) + return torch.sparse_coo_tensor( + torch.tensor([[0]]), + torch.tensor([1.0]), + (4,), + ) + + +def _reference_nearest_indices(codebook, values): + codebook = codebook.contiguous() + flat = values.reshape(-1).float() + insertion = torch.searchsorted(codebook, flat) + left = (insertion - 1).clamp(0, codebook.numel() - 1) + right = insertion.clamp(0, codebook.numel() - 1) + return torch.where( + (flat - codebook[left]).abs() <= (flat - codebook[right]).abs(), + left, + right, + ).to(torch.uint8).reshape(values.shape) + + +def _reference_recompress(momentum, codebook, period): + blocks = momentum.reshape(-1, period) + magnitudes = blocks.abs().amax(dim=1, keepdim=True) + normalized = blocks.clone() + nonzero = magnitudes > 0 + normalized.div_(magnitudes) + normalized.masked_fill_(~nonzero, 0.0) + return _reference_nearest_indices(codebook, normalized), magnitudes + + +def test_decode_quantized_momentum_validates_geometry_and_returns_dense_tight_fp32(): + codebook = torch.tensor([-1.0, -0.25, 0.25, 1.0], dtype=torch.float32) + indices = torch.tensor([[0, 2], [3, 1]], dtype=torch.uint8) + magnitudes = torch.tensor([[2.0], [4.0]], dtype=torch.float32) + codebook_before = codebook.clone() + indices_before = indices.clone() + magnitudes_before = magnitudes.clone() + + dense = _decode_quantized_momentum( + codebook, + indices, + magnitudes, + logical_shape=(2, 2), + period=2, + step=7, + ) + + _assert_tight(dense, shape=(2, 2)) + assert torch.equal(dense, torch.tensor([[-2.0, 0.5], [4.0, -1.0]])) + assert torch.equal(codebook, codebook_before) + assert torch.equal(indices, indices_before) + assert torch.equal(magnitudes, magnitudes_before) + assert dense.untyped_storage().data_ptr() != magnitudes.untyped_storage().data_ptr() + + +def test_recompression_uses_exact_lower_index_tie_semantics(): + momentum = torch.tensor([[-2.0, -1.5], [2.0, 0.0]], requires_grad=True) + indices, magnitudes = _recompress_dense_momentum( + momentum, + _codebook(), + period=2, + step=1, + ) + + _assert_tight(indices, shape=(2, 2), dtype=torch.uint8) + _assert_tight(magnitudes, shape=(2, 1)) + assert torch.equal(magnitudes, torch.tensor([[2.0], [2.0]])) + # -0.75 ties indices 0/1 and 0.0 ties indices 1/2. Both choose the lower index. + assert torch.equal(indices, torch.tensor([[0, 0], [3, 1]], dtype=torch.uint8)) + assert torch.equal(momentum.detach(), torch.tensor([[-2.0, -1.5], [2.0, 0.0]])) + + +def test_period_one_recompression_is_exact_for_every_finite_fp32_scale(): + tiny = torch.nextafter(torch.tensor(0.0), torch.tensor(1.0)) + maximum = torch.tensor(torch.finfo(torch.float32).max) + momentum = torch.stack( + ( + -maximum, + torch.tensor(-123.75), + -tiny, + torch.tensor(-0.0), + torch.tensor(0.0), + tiny, + torch.tensor(19.125), + maximum, + ) + ) + codebook = torch.tensor( + [-1.0, -1.0, -0.2, 0.3, 1.0, 1.0], + dtype=torch.float32, + ) + + indices, magnitudes = _recompress_dense_momentum( + momentum, + codebook, + period=1, + step=99, + ) + reconstructed = _decode_quantized_momentum( + codebook, + indices, + magnitudes, + logical_shape=momentum.shape, + period=1, + step=99, + ) + + assert torch.equal(reconstructed.view(torch.int32), momentum.view(torch.int32)) + assert torch.equal(magnitudes.reshape(-1), momentum.abs()) + assert torch.equal(indices[momentum < 0], torch.zeros_like(indices[momentum < 0])) + assert torch.equal(indices[momentum > 0], torch.full_like(indices[momentum > 0], 4)) + assert indices[3].item() == 0 + assert indices[4].item() == codebook.numel() - 1 + + +def test_period_one_block_reduction_preserves_every_fp32_bit(): + tiny = torch.nextafter(torch.tensor(0.0), torch.tensor(1.0)) + dense = torch.stack( + ( + torch.tensor(-0.0), + torch.tensor(0.0), + tiny, + torch.tensor(torch.finfo(torch.float32).tiny), + torch.tensor(19.125), + torch.tensor(torch.finfo(torch.float32).max), + ) + ) + + reduced = _reduce_block_second_moment(dense, period=1, step=99) + + assert torch.equal( + reduced.reshape(-1).view(torch.int32), + dense.view(torch.int32), + ) + + +def test_scalar_momentum_roundtrips_as_zero_dimensional_logical_state(): + momentum = torch.tensor(-123.75, requires_grad=True) + indices, magnitudes = _recompress_dense_momentum( + momentum, + _codebook(), + period=1, + step=2, + ) + decoded = _decode_quantized_momentum( + _codebook(), + indices, + magnitudes, + logical_shape=(), + period=1, + step=2, + ) + + _assert_tight(indices, shape=(1, 1), dtype=torch.uint8) + _assert_tight(magnitudes, shape=(1, 1)) + _assert_tight(decoded, shape=()) + assert decoded.ndim == 0 + assert torch.equal(decoded, momentum.detach()) + + +def test_expanded_zero_stride_inputs_preserve_logical_row_major_values(): + momentum = torch.tensor([-2.0, 0.0, 3.0]).reshape(3, 1).expand(3, 4) + assert momentum.stride() == (1, 0) + codebook = torch.tensor([-1.0, 0.0, 1.0]) + indices, magnitudes = _recompress_dense_momentum( + momentum, + codebook, + period=4, + step=1, + ) + decoded = _decode_quantized_momentum( + codebook, + indices, + magnitudes, + logical_shape=momentum.shape, + period=4, + step=1, + ) + assert torch.equal(decoded, momentum) + + dense_second_moment = torch.tensor([1.0, 3.0, 5.0]).reshape(3, 1).expand(3, 4) + reduced = _reduce_block_second_moment( + dense_second_moment, + period=4, + step=1, + ) + assert torch.equal(reduced, torch.tensor([[1.0], [3.0], [5.0]])) + expanded_blocks = _expand_block_second_moment( + torch.tensor([[2.0]]).expand(3, 1), + logical_shape=(3, 4), + period=4, + step=1, + ) + assert torch.equal(expanded_blocks, torch.full((3, 4), 2.0)) + + factored = _expand_factored_second_moment( + torch.tensor([2.0]).expand(3), + torch.tensor([4.0]).expand(5), + logical_shape=(3, 5), + step=1, + ) + assert torch.equal(factored, torch.full((3, 5), 4.0)) + projected_row, projected_column = _project_factored_second_moment( + torch.tensor([7.0]).expand(4, 6), + step=1, + ) + assert torch.equal(projected_row, torch.full((4,), 7.0)) + assert torch.equal(projected_column, torch.full((6,), 7.0)) + + +def test_noncontiguous_codebook_is_normalized_without_searchsorted_warning(): + codebook = torch.linspace(-1.0, 1.0, 17)[::2] + assert not codebook.is_contiguous() + values = torch.tensor([-1.0, -0.875, -0.75, 0.0, 0.875, 1.0]) + assert torch.equal( + portable_math._nearest_codebook_indices(codebook, values), + _reference_nearest_indices(codebook, values), + ) + indices, magnitudes = _recompress_dense_momentum( + values, + codebook, + period=2, + step=1, + ) + assert torch.equal( + _decode_quantized_momentum( + codebook, + indices, + magnitudes, + logical_shape=values.shape, + period=2, + step=1, + ), + (codebook[indices.long()] * magnitudes).reshape(values.shape), + ) + + +@pytest.mark.parametrize("period", [1, 2, 4, 8]) +def test_recompression_supports_every_divisor_as_a_target_period(period): + momentum = torch.tensor( + [[-8.0, -5.0, -3.0, -1.0], [0.0, 1.0, 2.0, 8.0]], + dtype=torch.float32, + ).t() + assert not momentum.is_contiguous() + codebook = torch.linspace(-1.0, 1.0, 17) + + indices, magnitudes = _recompress_dense_momentum( + momentum, + codebook, + period=period, + step=2, + ) + blocks = momentum.numel() // period + expected_magnitudes = momentum.reshape(blocks, period).abs().amax(dim=1, keepdim=True) + assert torch.equal(magnitudes, expected_magnitudes) + _assert_tight(indices, shape=(blocks, period), dtype=torch.uint8) + _assert_tight(magnitudes, shape=(blocks, 1)) + + +def test_tiny_chunk_budget_matches_whole_tensor_references_across_partial_tails(monkeypatch): + monkeypatch.setattr(portable_math, "_PORTABLE_STATE_CHUNK_ELEMENTS", 7) + codebook = torch.tensor([-1.0, -0.5, 0.0, 0.5, 1.0]) + normalized = torch.tensor( + [ + [-1.0, -0.75, -0.5, -0.25, 0.0, 0.25], + [0.5, 0.75, 1.0, -0.75, 0.25, -0.25], + [0.1, -0.1, 0.9, -0.9, 0.6, -0.6], + [0.2, -0.2, 0.8, -0.8, 0.4, -0.4], + [0.3, -0.3, 0.7, -0.7, 1.0, -1.0], + ] + ).t() + assert not normalized.is_contiguous() + expected_nearest = _reference_nearest_indices(codebook, normalized) + actual_nearest = portable_math._nearest_codebook_indices(codebook, normalized) + assert torch.equal(actual_nearest, expected_nearest) + + momentum = torch.linspace(-9.0, 7.0, 30).reshape(6, 5).t() + assert not momentum.is_contiguous() + expected_indices, expected_magnitudes = _reference_recompress( + momentum, + codebook, + 6, + ) + actual_indices, actual_magnitudes = _recompress_dense_momentum( + momentum, + codebook, + period=6, + step=3, + ) + assert torch.equal(actual_indices, expected_indices) + assert torch.equal(actual_magnitudes, expected_magnitudes) + decoded = _decode_quantized_momentum( + codebook, + actual_indices, + actual_magnitudes, + logical_shape=momentum.shape, + period=6, + step=3, + ) + expected_decoded = ( + codebook[expected_indices.long()] * expected_magnitudes + ).reshape(momentum.shape) + assert torch.equal(decoded, expected_decoded) + + dense_second_moment = torch.arange(1, 31, dtype=torch.float32).reshape(6, 5).t() + expected_blocks = dense_second_moment.reshape(-1, 6).to(torch.float64).mean( + dim=1, + keepdim=True, + ).to(torch.float32) + actual_blocks = _reduce_block_second_moment( + dense_second_moment, + period=6, + step=3, + ) + assert torch.equal(actual_blocks, expected_blocks) + expanded_blocks = _expand_block_second_moment( + actual_blocks, + logical_shape=dense_second_moment.shape, + period=6, + step=3, + ) + expected_expanded_blocks = torch.repeat_interleave( + expected_blocks.reshape(-1), + 6, + ).reshape(dense_second_moment.shape) + assert torch.equal(expanded_blocks, expected_expanded_blocks) + + row = torch.tensor([1.0, 2.0, 4.0, 5.0, 8.0]) + column = torch.tensor([0.5, 1.0, 2.0, 3.0, 4.0, 5.5]) + expected_factored = ( + torch.outer(row.to(torch.float64), column.to(torch.float64)) + / row.to(torch.float64).mean() + ).to(torch.float32) + actual_factored = _expand_factored_second_moment( + row, + column, + logical_shape=(5, 6), + step=3, + ) + assert torch.equal(actual_factored, expected_factored) + actual_row, actual_column = _project_factored_second_moment( + actual_factored.t(), + step=3, + ) + expected64 = actual_factored.t().to(torch.float64) + assert torch.equal(actual_row, expected64.mean(dim=1).to(torch.float32)) + assert torch.equal(actual_column, expected64.mean(dim=0).to(torch.float32)) + + +def test_period_larger_than_chunk_budget_keeps_whole_blocks_and_matches_reference(monkeypatch): + monkeypatch.setattr(portable_math, "_PORTABLE_STATE_CHUNK_ELEMENTS", 5) + period = 11 + codebook = torch.linspace(-1.0, 1.0, 9) + momentum = torch.linspace(-17.0, 13.0, 33).reshape(11, 3).t() + assert not momentum.is_contiguous() + + expected_indices, expected_magnitudes = _reference_recompress( + momentum, + codebook, + period, + ) + actual_indices, actual_magnitudes = _recompress_dense_momentum( + momentum, + codebook, + period=period, + step=2, + ) + decoded = _decode_quantized_momentum( + codebook, + actual_indices, + actual_magnitudes, + logical_shape=momentum.shape, + period=period, + step=2, + ) + + assert torch.equal(actual_indices, expected_indices) + assert torch.equal(actual_magnitudes, expected_magnitudes) + assert torch.equal( + decoded, + (codebook[expected_indices.long()] * expected_magnitudes).reshape(momentum.shape), + ) + second_moment = momentum.square() + actual_blocks = _reduce_block_second_moment( + second_moment, + period=period, + step=2, + ) + expected_blocks = second_moment.reshape(-1, period).to(torch.float64).mean( + dim=1, + keepdim=True, + ).to(torch.float32) + assert torch.equal(actual_blocks, expected_blocks) + + +def test_block_second_moment_expands_then_projects_to_a_new_period(): + blocks = torch.tensor([[1.0], [3.0], [5.0], [7.0]], requires_grad=True) + dense = _expand_block_second_moment( + blocks, + logical_shape=(2, 4), + period=2, + step=3, + ) + projected = _reduce_block_second_moment(dense, period=4, step=3) + + _assert_tight(dense, shape=(2, 4)) + _assert_tight(projected, shape=(2, 1)) + assert torch.equal( + dense, + torch.tensor([[1.0, 1.0, 3.0, 3.0], [5.0, 5.0, 7.0, 7.0]]), + ) + assert torch.equal(projected, torch.tensor([[2.0], [6.0]])) + assert torch.equal(blocks.detach(), torch.tensor([[1.0], [3.0], [5.0], [7.0]])) + + +def test_block_reduction_uses_overflow_safe_means(): + maximum = torch.finfo(torch.float32).max + dense = torch.full((2, 4), maximum) + + projected = _reduce_block_second_moment(dense, period=4, step=1) + + assert torch.equal(projected, torch.full((2, 1), maximum)) + _assert_tight(projected, shape=(2, 1)) + + +def test_factored_expansion_and_projection_follow_adafactor_geometry(): + row = torch.tensor([1.0, 3.0]) + column = torch.tensor([2.0, 4.0, 6.0]) + dense = _expand_factored_second_moment( + row, + column, + logical_shape=(2, 3), + step=4, + ) + + _assert_tight(dense, shape=(2, 3)) + assert torch.equal(dense, torch.tensor([[1.0, 2.0, 3.0], [3.0, 6.0, 9.0]])) + + source = torch.tensor([[1.0, 3.0, 5.0], [2.0, 4.0, 6.0]]).t() + assert not source.is_contiguous() + projected_row, projected_column = _project_factored_second_moment(source, step=4) + assert torch.equal(projected_row, torch.tensor([1.5, 3.5, 5.5])) + assert torch.equal(projected_column, torch.tensor([3.0, 4.0])) + _assert_tight(projected_row, shape=(3,)) + _assert_tight(projected_column, shape=(2,)) + + +def test_factored_projection_is_overflow_safe_and_roundtrips_consistent_rank_one_state(): + maximum = torch.finfo(torch.float32).max + projected_row, projected_column = _project_factored_second_moment( + torch.full((2, 2), maximum), + step=1, + ) + assert torch.equal(projected_row, torch.full((2,), maximum)) + assert torch.equal(projected_column, torch.full((2,), maximum)) + + row = torch.tensor([1.0, 3.0]) + column = torch.tensor([1.0, 2.0, 3.0]) + dense = _expand_factored_second_moment( + row, + column, + logical_shape=(2, 3), + step=2, + ) + roundtrip_row, roundtrip_column = _project_factored_second_moment(dense, step=2) + torch.testing.assert_close(roundtrip_row, row, rtol=0, atol=0) + torch.testing.assert_close(roundtrip_column, column, rtol=0, atol=0) + + +@pytest.mark.parametrize("shape", [(0,), (0, 3), (2, 0)]) +def test_zero_element_momentum_and_block_state_are_well_defined(shape): + momentum = torch.empty(shape, dtype=torch.float32, requires_grad=True) + codebook = _codebook() + indices, magnitudes = _recompress_dense_momentum( + momentum, + codebook, + period=7, + step=1, + ) + decoded = _decode_quantized_momentum( + codebook, + indices, + magnitudes, + logical_shape=shape, + period=7, + step=1, + ) + reduced = _reduce_block_second_moment(momentum.detach(), period=7, step=1) + expanded = _expand_block_second_moment( + reduced, + logical_shape=shape, + period=7, + step=1, + ) + + _assert_tight(indices, shape=(0, 7), dtype=torch.uint8) + _assert_tight(magnitudes, shape=(0, 1)) + _assert_tight(decoded, shape=shape) + _assert_tight(reduced, shape=(0, 1)) + _assert_tight(expanded, shape=shape) + + +def test_empty_index_decode_does_not_invoke_max(monkeypatch): + original_max = torch.Tensor.max + + def reject_empty_max(tensor, *args, **kwargs): + if tensor.numel() == 0: + raise AssertionError("max() was called on empty momentum indices") + return original_max(tensor, *args, **kwargs) + + monkeypatch.setattr(torch.Tensor, "max", reject_empty_max) + decoded = _decode_quantized_momentum( + _codebook(), + torch.empty((0, 13), dtype=torch.uint8), + torch.empty((0, 1), dtype=torch.float32), + logical_shape=(0, 9), + period=13, + step=1, + ) + _assert_tight(decoded, shape=(0, 9)) + + +@pytest.mark.parametrize("shape", [(0, 3), (3, 0), (0, 0)]) +def test_zero_element_factored_state_uses_finite_empty_geometry(shape): + dense = torch.empty(shape, dtype=torch.float32) + row, column = _project_factored_second_moment(dense, step=1) + expanded = _expand_factored_second_moment( + row, + column, + logical_shape=shape, + step=1, + ) + + _assert_tight(row, shape=(shape[0],)) + _assert_tight(column, shape=(shape[1],)) + _assert_tight(expanded, shape=shape) + assert torch.count_nonzero(row) == 0 + assert torch.count_nonzero(column) == 0 + + +@pytest.mark.parametrize("value", [True, 1.0, torch.tensor(1), None]) +def test_counter_validation_rejects_non_host_ints(value): + with pytest.raises(TypeError, match="host int"): + _validate_state_counter(value, name="step", minimum=1) + + +def test_counter_validation_is_strict_about_bounds_and_its_own_arguments(): + assert _validate_state_counter(0, name="global_step") == 0 + assert _validate_state_counter(3, name="step", minimum=1) == 3 + with pytest.raises(ValueError, match="at least 1"): + _validate_state_counter(0, name="step", minimum=1) + with pytest.raises(TypeError, match="name"): + _validate_state_counter(1, name=object()) + with pytest.raises(ValueError, match="must not be empty"): + _validate_state_counter(1, name="") + with pytest.raises(TypeError, match="minimum"): + _validate_state_counter(1, name="step", minimum=True) + with pytest.raises(ValueError, match="nonnegative"): + _validate_state_counter(1, name="step", minimum=-1) + + +@pytest.mark.parametrize("period", [True, 2.0, 0, -1, 3]) +def test_invalid_periods_fail_before_projection(period): + exception = TypeError if type(period) is not int else ValueError + with pytest.raises(exception, match="period"): + _recompress_dense_momentum( + torch.ones(4), + _codebook(), + period=period, + step=1, + ) + + +@pytest.mark.parametrize( + "shape,exception", + [([4], TypeError), ((True,), TypeError), ((-1, 4), ValueError), ((3,), ValueError)], +) +def test_invalid_logical_shapes_and_geometry_are_rejected(shape, exception): + with pytest.raises(exception): + _decode_quantized_momentum( + _codebook(), + torch.zeros((2, 2), dtype=torch.uint8), + torch.ones((2, 1)), + logical_shape=shape, + period=2, + step=1, + ) + + +@pytest.mark.parametrize( + "codebook,match", + [ + (torch.tensor([], dtype=torch.float32), "between 1 and 256"), + (torch.zeros(257, dtype=torch.float32), "between 1 and 256"), + (torch.tensor([-1.0, 0.5, 0.25, 1.0]), "sorted"), + (torch.tensor([-1.1, 0.0, 1.0]), r"\[-1, 1\]"), + (torch.tensor([-1.0, float("nan"), 1.0]), "finite"), + ], +) +def test_invalid_codebook_values_are_rejected(codebook, match): + with pytest.raises(ValueError, match=match): + _recompress_dense_momentum( + torch.ones(4), + codebook, + period=2, + step=1, + ) + + +@pytest.mark.parametrize( + "codebook", + [torch.ones(4, dtype=torch.float64), torch.ones((2, 2), dtype=torch.float32)], +) +def test_invalid_codebook_dtype_or_rank_is_rejected(codebook): + with pytest.raises((TypeError, ValueError), match="codebook"): + _recompress_dense_momentum( + torch.ones(4), + codebook, + period=2, + step=1, + ) + + +@pytest.mark.parametrize( + "momentum", + [ + torch.ones(4, dtype=torch.float64), + torch.tensor([1.0, float("inf"), 2.0, 3.0]), + torch.nn.Parameter(torch.ones(4)), + _sparse_tensor(), + torch.empty(4, device="meta"), + ], +) +def test_momentum_input_must_be_plain_strided_finite_fp32(momentum): + with pytest.raises((TypeError, ValueError), match="momentum"): + _recompress_dense_momentum( + momentum, + _codebook(), + period=2, + step=1, + ) + + +@pytest.mark.parametrize( + "indices,magnitudes,match", + [ + (torch.zeros((2, 2), dtype=torch.int64), torch.ones((2, 1)), "indices"), + (torch.zeros((4,), dtype=torch.uint8), torch.ones((2, 1)), "indices"), + (torch.full((2, 2), 4, dtype=torch.uint8), torch.ones((2, 1)), "out-of-range"), + (torch.zeros((2, 2), dtype=torch.uint8), torch.ones((2,), dtype=torch.float32), "magnitudes"), + (torch.zeros((2, 2), dtype=torch.uint8), -torch.ones((2, 1)), "nonnegative"), + ( + torch.zeros((2, 2), dtype=torch.uint8), + torch.tensor([[1.0], [float("nan")]]), + "finite", + ), + ], +) +def test_decode_rejects_malformed_quantized_state(indices, magnitudes, match): + with pytest.raises((TypeError, ValueError), match=match): + _decode_quantized_momentum( + _codebook(), + indices, + magnitudes, + logical_shape=(4,), + period=2, + step=1, + ) + + +def test_second_moment_projection_rejects_negative_or_nonfinite_state(): + with pytest.raises(ValueError, match="nonnegative"): + _reduce_block_second_moment(torch.tensor([1.0, -1.0]), period=1, step=1) + with pytest.raises(ValueError, match="finite"): + _project_factored_second_moment( + torch.tensor([[1.0, float("nan")]]), + step=1, + ) + + +def test_factored_expansion_rejects_bad_geometry_and_degenerate_inconsistent_state(): + with pytest.raises(ValueError, match="2-D"): + _expand_factored_second_moment( + torch.ones(4), + torch.ones(1), + logical_shape=(4,), + step=1, + ) + with pytest.raises(ValueError, match="do not match"): + _expand_factored_second_moment( + torch.ones(3), + torch.ones(2), + logical_shape=(2, 2), + step=1, + ) + with pytest.raises(ValueError, match="zero row mean"): + _expand_factored_second_moment( + torch.zeros(2), + torch.ones(2), + logical_shape=(2, 2), + step=1, + ) + + +def test_factored_expansion_rejects_a_finite_input_whose_result_overflows_fp32(): + maximum = torch.finfo(torch.float32).max + tiny = torch.finfo(torch.float32).tiny + with pytest.raises(ValueError, match="cannot be represented"): + _expand_factored_second_moment( + torch.tensor([tiny, maximum]), + torch.tensor([maximum, maximum]), + logical_shape=(2, 2), + step=1, + ) diff --git a/tests/test_portable_wire.py b/tests/test_portable_wire.py new file mode 100644 index 0000000..74c693a --- /dev/null +++ b/tests/test_portable_wire.py @@ -0,0 +1,480 @@ +import dataclasses +import hashlib +import math +import struct + +import pytest +import torch + +import gefen.portable_wire as wire +from gefen.portable_wire import ( + _CanonicalWireLimits, + _CanonicalWirePlan, + _CanonicalWireTensorSpec, + _PreparedCanonicalWireValue, + _parse_canonical_wire_metadata, + _prepare_canonical_wire_value, + _reconstruct_canonical_wire_value, +) + + +def _limits(**changes): + limits = _CanonicalWireLimits( + max_fragment_tensor_bytes=1 << 20, + max_collective_tensor_bytes=8 << 20, + max_collective_metadata_bytes=1 << 20, + max_metadata_bytes=1 << 20, + ) + return dataclasses.replace(limits, **changes) + + +def _assert_value_equal(left, right): + if torch.is_tensor(left) or torch.is_tensor(right): + assert type(left) is torch.Tensor + assert type(right) is torch.Tensor + assert left.dtype == right.dtype + assert tuple(left.shape) == tuple(right.shape) + assert torch.equal(left, right) + return + assert type(left) is type(right) + if type(left) is dict: + assert list(left) == list(right) + for key in left: + _assert_value_equal(left[key], right[key]) + elif type(left) in {list, tuple}: + assert len(left) == len(right) + for left_item, right_item in zip(left, right): + _assert_value_equal(left_item, right_item) + elif type(left) is float and left == 0.0: + assert math.copysign(1.0, left) == math.copysign(1.0, right) + else: + assert left == right + + +def _round_trip(value, limits=None): + limits = _limits() if limits is None else limits + plan = _prepare_canonical_wire_value(value, limits) + prepared = _parse_canonical_wire_metadata(plan.metadata, limits=limits) + assert prepared.metadata_digest == plan.metadata_digest + assert prepared.fragment_digest == plan.fragment_digest + assert prepared.tensor_specs == plan.tensor_specs + result = _reconstruct_canonical_wire_value( + prepared, + plan.tensors, + expected_fragment_digest=plan.fragment_digest, + ) + _assert_value_equal(value, result) + return plan, prepared, result + + +def _metadata(root, descriptors=()): + return ( + b"GFNCV1\0\0" + + struct.pack(">HHI", 1, 0, 0) + + root + + struct.pack(">Q", len(descriptors)) + + b"".join(descriptors) + ) + + +def _string(value): + encoded = value.encode("utf-8") + return b"\x05" + struct.pack(">Q", len(encoded)) + encoded + + +def _descriptor(dtype_code=1, shape=(), numel=1, nbytes=1, digest=b"\0" * 32, flags=0): + return ( + struct.pack(">HHI", dtype_code, len(shape), flags) + + b"".join(struct.pack(">Q", dimension) for dimension in shape) + + struct.pack(">QQ", numel, nbytes) + + digest + ) + + +def test_golden_metadata_tensor_and_fragment_digests(): + value = {"a": -1, "b": torch.tensor([1, -2], dtype=torch.int16)} + plan = _prepare_canonical_wire_value(value, _limits()) + assert plan.metadata.hex() == ( + "47464e435631000000010000000000000800000000000000020500000000000000016103010000000000000001010500000000000000016209000000000000000000000000000000010004000100000000000000000000000200000000000000020000000000000004" + "356f76987e72498bcbc4866de355fdcbad68f9a1eb8807f1073728ee463f1e27" + ) + assert plan.metadata_digest.hex() == "c5894b6b9cc4f0f8591ed4fbb66d943b0bb69362fb460df413c42b36def889d8" + assert plan.fragment_digest.hex() == "155e9c7c022131d7ea971fd7bc1d69e13c7f9dcbaf89301bb29c135a24de4846" + _round_trip(value) + + +def test_round_trip_primitives_nested_containers_and_edge_values(): + value = { + "bools": [False, True], + "floats": (0.0, -0.0, float.fromhex("0x0.0000000000001p-1022"), 1.25), + "integers": [0, 1, -1, (1 << 4095) - 1], + "none": None, + "strings": ["", "ASCII", "λ🙂"], + } + _round_trip(value) + + +@pytest.mark.parametrize( + ("dtype", "values"), + [ + (torch.bool, [False, True, True]), + (torch.uint8, [0, 1, 255]), + (torch.int8, [-128, 0, 127]), + (torch.int16, [-32768, 0, 32767]), + (torch.int32, [-(1 << 31), 0, (1 << 31) - 1]), + (torch.int64, [-(1 << 63), 0, (1 << 63) - 1]), + (torch.float16, [-1.5, 0.0, 2.25]), + (torch.bfloat16, [-1.5, 0.0, 2.25]), + (torch.float32, [-1.5, 0.0, 2.25]), + (torch.float64, [-1.5, 0.0, 2.25]), + (torch.complex64, [complex(-1.5, 2.0), 0j, complex(2.25, -3.5)]), + (torch.complex128, [complex(-1.5, 2.0), 0j, complex(2.25, -3.5)]), + ], +) +def test_round_trip_fixed_dtype_table(dtype, values): + plan, _, result = _round_trip(torch.tensor(values, dtype=dtype)) + assert plan.tensor_specs[0].dtype_code == 1 + [record[1] for record in wire._DTYPE_RECORDS].index(dtype) + assert result.is_contiguous() + assert result.storage_offset() == 0 + assert result.untyped_storage().nbytes() == result.numel() * result.element_size() + + +def test_round_trip_noncontiguous_conjugate_scalar_and_empty_tensors(): + noncontiguous = torch.arange(30, dtype=torch.float64).reshape(5, 6).t()[1:5:2] + conjugate = torch.tensor([1 + 2j, 3 - 4j], dtype=torch.complex128).conj() + scalar = torch.tensor(-7, dtype=torch.int64) + empty = torch.empty((3, 0, 2), dtype=torch.float32) + plan, _, result = _round_trip([noncontiguous, conjugate, scalar, empty]) + assert not noncontiguous.is_contiguous() + assert conjugate.is_conj() + assert all(tensor.is_contiguous() for tensor in plan.tensors) + assert not result[1].is_conj() + + +def test_tensor_digest_uses_canonical_little_endian_component_bytes(): + value = torch.tensor([0x01020304], dtype=torch.int32) + plan = _prepare_canonical_wire_value(value, _limits()) + expected = hashlib.sha256() + expected.update(b"gefen.canonical_wire.tensor.v1\0") + expected.update(struct.pack(">HHQ", 5, 1, 1)) + expected.update(struct.pack(">Q", 4)) + expected.update(b"\x04\x03\x02\x01") + assert plan.tensor_specs[0].digest == expected.digest() + + +def test_fragment_digest_frames_metadata_and_ordered_tensor_digests(): + plan = _prepare_canonical_wire_value( + [torch.tensor([1], dtype=torch.int16), torch.tensor([2.0], dtype=torch.float64)], + _limits(), + ) + expected = hashlib.sha256() + expected.update(b"gefen.canonical_wire.fragment.v1\0") + expected.update(hashlib.sha256(plan.metadata).digest()) + expected.update(struct.pack(">Q", 2)) + for spec in plan.tensor_specs: + expected.update(struct.pack(">Q", spec.nbytes)) + expected.update(spec.digest) + assert plan.fragment_digest == expected.digest() + + +def test_dictionary_encoding_is_deterministic_and_sorted(): + left = _prepare_canonical_wire_value({"z": 1, "a": 2}, _limits()) + right = _prepare_canonical_wire_value({"a": 2, "z": 1}, _limits()) + assert left.metadata == right.metadata + assert left.fragment_digest == right.fragment_digest + + +def test_prepare_and_reconstruct_own_tensor_storage_and_containers(): + source = torch.arange(6, dtype=torch.float32).reshape(2, 3).t() + value = {"tensor": [source], "stable": 3} + plan = _prepare_canonical_wire_value(value, _limits(chunk_bytes=3)) + source.fill_(99) + value["tensor"].append("late") + prepared = _parse_canonical_wire_metadata(plan.metadata, limits=_limits(chunk_bytes=3)) + result = _reconstruct_canonical_wire_value(prepared, plan.tensors) + assert torch.equal(result["tensor"][0], torch.arange(6, dtype=torch.float32).reshape(2, 3).t()) + assert result["tensor"][0].data_ptr() != plan.tensors[0].data_ptr() + plan.tensors[0].fill_(-1) + assert not torch.equal(result["tensor"][0], plan.tensors[0]) + result["tensor"].append("owned") + assert len(value["tensor"]) == 2 + + +def test_tiny_chunk_budget_is_wire_identical_for_real_and_complex_tensors(): + value = { + "complex": torch.tensor([1 + 2j, -3 + 4j, 5 - 6j], dtype=torch.complex128), + "view": torch.arange(60, dtype=torch.float64).reshape(6, 10)[:, ::3], + } + normal = _prepare_canonical_wire_value(value, _limits()) + tiny_limits = _limits(chunk_bytes=1) + tiny = _prepare_canonical_wire_value(value, tiny_limits) + assert tiny.metadata == normal.metadata + assert tiny.metadata_digest == normal.metadata_digest + assert tiny.fragment_digest == normal.fragment_digest + for tiny_tensor, normal_tensor in zip(tiny.tensors, normal.tensors): + assert torch.equal(tiny_tensor, normal_tensor) + prepared = _parse_canonical_wire_metadata(tiny.metadata, limits=tiny_limits) + result = _reconstruct_canonical_wire_value(prepared, tiny.tensors) + _assert_value_equal(value, result) + + +def test_noncontiguous_clone_scales_coordinate_scratch_to_chunk_budget(monkeypatch): + value = torch.arange(2 * 3 * 4 * 5, dtype=torch.float32).reshape(2, 3, 4, 5) + value = value.permute(3, 2, 1, 0) + counts = [] + original = wire.torch.arange + + def tracked(start, stop, **kwargs): + counts.append(stop - start) + return original(start, stop, **kwargs) + + monkeypatch.setattr(wire.torch, "arange", tracked) + limits = _limits(chunk_bytes=100) + plan = _prepare_canonical_wire_value(value, limits) + + assert counts + assert max(counts) <= 100 // (value.element_size() + 8 * (value.ndim + 2)) + assert torch.equal(plan.tensors[0], value) + + +@pytest.mark.parametrize( + "metadata", + [ + b"", + b"BADMAGIC" + struct.pack(">HHI", 1, 0, 0) + b"\0" + struct.pack(">Q", 0), + b"GFNCV1\0\0" + struct.pack(">HHI", 2, 0, 0) + b"\0" + struct.pack(">Q", 0), + b"GFNCV1\0\0" + struct.pack(">HHI", 1, 1, 0) + b"\0" + struct.pack(">Q", 0), + b"GFNCV1\0\0" + struct.pack(">HHI", 1, 0, 1) + b"\0" + struct.pack(">Q", 0), + _metadata(b"\xff"), + _metadata(b"\0") + b"trailing", + ], +) +def test_parser_rejects_headers_unknown_tags_truncation_and_trailing_bytes(metadata): + with pytest.raises(ValueError): + _parse_canonical_wire_metadata(metadata, limits=_limits()) + + +@pytest.mark.parametrize( + "root", + [ + b"\x03\x02" + struct.pack(">Q", 0), + b"\x03\x00" + struct.pack(">Q", 1) + b"\0", + b"\x03\x01" + struct.pack(">Q", 0), + b"\x04" + struct.pack(">d", math.inf), + b"\x05" + struct.pack(">Q", 1) + b"\xff", + ], +) +def test_parser_rejects_noncanonical_scalars(root): + with pytest.raises(ValueError): + _parse_canonical_wire_metadata(_metadata(root), limits=_limits()) + + +def test_parser_rejects_non_string_unsorted_and_duplicate_dictionary_keys(): + roots = [ + b"\x08" + struct.pack(">Q", 1) + b"\0\0", + b"\x08" + struct.pack(">Q", 2) + _string("b") + b"\0" + _string("a") + b"\0", + b"\x08" + struct.pack(">Q", 2) + _string("a") + b"\0" + _string("a") + b"\0", + ] + for root in roots: + with pytest.raises(ValueError): + _parse_canonical_wire_metadata(_metadata(root), limits=_limits()) + + +def test_parser_rejects_tensor_reference_and_descriptor_count_mismatches(): + cases = [ + _metadata(b"\x09" + struct.pack(">Q", 1), [_descriptor()]), + _metadata(b"\x09" + struct.pack(">Q", 0)), + _metadata(b"\0", [_descriptor()]), + ] + for metadata in cases: + with pytest.raises(ValueError): + _parse_canonical_wire_metadata(metadata, limits=_limits()) + + +@pytest.mark.parametrize( + "descriptor", + [ + _descriptor(dtype_code=99), + _descriptor(flags=1), + _descriptor(dtype_code=5, shape=(2,), numel=3, nbytes=12), + _descriptor(dtype_code=5, shape=(2,), numel=2, nbytes=7), + ], +) +def test_parser_rejects_invalid_tensor_descriptors(descriptor): + metadata = _metadata(b"\x09" + struct.pack(">Q", 0), [descriptor]) + with pytest.raises(ValueError): + _parse_canonical_wire_metadata(metadata, limits=_limits()) + + +def test_parser_rejects_tensor_shape_product_overflow(): + descriptor = _descriptor(dtype_code=1, shape=((1 << 62), 4), numel=0, nbytes=0) + metadata = _metadata(b"\x09" + struct.pack(">Q", 0), [descriptor]) + with pytest.raises(ValueError, match="overflows"): + _parse_canonical_wire_metadata(metadata, limits=_limits()) + + +def test_parser_rejects_unrepresentable_dimension_even_when_product_is_zero(): + descriptor = _descriptor(dtype_code=1, shape=((1 << 63), 0), numel=0, nbytes=0) + metadata = _metadata(b"\x09" + struct.pack(">Q", 0), [descriptor]) + with pytest.raises(ValueError, match="signed int64"): + _parse_canonical_wire_metadata(metadata, limits=_limits()) + + +def test_parser_validates_all_metadata_before_tensor_allocation(monkeypatch): + plan = _prepare_canonical_wire_value(torch.arange(4), _limits()) + + def fail_allocation(*args, **kwargs): + raise AssertionError("parser allocated a tensor") + + monkeypatch.setattr(wire.torch, "empty", fail_allocation) + prepared = _parse_canonical_wire_metadata(plan.metadata, limits=_limits()) + assert prepared.tensor_specs == plan.tensor_specs + with pytest.raises(ValueError, match="trailing"): + _parse_canonical_wire_metadata(plan.metadata + b"x", limits=_limits()) + + +@pytest.mark.parametrize( + "field", + [field.name for field in dataclasses.fields(_CanonicalWireLimits)], +) +def test_limits_require_strict_positive_ints(field): + limits = _limits() + with pytest.raises(ValueError): + dataclasses.replace(limits, **{field: 0}) + with pytest.raises(TypeError): + dataclasses.replace(limits, **{field: True}) + + +def test_limits_reject_unrepresentable_or_inconsistent_bounds(): + with pytest.raises(TypeError): + _CanonicalWireLimits() + with pytest.raises(ValueError): + _limits(max_tensor_rank=(1 << 16)) + with pytest.raises(ValueError, match="recursive parsing"): + _limits(max_tree_depth=257) + with pytest.raises(ValueError): + _CanonicalWireLimits( + max_fragment_tensor_bytes=(1 << 40) + 1, + max_collective_tensor_bytes=1 << 40, + max_collective_metadata_bytes=1 << 30, + ) + with pytest.raises(ValueError): + _CanonicalWireLimits( + max_fragment_tensor_bytes=1 << 20, + max_collective_tensor_bytes=8 << 20, + max_collective_metadata_bytes=1 << 30, + max_metadata_bytes=(1 << 30) + 1, + ) + + +def test_prepare_and_parse_enforce_tree_and_scalar_limits(): + cases = [ + ([None, None], _limits(max_container_items=1), "max_container_items"), + ([[None]], _limits(max_tree_depth=1), "max_tree_depth"), + ([None, None], _limits(max_tree_nodes=2), "max_tree_nodes"), + ("ab", _limits(max_string_bytes=1), "max_string_bytes"), + (256, _limits(max_integer_bytes=1), "max_integer_bytes"), + ] + for value, limits, match in cases: + with pytest.raises(ValueError, match=match): + _prepare_canonical_wire_value(value, limits) + encoded = [ + (_metadata(b"\x06" + struct.pack(">Q", 2) + b"\0\0"), _limits(max_container_items=1)), + (_metadata(b"\x06" + struct.pack(">Q", 1) + b"\x06" + struct.pack(">Q", 1) + b"\0"), _limits(max_tree_depth=1)), + (_metadata(b"\x06" + struct.pack(">Q", 2) + b"\0\0"), _limits(max_tree_nodes=2)), + (_metadata(_string("ab")), _limits(max_string_bytes=1)), + (_metadata(b"\x03\0" + struct.pack(">Q", 2) + b"\x01\x00"), _limits(max_integer_bytes=1)), + ] + for metadata, limits in encoded: + with pytest.raises(ValueError): + _parse_canonical_wire_metadata(metadata, limits=limits) + + +def test_prepare_and_parse_enforce_metadata_tensor_count_rank_and_byte_limits(): + tensor = torch.arange(4, dtype=torch.int32) + normal = _prepare_canonical_wire_value(tensor, _limits()) + two_tensors = _prepare_canonical_wire_value([tensor, tensor], _limits()) + rank_two = _prepare_canonical_wire_value(tensor.reshape(2, 2), _limits()) + with pytest.raises(ValueError, match="max_metadata_bytes"): + _prepare_canonical_wire_value(None, _limits(max_metadata_bytes=24)) + with pytest.raises(ValueError, match="max_metadata_bytes"): + _parse_canonical_wire_metadata(normal.metadata, limits=_limits(max_metadata_bytes=len(normal.metadata) - 1)) + with pytest.raises(ValueError, match="max_tensors"): + _prepare_canonical_wire_value([tensor, tensor], _limits(max_tensors=1)) + with pytest.raises(ValueError, match="max_tensors"): + _parse_canonical_wire_metadata(two_tensors.metadata, limits=_limits(max_tensors=1)) + with pytest.raises(ValueError, match="max_tensor_rank"): + _prepare_canonical_wire_value(tensor.reshape(2, 2), _limits(max_tensor_rank=1)) + with pytest.raises(ValueError, match="max_tensor_rank"): + _parse_canonical_wire_metadata(rank_two.metadata, limits=_limits(max_tensor_rank=1)) + with pytest.raises(ValueError, match="max_fragment_tensor_bytes"): + _prepare_canonical_wire_value(tensor, _limits(max_fragment_tensor_bytes=15)) + with pytest.raises(ValueError, match="max_fragment_tensor_bytes"): + _parse_canonical_wire_metadata(normal.metadata, limits=_limits(max_fragment_tensor_bytes=15)) + + +def test_reconstruction_rejects_fragment_payload_and_geometry_corruption(): + value = torch.tensor([1.0, 2.0], dtype=torch.float32) + plan = _prepare_canonical_wire_value(value, _limits()) + prepared = _parse_canonical_wire_metadata(plan.metadata, limits=_limits()) + with pytest.raises(ValueError, match="fragment digest"): + _reconstruct_canonical_wire_value(prepared, plan.tensors, expected_fragment_digest=b"x" * 32) + corrupted = (plan.tensors[0].clone(),) + corrupted[0][0] = 9 + with pytest.raises(ValueError, match="invalid digest"): + _reconstruct_canonical_wire_value(prepared, corrupted) + with pytest.raises(ValueError, match="payload count"): + _reconstruct_canonical_wire_value(prepared, ()) + with pytest.raises(ValueError, match="descriptor"): + _reconstruct_canonical_wire_value(prepared, (plan.tensors[0].to(torch.float64),)) + with pytest.raises(ValueError, match="descriptor"): + _reconstruct_canonical_wire_value(prepared, (plan.tensors[0].reshape(1, 2),)) + + +def test_reconstruction_rejects_nonfinite_nontight_and_noncontiguous_payloads(): + plan = _prepare_canonical_wire_value(torch.tensor([1.0, 2.0]), _limits()) + prepared = _parse_canonical_wire_metadata(plan.metadata, limits=_limits()) + nonfinite = plan.tensors[0].clone() + nonfinite[0] = math.inf + with pytest.raises(ValueError, match="finite"): + _reconstruct_canonical_wire_value(prepared, (nonfinite,)) + backing = torch.empty(3, dtype=torch.float32) + backing[:2].copy_(plan.tensors[0]) + with pytest.raises(ValueError, match="descriptor"): + _reconstruct_canonical_wire_value(prepared, (backing[:2],)) + noncontiguous = torch.stack((plan.tensors[0], plan.tensors[0]))[:, 0] + assert not noncontiguous.is_contiguous() + with pytest.raises(ValueError, match="descriptor"): + _reconstruct_canonical_wire_value(prepared, (noncontiguous,)) + + +def test_prepare_rejects_unsupported_values_nonfinite_and_tensor_subclasses(): + class TensorSubclass(torch.Tensor): + pass + + invalid = [object(), {1: "bad"}, float("nan"), torch.tensor([float("inf")]), torch.tensor([1], dtype=torch.uint16)] + invalid.append(torch.Tensor._make_subclass(TensorSubclass, torch.tensor([1.0]), False)) + for value in invalid: + with pytest.raises((TypeError, ValueError)): + _prepare_canonical_wire_value(value, _limits()) + + +def test_strict_internal_type_construction_and_api_arguments(): + limits = _limits() + with pytest.raises(TypeError): + _prepare_canonical_wire_value(None, object()) + with pytest.raises(TypeError): + _parse_canonical_wire_metadata(bytearray(_metadata(b"\0")), limits=limits) + with pytest.raises(TypeError): + _reconstruct_canonical_wire_value(object(), ()) + with pytest.raises(TypeError): + _reconstruct_canonical_wire_value( + _parse_canonical_wire_metadata(_metadata(b"\0"), limits=limits), + iter(()), + ) + with pytest.raises(ValueError): + _CanonicalWireTensorSpec(0, 1, torch.int8, (), 1, 1, b"x" * 32) + assert dataclasses.is_dataclass(_CanonicalWirePlan) + assert dataclasses.is_dataclass(_PreparedCanonicalWireValue) + with pytest.raises(dataclasses.FrozenInstanceError): + limits.chunk_bytes = 1 diff --git a/tests/test_rebinding_cpu.py b/tests/test_rebinding_cpu.py new file mode 100644 index 0000000..72e7ec6 --- /dev/null +++ b/tests/test_rebinding_cpu.py @@ -0,0 +1,1064 @@ +"""CPU coverage for atomic pre-initialization parameter rebinding.""" + +import copy +from dataclasses import FrozenInstanceError, fields, is_dataclass, replace +import gc +import weakref + +import pytest +import torch + +from gefen import ( + Gefen, + GefenMuon, + GefenMuonHybrid, + LogicalSlice, + ParameterIdentity, + ParameterLayout, + ParameterRebinding, + PlacementKind, + ProcessGroupIdentity, + ShardIdentity, + ShardPlacement, + ShardingManifest, +) +from gefen.rebinding import LogicalSlotBinding + + +def _replicated(parameter, fqn, shape=None): + identity = ParameterIdentity(fqn, tuple(parameter.shape) if shape is None else shape) + shard = ShardIdentity( + identity, + ParameterLayout.REPLICATED, + LogicalSlice.full(identity), + ) + return shard + + +def _flat_manifest(fqn, shape, lengths, local_member): + identity = ParameterIdentity(fqn, shape) + members = tuple("rank:{}".format(index) for index in range(len(lengths))) + group = ProcessGroupIdentity("data_parallel", members) + shards = [] + offset = 0 + for coordinate, (member, length) in enumerate(zip(members, lengths)): + shards.append( + ShardIdentity( + identity, + ParameterLayout.FLATTENED_ELEMENT_SHARD, + LogicalSlice(offset, length), + placements=( + ShardPlacement( + "data_parallel", + PlacementKind.FLAT_SHARD, + coordinate, + len(members), + ), + ), + process_group=group, + local_member=member, + ) + ) + offset += length + manifest = ShardingManifest(tuple(reversed(shards))) + local = next(item for item in manifest.shards if item.local_member == local_member) + return local, manifest + + +def _owner_shard(identity, group, local_member, owner): + coordinate = group.ordered_members.index(local_member) + return ShardIdentity( + identity, + ParameterLayout.WHOLE_PARAMETER_OWNER, + LogicalSlice.full(identity) if local_member == owner else LogicalSlice(0, 0), + placements=( + ShardPlacement( + "data_parallel", + PlacementKind.WHOLE_PARAMETER_OWNER, + coordinate, + len(group.ordered_members), + ), + ), + process_group=group, + local_member=local_member, + owner=owner, + ) + + +def _nested_equal(left, right): + assert type(left) is type(right) + if isinstance(left, dict): + assert set(left) == set(right) + for key in left: + _nested_equal(left[key], right[key]) + elif isinstance(left, (list, tuple)): + assert len(left) == len(right) + for left_item, right_item in zip(left, right): + _nested_equal(left_item, right_item) + elif torch.is_tensor(left): + assert torch.equal(left, right) + else: + assert left == right + + +def _contains_tensor(value): + if torch.is_tensor(value): + return True + if is_dataclass(value) and not isinstance(value, type): + return any( + _contains_tensor(getattr(value, field.name)) for field in fields(value) + ) + if isinstance(value, dict): + return any( + _contains_tensor(key) or _contains_tensor(item) + for key, item in value.items() + ) + if isinstance(value, (list, tuple, set, frozenset)): + return any(_contains_tensor(item) for item in value) + return False + + +def _snapshot(optimizer): + return { + "state": optimizer.state, + "state_items": tuple( + (parameter, pstate, copy.deepcopy(pstate)) for parameter, pstate in optimizer.state.items() + ), + "param_groups": optimizer.param_groups, + "groups": tuple( + ( + group, + group["params"], + tuple(group["params"]), + copy.deepcopy({key: value for key, value in group.items() if key != "params"}), + ) + for group in optimizer.param_groups + ), + "param_names": optimizer._param_names, + "param_names_value": dict(optimizer._param_names), + "bindings": optimizer._gefen_shard_bindings, + # Copy the mapping so an in-place add/remove/replace of a binding that + # keeps the registry object identity is still caught. + "bindings_contents": dict(optimizer._gefen_shard_bindings), + "local_bindings": optimizer._gefen_local_shard_bindings, + "manifest": optimizer._gefen_sharding_manifest, + "finalized": optimizer._gefen_post_sharding_finalized, + "finalized_slots": optimizer._gefen_finalized_slots, + "logical_slots": optimizer._gefen_logical_slots, + "caches": tuple( + (name, getattr(optimizer, name)) + for name in ( + "_gefen_codebook_by_device", + "_gefen_codebook_lut_by_device", + "_sr_seed_by_device", + "_gefen_global_step_by_device", + ) + ), + # Bitwise clone of each device-cache entry so an in-place copy_() into a + # cached tensor that preserves the cache identity is still caught. + "caches_contents": tuple( + ( + name, + { + device: value.detach().clone() if torch.is_tensor(value) else copy.deepcopy(value) + for device, value in getattr(optimizer, name).items() + }, + ) + for name in ( + "_gefen_codebook_by_device", + "_gefen_codebook_lut_by_device", + "_sr_seed_by_device", + "_gefen_global_step_by_device", + ) + ), + "capt_stacks": optimizer._capt_stacks, + "static_mark_sig": optimizer._static_mark_sig, + "global_step": optimizer._gefen_global_step, + "codebook": optimizer._gefen_codebook, + } + + +def _assert_snapshot(optimizer, snapshot): + assert optimizer.state is snapshot["state"] + assert optimizer.param_groups is snapshot["param_groups"] + assert optimizer._param_names is snapshot["param_names"] + assert optimizer._param_names == snapshot["param_names_value"] + assert optimizer._gefen_shard_bindings is snapshot["bindings"] + expected_bindings = snapshot["bindings_contents"] + assert optimizer._gefen_shard_bindings.keys() == expected_bindings.keys() + for key, expected_binding in expected_bindings.items(): + assert optimizer._gefen_shard_bindings[key] is expected_binding + assert optimizer._gefen_local_shard_bindings is snapshot["local_bindings"] + assert optimizer._gefen_sharding_manifest is snapshot["manifest"] + assert optimizer._gefen_post_sharding_finalized is snapshot["finalized"] + assert optimizer._gefen_finalized_slots is snapshot["finalized_slots"] + assert optimizer._gefen_logical_slots is snapshot["logical_slots"] + assert optimizer._capt_stacks is snapshot["capt_stacks"] + assert optimizer._static_mark_sig is snapshot["static_mark_sig"] + assert optimizer._gefen_global_step is snapshot["global_step"] + assert optimizer._gefen_codebook is snapshot["codebook"] + for parameter, state_ref, expected in snapshot["state_items"]: + assert optimizer.state[parameter] is state_ref + _nested_equal(optimizer.state[parameter], expected) + for group, params_ref, expected_params, expected_values in snapshot["groups"]: + assert group in optimizer.param_groups + assert group["params"] is params_ref + assert tuple(group["params"]) == expected_params + _nested_equal( + {key: value for key, value in group.items() if key != "params"}, + expected_values, + ) + for name, cache_ref in snapshot["caches"]: + assert getattr(optimizer, name) is cache_ref + for name, expected_cache in snapshot["caches_contents"]: + live_cache = getattr(optimizer, name) + assert live_cache.keys() == expected_cache.keys() + for device, expected_value in expected_cache.items(): + live_value = live_cache[device] + if torch.is_tensor(expected_value): + assert torch.is_tensor(live_value) + assert torch.equal(live_value, expected_value) + else: + assert live_value == expected_value + + +@pytest.mark.parametrize( + ("changes", "error", "match"), + [ + ({"group_index": True}, TypeError, "group_index must be an integer"), + ({"group_index": -1}, ValueError, "group_index must be nonnegative"), + ( + {"original_slot_index": 1.0}, + TypeError, + "original_slot_index must be an integer", + ), + ( + {"original_slot_index": -1}, + ValueError, + "original_slot_index must be nonnegative", + ), + ( + {"compatibility_name": object()}, + TypeError, + "compatibility_name must be a string", + ), + ( + {"compatibility_name": "MixedCase"}, + ValueError, + "compatibility_name must be lowercase", + ), + ({"shard": object()}, TypeError, "shard must be a ShardIdentity"), + ], +) +def test_logical_slot_binding_strict_validation(changes, error, match): + parameter = torch.nn.Parameter(torch.ones(4)) + values = { + "group_index": 0, + "original_slot_index": 0, + "compatibility_name": "", + "shard": _replicated(parameter, "Weight"), + } + values.update(changes) + + with pytest.raises(error, match=match): + LogicalSlotBinding(**values) + + +def _interleaved_logical_slot_optimizer(): + sources = [ + torch.nn.Parameter(torch.full((2, 2), float(index + 1))) + for index in range(5) + ] + source_refs = tuple(weakref.ref(parameter) for parameter in sources) + names = ( + "group0.live0", + "group0.remote", + "group0.live2", + "group1.remote", + "group1.live1", + ) + optimizer = GefenMuon( + [ + {"params": list(zip(names[:3], sources[:3]))}, + {"params": list(zip(names[3:], sources[3:]))}, + ], + fused=False, + ) + process_group = ProcessGroupIdentity("data_parallel", ("rank:0", "rank:1")) + owners = ("rank:0", "rank:1", "rank:0", "rank:1", "rank:0") + records = [] + local_shards = [] + for index, owner in enumerate(owners): + identity = ParameterIdentity("Model.Parameter{}".format(index), (2, 2)) + parameter_records = tuple( + _owner_shard(identity, process_group, member, owner) + for member in process_group.ordered_members + ) + records.extend(parameter_records) + local_shards.append( + next( + shard + for shard in parameter_records + if shard.local_member == "rank:0" + ) + ) + targets = { + index: torch.nn.Parameter(sources[index].detach().clone()) + for index in (0, 2, 4) + } + rebindings = [ + ParameterRebinding( + sources[index], + targets.get(index), + local_shards[index], + ) + for index in (4, 1, 3, 0, 2) + ] + optimizer.post_sharding( + rebindings, + manifest=ShardingManifest(tuple(records)), + ) + return optimizer, source_refs, targets, tuple(local_shards), names + + +def test_logical_slot_registry_preserves_original_interleaved_structure_without_tensors(): + optimizer, source_refs, targets, local_shards, names = ( + _interleaved_logical_slot_optimizer() + ) + gc.collect() + + assert all(reference() is None for reference in source_refs) + assert optimizer.param_groups[0]["params"] == [targets[0], targets[2]] + assert optimizer.param_groups[1]["params"] == [targets[4]] + assert optimizer.param_groups[0]["param_names"] == [names[0], names[2]] + assert optimizer.param_groups[1]["param_names"] == [names[4]] + assert tuple( + (slot.group_index, slot.original_slot_index) + for slot in optimizer._gefen_logical_slots + ) == ((0, 0), (0, 1), (0, 2), (1, 0), (1, 1)) + assert tuple( + slot.compatibility_name for slot in optimizer._gefen_logical_slots + ) == names + assert tuple( + slot.shard for slot in optimizer._gefen_logical_slots + ) == local_shards + assert not _contains_tensor(optimizer._gefen_logical_slots) + assert all( + not hasattr(slot, "__dict__") for slot in optimizer._gefen_logical_slots + ) + with pytest.raises(FrozenInstanceError): + optimizer._gefen_logical_slots[0].compatibility_name = "changed" + assert optimizer.optimizer_contract().capabilities.stable_shard_identity + + +@pytest.mark.parametrize( + "corruption", + [ + "logical_slots", + "manifest", + "finalized_slots", + "local_bindings", + "group_names", + "name_cache", + "state_name", + ], +) +def test_finalized_layout_guard_checks_logical_slot_caches_and_names(corruption): + parameters = [ + torch.nn.Parameter(torch.full((4,), float(index + 1))) + for index in range(2) + ] + optimizer = Gefen( + [("first", parameters[0]), ("second", parameters[1])], + fused=False, + factored_v_2d=False, + ) + shards = tuple( + _replicated(parameter, "Model.Parameter{}".format(index)) + for index, parameter in enumerate(parameters) + ) + optimizer.post_sharding( + tuple( + ParameterRebinding(parameter, parameter, shard) + for parameter, shard in zip(parameters, shards) + ), + manifest=ShardingManifest(shards), + ) + + if corruption == "logical_slots": + slots = optimizer._gefen_logical_slots + optimizer._gefen_logical_slots = ( + replace(slots[0], compatibility_name="changed"), + slots[1], + ) + elif corruption == "manifest": + optimizer._gefen_sharding_manifest = ShardingManifest( + ( + _replicated(parameters[0], "Model.Parameter0", shape=(5,)), + shards[1], + ) + ) + elif corruption == "finalized_slots": + optimizer._gefen_finalized_slots = (tuple(reversed(parameters)),) + elif corruption == "local_bindings": + optimizer._gefen_local_shard_bindings = tuple( + reversed(optimizer._gefen_local_shard_bindings) + ) + elif corruption == "group_names": + optimizer.param_groups[0]["param_names"].reverse() + elif corruption == "name_cache": + optimizer._param_names[parameters[0]] = "changed" + else: + optimizer.state[parameters[0]]["name"] = "changed" + + assert not optimizer.optimizer_contract().capabilities.stable_shard_identity + with pytest.raises(RuntimeError, match="changed outside post_sharding"): + optimizer.state_dict() + + +def test_replicated_rebind_preserves_legacy_name_and_enables_identity_contract(): + old = torch.nn.Parameter(torch.arange(16, dtype=torch.float32).reshape(4, 4)) + new = torch.nn.Parameter(old.detach().clone()) + optimizer = Gefen([("Encoder.MixedCase.Weight", old)], fused=False) + old_group_options = {key: value for key, value in optimizer.param_groups[0].items() if key != "params"} + identity = ParameterIdentity("Encoder.MixedCase.Weight", (4, 4)) + + optimizer.rebind_parameter(old, new, identity=identity) + + assert optimizer.param_groups[0]["params"] == [new] + assert {key: value for key, value in optimizer.param_groups[0].items() if key != "params"} == old_group_options + assert old not in optimizer.state + assert optimizer.state[new] == {"name": "encoder.mixedcase.weight"} + assert optimizer.parameter_identity(new) == identity + assert optimizer.shard_identity(new).parameter.fqn == "Encoder.MixedCase.Weight" + assert optimizer.shard_bindings() == ((new, optimizer.shard_identity(new)),) + assert optimizer.sharding_manifest().shards == (optimizer.shard_identity(new),) + contract = optimizer.optimizer_contract() + assert contract.capabilities.canonical_parameter_fqns + assert contract.capabilities.stable_shard_identity + assert contract.capabilities.shard_rebinding + assert contract.capabilities.post_sharding + assert contract.capabilities.canonical_state_io + assert contract.capabilities.explicit_process_group_codebook_scope + with pytest.raises(RuntimeError, match="already finalized"): + optimizer.rebind_parameter(new, new, identity=identity) + with pytest.raises(RuntimeError, match="cannot add"): + optimizer.add_param_group({"params": [torch.nn.Parameter(torch.ones(2))]}) + + +def test_same_tensor_identity_only_rebinding_is_supported(): + parameter = torch.nn.Parameter(torch.ones(4, 4)) + optimizer = Gefen([("legacy.name", parameter)], fused=False) + identity = ParameterIdentity("Exact.FQN", (4, 4)) + + optimizer.rebind_parameter(parameter, parameter, identity=identity) + + assert optimizer.param_groups[0]["params"][0] is parameter + assert optimizer.state[parameter]["name"] == "legacy.name" + assert optimizer.parameter_identity(parameter).fqn == "Exact.FQN" + + +def test_replicated_rebinding_accepts_nonoverlapping_strided_storage(): + old = torch.nn.Parameter(torch.ones(4)) + target = torch.nn.Parameter(torch.arange(8, dtype=torch.float32)[::2]) + optimizer = Gefen([("weight", old)], fused=False, factored_v_2d=False) + + optimizer.rebind_parameter( + old, + target, + identity=ParameterIdentity("Weight", (4,)), + ) + + assert optimizer.param_groups[0]["params"] == [target] + assert optimizer.optimizer_contract().capabilities.stable_shard_identity + + +def test_rebinding_after_external_parameter_replacement_removes_orphan_state(): + old = torch.nn.Parameter(torch.ones(4, 4)) + new = torch.nn.Parameter(torch.full((4, 4), 2.0)) + optimizer = Gefen([("layer.weight", old)], fused=False) + optimizer.param_groups[0]["params"][0] = new + + optimizer.rebind_parameter( + old, + new, + identity=ParameterIdentity("layer.weight", (4, 4)), + ) + + assert old not in optimizer.state + assert old not in optimizer._param_names + assert optimizer.state[new] == {"name": "layer.weight"} + + +def test_external_replacement_reorder_preserves_each_source_compatibility_name(): + old_a = torch.nn.Parameter(torch.ones(4)) + old_b = torch.nn.Parameter(torch.ones(4) * 2) + new_a = torch.nn.Parameter(torch.ones(4) * 3) + new_b = torch.nn.Parameter(torch.ones(4) * 4) + optimizer = Gefen( + [("source.a", old_a), ("source.b", old_b)], + fused=False, + factored_v_2d=False, + ) + optimizer.param_groups[0]["params"][:] = [new_b, new_a] + shard_a = _replicated(new_a, "Canonical.A") + shard_b = _replicated(new_b, "Canonical.B") + + optimizer.post_sharding( + ( + ParameterRebinding(old_a, new_a, shard_a), + ParameterRebinding(old_b, new_b, shard_b), + ), + manifest=ShardingManifest((shard_b, shard_a)), + ) + + assert optimizer.param_groups[0]["params"] == [new_b, new_a] + assert optimizer.param_groups[0]["param_names"] == ["source.b", "source.a"] + assert optimizer.state[new_a]["name"] == "source.a" + assert optimizer.state[new_b]["name"] == "source.b" + + +def test_flattened_plain_gefen_rebind_steps_like_direct_local_shard(): + old = torch.nn.Parameter(torch.arange(16, dtype=torch.float32).reshape(4, 4)) + rebound = torch.nn.Parameter(torch.arange(8, dtype=torch.float32)) + oracle = torch.nn.Parameter(rebound.detach().clone()) + optimizer = Gefen( + [("layer.weight", old)], + lr=1e-3, + fused=False, + factored_v_2d=False, + ) + reference = Gefen( + [("layer.weight", oracle)], + lr=1e-3, + fused=False, + factored_v_2d=False, + ) + local, manifest = _flat_manifest("Model.Layer.Weight", (4, 4), (8, 8), "rank:0") + optimizer.post_sharding((ParameterRebinding(old, rebound, local),), manifest=manifest) + grad = torch.linspace(-1, 1, 8) + + for parameter, current in ((rebound, optimizer), (oracle, reference)): + parameter.grad = grad.clone() + current.step() + current.zero_grad() + + assert torch.equal(rebound, oracle) + _nested_equal(optimizer.state[rebound], reference.state[oracle]) + assert torch.equal(optimizer._gefen_codebook, reference._gefen_codebook) + + +def test_native_resume_preserves_target_flat_shard_identity_and_continuation(): + def build(): + old = torch.nn.Parameter(torch.zeros(4, 4)) + local_param = torch.nn.Parameter(torch.arange(8, dtype=torch.float32)) + optimizer = Gefen( + [("layer.weight", old)], + lr=1e-3, + fused=False, + factored_v_2d=False, + ) + local, manifest = _flat_manifest("Model.Layer.Weight", (4, 4), (8, 8), "rank:0") + optimizer.post_sharding((ParameterRebinding(old, local_param, local),), manifest=manifest) + return optimizer, local_param, local + + source, source_param, _ = build() + target, target_param, target_identity = build() + for grad in (torch.linspace(-1, 1, 8), torch.linspace(1, -1, 8)): + source_param.grad = grad.clone() + source.step() + source.zero_grad() + checkpoint = copy.deepcopy(source.state_dict()) + with torch.no_grad(): + target_param.copy_(source_param) + + target.load_state_dict(checkpoint) + + assert target.shard_identity(target_param) == target_identity + assert target.optimizer_contract().capabilities.stable_shard_identity + next_grad = torch.linspace(-0.5, 0.5, 8) + for parameter, optimizer in ( + (source_param, source), + (target_param, target), + ): + parameter.grad = next_grad.clone() + optimizer.step() + optimizer.zero_grad() + assert torch.equal(source_param, target_param) + _nested_equal(source.state_dict(), target.state_dict()) + + +def test_factored_logical_matrix_flattening_rejects_atomically(): + old = torch.nn.Parameter(torch.ones(4, 4)) + new = torch.nn.Parameter(torch.ones(8)) + optimizer = Gefen([("layer.weight", old)], fused=False, factored_v_2d=True) + local, manifest = _flat_manifest("layer.weight", (4, 4), (8, 8), "rank:0") + snapshot = _snapshot(optimizer) + + with pytest.raises(ValueError, match="factored_v_2d=False"): + optimizer.post_sharding((ParameterRebinding(old, new, local),), manifest=manifest) + + _assert_snapshot(optimizer, snapshot) + + +def test_late_batch_failure_is_atomic_and_continuation_is_bit_exact(): + def build(): + params = [ + torch.nn.Parameter(torch.arange(8, dtype=torch.float32)), + torch.nn.Parameter(torch.arange(8, dtype=torch.float32) + 1), + ] + return Gefen( + [("first", params[0]), ("second", params[1])], + fused=False, + factored_v_2d=False, + ), params + + target, target_params = build() + control, control_params = build() + replacements = [ + torch.nn.Parameter(torch.arange(8, dtype=torch.float32) + 2), + torch.nn.Parameter(torch.arange(7, dtype=torch.float32) + 3), + ] + first = _replicated(replacements[0], "First", (8,)) + second = _replicated(torch.empty(8), "Second", (8,)) + manifest = ShardingManifest((second, first)) + snapshot = _snapshot(target) + + with pytest.raises(ValueError, match="complete parameter storage"): + target.post_sharding( + ( + ParameterRebinding(target_params[0], replacements[0], first), + ParameterRebinding(target_params[1], replacements[1], second), + ), + manifest=manifest, + ) + + _assert_snapshot(target, snapshot) + grads = (torch.linspace(-1, 1, 8), torch.linspace(1, -1, 8)) + for params, optimizer in ((target_params, target), (control_params, control)): + for parameter, grad in zip(params, grads): + parameter.grad = grad.clone() + optimizer.step() + optimizer.zero_grad() + for target_param, control_param in zip(target_params, control_params): + assert torch.equal(target_param, control_param) + _nested_equal(target.state_dict(), control.state_dict()) + + +def test_checkpoint_schema_preparation_failure_is_atomic(monkeypatch): + old = torch.nn.Parameter(torch.ones(4)) + new = torch.nn.Parameter(torch.ones(4) * 2) + optimizer = Gefen([("weight", old)], fused=False, factored_v_2d=False) + shard = _replicated(new, "Weight") + snapshot = _snapshot(optimizer) + + def reject_schema(): + raise RuntimeError("schema preparation failed") + + monkeypatch.setattr( + optimizer, + "_install_rank_local_checkpoint_schema", + reject_schema, + ) + with pytest.raises(RuntimeError, match="schema preparation failed"): + optimizer.post_sharding( + (ParameterRebinding(old, new, shard),), + manifest=ShardingManifest((shard,)), + ) + + _assert_snapshot(optimizer, snapshot) + + +def test_whole_optimizer_pristine_guard_rejects_after_partial_initialization(): + params = [ + torch.nn.Parameter(torch.ones(8)), + torch.nn.Parameter(torch.ones(8) * 2), + ] + optimizer = Gefen( + [("first", params[0]), ("second", params[1])], + fused=False, + factored_v_2d=False, + ) + params[0].grad = torch.ones_like(params[0]) + optimizer.step() + optimizer.zero_grad() + replacements = [torch.nn.Parameter(param.detach().clone()) for param in params] + shards = ( + _replicated(replacements[0], "First"), + _replicated(replacements[1], "Second"), + ) + snapshot = _snapshot(optimizer) + + with pytest.raises(RuntimeError, match="before optimizer mutation"): + optimizer.post_sharding( + tuple(ParameterRebinding(old, new, shard) for old, new, shard in zip(params, replacements, shards)), + manifest=ShardingManifest(shards), + ) + + _assert_snapshot(optimizer, snapshot) + + +@pytest.mark.parametrize("failure", ["duplicate_target", "unknown_source", "gradient", "manifest"]) +def test_rebinding_structural_failures_are_no_ops(failure): + old = torch.nn.Parameter(torch.ones(4)) + new = torch.nn.Parameter(torch.ones(4) * 2) + optimizer = Gefen([("weight", old)], fused=False, factored_v_2d=False) + shard = _replicated(new, "Weight") + manifest = ShardingManifest((shard,)) + binding = ParameterRebinding(old, new, shard) + if failure == "duplicate_target": + other = torch.nn.Parameter(torch.ones(4)) + optimizer = Gefen( + [("weight", old), ("other", other)], + fused=False, + factored_v_2d=False, + ) + other_shard = _replicated(new, "Other") + manifest = ShardingManifest((shard, other_shard)) + bindings = ( + binding, + ParameterRebinding(other, new, other_shard), + ) + elif failure == "unknown_source": + bindings = (ParameterRebinding(torch.nn.Parameter(torch.ones(4)), new, shard),) + elif failure == "gradient": + new.grad = torch.ones_like(new) + bindings = (binding,) + else: + other_shard = _replicated(new, "Other") + manifest = ShardingManifest((other_shard,)) + bindings = (binding,) + snapshot = _snapshot(optimizer) + + with pytest.raises((RuntimeError, ValueError)): + optimizer.post_sharding(bindings, manifest=manifest) + + _assert_snapshot(optimizer, snapshot) + + +def test_rebinding_rejects_distinct_targets_with_overlapping_storage(): + old = [ + torch.nn.Parameter(torch.ones(8)), + torch.nn.Parameter(torch.ones(8) * 2), + ] + optimizer = Gefen( + [("first", old[0]), ("second", old[1])], + fused=False, + factored_v_2d=False, + ) + storage = torch.arange(12, dtype=torch.float32) + targets = [ + torch.nn.Parameter(storage[:8]), + torch.nn.Parameter(storage[4:]), + ] + shards = ( + _replicated(targets[0], "First"), + _replicated(targets[1], "Second"), + ) + snapshot = _snapshot(optimizer) + + with pytest.raises(ValueError, match="must not overlap"): + optimizer.post_sharding( + ( + ParameterRebinding(old[0], targets[0], shards[0]), + ParameterRebinding(old[1], targets[1], shards[1]), + ), + manifest=ShardingManifest(shards), + ) + + _assert_snapshot(optimizer, snapshot) + + +@pytest.mark.parametrize("target_kind", ["noncontiguous", "integer", "internal_overlap"]) +def test_rebinding_rejects_invalid_target_storage_atomically(target_kind): + old = torch.nn.Parameter(torch.ones(4)) + optimizer = Gefen([("weight", old)], fused=False, factored_v_2d=False) + if target_kind == "noncontiguous": + target = torch.nn.Parameter(torch.arange(8, dtype=torch.float32)[::2]) + local, manifest = _flat_manifest("Weight", (8,), (4, 4), "rank:0") + error = "contiguous physical storage" + elif target_kind == "integer": + target = torch.nn.Parameter(torch.ones(4, dtype=torch.int64), requires_grad=False) + local = _replicated(target, "Weight") + manifest = ShardingManifest((local,)) + error = "floating-point storage" + else: + target = torch.nn.Parameter(torch.ones(1).expand(4)) + local = _replicated(target, "Weight") + manifest = ShardingManifest((local,)) + error = "internal storage overlap" + snapshot = _snapshot(optimizer) + before = target.detach().clone() + + with pytest.raises(ValueError, match=error): + optimizer.post_sharding((ParameterRebinding(old, target, local),), manifest=manifest) + + _assert_snapshot(optimizer, snapshot) + assert torch.equal(target, before) + + +def test_rebinding_rejects_mixed_local_members_for_one_process_group(): + old = [torch.nn.Parameter(torch.ones(2)), torch.nn.Parameter(torch.ones(2))] + targets = [torch.nn.Parameter(torch.ones(2)), torch.nn.Parameter(torch.ones(2))] + optimizer = Gefen( + [("first", old[0]), ("second", old[1])], + fused=False, + factored_v_2d=False, + ) + first_local, first_manifest = _flat_manifest("First", (4,), (2, 2), "rank:0") + second_local, second_manifest = _flat_manifest("Second", (4,), (2, 2), "rank:1") + manifest = ShardingManifest(first_manifest.shards + second_manifest.shards) + snapshot = _snapshot(optimizer) + + with pytest.raises(ValueError, match="one member"): + optimizer.post_sharding( + ( + ParameterRebinding(old[0], targets[0], first_local), + ParameterRebinding(old[1], targets[1], second_local), + ), + manifest=manifest, + ) + + _assert_snapshot(optimizer, snapshot) + + +def test_external_replacement_rejects_duplicate_stale_sources(): + old = [torch.nn.Parameter(torch.ones(4)), torch.nn.Parameter(torch.ones(4))] + targets = [torch.nn.Parameter(torch.ones(4)), torch.nn.Parameter(torch.ones(4))] + optimizer = Gefen( + [("first", old[0]), ("second", old[1])], + fused=False, + factored_v_2d=False, + ) + optimizer.param_groups[0]["params"][:] = targets + shards = ( + _replicated(targets[0], "First"), + _replicated(targets[1], "Second"), + ) + snapshot = _snapshot(optimizer) + + with pytest.raises(ValueError, match="source tensors must be unique"): + optimizer.post_sharding( + ( + ParameterRebinding(old[0], targets[0], shards[0]), + ParameterRebinding(old[0], targets[1], shards[1]), + ), + manifest=ShardingManifest(shards), + ) + + _assert_snapshot(optimizer, snapshot) + + +@pytest.mark.parametrize("mutation", ["replace", "reorder"]) +def test_finalized_layout_guard_rejects_direct_parameter_group_mutation(mutation): + old = [torch.nn.Parameter(torch.ones(4)), torch.nn.Parameter(torch.ones(4) * 2)] + new = [torch.nn.Parameter(torch.ones(4) * 3), torch.nn.Parameter(torch.ones(4) * 4)] + optimizer = Gefen( + [("first", old[0]), ("second", old[1])], + fused=False, + factored_v_2d=False, + ) + shards = (_replicated(new[0], "First"), _replicated(new[1], "Second")) + optimizer.post_sharding( + ( + ParameterRebinding(old[0], new[0], shards[0]), + ParameterRebinding(old[1], new[1], shards[1]), + ), + manifest=ShardingManifest(shards), + ) + before = [parameter.detach().clone() for parameter in new] + if mutation == "replace": + rogue = torch.nn.Parameter(torch.ones(4) * 5) + optimizer.param_groups[0]["params"][0] = rogue + else: + optimizer.param_groups[0]["params"].reverse() + assert not optimizer.optimizer_contract().capabilities.stable_shard_identity + + with pytest.raises(RuntimeError, match="changed outside post_sharding"): + optimizer.step() + with pytest.raises(RuntimeError, match="changed outside post_sharding"): + optimizer.state_dict() + with pytest.raises(RuntimeError, match="changed outside post_sharding"): + optimizer.shard_bindings() + with pytest.raises(RuntimeError, match="changed outside post_sharding"): + optimizer.sharding_manifest() + for parameter, expected in zip(new, before): + assert torch.equal(parameter, expected) + + +@pytest.mark.parametrize("optimizer_type", ["plain", "muon"]) +def test_finalized_layout_guard_rechecks_after_closure(optimizer_type): + shape = (4,) if optimizer_type == "plain" else (2, 2) + old = torch.nn.Parameter(torch.ones(shape)) + bound = torch.nn.Parameter(torch.ones(shape) * 2) + rogue = torch.nn.Parameter(torch.ones(shape) * 3) + if optimizer_type == "plain": + optimizer = Gefen([("weight", old)], fused=False, factored_v_2d=False) + else: + optimizer = GefenMuon([("weight", old)], fused=False) + optimizer.rebind_parameter( + old, + bound, + identity=ParameterIdentity("Weight", shape), + ) + bound_before = bound.detach().clone() + rogue_before = rogue.detach().clone() + + def mutate_layout(): + optimizer.param_groups[0]["params"][0] = rogue + rogue.grad = torch.ones_like(rogue) + return torch.tensor(1.0) + + with pytest.raises(RuntimeError, match="changed outside post_sharding"): + optimizer.step(mutate_layout) + + assert torch.equal(bound, bound_before) + assert torch.equal(rogue, rogue_before) + assert optimizer.state[bound] == {"name": "weight"} + + +def test_finalized_layout_guard_rechecks_after_load_pre_hook(): + old = torch.nn.Parameter(torch.ones(4)) + bound = torch.nn.Parameter(torch.ones(4) * 2) + rogue = torch.nn.Parameter(torch.ones(4) * 3) + optimizer = Gefen([("weight", old)], fused=False, factored_v_2d=False) + optimizer.rebind_parameter( + old, + bound, + identity=ParameterIdentity("Weight", (4,)), + ) + checkpoint = copy.deepcopy(optimizer.state_dict()) + bound_before = bound.detach().clone() + rogue_before = rogue.detach().clone() + + def mutate_layout(current, state_dict): + current.param_groups[0]["params"][0] = rogue + + optimizer.register_load_state_dict_pre_hook(mutate_layout) + with pytest.raises(RuntimeError, match="changed outside post_sharding"): + optimizer.load_state_dict(checkpoint) + + assert torch.equal(bound, bound_before) + assert torch.equal(rogue, rogue_before) + assert optimizer.state[bound] == {"name": "weight"} + + +def test_muon_replicated_rebind_steps_like_direct_optimizer(): + old = torch.nn.Parameter(torch.arange(16, dtype=torch.float32).reshape(4, 4)) + rebound = torch.nn.Parameter(old.detach().clone()) + oracle = torch.nn.Parameter(old.detach().clone()) + optimizer = GefenMuon([("layer.weight", old)], fused=False) + reference = GefenMuon([("layer.weight", oracle)], fused=False) + identity = ParameterIdentity("Layer.Weight", (4, 4)) + optimizer.rebind_parameter(old, rebound, identity=identity) + grad = torch.linspace(-1, 1, 16).reshape(4, 4) + + for parameter, current in ((rebound, optimizer), (oracle, reference)): + parameter.grad = grad.clone() + current.step() + current.zero_grad() + + assert torch.equal(rebound, oracle) + _nested_equal(optimizer.state[rebound], reference.state[oracle]) + + +def test_muon_whole_owner_post_sharding_prunes_nonowner_and_blocks_training(): + old_owner = torch.nn.Parameter(torch.ones(4, 4)) + old_nonowner = torch.nn.Parameter(torch.ones(4, 4) * 2) + new_owner = torch.nn.Parameter(torch.ones(4, 4) * 3) + optimizer = GefenMuon( + [("owner.weight", old_owner), ("remote.weight", old_nonowner)], + fused=False, + ) + group = ProcessGroupIdentity("data_parallel", ("rank:0", "rank:1")) + owner_identity = ParameterIdentity("Owner.Weight", (4, 4)) + remote_identity = ParameterIdentity("Remote.Weight", (4, 4)) + owner_records = tuple(_owner_shard(owner_identity, group, member, "rank:0") for member in group.ordered_members) + remote_records = tuple(_owner_shard(remote_identity, group, member, "rank:1") for member in group.ordered_members) + manifest = ShardingManifest(owner_records + remote_records) + local_owner = next(item for item in owner_records if item.local_member == "rank:0") + local_nonowner = next(item for item in remote_records if item.local_member == "rank:0") + + optimizer.post_sharding( + ( + ParameterRebinding(old_nonowner, None, local_nonowner), + ParameterRebinding(old_owner, new_owner, local_owner), + ), + manifest=manifest, + ) + + assert optimizer.param_groups[0]["params"] == [new_owner] + assert set(optimizer.state) == {new_owner} + assert old_owner not in optimizer._param_names + assert old_nonowner not in optimizer._param_names + assert optimizer.shard_bindings() == tuple( + sorted( + ((new_owner, local_owner), (None, local_nonowner)), + key=lambda item: item[1].sort_key, + ) + ) + contract = optimizer.optimizer_contract() + assert contract.capabilities.canonical_parameter_fqns + assert contract.capabilities.stable_shard_identity + assert contract.capabilities.explicit_process_group_codebook_scope + assert all( + support.layout is not ParameterLayout.WHOLE_PARAMETER_OWNER for support in contract.capabilities.training + ) + snapshot = copy.deepcopy(optimizer.state_dict()) + parameter_snapshot = new_owner.detach().clone() + closure_calls = [] + with pytest.raises(RuntimeError, match="codebook scope"): + optimizer.step(lambda: closure_calls.append(True)) + assert closure_calls == [] + assert torch.equal(new_owner, parameter_snapshot) + _nested_equal(optimizer.state_dict(), snapshot) + + +def test_muon_all_nonowner_post_sharding_retains_manifest_without_fake_state(): + old = torch.nn.Parameter(torch.ones(4, 4)) + optimizer = GefenMuon([("remote.weight", old)], fused=False) + group = ProcessGroupIdentity("data_parallel", ("rank:0", "rank:1")) + identity = ParameterIdentity("Remote.Weight", (4, 4)) + records = tuple(_owner_shard(identity, group, member, "rank:1") for member in group.ordered_members) + local = next(item for item in records if item.local_member == "rank:0") + manifest = ShardingManifest(records) + + optimizer.post_sharding((ParameterRebinding(old, None, local),), manifest=manifest) + + assert optimizer.param_groups[0]["params"] == [] + assert optimizer.param_groups[0]["param_names"] == [] + assert optimizer.state == {} + assert optimizer.shard_bindings() == ((None, local),) + assert optimizer.sharding_manifest() == manifest + assert optimizer.optimizer_contract().capabilities.stable_shard_identity + + +def test_muon_flattened_rebinding_rejects_without_mutation(): + old = torch.nn.Parameter(torch.ones(4, 4)) + new = torch.nn.Parameter(torch.ones(8)) + optimizer = GefenMuon([("layer.weight", old)], fused=False) + local, manifest = _flat_manifest("Layer.Weight", (4, 4), (8, 8), "rank:0") + snapshot = _snapshot(optimizer) + + with pytest.raises(ValueError, match="complete matrices"): + optimizer.post_sharding((ParameterRebinding(old, new, local),), manifest=manifest) + + _assert_snapshot(optimizer, snapshot) + + +def test_gefen_backed_hybrid_contract_claims_composite_rebinding(): + matrix = torch.nn.Parameter(torch.ones(4, 4)) + bias = torch.nn.Parameter(torch.ones(4)) + optimizer = GefenMuonHybrid( + [("layer.weight", matrix)], + [("layer.bias", bias)], + lr=1e-3, + fused=False, + ) + contract = optimizer.optimizer_contract() + assert contract.capabilities.shard_rebinding + assert contract.capabilities.post_sharding + assert contract.capabilities.explicit_process_group_codebook_scope + assert not contract.capabilities.canonical_parameter_fqns + assert not contract.capabilities.stable_shard_identity + assert optimizer.state_dict()["backup_optimizer"] == "gefen" diff --git a/tests/test_scoped_collective_agreement_fixes.py b/tests/test_scoped_collective_agreement_fixes.py new file mode 100644 index 0000000..51a27a4 --- /dev/null +++ b/tests/test_scoped_collective_agreement_fixes.py @@ -0,0 +1,430 @@ +"""Scope-agreed codebook decisions and canonical import parameter freshness.""" + +from datetime import timedelta +import multiprocessing as mp +import os +import tempfile + +import pytest +import torch +import torch.distributed as dist + +from gefen import ( + CodebookProcessGroupBinding, + Gefen, + GefenMuon, + LogicalSlice, + ParameterIdentity, + ParameterLayout, + ParameterRebinding, + PlacementKind, + ProcessGroupIdentity, + ShardIdentity, + ShardPlacement, + ShardingManifest, +) + +from _state_snapshot import assert_deep_state_snapshot, deep_state_snapshot + + +_WORLD = 2 + + +def _replicated(parameter, group, member): + coordinate = group.ordered_members.index(member) + return ShardIdentity( + parameter, + ParameterLayout.REPLICATED, + LogicalSlice.full(parameter), + process_group=group, + local_member=member, + placements=( + ShardPlacement( + "dp", + PlacementKind.REPLICATE, + coordinate, + len(group.ordered_members), + ), + ), + ) + + +def _whole_owner(parameter, group, member, owner): + coordinate = group.ordered_members.index(member) + return ShardIdentity( + parameter, + ParameterLayout.WHOLE_PARAMETER_OWNER, + LogicalSlice.full(parameter) if member == owner else LogicalSlice(0, 0), + process_group=group, + local_member=member, + owner=owner, + placements=( + ShardPlacement( + "dp", + PlacementKind.WHOLE_PARAMETER_OWNER, + coordinate, + len(group.ordered_members), + ), + ), + ) + + +def _init_worker(rank, init_file): + dist.init_process_group( + "gloo", + init_method="file://{}".format(init_file), + rank=rank, + world_size=_WORLD, + timeout=timedelta(seconds=90), + ) + members = tuple("rank:{}".format(index) for index in range(_WORLD)) + return ProcessGroupIdentity("data_parallel", members), members + + +def _reuse_decision_worker(rank, init_file, queue): + """Rank-asymmetric automatic_period presence must not gate collectives. + + Reproduces the documented FSDP optim-state consolidation resume: the + loader restores per-parameter state but strips the optimizer-common + codebook on every member. A whole-parameter owner then holds restored + periods while empty non-owner members hold no state at all, so the + rank-local _resuming_from_checkpoint() predicate diverges. The group + must still resolve ONE reuse_existing_periods decision before the + conditionally issued "initialize" operation header. + """ + + try: + group, members = _init_worker(rank, init_file) + runtime_group = dist.group.WORLD + + owner_source = torch.nn.Parameter(torch.ones(2, 2)) + optimizer = GefenMuon([("matrix", owner_source)], fused=False) + identity = ParameterIdentity("Matrix", (2, 2)) + records = tuple( + _whole_owner(identity, group, member, "rank:0") for member in members + ) + optimizer.post_sharding( + ( + ParameterRebinding( + owner_source, + owner_source if rank == 0 else None, + records[rank], + ), + ), + manifest=ShardingManifest(records), + codebook_process_group=CodebookProcessGroupBinding( + group, members[rank], runtime_group, torch.device("cpu") + ), + ) + if rank == 0: + owner_source.grad = torch.tensor([[1.0, -2.0], [3.0, -4.0]]) + optimizer.step() + first_period = ( + int(optimizer.state[owner_source]["automatic_period"]) + if rank == 0 + else None + ) + state_asymmetric = ( + bool(optimizer.state) if rank == 0 else not optimizer.state + ) + + # Simulate the consolidation loader on every member: per-parameter + # state survives, the optimizer-common codebook does not. + optimizer._gefen_codebook = None + optimizer._gefen_codebook_by_device.clear() + optimizer._gefen_codebook_lut_by_device.clear() + optimizer._gefen_codebook_scope_validated = False + + if rank == 0: + owner_source.grad = torch.tensor([[0.5, -1.5], [2.5, -3.5]]) + + # Capture the reuse decision each member would gate collectives on, + # then abort symmetrically before any divergent collective runs. + captured = [] + + def capture(reuse_existing_periods=False): + captured.append(bool(reuse_existing_periods)) + raise RuntimeError("captured reuse decision") + + optimizer._ensure_gefen_codebook = capture + try: + optimizer.step() + raise AssertionError("sentinel resume step unexpectedly completed") + except RuntimeError as exc: + if "captured reuse decision" not in str(exc): + raise + finally: + del optimizer._ensure_gefen_codebook + + decisions = [None] * _WORLD + dist.all_gather_object(decisions, captured[0]) + + resume_completed = False + codebook_agreement = False + period_preserved = False + if len(set(decisions)) == 1: + # The schedule is agreed; the real resume step must complete + # collectively and keep the restored periods on the owner. + step_before = optimizer._gefen_global_step + optimizer.step() + resume_completed = ( + optimizer._gefen_codebook is not None + and optimizer._gefen_global_step == step_before + 1 + ) + codebooks = [ + torch.empty_like(optimizer._gefen_codebook) for _ in range(_WORLD) + ] + dist.all_gather(codebooks, optimizer._gefen_codebook) + codebook_agreement = all( + torch.equal(item, codebooks[0]) for item in codebooks[1:] + ) + if rank == 0: + period_preserved = ( + int(optimizer.state[owner_source]["automatic_period"]) + == first_period + ) + else: + period_preserved = not optimizer.state + queue.put( + { + "rank": rank, + "state_asymmetric": state_asymmetric, + "decisions": decisions, + "resume_completed": resume_completed, + "codebook_agreement": codebook_agreement, + "period_preserved": period_preserved, + } + ) + except Exception as exc: + queue.put({"rank": rank, "error": repr(exc)}) + finally: + if dist.is_available() and dist.is_initialized(): + dist.destroy_process_group() + + +def _revalidation_decision_worker(rank, init_file, queue): + """A rank-local scope-validated reset must re-validate on every member. + + A native checkpoint reload is rank-local and resets + _gefen_codebook_scope_validated on the member that loaded it only, while + every fingerprint the step + header exchanges stays identical. The next step must make the SAME + early-return decision inside _ensure_codebook_scope_agreement on every + member. + """ + + try: + group, members = _init_worker(rank, init_file) + runtime_group = dist.group.WORLD + + parameter = torch.nn.Parameter(torch.zeros(4)) + optimizer = Gefen([("weight", parameter)], fused=False, factored_v_2d=False) + identity = ParameterIdentity("Weight", (4,)) + records = tuple(_replicated(identity, group, member) for member in members) + optimizer.post_sharding( + (ParameterRebinding(parameter, parameter, records[rank]),), + manifest=ShardingManifest(records), + codebook_process_group=CodebookProcessGroupBinding( + group, members[rank], runtime_group, torch.device("cpu") + ), + ) + optimizer._resolve_automatic_period = lambda *args: 4 + parameter.grad = torch.tensor([-1.0, -0.5, 0.25, 1.0]) + optimizer.step() + validated_after_first = bool(optimizer._gefen_codebook_scope_validated) + + # Reload native state on ONE member only; every state value stays + # bit-identical, so the step header fingerprints still agree across + # members. + checkpoint = optimizer.state_dict() + if rank == 0: + optimizer.load_state_dict(checkpoint) + reset_asymmetric = ( + not optimizer._gefen_codebook_scope_validated + if rank == 0 + else bool(optimizer._gefen_codebook_scope_validated) + ) + + parameter.grad = torch.tensor([0.75, -0.25, 0.5, -1.0]) + captured = [] + + def capture(): + captured.append(not optimizer._gefen_codebook_scope_validated) + raise RuntimeError("captured revalidation decision") + + optimizer._ensure_codebook_scope_agreement = capture + try: + optimizer.step() + raise AssertionError("sentinel step unexpectedly completed") + except RuntimeError as exc: + if "captured revalidation decision" not in str(exc): + raise + finally: + del optimizer._ensure_codebook_scope_agreement + + decisions = [None] * _WORLD + dist.all_gather_object(decisions, captured[0]) + + second_completed = False + revalidated = False + parameters_agree = False + if len(set(decisions)) == 1: + optimizer.step() + second_completed = True + revalidated = bool(optimizer._gefen_codebook_scope_validated) + gathered = [torch.empty_like(parameter.detach()) for _ in range(_WORLD)] + dist.all_gather(gathered, parameter.detach()) + parameters_agree = all( + torch.equal(item, gathered[0]) for item in gathered[1:] + ) + queue.put( + { + "rank": rank, + "validated_after_first": validated_after_first, + "reset_asymmetric": reset_asymmetric, + "decisions": decisions, + "second_completed": second_completed, + "revalidated": revalidated, + "parameters_agree": parameters_agree, + } + ) + except Exception as exc: + queue.put({"rank": rank, "error": repr(exc)}) + finally: + if dist.is_available() and dist.is_initialized(): + dist.destroy_process_group() + + +def _run_workers(target): + context = mp.get_context("spawn") + queue = context.Queue() + fd, init_file = tempfile.mkstemp(prefix="gefen-scope-agreement-") + os.close(fd) + os.unlink(init_file) + processes = [ + context.Process(target=target, args=(rank, init_file, queue)) + for rank in range(_WORLD) + ] + try: + for process in processes: + process.start() + results = [queue.get(timeout=120) for _ in processes] + for process in processes: + process.join(timeout=10) + for process in processes: + if process.is_alive(): + process.terminate() + process.join(timeout=5) + pytest.fail("scoped agreement worker hung") + assert process.exitcode == 0 + return sorted(results, key=lambda item: item["rank"]) + finally: + for process in processes: + if process.is_alive(): + process.terminate() + process.join(timeout=5) + if os.path.exists(init_file): + os.unlink(init_file) + + +def test_resume_reuse_decision_is_scope_agreed_for_empty_members(): + results = _run_workers(_reuse_decision_worker) + + assert all("error" not in item for item in results), results + for item in results: + assert item["state_asymmetric"], item + # Every member must gate the conditional "initialize" header on the + # same group-wide decision; a resumed scope reuses restored periods. + assert len(set(item["decisions"])) == 1, item + assert all(item["decisions"]), item + assert item["resume_completed"], item + assert item["codebook_agreement"], item + assert item["period_preserved"], item + + +def test_rank_local_scope_validated_reset_revalidates_on_every_member(): + results = _run_workers(_revalidation_decision_worker) + + assert all("error" not in item for item in results), results + for item in results: + assert item["validated_after_first"], item + assert item["reset_asymmetric"], item + # Every member must make the same re-validation decision after a + # collective-free rank-local reset. + assert len(set(item["decisions"])) == 1, item + assert item["second_completed"], item + assert item["revalidated"], item + assert item["parameters_agree"], item + + +def _build_finalized_canonical(): + parameter = torch.nn.Parameter(torch.linspace(-0.5, 0.7, 8)) + optimizer = Gefen( + [("layer.weight", parameter)], + fused=False, + factored_v_2d=False, + ) + optimizer.rebind_parameter( + parameter, + parameter, + identity=ParameterIdentity("Layer.Weight", (8,)), + ) + optimizer._resolve_automatic_period = lambda *args, **kwargs: 4 + return optimizer, parameter + + +def _exported_canonical_state(): + source, source_parameter = _build_finalized_canonical() + source_parameter.grad = torch.linspace(-1.0, 0.8, 8) + source.step() + return source.export_canonical_state() + + +def test_parameter_storage_retarget_invalidates_prepared_canonical_import(): + exported = _exported_canonical_state() + target, target_parameter = _build_finalized_canonical() + + prepared = target.prepare_canonical_state_import(exported) + with torch.no_grad(): + target_parameter.data = target_parameter.data.clone() + + before_rejection = deep_state_snapshot(target) + with pytest.raises( + RuntimeError, match="changed after canonical import preparation" + ): + target.commit_canonical_state_import(prepared) + # The rejected commit must be a no-op before any recovery import masks it. + assert_deep_state_snapshot(target, before_rejection) + + # The refusal must keep canonical I/O available with a fresh preparation. + target.import_canonical_state(exported) + assert "automatic_period" in target.state[target_parameter] + + +def test_parameter_inplace_mutation_invalidates_prepared_canonical_import(): + exported = _exported_canonical_state() + target, target_parameter = _build_finalized_canonical() + + prepared = target.prepare_canonical_state_import(exported) + with torch.no_grad(): + target_parameter.mul_(2.0) + + before_rejection = deep_state_snapshot(target) + with pytest.raises( + RuntimeError, match="changed after canonical import preparation" + ): + target.commit_canonical_state_import(prepared) + # The rejected commit must be a no-op before any recovery import masks it. + assert_deep_state_snapshot(target, before_rejection) + + target.import_canonical_state(exported) + assert "automatic_period" in target.state[target_parameter] + + +def test_undisturbed_prepared_canonical_import_still_commits(): + exported = _exported_canonical_state() + target, target_parameter = _build_finalized_canonical() + + prepared = target.prepare_canonical_state_import(exported) + target.commit_canonical_state_import(prepared) + assert "automatic_period" in target.state[target_parameter] diff --git a/tests/test_shard_identity_contracts.py b/tests/test_shard_identity_contracts.py new file mode 100644 index 0000000..3e58d56 --- /dev/null +++ b/tests/test_shard_identity_contracts.py @@ -0,0 +1,746 @@ +"""CPU coverage for canonical parameter and shard identity descriptors.""" + +from dataclasses import FrozenInstanceError + +import pytest + +import gefen +import gefen.contracts as contracts_module +from gefen import ( + IDENTITY_SCHEMA_VERSION, + LogicalSlice, + ParameterIdentity, + ParameterLayout, + PlacementKind, + ProcessGroupIdentity, + ShardIdentity, + ShardPlacement, + ShardingManifest, +) +from gefen.contracts import LogicalRegion + + +def _group(): + return ProcessGroupIdentity("data_parallel", ("rank:2", "rank:0", "rank:1")) + + +def _flat_shard(parameter, group, member, offset, length): + coordinate = group.ordered_members.index(member) + return ShardIdentity( + parameter, + ParameterLayout.FLATTENED_ELEMENT_SHARD, + LogicalSlice(offset, length), + placements=( + ShardPlacement( + "data_parallel", + PlacementKind.FLAT_SHARD, + coordinate, + len(group.ordered_members), + ), + ), + process_group=group, + local_member=member, + ) + + +def _owner_shard(parameter, group, member, owner): + coordinate = group.ordered_members.index(member) + logical_slice = LogicalSlice.full(parameter) if member == owner else LogicalSlice(0, 0) + return ShardIdentity( + parameter, + ParameterLayout.WHOLE_PARAMETER_OWNER, + logical_slice, + placements=( + ShardPlacement( + "data_parallel", + PlacementKind.WHOLE_PARAMETER_OWNER, + coordinate, + len(group.ordered_members), + ), + ), + process_group=group, + local_member=member, + owner=owner, + ) + + +def _dtensor_shard( + parameter, + group, + member, + offsets, + lengths, + *, + dimension=0, + replicate=False, +): + coordinate = group.ordered_members.index(member) + return ShardIdentity( + parameter, + ParameterLayout.DTENSOR_1D_DEFAULT_WORLD, + LogicalRegion(offsets, lengths), + placements=( + ShardPlacement( + "dp", + PlacementKind.REPLICATE + if replicate + else PlacementKind.DIMENSION_SHARD, + coordinate, + len(group.ordered_members), + None if replicate else dimension, + ), + ), + process_group=group, + local_member=member, + ) + + +def test_parameter_and_process_group_identities_are_exact_and_immutable(): + parameter = ParameterIdentity("Encoder.Block.Weight", [4, 8]) + group = ProcessGroupIdentity("pipeline:1/dp", ["worker:b", "worker:a"]) + + assert parameter.schema_version == IDENTITY_SCHEMA_VERSION + assert parameter.fqn == "Encoder.Block.Weight" + assert parameter.global_shape == (4, 8) + assert parameter.numel == 32 + assert group.ordered_members == ("worker:b", "worker:a") + with pytest.raises(FrozenInstanceError): + parameter.fqn = "changed" + with pytest.raises(FrozenInstanceError): + group.semantic_name = "changed" + + +@pytest.mark.parametrize( + "fqn", ["", " layer.weight", "layer.weight ", ".layer", "layer.", "layer..weight", "layer\x00weight"] +) +def test_parameter_identity_rejects_noncanonical_fqns(fqn): + with pytest.raises(ValueError): + ParameterIdentity(fqn, (4, 4)) + + +@pytest.mark.parametrize("shape", [(4, -1), (4, True), (4, 1.5)]) +def test_parameter_identity_rejects_invalid_global_shapes(shape): + with pytest.raises(ValueError): + ParameterIdentity("layer.weight", shape) + + +@pytest.mark.parametrize("shape", ["48", b"\x04\x08", bytearray((4, 8))]) +def test_parameter_identity_rejects_string_and_bytes_shape_containers(shape): + with pytest.raises(TypeError, match="sequence"): + ParameterIdentity("layer.weight", shape) + + +def test_process_group_identity_preserves_authoritative_member_order(): + group = _group() + assert group.ordered_members == ("rank:2", "rank:0", "rank:1") + with pytest.raises(ValueError, match="unique"): + ProcessGroupIdentity("dp", ("rank:0", "rank:0")) + with pytest.raises(ValueError, match="non-empty"): + ProcessGroupIdentity("dp", ()) + with pytest.raises(TypeError, match="sequence"): + ProcessGroupIdentity("dp", "ab") + + +def test_placement_and_logical_slice_validation(): + placement = ShardPlacement("dp", PlacementKind.DIMENSION_SHARD, 1, 2, 0) + assert placement.parameter_dimension == 0 + with pytest.raises(ValueError, match="within"): + ShardPlacement("dp", PlacementKind.FLAT_SHARD, 2, 2) + with pytest.raises(ValueError, match="only"): + ShardPlacement("dp", PlacementKind.REPLICATE, 0, 1, 0) + with pytest.raises(ValueError, match="requires"): + ShardPlacement("dp", PlacementKind.DIMENSION_SHARD, 0, 1) + with pytest.raises(ValueError): + LogicalSlice(-1, 1) + with pytest.raises(ValueError): + LogicalSlice(0, -1) + + +def test_released_logical_slice_sort_key_shape_has_a_golden_value(): + parameter = ParameterIdentity("layer.weight", (8,)) + group = ProcessGroupIdentity("dp", ("rank:0", "rank:1")) + shard = _flat_shard(parameter, group, "rank:1", 4, 4) + + assert shard.sort_key == ( + "layer.weight", + 4, + 4, + "flattened_element_shard", + "dp", + 1, + -1, + (("data_parallel", "flat_shard", 1, 2, -1),), + ) + + +def test_logical_region_is_immutable_and_validates_rank_bounds_and_values(): + offsets = [1, 0] + lengths = [2, 4] + region = LogicalRegion(offsets, lengths) + offsets[0] = 0 + lengths[0] = 3 + + assert region.offsets == (1, 0) + assert region.lengths == (2, 4) + assert region.rank == 2 + assert region.numel == 8 + with pytest.raises(FrozenInstanceError): + region.offsets = (0, 0) + + parameter = ParameterIdentity("layer.weight", (4, 4)) + assert LogicalRegion.full(parameter) == LogicalRegion((0, 0), (4, 4)) + region.validate_bounds(parameter) + with pytest.raises(ValueError, match="same rank"): + LogicalRegion((0,), (1, 1)) + with pytest.raises(ValueError, match="nonnegative"): + LogicalRegion((-1, 0), (1, 1)) + with pytest.raises(ValueError, match="nonnegative"): + LogicalRegion((False, 0), (1, 1)) + with pytest.raises(TypeError, match="sequence"): + LogicalRegion("00", (1, 1)) + with pytest.raises(ValueError, match="rank"): + LogicalRegion((0,), (4,)).validate_bounds(parameter) + with pytest.raises(ValueError, match="exceeds"): + LogicalRegion((3, 0), (2, 4)).validate_bounds(parameter) + + +def test_logical_region_intersection_overlap_and_exact_coverage_utilities(): + parameter = ParameterIdentity("layer.weight", (4, 4)) + first = LogicalRegion((0, 0), (2, 4)) + second = LogicalRegion((2, 0), (2, 4)) + overlapping = LogicalRegion((1, 0), (2, 4)) + + assert first.intersection(second) == LogicalRegion((2, 0), (0, 4)) + assert not first.overlaps(second) + assert first.intersection(overlapping) == LogicalRegion((1, 0), (1, 4)) + assert first.overlaps(overlapping) + LogicalRegion.validate_exact_coverage(parameter, (second, first)) + with pytest.raises(ValueError, match="overlap"): + LogicalRegion.validate_exact_coverage(parameter, (first, overlapping)) + with pytest.raises(ValueError, match="exactly cover"): + LogicalRegion.validate_exact_coverage(parameter, (first,)) + with pytest.raises(ValueError, match="equal ranks"): + first.intersection(LogicalRegion((0,), (1,))) + + +def test_replicated_identity_can_be_local_or_process_group_scoped(): + parameter = ParameterIdentity("layer.weight", (4, 4)) + local = ShardIdentity( + parameter, + ParameterLayout.REPLICATED, + LogicalSlice.full(parameter), + ) + assert local.process_group is None + + group = _group() + scoped = [ + ShardIdentity( + parameter, + ParameterLayout.REPLICATED, + LogicalSlice.full(parameter), + placements=( + ShardPlacement( + "data_parallel", + PlacementKind.REPLICATE, + group.ordered_members.index(member), + len(group.ordered_members), + ), + ), + process_group=group, + local_member=member, + ) + for member in group.ordered_members + ] + manifest = ShardingManifest(tuple(reversed(scoped))) + assert {item.local_member for item in manifest.shards} == set(group.ordered_members) + + +def test_contiguous_slice_schema_rejects_dtensor_identity_until_regions_exist(): + parameter = ParameterIdentity("layer.weight", (4, 4)) + group = ProcessGroupIdentity("dp", ("rank:0", "rank:1")) + with pytest.raises(ValueError, match="logical-region"): + ShardIdentity( + parameter, + ParameterLayout.DTENSOR_1D_DEFAULT_WORLD, + LogicalSlice(0, 8), + placements=(ShardPlacement("dp", PlacementKind.DIMENSION_SHARD, 0, 2, 0),), + process_group=group, + local_member="rank:0", + ) + + +def test_dtensor_row_shard_manifest_accepts_uneven_regions_and_sorts_them(): + parameter = ParameterIdentity("layer.weight", (5, 4)) + group = ProcessGroupIdentity("dp", ("rank:0", "rank:1", "rank:2")) + shards = ( + _dtensor_shard(parameter, group, "rank:2", (4, 0), (1, 4)), + _dtensor_shard(parameter, group, "rank:0", (0, 0), (2, 4)), + _dtensor_shard(parameter, group, "rank:1", (2, 0), (2, 4)), + ) + + manifest = ShardingManifest(shards) + + assert tuple(item.local_member for item in manifest.shards) == ( + "rank:0", + "rank:1", + "rank:2", + ) + assert tuple(item.logical_region for item in manifest.shards) == ( + LogicalRegion((0, 0), (2, 4)), + LogicalRegion((2, 0), (2, 4)), + LogicalRegion((4, 0), (1, 4)), + ) + + +def test_dtensor_column_shard_and_replicate_manifests_are_complete(): + parameter = ParameterIdentity("layer.weight", (3, 5)) + group = ProcessGroupIdentity("dp", ("rank:0", "rank:1")) + column_manifest = ShardingManifest( + ( + _dtensor_shard( + parameter, + group, + "rank:1", + (0, 3), + (3, 2), + dimension=1, + ), + _dtensor_shard( + parameter, + group, + "rank:0", + (0, 0), + (3, 3), + dimension=1, + ), + ) + ) + assert tuple(item.logical_region.offsets for item in column_manifest.shards) == ( + (0, 0), + (0, 3), + ) + + replicated = ShardingManifest( + tuple( + _dtensor_shard( + parameter, + group, + member, + (0, 0), + parameter.global_shape, + replicate=True, + ) + for member in reversed(group.ordered_members) + ) + ) + assert all( + item.logical_region == LogicalRegion.full(parameter) + for item in replicated.shards + ) + + +def test_dtensor_manifest_retains_explicit_empty_member_regions(): + parameter = ParameterIdentity("small.weight", (2, 4)) + group = ProcessGroupIdentity( + "dp", ("rank:0", "rank:1", "rank:2", "rank:3") + ) + shards = tuple( + _dtensor_shard(parameter, group, member, offsets, lengths) + for member, offsets, lengths in ( + ("rank:0", (0, 0), (1, 4)), + ("rank:1", (1, 0), (1, 4)), + ("rank:2", (2, 0), (0, 4)), + ("rank:3", (2, 0), (0, 4)), + ) + ) + + manifest = ShardingManifest(tuple(reversed(shards))) + + assert tuple(item.local_member for item in manifest.shards) == ( + "rank:0", + "rank:1", + "rank:2", + "rank:3", + ) + assert tuple(item.logical_region.numel for item in manifest.shards) == ( + 4, + 4, + 0, + 0, + ) + + +def test_dtensor_manifest_rejects_gap_overlap_and_missing_member(): + parameter = ParameterIdentity("layer.weight", (4, 4)) + group = ProcessGroupIdentity("dp", ("rank:0", "rank:1")) + with pytest.raises(ValueError, match="gapless"): + ShardingManifest( + ( + _dtensor_shard(parameter, group, "rank:0", (0, 0), (1, 4)), + _dtensor_shard(parameter, group, "rank:1", (2, 0), (2, 4)), + ) + ) + with pytest.raises(ValueError, match="gapless"): + ShardingManifest( + ( + _dtensor_shard(parameter, group, "rank:0", (0, 0), (3, 4)), + _dtensor_shard(parameter, group, "rank:1", (2, 0), (2, 4)), + ) + ) + with pytest.raises(ValueError, match="each process-group member"): + ShardingManifest( + (_dtensor_shard(parameter, group, "rank:0", (0, 0), (4, 4)),) + ) + + +def test_dtensor_identity_rejects_invalid_region_and_placement_geometry(): + parameter = ParameterIdentity("layer.weight", (4, 4)) + group = ProcessGroupIdentity("dp", ("rank:0", "rank:1")) + with pytest.raises(ValueError, match="LogicalRegion"): + ShardIdentity( + parameter, + ParameterLayout.DTENSOR_1D_DEFAULT_WORLD, + LogicalSlice(0, 8), + placements=( + ShardPlacement("dp", PlacementKind.DIMENSION_SHARD, 0, 2, 0), + ), + process_group=group, + local_member="rank:0", + ) + with pytest.raises(ValueError, match="unsharded"): + _dtensor_shard(parameter, group, "rank:0", (0, 1), (2, 3)) + with pytest.raises(ValueError, match="replicated"): + _dtensor_shard( + parameter, + group, + "rank:0", + (0, 0), + (2, 4), + replicate=True, + ) + with pytest.raises(ValueError, match="one dimension-shard"): + ShardIdentity( + parameter, + ParameterLayout.DTENSOR_1D_DEFAULT_WORLD, + LogicalRegion((0, 0), (2, 4)), + placements=(ShardPlacement("dp", PlacementKind.FLAT_SHARD, 0, 2),), + process_group=group, + local_member="rank:0", + ) + with pytest.raises(ValueError, match="coordinates"): + ShardIdentity( + parameter, + ParameterLayout.DTENSOR_1D_DEFAULT_WORLD, + LogicalRegion((0, 0), (2, 4)), + placements=( + ShardPlacement("dp", PlacementKind.DIMENSION_SHARD, 1, 2, 0), + ), + process_group=group, + local_member="rank:0", + ) + + +@pytest.mark.parametrize("version", [True, 1.0, 0, 2]) +def test_parameter_identity_schema_version_requires_exact_supported_int(version): + with pytest.raises(ValueError, match="schema version"): + ParameterIdentity("layer.weight", (4, 4), schema_version=version) + + +def test_every_versioned_identity_descriptor_checks_its_schema_version(): + parameter = ParameterIdentity("layer.weight", (4, 4)) + local = ShardIdentity( + parameter, + ParameterLayout.REPLICATED, + LogicalSlice.full(parameter), + ) + with pytest.raises(ValueError, match="schema version"): + ProcessGroupIdentity("dp", ("rank:0",), schema_version=True) + with pytest.raises(ValueError, match="schema version"): + ShardIdentity( + parameter, + ParameterLayout.REPLICATED, + LogicalSlice.full(parameter), + schema_version=1.0, + ) + with pytest.raises(ValueError, match="schema version"): + ShardingManifest((local,), schema_version=False) + + +def test_flattened_manifest_normalizes_order_and_requires_exact_coverage(): + parameter = ParameterIdentity("Encoder.Weight", (4, 4)) + group = _group() + shards = ( + _flat_shard(parameter, group, "rank:1", 16, 0), + _flat_shard(parameter, group, "rank:0", 8, 8), + _flat_shard(parameter, group, "rank:2", 0, 8), + ) + + manifest = ShardingManifest(shards) + + assert tuple((item.logical_slice.flat_offset, item.local_member) for item in manifest.shards) == ( + (0, "rank:2"), + (8, "rank:0"), + (16, "rank:1"), + ) + assert manifest.for_parameter("Encoder.Weight") == manifest.shards + assert manifest.for_parameter("missing") == () + + with pytest.raises(ValueError, match="gapless"): + ShardingManifest( + ( + _flat_shard(parameter, group, "rank:2", 0, 7), + _flat_shard(parameter, group, "rank:0", 8, 8), + _flat_shard(parameter, group, "rank:1", 16, 0), + ) + ) + with pytest.raises(ValueError, match="gapless"): + ShardingManifest( + ( + _flat_shard(parameter, group, "rank:2", 0, 9), + _flat_shard(parameter, group, "rank:0", 8, 8), + _flat_shard(parameter, group, "rank:1", 16, 0), + ) + ) + with pytest.raises(ValueError, match="each process-group member"): + ShardingManifest(shards[:-1]) + inconsistent_axis = ShardIdentity( + parameter, + ParameterLayout.FLATTENED_ELEMENT_SHARD, + LogicalSlice(8, 8), + placements=(ShardPlacement("zero", PlacementKind.FLAT_SHARD, 1, 3),), + process_group=group, + local_member="rank:0", + ) + with pytest.raises(ValueError, match="one topology"): + ShardingManifest((shards[0], inconsistent_axis, shards[2])) + with pytest.raises(ValueError, match="unique"): + ShardingManifest((shards[0], shards[0])) + + +def test_flattened_manifest_empty_shard_is_order_independent_at_middle_boundary(): + parameter = ParameterIdentity("layer.weight", (4, 4)) + group = _group() + manifest = ShardingManifest( + ( + _flat_shard(parameter, group, "rank:2", 0, 8), + _flat_shard(parameter, group, "rank:0", 8, 8), + _flat_shard(parameter, group, "rank:1", 8, 0), + ) + ) + assert sum(item.logical_slice.length for item in manifest.shards) == 16 + with pytest.raises(ValueError, match="partition boundary"): + ShardingManifest( + ( + _flat_shard(parameter, group, "rank:2", 0, 8), + _flat_shard(parameter, group, "rank:0", 8, 8), + _flat_shard(parameter, group, "rank:1", 7, 0), + ) + ) + + +def test_flattened_shard_requires_explicit_group_and_in_bounds_slice(): + parameter = ParameterIdentity("layer.weight", (4, 4)) + group = _group() + placement = ShardPlacement("dp", PlacementKind.FLAT_SHARD, 0, 3) + with pytest.raises(ValueError, match="process group"): + ShardIdentity( + parameter, + ParameterLayout.FLATTENED_ELEMENT_SHARD, + LogicalSlice(0, 8), + placements=(placement,), + ) + with pytest.raises(ValueError, match="exceeds"): + _flat_shard(parameter, group, "rank:2", 12, 8) + with pytest.raises(ValueError, match="coordinates"): + ShardIdentity( + parameter, + ParameterLayout.FLATTENED_ELEMENT_SHARD, + LogicalSlice(0, 8), + placements=(ShardPlacement("dp", PlacementKind.FLAT_SHARD, 1, 3),), + process_group=group, + local_member="rank:2", + ) + + +def test_replicated_identity_rejects_ambiguous_multi_axis_placements(): + parameter = ParameterIdentity("layer.weight", (4, 4)) + group = ProcessGroupIdentity("mesh", ("rank:0",)) + placements = ( + ShardPlacement("tensor", PlacementKind.REPLICATE, 0, 1), + ShardPlacement("data", PlacementKind.REPLICATE, 0, 1), + ) + with pytest.raises(ValueError, match="one placement"): + ShardIdentity( + parameter, + ParameterLayout.REPLICATED, + LogicalSlice.full(parameter), + placements=placements, + process_group=group, + local_member="rank:0", + ) + + +def test_whole_parameter_owner_manifest_has_one_complete_owner(): + parameter = ParameterIdentity("layer.weight", (4, 4)) + group = _group() + shards = tuple(_owner_shard(parameter, group, member, "rank:0") for member in group.ordered_members) + + manifest = ShardingManifest(tuple(reversed(shards))) + + owner = next(item for item in manifest.shards if item.local_member == "rank:0") + assert owner.logical_slice == LogicalSlice.full(parameter) + assert all(item.logical_slice.length == 0 for item in manifest.shards if item.local_member != "rank:0") + with pytest.raises(ValueError, match="non-owner"): + ShardIdentity( + parameter, + ParameterLayout.WHOLE_PARAMETER_OWNER, + LogicalSlice.full(parameter), + placements=shards[0].placements, + process_group=group, + local_member="rank:2", + owner="rank:0", + ) + + +def test_zero_numel_whole_parameter_still_distinguishes_owner_by_member(): + parameter = ParameterIdentity("empty.weight", (0, 4)) + group = ProcessGroupIdentity("dp", ("rank:0", "rank:1")) + manifest = ShardingManifest( + tuple(_owner_shard(parameter, group, member, "rank:0") for member in group.ordered_members) + ) + assert tuple(item.owner for item in manifest.shards) == ("rank:0", "rank:0") + + +def test_whole_owner_manifest_rejects_owner_disagreement(): + parameter = ParameterIdentity("layer.weight", (4, 4)) + group = ProcessGroupIdentity("dp", ("rank:0", "rank:1")) + with pytest.raises(ValueError, match="one owner"): + ShardingManifest( + ( + _owner_shard(parameter, group, "rank:0", "rank:0"), + _owner_shard(parameter, group, "rank:1", "rank:1"), + ) + ) + + +def test_manifest_rejects_parameter_or_layout_disagreement(): + group = _group() + first = ParameterIdentity("layer.weight", (4, 4)) + different = ParameterIdentity("layer.weight", (2, 8)) + with pytest.raises(ValueError, match="parameter identity"): + ShardingManifest( + ( + _flat_shard(first, group, "rank:2", 0, 8), + _flat_shard(different, group, "rank:0", 8, 8), + _flat_shard(first, group, "rank:1", 16, 0), + ) + ) + flat = tuple( + _flat_shard(first, group, member, offset, length) + for member, offset, length in ( + ("rank:2", 0, 8), + ("rank:0", 8, 8), + ("rank:1", 16, 0), + ) + ) + replicated = ShardIdentity( + first, + ParameterLayout.REPLICATED, + LogicalSlice.full(first), + placements=( + ShardPlacement("data_parallel", PlacementKind.REPLICATE, 0, 3), + ), + process_group=group, + local_member="rank:2", + ) + with pytest.raises(ValueError, match="layout or process group"): + ShardingManifest((replicated, flat[1], flat[2])) + + other_group = ProcessGroupIdentity( + "other_data_parallel", group.ordered_members + ) + with pytest.raises(ValueError, match="layout or process group"): + ShardingManifest( + ( + flat[0], + _flat_shard(first, other_group, "rank:0", 8, 8), + flat[2], + ) + ) + + +def test_identity_descriptors_copy_mutable_input_sequences_deeply_enough(): + shape = [4, 4] + members = ["rank:0"] + placements = [ + ShardPlacement("dp", PlacementKind.REPLICATE, 0, 1), + ] + parameter = ParameterIdentity("layer.weight", shape) + group = ProcessGroupIdentity("dp", members) + shard = ShardIdentity( + parameter, + ParameterLayout.REPLICATED, + LogicalSlice.full(parameter), + placements=placements, + process_group=group, + local_member="rank:0", + ) + source_shards = [shard] + manifest = ShardingManifest(source_shards) + + shape.append(8) + members.append("rank:1") + placements.clear() + source_shards.clear() + + assert parameter.global_shape == (4, 4) + assert group.ordered_members == ("rank:0",) + assert shard.placements == ( + ShardPlacement("dp", PlacementKind.REPLICATE, 0, 1), + ) + assert manifest.shards == (shard,) + + +def test_manifest_order_is_fqn_then_logical_slice_then_member(): + group = ProcessGroupIdentity("dp", ("rank:1", "rank:0")) + alpha = ParameterIdentity("Alpha.Weight", (0,)) + zeta = ParameterIdentity("Zeta.Weight", (2,)) + shards = ( + _flat_shard(zeta, group, "rank:0", 2, 0), + _flat_shard(alpha, group, "rank:0", 0, 0), + _flat_shard(zeta, group, "rank:1", 0, 2), + _flat_shard(alpha, group, "rank:1", 0, 0), + ) + + manifest = ShardingManifest(tuple(reversed(shards))) + + assert tuple( + ( + item.parameter.fqn, + item.logical_slice.flat_offset, + item.local_member, + ) + for item in manifest.shards + ) == ( + ("Alpha.Weight", 0, "rank:1"), + ("Alpha.Weight", 0, "rank:0"), + ("Zeta.Weight", 0, "rank:1"), + ("Zeta.Weight", 2, "rank:0"), + ) + + +def test_identity_contracts_are_public_lazy_exports(): + for name in ( + "ParameterIdentity", + "ProcessGroupIdentity", + "PlacementKind", + "ShardPlacement", + "LogicalSlice", + "ShardIdentity", + "ShardingManifest", + ): + assert name in gefen.__all__ + assert getattr(gefen, name).__module__ == "gefen.contracts" + assert "ParameterRebinding" in gefen.__all__ + assert gefen.ParameterRebinding.__module__ == "gefen.rebinding" + assert "LogicalRegion" in contracts_module.__all__