diff --git a/CHANGELOG.md b/CHANGELOG.md index 20b6b0f..fb1be69 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,15 @@ All notable changes to this project are documented here. This project adheres to [Semantic Versioning](https://semver.org/). +## Unreleased + +- Add stable post-sharding rebinding, local training, and exact same-topology native checkpoint continuation for narrow one-dimensional default-world DTensors, including uneven and empty shards. +- Freeze GefenMuon's DTensor parameter membership, order, rebinding, mesh route, and `sharded_mode` after the first collective routing validation; later edits fail before optimizer or parameter mutation because initialized state has no mode-migration rule. +- Extend `GefenMuonHybrid` atomic rebinding and native load staging to exact AdamW backups and mixed DTensor layouts while preserving nested child state and fail-before-mutation semantics. +- Bind finalized scoped native checkpoints to their exact process-group identity and local layout, including direct distributed Muon; unscoped distributed Muon retains its complete cross-world-size native checkpoint path. +- Add representation-qualified checkpoint transitions, including the directional `defined_projection_factored_to_block_live_fp32_target_period_one_v1` portable projection; reverse block-to-factored conversion remains unsupported. +- Replace the dense-everywhere portable DCP save format with a load-compatible sharded version that stores multi-member canonical fields as DCP DTensors and singleton fields as ordinary full tensors, supports plain-Gefen DTensor resharding and Gefen-backed Hybrid child namespaces, and bounds metadata, chunks, and aggregate payloads. + ## [0.4.0] - 2026-07-12 Correctness and compatibility: diff --git a/COMPATIBILITY.md b/COMPATIBILITY.md index ecb244f..9decf00 100644 --- a/COMPATIBILITY.md +++ b/COMPATIBILITY.md @@ -95,7 +95,7 @@ Native single-process optimizer `state_dict()`/`load_state_dict()` and the expli This format is deliberately same-topology only and currently requires one 1-D DeviceMesh spanning the default process-group world. Every rank must participate in both save and restore; each process temporarily holds all serialized rank payloads on CPU, so the leading checkpoint-time CPU cost approaches `world_size ×` its local optimizer-state size plus local serialization scratch. Loading validates world size, parameter order and names, global and local shapes and dtypes, mesh membership and names, structural placements, rank coordinates, global step, deterministic policy, frozen codebook, and sharded mode before mutation. Multidimensional meshes, subgroups, pipeline-local optimizers, world-size/topology changes, and old unsafe untagged full checkpoints fail closed rather than silently applying rank 0's state to every shard. No optimizer-state reshard portability is claimed, model-only DCP is unaffected, and the full-state DCP support described here does not extend beyond plain Gefen and Muon `approx`. -Finalized exact period-one plain Gefen, `GefenMuon(sharded_mode="distributed")`, and Gefen-backed `GefenMuonHybrid` additionally expose portable global-state DCP through `save_portable_dcp(...)` and `load_portable_dcp(...)`. This separate synchronous path gathers and verifies complete logical optimizer state, stores a tensor-only bounded canonical-wire envelope through PyTorch DCP, and projects it onto supported replicated, flattened, or whole-owner targets after load. It can change flattened placement and redistribute Muon owners across checkpoint world sizes; factored second moments remain replicated and same-topology. Gefen-backed Hybrid uses a separately versioned, digested composite wrapper around unchanged child v3 documents, exact disjoint FQN routing, and one all-child freshness/commit boundary. Every checkpoint member temporarily holds the complete dense global document and encoded CPU payload. The target import is fail-before-local-mutation after DCP has completed the read, but DCP storage publication itself is not transactionally atomic. AdamW-backed Hybrid, DTensor targets, asynchronous saving, mixed model/optimizer `Stateful` composition, singleton checkpoint scopes inside a larger initialized default world, and multi-member checkpoint groups whose coordinate zero is not global rank zero are not supported by this path. +Finalized exact period-one plain Gefen, `GefenMuon(sharded_mode="distributed")`, and Gefen-backed `GefenMuonHybrid` additionally expose portable global-state DCP through `save_portable_dcp(...)` and `load_portable_dcp(...)`. This separate synchronous path stores bounded semantic metadata plus DCP-native sharded fp32 fields, so checkpoint ranks retain canonical local chunks instead of all materializing one dense optimizer document. Plain Gefen can reshard exact block state among replicated, flattened, and narrow one-dimensional default-world DTensor layouts; replicated factored state restores exactly to replicated factored state or follows the versioned `defined_projection_factored_to_block_live_fp32_target_period_one_v1` rule to a replicated, flattened, or narrow-DTensor block target using compact factor slices and the saved exact fp32 row denominator. The standard DCP reader can transiently deserialize one saved source storage chunk before narrowing it, so each source chunk must fit the load-time fragment ceiling and a lower-limit reader may reject a checkpoint written with larger chunks. Distributed Muon remains limited to replicated and whole-owner portable layouts and can redistribute owners across checkpoint world sizes. Gefen-backed Hybrid stores each present child in a derived sharded namespace and uses exact disjoint FQN routing plus one all-child freshness/commit boundary. The target import is fail-before-local-mutation after DCP completes the read, but DCP storage publication itself is not transactionally atomic. The earlier dense DCP envelope remains load-compatible and retains its original complete-document memory cost. AdamW-backed Hybrid portable state, Muon DTensor resharding, block-to-factored conversion, asynchronous saving, mixed model/optimizer `Stateful` composition, singleton checkpoint scopes inside a larger initialized default world, and multi-member checkpoint groups whose coordinate zero is not global rank zero are not supported by this path. ## Transformers Trainer DDP diff --git a/README.md b/README.md index a57330c..31c2dc2 100644 --- a/README.md +++ b/README.md @@ -604,6 +604,8 @@ Under FSDP2, Muon's orthogonalization needs each full weight matrix, but every G - **`"distributed"`** (experimental) — each matrix is assigned by its stable position in the full distributed parameter set to one GPU, which does the work and shares the result. Bit-identical to `"exact"` on homogeneous GPUs and faster as you add them. Momentum ownership is stable when the active gradient set varies, and `state_dict()` collectively gathers owner-local momentum so rank 0 can write a complete checkpoint after all ranks call it. - **`"approx"`** — each GPU works on just its slice. Fastest, but results genuinely differ — an accuracy trade. +Choose the DTensor route before its first optimizer collective. After GefenMuon validates a multi-rank DTensor route, its parameter membership and order, live tensor/mesh binding, and `sharded_mode` are frozen; even a unanimous later edit is rejected because initialized local, global, and owner momentum have no implicit migration rule. Construct a new optimizer and use a supported checkpoint path when changing those semantics. + ```python from gefen import GefenMuonHybrid @@ -615,7 +617,7 @@ opt = GefenMuonHybrid( # only takes effect under FSDP2 (DTensor params); no-op single-GPU ``` -> **`"distributed"` checkpointing is collective.** `state_dict()` gathers each owner's momentum across ranks, so **every rank must call it** (as in a standard FSDP full-state-dict flow). Calling `state_dict()` on rank 0 only — e.g. a rank-0-only save loop — **deadlocks**. Save and load also transiently materialize the full unsharded momentum on every rank, so peak memory at checkpoint time approaches `"exact"` mode's. The saved checkpoint is complete on every rank and resumes under any world size, including a single-process optimizer; incomplete or inconsistent owner state fails closed before any state is touched. +> **`"distributed"` checkpointing is collective.** `state_dict()` gathers each owner's momentum across ranks, so **every rank must call it** (as in a standard FSDP full-state-dict flow). Calling `state_dict()` on rank 0 only — e.g. a rank-0-only save loop — **deadlocks**. Save and load also transiently materialize the full unsharded momentum on every rank, so peak memory at checkpoint time approaches `"exact"` mode's. An unscoped `GefenMuon` checkpoint is complete on every rank and resumes under any world size, including a single-process optimizer; incomplete or inconsistent owner state fails closed before mutation. Once direct Muon or `GefenMuonHybrid` is finalized with an explicit codebook scope, its ordinary native checkpoint carries an exact scope, topology, and local-layout guard and is same-topology only; use the separate portable DCP path for supported Gefen-backed topology changes. ![Gefen-Muon exact / distributed / approx sharded — eval loss](https://raw.githubusercontent.com/thad0ctor/Gefen-X/main/docs/benchmarks/muon_shard_loss.png) ![Gefen-Muon exact / distributed / approx sharded — throughput & VRAM](https://raw.githubusercontent.com/thad0ctor/Gefen-X/main/docs/benchmarks/muon_shard_perf.png) @@ -632,13 +634,13 @@ Measured (Qwen3-0.6B, 2 and 4 GPUs): `"distributed"` matched `"exact"` in the re `Gefen`, `GefenMuon`, and `GefenMuonHybrid` preserve the parameter groups you pass to the optimizer, so list-indexed layer-wise LR recipes, per-group LR logging, and `state_dict()["param_groups"]` see the same group boundaries as conventional `torch.optim` optimizers. Per-parameter names are stored in optimizer state and mirrored as each group's `param_names` list for integrations that need name-level routing. Checkpoints from the older one-group-per-parameter layout are migrated on load when the total parameter order still matches and the old per-param hyperparameters can be represented by the new group layout. -Finalized exact period-one plain Gefen, distributed-owner GefenMuon, and Gefen-backed `GefenMuonHybrid` also provide a separate portable global optimizer-state path through `save_portable_dcp(...)` and `load_portable_dcp(...)`. It can reshard supported block-state parameters, redistribute Muon owners, and restore both Gefen-backed Hybrid children as one validated composite transaction across checkpoint topologies; it is synchronous, temporarily materializes the complete dense optimizer document on every checkpoint rank, and is distinct from ordinary FSDP2 optimizer checkpoints. See the [optimizer integration contracts](https://github.com/thad0ctor/Gefen-X/blob/main/docs/optimizer_contracts.md#portable-global-state-v3) for the supported layouts, setup, and exclusions. +Finalized exact period-one plain Gefen, distributed-owner GefenMuon, and Gefen-backed `GefenMuonHybrid` also provide a separate portable global optimizer-state path through `save_portable_dcp(...)` and `load_portable_dcp(...)`. It uses DCP-native sharded fields rather than materializing the complete dense optimizer document on every checkpoint rank, can reshard plain-Gefen replicated/flattened/narrow-DTensor state and Muon replicated/whole-owner state, defines a versioned bounded factored-to-block projection, and restores every present Gefen-backed Hybrid child as one validated composite transaction. Retained payload is target-canonical-local, although the standard DCP reader may transiently deserialize one saved source storage chunk and rejects it when that chunk exceeds the current fragment ceiling. It is synchronous and distinct from ordinary FSDP2 optimizer checkpoints; the earlier dense DCP envelope remains load-compatible. See the [optimizer integration contracts](https://github.com/thad0ctor/Gefen-X/blob/main/docs/optimizer_contracts.md#portable-global-state-v3) for the exact optimizer/layout matrix, setup, and exclusions. Plain `Gefen` with ordinary replicated CUDA parameters can keep its persistent per-parameter optimizer state on CPU between eager steps with `optimizer.offload_state_("cpu")`. Each step synchronously stages only the parameter currently being updated to its CUDA device, copies the updated state back to CPU, and releases the temporary device state; the small shared codebook remains CUDA-resident. `optimizer.restore_state_()` atomically returns all state to the parameter devices, while `move_state_()` also disables an active offload policy. This path intentionally excludes `GefenMuon`, `GefenMuonHybrid`, sharded or DTensor parameters, multi-member explicit codebook scopes, capturable optimizers, `torch.compile`, and CUDA graph capture. ## Known limitations -- **Hybrid checkpoint schema.** `GefenMuonHybrid`'s ordinary `state_dict()` uses its own nested `{"muon": ..., "backup": ..., "backup_optimizer": "gefen" | "adamw"}` layout. Resume from a checkpoint the hybrid itself saved—not one consolidated or converted to the flat torch `{state, param_groups}` layout. Cross-backend loads are rejected before either child is mutated; legacy untagged hybrid checkpoints are interpreted as Gefen-backed. The separate topology-neutral DCP path above supports only a finalized Gefen-backed Hybrid; AdamW-backed Hybrid remains same-topology through its ordinary nested checkpoint. +- **Hybrid checkpoint schema.** `GefenMuonHybrid`'s ordinary `state_dict()` uses its own nested `{"muon": ..., "backup": ..., "backup_optimizer": "gefen" | "adamw"}` layout and adds a versioned `finalized_binding` record after `post_sharding`. Resume from a checkpoint the Hybrid itself saved—not one consolidated or converted to the flat torch `{state, param_groups}` layout. Finalized loads require the exact child set, routing, shard identity, and local topology before either child is mutated; finalized and legacy-unfinalized schemas do not cross-load. Cross-backend loads reject, while legacy untagged Hybrid checkpoints are interpreted as Gefen-backed. The separate topology-neutral DCP path above supports only a finalized Gefen-backed Hybrid; AdamW-backed Hybrid remains same-topology through its ordinary nested checkpoint. - **FSDP2 optimizer checkpoints don't reshard.** Plain Gefen and Muon `approx` resume only on the same GPU count and layout; changing either refuses to load. Model weights are unaffected — [details](https://github.com/thad0ctor/Gefen-X/blob/main/COMPATIBILITY.md#optimizer-checkpoint-scope). - **CPU state offload is synchronous.** Plain-Gefen state offload reduces persistent CUDA optimizer-state residency by paging one parameter at a time, but it adds blocking CPU↔CUDA transfers to every updated parameter and is not an asynchronous overlap engine. - **True-FP16 overflow skips are invisible to Accelerate's `step_was_skipped` flag.** BF16 and standard AMP are unaffected and are the recommended modes in Trainer/Accelerate. diff --git a/docs/optimizer_contracts.md b/docs/optimizer_contracts.md index 270ce89..703f110 100644 --- a/docs/optimizer_contracts.md +++ b/docs/optimizer_contracts.md @@ -25,11 +25,11 @@ assert ParameterLayout.DTENSOR_1D_DEFAULT_WORLD in rank_local_dcp.same_topology - `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, state movement, and offload are independent capability fields. 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. +The current DTensor declaration is deliberately narrow: `DTENSOR_1D_DEFAULT_WORLD` means exact public `DTensor` values on one shared one-dimensional mesh spanning the default world, with one `Replicate()` or `Shard(parameter_dimension)` placement. Multidimensional meshes, subgroups, partial placements, local tensor subclasses, and arbitrary placement-changing native loads are not implied. Supported placement-changing restore is declared separately by `CANONICAL_GLOBAL`. ## 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. +`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 describe one default-world mesh axis with either replication or one parameter-dimension shard, including uneven and empty shards. The rebinding, training, rank-local checkpoint, and plain-Gefen portable data planes consume those regions only through their separately declared capability entries. 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, codebook scoping, state movement, or offload; those capabilities remain separate. @@ -39,7 +39,7 @@ These descriptors do not treat legacy `param_names`, generated names, Python ten 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. +Plain Gefen rebinds complete replicated parameters, contiguous physical 1-D flattened element shards, and exact narrow DTensors. A flattened logical matrix requires `factored_v_2d=False`; a DTensor logical matrix uses local block-second-moment state even when the optimizer policy enables factored state for replicated matrices. GefenMuon rebinds complete replicated matrices, narrow DTensor matrices, and 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. `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. Its backup may be exact `Gefen` or exact `torch.optim.AdamW`; the AdamW adapter supports replicated, flattened-element, whole-owner, and narrow-DTensor routes in fp16, bf16, fp32, and fp64. AdamW whole-owner training requires complete owner storage and adapter-managed post-step parameter synchronization. The portable global-state path described below remains Gefen-backed because AdamW has no topology-neutral codec. 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; state movement and offload activation; 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. @@ -47,13 +47,13 @@ After finalization, every entry point re-validates the published layout, and thi `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 `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. `GefenMuonHybrid` implements one atomic composite transaction that installs the same exact binding object on every present Gefen or GefenMuon child while retaining independent child codebooks. An AdamW backup participates in the same rebinding transaction but has no learned codebook or codebook collective of its own. -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. +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 and every nonempty narrow-DTensor region contributes its local logical storage after nonempty members 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. +`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 replicated behavior remains collective-free. Unscoped multi-rank GefenMuon DTensor operations first exchange one fixed-size routing header on the initialized default world, require every participating mesh to span that world, and fail closed on subgroup or partially overlapping meshes; an eagerly validated capturable step instead restores its frozen route inside capture. The first successful routing validation freezes DTensor parameter membership and order, live tensor/mesh bindings, and `sharded_mode`; later edits, including unanimous mode changes, reject because no migration exists among local-shard, global, and owner momentum. 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. +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. Narrow DTensor checkpoints use rank-local carrier tensors and collective validation, require the exact same default-world mesh, placement, and local region, and reject placement or world-size changes. GefenMuon's rank-local DTensor declaration is limited by its mode-specific checkpoint entries. Hybrid's nested native loader stages every present child plus its finalized routing/binding guard before publishing either child; finalized flattened, whole-owner, and narrow-DTensor entries are same-topology only. This includes AdamW-backed layouts and remains distinct from the topology-changing portable path. ## Exact-binding canonical local state @@ -65,9 +65,9 @@ This is an exact-binding transport-neutral local fragment, distinct from the den ## 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. +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, authoritative state, representation-qualified 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 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 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. +`export_portable_state(checkpoint_process_group=..., transaction_id=..., limits=...)` and `import_portable_state(state, checkpoint_process_group=..., transaction_id=..., limits=...)` implement the complete-document in-memory path on exact `Gefen`, `GefenMuon`, and supported Gefen-backed `GefenMuonHybrid` instances. Every member receives the complete dense document. 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. Use the sharded DCP helpers below when checkpoint ranks must not all materialize that document. ```python from gefen import PortableStateLimits @@ -92,9 +92,11 @@ target_optimizer.import_portable_state( ) ``` -`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. +`save_portable_dcp(...)` and `load_portable_dcp(...)` provide the scalable synchronous PyTorch Distributed Checkpoint path. On a multi-member binding, envelope version 2 persists one bounded tensor-only metadata value plus semantic fp32 fields as DCP-native one-dimensional `Shard(0)` DTensors; a singleton binding stores each field as an ordinary full tensor because no checkpoint sharding exists. Save redistributes live replicated, flattened, whole-owner, or supported dimension-sharded payloads directly into canonical chunks; load retains only the target rank's canonical chunks and redistributes them directly into the target layout. The indexed dense path uses memory proportional to local payload, bounded fp32/int64 routing tensors, and an O(world-size) piece table even for a non-leading-dimension shard, including uneven and empty shards. A rank holds a complete logical field only when its live source or target layout inherently does so, such as replication or whole ownership. Factored row and column vectors remain their own fields, the saved fp32 row denominator preserves the live projection's exact operator order, and factored-to-block restore routes only the compact factor slices needed to create the target rank's local dense block. -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. +The version-2 load planner reads and validates framework metadata up front, checking the exact namespace and key set, exact tensor metadata and chunk geometry, tensor and chunk counts, per-field and aggregate bytes, bounded semantic metadata, source manifests, per-source canonical-chunk SHA-256 values, and completion digests before optimizer mutation. Before any payload read, `max_fragment_tensor_bytes` rejects an oversized aggregate projected target fp32 state, per-FQN factor scratch, source-chunk integrity reassembly or digest scratch, and fp32/int64 projection routing or amplified send tensor. Retained optimizer payload is target-canonical-local, but the standard DCP reader may transiently deserialize one saved source storage chunk before narrowing it; every source chunk must therefore fit the same ceiling, and a low-limit load can reject a checkpoint written with larger source chunks. PyTorch DCP may consult the storage reader's metadata again while loading the bounded metadata and payload subsets. Metadata, target options, routing, and every rank-local preparation failure are voted before the next rank-dependent collective. The exact storage class and its bounded string or path-like `checkpoint_id` are included in collective preflight, so every member must address the same checkpoint; a custom storage plugin without that identity fails closed. Tensor-only optimizer data avoids DCP's opaque-object payload path, but the framework's own checkpoint metadata remains a trusted-storage boundary. The loader retains compatibility with the earlier dense envelope version 1; loading that legacy format necessarily retains its original complete-document memory behavior. Gefen versions the semantic and envelope formats independently, while the surrounding on-disk DCP format retains PyTorch's cross-release compatibility policy. + +A finalized Gefen-backed `GefenMuonHybrid` uses a separate `gefen.portable_composite_state` version-1 wrapper for the complete-document in-memory API 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 child global steps and deterministic settings to agree, and covers the complete wrapper with its own digest. DCP composite envelope version 2 instead stores each present child under a derived role namespace using the child's sharded envelope version 2, while parent metadata protects routing, child presence, aggregate limits, and child digests. 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 every child unchanged. Muon-only and Gefen-backup-only Hybrids are supported. AdamW-backed Hybrid rebinding and native same-topology checkpoints are supported, but its portable global path remains negative because AdamW lacks a topology-neutral moment codec for this composite schema. ```python import torch.distributed.checkpoint as dcp @@ -118,11 +120,13 @@ load_portable_dcp( ) ``` -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`, no active or poisoned state offload, 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 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`, no active or poisoned state offload, a complete declared native state variant, and period one for selected or initialized state. Plain Gefen supports replicated, contiguous flattened-element, and narrow DTensor layouts. Block-second-moment state restores exactly across those layouts. Replicated factored-second-moment state restores exactly to a replicated factored target or follows the separately declared directional projection `defined_projection_factored_to_block_live_fp32_target_period_one_v1` to a replicated, flattened, or narrow-DTensor block target; block-to-factored restore is unsupported. The versioned projection expands the row/column state using the live decomposed fp32 operator order, preserves the source step counters, and recompresses momentum for target period one, so it is a defined numerical projection rather than a bit-exact representation conversion. + +GefenMuon portable global state 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. Its narrow-DTensor training and native-checkpoint support do not imply Muon DTensor portable resharding. A ready Gefen-backed Hybrid declaration is the intersection of the present children's canonical-global layout and topology-change guarantees because one composite operation processes every child; per-routed-parameter training declarations remain the union of child claims. Adapters must inspect each role's child contract together with the finalized immutable `optimizer.parameter_routing()` result. 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 is not evidence of live state offload. +Portable global state currently excludes non-period-one initialized state, block-to-factored and unqualified representation conversions, Muon DTensor resharding, 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 portable document or DCP shard is a detached checkpoint artifact and is not evidence of live state offload. ## Quiescent optimizer-state movement and offload @@ -136,9 +140,9 @@ Activation and restore are fail-before-mutation. Activation also rejects persist `atomic_state_movement` is a dynamic instance capability: it is true only while a noncapturable Gefen or GefenMuon instance has a supported live binding and ordinary CPU/CUDA state representation. GefenMuonHybrid remains false at the composite level because it cannot coordinate an atomic transaction across arbitrary backup optimizers. Movement performs no collectives and its fail-before-mutation guarantee is per optimizer instance; a distributed adapter remains responsible for scheduling instances and coordinating rank-level readiness. `state_offload` is likewise a conservative dynamic readiness claim: it is true only when the live exact plain-Gefen instance can safely enter or retain the supported CPU policy, and false for poisoned or excluded configurations. -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. +Plain Gefen declares replicated, flattened-element, and narrow-DTensor training layouts. Its PyTorch rank-local DTensor checkpoint transport is same-topology only, while its separately qualified portable global transport covers the exact block and defined factored-to-block paths above. 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. Unscoped native Parallel-Muon checkpoints separately declare world-size owner redistribution, not arbitrary placement-changing resharding; a finalized explicit scope instead requires the exact saved scope, topology, and native local layout. Portable Muon resharding remains replicated/whole-owner. `GefenMuonHybrid` retains nested child namespaces for native and portable checkpoints and never flattens 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`. +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. Hybrid native loads stage every present child and the finalized routing guard before callback-free child publication. These transport entries report `atomic_load=True`. This is a per-optimizer-instance fail-before-mutation guarantee at the optimizer load boundary, not rollback after process death, a coordinated all-rank commit after backend failure, 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 outside the guarantee. ## Adapter requirements diff --git a/src/gefen/__init__.py b/src/gefen/__init__.py index e66eef9..c4bc41c 100644 --- a/src/gefen/__init__.py +++ b/src/gefen/__init__.py @@ -11,12 +11,17 @@ "Gefen", "GefenMuon", "GefenMuonHybrid", + "CHECKPOINT_STATE_TRANSITION_SCHEMA_VERSION", "CONTRACT_SCHEMA_VERSION", "CANONICAL_STATE_FORMAT_VERSION", "PORTABLE_STATE_FORMAT_VERSION", "IDENTITY_SCHEMA_VERSION", "CheckpointSupport", "CheckpointTransport", + "CheckpointProjectionQualifier", + "CheckpointStateRepresentation", + "CheckpointStateTransition", + "CheckpointStateTransitionKind", "CanonicalStateProvider", "CheckpointProcessGroupBinding", "CodebookProcessGroupBinding", @@ -117,10 +122,15 @@ def __getattr__(name): return getattr(portable_dcp, name) if name in ( + "CHECKPOINT_STATE_TRANSITION_SCHEMA_VERSION", "CONTRACT_SCHEMA_VERSION", "IDENTITY_SCHEMA_VERSION", "CheckpointSupport", "CheckpointTransport", + "CheckpointProjectionQualifier", + "CheckpointStateRepresentation", + "CheckpointStateTransition", + "CheckpointStateTransitionKind", "CanonicalStateProvider", "OptimizerCapabilities", "OptimizerChildContract", diff --git a/src/gefen/_adamw_adapter.py b/src/gefen/_adamw_adapter.py new file mode 100644 index 0000000..42c85a9 --- /dev/null +++ b/src/gefen/_adamw_adapter.py @@ -0,0 +1,416 @@ +"""Private stable-rebinding and atomic-load adapter for an exact AdamW child.""" + +from __future__ import annotations + +from collections import OrderedDict, defaultdict +import math + +import torch + +from gefen.contracts import ParameterLayout, ShardingManifest +from gefen.dtensor import ( + is_exact_dtensor, + resolve_local_tensor, + validate_dtensor_rebinding_plan, +) + + +def _require_exact_adamw(optimizer) -> None: + if type(optimizer) is not torch.optim.AdamW: + raise TypeError("stable Hybrid AdamW integration requires an exact torch.optim.AdamW child") + if type(optimizer.__dict__) is not dict: + raise TypeError("stable Hybrid AdamW integration requires an exact attribute dictionary") + + +def _target_has_internal_overlap(target) -> bool: + required_span = 1 + dimensions = sorted( + (stride, size) + for size, stride in zip(target.shape, target.stride()) + if size > 1 + ) + for stride, size in dimensions: + if stride < required_span: + return True + required_span += (size - 1) * stride + return False + + +def _validate_storage(target) -> torch.Tensor: + local = resolve_local_tensor(target) if is_exact_dtensor(target) else target + if not isinstance(local, torch.Tensor): + raise TypeError("AdamW rebound targets must be tensors") + if torch.is_complex(local): + raise ValueError("Hybrid AdamW rebinding does not support complex targets") + if local.layout is not torch.strided or local.dtype not in ( + torch.float16, + torch.bfloat16, + torch.float32, + torch.float64, + ): + raise ValueError("Hybrid AdamW rebound targets require strided floating-point storage") + if local.is_meta: + raise ValueError("Hybrid AdamW rebound targets require materialized storage") + if _target_has_internal_overlap(local): + raise ValueError("Hybrid AdamW rebound targets must be free of internal storage overlap") + return local + + +def _validate_rebinding_target(rebinding) -> None: + shard = rebinding.shard + target = rebinding.new_parameter + if rebinding.old_parameter.grad is not None or ( + target is not None and target.grad is not None + ): + raise RuntimeError("Hybrid AdamW post_sharding must run before source or target gradients exist") + owns_whole = ( + shard.layout is ParameterLayout.WHOLE_PARAMETER_OWNER + and shard.local_member == shard.owner + ) + if shard.layout is ParameterLayout.WHOLE_PARAMETER_OWNER and not owns_whole: + if target is not None: + raise ValueError("Hybrid AdamW whole-parameter nonowners must bind no storage") + return + if target is None: + raise ValueError("only a whole-parameter nonowner may bind no AdamW storage") + if not isinstance(target, torch.Tensor): + raise TypeError("Hybrid AdamW rebinding targets must be tensors or None") + local = _validate_storage(target) + if not target.is_leaf and not target.retains_grad: + raise ValueError("can't optimize a non-leaf rebound AdamW Tensor") + + if shard.layout is ParameterLayout.REPLICATED: + if tuple(target.shape) != shard.parameter.global_shape: + raise ValueError("replicated Hybrid AdamW rebinding requires complete parameter storage") + logical_numel = shard.logical_slice.length + physical_numel = target.numel() + elif shard.layout is ParameterLayout.FLATTENED_ELEMENT_SHARD: + if target.ndim != 1 or not target.is_contiguous(): + raise ValueError("flattened Hybrid AdamW rebinding requires a contiguous 1-D tensor shard") + logical_numel = shard.logical_slice.length + physical_numel = target.numel() + elif shard.layout is ParameterLayout.WHOLE_PARAMETER_OWNER: + if tuple(target.shape) != shard.parameter.global_shape: + raise ValueError("Hybrid AdamW whole-parameter owners require complete parameter storage") + logical_numel = shard.logical_slice.length + physical_numel = target.numel() + elif shard.layout is ParameterLayout.DTENSOR_1D_DEFAULT_WORLD: + if tuple(target.shape) != shard.parameter.global_shape: + raise ValueError("Hybrid AdamW DTensor global shape differs from its identity") + logical_numel = shard.logical_region.numel + physical_numel = local.numel() + else: + raise ValueError("Hybrid AdamW rebinding does not support this parameter layout") + if physical_numel != logical_numel: + raise ValueError("Hybrid AdamW rebound storage size differs from its logical shard") + + +def stage_adamw_post_sharding(optimizer, rebindings, manifest: ShardingManifest): + """Return an isolated pristine AdamW shadow with rebound parameter slots.""" + + _require_exact_adamw(optimizer) + rebindings = tuple(rebindings) + if not rebindings: + raise ValueError("Hybrid AdamW child has no routed rebinding") + validate_dtensor_rebinding_plan(rebindings, manifest) + + live_slots = [] + for group_index, group in enumerate(optimizer.param_groups): + if type(group) is not dict or type(group.get("params")) is not list: + raise TypeError("Hybrid AdamW rebinding requires exact parameter-group containers") + names = group.get("param_names") + if type(names) is not list or len(names) != len(group["params"]): + raise ValueError("Hybrid AdamW rebinding requires one stable name per parameter slot") + if any(type(name) is not str for name in names): + raise TypeError("Hybrid AdamW parameter names must be strings") + for parameter_index, (parameter, name) in enumerate(zip(group["params"], names)): + live_slots.append((group_index, parameter_index, parameter, name)) + + if len(rebindings) != len(live_slots): + raise ValueError("Hybrid AdamW post_sharding requires one rebinding per original slot") + recognized = [slot[2] for slot in live_slots] + [item.old_parameter for item in rebindings] + for parameter, state in optimizer.state.items(): + if not any(parameter is candidate for candidate in recognized): + raise RuntimeError("Hybrid AdamW rebinding found orphan optimizer state") + if type(state) is not dict or state: + raise RuntimeError("Hybrid AdamW rebinding is allowed only before optimizer state initialization") + + assigned_positions = {} + seen_sources = [] + seen_targets = [] + for binding_index, rebinding in enumerate(rebindings): + if not isinstance(rebinding.old_parameter, torch.Tensor): + raise TypeError("Hybrid AdamW rebinding sources must be tensors") + if any(rebinding.old_parameter is source for source in seen_sources): + raise ValueError("Hybrid AdamW rebinding sources must be unique") + seen_sources.append(rebinding.old_parameter) + if rebinding.shard not in manifest.shards: + raise ValueError("Hybrid AdamW local shard identity is absent from its child manifest") + _validate_rebinding_target(rebinding) + target = rebinding.new_parameter + if target is not None: + if any(target is previous for previous in seen_targets): + raise ValueError("Hybrid AdamW rebound targets must be unique") + seen_targets.append(target) + + source_positions = [ + index for index, slot in enumerate(live_slots) if slot[2] is rebinding.old_parameter + ] + target_positions = ( + [] + if target is None + else [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("Hybrid AdamW rebound target occupies another optimizer slot") + elif not source_positions and len(target_positions) == 1: + position = target_positions[0] + else: + raise ValueError("Hybrid AdamW source must identify exactly one optimizer slot") + if position in assigned_positions: + raise ValueError("Hybrid AdamW rebindings target the same optimizer slot") + assigned_positions[position] = binding_index + if len(assigned_positions) != len(live_slots): + raise ValueError("Hybrid AdamW post_sharding did not bind every optimizer slot") + + staged_groups = [] + slot_cursor = 0 + for group in optimizer.param_groups: + staged_group = dict(group) + staged_params = [] + staged_names = [] + for _ in group["params"]: + rebinding = rebindings[assigned_positions[slot_cursor]] + compatibility_name = live_slots[slot_cursor][3] + slot_cursor += 1 + if rebinding.new_parameter is not None: + staged_params.append(rebinding.new_parameter) + staged_names.append(compatibility_name) + staged_group["params"] = staged_params + staged_group["param_names"] = staged_names + staged_groups.append(staged_group) + + staged = object.__new__(type(optimizer)) + staged.__dict__ = optimizer.__dict__.copy() + staged.defaults = optimizer.defaults.copy() + staged.param_groups = staged_groups + staged.state = defaultdict(dict) + validate_adamw_state(staged) + return staged + + +def _scalar_value(name, value, *, minimum=0.0, strict_maximum=None) -> float: + if torch.is_tensor(value): + if value.numel() != 1 or value.dtype == torch.bool: + raise ValueError("AdamW {} must be a scalar".format(name)) + scalar = float(value.detach().cpu().item()) + elif type(value) in (int, float): + scalar = float(value) + else: + raise TypeError("AdamW {} must be numeric".format(name)) + if not math.isfinite(scalar) or scalar < minimum: + raise ValueError("AdamW {} is outside its supported range".format(name)) + if strict_maximum is not None and scalar >= strict_maximum: + raise ValueError("AdamW {} is outside its supported range".format(name)) + return scalar + + +def _validate_group(group, defaults, expected_names=None) -> None: + allowed = set(defaults) | {"params", "param_names"} + if type(group) is not dict or set(group) - {"initial_lr"} != allowed: + raise ValueError("AdamW parameter group has an invalid exact schema") + if type(group["params"]) is not list: + raise TypeError("AdamW parameter group params must be a list") + names = group["param_names"] + if type(names) is not list or len(names) != len(group["params"]): + raise ValueError("AdamW parameter group requires one name per parameter") + if any(type(name) is not str for name in names): + raise TypeError("AdamW parameter names must be strings") + if expected_names is not None and tuple(names) != tuple(expected_names): + raise ValueError("AdamW checkpoint parameter names differ from the finalized routing") + _scalar_value("lr", group["lr"]) + if "initial_lr" in group: + _scalar_value("initial_lr", group["initial_lr"]) + betas = group["betas"] + if type(betas) is not tuple or len(betas) != 2: + raise ValueError("AdamW betas must be an exact pair") + _scalar_value("beta1", betas[0], strict_maximum=1.0) + _scalar_value("beta2", betas[1], strict_maximum=1.0) + _scalar_value("eps", group["eps"]) + _scalar_value("weight_decay", group["weight_decay"]) + for key in ("amsgrad", "maximize", "capturable", "differentiable"): + if type(group[key]) is not bool: + raise TypeError("AdamW {} must be a bool".format(key)) + for key in ("foreach", "fused"): + if group[key] is not None and type(group[key]) is not bool: + raise TypeError("AdamW {} must be a bool or None".format(key)) + if "decoupled_weight_decay" in defaults and group["decoupled_weight_decay"] is not True: + raise ValueError("AdamW requires decoupled_weight_decay=True") + + +def _moment_local(moment, parameter, name): + if is_exact_dtensor(parameter): + if not is_exact_dtensor(moment): + raise TypeError("AdamW {} must be a DTensor for a DTensor parameter".format(name)) + if tuple(moment.shape) != tuple(parameter.shape): + raise ValueError("AdamW {} global shape differs from its parameter".format(name)) + moment_mesh = moment.device_mesh + parameter_mesh = parameter.device_mesh + moment_mesh_signature = ( + str(moment_mesh.device_type), + tuple(moment_mesh.shape), + tuple(moment_mesh.mesh_dim_names or ()), + tuple(int(rank) for rank in moment_mesh.mesh.detach().cpu().reshape(-1).tolist()), + ) + parameter_mesh_signature = ( + str(parameter_mesh.device_type), + tuple(parameter_mesh.shape), + tuple(parameter_mesh.mesh_dim_names or ()), + tuple(int(rank) for rank in parameter_mesh.mesh.detach().cpu().reshape(-1).tolist()), + ) + if moment_mesh_signature != parameter_mesh_signature or tuple(moment.placements) != tuple(parameter.placements): + raise ValueError("AdamW {} DTensor layout differs from its parameter".format(name)) + local = resolve_local_tensor(moment) + parameter_local = resolve_local_tensor(parameter) + if tuple(local.shape) != tuple(parameter_local.shape): + raise ValueError("AdamW {} local shape differs from its parameter".format(name)) + else: + if type(moment) is not torch.Tensor: + raise TypeError("AdamW {} must be an ordinary tensor".format(name)) + local = moment + parameter_local = parameter + if tuple(local.shape) != tuple(parameter_local.shape): + raise ValueError("AdamW {} shape differs from its parameter".format(name)) + if local.layout is not torch.strided or local.dtype != parameter_local.dtype or local.device != parameter_local.device: + raise ValueError("AdamW {} storage semantics differ from its parameter".format(name)) + if not bool(torch.isfinite(local.detach()).all()): + raise ValueError("AdamW {} must contain only finite values".format(name)) + return local + + +def _validate_parameter_state(parameter, state, group) -> None: + if type(state) is not dict: + raise TypeError("AdamW parameter state must be an exact dictionary") + if not state: + return + expected = {"step", "exp_avg", "exp_avg_sq"} + if group["amsgrad"]: + expected.add("max_exp_avg_sq") + if set(state) != expected: + raise ValueError("AdamW parameter state has an invalid exact schema") + step = state["step"] + if ( + type(step) is not torch.Tensor + or step.ndim != 0 + or step.dtype not in {torch.float32, torch.float64} + or step.layout is not torch.strided + ): + raise ValueError("AdamW step must be a floating-point scalar tensor") + if group["capturable"] or group["fused"]: + if step.device != parameter.device: + raise ValueError("capturable or fused AdamW step must live with its parameter") + elif step.device.type != "cpu": + raise ValueError("ordinary AdamW step must use CPU scalar storage") + step_value = float(step.detach().cpu().item()) + if not math.isfinite(step_value) or step_value <= 0 or step_value != math.floor(step_value): + raise ValueError("AdamW step must be a positive finite integer-valued scalar") + _moment_local(state["exp_avg"], parameter, "exp_avg") + exp_avg_sq = _moment_local(state["exp_avg_sq"], parameter, "exp_avg_sq") + if bool((exp_avg_sq.detach() < 0).any()): + raise ValueError("AdamW exp_avg_sq must be nonnegative") + if group["amsgrad"]: + maximum = _moment_local(state["max_exp_avg_sq"], parameter, "max_exp_avg_sq") + if bool((maximum.detach() < exp_avg_sq.detach()).any()): + raise ValueError("AdamW max_exp_avg_sq must dominate exp_avg_sq") + + +def validate_adamw_state(optimizer, *, expected_names=None) -> None: + """Validate exact live AdamW group and per-parameter state semantics.""" + + _require_exact_adamw(optimizer) + if type(optimizer.defaults) is not dict or type(optimizer.param_groups) is not list: + raise TypeError("AdamW defaults and parameter groups must use exact containers") + if expected_names is not None and len(expected_names) != len(optimizer.param_groups): + raise ValueError("AdamW finalized group count differs from its routing") + live_parameters = [] + group_by_parameter = {} + for index, group in enumerate(optimizer.param_groups): + _validate_group( + group, + optimizer.defaults, + None if expected_names is None else expected_names[index], + ) + for parameter in group["params"]: + if not isinstance(parameter, torch.Tensor): + raise TypeError("AdamW parameter groups must contain tensors") + parameter_id = id(parameter) + if parameter_id in group_by_parameter: + raise ValueError("AdamW parameter groups contain a duplicate parameter") + live_parameters.append(parameter) + group_by_parameter[parameter_id] = (parameter, group) + for parameter, state in optimizer.state.items(): + entry = group_by_parameter.get(id(parameter)) + if entry is None or entry[0] is not parameter: + raise ValueError("AdamW state contains an orphan parameter entry") + _validate_parameter_state(parameter, state, entry[1]) + + +def prepare_adamw_load_state_dict(optimizer, state_dict): + """Run AdamW pre-hooks and stage a validated load on an isolated shadow.""" + + _require_exact_adamw(optimizer) + expected_names = tuple(tuple(group["param_names"]) for group in optimizer.param_groups) + prepared = state_dict.copy() + for pre_hook in optimizer._optimizer_load_state_dict_pre_hooks.values(): + hook_result = pre_hook(optimizer, prepared) + if hook_result is not None: + prepared = hook_result + + staged = object.__new__(type(optimizer)) + staged.__dict__ = optimizer.__dict__.copy() + staged.defaults = optimizer.defaults.copy() + staged.param_groups = [ + { + **group, + "params": list(group["params"]), + "param_names": list(group["param_names"]), + } + for group in optimizer.param_groups + ] + staged.state = defaultdict(dict) + staged._optimizer_load_state_dict_pre_hooks = OrderedDict() + staged._optimizer_load_state_dict_post_hooks = OrderedDict() + torch.optim.Optimizer.load_state_dict(staged, prepared) + staged._optimizer_load_state_dict_pre_hooks = optimizer._optimizer_load_state_dict_pre_hooks + staged._optimizer_load_state_dict_post_hooks = optimizer._optimizer_load_state_dict_post_hooks + validate_adamw_state(staged, expected_names=expected_names) + return staged + + +def commit_adamw_stage(optimizer, staged) -> None: + """Publish a prepared AdamW rebinding/load through callback-free swaps.""" + + _require_exact_adamw(optimizer) + _require_exact_adamw(staged) + live_defaults = optimizer.defaults + dict.update(live_defaults, staged.defaults) + staged.defaults = live_defaults + dict.update(optimizer.__dict__, staged.__dict__) + + +def run_adamw_load_post_hooks(optimizer) -> None: + """Run AdamW load post-hooks after the composite core is fully published.""" + + for post_hook in optimizer._optimizer_load_state_dict_post_hooks.values(): + post_hook(optimizer) + + +__all__ = [ + "commit_adamw_stage", + "prepare_adamw_load_state_dict", + "run_adamw_load_post_hooks", + "stage_adamw_post_sharding", + "validate_adamw_state", +] diff --git a/src/gefen/contracts.py b/src/gefen/contracts.py index 37144fa..188764c 100644 --- a/src/gefen/contracts.py +++ b/src/gefen/contracts.py @@ -22,6 +22,7 @@ CONTRACT_SCHEMA_VERSION = 1 IDENTITY_SCHEMA_VERSION = 1 +CHECKPOINT_STATE_TRANSITION_SCHEMA_VERSION = 1 class StateScope(str, Enum): @@ -105,6 +106,28 @@ class CheckpointTransport(str, Enum): CANONICAL_GLOBAL = "canonical_global" +class CheckpointStateRepresentation(str, Enum): + """Persistent optimizer-state representation at one checkpoint endpoint.""" + + BLOCK_SECOND_MOMENT = "block_second_moment" + FACTORED_SECOND_MOMENT = "factored_second_moment" + + +class CheckpointStateTransitionKind(str, Enum): + """Numerical relationship between source and target checkpoint state.""" + + EXACT = "exact" + DEFINED_PROJECTION = "defined_projection" + + +class CheckpointProjectionQualifier(str, Enum): + """Versioned numerical rule for a defined checkpoint projection.""" + + FACTORED_TO_BLOCK_LIVE_FP32_TARGET_PERIOD_ONE_V1 = ( + "defined_projection_factored_to_block_live_fp32_target_period_one_v1" + ) + + class TopologyChange(str, Enum): """Specific topology mutation supported by a checkpoint transport.""" @@ -949,10 +972,118 @@ def __post_init__(self) -> None: raise TypeError("TrainingSupport.{} must be a bool".format(name)) +@dataclass(frozen=True) +class CheckpointStateTransition: + """One directional, representation-qualified restore path. + + Exact transitions refine the legacy layout sets on ``CheckpointSupport``. Defined projections are additive qualified paths and therefore do not broaden those legacy exact-restore sets. + """ + + source: CheckpointStateRepresentation + target: CheckpointStateRepresentation + kind: CheckpointStateTransitionKind + same_topology: AbstractSet[ParameterLayout] + topology_changing: AbstractSet[ParameterLayout] + topology_change_kinds: AbstractSet[TopologyChange] = frozenset() + qualifier: Optional[CheckpointProjectionQualifier] = None + schema_version: int = CHECKPOINT_STATE_TRANSITION_SCHEMA_VERSION + + @property + def target_layouts(self) -> FrozenSet[ParameterLayout]: + """Return every target layout covered by this directional path.""" + + return self.same_topology | self.topology_changing + + 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), + ) + if not isinstance(self.source, CheckpointStateRepresentation): + raise TypeError( + "CheckpointStateTransition.source must be a CheckpointStateRepresentation" + ) + if not isinstance(self.target, CheckpointStateRepresentation): + raise TypeError( + "CheckpointStateTransition.target must be a CheckpointStateRepresentation" + ) + if not isinstance(self.kind, CheckpointStateTransitionKind): + raise TypeError( + "CheckpointStateTransition.kind must be a CheckpointStateTransitionKind" + ) + if any( + not isinstance(layout, ParameterLayout) + for layout in self.same_topology | self.topology_changing + ): + raise TypeError( + "CheckpointStateTransition layouts must be ParameterLayout values" + ) + if any( + not isinstance(change, TopologyChange) + for change in self.topology_change_kinds + ): + raise TypeError( + "CheckpointStateTransition.topology_change_kinds must contain TopologyChange values" + ) + if not self.same_topology and not self.topology_changing: + raise ValueError( + "CheckpointStateTransition must declare at least one target layout" + ) + if bool(self.topology_changing) != bool(self.topology_change_kinds): + raise ValueError( + "transition topology-changing layouts and change kinds must be declared together" + ) + if ( + type(self.schema_version) is not int + or self.schema_version != CHECKPOINT_STATE_TRANSITION_SCHEMA_VERSION + ): + raise ValueError("unsupported checkpoint state-transition schema version") + if self.kind is CheckpointStateTransitionKind.EXACT: + if self.source is not self.target: + raise ValueError( + "exact checkpoint state transitions require identical representations" + ) + if self.qualifier is not None: + raise ValueError( + "exact checkpoint state transitions must not carry a projection qualifier" + ) + else: + if self.source is self.target: + raise ValueError( + "defined checkpoint projections require distinct representations" + ) + if not isinstance(self.qualifier, CheckpointProjectionQualifier): + raise TypeError( + "defined checkpoint projections require a CheckpointProjectionQualifier" + ) + if ( + self.qualifier + is CheckpointProjectionQualifier.FACTORED_TO_BLOCK_LIVE_FP32_TARGET_PERIOD_ONE_V1 + and ( + self.source + is not CheckpointStateRepresentation.FACTORED_SECOND_MOMENT + or self.target + is not CheckpointStateRepresentation.BLOCK_SECOND_MOMENT + or self.kind + is not CheckpointStateTransitionKind.DEFINED_PROJECTION + ) + ): + raise ValueError( + "factored-to-block live-fp32 projection qualifier requires its declared direction and kind" + ) + + @dataclass(frozen=True) class CheckpointSupport: """One transport's separately qualified checkpoint capability. + The legacy layout sets describe exact same-representation restores. ``state_transitions`` refines those exact claims and adds separately qualified projection paths without broadening the legacy sets. + ``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 @@ -968,6 +1099,7 @@ class CheckpointSupport: required_sharded_modes: AbstractSet[str] = frozenset() requires_collective: bool = False atomic_load: bool = False + state_transitions: Sequence[CheckpointStateTransition] = () def __post_init__(self) -> None: object.__setattr__(self, "same_topology", _frozenset(self.same_topology)) @@ -982,6 +1114,7 @@ def __post_init__(self) -> None: "required_sharded_modes", _frozenset(self.required_sharded_modes), ) + object.__setattr__(self, "state_transitions", _tuple(self.state_transitions)) 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) @@ -995,6 +1128,54 @@ def __post_init__(self) -> None: raise ValueError( "topology-changing layouts and change kinds must be declared together" ) + if any( + not isinstance(transition, CheckpointStateTransition) + for transition in self.state_transitions + ): + raise TypeError( + "CheckpointSupport.state_transitions must contain CheckpointStateTransition values" + ) + transition_keys = tuple( + ( + transition.source, + transition.target, + transition.kind, + transition.qualifier, + transition.schema_version, + ) + for transition in self.state_transitions + ) + if len(set(transition_keys)) != len(transition_keys): + raise ValueError( + "CheckpointSupport state-transition declarations must be unique" + ) + exact = tuple( + transition + for transition in self.state_transitions + if transition.kind is CheckpointStateTransitionKind.EXACT + ) + if exact: + exact_same = frozenset( + layout for transition in exact for layout in transition.same_topology + ) + exact_changing = frozenset( + layout + for transition in exact + for layout in transition.topology_changing + ) + exact_change_kinds = frozenset( + change + for transition in exact + for change in transition.topology_change_kinds + ) + if ( + exact_same != self.same_topology + or exact_changing != self.topology_changing + or exact_change_kinds != self.topology_change_kinds + ): + raise ValueError( + "exact state transitions must cover the CheckpointSupport layout and topology-change declarations" + ) @dataclass(frozen=True) @@ -1367,6 +1548,7 @@ def _gefen_contract( canonical_global_same_topology: AbstractSet[ParameterLayout] = frozenset(), canonical_global_topology_changing: AbstractSet[ParameterLayout] = frozenset(), canonical_global_topology_change_kinds: AbstractSet[TopologyChange] = frozenset(), + canonical_global_state_transitions: Sequence[CheckpointStateTransition] = (), atomic_state_movement: bool = False, state_offload: bool = False, ) -> OptimizerContract: @@ -1376,6 +1558,7 @@ def _gefen_contract( canonical_global_topology_change_kinds = _frozenset( canonical_global_topology_change_kinds ) + canonical_global_state_transitions = _tuple(canonical_global_state_transitions) block_fields = ( StateField("vmean", StateScope.PARAMETER, StateGeometry.BLOCK, True), StateField("vmean_step", StateScope.PARAMETER, StateGeometry.SCALAR, True), @@ -1567,6 +1750,7 @@ def _gefen_contract( topology_change_kinds=canonical_global_topology_change_kinds, requires_collective=True, atomic_load=True, + state_transitions=canonical_global_state_transitions, ) ) return OptimizerContract( @@ -1621,6 +1805,7 @@ def _gefen_muon_contract( stable_shard_identity: bool = False, explicit_process_group_codebook_scope: bool = False, whole_parameter_owner: bool = False, + native_dtensor_world_size_redistribution: bool = True, canonical_state_layouts: AbstractSet[ParameterLayout] = frozenset(), canonical_global_same_topology: AbstractSet[ParameterLayout] = frozenset(), canonical_global_topology_changing: AbstractSet[ParameterLayout] = frozenset(), @@ -1894,14 +2079,21 @@ def _gefen_muon_contract( ) ) if sharded_modes == frozenset({"distributed"}): + native_dtensor_topology_changing = ( + frozenset({_DTENSOR_LAYOUT}) + if native_dtensor_world_size_redistribution + else frozenset() + ) checkpoints.append( CheckpointSupport( CheckpointTransport.NATIVE_OPTIMIZER, frozenset({_DTENSOR_LAYOUT}), - frozenset({_DTENSOR_LAYOUT}), + native_dtensor_topology_changing, ProcessGroupScope.INFERRED_DEVICE_MESH, - topology_change_kinds=frozenset( - {TopologyChange.WORLD_SIZE_OWNER_REDISTRIBUTION} + topology_change_kinds=( + frozenset({TopologyChange.WORLD_SIZE_OWNER_REDISTRIBUTION}) + if native_dtensor_topology_changing + else frozenset() ), mesh_dimensions=(1,), required_sharded_modes=frozenset({"distributed"}), @@ -1959,6 +2151,7 @@ def _hybrid_training( muon: Optional[OptimizerContract], backup: Optional[OptimizerContract], backup_present: bool, + adamw_backup: bool, ) -> Tuple[TrainingSupport, ...]: """Compose hybrid training claims from the routed children's own claims. @@ -1966,9 +2159,8 @@ def _hybrid_training( 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 without a contract of its own (AdamW) contributes only the layouts + that the Hybrid's exact-type adapter validates directly. """ child_claims = [] @@ -1977,9 +2169,30 @@ def _hybrid_training( if backup is not None: child_claims.append(backup.capabilities.training) elif backup_present: - child_claims.append( - (TrainingSupport(ParameterLayout.REPLICATED, ProcessGroupScope.NONE),) - ) + backup_claims = [ + TrainingSupport(ParameterLayout.REPLICATED, ProcessGroupScope.NONE) + ] + if adamw_backup: + backup_claims.extend( + ( + TrainingSupport( + ParameterLayout.FLATTENED_ELEMENT_SHARD, + ProcessGroupScope.NONE, + ), + TrainingSupport( + ParameterLayout.WHOLE_PARAMETER_OWNER, + ProcessGroupScope.ADAPTER_DEFINED, + requires_complete_parameter_storage=True, + requires_post_step_parameter_sync=True, + ), + TrainingSupport( + _DTENSOR_LAYOUT, + ProcessGroupScope.INFERRED_DEVICE_MESH, + mesh_dimensions=(1,), + ), + ) + ) + child_claims.append(tuple(backup_claims)) training = [] for claims in child_claims: for support in claims: @@ -1993,6 +2206,8 @@ def _hybrid_contract( muon: Optional[OptimizerContract], backup: Optional[OptimizerContract], backup_implementation: str, + adamw_backup: bool = False, + native_finalized_layouts: AbstractSet[ParameterLayout] = frozenset(), canonical_parameter_fqns: bool = False, stable_shard_identity: bool = False, explicit_process_group_codebook_scope: bool = False, @@ -2009,13 +2224,36 @@ def _hybrid_contract( StateGeometry.OPAQUE, True, ), + StateField( + "finalized_binding", + StateScope.OPTIMIZER_COMMON, + StateGeometry.OPAQUE, + True, + optional=True, + description="Versioned exact shard-binding, routing, and child-slot guard for finalized native checkpoints.", + ), ) 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)) + training = _hybrid_training( + muon, + backup, + bool(backup_implementation), + adamw_backup, + ) + native_finalized_layouts = _frozenset(native_finalized_layouts) + native_local_layouts = frozenset({ParameterLayout.REPLICATED}) | ( + native_finalized_layouts + & frozenset( + { + ParameterLayout.FLATTENED_ELEMENT_SHARD, + ParameterLayout.WHOLE_PARAMETER_OWNER, + } + ) + ) canonical_global_same_topology = _frozenset(canonical_global_same_topology) canonical_global_topology_changing = _frozenset(canonical_global_topology_changing) canonical_global_topology_change_kinds = _frozenset( @@ -2024,11 +2262,24 @@ def _hybrid_contract( checkpoints = [ CheckpointSupport( CheckpointTransport.COMPOSITE_NATIVE, - frozenset({ParameterLayout.REPLICATED}), + native_local_layouts, frozenset(), ProcessGroupScope.NONE, + atomic_load=True, ), ] + if _DTENSOR_LAYOUT in native_finalized_layouts: + checkpoints.append( + CheckpointSupport( + CheckpointTransport.COMPOSITE_NATIVE, + frozenset({_DTENSOR_LAYOUT}), + frozenset(), + ProcessGroupScope.DEFAULT_WORLD, + mesh_dimensions=(1,), + requires_collective=True, + atomic_load=True, + ) + ) if canonical_global_same_topology or canonical_global_topology_changing: checkpoints.append( CheckpointSupport( @@ -2066,8 +2317,13 @@ def _hybrid_contract( __all__ = [ + "CHECKPOINT_STATE_TRANSITION_SCHEMA_VERSION", "CONTRACT_SCHEMA_VERSION", "IDENTITY_SCHEMA_VERSION", + "CheckpointProjectionQualifier", + "CheckpointStateRepresentation", + "CheckpointStateTransition", + "CheckpointStateTransitionKind", "CheckpointSupport", "CheckpointTransport", "CanonicalStateProvider", diff --git a/src/gefen/dtensor.py b/src/gefen/dtensor.py new file mode 100644 index 0000000..dac70a8 --- /dev/null +++ b/src/gefen/dtensor.py @@ -0,0 +1,199 @@ +"""PyTorch DTensor validation for stable optimizer rebinding.""" + +from __future__ import annotations + +from typing import Iterable, Optional, Tuple + +import torch +import torch.distributed as dist +from torch.distributed.tensor import DTensor, Replicate, Shard + +from gefen.contracts import ( + LogicalRegion, + ParameterLayout, + PlacementKind, + ShardIdentity, + ShardingManifest, +) + + +def is_exact_dtensor(value) -> bool: + """Return whether ``value`` is the public, exact PyTorch DTensor type.""" + + return type(value) is DTensor + + +def looks_like_dtensor(value) -> bool: + """Return whether a tensor advertises DTensor-like runtime attributes.""" + + return isinstance(value, DTensor) or ( + isinstance(value, torch.Tensor) + and hasattr(value, "to_local") + and hasattr(value, "placements") + and hasattr(value, "device_mesh") + ) + + +def resolve_local_tensor(value) -> torch.Tensor: + """Resolve an exact DTensor to an ordinary materialized local tensor.""" + + if not is_exact_dtensor(value): + raise TypeError("DTensor rebinding requires an exact torch.distributed.tensor.DTensor") + local = value.to_local() + if hasattr(local, "wait"): + local = local.wait() + if type(local) is not torch.Tensor: + raise TypeError("DTensor rebinding requires an ordinary local torch.Tensor") + if local.layout is not torch.strided: + raise ValueError("DTensor rebinding requires dense strided local storage") + if local.is_meta: + raise ValueError("DTensor rebinding requires materialized local storage") + return local + + +def expected_1d_region(shard: ShardIdentity) -> LogicalRegion: + """Return PyTorch's standard 1-D placement region for ``shard``.""" + + if shard.layout is not ParameterLayout.DTENSOR_1D_DEFAULT_WORLD: + raise ValueError("expected_1d_region requires a DTensor shard identity") + if len(shard.placements) != 1: + raise ValueError("one-dimensional DTensor identity requires one placement") + placement = shard.placements[0] + if placement.kind is PlacementKind.REPLICATE: + return LogicalRegion.full(shard.parameter) + if placement.kind is not PlacementKind.DIMENSION_SHARD: + raise ValueError("one-dimensional DTensor identity requires Shard or Replicate") + dimension = placement.parameter_dimension + if type(dimension) is not int: + raise ValueError("sharded DTensor identity requires a parameter dimension") + global_length = shard.parameter.global_shape[dimension] + chunk = (global_length + placement.parts - 1) // placement.parts + offset = min(placement.coordinate * chunk, global_length) + length = max(0, min(global_length, offset + chunk) - offset) + offsets = [0] * len(shard.parameter.global_shape) + lengths = list(shard.parameter.global_shape) + offsets[dimension] = offset + lengths[dimension] = length + return LogicalRegion(tuple(offsets), tuple(lengths)) + + +def _global_group_ranks(group) -> Tuple[int, ...]: + try: + size = dist.get_world_size(group) + return tuple(dist.get_global_rank(group, rank) for rank in range(size)) + except Exception as exc: + raise ValueError("DTensor rebinding could not resolve the DeviceMesh process group") from exc + + +def _validate_manifest_dtensor_shards(shards: Tuple[ShardIdentity, ...]) -> Tuple[object, str]: + process_group = shards[0].process_group + mesh_axis = shards[0].placements[0].mesh_axis + for shard in shards: + if shard.process_group != process_group: + raise ValueError("DTensor manifest shards must share one stable process-group identity") + if shard.placements[0].mesh_axis != mesh_axis: + raise ValueError("DTensor manifest shards must share one semantic mesh axis") + if shard.logical_region != expected_1d_region(shard): + raise ValueError("DTensor manifest logical region does not match standard one-dimensional chunk geometry") + return process_group, mesh_axis + + +def _validate_live_dtensor(target, shard: ShardIdentity, mesh_axis: str): + if not is_exact_dtensor(target): + raise TypeError("DTensor shard identity requires an exact live DTensor") + if tuple(target.shape) != shard.parameter.global_shape: + raise ValueError("DTensor rebinding global shape does not match its parameter identity") + + local = resolve_local_tensor(target) + if tuple(local.shape) != shard.logical_region.lengths: + raise ValueError("DTensor rebinding local shape does not match its logical region") + if local.numel() != shard.logical_region.numel: + raise ValueError("DTensor rebinding local storage size does not match its logical region") + + placements = tuple(target.placements) + if len(placements) != 1: + raise ValueError("DTensor rebinding requires exactly one live placement") + live_placement = placements[0] + descriptor = shard.placements[0] + if descriptor.kind is PlacementKind.REPLICATE: + if type(live_placement) is not Replicate: + raise ValueError("DTensor rebinding placement does not match the replicated identity") + elif ( + descriptor.kind is not PlacementKind.DIMENSION_SHARD + or type(live_placement) is not Shard + or type(live_placement.dim) is not int + or live_placement.dim != descriptor.parameter_dimension + ): + raise ValueError("DTensor rebinding placement does not match the dimension-shard identity") + + mesh = target.device_mesh + if tuple(mesh.shape) != (descriptor.parts,): + raise ValueError("DTensor rebinding requires one mesh dimension matching the identity parts") + dim_names = tuple(mesh.mesh_dim_names or ()) + if dim_names and dim_names != (mesh_axis,): + raise ValueError("DTensor rebinding DeviceMesh dimension name does not match the semantic mesh axis") + coordinate = mesh.get_coordinate() + if coordinate is None or list(coordinate) != [descriptor.coordinate]: + raise ValueError("DTensor rebinding mesh coordinate does not match the shard identity") + return local, mesh, mesh.get_group() + + +def validate_dtensor_rebinding_plan( + rebindings: Iterable, manifest: ShardingManifest +) -> Optional[Tuple[object, object]]: + """Validate all declared and live DTensors before rebinding publication.""" + + rebindings = tuple(rebindings) + for rebinding in rebindings: + target = rebinding.new_parameter + declared = rebinding.shard.layout is ParameterLayout.DTENSOR_1D_DEFAULT_WORLD + exact = is_exact_dtensor(target) + if looks_like_dtensor(target) and not exact: + raise TypeError("DTensor-like tensor subclasses are not supported by stable rebinding") + if exact != declared: + raise ValueError("live DTensor targets and DTensor shard identities must match exactly") + + manifest_shards = tuple( + shard for shard in manifest.shards if shard.layout is ParameterLayout.DTENSOR_1D_DEFAULT_WORLD + ) + if not manifest_shards: + return None + if not dist.is_available() or not dist.is_initialized(): + raise RuntimeError("DTensor rebinding requires an initialized default distributed process group") + + process_group, mesh_axis = _validate_manifest_dtensor_shards(manifest_shards) + world = dist.get_world_size() + if len(process_group.ordered_members) != world: + raise ValueError("DTensor stable process-group identity must span the default world") + + shared_mesh = None + shared_mesh_group = None + for rebinding in rebindings: + if rebinding.shard.layout is not ParameterLayout.DTENSOR_1D_DEFAULT_WORLD: + continue + descriptor = rebinding.shard.placements[0] + if descriptor.parts != world: + raise ValueError("DTensor placement parts must equal the default-world size") + _, mesh, mesh_group = _validate_live_dtensor(rebinding.new_parameter, rebinding.shard, mesh_axis) + if dist.get_world_size(mesh_group) != world: + raise ValueError("DTensor DeviceMesh process group must span the default world") + if set(_global_group_ranks(mesh_group)) != set(range(world)): + raise ValueError("DTensor DeviceMesh process group must contain exactly the default-world ranks") + mesh_ranks = tuple(int(rank) for rank in mesh.mesh.detach().cpu().reshape(-1).tolist()) + if set(mesh_ranks) != set(range(world)): + raise ValueError("DTensor DeviceMesh must contain exactly the default-world ranks") + if shared_mesh is None: + shared_mesh = mesh + shared_mesh_group = mesh_group + elif mesh is not shared_mesh or mesh_group is not shared_mesh_group: + raise ValueError("all rebound DTensors must share one DeviceMesh and process-group handle") + return shared_mesh, shared_mesh_group + + +__all__ = [ + "expected_1d_region", + "is_exact_dtensor", + "looks_like_dtensor", + "resolve_local_tensor", + "validate_dtensor_rebinding_plan", +] diff --git a/src/gefen/gefen.py b/src/gefen/gefen.py index 1240682..db18a67 100644 --- a/src/gefen/gefen.py +++ b/src/gefen/gefen.py @@ -34,6 +34,10 @@ ) from gefen.codebook import CodebookProcessGroupBinding from gefen.contracts import ( + CheckpointProjectionQualifier, + CheckpointStateRepresentation, + CheckpointStateTransition, + CheckpointStateTransitionKind, LogicalSlice, OptimizerContract, ParameterIdentity, @@ -47,6 +51,11 @@ TopologyChange, _gefen_contract, ) +from gefen.dtensor import ( + is_exact_dtensor, + resolve_local_tensor, + validate_dtensor_rebinding_plan, +) from gefen.partitioning import find_period_by_block_variance from gefen.portable_identity import ( _parse_shard_identity, @@ -1247,6 +1256,12 @@ def __init__( # per finalized manifest object instead of rehashing every identity on # every scoped step header. self._gefen_manifest_forensics_cache = None + # Stable intent/device for the default-world collectives used by the + # rank-indexed DTensor checkpoint carrier. A finalized optimizer keeps + # these even if user code later corrupts a public parameter container, + # so one rank cannot silently skip the failure-agreement collective. + self._gefen_rank_local_checkpoint_collective = False + self._gefen_rank_local_checkpoint_device = 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 @@ -1255,6 +1270,18 @@ def __init__( super().__init__(self._normalize_param_groups(params), defaults) self._gefen_checkpoint_schema_ready = True self._install_rank_local_checkpoint_schema() + if self._rank_local_checkpoint_context(fail=False) is not None: + exact_dtensors = tuple( + parameter + for group in self.param_groups + for parameter in group["params"] + if is_exact_dtensor(parameter) + ) + if exact_dtensors: + self._gefen_rank_local_checkpoint_collective = True + self._gefen_rank_local_checkpoint_device = resolve_local_tensor( + exact_dtensors[0] + ).device @property def _step_supports_amp_scaling(self) -> bool: @@ -1283,24 +1310,92 @@ def optimizer_contract(self) -> OptimizerContract: canonical_global_same_topology = _portable_runtime_layouts(self) except Exception: canonical_global_same_topology = frozenset() - has_factored_matrix = ( - canonical_global_same_topology + factored_same_topology = frozenset( + slot.shard.layout + for slot in self._gefen_logical_slots + if slot.shard.layout in 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 + and len(slot.shard.parameter.global_shape) == 2 + and slot.shard.layout is ParameterLayout.REPLICATED + ) + block_same_topology = frozenset( + slot.shard.layout + for slot in self._gefen_logical_slots + if slot.shard.layout in canonical_global_same_topology + and ( + not self._factored_v_2d + or len(slot.shard.parameter.global_shape) != 2 + or slot.shard.layout is not ParameterLayout.REPLICATED ) ) + has_factored_matrix = bool(factored_same_topology) + canonical_global_state_transitions = () if canonical_global_same_topology and not has_factored_matrix: canonical_global_topology_changing = frozenset( { ParameterLayout.REPLICATED, ParameterLayout.FLATTENED_ELEMENT_SHARD, + ParameterLayout.DTENSOR_1D_DEFAULT_WORLD, } ) canonical_global_topology_change_kinds = frozenset( {TopologyChange.PLACEMENT_RESHARD} ) + canonical_global_state_transitions = ( + CheckpointStateTransition( + CheckpointStateRepresentation.BLOCK_SECOND_MOMENT, + CheckpointStateRepresentation.BLOCK_SECOND_MOMENT, + CheckpointStateTransitionKind.EXACT, + canonical_global_same_topology, + canonical_global_topology_changing, + topology_change_kinds=canonical_global_topology_change_kinds, + ), + ) + elif canonical_global_same_topology: + canonical_global_topology_changing = frozenset() + canonical_global_topology_change_kinds = frozenset() + projection_topology_changing = frozenset( + { + ParameterLayout.REPLICATED, + ParameterLayout.FLATTENED_ELEMENT_SHARD, + ParameterLayout.DTENSOR_1D_DEFAULT_WORLD, + } + ) + canonical_global_state_transitions = ( + CheckpointStateTransition( + CheckpointStateRepresentation.FACTORED_SECOND_MOMENT, + CheckpointStateRepresentation.FACTORED_SECOND_MOMENT, + CheckpointStateTransitionKind.EXACT, + factored_same_topology, + frozenset(), + ), + *( + ( + CheckpointStateTransition( + CheckpointStateRepresentation.BLOCK_SECOND_MOMENT, + CheckpointStateRepresentation.BLOCK_SECOND_MOMENT, + CheckpointStateTransitionKind.EXACT, + block_same_topology, + frozenset(), + ), + ) + if block_same_topology + else () + ), + CheckpointStateTransition( + CheckpointStateRepresentation.FACTORED_SECOND_MOMENT, + CheckpointStateRepresentation.BLOCK_SECOND_MOMENT, + CheckpointStateTransitionKind.DEFINED_PROJECTION, + frozenset({ParameterLayout.REPLICATED}), + projection_topology_changing, + topology_change_kinds=frozenset( + {TopologyChange.PLACEMENT_RESHARD} + ), + qualifier=( + CheckpointProjectionQualifier.FACTORED_TO_BLOCK_LIVE_FP32_TARGET_PERIOD_ONE_V1 + ), + ), + ) else: canonical_global_topology_changing = frozenset() canonical_global_topology_change_kinds = frozenset() @@ -1313,6 +1408,7 @@ def optimizer_contract(self) -> OptimizerContract: 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, + canonical_global_state_transitions=canonical_global_state_transitions, atomic_state_movement=self._atomic_state_movement_supported(), state_offload=self._state_offload_supported(), native_flattened_checkpoint=( @@ -1479,6 +1575,8 @@ def _layout_forensics_fast_tokens(self): self._gefen_local_shard_bindings, self._gefen_shard_bindings, len(self._gefen_shard_bindings), + self._gefen_rank_local_checkpoint_collective, + self._gefen_rank_local_checkpoint_device, self._param_names, len(self._param_names), self.param_groups, @@ -1712,6 +1810,29 @@ def _finalized_binding_layout_matches_full(self) -> bool: or parameter_state.get("name") != expected_name ): return False + dtensor_plan = validate_dtensor_rebinding_plan( + tuple( + ParameterRebinding(parameter, parameter, shard) + for parameter, shard in self._gefen_local_shard_bindings + ), + self._gefen_sharding_manifest, + ) + if self._gefen_rank_local_checkpoint_collective != ( + dtensor_plan is not None + ): + return False + if dtensor_plan is not None: + dtensor_devices = { + resolve_local_tensor(parameter).device + for parameter, shard in self._gefen_local_shard_bindings + if shard.layout is ParameterLayout.DTENSOR_1D_DEFAULT_WORLD + } + if ( + len(dtensor_devices) != 1 + or self._gefen_rank_local_checkpoint_device + != next(iter(dtensor_devices)) + ): + return False return True except ( AttributeError, @@ -2182,6 +2303,7 @@ def _normalize_serialized_native_local_shards_v2(cls, value): ParameterLayout.REPLICATED, ParameterLayout.FLATTENED_ELEMENT_SHARD, ParameterLayout.WHOLE_PARAMETER_OWNER, + ParameterLayout.DTENSOR_1D_DEFAULT_WORLD, } for record in records: if type(record) is not dict or set(record) != { @@ -2288,15 +2410,18 @@ def _native_local_shards_matches_live(self, value): def _parameter_in(parameters, candidate) -> bool: return any(item is candidate for item in parameters) - @staticmethod - def _assert_rebound_storage_disjoint(parameters) -> None: + @classmethod + def _assert_rebound_storage_disjoint(cls, parameters) -> None: storage_ranges = [] for parameter in parameters: - if parameter.numel() == 0: + local = parameter + if is_exact_dtensor(parameter): + local = resolve_local_tensor(parameter) + if local.numel() == 0: continue - storage = parameter.untyped_storage() - storage_id = (str(parameter.device), storage.data_ptr()) - if not parameter.is_contiguous(): + storage = local.untyped_storage() + storage_id = (str(local.device), storage.data_ptr()) + if not local.is_contiguous(): for other_id, _, _ in storage_ranges: if other_id == storage_id: raise ValueError( @@ -2304,8 +2429,8 @@ def _assert_rebound_storage_disjoint(parameters) -> None: ) storage_ranges.append((storage_id, None, None)) continue - start = parameter.storage_offset() * parameter.element_size() - end = start + parameter.numel() * parameter.element_size() + start = local.storage_offset() * local.element_size() + end = start + local.numel() * local.element_size() for other_id, other_start, other_end in storage_ranges: if other_id != storage_id: continue @@ -2446,11 +2571,20 @@ def _validate_rebinding_layout(self, rebinding: ParameterRebinding) -> None: "flattened logical matrices require factored_v_2d=False until " "matrix-aware factored-state projection is implemented" ) + elif shard.layout is ParameterLayout.DTENSOR_1D_DEFAULT_WORLD: + local = resolve_local_tensor(target) + if local.numel() != shard.logical_region.numel: + raise ValueError( + "Gefen rebound DTensor local size does not match its logical region" + ) else: raise ValueError( - "plain Gefen rebinding supports replicated or flattened element shards only" + "plain Gefen rebinding supports replicated, flattened element-shard, or one-dimensional DTensor layouts only" ) - if target.numel() != shard.logical_slice.length: + if ( + shard.layout is not ParameterLayout.DTENSOR_1D_DEFAULT_WORLD + and target.numel() != shard.logical_slice.length + ): raise ValueError( "Gefen rebound tensor numel does not match its logical slice" ) @@ -2486,9 +2620,10 @@ def _validate_codebook_process_group_binding( ParameterLayout.REPLICATED, ParameterLayout.FLATTENED_ELEMENT_SHARD, ParameterLayout.WHOLE_PARAMETER_OWNER, + ParameterLayout.DTENSOR_1D_DEFAULT_WORLD, }: raise ValueError( - "explicit codebook scope supports replicated, flattened, or whole-parameter owner identities" + "explicit codebook scope supports replicated, flattened, whole-parameter owner, or one-dimensional DTensor identities" ) self._validate_codebook_runtime_binding(binding) @@ -2560,6 +2695,26 @@ def _assert_runtime_codebook_process_group(self, *, full: bool = False) -> None: self._assert_finalized_binding_layout(full=full) self._validate_codebook_runtime_binding(binding) + @staticmethod + def _codebook_shard_is_replicated(shard: ShardIdentity) -> bool: + return shard.layout is ParameterLayout.REPLICATED or ( + shard.layout is ParameterLayout.DTENSOR_1D_DEFAULT_WORLD + and len(shard.placements) == 1 + and shard.placements[0].kind is PlacementKind.REPLICATE + ) + + def _codebook_parameter_uses_replicated_scope_controls( + self, parameter, shard: ShardIdentity + ) -> bool: + """Whether one member contributes a complete logical parameter. + + The base implementation follows physical replication. Subclasses whose + data plane materializes a complete logical value from another physical + layout may narrow this hook without changing plain Gefen accounting. + """ + + return self._codebook_shard_is_replicated(shard) + def _codebook_parameter_contributes(self, parameter) -> bool: binding = self._gefen_codebook_process_group if binding is None: @@ -2569,7 +2724,9 @@ def _codebook_parameter_contributes(self, parameter) -> bool: raise RuntimeError( "scoped codebook parameter has no finalized shard identity" ) - if shard.layout is ParameterLayout.REPLICATED: + if self._codebook_parameter_uses_replicated_scope_controls( + parameter, shard + ): return shard.local_member == binding.identity.ordered_members[0] return True @@ -2596,6 +2753,7 @@ def _stage_post_sharding(self, rebindings, manifest, codebook_process_group=None raise ValueError( "post_sharding manifest FQNs must exactly match optimizer slots" ) + dtensor_plan = validate_dtensor_rebinding_plan(rebindings, manifest) assigned_positions = {} binding_names = {} @@ -2621,15 +2779,14 @@ def _stage_post_sharding(self, rebindings, manifest, codebook_process_group=None 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): + target_storage = target + if is_exact_dtensor(target): + target_storage = resolve_local_tensor(target) + if torch.is_complex(target_storage): raise ValueError( "Gefen does not support complex rebound parameters" ) - if target.layout is not torch.strided or target.dtype not in ( + if target_storage.layout is not torch.strided or target_storage.dtype not in ( torch.float16, torch.bfloat16, torch.float32, @@ -2638,16 +2795,16 @@ def _stage_post_sharding(self, rebindings, manifest, codebook_process_group=None raise ValueError( "rebound parameters require strided floating-point storage" ) - if target.is_meta: + if target_storage.is_meta: raise ValueError("rebound parameters require materialized storage") if ( rebinding.shard.layout is ParameterLayout.FLATTENED_ELEMENT_SHARD - and not target.is_contiguous() + and not target_storage.is_contiguous() ): raise ValueError( "flattened Gefen rebinding requires contiguous physical storage" ) - if self._target_may_have_internal_storage_overlap(target): + if self._target_may_have_internal_storage_overlap(target_storage): raise ValueError( "rebound target storage must be provably free of internal storage overlap" ) @@ -2780,6 +2937,17 @@ def _stage_post_sharding(self, rebindings, manifest, codebook_process_group=None staged._gefen_logical_slots = tuple(logical_slots) staged._gefen_sharding_manifest = manifest staged._gefen_post_sharding_finalized = True + staged._gefen_rank_local_checkpoint_collective = dtensor_plan is not None + staged._gefen_rank_local_checkpoint_device = ( + next( + resolve_local_tensor(rebinding.new_parameter).device + for rebinding in rebindings + if rebinding.shard.layout + is ParameterLayout.DTENSOR_1D_DEFAULT_WORLD + ) + if dtensor_plan is not None + else None + ) staged._gefen_codebook_process_group = codebook_process_group staged._gefen_codebook_scope_validated = False staged._gefen_codebook_by_device = {} @@ -5296,6 +5464,11 @@ def _codebook_scope_fingerprint(self): repr(self._serialized_codebook_scope()).encode("utf-8") ) + def _codebook_scope_collective_intent_fingerprint(self): + """Return a fixed-size subclass fingerprint for collective routing.""" + + return (0, 0, 0, 0) + def _codebook_value_fingerprint(self): if self._gefen_codebook is None: return (0, 0, 0, 0) @@ -5313,6 +5486,7 @@ def _validate_codebook_scope_operation_header(self, operation: str) -> None: "refresh": 2, "periodic_step": 3, "step": 4, + "checkpoint": 5, } if operation not in operation_codes: raise ValueError("unknown scoped codebook operation") @@ -5330,6 +5504,7 @@ def _validate_codebook_scope_operation_header(self, operation: str) -> None: int(hasattr(self, "found_inf") or hasattr(self, "grad_scale")), *self._codebook_scope_fingerprint(), *self._codebook_manifest_fingerprint(), + *self._codebook_scope_collective_intent_fingerprint(), *self._codebook_value_fingerprint(), # Trailing decision bit, deliberately excluded from the equality # check below: collective-free rank-local operations @@ -5370,10 +5545,20 @@ def _validate_codebook_scope_contribution_controls( } local_records = [] for parameter, shard in self._gefen_local_shard_bindings: + local_extent = ( + shard.logical_region.numel + if shard.layout is ParameterLayout.DTENSOR_1D_DEFAULT_WORLD + else shard.logical_slice.length + ) + replicated_controls = ( + self._codebook_parameter_uses_replicated_scope_controls( + parameter, shard + ) + ) active = int( parameter is not None and parameter.grad is not None - and parameter.numel() > 0 + and (replicated_controls or local_extent > 0) ) period = -1 if active: @@ -5381,13 +5566,16 @@ def _validate_codebook_scope_contribution_controls( period = int(self.state[parameter]["automatic_period"]) else: period = staged_by_parameter[id(parameter)] - layout_code = { - ParameterLayout.REPLICATED: 1, - ParameterLayout.FLATTENED_ELEMENT_SHARD: 2, - ParameterLayout.WHOLE_PARAMETER_OWNER: 3, - }[shard.layout] + if replicated_controls: + layout_code = 1 + else: + layout_code = { + ParameterLayout.FLATTENED_ELEMENT_SHARD: 2, + ParameterLayout.WHOLE_PARAMETER_OWNER: 3, + ParameterLayout.DTENSOR_1D_DEFAULT_WORLD: 4, + }[shard.layout] local_records.append( - (layout_code, active, period, shard.logical_slice.length) + (layout_code, active, period, local_extent) ) header = torch.tensor( @@ -5417,7 +5605,7 @@ def _validate_codebook_scope_contribution_controls( "scoped Gefen codebook local layouts differ across members" ) if layout_code != 1: - if layout_code == 2: + if layout_code in {2, 4}: nonempty_activity = { int(item[offset + 1].item()) for item in gathered @@ -5425,7 +5613,8 @@ def _validate_codebook_scope_contribution_controls( } if len(nonempty_activity) > 1: raise RuntimeError( - "scoped flattened parameters require every nonempty shard to agree on gradient presence" + "scoped flattened and DTensor parameters require every " + "nonempty shard to agree on gradient presence" ) continue replicated_controls = { @@ -5964,6 +6153,9 @@ def _maybe_refresh_gefen_codebook(self) -> None: reuse_existing_periods=self._scope_agreed_resuming_from_checkpoint() ) + def _preflight_codebook_gradient_collectives(self) -> None: + """Validate subclass gradient collectives before manual codebook I/O.""" + @torch.no_grad() def initialize_codebook(self) -> bool: """Collectively initialize the learned codebook without taking a step.""" @@ -5983,6 +6175,7 @@ def initialize_codebook(self) -> bool: self._synchronize_codebook_scope_failure(local_error, "gradient preflight") elif local_error is not None: raise local_error + self._preflight_codebook_gradient_collectives() self._ensure_codebook_scope_agreement() if self._gefen_codebook is not None: return False @@ -6019,6 +6212,7 @@ def refresh_codebook(self) -> bool: self._synchronize_codebook_scope_failure(local_error, "gradient preflight") elif local_error is not None: raise local_error + self._preflight_codebook_gradient_collectives() self._ensure_codebook_scope_agreement() before = self._gefen_codebook self._refresh_codebook_with_requant() @@ -7567,21 +7761,121 @@ def import_portable_state( canonical_state_dict = export_canonical_state load_canonical_state_dict = import_canonical_state + def _rank_local_checkpoint_collective_intent(self) -> bool: + return bool( + getattr(self, "_gefen_rank_local_checkpoint_collective", False) + ) or self._uses_rank_local_sharded_state() + + def _rank_local_checkpoint_collective_device(self) -> torch.device: + import torch.distributed as dist + + backend = str(dist.get_backend()).lower() + if "nccl" not in backend: + return torch.device("cpu") + device = getattr(self, "_gefen_rank_local_checkpoint_device", None) + if device is None: + for group in self.param_groups: + for parameter in group["params"]: + if is_exact_dtensor(parameter): + device = resolve_local_tensor(parameter).device + break + if device is not None: + break + if device is None or device.type != "cuda" or device.index is None: + raise RuntimeError( + "Gefen rank-local NCCL checkpoint agreement requires a concrete CUDA device" + ) + return device + + def _synchronize_rank_local_checkpoint_failure(self, error, phase: str) -> None: + if not self._rank_local_checkpoint_collective_intent(): + if error is not None: + raise error + return + if ( + not torch.distributed.is_available() + or not torch.distributed.is_initialized() + ): + if error is not None: + raise error + raise RuntimeError( + "Gefen rank-local DTensor checkpointing requires initialized torch.distributed" + ) + import torch.distributed as dist + + failed = torch.tensor( + int(error is not None), + dtype=torch.int32, + device=self._rank_local_checkpoint_collective_device(), + ) + dist.all_reduce(failed, op=dist.ReduceOp.MAX) + if int(failed.item()) == 0: + return + if error is not None: + raise RuntimeError( + "Gefen rank-local checkpoint {} failed on this rank: {}".format( + phase, error + ) + ) from error + raise RuntimeError( + "Gefen rank-local checkpoint {} failed on another rank".format(phase) + ) + + def _run_rank_local_checkpoint_phase(self, phase: str, callback): + result = None + error = None + try: + result = callback() + except Exception as exc: + error = exc + self._synchronize_rank_local_checkpoint_failure(error, phase) + return result + def state_dict(self): """Run optimizer state-dict hooks around Gefen's complete schema.""" - self._assert_finalized_binding_layout(full=True) - self._assert_state_export_safe() - for pre_hook in self._optimizer_state_dict_pre_hooks.values(): - pre_hook(self) - self._assert_finalized_binding_layout(full=True) - self._assert_state_export_safe() - state_dict = self._state_dict_impl() - for post_hook in self._optimizer_state_dict_post_hooks.values(): - hook_result = post_hook(self, state_dict) - if hook_result is not None: - state_dict = hook_result - return state_dict + def validate_layout(): + self._assert_finalized_binding_layout(full=True) + self._assert_state_export_safe() + + def run_pre_hooks(): + for pre_hook in self._optimizer_state_dict_pre_hooks.values(): + pre_hook(self) + + self._run_rank_local_checkpoint_phase( + "initial layout validation", validate_layout + ) + self._run_rank_local_checkpoint_phase( + "state-dict pre-hook", run_pre_hooks + ) + self._run_rank_local_checkpoint_phase( + "post-hook layout validation", validate_layout + ) + state_dict = self._run_rank_local_checkpoint_phase( + "local serialization", + lambda: self._state_dict_impl(consolidate_rank_local=False), + ) + if self._rank_local_checkpoint_collective_intent(): + checkpoint_metadata = state_dict["param_groups"][0][ + "_gefen_checkpoint_metadata" + ] + self._consolidate_rank_local_sharded_state( + state_dict, checkpoint_metadata + ) + for group in state_dict["param_groups"]: + group["_gefen_checkpoint_metadata"] = dict(checkpoint_metadata) + + def run_post_hooks(): + result = state_dict + for post_hook in self._optimizer_state_dict_post_hooks.values(): + hook_result = post_hook(self, result) + if hook_result is not None: + result = hook_result + return result + + return self._run_rank_local_checkpoint_phase( + "state-dict post-hook", run_post_hooks + ) def _base_state_dict_without_hooks(self): """Call ``Optimizer.state_dict`` without double-firing public hooks.""" @@ -7766,14 +8060,24 @@ def _install_rank_local_checkpoint_schema(self) -> None: ) for param in params[1:]: self.state[param][_RANK_LOCAL_MEMBER_KEY] = True + codebook_scope = self._serialized_codebook_scope() + native_local_shards = self._serialized_native_local_shards() placeholder = { - "format_version": _RANK_LOCAL_METADATA_VERSION, + "format_version": ( + _SCOPED_NATIVE_METADATA_VERSION + if codebook_scope is not None + else _RANK_LOCAL_METADATA_VERSION + ), "global_step": self._gefen_global_step, "codebook": None, "deterministic": self._deterministic, "device_anchor": self._checkpoint_device_anchor(), "rank_local_sharded_state": {"format": _RANK_LOCAL_FORMAT}, } + if codebook_scope is not None: + placeholder["codebook_scope"] = codebook_scope + if native_local_shards is not None: + placeholder["native_local_shards"] = native_local_shards for group in self.param_groups: group["_gefen_checkpoint_metadata"] = placeholder @@ -7785,51 +8089,62 @@ def _rank_local_sharded_signature(self, context=None): for param_index, (name, p) in enumerate( self._iter_group_params_with_names(group) ): + stable_shard = ( + self._gefen_shard_bindings.get(p) + if self._gefen_post_sharding_finalized + else None + ) identifier = "group_{}_param_{}_{}".format( group_index, param_index, str(name).lower() ) if not self._is_dtensor_parameter(p): - signature.append( - { - "identifier": identifier, - "group": group_index, - "param": param_index, - "name": str(name).lower(), - "sharded": False, - "shape": list(p.shape), - "local_shape": list(p.shape), - "dtype": str(p.dtype), - "requires_grad": bool(p.requires_grad), - "sharded_mode": group.get("sharded_mode"), - } - ) - continue - local = p.to_local() - if hasattr(local, "wait"): - local = local.wait() - mesh = p.device_mesh - coordinate = mesh.get_coordinate() - signature.append( - { + record = { "identifier": identifier, "group": group_index, "param": param_index, "name": str(name).lower(), - "sharded": True, + "sharded": False, "shape": list(p.shape), - "local_shape": list(local.shape), + "local_shape": list(p.shape), "dtype": str(p.dtype), - "local_dtype": str(local.dtype), "requires_grad": bool(p.requires_grad), - "placements": [ - self._placement_checkpoint_signature(item) - for item in p.placements - ], - "mesh": self._device_mesh_checkpoint_signature(mesh), - "coordinate": None if coordinate is None else list(coordinate), "sharded_mode": group.get("sharded_mode"), } - ) + if stable_shard is not None: + record["stable_shard_identity"] = _serialize_shard_identity( + stable_shard + ) + signature.append(record) + continue + local = p.to_local() + if hasattr(local, "wait"): + local = local.wait() + mesh = p.device_mesh + coordinate = mesh.get_coordinate() + record = { + "identifier": identifier, + "group": group_index, + "param": param_index, + "name": str(name).lower(), + "sharded": True, + "shape": list(p.shape), + "local_shape": list(local.shape), + "dtype": str(p.dtype), + "local_dtype": str(local.dtype), + "requires_grad": bool(p.requires_grad), + "placements": [ + self._placement_checkpoint_signature(item) + for item in p.placements + ], + "mesh": self._device_mesh_checkpoint_signature(mesh), + "coordinate": None if coordinate is None else list(coordinate), + "sharded_mode": group.get("sharded_mode"), + } + if stable_shard is not None: + record["stable_shard_identity"] = _serialize_shard_identity( + stable_shard + ) + signature.append(record) return signature @staticmethod @@ -7984,19 +8299,31 @@ def _consolidate_rank_local_sharded_state( import torch.distributed as dist - context = self._rank_local_checkpoint_context() + context = None + saved_ids = None + local_signature = None + local_manifest = None + local_control = None + local_error = None + try: + context = self._rank_local_checkpoint_context() + saved_ids = list(state_dict["state"]) + local_signature = self._rank_local_sharded_signature(context) + local_manifest = self._rank_local_parameter_manifest(local_signature) + local_control = { + "saved_ids": saved_ids, + "manifest": local_manifest, + "signature": local_signature, + "global_step": self._gefen_global_step, + "deterministic": self._deterministic, + } + except Exception as exc: + local_error = exc + self._synchronize_rank_local_checkpoint_failure( + local_error, "collective preparation" + ) world = context["world_size"] rank = context["global_rank"] - saved_ids = list(state_dict["state"]) - local_signature = self._rank_local_sharded_signature(context) - local_manifest = self._rank_local_parameter_manifest(local_signature) - local_control = { - "saved_ids": saved_ids, - "manifest": local_manifest, - "signature": local_signature, - "global_step": self._gefen_global_step, - "deterministic": self._deterministic, - } controls = [None] * world dist.all_gather_object(controls, local_control) if any(control["saved_ids"] != saved_ids for control in controls): @@ -8111,7 +8438,11 @@ def _consolidate_rank_local_sharded_state( "global_step": global_steps[0], "deterministic": deterministic_values[0], } - checkpoint_metadata["format_version"] = _RANK_LOCAL_METADATA_VERSION + checkpoint_metadata["format_version"] = ( + _SCOPED_NATIVE_METADATA_VERSION + if checkpoint_metadata.get("codebook_scope") is not None + else _RANK_LOCAL_METADATA_VERSION + ) checkpoint_metadata["global_step"] = global_steps[0] checkpoint_metadata["deterministic"] = deterministic_values[0] checkpoint_metadata["codebook"] = None @@ -8353,21 +8684,40 @@ 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) + self._run_rank_local_checkpoint_phase( + "initial load layout validation", + lambda: self._assert_finalized_binding_layout(full=True), + ) + + def run_pre_hooks(): + prepared = state_dict.copy() + for pre_hook in self._optimizer_load_state_dict_pre_hooks.values(): + hook_result = pre_hook(self, prepared) + if hook_result is not None: + prepared = hook_result + return prepared + + prepared = self._run_rank_local_checkpoint_phase( + "load pre-hook", run_pre_hooks + ) + self._run_rank_local_checkpoint_phase( + "post-hook load layout validation", + lambda: self._assert_finalized_binding_layout(full=True), + ) + return self._run_rank_local_checkpoint_phase( + "local load staging", lambda: self._stage_load_state_dict(prepared) + ) def _publish_load_state_dict(self, staged): """Commit a prepared restore through non-throwing swaps (phase two).""" self._commit_staged_load_state_dict(staged) - for post_hook in self._optimizer_load_state_dict_post_hooks.values(): - post_hook(self) + + def run_post_hooks(): + for post_hook in self._optimizer_load_state_dict_post_hooks.values(): + post_hook(self) + + self._run_rank_local_checkpoint_phase("load post-hook", run_post_hooks) def _stage_load_state_dict(self, state_dict): """Prepare a complete restore without mutating the live optimizer. diff --git a/src/gefen/gefen_muon.py b/src/gefen/gefen_muon.py index e967794..c7d71cc 100644 --- a/src/gefen/gefen_muon.py +++ b/src/gefen/gefen_muon.py @@ -10,6 +10,7 @@ from gefen.contracts import ( OptimizerContract, ParameterLayout, + PlacementKind, TopologyChange, _gefen_muon_contract, ) @@ -659,8 +660,9 @@ def __init__( # The NS/momentum compute (and the persistent momentum state) is therefore # cut ~world_size x while staying bit-for-bit identical to "exact". Only # the per-step gradient all-gather (a collective every rank must join) and - # one extra update broadcast are replicated. Falls back to the "exact" - # full-NS-everywhere path for non-1D meshes (e.g. HSDP x TP). + # one extra update broadcast are replicated. Multi-rank DTensor support + # is deliberately limited to one-dimensional DeviceMeshes; unsupported + # multidimensional routes are rejected during parameter registration. # # Homogeneity assumption: "exact" runs Newton-Schulz redundantly on # every rank's own GPU and relies on the results agreeing bit-for-bit, @@ -756,10 +758,17 @@ def __init__( capturable=capturable, verbose=verbose, ) + # Keep a non-mutating route anchor from construction onward. A later + # coordinated eager preflight replaces this with the validated baseline + # anchor, but retaining the initial default-world mesh lets that first + # preflight reject rank-asymmetric direct param-group edits with a + # fixed-size collective instead of letting the edited rank skip it. + self._set_muon_collective_intent_anchors_from_live() def optimizer_contract(self) -> OptimizerContract: """Return the immutable Muon state and capability contract.""" + self._assert_live_muon_contract_routing() sharded_modes = frozenset( group["sharded_mode"] for group in self.param_groups ) @@ -806,6 +815,9 @@ def optimizer_contract(self) -> OptimizerContract: canonical_parameter_fqns=self._canonical_identity_ready(), stable_shard_identity=self._canonical_identity_ready(), explicit_process_group_codebook_scope=True, + native_dtensor_world_size_redistribution=( + self._gefen_codebook_process_group is None + ), canonical_state_layouts=canonical_state_layouts, canonical_global_same_topology=canonical_global_same_topology, canonical_global_topology_changing=canonical_global_topology_changing, @@ -824,6 +836,7 @@ def _canonical_state_layout_supported(self, layout) -> bool: return layout is ParameterLayout.REPLICATED def _canonical_state_variant_layout(self): + self._assert_live_muon_contract_routing() sharded_modes = frozenset( group["sharded_mode"] for group in self.param_groups ) @@ -867,12 +880,25 @@ def _validate_rebinding_layout(self, rebinding) -> None: raise ValueError( "a GefenMuon whole-parameter non-owner must not retain storage" ) + elif shard.layout is ParameterLayout.DTENSOR_1D_DEFAULT_WORLD: + # The base staging path has already validated that this is an exact + # public DTensor on one default-world DeviceMesh, that its Shard or + # Replicate placement matches the stable identity, and that its + # local storage is dense, materialized, and the declared region + # size. Reuse the base local-storage gate here while retaining + # Muon's additional logical-matrix requirement above. + super()._validate_rebinding_layout(rebinding) else: raise ValueError( - "GefenMuon rebinding supports replicated complete matrices or " - "whole-parameter ownership only" + "GefenMuon rebinding supports replicated complete matrices, " + "whole-parameter ownership, or one-dimensional default-world " + "DTensor matrices only" ) - if target is not None and target.numel() != shard.logical_slice.length: + if ( + target is not None + and shard.layout is not ParameterLayout.DTENSOR_1D_DEFAULT_WORLD + and target.numel() != shard.logical_slice.length + ): raise ValueError( "GefenMuon rebound tensor numel does not match its logical slice" ) @@ -883,6 +909,20 @@ def _has_unscoped_whole_owner_bindings(self) -> bool: for _, shard in self._gefen_local_shard_bindings ) + def _stage_post_sharding( + self, rebindings, manifest, codebook_process_group=None + ): + staged = super()._stage_post_sharding( + rebindings, manifest, codebook_process_group + ) + # Commits update the live __dict__ rather than deleting absent keys, so + # publish explicit cleared sentinels instead of merely popping the + # staged attributes. + staged._gefen_muon_collective_intent_baseline = None + staged._gefen_muon_capture_plan = None + staged._set_muon_collective_intent_anchors_from_live() + return staged + def add_param_group(self, param_group): """Add a validated 2D Muon parameter group atomically. @@ -895,6 +935,13 @@ def add_param_group(self, param_group): # containers and missing ``params``. return super().add_param_group(param_group) + if getattr(self, "_gefen_muon_collective_intent_baseline", None) is not None: + raise RuntimeError( + "GefenMuon cannot add a parameter group after DTensor collective " + "routing has been validated; construct a new optimizer because " + "live optimizer state has no defined group-addition migration" + ) + group = dict(param_group) raw_params = group["params"] if isinstance(raw_params, torch.Tensor): @@ -906,12 +953,31 @@ def add_param_group(self, param_group): # Validate tensor type/complex support and Muon's dimensionality before # registration. Gefen.add_param_group repeats its own checks, but doing # this first keeps a mixed valid/invalid addition all-or-nothing. + validated_params = [] for _, param in self._iter_params_with_names(raw_params): if param.ndim != 2: raise ValueError( "GefenMuon only supports 2D parameters whereas we found a " "parameter with size: {}".format(param.size()) ) + validated_params.append(param) + for param in validated_params: + if self._is_sharded(param): + self._validate_muon_dtensor_mesh(param.device_mesh) + if ( + "_gefen_muon_collective_intent_anchors" in self.__dict__ + and not self._gefen_muon_collective_intent_anchors + and any( + self._is_sharded(param) + and param.device_mesh.size() > 1 + for param in validated_params + ) + ): + raise RuntimeError( + "GefenMuon cannot add the first multi-rank DTensor parameter " + "group after construction because no common collective routing " + "anchor exists; construct a new optimizer with that group" + ) defaults = self._muon_group_defaults momentum = group.get("momentum", defaults["momentum"]) @@ -1008,7 +1074,17 @@ def add_param_group(self, param_group): "betas": (momentum, 0.0), } ) - return super().add_param_group(group) + result = super().add_param_group(group) + # During Optimizer.__init__ the anchor attribute does not exist yet. + # A pristine public addition may introduce the first DTensor mesh; keep + # an already-established anchor unchanged so an asymmetric addition is + # still visible to the next fixed-size header on the original mesh. + if ( + "_gefen_muon_collective_intent_anchors" in self.__dict__ + and not self._gefen_muon_collective_intent_anchors + ): + self._set_muon_collective_intent_anchors_from_live() + return result def _init_gefen_muon_state(self, state, grad_view: torch.Tensor) -> None: self._init_gefen_state(state, grad_view) @@ -1016,6 +1092,61 @@ def _init_gefen_muon_state(self, state, grad_view: torch.Tensor) -> None: def _codebook_requires_2d_parameters(self) -> bool: return True + def _codebook_parameter_uses_replicated_scope_controls( + self, parameter, shard + ) -> bool: + if super()._codebook_parameter_uses_replicated_scope_controls( + parameter, shard + ): + return True + if ( + parameter is None + or shard.layout is not ParameterLayout.DTENSOR_1D_DEFAULT_WORLD + or len(shard.placements) != 1 + or shard.placements[0].kind is not PlacementKind.DIMENSION_SHARD + ): + return False + for group in self.param_groups: + if any(candidate is parameter for candidate in group["params"]): + return group["sharded_mode"] in {"exact", "distributed"} + return False + + def _preflight_codebook_gradient_collectives(self) -> None: + self._assert_sharded_grad_presence_consistent() + + def _preflight_unscoped_manual_codebook_routing(self) -> None: + if self._gefen_codebook_process_group is not None: + return + if self._restore_muon_capture_plan_if_capturing(): + return + # Base manual codebook APIs begin with local finalized-layout and + # gradient validation. Run the fixed mesh header first so a malformed + # direct group edit cannot make one rank exit while peers continue to a + # later DTensor codebook collective. + self._assert_sharded_mode_collective_intent_consistent() + + @torch.no_grad() + def initialize_codebook(self) -> bool: + self._preflight_unscoped_manual_codebook_routing() + return super().initialize_codebook() + + @torch.no_grad() + def refresh_codebook(self) -> bool: + self._preflight_unscoped_manual_codebook_routing() + return super().refresh_codebook() + + def _codebook_scope_collective_intent_fingerprint(self): + payload = tuple( + ( + slot.group_index, + slot.original_slot_index, + slot.shard.parameter.fqn, + self.param_groups[slot.group_index]["sharded_mode"], + ) + for slot in self._gefen_logical_slots + ) + return self._sha256_int64(repr(payload).encode("utf-8")) + def _iter_gefen_grad_periods( self, reuse_existing_periods: bool = False, staged_periods=None ): @@ -1440,8 +1571,435 @@ def _dist_available() -> bool: return False return torch.distributed.is_initialized() + @staticmethod + def _gradient_preflight_mesh_key(mesh): + process_groups = tuple(mesh.get_all_groups()) + return ( + ( + 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() + ), + tuple(str(pg.group_name) for pg in process_groups), + ), + process_groups, + ) + + @staticmethod + def _gradient_preflight_portable_mesh_key(mesh): + """Return rank-neutral mesh identity for cross-rank intent hashes. + + GefenMuon accepts only one-dimensional multi-rank meshes, whose sole + process group has one shared c10d name on every member. Include that + name so distinct same-topology DeviceMesh instances cannot be confused + across ranks while keeping the actual handle local to scheduling. + """ + + mesh_dim_names = getattr(mesh, "mesh_dim_names", None) + process_groups = tuple(mesh.get_all_groups()) + return ( + 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() + ), + ( + None + if mesh_dim_names is None + else tuple(str(item) for item in mesh_dim_names) + ), + tuple(str(group.group_name) for group in process_groups), + ) + + def _validate_muon_dtensor_mesh(self, mesh) -> None: + """Reject DTensor mesh routes outside Muon's published contract.""" + + if mesh.size() < 2: + return + if len(mesh.shape) != 1: + raise RuntimeError( + "GefenMuon supports multi-rank DTensor collective routing only " + "on a one-dimensional DeviceMesh; multidimensional meshes are " + "not supported" + ) + if ( + self._gefen_codebook_process_group is not None + or not self._dist_available() + ): + return + + import torch.distributed as dist + + mesh_members = tuple( + int(item) for item in mesh.mesh.detach().cpu().reshape(-1).tolist() + ) + world = dist.get_world_size() + if len(mesh_members) != world or set(mesh_members) != set(range(world)): + raise RuntimeError( + "unscoped GefenMuon requires every multi-rank DTensor mesh to " + "span the initialized default process group; subgroup and " + "partially overlapping meshes require an explicit codebook scope" + ) + + @staticmethod + def _muon_mode_token(mode): + if isinstance(mode, str) and mode in {"approx", "exact", "distributed"}: + return mode, True + if mode is None: + return ("invalid", "missing"), False + return ( + "invalid", + "{}.{}".format(type(mode).__module__, type(mode).__qualname__), + repr(mode), + ), False + + def _assert_live_muon_contract_routing(self) -> None: + """Validate live route metadata without collectives or mutation.""" + + invalid_groups = [ + index + for index, group in enumerate(self.param_groups) + if not self._muon_mode_token(group.get("sharded_mode"))[1] + ] + if invalid_groups: + raise RuntimeError( + "GefenMuon cannot report an optimizer contract while parameter " + "groups {} have an invalid sharded_mode; expected 'approx', " + "'exact', or 'distributed'".format(invalid_groups) + ) + intent = self._muon_collective_intent() + baseline = getattr( + self, "_gefen_muon_collective_intent_baseline", None + ) + if baseline is not None and baseline != intent["signature"]: + raise RuntimeError( + "GefenMuon cannot report an optimizer contract from live groups " + "that differ from its frozen DTensor collective routing; " + "construct a new optimizer for parameter-order, rebinding, or " + "sharded_mode transitions" + ) + + def _muon_collective_intent(self): + """Describe live DTensor routing without entering a collective.""" + + by_mesh = OrderedDict() + signature = [] + portable_structure = [] + portable_routing = [] + invalid_modes = [] + for group_index, group in enumerate(self.param_groups): + raw_mode = group.get("sharded_mode") + mode_token, mode_valid = self._muon_mode_token(raw_mode) + for parameter_index, (name, parameter) in enumerate( + self._iter_group_params_with_names(group) + ): + if not self._is_sharded(parameter): + continue + mesh = parameter.device_mesh + self._validate_muon_dtensor_mesh(mesh) + if mesh.get_coordinate() is None or mesh.size() < 2: + continue + mesh_key, process_groups = self._gradient_preflight_mesh_key(mesh) + portable_mesh_key = self._gradient_preflight_portable_mesh_key(mesh) + entry = by_mesh.get(mesh_key) + if entry is None: + entry = { + "process_groups": process_groups, + "device": self._state_tensor_device(parameter), + "items": [], + } + by_mesh[mesh_key] = entry + label = str(name) + shape = tuple(parameter.shape) + dtype = str(parameter.dtype) + placements = tuple(str(item) for item in parameter.placements) + entry["items"].append( + ( + group_index, + parameter_index, + label, + parameter, + raw_mode, + ) + ) + signature.append( + ( + group_index, + parameter_index, + id(parameter), + id(mesh), + label, + shape, + dtype, + placements, + mode_token, + mesh_key, + ) + ) + portable_structure.append( + ( + group_index, + parameter_index, + label, + shape, + dtype, + placements, + portable_mesh_key, + ) + ) + portable_routing.append( + (group_index, parameter_index, label, mode_token) + ) + if not mode_valid: + invalid_modes.append(label) + + # One retained device anchor makes the fixed optimizer-global header a + # single default-world collective even when a later direct edit adds, + # removes, or reorders secondary full-world DeviceMeshes. Distributed + # construction requires every world rank to begin with at least one + # multi-rank DTensor; adding the first one later is rejected separately. + anchors = () + if by_mesh: + mesh_key = next(iter(by_mesh)) + anchors = ( + ( + mesh_key, + by_mesh[mesh_key]["process_groups"], + by_mesh[mesh_key]["device"], + ), + ) + return { + "by_mesh": by_mesh, + "signature": tuple(signature), + "portable_structure": tuple(portable_structure), + "portable_routing": tuple(portable_routing), + "invalid_modes": tuple(invalid_modes), + "anchors": anchors, + } + + def _set_muon_collective_intent_anchors_from_live(self) -> None: + self._gefen_muon_collective_intent_anchors = self._muon_collective_intent()[ + "anchors" + ] + + def _freeze_muon_capture_plan(self, intent) -> None: + groups = [] + for group in self.param_groups: + names = group.get("param_names") + groups.append( + ( + group, + tuple(group["params"]), + None if names is None else tuple(names), + group["sharded_mode"], + ) + ) + self._gefen_muon_capture_plan = ( + tuple(groups), + tuple(self._param_names.items()), + ) + self._gefen_muon_collective_intent_anchors = intent["anchors"] + + def _muon_capture_has_cuda_route(self) -> bool: + plan = getattr(self, "_gefen_muon_capture_plan", None) + if plan is not None: + parameters = ( + parameter + for _, group_parameters, _, _ in plan[0] + for parameter in group_parameters + ) + if any(parameter.device.type == "cuda" for parameter in parameters): + return True + return any( + anchor[2].type == "cuda" + for anchor in getattr( + self, "_gefen_muon_collective_intent_anchors", () + ) + ) + + def _restore_muon_capture_plan_if_capturing(self) -> bool: + """Route capture from the eagerly validated immutable plan.""" + + if not self._muon_capture_has_cuda_route(): + return False + if not torch.cuda.is_current_stream_capturing(): + return False + baseline = getattr(self, "_gefen_muon_collective_intent_baseline", None) + plan = getattr(self, "_gefen_muon_capture_plan", None) + if baseline is None or plan is None: + raise RuntimeError( + "GefenMuon DTensor CUDA graph capture requires an eager warmup " + "step to validate and freeze sharded_mode routing before capture" + ) + + groups, param_names = plan + # Restore route-affecting public containers before finalized-layout + # checks, scoped operation headers, or distributed/regular partitioning. + # No rank branches on its mutable public mode or membership during + # capture; every rank uses the plan established by the eager handshake. + self.param_groups[:] = [record[0] for record in groups] + for group, parameters, names, mode in groups: + group["params"] = list(parameters) + if names is None: + group.pop("param_names", None) + else: + group["param_names"] = list(names) + group["sharded_mode"] = mode + self._param_names = dict(param_names) + return True + + def _muon_intent_digest(self, value): + mask = (1 << 63) - 1 + return tuple( + item & mask + for item in self._sha256_int64(repr(value).encode("utf-8")) + ) + + def _reduce_muon_intent_header( + self, intent, baseline_changed, *, preparation_failed + ): + """Exchange one fixed-size validation header on stable mesh anchors.""" + + import torch.distributed as dist + + local = ( + int(preparation_failed), + int(bool(intent["invalid_modes"])), + len(intent["signature"]), + *self._muon_intent_digest(intent["portable_structure"]), + *self._muon_intent_digest(intent["portable_routing"]), + int(baseline_changed), + ) + encoded = [value for item in local for value in (item, -item)] + anchors = getattr(self, "_gefen_muon_collective_intent_anchors", ()) + if not anchors: + anchors = intent["anchors"] + if not anchors: + return [] + _, _, device = anchors[0] + probe = torch.tensor(encoded, dtype=torch.int64, device=device) + dist.all_reduce(probe, op=dist.ReduceOp.MAX, group=dist.group.WORLD) + values = probe.cpu().tolist() + return [ + tuple( + (values[index], -values[index + 1]) + for index in range(0, len(values), 2) + ) + ] + @torch._dynamo.disable - def _assert_sharded_grad_presence_consistent(self) -> None: + def _assert_sharded_mode_collective_intent_consistent(self) -> None: + """Agree on and freeze each DTensor parameter's collective routing. + + Unscoped operations exchange one constant-size validation/count/hash + header on the initialized default process group. Invalid modes or + asymmetric membership, order, rebinding, and routing therefore fail on + every rank without a variable communicator schedule. The first + successful preflight freezes both eager identity and CUDA-capture + routing; later eager drift is rejected because live state has no defined + migration. + """ + + if not self._dist_available(): + return + try: + intent = self._muon_collective_intent() + preparation_error = None + except Exception as exc: + preparation_error = exc + intent = { + "by_mesh": OrderedDict(), + "signature": (), + "portable_structure": (), + "portable_routing": (), + "invalid_modes": (), + "anchors": (), + } + baseline = getattr(self, "_gefen_muon_collective_intent_baseline", None) + anchors = getattr(self, "_gefen_muon_collective_intent_anchors", ()) + if not anchors and not intent["anchors"]: + if preparation_error is not None: + raise RuntimeError( + "GefenMuon could not prepare DTensor collective routing " + "validation" + ) from preparation_error + return + scoped = self._gefen_codebook_process_group is not None + baseline_changed = baseline is not None and baseline != intent["signature"] + + if not scoped: + headers = self._reduce_muon_intent_header( + intent, + baseline_changed, + preparation_failed=preparation_error is not None, + ) + preparation_failed_any = any(header[0][0] != 0 for header in headers) + invalid_any = any(header[1][0] != 0 for header in headers) + structure_mismatch = any( + header[index][0] != header[index][1] + for header in headers + for index in range(2, 7) + ) + routing_mismatch = any( + header[index][0] != header[index][1] + for header in headers + for index in range(7, 11) + ) + baseline_changed = any(header[11][0] != 0 for header in headers) + if preparation_failed_any: + raise RuntimeError( + "GefenMuon requires structurally valid and identical DTensor " + "parameter groups on every rank before optimizer collectives" + ) + if invalid_any: + raise RuntimeError( + "GefenMuon requires sharded_mode to be 'approx', 'exact', " + "or 'distributed' on every rank of a DTensor/FSDP mesh" + ) + if structure_mismatch: + raise RuntimeError( + "GefenMuon requires identical DTensor parameter membership, " + "order, names, shapes, placements, and mesh routing on every " + "rank before optimizer collectives" + ) + if routing_mismatch: + raise RuntimeError( + "GefenMuon requires identical sharded_mode collective intent " + "on every rank of a DTensor/FSDP mesh before exact, " + "distributed, or approximate codebook operations" + ) + elif preparation_error is not None: + raise RuntimeError( + "GefenMuon requires structurally valid DTensor parameter groups" + ) from preparation_error + elif intent["invalid_modes"]: + # The mandatory scoped operation header has already agreed on the + # current fingerprint, so a unanimous invalid value can fail + # locally without splitting scope members. + raise RuntimeError( + "GefenMuon requires sharded_mode to be 'approx', 'exact', or " + "'distributed'" + ) + + if baseline_changed: + raise RuntimeError( + "GefenMuon rejects post-validation DTensor parameter-order, " + "rebinding, or sharded_mode transitions because live optimizer " + "state has no defined migration" + ) + + if baseline is None: + self._gefen_muon_collective_intent_baseline = intent["signature"] + if getattr(self, "_gefen_muon_capture_plan", None) is None: + self._freeze_muon_capture_plan(intent) + + @torch._dynamo.disable + def _assert_sharded_grad_presence_consistent( + self, *, routing_preflight_done: bool = False + ) -> None: """Fail before collectives when mesh ranks disagree on the step inputs. Exact and distributed Muon reconstruct full DTensor gradients with @@ -1450,30 +2008,24 @@ def _assert_sharded_grad_presence_consistent(self) -> None: index. Two rank-local properties therefore must agree mesh-wide before any of those collectives start: which parameters have gradients, and the order the parameters were registered in. Pack both per DeviceMesh - — an activity bit summed across every mesh dimension, and an insertion + — an activity bit summed across its one process group, and an insertion position reduced to its mesh-wide max and min — so every participating rank sees the same global result and raises on the same disagreement before codebook learning or optimizer state/parameter mutation begins. - Plain tensors/DDP and local-shard ``approx`` mode take no Muon gradient - collectives and therefore pay no preflight collective. Manual CUDA graph - capture also skips this host-branching check: capturable Gefen requires - eager warmup before capture, and those warmup steps establish that the - graph's fixed gradient-presence pattern is rank-consistent. + Plain tensors/DDP take no Muon gradient collectives. Local-shard + ``approx`` skips the larger activity/order probes but still joins the + fixed-size mode-intent handshake so asymmetric runtime mode changes fail + on every mesh rank before mutation. Manual CUDA graph capture first + restores the route frozen by eager warmup, then skips this host-branching + check; the graph can never partition from mutable public group modes. """ - if not self._dist_available(): + if self._restore_muon_capture_plan_if_capturing(): return - # Do not initialize or query CUDA for a CPU-only mesh optimizer. - # Spawned Gloo workers may run alongside a CUDA-heavy parent process, - # and an unrelated current-stream query can otherwise block before the - # CPU collective preflight. CUDA-backed optimizers retain the capture - # guard. - if any( - param.device.type == "cuda" - for group in self.param_groups - for param in group["params"] - ) and torch.cuda.is_current_stream_capturing(): + if not self._dist_available(): return + if not routing_preflight_done: + self._assert_sharded_mode_collective_intent_consistent() import torch.distributed as dist @@ -1492,23 +2044,11 @@ def _assert_sharded_grad_presence_consistent(self) -> None: mesh = p.device_mesh if mesh.get_coordinate() is None or mesh.size() < 2: continue - # Two equivalent DeviceMesh objects can still refer to distinct - # c10d groups, and a mis-built training script can register the - # same parameters in a different order per rank. Key each mesh - # by content instead of object identity, and sort meshes and - # items below, so this preflight's own collectives stay matched - # under both — and rank-divergent registration order is then - # detected positionally rather than corrupting the comparison. - process_groups = tuple(mesh.get_all_groups()) - key = ( - 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() - ), - tuple(str(pg.group_name) for pg in process_groups), - ) + # Keep rank-local c10d group handles in the scheduling key, but + # preserve first-seen mesh order. The fixed world header has + # already compared topology, the stable 1-D process-group token, + # and parameter order. Scheduling still uses the live handle. + key, process_groups = self._gradient_preflight_mesh_key(mesh) entry = by_mesh.get(key) if entry is None: entry = { @@ -1524,8 +2064,10 @@ def _assert_sharded_grad_presence_consistent(self) -> None: order_mismatches = [] duplicate_labels = [] duplicates_anywhere = False - for mesh_key in sorted(by_mesh): - entry = by_mesh[mesh_key] + # The optimizer-global fixed header has already proven identical + # parameter membership and order. Follow first-seen parameter/mesh order + # here; the stable 1-D process-group token already distinguishes routes. + for entry in by_mesh.values(): mesh = entry["mesh"] items = sorted( entry["items"], @@ -1612,10 +2154,8 @@ def _assert_sharded_grad_presence_consistent(self) -> None: dtype=torch.int32, device=device, ) - # Reducing the same vectors once along every Cartesian mesh - # dimension propagates mesh-global results to every rank. This is - # two collectives per mesh dimension, independent of the number of - # optimizer parameters. + # The supported one-dimensional mesh needs exactly two collectives + # for these vectors, independent of the optimizer parameter count. for process_group in entry["process_groups"]: if dist.get_world_size(process_group) > 1: dist.all_reduce( @@ -1625,11 +2165,11 @@ def _assert_sharded_grad_presence_consistent(self) -> None: order_probe, op=dist.ReduceOp.MAX, group=process_group ) - # Do not raise inside this loop: with DTensors on overlapping but - # non-identical meshes, a rank exiting early would strand peers - # that only share a later mesh inside that mesh's all_reduce. - # Accumulate every violation and raise after all local meshes - # have completed their probe collectives, like the presence path. + # Do not raise inside this loop. The fixed world header has + # validated the same ordered set of supported full-world meshes on + # every rank, so complete every first-seen mesh's probes before + # reporting accumulated violations. Otherwise a peer could be left + # inside a later validated mesh collective. order_flat = order_probe.cpu().tolist() if order_flat[-1] > 0: duplicates_anywhere = True @@ -1715,8 +2255,8 @@ def _step_distributed(self, items) -> None: # world), and the per-bucket broadcasts are a pure communication phase. # Eligibility (1-D mesh, world>=2) is a property of each param's mesh and # so is identical on every rank -> the eligible/fallback split, and thus - # the collective order, agrees globally. Non-eligible matrices (multi-dim - # HSDP x TP meshes, world==1) keep the replicated exact full-NS path. + # the collective order, agrees globally. Plain tensors and world-one + # DTensors keep the replicated exact full-NS path. eligible, fallback = [], [] for (group, name, p, grad) in items: pg = self._distributed_process_group(p) @@ -2021,45 +2561,145 @@ def _distributed_state_items(self, state_dict): by_pg.setdefault(pg, []).append((name, p, saved_id)) return by_pg - def _consolidate_distributed_state_dict(self, state_dict): - # In distributed mode the persistent momentum state is updated only on the - # stable owner rank for each matrix. Before checkpointing, broadcast each - # owner's serialized per-param state to every rank so rank-0-only writers - # can save a complete optimizer state_dict AFTER all ranks have called - # this method. This method is therefore collective for distributed - # sharded params, matching the DTensor collectives used by step(). - if not self._dist_available(): - return state_dict + def _preflight_muon_checkpoint_routing(self) -> None: + """Freeze or validate routing before checkpoint mode decisions.""" + + try: + capture_plan_restored = self._restore_muon_capture_plan_if_capturing() + capture_route_error = None + except Exception as exc: + capture_plan_restored = False + capture_route_error = exc + scoped = self._gefen_codebook_process_group is not None + if scoped: + self._synchronize_codebook_scope_failure( + capture_route_error, "checkpoint capture routing preflight" + ) + elif capture_route_error is not None: + raise capture_route_error + if capture_plan_restored: + # A capture must use the immutable eager route and cannot enter the + # host-side intent handshake. The downstream checkpoint path keeps + # its existing capture restrictions. + return + + baseline_before = getattr( + self, "_gefen_muon_collective_intent_baseline", None + ) + capture_plan_before = getattr(self, "_gefen_muon_capture_plan", None) + anchors_before = getattr( + self, "_gefen_muon_collective_intent_anchors", () + ) + try: + self._assert_sharded_mode_collective_intent_consistent() + routing_error = None + except Exception as exc: + routing_error = exc + if scoped: + self._synchronize_codebook_scope_failure( + routing_error, "checkpoint routing preflight" + ) + elif routing_error is not None: + raise routing_error + if not scoped: + return + try: + # Scoped mode intent is part of the fixed operation header. This + # second agreement catches a pristine optimizer whose first local + # baseline was frozen from rank-divergent public group modes. + self._validate_codebook_scope_operation_header("checkpoint") + except Exception: + # A failed first checkpoint must not leave different private route + # baselines or capture plans installed on the scope members. + self._gefen_muon_collective_intent_baseline = baseline_before + self._gefen_muon_capture_plan = capture_plan_before + self._gefen_muon_collective_intent_anchors = anchors_before + raise + + def _synchronize_muon_checkpoint_failure(self, error, phase: str) -> None: + """Propagate a checkpoint callback/preparation error before collectives.""" + + if self._gefen_codebook_process_group is not None: + self._synchronize_codebook_scope_failure( + error, "checkpoint {}".format(phase) + ) + return + anchors = getattr(self, "_gefen_muon_collective_intent_anchors", ()) + if not anchors or not self._dist_available(): + if error is not None: + raise error + return + + # Exact/distributed DTensor state is not rank-local, so the base + # checkpoint failure protocol intentionally has no vote for it. Use the + # retained full-world anchor solely for error agreement; do not change + # the base intent that controls approx payload consolidation. import torch.distributed as dist - by_pg = self._distributed_state_items(state_dict) - if not by_pg: - return state_dict + failed = torch.tensor( + int(error is not None), + dtype=torch.int32, + device=anchors[0][2], + ) + dist.all_reduce(failed, op=dist.ReduceOp.MAX, group=dist.group.WORLD) + if int(failed.item()) == 0: + return + if error is not None: + raise RuntimeError( + "GefenMuon unscoped DTensor checkpoint {} failed on this " + "rank: {}".format(phase, error) + ) from error + raise RuntimeError( + "GefenMuon unscoped DTensor checkpoint {} failed on another " + "rank".format(phase) + ) - marker_groups = [] - for pg, pg_items in by_pg.items(): + def _synchronize_rank_local_checkpoint_failure(self, error, phase: str) -> None: + # Base checkpoint phases decide whether to run their failure agreement + # only after a callback. Revalidate first so a state/load pre-hook that + # edits sharded_mode cannot make that very decision rank-divergent. + self._preflight_muon_checkpoint_routing() + self._synchronize_muon_checkpoint_failure(error, phase) + error = None + return super()._synchronize_rank_local_checkpoint_failure(error, phase) + + def state_dict(self): + """Export only after agreeing on the live Muon collective route.""" + + self._preflight_muon_checkpoint_routing() + return super().state_dict() + + def _prepare_distributed_state_dict_consolidation(self, state_dict): + """Build every Parallel-Muon owner payload before communication.""" + + import pickle + import torch.distributed as dist + + plans = [] + for pg, pg_items in self._distributed_state_items(state_dict).items(): world = dist.get_world_size(pg) my_coord = dist.get_group_rank(pg, dist.get_rank()) - global_rank = [dist.get_global_rank(pg, c) for c in range(world)] - owner_by_saved_id = { - saved_id: _stable_distributed_owner(idx, world) - for idx, (_, _, saved_id) in enumerate(pg_items) - } - - for name, p, saved_id in pg_items: - owner = owner_by_saved_id[saved_id] - src = global_rank[owner] + global_ranks = tuple( + dist.get_global_rank(pg, coordinate) + for coordinate in range(world) + ) + item_plans = [] + for index, (name, parameter, saved_id) in enumerate(pg_items): + owner = _stable_distributed_owner(index, world) tensor_values = [] + payload = None if my_coord == owner: - pstate = state_dict["state"].get(saved_id, {}) - meta = [] - for key, value in pstate.items(): + metadata = [] + state_keys = [] + for key, value in state_dict["state"].get(saved_id, {}).items(): + key_text = str(key) + state_keys.append(key_text) if torch.is_tensor(value): tensor = value.detach() if not tensor.is_contiguous(): tensor = tensor.contiguous() - meta.append( + metadata.append( ( key, "tensor", @@ -2069,64 +2709,160 @@ def _consolidate_distributed_state_dict(self, state_dict): ) tensor_values.append(tensor) else: - meta.append((key, "object", value)) - else: - meta = None - - obj = [meta] - dist.broadcast_object_list(obj, src=src, group=pg) - meta = obj[0] - - owner_state = {} - tensor_idx = 0 - for entry in meta: - key = entry[0] - kind = entry[1] - if kind == "object": - owner_state[key] = entry[2] - continue - - _, _, shape, dtype = entry - if my_coord == owner: - tensor = tensor_values[tensor_idx] - else: - tensor = torch.empty( - shape, - dtype=dtype, - device=self._state_tensor_device(p), - ) - dist.broadcast(tensor, src=src, group=pg) - owner_state[key] = tensor - tensor_idx += 1 + metadata.append((key, "object", value)) + marker = { + "state_keys": tuple(sorted(state_keys)), + "initialized": any(key != "name" for key in state_keys), + } + payload = (metadata, marker) + # broadcast_object_list serializes this payload internally. + # Exercise serialization before the failure vote so an + # owner-only object error cannot strand its peers. + pickle.dumps(payload) + item_plans.append( + { + "name": str(name), + "parameter": parameter, + "saved_id": saved_id, + "owner": owner, + "source": global_ranks[owner], + "is_owner": my_coord == owner, + "tensor_values": tuple(tensor_values), + "payload": payload, + "receive_device": self._state_tensor_device(parameter), + } + ) + plans.append( + { + "process_group": pg, + "world_size": world, + "items": tuple(item_plans), + } + ) + return tuple(plans) + + def _consolidate_distributed_state_dict(self, state_dict): + # In distributed mode the persistent momentum state is updated only on the + # stable owner rank for each matrix. Before checkpointing, broadcast each + # owner's serialized per-param state to every rank so rank-0-only writers + # can save a complete optimizer state_dict AFTER all ranks have called + # this method. This method is therefore collective for distributed + # sharded params, matching the DTensor collectives used by step(). + if not self._dist_available(): + return state_dict + + import torch.distributed as dist + + try: + plans = self._prepare_distributed_state_dict_consolidation(state_dict) + preparation_error = None + except Exception as exc: + plans = () + preparation_error = exc + self._synchronize_muon_checkpoint_failure( + preparation_error, "distributed serialization preparation" + ) + if not plans: + return state_dict - state_dict["state"][saved_id] = owner_state + # Phase one communicates every small owner descriptor. Do not allocate + # or inspect descriptor-dependent tensors yet: all ranks must finish the + # same metadata schedule even if one later rejects an allocation/schema. + received_groups = [] + for group_plan in plans: + pg = group_plan["process_group"] + received_items = [] + for item in group_plan["items"]: + obj = [item["payload"]] + dist.broadcast_object_list(obj, src=item["source"], group=pg) + received_items.append((item, obj[0])) + received_groups.append((group_plan, tuple(received_items))) + + # Phase two validates every received descriptor and allocates every + # non-owner tensor before the first tensor broadcast. A deterministic + # allocation/device/schema failure therefore takes one coordinated vote + # instead of stranding an owner inside the matching broadcast. + try: + execution_groups = [] + for group_plan, received_items in received_groups: + execution_items = [] + marker_params = [] + for item, payload in received_items: + metadata, marker = payload + if not isinstance(metadata, (list, tuple)): + raise RuntimeError( + "distributed owner metadata is not a sequence" + ) + if not isinstance(marker, dict): + raise RuntimeError( + "distributed owner marker is not a mapping" + ) + owner_state = {} + tensors = [] + tensor_idx = 0 + for entry in metadata: + if not isinstance(entry, (list, tuple)) or len(entry) < 2: + raise RuntimeError( + "distributed owner state metadata is malformed" + ) + key, kind = entry[:2] + if kind == "object" and len(entry) == 3: + owner_state[key] = entry[2] + continue + if kind != "tensor" or len(entry) != 4: + raise RuntimeError( + "distributed owner state kind is unsupported" + ) + _, _, shape, dtype = entry + if item["is_owner"]: + tensor = item["tensor_values"][tensor_idx] + else: + tensor = torch.empty( + shape, + dtype=dtype, + device=item["receive_device"], + ) + owner_state[key] = tensor + tensors.append(tensor) + tensor_idx += 1 + if tensor_idx != len(item["tensor_values"]) and item["is_owner"]: + raise RuntimeError( + "distributed owner tensor metadata count is inconsistent" + ) + marker_param = { + "saved_id": item["saved_id"], + "name": item["name"], + "shape": tuple(item["parameter"].shape), + "owner": item["owner"], + "state_keys": tuple(marker["state_keys"]), + "initialized": bool(marker["initialized"]), + } + execution_items.append((item, owner_state, tuple(tensors))) + marker_params.append(marker_param) + execution_groups.append( + (group_plan, tuple(execution_items), tuple(marker_params)) + ) + receive_error = None + except Exception as exc: + execution_groups = [] + receive_error = exc + self._synchronize_muon_checkpoint_failure( + receive_error, "distributed receive preparation" + ) + # Phase three contains only the preplanned tensor broadcasts followed + # by non-throwing publication into the already-validated dictionaries. + marker_groups = [] + for group_plan, execution_items, marker_params in execution_groups: + pg = group_plan["process_group"] + for item, owner_state, tensors in execution_items: + for tensor in tensors: + dist.broadcast(tensor, src=item["source"], group=pg) + state_dict["state"][item["saved_id"]] = owner_state marker_groups.append( { - "world_size": world, - "params": [ - { - "saved_id": saved_id, - "name": str(name), - "shape": tuple(p.shape), - "owner": owner_by_saved_id[saved_id], - "state_keys": tuple( - sorted( - str(key) - for key in state_dict["state"] - .get(saved_id, {}) - .keys() - ) - ), - "initialized": any( - str(key) != "name" - for key in state_dict["state"] - .get(saved_id, {}) - .keys() - ), - } - for name, p, saved_id in pg_items - ], + "world_size": group_plan["world_size"], + "params": list(marker_params), } ) @@ -2164,24 +2900,37 @@ def _drop_non_owned_distributed_state(self) -> None: pstate.clear() pstate["name"] = str(name).lower() - def _state_dict_impl(self): + def _state_dict_impl(self, *, consolidate_rank_local: bool = True): # Parallel-Muon owner state must be made complete before Gefen's # rank-local DTensor adapter serializes the per-rank payload. This # ordering matters for mixed optimizers carrying both ``approx`` and # ``distributed`` groups: wrapping first replaces every real state entry # with opaque transport tensors, which are not owner momentum. - state_dict = super()._state_dict_impl(consolidate_rank_local=False) + self._preflight_muon_checkpoint_routing() + try: + state_dict = super()._state_dict_impl(consolidate_rank_local=False) + serialization_error = None + except Exception as exc: + state_dict = None + serialization_error = exc + self._synchronize_muon_checkpoint_failure( + serialization_error, "local serialization preparation" + ) state_dict = self._consolidate_distributed_state_dict(state_dict) checkpoint_metadata = None groups = state_dict.get("param_groups", ()) if groups: checkpoint_metadata = groups[0].get("_gefen_checkpoint_metadata") - if self._uses_rank_local_sharded_state(): + if consolidate_rank_local and self._uses_rank_local_sharded_state(): + metadata_error = None if not isinstance(checkpoint_metadata, dict): - raise RuntimeError( + metadata_error = RuntimeError( "GefenMuon rank-local checkpoint metadata was not initialized" ) + self._synchronize_muon_checkpoint_failure( + metadata_error, "rank-local serialization preparation" + ) self._consolidate_rank_local_sharded_state( state_dict, checkpoint_metadata ) @@ -2615,7 +3364,7 @@ def _validate_distributed_checkpoint_load(self, state_dict, marker) -> None: # A sharded_mode='distributed' optimizer group may contain a mixture of # Parallel-Muon-eligible parameters and replicated fallbacks (plain - # tensors, multi-dimensional meshes, or world-one meshes). Save-side + # tensors or world-one meshes). Save-side # consolidation creates one owner manifest per eligible process group # and deliberately leaves fallback state to the ordinary Gefen loader. # Bind the proof to those same ordered process-group partitions instead @@ -2825,7 +3574,17 @@ def _validate_distributed_checkpoint_load(self, state_dict, marker) -> None: ) self._validate_distributed_codebook(state_dict, required=initialized_any) + def _prepare_load_state_dict(self, state_dict): + """Stage a load only after agreeing on the live Muon route.""" + + self._preflight_muon_checkpoint_routing() + return super()._prepare_load_state_dict(state_dict) + def _load_state_dict_impl(self, state_dict): + self._preflight_muon_checkpoint_routing() + routing_baseline_before = getattr( + self, "_gefen_muon_collective_intent_baseline", None + ) state_dict = dict(state_dict) state_dict = self._pack_legacy_param_groups_for_load(state_dict) marker = self._distributed_checkpoint_marker(state_dict) @@ -2852,11 +3611,61 @@ def _load_state_dict_impl(self, state_dict): "batched_ns_workspace_bytes", BATCHED_NS_DEFAULT_WORKSPACE_BYTES, ) + if not self._muon_mode_token(group.get("sharded_mode"))[1]: + raise ValueError( + "GefenMuon checkpoint sharded_mode must be 'approx', " + "'exact', or 'distributed'" + ) self._drop_non_owned_distributed_state() self._install_rank_local_checkpoint_schema() + # Optimizer.load_state_dict may replace param-group dictionaries. This + # method runs on the isolated staged optimizer, so discard capture plans + # that still reference pre-load group objects only after the complete + # load path above succeeded. After publication, the public load post-hook + # phase performs one eager routing preflight and freezes a fresh plan + # against the loaded group objects. The staged native-state validator + # runs next; publication remains atomic if the loaded route is invalid. + loaded_intent = self._muon_collective_intent() + if ( + routing_baseline_before is not None + and routing_baseline_before != loaded_intent["signature"] + ): + raise RuntimeError( + "GefenMuon rejects post-validation DTensor parameter-order, " + "rebinding, or sharded_mode transitions because live optimizer " + "state has no defined migration" + ) + self._gefen_muon_collective_intent_baseline = ( + loaded_intent["signature"] if loaded_intent["anchors"] else None + ) + self._gefen_muon_capture_plan = None + self._gefen_muon_collective_intent_anchors = loaded_intent["anchors"] @torch.no_grad() def step(self, closure=None): + try: + capture_plan_restored = self._restore_muon_capture_plan_if_capturing() + capture_route_error = None + except Exception as exc: + capture_plan_restored = False + capture_route_error = exc + if self._gefen_codebook_process_group is not None: + self._synchronize_codebook_scope_failure( + capture_route_error, "capture routing preflight" + ) + elif capture_route_error is not None: + raise capture_route_error + + initial_routing_preflight_done = ( + self._gefen_codebook_process_group is None + and not capture_plan_restored + ) + if initial_routing_preflight_done: + # Run before finalized-layout forensics so a direct asymmetric + # membership/rebinding edit is rejected by every mesh member, not + # just by the rank whose local layout cache notices it first. + self._assert_sharded_mode_collective_intent_consistent() + self._assert_finalized_binding_layout() self._assert_runtime_codebook_process_group() if self._has_unscoped_whole_owner_bindings(): @@ -2891,6 +3700,19 @@ def step(self, closure=None): self._assert_finalized_binding_layout() self._assert_runtime_codebook_process_group() + routing_preflight_done = ( + initial_routing_preflight_done and closure is None + ) + if ( + self._gefen_codebook_process_group is None + and not routing_preflight_done + and not capture_plan_restored + ): + # This fixed-size header must precede local structural validation: + # a malformed group on one rank is itself a coordinated rejection, + # not permission for peers to enter a later optimizer collective. + self._assert_sharded_mode_collective_intent_consistent() + routing_preflight_done = True try: _assert_optimizer_gradients_structurally_valid( self, require_2d_params=True @@ -2916,6 +3738,13 @@ def step(self, closure=None): ) and not self._prepare_scoped_amp_optimizer_step(): return loss + # Validate/freeze sharded routing before any rank-local mode branch. In + # manual CUDA capture the early capture-plan restore has already forced + # the coordinated route, so this helper performs no host collective. + self._assert_sharded_grad_presence_consistent( + routing_preflight_done=routing_preflight_done + ) + # Partition the work once so distributed-mode sharded params can take the # stable-owner Parallel-Muon path while every other param keeps the normal # per-param path. @@ -2941,14 +3770,6 @@ def step(self, closure=None): elif grad is not None: regular_items.append((group, name, p, grad)) - # This must precede _maybe_refresh_gefen_codebook(): the first-step - # Muon codebook iterator itself calls full_tensor() in exact/distributed - # modes, before either update dispatcher gets a chance to validate the - # active set. The preflight is mutation-free and gives every mesh rank - # the same clear error instead of leaving active ranks in an unmatched - # collective. - self._assert_sharded_grad_presence_consistent() - self._maybe_refresh_gefen_codebook() self._maybe_save_gefen_grad_histogram() diff --git a/src/gefen/hybrid.py b/src/gefen/hybrid.py index e023f9c..46eae11 100644 --- a/src/gefen/hybrid.py +++ b/src/gefen/hybrid.py @@ -60,6 +60,12 @@ import torch import torch.nn as nn +from gefen._adamw_adapter import ( + commit_adamw_stage, + prepare_adamw_load_state_dict, + run_adamw_load_post_hooks, + stage_adamw_post_sharding, +) from gefen.codebook import CodebookProcessGroupBinding from gefen.contracts import ( LogicalSlice, @@ -76,6 +82,11 @@ _assert_optimizer_gradients_structurally_valid, ) from gefen.gefen_muon import GefenMuon +from gefen.dtensor import ( + is_exact_dtensor, + resolve_local_tensor, + validate_dtensor_rebinding_plan, +) from gefen.params import ( DEFAULT_BACKUP_SUBSTRINGS, is_muon_param, @@ -83,9 +94,19 @@ validate_split, ) from gefen.rebinding import ParameterRebinding +from gefen.portable_identity import ( + _normalize_sharding_manifest, + _normalize_shard_identity, + _serialize_sharding_manifest, + _serialize_shard_identity, +) logger = logging.getLogger(__name__) +_HYBRID_NATIVE_BINDING_KEY = "finalized_binding" +_HYBRID_NATIVE_BINDING_FORMAT = "gefen.hybrid.finalized_native_binding" +_HYBRID_NATIVE_BINDING_VERSION = 1 + _UNKNOWN_STATE_KEY_MSG = ( "GefenMuonHybrid.state was accessed with a key that is not a parameter of " "either sub-optimizer (the Muon half or the backup half). A torch optimizer " @@ -230,8 +251,12 @@ class GefenMuonHybrid(torch.optim.Optimizer): * ``state_dict()`` schema: NOT the standard ``{"state", "param_groups"}`` layout, but ``{"muon": , "backup": - , "backup_optimizer": "gefen" | "adamw"}``. - ``load_state_dict`` only accepts that nested schema; checkpoints + , "backup_optimizer": "gefen" | "adamw"}``, + plus a versioned ``finalized_binding`` routing/layout guard after + ``post_sharding``. ``load_state_dict`` only accepts that nested schema; + finalized and legacy-unfinalized schemas do not cross-load, and a + finalized load requires the exact child set, FQN routing, shard identity, + and local topology before either child is mutated. Checkpoints consolidated/converted to the flat torch layout (e.g. by FSDP/DeepSpeed tooling) are rejected rather than silently ignored. Legacy nested checkpoints without ``backup_optimizer`` are treated as Gefen-backed. @@ -570,9 +595,12 @@ def _is_no_decay(name): self._hybrid_post_sharding_finalized = False self._hybrid_sharding_manifest = None self._hybrid_local_shard_bindings = () + self._hybrid_shard_bindings = {} self._hybrid_fqn_roles = () self._hybrid_codebook_process_group = None self._hybrid_finalized_slots = () + self._hybrid_rank_local_checkpoint_collective = False + self._hybrid_rank_local_checkpoint_device = None # Composite finalized-layout forensics cache, mirroring the O(local # params) scheme Gefen/GefenMuon use for their own step guards. The @@ -651,12 +679,14 @@ def _gefen_rebinding_children(self): 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" - ) + if self.backup_optimizer == "gefen": + if type(self.backup) is not Gefen: + raise TypeError("GefenMuonHybrid post_sharding requires an exact Gefen backup child") + elif self.backup_optimizer == "adamw": + if type(self.backup) is not torch.optim.AdamW: + raise TypeError("GefenMuonHybrid post_sharding requires an exact AdamW backup child") + else: + raise ValueError("GefenMuonHybrid has an invalid backup optimizer policy") children.append(("backup", self.backup)) if not children: raise RuntimeError("GefenMuonHybrid has no child optimizer to rebind") @@ -685,9 +715,12 @@ def _hybrid_identity_metadata_empty(self) -> bool: return ( self._hybrid_sharding_manifest is None and self._hybrid_local_shard_bindings == () + and self._hybrid_shard_bindings == {} and self._hybrid_fqn_roles == () and self._hybrid_codebook_process_group is None and self._hybrid_finalized_slots == () + and self._hybrid_rank_local_checkpoint_collective is False + and self._hybrid_rank_local_checkpoint_device is None ) def _assert_composite_rebinding_pristine(self, children) -> None: @@ -801,11 +834,37 @@ def _stage_post_sharding( 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") + dtensor_plan = validate_dtensor_rebinding_plan(rebindings, manifest) + rank_local_checkpoint_device = None + if dtensor_plan is not None: + rank_local_checkpoint_device = next( + resolve_local_tensor(item.new_parameter).device + for item in rebindings + if item.shard.layout is ParameterLayout.DTENSOR_1D_DEFAULT_WORLD + ) + gefen_roles = { + role + for role, child in children + if type(child) in {Gefen, GefenMuon} + } 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") + scoped_fqns = { + fqn for fqn, role in fqn_roles.items() if role in gefen_roles + } + if not scoped_fqns: + raise ValueError("a Hybrid codebook binding requires at least one Gefen-backed child") + if any( + shard.parameter.fqn in scoped_fqns + and shard.process_group != codebook_process_group.identity + for shard in manifest.shards + ): + raise ValueError("every Gefen-backed Hybrid shard must use the shared codebook process-group identity") + if any( + fqn_roles[item.shard.parameter.fqn] in gefen_roles + and item.shard.local_member != codebook_process_group.local_member + for item in rebindings + ): + raise ValueError("every Gefen-backed Hybrid local shard must match the shared codebook member") staged_children = [] for role, child in children: @@ -817,22 +876,30 @@ def _stage_post_sharding( 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") + if type(child) is torch.optim.AdamW: + staged = stage_adamw_post_sharding( + child, + child_rebindings, + child_manifest, + ) + else: + 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() + shard_bindings = {} new_state_param_owner = {} finalized_slots = [] for role, child, staged, _child_manifest in staged_children: @@ -842,12 +909,22 @@ def _stage_post_sharding( tuple(tuple(group["params"]) for group in staged.param_groups), ) ) - for parameter, shard in staged._gefen_local_shard_bindings: + child_bindings = ( + tuple( + (item.new_parameter, item.shard) + for item in by_role[role] + ) + if type(child) is torch.optim.AdamW + else staged._gefen_local_shard_bindings + ) + for parameter, shard in child_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)) + if parameter is not None: + shard_bindings[parameter] = shard for group in staged.param_groups: for parameter in group["params"]: parameter_id = id(parameter) @@ -864,9 +941,12 @@ def _stage_post_sharding( "state_param_owner": new_state_param_owner, "manifest": manifest, "local_bindings": tuple(local_bindings), + "shard_bindings": shard_bindings, "fqn_roles": tuple(sorted(fqn_roles.items())), "codebook_process_group": codebook_process_group, "finalized_slots": tuple(finalized_slots), + "rank_local_checkpoint_collective": dtensor_plan is not None, + "rank_local_checkpoint_device": rank_local_checkpoint_device, } def post_sharding( @@ -891,7 +971,10 @@ def post_sharding( codebook_process_group, ) for _role, child, staged_child, _child_manifest in staged["children"]: - dict.update(child.__dict__, staged_child.__dict__) + if type(child) is torch.optim.AdamW: + commit_adamw_stage(child, staged_child) + else: + dict.update(child.__dict__, staged_child.__dict__) dict.update( self.__dict__, { @@ -899,9 +982,16 @@ def post_sharding( "_hybrid_post_sharding_finalized": True, "_hybrid_sharding_manifest": staged["manifest"], "_hybrid_local_shard_bindings": staged["local_bindings"], + "_hybrid_shard_bindings": staged["shard_bindings"], "_hybrid_fqn_roles": staged["fqn_roles"], "_hybrid_codebook_process_group": staged["codebook_process_group"], "_hybrid_finalized_slots": staged["finalized_slots"], + "_hybrid_rank_local_checkpoint_collective": staged[ + "rank_local_checkpoint_collective" + ], + "_hybrid_rank_local_checkpoint_device": staged[ + "rank_local_checkpoint_device" + ], # 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 @@ -988,9 +1078,13 @@ def _hybrid_layout_forensics_fast_tokens(self): self._hybrid_post_sharding_finalized, self._hybrid_sharding_manifest, self._hybrid_local_shard_bindings, + self._hybrid_shard_bindings, + len(self._hybrid_shard_bindings), self._hybrid_fqn_roles, self._hybrid_finalized_slots, self._hybrid_codebook_process_group, + self._hybrid_rank_local_checkpoint_collective, + self._hybrid_rank_local_checkpoint_device, self._state_param_owner, len(self._state_param_owner), self.defaults, @@ -998,7 +1092,11 @@ def _hybrid_layout_forensics_fast_tokens(self): len(self._subopts), ] for child in self._subopts: - live.append(child._finalized_binding_layout_matches()) + live.append( + True + if type(child) is torch.optim.AdamW + else child._finalized_binding_layout_matches() + ) GefenMuonHybrid._hybrid_child_param_group_tokens(child, live) return tuple(live) @@ -1037,8 +1135,10 @@ def _finalized_binding_layout_matches_full(self) -> bool: 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_shard_bindings) is not dict or type(self._hybrid_fqn_roles) is not tuple or type(self._hybrid_finalized_slots) is not tuple + or type(self._hybrid_rank_local_checkpoint_collective) is not bool or type(self._state_param_owner) is not dict ): return False @@ -1064,56 +1164,110 @@ def _finalized_binding_layout_matches_full(self) -> bool: 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): + gefen_roles = { + role + for role, child in children + if type(child) in {Gefen, GefenMuon} + } + if not gefen_roles or any( + role_by_fqn[shard.parameter.fqn] in gefen_roles + and shard.process_group != binding.identity + for shard in self._hybrid_sharding_manifest.shards + ): return False - expected_local = [] + local_by_role = {role: [] for role, _child in children} + local_fqns = set() + for parameter, shard in self._hybrid_local_shard_bindings: + if shard not in self._hybrid_sharding_manifest.shards: + return False + fqn = shard.parameter.fqn + role = role_by_fqn.get(fqn) + if role not in local_by_role or fqn in local_fqns: + return False + local_fqns.add(fqn) + local_by_role[role].append((parameter, shard)) + if parameter is not None: + if self._hybrid_shard_bindings.get(parameter) != shard: + return False + if local_fqns != set(role_by_fqn): + return False + if tuple( + sorted(self._hybrid_local_shard_bindings, key=lambda item: item[1].sort_key) + ) != self._hybrid_local_shard_bindings: + return False + 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: + if type(child) is torch.optim.AdamW: + if type(child.param_groups) is not list: + return False + for group in child.param_groups: + if ( + type(group) is not dict + or type(group.get("params")) is not list + or type(group.get("param_names")) is not list + or len(group["params"]) != len(group["param_names"]) + ): + return False + child_live = [ + parameter + for group in child.param_groups + for parameter in group["params"] + ] + routed_live = [ + parameter + for parameter, _shard in local_by_role[role] + if parameter is not None + ] + if len(child_live) != len(routed_live) or any( + not any(parameter is routed for routed in routed_live) + for parameter in child_live + ): return False - if binding is not None and ( - shard.process_group != binding.identity or shard.local_member != binding.local_member + else: + 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 + child_bindings = tuple( + sorted( + child._gefen_local_shard_bindings, + key=lambda item: item[1].sort_key, + ) + ) + if len(child_bindings) != len(local_by_role[role]) or any( + not self._same_local_binding(live, expected) + for live, expected in zip(local_by_role[role], child_bindings) + ): + return False + if binding is not None and any( + shard.process_group != binding.identity + or shard.local_member != binding.local_member + for _parameter, shard in child_bindings ): 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): @@ -1122,6 +1276,39 @@ def _finalized_binding_layout_matches_full(self) -> bool: 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 + if self._hybrid_shard_bindings.get(parameter) is None: + return False + if len(self._hybrid_shard_bindings) != len(expected_owner): + return False + dtensor_plan = validate_dtensor_rebinding_plan( + tuple( + ParameterRebinding(parameter, parameter, shard) + for parameter, shard in self._hybrid_local_shard_bindings + if parameter is not None + ), + self._hybrid_sharding_manifest, + ) + if self._hybrid_rank_local_checkpoint_collective != ( + dtensor_plan is not None + ): + return False + if dtensor_plan is None: + if self._hybrid_rank_local_checkpoint_device is not None: + return False + else: + expected_checkpoint_device = next( + resolve_local_tensor(parameter).device + for parameter, shard in self._hybrid_local_shard_bindings + if parameter is not None + and shard.layout is ParameterLayout.DTENSOR_1D_DEFAULT_WORLD + ) + if ( + type(self._hybrid_rank_local_checkpoint_device) + is not torch.device + or self._hybrid_rank_local_checkpoint_device + != expected_checkpoint_device + ): + return False return True except ( AttributeError, @@ -1150,6 +1337,152 @@ def _assert_finalized_binding_layout(self, *, full: bool = False) -> None: elif not self._hybrid_identity_metadata_empty(): raise RuntimeError("GefenMuonHybrid found an incomplete post_sharding identity plan") + def _rank_local_checkpoint_collective_intent(self) -> bool: + """Return whether native checkpoint phases require default-world agreement.""" + + return bool( + getattr(self, "_hybrid_rank_local_checkpoint_collective", False) + ) + + def _rank_local_checkpoint_collective_device(self) -> torch.device: + import torch.distributed as dist + + backend = str(dist.get_backend()).lower() + if "nccl" not in backend: + return torch.device("cpu") + device = getattr(self, "_hybrid_rank_local_checkpoint_device", None) + if device is not None: + if device.type == "cuda" and device.index is not None: + return device + raise RuntimeError( + "GefenMuonHybrid rank-local NCCL checkpoint agreement requires a concrete CUDA device" + ) + try: + parameters = tuple( + parameter + for parameter, shard in self._hybrid_local_shard_bindings + if parameter is not None + and shard.layout is ParameterLayout.DTENSOR_1D_DEFAULT_WORLD + ) + except (AttributeError, TypeError, ValueError): + parameters = () + if not parameters: + try: + parameters = tuple( + parameter + for child in self._subopts + for group in child.param_groups + for parameter in group["params"] + if is_exact_dtensor(parameter) + ) + except (AttributeError, KeyError, TypeError): + parameters = () + for parameter in parameters: + if is_exact_dtensor(parameter): + device = resolve_local_tensor(parameter).device + if device.type == "cuda" and device.index is not None: + return device + raise RuntimeError( + "GefenMuonHybrid rank-local NCCL checkpoint agreement requires a concrete CUDA device" + ) + + def _synchronize_rank_local_checkpoint_failure(self, error, phase: str) -> None: + if self._rank_local_checkpoint_collective_intent(): + if ( + not torch.distributed.is_available() + or not torch.distributed.is_initialized() + ): + if error is not None: + raise error + raise RuntimeError( + "GefenMuonHybrid rank-local DTensor checkpointing requires initialized torch.distributed" + ) + import torch.distributed as dist + + failed = torch.tensor( + int(error is not None), + dtype=torch.int32, + device=self._rank_local_checkpoint_collective_device(), + ) + dist.all_reduce(failed, op=dist.ReduceOp.MAX) + if int(failed.item()) == 0: + return + if error is not None: + raise RuntimeError( + "GefenMuonHybrid rank-local checkpoint {} failed on this rank: {}".format( + phase, error + ) + ) from error + raise RuntimeError( + "GefenMuonHybrid rank-local checkpoint {} failed on another rank".format( + phase + ) + ) + + binding = getattr(self, "_hybrid_codebook_process_group", None) + if type(binding) is not CodebookProcessGroupBinding: + binding = None + try: + binding = next( + candidate + for child in self._subopts + for candidate in ( + getattr(child, "_gefen_codebook_process_group", None), + ) + if type(candidate) is CodebookProcessGroupBinding + ) + except (AttributeError, StopIteration, TypeError): + binding = None + if binding is not None and len(binding.identity.ordered_members) > 1: + if ( + not torch.distributed.is_available() + or not torch.distributed.is_initialized() + ): + if error is not None: + raise error + raise RuntimeError( + "GefenMuonHybrid scoped checkpointing requires initialized torch.distributed" + ) + import torch.distributed as dist + + failed = torch.tensor( + int(error is not None), + dtype=torch.int32, + device=binding.collective_device, + ) + dist.all_reduce( + failed, + op=dist.ReduceOp.MAX, + group=binding.process_group, + ) + if int(failed.item()) == 0: + return + if error is not None: + raise RuntimeError( + "GefenMuonHybrid scoped checkpoint {} failed on local member {}: {}".format( + phase, + binding.local_member, + error, + ) + ) from error + raise RuntimeError( + "GefenMuonHybrid scoped checkpoint {} failed on another process-group member".format( + phase + ) + ) + if error is not None: + raise error + + def _run_rank_local_checkpoint_phase(self, phase: str, callback): + result = None + error = None + try: + result = callback() + except Exception as exc: + error = exc + self._synchronize_rank_local_checkpoint_failure(error, phase) + return result + def parameter_identity(self, parameter) -> ParameterIdentity: """Return the canonical identity bound to one live Hybrid parameter.""" @@ -1159,10 +1492,14 @@ def shard_identity(self, parameter) -> ShardIdentity: """Return the stable shard identity bound to one live parameter.""" self._assert_finalized_binding_layout() + try: + shard = self._hybrid_shard_bindings[parameter] + except KeyError: + raise KeyError("parameter has no finalized GefenMuonHybrid shard identity") 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) + return shard def shard_bindings(self): """Return all local tensor/identity pairs in canonical order.""" @@ -1190,6 +1527,125 @@ def codebook_process_group_binding(self): self._assert_finalized_binding_layout() return self._hybrid_codebook_process_group + def _serialized_finalized_native_binding(self): + if not self._hybrid_post_sharding_finalized: + return None + self._assert_finalized_binding_layout(full=True) + role_by_fqn = dict(self._hybrid_fqn_roles) + child_groups = [] + for role, child in self._gefen_rebinding_children(): + child_groups.append( + { + "role": role, + "groups": [ + [ + self._hybrid_shard_bindings[parameter].parameter.fqn + for parameter in group["params"] + ] + for group in child.param_groups + ], + } + ) + return { + "format": _HYBRID_NATIVE_BINDING_FORMAT, + "format_version": _HYBRID_NATIVE_BINDING_VERSION, + "backup_optimizer": self.backup_optimizer, + "manifest": _serialize_sharding_manifest(self._hybrid_sharding_manifest), + "routing": [ + {"fqn": fqn, "role": role} + for fqn, role in self._hybrid_fqn_roles + ], + "local_shards": [ + { + "role": role_by_fqn[shard.parameter.fqn], + "shard": _serialize_shard_identity(shard), + } + for _parameter, shard in self._hybrid_local_shard_bindings + ], + "child_groups": child_groups, + } + + @staticmethod + def _normalize_finalized_native_binding(value): + keys = { + "format", + "format_version", + "backup_optimizer", + "manifest", + "routing", + "local_shards", + "child_groups", + } + if type(value) is not dict or set(value) != keys: + raise ValueError("Hybrid finalized native binding has an invalid schema") + if value["format"] != _HYBRID_NATIVE_BINDING_FORMAT: + raise ValueError("Hybrid finalized native binding has an unsupported format") + if type(value["format_version"]) is not int or value["format_version"] != _HYBRID_NATIVE_BINDING_VERSION: + raise ValueError("Hybrid finalized native binding has an unsupported format_version") + if value["backup_optimizer"] not in {"gefen", "adamw"}: + raise ValueError("Hybrid finalized native binding has an invalid backup policy") + if type(value["routing"]) is not list: + raise ValueError("Hybrid finalized native routing must be a list") + routing = [] + for item in value["routing"]: + if type(item) is not dict or set(item) != {"fqn", "role"}: + raise ValueError("Hybrid finalized native routing entry has an invalid schema") + if type(item["fqn"]) is not str or item["role"] not in {"muon", "backup"}: + raise ValueError("Hybrid finalized native routing entry is invalid") + routing.append({"fqn": item["fqn"], "role": item["role"]}) + if len({item["fqn"] for item in routing}) != len(routing): + raise ValueError("Hybrid finalized native routing contains duplicate FQNs") + if routing != sorted(routing, key=lambda item: item["fqn"]): + raise ValueError("Hybrid finalized native routing is not canonical") + + if type(value["local_shards"]) is not list: + raise ValueError("Hybrid finalized native local shards must be a list") + local_shards = [] + for item in value["local_shards"]: + if type(item) is not dict or set(item) != {"role", "shard"}: + raise ValueError("Hybrid finalized native local shard has an invalid schema") + if item["role"] not in {"muon", "backup"}: + raise ValueError("Hybrid finalized native local shard has an invalid role") + local_shards.append( + { + "role": item["role"], + "shard": _normalize_shard_identity(item["shard"]), + } + ) + + if type(value["child_groups"]) is not list: + raise ValueError("Hybrid finalized native child groups must be a list") + child_groups = [] + seen_roles = set() + for item in value["child_groups"]: + if type(item) is not dict or set(item) != {"role", "groups"}: + raise ValueError("Hybrid finalized native child group entry has an invalid schema") + role = item["role"] + groups = item["groups"] + if role not in {"muon", "backup"} or role in seen_roles or type(groups) is not list: + raise ValueError("Hybrid finalized native child group entry is invalid") + seen_roles.add(role) + normalized_groups = [] + for group in groups: + if type(group) is not list or any(type(fqn) is not str for fqn in group): + raise ValueError("Hybrid finalized native child slots must be FQN lists") + normalized_groups.append(list(group)) + child_groups.append({"role": role, "groups": normalized_groups}) + return { + "format": _HYBRID_NATIVE_BINDING_FORMAT, + "format_version": _HYBRID_NATIVE_BINDING_VERSION, + "backup_optimizer": value["backup_optimizer"], + "manifest": _normalize_sharding_manifest(value["manifest"]), + "routing": routing, + "local_shards": local_shards, + "child_groups": child_groups, + } + + def _validate_finalized_native_binding(self, value) -> None: + normalized = self._normalize_finalized_native_binding(value) + if normalized != self._serialized_finalized_native_binding(): + raise ValueError("Hybrid finalized native checkpoint binding or routing differs from the live optimizer") + @property def _gefen_codebook_process_group(self): # Read-only alias: the scoped-step protocol borrowed from Gefen @@ -1339,33 +1795,75 @@ def step(self, closure=None): return loss def state_dict(self): - self._assert_finalized_binding_layout() + self._run_rank_local_checkpoint_phase( + "initial state-dict layout validation", + 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() + def run_pre_hooks(): + for pre_hook in self._optimizer_state_dict_pre_hooks.values(): + pre_hook(self) + + self._run_rank_local_checkpoint_phase("state-dict pre-hook", run_pre_hooks) + self._run_rank_local_checkpoint_phase( + "post-hook state-dict layout validation", + 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, + "muon": ( + self._run_rank_local_checkpoint_phase( + "muon state-dict serialization", + self.muon.state_dict, + ) + if self.muon is not None + else None + ), + "backup": ( + self._run_rank_local_checkpoint_phase( + "backup state-dict serialization", + self.backup.state_dict, + ) + if self.backup is not None + else None + ), "backup_optimizer": self.backup_optimizer, } - for post_hook in self._optimizer_state_dict_post_hooks.values(): - hook_result = post_hook(self, state_dict) - if hook_result is not None: - state_dict = hook_result - return state_dict + if self._hybrid_post_sharding_finalized: + state_dict[_HYBRID_NATIVE_BINDING_KEY] = self._run_rank_local_checkpoint_phase( + "finalized state-dict binding serialization", + self._serialized_finalized_native_binding, + ) + + def run_post_hooks(): + prepared = state_dict + for post_hook in self._optimizer_state_dict_post_hooks.values(): + hook_result = post_hook(self, prepared) + if hook_result is not None: + prepared = hook_result + return prepared + + return self._run_rank_local_checkpoint_phase( + "state-dict post-hook", + run_post_hooks, + ) 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)) + rebinding_children = GefenMuonHybrid._gefen_rebinding_children(self) + rebinding_ready = bool(rebinding_children) + codebook_scope_supported = any( + type(child) in {Gefen, GefenMuon} + for _role, child in rebinding_children + ) identity_ready = GefenMuonHybrid._canonical_identity_ready(self) except Exception: rebinding_ready = False + codebook_scope_supported = False identity_ready = False try: from gefen.portable_hybrid import _hybrid_portable_contract_support @@ -1396,9 +1894,18 @@ def optimizer_contract(self) -> OptimizerContract: muon=muon_contract, backup=backup_contract, backup_implementation=backup_implementation, + adamw_backup=type(self.backup) is torch.optim.AdamW, + native_finalized_layouts=( + frozenset( + shard.layout + for _parameter, shard in self._hybrid_local_shard_bindings + ) + if identity_ready + else frozenset() + ), canonical_parameter_fqns=identity_ready, stable_shard_identity=identity_ready, - explicit_process_group_codebook_scope=rebinding_ready, + explicit_process_group_codebook_scope=codebook_scope_supported, shard_rebinding=rebinding_ready, post_sharding=rebinding_ready, canonical_global_same_topology=canonical_global_same_topology, @@ -1445,57 +1952,102 @@ def import_portable_state( ) def load_state_dict(self, state_dict): - self._assert_finalized_binding_layout() + self._run_rank_local_checkpoint_phase( + "initial load layout validation", + 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. - 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() + def run_pre_hooks(): + prepared = state_dict.copy() + for pre_hook in self._optimizer_load_state_dict_pre_hooks.values(): + hook_result = pre_hook(self, prepared) + if hook_result is not None: + prepared = hook_result + return prepared + + state_dict = self._run_rank_local_checkpoint_phase( + "load pre-hook", + run_pre_hooks, + ) + self._run_rank_local_checkpoint_phase( + "post-hook load layout validation", + 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"} # checkpoint (e.g. one consolidated/converted by FSDP or DeepSpeed # tooling) ran on with ZEROED momentum -- a quiet correctness bug. - if "muon" not in state_dict and "backup" not in state_dict: - raise ValueError( - "GefenMuonHybrid.load_state_dict expects the hybrid's own nested " - 'schema {{"muon": , "backup": ' - ', "backup_optimizer": "gefen" | ' - '"adamw"}} (what GefenMuonHybrid.state_dict() ' - "saves), but got keys {}. Consolidated/converted checkpoints in " - 'the standard {{"state", "param_groups"}} layout are not ' - "supported; resume from the hybrid-saved checkpoint " - "instead.".format(sorted(map(str, state_dict.keys()))) - ) - # A half present here but missing/None in the checkpoint (or vice - # versa) would silently resume that half from scratch (zeroed - # momentum), so reject those presence mismatches before loading either - # child. The backup-backend check below runs before child loads too. - for attr in ("muon", "backup"): - sub = getattr(self, attr) - sub_state = state_dict.get(attr) - if sub is None and sub_state is not None: + def validate_parent_schema(): + allowed_keys = { + "muon", + "backup", + "backup_optimizer", + _HYBRID_NATIVE_BINDING_KEY, + } + if "muon" not in state_dict and "backup" not in state_dict: + raise ValueError( + "GefenMuonHybrid.load_state_dict expects the hybrid's own nested " + 'schema {{"muon": , "backup": ' + ', "backup_optimizer": "gefen" | ' + '"adamw"}} (what GefenMuonHybrid.state_dict() ' + "saves), but got keys {}. Consolidated/converted checkpoints in " + 'the standard {{"state", "param_groups"}} layout are not ' + "supported; resume from the hybrid-saved checkpoint " + "instead.".format(sorted(map(str, state_dict.keys()))) + ) + missing_halves = {"muon", "backup"} - set(state_dict) + if missing_halves: raise ValueError( - "GefenMuonHybrid.load_state_dict: the checkpoint carries a " - '"{}" half but this optimizer has none (was it constructed ' - "with a different parameter split?)".format(attr) + "GefenMuonHybrid.load_state_dict outer schema is missing {}".format( + " and ".join(sorted(missing_halves)) + ) ) - if sub is not None and sub_state is None: + if not set(state_dict).issubset(allowed_keys): raise ValueError( - "GefenMuonHybrid.load_state_dict: this optimizer has a " - '"{}" half but the checkpoint carries none; loading would ' - "silently reset its momentum/state".format(attr) + "GefenMuonHybrid.load_state_dict received an invalid exact outer schema" ) - # Old hybrid checkpoints predate the selector and necessarily contain a - # Gefen backup. Guard the backend before loading either child: torch's - # 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") + finalized_binding = state_dict.get(_HYBRID_NATIVE_BINDING_KEY) + if self._hybrid_post_sharding_finalized: + if finalized_binding is None: + raise ValueError( + "a finalized Hybrid requires a versioned native binding guard" + ) + self._validate_finalized_native_binding(finalized_binding) + elif finalized_binding is not None: + raise ValueError( + "an unfinalized Hybrid cannot load finalized native shard state" + ) + # A half present here but missing/None in the checkpoint (or vice + # versa) would silently resume that half from scratch (zeroed + # momentum), so reject those presence mismatches before loading either + # child. The backup-backend check below runs before child loads too. + for attr in ("muon", "backup"): + sub = getattr(self, attr) + sub_state = state_dict.get(attr) + if sub is None and sub_state is not None: + raise ValueError( + "GefenMuonHybrid.load_state_dict: the checkpoint carries a " + '"{}" half but this optimizer has none (was it constructed ' + "with a different parameter split?)".format(attr) + ) + if sub is not None and sub_state is None: + raise ValueError( + "GefenMuonHybrid.load_state_dict: this optimizer has a " + '"{}" half but the checkpoint carries none; loading would ' + "silently reset its momentum/state".format(attr) + ) + # Old hybrid checkpoints predate the selector and necessarily contain a + # Gefen backup. Guard the backend before loading either child: torch's + # 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 None: + return + 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 " @@ -1504,55 +2056,60 @@ 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. - muon_staged = None - if self.muon is not None: - muon_staged = self.muon._prepare_load_state_dict(state_dict["muon"]) - - if self.backup is None: - # muon-only hybrid: publish the single staged child. - if muon_staged is not None: - self.muon._publish_load_state_dict(muon_staged) - elif isinstance(self.backup, Gefen): - # Both halves expose the Gefen staging primitive: stage both (all - # validation and rejection happens here, before any mutation), then - # publish through non-throwing swaps so neither can fail mid-commit. - backup_staged = self.backup._prepare_load_state_dict(state_dict["backup"]) - if muon_staged is not None: - self.muon._publish_load_state_dict(muon_staged) - self.backup._publish_load_state_dict(backup_staged) - else: - # Foreign backup (torch ``AdamW``): no non-throwing staging - # primitive, but ``torch.optim.Optimizer.load_state_dict`` validates - # the group structure and casts every state tensor before its single - # ``__setstate__``, so it is itself fail-before-mutation. Commit it - # first, while the muon child is only staged (never published); if it - # raises, the muon child is untouched. The muon child was already - # validated above, so publishing its swap afterwards cannot fail. - self.backup.load_state_dict(state_dict["backup"]) - if muon_staged is not None: - self.muon._publish_load_state_dict(muon_staged) - - for post_hook in self._optimizer_load_state_dict_post_hooks.values(): - post_hook(self) + + self._run_rank_local_checkpoint_phase( + "parent load validation", + validate_parent_schema, + ) + # Stage every child, including AdamW, before publishing any child. The + # commit loop below contains only core dict/container swaps and invokes + # no callbacks. Child post-hooks run only after the complete composite + # core is visible, so a callback can never observe a half-loaded Hybrid. + staged_children = [] + for role, child in self._gefen_rebinding_children(): + child_state = state_dict[role] + + def stage_child(child=child, child_state=child_state): + return ( + prepare_adamw_load_state_dict(child, child_state) + if type(child) is torch.optim.AdamW + else child._prepare_load_state_dict(child_state) + ) + + staged = self._run_rank_local_checkpoint_phase( + "{} load staging".format(role), + stage_child, + ) + staged_children.append((role, child, staged)) + + for _role, child, staged in staged_children: + if type(child) is torch.optim.AdamW: + commit_adamw_stage(child, staged) + else: + child._commit_staged_load_state_dict(staged) + self._hybrid_layout_forensics_verdict = None + + for role, child, _staged in staged_children: + def run_child_post_hooks(child=child): + if type(child) is torch.optim.AdamW: + run_adamw_load_post_hooks(child) + else: + for post_hook in child._optimizer_load_state_dict_post_hooks.values(): + post_hook(child) + + self._run_rank_local_checkpoint_phase( + "{} load post-hook".format(role), + run_child_post_hooks, + ) + + def run_parent_post_hooks(): + for post_hook in self._optimizer_load_state_dict_post_hooks.values(): + post_hook(self) + + self._run_rank_local_checkpoint_phase( + "load post-hook", + run_parent_post_hooks, + ) @staticmethod def _auto_split(params_or_model, backup_substrings): diff --git a/src/gefen/portable.py b/src/gefen/portable.py index d759a94..b0c3f87 100644 --- a/src/gefen/portable.py +++ b/src/gefen/portable.py @@ -480,6 +480,51 @@ def _expand_factored_second_moment( return _finish_output(dense, name="dense second moment") +def _expand_factored_second_moment_live_fp32_v1( + row, + column, + *, + logical_shape, + step, +) -> torch.Tensor: + """Expand factored state with the live decomposed fp32 operator ordering.""" + + _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") + if math.prod(shape) == 0: + return _finish_output( + _new_output(shape, dtype=torch.float32, device=row.device), + name="live-fp32 dense second moment", + ) + + # This ordering is the versioned numerical contract. It intentionally + # rounds the outer product before division and mirrors Gefen's live + # decomposed update, including its minimum-normal denominator clamp. + denominator = row.detach().mean().clamp_(min=torch.finfo(torch.float32).tiny) + dense = torch.outer(row.detach(), column.detach()).div_(denominator) + return _finish_output(dense, name="live-fp32 dense second moment") + + def _project_factored_second_moment( dense, *, diff --git a/src/gefen/portable_dcp.py b/src/gefen/portable_dcp.py index c698bcb..4cfb224 100644 --- a/src/gefen/portable_dcp.py +++ b/src/gefen/portable_dcp.py @@ -14,7 +14,6 @@ _CanonicalWireLimits, _CanonicalWirePlan, _parse_canonical_wire_metadata, - _prepare_canonical_wire_value, _reconstruct_canonical_wire_value, ) @@ -27,6 +26,12 @@ _DCP_LOAD_PREFLIGHT_TRANSACTION = "gefen-portable-dcp-load-preflight-v1" +def _prepare_canonical_wire_value(value, limits): + from gefen.portable_wire import _prepare_canonical_wire_value as prepare + + return prepare(value, limits) + + 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") @@ -206,7 +211,7 @@ def _preflight_dcp_operation( _validate_dcp_runtime(binding) context = { **base_context, - "dcp_envelope_version": 1, + "dcp_envelope_version": 2, "dcp_namespace": normalized_namespace, "dcp_storage": storage_identity, "transaction_id": normalized_transaction, @@ -451,7 +456,6 @@ def save_portable_dcp( ): """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( @@ -464,38 +468,33 @@ def save_portable_dcp( 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 + from gefen.hybrid import GefenMuonHybrid - _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, - ) + if type(optimizer) is GefenMuonHybrid: + from gefen.portable_dcp_hybrid_sharded import ( + _save_sharded_hybrid_dcp, + ) + + return _save_sharded_hybrid_dcp( + optimizer, + binding=binding, + storage_writer=storage_writer, + transaction_id=transaction_id, + limits=limits, + namespace=namespace, + ) + else: + from gefen.portable_dcp_sharded import _save_sharded_portable_dcp + + return _save_sharded_portable_dcp( + optimizer, + binding=binding, + storage_writer=storage_writer, + transaction_id=transaction_id, + limits=limits, + namespace=namespace, + ) + raise AssertionError("unreachable portable DCP optimizer dispatch") def load_portable_dcp( @@ -522,6 +521,68 @@ def load_portable_dcp( storage_type=StorageReader, operation="load", ) + from gefen.hybrid import GefenMuonHybrid + + modern = False + checkpoint_metadata = None + error = None + try: + checkpoint_metadata = storage_reader.read_metadata() + if type(optimizer) is GefenMuonHybrid: + from gefen.portable_dcp_hybrid_sharded import ( + _metadata_storage_spec, + ) + else: + from gefen.portable_dcp_sharded import _metadata_storage_spec + + modern = ( + _metadata_storage_spec( + checkpoint_metadata, + namespace=namespace, + limits=limits, + ) + is not None + ) + except Exception as exc: + error = exc + from gefen.portable_collective import _collective_unanimous_status + + _collective_unanimous_status( + binding, + error, + operation="portable_dcp_load_envelope", + transaction_id=transaction_id, + context_digest=bytes(31) + bytes((2 if modern else 1,)), + limits=limits._wire_limits(collective=True), + ) + assert checkpoint_metadata is not None + if modern: + if type(optimizer) is GefenMuonHybrid: + from gefen.portable_dcp_hybrid_sharded import ( + _load_sharded_hybrid_dcp, + ) + + return _load_sharded_hybrid_dcp( + optimizer, + binding=binding, + storage_reader=storage_reader, + checkpoint_metadata=checkpoint_metadata, + transaction_id=transaction_id, + limits=limits, + namespace=namespace, + ) + else: + from gefen.portable_dcp_sharded import _load_sharded_portable_dcp + + return _load_sharded_portable_dcp( + optimizer, + binding=binding, + storage_reader=storage_reader, + checkpoint_metadata=checkpoint_metadata, + transaction_id=transaction_id, + limits=limits, + namespace=namespace, + ) wire_limits = limits._wire_limits(collective=True) state = {namespace: {}} dcp.load( diff --git a/src/gefen/portable_dcp_hybrid_sharded.py b/src/gefen/portable_dcp_hybrid_sharded.py new file mode 100644 index 0000000..31dd00d --- /dev/null +++ b/src/gefen/portable_dcp_hybrid_sharded.py @@ -0,0 +1,972 @@ +"""Scalable sharded-DCP envelope for Gefen-backed Hybrid optimizers.""" + +from __future__ import annotations + +import hashlib + +import torch + +from gefen.checkpoint import CheckpointProcessGroupBinding + + +_COMPOSITE_FORMAT = "gefen.portable_dcp_sharded_composite" +_COMPOSITE_FORMAT_VERSION = 2 +_COMPOSITE_COVERAGE = "global_logical_composite_optimizer" +_COMPOSITE_IMPLEMENTATION = "gefen.GefenMuonHybrid" +_COMPOSITE_METADATA_KEY = "__gefen_portable_composite_metadata_v2__" +_ROLES = ("muon", "backup") +_CHILD_IMPLEMENTATIONS = { + "muon": "gefen.GefenMuon", + "backup": "gefen.Gefen", +} + + +def _child_namespace(namespace: str, role: str, limits) -> str: + if role not in _ROLES: + raise ValueError("invalid sharded Hybrid child role") + from gefen.portable_dcp import _require_namespace + + return _require_namespace( + "{}-gefen-{}".format(namespace, role), + limits, + ) + + +def _encode_metadata(value, *, limits) -> torch.Tensor: + from gefen.portable_wire import _prepare_canonical_wire_value + + plan = _prepare_canonical_wire_value( + value, + limits._wire_limits(collective=True), + ) + if plan.payload_tensors: + raise RuntimeError("sharded Hybrid metadata unexpectedly contains tensors") + return torch.frombuffer(bytearray(plan.metadata), dtype=torch.uint8) + + +def _completed_metadata(value, *, limits): + from gefen.portable_schema import portable_state_digest + + semantic = dict(value) + completed = { + **semantic, + "completion": { + "metadata_digest": portable_state_digest(semantic), + }, + } + return completed, _encode_metadata(completed, limits=limits) + + +def _validate_aggregate_child_limits( + child_metadata, + *, + binding: CheckpointProcessGroupBinding, + limits, +): + from gefen.portable_dcp_sharded import _chunk_segment + + fields = [ + field + for metadata in child_metadata.values() + for field in metadata["fields"] + ] + if ( + len(fields) > limits.max_tensors + or len(fields) > limits.max_container_items + ): + raise ValueError("sharded Hybrid fields exceed count limits") + coordinate = binding.identity.ordered_members.index(binding.local_member) + parts = len(binding.identity.ordered_members) + expected_chunks = 1 + len(child_metadata) + len(fields) * parts + if ( + expected_chunks > limits.max_container_items + or expected_chunks > limits.max_tree_nodes + ): + raise ValueError("sharded Hybrid DCP chunks exceed aggregate limits") + global_bytes = 0 + local_bytes = 0 + for field in fields: + field_bytes = field["numel"] * 4 + if global_bytes > limits.max_collective_tensor_bytes - field_bytes: + raise ValueError("sharded Hybrid fields exceed max_collective_tensor_bytes") + global_bytes += field_bytes + chunk_bytes = _chunk_segment( + field["numel"], + parts, + coordinate, + ).length * 4 + if local_bytes > limits.max_fragment_tensor_bytes - chunk_bytes: + raise ValueError("sharded Hybrid local fields exceed max_fragment_tensor_bytes") + local_bytes += chunk_bytes + + +def _decode_metadata(value, *, namespace: str, limits): + from gefen.portable_wire import ( + _parse_canonical_wire_metadata, + _reconstruct_canonical_wire_value, + ) + + if ( + type(value) is not torch.Tensor + or value.dtype is not torch.uint8 + or value.device.type != "cpu" + or value.ndim != 1 + or not value.is_contiguous() + or value.numel() < 1 + or value.numel() > limits.max_metadata_bytes + ): + raise ValueError("sharded Hybrid metadata tensor is invalid") + prepared = _parse_canonical_wire_metadata( + bytes(memoryview(value.numpy())), + limits=limits._wire_limits(collective=True), + ) + if prepared.tensor_specs: + raise ValueError("sharded Hybrid metadata must not contain tensor payloads") + return _normalize_metadata( + _reconstruct_canonical_wire_value(prepared, ()), + namespace=namespace, + limits=limits, + ) + + +def _normalize_metadata(value, *, namespace: str, limits): + from gefen.portable_dcp_sharded import _strict_counter, _strict_record + from gefen.portable_hybrid import _normalize_routing + from gefen.portable_identity import ( + _normalize_process_group_identity, + _parse_process_group_identity, + ) + from gefen.portable_schema import portable_state_digest + + _strict_record( + value, + { + "format", + "format_version", + "coverage", + "implementation", + "backup_optimizer", + "source_process_group", + "routing", + "common", + "children", + "completion", + }, + name="sharded Hybrid metadata", + ) + if ( + value["format"] != _COMPOSITE_FORMAT + or value["format_version"] != _COMPOSITE_FORMAT_VERSION + or value["coverage"] != _COMPOSITE_COVERAGE + or value["implementation"] != _COMPOSITE_IMPLEMENTATION + or value["backup_optimizer"] != "gefen" + ): + raise ValueError("unsupported sharded Hybrid metadata identity") + process_group_record = _normalize_process_group_identity( + value["source_process_group"] + ) + process_group = _parse_process_group_identity(process_group_record) + if len(process_group.ordered_members) > limits.max_members: + raise ValueError("sharded Hybrid source process group exceeds limits") + routing = _normalize_routing(value["routing"]) + if len(routing) > limits.max_container_items: + raise ValueError("sharded Hybrid routing exceeds limits") + common = _strict_record( + value["common"], + {"gefen_global_step", "gefen_deterministic"}, + name="sharded Hybrid common state", + ) + _strict_counter(common["gefen_global_step"], name="gefen_global_step") + if type(common["gefen_deterministic"]) is not bool: + raise ValueError("sharded Hybrid deterministic policy must be bool") + children = _strict_record( + value["children"], + set(_ROLES), + name="sharded Hybrid children", + ) + normalized_children = {} + for role in _ROLES: + child = children[role] + if child is None: + normalized_children[role] = None + continue + _strict_record( + child, + {"namespace", "implementation", "metadata_digest"}, + name="sharded Hybrid child", + ) + if ( + child["namespace"] != _child_namespace(namespace, role, limits) + or child["implementation"] != _CHILD_IMPLEMENTATIONS[role] + or type(child["metadata_digest"]) is not str + or len(child["metadata_digest"]) != 64 + ): + raise ValueError("sharded Hybrid child descriptor is invalid") + try: + bytes.fromhex(child["metadata_digest"]) + except ValueError as exc: + raise ValueError("sharded Hybrid child digest is invalid") from exc + normalized_children[role] = dict(child) + if all(child is None for child in normalized_children.values()): + raise ValueError("sharded Hybrid metadata requires at least one child") + completion = _strict_record( + value["completion"], + {"metadata_digest"}, + name="sharded Hybrid completion", + ) + digest = completion["metadata_digest"] + if type(digest) is not str or len(digest) != 64: + raise ValueError("sharded Hybrid completion digest is invalid") + try: + bytes.fromhex(digest) + except ValueError as exc: + raise ValueError("sharded Hybrid completion digest is invalid") from exc + semantic = {key: value[key] for key in value if key != "completion"} + if portable_state_digest(semantic) != digest: + raise ValueError("sharded Hybrid metadata digest mismatch") + return { + **value, + "source_process_group": process_group_record, + "routing": routing, + "common": dict(common), + "children": normalized_children, + "completion": dict(completion), + } + + +def _metadata_storage_spec(checkpoint_metadata, *, namespace: str, limits): + from torch.distributed.checkpoint.metadata import Metadata + from gefen.portable_dcp import _flat_key, _validate_tensor_metadata + + if type(checkpoint_metadata) is not Metadata: + raise TypeError("portable DCP requires exact checkpoint Metadata") + entries = checkpoint_metadata.state_dict_metadata + if type(entries) is not dict: + raise TypeError("portable DCP metadata entries must be a dict") + key = _flat_key(namespace, _COMPOSITE_METADATA_KEY) + if key not in entries: + return None + dtype, shape, nbytes = _validate_tensor_metadata( + entries[key], + name="sharded Hybrid DCP metadata", + limits=limits._wire_limits(collective=True), + ) + if ( + dtype is not torch.uint8 + or len(shape) != 1 + or nbytes < 1 + or nbytes > limits.max_metadata_bytes + ): + raise ValueError("sharded Hybrid DCP metadata exceeds limits") + return shape + + +def _read_metadata( + *, + storage_reader, + checkpoint_metadata, + binding: CheckpointProcessGroupBinding, + namespace: str, + limits, + transaction_id: str, +): + import torch.distributed.checkpoint as dcp + from gefen.portable_collective import _collective_unanimous_status + + shape = None + tensor = None + state = None + planner = None + error = None + try: + shape = _metadata_storage_spec( + checkpoint_metadata, + namespace=namespace, + limits=limits, + ) + if shape is not None: + tensor = torch.empty(shape, dtype=torch.uint8, device="cpu") + state = {namespace: {_COMPOSITE_METADATA_KEY: tensor}} + planner = dcp.DefaultLoadPlanner(allow_partial_load=True) + except Exception as exc: + error = exc + shape_bytes = repr(tuple(shape) if shape is not None else None).encode("ascii") + _collective_unanimous_status( + binding, + error, + operation="portable_dcp_hybrid_metadata_prepare", + transaction_id=transaction_id, + context_digest=hashlib.sha256( + b"gefen.portable_dcp_hybrid.metadata.v1\0" + + namespace.encode("utf-8") + + b"\0" + + shape_bytes + ).digest(), + limits=limits._wire_limits(collective=True), + ) + if shape is None: + return None + assert tensor is not None and state is not None and planner is not None + dcp.load( + state, + storage_reader=storage_reader, + planner=planner, + process_group=binding.process_group, + ) + return _decode_metadata(tensor, namespace=namespace, limits=limits) + + +def _prepare_save_state( + optimizer, + *, + binding: CheckpointProcessGroupBinding, + transaction_id: str, + limits, + namespace: str, +): + from gefen.portable_collective import _collective_unanimous_status + from gefen.portable_dcp_sharded import ( + _DCP_SHARDED_METADATA_KEY, + _decode_metadata_tensor, + _prepare_sharded_save_state, + ) + from gefen.portable_hybrid import ( + _hybrid_child_transaction_id, + _hybrid_portable_live_token, + _validate_hybrid_portable_readiness, + ) + from gefen.portable_identity import _serialize_process_group_identity + + children = None + routing = None + live_token = None + error = None + try: + children, routing, _layouts = _validate_hybrid_portable_readiness( + optimizer, + binding, + ) + live_token = _hybrid_portable_live_token(optimizer) + except Exception as exc: + error = exc + _collective_unanimous_status( + binding, + error, + operation="portable_dcp_hybrid_save_local", + transaction_id=transaction_id, + context_digest=bytes(32), + limits=limits._wire_limits(collective=True), + ) + assert children is not None and routing is not None and live_token is not None + + state = {} + child_descriptors = {role: None for role in _ROLES} + child_metadata = {} + for role, child in children: + child_transaction = _hybrid_child_transaction_id( + transaction_id, + role, + "export", + ) + child_namespace = None + error = None + try: + child_namespace = _child_namespace(namespace, role, limits) + except Exception as exc: + error = exc + child_context = hashlib.sha256( + b"gefen.portable_dcp_hybrid.save_child.v1\0" + + role.encode("ascii") + ).digest() + _collective_unanimous_status( + binding, + error, + operation="portable_dcp_hybrid_save_child_prep_{}".format(role), + transaction_id=child_transaction, + context_digest=child_context, + limits=limits._wire_limits(collective=True), + ) + assert child_namespace is not None + + child_state = None + metadata = None + error = None + try: + child_state = _prepare_sharded_save_state( + child, + binding=binding, + transaction_id=child_transaction, + limits=limits, + namespace=child_namespace, + ) + metadata = _decode_metadata_tensor( + child_state[child_namespace][_DCP_SHARDED_METADATA_KEY], + limits=limits, + ) + child_metadata[role] = metadata + child_descriptors[role] = { + "namespace": child_namespace, + "implementation": metadata["implementation"], + "metadata_digest": metadata["completion"]["metadata_digest"], + } + state.update(child_state) + except Exception as exc: + error = exc + _collective_unanimous_status( + binding, + error, + operation="portable_dcp_hybrid_save_child_{}".format(role), + transaction_id=child_transaction, + context_digest=child_context, + limits=limits._wire_limits(collective=True), + ) + assert child_state is not None and metadata is not None + + common_values = [ + ( + metadata["common"]["gefen_global_step"], + metadata["common"]["gefen_deterministic"], + ) + for metadata in child_metadata.values() + ] + _validate_aggregate_child_limits( + child_metadata, + binding=binding, + limits=limits, + ) + if any(value != common_values[0] for value in common_values[1:]): + raise RuntimeError("sharded Hybrid child common state disagrees") + semantic = { + "format": _COMPOSITE_FORMAT, + "format_version": _COMPOSITE_FORMAT_VERSION, + "coverage": _COMPOSITE_COVERAGE, + "implementation": _COMPOSITE_IMPLEMENTATION, + "backup_optimizer": "gefen", + "source_process_group": _serialize_process_group_identity( + binding.identity + ), + "routing": routing, + "common": { + "gefen_global_step": common_values[0][0], + "gefen_deterministic": common_values[0][1], + }, + "children": child_descriptors, + } + metadata, metadata_tensor = _completed_metadata(semantic, limits=limits) + _normalize_metadata(metadata, namespace=namespace, limits=limits) + if live_token != _hybrid_portable_live_token(optimizer): + raise RuntimeError("live Hybrid state changed during sharded DCP save preparation") + state[namespace] = {_COMPOSITE_METADATA_KEY: metadata_tensor} + metadata_bytes = metadata_tensor.numel() + for role, descriptor in child_descriptors.items(): + if descriptor is None: + continue + from gefen.portable_dcp_sharded import _DCP_SHARDED_METADATA_KEY + + metadata_bytes += state[descriptor["namespace"]][ + _DCP_SHARDED_METADATA_KEY + ].numel() + if metadata_bytes > limits.max_collective_metadata_bytes: + raise ValueError("sharded Hybrid metadata exceeds max_collective_metadata_bytes") + return state, metadata + + +def _save_sharded_hybrid_dcp( + optimizer, + *, + binding: CheckpointProcessGroupBinding, + storage_writer, + transaction_id: str, + limits, + namespace: str, +): + import torch.distributed.checkpoint as dcp + from gefen.portable_collective import _collective_unanimous_status + + state = None + metadata = None + planner = None + error = None + try: + state, metadata = _prepare_save_state( + optimizer, + binding=binding, + transaction_id=transaction_id, + limits=limits, + namespace=namespace, + ) + planner = dcp.DefaultSavePlanner() + except Exception as exc: + error = exc + digest = ( + bytes(32) + if metadata is None + else bytes.fromhex(metadata["completion"]["metadata_digest"]) + ) + _collective_unanimous_status( + binding, + error, + operation="portable_dcp_hybrid_save_prepare", + transaction_id=transaction_id, + context_digest=digest, + limits=limits._wire_limits(collective=True), + ) + assert state is not None and planner is not None + result = None + error = None + try: + result = dcp.save( + state, + storage_writer=storage_writer, + planner=planner, + process_group=binding.process_group, + ) + except Exception as exc: + error = exc + _collective_unanimous_status( + binding, + error, + operation="portable_dcp_hybrid_save_dcp_complete", + transaction_id=transaction_id, + context_digest=digest, + limits=limits._wire_limits(collective=True), + ) + return result + + +def _validate_children( + optimizer, + metadata, + child_metadata, + *, + binding: CheckpointProcessGroupBinding, +): + from gefen.portable_hybrid import _validate_hybrid_portable_readiness + from gefen.portable_identity import _parse_process_group_identity + + children, routing, _layouts = _validate_hybrid_portable_readiness( + optimizer, + binding, + ) + child_by_role = dict(children) + if metadata["routing"] != routing: + raise ValueError("sharded Hybrid routing does not match the target") + if { + role: metadata["children"][role] is not None for role in _ROLES + } != {role: role in child_by_role for role in _ROLES}: + raise ValueError("sharded Hybrid child presence does not match the target") + source_group = _parse_process_group_identity(metadata["source_process_group"]) + expected_routing = {} + for role, child in child_by_role.items(): + descriptor = metadata["children"][role] + child_state = child_metadata[role] + if ( + child_state["implementation"] != _CHILD_IMPLEMENTATIONS[role] + or child_state["completion"]["metadata_digest"] + != descriptor["metadata_digest"] + or _parse_process_group_identity(child_state["source_process_group"]) + != source_group + or child_state["common"]["gefen_global_step"] + != metadata["common"]["gefen_global_step"] + or child_state["common"]["gefen_deterministic"] + is not metadata["common"]["gefen_deterministic"] + ): + raise ValueError("sharded Hybrid child metadata conflicts with its envelope") + for fqn in child_state["parameters"]: + if fqn in expected_routing: + raise ValueError("sharded Hybrid child parameter FQNs overlap") + expected_routing[fqn] = role + if metadata["routing"] != { + fqn: expected_routing[fqn] for fqn in sorted(expected_routing) + }: + raise ValueError("sharded Hybrid routing does not match its children") + return tuple((role, child_by_role[role]) for role in _ROLES if role in child_by_role) + + +def _expected_checkpoint_keys(metadata, child_metadata, *, namespace: str): + from gefen.portable_dcp import _flat_key + + keys = {_flat_key(namespace, _COMPOSITE_METADATA_KEY)} + for role, child in child_metadata.items(): + child_namespace = metadata["children"][role]["namespace"] + from gefen.portable_dcp_sharded import _DCP_SHARDED_METADATA_KEY + + keys.add(_flat_key(child_namespace, _DCP_SHARDED_METADATA_KEY)) + keys.update( + _flat_key(child_namespace, field["key"]) + for field in child["fields"] + ) + return keys + + +def _validate_checkpoint_entries( + checkpoint_metadata, + metadata, + child_metadata, + *, + namespace: str, + binding: CheckpointProcessGroupBinding, + limits, +): + from gefen.portable_dcp import _flat_key + from gefen.portable_dcp_sharded import _validate_sharded_dcp_entries + + expected = _expected_checkpoint_keys( + metadata, + child_metadata, + namespace=namespace, + ) + if set(checkpoint_metadata.state_dict_metadata) != expected: + raise ValueError("sharded Hybrid checkpoint has unexpected tensor keys") + if _flat_key(namespace, _COMPOSITE_METADATA_KEY) not in expected: + raise ValueError("sharded Hybrid checkpoint is missing its envelope") + total_chunks = 0 + metadata_bytes = 0 + metadata_keys = {_flat_key(namespace, _COMPOSITE_METADATA_KEY)} + from gefen.portable_dcp_sharded import _DCP_SHARDED_METADATA_KEY + + metadata_keys.update( + _flat_key(metadata["children"][role]["namespace"], _DCP_SHARDED_METADATA_KEY) + for role in child_metadata + ) + for key in expected: + entry = checkpoint_metadata.state_dict_metadata[key] + chunks = getattr(entry, "chunks", None) + if type(chunks) is not list: + raise TypeError("sharded Hybrid DCP entries require chunk lists") + total_chunks += len(chunks) + if key in metadata_keys: + size = getattr(entry, "size", None) + if type(size) is not torch.Size or len(size) != 1: + raise ValueError("sharded Hybrid metadata geometry is invalid") + metadata_bytes += size[0] + if ( + total_chunks > limits.max_container_items + or total_chunks > limits.max_tree_nodes + ): + raise ValueError("sharded Hybrid DCP chunks exceed aggregate limits") + if metadata_bytes > limits.max_collective_metadata_bytes: + raise ValueError("sharded Hybrid metadata exceeds max_collective_metadata_bytes") + for role, child in child_metadata.items(): + _validate_sharded_dcp_entries( + checkpoint_metadata, + child, + namespace=metadata["children"][role]["namespace"], + binding=binding, + limits=limits, + checkpoint_keys=expected, + ) + return expected + + +def _load_sharded_hybrid_dcp( + optimizer, + *, + binding: CheckpointProcessGroupBinding, + storage_reader, + checkpoint_metadata, + transaction_id: str, + limits, + namespace: str, +): + from gefen.portable_collective import _collective_unanimous_status + from gefen.portable_dcp_sharded import ( + _load_sharded_payloads, + _metadata_storage_spec as _child_metadata_storage_spec, + _projected_target_byte_plan, + _read_sharded_metadata, + _stage_sharded_import, + ) + from gefen.portable_hybrid import ( + _hybrid_child_transaction_id, + _hybrid_portable_live_token, + _validate_hybrid_portable_readiness, + ) + + metadata = None + error = None + try: + metadata = _read_metadata( + storage_reader=storage_reader, + checkpoint_metadata=checkpoint_metadata, + binding=binding, + namespace=namespace, + limits=limits, + transaction_id=transaction_id, + ) + if metadata is None: + raise ValueError("checkpoint does not contain a sharded Hybrid envelope") + except Exception as exc: + error = exc + digest = ( + bytes(32) + if metadata is None + else bytes.fromhex(metadata["completion"]["metadata_digest"]) + ) + _collective_unanimous_status( + binding, + error, + operation="portable_dcp_hybrid_load_metadata", + transaction_id=transaction_id, + context_digest=digest, + limits=limits._wire_limits(collective=True), + ) + assert checkpoint_metadata is not None and metadata is not None + + child_metadata = {} + children = None + expected_keys = None + composite_live_token = None + child_shapes = {} + error = None + try: + composite_live_token = _hybrid_portable_live_token(optimizer) + root_shape = _metadata_storage_spec( + checkpoint_metadata, + namespace=namespace, + limits=limits, + ) + assert root_shape is not None + aggregate_metadata_bytes = root_shape[0] + for role in _ROLES: + descriptor = metadata["children"][role] + if descriptor is None: + continue + shape = _child_metadata_storage_spec( + checkpoint_metadata, + namespace=descriptor["namespace"], + limits=limits, + ) + if shape is None: + raise ValueError("sharded Hybrid child metadata is missing") + child_shapes[role] = shape + aggregate_metadata_bytes += shape[0] + if aggregate_metadata_bytes > limits.max_collective_metadata_bytes: + raise ValueError("sharded Hybrid metadata exceeds max_collective_metadata_bytes") + except Exception as exc: + error = exc + _collective_unanimous_status( + binding, + error, + operation="portable_dcp_hybrid_load_child_specs", + transaction_id=transaction_id, + context_digest=digest, + limits=limits._wire_limits(collective=True), + ) + assert composite_live_token is not None + + for role in _ROLES: + descriptor = metadata["children"][role] + if descriptor is None: + continue + child_state = None + error = None + try: + child_transaction = _hybrid_child_transaction_id( + transaction_id, + role, + "import", + ) + child_state = _read_sharded_metadata( + storage_reader=storage_reader, + checkpoint_metadata=checkpoint_metadata, + binding=binding, + namespace=descriptor["namespace"], + limits=limits, + transaction_id=child_transaction, + ) + if child_state is None: + raise ValueError("sharded Hybrid child metadata is missing") + except Exception as exc: + error = exc + role_digest = hashlib.sha256(digest + role.encode("ascii")).digest() + _collective_unanimous_status( + binding, + error, + operation="portable_dcp_hybrid_load_child_metadata", + transaction_id=transaction_id, + context_digest=role_digest, + limits=limits._wire_limits(collective=True), + ) + assert child_state is not None + child_metadata[role] = child_state + + error = None + try: + children = _validate_children( + optimizer, + metadata, + child_metadata, + binding=binding, + ) + _validate_aggregate_child_limits( + child_metadata, + binding=binding, + limits=limits, + ) + expected_keys = _validate_checkpoint_entries( + checkpoint_metadata, + metadata, + child_metadata, + namespace=namespace, + binding=binding, + limits=limits, + ) + except Exception as exc: + error = exc + _collective_unanimous_status( + binding, + error, + operation="portable_dcp_hybrid_load_children", + transaction_id=transaction_id, + context_digest=digest, + limits=limits._wire_limits(collective=True), + ) + assert ( + children is not None + and expected_keys is not None + and composite_live_token is not None + ) + + projected_digest = digest + error = None + try: + child_plans = [] + for role, child in children: + ( + child_digest, + child_projected, + child_scratch, + child_send, + ) = _projected_target_byte_plan( + child, + child_metadata[role], + binding=binding, + limits=limits, + ) + child_plans.append( + ( + role, + child_digest, + child_projected, + child_scratch, + child_send, + ) + ) + projected_totals = [0 for _member in binding.identity.ordered_members] + for ( + _role, + _child_digest, + child_projected, + _child_scratch, + _child_send, + ) in child_plans: + for coordinate, value in enumerate(child_projected): + projected_totals[coordinate] += value + if any(value > limits.max_fragment_tensor_bytes for value in projected_totals): + raise ValueError("sharded Hybrid projected target state exceeds max_fragment_tensor_bytes") + hasher = hashlib.sha256() + hasher.update(b"gefen.portable_dcp_hybrid.projected.v1\0") + hasher.update(digest) + for role, child_digest, child_projected, child_scratch, child_send in child_plans: + hasher.update(role.encode("ascii")) + hasher.update(child_digest) + for values in (child_projected, child_scratch, child_send): + for value in values: + hasher.update(value.to_bytes(16, "big", signed=False)) + projected_digest = hasher.digest() + except Exception as exc: + error = exc + _collective_unanimous_status( + binding, + error, + operation="portable_dcp_hybrid_load_projected_limits", + transaction_id=transaction_id, + context_digest=projected_digest, + limits=limits._wire_limits(collective=True), + ) + + staged_children = [] + for role, child in children: + child_state = child_metadata[role] + child_namespace = metadata["children"][role]["namespace"] + child_transaction = _hybrid_child_transaction_id( + transaction_id, + role, + "import", + ) + local_fields = None + staged = None + live_token = None + error = None + try: + shape = child_shapes[role] + local_fields = _load_sharded_payloads( + storage_reader=storage_reader, + checkpoint_metadata=checkpoint_metadata, + metadata=child_state, + metadata_shape=shape, + binding=binding, + namespace=child_namespace, + limits=limits, + transaction_id=child_transaction, + context_digest=bytes.fromhex( + child_state["completion"]["metadata_digest"] + ), + checkpoint_keys=expected_keys, + ) + staged, live_token = _stage_sharded_import( + child, + child_state, + local_fields, + binding=binding, + limits=limits, + transaction_id=child_transaction, + ) + except Exception as exc: + error = exc + role_digest = hashlib.sha256(digest + role.encode("ascii")).digest() + _collective_unanimous_status( + binding, + error, + operation="portable_dcp_hybrid_load_stage", + transaction_id=child_transaction, + context_digest=role_digest, + limits=limits._wire_limits(collective=True), + ) + assert local_fields is not None and staged is not None and live_token is not None + staged_children.append((child, staged, live_token)) + + 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 sharded DCP staging") + from gefen import portable_runtime as runtime + + for child, _staged, live_token in staged_children: + if live_token != runtime._portable_live_token(child): + raise RuntimeError("live Hybrid child changed after sharded DCP staging") + except Exception as exc: + error = exc + _collective_unanimous_status( + binding, + error, + operation="portable_dcp_hybrid_load_freshness", + transaction_id=transaction_id, + context_digest=digest, + limits=limits._wire_limits(collective=True), + ) + + from gefen.gefen import Gefen + + for child, staged, _live_token in staged_children: + Gefen._commit_staged_load_state_dict(child, staged) + dict.__setitem__( + optimizer.__dict__, + "_deterministic", + metadata["common"]["gefen_deterministic"], + ) + + +__all__ = [] diff --git a/src/gefen/portable_dcp_sharded.py b/src/gefen/portable_dcp_sharded.py new file mode 100644 index 0000000..db0ab1a --- /dev/null +++ b/src/gefen/portable_dcp_sharded.py @@ -0,0 +1,4016 @@ +"""DCP-native sharded transport for portable optimizer tensors. + +The helpers in this module keep logical payloads partitioned while converting +between optimizer-owned shard geometry and DCP's canonical one-dimensional +DTensor placement. Portable semantic validation and atomic publication remain +owned by :mod:`gefen.portable_runtime`. +""" + +from __future__ import annotations + +from dataclasses import dataclass +import hashlib +import math + +import torch +import torch.distributed as dist + +from gefen.checkpoint import CheckpointProcessGroupBinding +from gefen.contracts import ( + ParameterLayout, + PlacementKind, + ShardIdentity, +) + + +_DCP_SHARDED_FORMAT = "gefen.portable_dcp_sharded" +_DCP_SHARDED_FORMAT_VERSION = 2 +_DCP_SHARDED_METADATA_KEY = "__gefen_portable_metadata_v2__" +_DCP_SHARDED_FIELD_PREFIX = "__gefen_portable_field_" +_DCP_SHARDED_FIELD_SUFFIX = "__" +_DCP_SHARDED_FIELD_DIGITS = 16 + + +def _make_phase_voter( + binding: CheckpointProcessGroupBinding, + *, + transaction_id: str, + context_digest: bytes, + limits, + operation_prefix: str, +): + """Return a monotonic status voter whose digest commits to each phase label.""" + + from gefen.portable_collective import _collective_unanimous_status + + ordinal = 0 + + def vote(label: str, error: Exception | None) -> None: + nonlocal ordinal + + ordinal += 1 + if type(label) is not str or not label: + label_error = ValueError("collective phase label must be a nonempty string") + if error is None: + error = label_error + label_bytes = label.encode("utf-8") if type(label) is str else b"" + phase_digest = hashlib.sha256( + b"gefen.portable_dcp_sharded.phase.v1\0" + + context_digest + + ordinal.to_bytes(8, "big", signed=False) + + len(label_bytes).to_bytes(8, "big", signed=False) + + label_bytes + ).digest() + _collective_unanimous_status( + binding, + error, + operation="{}_{:08d}".format(operation_prefix, ordinal), + transaction_id=transaction_id, + context_digest=phase_digest, + limits=limits._wire_limits(collective=True), + ) + + return vote + + +@dataclass(frozen=True) +class _LogicalSegment: + """One contiguous mapping between global and rank-local flat storage.""" + + global_offset: int + length: int + local_offset: int + + def __post_init__(self) -> None: + for name, value in ( + ("global_offset", self.global_offset), + ("length", self.length), + ("local_offset", self.local_offset), + ): + if type(value) is not int or value < 0: + raise ValueError("{} must be a nonnegative exact int".format(name)) + + +def _chunk_segment(numel: int, parts: int, coordinate: int) -> _LogicalSegment: + """Return the standard ``torch.chunk`` interval for one mesh coordinate.""" + + if type(numel) is not int or numel < 0: + raise ValueError("numel must be a nonnegative exact int") + if type(parts) is not int or parts < 1: + raise ValueError("parts must be a positive exact int") + if type(coordinate) is not int or coordinate < 0 or coordinate >= parts: + raise ValueError("coordinate is outside the chunk partition") + chunk = (numel + parts - 1) // parts + offset = min(coordinate * chunk, numel) + length = max(0, min(numel, offset + chunk) - offset) + return _LogicalSegment(offset, length, 0) + + +def _canonical_segments(numel: int, parts: int): + return tuple((_chunk_segment(numel, parts, coordinate),) for coordinate in range(parts)) + + +def _full_segment(numel: int): + return (_LogicalSegment(0, numel, 0),) if numel else () + + +def _dimension_shard_segments(shard: ShardIdentity): + placement = shard.placements[0] + dimension = placement.parameter_dimension + if placement.kind is not PlacementKind.DIMENSION_SHARD or type(dimension) is not int: + raise ValueError("dimension-shard segments require one dimension placement") + shape = shard.parameter.global_shape + region = shard.logical_region + if len(shape) == 0: + raise ValueError("a scalar cannot use a dimension-shard placement") + if region.numel == 0: + return () + tail = math.prod(shape[dimension + 1 :]) + local_span = region.lengths[dimension] * tail + global_span = shape[dimension] * tail + prefix_count = math.prod(shape[:dimension]) + dimension_offset = region.offsets[dimension] * tail + return tuple( + _LogicalSegment( + prefix * global_span + dimension_offset, + local_span, + prefix * local_span, + ) + for prefix in range(prefix_count) + if local_span + ) + + +def _dense_segments_for_shard(shard: ShardIdentity, *, unique_replicas: bool): + """Map one parameter-shaped dense field to the live shard's flat storage.""" + + if type(shard) is not ShardIdentity: + raise TypeError("shard must be an exact ShardIdentity") + numel = shard.parameter.numel + placement = shard.placements[0] + if shard.layout is ParameterLayout.REPLICATED: + if unique_replicas and placement.coordinate != 0: + return () + return _full_segment(numel) + if shard.layout is ParameterLayout.FLATTENED_ELEMENT_SHARD: + logical_slice = shard.logical_slice + if logical_slice.length == 0: + return () + return ( + _LogicalSegment( + logical_slice.flat_offset, + logical_slice.length, + 0, + ), + ) + if shard.layout is ParameterLayout.WHOLE_PARAMETER_OWNER: + if shard.local_member != shard.owner: + return () + return _full_segment(numel) + if shard.layout is ParameterLayout.DTENSOR_1D_DEFAULT_WORLD: + if placement.kind is PlacementKind.REPLICATE: + if unique_replicas and placement.coordinate != 0: + return () + return _full_segment(numel) + return _dimension_shard_segments(shard) + raise ValueError("unsupported dense portable shard layout") + + +def _special_segments_for_shard( + shard: ShardIdentity, + *, + numel: int, + unique_replicas: bool, +): + """Map a full-vector field whose ownership follows its parameter shard.""" + + placement = shard.placements[0] + if shard.layout is ParameterLayout.REPLICATED or ( + shard.layout is ParameterLayout.DTENSOR_1D_DEFAULT_WORLD + and placement.kind is PlacementKind.REPLICATE + ): + if unique_replicas and placement.coordinate != 0: + return () + return _full_segment(numel) + if shard.layout is ParameterLayout.WHOLE_PARAMETER_OWNER: + if shard.local_member != shard.owner: + return () + return _full_segment(numel) + raise ValueError("special portable fields require replicas or a whole owner") + + +def _local_numel(segments) -> int: + return max((segment.local_offset + segment.length for segment in segments), default=0) + + +def _validate_segments(segments, *, global_numel: int, name: str): + try: + segments = tuple(segments) + except TypeError as exc: + raise TypeError("{} segments must be iterable".format(name)) from exc + for segment in segments: + if type(segment) is not _LogicalSegment: + raise TypeError("{} entries must be exact logical segments".format(name)) + if segment.global_offset + segment.length > global_numel: + raise ValueError("{} segment exceeds the global field".format(name)) + local_intervals = sorted( + (segment.local_offset, segment.local_offset + segment.length) + for segment in segments + if segment.length + ) + cursor = 0 + for start, stop in local_intervals: + if start != cursor: + raise ValueError("{} local segments must cover storage exactly once".format(name)) + cursor = stop + return segments + + +def _intersection(source: _LogicalSegment, target: _LogicalSegment): + start = max(source.global_offset, target.global_offset) + stop = min( + source.global_offset + source.length, + target.global_offset + target.length, + ) + if stop <= start: + return None + return ( + source.local_offset + start - source.global_offset, + target.local_offset + start - target.global_offset, + stop - start, + ) + + +def _canonical_projection_send_numel( + global_numel: int, + *, + target_segments_by_rank, + parts: int, +): + """Count exact canonical-source send elements, including target duplication.""" + + if len(target_segments_by_rank) != parts: + raise ValueError("projection target segments must match the checkpoint group") + targets = tuple( + _validate_segments( + segments, + global_numel=global_numel, + name="projection target", + ) + for segments in target_segments_by_rank + ) + sources = _canonical_segments(global_numel, parts) + full = _full_segment(global_numel) + if all(target == full for target in targets): + return tuple( + source[0].length * parts if source else 0 + for source in sources + ) + events = [] + for segments in targets: + for segment in segments: + if segment.length: + events.append((segment.global_offset, 1)) + events.append((segment.global_offset + segment.length, -1)) + events.sort() + counts = [] + event_index = 0 + multiplicity = 0 + for source_segments in sources: + count = 0 + for source in source_segments: + start = source.global_offset + stop = start + source.length + while event_index < len(events) and events[event_index][0] <= start: + position = events[event_index][0] + while event_index < len(events) and events[event_index][0] == position: + multiplicity += events[event_index][1] + event_index += 1 + cursor = start + while cursor < stop: + next_position = stop + if event_index < len(events): + next_position = min(next_position, events[event_index][0]) + count += (next_position - cursor) * multiplicity + cursor = next_position + while event_index < len(events) and events[event_index][0] == cursor: + multiplicity += events[event_index][1] + event_index += 1 + counts.append(count) + return tuple(counts) + + +def _dense_projection_send_numel(global_numel: int, *, target_shards, parts: int): + if len(target_shards) != parts: + raise ValueError("dense projection targets must match the checkpoint group") + replicated = all( + shard.layout is ParameterLayout.REPLICATED + or ( + shard.layout is ParameterLayout.DTENSOR_1D_DEFAULT_WORLD + and shard.placements[0].kind is PlacementKind.REPLICATE + ) + for shard in target_shards + ) + unique = all( + shard.layout in { + ParameterLayout.FLATTENED_ELEMENT_SHARD, + ParameterLayout.WHOLE_PARAMETER_OWNER, + } + or ( + shard.layout is ParameterLayout.DTENSOR_1D_DEFAULT_WORLD + and shard.placements[0].kind is PlacementKind.DIMENSION_SHARD + ) + for shard in target_shards + ) + if not replicated and not unique: + raise ValueError("dense projection targets have incompatible layouts") + multiplier = parts if replicated else 1 + return tuple( + _chunk_segment(global_numel, parts, coordinate).length * multiplier + for coordinate in range(parts) + ) + + +def _redistribute_flat_field( + local_tensor, + *, + source_segments_by_rank, + target_segments_by_rank, + global_numel: int, + binding: CheckpointProcessGroupBinding, + phase_vote=None, +) -> torch.Tensor: + """Redistribute one flat logical field without materializing it globally.""" + + def finish_phase(label, error): + if phase_vote is not None: + phase_vote(label, error) + if error is not None: + raise error + + members = None + parts = None + coordinate = None + sources = None + targets = None + local_sources = None + source_numel = None + dtype = None + dtype_codes = None + dtype_code = None + gathered_codes = None + error = None + try: + if type(binding) is not CheckpointProcessGroupBinding: + raise TypeError("binding must be a CheckpointProcessGroupBinding") + if phase_vote is not None and not callable(phase_vote): + raise TypeError("phase_vote must be callable or None") + members = binding.identity.ordered_members + parts = len(members) + if ( + len(source_segments_by_rank) != parts + or len(target_segments_by_rank) != parts + ): + raise ValueError("segment tables must match the checkpoint process group") + coordinate = members.index(binding.local_member) + sources = tuple( + _validate_segments( + segments, + global_numel=global_numel, + name="source", + ) + for segments in source_segments_by_rank + ) + targets = tuple( + _validate_segments( + segments, + global_numel=global_numel, + name="target", + ) + for segments in target_segments_by_rank + ) + local_sources = sources[coordinate] + source_numel = _local_numel(local_sources) + if source_numel: + if type(local_tensor) is not torch.Tensor: + raise TypeError( + "a nonempty source mapping requires an exact torch.Tensor" + ) + if ( + local_tensor.layout is not torch.strided + or local_tensor.is_meta + or local_tensor.is_nested + or local_tensor.is_quantized + or local_tensor.requires_grad + or local_tensor.numel() != source_numel + ): + raise ValueError( + "source tensor is incompatible with its logical segments" + ) + dtype = local_tensor.dtype + else: + if local_tensor is not None and ( + type(local_tensor) is not torch.Tensor + or local_tensor.numel() != 0 + ): + raise ValueError( + "an empty source mapping requires no payload or an empty tensor" + ) + dtype = torch.float32 if local_tensor is None else local_tensor.dtype + dtype_codes = { + torch.uint8: 1, + torch.int8: 2, + torch.int16: 3, + torch.int32: 4, + torch.int64: 5, + torch.float16: 6, + torch.bfloat16: 7, + torch.float32: 8, + torch.float64: 9, + } + if dtype not in dtype_codes: + raise TypeError( + "portable sharded redistribution has an unsupported dtype" + ) + dtype_code = torch.tensor( + [dtype_codes[dtype]], + dtype=torch.int64, + device=binding.collective_device, + ) + gathered_codes = [torch.empty_like(dtype_code) for _ in range(parts)] + except Exception as exc: + error = exc + finish_phase("flat_pre_dtype", error) + + error = None + if parts > 1: + try: + dist.all_gather( + gathered_codes, + dtype_code, + group=binding.process_group, + ) + if any( + int(item.item()) != dtype_codes[dtype] + for item in gathered_codes + ): + raise ValueError( + "portable sharded field dtypes disagree across ranks" + ) + except Exception as exc: + error = exc + finish_phase("flat_dtype", error) + + send = None + input_splits = None + output_splits = None + receive = None + error = None + try: + flat = None + if source_numel: + flat = local_tensor.detach().reshape(-1).to( + binding.collective_device + ) + send_pieces = [] + input_splits = [] + for destination in range(parts): + count = 0 + for source in local_sources: + for target in targets[destination]: + overlap = _intersection(source, target) + if overlap is None: + continue + source_offset, _target_offset, length = overlap + send_pieces.append( + flat[source_offset : source_offset + length] + ) + count += length + input_splits.append(count) + if send_pieces: + send = torch.cat(send_pieces) + else: + send = torch.empty( + 0, + dtype=dtype, + device=binding.collective_device, + ) + output_splits = [] + for source_rank in range(parts): + count = 0 + for source in sources[source_rank]: + for target in targets[coordinate]: + overlap = _intersection(source, target) + if overlap is not None: + count += overlap[2] + output_splits.append(count) + receive = torch.empty( + sum(output_splits), + dtype=dtype, + device=binding.collective_device, + ) + except Exception as exc: + error = exc + finish_phase("flat_send_receive_prep", error) + + error = None + try: + if parts == 1: + receive.copy_(send) + else: + dist.all_to_all_single( + receive, + send, + output_split_sizes=output_splits, + input_split_sizes=input_splits, + group=binding.process_group, + ) + except Exception as exc: + error = exc + finish_phase("flat_exchange", error) + + result = None + error = None + try: + target_numel = _local_numel(targets[coordinate]) + result = torch.empty( + target_numel, + dtype=dtype, + device=binding.collective_device, + ) + written = torch.zeros( + target_numel, + dtype=torch.bool, + device=binding.collective_device, + ) + receive_offset = 0 + for source_rank in range(parts): + for source in sources[source_rank]: + for target in targets[coordinate]: + overlap = _intersection(source, target) + if overlap is None: + continue + _source_offset, target_offset, length = overlap + if bool( + written[target_offset : target_offset + length].any() + ): + raise ValueError( + "portable sharded source segments overlap" + ) + result[target_offset : target_offset + length].copy_( + receive[receive_offset : receive_offset + length] + ) + written[target_offset : target_offset + length] = True + receive_offset += length + if receive_offset != receive.numel() or ( + target_numel and not bool(written.all()) + ): + raise ValueError( + "portable sharded source segments do not cover the target" + ) + except Exception as exc: + error = exc + finish_phase("flat_finish", error) + assert result is not None + return result + + +def _dense_local_numel(shard: ShardIdentity, *, unique_replicas: bool) -> int: + placement = shard.placements[0] + if shard.layout is ParameterLayout.REPLICATED: + return 0 if unique_replicas and placement.coordinate != 0 else shard.parameter.numel + if shard.layout is ParameterLayout.FLATTENED_ELEMENT_SHARD: + return shard.logical_slice.length + if shard.layout is ParameterLayout.WHOLE_PARAMETER_OWNER: + return shard.parameter.numel if shard.local_member == shard.owner else 0 + if shard.layout is ParameterLayout.DTENSOR_1D_DEFAULT_WORLD: + if placement.kind is PlacementKind.REPLICATE: + return 0 if unique_replicas and placement.coordinate != 0 else shard.parameter.numel + if placement.kind is PlacementKind.DIMENSION_SHARD: + return shard.logical_region.numel + raise ValueError("unsupported dense portable shard layout") + + +def _dense_global_indices( + shard: ShardIdentity, + *, + unique_replicas: bool, + device: torch.device, +) -> torch.Tensor: + local_numel = _dense_local_numel( + shard, + unique_replicas=unique_replicas, + ) + if local_numel == 0: + return torch.empty(0, dtype=torch.int64, device=device) + if shard.layout in { + ParameterLayout.REPLICATED, + ParameterLayout.WHOLE_PARAMETER_OWNER, + } or ( + shard.layout is ParameterLayout.DTENSOR_1D_DEFAULT_WORLD + and shard.placements[0].kind is PlacementKind.REPLICATE + ): + return torch.arange(shard.parameter.numel, dtype=torch.int64, device=device) + if shard.layout is ParameterLayout.FLATTENED_ELEMENT_SHARD: + logical_slice = shard.logical_slice + return torch.arange( + logical_slice.flat_offset, + logical_slice.flat_offset + logical_slice.length, + dtype=torch.int64, + device=device, + ) + placement = shard.placements[0] + dimension = placement.parameter_dimension + if ( + shard.layout is not ParameterLayout.DTENSOR_1D_DEFAULT_WORLD + or placement.kind is not PlacementKind.DIMENSION_SHARD + or type(dimension) is not int + ): + raise ValueError("dense index mapping requires a supported shard") + shape = shard.parameter.global_shape + region = shard.logical_region + tail = math.prod(shape[dimension + 1 :]) + local_span = region.lengths[dimension] * tail + global_span = shape[dimension] * tail + local = torch.arange(local_numel, dtype=torch.int64, device=device) + prefix = torch.div(local, local_span, rounding_mode="floor") + within = torch.remainder(local, local_span) + return prefix.mul_(global_span).add_( + region.offsets[dimension] * tail + ).add_(within) + + +def _dense_target_offsets(global_indices: torch.Tensor, shard: ShardIdentity): + placement = shard.placements[0] + if shard.layout is ParameterLayout.REPLICATED or ( + shard.layout is ParameterLayout.DTENSOR_1D_DEFAULT_WORLD + and placement.kind is PlacementKind.REPLICATE + ): + return torch.ones_like(global_indices, dtype=torch.bool), global_indices + if shard.layout is ParameterLayout.WHOLE_PARAMETER_OWNER: + if shard.local_member != shard.owner: + return ( + torch.zeros_like(global_indices, dtype=torch.bool), + torch.empty(0, dtype=torch.int64, device=global_indices.device), + ) + return torch.ones_like(global_indices, dtype=torch.bool), global_indices + if shard.layout is ParameterLayout.FLATTENED_ELEMENT_SHARD: + logical_slice = shard.logical_slice + mask = global_indices.ge(logical_slice.flat_offset) & global_indices.lt( + logical_slice.flat_offset + logical_slice.length + ) + return mask, global_indices[mask] - logical_slice.flat_offset + dimension = placement.parameter_dimension + if ( + shard.layout is not ParameterLayout.DTENSOR_1D_DEFAULT_WORLD + or placement.kind is not PlacementKind.DIMENSION_SHARD + or type(dimension) is not int + ): + raise ValueError("dense target mapping requires a supported shard") + shape = shard.parameter.global_shape + region = shard.logical_region + tail = math.prod(shape[dimension + 1 :]) + global_span = shape[dimension] * tail + within_span = torch.remainder(global_indices, global_span) + dimension_coordinate = torch.div( + within_span, + tail, + rounding_mode="floor", + ) + start = region.offsets[dimension] + length = region.lengths[dimension] + mask = dimension_coordinate.ge(start) & dimension_coordinate.lt(start + length) + selected = global_indices[mask] + selected_within = within_span[mask] + prefix = torch.div(selected, global_span, rounding_mode="floor") + local_span = length * tail + offsets = prefix.mul(local_span).add_( + selected_within - start * tail + ) + return mask, offsets + + +def _redistribute_indexed_dense( + local_tensor, + local_global_indices: torch.Tensor, + *, + target_shards, + global_numel: int, + binding: CheckpointProcessGroupBinding, + phase_vote=None, +) -> torch.Tensor: + def finish_phase(label, error): + if phase_vote is not None: + phase_vote(label, error) + if error is not None: + raise error + + members = None + parts = None + coordinate = None + source_numel = None + dtype = None + flat = None + indices = None + dtype_codes = None + dtype_code = None + gathered_codes = None + error = None + try: + if type(binding) is not CheckpointProcessGroupBinding: + raise TypeError("binding must be a CheckpointProcessGroupBinding") + if phase_vote is not None and not callable(phase_vote): + raise TypeError("phase_vote must be callable or None") + if type(global_numel) is not int or global_numel < 0: + raise ValueError("global_numel must be a nonnegative exact int") + members = binding.identity.ordered_members + parts = len(members) + coordinate = members.index(binding.local_member) + if target_shards is not None and len(target_shards) != parts: + raise ValueError("dense target shards do not match the checkpoint group") + if type(local_global_indices) is not torch.Tensor or local_global_indices.dtype is not torch.int64: + raise TypeError("dense redistribution indices must be int64 tensors") + source_numel = local_global_indices.numel() + if source_numel: + if ( + type(local_tensor) is not torch.Tensor + or local_tensor.layout is not torch.strided + or local_tensor.numel() != source_numel + or local_tensor.requires_grad + ): + raise ValueError("dense redistribution source tensor is incompatible") + dtype = local_tensor.dtype + flat = local_tensor.detach().reshape(-1).to(binding.collective_device) + else: + if local_tensor is not None and ( + type(local_tensor) is not torch.Tensor or local_tensor.numel() != 0 + ): + raise ValueError("empty dense redistribution source has a payload") + dtype = torch.float32 if local_tensor is None else local_tensor.dtype + flat = torch.empty(0, dtype=dtype, device=binding.collective_device) + indices = local_global_indices.to(binding.collective_device) + if bool((indices < 0).any()) or bool((indices >= global_numel).any()): + raise ValueError("dense redistribution indices exceed the logical field") + dtype_codes = { + torch.uint8: 1, + torch.int8: 2, + torch.int16: 3, + torch.int32: 4, + torch.int64: 5, + torch.float16: 6, + torch.bfloat16: 7, + torch.float32: 8, + torch.float64: 9, + } + if dtype not in dtype_codes: + raise TypeError("portable sharded redistribution has an unsupported dtype") + dtype_code = torch.tensor( + [dtype_codes[dtype]], + dtype=torch.int64, + device=binding.collective_device, + ) + gathered_codes = [torch.empty_like(dtype_code) for _ in range(parts)] + except Exception as exc: + error = exc + finish_phase("indexed_pre_dtype", error) + + error = None + if parts > 1: + try: + dist.all_gather(gathered_codes, dtype_code, group=binding.process_group) + if any(int(item.item()) != dtype_codes[dtype] for item in gathered_codes): + raise ValueError("portable sharded field dtypes disagree across ranks") + except Exception as exc: + error = exc + finish_phase("indexed_dtype", error) + + input_splits = None + send_values = None + send_indices = None + send_counts = None + receive_counts = None + error = None + try: + value_pieces = [] + index_pieces = [] + input_splits = [] + for destination in range(parts): + if target_shards is None: + segment = _chunk_segment(global_numel, parts, destination) + mask = indices.ge(segment.global_offset) & indices.lt(segment.global_offset + segment.length) + else: + mask, _offsets = _dense_target_offsets(indices, target_shards[destination]) + value_pieces.append(flat[mask]) + index_pieces.append(indices[mask]) + input_splits.append(int(mask.sum().item())) + send_values = torch.cat(value_pieces) if value_pieces else torch.empty(0, dtype=dtype, device=binding.collective_device) + send_indices = torch.cat(index_pieces) if index_pieces else torch.empty(0, dtype=torch.int64, device=binding.collective_device) + send_counts = torch.tensor(input_splits, dtype=torch.int64, device=binding.collective_device) + receive_counts = torch.empty_like(send_counts) + except Exception as exc: + error = exc + finish_phase("indexed_send_prep", error) + + error = None + try: + if parts == 1: + receive_counts.copy_(send_counts) + else: + dist.all_to_all_single(receive_counts, send_counts, group=binding.process_group) + except Exception as exc: + error = exc + finish_phase("indexed_count_exchange", error) + + output_splits = None + receive_values = None + receive_indices = None + error = None + try: + output_splits = [int(value.item()) for value in receive_counts] + if any(value < 0 for value in output_splits): + raise ValueError("dense redistribution received a negative count") + receive_values = torch.empty(sum(output_splits), dtype=dtype, device=binding.collective_device) + receive_indices = torch.empty(sum(output_splits), dtype=torch.int64, device=binding.collective_device) + except Exception as exc: + error = exc + finish_phase("indexed_receive_prep", error) + + error = None + try: + if parts == 1: + receive_values.copy_(send_values) + else: + dist.all_to_all_single( + receive_values, + send_values, + output_split_sizes=output_splits, + input_split_sizes=input_splits, + group=binding.process_group, + ) + except Exception as exc: + error = exc + finish_phase("indexed_values_exchange", error) + + error = None + try: + if parts == 1: + receive_indices.copy_(send_indices) + else: + dist.all_to_all_single( + receive_indices, + send_indices, + output_split_sizes=output_splits, + input_split_sizes=input_splits, + group=binding.process_group, + ) + except Exception as exc: + error = exc + finish_phase("indexed_indices_exchange", error) + + result = None + error = None + try: + if target_shards is None: + local_target = _chunk_segment(global_numel, parts, coordinate) + target_numel = local_target.length + offsets = receive_indices - local_target.global_offset + mask = offsets.ge(0) & offsets.lt(target_numel) + else: + target_numel = _dense_local_numel(target_shards[coordinate], unique_replicas=False) + mask, offsets = _dense_target_offsets(receive_indices, target_shards[coordinate]) + if not bool(mask.all()): + raise ValueError("dense redistribution received an unowned element") + result = torch.empty(target_numel, dtype=dtype, device=binding.collective_device) + written = torch.zeros(target_numel, dtype=torch.bool, device=binding.collective_device) + if offsets.numel(): + if offsets.unique().numel() != offsets.numel() or bool(written[offsets].any()): + raise ValueError("dense redistribution source elements overlap") + result[offsets] = receive_values + written[offsets] = True + if target_numel and not bool(written.all()): + raise ValueError("dense redistribution does not cover the target") + except Exception as exc: + error = exc + finish_phase("indexed_finish", error) + assert result is not None + return result + + +def _redistribute_dense_to_canonical( + local_tensor, + *, + source_shard: ShardIdentity, + global_numel: int, + binding: CheckpointProcessGroupBinding, + phase_vote=None, +): + indices = None + error = None + try: + indices = _dense_global_indices( + source_shard, + unique_replicas=True, + device=binding.collective_device, + ) + except Exception as exc: + error = exc + if phase_vote is not None: + phase_vote("dense_source_index_prep", error) + if error is not None: + raise error + assert indices is not None + return _redistribute_indexed_dense( + local_tensor, + indices, + target_shards=None, + global_numel=global_numel, + binding=binding, + phase_vote=phase_vote, + ) + + +def _redistribute_canonical_to_dense( + local_tensor, + *, + target_shards, + global_numel: int, + binding: CheckpointProcessGroupBinding, + phase_vote=None, +): + indices = None + error = None + try: + coordinate = binding.identity.ordered_members.index(binding.local_member) + source = _chunk_segment( + global_numel, + len(binding.identity.ordered_members), + coordinate, + ) + indices = torch.arange( + source.global_offset, + source.global_offset + source.length, + dtype=torch.int64, + device=binding.collective_device, + ) + except Exception as exc: + error = exc + if phase_vote is not None: + phase_vote("canonical_source_index_prep", error) + if error is not None: + raise error + assert indices is not None + return _redistribute_indexed_dense( + local_tensor, + indices, + target_shards=target_shards, + global_numel=global_numel, + binding=binding, + phase_vote=phase_vote, + ) + + +def _device_mesh(binding: CheckpointProcessGroupBinding): + """Build a 1-D DeviceMesh around the optimizer-owned checkpoint group.""" + + from torch.distributed.device_mesh import DeviceMesh + + if not dist.is_available() or not dist.is_initialized(): + raise RuntimeError("sharded portable DCP requires initialized torch.distributed") + process_group = binding.process_group + if process_group is None: + process_group = dist.group.WORLD + return DeviceMesh.from_group( + process_group, + binding.collective_device.type, + mesh_dim_names=("gefen_portable_dcp",), + ) + + +def _as_canonical_dtensor( + local_tensor: torch.Tensor, + *, + global_numel: int, + binding: CheckpointProcessGroupBinding, +): + """Wrap one standard canonical chunk for DCP without another collective.""" + + from torch.distributed.tensor import DTensor, Shard + + parts = len(binding.identity.ordered_members) + coordinate = binding.identity.ordered_members.index(binding.local_member) + expected = _chunk_segment(global_numel, parts, coordinate) + if type(local_tensor) is not torch.Tensor or local_tensor.numel() != expected.length: + raise ValueError("canonical DCP local tensor has invalid chunk geometry") + if parts == 1: + return local_tensor + return DTensor.from_local( + local_tensor.contiguous(), + _device_mesh(binding), + (Shard(0),), + run_check=False, + shape=torch.Size((global_numel,)), + stride=(1,), + ) + + +def _allocate_canonical_dtensor( + *, + global_numel: int, + dtype: torch.dtype, + binding: CheckpointProcessGroupBinding, +): + parts = len(binding.identity.ordered_members) + coordinate = binding.identity.ordered_members.index(binding.local_member) + local = torch.empty( + _chunk_segment(global_numel, parts, coordinate).length, + dtype=dtype, + device=binding.collective_device, + ) + return _as_canonical_dtensor( + local, + global_numel=global_numel, + binding=binding, + ) + + +def _canonical_local_tensor(value) -> torch.Tensor: + from torch.distributed.tensor import DTensor + + if type(value) is DTensor: + local = value.to_local() + if hasattr(local, "wait"): + local = local.wait() + if type(local) is not torch.Tensor: + raise TypeError("canonical DTensor produced a non-tensor local value") + return local + if type(value) is not torch.Tensor: + raise TypeError("canonical DCP value must be a tensor or exact DTensor") + return value + + +def _tensor_sha256(value: torch.Tensor) -> bytes: + if type(value) is not torch.Tensor: + raise TypeError("integrity hashing requires an exact torch.Tensor") + hasher = hashlib.sha256() + flat = value.detach().contiguous().reshape(-1).view(torch.uint8) + chunk_bytes = 1 << 20 + for start in range(0, flat.numel(), chunk_bytes): + chunk = flat[start : start + chunk_bytes].cpu() + hasher.update(memoryview(chunk.numpy())) + return hasher.digest() + + +def _field_key(index: int) -> str: + if type(index) is not int or index < 0 or index >= 10**_DCP_SHARDED_FIELD_DIGITS: + raise ValueError("portable DCP field index is out of range") + return "{}{:0{}d}{}".format( + _DCP_SHARDED_FIELD_PREFIX, + index, + _DCP_SHARDED_FIELD_DIGITS, + _DCP_SHARDED_FIELD_SUFFIX, + ) + + +def _tensor_descriptor(value, *, name: str): + if type(value) is not torch.Tensor: + raise TypeError("{} must be an exact torch.Tensor".format(name)) + if ( + 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.requires_grad + or value.dtype is not torch.float32 + or not value.is_contiguous() + ): + raise TypeError("{} must be a contiguous CPU fp32 tensor".format(name)) + return { + "dtype": "float32", + "shape": list(value.shape), + "sha256": _tensor_sha256(value).hex(), + } + + +def _structure_local_fragment(fragment): + """Remove tensor payloads while retaining bounded semantic descriptors.""" + + if type(fragment) is not dict: + raise TypeError("portable local fragment must be a dictionary") + common = fragment["common"] + codebook = common["gefen_codebook"] + structured_common = { + "gefen_global_step": common["gefen_global_step"], + "gefen_codebook": ( + None + if codebook is None + else _tensor_descriptor(codebook, name="gefen_codebook") + ), + "gefen_deterministic": common["gefen_deterministic"], + } + slots = [] + for slot in fragment["logical_slots"]: + state = {} + for key, value in slot["state"].items(): + state[key] = ( + _tensor_descriptor(value, name="{} {}".format(slot["compatibility_name"], key)) + if type(value) is torch.Tensor + else value + ) + slots.append({**slot, "state": state}) + return { + "format": fragment["format"], + "format_version": fragment["format_version"], + "coverage": fragment["coverage"], + "implementation": fragment["implementation"], + "member": fragment["member"], + "policy": fragment["policy"], + "common": structured_common, + "manifest": fragment["manifest"], + "catalog": fragment["catalog"], + "logical_slots": slots, + } + + +def _consensus(values, *, name: str): + from gefen.portable_state import _values_equal + + 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 sharded portable fragments".format(name)) + return reference + + +def _descriptor_shape(descriptor, *, name: str): + if type(descriptor) is not dict or set(descriptor) != {"dtype", "shape", "sha256"}: + raise ValueError("{} has an invalid tensor descriptor".format(name)) + if descriptor["dtype"] != "float32": + raise TypeError("{} must use fp32".format(name)) + shape = descriptor["shape"] + digest = descriptor["sha256"] + if ( + type(shape) is not list + or any(type(dimension) is not int or dimension < 0 for dimension in shape) + or type(digest) is not str + or len(digest) != 64 + ): + raise ValueError("{} has an invalid tensor descriptor".format(name)) + try: + bytes.fromhex(digest) + except ValueError as exc: + raise ValueError("{} has an invalid tensor digest".format(name)) from exc + return tuple(shape) + + +def _add_field(fields, *, fqn, name, shape, kind, nonnegative): + index = len(fields) + spec = { + "index": index, + "key": _field_key(index), + "fqn": fqn, + "name": name, + "shape": list(shape), + "numel": math.prod(shape), + "dtype": "float32", + "kind": kind, + "nonnegative": nonnegative, + } + fields.append(spec) + return index + + +def _validate_field_descriptors(slots, *, key: str, shape_by_shard, replicated: bool): + descriptors = [] + for slot in slots: + from gefen.portable_identity import _parse_shard_identity + + shard = _parse_shard_identity(slot["shard"]) + descriptor = slot["state"].get(key) + expected_shape = shape_by_shard(shard) + if expected_shape is None: + if descriptor is not None: + raise ValueError("a non-payload shard carries {}".format(key)) + continue + if descriptor is None or _descriptor_shape(descriptor, name=key) != tuple(expected_shape): + raise ValueError("{} descriptor does not match its shard".format(key)) + descriptors.append(descriptor) + if not descriptors: + raise ValueError("{} requires at least one payload descriptor".format(key)) + if replicated: + _consensus(descriptors, name="{} replicas".format(key)) + + +def _assemble_sharded_schema(structures, *, binding: CheckpointProcessGroupBinding): + """Assemble bounded payload metadata without assembling any dense field.""" + + from gefen.portable_identity import ( + _parse_parameter_identity, + _parse_shard_identity, + _parse_sharding_manifest, + _serialize_process_group_identity, + ) + from gefen.portable_state import _local_dense_shape + + members = binding.identity.ordered_members + structures = tuple(structures) + if len(structures) != len(members): + raise ValueError("sharded portable metadata requires every checkpoint member") + if tuple(fragment["member"] for fragment in structures) != members: + raise ValueError("sharded portable metadata is not in member order") + implementation = _consensus( + [fragment["implementation"] for fragment in structures], + name="implementation", + ) + policy = _consensus([fragment["policy"] for fragment in structures], name="policy") + manifest_record = _consensus( + [fragment["manifest"] for fragment in structures], + name="manifest", + ) + catalog = _consensus([fragment["catalog"] for fragment in structures], name="catalog") + manifest = _parse_sharding_manifest(manifest_record) + if any(shard.process_group != binding.identity for shard in manifest.shards): + raise ValueError("portable manifest process groups must match the checkpoint binding") + + common_values = [fragment["common"] for fragment in structures] + global_step = _consensus( + [common["gefen_global_step"] for common in common_values], + name="gefen_global_step", + ) + deterministic = _consensus( + [common["gefen_deterministic"] for common in common_values], + name="gefen_deterministic", + ) + codebooks = [common["gefen_codebook"] for common in common_values] + fields = [] + codebook_field = None + if any(codebook is not None for codebook in codebooks): + if any(codebook is None for codebook in codebooks): + raise ValueError("portable codebook presence differs across members") + _consensus(codebooks, name="gefen_codebook") + if _descriptor_shape(codebooks[0], name="gefen_codebook") != (256,): + raise ValueError("portable codebook has invalid geometry") + codebook_field = _add_field( + fields, + fqn=None, + name="gefen_codebook", + shape=(256,), + kind="common", + nonnegative=False, + ) + + all_slots = [slot for fragment in structures for slot in fragment["logical_slots"]] + slots_by_shard = {_parse_shard_identity(slot["shard"]): slot for slot in all_slots} + if len(slots_by_shard) != len(all_slots) or set(slots_by_shard) != set(manifest.shards): + raise ValueError("sharded portable slots must exactly cover the manifest") + parameters = {} + for fqn in sorted(catalog): + parameter = _parse_parameter_identity(catalog[fqn]["identity"]) + shards = manifest.for_parameter(fqn) + slots = [slots_by_shard[shard] for shard in shards] + _consensus( + [(slot["group_index"], slot["original_slot_index"]) for slot in slots], + name="logical slot position", + ) + compatibility_name = _consensus( + [slot["compatibility_name"] for slot in slots], + name="compatibility_name", + ) + required = [ + slot + for slot, shard in zip(slots, shards) + if _local_dense_shape(shard) is not None and parameter.numel > 0 + ] + variants = {slot["state_variant"] for slot in required} + if not required or 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 state_variant") + 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 + } + ) + state = {} + source_second = None + if variant.startswith("initialized_"): + step = _consensus( + [slot["state"]["step"] for slot in required], + name="parameter step", + ) + if type(step) is not int or step < 1 or step > global_step: + raise ValueError("portable parameter step is invalid") + replicated_dense = shards[0].layout is ParameterLayout.REPLICATED or ( + shards[0].layout is ParameterLayout.DTENSOR_1D_DEFAULT_WORLD + and shards[0].placements[0].kind is PlacementKind.REPLICATE + ) + _validate_field_descriptors( + slots, + key="momentum", + shape_by_shard=_local_dense_shape, + replicated=replicated_dense, + ) + state["step"] = step + state["momentum_field"] = _add_field( + fields, + fqn=fqn, + name="momentum", + shape=parameter.global_shape, + kind="dense", + nonnegative=False, + ) + 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", + ) + if type(second_step) is not int or second_step < 1 or second_step > step: + raise ValueError("portable second_moment_step is invalid") + _validate_field_descriptors( + slots, + key="second_moment", + shape_by_shard=_local_dense_shape, + replicated=replicated_dense, + ) + state.update( + { + "second_moment_field": _add_field( + fields, + fqn=fqn, + name="second_moment", + shape=parameter.global_shape, + kind="dense", + nonnegative=True, + ), + "second_moment_step": second_step, + } + ) + elif implementation == "gefen.Gefen" and variant == "initialized_factored": + source_second = "factored" + if any(shard.layout is not ParameterLayout.REPLICATED for shard in shards): + raise ValueError("factored portable state requires replicated shards") + factored_step = _consensus( + [slot["state"]["factored_step"] for slot in required], + name="factored_step", + ) + if type(factored_step) is not int or factored_step < 1 or factored_step > step: + raise ValueError("portable factored_step is invalid") + rows, columns = parameter.global_shape + for key, shape in (("v_row", (rows,)), ("v_col", (columns,))): + _validate_field_descriptors( + slots, + key=key, + shape_by_shard=lambda _shard, expected=shape: expected, + replicated=True, + ) + state["{}_field".format(key)] = _add_field( + fields, + fqn=fqn, + name=key, + shape=shape, + kind="special", + nonnegative=True, + ) + state["v_row_denominator_field"] = _add_field( + fields, + fqn=fqn, + name="v_row_denominator", + shape=(1,), + kind="special", + nonnegative=True, + ) + state["factored_step"] = factored_step + elif implementation == "gefen.GefenMuon" and variant == "initialized_dense_normuon": + shape = (parameter.global_shape[0], 1) + normuon_step = _consensus( + [slot["state"]["normuon_step"] for slot in required], + name="normuon_step", + ) + if type(normuon_step) is not int or normuon_step < 1 or normuon_step > step: + raise ValueError("portable normuon_step is invalid") + _validate_field_descriptors( + slots, + key="normuon_v", + shape_by_shard=lambda shard: ( + shape + if shard.layout is ParameterLayout.REPLICATED + or ( + shard.layout is ParameterLayout.WHOLE_PARAMETER_OWNER + and shard.local_member == shard.owner + ) + else None + ), + replicated=shards[0].layout is ParameterLayout.REPLICATED, + ) + state.update( + { + "normuon_v_field": _add_field( + fields, + fqn=fqn, + name="normuon_v", + shape=shape, + kind="special", + nonnegative=True, + ), + "normuon_step": normuon_step, + } + ) + parameters[fqn] = { + "identity": catalog[fqn]["identity"], + "algorithm_options": catalog[fqn]["algorithm_options"], + "group_index": slots[0]["group_index"], + "original_slot_index": slots[0]["original_slot_index"], + "compatibility_name": compatibility_name, + "state_variant": variant, + "state": state, + "projection_hints": ( + { + "source_periods": periods, + "source_second_moment": source_second, + "target_period": 1, + } + if implementation == "gefen.Gefen" + else {"source_periods": periods, "target_period": 1} + ), + } + return { + "format": _DCP_SHARDED_FORMAT, + "format_version": _DCP_SHARDED_FORMAT_VERSION, + "implementation": implementation, + "source_process_group": _serialize_process_group_identity(binding.identity), + "source_world_size": len(members), + "policy": policy, + "common": { + "gefen_global_step": global_step, + "gefen_codebook_field": codebook_field, + "gefen_deterministic": deterministic, + }, + "source_manifest": manifest_record, + "parameters": parameters, + "fields": fields, + } + + +def _local_fragment_field(fragment, spec): + if spec["kind"] == "common": + return fragment["common"]["gefen_codebook"] + slot = next( + item + for item in fragment["logical_slots"] + if item["shard"]["parameter"]["fqn"] == spec["fqn"] + ) + if spec["name"] == "v_row_denominator": + row = slot["state"].get("v_row") + if row is None: + return None + return row.mean().clamp_(min=torch.finfo(torch.float32).tiny).reshape(1) + return slot["state"].get(spec["name"]) + + +def _source_shards_in_member_order(manifest, fqn, members): + shards_by_member = { + shard.local_member: shard + for shard in manifest.for_parameter(fqn) + } + if set(shards_by_member) != set(members): + raise ValueError("portable field shards are not in checkpoint member order") + return tuple(shards_by_member[member] for member in members) + + +def _source_segments_for_field(spec, manifest, *, members): + if spec["kind"] == "common": + return tuple( + _full_segment(spec["numel"]) if coordinate == 0 else () + for coordinate in range(len(members)) + ) + shards = _source_shards_in_member_order(manifest, spec["fqn"], members) + if spec["kind"] == "dense": + return tuple( + _dense_segments_for_shard(shard, unique_replicas=True) + for shard in shards + ) + if spec["kind"] == "special": + return tuple( + _special_segments_for_shard( + shard, + numel=spec["numel"], + unique_replicas=True, + ) + for shard in shards + ) + raise ValueError("portable DCP field has an unknown geometry kind") + + +def _all_gather_digest( + digest: bytes, + binding: CheckpointProcessGroupBinding, + *, + phase_vote=None, +): + def finish_phase(label, error): + if phase_vote is not None: + phase_vote(label, error) + if error is not None: + raise error + + local = None + parts = None + gathered = None + error = None + try: + if type(digest) is not bytes or len(digest) != hashlib.sha256().digest_size: + raise ValueError("portable field digest must be one SHA-256 value") + local = torch.tensor( + list(digest), + dtype=torch.uint8, + device=binding.collective_device, + ) + parts = len(binding.identity.ordered_members) + gathered = [torch.empty_like(local) for _ in range(parts)] + except Exception as exc: + error = exc + finish_phase("digest_prep", error) + + error = None + try: + if parts == 1: + gathered[0].copy_(local) + else: + dist.all_gather(gathered, local, group=binding.process_group) + except Exception as exc: + error = exc + finish_phase("digest_exchange", error) + + result = None + error = None + try: + result = tuple(bytes(item.cpu().tolist()) for item in gathered) + except Exception as exc: + error = exc + finish_phase("digest_finish", error) + assert result is not None + return result + + +def _metadata_bytes(metadata, *, limits): + from gefen.portable_schema import portable_state_digest + from gefen.portable_wire import _prepare_canonical_wire_value + + semantic = {**metadata} + semantic_digest = portable_state_digest(semantic) + payload_hasher = hashlib.sha256() + for field in semantic["fields"]: + payload_hasher.update(field["key"].encode("ascii")) + for digest in field["source_chunk_sha256"]: + payload_hasher.update(bytes.fromhex(digest)) + completed = { + **semantic, + "completion": { + "metadata_digest": semantic_digest, + "payload_digest": payload_hasher.hexdigest(), + }, + } + plan = _prepare_canonical_wire_value(completed, limits._wire_limits(collective=True)) + if plan.payload_tensors: + raise RuntimeError("sharded portable DCP metadata unexpectedly contains tensor payloads") + return plan.metadata + + +def _gather_structures( + binding: CheckpointProcessGroupBinding, + local_structure, + *, + transaction_id: str, + limits, +): + from gefen.portable_collective import _collective_visit_canonical_fragments + + structures = [] + _collective_visit_canonical_fragments( + binding, + local_structure, + operation="portable_dcp_sharded_structures", + transaction_id=transaction_id, + context_digest=bytes(32), + limits=limits._wire_limits(collective=True), + consume=lambda _member, value: structures.append(value), + ) + return tuple(structures) + + +def _prepare_sharded_save_state( + optimizer, + *, + binding: CheckpointProcessGroupBinding, + transaction_id: str, + limits, + namespace: str, +): + """Prepare a DCP state dict whose large fields are sharded DTensors.""" + + from gefen.portable_identity import _parse_sharding_manifest + from gefen import portable_runtime as runtime + from gefen.portable_collective import _collective_unanimous_status + + implementation = runtime._optimizer_implementation(optimizer) + live_token = runtime._portable_live_token(optimizer) + local_fragment = None + local_structure = None + error = None + try: + prepared = runtime._prepare_local_structure( + optimizer, + implementation, + binding, + limits, + include_payload=True, + ) + runtime._validate_prepared_local_state( + optimizer, + implementation, + prepared, + ) + local_fragment = runtime._materialize_local_fragment( + optimizer, + implementation, + binding, + limits, + prepared, + ) + local_structure = _structure_local_fragment(local_fragment) + except Exception as exc: + error = exc + _collective_unanimous_status( + binding, + error, + operation="portable_dcp_sharded_save_local", + transaction_id=transaction_id, + context_digest=bytes(32), + limits=limits._wire_limits(collective=True), + ) + assert local_fragment is not None and local_structure is not None + structures = _gather_structures( + binding, + local_structure, + transaction_id=transaction_id, + limits=limits, + ) + metadata = None + manifest = None + source_segments_by_index = {} + source_shards_by_index = {} + local_values = {} + schema_digest = bytes(32) + error = None + try: + metadata = _assemble_sharded_schema(structures, binding=binding) + expected_chunks = 1 + len(metadata["fields"]) * len( + binding.identity.ordered_members + ) + if ( + expected_chunks > limits.max_container_items + or expected_chunks > limits.max_tree_nodes + ): + raise ValueError("sharded portable DCP chunks exceed aggregate limits") + _validate_field_byte_limits( + metadata["fields"], + binding=binding, + limits=limits, + ) + manifest = _parse_sharding_manifest(metadata["source_manifest"]) + coordinate = binding.identity.ordered_members.index(binding.local_member) + for spec in metadata["fields"]: + local_value = _local_fragment_field(local_fragment, spec) + if spec["kind"] == "dense": + source_shards = _source_shards_in_member_order( + manifest, + spec["fqn"], + binding.identity.ordered_members, + ) + expected_numel = _dense_local_numel( + source_shards[coordinate], + unique_replicas=True, + ) + canonical_numel = _chunk_segment( + spec["numel"], + len(binding.identity.ordered_members), + coordinate, + ).length + if max( + expected_numel * 8, + canonical_numel * 8, + len(binding.identity.ordered_members) * 8, + ) > limits.max_fragment_tensor_bytes: + raise ValueError( + "sharded portable save routing scratch exceeds max_fragment_tensor_bytes" + ) + source_shards_by_index[spec["index"]] = source_shards + else: + source_segments = _source_segments_for_field( + spec, + manifest, + members=binding.identity.ordered_members, + ) + expected_numel = _local_numel(source_segments[coordinate]) + source_segments_by_index[spec["index"]] = source_segments + if expected_numel == 0: + local_value = None + elif ( + type(local_value) is not torch.Tensor + or local_value.numel() != expected_numel + or local_value.dtype is not torch.float32 + ): + raise ValueError("sharded portable source field geometry is invalid") + local_values[spec["index"]] = local_value + schema_digest = runtime._context_digest(metadata) + except Exception as exc: + error = exc + _collective_unanimous_status( + binding, + error, + operation="portable_dcp_sharded_save_schema", + transaction_id=transaction_id, + context_digest=schema_digest, + limits=limits._wire_limits(collective=True), + ) + assert metadata is not None and manifest is not None + phase_vote = _make_phase_voter( + binding, + transaction_id=transaction_id, + context_digest=schema_digest, + limits=limits, + operation_prefix="portable_dcp_sharded_save_phase", + ) + state = {} + for spec in metadata["fields"]: + def field_vote(label, error, *, field_index=spec["index"]): + phase_vote("field.{}.{}".format(field_index, label), error) + + local_value = local_values[spec["index"]] + canonical = None + local_digest = None + error = None + try: + if spec["kind"] == "dense": + coordinate = binding.identity.ordered_members.index( + binding.local_member + ) + canonical = _redistribute_dense_to_canonical( + local_value, + source_shard=source_shards_by_index[spec["index"]][ + coordinate + ], + global_numel=spec["numel"], + binding=binding, + phase_vote=field_vote, + ) + else: + source_segments = source_segments_by_index[spec["index"]] + canonical = _redistribute_flat_field( + local_value, + source_segments_by_rank=source_segments, + target_segments_by_rank=_canonical_segments( + spec["numel"], + len(binding.identity.ordered_members), + ), + global_numel=spec["numel"], + binding=binding, + phase_vote=field_vote, + ) + local_digest = _tensor_sha256(canonical) + except Exception as exc: + error = exc + field_vote("projection_hash", error) + assert canonical is not None and local_digest is not None + chunk_digests = _all_gather_digest( + local_digest, + binding, + phase_vote=field_vote, + ) + error = None + try: + spec["source_chunk_sha256"] = [digest.hex() for digest in chunk_digests] + state[spec["key"]] = _as_canonical_dtensor( + canonical, + global_numel=spec["numel"], + binding=binding, + ) + except Exception as exc: + error = exc + field_vote("wrap_finish", error) + + error = None + try: + metadata = _metadata_bytes(metadata, limits=limits) + state[_DCP_SHARDED_METADATA_KEY] = torch.frombuffer( + bytearray(metadata), + dtype=torch.uint8, + ) + runtime._validate_live_readiness(optimizer, implementation, binding) + if live_token != runtime._portable_live_token(optimizer): + raise RuntimeError("live optimizer state changed during sharded DCP save preparation") + except Exception as exc: + error = exc + phase_vote("save.finalize", error) + return {namespace: state} + + +def _strict_record(value, keys, *, name: str): + if type(value) is not dict or set(value) != set(keys): + raise ValueError("{} has an invalid schema".format(name)) + return value + + +def _strict_counter(value, *, name: str, minimum=0, maximum=(1 << 53) - 1): + if type(value) is not int or value < minimum or value > maximum: + raise ValueError("{} is not a bounded exact counter".format(name)) + return value + + +def _validate_loaded_field_spec( + spec, + *, + index: int, + source_world_size: int, + limits, +): + _strict_record( + spec, + { + "index", + "key", + "fqn", + "name", + "shape", + "numel", + "dtype", + "kind", + "nonnegative", + "source_chunk_sha256", + }, + name="sharded portable field", + ) + if spec["index"] != index or spec["key"] != _field_key(index): + raise ValueError("sharded portable field ordering is invalid") + if spec["fqn"] is not None and ( + type(spec["fqn"]) is not str or not spec["fqn"] + ): + raise ValueError("sharded portable field FQN is invalid") + if type(spec["name"]) is not str or not spec["name"]: + raise ValueError("sharded portable field name is invalid") + shape = spec["shape"] + if ( + type(shape) is not list + or any(type(dimension) is not int or dimension < 0 for dimension in shape) + or len(shape) > limits.max_tensor_rank + or spec["numel"] != math.prod(shape) + or type(spec["numel"]) is not int + or spec["numel"] < 1 + or spec["dtype"] != "float32" + or spec["kind"] not in {"common", "dense", "special"} + or type(spec["nonnegative"]) is not bool + ): + raise ValueError("sharded portable field geometry is invalid") + digests = spec["source_chunk_sha256"] + if type(digests) is not list or len(digests) != source_world_size: + raise ValueError("sharded portable field chunk digests are incomplete") + for digest in digests: + if type(digest) is not str or len(digest) != 64: + raise ValueError("sharded portable field chunk digest is invalid") + try: + bytes.fromhex(digest) + except ValueError as exc: + raise ValueError("sharded portable field chunk digest is invalid") from exc + return spec + + +def _validate_loaded_parameter_record( + fqn, + record, + *, + implementation, + policy, + manifest, + fields, + global_step, + referenced_fields, +): + from gefen.portable_identity import ( + _normalize_parameter_identity, + _parse_parameter_identity, + ) + from gefen.portable_state import _normalize_options + + _strict_record( + record, + { + "identity", + "algorithm_options", + "group_index", + "original_slot_index", + "compatibility_name", + "state_variant", + "state", + "projection_hints", + }, + name="sharded portable parameter", + ) + 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("sharded portable parameter identity conflicts with its manifest") + options = _normalize_options(record["algorithm_options"], implementation) + for key in ("group_index", "original_slot_index"): + _strict_counter(record[key], name=key) + if ( + type(record["compatibility_name"]) is not str + or not record["compatibility_name"] + or record["compatibility_name"] != record["compatibility_name"].lower() + ): + raise ValueError("sharded portable compatibility_name is invalid") + variant = record["state_variant"] + state = record["state"] + hints = record["projection_hints"] + if implementation == "gefen.Gefen": + _strict_record( + hints, + {"source_periods", "source_second_moment", "target_period"}, + name="plain sharded projection hints", + ) + else: + _strict_record( + hints, + {"source_periods", "target_period"}, + name="Muon sharded projection hints", + ) + periods = hints["source_periods"] + if type(periods) is not list or periods not in ([], [1]) or hints["target_period"] != 1: + raise ValueError("sharded portable source periods are invalid") + + def field_reference(key, *, expected_name, expected_shape, expected_kind, nonnegative): + index = state[key] + if type(index) is not int or index < 0 or index >= len(fields): + raise ValueError("{} has an invalid field reference".format(key)) + if index in referenced_fields: + raise ValueError("sharded portable fields must be referenced exactly once") + field = fields[index] + if ( + field["fqn"] != fqn + or field["name"] != expected_name + or tuple(field["shape"]) != tuple(expected_shape) + or field["kind"] != expected_kind + or field["nonnegative"] is not nonnegative + ): + raise ValueError("{} refers to an incompatible field".format(key)) + referenced_fields.add(index) + + if variant == "pristine": + _strict_record(state, set(), name="pristine sharded state") + if periods: + raise ValueError("pristine sharded state has source periods") + elif variant == "period_selected": + _strict_record(state, set(), name="period-selected sharded state") + if periods != [1]: + raise ValueError("period-selected sharded state requires period one") + elif implementation == "gefen.Gefen" and variant == "initialized_dense": + _strict_record( + state, + {"step", "momentum_field", "second_moment_field", "second_moment_step"}, + name="dense sharded state", + ) + step = _strict_counter(state["step"], name="step", minimum=1) + second_step = _strict_counter( + state["second_moment_step"], + name="second_moment_step", + minimum=1, + ) + if step > global_step or second_step > step or periods != [1]: + raise ValueError("dense sharded state counters are invalid") + if hints["source_second_moment"] != "block" or options["second_moment_policy"] != "block": + raise ValueError("dense sharded state conflicts with its second-moment policy") + field_reference( + "momentum_field", + expected_name="momentum", + expected_shape=identity.global_shape, + expected_kind="dense", + nonnegative=False, + ) + field_reference( + "second_moment_field", + expected_name="second_moment", + expected_shape=identity.global_shape, + expected_kind="dense", + nonnegative=True, + ) + elif implementation == "gefen.Gefen" and variant == "initialized_factored": + _strict_record( + state, + { + "step", + "momentum_field", + "v_row_field", + "v_col_field", + "v_row_denominator_field", + "factored_step", + }, + name="factored sharded state", + ) + step = _strict_counter(state["step"], name="step", minimum=1) + factored_step = _strict_counter( + state["factored_step"], + name="factored_step", + minimum=1, + ) + if step > global_step or factored_step > step or periods != [1]: + raise ValueError("factored sharded state counters are invalid") + if ( + hints["source_second_moment"] != "factored" + or options["second_moment_policy"] != "factored" + or len(identity.global_shape) != 2 + ): + raise ValueError("factored sharded state conflicts with its policy") + field_reference( + "momentum_field", + expected_name="momentum", + expected_shape=identity.global_shape, + expected_kind="dense", + nonnegative=False, + ) + field_reference( + "v_row_field", + expected_name="v_row", + expected_shape=(identity.global_shape[0],), + expected_kind="special", + nonnegative=True, + ) + field_reference( + "v_col_field", + expected_name="v_col", + expected_shape=(identity.global_shape[1],), + expected_kind="special", + nonnegative=True, + ) + field_reference( + "v_row_denominator_field", + expected_name="v_row_denominator", + expected_shape=(1,), + expected_kind="special", + nonnegative=True, + ) + elif implementation == "gefen.GefenMuon" and variant in { + "initialized_dense", + "initialized_dense_normuon", + }: + expected = {"step", "momentum_field"} + if variant == "initialized_dense_normuon": + expected.update({"normuon_v_field", "normuon_step"}) + _strict_record(state, expected, name="Muon sharded state") + step = _strict_counter(state["step"], name="step", minimum=1) + if step > global_step or periods != [1]: + raise ValueError("Muon sharded state counters are invalid") + if options["normuon"] is not (variant == "initialized_dense_normuon"): + raise ValueError("Muon sharded state conflicts with its NorMuon policy") + field_reference( + "momentum_field", + expected_name="momentum", + expected_shape=identity.global_shape, + expected_kind="dense", + nonnegative=False, + ) + if variant == "initialized_dense_normuon": + normuon_step = _strict_counter( + state["normuon_step"], + name="normuon_step", + minimum=1, + ) + if normuon_step > step: + raise ValueError("normuon_step exceeds parameter step") + field_reference( + "normuon_v_field", + expected_name="normuon_v", + expected_shape=(identity.global_shape[0], 1), + expected_kind="special", + nonnegative=True, + ) + else: + raise ValueError("sharded portable parameter state_variant is unsupported") + return { + **record, + "identity": identity_record, + "algorithm_options": options, + } + + +def _normalize_sharded_metadata(value, *, limits): + from gefen.portable_identity import ( + _normalize_process_group_identity, + _normalize_sharding_manifest, + _parse_process_group_identity, + _parse_sharding_manifest, + ) + from gefen.portable_schema import portable_state_digest + from gefen.portable_state import ( + _IMPLEMENTATIONS, + _normalize_catalog, + _normalize_policy, + ) + + _strict_record( + value, + { + "format", + "format_version", + "implementation", + "source_process_group", + "source_world_size", + "policy", + "common", + "source_manifest", + "parameters", + "fields", + "completion", + }, + name="sharded portable metadata", + ) + if value["format"] != _DCP_SHARDED_FORMAT or value["format_version"] != _DCP_SHARDED_FORMAT_VERSION: + raise ValueError("unsupported sharded portable DCP format") + implementation = value["implementation"] + if implementation not in _IMPLEMENTATIONS: + raise ValueError("unsupported sharded portable implementation") + process_group_record = _normalize_process_group_identity(value["source_process_group"]) + process_group = _parse_process_group_identity(process_group_record) + source_world_size = value["source_world_size"] + if ( + type(source_world_size) is not int + or source_world_size < 1 + or source_world_size > limits.max_members + or source_world_size != len(process_group.ordered_members) + ): + raise ValueError("sharded portable source world size is invalid") + policy = _normalize_policy(value["policy"], implementation) + manifest_record = _normalize_sharding_manifest(value["source_manifest"]) + manifest = _parse_sharding_manifest(manifest_record) + if any(shard.process_group != process_group for shard in manifest.shards): + raise ValueError("sharded portable source manifest has a different process group") + common = _strict_record( + value["common"], + {"gefen_global_step", "gefen_codebook_field", "gefen_deterministic"}, + name="sharded portable common state", + ) + global_step = _strict_counter(common["gefen_global_step"], name="gefen_global_step") + if type(common["gefen_deterministic"]) is not bool: + raise ValueError("sharded portable deterministic policy must be bool") + fields_value = value["fields"] + if ( + type(fields_value) is not list + or len(fields_value) > limits.max_tensors + or len(fields_value) > limits.max_container_items + ): + raise ValueError("sharded portable field count exceeds limits") + fields = [ + _validate_loaded_field_spec( + spec, + index=index, + source_world_size=source_world_size, + limits=limits, + ) + for index, spec in enumerate(fields_value) + ] + codebook_field = common["gefen_codebook_field"] + referenced_fields = set() + if codebook_field is not None: + if type(codebook_field) is not int or codebook_field < 0 or codebook_field >= len(fields): + raise ValueError("sharded portable codebook field is invalid") + field = fields[codebook_field] + if ( + field["fqn"] is not None + or field["name"] != "gefen_codebook" + or field["shape"] != [256] + or field["kind"] != "common" + or field["nonnegative"] + ): + raise ValueError("sharded portable codebook field is incompatible") + referenced_fields.add(codebook_field) + parameters_value = value["parameters"] + manifest_fqns = {shard.parameter.fqn for shard in manifest.shards} + if type(parameters_value) is not dict or set(parameters_value) != manifest_fqns: + raise ValueError("sharded portable parameters do not match the manifest") + parameters = { + fqn: _validate_loaded_parameter_record( + fqn, + parameters_value[fqn], + implementation=implementation, + policy=policy, + manifest=manifest, + fields=fields, + global_step=global_step, + referenced_fields=referenced_fields, + ) + for fqn in sorted(parameters_value) + } + allowed_layouts = ( + { + ParameterLayout.REPLICATED, + ParameterLayout.FLATTENED_ELEMENT_SHARD, + ParameterLayout.DTENSOR_1D_DEFAULT_WORLD, + } + if implementation == "gefen.Gefen" + else { + ParameterLayout.REPLICATED, + ParameterLayout.WHOLE_PARAMETER_OWNER, + } + ) + if any(shard.layout not in allowed_layouts for shard in manifest.shards): + raise ValueError("sharded portable source layout is unsupported for its implementation") + _normalize_catalog( + { + fqn: { + "identity": record["identity"], + "algorithm_options": record["algorithm_options"], + } + for fqn, record in parameters.items() + }, + implementation, + manifest, + policy, + ) + for fqn, record in parameters.items(): + shards = manifest.for_parameter(fqn) + if record["state_variant"] == "initialized_factored" and any( + shard.layout is not ParameterLayout.REPLICATED for shard in shards + ): + raise ValueError("factored sharded portable state requires replicated source shards") + if ( + implementation == "gefen.GefenMuon" + and any(shard.layout is ParameterLayout.WHOLE_PARAMETER_OWNER for shard in shards) + and record["algorithm_options"]["sharded_mode"] != "distributed" + ): + raise ValueError("whole-parameter Muon state requires sharded_mode='distributed'") + if referenced_fields != set(range(len(fields))): + raise ValueError("sharded portable metadata contains unreferenced fields") + if codebook_field is None and any( + record["state_variant"].startswith("initialized_") + for record in parameters.values() + ): + raise ValueError("initialized sharded portable state requires a codebook") + completion = _strict_record( + value["completion"], + {"metadata_digest", "payload_digest"}, + name="sharded portable completion", + ) + for key in ("metadata_digest", "payload_digest"): + digest = completion[key] + if type(digest) is not str or len(digest) != 64: + raise ValueError("sharded portable completion digest is invalid") + try: + bytes.fromhex(digest) + except ValueError as exc: + raise ValueError("sharded portable completion digest is invalid") from exc + semantic = {key: value[key] for key in value if key != "completion"} + if portable_state_digest(semantic) != completion["metadata_digest"]: + raise ValueError("sharded portable metadata digest mismatch") + payload_hasher = hashlib.sha256() + for field in fields: + payload_hasher.update(field["key"].encode("ascii")) + for digest in field["source_chunk_sha256"]: + payload_hasher.update(bytes.fromhex(digest)) + if payload_hasher.hexdigest() != completion["payload_digest"]: + raise ValueError("sharded portable payload digest mismatch") + return { + **value, + "source_process_group": process_group_record, + "policy": policy, + "source_manifest": manifest_record, + "parameters": parameters, + "fields": fields, + } + + +def _validate_field_byte_limits( + fields, + *, + binding: CheckpointProcessGroupBinding, + limits, +): + coordinate = binding.identity.ordered_members.index(binding.local_member) + parts = len(binding.identity.ordered_members) + element_size = torch.tensor([], dtype=torch.float32).element_size() + global_bytes = 0 + local_bytes = 0 + for field in fields: + field_bytes = field["numel"] * element_size + if global_bytes > limits.max_collective_tensor_bytes - field_bytes: + raise ValueError("sharded portable fields exceed max_collective_tensor_bytes") + global_bytes += field_bytes + local_chunk_bytes = ( + _chunk_segment(field["numel"], parts, coordinate).length + * element_size + ) + if local_bytes > limits.max_fragment_tensor_bytes - local_chunk_bytes: + raise ValueError("sharded portable local fields exceed max_fragment_tensor_bytes") + local_bytes += local_chunk_bytes + return global_bytes, local_bytes + + +def _validate_source_storage_chunk_limit(length: int, *, limits): + if type(length) is not int or length < 0: + raise ValueError("sharded portable DCP source storage chunk has invalid size") + if length * torch.tensor([], dtype=torch.float32).element_size() > limits.max_fragment_tensor_bytes: + raise ValueError("sharded portable DCP source storage chunk exceeds max_fragment_tensor_bytes") + + +def _decode_metadata_tensor(value, *, limits): + from gefen.portable_wire import ( + _parse_canonical_wire_metadata, + _reconstruct_canonical_wire_value, + ) + + if ( + type(value) is not torch.Tensor + or value.dtype is not torch.uint8 + or value.device.type != "cpu" + or value.ndim != 1 + or not value.is_contiguous() + or value.numel() < 1 + or value.numel() > limits.max_metadata_bytes + ): + raise ValueError("sharded portable metadata tensor is invalid") + raw = bytes(memoryview(value.numpy())) + prepared = _parse_canonical_wire_metadata( + raw, + limits=limits._wire_limits(collective=True), + ) + if prepared.tensor_specs: + raise ValueError("sharded portable metadata must not contain tensor payloads") + document = _reconstruct_canonical_wire_value(prepared, ()) + return _normalize_sharded_metadata(document, limits=limits) + + +def _metadata_storage_spec(checkpoint_metadata, *, namespace: str, limits): + from torch.distributed.checkpoint.metadata import Metadata + from gefen.portable_dcp import _flat_key, _validate_tensor_metadata + + if type(checkpoint_metadata) is not Metadata: + raise TypeError("portable DCP requires exact checkpoint Metadata") + entries = checkpoint_metadata.state_dict_metadata + if type(entries) is not dict: + raise TypeError("portable DCP metadata entries must be a dict") + key = _flat_key(namespace, _DCP_SHARDED_METADATA_KEY) + if key not in entries: + return None + dtype, shape, nbytes = _validate_tensor_metadata( + entries[key], + name="sharded portable DCP metadata", + limits=limits._wire_limits(collective=True), + ) + if dtype is not torch.uint8 or len(shape) != 1 or nbytes < 1 or nbytes > limits.max_metadata_bytes: + raise ValueError("sharded portable DCP metadata exceeds limits") + return shape + + +def _read_sharded_metadata( + *, + storage_reader, + checkpoint_metadata, + binding: CheckpointProcessGroupBinding, + namespace: str, + limits, + transaction_id: str, +): + import torch.distributed.checkpoint as dcp + from gefen.portable_collective import _collective_unanimous_status + + shape = None + tensor = None + state = None + planner = None + error = None + try: + shape = _metadata_storage_spec( + checkpoint_metadata, + namespace=namespace, + limits=limits, + ) + if shape is not None: + tensor = torch.empty(shape, dtype=torch.uint8, device="cpu") + state = {namespace: {_DCP_SHARDED_METADATA_KEY: tensor}} + planner = dcp.DefaultLoadPlanner(allow_partial_load=True) + except Exception as exc: + error = exc + shape_bytes = repr(tuple(shape) if shape is not None else None).encode("ascii") + _collective_unanimous_status( + binding, + error, + operation="portable_dcp_sharded_metadata_prepare", + transaction_id=transaction_id, + context_digest=hashlib.sha256( + b"gefen.portable_dcp_sharded.metadata.v1\0" + + namespace.encode("utf-8") + + b"\0" + + shape_bytes + ).digest(), + limits=limits._wire_limits(collective=True), + ) + if shape is None: + return None + assert tensor is not None and state is not None and planner is not None + dcp.load( + state, + storage_reader=storage_reader, + planner=planner, + process_group=binding.process_group, + ) + return _decode_metadata_tensor(tensor, limits=limits) + + +def _validate_sharded_dcp_entries( + checkpoint_metadata, + metadata, + *, + namespace: str, + binding: CheckpointProcessGroupBinding, + limits, + checkpoint_keys=None, +): + from torch.distributed.checkpoint.metadata import ( + ChunkStorageMetadata, + TensorProperties, + TensorStorageMetadata, + ) + from gefen.portable_dcp import _flat_key + + entries = checkpoint_metadata.state_dict_metadata + metadata_key = _flat_key(namespace, _DCP_SHARDED_METADATA_KEY) + child_keys = {metadata_key} + for field in metadata["fields"]: + child_keys.add(_flat_key(namespace, field["key"])) + expected = child_keys if checkpoint_keys is None else set(checkpoint_keys) + if ( + not child_keys.issubset(expected) + or set(entries) != expected + or any(type(key) is not str for key in entries) + ): + raise ValueError("sharded portable DCP checkpoint has unexpected tensor keys") + planner_data = checkpoint_metadata.planner_data + expected_paths = {key: tuple(key.split(".", 1)) for key in expected} + if type(planner_data) is not dict or planner_data != expected_paths: + raise ValueError("sharded portable DCP checkpoint has incompatible planner paths") + _validate_field_byte_limits( + metadata["fields"], + binding=binding, + limits=limits, + ) + current_parts = len(binding.identity.ordered_members) + from gefen.portable_dcp import _validate_tensor_metadata + + _validate_tensor_metadata( + entries[metadata_key], + name="sharded portable DCP metadata", + limits=limits._wire_limits(collective=True), + ) + total_chunks = len(entries[metadata_key].chunks) + if ( + total_chunks > limits.max_container_items + or total_chunks > limits.max_tree_nodes + ): + raise ValueError("sharded portable DCP chunks exceed aggregate limits") + for field in metadata["fields"]: + _validate_integrity_assignment_limits( + field, + target_parts=current_parts, + limits=limits, + ) + entry = entries[_flat_key(namespace, field["key"])] + if type(entry) is not TensorStorageMetadata: + raise TypeError("sharded portable field requires tensor DCP metadata") + properties = entry.properties + if ( + type(properties) is not TensorProperties + or properties.dtype is not torch.float32 + 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 + or type(entry.size) is not torch.Size + or any(type(dimension) is not int for dimension in entry.size) + or tuple(entry.size) != (field["numel"],) + or type(entry.chunks) is not list + or not entry.chunks + or len(entry.chunks) > limits.max_container_items + ): + raise ValueError("sharded portable DCP field metadata is incompatible") + total_chunks += len(entry.chunks) + if ( + total_chunks > limits.max_container_items + or total_chunks > limits.max_tree_nodes + ): + raise ValueError("sharded portable DCP chunks exceed aggregate limits") + intervals = [] + for chunk in entry.chunks: + if ( + type(chunk) is not ChunkStorageMetadata + or type(chunk.offsets) is not torch.Size + or type(chunk.sizes) is not torch.Size + or len(chunk.offsets) != 1 + or len(chunk.sizes) != 1 + or type(chunk.offsets[0]) is not int + or type(chunk.sizes[0]) is not int + ): + raise ValueError("sharded portable DCP field chunk is invalid") + offset = chunk.offsets[0] + length = chunk.sizes[0] + if offset < 0 or length < 0 or offset + length > field["numel"]: + raise ValueError("sharded portable DCP field chunk exceeds its tensor") + _validate_source_storage_chunk_limit(length, limits=limits) + intervals.append((offset, offset + length)) + intervals.sort() + cursor = 0 + for start, stop in intervals: + if start != cursor: + raise ValueError("sharded portable DCP chunks must cover each field exactly once") + cursor = stop + if cursor != field["numel"]: + raise ValueError("sharded portable DCP chunks do not cover their field") + local_chunk = _chunk_segment( + field["numel"], + current_parts, + binding.identity.ordered_members.index(binding.local_member), + ) + if local_chunk.length * torch.tensor([], dtype=torch.float32).element_size() > limits.max_fragment_tensor_bytes: + raise ValueError("sharded portable DCP local field exceeds max_fragment_tensor_bytes") + + +def _allocate_sharded_load_state( + metadata, + *, + metadata_tensor: torch.Tensor, + namespace: str, + binding: CheckpointProcessGroupBinding, +): + envelope = {_DCP_SHARDED_METADATA_KEY: metadata_tensor} + for field in metadata["fields"]: + envelope[field["key"]] = _allocate_canonical_dtensor( + global_numel=field["numel"], + dtype=torch.float32, + binding=binding, + ) + return {namespace: envelope} + + +def _source_chunk_assignments(numel: int, source_parts: int, target_parts: int): + assigned = [[] for _ in range(target_parts)] + local_offsets = [0] * target_parts + for source_coordinate in range(source_parts): + target_coordinate = source_coordinate % target_parts + source_chunk = _chunk_segment(numel, source_parts, source_coordinate) + assigned[target_coordinate].append( + ( + source_coordinate, + _LogicalSegment( + source_chunk.global_offset, + source_chunk.length, + local_offsets[target_coordinate], + ), + ) + ) + local_offsets[target_coordinate] += source_chunk.length + return tuple(tuple(items) for items in assigned) + + +def _validate_integrity_assignment_limits(field, *, target_parts: int, limits): + assignments = _source_chunk_assignments( + field["numel"], + len(field["source_chunk_sha256"]), + target_parts, + ) + element_size = torch.tensor([], dtype=torch.float32).element_size() + source_parts = len(field["source_chunk_sha256"]) + if max( + source_parts * hashlib.sha256().digest_size * 4, + source_parts * 4, + ) > limits.max_fragment_tensor_bytes: + raise ValueError("sharded portable integrity digest scratch exceeds max_fragment_tensor_bytes") + for items in assignments: + local_numel = max( + ( + segment.local_offset + segment.length + for _source_coordinate, segment in items + ), + default=0, + ) + if local_numel * element_size > limits.max_fragment_tensor_bytes: + raise ValueError("sharded portable integrity scratch exceeds max_fragment_tensor_bytes") + + +def _verify_field_integrity( + local_canonical: torch.Tensor, + field, + *, + binding: CheckpointProcessGroupBinding, + phase_vote=None, +): + def finish_phase(label, error): + if phase_vote is not None: + phase_vote(label, error) + if error is not None: + raise error + + target_parts = None + source_parts = None + assignments = None + assigned_values = None + error = None + try: + target_parts = len(binding.identity.ordered_members) + source_parts = len(field["source_chunk_sha256"]) + assignments = _source_chunk_assignments(field["numel"], source_parts, target_parts) + target_segments = tuple(tuple(segment for _source_coordinate, segment in items) for items in assignments) + source_segments = _canonical_segments(field["numel"], target_parts) + assigned_values = _redistribute_flat_field( + local_canonical, + source_segments_by_rank=source_segments, + target_segments_by_rank=target_segments, + global_numel=field["numel"], + binding=binding, + phase_vote=phase_vote, + ) + except Exception as exc: + error = exc + finish_phase("integrity_projection", error) + + digests = None + owners = None + error = None + try: + coordinate = binding.identity.ordered_members.index(binding.local_member) + digests = torch.zeros( + (source_parts, hashlib.sha256().digest_size), + dtype=torch.int32, + device=binding.collective_device, + ) + owners = torch.zeros(source_parts, dtype=torch.int32, device=binding.collective_device) + for source_coordinate, segment in assignments[coordinate]: + value = assigned_values[segment.local_offset : segment.local_offset + segment.length] + digest = _tensor_sha256(value) + digests[source_coordinate].copy_(torch.tensor(list(digest), dtype=torch.int32, device=binding.collective_device)) + owners[source_coordinate] = 1 + except Exception as exc: + error = exc + finish_phase("integrity_reduce_prep", error) + + error = None + try: + if target_parts > 1: + dist.all_reduce(digests, op=dist.ReduceOp.SUM, group=binding.process_group) + except Exception as exc: + error = exc + finish_phase("integrity_digest_reduce", error) + + error = None + try: + if target_parts > 1: + dist.all_reduce(owners, op=dist.ReduceOp.SUM, group=binding.process_group) + except Exception as exc: + error = exc + finish_phase("integrity_owner_reduce", error) + + error = None + try: + if not bool((owners == 1).all()): + raise RuntimeError("portable DCP integrity chunks do not have one verifier") + actual = [bytes(row.to(torch.uint8).cpu().tolist()).hex() for row in digests] + if actual != field["source_chunk_sha256"]: + raise ValueError("sharded portable DCP field integrity digest mismatch") + except Exception as exc: + error = exc + finish_phase("integrity_finish", error) + + +def _load_sharded_payloads( + *, + storage_reader, + checkpoint_metadata, + metadata, + metadata_tensor: torch.Tensor | None = None, + metadata_shape=None, + binding: CheckpointProcessGroupBinding, + namespace: str, + limits, + transaction_id: str, + context_digest: bytes, + checkpoint_keys=None, +): + import torch.distributed.checkpoint as dcp + from gefen.portable_collective import _collective_unanimous_status + + state = None + planner = None + error = None + try: + if metadata_tensor is None: + if metadata_shape is None: + raise ValueError("sharded portable metadata allocation requires its shape") + metadata_tensor = torch.empty(metadata_shape, dtype=torch.uint8, device="cpu") + _validate_sharded_dcp_entries( + checkpoint_metadata, + metadata, + namespace=namespace, + binding=binding, + limits=limits, + checkpoint_keys=checkpoint_keys, + ) + state = _allocate_sharded_load_state( + metadata, + metadata_tensor=metadata_tensor, + namespace=namespace, + binding=binding, + ) + planner = dcp.DefaultLoadPlanner() + except Exception as exc: + error = exc + _collective_unanimous_status( + binding, + error, + operation="portable_dcp_sharded_load_allocate", + transaction_id=transaction_id, + context_digest=context_digest, + limits=limits._wire_limits(collective=True), + ) + assert state is not None and planner is not None + error = None + try: + dcp.load( + state, + storage_reader=storage_reader, + planner=planner, + process_group=binding.process_group, + ) + except Exception as exc: + error = exc + _collective_unanimous_status( + binding, + error, + operation="portable_dcp_sharded_load_dcp_complete", + transaction_id=transaction_id, + context_digest=context_digest, + limits=limits._wire_limits(collective=True), + ) + envelope = state[namespace] + local_fields = {} + phase_vote = _make_phase_voter( + binding, + transaction_id=transaction_id, + context_digest=context_digest, + limits=limits, + operation_prefix="portable_dcp_sharded_integrity_phase", + ) + for field in metadata["fields"]: + def field_vote(label, error, *, field_index=field["index"]): + phase_vote("field.{}.{}".format(field_index, label), error) + + local = None + error = None + try: + local = _canonical_local_tensor(envelope[field["key"]]) + if local.dtype is not torch.float32 or local.ndim != 1: + raise TypeError("sharded portable DCP loaded an incompatible field") + for start in range(0, local.numel(), 1 << 20): + chunk = local[start : start + (1 << 20)] + if not bool(torch.isfinite(chunk).all()): + raise ValueError("sharded portable DCP field must be finite") + if field["nonnegative"] and not bool((chunk >= 0).all()): + raise ValueError("sharded portable DCP field must be nonnegative") + except Exception as exc: + error = exc + field_digest = hashlib.sha256( + context_digest + + field["index"].to_bytes(8, "big", signed=False) + ).digest() + _collective_unanimous_status( + binding, + error, + operation="portable_dcp_sharded_load_field", + transaction_id=transaction_id, + context_digest=field_digest, + limits=limits._wire_limits(collective=True), + ) + assert local is not None + _verify_field_integrity( + local, + field, + binding=binding, + phase_vote=field_vote, + ) + local_fields[field["index"]] = local + return local_fields + + +def _target_shards_in_member_order(manifest, fqn, binding): + by_member = { + shard.local_member: shard for shard in manifest.for_parameter(fqn) + } + if set(by_member) != set(binding.identity.ordered_members): + raise ValueError("portable target shards do not match checkpoint members") + return tuple(by_member[member] for member in binding.identity.ordered_members) + + +def _project_loaded_field( + local_canonical, + field, + *, + target_shards, + binding: CheckpointProcessGroupBinding, + replicate=False, + phase_vote=None, +): + mode = None + target_segments = None + result = None + error = None + try: + if not replicate and field["kind"] == "dense": + mode = "dense" + else: + mode = "flat" + if replicate or field["kind"] == "common": + target_segments = tuple( + _full_segment(field["numel"]) + for _member in binding.identity.ordered_members + ) + elif field["kind"] == "special": + target_segments = tuple( + _special_segments_for_shard( + shard, + numel=field["numel"], + unique_replicas=False, + ) + for shard in target_shards + ) + else: + raise ValueError("portable loaded field has an unknown geometry kind") + except Exception as exc: + error = exc + if phase_vote is not None: + phase_vote("projection_prep", error) + if error is not None: + raise error + + error = None + try: + if mode == "dense": + result = _redistribute_canonical_to_dense( + local_canonical, + target_shards=target_shards, + global_numel=field["numel"], + binding=binding, + phase_vote=phase_vote, + ) + else: + assert target_segments is not None + result = _redistribute_flat_field( + local_canonical, + source_segments_by_rank=_canonical_segments( + field["numel"], + len(binding.identity.ordered_members), + ), + target_segments_by_rank=target_segments, + global_numel=field["numel"], + binding=binding, + phase_vote=phase_vote, + ) + result = result.cpu() + except Exception as exc: + error = exc + if phase_vote is not None: + phase_vote("projection_finish", error) + if error is not None: + raise error + assert result is not None + return result + + +def _factored_axis_segments(shard: ShardIdentity, *, axis: int): + """Return the sorted compact factor slice needed by one target shard.""" + + if type(shard) is not ShardIdentity or axis not in (0, 1): + raise ValueError("factored factor routing requires a target shard and axis") + shape = shard.parameter.global_shape + if len(shape) != 2: + raise ValueError("factored factor routing requires a matrix parameter") + axis_numel = shape[axis] + placement = shard.placements[0] + if shard.layout is ParameterLayout.REPLICATED or ( + shard.layout is ParameterLayout.DTENSOR_1D_DEFAULT_WORLD + and placement.kind is PlacementKind.REPLICATE + ): + return _full_segment(axis_numel) + if shard.layout is ParameterLayout.WHOLE_PARAMETER_OWNER: + return _full_segment(axis_numel) if shard.local_member == shard.owner else () + if shard.layout is ParameterLayout.DTENSOR_1D_DEFAULT_WORLD: + region = shard.logical_region + if region.numel == 0: + return () + start = region.offsets[axis] + length = region.lengths[axis] + return (_LogicalSegment(start, length, 0),) if length else () + if shard.layout is not ParameterLayout.FLATTENED_ELEMENT_SHARD: + raise ValueError("factored factor routing has an unsupported target layout") + logical_slice = shard.logical_slice + if logical_slice.length == 0: + return () + rows, columns = shape + start = logical_slice.flat_offset + stop = start + logical_slice.length + if axis == 0: + row_start = start // columns + row_stop = (stop - 1) // columns + 1 + return (_LogicalSegment(row_start, row_stop - row_start, 0),) + if logical_slice.length >= columns: + return _full_segment(columns) + first_column = start % columns + final_column = (stop - 1) % columns + 1 + if start // columns == (stop - 1) // columns: + return (_LogicalSegment(first_column, logical_slice.length, 0),) + head_length = final_column + tail_length = columns - first_column + segments = [] + local_offset = 0 + if head_length: + segments.append(_LogicalSegment(0, head_length, local_offset)) + local_offset += head_length + if tail_length: + segments.append(_LogicalSegment(first_column, tail_length, local_offset)) + return tuple(segments) + + +def _segment_global_indices(segments): + pieces = [ + torch.arange( + segment.global_offset, + segment.global_offset + segment.length, + dtype=torch.int64, + ) + for segment in segments + if segment.length + ] + return torch.cat(pieces) if pieces else torch.empty(0, dtype=torch.int64) + + +def _project_loaded_factored_axis( + local_canonical, + field, + *, + target_shards, + axis: int, + binding: CheckpointProcessGroupBinding, + phase_vote=None, +): + target_segments = None + result = None + indices = None + error = None + try: + target_segments = tuple( + _factored_axis_segments(shard, axis=axis) + for shard in target_shards + ) + except Exception as exc: + error = exc + if phase_vote is not None: + phase_vote("factored_axis_prep", error) + if error is not None: + raise error + + error = None + try: + result = _redistribute_flat_field( + local_canonical, + source_segments_by_rank=_canonical_segments( + field["numel"], + len(binding.identity.ordered_members), + ), + target_segments_by_rank=target_segments, + global_numel=field["numel"], + binding=binding, + phase_vote=phase_vote, + ).cpu() + coordinate = binding.identity.ordered_members.index(binding.local_member) + indices = _segment_global_indices(target_segments[coordinate]) + if result.numel() != indices.numel(): + raise ValueError("factored factor routing produced invalid local geometry") + except Exception as exc: + error = exc + if phase_vote is not None: + phase_vote("factored_axis_finish", error) + if error is not None: + raise error + assert result is not None and indices is not None + return result, indices + + +def _local_factored_projection( + row, + row_indices, + column, + column_indices, + denominator, + shard: ShardIdentity, +): + if ( + type(row) is not torch.Tensor + or type(row_indices) is not torch.Tensor + or type(column) is not torch.Tensor + or type(column_indices) is not torch.Tensor + or type(denominator) is not torch.Tensor + or row.dtype is not torch.float32 + or row_indices.dtype is not torch.int64 + or column.dtype is not torch.float32 + or column_indices.dtype is not torch.int64 + or denominator.dtype is not torch.float32 + or row.ndim != 1 + or row_indices.ndim != 1 + or column.ndim != 1 + or column_indices.ndim != 1 + or denominator.numel() != 1 + or row.numel() != row_indices.numel() + or column.numel() != column_indices.numel() + or len(shard.parameter.global_shape) != 2 + ): + raise ValueError("factored projection inputs do not match the target parameter") + if row_indices.numel() and not bool(row_indices[1:].gt(row_indices[:-1]).all()): + raise ValueError("factored row routing indices must be strictly increasing") + if column_indices.numel() and not bool(column_indices[1:].gt(column_indices[:-1]).all()): + raise ValueError("factored column routing indices must be strictly increasing") + if not bool(torch.isfinite(denominator).all()) or not bool((denominator > 0).all()): + raise ValueError("factored row denominator must be finite and positive") + local_numel = _dense_local_numel(shard, unique_replicas=False) + if local_numel == 0: + if row.numel() or column.numel(): + raise ValueError("empty target shard received factored factor payloads") + return None + global_indices = _dense_global_indices( + shard, + unique_replicas=False, + device=torch.device("cpu"), + ) + columns = shard.parameter.global_shape[1] + required_rows = torch.div(global_indices, columns, rounding_mode="floor") + required_columns = torch.remainder(global_indices, columns) + row_offsets = torch.searchsorted(row_indices, required_rows) + column_offsets = torch.searchsorted(column_indices, required_columns) + if ( + bool((row_offsets >= row_indices.numel()).any()) + or bool((column_offsets >= column_indices.numel()).any()) + or not bool(row_indices[row_offsets].eq(required_rows).all()) + or not bool(column_indices[column_offsets].eq(required_columns).all()) + ): + raise ValueError("factored factor routing does not cover the target shard") + return row[row_offsets].mul(column[column_offsets]).div_(denominator.reshape(())) + + +def _parameter_options_compatible( + source, + target, + *, + implementation: str, +): + from gefen.portable_state import _PLAIN_OPTION_KEYS, _values_equal + + if implementation != "gefen.Gefen": + return _values_equal(source, target) + invariant = _PLAIN_OPTION_KEYS - {"second_moment_policy"} + if any(not _values_equal(source[key], target[key]) for key in invariant): + return False + source_second = source["second_moment_policy"] + target_second = target["second_moment_policy"] + return source_second == target_second or ( + source_second == "factored" and target_second == "block" + ) + + +def _projected_target_byte_plan( + optimizer, + metadata, + *, + binding: CheckpointProcessGroupBinding, + limits, +): + """Validate aggregate target fp32 payloads before DCP reads any field.""" + + from gefen import portable_runtime as runtime + from gefen.portable_state import _validate_portable_projection_policy + + implementation = runtime._optimizer_implementation(optimizer) + if metadata["implementation"] != implementation: + raise ValueError("sharded portable implementation does not match the target") + prepared = runtime._prepare_local_structure( + optimizer, + implementation, + binding, + limits, + include_payload=False, + ) + runtime._validate_prepared_local_state(optimizer, implementation, prepared) + target_descriptor = prepared["target_descriptor"] + _validate_portable_projection_policy( + metadata["policy"], + target_descriptor["policy"], + implementation, + ) + if set(metadata["parameters"]) != set(target_descriptor["catalog"]): + raise ValueError("sharded portable parameter catalog does not match the target") + target_manifest = optimizer._gefen_sharding_manifest + members = binding.identity.ordered_members + projected_numel = [ + 256 if metadata["common"]["gefen_codebook_field"] is not None else 0 + for _member in members + ] + largest_factored_route_bytes = [0 for _member in members] + largest_projection_routing_bytes = [0 for _member in members] + + def include_projection_send( + global_numel, + target_segments_by_rank, + *, + element_size, + ): + send_numel = _canonical_projection_send_numel( + global_numel, + target_segments_by_rank=target_segments_by_rank, + parts=len(members), + ) + for source_coordinate, value in enumerate(send_numel): + largest_projection_routing_bytes[source_coordinate] = max( + largest_projection_routing_bytes[source_coordinate], + value * element_size, + ) + + if metadata["common"]["gefen_codebook_field"] is not None: + include_projection_send( + 256, + tuple(_full_segment(256) for _member in members), + element_size=4, + ) + for fqn in sorted(metadata["parameters"]): + source_record = metadata["parameters"][fqn] + target_record = target_descriptor["catalog"][fqn] + if source_record["identity"] != target_record["identity"]: + raise ValueError("sharded portable parameter identity does not match the target") + source_options = source_record["algorithm_options"] + target_options = target_record["algorithm_options"] + if not _parameter_options_compatible(source_options, target_options, implementation=implementation): + raise ValueError("sharded portable algorithm options do not match the target") + source_second = source_options.get("second_moment_policy") + target_second = target_options.get("second_moment_policy") + if source_second == "block" and target_second == "factored": + raise ValueError("portable block-to-factored second-moment migration is unsupported") + target_shards = _target_shards_in_member_order(target_manifest, fqn, binding) + variant = source_record["state_variant"] + if variant.startswith("initialized_"): + dense_send_numel = _dense_projection_send_numel( + target_shards[0].parameter.numel, + target_shards=target_shards, + parts=len(members), + ) + for coordinate, send_numel in enumerate(dense_send_numel): + canonical_numel = _chunk_segment( + target_shards[0].parameter.numel, + len(members), + coordinate, + ).length + target_numel = _dense_local_numel( + target_shards[coordinate], + unique_replicas=False, + ) + largest_projection_routing_bytes[coordinate] = max( + largest_projection_routing_bytes[coordinate], + len(members) * 8, + canonical_numel * 8, + send_numel * 4, + send_numel * 8, + target_numel * 8, + ) + for coordinate, target_shard in enumerate(target_shards): + local_numel = _dense_local_numel(target_shard, unique_replicas=False) + if variant.startswith("initialized_"): + projected_numel[coordinate] += local_numel + if implementation == "gefen.Gefen" and variant == "initialized_dense": + projected_numel[coordinate] += local_numel + elif implementation == "gefen.Gefen" and variant == "initialized_factored": + row_segments = _factored_axis_segments(target_shard, axis=0) + column_segments = _factored_axis_segments(target_shard, axis=1) + row_numel = _local_numel(row_segments) + column_numel = _local_numel(column_segments) + route_numel = row_numel + column_numel + 1 + if target_second != "factored": + route_numel += local_numel + largest_factored_route_bytes[coordinate] = max( + largest_factored_route_bytes[coordinate], + route_numel * 4, + row_numel * 8, + column_numel * 8, + local_numel * 8, + ) + if target_second == "factored": + projected_numel[coordinate] += ( + target_shard.parameter.global_shape[0] + + target_shard.parameter.global_shape[1] + ) + else: + projected_numel[coordinate] += local_numel + elif implementation == "gefen.GefenMuon" and variant == "initialized_dense_normuon": + normuon_numel = target_shard.parameter.global_shape[0] + projected_numel[coordinate] += _local_numel( + _special_segments_for_shard( + target_shard, + numel=normuon_numel, + unique_replicas=False, + ) + ) + if implementation == "gefen.Gefen" and variant == "initialized_factored": + rows, columns = target_shards[0].parameter.global_shape + include_projection_send( + rows, + tuple( + _factored_axis_segments(target_shard, axis=0) + for target_shard in target_shards + ), + element_size=4, + ) + include_projection_send( + columns, + tuple( + _factored_axis_segments(target_shard, axis=1) + for target_shard in target_shards + ), + element_size=4, + ) + include_projection_send( + 1, + tuple(_full_segment(1) for _member in members), + element_size=4, + ) + elif implementation == "gefen.GefenMuon" and variant == "initialized_dense_normuon": + normuon_numel = target_shards[0].parameter.global_shape[0] + include_projection_send( + normuon_numel, + tuple( + _special_segments_for_shard( + target_shard, + numel=normuon_numel, + unique_replicas=False, + ) + for target_shard in target_shards + ), + element_size=4, + ) + element_size = torch.tensor([], dtype=torch.float32).element_size() + projected_bytes = [numel * element_size for numel in projected_numel] + factored_route_bytes = largest_factored_route_bytes + projection_routing_bytes = largest_projection_routing_bytes + if any(value > limits.max_fragment_tensor_bytes for value in projected_bytes): + raise ValueError("sharded portable projected target state exceeds max_fragment_tensor_bytes") + if any(value > limits.max_fragment_tensor_bytes for value in factored_route_bytes): + raise ValueError("sharded portable factored routing scratch exceeds max_fragment_tensor_bytes") + if any(value > limits.max_fragment_tensor_bytes for value in projection_routing_bytes): + raise ValueError("sharded portable projection routing scratch exceeds max_fragment_tensor_bytes") + digest = runtime._context_digest( + { + "target_descriptor": target_descriptor, + "projected_fp32_bytes": projected_bytes, + "factored_route_fp32_bytes": factored_route_bytes, + "projection_routing_bytes": projection_routing_bytes, + } + ) + return ( + digest, + tuple(projected_bytes), + tuple(factored_route_bytes), + tuple(projection_routing_bytes), + ) + + +def _stage_sharded_import( + optimizer, + metadata, + local_fields, + *, + binding: CheckpointProcessGroupBinding, + limits, + transaction_id: str, +): + from gefen import portable_runtime as runtime + from gefen.portable import ( + _recompress_dense_momentum, + _reduce_block_second_moment, + ) + from gefen.portable_state import ( + _SECOND_MOMENT_PROJECTION_FACTORED_TO_BLOCK, + _normalize_common, + _validate_portable_projection_policy, + ) + from gefen.portable_collective import _collective_unanimous_status + + implementation = None + live_token = None + target_descriptor = None + target_manifest = None + local_by_fqn = None + fields = None + target_digest = bytes.fromhex(metadata["completion"]["metadata_digest"]) + error = None + try: + implementation = runtime._optimizer_implementation(optimizer) + if metadata["implementation"] != implementation: + raise ValueError("sharded portable implementation does not match the target") + live_token = runtime._portable_live_token(optimizer) + prepared = runtime._prepare_local_structure( + optimizer, + implementation, + binding, + limits, + include_payload=False, + ) + runtime._validate_prepared_local_state( + optimizer, + implementation, + prepared, + ) + target_descriptor = prepared["target_descriptor"] + _validate_portable_projection_policy( + metadata["policy"], + target_descriptor["policy"], + implementation, + ) + if set(metadata["parameters"]) != set(target_descriptor["catalog"]): + raise ValueError("sharded portable parameter catalog does not match the target") + target_manifest = optimizer._gefen_sharding_manifest + local_by_fqn = prepared["local_by_fqn"] + fields = {field["index"]: field for field in metadata["fields"]} + if set(local_fields) != set(fields): + raise ValueError("sharded portable local fields are incomplete") + coordinate = binding.identity.ordered_members.index(binding.local_member) + parts = len(binding.identity.ordered_members) + for index, field in fields.items(): + local = local_fields[index] + if ( + type(local) is not torch.Tensor + or local.dtype is not torch.float32 + or local.ndim != 1 + or local.numel() + != _chunk_segment(field["numel"], parts, coordinate).length + ): + raise ValueError("sharded portable local field geometry is invalid") + for fqn in sorted(local_by_fqn): + parameter, shard = local_by_fqn[fqn] + source_record = metadata["parameters"][fqn] + target_record = target_descriptor["catalog"][fqn] + if source_record["identity"] != target_record["identity"]: + raise ValueError( + "sharded portable parameter identity does not match the target" + ) + target_options = target_record["algorithm_options"] + source_options = source_record["algorithm_options"] + if not _parameter_options_compatible( + source_options, + target_options, + implementation=implementation, + ): + raise ValueError("sharded portable algorithm options do not match the target") + source_second = source_options.get("second_moment_policy") + target_second = target_options.get("second_moment_policy") + if source_second == "block" and target_second == "factored": + raise ValueError("portable block-to-factored second-moment migration is unsupported") + if ( + source_second == "factored" + and target_second == "block" + and metadata["policy"]["second_moment_projection"] + != _SECOND_MOMENT_PROJECTION_FACTORED_TO_BLOCK + ): + raise ValueError( + "sharded portable factored-to-block projection is not authorized by its policy" + ) + target_shards = _target_shards_in_member_order( + target_manifest, + fqn, + binding, + ) + variant = source_record["state_variant"] + if variant.startswith("initialized_"): + for target_shard in target_shards: + _dense_local_numel( + target_shard, + unique_replicas=False, + ) + if variant == "initialized_factored": + for target_shard in target_shards: + _factored_axis_segments(target_shard, axis=0) + _factored_axis_segments(target_shard, axis=1) + if ( + implementation == "gefen.GefenMuon" + and variant == "initialized_dense_normuon" + ): + normuon_field = fields[ + source_record["state"]["normuon_v_field"] + ] + for target_shard in target_shards: + _special_segments_for_shard( + target_shard, + numel=normuon_field["numel"], + unique_replicas=False, + ) + if parameter is not None: + runtime._local_dense_shape(shard) + target_context = runtime._target_context( + runtime._base_context(binding, implementation), + target_descriptor, + ) + target_context["document_digest"] = metadata["completion"][ + "metadata_digest" + ] + runtime._preflight_portable_value(target_context, limits) + target_digest = runtime._context_digest(target_context) + except Exception as exc: + error = exc + _collective_unanimous_status( + binding, + error, + operation="portable_dcp_sharded_load_target", + transaction_id=transaction_id, + context_digest=target_digest, + limits=limits._wire_limits(collective=True), + ) + assert ( + implementation is not None + and live_token is not None + and target_descriptor is not None + and target_manifest is not None + and local_by_fqn is not None + and fields is not None + ) + + phase_vote = _make_phase_voter( + binding, + transaction_id=transaction_id, + context_digest=target_digest, + limits=limits, + operation_prefix="portable_dcp_sharded_stage_phase", + ) + + def scoped_vote(prefix): + return lambda label, error: phase_vote("{}.{}".format(prefix, label), error) + + codebook = None + common = None + error = None + try: + codebook_index = metadata["common"]["gefen_codebook_field"] + if codebook_index is not None: + field = fields[codebook_index] + codebook = _project_loaded_field( + local_fields[codebook_index], + field, + target_shards=(), + binding=binding, + replicate=True, + phase_vote=scoped_vote("common.codebook"), + ).reshape(256) + common = _normalize_common( + { + "gefen_global_step": metadata["common"]["gefen_global_step"], + "gefen_codebook": codebook, + "gefen_deterministic": metadata["common"]["gefen_deterministic"], + } + ) + except Exception as exc: + error = exc + phase_vote("common.normalize", error) + assert common is not None + + canonical_parameters = {} + for fqn in sorted(local_by_fqn): + parameter = None + shard = None + source_record = None + target_options = None + target_second = None + target_shards = None + variant = None + source_state = None + local_shape = None + has_local_payload = None + projected = None + error = None + try: + parameter, shard = local_by_fqn[fqn] + source_record = metadata["parameters"][fqn] + target_options = target_descriptor["catalog"][fqn]["algorithm_options"] + source_options = source_record["algorithm_options"] + if not _parameter_options_compatible(source_options, target_options, implementation=implementation): + raise ValueError("sharded portable algorithm options do not match the target") + source_second = source_options.get("second_moment_policy") + target_second = target_options.get("second_moment_policy") + if source_second == "block" and target_second == "factored": + raise ValueError("portable block-to-factored second-moment migration is unsupported") + target_shards = _target_shards_in_member_order(target_manifest, fqn, binding) + variant = source_record["state_variant"] + source_state = source_record["state"] + local_shape = runtime._local_dense_shape(shard) if parameter is not None else None + has_local_payload = local_shape is not None + projected = {} + except Exception as exc: + error = exc + phase_vote("parameter.{}.setup".format(fqn), error) + assert shard is not None and source_record is not None and target_shards is not None and variant is not None and source_state is not None and projected is not None + + if variant == "period_selected": + error = None + try: + if has_local_payload: + projected = {"automatic_period": 1} + except Exception as exc: + error = exc + phase_vote("parameter.{}.period_selected".format(fqn), error) + elif variant.startswith("initialized_"): + error = None + try: + if codebook is None: + raise ValueError("initialized sharded portable state requires a codebook") + except Exception as exc: + error = exc + phase_vote("parameter.{}.initialized_prep".format(fqn), error) + + momentum = None + expected_momentum_numel = None + error = None + try: + momentum_index = source_state["momentum_field"] + momentum = _project_loaded_field( + local_fields[momentum_index], + fields[momentum_index], + target_shards=target_shards, + binding=binding, + phase_vote=scoped_vote("parameter.{}.momentum".format(fqn)), + ) + expected_momentum_numel = 0 if local_shape is None else math.prod(local_shape) + if momentum.numel() != expected_momentum_numel: + raise ValueError("sharded portable momentum does not match its target shard") + if has_local_payload: + momentum = momentum.reshape(local_shape) + except Exception as exc: + error = exc + phase_vote("parameter.{}.momentum_geometry".format(fqn), error) + assert momentum is not None and expected_momentum_numel is not None + + error = None + try: + if has_local_payload: + indices, magnitudes = _recompress_dense_momentum( + momentum, + codebook, + period=1, + step=source_state["step"], + ) + projected = { + "automatic_period": 1, + "step": source_state["step"], + "m_codebook": indices, + "m_magnitude": magnitudes, + } + del momentum + except Exception as exc: + error = exc + phase_vote("parameter.{}.momentum_recompress".format(fqn), error) + + if implementation == "gefen.Gefen" and variant == "initialized_dense": + second = None + error = None + try: + second_index = source_state["second_moment_field"] + second = _project_loaded_field( + local_fields[second_index], + fields[second_index], + target_shards=target_shards, + binding=binding, + phase_vote=scoped_vote("parameter.{}.second_moment".format(fqn)), + ) + if second.numel() != expected_momentum_numel: + raise ValueError("sharded portable second moment does not match its target shard") + if has_local_payload: + projected.update( + { + "vmean": _reduce_block_second_moment( + second.reshape(local_shape), + period=1, + step=source_state["second_moment_step"], + ), + "vmean_step": source_state["second_moment_step"], + } + ) + del second + except Exception as exc: + error = exc + phase_vote("parameter.{}.second_moment_finish".format(fqn), error) + elif implementation == "gefen.Gefen" and variant == "initialized_factored": + row = None + row_indices = None + error = None + try: + row_index = source_state["v_row_field"] + row, row_indices = _project_loaded_factored_axis( + local_fields[row_index], + fields[row_index], + target_shards=target_shards, + axis=0, + binding=binding, + phase_vote=scoped_vote("parameter.{}.factored_row".format(fqn)), + ) + except Exception as exc: + error = exc + phase_vote("parameter.{}.factored_row_finish".format(fqn), error) + assert row is not None and row_indices is not None + + column = None + column_indices = None + error = None + try: + column_index = source_state["v_col_field"] + column, column_indices = _project_loaded_factored_axis( + local_fields[column_index], + fields[column_index], + target_shards=target_shards, + axis=1, + binding=binding, + phase_vote=scoped_vote("parameter.{}.factored_column".format(fqn)), + ) + except Exception as exc: + error = exc + phase_vote("parameter.{}.factored_column_finish".format(fqn), error) + assert column is not None and column_indices is not None + + denominator = None + error = None + try: + denominator_index = source_state["v_row_denominator_field"] + denominator = _project_loaded_field( + local_fields[denominator_index], + fields[denominator_index], + target_shards=(), + binding=binding, + replicate=True, + phase_vote=scoped_vote("parameter.{}.factored_denominator".format(fqn)), + ) + if denominator.numel() != 1: + raise ValueError("sharded portable factored denominator has invalid geometry") + except Exception as exc: + error = exc + phase_vote("parameter.{}.factored_denominator_finish".format(fqn), error) + assert denominator is not None + + error = None + try: + if target_second == "factored" and has_local_payload: + if shard.layout is not ParameterLayout.REPLICATED: + raise ValueError("factored portable state requires a replicated target") + projected.update( + { + "v_row": row, + "v_col": column, + "factored_step": source_state["factored_step"], + } + ) + elif target_second != "factored": + second = _local_factored_projection( + row, + row_indices, + column, + column_indices, + denominator, + shard, + ) + del row, row_indices, column, column_indices, denominator + if has_local_payload: + if second is None: + raise ValueError("initialized target shard is missing factored projection state") + projected.update( + { + "vmean": _reduce_block_second_moment( + second, + period=1, + step=source_state["factored_step"], + ), + "vmean_step": source_state["factored_step"], + } + ) + elif second is not None: + raise ValueError("empty target shard received factored projection state") + del second + except Exception as exc: + error = exc + phase_vote("parameter.{}.factored_projection_finish".format(fqn), error) + elif implementation == "gefen.GefenMuon" and variant == "initialized_dense_normuon": + normuon = None + error = None + try: + normuon_index = source_state["normuon_v_field"] + normuon = _project_loaded_field( + local_fields[normuon_index], + fields[normuon_index], + target_shards=target_shards, + binding=binding, + phase_vote=scoped_vote("parameter.{}.normuon".format(fqn)), + ) + expected_normuon_numel = shard.parameter.global_shape[0] if has_local_payload else 0 + if normuon.numel() != expected_normuon_numel: + raise ValueError("sharded portable NorMuon state does not match its target shard") + if has_local_payload: + projected.update( + { + "normuon_v": normuon.reshape(shard.parameter.global_shape[0], 1), + "normuon_step": source_state["normuon_step"], + } + ) + except Exception as exc: + error = exc + phase_vote("parameter.{}.normuon_finish".format(fqn), error) + + error = None + try: + if parameter is not None: + canonical_parameters[fqn] = {"state": projected} + except Exception as exc: + error = exc + phase_vote("parameter.{}.finish".format(fqn), error) + + staged = None + error = None + try: + 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) + staging_owner = object.__new__(type(optimizer)) + staging_owner.__dict__ = optimizer.__dict__.copy() + staging_owner._deterministic = common["gefen_deterministic"] + has_dtensor_target = any( + target_shard.layout is ParameterLayout.DTENSOR_1D_DEFAULT_WORLD + for _parameter, target_shard in local_by_fqn.values() + ) + if has_dtensor_target: + native.pop("gefen_native_local_shards", None) + for group in native["param_groups"]: + checkpoint_metadata = dict(group[runtime._RANK_LOCAL_GROUP_METADATA_KEY]) + checkpoint_metadata.pop("native_local_shards", None) + group[runtime._RANK_LOCAL_GROUP_METADATA_KEY] = checkpoint_metadata + staging_owner._uses_rank_local_sharded_state = lambda: False + staging_owner._serialized_native_local_shards = lambda: None + staged = staging_owner._stage_load_state_dict(native) + if has_dtensor_target: + staged.__dict__.pop("_uses_rank_local_sharded_state", None) + staged.__dict__.pop("_serialized_native_local_shards", None) + staged._install_rank_local_checkpoint_schema() + runtime._validate_rank_local_transport_schema(staged, binding) + 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("sharded portable import staging produced unsafe publication containers") + if live_token != runtime._portable_live_token(optimizer): + raise RuntimeError("live optimizer state changed during sharded portable import preparation") + except Exception as exc: + error = exc + phase_vote("stage.finalize", error) + assert staged is not None + return staged, live_token + + +def _save_sharded_portable_dcp( + optimizer, + *, + binding: CheckpointProcessGroupBinding, + storage_writer, + transaction_id: str, + limits, + namespace: str, +): + import torch.distributed.checkpoint as dcp + + from gefen.portable_collective import _collective_unanimous_status + + state = None + planner = None + digest = bytes(32) + error = None + try: + state = _prepare_sharded_save_state( + optimizer, + binding=binding, + transaction_id=transaction_id, + limits=limits, + namespace=namespace, + ) + metadata = _decode_metadata_tensor( + state[namespace][_DCP_SHARDED_METADATA_KEY], + limits=limits, + ) + digest = bytes.fromhex(metadata["completion"]["metadata_digest"]) + planner = dcp.DefaultSavePlanner() + except Exception as exc: + error = exc + _collective_unanimous_status( + binding, + error, + operation="portable_dcp_sharded_save_prepare", + transaction_id=transaction_id, + context_digest=digest, + limits=limits._wire_limits(collective=True), + ) + assert state is not None and planner is not None + result = None + error = None + try: + result = dcp.save( + state, + storage_writer=storage_writer, + planner=planner, + process_group=binding.process_group, + ) + except Exception as exc: + error = exc + _collective_unanimous_status( + binding, + error, + operation="portable_dcp_sharded_save_dcp_complete", + transaction_id=transaction_id, + context_digest=digest, + limits=limits._wire_limits(collective=True), + ) + return result + + +def _load_sharded_portable_dcp( + optimizer, + *, + binding: CheckpointProcessGroupBinding, + storage_reader, + checkpoint_metadata, + transaction_id: str, + limits, + namespace: str, +): + from gefen.portable_collective import _collective_unanimous_status + + metadata = None + digest = bytes(32) + error = None + try: + metadata = _read_sharded_metadata( + storage_reader=storage_reader, + checkpoint_metadata=checkpoint_metadata, + binding=binding, + namespace=namespace, + limits=limits, + transaction_id=transaction_id, + ) + if metadata is None: + raise ValueError("checkpoint does not contain a sharded portable DCP envelope") + digest = bytes.fromhex(metadata["completion"]["metadata_digest"]) + except Exception as exc: + error = exc + _collective_unanimous_status( + binding, + error, + operation="portable_dcp_sharded_load_metadata", + transaction_id=transaction_id, + context_digest=digest, + limits=limits._wire_limits(collective=True), + ) + assert checkpoint_metadata is not None and metadata is not None + + projected_digest = digest + error = None + try: + ( + projected_digest, + _projected_bytes, + _scratch_bytes, + _send_bytes, + ) = _projected_target_byte_plan( + optimizer, + metadata, + binding=binding, + limits=limits, + ) + except Exception as exc: + error = exc + _collective_unanimous_status( + binding, + error, + operation="portable_dcp_sharded_load_projected_limits", + transaction_id=transaction_id, + context_digest=projected_digest, + limits=limits._wire_limits(collective=True), + ) + + local_fields = None + error = None + try: + metadata_shape = _metadata_storage_spec( + checkpoint_metadata, + namespace=namespace, + limits=limits, + ) + assert metadata_shape is not None + local_fields = _load_sharded_payloads( + storage_reader=storage_reader, + checkpoint_metadata=checkpoint_metadata, + metadata=metadata, + metadata_shape=metadata_shape, + binding=binding, + namespace=namespace, + limits=limits, + transaction_id=transaction_id, + context_digest=digest, + ) + except Exception as exc: + error = exc + _collective_unanimous_status( + binding, + error, + operation="portable_dcp_sharded_load_payloads", + transaction_id=transaction_id, + context_digest=digest, + limits=limits._wire_limits(collective=True), + ) + assert local_fields is not None + + staged = None + live_token = None + error = None + try: + staged, live_token = _stage_sharded_import( + optimizer, + metadata, + local_fields, + binding=binding, + limits=limits, + transaction_id=transaction_id, + ) + except Exception as exc: + error = exc + _collective_unanimous_status( + binding, + error, + operation="portable_dcp_sharded_load_stage", + transaction_id=transaction_id, + context_digest=digest, + limits=limits._wire_limits(collective=True), + ) + assert staged is not None and live_token is not None + + error = None + try: + from gefen import portable_runtime as runtime + + implementation = runtime._optimizer_implementation(optimizer) + runtime._validate_live_readiness(optimizer, implementation, binding) + if live_token != runtime._portable_live_token(optimizer): + raise RuntimeError("live optimizer state changed after sharded portable import preparation") + except Exception as exc: + error = exc + _collective_unanimous_status( + binding, + error, + operation="portable_dcp_sharded_load_freshness", + transaction_id=transaction_id, + context_digest=digest, + limits=limits._wire_limits(collective=True), + ) + + from gefen.gefen import Gefen + + Gefen._commit_staged_load_state_dict(optimizer, staged) + + +__all__ = [] diff --git a/src/gefen/portable_runtime.py b/src/gefen/portable_runtime.py index 84f1843..8435796 100644 --- a/src/gefen/portable_runtime.py +++ b/src/gefen/portable_runtime.py @@ -14,6 +14,7 @@ from gefen.codebook import CodebookProcessGroupBinding from gefen.contracts import ( IDENTITY_SCHEMA_VERSION, + LogicalRegion, LogicalSlice, ParameterIdentity, ParameterLayout, @@ -23,6 +24,7 @@ ShardPlacement, ShardingManifest, ) +from torch.distributed.tensor import DTensor, Replicate, Shard from gefen.portable import ( _decode_quantized_momentum, _expand_block_second_moment, @@ -42,7 +44,8 @@ from gefen.portable_state import ( PortableStateLimits, _MOMENTUM_PROJECTION, - _SECOND_MOMENT_PROJECTION, + _SECOND_MOMENT_PROJECTION_EXACT, + _SECOND_MOMENT_PROJECTION_FACTORED_TO_BLOCK, _assemble_portable_state_fragments, _build_portable_state_fragment, _derived_role, @@ -53,7 +56,7 @@ _normalize_policy, _normalize_portable_state_fragment, _project_portable_parameter_state, - _values_equal, + _validate_portable_projection_policy, ) from gefen.rebinding import LogicalSlotBinding @@ -71,6 +74,7 @@ )._wire_limits() _PLAIN_GROUP_REQUIRED = frozenset({"params", "param_names", "lr", "beta1", "beta2", "eps", "weight_decay"}) +_RANK_LOCAL_GROUP_METADATA_KEY = "_gefen_checkpoint_metadata" _PLAIN_GROUP_ALLOWED = _PLAIN_GROUP_REQUIRED | {"name"} _MUON_GROUP_REQUIRED = _PLAIN_GROUP_REQUIRED | frozenset( { @@ -134,6 +138,8 @@ ) _MUON_KEYS = frozenset({"name", "automatic_period", "step", "m_codebook", "m_magnitude"}) _NORMUON_KEYS = _MUON_KEYS | {"normuon_v", "normuon_step"} +_RANK_LOCAL_PAYLOAD_KEY_PREFIX = "_gefen_rank_local_payload_" +_RANK_LOCAL_MEMBER_KEY = "_gefen_rank_local_member" def _bounded_utf8_length(value: str, *, limit: int, name: str) -> int: @@ -315,20 +321,34 @@ def _validate_exact_shard_identity(shard) -> ShardIdentity: 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 shard.layout is ParameterLayout.DTENSOR_1D_DEFAULT_WORLD: + if type(shard.logical_slice) is not LogicalRegion: + raise TypeError("portable DTensor state requires exact LogicalRegion values") + if ( + type(shard.logical_slice.offsets) is not tuple + or type(shard.logical_slice.lengths) is not tuple + or any(type(value) is not int for value in shard.logical_slice.offsets + shard.logical_slice.lengths) + ): + raise TypeError("portable DTensor regions require exact primitive fields") + logical_extent = LogicalRegion( + shard.logical_slice.offsets, + shard.logical_slice.lengths, + ) + else: + 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: + raise TypeError("portable logical slices require exact primitive fields") + logical_extent = LogicalSlice( + shard.logical_slice.flat_offset, + shard.logical_slice.length, + ) 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 + 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 @@ -358,7 +378,7 @@ def _validate_exact_shard_identity(shard) -> ShardIdentity: validated = ShardIdentity( parameter, shard.layout, - logical_slice, + logical_extent, placements=placements, process_group=process_group, local_member=shard.local_member, @@ -535,16 +555,21 @@ 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") + factored_v_2d = raw.get("factored_v_2d") policy = { "schema_version": 1, - "factored_v_2d": raw.get("factored_v_2d"), + "factored_v_2d": 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, + "second_moment_projection": ( + _SECOND_MOMENT_PROJECTION_FACTORED_TO_BLOCK + if implementation == _PLAIN_IMPLEMENTATION and factored_v_2d is True + else _SECOND_MOMENT_PROJECTION_EXACT + ), } return _normalize_policy(policy, implementation) @@ -556,13 +581,25 @@ def _normalized_ns_schedule(group): return [[float(a), float(b), float(c)] for a, b, c in schedule] -def _group_options(group, implementation: str, *, second_moment_policy=None): +def _group_options( + group, + implementation: str, + *, + second_moment_policy=None, + allow_rank_local_metadata=False, +): 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 allow_rank_local_metadata: + if implementation != _PLAIN_IMPLEMENTATION: + raise ValueError("rank-local group metadata is only valid for plain Gefen") + allowed = allowed | {_RANK_LOCAL_GROUP_METADATA_KEY} if not required.issubset(group) or not set(group).issubset(allowed): raise ValueError("portable parameter group contains missing or unknown keys") + if _RANK_LOCAL_GROUP_METADATA_KEY in group and type(group[_RANK_LOCAL_GROUP_METADATA_KEY]) is not dict: + raise TypeError("portable rank-local group metadata must be a plain dictionary") 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") @@ -608,7 +645,75 @@ def _group_options(group, implementation: str, *, second_moment_policy=None): ) +def _dtensor_local_storage(parameter) -> torch.Tensor: + if type(parameter) is not DTensor: + raise TypeError("portable DTensor state requires an exact torch.distributed.tensor.DTensor") + local = parameter.to_local() + if hasattr(local, "wait"): + local = local.wait() + if type(local) is not torch.Tensor: + raise TypeError("portable DTensor state requires exact local torch.Tensor storage") + if ( + local.layout is not torch.strided + or local.device.type not in {"cpu", "cuda"} + or local.dtype not in {torch.float16, torch.bfloat16, torch.float32, torch.float64} + or local.dtype.is_complex + or local.is_meta + or local.is_nested + or local.is_quantized + or getattr(local, "fake_mode", None) is not None + ): + raise TypeError("portable DTensor state requires supported materialized local storage") + return local + + +def _parameter_local_storage(parameter, shard: ShardIdentity) -> torch.Tensor: + if shard.layout is ParameterLayout.DTENSOR_1D_DEFAULT_WORLD: + return _dtensor_local_storage(parameter) + if isinstance(parameter, DTensor) or ( + isinstance(parameter, torch.Tensor) + and hasattr(parameter, "to_local") + and hasattr(parameter, "placements") + and hasattr(parameter, "device_mesh") + ): + raise TypeError("portable state rejects DTensor-like storage for a non-DTensor identity") + return parameter + + def _parameter_supported(parameter, shard: ShardIdentity) -> bool: + if shard.layout is ParameterLayout.DTENSOR_1D_DEFAULT_WORLD: + if type(parameter) is not DTensor: + return False + local = _dtensor_local_storage(parameter) + if ( + tuple(parameter.shape) != shard.parameter.global_shape + or parameter.dtype != local.dtype + or tuple(local.shape) != shard.logical_region.lengths + or local.numel() != shard.logical_region.numel + or len(tuple(parameter.placements)) != 1 + ): + return False + descriptor = shard.placements[0] + placement = tuple(parameter.placements)[0] + if descriptor.kind is PlacementKind.REPLICATE: + if type(placement) is not Replicate: + return False + elif ( + descriptor.kind is not PlacementKind.DIMENSION_SHARD + or type(placement) is not Shard + or type(placement.dim) is not int + or placement.dim != descriptor.parameter_dimension + ): + return False + mesh = parameter.device_mesh + coordinate = mesh.get_coordinate() + dim_names = tuple(mesh.mesh_dim_names or ()) + return ( + tuple(mesh.shape) == (descriptor.parts,) + and (not dim_names or dim_names == (descriptor.mesh_axis,)) + and coordinate is not None + and list(coordinate) == [descriptor.coordinate] + ) if type(parameter) not in {torch.Tensor, nn.Parameter} or not ( parameter.layout is torch.strided and parameter.device.type in {"cpu", "cuda"} @@ -681,6 +786,7 @@ def _validate_optimizer_shell( if shard.layout not in { ParameterLayout.REPLICATED, ParameterLayout.FLATTENED_ELEMENT_SHARD, + ParameterLayout.DTENSOR_1D_DEFAULT_WORLD, }: raise ValueError("plain portable state has an unsupported layout") elif shard.layout not in { @@ -727,16 +833,77 @@ def _live_maps(optimizer, binding: CheckpointProcessGroupBinding): 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") + _validate_rank_local_transport_schema(optimizer, binding) return local_by_fqn -def _parameter_storage_token(parameter): +def _validate_rank_local_transport_schema( + optimizer, + binding: CheckpointProcessGroupBinding, +) -> None: + parameters = [ + parameter + for group in optimizer.param_groups + for parameter in group["params"] + ] + + def transport_keys(state): + return tuple( + key + for key in state + if type(key) is str + and ( + key == _RANK_LOCAL_MEMBER_KEY + or key.startswith(_RANK_LOCAL_PAYLOAD_KEY_PREFIX) + ) + ) + + has_dtensor = any( + shard.layout is ParameterLayout.DTENSOR_1D_DEFAULT_WORLD + for _parameter, shard in optimizer._gefen_local_shard_bindings + ) + if not has_dtensor: + if any(transport_keys(optimizer.state[parameter]) for parameter in parameters): + raise ValueError("portable state contains unexpected rank-local transport keys") + return + if optimizer._gefen_rank_local_checkpoint_collective is not True or not parameters: + raise ValueError("portable DTensor state requires an installed rank-local transport schema") + + expected_payload_keys = { + "{}{}".format(_RANK_LOCAL_PAYLOAD_KEY_PREFIX, rank) + for rank in range(len(binding.identity.ordered_members)) + } + for parameter_index, parameter in enumerate(parameters): + state = optimizer.state[parameter] + keys = transport_keys(state) + if parameter_index == 0: + if set(keys) != expected_payload_keys: + raise ValueError("portable rank-local carrier keys do not match the default world") + for key in expected_payload_keys: + value = state[key] + if ( + type(value) is not torch.Tensor + or value.layout is not torch.strided + or value.device.type != "cpu" + or value.dtype != torch.uint8 + or tuple(value.shape) != (1,) + or value.requires_grad + or not value.is_contiguous() + or value.storage_offset() != 0 + or value.untyped_storage().nbytes() != 1 + or int(value.item()) != 0 + ): + raise ValueError("portable rank-local carrier payload has an invalid schema") + elif keys != (_RANK_LOCAL_MEMBER_KEY,) or state[_RANK_LOCAL_MEMBER_KEY] is not True: + raise ValueError("portable rank-local member marker has an invalid schema") + + +def _tensor_storage_metadata_token(parameter): try: version = parameter._version except RuntimeError: version = None return ( - id(parameter), version, str(parameter.device), str(parameter.dtype), @@ -751,6 +918,49 @@ def _parameter_storage_token(parameter): ) +def _ordinary_parameter_storage_token(parameter): + return (id(parameter), _tensor_storage_metadata_token(parameter)) + + +def _parameter_storage_token(parameter): + if type(parameter) is DTensor: + local = _dtensor_local_storage(parameter) + try: + version = parameter._version + except RuntimeError: + version = None + placements = [] + for placement in tuple(parameter.placements): + if type(placement) is Shard: + placements.append(("shard", placement.dim)) + elif type(placement) is Replicate: + placements.append(("replicate",)) + else: + placements.append((type(placement).__module__, type(placement).__qualname__, repr(placement))) + mesh = parameter.device_mesh + coordinate = mesh.get_coordinate() + return ( + "dtensor", + id(parameter), + version, + str(parameter.dtype), + tuple(parameter.shape), + tuple(placements), + tuple(mesh.shape), + tuple(mesh.mesh_dim_names or ()), + None if coordinate is None else tuple(coordinate), + _tensor_storage_metadata_token(local), + ) + if isinstance(parameter, DTensor) or ( + isinstance(parameter, torch.Tensor) + and hasattr(parameter, "to_local") + and hasattr(parameter, "placements") + and hasattr(parameter, "device_mesh") + ): + raise TypeError("portable state rejects DTensor-like tensor subclasses") + return _ordinary_parameter_storage_token(parameter) + + def _portable_value_token(value): if type(value) is torch.Tensor: hasher = hashlib.sha256() @@ -881,12 +1091,25 @@ def live_group_options(group): def _catalog_and_options(optimizer, implementation: str, policy): + allow_rank_local_metadata = ( + implementation == _PLAIN_IMPLEMENTATION + and optimizer._gefen_rank_local_checkpoint_collective is True + and any( + slot.shard.layout is ParameterLayout.DTENSOR_1D_DEFAULT_WORLD + for slot in optimizer._gefen_logical_slots + ) + ) 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( + group, + implementation, + second_moment_policy="block", + allow_rank_local_metadata=allow_rank_local_metadata, + ) group_options.append(group) else: group_options.append(_group_options(group, implementation)) @@ -898,14 +1121,19 @@ def _catalog_and_options(optimizer, implementation: str, policy): 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" + configured_factored = policy["factored_v_2d"] and len(identity.global_shape) == 2 + second = "factored" if configured_factored and slot.shard.layout is ParameterLayout.REPLICATED else "block" options = _group_options( group_options[slot.group_index], implementation, second_moment_policy=second, + allow_rank_local_metadata=allow_rank_local_metadata, ) - if second == "factored" and slot.shard.layout is not ParameterLayout.REPLICATED: - raise ValueError("factored portable state supports replicated matrices only") + if configured_factored and slot.shard.layout not in { + ParameterLayout.REPLICATED, + ParameterLayout.DTENSOR_1D_DEFAULT_WORLD, + }: + raise ValueError("factored portable state supports replicated matrices or the DTensor block fallback") else: options = group_options[slot.group_index] if slot.shard.layout is ParameterLayout.WHOLE_PARAMETER_OWNER and options["sharded_mode"] != "distributed": @@ -974,9 +1202,24 @@ 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): + def is_transport_key(key): + return type(key) is str and ( + key == _RANK_LOCAL_MEMBER_KEY + or key.startswith(_RANK_LOCAL_PAYLOAD_KEY_PREFIX) + ) + + if any( + key not in _AUTHORITATIVE_STATE_KEYS + and key not in _IGNORED_DERIVED_STATE_KEYS + and not is_transport_key(key) + 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} + core = { + key: value + for key, value in state.items() + if key not in _IGNORED_DERIVED_STATE_KEYS and not is_transport_key(key) + } if core.get("name") != compatibility_name: raise ValueError("portable parameter state name does not match its logical slot") return core @@ -1480,8 +1723,11 @@ def _stage_portable_import( 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") + _validate_portable_projection_policy( + document["policy"], + target_fragment["policy"], + implementation, + ) common = document["common"] if set(document["parameters"]) != set(target_fragment["catalog"]): raise ValueError("portable parameter catalog does not match the target") @@ -1499,6 +1745,7 @@ def _stage_portable_import( implementation=implementation, global_step=common["gefen_global_step"], codebook=common["gefen_codebook"], + source_second_moment_projection=document["policy"]["second_moment_projection"], target_algorithm_options=target_options, target_second_moment=target_second, ) @@ -1521,7 +1768,30 @@ def _stage_portable_import( staging_owner = object.__new__(type(optimizer)) staging_owner.__dict__ = optimizer.__dict__.copy() staging_owner._deterministic = common["gefen_deterministic"] + has_dtensor_target = any( + shard.layout is ParameterLayout.DTENSOR_1D_DEFAULT_WORLD + for _parameter, shard in local_by_fqn.values() + ) + if has_dtensor_target: + # The native rank-local checkpoint envelope deliberately accepts only + # exact-topology restores. Portable state has already projected the + # canonical global record to this target topology, so stage that local + # result through the native semantic validator without presenting it as + # a rank-local checkpoint. These method shadows exist only on the + # isolated staging owner and are removed before publication. + native.pop("gefen_native_local_shards", None) + for group in native["param_groups"]: + metadata = dict(group[_RANK_LOCAL_GROUP_METADATA_KEY]) + metadata.pop("native_local_shards", None) + group[_RANK_LOCAL_GROUP_METADATA_KEY] = metadata + staging_owner._uses_rank_local_sharded_state = lambda: False + staging_owner._serialized_native_local_shards = lambda: None staged = staging_owner._stage_load_state_dict(native) + if has_dtensor_target: + staged.__dict__.pop("_uses_rank_local_sharded_state", None) + staged.__dict__.pop("_serialized_native_local_shards", None) + staged._install_rank_local_checkpoint_schema() + _validate_rank_local_transport_schema(staged, binding) optimizer._preserve_canonical_target_configuration(staged) if ( type(staged.__dict__) is not dict @@ -1669,6 +1939,7 @@ def _validate_readiness_slot( return if parameter is None: raise ValueError("a portable payload shard requires local storage") + parameter_storage = _parameter_local_storage(parameter, shard) core = _parameter_state_core(optimizer, parameter, slot.compatibility_name) keys = frozenset(core) @@ -1708,7 +1979,7 @@ def _validate_readiness_slot( ) if indices.device != magnitudes.device: raise ValueError("momentum indices and magnitudes must share a device") - if indices.device != parameter.device: + if indices.device != parameter_storage.device: raise ValueError("portable parameter state must share its parameter device") if implementation == _PLAIN_IMPLEMENTATION: @@ -1735,7 +2006,7 @@ def _validate_readiness_slot( nonnegative=True, validate_values=validate_values, ) - if v_row.device != parameter.device or v_col.device != parameter.device: + if v_row.device != parameter_storage.device or v_col.device != parameter_storage.device: raise ValueError("factored portable state must share its parameter device") return if keys != _PLAIN_BLOCK_KEYS: @@ -1751,7 +2022,7 @@ def _validate_readiness_slot( nonnegative=True, validate_values=validate_values, ) - if vmean.device != parameter.device: + if vmean.device != parameter_storage.device: raise ValueError("block portable state must share its parameter device") return @@ -1770,7 +2041,7 @@ def _validate_readiness_slot( nonnegative=True, validate_values=validate_values, ) - if normuon_v.device != parameter.device: + if normuon_v.device != parameter_storage.device: raise ValueError("NorMuon portable state must share its parameter device") diff --git a/src/gefen/portable_state.py b/src/gefen/portable_state.py index e28f722..0dda55a 100644 --- a/src/gefen/portable_state.py +++ b/src/gefen/portable_state.py @@ -14,7 +14,12 @@ ShardIdentity, ShardingManifest, ) -from gefen.portable import _recompress_dense_momentum, _reduce_block_second_moment, _validate_codebook +from gefen.portable import ( + _expand_factored_second_moment_live_fp32_v1, + _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, @@ -40,7 +45,14 @@ _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" +_SECOND_MOMENT_PROJECTION_EXACT = "exact_representation_target_period_one_v1" +_SECOND_MOMENT_PROJECTION_FACTORED_TO_BLOCK = "defined_projection_factored_to_block_live_fp32_target_period_one_v1" +_SECOND_MOMENT_PROJECTIONS = frozenset( + { + _SECOND_MOMENT_PROJECTION_EXACT, + _SECOND_MOMENT_PROJECTION_FACTORED_TO_BLOCK, + } +) _MAX_EXACT_COUNTER = (1 << 53) - 1 _FRAGMENT_KEYS = frozenset( @@ -283,8 +295,13 @@ def _normalize_policy(value, implementation: str): 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: + if value["second_moment_projection"] not in _SECOND_MOMENT_PROJECTIONS: raise ValueError("unsupported portable second-moment projection") + if ( + implementation == "gefen.GefenMuon" + and value["second_moment_projection"] != _SECOND_MOMENT_PROJECTION_EXACT + ): + raise ValueError("Muon portable policy requires exact 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)) @@ -296,6 +313,26 @@ def _normalize_policy(value, implementation: str): return {**value, "period_one_substrings": list(substrings)} +def _validate_portable_projection_policy(source_policy, target_policy, implementation: str): + """Validate exact policy compatibility plus the defined one-way migration.""" + + source = _normalize_policy(source_policy, implementation) + target = _normalize_policy(target_policy, implementation) + invariant_keys = _POLICY_KEYS - {"factored_v_2d", "second_moment_projection"} + if any(not _values_equal(source[key], target[key]) for key in invariant_keys): + raise ValueError("portable document policy does not match the target") + source_factored = source["factored_v_2d"] + target_factored = target["factored_v_2d"] + if source_factored == target_factored: + return + if implementation != "gefen.Gefen": + raise ValueError("portable document policy does not match the target") + if not source_factored and target_factored: + raise ValueError("portable block-to-factored second-moment migration is unsupported") + if source["second_moment_projection"] != _SECOND_MOMENT_PROJECTION_FACTORED_TO_BLOCK: + raise ValueError("portable source policy does not authorize factored-to-block second-moment projection") + + 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") @@ -379,8 +416,9 @@ def _normalize_complete_parameter_record(fqn, value, implementation: str, global 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: + configured_factored = policy["factored_v_2d"] and len(shape) == 2 + representation = options["second_moment_policy"] + if representation == "factored" and not configured_factored: raise ValueError("parameter second_moment_policy conflicts with the optimizer policy") if variant == "pristine": expected_keys = frozenset() @@ -392,11 +430,11 @@ def _normalize_complete_parameter_record(fqn, value, implementation: str, global 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": + if not periods or hints["source_second_moment"] != "block" or 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": + if not periods or hints["source_second_moment"] != "factored" or representation != "factored": raise ValueError("factored state conflicts with its policy or projection hints") else: raise ValueError("unsupported plain portable state_variant") @@ -544,6 +582,8 @@ def _derived_role(shard: ShardIdentity) -> str: 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" + if shard.layout is ParameterLayout.DTENSOR_1D_DEFAULT_WORLD: + return "live" if shard.logical_region.numel else "empty_dtensor" raise ValueError("unsupported portable fragment layout") @@ -555,6 +595,8 @@ def _local_dense_shape(shard: ShardIdentity): return shard.parameter.global_shape if shard.layout is ParameterLayout.FLATTENED_ELEMENT_SHARD: return (shard.logical_slice.length,) + if shard.layout is ParameterLayout.DTENSOR_1D_DEFAULT_WORLD: + return shard.logical_region.lengths raise ValueError("unsupported portable fragment layout") @@ -573,7 +615,17 @@ def _normalize_catalog(value, implementation: str, manifest: ShardingManifest, p 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" + configured_factored = policy["factored_v_2d"] and len(identity.global_shape) == 2 + parameter_shards = manifest.for_parameter(fqn) + layouts = {shard.layout for shard in parameter_shards} + if configured_factored and layouts == {ParameterLayout.REPLICATED}: + expected = "factored" + elif configured_factored and layouts == {ParameterLayout.DTENSOR_1D_DEFAULT_WORLD}: + expected = "block" + elif configured_factored: + raise ValueError("factored logical matrices require replicated storage or the DTensor block fallback") + else: + expected = "block" if options["second_moment_policy"] != expected: raise ValueError("catalog second_moment_policy conflicts with the optimizer policy") if expected == "factored" and any( @@ -599,7 +651,11 @@ def _normalize_fragment_slot(value, implementation: str, common, catalog, member shard_record = _normalize_shard_identity(value["shard"]) shard = _parse_shard_identity(shard_record) allowed_layouts = ( - {ParameterLayout.REPLICATED, ParameterLayout.FLATTENED_ELEMENT_SHARD} + { + ParameterLayout.REPLICATED, + ParameterLayout.FLATTENED_ELEMENT_SHARD, + ParameterLayout.DTENSOR_1D_DEFAULT_WORLD, + } if implementation == "gefen.Gefen" else {ParameterLayout.REPLICATED, ParameterLayout.WHOLE_PARAMETER_OWNER} ) @@ -748,7 +804,11 @@ def _normalize_portable_state_fragment(fragment, *, limits): manifest_record = _normalize_sharding_manifest(fragment["manifest"]) manifest = _parse_sharding_manifest(manifest_record) allowed_layouts = ( - {ParameterLayout.REPLICATED, ParameterLayout.FLATTENED_ELEMENT_SHARD} + { + ParameterLayout.REPLICATED, + ParameterLayout.FLATTENED_ELEMENT_SHARD, + ParameterLayout.DTENSOR_1D_DEFAULT_WORLD, + } if implementation == "gefen.Gefen" else {ParameterLayout.REPLICATED, ParameterLayout.WHOLE_PARAMETER_OWNER} ) @@ -926,10 +986,10 @@ def _assemble_parameter(fqn, slots, manifest: ShardingManifest, catalog_entry, i } ) 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": + representation = catalog_entry["algorithm_options"]["second_moment_policy"] + if variant == "initialized_dense" and representation != "block": raise ValueError("block portable state conflicts with factored_v_2d policy") - if variant == "initialized_factored" and expected != "factored": + if variant == "initialized_factored" and representation != "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: @@ -1042,6 +1102,7 @@ def _project_portable_parameter_state( implementation, global_step, codebook, + source_second_moment_projection, target_algorithm_options, target_second_moment=None, ): @@ -1054,7 +1115,11 @@ def _project_portable_parameter_state( if not isinstance(target_shard, ShardIdentity): raise TypeError("target_shard must be a ShardIdentity") target_layouts = ( - {ParameterLayout.REPLICATED, ParameterLayout.FLATTENED_ELEMENT_SHARD} + { + ParameterLayout.REPLICATED, + ParameterLayout.FLATTENED_ELEMENT_SHARD, + ParameterLayout.DTENSOR_1D_DEFAULT_WORLD, + } if implementation == "gefen.Gefen" else {ParameterLayout.REPLICATED, ParameterLayout.WHOLE_PARAMETER_OWNER} ) @@ -1077,30 +1142,40 @@ def _project_portable_parameter_state( "codebook_refresh_every": 0, "stochastic_round": False, "momentum_projection": _MOMENTUM_PROJECTION, - "second_moment_projection": _SECOND_MOMENT_PROJECTION, + "second_moment_projection": source_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") + source_options = record["algorithm_options"] + invariant_option_keys = _PLAIN_OPTION_KEYS - {"second_moment_policy"} + if any(not _values_equal(source_options[key], target_options[key]) for key in invariant_option_keys): + raise ValueError("portable algorithm options do not match the target") + source_representation = source_options["second_moment_policy"] + if source_representation == "block" and target_second_moment == "factored": + raise ValueError("portable block-to-factored second-moment migration is unsupported") + if ( + source_representation == "factored" + and target_second_moment == "block" + and source_second_moment_projection != _SECOND_MOMENT_PROJECTION_FACTORED_TO_BLOCK + ): + raise ValueError("portable source policy does not authorize factored-to-block second-moment projection") 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 not _values_equal(record["algorithm_options"], target_options): + raise ValueError("portable algorithm options do not match the target") if target_second_moment is not None: raise ValueError("Muon projection does not accept target_second_moment") if ( @@ -1147,15 +1222,37 @@ def _project_portable_parameter_state( } ) 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"], - } - ) + second_step = record["state"]["factored_step"] + if target_second_moment == "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": second_step, + } + ) + else: + dense_second = _expand_factored_second_moment_live_fp32_v1( + record["state"]["v_row"], + record["state"]["v_col"], + logical_shape=target_shard.parameter.global_shape, + step=second_step, + ) + local_second = _project_dense_logical_field( + target_shard.parameter, + dense_second, + target_shard, + ) + if local_second is None: + return {} + result.update( + { + "vmean": _reduce_block_second_moment(local_second, period=1, step=second_step), + "vmean_step": second_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") diff --git a/tests/test_capturable_fsdp2.py b/tests/test_capturable_fsdp2.py index 444ea05..78b9640 100644 --- a/tests/test_capturable_fsdp2.py +++ b/tests/test_capturable_fsdp2.py @@ -362,6 +362,156 @@ def _worker(rank, world, case_name, port, q): dist.destroy_process_group() +def _frozen_routing_worker(rank, world, port, q): + import torch.distributed as dist + from torch.distributed.tensor import init_device_mesh + + os.environ["MASTER_ADDR"] = "127.0.0.1" + os.environ["MASTER_PORT"] = port + os.environ["RANK"] = str(rank) + os.environ["WORLD_SIZE"] = str(world) + torch.cuda.set_device(rank) + device = torch.device("cuda", rank) + dist.init_process_group( + "nccl", rank=rank, world_size=world, timeout=timedelta(seconds=180) + ) + try: + mesh = init_device_mesh("cuda", (world,)) + case = { + "kind": "muon", + "sharded_mode": "approx", + "specs": [("route_w", (32, 32))], + "dtype": torch.bfloat16, + } + full_init = _make_full_tensors( + case["specs"], + seed=8100, + device=device, + dtype=case["dtype"], + scale=0.1, + ) + full_grad = _make_full_tensors( + case["specs"], + seed=8200, + device=device, + dtype=case["dtype"], + scale=0.01, + ) + params = _build_params(case["specs"], full_init, mesh) + lr = torch.tensor(1e-1, device=device) + optimizer = _build_optimizer(case, params, lr) + static_grads = _attach_static_grads(params, full_grad, mesh) + _fill_static_grads(static_grads, full_grad, rank, world) + + side = torch.cuda.Stream(device=rank) + side.wait_stream(torch.cuda.current_stream(rank)) + with torch.cuda.stream(side): + for _ in range(WARMUP): + optimizer.step() + torch.cuda.current_stream(rank).wait_stream(side) + torch.cuda.synchronize(rank) + parameter = params[0][1] + before = parameter.to_local().detach().clone() + warm_mode = optimizer.param_groups[0]["sharded_mode"] + warm_state_local = ( + optimizer.state[parameter]["m_codebook"].numel() + == parameter.to_local().numel() + ) + + if rank == 0: + optimizer.param_groups[0]["sharded_mode"] = "exact" + dist.barrier() + asymmetric_graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(asymmetric_graph): + optimizer.step() + torch.cuda.synchronize(rank) + asymmetric_restored = optimizer.param_groups[0]["sharded_mode"] == "approx" + dist.barrier() + asymmetric_graph.replay() + torch.cuda.synchronize(rank) + dist.barrier() + + optimizer.param_groups[0]["sharded_mode"] = "distributed" + unanimous_graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(unanimous_graph): + optimizer.step() + torch.cuda.synchronize(rank) + unanimous_restored = optimizer.param_groups[0]["sharded_mode"] == "approx" + dist.barrier() + unanimous_graph.replay() + torch.cuda.synchronize(rank) + dist.barrier() + + after = parameter.to_local().detach().clone() + max_parameter_change = (after.float() - before.float()).abs().max().item() + # Recording does not execute the captured step; only the two explicit + # replays advance device-side counters beyond eager warmup. + counters_ok, counter_message = _assert_counters(optimizer, WARMUP + 2) + q.put( + ( + "frozen-routing", + rank, + { + "warm_mode": warm_mode, + "warm_state_local": warm_state_local, + "asymmetric_restored": asymmetric_restored, + "unanimous_restored": unanimous_restored, + "parameter_changed": max_parameter_change > 0.0, + "max_parameter_change": max_parameter_change, + "counters_ok": counters_ok, + "counter_message": counter_message, + }, + ) + ) + except Exception: + q.put(("error", rank, traceback.format_exc())) + finally: + if dist.is_initialized(): + dist.destroy_process_group() + + +def _run_frozen_routing_case(): + import torch.multiprocessing as mp + + ctx = mp.get_context("spawn") + q = ctx.Queue() + port = _get_free_port() + procs = [ + ctx.Process(target=_frozen_routing_worker, args=(rank, WORLD, port, q)) + for rank in range(WORLD) + ] + for proc in procs: + proc.start() + + messages = [] + deadline = time.monotonic() + 240 + while time.monotonic() < deadline and len(messages) < WORLD: + try: + messages.append(q.get(timeout=0.5)) + except queue.Empty: + if not any(proc.is_alive() for proc in procs): + break + for proc in procs: + proc.join(timeout=10) + for proc in procs: + if proc.is_alive(): + proc.terminate() + proc.join(timeout=10) + + assert all(proc.exitcode == 0 for proc in procs), [ + proc.exitcode for proc in procs + ] + errors = [message[2] for message in messages if message[0] == "error"] + assert not errors, "\n".join(errors) + results = { + message[1]: message[2] + for message in messages + if message[0] == "frozen-routing" + } + assert set(results) == set(range(WORLD)), messages + return results + + def _run_case(case_name): import torch.multiprocessing as mp @@ -468,6 +618,23 @@ def run(): def test_capturable_fsdp2_capture(case_name): assert _run_case(case_name) + @pytest.mark.skipif( + not torch.cuda.is_available() + or not torch.distributed.is_available() + or not torch.distributed.is_nccl_available() + or torch.cuda.device_count() < WORLD, + reason="capturable FSDP2 capture needs >=2 CUDA GPUs and NCCL", + ) + def test_capturable_muon_freezes_warmup_routing_before_mode_partition(): + results = _run_frozen_routing_case() + for result in results.values(): + assert result["warm_mode"] == "approx", result + assert result["warm_state_local"], result + assert result["asymmetric_restored"], result + assert result["unanimous_restored"], result + assert result["counters_ok"], result["counter_message"] + assert result["parameter_changed"], result["max_parameter_change"] + if __name__ == "__main__": ok = run() diff --git a/tests/test_checkpoint_projection_contracts.py b/tests/test_checkpoint_projection_contracts.py new file mode 100644 index 0000000..0306542 --- /dev/null +++ b/tests/test_checkpoint_projection_contracts.py @@ -0,0 +1,237 @@ +"""Focused capability coverage for directional checkpoint-state projections.""" + +from dataclasses import replace + +import pytest +import torch + +import gefen +from gefen.codebook import CodebookProcessGroupBinding +from gefen.contracts import ( + CHECKPOINT_STATE_TRANSITION_SCHEMA_VERSION, + CheckpointProjectionQualifier, + CheckpointStateRepresentation, + CheckpointStateTransition, + CheckpointStateTransitionKind, + CheckpointSupport, + CheckpointTransport, + LogicalSlice, + ParameterIdentity, + ParameterLayout, + PlacementKind, + ProcessGroupIdentity, + ProcessGroupScope, + ShardIdentity, + ShardPlacement, + ShardingManifest, + TopologyChange, +) +from gefen.gefen import Gefen +from gefen.portable_state import _SECOND_MOMENT_PROJECTION_FACTORED_TO_BLOCK +from gefen.rebinding import ParameterRebinding + + +_MEMBER = "rank:0" + + +def _finalized_plain(*, factored, include_block_bias=False): + weight = torch.nn.Parameter( + torch.arange(1, 7, dtype=torch.float32).reshape(2, 3) + ) + named_parameters = [("layer.weight", weight)] + if include_block_bias: + named_parameters.append( + ("layer.bias", torch.nn.Parameter(torch.arange(3, dtype=torch.float32))) + ) + optimizer = Gefen( + named_parameters, + fused=False, + factored_v_2d=factored, + force_2d_period_one=True, + ) + group = ProcessGroupIdentity("checkpoint", (_MEMBER,)) + shards = tuple( + ShardIdentity( + identity, + ParameterLayout.REPLICATED, + LogicalSlice.full(identity), + placements=( + ShardPlacement("checkpoint", PlacementKind.REPLICATE, 0, 1), + ), + process_group=group, + local_member=_MEMBER, + ) + for identity in ( + ParameterIdentity(name, tuple(parameter.shape)) + for name, parameter in named_parameters + ) + ) + optimizer.post_sharding( + tuple( + ParameterRebinding(parameter, parameter, shard) + for (_, parameter), shard in zip(named_parameters, shards) + ), + manifest=ShardingManifest(shards), + codebook_process_group=CodebookProcessGroupBinding( + group, _MEMBER, None, torch.device("cpu") + ), + ) + support = next( + item + for item in optimizer.optimizer_contract().capabilities.checkpoints + if item.transport is CheckpointTransport.CANONICAL_GLOBAL + ) + return support + + +def _transition(support, source, target): + matches = tuple( + transition + for transition in support.state_transitions + if transition.source is source and transition.target is target + ) + assert len(matches) == 1 + return matches[0] + + +def test_factored_source_separates_exact_restore_from_versioned_block_projection(): + support = _finalized_plain(factored=True) + factored = CheckpointStateRepresentation.FACTORED_SECOND_MOMENT + block = CheckpointStateRepresentation.BLOCK_SECOND_MOMENT + + assert support.same_topology == frozenset({ParameterLayout.REPLICATED}) + assert not support.topology_changing + exact = _transition(support, factored, factored) + assert exact.kind is CheckpointStateTransitionKind.EXACT + assert exact.same_topology == support.same_topology + assert not exact.topology_changing + assert exact.qualifier is None + + projection = _transition(support, factored, block) + assert projection.kind is CheckpointStateTransitionKind.DEFINED_PROJECTION + assert ( + projection.qualifier + is CheckpointProjectionQualifier.FACTORED_TO_BLOCK_LIVE_FP32_TARGET_PERIOD_ONE_V1 + ) + assert ( + projection.qualifier.value + == _SECOND_MOMENT_PROJECTION_FACTORED_TO_BLOCK + ) + assert projection.schema_version == CHECKPOINT_STATE_TRANSITION_SCHEMA_VERSION + assert projection.same_topology == frozenset({ParameterLayout.REPLICATED}) + assert projection.topology_changing == frozenset( + { + ParameterLayout.REPLICATED, + ParameterLayout.FLATTENED_ELEMENT_SHARD, + ParameterLayout.DTENSOR_1D_DEFAULT_WORLD, + } + ) + assert projection.target_layouts == frozenset( + { + ParameterLayout.REPLICATED, + ParameterLayout.FLATTENED_ELEMENT_SHARD, + ParameterLayout.DTENSOR_1D_DEFAULT_WORLD, + } + ) + assert projection.topology_change_kinds == frozenset( + {TopologyChange.PLACEMENT_RESHARD} + ) + + +def test_mixed_replicated_factored_weight_and_block_bias_declare_both_exact_paths(): + support = _finalized_plain(factored=True, include_block_bias=True) + factored = CheckpointStateRepresentation.FACTORED_SECOND_MOMENT + block = CheckpointStateRepresentation.BLOCK_SECOND_MOMENT + + replicated = frozenset({ParameterLayout.REPLICATED}) + assert support.same_topology == replicated + factored_exact = _transition(support, factored, factored) + block_exact = _transition(support, block, block) + assert factored_exact.kind is CheckpointStateTransitionKind.EXACT + assert factored_exact.same_topology == replicated + assert block_exact.kind is CheckpointStateTransitionKind.EXACT + assert block_exact.same_topology == replicated + projection = _transition(support, factored, block) + assert ( + projection.qualifier + is CheckpointProjectionQualifier.FACTORED_TO_BLOCK_LIVE_FP32_TARGET_PERIOD_ONE_V1 + ) + + +def test_block_source_advertises_only_exact_block_restore(): + support = _finalized_plain(factored=False) + block = CheckpointStateRepresentation.BLOCK_SECOND_MOMENT + factored = CheckpointStateRepresentation.FACTORED_SECOND_MOMENT + + assert support.same_topology == frozenset({ParameterLayout.REPLICATED}) + assert support.topology_changing == frozenset( + { + ParameterLayout.REPLICATED, + ParameterLayout.FLATTENED_ELEMENT_SHARD, + ParameterLayout.DTENSOR_1D_DEFAULT_WORLD, + } + ) + exact = _transition(support, block, block) + assert exact.kind is CheckpointStateTransitionKind.EXACT + assert exact.same_topology == support.same_topology + assert exact.topology_changing == support.topology_changing + assert exact.topology_change_kinds == support.topology_change_kinds + assert all( + transition.target is not factored + for transition in support.state_transitions + ) + + +def test_transition_schema_is_additive_immutable_and_validates_direction(): + legacy = CheckpointSupport( + CheckpointTransport.NATIVE_OPTIMIZER, + frozenset({ParameterLayout.REPLICATED}), + frozenset(), + ProcessGroupScope.NONE, + ) + assert legacy.state_transitions == () + assert gefen.CheckpointStateTransition is CheckpointStateTransition + assert ( + gefen.CHECKPOINT_STATE_TRANSITION_SCHEMA_VERSION + == CHECKPOINT_STATE_TRANSITION_SCHEMA_VERSION + ) + + factored = CheckpointStateRepresentation.FACTORED_SECOND_MOMENT + block = CheckpointStateRepresentation.BLOCK_SECOND_MOMENT + source_layouts = {ParameterLayout.REPLICATED} + projection = CheckpointStateTransition( + factored, + block, + CheckpointStateTransitionKind.DEFINED_PROJECTION, + source_layouts, + set(), + qualifier=( + CheckpointProjectionQualifier.FACTORED_TO_BLOCK_LIVE_FP32_TARGET_PERIOD_ONE_V1 + ), + ) + source_layouts.clear() + assert projection.same_topology == frozenset({ParameterLayout.REPLICATED}) + with pytest.raises(TypeError, match="require a CheckpointProjectionQualifier"): + replace(projection, qualifier=None) + with pytest.raises(ValueError, match="declared direction and kind"): + replace(projection, source=block, target=factored) + with pytest.raises(ValueError, match="unsupported checkpoint state-transition schema"): + replace(projection, schema_version=CHECKPOINT_STATE_TRANSITION_SCHEMA_VERSION + 1) + + +def test_exact_transition_must_cover_legacy_exact_layout_claims(): + exact = CheckpointStateTransition( + CheckpointStateRepresentation.BLOCK_SECOND_MOMENT, + CheckpointStateRepresentation.BLOCK_SECOND_MOMENT, + CheckpointStateTransitionKind.EXACT, + frozenset({ParameterLayout.REPLICATED}), + frozenset(), + ) + with pytest.raises(ValueError, match="exact state transitions must cover"): + CheckpointSupport( + CheckpointTransport.CANONICAL_GLOBAL, + frozenset({ParameterLayout.FLATTENED_ELEMENT_SHARD}), + frozenset(), + ProcessGroupScope.ADAPTER_DEFINED, + state_transitions=(exact,), + ) diff --git a/tests/test_codebook_scope_cpu.py b/tests/test_codebook_scope_cpu.py index b427d4a..4ca07a9 100644 --- a/tests/test_codebook_scope_cpu.py +++ b/tests/test_codebook_scope_cpu.py @@ -480,6 +480,21 @@ def test_native_v2_guard_retains_pruned_logical_positions_and_rejects_corruption "_validate_codebook_runtime_binding", lambda self, binding: None, ) + monkeypatch.setattr( + GefenMuon, + "_preflight_muon_checkpoint_routing", + lambda self: None, + ) + + def synchronize_checkpoint_failure_locally(_self, error, _phase): + if error is not None: + raise error + + monkeypatch.setattr( + GefenMuon, + "_synchronize_muon_checkpoint_failure", + synchronize_checkpoint_failure_locally, + ) group = ProcessGroupIdentity("data_parallel", ("rank:0", "rank:1")) def build(): diff --git a/tests/test_dtensor_rebinding.py b/tests/test_dtensor_rebinding.py new file mode 100644 index 0000000..3c6d0a9 --- /dev/null +++ b/tests/test_dtensor_rebinding.py @@ -0,0 +1,2002 @@ +"""Two-rank CPU coverage for Gefen-family 1-D DTensor rebinding.""" + +from __future__ import annotations + +import copy +from datetime import timedelta +import multiprocessing as mp +import os +import queue as queue_module +import shutil +import tempfile +import traceback + +import pytest +import torch +import torch.distributed as dist + + +_WORLD = 2 +_MEMBERS = ("rank:0", "rank:1") +_STATE_KEYS = { + "name", + "automatic_period", + "step", + "m_codebook", + "m_magnitude", + "vmean", + "vmean_step", + "v_row", + "v_col", + "factored_step", +} + + +def _local(value): + if hasattr(value, "to_local"): + value = value.to_local() + if hasattr(value, "wait"): + value = value.wait() + return value + + +def _raw_bytes(value): + value = _local(value).detach().contiguous() + return value.reshape(-1).view(torch.uint8).clone() + + +def _tensor_unchanged(value, snapshot): + local = _local(value) + return ( + local.dtype == snapshot[0] + and local.layout == snapshot[1] + and tuple(local.shape) == snapshot[2] + and torch.equal(_raw_bytes(local), snapshot[3]) + ) + + +def _tensor_snapshot(value): + local = _local(value) + return local.dtype, local.layout, tuple(local.shape), _raw_bytes(local) + + +def _assert_value_exact(actual, expected, path="value"): + if torch.is_tensor(expected): + assert torch.is_tensor(actual), path + actual = _local(actual) + expected = _local(expected) + assert actual.dtype == expected.dtype, path + assert actual.layout == expected.layout, path + assert tuple(actual.shape) == tuple(expected.shape), path + assert torch.equal(_raw_bytes(actual), _raw_bytes(expected)), path + return + assert type(actual) is type(expected), path + if isinstance(expected, dict): + assert set(actual) == set(expected), path + for key in expected: + _assert_value_exact(actual[key], expected[key], "{}.{}".format(path, key)) + elif isinstance(expected, (tuple, list)): + assert len(actual) == len(expected), path + for index, (actual_item, expected_item) in enumerate(zip(actual, expected)): + _assert_value_exact(actual_item, expected_item, "{}[{}]".format(path, index)) + else: + assert actual == expected, path + + +def _persistent_snapshot(optimizer, parameter): + return { + "global_step": optimizer._gefen_global_step, + "codebook": None if optimizer._gefen_codebook is None else optimizer._gefen_codebook.detach().clone(), + "state": { + key: value.detach().clone() if torch.is_tensor(value) else copy.deepcopy(value) + for key, value in optimizer.state[parameter].items() + if key in _STATE_KEYS + }, + } + + +def _assert_optimizer_exact(actual_optimizer, actual_parameter, expected_optimizer, expected_parameter, path): + _assert_value_exact( + _persistent_snapshot(actual_optimizer, actual_parameter), + _persistent_snapshot(expected_optimizer, expected_parameter), + path, + ) + + +def _region(shape, coordinate, parts, dimension): + from gefen import LogicalRegion + + global_length = shape[dimension] + chunk = (global_length + parts - 1) // parts + offset = min(coordinate * chunk, global_length) + length = max(0, min(global_length, offset + chunk) - offset) + offsets = [0] * len(shape) + lengths = list(shape) + offsets[dimension] = offset + lengths[dimension] = length + return LogicalRegion(tuple(offsets), tuple(lengths)) + + +def _manifest( + fqn, + shape, + *, + kind, + dimension=None, + members=_MEMBERS, + semantic_name="data_parallel", + layout=None, +): + from gefen import ( + LogicalRegion, + LogicalSlice, + ParameterIdentity, + ParameterLayout, + PlacementKind, + ProcessGroupIdentity, + ShardIdentity, + ShardPlacement, + ShardingManifest, + ) + + identity = ParameterIdentity(fqn, shape) + group = ProcessGroupIdentity(semantic_name, members) + if layout is None: + layout = ParameterLayout.DTENSOR_1D_DEFAULT_WORLD + shards = [] + for coordinate, member in enumerate(members): + if layout is ParameterLayout.REPLICATED: + logical_slice = LogicalSlice.full(identity) + elif kind is PlacementKind.REPLICATE: + logical_slice = LogicalRegion.full(identity) + else: + logical_slice = _region(shape, coordinate, len(members), dimension) + shards.append( + ShardIdentity( + identity, + layout, + logical_slice, + placements=( + ShardPlacement( + "dp", + kind, + coordinate, + len(members), + parameter_dimension=dimension, + ), + ), + process_group=group, + local_member=member, + ) + ) + return ShardingManifest(tuple(shards)), tuple(shards) + + +def _combined_manifest(*manifests): + from gefen import ShardingManifest + + return ShardingManifest(tuple(shard for manifest in manifests for shard in manifest.shards)) + + +def _optimizer(parameter, name="legacy.weight"): + from gefen import Gefen + + return Gefen( + [(name, parameter)], + lr=3e-3, + betas=(0.8, 0.97), + eps=2e-8, + weight_decay=0.02, + fused=False, + deterministic=True, + force_2d_period_one=True, + factored_v_2d=False, + ) + + +def _two_parameter_optimizer(first, second): + from gefen import Gefen + + return Gefen( + [("legacy.first", first), ("legacy.second", second)], + lr=3e-3, + betas=(0.8, 0.97), + eps=2e-8, + weight_decay=0.02, + fused=False, + deterministic=True, + force_2d_period_one=True, + factored_v_2d=False, + ) + + +def _bind(optimizer, old_parameter, target, shard, manifest): + from gefen import ParameterRebinding + + optimizer.post_sharding( + (ParameterRebinding(old_parameter, target, shard),), + manifest=manifest, + ) + + +def _model(parameter): + from torch import nn + + model = nn.Module() + model.register_parameter("weight", parameter) + return model + + +def _local_slice(value, region): + slices = tuple(slice(offset, offset + length) for offset, length in zip(region.offsets, region.lengths)) + return value[slices].clone() + + +def _placement(kind, dimension): + from torch.distributed.tensor import Replicate, Shard + + return Replicate() if kind == "replicate" else Shard(dimension) + + +def _placement_kind(kind): + from gefen import PlacementKind + + return PlacementKind.REPLICATE if kind == "replicate" else PlacementKind.DIMENSION_SHARD + + +def _contract_is_exact(optimizer): + from gefen import CheckpointTransport, ParameterLayout, ProcessGroupScope + + contract = optimizer.optimizer_contract() + rank_local = [ + support + for support in contract.capabilities.checkpoints + if support.transport is CheckpointTransport.PYTORCH_RANK_LOCAL + ] + native_dtensor = [ + support + for support in contract.capabilities.checkpoints + if support.transport is CheckpointTransport.NATIVE_OPTIMIZER + and ParameterLayout.DTENSOR_1D_DEFAULT_WORLD in support.same_topology + ] + return ( + contract.capabilities.canonical_parameter_fqns + and contract.capabilities.stable_shard_identity + and contract.capabilities.shard_rebinding + and contract.capabilities.post_sharding + and optimizer.codebook_process_group_binding() is None + and len(rank_local) == 1 + and rank_local[0].same_topology == frozenset({ParameterLayout.DTENSOR_1D_DEFAULT_WORLD}) + and rank_local[0].topology_changing == frozenset() + and rank_local[0].process_group_scope is ProcessGroupScope.DEFAULT_WORLD + and rank_local[0].requires_collective + and rank_local[0].atomic_load + and not native_dtensor + ) + + +def _run_valid_case(rank, mesh, *, shape, kind, dimension): + from torch import nn + from torch.distributed.checkpoint.state_dict import ( + StateDictOptions, + get_optimizer_state_dict, + set_optimizer_state_dict, + ) + from torch.distributed.tensor import distribute_tensor + + placement = _placement(kind, dimension) + placement_kind = _placement_kind(kind) + full = torch.linspace(-0.75, 0.5, torch.tensor(shape).prod().item(), dtype=torch.float32).reshape(shape) + first_gradient = torch.linspace(0.6, -0.4, full.numel(), dtype=torch.float32).reshape(shape) + second_gradient = torch.arange(full.numel(), dtype=torch.float32).reshape(shape).cos() + manifest, shards = _manifest( + "Model.Weight", + shape, + kind=placement_kind, + dimension=dimension, + ) + shard = shards[rank] + + old = nn.Parameter(full.clone()) + target = nn.Parameter(distribute_tensor(full.clone(), mesh, [placement])) + optimizer = _optimizer(old) + _bind(optimizer, old, target, shard, manifest) + model = _model(target) + assert optimizer.param_groups[0]["params"] == [target] + assert old not in optimizer.state + assert optimizer.shard_identity(target) == shard + assert optimizer.shard_bindings() == ((target, shard),) + assert optimizer.sharding_manifest() == manifest + assert _contract_is_exact(optimizer) + + local_initial = _local(target).detach().clone() + local_gradient = ( + first_gradient.clone() if kind == "replicate" else _local_slice(first_gradient, shard.logical_region) + ) + local_reference_parameter = nn.Parameter(local_initial) + local_reference = _optimizer(local_reference_parameter) + + target.grad = distribute_tensor(first_gradient.clone(), mesh, [placement]) + local_reference_parameter.grad = local_gradient + optimizer.step() + local_reference.step() + assert torch.equal(_local(target.detach()), local_reference_parameter.detach()) + _assert_optimizer_exact(optimizer, target, local_reference, local_reference_parameter, "first local oracle") + + native_checkpoint = copy.deepcopy(optimizer.state_dict()) + public_checkpoint = get_optimizer_state_dict( + model, + optimizer, + options=StateDictOptions(full_state_dict=True, cpu_offload=True), + ) + current_full = target.detach().full_tensor().clone() + + native_old = nn.Parameter(current_full.clone()) + native_target = nn.Parameter(distribute_tensor(current_full.clone(), mesh, [placement])) + native_optimizer = _optimizer(native_old) + _bind(native_optimizer, native_old, native_target, shard, manifest) + native_optimizer.load_state_dict(native_checkpoint) + + public_old = nn.Parameter(current_full.clone()) + public_target = nn.Parameter(distribute_tensor(current_full.clone(), mesh, [placement])) + public_optimizer = _optimizer(public_old) + _bind(public_optimizer, public_old, public_target, shard, manifest) + public_model = _model(public_target) + set_optimizer_state_dict( + public_model, + public_optimizer, + public_checkpoint if rank == 0 else {}, + options=StateDictOptions(full_state_dict=True, broadcast_from_rank0=True), + ) + + assert torch.equal(_local(native_target.detach()), _local(target.detach())) + assert torch.equal(_local(public_target.detach()), _local(target.detach())) + _assert_optimizer_exact(native_optimizer, native_target, optimizer, target, "native restore") + _assert_optimizer_exact(public_optimizer, public_target, optimizer, target, "public DCP restore") + + second_local_gradient = ( + second_gradient.clone() if kind == "replicate" else _local_slice(second_gradient, shard.logical_region) + ) + target.grad = distribute_tensor(second_gradient.clone(), mesh, [placement]) + native_target.grad = distribute_tensor(second_gradient.clone(), mesh, [placement]) + public_target.grad = distribute_tensor(second_gradient.clone(), mesh, [placement]) + local_reference_parameter.grad = second_local_gradient + optimizer.step() + native_optimizer.step() + public_optimizer.step() + local_reference.step() + + for label, candidate_optimizer, candidate_parameter in ( + ("source continuation", optimizer, target), + ("native continuation", native_optimizer, native_target), + ("public DCP continuation", public_optimizer, public_target), + ): + assert torch.equal(_local(candidate_parameter.detach()), local_reference_parameter.detach()), label + _assert_optimizer_exact( + candidate_optimizer, candidate_parameter, local_reference, local_reference_parameter, label + ) + + empty = shard.logical_region.numel == 0 + if empty: + assert set(_persistent_snapshot(optimizer, target)["state"]) == {"name"} + assert optimizer._gefen_codebook is None + return { + "shape": shape, + "kind": kind, + "dimension": dimension, + "empty": empty, + } + + +def _valid_worker(rank, init_file, result_queue): + from torch.distributed.tensor import init_device_mesh + + try: + dist.init_process_group( + "gloo", + init_method="file://{}".format(init_file), + rank=rank, + world_size=_WORLD, + timeout=timedelta(seconds=90), + ) + mesh = init_device_mesh("cpu", (_WORLD,), mesh_dim_names=("dp",)) + cases = [ + ((5, 4), "shard", 0), + ((1, 4), "shard", 0), + ((4, 5), "shard", 1), + ((4, 1), "shard", 1), + ((5, 4), "replicate", None), + ] + results = [ + _run_valid_case(rank, mesh, shape=shape, kind=kind, dimension=dimension) for shape, kind, dimension in cases + ] + result_queue.put({"rank": rank, "cases": results}) + except BaseException: + result_queue.put({"rank": rank, "traceback": traceback.format_exc()}) + raise + finally: + if dist.is_available() and dist.is_initialized(): + dist.destroy_process_group() + + +def _scoped_distributed_muon_contract_worker(rank, init_file, result_queue): + from gefen import ( + CheckpointTransport, + CodebookProcessGroupBinding, + GefenMuon, + ParameterLayout, + ParameterRebinding, + PlacementKind, + ProcessGroupScope, + ) + from torch import nn + from torch.distributed.tensor import Shard, distribute_tensor, init_device_mesh + + try: + dist.init_process_group( + "gloo", + init_method="file://{}".format(init_file), + rank=rank, + world_size=_WORLD, + timeout=timedelta(seconds=90), + ) + mesh = init_device_mesh("cpu", (_WORLD,), mesh_dim_names=("dp",)) + full = torch.linspace(-0.75, 0.5, 16, dtype=torch.float32).reshape(4, 4) + gradient = torch.arange(16, dtype=torch.float32).reshape(4, 4).sin() + manifest, shards = _manifest( + "Model.Weight", + full.shape, + kind=PlacementKind.DIMENSION_SHARD, + dimension=0, + ) + + def build(values): + old = nn.Parameter(values.clone()) + target = nn.Parameter(distribute_tensor(values.clone(), mesh, [Shard(0)])) + optimizer = GefenMuon( + [("legacy.weight", old)], + lr=3e-3, + fused=False, + deterministic=True, + sharded_mode="distributed", + ) + optimizer.post_sharding( + (ParameterRebinding(old, target, shards[rank]),), + manifest=manifest, + codebook_process_group=CodebookProcessGroupBinding( + shards[rank].process_group, + _MEMBERS[rank], + dist.group.WORLD, + torch.device("cpu"), + ), + ) + return optimizer, target + + source, source_parameter = build(full) + supports = [ + support + for support in source.optimizer_contract().capabilities.checkpoints + if support.transport is CheckpointTransport.NATIVE_OPTIMIZER + and ParameterLayout.DTENSOR_1D_DEFAULT_WORLD in support.same_topology + ] + contract_exact = ( + len(supports) == 1 + and supports[0].same_topology == frozenset({ParameterLayout.DTENSOR_1D_DEFAULT_WORLD}) + and not supports[0].topology_changing + and not supports[0].topology_change_kinds + and supports[0].process_group_scope is ProcessGroupScope.INFERRED_DEVICE_MESH + and supports[0].requires_collective + and supports[0].atomic_load + ) + + source_parameter.grad = distribute_tensor(gradient.clone(), mesh, [Shard(0)]) + source.step() + checkpoint = copy.deepcopy(source.state_dict()) + current = source_parameter.detach().full_tensor().clone() + resumed, resumed_parameter = build(current) + resumed.load_state_dict(checkpoint) + common_state_restored = resumed._gefen_global_step == source._gefen_global_step and torch.equal( + resumed._gefen_codebook, source._gefen_codebook + ) + owner_state_restored = True + if rank == 0: + try: + _assert_optimizer_exact( + resumed, + resumed_parameter, + source, + source_parameter, + "scoped distributed Muon native restore", + ) + except AssertionError: + owner_state_restored = False + + continuation_gradient = gradient.cos() + source_parameter.grad = distribute_tensor(continuation_gradient.clone(), mesh, [Shard(0)]) + resumed_parameter.grad = distribute_tensor(continuation_gradient.clone(), mesh, [Shard(0)]) + source.step() + resumed.step() + continued_exactly = torch.equal( + _local(source_parameter.detach()), + _local(resumed_parameter.detach()), + ) + if rank == 0: + try: + _assert_optimizer_exact( + resumed, + resumed_parameter, + source, + source_parameter, + "scoped distributed Muon native continuation", + ) + except AssertionError: + continued_exactly = False + result_queue.put( + { + "rank": rank, + "contract_exact": contract_exact, + "common_state_restored": common_state_restored, + "owner_state_restored": owner_state_restored, + "continued_exactly": continued_exactly, + } + ) + except BaseException: + result_queue.put({"rank": rank, "traceback": traceback.format_exc()}) + raise + finally: + if dist.is_available() and dist.is_initialized(): + dist.destroy_process_group() + + +def _rejection_result(optimizer, targets, operation): + from _state_snapshot import assert_deep_state_snapshot, deep_state_snapshot + + if torch.is_tensor(targets): + targets = (targets,) + else: + targets = tuple(targets) + optimizer_snapshot = deep_state_snapshot(optimizer) + target_snapshots = tuple(_tensor_snapshot(target) for target in targets) + error = None + try: + operation() + except (TypeError, ValueError, RuntimeError) as exc: + error = "{}: {}".format(type(exc).__name__, exc) + try: + assert_deep_state_snapshot(optimizer, optimizer_snapshot) + unchanged = all(_tensor_unchanged(target, snapshot) for target, snapshot in zip(targets, target_snapshots)) + except AssertionError: + unchanged = False + return {"rejected": error is not None, "unchanged": unchanged, "error": error} + + +def _run_rejections(rank, mesh): + from gefen import ParameterLayout, ParameterRebinding, PlacementKind + from torch import nn + from torch.distributed.tensor import DeviceMesh, Replicate, Shard, distribute_tensor + + full = torch.arange(16, dtype=torch.float32).reshape(4, 4) + results = {} + + manifest, shards = _manifest( + "Model.Weight", + full.shape, + kind=PlacementKind.REPLICATE, + layout=ParameterLayout.REPLICATED, + ) + old = nn.Parameter(full.clone()) + target = nn.Parameter(distribute_tensor(full.clone(), mesh, [Shard(0)])) + optimizer = _optimizer(old) + results["wrong_layout"] = _rejection_result( + optimizer, + target, + lambda: _bind(optimizer, old, target, shards[rank], manifest), + ) + + manifest, shards = _manifest( + "Model.Weight", + full.shape, + kind=PlacementKind.DIMENSION_SHARD, + dimension=0, + ) + old = nn.Parameter(full.clone()) + target = nn.Parameter(full[rank * 2 : (rank + 1) * 2].clone()) + optimizer = _optimizer(old) + results["wrong_type"] = _rejection_result( + optimizer, + target, + lambda: _bind(optimizer, old, target, shards[rank], manifest), + ) + + three_members = ("rank:0", "rank:1", "rank:2") + manifest, shards = _manifest( + "Model.Weight", + full.shape, + kind=PlacementKind.REPLICATE, + members=three_members, + ) + old = nn.Parameter(full.clone()) + target = nn.Parameter(distribute_tensor(full.clone(), mesh, [Replicate()])) + optimizer = _optimizer(old) + results["wrong_world"] = _rejection_result( + optimizer, + target, + lambda: _bind(optimizer, old, target, shards[rank], manifest), + ) + + manifest, shards = _manifest( + "Model.Weight", + full.shape, + kind=PlacementKind.DIMENSION_SHARD, + dimension=0, + ) + old = nn.Parameter(full.clone()) + target = nn.Parameter(distribute_tensor(full.clone(), mesh, [Shard(1)])) + optimizer = _optimizer(old) + results["wrong_placement"] = _rejection_result( + optimizer, + target, + lambda: _bind(optimizer, old, target, shards[rank], manifest), + ) + + old = nn.Parameter(full.clone()) + target = nn.Parameter(distribute_tensor(full.clone(), mesh, [Shard(0)])) + optimizer = _optimizer(old) + results["wrong_coordinate"] = _rejection_result( + optimizer, + target, + lambda: _bind(optimizer, old, target, shards[1 - rank], manifest), + ) + + reverse_mesh = DeviceMesh("cpu", [1, 0], mesh_dim_names=("reverse",)) + old = nn.Parameter(full.clone()) + target = nn.Parameter(distribute_tensor(full.clone(), reverse_mesh, [Shard(0)])) + optimizer = _optimizer(old) + results["wrong_mesh"] = _rejection_result( + optimizer, + target, + lambda: _bind(optimizer, old, target, shards[rank], manifest), + ) + + first_manifest, first_shards = _manifest( + "Model.First", + full.shape, + kind=PlacementKind.DIMENSION_SHARD, + dimension=0, + semantic_name="shared_mesh", + ) + second_manifest, second_shards = _manifest( + "Model.Second", + full.shape, + kind=PlacementKind.DIMENSION_SHARD, + dimension=0, + semantic_name="shared_mesh", + ) + combined = _combined_manifest(first_manifest, second_manifest) + first_old = nn.Parameter(full.clone()) + second_old = nn.Parameter(full.neg().clone()) + first_target = nn.Parameter(distribute_tensor(full.clone(), mesh, [Shard(0)])) + second_target = nn.Parameter(distribute_tensor(full.neg().clone(), reverse_mesh, [Shard(0)])) + mixed = _two_parameter_optimizer(first_old, second_old) + results["shared_mesh"] = _rejection_result( + mixed, + (first_target, second_target), + lambda: mixed.post_sharding( + ( + ParameterRebinding(first_old, first_target, first_shards[rank]), + ParameterRebinding(second_old, second_target, second_shards[rank]), + ), + manifest=combined, + ), + ) + return results + + +def _rejection_worker(rank, init_file, result_queue): + from torch.distributed.tensor import init_device_mesh + + try: + dist.init_process_group( + "gloo", + init_method="file://{}".format(init_file), + rank=rank, + world_size=_WORLD, + timeout=timedelta(seconds=90), + ) + mesh = init_device_mesh("cpu", (_WORLD,), mesh_dim_names=("dp",)) + result_queue.put({"rank": rank, "results": _run_rejections(rank, mesh)}) + except BaseException: + result_queue.put({"rank": rank, "traceback": traceback.format_exc()}) + raise + finally: + if dist.is_available() and dist.is_initialized(): + dist.destroy_process_group() + + +def _scoped_replicated_dtensor_worker(rank, init_file, result_queue): + from gefen import ( + CodebookProcessGroupBinding, + ParameterRebinding, + PlacementKind, + ) + from gefen.gefen import learn_gefen_exact_codebook_from_grad_periods + from torch import nn + from torch.distributed.tensor import DTensor, Replicate, Shard, distribute_tensor, init_device_mesh + + try: + dist.init_process_group( + "gloo", + init_method="file://{}".format(init_file), + rank=rank, + world_size=_WORLD, + timeout=timedelta(seconds=90), + ) + mesh = init_device_mesh("cpu", (_WORLD,), mesh_dim_names=("dp",)) + replicated_manifest, replicated_shards = _manifest( + "Model.Replicated", + (4,), + kind=PlacementKind.REPLICATE, + semantic_name="scoped_dtensor", + ) + sharded_manifest, sharded_shards = _manifest( + "Model.Sharded", + (4,), + kind=PlacementKind.DIMENSION_SHARD, + dimension=0, + semantic_name="scoped_dtensor", + ) + manifest = _combined_manifest(replicated_manifest, sharded_manifest) + binding = CodebookProcessGroupBinding( + replicated_shards[rank].process_group, + _MEMBERS[rank], + dist.group.WORLD, + torch.device("cpu"), + ) + + replicated_old = nn.Parameter(torch.zeros(4)) + sharded_old = nn.Parameter(torch.zeros(4)) + replicated = nn.Parameter(distribute_tensor(torch.zeros(4), mesh, [Replicate()])) + sharded = nn.Parameter(distribute_tensor(torch.zeros(4), mesh, [Shard(0)])) + optimizer = _two_parameter_optimizer(replicated_old, sharded_old) + optimizer.post_sharding( + ( + ParameterRebinding(replicated_old, replicated, replicated_shards[rank]), + ParameterRebinding(sharded_old, sharded, sharded_shards[rank]), + ), + manifest=manifest, + codebook_process_group=binding, + ) + optimizer._resolve_automatic_period = lambda *args: 2 + replicated_gradients = ( + torch.tensor([-4.0, -1.0, 2.0, 8.0]), + torch.tensor([-9.0, 3.0, 5.0, 7.0]), + ) + sharded_gradients = ( + torch.tensor([-3.0, 6.0]), + torch.tensor([2.0, 11.0]), + ) + replicated.grad = DTensor.from_local( + replicated_gradients[rank].clone(), + mesh, + [Replicate()], + run_check=False, + ) + sharded.grad = DTensor.from_local( + sharded_gradients[rank].clone(), + mesh, + [Shard(0)], + run_check=False, + ) + initialized = optimizer.initialize_codebook() + intended_oracle = learn_gefen_exact_codebook_from_grad_periods( + grad_periods=( + ( + "Model.Replicated", + replicated_gradients[0], + 2, + replicated_gradients[0], + ), + *tuple(("Model.Sharded", gradient, 2, gradient) for gradient in sharded_gradients), + ), + codebook_device=torch.device("cpu"), + num_codebooks=256, + force_endpoints=True, + verbose=False, + compute_mse_logging=False, + use_fused_histogram=False, + ) + duplicated_oracle = learn_gefen_exact_codebook_from_grad_periods( + grad_periods=( + *tuple(("Model.Replicated", gradient, 2, gradient) for gradient in replicated_gradients), + *tuple(("Model.Sharded", gradient, 2, gradient) for gradient in sharded_gradients), + ), + codebook_device=torch.device("cpu"), + num_codebooks=256, + force_endpoints=True, + verbose=False, + compute_mse_logging=False, + use_fused_histogram=False, + ) + logical_once = torch.equal(optimizer._gefen_codebook, intended_oracle) + duplicate_would_differ = not torch.equal(intended_oracle, duplicated_oracle) + gathered = [torch.empty_like(optimizer._gefen_codebook) for _ in range(_WORLD)] + dist.all_gather(gathered, optimizer._gefen_codebook) + agreed = all(torch.equal(item, gathered[0]) for item in gathered[1:]) + + mismatch_replicated_old = nn.Parameter(torch.zeros(4)) + mismatch_sharded_old = nn.Parameter(torch.zeros(4)) + mismatch_replicated = nn.Parameter(distribute_tensor(torch.zeros(4), mesh, [Replicate()])) + mismatch_sharded = nn.Parameter(distribute_tensor(torch.zeros(4), mesh, [Shard(0)])) + mismatch_optimizer = _two_parameter_optimizer( + mismatch_replicated_old, + mismatch_sharded_old, + ) + mismatch_optimizer.post_sharding( + ( + ParameterRebinding( + mismatch_replicated_old, + mismatch_replicated, + replicated_shards[rank], + ), + ParameterRebinding( + mismatch_sharded_old, + mismatch_sharded, + sharded_shards[rank], + ), + ), + manifest=manifest, + codebook_process_group=binding, + ) + + def mismatched_period(_name, parameter, _gradient): + if parameter is mismatch_replicated: + return 2 if rank == 0 else 4 + return 2 + + mismatch_optimizer._resolve_automatic_period = mismatched_period + mismatch_replicated.grad = DTensor.from_local( + replicated_gradients[0].clone(), + mesh, + [Replicate()], + run_check=False, + ) + mismatch_sharded.grad = DTensor.from_local( + sharded_gradients[rank].clone(), + mesh, + [Shard(0)], + run_check=False, + ) + mismatch_before = ( + _persistent_snapshot(mismatch_optimizer, mismatch_replicated), + _persistent_snapshot(mismatch_optimizer, mismatch_sharded), + ) + try: + mismatch_optimizer.initialize_codebook() + period_mismatch_rejected = False + except RuntimeError as exc: + period_mismatch_rejected = "automatic periods" in str(exc) + try: + _assert_value_exact( + ( + _persistent_snapshot(mismatch_optimizer, mismatch_replicated), + _persistent_snapshot(mismatch_optimizer, mismatch_sharded), + ), + mismatch_before, + "period mismatch atomicity", + ) + mismatch_atomic = True + except AssertionError: + mismatch_atomic = False + result_queue.put( + { + "rank": rank, + "initialized": initialized, + "logical_once": logical_once, + "duplicate_would_differ": duplicate_would_differ, + "agreed": agreed, + "period_mismatch_rejected": period_mismatch_rejected, + "mismatch_atomic": mismatch_atomic, + } + ) + except BaseException: + result_queue.put({"rank": rank, "traceback": traceback.format_exc()}) + raise + finally: + if dist.is_available() and dist.is_initialized(): + dist.destroy_process_group() + + +def _scoped_muon_codebook_worker(rank, init_file, result_queue): + from gefen import ( + CodebookProcessGroupBinding, + GefenMuon, + ParameterRebinding, + PlacementKind, + ) + from gefen.gefen import learn_gefen_exact_codebook_from_grad_periods + from torch import nn + from torch.distributed.tensor import Replicate, Shard, distribute_tensor, init_device_mesh + + try: + dist.init_process_group( + "gloo", + init_method="file://{}".format(init_file), + rank=rank, + world_size=_WORLD, + timeout=timedelta(seconds=90), + ) + mesh = init_device_mesh("cpu", (_WORLD,), mesh_dim_names=("dp",)) + shape = (16, 32) + generator = torch.Generator().manual_seed(0) + replicated_gradient = torch.randn(shape, generator=generator) + sharded_gradient = torch.randn(shape, generator=generator).mul_(3.0) + replicated_manifest, replicated_shards = _manifest( + "Model.Replicated", + shape, + kind=PlacementKind.REPLICATE, + semantic_name="scoped_muon_dtensor", + ) + sharded_manifest, sharded_shards = _manifest( + "Model.Sharded", + shape, + kind=PlacementKind.DIMENSION_SHARD, + dimension=0, + semantic_name="scoped_muon_dtensor", + ) + manifest = _combined_manifest(replicated_manifest, sharded_manifest) + + def build(mode, *, scoped=True): + replicated_old = nn.Parameter(torch.zeros(shape)) + sharded_old = nn.Parameter(torch.zeros(shape)) + replicated = nn.Parameter(distribute_tensor(torch.zeros(shape), mesh, [Replicate()])) + sharded = nn.Parameter(distribute_tensor(torch.zeros(shape), mesh, [Shard(0)])) + optimizer = GefenMuon( + [ + ("legacy.replicated", replicated_old), + ("legacy.sharded", sharded_old), + ], + lr=3e-3, + fused=False, + deterministic=True, + sharded_mode=mode, + ) + post_sharding_options = {} + if scoped: + post_sharding_options["codebook_process_group"] = ( + CodebookProcessGroupBinding( + replicated_shards[rank].process_group, + _MEMBERS[rank], + dist.group.WORLD, + torch.device("cpu"), + ) + ) + optimizer.post_sharding( + ( + ParameterRebinding( + replicated_old, + replicated, + replicated_shards[rank], + ), + ParameterRebinding( + sharded_old, + sharded, + sharded_shards[rank], + ), + ), + manifest=manifest, + **post_sharding_options, + ) + return optimizer, replicated, sharded + + def set_gradients(replicated, sharded, *, sharded_active=True): + replicated.grad = distribute_tensor( + replicated_gradient.clone(), mesh, [Replicate()] + ) + sharded_value = distribute_tensor( + sharded_gradient.clone(), mesh, [Shard(0)] + ) + sharded.grad = sharded_value if sharded_active else None + + def solve(items): + return learn_gefen_exact_codebook_from_grad_periods( + grad_periods=items, + codebook_device=torch.device("cpu"), + num_codebooks=256, + force_endpoints=True, + verbose=False, + compute_mse_logging=False, + use_fused_histogram=False, + ) + + def values_equal(actual, expected, label): + try: + _assert_value_exact(actual, expected, label) + except AssertionError: + return False + return True + + exact, exact_replicated, exact_sharded = build("exact") + exact._predict_period_from_grad_sq = lambda *args: 16 + set_gradients(exact_replicated, exact_sharded) + exact_initialized = exact.initialize_codebook() + exact_oracle = solve( + ( + ("Model.Replicated", replicated_gradient, 16, replicated_gradient), + ("Model.Sharded", sharded_gradient, 16, sharded_gradient), + ) + ) + duplicated_shard_oracle = solve( + ( + ("Model.Replicated", replicated_gradient, 16, replicated_gradient), + ("Model.Sharded", sharded_gradient, 16, sharded_gradient), + ("Model.Sharded", sharded_gradient, 16, sharded_gradient), + ) + ) + checkpoint_hook = None + if rank == 0: + def reject_checkpoint(_optimizer): + raise RuntimeError("scoped checkpoint hook failure") + + checkpoint_hook = exact.register_state_dict_pre_hook( + reject_checkpoint + ) + try: + exact.state_dict() + scoped_checkpoint_hook_rejected = False + except RuntimeError as exc: + scoped_checkpoint_hook_rejected = ( + "checkpoint state-dict pre-hook failed" in str(exc) + ) + finally: + if checkpoint_hook is not None: + checkpoint_hook.remove() + + approximate, approximate_replicated, approximate_sharded = build("approx") + approximate._predict_period_from_grad_sq = lambda *args: 16 + set_gradients(approximate_replicated, approximate_sharded) + approximate_initialized = approximate.initialize_codebook() + approximate_oracle = solve( + ( + ("Model.Replicated", replicated_gradient, 16, replicated_gradient), + ( + "Model.Sharded", + sharded_gradient[: shape[0] // 2], + 16, + sharded_gradient[: shape[0] // 2], + ), + ( + "Model.Sharded", + sharded_gradient[shape[0] // 2 :], + 16, + sharded_gradient[shape[0] // 2 :], + ), + ) + ) + + mismatch, mismatch_replicated, mismatch_sharded = build("distributed") + + def mismatched_period(_name, parameter, _gradient): + if parameter is mismatch_sharded: + return 16 if rank == 0 else 32 + return 16 + + mismatch._predict_period_from_grad_sq = mismatched_period + set_gradients(mismatch_replicated, mismatch_sharded) + mismatch_before = ( + _persistent_snapshot(mismatch, mismatch_replicated), + _persistent_snapshot(mismatch, mismatch_sharded), + ) + try: + mismatch.initialize_codebook() + period_mismatch_rejected = False + except RuntimeError as exc: + period_mismatch_rejected = "automatic periods" in str(exc) + mismatch_after = ( + _persistent_snapshot(mismatch, mismatch_replicated), + _persistent_snapshot(mismatch, mismatch_sharded), + ) + + presence, presence_replicated, presence_sharded = build("exact") + presence._predict_period_from_grad_sq = lambda *args: 16 + set_gradients( + presence_replicated, + presence_sharded, + sharded_active=rank == 0, + ) + presence_before = ( + _persistent_snapshot(presence, presence_replicated), + _persistent_snapshot(presence, presence_sharded), + ) + try: + presence.initialize_codebook() + initialize_presence_rejected = False + except RuntimeError as exc: + initialize_presence_rejected = "identical gradient presence" in str(exc) + presence_after = ( + _persistent_snapshot(presence, presence_replicated), + _persistent_snapshot(presence, presence_sharded), + ) + + refresh, refresh_replicated, refresh_sharded = build("exact") + refresh._predict_period_from_grad_sq = lambda *args: 16 + set_gradients(refresh_replicated, refresh_sharded) + refresh.initialize_codebook() + refresh_before = ( + _persistent_snapshot(refresh, refresh_replicated), + _persistent_snapshot(refresh, refresh_sharded), + ) + set_gradients( + refresh_replicated, + refresh_sharded, + sharded_active=rank == 0, + ) + try: + refresh.refresh_codebook() + refresh_presence_rejected = False + except RuntimeError as exc: + refresh_presence_rejected = "identical gradient presence" in str(exc) + refresh_after = ( + _persistent_snapshot(refresh, refresh_replicated), + _persistent_snapshot(refresh, refresh_sharded), + ) + + divergent, divergent_replicated, divergent_sharded = build( + "approx" if rank else "exact" + ) + divergent._predict_period_from_grad_sq = lambda *args: 16 + set_gradients(divergent_replicated, divergent_sharded) + try: + divergent.initialize_codebook() + mode_mismatch_rejected = False + except RuntimeError as exc: + mode_mismatch_rejected = "policy" in str(exc) + + unscoped_divergent, unscoped_replicated, unscoped_sharded = build( + "approx" if rank else "exact", + scoped=False, + ) + unscoped_divergent._predict_period_from_grad_sq = lambda *args: 16 + set_gradients(unscoped_replicated, unscoped_sharded) + try: + unscoped_divergent.initialize_codebook() + unscoped_mode_mismatch_rejected = False + except RuntimeError as exc: + unscoped_mode_mismatch_rejected = ( + "identical sharded_mode collective intent" in str(exc) + ) + + result_queue.put( + { + "rank": rank, + "exact_initialized": exact_initialized, + "exact_logical_once": torch.equal(exact._gefen_codebook, exact_oracle), + "duplicate_would_differ": not torch.equal( + exact_oracle, duplicated_shard_oracle + ), + "scoped_checkpoint_hook_rejected": ( + scoped_checkpoint_hook_rejected + ), + "approximate_initialized": approximate_initialized, + "approximate_local_shards": torch.equal( + approximate._gefen_codebook, approximate_oracle + ), + "period_mismatch_rejected": period_mismatch_rejected, + "period_mismatch_atomic": values_equal( + mismatch_after, mismatch_before, "period mismatch atomicity" + ), + "initialize_presence_rejected": initialize_presence_rejected, + "initialize_presence_atomic": values_equal( + presence_after, presence_before, "initialize presence atomicity" + ), + "refresh_presence_rejected": refresh_presence_rejected, + "refresh_presence_atomic": values_equal( + refresh_after, refresh_before, "refresh presence atomicity" + ), + "mode_mismatch_rejected": mode_mismatch_rejected, + "unscoped_mode_mismatch_rejected": ( + unscoped_mode_mismatch_rejected + ), + } + ) + except BaseException: + result_queue.put({"rank": rank, "traceback": traceback.format_exc()}) + raise + finally: + if dist.is_available() and dist.is_initialized(): + dist.destroy_process_group() + + +def _unscoped_muon_mode_transition_worker(rank, init_file, result_queue): + from gefen import GefenMuon, ParameterRebinding, PlacementKind + from torch import nn + from torch.distributed.tensor import Shard, distribute_tensor, init_device_mesh + + try: + dist.init_process_group( + "gloo", + init_method="file://{}".format(init_file), + rank=rank, + world_size=_WORLD, + timeout=timedelta(seconds=90), + ) + mesh = init_device_mesh("cpu", (_WORLD,), mesh_dim_names=("dp",)) + shape = (4, 4) + full = torch.arange(16, dtype=torch.float32).reshape(shape).div_(16) + gradient = torch.arange(16, dtype=torch.float32).reshape(shape).sin() + manifest, shards = _manifest( + "Model.Weight", + shape, + kind=PlacementKind.DIMENSION_SHARD, + dimension=0, + semantic_name="unscoped_muon_transition", + ) + old = nn.Parameter(full.clone()) + parameter = nn.Parameter(distribute_tensor(full.clone(), mesh, [Shard(0)])) + optimizer = GefenMuon( + [("legacy.weight", old)], + lr=3e-3, + fused=False, + deterministic=True, + sharded_mode="approx", + ) + optimizer.post_sharding( + (ParameterRebinding(old, parameter, shards[rank]),), + manifest=manifest, + ) + optimizer._predict_period_from_grad_sq = lambda *args: 1 + parameter.grad = distribute_tensor(gradient.clone(), mesh, [Shard(0)]) + initialized = optimizer.initialize_codebook() + parameter.grad = distribute_tensor(gradient.cos(), mesh, [Shard(0)]) + optimizer.step() + warm_global_step = optimizer._gefen_global_step + warm_state_initialized = "m_codebook" in optimizer.state[parameter] + parameter_before = _tensor_snapshot(parameter) + state_before = _persistent_snapshot(optimizer, parameter) + if rank == 0: + optimizer.param_groups[0]["sharded_mode"] = "exact" + try: + optimizer.step() + asymmetric_rejected = False + except RuntimeError as exc: + asymmetric_rejected = ( + "identical sharded_mode collective intent" in str(exc) + ) + asymmetric_parameter_unchanged = _tensor_unchanged( + parameter, parameter_before + ) + try: + _assert_value_exact( + _persistent_snapshot(optimizer, parameter), + state_before, + "asymmetric post-warmup mode transition atomicity", + ) + asymmetric_state_unchanged = True + except AssertionError: + asymmetric_state_unchanged = False + + optimizer.param_groups[0]["sharded_mode"] = "distributed" + try: + optimizer.step() + unanimous_rejected = False + except RuntimeError as exc: + unanimous_rejected = "post-validation DTensor" in str(exc) + unanimous_parameter_unchanged = _tensor_unchanged( + parameter, parameter_before + ) + try: + _assert_value_exact( + _persistent_snapshot(optimizer, parameter), + state_before, + "unanimous post-warmup mode transition atomicity", + ) + unanimous_state_unchanged = True + except AssertionError: + unanimous_state_unchanged = False + + optimizer.param_groups[0]["sharded_mode"] = "approx" + replacement = nn.Parameter( + distribute_tensor((full + 1).clone(), mesh, [Shard(0)]) + ) + group = optimizer.param_groups[0] + original_params = list(group["params"]) + original_names = list(group["param_names"]) + + def unchanged(label): + if not _tensor_unchanged(parameter, parameter_before): + return False + try: + state_after = _persistent_snapshot(optimizer, parameter) + _assert_value_exact( + state_after, + state_before, + label, + ) + return True + except AssertionError: + return False + + if rank == 0: + group["sharded_mode"] = "invalid" + try: + optimizer.step() + invalid_mode_rejected = False + except RuntimeError as exc: + invalid_mode_rejected = "requires sharded_mode" in str(exc) + invalid_mode_unchanged = unchanged("invalid-mode header atomicity") + group["sharded_mode"] = "approx" + + if rank == 0: + group["params"] = [] + group["param_names"] = [] + try: + optimizer.step() + membership_rejected = False + except RuntimeError as exc: + membership_rejected = "identical DTensor parameter membership" in str( + exc + ) + membership_unchanged = unchanged("membership header atomicity") + group["params"] = list(original_params) + group["param_names"] = list(original_names) + + if rank == 0: + group["params"] = None + try: + optimizer.refresh_codebook() + malformed_rejected = False + except RuntimeError as exc: + malformed_rejected = "structurally valid and identical" in str(exc) + malformed_unchanged = unchanged("malformed-group header atomicity") + group["params"] = list(original_params) + + if rank == 0: + group["params"] = [replacement] + try: + optimizer.step() + rebind_rejected = False + except RuntimeError as exc: + rebind_rejected = "post-validation DTensor" in str(exc) + rebind_unchanged = unchanged("rebind header atomicity") + group["params"] = list(original_params) + + group["sharded_mode"] = "exact" + baseline_before_add = optimizer._gefen_muon_collective_intent_baseline + capture_plan_before_add = optimizer._gefen_muon_capture_plan + try: + optimizer.add_param_group( + { + "params": [("extra", replacement)], + "sharded_mode": "approx", + } + ) + warm_add_rejected = False + except RuntimeError as exc: + warm_add_rejected = "cannot add a parameter group" in str(exc) + warm_add_preserved_plan = ( + optimizer._gefen_muon_collective_intent_baseline + is baseline_before_add + and optimizer._gefen_muon_capture_plan is capture_plan_before_add + and len(optimizer.param_groups) == 1 + ) + try: + optimizer.step() + add_bypass_transition_rejected = False + except RuntimeError as exc: + add_bypass_transition_rejected = "post-validation DTensor" in str(exc) + add_bypass_unchanged = unchanged("warm add-group bypass atomicity") + group["sharded_mode"] = "approx" + + checkpoint = optimizer.state_dict() + old_capture_plan = optimizer._gefen_muon_capture_plan + optimizer.load_state_dict(checkpoint) + loaded_group = optimizer.param_groups[0] + fresh_loaded_plan = optimizer._gefen_muon_capture_plan + load_refroze_fresh_plan = ( + optimizer._gefen_muon_collective_intent_baseline is not None + and fresh_loaded_plan is not None + and fresh_loaded_plan is not old_capture_plan + and fresh_loaded_plan[0][0][0] is loaded_group + ) + loaded_group["sharded_mode"] = "exact" + try: + optimizer.add_param_group( + { + "params": [("post_load_extra", replacement)], + "sharded_mode": "approx", + } + ) + post_load_add_rejected = False + except RuntimeError as exc: + post_load_add_rejected = "cannot add a parameter group" in str(exc) + post_load_add_preserved_plan = ( + optimizer._gefen_muon_capture_plan is fresh_loaded_plan + and len(optimizer.param_groups) == 1 + ) + loaded_group["sharded_mode"] = "approx" + load_plan_unchanged = unchanged("post-load route-plan atomicity") + result_queue.put( + { + "rank": rank, + "initialized": initialized, + "warm_global_step": warm_global_step, + "warm_state_initialized": warm_state_initialized, + "asymmetric_rejected": asymmetric_rejected, + "asymmetric_parameter_unchanged": ( + asymmetric_parameter_unchanged + ), + "asymmetric_state_unchanged": asymmetric_state_unchanged, + "unanimous_rejected": unanimous_rejected, + "unanimous_parameter_unchanged": unanimous_parameter_unchanged, + "unanimous_state_unchanged": unanimous_state_unchanged, + "invalid_mode_rejected": invalid_mode_rejected, + "invalid_mode_unchanged": invalid_mode_unchanged, + "membership_rejected": membership_rejected, + "membership_unchanged": membership_unchanged, + "malformed_rejected": malformed_rejected, + "malformed_unchanged": malformed_unchanged, + "rebind_rejected": rebind_rejected, + "rebind_unchanged": rebind_unchanged, + "warm_add_rejected": warm_add_rejected, + "warm_add_preserved_plan": warm_add_preserved_plan, + "add_bypass_transition_rejected": ( + add_bypass_transition_rejected + ), + "add_bypass_unchanged": add_bypass_unchanged, + "load_refroze_fresh_plan": load_refroze_fresh_plan, + "post_load_add_rejected": post_load_add_rejected, + "post_load_add_preserved_plan": ( + post_load_add_preserved_plan + ), + "load_plan_unchanged": load_plan_unchanged, + "global_step": optimizer._gefen_global_step, + } + ) + except BaseException: + result_queue.put({"rank": rank, "traceback": traceback.format_exc()}) + raise + finally: + if dist.is_available() and dist.is_initialized(): + dist.destroy_process_group() + + +def _muon_checkpoint_routing_worker(rank, init_file, result_queue): + from gefen import GefenMuon + from torch import nn + from torch.distributed.tensor import Shard, distribute_tensor, init_device_mesh + + try: + dist.init_process_group( + "gloo", + init_method="file://{}".format(init_file), + rank=rank, + world_size=_WORLD, + timeout=timedelta(seconds=90), + ) + mesh = init_device_mesh("cpu", (_WORLD,), mesh_dim_names=("dp",)) + full = torch.arange(16, dtype=torch.float32).reshape(4, 4).div_(16) + gradient = torch.arange(16, dtype=torch.float32).reshape(4, 4).sin() + parameter = nn.Parameter( + distribute_tensor(full.clone(), mesh, [Shard(0)]) + ) + optimizer = GefenMuon( + [("weight", parameter)], + lr=3e-3, + fused=False, + deterministic=True, + ns_steps=1, + sharded_mode="approx", + ) + optimizer._predict_period_from_grad_sq = lambda *args: 1 + parameter.grad = distribute_tensor(gradient.clone(), mesh, [Shard(0)]) + optimizer.step() + parameter.grad = None + + checkpoint = copy.deepcopy(optimizer.state_dict()) + valid_export = ( + checkpoint["param_groups"][0]["sharded_mode"] == "approx" + and optimizer._gefen_global_step == 1 + ) + optimizer.param_groups[0]["sharded_mode"] = "exact" + try: + optimizer.optimizer_contract() + frozen_contract_rejected = False + except RuntimeError as exc: + frozen_contract_rejected = "frozen DTensor collective routing" in str( + exc + ) + try: + optimizer._canonical_state_variant_layout() + frozen_layout_rejected = False + except RuntimeError as exc: + frozen_layout_rejected = "frozen DTensor collective routing" in str( + exc + ) + optimizer.param_groups[0]["sharded_mode"] = "invalid" + try: + optimizer.optimizer_contract() + invalid_contract_rejected = False + except RuntimeError as exc: + invalid_contract_rejected = "invalid sharded_mode" in str(exc) + optimizer.param_groups[0]["sharded_mode"] = "approx" + valid_contract = optimizer.optimizer_contract() is not None + valid_layout = optimizer._canonical_state_variant_layout() is not None + parameter_before = _tensor_snapshot(parameter) + state_before = _persistent_snapshot(optimizer, parameter) + + def unchanged(label): + if not _tensor_unchanged(parameter, parameter_before): + return False + try: + _assert_value_exact( + _persistent_snapshot(optimizer, parameter), + state_before, + label, + ) + return True + except AssertionError: + return False + + if rank == 0: + optimizer.param_groups[0]["sharded_mode"] = "distributed" + try: + optimizer.state_dict() + asymmetric_export_rejected = False + except RuntimeError as exc: + asymmetric_export_rejected = ( + "identical sharded_mode collective intent" in str(exc) + ) + asymmetric_export_atomic = unchanged("asymmetric export atomicity") + + optimizer.param_groups[0]["sharded_mode"] = "distributed" + try: + optimizer.state_dict() + unanimous_export_rejected = False + except RuntimeError as exc: + unanimous_export_rejected = "post-validation DTensor" in str(exc) + unanimous_export_atomic = unchanged("unanimous export atomicity") + optimizer.param_groups[0]["sharded_mode"] = "approx" + + def mutate_route_in_pre_hook(_optimizer): + if rank == 0: + _optimizer.param_groups[0]["sharded_mode"] = "distributed" + + hook = optimizer.register_state_dict_pre_hook(mutate_route_in_pre_hook) + try: + optimizer.state_dict() + hook_export_rejected = False + except RuntimeError as exc: + hook_export_rejected = ( + "identical sharded_mode collective intent" in str(exc) + ) + finally: + hook.remove() + hook_export_atomic = unchanged("pre-hook export atomicity") + optimizer.param_groups[0]["sharded_mode"] = "approx" + + if rank == 0: + optimizer.param_groups[0]["sharded_mode"] = "distributed" + try: + optimizer.load_state_dict(copy.deepcopy(checkpoint)) + asymmetric_load_rejected = False + except RuntimeError as exc: + asymmetric_load_rejected = ( + "identical sharded_mode collective intent" in str(exc) + ) + asymmetric_load_atomic = unchanged("asymmetric load atomicity") + + optimizer.param_groups[0]["sharded_mode"] = "distributed" + try: + optimizer.load_state_dict(copy.deepcopy(checkpoint)) + unanimous_load_rejected = False + except RuntimeError as exc: + unanimous_load_rejected = "post-validation DTensor" in str(exc) + unanimous_load_atomic = unchanged("unanimous load atomicity") + optimizer.param_groups[0]["sharded_mode"] = "approx" + + drift_checkpoint = copy.deepcopy(checkpoint) + drift_checkpoint["param_groups"][0]["sharded_mode"] = "exact" + try: + optimizer.load_state_dict(drift_checkpoint) + checkpoint_route_rejected = False + except RuntimeError as exc: + checkpoint_route_rejected = "post-validation DTensor" in str(exc) + checkpoint_route_atomic = unchanged("checkpoint route atomicity") + + optimizer.load_state_dict(copy.deepcopy(checkpoint)) + valid_load = ( + optimizer.param_groups[0]["sharded_mode"] == "approx" + and unchanged("valid unchanged load") + ) + result_queue.put( + { + "rank": rank, + "valid_export": valid_export, + "frozen_contract_rejected": frozen_contract_rejected, + "frozen_layout_rejected": frozen_layout_rejected, + "invalid_contract_rejected": invalid_contract_rejected, + "valid_contract": valid_contract, + "valid_layout": valid_layout, + "asymmetric_export_rejected": asymmetric_export_rejected, + "asymmetric_export_atomic": asymmetric_export_atomic, + "unanimous_export_rejected": unanimous_export_rejected, + "unanimous_export_atomic": unanimous_export_atomic, + "hook_export_rejected": hook_export_rejected, + "hook_export_atomic": hook_export_atomic, + "asymmetric_load_rejected": asymmetric_load_rejected, + "asymmetric_load_atomic": asymmetric_load_atomic, + "unanimous_load_rejected": unanimous_load_rejected, + "unanimous_load_atomic": unanimous_load_atomic, + "checkpoint_route_rejected": checkpoint_route_rejected, + "checkpoint_route_atomic": checkpoint_route_atomic, + "valid_load": valid_load, + } + ) + except BaseException: + result_queue.put({"rank": rank, "traceback": traceback.format_exc()}) + raise + finally: + if dist.is_available() and dist.is_initialized(): + dist.destroy_process_group() + + +def _muon_checkpoint_callback_failure_worker(rank, init_file, result_queue): + from gefen import GefenMuon + from torch import nn + from torch.distributed.tensor import Shard, distribute_tensor, init_device_mesh + + try: + dist.init_process_group( + "gloo", + init_method="file://{}".format(init_file), + rank=rank, + world_size=_WORLD, + timeout=timedelta(seconds=90), + ) + mesh = init_device_mesh("cpu", (_WORLD,), mesh_dim_names=("dp",)) + full = torch.arange(16, dtype=torch.float32).reshape(4, 4).div_(16) + gradient = torch.arange(16, dtype=torch.float32).reshape(4, 4).sin() + cases = [] + for mode in ("exact", "distributed"): + parameter = nn.Parameter( + distribute_tensor(full.clone(), mesh, [Shard(0)]) + ) + optimizer = GefenMuon( + [("weight", parameter)], + lr=3e-3, + fused=False, + deterministic=True, + ns_steps=1, + sharded_mode=mode, + ) + optimizer._predict_period_from_grad_sq = lambda *args: 1 + parameter.grad = distribute_tensor( + gradient.clone(), mesh, [Shard(0)] + ) + optimizer.step() + parameter.grad = None + checkpoint = copy.deepcopy(optimizer.state_dict()) + parameter_before = _tensor_snapshot(parameter) + state_before = _persistent_snapshot(optimizer, parameter) + + def rejected(register, hook, action, phase): + handle = register(hook) if rank == 0 else None + try: + action() + matched = False + except RuntimeError as exc: + matched = "checkpoint {} failed".format(phase) in str(exc) + finally: + if handle is not None: + handle.remove() + dist.barrier() + return matched + + state_pre = rejected( + optimizer.register_state_dict_pre_hook, + lambda _optimizer: (_ for _ in ()).throw( + RuntimeError("state pre-hook failure") + ), + optimizer.state_dict, + "state-dict pre-hook", + ) + state_post = rejected( + optimizer.register_state_dict_post_hook, + lambda _optimizer, _state: (_ for _ in ()).throw( + RuntimeError("state post-hook failure") + ), + optimizer.state_dict, + "state-dict post-hook", + ) + load_pre = rejected( + optimizer.register_load_state_dict_pre_hook, + lambda _optimizer, _state: (_ for _ in ()).throw( + RuntimeError("load pre-hook failure") + ), + lambda: optimizer.load_state_dict(copy.deepcopy(checkpoint)), + "load pre-hook", + ) + try: + _assert_value_exact( + _persistent_snapshot(optimizer, parameter), + state_before, + "pre-commit checkpoint callback atomicity", + ) + precommit_state_atomic = True + except AssertionError: + precommit_state_atomic = False + load_post = rejected( + optimizer.register_load_state_dict_post_hook, + lambda _optimizer: (_ for _ in ()).throw( + RuntimeError("load post-hook failure") + ), + lambda: optimizer.load_state_dict(copy.deepcopy(checkpoint)), + "load post-hook", + ) + if rank == 0: + optimizer._gefen_global_step_by_device = { + torch.device("cpu"): torch.tensor([1, 2]) + } + try: + optimizer.state_dict() + counter_preparation_rejected = False + except RuntimeError as exc: + counter_preparation_rejected = ( + "checkpoint local serialization preparation failed" + in str(exc) + ) + optimizer._gefen_global_step_by_device = {} + dist.barrier() + + if mode == "distributed": + if rank == 0: + optimizer.state[parameter]["unpickleable"] = lambda: None + try: + optimizer.state_dict() + owner_preparation_rejected = False + except RuntimeError as exc: + owner_preparation_rejected = ( + "checkpoint distributed serialization preparation failed" + in str(exc) + ) + optimizer.state[parameter].pop("unpickleable", None) + dist.barrier() + else: + owner_preparation_rejected = True + valid_checkpoint = optimizer.state_dict() + optimizer.load_state_dict(copy.deepcopy(valid_checkpoint)) + parameter_atomic = _tensor_unchanged(parameter, parameter_before) + cases.append( + { + "mode": mode, + "state_pre": state_pre, + "state_post": state_post, + "load_pre": load_pre, + "load_post": load_post, + "counter_preparation": counter_preparation_rejected, + "owner_preparation": owner_preparation_rejected, + "valid": ( + valid_checkpoint["param_groups"][0]["sharded_mode"] + == mode + ), + "precommit_state_atomic": precommit_state_atomic, + "parameter_atomic": parameter_atomic, + } + ) + dist.barrier() + result_queue.put({"rank": rank, "cases": cases}) + except BaseException: + result_queue.put({"rank": rank, "traceback": traceback.format_exc()}) + raise + finally: + if dist.is_available() and dist.is_initialized(): + dist.destroy_process_group() + + +def _spawn(worker): + context = mp.get_context("spawn") + result_queue = context.Queue() + directory = tempfile.mkdtemp(prefix="gefen-dtensor-rebinding-") + init_file = os.path.join(directory, "store") + 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=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() + shutil.rmtree(directory, ignore_errors=True) + assert len(results) == _WORLD, (results, [process.exitcode for process in processes]) + assert all(process.exitcode == 0 for process in processes), results + return sorted(results, key=lambda item: item["rank"]) + + +@pytest.mark.skipif( + not dist.is_available() or not dist.is_gloo_available(), + reason="DTensor rebinding requires Gloo", +) +def test_dtensor_rebinding_matches_rank_local_oracles_and_continues_exactly(): + results = _spawn(_valid_worker) + assert all("traceback" not in result for result in results), results + assert all(len(result["cases"]) == 5 for result in results), results + assert any(case["empty"] for result in results for case in result["cases"]), results + + +@pytest.mark.skipif( + not dist.is_available() or not dist.is_gloo_available(), + reason="scoped distributed Muon DTensor coverage requires Gloo", +) +def test_scoped_distributed_muon_native_contract_is_same_topology_and_loads_exactly(): + results = _spawn(_scoped_distributed_muon_contract_worker) + assert all("traceback" not in result for result in results), results + assert all( + result["contract_exact"] + and result["common_state_restored"] + and result["owner_state_restored"] + and result["continued_exactly"] + for result in results + ), results + + +@pytest.mark.skipif( + not dist.is_available() or not dist.is_gloo_available(), + reason="DTensor rebinding requires Gloo", +) +def test_dtensor_rebinding_rejects_invalid_runtime_topology_without_mutation(): + results = _spawn(_rejection_worker) + assert all("traceback" not in result for result in results), results + failures = { + (result["rank"], name): details + for result in results + for name, details in result["results"].items() + if not details["rejected"] or not details["unchanged"] + } + assert not failures, failures + + +@pytest.mark.skipif( + not dist.is_available() or not dist.is_gloo_available(), + reason="scoped DTensor codebook learning requires Gloo", +) +def test_scoped_replicated_dtensor_contributes_once_and_requires_period_agreement(): + results = _spawn(_scoped_replicated_dtensor_worker) + assert all("traceback" not in result for result in results), results + for result in results: + assert result["initialized"], result + assert result["logical_once"], result + assert result["duplicate_would_differ"], result + assert result["agreed"], result + assert result["period_mismatch_rejected"], result + assert result["mismatch_atomic"], result + + +@pytest.mark.skipif( + not dist.is_available() or not dist.is_gloo_available(), + reason="scoped Muon DTensor codebook coverage requires Gloo", +) +def test_scoped_muon_dtensor_codebook_uses_one_logical_input_and_safe_manual_preflights(): + results = _spawn(_scoped_muon_codebook_worker) + assert all("traceback" not in result for result in results), results + for result in results: + assert result["exact_initialized"], result + assert result["exact_logical_once"], result + assert result["duplicate_would_differ"], result + assert result["scoped_checkpoint_hook_rejected"], result + assert result["approximate_initialized"], result + assert result["approximate_local_shards"], result + assert result["period_mismatch_rejected"], result + assert result["period_mismatch_atomic"], result + assert result["initialize_presence_rejected"], result + assert result["initialize_presence_atomic"], result + assert result["refresh_presence_rejected"], result + assert result["refresh_presence_atomic"], result + assert result["mode_mismatch_rejected"], result + assert result["unscoped_mode_mismatch_rejected"], result + + +@pytest.mark.skipif( + not dist.is_available() or not dist.is_gloo_available(), + reason="unscoped Muon DTensor transition coverage requires Gloo", +) +def test_unscoped_muon_post_warmup_mode_transition_rejects_symmetrically_without_mutation(): + results = _spawn(_unscoped_muon_mode_transition_worker) + assert all("traceback" not in result for result in results), results + assert all( + result["initialized"] + and result["warm_global_step"] == 1 + and result["warm_state_initialized"] + and result["asymmetric_rejected"] + and result["asymmetric_parameter_unchanged"] + and result["asymmetric_state_unchanged"] + and result["unanimous_rejected"] + and result["unanimous_parameter_unchanged"] + and result["unanimous_state_unchanged"] + and result["invalid_mode_rejected"] + and result["invalid_mode_unchanged"] + and result["membership_rejected"] + and result["membership_unchanged"] + and result["malformed_rejected"] + and result["malformed_unchanged"] + and result["rebind_rejected"] + and result["rebind_unchanged"] + and result["warm_add_rejected"] + and result["warm_add_preserved_plan"] + and result["add_bypass_transition_rejected"] + and result["add_bypass_unchanged"] + and result["load_refroze_fresh_plan"] + and result["post_load_add_rejected"] + and result["post_load_add_preserved_plan"] + and result["load_plan_unchanged"] + and result["global_step"] == 1 + for result in results + ), results + + +@pytest.mark.skipif( + not dist.is_available() or not dist.is_gloo_available(), + reason="Muon checkpoint routing coverage requires Gloo", +) +def test_muon_checkpoint_rejects_post_warmup_route_drift_before_save_and_load(): + results = _spawn(_muon_checkpoint_routing_worker) + assert all("traceback" not in result for result in results), results + assert all( + result["valid_export"] + and result["frozen_contract_rejected"] + and result["frozen_layout_rejected"] + and result["invalid_contract_rejected"] + and result["valid_contract"] + and result["valid_layout"] + and result["asymmetric_export_rejected"] + and result["asymmetric_export_atomic"] + and result["unanimous_export_rejected"] + and result["unanimous_export_atomic"] + and result["hook_export_rejected"] + and result["hook_export_atomic"] + and result["asymmetric_load_rejected"] + and result["asymmetric_load_atomic"] + and result["unanimous_load_rejected"] + and result["unanimous_load_atomic"] + and result["checkpoint_route_rejected"] + and result["checkpoint_route_atomic"] + and result["valid_load"] + for result in results + ), results + + +@pytest.mark.skipif( + not dist.is_available() or not dist.is_gloo_available(), + reason="Muon checkpoint callback failure coverage requires Gloo", +) +def test_muon_checkpoint_synchronizes_exact_and_distributed_callback_failures(): + results = _spawn(_muon_checkpoint_callback_failure_worker) + assert all("traceback" not in result for result in results), results + for result in results: + assert [case["mode"] for case in result["cases"]] == [ + "exact", + "distributed", + ] + for case in result["cases"]: + assert case["state_pre"], (result["rank"], case) + assert case["state_post"], (result["rank"], case) + assert case["load_pre"], (result["rank"], case) + assert case["load_post"], (result["rank"], case) + assert case["counter_preparation"], (result["rank"], case) + assert case["owner_preparation"], (result["rank"], case) + assert case["valid"], (result["rank"], case) + assert case["precommit_state_atomic"], (result["rank"], case) + assert case["parameter_atomic"], (result["rank"], case) diff --git a/tests/test_dtensor_rebinding_acceptance.py b/tests/test_dtensor_rebinding_acceptance.py new file mode 100644 index 0000000..9b63ea8 --- /dev/null +++ b/tests/test_dtensor_rebinding_acceptance.py @@ -0,0 +1,628 @@ +"""Focused acceptance coverage for the plain Gefen DTensor data plane.""" + +from __future__ import annotations + +import copy +from datetime import timedelta +import multiprocessing as mp +import os +import queue as queue_module +import shutil +import tempfile +import time +import traceback + +import pytest +import torch +import torch.distributed as dist + + +_CANONICAL_STATE_KEYS = { + "name", + "automatic_period", + "step", + "m_codebook", + "m_magnitude", + "vmean", + "vmean_step", + "v_row", + "v_col", + "factored_step", +} + + +def _local(value): + if hasattr(value, "to_local"): + value = value.to_local() + if hasattr(value, "wait"): + value = value.wait() + return value + + +def _raw_bytes(value): + value = _local(value).detach().contiguous() + return value.reshape(-1).view(torch.uint8).clone() + + +def _clone(value): + if torch.is_tensor(value): + return _local(value).detach().clone() + if isinstance(value, dict): + return {key: _clone(item) for key, item in value.items()} + if isinstance(value, list): + return [_clone(item) for item in value] + if isinstance(value, tuple): + return tuple(_clone(item) for item in value) + return copy.deepcopy(value) + + +def _assert_exact(actual, expected, path="value"): + if torch.is_tensor(expected): + assert torch.is_tensor(actual), path + actual = _local(actual) + expected = _local(expected) + assert actual.dtype == expected.dtype, path + assert actual.layout == expected.layout, path + assert tuple(actual.shape) == tuple(expected.shape), path + assert torch.equal(_raw_bytes(actual), _raw_bytes(expected)), path + return + assert type(actual) is type(expected), path + if isinstance(expected, dict): + assert set(actual) == set(expected), (path, tuple(actual), tuple(expected)) + for key in expected: + _assert_exact(actual[key], expected[key], "{}.{}".format(path, key)) + elif isinstance(expected, (list, tuple)): + assert len(actual) == len(expected), path + for index, (actual_item, expected_item) in enumerate(zip(actual, expected)): + _assert_exact(actual_item, expected_item, "{}[{}]".format(path, index)) + else: + assert actual == expected, path + + +def _region(shape, coordinate, parts, dimension=0): + from gefen import LogicalRegion + + global_length = shape[dimension] + chunk = (global_length + parts - 1) // parts + offset = min(coordinate * chunk, global_length) + length = max(0, min(global_length, offset + chunk) - offset) + offsets = [0] * len(shape) + lengths = list(shape) + offsets[dimension] = offset + lengths[dimension] = length + return LogicalRegion(tuple(offsets), tuple(lengths)) + + +def _manifest(fqn, shape, world, *, axis_lengths=None): + from gefen import ( + LogicalRegion, + ParameterIdentity, + ParameterLayout, + PlacementKind, + ProcessGroupIdentity, + ShardIdentity, + ShardPlacement, + ShardingManifest, + ) + + members = tuple("rank:{}".format(rank) for rank in range(world)) + identity = ParameterIdentity(fqn, tuple(shape)) + process_group = ProcessGroupIdentity("data_parallel", members) + regions = [] + if axis_lengths is None: + regions = [_region(shape, rank, world) for rank in range(world)] + else: + cursor = 0 + for length in axis_lengths: + regions.append(LogicalRegion((cursor, 0), (length, shape[1]))) + cursor += length + shards = tuple( + ShardIdentity( + identity, + ParameterLayout.DTENSOR_1D_DEFAULT_WORLD, + regions[rank], + placements=( + ShardPlacement( + "dp", + PlacementKind.DIMENSION_SHARD, + rank, + world, + parameter_dimension=0, + ), + ), + process_group=process_group, + local_member=members[rank], + ) + for rank in range(world) + ) + return ShardingManifest(shards), shards + + +def _optimizer(parameter, *, factored_v_2d=False): + from gefen import Gefen + + return Gefen( + [("legacy.weight", parameter)], + lr=3e-3, + betas=(0.8, 0.97), + eps=2e-8, + weight_decay=0.02, + fused=False, + deterministic=True, + force_2d_period_one=True, + factored_v_2d=factored_v_2d, + ) + + +def _make_finalized(mesh, rank, world, full, *, fqn="Model.Weight", factored_v_2d=False): + from gefen import ParameterRebinding + from torch import nn + from torch.distributed.tensor import Shard, distribute_tensor + + parameter = nn.Parameter(distribute_tensor(full.clone(), mesh, [Shard(0)])) + optimizer = _optimizer(parameter, factored_v_2d=factored_v_2d) + manifest, shards = _manifest(fqn, full.shape, world) + optimizer.post_sharding( + (ParameterRebinding(parameter, parameter, shards[rank]),), + manifest=manifest, + ) + return optimizer, parameter, manifest, shards[rank] + + +def _assign_gradient(parameter, mesh, gradient): + from torch.distributed.tensor import Shard, distribute_tensor + + parameter.grad = distribute_tensor(gradient.clone(), mesh, [Shard(0)]) + + +def _persistent_snapshot(optimizer, parameter): + return { + "global_step": optimizer._gefen_global_step, + "codebook": _clone(optimizer._gefen_codebook), + "parameter": _clone(parameter), + "state": _clone( + { + key: value + for key, value in optimizer.state[parameter].items() + if key in _CANONICAL_STATE_KEYS + } + ), + } + + +def _observable_snapshot(optimizer, parameter): + parameter_state = optimizer.state.get(parameter) + return { + "containers": ( + id(optimizer.__dict__), + id(optimizer.state), + id(optimizer.param_groups), + tuple(id(group) for group in optimizer.param_groups), + tuple(id(group.get("params")) for group in optimizer.param_groups), + id(optimizer._param_names), + id(optimizer._gefen_shard_bindings), + None if parameter_state is None else id(parameter_state), + ), + "parameter": _clone(parameter), + "global_step": optimizer._gefen_global_step, + "codebook": _clone(optimizer._gefen_codebook), + "state": None if parameter_state is None else _clone(dict(parameter_state)), + "groups": tuple( + ( + tuple(id(item) for item in group["params"]), + _clone({key: value for key, value in group.items() if key != "params"}), + ) + for group in optimizer.param_groups + ), + "names": tuple((id(item), value) for item, value in optimizer._param_names.items()), + "bindings": tuple((id(item), value) for item, value in optimizer._gefen_shard_bindings.items()), + "local_bindings": tuple( + (None if item is None else id(item), shard) + for item, shard in optimizer._gefen_local_shard_bindings + ), + "manifest": optimizer._gefen_sharding_manifest, + "finalized": optimizer._gefen_post_sharding_finalized, + } + + +def _zero_tensors(value): + changed = False + if torch.is_tensor(value): + changed = bool(value.numel() and torch.count_nonzero(value).item()) + value.zero_() + return changed + if isinstance(value, dict): + for item in value.values(): + changed = _zero_tensors(item) or changed + return changed + if isinstance(value, (list, tuple)): + for item in value: + changed = _zero_tensors(item) or changed + return changed + return False + + +def _init_process_group(rank, world, init_file): + torch.set_num_threads(1) + dist.init_process_group( + "gloo", + init_method="file://{}".format(init_file), + rank=rank, + world_size=world, + timeout=timedelta(seconds=30), + ) + + +def _finish_process_group(): + if dist.is_available() and dist.is_initialized(): + dist.destroy_process_group() + + +def _routes_worker(rank, world, init_file, checkpoint_dir, result_queue): + del checkpoint_dir + from torch.distributed.tensor import init_device_mesh + + try: + _init_process_group(rank, world, init_file) + mesh = init_device_mesh("cpu", (world,), mesh_dim_names=("dp",)) + full = torch.linspace(-0.75, 0.5, 20, dtype=torch.float32).reshape(5, 4) + gradient = torch.linspace(0.6, -0.4, 20, dtype=torch.float32).reshape(5, 4) + + identity_optimizer, identity_parameter, manifest, shard = _make_finalized(mesh, rank, world, full) + identity_binding = ( + identity_optimizer.param_groups[0]["params"][0] is identity_parameter + and identity_optimizer.shard_identity(identity_parameter) == shard + and identity_optimizer.shard_bindings() == ((identity_parameter, shard),) + and identity_optimizer.sharding_manifest() == manifest + ) + _assign_gradient(identity_parameter, mesh, gradient) + identity_optimizer.step() + identity_stepped = identity_optimizer.state[identity_parameter].get("step") == 1 + + factored_optimizer, factored_parameter, _, _ = _make_finalized( + mesh, + rank, + world, + full.neg(), + fqn="Model.FactoredWeight", + factored_v_2d=True, + ) + _assign_gradient(factored_parameter, mesh, gradient.cos()) + factored_optimizer.step() + local = _local(factored_parameter) + state = factored_optimizer.state[factored_parameter] + factored_uses_block_state = ( + "vmean" in state + and "vmean_step" in state + and tuple(state["vmean"].shape) == (local.numel(), 1) + and state["vmean_step"] == 1 + and "v_row" not in state + and "v_col" not in state + and "factored_step" not in state + ) + result_queue.put( + { + "rank": rank, + "identity_binding": identity_binding, + "identity_stepped": identity_stepped, + "factored_uses_block_state": factored_uses_block_state, + "factored_keys": sorted(key for key in state if not key.startswith("_gefen_rank_local_")), + } + ) + except BaseException: + result_queue.put({"rank": rank, "fatal": traceback.format_exc()}) + finally: + _finish_process_group() + + +def _remote_manifest_worker(rank, world, init_file, checkpoint_dir, result_queue): + del checkpoint_dir + from gefen import ParameterRebinding + from torch import nn + from torch.distributed.tensor import Shard, distribute_tensor, init_device_mesh + + try: + _init_process_group(rank, world, init_file) + mesh = init_device_mesh("cpu", (world,), mesh_dim_names=("dp",)) + full = torch.arange(20, dtype=torch.float32).reshape(5, 4) + parameter = nn.Parameter(distribute_tensor(full.clone(), mesh, [Shard(0)])) + optimizer = _optimizer(parameter) + manifest, shards = _manifest("Model.Weight", full.shape, world, axis_lengths=(2, 1, 2)) + local_descriptor_matches = tuple(_local(parameter).shape) == shards[rank].logical_region.lengths + before = _observable_snapshot(optimizer, parameter) + error = None + try: + optimizer.post_sharding( + (ParameterRebinding(parameter, parameter, shards[rank]),), + manifest=manifest, + ) + except (TypeError, ValueError, RuntimeError) as exc: + error = "{}: {}".format(type(exc).__name__, exc) + unchanged = True + try: + _assert_exact(_observable_snapshot(optimizer, parameter), before, "remote manifest rejection") + except AssertionError: + unchanged = False + result_queue.put( + { + "rank": rank, + "rejected": error is not None, + "error": error, + "unchanged": unchanged, + "local_descriptor_matches": local_descriptor_matches, + } + ) + except BaseException: + result_queue.put({"rank": rank, "fatal": traceback.format_exc()}) + finally: + _finish_process_group() + + +def _checkpoint_worker(rank, world, init_file, checkpoint_dir, result_queue): + import torch.distributed.checkpoint as dcp + from torch.distributed.checkpoint import FileSystemReader, FileSystemWriter + from torch.distributed.checkpoint.state_dict import ( + StateDictOptions, + get_optimizer_state_dict, + set_optimizer_state_dict, + ) + from torch import nn + from torch.distributed.tensor import init_device_mesh + + try: + _init_process_group(rank, world, init_file) + mesh = init_device_mesh("cpu", (world,), mesh_dim_names=("dp",)) + full = torch.linspace(-1.0, 1.0, 20, dtype=torch.float32).reshape(5, 4) + first_gradient = torch.arange(20, dtype=torch.float32).reshape(5, 4).sin() + second_gradient = torch.arange(20, dtype=torch.float32).reshape(5, 4).cos() + source, source_parameter, _, _ = _make_finalized(mesh, rank, world, full) + _assign_gradient(source_parameter, mesh, first_gradient) + source.step() + native_state = source.state_dict() + + current_full = source_parameter.detach().full_tensor().clone() + mismatch, mismatch_parameter, _, _ = _make_finalized( + mesh, + rank, + world, + current_full, + fqn="Model.DifferentWeight", + ) + before_mismatch = _observable_snapshot(mismatch, mismatch_parameter) + mismatch_error = None + try: + mismatch.load_state_dict(native_state) + except (TypeError, ValueError, RuntimeError) as exc: + mismatch_error = "{}: {}".format(type(exc).__name__, exc) + mismatch_unchanged = True + try: + _assert_exact( + _observable_snapshot(mismatch, mismatch_parameter), + before_mismatch, + "stable identity mismatch", + ) + except AssertionError: + mismatch_unchanged = False + + source_model = nn.Module() + source_model.register_parameter("weight", source_parameter) + public_state = get_optimizer_state_dict( + source_model, + source, + options=StateDictOptions(full_state_dict=True, cpu_offload=True), + ) + public_handoff = [public_state if rank == 0 else None] + dist.broadcast_object_list(public_handoff, src=0) + persisted = {"optimizer": public_handoff[0]} + metadata = dcp.save( + state_dict=persisted, + storage_writer=FileSystemWriter(checkpoint_dir), + process_group=dist.group.WORLD, + ) + on_disk_metadata = FileSystemReader(checkpoint_dir).read_metadata() + load_target = _clone(persisted) + zeroing_changed_payload = _zero_tensors(load_target) + dcp.load( + state_dict=load_target, + storage_reader=FileSystemReader(checkpoint_dir), + process_group=dist.group.WORLD, + ) + _assert_exact(load_target, persisted, "filesystem DCP round trip") + + resumed, resumed_parameter, _, _ = _make_finalized(mesh, rank, world, current_full) + resumed_model = nn.Module() + resumed_model.register_parameter("weight", resumed_parameter) + set_optimizer_state_dict( + resumed_model, + resumed, + load_target["optimizer"] if rank == 0 else {}, + options=StateDictOptions(full_state_dict=True, broadcast_from_rank0=True), + ) + restored = True + try: + _assert_exact( + _persistent_snapshot(resumed, resumed_parameter), + _persistent_snapshot(source, source_parameter), + "DCP restored optimizer", + ) + except AssertionError: + restored = False + + _assign_gradient(source_parameter, mesh, second_gradient) + _assign_gradient(resumed_parameter, mesh, second_gradient) + source.step() + resumed.step() + continued = True + try: + _assert_exact( + _persistent_snapshot(resumed, resumed_parameter), + _persistent_snapshot(source, source_parameter), + "DCP exact continuation", + ) + except AssertionError: + continued = False + result_queue.put( + { + "rank": rank, + "mismatch_rejected": mismatch_error is not None, + "mismatch_error": mismatch_error, + "mismatch_unchanged": mismatch_unchanged, + "dcp_metadata": bool(metadata.state_dict_metadata), + "reader_metadata": set(on_disk_metadata.state_dict_metadata) == set(metadata.state_dict_metadata), + "zeroing_changed_payload": zeroing_changed_payload, + "restored": restored, + "continued": continued, + } + ) + except BaseException: + result_queue.put({"rank": rank, "fatal": traceback.format_exc()}) + finally: + _finish_process_group() + + +def _checkpoint_failure_worker(rank, world, init_file, checkpoint_dir, result_queue): + del checkpoint_dir + from torch import nn + from torch.distributed.tensor import init_device_mesh + + try: + _init_process_group(rank, world, init_file) + mesh = init_device_mesh("cpu", (world,), mesh_dim_names=("dp",)) + full = torch.arange(16, dtype=torch.float32).reshape(4, 4) + optimizer, parameter, _, _ = _make_finalized(mesh, rank, world, full) + + def fail_pre_hook(_optimizer): + raise RuntimeError("injected rank-zero state-dict pre-hook failure") + + handle = optimizer.register_state_dict_pre_hook(fail_pre_hook) if rank == 0 else None + pre_hook_error = None + try: + optimizer.state_dict() + except RuntimeError as exc: + pre_hook_error = str(exc) + if handle is not None: + handle.remove() + dist.barrier() + + original_parameter = optimizer.param_groups[0]["params"][0] + if rank == 0: + optimizer.param_groups[0]["params"][0] = nn.Parameter(torch.zeros_like(_local(parameter))) + layout_error = None + try: + optimizer.state_dict() + except RuntimeError as exc: + layout_error = str(exc) + if rank == 0: + optimizer.param_groups[0]["params"][0] = original_parameter + result_queue.put( + { + "rank": rank, + "pre_hook_error": pre_hook_error, + "layout_error": layout_error, + } + ) + except BaseException: + result_queue.put({"rank": rank, "fatal": traceback.format_exc()}) + finally: + _finish_process_group() + + +def _spawn(worker, world, *, timeout): + context = mp.get_context("spawn") + result_queue = context.Queue() + directory = tempfile.mkdtemp(prefix="gefen-dtensor-acceptance-") + init_file = os.path.join(directory, "store") + checkpoint_dir = os.path.join(directory, "checkpoint") + processes = [ + context.Process( + target=worker, + args=(rank, world, init_file, checkpoint_dir, result_queue), + ) + for rank in range(world) + ] + results = [] + deadline = time.monotonic() + timeout + try: + for process in processes: + process.start() + while len(results) < world: + remaining = deadline - time.monotonic() + if remaining <= 0: + break + try: + results.append(result_queue.get(timeout=remaining)) + except queue_module.Empty: + break + for process in processes: + process.join(timeout=max(0.0, min(5.0, deadline - time.monotonic()))) + 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() + shutil.rmtree(directory, ignore_errors=True) + assert len(results) == world, (results, [process.exitcode for process in processes]) + assert all(process.exitcode == 0 for process in processes), (results, [process.exitcode for process in processes]) + results = sorted(results, key=lambda item: item["rank"]) + assert all("fatal" not in result for result in results), results + return results + + +@pytest.fixture(scope="module") +def _route_results(): + return _spawn(_routes_worker, 2, timeout=60) + + +@pytest.fixture(scope="module") +def _checkpoint_results(): + return _spawn(_checkpoint_worker, 2, timeout=90) + + +@pytest.mark.skipif(not dist.is_available() or not dist.is_gloo_available(), reason="DTensor acceptance coverage requires Gloo") +def test_identity_only_rebinding_finalizes_an_existing_dtensor_slot(_route_results): + assert all(result["identity_binding"] and result["identity_stepped"] for result in _route_results), _route_results + + +@pytest.mark.skipif(not dist.is_available() or not dist.is_gloo_available(), reason="DTensor acceptance coverage requires Gloo") +def test_factored_v_2d_routes_sharded_matrices_to_local_block_vmean(_route_results): + assert all(result["factored_uses_block_state"] for result in _route_results), _route_results + + +@pytest.mark.skipif(not dist.is_available() or not dist.is_gloo_available(), reason="DTensor acceptance coverage requires Gloo") +def test_malformed_remote_manifest_region_fails_locally_before_mutation(): + results = _spawn(_remote_manifest_worker, 3, timeout=60) + assert results[0]["local_descriptor_matches"], results + assert all(result["rejected"] and result["unchanged"] for result in results), results + + +@pytest.mark.skipif(not dist.is_available() or not dist.is_gloo_available(), reason="DTensor acceptance coverage requires Gloo") +def test_rank_local_load_requires_exact_stable_binding(_checkpoint_results): + assert all(result["mismatch_rejected"] and result["mismatch_unchanged"] for result in _checkpoint_results), _checkpoint_results + + +@pytest.mark.skipif(not dist.is_available() or not dist.is_gloo_available(), reason="DTensor acceptance coverage requires Gloo") +def test_rank_local_state_dict_failures_raise_symmetrically_without_hanging(): + results = _spawn(_checkpoint_failure_worker, 2, timeout=45) + assert all(result["pre_hook_error"] is not None and result["layout_error"] is not None for result in results), results + assert "failed on this rank" in results[0]["pre_hook_error"], results + assert "failed on another rank" in results[1]["pre_hook_error"], results + assert "failed on this rank" in results[0]["layout_error"], results + assert "failed on another rank" in results[1]["layout_error"], results + + +@pytest.mark.skipif(not dist.is_available() or not dist.is_gloo_available(), reason="DTensor acceptance coverage requires Gloo") +def test_rank_local_checkpoint_persists_through_filesystem_dcp(_checkpoint_results): + assert all( + result["dcp_metadata"] + and result["reader_metadata"] + and result["zeroing_changed_payload"] + and result["restored"] + and result["continued"] + for result in _checkpoint_results + ), _checkpoint_results diff --git a/tests/test_gefen_fsdp2_checkpoint.py b/tests/test_gefen_fsdp2_checkpoint.py index 2fcbbb0..7586afd 100644 --- a/tests/test_gefen_fsdp2_checkpoint.py +++ b/tests/test_gefen_fsdp2_checkpoint.py @@ -736,3 +736,186 @@ def test_rank_local_full_dcp_set_optimizer_state_is_exact_under_fully_shard( assert checks is not None, "fully_shard checkpoint workers timed out" assert all(process.exitcode == 0 for process in processes) assert all(all(rank_check) for rank_check in checks), checks + + +def _fully_shard_rebinding_worker( + rank: int, world: int, port: str, fused: bool, result_queue +) -> None: + import torch.distributed as dist + import torch.nn as nn + from torch.distributed.checkpoint.state_dict import ( + StateDictOptions, + get_optimizer_state_dict, + set_optimizer_state_dict, + ) + from torch.distributed.device_mesh import init_device_mesh + from torch.distributed.fsdp import fully_shard + from torch.distributed.tensor import Shard, distribute_tensor + + from gefen import ( + Gefen, + LogicalRegion, + ParameterIdentity, + ParameterLayout, + ParameterRebinding, + PlacementKind, + ProcessGroupIdentity, + ShardIdentity, + ShardPlacement, + ShardingManifest, + ) + + os.environ["MASTER_ADDR"] = "127.0.0.1" + os.environ["MASTER_PORT"] = port + os.environ["RANK"] = str(rank) + os.environ["WORLD_SIZE"] = str(world) + torch.cuda.set_device(rank) + dist.init_process_group("nccl", rank=rank, world_size=world) + try: + mesh = init_device_mesh("cuda", (world,), mesh_dim_names=("dp",)) + dtype = torch.bfloat16 if fused else torch.float32 + full = torch.linspace(-1, 1, 64, device="cuda", dtype=dtype).reshape(8, 8) + first_grad = torch.arange(64, device="cuda", dtype=torch.float32).sin().reshape(8, 8).to(dtype) + second_grad = torch.arange(64, device="cuda", dtype=torch.float32).cos().reshape(8, 8).to(dtype) + + identity = ParameterIdentity("Model.Weight", (8, 8)) + group = ProcessGroupIdentity( + "data_parallel", tuple("rank:{}".format(item) for item in range(world)) + ) + chunk = (identity.global_shape[0] + world - 1) // world + shards = [] + for coordinate, member in enumerate(group.ordered_members): + offset = min(coordinate * chunk, identity.global_shape[0]) + length = max( + 0, + min(identity.global_shape[0], offset + chunk) - offset, + ) + shards.append( + ShardIdentity( + identity, + ParameterLayout.DTENSOR_1D_DEFAULT_WORLD, + LogicalRegion((offset, 0), (length, identity.global_shape[1])), + placements=( + ShardPlacement( + "dp", + PlacementKind.DIMENSION_SHARD, + coordinate, + world, + parameter_dimension=0, + ), + ), + process_group=group, + local_member=member, + ) + ) + manifest = ShardingManifest(tuple(shards)) + + def construct(values): + model = nn.Linear(8, 8, bias=False, device="cuda", dtype=dtype) + with torch.no_grad(): + model.weight.copy_(values) + old = model.weight + optimizer = Gefen( + model.named_parameters(), + lr=1e-3, + fused=fused, + deterministic=True, + factored_v_2d=False, + force_2d_period_one=True, + ) + fully_shard(model, mesh=mesh) + optimizer.post_sharding( + (ParameterRebinding(old, model.weight, shards[rank]),), + manifest=manifest, + ) + return model, optimizer + + model, optimizer = construct(full) + identity_ok = ( + optimizer.shard_identity(model.weight) == shards[rank] + and optimizer.sharding_manifest() == manifest + ) + model.weight.grad = distribute_tensor(first_grad, mesh, [Shard(0)]) + optimizer.step() + current = model.weight.detach().full_tensor().clone() + checkpoint = get_optimizer_state_dict( + model, + optimizer, + options=StateDictOptions(full_state_dict=True, cpu_offload=True), + ) + + resumed_model, resumed = construct(current) + set_optimizer_state_dict( + resumed_model, + resumed, + checkpoint if rank == 0 else {}, + options=StateDictOptions( + full_state_dict=True, + broadcast_from_rank0=True, + ), + ) + restored = ( + torch.equal( + resumed._gefen_codebook.detach().cpu(), + optimizer._gefen_codebook.detach().cpu(), + ) + and torch.equal( + resumed.state[resumed_model.weight]["m_codebook"].detach().cpu(), + optimizer.state[model.weight]["m_codebook"].detach().cpu(), + ) + ) + model.weight.grad = distribute_tensor(second_grad, mesh, [Shard(0)]) + resumed_model.weight.grad = distribute_tensor(second_grad, mesh, [Shard(0)]) + optimizer.step() + resumed.step() + continued = torch.equal( + model.weight.detach().full_tensor(), + resumed_model.weight.detach().full_tensor(), + ) + checks = [None] * world + dist.all_gather_object(checks, (identity_ok, restored, continued)) + if rank == 0: + result_queue.put(checks) + except BaseException: + result_queue.put([{"rank": rank, "traceback": traceback.format_exc()}]) + raise + finally: + dist.destroy_process_group() + + +@pytest.mark.skipif( + not torch.cuda.is_available() + or torch.cuda.device_count() < 2 + or not torch.distributed.is_nccl_available(), + reason="fully_shard DTensor rebinding requires two CUDA GPUs and NCCL", +) +@pytest.mark.parametrize("fused", [False, True]) +def test_post_sharding_rebinding_and_dcp_are_exact_under_fully_shard(fused): + import torch.multiprocessing as mp + + context = mp.get_context("spawn") + result_queue = context.Queue() + port = _free_port() + processes = [ + context.Process( + target=_fully_shard_rebinding_worker, + args=(rank, 2, port, fused, result_queue), + ) + for rank in range(2) + ] + for process in processes: + process.start() + try: + checks = result_queue.get(timeout=300) + except queue.Empty: + checks = None + for process in processes: + process.join(timeout=300) + for process in processes: + if process.is_alive(): + process.terminate() + process.join(timeout=10) + assert checks is not None, "fully_shard rebinding workers timed out" + assert all(process.exitcode == 0 for process in processes) + assert all(isinstance(rank_check, tuple) for rank_check in checks), checks + assert all(all(rank_check) for rank_check in checks), checks diff --git a/tests/test_hybrid_adamw_dtensor.py b/tests/test_hybrid_adamw_dtensor.py new file mode 100644 index 0000000..05aed42 --- /dev/null +++ b/tests/test_hybrid_adamw_dtensor.py @@ -0,0 +1,235 @@ +"""Two-rank same-topology Hybrid rebinding with a DTensor AdamW backup.""" + +from datetime import timedelta +import copy +import multiprocessing as mp +import os +import queue as queue_module +import shutil +import tempfile +import traceback + +import pytest +import torch +import torch.distributed as dist + + +_WORLD = 2 +_MEMBERS = ("rank:0", "rank:1") +_DTYPES = ( + ("float16", torch.float16), + ("bfloat16", torch.bfloat16), + ("float32", torch.float32), + ("float64", torch.float64), +) + + +def _region(shape, coordinate, dimension): + from gefen import LogicalRegion + + global_length = shape[dimension] + chunk = (global_length + _WORLD - 1) // _WORLD + offset = min(coordinate * chunk, global_length) + length = max(0, min(global_length, offset + chunk) - offset) + offsets = [0] * len(shape) + lengths = list(shape) + offsets[dimension] = offset + lengths[dimension] = length + return LogicalRegion(tuple(offsets), tuple(lengths)) + + +def _manifest(): + from gefen import ( + ParameterIdentity, + ParameterLayout, + PlacementKind, + ProcessGroupIdentity, + ShardIdentity, + ShardPlacement, + ShardingManifest, + ) + + group = ProcessGroupIdentity("data_parallel", _MEMBERS) + shards = [] + by_fqn = {} + for fqn, shape in (("Model.Bias", (4,)),): + identity = ParameterIdentity(fqn, shape) + parameter_shards = [] + for coordinate, member in enumerate(_MEMBERS): + shard = ShardIdentity( + identity, + ParameterLayout.DTENSOR_1D_DEFAULT_WORLD, + _region(shape, coordinate, 0), + placements=( + ShardPlacement( + "dp", + PlacementKind.DIMENSION_SHARD, + coordinate, + _WORLD, + parameter_dimension=0, + ), + ), + process_group=group, + local_member=member, + ) + shards.append(shard) + parameter_shards.append(shard) + by_fqn[fqn] = tuple(parameter_shards) + return ShardingManifest(tuple(shards)), by_fqn + + +def _optimizer(old_bias): + from gefen import GefenMuonHybrid + + return GefenMuonHybrid( + [], + [("layer.bias", old_bias)], + lr=1e-3, + fused=False, + backup_optimizer="adamw", + ) + + +def _local(value): + value = value.to_local() + return value.wait() if hasattr(value, "wait") else value + + +def _state_exact(left, left_parameter, right, right_parameter): + left_state = left.backup.state[left_parameter] + right_state = right.backup.state[right_parameter] + if set(left_state) != set(right_state): + return False + for key in left_state: + left_value = left_state[key] + right_value = right_state[key] + if torch.is_tensor(left_value): + if hasattr(left_value, "to_local"): + left_value = _local(left_value) + right_value = _local(right_value) + if left_value.dtype != right_value.dtype or not torch.equal(left_value, right_value): + return False + elif left_value != right_value: + return False + return True + + +def _exercise(rank, mesh, dtype): + from torch import nn + from torch.distributed.tensor import Shard, distribute_tensor + + from gefen import ParameterRebinding + + full_bias = torch.arange(4, dtype=dtype) + manifest, shards = _manifest() + + old_bias = nn.Parameter(full_bias.clone()) + bias = nn.Parameter(distribute_tensor(full_bias.clone(), mesh, [Shard(0)])) + optimizer = _optimizer(old_bias) + optimizer.post_sharding( + ( + ParameterRebinding(old_bias, bias, shards["Model.Bias"][rank]), + ), + manifest=manifest, + ) + + first_gradient = torch.tensor([0.5, -0.25, 0.75, -1.0], dtype=dtype) + bias.grad = distribute_tensor(first_gradient, mesh, [Shard(0)]) + optimizer.step() + optimizer.zero_grad() + checkpoint = copy.deepcopy(optimizer.state_dict()) + current_bias = bias.detach().full_tensor().clone() + + resumed_old_bias = nn.Parameter(current_bias.clone()) + resumed_bias = nn.Parameter(distribute_tensor(current_bias, mesh, [Shard(0)])) + resumed = _optimizer(resumed_old_bias) + resumed.post_sharding( + ( + ParameterRebinding(resumed_old_bias, resumed_bias, shards["Model.Bias"][rank]), + ), + manifest=manifest, + ) + resumed.load_state_dict(checkpoint) + assert _state_exact(optimizer, bias, resumed, resumed_bias) + + second_gradient = torch.tensor([-0.5, 0.125, 0.25, 0.875], dtype=dtype) + bias.grad = distribute_tensor(second_gradient, mesh, [Shard(0)]) + resumed_bias.grad = distribute_tensor(second_gradient, mesh, [Shard(0)]) + optimizer.step() + resumed.step() + assert torch.equal(_local(bias.detach()), _local(resumed_bias.detach())) + assert _state_exact(optimizer, bias, resumed, resumed_bias) + + +def _worker(rank, init_file, result_queue): + from torch.distributed.tensor import init_device_mesh + + try: + torch.set_num_threads(1) + dist.init_process_group( + "gloo", + init_method="file://{}".format(init_file), + rank=rank, + world_size=_WORLD, + timeout=timedelta(seconds=60), + ) + mesh = init_device_mesh("cpu", (_WORLD,), mesh_dim_names=("dp",)) + completed = [] + for dtype_name, dtype in _DTYPES: + _exercise(rank, mesh, dtype) + completed.append(dtype_name) + result_queue.put({"rank": rank, "completed": tuple(completed)}) + except BaseException: + result_queue.put({"rank": rank, "traceback": traceback.format_exc()}) + raise + finally: + if dist.is_available() and dist.is_initialized(): + dist.destroy_process_group() + + +def _run(): + context = mp.get_context("spawn") + result_queue = context.Queue() + directory = tempfile.mkdtemp(prefix="gefen-hybrid-adamw-dtensor-") + init_file = os.path.join(directory, "store") + 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=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() + shutil.rmtree(directory, ignore_errors=True) + assert len(results) == _WORLD, (results, [process.exitcode for process in processes]) + assert all(process.exitcode == 0 for process in processes), results + return results + + +@pytest.mark.skipif( + not dist.is_available() or not dist.is_gloo_available(), + reason="DTensor Hybrid rebinding requires Gloo", +) +def test_hybrid_adamw_dtensor_rebinding_and_native_continuation_are_exact(): + results = _run() + assert all("traceback" not in result for result in results), results + expected = tuple(name for name, _dtype in _DTYPES) + assert all(result["completed"] == expected for result in results), results diff --git a/tests/test_hybrid_adamw_layouts_distributed.py b/tests/test_hybrid_adamw_layouts_distributed.py new file mode 100644 index 0000000..df1e011 --- /dev/null +++ b/tests/test_hybrid_adamw_layouts_distributed.py @@ -0,0 +1,420 @@ +"""Two-rank AdamW Hybrid coverage for flattened and whole-owner layouts.""" + +from datetime import timedelta +import copy +import multiprocessing as mp +import os +import queue as queue_module +import shutil +import tempfile +import traceback + +import pytest +import torch +import torch.distributed as dist + +from _state_snapshot import assert_deep_state_snapshot, deep_state_snapshot + + +_WORLD = 2 +_MEMBERS = ("rank:0", "rank:1") +_OWNER = "rank:1" +_SHAPE = (2, 3) +_DTYPES = ( + ("float16", torch.float16), + ("bfloat16", torch.bfloat16), + ("float32", torch.float32), + ("float64", torch.float64), +) + + +def _manifest(layout_name): + from gefen import ( + LogicalSlice, + ParameterIdentity, + ParameterLayout, + PlacementKind, + ProcessGroupIdentity, + ShardIdentity, + ShardPlacement, + ShardingManifest, + ) + + group = ProcessGroupIdentity("data_parallel", _MEMBERS) + identity = ParameterIdentity("Model.Backup.Weight", _SHAPE) + shards = [] + for coordinate, member in enumerate(_MEMBERS): + if layout_name == "flattened": + layout = ParameterLayout.FLATTENED_ELEMENT_SHARD + logical_slice = LogicalSlice(coordinate * 3, 3) + placement_kind = PlacementKind.FLAT_SHARD + owner = None + else: + layout = ParameterLayout.WHOLE_PARAMETER_OWNER + logical_slice = LogicalSlice.full(identity) if member == _OWNER else LogicalSlice(0, 0) + placement_kind = PlacementKind.WHOLE_PARAMETER_OWNER + owner = _OWNER + shards.append( + ShardIdentity( + identity, + layout, + logical_slice, + placements=( + ShardPlacement( + "dp", + placement_kind, + coordinate, + _WORLD, + ), + ), + process_group=group, + local_member=member, + owner=owner, + ) + ) + return ShardingManifest(tuple(shards)), tuple(shards) + + +def _optimizer(parameter): + from gefen import GefenMuonHybrid + + return GefenMuonHybrid( + [], + [("backup.weight", parameter)], + lr=2e-3, + betas=(0.8, 0.95), + eps=1e-6, + weight_decay=0.03, + fused=False, + backup_optimizer="adamw", + ) + + +def _reference(parameter): + return torch.optim.AdamW( + [parameter], + lr=2e-3, + betas=(0.8, 0.95), + eps=1e-6, + weight_decay=0.03, + fused=False, + ) + + +def _build(rank, full_value, layout_name): + from torch import nn + + from gefen import ParameterRebinding + + manifest, shards = _manifest(layout_name) + shard = shards[rank] + old_parameter = nn.Parameter(full_value.clone()) + if layout_name == "flattened": + start = shard.logical_slice.flat_offset + target = nn.Parameter(full_value.reshape(-1)[start : start + shard.logical_slice.length].clone()) + elif _MEMBERS[rank] == _OWNER: + target = nn.Parameter(full_value.clone()) + else: + target = None + optimizer = _optimizer(old_parameter) + optimizer.post_sharding( + (ParameterRebinding(old_parameter, target, shard),), + manifest=manifest, + ) + return optimizer, target, shard + + +def _training_support(contract, layout): + matches = [item for item in contract.capabilities.training if item.layout is layout] + assert len(matches) == 1 + return matches[0] + + +def _assert_contract(optimizer, layout_name): + from gefen import ( + CheckpointTransport, + ParameterLayout, + Precision, + ProcessGroupScope, + ) + + contract = optimizer.optimizer_contract() + assert contract.capabilities.precisions == frozenset( + {Precision.FLOAT16, Precision.BFLOAT16, Precision.FLOAT32, Precision.FLOAT64} + ) + assert {item.layout for item in contract.capabilities.training} == { + ParameterLayout.REPLICATED, + ParameterLayout.FLATTENED_ELEMENT_SHARD, + ParameterLayout.WHOLE_PARAMETER_OWNER, + ParameterLayout.DTENSOR_1D_DEFAULT_WORLD, + } + flattened = _training_support(contract, ParameterLayout.FLATTENED_ELEMENT_SHARD) + assert flattened.process_group_scope is ProcessGroupScope.NONE + assert not flattened.requires_complete_parameter_storage + assert not flattened.requires_complete_logical_matrix + assert not flattened.requires_post_step_parameter_sync + owner = _training_support(contract, ParameterLayout.WHOLE_PARAMETER_OWNER) + assert owner.process_group_scope is ProcessGroupScope.ADAPTER_DEFINED + assert owner.requires_complete_parameter_storage + assert not owner.requires_complete_logical_matrix + assert owner.requires_post_step_parameter_sync + + expected_layout = ( + ParameterLayout.FLATTENED_ELEMENT_SHARD + if layout_name == "flattened" + else ParameterLayout.WHOLE_PARAMETER_OWNER + ) + native = [ + item + for item in contract.capabilities.checkpoints + if item.transport is CheckpointTransport.COMPOSITE_NATIVE + ] + assert len(native) == 1 + assert native[0].same_topology == frozenset({ParameterLayout.REPLICATED, expected_layout}) + assert not native[0].topology_changing + assert native[0].process_group_scope is ProcessGroupScope.NONE + assert not native[0].requires_collective + assert native[0].atomic_load + assert all( + item.transport is not CheckpointTransport.CANONICAL_GLOBAL + for item in contract.capabilities.checkpoints + ) + + +def _assign_gradient(target, full_gradient, shard): + if target is None: + return + if target.ndim == 1: + start = shard.logical_slice.flat_offset + target.grad = full_gradient.reshape(-1)[start : start + shard.logical_slice.length].clone() + else: + target.grad = full_gradient.clone() + + +def _step(optimizer, target, full_gradient, shard): + _assign_gradient(target, full_gradient, shard) + optimizer.step() + optimizer.zero_grad() + + +def _tree_exact(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 left.layout == right.layout + and tuple(left.shape) == tuple(right.shape) + 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(_tree_exact(left[key], right[key]) for key in left) + if isinstance(left, (tuple, list)): + return len(left) == len(right) and all( + _tree_exact(left_item, right_item) + for left_item, right_item in zip(left, right) + ) + return left == right + + +def _assert_reference_exact(optimizer, target, shard, reference_optimizer, reference_parameter): + if target is None: + assert optimizer.backup.param_groups[0]["params"] == [] + assert not optimizer.backup.state + return + reference_value = reference_parameter.detach() + if target.ndim == 1: + start = shard.logical_slice.flat_offset + reference_value = reference_value.reshape(-1)[start : start + shard.logical_slice.length] + assert torch.equal(target.detach(), reference_value) + + state = optimizer.backup.state[target] + reference_state = reference_optimizer.state[reference_parameter] + assert torch.equal(state["step"], reference_state["step"]) + for name in ("exp_avg", "exp_avg_sq"): + expected = reference_state[name] + if target.ndim == 1: + start = shard.logical_slice.flat_offset + expected = expected.reshape(-1)[start : start + shard.logical_slice.length] + assert torch.equal(state[name], expected) + + +def _corrupt_checkpoint(checkpoint, target): + invalid = copy.deepcopy(checkpoint) + backup = invalid["backup"] + if target is None: + backup["param_groups"].append(copy.deepcopy(backup["param_groups"][0])) + else: + (parameter_id,) = backup["state"] + backup["state"][parameter_id]["exp_avg"] = torch.zeros(1, dtype=target.dtype) + return invalid + + +def _exercise_case(rank, dtype, layout_name): + full_value = torch.linspace(-0.75, 0.625, 6, dtype=dtype).reshape(_SHAPE) + first_gradient = torch.tensor( + [[0.5, -0.25, 0.75], [-1.0, 0.375, -0.625]], + dtype=dtype, + ) + second_gradient = torch.tensor( + [[-0.5, 0.125, 0.25], [0.875, -0.75, 0.625]], + dtype=dtype, + ) + + optimizer, target, shard = _build(rank, full_value, layout_name) + _assert_contract(optimizer, layout_name) + reference_parameter = torch.nn.Parameter(full_value.clone()) + reference_optimizer = _reference(reference_parameter) + + _step(optimizer, target, first_gradient, shard) + reference_parameter.grad = first_gradient.clone() + reference_optimizer.step() + reference_optimizer.zero_grad() + _assert_reference_exact( + optimizer, + target, + shard, + reference_optimizer, + reference_parameter, + ) + checkpoint = copy.deepcopy(optimizer.state_dict()) + + resumed, resumed_target, resumed_shard = _build( + rank, + reference_parameter.detach(), + layout_name, + ) + resumed.load_state_dict(checkpoint) + assert _tree_exact(resumed.state_dict(), optimizer.state_dict()) + if target is not None: + assert torch.equal(resumed_target, target) + + _step(optimizer, target, second_gradient, shard) + _step(resumed, resumed_target, second_gradient, resumed_shard) + reference_parameter.grad = second_gradient.clone() + reference_optimizer.step() + reference_optimizer.zero_grad() + assert _tree_exact(resumed.state_dict(), optimizer.state_dict()) + if target is not None: + assert torch.equal(resumed_target, target) + _assert_reference_exact( + optimizer, + target, + shard, + reference_optimizer, + reference_parameter, + ) + _assert_reference_exact( + resumed, + resumed_target, + resumed_shard, + reference_optimizer, + reference_parameter, + ) + + invalid = _corrupt_checkpoint(checkpoint, resumed_target) + child_snapshot = deep_state_snapshot(resumed.backup) + state_before = copy.deepcopy(resumed.state_dict()) + target_before = resumed_target.detach().clone() if resumed_target is not None else None + finalized_before = ( + resumed._hybrid_sharding_manifest, + resumed._hybrid_local_shard_bindings, + resumed._hybrid_shard_bindings, + resumed._hybrid_fqn_roles, + resumed._hybrid_finalized_slots, + ) + try: + resumed.load_state_dict(invalid) + except (RuntimeError, TypeError, ValueError): + pass + else: + raise AssertionError("corrupt AdamW checkpoint was accepted") + assert_deep_state_snapshot(resumed.backup, child_snapshot) + assert _tree_exact(resumed.state_dict(), state_before) + if resumed_target is not None: + assert torch.equal(resumed_target, target_before) + finalized_after = ( + resumed._hybrid_sharding_manifest, + resumed._hybrid_local_shard_bindings, + resumed._hybrid_shard_bindings, + resumed._hybrid_fqn_roles, + resumed._hybrid_finalized_slots, + ) + assert all(after is before for after, before in zip(finalized_after, finalized_before)) + + +def _worker(rank, init_file, layout_name, result_queue): + try: + torch.set_num_threads(1) + dist.init_process_group( + "gloo", + init_method="file://{}".format(init_file), + rank=rank, + world_size=_WORLD, + timeout=timedelta(seconds=120), + ) + completed = [] + for dtype_name, dtype in _DTYPES: + _exercise_case(rank, dtype, layout_name) + completed.append(dtype_name) + result_queue.put({"rank": rank, "completed": tuple(completed)}) + except BaseException: + result_queue.put({"rank": rank, "traceback": traceback.format_exc()}) + raise + finally: + if dist.is_available() and dist.is_initialized(): + dist.destroy_process_group() + + +def _run(layout_name): + context = mp.get_context("spawn") + result_queue = context.Queue() + directory = tempfile.mkdtemp(prefix="gefen-hybrid-adamw-{}-".format(layout_name)) + init_file = os.path.join(directory, "store") + processes = [ + context.Process( + target=_worker, + args=(rank, init_file, layout_name, 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() + shutil.rmtree(directory, ignore_errors=True) + assert len(results) == _WORLD, (results, [process.exitcode for process in processes]) + assert all(process.exitcode == 0 for process in processes), results + return results + + +@pytest.mark.skipif( + not dist.is_available() or not dist.is_gloo_available(), + reason="distributed Hybrid AdamW rebinding requires Gloo", +) +@pytest.mark.parametrize("layout_name", ["flattened", "whole_owner"]) +def test_hybrid_adamw_sharded_native_continuation_is_exact_and_atomic(layout_name): + results = _run(layout_name) + assert all("traceback" not in result for result in results), results + expected = tuple(name for name, _dtype in _DTYPES) + assert all(result["completed"] == expected for result in results), results diff --git a/tests/test_hybrid_adamw_staging.py b/tests/test_hybrid_adamw_staging.py new file mode 100644 index 0000000..3f22001 --- /dev/null +++ b/tests/test_hybrid_adamw_staging.py @@ -0,0 +1,421 @@ +"""Focused stable-binding and native-load staging tests for Hybrid AdamW.""" + +import copy + +import pytest +import torch + +from gefen import GefenMuonHybrid +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 + + +def _shard(fqn, shape): + identity = ParameterIdentity(fqn, shape) + return ShardIdentity( + identity, + ParameterLayout.REPLICATED, + LogicalSlice.full(identity), + ) + + +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) + 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 _finalized(*, backup_optimizer="adamw", prefix="Model", matrix=None, bias=None): + old_matrix = torch.nn.Parameter(torch.arange(4, dtype=torch.float32).reshape(2, 2)) + old_bias = torch.nn.Parameter(torch.arange(4, dtype=torch.float32)) + optimizer = GefenMuonHybrid( + [("layer.weight", old_matrix)], + [("layer.bias", old_bias)], + lr=1e-3, + fused=False, + backup_optimizer=backup_optimizer, + ) + matrix = torch.nn.Parameter(old_matrix.detach().clone() if matrix is None else matrix.detach().clone()) + bias = torch.nn.Parameter(old_bias.detach().clone() if bias is None else bias.detach().clone()) + matrix_shard = _shard(prefix + ".Layer.Weight", (2, 2)) + bias_shard = _shard(prefix + ".Layer.Bias", (4,)) + manifest = ShardingManifest((matrix_shard, bias_shard)) + optimizer.post_sharding( + ( + ParameterRebinding(old_matrix, matrix, matrix_shard), + ParameterRebinding(old_bias, bias, bias_shard), + ), + manifest=manifest, + ) + return optimizer, matrix, bias + + +def _unfinalized(*, backup_optimizer="adamw", matrix=None, bias=None): + matrix = torch.nn.Parameter( + torch.arange(4, dtype=torch.float32).reshape(2, 2) + if matrix is None + else matrix.detach().clone() + ) + bias = torch.nn.Parameter( + torch.arange(4, dtype=torch.float32) + if bias is None + else bias.detach().clone() + ) + return ( + GefenMuonHybrid( + [("layer.weight", matrix)], + [("layer.bias", bias)], + lr=1e-3, + fused=False, + backup_optimizer=backup_optimizer, + ), + matrix, + bias, + ) + + +def _step_backup(optimizer, bias, gradient): + optimizer.zero_grad() + bias.grad = gradient.clone() + optimizer.step() + optimizer.zero_grad() + + +def _assert_nested_exact(actual, expected, path="value"): + assert type(actual) is type(expected), path + if torch.is_tensor(expected): + assert actual.dtype == expected.dtype, path + assert tuple(actual.shape) == tuple(expected.shape), path + assert torch.equal(actual, expected), path + elif isinstance(expected, dict): + assert set(actual) == set(expected), path + for key in expected: + _assert_nested_exact(actual[key], expected[key], "{}[{!r}]".format(path, key)) + elif isinstance(expected, (tuple, list)): + assert len(actual) == len(expected), path + for index, (live, saved) in enumerate(zip(actual, expected)): + _assert_nested_exact(live, saved, "{}[{}]".format(path, index)) + else: + assert actual == expected, path + + +def _hybrid_snapshot(optimizer): + return { + "children": tuple( + (child, deep_state_snapshot(child)) for child in optimizer._subopts + ), + "owner": optimizer._state_param_owner, + "owner_items": tuple(optimizer._state_param_owner.items()), + "shards": optimizer._hybrid_shard_bindings, + "shard_items": tuple(optimizer._hybrid_shard_bindings.items()), + "manifest": optimizer._hybrid_sharding_manifest, + "local": optimizer._hybrid_local_shard_bindings, + "roles": optimizer._hybrid_fqn_roles, + "slots": optimizer._hybrid_finalized_slots, + } + + +def _assert_hybrid_snapshot(optimizer, snapshot): + assert optimizer._state_param_owner is snapshot["owner"] + assert optimizer._hybrid_shard_bindings is snapshot["shards"] + 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_finalized_slots is snapshot["slots"] + assert len(optimizer._state_param_owner) == len(snapshot["owner_items"]) + for (live_key, live_value), (saved_key, saved_value) in zip( + optimizer._state_param_owner.items(), snapshot["owner_items"] + ): + assert live_key == saved_key + assert live_value[0] is saved_value[0] + assert live_value[1] is saved_value[1] + for (live_parameter, live_shard), (saved_parameter, saved_shard) in zip( + optimizer._hybrid_shard_bindings.items(), snapshot["shard_items"] + ): + assert live_parameter is saved_parameter + assert live_shard == saved_shard + for live_child, (saved_child, child_snapshot) in zip( + optimizer._subopts, snapshot["children"] + ): + assert live_child is saved_child + assert_deep_state_snapshot(live_child, child_snapshot) + + +def test_finalized_adamw_native_checkpoint_continues_exactly(): + source, source_matrix, source_bias = _finalized() + _step_backup(source, source_bias, torch.tensor([0.5, -0.25, 0.75, -1.0])) + checkpoint = copy.deepcopy(source.state_dict()) + assert set(checkpoint) == { + "muon", + "backup", + "backup_optimizer", + "finalized_binding", + } + assert checkpoint["finalized_binding"]["format_version"] == 1 + + target, target_matrix, target_bias = _finalized( + matrix=source_matrix, + bias=source_bias, + ) + target.load_state_dict(checkpoint) + _assert_nested_exact(target.state_dict(), source.state_dict()) + + next_gradient = torch.tensor([-0.5, 0.125, 0.25, 0.875]) + _step_backup(source, source_bias, next_gradient) + _step_backup(target, target_bias, next_gradient) + assert torch.equal(target_matrix, source_matrix) + assert torch.equal(target_bias, source_bias) + _assert_nested_exact(target.state_dict(), source.state_dict()) + + +def test_adamw_native_checkpoint_accepts_scheduler_initial_lr(): + source, source_matrix, source_bias = _finalized() + torch.optim.lr_scheduler.LambdaLR(source, lambda _step: 1.0) + assert source.backup.param_groups[0]["initial_lr"] == 1e-3 + _step_backup(source, source_bias, torch.tensor([0.5, -0.25, 0.75, -1.0])) + checkpoint = copy.deepcopy(source.state_dict()) + + target, target_matrix, target_bias = _finalized( + matrix=source_matrix, + bias=source_bias, + ) + torch.optim.lr_scheduler.LambdaLR(target, lambda _step: 1.0) + target.load_state_dict(checkpoint) + _assert_nested_exact(target.state_dict(), source.state_dict()) + + next_gradient = torch.tensor([-0.5, 0.125, 0.25, 0.875]) + _step_backup(source, source_bias, next_gradient) + _step_backup(target, target_bias, next_gradient) + assert torch.equal(target_matrix, source_matrix) + assert torch.equal(target_bias, source_bias) + + +def test_adamw_backup_supports_flat_shards_while_muon_nonowner_is_pruned(): + old_matrix = torch.nn.Parameter(torch.ones(2, 2)) + old_bias = torch.nn.Parameter(torch.ones(4)) + optimizer = GefenMuonHybrid( + [("layer.weight", old_matrix)], + [("layer.bias", old_bias)], + lr=1e-3, + fused=False, + backup_optimizer="adamw", + ) + 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,), 3.0)) + optimizer.post_sharding( + ( + ParameterRebinding(old_matrix, None, local_matrix), + ParameterRebinding(old_bias, bias, local_bias), + ), + manifest=ShardingManifest(matrix_shards + bias_shards), + ) + + assert optimizer._canonical_identity_ready() + assert optimizer.muon.param_groups[0]["params"] == [] + assert optimizer.backup.param_groups[0]["params"] == [bias] + assert optimizer.shard_identity(bias) == local_bias + assert optimizer.shard_bindings() == ((bias, local_bias), (None, local_matrix)) + assert set(optimizer._hybrid_shard_bindings) == {bias} + + +@pytest.mark.parametrize( + "damage", + ["extra_key", "wrong_shape", "negative_variance", "fractional_step"], +) +def test_invalid_adamw_state_rejects_before_any_child_publication(damage): + source, _matrix, source_bias = _finalized() + _step_backup(source, source_bias, torch.tensor([0.5, -0.25, 0.75, -1.0])) + checkpoint = copy.deepcopy(source.state_dict()) + (parameter_id,) = checkpoint["backup"]["state"] + state = checkpoint["backup"]["state"][parameter_id] + if damage == "extra_key": + state["foreign"] = torch.ones(1) + elif damage == "wrong_shape": + state["exp_avg"] = torch.ones(3) + elif damage == "negative_variance": + state["exp_avg_sq"][0] = -1.0 + else: + state["step"] = torch.tensor(1.5) + + target, _target_matrix, target_bias = _finalized() + _step_backup(target, target_bias, torch.tensor([-0.5, 0.25, -0.75, 1.0])) + snapshot = _hybrid_snapshot(target) + with pytest.raises((TypeError, ValueError), match="AdamW"): + target.load_state_dict(checkpoint) + _assert_hybrid_snapshot(target, snapshot) + + +def test_finalized_binding_guard_rejects_changed_identity_before_child_hooks(): + source, _matrix, source_bias = _finalized(prefix="Source") + _step_backup(source, source_bias, torch.ones(4)) + checkpoint = copy.deepcopy(source.state_dict()) + target, _target_matrix, _target_bias = _finalized(prefix="Target") + snapshot = _hybrid_snapshot(target) + calls = [] + target.muon.register_load_state_dict_pre_hook(lambda *_args: calls.append("muon")) + target.backup.register_load_state_dict_pre_hook(lambda *_args: calls.append("backup")) + + with pytest.raises(ValueError, match="binding or routing differs"): + target.load_state_dict(checkpoint) + + assert calls == [] + _assert_hybrid_snapshot(target, snapshot) + + +@pytest.mark.parametrize("backup_optimizer", ["gefen", "adamw"]) +def test_all_child_core_is_published_before_any_load_post_hook(backup_optimizer): + source, _source_matrix, _source_bias = _finalized( + backup_optimizer=backup_optimizer + ) + source.muon.param_groups[0]["lr"] = 0.006 + source.backup.param_groups[0]["lr"] = 0.007 + checkpoint = copy.deepcopy(source.state_dict()) + target, _target_matrix, _target_bias = _finalized( + backup_optimizer=backup_optimizer + ) + events = [] + + def muon_post(_child): + assert target.muon.param_groups[0]["lr"] == 0.006 + assert target.backup.param_groups[0]["lr"] == 0.007 + events.append("muon") + + def backup_post(_child): + assert target.muon.param_groups[0]["lr"] == 0.006 + assert target.backup.param_groups[0]["lr"] == 0.007 + events.append("backup") + + target.muon.register_load_state_dict_post_hook(muon_post) + target.backup.register_load_state_dict_post_hook(backup_post) + target.register_load_state_dict_post_hook(lambda _optimizer: events.append("hybrid")) + target.load_state_dict(checkpoint) + assert events == ["muon", "backup", "hybrid"] + + +@pytest.mark.parametrize("backup_optimizer", ["gefen", "adamw"]) +def test_throwing_first_child_post_hook_observes_fully_committed_core(backup_optimizer): + source, _source_matrix, _source_bias = _finalized( + backup_optimizer=backup_optimizer + ) + source.muon.param_groups[0]["lr"] = 0.004 + source.backup.param_groups[0]["lr"] = 0.008 + checkpoint = copy.deepcopy(source.state_dict()) + target, _target_matrix, _target_bias = _finalized( + backup_optimizer=backup_optimizer + ) + + def stop_after_observation(_child): + assert target.muon.param_groups[0]["lr"] == 0.004 + assert target.backup.param_groups[0]["lr"] == 0.008 + raise RuntimeError("post-hook sentinel") + + target.muon.register_load_state_dict_post_hook(stop_after_observation) + with pytest.raises(RuntimeError, match="post-hook sentinel"): + target.load_state_dict(checkpoint) + assert target.muon.param_groups[0]["lr"] == 0.004 + assert target.backup.param_groups[0]["lr"] == 0.008 + + +def test_unfinalized_legacy_adamw_checkpoint_schema_and_continuation_are_preserved(): + source, source_matrix, source_bias = _unfinalized() + _step_backup(source, source_bias, torch.tensor([0.5, -0.25, 0.75, -1.0])) + checkpoint = copy.deepcopy(source.state_dict()) + assert set(checkpoint) == {"muon", "backup", "backup_optimizer"} + + target, target_matrix, target_bias = _unfinalized( + matrix=source_matrix, + bias=source_bias, + ) + target.load_state_dict(checkpoint) + next_gradient = torch.tensor([-0.5, 0.125, 0.25, 0.875]) + _step_backup(source, source_bias, next_gradient) + _step_backup(target, target_bias, next_gradient) + assert torch.equal(target_matrix, source_matrix) + assert torch.equal(target_bias, source_bias) + _assert_nested_exact(target.state_dict(), source.state_dict()) + + +def test_finalized_and_legacy_native_checkpoint_schemas_do_not_cross_load(): + finalized, _matrix, _bias = _finalized() + legacy, _legacy_matrix, _legacy_bias = _unfinalized() + finalized_checkpoint = copy.deepcopy(finalized.state_dict()) + legacy_checkpoint = copy.deepcopy(legacy.state_dict()) + + with pytest.raises(ValueError, match="requires a versioned native binding guard"): + finalized.load_state_dict(legacy_checkpoint) + with pytest.raises(ValueError, match="cannot load finalized native shard state"): + legacy.load_state_dict(finalized_checkpoint) diff --git a/tests/test_hybrid_mixed_dtensor.py b/tests/test_hybrid_mixed_dtensor.py new file mode 100644 index 0000000..b9e45ab --- /dev/null +++ b/tests/test_hybrid_mixed_dtensor.py @@ -0,0 +1,652 @@ +"""Two-rank mixed-child Hybrid DTensor rebinding and native-load coverage.""" + +from datetime import timedelta +import copy +import multiprocessing as mp +import os +import queue as queue_module +import shutil +import tempfile +import traceback + +import pytest +import torch +import torch.distributed as dist + + +_WORLD = 2 +_MEMBERS = ("rank:0", "rank:1") +_MUON_FQN = "Model.Muon.Weight" +_BACKUP_FQN = "Model.Backup.Weight" + + +def _region(shape, coordinate, dimension): + from gefen import LogicalRegion + + global_length = shape[dimension] + chunk = (global_length + _WORLD - 1) // _WORLD + offset = min(coordinate * chunk, global_length) + length = max(0, min(global_length, offset + chunk) - offset) + offsets = [0] * len(shape) + lengths = list(shape) + offsets[dimension] = offset + lengths[dimension] = length + return LogicalRegion(tuple(offsets), tuple(lengths)) + + +def _manifest(): + from gefen import ( + ParameterIdentity, + ParameterLayout, + PlacementKind, + ProcessGroupIdentity, + ShardIdentity, + ShardPlacement, + ShardingManifest, + ) + + group = ProcessGroupIdentity("data_parallel", _MEMBERS) + shards = [] + by_fqn = {} + for fqn, shape, dimension in ( + (_MUON_FQN, (4, 4), 0), + (_BACKUP_FQN, (3, 4), 1), + ): + identity = ParameterIdentity(fqn, shape) + parameter_shards = [] + for coordinate, member in enumerate(_MEMBERS): + shard = ShardIdentity( + identity, + ParameterLayout.DTENSOR_1D_DEFAULT_WORLD, + _region(shape, coordinate, dimension), + placements=( + ShardPlacement( + "dp", + PlacementKind.DIMENSION_SHARD, + coordinate, + _WORLD, + parameter_dimension=dimension, + ), + ), + process_group=group, + local_member=member, + ) + shards.append(shard) + parameter_shards.append(shard) + by_fqn[fqn] = tuple(parameter_shards) + return ShardingManifest(tuple(shards)), by_fqn + + +def _optimizer(old_muon, old_backup, backup_optimizer, sharded_mode): + from gefen import GefenMuonHybrid + + return GefenMuonHybrid( + [("muon.weight", old_muon)], + [("backup.weight", old_backup)], + lr=1e-3, + fused=False, + deterministic=True, + sharded_mode=sharded_mode, + backup_optimizer=backup_optimizer, + backup_2d_period_one=True, + ) + + +def _local(value): + value = value.to_local() + return value.wait() if hasattr(value, "wait") else value + + +def _tree_exact(left, right): + if torch.is_tensor(left) or torch.is_tensor(right): + if not torch.is_tensor(left) or not torch.is_tensor(right): + return False + left_dtensor = hasattr(left, "to_local") + right_dtensor = hasattr(right, "to_local") + if left_dtensor != right_dtensor: + return False + if left_dtensor: + if ( + tuple(left.shape) != tuple(right.shape) + or tuple(map(str, left.placements)) != tuple(map(str, right.placements)) + ): + return False + left = _local(left) + right = _local(right) + return ( + left.dtype == right.dtype + and tuple(left.shape) == tuple(right.shape) + 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( + _tree_exact(left[key], right[key]) for key in left + ) + if isinstance(left, (list, tuple)): + return len(left) == len(right) and all( + _tree_exact(left_item, right_item) + for left_item, right_item in zip(left, right) + ) + return left == right + + +def _first_mismatch(left, right, path="root"): + if _tree_exact(left, right): + return None + if torch.is_tensor(left) or torch.is_tensor(right): + return "{}: tensor mismatch {} {}".format( + path, + tuple(left.shape) if torch.is_tensor(left) else type(left).__name__, + tuple(right.shape) if torch.is_tensor(right) else type(right).__name__, + ) + if type(left) is not type(right): + return "{}: type {} != {}".format(path, type(left).__name__, type(right).__name__) + if isinstance(left, dict): + if set(left) != set(right): + return "{}: keys {} != {}".format(path, sorted(map(str, left)), sorted(map(str, right))) + for key in left: + mismatch = _first_mismatch(left[key], right[key], "{}[{!r}]".format(path, key)) + if mismatch is not None: + return mismatch + elif isinstance(left, (list, tuple)): + if len(left) != len(right): + return "{}: length {} != {}".format(path, len(left), len(right)) + for index, (left_item, right_item) in enumerate(zip(left, right)): + mismatch = _first_mismatch(left_item, right_item, "{}[{}]".format(path, index)) + if mismatch is not None: + return mismatch + return "{}: {!r} != {!r}".format(path, left, right) + + +def _semantic_checkpoint(value): + if isinstance(value, dict): + return { + key: _semantic_checkpoint(item) + for key, item in value.items() + if key + not in { + "_gefen_rank_local_member", + "_gefen_rank_local_payload_0", + "_gefen_rank_local_payload_1", + } + } + if isinstance(value, list): + return [_semantic_checkpoint(item) for item in value] + if isinstance(value, tuple): + return tuple(_semantic_checkpoint(item) for item in value) + return value + + +def _build(rank, mesh, full_muon, full_backup, backup_optimizer, sharded_mode): + from torch import nn + from torch.distributed.tensor import Shard, distribute_tensor + + from gefen import ( + CheckpointTransport, + CodebookProcessGroupBinding, + ParameterLayout, + ParameterRebinding, + ProcessGroupScope, + ) + + manifest, shards = _manifest() + old_muon = nn.Parameter(full_muon.clone()) + old_backup = nn.Parameter(full_backup.clone()) + muon = nn.Parameter(distribute_tensor(full_muon.clone(), mesh, [Shard(0)])) + backup = nn.Parameter(distribute_tensor(full_backup.clone(), mesh, [Shard(1)])) + optimizer = _optimizer( + old_muon, + old_backup, + backup_optimizer, + sharded_mode, + ) + codebook_binding = CodebookProcessGroupBinding( + shards[_MUON_FQN][rank].process_group, + _MEMBERS[rank], + dist.group.WORLD, + torch.device("cpu"), + ) + optimizer.post_sharding( + ( + ParameterRebinding(old_muon, muon, shards[_MUON_FQN][rank]), + ParameterRebinding(old_backup, backup, shards[_BACKUP_FQN][rank]), + ), + manifest=manifest, + codebook_process_group=codebook_binding, + ) + assert optimizer.shard_identity(muon) == shards[_MUON_FQN][rank] + assert optimizer.shard_identity(backup) == shards[_BACKUP_FQN][rank] + contract = optimizer.optimizer_contract() + dtensor_native = [ + item + for item in contract.capabilities.checkpoints + if item.transport is CheckpointTransport.COMPOSITE_NATIVE + and item.same_topology + == frozenset({ParameterLayout.DTENSOR_1D_DEFAULT_WORLD}) + ] + assert len(dtensor_native) == 1 + assert dtensor_native[0].process_group_scope is ProcessGroupScope.DEFAULT_WORLD + assert dtensor_native[0].mesh_dimensions == (1,) + assert dtensor_native[0].requires_collective + assert dtensor_native[0].atomic_load + if backup_optimizer == "adamw": + assert all( + item.transport is not CheckpointTransport.CANONICAL_GLOBAL + for item in contract.capabilities.checkpoints + ) + return optimizer, muon, backup + + +def _set_gradients(muon, backup, mesh, muon_gradient, backup_gradient): + from torch.distributed.tensor import Shard, distribute_tensor + + muon.grad = distribute_tensor(muon_gradient.clone(), mesh, [Shard(0)]) + backup.grad = distribute_tensor(backup_gradient.clone(), mesh, [Shard(1)]) + + +def _worker(rank, init_file, backup_optimizer, sharded_mode, result_queue): + from torch.distributed.tensor import Shard, distribute_tensor, init_device_mesh + + try: + torch.set_num_threads(1) + dist.init_process_group( + "gloo", + init_method="file://{}".format(init_file), + rank=rank, + world_size=_WORLD, + timeout=timedelta(seconds=120), + ) + mesh = init_device_mesh("cpu", (_WORLD,), mesh_dim_names=("dp",)) + full_muon = torch.linspace(-0.75, 0.75, 16, dtype=torch.float32).reshape(4, 4) + full_backup = torch.linspace(-0.5, 0.625, 12, dtype=torch.float32).reshape(3, 4) + first_muon_gradient = torch.arange(16, dtype=torch.float32).reshape(4, 4).sin() + first_backup_gradient = torch.arange(12, dtype=torch.float32).reshape(3, 4).cos() + + asymmetric, asymmetric_muon, asymmetric_backup = _build( + rank, + mesh, + full_muon, + full_backup, + backup_optimizer, + sharded_mode, + ) + asymmetric_before = copy.deepcopy(asymmetric.state_dict()) + asymmetric_muon_before = _local(asymmetric_muon.detach()).clone() + asymmetric_backup_before = _local(asymmetric_backup.detach()).clone() + asymmetric_gradient = distribute_tensor( + first_muon_gradient.clone(), + mesh, + [Shard(0)], + ) + if rank == 0: + asymmetric_muon.grad = asymmetric_gradient + try: + asymmetric.step() + except RuntimeError as exc: + asymmetric_error = "gradient presence" in str(exc) + else: + asymmetric_error = False + asymmetric_rejected = ( + asymmetric_error + and _tree_exact(asymmetric_before, asymmetric.state_dict()) + and torch.equal( + _local(asymmetric_muon.detach()), + asymmetric_muon_before, + ) + and torch.equal( + _local(asymmetric_backup.detach()), + asymmetric_backup_before, + ) + ) + + optimizer, muon, backup = _build( + rank, + mesh, + full_muon, + full_backup, + backup_optimizer, + sharded_mode, + ) + _set_gradients( + muon, + backup, + mesh, + first_muon_gradient, + first_backup_gradient, + ) + optimizer.step() + optimizer.zero_grad() + checkpoint = copy.deepcopy(optimizer.state_dict()) + current_muon = muon.detach().full_tensor().clone() + current_backup = backup.detach().full_tensor().clone() + + rejected, rejected_muon, rejected_backup = _build( + rank, + mesh, + current_muon, + current_backup, + backup_optimizer, + sharded_mode, + ) + rejected_before = copy.deepcopy(rejected.state_dict()) + rejected_errors = [] + for rejected_role in ("muon", "backup"): + invalid = copy.deepcopy(checkpoint) + invalid[rejected_role]["param_groups"].append( + copy.deepcopy(invalid[rejected_role]["param_groups"][0]) + ) + try: + rejected.load_state_dict(invalid) + except (RuntimeError, ValueError): + rejected_errors.append(True) + else: + rejected_errors.append(False) + rejected_unchanged = ( + all(rejected_errors) + and _tree_exact(rejected_before, rejected.state_dict()) + and torch.equal(_local(rejected_muon.detach()), _local(muon.detach())) + and torch.equal(_local(rejected_backup.detach()), _local(backup.detach())) + ) + + resumed, resumed_muon, resumed_backup = _build( + rank, + mesh, + current_muon, + current_backup, + backup_optimizer, + sharded_mode, + ) + resumed.load_state_dict(checkpoint) + restored_source_state = _semantic_checkpoint(optimizer.state_dict()) + restored_resumed_state = _semantic_checkpoint(resumed.state_dict()) + restored = _tree_exact(restored_source_state, restored_resumed_state) + restored_mismatch = _first_mismatch( + restored_source_state, + restored_resumed_state, + ) + + second_muon_gradient = torch.arange(16, dtype=torch.float32).reshape(4, 4).cos() + second_backup_gradient = torch.arange(12, dtype=torch.float32).reshape(3, 4).sin() + _set_gradients( + muon, + backup, + mesh, + second_muon_gradient, + second_backup_gradient, + ) + _set_gradients( + resumed_muon, + resumed_backup, + mesh, + second_muon_gradient, + second_backup_gradient, + ) + optimizer.step() + resumed.step() + muon_parameter_exact = torch.equal( + _local(muon.detach()), _local(resumed_muon.detach()) + ) + backup_parameter_exact = torch.equal( + _local(backup.detach()), _local(resumed_backup.detach()) + ) + continued_source_state = _semantic_checkpoint(optimizer.state_dict()) + continued_resumed_state = _semantic_checkpoint(resumed.state_dict()) + continued_state_exact = _tree_exact( + continued_source_state, + continued_resumed_state, + ) + continued = ( + muon_parameter_exact + and backup_parameter_exact + and continued_state_exact + ) + result_queue.put( + { + "rank": rank, + "restored": restored, + "continued": continued, + "rejected_unchanged": rejected_unchanged, + "asymmetric_rejected": asymmetric_rejected, + "restored_mismatch": restored_mismatch, + "continued_mismatch": _first_mismatch( + continued_source_state, + continued_resumed_state, + ), + "muon_parameter_exact": muon_parameter_exact, + "backup_parameter_exact": backup_parameter_exact, + } + ) + except BaseException: + result_queue.put({"rank": rank, "traceback": traceback.format_exc()}) + raise + finally: + if dist.is_available() and dist.is_initialized(): + dist.destroy_process_group() + + +def _checkpoint_failure_worker(rank, init_file, result_queue): + from torch.distributed.tensor import init_device_mesh + + try: + torch.set_num_threads(1) + dist.init_process_group( + "gloo", + init_method="file://{}".format(init_file), + rank=rank, + world_size=_WORLD, + timeout=timedelta(seconds=60), + ) + mesh = init_device_mesh("cpu", (_WORLD,), mesh_dim_names=("dp",)) + full_muon = torch.linspace(-0.75, 0.75, 16, dtype=torch.float32).reshape(4, 4) + full_backup = torch.linspace(-0.5, 0.625, 12, dtype=torch.float32).reshape(3, 4) + muon_gradient = torch.arange(16, dtype=torch.float32).reshape(4, 4).sin() + backup_gradient = torch.arange(12, dtype=torch.float32).reshape(3, 4).cos() + + optimizer, muon, backup = _build( + rank, + mesh, + full_muon, + full_backup, + "adamw", + "approx", + ) + _set_gradients(muon, backup, mesh, muon_gradient, backup_gradient) + optimizer.step() + optimizer.zero_grad() + checkpoint = copy.deepcopy(optimizer.state_dict()) + save_before = _semantic_checkpoint(checkpoint) + save_muon_before = _local(muon.detach()).clone() + save_backup_before = _local(backup.detach()).clone() + + hook_handle = None + if rank == 0: + def reject_save(_optimizer): + raise ValueError("injected rank-local Hybrid save rejection") + + hook_handle = optimizer.register_state_dict_pre_hook(reject_save) + try: + optimizer.state_dict() + except RuntimeError as exc: + save_error = "state-dict pre-hook" in str(exc) + else: + save_error = False + finally: + if hook_handle is not None: + hook_handle.remove() + save_after = _semantic_checkpoint(optimizer.state_dict()) + save_rejected = ( + save_error + and _tree_exact(save_before, save_after) + and torch.equal(_local(muon.detach()), save_muon_before) + and torch.equal(_local(backup.detach()), save_backup_before) + ) + + current_muon = muon.detach().full_tensor().clone() + current_backup = backup.detach().full_tensor().clone() + rejected, rejected_muon, rejected_backup = _build( + rank, + mesh, + current_muon, + current_backup, + "adamw", + "approx", + ) + load_before = copy.deepcopy(rejected.state_dict()) + load_muon_before = _local(rejected_muon.detach()).clone() + load_backup_before = _local(rejected_backup.detach()).clone() + invalid = copy.deepcopy(checkpoint) + if rank == 0: + invalid["backup"]["param_groups"].append( + copy.deepcopy(invalid["backup"]["param_groups"][0]) + ) + try: + rejected.load_state_dict(invalid) + except RuntimeError as exc: + load_error = "backup load staging" in str(exc) + else: + load_error = False + load_after = rejected.state_dict() + load_rejected = ( + load_error + and _tree_exact(load_before, load_after) + and torch.equal(_local(rejected_muon.detach()), load_muon_before) + and torch.equal(_local(rejected_backup.detach()), load_backup_before) + ) + result_queue.put( + { + "rank": rank, + "save_rejected": save_rejected, + "load_rejected": load_rejected, + } + ) + except BaseException: + result_queue.put({"rank": rank, "traceback": traceback.format_exc()}) + raise + finally: + if dist.is_available() and dist.is_initialized(): + dist.destroy_process_group() + + +def _run(backup_optimizer, sharded_mode): + context = mp.get_context("spawn") + result_queue = context.Queue() + directory = tempfile.mkdtemp(prefix="gefen-hybrid-mixed-dtensor-") + init_file = os.path.join(directory, "store") + processes = [ + context.Process( + target=_worker, + args=(rank, init_file, backup_optimizer, sharded_mode, 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=300)) + except queue_module.Empty: + pass + for process in processes: + process.join(timeout=15) + 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() + shutil.rmtree(directory, ignore_errors=True) + assert len(results) == _WORLD, (results, [process.exitcode for process in processes]) + assert all(process.exitcode == 0 for process in processes), results + return results + + +def _run_checkpoint_failures(): + context = mp.get_context("spawn") + result_queue = context.Queue() + directory = tempfile.mkdtemp(prefix="gefen-hybrid-mixed-dtensor-failure-") + init_file = os.path.join(directory, "store") + processes = [ + context.Process( + target=_checkpoint_failure_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=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() + shutil.rmtree(directory, ignore_errors=True) + assert len(results) == _WORLD, (results, [process.exitcode for process in processes]) + assert all(process.exitcode == 0 for process in processes), results + return results + + +@pytest.mark.skipif( + not dist.is_available() or not dist.is_gloo_available(), + reason="mixed Hybrid DTensor rebinding requires Gloo", +) +@pytest.mark.parametrize( + ("backup_optimizer", "sharded_mode"), + [ + ("gefen", "approx"), + ("gefen", "exact"), + ("gefen", "distributed"), + ("adamw", "approx"), + ("adamw", "exact"), + ("adamw", "distributed"), + ], +) +def test_mixed_hybrid_dtensor_native_load_is_atomic_and_continues_exactly( + backup_optimizer, + sharded_mode, +): + results = _run(backup_optimizer, sharded_mode) + assert all("traceback" not in result for result in results), results + assert all( + result["restored"] + and result["continued"] + and result["rejected_unchanged"] + and result["asymmetric_rejected"] + for result in results + ), results + + +@pytest.mark.skipif( + not dist.is_available() or not dist.is_gloo_available(), + reason="mixed Hybrid DTensor checkpoint failure agreement requires Gloo", +) +def test_mixed_hybrid_dtensor_checkpoint_failures_are_rank_symmetric_and_atomic(): + results = _run_checkpoint_failures() + assert all("traceback" not in result for result in results), results + assert all( + result["save_rejected"] and result["load_rejected"] + for result in results + ), results diff --git a/tests/test_hybrid_rebinding.py b/tests/test_hybrid_rebinding.py index 1d91932..897be40 100644 --- a/tests/test_hybrid_rebinding.py +++ b/tests/test_hybrid_rebinding.py @@ -131,6 +131,8 @@ def _snapshot(optimizer): "finalized": optimizer._hybrid_post_sharding_finalized, "manifest": optimizer._hybrid_sharding_manifest, "local": optimizer._hybrid_local_shard_bindings, + "shards": optimizer._hybrid_shard_bindings, + "shard_contents": tuple(optimizer._hybrid_shard_bindings.items()), "roles": optimizer._hybrid_fqn_roles, "binding": optimizer._hybrid_codebook_process_group, "slots": optimizer._hybrid_finalized_slots, @@ -148,6 +150,13 @@ def _assert_snapshot(optimizer, snapshot): 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_shard_bindings is snapshot["shards"] + assert len(optimizer._hybrid_shard_bindings) == len(snapshot["shard_contents"]) + for (live_parameter, live_shard), (expected_parameter, expected_shard) in zip( + optimizer._hybrid_shard_bindings.items(), snapshot["shard_contents"] + ): + assert live_parameter is expected_parameter + assert live_shard == expected_shard 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"] @@ -498,31 +507,47 @@ def test_one_child_hybrids_support_convenience_rebinding(role): assert set(optimizer._state_param_owner) == {id(new)} -def test_adamw_backup_rejects_before_any_child_or_hybrid_mutation(): +def test_adamw_backup_rebinds_with_parent_owned_identity_and_muon_only_codebook_scope(): optimizer, old_matrix, old_bias = _optimizer(backup_optimizer="adamw") - matrix_shard = _ungrouped_replicated("Model.Layer.Weight", (2, 2)) + 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 = _ungrouped_replicated("Model.Layer.Bias", (4,)) - snapshot = _snapshot(optimizer) + matrix = torch.nn.Parameter(torch.full((2, 2), 3.0)) + bias = torch.nn.Parameter(torch.full((4,), 5.0)) - 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)), - ) + optimizer.post_sharding( + ( + ParameterRebinding(old_matrix, matrix, matrix_shard), + ParameterRebinding(old_bias, bias, bias_shard), + ), + manifest=ShardingManifest((matrix_shard, bias_shard)), + codebook_process_group=binding, + ) - _assert_snapshot(optimizer, snapshot) - assert not optimizer._canonical_identity_ready() + assert optimizer._canonical_identity_ready() + assert optimizer.shard_identity(matrix) == matrix_shard + assert optimizer.shard_identity(bias) == bias_shard + assert optimizer._hybrid_shard_bindings == { + matrix: matrix_shard, + bias: bias_shard, + } + assert optimizer.muon.codebook_process_group_binding() is binding + assert not hasattr(optimizer.backup, "_gefen_codebook_process_group") + assert optimizer.backup.param_groups[0]["params"] == [bias] + assert optimizer.backup.param_groups[0]["param_names"] == ["layer.bias"] + assert optimizer.state[bias] is optimizer.backup.state[bias] + assert "finalized_binding" in optimizer.state_dict() def test_composite_guard_fails_closed_after_owner_metadata_corruption(): diff --git a/tests/test_hybrid_scoped_failure_protocol.py b/tests/test_hybrid_scoped_failure_protocol.py index dc0f73e..22c4606 100644 --- a/tests/test_hybrid_scoped_failure_protocol.py +++ b/tests/test_hybrid_scoped_failure_protocol.py @@ -378,6 +378,33 @@ def closure(): } +def _checkpoint_parent_failure_result(rank, group): + optimizer, muon_parameter, backup_parameter, backup_shard = _make_scoped_hybrid( + rank, group + ) + _seed_hybrid_state( + optimizer, + muon_parameter, + backup_parameter, + backup_shard, + ) + hook_handle = None + if rank == 0: + def reject_save(_optimizer): + raise RuntimeError("checkpoint pre-hook boom on rank:0") + + hook_handle = optimizer.register_state_dict_pre_hook(reject_save) + try: + optimizer.state_dict() + message = None + except RuntimeError as exc: + message = str(exc) + finally: + if hook_handle is not None: + hook_handle.remove() + return {"message": message} + + def _distributed_worker(rank, init_file, result_queue): try: dist.init_process_group( @@ -398,6 +425,8 @@ def _distributed_worker(rank, init_file, result_queue): dist.barrier() closure_divergent = _closure_divergent_result(rank, group) dist.barrier() + checkpoint_parent = _checkpoint_parent_failure_result(rank, group) + dist.barrier() result_queue.put( { "rank": rank, @@ -406,6 +435,7 @@ def _distributed_worker(rank, init_file, result_queue): "amp_finite": amp_finite, "preflight": preflight, "closure_divergent": closure_divergent, + "checkpoint_parent": checkpoint_parent, } ) except BaseException: @@ -515,6 +545,15 @@ def test_hybrid_step_synchronizes_amp_and_preflight_across_the_scope(): 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 + # Hybrid parent checkpoint phases vote through the installed scope before + # the distributed-Muon child enters owner-state consolidation. A one-rank + # pre-hook failure therefore raises on both members instead of hanging the + # peer in the child's first broadcast. + checkpoint_parent = [result["checkpoint_parent"] for result in results] + assert all(item["message"] is not None for item in checkpoint_parent), checkpoint_parent + assert "checkpoint pre-hook boom on rank:0" in checkpoint_parent[0]["message"] + assert "state-dict pre-hook failed on another process-group member" in checkpoint_parent[1]["message"] + def _make_plain_hybrid(): matrix = torch.nn.Parameter(_muon_initial().clone()) diff --git a/tests/test_muon_distributed_checkpoint_safety.py b/tests/test_muon_distributed_checkpoint_safety.py index 66fa37f..7aaad73 100644 --- a/tests/test_muon_distributed_checkpoint_safety.py +++ b/tests/test_muon_distributed_checkpoint_safety.py @@ -523,10 +523,12 @@ def make_fallback_pair(parallel_value, fallback_value): rejection_message = None try: rejected.load_state_dict(corrupted) - except ValueError as exc: + except RuntimeError as exc: rejection_message = str(exc) fallback_atomic = ( rejection_message is not None + and "unscoped DTensor checkpoint local load staging failed on this rank" + in rejection_message and "ordered eligible" in rejection_message and snapshot_equal(rejected, rejected_params, before) ) @@ -767,10 +769,12 @@ def make_optimizer(params): error = None try: invalid_target.load_state_dict(invalid_checkpoint) - except ValueError as exc: + except RuntimeError as exc: error = str(exc) rejection_ok = ( error is not None + and "unscoped DTensor checkpoint local load staging failed on this rank" + in error and "refused populated" in error and invalid_target._gefen_global_step == step_before and torch.equal(invalid_target._gefen_codebook, codebook_before) diff --git a/tests/test_muon_grad_presence.py b/tests/test_muon_grad_presence.py index 7899f46..097ed79 100644 --- a/tests/test_muon_grad_presence.py +++ b/tests/test_muon_grad_presence.py @@ -569,8 +569,12 @@ def _reversed_order_worker(rank, world, port, result_queue): world_size=world, timeout=timedelta(seconds=12), ) - mesh = init_device_mesh("cpu", (world,)) - second_mesh = init_device_mesh("cpu", (world,)) + mesh = init_device_mesh( + "cpu", (world,), mesh_dim_names=("primary",) + ) + second_mesh = init_device_mesh( + "cpu", (world,), mesh_dim_names=("secondary",) + ) results = [] import warnings @@ -583,7 +587,7 @@ def _reversed_order_worker(rank, world, port, result_queue): "bare_ambiguous", "duplicate_names", "duplicate_one_rank", - "cross_mesh_reversed", + "asymmetric_second_mesh", ): generator = torch.Generator(device="cpu").manual_seed( 6100 + len(results) @@ -597,7 +601,7 @@ def _reversed_order_worker(rank, world, port, result_queue): in ("bare_ambiguous", "duplicate_names", "duplicate_one_rank") else (6, 10) ) - b_mesh = second_mesh if case == "cross_mesh_reversed" else mesh + b_mesh = second_mesh if case == "asymmetric_second_mesh" else mesh full_a = torch.randn(8, 8, generator=generator) full_a_grad = torch.randn(8, 8, generator=generator) * 0.01 full_b = torch.randn(*b_shape, generator=generator) @@ -618,11 +622,15 @@ def _reversed_order_worker(rank, world, port, result_queue): params = ( [("w", a), ("w", b)] if rank == 1 else [("a", a), ("b", b)] ) + elif case == "asymmetric_second_mesh": + # Both ranks retain A as the common first consensus anchor, but + # only rank 0 registers B on a second mesh. The fixed header + # must reject this globally and stop; rank 0 must never enter an + # unmatched second anchor after rank 1 has returned. + params = [("a", a), ("b", b)] if rank == 0 else [("a", a)] else: params = [("a", a), ("b", b)] - if ( - case.startswith("reversed") or case == "cross_mesh_reversed" - ) and rank == 1: + if case.startswith("reversed") and rank == 1: params = list(reversed(params)) optimizer = GefenMuon( params, @@ -740,32 +748,44 @@ def test_sharded_grad_presence_check_rejects_rank_divergent_order(): "bare_ambiguous", "duplicate_names", "duplicate_one_rank", - "cross_mesh_reversed", + "asymmetric_second_mesh", }, (rank, cases) # Duplicate explicit labels defeat cross-rank identification, so they # fail closed even with aligned order and full gradient presence — # and the rank whose own labels are clean must raise too, via the # reduced flag, instead of blocking in an abandoned collective. - for case in ("duplicate_names", "duplicate_one_rank"): - message = cases[case]["message"] - assert message is not None, (rank, case) - assert "unique name" in message, (rank, case, message) - # Order swaps between parameters on different meshes are step-fatal - # too; the optimizer-wide position probe must catch them. - message = cases["cross_mesh_reversed"]["message"] - assert message is not None, (rank, "cross_mesh_reversed") - assert "identical parameter-group order" in message, (rank, message) + message = cases["duplicate_names"]["message"] + assert message is not None, (rank, "duplicate_names") + assert "unique name" in message, (rank, "duplicate_names", message) + # A duplicate introduced on only one rank now fails one layer earlier: + # the constant-size intent header detects the asymmetric identity stream + # before the detailed name validator. Both diagnostics are fail-closed; + # this one is what guarantees the clean-label peer rejects too. + message = cases["duplicate_one_rank"]["message"] + assert message is not None, (rank, "duplicate_one_rank") + assert "identical DTensor parameter membership" in message, ( + rank, + "duplicate_one_rank", + message, + ) + # A secondary mesh present on only one rank must stop after the one + # common fixed header; no rank may enter a variable second anchor. + message = cases["asymmetric_second_mesh"]["message"] + assert message is not None, (rank, "asymmetric_second_mesh") + assert "identical DTensor parameter membership" in message, ( + rank, + message, + ) # Rank-divergent order is rejected up front, with or without a # gradient-presence mismatch layered on top. for case in ("reversed_mismatch", "reversed_consistent"): message = cases[case]["message"] assert message is not None, (rank, case) - assert "identical parameter-group order" in message, ( + assert "identical DTensor parameter membership" in message, ( rank, case, message, ) - assert "a" in message and "b" in message, (rank, case, message) message = cases["aligned_mismatch"]["message"] assert message is not None, (rank, "aligned_mismatch") assert "identical gradient presence" in message, (rank, message) @@ -790,7 +810,7 @@ def test_sharded_grad_presence_check_rejects_rank_divergent_order(): "aligned_consistent", "duplicate_names", "duplicate_one_rank", - "cross_mesh_reversed", + "asymmetric_second_mesh", ): assert not cases[case]["warned"], (rank, case) @@ -801,7 +821,7 @@ def test_sharded_grad_presence_check_rejects_rank_divergent_order(): "reversed_consistent", "aligned_mismatch", "duplicate_names", - "cross_mesh_reversed", + "asymmetric_second_mesh", ): assert ( rank_results[0][case]["message"] == rank_results[1][case]["message"] @@ -938,3 +958,402 @@ def test_cpu_mesh_preflights_never_query_cuda_capture_state(): # protocol; skipping the CUDA query must not change the decision. for rank, payload in rank_results.items(): assert payload["native_amp"] is True, (rank, payload) + + +def _two_dimensional_mesh_rejection_worker(rank, world, port, result_queue): + import torch.distributed as dist + from torch.distributed.tensor import ( + Replicate, + Shard, + distribute_tensor, + init_device_mesh, + ) + + try: + os.environ["MASTER_ADDR"] = "127.0.0.1" + os.environ["MASTER_PORT"] = port + os.environ["RANK"] = str(rank) + os.environ["WORLD_SIZE"] = str(world) + dist.init_process_group( + "gloo", + rank=rank, + world_size=world, + timeout=timedelta(seconds=30), + ) + mesh = init_device_mesh( + "cpu", + (2, 2), + mesh_dim_names=("shard", "replicate"), + ) + placements = [Shard(0), Replicate()] + full = torch.arange(64, dtype=torch.float32).reshape(8, 8).div_(64) + full_grad = torch.arange(64, dtype=torch.float32).reshape(8, 8).sin() + parameter = nn.Parameter( + distribute_tensor(full.clone(), mesh, placements) + ) + parameter.grad = distribute_tensor(full_grad.clone(), mesh, placements) + before = parameter.to_local().detach().clone() + message = None + try: + GefenMuon( + [("weight", parameter)], + lr=5e-2, + weight_decay=0.0, + fused=False, + ns_steps=1, + sharded_mode="exact", + ) + except RuntimeError as exc: + message = str(exc) + result_queue.put( + ( + "result", + rank, + { + "message": message, + "unchanged": torch.equal( + before, parameter.to_local().detach() + ), + }, + ) + ) + except Exception: + result_queue.put(("error", rank, traceback.format_exc())) + finally: + if dist.is_initialized(): + dist.destroy_process_group() + + +@pytest.mark.skipif( + not torch.distributed.is_available() or not torch.distributed.is_gloo_available(), + reason="2-D DeviceMesh intent regression needs Gloo", +) +def test_muon_rejects_multidimensional_mesh_before_optimizer_collectives(): + import torch.multiprocessing as mp + + world = 4 + context = mp.get_context("spawn") + result_queue = context.Queue() + port = _free_port() + processes = [ + context.Process( + target=_two_dimensional_mesh_rejection_worker, + args=(rank, world, port, result_queue), + ) + for rank in range(world) + ] + for process in processes: + process.start() + + messages = [] + try: + for _ in range(world): + messages.append(result_queue.get(timeout=90)) + except queue.Empty: + pass + finally: + for process in processes: + process.join(timeout=10) + for process in processes: + if process.is_alive(): + process.terminate() + process.join(timeout=5) + + assert all(process.exitcode == 0 for process in processes), [ + process.exitcode for process in processes + ] + errors = [payload for kind, _, payload in messages if kind == "error"] + assert not errors, "\n".join(errors) + results = { + rank: payload for kind, rank, payload in messages if kind == "result" + } + assert set(results) == set(range(world)), messages + for rank, payload in results.items(): + assert "one-dimensional DeviceMesh" in payload["message"], ( + rank, + payload, + ) + assert payload["unchanged"], (rank, payload) + + +def _subgroup_nonmember_late_add_worker(rank, world, port, result_queue): + import torch.distributed as dist + from torch.distributed.tensor import Shard, distribute_tensor, init_device_mesh + + try: + os.environ["MASTER_ADDR"] = "127.0.0.1" + os.environ["MASTER_PORT"] = port + os.environ["RANK"] = str(rank) + os.environ["WORLD_SIZE"] = str(world) + dist.init_process_group( + "gloo", + rank=rank, + world_size=world, + timeout=timedelta(seconds=20), + ) + full_mesh = init_device_mesh("cpu", (world,)) + subgroup_mesh = init_device_mesh("cpu", (world - 1,)) + full = torch.arange(16, dtype=torch.float32).reshape(4, 4) + full_parameter = nn.Parameter( + distribute_tensor(full.clone(), full_mesh, [Shard(0)]) + ) + subgroup_parameter = nn.Parameter( + distribute_tensor(full, subgroup_mesh, [Shard(0)]) + ) + base_parameter = nn.Parameter(torch.ones(4, 4)) + plain_optimizer = GefenMuon( + [("base", base_parameter)], + lr=1e-3, + fused=False, + ns_steps=1, + sharded_mode="exact", + ) + anchored_optimizer = GefenMuon( + [("full", full_parameter)], + lr=1e-3, + fused=False, + ns_steps=1, + sharded_mode="exact", + ) + + plain_message = None + try: + plain_optimizer.add_param_group( + {"params": [("subgroup", subgroup_parameter)]} + ) + except RuntimeError as exc: + plain_message = str(exc) + anchored_message = None + try: + anchored_optimizer.add_param_group( + {"params": [("subgroup", subgroup_parameter)]} + ) + except RuntimeError as exc: + anchored_message = str(exc) + result_queue.put( + ( + "result", + rank, + { + "coordinate": subgroup_mesh.get_coordinate(), + "plain_message": plain_message, + "plain_group_count": len(plain_optimizer.param_groups), + "anchored_message": anchored_message, + "anchored_group_count": len(anchored_optimizer.param_groups), + }, + ) + ) + except Exception: + result_queue.put(("error", rank, traceback.format_exc())) + finally: + if dist.is_initialized(): + dist.destroy_process_group() + + +@pytest.mark.skipif( + not torch.distributed.is_available() or not torch.distributed.is_gloo_available(), + reason="subgroup nonmember late-add regression needs Gloo", +) +def test_subgroup_dtensor_add_rejects_before_registration_on_all_world_ranks(): + import torch.multiprocessing as mp + + world = 3 + context = mp.get_context("spawn") + result_queue = context.Queue() + port = _free_port() + processes = [ + context.Process( + target=_subgroup_nonmember_late_add_worker, + args=(rank, world, port, result_queue), + ) + for rank in range(world) + ] + for process in processes: + process.start() + + messages = [] + try: + for _ in range(world): + messages.append(result_queue.get(timeout=60)) + except queue.Empty: + pass + finally: + for process in processes: + process.join(timeout=10) + for process in processes: + if process.is_alive(): + process.terminate() + process.join(timeout=5) + + assert all(process.exitcode == 0 for process in processes), [ + process.exitcode for process in processes + ] + errors = [payload for kind, _, payload in messages if kind == "error"] + assert not errors, "\n".join(errors) + results = { + rank: payload for kind, rank, payload in messages if kind == "result" + } + assert set(results) == set(range(world)), messages + assert results[world - 1]["coordinate"] is None + for rank, payload in results.items(): + for key in ("plain_message", "anchored_message"): + assert "span the initialized default process group" in payload[key], ( + rank, + payload, + ) + assert payload["plain_group_count"] == 1, (rank, payload) + assert payload["anchored_group_count"] == 1, (rank, payload) + + +def _distinct_mesh_route_worker(rank, world, port, result_queue): + import torch.distributed as dist + from torch.distributed.tensor import Shard, distribute_tensor, init_device_mesh + + try: + os.environ["MASTER_ADDR"] = "127.0.0.1" + os.environ["MASTER_PORT"] = port + os.environ["RANK"] = str(rank) + os.environ["WORLD_SIZE"] = str(world) + dist.init_process_group( + "gloo", + rank=rank, + world_size=world, + timeout=timedelta(seconds=20), + ) + mesh_a = init_device_mesh( + "cpu", (world,), mesh_dim_names=("dp",) + ) + mesh_b = init_device_mesh( + "cpu", (world,), mesh_dim_names=("dp",) + ) + full = torch.arange(64, dtype=torch.float32).reshape(8, 8).div_(64) + # Wrapping an equal-topology DTensor in nn.Parameter canonicalizes its + # DeviceMesh on recent PyTorch. Raw detached leaf DTensors preserve the + # distinct live c10d routes that the optimizer must distinguish. + parameter_a = distribute_tensor( + full.clone(), mesh_a, [Shard(0)] + ).requires_grad_() + parameter_b = distribute_tensor( + full.clone(), mesh_b, [Shard(0)] + ).requires_grad_() + parameter = parameter_a if rank == 0 else parameter_b + before = parameter.to_local().detach().clone() + optimizer = GefenMuon( + [("weight", parameter)], + lr=1e-3, + fused=False, + ns_steps=1, + sharded_mode="exact", + ) + message = None + try: + optimizer._assert_sharded_mode_collective_intent_consistent() + except RuntimeError as exc: + message = str(exc) + result_queue.put( + ( + "result", + rank, + { + "message": message, + "group_a": str(mesh_a.get_group().group_name), + "group_b": str(mesh_b.get_group().group_name), + "parameter_a_group": str( + parameter_a.device_mesh.get_group().group_name + ), + "parameter_b_group": str( + parameter_b.device_mesh.get_group().group_name + ), + "selected_group": str( + parameter.device_mesh.get_group().group_name + ), + "is_leaf": parameter.is_leaf, + "portable": optimizer._gradient_preflight_portable_mesh_key( + parameter.device_mesh + ), + "unchanged": torch.equal( + before, parameter.to_local().detach() + ), + "baseline": ( + getattr( + optimizer, + "_gefen_muon_collective_intent_baseline", + None, + ) + is None + ), + }, + ) + ) + except Exception: + result_queue.put(("error", rank, traceback.format_exc())) + finally: + if dist.is_initialized(): + dist.destroy_process_group() + + +@pytest.mark.skipif( + not torch.distributed.is_available() or not torch.distributed.is_gloo_available(), + reason="distinct DeviceMesh routing regression needs Gloo", +) +def test_intent_distinguishes_distinct_mesh_routes_or_accepts_shared_group(): + import torch.multiprocessing as mp + + world = 2 + context = mp.get_context("spawn") + result_queue = context.Queue() + port = _free_port() + processes = [ + context.Process( + target=_distinct_mesh_route_worker, + args=(rank, world, port, result_queue), + ) + for rank in range(world) + ] + for process in processes: + process.start() + + messages = [] + try: + for _ in range(world): + messages.append(result_queue.get(timeout=60)) + except queue.Empty: + pass + finally: + for process in processes: + process.join(timeout=10) + for process in processes: + if process.is_alive(): + process.terminate() + process.join(timeout=5) + + assert all(process.exitcode == 0 for process in processes), [ + process.exitcode for process in processes + ] + errors = [payload for kind, _, payload in messages if kind == "error"] + assert not errors, "\n".join(errors) + results = { + rank: payload for kind, rank, payload in messages if kind == "result" + } + assert set(results) == set(range(world)), messages + assert results[0]["group_a"] == results[1]["group_a"] + assert results[0]["group_b"] == results[1]["group_b"] + for payload in results.values(): + assert payload["parameter_a_group"] == payload["group_a"] + assert payload["parameter_b_group"] == payload["group_b"] + assert results[0]["selected_group"] == results[0]["group_a"] + assert results[1]["selected_group"] == results[1]["group_b"] + for rank, payload in results.items(): + assert payload["is_leaf"], (rank, payload) + assert payload["unchanged"], (rank, payload) + if payload["group_a"] != payload["group_b"]: + assert payload["message"] is not None, (rank, payload) + assert ( + "identical DTensor parameter membership" in payload["message"] + ), (rank, payload) + assert payload["baseline"], (rank, payload) + else: + # PyTorch 2.5 reuses the default-world group for equivalent 1-D + # DeviceMeshes. There is no distinct live route to reject. + assert payload["message"] is None, (rank, payload) + assert not payload["baseline"], (rank, payload) diff --git a/tests/test_optimizer_contracts.py b/tests/test_optimizer_contracts.py index 5be414a..e0ad45d 100644 --- a/tests/test_optimizer_contracts.py +++ b/tests/test_optimizer_contracts.py @@ -124,6 +124,27 @@ def _checkpoint_support(contract, transport, layout): ] +def _adamw_adapter_training(): + return ( + TrainingSupport(ParameterLayout.REPLICATED, ProcessGroupScope.NONE), + TrainingSupport( + ParameterLayout.FLATTENED_ELEMENT_SHARD, + ProcessGroupScope.NONE, + ), + TrainingSupport( + ParameterLayout.WHOLE_PARAMETER_OWNER, + ProcessGroupScope.ADAPTER_DEFINED, + requires_complete_parameter_storage=True, + requires_post_step_parameter_sync=True, + ), + TrainingSupport( + _DTENSOR, + ProcessGroupScope.INFERRED_DEVICE_MESH, + mesh_dimensions=(1,), + ), + ) + + @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)) @@ -486,10 +507,9 @@ def test_hybrid_contract_preserves_child_namespaces(backup_optimizer): 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 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 @@ -501,14 +521,15 @@ def test_hybrid_contract_preserves_child_namespaces(backup_optimizer): 1 ].contract.capabilities.explicit_process_group_codebook_scope assert {field.name for field in contract.state_layout.fields} == { - "backup_optimizer" + "backup_optimizer", + "finalized_binding", } 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 + assert checkpoint[0].atomic_load def test_muon_contract_keeps_mixed_normuon_variants_in_one_mode(): @@ -530,6 +551,22 @@ def test_muon_contract_keeps_mixed_normuon_variants_in_one_mode(): assert "quantized_normuon_replicated_exact" in variants +def test_muon_contract_rejects_invalid_live_mode_without_normalizing_it(): + parameter = torch.nn.Parameter(torch.ones(4, 4)) + optimizer = GefenMuon( + [("weight", parameter)], + fused=False, + sharded_mode="approx", + ) + optimizer.param_groups[0]["sharded_mode"] = "invalid" + + with pytest.raises(RuntimeError, match="invalid sharded_mode"): + optimizer.optimizer_contract() + with pytest.raises(RuntimeError, match="invalid sharded_mode"): + optimizer._canonical_state_variant_layout() + assert optimizer.param_groups[0]["sharded_mode"] == "invalid" + + @pytest.mark.parametrize("implementation", ["gefen", "muon"]) def test_capturable_contract_declines_atomic_state_movement(implementation): shape = (4,) if implementation == "gefen" else (2, 2) @@ -783,13 +820,18 @@ def test_backup_only_hybrid_training_claims_come_from_backup_child(backup_optimi 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. + # The Hybrid adapter owns and validates exact AdamW DTensor rebinding, + # so the backup-only composite can make this narrowly scoped claim even + # though AdamW does not publish a child contract of its own. assert contract.children[0].contract is None - assert contract.capabilities.training == ( - TrainingSupport(ParameterLayout.REPLICATED, ProcessGroupScope.NONE), + assert contract.capabilities.training == _adamw_adapter_training() + checkpoints = contract.capabilities.checkpoints + assert len(checkpoints) == 1 + assert all( + item.transport is not CheckpointTransport.CANONICAL_GLOBAL + for item in checkpoints ) + assert not contract.capabilities.explicit_process_group_codebook_scope @pytest.mark.parametrize("backup_optimizer", ["gefen", "adamw"]) @@ -809,9 +851,7 @@ def test_hybrid_training_claims_are_the_union_of_routed_children(backup_optimize if backup_optimizer == "gefen": backup_training = contract.children[1].contract.capabilities.training else: - backup_training = ( - TrainingSupport(ParameterLayout.REPLICATED, ProcessGroupScope.NONE), - ) + backup_training = _adamw_adapter_training() assert set(training) == set(muon_training) | set(backup_training) assert len(set(training)) == len(training) diff --git a/tests/test_portable_dcp.py b/tests/test_portable_dcp.py index ac0fd8a..a1ae30a 100644 --- a/tests/test_portable_dcp.py +++ b/tests/test_portable_dcp.py @@ -329,6 +329,104 @@ def test_filesystem_dcp_round_trip_is_tensor_only(tmp_path, variant): ) +def test_public_loader_accepts_legacy_v1_plain_checkpoint(tmp_path): + ( + source, + source_parameter, + source_binding, + target, + target_parameter, + target_binding, + ) = _optimizer_pair("plain-block") + _initialize(source, source_parameter, variant="plain-block") + limits = _limits() + document = source.export_portable_state( + checkpoint_process_group=source_binding, + transaction_id="legacy-v1-plain-export", + limits=limits, + ) + wire_limits = limits._wire_limits(collective=True) + plan = portable_dcp._prepare_canonical_wire_value(document, wire_limits) + state = portable_dcp._state_from_plan("optimizer", plan, wire_limits) + checkpoint = tmp_path / "legacy-v1-plain" + dcp.save(state, storage_writer=dcp.FileSystemWriter(checkpoint)) + keys = dcp.FileSystemReader(checkpoint).read_metadata().state_dict_metadata + assert "optimizer.__gefen_portable_metadata_v1__" in keys + assert "optimizer.__gefen_portable_metadata_v2__" not in keys + load_portable_dcp( + target, + checkpoint_process_group=target_binding, + storage_reader=dcp.FileSystemReader(checkpoint), + transaction_id="legacy-v1-plain-load", + limits=limits, + ) + _assert_loaded_state( + target, + target_parameter, + document, + variant="plain-block", + ) + + +def test_sharded_dcp_round_trips_initialized_scalar_state(tmp_path): + def scalar_optimizer(*, deterministic): + parameter = torch.nn.Parameter(torch.tensor(2.0)) + optimizer = Gefen( + [("scalar", parameter)], + fused=False, + factored_v_2d=False, + period_one_substrings=("scalar",), + deterministic=deterministic, + ) + return optimizer, parameter, _finalize( + optimizer, + parameter, + layout=ParameterLayout.REPLICATED, + ) + + source, source_parameter, source_binding = scalar_optimizer( + deterministic=True + ) + source._gefen_global_step = 4 + source._gefen_codebook = torch.linspace(-1.0, 1.0, 256) + source.state[source_parameter].update( + { + "automatic_period": 1, + "step": 4, + "m_codebook": torch.tensor([[255]], dtype=torch.uint8), + "m_magnitude": torch.tensor([[2.0]]), + "vmean": torch.tensor([[3.0]]), + "vmean_step": 3, + } + ) + target, target_parameter, target_binding = scalar_optimizer( + deterministic=False + ) + checkpoint = tmp_path / "scalar-v2" + save_portable_dcp( + source, + checkpoint_process_group=source_binding, + storage_writer=dcp.FileSystemWriter(checkpoint), + transaction_id="scalar-v2-save", + limits=_limits(), + ) + load_portable_dcp( + target, + checkpoint_process_group=target_binding, + storage_reader=dcp.FileSystemReader(checkpoint), + transaction_id="scalar-v2-load", + limits=_limits(), + ) + assert target._gefen_global_step == 4 + assert target._deterministic is True + assert target.state[target_parameter]["step"] == 4 + assert target.state[target_parameter]["vmean_step"] == 3 + assert torch.equal( + target.state[target_parameter]["vmean"].reshape(()), + torch.tensor(3.0), + ) + + @pytest.mark.parametrize( "namespace", ["", " optimizer", "optimizer.", "optim\x00izer", "optim/izer", "optimé"], @@ -592,15 +690,44 @@ def test_corrupted_tensor_payload_leaves_target_unchanged(tmp_path): limits=limits, ) - wire_limits = limits._wire_limits(collective=True) - state = {"optimizer": {}} + from gefen.portable_dcp_sharded import ( + _DCP_SHARDED_METADATA_KEY, + _allocate_sharded_load_state, + _metadata_storage_spec, + _read_sharded_metadata, + ) + + reader = dcp.FileSystemReader(original) + checkpoint_metadata = reader.read_metadata() + metadata = _read_sharded_metadata( + storage_reader=reader, + checkpoint_metadata=checkpoint_metadata, + binding=source_binding, + namespace="optimizer", + limits=limits, + transaction_id="dcp-read-corruption-source-metadata", + ) + metadata_shape = _metadata_storage_spec( + checkpoint_metadata, + namespace="optimizer", + limits=limits, + ) + state = _allocate_sharded_load_state( + metadata, + metadata_tensor=torch.empty(metadata_shape, dtype=torch.uint8), + namespace="optimizer", + binding=source_binding, + ) dcp.load( state, - storage_reader=dcp.FileSystemReader(original), - planner=portable_dcp._load_planner("optimizer", wire_limits), + storage_reader=reader, process_group=None, ) - payload_keys = sorted(key for key in state["optimizer"] if key.startswith(portable_dcp._DCP_PAYLOAD_PREFIX)) + payload_keys = sorted( + key + for key in state["optimizer"] + if key != _DCP_SHARDED_METADATA_KEY + ) 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) @@ -612,7 +739,7 @@ def test_corrupted_tensor_payload_leaves_target_unchanged(tmp_path): ) before = target._canonical_import_live_token() - with pytest.raises(RuntimeError, match="invalid digest"): + with pytest.raises(RuntimeError, match="integrity digest mismatch"): load_portable_dcp( target, checkpoint_process_group=target_binding, diff --git a/tests/test_portable_dcp_sharded_segments.py b/tests/test_portable_dcp_sharded_segments.py new file mode 100644 index 0000000..a20a829 --- /dev/null +++ b/tests/test_portable_dcp_sharded_segments.py @@ -0,0 +1,547 @@ +"""CPU/Gloo coverage for bounded sharded portable-DCP redistribution.""" + +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 +import torch.distributed.checkpoint as dcp +from torch.distributed.checkpoint import FileSystemReader, FileSystemWriter + +from gefen.checkpoint import CheckpointProcessGroupBinding +from gefen.codebook import CodebookProcessGroupBinding +from gefen.contracts import ( + LogicalRegion, + LogicalSlice, + ParameterIdentity, + ParameterLayout, + PlacementKind, + ProcessGroupIdentity, + ShardIdentity, + ShardPlacement, + ShardingManifest, +) +from gefen.gefen import Gefen +from gefen.portable_dcp_sharded import ( + _DCP_SHARDED_METADATA_KEY, + _allocate_canonical_dtensor, + _as_canonical_dtensor, + _canonical_local_tensor, + _canonical_projection_send_numel, + _canonical_segments, + _decode_metadata_tensor, + _dense_local_numel, + _dense_projection_send_numel, + _dense_segments_for_shard, + _factored_axis_segments, + _full_segment, + _load_sharded_payloads, + _local_factored_projection, + _prepare_sharded_save_state, + _read_sharded_metadata, + _redistribute_canonical_to_dense, + _redistribute_dense_to_canonical, + _redistribute_flat_field, + _segment_global_indices, + _validate_sharded_dcp_entries, +) +from gefen.portable_state import PortableStateLimits +from gefen.rebinding import ParameterRebinding + + +def _limits(): + return PortableStateLimits( + max_fragment_tensor_bytes=1 << 20, + max_collective_tensor_bytes=4 << 20, + max_collective_metadata_bytes=4 << 20, + max_metadata_bytes=1 << 20, + max_members=4, + 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 _dtensor_shards(identity, group, *, dimension): + chunk = (identity.global_shape[dimension] + len(group.ordered_members) - 1) // len( + group.ordered_members + ) + shards = [] + for coordinate, member in enumerate(group.ordered_members): + offset = min(coordinate * chunk, identity.global_shape[dimension]) + length = max( + 0, + min(identity.global_shape[dimension], offset + chunk) - offset, + ) + offsets = [0] * len(identity.global_shape) + lengths = list(identity.global_shape) + offsets[dimension] = offset + lengths[dimension] = length + shards.append( + ShardIdentity( + identity, + ParameterLayout.DTENSOR_1D_DEFAULT_WORLD, + LogicalRegion(tuple(offsets), tuple(lengths)), + placements=( + ShardPlacement( + "model", + PlacementKind.DIMENSION_SHARD, + coordinate, + len(group.ordered_members), + dimension, + ), + ), + process_group=group, + local_member=member, + ) + ) + return tuple(shards) + + +def _flat_shards(identity, group, lengths): + 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( + "model", + PlacementKind.FLAT_SHARD, + coordinate, + len(group.ordered_members), + ), + ), + process_group=group, + local_member=member, + ) + ) + offset += length + return tuple(shards) + + +def test_factored_axis_routes_only_compact_flat_head_and_tail(): + group = ProcessGroupIdentity("factored_flat_segments", ("rank:0", "rank:1", "rank:2")) + identity = ParameterIdentity("model.matrix", (3, 5)) + shard = _flat_shards(identity, group, (4, 3, 8))[1] + row_segments = _factored_axis_segments(shard, axis=0) + column_segments = _factored_axis_segments(shard, axis=1) + assert _segment_global_indices(row_segments).tolist() == [0, 1] + assert _segment_global_indices(column_segments).tolist() == [0, 1, 4] + assert sum(segment.length for segment in column_segments) < identity.global_shape[1] + + +@pytest.mark.parametrize( + ("dimension", "coordinate", "expected_rows", "expected_columns"), + [ + (0, 0, [0, 1], [0, 1, 2, 3, 4, 5]), + (1, 1, [0, 1, 2, 3], [3, 4, 5]), + ], +) +def test_factored_axis_routes_dtensor_dimension_slices( + dimension, + coordinate, + expected_rows, + expected_columns, +): + group = ProcessGroupIdentity("factored_dtensor_segments", ("rank:0", "rank:1")) + identity = ParameterIdentity("model.matrix", (4, 6)) + shard = _dtensor_shards(identity, group, dimension=dimension)[coordinate] + assert _segment_global_indices(_factored_axis_segments(shard, axis=0)).tolist() == expected_rows + assert _segment_global_indices(_factored_axis_segments(shard, axis=1)).tolist() == expected_columns + + +def test_factored_axis_routes_empty_dtensor_shard_without_factors(): + group = ProcessGroupIdentity("factored_empty_segments", ("rank:0", "rank:1")) + identity = ParameterIdentity("model.matrix", (2, 1)) + empty = _dtensor_shards(identity, group, dimension=1)[1] + assert _factored_axis_segments(empty, axis=0) == () + assert _factored_axis_segments(empty, axis=1) == () + + +def test_compact_factored_projection_uses_saved_fp32_denominator_exactly(): + group = ProcessGroupIdentity("factored_projection_segments", ("rank:0", "rank:1", "rank:2")) + identity = ParameterIdentity("model.matrix", (3, 5)) + shard = _flat_shards(identity, group, (4, 3, 8))[1] + full_row = torch.tensor([1.0, 1.0000001192092896, 16.0], dtype=torch.float32) + full_column = torch.tensor([2.0, 4.0, 8.0, 16.0, 32.0], dtype=torch.float32) + row_indices = _segment_global_indices(_factored_axis_segments(shard, axis=0)) + column_indices = _segment_global_indices(_factored_axis_segments(shard, axis=1)) + denominator = full_row.mean().clamp_(min=torch.finfo(torch.float32).tiny).reshape(1) + projected = _local_factored_projection( + full_row[row_indices], + row_indices, + full_column[column_indices], + column_indices, + denominator, + shard, + ) + expected = torch.outer(full_row, full_column).div_(denominator.reshape(())).reshape(-1)[4:7] + assert torch.equal(projected.view(torch.uint8), expected.view(torch.uint8)) + + +def test_factored_router_keeps_one_by_four_thousand_flat_scratch_below_fragment_limit(): + members = tuple("rank:{}".format(index) for index in range(4)) + group = ProcessGroupIdentity("factored_large_flat_segments", members) + identity = ParameterIdentity("model.matrix", (1, 4000)) + shards = _flat_shards(identity, group, (1000, 1000, 1000, 1000)) + scratch_bytes = [] + for shard in shards: + row_numel = sum(segment.length for segment in _factored_axis_segments(shard, axis=0)) + column_numel = sum(segment.length for segment in _factored_axis_segments(shard, axis=1)) + local_dense_numel = _dense_local_numel(shard, unique_replicas=False) + scratch_bytes.append((row_numel + column_numel + 1 + local_dense_numel) * 4) + assert scratch_bytes == [8008, 8008, 8008, 8008] + assert max(scratch_bytes) < 9000 + assert identity.global_shape[1] * 4 > 9000 + + +def test_projection_send_plan_counts_replicated_scalar_amplification(): + parts = 1024 + counts = _canonical_projection_send_numel( + 1, + target_segments_by_rank=tuple(_full_segment(1) for _coordinate in range(parts)), + parts=parts, + ) + assert counts[0] * 4 == 4096 + assert all(value == 0 for value in counts[1:]) + assert counts[0] * 4 > 2048 + + +def test_dense_send_plan_is_closed_form_for_large_nonleading_dtensor_shards(): + group = ProcessGroupIdentity("dense_large_dtensor_send", ("rank:0", "rank:1")) + identity = ParameterIdentity("model.matrix", (1_000_000, 2)) + shards = _dtensor_shards(identity, group, dimension=1) + assert _dense_projection_send_numel( + identity.numel, + target_shards=shards, + parts=2, + ) == (1_000_000, 1_000_000) + + +def _worker(rank, world_size, init_file, checkpoint_dir, result_queue): + try: + dist.init_process_group( + "gloo", + init_method="file://{}".format(init_file), + rank=rank, + world_size=world_size, + timeout=timedelta(seconds=60), + ) + members = tuple("rank:{}".format(index) for index in range(world_size)) + group = ProcessGroupIdentity("portable_sharded_segments", members) + binding = CheckpointProcessGroupBinding( + group, + members[rank], + dist.group.WORLD, + torch.device("cpu"), + ) + + identity = ParameterIdentity("model.weight", (3, 5)) + dtensor_shards = _dtensor_shards(identity, group, dimension=1) + source_segments = tuple( + _dense_segments_for_shard(shard, unique_replicas=True) + for shard in dtensor_shards + ) + canonical_segments = _canonical_segments(identity.numel, world_size) + global_value = torch.arange(identity.numel, dtype=torch.float32).reshape(identity.global_shape) + region = dtensor_shards[rank].logical_region + index = tuple( + slice(offset, offset + length) + for offset, length in zip(region.offsets, region.lengths) + ) + local = global_value[index].contiguous() + canonical = _redistribute_flat_field( + local, + source_segments_by_rank=source_segments, + target_segments_by_rank=canonical_segments, + global_numel=identity.numel, + binding=binding, + ) + canonical_segment = canonical_segments[rank][0] + expected_canonical = global_value.reshape(-1)[ + canonical_segment.global_offset : canonical_segment.global_offset + canonical_segment.length + ] + checkpoint_value = _as_canonical_dtensor( + canonical, + global_numel=identity.numel, + binding=binding, + ) + dcp.save( + {"field": checkpoint_value}, + storage_writer=FileSystemWriter(checkpoint_dir), + process_group=binding.process_group, + ) + loaded_value = _allocate_canonical_dtensor( + global_numel=identity.numel, + dtype=torch.float32, + binding=binding, + ) + dcp.load( + {"field": loaded_value}, + storage_reader=FileSystemReader(checkpoint_dir), + process_group=binding.process_group, + ) + loaded_canonical = _canonical_local_tensor(loaded_value) + + flat_shards = _flat_shards(identity, group, (4, 11)) + flat_segments = tuple( + _dense_segments_for_shard(shard, unique_replicas=False) + for shard in flat_shards + ) + projected = _redistribute_flat_field( + canonical, + source_segments_by_rank=canonical_segments, + target_segments_by_rank=flat_segments, + global_numel=identity.numel, + binding=binding, + ) + flat_slice = flat_shards[rank].logical_slice + expected_projected = global_value.reshape(-1)[ + flat_slice.flat_offset : flat_slice.flat_offset + flat_slice.length + ] + + empty_identity = ParameterIdentity("model.tiny", (2, 1)) + empty_shards = _dtensor_shards(empty_identity, group, dimension=1) + empty_sources = tuple( + _dense_segments_for_shard(shard, unique_replicas=True) + for shard in empty_shards + ) + empty_global = torch.tensor([[3.0], [7.0]]) + empty_region = empty_shards[rank].logical_region + empty_index = tuple( + slice(offset, offset + length) + for offset, length in zip(empty_region.offsets, empty_region.lengths) + ) + empty_local = empty_global[empty_index].contiguous() + empty_canonical = _redistribute_flat_field( + empty_local, + source_segments_by_rank=empty_sources, + target_segments_by_rank=_canonical_segments(empty_identity.numel, world_size), + global_numel=empty_identity.numel, + binding=binding, + ) + + large_identity = ParameterIdentity("model.large", (100_000, 2)) + large_shards = _dtensor_shards(large_identity, group, dimension=1) + large_region = large_shards[rank].logical_region + large_index = tuple( + slice(offset, offset + length) + for offset, length in zip( + large_region.offsets, + large_region.lengths, + ) + ) + large_global = torch.arange( + large_identity.numel, + dtype=torch.float32, + ).reshape(large_identity.global_shape) + large_local = large_global[large_index].contiguous() + large_canonical = _redistribute_dense_to_canonical( + large_local, + source_shard=large_shards[rank], + global_numel=large_identity.numel, + binding=binding, + ) + large_round_trip = _redistribute_canonical_to_dense( + large_canonical, + target_shards=large_shards, + global_numel=large_identity.numel, + binding=binding, + ) + + flat_manifest_shards = _flat_shards(identity, group, (8, 7)) + source_shard = flat_manifest_shards[rank] + source_start = source_shard.logical_slice.flat_offset + source_stop = source_start + source_shard.logical_slice.length + original_parameter = torch.nn.Parameter(global_value.clone()) + local_parameter = torch.nn.Parameter( + global_value.reshape(-1)[source_start:source_stop].clone() + ) + optimizer = Gefen( + [("weight", original_parameter)], + fused=False, + factored_v_2d=False, + force_2d_period_one=True, + deterministic=True, + ) + optimizer.post_sharding( + (ParameterRebinding(original_parameter, local_parameter, source_shard),), + manifest=ShardingManifest(flat_manifest_shards), + codebook_process_group=CodebookProcessGroupBinding( + group, + members[rank], + dist.group.WORLD, + torch.device("cpu"), + ), + ) + optimizer._gefen_global_step = 4 + optimizer._gefen_codebook = torch.linspace(-1.0, 1.0, 256) + local_momentum = torch.linspace(-0.5, 0.75, identity.numel)[source_start:source_stop] + optimizer.state[local_parameter].update( + { + "automatic_period": 1, + "step": 3, + "m_codebook": torch.where( + torch.signbit(local_momentum), + torch.zeros(local_momentum.numel(), dtype=torch.uint8), + torch.full((local_momentum.numel(),), 255, dtype=torch.uint8), + ).reshape(-1, 1), + "m_magnitude": local_momentum.abs().reshape(-1, 1), + "vmean": torch.linspace(0.25, 1.25, identity.numel)[ + source_start:source_stop + ].reshape(-1, 1), + "vmean_step": 2, + } + ) + prepared_state = _prepare_sharded_save_state( + optimizer, + binding=binding, + transaction_id="prepare-sharded-save", + limits=_limits(), + namespace="optimizer", + ) + prepared_envelope = prepared_state["optimizer"] + prepared_field_keys = [ + key for key in prepared_envelope if key != _DCP_SHARDED_METADATA_KEY + ] + prepared_local_sizes = [ + _canonical_local_tensor(prepared_envelope[key]).numel() + for key in prepared_field_keys + ] + decoded_metadata = _decode_metadata_tensor( + prepared_envelope[_DCP_SHARDED_METADATA_KEY], + limits=_limits(), + ) + optimizer_checkpoint_dir = "{}-optimizer".format(checkpoint_dir) + dcp.save( + prepared_state, + storage_writer=FileSystemWriter(optimizer_checkpoint_dir), + process_group=binding.process_group, + ) + optimizer_reader = FileSystemReader(optimizer_checkpoint_dir) + optimizer_checkpoint_metadata = optimizer_reader.read_metadata() + loaded_metadata = _read_sharded_metadata( + storage_reader=optimizer_reader, + checkpoint_metadata=optimizer_checkpoint_metadata, + binding=binding, + namespace="optimizer", + limits=_limits(), + transaction_id="sharded-segment-metadata", + ) + _validate_sharded_dcp_entries( + optimizer_checkpoint_metadata, + loaded_metadata, + namespace="optimizer", + binding=binding, + limits=_limits(), + ) + loaded_fields = _load_sharded_payloads( + storage_reader=optimizer_reader, + checkpoint_metadata=optimizer_checkpoint_metadata, + metadata=loaded_metadata, + metadata_tensor=prepared_envelope[_DCP_SHARDED_METADATA_KEY].clone(), + binding=binding, + namespace="optimizer", + limits=_limits(), + transaction_id="sharded-segment-load", + context_digest=bytes.fromhex( + loaded_metadata["completion"]["metadata_digest"] + ), + ) + + result_queue.put( + { + "rank": rank, + "canonical": torch.equal(canonical.cpu(), expected_canonical), + "dcp": torch.equal(loaded_canonical.cpu(), expected_canonical), + "projected": torch.equal(projected.cpu(), expected_projected), + "empty_source": empty_local.numel() == (2 if rank == 0 else 0), + "empty_canonical": torch.equal( + empty_canonical.cpu(), + empty_global.reshape(-1)[rank : rank + 1], + ), + "large_shard1": torch.equal( + large_round_trip.cpu(), + large_local.reshape(-1), + ), + "prepared_schema": ( + len(prepared_field_keys) == 3 + and prepared_envelope[_DCP_SHARDED_METADATA_KEY].dtype is torch.uint8 + and prepared_envelope[_DCP_SHARDED_METADATA_KEY].numel() > 0 + and all(size <= 256 for size in prepared_local_sizes) + and len(decoded_metadata["fields"]) == 3 + and decoded_metadata["parameters"]["model.weight"]["state_variant"] + == "initialized_dense" + and loaded_metadata["completion"] == decoded_metadata["completion"] + and len(loaded_fields) == 3 + ), + } + ) + 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="sharded portable-DCP redistribution requires Gloo", +) +def test_dtensor_dimension_shards_redistribute_through_canonical_chunks(tmp_path): + context = mp.get_context("spawn") + result_queue = context.Queue() + descriptor, init_file = tempfile.mkstemp(prefix="gefen-portable-sharded-segments-") + os.close(descriptor) + os.unlink(init_file) + checkpoint_dir = str(tmp_path / "checkpoint") + processes = [ + context.Process( + target=_worker, + args=(rank, 2, init_file, checkpoint_dir, result_queue), + ) + for rank in range(2) + ] + results = [] + try: + for process in processes: + process.start() + try: + for _ in processes: + results.append(result_queue.get(timeout=90)) + 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) == 2, (results, [process.exitcode for process in processes]) + assert all(process.exitcode == 0 for process in processes) + assert all("fatal_error" not in result for result in results), results + assert all(all(value is True for key, value in result.items() if key != "rank") for result in results) diff --git a/tests/test_portable_dcp_sharded_validation.py b/tests/test_portable_dcp_sharded_validation.py new file mode 100644 index 0000000..2c1f17a --- /dev/null +++ b/tests/test_portable_dcp_sharded_validation.py @@ -0,0 +1,685 @@ +"""Strict resource and DCP-metadata validation for sharded portable state.""" + +import copy +from dataclasses import replace + +import pytest +import torch +from torch.distributed.checkpoint.metadata import ( + ChunkStorageMetadata, + Metadata, + TensorProperties, + TensorStorageMetadata, +) + +from gefen import ( + CheckpointProcessGroupBinding, + CodebookProcessGroupBinding, + Gefen, + GefenMuon, + LogicalSlice, + ParameterIdentity, + ParameterLayout, + ParameterRebinding, + PlacementKind, + ProcessGroupIdentity, + ShardIdentity, + ShardPlacement, + ShardingManifest, +) +from gefen.portable_dcp import _flat_key +from gefen.portable_dcp_sharded import ( + _DCP_SHARDED_METADATA_KEY, + _decode_metadata_tensor, + _field_key, + _normalize_sharded_metadata, + _prepare_sharded_save_state, + _source_chunk_assignments, + _tensor_sha256, + _validate_field_byte_limits, + _validate_integrity_assignment_limits, + _validate_sharded_dcp_entries, + _validate_source_storage_chunk_limit, +) +from gefen.portable_state import PortableStateLimits +from gefen.portable_schema import portable_state_digest + + +def _binding(): + identity = ProcessGroupIdentity("sharded_validation", ("rank:0",)) + return CheckpointProcessGroupBinding( + identity, + "rank:0", + None, + torch.device("cpu"), + ) + + +def _limits(*, fragment=4096, collective=4096, containers=16): + return PortableStateLimits( + max_fragment_tensor_bytes=fragment, + max_collective_tensor_bytes=collective, + max_collective_metadata_bytes=4096, + max_metadata_bytes=4096, + max_container_items=containers, + max_tree_nodes=64, + max_tensors=16, + ) + + +def _semantic_limits(): + return PortableStateLimits( + max_fragment_tensor_bytes=4096, + max_collective_tensor_bytes=4096, + max_collective_metadata_bytes=16 << 10, + max_metadata_bytes=16 << 10, + max_container_items=1024, + max_tree_nodes=4096, + max_tensors=16, + ) + + +def _field(index, numel): + return { + "index": index, + "key": _field_key(index), + "fqn": "model.parameter{}".format(index), + "name": "momentum", + "shape": [numel], + "numel": numel, + "dtype": "float32", + "kind": "dense", + "nonnegative": False, + "source_chunk_sha256": ["00" * 32], + } + + +def _entry(dtype, numel, chunks): + return TensorStorageMetadata( + properties=TensorProperties(dtype=dtype), + size=torch.Size((numel,)), + chunks=chunks, + ) + + +def _checkpoint(fields, field_chunks): + namespace = "optimizer" + metadata_key = _flat_key(namespace, _DCP_SHARDED_METADATA_KEY) + entries = { + metadata_key: _entry( + torch.uint8, + 1, + [ + ChunkStorageMetadata( + offsets=torch.Size((0,)), + sizes=torch.Size((1,)), + ) + ], + ) + } + for field, chunks in zip(fields, field_chunks): + entries[_flat_key(namespace, field["key"])] = _entry( + torch.float32, + field["numel"], + chunks, + ) + return Metadata( + state_dict_metadata=entries, + planner_data={key: tuple(key.split(".", 1)) for key in entries}, + ) + + +@pytest.fixture(scope="module") +def _initialized_v2_metadata(): + group = ProcessGroupIdentity("sharded_validation", ("rank:0",)) + parameter = torch.nn.Parameter(torch.tensor([1.0, 2.0])) + optimizer = Gefen( + [("weight", parameter)], + fused=False, + factored_v_2d=False, + period_one_substrings=("weight",), + deterministic=True, + ) + identity = ParameterIdentity("weight", (2,)) + shard = ShardIdentity( + identity, + ParameterLayout.REPLICATED, + LogicalSlice.full(identity), + placements=( + ShardPlacement( + "checkpoint", + PlacementKind.REPLICATE, + 0, + 1, + ), + ), + process_group=group, + local_member="rank:0", + ) + optimizer.post_sharding( + (ParameterRebinding(parameter, parameter, shard),), + manifest=ShardingManifest((shard,)), + codebook_process_group=CodebookProcessGroupBinding( + group, + "rank:0", + None, + torch.device("cpu"), + ), + ) + optimizer._gefen_global_step = 4 + optimizer._gefen_codebook = torch.linspace(-1.0, 1.0, 256) + optimizer.state[parameter].update( + { + "automatic_period": 1, + "step": 4, + "m_codebook": torch.tensor([[0], [255]], dtype=torch.uint8), + "m_magnitude": torch.tensor([[1.0], [2.0]]), + "vmean": torch.tensor([[3.0], [4.0]]), + "vmean_step": 3, + } + ) + limits = _semantic_limits() + state = _prepare_sharded_save_state( + optimizer, + binding=CheckpointProcessGroupBinding( + group, + "rank:0", + None, + torch.device("cpu"), + ), + transaction_id="sharded-validation-v2", + limits=limits, + namespace="optimizer", + ) + return _decode_metadata_tensor( + state["optimizer"][_DCP_SHARDED_METADATA_KEY], + limits=limits, + ) + + +def _different_digest(value): + replacement = "0" if value[0] != "0" else "1" + return replacement + value[1:] + + +def _refresh_metadata_digest(metadata): + metadata["completion"]["metadata_digest"] = portable_state_digest( + {key: value for key, value in metadata.items() if key != "completion"} + ) + + +@pytest.fixture(scope="module") +def _factored_v2_metadata(): + group = ProcessGroupIdentity("sharded_factored_validation", ("rank:0",)) + parameter = torch.nn.Parameter(torch.arange(1, 7, dtype=torch.float32).reshape(2, 3)) + optimizer = Gefen( + [("matrix", parameter)], + fused=False, + factored_v_2d=True, + force_2d_period_one=True, + deterministic=True, + ) + identity = ParameterIdentity("matrix", (2, 3)) + shard = ShardIdentity( + identity, + ParameterLayout.REPLICATED, + LogicalSlice.full(identity), + placements=(ShardPlacement("checkpoint", PlacementKind.REPLICATE, 0, 1),), + process_group=group, + local_member="rank:0", + ) + optimizer.post_sharding( + (ParameterRebinding(parameter, parameter, shard),), + manifest=ShardingManifest((shard,)), + codebook_process_group=CodebookProcessGroupBinding(group, "rank:0", None, torch.device("cpu")), + ) + optimizer._gefen_global_step = 4 + optimizer._gefen_codebook = torch.linspace(-1.0, 1.0, 256) + optimizer.state[parameter].update( + { + "automatic_period": 1, + "step": 4, + "m_codebook": torch.tensor([[0], [255], [0], [255], [0], [255]], dtype=torch.uint8), + "m_magnitude": torch.arange(1, 7, dtype=torch.float32).reshape(-1, 1), + "v_row": torch.tensor([1.0, 3.0], dtype=torch.float32), + "v_col": torch.tensor([2.0, 4.0, 8.0], dtype=torch.float32), + "factored_step": 3, + } + ) + limits = _semantic_limits() + state = _prepare_sharded_save_state( + optimizer, + binding=CheckpointProcessGroupBinding(group, "rank:0", None, torch.device("cpu")), + transaction_id="sharded-factored-validation-v2", + limits=limits, + namespace="optimizer", + ) + return _decode_metadata_tensor(state["optimizer"][_DCP_SHARDED_METADATA_KEY], limits=limits) + + +@pytest.fixture(scope="module") +def _muon_pristine_v2_metadata(): + group = ProcessGroupIdentity("sharded_muon_validation", ("rank:0",)) + parameter = torch.nn.Parameter(torch.arange(1, 7, dtype=torch.float32).reshape(2, 3)) + optimizer = GefenMuon( + [("matrix", parameter)], + fused=False, + sharded_mode="exact", + deterministic=True, + ) + identity = ParameterIdentity("matrix", (2, 3)) + shard = ShardIdentity( + identity, + ParameterLayout.REPLICATED, + LogicalSlice.full(identity), + placements=(ShardPlacement("checkpoint", PlacementKind.REPLICATE, 0, 1),), + process_group=group, + local_member="rank:0", + ) + optimizer.post_sharding( + (ParameterRebinding(parameter, parameter, shard),), + manifest=ShardingManifest((shard,)), + codebook_process_group=CodebookProcessGroupBinding(group, "rank:0", None, torch.device("cpu")), + ) + limits = _semantic_limits() + state = _prepare_sharded_save_state( + optimizer, + binding=CheckpointProcessGroupBinding(group, "rank:0", None, torch.device("cpu")), + transaction_id="sharded-muon-validation-v2", + limits=limits, + namespace="optimizer", + ) + return _decode_metadata_tensor(state["optimizer"][_DCP_SHARDED_METADATA_KEY], limits=limits) + + +def _mutate_v2_metadata(metadata, case): + parameter = metadata["parameters"]["weight"] + if case == "unknown-key": + metadata["unknown"] = None + elif case == "negative-global-step": + metadata["common"]["gefen_global_step"] = -1 + elif case == "step-exceeds-global": + parameter["state"]["step"] = metadata["common"]["gefen_global_step"] + 1 + elif case == "unsupported-variant": + parameter["state_variant"] = "initialized_unknown" + elif case == "duplicate-field-reference": + parameter["state"]["second_moment_field"] = parameter["state"]["momentum_field"] + elif case == "unreferenced-field": + field = copy.deepcopy(metadata["fields"][1]) + field.update( + { + "index": len(metadata["fields"]), + "key": _field_key(len(metadata["fields"])), + "name": "unused", + } + ) + metadata["fields"].append(field) + elif case == "incompatible-field-reference": + parameter["state"]["momentum_field"] = parameter["state"]["second_moment_field"] + elif case == "manifest-process-group-mismatch": + metadata["source_manifest"]["shards"][0]["process_group"]["semantic_name"] = "different_group" + elif case == "short-chunk-digest": + metadata["fields"][1]["source_chunk_sha256"][0] = "00" * 31 + elif case == "nonhex-chunk-digest": + metadata["fields"][1]["source_chunk_sha256"][0] = "gg" * 32 + elif case == "metadata-digest-mismatch": + completion = metadata["completion"] + completion["metadata_digest"] = _different_digest(completion["metadata_digest"]) + elif case == "payload-digest-mismatch": + completion = metadata["completion"] + completion["payload_digest"] = _different_digest(completion["payload_digest"]) + else: + raise AssertionError("unknown mutation case") + + +def test_v2_semantic_metadata_fixture_is_valid(_initialized_v2_metadata): + metadata = copy.deepcopy(_initialized_v2_metadata) + assert _normalize_sharded_metadata(metadata, limits=_semantic_limits()) == metadata + + +@pytest.mark.parametrize( + ("case", "message"), + [ + ("unknown-key", "invalid schema"), + ("negative-global-step", "bounded exact counter"), + ("step-exceeds-global", "counters are invalid"), + ("unsupported-variant", "state_variant is unsupported"), + ("duplicate-field-reference", "referenced exactly once"), + ("unreferenced-field", "unreferenced fields"), + ("incompatible-field-reference", "incompatible field"), + ("manifest-process-group-mismatch", "different process group"), + ("short-chunk-digest", "chunk digest is invalid"), + ("nonhex-chunk-digest", "chunk digest is invalid"), + ("metadata-digest-mismatch", "metadata digest mismatch"), + ("payload-digest-mismatch", "payload digest mismatch"), + ], +) +def test_v2_semantic_metadata_rejects_malformed_envelopes( + _initialized_v2_metadata, + case, + message, +): + metadata = copy.deepcopy(_initialized_v2_metadata) + _mutate_v2_metadata(metadata, case) + with pytest.raises(ValueError, match=message): + _normalize_sharded_metadata(metadata, limits=_semantic_limits()) + + +@pytest.mark.parametrize( + ("mutation", "message"), + [ + ("missing-reference", "invalid schema"), + ("duplicate-reference", "referenced exactly once"), + ("wrong-name", "incompatible field"), + ("wrong-shape", "incompatible field"), + ], +) +def test_v2_factored_denominator_reference_is_strict( + _factored_v2_metadata, + mutation, + message, +): + metadata = copy.deepcopy(_factored_v2_metadata) + state = metadata["parameters"]["matrix"]["state"] + denominator_index = state["v_row_denominator_field"] + if mutation == "missing-reference": + state.pop("v_row_denominator_field") + elif mutation == "duplicate-reference": + state["v_row_denominator_field"] = state["v_row_field"] + elif mutation == "wrong-name": + metadata["fields"][denominator_index]["name"] = "other_denominator" + elif mutation == "wrong-shape": + metadata["fields"][denominator_index]["shape"] = [2] + metadata["fields"][denominator_index]["numel"] = 2 + else: + raise AssertionError("unknown denominator mutation") + _refresh_metadata_digest(metadata) + with pytest.raises(ValueError, match=message): + _normalize_sharded_metadata(metadata, limits=_semantic_limits()) + + +def test_v2_factored_metadata_emits_digest_protected_row_denominator(_factored_v2_metadata): + metadata = copy.deepcopy(_factored_v2_metadata) + state = metadata["parameters"]["matrix"]["state"] + field = metadata["fields"][state["v_row_denominator_field"]] + assert field["name"] == "v_row_denominator" + assert field["shape"] == [1] + assert field["kind"] == "special" + assert field["nonnegative"] is True + assert _normalize_sharded_metadata(metadata, limits=_semantic_limits()) == metadata + + +def test_v2_plain_source_rejects_whole_owner_layout(_initialized_v2_metadata): + metadata = copy.deepcopy(_initialized_v2_metadata) + shard = metadata["source_manifest"]["shards"][0] + shard["layout"] = "whole_parameter_owner" + shard["placements"][0]["kind"] = "whole_parameter_owner" + shard["owner"] = "rank:0" + _refresh_metadata_digest(metadata) + with pytest.raises(ValueError, match="source layout is unsupported"): + _normalize_sharded_metadata(metadata, limits=_semantic_limits()) + + +def test_v2_plain_factored_source_requires_replicated_layout(_factored_v2_metadata): + metadata = copy.deepcopy(_factored_v2_metadata) + shard = metadata["source_manifest"]["shards"][0] + shard["layout"] = "flattened_element_shard" + shard["placements"][0]["kind"] = "flat_shard" + _refresh_metadata_digest(metadata) + with pytest.raises(ValueError, match="factored logical matrices|factored sharded portable"): + _normalize_sharded_metadata(metadata, limits=_semantic_limits()) + + +def test_v2_plain_factored_source_policy_must_match_catalog(_factored_v2_metadata): + metadata = copy.deepcopy(_factored_v2_metadata) + metadata["policy"]["factored_v_2d"] = False + _refresh_metadata_digest(metadata) + with pytest.raises(ValueError, match="second_moment_policy conflicts"): + _normalize_sharded_metadata(metadata, limits=_semantic_limits()) + + +def test_v2_muon_source_rejects_dtensor_even_when_pristine(_muon_pristine_v2_metadata): + metadata = copy.deepcopy(_muon_pristine_v2_metadata) + shard = metadata["source_manifest"]["shards"][0] + shard["layout"] = "dtensor_1d_default_world" + shard["logical_extent"] = { + "kind": "logical_region", + "offsets": [0, 0], + "lengths": [2, 3], + } + _refresh_metadata_digest(metadata) + with pytest.raises(ValueError, match="source layout is unsupported"): + _normalize_sharded_metadata(metadata, limits=_semantic_limits()) + + +def test_v2_muon_whole_owner_source_requires_distributed_mode(_muon_pristine_v2_metadata): + metadata = copy.deepcopy(_muon_pristine_v2_metadata) + shard = metadata["source_manifest"]["shards"][0] + shard["layout"] = "whole_parameter_owner" + shard["logical_extent"] = { + "kind": "logical_slice", + "flat_offset": 0, + "length": 6, + } + shard["placements"][0]["kind"] = "whole_parameter_owner" + shard["owner"] = "rank:0" + _refresh_metadata_digest(metadata) + with pytest.raises(ValueError, match="requires sharded_mode='distributed'"): + _normalize_sharded_metadata(metadata, limits=_semantic_limits()) + + +def test_v2_muon_source_parameters_must_be_logical_matrices(_muon_pristine_v2_metadata): + metadata = copy.deepcopy(_muon_pristine_v2_metadata) + metadata["parameters"]["matrix"]["identity"]["global_shape"] = [6] + shard = metadata["source_manifest"]["shards"][0] + shard["parameter"]["global_shape"] = [6] + shard["logical_extent"]["length"] = 6 + _refresh_metadata_digest(metadata) + with pytest.raises(ValueError, match="logical matrices"): + _normalize_sharded_metadata(metadata, limits=_semantic_limits()) + + +def test_sharded_save_rejects_int64_index_route_above_fragment_limit(): + group = ProcessGroupIdentity("sharded_save_route_validation", ("rank:0",)) + parameter = torch.nn.Parameter(torch.ones((20, 20), dtype=torch.float32)) + optimizer = GefenMuon( + [("matrix", parameter)], + fused=False, + sharded_mode="exact", + deterministic=True, + ) + identity = ParameterIdentity("matrix", (20, 20)) + shard = ShardIdentity( + identity, + ParameterLayout.REPLICATED, + LogicalSlice.full(identity), + placements=(ShardPlacement("checkpoint", PlacementKind.REPLICATE, 0, 1),), + process_group=group, + local_member="rank:0", + ) + optimizer.post_sharding( + (ParameterRebinding(parameter, parameter, shard),), + manifest=ShardingManifest((shard,)), + codebook_process_group=CodebookProcessGroupBinding(group, "rank:0", None, torch.device("cpu")), + ) + optimizer._gefen_global_step = 2 + optimizer._gefen_codebook = torch.linspace(-1.0, 1.0, 256) + optimizer.state[parameter].update( + { + "automatic_period": 1, + "step": 2, + "m_codebook": torch.zeros((400, 1), dtype=torch.uint8), + "m_magnitude": torch.ones((400, 1), dtype=torch.float32), + } + ) + limits = replace(_semantic_limits(), max_fragment_tensor_bytes=2800) + with pytest.raises(RuntimeError, match="save routing scratch"): + _prepare_sharded_save_state( + optimizer, + binding=CheckpointProcessGroupBinding(group, "rank:0", None, torch.device("cpu")), + transaction_id="sharded-save-route-limit", + limits=limits, + namespace="optimizer", + ) + + +def test_field_byte_limits_are_aggregate_not_per_tensor(): + fields = [_field(0, 200), _field(1, 200)] + with pytest.raises(ValueError, match="max_collective_tensor_bytes"): + _validate_field_byte_limits( + fields, + binding=_binding(), + limits=_limits(fragment=1024, collective=1024), + ) + + +def test_dcp_chunk_limit_is_aggregate_across_fields(): + fields = [_field(0, 4), _field(1, 4)] + split = [ + ChunkStorageMetadata( + offsets=torch.Size((0,)), + sizes=torch.Size((2,)), + ), + ChunkStorageMetadata( + offsets=torch.Size((2,)), + sizes=torch.Size((2,)), + ), + ] + checkpoint = _checkpoint(fields, [split, split]) + with pytest.raises(ValueError, match="aggregate limits"): + _validate_sharded_dcp_entries( + checkpoint, + {"fields": fields}, + namespace="optimizer", + binding=_binding(), + limits=_limits(containers=3), + ) + + +def test_dcp_chunk_limit_includes_mandatory_metadata_chunk(): + fields = [_field(0, 4), _field(1, 4)] + complete = [ + ChunkStorageMetadata( + offsets=torch.Size((0,)), + sizes=torch.Size((4,)), + ) + ] + checkpoint = _checkpoint(fields, [complete, complete]) + with pytest.raises(ValueError, match="aggregate limits"): + _validate_sharded_dcp_entries( + checkpoint, + {"fields": fields}, + namespace="optimizer", + binding=_binding(), + limits=_limits(containers=2), + ) + + +def test_dcp_chunk_geometry_accepts_trailing_empty_shards(): + fields = [_field(0, 1)] + chunks = [ + ChunkStorageMetadata( + offsets=torch.Size((0,)), + sizes=torch.Size((1,)), + ), + ChunkStorageMetadata( + offsets=torch.Size((1,)), + sizes=torch.Size((0,)), + ), + ] + checkpoint = _checkpoint(fields, [chunks]) + _validate_sharded_dcp_entries( + checkpoint, + {"fields": fields}, + namespace="optimizer", + binding=_binding(), + limits=_limits(), + ) + + +@pytest.mark.parametrize( + ("geometry", "message"), + [ + (((0, -1),), "exceeds its tensor"), + (((-1, 0), (0, 4)), "exceeds its tensor"), + (((0, 4), (5, 0)), "exceeds its tensor"), + (((0, 2), (3, 1)), "cover each field exactly once"), + (((0, 3), (2, 2)), "cover each field exactly once"), + ], +) +def test_dcp_chunk_geometry_rejects_invalid_boundaries( + geometry, + message, +): + fields = [_field(0, 4)] + chunks = [ + ChunkStorageMetadata( + offsets=torch.Size((offset,)), + sizes=torch.Size((length,)), + ) + for offset, length in geometry + ] + checkpoint = _checkpoint(fields, [chunks]) + with pytest.raises(ValueError, match=message): + _validate_sharded_dcp_entries( + checkpoint, + {"fields": fields}, + namespace="optimizer", + binding=_binding(), + limits=_limits(), + ) + + +def test_dcp_chunk_geometry_requires_exact_torch_size_and_ints(): + fields = [_field(0, 4)] + inexact = [ + ChunkStorageMetadata( + offsets=[0.0], + sizes=[4.0], + ) + ] + checkpoint = _checkpoint(fields, [inexact]) + with pytest.raises(ValueError, match="chunk is invalid"): + _validate_sharded_dcp_entries( + checkpoint, + {"fields": fields}, + namespace="optimizer", + binding=_binding(), + limits=_limits(), + ) + + +def test_integrity_reassembly_rejects_packed_assignment_above_fragment_limit(): + field = _field(0, 200) + field["source_chunk_sha256"] = ["00" * 32, "11" * 32, "22" * 32] + with pytest.raises(ValueError, match="integrity scratch"): + _validate_integrity_assignment_limits( + field, + target_parts=2, + limits=_limits(fragment=400, collective=4096), + ) + + +def test_integrity_assignment_accounts_for_six_element_packed_reassembly(): + assignments = _source_chunk_assignments(10, 3, 2) + packed_numel = [ + max( + (segment.local_offset + segment.length for _source, segment in items), + default=0, + ) + for items in assignments + ] + assert packed_numel == [6, 4] + + +def test_dcp_source_storage_chunk_must_fit_current_fragment_limit(): + with pytest.raises(ValueError, match="source storage chunk"): + _validate_source_storage_chunk_limit( + 8, + limits=_limits(fragment=16, collective=4096), + ) + + +def test_integrity_hash_accepts_scalar_and_empty_tensors(): + assert len(_tensor_sha256(torch.tensor(3.5))) == 32 + assert len(_tensor_sha256(torch.empty(0))) == 32 diff --git a/tests/test_portable_dcp_topologies.py b/tests/test_portable_dcp_topologies.py index 9dc6a44..3edcab9 100644 --- a/tests/test_portable_dcp_topologies.py +++ b/tests/test_portable_dcp_topologies.py @@ -1,5 +1,7 @@ """Warning-strict CPU/Gloo topology coverage for portable DCP.""" +import copy +from dataclasses import replace from datetime import timedelta import multiprocessing as mp import os @@ -17,6 +19,10 @@ from gefen.checkpoint import CheckpointProcessGroupBinding from gefen.codebook import CodebookProcessGroupBinding from gefen.contracts import ( + CheckpointProjectionQualifier, + CheckpointStateRepresentation, + CheckpointStateTransitionKind, + CheckpointTransport, LogicalSlice, ParameterIdentity, ParameterLayout, @@ -28,6 +34,7 @@ ) from gefen.gefen import Gefen from gefen.gefen_muon import GefenMuon +from gefen.hybrid import GefenMuonHybrid from gefen.portable import _decode_quantized_momentum from gefen.portable_state import PortableStateLimits from gefen.rebinding import ParameterRebinding @@ -37,6 +44,34 @@ _MUON_FQN = "model.matrix" +class _NoTouchFileSystemWriter(FileSystemWriter): + def __init__(self, path): + super().__init__(path) + self.touched = False + + def _unexpected_io(self): + self.touched = True + raise AssertionError("portable DCP writer was touched during save preparation") + + def reset(self, *_args, **_kwargs): + self._unexpected_io() + + def set_up_storage_writer(self, *_args, **_kwargs): + self._unexpected_io() + + def prepare_local_plan(self, *_args, **_kwargs): + self._unexpected_io() + + def prepare_global_plan(self, *_args, **_kwargs): + self._unexpected_io() + + def write_data(self, *_args, **_kwargs): + self._unexpected_io() + + def finish(self, *_args, **_kwargs): + self._unexpected_io() + + def _limits(): return PortableStateLimits( max_fragment_tensor_bytes=1 << 20, @@ -56,6 +91,10 @@ def _limits(): ) +def _container_limits(max_container_items): + return replace(_limits(), max_container_items=max_container_items) + + def _strict_worker_warnings(): warnings.simplefilter("error") warnings.filterwarnings( @@ -169,6 +208,35 @@ def _plain_optimizer(parameter, *, deterministic): ) +def _scalar_optimizer(parameter, *, deterministic): + return Gefen( + [("scalar", parameter)], + lr=2.5e-3, + betas=(0.8, 0.97), + eps=2.0e-8, + weight_decay=0.03, + fused=False, + factored_v_2d=False, + period_one_substrings=("scalar",), + deterministic=deterministic, + ) + + +def _projection_optimizer(parameter, *, factored, 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_1d_period_one=True, + force_2d_period_one=True, + factored_v_2d=factored, + deterministic=deterministic, + ) + + def _muon_optimizer(parameter, *, deterministic, normuon): return GefenMuon( [("matrix", parameter)], @@ -291,6 +359,38 @@ def _seed_plain_state(optimizer, parameter, momentum, second_moment): ) +def _seed_scalar_state(optimizer, parameter): + optimizer._gefen_global_step = 4 + optimizer._gefen_codebook = _codebook() + optimizer.state[parameter].update( + { + "automatic_period": 1, + "step": 4, + "m_codebook": torch.tensor([[255]], dtype=torch.uint8), + "m_magnitude": torch.tensor([[2.0]], dtype=torch.float32), + "vmean": torch.tensor([[3.0]], dtype=torch.float32), + "vmean_step": 3, + } + ) + + +def _seed_factored_state(optimizer, parameter, momentum, row, column): + 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, + "v_row": row.clone(), + "v_col": column.clone(), + "factored_step": 6, + } + ) + + def _seed_muon_state(optimizer, parameter, *, normuon): indices, magnitudes = _quantized_period_one(_muon_momentum()) optimizer._gefen_global_step = 13 @@ -320,8 +420,8 @@ def _bits_equal(left, 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), + left.detach().contiguous().reshape(-1).view(torch.uint8), + right.detach().contiguous().reshape(-1).view(torch.uint8), ) ) @@ -358,6 +458,14 @@ def _muon_reference(*, normuon): return optimizer, parameter +def _capture_runtime_error(operation): + try: + operation() + except RuntimeError as exc: + return str(exc) + raise AssertionError("portable DCP operation unexpectedly succeeded") + + def _run_phase(worker, world_size, *worker_args, timeout=180): context = mp.get_context("spawn") result_queue = context.Queue() @@ -722,6 +830,973 @@ def _muon_owner_load_worker( dist.destroy_process_group() +def _scalar_empty_chunk_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("rank:{}".format(index) for index in range(world_size)) + group = ProcessGroupIdentity("portable_dcp_scalar_empty_chunk", members) + identity = ParameterIdentity("model.scalar", ()) + manifest, shards = _replicated_manifest(identity, group) + codebook_binding, checkpoint_binding = _bindings( + group, + members[rank], + dist.group.WORLD, + ) + + source_parameter = torch.nn.Parameter(torch.tensor(2.0)) + source = _scalar_optimizer(source_parameter, deterministic=True) + source.post_sharding( + (ParameterRebinding(source_parameter, source_parameter, shards[rank]),), + manifest=manifest, + codebook_process_group=codebook_binding, + ) + _seed_scalar_state(source, source_parameter) + save_portable_dcp( + source, + checkpoint_process_group=checkpoint_binding, + storage_writer=FileSystemWriter(checkpoint_dir), + transaction_id="scalar-empty-chunk-save", + limits=_limits(), + ) + dist.barrier() + + entries = FileSystemReader(checkpoint_dir).read_metadata().state_dict_metadata + field_entries = [ + entry + for key, entry in entries.items() + if key.startswith("optimizer.__gefen_portable_field_") + ] + small_geometries = [ + sorted( + (int(chunk.offsets[0]), int(chunk.sizes[0])) + for chunk in entry.chunks + ) + for entry in field_entries + if tuple(entry.size) == (1,) + ] + v2_sharded = ( + "optimizer.__gefen_portable_metadata_v2__" in entries + and "optimizer.__gefen_portable_metadata_v1__" not in entries + and bool(field_entries) + and all(len(entry.chunks) == world_size for entry in field_entries) + ) + trailing_empty_chunks = bool(small_geometries) and all( + geometry == [(0, 1), (1, 0)] for geometry in small_geometries + ) + + target_parameter = torch.nn.Parameter(torch.tensor(2.0)) + target = _scalar_optimizer(target_parameter, deterministic=False) + target.post_sharding( + (ParameterRebinding(target_parameter, target_parameter, shards[rank]),), + manifest=manifest, + codebook_process_group=codebook_binding, + ) + load_portable_dcp( + target, + checkpoint_process_group=checkpoint_binding, + storage_reader=FileSystemReader(checkpoint_dir), + transaction_id="scalar-empty-chunk-load", + limits=_limits(), + ) + restored_exact = ( + _state_equal(target.state[target_parameter], source.state[source_parameter]) + and target._gefen_global_step == source._gefen_global_step == 4 + and target._deterministic is source._deterministic is True + and _bits_equal(target._gefen_codebook, source._gefen_codebook) + ) + + gradient = torch.tensor(0.25) + target_parameter.grad = gradient.clone() + source_parameter.grad = gradient.clone() + target.step() + source.step() + continuation_exact = ( + _bits_equal(target_parameter, source_parameter) + and _state_equal(target.state[target_parameter], source.state[source_parameter]) + and target._gefen_global_step == source._gefen_global_step == 5 + and _bits_equal(target._gefen_codebook, source._gefen_codebook) + ) + + mismatch_parameter = torch.nn.Parameter(torch.tensor(-7.0)) + mismatch = _scalar_optimizer(mismatch_parameter, deterministic=False) + mismatch.post_sharding( + (ParameterRebinding(mismatch_parameter, mismatch_parameter, shards[rank]),), + manifest=manifest, + codebook_process_group=codebook_binding, + ) + if rank == 0: + mismatch.param_groups[0]["eps"] = 1.0e-7 + before = mismatch._canonical_import_live_token() + state_object = mismatch.state[mismatch_parameter] + mismatch_message = _capture_runtime_error( + lambda: load_portable_dcp( + mismatch, + checkpoint_process_group=checkpoint_binding, + storage_reader=FileSystemReader(checkpoint_dir), + transaction_id="scalar-empty-chunk-asymmetric-target", + limits=_limits(), + ) + ) + mismatch_unchanged = ( + mismatch._canonical_import_live_token() == before + and mismatch.state[mismatch_parameter] is state_object + and mismatch.state[mismatch_parameter] == {"name": "scalar"} + and mismatch._gefen_global_step == 0 + and mismatch._gefen_codebook is None + ) + + identity_mismatch = ParameterIdentity("model.scalar", (1,)) + identity_mismatch_manifest, identity_mismatch_shards = _replicated_manifest( + identity_mismatch, + group, + ) + identity_mismatch_parameter = torch.nn.Parameter(torch.tensor([-11.0])) + identity_mismatch_target = _scalar_optimizer( + identity_mismatch_parameter, + deterministic=False, + ) + identity_mismatch_target.post_sharding( + ( + ParameterRebinding( + identity_mismatch_parameter, + identity_mismatch_parameter, + identity_mismatch_shards[rank], + ), + ), + manifest=identity_mismatch_manifest, + codebook_process_group=codebook_binding, + ) + identity_before = identity_mismatch_target._canonical_import_live_token() + identity_state_object = identity_mismatch_target.state[ + identity_mismatch_parameter + ] + identity_parameter_before = identity_mismatch_parameter.detach().clone() + identity_mismatch_message = _capture_runtime_error( + lambda: load_portable_dcp( + identity_mismatch_target, + checkpoint_process_group=checkpoint_binding, + storage_reader=FileSystemReader(checkpoint_dir), + transaction_id="scalar-empty-chunk-identity-mismatch", + limits=_limits(), + ) + ) + identity_mismatch_unchanged = ( + identity_mismatch_target._canonical_import_live_token() + == identity_before + and identity_mismatch_target.state[identity_mismatch_parameter] + is identity_state_object + and identity_mismatch_target.state[identity_mismatch_parameter] + == {"name": "scalar"} + and _bits_equal( + identity_mismatch_parameter, + identity_parameter_before, + ) + and identity_mismatch_target._gefen_global_step == 0 + and identity_mismatch_target._gefen_codebook is None + ) + result_queue.put( + { + "rank": rank, + "v2_sharded": v2_sharded, + "small_field_count": len(small_geometries), + "trailing_empty_chunks": trailing_empty_chunks, + "restored_exact": restored_exact, + "continuation_exact": continuation_exact, + "mismatch_message": mismatch_message, + "mismatch_unchanged": mismatch_unchanged, + "identity_mismatch_message": identity_mismatch_message, + "identity_mismatch_unchanged": identity_mismatch_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 _factored_projection_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("rank:{}".format(index) for index in range(world_size)) + group = ProcessGroupIdentity("portable_dcp_factored_projection", members) + identity = ParameterIdentity(_PLAIN_FQN, (2, 4)) + replicated_manifest, replicated_shards = _replicated_manifest(identity, group) + codebook_binding, checkpoint_binding = _bindings( + group, + members[rank], + dist.group.WORLD, + ) + full_parameter = _plain_initial_parameter() + momentum = torch.tensor( + [[-0.5, 0.75, -1.0, 1.25], [1.5, -1.75, 2.0, -2.25]], + dtype=torch.float32, + ) + row = torch.tensor([1.0, 3.0], dtype=torch.float32) + column = torch.tensor([2.0, 4.0, 8.0, 16.0], dtype=torch.float32) + gradient = _plain_gradient() + + source_parameter = torch.nn.Parameter(full_parameter.clone()) + source = _projection_optimizer( + source_parameter, + factored=True, + deterministic=True, + ) + source.post_sharding( + ( + ParameterRebinding( + source_parameter, + source_parameter, + replicated_shards[rank], + ), + ), + manifest=replicated_manifest, + codebook_process_group=codebook_binding, + ) + _seed_factored_state(source, source_parameter, momentum, row, column) + + canonical_support = next( + support + for support in source.optimizer_contract().capabilities.checkpoints + if support.transport is CheckpointTransport.CANONICAL_GLOBAL + ) + transition = next( + item + for item in canonical_support.state_transitions + if item.source is CheckpointStateRepresentation.FACTORED_SECOND_MOMENT + and item.target is CheckpointStateRepresentation.BLOCK_SECOND_MOMENT + ) + declared_projection = ( + transition.kind is CheckpointStateTransitionKind.DEFINED_PROJECTION + and transition.qualifier + is CheckpointProjectionQualifier.FACTORED_TO_BLOCK_LIVE_FP32_TARGET_PERIOD_ONE_V1 + and ParameterLayout.FLATTENED_ELEMENT_SHARD + in transition.topology_changing + ) + save_portable_dcp( + source, + checkpoint_process_group=checkpoint_binding, + storage_writer=FileSystemWriter(checkpoint_dir), + transaction_id="factored-to-flat-block-save", + limits=_limits(), + ) + dist.barrier() + entries = FileSystemReader(checkpoint_dir).read_metadata().state_dict_metadata + v2_sharded = ( + "optimizer.__gefen_portable_metadata_v2__" in entries + and "optimizer.__gefen_portable_metadata_v1__" not in entries + and any( + key.startswith("optimizer.__gefen_portable_field_") + for key in entries + ) + ) + + lengths = (3, 5) + flat_manifest, flat_shards = _flat_manifest(identity, group, lengths) + shard = flat_shards[rank] + start = shard.logical_slice.flat_offset + stop = start + shard.logical_slice.length + local_initial = full_parameter.reshape(-1)[start:stop].clone() + original = torch.nn.Parameter(full_parameter.clone()) + target_parameter = torch.nn.Parameter(local_initial.clone()) + target = _projection_optimizer( + original, + factored=False, + deterministic=False, + ) + target.post_sharding( + (ParameterRebinding(original, target_parameter, shard),), + manifest=flat_manifest, + codebook_process_group=codebook_binding, + ) + + from gefen.portable_dcp_sharded import ( + _load_sharded_payloads, + _metadata_storage_spec, + _normalize_sharded_metadata, + _read_sharded_metadata, + _stage_sharded_import, + ) + from gefen.portable_schema import portable_state_digest + from gefen.portable_state import _SECOND_MOMENT_PROJECTION_EXACT + + invalid_reader = FileSystemReader(checkpoint_dir) + checkpoint_metadata = invalid_reader.read_metadata() + source_metadata = _read_sharded_metadata( + storage_reader=invalid_reader, + checkpoint_metadata=checkpoint_metadata, + binding=checkpoint_binding, + namespace="optimizer", + limits=_limits(), + transaction_id="factored-to-flat-block-invalid-qualifier-metadata", + ) + source_factored_state = source_metadata["parameters"][_PLAIN_FQN]["state"] + denominator_field = source_metadata["fields"][ + source_factored_state["v_row_denominator_field"] + ] + v2_denominator = ( + denominator_field["name"] == "v_row_denominator" + and denominator_field["shape"] == [1] + and denominator_field["kind"] == "special" + and denominator_field["nonnegative"] is True + ) + metadata_shape = _metadata_storage_spec( + checkpoint_metadata, + namespace="optimizer", + limits=_limits(), + ) + invalid_local_fields = _load_sharded_payloads( + storage_reader=invalid_reader, + checkpoint_metadata=checkpoint_metadata, + metadata=source_metadata, + metadata_tensor=torch.empty(metadata_shape, dtype=torch.uint8), + binding=checkpoint_binding, + namespace="optimizer", + limits=_limits(), + transaction_id="factored-to-flat-block-invalid-qualifier-payload", + context_digest=bytes.fromhex( + source_metadata["completion"]["metadata_digest"] + ), + ) + invalid_metadata = copy.deepcopy(source_metadata) + invalid_metadata["policy"][ + "second_moment_projection" + ] = _SECOND_MOMENT_PROJECTION_EXACT + invalid_semantic = { + key: value + for key, value in invalid_metadata.items() + if key != "completion" + } + invalid_metadata["completion"]["metadata_digest"] = portable_state_digest( + invalid_semantic + ) + invalid_metadata = _normalize_sharded_metadata( + invalid_metadata, + limits=_limits(), + ) + invalid_before = target._canonical_import_live_token() + invalid_state_object = target.state[target_parameter] + invalid_parameter_before = target_parameter.detach().clone() + invalid_message = _capture_runtime_error( + lambda: _stage_sharded_import( + target, + invalid_metadata, + invalid_local_fields, + binding=checkpoint_binding, + limits=_limits(), + transaction_id="factored-to-flat-block-invalid-qualifier", + ) + ) + invalid_unchanged = ( + target._canonical_import_live_token() == invalid_before + and target.state[target_parameter] is invalid_state_object + and target.state[target_parameter] == {"name": "weight"} + and _bits_equal(target_parameter, invalid_parameter_before) + and target._gefen_global_step == 0 + and target._gefen_codebook is None + and target._deterministic is False + ) + + load_portable_dcp( + target, + checkpoint_process_group=checkpoint_binding, + storage_reader=FileSystemReader(checkpoint_dir), + transaction_id="factored-to-flat-block-load", + limits=_limits(), + ) + + expected_second = torch.outer(row, column).div_( + row.mean().clamp_(min=torch.finfo(torch.float32).tiny) + ) + expected_local_momentum = momentum.reshape(-1)[start:stop].clone() + expected_local_second = expected_second.reshape(-1)[start:stop].clone() + target_state = target.state[target_parameter] + decoded = _decode_quantized_momentum( + target._gefen_codebook, + target_state["m_codebook"], + target_state["m_magnitude"], + logical_shape=(lengths[rank],), + period=1, + step=target_state["step"], + ) + projection_exact = ( + _bits_equal(decoded, expected_local_momentum) + and _bits_equal( + target_state["vmean"].reshape(-1), + expected_local_second, + ) + and target_state["step"] == 7 + and target_state["vmean_step"] == 6 + and target._gefen_global_step == 9 + and target._deterministic is True + ) + + oracle_original = torch.nn.Parameter(full_parameter.clone()) + oracle_parameter = torch.nn.Parameter(local_initial.clone()) + oracle = _projection_optimizer( + oracle_original, + factored=False, + deterministic=True, + ) + oracle.post_sharding( + (ParameterRebinding(oracle_original, oracle_parameter, shard),), + manifest=flat_manifest, + codebook_process_group=codebook_binding, + ) + _seed_plain_state( + oracle, + oracle_parameter, + expected_local_momentum, + expected_local_second, + ) + imported_state_exact = _state_equal( + target.state[target_parameter], + oracle.state[oracle_parameter], + ) + local_gradient = gradient.reshape(-1)[start:stop].clone() + target_parameter.grad = local_gradient.clone() + oracle_parameter.grad = local_gradient.clone() + target.step() + oracle.step() + continuation_exact = ( + _bits_equal(target_parameter, oracle_parameter) + and _state_equal( + target.state[target_parameter], + oracle.state[oracle_parameter], + ) + and target._gefen_global_step == oracle._gefen_global_step == 10 + and _bits_equal(target._gefen_codebook, oracle._gefen_codebook) + ) + result_queue.put( + { + "rank": rank, + "v2_sharded": v2_sharded, + "v2_denominator": v2_denominator, + "declared_projection": declared_projection, + "length": lengths[rank], + "invalid_message": invalid_message, + "invalid_unchanged": invalid_unchanged, + "projection_exact": projection_exact, + "imported_state_exact": imported_state_exact, + "continuation_exact": continuation_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 _stage_projection_failure_worker( + rank, + world_size, + init_file, + checkpoint_dir, + failure_site, + 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("rank:{}".format(index) for index in range(world_size)) + group = ProcessGroupIdentity("portable_dcp_stage_failure_{}".format(failure_site), members) + identity = ParameterIdentity(_PLAIN_FQN, (2, 4)) + replicated_manifest, replicated_shards = _replicated_manifest(identity, group) + codebook_binding, checkpoint_binding = _bindings(group, members[rank], dist.group.WORLD) + full_parameter = _plain_initial_parameter() + momentum = _plain_momentum() + + source_parameter = torch.nn.Parameter(full_parameter.clone()) + source = _projection_optimizer( + source_parameter, + factored=failure_site == "factored_projection", + deterministic=True, + ) + source.post_sharding( + (ParameterRebinding(source_parameter, source_parameter, replicated_shards[rank]),), + manifest=replicated_manifest, + codebook_process_group=codebook_binding, + ) + if failure_site == "factored_projection": + _seed_factored_state( + source, + source_parameter, + momentum, + torch.tensor([1.0, 3.0], dtype=torch.float32), + torch.tensor([2.0, 4.0, 8.0, 16.0], dtype=torch.float32), + ) + else: + _seed_plain_state(source, source_parameter, momentum, _plain_second_moment()) + save_portable_dcp( + source, + checkpoint_process_group=checkpoint_binding, + storage_writer=FileSystemWriter(checkpoint_dir), + transaction_id="stage-failure-{}-save".format(failure_site), + limits=_limits(), + ) + dist.barrier() + + flat_manifest, flat_shards = _flat_manifest(identity, group, (3, 5)) + shard = flat_shards[rank] + start = shard.logical_slice.flat_offset + stop = start + shard.logical_slice.length + original = torch.nn.Parameter(full_parameter.clone()) + target_parameter = torch.nn.Parameter(full_parameter.reshape(-1)[start:stop].clone()) + target = _projection_optimizer(original, factored=False, deterministic=False) + target.post_sharding( + (ParameterRebinding(original, target_parameter, shard),), + manifest=flat_manifest, + codebook_process_group=codebook_binding, + ) + before_token = target._canonical_import_live_token() + before_state = target.state[target_parameter] + before_parameter = target_parameter.detach().clone() + + if rank == 0: + def injected_failure(*_args, **_kwargs): + raise ValueError("injected {} failure".format(failure_site)) + + if failure_site == "momentum_recompress": + from gefen import portable as portable_module + + portable_module._recompress_dense_momentum = injected_failure + else: + from gefen import portable_dcp_sharded as sharded_module + + sharded_module._local_factored_projection = injected_failure + + message = _capture_runtime_error( + lambda: load_portable_dcp( + target, + checkpoint_process_group=checkpoint_binding, + storage_reader=FileSystemReader(checkpoint_dir), + transaction_id="stage-failure-{}-load".format(failure_site), + limits=_limits(), + ) + ) + heartbeat = torch.tensor([rank + 1], dtype=torch.int64) + dist.all_reduce(heartbeat) + unchanged = ( + target._canonical_import_live_token() == before_token + and target.state[target_parameter] is before_state + and target.state[target_parameter] == {"name": "weight"} + and _bits_equal(target_parameter, before_parameter) + and target._gefen_global_step == 0 + and target._gefen_codebook is None + and target._deterministic is False + ) + result_queue.put( + { + "rank": rank, + "message": message, + "unchanged": unchanged, + "heartbeat": heartbeat.item() == 3, + } + ) + 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 _projected_target_limit_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("rank:{}".format(index) for index in range(world_size)) + group = ProcessGroupIdentity("portable_dcp_projected_target_limit", members) + identity = ParameterIdentity(_PLAIN_FQN, (2, 4)) + replicated_manifest, replicated_shards = _replicated_manifest(identity, group) + codebook_binding, checkpoint_binding = _bindings(group, members[rank], dist.group.WORLD) + full_parameter = _plain_initial_parameter() + source_parameter = torch.nn.Parameter(full_parameter.clone()) + source = _plain_optimizer(source_parameter, deterministic=True) + source.post_sharding( + (ParameterRebinding(source_parameter, source_parameter, replicated_shards[rank]),), + manifest=replicated_manifest, + codebook_process_group=codebook_binding, + ) + _seed_plain_state(source, source_parameter, _plain_momentum(), _plain_second_moment()) + save_portable_dcp( + source, + checkpoint_process_group=checkpoint_binding, + storage_writer=FileSystemWriter(checkpoint_dir), + transaction_id="projected-target-limit-save", + limits=_limits(), + ) + dist.barrier() + + flat_manifest, flat_shards = _flat_manifest(identity, group, (8, 0)) + shard = flat_shards[rank] + start = shard.logical_slice.flat_offset + stop = start + shard.logical_slice.length + original = torch.nn.Parameter(full_parameter.clone()) + target_parameter = torch.nn.Parameter(full_parameter.reshape(-1)[start:stop].clone()) + target = _plain_optimizer(original, deterministic=False) + target.post_sharding( + (ParameterRebinding(original, target_parameter, shard),), + manifest=flat_manifest, + codebook_process_group=codebook_binding, + ) + before_token = target._canonical_import_live_token() + before_state = target.state[target_parameter] + before_parameter = target_parameter.detach().clone() + message = _capture_runtime_error( + lambda: load_portable_dcp( + target, + checkpoint_process_group=checkpoint_binding, + storage_reader=FileSystemReader(checkpoint_dir), + transaction_id="projected-target-limit-load", + limits=replace(_limits(), max_fragment_tensor_bytes=800), + ) + ) + heartbeat = torch.tensor([rank + 1], dtype=torch.int64) + dist.all_reduce(heartbeat) + result_queue.put( + { + "rank": rank, + "message": message, + "unchanged": ( + target._canonical_import_live_token() == before_token + and target.state[target_parameter] is before_state + and target.state[target_parameter] == {"name": "weight"} + and _bits_equal(target_parameter, before_parameter) + and target._gefen_global_step == 0 + and target._gefen_codebook is None + ), + "heartbeat": heartbeat.item() == 3, + } + ) + 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 _save_chunk_limit_worker( + rank, + world_size, + init_file, + checkpoint_root, + 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), + ) + members = tuple("rank:{}".format(index) for index in range(world_size)) + group = ProcessGroupIdentity("portable_dcp_save_chunk_limits", members) + codebook_binding, checkpoint_binding = _bindings( + group, + members[rank], + dist.group.WORLD, + ) + + first = torch.nn.Parameter( + torch.tensor([[0.5, -1.0], [1.5, -2.0]], dtype=torch.float32) + ) + second = torch.nn.Parameter( + torch.tensor([[2.5, -3.0], [3.5, -4.0]], dtype=torch.float32) + ) + plain = Gefen( + [("first", first), ("second", second)], + 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=True, + ) + first_manifest, first_shards = _replicated_manifest( + ParameterIdentity("model.first", (2, 2)), + group, + ) + second_manifest, second_shards = _replicated_manifest( + ParameterIdentity("model.second", (2, 2)), + group, + ) + plain.post_sharding( + ( + ParameterRebinding(first, first, first_shards[rank]), + ParameterRebinding(second, second, second_shards[rank]), + ), + manifest=ShardingManifest( + first_manifest.shards + second_manifest.shards + ), + codebook_process_group=codebook_binding, + ) + _seed_plain_state( + plain, + first, + torch.tensor( + [[-0.5, 0.75], [-1.0, 1.25]], + dtype=torch.float32, + ), + torch.tensor( + [[0.25, 0.5], [0.75, 1.0]], + dtype=torch.float32, + ), + ) + _seed_plain_state( + plain, + second, + torch.tensor( + [[1.5, -1.75], [2.0, -2.25]], + dtype=torch.float32, + ), + torch.tensor( + [[1.25, 1.5], [1.75, 2.0]], + dtype=torch.float32, + ), + ) + plain_writer_path = "{}-plain".format(checkpoint_root) + plain_writer = _NoTouchFileSystemWriter(plain_writer_path) + plain_before = plain._canonical_import_live_token() + plain_message = _capture_runtime_error( + lambda: save_portable_dcp( + plain, + checkpoint_process_group=checkpoint_binding, + storage_writer=plain_writer, + transaction_id="plain-save-expected-chunk-limit", + limits=_container_limits(10), + ) + ) + plain_unchanged = plain._canonical_import_live_token() == plain_before + dist.barrier() + + matrix = torch.nn.Parameter(_muon_initial_parameter().clone()) + bias_a = torch.nn.Parameter( + torch.tensor([0.5, -1.0, 1.5, -2.0], dtype=torch.float32) + ) + bias_b = torch.nn.Parameter( + torch.tensor([2.5, -3.0, 3.5, -4.0], dtype=torch.float32) + ) + hybrid = GefenMuonHybrid( + [("matrix", matrix)], + [("bias_a", bias_a), ("bias_b", bias_b)], + 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=True, + normuon=True, + normuon_beta2=0.9, + normuon_eps=3.0e-8, + ) + matrix_manifest, matrix_shards = _replicated_manifest( + ParameterIdentity("model.matrix", (3, 2)), + group, + ) + bias_a_manifest, bias_a_shards = _replicated_manifest( + ParameterIdentity("model.bias_a", (4,)), + group, + ) + bias_b_manifest, bias_b_shards = _replicated_manifest( + ParameterIdentity("model.bias_b", (4,)), + group, + ) + hybrid.post_sharding( + ( + ParameterRebinding(matrix, matrix, matrix_shards[rank]), + ParameterRebinding(bias_a, bias_a, bias_a_shards[rank]), + ParameterRebinding(bias_b, bias_b, bias_b_shards[rank]), + ), + manifest=ShardingManifest( + matrix_manifest.shards + + bias_a_manifest.shards + + bias_b_manifest.shards + ), + codebook_process_group=codebook_binding, + ) + _seed_muon_state(hybrid.muon, matrix, normuon=True) + _seed_plain_state( + hybrid.backup, + bias_a, + torch.tensor([-0.5, 0.75, -1.0, 1.25], dtype=torch.float32), + torch.tensor([0.25, 0.5, 0.75, 1.0], dtype=torch.float32), + ) + _seed_plain_state( + hybrid.backup, + bias_b, + torch.tensor([1.5, -1.75, 2.0, -2.25], dtype=torch.float32), + torch.tensor([1.25, 1.5, 1.75, 2.0], dtype=torch.float32), + ) + hybrid.backup._gefen_global_step = 13 + from gefen import portable_dcp_sharded as sharded_module + + original_decode = sharded_module._decode_metadata_tensor + if rank == 0: + def fail_first_child_decode(*_args, **_kwargs): + raise ValueError("injected first Hybrid child decode failure") + + sharded_module._decode_metadata_tensor = fail_first_child_decode + phase_writer_path = "{}-hybrid-phase".format(checkpoint_root) + phase_writer = _NoTouchFileSystemWriter(phase_writer_path) + phase_before = ( + hybrid.muon._canonical_import_live_token(), + hybrid.backup._canonical_import_live_token(), + ) + hybrid_phase_message = _capture_runtime_error( + lambda: save_portable_dcp( + hybrid, + checkpoint_process_group=checkpoint_binding, + storage_writer=phase_writer, + transaction_id="hybrid-save-first-child-decode-failure", + limits=_limits(), + ) + ) + sharded_module._decode_metadata_tensor = original_decode + phase_heartbeat = torch.tensor([rank + 1], dtype=torch.int64) + dist.all_reduce(phase_heartbeat) + hybrid_phase_unchanged = phase_before == ( + hybrid.muon._canonical_import_live_token(), + hybrid.backup._canonical_import_live_token(), + ) + + projected_path = "{}-hybrid-projected".format(checkpoint_root) + save_portable_dcp( + hybrid, + checkpoint_process_group=checkpoint_binding, + storage_writer=FileSystemWriter(projected_path), + transaction_id="hybrid-projected-limit-save", + limits=_limits(), + ) + dist.barrier() + projected_before = ( + hybrid.muon._canonical_import_live_token(), + hybrid.backup._canonical_import_live_token(), + ) + hybrid_projected_message = _capture_runtime_error( + lambda: load_portable_dcp( + hybrid, + checkpoint_process_group=checkpoint_binding, + storage_reader=FileSystemReader(projected_path), + transaction_id="hybrid-projected-limit-load", + limits=replace(_limits(), max_fragment_tensor_bytes=1500), + ) + ) + hybrid_projected_unchanged = projected_before == ( + hybrid.muon._canonical_import_live_token(), + hybrid.backup._canonical_import_live_token(), + ) + + hybrid_writer_path = "{}-hybrid".format(checkpoint_root) + hybrid_writer = _NoTouchFileSystemWriter(hybrid_writer_path) + hybrid_before = ( + hybrid.muon._canonical_import_live_token(), + hybrid.backup._canonical_import_live_token(), + ) + hybrid_message = _capture_runtime_error( + lambda: save_portable_dcp( + hybrid, + checkpoint_process_group=checkpoint_binding, + storage_writer=hybrid_writer, + transaction_id="hybrid-save-expected-chunk-limit", + limits=_container_limits(18), + ) + ) + hybrid_unchanged = hybrid_before == ( + hybrid.muon._canonical_import_live_token(), + hybrid.backup._canonical_import_live_token(), + ) + result_queue.put( + { + "rank": rank, + "plain_message": plain_message, + "plain_writer_untouched": ( + not plain_writer.touched + and not os.path.exists(plain_writer_path) + ), + "plain_unchanged": plain_unchanged, + "hybrid_message": hybrid_message, + "hybrid_phase_message": hybrid_phase_message, + "hybrid_phase_writer_untouched": ( + not phase_writer.touched + and not os.path.exists(phase_writer_path) + ), + "hybrid_phase_unchanged": hybrid_phase_unchanged, + "hybrid_phase_heartbeat": phase_heartbeat.item() == 3, + "hybrid_projected_message": hybrid_projected_message, + "hybrid_projected_unchanged": hybrid_projected_unchanged, + "hybrid_writer_untouched": ( + not hybrid_writer.touched + and not os.path.exists(hybrid_writer_path) + ), + "hybrid_unchanged": hybrid_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 _subgroup_worker( rank, world_size, @@ -902,6 +1977,133 @@ def test_portable_dcp_muon_replicated_singleton_to_world_owner(tmp_path, normuon 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 empty-chunk coverage requires Gloo", +) +def test_portable_dcp_two_rank_scalar_fields_accept_trailing_empty_chunks_and_continue(tmp_path): + results = _run_phase( + _scalar_empty_chunk_worker, + 2, + str(tmp_path / "scalar-empty-chunk"), + ) + assert all("fatal_error" not in result for result in results), results + assert all(result["v2_sharded"] for result in results), results + assert all(result["small_field_count"] == 2 for result in results), results + assert all(result["trailing_empty_chunks"] for result in results), results + assert all(result["restored_exact"] and result["continuation_exact"] for result in results), results + assert all(result["mismatch_unchanged"] for result in results), results + mismatch_messages = [result["mismatch_message"] for result in results] + assert mismatch_messages[0] == mismatch_messages[1] + assert "algorithm options do not match" in mismatch_messages[0] + assert all(result["identity_mismatch_unchanged"] for result in results), results + identity_mismatch_messages = [ + result["identity_mismatch_message"] for result in results + ] + assert identity_mismatch_messages[0] == identity_mismatch_messages[1] + assert "parameter identity does not match the target" in identity_mismatch_messages[0] + + +@pytest.mark.skipif( + not dist.is_available() or not dist.is_gloo_available(), + reason="portable DCP defined-projection coverage requires Gloo", +) +def test_portable_dcp_projects_replicated_factored_state_to_uneven_flat_block_shards(tmp_path): + results = _run_phase( + _factored_projection_worker, + 2, + str(tmp_path / "factored-to-flat-block"), + ) + assert all("fatal_error" not in result for result in results), results + assert [result["length"] for result in results] == [3, 5] + invalid_messages = [result["invalid_message"] for result in results] + assert invalid_messages[0] == invalid_messages[1] + assert "source policy does not authorize factored-to-block" in invalid_messages[0] + assert all( + result["v2_sharded"] + and result["v2_denominator"] + and result["declared_projection"] + and result["invalid_unchanged"] + and result["projection_exact"] + and result["imported_state_exact"] + and result["continuation_exact"] + for result in results + ), results + + +@pytest.mark.parametrize("failure_site", ["momentum_recompress", "factored_projection"]) +@pytest.mark.skipif( + not dist.is_available() or not dist.is_gloo_available(), + reason="portable DCP stage failure-protocol coverage requires Gloo", +) +def test_portable_dcp_stage_projection_failures_are_collective_and_atomic(tmp_path, failure_site): + results = _run_phase( + _stage_projection_failure_worker, + 2, + str(tmp_path / failure_site), + failure_site, + ) + assert all("fatal_error" not in result for result in results), results + messages = [result["message"] for result in results] + assert messages[0] == messages[1] + assert "injected {} failure".format(failure_site) in messages[0] + assert all(result["unchanged"] and result["heartbeat"] for result in results), results + + +@pytest.mark.skipif( + not dist.is_available() or not dist.is_gloo_available(), + reason="portable DCP projected target-limit coverage requires Gloo", +) +def test_portable_dcp_rejects_uneven_target_projected_state_before_payload_read(tmp_path): + results = _run_phase( + _projected_target_limit_worker, + 2, + str(tmp_path / "projected-target-limit"), + ) + assert all("fatal_error" not in result for result in results), results + messages = [result["message"] for result in results] + assert messages[0] == messages[1] + assert "projected target state exceeds max_fragment_tensor_bytes" in messages[0] + assert all(result["unchanged"] and result["heartbeat"] for result in results), results + + +@pytest.mark.skipif( + not dist.is_available() or not dist.is_gloo_available(), + reason="portable DCP save-limit coverage requires Gloo", +) +def test_portable_dcp_save_rejects_expected_plain_and_hybrid_chunks_before_writer_io(tmp_path): + results = _run_phase( + _save_chunk_limit_worker, + 2, + str(tmp_path / "save-chunk-limits"), + timeout=240, + ) + assert all("fatal_error" not in result for result in results), results + plain_messages = [result["plain_message"] for result in results] + assert plain_messages[0] == plain_messages[1] + assert "sharded portable DCP chunks exceed aggregate limits" in plain_messages[0] + hybrid_messages = [result["hybrid_message"] for result in results] + assert hybrid_messages[0] == hybrid_messages[1] + assert "sharded Hybrid DCP chunks exceed aggregate limits" in hybrid_messages[0] + phase_messages = [result["hybrid_phase_message"] for result in results] + assert phase_messages[0] == phase_messages[1] + assert "injected first Hybrid child decode failure" in phase_messages[0] + projected_messages = [result["hybrid_projected_message"] for result in results] + assert projected_messages[0] == projected_messages[1] + assert "Hybrid projected target state exceeds max_fragment_tensor_bytes" in projected_messages[0] + assert all( + result["plain_writer_untouched"] + and result["plain_unchanged"] + and result["hybrid_writer_untouched"] + and result["hybrid_unchanged"] + and result["hybrid_phase_writer_untouched"] + and result["hybrid_phase_unchanged"] + and result["hybrid_phase_heartbeat"] + and result["hybrid_projected_unchanged"] + for result in results + ), results + + @pytest.mark.skipif( not dist.is_available() or not dist.is_gloo_available(), reason="portable DCP subgroup coverage requires Gloo", diff --git a/tests/test_portable_dtensor.py b/tests/test_portable_dtensor.py new file mode 100644 index 0000000..c3552b9 --- /dev/null +++ b/tests/test_portable_dtensor.py @@ -0,0 +1,1152 @@ +"""Two-rank portable-v3 coverage for exact public DTensor storage.""" + +from __future__ import annotations + +import copy +from datetime import timedelta +import math +import multiprocessing as mp +import os +import queue as queue_module +import shutil +import tempfile +import time +import traceback + +import pytest +import torch +import torch.distributed as dist + + +_WORLD = 2 +_FQN = "model.weight" + + +def _limits(): + from gefen import PortableStateLimits + + return PortableStateLimits( + max_fragment_tensor_bytes=2 << 20, + max_collective_tensor_bytes=8 << 20, + max_collective_metadata_bytes=8 << 20, + chunk_bytes=29, + max_members=4, + max_metadata_bytes=2 << 20, + max_tree_nodes=20_000, + max_tree_depth=32, + max_container_items=20_000, + max_string_bytes=16 << 10, + max_integer_bytes=128, + max_tensors=512, + max_tensor_rank=8, + diagnostic_bytes=1024, + ) + + +def _members(): + return tuple("rank:{}".format(rank) for rank in range(_WORLD)) + + +def _bindings(group, rank): + from gefen import CheckpointProcessGroupBinding, CodebookProcessGroupBinding + + member = _members()[rank] + return ( + CodebookProcessGroupBinding(group, member, dist.group.WORLD, torch.device("cpu")), + CheckpointProcessGroupBinding(group, member, dist.group.WORLD, torch.device("cpu")), + ) + + +def _standard_region(shape, coordinate, dimension): + from gefen import LogicalRegion + + chunk = (shape[dimension] + _WORLD - 1) // _WORLD + offset = min(coordinate * chunk, shape[dimension]) + length = max(0, min(shape[dimension], offset + chunk) - offset) + offsets = [0] * len(shape) + lengths = list(shape) + offsets[dimension] = offset + lengths[dimension] = length + return LogicalRegion(tuple(offsets), tuple(lengths)) + + +def _dtensor_manifest(identity, group, *, placement_kind, dimension=None): + from gefen import ( + LogicalRegion, + ParameterLayout, + PlacementKind, + ShardIdentity, + ShardPlacement, + ShardingManifest, + ) + + shards = [] + for coordinate, member in enumerate(group.ordered_members): + region = ( + LogicalRegion.full(identity) + if placement_kind is PlacementKind.REPLICATE + else _standard_region(identity.global_shape, coordinate, dimension) + ) + shards.append( + ShardIdentity( + identity, + ParameterLayout.DTENSOR_1D_DEFAULT_WORLD, + region, + placements=( + ShardPlacement( + "dp", + placement_kind, + coordinate, + _WORLD, + parameter_dimension=dimension, + ), + ), + process_group=group, + local_member=member, + ) + ) + return ShardingManifest(tuple(shards)), tuple(shards) + + +def _replicated_manifest(identity, group): + from gefen import LogicalSlice, ParameterLayout, PlacementKind, ShardIdentity, ShardPlacement, ShardingManifest + + shards = tuple( + ShardIdentity( + identity, + ParameterLayout.REPLICATED, + LogicalSlice.full(identity), + placements=(ShardPlacement("dp", PlacementKind.REPLICATE, coordinate, _WORLD),), + process_group=group, + local_member=member, + ) + for coordinate, member in enumerate(group.ordered_members) + ) + return ShardingManifest(shards), shards + + +def _flat_manifest(identity, group): + from gefen import LogicalSlice, ParameterLayout, PlacementKind, ShardIdentity, ShardPlacement, ShardingManifest + + first = max(0, identity.numel // 3) + lengths = (first, identity.numel - first) + 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, _WORLD),), + process_group=group, + local_member=member, + ) + ) + offset += length + return ShardingManifest(tuple(shards)), tuple(shards) + + +def _topology_manifest(identity, group, topology): + from gefen import PlacementKind + + if topology == "replicated": + return _replicated_manifest(identity, group) + if topology == "flat": + return _flat_manifest(identity, group) + if topology == "dtensor_shard0": + return _dtensor_manifest(identity, group, placement_kind=PlacementKind.DIMENSION_SHARD, dimension=0) + if topology == "dtensor_shard1": + return _dtensor_manifest(identity, group, placement_kind=PlacementKind.DIMENSION_SHARD, dimension=1) + if topology == "dtensor_replicate": + return _dtensor_manifest(identity, group, placement_kind=PlacementKind.REPLICATE) + raise AssertionError("unknown topology") + + +def _region_index(region): + return tuple(slice(offset, offset + length) for offset, length in zip(region.offsets, region.lengths)) + + +def _local_projection(value, shard): + from gefen import ParameterLayout, PlacementKind + + if shard.layout is ParameterLayout.REPLICATED: + return value.clone() + if shard.layout is ParameterLayout.FLATTENED_ELEMENT_SHARD: + start = shard.logical_slice.flat_offset + return value.reshape(-1)[start : start + shard.logical_slice.length].clone() + if shard.layout is ParameterLayout.DTENSOR_1D_DEFAULT_WORLD: + if shard.placements[0].kind is PlacementKind.REPLICATE: + return value.clone() + return value[_region_index(shard.logical_region)].clone() + raise AssertionError("unsupported test projection") + + +def _local_tensor(parameter): + value = parameter.to_local() if hasattr(parameter, "to_local") else parameter + if hasattr(value, "wait"): + value = value.wait() + return value + + +def _make_optimizer(rank, mesh, group, topology, full_parameter, *, factored_v_2d=False, deterministic=True): + from gefen import Gefen, ParameterIdentity, ParameterLayout, ParameterRebinding + from torch import nn + from torch.distributed.tensor import Replicate, Shard, distribute_tensor + + identity = ParameterIdentity(_FQN, tuple(full_parameter.shape)) + manifest, shards = _topology_manifest(identity, group, topology) + shard = shards[rank] + if topology.startswith("dtensor_"): + if topology == "dtensor_shard0": + placement = Shard(0) + elif topology == "dtensor_shard1": + placement = Shard(1) + else: + placement = Replicate() + parameter = nn.Parameter(distribute_tensor(full_parameter.clone(), mesh, [placement])) + original = parameter + elif topology == "replicated": + parameter = nn.Parameter(full_parameter.clone()) + original = parameter + else: + original = nn.Parameter(full_parameter.clone()) + parameter = nn.Parameter(_local_projection(full_parameter, shard)) + 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=factored_v_2d, + deterministic=deterministic, + ) + codebook_binding, checkpoint_binding = _bindings(group, rank) + optimizer.post_sharding( + (ParameterRebinding(original, parameter, shard),), + manifest=manifest, + codebook_process_group=codebook_binding, + ) + if shard.layout is ParameterLayout.DTENSOR_1D_DEFAULT_WORLD: + assert tuple(_local_tensor(parameter).shape) == shard.logical_region.lengths + return optimizer, parameter, shard, checkpoint_binding + + +def _codebook(): + return torch.linspace(-1.0, 1.0, 256, dtype=torch.float32) + + +def _quantize_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_block_state(optimizer, parameter, shard, momentum, second_moment): + optimizer._gefen_global_step = 9 + optimizer._gefen_codebook = _codebook() + local_momentum = _local_projection(momentum, shard) + if local_momentum.numel() == 0: + return + local_second = _local_projection(second_moment, shard) + indices, magnitudes = _quantize_period_one(local_momentum) + optimizer.state[parameter].update( + { + "automatic_period": 1, + "step": 7, + "m_codebook": indices, + "m_magnitude": magnitudes, + "vmean": local_second.reshape(-1, 1), + "vmean_step": 6, + } + ) + + +def _seed_factored_state(optimizer, parameter, momentum, row, column): + optimizer._gefen_global_step = 9 + optimizer._gefen_codebook = _codebook() + indices, magnitudes = _quantize_period_one(momentum) + optimizer.state[parameter].update( + { + "automatic_period": 1, + "step": 7, + "m_codebook": indices, + "m_magnitude": magnitudes, + "v_row": row.clone(), + "v_col": column.clone(), + "factored_step": 6, + } + ) + + +def _decode_local_state(optimizer, parameter, shard): + from gefen import ParameterLayout + from gefen.portable import _decode_quantized_momentum + + state = optimizer.state[parameter] + if shard.layout is ParameterLayout.DTENSOR_1D_DEFAULT_WORLD and shard.logical_region.numel == 0: + return None, None + momentum = _decode_quantized_momentum( + optimizer._gefen_codebook, + state["m_codebook"], + state["m_magnitude"], + logical_shape=tuple(_local_tensor(parameter).shape), + period=1, + step=state["step"], + ) + return momentum, state["vmean"].reshape(_local_tensor(parameter).shape) + + +def _bits_equal(left, right): + return ( + left.dtype == right.dtype + and tuple(left.shape) == tuple(right.shape) + and torch.equal(left.detach().contiguous().reshape(-1).view(torch.uint8), right.detach().contiguous().reshape(-1).view(torch.uint8)) + ) + + +def _assign_gradient(parameter, shard, mesh, full_gradient): + from gefen import ParameterLayout + from torch.distributed.tensor import Replicate, Shard, distribute_tensor + + if shard.layout is ParameterLayout.DTENSOR_1D_DEFAULT_WORLD: + descriptor = shard.placements[0] + placement = Replicate() if descriptor.parameter_dimension is None else Shard(descriptor.parameter_dimension) + parameter.grad = distribute_tensor(full_gradient.clone(), mesh, [placement]) + else: + parameter.grad = _local_projection(full_gradient, shard) + + +def _state_snapshot(optimizer, parameter, *, include_containers=False): + def clone(value): + if torch.is_tensor(value): + return value.detach().clone() + if type(value) is dict: + return {key: clone(item) for key, item in value.items()} + if type(value) is list: + return [clone(item) for item in value] + if type(value) is tuple: + return tuple(clone(item) for item in value) + return copy.deepcopy(value) + + state = { + key: clone(value) + for key, value in optimizer.state[parameter].items() + } + snapshot = { + "global_step": optimizer._gefen_global_step, + "codebook": None if optimizer._gefen_codebook is None else optimizer._gefen_codebook.detach().clone(), + "deterministic": optimizer._deterministic, + "parameter": _local_tensor(parameter).detach().clone(), + "state": state, + } + if include_containers: + snapshot["groups"] = tuple( + {key: clone(value) for key, value in group.items() if key != "params"} + for group in optimizer.param_groups + ) + snapshot["containers"] = ( + id(optimizer.state), + id(optimizer.param_groups), + id(optimizer.state[parameter]), + id(optimizer._param_names), + id(optimizer._gefen_shard_bindings), + ) + return snapshot + + +def _rank_local_checkpoint_schema_is_valid(optimizer): + state_dict = optimizer.state_dict() + metadata = state_dict["param_groups"][0]["_gefen_checkpoint_metadata"] + marker = metadata.get("rank_local_sharded_state") + expected_carriers = { + "_gefen_rank_local_payload_{}".format(rank) + for rank in range(_WORLD) + } + carrier_sets = [ + { + key + for key in state + if type(key) is str and key.startswith("_gefen_rank_local_payload_") + } + for state in state_dict["state"].values() + ] + return ( + type(marker) is dict + and marker.get("format") == "rank_local_dtensor_v2" + and carrier_sets.count(expected_carriers) == 1 + and all(not keys or keys == expected_carriers for keys in carrier_sets) + ) + + +def _live_carrier_schema_is_valid(state): + expected_keys = { + "name", + "_gefen_rank_local_payload_0", + "_gefen_rank_local_payload_1", + } + if set(state) != expected_keys: + return False + for rank in range(_WORLD): + value = state["_gefen_rank_local_payload_{}".format(rank)] + if ( + type(value) is not torch.Tensor + or value.layout is not torch.strided + or value.device.type != "cpu" + or value.dtype != torch.uint8 + or tuple(value.shape) != (1,) + or value.requires_grad + or not value.is_contiguous() + or value.storage_offset() != 0 + or value.untyped_storage().nbytes() != 1 + or int(value.item()) != 0 + ): + return False + return True + + +def _assert_nested_exact(actual, expected, path="value"): + if torch.is_tensor(expected): + assert torch.is_tensor(actual), path + assert _bits_equal(actual, expected), path + return + assert type(actual) is type(expected), path + if isinstance(expected, dict): + assert set(actual) == set(expected), path + for key in expected: + _assert_nested_exact(actual[key], expected[key], "{}.{}".format(path, key)) + elif isinstance(expected, (tuple, list)): + assert len(actual) == len(expected), path + for index, (actual_item, expected_item) in enumerate(zip(actual, expected)): + _assert_nested_exact(actual_item, expected_item, "{}[{}]".format(path, index)) + else: + assert actual == expected, path + + +def _run_transition(rank, mesh, group, *, source_topology, target_topology, shape, label, factored_projection=False): + from gefen import ParameterLayout + + full_parameter = torch.linspace(-0.8, 0.7, torch.tensor(shape).prod().item(), dtype=torch.float32).reshape(shape) + momentum = torch.linspace(-1.5, 1.25, full_parameter.numel(), dtype=torch.float32).reshape(shape) + second = torch.arange(1, full_parameter.numel() + 1, dtype=torch.float32).reshape(shape).div_(7.0) + gradient = torch.arange(full_parameter.numel(), dtype=torch.float32).reshape(shape).cos() + + source, source_parameter, source_shard, source_binding = _make_optimizer( + rank, + mesh, + group, + source_topology, + full_parameter, + factored_v_2d=factored_projection, + deterministic=True, + ) + if factored_projection: + row = torch.linspace(1.0, 2.0, shape[0], dtype=torch.float32) + column = torch.linspace(0.5, 1.5, shape[1], dtype=torch.float32) + _seed_factored_state(source, source_parameter, momentum, row, column) + projected_second = torch.outer(row, column).div_(row.mean().clamp_(min=torch.finfo(torch.float32).tiny)) + else: + _seed_block_state(source, source_parameter, source_shard, momentum, second) + projected_second = second + document = source.export_portable_state( + checkpoint_process_group=source_binding, + transaction_id="{}-export".format(label), + limits=_limits(), + ) + record = document["parameters"][_FQN] + document_exact = _bits_equal(record["state"]["momentum"], momentum) + if factored_projection: + document_exact = document_exact and _bits_equal(record["state"]["v_row"], row) and _bits_equal(record["state"]["v_col"], column) + else: + document_exact = document_exact and _bits_equal(record["state"]["second_moment"], second) + + target, target_parameter, target_shard, target_binding = _make_optimizer( + rank, + mesh, + group, + target_topology, + full_parameter, + factored_v_2d=factored_projection, + deterministic=False, + ) + target.import_portable_state( + document, + checkpoint_process_group=target_binding, + transaction_id="{}-import".format(label), + limits=_limits(), + ) + expected_local_momentum = _local_projection(momentum, target_shard) + expected_local_second = _local_projection(projected_second, target_shard) + if expected_local_momentum.numel() == 0: + projected_exact = _live_carrier_schema_is_valid( + target.state[target_parameter] + ) + else: + decoded, local_second = _decode_local_state(target, target_parameter, target_shard) + projected_exact = ( + _bits_equal(decoded, expected_local_momentum) + and _bits_equal(local_second, expected_local_second) + and target.state[target_parameter]["step"] == 7 + and target.state[target_parameter]["vmean_step"] == 6 + and target._gefen_global_step == 9 + and target._deterministic is True + ) + checkpoint_schema_exact = ( + _rank_local_checkpoint_schema_is_valid(target) + if target_shard.layout is ParameterLayout.DTENSOR_1D_DEFAULT_WORLD + else True + ) + + reference, reference_parameter, reference_shard, _ = _make_optimizer( + rank, + mesh, + group, + target_topology, + full_parameter, + factored_v_2d=factored_projection, + deterministic=True, + ) + _seed_block_state(reference, reference_parameter, reference_shard, momentum, projected_second) + _assert_nested_exact( + _state_snapshot(target, target_parameter), + _state_snapshot(reference, reference_parameter), + "{} imported state".format(label), + ) + _assign_gradient(target_parameter, target_shard, mesh, gradient) + _assign_gradient(reference_parameter, reference_shard, mesh, gradient) + target.step() + reference.step() + continuation_exact = True + try: + _assert_nested_exact( + _state_snapshot(target, target_parameter), + _state_snapshot(reference, reference_parameter), + "{} continuation".format(label), + ) + except AssertionError: + continuation_exact = False + return { + "label": label, + "document_exact": document_exact, + "projected_exact": projected_exact, + "checkpoint_schema_exact": checkpoint_schema_exact, + "continuation_exact": continuation_exact, + "empty": expected_local_momentum.numel() == 0, + }, document + + +def _mismatch_result(rank, mesh, group, document): + target_values = torch.arange(20, dtype=torch.float32).reshape(4, 5) + target, parameter, _, binding = _make_optimizer( + rank, + mesh, + group, + "dtensor_shard0", + target_values, + deterministic=False, + ) + before = _state_snapshot(target, parameter, include_containers=True) + error = None + try: + target.import_portable_state( + document, + checkpoint_process_group=binding, + transaction_id="dtensor-identity-mismatch", + limits=_limits(), + ) + except RuntimeError as exc: + error = str(exc) + unchanged = True + try: + _assert_nested_exact( + _state_snapshot(target, parameter, include_containers=True), + before, + "mismatched import", + ) + except AssertionError: + unchanged = False + return {"rejected": error is not None, "unchanged": unchanged, "error": error} + + +def _dcp_transition_result(rank, mesh, group, checkpoint_dir): + from gefen import load_portable_dcp, save_portable_dcp + from torch.distributed.checkpoint import FileSystemReader, FileSystemWriter + + shape = (4, 1) + full_parameter = torch.linspace(-0.8, 0.7, math.prod(shape), dtype=torch.float32).reshape(shape) + momentum = torch.linspace(-1.5, 1.25, math.prod(shape), dtype=torch.float32).reshape(shape) + second = torch.arange(1, math.prod(shape) + 1, dtype=torch.float32).reshape(shape).div_(7.0) + gradient = torch.arange(math.prod(shape), dtype=torch.float32).reshape(shape).cos() + source, source_parameter, source_shard, source_binding = _make_optimizer( + rank, + mesh, + group, + "dtensor_shard0", + full_parameter, + deterministic=True, + ) + _seed_block_state(source, source_parameter, source_shard, momentum, second) + + optimizer_type = type(source) + original_export = optimizer_type.export_portable_state + + def reject_dense_export(*_args, **_kwargs): + raise AssertionError("sharded DCP must not call dense portable export") + + optimizer_type.export_portable_state = reject_dense_export + try: + save_portable_dcp( + source, + checkpoint_process_group=source_binding, + storage_writer=FileSystemWriter(checkpoint_dir), + transaction_id="dtensor-dcp-save", + limits=_limits(), + ) + finally: + optimizer_type.export_portable_state = original_export + + checkpoint_metadata = FileSystemReader(checkpoint_dir).read_metadata() + entries = checkpoint_metadata.state_dict_metadata + v2_key = "optimizer.__gefen_portable_metadata_v2__" + field_entries = [ + value + for key, value in entries.items() + if key.startswith("optimizer.__gefen_portable_field_") + ] + sharded_envelope = ( + v2_key in entries + and "optimizer.__gefen_portable_metadata_v1__" not in entries + and bool(field_entries) + and all(len(entry.chunks) == _WORLD for entry in field_entries) + ) + + target, target_parameter, target_shard, target_binding = _make_optimizer( + rank, + mesh, + group, + "dtensor_shard1", + full_parameter, + deterministic=False, + ) + load_portable_dcp( + target, + checkpoint_process_group=target_binding, + storage_reader=FileSystemReader(checkpoint_dir), + transaction_id="dtensor-dcp-load", + limits=_limits(), + ) + reference, reference_parameter, reference_shard, _ = _make_optimizer( + rank, + mesh, + group, + "dtensor_shard1", + full_parameter, + deterministic=True, + ) + _seed_block_state(reference, reference_parameter, reference_shard, momentum, second) + imported_exact = True + try: + _assert_nested_exact( + _state_snapshot(target, target_parameter), + _state_snapshot(reference, reference_parameter), + "DTensor DCP imported state", + ) + except AssertionError: + imported_exact = False + _assign_gradient(target_parameter, target_shard, mesh, gradient) + _assign_gradient(reference_parameter, reference_shard, mesh, gradient) + target.step() + reference.step() + continuation_exact = True + try: + _assert_nested_exact( + _state_snapshot(target, target_parameter), + _state_snapshot(reference, reference_parameter), + "DTensor DCP continuation", + ) + except AssertionError: + continuation_exact = False + mismatch, mismatch_parameter, _mismatch_shard, mismatch_binding = _make_optimizer( + rank, + mesh, + group, + "dtensor_shard1", + full_parameter, + deterministic=False, + ) + if rank == 0: + mismatch.param_groups[0]["eps"] = 1.0e-7 + mismatch_before = _state_snapshot( + mismatch, + mismatch_parameter, + include_containers=True, + ) + mismatch_error = None + try: + load_portable_dcp( + mismatch, + checkpoint_process_group=mismatch_binding, + storage_reader=FileSystemReader(checkpoint_dir), + transaction_id="dtensor-dcp-target-mismatch", + limits=_limits(), + ) + except RuntimeError as exc: + mismatch_error = str(exc) + mismatch_unchanged = True + try: + _assert_nested_exact( + _state_snapshot( + mismatch, + mismatch_parameter, + include_containers=True, + ), + mismatch_before, + "DTensor DCP mismatched target", + ) + except AssertionError: + mismatch_unchanged = False + corrupted, corrupted_parameter, _corrupted_shard, corrupted_binding = _make_optimizer( + rank, + mesh, + group, + "dtensor_shard1", + full_parameter, + deterministic=False, + ) + corrupted_before = _state_snapshot( + corrupted, + corrupted_parameter, + include_containers=True, + ) + from gefen import portable_dcp_sharded + + original_local_tensor = portable_dcp_sharded._canonical_local_tensor + corrupted_once = False + + def corrupt_one_rank(value): + nonlocal corrupted_once + local = original_local_tensor(value) + if rank == 0 and not corrupted_once and local.dtype is torch.float32 and local.numel(): + local.reshape(-1)[0] = float("nan") + corrupted_once = True + return local + + portable_dcp_sharded._canonical_local_tensor = corrupt_one_rank + corruption_error = None + try: + load_portable_dcp( + corrupted, + checkpoint_process_group=corrupted_binding, + storage_reader=FileSystemReader(checkpoint_dir), + transaction_id="dtensor-dcp-asymmetric-corruption", + limits=_limits(), + ) + except RuntimeError as exc: + corruption_error = str(exc) + finally: + portable_dcp_sharded._canonical_local_tensor = original_local_tensor + corruption_unchanged = True + try: + _assert_nested_exact( + _state_snapshot( + corrupted, + corrupted_parameter, + include_containers=True, + ), + corrupted_before, + "DTensor DCP asymmetric corruption", + ) + except AssertionError: + corruption_unchanged = False + return { + "sharded_envelope": sharded_envelope, + "imported_exact": imported_exact, + "continuation_exact": continuation_exact, + "empty": target_shard.logical_region.numel == 0, + "mismatch_error": mismatch_error, + "mismatch_unchanged": mismatch_unchanged, + "corruption_error": corruption_error, + "corruption_unchanged": corruption_unchanged, + } + + +def _malformed_transport_result(rank, mesh, group): + values = torch.arange(20, dtype=torch.float32).reshape(5, 4) + optimizer, parameter, _, binding = _make_optimizer( + rank, + mesh, + group, + "dtensor_shard0", + values, + ) + optimizer.state[parameter]["_gefen_rank_local_payload_evil"] = torch.tensor( + [17], + dtype=torch.uint8, + ) + before = _state_snapshot(optimizer, parameter, include_containers=True) + error = None + try: + optimizer.export_portable_state( + checkpoint_process_group=binding, + transaction_id="dtensor-malformed-transport", + limits=_limits(), + ) + except RuntimeError as exc: + error = str(exc) + unchanged = True + try: + _assert_nested_exact( + _state_snapshot(optimizer, parameter, include_containers=True), + before, + "malformed transport export", + ) + except AssertionError: + unchanged = False + return {"rejected": error is not None, "unchanged": unchanged, "error": error} + + +def _mixed_factored_contract_result(rank, mesh, group): + from gefen import ( + CheckpointProjectionQualifier, + CheckpointStateRepresentation, + CheckpointStateTransitionKind, + CheckpointTransport, + Gefen, + ParameterIdentity, + ParameterLayout, + ParameterRebinding, + PlacementKind, + ShardingManifest, + ) + from torch import nn + from torch.distributed.tensor import Shard, distribute_tensor + + factored_identity = ParameterIdentity("model.factored", (2, 4)) + block_identity = ParameterIdentity("model.block", (4, 4)) + factored_manifest, factored_shards = _replicated_manifest( + factored_identity, group + ) + block_manifest, block_shards = _dtensor_manifest( + block_identity, + group, + placement_kind=PlacementKind.DIMENSION_SHARD, + dimension=0, + ) + manifest = ShardingManifest( + tuple(factored_manifest.shards) + tuple(block_manifest.shards) + ) + factored = nn.Parameter(torch.arange(8, dtype=torch.float32).reshape(2, 4)) + block_old = nn.Parameter(torch.arange(16, dtype=torch.float32).reshape(4, 4)) + block = nn.Parameter( + distribute_tensor(block_old.detach().clone(), mesh, [Shard(0)]) + ) + optimizer = Gefen( + [("factored", factored), ("block", block_old)], + fused=False, + deterministic=True, + force_2d_period_one=True, + factored_v_2d=True, + ) + codebook_binding, _ = _bindings(group, rank) + optimizer.post_sharding( + ( + ParameterRebinding(factored, factored, factored_shards[rank]), + ParameterRebinding(block_old, block, block_shards[rank]), + ), + manifest=manifest, + codebook_process_group=codebook_binding, + ) + support = next( + item + for item in optimizer.optimizer_contract().capabilities.checkpoints + if item.transport is CheckpointTransport.CANONICAL_GLOBAL + ) + factored_exact = next( + item + for item in support.state_transitions + if item.source is CheckpointStateRepresentation.FACTORED_SECOND_MOMENT + and item.target is CheckpointStateRepresentation.FACTORED_SECOND_MOMENT + ) + block_exact = next( + item + for item in support.state_transitions + if item.source is CheckpointStateRepresentation.BLOCK_SECOND_MOMENT + and item.target is CheckpointStateRepresentation.BLOCK_SECOND_MOMENT + ) + projection = next( + item + for item in support.state_transitions + if item.source is CheckpointStateRepresentation.FACTORED_SECOND_MOMENT + and item.target is CheckpointStateRepresentation.BLOCK_SECOND_MOMENT + ) + return { + "layouts": support.same_topology + == frozenset( + { + ParameterLayout.REPLICATED, + ParameterLayout.DTENSOR_1D_DEFAULT_WORLD, + } + ), + "factored_exact": ( + factored_exact.kind is CheckpointStateTransitionKind.EXACT + and factored_exact.same_topology + == frozenset({ParameterLayout.REPLICATED}) + ), + "block_exact": ( + block_exact.kind is CheckpointStateTransitionKind.EXACT + and block_exact.same_topology + == frozenset({ParameterLayout.DTENSOR_1D_DEFAULT_WORLD}) + ), + "exact_union": ( + factored_exact.same_topology | block_exact.same_topology + == support.same_topology + ), + "projection": ( + projection.kind is CheckpointStateTransitionKind.DEFINED_PROJECTION + and projection.qualifier + is CheckpointProjectionQualifier.FACTORED_TO_BLOCK_LIVE_FP32_TARGET_PERIOD_ONE_V1 + ), + } + + +def _worker(rank, init_file, result_queue): + from gefen import ProcessGroupIdentity + from torch.distributed.tensor import init_device_mesh + + try: + torch.set_num_threads(1) + dist.init_process_group( + "gloo", + init_method="file://{}".format(init_file), + rank=rank, + world_size=_WORLD, + timeout=timedelta(seconds=45), + ) + mesh = init_device_mesh("cpu", (_WORLD,), mesh_dim_names=("dp",)) + group = ProcessGroupIdentity("portable_dtensor", _members()) + cases = [] + last_document = None + for transition in ( + { + "source_topology": "dtensor_shard0", + "target_topology": "replicated", + "shape": (5, 4), + "label": "dtensor-shard0-to-replicated", + }, + { + "source_topology": "replicated", + "target_topology": "dtensor_shard1", + "shape": (4, 1), + "label": "replicated-to-dtensor-shard1-empty", + }, + { + "source_topology": "dtensor_replicate", + "target_topology": "flat", + "shape": (5, 4), + "label": "dtensor-replicate-to-flat", + }, + { + "source_topology": "replicated", + "target_topology": "dtensor_replicate", + "shape": (5, 4), + "label": "replicated-to-dtensor-replicate", + }, + { + "source_topology": "flat", + "target_topology": "dtensor_shard0", + "shape": (5, 4), + "label": "flat-to-dtensor-shard0", + }, + { + "source_topology": "dtensor_shard0", + "target_topology": "dtensor_shard1", + "shape": (5, 4), + "label": "dtensor-shard0-to-shard1", + }, + { + "source_topology": "replicated", + "target_topology": "dtensor_shard0", + "shape": (5, 4), + "label": "factored-replicated-to-dtensor-block", + "factored_projection": True, + }, + ): + result, last_document = _run_transition(rank, mesh, group, **transition) + cases.append(result) + dist.barrier() + dcp = _dcp_transition_result( + rank, + mesh, + group, + os.path.join(os.path.dirname(init_file), "portable-dcp"), + ) + dist.barrier() + malformed_transport = _malformed_transport_result(rank, mesh, group) + dist.barrier() + mismatch = _mismatch_result(rank, mesh, group, last_document) + mixed_factored_contract = _mixed_factored_contract_result( + rank, mesh, group + ) + result_queue.put( + { + "rank": rank, + "cases": cases, + "dcp": dcp, + "malformed_transport": malformed_transport, + "mismatch": mismatch, + "mixed_factored_contract": mixed_factored_contract, + } + ) + except BaseException: + result_queue.put({"rank": rank, "fatal": traceback.format_exc()}) + finally: + if dist.is_available() and dist.is_initialized(): + dist.destroy_process_group() + + +def _run_workers(): + context = mp.get_context("spawn") + result_queue = context.Queue() + directory = tempfile.mkdtemp(prefix="gefen-portable-dtensor-") + init_file = os.path.join(directory, "store") + processes = [context.Process(target=_worker, args=(rank, init_file, result_queue)) for rank in range(_WORLD)] + results = [] + deadline = time.monotonic() + 120 + try: + for process in processes: + process.start() + while len(results) < _WORLD: + remaining = deadline - time.monotonic() + if remaining <= 0: + break + try: + results.append(result_queue.get(timeout=remaining)) + except queue_module.Empty: + break + for process in processes: + process.join(timeout=max(0.0, min(5.0, deadline - time.monotonic()))) + 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() + shutil.rmtree(directory, ignore_errors=True) + assert len(results) == _WORLD, (results, [process.exitcode for process in processes]) + assert all(process.exitcode == 0 for process in processes), (results, [process.exitcode for process in processes]) + results = sorted(results, key=lambda item: item["rank"]) + assert all("fatal" not in result for result in results), results + return results + + +@pytest.mark.skipif(not dist.is_available() or not dist.is_gloo_available(), reason="portable DTensor coverage requires Gloo") +def test_portable_dtensor_reshards_all_public_1d_placements_and_continues_exactly(): + results = _run_workers() + assert all(len(result["cases"]) == 7 for result in results), results + assert any(case["empty"] for result in results for case in result["cases"]), results + assert all( + case["document_exact"] + and case["projected_exact"] + and case["checkpoint_schema_exact"] + and case["continuation_exact"] + for result in results + for case in result["cases"] + ), results + assert any(result["dcp"]["empty"] for result in results), results + assert all( + result["dcp"]["sharded_envelope"] + and result["dcp"]["imported_exact"] + and result["dcp"]["continuation_exact"] + and result["dcp"]["mismatch_error"] is not None + and result["dcp"]["mismatch_unchanged"] + and result["dcp"]["corruption_error"] is not None + and result["dcp"]["corruption_unchanged"] + for result in results + ), results + mismatch_messages = [result["dcp"]["mismatch_error"] for result in results] + assert mismatch_messages[0] == mismatch_messages[1] + corruption_messages = [result["dcp"]["corruption_error"] for result in results] + assert corruption_messages[0] == corruption_messages[1] + assert all( + result["malformed_transport"]["rejected"] + and result["malformed_transport"]["unchanged"] + for result in results + ), results + assert all(result["mismatch"]["rejected"] and result["mismatch"]["unchanged"] for result in results), results + assert all( + all(result["mixed_factored_contract"].values()) for result in results + ), results + + +def test_public_rebinding_and_portable_runtime_reject_dtensor_like_tensor_subclasses(): + from gefen import Gefen, LogicalRegion, ParameterIdentity, ParameterLayout, ParameterRebinding, PlacementKind, ProcessGroupIdentity, ShardIdentity, ShardPlacement, ShardingManifest + from gefen.portable_runtime import _parameter_storage_token, _parameter_supported + from torch import nn + + class DTensorLike(torch.Tensor): + @staticmethod + def __new__(cls): + return torch.Tensor._make_subclass(cls, torch.zeros(2, 2), False) + + def to_local(self): + return self.as_subclass(torch.Tensor) + + @property + def placements(self): + return () + + @property + def device_mesh(self): + return None + + identity = ParameterIdentity(_FQN, (2, 2)) + group = ProcessGroupIdentity("portable_dtensor_subclass", ("rank:0",)) + shard = ShardIdentity( + identity, + ParameterLayout.DTENSOR_1D_DEFAULT_WORLD, + LogicalRegion.full(identity), + placements=(ShardPlacement("dp", PlacementKind.REPLICATE, 0, 1),), + process_group=group, + local_member="rank:0", + ) + value = DTensorLike() + original = nn.Parameter(torch.zeros(2, 2)) + optimizer = Gefen( + [("weight", original)], + fused=False, + force_2d_period_one=True, + ) + before = ( + id(optimizer.state), + id(optimizer.param_groups), + id(optimizer.param_groups[0]), + optimizer.param_groups[0]["params"][0], + tuple(optimizer.state), + ) + with pytest.raises(TypeError, match="DTensor-like tensor subclasses"): + optimizer.post_sharding( + (ParameterRebinding(original, value, shard),), + manifest=ShardingManifest((shard,)), + ) + after = ( + id(optimizer.state), + id(optimizer.param_groups), + id(optimizer.param_groups[0]), + optimizer.param_groups[0]["params"][0], + tuple(optimizer.state), + ) + assert after == before + assert not _parameter_supported(value, shard) + with pytest.raises(TypeError, match="DTensor-like tensor subclasses"): + _parameter_storage_token(value) diff --git a/tests/test_portable_hybrid_distributed.py b/tests/test_portable_hybrid_distributed.py index a239593..ad5185b 100644 --- a/tests/test_portable_hybrid_distributed.py +++ b/tests/test_portable_hybrid_distributed.py @@ -17,6 +17,7 @@ from gefen.checkpoint import CheckpointProcessGroupBinding from gefen.codebook import CodebookProcessGroupBinding from gefen.contracts import ( + CheckpointTransport, LogicalSlice, ParameterIdentity, ParameterLayout, @@ -25,6 +26,7 @@ ShardIdentity, ShardPlacement, ShardingManifest, + TopologyChange, ) from gefen.portable_hybrid import _hybrid_portable_live_token from gefen.portable_state import PortableStateLimits @@ -103,6 +105,27 @@ def _owner_shards(identity, group, owner): ) +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 _flat_shards(identity, group, lengths): if len(lengths) != len(group.ordered_members) or sum(lengths) != identity.numel: raise AssertionError("invalid flattened test partition") @@ -641,6 +664,348 @@ def _run_distributed_workers(checkpoint_dir): return sorted(results, key=lambda item: item["rank"]) +def _make_singleton_hybrid(*, deterministic=True, replicated=False): + member = "load:0" if replicated else "save:0" + group = ProcessGroupIdentity( + ( + "hybrid_portable_reverse_singleton" + if replicated + else "hybrid_portable_singleton" + ), + (member,), + ) + muon_identity = ParameterIdentity(_MUON_FQN, (3, 2)) + backup_identity = ParameterIdentity(_BACKUP_FQN, (8,)) + muon_parameter = torch.nn.Parameter(_muon_initial().clone()) + backup_parameter = torch.nn.Parameter(_backup_initial().clone()) + optimizer = GefenMuonHybrid( + [("matrix", muon_parameter)], + [("vector", backup_parameter)], + 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_shard = ( + _replicated_shards(muon_identity, group)[0] + if replicated + else _owner_shards(muon_identity, group, member)[0] + ) + backup_shard = ( + _replicated_shards(backup_identity, group)[0] + if replicated + else _flat_shards(backup_identity, group, (8,))[0] + ) + optimizer.post_sharding( + ( + ParameterRebinding(muon_parameter, muon_parameter, muon_shard), + ParameterRebinding(backup_parameter, backup_parameter, backup_shard), + ), + manifest=ShardingManifest((muon_shard, backup_shard)), + codebook_process_group=CodebookProcessGroupBinding( + group, + member, + None, + torch.device("cpu"), + ), + ) + return ( + optimizer, + muon_parameter, + backup_parameter, + backup_shard, + CheckpointProcessGroupBinding( + group, + member, + None, + torch.device("cpu"), + ), + ) + + +def _has_replicated_portable_intersection(optimizer, *, same_topology): + supports = tuple( + support + for support in optimizer.optimizer_contract().capabilities.checkpoints + if support.transport is CheckpointTransport.CANONICAL_GLOBAL + ) + return ( + len(supports) == 1 + and supports[0].same_topology + == ( + frozenset({ParameterLayout.REPLICATED}) + if same_topology + else frozenset() + ) + and supports[0].topology_changing + == frozenset({ParameterLayout.REPLICATED}) + and supports[0].topology_change_kinds + == frozenset({TopologyChange.PLACEMENT_RESHARD}) + and supports[0].requires_collective + and supports[0].atomic_load + ) + + +def _singleton_save_worker(_rank, init_file, checkpoint_dir, result_queue): + _strict_worker_warnings() + try: + dist.init_process_group( + "gloo", + init_method="file://{}".format(init_file), + rank=0, + world_size=1, + timeout=timedelta(seconds=120), + ) + source = _make_singleton_hybrid() + _seed_hybrid_state(source[0], source[1], source[2], source[3]) + save_portable_dcp( + source[0], + checkpoint_process_group=source[4], + storage_writer=FileSystemWriter(checkpoint_dir), + transaction_id="hybrid-world-change-save", + limits=_limits(), + ) + keys = FileSystemReader(checkpoint_dir).read_metadata().state_dict_metadata + result_queue.put( + { + "rank": 0, + "v2": ( + "optimizer.__gefen_portable_composite_metadata_v2__" in keys + and "optimizer.__gefen_portable_metadata_v1__" not in keys + ), + } + ) + except BaseException: + result_queue.put({"rank": 0, "fatal_error": traceback.format_exc()}) + finally: + if dist.is_available() and dist.is_initialized(): + dist.destroy_process_group() + + +def _world_change_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=_WORLD, + timeout=timedelta(seconds=120), + ) + group = ProcessGroupIdentity( + "hybrid_portable_world_change_target", + _members(), + ) + 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-world-change-load", + 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], + ) + 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 _reverse_world_change_save_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_reverse_source", + _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]) + intersection_exact = _has_replicated_portable_intersection( + source[0], + same_topology=False, + ) + save_portable_dcp( + source[0], + checkpoint_process_group=source[4], + storage_writer=FileSystemWriter(checkpoint_dir), + transaction_id="hybrid-reverse-world-change-save", + limits=_limits(), + ) + keys = FileSystemReader(checkpoint_dir).read_metadata().state_dict_metadata + result_queue.put( + { + "rank": rank, + "intersection_exact": intersection_exact, + "v2": ( + "optimizer.__gefen_portable_composite_metadata_v2__" in keys + and "optimizer.__gefen_portable_metadata_v1__" not in keys + ), + } + ) + 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 _reverse_world_change_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=0, + world_size=1, + timeout=timedelta(seconds=120), + ) + target = _make_singleton_hybrid(deterministic=False, replicated=True) + reference = _make_singleton_hybrid(deterministic=True, replicated=True) + _seed_hybrid_state( + reference[0], + reference[1], + reference[2], + reference[3], + ) + intersection_exact = _has_replicated_portable_intersection( + target[0], + same_topology=True, + ) + load_portable_dcp( + target[0], + checkpoint_process_group=target[4], + storage_reader=FileSystemReader(checkpoint_dir), + transaction_id="hybrid-reverse-world-change-load", + 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], + ) + result_queue.put( + { + "rank": 0, + "intersection_exact": intersection_exact, + "restored_exact": restored_exact, + "next_step_exact": next_step_exact, + } + ) + except BaseException: + result_queue.put({"rank": 0, "fatal_error": traceback.format_exc()}) + finally: + if dist.is_available() and dist.is_initialized(): + dist.destroy_process_group() + + +def _run_phase(worker, world_size, checkpoint_dir): + context = mp.get_context("spawn") + result_queue = context.Queue() + descriptor, init_file = tempfile.mkstemp(prefix="gefen-portable-hybrid-phase-") + 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_size) + ] + 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_size, (results, [process.exitcode for process in processes]) + assert all(process.exitcode == 0 for process in processes) + results = sorted(results, key=lambda item: item["rank"]) + assert all("fatal_error" not in result for result in results), results + return results + + @pytest.mark.skipif( not dist.is_available() or not dist.is_gloo_available(), reason="portable Hybrid topology coverage requires Gloo", @@ -666,7 +1031,6 @@ def test_two_process_portable_hybrid_topology_change_and_dcp(tmp_path): 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] @@ -677,3 +1041,48 @@ def test_two_process_portable_hybrid_topology_change_and_dcp(tmp_path): and result["asymmetric"]["common_unchanged"] for result in results ), results + + +@pytest.mark.skipif( + not dist.is_available() or not dist.is_gloo_available(), + reason="portable Hybrid world-size coverage requires Gloo", +) +def test_sharded_hybrid_dcp_changes_world_size_and_continues_exactly(tmp_path): + checkpoint_dir = str(tmp_path / "hybrid-world-change") + saved = _run_phase(_singleton_save_worker, 1, checkpoint_dir) + assert saved == [{"rank": 0, "v2": True}] + loaded = _run_phase(_world_change_load_worker, _WORLD, checkpoint_dir) + assert all( + result["restored_exact"] and result["next_step_exact"] + for result in loaded + ), loaded + + +@pytest.mark.skipif( + not dist.is_available() or not dist.is_gloo_available(), + reason="portable Hybrid reverse world-size coverage requires Gloo", +) +def test_sharded_hybrid_dcp_collapses_to_singleton_and_continues_exactly(tmp_path): + checkpoint_dir = str(tmp_path / "hybrid-reverse-world-change") + saved = _run_phase( + _reverse_world_change_save_worker, + _WORLD, + checkpoint_dir, + ) + assert all( + result["intersection_exact"] and result["v2"] + for result in saved + ), saved + loaded = _run_phase( + _reverse_world_change_load_worker, + 1, + checkpoint_dir, + ) + assert loaded == [ + { + "rank": 0, + "intersection_exact": True, + "restored_exact": True, + "next_step_exact": True, + } + ] diff --git a/tests/test_portable_hybrid_runtime.py b/tests/test_portable_hybrid_runtime.py index 0d43303..5ccc9a2 100644 --- a/tests/test_portable_hybrid_runtime.py +++ b/tests/test_portable_hybrid_runtime.py @@ -327,9 +327,18 @@ def test_hybrid_portable_composite_uses_collective_tensor_limit(): _assert_same_optimizer_state(source, target) -def test_hybrid_portable_dcp_round_trip(tmp_path): +def test_hybrid_portable_dcp_round_trip(tmp_path, monkeypatch): source, source_matrix, source_bias, source_binding = _initialized_source() checkpoint = tmp_path / "hybrid-portable" + + def reject_dense_export(*_args, **_kwargs): + raise AssertionError("sharded Hybrid DCP must not call dense export") + + monkeypatch.setattr( + type(source), + "export_portable_state", + reject_dense_export, + ) save_portable_dcp( source, checkpoint_process_group=source_binding, @@ -337,6 +346,11 @@ def test_hybrid_portable_dcp_round_trip(tmp_path): transaction_id="hybrid-dcp-save-v1", limits=_limits(), ) + keys = dcp.FileSystemReader(checkpoint).read_metadata().state_dict_metadata + assert "optimizer.__gefen_portable_composite_metadata_v2__" in keys + assert "optimizer.__gefen_portable_metadata_v1__" not in keys + assert any(key.startswith("optimizer-gefen-muon.__gefen_portable_field_") for key in keys) + assert any(key.startswith("optimizer-gefen-backup.__gefen_portable_field_") for key in keys) target, target_matrix, target_bias, target_binding = _target_from_source( source_matrix, source_bias, @@ -360,6 +374,38 @@ def test_hybrid_portable_dcp_round_trip(tmp_path): _assert_same_optimizer_state(source, target) +def test_public_loader_accepts_legacy_v1_hybrid_checkpoint(tmp_path): + from gefen import portable_dcp + + source, source_matrix, source_bias, source_binding = _initialized_source() + limits = _limits() + document = source.export_portable_state( + checkpoint_process_group=source_binding, + transaction_id="legacy-v1-hybrid-export", + limits=limits, + ) + wire_limits = limits._wire_limits(collective=True) + plan = portable_dcp._prepare_canonical_wire_value(document, wire_limits) + state = portable_dcp._state_from_plan("optimizer", plan, wire_limits) + checkpoint = tmp_path / "legacy-v1-hybrid" + dcp.save(state, storage_writer=dcp.FileSystemWriter(checkpoint)) + keys = dcp.FileSystemReader(checkpoint).read_metadata().state_dict_metadata + assert "optimizer.__gefen_portable_metadata_v1__" in keys + assert "optimizer.__gefen_portable_composite_metadata_v2__" not in keys + 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="legacy-v1-hybrid-load", + limits=limits, + ) + _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)) diff --git a/tests/test_portable_runtime.py b/tests/test_portable_runtime.py index df657bf..6f0f9a7 100644 --- a/tests/test_portable_runtime.py +++ b/tests/test_portable_runtime.py @@ -23,7 +23,7 @@ ) from gefen.gefen import Gefen from gefen.gefen_muon import GefenMuon -from gefen.portable import _decode_quantized_momentum +from gefen.portable import _decode_quantized_momentum, _recompress_dense_momentum from gefen.portable_runtime import ( _export_portable_state, _import_portable_state, @@ -52,15 +52,23 @@ def _bindings(group): ) -def _finalize(optimizer, parameter, *, layout): +def _finalize(optimizer, parameter, *, layout, global_shape=None): group = ProcessGroupIdentity("checkpoint", (_MEMBER,)) - identity = ParameterIdentity("layer.weight", tuple(parameter.shape)) + identity = ParameterIdentity( + "layer.weight", + tuple(parameter.shape) if global_shape is None else global_shape, + ) if layout is ParameterLayout.REPLICATED: kind = PlacementKind.REPLICATE owner = None - else: + elif layout is ParameterLayout.FLATTENED_ELEMENT_SHARD: + kind = PlacementKind.FLAT_SHARD + owner = None + elif layout is ParameterLayout.WHOLE_PARAMETER_OWNER: kind = PlacementKind.WHOLE_PARAMETER_OWNER owner = _MEMBER + else: + raise AssertionError("unsupported test layout") shard = ShardIdentity( identity, layout, @@ -79,8 +87,11 @@ def _finalize(optimizer, parameter, *, layout): return checkpoint_binding -def _plain(*, factored, deterministic): - parameter = torch.nn.Parameter(torch.arange(1, 7, dtype=torch.float32).reshape(2, 3)) +def _plain(*, factored, deterministic, layout=ParameterLayout.REPLICATED): + values = torch.arange(1, 7, dtype=torch.float32).reshape(2, 3) + parameter = torch.nn.Parameter( + values.reshape(-1) if layout is ParameterLayout.FLATTENED_ELEMENT_SHARD else values + ) optimizer = Gefen( [("weight", parameter)], fused=False, @@ -88,7 +99,12 @@ def _plain(*, factored, deterministic): force_2d_period_one=True, deterministic=deterministic, ) - binding = _finalize(optimizer, parameter, layout=ParameterLayout.REPLICATED) + binding = _finalize( + optimizer, + parameter, + layout=layout, + global_shape=(2, 3), + ) return optimizer, parameter, binding @@ -178,10 +194,11 @@ def test_initialized_singleton_export_import_round_trip(variant): elif variant == "block": assert support.topology_changing == frozenset( { - ParameterLayout.REPLICATED, - ParameterLayout.FLATTENED_ELEMENT_SHARD, - } - ) + ParameterLayout.REPLICATED, + ParameterLayout.FLATTENED_ELEMENT_SHARD, + ParameterLayout.DTENSOR_1D_DEFAULT_WORLD, + } + ) assert support.topology_change_kinds == frozenset({TopologyChange.PLACEMENT_RESHARD}) else: assert support.topology_changing == frozenset( @@ -237,6 +254,132 @@ def test_initialized_singleton_export_import_round_trip(variant): assert target_state["normuon_step"] == 2 +def test_factored_to_flat_block_import_matches_target_topology_reference_continuation(): + source, source_parameter, source_binding = _plain(factored=True, deterministic=True) + target, target_parameter, target_binding = _plain( + factored=False, + deterministic=False, + layout=ParameterLayout.FLATTENED_ELEMENT_SHARD, + ) + reference, reference_parameter, _ = _plain( + factored=False, + deterministic=True, + layout=ParameterLayout.FLATTENED_ELEMENT_SHARD, + ) + _initialize(source, source_parameter, variant="factored") + document = _export_portable_state( + source, + checkpoint_process_group=source_binding, + transaction_id="export-factored-to-block", + limits=_limits(), + ) + assert "factored_to_block_live_fp32" in document["policy"]["second_moment_projection"] + + _import_portable_state( + target, + document, + checkpoint_process_group=target_binding, + transaction_id="import-factored-to-block", + limits=_limits(), + ) + + record = document["parameters"]["layer.weight"] + expected_second = torch.outer(record["state"]["v_row"], record["state"]["v_col"]).div_( + record["state"]["v_row"].mean().clamp_(min=torch.finfo(torch.float32).tiny) + ) + target_state = target.state[target_parameter] + assert torch.equal(target_state["vmean"].reshape(2, 3).view(torch.int32), expected_second.view(torch.int32)) + assert target_state["vmean_step"] == record["state"]["factored_step"] + + reference._gefen_global_step = document["common"]["gefen_global_step"] + reference._gefen_codebook = document["common"]["gefen_codebook"].clone() + reference._deterministic = document["common"]["gefen_deterministic"] + reference_indices, reference_magnitudes = _recompress_dense_momentum( + record["state"]["momentum"], + reference._gefen_codebook, + period=1, + step=record["state"]["step"], + ) + reference.state[reference_parameter].update( + { + "automatic_period": 1, + "step": record["state"]["step"], + "m_codebook": reference_indices, + "m_magnitude": reference_magnitudes, + "vmean": expected_second.reshape(-1, 1).clone(), + "vmean_step": record["state"]["factored_step"], + } + ) + + gradient = torch.tensor([[-0.75, 0.5, 2.0], [3.25, -1.5, 0.125]], dtype=torch.float32) + target_parameter.grad = gradient.reshape(-1).clone() + reference_parameter.grad = gradient.reshape(-1).clone() + target.step() + reference.step() + + assert torch.equal(target_parameter.view(torch.int32), reference_parameter.view(torch.int32)) + assert target._gefen_global_step == reference._gefen_global_step == 5 + reference_state = reference.state[reference_parameter] + for key in ("automatic_period", "step", "vmean_step"): + assert target_state[key] == reference_state[key] + assert torch.equal(target_state["m_codebook"], reference_state["m_codebook"]) + for key in ("m_magnitude", "vmean"): + assert torch.equal(target_state[key].view(torch.int32), reference_state[key].view(torch.int32)) + + +def test_factored_to_block_projection_overflow_fails_before_live_mutation(): + source, source_parameter, source_binding = _plain(factored=True, deterministic=True) + target, target_parameter, target_binding = _plain(factored=False, deterministic=False) + _initialize(source, source_parameter, variant="factored") + _initialize(target, target_parameter, variant="block") + maximum = torch.finfo(torch.float32).max + source.state[source_parameter]["v_row"] = torch.full((2,), maximum) + source.state[source_parameter]["v_col"] = torch.full((3,), maximum) + document = _export_portable_state( + source, + checkpoint_process_group=source_binding, + transaction_id="export-factored-overflow", + limits=_limits(), + ) + before = target._canonical_import_live_token() + + with pytest.raises(RuntimeError, match="cannot be represented"): + _import_portable_state( + target, + document, + checkpoint_process_group=target_binding, + transaction_id="import-factored-overflow", + limits=_limits(), + ) + + assert target._canonical_import_live_token() == before + + +def test_block_to_factored_import_is_rejected_before_live_mutation(): + source, source_parameter, source_binding = _plain(factored=False, deterministic=True) + target, target_parameter, target_binding = _plain(factored=True, deterministic=False) + _initialize(source, source_parameter, variant="block") + _initialize(target, target_parameter, variant="factored") + document = _export_portable_state( + source, + checkpoint_process_group=source_binding, + transaction_id="export-block-for-reverse-rejection", + limits=_limits(), + ) + before = target._canonical_import_live_token() + + with pytest.raises(RuntimeError, match="block-to-factored"): + _import_portable_state( + target, + document, + checkpoint_process_group=target_binding, + transaction_id="import-block-to-factored", + limits=_limits(), + ) + + assert target._canonical_import_live_token() == before + + 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) diff --git a/tests/test_portable_state.py b/tests/test_portable_state.py index 1894cd5..91c4e0d 100644 --- a/tests/test_portable_state.py +++ b/tests/test_portable_state.py @@ -26,6 +26,8 @@ from gefen.portable_schema import build_portable_state_document from gefen.portable_state import ( PortableStateLimits, + _SECOND_MOMENT_PROJECTION_EXACT, + _SECOND_MOMENT_PROJECTION_FACTORED_TO_BLOCK, _assemble_portable_state_fragments, _build_portable_state_fragment, _normalize_gefen_portable_state_document, @@ -57,7 +59,9 @@ def _policy(*, factored=False): "stochastic_round": False, "codebook_refresh_every": 0, "momentum_projection": "dense_fp32_target_period_one_v1", - "second_moment_projection": "exact_representation_target_period_one_v1", + "second_moment_projection": ( + _SECOND_MOMENT_PROJECTION_FACTORED_TO_BLOCK if factored else _SECOND_MOMENT_PROJECTION_EXACT + ), } @@ -294,6 +298,7 @@ def test_plain_initialized_block_assembles_and_projects_exact_period_one(layout) implementation="gefen.Gefen", global_step=4, codebook=document["common"]["gefen_codebook"], + source_second_moment_projection=document["policy"]["second_moment_projection"], target_algorithm_options=_plain_options(), target_second_moment="block", ) @@ -333,7 +338,7 @@ def test_empty_flat_fragment_is_pristine_and_ownership_is_enforced(): _normalize_portable_state_fragment(corrupted, limits=_limits()) -def test_factored_replicas_preserve_authoritative_factor_bytes_and_reject_migration(): +def test_factored_replicas_preserve_factor_bytes_and_define_live_fp32_block_projection(): parameter = ParameterIdentity("matrix.weight", (2, 3)) group = _group() shards = _replicas(parameter, group) @@ -360,23 +365,91 @@ def test_factored_replicas_preserve_authoritative_factor_bytes_and_reject_migrat implementation="gefen.Gefen", global_step=4, codebook=_codebook(), + source_second_moment_projection=document["policy"]["second_moment_projection"], 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"): + assert torch.equal(projected["v_col"].view(torch.int32), v_col.view(torch.int32)) + block = _project_portable_parameter_state( + record, + shards[0], + implementation="gefen.Gefen", + global_step=4, + codebook=_codebook(), + source_second_moment_projection=document["policy"]["second_moment_projection"], + target_algorithm_options=_plain_options(second="block"), + target_second_moment="block", + ) + denominator = v_row.mean().clamp_(min=torch.finfo(torch.float32).tiny) + expected = torch.outer(v_row, v_col).div_(denominator) + assert torch.equal(block["vmean"].reshape(2, 3).view(torch.int32), expected.view(torch.int32)) + assert block["vmean_step"] == 3 + assert block["step"] == 4 + assert block["automatic_period"] == 1 + 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_second_moment_projection_rejects_reverse_and_exact_only_factored_migration(): + parameter = ParameterIdentity("matrix.weight", (2, 3)) + group = _group() + shards = _replicas(parameter, group) + block_document = _assemble_portable_state_fragments( + _fragments( + "gefen.Gefen", + parameter, + shards, + options=_plain_options(second="block"), + policy=_policy(), + variant="initialized_dense", + momentum=torch.arange(6, dtype=torch.float32).reshape(2, 3), + second=torch.arange(1, 7, dtype=torch.float32).reshape(2, 3), + ), + process_group_identity=group, + limits=_limits(), + ) + with pytest.raises(ValueError, match="block-to-factored"): _project_portable_parameter_state( - record, + block_document["parameters"][parameter.fqn], shards[0], implementation="gefen.Gefen", global_step=4, - codebook=_codebook(), + codebook=block_document["common"]["gefen_codebook"], + source_second_moment_projection=block_document["policy"]["second_moment_projection"], target_algorithm_options=_plain_options(second="factored"), + target_second_moment="factored", + ) + + exact_policy = _policy(factored=True) + exact_policy["second_moment_projection"] = _SECOND_MOMENT_PROJECTION_EXACT + exact_factored_document = _assemble_portable_state_fragments( + _fragments( + "gefen.Gefen", + parameter, + shards, + options=_plain_options(second="factored"), + policy=exact_policy, + variant="initialized_factored", + momentum=torch.arange(6, dtype=torch.float32).reshape(2, 3), + v_row=torch.tensor([2.0, 3.0]), + v_col=torch.tensor([5.0, 7.0, 11.0]), + ), + process_group_identity=group, + limits=_limits(), + ) + with pytest.raises(ValueError, match="does not authorize"): + _project_portable_parameter_state( + exact_factored_document["parameters"][parameter.fqn], + shards[0], + implementation="gefen.Gefen", + global_step=4, + codebook=exact_factored_document["common"]["gefen_codebook"], + source_second_moment_projection=exact_factored_document["policy"]["second_moment_projection"], + target_algorithm_options=_plain_options(second="block"), 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(): @@ -395,6 +468,17 @@ def test_period_selected_projection_binds_target_second_moment_policy(): process_group_identity=group, limits=_limits(), ) + projected = _project_portable_parameter_state( + document["parameters"][parameter.fqn], + shards[0], + implementation="gefen.Gefen", + global_step=0, + codebook=document["common"]["gefen_codebook"], + source_second_moment_projection=document["policy"]["second_moment_projection"], + target_algorithm_options=_plain_options(second="block"), + target_second_moment="block", + ) + assert projected == {"automatic_period": 1} with pytest.raises(ValueError, match="conflicts with target algorithm options"): _project_portable_parameter_state( document["parameters"][parameter.fqn], @@ -402,6 +486,7 @@ def test_period_selected_projection_binds_target_second_moment_policy(): implementation="gefen.Gefen", global_step=0, codebook=document["common"]["gefen_codebook"], + source_second_moment_projection=document["policy"]["second_moment_projection"], target_algorithm_options=_plain_options(second="factored"), target_second_moment="block", ) @@ -438,6 +523,7 @@ def test_factored_matrix_rejects_flat_fragments_before_initialization(): implementation="gefen.Gefen", global_step=document["common"]["gefen_global_step"], codebook=document["common"]["gefen_codebook"], + source_second_moment_projection=document["policy"]["second_moment_projection"], target_algorithm_options=_plain_options(second="factored"), target_second_moment="factored", ) @@ -469,6 +555,7 @@ def test_muon_whole_owner_assembles_and_projects_after_owner_move(normuon): implementation="gefen.GefenMuon", global_step=4, codebook=_codebook(), + source_second_moment_projection=document["policy"]["second_moment_projection"], target_algorithm_options=_muon_options(normuon=normuon), ) owner = _project_portable_parameter_state( @@ -477,6 +564,7 @@ def test_muon_whole_owner_assembles_and_projects_after_owner_move(normuon): implementation="gefen.GefenMuon", global_step=4, codebook=_codebook(), + source_second_moment_projection=document["policy"]["second_moment_projection"], target_algorithm_options=_muon_options(normuon=normuon), ) assert nonowner == {} @@ -508,6 +596,7 @@ def test_muon_whole_owner_projection_requires_distributed_mode(): implementation="gefen.GefenMuon", global_step=document["common"]["gefen_global_step"], codebook=document["common"]["gefen_codebook"], + source_second_moment_projection=document["policy"]["second_moment_projection"], target_algorithm_options=_muon_options(mode="exact"), ) @@ -667,6 +756,7 @@ def test_fragment_rejects_ungrouped_shards_and_empty_period_projection_stays_emp implementation="gefen.Gefen", global_step=4, codebook=document["common"]["gefen_codebook"], + source_second_moment_projection=document["policy"]["second_moment_projection"], target_algorithm_options=_plain_options(), target_second_moment="block", ) diff --git a/tests/test_portable_state_math.py b/tests/test_portable_state_math.py index b5bfdea..5b60b4f 100644 --- a/tests/test_portable_state_math.py +++ b/tests/test_portable_state_math.py @@ -8,6 +8,7 @@ _decode_quantized_momentum, _expand_block_second_moment, _expand_factored_second_moment, + _expand_factored_second_moment_live_fp32_v1, _project_factored_second_moment, _recompress_dense_momentum, _reduce_block_second_moment, @@ -482,6 +483,54 @@ def test_factored_expansion_and_projection_follow_adafactor_geometry(): _assert_tight(projected_column, shape=(2,)) +def test_live_fp32_factored_expansion_matches_optimizer_operator_bits(): + row = torch.tensor([0.10000000149011612, 3.75, 19.125], dtype=torch.float32) + column = torch.tensor([0.30000001192092896, 7.25, 23.5, 101.0], dtype=torch.float32) + expected = torch.outer(row, column).div_( + row.mean().clamp_(min=torch.finfo(torch.float32).tiny) + ) + + actual = _expand_factored_second_moment_live_fp32_v1( + row, + column, + logical_shape=(3, 4), + step=7, + ) + + _assert_tight(actual, shape=(3, 4)) + assert torch.equal(actual.view(torch.int32), expected.view(torch.int32)) + + +def test_live_fp32_factored_expansion_defines_zero_and_subnormal_clamp_edges(): + subnormal = torch.nextafter(torch.tensor(0.0), torch.tensor(1.0)) + subnormal_result = _expand_factored_second_moment_live_fp32_v1( + torch.stack((subnormal, subnormal)), + torch.stack((subnormal, subnormal)), + logical_shape=(2, 2), + step=1, + ) + assert torch.equal(subnormal_result, torch.zeros((2, 2))) + + zero_row_result = _expand_factored_second_moment_live_fp32_v1( + torch.zeros(2), + torch.tensor([1.0, torch.finfo(torch.float32).max]), + logical_shape=(2, 2), + step=1, + ) + assert torch.equal(zero_row_result, torch.zeros((2, 2))) + + +def test_live_fp32_factored_expansion_rejects_operator_overflow(): + maximum = torch.finfo(torch.float32).max + with pytest.raises(ValueError, match="cannot be represented"): + _expand_factored_second_moment_live_fp32_v1( + torch.full((2,), maximum), + torch.full((2,), maximum), + logical_shape=(2, 2), + step=1, + ) + + 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(