From 9a5add95081665439f98fc6f1bf578ff6f01780e Mon Sep 17 00:00:00 2001 From: thad0ctor Date: Sun, 12 Jul 2026 17:21:23 -0700 Subject: [PATCH 01/52] Add platform-agnostic optimizer contracts --- README.md | 2 + docs/optimizer_contracts.md | 36 + src/gefen/__init__.py | 45 + src/gefen/contracts.py | 1040 +++++++++++++++++ src/gefen/gefen.py | 6 + src/gefen/gefen_muon.py | 23 + src/gefen/hybrid.py | 25 + ...test_muon_distributed_checkpoint_safety.py | 34 + tests/test_muon_fsdp2_approx.py | 19 + tests/test_muon_fsdp2_parity.py | 19 + tests/test_optimizer_contracts.py | 630 ++++++++++ 11 files changed, 1879 insertions(+) create mode 100644 docs/optimizer_contracts.md create mode 100644 src/gefen/contracts.py create mode 100644 tests/test_optimizer_contracts.py diff --git a/README.md b/README.md index 92ca5f8..21f676b 100644 --- a/README.md +++ b/README.md @@ -131,6 +131,8 @@ Mixed precision works out of the box: BF16 and standard AMP behave exactly as wi > **DeepSpeed ZeRO config.** Set `"zero_allow_untested_optimizer": true` and leave the config's `optimizer` section unset. With optimizer CPU-offload, also set `"zero_force_ds_cpu_optimizer": false` — otherwise raw DeepSpeed refuses to initialize, and accelerate-based launchers (axolotl) silently swap in DeepSpeed's own CPU Adam. +Platform adapters can query the immutable [`optimizer_contract()` capability and state-layout descriptors](https://github.com/thad0ctor/Gefen-X/blob/main/docs/optimizer_contracts.md) instead of depending on Gefen's private optimizer attributes. + ## Determinism (`deterministic`) Set `deterministic=True` when every GPU replica must produce bit-identical results on matching GPUs. It is off by default (the fastest routing), checkpoints remember the setting, and older checkpoints without it still load. One combination is rejected: plain Gefen with `deterministic=True`, `factored_v_2d=True`, and `stochastic_round=True` together. diff --git a/docs/optimizer_contracts.md b/docs/optimizer_contracts.md new file mode 100644 index 0000000..0e3de25 --- /dev/null +++ b/docs/optimizer_contracts.md @@ -0,0 +1,36 @@ +# Optimizer integration contracts + +Gefen exposes immutable, versioned optimizer contracts for platform adapters that need to inspect state ownership and validated distributed capabilities without depending on private attributes. Calling `optimizer.optimizer_contract()` is read-only and does not change parameters, optimizer state, checkpoint schemas, or step behavior. + +```python +from gefen import CheckpointTransport, Gefen, ParameterLayout + +optimizer = Gefen(model.named_parameters(), lr=3e-5) +contract = optimizer.optimizer_contract() + +assert contract.schema_version == 1 +rank_local_dcp = next( + support + for support in contract.capabilities.checkpoints + if support.transport is CheckpointTransport.PYTORCH_RANK_LOCAL +) +assert ParameterLayout.DTENSOR_1D_DEFAULT_WORLD in rank_local_dcp.same_topology +``` + +## Contract boundaries + +- `OptimizerStateLayout` separates optimizer-common authoritative state, per-parameter authoritative state, derived caches, checkpoint transport fields, and composite child namespaces. +- `StateVariant` identifies valid lazy, initialized, local-shard, global-parameter, owner, non-owner, and migrated state combinations using structured layout, mode, rank, extent, ownership, and inactive-field declarations. +- `TrainingSupport` qualifies each validated parameter layout by process-group source, mesh dimensionality, sharded mode, and whether the update needs complete parameter storage or a transient complete logical matrix. +- `CheckpointSupport` reports same-topology, topology-changing, and fail-before-mutation load support separately for native, PyTorch rank-local, 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. + +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. + +## Adapter requirements + +An adapter should match the exact `TrainingSupport` or `CheckpointSupport` entry it intends to use, including transport, layout, process-group scope, mesh dimensions, and sharded mode. It should not treat successful training as checkpoint support, same-topology checkpointing as resharding support, or accepted caller names as canonical fully qualified parameter identity. + +The contract types do not import a distributed platform and do not perform collectives. Platform adapters remain responsible for topology discovery, deterministic scheduling, and lifecycle orchestration; future core mutation APIs can consume these declarations without changing their meaning. diff --git a/src/gefen/__init__.py b/src/gefen/__init__.py index d2eb309..c427b95 100644 --- a/src/gefen/__init__.py +++ b/src/gefen/__init__.py @@ -11,6 +11,26 @@ "Gefen", "GefenMuon", "GefenMuonHybrid", + "CONTRACT_SCHEMA_VERSION", + "CheckpointSupport", + "CheckpointTransport", + "OptimizerCapabilities", + "OptimizerChildContract", + "OptimizerContract", + "OptimizerContractProvider", + "OptimizerStateLayout", + "ParameterLayout", + "ParameterStateRole", + "Precision", + "ProcessGroupScope", + "StateExtent", + "StateField", + "StateGeometry", + "StateKeyMatch", + "StateScope", + "StateVariant", + "TopologyChange", + "TrainingSupport", "split_params_for_muon", "validate_split", "kernels", @@ -35,6 +55,31 @@ def __getattr__(name): from . import params return getattr(params, name) + if name in ( + "CONTRACT_SCHEMA_VERSION", + "CheckpointSupport", + "CheckpointTransport", + "OptimizerCapabilities", + "OptimizerChildContract", + "OptimizerContract", + "OptimizerContractProvider", + "OptimizerStateLayout", + "ParameterLayout", + "ParameterStateRole", + "Precision", + "ProcessGroupScope", + "StateExtent", + "StateField", + "StateGeometry", + "StateKeyMatch", + "StateScope", + "StateVariant", + "TopologyChange", + "TrainingSupport", + ): + from . import contracts + + return getattr(contracts, name) if name == "kernels": # NOT `from . import kernels`: its fromlist handling re-enters this # __getattr__ before the submodule import runs, recursing forever. diff --git a/src/gefen/contracts.py b/src/gefen/contracts.py new file mode 100644 index 0000000..3149835 --- /dev/null +++ b/src/gefen/contracts.py @@ -0,0 +1,1040 @@ +"""Platform-agnostic optimizer capability and state-layout contracts. + +The descriptors in this module are read-only declarations. They describe the +optimizer state Gefen owns and the integration surfaces implemented today +without importing DCP, DTensor, Megatron, DeepSpeed, or another adapter. +""" + +from dataclasses import dataclass +from enum import Enum +from typing import ( + AbstractSet, + FrozenSet, + Optional, + Protocol, + Sequence, + Tuple, + runtime_checkable, +) + + +CONTRACT_SCHEMA_VERSION = 1 + + +class StateScope(str, Enum): + """Ownership category for one optimizer-state field.""" + + OPTIMIZER_COMMON = "optimizer_common" + PARAMETER = "parameter" + DERIVED = "derived" + + +class StateGeometry(str, Enum): + """Logical geometry kind, independent of local/global extent.""" + + SCALAR = "scalar" + CODEBOOK = "codebook" + PARAMETER = "parameter" + BLOCK = "block" + ROW = "row" + COLUMN = "column" + OPAQUE = "opaque" + + +class StateKeyMatch(str, Enum): + """How a field declaration matches live or serialized dictionary keys.""" + + EXACT = "exact" + PREFIX = "prefix" + + +class StateExtent(str, Enum): + """Logical extent represented by a per-parameter state variant.""" + + METADATA_ONLY = "metadata_only" + LOCAL_STORAGE = "local_storage" + GLOBAL_PARAMETER = "global_parameter" + OWNER_PARAMETER = "owner_parameter" + + +class ParameterStateRole(str, Enum): + """Ownership role on which a per-parameter variant is present.""" + + ANY = "any" + OWNER = "owner" + NON_OWNER = "non_owner" + + +class ParameterLayout(str, Enum): + """Concrete parameter layouts exposed by the current implementation.""" + + REPLICATED = "replicated" + FLATTENED_ELEMENT_SHARD = "flattened_element_shard" + WHOLE_PARAMETER_OWNER = "whole_parameter_owner" + DTENSOR_1D_DEFAULT_WORLD = "dtensor_1d_default_world" + + +class ProcessGroupScope(str, Enum): + """How an implemented path obtains its collective process group.""" + + NONE = "none" + DEFAULT_WORLD = "default_world" + INFERRED_DEVICE_MESH = "inferred_device_mesh" + ADAPTER_DEFINED = "adapter_defined" + + +class CheckpointTransport(str, Enum): + """Checkpoint transport with independently scoped topology support.""" + + NATIVE_OPTIMIZER = "native_optimizer" + PYTORCH_RANK_LOCAL = "pytorch_rank_local" + COMPOSITE_NATIVE = "composite_native" + + +class TopologyChange(str, Enum): + """Specific topology mutation supported by a checkpoint transport.""" + + WORLD_SIZE_OWNER_REDISTRIBUTION = "world_size_owner_redistribution" + PLACEMENT_RESHARD = "placement_reshard" + + +class Precision(str, Enum): + """Parameter/gradient storage precision accepted by validated core paths.""" + + FLOAT32 = "float32" + BFLOAT16 = "bfloat16" + FLOAT16 = "float16" + FLOAT64 = "float64" + + +def _tuple(value): + return tuple(value) + + +def _frozenset(value): + return frozenset(value) + + +def _validate_dimensions(name, values, *, positive): + if len(set(values)) != len(values): + raise ValueError("{} must be unique".format(name)) + minimum = 1 if positive else 0 + if any(type(value) is not int or value < minimum for value in values): + relation = "positive" if positive else "nonnegative" + raise ValueError("{} must contain {} integers".format(name, relation)) + + +@dataclass(frozen=True) +class StateField: + """One named authoritative field, runtime cache, or transport field.""" + + name: str + scope: StateScope + geometry: StateGeometry + checkpointed: bool + key_match: StateKeyMatch = StateKeyMatch.EXACT + applicable_sharded_modes: AbstractSet[str] = frozenset() + description: str = "" + + def __post_init__(self) -> None: + object.__setattr__( + self, + "applicable_sharded_modes", + _frozenset(self.applicable_sharded_modes), + ) + if not self.name: + raise ValueError("StateField.name must be non-empty") + if not isinstance(self.scope, StateScope): + raise TypeError("StateField.scope must be a StateScope") + if not isinstance(self.geometry, StateGeometry): + raise TypeError("StateField.geometry must be a StateGeometry") + if not isinstance(self.key_match, StateKeyMatch): + raise TypeError("StateField.key_match must be a StateKeyMatch") + + @property + def authoritative(self) -> bool: + """Whether the field carries optimizer meaning rather than a cache.""" + + return self.scope is not StateScope.DERIVED + + def matches(self, key: str) -> bool: + """Return whether a live or serialized key matches this declaration.""" + + if self.key_match is StateKeyMatch.PREFIX: + return isinstance(key, str) and key.startswith(self.name) + return key == self.name + + +@dataclass(frozen=True) +class StateVariant: + """A machine-selectable valid per-parameter state combination.""" + + name: str + fields: Sequence[str] + layouts: AbstractSet[ParameterLayout] + extent: StateExtent + role: ParameterStateRole = ParameterStateRole.ANY + initialized: bool = True + parameter_ranks: Optional[Sequence[int]] = None + excluded_parameter_ranks: Sequence[int] = () + sharded_mode: Optional[str] = None + inactive_fields: Sequence[str] = () + migration_only: bool = False + description: str = "" + + def __post_init__(self) -> None: + object.__setattr__(self, "fields", _tuple(self.fields)) + object.__setattr__(self, "layouts", _frozenset(self.layouts)) + object.__setattr__(self, "inactive_fields", _tuple(self.inactive_fields)) + object.__setattr__( + self, "excluded_parameter_ranks", _tuple(self.excluded_parameter_ranks) + ) + if self.parameter_ranks is not None: + object.__setattr__(self, "parameter_ranks", _tuple(self.parameter_ranks)) + if not self.name: + raise ValueError("StateVariant.name must be non-empty") + if not self.fields: + raise ValueError("StateVariant.fields must be non-empty") + if len(set(self.fields)) != len(self.fields): + raise ValueError("StateVariant.fields must not contain duplicates") + if not self.layouts: + raise ValueError("StateVariant.layouts must be non-empty") + if not isinstance(self.extent, StateExtent): + raise TypeError("StateVariant.extent must be a StateExtent") + if not isinstance(self.role, ParameterStateRole): + raise TypeError("StateVariant.role must be a ParameterStateRole") + if not set(self.inactive_fields).issubset(self.fields): + raise ValueError("StateVariant.inactive_fields must be present in fields") + if self.parameter_ranks is not None and set( + self.parameter_ranks + ) & set(self.excluded_parameter_ranks): + raise ValueError("included and excluded parameter ranks must be disjoint") + if self.parameter_ranks is not None: + _validate_dimensions("parameter_ranks", self.parameter_ranks, positive=False) + _validate_dimensions( + "excluded_parameter_ranks", + self.excluded_parameter_ranks, + positive=False, + ) + + +@dataclass(frozen=True) +class OptimizerStateLayout: + """Declared fields, per-parameter variants, and composite namespaces.""" + + fields: Sequence[StateField] + parameter_variants: Sequence[StateVariant] + composite_namespaces: Sequence[str] = () + + def __post_init__(self) -> None: + object.__setattr__(self, "fields", _tuple(self.fields)) + object.__setattr__(self, "parameter_variants", _tuple(self.parameter_variants)) + object.__setattr__(self, "composite_namespaces", _tuple(self.composite_namespaces)) + names = tuple(field.name for field in self.fields) + if len(set(names)) != len(names): + raise ValueError("OptimizerStateLayout fields must have unique names") + if len(set(self.composite_namespaces)) != len(self.composite_namespaces): + raise ValueError("composite namespaces must be unique") + variant_names = tuple(variant.name for variant in self.parameter_variants) + if len(set(variant_names)) != len(variant_names): + raise ValueError("StateVariant names must be unique") + parameter_fields = { + field.name for field in self.fields if field.scope is StateScope.PARAMETER + } + for variant in self.parameter_variants: + unknown = set(variant.fields) - parameter_fields + if unknown: + raise ValueError( + "StateVariant {!r} references undeclared parameter fields: {}".format( + variant.name, sorted(unknown) + ) + ) + + def fields_for_scope(self, scope: StateScope) -> Tuple[StateField, ...]: + """Return fields in declaration order for one ownership scope.""" + + return tuple(field for field in self.fields if field.scope is scope) + + def field(self, name: str) -> StateField: + """Return one declared field by name.""" + + for field in self.fields: + if field.matches(name): + return field + raise KeyError(name) + + +@dataclass(frozen=True) +class TrainingSupport: + """One qualified training layout capability.""" + + layout: ParameterLayout + process_group_scope: ProcessGroupScope + mesh_dimensions: Optional[Sequence[int]] = None + sharded_mode: Optional[str] = None + requires_complete_parameter_storage: bool = False + requires_complete_logical_matrix: bool = False + + def __post_init__(self) -> None: + if self.mesh_dimensions is not None: + object.__setattr__(self, "mesh_dimensions", _tuple(self.mesh_dimensions)) + _validate_dimensions( + "mesh_dimensions", self.mesh_dimensions, positive=True + ) + if not isinstance(self.layout, ParameterLayout): + raise TypeError("TrainingSupport.layout must be a ParameterLayout") + if not isinstance(self.process_group_scope, ProcessGroupScope): + raise TypeError( + "TrainingSupport.process_group_scope must be a ProcessGroupScope" + ) + + +@dataclass(frozen=True) +class CheckpointSupport: + """One transport's separately qualified checkpoint capability.""" + + transport: CheckpointTransport + same_topology: AbstractSet[ParameterLayout] + topology_changing: AbstractSet[ParameterLayout] + process_group_scope: ProcessGroupScope + topology_change_kinds: AbstractSet[TopologyChange] = frozenset() + mesh_dimensions: Optional[Sequence[int]] = None + required_sharded_modes: AbstractSet[str] = frozenset() + requires_collective: bool = False + atomic_load: bool = False + + def __post_init__(self) -> None: + object.__setattr__(self, "same_topology", _frozenset(self.same_topology)) + object.__setattr__(self, "topology_changing", _frozenset(self.topology_changing)) + object.__setattr__( + self, "topology_change_kinds", _frozenset(self.topology_change_kinds) + ) + object.__setattr__( + self, + "required_sharded_modes", + _frozenset(self.required_sharded_modes), + ) + if self.mesh_dimensions is not None: + object.__setattr__(self, "mesh_dimensions", _tuple(self.mesh_dimensions)) + _validate_dimensions( + "mesh_dimensions", self.mesh_dimensions, positive=True + ) + if not isinstance(self.transport, CheckpointTransport): + raise TypeError("CheckpointSupport.transport must be a CheckpointTransport") + if not isinstance(self.process_group_scope, ProcessGroupScope): + raise TypeError( + "CheckpointSupport.process_group_scope must be a ProcessGroupScope" + ) + if bool(self.topology_changing) != bool(self.topology_change_kinds): + raise ValueError( + "topology-changing layouts and change kinds must be declared together" + ) + + +@dataclass(frozen=True) +class OptimizerCapabilities: + """Implemented integration capabilities, including explicit negative claims.""" + + training: Sequence[TrainingSupport] + checkpoints: Sequence[CheckpointSupport] + precisions: AbstractSet[Precision] + supported_parameter_ranks: Optional[Sequence[int]] + accepts_semantic_parameter_names: bool + canonical_parameter_fqns: bool + stable_shard_identity: bool + explicit_process_group_codebook_scope: bool + shard_rebinding: bool + post_sharding: bool + canonical_state_io: bool + atomic_state_movement: bool + state_offload: bool + + def __post_init__(self) -> None: + object.__setattr__(self, "training", _tuple(self.training)) + object.__setattr__(self, "checkpoints", _tuple(self.checkpoints)) + object.__setattr__(self, "precisions", _frozenset(self.precisions)) + if self.supported_parameter_ranks is not None: + object.__setattr__( + self, + "supported_parameter_ranks", + _tuple(self.supported_parameter_ranks), + ) + _validate_dimensions( + "supported_parameter_ranks", + self.supported_parameter_ranks, + positive=False, + ) + + +@dataclass(frozen=True) +class OptimizerChildContract: + """One named child of a composite optimizer contract.""" + + role: str + implementation: str + contract: Optional["OptimizerContract"] + + def __post_init__(self) -> None: + if not self.role: + raise ValueError("OptimizerChildContract.role must be non-empty") + if not self.implementation: + raise ValueError("OptimizerChildContract.implementation must be non-empty") + + +@dataclass(frozen=True) +class OptimizerContract: + """Complete read-only state and capability contract for one optimizer.""" + + implementation: str + state_layout: OptimizerStateLayout + capabilities: OptimizerCapabilities + children: Sequence[OptimizerChildContract] = () + schema_version: int = CONTRACT_SCHEMA_VERSION + + def __post_init__(self) -> None: + object.__setattr__(self, "children", _tuple(self.children)) + if not self.implementation: + raise ValueError("OptimizerContract.implementation must be non-empty") + if self.schema_version != CONTRACT_SCHEMA_VERSION: + raise ValueError( + "unsupported optimizer contract schema version: {}".format( + self.schema_version + ) + ) + roles = tuple(child.role for child in self.children) + if len(set(roles)) != len(roles): + raise ValueError("OptimizerContract child roles must be unique") + + +@runtime_checkable +class OptimizerContractProvider(Protocol): + """Structural protocol implemented by optimizers that expose a contract.""" + + def optimizer_contract(self) -> OptimizerContract: + """Return the optimizer's immutable state and capability declaration.""" + + +_ALL_PRECISIONS = frozenset( + {Precision.FLOAT32, Precision.BFLOAT16, Precision.FLOAT16, Precision.FLOAT64} +) +_DTENSOR_LAYOUT = ParameterLayout.DTENSOR_1D_DEFAULT_WORLD +_MUON_MODES = frozenset({"exact", "approx", "distributed"}) + + +def _common_fields() -> Tuple[StateField, ...]: + return ( + StateField( + "gefen_global_step", + StateScope.OPTIMIZER_COMMON, + StateGeometry.SCALAR, + True, + description="Canonical optimizer-wide update counter.", + ), + StateField( + "gefen_codebook", + StateScope.OPTIMIZER_COMMON, + StateGeometry.CODEBOOK, + True, + description="Canonical learned 256-entry momentum codebook.", + ), + StateField( + "gefen_deterministic", + StateScope.OPTIMIZER_COMMON, + StateGeometry.SCALAR, + True, + description="Checkpoint-bound deterministic execution policy.", + ), + ) + + +def _base_parameter_fields() -> Tuple[StateField, ...]: + return ( + StateField("name", StateScope.PARAMETER, StateGeometry.OPAQUE, True), + StateField("automatic_period", StateScope.PARAMETER, StateGeometry.SCALAR, True), + StateField("step", StateScope.PARAMETER, StateGeometry.SCALAR, True), + StateField("m_codebook", StateScope.PARAMETER, StateGeometry.PARAMETER, True), + StateField("m_magnitude", StateScope.PARAMETER, StateGeometry.BLOCK, True), + ) + + +def _derived_fields() -> Tuple[StateField, ...]: + return ( + StateField("stepsize", StateScope.DERIVED, StateGeometry.BLOCK, False), + StateField("_h_buf", StateScope.DERIVED, StateGeometry.BLOCK, False), + StateField("_capt_scalars", StateScope.DERIVED, StateGeometry.OPAQUE, False), + StateField("_capt_consts", StateScope.DERIVED, StateGeometry.OPAQUE, False), + StateField("_capt_consts_key", StateScope.DERIVED, StateGeometry.OPAQUE, False), + StateField("_capt_stack", StateScope.DERIVED, StateGeometry.OPAQUE, False), + StateField("_capt_row", StateScope.DERIVED, StateGeometry.OPAQUE, False), + StateField("_param_names", StateScope.DERIVED, StateGeometry.OPAQUE, False), + StateField("_lr_scalar_cache", StateScope.DERIVED, StateGeometry.SCALAR, False), + StateField("_fused_build_ok", StateScope.DERIVED, StateGeometry.SCALAR, False), + StateField("_static_mark_sig", StateScope.DERIVED, StateGeometry.OPAQUE, False), + StateField( + "m_codebook_shape", + StateScope.DERIVED, + StateGeometry.OPAQUE, + True, + description="Inert legacy checkpoint metadata retained for compatibility.", + ), + StateField( + "_gefen_codebook_by_device", + StateScope.DERIVED, + StateGeometry.CODEBOOK, + False, + ), + StateField( + "_gefen_codebook_lut_by_device", + StateScope.DERIVED, + StateGeometry.OPAQUE, + False, + ), + StateField("_sr_seed_by_device", StateScope.DERIVED, StateGeometry.SCALAR, False), + StateField( + "_gefen_global_step_by_device", + StateScope.DERIVED, + StateGeometry.SCALAR, + False, + ), + StateField("_capt_stacks", StateScope.DERIVED, StateGeometry.OPAQUE, False), + StateField( + "_gefen_rank_local_payload_", + StateScope.DERIVED, + StateGeometry.OPAQUE, + True, + key_match=StateKeyMatch.PREFIX, + description="PyTorch rank-local checkpoint transport.", + ), + StateField( + "_gefen_rank_local_member", + StateScope.DERIVED, + StateGeometry.OPAQUE, + True, + description="PyTorch rank-local checkpoint transport marker.", + ), + StateField( + "_gefen_checkpoint_metadata", + StateScope.DERIVED, + StateGeometry.OPAQUE, + True, + description="PyTorch transport mirror of optimizer-common state.", + ), + ) + + +def _base_training() -> Tuple[TrainingSupport, ...]: + return ( + TrainingSupport(ParameterLayout.REPLICATED, ProcessGroupScope.NONE), + TrainingSupport( + _DTENSOR_LAYOUT, + ProcessGroupScope.INFERRED_DEVICE_MESH, + mesh_dimensions=(1,), + ), + ) + + +def _negative_capabilities( + *, + training: Tuple[TrainingSupport, ...], + checkpoints: Tuple[CheckpointSupport, ...], + supported_parameter_ranks: Optional[Tuple[int, ...]], +) -> OptimizerCapabilities: + return OptimizerCapabilities( + training=training, + checkpoints=checkpoints, + precisions=_ALL_PRECISIONS, + supported_parameter_ranks=supported_parameter_ranks, + accepts_semantic_parameter_names=True, + canonical_parameter_fqns=False, + stable_shard_identity=False, + explicit_process_group_codebook_scope=False, + shard_rebinding=False, + post_sharding=False, + canonical_state_io=False, + atomic_state_movement=False, + state_offload=False, + ) + + +def _gefen_contract(*, factored_v_2d: bool) -> OptimizerContract: + block_fields = ( + StateField("vmean", StateScope.PARAMETER, StateGeometry.BLOCK, True), + StateField("vmean_step", StateScope.PARAMETER, StateGeometry.SCALAR, True), + ) + factored_fields = ( + StateField("v_row", StateScope.PARAMETER, StateGeometry.ROW, True), + StateField("v_col", StateScope.PARAMETER, StateGeometry.COLUMN, True), + StateField("factored_step", StateScope.PARAMETER, StateGeometry.SCALAR, True), + ) + layouts = frozenset( + { + ParameterLayout.REPLICATED, + ParameterLayout.FLATTENED_ELEMENT_SHARD, + _DTENSOR_LAYOUT, + } + ) + base_names = tuple(field.name for field in _base_parameter_fields()) + variants = [ + StateVariant( + "name_only", + ("name",), + layouts, + StateExtent.METADATA_ONLY, + initialized=False, + ), + ] + factored_names = tuple(field.name for field in factored_fields) + block_names = tuple(field.name for field in block_fields) + legacy_block_names = ("vmean",) + if factored_v_2d: + variants.extend( + ( + StateVariant( + "block_second_moment_non_2d", + base_names + block_names, + frozenset({ParameterLayout.REPLICATED}), + StateExtent.LOCAL_STORAGE, + excluded_parameter_ranks=(2,), + ), + StateVariant( + "block_second_moment_sharded", + base_names + block_names, + frozenset( + { + ParameterLayout.FLATTENED_ELEMENT_SHARD, + _DTENSOR_LAYOUT, + } + ), + StateExtent.LOCAL_STORAGE, + ), + StateVariant( + "factored_second_moment", + base_names + factored_names, + frozenset({ParameterLayout.REPLICATED}), + StateExtent.GLOBAL_PARAMETER, + parameter_ranks=(2,), + ), + StateVariant( + "block_state_pending_factored_initialization", + base_names + block_names, + frozenset({ParameterLayout.REPLICATED}), + StateExtent.GLOBAL_PARAMETER, + parameter_ranks=(2,), + migration_only=True, + ), + StateVariant( + "legacy_block_state_pending_factored_initialization", + base_names + legacy_block_names, + frozenset({ParameterLayout.REPLICATED}), + StateExtent.GLOBAL_PARAMETER, + parameter_ranks=(2,), + migration_only=True, + ), + StateVariant( + "factored_with_retained_block_state", + base_names + factored_names + block_names, + frozenset({ParameterLayout.REPLICATED}), + StateExtent.GLOBAL_PARAMETER, + parameter_ranks=(2,), + inactive_fields=block_names, + migration_only=True, + description="Factored-v state after loading a block-v checkpoint.", + ), + StateVariant( + "factored_with_retained_legacy_block_state", + base_names + factored_names + legacy_block_names, + frozenset({ParameterLayout.REPLICATED}), + StateExtent.GLOBAL_PARAMETER, + parameter_ranks=(2,), + inactive_fields=legacy_block_names, + migration_only=True, + description="Factored-v state after loading a pre-vmean-counter checkpoint.", + ), + ) + ) + else: + variants.extend( + ( + StateVariant( + "block_second_moment", + base_names + block_names, + layouts, + StateExtent.LOCAL_STORAGE, + ), + StateVariant( + "factored_state_pending_block_initialization", + base_names + factored_names, + frozenset({ParameterLayout.REPLICATED}), + StateExtent.GLOBAL_PARAMETER, + parameter_ranks=(2,), + migration_only=True, + ), + StateVariant( + "legacy_block_without_vmean_counter", + base_names + legacy_block_names, + layouts, + StateExtent.LOCAL_STORAGE, + migration_only=True, + ), + StateVariant( + "block_with_retained_factored_state", + base_names + block_names + factored_names, + frozenset({ParameterLayout.REPLICATED}), + StateExtent.LOCAL_STORAGE, + parameter_ranks=(2,), + inactive_fields=factored_names, + migration_only=True, + description="Block-v state after loading a factored-v checkpoint.", + ), + ) + ) + fields = ( + _common_fields() + + _base_parameter_fields() + + block_fields + + factored_fields + ) + fields += _derived_fields() + training = _base_training() + ( + TrainingSupport( + ParameterLayout.FLATTENED_ELEMENT_SHARD, + ProcessGroupScope.NONE, + ), + ) + checkpoints = ( + CheckpointSupport( + CheckpointTransport.NATIVE_OPTIMIZER, + frozenset( + { + ParameterLayout.REPLICATED, + ParameterLayout.FLATTENED_ELEMENT_SHARD, + } + ), + frozenset(), + ProcessGroupScope.NONE, + ), + CheckpointSupport( + CheckpointTransport.PYTORCH_RANK_LOCAL, + frozenset({_DTENSOR_LAYOUT}), + frozenset(), + ProcessGroupScope.DEFAULT_WORLD, + mesh_dimensions=(1,), + requires_collective=True, + atomic_load=True, + ), + ) + return OptimizerContract( + implementation="gefen.Gefen", + state_layout=OptimizerStateLayout(fields, tuple(variants)), + capabilities=_negative_capabilities( + training=training, + checkpoints=checkpoints, + supported_parameter_ranks=None, + ), + ) + + +def _muon_initialized_variant( + *, + name: str, + field_names: Tuple[str, ...], + layout: ParameterLayout, + extent: StateExtent, + mode: str, + role: ParameterStateRole = ParameterStateRole.ANY, +) -> StateVariant: + return StateVariant( + name, + field_names, + frozenset({layout}), + extent, + role=role, + parameter_ranks=(2,), + sharded_mode=mode, + ) + + +def _gefen_muon_contract( + *, + sharded_modes: FrozenSet[str], + normuon_modes: FrozenSet[str], + non_normuon_modes: FrozenSet[str], +) -> OptimizerContract: + sharded_modes = _frozenset(sharded_modes) + normuon_modes = _frozenset(normuon_modes) + non_normuon_modes = _frozenset(non_normuon_modes) + normuon_fields = ( + StateField("normuon_v", StateScope.PARAMETER, StateGeometry.ROW, True), + StateField("normuon_step", StateScope.PARAMETER, StateGeometry.SCALAR, True), + ) + fields = _common_fields() + _base_parameter_fields() + if "distributed" in sharded_modes: + fields += ( + StateField( + "gefen_muon_distributed", + StateScope.DERIVED, + StateGeometry.OPAQUE, + True, + applicable_sharded_modes=frozenset({"distributed"}), + description="Native Parallel-Muon ownership manifest.", + ), + ) + if normuon_modes: + fields += normuon_fields + fields += _derived_fields() + + base_names = tuple(field.name for field in _base_parameter_fields()) + normuon_names = base_names + tuple(field.name for field in normuon_fields) + variants = [] + for mode in sorted(sharded_modes): + variants.append( + StateVariant( + "name_only_replicated_" + mode, + ("name",), + frozenset({ParameterLayout.REPLICATED}), + StateExtent.METADATA_ONLY, + initialized=False, + parameter_ranks=(2,), + sharded_mode=mode, + ) + ) + if mode == "distributed": + variants.append( + StateVariant( + "distributed_owner_name_only", + ("name",), + frozenset({_DTENSOR_LAYOUT}), + StateExtent.METADATA_ONLY, + role=ParameterStateRole.OWNER, + initialized=False, + parameter_ranks=(2,), + sharded_mode=mode, + ) + ) + else: + variants.append( + StateVariant( + "name_only_dtensor_" + mode, + ("name",), + frozenset({_DTENSOR_LAYOUT}), + StateExtent.METADATA_ONLY, + initialized=False, + parameter_ranks=(2,), + sharded_mode=mode, + ) + ) + for mode_set, field_names, prefix in ( + (non_normuon_modes, base_names, "quantized_muon"), + (normuon_modes, normuon_names, "quantized_normuon"), + ): + if not mode_set: + continue + for mode in sorted(mode_set): + variants.append( + _muon_initialized_variant( + name=prefix + "_replicated_" + mode, + field_names=field_names, + layout=ParameterLayout.REPLICATED, + extent=StateExtent.GLOBAL_PARAMETER, + mode=mode, + ) + ) + if "approx" in mode_set: + variants.append( + _muon_initialized_variant( + name=prefix + "_local", + field_names=field_names, + layout=_DTENSOR_LAYOUT, + extent=StateExtent.LOCAL_STORAGE, + mode="approx", + ) + ) + if "exact" in mode_set: + variants.append( + _muon_initialized_variant( + name=prefix + "_global", + field_names=field_names, + layout=_DTENSOR_LAYOUT, + extent=StateExtent.GLOBAL_PARAMETER, + mode="exact", + ) + ) + if "distributed" in mode_set: + variants.append( + _muon_initialized_variant( + name=prefix + "_owner", + field_names=field_names, + layout=_DTENSOR_LAYOUT, + extent=StateExtent.OWNER_PARAMETER, + mode="distributed", + role=ParameterStateRole.OWNER, + ) + ) + if "distributed" in sharded_modes: + variants.extend( + ( + StateVariant( + "distributed_non_owner_name_only", + ("name",), + frozenset({_DTENSOR_LAYOUT}), + StateExtent.METADATA_ONLY, + role=ParameterStateRole.NON_OWNER, + initialized=False, + parameter_ranks=(2,), + sharded_mode="distributed", + ), + StateVariant( + "distributed_non_owner", + ("name", "automatic_period"), + frozenset({_DTENSOR_LAYOUT}), + StateExtent.METADATA_ONLY, + role=ParameterStateRole.NON_OWNER, + initialized=False, + parameter_ranks=(2,), + sharded_mode="distributed", + ), + ) + ) + + training = tuple( + support + for mode in sorted(sharded_modes) + for support in ( + TrainingSupport( + ParameterLayout.REPLICATED, + ProcessGroupScope.NONE, + sharded_mode=mode, + requires_complete_logical_matrix=True, + ), + TrainingSupport( + _DTENSOR_LAYOUT, + ProcessGroupScope.INFERRED_DEVICE_MESH, + mesh_dimensions=(1,), + sharded_mode=mode, + requires_complete_logical_matrix=mode in {"exact", "distributed"}, + ), + ) + ) + checkpoints = [ + CheckpointSupport( + CheckpointTransport.NATIVE_OPTIMIZER, + frozenset({ParameterLayout.REPLICATED}), + frozenset(), + ProcessGroupScope.NONE, + required_sharded_modes=sharded_modes, + ) + ] + if ( + "approx" in sharded_modes + and sharded_modes.issubset({"approx", "distributed"}) + ): + checkpoints.append( + CheckpointSupport( + CheckpointTransport.PYTORCH_RANK_LOCAL, + frozenset({_DTENSOR_LAYOUT}), + frozenset(), + ProcessGroupScope.DEFAULT_WORLD, + mesh_dimensions=(1,), + required_sharded_modes=sharded_modes, + requires_collective=True, + atomic_load=True, + ) + ) + if sharded_modes == frozenset({"distributed"}): + checkpoints.append( + CheckpointSupport( + CheckpointTransport.NATIVE_OPTIMIZER, + frozenset({_DTENSOR_LAYOUT}), + frozenset({_DTENSOR_LAYOUT}), + ProcessGroupScope.INFERRED_DEVICE_MESH, + topology_change_kinds=frozenset( + {TopologyChange.WORLD_SIZE_OWNER_REDISTRIBUTION} + ), + mesh_dimensions=(1,), + required_sharded_modes=frozenset({"distributed"}), + requires_collective=True, + atomic_load=True, + ) + ) + return OptimizerContract( + implementation="gefen.GefenMuon", + state_layout=OptimizerStateLayout(fields, tuple(variants)), + capabilities=_negative_capabilities( + training=training, + checkpoints=tuple(checkpoints), + supported_parameter_ranks=(2,), + ), + ) + + +def _hybrid_contract( + *, + muon: Optional[OptimizerContract], + backup: Optional[OptimizerContract], + backup_implementation: str, +) -> OptimizerContract: + fields = ( + StateField( + "backup_optimizer", + StateScope.OPTIMIZER_COMMON, + StateGeometry.OPAQUE, + True, + ), + ) + children = [] + if muon is not None: + children.append(OptimizerChildContract("muon", muon.implementation, muon)) + if backup_implementation: + children.append( + OptimizerChildContract("backup", backup_implementation, backup) + ) + if muon is None: + training = _base_training() + else: + training = muon.capabilities.training + checkpoints = ( + CheckpointSupport( + CheckpointTransport.COMPOSITE_NATIVE, + frozenset({ParameterLayout.REPLICATED}), + frozenset(), + ProcessGroupScope.NONE, + ), + ) + return OptimizerContract( + implementation="gefen.GefenMuonHybrid", + state_layout=OptimizerStateLayout( + fields, + (), + composite_namespaces=("muon", "backup"), + ), + capabilities=_negative_capabilities( + training=training, + checkpoints=checkpoints, + supported_parameter_ranks=None, + ), + children=tuple(children), + ) + + +__all__ = [ + "CONTRACT_SCHEMA_VERSION", + "CheckpointSupport", + "CheckpointTransport", + "OptimizerCapabilities", + "OptimizerChildContract", + "OptimizerContract", + "OptimizerContractProvider", + "OptimizerStateLayout", + "ParameterLayout", + "ParameterStateRole", + "Precision", + "ProcessGroupScope", + "StateExtent", + "StateField", + "StateGeometry", + "StateKeyMatch", + "StateScope", + "StateVariant", + "TrainingSupport", + "TopologyChange", +] diff --git a/src/gefen/gefen.py b/src/gefen/gefen.py index 6cfd3e5..c6e5251 100644 --- a/src/gefen/gefen.py +++ b/src/gefen/gefen.py @@ -22,6 +22,7 @@ import torch import torch.nn as nn +from gefen.contracts import OptimizerContract, _gefen_contract from gefen.partitioning import find_period_by_block_variance import gefen.quantization as quantization_module from gefen.kernels.automatic_vmean import ( @@ -1068,6 +1069,11 @@ def _step_supports_amp_scaling(self) -> bool: # true-FP16 storage with DTensors after a collective presence preflight. return _amp_native_scaling_required(self) + def optimizer_contract(self) -> OptimizerContract: + """Return the immutable state-layout and integration capability contract.""" + + return _gefen_contract(factored_v_2d=self._factored_v_2d) + @staticmethod def _normalize_param_groups(params): if isinstance(params, torch.Tensor): diff --git a/src/gefen/gefen_muon.py b/src/gefen/gefen_muon.py index 4196895..ffb275b 100644 --- a/src/gefen/gefen_muon.py +++ b/src/gefen/gefen_muon.py @@ -6,6 +6,7 @@ import torch import torch.nn as nn +from gefen.contracts import OptimizerContract, _gefen_muon_contract from gefen.gefen import ( Gefen, _amp_prepare_optimizer_step, @@ -745,6 +746,28 @@ def __init__( verbose=verbose, ) + def optimizer_contract(self) -> OptimizerContract: + """Return the immutable Muon state and capability contract.""" + + sharded_modes = frozenset( + group["sharded_mode"] for group in self.param_groups + ) + normuon_modes = frozenset( + group["sharded_mode"] + for group in self.param_groups + if group.get("normuon", False) + ) + non_normuon_modes = frozenset( + group["sharded_mode"] + for group in self.param_groups + if not group.get("normuon", False) + ) + return _gefen_muon_contract( + sharded_modes=sharded_modes, + normuon_modes=normuon_modes, + non_normuon_modes=non_normuon_modes, + ) + def add_param_group(self, param_group): """Add a validated 2D Muon parameter group atomically. diff --git a/src/gefen/hybrid.py b/src/gefen/hybrid.py index 8837e18..72e3b50 100644 --- a/src/gefen/hybrid.py +++ b/src/gefen/hybrid.py @@ -60,6 +60,7 @@ import torch import torch.nn as nn +from gefen.contracts import OptimizerContract, _hybrid_contract from gefen.gefen import ( Gefen, _amp_native_scaling_required, @@ -725,6 +726,30 @@ def state_dict(self): state_dict = hook_result return state_dict + def optimizer_contract(self) -> OptimizerContract: + """Return the composite contract without flattening either child schema.""" + + muon_contract = ( + self.muon.optimizer_contract() if self.muon is not None else None + ) + backup_contract = ( + self.backup.optimizer_contract() + if self.backup is not None and hasattr(self.backup, "optimizer_contract") + else None + ) + backup_implementation = "" + if self.backup is not None: + backup_implementation = ( + backup_contract.implementation + if backup_contract is not None + else type(self.backup).__module__ + "." + type(self.backup).__qualname__ + ) + return _hybrid_contract( + muon=muon_contract, + backup=backup_contract, + backup_implementation=backup_implementation, + ) + def load_state_dict(self, state_dict): # Instance load pre-hooks first (a pre-hook may return a replacement # dict -- e.g. one that converts a foreign schema), mirroring diff --git a/tests/test_muon_distributed_checkpoint_safety.py b/tests/test_muon_distributed_checkpoint_safety.py index 8688999..66fa37f 100644 --- a/tests/test_muon_distributed_checkpoint_safety.py +++ b/tests/test_muon_distributed_checkpoint_safety.py @@ -339,6 +339,8 @@ def _mixed_cpu_checkpoint_worker(rank, world, port, result_queue): import torch.distributed as dist from torch.distributed.tensor import Shard, distribute_tensor, init_device_mesh + from gefen import ParameterLayout, ParameterStateRole, StateExtent, StateScope + try: os.environ["MASTER_ADDR"] = "127.0.0.1" os.environ["MASTER_PORT"] = port @@ -434,6 +436,38 @@ def make_fallback_pair(parallel_value, fallback_value): assign(source_params[0], full(922) * 0.01) source_params[1].grad = full(923) * 0.01 source.step() + contract = source.optimizer_contract() + if rank == 0: + parallel_variant = next( + item + for item in contract.state_layout.parameter_variants + if item.name == "quantized_muon_owner" + ) + assert "m_codebook" in source.state[source_params[0]] + assert ( + source.state[source_params[0]]["m_codebook"].numel() + == initial_parallel.numel() + ) + assert parallel_variant.role is ParameterStateRole.OWNER + assert parallel_variant.extent is StateExtent.OWNER_PARAMETER + else: + parallel_variant = next( + item + for item in contract.state_layout.parameter_variants + if item.name == "distributed_non_owner" + ) + authoritative_keys = { + key + for key in source.state[source_params[0]] + if contract.state_layout.field(key).scope is StateScope.PARAMETER + } + assert authoritative_keys == {"name", "automatic_period"} + assert parallel_variant.role is ParameterStateRole.NON_OWNER + assert parallel_variant.extent is StateExtent.METADATA_ONLY + assert parallel_variant.layouts == frozenset( + {ParameterLayout.DTENSOR_1D_DEFAULT_WORLD} + ) + assert parallel_variant.sharded_mode == "distributed" checkpoint = _clone(source.state_dict()) saved_ids = checkpoint["param_groups"][0]["params"] marker = checkpoint["gefen_muon_distributed"] diff --git a/tests/test_muon_fsdp2_approx.py b/tests/test_muon_fsdp2_approx.py index 52a0b68..af378b7 100644 --- a/tests/test_muon_fsdp2_approx.py +++ b/tests/test_muon_fsdp2_approx.py @@ -43,6 +43,7 @@ def _worker(rank, world, case, port, q): import torch.distributed as dist from torch.distributed.tensor import Shard, distribute_tensor, init_device_mesh + from gefen import ParameterLayout, StateExtent from gefen.gefen_muon import GefenMuon os.environ["MASTER_ADDR"] = "127.0.0.1" @@ -64,6 +65,24 @@ def _worker(rank, world, case, port, q): p.grad = distribute_tensor(full_grad.clone(), mesh, [Shard(0)]) opt.step() local_rows = p.to_local().shape[0] + contract = opt.optimizer_contract() + if p.to_local().numel() == 0: + variant_name = "name_only_dtensor_approx" + expected_extent = StateExtent.METADATA_ONLY + else: + variant_name = "quantized_muon_local" + expected_extent = StateExtent.LOCAL_STORAGE + assert opt.state[p]["m_codebook"].numel() == p.to_local().numel() + variant = next( + item + for item in contract.state_layout.parameter_variants + if item.name == variant_name + ) + assert variant.layouts == frozenset( + {ParameterLayout.DTENSOR_1D_DEFAULT_WORLD} + ) + assert variant.sharded_mode == "approx" + assert variant.extent is expected_extent gathered = p.detach().full_tensor() # collective; all ranks must call if rank == 0: diff --git a/tests/test_muon_fsdp2_parity.py b/tests/test_muon_fsdp2_parity.py index 9ddf216..1313800 100644 --- a/tests/test_muon_fsdp2_parity.py +++ b/tests/test_muon_fsdp2_parity.py @@ -82,6 +82,7 @@ def _worker(rank, world, case, port, q): import torch.distributed as dist from torch.distributed.tensor import Shard, distribute_tensor, init_device_mesh + from gefen import ParameterLayout, StateExtent from gefen.gefen_muon import GefenMuon os.environ["MASTER_ADDR"] = "127.0.0.1" @@ -109,6 +110,24 @@ def _worker(rank, world, case, port, q): ) opt.step() + contract = opt.optimizer_contract() + variant_name = ( + "quantized_normuon_global" + if case.startswith("normuon-") + else "quantized_muon_global" + ) + variant = next( + item + for item in contract.state_layout.parameter_variants + if item.name == variant_name + ) + assert variant.layouts == frozenset( + {ParameterLayout.DTENSOR_1D_DEFAULT_WORLD} + ) + assert variant.sharded_mode == "exact" + assert variant.extent is StateExtent.GLOBAL_PARAMETER + assert opt.state[p]["m_codebook"].numel() == full_init.numel() + gathered = p.detach().full_tensor() # collective; all ranks must call if rank == 0: diff --git a/tests/test_optimizer_contracts.py b/tests/test_optimizer_contracts.py new file mode 100644 index 0000000..5b3d81c --- /dev/null +++ b/tests/test_optimizer_contracts.py @@ -0,0 +1,630 @@ +"""CPU coverage for public optimizer capability and state-layout contracts.""" + +import copy +from dataclasses import FrozenInstanceError + +import pytest +import torch + +import gefen +from gefen import ( + CONTRACT_SCHEMA_VERSION, + CheckpointTransport, + Gefen, + GefenMuon, + GefenMuonHybrid, + OptimizerContractProvider, + OptimizerStateLayout, + ParameterLayout, + ParameterStateRole, + Precision, + ProcessGroupScope, + StateExtent, + StateField, + StateGeometry, + StateKeyMatch, + StateScope, + StateVariant, + TopologyChange, +) + + +_DTENSOR = ParameterLayout.DTENSOR_1D_DEFAULT_WORLD + + +def _values_equal(left, right): + if torch.is_tensor(left) or torch.is_tensor(right): + return torch.is_tensor(left) and torch.is_tensor(right) and torch.equal( + left, right + ) + if type(left) is not type(right): + return False + if isinstance(left, dict): + return set(left) == set(right) and all( + _values_equal(left[key], right[key]) for key in left + ) + if isinstance(left, (list, tuple)): + return len(left) == len(right) and all( + _values_equal(a, b) for a, b in zip(left, right) + ) + return left == right + + +def _declares_state_key(contract, key): + try: + contract.state_layout.field(key) + except KeyError: + return False + return True + + +def _step_with_fixed_grad(optimizer, params): + for index, param in enumerate(params): + values = torch.arange(param.numel(), dtype=param.dtype).reshape_as(param) + param.grad = (values + index + 1) / max(1, param.numel()) + optimizer.step() + + +def _persistent_parameter_keys(contract, state): + return { + key + for key in state + if _declares_state_key(contract, key) + and contract.state_layout.field(key).scope is StateScope.PARAMETER + } + + +def _matching_variants( + contract, + state, + *, + layout=ParameterLayout.REPLICATED, + parameter_rank=2, + sharded_mode=None, + role=ParameterStateRole.ANY, +): + keys = _persistent_parameter_keys(contract, state) + return [ + variant + for variant in contract.state_layout.parameter_variants + if set(variant.fields) == keys + and layout in variant.layouts + and ( + variant.parameter_ranks is None + or parameter_rank in variant.parameter_ranks + ) + and parameter_rank not in variant.excluded_parameter_ranks + and variant.sharded_mode == sharded_mode + and (variant.role is ParameterStateRole.ANY or variant.role is role) + ] + + +def _training_support(contract, layout, mode=None): + matches = [ + item + for item in contract.capabilities.training + if item.layout is layout + and (mode is None or item.sharded_mode == mode) + ] + assert len(matches) == 1 + return matches[0] + + +def _checkpoint_support(contract, transport, layout): + return [ + item + for item in contract.capabilities.checkpoints + if item.transport is transport + and (layout in item.same_topology or layout in item.topology_changing) + ] + + +@pytest.mark.parametrize("factored_v_2d", [False, True]) +def test_plain_contract_matches_live_persistent_state(factored_v_2d): + param = torch.nn.Parameter(torch.arange(16, dtype=torch.float32).reshape(4, 4)) + optimizer = Gefen( + [("layer.weight", param)], + fused=False, + factored_v_2d=factored_v_2d, + ) + assert isinstance(optimizer, OptimizerContractProvider) + contract = optimizer.optimizer_contract() + + assert contract.schema_version == CONTRACT_SCHEMA_VERSION + assert contract.implementation == "gefen.Gefen" + assert contract.capabilities.supported_parameter_ranks is None + assert { + item.layout for item in contract.capabilities.training + } == { + ParameterLayout.REPLICATED, + ParameterLayout.FLATTENED_ELEMENT_SHARD, + _DTENSOR, + } + dtensor_training = _training_support(contract, _DTENSOR) + assert dtensor_training.mesh_dimensions == (1,) + assert dtensor_training.process_group_scope is ProcessGroupScope.INFERRED_DEVICE_MESH + dcp = _checkpoint_support( + contract, CheckpointTransport.PYTORCH_RANK_LOCAL, _DTENSOR + ) + assert len(dcp) == 1 + assert dcp[0].process_group_scope is ProcessGroupScope.DEFAULT_WORLD + assert dcp[0].mesh_dimensions == (1,) + assert dcp[0].requires_collective + assert dcp[0].atomic_load + assert not dcp[0].topology_changing + native = next( + item + for item in contract.capabilities.checkpoints + if item.transport is CheckpointTransport.NATIVE_OPTIMIZER + ) + assert not native.atomic_load + assert contract.capabilities.accepts_semantic_parameter_names + assert not contract.capabilities.canonical_parameter_fqns + assert not contract.capabilities.stable_shard_identity + assert not contract.capabilities.explicit_process_group_codebook_scope + assert not contract.capabilities.shard_rebinding + assert not contract.capabilities.post_sharding + assert not contract.capabilities.canonical_state_io + assert not contract.capabilities.atomic_state_movement + assert not contract.capabilities.state_offload + assert Precision.FLOAT64 in contract.capabilities.precisions + flattened = _training_support( + contract, ParameterLayout.FLATTENED_ELEMENT_SHARD + ) + assert flattened.process_group_scope is ProcessGroupScope.NONE + + assert optimizer.state[param] == {"name": "layer.weight"} + name_only = next( + item + for item in contract.state_layout.parameter_variants + if item.name == "name_only" + ) + assert name_only.fields == ("name",) + assert name_only.extent is StateExtent.METADATA_ONLY + assert not name_only.initialized + + _step_with_fixed_grad(optimizer, [param]) + state = optimizer.state[param] + assert all(_declares_state_key(contract, key) for key in state) + expected_variant = ( + "factored_second_moment" if factored_v_2d else "block_second_moment" + ) + variant = next( + item + for item in contract.state_layout.parameter_variants + if item.name == expected_variant + ) + persistent_keys = { + key + for key in state + if contract.state_layout.field(key).scope is StateScope.PARAMETER + } + assert persistent_keys == set(variant.fields) + + state_dict = optimizer.state_dict() + common = { + field.name + for field in contract.state_layout.fields_for_scope( + StateScope.OPTIMIZER_COMMON + ) + } + assert common == {"gefen_global_step", "gefen_codebook", "gefen_deterministic"} + assert common.issubset(state_dict) + + +@pytest.mark.parametrize( + ("source_factored", "target_factored", "migrated_variant"), + [ + (True, False, "block_with_retained_factored_state"), + (False, True, "factored_with_retained_block_state"), + ], +) +def test_plain_contract_declares_factored_v_checkpoint_migration_states( + source_factored, target_factored, migrated_variant +): + initial = torch.arange(16, dtype=torch.float32).reshape(4, 4) + source_param = torch.nn.Parameter(initial.clone()) + source = Gefen( + [("layer.weight", source_param)], + fused=False, + factored_v_2d=source_factored, + ) + _step_with_fixed_grad(source, [source_param]) + checkpoint = copy.deepcopy(source.state_dict()) + + target_param = torch.nn.Parameter(initial.clone()) + target = Gefen( + [("layer.weight", target_param)], + fused=False, + factored_v_2d=target_factored, + ) + target.load_state_dict(checkpoint) + contract = target.optimizer_contract() + pending = _matching_variants(contract, target.state[target_param]) + expected_pending = ( + "factored_state_pending_block_initialization" + if source_factored + else "block_state_pending_factored_initialization" + ) + assert [variant.name for variant in pending] == [expected_pending] + assert pending[0].migration_only + + _step_with_fixed_grad(target, [target_param]) + matches = _matching_variants(contract, target.state[target_param]) + assert [variant.name for variant in matches] == [migrated_variant] + assert matches[0].inactive_fields + + +@pytest.mark.parametrize("target_factored", [False, True]) +def test_plain_contract_declares_legacy_vmean_counter_transition(target_factored): + initial = torch.arange(16, dtype=torch.float32).reshape(4, 4) + source_param = torch.nn.Parameter(initial.clone()) + source = Gefen( + [("layer.weight", source_param)], + fused=False, + factored_v_2d=False, + ) + _step_with_fixed_grad(source, [source_param]) + checkpoint = copy.deepcopy(source.state_dict()) + checkpoint["state"][0].pop("vmean_step") + + target_param = torch.nn.Parameter(initial.clone()) + target = Gefen( + [("layer.weight", target_param)], + fused=False, + factored_v_2d=target_factored, + ) + target.load_state_dict(checkpoint) + contract = target.optimizer_contract() + pending = _matching_variants(contract, target.state[target_param]) + expected_pending = ( + "legacy_block_state_pending_factored_initialization" + if target_factored + else "legacy_block_without_vmean_counter" + ) + assert [variant.name for variant in pending] == [expected_pending] + + _step_with_fixed_grad(target, [target_param]) + matches = _matching_variants(contract, target.state[target_param]) + expected_after = ( + "factored_with_retained_legacy_block_state" + if target_factored + else "block_second_moment" + ) + assert [variant.name for variant in matches] == [expected_after] + + +def test_plain_contract_classifies_legacy_m_codebook_shape_as_inert(): + initial = torch.arange(16, dtype=torch.float32).reshape(4, 4) + source_param = torch.nn.Parameter(initial.clone()) + source = Gefen([("layer.weight", source_param)], fused=False) + _step_with_fixed_grad(source, [source_param]) + checkpoint = copy.deepcopy(source.state_dict()) + checkpoint["state"][0]["m_codebook_shape"] = tuple( + checkpoint["state"][0]["m_codebook"].shape + ) + + target_param = torch.nn.Parameter(initial.clone()) + target = Gefen([("layer.weight", target_param)], fused=False) + target.load_state_dict(checkpoint) + contract = target.optimizer_contract() + field = contract.state_layout.field("m_codebook_shape") + assert field.scope is StateScope.DERIVED + assert not field.authoritative + assert "m_codebook_shape" in target.state[target_param] + assert _matching_variants(contract, target.state[target_param]) + + +@pytest.mark.parametrize( + ("sharded_mode", "normuon", "same_topology_dtensor", "reshard_dtensor"), + [ + ("exact", False, False, False), + ("approx", False, True, False), + ("distributed", True, True, True), + ], +) +def test_muon_contract_separates_mode_topology_and_state_extent( + sharded_mode, normuon, same_topology_dtensor, reshard_dtensor +): + param = torch.nn.Parameter(torch.arange(16, dtype=torch.float32).reshape(4, 4)) + optimizer = GefenMuon( + [("layer.weight", param)], + fused=False, + sharded_mode=sharded_mode, + normuon=normuon, + ) + contract = optimizer.optimizer_contract() + + assert contract.implementation == "gefen.GefenMuon" + assert contract.capabilities.supported_parameter_ranks == (2,) + replicated = _training_support(contract, ParameterLayout.REPLICATED) + assert replicated.requires_complete_logical_matrix + assert not replicated.requires_complete_parameter_storage + dtensor = _training_support(contract, _DTENSOR, sharded_mode) + assert dtensor.mesh_dimensions == (1,) + assert dtensor.requires_complete_logical_matrix == ( + sharded_mode in ("exact", "distributed") + ) + assert not dtensor.requires_complete_parameter_storage + + same_topology = any( + _DTENSOR in support.same_topology + for support in contract.capabilities.checkpoints + ) + topology_changing = any( + _DTENSOR in support.topology_changing + for support in contract.capabilities.checkpoints + ) + assert same_topology is same_topology_dtensor + assert topology_changing is reshard_dtensor + if same_topology_dtensor: + support = next( + item + for item in contract.capabilities.checkpoints + if _DTENSOR in item.same_topology + ) + assert support.mesh_dimensions == (1,) + assert support.requires_collective + assert support.atomic_load + if reshard_dtensor: + support = next( + item + for item in contract.capabilities.checkpoints + if _DTENSOR in item.topology_changing + ) + assert support.topology_change_kinds == frozenset( + {TopologyChange.WORLD_SIZE_OWNER_REDISTRIBUTION} + ) + assert TopologyChange.PLACEMENT_RESHARD not in support.topology_change_kinds + + common_names = { + field.name + for field in contract.state_layout.fields_for_scope( + StateScope.OPTIMIZER_COMMON + ) + } + assert "gefen_muon_distributed" not in common_names + if sharded_mode == "distributed": + manifest = contract.state_layout.field("gefen_muon_distributed") + assert manifest.scope is StateScope.DERIVED + assert not manifest.authoritative + owner = next( + item + for item in contract.state_layout.parameter_variants + if item.name == "quantized_normuon_owner" + ) + assert owner.extent is StateExtent.OWNER_PARAMETER + assert owner.role is ParameterStateRole.OWNER + non_owner = next( + item + for item in contract.state_layout.parameter_variants + if item.name == "distributed_non_owner" + ) + assert non_owner.fields == ("name", "automatic_period") + assert non_owner.role is ParameterStateRole.NON_OWNER + owner_lazy = _matching_variants( + contract, + {"name": "layer.weight"}, + layout=_DTENSOR, + sharded_mode="distributed", + role=ParameterStateRole.OWNER, + ) + non_owner_lazy = _matching_variants( + contract, + {"name": "layer.weight"}, + layout=_DTENSOR, + sharded_mode="distributed", + role=ParameterStateRole.NON_OWNER, + ) + assert [item.name for item in owner_lazy] == ["distributed_owner_name_only"] + assert [item.name for item in non_owner_lazy] == [ + "distributed_non_owner_name_only" + ] + + assert optimizer.state[param] == {"name": "layer.weight"} + _step_with_fixed_grad(optimizer, [param]) + state = optimizer.state[param] + assert all(_declares_state_key(contract, key) for key in state) + expected_variant = ( + "quantized_normuon_replicated_" + sharded_mode + if normuon + else "quantized_muon_replicated_" + sharded_mode + ) + variant = next( + item + for item in contract.state_layout.parameter_variants + if item.name == expected_variant + ) + persistent_keys = _persistent_parameter_keys(contract, state) + assert persistent_keys == set(variant.fields) + + +@pytest.mark.parametrize("backup_optimizer", ["gefen", "adamw"]) +def test_hybrid_contract_preserves_child_namespaces(backup_optimizer): + matrix = torch.nn.Parameter(torch.arange(16, dtype=torch.float32).reshape(4, 4)) + bias = torch.nn.Parameter(torch.arange(4, dtype=torch.float32)) + optimizer = GefenMuonHybrid( + [("layer.weight", matrix)], + [("layer.bias", bias)], + lr=1e-3, + fused=False, + backup_optimizer=backup_optimizer, + ) + contract = optimizer.optimizer_contract() + + assert contract.implementation == "gefen.GefenMuonHybrid" + assert tuple(child.role for child in contract.children) == ("muon", "backup") + assert contract.children[0].contract.implementation == "gefen.GefenMuon" + if backup_optimizer == "gefen": + assert contract.children[1].contract.implementation == "gefen.Gefen" + else: + assert contract.children[1].implementation == "torch.optim.adamw.AdamW" + assert contract.children[1].contract is None + assert contract.state_layout.composite_namespaces == ("muon", "backup") + assert {field.name for field in contract.state_layout.fields} == { + "backup_optimizer" + } + checkpoint = contract.capabilities.checkpoints + assert len(checkpoint) == 1 + assert checkpoint[0].transport is CheckpointTransport.COMPOSITE_NATIVE + assert checkpoint[0].same_topology == frozenset({ParameterLayout.REPLICATED}) + assert not checkpoint[0].topology_changing + + +def test_muon_contract_keeps_mixed_normuon_variants_in_one_mode(): + plain = torch.nn.Parameter(torch.ones(4, 4)) + normuon = torch.nn.Parameter(torch.ones(4, 4)) + optimizer = GefenMuon( + [ + {"params": [("plain", plain)], "normuon": False}, + {"params": [("normuon", normuon)], "normuon": True}, + ], + fused=False, + sharded_mode="exact", + ) + variants = { + item.name for item in optimizer.optimizer_contract().state_layout.parameter_variants + } + assert "quantized_muon_replicated_exact" in variants + assert "quantized_normuon_replicated_exact" in variants + + +def test_muon_mixed_approx_distributed_checkpoint_is_same_topology_only(): + first = torch.nn.Parameter(torch.ones(4, 4)) + second = torch.nn.Parameter(torch.ones(4, 4)) + optimizer = GefenMuon( + [ + {"params": [("first", first)], "sharded_mode": "approx"}, + {"params": [("second", second)], "sharded_mode": "distributed"}, + ], + fused=False, + ) + support = _checkpoint_support( + optimizer.optimizer_contract(), + CheckpointTransport.PYTORCH_RANK_LOCAL, + _DTENSOR, + ) + assert len(support) == 1 + assert support[0].required_sharded_modes == frozenset( + {"approx", "distributed"} + ) + assert not support[0].topology_changing + + +@pytest.mark.parametrize("backup_optimizer", ["gefen", "adamw"]) +def test_hybrid_approx_does_not_require_complete_dtensor_matrix(backup_optimizer): + matrix = torch.nn.Parameter(torch.ones(4, 4)) + bias = torch.nn.Parameter(torch.ones(4)) + optimizer = GefenMuonHybrid( + [("layer.weight", matrix)], + [("layer.bias", bias)], + lr=1e-3, + fused=False, + sharded_mode="approx", + backup_optimizer=backup_optimizer, + ) + contract = optimizer.optimizer_contract() + support = _training_support(contract, _DTENSOR, "approx") + assert not support.requires_complete_parameter_storage + assert not support.requires_complete_logical_matrix + + +def test_bare_parameters_do_not_claim_canonical_fqns(): + param = torch.nn.Parameter(torch.ones(4, 4)) + optimizer = Gefen([param], fused=False) + contract = optimizer.optimizer_contract() + assert optimizer.state[param]["name"] == "param_0" + assert contract.capabilities.accepts_semantic_parameter_names + assert not contract.capabilities.canonical_parameter_fqns + assert not contract.capabilities.stable_shard_identity + + +def test_contract_is_deeply_immutable_and_query_is_behavior_neutral(): + initial = torch.linspace(-1, 1, 16).reshape(4, 4) + param = torch.nn.Parameter(initial.clone()) + reference_param = torch.nn.Parameter(initial.clone()) + optimizer = Gefen([("layer.weight", param)], fused=False, factored_v_2d=False) + reference = Gefen( + [("layer.weight", reference_param)], fused=False, factored_v_2d=False + ) + + before = copy.deepcopy(optimizer.state_dict()) + contract = optimizer.optimizer_contract() + assert contract == optimizer.optimizer_contract() + assert _values_equal(before, optimizer.state_dict()) + with pytest.raises(FrozenInstanceError): + contract.implementation = "changed" + + source_fields = [ + StateField("name", StateScope.PARAMETER, StateGeometry.OPAQUE, True) + ] + source_variant_fields = ["name"] + source_layouts = {ParameterLayout.REPLICATED} + source_variants = [ + StateVariant( + "name_only", + source_variant_fields, + source_layouts, + StateExtent.METADATA_ONLY, + initialized=False, + ) + ] + layout = OptimizerStateLayout(source_fields, source_variants) + source_fields.clear() + source_variant_fields.clear() + source_layouts.clear() + source_variants.clear() + assert tuple(field.name for field in layout.fields) == ("name",) + assert layout.parameter_variants[0].fields == ("name",) + assert layout.parameter_variants[0].layouts == frozenset( + {ParameterLayout.REPLICATED} + ) + + grad = torch.arange(16, dtype=torch.float32).reshape(4, 4) / 16 + param.grad = grad.clone() + reference_param.grad = grad.clone() + optimizer.step() + reference.step() + assert torch.equal(param, reference_param) + assert _values_equal(optimizer.state_dict(), reference.state_dict()) + + +def test_contract_rejects_duplicate_variant_names(): + field = StateField("name", StateScope.PARAMETER, StateGeometry.OPAQUE, True) + variant = StateVariant( + "name_only", + ("name",), + frozenset({ParameterLayout.REPLICATED}), + StateExtent.METADATA_ONLY, + initialized=False, + ) + with pytest.raises(ValueError, match="names must be unique"): + OptimizerStateLayout((field,), (variant, variant)) + + +def test_derived_fields_are_explicitly_non_authoritative(): + optimizer = Gefen( + [("layer.weight", torch.nn.Parameter(torch.ones(4, 4)))], + fused=False, + ) + derived = optimizer.optimizer_contract().state_layout.fields_for_scope( + StateScope.DERIVED + ) + assert derived + assert all(not field.authoritative for field in derived) + assert {field.name for field in derived if field.checkpointed} == { + "m_codebook_shape", + "_gefen_rank_local_payload_", + "_gefen_rank_local_member", + "_gefen_checkpoint_metadata", + } + payload_field = optimizer.optimizer_contract().state_layout.field( + "_gefen_rank_local_payload_3" + ) + assert payload_field.key_match is StateKeyMatch.PREFIX + + +def test_all_public_contract_exports_resolve(): + from gefen import contracts + + assert all(getattr(gefen, name) is not None for name in contracts.__all__) From 4e2d967ef0e05ff8b2edd8acfb01093c0fd5506d Mon Sep 17 00:00:00 2001 From: thad0ctor Date: Sun, 12 Jul 2026 17:56:54 -0700 Subject: [PATCH 02/52] Make native optimizer loads atomic --- docs/optimizer_contracts.md | 2 + src/gefen/contracts.py | 10 +- src/gefen/gefen.py | 140 +++++++++- tests/test_native_load_atomicity.py | 403 ++++++++++++++++++++++++++++ tests/test_optimizer_contracts.py | 10 +- 5 files changed, 556 insertions(+), 9 deletions(-) create mode 100644 tests/test_native_load_atomicity.py diff --git a/docs/optimizer_contracts.md b/docs/optimizer_contracts.md index 0e3de25..f1ca346 100644 --- a/docs/optimizer_contracts.md +++ b/docs/optimizer_contracts.md @@ -29,6 +29,8 @@ The current DTensor declaration is deliberately narrow: `DTENSOR_1D_DEFAULT_WORL Plain Gefen declares replicated, flattened element-shard, and the narrow DTensor training layouts. Its PyTorch rank-local checkpoint transport is same-topology only. GefenMuon declares replicated and narrow DTensor training, with mode-specific state extents: `approx` state is local, `exact` state is logically global, and `distributed` momentum is held by the parameter owner while non-owners retain metadata only. Native Parallel-Muon checkpoints separately declare world-size owner redistribution, not placement-changing resharding. `GefenMuonHybrid` retains its nested child namespaces and does not flatten AdamW or Gefen child state into a fabricated common schema. +Gefen and GefenMuon native loads, rank-local payload restoration through `load_state_dict`, and distributed-owner payload restoration prepare their complete core restore before changing local live optimizer state. Their transport entries report `atomic_load=True`. This is a per-optimizer-instance fail-before-mutation guarantee at the optimizer load boundary, not a coordinated all-rank commit or a guarantee over work an external checkpoint orchestrator performs before calling the optimizer. Load pre-hooks run before that boundary and load post-hooks run afterward, so arbitrary side effects in user hooks are also outside the guarantee. Hybrid composite loads do not yet provide the same guarantee and report `atomic_load=False`. + ## Adapter requirements An adapter should match the exact `TrainingSupport` or `CheckpointSupport` entry it intends to use, including transport, layout, process-group scope, mesh dimensions, and sharded mode. It should not treat successful training as checkpoint support, same-topology checkpointing as resharding support, or accepted caller names as canonical fully qualified parameter identity. diff --git a/src/gefen/contracts.py b/src/gefen/contracts.py index 3149835..0b1b6a3 100644 --- a/src/gefen/contracts.py +++ b/src/gefen/contracts.py @@ -291,7 +291,13 @@ def __post_init__(self) -> None: @dataclass(frozen=True) class CheckpointSupport: - """One transport's separately qualified checkpoint capability.""" + """One transport's separately qualified checkpoint capability. + + ``atomic_load`` means each participating optimizer instance validates and + prepares its core restore before local mutation. It does not claim a + coordinated all-rank commit after failures outside the declared process + group or arbitrary user hook side effects. + """ transport: CheckpointTransport same_topology: AbstractSet[ParameterLayout] @@ -712,6 +718,7 @@ def _gefen_contract(*, factored_v_2d: bool) -> OptimizerContract: ), frozenset(), ProcessGroupScope.NONE, + atomic_load=True, ), CheckpointSupport( CheckpointTransport.PYTORCH_RANK_LOCAL, @@ -922,6 +929,7 @@ def _gefen_muon_contract( frozenset(), ProcessGroupScope.NONE, required_sharded_modes=sharded_modes, + atomic_load=True, ) ] if ( diff --git a/src/gefen/gefen.py b/src/gefen/gefen.py index c6e5251..eecb73f 100644 --- a/src/gefen/gefen.py +++ b/src/gefen/gefen.py @@ -4464,17 +4464,84 @@ def _pack_legacy_param_groups_for_load(self, state_dict): return migrated def load_state_dict(self, state_dict): - """Run load hooks around Gefen's complete restore transaction.""" + """Atomically restore Gefen state between the public load hooks.""" 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._load_state_dict_impl(state_dict) + staged = self._stage_load_state_dict(state_dict) + self._commit_staged_load_state_dict(staged) for post_hook in self._optimizer_load_state_dict_post_hooks.values(): post_hook(self) + def _stage_load_state_dict(self, state_dict): + """Prepare a complete restore without mutating the live optimizer. + + ``Optimizer.load_state_dict`` commits before Gefen restores auxiliary + dtypes, normalizes counters, rebuilds names, and prepares checkpoint + schema. Run that complete dynamically dispatched path on an isolated + shadow first so a failure in any of those operations is a true no-op. + ``copy.copy`` cannot be used here because ``Optimizer.__getstate__`` + intentionally omits Gefen's private runtime attributes. + """ + + staged = object.__new__(type(self)) + staged.__dict__ = self.__dict__.copy() + + # ``Optimizer.__setstate__`` mutates defaults, while the remaining + # containers are cleared or invalidated by Gefen before/after the base + # load. Isolate them so even a late staging failure cannot disturb live + # caches, capturable views, or device counters. + staged.defaults = self.defaults.copy() + staged._gefen_codebook_by_device = {} + staged._gefen_codebook_lut_by_device = {} + staged._sr_seed_by_device = {} + staged._gefen_global_step_by_device = {} + staged._capt_stacks = None + + staged._load_state_dict_impl(state_dict) + staged._validate_loaded_native_state() + return staged + + def _commit_staged_load_state_dict(self, staged) -> None: + """Publish an already prepared restore through non-throwing swaps.""" + + # The base loader mutates the existing defaults mapping via setdefault; + # preserve that public object identity while publishing the staged + # value. Both mappings are ordinary built-in dicts created by Optimizer. + live_defaults = self.defaults + live_defaults.update(staged.defaults) + staged.defaults = live_defaults + self.__dict__.update(staged.__dict__) + + def _validate_loaded_native_state(self) -> None: + """Validate the complete prepared native state before publication.""" + + self._validate_rank_local_counter( + "gefen_global_step", self._gefen_global_step + ) + signature = self._rank_local_sharded_signature(context={}) + states = [ + self.state[param] + for group in self.param_groups + for param in group["params"] + ] + self._validate_rank_local_states( + states, + signature, + self._gefen_codebook, + allow_legacy_vmean_counter=True, + ) + for group in self.param_groups: + self._validate_group_options( + group["lr"], + (group["beta1"], group["beta2"]), + group["eps"], + group["weight_decay"], + ) + def _base_load_state_dict_without_hooks(self, state_dict): pre_hooks = self._optimizer_load_state_dict_pre_hooks post_hooks = self._optimizer_load_state_dict_post_hooks @@ -4543,7 +4610,14 @@ def _validate_rank_local_counter(name, value) -> float: ) return scalar - def _validate_rank_local_states(self, states, signature, codebook) -> None: + def _validate_rank_local_states( + self, + states, + signature, + codebook, + *, + allow_legacy_vmean_counter: bool = False, + ) -> None: if not isinstance(states, list) or len(states) != len(signature): raise ValueError( "Gefen rank-local checkpoint payload has the wrong parameter count" @@ -4591,6 +4665,33 @@ def _validate_rank_local_states(self, states, signature, codebook) -> None: momentum_keys = ("m_codebook", "m_magnitude") carries_momentum = any(key in pstate for key in momentum_keys) + initialized_keys = ( + "step", + "vmean", + "vmean_step", + "v_row", + "v_col", + "factored_step", + "normuon_v", + "normuon_step", + ) + if any(key in pstate for key in initialized_keys) and not carries_momentum: + raise ValueError( + "Gefen rank-local checkpoint initialized state is missing " + "quantized momentum" + ) + if ( + period is not None + and not carries_momentum + and not ( + param_signature.get("sharded") + and param_signature.get("sharded_mode") == "distributed" + ) + ): + raise ValueError( + "Gefen rank-local checkpoint automatic_period is invalid without " + "initialized momentum" + ) if carries_momentum: has_quantized_momentum = True if not all(key in pstate for key in momentum_keys) or period is None: @@ -4640,6 +4741,22 @@ def _validate_rank_local_states(self, states, signature, codebook) -> None: "are invalid" ) + has_vmean = "vmean" in pstate + has_vmean_step = "vmean_step" in pstate + if has_vmean_step and not has_vmean: + raise ValueError( + "Gefen rank-local checkpoint block second moment is incomplete" + ) + if ( + has_vmean + and not has_vmean_step + and not allow_legacy_vmean_counter + ): + raise ValueError( + "Gefen rank-local checkpoint block second moment is missing " + "vmean_step" + ) + # GefenMuon groups carry a sharded_mode and intentionally use # quantized momentum without Adam's second moment. Plain Gefen # must carry one complete second-moment representation. @@ -4656,12 +4773,15 @@ def _validate_rank_local_states(self, states, signature, codebook) -> None: "Gefen rank-local checkpoint initialized factored " "state requires factored_step >= 1" ) - elif "vmean" not in pstate or "vmean_step" not in pstate: + elif "vmean" not in pstate: raise ValueError( - "Gefen rank-local checkpoint plain momentum is missing " - "vmean/vmean_step" + "Gefen rank-local checkpoint plain momentum is missing a " + "second moment" ) - elif counters["vmean_step"] < 1: + elif ( + "vmean_step" in counters + and counters["vmean_step"] < 1 + ): raise ValueError( "Gefen rank-local checkpoint initialized vmean requires " "vmean_step >= 1" @@ -4672,6 +4792,11 @@ def _validate_rank_local_states(self, states, signature, codebook) -> None: raise ValueError( "Gefen rank-local checkpoint factored second moment is incomplete" ) + has_factored_step = "factored_step" in pstate + if has_factored_step != (factored[0] is not None): + raise ValueError( + "Gefen rank-local checkpoint factored second moment is incomplete" + ) if factored[0] is not None: global_shape = param_signature["shape"] if len(global_shape) != 2: @@ -5214,6 +5339,7 @@ def _load_state_dict_impl(self, state_dict): if key not in pstate: continue value = pstate[key] + self._validate_rank_local_counter(key, value) if self.capturable: if not torch.is_tensor(value): pstate[key] = torch.full( diff --git a/tests/test_native_load_atomicity.py b/tests/test_native_load_atomicity.py new file mode 100644 index 0000000..2041eea --- /dev/null +++ b/tests/test_native_load_atomicity.py @@ -0,0 +1,403 @@ +"""Fail-before-mutation coverage for native Gefen checkpoint restores.""" + +import copy + +import pytest +import torch + +from gefen import Gefen, GefenMuon + + +def _build_optimizer(kind, seed): + generator = torch.Generator().manual_seed(seed) + param = torch.nn.Parameter(torch.randn(8, 8, generator=generator)) + named = [("layer.weight", param)] + if kind == "block": + optimizer = Gefen(named, lr=1e-3, fused=False, factored_v_2d=False) + elif kind == "factored": + optimizer = Gefen(named, lr=1e-3, fused=False, factored_v_2d=True) + elif kind == "muon": + optimizer = GefenMuon(named, lr=1e-3, fused=False, normuon=False) + elif kind == "normuon": + optimizer = GefenMuon(named, lr=1e-3, fused=False, normuon=True) + else: + raise AssertionError("unknown optimizer kind: {}".format(kind)) + param.grad = torch.randn(param.shape, generator=generator) + optimizer.step() + optimizer.zero_grad() + return optimizer, param + + +def _seed_live_runtime_state(optimizer, param): + device = torch.device("cpu") + pstate = optimizer.state[param] + pstate["_capt_stack"] = torch.tensor([11.0]) + pstate["_capt_row"] = 3 + pstate["_capt_scalars"] = torch.tensor([12.0]) + optimizer._capt_stacks = {device: ({"rows": [param]},)} + optimizer._gefen_codebook_by_device[device] = torch.tensor([13.0]) + optimizer._gefen_codebook_lut_by_device[device] = torch.tensor([14.0]) + optimizer._sr_seed_by_device[device] = torch.tensor(15, dtype=torch.int64) + optimizer._gefen_global_step_by_device[device] = torch.tensor(16, dtype=torch.int64) + optimizer._static_mark_sig = ("sentinel", 17) + + +def _snapshot_live_optimizer(optimizer): + state_refs = dict(optimizer.state) + state_values = {} + state_tensor_refs = {} + for param, pstate in state_refs.items(): + state_values[param] = {} + state_tensor_refs[param] = {} + for key, value in pstate.items(): + if torch.is_tensor(value): + state_values[param][key] = value.detach().clone() + state_tensor_refs[param][key] = (value, value._version) + else: + state_values[param][key] = copy.deepcopy(value) + + cache_attrs = ( + "_gefen_codebook_by_device", + "_gefen_codebook_lut_by_device", + "_sr_seed_by_device", + "_gefen_global_step_by_device", + ) + caches = {} + for attr in cache_attrs: + cache = getattr(optimizer, attr) + caches[attr] = ( + cache, + { + key: (value, value.detach().clone(), value._version) + for key, value in cache.items() + }, + ) + return { + "state": optimizer.state, + "state_refs": state_refs, + "state_values": state_values, + "state_tensor_refs": state_tensor_refs, + "param_groups": optimizer.param_groups, + "group_refs": tuple(optimizer.param_groups), + "group_param_refs": tuple(group["params"] for group in optimizer.param_groups), + "group_values": tuple( + { + key: copy.deepcopy(value) + for key, value in group.items() + if key != "params" + } + for group in optimizer.param_groups + ), + "defaults": optimizer.defaults, + "defaults_value": copy.deepcopy(optimizer.defaults), + "param_names": optimizer._param_names, + "param_names_value": dict(optimizer._param_names), + "global_step": optimizer._gefen_global_step, + "codebook": optimizer._gefen_codebook, + "capt_stacks": optimizer._capt_stacks, + "static_mark_sig": optimizer._static_mark_sig, + "caches": caches, + } + + +def _assert_live_optimizer_unchanged(optimizer, snapshot): + assert optimizer.state is snapshot["state"] + assert optimizer.param_groups is snapshot["param_groups"] + assert optimizer.defaults is snapshot["defaults"] + assert optimizer.defaults == snapshot["defaults_value"] + assert optimizer._param_names is snapshot["param_names"] + assert optimizer._param_names == snapshot["param_names_value"] + assert optimizer._gefen_global_step is snapshot["global_step"] + assert optimizer._gefen_codebook is snapshot["codebook"] + assert optimizer._capt_stacks is snapshot["capt_stacks"] + assert optimizer._static_mark_sig is snapshot["static_mark_sig"] + + assert tuple(optimizer.param_groups) == snapshot["group_refs"] + for group, group_ref, params_ref, values in zip( + optimizer.param_groups, + snapshot["group_refs"], + snapshot["group_param_refs"], + snapshot["group_values"], + ): + assert group is group_ref + assert group["params"] is params_ref + assert {key: value for key, value in group.items() if key != "params"} == values + + assert set(optimizer.state) == set(snapshot["state_refs"]) + for param, pstate_ref in snapshot["state_refs"].items(): + pstate = optimizer.state[param] + assert pstate is pstate_ref + assert set(pstate) == set(snapshot["state_values"][param]) + for key, expected in snapshot["state_values"][param].items(): + value = pstate[key] + if torch.is_tensor(value): + tensor_ref, version = snapshot["state_tensor_refs"][param][key] + assert value is tensor_ref + assert value._version == version + assert torch.equal(value, expected) + else: + assert value == expected + + for attr, (cache_ref, entries) in snapshot["caches"].items(): + cache = getattr(optimizer, attr) + assert cache is cache_ref + assert set(cache) == set(entries) + for key, (tensor_ref, expected, version) in entries.items(): + assert cache[key] is tensor_ref + assert cache[key]._version == version + assert torch.equal(cache[key], expected) + + +def _assert_nested_equal(actual, expected): + assert type(actual) is type(expected) + if isinstance(actual, dict): + assert set(actual) == set(expected) + for key in actual: + _assert_nested_equal(actual[key], expected[key]) + elif isinstance(actual, (list, tuple)): + assert len(actual) == len(expected) + for left, right in zip(actual, expected): + _assert_nested_equal(left, right) + elif torch.is_tensor(actual): + assert actual.dtype == expected.dtype + torch.testing.assert_close(actual, expected, rtol=0, atol=0, equal_nan=True) + else: + assert actual == expected + + +def _set_all_codebook_copies(checkpoint, codebook): + checkpoint["gefen_codebook"] = codebook + for group in checkpoint["param_groups"]: + metadata = group.get("_gefen_checkpoint_metadata") + if metadata is not None: + metadata["codebook"] = codebook + + +def _corrupt_automatic_period(checkpoint): + next(iter(checkpoint["state"].values()))["automatic_period"] = 0 + + +def _corrupt_momentum_indices(checkpoint): + pstate = next(iter(checkpoint["state"].values())) + pstate["m_codebook"] = pstate["m_codebook"].to(torch.float32) + + +def _corrupt_momentum_magnitude(checkpoint): + pstate = next(iter(checkpoint["state"].values())) + pstate["m_magnitude"] = torch.full_like(pstate["m_magnitude"], -1.0) + + +def _remove_quantized_momentum(checkpoint): + pstate = next(iter(checkpoint["state"].values())) + pstate.pop("m_codebook") + pstate.pop("m_magnitude") + + +def _corrupt_block_second_moment(checkpoint): + pstate = next(iter(checkpoint["state"].values())) + pstate["vmean"] = torch.full_like(pstate["vmean"], float("nan")) + + +def _corrupt_factored_second_moment(checkpoint): + next(iter(checkpoint["state"].values())).pop("v_col") + + +def _corrupt_normuon_second_moment(checkpoint): + pstate = next(iter(checkpoint["state"].values())) + pstate["normuon_v"] = torch.full_like(pstate["normuon_v"], -1.0) + + +def _corrupt_codebook(checkpoint): + codebook = torch.linspace(-1.0, 1.0, 255, dtype=torch.float32) + _set_all_codebook_copies(checkpoint, codebook) + + +@pytest.mark.parametrize( + ("kind", "counter_key", "invalid_value"), + [ + ("block", "step", torch.tensor([1.0, 2.0])), + ("block", "step", torch.tensor(1.5)), + ("block", "step", torch.tensor(True)), + ("block", "step", torch.tensor(-1.0)), + ("block", "step", torch.tensor(float("nan"))), + ("block", "vmean_step", torch.tensor([1.0, 2.0])), + ("factored", "factored_step", torch.tensor([1.0, 2.0])), + ("muon", "step", torch.tensor([1.0, 2.0])), + ("normuon", "normuon_step", torch.tensor([1.0, 2.0])), + ], +) +def test_native_counter_normalization_failure_is_fail_before_mutation( + kind, counter_key, invalid_value +): + source, _ = _build_optimizer(kind, seed=1) + target, target_param = _build_optimizer(kind, seed=2) + checkpoint = copy.deepcopy(source.state_dict()) + next(iter(checkpoint["state"].values()))[counter_key] = invalid_value + checkpoint_before = copy.deepcopy(checkpoint) + _seed_live_runtime_state(target, target_param) + snapshot = _snapshot_live_optimizer(target) + post_calls = [] + target.register_load_state_dict_post_hook(lambda optimizer: post_calls.append(optimizer)) + + with pytest.raises((RuntimeError, ValueError)): + target.load_state_dict(checkpoint) + + assert post_calls == [] + _assert_live_optimizer_unchanged(target, snapshot) + _assert_nested_equal(checkpoint, checkpoint_before) + + +@pytest.mark.parametrize("kind", ["block", "muon"]) +def test_native_base_layout_failure_preserves_live_caches_and_objects(kind): + source, _ = _build_optimizer(kind, seed=3) + target, target_param = _build_optimizer(kind, seed=4) + checkpoint = copy.deepcopy(source.state_dict()) + checkpoint["param_groups"].append(copy.deepcopy(checkpoint["param_groups"][0])) + checkpoint_before = copy.deepcopy(checkpoint) + _seed_live_runtime_state(target, target_param) + snapshot = _snapshot_live_optimizer(target) + + with pytest.raises(ValueError, match="different number of parameter groups"): + target.load_state_dict(checkpoint) + + _assert_live_optimizer_unchanged(target, snapshot) + _assert_nested_equal(checkpoint, checkpoint_before) + + +@pytest.mark.parametrize( + ("kind", "corrupt"), + [ + ("block", _corrupt_automatic_period), + ("block", _corrupt_momentum_indices), + ("block", _corrupt_momentum_magnitude), + ("block", _remove_quantized_momentum), + ("block", _corrupt_block_second_moment), + ("block", _corrupt_codebook), + ("factored", _corrupt_factored_second_moment), + ("normuon", _corrupt_normuon_second_moment), + ], + ids=( + "automatic-period", + "momentum-indices", + "momentum-magnitude", + "missing-momentum", + "block-second-moment", + "codebook", + "factored-second-moment", + "normuon-second-moment", + ), +) +def test_invalid_native_state_is_rejected_before_live_mutation(kind, corrupt): + source, _ = _build_optimizer(kind, seed=9) + target, target_param = _build_optimizer(kind, seed=10) + checkpoint = copy.deepcopy(source.state_dict()) + corrupt(checkpoint) + checkpoint_before = copy.deepcopy(checkpoint) + _seed_live_runtime_state(target, target_param) + snapshot = _snapshot_live_optimizer(target) + + with pytest.raises(ValueError): + target.load_state_dict(checkpoint) + + _assert_live_optimizer_unchanged(target, snapshot) + _assert_nested_equal(checkpoint, checkpoint_before) + + +def test_native_load_hooks_run_once_around_transaction(): + source, _ = _build_optimizer("block", seed=5) + target, _ = _build_optimizer("block", seed=6) + checkpoint = copy.deepcopy(source.state_dict()) + calls = [] + + def replace_with_invalid_counter(optimizer, state_dict): + calls.append(("pre", optimizer)) + replacement = copy.deepcopy(state_dict) + next(iter(replacement["state"].values()))["step"] = torch.tensor([1.0, 2.0]) + return replacement + + target.register_load_state_dict_pre_hook(replace_with_invalid_counter) + target.register_load_state_dict_post_hook( + lambda optimizer: calls.append(("post", optimizer)) + ) + snapshot = _snapshot_live_optimizer(target) + + with pytest.raises((RuntimeError, ValueError)): + target.load_state_dict(checkpoint) + + assert calls == [("pre", target)] + _assert_live_optimizer_unchanged(target, snapshot) + + +def test_successful_native_load_hooks_run_once_and_commit(): + source, source_param = _build_optimizer("factored", seed=7) + source_param.grad = torch.full_like(source_param, 0.25) + source.step() + source.zero_grad() + target, target_param = _build_optimizer("factored", seed=8) + checkpoint = copy.deepcopy(source.state_dict()) + calls = [] + target.register_load_state_dict_pre_hook( + lambda optimizer, state_dict: calls.append(("pre", optimizer)) + ) + target.register_load_state_dict_post_hook( + lambda optimizer: calls.append(("post", optimizer)) + ) + old_state = target.state + old_defaults = target.defaults + + target.load_state_dict(checkpoint) + + assert calls == [("pre", target), ("post", target)] + assert target.state is not old_state + assert target.defaults is old_defaults + assert target._gefen_global_step == source._gefen_global_step == 2 + assert target.state[target_param]["step"] == source.state[source_param]["step"] == 2 + assert torch.equal(target._gefen_codebook, source._gefen_codebook) + + +def test_rejected_late_parameter_corruption_preserves_bit_exact_continuation(): + def build(seed): + generator = torch.Generator().manual_seed(seed) + params = [ + torch.nn.Parameter(torch.randn(4, 8, generator=generator)), + torch.nn.Parameter(torch.randn(8, 4, generator=generator)), + ] + optimizer = Gefen( + [("first.weight", params[0]), ("second.weight", params[1])], + lr=1e-3, + fused=False, + factored_v_2d=False, + ) + for param in params: + param.grad = torch.randn(param.shape, generator=generator) + optimizer.step() + optimizer.zero_grad() + return optimizer, params + + source, _ = build(11) + target, target_params = build(12) + control, control_params = build(12) + checkpoint = copy.deepcopy(source.state_dict()) + list(checkpoint["state"].values())[-1]["step"] = torch.tensor([1.0, 2.0]) + snapshot = _snapshot_live_optimizer(target) + + with pytest.raises(ValueError): + target.load_state_dict(checkpoint) + + _assert_live_optimizer_unchanged(target, snapshot) + generator = torch.Generator().manual_seed(13) + next_grads = [ + torch.randn(param.shape, generator=generator) for param in target_params + ] + for params, optimizer in ( + (target_params, target), + (control_params, control), + ): + for param, grad in zip(params, next_grads): + param.grad = grad.clone() + optimizer.step() + optimizer.zero_grad() + + for target_param, control_param in zip(target_params, control_params): + assert torch.equal(target_param, control_param) + _assert_nested_equal(target.state_dict(), control.state_dict()) diff --git a/tests/test_optimizer_contracts.py b/tests/test_optimizer_contracts.py index 5b3d81c..206f9e4 100644 --- a/tests/test_optimizer_contracts.py +++ b/tests/test_optimizer_contracts.py @@ -157,7 +157,7 @@ def test_plain_contract_matches_live_persistent_state(factored_v_2d): for item in contract.capabilities.checkpoints if item.transport is CheckpointTransport.NATIVE_OPTIMIZER ) - assert not native.atomic_load + assert native.atomic_load assert contract.capabilities.accepts_semantic_parameter_names assert not contract.capabilities.canonical_parameter_fqns assert not contract.capabilities.stable_shard_identity @@ -337,6 +337,13 @@ def test_muon_contract_separates_mode_topology_and_state_extent( assert contract.implementation == "gefen.GefenMuon" assert contract.capabilities.supported_parameter_ranks == (2,) + native = next( + item + for item in contract.capabilities.checkpoints + if item.transport is CheckpointTransport.NATIVE_OPTIMIZER + and ParameterLayout.REPLICATED in item.same_topology + ) + assert native.atomic_load replicated = _training_support(contract, ParameterLayout.REPLICATED) assert replicated.requires_complete_logical_matrix assert not replicated.requires_complete_parameter_storage @@ -469,6 +476,7 @@ def test_hybrid_contract_preserves_child_namespaces(backup_optimizer): assert checkpoint[0].transport is CheckpointTransport.COMPOSITE_NATIVE assert checkpoint[0].same_topology == frozenset({ParameterLayout.REPLICATED}) assert not checkpoint[0].topology_changing + assert not checkpoint[0].atomic_load def test_muon_contract_keeps_mixed_normuon_variants_in_one_mode(): From 102b7f2da520c7b526b2d0a6606ada15514c11c3 Mon Sep 17 00:00:00 2001 From: thad0ctor Date: Sun, 12 Jul 2026 18:23:52 -0700 Subject: [PATCH 03/52] Add canonical shard identity descriptors --- docs/optimizer_contracts.md | 6 + src/gefen/__init__.py | 16 + src/gefen/contracts.py | 341 ++++++++++++++++++ tests/test_shard_identity_contracts.py | 472 +++++++++++++++++++++++++ 4 files changed, 835 insertions(+) create mode 100644 tests/test_shard_identity_contracts.py diff --git a/docs/optimizer_contracts.md b/docs/optimizer_contracts.md index f1ca346..6e2b7ee 100644 --- a/docs/optimizer_contracts.md +++ b/docs/optimizer_contracts.md @@ -27,6 +27,12 @@ assert ParameterLayout.DTENSOR_1D_DEFAULT_WORLD in rank_local_dcp.same_topology The current DTensor declaration is deliberately narrow: `DTENSOR_1D_DEFAULT_WORLD` means one shared one-dimensional mesh spanning the default world. Multidimensional meshes, subgroups, and placement-changing loads are not implied by that declaration. +## Canonical parameter and shard identity + +`ParameterIdentity` records an exact, case-preserving model FQN and global logical shape independently of any live tensor object. `ProcessGroupIdentity` records an adapter-defined semantic group name and authoritative ordered member IDs without importing a framework process-group type. `ShardIdentity` combines those values with a contiguous row-major logical range, an explicit parameter layout, structured placements, the local member, and an optional whole-parameter owner. `ShardingManifest` validates and deterministically orders the complete identity set; flattened manifests 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. The contiguous-range schema deliberately rejects DTensor identities because column and multidimensional shards require a richer logical-region descriptor; the existing narrow DTensor training and rank-local checkpoint paths remain independently declared. + +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. + Plain Gefen declares replicated, flattened element-shard, and the narrow DTensor training layouts. Its PyTorch rank-local checkpoint transport is same-topology only. GefenMuon declares replicated and narrow DTensor training, with mode-specific state extents: `approx` state is local, `exact` state is logically global, and `distributed` momentum is held by the parameter owner while non-owners retain metadata only. Native Parallel-Muon checkpoints separately declare world-size owner redistribution, not placement-changing resharding. `GefenMuonHybrid` retains its nested child namespaces and does not flatten AdamW or Gefen child state into a fabricated common schema. Gefen and GefenMuon native loads, rank-local payload restoration through `load_state_dict`, and distributed-owner payload restoration prepare their complete core restore before changing local live optimizer state. Their transport entries report `atomic_load=True`. This is a per-optimizer-instance fail-before-mutation guarantee at the optimizer load boundary, not a coordinated all-rank commit or a guarantee over work an external checkpoint orchestrator performs before calling the optimizer. Load pre-hooks run before that boundary and load post-hooks run afterward, so arbitrary side effects in user hooks are also outside the guarantee. Hybrid composite loads do not yet provide the same guarantee and report `atomic_load=False`. diff --git a/src/gefen/__init__.py b/src/gefen/__init__.py index c427b95..416a767 100644 --- a/src/gefen/__init__.py +++ b/src/gefen/__init__.py @@ -12,6 +12,7 @@ "GefenMuon", "GefenMuonHybrid", "CONTRACT_SCHEMA_VERSION", + "IDENTITY_SCHEMA_VERSION", "CheckpointSupport", "CheckpointTransport", "OptimizerCapabilities", @@ -19,10 +20,17 @@ "OptimizerContract", "OptimizerContractProvider", "OptimizerStateLayout", + "LogicalSlice", "ParameterLayout", + "ParameterIdentity", "ParameterStateRole", + "PlacementKind", "Precision", "ProcessGroupScope", + "ProcessGroupIdentity", + "ShardIdentity", + "ShardPlacement", + "ShardingManifest", "StateExtent", "StateField", "StateGeometry", @@ -57,6 +65,7 @@ def __getattr__(name): return getattr(params, name) if name in ( "CONTRACT_SCHEMA_VERSION", + "IDENTITY_SCHEMA_VERSION", "CheckpointSupport", "CheckpointTransport", "OptimizerCapabilities", @@ -64,10 +73,17 @@ def __getattr__(name): "OptimizerContract", "OptimizerContractProvider", "OptimizerStateLayout", + "LogicalSlice", "ParameterLayout", + "ParameterIdentity", "ParameterStateRole", + "PlacementKind", "Precision", "ProcessGroupScope", + "ProcessGroupIdentity", + "ShardIdentity", + "ShardPlacement", + "ShardingManifest", "StateExtent", "StateField", "StateGeometry", diff --git a/src/gefen/contracts.py b/src/gefen/contracts.py index 0b1b6a3..f861597 100644 --- a/src/gefen/contracts.py +++ b/src/gefen/contracts.py @@ -7,6 +7,7 @@ from dataclasses import dataclass from enum import Enum +import math from typing import ( AbstractSet, FrozenSet, @@ -19,6 +20,7 @@ CONTRACT_SCHEMA_VERSION = 1 +IDENTITY_SCHEMA_VERSION = 1 class StateScope(str, Enum): @@ -74,6 +76,15 @@ class ParameterLayout(str, Enum): DTENSOR_1D_DEFAULT_WORLD = "dtensor_1d_default_world" +class PlacementKind(str, Enum): + """Framework-neutral placement carried by a stable shard identity.""" + + REPLICATE = "replicate" + FLAT_SHARD = "flat_shard" + DIMENSION_SHARD = "dimension_shard" + WHOLE_PARAMETER_OWNER = "whole_parameter_owner" + + class ProcessGroupScope(str, Enum): """How an implemented path obtains its collective process group.""" @@ -124,6 +135,328 @@ def _validate_dimensions(name, values, *, positive): raise ValueError("{} must contain {} integers".format(name, relation)) +def _validate_identity_name(name, value): + if not isinstance(value, str) or not value or value != value.strip(): + raise ValueError("{} must be a non-empty string without outer whitespace".format(name)) + if "\x00" in value: + raise ValueError("{} must not contain NUL".format(name)) + + +def _validate_identity_schema_version(name, value): + if type(value) is not int or value != IDENTITY_SCHEMA_VERSION: + raise ValueError("unsupported {} schema version".format(name)) + + +@dataclass(frozen=True) +class ParameterIdentity: + """Canonical logical parameter identity, independent of tensor storage.""" + + fqn: str + global_shape: Sequence[int] + schema_version: int = IDENTITY_SCHEMA_VERSION + + def __post_init__(self) -> None: + _validate_identity_name("ParameterIdentity.fqn", self.fqn) + if self.fqn.startswith(".") or self.fqn.endswith(".") or ".." in self.fqn: + raise ValueError("ParameterIdentity.fqn must contain non-empty dot-separated components") + if isinstance(self.global_shape, (str, bytes, bytearray)): + raise TypeError("ParameterIdentity.global_shape must be a sequence of dimensions") + object.__setattr__(self, "global_shape", _tuple(self.global_shape)) + if any(type(dim) is not int or dim < 0 for dim in self.global_shape): + raise ValueError("ParameterIdentity.global_shape must contain nonnegative integers") + _validate_identity_schema_version("parameter identity", self.schema_version) + + @property + def numel(self) -> int: + """Return the canonical logical element count.""" + + return math.prod(self.global_shape) + + +@dataclass(frozen=True) +class ProcessGroupIdentity: + """Stable semantic process-group identity supplied by an adapter.""" + + semantic_name: str + ordered_members: Sequence[str] + schema_version: int = IDENTITY_SCHEMA_VERSION + + def __post_init__(self) -> None: + _validate_identity_name("ProcessGroupIdentity.semantic_name", self.semantic_name) + if isinstance(self.ordered_members, (str, bytes)): + raise TypeError("ProcessGroupIdentity.ordered_members must be a sequence of member IDs") + object.__setattr__(self, "ordered_members", _tuple(self.ordered_members)) + if not self.ordered_members: + raise ValueError("ProcessGroupIdentity.ordered_members must be non-empty") + for member in self.ordered_members: + _validate_identity_name("process-group member", member) + if len(set(self.ordered_members)) != len(self.ordered_members): + raise ValueError("ProcessGroupIdentity.ordered_members must be unique") + _validate_identity_schema_version("process-group identity", self.schema_version) + + +@dataclass(frozen=True) +class ShardPlacement: + """One explicit mesh-axis placement for a logical parameter shard.""" + + mesh_axis: str + kind: PlacementKind + coordinate: int + parts: int + parameter_dimension: Optional[int] = None + + def __post_init__(self) -> None: + _validate_identity_name("ShardPlacement.mesh_axis", self.mesh_axis) + if not isinstance(self.kind, PlacementKind): + raise TypeError("ShardPlacement.kind must be a PlacementKind") + if type(self.parts) is not int or self.parts <= 0: + raise ValueError("ShardPlacement.parts must be a positive integer") + if type(self.coordinate) is not int or self.coordinate < 0 or self.coordinate >= self.parts: + raise ValueError("ShardPlacement.coordinate must be within parts") + if self.kind is PlacementKind.DIMENSION_SHARD: + if type(self.parameter_dimension) is not int or self.parameter_dimension < 0: + raise ValueError("dimension-shard placement requires a nonnegative parameter dimension") + elif self.parameter_dimension is not None: + raise ValueError("only a dimension-shard placement may name a parameter dimension") + + +@dataclass(frozen=True) +class LogicalSlice: + """Contiguous range in canonical row-major flattened parameter order.""" + + flat_offset: int + length: int + + def __post_init__(self) -> None: + if type(self.flat_offset) is not int or self.flat_offset < 0: + raise ValueError("LogicalSlice.flat_offset must be a nonnegative integer") + if type(self.length) is not int or self.length < 0: + raise ValueError("LogicalSlice.length must be a nonnegative integer") + + @classmethod + def full(cls, parameter: ParameterIdentity) -> "LogicalSlice": + """Return the complete logical range for ``parameter``.""" + + if not isinstance(parameter, ParameterIdentity): + raise TypeError("parameter must be a ParameterIdentity") + return cls(0, parameter.numel) + + +@dataclass(frozen=True) +class ShardIdentity: + """Stable identity of one process-group member's logical parameter shard.""" + + parameter: ParameterIdentity + layout: ParameterLayout + logical_slice: LogicalSlice + placements: Sequence[ShardPlacement] = () + process_group: Optional[ProcessGroupIdentity] = None + local_member: Optional[str] = None + owner: Optional[str] = None + schema_version: int = IDENTITY_SCHEMA_VERSION + + def __post_init__(self) -> None: + if not isinstance(self.parameter, ParameterIdentity): + raise TypeError("ShardIdentity.parameter must be a ParameterIdentity") + if not isinstance(self.layout, ParameterLayout): + raise TypeError("ShardIdentity.layout must be a ParameterLayout") + if not isinstance(self.logical_slice, LogicalSlice): + raise TypeError("ShardIdentity.logical_slice must be a LogicalSlice") + placements = _tuple(self.placements) + if any(not isinstance(item, ShardPlacement) for item in placements): + raise TypeError("ShardIdentity.placements must contain ShardPlacement values") + axes = tuple(item.mesh_axis for item in placements) + if len(set(axes)) != len(axes): + raise ValueError("ShardIdentity placement mesh axes must be unique") + object.__setattr__( + self, + "placements", + tuple(sorted(placements, key=lambda item: item.mesh_axis)), + ) + if self.logical_slice.flat_offset + self.logical_slice.length > self.parameter.numel: + raise ValueError("ShardIdentity.logical_slice exceeds the global parameter") + if self.process_group is None: + if self.local_member is not None or self.owner is not None: + raise ValueError("ShardIdentity members and owners require a process-group identity") + else: + if not isinstance(self.process_group, ProcessGroupIdentity): + raise TypeError("ShardIdentity.process_group must be a ProcessGroupIdentity") + if self.local_member not in self.process_group.ordered_members: + raise ValueError("ShardIdentity.local_member must belong to the process group") + if self.owner is not None and self.owner not in self.process_group.ordered_members: + raise ValueError("ShardIdentity.owner must belong to the process group") + if self.layout is not ParameterLayout.WHOLE_PARAMETER_OWNER and self.owner is not None: + raise ValueError("ShardIdentity.owner is valid only for whole-parameter ownership") + + full = self.logical_slice == LogicalSlice.full(self.parameter) + kinds = tuple(item.kind for item in self.placements) + if self.process_group is not None: + member_index = self.process_group.ordered_members.index(self.local_member) + for placement in self.placements: + if placement.parts != len(self.process_group.ordered_members) or placement.coordinate != member_index: + raise ValueError("ShardIdentity placement coordinates must match the ordered process-group members") + for placement in self.placements: + if placement.parameter_dimension is not None and placement.parameter_dimension >= len( + self.parameter.global_shape + ): + raise ValueError("ShardIdentity placement dimension exceeds parameter rank") + if self.layout is ParameterLayout.REPLICATED: + if not full or self.owner is not None: + raise ValueError("replicated identity must cover the full parameter") + if any(kind is not PlacementKind.REPLICATE for kind in kinds): + raise ValueError("replicated identity has a non-replicated placement") + if self.process_group is None and self.placements: + raise ValueError("an ungrouped replicated identity has no placements") + if self.process_group is not None and len(self.placements) != 1: + raise ValueError("a process-group replicated identity requires one placement") + elif self.layout is ParameterLayout.FLATTENED_ELEMENT_SHARD: + if self.process_group is None or self.owner is not None: + raise ValueError("flattened element shards require a process group and no owner") + if len(kinds) != 1 or kinds[0] is not PlacementKind.FLAT_SHARD: + raise ValueError("flattened element shards require one flat-shard placement") + elif self.layout is ParameterLayout.WHOLE_PARAMETER_OWNER: + if self.process_group is None or self.owner is None: + raise ValueError("whole-parameter ownership requires a process group and owner") + owns_parameter = self.local_member == self.owner + if owns_parameter and not full: + raise ValueError("the owner must carry the full whole-parameter logical slice") + if not owns_parameter and self.logical_slice != LogicalSlice(0, 0): + raise ValueError("a non-owner whole-parameter slice must be empty") + if len(kinds) != 1 or kinds[0] is not PlacementKind.WHOLE_PARAMETER_OWNER: + raise ValueError("whole-parameter ownership requires one owner placement") + elif self.layout is ParameterLayout.DTENSOR_1D_DEFAULT_WORLD: + raise ValueError( + "stable DTensor identity requires a logical-region descriptor and is " + "not implemented by the contiguous-slice identity schema" + ) + _validate_identity_schema_version("shard identity", self.schema_version) + + @property + def sort_key(self): + """Return a deterministic structural ordering key.""" + + member_index = -1 + group_name = "" + if self.process_group is not None: + group_name = self.process_group.semantic_name + member_index = self.process_group.ordered_members.index(self.local_member) + owner_index = -1 + if self.process_group is not None and self.owner is not None: + owner_index = self.process_group.ordered_members.index(self.owner) + placement_key = tuple( + ( + item.mesh_axis, + item.kind.value, + item.coordinate, + item.parts, + -1 if item.parameter_dimension is None else item.parameter_dimension, + ) + for item in self.placements + ) + return ( + self.parameter.fqn, + self.logical_slice.flat_offset, + self.logical_slice.length, + self.layout.value, + group_name, + member_index, + owner_index, + placement_key, + ) + + +@dataclass(frozen=True) +class ShardingManifest: + """Complete deterministic shard identity set supplied by one adapter.""" + + shards: Sequence[ShardIdentity] + schema_version: int = IDENTITY_SCHEMA_VERSION + + def __post_init__(self) -> None: + shards = _tuple(self.shards) + if not shards: + raise ValueError("ShardingManifest.shards must be non-empty") + if any(not isinstance(item, ShardIdentity) for item in shards): + raise TypeError("ShardingManifest.shards must contain ShardIdentity values") + ordered = tuple(sorted(shards, key=lambda item: item.sort_key)) + if len(set(ordered)) != len(ordered): + raise ValueError("ShardingManifest.shards must be unique") + object.__setattr__(self, "shards", ordered) + _validate_identity_schema_version("sharding manifest", self.schema_version) + + by_fqn = {} + for shard in ordered: + by_fqn.setdefault(shard.parameter.fqn, []).append(shard) + for fqn, parameter_shards in by_fqn.items(): + parameter = parameter_shards[0].parameter + if any(item.parameter != parameter for item in parameter_shards[1:]): + raise ValueError("manifest shards for {!r} disagree on parameter identity".format(fqn)) + layouts = {item.layout for item in parameter_shards} + groups = {item.process_group for item in parameter_shards} + if len(layouts) != 1 or len(groups) != 1: + raise ValueError("manifest shards for {!r} disagree on layout or process group".format(fqn)) + layout = parameter_shards[0].layout + group = parameter_shards[0].process_group + members = tuple(item.local_member for item in parameter_shards) + if group is None: + if len(parameter_shards) != 1: + raise ValueError("an ungrouped parameter must have exactly one manifest shard") + elif set(members) != set(group.ordered_members) or len(members) != len(group.ordered_members): + raise ValueError("manifest shards must contain each process-group member exactly once") + if group is not None: + placement_shapes = { + tuple( + ( + placement.mesh_axis, + placement.kind, + placement.parts, + placement.parameter_dimension, + ) + for placement in item.placements + ) + for item in parameter_shards + } + if len(placement_shapes) != 1: + raise ValueError("manifest member placements must agree on one topology") + + if layout is ParameterLayout.FLATTENED_ELEMENT_SHARD: + cursor = 0 + boundaries = {0} + empty_offsets = [] + for item in sorted( + parameter_shards, + key=lambda shard: ( + shard.logical_slice.flat_offset, + group.ordered_members.index(shard.local_member), + ), + ): + if item.logical_slice.length == 0: + empty_offsets.append(item.logical_slice.flat_offset) + continue + if item.logical_slice.flat_offset != cursor: + raise ValueError("flattened manifest slices must be gapless and non-overlapping") + cursor += item.logical_slice.length + boundaries.add(cursor) + if cursor != parameter.numel: + raise ValueError("flattened manifest slices must cover the global parameter") + if any(offset not in boundaries for offset in empty_offsets): + raise ValueError( + "empty flattened manifest slices must use a partition boundary" + ) + elif layout is ParameterLayout.REPLICATED: + if any(item.logical_slice != LogicalSlice.full(parameter) for item in parameter_shards): + raise ValueError("replicated manifest shards must all be complete") + elif layout is ParameterLayout.WHOLE_PARAMETER_OWNER: + owners = {item.owner for item in parameter_shards} + if len(owners) != 1: + raise ValueError("whole-parameter manifest shards must agree on one owner") + + def for_parameter(self, fqn: str) -> Tuple[ShardIdentity, ...]: + """Return one canonical parameter's shards in deterministic order.""" + + return tuple(item for item in self.shards if item.parameter.fqn == fqn) + + @dataclass(frozen=True) class StateField: """One named authoritative field, runtime cache, or transport field.""" @@ -1026,6 +1359,7 @@ def _hybrid_contract( __all__ = [ "CONTRACT_SCHEMA_VERSION", + "IDENTITY_SCHEMA_VERSION", "CheckpointSupport", "CheckpointTransport", "OptimizerCapabilities", @@ -1033,10 +1367,17 @@ def _hybrid_contract( "OptimizerContract", "OptimizerContractProvider", "OptimizerStateLayout", + "LogicalSlice", "ParameterLayout", + "ParameterIdentity", "ParameterStateRole", + "PlacementKind", "Precision", "ProcessGroupScope", + "ProcessGroupIdentity", + "ShardIdentity", + "ShardPlacement", + "ShardingManifest", "StateExtent", "StateField", "StateGeometry", diff --git a/tests/test_shard_identity_contracts.py b/tests/test_shard_identity_contracts.py new file mode 100644 index 0000000..e3c1823 --- /dev/null +++ b/tests/test_shard_identity_contracts.py @@ -0,0 +1,472 @@ +"""CPU coverage for canonical parameter and shard identity descriptors.""" + +from dataclasses import FrozenInstanceError + +import pytest + +import gefen +from gefen import ( + IDENTITY_SCHEMA_VERSION, + LogicalSlice, + ParameterIdentity, + ParameterLayout, + PlacementKind, + ProcessGroupIdentity, + ShardIdentity, + ShardPlacement, + ShardingManifest, +) + + +def _group(): + return ProcessGroupIdentity("data_parallel", ("rank:2", "rank:0", "rank:1")) + + +def _flat_shard(parameter, group, member, offset, length): + coordinate = group.ordered_members.index(member) + return ShardIdentity( + parameter, + ParameterLayout.FLATTENED_ELEMENT_SHARD, + LogicalSlice(offset, length), + placements=( + ShardPlacement( + "data_parallel", + PlacementKind.FLAT_SHARD, + coordinate, + len(group.ordered_members), + ), + ), + process_group=group, + local_member=member, + ) + + +def _owner_shard(parameter, group, member, owner): + coordinate = group.ordered_members.index(member) + logical_slice = LogicalSlice.full(parameter) if member == owner else LogicalSlice(0, 0) + return ShardIdentity( + parameter, + ParameterLayout.WHOLE_PARAMETER_OWNER, + logical_slice, + placements=( + ShardPlacement( + "data_parallel", + PlacementKind.WHOLE_PARAMETER_OWNER, + coordinate, + len(group.ordered_members), + ), + ), + process_group=group, + local_member=member, + owner=owner, + ) + + +def test_parameter_and_process_group_identities_are_exact_and_immutable(): + parameter = ParameterIdentity("Encoder.Block.Weight", [4, 8]) + group = ProcessGroupIdentity("pipeline:1/dp", ["worker:b", "worker:a"]) + + assert parameter.schema_version == IDENTITY_SCHEMA_VERSION + assert parameter.fqn == "Encoder.Block.Weight" + assert parameter.global_shape == (4, 8) + assert parameter.numel == 32 + assert group.ordered_members == ("worker:b", "worker:a") + with pytest.raises(FrozenInstanceError): + parameter.fqn = "changed" + with pytest.raises(FrozenInstanceError): + group.semantic_name = "changed" + + +@pytest.mark.parametrize( + "fqn", ["", " layer.weight", "layer.weight ", ".layer", "layer.", "layer..weight", "layer\x00weight"] +) +def test_parameter_identity_rejects_noncanonical_fqns(fqn): + with pytest.raises(ValueError): + ParameterIdentity(fqn, (4, 4)) + + +@pytest.mark.parametrize("shape", [(4, -1), (4, True), (4, 1.5)]) +def test_parameter_identity_rejects_invalid_global_shapes(shape): + with pytest.raises(ValueError): + ParameterIdentity("layer.weight", shape) + + +@pytest.mark.parametrize("shape", ["48", b"\x04\x08", bytearray((4, 8))]) +def test_parameter_identity_rejects_string_and_bytes_shape_containers(shape): + with pytest.raises(TypeError, match="sequence"): + ParameterIdentity("layer.weight", shape) + + +def test_process_group_identity_preserves_authoritative_member_order(): + group = _group() + assert group.ordered_members == ("rank:2", "rank:0", "rank:1") + with pytest.raises(ValueError, match="unique"): + ProcessGroupIdentity("dp", ("rank:0", "rank:0")) + with pytest.raises(ValueError, match="non-empty"): + ProcessGroupIdentity("dp", ()) + with pytest.raises(TypeError, match="sequence"): + ProcessGroupIdentity("dp", "ab") + + +def test_placement_and_logical_slice_validation(): + placement = ShardPlacement("dp", PlacementKind.DIMENSION_SHARD, 1, 2, 0) + assert placement.parameter_dimension == 0 + with pytest.raises(ValueError, match="within"): + ShardPlacement("dp", PlacementKind.FLAT_SHARD, 2, 2) + with pytest.raises(ValueError, match="only"): + ShardPlacement("dp", PlacementKind.REPLICATE, 0, 1, 0) + with pytest.raises(ValueError, match="requires"): + ShardPlacement("dp", PlacementKind.DIMENSION_SHARD, 0, 1) + with pytest.raises(ValueError): + LogicalSlice(-1, 1) + with pytest.raises(ValueError): + LogicalSlice(0, -1) + + +def test_replicated_identity_can_be_local_or_process_group_scoped(): + parameter = ParameterIdentity("layer.weight", (4, 4)) + local = ShardIdentity( + parameter, + ParameterLayout.REPLICATED, + LogicalSlice.full(parameter), + ) + assert local.process_group is None + + group = _group() + scoped = [ + ShardIdentity( + parameter, + ParameterLayout.REPLICATED, + LogicalSlice.full(parameter), + placements=( + ShardPlacement( + "data_parallel", + PlacementKind.REPLICATE, + group.ordered_members.index(member), + len(group.ordered_members), + ), + ), + process_group=group, + local_member=member, + ) + for member in group.ordered_members + ] + manifest = ShardingManifest(tuple(reversed(scoped))) + assert {item.local_member for item in manifest.shards} == set(group.ordered_members) + + +def test_contiguous_slice_schema_rejects_dtensor_identity_until_regions_exist(): + parameter = ParameterIdentity("layer.weight", (4, 4)) + group = ProcessGroupIdentity("dp", ("rank:0", "rank:1")) + with pytest.raises(ValueError, match="logical-region"): + ShardIdentity( + parameter, + ParameterLayout.DTENSOR_1D_DEFAULT_WORLD, + LogicalSlice(0, 8), + placements=(ShardPlacement("dp", PlacementKind.DIMENSION_SHARD, 0, 2, 0),), + process_group=group, + local_member="rank:0", + ) + + +@pytest.mark.parametrize("version", [True, 1.0, 0, 2]) +def test_parameter_identity_schema_version_requires_exact_supported_int(version): + with pytest.raises(ValueError, match="schema version"): + ParameterIdentity("layer.weight", (4, 4), schema_version=version) + + +def test_every_versioned_identity_descriptor_checks_its_schema_version(): + parameter = ParameterIdentity("layer.weight", (4, 4)) + local = ShardIdentity( + parameter, + ParameterLayout.REPLICATED, + LogicalSlice.full(parameter), + ) + with pytest.raises(ValueError, match="schema version"): + ProcessGroupIdentity("dp", ("rank:0",), schema_version=True) + with pytest.raises(ValueError, match="schema version"): + ShardIdentity( + parameter, + ParameterLayout.REPLICATED, + LogicalSlice.full(parameter), + schema_version=1.0, + ) + with pytest.raises(ValueError, match="schema version"): + ShardingManifest((local,), schema_version=False) + + +def test_flattened_manifest_normalizes_order_and_requires_exact_coverage(): + parameter = ParameterIdentity("Encoder.Weight", (4, 4)) + group = _group() + shards = ( + _flat_shard(parameter, group, "rank:1", 16, 0), + _flat_shard(parameter, group, "rank:0", 8, 8), + _flat_shard(parameter, group, "rank:2", 0, 8), + ) + + manifest = ShardingManifest(shards) + + assert tuple((item.logical_slice.flat_offset, item.local_member) for item in manifest.shards) == ( + (0, "rank:2"), + (8, "rank:0"), + (16, "rank:1"), + ) + assert manifest.for_parameter("Encoder.Weight") == manifest.shards + assert manifest.for_parameter("missing") == () + + with pytest.raises(ValueError, match="gapless"): + ShardingManifest( + ( + _flat_shard(parameter, group, "rank:2", 0, 7), + _flat_shard(parameter, group, "rank:0", 8, 8), + _flat_shard(parameter, group, "rank:1", 16, 0), + ) + ) + with pytest.raises(ValueError, match="gapless"): + ShardingManifest( + ( + _flat_shard(parameter, group, "rank:2", 0, 9), + _flat_shard(parameter, group, "rank:0", 8, 8), + _flat_shard(parameter, group, "rank:1", 16, 0), + ) + ) + with pytest.raises(ValueError, match="each process-group member"): + ShardingManifest(shards[:-1]) + inconsistent_axis = ShardIdentity( + parameter, + ParameterLayout.FLATTENED_ELEMENT_SHARD, + LogicalSlice(8, 8), + placements=(ShardPlacement("zero", PlacementKind.FLAT_SHARD, 1, 3),), + process_group=group, + local_member="rank:0", + ) + with pytest.raises(ValueError, match="one topology"): + ShardingManifest((shards[0], inconsistent_axis, shards[2])) + with pytest.raises(ValueError, match="unique"): + ShardingManifest((shards[0], shards[0])) + + +def test_flattened_manifest_empty_shard_is_order_independent_at_middle_boundary(): + parameter = ParameterIdentity("layer.weight", (4, 4)) + group = _group() + manifest = ShardingManifest( + ( + _flat_shard(parameter, group, "rank:2", 0, 8), + _flat_shard(parameter, group, "rank:0", 8, 8), + _flat_shard(parameter, group, "rank:1", 8, 0), + ) + ) + assert sum(item.logical_slice.length for item in manifest.shards) == 16 + with pytest.raises(ValueError, match="partition boundary"): + ShardingManifest( + ( + _flat_shard(parameter, group, "rank:2", 0, 8), + _flat_shard(parameter, group, "rank:0", 8, 8), + _flat_shard(parameter, group, "rank:1", 7, 0), + ) + ) + + +def test_flattened_shard_requires_explicit_group_and_in_bounds_slice(): + parameter = ParameterIdentity("layer.weight", (4, 4)) + group = _group() + placement = ShardPlacement("dp", PlacementKind.FLAT_SHARD, 0, 3) + with pytest.raises(ValueError, match="process group"): + ShardIdentity( + parameter, + ParameterLayout.FLATTENED_ELEMENT_SHARD, + LogicalSlice(0, 8), + placements=(placement,), + ) + with pytest.raises(ValueError, match="exceeds"): + _flat_shard(parameter, group, "rank:2", 12, 8) + with pytest.raises(ValueError, match="coordinates"): + ShardIdentity( + parameter, + ParameterLayout.FLATTENED_ELEMENT_SHARD, + LogicalSlice(0, 8), + placements=(ShardPlacement("dp", PlacementKind.FLAT_SHARD, 1, 3),), + process_group=group, + local_member="rank:2", + ) + + +def test_replicated_identity_rejects_ambiguous_multi_axis_placements(): + parameter = ParameterIdentity("layer.weight", (4, 4)) + group = ProcessGroupIdentity("mesh", ("rank:0",)) + placements = ( + ShardPlacement("tensor", PlacementKind.REPLICATE, 0, 1), + ShardPlacement("data", PlacementKind.REPLICATE, 0, 1), + ) + with pytest.raises(ValueError, match="one placement"): + ShardIdentity( + parameter, + ParameterLayout.REPLICATED, + LogicalSlice.full(parameter), + placements=placements, + process_group=group, + local_member="rank:0", + ) + + +def test_whole_parameter_owner_manifest_has_one_complete_owner(): + parameter = ParameterIdentity("layer.weight", (4, 4)) + group = _group() + shards = tuple(_owner_shard(parameter, group, member, "rank:0") for member in group.ordered_members) + + manifest = ShardingManifest(tuple(reversed(shards))) + + owner = next(item for item in manifest.shards if item.local_member == "rank:0") + assert owner.logical_slice == LogicalSlice.full(parameter) + assert all(item.logical_slice.length == 0 for item in manifest.shards if item.local_member != "rank:0") + with pytest.raises(ValueError, match="non-owner"): + ShardIdentity( + parameter, + ParameterLayout.WHOLE_PARAMETER_OWNER, + LogicalSlice.full(parameter), + placements=shards[0].placements, + process_group=group, + local_member="rank:2", + owner="rank:0", + ) + + +def test_zero_numel_whole_parameter_still_distinguishes_owner_by_member(): + parameter = ParameterIdentity("empty.weight", (0, 4)) + group = ProcessGroupIdentity("dp", ("rank:0", "rank:1")) + manifest = ShardingManifest( + tuple(_owner_shard(parameter, group, member, "rank:0") for member in group.ordered_members) + ) + assert tuple(item.owner for item in manifest.shards) == ("rank:0", "rank:0") + + +def test_whole_owner_manifest_rejects_owner_disagreement(): + parameter = ParameterIdentity("layer.weight", (4, 4)) + group = ProcessGroupIdentity("dp", ("rank:0", "rank:1")) + with pytest.raises(ValueError, match="one owner"): + ShardingManifest( + ( + _owner_shard(parameter, group, "rank:0", "rank:0"), + _owner_shard(parameter, group, "rank:1", "rank:1"), + ) + ) + + +def test_manifest_rejects_parameter_or_layout_disagreement(): + group = _group() + first = ParameterIdentity("layer.weight", (4, 4)) + different = ParameterIdentity("layer.weight", (2, 8)) + with pytest.raises(ValueError, match="parameter identity"): + ShardingManifest( + ( + _flat_shard(first, group, "rank:2", 0, 8), + _flat_shard(different, group, "rank:0", 8, 8), + _flat_shard(first, group, "rank:1", 16, 0), + ) + ) + flat = tuple( + _flat_shard(first, group, member, offset, length) + for member, offset, length in ( + ("rank:2", 0, 8), + ("rank:0", 8, 8), + ("rank:1", 16, 0), + ) + ) + replicated = ShardIdentity( + first, + ParameterLayout.REPLICATED, + LogicalSlice.full(first), + placements=( + ShardPlacement("data_parallel", PlacementKind.REPLICATE, 0, 3), + ), + process_group=group, + local_member="rank:2", + ) + with pytest.raises(ValueError, match="layout or process group"): + ShardingManifest((replicated, flat[1], flat[2])) + + other_group = ProcessGroupIdentity( + "other_data_parallel", group.ordered_members + ) + with pytest.raises(ValueError, match="layout or process group"): + ShardingManifest( + ( + flat[0], + _flat_shard(first, other_group, "rank:0", 8, 8), + flat[2], + ) + ) + + +def test_identity_descriptors_copy_mutable_input_sequences_deeply_enough(): + shape = [4, 4] + members = ["rank:0"] + placements = [ + ShardPlacement("dp", PlacementKind.REPLICATE, 0, 1), + ] + parameter = ParameterIdentity("layer.weight", shape) + group = ProcessGroupIdentity("dp", members) + shard = ShardIdentity( + parameter, + ParameterLayout.REPLICATED, + LogicalSlice.full(parameter), + placements=placements, + process_group=group, + local_member="rank:0", + ) + source_shards = [shard] + manifest = ShardingManifest(source_shards) + + shape.append(8) + members.append("rank:1") + placements.clear() + source_shards.clear() + + assert parameter.global_shape == (4, 4) + assert group.ordered_members == ("rank:0",) + assert shard.placements == ( + ShardPlacement("dp", PlacementKind.REPLICATE, 0, 1), + ) + assert manifest.shards == (shard,) + + +def test_manifest_order_is_fqn_then_logical_slice_then_member(): + group = ProcessGroupIdentity("dp", ("rank:1", "rank:0")) + alpha = ParameterIdentity("Alpha.Weight", (0,)) + zeta = ParameterIdentity("Zeta.Weight", (2,)) + shards = ( + _flat_shard(zeta, group, "rank:0", 2, 0), + _flat_shard(alpha, group, "rank:0", 0, 0), + _flat_shard(zeta, group, "rank:1", 0, 2), + _flat_shard(alpha, group, "rank:1", 0, 0), + ) + + manifest = ShardingManifest(tuple(reversed(shards))) + + assert tuple( + ( + item.parameter.fqn, + item.logical_slice.flat_offset, + item.local_member, + ) + for item in manifest.shards + ) == ( + ("Alpha.Weight", 0, "rank:1"), + ("Alpha.Weight", 0, "rank:0"), + ("Zeta.Weight", 0, "rank:1"), + ("Zeta.Weight", 2, "rank:0"), + ) + + +def test_identity_contracts_are_public_lazy_exports(): + for name in ( + "ParameterIdentity", + "ProcessGroupIdentity", + "PlacementKind", + "ShardPlacement", + "LogicalSlice", + "ShardIdentity", + "ShardingManifest", + ): + assert name in gefen.__all__ + assert getattr(gefen, name).__module__ == "gefen.contracts" From 8e0bb97390d4995e25a26644dc433aee307f47af Mon Sep 17 00:00:00 2001 From: thad0ctor Date: Sun, 12 Jul 2026 19:01:12 -0700 Subject: [PATCH 04/52] Add atomic post-sharding rebinding --- docs/optimizer_contracts.md | 8 + src/gefen/__init__.py | 5 + src/gefen/contracts.py | 29 +- src/gefen/gefen.py | 545 ++++++++++++++++- src/gefen/gefen_muon.py | 51 +- src/gefen/rebinding.py | 22 + tests/test_optimizer_contracts.py | 4 +- tests/test_rebinding_cpu.py | 810 +++++++++++++++++++++++++ tests/test_shard_identity_contracts.py | 2 + 9 files changed, 1465 insertions(+), 11 deletions(-) create mode 100644 src/gefen/rebinding.py create mode 100644 tests/test_rebinding_cpu.py diff --git a/docs/optimizer_contracts.md b/docs/optimizer_contracts.md index 6e2b7ee..cf7b358 100644 --- a/docs/optimizer_contracts.md +++ b/docs/optimizer_contracts.md @@ -33,6 +33,14 @@ The current DTensor declaration is deliberately narrow: `DTENSOR_1D_DEFAULT_WORL 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. +## Atomic post-sharding rebinding + +`Gefen.post_sharding(rebindings, manifest=...)` finalizes the complete local optimizer layout as one pre-initialization transaction. Each `ParameterRebinding` maps an existing optimizer slot to its local live tensor, or to `None` for an explicit whole-parameter non-owner. The global `ShardingManifest` remains descriptive: the core compares its FQN set with every local optimizer slot and validates each supplied local shard, but it does not infer live tensors, current members, owners, or runtime collective handles from manifest order. `rebind_parameter` and `rebind_shard` are one-slot conveniences and therefore apply only when they describe the optimizer's complete plan. + +Rebinding is allowed only while the entire optimizer is pristine: global step zero, no learned codebook, no gradients, no authoritative parameter state, no active capture stacks, and no nonzero device counters. The core stages every group, compatibility name, constructor-only state removal, canonical binding, cache invalidation, device counter, and checkpoint-schema update before publishing the result. A failed batch leaves the exact live optimizer objects unchanged. A successful batch preserves group order, group options, and released lowercase compatibility names while storing exact FQNs separately; it seals the layout against later incremental groups or rebindings. Targets must have no internal storage overlap and distinct targets may not overlap one another. Schema version 1 conservatively rejects multidimensional strided layouts whose element disjointness cannot be proven from dense stride spans, as well as distinct noncontiguous targets that share one storage even when their logical elements are disjoint. Tied aliases must already be collapsed to one optimizer slot. + +Plain Gefen currently rebinds complete replicated parameters and contiguous physical 1-D flattened element shards. A flattened logical matrix requires `factored_v_2d=False` because canonical row/column factored-state projection is not implemented. GefenMuon rebinds complete replicated matrices and can finalize/prune whole-parameter owner manifests, but stepping a whole-owner binding remains explicitly disabled until the independent adapter-defined process-group codebook scope is implemented. Whole-owner training, DTensor stable identity, Hybrid composite rebinding, canonical checkpoint I/O, state movement, and offload therefore remain unclaimed. + Plain Gefen declares replicated, flattened element-shard, and the narrow DTensor training layouts. Its PyTorch rank-local checkpoint transport is same-topology only. GefenMuon declares replicated and narrow DTensor training, with mode-specific state extents: `approx` state is local, `exact` state is logically global, and `distributed` momentum is held by the parameter owner while non-owners retain metadata only. Native Parallel-Muon checkpoints separately declare world-size owner redistribution, not placement-changing resharding. `GefenMuonHybrid` retains its nested child namespaces and does not flatten AdamW or Gefen child state into a fabricated common schema. Gefen and GefenMuon native loads, rank-local payload restoration through `load_state_dict`, and distributed-owner payload restoration prepare their complete core restore before changing local live optimizer state. Their transport entries report `atomic_load=True`. This is a per-optimizer-instance fail-before-mutation guarantee at the optimizer load boundary, not a coordinated all-rank commit or a guarantee over work an external checkpoint orchestrator performs before calling the optimizer. Load pre-hooks run before that boundary and load post-hooks run afterward, so arbitrary side effects in user hooks are also outside the guarantee. Hybrid composite loads do not yet provide the same guarantee and report `atomic_load=False`. diff --git a/src/gefen/__init__.py b/src/gefen/__init__.py index 416a767..3401c6c 100644 --- a/src/gefen/__init__.py +++ b/src/gefen/__init__.py @@ -24,6 +24,7 @@ "ParameterLayout", "ParameterIdentity", "ParameterStateRole", + "ParameterRebinding", "PlacementKind", "Precision", "ProcessGroupScope", @@ -63,6 +64,10 @@ def __getattr__(name): from . import params return getattr(params, name) + if name == "ParameterRebinding": + from .rebinding import ParameterRebinding + + return ParameterRebinding if name in ( "CONTRACT_SCHEMA_VERSION", "IDENTITY_SCHEMA_VERSION", diff --git a/src/gefen/contracts.py b/src/gefen/contracts.py index f861597..e3037f3 100644 --- a/src/gefen/contracts.py +++ b/src/gefen/contracts.py @@ -877,6 +877,10 @@ def _negative_capabilities( training: Tuple[TrainingSupport, ...], checkpoints: Tuple[CheckpointSupport, ...], supported_parameter_ranks: Optional[Tuple[int, ...]], + canonical_parameter_fqns: bool = False, + stable_shard_identity: bool = False, + shard_rebinding: bool = False, + post_sharding: bool = False, ) -> OptimizerCapabilities: return OptimizerCapabilities( training=training, @@ -884,18 +888,23 @@ def _negative_capabilities( precisions=_ALL_PRECISIONS, supported_parameter_ranks=supported_parameter_ranks, accepts_semantic_parameter_names=True, - canonical_parameter_fqns=False, - stable_shard_identity=False, + canonical_parameter_fqns=canonical_parameter_fqns, + stable_shard_identity=stable_shard_identity, explicit_process_group_codebook_scope=False, - shard_rebinding=False, - post_sharding=False, + shard_rebinding=shard_rebinding, + post_sharding=post_sharding, canonical_state_io=False, atomic_state_movement=False, state_offload=False, ) -def _gefen_contract(*, factored_v_2d: bool) -> OptimizerContract: +def _gefen_contract( + *, + factored_v_2d: bool, + canonical_parameter_fqns: bool = False, + stable_shard_identity: bool = False, +) -> OptimizerContract: block_fields = ( StateField("vmean", StateScope.PARAMETER, StateGeometry.BLOCK, True), StateField("vmean_step", StateScope.PARAMETER, StateGeometry.SCALAR, True), @@ -1070,6 +1079,10 @@ def _gefen_contract(*, factored_v_2d: bool) -> OptimizerContract: training=training, checkpoints=checkpoints, supported_parameter_ranks=None, + canonical_parameter_fqns=canonical_parameter_fqns, + stable_shard_identity=stable_shard_identity, + shard_rebinding=True, + post_sharding=True, ), ) @@ -1099,6 +1112,8 @@ def _gefen_muon_contract( sharded_modes: FrozenSet[str], normuon_modes: FrozenSet[str], non_normuon_modes: FrozenSet[str], + canonical_parameter_fqns: bool = False, + stable_shard_identity: bool = False, ) -> OptimizerContract: sharded_modes = _frozenset(sharded_modes) normuon_modes = _frozenset(normuon_modes) @@ -1304,6 +1319,10 @@ def _gefen_muon_contract( training=training, checkpoints=tuple(checkpoints), supported_parameter_ranks=(2,), + canonical_parameter_fqns=canonical_parameter_fqns, + stable_shard_identity=stable_shard_identity, + shard_rebinding=True, + post_sharding=True, ), ) diff --git a/src/gefen/gefen.py b/src/gefen/gefen.py index eecb73f..7d1fa62 100644 --- a/src/gefen/gefen.py +++ b/src/gefen/gefen.py @@ -15,15 +15,24 @@ import math import os import warnings -from collections import OrderedDict +from collections import defaultdict, OrderedDict from itertools import chain from typing import Iterable, Optional, Tuple, Union import torch import torch.nn as nn -from gefen.contracts import OptimizerContract, _gefen_contract +from gefen.contracts import ( + LogicalSlice, + OptimizerContract, + ParameterIdentity, + ParameterLayout, + ShardIdentity, + ShardingManifest, + _gefen_contract, +) from gefen.partitioning import find_period_by_block_variance +from gefen.rebinding import ParameterRebinding import gefen.quantization as quantization_module from gefen.kernels.automatic_vmean import ( automatic_vmean_update_cuda as _automatic_vmean_update_cuda, @@ -1052,6 +1061,14 @@ def __init__( weight_decay=weight_decay, ) self._param_names = {} + # Exact canonical identity remains separate from compatibility + # ``param_names``. A full post_sharding transaction publishes these + # registries only after validating every optimizer slot and manifest. + self._gefen_shard_bindings = {} + self._gefen_local_shard_bindings = () + self._gefen_sharding_manifest = None + self._gefen_post_sharding_finalized = False + self._gefen_finalized_slots = () # ``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 @@ -1072,7 +1089,519 @@ def _step_supports_amp_scaling(self) -> bool: def optimizer_contract(self) -> OptimizerContract: """Return the immutable state-layout and integration capability contract.""" - return _gefen_contract(factored_v_2d=self._factored_v_2d) + identity_ready = self._canonical_identity_ready() + return _gefen_contract( + factored_v_2d=self._factored_v_2d, + canonical_parameter_fqns=identity_ready, + stable_shard_identity=identity_ready, + ) + + def _canonical_identity_ready(self) -> bool: + if ( + not self._gefen_post_sharding_finalized + or self._gefen_sharding_manifest is None + ): + return False + return self._finalized_binding_layout_matches() + + def _finalized_binding_layout_matches(self) -> bool: + if len(self.param_groups) != len(self._gefen_finalized_slots): + return False + for group, expected in zip(self.param_groups, self._gefen_finalized_slots): + params = group.get("params") + if not isinstance(params, (list, tuple)) or len(params) != len(expected): + return False + if any(live is not bound for live, bound in zip(params, expected)): + return False + live_params = [ + param for group in self.param_groups for param in group["params"] + ] + bound = [ + (parameter, shard) + for parameter, shard in self._gefen_local_shard_bindings + if parameter is not None + ] + if len(live_params) != len(bound) or len(self._gefen_shard_bindings) != len( + bound + ): + return False + for parameter, shard in bound: + if not self._parameter_in(live_params, parameter): + return False + if self._gefen_shard_bindings.get(parameter) != shard: + return False + return True + + def _assert_finalized_binding_layout(self) -> None: + if self._gefen_post_sharding_finalized and not self._finalized_binding_layout_matches(): + raise RuntimeError( + "Gefen finalized parameter layout changed outside post_sharding" + ) + + def parameter_identity(self, parameter) -> ParameterIdentity: + """Return the exact canonical identity bound to one live parameter.""" + + return self.shard_identity(parameter).parameter + + def shard_identity(self, parameter) -> ShardIdentity: + """Return the stable shard identity bound to one live parameter.""" + + self._assert_finalized_binding_layout() + try: + return self._gefen_shard_bindings[parameter] + except KeyError: + raise KeyError("parameter has no finalized Gefen shard identity") from None + + def shard_bindings(self): + """Return local tensor/identity pairs in canonical structural order.""" + + self._assert_finalized_binding_layout() + return self._gefen_local_shard_bindings + + def sharding_manifest(self): + """Return the finalized global identity manifest, or ``None``.""" + + self._assert_finalized_binding_layout() + return self._gefen_sharding_manifest + + @staticmethod + def _parameter_in(parameters, candidate) -> bool: + return any(item is candidate for item in parameters) + + @staticmethod + def _assert_rebound_storage_disjoint(parameters) -> None: + storage_ranges = [] + for parameter in parameters: + if parameter.numel() == 0: + continue + storage = parameter.untyped_storage() + storage_id = (str(parameter.device), storage.data_ptr()) + if not parameter.is_contiguous(): + for other_id, _, _ in storage_ranges: + if other_id == storage_id: + raise ValueError( + "rebound target tensors must not share noncontiguous storage" + ) + storage_ranges.append((storage_id, None, None)) + continue + start = parameter.storage_offset() * parameter.element_size() + end = start + parameter.numel() * parameter.element_size() + for other_id, other_start, other_end in storage_ranges: + if other_id != storage_id: + continue + if other_start is None or max(start, other_start) < min(end, other_end): + raise ValueError( + "rebound target tensor storage ranges must not overlap" + ) + storage_ranges.append((storage_id, start, end)) + + @staticmethod + def _target_may_have_internal_storage_overlap(parameter) -> bool: + """Conservatively prove disjoint positive-stride tensor elements.""" + + required_span = 1 + dimensions = sorted( + (stride, size) + for size, stride in zip(parameter.shape, parameter.stride()) + if size > 1 + ) + for stride, size in dimensions: + if stride < required_span: + return True + required_span += (size - 1) * stride + return False + + def _assert_rebinding_pristine(self, rebindings) -> None: + if self._gefen_post_sharding_finalized: + raise RuntimeError("Gefen post-sharding identity is already finalized") + if ( + self._gefen_shard_bindings + or self._gefen_local_shard_bindings + or self._gefen_sharding_manifest is not None + or self._gefen_finalized_slots + ): + raise RuntimeError( + "Gefen parameter rebinding found an incomplete prior identity plan" + ) + global_step = self._validate_rank_local_counter( + "gefen_global_step", self._gefen_global_step + ) + if global_step != 0 or self._gefen_codebook is not None: + raise RuntimeError( + "Gefen parameter rebinding is allowed only before optimizer mutation" + ) + if self._capt_stacks is not None: + raise RuntimeError( + "Gefen parameter rebinding cannot discard active capturable stacks" + ) + if self._gefen_codebook_by_device or self._gefen_codebook_lut_by_device: + raise RuntimeError( + "Gefen parameter rebinding cannot discard initialized codebook caches" + ) + for cache_name in ( + "_gefen_global_step_by_device", + "_sr_seed_by_device", + ): + cache = getattr(self, cache_name) + for value in cache.values(): + if ( + not torch.is_tensor(value) + or value.numel() != 1 + or value.dtype == torch.bool + or not bool(torch.isfinite(value.detach()).all()) + or float(value.detach().cpu().item()) != 0.0 + ): + raise RuntimeError( + "Gefen parameter rebinding found non-pristine device counters" + ) + + live = [param for group in self.param_groups for param in group["params"]] + sources = [item.old_parameter for item in rebindings] + allowed_state_keys = tuple(live) + tuple(sources) + for parameter in self._param_names: + if not self._parameter_in(allowed_state_keys, parameter): + raise RuntimeError( + "Gefen parameter rebinding found orphan parameter-name state" + ) + compatibility_name = self._param_names[parameter] + if ( + not isinstance(compatibility_name, str) + or compatibility_name != compatibility_name.lower() + ): + raise RuntimeError( + "Gefen parameter rebinding found an invalid compatibility name" + ) + allowed_names = {"name", _RANK_LOCAL_MEMBER_KEY} + for parameter, pstate in self.state.items(): + if not self._parameter_in(allowed_state_keys, parameter): + raise RuntimeError( + "Gefen parameter rebinding found orphan optimizer state" + ) + if not isinstance(pstate, dict): + raise RuntimeError( + "Gefen parameter rebinding found invalid constructor state" + ) + for key in pstate: + if key in allowed_names: + continue + if isinstance(key, str) and key.startswith( + _RANK_LOCAL_PAYLOAD_KEY_PREFIX + ): + continue + raise RuntimeError( + "Gefen parameter rebinding found initialized optimizer state" + ) + state_name = pstate.get("name") + if state_name is not None and ( + not isinstance(state_name, str) or state_name != state_name.lower() + ): + raise RuntimeError( + "Gefen parameter rebinding found an invalid state name" + ) + + def _validate_rebinding_layout(self, rebinding: ParameterRebinding) -> None: + shard = rebinding.shard + target = rebinding.new_parameter + if shard.layout is ParameterLayout.REPLICATED: + if target is None or tuple(target.shape) != shard.parameter.global_shape: + raise ValueError( + "replicated Gefen rebinding requires complete parameter storage" + ) + elif shard.layout is ParameterLayout.FLATTENED_ELEMENT_SHARD: + if target is None or target.ndim != 1: + raise ValueError( + "flattened Gefen rebinding requires a live 1-D tensor shard" + ) + if not target.is_contiguous(): + raise ValueError( + "flattened Gefen rebinding requires contiguous physical storage" + ) + if self._factored_v_2d and len(shard.parameter.global_shape) == 2: + raise ValueError( + "flattened logical matrices require factored_v_2d=False until " + "matrix-aware factored-state projection is implemented" + ) + else: + raise ValueError( + "plain Gefen rebinding supports replicated or flattened element " + "shards only" + ) + if target.numel() != shard.logical_slice.length: + raise ValueError( + "Gefen rebound tensor numel does not match its logical slice" + ) + + def _stage_post_sharding(self, rebindings, manifest): + self._assert_rebinding_pristine(rebindings) + live_slots = [] + for group_index, group in enumerate(self.param_groups): + params = list(group["params"]) + names = list(group.get("param_names", ())) + if len(names) != len(params): + names = [self._param_name(param) for param in params] + for parameter_index, (parameter, name) in enumerate(zip(params, names)): + live_slots.append( + (group_index, parameter_index, parameter, str(name)) + ) + if len(rebindings) != len(live_slots): + raise ValueError( + "post_sharding requires exactly one rebinding for every optimizer slot" + ) + + manifest_fqns = { + shard.parameter.fqn for shard in manifest.shards + } + local_fqns = [item.shard.parameter.fqn for item in rebindings] + if len(set(local_fqns)) != len(local_fqns): + raise ValueError("local canonical parameter FQNs must be unique") + if set(local_fqns) != manifest_fqns: + raise ValueError( + "post_sharding manifest FQNs must exactly match optimizer slots" + ) + + assigned_positions = {} + binding_names = {} + final_targets = [] + seen_sources = [] + local_member_by_group = {} + for binding_index, rebinding in enumerate(rebindings): + if not isinstance(rebinding.old_parameter, torch.Tensor): + raise TypeError("rebinding source must be a Tensor") + if self._parameter_in(seen_sources, rebinding.old_parameter): + raise ValueError("rebinding source tensors must be unique") + seen_sources.append(rebinding.old_parameter) + process_group = rebinding.shard.process_group + if process_group is not None: + previous_member = local_member_by_group.setdefault( + process_group, rebinding.shard.local_member + ) + if previous_member != rebinding.shard.local_member: + raise ValueError( + "local shard bindings must use one member per process group" + ) + target = rebinding.new_parameter + if target is not None: + if not isinstance(target, torch.Tensor): + raise TypeError("rebinding target must be a Tensor or None") + if self._is_dtensor_parameter(target): + raise ValueError( + "stable DTensor rebinding requires the deferred logical-region " + "identity schema" + ) + if torch.is_complex(target): + raise ValueError("Gefen does not support complex rebound parameters") + if target.layout is not torch.strided or target.dtype not in ( + torch.float16, + torch.bfloat16, + torch.float32, + torch.float64, + ): + raise ValueError( + "rebound parameters require strided floating-point storage" + ) + if target.is_meta: + raise ValueError("rebound parameters require materialized storage") + if ( + rebinding.shard.layout + is ParameterLayout.FLATTENED_ELEMENT_SHARD + and not target.is_contiguous() + ): + raise ValueError( + "flattened Gefen rebinding requires contiguous physical storage" + ) + if self._target_may_have_internal_storage_overlap(target): + raise ValueError( + "rebound target storage must be provably free of internal " + "storage overlap" + ) + if not target.is_leaf and not target.retains_grad: + raise ValueError("can't optimize a non-leaf rebound Tensor") + if rebinding.old_parameter.grad is not None or ( + target is not None and target.grad is not None + ): + raise RuntimeError( + "post_sharding must run before source or target gradients exist" + ) + if rebinding.shard not in manifest.shards: + raise ValueError("local shard identity is absent from the manifest") + whole_owner = ( + rebinding.shard.layout is ParameterLayout.WHOLE_PARAMETER_OWNER + ) + local_owns = rebinding.shard.local_member == rebinding.shard.owner + if whole_owner and local_owns != (target is not None): + raise ValueError( + "whole-parameter owner bindings require storage only on the owner" + ) + if not whole_owner and target is None: + raise ValueError("only a whole-parameter non-owner may bind None") + + source_positions = [ + index + for index, slot in enumerate(live_slots) + if slot[2] is rebinding.old_parameter + ] + target_positions = [] + if target is not None: + target_positions = [ + index + for index, slot in enumerate(live_slots) + if slot[2] is target + ] + if len(source_positions) == 1: + position = source_positions[0] + if target_positions and target_positions != [position]: + raise ValueError( + "rebinding target already occupies another optimizer slot" + ) + elif len(source_positions) == 0 and len(target_positions) == 1: + position = target_positions[0] + source_known = ( + rebinding.old_parameter in self.state + or rebinding.old_parameter in self._param_names + ) + if not source_known: + raise ValueError("rebinding source is not registered with Gefen") + else: + raise ValueError( + "rebinding source must identify exactly one optimizer slot" + ) + if position in assigned_positions: + raise ValueError("multiple rebindings target one optimizer slot") + assigned_positions[position] = binding_index + compatibility_name = self._param_names.get(rebinding.old_parameter) + source_state = self.state.get(rebinding.old_parameter) + if compatibility_name is None and isinstance(source_state, dict): + compatibility_name = source_state.get("name") + if compatibility_name is None: + compatibility_name = live_slots[position][3] + if ( + not isinstance(compatibility_name, str) + or compatibility_name != compatibility_name.lower() + ): + raise ValueError( + "rebinding source compatibility name must be a lowercase string" + ) + binding_names[binding_index] = compatibility_name + if target is not None: + if self._parameter_in(final_targets, target): + raise ValueError("rebound target tensors must be unique") + final_targets.append(target) + self._validate_rebinding_layout(rebinding) + + if len(assigned_positions) != len(live_slots): + raise ValueError("post_sharding did not bind every optimizer slot") + self._assert_rebound_storage_disjoint(final_targets) + + staged = object.__new__(type(self)) + staged.__dict__ = self.__dict__.copy() + staged.param_groups = [] + staged.state = defaultdict(dict) + staged._param_names = {} + staged._gefen_shard_bindings = {} + local_bindings = [] + slot_cursor = 0 + for group in self.param_groups: + staged_group = dict(group) + staged_group.pop("_gefen_checkpoint_metadata", None) + staged_params = [] + staged_names = [] + names = list(group.get("param_names", ())) + if len(names) != len(group["params"]): + names = [self._param_name(param) for param in group["params"]] + for _ in names: + binding_index = assigned_positions[slot_cursor] + rebinding = rebindings[binding_index] + compatibility_name = binding_names[binding_index] + slot_cursor += 1 + target = rebinding.new_parameter + local_bindings.append((target, rebinding.shard)) + if target is None: + continue + staged_params.append(target) + staged_names.append(compatibility_name) + staged._param_names[target] = compatibility_name + staged.state[target]["name"] = compatibility_name + staged._gefen_shard_bindings[target] = rebinding.shard + staged_group["params"] = staged_params + staged_group["param_names"] = staged_names + staged.param_groups.append(staged_group) + + staged._gefen_local_shard_bindings = tuple( + sorted(local_bindings, key=lambda item: item[1].sort_key) + ) + staged._gefen_sharding_manifest = manifest + staged._gefen_post_sharding_finalized = True + staged._gefen_codebook_by_device = {} + staged._gefen_codebook_lut_by_device = {} + staged._sr_seed_by_device = {} + staged._gefen_global_step_by_device = {} + staged._capt_stacks = None + staged._static_mark_sig = None + staged._lr_scalar_cache = None + staged._ensure_gefen_global_step_devices() + staged._install_rank_local_checkpoint_schema() + staged._gefen_finalized_slots = tuple( + tuple(group["params"]) for group in staged.param_groups + ) + return staged + + def post_sharding(self, rebindings, *, manifest: ShardingManifest) -> None: + """Atomically finalize every local optimizer slot after sharding.""" + + if not isinstance(manifest, ShardingManifest): + raise TypeError("manifest must be a ShardingManifest") + if isinstance(rebindings, (str, bytes)): + raise TypeError("rebindings must be a sequence") + rebindings = tuple(rebindings) + if not rebindings or any( + not isinstance(item, ParameterRebinding) for item in rebindings + ): + raise TypeError( + "rebindings must contain one or more ParameterRebinding values" + ) + sources = [] + for rebinding in rebindings: + if self._parameter_in(sources, rebinding.old_parameter): + raise ValueError("rebinding source tensors must be unique") + sources.append(rebinding.old_parameter) + staged = self._stage_post_sharding(rebindings, manifest) + self.__dict__.update(staged.__dict__) + + def rebind_shard( + self, + old_parameter, + new_parameter, + *, + shard: ShardIdentity, + manifest: ShardingManifest, + ) -> None: + """Apply a complete one-slot shard plan through ``post_sharding``.""" + + self.post_sharding( + (ParameterRebinding(old_parameter, new_parameter, shard),), + manifest=manifest, + ) + + def rebind_parameter( + self, + old_parameter, + new_parameter, + *, + identity: ParameterIdentity, + ) -> None: + """Bind one complete replicated parameter with a canonical identity.""" + + shard = ShardIdentity( + identity, + ParameterLayout.REPLICATED, + LogicalSlice.full(identity), + ) + self.rebind_shard( + old_parameter, + new_parameter, + shard=shard, + manifest=ShardingManifest((shard,)), + ) @staticmethod def _normalize_param_groups(params): @@ -1181,6 +1710,10 @@ def add_param_group(self, param_group): stable lowercase name is stored in its per-param state and in the group's ``param_names`` list for introspection. """ + if getattr(self, "_gefen_post_sharding_finalized", False): + raise RuntimeError( + "Gefen cannot add parameter groups after post_sharding finalization" + ) if not isinstance(param_group, dict): raise TypeError( "param_group must be a dict, got {}".format( @@ -3710,8 +4243,10 @@ def _step_automatic_merged(self, items) -> None: def state_dict(self): """Run optimizer state-dict hooks around Gefen's complete schema.""" + self._assert_finalized_binding_layout() for pre_hook in self._optimizer_state_dict_pre_hooks.values(): pre_hook(self) + self._assert_finalized_binding_layout() state_dict = self._state_dict_impl() for post_hook in self._optimizer_state_dict_post_hooks.values(): hook_result = post_hook(self, state_dict) @@ -4466,11 +5001,13 @@ def _pack_legacy_param_groups_for_load(self, state_dict): def load_state_dict(self, state_dict): """Atomically restore Gefen state between the public load hooks.""" + self._assert_finalized_binding_layout() 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() staged = self._stage_load_state_dict(state_dict) self._commit_staged_load_state_dict(staged) for post_hook in self._optimizer_load_state_dict_post_hooks.values(): @@ -5370,6 +5907,7 @@ def step(self, closure=None): closure feed the first step's codebook learning correctly. The returned loss is passed through. """ + self._assert_finalized_binding_layout() self._assert_capturable_if_capturing() loss = None @@ -5377,6 +5915,7 @@ def step(self, closure=None): with torch.enable_grad(): loss = closure() + self._assert_finalized_binding_layout() _assert_optimizer_gradients_structurally_valid(self) # GradScaler invokes native-AMP optimizers even on overflow. Decide diff --git a/src/gefen/gefen_muon.py b/src/gefen/gefen_muon.py index ffb275b..3d32a54 100644 --- a/src/gefen/gefen_muon.py +++ b/src/gefen/gefen_muon.py @@ -6,7 +6,7 @@ import torch import torch.nn as nn -from gefen.contracts import OptimizerContract, _gefen_muon_contract +from gefen.contracts import OptimizerContract, ParameterLayout, _gefen_muon_contract from gefen.gefen import ( Gefen, _amp_prepare_optimizer_step, @@ -766,6 +766,48 @@ def optimizer_contract(self) -> OptimizerContract: sharded_modes=sharded_modes, normuon_modes=normuon_modes, non_normuon_modes=non_normuon_modes, + canonical_parameter_fqns=self._canonical_identity_ready(), + stable_shard_identity=self._canonical_identity_ready(), + ) + + def _validate_rebinding_layout(self, rebinding) -> None: + shard = rebinding.shard + target = rebinding.new_parameter + if len(shard.parameter.global_shape) != 2: + raise ValueError("GefenMuon canonical parameters must be logical matrices") + if shard.layout is ParameterLayout.REPLICATED: + if target is None or tuple(target.shape) != shard.parameter.global_shape: + raise ValueError( + "replicated GefenMuon rebinding requires one complete 2-D matrix" + ) + elif shard.layout is ParameterLayout.WHOLE_PARAMETER_OWNER: + local_owns = shard.local_member == shard.owner + if local_owns and ( + target is None + or target.ndim != 2 + or tuple(target.shape) != shard.parameter.global_shape + ): + raise ValueError( + "a GefenMuon whole-parameter owner requires one complete 2-D matrix" + ) + if not local_owns and target is not None: + raise ValueError( + "a GefenMuon whole-parameter non-owner must not retain storage" + ) + else: + raise ValueError( + "GefenMuon rebinding supports replicated complete matrices or " + "whole-parameter ownership only" + ) + if target is not None and target.numel() != shard.logical_slice.length: + raise ValueError( + "GefenMuon rebound tensor numel does not match its logical slice" + ) + + def _has_unscoped_whole_owner_bindings(self) -> bool: + return any( + shard.layout is ParameterLayout.WHOLE_PARAMETER_OWNER + for _, shard in self._gefen_local_shard_bindings ) def add_param_group(self, param_group): @@ -2619,12 +2661,19 @@ def _load_state_dict_impl(self, state_dict): @torch.no_grad() def step(self, closure=None): + self._assert_finalized_binding_layout() + if self._has_unscoped_whole_owner_bindings(): + raise RuntimeError( + "GefenMuon whole-parameter owner stepping requires the separate " + "explicit process-group codebook scope, which is not implemented" + ) self._assert_capturable_if_capturing() loss = None if closure is not None: with torch.enable_grad(): loss = closure() + self._assert_finalized_binding_layout() _assert_optimizer_gradients_structurally_valid( self, require_2d_params=True ) diff --git a/src/gefen/rebinding.py b/src/gefen/rebinding.py new file mode 100644 index 0000000..657c592 --- /dev/null +++ b/src/gefen/rebinding.py @@ -0,0 +1,22 @@ +"""Runtime requests for atomic parameter/shard rebinding.""" + +from dataclasses import dataclass +from typing import Optional + +from gefen.contracts import ShardIdentity + + +@dataclass(frozen=True, eq=False) +class ParameterRebinding: + """Bind one optimizer slot to a local tensor or prune it as a non-owner.""" + + old_parameter: object + new_parameter: Optional[object] + shard: ShardIdentity + + def __post_init__(self) -> None: + if not isinstance(self.shard, ShardIdentity): + raise TypeError("ParameterRebinding.shard must be a ShardIdentity") + + +__all__ = ["ParameterRebinding"] diff --git a/tests/test_optimizer_contracts.py b/tests/test_optimizer_contracts.py index 206f9e4..4157520 100644 --- a/tests/test_optimizer_contracts.py +++ b/tests/test_optimizer_contracts.py @@ -162,8 +162,8 @@ def test_plain_contract_matches_live_persistent_state(factored_v_2d): assert not contract.capabilities.canonical_parameter_fqns assert not contract.capabilities.stable_shard_identity assert not contract.capabilities.explicit_process_group_codebook_scope - assert not contract.capabilities.shard_rebinding - assert not contract.capabilities.post_sharding + 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 diff --git a/tests/test_rebinding_cpu.py b/tests/test_rebinding_cpu.py new file mode 100644 index 0000000..200ae70 --- /dev/null +++ b/tests/test_rebinding_cpu.py @@ -0,0 +1,810 @@ +"""CPU coverage for atomic pre-initialization parameter rebinding.""" + +import copy + +import pytest +import torch + +from gefen import ( + Gefen, + GefenMuon, + GefenMuonHybrid, + LogicalSlice, + ParameterIdentity, + ParameterLayout, + ParameterRebinding, + PlacementKind, + ProcessGroupIdentity, + ShardIdentity, + ShardPlacement, + ShardingManifest, +) + + +def _replicated(parameter, fqn, shape=None): + identity = ParameterIdentity(fqn, tuple(parameter.shape) if shape is None else shape) + shard = ShardIdentity( + identity, + ParameterLayout.REPLICATED, + LogicalSlice.full(identity), + ) + return shard + + +def _flat_manifest(fqn, shape, lengths, local_member): + identity = ParameterIdentity(fqn, shape) + members = tuple("rank:{}".format(index) for index in range(len(lengths))) + group = ProcessGroupIdentity("data_parallel", members) + shards = [] + offset = 0 + for coordinate, (member, length) in enumerate(zip(members, lengths)): + shards.append( + ShardIdentity( + identity, + ParameterLayout.FLATTENED_ELEMENT_SHARD, + LogicalSlice(offset, length), + placements=( + ShardPlacement( + "data_parallel", + PlacementKind.FLAT_SHARD, + coordinate, + len(members), + ), + ), + process_group=group, + local_member=member, + ) + ) + offset += length + manifest = ShardingManifest(tuple(reversed(shards))) + local = next(item for item in manifest.shards if item.local_member == local_member) + return local, manifest + + +def _owner_shard(identity, group, local_member, owner): + coordinate = group.ordered_members.index(local_member) + return ShardIdentity( + identity, + ParameterLayout.WHOLE_PARAMETER_OWNER, + LogicalSlice.full(identity) if local_member == owner else LogicalSlice(0, 0), + placements=( + ShardPlacement( + "data_parallel", + PlacementKind.WHOLE_PARAMETER_OWNER, + coordinate, + len(group.ordered_members), + ), + ), + process_group=group, + local_member=local_member, + owner=owner, + ) + + +def _nested_equal(left, right): + assert type(left) is type(right) + if isinstance(left, dict): + assert set(left) == set(right) + for key in left: + _nested_equal(left[key], right[key]) + elif isinstance(left, (list, tuple)): + assert len(left) == len(right) + for left_item, right_item in zip(left, right): + _nested_equal(left_item, right_item) + elif torch.is_tensor(left): + assert torch.equal(left, right) + else: + assert left == right + + +def _snapshot(optimizer): + return { + "state": optimizer.state, + "state_items": tuple( + (parameter, pstate, copy.deepcopy(pstate)) for parameter, pstate in optimizer.state.items() + ), + "param_groups": optimizer.param_groups, + "groups": tuple( + ( + group, + group["params"], + tuple(group["params"]), + copy.deepcopy({key: value for key, value in group.items() if key != "params"}), + ) + for group in optimizer.param_groups + ), + "param_names": optimizer._param_names, + "param_names_value": dict(optimizer._param_names), + "bindings": optimizer._gefen_shard_bindings, + "local_bindings": optimizer._gefen_local_shard_bindings, + "manifest": optimizer._gefen_sharding_manifest, + "finalized": optimizer._gefen_post_sharding_finalized, + "finalized_slots": optimizer._gefen_finalized_slots, + "caches": tuple( + (name, getattr(optimizer, name)) + for name in ( + "_gefen_codebook_by_device", + "_gefen_codebook_lut_by_device", + "_sr_seed_by_device", + "_gefen_global_step_by_device", + ) + ), + "capt_stacks": optimizer._capt_stacks, + "static_mark_sig": optimizer._static_mark_sig, + "global_step": optimizer._gefen_global_step, + "codebook": optimizer._gefen_codebook, + } + + +def _assert_snapshot(optimizer, snapshot): + assert optimizer.state is snapshot["state"] + assert optimizer.param_groups is snapshot["param_groups"] + assert optimizer._param_names is snapshot["param_names"] + assert optimizer._param_names == snapshot["param_names_value"] + assert optimizer._gefen_shard_bindings is snapshot["bindings"] + assert optimizer._gefen_local_shard_bindings is snapshot["local_bindings"] + assert optimizer._gefen_sharding_manifest is snapshot["manifest"] + assert optimizer._gefen_post_sharding_finalized is snapshot["finalized"] + assert optimizer._gefen_finalized_slots is snapshot["finalized_slots"] + assert optimizer._capt_stacks is snapshot["capt_stacks"] + assert optimizer._static_mark_sig is snapshot["static_mark_sig"] + assert optimizer._gefen_global_step is snapshot["global_step"] + assert optimizer._gefen_codebook is snapshot["codebook"] + for parameter, state_ref, expected in snapshot["state_items"]: + assert optimizer.state[parameter] is state_ref + _nested_equal(optimizer.state[parameter], expected) + for group, params_ref, expected_params, expected_values in snapshot["groups"]: + assert group in optimizer.param_groups + assert group["params"] is params_ref + assert tuple(group["params"]) == expected_params + _nested_equal( + {key: value for key, value in group.items() if key != "params"}, + expected_values, + ) + for name, cache_ref in snapshot["caches"]: + assert getattr(optimizer, name) is cache_ref + + +def test_replicated_rebind_preserves_legacy_name_and_enables_identity_contract(): + old = torch.nn.Parameter(torch.arange(16, dtype=torch.float32).reshape(4, 4)) + new = torch.nn.Parameter(old.detach().clone()) + optimizer = Gefen([("Encoder.MixedCase.Weight", old)], fused=False) + old_group_options = {key: value for key, value in optimizer.param_groups[0].items() if key != "params"} + identity = ParameterIdentity("Encoder.MixedCase.Weight", (4, 4)) + + optimizer.rebind_parameter(old, new, identity=identity) + + assert optimizer.param_groups[0]["params"] == [new] + assert {key: value for key, value in optimizer.param_groups[0].items() if key != "params"} == old_group_options + assert old not in optimizer.state + assert optimizer.state[new] == {"name": "encoder.mixedcase.weight"} + assert optimizer.parameter_identity(new) == identity + assert optimizer.shard_identity(new).parameter.fqn == "Encoder.MixedCase.Weight" + assert optimizer.shard_bindings() == ((new, optimizer.shard_identity(new)),) + assert optimizer.sharding_manifest().shards == (optimizer.shard_identity(new),) + contract = optimizer.optimizer_contract() + assert contract.capabilities.canonical_parameter_fqns + assert contract.capabilities.stable_shard_identity + assert contract.capabilities.shard_rebinding + assert contract.capabilities.post_sharding + assert not contract.capabilities.canonical_state_io + assert not contract.capabilities.explicit_process_group_codebook_scope + with pytest.raises(RuntimeError, match="already finalized"): + optimizer.rebind_parameter(new, new, identity=identity) + with pytest.raises(RuntimeError, match="cannot add"): + optimizer.add_param_group({"params": [torch.nn.Parameter(torch.ones(2))]}) + + +def test_same_tensor_identity_only_rebinding_is_supported(): + parameter = torch.nn.Parameter(torch.ones(4, 4)) + optimizer = Gefen([("legacy.name", parameter)], fused=False) + identity = ParameterIdentity("Exact.FQN", (4, 4)) + + optimizer.rebind_parameter(parameter, parameter, identity=identity) + + assert optimizer.param_groups[0]["params"][0] is parameter + assert optimizer.state[parameter]["name"] == "legacy.name" + assert optimizer.parameter_identity(parameter).fqn == "Exact.FQN" + + +def test_replicated_rebinding_accepts_nonoverlapping_strided_storage(): + old = torch.nn.Parameter(torch.ones(4)) + target = torch.nn.Parameter(torch.arange(8, dtype=torch.float32)[::2]) + optimizer = Gefen([("weight", old)], fused=False, factored_v_2d=False) + + optimizer.rebind_parameter( + old, + target, + identity=ParameterIdentity("Weight", (4,)), + ) + + assert optimizer.param_groups[0]["params"] == [target] + assert optimizer.optimizer_contract().capabilities.stable_shard_identity + + +def test_rebinding_after_external_parameter_replacement_removes_orphan_state(): + old = torch.nn.Parameter(torch.ones(4, 4)) + new = torch.nn.Parameter(torch.full((4, 4), 2.0)) + optimizer = Gefen([("layer.weight", old)], fused=False) + optimizer.param_groups[0]["params"][0] = new + + optimizer.rebind_parameter( + old, + new, + identity=ParameterIdentity("layer.weight", (4, 4)), + ) + + assert old not in optimizer.state + assert old not in optimizer._param_names + assert optimizer.state[new] == {"name": "layer.weight"} + + +def test_external_replacement_reorder_preserves_each_source_compatibility_name(): + old_a = torch.nn.Parameter(torch.ones(4)) + old_b = torch.nn.Parameter(torch.ones(4) * 2) + new_a = torch.nn.Parameter(torch.ones(4) * 3) + new_b = torch.nn.Parameter(torch.ones(4) * 4) + optimizer = Gefen( + [("source.a", old_a), ("source.b", old_b)], + fused=False, + factored_v_2d=False, + ) + optimizer.param_groups[0]["params"][:] = [new_b, new_a] + shard_a = _replicated(new_a, "Canonical.A") + shard_b = _replicated(new_b, "Canonical.B") + + optimizer.post_sharding( + ( + ParameterRebinding(old_a, new_a, shard_a), + ParameterRebinding(old_b, new_b, shard_b), + ), + manifest=ShardingManifest((shard_b, shard_a)), + ) + + assert optimizer.param_groups[0]["params"] == [new_b, new_a] + assert optimizer.param_groups[0]["param_names"] == ["source.b", "source.a"] + assert optimizer.state[new_a]["name"] == "source.a" + assert optimizer.state[new_b]["name"] == "source.b" + + +def test_flattened_plain_gefen_rebind_steps_like_direct_local_shard(): + old = torch.nn.Parameter(torch.arange(16, dtype=torch.float32).reshape(4, 4)) + rebound = torch.nn.Parameter(torch.arange(8, dtype=torch.float32)) + oracle = torch.nn.Parameter(rebound.detach().clone()) + optimizer = Gefen( + [("layer.weight", old)], + lr=1e-3, + fused=False, + factored_v_2d=False, + ) + reference = Gefen( + [("layer.weight", oracle)], + lr=1e-3, + fused=False, + factored_v_2d=False, + ) + local, manifest = _flat_manifest("Model.Layer.Weight", (4, 4), (8, 8), "rank:0") + optimizer.post_sharding((ParameterRebinding(old, rebound, local),), manifest=manifest) + grad = torch.linspace(-1, 1, 8) + + for parameter, current in ((rebound, optimizer), (oracle, reference)): + parameter.grad = grad.clone() + current.step() + current.zero_grad() + + assert torch.equal(rebound, oracle) + _nested_equal(optimizer.state[rebound], reference.state[oracle]) + assert torch.equal(optimizer._gefen_codebook, reference._gefen_codebook) + + +def test_native_resume_preserves_target_flat_shard_identity_and_continuation(): + def build(): + old = torch.nn.Parameter(torch.zeros(4, 4)) + local_param = torch.nn.Parameter(torch.arange(8, dtype=torch.float32)) + optimizer = Gefen( + [("layer.weight", old)], + lr=1e-3, + fused=False, + factored_v_2d=False, + ) + local, manifest = _flat_manifest("Model.Layer.Weight", (4, 4), (8, 8), "rank:0") + optimizer.post_sharding((ParameterRebinding(old, local_param, local),), manifest=manifest) + return optimizer, local_param, local + + source, source_param, _ = build() + target, target_param, target_identity = build() + for grad in (torch.linspace(-1, 1, 8), torch.linspace(1, -1, 8)): + source_param.grad = grad.clone() + source.step() + source.zero_grad() + checkpoint = copy.deepcopy(source.state_dict()) + with torch.no_grad(): + target_param.copy_(source_param) + + target.load_state_dict(checkpoint) + + assert target.shard_identity(target_param) == target_identity + assert target.optimizer_contract().capabilities.stable_shard_identity + next_grad = torch.linspace(-0.5, 0.5, 8) + for parameter, optimizer in ( + (source_param, source), + (target_param, target), + ): + parameter.grad = next_grad.clone() + optimizer.step() + optimizer.zero_grad() + assert torch.equal(source_param, target_param) + _nested_equal(source.state_dict(), target.state_dict()) + + +def test_factored_logical_matrix_flattening_rejects_atomically(): + old = torch.nn.Parameter(torch.ones(4, 4)) + new = torch.nn.Parameter(torch.ones(8)) + optimizer = Gefen([("layer.weight", old)], fused=False, factored_v_2d=True) + local, manifest = _flat_manifest("layer.weight", (4, 4), (8, 8), "rank:0") + snapshot = _snapshot(optimizer) + + with pytest.raises(ValueError, match="factored_v_2d=False"): + optimizer.post_sharding((ParameterRebinding(old, new, local),), manifest=manifest) + + _assert_snapshot(optimizer, snapshot) + + +def test_late_batch_failure_is_atomic_and_continuation_is_bit_exact(): + def build(): + params = [ + torch.nn.Parameter(torch.arange(8, dtype=torch.float32)), + torch.nn.Parameter(torch.arange(8, dtype=torch.float32) + 1), + ] + return Gefen( + [("first", params[0]), ("second", params[1])], + fused=False, + factored_v_2d=False, + ), params + + target, target_params = build() + control, control_params = build() + replacements = [ + torch.nn.Parameter(torch.arange(8, dtype=torch.float32) + 2), + torch.nn.Parameter(torch.arange(7, dtype=torch.float32) + 3), + ] + first = _replicated(replacements[0], "First", (8,)) + second = _replicated(torch.empty(8), "Second", (8,)) + manifest = ShardingManifest((second, first)) + snapshot = _snapshot(target) + + with pytest.raises(ValueError, match="complete parameter storage"): + target.post_sharding( + ( + ParameterRebinding(target_params[0], replacements[0], first), + ParameterRebinding(target_params[1], replacements[1], second), + ), + manifest=manifest, + ) + + _assert_snapshot(target, snapshot) + grads = (torch.linspace(-1, 1, 8), torch.linspace(1, -1, 8)) + for params, optimizer in ((target_params, target), (control_params, control)): + for parameter, grad in zip(params, grads): + parameter.grad = grad.clone() + optimizer.step() + optimizer.zero_grad() + for target_param, control_param in zip(target_params, control_params): + assert torch.equal(target_param, control_param) + _nested_equal(target.state_dict(), control.state_dict()) + + +def test_checkpoint_schema_preparation_failure_is_atomic(monkeypatch): + old = torch.nn.Parameter(torch.ones(4)) + new = torch.nn.Parameter(torch.ones(4) * 2) + optimizer = Gefen([("weight", old)], fused=False, factored_v_2d=False) + shard = _replicated(new, "Weight") + snapshot = _snapshot(optimizer) + + def reject_schema(): + raise RuntimeError("schema preparation failed") + + monkeypatch.setattr( + optimizer, + "_install_rank_local_checkpoint_schema", + reject_schema, + ) + with pytest.raises(RuntimeError, match="schema preparation failed"): + optimizer.post_sharding( + (ParameterRebinding(old, new, shard),), + manifest=ShardingManifest((shard,)), + ) + + _assert_snapshot(optimizer, snapshot) + + +def test_whole_optimizer_pristine_guard_rejects_after_partial_initialization(): + params = [ + torch.nn.Parameter(torch.ones(8)), + torch.nn.Parameter(torch.ones(8) * 2), + ] + optimizer = Gefen( + [("first", params[0]), ("second", params[1])], + fused=False, + factored_v_2d=False, + ) + params[0].grad = torch.ones_like(params[0]) + optimizer.step() + optimizer.zero_grad() + replacements = [torch.nn.Parameter(param.detach().clone()) for param in params] + shards = ( + _replicated(replacements[0], "First"), + _replicated(replacements[1], "Second"), + ) + snapshot = _snapshot(optimizer) + + with pytest.raises(RuntimeError, match="before optimizer mutation"): + optimizer.post_sharding( + tuple(ParameterRebinding(old, new, shard) for old, new, shard in zip(params, replacements, shards)), + manifest=ShardingManifest(shards), + ) + + _assert_snapshot(optimizer, snapshot) + + +@pytest.mark.parametrize("failure", ["duplicate_target", "unknown_source", "gradient", "manifest"]) +def test_rebinding_structural_failures_are_no_ops(failure): + old = torch.nn.Parameter(torch.ones(4)) + new = torch.nn.Parameter(torch.ones(4) * 2) + optimizer = Gefen([("weight", old)], fused=False, factored_v_2d=False) + shard = _replicated(new, "Weight") + manifest = ShardingManifest((shard,)) + binding = ParameterRebinding(old, new, shard) + if failure == "duplicate_target": + other = torch.nn.Parameter(torch.ones(4)) + optimizer = Gefen( + [("weight", old), ("other", other)], + fused=False, + factored_v_2d=False, + ) + other_shard = _replicated(new, "Other") + manifest = ShardingManifest((shard, other_shard)) + bindings = ( + binding, + ParameterRebinding(other, new, other_shard), + ) + elif failure == "unknown_source": + bindings = (ParameterRebinding(torch.nn.Parameter(torch.ones(4)), new, shard),) + elif failure == "gradient": + new.grad = torch.ones_like(new) + bindings = (binding,) + else: + other_shard = _replicated(new, "Other") + manifest = ShardingManifest((other_shard,)) + bindings = (binding,) + snapshot = _snapshot(optimizer) + + with pytest.raises((RuntimeError, ValueError)): + optimizer.post_sharding(bindings, manifest=manifest) + + _assert_snapshot(optimizer, snapshot) + + +def test_rebinding_rejects_distinct_targets_with_overlapping_storage(): + old = [ + torch.nn.Parameter(torch.ones(8)), + torch.nn.Parameter(torch.ones(8) * 2), + ] + optimizer = Gefen( + [("first", old[0]), ("second", old[1])], + fused=False, + factored_v_2d=False, + ) + storage = torch.arange(12, dtype=torch.float32) + targets = [ + torch.nn.Parameter(storage[:8]), + torch.nn.Parameter(storage[4:]), + ] + shards = ( + _replicated(targets[0], "First"), + _replicated(targets[1], "Second"), + ) + snapshot = _snapshot(optimizer) + + with pytest.raises(ValueError, match="must not overlap"): + optimizer.post_sharding( + ( + ParameterRebinding(old[0], targets[0], shards[0]), + ParameterRebinding(old[1], targets[1], shards[1]), + ), + manifest=ShardingManifest(shards), + ) + + _assert_snapshot(optimizer, snapshot) + + +@pytest.mark.parametrize("target_kind", ["noncontiguous", "integer", "internal_overlap"]) +def test_rebinding_rejects_invalid_target_storage_atomically(target_kind): + old = torch.nn.Parameter(torch.ones(4)) + optimizer = Gefen([("weight", old)], fused=False, factored_v_2d=False) + if target_kind == "noncontiguous": + target = torch.nn.Parameter(torch.arange(8, dtype=torch.float32)[::2]) + local, manifest = _flat_manifest("Weight", (8,), (4, 4), "rank:0") + error = "contiguous physical storage" + elif target_kind == "integer": + target = torch.nn.Parameter(torch.ones(4, dtype=torch.int64), requires_grad=False) + local = _replicated(target, "Weight") + manifest = ShardingManifest((local,)) + error = "floating-point storage" + else: + target = torch.nn.Parameter(torch.ones(1).expand(4)) + local = _replicated(target, "Weight") + manifest = ShardingManifest((local,)) + error = "internal storage overlap" + snapshot = _snapshot(optimizer) + before = target.detach().clone() + + with pytest.raises(ValueError, match=error): + optimizer.post_sharding((ParameterRebinding(old, target, local),), manifest=manifest) + + _assert_snapshot(optimizer, snapshot) + assert torch.equal(target, before) + + +def test_rebinding_rejects_mixed_local_members_for_one_process_group(): + old = [torch.nn.Parameter(torch.ones(2)), torch.nn.Parameter(torch.ones(2))] + targets = [torch.nn.Parameter(torch.ones(2)), torch.nn.Parameter(torch.ones(2))] + optimizer = Gefen( + [("first", old[0]), ("second", old[1])], + fused=False, + factored_v_2d=False, + ) + first_local, first_manifest = _flat_manifest("First", (4,), (2, 2), "rank:0") + second_local, second_manifest = _flat_manifest("Second", (4,), (2, 2), "rank:1") + manifest = ShardingManifest(first_manifest.shards + second_manifest.shards) + snapshot = _snapshot(optimizer) + + with pytest.raises(ValueError, match="one member"): + optimizer.post_sharding( + ( + ParameterRebinding(old[0], targets[0], first_local), + ParameterRebinding(old[1], targets[1], second_local), + ), + manifest=manifest, + ) + + _assert_snapshot(optimizer, snapshot) + + +def test_external_replacement_rejects_duplicate_stale_sources(): + old = [torch.nn.Parameter(torch.ones(4)), torch.nn.Parameter(torch.ones(4))] + targets = [torch.nn.Parameter(torch.ones(4)), torch.nn.Parameter(torch.ones(4))] + optimizer = Gefen( + [("first", old[0]), ("second", old[1])], + fused=False, + factored_v_2d=False, + ) + optimizer.param_groups[0]["params"][:] = targets + shards = ( + _replicated(targets[0], "First"), + _replicated(targets[1], "Second"), + ) + snapshot = _snapshot(optimizer) + + with pytest.raises(ValueError, match="source tensors must be unique"): + optimizer.post_sharding( + ( + ParameterRebinding(old[0], targets[0], shards[0]), + ParameterRebinding(old[0], targets[1], shards[1]), + ), + manifest=ShardingManifest(shards), + ) + + _assert_snapshot(optimizer, snapshot) + + +@pytest.mark.parametrize("mutation", ["replace", "reorder"]) +def test_finalized_layout_guard_rejects_direct_parameter_group_mutation(mutation): + old = [torch.nn.Parameter(torch.ones(4)), torch.nn.Parameter(torch.ones(4) * 2)] + new = [torch.nn.Parameter(torch.ones(4) * 3), torch.nn.Parameter(torch.ones(4) * 4)] + optimizer = Gefen( + [("first", old[0]), ("second", old[1])], + fused=False, + factored_v_2d=False, + ) + shards = (_replicated(new[0], "First"), _replicated(new[1], "Second")) + optimizer.post_sharding( + ( + ParameterRebinding(old[0], new[0], shards[0]), + ParameterRebinding(old[1], new[1], shards[1]), + ), + manifest=ShardingManifest(shards), + ) + before = [parameter.detach().clone() for parameter in new] + if mutation == "replace": + rogue = torch.nn.Parameter(torch.ones(4) * 5) + optimizer.param_groups[0]["params"][0] = rogue + else: + optimizer.param_groups[0]["params"].reverse() + assert not optimizer.optimizer_contract().capabilities.stable_shard_identity + + with pytest.raises(RuntimeError, match="changed outside post_sharding"): + optimizer.step() + with pytest.raises(RuntimeError, match="changed outside post_sharding"): + optimizer.state_dict() + with pytest.raises(RuntimeError, match="changed outside post_sharding"): + optimizer.shard_bindings() + with pytest.raises(RuntimeError, match="changed outside post_sharding"): + optimizer.sharding_manifest() + for parameter, expected in zip(new, before): + assert torch.equal(parameter, expected) + + +@pytest.mark.parametrize("optimizer_type", ["plain", "muon"]) +def test_finalized_layout_guard_rechecks_after_closure(optimizer_type): + shape = (4,) if optimizer_type == "plain" else (2, 2) + old = torch.nn.Parameter(torch.ones(shape)) + bound = torch.nn.Parameter(torch.ones(shape) * 2) + rogue = torch.nn.Parameter(torch.ones(shape) * 3) + if optimizer_type == "plain": + optimizer = Gefen([("weight", old)], fused=False, factored_v_2d=False) + else: + optimizer = GefenMuon([("weight", old)], fused=False) + optimizer.rebind_parameter( + old, + bound, + identity=ParameterIdentity("Weight", shape), + ) + bound_before = bound.detach().clone() + rogue_before = rogue.detach().clone() + + def mutate_layout(): + optimizer.param_groups[0]["params"][0] = rogue + rogue.grad = torch.ones_like(rogue) + return torch.tensor(1.0) + + with pytest.raises(RuntimeError, match="changed outside post_sharding"): + optimizer.step(mutate_layout) + + assert torch.equal(bound, bound_before) + assert torch.equal(rogue, rogue_before) + assert optimizer.state[bound] == {"name": "weight"} + + +def test_finalized_layout_guard_rechecks_after_load_pre_hook(): + old = torch.nn.Parameter(torch.ones(4)) + bound = torch.nn.Parameter(torch.ones(4) * 2) + rogue = torch.nn.Parameter(torch.ones(4) * 3) + optimizer = Gefen([("weight", old)], fused=False, factored_v_2d=False) + optimizer.rebind_parameter( + old, + bound, + identity=ParameterIdentity("Weight", (4,)), + ) + checkpoint = copy.deepcopy(optimizer.state_dict()) + bound_before = bound.detach().clone() + rogue_before = rogue.detach().clone() + + def mutate_layout(current, state_dict): + current.param_groups[0]["params"][0] = rogue + + optimizer.register_load_state_dict_pre_hook(mutate_layout) + with pytest.raises(RuntimeError, match="changed outside post_sharding"): + optimizer.load_state_dict(checkpoint) + + assert torch.equal(bound, bound_before) + assert torch.equal(rogue, rogue_before) + assert optimizer.state[bound] == {"name": "weight"} + + +def test_muon_replicated_rebind_steps_like_direct_optimizer(): + old = torch.nn.Parameter(torch.arange(16, dtype=torch.float32).reshape(4, 4)) + rebound = torch.nn.Parameter(old.detach().clone()) + oracle = torch.nn.Parameter(old.detach().clone()) + optimizer = GefenMuon([("layer.weight", old)], fused=False) + reference = GefenMuon([("layer.weight", oracle)], fused=False) + identity = ParameterIdentity("Layer.Weight", (4, 4)) + optimizer.rebind_parameter(old, rebound, identity=identity) + grad = torch.linspace(-1, 1, 16).reshape(4, 4) + + for parameter, current in ((rebound, optimizer), (oracle, reference)): + parameter.grad = grad.clone() + current.step() + current.zero_grad() + + assert torch.equal(rebound, oracle) + _nested_equal(optimizer.state[rebound], reference.state[oracle]) + + +def test_muon_whole_owner_post_sharding_prunes_nonowner_and_blocks_training(): + old_owner = torch.nn.Parameter(torch.ones(4, 4)) + old_nonowner = torch.nn.Parameter(torch.ones(4, 4) * 2) + new_owner = torch.nn.Parameter(torch.ones(4, 4) * 3) + optimizer = GefenMuon( + [("owner.weight", old_owner), ("remote.weight", old_nonowner)], + fused=False, + ) + group = ProcessGroupIdentity("data_parallel", ("rank:0", "rank:1")) + owner_identity = ParameterIdentity("Owner.Weight", (4, 4)) + remote_identity = ParameterIdentity("Remote.Weight", (4, 4)) + owner_records = tuple(_owner_shard(owner_identity, group, member, "rank:0") for member in group.ordered_members) + remote_records = tuple(_owner_shard(remote_identity, group, member, "rank:1") for member in group.ordered_members) + manifest = ShardingManifest(owner_records + remote_records) + local_owner = next(item for item in owner_records if item.local_member == "rank:0") + local_nonowner = next(item for item in remote_records if item.local_member == "rank:0") + + optimizer.post_sharding( + ( + ParameterRebinding(old_nonowner, None, local_nonowner), + ParameterRebinding(old_owner, new_owner, local_owner), + ), + manifest=manifest, + ) + + assert optimizer.param_groups[0]["params"] == [new_owner] + assert set(optimizer.state) == {new_owner} + assert old_owner not in optimizer._param_names + assert old_nonowner not in optimizer._param_names + assert optimizer.shard_bindings() == tuple( + sorted( + ((new_owner, local_owner), (None, local_nonowner)), + key=lambda item: item[1].sort_key, + ) + ) + contract = optimizer.optimizer_contract() + assert contract.capabilities.canonical_parameter_fqns + assert contract.capabilities.stable_shard_identity + assert not contract.capabilities.explicit_process_group_codebook_scope + assert all( + support.layout is not ParameterLayout.WHOLE_PARAMETER_OWNER for support in contract.capabilities.training + ) + snapshot = copy.deepcopy(optimizer.state_dict()) + parameter_snapshot = new_owner.detach().clone() + closure_calls = [] + with pytest.raises(RuntimeError, match="codebook scope"): + optimizer.step(lambda: closure_calls.append(True)) + assert closure_calls == [] + assert torch.equal(new_owner, parameter_snapshot) + _nested_equal(optimizer.state_dict(), snapshot) + + +def test_muon_all_nonowner_post_sharding_retains_manifest_without_fake_state(): + old = torch.nn.Parameter(torch.ones(4, 4)) + optimizer = GefenMuon([("remote.weight", old)], fused=False) + group = ProcessGroupIdentity("data_parallel", ("rank:0", "rank:1")) + identity = ParameterIdentity("Remote.Weight", (4, 4)) + records = tuple(_owner_shard(identity, group, member, "rank:1") for member in group.ordered_members) + local = next(item for item in records if item.local_member == "rank:0") + manifest = ShardingManifest(records) + + optimizer.post_sharding((ParameterRebinding(old, None, local),), manifest=manifest) + + assert optimizer.param_groups[0]["params"] == [] + assert optimizer.param_groups[0]["param_names"] == [] + assert optimizer.state == {} + assert optimizer.shard_bindings() == ((None, local),) + assert optimizer.sharding_manifest() == manifest + assert optimizer.optimizer_contract().capabilities.stable_shard_identity + + +def test_muon_flattened_rebinding_rejects_without_mutation(): + old = torch.nn.Parameter(torch.ones(4, 4)) + new = torch.nn.Parameter(torch.ones(8)) + optimizer = GefenMuon([("layer.weight", old)], fused=False) + local, manifest = _flat_manifest("Layer.Weight", (4, 4), (8, 8), "rank:0") + snapshot = _snapshot(optimizer) + + with pytest.raises(ValueError, match="complete matrices"): + optimizer.post_sharding((ParameterRebinding(old, new, local),), manifest=manifest) + + _assert_snapshot(optimizer, snapshot) + + +def test_hybrid_contract_does_not_claim_composite_rebinding(): + matrix = torch.nn.Parameter(torch.ones(4, 4)) + bias = torch.nn.Parameter(torch.ones(4)) + optimizer = GefenMuonHybrid( + [("layer.weight", matrix)], + [("layer.bias", bias)], + lr=1e-3, + fused=False, + ) + contract = optimizer.optimizer_contract() + assert not contract.capabilities.shard_rebinding + assert not contract.capabilities.post_sharding + assert not contract.capabilities.canonical_parameter_fqns diff --git a/tests/test_shard_identity_contracts.py b/tests/test_shard_identity_contracts.py index e3c1823..bc3eb1f 100644 --- a/tests/test_shard_identity_contracts.py +++ b/tests/test_shard_identity_contracts.py @@ -470,3 +470,5 @@ def test_identity_contracts_are_public_lazy_exports(): ): assert name in gefen.__all__ assert getattr(gefen, name).__module__ == "gefen.contracts" + assert "ParameterRebinding" in gefen.__all__ + assert gefen.ParameterRebinding.__module__ == "gefen.rebinding" From 750905a0397a07384a0d21ce3e4fc25d0296d343 Mon Sep 17 00:00:00 2001 From: thad0ctor Date: Sun, 12 Jul 2026 20:44:11 -0700 Subject: [PATCH 05/52] Add explicit codebook process-group scopes --- docs/optimizer_contracts.md | 14 +- src/gefen/__init__.py | 5 + src/gefen/codebook.py | 46 + src/gefen/contracts.py | 95 +- src/gefen/gefen.py | 1334 ++++++++++++++++++++-- src/gefen/gefen_muon.py | 133 ++- tests/test_codebook_scope_cpu.py | 508 ++++++++ tests/test_codebook_scope_distributed.py | 920 +++++++++++++++ tests/test_optimizer_contracts.py | 27 +- tests/test_rebinding_cpu.py | 4 +- 10 files changed, 2921 insertions(+), 165 deletions(-) create mode 100644 src/gefen/codebook.py create mode 100644 tests/test_codebook_scope_cpu.py create mode 100644 tests/test_codebook_scope_distributed.py diff --git a/docs/optimizer_contracts.md b/docs/optimizer_contracts.md index cf7b358..7b254b5 100644 --- a/docs/optimizer_contracts.md +++ b/docs/optimizer_contracts.md @@ -39,7 +39,19 @@ 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, but stepping a whole-owner binding remains explicitly disabled until the independent adapter-defined process-group codebook scope is implemented. Whole-owner training, DTensor stable identity, Hybrid composite rebinding, canonical checkpoint I/O, state movement, and offload therefore remain unclaimed. +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. DTensor stable identity, Hybrid composite rebinding, canonical checkpoint I/O, state movement, and offload remain unclaimed. + +## Explicit learned-codebook process groups + +`CodebookProcessGroupBinding` maps one stable `ProcessGroupIdentity` and local semantic member to an opaque PyTorch process-group handle plus an explicit collective device. It is accepted only through the complete `post_sharding(..., codebook_process_group=...)` transaction: every manifest shard must use that one semantic group, each local shard must name the binding's local member, the runtime group size and coordinate must match `ordered_members`, and the backend must support the supplied device. A one-member scope uses `process_group=None`; multi-member scopes must pass a real handle, including `dist.group.WORLD` when the default world is intentionally the semantic scope. Gefen never treats `None` as an implicit default-world selection. + +The `explicit_process_group_codebook_scope` capability reports that Gefen or GefenMuon implements this API; `codebook_process_group_binding()` separately reports whether a particular finalized instance has an active binding. GefenMuonHybrid remains negative at the composite level because its independent children do not form one atomically coordinated codebook scope. + +The optimizer owns one learned codebook and therefore accepts one scope. Histogram accounting represents each logical parameter once: the first ordered member contributes a replicated parameter after all members agree on gradient presence and its automatic period, every flattened shard contributes its local logical slice after all nonempty slices agree on gradient presence, and only the declared whole-parameter owner contributes an owned matrix. Local inputs are visited in canonical manifest order. Members first stage periods and an integer histogram without touching live optimizer state, exchange operation, step, scope, manifest, active-slice, period, policy, and old-codebook controls, sum the fixed-size `int64` histogram in the supplied group, solve the same exact-DP problem, verify codebook agreement, and only then publish periods and the codebook. Refresh additionally stages every replacement momentum-index tensor in canonical order and exchanges readiness before replacing indices or invalidating derived codebook/LUT caches. A local preparation, solve, or requantization failure is reported by every participant before optimizer-state commit; this is not rollback after process death or a collective-backend failure during the final commit window. + +`initialize_codebook()` and `refresh_codebook()` expose these operations for adapters that enter their optimizers in a deterministic order; `binding.sort_key` supplies the stable process-group portion of that schedule. Normal `step()` still initializes automatically and plain Gefen still honors `codebook_refresh_every`. Every scoped step exchanges a common operation header before any rank-dependent branch, and the first step after binding or native load additionally verifies codebook bytes and the complete manifest. Scoped native AMP requires every member to select the same protocol and present identical `found_inf` and `grad_scale` values; a mismatch raises collectively and requires a group-aware gradient scaler rather than changing external scaler state behind its back. Multi-member explicit scopes reject `capturable=True` because their host validation and process-group collectives are not CUDA-graph-safe. A one-member local scope may initialize during eager warmup and then use ordinary capturable stepping, but manual codebook replacement remains rejected. Ordinary unscoped behavior remains collective-free. Explicit scope does not replace DTensor mesh collectives, AMP mesh preflights, Parallel-Muon ownership collectives, or checkpoint transport groups. + +Native checkpoints store a primitive rank-neutral scope record—format version, semantic name, ordered members, and refresh schedule—at the top level and in the existing transport mirror. Scoped group metadata uses a new outer format version so an older loader rejects rather than silently discarding the scope. Runtime process-group handles and collective devices are live adapter configuration and are never serialized or reconstructed. Flattened and whole-owner local payloads additionally carry a primitive rank-local shard-identity guard ordered by native parameter group and slot, including replicated slots in a mixed optimizer and separately listed pruned nonowners; this prevents an equal-shaped checkpoint from another member, logical slice, or parameter ordering from being reinterpreted positionally. 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. Whole-owner checkpoint completeness, scoped DTensor rank-local transport, scope migration, topology-changing canonical I/O, and Hybrid-wide coordination are not claimed. 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. diff --git a/src/gefen/__init__.py b/src/gefen/__init__.py index 3401c6c..adabf3f 100644 --- a/src/gefen/__init__.py +++ b/src/gefen/__init__.py @@ -15,6 +15,7 @@ "IDENTITY_SCHEMA_VERSION", "CheckpointSupport", "CheckpointTransport", + "CodebookProcessGroupBinding", "OptimizerCapabilities", "OptimizerChildContract", "OptimizerContract", @@ -68,6 +69,10 @@ def __getattr__(name): from .rebinding import ParameterRebinding return ParameterRebinding + if name == "CodebookProcessGroupBinding": + from .codebook import CodebookProcessGroupBinding + + return CodebookProcessGroupBinding if name in ( "CONTRACT_SCHEMA_VERSION", "IDENTITY_SCHEMA_VERSION", diff --git a/src/gefen/codebook.py b/src/gefen/codebook.py new file mode 100644 index 0000000..7dced8c --- /dev/null +++ b/src/gefen/codebook.py @@ -0,0 +1,46 @@ +"""Runtime binding for one optimizer-wide learned-codebook process group.""" + +from dataclasses import dataclass +from typing import Optional + +import torch + +from gefen.contracts import ProcessGroupIdentity + + +@dataclass(frozen=True, eq=False) +class CodebookProcessGroupBinding: + """Bind a stable semantic group to one framework runtime group handle. + + ``process_group`` is deliberately opaque to the descriptor. Gefen validates + and consumes it through PyTorch's public distributed APIs when the binding + is installed by ``post_sharding``. A one-member scope uses ``None`` rather + than implicitly selecting the default world. + """ + + identity: ProcessGroupIdentity + local_member: str + process_group: Optional[object] + collective_device: torch.device + + def __post_init__(self) -> None: + if not isinstance(self.identity, ProcessGroupIdentity): + raise TypeError("CodebookProcessGroupBinding.identity must be a ProcessGroupIdentity") + if self.local_member not in self.identity.ordered_members: + raise ValueError("CodebookProcessGroupBinding.local_member must belong to the identity") + try: + device = torch.device(self.collective_device) + except (TypeError, RuntimeError) as exc: + raise TypeError("CodebookProcessGroupBinding.collective_device must be a torch device") from exc + if device.type == "meta": + raise ValueError("codebook collectives require a materialized device") + object.__setattr__(self, "collective_device", device) + + @property + def sort_key(self): + """Return the stable adapter scheduling key for this collective scope.""" + + return (self.identity.semantic_name, self.identity.ordered_members) + + +__all__ = ["CodebookProcessGroupBinding"] diff --git a/src/gefen/contracts.py b/src/gefen/contracts.py index e3037f3..f5a0cea 100644 --- a/src/gefen/contracts.py +++ b/src/gefen/contracts.py @@ -468,6 +468,7 @@ class StateField: key_match: StateKeyMatch = StateKeyMatch.EXACT applicable_sharded_modes: AbstractSet[str] = frozenset() description: str = "" + optional: bool = False def __post_init__(self) -> None: object.__setattr__( @@ -483,6 +484,8 @@ def __post_init__(self) -> None: raise TypeError("StateField.geometry must be a StateGeometry") if not isinstance(self.key_match, StateKeyMatch): raise TypeError("StateField.key_match must be a StateKeyMatch") + if type(self.optional) is not bool: + raise TypeError("StateField.optional must be a bool") @property def authoritative(self) -> bool: @@ -607,6 +610,7 @@ class TrainingSupport: sharded_mode: Optional[str] = None requires_complete_parameter_storage: bool = False requires_complete_logical_matrix: bool = False + requires_post_step_parameter_sync: bool = False def __post_init__(self) -> None: if self.mesh_dimensions is not None: @@ -620,6 +624,13 @@ def __post_init__(self) -> None: raise TypeError( "TrainingSupport.process_group_scope must be a ProcessGroupScope" ) + for name in ( + "requires_complete_parameter_storage", + "requires_complete_logical_matrix", + "requires_post_step_parameter_sync", + ): + if type(getattr(self, name)) is not bool: + raise TypeError("TrainingSupport.{} must be a bool".format(name)) @dataclass(frozen=True) @@ -783,6 +794,14 @@ def _common_fields() -> Tuple[StateField, ...]: True, description="Checkpoint-bound deterministic execution policy.", ), + StateField( + "gefen_codebook_scope", + StateScope.OPTIMIZER_COMMON, + StateGeometry.OPAQUE, + True, + description="Stable adapter-defined learned-codebook process-group scope.", + optional=True, + ), ) @@ -828,6 +847,20 @@ def _derived_fields() -> Tuple[StateField, ...]: StateGeometry.OPAQUE, False, ), + StateField( + "_gefen_codebook_process_group", + StateScope.DERIVED, + StateGeometry.OPAQUE, + False, + description="Live runtime handle for the adapter-defined codebook scope.", + ), + StateField( + "_gefen_codebook_scope_validated", + StateScope.DERIVED, + StateGeometry.SCALAR, + False, + description="Rebuildable cross-member scope-agreement cache.", + ), StateField("_sr_seed_by_device", StateScope.DERIVED, StateGeometry.SCALAR, False), StateField( "_gefen_global_step_by_device", @@ -858,6 +891,13 @@ def _derived_fields() -> Tuple[StateField, ...]: True, description="PyTorch transport mirror of optimizer-common state.", ), + StateField( + "gefen_native_local_shards", + StateScope.DERIVED, + StateGeometry.OPAQUE, + True, + description="Native rank-local identity guard for scoped physical shards.", + ), ) @@ -879,6 +919,7 @@ def _negative_capabilities( supported_parameter_ranks: Optional[Tuple[int, ...]], canonical_parameter_fqns: bool = False, stable_shard_identity: bool = False, + explicit_process_group_codebook_scope: bool = False, shard_rebinding: bool = False, post_sharding: bool = False, ) -> OptimizerCapabilities: @@ -890,7 +931,7 @@ def _negative_capabilities( accepts_semantic_parameter_names=True, canonical_parameter_fqns=canonical_parameter_fqns, stable_shard_identity=stable_shard_identity, - explicit_process_group_codebook_scope=False, + explicit_process_group_codebook_scope=explicit_process_group_codebook_scope, shard_rebinding=shard_rebinding, post_sharding=post_sharding, canonical_state_io=False, @@ -904,6 +945,8 @@ def _gefen_contract( factored_v_2d: bool, canonical_parameter_fqns: bool = False, stable_shard_identity: bool = False, + explicit_process_group_codebook_scope: bool = False, + native_flattened_checkpoint: bool = False, ) -> OptimizerContract: block_fields = ( StateField("vmean", StateScope.PARAMETER, StateGeometry.BLOCK, True), @@ -1052,11 +1095,11 @@ def _gefen_contract( checkpoints = ( CheckpointSupport( CheckpointTransport.NATIVE_OPTIMIZER, - frozenset( - { - ParameterLayout.REPLICATED, - ParameterLayout.FLATTENED_ELEMENT_SHARD, - } + frozenset({ParameterLayout.REPLICATED}) + | ( + frozenset({ParameterLayout.FLATTENED_ELEMENT_SHARD}) + if native_flattened_checkpoint + else frozenset() ), frozenset(), ProcessGroupScope.NONE, @@ -1081,6 +1124,7 @@ def _gefen_contract( supported_parameter_ranks=None, canonical_parameter_fqns=canonical_parameter_fqns, stable_shard_identity=stable_shard_identity, + explicit_process_group_codebook_scope=explicit_process_group_codebook_scope, shard_rebinding=True, post_sharding=True, ), @@ -1114,6 +1158,8 @@ def _gefen_muon_contract( non_normuon_modes: FrozenSet[str], canonical_parameter_fqns: bool = False, stable_shard_identity: bool = False, + explicit_process_group_codebook_scope: bool = False, + whole_parameter_owner: bool = False, ) -> OptimizerContract: sharded_modes = _frozenset(sharded_modes) normuon_modes = _frozenset(normuon_modes) @@ -1178,6 +1224,19 @@ def _gefen_muon_contract( sharded_mode=mode, ) ) + if whole_parameter_owner: + variants.append( + StateVariant( + "name_only_whole_owner_" + mode, + ("name",), + frozenset({ParameterLayout.WHOLE_PARAMETER_OWNER}), + StateExtent.METADATA_ONLY, + role=ParameterStateRole.OWNER, + initialized=False, + parameter_ranks=(2,), + sharded_mode=mode, + ) + ) for mode_set, field_names, prefix in ( (non_normuon_modes, base_names, "quantized_muon"), (normuon_modes, normuon_names, "quantized_normuon"), @@ -1194,6 +1253,17 @@ def _gefen_muon_contract( mode=mode, ) ) + if whole_parameter_owner: + variants.append( + _muon_initialized_variant( + name=prefix + "_whole_owner_" + mode, + field_names=field_names, + layout=ParameterLayout.WHOLE_PARAMETER_OWNER, + extent=StateExtent.OWNER_PARAMETER, + mode=mode, + role=ParameterStateRole.OWNER, + ) + ) if "approx" in mode_set: variants.append( _muon_initialized_variant( @@ -1270,6 +1340,18 @@ def _gefen_muon_contract( ), ) ) + if whole_parameter_owner: + training += tuple( + TrainingSupport( + ParameterLayout.WHOLE_PARAMETER_OWNER, + ProcessGroupScope.ADAPTER_DEFINED, + sharded_mode=mode, + requires_complete_parameter_storage=True, + requires_complete_logical_matrix=True, + requires_post_step_parameter_sync=True, + ) + for mode in sorted(sharded_modes) + ) checkpoints = [ CheckpointSupport( CheckpointTransport.NATIVE_OPTIMIZER, @@ -1321,6 +1403,7 @@ def _gefen_muon_contract( supported_parameter_ranks=(2,), canonical_parameter_fqns=canonical_parameter_fqns, stable_shard_identity=stable_shard_identity, + explicit_process_group_codebook_scope=explicit_process_group_codebook_scope, shard_rebinding=True, post_sharding=True, ), diff --git a/src/gefen/gefen.py b/src/gefen/gefen.py index 7d1fa62..4cc3246 100644 --- a/src/gefen/gefen.py +++ b/src/gefen/gefen.py @@ -10,6 +10,7 @@ helpers it steps through; ``gefen_muon``/``hybrid`` build on it. """ +import hashlib import io import logging import math @@ -22,12 +23,16 @@ import torch import torch.nn as nn +from gefen.codebook import CodebookProcessGroupBinding from gefen.contracts import ( LogicalSlice, OptimizerContract, ParameterIdentity, ParameterLayout, + PlacementKind, + ProcessGroupIdentity, ShardIdentity, + ShardPlacement, ShardingManifest, _gefen_contract, ) @@ -65,6 +70,9 @@ _RANK_LOCAL_MEMBER_KEY = "_gefen_rank_local_member" _RANK_LOCAL_FORMAT = "rank_local_dtensor_v2" _RANK_LOCAL_METADATA_VERSION = 3 +_CODEBOOK_SCOPE_FORMAT_VERSION = 1 +_SCOPED_NATIVE_METADATA_VERSION = 4 +_NATIVE_LOCAL_SHARDS_FORMAT_VERSION = 1 def _rank_local_payload_key(global_rank: int) -> str: @@ -390,24 +398,15 @@ def gefen_automatic_fused_update( ) -def learn_gefen_exact_codebook_from_grad_periods( - *, - grad_periods, - codebook_device: torch.device, - num_codebooks: int, - force_endpoints: bool, - verbose: bool, - compute_mse_logging: bool, - use_fused_histogram: bool = FUSE_HISTOGRAM_FOR_EXACT, -) -> Optional[torch.Tensor]: +def _gefen_exact_histogram_from_grad_periods( + *, grad_periods, num_codebooks: int, use_fused_histogram: bool +) -> torch.Tensor: histogram_bins = num_codebooks * 16 bin_width = 2.0 / float(histogram_bins) # Accumulate counts as int64 and cast to float32 once at assembly: repeated # fp32 adds round for bins past 2^24 counts, so chunked accumulation could # otherwise drift from the whole-tensor form's single cast. bin_counts_cpu = torch.zeros(histogram_bins, dtype=torch.int64, device="cpu") - total_numel = 0 - prev_mode = torch.get_deterministic_debug_mode() torch.set_deterministic_debug_mode(0) try: @@ -426,7 +425,6 @@ def learn_gefen_exact_codebook_from_grad_periods( ) gefen_exact_histogram_cuda(flat_float, period, local_counts_cuda) bin_counts_cpu.add_(local_counts_cuda.cpu()) - total_numel += flat_float.numel() else: blocks = automatic_partition_view(flat_float, period) absmax = blocks.abs().amax(dim=1, keepdim=True) @@ -442,9 +440,33 @@ def learn_gefen_exact_codebook_from_grad_periods( bin_indices, minlength=histogram_bins ) bin_counts_cpu.add_(local_counts.cpu()) - total_numel += normalized_flat.numel() finally: torch.set_deterministic_debug_mode(prev_mode) + return bin_counts_cpu + + +def _learn_gefen_exact_codebook_from_histogram( + *, + bin_counts_cpu: torch.Tensor, + codebook_device: torch.device, + num_codebooks: int, + force_endpoints: bool, + verbose: bool, + compute_mse_logging: bool, +) -> Optional[torch.Tensor]: + histogram_bins = num_codebooks * 16 + if ( + not torch.is_tensor(bin_counts_cpu) + or bin_counts_cpu.device.type != "cpu" + or bin_counts_cpu.dtype != torch.int64 + or tuple(bin_counts_cpu.shape) != (histogram_bins,) + ): + raise ValueError( + "Gefen exact-codebook histogram must be a CPU int64 vector with {} bins".format( + histogram_bins + ) + ) + total_numel = int(bin_counts_cpu.sum().item()) if total_numel == 0: return None @@ -493,6 +515,33 @@ def learn_gefen_exact_codebook_from_grad_periods( return codebook +def learn_gefen_exact_codebook_from_grad_periods( + *, + grad_periods, + codebook_device: torch.device, + num_codebooks: int, + force_endpoints: bool, + verbose: bool, + compute_mse_logging: bool, + use_fused_histogram: bool = FUSE_HISTOGRAM_FOR_EXACT, +) -> Optional[torch.Tensor]: + """Learn one exact-DP codebook from an unscoped local gradient stream.""" + + histogram = _gefen_exact_histogram_from_grad_periods( + grad_periods=grad_periods, + num_codebooks=num_codebooks, + use_fused_histogram=use_fused_histogram, + ) + return _learn_gefen_exact_codebook_from_histogram( + bin_counts_cpu=histogram, + codebook_device=codebook_device, + num_codebooks=num_codebooks, + force_endpoints=force_endpoints, + verbose=verbose, + compute_mse_logging=compute_mse_logging, + ) + + def _resolve_find_period_backend(grad: torch.Tensor) -> str: # An explicit FIND_PERIOD_BACKEND (set by config/tests) is an override and # wins. Otherwise resolve PER CALL from the tensor's own device. Do NOT cache @@ -1024,6 +1073,12 @@ def __init__( # cudagraph-trees recording rejects whenever the codebook input gets # copied into the graph pool). Cleared with _gefen_codebook_by_device. self._gefen_codebook_lut_by_device = {} + # An explicit learned-codebook scope is installed only as part of the + # same full post_sharding transaction that publishes canonical shard + # identities. It is live runtime configuration, never inferred from a + # default process group or from tensor storage. + self._gefen_codebook_process_group = None + self._gefen_codebook_scope_validated = False # Capturable stochastic rounding: ONE 0-dim int64 device tensor per # device holding the per-step rounding seed (a device-side mirror of # _gefen_global_step). Created lazily by _sr_seed_on and advanced on @@ -1084,6 +1139,11 @@ def _step_supports_amp_scaling(self) -> bool: # FP32-master AMP. Native handling is needed for active local FP16 # gradients, or statically for distributed optimizers that combine any # true-FP16 storage with DTensors after a collective presence preflight. + # An explicit codebook scope is also a topology property: whole owners + # and empty nonowners must enter the same GradScaler protocol before + # Gefen can compare found_inf/grad_scale collectively inside step(). + if self._gefen_codebook_process_group is not None: + return True return _amp_native_scaling_required(self) def optimizer_contract(self) -> OptimizerContract: @@ -1094,6 +1154,20 @@ def optimizer_contract(self) -> OptimizerContract: factored_v_2d=self._factored_v_2d, canonical_parameter_fqns=identity_ready, stable_shard_identity=identity_ready, + explicit_process_group_codebook_scope=True, + native_flattened_checkpoint=( + self._codebook_scope_ready() + and any( + shard.layout is ParameterLayout.FLATTENED_ELEMENT_SHARD + for _, shard in self._gefen_local_shard_bindings + ) + ), + ) + + def _codebook_scope_ready(self) -> bool: + return ( + self._gefen_codebook_process_group is not None + and self._canonical_identity_ready() ) def _canonical_identity_ready(self) -> bool: @@ -1164,6 +1238,273 @@ def sharding_manifest(self): self._assert_finalized_binding_layout() return self._gefen_sharding_manifest + def codebook_process_group_binding(self): + """Return the finalized explicit learned-codebook binding, or ``None``.""" + + self._assert_finalized_binding_layout() + return self._gefen_codebook_process_group + + def _serialized_codebook_scope(self): + binding = self._gefen_codebook_process_group + if binding is None: + return None + return { + "format_version": _CODEBOOK_SCOPE_FORMAT_VERSION, + "semantic_name": binding.identity.semantic_name, + "ordered_members": list(binding.identity.ordered_members), + "refresh_every": self._codebook_refresh_every, + } + + @staticmethod + def _normalize_serialized_codebook_scope(value): + if value is None: + return None + if not isinstance(value, dict) or set(value) != { + "format_version", + "semantic_name", + "ordered_members", + "refresh_every", + }: + raise ValueError("Gefen checkpoint codebook scope has an invalid schema") + if ( + type(value["format_version"]) is not int + or value["format_version"] != _CODEBOOK_SCOPE_FORMAT_VERSION + ): + raise ValueError( + "Gefen checkpoint codebook scope has an unsupported format version" + ) + if not isinstance(value["ordered_members"], list): + raise ValueError( + "Gefen checkpoint codebook scope ordered_members must be a list" + ) + from gefen.contracts import ProcessGroupIdentity + + try: + identity = ProcessGroupIdentity( + value["semantic_name"], tuple(value["ordered_members"]) + ) + except (TypeError, ValueError) as exc: + raise ValueError( + "Gefen checkpoint codebook scope identity is invalid" + ) from exc + refresh_every = value["refresh_every"] + if type(refresh_every) is not int or refresh_every < 0: + raise ValueError( + "Gefen checkpoint codebook scope refresh_every must be a nonnegative integer" + ) + return { + "format_version": _CODEBOOK_SCOPE_FORMAT_VERSION, + "semantic_name": identity.semantic_name, + "ordered_members": list(identity.ordered_members), + "refresh_every": refresh_every, + } + + @staticmethod + def _serialized_native_local_shard(shard): + return { + "fqn": shard.parameter.fqn, + "global_shape": list(shard.parameter.global_shape), + "layout": shard.layout.value, + "flat_offset": shard.logical_slice.flat_offset, + "length": shard.logical_slice.length, + "process_group": { + "semantic_name": shard.process_group.semantic_name, + "ordered_members": list(shard.process_group.ordered_members), + }, + "local_member": shard.local_member, + "owner": shard.owner, + "placements": [ + { + "mesh_axis": placement.mesh_axis, + "kind": placement.kind.value, + "coordinate": placement.coordinate, + "parts": placement.parts, + "parameter_dimension": placement.parameter_dimension, + } + for placement in shard.placements + ], + } + + def _serialized_native_local_shards(self): + if self._gefen_codebook_process_group is None: + return None + if not any( + shard.layout is not ParameterLayout.REPLICATED + for _, shard in self._gefen_local_shard_bindings + ): + return None + + # Native Optimizer.load_state_dict maps per-parameter state by + # parameter-group and slot position, not by canonical FQN. Bind every + # serialized live slot to its shard identity (including replicated + # slots in a mixed optimizer), and retain pruned whole-owner records + # separately. A canonical shard set alone would not distinguish A/B + # from B/A when equal-shaped parameters exchange positions. + param_groups = [] + for group in self.param_groups: + param_groups.append( + [ + self._serialized_native_local_shard( + self._gefen_shard_bindings[parameter] + ) + for parameter in group["params"] + ] + ) + pruned_shards = [ + self._serialized_native_local_shard(shard) + for parameter, shard in self._gefen_local_shard_bindings + if parameter is None + ] + return { + "format_version": _NATIVE_LOCAL_SHARDS_FORMAT_VERSION, + "param_groups": param_groups, + "pruned_shards": pruned_shards, + } + + @classmethod + def _normalize_serialized_native_local_shard(cls, record): + expected = { + "fqn", + "global_shape", + "layout", + "flat_offset", + "length", + "process_group", + "local_member", + "owner", + "placements", + } + if not isinstance(record, dict) or set(record) != expected: + raise ValueError( + "Gefen native local-shard metadata has an invalid schema" + ) + group_record = record["process_group"] + if not isinstance(group_record, dict) or set(group_record) != { + "semantic_name", + "ordered_members", + }: + raise ValueError( + "Gefen native local-shard process-group metadata is invalid" + ) + if not isinstance(group_record["ordered_members"], list): + raise ValueError( + "Gefen native local-shard ordered_members must be a list" + ) + if not isinstance(record["global_shape"], list) or not isinstance( + record["placements"], list + ): + raise ValueError( + "Gefen native local-shard shapes and placements must be lists" + ) + try: + parameter = ParameterIdentity( + record["fqn"], tuple(record["global_shape"]) + ) + group = ProcessGroupIdentity( + group_record["semantic_name"], + tuple(group_record["ordered_members"]), + ) + placements = [] + for placement_record in record["placements"]: + if not isinstance(placement_record, dict) or set( + placement_record + ) != { + "mesh_axis", + "kind", + "coordinate", + "parts", + "parameter_dimension", + }: + raise ValueError("invalid placement schema") + placements.append( + ShardPlacement( + placement_record["mesh_axis"], + PlacementKind(placement_record["kind"]), + placement_record["coordinate"], + placement_record["parts"], + placement_record["parameter_dimension"], + ) + ) + shard = ShardIdentity( + parameter, + ParameterLayout(record["layout"]), + LogicalSlice(record["flat_offset"], record["length"]), + process_group=group, + local_member=record["local_member"], + owner=record["owner"], + placements=tuple(placements), + ) + except (TypeError, ValueError) as exc: + raise ValueError( + "Gefen native local-shard metadata is invalid" + ) from exc + if shard.layout not in { + ParameterLayout.REPLICATED, + ParameterLayout.FLATTENED_ELEMENT_SHARD, + ParameterLayout.WHOLE_PARAMETER_OWNER, + }: + raise ValueError( + "Gefen native local-shard metadata has an unsupported layout" + ) + return cls._serialized_native_local_shard(shard) + + @classmethod + def _normalize_serialized_native_local_shards(cls, value): + if value is None: + return None + if not isinstance(value, dict) or set(value) != { + "format_version", + "param_groups", + "pruned_shards", + }: + raise ValueError( + "Gefen native local-shard metadata has an invalid schema" + ) + format_version = value["format_version"] + if ( + type(format_version) is not int + or format_version != _NATIVE_LOCAL_SHARDS_FORMAT_VERSION + ): + raise ValueError( + "Unsupported Gefen native local-shard format_version: {}".format( + format_version + ) + ) + param_groups = value["param_groups"] + pruned_shards = value["pruned_shards"] + if not isinstance(param_groups, list) or not isinstance(pruned_shards, list): + raise ValueError( + "Gefen native local-shard parameter groups and pruned shards must be lists" + ) + normalized_groups = [] + for group in param_groups: + if not isinstance(group, list): + raise ValueError( + "Gefen native local-shard parameter groups must contain lists" + ) + normalized_groups.append( + [ + cls._normalize_serialized_native_local_shard(record) + for record in group + ] + ) + normalized_pruned = [ + cls._normalize_serialized_native_local_shard(record) + for record in pruned_shards + ] + if any( + record["layout"] != ParameterLayout.WHOLE_PARAMETER_OWNER.value + for record in normalized_pruned + ): + raise ValueError( + "Gefen native pruned-shard metadata must describe whole-parameter nonowners" + ) + return { + "format_version": _NATIVE_LOCAL_SHARDS_FORMAT_VERSION, + "param_groups": normalized_groups, + "pruned_shards": normalized_pruned, + } + @staticmethod def _parameter_in(parameters, candidate) -> bool: return any(item is candidate for item in parameters) @@ -1331,7 +1672,127 @@ def _validate_rebinding_layout(self, rebinding: ParameterRebinding) -> None: "Gefen rebound tensor numel does not match its logical slice" ) - def _stage_post_sharding(self, rebindings, manifest): + @staticmethod + def _codebook_scope_backend_device_is_valid(backend: str, device) -> bool: + backend = str(backend).lower() + if "nccl" in backend: + return device.type == "cuda" and device.index is not None + if "gloo" in backend or "mpi" in backend: + return device.type == "cpu" + return device.type in {"cpu", "cuda"} + + def _validate_codebook_process_group_binding( + self, binding, rebindings, manifest + ) -> None: + if not isinstance(binding, CodebookProcessGroupBinding): + raise TypeError( + "codebook_process_group must be a CodebookProcessGroupBinding" + ) + identity = binding.identity + if any(shard.process_group != identity for shard in manifest.shards): + raise ValueError( + "every scoped manifest shard must use the codebook process-group identity" + ) + for rebinding in rebindings: + shard = rebinding.shard + if shard.local_member != binding.local_member: + raise ValueError( + "local shard members must match the codebook binding member" + ) + if shard.layout not in { + ParameterLayout.REPLICATED, + ParameterLayout.FLATTENED_ELEMENT_SHARD, + ParameterLayout.WHOLE_PARAMETER_OWNER, + }: + raise ValueError( + "explicit codebook scope supports replicated, flattened, or " + "whole-parameter owner identities" + ) + + self._validate_codebook_runtime_binding(binding) + + def _validate_codebook_runtime_binding(self, binding) -> None: + members = binding.identity.ordered_members + if len(members) == 1: + if binding.process_group is not None: + raise ValueError( + "a one-member codebook scope must use process_group=None" + ) + return + + if self.capturable: + raise ValueError( + "capturable=True does not support a multi-member explicit codebook scope" + ) + + if binding.process_group is None: + raise ValueError( + "a multi-member codebook scope requires an explicit runtime process group" + ) + if not torch.distributed.is_available() or not torch.distributed.is_initialized(): + raise RuntimeError( + "a multi-member codebook scope requires initialized torch.distributed" + ) + import torch.distributed as dist + + try: + world = dist.get_world_size(binding.process_group) + group_rank = dist.get_group_rank( + binding.process_group, dist.get_rank() + ) + backend = dist.get_backend(binding.process_group) + except Exception as exc: + raise ValueError( + "the current rank must belong to the explicit codebook process group" + ) from exc + if world != len(members): + raise ValueError( + "codebook process-group world size does not match its stable identity" + ) + if group_rank < 0 or members[group_rank] != binding.local_member: + raise ValueError( + "runtime codebook group order does not match ordered semantic members" + ) + if not self._codebook_scope_backend_device_is_valid( + backend, binding.collective_device + ): + raise ValueError( + "codebook collective device is incompatible with the runtime backend" + ) + if binding.collective_device.type == "cuda": + index = binding.collective_device.index + if ( + not torch.cuda.is_available() + or index is None + or index < 0 + or index >= torch.cuda.device_count() + ): + raise ValueError("codebook collective CUDA device is unavailable") + + @torch._dynamo.disable + def _assert_runtime_codebook_process_group(self) -> None: + binding = self._gefen_codebook_process_group + if binding is None: + return + self._assert_finalized_binding_layout() + self._validate_codebook_runtime_binding(binding) + + def _codebook_parameter_contributes(self, parameter) -> bool: + binding = self._gefen_codebook_process_group + if binding is None: + return True + shard = self._gefen_shard_bindings.get(parameter) + if shard is None: + raise RuntimeError( + "scoped codebook parameter has no finalized shard identity" + ) + if shard.layout is ParameterLayout.REPLICATED: + return shard.local_member == binding.identity.ordered_members[0] + return True + + def _stage_post_sharding( + self, rebindings, manifest, codebook_process_group=None + ): self._assert_rebinding_pristine(rebindings) live_slots = [] for group_index, group in enumerate(self.param_groups): @@ -1491,6 +1952,10 @@ def _stage_post_sharding(self, rebindings, manifest): if len(assigned_positions) != len(live_slots): raise ValueError("post_sharding did not bind every optimizer slot") self._assert_rebound_storage_disjoint(final_targets) + if codebook_process_group is not None: + self._validate_codebook_process_group_binding( + codebook_process_group, rebindings, manifest + ) staged = object.__new__(type(self)) staged.__dict__ = self.__dict__.copy() @@ -1531,6 +1996,8 @@ def _stage_post_sharding(self, rebindings, manifest): ) staged._gefen_sharding_manifest = manifest staged._gefen_post_sharding_finalized = True + staged._gefen_codebook_process_group = codebook_process_group + staged._gefen_codebook_scope_validated = False staged._gefen_codebook_by_device = {} staged._gefen_codebook_lut_by_device = {} staged._sr_seed_by_device = {} @@ -1545,7 +2012,13 @@ def _stage_post_sharding(self, rebindings, manifest): ) return staged - def post_sharding(self, rebindings, *, manifest: ShardingManifest) -> None: + def post_sharding( + self, + rebindings, + *, + manifest: ShardingManifest, + codebook_process_group=None, + ) -> None: """Atomically finalize every local optimizer slot after sharding.""" if not isinstance(manifest, ShardingManifest): @@ -1564,7 +2037,9 @@ def post_sharding(self, rebindings, *, manifest: ShardingManifest) -> None: if self._parameter_in(sources, rebinding.old_parameter): raise ValueError("rebinding source tensors must be unique") sources.append(rebinding.old_parameter) - staged = self._stage_post_sharding(rebindings, manifest) + staged = self._stage_post_sharding( + rebindings, manifest, codebook_process_group + ) self.__dict__.update(staged.__dict__) def rebind_shard( @@ -1675,6 +2150,28 @@ def _iter_group_params_with_names(self, group): else: yield self._param_name(param), param + def _iter_codebook_params_with_names(self): + """Yield codebook inputs in canonical order when a scope is bound.""" + + if self._gefen_codebook_process_group is None: + for group in self.param_groups: + for name, parameter in self._iter_group_params_with_names(group): + yield group, name, parameter + return + by_identity = {} + for group in self.param_groups: + for name, parameter in self._iter_group_params_with_names(group): + by_identity[id(parameter)] = (group, name, parameter) + for parameter, _ in self._gefen_local_shard_bindings: + if parameter is not None: + yield by_identity[id(parameter)] + + def _has_local_codebook_gradients(self) -> bool: + return any( + parameter.grad is not None + for _, _, parameter in self._iter_codebook_params_with_names() + ) + @staticmethod def _unique_name(base: str, existing_names) -> str: name = str(base).lower() @@ -2752,6 +3249,11 @@ def _gefen_codebook_device(self) -> torch.device: if p.grad is not None: return p.grad.device return p.device + # A whole-parameter non-owner deliberately has no local parameter + # storage but must still materialize the group-agreed codebook so its + # optimizer-common checkpoint state remains complete. + if self._gefen_codebook_process_group is not None: + return torch.device("cpu") raise ValueError( "Expected at least one parameter when choosing the Gefen codebook device." ) @@ -2948,101 +3450,521 @@ def _print_v_period( ) ) - def _iter_gefen_grad_periods(self, reuse_existing_periods: bool = False): + def _iter_gefen_grad_periods( + self, reuse_existing_periods: bool = False, staged_periods=None + ): + for _, param_name, p in self._iter_codebook_params_with_names(): + if p.grad is None: + continue + grad = p.grad + # Codebook learning runs before the step loop's own guard, so + # reject sparse grads here too with the same clear error. + if getattr(grad, "is_sparse", False): + raise RuntimeError("Gefen does not support sparse gradients") + if torch.is_complex(grad): + raise RuntimeError("Gefen does not support complex gradients") + if hasattr(grad, "to_local"): + grad = grad.to_local() + if hasattr(grad, "wait"): + grad = grad.wait() + grad = grad.detach() + flat = grad.reshape(-1) + if flat.numel() == 0: + continue - for group in self.param_groups: - for param_name, p in self._iter_group_params_with_names(group): - if p.grad is None: - continue - grad = p.grad - # Codebook learning runs before the step loop's own guard, so - # reject sparse grads here too with the same clear error. - if getattr(grad, "is_sparse", False): - raise RuntimeError("Gefen does not support sparse gradients") - if torch.is_complex(grad): - raise RuntimeError("Gefen does not support complex gradients") - if hasattr(grad, "to_local"): - grad = grad.to_local() - if hasattr(grad, "wait"): - grad = grad.wait() - grad = grad.detach() - flat = grad.reshape(-1) - if flat.numel() == 0: - continue + if reuse_existing_periods: + state = self.state[p] + if "automatic_period" not in state: + raise ValueError( + "Expected automatic_period to exist for {} before refreshing Gefen codebook at optimizer step {}".format( + param_name, + self._gefen_global_step, + ) + ) + period = state["automatic_period"] + elif flat.numel() == 1: + period = 1 + else: + period = self._resolve_automatic_period(param_name, p, grad) + + if flat.numel() % period != 0: + raise ValueError( + "Automatic partition period {} does not divide parameter {} with numel {} while learning Gefen codebook".format( + period, + param_name, + flat.numel(), + ) + ) + + if staged_periods is None: + self.state[p]["automatic_period"] = period + elif not reuse_existing_periods: + staged_periods.append((p, period)) + if not self._codebook_parameter_contributes(p): + continue + + yield param_name, flat, period, grad + + def _synchronize_codebook_scope_failure(self, error, phase: str) -> None: + binding = self._gefen_codebook_process_group + if binding is None or len(binding.identity.ordered_members) == 1: + if error is not None: + raise error + return + self._assert_runtime_codebook_process_group() + 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( + "scoped Gefen codebook {} failed on local member {}: {}".format( + phase, binding.local_member, error + ) + ) from error + raise RuntimeError( + "scoped Gefen codebook {} failed on another process-group member".format( + phase + ) + ) + + def _prepare_scoped_amp_optimizer_step(self) -> bool: + binding = self._gefen_codebook_process_group + if binding is None: + return _amp_prepare_optimizer_step(self) + found_inf = getattr(self, "found_inf", None) + grad_scale = getattr(self, "grad_scale", None) + try: + if found_inf is None: + local_overflow = False + elif torch.is_tensor(found_inf): + if found_inf.numel() != 1: + raise RuntimeError( + "GradScaler supplied a non-scalar found_inf tensor with " + "shape {}".format(tuple(found_inf.shape)) + ) + local_overflow = bool(found_inf.detach().item()) + else: + if len(binding.identity.ordered_members) > 1: + raise RuntimeError( + "a multi-member scoped optimizer requires a group-aware " + "gradient scaler to provide tensor found_inf on every member" + ) + local_overflow = bool(found_inf) + if grad_scale is None: + scale_present = False + local_scale = 0.0 + elif torch.is_tensor(grad_scale): + if grad_scale.numel() != 1: + raise RuntimeError( + "GradScaler supplied a non-scalar grad_scale tensor with " + "shape {}".format(tuple(grad_scale.shape)) + ) + scale_present = True + local_scale = float(grad_scale.detach().item()) + else: + scale_present = True + local_scale = float(grad_scale) + if scale_present and ( + not math.isfinite(local_scale) or local_scale <= 0.0 + ): + raise RuntimeError( + "GradScaler supplied a non-finite or non-positive grad_scale" + ) + local_error = None + except Exception as exc: + local_overflow = False + scale_present = False + local_scale = 0.0 + local_error = exc + self._synchronize_codebook_scope_failure( + local_error, "AMP overflow preflight" + ) + if len(binding.identity.ordered_members) > 1: + import torch.distributed as dist + + amp_control = torch.tensor( + [int(local_overflow), int(scale_present), local_scale], + dtype=torch.float64, + device=binding.collective_device, + ) + controls = [ + torch.empty_like(amp_control) + for _ in binding.identity.ordered_members + ] + dist.all_gather(controls, amp_control, group=binding.process_group) + if any(not torch.equal(item, controls[0]) for item in controls[1:]): + raise RuntimeError( + "scoped Gefen AMP requires identical found_inf and grad_scale " + "on every process-group member; use a group-aware gradient scaler" + ) + if local_overflow: + return False + elif local_overflow: + return False + return _amp_prepare_optimizer_step(self) + + def _reduce_codebook_scope_histogram(self, histogram: torch.Tensor) -> torch.Tensor: + binding = self._gefen_codebook_process_group + if binding is None or len(binding.identity.ordered_members) == 1: + return histogram + self._assert_runtime_codebook_process_group() + import torch.distributed as dist + + reduced = histogram.to(binding.collective_device).clone() + dist.all_reduce(reduced, op=dist.ReduceOp.SUM, group=binding.process_group) + return reduced.cpu() + + def _codebook_manifest_fingerprint(self): + manifest = self._gefen_sharding_manifest + payload = tuple( + ( + shard.parameter.fqn, + shard.parameter.global_shape, + shard.layout.value, + shard.logical_slice.flat_offset, + shard.logical_slice.length, + shard.process_group.semantic_name, + shard.process_group.ordered_members, + shard.local_member, + shard.owner, + tuple( + ( + placement.mesh_axis, + placement.kind.value, + placement.coordinate, + placement.parts, + placement.parameter_dimension, + ) + for placement in shard.placements + ), + ) + for shard in manifest.shards + ) + return self._sha256_int64(repr(payload).encode("utf-8")) + + @staticmethod + def _sha256_int64(payload): + digest = hashlib.sha256(payload).digest() + return tuple( + int.from_bytes(digest[index : index + 8], "big", signed=True) + for index in range(0, len(digest), 8) + ) + def _codebook_scope_fingerprint(self): + return self._sha256_int64( + repr(self._serialized_codebook_scope()).encode("utf-8") + ) + + def _codebook_value_fingerprint(self): + if self._gefen_codebook is None: + return (0, 0, 0, 0) + raw = bytes( + self._gefen_codebook.detach() + .cpu() + .contiguous() + .view(torch.uint8) + .tolist() + ) + return self._sha256_int64(raw) + + def _validate_codebook_scope_operation_header(self, operation: str) -> None: + binding = self._gefen_codebook_process_group + if binding is None or len(binding.identity.ordered_members) == 1: + return + operation_codes = { + "initialize": 1, + "refresh": 2, + "periodic_step": 3, + "step": 4, + } + if operation not in operation_codes: + raise ValueError("unknown scoped codebook operation") + self._assert_runtime_codebook_process_group() + import torch.distributed as dist + + header = torch.tensor( + [ + operation_codes[operation], + int(self._gefen_global_step), + int(self._gefen_codebook is not None), + int(self._deterministic), + int(self._codebook_refresh_every), + int(self.capturable), + int(hasattr(self, "found_inf") or hasattr(self, "grad_scale")), + *self._codebook_scope_fingerprint(), + *self._codebook_manifest_fingerprint(), + *self._codebook_value_fingerprint(), + ], + dtype=torch.int64, + device=binding.collective_device, + ) + headers = [torch.empty_like(header) for _ in binding.identity.ordered_members] + dist.all_gather(headers, header, group=binding.process_group) + if any(not torch.equal(item, headers[0]) for item in headers[1:]): + raise RuntimeError( + "scoped Gefen codebook operation, step, scope, manifest, policy, " + "or old codebook differs across process-group members" + ) + if operation != "step" and self._gefen_codebook is not None: + self._verify_codebook_scope_agreement(self._gefen_codebook) + + def _validate_codebook_scope_contribution_controls( + self, staged_periods, *, reuse_existing_periods: bool + ) -> None: + binding = self._gefen_codebook_process_group + if binding is None or len(binding.identity.ordered_members) == 1: + return + self._assert_runtime_codebook_process_group() + import torch.distributed as dist + + staged_by_parameter = { + id(parameter): int(period) for parameter, period in staged_periods + } + local_records = [] + for parameter, shard in self._gefen_local_shard_bindings: + active = int( + parameter is not None + and parameter.grad is not None + and parameter.numel() > 0 + ) + period = -1 + if active: if reuse_existing_periods: - state = self.state[p] - if "automatic_period" not in state: - raise ValueError( - "Expected automatic_period to exist for {} before refreshing Gefen codebook at optimizer step {}".format( - param_name, - self._gefen_global_step, - ) - ) - period = state["automatic_period"] - elif flat.numel() == 1: - period = 1 + period = int(self.state[parameter]["automatic_period"]) else: - period = self._resolve_automatic_period(param_name, p, grad) + period = staged_by_parameter[id(parameter)] + layout_code = { + ParameterLayout.REPLICATED: 1, + ParameterLayout.FLATTENED_ELEMENT_SHARD: 2, + ParameterLayout.WHOLE_PARAMETER_OWNER: 3, + }[shard.layout] + local_records.append( + (layout_code, active, period, shard.logical_slice.length) + ) - self.state[p]["automatic_period"] = period + header = torch.tensor( + list(self._codebook_manifest_fingerprint()) + [len(local_records)], + dtype=torch.int64, + device=binding.collective_device, + ) + headers = [torch.empty_like(header) for _ in binding.identity.ordered_members] + dist.all_gather(headers, header, group=binding.process_group) + if any(not torch.equal(item, headers[0]) for item in headers[1:]): + raise RuntimeError( + "scoped Gefen codebook manifests differ across process-group members" + ) - if flat.numel() % period != 0: - raise ValueError( - "Automatic partition period {} does not divide parameter {} with numel {} while learning Gefen codebook".format( - period, - param_name, - flat.numel(), + flattened = [] + for record in local_records: + flattened.extend(record) + local = torch.tensor( + flattened, dtype=torch.int64, device=binding.collective_device + ) + gathered = [torch.empty_like(local) for _ in binding.identity.ordered_members] + dist.all_gather(gathered, local, group=binding.process_group) + for record_index, (layout_code, _, _, _) in enumerate(local_records): + offset = record_index * 4 + if any(int(item[offset].item()) != layout_code for item in gathered): + raise RuntimeError( + "scoped Gefen codebook local layouts differ across members" + ) + if layout_code != 1: + if layout_code == 2: + nonempty_activity = { + int(item[offset + 1].item()) + for item in gathered + if int(item[offset + 3].item()) > 0 + } + if len(nonempty_activity) > 1: + raise RuntimeError( + "scoped flattened parameters require every nonempty " + "shard to agree on gradient presence" ) - ) + continue + replicated_controls = { + (int(item[offset + 1].item()), int(item[offset + 2].item())) + for item in gathered + } + if len(replicated_controls) != 1: + raise RuntimeError( + "scoped replicated parameters require identical gradient " + "presence and automatic periods on every member" + ) + + def _verify_codebook_scope_agreement(self, codebook: torch.Tensor) -> None: + binding = self._gefen_codebook_process_group + if binding is None or len(binding.identity.ordered_members) == 1: + return + self._assert_runtime_codebook_process_group() + import torch.distributed as dist - yield param_name, flat, period, grad + local = codebook.detach().to( + device=binding.collective_device, dtype=torch.float32 + ).contiguous() + gathered = [torch.empty_like(local) for _ in binding.identity.ordered_members] + dist.all_gather(gathered, local, group=binding.process_group) + if any(not torch.equal(item, gathered[0]) for item in gathered[1:]): + raise RuntimeError( + "scoped Gefen codebook values differ across process-group members" + ) + + def _ensure_codebook_scope_agreement(self) -> None: + binding = self._gefen_codebook_process_group + if binding is None or self._gefen_codebook_scope_validated: + return + self._assert_runtime_codebook_process_group() + if len(binding.identity.ordered_members) == 1: + self._gefen_codebook_scope_validated = self._gefen_codebook is not None + return + import torch.distributed as dist + + control = torch.tensor( + [ + int(self._gefen_global_step), + int(self._gefen_codebook is not None), + int(self._deterministic), + int(self._codebook_refresh_every), + int(hasattr(self, "found_inf") or hasattr(self, "grad_scale")), + int(self.capturable), + *self._codebook_scope_fingerprint(), + *self._codebook_manifest_fingerprint(), + ], + dtype=torch.int64, + device=binding.collective_device, + ) + controls = [ + torch.empty_like(control) for _ in binding.identity.ordered_members + ] + dist.all_gather(controls, control, group=binding.process_group) + if any(not torch.equal(item, controls[0]) for item in controls[1:]): + raise RuntimeError( + "scoped Gefen codebook step, presence, policy, scope, manifest, " + "or AMP protocol differs across process-group members" + ) + if self._gefen_codebook is not None: + self._verify_codebook_scope_agreement(self._gefen_codebook) + self._gefen_codebook_scope_validated = True + + def _codebook_requires_2d_parameters(self) -> bool: + return False + + def _assert_codebook_capture_ready(self) -> None: + if ( + torch.cuda.is_available() + and torch.cuda.is_current_stream_capturing() + and self._gefen_codebook is None + and self._has_local_codebook_gradients() + ): + raise RuntimeError( + "Gefen codebook initialization is host-driven and must complete " + "during eager warmup before CUDA graph capture" + ) def _learn_gefen_exact_codebook( self, reuse_existing_periods: bool = False, compute_mse_logging: bool = True, - ) -> Optional[torch.Tensor]: + ): + fused_build_ok = self._fused_build_ok + try: + return self._prepare_gefen_exact_codebook( + reuse_existing_periods=reuse_existing_periods, + compute_mse_logging=compute_mse_logging, + ) + except Exception: + self._fused_build_ok = fused_build_ok + raise + + def _prepare_gefen_exact_codebook( + self, + reuse_existing_periods: bool = False, + compute_mse_logging: bool = True, + ): + if torch.cuda.is_available() and torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "Gefen codebook initialization is host-driven and cannot run " + "inside CUDA graph capture; initialize during eager warmup" + ) # Fire the fused-toolchain probe ONCE here, at the first fused CUDA # extension boundary (step-1 codebook learning), so a broken toolchain # is discovered and downgraded before the histogram / period-finding / # update paths run. On a healthy box this just builds the extension a # little earlier; for fused=False the probe is a no-op (never builds). - if self.fused: - self._gefen_fused_toolchain_ok() - codebook = learn_gefen_exact_codebook_from_grad_periods( - grad_periods=self._iter_gefen_grad_periods( - reuse_existing_periods=reuse_existing_periods - ), - codebook_device=self._gefen_codebook_device(), - num_codebooks=256, - force_endpoints=True, - verbose=self.verbose, - compute_mse_logging=compute_mse_logging, - # ``fused=False`` promises a JIT-free pure-PyTorch path. Do not load - # the separate histogram extension in that mode (or after the fused - # toolchain probe failed). An explicit global period-backend choice - # remains independent of this histogram implementation decision. - use_fused_histogram=( - FUSE_HISTOGRAM_FOR_EXACT - and self.fused - and self._fused_build_ok is not False - ), + if not reuse_existing_periods: + self._validate_codebook_scope_operation_header("initialize") + staged_periods = [] + try: + if self.fused: + self._gefen_fused_toolchain_ok() + histogram = _gefen_exact_histogram_from_grad_periods( + grad_periods=self._iter_gefen_grad_periods( + reuse_existing_periods=reuse_existing_periods, + staged_periods=staged_periods, + ), + num_codebooks=256, + # ``fused=False`` promises a JIT-free pure-PyTorch path. Do not + # load the separate histogram extension in that mode (or after + # the fused toolchain probe failed). + use_fused_histogram=( + FUSE_HISTOGRAM_FOR_EXACT + and self.fused + and self._fused_build_ok is not False + ), + ) + local_error = None + except Exception as exc: + histogram = None + local_error = exc + self._synchronize_codebook_scope_failure( + local_error, "local histogram preparation" ) - if codebook is None: - return None + self._validate_codebook_scope_contribution_controls( + staged_periods, + reuse_existing_periods=reuse_existing_periods, + ) + histogram = self._reduce_codebook_scope_histogram(histogram) - return codebook + try: + codebook = _learn_gefen_exact_codebook_from_histogram( + bin_counts_cpu=histogram, + codebook_device=self._gefen_codebook_device(), + num_codebooks=256, + force_endpoints=True, + verbose=self.verbose, + compute_mse_logging=compute_mse_logging, + ) + local_error = None + except Exception as exc: + codebook = None + local_error = exc + self._synchronize_codebook_scope_failure( + local_error, "exact-DP solve and placement" + ) + if codebook is not None: + self._verify_codebook_scope_agreement(codebook) + return codebook, staged_periods def _ensure_gefen_codebook(self, reuse_existing_periods: bool = False) -> None: - codebook = self._learn_gefen_exact_codebook( + codebook, staged_periods = self._learn_gefen_exact_codebook( reuse_existing_periods=reuse_existing_periods, compute_mse_logging=False, ) if codebook is not None: + for parameter, period in staged_periods: + self.state[parameter]["automatic_period"] = period self._gefen_codebook = codebook self._gefen_codebook_by_device.clear() self._gefen_codebook_lut_by_device.clear() + self._gefen_codebook_scope_validated = ( + self._gefen_codebook_process_group is not None + ) def _step_automatic_factored( self, group, param_name: str, p: torch.Tensor, grad: torch.Tensor @@ -3276,17 +4198,36 @@ def _refresh_codebook_with_requant(self) -> None: # the codebook that WROTE them, so the momentum survives the swap with # only nearest-codeword rounding error; m_magnitude is unchanged # (coefficients stay normalized to [-1, 1]). + if self.capturable: + raise RuntimeError( + "capturable=True does not support replacing the learned codebook" + ) + self._validate_codebook_scope_operation_header("refresh") old_codebook = self._gefen_codebook if old_codebook is None: return - new_codebook = self._learn_gefen_exact_codebook( + new_codebook, staged_periods = self._learn_gefen_exact_codebook( reuse_existing_periods=True, compute_mse_logging=False ) + if staged_periods: + raise RuntimeError( + "codebook refresh unexpectedly attempted to replace block periods" + ) if new_codebook is None: return staged_indices = [] - for pgroup in self.param_groups: - for p in pgroup["params"]: + if self._gefen_codebook_process_group is None: + parameters = [ + p for pgroup in self.param_groups for p in pgroup["params"] + ] + else: + parameters = [ + p + for p, _ in self._gefen_local_shard_bindings + if p is not None + ] + try: + for p in parameters: pstate = self.state.get(p) if not pstate or "m_codebook" not in pstate: continue @@ -3301,8 +4242,14 @@ def _refresh_codebook_with_requant(self) -> None: ) indices = gefen_nearest_codebook_indices(new_codebook_local, coeffs) staged_indices.append((pstate["m_codebook"], indices)) + local_error = None + except Exception as exc: + local_error = exc + self._synchronize_codebook_scope_failure( + local_error, "momentum requantization staging" + ) - # Commit only after every allocation/dequantization/search succeeded. + # Commit only after every member prepared every local replacement. # The staged tensors already have the destination shape/device, so this # tail is a deterministic sequence of infallible in-place copies. for stored_indices, indices in staged_indices: @@ -3331,6 +4278,18 @@ def _maybe_refresh_gefen_codebook(self) -> None: self._gefen_codebook_lut_by_device.clear() return + binding = self._gefen_codebook_process_group + if ( + (binding is None or len(binding.identity.ordered_members) == 1) + and not self._has_local_codebook_gradients() + ): + # No local histogram exists, so learning would return None. This + # fast no-op is also what keeps a no-gradient capturable step free + # of the host-driven exact-DP preparation path. Multi-member scopes + # still enter collectively because another member (for example a + # whole-parameter owner) may hold the only active gradient. + return + # No codebook yet. On a fresh run this learns it and predicts periods. # On resume the persisted codebook is normally restored in # load_state_dict, but FSDP optim-state consolidation strips custom @@ -3343,6 +4302,70 @@ def _maybe_refresh_gefen_codebook(self) -> None: reuse_existing_periods=self._resuming_from_checkpoint() ) + @torch.no_grad() + def initialize_codebook(self) -> bool: + """Collectively initialize the learned codebook without taking a step.""" + + self._assert_finalized_binding_layout() + self._assert_runtime_codebook_process_group() + self._assert_codebook_capture_ready() + self._validate_codebook_scope_operation_header("initialize") + try: + _assert_optimizer_gradients_structurally_valid( + self, require_2d_params=self._codebook_requires_2d_parameters() + ) + local_error = None + except Exception as exc: + local_error = exc + if self._gefen_codebook_process_group is not None: + self._synchronize_codebook_scope_failure( + local_error, "gradient preflight" + ) + elif local_error is not None: + raise local_error + self._ensure_codebook_scope_agreement() + if self._gefen_codebook is not None: + return False + self._ensure_gefen_codebook( + reuse_existing_periods=self._resuming_from_checkpoint() + ) + return self._gefen_codebook is not None + + @torch.no_grad() + def refresh_codebook(self) -> bool: + """Collectively relearn and atomically requantize the current codebook.""" + + self._assert_finalized_binding_layout() + self._assert_runtime_codebook_process_group() + if torch.cuda.is_available() and torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "capturable=True does not support replacing the learned codebook" + ) + self._validate_codebook_scope_operation_header("refresh") + if self.capturable: + raise RuntimeError( + "capturable=True does not support replacing the learned codebook" + ) + if self._gefen_codebook is None: + raise RuntimeError("Gefen codebook must be initialized before refresh") + try: + _assert_optimizer_gradients_structurally_valid( + self, require_2d_params=self._codebook_requires_2d_parameters() + ) + local_error = None + except Exception as exc: + local_error = exc + if self._gefen_codebook_process_group is not None: + self._synchronize_codebook_scope_failure( + local_error, "gradient preflight" + ) + elif local_error is not None: + raise local_error + self._ensure_codebook_scope_agreement() + before = self._gefen_codebook + self._refresh_codebook_with_requant() + return self._gefen_codebook is not before + def _maybe_save_gefen_grad_histogram(self) -> None: if not hasattr(quantization_module, "LIST_STEPS_SAVE_HIST_GRAD"): return @@ -4899,6 +5922,12 @@ def _compact(value): } state_dict["gefen_global_step"] = self._gefen_global_step state_dict["gefen_deterministic"] = self._deterministic + codebook_scope = self._serialized_codebook_scope() + if codebook_scope is not None: + state_dict["gefen_codebook_scope"] = codebook_scope + native_local_shards = self._serialized_native_local_shards() + if native_local_shards is not None: + state_dict["gefen_native_local_shards"] = native_local_shards # The exact-DP codebook is learned once on the first step and then frozen # for the rest of the run. It is not per-param state, so persist it # explicitly; without it resume would reinterpret the restored uint8 @@ -4916,7 +5945,11 @@ def _compact(value): # learned one. The loader removes this private transport metadata before # handing groups to torch, so it never leaks into live scheduler groups. checkpoint_metadata = { - "format_version": 1, + "format_version": ( + _SCOPED_NATIVE_METADATA_VERSION + if codebook_scope is not None + else 1 + ), "global_step": self._gefen_global_step, "codebook": self._gefen_codebook, "deterministic": self._deterministic, @@ -4928,6 +5961,12 @@ def _compact(value): # of the private group metadata during load. "device_anchor": self._checkpoint_device_anchor(), } + if codebook_scope is not None: + checkpoint_metadata["codebook_scope"] = self._serialized_codebook_scope() + if native_local_shards is not None: + checkpoint_metadata["native_local_shards"] = ( + self._serialized_native_local_shards() + ) if consolidate_rank_local: self._consolidate_rank_local_sharded_state( state_dict, checkpoint_metadata @@ -5675,6 +6714,22 @@ def _load_state_dict_impl(self, state_dict): gefen_global_step = state_dict.pop("gefen_global_step", None) gefen_codebook = state_dict.pop("gefen_codebook", None) + has_top_level_codebook_scope = "gefen_codebook_scope" in state_dict + gefen_codebook_scope = state_dict.pop("gefen_codebook_scope", None) + if has_top_level_codebook_scope: + gefen_codebook_scope = self._normalize_serialized_codebook_scope( + gefen_codebook_scope + ) + has_top_level_native_local_shards = ( + "gefen_native_local_shards" in state_dict + ) + native_local_shards = state_dict.pop( + "gefen_native_local_shards", None + ) + if has_top_level_native_local_shards: + native_local_shards = self._normalize_serialized_native_local_shards( + native_local_shards + ) has_top_level_deterministic = "gefen_deterministic" in state_dict gefen_deterministic = state_dict.pop("gefen_deterministic", None) if has_top_level_deterministic and type(gefen_deterministic) is not bool: @@ -5688,17 +6743,27 @@ def _load_state_dict_impl(self, state_dict): "Gefen checkpoint metadata is present on only some parameter groups" ) first_metadata = group_metadata[0] - if first_metadata.get("format_version") not in ( - 1, - 2, - _RANK_LOCAL_METADATA_VERSION, - ): - raise ValueError( - "Unsupported Gefen checkpoint metadata format_version: {}".format( - first_metadata.get("format_version") - ) - ) for metadata in group_metadata: + metadata_version = metadata.get("format_version") + if metadata_version not in ( + 1, + 2, + _RANK_LOCAL_METADATA_VERSION, + _SCOPED_NATIVE_METADATA_VERSION, + ): + raise ValueError( + "Unsupported Gefen checkpoint metadata format_version: {}".format( + metadata_version + ) + ) + if ("codebook_scope" in metadata) != ( + metadata_version == _SCOPED_NATIVE_METADATA_VERSION + ): + raise ValueError( + "Gefen scoped checkpoint metadata must use format_version {}".format( + _SCOPED_NATIVE_METADATA_VERSION + ) + ) if ( "deterministic" in metadata and type(metadata["deterministic"]) is not bool @@ -5708,6 +6773,9 @@ def _load_state_dict_impl(self, state_dict): "must be a bool, got {!r}".format(metadata["deterministic"]) ) for metadata in group_metadata[1:]: + same_version = metadata.get( + "format_version" + ) == first_metadata.get("format_version") same_step = metadata.get("global_step") == first_metadata.get( "global_step" ) @@ -5724,7 +6792,20 @@ def _load_state_dict_impl(self, state_dict): same_deterministic = metadata.get( "deterministic" ) == first_metadata.get("deterministic") - if not same_step or not same_codebook or not same_deterministic: + same_codebook_scope = metadata.get( + "codebook_scope" + ) == first_metadata.get("codebook_scope") + same_native_local_shards = metadata.get( + "native_local_shards" + ) == first_metadata.get("native_local_shards") + if ( + not same_version + or not same_step + or not same_codebook + or not same_deterministic + or not same_codebook_scope + or not same_native_local_shards + ): raise ValueError( "Gefen checkpoint parameter groups carry inconsistent " "optimizer metadata" @@ -5732,6 +6813,14 @@ def _load_state_dict_impl(self, state_dict): metadata_step = first_metadata.get("global_step", 0) metadata_codebook = first_metadata.get("codebook") metadata_deterministic = first_metadata.get("deterministic") + metadata_codebook_scope = self._normalize_serialized_codebook_scope( + first_metadata.get("codebook_scope") + ) + metadata_native_local_shards = ( + self._normalize_serialized_native_local_shards( + first_metadata.get("native_local_shards") + ) + ) if gefen_global_step is None: gefen_global_step = metadata_step elif gefen_global_step != metadata_step: @@ -5761,6 +6850,27 @@ def _load_state_dict_impl(self, state_dict): "Gefen checkpoint top-level and parameter-group deterministic " "policies disagree" ) + if not has_top_level_codebook_scope: + gefen_codebook_scope = metadata_codebook_scope + elif gefen_codebook_scope != metadata_codebook_scope: + raise ValueError( + "Gefen checkpoint top-level and parameter-group codebook scopes disagree" + ) + if not has_top_level_native_local_shards: + native_local_shards = metadata_native_local_shards + elif native_local_shards != metadata_native_local_shards: + raise ValueError( + "Gefen checkpoint top-level and parameter-group local shards disagree" + ) + live_codebook_scope = self._serialized_codebook_scope() + if gefen_codebook_scope != live_codebook_scope: + raise ValueError( + "Gefen checkpoint codebook scope does not match the live explicit binding" + ) + if native_local_shards != self._serialized_native_local_shards(): + raise ValueError( + "Gefen checkpoint native local-shard identity does not match the live binding" + ) if gefen_deterministic is not None: if type(gefen_deterministic) is not bool: raise ValueError( @@ -5829,6 +6939,7 @@ def _load_state_dict_impl(self, state_dict): self._gefen_codebook = gefen_codebook self._gefen_codebook_by_device.clear() self._gefen_codebook_lut_by_device.clear() + self._gefen_codebook_scope_validated = False # The identity fingerprint would catch the wholesale state swap anyway; # reset it explicitly so the first post-load capturable step re-marks. self._static_mark_sig = None @@ -5908,7 +7019,9 @@ def step(self, closure=None): returned loss is passed through. """ self._assert_finalized_binding_layout() + self._assert_runtime_codebook_process_group() self._assert_capturable_if_capturing() + self._assert_codebook_capture_ready() loss = None if closure is not None: @@ -5916,7 +7029,20 @@ def step(self, closure=None): loss = closure() self._assert_finalized_binding_layout() - _assert_optimizer_gradients_structurally_valid(self) + self._assert_runtime_codebook_process_group() + try: + _assert_optimizer_gradients_structurally_valid(self) + local_preflight_error = None + except Exception as exc: + local_preflight_error = exc + self._validate_codebook_scope_operation_header("step") + if self._gefen_codebook_process_group is not None: + self._synchronize_codebook_scope_failure( + local_preflight_error, "gradient preflight" + ) + elif local_preflight_error is not None: + raise local_preflight_error + self._ensure_codebook_scope_agreement() # GradScaler invokes native-AMP optimizers even on overflow. Decide # before codebook learning, periodic refresh, capturable counters, or @@ -5924,9 +7050,15 @@ def step(self, closure=None): # ordinary BF16/FP32 path where GradScaler attaches nothing. if ( hasattr(self, "found_inf") or hasattr(self, "grad_scale") - ) and not _amp_prepare_optimizer_step(self): + ) and not self._prepare_scoped_amp_optimizer_step(): return loss + if ( + self._gefen_codebook_process_group is not None + and self._codebook_refresh_every + ): + self._validate_codebook_scope_operation_header("periodic_step") + self._maybe_refresh_gefen_codebook() self._maybe_save_gefen_grad_histogram() # Periodic codebook re-learn (opt-in): every N steps, refit the exact-DP diff --git a/src/gefen/gefen_muon.py b/src/gefen/gefen_muon.py index 3d32a54..7880fbd 100644 --- a/src/gefen/gefen_muon.py +++ b/src/gefen/gefen_muon.py @@ -9,7 +9,6 @@ from gefen.contracts import OptimizerContract, ParameterLayout, _gefen_muon_contract from gefen.gefen import ( Gefen, - _amp_prepare_optimizer_step, _assert_optimizer_gradients_structurally_valid, ) @@ -768,6 +767,14 @@ def optimizer_contract(self) -> OptimizerContract: non_normuon_modes=non_normuon_modes, canonical_parameter_fqns=self._canonical_identity_ready(), stable_shard_identity=self._canonical_identity_ready(), + explicit_process_group_codebook_scope=True, + whole_parameter_owner=( + self._codebook_scope_ready() + and any( + shard.layout is ParameterLayout.WHOLE_PARAMETER_OWNER + for _, shard in self._gefen_local_shard_bindings + ) + ), ) def _validate_rebinding_layout(self, rebinding) -> None: @@ -805,7 +812,7 @@ def _validate_rebinding_layout(self, rebinding) -> None: ) def _has_unscoped_whole_owner_bindings(self) -> bool: - return any( + return self._gefen_codebook_process_group is None and any( shard.layout is ParameterLayout.WHOLE_PARAMETER_OWNER for _, shard in self._gefen_local_shard_bindings ) @@ -940,7 +947,12 @@ def add_param_group(self, param_group): def _init_gefen_muon_state(self, state, grad_view: torch.Tensor) -> None: self._init_gefen_state(state, grad_view) - def _iter_gefen_grad_periods(self, reuse_existing_periods: bool = False): + def _codebook_requires_2d_parameters(self) -> bool: + return True + + def _iter_gefen_grad_periods( + self, reuse_existing_periods: bool = False, staged_periods=None + ): # Same as Gefen._iter_gefen_grad_periods, but for sharded (DTensor) # gradients gather the FULL matrix (full_tensor) instead of taking the # local shard. The exact-DP codebook and the per-param block period are @@ -949,56 +961,58 @@ def _iter_gefen_grad_periods(self, reuse_existing_periods: bool = False): # matches). With full grads, flat.numel() is the global numel and is # never 0, so every rank iterates every parameter in the same order and # the full_tensor() collective is matched across ranks. - for group in self.param_groups: - for param_name, p in self._iter_group_params_with_names(group): - if p.grad is None: - continue - grad = p.grad - # approx mode learns the codebook/period from the LOCAL shard - # (no all-gather) so periods divide the local numel that the - # approximate step operates on; exact mode gathers the full matrix. - if group["sharded_mode"] == "approx" and hasattr( - grad, "to_local" - ): - grad = grad.to_local() - elif hasattr(grad, "full_tensor"): - grad = grad.full_tensor() - elif hasattr(grad, "to_local"): - grad = grad.to_local() - if hasattr(grad, "wait"): - grad = grad.wait() - grad = grad.detach() - flat = grad.reshape(-1) - if flat.numel() == 0: - continue - - if reuse_existing_periods: - state = self.state[p] - if "automatic_period" not in state: - raise ValueError( - "Expected automatic_period to exist for {} before refreshing Gefen codebook at optimizer step {}".format( - param_name, - self._gefen_global_step, - ) - ) - period = state["automatic_period"] - elif flat.numel() == 1: - period = 1 - else: - period = self._predict_period_from_grad_sq(param_name, p, grad) - - self.state[p]["automatic_period"] = period + for group, param_name, p in self._iter_codebook_params_with_names(): + if p.grad is None: + continue + grad = p.grad + # approx mode learns the codebook/period from the LOCAL shard + # (no all-gather) so periods divide the local numel that the + # approximate step operates on; exact mode gathers the full matrix. + if group["sharded_mode"] == "approx" and hasattr(grad, "to_local"): + grad = grad.to_local() + elif hasattr(grad, "full_tensor"): + grad = grad.full_tensor() + elif hasattr(grad, "to_local"): + grad = grad.to_local() + if hasattr(grad, "wait"): + grad = grad.wait() + grad = grad.detach() + flat = grad.reshape(-1) + if flat.numel() == 0: + continue - if flat.numel() % period != 0: + if reuse_existing_periods: + state = self.state[p] + if "automatic_period" not in state: raise ValueError( - "Automatic partition period {} does not divide parameter {} with numel {} while learning Gefen codebook".format( - period, + "Expected automatic_period to exist for {} before refreshing Gefen codebook at optimizer step {}".format( param_name, - flat.numel(), + self._gefen_global_step, ) ) + period = state["automatic_period"] + elif flat.numel() == 1: + period = 1 + else: + period = self._predict_period_from_grad_sq(param_name, p, grad) + + if flat.numel() % period != 0: + raise ValueError( + "Automatic partition period {} does not divide parameter {} with numel {} while learning Gefen codebook".format( + period, + param_name, + flat.numel(), + ) + ) + + if staged_periods is None: + self.state[p]["automatic_period"] = period + elif not reuse_existing_periods: + staged_periods.append((p, period)) + if not self._codebook_parameter_contributes(p): + continue - yield param_name, flat, period, grad + yield param_name, flat, period, grad def _quantize_momentum_(self, state, momentum_view: torch.Tensor) -> None: period = state["automatic_period"] @@ -2662,21 +2676,36 @@ def _load_state_dict_impl(self, state_dict): @torch.no_grad() def step(self, closure=None): self._assert_finalized_binding_layout() + self._assert_runtime_codebook_process_group() if self._has_unscoped_whole_owner_bindings(): raise RuntimeError( "GefenMuon whole-parameter owner stepping requires the separate " - "explicit process-group codebook scope, which is not implemented" + "explicit process-group codebook scope" ) self._assert_capturable_if_capturing() + self._assert_codebook_capture_ready() loss = None if closure is not None: with torch.enable_grad(): loss = closure() self._assert_finalized_binding_layout() - _assert_optimizer_gradients_structurally_valid( - self, require_2d_params=True - ) + self._assert_runtime_codebook_process_group() + try: + _assert_optimizer_gradients_structurally_valid( + self, require_2d_params=True + ) + local_preflight_error = None + except Exception as exc: + local_preflight_error = exc + self._validate_codebook_scope_operation_header("step") + if self._gefen_codebook_process_group is not None: + self._synchronize_codebook_scope_failure( + local_preflight_error, "gradient preflight" + ) + elif local_preflight_error is not None: + raise local_preflight_error + self._ensure_codebook_scope_agreement() # Native GradScaler calls AMP-aware optimizers even when its non-finite # scan found an overflow. Skip before the sharded preflight, first-step @@ -2684,7 +2713,7 @@ def step(self, closure=None): # gradients are unscaled once here for every existing Muon path. if ( hasattr(self, "found_inf") or hasattr(self, "grad_scale") - ) and not _amp_prepare_optimizer_step(self): + ) and not self._prepare_scoped_amp_optimizer_step(): return loss # Partition the work once so distributed-mode sharded params can take the diff --git a/tests/test_codebook_scope_cpu.py b/tests/test_codebook_scope_cpu.py new file mode 100644 index 0000000..bf0d668 --- /dev/null +++ b/tests/test_codebook_scope_cpu.py @@ -0,0 +1,508 @@ +import copy +from dataclasses import FrozenInstanceError +import io + +import pytest +import torch + +from gefen import ( + CodebookProcessGroupBinding, + Gefen, + GefenMuon, + LogicalSlice, + ParameterIdentity, + ParameterLayout, + ParameterRebinding, + ParameterStateRole, + PlacementKind, + ProcessGroupIdentity, + ProcessGroupScope, + ShardIdentity, + ShardPlacement, + ShardingManifest, + StateExtent, +) +import gefen.gefen as gefen_module + + +def _replicated_shard(parameter, group, member): + coordinate = group.ordered_members.index(member) + return ShardIdentity( + parameter, + ParameterLayout.REPLICATED, + LogicalSlice.full(parameter), + process_group=group, + local_member=member, + placements=( + ShardPlacement( + "data_parallel", + PlacementKind.REPLICATE, + coordinate, + len(group.ordered_members), + ), + ), + ) + + +def _flat_shard(parameter, group, member): + coordinate = group.ordered_members.index(member) + return ShardIdentity( + parameter, + ParameterLayout.FLATTENED_ELEMENT_SHARD, + LogicalSlice.full(parameter), + process_group=group, + local_member=member, + placements=( + ShardPlacement( + "data_parallel", + PlacementKind.FLAT_SHARD, + coordinate, + len(group.ordered_members), + ), + ), + ) + + +def _whole_owner_shard(parameter, group, member, owner): + coordinate = group.ordered_members.index(member) + return ShardIdentity( + parameter, + ParameterLayout.WHOLE_PARAMETER_OWNER, + LogicalSlice.full(parameter) if member == owner else LogicalSlice(0, 0), + process_group=group, + local_member=member, + owner=owner, + placements=( + ShardPlacement( + "data_parallel", + PlacementKind.WHOLE_PARAMETER_OWNER, + coordinate, + len(group.ordered_members), + ), + ), + ) + + +def _single_member_binding(group, member="rank:0"): + return CodebookProcessGroupBinding(group, member, None, torch.device("cpu")) + + +def _finalize_replicated(optimizer, parameter, *, fqn="Layer.Weight"): + group = ProcessGroupIdentity("data_parallel", ("rank:0",)) + identity = ParameterIdentity(fqn, tuple(parameter.shape)) + shard = _replicated_shard(identity, group, "rank:0") + binding = _single_member_binding(group) + optimizer.post_sharding( + (ParameterRebinding(parameter, parameter, shard),), + manifest=ShardingManifest((shard,)), + codebook_process_group=binding, + ) + return binding + + +def _snapshot(optimizer): + return { + "dict": optimizer.__dict__.copy(), + "groups": optimizer.param_groups, + "state": optimizer.state, + "codebook": optimizer._gefen_codebook, + "binding": optimizer._gefen_codebook_process_group, + "validated": optimizer._gefen_codebook_scope_validated, + } + + +def _assert_snapshot_identity(optimizer, snapshot): + assert optimizer.param_groups is snapshot["groups"] + assert optimizer.state is snapshot["state"] + assert optimizer._gefen_codebook is snapshot["codebook"] + assert optimizer._gefen_codebook_process_group is snapshot["binding"] + assert optimizer._gefen_codebook_scope_validated is snapshot["validated"] + assert optimizer.__dict__.keys() == snapshot["dict"].keys() + for key, value in snapshot["dict"].items(): + assert optimizer.__dict__[key] is value + + +def test_codebook_process_group_binding_is_public_frozen_and_ordered(): + group = ProcessGroupIdentity("replica", ("worker:b", "worker:a")) + binding = CodebookProcessGroupBinding(group, "worker:b", object(), torch.device("cpu")) + + assert binding.identity is group + assert binding.local_member == "worker:b" + assert binding.sort_key == ("replica", ("worker:b", "worker:a")) + with pytest.raises(FrozenInstanceError): + binding.local_member = "worker:a" + with pytest.raises(TypeError, match="ProcessGroupIdentity"): + CodebookProcessGroupBinding(object(), "worker:b", None, "cpu") + with pytest.raises(ValueError, match="local_member"): + CodebookProcessGroupBinding(group, "missing", None, "cpu") + with pytest.raises(ValueError, match="materialized"): + CodebookProcessGroupBinding(group, "worker:b", None, "meta") + + +def test_scoped_one_member_initialization_matches_unscoped_first_step(): + initial = torch.arange(1, 17, dtype=torch.float32).reshape(4, 4) + scoped_param = torch.nn.Parameter(initial.clone()) + reference_param = torch.nn.Parameter(initial.clone()) + scoped = Gefen([("Layer.Weight", scoped_param)], fused=False, factored_v_2d=False) + reference = Gefen([("Layer.Weight", reference_param)], fused=False, factored_v_2d=False) + binding = _finalize_replicated(scoped, scoped_param) + gradient = torch.linspace(-2, 3, initial.numel()).reshape_as(initial) + scoped_param.grad = gradient.clone() + reference_param.grad = gradient.clone() + + assert scoped.initialize_codebook() + assert not scoped.initialize_codebook() + scoped.step() + reference.step() + + assert scoped.codebook_process_group_binding() is binding + assert scoped.optimizer_contract().capabilities.explicit_process_group_codebook_scope + assert torch.equal(scoped._gefen_codebook, reference._gefen_codebook) + assert torch.equal(scoped_param, reference_param) + for key in scoped.state[scoped_param]: + left = scoped.state[scoped_param][key] + right = reference.state[reference_param][key] + if torch.is_tensor(left): + assert torch.equal(left, right) + else: + assert left == right + frozen_codebook = scoped._gefen_codebook + frozen_indices = scoped.state[scoped_param]["m_codebook"].detach().clone() + scoped_param.grad = None + assert not scoped.refresh_codebook() + assert scoped._gefen_codebook is frozen_codebook + assert torch.equal(scoped.state[scoped_param]["m_codebook"], frozen_indices) + + +def test_scope_binding_validation_is_part_of_atomic_post_sharding(): + parameter = torch.nn.Parameter(torch.ones(8)) + optimizer = Gefen([("p", parameter)], fused=False, factored_v_2d=False) + group = ProcessGroupIdentity("dp", ("rank:0",)) + wrong_group = ProcessGroupIdentity("other", ("rank:0",)) + identity = ParameterIdentity("P", (8,)) + shard = _replicated_shard(identity, group, "rank:0") + manifest = ShardingManifest((shard,)) + snapshot = _snapshot(optimizer) + + with pytest.raises(ValueError, match="process-group identity"): + optimizer.post_sharding( + (ParameterRebinding(parameter, parameter, shard),), + manifest=manifest, + codebook_process_group=_single_member_binding(wrong_group), + ) + + _assert_snapshot_identity(optimizer, snapshot) + assert not optimizer._gefen_post_sharding_finalized + + +def test_scope_binding_rejects_implicit_default_world_for_one_member(): + parameter = torch.nn.Parameter(torch.ones(4)) + optimizer = Gefen([parameter], fused=False, factored_v_2d=False) + group = ProcessGroupIdentity("dp", ("rank:0",)) + identity = ParameterIdentity("P", (4,)) + shard = _replicated_shard(identity, group, "rank:0") + + with pytest.raises(ValueError, match="process_group=None"): + optimizer.post_sharding( + (ParameterRebinding(parameter, parameter, shard),), + manifest=ShardingManifest((shard,)), + codebook_process_group=CodebookProcessGroupBinding(group, "rank:0", object(), "cpu"), + ) + + assert not optimizer._gefen_post_sharding_finalized + + +def test_initialization_failure_does_not_publish_periods_or_codebook(monkeypatch): + parameter = torch.nn.Parameter(torch.arange(8, dtype=torch.float32)) + optimizer = Gefen([("p", parameter)], fused=False, factored_v_2d=False) + _finalize_replicated(optimizer, parameter, fqn="P") + parameter.grad = torch.arange(1, 9, dtype=torch.float32) + before_state = copy.deepcopy(optimizer.state[parameter]) + before_cache = optimizer._gefen_codebook_by_device + + def fail_exact_dp(*args, **kwargs): + raise RuntimeError("injected exact-DP failure") + + monkeypatch.setattr(gefen_module.quantization_module, "exact_dp", fail_exact_dp) + with pytest.raises(RuntimeError, match="injected exact-DP failure"): + optimizer.initialize_codebook() + + assert optimizer.state[parameter] == before_state + assert optimizer._gefen_codebook is None + assert optimizer._gefen_codebook_by_device is before_cache + assert not optimizer._gefen_codebook_scope_validated + assert optimizer._gefen_global_step == 0 + + +def test_scoped_native_checkpoint_uses_primitive_scope_and_requires_exact_binding(): + source_param = torch.nn.Parameter(torch.arange(8, dtype=torch.float32)) + source = Gefen([("p", source_param)], fused=False, factored_v_2d=False) + _finalize_replicated(source, source_param, fqn="P") + source_param.grad = torch.arange(1, 9, dtype=torch.float32) + source.step() + checkpoint = source.state_dict() + + scope = checkpoint["gefen_codebook_scope"] + assert scope == { + "format_version": 1, + "semantic_name": "data_parallel", + "ordered_members": ["rank:0"], + "refresh_every": 0, + } + assert all(group["_gefen_checkpoint_metadata"]["codebook_scope"] == scope for group in checkpoint["param_groups"]) + assert all(group["_gefen_checkpoint_metadata"]["format_version"] == 4 for group in checkpoint["param_groups"]) + buffer = io.BytesIO() + torch.save(checkpoint, buffer) + buffer.seek(0) + loaded = torch.load(buffer, weights_only=True) + assert loaded["gefen_codebook_scope"] == scope + + target_param = torch.nn.Parameter(source_param.detach().clone()) + target = Gefen([("p", target_param)], fused=False, factored_v_2d=False) + _finalize_replicated(target, target_param, fqn="P") + live_binding = target.codebook_process_group_binding() + target.load_state_dict(copy.deepcopy(checkpoint)) + assert target.codebook_process_group_binding() is live_binding + assert not target._gefen_codebook_scope_validated + assert torch.equal(target._gefen_codebook, source._gefen_codebook) + continuation_grad = torch.linspace(-1, 1, 8) + source_param.grad = continuation_grad.clone() + target_param.grad = continuation_grad.clone() + source.step() + target.step() + assert torch.equal(target_param, source_param) + assert torch.equal(target._gefen_codebook, source._gefen_codebook) + + conflicting_param = torch.nn.Parameter(torch.arange(8, dtype=torch.float32)) + conflicting_target = Gefen([("p", conflicting_param)], fused=False, factored_v_2d=False) + _finalize_replicated(conflicting_target, conflicting_param, fqn="P") + conflicting = copy.deepcopy(checkpoint) + conflicting["gefen_codebook_scope"] = dict(conflicting["gefen_codebook_scope"]) + conflicting["gefen_codebook_scope"]["semantic_name"] = "other" + conflicting_before = _snapshot(conflicting_target) + with pytest.raises(ValueError, match="scopes disagree"): + conflicting_target.load_state_dict(conflicting) + _assert_snapshot_identity(conflicting_target, conflicting_before) + + unsafe_version = copy.deepcopy(checkpoint) + unsafe_version["param_groups"][0]["_gefen_checkpoint_metadata"]["format_version"] = 1 + with pytest.raises(ValueError, match="format_version 4"): + conflicting_target.load_state_dict(unsafe_version) + _assert_snapshot_identity(conflicting_target, conflicting_before) + + unbound_param = torch.nn.Parameter(torch.arange(8, dtype=torch.float32)) + unbound = Gefen([("p", unbound_param)], fused=False, factored_v_2d=False) + before = _snapshot(unbound) + with pytest.raises(ValueError, match="does not match"): + unbound.load_state_dict(checkpoint) + _assert_snapshot_identity(unbound, before) + + schedule_param = torch.nn.Parameter(torch.arange(8, dtype=torch.float32)) + schedule_target = Gefen( + [("p", schedule_param)], + fused=False, + factored_v_2d=False, + codebook_refresh_every=1, + ) + _finalize_replicated(schedule_target, schedule_param, fqn="P") + schedule_before = _snapshot(schedule_target) + with pytest.raises(ValueError, match="does not match"): + schedule_target.load_state_dict(checkpoint) + _assert_snapshot_identity(schedule_target, schedule_before) + + +@pytest.mark.parametrize( + "second_layout", + [ + ParameterLayout.FLATTENED_ELEMENT_SHARD, + ParameterLayout.REPLICATED, + ], +) +def test_native_flat_checkpoint_guard_rejects_reordered_equal_shape_slots_atomically( + second_layout, +): + group = ProcessGroupIdentity("data_parallel", ("rank:0",)) + first_identity = ParameterIdentity("Model.First", (4,)) + second_identity = ParameterIdentity("Model.Second", (4,)) + first_shard = _flat_shard(first_identity, group, "rank:0") + second_shard = ( + _flat_shard(second_identity, group, "rank:0") + if second_layout is ParameterLayout.FLATTENED_ELEMENT_SHARD + else _replicated_shard(second_identity, group, "rank:0") + ) + manifest = ShardingManifest((first_shard, second_shard)) + + source_first = torch.nn.Parameter(torch.zeros(4)) + source_second = torch.nn.Parameter(torch.zeros(4)) + source = Gefen( + [("first", source_first), ("second", source_second)], + fused=False, + factored_v_2d=False, + ) + source.post_sharding( + ( + ParameterRebinding(source_first, source_first, first_shard), + ParameterRebinding(source_second, source_second, second_shard), + ), + manifest=manifest, + codebook_process_group=_single_member_binding(group), + ) + source._resolve_automatic_period = lambda *args: 4 + source_first.grad = torch.tensor([1.0, 2.0, 3.0, 4.0]) + source_second.grad = torch.tensor([9.0, -1.0, -2.0, -3.0]) + source.step() + checkpoint = source.state_dict() + source_slots = checkpoint["gefen_native_local_shards"]["param_groups"] + assert ( + checkpoint["param_groups"][0]["_gefen_checkpoint_metadata"]["native_local_shards"] + == checkpoint["gefen_native_local_shards"] + ) + buffer = io.BytesIO() + torch.save(checkpoint, buffer) + buffer.seek(0) + assert torch.load(buffer, weights_only=True)["gefen_native_local_shards"] == checkpoint["gefen_native_local_shards"] + malformed_version = copy.deepcopy(checkpoint) + malformed_version["gefen_native_local_shards"]["format_version"] = True + for group_record in malformed_version["param_groups"]: + group_record["_gefen_checkpoint_metadata"]["native_local_shards"]["format_version"] = True + source_before = _snapshot(source) + with pytest.raises(ValueError, match="format_version"): + source.load_state_dict(malformed_version) + _assert_snapshot_identity(source, source_before) + assert [[record["fqn"] for record in records] for records in source_slots] == [["Model.First", "Model.Second"]] + if second_layout is ParameterLayout.REPLICATED: + assert source_slots[0][1]["layout"] == ParameterLayout.REPLICATED.value + + target_second = torch.nn.Parameter(torch.zeros(4)) + target_first = torch.nn.Parameter(torch.zeros(4)) + target = Gefen( + [("second", target_second), ("first", target_first)], + fused=False, + factored_v_2d=False, + ) + target.post_sharding( + ( + ParameterRebinding(target_second, target_second, second_shard), + ParameterRebinding(target_first, target_first, first_shard), + ), + manifest=manifest, + codebook_process_group=_single_member_binding(group), + ) + target_before = _snapshot(target) + + with pytest.raises(ValueError, match="local-shard identity"): + target.load_state_dict(copy.deepcopy(checkpoint)) + + _assert_snapshot_identity(target, target_before) + + +def test_unscoped_native_checkpoint_does_not_serialize_a_none_scope(): + parameter = torch.nn.Parameter(torch.arange(4, dtype=torch.float32)) + optimizer = Gefen([parameter], fused=False) + checkpoint = optimizer.state_dict() + + assert "gefen_codebook_scope" not in checkpoint + assert all("codebook_scope" not in group["_gefen_checkpoint_metadata"] for group in checkpoint["param_groups"]) + + +def test_capturable_optimizer_rejects_manual_codebook_replacement(): + parameter = torch.nn.Parameter(torch.arange(8, dtype=torch.float32)) + optimizer = Gefen( + [("p", parameter)], + fused=False, + factored_v_2d=False, + capturable=True, + ) + _finalize_replicated(optimizer, parameter, fqn="P") + parameter.grad = torch.arange(1, 9, dtype=torch.float32) + assert optimizer.initialize_codebook() + + with pytest.raises(RuntimeError, match="capturable=True"): + optimizer.refresh_codebook() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +def test_capturable_scope_rejects_first_active_step_inside_graph_capture(): + parameter = torch.nn.Parameter(torch.arange(8, dtype=torch.float32, device="cuda")) + optimizer = Gefen( + [("p", parameter)], + fused=False, + factored_v_2d=False, + capturable=True, + ) + _finalize_replicated(optimizer, parameter, fqn="P") + parameter.grad = torch.arange(1, 9, dtype=torch.float32, device="cuda") + graph = torch.cuda.CUDAGraph() + + with pytest.raises(RuntimeError, match="eager warmup"): + with torch.cuda.graph(graph): + optimizer.step() + + assert optimizer._gefen_codebook is None + assert optimizer._gefen_global_step == 0 + + +def test_capturable_multi_member_scope_binding_is_atomic_rejection(): + parameter = torch.nn.Parameter(torch.arange(8, dtype=torch.float32)) + optimizer = Gefen( + [("p", parameter)], + fused=False, + factored_v_2d=False, + capturable=True, + ) + group = ProcessGroupIdentity("dp", ("rank:0", "rank:1")) + identity = ParameterIdentity("P", (8,)) + records = tuple(_replicated_shard(identity, group, member) for member in group.ordered_members) + before = _snapshot(optimizer) + + with pytest.raises(ValueError, match="capturable=True"): + optimizer.post_sharding( + (ParameterRebinding(parameter, parameter, records[0]),), + manifest=ShardingManifest(records), + codebook_process_group=CodebookProcessGroupBinding(group, "rank:0", object(), "cpu"), + ) + + _assert_snapshot_identity(optimizer, before) + + +def test_muon_whole_owner_scope_enables_owner_update_with_sync_requirement(): + parameter = torch.nn.Parameter(torch.ones(4, 4)) + optimizer = GefenMuon([("matrix", parameter)], fused=False) + group = ProcessGroupIdentity("owner_group", ("rank:0",)) + identity = ParameterIdentity("Matrix", (4, 4)) + shard = _whole_owner_shard(identity, group, "rank:0", "rank:0") + binding = _single_member_binding(group) + optimizer.post_sharding( + (ParameterRebinding(parameter, parameter, shard),), + manifest=ShardingManifest((shard,)), + codebook_process_group=binding, + ) + before = parameter.detach().clone() + parameter.grad = torch.arange(1, 17, dtype=torch.float32).reshape(4, 4) + + optimizer.step() + + support = next( + item + for item in optimizer.optimizer_contract().capabilities.training + if item.layout is ParameterLayout.WHOLE_PARAMETER_OWNER + ) + assert support.process_group_scope is ProcessGroupScope.ADAPTER_DEFINED + assert support.requires_complete_parameter_storage + assert support.requires_post_step_parameter_sync + assert all( + ParameterLayout.WHOLE_PARAMETER_OWNER not in checkpoint.same_topology + for checkpoint in optimizer.optimizer_contract().capabilities.checkpoints + ) + owner_variants = [ + variant + for variant in optimizer.optimizer_contract().state_layout.parameter_variants + if ParameterLayout.WHOLE_PARAMETER_OWNER in variant.layouts and variant.initialized + ] + assert owner_variants + assert all( + variant.extent is StateExtent.OWNER_PARAMETER and variant.role is ParameterStateRole.OWNER + for variant in owner_variants + ) + assert not torch.equal(parameter, before) + assert optimizer._gefen_global_step == 1 diff --git a/tests/test_codebook_scope_distributed.py b/tests/test_codebook_scope_distributed.py new file mode 100644 index 0000000..cb0c048 --- /dev/null +++ b/tests/test_codebook_scope_distributed.py @@ -0,0 +1,920 @@ +from datetime import timedelta +import copy +import multiprocessing as mp +import os +import tempfile + +import pytest +import torch +import torch.distributed as dist + +from gefen import ( + CodebookProcessGroupBinding, + Gefen, + GefenMuon, + LogicalSlice, + ParameterIdentity, + ParameterLayout, + ParameterRebinding, + PlacementKind, + ProcessGroupIdentity, + ShardIdentity, + ShardPlacement, + ShardingManifest, +) +from gefen.gefen import learn_gefen_exact_codebook_from_grad_periods +import gefen.gefen as gefen_module + + +def _replicated(parameter, group, member): + coordinate = group.ordered_members.index(member) + return ShardIdentity( + parameter, + ParameterLayout.REPLICATED, + LogicalSlice.full(parameter), + process_group=group, + local_member=member, + placements=( + ShardPlacement( + "dp", + PlacementKind.REPLICATE, + coordinate, + len(group.ordered_members), + ), + ), + ) + + +def _flat(parameter, group, member, offset, length): + coordinate = group.ordered_members.index(member) + return ShardIdentity( + parameter, + ParameterLayout.FLATTENED_ELEMENT_SHARD, + LogicalSlice(offset, length), + process_group=group, + local_member=member, + placements=( + ShardPlacement( + "dp", + PlacementKind.FLAT_SHARD, + coordinate, + len(group.ordered_members), + ), + ), + ) + + +def _whole_owner(parameter, group, member, owner): + coordinate = group.ordered_members.index(member) + return ShardIdentity( + parameter, + ParameterLayout.WHOLE_PARAMETER_OWNER, + LogicalSlice.full(parameter) if member == owner else LogicalSlice(0, 0), + process_group=group, + local_member=member, + owner=owner, + placements=( + ShardPlacement( + "dp", + PlacementKind.WHOLE_PARAMETER_OWNER, + coordinate, + len(group.ordered_members), + ), + ), + ) + + +def _binding(group, rank, process_group): + return CodebookProcessGroupBinding(group, "rank:{}".format(rank), process_group, torch.device("cpu")) + + +def _finalize(optimizer, parameter, local_shard, manifest, binding): + optimizer.post_sharding( + (ParameterRebinding(parameter, parameter, local_shard),), + manifest=manifest, + codebook_process_group=binding, + ) + + +def _distributed_worker(rank, world, init_file, queue): + try: + dist.init_process_group( + "gloo", + init_method="file://{}".format(init_file), + rank=rank, + world_size=world, + timeout=timedelta(seconds=45), + ) + members = tuple("rank:{}".format(index) for index in range(world)) + group = ProcessGroupIdentity("data_parallel", members) + runtime_group = dist.group.WORLD + + # Runtime coordinate/member validation must fail before post_sharding + # publication, then leave the same optimizer retryable with a valid map. + invalid_param = torch.nn.Parameter(torch.ones(4)) + invalid_optimizer = Gefen([("invalid", invalid_param)], fused=False, factored_v_2d=False) + invalid_identity = ParameterIdentity("Invalid", (4,)) + reversed_group = ProcessGroupIdentity("data_parallel", tuple(reversed(members))) + invalid_records = tuple( + _replicated(invalid_identity, reversed_group, member) for member in reversed_group.ordered_members + ) + invalid_local = next(item for item in invalid_records if item.local_member == "rank:{}".format(rank)) + bad_binding = CodebookProcessGroupBinding( + reversed_group, + "rank:{}".format(rank), + runtime_group, + torch.device("cpu"), + ) + try: + invalid_optimizer.post_sharding( + (ParameterRebinding(invalid_param, invalid_param, invalid_local),), + manifest=ShardingManifest(invalid_records), + codebook_process_group=bad_binding, + ) + invalid_rejected = False + except ValueError: + invalid_rejected = True + invalid_atomic = ( + not invalid_optimizer._gefen_post_sharding_finalized + and invalid_optimizer._gefen_codebook_process_group is None + and invalid_optimizer.param_groups[0]["params"] == [invalid_param] + ) + + # Replicated logical state contributes exactly once from ordered member + # zero. Rank 1 intentionally has different values but the same period; + # including both replicas would produce a different histogram oracle. + replicated_param = torch.nn.Parameter(torch.zeros(4)) + replicated_optimizer = Gefen( + [("replicated", replicated_param)], + fused=False, + factored_v_2d=False, + ) + replicated_identity = ParameterIdentity("Replicated", (4,)) + replicated_records = tuple(_replicated(replicated_identity, group, member) for member in members) + _finalize( + replicated_optimizer, + replicated_param, + replicated_records[rank], + ShardingManifest(replicated_records), + _binding(group, rank, runtime_group), + ) + replicated_optimizer._resolve_automatic_period = lambda *args: 4 + canonical_grad = torch.tensor([-4.0, -1.0, 2.0, 8.0]) + other_grad = torch.tensor([-8.0, 3.0, 4.0, 4.0]) + replicated_param.grad = canonical_grad.clone() if rank == 0 else other_grad + replicated_initialized = replicated_optimizer.initialize_codebook() + replicated_oracle = learn_gefen_exact_codebook_from_grad_periods( + grad_periods=(("Replicated", canonical_grad, 4, canonical_grad),), + codebook_device=torch.device("cpu"), + num_codebooks=256, + force_endpoints=True, + verbose=False, + compute_mse_logging=False, + use_fused_histogram=False, + ) + replicated_logical_once = torch.equal(replicated_optimizer._gefen_codebook, replicated_oracle) + + mismatch_param = torch.nn.Parameter(torch.zeros(4)) + mismatch_optimizer = Gefen([("mismatch", mismatch_param)], fused=False, factored_v_2d=False) + mismatch_identity = ParameterIdentity("Mismatch", (4,)) + mismatch_records = tuple(_replicated(mismatch_identity, group, member) for member in members) + _finalize( + mismatch_optimizer, + mismatch_param, + mismatch_records[rank], + ShardingManifest(mismatch_records), + _binding(group, rank, runtime_group), + ) + mismatch_optimizer._resolve_automatic_period = lambda *args: 4 + mismatch_param.grad = canonical_grad.clone() if rank == 0 else None + try: + mismatch_optimizer.initialize_codebook() + mismatch_rejected = False + except RuntimeError as exc: + mismatch_rejected = "identical gradient presence" in str(exc) + mismatch_atomic = ( + mismatch_optimizer._gefen_codebook is None + and mismatch_optimizer._gefen_global_step == 0 + and mismatch_optimizer.state[mismatch_param] == {"name": "mismatch"} + ) + + amp_param = torch.nn.Parameter(torch.zeros(4)) + amp_optimizer = Gefen([("amp", amp_param)], fused=False, factored_v_2d=False) + amp_identity = ParameterIdentity("Amp", (8,)) + amp_records = tuple(_flat(amp_identity, group, member, index * 4, 4) for index, member in enumerate(members)) + _finalize( + amp_optimizer, + amp_param, + amp_records[rank], + ShardingManifest(amp_records), + _binding(group, rank, runtime_group), + ) + amp_optimizer._resolve_automatic_period = lambda *args: 4 + amp_param.grad = canonical_grad.clone() if rank == 0 else other_grad.clone() + amp_grad_before = amp_param.grad.detach().clone() + amp_optimizer.found_inf = torch.tensor(float(rank == 0)) + amp_optimizer.grad_scale = torch.tensor(8.0) + try: + amp_optimizer.step() + amp_mismatch_rejected = False + except RuntimeError as exc: + amp_mismatch_rejected = "group-aware gradient scaler" in str(exc) + amp_optimizer.found_inf = torch.tensor(1.0) + amp_optimizer.step() + amp_overflow_atomic = ( + amp_optimizer._gefen_global_step == 0 + and amp_optimizer._gefen_codebook is None + and amp_optimizer.state[amp_param] == {"name": "amp"} + and torch.equal(amp_param.grad, amp_grad_before) + and torch.equal(amp_param, torch.zeros_like(amp_param)) + ) + if rank == 0: + amp_optimizer.found_inf = torch.tensor(0.0) + else: + del amp_optimizer.found_inf + del amp_optimizer.grad_scale + try: + amp_optimizer.step() + amp_protocol_rejected = False + except RuntimeError as exc: + amp_protocol_rejected = "policy" in str(exc) + + # Flattened shards each contribute once. The global oracle uses two + # period-4 blocks, exactly matching the two physical local shards. + flat_param = torch.nn.Parameter(torch.zeros(4)) + flat_optimizer = Gefen([("flat", flat_param)], fused=False, factored_v_2d=False) + flat_identity = ParameterIdentity("Flat", (8,)) + flat_records = tuple(_flat(flat_identity, group, member, index * 4, 4) for index, member in enumerate(members)) + _finalize( + flat_optimizer, + flat_param, + flat_records[rank], + ShardingManifest(flat_records), + _binding(group, rank, runtime_group), + ) + flat_optimizer._resolve_automatic_period = lambda *args: 4 + flat_grads = ( + torch.tensor([-3.0, -2.0, 1.0, 7.0]), + torch.tensor([-9.0, 2.0, 5.0, 6.0]), + ) + flat_param.grad = flat_grads[rank].clone() + flat_initialized = flat_optimizer.initialize_codebook() + flat_oracle = learn_gefen_exact_codebook_from_grad_periods( + grad_periods=tuple(("Flat", gradient, 4, gradient) for gradient in flat_grads), + codebook_device=torch.device("cpu"), + num_codebooks=256, + force_endpoints=True, + verbose=False, + compute_mse_logging=False, + use_fused_histogram=False, + ) + flat_global = torch.equal(flat_optimizer._gefen_codebook, flat_oracle) + flat_native_claim = any( + ParameterLayout.FLATTENED_ELEMENT_SHARD in item.same_topology + for item in flat_optimizer.optimizer_contract().capabilities.checkpoints + ) + + flat_mismatch_param = torch.nn.Parameter(torch.zeros(4)) + flat_mismatch_optimizer = Gefen( + [("flat_mismatch", flat_mismatch_param)], + fused=False, + factored_v_2d=False, + ) + flat_mismatch_identity = ParameterIdentity("FlatMismatch", (8,)) + flat_mismatch_records = tuple( + _flat(flat_mismatch_identity, group, member, index * 4, 4) for index, member in enumerate(members) + ) + _finalize( + flat_mismatch_optimizer, + flat_mismatch_param, + flat_mismatch_records[rank], + ShardingManifest(flat_mismatch_records), + _binding(group, rank, runtime_group), + ) + flat_mismatch_optimizer._resolve_automatic_period = lambda *args: 4 + flat_mismatch_param.grad = flat_grads[rank].clone() if rank == 0 else None + try: + flat_mismatch_optimizer.initialize_codebook() + flat_mismatch_rejected = False + except RuntimeError as exc: + flat_mismatch_rejected = "every nonempty shard" in str(exc) + flat_mismatch_atomic = flat_mismatch_optimizer._gefen_codebook is None and flat_mismatch_optimizer.state[ + flat_mismatch_param + ] == {"name": "flat_mismatch"} + + gathered_codebooks = [torch.empty_like(flat_optimizer._gefen_codebook) for _ in range(world)] + dist.all_gather(gathered_codebooks, flat_optimizer._gefen_codebook) + flat_agreement = all(torch.equal(item, gathered_codebooks[0]) for item in gathered_codebooks[1:]) + flat_optimizer.step() + flat_checkpoint = copy.deepcopy(flat_optimizer.state_dict()) + checkpoint_scopes = [None] * world + dist.all_gather_object( + checkpoint_scopes, + flat_checkpoint["gefen_codebook_scope"], + group=runtime_group, + ) + rank_neutral_checkpoint_scope = all(item == checkpoint_scopes[0] for item in checkpoint_scopes[1:]) + local_shard_records = [None] * world + dist.all_gather_object( + local_shard_records, + flat_checkpoint["gefen_native_local_shards"], + group=runtime_group, + ) + rank_local_checkpoint_identity = local_shard_records[0] != local_shard_records[1] + rank_zero_checkpoint = [flat_checkpoint if rank == 0 else None] + dist.broadcast_object_list(rank_zero_checkpoint, src=0, group=runtime_group) + cross_param = torch.nn.Parameter(flat_param.detach().clone()) + cross_target = Gefen([("flat", cross_param)], fused=False, factored_v_2d=False) + _finalize( + cross_target, + cross_param, + flat_records[rank], + ShardingManifest(flat_records), + _binding(group, rank, runtime_group), + ) + try: + cross_target.load_state_dict(rank_zero_checkpoint[0]) + cross_member_guard = rank == 0 + except ValueError as exc: + cross_member_guard = rank == 1 and "local-shard identity" in str(exc) + resumed_param = torch.nn.Parameter(flat_param.detach().clone()) + resumed = Gefen([("flat", resumed_param)], fused=False, factored_v_2d=False) + resumed_binding = _binding(group, rank, runtime_group) + _finalize( + resumed, + resumed_param, + flat_records[rank], + ShardingManifest(flat_records), + resumed_binding, + ) + resumed.load_state_dict(flat_checkpoint) + continuation_grad = flat_grads[rank].flip(0).clone() + resumed_param.grad = continuation_grad.clone() + flat_param.grad = continuation_grad.clone() + resumed.step() + flat_optimizer.step() + flat_checkpoint_continuation = ( + torch.equal(resumed_param, flat_param) + and resumed.codebook_process_group_binding() is resumed_binding + and torch.equal(resumed._gefen_codebook, flat_optimizer._gefen_codebook) + ) + refresh_succeeded = flat_optimizer.refresh_codebook() + continuation_grads = tuple(gradient.flip(0) for gradient in flat_grads) + refresh_oracle = learn_gefen_exact_codebook_from_grad_periods( + grad_periods=tuple(("Flat", gradient, 4, gradient) for gradient in continuation_grads), + codebook_device=torch.device("cpu"), + num_codebooks=256, + force_endpoints=True, + verbose=False, + compute_mse_logging=False, + use_fused_histogram=False, + ) + refresh_matches_oracle = torch.equal(flat_optimizer._gefen_codebook, refresh_oracle) + refreshed_codebooks = [torch.empty_like(flat_optimizer._gefen_codebook) for _ in range(world)] + dist.all_gather(refreshed_codebooks, flat_optimizer._gefen_codebook) + refresh_agreement = all(torch.equal(item, refreshed_codebooks[0]) for item in refreshed_codebooks[1:]) + refresh_failure_grad = flat_grads[rank].roll(1).clone() + flat_param.grad = refresh_failure_grad + old_refresh_codebook_object = flat_optimizer._gefen_codebook + old_refresh_codebook = flat_optimizer._gefen_codebook.detach().clone() + old_refresh_indices = flat_optimizer.state[flat_param]["m_codebook"].detach().clone() + old_refresh_magnitude = flat_optimizer.state[flat_param]["m_magnitude"].detach().clone() + old_refresh_period = flat_optimizer.state[flat_param]["automatic_period"] + old_refresh_step = flat_optimizer._gefen_global_step + codebook_cache = flat_optimizer._gefen_codebook_by_device + lut_cache = flat_optimizer._gefen_codebook_lut_by_device + cache_marker = torch.tensor([17.0]) + lut_marker = torch.tensor([19.0]) + codebook_cache[torch.device("cpu")] = cache_marker + lut_cache[torch.device("cpu")] = lut_marker + original_nearest = gefen_module.gefen_nearest_codebook_indices + if rank == 1: + + def fail_nearest(*args, **kwargs): + raise RuntimeError("rank-local requantization failure") + + gefen_module.gefen_nearest_codebook_indices = fail_nearest + try: + flat_optimizer.refresh_codebook() + refresh_failure_seen = False + except RuntimeError as exc: + refresh_failure_seen = "requantization" in str(exc) + finally: + gefen_module.gefen_nearest_codebook_indices = original_nearest + refresh_failure_atomic = ( + flat_optimizer._gefen_codebook is old_refresh_codebook_object + and torch.equal(flat_optimizer._gefen_codebook, old_refresh_codebook) + and torch.equal( + flat_optimizer.state[flat_param]["m_codebook"], + old_refresh_indices, + ) + and torch.equal( + flat_optimizer.state[flat_param]["m_magnitude"], + old_refresh_magnitude, + ) + and flat_optimizer.state[flat_param]["automatic_period"] == old_refresh_period + and flat_optimizer._gefen_global_step == old_refresh_step + and flat_optimizer._gefen_codebook_by_device is codebook_cache + and flat_optimizer._gefen_codebook_lut_by_device is lut_cache + and codebook_cache[torch.device("cpu")] is cache_marker + and lut_cache[torch.device("cpu")] is lut_marker + ) + refresh_retry_succeeded = flat_optimizer.refresh_codebook() + + periodic_param = torch.nn.Parameter(torch.zeros(4)) + periodic_optimizer = Gefen( + [("periodic", periodic_param)], + fused=False, + factored_v_2d=False, + codebook_refresh_every=2, + ) + periodic_identity = ParameterIdentity("Periodic", (8,)) + periodic_records = tuple( + _flat(periodic_identity, group, member, index * 4, 4) for index, member in enumerate(members) + ) + _finalize( + periodic_optimizer, + periodic_param, + periodic_records[rank], + ShardingManifest(periodic_records), + _binding(group, rank, runtime_group), + ) + periodic_optimizer._resolve_automatic_period = lambda *args: 4 + periodic_param.grad = flat_grads[rank].clone() + periodic_optimizer.step() + first_periodic_codebook = periodic_optimizer._gefen_codebook + + if rank == 0: + periodic_param.grad = torch.sparse_coo_tensor( + torch.tensor([[0]]), + torch.tensor([1.0]), + size=(4,), + ) + else: + periodic_param.grad = flat_grads[rank].roll(1).clone() + try: + periodic_optimizer.step() + periodic_nondue_failure_rejected = False + except RuntimeError as exc: + periodic_nondue_failure_rejected = "gradient preflight" in str(exc) + periodic_nondue_failure_rejected = ( + periodic_nondue_failure_rejected + and periodic_optimizer._gefen_global_step == 1 + and periodic_optimizer._gefen_codebook is first_periodic_codebook + ) + + periodic_param.grad = flat_grads[rank].roll(1).clone() + periodic_optimizer.step() + periodic_nondue_failure_rejected = ( + periodic_nondue_failure_rejected + and periodic_optimizer._gefen_global_step == 2 + and periodic_optimizer._gefen_codebook is first_periodic_codebook + ) + periodic_param.grad = flat_grads[rank].flip(0).clone() + periodic_optimizer.step() + periodic_codebooks = [torch.empty_like(periodic_optimizer._gefen_codebook) for _ in range(world)] + dist.all_gather(periodic_codebooks, periodic_optimizer._gefen_codebook) + periodic_refresh_valid = ( + periodic_optimizer._gefen_global_step == 3 + and periodic_optimizer._gefen_codebook is not first_periodic_codebook + and all(torch.equal(item, periodic_codebooks[0]) for item in periodic_codebooks[1:]) + and periodic_optimizer.state_dict()["gefen_codebook_scope"]["refresh_every"] == 2 + ) + + # Whole-matrix owner mode keeps no fake tensor/state on nonowners. Every + # member still joins codebook collectives and receives common state; + # the adapter performs the separately declared post-step matrix sync. + owner_source = torch.nn.Parameter(torch.ones(2, 2)) + owner_optimizer = GefenMuon([("matrix", owner_source)], fused=False) + owner_identity = ParameterIdentity("Matrix", (2, 2)) + owner_records = tuple(_whole_owner(owner_identity, group, member, "rank:0") for member in members) + owner_optimizer.post_sharding( + ( + ParameterRebinding( + owner_source, + owner_source if rank == 0 else None, + owner_records[rank], + ), + ), + manifest=ShardingManifest(owner_records), + codebook_process_group=_binding(group, rank, runtime_group), + ) + if rank == 0: + owner_source.grad = torch.tensor([[1.0, -2.0], [3.0, -4.0]]) + owner_initialized = owner_optimizer.initialize_codebook() + owner_codebooks = [torch.empty_like(owner_optimizer._gefen_codebook) for _ in range(world)] + dist.all_gather(owner_codebooks, owner_optimizer._gefen_codebook) + owner_agreement = all(torch.equal(item, owner_codebooks[0]) for item in owner_codebooks[1:]) + owner_optimizer.step() + synchronized_matrix = owner_source.detach().clone() if rank == 0 else torch.empty(2, 2, dtype=torch.float32) + dist.broadcast(synchronized_matrix, src=0, group=runtime_group) + owner_step_valid = ( + owner_optimizer._gefen_global_step == 1 + and not torch.equal(synchronized_matrix, torch.ones_like(synchronized_matrix)) + and (rank != 0 or torch.equal(owner_source, synchronized_matrix)) + and (rank == 0 or (not owner_optimizer.param_groups[0]["params"] and not owner_optimizer.state)) + ) + owner_support = next( + item + for item in owner_optimizer.optimizer_contract().capabilities.training + if item.layout is ParameterLayout.WHOLE_PARAMETER_OWNER + ) + owner_contract_valid = owner_support.requires_post_step_parameter_sync + owner_amp_protocol = owner_optimizer._step_supports_amp_scaling + + # A failure after the shared histogram reduction must be observed by + # every member before staged periods or the codebook are committed. + failure_param = torch.nn.Parameter(torch.zeros(4)) + failure_optimizer = Gefen([("failure", failure_param)], fused=False, factored_v_2d=False) + failure_identity = ParameterIdentity("Failure", (8,)) + failure_records = tuple( + _flat(failure_identity, group, member, index * 4, 4) for index, member in enumerate(members) + ) + _finalize( + failure_optimizer, + failure_param, + failure_records[rank], + ShardingManifest(failure_records), + _binding(group, rank, runtime_group), + ) + failure_optimizer._resolve_automatic_period = lambda *args: 4 + failure_param.grad = flat_grads[rank].clone() + original_exact_dp = gefen_module.quantization_module.exact_dp + if rank == 1: + + def fail_exact_dp(*args, **kwargs): + raise RuntimeError("rank-local exact-DP failure") + + gefen_module.quantization_module.exact_dp = fail_exact_dp + try: + failure_optimizer.initialize_codebook() + failure_seen = False + except RuntimeError as exc: + failure_seen = "exact-DP" in str(exc) + finally: + gefen_module.quantization_module.exact_dp = original_exact_dp + failure_atomic = ( + failure_optimizer._gefen_codebook is None + and failure_optimizer._gefen_global_step == 0 + and failure_optimizer.state[failure_param] == {"name": "failure"} + ) + retry_succeeded = failure_optimizer.initialize_codebook() + restored_codebook = failure_optimizer._gefen_codebook + if rank == 0: + failure_optimizer._gefen_codebook = None + try: + failure_optimizer.refresh_codebook() + presence_mismatch_rejected = False + except RuntimeError as exc: + presence_mismatch_rejected = "old codebook differs" in str(exc) + if rank == 0: + failure_optimizer._gefen_codebook = restored_codebook + original_manifest = failure_optimizer._gefen_sharding_manifest + failure_optimizer._gefen_codebook_scope_validated = False + if rank == 0: + alternate_identity = ParameterIdentity("Alternate", (8,)) + alternate_records = tuple( + _flat(alternate_identity, group, member, index * 4, 4) for index, member in enumerate(members) + ) + failure_optimizer._gefen_sharding_manifest = ShardingManifest(alternate_records) + try: + failure_optimizer.initialize_codebook() + manifest_mismatch_rejected = False + except RuntimeError as exc: + manifest_mismatch_rejected = "manifest" in str(exc) + failure_optimizer._gefen_sharding_manifest = original_manifest + + queue.put( + { + "rank": rank, + "invalid_rejected": invalid_rejected, + "invalid_atomic": invalid_atomic, + "replicated_initialized": replicated_initialized, + "replicated_logical_once": replicated_logical_once, + "mismatch_rejected": mismatch_rejected, + "mismatch_atomic": mismatch_atomic, + "amp_overflow_atomic": amp_overflow_atomic, + "amp_mismatch_rejected": amp_mismatch_rejected, + "amp_protocol_rejected": amp_protocol_rejected, + "flat_initialized": flat_initialized, + "flat_global": flat_global, + "flat_native_claim": flat_native_claim, + "flat_mismatch_rejected": flat_mismatch_rejected, + "flat_mismatch_atomic": flat_mismatch_atomic, + "flat_agreement": flat_agreement, + "flat_checkpoint_continuation": flat_checkpoint_continuation, + "rank_neutral_checkpoint_scope": rank_neutral_checkpoint_scope, + "rank_local_checkpoint_identity": rank_local_checkpoint_identity, + "cross_member_guard": cross_member_guard, + "refresh_succeeded": refresh_succeeded, + "refresh_matches_oracle": refresh_matches_oracle, + "refresh_agreement": refresh_agreement, + "refresh_failure_seen": refresh_failure_seen, + "refresh_failure_atomic": refresh_failure_atomic, + "refresh_retry_succeeded": refresh_retry_succeeded, + "periodic_nondue_failure_rejected": periodic_nondue_failure_rejected, + "periodic_refresh_valid": periodic_refresh_valid, + "owner_initialized": owner_initialized, + "owner_agreement": owner_agreement, + "owner_step_valid": owner_step_valid, + "owner_contract_valid": owner_contract_valid, + "owner_amp_protocol": owner_amp_protocol, + "failure_seen": failure_seen, + "failure_atomic": failure_atomic, + "retry_succeeded": retry_succeeded, + "presence_mismatch_rejected": presence_mismatch_rejected, + "manifest_mismatch_rejected": manifest_mismatch_rejected, + } + ) + except Exception as exc: + queue.put({"rank": rank, "error": repr(exc)}) + finally: + if dist.is_available() and dist.is_initialized(): + dist.destroy_process_group() + + +def _run_workers(world=2): + context = mp.get_context("spawn") + queue = context.Queue() + fd, init_file = tempfile.mkstemp(prefix="gefen-codebook-scope-") + os.close(fd) + os.unlink(init_file) + processes = [ + context.Process( + target=_distributed_worker, + args=(rank, world, init_file, queue), + ) + for rank in range(world) + ] + try: + for process in processes: + process.start() + results = [queue.get(timeout=60) for _ in processes] + for process in processes: + process.join(timeout=10) + for process in processes: + if process.is_alive(): + process.terminate() + process.join(timeout=5) + pytest.fail("scoped codebook worker hung") + assert process.exitcode == 0 + return sorted(results, key=lambda item: item["rank"]) + finally: + for process in processes: + if process.is_alive(): + process.terminate() + process.join(timeout=5) + if os.path.exists(init_file): + os.unlink(init_file) + + +def test_explicit_gloo_scope_aggregates_logical_state_and_fails_atomically(): + results = _run_workers() + + assert all("error" not in item for item in results), results + for item in results: + assert item["invalid_rejected"], item + assert item["invalid_atomic"], item + assert item["replicated_initialized"], item + assert item["replicated_logical_once"], item + assert item["mismatch_rejected"], item + assert item["mismatch_atomic"], item + assert item["amp_overflow_atomic"], item + assert item["amp_mismatch_rejected"], item + assert item["amp_protocol_rejected"], item + assert item["flat_initialized"], item + assert item["flat_global"], item + assert item["flat_native_claim"], item + assert item["flat_mismatch_rejected"], item + assert item["flat_mismatch_atomic"], item + assert item["flat_agreement"], item + assert item["flat_checkpoint_continuation"], item + assert item["rank_neutral_checkpoint_scope"], item + assert item["rank_local_checkpoint_identity"], item + assert item["cross_member_guard"], item + assert item["refresh_succeeded"], item + assert item["refresh_matches_oracle"], item + assert item["refresh_agreement"], item + assert item["refresh_failure_seen"], item + assert item["refresh_failure_atomic"], item + assert item["refresh_retry_succeeded"], item + assert item["periodic_nondue_failure_rejected"], item + assert item["periodic_refresh_valid"], item + assert item["owner_initialized"], item + assert item["owner_agreement"], item + assert item["owner_step_valid"], item + assert item["owner_contract_valid"], item + assert item["owner_amp_protocol"], item + assert item["failure_seen"], item + assert item["failure_atomic"], item + assert item["retry_succeeded"], item + assert item["presence_mismatch_rejected"], item + assert item["manifest_mismatch_rejected"], item + + +def _subgroup_worker(rank, world, init_file, queue): + try: + dist.init_process_group( + "gloo", + init_method="file://{}".format(init_file), + rank=rank, + world_size=world, + timeout=timedelta(seconds=45), + ) + first_ranks = (0, 2) + second_ranks = (1, 3) + first_group = dist.new_group(ranks=list(first_ranks)) + second_group = dist.new_group(ranks=list(second_ranks)) + if rank in first_ranks: + group_index = 0 + group_ranks = first_ranks + runtime_group = first_group + canonical_grad = torch.tensor([-1.0, -0.5, 0.25, 1.0]) + else: + group_index = 1 + group_ranks = second_ranks + runtime_group = second_group + canonical_grad = torch.tensor([-1.0, -0.9, 0.8, 1.0]) + coordinate = group_ranks.index(rank) + members = ("slot:0", "slot:1") + identity = ProcessGroupIdentity("replica_subgroup:{}".format(group_index), members) + parameter_identity = ParameterIdentity("Subgroup{}.Weight".format(group_index), (4,)) + records = tuple(_replicated(parameter_identity, identity, member) for member in members) + parameter = torch.nn.Parameter(torch.zeros(4)) + optimizer = Gefen([("weight", parameter)], fused=False, factored_v_2d=False) + optimizer.post_sharding( + (ParameterRebinding(parameter, parameter, records[coordinate]),), + manifest=ShardingManifest(records), + codebook_process_group=CodebookProcessGroupBinding( + identity, + members[coordinate], + runtime_group, + torch.device("cpu"), + ), + ) + optimizer._resolve_automatic_period = lambda *args: 4 + parameter.grad = canonical_grad.clone() if coordinate == 0 else torch.tensor([-1.0, -0.2, 0.1, 1.0]) + optimizer.initialize_codebook() + queue.put( + { + "rank": rank, + "group": group_index, + "codebook": optimizer._gefen_codebook.tolist(), + } + ) + except Exception as exc: + queue.put({"rank": rank, "error": repr(exc)}) + finally: + if dist.is_available() and dist.is_initialized(): + dist.destroy_process_group() + + +def _run_subgroup_workers(): + context = mp.get_context("spawn") + queue = context.Queue() + fd, init_file = tempfile.mkstemp(prefix="gefen-codebook-subgroups-") + os.close(fd) + os.unlink(init_file) + processes = [ + context.Process( + target=_subgroup_worker, + args=(rank, 4, init_file, queue), + ) + for rank in range(4) + ] + try: + for process in processes: + process.start() + results = [queue.get(timeout=60) for _ in processes] + for process in processes: + process.join(timeout=10) + for process in processes: + if process.is_alive(): + process.terminate() + process.join(timeout=5) + pytest.fail("scoped codebook subgroup worker hung") + assert process.exitcode == 0 + return sorted(results, key=lambda item: item["rank"]) + finally: + for process in processes: + if process.is_alive(): + process.terminate() + process.join(timeout=5) + if os.path.exists(init_file): + os.unlink(init_file) + + +def test_explicit_gloo_subgroups_are_isolated_from_default_world(): + results = _run_subgroup_workers() + + assert all("error" not in item for item in results), results + first = [item["codebook"] for item in results if item["group"] == 0] + second = [item["codebook"] for item in results if item["group"] == 1] + assert first[0] == first[1] + assert second[0] == second[1] + assert first[0] != second[0] + + +def _nccl_empty_owner_worker(rank, world, init_file, queue): + try: + torch.cuda.set_device(rank) + dist.init_process_group( + "nccl", + init_method="file://{}".format(init_file), + rank=rank, + world_size=world, + timeout=timedelta(seconds=60), + ) + members = tuple("rank:{}".format(index) for index in range(world)) + group = ProcessGroupIdentity("nccl_owner", members) + identity = ParameterIdentity("Matrix", (2, 2)) + records = tuple(_whole_owner(identity, group, member, "rank:0") for member in members) + source = torch.nn.Parameter(torch.ones(2, 2, device=torch.device("cuda", rank))) + optimizer = GefenMuon([("matrix", source)], fused=False) + optimizer.post_sharding( + ( + ParameterRebinding( + source, + source if rank == 0 else None, + records[rank], + ), + ), + manifest=ShardingManifest(records), + codebook_process_group=CodebookProcessGroupBinding( + group, + members[rank], + dist.group.WORLD, + torch.device("cuda", rank), + ), + ) + if rank == 0: + source.grad = torch.tensor([[1.0, -2.0], [3.0, -4.0]], device=source.device) + initialized = optimizer.initialize_codebook() + local = optimizer._gefen_codebook.to(torch.device("cuda", rank)) + gathered = [torch.empty_like(local) for _ in range(world)] + dist.all_gather(gathered, local) + agreement = all(torch.equal(item, gathered[0]) for item in gathered[1:]) + optimizer.step() + queue.put( + { + "rank": rank, + "initialized": initialized, + "agreement": agreement, + "empty_nonowner": rank == 0 + or ( + not optimizer.param_groups[0]["params"] + and not optimizer.state + and optimizer._gefen_codebook.device.type == "cpu" + ), + "step": optimizer._gefen_global_step, + } + ) + except Exception as exc: + queue.put({"rank": rank, "error": repr(exc)}) + finally: + if dist.is_available() and dist.is_initialized(): + dist.destroy_process_group() + + +@pytest.mark.skipif( + not dist.is_available() or not dist.is_nccl_available() or torch.cuda.device_count() < 2, + reason="requires NCCL and two CUDA devices", +) +def test_nccl_scope_uses_explicit_collective_device_with_empty_nonowner(): + context = mp.get_context("spawn") + queue = context.Queue() + fd, init_file = tempfile.mkstemp(prefix="gefen-codebook-nccl-") + os.close(fd) + os.unlink(init_file) + processes = [ + context.Process( + target=_nccl_empty_owner_worker, + args=(rank, 2, init_file, queue), + ) + for rank in range(2) + ] + try: + for process in processes: + process.start() + results = [queue.get(timeout=90) for _ in processes] + for process in processes: + process.join(timeout=15) + for process in processes: + if process.is_alive(): + process.terminate() + process.join(timeout=5) + pytest.fail("NCCL scoped codebook worker hung") + assert process.exitcode == 0 + results.sort(key=lambda item: item["rank"]) + assert all("error" not in item for item in results), results + assert all(item["initialized"] for item in results) + assert all(item["agreement"] for item in results) + assert all(item["empty_nonowner"] for item in results) + assert all(item["step"] == 1 for item in results) + finally: + for process in processes: + if process.is_alive(): + process.terminate() + process.join(timeout=5) + if os.path.exists(init_file): + os.unlink(init_file) diff --git a/tests/test_optimizer_contracts.py b/tests/test_optimizer_contracts.py index 4157520..7dd2c31 100644 --- a/tests/test_optimizer_contracts.py +++ b/tests/test_optimizer_contracts.py @@ -158,10 +158,11 @@ def test_plain_contract_matches_live_persistent_state(factored_v_2d): if item.transport is CheckpointTransport.NATIVE_OPTIMIZER ) assert native.atomic_load + assert ParameterLayout.FLATTENED_ELEMENT_SHARD not in native.same_topology assert contract.capabilities.accepts_semantic_parameter_names assert not contract.capabilities.canonical_parameter_fqns assert not contract.capabilities.stable_shard_identity - assert not contract.capabilities.explicit_process_group_codebook_scope + 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 @@ -208,8 +209,22 @@ def test_plain_contract_matches_live_persistent_state(factored_v_2d): StateScope.OPTIMIZER_COMMON ) } - assert common == {"gefen_global_step", "gefen_codebook", "gefen_deterministic"} - assert common.issubset(state_dict) + assert common == { + "gefen_global_step", + "gefen_codebook", + "gefen_deterministic", + "gefen_codebook_scope", + } + required_common = { + field.name + for field in contract.state_layout.fields_for_scope( + StateScope.OPTIMIZER_COMMON + ) + if not field.optional + } + assert required_common.issubset(state_dict) + assert contract.state_layout.field("gefen_codebook_scope").optional + assert "gefen_codebook_scope" not in state_dict @pytest.mark.parametrize( @@ -337,6 +352,7 @@ def test_muon_contract_separates_mode_topology_and_state_extent( assert contract.implementation == "gefen.GefenMuon" assert contract.capabilities.supported_parameter_ranks == (2,) + assert contract.capabilities.explicit_process_group_codebook_scope native = next( item for item in contract.capabilities.checkpoints @@ -468,6 +484,10 @@ 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") + assert not contract.capabilities.explicit_process_group_codebook_scope + assert contract.children[0].contract.capabilities.explicit_process_group_codebook_scope + if backup_optimizer == "gefen": + assert contract.children[1].contract.capabilities.explicit_process_group_codebook_scope assert {field.name for field in contract.state_layout.fields} == { "backup_optimizer" } @@ -625,6 +645,7 @@ def test_derived_fields_are_explicitly_non_authoritative(): "_gefen_rank_local_payload_", "_gefen_rank_local_member", "_gefen_checkpoint_metadata", + "gefen_native_local_shards", } payload_field = optimizer.optimizer_contract().state_layout.field( "_gefen_rank_local_payload_3" diff --git a/tests/test_rebinding_cpu.py b/tests/test_rebinding_cpu.py index 200ae70..2322f77 100644 --- a/tests/test_rebinding_cpu.py +++ b/tests/test_rebinding_cpu.py @@ -188,7 +188,7 @@ def test_replicated_rebind_preserves_legacy_name_and_enables_identity_contract() assert contract.capabilities.shard_rebinding assert contract.capabilities.post_sharding assert not contract.capabilities.canonical_state_io - assert not contract.capabilities.explicit_process_group_codebook_scope + assert contract.capabilities.explicit_process_group_codebook_scope with pytest.raises(RuntimeError, match="already finalized"): optimizer.rebind_parameter(new, new, identity=identity) with pytest.raises(RuntimeError, match="cannot add"): @@ -749,7 +749,7 @@ def test_muon_whole_owner_post_sharding_prunes_nonowner_and_blocks_training(): contract = optimizer.optimizer_contract() assert contract.capabilities.canonical_parameter_fqns assert contract.capabilities.stable_shard_identity - assert not contract.capabilities.explicit_process_group_codebook_scope + assert contract.capabilities.explicit_process_group_codebook_scope assert all( support.layout is not ParameterLayout.WHOLE_PARAMETER_OWNER for support in contract.capabilities.training ) From ceeef2bb8e334105a5aa1d62f3913218f8bb1152 Mon Sep 17 00:00:00 2001 From: thad0ctor Date: Sun, 12 Jul 2026 21:47:36 -0700 Subject: [PATCH 06/52] Add exact-binding canonical state I/O --- docs/optimizer_contracts.md | 10 +- src/gefen/__init__.py | 11 + src/gefen/canonical.py | 158 ++++ src/gefen/contracts.py | 108 ++- src/gefen/gefen.py | 882 ++++++++++++++++++++++- src/gefen/gefen_muon.py | 24 + tests/test_canonical_state_cpu.py | 837 +++++++++++++++++++++ tests/test_codebook_scope_distributed.py | 59 ++ tests/test_rebinding_cpu.py | 2 +- 9 files changed, 2083 insertions(+), 8 deletions(-) create mode 100644 src/gefen/canonical.py create mode 100644 tests/test_canonical_state_cpu.py diff --git a/docs/optimizer_contracts.md b/docs/optimizer_contracts.md index 7b254b5..c5bca6b 100644 --- a/docs/optimizer_contracts.md +++ b/docs/optimizer_contracts.md @@ -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. DTensor stable identity, Hybrid composite rebinding, canonical checkpoint I/O, state movement, and offload remain unclaimed. +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. DTensor stable identity, Hybrid composite rebinding, topology-changing canonical checkpoint I/O, state movement, and offload remain unclaimed. ## Explicit learned-codebook process groups @@ -53,6 +53,14 @@ The optimizer owns one learned codebook and therefore accepts one scope. Histogr 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 ordered by native parameter group and slot, including replicated slots in a mixed optimizer and separately listed pruned nonowners; this prevents an equal-shaped checkpoint from another member, logical slice, or parameter ordering from being reinterpreted positionally. 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. Whole-owner checkpoint completeness, scoped DTensor rank-local transport, scope migration, topology-changing canonical I/O, and Hybrid-wide coordination are not claimed. +## Exact-binding canonical local state + +`export_canonical_state()` returns the versioned `gefen.bound_state` local-fragment format and `import_canonical_state()` restores it by exact FQN rather than native parameter IDs or group/slot position. `prepare_canonical_state_import()` performs complete schema, implementation, policy, manifest, shard, group-option, common-state, and tensor-geometry validation on a shadow optimizer without mutating the live instance; an adapter may synchronize preparation status in its own process group and then call `commit_canonical_state_import()`. Prepared imports are optimizer-specific, single-use, and rejected if live optimizer state changed after preparation. The convenience import is the local prepare-plus-commit transaction. Canonical methods bypass ordinary native state-dict hooks so arbitrary hook side effects cannot weaken the core fail-before-mutation boundary. + +The v1 document contains only primitive containers and detached tight finite CPU leaves of the ordinary dense `torch.Tensor` type, round-trips through `torch.load(weights_only=True)`, records the complete stable manifest, and maps each locally present parameter's authoritative state and algorithm-shaping group options by exact case-preserving FQN. After finalization the exact FQN is also the name used by period-routing policy, so compatibility names, devices, runtime process-group handles, derived lookup tables, capturable buffers, checkpoint carriers, and native positional IDs are not optimizer meaning in this format; the target retains its own lowercase compatibility names after import. Codebook scope identity is recorded, but its runtime handle and collective device must already be installed through `post_sharding`. + +This is an exact-binding transport-neutral local fragment, not the dense global logical state planned for portable DCP v3. Its dynamic `CANONICAL_LOCAL` checkpoint entry covers finalized plain-Gefen replicated and flattened local shards and finalized replicated GefenMuon, performs no collectives, reports atomic local import, and has an empty topology-changing set. A different member, slice, manifest, algorithm policy, group option, or declared state variant rejects. Export, preparation, and commit are quiescent checkpoint-boundary operations; prepared imports use content-bearing freshness tokens, including device counters, to reject intervening mutation. Export remains available after capturable warmup, but a capturable import target must still be fresh, before authoritative device state or a CUDA graph exists; importing replaces state identities, so an already captured graph cannot safely remain attached. Configurations with `stochastic_round=True` do not claim canonical v1 because the decomposed path intentionally lacks the fused stochastic quantizer, so changing effective fused availability would change the algorithm. DTensor, whole-owner completeness, Hybrid composition, rank-fragment gathering, resharding, world-size change, dense momentum decoding, and target-topology recompression remain unclaimed. + Plain Gefen declares replicated, flattened element-shard, and the narrow DTensor training layouts. Its PyTorch rank-local checkpoint transport is same-topology only. GefenMuon declares replicated and narrow DTensor training, with mode-specific state extents: `approx` state is local, `exact` state is logically global, and `distributed` momentum is held by the parameter owner while non-owners retain metadata only. Native Parallel-Muon checkpoints separately declare world-size owner redistribution, not placement-changing resharding. `GefenMuonHybrid` retains its nested child namespaces and does not flatten AdamW or Gefen child state into a fabricated common schema. Gefen and GefenMuon native loads, rank-local payload restoration through `load_state_dict`, and distributed-owner payload restoration prepare their complete core restore before changing local live optimizer state. Their transport entries report `atomic_load=True`. This is a per-optimizer-instance fail-before-mutation guarantee at the optimizer load boundary, not a coordinated all-rank commit or a guarantee over work an external checkpoint orchestrator performs before calling the optimizer. Load pre-hooks run before that boundary and load post-hooks run afterward, so arbitrary side effects in user hooks are also outside the guarantee. Hybrid composite loads do not yet provide the same guarantee and report `atomic_load=False`. diff --git a/src/gefen/__init__.py b/src/gefen/__init__.py index adabf3f..20951da 100644 --- a/src/gefen/__init__.py +++ b/src/gefen/__init__.py @@ -12,10 +12,13 @@ "GefenMuon", "GefenMuonHybrid", "CONTRACT_SCHEMA_VERSION", + "CANONICAL_STATE_FORMAT_VERSION", "IDENTITY_SCHEMA_VERSION", "CheckpointSupport", "CheckpointTransport", + "CanonicalStateProvider", "CodebookProcessGroupBinding", + "PreparedCanonicalStateImport", "OptimizerCapabilities", "OptimizerChildContract", "OptimizerContract", @@ -73,11 +76,19 @@ def __getattr__(name): from .codebook import CodebookProcessGroupBinding return CodebookProcessGroupBinding + if name in ( + "CANONICAL_STATE_FORMAT_VERSION", + "PreparedCanonicalStateImport", + ): + from . import canonical + + return getattr(canonical, name) if name in ( "CONTRACT_SCHEMA_VERSION", "IDENTITY_SCHEMA_VERSION", "CheckpointSupport", "CheckpointTransport", + "CanonicalStateProvider", "OptimizerCapabilities", "OptimizerChildContract", "OptimizerContract", diff --git a/src/gefen/canonical.py b/src/gefen/canonical.py new file mode 100644 index 0000000..2f3f45c --- /dev/null +++ b/src/gefen/canonical.py @@ -0,0 +1,158 @@ +"""Primitive, device-neutral helpers for canonical optimizer state.""" + +import math + +import torch + + +CANONICAL_STATE_FORMAT_VERSION = 1 +_IMPORT_PLAN_TOKEN = object() + + +class PreparedCanonicalStateImport: + """Opaque, single-use result of canonical import preparation.""" + + __slots__ = ("_optimizer", "_live_token", "_staged", "_consumed") + + def __init__(self, optimizer, live_token, staged, *, _token=None): + if _token is not _IMPORT_PLAN_TOKEN: + raise TypeError("PreparedCanonicalStateImport values are created by an optimizer") + self._optimizer = optimizer + self._live_token = live_token + self._staged = staged + self._consumed = False + + +def make_prepared_canonical_state_import(optimizer, live_token, staged): + return PreparedCanonicalStateImport( + optimizer, + live_token, + staged, + _token=_IMPORT_PLAN_TOKEN, + ) + + +def canonical_value_supported(value, *, finite_tensors=False) -> bool: + """Return whether ``value`` has a deterministic weights-only wire form.""" + + if type(value) is torch.Tensor: + supported = ( + value.layout is torch.strided + and not value.is_meta + and not value.is_nested + and not value.is_quantized + ) + if ( + supported + and finite_tensors + and (value.is_floating_point() or value.is_complex()) + ): + try: + supported = bool(torch.isfinite(value.detach()).all()) + except (NotImplementedError, RuntimeError, TypeError): + supported = False + return supported + if torch.is_tensor(value): + return False + if type(value) is float: + return math.isfinite(value) + if value is None or type(value) in {bool, int, str}: + return True + if type(value) in {list, tuple}: + return all( + canonical_value_supported( + item, finite_tensors=finite_tensors + ) + for item in value + ) + if type(value) is dict: + return all( + type(key) is str + and canonical_value_supported( + item, finite_tensors=finite_tensors + ) + for key, item in value.items() + ) + return False + + +def clone_canonical_value(value, *, path="value"): + """Clone one supported value into a CPU, weights-only-safe wire value.""" + + if torch.is_tensor(value): + if ( + type(value) is not torch.Tensor + or value.layout is not torch.strided + or value.is_meta + or value.is_nested + or value.is_quantized + ): + raise TypeError( + "{} must be a plain materialized strided tensor for canonical state".format( + path + ) + ) + cloned = ( + value.detach() + .to(device="cpu") + .resolve_conj() + .resolve_neg() + .contiguous() + .clone() + ) + if cloned.is_floating_point() or cloned.is_complex(): + try: + finite = bool(torch.isfinite(cloned).all()) + except (NotImplementedError, RuntimeError, TypeError) as exc: + raise ValueError( + "{} tensor dtype does not support finite canonical state".format( + path + ) + ) from exc + if not finite: + raise ValueError("{} tensor must be finite".format(path)) + return cloned + if type(value) is float: + if not math.isfinite(value): + raise ValueError("{} must be finite".format(path)) + return value + if value is None or type(value) in {bool, int, str}: + return value + if type(value) is list: + return [clone_canonical_value(item, path="{}[{}]".format(path, index)) for index, item in enumerate(value)] + if type(value) is tuple: + return tuple(clone_canonical_value(item, path="{}[{}]".format(path, index)) for index, item in enumerate(value)) + if type(value) is dict: + if any(type(key) is not str for key in value): + raise TypeError("{} dictionary keys must be strings".format(path)) + return {key: clone_canonical_value(value[key], path="{}.{}".format(path, key)) for key in sorted(value)} + raise TypeError("{} has unsupported canonical-state type {}".format(path, type(value).__name__)) + + +def canonical_values_equal(left, right) -> bool: + """Compare canonical values exactly while ignoring tensor device.""" + + if torch.is_tensor(left) or torch.is_tensor(right): + return ( + torch.is_tensor(left) + and torch.is_tensor(right) + and left.dtype == right.dtype + and tuple(left.shape) == tuple(right.shape) + and torch.equal( + left.detach().cpu().contiguous(), + right.detach().cpu().contiguous(), + ) + ) + if type(left) is not type(right): + return False + if type(left) is dict: + return set(left) == set(right) and all(canonical_values_equal(left[key], right[key]) for key in left) + if type(left) in {list, tuple}: + return len(left) == len(right) and all(canonical_values_equal(a, b) for a, b in zip(left, right)) + return left == right + + +__all__ = [ + "CANONICAL_STATE_FORMAT_VERSION", + "PreparedCanonicalStateImport", +] diff --git a/src/gefen/contracts.py b/src/gefen/contracts.py index f5a0cea..6f05b7a 100644 --- a/src/gefen/contracts.py +++ b/src/gefen/contracts.py @@ -100,6 +100,7 @@ class CheckpointTransport(str, Enum): NATIVE_OPTIMIZER = "native_optimizer" PYTORCH_RANK_LOCAL = "pytorch_rank_local" COMPOSITE_NATIVE = "composite_native" + CANONICAL_LOCAL = "canonical_local" class TopologyChange(str, Enum): @@ -764,6 +765,23 @@ def optimizer_contract(self) -> OptimizerContract: """Return the optimizer's immutable state and capability declaration.""" +@runtime_checkable +class CanonicalStateProvider(Protocol): + """Structural protocol for exact-binding canonical local state.""" + + def export_canonical_state(self): + """Return a versioned local canonical state fragment.""" + + def prepare_canonical_state_import(self, state): + """Validate and stage a canonical fragment without live mutation.""" + + def commit_canonical_state_import(self, prepared) -> None: + """Commit one still-current prepared import.""" + + def import_canonical_state(self, state) -> None: + """Prepare and commit a canonical fragment atomically.""" + + _ALL_PRECISIONS = frozenset( {Precision.FLOAT32, Precision.BFLOAT16, Precision.FLOAT16, Precision.FLOAT64} ) @@ -922,6 +940,7 @@ def _negative_capabilities( explicit_process_group_codebook_scope: bool = False, shard_rebinding: bool = False, post_sharding: bool = False, + canonical_state_io: bool = False, ) -> OptimizerCapabilities: return OptimizerCapabilities( training=training, @@ -934,7 +953,7 @@ def _negative_capabilities( explicit_process_group_codebook_scope=explicit_process_group_codebook_scope, shard_rebinding=shard_rebinding, post_sharding=post_sharding, - canonical_state_io=False, + canonical_state_io=canonical_state_io, atomic_state_movement=False, state_offload=False, ) @@ -947,7 +966,9 @@ def _gefen_contract( stable_shard_identity: bool = False, explicit_process_group_codebook_scope: bool = False, native_flattened_checkpoint: bool = False, + canonical_state_layouts: AbstractSet[ParameterLayout] = frozenset(), ) -> OptimizerContract: + canonical_state_layouts = _frozenset(canonical_state_layouts) block_fields = ( StateField("vmean", StateScope.PARAMETER, StateGeometry.BLOCK, True), StateField("vmean_step", StateScope.PARAMETER, StateGeometry.SCALAR, True), @@ -973,6 +994,13 @@ def _gefen_contract( StateExtent.METADATA_ONLY, initialized=False, ), + StateVariant( + "period_selected", + ("name", "automatic_period"), + layouts, + StateExtent.METADATA_ONLY, + initialized=False, + ), ] factored_names = tuple(field.name for field in factored_fields) block_names = tuple(field.name for field in block_fields) @@ -1092,7 +1120,7 @@ def _gefen_contract( ProcessGroupScope.NONE, ), ) - checkpoints = ( + checkpoints = [ CheckpointSupport( CheckpointTransport.NATIVE_OPTIMIZER, frozenset({ParameterLayout.REPLICATED}) @@ -1114,19 +1142,30 @@ def _gefen_contract( requires_collective=True, atomic_load=True, ), - ) + ] + if canonical_state_layouts: + checkpoints.append( + CheckpointSupport( + CheckpointTransport.CANONICAL_LOCAL, + canonical_state_layouts, + frozenset(), + ProcessGroupScope.NONE, + atomic_load=True, + ) + ) return OptimizerContract( implementation="gefen.Gefen", state_layout=OptimizerStateLayout(fields, tuple(variants)), capabilities=_negative_capabilities( training=training, - checkpoints=checkpoints, + checkpoints=tuple(checkpoints), supported_parameter_ranks=None, canonical_parameter_fqns=canonical_parameter_fqns, stable_shard_identity=stable_shard_identity, explicit_process_group_codebook_scope=explicit_process_group_codebook_scope, shard_rebinding=True, post_sharding=True, + canonical_state_io=bool(canonical_state_layouts), ), ) @@ -1160,7 +1199,9 @@ def _gefen_muon_contract( stable_shard_identity: bool = False, explicit_process_group_codebook_scope: bool = False, whole_parameter_owner: bool = False, + canonical_state_layouts: AbstractSet[ParameterLayout] = frozenset(), ) -> OptimizerContract: + canonical_state_layouts = _frozenset(canonical_state_layouts) sharded_modes = _frozenset(sharded_modes) normuon_modes = _frozenset(normuon_modes) non_normuon_modes = _frozenset(non_normuon_modes) @@ -1199,6 +1240,17 @@ def _gefen_muon_contract( sharded_mode=mode, ) ) + variants.append( + StateVariant( + "period_selected_replicated_" + mode, + ("name", "automatic_period"), + frozenset({ParameterLayout.REPLICATED}), + StateExtent.METADATA_ONLY, + initialized=False, + parameter_ranks=(2,), + sharded_mode=mode, + ) + ) if mode == "distributed": variants.append( StateVariant( @@ -1212,6 +1264,18 @@ def _gefen_muon_contract( sharded_mode=mode, ) ) + variants.append( + StateVariant( + "distributed_owner_period_selected", + ("name", "automatic_period"), + frozenset({_DTENSOR_LAYOUT}), + StateExtent.METADATA_ONLY, + role=ParameterStateRole.OWNER, + initialized=False, + parameter_ranks=(2,), + sharded_mode=mode, + ) + ) else: variants.append( StateVariant( @@ -1224,6 +1288,17 @@ def _gefen_muon_contract( sharded_mode=mode, ) ) + variants.append( + StateVariant( + "period_selected_dtensor_" + mode, + ("name", "automatic_period"), + frozenset({_DTENSOR_LAYOUT}), + StateExtent.METADATA_ONLY, + initialized=False, + parameter_ranks=(2,), + sharded_mode=mode, + ) + ) if whole_parameter_owner: variants.append( StateVariant( @@ -1237,6 +1312,18 @@ def _gefen_muon_contract( sharded_mode=mode, ) ) + variants.append( + StateVariant( + "period_selected_whole_owner_" + mode, + ("name", "automatic_period"), + frozenset({ParameterLayout.WHOLE_PARAMETER_OWNER}), + StateExtent.METADATA_ONLY, + role=ParameterStateRole.OWNER, + initialized=False, + parameter_ranks=(2,), + sharded_mode=mode, + ) + ) for mode_set, field_names, prefix in ( (non_normuon_modes, base_names, "quantized_muon"), (normuon_modes, normuon_names, "quantized_normuon"), @@ -1394,6 +1481,17 @@ def _gefen_muon_contract( atomic_load=True, ) ) + if canonical_state_layouts: + checkpoints.append( + CheckpointSupport( + CheckpointTransport.CANONICAL_LOCAL, + canonical_state_layouts, + frozenset(), + ProcessGroupScope.NONE, + required_sharded_modes=sharded_modes, + atomic_load=True, + ) + ) return OptimizerContract( implementation="gefen.GefenMuon", state_layout=OptimizerStateLayout(fields, tuple(variants)), @@ -1406,6 +1504,7 @@ def _gefen_muon_contract( explicit_process_group_codebook_scope=explicit_process_group_codebook_scope, shard_rebinding=True, post_sharding=True, + canonical_state_io=bool(canonical_state_layouts), ), ) @@ -1464,6 +1563,7 @@ def _hybrid_contract( "IDENTITY_SCHEMA_VERSION", "CheckpointSupport", "CheckpointTransport", + "CanonicalStateProvider", "OptimizerCapabilities", "OptimizerChildContract", "OptimizerContract", diff --git a/src/gefen/gefen.py b/src/gefen/gefen.py index 4cc3246..df81cf4 100644 --- a/src/gefen/gefen.py +++ b/src/gefen/gefen.py @@ -23,12 +23,21 @@ import torch import torch.nn as nn +from gefen.canonical import ( + CANONICAL_STATE_FORMAT_VERSION, + PreparedCanonicalStateImport, + canonical_value_supported, + canonical_values_equal, + clone_canonical_value, + make_prepared_canonical_state_import, +) from gefen.codebook import CodebookProcessGroupBinding from gefen.contracts import ( LogicalSlice, OptimizerContract, ParameterIdentity, ParameterLayout, + ParameterStateRole, PlacementKind, ProcessGroupIdentity, ShardIdentity, @@ -73,6 +82,34 @@ _CODEBOOK_SCOPE_FORMAT_VERSION = 1 _SCOPED_NATIVE_METADATA_VERSION = 4 _NATIVE_LOCAL_SHARDS_FORMAT_VERSION = 1 +_CANONICAL_PARAMETER_STATE_KEYS = frozenset( + { + "name", + "automatic_period", + "step", + "m_codebook", + "m_magnitude", + "vmean", + "vmean_step", + "v_row", + "v_col", + "factored_step", + "normuon_v", + "normuon_step", + } +) +_CANONICAL_DERIVED_PARAMETER_STATE_KEYS = frozenset( + { + "stepsize", + "_h_buf", + "_capt_scalars", + "_capt_consts", + "_capt_consts_key", + "_capt_stack", + "_capt_row", + "m_codebook_shape", + } +) def _rank_local_payload_key(global_rank: int) -> str: @@ -1022,6 +1059,8 @@ def __init__( ) self._factored_v_2d = factored_v_2d self._deterministic = deterministic + if type(codebook_refresh_every) is not int: + raise TypeError("codebook_refresh_every must be an integer") if codebook_refresh_every < 0: raise ValueError( "codebook_refresh_every must be >= 0 but is: {}".format( @@ -1150,11 +1189,13 @@ def optimizer_contract(self) -> OptimizerContract: """Return the immutable state-layout and integration capability contract.""" identity_ready = self._canonical_identity_ready() + canonical_state_layouts = self._canonical_state_layouts() return _gefen_contract( factored_v_2d=self._factored_v_2d, canonical_parameter_fqns=identity_ready, stable_shard_identity=identity_ready, explicit_process_group_codebook_scope=True, + canonical_state_layouts=canonical_state_layouts, native_flattened_checkpoint=( self._codebook_scope_ready() and any( @@ -1164,6 +1205,108 @@ def optimizer_contract(self) -> OptimizerContract: ), ) + def _canonical_state_variant_layout(self): + return _gefen_contract( + factored_v_2d=self._factored_v_2d, + ).state_layout + + def _canonical_state_layout_supported(self, layout) -> bool: + return layout in { + ParameterLayout.REPLICATED, + ParameterLayout.FLATTENED_ELEMENT_SHARD, + } + + @staticmethod + def _canonical_group_options_value(group): + return { + key: value + for key, value in group.items() + if key + not in { + "params", + "param_names", + "name", + "_gefen_checkpoint_metadata", + } + } + + def _canonical_state_layouts(self): + if not self._canonical_identity_ready(): + return frozenset() + if self._stochastic_round: + return frozenset() + if any( + parameter is None + or not self._canonical_state_layout_supported(shard.layout) + for parameter, shard in self._gefen_local_shard_bindings + ): + return frozenset() + if any( + not canonical_value_supported( + self._canonical_group_options_value(group), + finite_tensors=True, + ) + for group in self.param_groups + ): + return frozenset() + if not canonical_value_supported( + self._canonical_policy(), + finite_tensors=True, + ): + return frozenset() + if not canonical_value_supported( + self._gefen_codebook, + finite_tensors=True, + ): + return frozenset() + try: + global_step = self._canonical_common_global_step() + if type(global_step) is not int or global_step < 0: + return frozenset() + self._validate_loaded_native_state() + except Exception: + return frozenset() + state_layout = self._canonical_state_variant_layout() + for group in self.param_groups: + group_options = self._canonical_group_options_value(group) + for parameter in group["params"]: + state = self.state.get(parameter, {}) + if any( + key not in _CANONICAL_PARAMETER_STATE_KEYS + and key not in _CANONICAL_DERIVED_PARAMETER_STATE_KEYS + for key in state + ): + return frozenset() + if any( + not canonical_value_supported( + value, + finite_tensors=True, + ) + for key, value in state.items() + if key in _CANONICAL_PARAMETER_STATE_KEYS + ): + return frozenset() + canonical_state = { + key: value + for key, value in state.items() + if key in _CANONICAL_PARAMETER_STATE_KEYS and key != "name" + } + shard = self._gefen_shard_bindings[parameter] + try: + self._validate_canonical_parameter_semantics( + parameter, + shard, + group_options, + canonical_state, + state_layout, + global_step, + ) + except (TypeError, ValueError, RuntimeError): + return frozenset() + return frozenset( + shard.layout for _, shard in self._gefen_local_shard_bindings + ) + def _codebook_scope_ready(self) -> bool: return ( self._gefen_codebook_process_group is not None @@ -1325,6 +1468,130 @@ def _serialized_native_local_shard(shard): ], } + @staticmethod + def _serialized_canonical_shard(shard): + process_group = None + if shard.process_group is not None: + process_group = { + "schema_version": shard.process_group.schema_version, + "semantic_name": shard.process_group.semantic_name, + "ordered_members": list(shard.process_group.ordered_members), + } + return { + "schema_version": shard.schema_version, + "parameter": { + "schema_version": shard.parameter.schema_version, + "fqn": shard.parameter.fqn, + "global_shape": list(shard.parameter.global_shape), + }, + "layout": shard.layout.value, + "logical_slice": { + "flat_offset": shard.logical_slice.flat_offset, + "length": shard.logical_slice.length, + }, + "process_group": process_group, + "local_member": shard.local_member, + "owner": shard.owner, + "placements": [ + { + "mesh_axis": placement.mesh_axis, + "kind": placement.kind.value, + "coordinate": placement.coordinate, + "parts": placement.parts, + "parameter_dimension": placement.parameter_dimension, + } + for placement in shard.placements + ], + } + + @classmethod + def _normalize_serialized_canonical_shard(cls, record): + if not isinstance(record, dict) or set(record) != { + "schema_version", + "parameter", + "layout", + "logical_slice", + "process_group", + "local_member", + "owner", + "placements", + }: + raise ValueError("Gefen canonical shard has an invalid schema") + parameter_record = record["parameter"] + if not isinstance(parameter_record, dict) or set(parameter_record) != { + "schema_version", + "fqn", + "global_shape", + }: + raise ValueError("Gefen canonical parameter identity is invalid") + if not isinstance(parameter_record["global_shape"], list): + raise ValueError("Gefen canonical global_shape must be a list") + slice_record = record["logical_slice"] + if not isinstance(slice_record, dict) or set(slice_record) != { + "flat_offset", + "length", + }: + raise ValueError("Gefen canonical logical slice is invalid") + group_record = record["process_group"] + if group_record is not None and ( + not isinstance(group_record, dict) + or set(group_record) + != {"schema_version", "semantic_name", "ordered_members"} + or not isinstance(group_record["ordered_members"], list) + ): + raise ValueError("Gefen canonical process-group identity is invalid") + if not isinstance(record["placements"], list): + raise ValueError("Gefen canonical placements must be a list") + try: + parameter = ParameterIdentity( + parameter_record["fqn"], + tuple(parameter_record["global_shape"]), + schema_version=parameter_record["schema_version"], + ) + process_group = None + if group_record is not None: + process_group = ProcessGroupIdentity( + group_record["semantic_name"], + tuple(group_record["ordered_members"]), + schema_version=group_record["schema_version"], + ) + placements = [] + for placement_record in record["placements"]: + if not isinstance(placement_record, dict) or set( + placement_record + ) != { + "mesh_axis", + "kind", + "coordinate", + "parts", + "parameter_dimension", + }: + raise ValueError("invalid placement schema") + placements.append( + ShardPlacement( + placement_record["mesh_axis"], + PlacementKind(placement_record["kind"]), + placement_record["coordinate"], + placement_record["parts"], + placement_record["parameter_dimension"], + ) + ) + shard = ShardIdentity( + parameter, + ParameterLayout(record["layout"]), + LogicalSlice( + slice_record["flat_offset"], slice_record["length"] + ), + process_group=process_group, + local_member=record["local_member"], + owner=record["owner"], + placements=tuple(placements), + schema_version=record["schema_version"], + ) + except (TypeError, ValueError) as exc: + raise ValueError("Gefen canonical shard identity is invalid") from exc + return cls._serialized_canonical_shard(shard) + def _serialized_native_local_shards(self): if self._gefen_codebook_process_group is None: return None @@ -3331,15 +3598,19 @@ def _resolve_automatic_period( # Single decision point for a parameter's automatic period so the # codebook-learning pass and the step agree. force_1d_period_one short- # circuits 1D params to per-element (period==1) before the block search. + routing_name = param_name + shard = self._gefen_shard_bindings.get(param) + if shard is not None: + routing_name = shard.parameter.fqn if self._force_1d_period_one and param.ndim == 1: return 1 if self._force_2d_period_one and param.ndim == 2: return 1 if self._period_one_substrings: - lname = str(param_name).lower() + lname = str(routing_name).lower() if any(sub in lname for sub in self._period_one_substrings): return 1 - return self._predict_period_from_grad_sq(param_name, param, grad) + return self._predict_period_from_grad_sq(routing_name, param, grad) def _predict_period_from_grad_sq( self, param_name: str, param: torch.Tensor, grad: torch.Tensor @@ -5263,6 +5534,604 @@ def _step_automatic_merged(self, items) -> None: codebooks, list(merged_codebook.reshape(k, nblocks, period).unbind(0)) ) + def _canonical_policy(self): + return { + "factored_v_2d": self._factored_v_2d, + "force_1d_period_one": self._force_1d_period_one, + "force_2d_period_one": self._force_2d_period_one, + "period_one_substrings": list(self._period_one_substrings), + "stochastic_round": self._stochastic_round, + "codebook_refresh_every": self._codebook_refresh_every, + } + + def _canonical_common_global_step(self) -> int: + device_step = self._device_gefen_global_step() + step = self._gefen_global_step if device_step is None else device_step + if self.capturable and self._stochastic_round: + seed_steps = [ + int(seed.item()) for seed in self._sr_seed_by_device.values() + ] + if any(seed_step != step for seed_step in seed_steps): + raise RuntimeError( + "Gefen capturable stochastic-round seeds disagree with the " + "canonical optimizer global step" + ) + return step + + def _serialized_sharding_manifest(self): + return [ + self._serialized_canonical_shard(shard) + for shard in self._gefen_sharding_manifest.shards + ] + + def _canonical_live_entries(self): + entries = {} + for group in self.param_groups: + options = clone_canonical_value( + self._canonical_group_options_value(group), + path="parameter group options", + ) + for compatibility_name, parameter in self._iter_group_params_with_names( + group + ): + shard = self._gefen_shard_bindings[parameter] + entries[shard.parameter.fqn] = ( + parameter, + shard, + str(compatibility_name), + options, + ) + return entries + + @staticmethod + def _canonical_value_token(value): + if torch.is_tensor(value): + raw = ( + value.detach() + .to(device="cpu") + .resolve_conj() + .resolve_neg() + .contiguous() + .reshape(-1) + .view(torch.uint8) + .numpy() + .tobytes() + ) + try: + version = value._version + except RuntimeError: + version = None + return ( + "tensor", + id(value), + version, + hashlib.sha256(raw).digest(), + str(value.device), + str(value.dtype), + tuple(value.shape), + ) + if type(value) is dict: + return ( + "dict", + tuple( + (key, Gefen._canonical_value_token(value[key])) + for key in sorted(value, key=repr) + ), + ) + if type(value) in {list, tuple}: + return ( + type(value).__name__, + tuple(Gefen._canonical_value_token(item) for item in value), + ) + try: + hash(value) + token = value + except TypeError: + token = (id(value), repr(value)) + return (type(value).__name__, token) + + def _canonical_import_live_token(self): + groups = tuple( + ( + id(group), + tuple(id(parameter) for parameter in group["params"]), + tuple( + str(name) + for name, _ in self._iter_group_params_with_names(group) + ), + self._canonical_value_token( + self._canonical_group_options_value(group) + ), + ) + for group in self.param_groups + ) + states = tuple( + ( + id(parameter), + id(self.state.get(parameter)), + self._canonical_value_token(self.state.get(parameter, {})), + ) + for group in self.param_groups + for parameter in group["params"] + ) + return ( + id(self.param_groups), + id(self.state), + id(self.defaults), + self._canonical_value_token(self.defaults), + groups, + states, + self._gefen_global_step, + self._device_gefen_global_step(), + self._canonical_value_token(self._gefen_codebook), + self._canonical_value_token(self._gefen_global_step_by_device), + self._canonical_value_token(self._sr_seed_by_device), + self._canonical_value_token(self._canonical_policy()), + self._deterministic, + self.capturable, + self.fused, + self.verbose, + self._fused_build_ok, + id(self._gefen_codebook_process_group), + self._canonical_value_token(self._serialized_codebook_scope()), + id(self._gefen_sharding_manifest), + tuple( + (id(parameter), shard.sort_key) + for parameter, shard in self._gefen_local_shard_bindings + ), + ) + + def _canonical_parameter_state_matches_variant( + self, shard, group_options, parameter_state, state_layout + ) -> bool: + carries_momentum = any( + key in parameter_state for key in ("m_codebook", "m_magnitude") + ) + carries_normuon = any( + key in parameter_state for key in ("normuon_v", "normuon_step") + ) + if bool(group_options.get("normuon", False)): + if carries_momentum and not carries_normuon: + return False + elif carries_normuon: + return False + fields = frozenset({"name", *parameter_state}) + parameter_rank = len(shard.parameter.global_shape) + sharded_mode = group_options.get("sharded_mode") + for variant in state_layout.parameter_variants: + if frozenset(variant.fields) != fields: + continue + if shard.layout not in variant.layouts: + continue + if variant.role is not ParameterStateRole.ANY: + continue + if variant.sharded_mode != sharded_mode: + continue + if ( + variant.parameter_ranks is not None + and parameter_rank not in variant.parameter_ranks + ): + continue + if parameter_rank in variant.excluded_parameter_ranks: + continue + return True + return False + + def _canonical_period_must_be_one(self, parameter, shard) -> bool: + if self._force_1d_period_one and parameter.ndim == 1: + return True + if self._force_2d_period_one and parameter.ndim == 2: + return True + fqn = shard.parameter.fqn.lower() + return any( + substring in fqn for substring in self._period_one_substrings + ) + + def _validate_canonical_parameter_semantics( + self, + parameter, + shard, + group_options, + parameter_state, + state_layout, + global_step, + ) -> None: + carries_momentum = any( + key in parameter_state for key in ("m_codebook", "m_magnitude") + ) + carries_normuon = any( + key in parameter_state for key in ("normuon_v", "normuon_step") + ) + counters = {} + for counter_key in ( + "step", + "vmean_step", + "factored_step", + "normuon_step", + ): + if counter_key not in parameter_state: + continue + counter = self._validate_rank_local_counter( + counter_key, parameter_state[counter_key] + ) + counters[counter_key] = counter + if counter > global_step: + raise ValueError( + "Gefen canonical parameter counter {} exceeds the optimizer " + "global step".format(counter_key) + ) + if "step" in counters and any( + counters[counter_key] > counters["step"] + for counter_key in ( + "vmean_step", + "factored_step", + "normuon_step", + ) + if counter_key in counters + ): + raise ValueError( + "Gefen canonical secondary counter exceeds the parameter step" + ) + period = parameter_state.get("automatic_period") + if ( + period is not None + and self._canonical_period_must_be_one(parameter, shard) + and (type(period) is not int or period != 1) + ): + raise ValueError( + "Gefen canonical automatic period violates period-one policy" + ) + if bool(group_options.get("normuon", False)): + if carries_momentum and not carries_normuon: + raise ValueError( + "Gefen canonical initialized NorMuon state is incomplete" + ) + elif carries_normuon: + raise ValueError( + "Gefen canonical NorMuon state is invalid for the target policy" + ) + if not self._canonical_parameter_state_matches_variant( + shard, group_options, parameter_state, state_layout + ): + raise ValueError( + "Gefen canonical parameter state does not match a declared state " + "variant" + ) + + @staticmethod + def _assert_canonical_state_outside_cuda_capture(operation) -> None: + if torch.cuda.is_available() and torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "canonical state {} cannot run during CUDA capture".format( + operation + ) + ) + + def _assert_canonical_import_target_safe(self) -> None: + if not self.capturable: + return + device_step = self._device_gefen_global_step() + initialized_state = any( + any(key != "name" for key in self.state.get(parameter, {})) + for group in self.param_groups + for parameter in group["params"] + ) + if ( + self._capt_stacks is not None + or self._gefen_codebook is not None + or self._gefen_global_step != 0 + or (device_step is not None and device_step != 0) + or initialized_state + ): + raise RuntimeError( + "canonical import into a capturable optimizer requires a fresh " + "target before device state or CUDA graphs are initialized" + ) + + def export_canonical_state(self): + """Export an exact-binding, device-neutral local state fragment.""" + + self._assert_finalized_binding_layout() + self._assert_canonical_state_outside_cuda_capture("export") + if not self._canonical_state_layouts(): + raise RuntimeError( + "canonical state export requires a supported finalized identity " + "layout and primitive algorithm policy" + ) + entries = self._canonical_live_entries() + parameters = {} + for fqn in sorted(entries): + parameter, shard, compatibility_name, options = entries[fqn] + state = {} + for key, value in self.state.get(parameter, {}).items(): + if key in _CANONICAL_DERIVED_PARAMETER_STATE_KEYS: + continue + if key == "name": + continue + if key not in _CANONICAL_PARAMETER_STATE_KEYS: + raise RuntimeError( + "canonical state export found undeclared parameter state key " + "{!r}".format(key) + ) + state[key] = clone_canonical_value( + value, path="parameters.{}.state.{}".format(fqn, key) + ) + parameters[fqn] = { + "compatibility_name": compatibility_name, + "shard": self._serialized_canonical_shard(shard), + "group_options": clone_canonical_value( + options, path="parameters.{}.group_options".format(fqn) + ), + "state": state, + } + return { + "format": "gefen.bound_state", + "format_version": CANONICAL_STATE_FORMAT_VERSION, + "coverage": "local_optimizer_fragment", + "implementation": self.optimizer_contract().implementation, + "policy": clone_canonical_value( + self._canonical_policy(), path="policy" + ), + "common": { + "gefen_global_step": self._canonical_common_global_step(), + "gefen_codebook": clone_canonical_value( + self._gefen_codebook, path="common.gefen_codebook" + ), + "gefen_deterministic": self._deterministic, + "gefen_codebook_scope": clone_canonical_value( + self._serialized_codebook_scope(), + path="common.gefen_codebook_scope", + ), + }, + "manifest": clone_canonical_value( + self._serialized_sharding_manifest(), path="manifest" + ), + "parameters": parameters, + } + + def _normalize_canonical_state_import(self, state): + state = clone_canonical_value(state, path="canonical state") + expected_top = { + "format", + "format_version", + "coverage", + "implementation", + "policy", + "common", + "manifest", + "parameters", + } + if type(state) is not dict or set(state) != expected_top: + raise ValueError("Gefen canonical state has an invalid top-level schema") + if state["format"] != "gefen.bound_state": + raise ValueError("Unsupported Gefen canonical state format") + if ( + type(state["format_version"]) is not int + or state["format_version"] != CANONICAL_STATE_FORMAT_VERSION + ): + raise ValueError( + "Unsupported Gefen canonical state format_version: {}".format( + state["format_version"] + ) + ) + if state["coverage"] != "local_optimizer_fragment": + raise ValueError("Unsupported Gefen canonical state coverage") + implementation = self.optimizer_contract().implementation + if state["implementation"] != implementation: + raise ValueError( + "Gefen canonical state implementation does not match the target" + ) + live_policy = clone_canonical_value( + self._canonical_policy(), path="live policy" + ) + if not canonical_values_equal(state["policy"], live_policy): + raise ValueError( + "Gefen canonical state algorithm policy does not match the target" + ) + + common = state["common"] + if type(common) is not dict or set(common) != { + "gefen_global_step", + "gefen_codebook", + "gefen_deterministic", + "gefen_codebook_scope", + }: + raise ValueError("Gefen canonical common state has an invalid schema") + if ( + type(common["gefen_global_step"]) is not int + or common["gefen_global_step"] < 0 + ): + raise ValueError( + "Gefen canonical global step must be a nonnegative integer" + ) + if type(common["gefen_deterministic"]) is not bool: + raise ValueError("Gefen canonical deterministic policy must be a bool") + normalized_scope = self._normalize_serialized_codebook_scope( + common["gefen_codebook_scope"] + ) + if normalized_scope != self._serialized_codebook_scope(): + raise ValueError( + "Gefen canonical codebook scope does not match the live binding" + ) + common["gefen_codebook_scope"] = normalized_scope + + manifest = state["manifest"] + if type(manifest) is not list: + raise ValueError("Gefen canonical manifest must be a list") + normalized_manifest = [ + self._normalize_serialized_canonical_shard(record) + for record in manifest + ] + if normalized_manifest != self._serialized_sharding_manifest(): + raise ValueError( + "Gefen canonical manifest does not match the finalized target" + ) + state["manifest"] = normalized_manifest + + parameters = state["parameters"] + if type(parameters) is not dict: + raise ValueError("Gefen canonical parameters must be an FQN mapping") + live_entries = self._canonical_live_entries() + state_layout = self._canonical_state_variant_layout() + if set(parameters) != set(live_entries): + raise ValueError( + "Gefen canonical parameter FQNs do not match the finalized target" + ) + for fqn, record in parameters.items(): + if type(record) is not dict or set(record) != { + "compatibility_name", + "shard", + "group_options", + "state", + }: + raise ValueError( + "Gefen canonical parameter {!r} has an invalid schema".format( + fqn + ) + ) + if type(record["compatibility_name"]) is not str: + raise ValueError( + "Gefen canonical compatibility names must be strings" + ) + parameter, shard, _, live_options = live_entries[fqn] + normalized_shard = self._normalize_serialized_canonical_shard( + record["shard"] + ) + if normalized_shard != self._serialized_canonical_shard(shard): + raise ValueError( + "Gefen canonical parameter shard does not match the live binding" + ) + record["shard"] = normalized_shard + if not canonical_values_equal(record["group_options"], live_options): + raise ValueError( + "Gefen canonical parameter-group options do not match the target" + ) + parameter_state = record["state"] + if type(parameter_state) is not dict or any( + key == "name" or key not in _CANONICAL_PARAMETER_STATE_KEYS + for key in parameter_state + ): + raise ValueError( + "Gefen canonical parameter state has undeclared fields" + ) + self._validate_canonical_parameter_semantics( + parameter, + shard, + live_options, + parameter_state, + state_layout, + common["gefen_global_step"], + ) + record["state"] = clone_canonical_value( + parameter_state, path="parameters.{}.state".format(fqn) + ) + del parameter + return state + + def _canonical_native_state_dict(self, canonical_state): + native = self._base_state_dict_without_hooks() + native["state"] = {} + live_entries = self._canonical_live_entries() + for saved_group, live_group in zip(native["param_groups"], self.param_groups): + for parameter_id, parameter in zip( + saved_group["params"], live_group["params"] + ): + shard = self._gefen_shard_bindings[parameter] + record = canonical_state["parameters"][shard.parameter.fqn] + parameter_state = { + "name": self._param_name(parameter), + **clone_canonical_value( + record["state"], + path="parameters.{}.state".format(shard.parameter.fqn), + ), + } + native["state"][parameter_id] = parameter_state + + common = canonical_state["common"] + native["gefen_global_step"] = common["gefen_global_step"] + native["gefen_codebook"] = common["gefen_codebook"] + native["gefen_deterministic"] = common["gefen_deterministic"] + scope = common["gefen_codebook_scope"] + if scope is not None: + native["gefen_codebook_scope"] = scope + else: + native.pop("gefen_codebook_scope", None) + local_shards = self._serialized_native_local_shards() + if local_shards is not None: + native["gefen_native_local_shards"] = local_shards + else: + native.pop("gefen_native_local_shards", None) + metadata = { + "format_version": ( + _SCOPED_NATIVE_METADATA_VERSION if scope is not None else 1 + ), + "global_step": common["gefen_global_step"], + "codebook": common["gefen_codebook"], + "deterministic": common["gefen_deterministic"], + "device_anchor": self._checkpoint_device_anchor(), + } + if scope is not None: + metadata["codebook_scope"] = scope + if local_shards is not None: + metadata["native_local_shards"] = local_shards + for group in native["param_groups"]: + group["_gefen_checkpoint_metadata"] = metadata + del live_entries + return native + + def _preserve_canonical_target_configuration(self, staged) -> None: + staged.defaults = self.defaults.copy() + staged.param_groups = self.param_groups + + def prepare_canonical_state_import(self, state): + """Validate and stage a canonical local import without live mutation.""" + + self._assert_finalized_binding_layout() + self._assert_canonical_state_outside_cuda_capture("import preparation") + self._assert_canonical_import_target_safe() + if not self._canonical_state_layouts(): + raise RuntimeError( + "canonical state import requires a supported finalized identity layout" + ) + normalized = self._normalize_canonical_state_import(state) + live_token = self._canonical_import_live_token() + staged = self._stage_load_state_dict( + self._canonical_native_state_dict(normalized) + ) + self._preserve_canonical_target_configuration(staged) + return make_prepared_canonical_state_import(self, live_token, staged) + + def commit_canonical_state_import(self, prepared) -> None: + """Commit one still-current prepared canonical import exactly once.""" + + if not isinstance(prepared, PreparedCanonicalStateImport): + raise TypeError( + "prepared must be a PreparedCanonicalStateImport from this optimizer" + ) + if prepared._optimizer is not self: + raise ValueError("prepared canonical state belongs to another optimizer") + if prepared._consumed: + raise RuntimeError("prepared canonical state import was already consumed") + self._assert_finalized_binding_layout() + self._assert_canonical_state_outside_cuda_capture("import commit") + self._assert_canonical_import_target_safe() + if prepared._live_token != self._canonical_import_live_token(): + raise RuntimeError( + "live optimizer state changed after canonical import preparation" + ) + prepared._consumed = True + self._commit_staged_load_state_dict(prepared._staged) + + def import_canonical_state(self, state) -> None: + """Atomically prepare and commit an exact-binding canonical fragment.""" + + self.commit_canonical_state_import( + self.prepare_canonical_state_import(state) + ) + + canonical_state_dict = export_canonical_state + load_canonical_state_dict = import_canonical_state + def state_dict(self): """Run optimizer state-dict hooks around Gefen's complete schema.""" @@ -6109,6 +6978,10 @@ def _validate_loaded_native_state(self) -> None: signature, self._gefen_codebook, allow_legacy_vmean_counter=True, + allow_preinitialized_periods=( + self._gefen_global_step == 0 + and self._gefen_codebook is not None + ), ) for group in self.param_groups: self._validate_group_options( @@ -6193,6 +7066,7 @@ def _validate_rank_local_states( codebook, *, allow_legacy_vmean_counter: bool = False, + allow_preinitialized_periods: bool = False, ) -> None: if not isinstance(states, list) or len(states) != len(signature): raise ValueError( @@ -6263,6 +7137,10 @@ def _validate_rank_local_states( param_signature.get("sharded") and param_signature.get("sharded_mode") == "distributed" ) + and not ( + allow_preinitialized_periods + and not any(key in pstate for key in initialized_keys) + ) ): raise ValueError( "Gefen rank-local checkpoint automatic_period is invalid without " diff --git a/src/gefen/gefen_muon.py b/src/gefen/gefen_muon.py index 7880fbd..9f179af 100644 --- a/src/gefen/gefen_muon.py +++ b/src/gefen/gefen_muon.py @@ -768,6 +768,7 @@ 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, + canonical_state_layouts=self._canonical_state_layouts(), whole_parameter_owner=( self._codebook_scope_ready() and any( @@ -777,6 +778,29 @@ def optimizer_contract(self) -> OptimizerContract: ), ) + def _canonical_state_layout_supported(self, layout) -> bool: + return layout is ParameterLayout.REPLICATED + + def _canonical_state_variant_layout(self): + sharded_modes = frozenset( + group["sharded_mode"] for group in self.param_groups + ) + normuon_modes = frozenset( + group["sharded_mode"] + for group in self.param_groups + if group.get("normuon", False) + ) + non_normuon_modes = frozenset( + group["sharded_mode"] + for group in self.param_groups + if not group.get("normuon", False) + ) + return _gefen_muon_contract( + sharded_modes=sharded_modes, + normuon_modes=normuon_modes, + non_normuon_modes=non_normuon_modes, + ).state_layout + def _validate_rebinding_layout(self, rebinding) -> None: shard = rebinding.shard target = rebinding.new_parameter diff --git a/tests/test_canonical_state_cpu.py b/tests/test_canonical_state_cpu.py new file mode 100644 index 0000000..fa4c7a1 --- /dev/null +++ b/tests/test_canonical_state_cpu.py @@ -0,0 +1,837 @@ +import copy +import io + +import pytest +import torch + +from gefen import ( + CANONICAL_STATE_FORMAT_VERSION, + CanonicalStateProvider, + CheckpointTransport, + Gefen, + GefenMuon, + LogicalSlice, + ParameterIdentity, + ParameterLayout, + ParameterRebinding, + PlacementKind, + PreparedCanonicalStateImport, + ProcessGroupIdentity, + ProcessGroupScope, + ShardIdentity, + ShardPlacement, + ShardingManifest, +) + + +def _replicated(identity, group=None, member=None): + placements = () + if group is not None: + coordinate = group.ordered_members.index(member) + placements = ( + ShardPlacement( + "dp", + PlacementKind.REPLICATE, + coordinate, + len(group.ordered_members), + ), + ) + return ShardIdentity( + identity, + ParameterLayout.REPLICATED, + LogicalSlice.full(identity), + process_group=group, + local_member=member, + placements=placements, + ) + + +def _flat(identity, group, member, offset, length): + coordinate = group.ordered_members.index(member) + return ShardIdentity( + identity, + ParameterLayout.FLATTENED_ELEMENT_SHARD, + LogicalSlice(offset, length), + process_group=group, + local_member=member, + placements=( + ShardPlacement( + "dp", + PlacementKind.FLAT_SHARD, + coordinate, + len(group.ordered_members), + ), + ), + ) + + +def _finalize(optimizer, bindings, manifest): + optimizer.post_sharding( + tuple(ParameterRebinding(parameter, parameter, shard) for parameter, shard in bindings), + manifest=manifest, + ) + + +def _snapshot(optimizer): + return { + "dict": optimizer.__dict__.copy(), + "groups": optimizer.param_groups, + "state": optimizer.state, + "codebook": optimizer._gefen_codebook, + "global_step": optimizer._gefen_global_step, + } + + +def _assert_snapshot_identity(optimizer, snapshot): + assert optimizer.param_groups is snapshot["groups"] + assert optimizer.state is snapshot["state"] + assert optimizer._gefen_codebook is snapshot["codebook"] + assert optimizer._gefen_global_step is snapshot["global_step"] + assert optimizer.__dict__.keys() == snapshot["dict"].keys() + for key, value in snapshot["dict"].items(): + assert optimizer.__dict__[key] is value + + +def _two_parameter_source(): + first = torch.nn.Parameter(torch.arange(1, 5, dtype=torch.float32)) + second = torch.nn.Parameter(torch.arange(11, 15, dtype=torch.float32)) + optimizer = Gefen( + [("first", first), ("second", second)], + fused=False, + factored_v_2d=False, + ) + first_identity = ParameterIdentity("Model.First", (4,)) + second_identity = ParameterIdentity("Model.Second", (4,)) + first_shard = _replicated(first_identity) + second_shard = _replicated(second_identity) + manifest = ShardingManifest((first_shard, second_shard)) + _finalize( + optimizer, + ((first, first_shard), (second, second_shard)), + manifest, + ) + optimizer._resolve_automatic_period = lambda *args: 4 + first.grad = torch.tensor([1.0, 2.0, 3.0, 4.0]) + second.grad = torch.tensor([-8.0, 2.0, 1.0, 3.0]) + optimizer.step() + return optimizer, first, second, first_shard, second_shard, manifest + + +def _reordered_target(first_value, second_value, first_shard, second_shard, manifest): + second = torch.nn.Parameter(second_value.detach().clone()) + first = torch.nn.Parameter(first_value.detach().clone()) + optimizer = Gefen( + [ + {"params": [("second_target", second)]}, + {"params": [("first_target", first)]}, + ], + fused=False, + factored_v_2d=False, + ) + _finalize( + optimizer, + ((second, second_shard), (first, first_shard)), + manifest, + ) + return optimizer, first, second + + +def test_canonical_local_capability_is_dynamic_and_exact_binding_only(): + parameter = torch.nn.Parameter(torch.ones(4)) + optimizer = Gefen([("p", parameter)], fused=False, factored_v_2d=False) + assert isinstance(optimizer, CanonicalStateProvider) + assert not optimizer.optimizer_contract().capabilities.canonical_state_io + assert all( + support.transport is not CheckpointTransport.CANONICAL_LOCAL + for support in optimizer.optimizer_contract().capabilities.checkpoints + ) + + identity = ParameterIdentity("Model.P", (4,)) + shard = _replicated(identity) + _finalize(optimizer, ((parameter, shard),), ShardingManifest((shard,))) + contract = optimizer.optimizer_contract() + support = next( + item for item in contract.capabilities.checkpoints if item.transport is CheckpointTransport.CANONICAL_LOCAL + ) + assert contract.capabilities.canonical_state_io + assert support.same_topology == frozenset({ParameterLayout.REPLICATED}) + assert not support.topology_changing + assert support.process_group_scope is ProcessGroupScope.NONE + assert support.atomic_load + assert not support.requires_collective + + +def test_pristine_export_is_primitive_device_neutral_and_excludes_derived_state(): + parameter = torch.nn.Parameter(torch.ones(8)) + optimizer = Gefen([("p", parameter)], fused=False, factored_v_2d=False) + identity = ParameterIdentity("Model.P", (8,)) + shard = _replicated(identity) + _finalize(optimizer, ((parameter, shard),), ShardingManifest((shard,))) + optimizer.state[parameter]["stepsize"] = torch.ones(1) + + exported = optimizer.export_canonical_state() + + assert exported["format"] == "gefen.bound_state" + assert exported["format_version"] == CANONICAL_STATE_FORMAT_VERSION + assert exported["coverage"] == "local_optimizer_fragment" + assert exported["manifest"][0]["process_group"] is None + assert set(exported["parameters"]) == {"Model.P"} + record = exported["parameters"]["Model.P"] + assert record["state"] == {} + assert "stepsize" not in record["state"] + assert "name" not in record["state"] + assert record["compatibility_name"] == "p" + + buffer = io.BytesIO() + torch.save(exported, buffer) + buffer.seek(0) + loaded = torch.load(buffer, weights_only=True) + assert loaded["format_version"] == CANONICAL_STATE_FORMAT_VERSION + assert loaded["parameters"]["Model.P"]["state"] == {} + + +def test_pristine_continuation_routes_period_policy_by_canonical_fqn(): + source_parameter = torch.nn.Parameter(torch.arange(1, 9, dtype=torch.float32)) + source = Gefen( + [("special_weight", source_parameter)], + fused=False, + factored_v_2d=False, + period_one_substrings=("special",), + ) + identity = ParameterIdentity("Model.SpecialWeight", (8,)) + shard = _replicated(identity) + manifest = ShardingManifest((shard,)) + _finalize(source, ((source_parameter, shard),), manifest) + + target_parameter = torch.nn.Parameter(source_parameter.detach().clone()) + target = Gefen( + [("unrelated_target_name", target_parameter)], + fused=False, + factored_v_2d=False, + period_one_substrings=("special",), + ) + _finalize(target, ((target_parameter, shard),), manifest) + target.import_canonical_state(source.export_canonical_state()) + + gradient = torch.arange(1, 9, dtype=torch.float32) + source_parameter.grad = gradient.clone() + target_parameter.grad = gradient.clone() + source.step() + target.step() + assert source.state[source_parameter]["automatic_period"] == 1 + assert target.state[target_parameter]["automatic_period"] == 1 + assert torch.equal(target_parameter, source_parameter) + assert torch.equal(target._gefen_codebook, source._gefen_codebook) + + +def test_informational_empty_compatibility_name_round_trips_but_stochastic_policy_is_unclaimed(): + parameter = torch.nn.Parameter(torch.ones(4)) + optimizer = Gefen([("", parameter)], fused=False, factored_v_2d=False) + identity = ParameterIdentity("Model.P", (4,)) + shard = _replicated(identity) + manifest = ShardingManifest((shard,)) + _finalize(optimizer, ((parameter, shard),), manifest) + exported = optimizer.export_canonical_state() + assert exported["parameters"]["Model.P"]["compatibility_name"] == "" + optimizer.import_canonical_state(exported) + assert optimizer.state[parameter]["name"] == "" + + stochastic_parameter = torch.nn.Parameter(torch.ones(4)) + with pytest.warns(RuntimeWarning, match="stochastic_round=True"): + stochastic = Gefen( + [("p", stochastic_parameter)], + fused=False, + factored_v_2d=False, + stochastic_round=True, + ) + _finalize(stochastic, ((stochastic_parameter, shard),), manifest) + assert not stochastic.optimizer_contract().capabilities.canonical_state_io + with pytest.raises(RuntimeError, match="supported finalized"): + stochastic.export_canonical_state() + + +def test_initialized_import_maps_by_fqn_across_order_and_group_boundaries(): + source, source_first, source_second, first_shard, second_shard, manifest = _two_parameter_source() + exported = source.export_canonical_state() + exported_magnitude = exported["parameters"]["Model.First"]["state"][ + "m_magnitude" + ] + assert exported_magnitude is not source.state[source_first]["m_magnitude"] + assert ( + exported_magnitude.untyped_storage().nbytes() + == exported_magnitude.numel() * exported_magnitude.element_size() + ) + target, target_first, target_second = _reordered_target( + source_first, + source_second, + first_shard, + second_shard, + manifest, + ) + load_hooks = [] + target.register_load_state_dict_pre_hook(lambda *args: load_hooks.append("pre")) + target.register_load_state_dict_post_hook(lambda *args: load_hooks.append("post")) + input_codebook = exported["common"]["gefen_codebook"] + input_magnitude = exported["parameters"]["Model.First"]["state"]["m_magnitude"] + input_codebook_value = input_codebook.clone() + input_magnitude_value = input_magnitude.clone() + + prepared = target.prepare_canonical_state_import(exported) + assert isinstance(prepared, PreparedCanonicalStateImport) + assert target._gefen_codebook is None + assert exported["common"]["gefen_codebook"] is input_codebook + assert exported["parameters"]["Model.First"]["state"]["m_magnitude"] is input_magnitude + assert torch.equal(input_codebook, input_codebook_value) + assert torch.equal(input_magnitude, input_magnitude_value) + target.commit_canonical_state_import(prepared) + + assert not load_hooks + assert torch.equal( + target.state[target_first]["m_magnitude"], + source.state[source_first]["m_magnitude"], + ) + assert torch.equal( + target.state[target_second]["m_magnitude"], + source.state[source_second]["m_magnitude"], + ) + assert target.state[target_first]["name"] == "first_target" + assert target.state[target_second]["name"] == "second_target" + + first_grad = torch.tensor([4.0, -3.0, 2.0, -1.0]) + second_grad = torch.tensor([1.0, 3.0, -5.0, 7.0]) + source_first.grad = first_grad.clone() + target_first.grad = first_grad.clone() + source_second.grad = second_grad.clone() + target_second.grad = second_grad.clone() + source.step() + target.step() + assert torch.equal(target_first, source_first) + assert torch.equal(target_second, source_second) + assert torch.equal(target._gefen_codebook, source._gefen_codebook) + + with pytest.raises(RuntimeError, match="already consumed"): + target.commit_canonical_state_import(prepared) + + +def test_prepared_import_rejects_wrong_optimizer_and_stale_live_state(): + source, source_first, source_second, first_shard, second_shard, manifest = _two_parameter_source() + exported = source.export_canonical_state() + target, _, _ = _reordered_target( + source_first, + source_second, + first_shard, + second_shard, + manifest, + ) + other, _, _ = _reordered_target( + source_first, + source_second, + first_shard, + second_shard, + manifest, + ) + prepared = target.prepare_canonical_state_import(exported) + + other_before = _snapshot(other) + with pytest.raises(ValueError, match="another optimizer"): + other.commit_canonical_state_import(prepared) + _assert_snapshot_identity(other, other_before) + + target._gefen_global_step += 1 + target_before = _snapshot(target) + with pytest.raises(RuntimeError, match="changed after"): + target.commit_canonical_state_import(prepared) + _assert_snapshot_identity(target, target_before) + + +@pytest.mark.parametrize( + "corrupt", + [ + "missing_parameter", + "wrong_manifest", + "wrong_shard", + "wrong_policy", + "wrong_group_options", + "unknown_state", + "bad_state_geometry", + "counter_ahead", + "secondary_counter_ahead", + "bool_version", + ], +) +def test_canonical_corruption_rejects_before_live_mutation(corrupt): + source, source_first, source_second, first_shard, second_shard, manifest = _two_parameter_source() + exported = source.export_canonical_state() + target, _, _ = _reordered_target( + source_first, + source_second, + first_shard, + second_shard, + manifest, + ) + damaged = copy.deepcopy(exported) + if corrupt == "missing_parameter": + damaged["parameters"].pop("Model.Second") + elif corrupt == "wrong_manifest": + damaged["manifest"][0]["parameter"]["fqn"] = "Wrong.First" + elif corrupt == "wrong_shard": + damaged["parameters"]["Model.First"]["shard"] = copy.deepcopy(damaged["parameters"]["Model.Second"]["shard"]) + elif corrupt == "wrong_policy": + damaged["policy"]["stochastic_round"] = True + elif corrupt == "wrong_group_options": + damaged["parameters"]["Model.First"]["group_options"]["lr"] = 9.0 + elif corrupt == "unknown_state": + damaged["parameters"]["Model.First"]["state"]["exp_avg"] = torch.ones(4) + elif corrupt == "bad_state_geometry": + damaged["parameters"]["Model.First"]["state"]["m_magnitude"] = torch.ones(2, 1) + elif corrupt == "counter_ahead": + damaged["parameters"]["Model.First"]["state"]["step"] = ( + damaged["common"]["gefen_global_step"] + 1 + ) + elif corrupt == "secondary_counter_ahead": + damaged["common"]["gefen_global_step"] = 100 + damaged["parameters"]["Model.First"]["state"]["step"] = 1 + damaged["parameters"]["Model.First"]["state"]["vmean_step"] = 2 + else: + damaged["format_version"] = True + before = _snapshot(target) + + with pytest.raises((TypeError, ValueError, RuntimeError)): + target.import_canonical_state(damaged) + + _assert_snapshot_identity(target, before) + + +def test_import_rejects_initialized_period_that_violates_forced_policy(): + source_parameter = torch.nn.Parameter(torch.arange(8, dtype=torch.float32)) + source = Gefen( + [("p", source_parameter)], + fused=False, + factored_v_2d=False, + force_1d_period_one=True, + ) + identity = ParameterIdentity("Model.P", (8,)) + shard = _replicated(identity) + manifest = ShardingManifest((shard,)) + _finalize(source, ((source_parameter, shard),), manifest) + source_parameter.grad = torch.arange(1, 9, dtype=torch.float32) + source.step() + damaged = source.export_canonical_state() + parameter_state = damaged["parameters"]["Model.P"]["state"] + parameter_state["automatic_period"] = 2 + parameter_state["m_codebook"] = torch.zeros(4, 2, dtype=torch.uint8) + parameter_state["m_magnitude"] = torch.ones(4, 1, dtype=torch.float32) + parameter_state["vmean"] = torch.ones(4, 1, dtype=torch.float32) + + target_parameter = torch.nn.Parameter(source_parameter.detach().clone()) + target = Gefen( + [("target", target_parameter)], + fused=False, + factored_v_2d=False, + force_1d_period_one=True, + ) + _finalize(target, ((target_parameter, shard),), manifest) + before = _snapshot(target) + with pytest.raises(ValueError, match="violates period-one policy"): + target.import_canonical_state(damaged) + _assert_snapshot_identity(target, before) + + +def test_codebook_initialized_before_first_step_round_trips_canonically(): + parameter = torch.nn.Parameter(torch.arange(8, dtype=torch.float32)) + source = Gefen([("p", parameter)], fused=False, factored_v_2d=False) + identity = ParameterIdentity("Model.P", (8,)) + shard = _replicated(identity) + manifest = ShardingManifest((shard,)) + _finalize(source, ((parameter, shard),), manifest) + source._resolve_automatic_period = lambda *args: 4 + parameter.grad = torch.arange(1, 9, dtype=torch.float32) + assert source.initialize_codebook() + assert set(source.state[parameter]) == {"name", "automatic_period"} + + target_parameter = torch.nn.Parameter(parameter.detach().clone()) + target = Gefen([("target", target_parameter)], fused=False, factored_v_2d=False) + _finalize(target, ((target_parameter, shard),), manifest) + target.import_canonical_state(source.export_canonical_state()) + + assert target._gefen_global_step == 0 + assert torch.equal(target._gefen_codebook, source._gefen_codebook) + assert target.state[target_parameter] == { + "name": "target", + "automatic_period": 4, + } + + +def test_flat_fragment_is_fqn_keyed_but_rejects_another_member_slice(): + group = ProcessGroupIdentity("flat", ("rank:0", "rank:1")) + identity = ParameterIdentity("Model.Flat", (8,)) + shards = ( + _flat(identity, group, "rank:0", 0, 4), + _flat(identity, group, "rank:1", 4, 4), + ) + manifest = ShardingManifest(shards) + source_parameter = torch.nn.Parameter(torch.arange(4, dtype=torch.float32)) + source = Gefen([("flat", source_parameter)], fused=False, factored_v_2d=False) + _finalize(source, ((source_parameter, shards[0]),), manifest) + source._resolve_automatic_period = lambda *args: 4 + source_parameter.grad = torch.tensor([1.0, 2.0, 3.0, 4.0]) + source.step() + exported = source.export_canonical_state() + support = next( + item + for item in source.optimizer_contract().capabilities.checkpoints + if item.transport is CheckpointTransport.CANONICAL_LOCAL + ) + assert support.same_topology == frozenset({ParameterLayout.FLATTENED_ELEMENT_SHARD}) + + target_parameter = torch.nn.Parameter(torch.arange(4, dtype=torch.float32)) + target = Gefen([("flat", target_parameter)], fused=False, factored_v_2d=False) + _finalize(target, ((target_parameter, shards[1]),), manifest) + before = _snapshot(target) + with pytest.raises(ValueError, match="shard"): + target.import_canonical_state(exported) + _assert_snapshot_identity(target, before) + + +def test_muon_replicated_supports_canonical_state_but_whole_owner_does_not(): + parameter = torch.nn.Parameter(torch.ones(2, 2)) + optimizer = GefenMuon([("matrix", parameter)], fused=False) + identity = ParameterIdentity("Model.Matrix", (2, 2)) + shard = _replicated(identity) + _finalize(optimizer, ((parameter, shard),), ShardingManifest((shard,))) + assert optimizer.optimizer_contract().capabilities.canonical_state_io + + group = ProcessGroupIdentity("owner", ("rank:0",)) + owner_shard = ShardIdentity( + identity, + ParameterLayout.WHOLE_PARAMETER_OWNER, + LogicalSlice.full(identity), + process_group=group, + local_member="rank:0", + owner="rank:0", + placements=( + ShardPlacement( + "dp", + PlacementKind.WHOLE_PARAMETER_OWNER, + 0, + 1, + ), + ), + ) + owner_parameter = torch.nn.Parameter(torch.ones(2, 2)) + owner = GefenMuon([("matrix", owner_parameter)], fused=False) + _finalize(owner, ((owner_parameter, owner_shard),), ShardingManifest((owner_shard,))) + assert not owner.optimizer_contract().capabilities.canonical_state_io + with pytest.raises(RuntimeError, match="supported finalized"): + owner.export_canonical_state() + + +def test_initialized_normuon_requires_its_authoritative_pair_and_continues_exactly(): + source_parameter = torch.nn.Parameter( + torch.tensor([[1.0, -2.0], [3.0, -4.0]]) + ) + source = GefenMuon( + [("matrix", source_parameter)], + fused=False, + normuon=True, + ) + identity = ParameterIdentity("Model.Matrix", (2, 2)) + shard = _replicated(identity) + manifest = ShardingManifest((shard,)) + _finalize(source, ((source_parameter, shard),), manifest) + source._resolve_automatic_period = lambda *args: 4 + source_parameter.grad = torch.tensor([[0.5, -1.0], [1.5, -2.0]]) + source.step() + exported = source.export_canonical_state() + assert {"normuon_v", "normuon_step"}.issubset( + exported["parameters"]["Model.Matrix"]["state"] + ) + + target_parameter = torch.nn.Parameter(source_parameter.detach().clone()) + target = GefenMuon( + [("target", target_parameter)], + fused=False, + normuon=True, + ) + _finalize(target, ((target_parameter, shard),), manifest) + damaged = copy.deepcopy(exported) + damaged_state = damaged["parameters"]["Model.Matrix"]["state"] + damaged_state.pop("normuon_v") + damaged_state.pop("normuon_step") + before = _snapshot(target) + with pytest.raises(ValueError, match="NorMuon state is incomplete"): + target.import_canonical_state(damaged) + _assert_snapshot_identity(target, before) + + target.import_canonical_state(exported) + next_grad = torch.tensor([[-2.0, 0.25], [0.75, -1.25]]) + source_parameter.grad = next_grad.clone() + target_parameter.grad = next_grad.clone() + source.step() + target.step() + assert torch.equal(target_parameter, source_parameter) + assert torch.equal( + target.state[target_parameter]["normuon_v"], + source.state[source_parameter]["normuon_v"], + ) + + +@pytest.mark.parametrize("foreign_state", ["block", "factored"]) +def test_muon_rejects_plain_gefen_authoritative_state_variants(foreign_state): + source_parameter = torch.nn.Parameter(torch.ones(2, 2)) + source = GefenMuon([("matrix", source_parameter)], fused=False) + identity = ParameterIdentity("Model.Matrix", (2, 2)) + shard = _replicated(identity) + manifest = ShardingManifest((shard,)) + _finalize(source, ((source_parameter, shard),), manifest) + source._resolve_automatic_period = lambda *args: 4 + source_parameter.grad = torch.tensor([[1.0, -2.0], [3.0, -4.0]]) + source.step() + damaged = source.export_canonical_state() + parameter_state = damaged["parameters"]["Model.Matrix"]["state"] + if foreign_state == "block": + parameter_state["vmean"] = torch.zeros_like( + parameter_state["m_magnitude"] + ) + parameter_state["vmean_step"] = 1 + else: + parameter_state["v_row"] = torch.zeros(2, dtype=torch.float32) + parameter_state["v_col"] = torch.zeros(2, dtype=torch.float32) + parameter_state["factored_step"] = 1 + + target_parameter = torch.nn.Parameter(source_parameter.detach().clone()) + target = GefenMuon([("target", target_parameter)], fused=False) + _finalize(target, ((target_parameter, shard),), manifest) + before = _snapshot(target) + with pytest.raises(ValueError, match="declared state variant"): + target.import_canonical_state(damaged) + _assert_snapshot_identity(target, before) + + +def test_unsupported_custom_state_or_callable_policy_disables_canonical_claim(): + parameter = torch.nn.Parameter(torch.ones(2, 2)) + optimizer = GefenMuon( + [("matrix", parameter)], + fused=False, + ) + identity = ParameterIdentity("Model.Matrix", (2, 2)) + shard = _replicated(identity) + _finalize(optimizer, ((parameter, shard),), ShardingManifest((shard,))) + optimizer.param_groups[0]["custom_callable"] = lambda: None + assert not optimizer.optimizer_contract().capabilities.canonical_state_io + + plain_parameter = torch.nn.Parameter(torch.ones(4)) + plain = Gefen([("p", plain_parameter)], fused=False) + plain_identity = ParameterIdentity("Model.P", (4,)) + plain_shard = _replicated(plain_identity) + _finalize( + plain, + ((plain_parameter, plain_shard),), + ShardingManifest((plain_shard,)), + ) + plain.state[plain_parameter]["custom_tensor"] = torch.ones(1) + assert not plain.optimizer_contract().capabilities.canonical_state_io + with pytest.raises(RuntimeError, match="primitive algorithm policy"): + plain.export_canonical_state() + + class TensorSubclass(torch.Tensor): + pass + + subclass_parameter = torch.nn.Parameter(torch.ones(4)) + subclassed = Gefen([("p", subclass_parameter)], fused=False) + subclass_identity = ParameterIdentity("Model.Subclass", (4,)) + subclass_shard = _replicated(subclass_identity) + _finalize( + subclassed, + ((subclass_parameter, subclass_shard),), + ShardingManifest((subclass_shard,)), + ) + subclassed.param_groups[0]["custom_tensor"] = torch.ones(1).as_subclass( + TensorSubclass + ) + assert not subclassed.optimizer_contract().capabilities.canonical_state_io + with pytest.raises(RuntimeError, match="primitive algorithm policy"): + subclassed.export_canonical_state() + + subclassed.param_groups[0].pop("custom_tensor") + with torch.no_grad(): + nested = torch.nested.nested_tensor( + [torch.ones(2), torch.ones(3)] + ) + subclassed.param_groups[0]["custom_nested"] = nested + assert not subclassed.optimizer_contract().capabilities.canonical_state_io + with pytest.raises(RuntimeError, match="primitive algorithm policy"): + subclassed.export_canonical_state() + + subclassed.param_groups[0].pop("custom_nested") + subclassed.param_groups[0]["custom_nan"] = torch.tensor([float("nan")]) + assert not subclassed.optimizer_contract().capabilities.canonical_state_io + with pytest.raises(RuntimeError, match="primitive algorithm policy"): + subclassed.export_canonical_state() + + subclassed.param_groups[0].pop("custom_nan") + if hasattr(torch, "float8_e4m3fn"): + subclassed.param_groups[0]["custom_float8"] = torch.ones( + 2, dtype=torch.float8_e4m3fn + ) + assert not subclassed.optimizer_contract().capabilities.canonical_state_io + with pytest.raises(RuntimeError, match="primitive algorithm policy"): + subclassed.export_canonical_state() + subclassed.param_groups[0].pop("custom_float8") + + subclassed.param_groups[0]["custom_conjugate"] = torch.tensor( + [1.0 + 2.0j] + ).conj() + assert subclassed.optimizer_contract().capabilities.canonical_state_io + prepared = subclassed.prepare_canonical_state_import( + subclassed.export_canonical_state() + ) + subclassed.commit_canonical_state_import(prepared) + + +def test_inference_mode_tensor_lr_has_a_stable_canonical_freshness_token(): + with torch.inference_mode(): + learning_rate = torch.tensor(1e-3) + parameter = torch.nn.Parameter(torch.ones(4)) + optimizer = Gefen( + [("p", parameter)], + lr=learning_rate, + fused=False, + factored_v_2d=False, + ) + identity = ParameterIdentity("Model.P", (4,)) + shard = _replicated(identity) + _finalize(optimizer, ((parameter, shard),), ShardingManifest((shard,))) + exported = optimizer.export_canonical_state() + groups = optimizer.param_groups + group = optimizer.param_groups[0] + + prepared = optimizer.prepare_canonical_state_import(exported) + optimizer.commit_canonical_state_import(prepared) + assert optimizer.param_groups is groups + assert optimizer.param_groups[0] is group + assert optimizer.param_groups[0]["lr"] is learning_rate + assert optimizer.defaults["lr"] is learning_rate + + +def test_prepared_import_detects_storage_level_live_state_mutation(): + optimizer, _, _, _, _, _ = _two_parameter_source() + prepared = optimizer.prepare_canonical_state_import( + optimizer.export_canonical_state() + ) + magnitude = next( + state["m_magnitude"] + for state in optimizer.state.values() + if "m_magnitude" in state + ) + version = magnitude._version + magnitude.numpy()[0, 0] += 7.0 + assert magnitude._version == version + before = _snapshot(optimizer) + + with pytest.raises(RuntimeError, match="changed after"): + optimizer.commit_canonical_state_import(prepared) + _assert_snapshot_identity(optimizer, before) + + +def test_nonfinite_authoritative_state_disables_the_export_capability(): + optimizer, first, _, _, _, _ = _two_parameter_source() + optimizer.state[first]["vmean"].fill_(float("inf")) + + assert not optimizer.optimizer_contract().capabilities.canonical_state_io + with pytest.raises(RuntimeError, match="primitive algorithm policy"): + optimizer.export_canonical_state() + + +def test_invalid_live_counter_relationship_disables_the_export_capability(): + optimizer, first, _, _, _, _ = _two_parameter_source() + optimizer.state[first]["step"] = optimizer._gefen_global_step + 1 + + assert not optimizer.optimizer_contract().capabilities.canonical_state_io + with pytest.raises(RuntimeError, match="primitive algorithm policy"): + optimizer.export_canonical_state() + + +def test_canonical_policy_is_constructor_validated_and_capability_checked(): + parameter = torch.nn.Parameter(torch.ones(4)) + with pytest.raises(TypeError, match="codebook_refresh_every must be an integer"): + Gefen( + [("p", parameter)], + fused=False, + codebook_refresh_every=float("nan"), + ) + + optimizer = Gefen([("p", parameter)], fused=False) + identity = ParameterIdentity("Model.P", (4,)) + shard = _replicated(identity) + _finalize(optimizer, ((parameter, shard),), ShardingManifest((shard,))) + optimizer._codebook_refresh_every = float("nan") + assert not optimizer.optimizer_contract().capabilities.canonical_state_io + with pytest.raises(RuntimeError, match="primitive algorithm policy"): + optimizer.export_canonical_state() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +def test_capturable_import_requires_a_fresh_target_before_graph_state_exists(): + parameter = torch.nn.Parameter(torch.arange(8, dtype=torch.float32, device="cuda")) + optimizer = Gefen( + [("p", parameter)], + fused=False, + factored_v_2d=False, + capturable=True, + ) + identity = ParameterIdentity("Model.P", (8,)) + shard = _replicated(identity) + _finalize(optimizer, ((parameter, shard),), ShardingManifest((shard,))) + optimizer._resolve_automatic_period = lambda *args: 4 + parameter.grad = torch.arange(1, 9, dtype=torch.float32, device="cuda") + optimizer.step() + + graph = torch.cuda.CUDAGraph() + torch.cuda.synchronize() + with torch.cuda.graph(graph): + optimizer.step() + exported = optimizer.export_canonical_state() + before = _snapshot(optimizer) + with pytest.raises(RuntimeError, match="requires a fresh target"): + optimizer.prepare_canonical_state_import(exported) + _assert_snapshot_identity(optimizer, before) + + fresh_parameter = torch.nn.Parameter(parameter.detach().clone()) + fresh = Gefen( + [("p", fresh_parameter)], + fused=False, + factored_v_2d=False, + capturable=True, + ) + _finalize(fresh, ((fresh_parameter, shard),), ShardingManifest((shard,))) + fresh.import_canonical_state(exported) + assert fresh._device_gefen_global_step() == exported["common"][ + "gefen_global_step" + ] + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +def test_canonical_cpu_fragment_imports_state_to_cuda_parameter_devices(): + source_parameter = torch.nn.Parameter(torch.arange(8, dtype=torch.float32)) + source = Gefen([("p", source_parameter)], fused=False, factored_v_2d=False) + identity = ParameterIdentity("Model.P", (8,)) + shard = _replicated(identity) + manifest = ShardingManifest((shard,)) + _finalize(source, ((source_parameter, shard),), manifest) + source._resolve_automatic_period = lambda *args: 4 + source_parameter.grad = torch.arange(1, 9, dtype=torch.float32) + source.step() + exported = source.export_canonical_state() + assert all( + not torch.is_tensor(value) or value.device.type == "cpu" + for value in exported["parameters"]["Model.P"]["state"].values() + ) + + target_parameter = torch.nn.Parameter(source_parameter.detach().cuda()) + target = Gefen([("p", target_parameter)], fused=False, factored_v_2d=False) + _finalize(target, ((target_parameter, shard),), manifest) + target.import_canonical_state(exported) + assert all( + not torch.is_tensor(value) or value.device.type == "cuda" for value in target.state[target_parameter].values() + ) diff --git a/tests/test_codebook_scope_distributed.py b/tests/test_codebook_scope_distributed.py index cb0c048..5f175b6 100644 --- a/tests/test_codebook_scope_distributed.py +++ b/tests/test_codebook_scope_distributed.py @@ -307,6 +307,7 @@ def _distributed_worker(rank, world, init_file, queue): flat_agreement = all(torch.equal(item, gathered_codebooks[0]) for item in gathered_codebooks[1:]) flat_optimizer.step() flat_checkpoint = copy.deepcopy(flat_optimizer.state_dict()) + flat_canonical = flat_optimizer.export_canonical_state() checkpoint_scopes = [None] * world dist.all_gather_object( checkpoint_scopes, @@ -337,6 +338,33 @@ def _distributed_worker(rank, world, init_file, queue): cross_member_guard = rank == 0 except ValueError as exc: cross_member_guard = rank == 1 and "local-shard identity" in str(exc) + canonical_shards = [None] * world + dist.all_gather_object( + canonical_shards, + flat_canonical["parameters"]["Flat"]["shard"], + group=runtime_group, + ) + canonical_rank_local_identity = canonical_shards[0] != canonical_shards[1] + rank_zero_canonical = [flat_canonical if rank == 0 else None] + dist.broadcast_object_list(rank_zero_canonical, src=0, group=runtime_group) + cross_canonical_param = torch.nn.Parameter(flat_param.detach().clone()) + cross_canonical = Gefen( + [("flat", cross_canonical_param)], + fused=False, + factored_v_2d=False, + ) + _finalize( + cross_canonical, + cross_canonical_param, + flat_records[rank], + ShardingManifest(flat_records), + _binding(group, rank, runtime_group), + ) + try: + cross_canonical.import_canonical_state(rank_zero_canonical[0]) + canonical_cross_member_guard = rank == 0 + except ValueError as exc: + canonical_cross_member_guard = rank == 1 and "shard" in str(exc) resumed_param = torch.nn.Parameter(flat_param.detach().clone()) resumed = Gefen([("flat", resumed_param)], fused=False, factored_v_2d=False) resumed_binding = _binding(group, rank, runtime_group) @@ -348,9 +376,26 @@ def _distributed_worker(rank, world, init_file, queue): resumed_binding, ) resumed.load_state_dict(flat_checkpoint) + canonical_param = torch.nn.Parameter(flat_param.detach().clone()) + canonical_resumed = Gefen( + [("flat", canonical_param)], + fused=False, + factored_v_2d=False, + ) + canonical_binding = _binding(group, rank, runtime_group) + _finalize( + canonical_resumed, + canonical_param, + flat_records[rank], + ShardingManifest(flat_records), + canonical_binding, + ) + canonical_resumed.import_canonical_state(flat_canonical) continuation_grad = flat_grads[rank].flip(0).clone() resumed_param.grad = continuation_grad.clone() + canonical_param.grad = continuation_grad.clone() flat_param.grad = continuation_grad.clone() + canonical_resumed.step() resumed.step() flat_optimizer.step() flat_checkpoint_continuation = ( @@ -358,6 +403,14 @@ def _distributed_worker(rank, world, init_file, queue): and resumed.codebook_process_group_binding() is resumed_binding and torch.equal(resumed._gefen_codebook, flat_optimizer._gefen_codebook) ) + flat_canonical_continuation = ( + torch.equal(canonical_param, flat_param) + and canonical_resumed.codebook_process_group_binding() is canonical_binding + and torch.equal( + canonical_resumed._gefen_codebook, + flat_optimizer._gefen_codebook, + ) + ) refresh_succeeded = flat_optimizer.refresh_codebook() continuation_grads = tuple(gradient.flip(0) for gradient in flat_grads) refresh_oracle = learn_gefen_exact_codebook_from_grad_periods( @@ -606,6 +659,9 @@ def fail_exact_dp(*args, **kwargs): "rank_neutral_checkpoint_scope": rank_neutral_checkpoint_scope, "rank_local_checkpoint_identity": rank_local_checkpoint_identity, "cross_member_guard": cross_member_guard, + "canonical_rank_local_identity": canonical_rank_local_identity, + "canonical_cross_member_guard": canonical_cross_member_guard, + "flat_canonical_continuation": flat_canonical_continuation, "refresh_succeeded": refresh_succeeded, "refresh_matches_oracle": refresh_matches_oracle, "refresh_agreement": refresh_agreement, @@ -692,6 +748,9 @@ def test_explicit_gloo_scope_aggregates_logical_state_and_fails_atomically(): assert item["rank_neutral_checkpoint_scope"], item assert item["rank_local_checkpoint_identity"], item assert item["cross_member_guard"], item + assert item["canonical_rank_local_identity"], item + assert item["canonical_cross_member_guard"], item + assert item["flat_canonical_continuation"], item assert item["refresh_succeeded"], item assert item["refresh_matches_oracle"], item assert item["refresh_agreement"], item diff --git a/tests/test_rebinding_cpu.py b/tests/test_rebinding_cpu.py index 2322f77..76d36f3 100644 --- a/tests/test_rebinding_cpu.py +++ b/tests/test_rebinding_cpu.py @@ -187,7 +187,7 @@ def test_replicated_rebind_preserves_legacy_name_and_enables_identity_contract() assert contract.capabilities.stable_shard_identity assert contract.capabilities.shard_rebinding assert contract.capabilities.post_sharding - assert not contract.capabilities.canonical_state_io + assert contract.capabilities.canonical_state_io assert contract.capabilities.explicit_process_group_codebook_scope with pytest.raises(RuntimeError, match="already finalized"): optimizer.rebind_parameter(new, new, identity=identity) From 08af6ae19ef8856b257a85569c8b2b37465feefd Mon Sep 17 00:00:00 2001 From: thad0ctor Date: Sun, 12 Jul 2026 22:43:48 -0700 Subject: [PATCH 07/52] Add atomic optimizer state movement --- docs/optimizer_contracts.md | 10 +- src/gefen/__init__.py | 2 + src/gefen/contracts.py | 16 +- src/gefen/gefen.py | 381 +++++++++- src/gefen/gefen_muon.py | 16 +- tests/test_optimizer_contracts.py | 94 ++- tests/test_state_movement.py | 889 +++++++++++++++++++++++ tests/test_state_movement_distributed.py | 579 +++++++++++++++ 8 files changed, 1971 insertions(+), 16 deletions(-) create mode 100644 tests/test_state_movement.py create mode 100644 tests/test_state_movement_distributed.py diff --git a/docs/optimizer_contracts.md b/docs/optimizer_contracts.md index c5bca6b..efd204a 100644 --- a/docs/optimizer_contracts.md +++ b/docs/optimizer_contracts.md @@ -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. DTensor stable identity, Hybrid composite rebinding, topology-changing canonical checkpoint I/O, state movement, and offload remain unclaimed. +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. DTensor stable identity, Hybrid composite rebinding, topology-changing canonical checkpoint I/O, and offload remain unclaimed. ## Explicit learned-codebook process groups @@ -61,6 +61,14 @@ The v1 document contains only primitive containers and detached tight finite CPU This is an exact-binding transport-neutral local fragment, not the dense global logical state planned for portable DCP v3. Its dynamic `CANONICAL_LOCAL` checkpoint entry covers finalized plain-Gefen replicated and flattened local shards and finalized replicated GefenMuon, performs no collectives, reports atomic local import, and has an empty topology-changing set. A different member, slice, manifest, algorithm policy, group option, or declared state variant rejects. Export, preparation, and commit are quiescent checkpoint-boundary operations; prepared imports use content-bearing freshness tokens, including device counters, to reject intervening mutation. Export remains available after capturable warmup, but a capturable import target must still be fresh, before authoritative device state or a CUDA graph exists; importing replaces state identities, so an already captured graph cannot safely remain attached. Configurations with `stochastic_round=True` do not claim canonical v1 because the decomposed path intentionally lacks the fused stochastic quantizer, so changing effective fused availability would change the algorithm. DTensor, whole-owner completeness, Hybrid composition, rank-fragment gathering, resharding, world-size change, dense momentum decoding, and target-topology recompression remain unclaimed. +## Quiescent optimizer-state movement + +`StateMovementProvider.move_state_(device=None)` performs blocking CPU/CUDA co-location movement for Gefen and GefenMuon. With `device=None`, each declared authoritative per-parameter tensor moves to that parameter's current local device, including declared state attached to wrapper-orphaned parameter keys, while the canonical learned codebook moves to the first live local parameter device in parameter-group order. An optimizer with no local parameter storage keeps common state on CPU. An explicit device is accepted only after every live local parameter already resides there; an unindexed CUDA target is resolved from the one co-located live parameter device. The adapter must therefore move the model parameters first and invoke `move_state_` at a quiescent boundary before the next optimizer step. + +The core validates the finalized binding and complete declared state representation, allocates detached tight copies of the codebook and every authoritative tensor, and waits for all participating CUDA devices before one local publication. A preparation, transfer, synchronization, or validation failure leaves the exact live optimizer state, caches, parameters, gradients, groups, tensor learning rates, names, bindings, and metadata unchanged. Successful movement replaces the public `optimizer.state` mapping, every reachable or orphan per-parameter state dictionary, the canonical codebook identity, and every moved tensor identity; preserved non-tensor values and rank-local carrier tensors retain their identities, so adapters must not retain the replaced containers. Successful publication preserves host counters and metadata plus rank-local checkpoint carriers, discards only the rebuildable `stepsize` and `_h_buf` buffers, and invalidates per-device codebook/LUT copies, codebook-scope validation, the compiled static-address signature, and the tensor-learning-rate scalar cache. Preserved extension metadata is limited to provably tensor-free trees made from `None`, exact `bool`, `int`, `float`, `complex`, `str`, `bytes`, `torch.device`, `torch.dtype`, `torch.layout`, or `torch.memory_format` leaves and exact `dict`, `list`, `tuple`, `set`, `frozenset`, `deque`, or `torch.Size` containers. Cyclic or multiply referenced container graphs are outside that tree form. Arbitrary opaque objects, `defaultdict` or `OrderedDict` extension values, custom container subclasses, and non-dictionary per-parameter state mappings are rejected even when they appear tensor-free; the optimizer-owned top-level `state` may use its normal exact `defaultdict(dict)` representation. Undeclared tensor-bearing state, meta/nested/subclassed state tensors, FakeTensor parameters, capturable state, active compilation, and CUDA graph capture are likewise rejected rather than partially moved. + +`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` remains false because Gefen does not support stepping while authoritative state is parked away from its parameter, asynchronous paging, or transparent CPU offload. + Plain Gefen declares replicated, flattened element-shard, and the narrow DTensor training layouts. Its PyTorch rank-local checkpoint transport is same-topology only. GefenMuon declares replicated and narrow DTensor training, with mode-specific state extents: `approx` state is local, `exact` state is logically global, and `distributed` momentum is held by the parameter owner while non-owners retain metadata only. Native Parallel-Muon checkpoints separately declare world-size owner redistribution, not placement-changing resharding. `GefenMuonHybrid` retains its nested child namespaces and does not flatten AdamW or Gefen child state into a fabricated common schema. Gefen and GefenMuon native loads, rank-local payload restoration through `load_state_dict`, and distributed-owner payload restoration prepare their complete core restore before changing local live optimizer state. Their transport entries report `atomic_load=True`. This is a per-optimizer-instance fail-before-mutation guarantee at the optimizer load boundary, not a coordinated all-rank commit or a guarantee over work an external checkpoint orchestrator performs before calling the optimizer. Load pre-hooks run before that boundary and load post-hooks run afterward, so arbitrary side effects in user hooks are also outside the guarantee. Hybrid composite loads do not yet provide the same guarantee and report `atomic_load=False`. diff --git a/src/gefen/__init__.py b/src/gefen/__init__.py index 20951da..0829915 100644 --- a/src/gefen/__init__.py +++ b/src/gefen/__init__.py @@ -40,6 +40,7 @@ "StateField", "StateGeometry", "StateKeyMatch", + "StateMovementProvider", "StateScope", "StateVariant", "TopologyChange", @@ -109,6 +110,7 @@ def __getattr__(name): "StateField", "StateGeometry", "StateKeyMatch", + "StateMovementProvider", "StateScope", "StateVariant", "TopologyChange", diff --git a/src/gefen/contracts.py b/src/gefen/contracts.py index 6f05b7a..f71caf4 100644 --- a/src/gefen/contracts.py +++ b/src/gefen/contracts.py @@ -782,6 +782,14 @@ def import_canonical_state(self, state) -> None: """Prepare and commit a canonical fragment atomically.""" +@runtime_checkable +class StateMovementProvider(Protocol): + """Structural protocol for quiescent atomic optimizer-state movement.""" + + def move_state_(self, device=None) -> None: + """Co-locate authoritative state with the optimizer's live parameters.""" + + _ALL_PRECISIONS = frozenset( {Precision.FLOAT32, Precision.BFLOAT16, Precision.FLOAT16, Precision.FLOAT64} ) @@ -941,6 +949,7 @@ def _negative_capabilities( shard_rebinding: bool = False, post_sharding: bool = False, canonical_state_io: bool = False, + atomic_state_movement: bool = False, ) -> OptimizerCapabilities: return OptimizerCapabilities( training=training, @@ -954,7 +963,7 @@ def _negative_capabilities( shard_rebinding=shard_rebinding, post_sharding=post_sharding, canonical_state_io=canonical_state_io, - atomic_state_movement=False, + atomic_state_movement=atomic_state_movement, state_offload=False, ) @@ -967,6 +976,7 @@ def _gefen_contract( explicit_process_group_codebook_scope: bool = False, native_flattened_checkpoint: bool = False, canonical_state_layouts: AbstractSet[ParameterLayout] = frozenset(), + atomic_state_movement: bool = False, ) -> OptimizerContract: canonical_state_layouts = _frozenset(canonical_state_layouts) block_fields = ( @@ -1166,6 +1176,7 @@ def _gefen_contract( shard_rebinding=True, post_sharding=True, canonical_state_io=bool(canonical_state_layouts), + atomic_state_movement=atomic_state_movement, ), ) @@ -1200,6 +1211,7 @@ def _gefen_muon_contract( explicit_process_group_codebook_scope: bool = False, whole_parameter_owner: bool = False, canonical_state_layouts: AbstractSet[ParameterLayout] = frozenset(), + atomic_state_movement: bool = False, ) -> OptimizerContract: canonical_state_layouts = _frozenset(canonical_state_layouts) sharded_modes = _frozenset(sharded_modes) @@ -1505,6 +1517,7 @@ def _gefen_muon_contract( shard_rebinding=True, post_sharding=True, canonical_state_io=bool(canonical_state_layouts), + atomic_state_movement=atomic_state_movement, ), ) @@ -1584,6 +1597,7 @@ def _hybrid_contract( "StateField", "StateGeometry", "StateKeyMatch", + "StateMovementProvider", "StateScope", "StateVariant", "TrainingSupport", diff --git a/src/gefen/gefen.py b/src/gefen/gefen.py index df81cf4..090a265 100644 --- a/src/gefen/gefen.py +++ b/src/gefen/gefen.py @@ -16,7 +16,7 @@ import math import os import warnings -from collections import defaultdict, OrderedDict +from collections import defaultdict, deque, OrderedDict from itertools import chain from typing import Iterable, Optional, Tuple, Union @@ -110,6 +110,55 @@ "m_codebook_shape", } ) +_STATE_MOVEMENT_TENSOR_KEYS = frozenset( + { + "step", + "m_codebook", + "m_magnitude", + "vmean", + "vmean_step", + "v_row", + "v_col", + "factored_step", + "normuon_v", + "normuon_step", + } +) +_STATE_MOVEMENT_COUNTER_KEYS = frozenset( + {"step", "vmean_step", "factored_step", "normuon_step"} +) +_STATE_MOVEMENT_SCRATCH_KEYS = frozenset({"stepsize", "_h_buf"}) +_STATE_MOVEMENT_CAPTURABLE_KEYS = frozenset( + { + "_capt_scalars", + "_capt_consts", + "_capt_consts_key", + "_capt_stack", + "_capt_row", + } +) +_STATE_MOVEMENT_DEVICE_TYPES = frozenset({"cpu", "cuda"}) +_STATE_MOVEMENT_METADATA_LEAF_TYPES = ( + bool, + int, + float, + complex, + str, + bytes, + torch.device, + torch.dtype, + torch.layout, + torch.memory_format, +) +_STATE_MOVEMENT_METADATA_MAPPING_TYPES = (dict,) +_STATE_MOVEMENT_METADATA_SEQUENCE_TYPES = ( + list, + tuple, + set, + frozenset, + deque, + torch.Size, +) def _rank_local_payload_key(global_rank: int) -> str: @@ -1189,13 +1238,17 @@ def optimizer_contract(self) -> OptimizerContract: """Return the immutable state-layout and integration capability contract.""" identity_ready = self._canonical_identity_ready() - canonical_state_layouts = self._canonical_state_layouts() + try: + canonical_state_layouts = self._canonical_state_layouts() + except Exception: + canonical_state_layouts = frozenset() return _gefen_contract( factored_v_2d=self._factored_v_2d, canonical_parameter_fqns=identity_ready, stable_shard_identity=identity_ready, explicit_process_group_codebook_scope=True, canonical_state_layouts=canonical_state_layouts, + atomic_state_movement=self._atomic_state_movement_supported(), native_flattened_checkpoint=( self._codebook_scope_ready() and any( @@ -2345,6 +2398,330 @@ def rebind_parameter( manifest=ShardingManifest((shard,)), ) + @staticmethod + def _state_value_is_movement_safe_metadata(value, seen=None) -> bool: + value_type = type(value) + if value is None or value_type in _STATE_MOVEMENT_METADATA_LEAF_TYPES: + return True + is_mapping = value_type in _STATE_MOVEMENT_METADATA_MAPPING_TYPES + if ( + not is_mapping + and value_type not in _STATE_MOVEMENT_METADATA_SEQUENCE_TYPES + ): + return False + if seen is None: + seen = set() + value_id = id(value) + if value_id in seen: + return False + seen.add(value_id) + items = value.items() if is_mapping else value + if is_mapping: + return all( + Gefen._state_value_is_movement_safe_metadata(key, seen) + and Gefen._state_value_is_movement_safe_metadata(item, seen) + for key, item in items + ) + return all( + Gefen._state_value_is_movement_safe_metadata(item, seen) + for item in items + ) + + @staticmethod + def _state_movement_tensor_supported(value) -> bool: + return ( + type(value) is torch.Tensor + and not ( + hasattr(value, "to_local") + and hasattr(value, "placements") + ) + and value.layout is torch.strided + and not value.is_nested + and not value.is_quantized + and getattr(value, "fake_mode", None) is None + and value.device.type in _STATE_MOVEMENT_DEVICE_TYPES + ) + + def _state_movement_rejection_reason(self): + if ( + self._gefen_post_sharding_finalized + and not self._finalized_binding_layout_matches() + ): + return "the finalized parameter binding no longer matches live groups" + if self.capturable: + return "capturable optimizers have device-authoritative replay state" + if self._capt_stacks is not None: + return "capturable state stacks are active" + if self._gefen_global_step_by_device or self._sr_seed_by_device: + return "capturable device counters or stochastic-rounding seeds are active" + + try: + parameters = [ + parameter + for group in self.param_groups + for parameter in group["params"] + ] + except (KeyError, TypeError): + return "parameter groups have an invalid structure" + for parameter in parameters: + if not torch.is_tensor(parameter): + return "parameter groups contain a non-tensor value" + if getattr(parameter, "fake_mode", None) is not None: + return "FakeTensor parameters do not have movable storage" + try: + self._state_move_parameter_device(parameter) + except Exception: + return "parameters must use CPU or CUDA local storage" + + state_type = type(self.state) + if not ( + state_type is dict + or ( + state_type is defaultdict + and self.state.default_factory is dict + ) + ): + return "optimizer state must use a supported standard mapping" + for parameter, parameter_state in self.state.items(): + if not torch.is_tensor(parameter): + return "optimizer state is keyed by a non-tensor value" + if getattr(parameter, "fake_mode", None) is not None: + return "FakeTensor state keys do not have movable storage" + try: + self._state_move_parameter_device(parameter) + except Exception: + return "optimizer state is keyed by a parameter without CPU or CUDA local storage" + if type(parameter_state) is not dict: + return "per-parameter optimizer state must use a plain dictionary" + for key, value in parameter_state.items(): + if not self._state_value_is_movement_safe_metadata(key): + return "optimizer state contains an unsupported key" + if key in _STATE_MOVEMENT_SCRATCH_KEYS: + continue + if key in _STATE_MOVEMENT_CAPTURABLE_KEYS: + return "capturable per-parameter state is active" + if key in _STATE_MOVEMENT_TENSOR_KEYS: + if torch.is_tensor(value): + if not self._state_movement_tensor_supported(value): + return "authoritative tensor state has an unsupported representation" + elif key not in _STATE_MOVEMENT_COUNTER_KEYS or type(value) is not int: + return "authoritative tensor state has an invalid value" + continue + if isinstance(key, str) and key.startswith( + _RANK_LOCAL_PAYLOAD_KEY_PREFIX + ): + if torch.is_tensor(value): + if not self._state_movement_tensor_supported(value): + return "rank-local transport state has an unsupported representation" + elif not self._state_value_is_movement_safe_metadata(value): + return "rank-local transport state contains unsupported metadata" + continue + if not self._state_value_is_movement_safe_metadata(value): + return "undeclared optimizer state is not provably tensor-free metadata" + + codebook = self._gefen_codebook + if codebook is not None and not self._state_movement_tensor_supported(codebook): + return "the canonical codebook has an unsupported tensor representation" + return None + + def _atomic_state_movement_supported(self) -> bool: + try: + reason = self._state_movement_rejection_reason() + except Exception: + return False + if reason is not None: + return False + if torch.compiler.is_compiling(): + return False + if torch.cuda.is_available(): + try: + if torch.cuda.is_current_stream_capturing(): + return False + except RuntimeError: + return False + return True + + @staticmethod + def _state_tensor_device(parameter: torch.Tensor) -> torch.device: + local = parameter.to_local() if hasattr(parameter, "to_local") else parameter + if hasattr(local, "wait"): + local = local.wait() + return local.device + + @classmethod + def _state_move_parameter_device(cls, parameter: torch.Tensor) -> torch.device: + local = parameter.to_local() if hasattr(parameter, "to_local") else parameter + if hasattr(local, "wait"): + local = local.wait() + if not torch.is_tensor(local): + raise RuntimeError("parameter local storage must be a tensor") + device = local.device + if device.type not in _STATE_MOVEMENT_DEVICE_TYPES: + raise RuntimeError("state movement supports only CPU and CUDA parameters") + return torch.device("cpu") if device.type == "cpu" else device + + @staticmethod + def _normalize_state_move_target(device, live_devices) -> torch.device: + try: + target = torch.device(device) + except (TypeError, RuntimeError) as exc: + raise TypeError("device must identify a CPU or CUDA device") from exc + if target.type not in _STATE_MOVEMENT_DEVICE_TYPES: + raise ValueError("Gefen state movement supports only CPU and CUDA devices") + if target.type == "cpu": + target = torch.device("cpu") + else: + if not torch.cuda.is_available(): + raise ValueError("CUDA state movement requires an available CUDA device") + if target.index is None: + unique_devices = set(live_devices) + if len(unique_devices) != 1: + raise ValueError( + "an unindexed CUDA target requires one co-located live parameter device" + ) + candidate = next(iter(unique_devices)) + if candidate.type != "cuda": + raise ValueError("CUDA state movement requires CUDA-resident parameters") + target = candidate + if target.index < 0 or target.index >= torch.cuda.device_count(): + raise ValueError("CUDA state movement target is not an available device") + + if not live_devices: + if target.type != "cpu": + raise ValueError( + "an optimizer without local parameters keeps common state on CPU" + ) + elif any(parameter_device != target for parameter_device in live_devices): + raise ValueError( + "explicit state movement requires every live parameter to already be " + "co-located on {}".format(target) + ) + return target + + @staticmethod + def _copy_state_tensor_for_move( + tensor: torch.Tensor, device: torch.device + ) -> torch.Tensor: + return tensor.to( + device=device, + dtype=tensor.dtype, + non_blocking=False, + copy=True, + memory_format=torch.contiguous_format, + ).detach() + + @classmethod + def _validate_staged_state_tensor( + cls, source: torch.Tensor, staged, device: torch.device + ) -> None: + if ( + not cls._state_movement_tensor_supported(staged) + or staged is source + or staged.device != device + or staged.dtype != source.dtype + or tuple(staged.shape) != tuple(source.shape) + or staged.requires_grad + or not staged.is_contiguous() + or staged.storage_offset() != 0 + or staged.untyped_storage().nbytes() + != staged.numel() * staged.element_size() + ): + raise RuntimeError("state movement produced an invalid staged tensor") + + def _stage_state_move(self, device): + live_parameters = [ + parameter + for group in self.param_groups + for parameter in group["params"] + ] + live_devices = [ + self._state_move_parameter_device(parameter) for parameter in live_parameters + ] + live_parameter_ids = {id(parameter) for parameter in live_parameters} + explicit_target = ( + None + if device is None + else self._normalize_state_move_target(device, live_devices) + ) + codebook_target = ( + explicit_target + if explicit_target is not None + else (live_devices[0] if live_devices else torch.device("cpu")) + ) + staged_state = defaultdict(dict) + cuda_devices = set() + + def stage_tensor(value, target): + staged = self._copy_state_tensor_for_move(value, target) + self._validate_staged_state_tensor(value, staged, target) + if value.device.type == "cuda": + cuda_devices.add(value.device) + if target.type == "cuda": + cuda_devices.add(target) + return staged + + staged_codebook = ( + None + if self._gefen_codebook is None + else stage_tensor(self._gefen_codebook, codebook_target) + ) + for parameter, parameter_state in self.state.items(): + staged_parameter_state = {} + parameter_target = ( + explicit_target if id(parameter) in live_parameter_ids else None + ) + for key, value in parameter_state.items(): + if key in _STATE_MOVEMENT_SCRATCH_KEYS: + continue + if key in _STATE_MOVEMENT_CAPTURABLE_KEYS: + raise RuntimeError( + "Gefen state movement cannot migrate capturable scratch state" + ) + if key in _STATE_MOVEMENT_TENSOR_KEYS and torch.is_tensor(value): + if parameter_target is None: + parameter_target = self._state_move_parameter_device(parameter) + staged_parameter_state[key] = stage_tensor(value, parameter_target) + else: + staged_parameter_state[key] = value + staged_state[parameter] = staged_parameter_state + + for cuda_device in sorted( + cuda_devices, key=lambda item: -1 if item.index is None else item.index + ): + torch.cuda.synchronize(cuda_device) + return staged_state, staged_codebook + + @torch.no_grad() + def move_state_(self, device=None) -> None: + """Atomically co-locate authoritative optimizer state with live parameters.""" + + self._assert_finalized_binding_layout() + try: + reason = self._state_movement_rejection_reason() + except Exception as exc: + raise RuntimeError( + "Gefen atomic state movement could not inspect live optimizer state" + ) from exc + if reason is not None: + raise RuntimeError("Gefen atomic state movement is unavailable: {}".format(reason)) + if torch.compiler.is_compiling(): + raise RuntimeError("Gefen state movement cannot run during torch.compile") + if torch.cuda.is_available() and torch.cuda.is_current_stream_capturing(): + raise RuntimeError("Gefen state movement cannot run during CUDA graph capture") + + staged_state, staged_codebook = self._stage_state_move(device) + self.__dict__.update( + { + "state": staged_state, + "_gefen_codebook": staged_codebook, + "_gefen_codebook_by_device": {}, + "_gefen_codebook_lut_by_device": {}, + "_gefen_codebook_scope_validated": False, + "_static_mark_sig": None, + "_lr_scalar_cache": None, + } + ) + @staticmethod def _normalize_param_groups(params): if isinstance(params, torch.Tensor): diff --git a/src/gefen/gefen_muon.py b/src/gefen/gefen_muon.py index 9f179af..3624f89 100644 --- a/src/gefen/gefen_muon.py +++ b/src/gefen/gefen_muon.py @@ -761,6 +761,10 @@ def optimizer_contract(self) -> OptimizerContract: for group in self.param_groups if not group.get("normuon", False) ) + try: + canonical_state_layouts = self._canonical_state_layouts() + except Exception: + canonical_state_layouts = frozenset() return _gefen_muon_contract( sharded_modes=sharded_modes, normuon_modes=normuon_modes, @@ -768,7 +772,8 @@ 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, - canonical_state_layouts=self._canonical_state_layouts(), + canonical_state_layouts=canonical_state_layouts, + atomic_state_movement=self._atomic_state_movement_supported(), whole_parameter_owner=( self._codebook_scope_ready() and any( @@ -1830,15 +1835,6 @@ def _step_automatic( update = self._compute_muon_update(group, param_name, p, grad, eff_numel) self._apply_muon_update(group, p, update, is_sharded, approx) - @staticmethod - def _state_tensor_device(p: torch.Tensor) -> torch.device: - if hasattr(p, "to_local"): - local = p.to_local() - if hasattr(local, "wait"): - local = local.wait() - return local.device - return p.device - def _distributed_state_items(self, state_dict): saved_ids = [] for saved_group in state_dict["param_groups"]: diff --git a/tests/test_optimizer_contracts.py b/tests/test_optimizer_contracts.py index 7dd2c31..9fd35ed 100644 --- a/tests/test_optimizer_contracts.py +++ b/tests/test_optimizer_contracts.py @@ -1,7 +1,8 @@ """CPU coverage for public optimizer capability and state-layout contracts.""" import copy -from dataclasses import FrozenInstanceError +from collections import defaultdict, OrderedDict +from dataclasses import FrozenInstanceError, dataclass import pytest import torch @@ -23,6 +24,7 @@ StateField, StateGeometry, StateKeyMatch, + StateMovementProvider, StateScope, StateVariant, TopologyChange, @@ -128,6 +130,7 @@ def test_plain_contract_matches_live_persistent_state(factored_v_2d): factored_v_2d=factored_v_2d, ) assert isinstance(optimizer, OptimizerContractProvider) + assert isinstance(optimizer, StateMovementProvider) contract = optimizer.optimizer_contract() assert contract.schema_version == CONTRACT_SCHEMA_VERSION @@ -166,7 +169,7 @@ def test_plain_contract_matches_live_persistent_state(factored_v_2d): 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 contract.capabilities.atomic_state_movement assert not contract.capabilities.state_offload assert Precision.FLOAT64 in contract.capabilities.precisions flattened = _training_support( @@ -348,11 +351,14 @@ def test_muon_contract_separates_mode_topology_and_state_extent( sharded_mode=sharded_mode, normuon=normuon, ) + assert isinstance(optimizer, StateMovementProvider) contract = optimizer.optimizer_contract() assert contract.implementation == "gefen.GefenMuon" assert contract.capabilities.supported_parameter_ranks == (2,) assert contract.capabilities.explicit_process_group_codebook_scope + assert contract.capabilities.atomic_state_movement + assert not contract.capabilities.state_offload native = next( item for item in contract.capabilities.checkpoints @@ -473,6 +479,7 @@ def test_hybrid_contract_preserves_child_namespaces(backup_optimizer): fused=False, backup_optimizer=backup_optimizer, ) + assert not isinstance(optimizer, StateMovementProvider) contract = optimizer.optimizer_contract() assert contract.implementation == "gefen.GefenMuonHybrid" @@ -485,6 +492,8 @@ def test_hybrid_contract_preserves_child_namespaces(backup_optimizer): assert contract.children[1].contract is None assert contract.state_layout.composite_namespaces == ("muon", "backup") assert not contract.capabilities.explicit_process_group_codebook_scope + assert not contract.capabilities.atomic_state_movement + assert not contract.capabilities.state_offload assert contract.children[0].contract.capabilities.explicit_process_group_codebook_scope if backup_optimizer == "gefen": assert contract.children[1].contract.capabilities.explicit_process_group_codebook_scope @@ -517,6 +526,86 @@ def test_muon_contract_keeps_mixed_normuon_variants_in_one_mode(): assert "quantized_normuon_replicated_exact" in variants +@pytest.mark.parametrize("implementation", ["gefen", "muon"]) +def test_capturable_contract_declines_atomic_state_movement(implementation): + shape = (4,) if implementation == "gefen" else (2, 2) + parameter = torch.nn.Parameter(torch.ones(shape)) + optimizer_type = Gefen if implementation == "gefen" else GefenMuon + optimizer = optimizer_type( + [("parameter", parameter)], + fused=False, + capturable=True, + ) + + capabilities = optimizer.optimizer_contract().capabilities + assert not capabilities.atomic_state_movement + assert not capabilities.state_offload + + +@pytest.mark.parametrize("implementation", ["gefen", "muon"]) +def test_undeclared_tensor_state_disables_atomic_state_movement(implementation): + shape = (4,) if implementation == "gefen" else (2, 2) + parameter = torch.nn.Parameter(torch.ones(shape)) + optimizer_type = Gefen if implementation == "gefen" else GefenMuon + optimizer = optimizer_type([("parameter", parameter)], fused=False) + assert optimizer.optimizer_contract().capabilities.atomic_state_movement + + optimizer.state[parameter]["undeclared_tensor"] = torch.ones(1) + capabilities = optimizer.optimizer_contract().capabilities + assert not capabilities.atomic_state_movement + assert not capabilities.state_offload + + +@pytest.mark.parametrize("implementation", ["gefen", "muon"]) +@pytest.mark.parametrize("contains_tensor", [False, True]) +def test_opaque_extension_state_disables_atomic_state_movement( + implementation, contains_tensor +): + @dataclass + class ExtensionState: + payload: object + + shape = (4,) if implementation == "gefen" else (2, 2) + parameter = torch.nn.Parameter(torch.ones(shape)) + optimizer_type = Gefen if implementation == "gefen" else GefenMuon + optimizer = optimizer_type([("parameter", parameter)], fused=False) + payload = torch.ones(1) if contains_tensor else "tensor-free" + optimizer.state[parameter]["extension"] = ExtensionState(payload) + + capabilities = optimizer.optimizer_contract().capabilities + assert not capabilities.atomic_state_movement + assert not capabilities.state_offload + + +@pytest.mark.parametrize("implementation", ["gefen", "muon"]) +@pytest.mark.parametrize("container_type", ["defaultdict", "ordered_dict"]) +def test_extension_mapping_with_hidden_tensor_disables_atomic_state_movement( + implementation, container_type +): + class TensorFactory: + def __init__(self, tensor): + self.tensor = tensor + + def __call__(self): + return self.tensor + + shape = (4,) if implementation == "gefen" else (2, 2) + parameter = torch.nn.Parameter(torch.ones(shape)) + optimizer_type = Gefen if implementation == "gefen" else GefenMuon + optimizer = optimizer_type([("parameter", parameter)], fused=False) + hidden_tensor = torch.ones(1) + if container_type == "defaultdict": + extension = defaultdict(TensorFactory(hidden_tensor)) + else: + extension = OrderedDict() + extension.hidden_tensor = hidden_tensor + optimizer.state[parameter]["extension"] = extension + + capabilities = optimizer.optimizer_contract().capabilities + assert not capabilities.atomic_state_movement + assert not capabilities.state_offload + + def test_muon_mixed_approx_distributed_checkpoint_is_same_topology_only(): first = torch.nn.Parameter(torch.ones(4, 4)) second = torch.nn.Parameter(torch.ones(4, 4)) @@ -657,3 +746,4 @@ def test_all_public_contract_exports_resolve(): from gefen import contracts assert all(getattr(gefen, name) is not None for name in contracts.__all__) + assert gefen.StateMovementProvider is StateMovementProvider diff --git a/tests/test_state_movement.py b/tests/test_state_movement.py new file mode 100644 index 0000000..a290acc --- /dev/null +++ b/tests/test_state_movement.py @@ -0,0 +1,889 @@ +"""Atomic optimizer-state movement coverage for Gefen and GefenMuon.""" + +import copy +from collections import defaultdict, OrderedDict +import warnings + +import pytest +import torch + +from gefen import Gefen, GefenMuon, ParameterIdentity + + +_KINDS = ("block", "factored", "muon", "normuon") +_MOVABLE_STATE_KEYS = frozenset( + { + "step", + "m_codebook", + "m_magnitude", + "vmean", + "vmean_step", + "v_row", + "v_col", + "factored_step", + "normuon_v", + "normuon_step", + } +) +_SCRATCH_KEYS = frozenset({"stepsize", "_h_buf"}) +_CACHE_ATTRS = ( + "_gefen_codebook_by_device", + "_gefen_codebook_lut_by_device", + "_sr_seed_by_device", + "_gefen_global_step_by_device", +) + + +def _build_initialized(kind): + shape = (8,) if kind == "block" else (2, 4) + values = torch.linspace(-0.4, 0.7, 8, dtype=torch.float32).reshape(shape) + parameter = torch.nn.Parameter(values.clone()) + tensor_lr = torch.tensor(2.0e-3, dtype=torch.float32) + group_metadata = {"labels": ["preserve", kind]} + group = { + "params": [("layer.weight", parameter)], + "lr": tensor_lr, + "movement_metadata": group_metadata, + } + if kind == "block": + optimizer = Gefen( + [group], + lr=tensor_lr, + fused=False, + factored_v_2d=False, + ) + elif kind == "factored": + optimizer = Gefen( + [group], + lr=tensor_lr, + fused=False, + factored_v_2d=True, + ) + elif kind in ("muon", "normuon"): + optimizer = GefenMuon( + [group], + lr=tensor_lr, + weight_decay=0.0, + fused=False, + ns_steps=1, + normuon=kind == "normuon", + ) + else: + raise AssertionError("unknown optimizer kind: {}".format(kind)) + + optimizer._resolve_automatic_period = lambda *args, **kwargs: 4 + optimizer._predict_period_from_grad_sq = lambda *args, **kwargs: 4 + parameter.grad = torch.linspace(-1.25, 0.75, 8).reshape_as(parameter) + optimizer.step() + return optimizer, parameter, tensor_lr, group_metadata + + +def _oversized_copy(tensor): + flat_size = tensor.numel() + backing = torch.empty( + flat_size + 11, + dtype=tensor.dtype, + device=tensor.device, + ) + result = backing.narrow(0, 5, flat_size).view(tensor.shape) + result.copy_(tensor) + assert result.untyped_storage().nbytes() > flat_size * tensor.element_size() + return result + + +def _make_persistent_state_oversized(optimizer, parameter): + state = optimizer.state[parameter] + for key, value in tuple(state.items()): + if key in _MOVABLE_STATE_KEYS and torch.is_tensor(value): + state[key] = _oversized_copy(value) + optimizer._gefen_codebook = _oversized_copy(optimizer._gefen_codebook) + + +def _persistent_tensor_snapshot(optimizer, parameter): + state = optimizer.state[parameter] + result = { + key: (value, value.detach().clone()) + for key, value in state.items() + if key in _MOVABLE_STATE_KEYS and torch.is_tensor(value) + } + result["_gefen_codebook"] = ( + optimizer._gefen_codebook, + optimizer._gefen_codebook.detach().clone(), + ) + return result + + +def _assert_fresh_tight_copy(actual, old, expected, device): + assert actual is not old + assert actual.device == device + assert actual.dtype == expected.dtype + assert actual.shape == expected.shape + assert actual.layout == torch.strided + assert actual.is_contiguous() + assert actual.storage_offset() == 0 + assert actual.untyped_storage().nbytes() == actual.numel() * actual.element_size() + torch.testing.assert_close(actual, expected, rtol=0, atol=0, equal_nan=True) + + +def _seed_discardable_runtime_state(optimizer, parameter, tensor_lr): + state = optimizer.state[parameter] + state["stepsize"] = _oversized_copy(torch.tensor([91.0])) + state["_h_buf"] = _oversized_copy(torch.tensor([92.0])) + carrier = _oversized_copy(torch.tensor([7, 11, 13], dtype=torch.int64)) + state["_gefen_rank_local_payload_0"] = carrier + state_metadata = { + "labels": ["persistent", "metadata"], + "device": torch.device("cpu"), + "dtype": torch.float32, + "layout": torch.strided, + "memory_format": torch.contiguous_format, + "shape": torch.Size((2, 3)), + "complex": 1.0 + 2.0j, + "bytes": b"metadata", + "optional": None, + } + state["movement_metadata"] = state_metadata + + cpu = torch.device("cpu") + optimizer._gefen_codebook_by_device[cpu] = optimizer._gefen_codebook + optimizer._gefen_codebook_lut_by_device[cpu] = torch.tensor([19.0]) + optimizer._gefen_codebook_scope_validated = True + optimizer._static_mark_sig = ("stale",) + optimizer._lr_scalar_cache = ( + tensor_lr, + tensor_lr._version, + float(tensor_lr.item()), + ) + return carrier, state_metadata + + +def _tensor_payload(tensor): + if tensor.device.type == "meta": + return None + if tensor.is_nested: + return tuple(item.detach().clone() for item in tensor.unbind()) + return tensor.detach().clone() + + +def _snapshot_tree(value): + if torch.is_tensor(value): + return ( + "tensor", + value, + value._version, + value.dtype, + value.device, + value.layout, + _tensor_payload(value), + ) + if isinstance(value, dict): + return ( + "dict", + value, + tuple((key, _snapshot_tree(item)) for key, item in value.items()), + ) + if isinstance(value, list): + return ("list", value, tuple(_snapshot_tree(item) for item in value)) + if isinstance(value, tuple): + return ("tuple", value, tuple(_snapshot_tree(item) for item in value)) + return ("leaf", value, copy.deepcopy(value)) + + +def _assert_tensor_payload(actual, payload): + if payload is None: + return + if actual.is_nested: + actual_items = actual.unbind() + assert len(actual_items) == len(payload) + for item, expected in zip(actual_items, payload): + torch.testing.assert_close(item, expected, rtol=0, atol=0, equal_nan=True) + return + torch.testing.assert_close(actual, payload, rtol=0, atol=0, equal_nan=True) + + +def _assert_tree_exact(actual, snapshot): + kind = snapshot[0] + if kind == "tensor": + _, reference, version, dtype, device, layout, payload = snapshot + assert actual is reference + assert actual._version == version + assert actual.dtype == dtype + assert actual.device == device + assert actual.layout == layout + _assert_tensor_payload(actual, payload) + return + if kind == "dict": + _, reference, entries = snapshot + assert actual is reference + assert tuple(actual) == tuple(key for key, _ in entries) + for key, child in entries: + _assert_tree_exact(actual[key], child) + return + if kind in ("list", "tuple"): + _, reference, entries = snapshot + assert actual is reference + assert len(actual) == len(entries) + for item, child in zip(actual, entries): + _assert_tree_exact(item, child) + return + _, reference, expected = snapshot + assert actual is reference + assert actual == expected + + +def _snapshot_exact_optimizer(optimizer): + parameters = tuple( + parameter + for group in optimizer.param_groups + for parameter in group["params"] + ) + return { + "top": optimizer.__dict__.copy(), + "state": _snapshot_tree(optimizer.state), + "groups": _snapshot_tree(optimizer.param_groups), + "defaults": _snapshot_tree(optimizer.defaults), + "param_names": _snapshot_tree(optimizer._param_names), + "codebook": _snapshot_tree(optimizer._gefen_codebook), + "caches": { + name: _snapshot_tree(getattr(optimizer, name)) for name in _CACHE_ATTRS + }, + "grads": tuple( + (parameter, _snapshot_tree(parameter.grad)) for parameter in parameters + ), + } + + +def _assert_exact_optimizer_snapshot(optimizer, snapshot): + assert optimizer.__dict__.keys() == snapshot["top"].keys() + for key, value in snapshot["top"].items(): + assert optimizer.__dict__[key] is value + _assert_tree_exact(optimizer.state, snapshot["state"]) + _assert_tree_exact(optimizer.param_groups, snapshot["groups"]) + _assert_tree_exact(optimizer.defaults, snapshot["defaults"]) + _assert_tree_exact(optimizer._param_names, snapshot["param_names"]) + _assert_tree_exact(optimizer._gefen_codebook, snapshot["codebook"]) + for name, expected in snapshot["caches"].items(): + _assert_tree_exact(getattr(optimizer, name), expected) + for parameter, expected in snapshot["grads"]: + _assert_tree_exact(parameter.grad, expected) + + +def _movement_candidates(optimizer, parameter): + candidates = {_tensor_storage_token(optimizer._gefen_codebook)} + candidates.update( + _tensor_storage_token(value) + for key, value in optimizer.state[parameter].items() + if key in _MOVABLE_STATE_KEYS and torch.is_tensor(value) + ) + return candidates + + +def _tensor_storage_token(tensor): + return ( + tensor.device, + tensor.untyped_storage().data_ptr(), + tensor.storage_offset(), + tensor.numel(), + ) + + +def _install_late_to_failure(monkeypatch, candidates, *, destination_type): + original_to = torch.Tensor.to + completed_copies = [] + + def flaky_to(tensor, *args, **kwargs): + result = original_to(tensor, *args, **kwargs) + if ( + _tensor_storage_token(tensor) in candidates + and result is not tensor + and result.device.type == destination_type + ): + completed_copies.append(result) + if len(completed_copies) == 3: + raise RuntimeError("injected late state-copy failure") + return result + + monkeypatch.setattr(torch.Tensor, "to", flaky_to) + return completed_copies + + +def _move_parameter_module(parameter, device): + module = torch.nn.Module() + module.register_parameter("weight", parameter) + module.to(device) + assert module.weight is parameter + return module + + +def _persistent_values(optimizer, parameter): + values = { + key: value.detach().cpu().clone() if torch.is_tensor(value) else copy.deepcopy(value) + for key, value in optimizer.state[parameter].items() + if key in _MOVABLE_STATE_KEYS + } + values["_gefen_codebook"] = optimizer._gefen_codebook.detach().cpu().clone() + return values + + +def _assert_persistent_values_equal(left, right): + assert set(left) == set(right) + for key in left: + if torch.is_tensor(left[key]): + torch.testing.assert_close(left[key], right[key], rtol=0, atol=0, equal_nan=True) + else: + assert left[key] == right[key] + + +@pytest.mark.parametrize("kind", _KINDS) +@pytest.mark.parametrize("explicit_destination", [False, True]) +def test_cpu_state_movement_is_fresh_tight_and_preserves_live_training_objects( + kind, explicit_destination +): + optimizer, parameter, tensor_lr, group_metadata = _build_initialized(kind) + assert optimizer.optimizer_contract().capabilities.atomic_state_movement + assert not optimizer.optimizer_contract().capabilities.state_offload + + _make_persistent_state_oversized(optimizer, parameter) + carrier, state_metadata = _seed_discardable_runtime_state( + optimizer, parameter, tensor_lr + ) + tensors_before = _persistent_tensor_snapshot(optimizer, parameter) + state_before = optimizer.state[parameter] + non_tensor_before = { + key: value + for key, value in state_before.items() + if not torch.is_tensor(value) and key not in _SCRATCH_KEYS + } + groups_before = optimizer.param_groups + group_before = optimizer.param_groups[0] + group_params_before = group_before["params"] + defaults_before = optimizer.defaults + param_names_before = optimizer._param_names + grad_before = parameter.grad + grad_value_before = grad_before.detach().clone() + parameter_value_before = parameter.detach().clone() + global_step_before = optimizer._gefen_global_step + + destination = torch.device("cpu") if explicit_destination else None + optimizer.move_state_(destination) + + assert optimizer.param_groups is groups_before + assert optimizer.param_groups[0] is group_before + assert optimizer.param_groups[0]["params"] is group_params_before + assert optimizer.param_groups[0]["params"][0] is parameter + assert optimizer.param_groups[0]["lr"] is tensor_lr + assert optimizer.param_groups[0]["movement_metadata"] is group_metadata + assert optimizer.defaults is defaults_before + assert optimizer.defaults["lr"] is tensor_lr + assert optimizer._param_names is param_names_before + assert parameter.grad is grad_before + torch.testing.assert_close(parameter.grad, grad_value_before, rtol=0, atol=0) + torch.testing.assert_close(parameter, parameter_value_before, rtol=0, atol=0) + assert optimizer._gefen_global_step == global_step_before + + state = optimizer.state[parameter] + for key, value in non_tensor_before.items(): + assert state[key] is value + assert state["movement_metadata"] is state_metadata + assert state["_gefen_rank_local_payload_0"] is carrier + assert "stepsize" not in state + assert "_h_buf" not in state + + for key, (old, expected) in tensors_before.items(): + actual = optimizer._gefen_codebook if key == "_gefen_codebook" else state[key] + _assert_fresh_tight_copy(actual, old, expected, torch.device("cpu")) + + assert optimizer._gefen_codebook_by_device == {} + assert optimizer._gefen_codebook_lut_by_device == {} + assert optimizer._sr_seed_by_device == {} + assert optimizer._gefen_global_step_by_device == {} + assert optimizer._gefen_codebook_scope_validated is False + assert optimizer._capt_stacks is None + assert optimizer._static_mark_sig is None + assert optimizer._lr_scalar_cache is None + assert optimizer.optimizer_contract().capabilities.atomic_state_movement + assert not optimizer.optimizer_contract().capabilities.state_offload + + +def test_pristine_state_movement_is_repeatable_and_returns_none(): + parameter = torch.nn.Parameter(torch.arange(8, dtype=torch.float32)) + optimizer = Gefen( + [("layer.weight", parameter)], + fused=False, + factored_v_2d=False, + ) + original_state = optimizer.state + original_parameter_state = optimizer.state[parameter] + parameter_value = parameter.detach().clone() + + assert optimizer.move_state_() is None + first_state = optimizer.state + first_parameter_state = optimizer.state[parameter] + assert first_state is not original_state + assert first_parameter_state is not original_parameter_state + assert first_parameter_state == {"name": "layer.weight"} + assert optimizer._gefen_codebook is None + + assert optimizer.move_state_(torch.device("cpu")) is None + assert optimizer.state is not first_state + assert optimizer.state[parameter] is not first_parameter_state + assert optimizer.state[parameter] == {"name": "layer.weight"} + torch.testing.assert_close(parameter, parameter_value, rtol=0, atol=0) + assert optimizer.optimizer_contract().capabilities.atomic_state_movement + + +def test_codebook_preinitialized_state_moves_without_initializing_parameter_tensors(): + parameter = torch.nn.Parameter(torch.arange(1, 9, dtype=torch.float32)) + tensor_lr = torch.tensor(3.0e-3) + optimizer = Gefen( + [("layer.weight", parameter)], + lr=tensor_lr, + fused=False, + factored_v_2d=False, + ) + optimizer._resolve_automatic_period = lambda *args, **kwargs: 4 + parameter.grad = torch.linspace(-2.0, 1.0, parameter.numel()) + optimizer._ensure_gefen_codebook(reuse_existing_periods=False) + assert set(optimizer.state[parameter]) == {"name", "automatic_period"} + optimizer._gefen_codebook = _oversized_copy(optimizer._gefen_codebook) + old_codebook = optimizer._gefen_codebook + expected_codebook = old_codebook.detach().clone() + grad_before = parameter.grad + state_metadata = {"preinitialized": True} + optimizer.state[parameter]["movement_metadata"] = state_metadata + + optimizer.move_state_() + + assert set(optimizer.state[parameter]) == { + "name", + "automatic_period", + "movement_metadata", + } + assert optimizer.state[parameter]["movement_metadata"] is state_metadata + assert parameter.grad is grad_before + assert optimizer._gefen_global_step == 0 + _assert_fresh_tight_copy( + optimizer._gefen_codebook, + old_codebook, + expected_codebook, + torch.device("cpu"), + ) + + +def test_noncontiguous_declared_state_is_normalized_to_a_tight_copy(): + optimizer, parameter, _, _ = _build_initialized("block") + source = torch.arange(12, dtype=torch.float32).reshape(3, 4).t() + assert not source.is_contiguous() + optimizer.state[parameter]["m_magnitude"] = source + expected = source.detach().clone() + + optimizer.move_state_() + + _assert_fresh_tight_copy( + optimizer.state[parameter]["m_magnitude"], + source, + expected, + torch.device("cpu"), + ) + + +def test_wrapper_orphan_state_is_preserved_and_co_located_with_its_key(): + optimizer, _, _, _ = _build_initialized("block") + orphan = torch.nn.Parameter(torch.arange(6, dtype=torch.float32)) + orphan_tensor = _oversized_copy(torch.linspace(1.0, 2.0, 3)) + orphan_carrier = _oversized_copy(torch.tensor([3, 5, 8], dtype=torch.uint8)) + orphan_metadata = {"owner": "wrapper"} + old_orphan_state = { + "name": "orphan.weight", + "automatic_period": 2, + "step": 7, + "m_magnitude": orphan_tensor, + "stepsize": torch.tensor([99.0]), + "_gefen_rank_local_payload_0": orphan_carrier, + "movement_metadata": orphan_metadata, + } + optimizer.state[orphan] = old_orphan_state + expected_tensor = orphan_tensor.detach().clone() + + optimizer.move_state_(torch.device("cpu")) + + assert orphan in optimizer.state + assert optimizer.state[orphan] is not old_orphan_state + moved = optimizer.state[orphan] + assert moved["name"] == "orphan.weight" + assert moved["automatic_period"] == 2 + assert moved["step"] == 7 + assert moved["movement_metadata"] is orphan_metadata + assert moved["_gefen_rank_local_payload_0"] is orphan_carrier + assert "stepsize" not in moved + _assert_fresh_tight_copy( + moved["m_magnitude"], + orphan_tensor, + expected_tensor, + torch.device("cpu"), + ) + + +@pytest.mark.parametrize("kind", _KINDS) +def test_native_checkpoint_continuation_remains_exact_after_movement(kind): + optimizer, parameter, _, _ = _build_initialized(kind) + resumed, resumed_parameter, _, _ = _build_initialized(kind) + + optimizer.move_state_() + resumed.load_state_dict(copy.deepcopy(optimizer.state_dict())) + _assert_persistent_values_equal( + _persistent_values(optimizer, parameter), + _persistent_values(resumed, resumed_parameter), + ) + + continuation_grad = torch.linspace(0.45, -0.85, parameter.numel()).reshape_as( + parameter + ) + parameter.grad = continuation_grad.clone() + resumed_parameter.grad = continuation_grad.clone() + optimizer.step() + resumed.step() + + torch.testing.assert_close(parameter, resumed_parameter, rtol=0, atol=0) + _assert_persistent_values_equal( + _persistent_values(optimizer, parameter), + _persistent_values(resumed, resumed_parameter), + ) + + +def test_movement_invalidates_prepared_canonical_import_but_keeps_io_available(): + def build_finalized(): + parameter = torch.nn.Parameter(torch.linspace(-0.5, 0.7, 8)) + optimizer = Gefen( + [("layer.weight", parameter)], + fused=False, + factored_v_2d=False, + ) + optimizer.rebind_parameter( + parameter, + parameter, + identity=ParameterIdentity("Layer.Weight", (8,)), + ) + optimizer._resolve_automatic_period = lambda *args, **kwargs: 4 + return optimizer, parameter + + source, source_parameter = build_finalized() + source_parameter.grad = torch.linspace(-1.0, 0.8, 8) + source.step() + source.move_state_() + exported = source.export_canonical_state() + + target, target_parameter = build_finalized() + prepared = target.prepare_canonical_state_import(exported) + target.move_state_() + with pytest.raises(RuntimeError, match="changed after canonical import preparation"): + target.commit_canonical_state_import(prepared) + + assert target.optimizer_contract().capabilities.canonical_state_io + assert target.optimizer_contract().capabilities.atomic_state_movement + target.import_canonical_state(exported) + _assert_persistent_values_equal( + _persistent_values(source, source_parameter), + _persistent_values(target, target_parameter), + ) + + +@pytest.mark.parametrize("kind", _KINDS) +def test_late_cpu_copy_failure_is_exactly_atomic(kind, monkeypatch): + optimizer, parameter, _, _ = _build_initialized(kind) + _make_persistent_state_oversized(optimizer, parameter) + candidates = _movement_candidates(optimizer, parameter) + assert len(candidates) >= 3 + snapshot = _snapshot_exact_optimizer(optimizer) + completed = _install_late_to_failure( + monkeypatch, + candidates, + destination_type="cpu", + ) + + with pytest.raises(RuntimeError, match="injected late state-copy failure"): + optimizer.move_state_() + + assert len(completed) == 3 + _assert_exact_optimizer_snapshot(optimizer, snapshot) + + +def _build_rejection_case(case): + if case == "capturable": + parameter = torch.nn.Parameter(torch.ones(8)) + optimizer = Gefen( + [("layer.weight", parameter)], + fused=False, + factored_v_2d=False, + capturable=True, + ) + return optimizer, parameter, None + if case == "meta_parameter": + parameter = torch.nn.Parameter(torch.empty(8, device="meta")) + optimizer = Gefen( + [("layer.weight", parameter)], + fused=False, + factored_v_2d=False, + ) + return optimizer, parameter, None + + optimizer, parameter, _, _ = _build_initialized("block") + if case == "meta_destination": + return optimizer, parameter, torch.device("meta") + if case == "mismatched_destination": + return optimizer, parameter, torch.device("cuda:0") + if case == "meta_state": + optimizer.state[parameter]["m_magnitude"] = optimizer.state[parameter][ + "m_magnitude" + ].to("meta") + elif case == "undeclared_tensor": + optimizer.state[parameter]["extension"] = torch.ones(2) + elif case == "tensor_in_extension_container": + optimizer.state[parameter]["extension"] = { + "nested": [torch.ones(2)] + } + elif case == "tensor_in_ordered_extension": + optimizer.state[parameter]["extension"] = OrderedDict( + (("nested", torch.ones(2)),) + ) + elif case == "tensor_in_rank_local_carrier": + optimizer.state[parameter]["_gefen_rank_local_payload_0"] = { + "nested": torch.ones(2) + } + elif case == "nested_tensor_layout": + with warnings.catch_warnings(): + warnings.simplefilter("ignore", UserWarning) + optimizer.state[parameter]["m_magnitude"] = torch.nested.nested_tensor( + [torch.ones(1), torch.ones(2)] + ) + elif case == "tensor_subclass_state": + class StateTensor(torch.Tensor): + pass + + optimizer.state[parameter]["m_magnitude"] = optimizer.state[parameter][ + "m_magnitude" + ].as_subclass(StateTensor) + elif case in {"opaque_tensor_extension", "opaque_scalar_extension"}: + class OpaqueExtension: + def __init__(self, payload): + self.payload = payload + + def __eq__(self, other): + if not isinstance(other, OpaqueExtension): + return False + if torch.is_tensor(self.payload): + return torch.equal(self.payload, other.payload) + return self.payload == other.payload + + payload = torch.ones(1) if case == "opaque_tensor_extension" else "metadata" + optimizer.state[parameter]["extension"] = OpaqueExtension(payload) + elif case == "defaultdict_factory_extension": + hidden = torch.ones(1) + optimizer.state[parameter]["extension"] = defaultdict( + lambda: hidden, + {"metadata": "value"}, + ) + elif case == "ordered_hidden_extension": + extension = OrderedDict((("metadata", "value"),)) + extension.hidden_tensor = torch.ones(1) + optimizer.state[parameter]["extension"] = extension + elif case == "custom_parameter_state_mapping": + custom_state = OrderedDict(optimizer.state[parameter]) + custom_state.hidden_tensor = torch.ones(1) + optimizer.state[parameter] = custom_state + elif case == "custom_top_level_state_mapping": + custom_state = OrderedDict(optimizer.state) + custom_state.hidden_tensor = torch.ones(1) + optimizer.state = custom_state + else: + raise AssertionError("unknown rejection case: {}".format(case)) + return optimizer, parameter, None + + +@pytest.mark.parametrize( + "case", + ( + "capturable", + "meta_parameter", + "meta_destination", + "mismatched_destination", + "meta_state", + "undeclared_tensor", + "tensor_in_extension_container", + "tensor_in_ordered_extension", + "tensor_in_rank_local_carrier", + "nested_tensor_layout", + "tensor_subclass_state", + "opaque_tensor_extension", + "opaque_scalar_extension", + "defaultdict_factory_extension", + "ordered_hidden_extension", + "custom_parameter_state_mapping", + "custom_top_level_state_mapping", + ), +) +def test_invalid_state_movement_is_rejected_before_any_live_mutation(case): + optimizer, _, destination = _build_rejection_case(case) + snapshot = _snapshot_exact_optimizer(optimizer) + + if case not in {"meta_destination", "mismatched_destination"}: + assert not optimizer.optimizer_contract().capabilities.atomic_state_movement + + with pytest.raises((TypeError, ValueError, RuntimeError)): + optimizer.move_state_(destination) + + _assert_exact_optimizer_snapshot(optimizer, snapshot) + + +@pytest.mark.parametrize("graph_type", ("cycle", "shared_container")) +def test_non_tree_extension_metadata_is_rejected_without_mutation(graph_type): + optimizer, parameter, _, _ = _build_initialized("block") + if graph_type == "cycle": + extension = [] + extension.append(extension) + else: + shared = ["metadata"] + extension = [shared, shared] + optimizer.state[parameter]["extension"] = extension + state_before = optimizer.state + parameter_state_before = optimizer.state[parameter] + items_before = tuple(parameter_state_before.items()) + codebook_before = optimizer._gefen_codebook + + assert not optimizer.optimizer_contract().capabilities.atomic_state_movement + with pytest.raises(RuntimeError, match="not provably tensor-free metadata"): + optimizer.move_state_() + + assert optimizer.state is state_before + assert optimizer.state[parameter] is parameter_state_before + assert tuple(optimizer.state[parameter]) == tuple(key for key, _ in items_before) + for key, value in items_before: + assert optimizer.state[parameter][key] is value + assert optimizer._gefen_codebook is codebook_before + + +@pytest.mark.parametrize( + "runtime_state", + ("compile", "capture", "capturable_stacks", "device_counter", "sr_seed"), +) +def test_active_runtime_state_disables_movement_without_mutation( + runtime_state, monkeypatch +): + optimizer, _, _, _ = _build_initialized("block") + if runtime_state == "compile": + monkeypatch.setattr(torch.compiler, "is_compiling", lambda: True) + elif runtime_state == "capture": + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: True) + elif runtime_state == "capturable_stacks": + optimizer._capt_stacks = {} + elif runtime_state == "device_counter": + optimizer._gefen_global_step_by_device[torch.device("cpu")] = torch.tensor(1) + elif runtime_state == "sr_seed": + optimizer._sr_seed_by_device[torch.device("cpu")] = torch.tensor(1) + else: + raise AssertionError("unknown runtime state: {}".format(runtime_state)) + + assert not optimizer.optimizer_contract().capabilities.atomic_state_movement + snapshot = _snapshot_exact_optimizer(optimizer) + with pytest.raises(RuntimeError): + optimizer.move_state_() + _assert_exact_optimizer_snapshot(optimizer, snapshot) + + +def test_stale_finalized_binding_disables_movement_without_mutation(): + parameter = torch.nn.Parameter(torch.ones(8)) + optimizer = Gefen( + [("layer.weight", parameter)], + fused=False, + factored_v_2d=False, + ) + optimizer.rebind_parameter( + parameter, + parameter, + identity=ParameterIdentity("Layer.Weight", (8,)), + ) + replacement = torch.nn.Parameter(torch.full((8,), 2.0)) + optimizer.param_groups[0]["params"][0] = replacement + + assert not optimizer.optimizer_contract().capabilities.atomic_state_movement + snapshot = _snapshot_exact_optimizer(optimizer) + with pytest.raises(RuntimeError, match="finalized parameter layout changed"): + optimizer.move_state_() + _assert_exact_optimizer_snapshot(optimizer, snapshot) + + +def test_movement_does_not_narrow_muons_generic_local_state_device_helper(): + meta_parameter = torch.empty(4, device="meta") + + assert GefenMuon._state_tensor_device(meta_parameter) == torch.device("meta") + with pytest.raises(RuntimeError, match="only CPU and CUDA"): + GefenMuon._state_move_parameter_device(meta_parameter) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +@pytest.mark.parametrize("kind", _KINDS) +def test_cpu_cuda_cpu_state_round_trip_preserves_cpu_continuation(kind): + optimizer, parameter, tensor_lr, _ = _build_initialized(kind) + reference, reference_parameter, reference_lr, _ = _build_initialized(kind) + initial_reference = _persistent_values(reference, reference_parameter) + _assert_persistent_values_equal( + _persistent_values(optimizer, parameter), initial_reference + ) + module = _move_parameter_module(parameter, "cuda") + cuda_grad = parameter.grad + cuda_lr = optimizer.param_groups[0]["lr"] + + optimizer.move_state_() + + assert parameter.grad is cuda_grad + assert optimizer.param_groups[0]["lr"] is cuda_lr is tensor_lr + assert optimizer._gefen_codebook.device.type == "cuda" + for key, value in optimizer.state[parameter].items(): + if key in _MOVABLE_STATE_KEYS and torch.is_tensor(value): + assert value.device.type == "cuda" + + module.to("cpu") + assert module.weight is parameter + cpu_grad = parameter.grad + optimizer.move_state_(torch.device("cpu")) + assert parameter.grad is cpu_grad + assert optimizer.param_groups[0]["lr"] is tensor_lr + assert optimizer._gefen_codebook.device.type == "cpu" + + continuation_grad = torch.linspace(0.6, -0.9, parameter.numel()).reshape_as(parameter) + parameter.grad = continuation_grad.clone() + reference_parameter.grad = continuation_grad.clone() + optimizer.step() + reference.step() + + assert optimizer.param_groups[0]["lr"] is tensor_lr + assert reference.param_groups[0]["lr"] is reference_lr + torch.testing.assert_close(parameter, reference_parameter, rtol=0, atol=0) + _assert_persistent_values_equal( + _persistent_values(optimizer, parameter), + _persistent_values(reference, reference_parameter), + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +def test_real_cpu_to_cuda_copy_failure_is_exactly_atomic(monkeypatch): + optimizer, parameter, _, _ = _build_initialized("block") + module = _move_parameter_module(parameter, "cuda") + assert optimizer._gefen_codebook.device.type == "cpu" + assert optimizer.state[parameter]["m_codebook"].device.type == "cpu" + candidates = _movement_candidates(optimizer, parameter) + snapshot = _snapshot_exact_optimizer(optimizer) + completed = _install_late_to_failure( + monkeypatch, + candidates, + destination_type="cuda", + ) + + with pytest.raises(RuntimeError, match="injected late state-copy failure"): + optimizer.move_state_() + + assert len(completed) == 3 + assert all(tensor.device.type == "cuda" for tensor in completed) + _assert_exact_optimizer_snapshot(optimizer, snapshot) + module.to("cpu") diff --git a/tests/test_state_movement_distributed.py b/tests/test_state_movement_distributed.py new file mode 100644 index 0000000..c109ce7 --- /dev/null +++ b/tests/test_state_movement_distributed.py @@ -0,0 +1,579 @@ +"""Real-process-group coverage for atomic optimizer-state movement.""" + +from __future__ import annotations + +from datetime import timedelta +import os +import queue +import socket +import traceback + +import pytest +import torch + + +_AUTHORITATIVE_TENSOR_KEYS = frozenset( + { + "step", + "m_codebook", + "m_magnitude", + "vmean", + "vmean_step", + "v_row", + "v_col", + "factored_step", + "normuon_v", + "normuon_step", + } +) + + +def _free_port() -> str: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.bind(("127.0.0.1", 0)) + return str(sock.getsockname()[1]) + + +def _oversized_copy(tensor: torch.Tensor) -> torch.Tensor: + backing = torch.empty( + tensor.numel() + 13, + dtype=tensor.dtype, + device=tensor.device, + ) + result = backing.narrow(0, 7, tensor.numel()).view(tensor.shape) + result.copy_(tensor) + assert result.untyped_storage().nbytes() > tensor.numel() * tensor.element_size() + return result + + +def _assert_fresh_tight_copy( + actual: torch.Tensor, + old: torch.Tensor, + expected: torch.Tensor, +) -> None: + assert type(actual) is torch.Tensor + assert actual is not old + assert actual.device == torch.device("cpu") + assert actual.dtype == expected.dtype + assert actual.shape == expected.shape + assert actual.layout is torch.strided + assert actual.is_contiguous() + assert actual.storage_offset() == 0 + assert actual.untyped_storage().nbytes() == actual.numel() * actual.element_size() + torch.testing.assert_close(actual, expected, rtol=0, atol=0, equal_nan=True) + + +def _seed_and_snapshot_movable_state(optimizer, parameters, rank: int, label: str): + first_state = optimizer.state[parameters[0]] + carrier = torch.tensor([rank, len(label), 41], dtype=torch.int64) + metadata = {"label": label, "members": ["preserve", rank]} + first_state["_gefen_rank_local_payload_{}".format(rank)] = carrier + first_state["movement_metadata"] = metadata + + records = [] + for parameter_index, parameter in enumerate(parameters): + for key, value in tuple(optimizer.state[parameter].items()): + if key not in _AUTHORITATIVE_TENSOR_KEYS or not torch.is_tensor(value): + continue + assert type(value) is torch.Tensor + oversized = _oversized_copy(value) + optimizer.state[parameter][key] = oversized + records.append( + ( + parameter_index, + key, + oversized, + oversized.detach().clone(), + ) + ) + + assert type(optimizer._gefen_codebook) is torch.Tensor + optimizer._gefen_codebook = _oversized_copy(optimizer._gefen_codebook) + codebook_record = ( + optimizer._gefen_codebook, + optimizer._gefen_codebook.detach().clone(), + ) + return records, codebook_record, carrier, metadata + + +def _assert_authoritative_state_equal(optimizer, reference, parameters, reference_parameters): + for parameter, reference_parameter in zip(parameters, reference_parameters): + state = optimizer.state[parameter] + reference_state = reference.state[reference_parameter] + keys = { + key + for key in state + if key in _AUTHORITATIVE_TENSOR_KEYS + } | { + key + for key in reference_state + if key in _AUTHORITATIVE_TENSOR_KEYS + } + for key in keys: + assert key in state and key in reference_state + actual = state[key] + expected = reference_state[key] + if torch.is_tensor(expected): + assert torch.is_tensor(actual) + torch.testing.assert_close(actual, expected, rtol=0, atol=0, equal_nan=True) + else: + assert actual == expected + torch.testing.assert_close( + optimizer._gefen_codebook, + reference._gefen_codebook, + rtol=0, + atol=0, + equal_nan=True, + ) + + +def _dtensor_case(rank, world, mesh, kind: str) -> None: + import torch.nn as nn + from torch.distributed.tensor import Shard, distribute_tensor + + from gefen import Gefen, GefenMuon + + shapes = ((4, 4), (4, 6)) if kind == "gefen" else ((4, 4), (1, 4)) + generator = torch.Generator().manual_seed(7100 + sum(ord(item) for item in kind)) + initial = [torch.randn(shape, generator=generator) * 0.05 for shape in shapes] + + def build(): + parameters = [ + nn.Parameter(distribute_tensor(value.clone(), mesh, [Shard(0)])) + for value in initial + ] + tensor_lr = torch.tensor(2.0e-3) + group_metadata = {"kind": kind, "ordered_members": list(range(world))} + group = { + "params": [ + ("{}.{}".format(kind, index), parameter) + for index, parameter in enumerate(parameters) + ], + "lr": tensor_lr, + "movement_metadata": group_metadata, + } + if kind == "gefen": + optimizer = Gefen( + [group], + lr=tensor_lr, + fused=False, + factored_v_2d=False, + ) + else: + optimizer = GefenMuon( + [group], + lr=tensor_lr, + fused=False, + ns_steps=1, + weight_decay=0.0, + sharded_mode=kind, + ) + optimizer._resolve_automatic_period = lambda *args, **kwargs: 4 + optimizer._predict_period_from_grad_sq = lambda *args, **kwargs: 4 + return optimizer, parameters, tensor_lr, group_metadata + + def assign_grads(parameters, seed): + grad_generator = torch.Generator().manual_seed(seed) + for parameter, shape in zip(parameters, shapes): + full_grad = torch.randn(shape, generator=grad_generator) * 0.01 + parameter.grad = distribute_tensor(full_grad, mesh, [Shard(0)]) + + optimizer, parameters, tensor_lr, group_metadata = build() + reference, reference_parameters, _, _ = build() + assign_grads(parameters, 7200) + assign_grads(reference_parameters, 7200) + optimizer.step() + reference.step() + + if kind == "distributed": + stateful = { + index + for index, parameter in enumerate(parameters) + if any( + key in _AUTHORITATIVE_TENSOR_KEYS + for key in optimizer.state[parameter] + ) + } + assert stateful == {rank} + if rank == 1: + assert parameters[1].to_local().numel() == 0 + assert optimizer.state[parameters[1]]["m_codebook"].numel() == shapes[1][0] * shapes[1][1] + elif kind == "approx": + stateful = { + index + for index, parameter in enumerate(parameters) + if any( + key in _AUTHORITATIVE_TENSOR_KEYS + for key in optimizer.state[parameter] + ) + } + assert stateful == ({0, 1} if rank == 0 else {0}) + else: + assert all( + any(key in _AUTHORITATIVE_TENSOR_KEYS for key in optimizer.state[parameter]) + for parameter in parameters + ) + + records, codebook_record, carrier, state_metadata = _seed_and_snapshot_movable_state( + optimizer, + parameters, + rank, + kind, + ) + assert records + contract_before = optimizer.optimizer_contract() + assert contract_before.capabilities.atomic_state_movement + assert not contract_before.capabilities.state_offload + + groups_before = optimizer.param_groups + group_before = optimizer.param_groups[0] + group_params_before = group_before["params"] + defaults_before = optimizer.defaults + names_before = optimizer._param_names + grads_before = [parameter.grad for parameter in parameters] + local_values_before = [parameter.detach().to_local().clone() for parameter in parameters] + meshes_before = [parameter.device_mesh for parameter in parameters] + placements_before = [parameter.placements for parameter in parameters] + mesh_groups_before = [parameter.device_mesh.get_group() for parameter in parameters] + codebook_binding_before = optimizer._gefen_codebook_process_group + shard_bindings_before = optimizer._gefen_shard_bindings + local_bindings_before = optimizer._gefen_local_shard_bindings + manifest_before = optimizer._gefen_sharding_manifest + + optimizer.move_state_() + + assert optimizer.optimizer_contract() == contract_before + assert optimizer.param_groups is groups_before + assert optimizer.param_groups[0] is group_before + assert optimizer.param_groups[0]["params"] is group_params_before + assert optimizer.param_groups[0]["lr"] is tensor_lr + assert optimizer.param_groups[0]["movement_metadata"] is group_metadata + assert optimizer.defaults is defaults_before + assert optimizer.defaults["lr"] is tensor_lr + assert optimizer._param_names is names_before + assert optimizer._gefen_codebook_process_group is codebook_binding_before + assert optimizer._gefen_shard_bindings is shard_bindings_before + assert optimizer._gefen_local_shard_bindings is local_bindings_before + assert optimizer._gefen_sharding_manifest is manifest_before + + for index, parameter in enumerate(parameters): + assert optimizer.param_groups[0]["params"][index] is parameter + assert parameter.grad is grads_before[index] + assert parameter.device_mesh is meshes_before[index] + assert parameter.placements == placements_before[index] + assert parameter.device_mesh.get_group() is mesh_groups_before[index] + torch.testing.assert_close( + parameter.detach().to_local(), + local_values_before[index], + rtol=0, + atol=0, + ) + + for parameter_index, key, old, expected in records: + _assert_fresh_tight_copy( + optimizer.state[parameters[parameter_index]][key], + old, + expected, + ) + _assert_fresh_tight_copy( + optimizer._gefen_codebook, + codebook_record[0], + codebook_record[1], + ) + assert optimizer.state[parameters[0]]["_gefen_rank_local_payload_{}".format(rank)] is carrier + assert optimizer.state[parameters[0]]["movement_metadata"] is state_metadata + + assign_grads(parameters, 7300) + assign_grads(reference_parameters, 7300) + optimizer.step() + reference.step() + for parameter, reference_parameter in zip(parameters, reference_parameters): + torch.testing.assert_close( + parameter.detach().to_local(), + reference_parameter.detach().to_local(), + rtol=0, + atol=0, + ) + _assert_authoritative_state_equal( + optimizer, + reference, + parameters, + reference_parameters, + ) + + +def _whole_owner_case(rank, world) -> None: + import torch.distributed as dist + import torch.nn as nn + + from gefen import ( + CodebookProcessGroupBinding, + GefenMuon, + LogicalSlice, + ParameterIdentity, + ParameterLayout, + ParameterRebinding, + PlacementKind, + ProcessGroupIdentity, + ShardIdentity, + ShardPlacement, + ShardingManifest, + ) + + members = tuple("rank:{}".format(index) for index in range(world)) + group_identity = ProcessGroupIdentity("movement_owner", members) + identities = ( + ParameterIdentity("Owner.First", (4, 4)), + ParameterIdentity("Owner.Second", (4, 4)), + ) + + def owner_shard(identity, member, owner): + return ShardIdentity( + identity, + ParameterLayout.WHOLE_PARAMETER_OWNER, + LogicalSlice.full(identity) if member == owner else LogicalSlice(0, 0), + process_group=group_identity, + local_member=member, + owner=owner, + placements=( + ShardPlacement( + "dp", + PlacementKind.WHOLE_PARAMETER_OWNER, + members.index(member), + world, + ), + ), + ) + + records_by_parameter = tuple( + tuple( + owner_shard(identity, member, members[index]) + for member in members + ) + for index, identity in enumerate(identities) + ) + manifest = ShardingManifest( + tuple( + shard + for parameter_records in records_by_parameter + for shard in parameter_records + ) + ) + + def build(): + generator = torch.Generator().manual_seed(7400) + old_parameters = [ + nn.Parameter(torch.randn(identity.global_shape, generator=generator) * 0.05) + for identity in identities + ] + tensor_lr = torch.tensor(2.0e-3) + group_metadata = {"layout": "whole-owner", "rank": rank} + optimizer = GefenMuon( + [ + { + "params": [ + ("owner.{}".format(index), parameter) + for index, parameter in enumerate(old_parameters) + ], + "lr": tensor_lr, + "movement_metadata": group_metadata, + } + ], + lr=tensor_lr, + fused=False, + ns_steps=1, + weight_decay=0.0, + ) + local_member = members[rank] + local_records = [ + next( + shard + for shard in parameter_records + if shard.local_member == local_member + ) + for parameter_records in records_by_parameter + ] + binding = CodebookProcessGroupBinding( + group_identity, + local_member, + dist.group.WORLD, + torch.device("cpu"), + ) + optimizer.post_sharding( + tuple( + ParameterRebinding( + parameter, + parameter if local_record.owner == local_member else None, + local_record, + ) + for parameter, local_record in zip(old_parameters, local_records) + ), + manifest=manifest, + codebook_process_group=binding, + ) + optimizer._resolve_automatic_period = lambda *args, **kwargs: 4 + optimizer._predict_period_from_grad_sq = lambda *args, **kwargs: 4 + return ( + optimizer, + list(optimizer.param_groups[0]["params"]), + tensor_lr, + group_metadata, + binding, + ) + + def assign_grad(parameters, seed): + generator = torch.Generator().manual_seed(seed + rank) + assert len(parameters) == 1 + parameters[0].grad = torch.randn(parameters[0].shape, generator=generator) * 0.01 + + optimizer, parameters, tensor_lr, group_metadata, binding = build() + reference, reference_parameters, _, _, _ = build() + assert len(parameters) == 1 + assert len(optimizer.shard_bindings()) == 2 + assert sum(parameter is None for parameter, _ in optimizer.shard_bindings()) == 1 + + assign_grad(parameters, 7500) + assign_grad(reference_parameters, 7500) + optimizer.step() + reference.step() + + records, codebook_record, carrier, state_metadata = _seed_and_snapshot_movable_state( + optimizer, + parameters, + rank, + "whole-owner", + ) + assert records + contract_before = optimizer.optimizer_contract() + assert contract_before.capabilities.atomic_state_movement + assert not contract_before.capabilities.state_offload + bindings_before = optimizer.shard_bindings() + manifest_before = optimizer.sharding_manifest() + binding_before = optimizer.codebook_process_group_binding() + groups_before = optimizer.param_groups + group_before = optimizer.param_groups[0] + group_params_before = group_before["params"] + defaults_before = optimizer.defaults + names_before = optimizer._param_names + grad_before = parameters[0].grad + value_before = parameters[0].detach().clone() + + optimizer.move_state_() + + assert optimizer.optimizer_contract() == contract_before + assert optimizer.param_groups is groups_before + assert optimizer.param_groups[0] is group_before + assert optimizer.param_groups[0]["params"] is group_params_before + assert optimizer.param_groups[0]["params"][0] is parameters[0] + assert optimizer.param_groups[0]["lr"] is tensor_lr + assert optimizer.param_groups[0]["movement_metadata"] is group_metadata + assert optimizer.defaults is defaults_before + assert optimizer.defaults["lr"] is tensor_lr + assert optimizer._param_names is names_before + assert optimizer.shard_bindings() is bindings_before + assert optimizer.sharding_manifest() is manifest_before + assert optimizer.codebook_process_group_binding() is binding_before is binding + assert parameters[0].grad is grad_before + torch.testing.assert_close(parameters[0], value_before, rtol=0, atol=0) + + for parameter_index, key, old, expected in records: + _assert_fresh_tight_copy( + optimizer.state[parameters[parameter_index]][key], + old, + expected, + ) + _assert_fresh_tight_copy( + optimizer._gefen_codebook, + codebook_record[0], + codebook_record[1], + ) + assert optimizer.state[parameters[0]]["_gefen_rank_local_payload_{}".format(rank)] is carrier + assert optimizer.state[parameters[0]]["movement_metadata"] is state_metadata + + assign_grad(parameters, 7600) + assign_grad(reference_parameters, 7600) + optimizer.step() + reference.step() + torch.testing.assert_close( + parameters[0], + reference_parameters[0], + rtol=0, + atol=0, + ) + _assert_authoritative_state_equal( + optimizer, + reference, + parameters, + reference_parameters, + ) + + +def _distributed_worker(rank, world, port, result_queue) -> None: + import torch.distributed as dist + from torch.distributed.tensor import 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=90), + ) + mesh = init_device_mesh("cpu", (world,), mesh_dim_names=("dp",)) + for kind in ("gefen", "exact", "approx", "distributed"): + _dtensor_case(rank, world, mesh, kind) + _whole_owner_case(rank, world) + result_queue.put(("result", rank)) + 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="distributed state movement coverage needs Gloo", +) +def test_atomic_state_movement_across_distributed_cpu_representations(): + import torch.multiprocessing as mp + + world = 2 + context = mp.get_context("spawn") + result_queue = context.Queue() + port = _free_port() + processes = [ + context.Process( + target=_distributed_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=180)) + 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 = [item[2] for item in messages if item[0] == "error"] + assert not errors, "\n".join(errors) + assert {item[1] for item in messages if item[0] == "result"} == set(range(world)) From 9a6e33bc40f2aa8ef7d6bd57951313bf3c46c074 Mon Sep 17 00:00:00 2001 From: thad0ctor Date: Sun, 12 Jul 2026 23:38:06 -0700 Subject: [PATCH 08/52] Add portable global-state foundations --- docs/optimizer_contracts.md | 8 +- src/gefen/__init__.py | 20 + src/gefen/checkpoint.py | 142 +++++ src/gefen/contracts.py | 263 ++++++++- src/gefen/portable.py | 498 +++++++++++++++++ src/gefen/portable_schema.py | 395 ++++++++++++++ tests/test_checkpoint_binding.py | 227 ++++++++ tests/test_optimizer_contracts.py | 10 + tests/test_portable_schema.py | 285 ++++++++++ tests/test_portable_state_math.py | 728 +++++++++++++++++++++++++ tests/test_shard_identity_contracts.py | 272 +++++++++ 11 files changed, 2836 insertions(+), 12 deletions(-) create mode 100644 src/gefen/checkpoint.py create mode 100644 src/gefen/portable.py create mode 100644 src/gefen/portable_schema.py create mode 100644 tests/test_checkpoint_binding.py create mode 100644 tests/test_portable_schema.py create mode 100644 tests/test_portable_state_math.py diff --git a/docs/optimizer_contracts.md b/docs/optimizer_contracts.md index efd204a..0d24884 100644 --- a/docs/optimizer_contracts.md +++ b/docs/optimizer_contracts.md @@ -29,7 +29,7 @@ The current DTensor declaration is deliberately narrow: `DTENSOR_1D_DEFAULT_WORL ## 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 logical range, an explicit parameter layout, structured placements, the local member, and an optional whole-parameter owner. `ShardingManifest` validates and deterministically orders the complete identity set; flattened manifests 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. The contiguous-range schema deliberately rejects DTensor identities because column and multidimensional shards require a richer logical-region descriptor; the existing narrow DTensor training and rank-local checkpoint paths remain independently declared. +`ParameterIdentity` records an exact, case-preserving model FQN and global logical shape independently of any live tensor object. `ProcessGroupIdentity` records an adapter-defined semantic group name and authoritative ordered member IDs without importing a framework process-group type. `ShardIdentity` combines those values with a contiguous row-major `LogicalSlice` for replicated, flattened, and whole-owner layouts or an axis-aligned `LogicalRegion` for the narrow one-dimensional DTensor layout, plus structured placements, the local member, and an optional whole-parameter owner. `ShardingManifest` validates and deterministically orders the complete identity set; flattened slices and dimension-sharded regions must cover each logical parameter exactly once without gaps or overlaps, replicated manifests carry one complete identity per declared member, and whole-parameter manifests identify one complete owner while retaining empty non-owner records. DTensor regions currently describe one default-world mesh axis with either replication or one parameter-dimension shard, including uneven and empty shards. This identity vocabulary does not by itself claim DTensor post-sharding rebinding or portable checkpoint support; those remain negative until the optimizer data plane consumes the regions. These descriptors do not treat legacy `param_names`, generated names, Python tensor identity, rank-local parameter IDs, devices, or dtypes as canonical identity. They also do not contain runtime collective handles. An adapter remains responsible for mapping a stable `ProcessGroupIdentity` to its framework process group and for canonicalizing tied aliases to one primary FQN and one optimizer slot; alias-rich identity is not part of schema version 1. Declaring identity metadata alone does not enable rebinding, canonical checkpoint I/O, topology-changing load, codebook scoping, state movement, or offload; those capabilities remain separate. @@ -61,6 +61,12 @@ The v1 document contains only primitive containers and detached tight finite CPU This is an exact-binding transport-neutral local fragment, not the dense global logical state planned for portable DCP v3. Its dynamic `CANONICAL_LOCAL` checkpoint entry covers finalized plain-Gefen replicated and flattened local shards and finalized replicated GefenMuon, performs no collectives, reports atomic local import, and has an empty topology-changing set. A different member, slice, manifest, algorithm policy, group option, or declared state variant rejects. Export, preparation, and commit are quiescent checkpoint-boundary operations; prepared imports use content-bearing freshness tokens, including device counters, to reject intervening mutation. Export remains available after capturable warmup, but a capturable import target must still be fresh, before authoritative device state or a CUDA graph exists; importing replaces state identities, so an already captured graph cannot safely remain attached. Configurations with `stochastic_round=True` do not claim canonical v1 because the decomposed path intentionally lacks the fused stochastic quantizer, so changing effective fused availability would change the algorithm. DTensor, whole-owner completeness, Hybrid composition, rank-fragment gathering, resharding, world-size change, dense momentum decoding, and target-topology recompression remain unclaimed. +## Portable global-state v3 envelope + +The transport-neutral `gefen.portable_state` version-3 format is the structural wire envelope for a future complete `global_logical_optimizer` artifact. 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, optional source provenance, and a completion marker whose deterministic SHA-256 covers every preceding value including tensor dtype, shape, and canonical little-endian bytes. The generic builder and normalizer validate the exact envelope, canonical value grammar, parameter identity records, completion marker, and digest; they do not by themselves prove model-catalog completeness or optimizer-specific state variants, tensor geometries, and projection hints. They accept only finite weights-only-safe primitive values, stream tensor cloning, finite validation, and hashing in bounded chunks, produce detached tight CPU tensors, and reject an incomplete or corrupted digest. Runtime global ranks are deliberately absent from the durable identity. + +`CheckpointProcessGroupBinding` separately binds an adapter-defined `ProcessGroupIdentity` and semantic local member to an explicit live PyTorch process group and CPU/CUDA collective device. Multi-member scopes, including the default world, must pass an explicit handle; `None` means exactly one member. Runtime validation checks initialized membership, group size and coordinate order, and backend/device compatibility without serializing global ranks or executing a collective. The `CANONICAL_GLOBAL` transport enum is defined for this v3 path, but Gefen and GefenMuon do not yet advertise it: a positive capability requires collective fragment registration, dense aggregation, target projection/recompression, unanimous preparation, and atomic local publication to be connected end to end. + ## Quiescent optimizer-state movement `StateMovementProvider.move_state_(device=None)` performs blocking CPU/CUDA co-location movement for Gefen and GefenMuon. With `device=None`, each declared authoritative per-parameter tensor moves to that parameter's current local device, including declared state attached to wrapper-orphaned parameter keys, while the canonical learned codebook moves to the first live local parameter device in parameter-group order. An optimizer with no local parameter storage keeps common state on CPU. An explicit device is accepted only after every live local parameter already resides there; an unindexed CUDA target is resolved from the one co-located live parameter device. The adapter must therefore move the model parameters first and invoke `move_state_` at a quiescent boundary before the next optimizer step. diff --git a/src/gefen/__init__.py b/src/gefen/__init__.py index 0829915..e52302e 100644 --- a/src/gefen/__init__.py +++ b/src/gefen/__init__.py @@ -13,10 +13,12 @@ "GefenMuonHybrid", "CONTRACT_SCHEMA_VERSION", "CANONICAL_STATE_FORMAT_VERSION", + "PORTABLE_STATE_FORMAT_VERSION", "IDENTITY_SCHEMA_VERSION", "CheckpointSupport", "CheckpointTransport", "CanonicalStateProvider", + "CheckpointProcessGroupBinding", "CodebookProcessGroupBinding", "PreparedCanonicalStateImport", "OptimizerCapabilities", @@ -24,6 +26,7 @@ "OptimizerContract", "OptimizerContractProvider", "OptimizerStateLayout", + "LogicalRegion", "LogicalSlice", "ParameterLayout", "ParameterIdentity", @@ -45,6 +48,9 @@ "StateVariant", "TopologyChange", "TrainingSupport", + "build_portable_state_document", + "normalize_portable_state_document", + "portable_state_digest", "split_params_for_muon", "validate_split", "kernels", @@ -77,6 +83,10 @@ def __getattr__(name): from .codebook import CodebookProcessGroupBinding return CodebookProcessGroupBinding + if name == "CheckpointProcessGroupBinding": + from .checkpoint import CheckpointProcessGroupBinding + + return CheckpointProcessGroupBinding if name in ( "CANONICAL_STATE_FORMAT_VERSION", "PreparedCanonicalStateImport", @@ -84,6 +94,15 @@ def __getattr__(name): from . import canonical return getattr(canonical, name) + if name in ( + "PORTABLE_STATE_FORMAT_VERSION", + "build_portable_state_document", + "normalize_portable_state_document", + "portable_state_digest", + ): + from . import portable_schema + + return getattr(portable_schema, name) if name in ( "CONTRACT_SCHEMA_VERSION", "IDENTITY_SCHEMA_VERSION", @@ -95,6 +114,7 @@ def __getattr__(name): "OptimizerContract", "OptimizerContractProvider", "OptimizerStateLayout", + "LogicalRegion", "LogicalSlice", "ParameterLayout", "ParameterIdentity", diff --git a/src/gefen/checkpoint.py b/src/gefen/checkpoint.py new file mode 100644 index 0000000..945d8f3 --- /dev/null +++ b/src/gefen/checkpoint.py @@ -0,0 +1,142 @@ +"""Runtime process-group binding for checkpoint adapters. + +The stable identity contains adapter-defined semantic member IDs. Runtime +global ranks are inspected only while validating the live PyTorch process +group; they are deliberately not part of this binding's durable identity. +""" + +from dataclasses import dataclass +from typing import Optional + +import torch + +from gefen.contracts import ProcessGroupIdentity + + +def _is_process_group(value: object) -> bool: + if not torch.distributed.is_available(): + return False + process_group_type = getattr(torch.distributed, "ProcessGroup", None) + return process_group_type is not None and isinstance(value, process_group_type) + + +@dataclass(frozen=True, eq=False) +class CheckpointProcessGroupBinding: + """Bind a stable checkpoint scope to one live PyTorch process group. + + Multi-member scopes always carry an explicit runtime handle, including + ``torch.distributed.group.WORLD`` when the default world is the intended + scope. A one-member scope is local and therefore uses ``None``. The runtime + handle and runtime global ranks are not canonical checkpoint metadata. + """ + + identity: ProcessGroupIdentity + local_member: str + process_group: Optional[object] + collective_device: torch.device + + def __post_init__(self) -> None: + if not isinstance(self.identity, ProcessGroupIdentity): + raise TypeError("CheckpointProcessGroupBinding.identity must be a ProcessGroupIdentity") + if not isinstance(self.local_member, str): + raise TypeError("CheckpointProcessGroupBinding.local_member must be a string") + if self.local_member not in self.identity.ordered_members: + raise ValueError("CheckpointProcessGroupBinding.local_member must belong to the identity") + if not isinstance(self.collective_device, torch.device): + raise TypeError("CheckpointProcessGroupBinding.collective_device must be a torch.device") + if self.collective_device.type not in {"cpu", "cuda"}: + raise ValueError("checkpoint collective device must be CPU or CUDA") + + member_count = len(self.identity.ordered_members) + if member_count == 1: + if self.process_group is not None: + raise ValueError("a one-member checkpoint scope must use process_group=None") + else: + if self.process_group is None: + raise ValueError( + "a multi-member checkpoint scope requires an explicit runtime process group" + ) + if not _is_process_group(self.process_group): + raise TypeError( + "CheckpointProcessGroupBinding.process_group must be a torch.distributed.ProcessGroup" + ) + + @property + def sort_key(self): + """Return the rank-invariant key used to order checkpoint scopes.""" + + return (self.identity.semantic_name, self.identity.ordered_members) + + def validate_runtime(self) -> None: + """Fail unless this process belongs at the declared semantic coordinate. + + This performs only local process-group metadata queries. It does not + execute a collective or retain the runtime group's global-rank list. + """ + + members = self.identity.ordered_members + if len(members) == 1: + return + if not torch.distributed.is_available() or not torch.distributed.is_initialized(): + raise RuntimeError( + "a multi-member checkpoint scope requires initialized torch.distributed" + ) + + import torch.distributed as dist + + try: + world_size = dist.get_world_size(self.process_group) + runtime_global_members = tuple(dist.get_process_group_ranks(self.process_group)) + current_global_rank = dist.get_rank() + group_rank = dist.get_group_rank(self.process_group, current_global_rank) + roundtrip_global_rank = dist.get_global_rank(self.process_group, group_rank) + backend = dist.get_backend(self.process_group) + except Exception as exc: + raise ValueError( + "the current rank must belong to the explicit checkpoint process group" + ) from exc + + if world_size != len(members) or len(runtime_global_members) != world_size: + raise ValueError( + "checkpoint process-group world size does not match its stable identity" + ) + if ( + group_rank < 0 + or group_rank >= world_size + or runtime_global_members[group_rank] != current_global_rank + or roundtrip_global_rank != current_global_rank + ): + raise ValueError( + "the current global rank is not a consistent member of the checkpoint process group" + ) + if members[group_rank] != self.local_member: + raise ValueError( + "runtime checkpoint group order does not match ordered semantic members" + ) + self._validate_collective_device(backend) + + def _validate_collective_device(self, backend: object) -> None: + backend_name = str(backend).lower() + if "nccl" in backend_name: + if self.collective_device.type != "cuda": + raise ValueError( + "checkpoint collective device is incompatible with the runtime backend" + ) + elif "gloo" in backend_name or "mpi" in backend_name: + if self.collective_device.type != "cpu": + raise ValueError( + "checkpoint collective device is incompatible with the runtime backend" + ) + + if self.collective_device.type == "cuda": + index = self.collective_device.index + if ( + not torch.cuda.is_available() + or index is None + or index < 0 + or index >= torch.cuda.device_count() + ): + raise ValueError("checkpoint collective CUDA device is unavailable") + + +__all__ = ["CheckpointProcessGroupBinding"] diff --git a/src/gefen/contracts.py b/src/gefen/contracts.py index f71caf4..7b80234 100644 --- a/src/gefen/contracts.py +++ b/src/gefen/contracts.py @@ -15,6 +15,7 @@ Protocol, Sequence, Tuple, + Union, runtime_checkable, ) @@ -101,6 +102,7 @@ class CheckpointTransport(str, Enum): PYTORCH_RANK_LOCAL = "pytorch_rank_local" COMPOSITE_NATIVE = "composite_native" CANONICAL_LOCAL = "canonical_local" + CANONICAL_GLOBAL = "canonical_global" class TopologyChange(str, Enum): @@ -243,13 +245,139 @@ def full(cls, parameter: ParameterIdentity) -> "LogicalSlice": return cls(0, parameter.numel) +@dataclass(frozen=True) +class LogicalRegion: + """One axis-aligned region in canonical logical parameter coordinates.""" + + offsets: Sequence[int] + lengths: Sequence[int] + + def __post_init__(self) -> None: + for name, values in ( + ("LogicalRegion.offsets", self.offsets), + ("LogicalRegion.lengths", self.lengths), + ): + if isinstance(values, (str, bytes, bytearray)): + raise TypeError("{} must be a sequence of dimensions".format(name)) + try: + normalized = tuple(values) + except TypeError as exc: + raise TypeError( + "{} must be a sequence of dimensions".format(name) + ) from exc + if any(type(value) is not int or value < 0 for value in normalized): + raise ValueError( + "{} must contain nonnegative integers".format(name) + ) + object.__setattr__(self, name.rsplit(".", 1)[1], normalized) + if len(self.offsets) != len(self.lengths): + raise ValueError( + "LogicalRegion offsets and lengths must have the same rank" + ) + + @property + def rank(self) -> int: + """Return the logical tensor rank described by this region.""" + + return len(self.offsets) + + @property + def numel(self) -> int: + """Return the number of logical elements in this region.""" + + return math.prod(self.lengths) + + @classmethod + def full(cls, parameter: ParameterIdentity) -> "LogicalRegion": + """Return the complete axis-aligned region for ``parameter``.""" + + if not isinstance(parameter, ParameterIdentity): + raise TypeError("parameter must be a ParameterIdentity") + return cls((0,) * len(parameter.global_shape), parameter.global_shape) + + def validate_bounds(self, parameter: ParameterIdentity) -> None: + """Raise when this region is not contained by ``parameter``.""" + + if not isinstance(parameter, ParameterIdentity): + raise TypeError("parameter must be a ParameterIdentity") + if self.rank != len(parameter.global_shape): + raise ValueError( + "LogicalRegion rank must match the global parameter rank" + ) + if any( + offset + length > dimension + for offset, length, dimension in zip( + self.offsets, self.lengths, parameter.global_shape + ) + ): + raise ValueError("LogicalRegion exceeds the global parameter") + + def intersection(self, other: "LogicalRegion") -> "LogicalRegion": + """Return the axis-aligned intersection with another same-rank region.""" + + if not isinstance(other, LogicalRegion): + raise TypeError("other must be a LogicalRegion") + if self.rank != other.rank: + raise ValueError("LogicalRegion intersection requires equal ranks") + offsets = tuple( + max(left, right) + for left, right in zip(self.offsets, other.offsets) + ) + ends = tuple( + min(left_offset + left_length, right_offset + right_length) + for left_offset, left_length, right_offset, right_length in zip( + self.offsets, + self.lengths, + other.offsets, + other.lengths, + ) + ) + return LogicalRegion( + offsets, + tuple(max(0, end - offset) for offset, end in zip(offsets, ends)), + ) + + def overlaps(self, other: "LogicalRegion") -> bool: + """Return whether two regions share at least one logical element.""" + + return self.intersection(other).numel > 0 + + @staticmethod + def validate_exact_coverage( + parameter: ParameterIdentity, regions: Sequence["LogicalRegion"] + ) -> None: + """Validate that bounded regions cover a parameter exactly once.""" + + if not isinstance(parameter, ParameterIdentity): + raise TypeError("parameter must be a ParameterIdentity") + if isinstance(regions, (str, bytes, bytearray)): + raise TypeError("regions must be a sequence of LogicalRegion values") + try: + regions = tuple(regions) + except TypeError as exc: + raise TypeError( + "regions must be a sequence of LogicalRegion values" + ) from exc + if any(not isinstance(region, LogicalRegion) for region in regions): + raise TypeError("regions must contain LogicalRegion values") + for region in regions: + region.validate_bounds(parameter) + for index, region in enumerate(regions): + if any(region.overlaps(other) for other in regions[index + 1 :]): + raise ValueError("LogicalRegions must not overlap") + if sum(region.numel for region in regions) != parameter.numel: + raise ValueError( + "LogicalRegions must exactly cover the global parameter" + ) + + @dataclass(frozen=True) class ShardIdentity: """Stable identity of one process-group member's logical parameter shard.""" parameter: ParameterIdentity layout: ParameterLayout - logical_slice: LogicalSlice + logical_slice: Union[LogicalSlice, LogicalRegion] placements: Sequence[ShardPlacement] = () process_group: Optional[ProcessGroupIdentity] = None local_member: Optional[str] = None @@ -261,8 +389,18 @@ def __post_init__(self) -> None: raise TypeError("ShardIdentity.parameter must be a ParameterIdentity") if not isinstance(self.layout, ParameterLayout): raise TypeError("ShardIdentity.layout must be a ParameterLayout") - if not isinstance(self.logical_slice, LogicalSlice): - raise TypeError("ShardIdentity.logical_slice must be a LogicalSlice") + uses_logical_region = isinstance(self.logical_slice, LogicalRegion) + if self.layout is ParameterLayout.DTENSOR_1D_DEFAULT_WORLD: + if not uses_logical_region: + raise ValueError( + "stable DTensor identity requires a logical-region descriptor " + "(LogicalRegion)" + ) + self.logical_slice.validate_bounds(self.parameter) + elif not isinstance(self.logical_slice, LogicalSlice): + raise TypeError( + "non-DTensor ShardIdentity.logical_slice must be a LogicalSlice" + ) placements = _tuple(self.placements) if any(not isinstance(item, ShardPlacement) for item in placements): raise TypeError("ShardIdentity.placements must contain ShardPlacement values") @@ -274,7 +412,11 @@ def __post_init__(self) -> None: "placements", tuple(sorted(placements, key=lambda item: item.mesh_axis)), ) - if self.logical_slice.flat_offset + self.logical_slice.length > self.parameter.numel: + if ( + not uses_logical_region + and self.logical_slice.flat_offset + self.logical_slice.length + > self.parameter.numel + ): raise ValueError("ShardIdentity.logical_slice exceeds the global parameter") if self.process_group is None: if self.local_member is not None or self.owner is not None: @@ -289,7 +431,11 @@ def __post_init__(self) -> None: if self.layout is not ParameterLayout.WHOLE_PARAMETER_OWNER and self.owner is not None: raise ValueError("ShardIdentity.owner is valid only for whole-parameter ownership") - full = self.logical_slice == LogicalSlice.full(self.parameter) + full = ( + self.logical_slice == LogicalRegion.full(self.parameter) + if uses_logical_region + else self.logical_slice == LogicalSlice.full(self.parameter) + ) kinds = tuple(item.kind for item in self.placements) if self.process_group is not None: member_index = self.process_group.ordered_members.index(self.local_member) @@ -326,12 +472,50 @@ def __post_init__(self) -> None: if len(kinds) != 1 or kinds[0] is not PlacementKind.WHOLE_PARAMETER_OWNER: raise ValueError("whole-parameter ownership requires one owner placement") elif self.layout is ParameterLayout.DTENSOR_1D_DEFAULT_WORLD: - raise ValueError( - "stable DTensor identity requires a logical-region descriptor and is " - "not implemented by the contiguous-slice identity schema" - ) + if self.process_group is None or self.owner is not None: + raise ValueError( + "DTensor identities require a process group and no owner" + ) + if len(self.placements) != 1 or kinds[0] not in { + PlacementKind.DIMENSION_SHARD, + PlacementKind.REPLICATE, + }: + raise ValueError( + "one-dimensional DTensor identities require one dimension-shard " + "or replicate placement" + ) + placement = self.placements[0] + if placement.kind is PlacementKind.REPLICATE: + if not full: + raise ValueError( + "replicated DTensor identities must cover the full parameter" + ) + else: + shard_dimension = placement.parameter_dimension + for dimension, (offset, length, global_length) in enumerate( + zip( + self.logical_slice.offsets, + self.logical_slice.lengths, + self.parameter.global_shape, + ) + ): + if dimension != shard_dimension and ( + offset != 0 or length != global_length + ): + raise ValueError( + "a dimension-sharded DTensor region must cover every " + "unsharded parameter dimension" + ) _validate_identity_schema_version("shard identity", self.schema_version) + @property + def logical_region(self) -> Optional[LogicalRegion]: + """Return the axis-aligned logical region, when this identity has one.""" + + if isinstance(self.logical_slice, LogicalRegion): + return self.logical_slice + return None + @property def sort_key(self): """Return a deterministic structural ordering key.""" @@ -354,15 +538,39 @@ def sort_key(self): ) for item in self.placements ) + if isinstance(self.logical_slice, LogicalSlice): + # Preserve the released contiguous-slice key exactly. Besides + # retaining its public structural shape, this leaves all existing + # manifest ordering and freshness tokens byte-for-byte stable. + return ( + self.parameter.fqn, + self.logical_slice.flat_offset, + self.logical_slice.length, + self.layout.value, + group_name, + member_index, + owner_index, + placement_key, + ) + + flat_offset = 0 + stride = 1 + for offset, dimension in reversed( + tuple(zip(self.logical_slice.offsets, self.parameter.global_shape)) + ): + flat_offset += offset * stride + stride *= dimension return ( self.parameter.fqn, - self.logical_slice.flat_offset, - self.logical_slice.length, + flat_offset, + self.logical_slice.numel, self.layout.value, group_name, member_index, owner_index, placement_key, + self.logical_slice.offsets, + self.logical_slice.lengths, ) @@ -451,6 +659,38 @@ def __post_init__(self) -> None: owners = {item.owner for item in parameter_shards} if len(owners) != 1: raise ValueError("whole-parameter manifest shards must agree on one owner") + elif layout is ParameterLayout.DTENSOR_1D_DEFAULT_WORLD: + placements = tuple(item.placements[0] for item in parameter_shards) + placement_kind = placements[0].kind + if placement_kind is PlacementKind.REPLICATE: + if any( + item.logical_region != LogicalRegion.full(parameter) + for item in parameter_shards + ): + raise ValueError( + "replicated DTensor manifest regions must all be complete" + ) + continue + + shard_dimension = placements[0].parameter_dimension + by_coordinate = sorted( + parameter_shards, + key=lambda item: item.placements[0].coordinate, + ) + cursor = 0 + for item in by_coordinate: + region = item.logical_region + if region.offsets[shard_dimension] != cursor: + raise ValueError( + "dimension-sharded DTensor manifest regions must be " + "gapless and non-overlapping in coordinate order" + ) + cursor += region.lengths[shard_dimension] + if cursor != parameter.global_shape[shard_dimension]: + raise ValueError( + "dimension-sharded DTensor manifest regions must cover " + "the global parameter" + ) def for_parameter(self, fqn: str) -> Tuple[ShardIdentity, ...]: """Return one canonical parameter's shards in deterministic order.""" @@ -1582,6 +1822,7 @@ def _hybrid_contract( "OptimizerContract", "OptimizerContractProvider", "OptimizerStateLayout", + "LogicalRegion", "LogicalSlice", "ParameterLayout", "ParameterIdentity", diff --git a/src/gefen/portable.py b/src/gefen/portable.py new file mode 100644 index 0000000..118ed56 --- /dev/null +++ b/src/gefen/portable.py @@ -0,0 +1,498 @@ +"""Strict pure math for topology-independent Gefen optimizer state.""" + +import math + +import torch + + +_PORTABLE_STATE_CHUNK_ELEMENTS = 1 << 20 + + +def _chunk_element_budget() -> int: + budget = _PORTABLE_STATE_CHUNK_ELEMENTS + if type(budget) is not int or budget <= 0: + raise RuntimeError("portable state chunk budget must be a positive int") + return budget + + +def _element_chunks(numel: int): + budget = _chunk_element_budget() + for start in range(0, numel, budget): + yield start, min(start + budget, numel) + + +def _whole_row_chunks(rows: int, width: int): + if rows == 0: + return + rows_per_chunk = max(1, _chunk_element_budget() // max(1, width)) + for start in range(0, rows, rows_per_chunk): + yield start, min(start + rows_per_chunk, rows) + + +def _read_flat_chunk(value: torch.Tensor, start: int, stop: int) -> torch.Tensor: + """Read one logical row-major flat slice without flattening the whole input.""" + + detached = value.detach() + if detached.is_contiguous(): + return detached.reshape(-1)[start:stop] + linear = torch.arange(start, stop, dtype=torch.int64, device=detached.device) + remainder = linear + reversed_coordinates = [] + for dimension in reversed(detached.shape): + reversed_coordinates.append(torch.remainder(remainder, dimension)) + remainder = torch.div(remainder, dimension, rounding_mode="floor") + coordinates = tuple(reversed(reversed_coordinates)) + return detached[coordinates] + + +def _tensor_values_valid(value: torch.Tensor, *, nonnegative: bool) -> tuple[bool, bool]: + finite = True + nonnegative_values = True + for start, stop in _element_chunks(value.numel()): + chunk = _read_flat_chunk(value, start, stop) + try: + if not bool(torch.isfinite(chunk).all()): + finite = False + break + except (NotImplementedError, RuntimeError, TypeError) as exc: + raise TypeError("tensor dtype does not support finite-state validation") from exc + if nonnegative and not bool((chunk >= 0).all()): + nonnegative_values = False + break + return finite, nonnegative_values + + +def _validate_state_counter(value, *, name: str, minimum: int = 0) -> int: + """Return one strict host counter after validating its lower bound.""" + + if type(name) is not str: + raise TypeError("counter name must be a string") + if not name: + raise ValueError("counter name must not be empty") + if type(minimum) is not int: + raise TypeError("counter minimum must be an int") + if minimum < 0: + raise ValueError("counter minimum must be nonnegative") + if type(value) is not int: + raise TypeError("{} must be a host int".format(name)) + if value < minimum: + raise ValueError("{} must be at least {}".format(name, minimum)) + return value + + +def _validate_plain_tensor( + value, + *, + name: str, + dtype: torch.dtype, + ndim: int | None = None, + nonnegative: bool = False, +) -> torch.Tensor: + if type(value) is not torch.Tensor: + raise TypeError("{} must be a plain torch.Tensor".format(name)) + if ( + value.layout is not torch.strided + or value.is_meta + or value.is_nested + or value.is_quantized + ): + raise TypeError("{} must be a materialized strided tensor".format(name)) + if value.dtype != dtype: + raise TypeError("{} must have dtype {}".format(name, dtype)) + if ndim is not None and value.ndim != ndim: + raise ValueError("{} must be {}-D".format(name, ndim)) + try: + finite, nonnegative_values = _tensor_values_valid( + value, + nonnegative=nonnegative, + ) + except TypeError as exc: + raise TypeError("{} dtype does not support finite-state validation".format(name)) from exc + if not finite: + raise ValueError("{} must be finite".format(name)) + if nonnegative and not nonnegative_values: + raise ValueError("{} must be nonnegative".format(name)) + return value + + +def _validate_logical_shape(logical_shape) -> tuple[int, ...]: + if type(logical_shape) not in {tuple, torch.Size}: + raise TypeError("logical_shape must be a tuple or torch.Size") + shape = tuple(logical_shape) + for dimension in shape: + if type(dimension) is not int: + raise TypeError("logical_shape dimensions must be ints") + if dimension < 0: + raise ValueError("logical_shape dimensions must be nonnegative") + return shape + + +def _validate_period(period, *, numel: int) -> int: + if type(period) is not int: + raise TypeError("period must be an int") + if period <= 0: + raise ValueError("period must be positive") + if numel % period != 0: + raise ValueError( + "logical element count {} is not divisible by period {}".format( + numel, + period, + ) + ) + return period + + +def _validate_codebook(codebook) -> torch.Tensor: + codebook = _validate_plain_tensor( + codebook, + name="codebook", + dtype=torch.float32, + ndim=1, + ) + if not 1 <= codebook.numel() <= 256: + raise ValueError("codebook must contain between 1 and 256 entries") + if codebook.numel() > 1 and not bool( + torch.all(codebook.detach()[1:] >= codebook.detach()[:-1]) + ): + raise ValueError("codebook must be sorted in nondecreasing order") + if not bool(((codebook.detach() >= -1.0) & (codebook.detach() <= 1.0)).all()): + raise ValueError("codebook entries must lie in [-1, 1]") + # ``searchsorted`` warns and performs its own hidden copy for a strided + # boundary tensor. A Gefen codebook has at most 256 elements, so normalize + # this bounded input once and reuse it across every logical-state chunk. + return codebook.detach().contiguous() + + +def _new_output(shape, *, dtype: torch.dtype, device: torch.device) -> torch.Tensor: + return torch.empty(shape, dtype=dtype, device=device) + + +def _finish_output(result: torch.Tensor, *, name: str) -> torch.Tensor: + if result.is_floating_point() and not _tensor_values_valid( + result, + nonnegative=False, + )[0]: + raise ValueError("{} cannot be represented as finite {} state".format(name, result.dtype)) + if ( + result.requires_grad + or not result.is_contiguous() + or result.storage_offset() != 0 + or result.untyped_storage().nbytes() != result.numel() * result.element_size() + ): + raise RuntimeError("{} projection did not produce tight detached state".format(name)) + return result + + +def _nearest_codebook_indices( + codebook: torch.Tensor, + normalized_values: torch.Tensor, +) -> torch.Tensor: + """Return nearest codewords, resolving equal distances to the lower index.""" + + if codebook.device != normalized_values.device: + raise ValueError("codebook and normalized values must be on the same device") + codebook = codebook.detach().contiguous() + output = _new_output( + normalized_values.shape, + dtype=torch.uint8, + device=normalized_values.device, + ) + if codebook.numel() == 1: + output.zero_() + return _finish_output(output, name="momentum indices") + + output_flat = output.reshape(-1) + for start, stop in _element_chunks(normalized_values.numel()): + segment = _read_flat_chunk(normalized_values, start, stop).float() + insertion = torch.searchsorted(codebook, segment) + left = (insertion - 1).clamp_(0, codebook.numel() - 1) + right = insertion.clamp_(0, codebook.numel() - 1) + indices = torch.where( + (segment - codebook[left]).abs() + <= (segment - codebook[right]).abs(), + left, + right, + ) + output_flat[start:stop].copy_(indices.to(torch.uint8)) + return _finish_output(output, name="momentum indices") + + +def _decode_quantized_momentum( + codebook, + indices, + magnitudes, + *, + logical_shape, + period, + step, +) -> torch.Tensor: + """Validate and decode block-quantized momentum to dense logical fp32 state.""" + + _validate_state_counter(step, name="step", minimum=1) + shape = _validate_logical_shape(logical_shape) + numel = math.prod(shape) + period = _validate_period(period, numel=numel) + codebook = _validate_codebook(codebook) + indices = _validate_plain_tensor( + indices, + name="momentum indices", + dtype=torch.uint8, + ndim=2, + ) + magnitudes = _validate_plain_tensor( + magnitudes, + name="momentum magnitudes", + dtype=torch.float32, + ndim=2, + nonnegative=True, + ) + if indices.device != codebook.device or magnitudes.device != codebook.device: + raise ValueError("codebook, momentum indices, and magnitudes must share a device") + + blocks = numel // period + if tuple(indices.shape) != (blocks, period): + raise ValueError( + "momentum indices must have shape ({}, {})".format(blocks, period) + ) + if tuple(magnitudes.shape) != (blocks, 1): + raise ValueError("momentum magnitudes must have shape ({}, 1)".format(blocks)) + for start, stop in _element_chunks(indices.numel()): + if int(_read_flat_chunk(indices, start, stop).max().item()) >= codebook.numel(): + raise ValueError("momentum indices contain an out-of-range codebook entry") + + dense = _new_output(shape, dtype=torch.float32, device=codebook.device) + dense_flat = dense.reshape(-1) + for row_start, row_stop in _whole_row_chunks(blocks, period): + flat_start = row_start * period + flat_stop = row_stop * period + index_chunk = _read_flat_chunk(indices, flat_start, flat_stop).long() + coefficients = codebook.index_select(0, index_chunk) + magnitude_chunk = _read_flat_chunk( + magnitudes, + row_start, + row_stop, + ).reshape(-1, 1) + decoded = coefficients.reshape(-1, period) * magnitude_chunk + dense_flat[flat_start:flat_stop].copy_(decoded.reshape(-1)) + return _finish_output(dense, name="dense momentum") + + +def _recompress_dense_momentum( + momentum, + codebook, + *, + period, + step, +) -> tuple[torch.Tensor, torch.Tensor]: + """Compress dense fp32 momentum into one magnitude and byte indices per block.""" + + _validate_state_counter(step, name="step", minimum=1) + momentum = _validate_plain_tensor( + momentum, + name="dense momentum", + dtype=torch.float32, + ) + codebook = _validate_codebook(codebook) + if momentum.device != codebook.device: + raise ValueError("dense momentum and codebook must share a device") + period = _validate_period(period, numel=momentum.numel()) + + blocks = momentum.numel() // period + indices = _new_output( + (blocks, period), + dtype=torch.uint8, + device=momentum.device, + ) + magnitudes = _new_output( + (blocks, 1), + dtype=torch.float32, + device=momentum.device, + ) + for row_start, row_stop in _whole_row_chunks(blocks, period): + flat_start = row_start * period + flat_stop = row_stop * period + block_values = _read_flat_chunk( + momentum, + flat_start, + flat_stop, + ).reshape(-1, period) + magnitude_chunk = block_values.abs().amax(dim=1, keepdim=True) + normalized = block_values.clone() + nonzero = magnitude_chunk > 0 + normalized.div_(magnitude_chunk) + normalized.masked_fill_(~nonzero, 0.0) + index_chunk = _nearest_codebook_indices(codebook, normalized) + indices[row_start:row_stop].copy_(index_chunk) + magnitudes[row_start:row_stop].copy_(magnitude_chunk) + return ( + _finish_output(indices, name="momentum indices"), + _finish_output(magnitudes, name="momentum magnitudes"), + ) + + +def _expand_block_second_moment( + block_values, + *, + logical_shape, + period, + step, +) -> torch.Tensor: + """Expand one nonnegative fp32 scalar per block to dense logical state.""" + + _validate_state_counter(step, name="vmean_step", minimum=1) + shape = _validate_logical_shape(logical_shape) + numel = math.prod(shape) + period = _validate_period(period, numel=numel) + block_values = _validate_plain_tensor( + block_values, + name="block second moment", + dtype=torch.float32, + ndim=2, + nonnegative=True, + ) + blocks = numel // period + if tuple(block_values.shape) != (blocks, 1): + raise ValueError("block second moment must have shape ({}, 1)".format(blocks)) + + dense = _new_output(shape, dtype=torch.float32, device=block_values.device) + dense_flat = dense.reshape(-1) + for row_start, row_stop in _whole_row_chunks(blocks, period): + flat_start = row_start * period + flat_stop = row_stop * period + block_chunk = _read_flat_chunk( + block_values, + row_start, + row_stop, + ).reshape(-1, 1) + dense_flat[flat_start:flat_stop].reshape(-1, period).copy_( + block_chunk.expand(-1, period) + ) + return _finish_output(dense, name="dense second moment") + + +def _reduce_block_second_moment( + dense, + *, + period, + step, +) -> torch.Tensor: + """Reduce dense nonnegative fp32 state to target-block arithmetic means.""" + + _validate_state_counter(step, name="vmean_step", minimum=1) + dense = _validate_plain_tensor( + dense, + name="dense second moment", + dtype=torch.float32, + nonnegative=True, + ) + period = _validate_period(period, numel=dense.numel()) + blocks = dense.numel() // period + reduced = _new_output( + (blocks, 1), + dtype=torch.float32, + device=dense.device, + ) + for row_start, row_stop in _whole_row_chunks(blocks, period): + flat_start = row_start * period + flat_stop = row_stop * period + values64 = _read_flat_chunk(dense, flat_start, flat_stop).reshape( + -1, + period, + ).to(torch.float64) + reduced[row_start:row_stop].copy_( + values64.mean(dim=1, keepdim=True).to(torch.float32) + ) + return _finish_output(reduced, name="block second moment") + + +def _expand_factored_second_moment( + row, + column, + *, + logical_shape, + step, +) -> torch.Tensor: + """Expand Adafactor row/column state as ``outer(row, column) / mean(row)``.""" + + _validate_state_counter(step, name="factored_step", minimum=1) + shape = _validate_logical_shape(logical_shape) + if len(shape) != 2: + raise ValueError("factored second moment requires a 2-D logical shape") + row = _validate_plain_tensor( + row, + name="row second moment", + dtype=torch.float32, + ndim=1, + nonnegative=True, + ) + column = _validate_plain_tensor( + column, + name="column second moment", + dtype=torch.float32, + ndim=1, + nonnegative=True, + ) + if row.device != column.device: + raise ValueError("row and column second moments must share a device") + if tuple(row.shape) != (shape[0],) or tuple(column.shape) != (shape[1],): + raise ValueError("factored second-moment vectors do not match logical_shape") + dense = _new_output(shape, dtype=torch.float32, device=row.device) + if math.prod(shape) == 0: + return _finish_output(dense, name="dense second moment") + + row_mean = row.detach().mean(dtype=torch.float64) + if bool(row_mean == 0): + has_nonzero_column = any( + bool((_read_flat_chunk(column, start, stop) != 0).any()) + for start, stop in _element_chunks(column.numel()) + ) + if has_nonzero_column: + raise ValueError( + "zero row mean is inconsistent with a nonzero column second moment" + ) + dense.zero_() + else: + column64 = column.detach().to(torch.float64) + for row_start, row_stop in _whole_row_chunks(shape[0], shape[1]): + row64 = _read_flat_chunk(row, row_start, row_stop).to(torch.float64) + expanded = torch.outer(row64, column64).div_(row_mean).to(torch.float32) + dense[row_start:row_stop].copy_(expanded) + return _finish_output(dense, name="dense second moment") + + +def _project_factored_second_moment( + dense, + *, + step, +) -> tuple[torch.Tensor, torch.Tensor]: + """Project dense 2-D second moment to its row and column arithmetic means.""" + + _validate_state_counter(step, name="factored_step", minimum=1) + dense = _validate_plain_tensor( + dense, + name="dense second moment", + dtype=torch.float32, + ndim=2, + nonnegative=True, + ) + rows, columns = dense.shape + row = _new_output((rows,), dtype=torch.float32, device=dense.device) + column = _new_output((columns,), dtype=torch.float32, device=dense.device) + if dense.numel() == 0: + row.zero_() + column.zero_() + else: + for matrix, output in ((dense.detach(), row), (dense.detach().t(), column)): + matrix_rows, width = matrix.shape + for row_start, row_stop in _whole_row_chunks(matrix_rows, width): + values64 = matrix[row_start:row_stop].to(torch.float64) + output[row_start:row_stop].copy_( + values64.mean(dim=1).to(torch.float32) + ) + return ( + _finish_output(row, name="row second moment"), + _finish_output(column, name="column second moment"), + ) + + +__all__ = [] diff --git a/src/gefen/portable_schema.py b/src/gefen/portable_schema.py new file mode 100644 index 0000000..a50244b --- /dev/null +++ b/src/gefen/portable_schema.py @@ -0,0 +1,395 @@ +"""Strict source-topology-neutral envelope for portable optimizer state.""" + +import hashlib +import hmac +import math +import struct +import sys + +import torch + +from gefen.contracts import ParameterIdentity + + +PORTABLE_STATE_FORMAT = "gefen.portable_state" +PORTABLE_STATE_FORMAT_VERSION = 3 +PORTABLE_STATE_COVERAGE = "global_logical_optimizer" +PORTABLE_STATE_DIGEST_ALGORITHM = "sha256" +_PORTABLE_DIGEST_CHUNK_BYTES = 1 << 23 +_PORTABLE_CLONE_CHUNK_BYTES = 1 << 23 +_PORTABLE_NATIVE_BYTEORDER = sys.byteorder + +_PORTABLE_STATE_TOP_LEVEL_KEYS = frozenset( + { + "format", + "format_version", + "coverage", + "implementation", + "policy", + "common", + "parameters", + "provenance", + "completion", + } +) +_PORTABLE_PARAMETER_KEYS = frozenset( + { + "identity", + "algorithm_options", + "state_variant", + "state", + "projection_hints", + } +) +_PORTABLE_IDENTITY_KEYS = frozenset( + {"schema_version", "fqn", "global_shape"} +) +_PORTABLE_COMPLETION_KEYS = frozenset( + {"status", "digest_algorithm", "digest"} +) + + +def _portable_tensor_chunk_elements(value: torch.Tensor) -> int: + budget = _PORTABLE_CLONE_CHUNK_BYTES + if type(budget) is not int or budget <= 0: + raise RuntimeError("portable clone chunk budget must be a positive int") + return max(1, budget // value.element_size()) + + +def _read_portable_tensor_chunk( + value: torch.Tensor, start: int, stop: int +) -> torch.Tensor: + detached = value.detach() + if detached.is_contiguous(): + return detached.reshape(-1)[start:stop] + linear = torch.arange(start, stop, dtype=torch.int64, device=detached.device) + remainder = linear + reversed_coordinates = [] + for dimension in reversed(detached.shape): + reversed_coordinates.append(torch.remainder(remainder, dimension)) + remainder = torch.div(remainder, dimension, rounding_mode="floor") + return detached[tuple(reversed(reversed_coordinates))] + + +def _clone_portable_value(value, *, path: str): + if torch.is_tensor(value): + if ( + type(value) is not torch.Tensor + or value.layout is not torch.strided + or value.is_meta + or value.is_nested + or value.is_quantized + ): + raise TypeError( + "{} must be a plain materialized strided tensor for portable state".format( + path + ) + ) + cloned = torch.empty(tuple(value.shape), dtype=value.dtype, device="cpu") + cloned_flat = cloned.reshape(-1) + chunk_elements = _portable_tensor_chunk_elements(value) + for start in range(0, value.numel(), chunk_elements): + stop = min(start + chunk_elements, value.numel()) + source = ( + _read_portable_tensor_chunk(value, start, stop) + .resolve_conj() + .resolve_neg() + .reshape(-1) + ) + destination = cloned_flat[start:stop] + destination.copy_(source) + if cloned.is_floating_point() or cloned.is_complex(): + try: + finite = bool(torch.isfinite(destination).all()) + except (NotImplementedError, RuntimeError, TypeError) as exc: + raise ValueError( + "{} tensor dtype does not support finite portable state".format( + path + ) + ) from exc + if not finite: + raise ValueError("{} tensor must be finite".format(path)) + return cloned + if type(value) is float: + if not math.isfinite(value): + raise ValueError("{} must be finite".format(path)) + return value + if value is None or type(value) in {bool, int, str}: + return value + if type(value) is list: + return [ + _clone_portable_value(item, path="{}[{}]".format(path, index)) + for index, item in enumerate(value) + ] + if type(value) is tuple: + return tuple( + _clone_portable_value(item, path="{}[{}]".format(path, index)) + for index, item in enumerate(value) + ) + if type(value) is dict: + if any(type(key) is not str for key in value): + raise TypeError("{} dictionary keys must be strings".format(path)) + return { + key: _clone_portable_value( + value[key], path="{}.{}".format(path, key) + ) + for key in sorted(value) + } + raise TypeError( + "{} has unsupported portable-state type {}".format( + path, type(value).__name__ + ) + ) + + +def _digest_bytes(hasher, value: bytes) -> None: + hasher.update(struct.pack(">Q", len(value))) + hasher.update(value) + + +def _update_portable_digest(hasher, value) -> None: + value_type = type(value) + if value is None: + hasher.update(b"N") + return + if value_type is bool: + hasher.update(b"B1" if value else b"B0") + return + if value_type is int: + hasher.update(b"I") + _digest_bytes(hasher, str(value).encode("ascii")) + return + if value_type is float: + hasher.update(b"F") + hasher.update(struct.pack(">d", value)) + return + if value_type is str: + hasher.update(b"S") + _digest_bytes(hasher, value.encode("utf-8")) + return + if value_type is list: + hasher.update(b"L") + hasher.update(struct.pack(">Q", len(value))) + for item in value: + _update_portable_digest(hasher, item) + return + if value_type is tuple: + hasher.update(b"T") + hasher.update(struct.pack(">Q", len(value))) + for item in value: + _update_portable_digest(hasher, item) + return + if value_type is dict: + hasher.update(b"D") + hasher.update(struct.pack(">Q", len(value))) + for key in sorted(value): + if type(key) is not str: + raise TypeError("portable state dictionary keys must be strings") + _update_portable_digest(hasher, key) + _update_portable_digest(hasher, value[key]) + return + if value_type is torch.Tensor: + if value.device.type != "cpu" or not value.is_contiguous(): + raise ValueError("portable digest tensors must be contiguous CPU tensors") + hasher.update(b"R") + _digest_bytes(hasher, str(value.dtype).encode("ascii")) + _update_portable_digest(hasher, list(value.shape)) + element_size = value.element_size() + hasher.update(struct.pack(">Q", value.numel() * element_size)) + elements_per_chunk = max( + 1, _PORTABLE_DIGEST_CHUNK_BYTES // element_size + ) + flat = value.reshape(-1) + for start in range(0, value.numel(), elements_per_chunk): + stop = min(start + elements_per_chunk, value.numel()) + raw = flat[start:stop].view(torch.uint8).reshape(-1) + component_size = element_size // 2 if value.is_complex() else element_size + if _PORTABLE_NATIVE_BYTEORDER not in {"little", "big"}: + raise RuntimeError("unsupported native byte order") + if _PORTABLE_NATIVE_BYTEORDER == "big" and component_size > 1: + raw = ( + raw.reshape(-1, component_size) + .flip(1) + .contiguous() + .reshape(-1) + ) + hasher.update(memoryview(raw.numpy())) + return + raise TypeError( + "portable state digest does not support {}".format(value_type.__name__) + ) + + +def _canonical_portable_state_digest(canonical) -> str: + hasher = hashlib.sha256() + _update_portable_digest(hasher, canonical) + return hasher.hexdigest() + + +def portable_state_digest(payload) -> str: + """Return the deterministic SHA-256 digest of one canonical wire value.""" + + canonical = _clone_portable_value(payload, path="portable digest payload") + return _canonical_portable_state_digest(canonical) + + +def _normalize_portable_parameter(fqn, record): + if type(record) is not dict or set(record) != _PORTABLE_PARAMETER_KEYS: + raise ValueError( + "portable parameter {!r} has an invalid schema".format(fqn) + ) + identity_record = record["identity"] + if ( + type(identity_record) is not dict + or set(identity_record) != _PORTABLE_IDENTITY_KEYS + or type(identity_record["global_shape"]) is not list + ): + raise ValueError( + "portable parameter {!r} has an invalid identity".format(fqn) + ) + try: + identity = ParameterIdentity( + identity_record["fqn"], + tuple(identity_record["global_shape"]), + schema_version=identity_record["schema_version"], + ) + except (TypeError, ValueError) as exc: + raise ValueError( + "portable parameter {!r} has an invalid identity".format(fqn) + ) from exc + if identity.fqn != fqn: + raise ValueError("portable parameter identity does not match its FQN key") + for key in ("algorithm_options", "state", "projection_hints"): + if type(record[key]) is not dict: + raise ValueError( + "portable parameter {!r} {} must be a dictionary".format( + fqn, key + ) + ) + if type(record["state_variant"]) is not str or not record["state_variant"]: + raise ValueError( + "portable parameter {!r} state_variant must be a non-empty string".format( + fqn + ) + ) + return record + + +def _portable_payload(document): + return { + key: document[key] + for key in sorted(_PORTABLE_STATE_TOP_LEVEL_KEYS - {"completion"}) + } + + +def _validate_portable_payload(document) -> None: + if document["format"] != PORTABLE_STATE_FORMAT: + raise ValueError("unsupported portable state format") + if ( + type(document["format_version"]) is not int + or document["format_version"] != PORTABLE_STATE_FORMAT_VERSION + ): + raise ValueError( + "unsupported portable state format_version: {}".format( + document["format_version"] + ) + ) + if document["coverage"] != PORTABLE_STATE_COVERAGE: + raise ValueError("unsupported portable state coverage") + implementation = document["implementation"] + if type(implementation) is not str or not implementation: + raise ValueError("portable state implementation must be a non-empty string") + if type(document["policy"]) is not dict: + raise ValueError("portable state policy must be a dictionary") + if type(document["common"]) is not dict: + raise ValueError("portable common state must be a dictionary") + if document["provenance"] is not None and type(document["provenance"]) is not dict: + raise ValueError("portable state provenance must be a dictionary or None") + parameters = document["parameters"] + if type(parameters) is not dict: + raise ValueError("portable parameters must be an FQN mapping") + for fqn, record in parameters.items(): + if type(fqn) is not str: + raise ValueError("portable parameter FQN keys must be strings") + _normalize_portable_parameter(fqn, record) + + +def _validate_portable_completion(document) -> None: + completion = document["completion"] + if type(completion) is not dict or set(completion) != _PORTABLE_COMPLETION_KEYS: + raise ValueError("portable state completion marker has an invalid schema") + if completion["status"] != "complete": + raise ValueError("portable state is not marked complete") + if completion["digest_algorithm"] != PORTABLE_STATE_DIGEST_ALGORITHM: + raise ValueError("unsupported portable state digest algorithm") + digest = completion["digest"] + if ( + type(digest) is not str + or len(digest) != hashlib.sha256().digest_size * 2 + or any(character not in "0123456789abcdef" for character in digest) + ): + raise ValueError("portable state completion digest is invalid") + expected_digest = _canonical_portable_state_digest(_portable_payload(document)) + if not hmac.compare_digest(digest, expected_digest): + raise ValueError("portable state completion digest does not match its payload") + + +def normalize_portable_state_document(state, *, expected_implementation=None): + """Clone and validate a complete portable v3 state document.""" + + document = _clone_portable_value(state, path="portable state") + if type(document) is not dict or set(document) != _PORTABLE_STATE_TOP_LEVEL_KEYS: + raise ValueError("portable state has an invalid top-level schema") + _validate_portable_payload(document) + _validate_portable_completion(document) + if ( + expected_implementation is not None + and document["implementation"] != expected_implementation + ): + raise ValueError("portable state implementation does not match the target") + return document + + +def build_portable_state_document( + *, + implementation, + policy, + common, + parameters, + provenance=None, +): + """Build and validate a complete source-topology-neutral v3 document.""" + + document = { + "format": PORTABLE_STATE_FORMAT, + "format_version": PORTABLE_STATE_FORMAT_VERSION, + "coverage": PORTABLE_STATE_COVERAGE, + "implementation": implementation, + "policy": policy, + "common": common, + "parameters": parameters, + "provenance": provenance, + } + canonical = _clone_portable_value(document, path="portable state") + if type(canonical) is not dict or set(canonical) != ( + _PORTABLE_STATE_TOP_LEVEL_KEYS - {"completion"} + ): + raise ValueError("portable state has an invalid top-level schema") + _validate_portable_payload(canonical) + canonical["completion"] = { + "status": "complete", + "digest_algorithm": PORTABLE_STATE_DIGEST_ALGORITHM, + "digest": _canonical_portable_state_digest(canonical), + } + return canonical + + +__all__ = [ + "PORTABLE_STATE_COVERAGE", + "PORTABLE_STATE_DIGEST_ALGORITHM", + "PORTABLE_STATE_FORMAT", + "PORTABLE_STATE_FORMAT_VERSION", + "build_portable_state_document", + "normalize_portable_state_document", + "portable_state_digest", +] diff --git a/tests/test_checkpoint_binding.py b/tests/test_checkpoint_binding.py new file mode 100644 index 0000000..e1c0e49 --- /dev/null +++ b/tests/test_checkpoint_binding.py @@ -0,0 +1,227 @@ +from dataclasses import FrozenInstanceError, fields +from datetime import timedelta +import multiprocessing as mp +import os +import tempfile + +import pytest +import torch +import torch.distributed as dist + +from gefen.checkpoint import CheckpointProcessGroupBinding +from gefen.contracts import ProcessGroupIdentity + + +def test_checkpoint_binding_is_frozen_rank_neutral_and_deterministically_ordered(): + identity = ProcessGroupIdentity("pipeline:1/data", ("worker:b",)) + binding = CheckpointProcessGroupBinding(identity, "worker:b", None, torch.device("cpu")) + + assert binding.identity is identity + assert binding.local_member == "worker:b" + assert binding.process_group is None + assert binding.collective_device == torch.device("cpu") + assert binding.sort_key == ("pipeline:1/data", ("worker:b",)) + assert tuple(field.name for field in fields(binding)) == ( + "identity", + "local_member", + "process_group", + "collective_device", + ) + with pytest.raises(FrozenInstanceError): + binding.local_member = "worker:a" + + +def test_checkpoint_binding_validates_descriptor_types_membership_and_device(): + identity = ProcessGroupIdentity("local", ("worker:0",)) + + with pytest.raises(TypeError, match="ProcessGroupIdentity"): + CheckpointProcessGroupBinding(object(), "worker:0", None, torch.device("cpu")) + with pytest.raises(TypeError, match="local_member"): + CheckpointProcessGroupBinding(identity, 0, None, torch.device("cpu")) + with pytest.raises(ValueError, match="belong"): + CheckpointProcessGroupBinding(identity, "worker:1", None, torch.device("cpu")) + with pytest.raises(TypeError, match="torch.device"): + CheckpointProcessGroupBinding(identity, "worker:0", None, "cpu") + with pytest.raises(ValueError, match="CPU or CUDA"): + CheckpointProcessGroupBinding(identity, "worker:0", None, torch.device("meta")) + + +def test_checkpoint_binding_requires_explicit_multi_member_handle_and_local_singleton(): + singleton = ProcessGroupIdentity("local", ("worker:0",)) + multiple = ProcessGroupIdentity("data", ("worker:0", "worker:1")) + + with pytest.raises(ValueError, match="one-member"): + CheckpointProcessGroupBinding(singleton, "worker:0", object(), torch.device("cpu")) + with pytest.raises(ValueError, match="explicit"): + CheckpointProcessGroupBinding(multiple, "worker:0", None, torch.device("cpu")) + with pytest.raises(TypeError, match="ProcessGroup"): + CheckpointProcessGroupBinding(multiple, "worker:0", object(), torch.device("cpu")) + + +def _expect_rejection(callable_object, exception_type, message): + try: + callable_object() + except exception_type as exc: + return message in str(exc) + return False + + +def _distributed_binding_worker(rank, world_size, init_file, queue): + try: + dist.init_process_group( + "gloo", + init_method="file://{}".format(init_file), + rank=rank, + world_size=world_size, + timeout=timedelta(seconds=45), + ) + world_members = tuple("world:{}".format(index) for index in range(world_size)) + world_identity = ProcessGroupIdentity("world", world_members) + world_binding = CheckpointProcessGroupBinding( + world_identity, + world_members[rank], + dist.group.WORLD, + torch.device("cpu"), + ) + world_binding.validate_runtime() + world_validated = world_binding.process_group is dist.group.WORLD + + wrong_member_binding = CheckpointProcessGroupBinding( + world_identity, + world_members[(rank + 1) % world_size], + dist.group.WORLD, + torch.device("cpu"), + ) + order_mismatch_rejected = _expect_rejection( + wrong_member_binding.validate_runtime, + ValueError, + "ordered semantic members", + ) + + short_identity = ProcessGroupIdentity("short", ("short:0", "short:1")) + short_binding = CheckpointProcessGroupBinding( + short_identity, + short_identity.ordered_members[rank % 2], + dist.group.WORLD, + torch.device("cpu"), + ) + size_mismatch_rejected = _expect_rejection( + short_binding.validate_runtime, + ValueError, + "world size", + ) + + wrong_device_binding = CheckpointProcessGroupBinding( + world_identity, + world_members[rank], + dist.group.WORLD, + torch.device("cuda:0"), + ) + device_mismatch_rejected = _expect_rejection( + wrong_device_binding.validate_runtime, + ValueError, + "runtime backend", + ) + + subgroup_global_ranks = (0, 2) + subgroup = dist.new_group(ranks=list(subgroup_global_ranks), backend="gloo") + subgroup_members = ("subgroup:left", "subgroup:right") + subgroup_validated = False + nonmember_rejected = False + if rank in subgroup_global_ranks: + coordinate = subgroup_global_ranks.index(rank) + subgroup_binding = CheckpointProcessGroupBinding( + ProcessGroupIdentity("subgroup", subgroup_members), + subgroup_members[coordinate], + subgroup, + torch.device("cpu"), + ) + subgroup_binding.validate_runtime() + subgroup_validated = True + else: + nonmember_rejected = _expect_rejection( + lambda: CheckpointProcessGroupBinding( + ProcessGroupIdentity("subgroup", subgroup_members), + subgroup_members[0], + subgroup, + torch.device("cpu"), + ), + TypeError, + "ProcessGroup", + ) + + dist.destroy_process_group() + uninitialized_rejected = _expect_rejection( + world_binding.validate_runtime, + RuntimeError, + "initialized torch.distributed", + ) + queue.put( + { + "rank": rank, + "world_validated": world_validated, + "order_mismatch_rejected": order_mismatch_rejected, + "size_mismatch_rejected": size_mismatch_rejected, + "device_mismatch_rejected": device_mismatch_rejected, + "subgroup_validated": subgroup_validated, + "nonmember_rejected": nonmember_rejected, + "uninitialized_rejected": uninitialized_rejected, + } + ) + except Exception as exc: + queue.put({"rank": rank, "error": repr(exc)}) + finally: + if dist.is_available() and dist.is_initialized(): + dist.destroy_process_group() + + +def _run_distributed_binding_workers(): + context = mp.get_context("spawn") + queue = context.Queue() + fd, init_file = tempfile.mkstemp(prefix="gefen-checkpoint-binding-") + os.close(fd) + os.unlink(init_file) + processes = [ + context.Process( + target=_distributed_binding_worker, + args=(rank, 3, init_file, queue), + ) + for rank in range(3) + ] + try: + for process in processes: + process.start() + results = [queue.get(timeout=60) for _ in processes] + for process in processes: + process.join(timeout=10) + for process in processes: + if process.is_alive(): + process.terminate() + process.join(timeout=5) + pytest.fail("checkpoint-binding worker hung") + assert process.exitcode == 0 + return sorted(results, key=lambda item: item["rank"]) + finally: + for process in processes: + if process.is_alive(): + process.terminate() + process.join(timeout=5) + if os.path.exists(init_file): + os.unlink(init_file) + + +@pytest.mark.skipif( + not dist.is_available() or not dist.is_gloo_available(), + reason="Gloo is required for process-group binding validation", +) +def test_checkpoint_binding_validates_real_world_and_subgroup_membership(): + results = _run_distributed_binding_workers() + + assert all("error" not in item for item in results), results + assert all(item["world_validated"] for item in results) + assert all(item["order_mismatch_rejected"] for item in results) + assert all(item["size_mismatch_rejected"] for item in results) + assert all(item["device_mismatch_rejected"] for item in results) + assert all(item["subgroup_validated"] for item in results if item["rank"] in (0, 2)) + assert results[1]["nonmember_rejected"] + assert all(item["uninitialized_rejected"] for item in results) diff --git a/tests/test_optimizer_contracts.py b/tests/test_optimizer_contracts.py index 9fd35ed..e7348e9 100644 --- a/tests/test_optimizer_contracts.py +++ b/tests/test_optimizer_contracts.py @@ -747,3 +747,13 @@ def test_all_public_contract_exports_resolve(): assert all(getattr(gefen, name) is not None for name in contracts.__all__) assert gefen.StateMovementProvider is StateMovementProvider + + +def test_portable_global_transport_is_defined_but_not_claimed_before_integration(): + parameter = torch.nn.Parameter(torch.ones(4)) + optimizer = Gefen([("parameter", parameter)], fused=False) + + assert all( + support.transport is not CheckpointTransport.CANONICAL_GLOBAL + for support in optimizer.optimizer_contract().capabilities.checkpoints + ) diff --git a/tests/test_portable_schema.py b/tests/test_portable_schema.py new file mode 100644 index 0000000..6dc3115 --- /dev/null +++ b/tests/test_portable_schema.py @@ -0,0 +1,285 @@ +"""Strict portable-state v3 envelope coverage.""" + +import copy +import io + +import pytest +import torch + +import gefen +import gefen.portable_schema as portable_schema_module +from gefen.portable_schema import ( + PORTABLE_STATE_COVERAGE, + PORTABLE_STATE_DIGEST_ALGORITHM, + PORTABLE_STATE_FORMAT, + PORTABLE_STATE_FORMAT_VERSION, + build_portable_state_document, + normalize_portable_state_document, + portable_state_digest, +) + + +def _parameter_record(*, tensor=None): + if tensor is None: + tensor = torch.arange(6, dtype=torch.float32).reshape(2, 3) + return { + "identity": { + "schema_version": 1, + "fqn": "Model.Weight", + "global_shape": [2, 3], + }, + "algorithm_options": {"lr": 1.0e-3, "betas": (0.9, 0.999)}, + "state_variant": "quantized_momentum", + "state": {"step": 3, "momentum": tensor}, + "projection_hints": {"source_periods": [3]}, + } + + +def _document(*, tensor=None): + return build_portable_state_document( + implementation="gefen.Gefen", + policy={"factored_v_2d": False}, + common={ + "gefen_global_step": 3, + "gefen_deterministic": False, + }, + parameters={"Model.Weight": _parameter_record(tensor=tensor)}, + provenance={"source_layouts": ["replicated"]}, + ) + + +def test_portable_document_is_complete_device_neutral_and_weights_only_safe(): + backing = torch.arange(30, dtype=torch.float32) + view = backing[5:17:2].reshape(2, 3) + document = _document(tensor=view) + + assert document["format"] == PORTABLE_STATE_FORMAT + assert document["format_version"] == PORTABLE_STATE_FORMAT_VERSION + assert document["coverage"] == PORTABLE_STATE_COVERAGE + assert document["completion"]["status"] == "complete" + assert ( + document["completion"]["digest_algorithm"] + == PORTABLE_STATE_DIGEST_ALGORITHM + ) + momentum = document["parameters"]["Model.Weight"]["state"]["momentum"] + assert momentum.device.type == "cpu" + assert momentum.is_contiguous() + assert momentum.storage_offset() == 0 + assert momentum.untyped_storage().nbytes() == momentum.numel() * momentum.element_size() + assert torch.equal(momentum, view) + assert momentum is not view + + buffer = io.BytesIO() + torch.save(document, buffer) + buffer.seek(0) + loaded = torch.load(buffer, weights_only=True) + normalized = normalize_portable_state_document(loaded) + assert normalized["completion"] == document["completion"] + + +def test_portable_digest_is_deterministic_and_type_shape_dtype_sensitive(): + left = {"b": [1, 2], "a": torch.tensor([1.0, 2.0])} + right = {"a": torch.tensor([1.0, 2.0]), "b": [1, 2]} + baseline = portable_state_digest(left) + + assert baseline == portable_state_digest(right) + assert baseline != portable_state_digest({"a": torch.tensor([[1.0, 2.0]]), "b": [1, 2]}) + assert baseline != portable_state_digest({"a": torch.tensor([1.0, 2.0], dtype=torch.float64), "b": [1, 2]}) + assert baseline != portable_state_digest({"a": torch.tensor([1.0, 3.0]), "b": [1, 2]}) + assert baseline != portable_state_digest({"a": torch.tensor([1.0, 2.0]), "b": (1, 2)}) + + +def test_portable_digest_streams_tensor_bytes_without_changing_the_digest(monkeypatch): + payload = {"tensor": torch.arange(257, dtype=torch.float32)} + baseline = portable_state_digest(payload) + + monkeypatch.setattr(portable_schema_module, "_PORTABLE_DIGEST_CHUNK_BYTES", 17) + + assert portable_state_digest(payload) == baseline + + +def test_portable_digest_v3_grammar_has_a_cross_version_golden_vector(): + payload = { + "a": None, + "b": True, + "c": -12345678901234567890, + "d": -0.0, + "e": "x\u2603", + "f": [1, (2, 3)], + "g": torch.tensor([[1, -2], [3, -4]], dtype=torch.int16), + "h": torch.tensor([1.5, -2.25], dtype=torch.bfloat16), + "i": torch.tensor(3 + 4j, dtype=torch.complex64), + "j": torch.empty((0, 2), dtype=torch.float64), + } + + assert portable_state_digest(payload) == ( + "348e09a77f1b3eae286adda1573722df9" + "38be6d8100631b272665f8c22c6e23a" + ) + + +def test_portable_clone_streams_noncontiguous_values_and_finite_checks(monkeypatch): + source = torch.arange(514, dtype=torch.float32)[1::2] + assert not source.is_contiguous() + calls = [] + original = portable_schema_module._read_portable_tensor_chunk + + def tracked(value, start, stop): + calls.append((start, stop)) + return original(value, start, stop) + + monkeypatch.setattr(portable_schema_module, "_PORTABLE_CLONE_CHUNK_BYTES", 17) + monkeypatch.setattr(portable_schema_module, "_read_portable_tensor_chunk", tracked) + + cloned = portable_schema_module._clone_portable_value(source, path="tensor") + + assert torch.equal(cloned, source) + assert cloned.is_contiguous() + assert cloned.storage_offset() == 0 + assert cloned.untyped_storage().nbytes() == cloned.numel() * cloned.element_size() + assert len(calls) > 1 + assert max(stop - start for start, stop in calls) <= 4 + + +def test_builder_and_normalizer_each_hash_the_tensor_tree_once(monkeypatch): + calls = [] + original = portable_schema_module._canonical_portable_state_digest + + def tracked(value): + calls.append(value) + return original(value) + + monkeypatch.setattr(portable_schema_module, "_canonical_portable_state_digest", tracked) + + document = _document() + assert len(calls) == 1 + normalize_portable_state_document(document) + assert len(calls) == 2 + + +def test_multibyte_tensor_digest_is_canonical_little_endian(monkeypatch): + little_endian_values = torch.tensor([0x0102, 0x0304], dtype=torch.int16) + simulated_big_endian_storage = torch.tensor([0x0201, 0x0403], dtype=torch.int16) + baseline = portable_state_digest(little_endian_values) + + monkeypatch.setattr(portable_schema_module, "_PORTABLE_NATIVE_BYTEORDER", "big") + + assert portable_state_digest(simulated_big_endian_storage) == baseline + + +@pytest.mark.parametrize( + "corruption", + ( + "payload", + "digest", + "status", + "algorithm", + "missing_top", + "bool_version", + "coverage", + "implementation", + "fqn", + "shape_type", + "state_variant", + "state_type", + "provenance", + ), +) +def test_portable_schema_and_completion_corruption_are_rejected(corruption): + document = _document() + damaged = copy.deepcopy(document) + if corruption == "payload": + damaged["common"]["gefen_global_step"] = 4 + elif corruption == "digest": + damaged["completion"]["digest"] = "0" * 64 + elif corruption == "status": + damaged["completion"]["status"] = "preparing" + elif corruption == "algorithm": + damaged["completion"]["digest_algorithm"] = "md5" + elif corruption == "missing_top": + damaged.pop("policy") + elif corruption == "bool_version": + damaged["format_version"] = True + elif corruption == "coverage": + damaged["coverage"] = "local_optimizer_fragment" + elif corruption == "implementation": + damaged["implementation"] = "" + elif corruption == "fqn": + damaged["parameters"]["Model.Weight"]["identity"]["fqn"] = "Other.Weight" + elif corruption == "shape_type": + damaged["parameters"]["Model.Weight"]["identity"]["global_shape"] = (2, 3) + elif corruption == "state_variant": + damaged["parameters"]["Model.Weight"]["state_variant"] = "" + elif corruption == "state_type": + damaged["parameters"]["Model.Weight"]["state"] = [] + else: + damaged["provenance"] = [] + + with pytest.raises((TypeError, ValueError)): + normalize_portable_state_document(damaged) + + +def test_expected_implementation_is_checked_after_digest_validation(): + document = _document() + + with pytest.raises(ValueError, match="does not match the target"): + normalize_portable_state_document( + document, expected_implementation="gefen.GefenMuon" + ) + + document["completion"]["digest"] = "0" * 64 + with pytest.raises(ValueError, match="digest does not match"): + normalize_portable_state_document( + document, expected_implementation="gefen.GefenMuon" + ) + + +@pytest.mark.parametrize( + "bad_value", + ( + torch.tensor([float("nan")]), + torch.tensor([float("inf")]), + lambda: None, + ), +) +def test_builder_rejects_nonportable_values(bad_value): + record = _parameter_record() + record["state"]["bad"] = bad_value + + with pytest.raises((TypeError, ValueError)): + build_portable_state_document( + implementation="gefen.Gefen", + policy={}, + common={}, + parameters={"Model.Weight": record}, + ) + + +def test_builder_and_normalizer_do_not_alias_or_mutate_inputs(): + record = _parameter_record() + parameters = {"Model.Weight": record} + source_momentum = record["state"]["momentum"] + + document = build_portable_state_document( + implementation="gefen.Gefen", + policy={}, + common={}, + parameters=parameters, + ) + normalized = normalize_portable_state_document(document) + + assert record["state"]["momentum"] is source_momentum + assert document is not normalized + assert document["parameters"] is not parameters + assert document["parameters"]["Model.Weight"] is not record + assert document["parameters"]["Model.Weight"]["state"]["momentum"] is not source_momentum + assert normalized["parameters"]["Model.Weight"]["state"]["momentum"] is not document["parameters"]["Model.Weight"]["state"]["momentum"] + + +def test_portable_schema_and_checkpoint_binding_exports_are_public(): + assert gefen.PORTABLE_STATE_FORMAT_VERSION == PORTABLE_STATE_FORMAT_VERSION + assert gefen.build_portable_state_document is build_portable_state_document + assert gefen.normalize_portable_state_document is normalize_portable_state_document + assert gefen.portable_state_digest is portable_state_digest + assert gefen.CheckpointProcessGroupBinding.__module__ == "gefen.checkpoint" + assert gefen.LogicalRegion.__module__ == "gefen.contracts" diff --git a/tests/test_portable_state_math.py b/tests/test_portable_state_math.py new file mode 100644 index 0000000..b7c83d0 --- /dev/null +++ b/tests/test_portable_state_math.py @@ -0,0 +1,728 @@ +import warnings + +import pytest +import torch + +import gefen.portable as portable_math +from gefen.portable import ( + _decode_quantized_momentum, + _expand_block_second_moment, + _expand_factored_second_moment, + _project_factored_second_moment, + _recompress_dense_momentum, + _reduce_block_second_moment, + _validate_state_counter, +) + + +def _assert_tight(tensor, *, shape, dtype=torch.float32): + assert type(tensor) is torch.Tensor + assert tuple(tensor.shape) == tuple(shape) + assert tensor.dtype == dtype + assert tensor.layout is torch.strided + assert not tensor.requires_grad + assert tensor.grad_fn is None + assert tensor.is_contiguous() + assert tensor.storage_offset() == 0 + assert tensor.untyped_storage().nbytes() == tensor.numel() * tensor.element_size() + assert bool(torch.isfinite(tensor).all()) + + +def _codebook(): + return torch.tensor([-1.0, -0.5, 0.5, 1.0], dtype=torch.float32) + + +def _sparse_tensor(): + with warnings.catch_warnings(): + warnings.simplefilter("ignore", UserWarning) + return torch.sparse_coo_tensor( + torch.tensor([[0]]), + torch.tensor([1.0]), + (4,), + ) + + +def _reference_nearest_indices(codebook, values): + codebook = codebook.contiguous() + flat = values.reshape(-1).float() + insertion = torch.searchsorted(codebook, flat) + left = (insertion - 1).clamp(0, codebook.numel() - 1) + right = insertion.clamp(0, codebook.numel() - 1) + return torch.where( + (flat - codebook[left]).abs() <= (flat - codebook[right]).abs(), + left, + right, + ).to(torch.uint8).reshape(values.shape) + + +def _reference_recompress(momentum, codebook, period): + blocks = momentum.reshape(-1, period) + magnitudes = blocks.abs().amax(dim=1, keepdim=True) + normalized = blocks.clone() + nonzero = magnitudes > 0 + normalized.div_(magnitudes) + normalized.masked_fill_(~nonzero, 0.0) + return _reference_nearest_indices(codebook, normalized), magnitudes + + +def test_decode_quantized_momentum_validates_geometry_and_returns_dense_tight_fp32(): + codebook = torch.tensor([-1.0, -0.25, 0.25, 1.0], dtype=torch.float32) + indices = torch.tensor([[0, 2], [3, 1]], dtype=torch.uint8) + magnitudes = torch.tensor([[2.0], [4.0]], dtype=torch.float32) + codebook_before = codebook.clone() + indices_before = indices.clone() + magnitudes_before = magnitudes.clone() + + dense = _decode_quantized_momentum( + codebook, + indices, + magnitudes, + logical_shape=(2, 2), + period=2, + step=7, + ) + + _assert_tight(dense, shape=(2, 2)) + assert torch.equal(dense, torch.tensor([[-2.0, 0.5], [4.0, -1.0]])) + assert torch.equal(codebook, codebook_before) + assert torch.equal(indices, indices_before) + assert torch.equal(magnitudes, magnitudes_before) + assert dense.untyped_storage().data_ptr() != magnitudes.untyped_storage().data_ptr() + + +def test_recompression_uses_exact_lower_index_tie_semantics(): + momentum = torch.tensor([[-2.0, -1.5], [2.0, 0.0]], requires_grad=True) + indices, magnitudes = _recompress_dense_momentum( + momentum, + _codebook(), + period=2, + step=1, + ) + + _assert_tight(indices, shape=(2, 2), dtype=torch.uint8) + _assert_tight(magnitudes, shape=(2, 1)) + assert torch.equal(magnitudes, torch.tensor([[2.0], [2.0]])) + # -0.75 ties indices 0/1 and 0.0 ties indices 1/2. Both choose the lower index. + assert torch.equal(indices, torch.tensor([[0, 0], [3, 1]], dtype=torch.uint8)) + assert torch.equal(momentum.detach(), torch.tensor([[-2.0, -1.5], [2.0, 0.0]])) + + +def test_period_one_recompression_is_exact_for_every_finite_fp32_scale(): + tiny = torch.nextafter(torch.tensor(0.0), torch.tensor(1.0)) + maximum = torch.tensor(torch.finfo(torch.float32).max) + momentum = torch.stack( + ( + -maximum, + torch.tensor(-123.75), + -tiny, + torch.tensor(-0.0), + torch.tensor(0.0), + tiny, + torch.tensor(19.125), + maximum, + ) + ) + codebook = torch.tensor( + [-1.0, -1.0, -0.2, 0.3, 1.0, 1.0], + dtype=torch.float32, + ) + + indices, magnitudes = _recompress_dense_momentum( + momentum, + codebook, + period=1, + step=99, + ) + reconstructed = _decode_quantized_momentum( + codebook, + indices, + magnitudes, + logical_shape=momentum.shape, + period=1, + step=99, + ) + + torch.testing.assert_close(reconstructed, momentum, rtol=0, atol=0) + assert torch.equal(magnitudes.reshape(-1), momentum.abs()) + assert torch.equal(indices[momentum < 0], torch.zeros_like(indices[momentum < 0])) + assert torch.equal(indices[momentum > 0], torch.full_like(indices[momentum > 0], 4)) + + +def test_scalar_momentum_roundtrips_as_zero_dimensional_logical_state(): + momentum = torch.tensor(-123.75, requires_grad=True) + indices, magnitudes = _recompress_dense_momentum( + momentum, + _codebook(), + period=1, + step=2, + ) + decoded = _decode_quantized_momentum( + _codebook(), + indices, + magnitudes, + logical_shape=(), + period=1, + step=2, + ) + + _assert_tight(indices, shape=(1, 1), dtype=torch.uint8) + _assert_tight(magnitudes, shape=(1, 1)) + _assert_tight(decoded, shape=()) + assert decoded.ndim == 0 + assert torch.equal(decoded, momentum.detach()) + + +def test_expanded_zero_stride_inputs_preserve_logical_row_major_values(): + momentum = torch.tensor([-2.0, 0.0, 3.0]).reshape(3, 1).expand(3, 4) + assert momentum.stride() == (1, 0) + codebook = torch.tensor([-1.0, 0.0, 1.0]) + indices, magnitudes = _recompress_dense_momentum( + momentum, + codebook, + period=4, + step=1, + ) + decoded = _decode_quantized_momentum( + codebook, + indices, + magnitudes, + logical_shape=momentum.shape, + period=4, + step=1, + ) + assert torch.equal(decoded, momentum) + + dense_second_moment = torch.tensor([1.0, 3.0, 5.0]).reshape(3, 1).expand(3, 4) + reduced = _reduce_block_second_moment( + dense_second_moment, + period=4, + step=1, + ) + assert torch.equal(reduced, torch.tensor([[1.0], [3.0], [5.0]])) + expanded_blocks = _expand_block_second_moment( + torch.tensor([[2.0]]).expand(3, 1), + logical_shape=(3, 4), + period=4, + step=1, + ) + assert torch.equal(expanded_blocks, torch.full((3, 4), 2.0)) + + factored = _expand_factored_second_moment( + torch.tensor([2.0]).expand(3), + torch.tensor([4.0]).expand(5), + logical_shape=(3, 5), + step=1, + ) + assert torch.equal(factored, torch.full((3, 5), 4.0)) + projected_row, projected_column = _project_factored_second_moment( + torch.tensor([7.0]).expand(4, 6), + step=1, + ) + assert torch.equal(projected_row, torch.full((4,), 7.0)) + assert torch.equal(projected_column, torch.full((6,), 7.0)) + + +def test_noncontiguous_codebook_is_normalized_without_searchsorted_warning(): + codebook = torch.linspace(-1.0, 1.0, 17)[::2] + assert not codebook.is_contiguous() + values = torch.tensor([-1.0, -0.875, -0.75, 0.0, 0.875, 1.0]) + assert torch.equal( + portable_math._nearest_codebook_indices(codebook, values), + _reference_nearest_indices(codebook, values), + ) + indices, magnitudes = _recompress_dense_momentum( + values, + codebook, + period=2, + step=1, + ) + assert torch.equal( + _decode_quantized_momentum( + codebook, + indices, + magnitudes, + logical_shape=values.shape, + period=2, + step=1, + ), + (codebook[indices.long()] * magnitudes).reshape(values.shape), + ) + + +@pytest.mark.parametrize("period", [1, 2, 4, 8]) +def test_recompression_supports_every_divisor_as_a_target_period(period): + momentum = torch.tensor( + [[-8.0, -5.0, -3.0, -1.0], [0.0, 1.0, 2.0, 8.0]], + dtype=torch.float32, + ).t() + assert not momentum.is_contiguous() + codebook = torch.linspace(-1.0, 1.0, 17) + + indices, magnitudes = _recompress_dense_momentum( + momentum, + codebook, + period=period, + step=2, + ) + blocks = momentum.numel() // period + expected_magnitudes = momentum.reshape(blocks, period).abs().amax(dim=1, keepdim=True) + assert torch.equal(magnitudes, expected_magnitudes) + _assert_tight(indices, shape=(blocks, period), dtype=torch.uint8) + _assert_tight(magnitudes, shape=(blocks, 1)) + + +def test_tiny_chunk_budget_matches_whole_tensor_references_across_partial_tails(monkeypatch): + monkeypatch.setattr(portable_math, "_PORTABLE_STATE_CHUNK_ELEMENTS", 7) + codebook = torch.tensor([-1.0, -0.5, 0.0, 0.5, 1.0]) + normalized = torch.tensor( + [ + [-1.0, -0.75, -0.5, -0.25, 0.0, 0.25], + [0.5, 0.75, 1.0, -0.75, 0.25, -0.25], + [0.1, -0.1, 0.9, -0.9, 0.6, -0.6], + [0.2, -0.2, 0.8, -0.8, 0.4, -0.4], + [0.3, -0.3, 0.7, -0.7, 1.0, -1.0], + ] + ).t() + assert not normalized.is_contiguous() + expected_nearest = _reference_nearest_indices(codebook, normalized) + actual_nearest = portable_math._nearest_codebook_indices(codebook, normalized) + assert torch.equal(actual_nearest, expected_nearest) + + momentum = torch.linspace(-9.0, 7.0, 30).reshape(6, 5).t() + assert not momentum.is_contiguous() + expected_indices, expected_magnitudes = _reference_recompress( + momentum, + codebook, + 6, + ) + actual_indices, actual_magnitudes = _recompress_dense_momentum( + momentum, + codebook, + period=6, + step=3, + ) + assert torch.equal(actual_indices, expected_indices) + assert torch.equal(actual_magnitudes, expected_magnitudes) + decoded = _decode_quantized_momentum( + codebook, + actual_indices, + actual_magnitudes, + logical_shape=momentum.shape, + period=6, + step=3, + ) + expected_decoded = ( + codebook[expected_indices.long()] * expected_magnitudes + ).reshape(momentum.shape) + assert torch.equal(decoded, expected_decoded) + + dense_second_moment = torch.arange(1, 31, dtype=torch.float32).reshape(6, 5).t() + expected_blocks = dense_second_moment.reshape(-1, 6).to(torch.float64).mean( + dim=1, + keepdim=True, + ).to(torch.float32) + actual_blocks = _reduce_block_second_moment( + dense_second_moment, + period=6, + step=3, + ) + assert torch.equal(actual_blocks, expected_blocks) + expanded_blocks = _expand_block_second_moment( + actual_blocks, + logical_shape=dense_second_moment.shape, + period=6, + step=3, + ) + expected_expanded_blocks = torch.repeat_interleave( + expected_blocks.reshape(-1), + 6, + ).reshape(dense_second_moment.shape) + assert torch.equal(expanded_blocks, expected_expanded_blocks) + + row = torch.tensor([1.0, 2.0, 4.0, 5.0, 8.0]) + column = torch.tensor([0.5, 1.0, 2.0, 3.0, 4.0, 5.5]) + expected_factored = ( + torch.outer(row.to(torch.float64), column.to(torch.float64)) + / row.to(torch.float64).mean() + ).to(torch.float32) + actual_factored = _expand_factored_second_moment( + row, + column, + logical_shape=(5, 6), + step=3, + ) + assert torch.equal(actual_factored, expected_factored) + actual_row, actual_column = _project_factored_second_moment( + actual_factored.t(), + step=3, + ) + expected64 = actual_factored.t().to(torch.float64) + assert torch.equal(actual_row, expected64.mean(dim=1).to(torch.float32)) + assert torch.equal(actual_column, expected64.mean(dim=0).to(torch.float32)) + + +def test_period_larger_than_chunk_budget_keeps_whole_blocks_and_matches_reference(monkeypatch): + monkeypatch.setattr(portable_math, "_PORTABLE_STATE_CHUNK_ELEMENTS", 5) + period = 11 + codebook = torch.linspace(-1.0, 1.0, 9) + momentum = torch.linspace(-17.0, 13.0, 33).reshape(11, 3).t() + assert not momentum.is_contiguous() + + expected_indices, expected_magnitudes = _reference_recompress( + momentum, + codebook, + period, + ) + actual_indices, actual_magnitudes = _recompress_dense_momentum( + momentum, + codebook, + period=period, + step=2, + ) + decoded = _decode_quantized_momentum( + codebook, + actual_indices, + actual_magnitudes, + logical_shape=momentum.shape, + period=period, + step=2, + ) + + assert torch.equal(actual_indices, expected_indices) + assert torch.equal(actual_magnitudes, expected_magnitudes) + assert torch.equal( + decoded, + (codebook[expected_indices.long()] * expected_magnitudes).reshape(momentum.shape), + ) + second_moment = momentum.square() + actual_blocks = _reduce_block_second_moment( + second_moment, + period=period, + step=2, + ) + expected_blocks = second_moment.reshape(-1, period).to(torch.float64).mean( + dim=1, + keepdim=True, + ).to(torch.float32) + assert torch.equal(actual_blocks, expected_blocks) + + +def test_block_second_moment_expands_then_projects_to_a_new_period(): + blocks = torch.tensor([[1.0], [3.0], [5.0], [7.0]], requires_grad=True) + dense = _expand_block_second_moment( + blocks, + logical_shape=(2, 4), + period=2, + step=3, + ) + projected = _reduce_block_second_moment(dense, period=4, step=3) + + _assert_tight(dense, shape=(2, 4)) + _assert_tight(projected, shape=(2, 1)) + assert torch.equal( + dense, + torch.tensor([[1.0, 1.0, 3.0, 3.0], [5.0, 5.0, 7.0, 7.0]]), + ) + assert torch.equal(projected, torch.tensor([[2.0], [6.0]])) + assert torch.equal(blocks.detach(), torch.tensor([[1.0], [3.0], [5.0], [7.0]])) + + +def test_block_reduction_uses_overflow_safe_means(): + maximum = torch.finfo(torch.float32).max + dense = torch.full((2, 4), maximum) + + projected = _reduce_block_second_moment(dense, period=4, step=1) + + assert torch.equal(projected, torch.full((2, 1), maximum)) + _assert_tight(projected, shape=(2, 1)) + + +def test_factored_expansion_and_projection_follow_adafactor_geometry(): + row = torch.tensor([1.0, 3.0]) + column = torch.tensor([2.0, 4.0, 6.0]) + dense = _expand_factored_second_moment( + row, + column, + logical_shape=(2, 3), + step=4, + ) + + _assert_tight(dense, shape=(2, 3)) + assert torch.equal(dense, torch.tensor([[1.0, 2.0, 3.0], [3.0, 6.0, 9.0]])) + + source = torch.tensor([[1.0, 3.0, 5.0], [2.0, 4.0, 6.0]]).t() + assert not source.is_contiguous() + projected_row, projected_column = _project_factored_second_moment(source, step=4) + assert torch.equal(projected_row, torch.tensor([1.5, 3.5, 5.5])) + assert torch.equal(projected_column, torch.tensor([3.0, 4.0])) + _assert_tight(projected_row, shape=(3,)) + _assert_tight(projected_column, shape=(2,)) + + +def test_factored_projection_is_overflow_safe_and_roundtrips_consistent_rank_one_state(): + maximum = torch.finfo(torch.float32).max + projected_row, projected_column = _project_factored_second_moment( + torch.full((2, 2), maximum), + step=1, + ) + assert torch.equal(projected_row, torch.full((2,), maximum)) + assert torch.equal(projected_column, torch.full((2,), maximum)) + + row = torch.tensor([1.0, 3.0]) + column = torch.tensor([1.0, 2.0, 3.0]) + dense = _expand_factored_second_moment( + row, + column, + logical_shape=(2, 3), + step=2, + ) + roundtrip_row, roundtrip_column = _project_factored_second_moment(dense, step=2) + torch.testing.assert_close(roundtrip_row, row, rtol=0, atol=0) + torch.testing.assert_close(roundtrip_column, column, rtol=0, atol=0) + + +@pytest.mark.parametrize("shape", [(0,), (0, 3), (2, 0)]) +def test_zero_element_momentum_and_block_state_are_well_defined(shape): + momentum = torch.empty(shape, dtype=torch.float32, requires_grad=True) + codebook = _codebook() + indices, magnitudes = _recompress_dense_momentum( + momentum, + codebook, + period=7, + step=1, + ) + decoded = _decode_quantized_momentum( + codebook, + indices, + magnitudes, + logical_shape=shape, + period=7, + step=1, + ) + reduced = _reduce_block_second_moment(momentum.detach(), period=7, step=1) + expanded = _expand_block_second_moment( + reduced, + logical_shape=shape, + period=7, + step=1, + ) + + _assert_tight(indices, shape=(0, 7), dtype=torch.uint8) + _assert_tight(magnitudes, shape=(0, 1)) + _assert_tight(decoded, shape=shape) + _assert_tight(reduced, shape=(0, 1)) + _assert_tight(expanded, shape=shape) + + +def test_empty_index_decode_does_not_invoke_max(monkeypatch): + original_max = torch.Tensor.max + + def reject_empty_max(tensor, *args, **kwargs): + if tensor.numel() == 0: + raise AssertionError("max() was called on empty momentum indices") + return original_max(tensor, *args, **kwargs) + + monkeypatch.setattr(torch.Tensor, "max", reject_empty_max) + decoded = _decode_quantized_momentum( + _codebook(), + torch.empty((0, 13), dtype=torch.uint8), + torch.empty((0, 1), dtype=torch.float32), + logical_shape=(0, 9), + period=13, + step=1, + ) + _assert_tight(decoded, shape=(0, 9)) + + +@pytest.mark.parametrize("shape", [(0, 3), (3, 0), (0, 0)]) +def test_zero_element_factored_state_uses_finite_empty_geometry(shape): + dense = torch.empty(shape, dtype=torch.float32) + row, column = _project_factored_second_moment(dense, step=1) + expanded = _expand_factored_second_moment( + row, + column, + logical_shape=shape, + step=1, + ) + + _assert_tight(row, shape=(shape[0],)) + _assert_tight(column, shape=(shape[1],)) + _assert_tight(expanded, shape=shape) + assert torch.count_nonzero(row) == 0 + assert torch.count_nonzero(column) == 0 + + +@pytest.mark.parametrize("value", [True, 1.0, torch.tensor(1), None]) +def test_counter_validation_rejects_non_host_ints(value): + with pytest.raises(TypeError, match="host int"): + _validate_state_counter(value, name="step", minimum=1) + + +def test_counter_validation_is_strict_about_bounds_and_its_own_arguments(): + assert _validate_state_counter(0, name="global_step") == 0 + assert _validate_state_counter(3, name="step", minimum=1) == 3 + with pytest.raises(ValueError, match="at least 1"): + _validate_state_counter(0, name="step", minimum=1) + with pytest.raises(TypeError, match="name"): + _validate_state_counter(1, name=object()) + with pytest.raises(ValueError, match="must not be empty"): + _validate_state_counter(1, name="") + with pytest.raises(TypeError, match="minimum"): + _validate_state_counter(1, name="step", minimum=True) + with pytest.raises(ValueError, match="nonnegative"): + _validate_state_counter(1, name="step", minimum=-1) + + +@pytest.mark.parametrize("period", [True, 2.0, 0, -1, 3]) +def test_invalid_periods_fail_before_projection(period): + exception = TypeError if type(period) is not int else ValueError + with pytest.raises(exception, match="period"): + _recompress_dense_momentum( + torch.ones(4), + _codebook(), + period=period, + step=1, + ) + + +@pytest.mark.parametrize( + "shape,exception", + [([4], TypeError), ((True,), TypeError), ((-1, 4), ValueError), ((3,), ValueError)], +) +def test_invalid_logical_shapes_and_geometry_are_rejected(shape, exception): + with pytest.raises(exception): + _decode_quantized_momentum( + _codebook(), + torch.zeros((2, 2), dtype=torch.uint8), + torch.ones((2, 1)), + logical_shape=shape, + period=2, + step=1, + ) + + +@pytest.mark.parametrize( + "codebook,match", + [ + (torch.tensor([], dtype=torch.float32), "between 1 and 256"), + (torch.zeros(257, dtype=torch.float32), "between 1 and 256"), + (torch.tensor([-1.0, 0.5, 0.25, 1.0]), "sorted"), + (torch.tensor([-1.1, 0.0, 1.0]), r"\[-1, 1\]"), + (torch.tensor([-1.0, float("nan"), 1.0]), "finite"), + ], +) +def test_invalid_codebook_values_are_rejected(codebook, match): + with pytest.raises(ValueError, match=match): + _recompress_dense_momentum( + torch.ones(4), + codebook, + period=2, + step=1, + ) + + +@pytest.mark.parametrize( + "codebook", + [torch.ones(4, dtype=torch.float64), torch.ones((2, 2), dtype=torch.float32)], +) +def test_invalid_codebook_dtype_or_rank_is_rejected(codebook): + with pytest.raises((TypeError, ValueError), match="codebook"): + _recompress_dense_momentum( + torch.ones(4), + codebook, + period=2, + step=1, + ) + + +@pytest.mark.parametrize( + "momentum", + [ + torch.ones(4, dtype=torch.float64), + torch.tensor([1.0, float("inf"), 2.0, 3.0]), + torch.nn.Parameter(torch.ones(4)), + _sparse_tensor(), + torch.empty(4, device="meta"), + ], +) +def test_momentum_input_must_be_plain_strided_finite_fp32(momentum): + with pytest.raises((TypeError, ValueError), match="momentum"): + _recompress_dense_momentum( + momentum, + _codebook(), + period=2, + step=1, + ) + + +@pytest.mark.parametrize( + "indices,magnitudes,match", + [ + (torch.zeros((2, 2), dtype=torch.int64), torch.ones((2, 1)), "indices"), + (torch.zeros((4,), dtype=torch.uint8), torch.ones((2, 1)), "indices"), + (torch.full((2, 2), 4, dtype=torch.uint8), torch.ones((2, 1)), "out-of-range"), + (torch.zeros((2, 2), dtype=torch.uint8), torch.ones((2,), dtype=torch.float32), "magnitudes"), + (torch.zeros((2, 2), dtype=torch.uint8), -torch.ones((2, 1)), "nonnegative"), + ( + torch.zeros((2, 2), dtype=torch.uint8), + torch.tensor([[1.0], [float("nan")]]), + "finite", + ), + ], +) +def test_decode_rejects_malformed_quantized_state(indices, magnitudes, match): + with pytest.raises((TypeError, ValueError), match=match): + _decode_quantized_momentum( + _codebook(), + indices, + magnitudes, + logical_shape=(4,), + period=2, + step=1, + ) + + +def test_second_moment_projection_rejects_negative_or_nonfinite_state(): + with pytest.raises(ValueError, match="nonnegative"): + _reduce_block_second_moment(torch.tensor([1.0, -1.0]), period=1, step=1) + with pytest.raises(ValueError, match="finite"): + _project_factored_second_moment( + torch.tensor([[1.0, float("nan")]]), + step=1, + ) + + +def test_factored_expansion_rejects_bad_geometry_and_degenerate_inconsistent_state(): + with pytest.raises(ValueError, match="2-D"): + _expand_factored_second_moment( + torch.ones(4), + torch.ones(1), + logical_shape=(4,), + step=1, + ) + with pytest.raises(ValueError, match="do not match"): + _expand_factored_second_moment( + torch.ones(3), + torch.ones(2), + logical_shape=(2, 2), + step=1, + ) + with pytest.raises(ValueError, match="zero row mean"): + _expand_factored_second_moment( + torch.zeros(2), + torch.ones(2), + logical_shape=(2, 2), + step=1, + ) + + +def test_factored_expansion_rejects_a_finite_input_whose_result_overflows_fp32(): + maximum = torch.finfo(torch.float32).max + tiny = torch.finfo(torch.float32).tiny + with pytest.raises(ValueError, match="cannot be represented"): + _expand_factored_second_moment( + torch.tensor([tiny, maximum]), + torch.tensor([maximum, maximum]), + logical_shape=(2, 2), + step=1, + ) diff --git a/tests/test_shard_identity_contracts.py b/tests/test_shard_identity_contracts.py index bc3eb1f..3e58d56 100644 --- a/tests/test_shard_identity_contracts.py +++ b/tests/test_shard_identity_contracts.py @@ -5,6 +5,7 @@ import pytest import gefen +import gefen.contracts as contracts_module from gefen import ( IDENTITY_SCHEMA_VERSION, LogicalSlice, @@ -16,6 +17,7 @@ ShardPlacement, ShardingManifest, ) +from gefen.contracts import LogicalRegion def _group(): @@ -62,6 +64,37 @@ def _owner_shard(parameter, group, member, owner): ) +def _dtensor_shard( + parameter, + group, + member, + offsets, + lengths, + *, + dimension=0, + replicate=False, +): + coordinate = group.ordered_members.index(member) + return ShardIdentity( + parameter, + ParameterLayout.DTENSOR_1D_DEFAULT_WORLD, + LogicalRegion(offsets, lengths), + placements=( + ShardPlacement( + "dp", + PlacementKind.REPLICATE + if replicate + else PlacementKind.DIMENSION_SHARD, + coordinate, + len(group.ordered_members), + None if replicate else dimension, + ), + ), + process_group=group, + local_member=member, + ) + + def test_parameter_and_process_group_identities_are_exact_and_immutable(): parameter = ParameterIdentity("Encoder.Block.Weight", [4, 8]) group = ProcessGroupIdentity("pipeline:1/dp", ["worker:b", "worker:a"]) @@ -123,6 +156,73 @@ def test_placement_and_logical_slice_validation(): LogicalSlice(0, -1) +def test_released_logical_slice_sort_key_shape_has_a_golden_value(): + parameter = ParameterIdentity("layer.weight", (8,)) + group = ProcessGroupIdentity("dp", ("rank:0", "rank:1")) + shard = _flat_shard(parameter, group, "rank:1", 4, 4) + + assert shard.sort_key == ( + "layer.weight", + 4, + 4, + "flattened_element_shard", + "dp", + 1, + -1, + (("data_parallel", "flat_shard", 1, 2, -1),), + ) + + +def test_logical_region_is_immutable_and_validates_rank_bounds_and_values(): + offsets = [1, 0] + lengths = [2, 4] + region = LogicalRegion(offsets, lengths) + offsets[0] = 0 + lengths[0] = 3 + + assert region.offsets == (1, 0) + assert region.lengths == (2, 4) + assert region.rank == 2 + assert region.numel == 8 + with pytest.raises(FrozenInstanceError): + region.offsets = (0, 0) + + parameter = ParameterIdentity("layer.weight", (4, 4)) + assert LogicalRegion.full(parameter) == LogicalRegion((0, 0), (4, 4)) + region.validate_bounds(parameter) + with pytest.raises(ValueError, match="same rank"): + LogicalRegion((0,), (1, 1)) + with pytest.raises(ValueError, match="nonnegative"): + LogicalRegion((-1, 0), (1, 1)) + with pytest.raises(ValueError, match="nonnegative"): + LogicalRegion((False, 0), (1, 1)) + with pytest.raises(TypeError, match="sequence"): + LogicalRegion("00", (1, 1)) + with pytest.raises(ValueError, match="rank"): + LogicalRegion((0,), (4,)).validate_bounds(parameter) + with pytest.raises(ValueError, match="exceeds"): + LogicalRegion((3, 0), (2, 4)).validate_bounds(parameter) + + +def test_logical_region_intersection_overlap_and_exact_coverage_utilities(): + parameter = ParameterIdentity("layer.weight", (4, 4)) + first = LogicalRegion((0, 0), (2, 4)) + second = LogicalRegion((2, 0), (2, 4)) + overlapping = LogicalRegion((1, 0), (2, 4)) + + assert first.intersection(second) == LogicalRegion((2, 0), (0, 4)) + assert not first.overlaps(second) + assert first.intersection(overlapping) == LogicalRegion((1, 0), (1, 4)) + assert first.overlaps(overlapping) + LogicalRegion.validate_exact_coverage(parameter, (second, first)) + with pytest.raises(ValueError, match="overlap"): + LogicalRegion.validate_exact_coverage(parameter, (first, overlapping)) + with pytest.raises(ValueError, match="exactly cover"): + LogicalRegion.validate_exact_coverage(parameter, (first,)) + with pytest.raises(ValueError, match="equal ranks"): + first.intersection(LogicalRegion((0,), (1,))) + + def test_replicated_identity_can_be_local_or_process_group_scoped(): parameter = ParameterIdentity("layer.weight", (4, 4)) local = ShardIdentity( @@ -169,6 +269,177 @@ def test_contiguous_slice_schema_rejects_dtensor_identity_until_regions_exist(): ) +def test_dtensor_row_shard_manifest_accepts_uneven_regions_and_sorts_them(): + parameter = ParameterIdentity("layer.weight", (5, 4)) + group = ProcessGroupIdentity("dp", ("rank:0", "rank:1", "rank:2")) + shards = ( + _dtensor_shard(parameter, group, "rank:2", (4, 0), (1, 4)), + _dtensor_shard(parameter, group, "rank:0", (0, 0), (2, 4)), + _dtensor_shard(parameter, group, "rank:1", (2, 0), (2, 4)), + ) + + manifest = ShardingManifest(shards) + + assert tuple(item.local_member for item in manifest.shards) == ( + "rank:0", + "rank:1", + "rank:2", + ) + assert tuple(item.logical_region for item in manifest.shards) == ( + LogicalRegion((0, 0), (2, 4)), + LogicalRegion((2, 0), (2, 4)), + LogicalRegion((4, 0), (1, 4)), + ) + + +def test_dtensor_column_shard_and_replicate_manifests_are_complete(): + parameter = ParameterIdentity("layer.weight", (3, 5)) + group = ProcessGroupIdentity("dp", ("rank:0", "rank:1")) + column_manifest = ShardingManifest( + ( + _dtensor_shard( + parameter, + group, + "rank:1", + (0, 3), + (3, 2), + dimension=1, + ), + _dtensor_shard( + parameter, + group, + "rank:0", + (0, 0), + (3, 3), + dimension=1, + ), + ) + ) + assert tuple(item.logical_region.offsets for item in column_manifest.shards) == ( + (0, 0), + (0, 3), + ) + + replicated = ShardingManifest( + tuple( + _dtensor_shard( + parameter, + group, + member, + (0, 0), + parameter.global_shape, + replicate=True, + ) + for member in reversed(group.ordered_members) + ) + ) + assert all( + item.logical_region == LogicalRegion.full(parameter) + for item in replicated.shards + ) + + +def test_dtensor_manifest_retains_explicit_empty_member_regions(): + parameter = ParameterIdentity("small.weight", (2, 4)) + group = ProcessGroupIdentity( + "dp", ("rank:0", "rank:1", "rank:2", "rank:3") + ) + shards = tuple( + _dtensor_shard(parameter, group, member, offsets, lengths) + for member, offsets, lengths in ( + ("rank:0", (0, 0), (1, 4)), + ("rank:1", (1, 0), (1, 4)), + ("rank:2", (2, 0), (0, 4)), + ("rank:3", (2, 0), (0, 4)), + ) + ) + + manifest = ShardingManifest(tuple(reversed(shards))) + + assert tuple(item.local_member for item in manifest.shards) == ( + "rank:0", + "rank:1", + "rank:2", + "rank:3", + ) + assert tuple(item.logical_region.numel for item in manifest.shards) == ( + 4, + 4, + 0, + 0, + ) + + +def test_dtensor_manifest_rejects_gap_overlap_and_missing_member(): + parameter = ParameterIdentity("layer.weight", (4, 4)) + group = ProcessGroupIdentity("dp", ("rank:0", "rank:1")) + with pytest.raises(ValueError, match="gapless"): + ShardingManifest( + ( + _dtensor_shard(parameter, group, "rank:0", (0, 0), (1, 4)), + _dtensor_shard(parameter, group, "rank:1", (2, 0), (2, 4)), + ) + ) + with pytest.raises(ValueError, match="gapless"): + ShardingManifest( + ( + _dtensor_shard(parameter, group, "rank:0", (0, 0), (3, 4)), + _dtensor_shard(parameter, group, "rank:1", (2, 0), (2, 4)), + ) + ) + with pytest.raises(ValueError, match="each process-group member"): + ShardingManifest( + (_dtensor_shard(parameter, group, "rank:0", (0, 0), (4, 4)),) + ) + + +def test_dtensor_identity_rejects_invalid_region_and_placement_geometry(): + parameter = ParameterIdentity("layer.weight", (4, 4)) + group = ProcessGroupIdentity("dp", ("rank:0", "rank:1")) + with pytest.raises(ValueError, match="LogicalRegion"): + ShardIdentity( + parameter, + ParameterLayout.DTENSOR_1D_DEFAULT_WORLD, + LogicalSlice(0, 8), + placements=( + ShardPlacement("dp", PlacementKind.DIMENSION_SHARD, 0, 2, 0), + ), + process_group=group, + local_member="rank:0", + ) + with pytest.raises(ValueError, match="unsharded"): + _dtensor_shard(parameter, group, "rank:0", (0, 1), (2, 3)) + with pytest.raises(ValueError, match="replicated"): + _dtensor_shard( + parameter, + group, + "rank:0", + (0, 0), + (2, 4), + replicate=True, + ) + with pytest.raises(ValueError, match="one dimension-shard"): + ShardIdentity( + parameter, + ParameterLayout.DTENSOR_1D_DEFAULT_WORLD, + LogicalRegion((0, 0), (2, 4)), + placements=(ShardPlacement("dp", PlacementKind.FLAT_SHARD, 0, 2),), + process_group=group, + local_member="rank:0", + ) + with pytest.raises(ValueError, match="coordinates"): + ShardIdentity( + parameter, + ParameterLayout.DTENSOR_1D_DEFAULT_WORLD, + LogicalRegion((0, 0), (2, 4)), + placements=( + ShardPlacement("dp", PlacementKind.DIMENSION_SHARD, 1, 2, 0), + ), + process_group=group, + local_member="rank:0", + ) + + @pytest.mark.parametrize("version", [True, 1.0, 0, 2]) def test_parameter_identity_schema_version_requires_exact_supported_int(version): with pytest.raises(ValueError, match="schema version"): @@ -472,3 +743,4 @@ def test_identity_contracts_are_public_lazy_exports(): assert getattr(gefen, name).__module__ == "gefen.contracts" assert "ParameterRebinding" in gefen.__all__ assert gefen.ParameterRebinding.__module__ == "gefen.rebinding" + assert "LogicalRegion" in contracts_module.__all__ From 251ee716a5c0706a5f9c3efaf4aa34ddfe95276e Mon Sep 17 00:00:00 2001 From: thad0ctor Date: Mon, 13 Jul 2026 00:05:33 -0700 Subject: [PATCH 09/52] Retain logical optimizer slot identities --- src/gefen/gefen.py | 256 ++++++++++++++++++++++++++---- src/gefen/rebinding.py | 32 ++++ tests/test_canonical_state_cpu.py | 45 ++++++ tests/test_rebinding_cpu.py | 217 +++++++++++++++++++++++++ 4 files changed, 521 insertions(+), 29 deletions(-) diff --git a/src/gefen/gefen.py b/src/gefen/gefen.py index 090a265..1b5a25c 100644 --- a/src/gefen/gefen.py +++ b/src/gefen/gefen.py @@ -46,7 +46,7 @@ _gefen_contract, ) from gefen.partitioning import find_period_by_block_variance -from gefen.rebinding import ParameterRebinding +from gefen.rebinding import LogicalSlotBinding, ParameterRebinding import gefen.quantization as quantization_module from gefen.kernels.automatic_vmean import ( automatic_vmean_update_cuda as _automatic_vmean_update_cuda, @@ -1212,6 +1212,7 @@ def __init__( self._gefen_sharding_manifest = None self._gefen_post_sharding_finalized = False self._gefen_finalized_slots = () + self._gefen_logical_slots = () # ``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 @@ -1375,32 +1376,184 @@ def _canonical_identity_ready(self) -> bool: return self._finalized_binding_layout_matches() def _finalized_binding_layout_matches(self) -> bool: - if len(self.param_groups) != len(self._gefen_finalized_slots): - return False - for group, expected in zip(self.param_groups, self._gefen_finalized_slots): - params = group.get("params") - if not isinstance(params, (list, tuple)) or len(params) != len(expected): + try: + if ( + type(self._gefen_logical_slots) is not tuple + or not self._gefen_logical_slots + or type(self._gefen_finalized_slots) is not tuple + or type(self._gefen_local_shard_bindings) is not tuple + or type(self._gefen_shard_bindings) is not dict + or type(self._param_names) is not dict + or not isinstance(self._gefen_sharding_manifest, ShardingManifest) + or len(self.param_groups) != len(self._gefen_finalized_slots) + ): return False - if any(live is not bound for live, bound in zip(params, expected)): + + logical_groups = [[] for _ in self.param_groups] + logical_fqns = set() + previous_position = None + for logical_slot in self._gefen_logical_slots: + if ( + type(logical_slot) is not LogicalSlotBinding + or type(logical_slot.group_index) is not int + or type(logical_slot.original_slot_index) is not int + or type(logical_slot.compatibility_name) is not str + or logical_slot.compatibility_name + != logical_slot.compatibility_name.lower() + or not isinstance(logical_slot.shard, ShardIdentity) + ): + return False + position = ( + logical_slot.group_index, + logical_slot.original_slot_index, + ) + if ( + logical_slot.group_index < 0 + or logical_slot.group_index >= len(logical_groups) + or logical_slot.original_slot_index + != len(logical_groups[logical_slot.group_index]) + ): + return False + if previous_position is None: + if position != (0, 0): + return False + elif logical_slot.group_index == previous_position[0]: + if logical_slot.original_slot_index != previous_position[1] + 1: + return False + elif position != (previous_position[0] + 1, 0): + return False + fqn = logical_slot.shard.parameter.fqn + if fqn in logical_fqns: + return False + logical_fqns.add(fqn) + logical_groups[logical_slot.group_index].append(logical_slot) + previous_position = position + + if any(not group for group in logical_groups): return False - live_params = [ - param for group in self.param_groups for param in group["params"] - ] - bound = [ - (parameter, shard) - for parameter, shard in self._gefen_local_shard_bindings - if parameter is not None - ] - if len(live_params) != len(bound) or len(self._gefen_shard_bindings) != len( - bound - ): - return False - for parameter, shard in bound: - if not self._parameter_in(live_params, parameter): + + manifest_shards = frozenset(self._gefen_sharding_manifest.shards) + if { + shard.parameter.fqn for shard in manifest_shards + } != logical_fqns or any( + logical_slot.shard not in manifest_shards + for logical_slot in self._gefen_logical_slots + ): return False - if self._gefen_shard_bindings.get(parameter) != shard: + + local_bindings = {} + previous_sort_key = None + for item in self._gefen_local_shard_bindings: + if type(item) is not tuple or len(item) != 2: + return False + parameter, shard = item + if parameter is not None and not isinstance(parameter, torch.Tensor): + return False + if not isinstance(shard, ShardIdentity): + return False + if previous_sort_key is not None and shard.sort_key < previous_sort_key: + return False + previous_sort_key = shard.sort_key + fqn = shard.parameter.fqn + if fqn in local_bindings: + return False + local_bindings[fqn] = (parameter, shard) + if set(local_bindings) != logical_fqns: return False - return True + + expected_live_groups = [] + expected_name_groups = [] + expected_live_bindings = [] + for logical_group in logical_groups: + live_group = [] + name_group = [] + for logical_slot in logical_group: + parameter, shard = local_bindings[ + logical_slot.shard.parameter.fqn + ] + if shard != logical_slot.shard: + return False + pruned = ( + shard.layout is ParameterLayout.WHOLE_PARAMETER_OWNER + and shard.local_member != shard.owner + ) + if pruned: + if parameter is not None: + return False + continue + if parameter is None: + return False + live_group.append(parameter) + name_group.append(logical_slot.compatibility_name) + expected_live_bindings.append((parameter, shard)) + expected_live_groups.append(tuple(live_group)) + expected_name_groups.append(tuple(name_group)) + + if len(self._gefen_shard_bindings) != len(expected_live_bindings): + return False + expected_live_ids = { + id(parameter) for parameter, _ in expected_live_bindings + } + if len(expected_live_ids) != len(expected_live_bindings): + return False + if { + id(parameter) for parameter in self._gefen_shard_bindings + } != expected_live_ids: + return False + for parameter, shard in expected_live_bindings: + if self._gefen_shard_bindings.get(parameter) != shard: + return False + + if len(self._param_names) != len(expected_live_bindings): + return False + if { + id(parameter) for parameter in self._param_names + } != expected_live_ids: + return False + for group_index, (group, expected_params, expected_names) in enumerate( + zip(self.param_groups, expected_live_groups, expected_name_groups) + ): + params = group.get("params") + names = group.get("param_names") + finalized = self._gefen_finalized_slots[group_index] + if ( + not isinstance(params, (list, tuple)) + or not isinstance(names, (list, tuple)) + or type(finalized) is not tuple + or len(params) != len(expected_params) + or len(names) != len(expected_names) + or len(finalized) != len(expected_params) + ): + return False + if any( + live is not expected + for live, expected in zip(params, expected_params) + ) or any( + bound is not expected + for bound, expected in zip(finalized, expected_params) + ): + return False + if tuple(names) != expected_names: + return False + for parameter, expected_name in zip(expected_params, expected_names): + if self._param_names.get(parameter) != expected_name: + return False + parameter_state = self.state.get(parameter) + if ( + type(parameter_state) is not dict + or parameter_state.get("name") != expected_name + ): + return False + return True + except ( + AttributeError, + IndexError, + KeyError, + RuntimeError, + TypeError, + ValueError, + ): + return False def _assert_finalized_binding_layout(self) -> None: if self._gefen_post_sharding_finalized and not self._finalized_binding_layout_matches(): @@ -1880,6 +2033,7 @@ def _assert_rebinding_pristine(self, rebindings) -> None: or self._gefen_local_shard_bindings or self._gefen_sharding_manifest is not None or self._gefen_finalized_slots + or self._gefen_logical_slots ): raise RuntimeError( "Gefen parameter rebinding found an incomplete prior identity plan" @@ -2284,8 +2438,9 @@ def _stage_post_sharding( staged._param_names = {} staged._gefen_shard_bindings = {} local_bindings = [] + logical_slots = [] slot_cursor = 0 - for group in self.param_groups: + for group_index, group in enumerate(self.param_groups): staged_group = dict(group) staged_group.pop("_gefen_checkpoint_metadata", None) staged_params = [] @@ -2293,13 +2448,21 @@ def _stage_post_sharding( names = list(group.get("param_names", ())) if len(names) != len(group["params"]): names = [self._param_name(param) for param in group["params"]] - for _ in names: + for original_slot_index, _ in enumerate(names): binding_index = assigned_positions[slot_cursor] rebinding = rebindings[binding_index] compatibility_name = binding_names[binding_index] slot_cursor += 1 target = rebinding.new_parameter local_bindings.append((target, rebinding.shard)) + logical_slots.append( + LogicalSlotBinding( + group_index, + original_slot_index, + compatibility_name, + rebinding.shard, + ) + ) if target is None: continue staged_params.append(target) @@ -2314,6 +2477,7 @@ def _stage_post_sharding( staged._gefen_local_shard_bindings = tuple( sorted(local_bindings, key=lambda item: item[1].sort_key) ) + staged._gefen_logical_slots = tuple(logical_slots) staged._gefen_sharding_manifest = manifest staged._gefen_post_sharding_finalized = True staged._gefen_codebook_process_group = codebook_process_group @@ -2828,11 +2992,32 @@ def _unique_name(base: str, existing_names) -> str: def _sync_param_names_to_state(self) -> None: self._param_names = {} - for group in self.param_groups: + logical_groups = None + if self._gefen_post_sharding_finalized: + logical_groups = [[] for _ in self.param_groups] + for logical_slot in self._gefen_logical_slots: + if ( + logical_slot.shard.layout + is ParameterLayout.WHOLE_PARAMETER_OWNER + and logical_slot.shard.local_member != logical_slot.shard.owner + ): + continue + logical_groups[logical_slot.group_index].append( + logical_slot.compatibility_name + ) + for group_index, group in enumerate(self.param_groups): params = group["params"] - names = list(group.get("param_names", ())) - if len(names) != len(params): - names = [self._param_name(p) for p in params] + if logical_groups is None: + names = list(group.get("param_names", ())) + if len(names) != len(params): + names = [self._param_name(p) for p in params] + else: + names = logical_groups[group_index] + if len(names) != len(params): + raise RuntimeError( + "Gefen finalized logical-slot registry does not match the " + "loaded layout" + ) normalized_names = [] for param, name in zip(params, names): name = str(name).lower() @@ -6052,6 +6237,19 @@ def _canonical_import_live_token(self): id(self._gefen_codebook_process_group), self._canonical_value_token(self._serialized_codebook_scope()), id(self._gefen_sharding_manifest), + id(self._gefen_logical_slots), + tuple( + ( + id(logical_slot), + logical_slot.group_index, + logical_slot.original_slot_index, + logical_slot.compatibility_name, + self._canonical_value_token( + self._serialized_canonical_shard(logical_slot.shard) + ), + ) + for logical_slot in self._gefen_logical_slots + ), tuple( (id(parameter), shard.sort_key) for parameter, shard in self._gefen_local_shard_bindings diff --git a/src/gefen/rebinding.py b/src/gefen/rebinding.py index 657c592..804e147 100644 --- a/src/gefen/rebinding.py +++ b/src/gefen/rebinding.py @@ -6,6 +6,38 @@ from gefen.contracts import ShardIdentity +@dataclass(frozen=True, slots=True) +class LogicalSlotBinding: + """Tensor-free identity for one original optimizer parameter slot.""" + + group_index: int + original_slot_index: int + compatibility_name: str + shard: ShardIdentity + + def __post_init__(self) -> None: + if type(self.group_index) is not int: + raise TypeError("LogicalSlotBinding.group_index must be an integer") + if self.group_index < 0: + raise ValueError("LogicalSlotBinding.group_index must be nonnegative") + if type(self.original_slot_index) is not int: + raise TypeError( + "LogicalSlotBinding.original_slot_index must be an integer" + ) + if self.original_slot_index < 0: + raise ValueError( + "LogicalSlotBinding.original_slot_index must be nonnegative" + ) + if type(self.compatibility_name) is not str: + raise TypeError("LogicalSlotBinding.compatibility_name must be a string") + if self.compatibility_name != self.compatibility_name.lower(): + raise ValueError( + "LogicalSlotBinding.compatibility_name must be lowercase" + ) + if not isinstance(self.shard, ShardIdentity): + raise TypeError("LogicalSlotBinding.shard must be a ShardIdentity") + + @dataclass(frozen=True, eq=False) class ParameterRebinding: """Bind one optimizer slot to a local tensor or prune it as a non-owner.""" diff --git a/tests/test_canonical_state_cpu.py b/tests/test_canonical_state_cpu.py index fa4c7a1..3e5becc 100644 --- a/tests/test_canonical_state_cpu.py +++ b/tests/test_canonical_state_cpu.py @@ -22,6 +22,7 @@ ShardPlacement, ShardingManifest, ) +from gefen.rebinding import LogicalSlotBinding def _replicated(identity, group=None, member=None): @@ -277,6 +278,7 @@ def test_initialized_import_maps_by_fqn_across_order_and_group_boundaries(): input_magnitude_value = input_magnitude.clone() prepared = target.prepare_canonical_state_import(exported) + logical_slots = target._gefen_logical_slots assert isinstance(prepared, PreparedCanonicalStateImport) assert target._gefen_codebook is None assert exported["common"]["gefen_codebook"] is input_codebook @@ -286,6 +288,7 @@ def test_initialized_import_maps_by_fqn_across_order_and_group_boundaries(): target.commit_canonical_state_import(prepared) assert not load_hooks + assert target._gefen_logical_slots is logical_slots assert torch.equal( target.state[target_first]["m_magnitude"], source.state[source_first]["m_magnitude"], @@ -309,6 +312,10 @@ def test_initialized_import_maps_by_fqn_across_order_and_group_boundaries(): assert torch.equal(target_second, source_second) assert torch.equal(target._gefen_codebook, source._gefen_codebook) + target.move_state_("cpu") + assert target._gefen_logical_slots is logical_slots + assert target.optimizer_contract().capabilities.canonical_state_io + with pytest.raises(RuntimeError, match="already consumed"): target.commit_canonical_state_import(prepared) @@ -344,6 +351,44 @@ def test_prepared_import_rejects_wrong_optimizer_and_stale_live_state(): _assert_snapshot_identity(target, target_before) +@pytest.mark.parametrize("mutation", ["registry_identity", "registry_structure"]) +def test_prepared_import_freshness_covers_logical_slot_registry(mutation): + parameter = torch.nn.Parameter(torch.ones(4)) + optimizer = Gefen( + [("p", parameter)], + fused=False, + factored_v_2d=False, + ) + identity = ParameterIdentity("Model.P", (4,)) + shard = _replicated(identity) + _finalize(optimizer, ((parameter, shard),), ShardingManifest((shard,))) + prepared = optimizer.prepare_canonical_state_import( + optimizer.export_canonical_state() + ) + slot = optimizer._gefen_logical_slots[0] + if mutation == "registry_identity": + optimizer._gefen_logical_slots = ( + LogicalSlotBinding( + slot.group_index, + slot.original_slot_index, + slot.compatibility_name, + slot.shard, + ), + ) + else: + object.__setattr__( + slot.shard, + "parameter", + ParameterIdentity("Model.Changed", (4,)), + ) + assert optimizer._finalized_binding_layout_matches() + before = _snapshot(optimizer) + + with pytest.raises(RuntimeError, match="changed after"): + optimizer.commit_canonical_state_import(prepared) + _assert_snapshot_identity(optimizer, before) + + @pytest.mark.parametrize( "corrupt", [ diff --git a/tests/test_rebinding_cpu.py b/tests/test_rebinding_cpu.py index 76d36f3..b852a69 100644 --- a/tests/test_rebinding_cpu.py +++ b/tests/test_rebinding_cpu.py @@ -1,6 +1,9 @@ """CPU coverage for atomic pre-initialization parameter rebinding.""" import copy +from dataclasses import FrozenInstanceError, fields, is_dataclass, replace +import gc +import weakref import pytest import torch @@ -19,6 +22,7 @@ ShardPlacement, ShardingManifest, ) +from gefen.rebinding import LogicalSlotBinding def _replicated(parameter, fqn, shape=None): @@ -97,6 +101,23 @@ def _nested_equal(left, right): assert left == right +def _contains_tensor(value): + if torch.is_tensor(value): + return True + if is_dataclass(value) and not isinstance(value, type): + return any( + _contains_tensor(getattr(value, field.name)) for field in fields(value) + ) + if isinstance(value, dict): + return any( + _contains_tensor(key) or _contains_tensor(item) + for key, item in value.items() + ) + if isinstance(value, (list, tuple, set, frozenset)): + return any(_contains_tensor(item) for item in value) + return False + + def _snapshot(optimizer): return { "state": optimizer.state, @@ -120,6 +141,7 @@ def _snapshot(optimizer): "manifest": optimizer._gefen_sharding_manifest, "finalized": optimizer._gefen_post_sharding_finalized, "finalized_slots": optimizer._gefen_finalized_slots, + "logical_slots": optimizer._gefen_logical_slots, "caches": tuple( (name, getattr(optimizer, name)) for name in ( @@ -146,6 +168,7 @@ def _assert_snapshot(optimizer, snapshot): assert optimizer._gefen_sharding_manifest is snapshot["manifest"] assert optimizer._gefen_post_sharding_finalized is snapshot["finalized"] assert optimizer._gefen_finalized_slots is snapshot["finalized_slots"] + assert optimizer._gefen_logical_slots is snapshot["logical_slots"] assert optimizer._capt_stacks is snapshot["capt_stacks"] assert optimizer._static_mark_sig is snapshot["static_mark_sig"] assert optimizer._gefen_global_step is snapshot["global_step"] @@ -165,6 +188,200 @@ def _assert_snapshot(optimizer, snapshot): assert getattr(optimizer, name) is cache_ref +@pytest.mark.parametrize( + ("changes", "error", "match"), + [ + ({"group_index": True}, TypeError, "group_index must be an integer"), + ({"group_index": -1}, ValueError, "group_index must be nonnegative"), + ( + {"original_slot_index": 1.0}, + TypeError, + "original_slot_index must be an integer", + ), + ( + {"original_slot_index": -1}, + ValueError, + "original_slot_index must be nonnegative", + ), + ( + {"compatibility_name": object()}, + TypeError, + "compatibility_name must be a string", + ), + ( + {"compatibility_name": "MixedCase"}, + ValueError, + "compatibility_name must be lowercase", + ), + ({"shard": object()}, TypeError, "shard must be a ShardIdentity"), + ], +) +def test_logical_slot_binding_strict_validation(changes, error, match): + parameter = torch.nn.Parameter(torch.ones(4)) + values = { + "group_index": 0, + "original_slot_index": 0, + "compatibility_name": "", + "shard": _replicated(parameter, "Weight"), + } + values.update(changes) + + with pytest.raises(error, match=match): + LogicalSlotBinding(**values) + + +def _interleaved_logical_slot_optimizer(): + sources = [ + torch.nn.Parameter(torch.full((2, 2), float(index + 1))) + for index in range(5) + ] + source_refs = tuple(weakref.ref(parameter) for parameter in sources) + names = ( + "group0.live0", + "group0.remote", + "group0.live2", + "group1.remote", + "group1.live1", + ) + optimizer = GefenMuon( + [ + {"params": list(zip(names[:3], sources[:3]))}, + {"params": list(zip(names[3:], sources[3:]))}, + ], + fused=False, + ) + process_group = ProcessGroupIdentity("data_parallel", ("rank:0", "rank:1")) + owners = ("rank:0", "rank:1", "rank:0", "rank:1", "rank:0") + records = [] + local_shards = [] + for index, owner in enumerate(owners): + identity = ParameterIdentity("Model.Parameter{}".format(index), (2, 2)) + parameter_records = tuple( + _owner_shard(identity, process_group, member, owner) + for member in process_group.ordered_members + ) + records.extend(parameter_records) + local_shards.append( + next( + shard + for shard in parameter_records + if shard.local_member == "rank:0" + ) + ) + targets = { + index: torch.nn.Parameter(sources[index].detach().clone()) + for index in (0, 2, 4) + } + rebindings = [ + ParameterRebinding( + sources[index], + targets.get(index), + local_shards[index], + ) + for index in (4, 1, 3, 0, 2) + ] + optimizer.post_sharding( + rebindings, + manifest=ShardingManifest(tuple(records)), + ) + return optimizer, source_refs, targets, tuple(local_shards), names + + +def test_logical_slot_registry_preserves_original_interleaved_structure_without_tensors(): + optimizer, source_refs, targets, local_shards, names = ( + _interleaved_logical_slot_optimizer() + ) + gc.collect() + + assert all(reference() is None for reference in source_refs) + assert optimizer.param_groups[0]["params"] == [targets[0], targets[2]] + assert optimizer.param_groups[1]["params"] == [targets[4]] + assert optimizer.param_groups[0]["param_names"] == [names[0], names[2]] + assert optimizer.param_groups[1]["param_names"] == [names[4]] + assert tuple( + (slot.group_index, slot.original_slot_index) + for slot in optimizer._gefen_logical_slots + ) == ((0, 0), (0, 1), (0, 2), (1, 0), (1, 1)) + assert tuple( + slot.compatibility_name for slot in optimizer._gefen_logical_slots + ) == names + assert tuple( + slot.shard for slot in optimizer._gefen_logical_slots + ) == local_shards + assert not _contains_tensor(optimizer._gefen_logical_slots) + assert all( + not hasattr(slot, "__dict__") for slot in optimizer._gefen_logical_slots + ) + with pytest.raises(FrozenInstanceError): + optimizer._gefen_logical_slots[0].compatibility_name = "changed" + assert optimizer.optimizer_contract().capabilities.stable_shard_identity + + +@pytest.mark.parametrize( + "corruption", + [ + "logical_slots", + "manifest", + "finalized_slots", + "local_bindings", + "group_names", + "name_cache", + "state_name", + ], +) +def test_finalized_layout_guard_checks_logical_slot_caches_and_names(corruption): + parameters = [ + torch.nn.Parameter(torch.full((4,), float(index + 1))) + for index in range(2) + ] + optimizer = Gefen( + [("first", parameters[0]), ("second", parameters[1])], + fused=False, + factored_v_2d=False, + ) + shards = tuple( + _replicated(parameter, "Model.Parameter{}".format(index)) + for index, parameter in enumerate(parameters) + ) + optimizer.post_sharding( + tuple( + ParameterRebinding(parameter, parameter, shard) + for parameter, shard in zip(parameters, shards) + ), + manifest=ShardingManifest(shards), + ) + + if corruption == "logical_slots": + slots = optimizer._gefen_logical_slots + optimizer._gefen_logical_slots = ( + replace(slots[0], compatibility_name="changed"), + slots[1], + ) + elif corruption == "manifest": + optimizer._gefen_sharding_manifest = ShardingManifest( + ( + _replicated(parameters[0], "Model.Parameter0", shape=(5,)), + shards[1], + ) + ) + elif corruption == "finalized_slots": + optimizer._gefen_finalized_slots = (tuple(reversed(parameters)),) + elif corruption == "local_bindings": + optimizer._gefen_local_shard_bindings = tuple( + reversed(optimizer._gefen_local_shard_bindings) + ) + elif corruption == "group_names": + optimizer.param_groups[0]["param_names"].reverse() + elif corruption == "name_cache": + optimizer._param_names[parameters[0]] = "changed" + else: + optimizer.state[parameters[0]]["name"] = "changed" + + assert not optimizer.optimizer_contract().capabilities.stable_shard_identity + with pytest.raises(RuntimeError, match="changed outside post_sharding"): + optimizer.state_dict() + + def test_replicated_rebind_preserves_legacy_name_and_enables_identity_contract(): old = torch.nn.Parameter(torch.arange(16, dtype=torch.float32).reshape(4, 4)) new = torch.nn.Parameter(old.detach().clone()) From 81ced575b05a1b2c0cec5301638b1404bb2ac087 Mon Sep 17 00:00:00 2001 From: thad0ctor Date: Mon, 13 Jul 2026 00:08:48 -0700 Subject: [PATCH 10/52] Add portable state transport primitives --- src/gefen/portable_fields.py | 353 +++++++++++++++ src/gefen/portable_identity.py | 289 ++++++++++++ src/gefen/portable_wire.py | 747 ++++++++++++++++++++++++++++++++ tests/test_portable_fields.py | 538 +++++++++++++++++++++++ tests/test_portable_identity.py | 517 ++++++++++++++++++++++ tests/test_portable_wire.py | 480 ++++++++++++++++++++ 6 files changed, 2924 insertions(+) create mode 100644 src/gefen/portable_fields.py create mode 100644 src/gefen/portable_identity.py create mode 100644 src/gefen/portable_wire.py create mode 100644 tests/test_portable_fields.py create mode 100644 tests/test_portable_identity.py create mode 100644 tests/test_portable_wire.py diff --git a/src/gefen/portable_fields.py b/src/gefen/portable_fields.py new file mode 100644 index 0000000..8a5441c --- /dev/null +++ b/src/gefen/portable_fields.py @@ -0,0 +1,353 @@ +"""Pure assembly and projection for dense logical optimizer-state fields.""" + +from collections.abc import Sequence + +import torch + +from gefen.contracts import ( + LogicalRegion, + ParameterIdentity, + ParameterLayout, + PlacementKind, + ShardIdentity, + ShardingManifest, +) +from gefen.portable import _element_chunks, _read_flat_chunk, _validate_plain_tensor + + +def _validate_parameter(parameter) -> ParameterIdentity: + if not isinstance(parameter, ParameterIdentity): + raise TypeError("parameter must be a ParameterIdentity") + # Reconstruct the frozen descriptor so an object modified through low-level + # attribute access cannot bypass its invariants at a checkpoint boundary. + return ParameterIdentity( + parameter.fqn, + parameter.global_shape, + schema_version=parameter.schema_version, + ) + + +def _validate_manifest(manifest) -> ShardingManifest: + if not isinstance(manifest, ShardingManifest): + raise TypeError("manifest must be a ShardingManifest") + return ShardingManifest( + manifest.shards, + schema_version=manifest.schema_version, + ) + + +def _validate_dense_tensor(value, *, name: str, shape) -> torch.Tensor: + value = _validate_plain_tensor( + value, + name=name, + dtype=torch.float32, + ) + if value.device.type != "cpu": + raise ValueError("{} must be on CPU".format(name)) + if tuple(value.shape) != tuple(shape): + raise ValueError("{} must have shape {}".format(name, tuple(shape))) + return value + + +def _finish_tight(value: torch.Tensor, *, name: str) -> torch.Tensor: + if ( + value.device.type != "cpu" + or value.dtype != torch.float32 + or value.requires_grad + or not value.is_contiguous() + or value.storage_offset() != 0 + or value.untyped_storage().nbytes() != value.numel() * value.element_size() + ): + raise RuntimeError("{} did not produce tight detached CPU fp32 state".format(name)) + return value + + +def _tight_clone(value: torch.Tensor, *, name: str) -> torch.Tensor: + result = torch.empty(tuple(value.shape), dtype=torch.float32, device="cpu") + flat = result.reshape(-1) + for start, stop in _element_chunks(value.numel()): + flat[start:stop].copy_(_read_flat_chunk(value, start, stop).resolve_conj().resolve_neg()) + return _finish_tight(result, name=name) + + +def _tensor_bits_equal(left: torch.Tensor, right: torch.Tensor) -> bool: + for start, stop in _element_chunks(left.numel()): + left_bytes = ( + _read_flat_chunk(left, start, stop) + .resolve_conj() + .resolve_neg() + .contiguous() + .view(torch.uint8) + ) + right_bytes = ( + _read_flat_chunk(right, start, stop) + .resolve_conj() + .resolve_neg() + .contiguous() + .view(torch.uint8) + ) + if not torch.equal(left_bytes, right_bytes): + return False + return True + + +def _tight_flat_slice( + value: torch.Tensor, + *, + flat_offset: int, + length: int, + name: str, +) -> torch.Tensor: + result = torch.empty((length,), dtype=torch.float32, device="cpu") + for start, stop in _element_chunks(length): + result[start:stop].copy_( + _read_flat_chunk( + value, + flat_offset + start, + flat_offset + stop, + ) + .resolve_conj() + .resolve_neg() + ) + return _finish_tight(result, name=name) + + +def _full_output(parameter: ParameterIdentity) -> torch.Tensor: + return torch.empty( + parameter.global_shape, + dtype=torch.float32, + device="cpu", + ) + + +def _region_index(region: LogicalRegion): + return tuple(slice(offset, offset + length) for offset, length in zip(region.offsets, region.lengths)) + + +def _normalize_fragments(fragments): + if isinstance(fragments, (str, bytes, bytearray)): + raise TypeError("fragments must be a sequence of shard/payload pairs") + try: + fragments = tuple(fragments) + except TypeError as exc: + raise TypeError("fragments must be a sequence of shard/payload pairs") from exc + + normalized = [] + for item in fragments: + if not isinstance(item, Sequence) or isinstance(item, (str, bytes, bytearray)) or len(item) != 2: + raise TypeError("each fragment must be a shard/payload pair") + shard, payload = item + if not isinstance(shard, ShardIdentity): + raise TypeError("each fragment identity must be a ShardIdentity") + normalized.append((shard, payload)) + return tuple(normalized) + + +def _validate_fragment_set( + manifest: ShardingManifest, + parameter: ParameterIdentity, + fragments, +): + expected = manifest.for_parameter(parameter.fqn) + if not expected: + raise ValueError("manifest does not contain the requested parameter") + if any(shard.parameter != parameter for shard in expected): + raise ValueError("manifest parameter identity does not match parameter") + + fragments = _normalize_fragments(fragments) + identities = tuple(shard for shard, _ in fragments) + if len(set(identities)) != len(identities): + raise ValueError("fragments must not contain duplicate shard identities") + if set(identities) != set(expected): + raise ValueError("fragments must exactly cover the requested parameter manifest") + return expected, {shard: payload for shard, payload in fragments} + + +def _assemble_replicated( + parameter: ParameterIdentity, + shards, + payload_by_shard, + *, + name: str, +) -> torch.Tensor: + reference = None + for shard in shards: + payload = payload_by_shard[shard] + if payload is None: + raise ValueError("{} replicas require a payload on every shard".format(name)) + payload = _validate_dense_tensor( + payload, + name="{} fragment".format(name), + shape=parameter.global_shape, + ) + if reference is None: + reference = payload + elif not _tensor_bits_equal(reference, payload): + raise ValueError("{} replicas disagree".format(name)) + if reference is None: # ShardingManifest forbids this, but keep failure local. + raise ValueError("{} requires at least one replica".format(name)) + return _tight_clone(reference, name="assembled {}".format(name)) + + +def _assemble_flattened( + parameter: ParameterIdentity, + shards, + payload_by_shard, +) -> torch.Tensor: + result = _full_output(parameter) + flat = result.reshape(-1) + for shard in shards: + logical_slice = shard.logical_slice + payload = payload_by_shard[shard] + if logical_slice.length == 0: + if payload is not None: + raise ValueError("an empty flattened fragment must have payload None") + continue + if payload is None: + raise ValueError("a nonempty flattened fragment requires a payload") + payload = _validate_dense_tensor( + payload, + name="flattened fragment", + shape=(logical_slice.length,), + ) + start = logical_slice.flat_offset + flat[start : start + logical_slice.length].copy_(payload.detach()) + return _finish_tight(result, name="assembled flattened field") + + +def _assemble_owner( + parameter: ParameterIdentity, + shards, + payload_by_shard, +) -> torch.Tensor: + owner_payload = None + for shard in shards: + payload = payload_by_shard[shard] + if shard.local_member == shard.owner: + if payload is None: + raise ValueError("the whole-parameter owner requires a payload") + owner_payload = _validate_dense_tensor( + payload, + name="whole-parameter owner fragment", + shape=parameter.global_shape, + ) + elif payload is not None: + raise ValueError("whole-parameter non-owners must have payload None") + if owner_payload is None: + raise ValueError("whole-parameter fragments do not contain the owner payload") + return _tight_clone(owner_payload, name="assembled whole-parameter field") + + +def _assemble_dtensor_dimension_shards( + parameter: ParameterIdentity, + shards, + payload_by_shard, +) -> torch.Tensor: + result = _full_output(parameter) + for shard in shards: + region = shard.logical_region + payload = payload_by_shard[shard] + if region.numel == 0: + if payload is not None: + raise ValueError("an empty DTensor fragment must have payload None") + continue + if payload is None: + raise ValueError("a nonempty DTensor fragment requires a payload") + payload = _validate_dense_tensor( + payload, + name="DTensor fragment", + shape=region.lengths, + ) + result[_region_index(region)].copy_(payload.detach()) + return _finish_tight(result, name="assembled DTensor field") + + +def _assemble_dense_logical_field( + manifest, + parameter, + fragments, +) -> torch.Tensor: + """Assemble one complete topology-neutral fp32 CPU logical field.""" + + manifest = _validate_manifest(manifest) + parameter = _validate_parameter(parameter) + shards, payload_by_shard = _validate_fragment_set( + manifest, + parameter, + fragments, + ) + layout = shards[0].layout + if layout is ParameterLayout.REPLICATED: + return _assemble_replicated( + parameter, + shards, + payload_by_shard, + name="replicated", + ) + if layout is ParameterLayout.FLATTENED_ELEMENT_SHARD: + return _assemble_flattened(parameter, shards, payload_by_shard) + if layout is ParameterLayout.WHOLE_PARAMETER_OWNER: + return _assemble_owner(parameter, shards, payload_by_shard) + if layout is ParameterLayout.DTENSOR_1D_DEFAULT_WORLD: + placement_kind = shards[0].placements[0].kind + if placement_kind is PlacementKind.REPLICATE: + return _assemble_replicated( + parameter, + shards, + payload_by_shard, + name="replicated DTensor", + ) + if placement_kind is PlacementKind.DIMENSION_SHARD: + return _assemble_dtensor_dimension_shards( + parameter, + shards, + payload_by_shard, + ) + raise ValueError("unsupported portable field layout") + + +def _project_dense_logical_field( + parameter, + tensor, + shard, +): + """Project one complete logical fp32 CPU field to a target shard identity.""" + + parameter = _validate_parameter(parameter) + if not isinstance(shard, ShardIdentity): + raise TypeError("shard must be a ShardIdentity") + if shard.parameter != parameter: + raise ValueError("target shard parameter identity does not match parameter") + tensor = _validate_dense_tensor( + tensor, + name="global logical field", + shape=parameter.global_shape, + ) + + if shard.layout is ParameterLayout.REPLICATED: + return _tight_clone(tensor, name="projected replicated field") + if shard.layout is ParameterLayout.FLATTENED_ELEMENT_SHARD: + logical_slice = shard.logical_slice + if logical_slice.length == 0: + return None + return _tight_flat_slice( + tensor, + flat_offset=logical_slice.flat_offset, + length=logical_slice.length, + name="projected flattened field", + ) + if shard.layout is ParameterLayout.WHOLE_PARAMETER_OWNER: + if shard.local_member != shard.owner: + return None + return _tight_clone(tensor, name="projected whole-parameter field") + if shard.layout is ParameterLayout.DTENSOR_1D_DEFAULT_WORLD: + placement = shard.placements[0] + if placement.kind is PlacementKind.REPLICATE: + return _tight_clone(tensor, name="projected replicated DTensor field") + if placement.kind is PlacementKind.DIMENSION_SHARD: + region = shard.logical_region + if region.numel == 0: + return None + values = tensor[_region_index(region)] + return _tight_clone(values, name="projected DTensor field") + raise ValueError("unsupported portable field layout") diff --git a/src/gefen/portable_identity.py b/src/gefen/portable_identity.py new file mode 100644 index 0000000..1a1872d --- /dev/null +++ b/src/gefen/portable_identity.py @@ -0,0 +1,289 @@ +"""Strict primitive wire codecs for portable optimizer identities.""" + +from gefen.contracts import ( + LogicalRegion, + LogicalSlice, + ParameterIdentity, + ParameterLayout, + PlacementKind, + ProcessGroupIdentity, + ShardIdentity, + ShardPlacement, + ShardingManifest, +) + + +_PARAMETER_IDENTITY_KEYS = frozenset({"schema_version", "fqn", "global_shape"}) +_PROCESS_GROUP_IDENTITY_KEYS = frozenset({"schema_version", "semantic_name", "ordered_members"}) +_SHARD_IDENTITY_KEYS = frozenset( + { + "schema_version", + "parameter", + "layout", + "logical_extent", + "placements", + "process_group", + "local_member", + "owner", + } +) +_LOGICAL_SLICE_KEYS = frozenset({"kind", "flat_offset", "length"}) +_LOGICAL_REGION_KEYS = frozenset({"kind", "offsets", "lengths"}) +_SHARD_PLACEMENT_KEYS = frozenset( + { + "mesh_axis", + "kind", + "coordinate", + "parts", + "parameter_dimension", + } +) +_SHARDING_MANIFEST_KEYS = frozenset({"schema_version", "shards"}) + + +def _require_exact_record(value, keys, *, name): + if type(value) is not dict or set(value) != keys: + raise ValueError("{} has an invalid schema".format(name)) + return value + + +def _require_exact_list(value, *, name): + if type(value) is not list: + raise ValueError("{} must be a list".format(name)) + return value + + +def _require_optional_string(value, *, name): + if value is not None and type(value) is not str: + raise ValueError("{} must be a string or None".format(name)) + return value + + +def _serialize_parameter_identity(identity): + if not isinstance(identity, ParameterIdentity): + raise TypeError("identity must be a ParameterIdentity") + return { + "schema_version": identity.schema_version, + "fqn": identity.fqn, + "global_shape": list(identity.global_shape), + } + + +def _parse_parameter_identity(record): + record = _require_exact_record(record, _PARAMETER_IDENTITY_KEYS, name="parameter identity") + shape = _require_exact_list(record["global_shape"], name="parameter identity global_shape") + if type(record["schema_version"]) is not int: + raise ValueError("parameter identity schema_version must be an int") + if type(record["fqn"]) is not str: + raise ValueError("parameter identity fqn must be a string") + if any(type(dimension) is not int for dimension in shape): + raise ValueError("parameter identity global_shape must contain only integers") + try: + return ParameterIdentity( + fqn=record["fqn"], + global_shape=tuple(shape), + schema_version=record["schema_version"], + ) + except (TypeError, ValueError) as exc: + raise ValueError("parameter identity is invalid") from exc + + +def _normalize_parameter_identity(record): + return _serialize_parameter_identity(_parse_parameter_identity(record)) + + +def _serialize_process_group_identity(identity): + if not isinstance(identity, ProcessGroupIdentity): + raise TypeError("identity must be a ProcessGroupIdentity") + return { + "schema_version": identity.schema_version, + "semantic_name": identity.semantic_name, + "ordered_members": list(identity.ordered_members), + } + + +def _parse_process_group_identity(record): + record = _require_exact_record( + record, + _PROCESS_GROUP_IDENTITY_KEYS, + name="process-group identity", + ) + members = _require_exact_list( + record["ordered_members"], + name="process-group identity ordered_members", + ) + if type(record["schema_version"]) is not int: + raise ValueError("process-group identity schema_version must be an int") + if type(record["semantic_name"]) is not str: + raise ValueError("process-group identity semantic_name must be a string") + if any(type(member) is not str for member in members): + raise ValueError("process-group identity ordered_members must contain only strings") + try: + return ProcessGroupIdentity( + semantic_name=record["semantic_name"], + ordered_members=tuple(members), + schema_version=record["schema_version"], + ) + except (TypeError, ValueError) as exc: + raise ValueError("process-group identity is invalid") from exc + + +def _normalize_process_group_identity(record): + return _serialize_process_group_identity(_parse_process_group_identity(record)) + + +def _serialize_logical_extent(logical_extent): + if isinstance(logical_extent, LogicalSlice): + return { + "kind": "logical_slice", + "flat_offset": logical_extent.flat_offset, + "length": logical_extent.length, + } + if isinstance(logical_extent, LogicalRegion): + return { + "kind": "logical_region", + "offsets": list(logical_extent.offsets), + "lengths": list(logical_extent.lengths), + } + raise TypeError("logical extent must be a LogicalSlice or LogicalRegion") + + +def _parse_logical_extent(record): + if type(record) is not dict or type(record.get("kind")) is not str: + raise ValueError("shard logical_extent has an invalid schema") + kind = record["kind"] + if kind == "logical_slice": + _require_exact_record(record, _LOGICAL_SLICE_KEYS, name="shard logical_slice") + if type(record["flat_offset"]) is not int or type(record["length"]) is not int: + raise ValueError("shard logical_slice offsets and lengths must be integers") + try: + return LogicalSlice(record["flat_offset"], record["length"]) + except (TypeError, ValueError) as exc: + raise ValueError("shard logical_slice is invalid") from exc + if kind == "logical_region": + _require_exact_record(record, _LOGICAL_REGION_KEYS, name="shard logical_region") + offsets = _require_exact_list(record["offsets"], name="shard logical_region offsets") + lengths = _require_exact_list(record["lengths"], name="shard logical_region lengths") + if any(type(value) is not int for value in offsets + lengths): + raise ValueError("shard logical_region offsets and lengths must contain only integers") + try: + return LogicalRegion(tuple(offsets), tuple(lengths)) + except (TypeError, ValueError) as exc: + raise ValueError("shard logical_region is invalid") from exc + raise ValueError("unsupported shard logical_extent kind") + + +def _serialize_shard_placement(placement): + if not isinstance(placement, ShardPlacement): + raise TypeError("placement must be a ShardPlacement") + return { + "mesh_axis": placement.mesh_axis, + "kind": placement.kind.value, + "coordinate": placement.coordinate, + "parts": placement.parts, + "parameter_dimension": placement.parameter_dimension, + } + + +def _parse_shard_placement(record): + record = _require_exact_record(record, _SHARD_PLACEMENT_KEYS, name="shard placement") + if type(record["mesh_axis"]) is not str: + raise ValueError("shard placement mesh_axis must be a string") + if type(record["kind"]) is not str: + raise ValueError("shard placement kind must be a string") + if type(record["coordinate"]) is not int: + raise ValueError("shard placement coordinate must be an int") + if type(record["parts"]) is not int: + raise ValueError("shard placement parts must be an int") + if record["parameter_dimension"] is not None and type(record["parameter_dimension"]) is not int: + raise ValueError("shard placement parameter_dimension must be an int or None") + try: + kind = PlacementKind(record["kind"]) + return ShardPlacement( + mesh_axis=record["mesh_axis"], + kind=kind, + coordinate=record["coordinate"], + parts=record["parts"], + parameter_dimension=record["parameter_dimension"], + ) + except (TypeError, ValueError) as exc: + raise ValueError("shard placement is invalid") from exc + + +def _serialize_shard_identity(identity): + if not isinstance(identity, ShardIdentity): + raise TypeError("identity must be a ShardIdentity") + return { + "schema_version": identity.schema_version, + "parameter": _serialize_parameter_identity(identity.parameter), + "layout": identity.layout.value, + "logical_extent": _serialize_logical_extent(identity.logical_slice), + "placements": [_serialize_shard_placement(placement) for placement in identity.placements], + "process_group": ( + None if identity.process_group is None else _serialize_process_group_identity(identity.process_group) + ), + "local_member": identity.local_member, + "owner": identity.owner, + } + + +def _parse_shard_identity(record): + record = _require_exact_record(record, _SHARD_IDENTITY_KEYS, name="shard identity") + if type(record["schema_version"]) is not int: + raise ValueError("shard identity schema_version must be an int") + if type(record["layout"]) is not str: + raise ValueError("shard identity layout must be a string") + placements = _require_exact_list(record["placements"], name="shard identity placements") + _require_optional_string(record["local_member"], name="shard identity local_member") + _require_optional_string(record["owner"], name="shard identity owner") + try: + parameter = _parse_parameter_identity(record["parameter"]) + logical_extent = _parse_logical_extent(record["logical_extent"]) + parsed_placements = tuple(_parse_shard_placement(placement) for placement in placements) + process_group = ( + None if record["process_group"] is None else _parse_process_group_identity(record["process_group"]) + ) + layout = ParameterLayout(record["layout"]) + return ShardIdentity( + parameter=parameter, + layout=layout, + logical_slice=logical_extent, + placements=parsed_placements, + process_group=process_group, + local_member=record["local_member"], + owner=record["owner"], + schema_version=record["schema_version"], + ) + except (TypeError, ValueError) as exc: + raise ValueError("shard identity is invalid") from exc + + +def _normalize_shard_identity(record): + return _serialize_shard_identity(_parse_shard_identity(record)) + + +def _serialize_sharding_manifest(manifest): + if not isinstance(manifest, ShardingManifest): + raise TypeError("manifest must be a ShardingManifest") + return { + "schema_version": manifest.schema_version, + "shards": [_serialize_shard_identity(shard) for shard in manifest.shards], + } + + +def _parse_sharding_manifest(record): + record = _require_exact_record(record, _SHARDING_MANIFEST_KEYS, name="sharding manifest") + if type(record["schema_version"]) is not int: + raise ValueError("sharding manifest schema_version must be an int") + shards = _require_exact_list(record["shards"], name="sharding manifest shards") + try: + return ShardingManifest( + shards=tuple(_parse_shard_identity(shard) for shard in shards), + schema_version=record["schema_version"], + ) + except (TypeError, ValueError) as exc: + raise ValueError("sharding manifest is invalid") from exc + + +def _normalize_sharding_manifest(record): + return _serialize_sharding_manifest(_parse_sharding_manifest(record)) diff --git a/src/gefen/portable_wire.py b/src/gefen/portable_wire.py new file mode 100644 index 0000000..9c4a02e --- /dev/null +++ b/src/gefen/portable_wire.py @@ -0,0 +1,747 @@ +"""Bounded binary wire codec for canonical optimizer-state fragments.""" + +from __future__ import annotations + +from dataclasses import dataclass +import hashlib +import hmac +import math +import struct +import sys + +import torch + + +_CANONICAL_WIRE_PREAMBLE = b"GFNCV1\0\0" +_CANONICAL_WIRE_CODEC = 1 +_CANONICAL_WIRE_TENSOR_DOMAIN = b"gefen.canonical_wire.tensor.v1\0" +_CANONICAL_WIRE_FRAGMENT_DOMAIN = b"gefen.canonical_wire.fragment.v1\0" +_SHA256_BYTES = hashlib.sha256().digest_size +_U16_MAX = (1 << 16) - 1 +_I64_MAX = (1 << 63) - 1 +_U64_MAX = (1 << 64) - 1 +_MAX_SAFE_TREE_DEPTH = 256 + +_TAG_NONE = 0x00 +_TAG_FALSE = 0x01 +_TAG_TRUE = 0x02 +_TAG_INTEGER = 0x03 +_TAG_FLOAT = 0x04 +_TAG_STRING = 0x05 +_TAG_LIST = 0x06 +_TAG_TUPLE = 0x07 +_TAG_DICTIONARY = 0x08 +_TAG_TENSOR = 0x09 + +_DTYPE_RECORDS = ( + (1, torch.bool, 1), + (2, torch.uint8, 1), + (3, torch.int8, 1), + (4, torch.int16, 2), + (5, torch.int32, 4), + (6, torch.int64, 8), + (7, torch.float16, 2), + (8, torch.bfloat16, 2), + (9, torch.float32, 4), + (10, torch.float64, 8), + (11, torch.complex64, 8), + (12, torch.complex128, 16), +) +_DTYPE_BY_CODE = {code: (dtype, size) for code, dtype, size in _DTYPE_RECORDS} +_DTYPE_BY_VALUE = {dtype: (code, size) for code, dtype, size in _DTYPE_RECORDS} + + +def _require_positive_limit(name: str, value: int, *, maximum: int = _U64_MAX) -> None: + if type(value) is not int: + raise TypeError("{} must be an int".format(name)) + if value <= 0 or value > maximum: + raise ValueError("{} must be in [1, {}]".format(name, maximum)) + + +@dataclass(frozen=True, slots=True) +class _CanonicalWireLimits: + max_fragment_tensor_bytes: int + max_collective_tensor_bytes: int + max_collective_metadata_bytes: int + chunk_bytes: int = 8 << 20 + max_members: int = 4096 + max_metadata_bytes: int = 64 << 20 + max_tree_nodes: int = 1_000_000 + max_tree_depth: int = 64 + max_container_items: int = 1_000_000 + max_string_bytes: int = 1 << 20 + max_integer_bytes: int = 4096 + max_tensors: int = 262_144 + max_tensor_rank: int = 64 + diagnostic_bytes: int = 2048 + + def __post_init__(self) -> None: + for name in ( + "max_fragment_tensor_bytes", + "max_collective_tensor_bytes", + "max_collective_metadata_bytes", + "chunk_bytes", + "max_members", + "max_metadata_bytes", + "max_tree_nodes", + "max_tree_depth", + "max_container_items", + "max_string_bytes", + "max_integer_bytes", + "max_tensors", + "diagnostic_bytes", + ): + _require_positive_limit(name, getattr(self, name)) + _require_positive_limit("max_tensor_rank", self.max_tensor_rank, maximum=_U16_MAX) + if self.max_tree_depth > _MAX_SAFE_TREE_DEPTH: + raise ValueError( + "max_tree_depth cannot exceed {} for recursive parsing".format( + _MAX_SAFE_TREE_DEPTH + ) + ) + if self.max_fragment_tensor_bytes > self.max_collective_tensor_bytes: + raise ValueError("max_fragment_tensor_bytes cannot exceed max_collective_tensor_bytes") + if self.max_metadata_bytes > self.max_collective_metadata_bytes: + raise ValueError("max_metadata_bytes cannot exceed max_collective_metadata_bytes") + + +@dataclass(frozen=True, slots=True) +class _CanonicalWireTensorSpec: + index: int + dtype_code: int + dtype: torch.dtype + shape: tuple[int, ...] + numel: int + nbytes: int + digest: bytes + + def __post_init__(self) -> None: + if type(self.index) is not int or self.index < 0 or self.index > _U64_MAX: + raise ValueError("tensor index must be a nonnegative int") + if type(self.dtype_code) is not int: + raise TypeError("tensor dtype code must be an int") + record = _DTYPE_BY_CODE.get(self.dtype_code) + if record is None or record[0] != self.dtype: + raise ValueError("tensor dtype code and dtype do not agree") + if ( + type(self.shape) is not tuple + or len(self.shape) > _U16_MAX + or any( + type(dimension) is not int + or dimension < 0 + or dimension > _I64_MAX + for dimension in self.shape + ) + ): + raise ValueError("tensor shape must be a tuple of nonnegative ints") + if type(self.numel) is not int or self.numel < 0 or self.numel > _U64_MAX: + raise ValueError("tensor numel must be a nonnegative int") + if type(self.nbytes) is not int or self.nbytes < 0 or self.nbytes > _U64_MAX: + raise ValueError("tensor nbytes must be a nonnegative int") + if type(self.digest) is not bytes or len(self.digest) != _SHA256_BYTES: + raise ValueError("tensor digest must be a 32-byte value") + if _shape_numel(self.shape) != self.numel or self.numel * record[1] != self.nbytes: + raise ValueError("tensor shape, numel, and nbytes do not agree") + + +@dataclass(frozen=True, slots=True) +class _CanonicalWirePlan: + metadata: bytes + tensors: tuple[torch.Tensor, ...] + tensor_specs: tuple[_CanonicalWireTensorSpec, ...] + metadata_digest: bytes + fragment_digest: bytes + total_tensor_bytes: int + + def __post_init__(self) -> None: + if type(self.metadata) is not bytes: + raise TypeError("canonical wire metadata must be bytes") + if type(self.tensors) is not tuple or any(type(value) is not torch.Tensor for value in self.tensors): + raise TypeError("canonical wire tensors must be a tuple of plain tensors") + if type(self.tensor_specs) is not tuple or any(type(value) is not _CanonicalWireTensorSpec for value in self.tensor_specs): + raise TypeError("canonical wire tensor specs must be a tuple of tensor specs") + if len(self.tensors) != len(self.tensor_specs): + raise ValueError("canonical wire tensors and specs must have equal lengths") + if any(spec.index != index for index, spec in enumerate(self.tensor_specs)): + raise ValueError("canonical wire tensor specs must have consecutive indices") + if any(not _tight_cpu_tensor(tensor, spec) for tensor, spec in zip(self.tensors, self.tensor_specs)): + raise ValueError("canonical wire tensors must be tight CPU tensors matching their specs") + _validate_digest_field("metadata_digest", self.metadata_digest) + _validate_digest_field("fragment_digest", self.fragment_digest) + if type(self.total_tensor_bytes) is not int or self.total_tensor_bytes < 0: + raise ValueError("total_tensor_bytes must be a nonnegative int") + if sum(spec.nbytes for spec in self.tensor_specs) != self.total_tensor_bytes: + raise ValueError("total_tensor_bytes does not match the tensor specs") + if not hmac.compare_digest(hashlib.sha256(self.metadata).digest(), self.metadata_digest): + raise ValueError("metadata_digest does not match canonical wire metadata") + if not hmac.compare_digest(_canonical_fragment_digest(self.metadata_digest, self.tensor_specs), self.fragment_digest): + raise ValueError("fragment_digest does not match canonical wire metadata and tensor specs") + + @property + def payload_tensors(self) -> tuple[torch.Tensor, ...]: + return self.tensors + + +@dataclass(frozen=True, slots=True) +class _TensorReference: + index: int + + +@dataclass(frozen=True, slots=True) +class _ListValue: + values: tuple[object, ...] + + +@dataclass(frozen=True, slots=True) +class _TupleValue: + values: tuple[object, ...] + + +@dataclass(frozen=True, slots=True) +class _DictionaryValue: + values: tuple[tuple[str, object], ...] + + +@dataclass(frozen=True, slots=True) +class _PreparedCanonicalWireValue: + tensor_specs: tuple[_CanonicalWireTensorSpec, ...] + metadata_digest: bytes + fragment_digest: bytes + total_tensor_bytes: int + _value: object + _limits: _CanonicalWireLimits + + def __post_init__(self) -> None: + if type(self.tensor_specs) is not tuple or any(type(value) is not _CanonicalWireTensorSpec for value in self.tensor_specs): + raise TypeError("canonical wire tensor specs must be a tuple of tensor specs") + if any(spec.index != index for index, spec in enumerate(self.tensor_specs)): + raise ValueError("canonical wire tensor specs must have consecutive indices") + _validate_digest_field("metadata_digest", self.metadata_digest) + _validate_digest_field("fragment_digest", self.fragment_digest) + if type(self.total_tensor_bytes) is not int or self.total_tensor_bytes < 0: + raise ValueError("total_tensor_bytes must be a nonnegative int") + if sum(spec.nbytes for spec in self.tensor_specs) != self.total_tensor_bytes: + raise ValueError("total_tensor_bytes does not match the tensor specs") + if not hmac.compare_digest(_canonical_fragment_digest(self.metadata_digest, self.tensor_specs), self.fragment_digest): + raise ValueError("fragment_digest does not match canonical wire metadata and tensor specs") + if type(self._limits) is not _CanonicalWireLimits: + raise TypeError("prepared canonical wire values require canonical wire limits") + + +def _validate_digest_field(name: str, value: bytes) -> None: + if type(value) is not bytes: + raise TypeError("{} must be bytes".format(name)) + if len(value) != _SHA256_BYTES: + raise ValueError("{} must be a 32-byte value".format(name)) + + +def _shape_numel(shape: tuple[int, ...]) -> int: + result = 1 + for dimension in shape: + if dimension != 0 and result > _U64_MAX // dimension: + raise ValueError("tensor shape product overflows u64") + result *= dimension + return result + + +class _MetadataWriter: + __slots__ = ("_data", "_limit") + + def __init__(self, limit: int): + self._data = bytearray() + self._limit = limit + + def write(self, value: bytes) -> None: + if len(self._data) + len(value) > self._limit: + raise ValueError("canonical wire metadata exceeds max_metadata_bytes") + self._data.extend(value) + + def finish(self) -> bytes: + return bytes(self._data) + + +class _PreparationState: + __slots__ = ("limits", "writer", "nodes", "tensors", "specs", "tensor_bytes") + + def __init__(self, limits: _CanonicalWireLimits): + self.limits = limits + self.writer = _MetadataWriter(limits.max_metadata_bytes) + self.nodes = 0 + self.tensors: list[torch.Tensor] = [] + self.specs: list[_CanonicalWireTensorSpec] = [] + self.tensor_bytes = 0 + + def visit(self, value, *, depth: int, path: str) -> None: + if depth > self.limits.max_tree_depth: + raise ValueError("canonical wire value exceeds max_tree_depth at {}".format(path)) + self.nodes += 1 + if self.nodes > self.limits.max_tree_nodes: + raise ValueError("canonical wire value exceeds max_tree_nodes") + value_type = type(value) + if value is None: + self.writer.write(bytes((_TAG_NONE,))) + return + if value_type is bool: + self.writer.write(bytes((_TAG_TRUE if value else _TAG_FALSE,))) + return + if value_type is int: + magnitude = abs(value) + magnitude_size = (magnitude.bit_length() + 7) // 8 + if magnitude_size > self.limits.max_integer_bytes: + raise ValueError("integer exceeds max_integer_bytes at {}".format(path)) + magnitude_bytes = magnitude.to_bytes(magnitude_size, "big") + self.writer.write(bytes((_TAG_INTEGER, 1 if value < 0 else 0))) + self.writer.write(struct.pack(">Q", len(magnitude_bytes))) + self.writer.write(magnitude_bytes) + return + if value_type is float: + if not math.isfinite(value): + raise ValueError("float must be finite at {}".format(path)) + self.writer.write(bytes((_TAG_FLOAT,))) + self.writer.write(struct.pack(">d", value)) + return + if value_type is str: + self._write_string(value, path=path) + return + if torch.is_tensor(value): + self._write_tensor(value, path=path) + return + if value_type is list or value_type is tuple: + if len(value) > self.limits.max_container_items: + raise ValueError("container exceeds max_container_items at {}".format(path)) + self.writer.write(bytes((_TAG_LIST if value_type is list else _TAG_TUPLE,))) + self.writer.write(struct.pack(">Q", len(value))) + for index, item in enumerate(value): + self.visit(item, depth=depth + 1, path="{}[{}]".format(path, index)) + return + if value_type is dict: + if len(value) > self.limits.max_container_items: + raise ValueError("container exceeds max_container_items at {}".format(path)) + if any(type(key) is not str for key in value): + raise TypeError("dictionary keys must be strings at {}".format(path)) + self.writer.write(bytes((_TAG_DICTIONARY,))) + self.writer.write(struct.pack(">Q", len(value))) + for key in sorted(value): + self.visit(key, depth=depth + 1, path="{}.".format(path)) + self.visit(value[key], depth=depth + 1, path="{}.{}".format(path, key)) + return + raise TypeError("{} has unsupported canonical wire type {}".format(path, value_type.__name__)) + + def _write_string(self, value: str, *, path: str) -> None: + encoded = _bounded_utf8(value, limit=self.limits.max_string_bytes, path=path) + self.writer.write(bytes((_TAG_STRING,))) + self.writer.write(struct.pack(">Q", len(encoded))) + self.writer.write(encoded) + + def _write_tensor(self, value: torch.Tensor, *, path: str) -> None: + if type(value) is not torch.Tensor: + raise TypeError("tensor must be a plain torch.Tensor at {}".format(path)) + if value.layout is not torch.strided or value.is_meta or value.is_nested or value.is_quantized: + raise TypeError("tensor must be a materialized strided tensor at {}".format(path)) + dtype_record = _DTYPE_BY_VALUE.get(value.dtype) + if dtype_record is None: + raise TypeError("tensor dtype {} is not supported at {}".format(value.dtype, path)) + shape = tuple(value.shape) + if len(shape) > self.limits.max_tensor_rank: + raise ValueError("tensor rank exceeds max_tensor_rank at {}".format(path)) + if len(self.tensors) >= self.limits.max_tensors: + raise ValueError("canonical wire value exceeds max_tensors") + dtype_code, element_size = dtype_record + numel = value.numel() + nbytes = numel * element_size + if nbytes > _U64_MAX or self.tensor_bytes + nbytes > self.limits.max_fragment_tensor_bytes: + raise ValueError("canonical wire tensors exceed max_fragment_tensor_bytes") + index = len(self.tensors) + cloned = _clone_tensor(value, chunk_bytes=self.limits.chunk_bytes, path=path) + digest = _canonical_tensor_digest( + cloned, + dtype_code=dtype_code, + shape=shape, + nbytes=nbytes, + chunk_bytes=self.limits.chunk_bytes, + ) + spec = _CanonicalWireTensorSpec(index, dtype_code, value.dtype, shape, numel, nbytes, digest) + self.tensors.append(cloned) + self.specs.append(spec) + self.tensor_bytes += nbytes + self.writer.write(bytes((_TAG_TENSOR,))) + self.writer.write(struct.pack(">Q", index)) + + +def _bounded_utf8(value: str, *, limit: int, path: str) -> bytes: + encoded = bytearray() + characters_per_chunk = max(1, min(1 << 14, limit // 4)) + for start in range(0, len(value), characters_per_chunk): + chunk = value[start : start + characters_per_chunk].encode("utf-8") + if len(encoded) + len(chunk) > limit: + raise ValueError("string exceeds max_string_bytes at {}".format(path)) + encoded.extend(chunk) + return bytes(encoded) + + +def _read_tensor_chunk(value: torch.Tensor, start: int, stop: int) -> torch.Tensor: + detached = value.detach() + if detached.is_contiguous(): + return detached.reshape(-1)[start:stop] + linear = torch.arange(start, stop, dtype=torch.int64, device=detached.device) + remainder = linear + reversed_coordinates = [] + for dimension in reversed(detached.shape): + reversed_coordinates.append(torch.remainder(remainder, dimension)) + remainder = torch.div(remainder, dimension, rounding_mode="floor") + return detached[tuple(reversed(reversed_coordinates))] + + +def _clone_tensor(value: torch.Tensor, *, chunk_bytes: int, path: str) -> torch.Tensor: + cloned = torch.empty(tuple(value.shape), dtype=value.dtype, device="cpu") + flat = cloned.reshape(-1) + scratch_bytes_per_element = value.element_size() + if not value.is_contiguous(): + # Logical indexing retains one int64 coordinate vector per dimension, + # plus the linear index and current quotient. Keep that scratch within + # the same caller-supplied chunk budget as the tensor copy. + scratch_bytes_per_element += 8 * (value.ndim + 2) + chunk_elements = max(1, chunk_bytes // scratch_bytes_per_element) + for start in range(0, value.numel(), chunk_elements): + stop = min(start + chunk_elements, value.numel()) + source = _read_tensor_chunk(value, start, stop).resolve_conj().resolve_neg().reshape(-1) + destination = flat[start:stop] + destination.copy_(source) + if cloned.is_floating_point() or cloned.is_complex(): + try: + finite = bool(torch.isfinite(destination).all()) + except (NotImplementedError, RuntimeError, TypeError) as exc: + raise ValueError("tensor dtype does not support finite values at {}".format(path)) from exc + if not finite: + raise ValueError("tensor must be finite at {}".format(path)) + return cloned + + +def _canonical_tensor_digest( + value: torch.Tensor, + *, + dtype_code: int, + shape: tuple[int, ...], + nbytes: int, + chunk_bytes: int, +) -> bytes: + hasher = hashlib.sha256() + hasher.update(_CANONICAL_WIRE_TENSOR_DOMAIN) + hasher.update(struct.pack(">HH", dtype_code, len(shape))) + for dimension in shape: + hasher.update(struct.pack(">Q", dimension)) + hasher.update(struct.pack(">Q", nbytes)) + element_size = value.element_size() + component_size = element_size // 2 if value.is_complex() else element_size + chunk_elements = max(1, chunk_bytes // element_size) + flat = value.reshape(-1) + for start in range(0, value.numel(), chunk_elements): + stop = min(start + chunk_elements, value.numel()) + raw = flat[start:stop].view(torch.uint8).reshape(-1) + if sys.byteorder == "big" and component_size > 1: + raw = raw.reshape(-1, component_size).flip(1).contiguous().reshape(-1) + elif sys.byteorder != "little": + raise RuntimeError("unsupported native byte order") + hasher.update(memoryview(raw.numpy())) + return hasher.digest() + + +def _canonical_fragment_digest(metadata_digest: bytes, specs: tuple[_CanonicalWireTensorSpec, ...]) -> bytes: + hasher = hashlib.sha256() + hasher.update(_CANONICAL_WIRE_FRAGMENT_DOMAIN) + hasher.update(metadata_digest) + hasher.update(struct.pack(">Q", len(specs))) + for spec in specs: + hasher.update(struct.pack(">Q", spec.nbytes)) + hasher.update(spec.digest) + return hasher.digest() + + +def _prepare_canonical_wire_value(value, limits: _CanonicalWireLimits) -> _CanonicalWirePlan: + """Clone and encode one canonical value into metadata and tensor payloads.""" + + if type(limits) is not _CanonicalWireLimits: + raise TypeError("limits must be _CanonicalWireLimits") + state = _PreparationState(limits) + state.writer.write(_CANONICAL_WIRE_PREAMBLE) + state.writer.write(struct.pack(">HHI", _CANONICAL_WIRE_CODEC, 0, 0)) + state.visit(value, depth=0, path="value") + state.writer.write(struct.pack(">Q", len(state.specs))) + for spec in state.specs: + state.writer.write(struct.pack(">HHI", spec.dtype_code, len(spec.shape), 0)) + for dimension in spec.shape: + state.writer.write(struct.pack(">Q", dimension)) + state.writer.write(struct.pack(">QQ", spec.numel, spec.nbytes)) + state.writer.write(spec.digest) + metadata = state.writer.finish() + specs = tuple(state.specs) + metadata_digest = hashlib.sha256(metadata).digest() + fragment_digest = _canonical_fragment_digest(metadata_digest, specs) + return _CanonicalWirePlan( + metadata=metadata, + tensors=tuple(state.tensors), + tensor_specs=specs, + metadata_digest=metadata_digest, + fragment_digest=fragment_digest, + total_tensor_bytes=state.tensor_bytes, + ) + + +class _MetadataReader: + __slots__ = ("metadata", "limits", "position", "nodes", "next_tensor") + + def __init__(self, metadata: bytes, limits: _CanonicalWireLimits): + self.metadata = metadata + self.limits = limits + self.position = 0 + self.nodes = 0 + self.next_tensor = 0 + + def read(self, size: int) -> bytes: + if size < 0 or size > len(self.metadata) - self.position: + raise ValueError("canonical wire metadata is truncated") + start = self.position + self.position += size + return self.metadata[start : self.position] + + def u8(self) -> int: + return self.read(1)[0] + + def u16(self) -> int: + return struct.unpack(">H", self.read(2))[0] + + def u32(self) -> int: + return struct.unpack(">I", self.read(4))[0] + + def u64(self) -> int: + return struct.unpack(">Q", self.read(8))[0] + + def value(self, *, depth: int, dictionary_key: bool = False): + if depth > self.limits.max_tree_depth: + raise ValueError("canonical wire metadata exceeds max_tree_depth") + self.nodes += 1 + if self.nodes > self.limits.max_tree_nodes: + raise ValueError("canonical wire metadata exceeds max_tree_nodes") + tag = self.u8() + if tag == _TAG_NONE: + result = None + elif tag == _TAG_FALSE: + result = False + elif tag == _TAG_TRUE: + result = True + elif tag == _TAG_INTEGER: + sign = self.u8() + if sign not in {0, 1}: + raise ValueError("canonical wire integer has an invalid sign") + size = self.u64() + if size > self.limits.max_integer_bytes: + raise ValueError("canonical wire integer exceeds max_integer_bytes") + encoded = self.read(size) + if encoded[:1] == b"\0": + raise ValueError("canonical wire integer magnitude is not canonical") + magnitude = int.from_bytes(encoded, "big") + if sign == 1 and magnitude == 0: + raise ValueError("canonical wire integer encodes negative zero") + result = -magnitude if sign else magnitude + elif tag == _TAG_FLOAT: + result = struct.unpack(">d", self.read(8))[0] + if not math.isfinite(result): + raise ValueError("canonical wire float must be finite") + elif tag == _TAG_STRING: + result = self.string_after_tag() + elif tag == _TAG_LIST or tag == _TAG_TUPLE: + count = self.u64() + if count > self.limits.max_container_items: + raise ValueError("canonical wire container exceeds max_container_items") + values = tuple(self.value(depth=depth + 1) for _ in range(count)) + result = _ListValue(values) if tag == _TAG_LIST else _TupleValue(values) + elif tag == _TAG_DICTIONARY: + count = self.u64() + if count > self.limits.max_container_items: + raise ValueError("canonical wire dictionary exceeds max_container_items") + values = [] + previous = None + for _ in range(count): + key = self.value(depth=depth + 1, dictionary_key=True) + if type(key) is not str: + raise ValueError("canonical wire dictionary key must be a string") + if previous is not None and key <= previous: + raise ValueError("canonical wire dictionary keys are not strictly increasing") + previous = key + values.append((key, self.value(depth=depth + 1))) + result = _DictionaryValue(tuple(values)) + elif tag == _TAG_TENSOR: + index = self.u64() + if index != self.next_tensor: + raise ValueError("canonical wire tensor references are not ordered") + if index >= self.limits.max_tensors: + raise ValueError("canonical wire metadata exceeds max_tensors") + self.next_tensor += 1 + result = _TensorReference(index) + else: + raise ValueError("canonical wire metadata has an unknown value tag") + if dictionary_key and type(result) is not str: + raise ValueError("canonical wire dictionary key must be a string") + return result + + def string_after_tag(self) -> str: + size = self.u64() + if size > self.limits.max_string_bytes: + raise ValueError("canonical wire string exceeds max_string_bytes") + encoded = self.read(size) + try: + result = encoded.decode("utf-8") + except UnicodeDecodeError as exc: + raise ValueError("canonical wire string is not valid UTF-8") from exc + if result.encode("utf-8") != encoded: + raise ValueError("canonical wire string is not canonically encoded") + return result + + +def _parse_canonical_wire_metadata( + metadata: bytes, + *, + limits: _CanonicalWireLimits, +) -> _PreparedCanonicalWireValue: + """Validate bounded metadata without allocating any tensor payloads.""" + + if type(metadata) is not bytes: + raise TypeError("metadata must be bytes") + if type(limits) is not _CanonicalWireLimits: + raise TypeError("limits must be _CanonicalWireLimits") + if len(metadata) > limits.max_metadata_bytes: + raise ValueError("canonical wire metadata exceeds max_metadata_bytes") + reader = _MetadataReader(metadata, limits) + if reader.read(len(_CANONICAL_WIRE_PREAMBLE)) != _CANONICAL_WIRE_PREAMBLE: + raise ValueError("canonical wire metadata has an invalid preamble") + if reader.u16() != _CANONICAL_WIRE_CODEC: + raise ValueError("canonical wire metadata uses an unsupported codec") + if reader.u16() != 0 or reader.u32() != 0: + raise ValueError("canonical wire metadata has nonzero reserved header fields") + value = reader.value(depth=0) + descriptor_count = reader.u64() + if descriptor_count != reader.next_tensor: + raise ValueError("canonical wire descriptor count does not match tensor references") + if descriptor_count > limits.max_tensors: + raise ValueError("canonical wire metadata exceeds max_tensors") + specs = [] + total_tensor_bytes = 0 + for index in range(descriptor_count): + dtype_code = reader.u16() + dtype_record = _DTYPE_BY_CODE.get(dtype_code) + if dtype_record is None: + raise ValueError("canonical wire tensor descriptor has an unknown dtype") + rank = reader.u16() + if rank > limits.max_tensor_rank: + raise ValueError("canonical wire tensor rank exceeds max_tensor_rank") + if reader.u32() != 0: + raise ValueError("canonical wire tensor descriptor has nonzero flags") + shape = tuple(reader.u64() for _ in range(rank)) + if any(dimension > _I64_MAX for dimension in shape): + raise ValueError("canonical wire tensor dimension exceeds signed int64") + numel = reader.u64() + nbytes = reader.u64() + digest = reader.read(_SHA256_BYTES) + expected_numel = 1 + for dimension in shape: + if dimension != 0 and expected_numel > _U64_MAX // dimension: + raise ValueError("canonical wire tensor shape product overflows u64") + expected_numel *= dimension + if expected_numel != numel: + raise ValueError("canonical wire tensor shape does not match numel") + dtype, element_size = dtype_record + if numel != 0 and numel > _U64_MAX // element_size: + raise ValueError("canonical wire tensor byte count overflows u64") + if numel * element_size != nbytes: + raise ValueError("canonical wire tensor numel does not match nbytes") + if total_tensor_bytes + nbytes > limits.max_fragment_tensor_bytes: + raise ValueError("canonical wire tensors exceed max_fragment_tensor_bytes") + specs.append(_CanonicalWireTensorSpec(index, dtype_code, dtype, shape, numel, nbytes, digest)) + total_tensor_bytes += nbytes + if reader.position != len(metadata): + raise ValueError("canonical wire metadata has trailing bytes") + specs_tuple = tuple(specs) + metadata_digest = hashlib.sha256(metadata).digest() + fragment_digest = _canonical_fragment_digest(metadata_digest, specs_tuple) + return _PreparedCanonicalWireValue( + tensor_specs=specs_tuple, + metadata_digest=metadata_digest, + fragment_digest=fragment_digest, + total_tensor_bytes=total_tensor_bytes, + _value=value, + _limits=limits, + ) + + +def _tight_cpu_tensor(value: torch.Tensor, spec: _CanonicalWireTensorSpec) -> bool: + if ( + type(value) is not torch.Tensor + or value.layout is not torch.strided + or value.device.type != "cpu" + or value.is_meta + or value.is_nested + or value.is_quantized + or value.dtype != spec.dtype + or tuple(value.shape) != spec.shape + or value.numel() != spec.numel + or value.element_size() * value.numel() != spec.nbytes + or not value.is_contiguous() + or value.is_conj() + or value.is_neg() + or value.storage_offset() != 0 + ): + return False + try: + return value.untyped_storage().nbytes() == spec.nbytes + except (AttributeError, RuntimeError): + return False + + +def _materialize_value(value, tensors: tuple[torch.Tensor, ...]): + if type(value) is _TensorReference: + return tensors[value.index] + if type(value) is _ListValue: + return [_materialize_value(item, tensors) for item in value.values] + if type(value) is _TupleValue: + return tuple(_materialize_value(item, tensors) for item in value.values) + if type(value) is _DictionaryValue: + return {key: _materialize_value(item, tensors) for key, item in value.values} + return value + + +def _reconstruct_canonical_wire_value( + prepared: _PreparedCanonicalWireValue, + payload_tensors, + *, + expected_fragment_digest: bytes | None = None, +): + """Verify separate tensor payloads and reconstruct one owned canonical value.""" + + if type(prepared) is not _PreparedCanonicalWireValue: + raise TypeError("prepared must be _PreparedCanonicalWireValue") + if expected_fragment_digest is not None: + _validate_digest_field("expected_fragment_digest", expected_fragment_digest) + if not hmac.compare_digest(prepared.fragment_digest, expected_fragment_digest): + raise ValueError("canonical wire fragment digest does not match") + if type(payload_tensors) not in {tuple, list}: + raise TypeError("payload_tensors must be a tuple or list") + if len(payload_tensors) != len(prepared.tensor_specs): + raise ValueError("canonical wire payload count does not match tensor descriptors") + payloads = tuple(payload_tensors) + for index, (payload, spec) in enumerate(zip(payloads, prepared.tensor_specs)): + if not torch.is_tensor(payload) or not _tight_cpu_tensor(payload, spec): + raise ValueError("canonical wire tensor payload {} does not match its descriptor".format(index)) + cloned = [] + for index, (payload, spec) in enumerate(zip(payloads, prepared.tensor_specs)): + owned = _clone_tensor(payload, chunk_bytes=prepared._limits.chunk_bytes, path="payload_tensors[{}]".format(index)) + digest = _canonical_tensor_digest( + owned, + dtype_code=spec.dtype_code, + shape=spec.shape, + nbytes=spec.nbytes, + chunk_bytes=prepared._limits.chunk_bytes, + ) + if not hmac.compare_digest(digest, spec.digest): + raise ValueError("canonical wire tensor payload {} has an invalid digest".format(index)) + cloned.append(owned) + return _materialize_value(prepared._value, tuple(cloned)) diff --git a/tests/test_portable_fields.py b/tests/test_portable_fields.py new file mode 100644 index 0000000..abc15d8 --- /dev/null +++ b/tests/test_portable_fields.py @@ -0,0 +1,538 @@ +"""CPU coverage for pure dense logical-field assembly and projection.""" + +import pytest +import torch + +from gefen.contracts import ( + LogicalRegion, + LogicalSlice, + ParameterIdentity, + ParameterLayout, + PlacementKind, + ProcessGroupIdentity, + ShardIdentity, + ShardPlacement, + ShardingManifest, +) +from gefen.portable_fields import ( + _assemble_dense_logical_field, + _project_dense_logical_field, +) + + +def _group(members=("worker:c", "worker:a", "worker:b")): + return ProcessGroupIdentity("data_parallel", members) + + +def _placement(group, member, kind, *, dimension=None): + return ShardPlacement( + "data_parallel", + kind, + group.ordered_members.index(member), + len(group.ordered_members), + dimension, + ) + + +def _local_replicated(parameter): + return ( + ShardIdentity( + parameter, + ParameterLayout.REPLICATED, + LogicalSlice.full(parameter), + ), + ) + + +def _grouped_replicated(parameter, group): + return tuple( + ShardIdentity( + parameter, + ParameterLayout.REPLICATED, + LogicalSlice.full(parameter), + placements=(_placement(group, member, PlacementKind.REPLICATE),), + process_group=group, + local_member=member, + ) + for member in group.ordered_members + ) + + +def _flat_shards(parameter, group, lengths): + offset = 0 + result = [] + for member, length in zip(group.ordered_members, lengths): + result.append( + ShardIdentity( + parameter, + ParameterLayout.FLATTENED_ELEMENT_SHARD, + LogicalSlice(offset, length), + placements=(_placement(group, member, PlacementKind.FLAT_SHARD),), + process_group=group, + local_member=member, + ) + ) + offset += length + return tuple(result) + + +def _owner_shards(parameter, group, owner): + return tuple( + ShardIdentity( + parameter, + ParameterLayout.WHOLE_PARAMETER_OWNER, + (LogicalSlice.full(parameter) if member == owner else LogicalSlice(0, 0)), + placements=( + _placement( + group, + member, + PlacementKind.WHOLE_PARAMETER_OWNER, + ), + ), + process_group=group, + local_member=member, + owner=owner, + ) + for member in group.ordered_members + ) + + +def _dtensor_shards(parameter, group, lengths, *, dimension): + offset = 0 + result = [] + for member, length in zip(group.ordered_members, lengths): + offsets = [0] * len(parameter.global_shape) + region_lengths = list(parameter.global_shape) + offsets[dimension] = offset + region_lengths[dimension] = length + result.append( + ShardIdentity( + parameter, + ParameterLayout.DTENSOR_1D_DEFAULT_WORLD, + LogicalRegion(offsets, region_lengths), + placements=( + _placement( + group, + member, + PlacementKind.DIMENSION_SHARD, + dimension=dimension, + ), + ), + process_group=group, + local_member=member, + ) + ) + offset += length + return tuple(result) + + +def _dtensor_replicated(parameter, group): + return tuple( + ShardIdentity( + parameter, + ParameterLayout.DTENSOR_1D_DEFAULT_WORLD, + LogicalRegion.full(parameter), + placements=(_placement(group, member, PlacementKind.REPLICATE),), + process_group=group, + local_member=member, + ) + for member in group.ordered_members + ) + + +def _assert_tight_cpu_fp32(value, shape): + assert type(value) is torch.Tensor + assert value.device.type == "cpu" + assert value.dtype == torch.float32 + assert tuple(value.shape) == tuple(shape) + assert not value.requires_grad + assert value.is_contiguous() + assert value.storage_offset() == 0 + assert value.untyped_storage().nbytes() == value.numel() * value.element_size() + assert bool(torch.isfinite(value).all()) + + +def _sparse_tensor(shape): + with torch.sparse.check_sparse_tensor_invariants(enable=True): + return torch.sparse_coo_tensor( + torch.tensor([[0], [1]]), + torch.tensor([1.0]), + size=shape, + ) + + +def _round_trip(parameter, shards, logical): + manifest = ShardingManifest(tuple(reversed(shards))) + fragments = tuple( + ( + shard, + _project_dense_logical_field(parameter, logical, shard), + ) + for shard in reversed(shards) + ) + assembled = _assemble_dense_logical_field( + manifest, + parameter, + fragments, + ) + _assert_tight_cpu_fp32(assembled, parameter.global_shape) + assert torch.equal(assembled, logical) + return assembled, fragments + + +def test_local_and_grouped_replicated_fields_round_trip_without_aliasing(): + parameter = ParameterIdentity("layer.weight", (2, 3)) + logical = torch.arange(12, dtype=torch.float32).reshape(2, 6)[:, ::2] + assert not logical.is_contiguous() + + for shards in ( + _local_replicated(parameter), + _grouped_replicated(parameter, _group()), + ): + assembled, fragments = _round_trip(parameter, shards, logical) + snapshots = tuple(payload.clone() for _, payload in fragments if payload is not None) + assembled.add_(1000) + assert torch.equal(logical, torch.tensor([[0.0, 2.0, 4.0], [6.0, 8.0, 10.0]])) + assert all( + torch.equal(payload, snapshot) + for (_, payload), snapshot in zip( + ((shard, payload) for shard, payload in fragments if payload is not None), + snapshots, + ) + ) + + +def test_flattened_uneven_and_empty_shards_preserve_row_major_order(): + parameter = ParameterIdentity("layer.weight", (2, 3)) + group = _group(("rank:0", "rank:1", "rank:2", "rank:3")) + shards = _flat_shards(parameter, group, (2, 0, 3, 1)) + logical = torch.tensor([[10.0, 11.0, 12.0], [20.0, 21.0, 22.0]]) + + assembled, fragments = _round_trip(parameter, shards, logical) + + assert torch.equal(assembled, logical) + by_member = {shard.local_member: payload for shard, payload in fragments} + assert torch.equal(by_member["rank:0"], torch.tensor([10.0, 11.0])) + assert by_member["rank:1"] is None + assert torch.equal(by_member["rank:2"], torch.tensor([12.0, 20.0, 21.0])) + assert torch.equal(by_member["rank:3"], torch.tensor([22.0])) + + +def test_whole_parameter_owner_projects_only_to_owner_and_round_trips(): + parameter = ParameterIdentity("layer.weight", (2, 3)) + group = _group() + shards = _owner_shards(parameter, group, "worker:a") + logical = torch.arange(6, dtype=torch.float32).reshape(2, 3).requires_grad_() + + assembled, fragments = _round_trip(parameter, shards, logical) + + assert torch.equal(assembled, logical.detach()) + by_member = {shard.local_member: payload for shard, payload in fragments} + assert by_member["worker:c"] is None + assert by_member["worker:b"] is None + _assert_tight_cpu_fp32(by_member["worker:a"], parameter.global_shape) + + +@pytest.mark.parametrize( + "shape,lengths,dimension", + ( + ((5, 4), (2, 0, 3), 0), + ((3, 5), (2, 2, 1), 1), + ), +) +def test_dtensor_row_and_column_regions_round_trip_uneven_and_empty_shards( + shape, + lengths, + dimension, +): + parameter = ParameterIdentity("layer.weight", shape) + group = _group() + shards = _dtensor_shards( + parameter, + group, + lengths, + dimension=dimension, + ) + base = torch.arange(shape[0] * (shape[1] * 2), dtype=torch.float32).reshape( + shape[0], + shape[1] * 2, + ) + logical = base[:, ::2] + assert not logical.is_contiguous() + + assembled, fragments = _round_trip(parameter, shards, logical) + + assert torch.equal(assembled, logical) + for shard, payload in fragments: + region = shard.logical_region + if region.numel == 0: + assert payload is None + else: + _assert_tight_cpu_fp32(payload, region.lengths) + + +def test_replicated_dtensor_fields_require_equal_full_replicas(): + parameter = ParameterIdentity("layer.weight", (2, 3)) + shards = _dtensor_replicated(parameter, _group()) + logical = torch.arange(6, dtype=torch.float32).reshape(2, 3) + + assembled, fragments = _round_trip(parameter, shards, logical) + + assert torch.equal(assembled, logical) + assert all(torch.equal(payload, logical) for _, payload in fragments) + + +def test_manifest_may_contain_other_parameters_but_fragments_must_not(): + parameter = ParameterIdentity("wanted.weight", (4,)) + shards = _flat_shards(parameter, _group(), (1, 1, 2)) + other_parameter = ParameterIdentity("other.weight", (1,)) + other_shard = _local_replicated(other_parameter)[0] + manifest = ShardingManifest(shards + (other_shard,)) + logical = torch.arange(4, dtype=torch.float32) + fragments = tuple((shard, _project_dense_logical_field(parameter, logical, shard)) for shard in shards) + + assert torch.equal( + _assemble_dense_logical_field(manifest, parameter, fragments), + logical, + ) + with pytest.raises(ValueError, match="exactly cover"): + _assemble_dense_logical_field( + manifest, + parameter, + fragments + ((other_shard, torch.ones(1)),), + ) + + +def test_fragments_reject_missing_extra_duplicate_and_malformed_identities(): + parameter = ParameterIdentity("layer.weight", (3,)) + shards = _flat_shards(parameter, _group(), (1, 1, 1)) + manifest = ShardingManifest(shards) + logical = torch.arange(3, dtype=torch.float32) + fragments = tuple((shard, _project_dense_logical_field(parameter, logical, shard)) for shard in shards) + + with pytest.raises(ValueError, match="exactly cover"): + _assemble_dense_logical_field(manifest, parameter, fragments[:-1]) + with pytest.raises(ValueError, match="duplicate"): + _assemble_dense_logical_field(manifest, parameter, fragments + (fragments[0],)) + different = _local_replicated(ParameterIdentity("different.weight", (3,)))[0] + with pytest.raises(ValueError, match="exactly cover"): + _assemble_dense_logical_field( + manifest, + parameter, + fragments[:-1] + ((different, torch.arange(3, dtype=torch.float32)),), + ) + with pytest.raises(TypeError, match="pair"): + _assemble_dense_logical_field(manifest, parameter, ((shards[0],),)) + with pytest.raises(TypeError, match="ShardIdentity"): + _assemble_dense_logical_field(manifest, parameter, (("not-a-shard", None),)) + + +def test_requested_parameter_must_match_the_manifest_identity_exactly(): + parameter = ParameterIdentity("layer.weight", (3,)) + shard = _local_replicated(parameter)[0] + manifest = ShardingManifest((shard,)) + + with pytest.raises(ValueError, match="does not contain"): + _assemble_dense_logical_field( + manifest, + ParameterIdentity("other.weight", (3,)), + ((shard, torch.arange(3, dtype=torch.float32)),), + ) + with pytest.raises(ValueError, match="does not match"): + _assemble_dense_logical_field( + manifest, + ParameterIdentity("layer.weight", (4,)), + ((shard, torch.arange(3, dtype=torch.float32)),), + ) + + +def test_replicas_reject_absent_and_disagreeing_payloads(): + parameter = ParameterIdentity("layer.weight", (2, 2)) + shards = _grouped_replicated(parameter, _group()) + manifest = ShardingManifest(shards) + logical = torch.arange(4, dtype=torch.float32).reshape(2, 2) + fragments = [(shard, logical.clone()) for shard in shards] + + fragments[1] = (shards[1], None) + with pytest.raises(ValueError, match="every shard"): + _assemble_dense_logical_field(manifest, parameter, fragments) + fragments[1] = (shards[1], logical.clone()) + fragments[2][1][1, 1] += 1 + with pytest.raises(ValueError, match="disagree"): + _assemble_dense_logical_field(manifest, parameter, fragments) + + +def test_replicas_require_bit_exact_signed_zero_payloads(): + parameter = ParameterIdentity("layer.weight", (1,)) + shards = _grouped_replicated(parameter, _group()) + manifest = ShardingManifest(shards) + fragments = [(shard, torch.tensor([0.0])) for shard in shards] + fragments[1] = (shards[1], torch.tensor([-0.0])) + + with pytest.raises(ValueError, match="disagree"): + _assemble_dense_logical_field(manifest, parameter, fragments) + + +def test_flattened_and_dtensor_empty_payloads_are_explicit_none(): + flat_parameter = ParameterIdentity("flat.weight", (2,)) + flat_group = _group() + flat_shards = _flat_shards(flat_parameter, flat_group, (1, 0, 1)) + flat_manifest = ShardingManifest(flat_shards) + flat_fragments = ( + (flat_shards[0], torch.tensor([1.0])), + (flat_shards[1], torch.empty(0)), + (flat_shards[2], torch.tensor([2.0])), + ) + with pytest.raises(ValueError, match="empty flattened"): + _assemble_dense_logical_field( + flat_manifest, + flat_parameter, + flat_fragments, + ) + missing_flat = list(flat_fragments) + missing_flat[0] = (flat_shards[0], None) + missing_flat[1] = (flat_shards[1], None) + with pytest.raises(ValueError, match="nonempty flattened"): + _assemble_dense_logical_field( + flat_manifest, + flat_parameter, + missing_flat, + ) + + dtensor_parameter = ParameterIdentity("dtensor.weight", (2, 2)) + dtensor_shards = _dtensor_shards( + dtensor_parameter, + _group(), + (1, 1, 0), + dimension=1, + ) + dtensor_manifest = ShardingManifest(dtensor_shards) + logical = torch.arange(4, dtype=torch.float32).reshape(2, 2) + dtensor_fragments = [ + ( + shard, + _project_dense_logical_field(dtensor_parameter, logical, shard), + ) + for shard in dtensor_shards + ] + dtensor_fragments[2] = (dtensor_shards[2], torch.empty((2, 0))) + with pytest.raises(ValueError, match="empty DTensor"): + _assemble_dense_logical_field( + dtensor_manifest, + dtensor_parameter, + dtensor_fragments, + ) + dtensor_fragments[2] = (dtensor_shards[2], None) + dtensor_fragments[0] = (dtensor_shards[0], None) + with pytest.raises(ValueError, match="nonempty DTensor"): + _assemble_dense_logical_field( + dtensor_manifest, + dtensor_parameter, + dtensor_fragments, + ) + + +def test_whole_parameter_payload_is_present_exactly_on_the_owner(): + parameter = ParameterIdentity("layer.weight", (2, 2)) + shards = _owner_shards(parameter, _group(), "worker:a") + manifest = ShardingManifest(shards) + logical = torch.arange(4, dtype=torch.float32).reshape(2, 2) + fragments = [ + ( + shard, + logical.clone() if shard.local_member == shard.owner else None, + ) + for shard in shards + ] + owner_index = next(index for index, shard in enumerate(shards) if shard.local_member == shard.owner) + + missing = list(fragments) + missing[owner_index] = (shards[owner_index], None) + with pytest.raises(ValueError, match="owner requires"): + _assemble_dense_logical_field(manifest, parameter, missing) + nonowner_index = next(index for index, shard in enumerate(shards) if shard.local_member != shard.owner) + extra = list(fragments) + extra[nonowner_index] = (shards[nonowner_index], torch.empty(0)) + with pytest.raises(ValueError, match="non-owners"): + _assemble_dense_logical_field(manifest, parameter, extra) + + +@pytest.mark.parametrize( + "bad_value,error", + ( + (torch.ones((2, 3), dtype=torch.float64), TypeError), + (torch.ones(6, dtype=torch.float32), ValueError), + (torch.tensor([[1.0, 2.0, 3.0], [4.0, float("inf"), 6.0]]), ValueError), + (torch.nn.Parameter(torch.ones((2, 3), dtype=torch.float32)), TypeError), + (_sparse_tensor((2, 3)), TypeError), + (torch.empty((2, 3), device="meta"), TypeError), + ), +) +def test_assembly_rejects_wrong_dtype_shape_nonfinite_subclass_and_sparse( + bad_value, + error, +): + parameter = ParameterIdentity("layer.weight", (2, 3)) + shard = _local_replicated(parameter)[0] + with pytest.raises(error): + _assemble_dense_logical_field( + ShardingManifest((shard,)), + parameter, + ((shard, bad_value),), + ) + + +@pytest.mark.parametrize( + "bad_value,error", + ( + (torch.ones((2, 3), dtype=torch.float64), TypeError), + (torch.ones(6, dtype=torch.float32), ValueError), + (torch.tensor([[1.0, 2.0, 3.0], [4.0, float("nan"), 6.0]]), ValueError), + (torch.nn.Parameter(torch.ones((2, 3), dtype=torch.float32)), TypeError), + (_sparse_tensor((2, 3)), TypeError), + ), +) +def test_projection_rejects_wrong_dtype_shape_nonfinite_subclass_and_sparse( + bad_value, + error, +): + parameter = ParameterIdentity("layer.weight", (2, 3)) + shard = _local_replicated(parameter)[0] + with pytest.raises(error): + _project_dense_logical_field(parameter, bad_value, shard) + + +def test_projection_requires_an_exact_target_parameter_identity(): + parameter = ParameterIdentity("layer.weight", (2, 3)) + logical = torch.ones((2, 3), dtype=torch.float32) + other_shard = _local_replicated(ParameterIdentity("other.weight", (2, 3)))[0] + + with pytest.raises(ValueError, match="does not match"): + _project_dense_logical_field(parameter, logical, other_shard) + with pytest.raises(TypeError, match="ShardIdentity"): + _project_dense_logical_field(parameter, logical, "not-a-shard") + with pytest.raises(TypeError, match="ParameterIdentity"): + _project_dense_logical_field("layer.weight", logical, other_shard) + + +def test_api_requires_valid_manifest_parameter_and_fragment_containers(): + parameter = ParameterIdentity("layer.weight", (1,)) + shard = _local_replicated(parameter)[0] + manifest = ShardingManifest((shard,)) + + with pytest.raises(TypeError, match="ShardingManifest"): + _assemble_dense_logical_field( + "not-a-manifest", + parameter, + ((shard, torch.ones(1)),), + ) + with pytest.raises(TypeError, match="ParameterIdentity"): + _assemble_dense_logical_field( + manifest, + "layer.weight", + ((shard, torch.ones(1)),), + ) + with pytest.raises(TypeError, match="sequence"): + _assemble_dense_logical_field(manifest, parameter, None) + with pytest.raises(TypeError, match="sequence"): + _assemble_dense_logical_field(manifest, parameter, "fragment") diff --git a/tests/test_portable_identity.py b/tests/test_portable_identity.py new file mode 100644 index 0000000..0a6e98d --- /dev/null +++ b/tests/test_portable_identity.py @@ -0,0 +1,517 @@ +"""Strict primitive wire coverage for portable optimizer identities.""" + +import copy +import io + +import pytest +import torch + +from gefen.contracts import ( + LogicalRegion, + LogicalSlice, + ParameterIdentity, + ParameterLayout, + PlacementKind, + ProcessGroupIdentity, + ShardIdentity, + ShardPlacement, + ShardingManifest, +) +from gefen.portable_identity import ( + _normalize_parameter_identity, + _normalize_process_group_identity, + _normalize_shard_identity, + _normalize_sharding_manifest, + _serialize_parameter_identity, + _serialize_process_group_identity, + _serialize_shard_identity, + _serialize_sharding_manifest, +) + + +def _group(): + return ProcessGroupIdentity("pipeline:1/data_parallel", ("worker:c", "worker:a", "worker:b")) + + +def _placement(group, member, kind, *, dimension=None): + return ShardPlacement( + "data_parallel", + kind, + group.ordered_members.index(member), + len(group.ordered_members), + dimension, + ) + + +def _replicated_shards(parameter, group): + return tuple( + ShardIdentity( + parameter, + ParameterLayout.REPLICATED, + LogicalSlice.full(parameter), + placements=(_placement(group, member, PlacementKind.REPLICATE),), + process_group=group, + local_member=member, + ) + for member in group.ordered_members + ) + + +def _flat_shards(parameter, group): + boundaries = ((0, 2), (2, 2), (4, 1)) + return tuple( + ShardIdentity( + parameter, + ParameterLayout.FLATTENED_ELEMENT_SHARD, + LogicalSlice(offset, length), + placements=(_placement(group, member, PlacementKind.FLAT_SHARD),), + process_group=group, + local_member=member, + ) + for member, (offset, length) in zip(group.ordered_members, boundaries) + ) + + +def _owner_shards(parameter, group, owner): + return tuple( + ShardIdentity( + parameter, + ParameterLayout.WHOLE_PARAMETER_OWNER, + (LogicalSlice.full(parameter) if member == owner else LogicalSlice(0, 0)), + placements=(_placement(group, member, PlacementKind.WHOLE_PARAMETER_OWNER),), + process_group=group, + local_member=member, + owner=owner, + ) + for member in group.ordered_members + ) + + +def _dtensor_shards(parameter, group, lengths, *, dimension): + offset = 0 + shards = [] + for member, length in zip(group.ordered_members, lengths): + offsets = [0] * len(parameter.global_shape) + region_lengths = list(parameter.global_shape) + offsets[dimension] = offset + region_lengths[dimension] = length + shards.append( + ShardIdentity( + parameter, + ParameterLayout.DTENSOR_1D_DEFAULT_WORLD, + LogicalRegion(offsets, region_lengths), + placements=( + _placement( + group, + member, + PlacementKind.DIMENSION_SHARD, + dimension=dimension, + ), + ), + process_group=group, + local_member=member, + ) + ) + offset += length + return tuple(shards) + + +def _dtensor_replicated_shards(parameter, group): + return tuple( + ShardIdentity( + parameter, + ParameterLayout.DTENSOR_1D_DEFAULT_WORLD, + LogicalRegion.full(parameter), + placements=(_placement(group, member, PlacementKind.REPLICATE),), + process_group=group, + local_member=member, + ) + for member in group.ordered_members + ) + + +def _complete_manifest(): + group = _group() + shards = ( + ShardIdentity( + ParameterIdentity("local.weight", (2, 3)), + ParameterLayout.REPLICATED, + LogicalSlice(0, 6), + ), + ) + shards += _replicated_shards(ParameterIdentity("replicated.weight", (2, 3)), group) + shards += _flat_shards(ParameterIdentity("flat.weight", (5,)), group) + shards += _owner_shards(ParameterIdentity("owner.weight", (2, 3)), group, "worker:a") + shards += _dtensor_shards( + ParameterIdentity("row.weight", (5, 4)), + group, + (2, 2, 1), + dimension=0, + ) + shards += _dtensor_shards( + ParameterIdentity("column_with_empty.weight", (2, 2)), + group, + (1, 1, 0), + dimension=1, + ) + shards += _dtensor_replicated_shards(ParameterIdentity("dtensor_replicated.weight", (2, 3)), group) + return ShardingManifest(tuple(reversed(shards))) + + +def _assert_primitive_tree(value): + assert type(value) in {dict, list, str, int, type(None)} + if type(value) is dict: + assert all(type(key) is str for key in value) + for item in value.values(): + _assert_primitive_tree(item) + elif type(value) is list: + for item in value: + _assert_primitive_tree(item) + + +def test_parameter_and_process_group_codecs_have_exact_canonical_wire_shapes(): + parameter = ParameterIdentity("Encoder.Block.Weight", (0, 4, 7)) + group = _group() + + parameter_wire = _serialize_parameter_identity(parameter) + group_wire = _serialize_process_group_identity(group) + + assert parameter_wire == { + "schema_version": 1, + "fqn": "Encoder.Block.Weight", + "global_shape": [0, 4, 7], + } + assert group_wire == { + "schema_version": 1, + "semantic_name": "pipeline:1/data_parallel", + "ordered_members": ["worker:c", "worker:a", "worker:b"], + } + assert _normalize_parameter_identity(parameter_wire) == parameter_wire + assert _normalize_process_group_identity(group_wire) == group_wire + _assert_primitive_tree(parameter_wire) + _assert_primitive_tree(group_wire) + + +@pytest.mark.parametrize( + "identity,extent_kind", + ( + ( + ShardIdentity( + ParameterIdentity("local.weight", (2, 3)), + ParameterLayout.REPLICATED, + LogicalSlice(0, 6), + ), + "logical_slice", + ), + (_flat_shards(ParameterIdentity("flat.weight", (5,)), _group())[1], "logical_slice"), + ( + _owner_shards( + ParameterIdentity("owner.weight", (2, 3)), + _group(), + "worker:a", + )[0], + "logical_slice", + ), + ( + _dtensor_shards( + ParameterIdentity("row.weight", (5, 4)), + _group(), + (2, 2, 1), + dimension=0, + )[2], + "logical_region", + ), + ( + _dtensor_shards( + ParameterIdentity("empty.weight", (2, 2)), + _group(), + (1, 1, 0), + dimension=1, + )[2], + "logical_region", + ), + ( + _dtensor_replicated_shards( + ParameterIdentity("dtensor_replicated.weight", (2, 3)), + _group(), + )[1], + "logical_region", + ), + ), +) +def test_shard_codec_losslessly_round_trips_every_supported_identity(identity, extent_kind): + wire = _serialize_shard_identity(identity) + + assert wire["logical_extent"]["kind"] == extent_kind + assert _normalize_shard_identity(wire) == wire + _assert_primitive_tree(wire) + + +def test_shard_wire_explicitly_discriminates_slice_and_region_fields(): + legacy = _serialize_shard_identity(_flat_shards(ParameterIdentity("flat.weight", (5,)), _group())[0]) + dtensor = _serialize_shard_identity( + _dtensor_shards( + ParameterIdentity("row.weight", (5, 4)), + _group(), + (2, 2, 1), + dimension=0, + )[0] + ) + + assert legacy["logical_extent"] == { + "kind": "logical_slice", + "flat_offset": 0, + "length": 2, + } + assert dtensor["logical_extent"] == { + "kind": "logical_region", + "offsets": [0, 0], + "lengths": [2, 4], + } + + +def test_manifest_codec_is_deterministic_complete_and_weights_only_safe(): + manifest = _complete_manifest() + wire = _serialize_sharding_manifest(manifest) + reversed_wire = copy.deepcopy(wire) + reversed_wire["shards"].reverse() + + assert [record["parameter"]["fqn"] for record in wire["shards"]] == [ + shard.parameter.fqn for shard in manifest.shards + ] + assert _normalize_sharding_manifest(wire) == wire + assert _normalize_sharding_manifest(reversed_wire) == wire + assert any( + record["logical_extent"] + == { + "kind": "logical_region", + "offsets": [0, 2], + "lengths": [2, 0], + } + for record in wire["shards"] + ) + _assert_primitive_tree(wire) + + buffer = io.BytesIO() + torch.save(wire, buffer) + buffer.seek(0) + loaded = torch.load(buffer, weights_only=True) + assert _normalize_sharding_manifest(loaded) == wire + + +def test_normalizers_rebuild_fresh_canonical_container_trees(): + source = _serialize_sharding_manifest(_complete_manifest()) + normalized = _normalize_sharding_manifest(source) + + assert normalized == source + assert normalized is not source + assert normalized["shards"] is not source["shards"] + assert normalized["shards"][0] is not source["shards"][0] + + source["shards"][0]["parameter"]["global_shape"].append(99) + assert normalized != source + + +@pytest.mark.parametrize( + "corruption", + ( + "not_dict", + "missing_key", + "extra_key", + "bool_version", + "tuple_shape", + "bool_dimension", + "invalid_fqn", + ), +) +def test_parameter_identity_schema_corruption_is_rejected(corruption): + wire = _serialize_parameter_identity(ParameterIdentity("layer.weight", (2, 3))) + if corruption == "not_dict": + wire = [] + elif corruption == "missing_key": + wire.pop("fqn") + elif corruption == "extra_key": + wire["global_rank"] = 0 + elif corruption == "bool_version": + wire["schema_version"] = True + elif corruption == "tuple_shape": + wire["global_shape"] = (2, 3) + elif corruption == "bool_dimension": + wire["global_shape"][0] = True + else: + wire["fqn"] = ".layer" + + with pytest.raises((TypeError, ValueError)): + _normalize_parameter_identity(wire) + + +@pytest.mark.parametrize( + "corruption", + ( + "missing_key", + "extra_handle", + "bool_version", + "tuple_members", + "non_string_member", + "duplicate_member", + "invalid_name", + ), +) +def test_process_group_identity_schema_corruption_is_rejected(corruption): + wire = _serialize_process_group_identity(_group()) + if corruption == "missing_key": + wire.pop("semantic_name") + elif corruption == "extra_handle": + wire["process_group_handle"] = 123 + elif corruption == "bool_version": + wire["schema_version"] = True + elif corruption == "tuple_members": + wire["ordered_members"] = tuple(wire["ordered_members"]) + elif corruption == "non_string_member": + wire["ordered_members"][0] = 0 + elif corruption == "duplicate_member": + wire["ordered_members"][1] = wire["ordered_members"][0] + else: + wire["semantic_name"] = " dp" + + with pytest.raises((TypeError, ValueError)): + _normalize_process_group_identity(wire) + + +@pytest.mark.parametrize( + "corruption", + ( + "missing_key", + "runtime_rank", + "bool_version", + "enum_layout", + "tuple_placements", + "bad_process_group_type", + "non_string_member", + "non_string_owner", + "unknown_extent_kind", + "mixed_slice_region_fields", + "tuple_region_offsets", + "bool_region_length", + "slice_for_dtensor", + "region_for_flat", + "placement_missing_key", + "enum_placement_kind", + "bool_coordinate", + "bool_parts", + "bool_parameter_dimension", + "coordinate_mismatch", + ), +) +def test_shard_identity_schema_and_semantic_corruption_is_rejected(corruption): + identity = _dtensor_shards( + ParameterIdentity("row.weight", (5, 4)), + _group(), + (2, 2, 1), + dimension=0, + )[0] + wire = _serialize_shard_identity(identity) + if corruption == "missing_key": + wire.pop("owner") + elif corruption == "runtime_rank": + wire["global_rank"] = 7 + elif corruption == "bool_version": + wire["schema_version"] = True + elif corruption == "enum_layout": + wire["layout"] = ParameterLayout.DTENSOR_1D_DEFAULT_WORLD + elif corruption == "tuple_placements": + wire["placements"] = tuple(wire["placements"]) + elif corruption == "bad_process_group_type": + wire["process_group"] = [] + elif corruption == "non_string_member": + wire["local_member"] = 0 + elif corruption == "non_string_owner": + wire["owner"] = 0 + elif corruption == "unknown_extent_kind": + wire["logical_extent"]["kind"] = "flat" + elif corruption == "mixed_slice_region_fields": + wire["logical_extent"]["flat_offset"] = 0 + elif corruption == "tuple_region_offsets": + wire["logical_extent"]["offsets"] = (0, 0) + elif corruption == "bool_region_length": + wire["logical_extent"]["lengths"][0] = True + elif corruption == "slice_for_dtensor": + wire["logical_extent"] = { + "kind": "logical_slice", + "flat_offset": 0, + "length": 8, + } + elif corruption == "region_for_flat": + wire["layout"] = "flattened_element_shard" + elif corruption == "placement_missing_key": + wire["placements"][0].pop("parts") + elif corruption == "enum_placement_kind": + wire["placements"][0]["kind"] = PlacementKind.DIMENSION_SHARD + elif corruption == "bool_coordinate": + wire["placements"][0]["coordinate"] = True + elif corruption == "bool_parts": + wire["placements"][0]["parts"] = True + elif corruption == "bool_parameter_dimension": + wire["placements"][0]["parameter_dimension"] = True + else: + wire["placements"][0]["coordinate"] = 1 + + with pytest.raises((TypeError, ValueError)): + _normalize_shard_identity(wire) + + +@pytest.mark.parametrize( + "corruption", + ( + "not_dict", + "missing_key", + "extra_key", + "bool_version", + "tuple_shards", + "empty_shards", + "duplicate_shard", + "incomplete_group", + ), +) +def test_sharding_manifest_schema_and_completeness_corruption_is_rejected( + corruption, +): + wire = _serialize_sharding_manifest(_complete_manifest()) + if corruption == "not_dict": + wire = [] + elif corruption == "missing_key": + wire.pop("schema_version") + elif corruption == "extra_key": + wire["default_process_group"] = "world" + elif corruption == "bool_version": + wire["schema_version"] = True + elif corruption == "tuple_shards": + wire["shards"] = tuple(wire["shards"]) + elif corruption == "empty_shards": + wire["shards"] = [] + elif corruption == "duplicate_shard": + wire["shards"].append(copy.deepcopy(wire["shards"][0])) + else: + fqn = "flat.weight" + wire["shards"] = [ + shard + for shard in wire["shards"] + if not (shard["parameter"]["fqn"] == fqn and shard["local_member"] == "worker:b") + ] + + with pytest.raises((TypeError, ValueError)): + _normalize_sharding_manifest(wire) + + +@pytest.mark.parametrize( + "serializer,value", + ( + (_serialize_parameter_identity, object()), + (_serialize_process_group_identity, object()), + (_serialize_shard_identity, object()), + (_serialize_sharding_manifest, object()), + ), +) +def test_serializers_require_contract_dataclasses(serializer, value): + with pytest.raises(TypeError): + serializer(value) diff --git a/tests/test_portable_wire.py b/tests/test_portable_wire.py new file mode 100644 index 0000000..74c693a --- /dev/null +++ b/tests/test_portable_wire.py @@ -0,0 +1,480 @@ +import dataclasses +import hashlib +import math +import struct + +import pytest +import torch + +import gefen.portable_wire as wire +from gefen.portable_wire import ( + _CanonicalWireLimits, + _CanonicalWirePlan, + _CanonicalWireTensorSpec, + _PreparedCanonicalWireValue, + _parse_canonical_wire_metadata, + _prepare_canonical_wire_value, + _reconstruct_canonical_wire_value, +) + + +def _limits(**changes): + limits = _CanonicalWireLimits( + max_fragment_tensor_bytes=1 << 20, + max_collective_tensor_bytes=8 << 20, + max_collective_metadata_bytes=1 << 20, + max_metadata_bytes=1 << 20, + ) + return dataclasses.replace(limits, **changes) + + +def _assert_value_equal(left, right): + if torch.is_tensor(left) or torch.is_tensor(right): + assert type(left) is torch.Tensor + assert type(right) is torch.Tensor + assert left.dtype == right.dtype + assert tuple(left.shape) == tuple(right.shape) + assert torch.equal(left, right) + return + assert type(left) is type(right) + if type(left) is dict: + assert list(left) == list(right) + for key in left: + _assert_value_equal(left[key], right[key]) + elif type(left) in {list, tuple}: + assert len(left) == len(right) + for left_item, right_item in zip(left, right): + _assert_value_equal(left_item, right_item) + elif type(left) is float and left == 0.0: + assert math.copysign(1.0, left) == math.copysign(1.0, right) + else: + assert left == right + + +def _round_trip(value, limits=None): + limits = _limits() if limits is None else limits + plan = _prepare_canonical_wire_value(value, limits) + prepared = _parse_canonical_wire_metadata(plan.metadata, limits=limits) + assert prepared.metadata_digest == plan.metadata_digest + assert prepared.fragment_digest == plan.fragment_digest + assert prepared.tensor_specs == plan.tensor_specs + result = _reconstruct_canonical_wire_value( + prepared, + plan.tensors, + expected_fragment_digest=plan.fragment_digest, + ) + _assert_value_equal(value, result) + return plan, prepared, result + + +def _metadata(root, descriptors=()): + return ( + b"GFNCV1\0\0" + + struct.pack(">HHI", 1, 0, 0) + + root + + struct.pack(">Q", len(descriptors)) + + b"".join(descriptors) + ) + + +def _string(value): + encoded = value.encode("utf-8") + return b"\x05" + struct.pack(">Q", len(encoded)) + encoded + + +def _descriptor(dtype_code=1, shape=(), numel=1, nbytes=1, digest=b"\0" * 32, flags=0): + return ( + struct.pack(">HHI", dtype_code, len(shape), flags) + + b"".join(struct.pack(">Q", dimension) for dimension in shape) + + struct.pack(">QQ", numel, nbytes) + + digest + ) + + +def test_golden_metadata_tensor_and_fragment_digests(): + value = {"a": -1, "b": torch.tensor([1, -2], dtype=torch.int16)} + plan = _prepare_canonical_wire_value(value, _limits()) + assert plan.metadata.hex() == ( + "47464e435631000000010000000000000800000000000000020500000000000000016103010000000000000001010500000000000000016209000000000000000000000000000000010004000100000000000000000000000200000000000000020000000000000004" + "356f76987e72498bcbc4866de355fdcbad68f9a1eb8807f1073728ee463f1e27" + ) + assert plan.metadata_digest.hex() == "c5894b6b9cc4f0f8591ed4fbb66d943b0bb69362fb460df413c42b36def889d8" + assert plan.fragment_digest.hex() == "155e9c7c022131d7ea971fd7bc1d69e13c7f9dcbaf89301bb29c135a24de4846" + _round_trip(value) + + +def test_round_trip_primitives_nested_containers_and_edge_values(): + value = { + "bools": [False, True], + "floats": (0.0, -0.0, float.fromhex("0x0.0000000000001p-1022"), 1.25), + "integers": [0, 1, -1, (1 << 4095) - 1], + "none": None, + "strings": ["", "ASCII", "λ🙂"], + } + _round_trip(value) + + +@pytest.mark.parametrize( + ("dtype", "values"), + [ + (torch.bool, [False, True, True]), + (torch.uint8, [0, 1, 255]), + (torch.int8, [-128, 0, 127]), + (torch.int16, [-32768, 0, 32767]), + (torch.int32, [-(1 << 31), 0, (1 << 31) - 1]), + (torch.int64, [-(1 << 63), 0, (1 << 63) - 1]), + (torch.float16, [-1.5, 0.0, 2.25]), + (torch.bfloat16, [-1.5, 0.0, 2.25]), + (torch.float32, [-1.5, 0.0, 2.25]), + (torch.float64, [-1.5, 0.0, 2.25]), + (torch.complex64, [complex(-1.5, 2.0), 0j, complex(2.25, -3.5)]), + (torch.complex128, [complex(-1.5, 2.0), 0j, complex(2.25, -3.5)]), + ], +) +def test_round_trip_fixed_dtype_table(dtype, values): + plan, _, result = _round_trip(torch.tensor(values, dtype=dtype)) + assert plan.tensor_specs[0].dtype_code == 1 + [record[1] for record in wire._DTYPE_RECORDS].index(dtype) + assert result.is_contiguous() + assert result.storage_offset() == 0 + assert result.untyped_storage().nbytes() == result.numel() * result.element_size() + + +def test_round_trip_noncontiguous_conjugate_scalar_and_empty_tensors(): + noncontiguous = torch.arange(30, dtype=torch.float64).reshape(5, 6).t()[1:5:2] + conjugate = torch.tensor([1 + 2j, 3 - 4j], dtype=torch.complex128).conj() + scalar = torch.tensor(-7, dtype=torch.int64) + empty = torch.empty((3, 0, 2), dtype=torch.float32) + plan, _, result = _round_trip([noncontiguous, conjugate, scalar, empty]) + assert not noncontiguous.is_contiguous() + assert conjugate.is_conj() + assert all(tensor.is_contiguous() for tensor in plan.tensors) + assert not result[1].is_conj() + + +def test_tensor_digest_uses_canonical_little_endian_component_bytes(): + value = torch.tensor([0x01020304], dtype=torch.int32) + plan = _prepare_canonical_wire_value(value, _limits()) + expected = hashlib.sha256() + expected.update(b"gefen.canonical_wire.tensor.v1\0") + expected.update(struct.pack(">HHQ", 5, 1, 1)) + expected.update(struct.pack(">Q", 4)) + expected.update(b"\x04\x03\x02\x01") + assert plan.tensor_specs[0].digest == expected.digest() + + +def test_fragment_digest_frames_metadata_and_ordered_tensor_digests(): + plan = _prepare_canonical_wire_value( + [torch.tensor([1], dtype=torch.int16), torch.tensor([2.0], dtype=torch.float64)], + _limits(), + ) + expected = hashlib.sha256() + expected.update(b"gefen.canonical_wire.fragment.v1\0") + expected.update(hashlib.sha256(plan.metadata).digest()) + expected.update(struct.pack(">Q", 2)) + for spec in plan.tensor_specs: + expected.update(struct.pack(">Q", spec.nbytes)) + expected.update(spec.digest) + assert plan.fragment_digest == expected.digest() + + +def test_dictionary_encoding_is_deterministic_and_sorted(): + left = _prepare_canonical_wire_value({"z": 1, "a": 2}, _limits()) + right = _prepare_canonical_wire_value({"a": 2, "z": 1}, _limits()) + assert left.metadata == right.metadata + assert left.fragment_digest == right.fragment_digest + + +def test_prepare_and_reconstruct_own_tensor_storage_and_containers(): + source = torch.arange(6, dtype=torch.float32).reshape(2, 3).t() + value = {"tensor": [source], "stable": 3} + plan = _prepare_canonical_wire_value(value, _limits(chunk_bytes=3)) + source.fill_(99) + value["tensor"].append("late") + prepared = _parse_canonical_wire_metadata(plan.metadata, limits=_limits(chunk_bytes=3)) + result = _reconstruct_canonical_wire_value(prepared, plan.tensors) + assert torch.equal(result["tensor"][0], torch.arange(6, dtype=torch.float32).reshape(2, 3).t()) + assert result["tensor"][0].data_ptr() != plan.tensors[0].data_ptr() + plan.tensors[0].fill_(-1) + assert not torch.equal(result["tensor"][0], plan.tensors[0]) + result["tensor"].append("owned") + assert len(value["tensor"]) == 2 + + +def test_tiny_chunk_budget_is_wire_identical_for_real_and_complex_tensors(): + value = { + "complex": torch.tensor([1 + 2j, -3 + 4j, 5 - 6j], dtype=torch.complex128), + "view": torch.arange(60, dtype=torch.float64).reshape(6, 10)[:, ::3], + } + normal = _prepare_canonical_wire_value(value, _limits()) + tiny_limits = _limits(chunk_bytes=1) + tiny = _prepare_canonical_wire_value(value, tiny_limits) + assert tiny.metadata == normal.metadata + assert tiny.metadata_digest == normal.metadata_digest + assert tiny.fragment_digest == normal.fragment_digest + for tiny_tensor, normal_tensor in zip(tiny.tensors, normal.tensors): + assert torch.equal(tiny_tensor, normal_tensor) + prepared = _parse_canonical_wire_metadata(tiny.metadata, limits=tiny_limits) + result = _reconstruct_canonical_wire_value(prepared, tiny.tensors) + _assert_value_equal(value, result) + + +def test_noncontiguous_clone_scales_coordinate_scratch_to_chunk_budget(monkeypatch): + value = torch.arange(2 * 3 * 4 * 5, dtype=torch.float32).reshape(2, 3, 4, 5) + value = value.permute(3, 2, 1, 0) + counts = [] + original = wire.torch.arange + + def tracked(start, stop, **kwargs): + counts.append(stop - start) + return original(start, stop, **kwargs) + + monkeypatch.setattr(wire.torch, "arange", tracked) + limits = _limits(chunk_bytes=100) + plan = _prepare_canonical_wire_value(value, limits) + + assert counts + assert max(counts) <= 100 // (value.element_size() + 8 * (value.ndim + 2)) + assert torch.equal(plan.tensors[0], value) + + +@pytest.mark.parametrize( + "metadata", + [ + b"", + b"BADMAGIC" + struct.pack(">HHI", 1, 0, 0) + b"\0" + struct.pack(">Q", 0), + b"GFNCV1\0\0" + struct.pack(">HHI", 2, 0, 0) + b"\0" + struct.pack(">Q", 0), + b"GFNCV1\0\0" + struct.pack(">HHI", 1, 1, 0) + b"\0" + struct.pack(">Q", 0), + b"GFNCV1\0\0" + struct.pack(">HHI", 1, 0, 1) + b"\0" + struct.pack(">Q", 0), + _metadata(b"\xff"), + _metadata(b"\0") + b"trailing", + ], +) +def test_parser_rejects_headers_unknown_tags_truncation_and_trailing_bytes(metadata): + with pytest.raises(ValueError): + _parse_canonical_wire_metadata(metadata, limits=_limits()) + + +@pytest.mark.parametrize( + "root", + [ + b"\x03\x02" + struct.pack(">Q", 0), + b"\x03\x00" + struct.pack(">Q", 1) + b"\0", + b"\x03\x01" + struct.pack(">Q", 0), + b"\x04" + struct.pack(">d", math.inf), + b"\x05" + struct.pack(">Q", 1) + b"\xff", + ], +) +def test_parser_rejects_noncanonical_scalars(root): + with pytest.raises(ValueError): + _parse_canonical_wire_metadata(_metadata(root), limits=_limits()) + + +def test_parser_rejects_non_string_unsorted_and_duplicate_dictionary_keys(): + roots = [ + b"\x08" + struct.pack(">Q", 1) + b"\0\0", + b"\x08" + struct.pack(">Q", 2) + _string("b") + b"\0" + _string("a") + b"\0", + b"\x08" + struct.pack(">Q", 2) + _string("a") + b"\0" + _string("a") + b"\0", + ] + for root in roots: + with pytest.raises(ValueError): + _parse_canonical_wire_metadata(_metadata(root), limits=_limits()) + + +def test_parser_rejects_tensor_reference_and_descriptor_count_mismatches(): + cases = [ + _metadata(b"\x09" + struct.pack(">Q", 1), [_descriptor()]), + _metadata(b"\x09" + struct.pack(">Q", 0)), + _metadata(b"\0", [_descriptor()]), + ] + for metadata in cases: + with pytest.raises(ValueError): + _parse_canonical_wire_metadata(metadata, limits=_limits()) + + +@pytest.mark.parametrize( + "descriptor", + [ + _descriptor(dtype_code=99), + _descriptor(flags=1), + _descriptor(dtype_code=5, shape=(2,), numel=3, nbytes=12), + _descriptor(dtype_code=5, shape=(2,), numel=2, nbytes=7), + ], +) +def test_parser_rejects_invalid_tensor_descriptors(descriptor): + metadata = _metadata(b"\x09" + struct.pack(">Q", 0), [descriptor]) + with pytest.raises(ValueError): + _parse_canonical_wire_metadata(metadata, limits=_limits()) + + +def test_parser_rejects_tensor_shape_product_overflow(): + descriptor = _descriptor(dtype_code=1, shape=((1 << 62), 4), numel=0, nbytes=0) + metadata = _metadata(b"\x09" + struct.pack(">Q", 0), [descriptor]) + with pytest.raises(ValueError, match="overflows"): + _parse_canonical_wire_metadata(metadata, limits=_limits()) + + +def test_parser_rejects_unrepresentable_dimension_even_when_product_is_zero(): + descriptor = _descriptor(dtype_code=1, shape=((1 << 63), 0), numel=0, nbytes=0) + metadata = _metadata(b"\x09" + struct.pack(">Q", 0), [descriptor]) + with pytest.raises(ValueError, match="signed int64"): + _parse_canonical_wire_metadata(metadata, limits=_limits()) + + +def test_parser_validates_all_metadata_before_tensor_allocation(monkeypatch): + plan = _prepare_canonical_wire_value(torch.arange(4), _limits()) + + def fail_allocation(*args, **kwargs): + raise AssertionError("parser allocated a tensor") + + monkeypatch.setattr(wire.torch, "empty", fail_allocation) + prepared = _parse_canonical_wire_metadata(plan.metadata, limits=_limits()) + assert prepared.tensor_specs == plan.tensor_specs + with pytest.raises(ValueError, match="trailing"): + _parse_canonical_wire_metadata(plan.metadata + b"x", limits=_limits()) + + +@pytest.mark.parametrize( + "field", + [field.name for field in dataclasses.fields(_CanonicalWireLimits)], +) +def test_limits_require_strict_positive_ints(field): + limits = _limits() + with pytest.raises(ValueError): + dataclasses.replace(limits, **{field: 0}) + with pytest.raises(TypeError): + dataclasses.replace(limits, **{field: True}) + + +def test_limits_reject_unrepresentable_or_inconsistent_bounds(): + with pytest.raises(TypeError): + _CanonicalWireLimits() + with pytest.raises(ValueError): + _limits(max_tensor_rank=(1 << 16)) + with pytest.raises(ValueError, match="recursive parsing"): + _limits(max_tree_depth=257) + with pytest.raises(ValueError): + _CanonicalWireLimits( + max_fragment_tensor_bytes=(1 << 40) + 1, + max_collective_tensor_bytes=1 << 40, + max_collective_metadata_bytes=1 << 30, + ) + with pytest.raises(ValueError): + _CanonicalWireLimits( + max_fragment_tensor_bytes=1 << 20, + max_collective_tensor_bytes=8 << 20, + max_collective_metadata_bytes=1 << 30, + max_metadata_bytes=(1 << 30) + 1, + ) + + +def test_prepare_and_parse_enforce_tree_and_scalar_limits(): + cases = [ + ([None, None], _limits(max_container_items=1), "max_container_items"), + ([[None]], _limits(max_tree_depth=1), "max_tree_depth"), + ([None, None], _limits(max_tree_nodes=2), "max_tree_nodes"), + ("ab", _limits(max_string_bytes=1), "max_string_bytes"), + (256, _limits(max_integer_bytes=1), "max_integer_bytes"), + ] + for value, limits, match in cases: + with pytest.raises(ValueError, match=match): + _prepare_canonical_wire_value(value, limits) + encoded = [ + (_metadata(b"\x06" + struct.pack(">Q", 2) + b"\0\0"), _limits(max_container_items=1)), + (_metadata(b"\x06" + struct.pack(">Q", 1) + b"\x06" + struct.pack(">Q", 1) + b"\0"), _limits(max_tree_depth=1)), + (_metadata(b"\x06" + struct.pack(">Q", 2) + b"\0\0"), _limits(max_tree_nodes=2)), + (_metadata(_string("ab")), _limits(max_string_bytes=1)), + (_metadata(b"\x03\0" + struct.pack(">Q", 2) + b"\x01\x00"), _limits(max_integer_bytes=1)), + ] + for metadata, limits in encoded: + with pytest.raises(ValueError): + _parse_canonical_wire_metadata(metadata, limits=limits) + + +def test_prepare_and_parse_enforce_metadata_tensor_count_rank_and_byte_limits(): + tensor = torch.arange(4, dtype=torch.int32) + normal = _prepare_canonical_wire_value(tensor, _limits()) + two_tensors = _prepare_canonical_wire_value([tensor, tensor], _limits()) + rank_two = _prepare_canonical_wire_value(tensor.reshape(2, 2), _limits()) + with pytest.raises(ValueError, match="max_metadata_bytes"): + _prepare_canonical_wire_value(None, _limits(max_metadata_bytes=24)) + with pytest.raises(ValueError, match="max_metadata_bytes"): + _parse_canonical_wire_metadata(normal.metadata, limits=_limits(max_metadata_bytes=len(normal.metadata) - 1)) + with pytest.raises(ValueError, match="max_tensors"): + _prepare_canonical_wire_value([tensor, tensor], _limits(max_tensors=1)) + with pytest.raises(ValueError, match="max_tensors"): + _parse_canonical_wire_metadata(two_tensors.metadata, limits=_limits(max_tensors=1)) + with pytest.raises(ValueError, match="max_tensor_rank"): + _prepare_canonical_wire_value(tensor.reshape(2, 2), _limits(max_tensor_rank=1)) + with pytest.raises(ValueError, match="max_tensor_rank"): + _parse_canonical_wire_metadata(rank_two.metadata, limits=_limits(max_tensor_rank=1)) + with pytest.raises(ValueError, match="max_fragment_tensor_bytes"): + _prepare_canonical_wire_value(tensor, _limits(max_fragment_tensor_bytes=15)) + with pytest.raises(ValueError, match="max_fragment_tensor_bytes"): + _parse_canonical_wire_metadata(normal.metadata, limits=_limits(max_fragment_tensor_bytes=15)) + + +def test_reconstruction_rejects_fragment_payload_and_geometry_corruption(): + value = torch.tensor([1.0, 2.0], dtype=torch.float32) + plan = _prepare_canonical_wire_value(value, _limits()) + prepared = _parse_canonical_wire_metadata(plan.metadata, limits=_limits()) + with pytest.raises(ValueError, match="fragment digest"): + _reconstruct_canonical_wire_value(prepared, plan.tensors, expected_fragment_digest=b"x" * 32) + corrupted = (plan.tensors[0].clone(),) + corrupted[0][0] = 9 + with pytest.raises(ValueError, match="invalid digest"): + _reconstruct_canonical_wire_value(prepared, corrupted) + with pytest.raises(ValueError, match="payload count"): + _reconstruct_canonical_wire_value(prepared, ()) + with pytest.raises(ValueError, match="descriptor"): + _reconstruct_canonical_wire_value(prepared, (plan.tensors[0].to(torch.float64),)) + with pytest.raises(ValueError, match="descriptor"): + _reconstruct_canonical_wire_value(prepared, (plan.tensors[0].reshape(1, 2),)) + + +def test_reconstruction_rejects_nonfinite_nontight_and_noncontiguous_payloads(): + plan = _prepare_canonical_wire_value(torch.tensor([1.0, 2.0]), _limits()) + prepared = _parse_canonical_wire_metadata(plan.metadata, limits=_limits()) + nonfinite = plan.tensors[0].clone() + nonfinite[0] = math.inf + with pytest.raises(ValueError, match="finite"): + _reconstruct_canonical_wire_value(prepared, (nonfinite,)) + backing = torch.empty(3, dtype=torch.float32) + backing[:2].copy_(plan.tensors[0]) + with pytest.raises(ValueError, match="descriptor"): + _reconstruct_canonical_wire_value(prepared, (backing[:2],)) + noncontiguous = torch.stack((plan.tensors[0], plan.tensors[0]))[:, 0] + assert not noncontiguous.is_contiguous() + with pytest.raises(ValueError, match="descriptor"): + _reconstruct_canonical_wire_value(prepared, (noncontiguous,)) + + +def test_prepare_rejects_unsupported_values_nonfinite_and_tensor_subclasses(): + class TensorSubclass(torch.Tensor): + pass + + invalid = [object(), {1: "bad"}, float("nan"), torch.tensor([float("inf")]), torch.tensor([1], dtype=torch.uint16)] + invalid.append(torch.Tensor._make_subclass(TensorSubclass, torch.tensor([1.0]), False)) + for value in invalid: + with pytest.raises((TypeError, ValueError)): + _prepare_canonical_wire_value(value, _limits()) + + +def test_strict_internal_type_construction_and_api_arguments(): + limits = _limits() + with pytest.raises(TypeError): + _prepare_canonical_wire_value(None, object()) + with pytest.raises(TypeError): + _parse_canonical_wire_metadata(bytearray(_metadata(b"\0")), limits=limits) + with pytest.raises(TypeError): + _reconstruct_canonical_wire_value(object(), ()) + with pytest.raises(TypeError): + _reconstruct_canonical_wire_value( + _parse_canonical_wire_metadata(_metadata(b"\0"), limits=limits), + iter(()), + ) + with pytest.raises(ValueError): + _CanonicalWireTensorSpec(0, 1, torch.int8, (), 1, 1, b"x" * 32) + assert dataclasses.is_dataclass(_CanonicalWirePlan) + assert dataclasses.is_dataclass(_PreparedCanonicalWireValue) + with pytest.raises(dataclasses.FrozenInstanceError): + limits.chunk_bytes = 1 From b3bed6f97c1a2fc2a8da2f1d8dd67a3ab51b19a9 Mon Sep 17 00:00:00 2001 From: thad0ctor Date: Mon, 13 Jul 2026 00:29:02 -0700 Subject: [PATCH 11/52] Add lossless native shard guards --- docs/optimizer_contracts.md | 2 +- src/gefen/gefen.py | 207 +++++++++++++++-- tests/test_codebook_scope_cpu.py | 275 ++++++++++++++++++++++- tests/test_codebook_scope_distributed.py | 65 +++++- 4 files changed, 521 insertions(+), 28 deletions(-) diff --git a/docs/optimizer_contracts.md b/docs/optimizer_contracts.md index 0d24884..7826cc7 100644 --- a/docs/optimizer_contracts.md +++ b/docs/optimizer_contracts.md @@ -51,7 +51,7 @@ The optimizer owns one learned codebook and therefore accepts one scope. Histogr `initialize_codebook()` and `refresh_codebook()` expose these operations for adapters that enter their optimizers in a deterministic order; `binding.sort_key` supplies the stable process-group portion of that schedule. Normal `step()` still initializes automatically and plain Gefen still honors `codebook_refresh_every`. Every scoped step exchanges a common operation header before any rank-dependent branch, and the first step after binding or native load additionally verifies codebook bytes and the complete manifest. Scoped native AMP requires every member to select the same protocol and present identical `found_inf` and `grad_scale` values; a mismatch raises collectively and requires a group-aware gradient scaler rather than changing external scaler state behind its back. Multi-member explicit scopes reject `capturable=True` because their host validation and process-group collectives are not CUDA-graph-safe. A one-member local scope may initialize during eager warmup and then use ordinary capturable stepping, but manual codebook replacement remains rejected. Ordinary unscoped behavior remains collective-free. Explicit scope does not replace DTensor mesh collectives, AMP mesh preflights, Parallel-Muon ownership collectives, or checkpoint transport groups. -Native checkpoints store a primitive rank-neutral scope record—format version, semantic name, ordered members, and refresh schedule—at the top level and in the existing transport mirror. Scoped group metadata uses a new outer format version so an older loader rejects rather than silently discarding the scope. Runtime process-group handles and collective devices are live adapter configuration and are never serialized or reconstructed. Flattened and whole-owner local payloads additionally carry a primitive rank-local shard-identity guard ordered by native parameter group and slot, including replicated slots in a mixed optimizer and separately listed pruned nonowners; this prevents an equal-shaped checkpoint from another member, logical slice, or parameter ordering from being reinterpreted positionally. 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. Whole-owner checkpoint completeness, scoped DTensor rank-local transport, scope migration, topology-changing canonical I/O, and Hybrid-wide coordination are not claimed. +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. Whole-owner checkpoint completeness, scoped DTensor rank-local transport, scope migration, topology-changing canonical I/O, and Hybrid-wide coordination are not claimed. ## Exact-binding canonical local state diff --git a/src/gefen/gefen.py b/src/gefen/gefen.py index 1b5a25c..3d982b9 100644 --- a/src/gefen/gefen.py +++ b/src/gefen/gefen.py @@ -46,6 +46,10 @@ _gefen_contract, ) from gefen.partitioning import find_period_by_block_variance +from gefen.portable_identity import ( + _parse_shard_identity, + _serialize_shard_identity, +) from gefen.rebinding import LogicalSlotBinding, ParameterRebinding import gefen.quantization as quantization_module from gefen.kernels.automatic_vmean import ( @@ -81,7 +85,8 @@ _RANK_LOCAL_METADATA_VERSION = 3 _CODEBOOK_SCOPE_FORMAT_VERSION = 1 _SCOPED_NATIVE_METADATA_VERSION = 4 -_NATIVE_LOCAL_SHARDS_FORMAT_VERSION = 1 +_NATIVE_LOCAL_SHARDS_LEGACY_FORMAT_VERSION = 1 +_NATIVE_LOCAL_SHARDS_FORMAT_VERSION = 2 _CANONICAL_PARAMETER_STATE_KEYS = frozenset( { "name", @@ -1798,7 +1803,7 @@ def _normalize_serialized_canonical_shard(cls, record): raise ValueError("Gefen canonical shard identity is invalid") from exc return cls._serialized_canonical_shard(shard) - def _serialized_native_local_shards(self): + def _serialized_native_local_shards_v1(self): if self._gefen_codebook_process_group is None: return None if not any( @@ -1829,11 +1834,32 @@ def _serialized_native_local_shards(self): if parameter is None ] return { - "format_version": _NATIVE_LOCAL_SHARDS_FORMAT_VERSION, + "format_version": _NATIVE_LOCAL_SHARDS_LEGACY_FORMAT_VERSION, "param_groups": param_groups, "pruned_shards": pruned_shards, } + def _serialized_native_local_shards(self): + if self._gefen_codebook_process_group is None: + return None + if not any( + shard.layout is not ParameterLayout.REPLICATED + for _, shard in self._gefen_local_shard_bindings + ): + return None + return { + "format_version": _NATIVE_LOCAL_SHARDS_FORMAT_VERSION, + "logical_slots": [ + { + "group_index": slot.group_index, + "original_slot_index": slot.original_slot_index, + "compatibility_name": slot.compatibility_name, + "shard": _serialize_shard_identity(slot.shard), + } + for slot in self._gefen_logical_slots + ], + } + @classmethod def _normalize_serialized_native_local_shard(cls, record): expected = { @@ -1922,9 +1948,7 @@ def _normalize_serialized_native_local_shard(cls, record): return cls._serialized_native_local_shard(shard) @classmethod - def _normalize_serialized_native_local_shards(cls, value): - if value is None: - return None + def _normalize_serialized_native_local_shards_v1(cls, value): if not isinstance(value, dict) or set(value) != { "format_version", "param_groups", @@ -1936,7 +1960,7 @@ def _normalize_serialized_native_local_shards(cls, value): format_version = value["format_version"] if ( type(format_version) is not int - or format_version != _NATIVE_LOCAL_SHARDS_FORMAT_VERSION + or format_version != _NATIVE_LOCAL_SHARDS_LEGACY_FORMAT_VERSION ): raise ValueError( "Unsupported Gefen native local-shard format_version: {}".format( @@ -1973,11 +1997,147 @@ def _normalize_serialized_native_local_shards(cls, value): "Gefen native pruned-shard metadata must describe whole-parameter nonowners" ) return { - "format_version": _NATIVE_LOCAL_SHARDS_FORMAT_VERSION, + "format_version": _NATIVE_LOCAL_SHARDS_LEGACY_FORMAT_VERSION, "param_groups": normalized_groups, "pruned_shards": normalized_pruned, } + @classmethod + def _normalize_serialized_native_local_shards_v2(cls, value): + if type(value) is not dict or set(value) != { + "format_version", + "logical_slots", + }: + raise ValueError( + "Gefen native local-shard metadata has an invalid schema" + ) + format_version = value["format_version"] + if ( + type(format_version) is not int + or format_version != _NATIVE_LOCAL_SHARDS_FORMAT_VERSION + ): + raise ValueError( + "Unsupported Gefen native local-shard format_version: {}".format( + format_version + ) + ) + records = value["logical_slots"] + if type(records) is not list or not records: + raise ValueError( + "Gefen native local-shard logical_slots must be a nonempty list" + ) + + normalized_records = [] + seen_fqns = set() + previous_position = None + supported_layouts = { + ParameterLayout.REPLICATED, + ParameterLayout.FLATTENED_ELEMENT_SHARD, + ParameterLayout.WHOLE_PARAMETER_OWNER, + } + for record in records: + if type(record) is not dict or set(record) != { + "group_index", + "original_slot_index", + "compatibility_name", + "shard", + }: + raise ValueError( + "Gefen native local-shard logical slot has an invalid schema" + ) + if ( + type(record["shard"]) is dict + and "process_group" in record["shard"] + and record["shard"]["process_group"] is None + ): + raise ValueError( + "Gefen native local-shard metadata requires a process-group identity" + ) + try: + shard = _parse_shard_identity(record["shard"]) + slot = LogicalSlotBinding( + record["group_index"], + record["original_slot_index"], + record["compatibility_name"], + shard, + ) + except (TypeError, ValueError) as exc: + raise ValueError( + "Gefen native local-shard logical slot is invalid" + ) from exc + position = (slot.group_index, slot.original_slot_index) + if previous_position is None: + valid_position = position == (0, 0) + elif position[0] == previous_position[0]: + valid_position = position[1] == previous_position[1] + 1 + else: + valid_position = position == (previous_position[0] + 1, 0) + if not valid_position: + raise ValueError( + "Gefen native local-shard logical slot positions must be contiguous" + ) + previous_position = position + if slot.shard.layout not in supported_layouts: + raise ValueError( + "Gefen native local-shard metadata has an unsupported layout" + ) + if slot.shard.process_group is None: + raise ValueError( + "Gefen native local-shard metadata requires a process-group identity" + ) + fqn = slot.shard.parameter.fqn + if fqn in seen_fqns: + raise ValueError( + "Gefen native local-shard logical slot FQNs must be unique" + ) + seen_fqns.add(fqn) + normalized_records.append( + { + "group_index": slot.group_index, + "original_slot_index": slot.original_slot_index, + "compatibility_name": slot.compatibility_name, + "shard": _serialize_shard_identity(slot.shard), + } + ) + return { + "format_version": _NATIVE_LOCAL_SHARDS_FORMAT_VERSION, + "logical_slots": normalized_records, + } + + @classmethod + def _normalize_serialized_native_local_shards(cls, value): + if value is None: + return None + if type(value) is not dict or "format_version" not in value: + raise ValueError( + "Gefen native local-shard metadata has an invalid schema" + ) + format_version = value["format_version"] + if type(format_version) is not int: + raise ValueError( + "Unsupported Gefen native local-shard format_version: {}".format( + format_version + ) + ) + if format_version == _NATIVE_LOCAL_SHARDS_LEGACY_FORMAT_VERSION: + return cls._normalize_serialized_native_local_shards_v1(value) + if format_version == _NATIVE_LOCAL_SHARDS_FORMAT_VERSION: + return cls._normalize_serialized_native_local_shards_v2(value) + raise ValueError( + "Unsupported Gefen native local-shard format_version: {}".format( + format_version + ) + ) + + def _native_local_shards_matches_live(self, value): + if value is None: + return self._serialized_native_local_shards() is None + if value["format_version"] == _NATIVE_LOCAL_SHARDS_LEGACY_FORMAT_VERSION: + return value == self._serialized_native_local_shards_v1() + if value["format_version"] == _NATIVE_LOCAL_SHARDS_FORMAT_VERSION: + return value == self._serialized_native_local_shards() + return False + @staticmethod def _parameter_in(parameters, candidate) -> bool: return any(item is candidate for item in parameters) @@ -7408,9 +7568,7 @@ def _compact(value): if codebook_scope is not None: checkpoint_metadata["codebook_scope"] = self._serialized_codebook_scope() if native_local_shards is not None: - checkpoint_metadata["native_local_shards"] = ( - self._serialized_native_local_shards() - ) + checkpoint_metadata["native_local_shards"] = native_local_shards if consolidate_rank_local: self._consolidate_rank_local_sharded_state( state_dict, checkpoint_metadata @@ -8196,6 +8354,7 @@ def _load_state_dict_impl(self, state_dict): "Gefen checkpoint metadata is present on only some parameter groups" ) first_metadata = group_metadata[0] + normalized_metadata_native_local_shards = [] for metadata in group_metadata: metadata_version = metadata.get("format_version") if metadata_version not in ( @@ -8225,7 +8384,14 @@ def _load_state_dict_impl(self, state_dict): "Gefen checkpoint parameter-group deterministic policy " "must be a bool, got {!r}".format(metadata["deterministic"]) ) - for metadata in group_metadata[1:]: + normalized_metadata_native_local_shards.append( + self._normalize_serialized_native_local_shards( + metadata.get("native_local_shards") + ) + ) + for metadata_index, metadata in enumerate( + group_metadata[1:], start=1 + ): same_version = metadata.get( "format_version" ) == first_metadata.get("format_version") @@ -8248,9 +8414,10 @@ def _load_state_dict_impl(self, state_dict): same_codebook_scope = metadata.get( "codebook_scope" ) == first_metadata.get("codebook_scope") - same_native_local_shards = metadata.get( - "native_local_shards" - ) == first_metadata.get("native_local_shards") + same_native_local_shards = ( + normalized_metadata_native_local_shards[metadata_index] + == normalized_metadata_native_local_shards[0] + ) if ( not same_version or not same_step @@ -8269,11 +8436,9 @@ def _load_state_dict_impl(self, state_dict): metadata_codebook_scope = self._normalize_serialized_codebook_scope( first_metadata.get("codebook_scope") ) - metadata_native_local_shards = ( - self._normalize_serialized_native_local_shards( - first_metadata.get("native_local_shards") - ) - ) + metadata_native_local_shards = normalized_metadata_native_local_shards[ + 0 + ] if gefen_global_step is None: gefen_global_step = metadata_step elif gefen_global_step != metadata_step: @@ -8320,7 +8485,7 @@ def _load_state_dict_impl(self, state_dict): raise ValueError( "Gefen checkpoint codebook scope does not match the live explicit binding" ) - if native_local_shards != self._serialized_native_local_shards(): + if not self._native_local_shards_matches_live(native_local_shards): raise ValueError( "Gefen checkpoint native local-shard identity does not match the live binding" ) diff --git a/tests/test_codebook_scope_cpu.py b/tests/test_codebook_scope_cpu.py index bf0d668..e3c7c13 100644 --- a/tests/test_codebook_scope_cpu.py +++ b/tests/test_codebook_scope_cpu.py @@ -122,6 +122,12 @@ def _assert_snapshot_identity(optimizer, snapshot): assert optimizer.__dict__[key] is value +def _replace_native_guard(checkpoint, guard): + checkpoint["gefen_native_local_shards"] = copy.deepcopy(guard) + for group in checkpoint["param_groups"]: + group["_gefen_checkpoint_metadata"]["native_local_shards"] = copy.deepcopy(guard) + + def test_codebook_process_group_binding_is_public_frozen_and_ordered(): group = ProcessGroupIdentity("replica", ("worker:b", "worker:a")) binding = CodebookProcessGroupBinding(group, "worker:b", object(), torch.device("cpu")) @@ -352,10 +358,12 @@ def test_native_flat_checkpoint_guard_rejects_reordered_equal_shape_slots_atomic source_second.grad = torch.tensor([9.0, -1.0, -2.0, -3.0]) source.step() checkpoint = source.state_dict() - source_slots = checkpoint["gefen_native_local_shards"]["param_groups"] + source_guard = checkpoint["gefen_native_local_shards"] + source_slots = source_guard["logical_slots"] + assert source_guard["format_version"] == 2 assert ( checkpoint["param_groups"][0]["_gefen_checkpoint_metadata"]["native_local_shards"] - == checkpoint["gefen_native_local_shards"] + is source_guard ) buffer = io.BytesIO() torch.save(checkpoint, buffer) @@ -369,9 +377,70 @@ def test_native_flat_checkpoint_guard_rejects_reordered_equal_shape_slots_atomic with pytest.raises(ValueError, match="format_version"): source.load_state_dict(malformed_version) _assert_snapshot_identity(source, source_before) - assert [[record["fqn"] for record in records] for records in source_slots] == [["Model.First", "Model.Second"]] + assert [ + ( + record["group_index"], + record["original_slot_index"], + record["compatibility_name"], + record["shard"]["parameter"]["fqn"], + ) + for record in source_slots + ] == [ + (0, 0, "first", "Model.First"), + (0, 1, "second", "Model.Second"), + ] if second_layout is ParameterLayout.REPLICATED: - assert source_slots[0][1]["layout"] == ParameterLayout.REPLICATED.value + assert source_slots[1]["shard"]["layout"] == ParameterLayout.REPLICATED.value + + legacy_checkpoint = copy.deepcopy(checkpoint) + legacy_guard = source._serialized_native_local_shards_v1() + legacy_checkpoint["gefen_native_local_shards"] = legacy_guard + for group_record in legacy_checkpoint["param_groups"]: + group_record["_gefen_checkpoint_metadata"]["native_local_shards"] = copy.deepcopy(legacy_guard) + legacy_first = torch.nn.Parameter(torch.zeros(4)) + legacy_second = torch.nn.Parameter(torch.zeros(4)) + legacy_target = Gefen( + [("target_first", legacy_first), ("target_second", legacy_second)], + fused=False, + factored_v_2d=False, + ) + legacy_target.post_sharding( + ( + ParameterRebinding(legacy_first, legacy_first, first_shard), + ParameterRebinding(legacy_second, legacy_second, second_shard), + ), + manifest=manifest, + codebook_process_group=_single_member_binding(group), + ) + legacy_target.load_state_dict(copy.deepcopy(legacy_checkpoint)) + assert torch.equal(legacy_target._gefen_codebook, source._gefen_codebook) + assert legacy_target.param_groups[0]["param_names"] == [ + "target_first", + "target_second", + ] + assert legacy_target.state[legacy_first]["name"] == "target_first" + assert legacy_target.state[legacy_second]["name"] == "target_second" + + mirror_only_checkpoint = copy.deepcopy(legacy_checkpoint) + mirror_only_checkpoint.pop("gefen_native_local_shards") + legacy_target.load_state_dict(mirror_only_checkpoint) + assert torch.equal(legacy_target._gefen_codebook, source._gefen_codebook) + + top_only_checkpoint = copy.deepcopy(legacy_checkpoint) + for group_record in top_only_checkpoint["param_groups"]: + group_record.pop("_gefen_checkpoint_metadata") + legacy_target.load_state_dict(top_only_checkpoint) + assert torch.equal(legacy_target._gefen_codebook, source._gefen_codebook) + + mixed_version_checkpoint = copy.deepcopy(checkpoint) + for group_record in mixed_version_checkpoint["param_groups"]: + group_record["_gefen_checkpoint_metadata"]["native_local_shards"] = ( + copy.deepcopy(legacy_guard) + ) + source_before = _snapshot(source) + with pytest.raises(ValueError, match="local shards disagree"): + source.load_state_dict(mixed_version_checkpoint) + _assert_snapshot_identity(source, source_before) target_second = torch.nn.Parameter(torch.zeros(4)) target_first = torch.nn.Parameter(torch.zeros(4)) @@ -396,6 +465,204 @@ def test_native_flat_checkpoint_guard_rejects_reordered_equal_shape_slots_atomic _assert_snapshot_identity(target, target_before) +def test_native_v2_guard_retains_pruned_logical_positions_and_rejects_corruption_atomically( + monkeypatch, +): + monkeypatch.setattr( + Gefen, + "_validate_codebook_runtime_binding", + lambda self, binding: None, + ) + group = ProcessGroupIdentity("data_parallel", ("rank:0", "rank:1")) + + def build(): + owned = torch.nn.Parameter(torch.ones(2, 2)) + remote = torch.nn.Parameter(torch.full((2, 2), 2.0)) + replicated = torch.nn.Parameter(torch.full((2, 2), 3.0)) + optimizer = GefenMuon( + [ + {"params": [("owned", owned), ("remote", remote)]}, + {"params": [("replicated", replicated)]}, + ], + fused=False, + ) + owned_identity = ParameterIdentity("Model.Owned", (2, 2)) + remote_identity = ParameterIdentity("Model.Remote", (2, 2)) + replicated_identity = ParameterIdentity("Model.Replicated", (2, 2)) + owned_records = tuple( + _whole_owner_shard(owned_identity, group, member, "rank:0") + for member in group.ordered_members + ) + remote_records = tuple( + _whole_owner_shard(remote_identity, group, member, "rank:1") + for member in group.ordered_members + ) + replicated_records = tuple( + _replicated_shard(replicated_identity, group, member) + for member in group.ordered_members + ) + local_owned = owned_records[0] + local_remote = remote_records[0] + local_replicated = replicated_records[0] + optimizer.post_sharding( + ( + ParameterRebinding(remote, None, local_remote), + ParameterRebinding(replicated, replicated, local_replicated), + ParameterRebinding(owned, owned, local_owned), + ), + manifest=ShardingManifest( + owned_records + remote_records + replicated_records + ), + codebook_process_group=CodebookProcessGroupBinding( + group, + "rank:0", + object(), + torch.device("cpu"), + ), + ) + return optimizer + + source = build() + checkpoint = source.state_dict() + guard = checkpoint["gefen_native_local_shards"] + assert guard["format_version"] == 2 + assert [ + ( + record["group_index"], + record["original_slot_index"], + record["compatibility_name"], + record["shard"]["parameter"]["fqn"], + ) + for record in guard["logical_slots"] + ] == [ + (0, 0, "owned", "Model.Owned"), + (0, 1, "remote", "Model.Remote"), + (1, 0, "replicated", "Model.Replicated"), + ] + assert source.param_groups[0]["param_names"] == ["owned"] + + matching = build() + matching.load_state_dict(copy.deepcopy(checkpoint)) + assert matching._gefen_logical_slots == source._gefen_logical_slots + + malformed_second_mirror = copy.deepcopy(checkpoint) + second_metadata = copy.deepcopy( + malformed_second_mirror["param_groups"][1][ + "_gefen_checkpoint_metadata" + ] + ) + second_metadata["native_local_shards"]["logical_slots"] = tuple( + second_metadata["native_local_shards"]["logical_slots"] + ) + malformed_second_mirror["param_groups"][1][ + "_gefen_checkpoint_metadata" + ] = second_metadata + malformed_target = build() + malformed_before = _snapshot(malformed_target) + with pytest.raises(ValueError, match="logical_slots must be a nonempty list"): + malformed_target.load_state_dict(malformed_second_mirror) + _assert_snapshot_identity(malformed_target, malformed_before) + + corruptions = [] + + class GuardDict(dict): + pass + + corruptions.append((GuardDict(copy.deepcopy(guard)), "invalid schema")) + + invalid_schema = copy.deepcopy(guard) + invalid_schema["logical_slots"][0]["extra"] = None + corruptions.append((invalid_schema, "invalid schema")) + + bool_position = copy.deepcopy(guard) + bool_position["logical_slots"][0]["group_index"] = False + corruptions.append((bool_position, "logical slot is invalid")) + + skipped_position = copy.deepcopy(guard) + skipped_position["logical_slots"][2]["group_index"] = 2 + corruptions.append((skipped_position, "positions must be contiguous")) + + uppercase_name = copy.deepcopy(guard) + uppercase_name["logical_slots"][1]["compatibility_name"] = "Remote" + corruptions.append((uppercase_name, "logical slot is invalid")) + + duplicate_fqn = copy.deepcopy(guard) + duplicate_fqn["logical_slots"][1]["shard"]["parameter"] = copy.deepcopy( + duplicate_fqn["logical_slots"][0]["shard"]["parameter"] + ) + corruptions.append((duplicate_fqn, "FQNs must be unique")) + + missing_group = copy.deepcopy(guard) + missing_group["logical_slots"][0]["shard"]["process_group"] = None + corruptions.append((missing_group, "requires a process-group identity")) + + changed_layout = copy.deepcopy(guard) + changed_layout["logical_slots"][0]["shard"]["layout"] = ( + ParameterLayout.REPLICATED.value + ) + changed_layout["logical_slots"][0]["shard"]["owner"] = None + changed_layout["logical_slots"][0]["shard"]["placements"][0]["kind"] = ( + PlacementKind.REPLICATE.value + ) + corruptions.append((changed_layout, "does not match")) + + changed_fqn = copy.deepcopy(guard) + changed_fqn["logical_slots"][1]["shard"]["parameter"]["fqn"] = ( + "Model.OtherRemote" + ) + corruptions.append((changed_fqn, "does not match")) + + changed_name = copy.deepcopy(guard) + changed_name["logical_slots"][1]["compatibility_name"] = "other" + corruptions.append((changed_name, "does not match")) + + for corrupted_guard, match in corruptions: + target = build() + before = _snapshot(target) + corrupted = copy.deepcopy(checkpoint) + _replace_native_guard(corrupted, corrupted_guard) + with pytest.raises(ValueError, match=match): + target.load_state_dict(corrupted) + _assert_snapshot_identity(target, before) + + def build_all_nonowner(): + remote = torch.nn.Parameter(torch.full((2, 2), 2.0)) + optimizer = GefenMuon([("remote", remote)], fused=False) + identity = ParameterIdentity("Model.Remote", (2, 2)) + records = tuple( + _whole_owner_shard(identity, group, member, "rank:1") + for member in group.ordered_members + ) + optimizer.post_sharding( + (ParameterRebinding(remote, None, records[0]),), + manifest=ShardingManifest(records), + codebook_process_group=CodebookProcessGroupBinding( + group, + "rank:0", + object(), + torch.device("cpu"), + ), + ) + return optimizer + + all_nonowner = build_all_nonowner() + all_nonowner_checkpoint = all_nonowner.state_dict() + all_nonowner_guard = all_nonowner_checkpoint["gefen_native_local_shards"] + assert all_nonowner.param_groups[0]["params"] == [] + assert [ + ( + record["group_index"], + record["original_slot_index"], + record["compatibility_name"], + record["shard"]["parameter"]["fqn"], + ) + for record in all_nonowner_guard["logical_slots"] + ] == [(0, 0, "remote", "Model.Remote")] + all_nonowner_target = build_all_nonowner() + all_nonowner_target.load_state_dict(copy.deepcopy(all_nonowner_checkpoint)) + assert all_nonowner_target._gefen_logical_slots == all_nonowner._gefen_logical_slots + + def test_unscoped_native_checkpoint_does_not_serialize_a_none_scope(): parameter = torch.nn.Parameter(torch.arange(4, dtype=torch.float32)) optimizer = Gefen([parameter], fused=False) diff --git a/tests/test_codebook_scope_distributed.py b/tests/test_codebook_scope_distributed.py index 5f175b6..4120be4 100644 --- a/tests/test_codebook_scope_distributed.py +++ b/tests/test_codebook_scope_distributed.py @@ -322,6 +322,27 @@ def _distributed_worker(rank, world, init_file, queue): group=runtime_group, ) rank_local_checkpoint_identity = local_shard_records[0] != local_shard_records[1] + flat_v2_guard_valid = ( + flat_checkpoint["gefen_native_local_shards"]["format_version"] == 2 + and len(flat_checkpoint["gefen_native_local_shards"]["logical_slots"]) + == 1 + and flat_checkpoint["gefen_native_local_shards"]["logical_slots"][0][ + "group_index" + ] + == 0 + and flat_checkpoint["gefen_native_local_shards"]["logical_slots"][0][ + "original_slot_index" + ] + == 0 + and flat_checkpoint["gefen_native_local_shards"]["logical_slots"][0][ + "compatibility_name" + ] + == "flat" + and flat_checkpoint["gefen_native_local_shards"]["logical_slots"][0][ + "shard" + ]["local_member"] + == "rank:{}".format(rank) + ) rank_zero_checkpoint = [flat_checkpoint if rank == 0 else None] dist.broadcast_object_list(rank_zero_checkpoint, src=0, group=runtime_group) cross_param = torch.nn.Parameter(flat_param.detach().clone()) @@ -376,6 +397,35 @@ def _distributed_worker(rank, world, init_file, queue): resumed_binding, ) resumed.load_state_dict(flat_checkpoint) + legacy_checkpoint = copy.deepcopy(flat_checkpoint) + legacy_guard = flat_optimizer._serialized_native_local_shards_v1() + legacy_checkpoint["gefen_native_local_shards"] = legacy_guard + for legacy_group in legacy_checkpoint["param_groups"]: + legacy_group["_gefen_checkpoint_metadata"][ + "native_local_shards" + ] = copy.deepcopy(legacy_guard) + legacy_param = torch.nn.Parameter(flat_param.detach().clone()) + legacy_resumed = Gefen( + [("flat", legacy_param)], + fused=False, + factored_v_2d=False, + ) + legacy_binding = _binding(group, rank, runtime_group) + _finalize( + legacy_resumed, + legacy_param, + flat_records[rank], + ShardingManifest(flat_records), + legacy_binding, + ) + legacy_resumed.load_state_dict(legacy_checkpoint) + legacy_v1_guard_accepted = ( + legacy_resumed.codebook_process_group_binding() is legacy_binding + and torch.equal( + legacy_resumed._gefen_codebook, + flat_optimizer._gefen_codebook, + ) + ) canonical_param = torch.nn.Parameter(flat_param.detach().clone()) canonical_resumed = Gefen( [("flat", canonical_param)], @@ -624,6 +674,8 @@ def fail_exact_dp(*args, **kwargs): failure_optimizer._gefen_codebook = restored_codebook original_manifest = failure_optimizer._gefen_sharding_manifest failure_optimizer._gefen_codebook_scope_validated = False + manifest_guard_codebook = failure_optimizer._gefen_codebook + manifest_guard_step = failure_optimizer._gefen_global_step if rank == 0: alternate_identity = ParameterIdentity("Alternate", (8,)) alternate_records = tuple( @@ -633,8 +685,13 @@ def fail_exact_dp(*args, **kwargs): try: failure_optimizer.initialize_codebook() manifest_mismatch_rejected = False - except RuntimeError as exc: - manifest_mismatch_rejected = "manifest" in str(exc) + except RuntimeError: + manifest_mismatch_rejected = True + manifest_mismatch_rejected = ( + manifest_mismatch_rejected + and failure_optimizer._gefen_codebook is manifest_guard_codebook + and failure_optimizer._gefen_global_step == manifest_guard_step + ) failure_optimizer._gefen_sharding_manifest = original_manifest queue.put( @@ -658,6 +715,8 @@ def fail_exact_dp(*args, **kwargs): "flat_checkpoint_continuation": flat_checkpoint_continuation, "rank_neutral_checkpoint_scope": rank_neutral_checkpoint_scope, "rank_local_checkpoint_identity": rank_local_checkpoint_identity, + "flat_v2_guard_valid": flat_v2_guard_valid, + "legacy_v1_guard_accepted": legacy_v1_guard_accepted, "cross_member_guard": cross_member_guard, "canonical_rank_local_identity": canonical_rank_local_identity, "canonical_cross_member_guard": canonical_cross_member_guard, @@ -747,6 +806,8 @@ def test_explicit_gloo_scope_aggregates_logical_state_and_fails_atomically(): assert item["flat_checkpoint_continuation"], item assert item["rank_neutral_checkpoint_scope"], item assert item["rank_local_checkpoint_identity"], item + assert item["flat_v2_guard_valid"], item + assert item["legacy_v1_guard_accepted"], item assert item["cross_member_guard"], item assert item["canonical_rank_local_identity"], item assert item["canonical_cross_member_guard"], item From 599ba33ac0bd33a1e1eeaa43336b429be3384ecd Mon Sep 17 00:00:00 2001 From: thad0ctor Date: Mon, 13 Jul 2026 00:33:14 -0700 Subject: [PATCH 12/52] Add symmetric portable state collectives --- src/gefen/portable_collective.py | 826 ++++++++++++++++++++++++++++++ tests/test_portable_collective.py | 575 +++++++++++++++++++++ 2 files changed, 1401 insertions(+) create mode 100644 src/gefen/portable_collective.py create mode 100644 tests/test_portable_collective.py diff --git a/src/gefen/portable_collective.py b/src/gefen/portable_collective.py new file mode 100644 index 0000000..246934a --- /dev/null +++ b/src/gefen/portable_collective.py @@ -0,0 +1,826 @@ +"""Symmetric tensor-only collectives for bounded canonical-state fragments. + +This module deliberately keeps the transport private. Callers stage values in +callbacks and publish them only after the complete operation returns. +""" + +from __future__ import annotations + +from dataclasses import fields +import hashlib +import hmac +import struct +import sys + +import torch + +from gefen.checkpoint import CheckpointProcessGroupBinding +from gefen.portable_wire import ( + _CanonicalWireLimits, + _CanonicalWirePlan, + _PreparedCanonicalWireValue, + _parse_canonical_wire_metadata, + _prepare_canonical_wire_value, + _reconstruct_canonical_wire_value, +) + + +_I64_MAX = (1 << 63) - 1 +_HEADER_MAGIC = 0x47464E434F4C4C31 # ``GFNCOLL1`` +_VOTE_MAGIC = 0x47464E564F544531 # ``GFNVOTE1`` +_PROTOCOL_VERSION = 1 +_DIGEST_FIELDS = 4 +_DIAGNOSTIC_BYTES = 2048 +_LABEL_BYTES = 1024 +_LIMIT_NAMES = tuple(field.name for field in fields(_CanonicalWireLimits)) + +_H_MAGIC = 0 +_H_VERSION = 1 +_H_STATUS = 2 +_H_COORDINATE = 3 +_H_MEMBERS = 4 +_H_METADATA_BYTES = 5 +_H_TENSOR_COUNT = 6 +_H_TENSOR_BYTES = 7 +_H_TRANSPORT = 8 +_H_OPERATION_BYTES = 9 +_H_TRANSACTION_BYTES = 10 +_H_CALLBACK_FLAGS = 11 +_H_OPERATION_DIGEST = 12 +_H_TRANSACTION_DIGEST = _H_OPERATION_DIGEST + _DIGEST_FIELDS +_H_CONTEXT_DIGEST = _H_TRANSACTION_DIGEST + _DIGEST_FIELDS +_H_IDENTITY_DIGEST = _H_CONTEXT_DIGEST + _DIGEST_FIELDS +_H_METADATA_DIGEST = _H_IDENTITY_DIGEST + _DIGEST_FIELDS +_H_FRAGMENT_DIGEST = _H_METADATA_DIGEST + _DIGEST_FIELDS +_H_LIMITS = _H_FRAGMENT_DIGEST + _DIGEST_FIELDS +_HEADER_FIELDS = _H_LIMITS + len(_LIMIT_NAMES) + +_V_MAGIC = 0 +_V_VERSION = 1 +_V_PHASE = 2 +_V_SOURCE = 3 +_V_STATUS = 4 +_V_COORDINATE = 5 +_V_RESERVED_0 = 6 +_V_RESERVED_1 = 7 +_VOTE_FIELDS = 8 + +_PHASE_METADATA_BUFFER = 1 +_PHASE_METADATA_PARSE = 2 +_PHASE_PAYLOAD_ALLOCATION = 3 +_PHASE_CHUNK_PREPARE = 4 +_PHASE_CHUNK_STAGE = 5 +_PHASE_RECONSTRUCT = 6 +_PHASE_CONSUME = 7 + + +class _CanonicalCollectiveError(RuntimeError): + """A deterministic, live-member-wide canonical transport rejection.""" + + +def _require_signed_header_int(name: str, value: int) -> None: + if type(value) is not int: + raise TypeError("{} must be an int".format(name)) + if value < 0 or value > _I64_MAX: + raise ValueError("{} must fit a nonnegative signed int64 header field".format(name)) + + +def _label_bytes(name: str, value: str) -> bytes: + if type(value) is not str: + raise TypeError("{} must be a string".format(name)) + if not value or value != value.strip() or "\x00" in value: + raise ValueError("{} must be nonempty, trimmed, and contain no NUL".format(name)) + encoded = value.encode("utf-8") + if len(encoded) > _LABEL_BYTES: + raise ValueError("{} exceeds {} UTF-8 bytes".format(name, _LABEL_BYTES)) + return encoded + + +def _digest_fields(value: bytes) -> tuple[int, int, int, int]: + if type(value) is not bytes or len(value) != hashlib.sha256().digest_size: + raise ValueError("digest header fields require exactly 32 bytes") + return tuple( + int.from_bytes(value[start : start + 8], "big", signed=True) + for start in range(0, len(value), 8) + ) + + +def _fields_digest(values: tuple[int, ...] | list[int]) -> bytes: + if len(values) != _DIGEST_FIELDS: + raise ValueError("a digest requires four signed int64 fields") + return b"".join(int(value).to_bytes(8, "big", signed=True) for value in values) + + +def _identity_digest(binding: CheckpointProcessGroupBinding) -> bytes: + identity = binding.identity + hasher = hashlib.sha256() + hasher.update(b"gefen.canonical_collective.identity.v1\0") + hasher.update(struct.pack(">Q", identity.schema_version)) + semantic_name = identity.semantic_name.encode("utf-8") + hasher.update(struct.pack(">Q", len(semantic_name))) + hasher.update(semantic_name) + hasher.update(struct.pack(">Q", len(identity.ordered_members))) + for member in identity.ordered_members: + encoded = member.encode("utf-8") + hasher.update(struct.pack(">Q", len(encoded))) + hasher.update(encoded) + return hasher.digest() + + +def _transport_code(binding: CheckpointProcessGroupBinding) -> int: + if sys.byteorder == "little": + byteorder = 0 + elif sys.byteorder == "big": + byteorder = 1 + else: + raise RuntimeError("unsupported native byte order") + return (0 if binding.collective_device.type == "cpu" else 2) + byteorder + + +def _diagnostic(exc: Exception | None, *, limit: int = _DIAGNOSTIC_BYTES) -> bytes: + if exc is None: + return bytes(_DIAGNOSTIC_BYTES) + usable = max(1, min(limit, _DIAGNOSTIC_BYTES)) + message = "{}: {}".format(type(exc).__name__, str(exc)).replace("\x00", "\\0") + encoded = message.encode("utf-8", errors="replace")[: usable - 1] + return encoded + bytes(_DIAGNOSTIC_BYTES - len(encoded)) + + +def _decode_diagnostic(value: bytes) -> str: + return value.split(b"\0", 1)[0].decode("utf-8", errors="replace") or "unspecified failure" + + +def _runtime_size(binding: CheckpointProcessGroupBinding) -> int: + if len(binding.identity.ordered_members) == 1: + return 1 + import torch.distributed as dist + + if not dist.is_available() or not dist.is_initialized(): + raise RuntimeError("a multi-member canonical collective requires initialized torch.distributed") + return dist.get_world_size(binding.process_group) + + +def _all_gather_fixed( + binding: CheckpointProcessGroupBinding, + values: torch.Tensor, + *, + member_count: int, +) -> tuple[torch.Tensor, ...]: + if member_count == 1: + return (values.detach().cpu(),) + import torch.distributed as dist + + collective = values.to(binding.collective_device) + gathered = [torch.empty_like(collective) for _ in range(member_count)] + dist.all_gather(gathered, collective, group=binding.process_group) + return tuple(value.cpu() for value in gathered) + + +def _exchange_header( + binding: CheckpointProcessGroupBinding, + header: tuple[int, ...], + diagnostic: bytes, + *, + member_count: int, +) -> tuple[tuple[tuple[int, ...], ...], tuple[bytes, ...]]: + header_tensor = torch.tensor(header, dtype=torch.int64) + diagnostic_tensor = torch.frombuffer(bytearray(diagnostic), dtype=torch.uint8) + gathered_headers = _all_gather_fixed(binding, header_tensor, member_count=member_count) + gathered_diagnostics = _all_gather_fixed(binding, diagnostic_tensor, member_count=member_count) + return ( + tuple(tuple(int(item) for item in value.tolist()) for value in gathered_headers), + tuple(bytes(value.numpy().tobytes()) for value in gathered_diagnostics), + ) + + +def _error_from_diagnostics( + phase: str, + diagnostics: tuple[bytes, ...], + failed: tuple[int, ...], +) -> _CanonicalCollectiveError: + details = "; ".join( + "semantic-coordinate[{}]: {}".format(coordinate, _decode_diagnostic(diagnostics[coordinate])) + for coordinate in failed + ) + return _CanonicalCollectiveError("canonical collective {} failed: {}".format(phase, details)) + + +def _phase_vote( + binding: CheckpointProcessGroupBinding, + *, + member_count: int, + coordinate: int, + source: int, + phase: int, + phase_name: str, + error: Exception | None, + diagnostic_limit: int, +) -> None: + vote = [0] * _VOTE_FIELDS + vote[_V_MAGIC] = _VOTE_MAGIC + vote[_V_VERSION] = _PROTOCOL_VERSION + vote[_V_PHASE] = phase + vote[_V_SOURCE] = source + vote[_V_STATUS] = 1 if error is not None else 0 + vote[_V_COORDINATE] = coordinate + headers, diagnostics = _exchange_header( + binding, + tuple(vote), + _diagnostic(error, limit=diagnostic_limit), + member_count=member_count, + ) + for expected_coordinate, header in enumerate(headers): + if len(header) != _VOTE_FIELDS or header[_V_MAGIC] != _VOTE_MAGIC or header[_V_VERSION] != _PROTOCOL_VERSION: + raise _CanonicalCollectiveError("canonical collective received an invalid fixed-size phase vote") + if header[_V_PHASE] != phase or header[_V_SOURCE] != source: + raise _CanonicalCollectiveError("canonical collective phase/source order diverged") + if header[_V_COORDINATE] != expected_coordinate or header[_V_RESERVED_0] != 0 or header[_V_RESERVED_1] != 0: + raise _CanonicalCollectiveError("canonical collective phase vote has an invalid semantic coordinate") + if header[_V_STATUS] not in {0, 1}: + raise _CanonicalCollectiveError("canonical collective phase vote has an invalid status") + failed = tuple(index for index, header in enumerate(headers) if header[_V_STATUS]) + if failed: + raise _error_from_diagnostics( + phase_name, + diagnostics, + failed, + ) + + +def _failure_header(coordinate: int, member_count: int) -> tuple[int, ...]: + header = [0] * _HEADER_FIELDS + header[_H_MAGIC] = _HEADER_MAGIC + header[_H_VERSION] = _PROTOCOL_VERSION + header[_H_STATUS] = 1 + header[_H_COORDINATE] = coordinate + header[_H_MEMBERS] = member_count + return tuple(header) + + +def _success_header( + binding: CheckpointProcessGroupBinding, + plan: _CanonicalWirePlan | None, + limits: _CanonicalWireLimits, + *, + coordinate: int, + member_count: int, + operation: bytes, + transaction: bytes, + context_digest: bytes, + callback_flags: int, +) -> tuple[int, ...]: + if plan is not None and type(plan) is not _CanonicalWirePlan: + raise TypeError("plan must be a _CanonicalWirePlan or None") + metadata_bytes = len(plan.metadata) if plan is not None else 0 + tensor_count = len(plan.tensor_specs) if plan is not None else 0 + tensor_bytes = plan.total_tensor_bytes if plan is not None else 0 + for name in _LIMIT_NAMES: + _require_signed_header_int("limits.{}".format(name), getattr(limits, name)) + if limits.diagnostic_bytes > _DIAGNOSTIC_BYTES: + raise ValueError("limits.diagnostic_bytes exceeds the collective diagnostic capacity") + for name, value in ( + ("member_count", member_count), + ("coordinate", coordinate), + ("metadata bytes", metadata_bytes), + ("tensor count", tensor_count), + ("tensor bytes", tensor_bytes), + ("operation bytes", len(operation)), + ("transaction bytes", len(transaction)), + ("callback flags", callback_flags), + ): + _require_signed_header_int(name, value) + if member_count > limits.max_members: + raise ValueError("checkpoint member count exceeds limits.max_members") + header = [0] * _HEADER_FIELDS + header[_H_MAGIC] = _HEADER_MAGIC + header[_H_VERSION] = _PROTOCOL_VERSION + header[_H_STATUS] = 0 + header[_H_COORDINATE] = coordinate + header[_H_MEMBERS] = member_count + header[_H_METADATA_BYTES] = metadata_bytes + header[_H_TENSOR_COUNT] = tensor_count + header[_H_TENSOR_BYTES] = tensor_bytes + header[_H_TRANSPORT] = _transport_code(binding) + header[_H_OPERATION_BYTES] = len(operation) + header[_H_TRANSACTION_BYTES] = len(transaction) + if callback_flags > 3: + raise ValueError("callback flags contain unknown bits") + header[_H_CALLBACK_FLAGS] = callback_flags + for offset, digest in ( + (_H_OPERATION_DIGEST, hashlib.sha256(b"operation\0" + operation).digest()), + (_H_TRANSACTION_DIGEST, hashlib.sha256(b"transaction\0" + transaction).digest()), + (_H_CONTEXT_DIGEST, context_digest), + (_H_IDENTITY_DIGEST, _identity_digest(binding)), + (_H_METADATA_DIGEST, plan.metadata_digest if plan is not None else bytes(32)), + (_H_FRAGMENT_DIGEST, plan.fragment_digest if plan is not None else bytes(32)), + ): + header[offset : offset + _DIGEST_FIELDS] = _digest_fields(digest) + for index, name in enumerate(_LIMIT_NAMES): + header[_H_LIMITS + index] = getattr(limits, name) + return tuple(header) + + +def _validate_initial_exchange( + binding: CheckpointProcessGroupBinding, + headers: tuple[tuple[int, ...], ...], + diagnostics: tuple[bytes, ...], + *, + member_count: int, +) -> None: + for expected_coordinate, header in enumerate(headers): + if len(header) != _HEADER_FIELDS or header[_H_MAGIC] != _HEADER_MAGIC or header[_H_VERSION] != _PROTOCOL_VERSION: + raise _CanonicalCollectiveError("canonical collective received an invalid fixed-size header") + if header[_H_STATUS] not in {0, 1}: + raise _CanonicalCollectiveError("canonical collective header has an invalid status") + if header[_H_COORDINATE] != expected_coordinate or header[_H_MEMBERS] != member_count: + raise _CanonicalCollectiveError("canonical collective header has an invalid semantic coordinate/order") + if header[_H_TRANSPORT] not in {0, 1, 2, 3} or header[_H_CALLBACK_FLAGS] not in {0, 1, 2, 3}: + raise _CanonicalCollectiveError("canonical collective header has invalid transport/callback flags") + failed = tuple(index for index, header in enumerate(headers) if header[_H_STATUS]) + if failed: + raise _error_from_diagnostics( + "preparation", + diagnostics, + failed, + ) + consensus_ranges = ( + (_H_OPERATION_BYTES, _H_CONTEXT_DIGEST), + (_H_CONTEXT_DIGEST, _H_METADATA_DIGEST), + (_H_LIMITS, _HEADER_FIELDS), + ) + reference = headers[0] + if any( + header[start:stop] != reference[start:stop] + for header in headers[1:] + for start, stop in consensus_ranges + ): + raise _CanonicalCollectiveError( + "canonical collective operation, transaction, context, identity, callback configuration, or limits diverged" + ) + limits = reference[_H_LIMITS:_HEADER_FIELDS] + limit_by_name = dict(zip(_LIMIT_NAMES, limits)) + total_metadata = 0 + total_tensors = 0 + for header in headers: + metadata_bytes = header[_H_METADATA_BYTES] + tensor_count = header[_H_TENSOR_COUNT] + tensor_bytes = header[_H_TENSOR_BYTES] + for name, value in ( + ("metadata bytes", metadata_bytes), + ("tensor count", tensor_count), + ("tensor bytes", tensor_bytes), + ): + _require_signed_header_int(name, value) + if metadata_bytes > limit_by_name["max_metadata_bytes"]: + raise _CanonicalCollectiveError("canonical collective fragment metadata exceeds max_metadata_bytes") + if tensor_count > limit_by_name["max_tensors"]: + raise _CanonicalCollectiveError("canonical collective fragment exceeds max_tensors") + if tensor_bytes > limit_by_name["max_fragment_tensor_bytes"]: + raise _CanonicalCollectiveError("canonical collective fragment exceeds max_fragment_tensor_bytes") + if total_metadata > limit_by_name["max_collective_metadata_bytes"] - metadata_bytes: + raise _CanonicalCollectiveError("canonical collective metadata exceeds max_collective_metadata_bytes") + if total_tensors > limit_by_name["max_collective_tensor_bytes"] - tensor_bytes: + raise _CanonicalCollectiveError("canonical collective tensors exceed max_collective_tensor_bytes") + total_metadata += metadata_bytes + total_tensors += tensor_bytes + + +def _broadcast_tensor( + binding: CheckpointProcessGroupBinding, + value: torch.Tensor, + *, + source: int, + member_count: int, +) -> None: + if member_count == 1: + return + import torch.distributed as dist + + global_source = dist.get_global_rank(binding.process_group, source) + dist.broadcast(value, src=global_source, group=binding.process_group) + + +def _metadata_from_tensor(value: torch.Tensor) -> bytes: + return bytes(value.detach().cpu().numpy().tobytes()) + + +def _canonical_byte_indices(start: int, stop: int, component_bytes: int) -> torch.Tensor: + canonical = torch.arange(start, stop, dtype=torch.int64) + base = torch.div(canonical, component_bytes, rounding_mode="floor") * component_bytes + return base + (component_bytes - 1 - torch.remainder(canonical, component_bytes)) + + +def _flat_native_bytes(value: torch.Tensor) -> torch.Tensor: + return value.reshape(-1).view(torch.uint8).reshape(-1) + + +def _canonical_cpu_chunk(value: torch.Tensor, start: int, stop: int) -> torch.Tensor: + raw = _flat_native_bytes(value) + if sys.byteorder == "little": + return raw[start:stop].clone() + if sys.byteorder != "big": + raise RuntimeError("unsupported native byte order") + component_bytes = value.element_size() // 2 if value.is_complex() else value.element_size() + if component_bytes == 1: + return raw[start:stop].clone() + return raw.index_select(0, _canonical_byte_indices(start, stop, component_bytes)) + + +def _store_canonical_cpu_chunk( + destination: torch.Tensor, + start: int, + stop: int, + canonical: torch.Tensor, +) -> None: + raw = _flat_native_bytes(destination) + if sys.byteorder == "little": + raw[start:stop].copy_(canonical) + return + if sys.byteorder != "big": + raise RuntimeError("unsupported native byte order") + component_bytes = destination.element_size() // 2 if destination.is_complex() else destination.element_size() + if component_bytes == 1: + raw[start:stop].copy_(canonical) + return + raw.index_copy_(0, _canonical_byte_indices(start, stop, component_bytes), canonical) + + +def _header_digest(header: tuple[int, ...], offset: int) -> bytes: + return _fields_digest(header[offset : offset + _DIGEST_FIELDS]) + + +def _receive_metadata( + binding: CheckpointProcessGroupBinding, + local_plan: _CanonicalWirePlan, + headers: tuple[tuple[int, ...], ...], + limits: _CanonicalWireLimits, + *, + coordinate: int, + source: int, + member_count: int, + validate_plan, +) -> _PreparedCanonicalWireValue: + header = headers[source] + metadata_bytes = header[_H_METADATA_BYTES] + metadata_buffer = None + error = None + try: + if coordinate == source: + metadata_buffer = torch.frombuffer(bytearray(local_plan.metadata), dtype=torch.uint8).to( + binding.collective_device + ) + else: + metadata_buffer = torch.empty( + metadata_bytes, + dtype=torch.uint8, + device=binding.collective_device, + ) + if metadata_buffer.numel() != metadata_bytes: + raise ValueError("canonical collective metadata buffer has an invalid size") + except Exception as exc: + error = exc + _phase_vote( + binding, + member_count=member_count, + coordinate=coordinate, + source=source, + phase=_PHASE_METADATA_BUFFER, + phase_name="metadata-buffer allocation", + error=error, + diagnostic_limit=limits.diagnostic_bytes, + ) + assert metadata_buffer is not None + _broadcast_tensor(binding, metadata_buffer, source=source, member_count=member_count) + + prepared = None + error = None + try: + metadata = _metadata_from_tensor(metadata_buffer) + if not hmac.compare_digest(hashlib.sha256(metadata).digest(), _header_digest(header, _H_METADATA_DIGEST)): + raise ValueError("canonical collective metadata digest does not match its source header") + prepared = _parse_canonical_wire_metadata(metadata, limits=limits) + if len(prepared.tensor_specs) != header[_H_TENSOR_COUNT]: + raise ValueError("canonical collective tensor count does not match its source header") + if prepared.total_tensor_bytes != header[_H_TENSOR_BYTES]: + raise ValueError("canonical collective tensor bytes do not match its source header") + if not hmac.compare_digest(prepared.fragment_digest, _header_digest(header, _H_FRAGMENT_DIGEST)): + raise ValueError("canonical collective fragment digest does not match its source header") + if validate_plan is not None: + validate_plan(binding.identity.ordered_members[source], prepared) + except Exception as exc: + error = exc + _phase_vote( + binding, + member_count=member_count, + coordinate=coordinate, + source=source, + phase=_PHASE_METADATA_PARSE, + phase_name="metadata validation", + error=error, + diagnostic_limit=limits.diagnostic_bytes, + ) + assert prepared is not None + return prepared + + +def _allocate_payloads( + binding: CheckpointProcessGroupBinding, + prepared: _PreparedCanonicalWireValue, + limits: _CanonicalWireLimits, + *, + coordinate: int, + source: int, + member_count: int, +) -> tuple[torch.Tensor, ...]: + payloads = None + error = None + try: + payloads = tuple( + torch.empty(spec.shape, dtype=spec.dtype, device="cpu") + for spec in prepared.tensor_specs + ) + except Exception as exc: + error = exc + _phase_vote( + binding, + member_count=member_count, + coordinate=coordinate, + source=source, + phase=_PHASE_PAYLOAD_ALLOCATION, + phase_name="payload allocation", + error=error, + diagnostic_limit=limits.diagnostic_bytes, + ) + assert payloads is not None + return payloads + + +def _transfer_payloads( + binding: CheckpointProcessGroupBinding, + local_plan: _CanonicalWirePlan, + prepared: _PreparedCanonicalWireValue, + payloads: tuple[torch.Tensor, ...], + limits: _CanonicalWireLimits, + *, + coordinate: int, + source: int, + member_count: int, + direct_cpu: bool, +) -> None: + for tensor_index, (spec, destination) in enumerate(zip(prepared.tensor_specs, payloads)): + for start in range(0, spec.nbytes, limits.chunk_bytes): + stop = min(start + limits.chunk_bytes, spec.nbytes) + wire_buffer = None + error = None + try: + if direct_cpu: + wire_buffer = _flat_native_bytes(destination)[start:stop] + if coordinate == source: + source_bytes = _flat_native_bytes(local_plan.tensors[tensor_index]) + wire_buffer.copy_(source_bytes[start:stop]) + else: + wire_buffer = torch.empty( + stop - start, + dtype=torch.uint8, + device=binding.collective_device, + ) + if coordinate == source: + canonical = _canonical_cpu_chunk(local_plan.tensors[tensor_index], start, stop) + wire_buffer.copy_(canonical) + except Exception as exc: + error = exc + _phase_vote( + binding, + member_count=member_count, + coordinate=coordinate, + source=source, + phase=_PHASE_CHUNK_PREPARE, + phase_name="payload chunk {}:{} preparation".format(tensor_index, start), + error=error, + diagnostic_limit=limits.diagnostic_bytes, + ) + assert wire_buffer is not None + _broadcast_tensor(binding, wire_buffer, source=source, member_count=member_count) + if not direct_cpu: + error = None + try: + canonical = wire_buffer.detach().cpu() + _store_canonical_cpu_chunk(destination, start, stop, canonical) + except Exception as exc: + error = exc + _phase_vote( + binding, + member_count=member_count, + coordinate=coordinate, + source=source, + phase=_PHASE_CHUNK_STAGE, + phase_name="payload chunk {}:{} staging".format(tensor_index, start), + error=error, + diagnostic_limit=limits.diagnostic_bytes, + ) + + +def _collective_unanimous_status( + binding: CheckpointProcessGroupBinding, + local_error: Exception | None, + *, + operation: str, + transaction_id: str, + context_digest: bytes, + limits: _CanonicalWireLimits, +) -> None: + """Exchange one fixed-size pre-publication status across all live members. + + This is the lightweight companion to the fragment visitor for import + readiness and freshness gates. It transfers only fixed-size status/header + tensors and diagnostics; it never broadcasts canonical metadata or payloads. + """ + + if not isinstance(binding, CheckpointProcessGroupBinding): + raise TypeError("binding must be a CheckpointProcessGroupBinding") + member_count = _runtime_size(binding) + coordinate = binding.identity.ordered_members.index(binding.local_member) + header = _failure_header(coordinate, member_count) + error = None + diagnostic_limit = _DIAGNOSTIC_BYTES + try: + binding.validate_runtime() + if type(limits) is not _CanonicalWireLimits: + raise TypeError("limits must be _CanonicalWireLimits") + diagnostic_limit = limits.diagnostic_bytes + if local_error is not None and not isinstance(local_error, Exception): + raise TypeError("local_error must be an Exception or None") + operation_bytes = _label_bytes("operation", operation) + transaction_bytes = _label_bytes("transaction_id", transaction_id) + if type(context_digest) is not bytes or len(context_digest) != hashlib.sha256().digest_size: + raise ValueError("context_digest must be exactly 32 bytes") + header = _success_header( + binding, + None, + limits, + coordinate=coordinate, + member_count=member_count, + operation=operation_bytes, + transaction=transaction_bytes, + context_digest=context_digest, + callback_flags=0, + ) + if local_error is not None: + error = local_error + header = _failure_header(coordinate, member_count) + except Exception as exc: + error = exc + headers, diagnostics = _exchange_header( + binding, + header, + _diagnostic(error, limit=diagnostic_limit), + member_count=member_count, + ) + _validate_initial_exchange( + binding, + headers, + diagnostics, + member_count=member_count, + ) + + +def _collective_visit_canonical_fragments( + binding: CheckpointProcessGroupBinding, + local_fragment, + *, + operation: str, + transaction_id: str, + context_digest: bytes, + limits: _CanonicalWireLimits, + validate_plan=None, + consume=None, +) -> None: + """Visit one staged canonical fragment from every semantic member in order. + + ``validate_plan(member, prepared)`` runs after bounded metadata validation + and before payload allocation. ``consume(member, value)`` receives an owned, + digest-verified value. Callbacks must stage only reversible local work; this + transport does not publish or roll back callback side effects. + """ + + if not isinstance(binding, CheckpointProcessGroupBinding): + raise TypeError("binding must be a CheckpointProcessGroupBinding") + member_count = _runtime_size(binding) + members = binding.identity.ordered_members + coordinate = members.index(binding.local_member) + plan = None + header = _failure_header(coordinate, member_count) + error = None + diagnostic_limit = _DIAGNOSTIC_BYTES + try: + binding.validate_runtime() + if type(limits) is not _CanonicalWireLimits: + raise TypeError("limits must be _CanonicalWireLimits") + diagnostic_limit = limits.diagnostic_bytes + if validate_plan is not None and not callable(validate_plan): + raise TypeError("validate_plan must be callable or None") + if consume is not None and not callable(consume): + raise TypeError("consume must be callable or None") + operation_bytes = _label_bytes("operation", operation) + transaction_bytes = _label_bytes("transaction_id", transaction_id) + if type(context_digest) is not bytes or len(context_digest) != hashlib.sha256().digest_size: + raise ValueError("context_digest must be exactly 32 bytes") + plan = _prepare_canonical_wire_value(local_fragment, limits) + header = _success_header( + binding, + plan, + limits, + coordinate=coordinate, + member_count=member_count, + operation=operation_bytes, + transaction=transaction_bytes, + context_digest=context_digest, + callback_flags=(1 if validate_plan is not None else 0) | (2 if consume is not None else 0), + ) + except Exception as exc: + error = exc + headers, diagnostics = _exchange_header( + binding, + header, + _diagnostic(error, limit=diagnostic_limit), + member_count=member_count, + ) + _validate_initial_exchange( + binding, + headers, + diagnostics, + member_count=member_count, + ) + assert plan is not None + assert type(limits) is _CanonicalWireLimits + direct_cpu = all(header[_H_TRANSPORT] == 0 for header in headers) + + for source in range(member_count): + prepared = _receive_metadata( + binding, + plan, + headers, + limits, + coordinate=coordinate, + source=source, + member_count=member_count, + validate_plan=validate_plan, + ) + payloads = _allocate_payloads( + binding, + prepared, + limits, + coordinate=coordinate, + source=source, + member_count=member_count, + ) + _transfer_payloads( + binding, + plan, + prepared, + payloads, + limits, + coordinate=coordinate, + source=source, + member_count=member_count, + direct_cpu=direct_cpu, + ) + value = None + error = None + try: + value = _reconstruct_canonical_wire_value( + prepared, + payloads, + expected_fragment_digest=_header_digest(headers[source], _H_FRAGMENT_DIGEST), + ) + except Exception as exc: + error = exc + _phase_vote( + binding, + member_count=member_count, + coordinate=coordinate, + source=source, + phase=_PHASE_RECONSTRUCT, + phase_name="fragment reconstruction", + error=error, + diagnostic_limit=limits.diagnostic_bytes, + ) + error = None + if consume is not None: + try: + consume(members[source], value) + except Exception as exc: + error = exc + _phase_vote( + binding, + member_count=member_count, + coordinate=coordinate, + source=source, + phase=_PHASE_CONSUME, + phase_name="fragment consumer", + error=error, + diagnostic_limit=limits.diagnostic_bytes, + ) + + +__all__ = [] diff --git a/tests/test_portable_collective.py b/tests/test_portable_collective.py new file mode 100644 index 0000000..cc0ed92 --- /dev/null +++ b/tests/test_portable_collective.py @@ -0,0 +1,575 @@ +from datetime import timedelta +import hashlib +import multiprocessing as mp +import os +import tempfile + +import pytest +import torch +import torch.distributed as dist + +from gefen.checkpoint import CheckpointProcessGroupBinding +from gefen.contracts import ProcessGroupIdentity +import gefen.portable_collective as portable_collective +from gefen.portable_collective import ( + _CanonicalCollectiveError, + _collective_unanimous_status, + _collective_visit_canonical_fragments, +) +from gefen.portable_wire import _CanonicalWireLimits, _PreparedCanonicalWireValue + + +def _limits(**overrides): + values = { + "max_fragment_tensor_bytes": 1 << 20, + "max_collective_tensor_bytes": 4 << 20, + "max_collective_metadata_bytes": 4 << 20, + "chunk_bytes": 5, + "max_members": 8, + "max_metadata_bytes": 1 << 20, + "max_tree_nodes": 1000, + "max_tree_depth": 16, + "max_container_items": 1000, + "max_string_bytes": 4096, + "max_integer_bytes": 128, + "max_tensors": 128, + "max_tensor_rank": 8, + "diagnostic_bytes": 512, + } + values.update(overrides) + return _CanonicalWireLimits(**values) + + +def _singleton_binding(): + identity = ProcessGroupIdentity("checkpoint", ("worker:solo",)) + return CheckpointProcessGroupBinding( + identity, + "worker:solo", + None, + torch.device("cpu"), + ) + + +def _context(value=b"context"): + return hashlib.sha256(value).digest() + + +def test_singleton_visits_owned_digest_verified_value_without_distributed_calls(monkeypatch): + def forbidden(*args, **kwargs): + raise AssertionError("object collectives and distributed calls are forbidden") + + monkeypatch.setattr(dist, "all_gather_object", forbidden) + monkeypatch.setattr(dist, "broadcast_object_list", forbidden) + monkeypatch.setattr(dist, "all_gather", forbidden) + monkeypatch.setattr(dist, "broadcast", forbidden) + original = torch.arange(6, dtype=torch.float32).reshape(2, 3) + plans = [] + consumed = [] + + def validate(member, prepared): + assert member == "worker:solo" + assert type(prepared) is _PreparedCanonicalWireValue + plans.append((member, prepared.total_tensor_bytes, len(prepared.tensor_specs))) + + def consume(member, value): + consumed.append((member, value)) + + _collective_visit_canonical_fragments( + _singleton_binding(), + {"flag": True, "payload": [original]}, + operation="export", + transaction_id="transaction-1", + context_digest=_context(), + limits=_limits(), + validate_plan=validate, + consume=consume, + ) + + assert plans == [("worker:solo", original.numel() * original.element_size(), 1)] + assert [member for member, _ in consumed] == ["worker:solo"] + result = consumed[0][1] + torch.testing.assert_close(result["payload"][0], original, rtol=0, atol=0) + assert result["payload"][0] is not original + assert result["payload"][0].device.type == "cpu" + assert result["payload"][0].is_contiguous() + original.add_(100) + torch.testing.assert_close( + result["payload"][0], + torch.arange(6, dtype=torch.float32).reshape(2, 3), + rtol=0, + atol=0, + ) + + +def test_singleton_transfers_scalar_empty_complex_and_small_chunk_payloads(): + value = { + "bool_scalar": torch.tensor(True), + "complex": torch.tensor([1 + 2j, -3 + 4j], dtype=torch.complex64), + "empty": torch.empty((2, 0, 3), dtype=torch.uint8), + "integer_scalar": torch.tensor(-7, dtype=torch.int64), + } + consumed = [] + _collective_visit_canonical_fragments( + _singleton_binding(), + value, + operation="export", + transaction_id="transaction-dtypes", + context_digest=_context(), + limits=_limits(chunk_bytes=3), + consume=lambda member, result: consumed.append(result), + ) + + assert len(consumed) == 1 + for name, expected in value.items(): + torch.testing.assert_close(consumed[0][name], expected, rtol=0, atol=0) + + +@pytest.mark.parametrize( + ("fragment", "validate_plan", "consume", "message"), + [ + (object(), None, None, "unsupported canonical wire type"), + (None, lambda member, plan: (_ for _ in ()).throw(ValueError("invalid plan")), None, "invalid plan"), + (None, None, lambda member, value: (_ for _ in ()).throw(RuntimeError("cannot stage")), "cannot stage"), + ], +) +def test_singleton_converts_local_preparation_and_callback_failures_to_collective_rejections( + fragment, + validate_plan, + consume, + message, +): + with pytest.raises(_CanonicalCollectiveError, match=message) as caught: + _collective_visit_canonical_fragments( + _singleton_binding(), + fragment, + operation="import", + transaction_id="transaction-2", + context_digest=_context(), + limits=_limits(), + validate_plan=validate_plan, + consume=consume, + ) + assert "semantic-coordinate[0]" in str(caught.value) + + +def test_singleton_does_not_invoke_consumer_after_metadata_corruption(monkeypatch): + original = portable_collective._metadata_from_tensor + consumed = [] + + def corrupt(value): + metadata = bytearray(original(value)) + metadata[0] ^= 1 + return bytes(metadata) + + monkeypatch.setattr(portable_collective, "_metadata_from_tensor", corrupt) + with pytest.raises(_CanonicalCollectiveError, match="metadata digest"): + _collective_visit_canonical_fragments( + _singleton_binding(), + torch.arange(4, dtype=torch.float32), + operation="export", + transaction_id="transaction-corrupt", + context_digest=_context(), + limits=_limits(), + consume=lambda member, value: consumed.append(value), + ) + assert consumed == [] + + +@pytest.mark.parametrize( + ("changes", "message"), + [ + ({"operation": ""}, "operation"), + ({"transaction_id": " bad"}, "transaction_id"), + ({"context_digest": b"short"}, "context_digest"), + ({"validate_plan": object()}, "validate_plan"), + ({"consume": object()}, "consume"), + ], +) +def test_singleton_strict_invocation_fields_fail_through_fixed_status_exchange(changes, message): + arguments = { + "operation": "export", + "transaction_id": "transaction-strict", + "context_digest": _context(), + "limits": _limits(), + } + arguments.update(changes) + with pytest.raises(_CanonicalCollectiveError, match=message): + _collective_visit_canonical_fragments( + _singleton_binding(), + None, + **arguments, + ) + + +def test_singleton_rejects_limits_that_do_not_fit_signed_header_fields(): + too_large = (1 << 63) + limits = _limits( + max_fragment_tensor_bytes=too_large, + max_collective_tensor_bytes=too_large, + ) + with pytest.raises(_CanonicalCollectiveError, match="signed int64"): + _collective_visit_canonical_fragments( + _singleton_binding(), + None, + operation="export", + transaction_id="transaction-overflow", + context_digest=_context(), + limits=limits, + ) + + +def test_unanimous_status_uses_only_fixed_exchanges_and_reports_local_error(monkeypatch): + def forbidden(*args, **kwargs): + raise AssertionError("no distributed collective is allowed for a singleton") + + monkeypatch.setattr(dist, "all_gather", forbidden) + monkeypatch.setattr(dist, "broadcast", forbidden) + monkeypatch.setattr(portable_collective, "_prepare_canonical_wire_value", forbidden) + _collective_unanimous_status( + _singleton_binding(), + None, + operation="import-ready", + transaction_id="transaction-ready", + context_digest=_context(), + limits=_limits(), + ) + with pytest.raises(_CanonicalCollectiveError, match="stale target"): + _collective_unanimous_status( + _singleton_binding(), + ValueError("stale target"), + operation="import-fresh", + transaction_id="transaction-ready", + context_digest=_context(), + limits=_limits(), + ) + + +def _capture_collective_error(callable_object): + try: + callable_object() + except _CanonicalCollectiveError as exc: + return str(exc) + return None + + +def _distributed_worker(rank, init_file, queue): + try: + dist.init_process_group( + "gloo", + init_method="file://{}".format(init_file), + rank=rank, + world_size=2, + timeout=timedelta(seconds=60), + ) + members = ("worker:left", "worker:right") + binding = CheckpointProcessGroupBinding( + ProcessGroupIdentity("checkpoint", members), + members[rank], + dist.group.WORLD, + torch.device("cpu"), + ) + original_all_gather_object = dist.all_gather_object + original_broadcast_object_list = dist.broadcast_object_list + + def forbidden(*args, **kwargs): + raise AssertionError("object collectives are forbidden") + + dist.all_gather_object = forbidden + dist.broadcast_object_list = forbidden + successes = [] + validation_order = [] + + _collective_visit_canonical_fragments( + binding, + { + "coordinate": rank, + "payload": torch.arange(4, dtype=torch.float32) + rank * 10, + }, + operation="export", + transaction_id="distributed-success", + context_digest=_context(), + limits=_limits(chunk_bytes=7), + validate_plan=lambda member, plan: validation_order.append((member, len(plan.tensor_specs))), + consume=lambda member, value: successes.append( + (member, value["coordinate"], value["payload"].tolist()) + ), + ) + dist.barrier() + + context_error = _capture_collective_error( + lambda: _collective_visit_canonical_fragments( + binding, + None, + operation="export", + transaction_id="divergent-context", + context_digest=_context(b"left" if rank == 0 else b"right"), + limits=_limits(), + ) + ) + dist.barrier() + + operation_error = _capture_collective_error( + lambda: _collective_visit_canonical_fragments( + binding, + None, + operation="export" if rank == 0 else "import", + transaction_id="divergent-operation", + context_digest=_context(), + limits=_limits(), + ) + ) + dist.barrier() + + transaction_error = _capture_collective_error( + lambda: _collective_visit_canonical_fragments( + binding, + None, + operation="export", + transaction_id="transaction-left" if rank == 0 else "transaction-right", + context_digest=_context(), + limits=_limits(), + ) + ) + dist.barrier() + + limits_error = _capture_collective_error( + lambda: _collective_visit_canonical_fragments( + binding, + None, + operation="export", + transaction_id="divergent-limits", + context_digest=_context(), + limits=_limits(chunk_bytes=5 + rank), + ) + ) + dist.barrier() + + wrong_binding = CheckpointProcessGroupBinding( + ProcessGroupIdentity("checkpoint", members), + members[1] if rank == 0 else members[rank], + dist.group.WORLD, + torch.device("cpu"), + ) + coordinate_error = _capture_collective_error( + lambda: _collective_visit_canonical_fragments( + wrong_binding, + None, + operation="export", + transaction_id="coordinate-mismatch", + context_digest=_context(), + limits=_limits(), + ) + ) + dist.barrier() + + callback_presence_error = _capture_collective_error( + lambda: _collective_visit_canonical_fragments( + binding, + None, + operation="import", + transaction_id="callback-presence", + context_digest=_context(), + limits=_limits(), + consume=None if rank == 0 else (lambda member, value: None), + ) + ) + dist.barrier() + + prepare_error = _capture_collective_error( + lambda: _collective_visit_canonical_fragments( + binding, + object() if rank == 0 else None, + operation="export", + transaction_id="prepare-failure", + context_digest=_context(), + limits=_limits(), + ) + ) + dist.barrier() + + def validate(member, plan): + if rank == 1 and member == members[0]: + raise ValueError("receiver rejected plan") + + validate_error = _capture_collective_error( + lambda: _collective_visit_canonical_fragments( + binding, + torch.arange(2, dtype=torch.float32), + operation="import", + transaction_id="validate-failure", + context_digest=_context(), + limits=_limits(), + validate_plan=validate, + ) + ) + dist.barrier() + + def consume(member, value): + if rank == 0 and member == members[1]: + raise RuntimeError("receiver could not stage value") + + consume_error = _capture_collective_error( + lambda: _collective_visit_canonical_fragments( + binding, + torch.arange(2, dtype=torch.float32), + operation="import", + transaction_id="consume-failure", + context_digest=_context(), + limits=_limits(), + consume=consume, + ) + ) + dist.barrier() + + budget_error = _capture_collective_error( + lambda: _collective_visit_canonical_fragments( + binding, + torch.arange(4, dtype=torch.float32), + operation="export", + transaction_id="aggregate-budget", + context_digest=_context(), + limits=_limits( + max_fragment_tensor_bytes=16, + max_collective_tensor_bytes=24, + ), + ) + ) + dist.barrier() + + original_broadcast = portable_collective._broadcast_tensor + broadcast_count = 0 + + def corrupt_broadcast(bound_binding, value, *, source, member_count): + nonlocal broadcast_count + broadcast_count += 1 + if rank == 0 and source == 0 and broadcast_count == 2: + value[0] ^= 1 + return original_broadcast( + bound_binding, + value, + source=source, + member_count=member_count, + ) + + portable_collective._broadcast_tensor = corrupt_broadcast + try: + corruption_error = _capture_collective_error( + lambda: _collective_visit_canonical_fragments( + binding, + torch.arange(2, dtype=torch.float32), + operation="export", + transaction_id="payload-corruption", + context_digest=_context(), + limits=_limits(), + ) + ) + finally: + portable_collective._broadcast_tensor = original_broadcast + dist.barrier() + + status_error = _capture_collective_error( + lambda: _collective_unanimous_status( + binding, + ValueError("target changed") if rank == 1 else None, + operation="import-fresh", + transaction_id="status-helper", + context_digest=_context(), + limits=_limits(), + ) + ) + dist.barrier() + + dist.all_gather_object = original_all_gather_object + dist.broadcast_object_list = original_broadcast_object_list + queue.put( + { + "rank": rank, + "successes": successes, + "validation_order": validation_order, + "context_error": context_error, + "operation_error": operation_error, + "transaction_error": transaction_error, + "limits_error": limits_error, + "coordinate_error": coordinate_error, + "callback_presence_error": callback_presence_error, + "prepare_error": prepare_error, + "validate_error": validate_error, + "consume_error": consume_error, + "budget_error": budget_error, + "corruption_error": corruption_error, + "status_error": status_error, + } + ) + except Exception as exc: + queue.put({"rank": rank, "fatal_error": repr(exc)}) + finally: + if dist.is_available() and dist.is_initialized(): + dist.destroy_process_group() + + +def _run_distributed_workers(): + context = mp.get_context("spawn") + queue = context.Queue() + fd, init_file = tempfile.mkstemp(prefix="gefen-portable-collective-") + os.close(fd) + os.unlink(init_file) + processes = [ + context.Process(target=_distributed_worker, args=(rank, init_file, queue)) + for rank in range(2) + ] + try: + for process in processes: + process.start() + results = [queue.get(timeout=120) for _ in processes] + for process in processes: + process.join(timeout=15) + for process in processes: + if process.is_alive(): + process.terminate() + process.join(timeout=5) + pytest.fail("portable-collective worker hung") + assert process.exitcode == 0 + return sorted(results, key=lambda item: item["rank"]) + finally: + for process in processes: + if process.is_alive(): + process.terminate() + process.join(timeout=5) + if os.path.exists(init_file): + os.unlink(init_file) + + +@pytest.mark.skipif( + not dist.is_available() or not dist.is_gloo_available(), + reason="Gloo is required for the canonical collective transport test", +) +def test_two_process_gloo_transport_is_symmetric_bounded_and_tensor_only(): + results = _run_distributed_workers() + + assert all("fatal_error" not in result for result in results), results + expected_successes = [ + ("worker:left", 0, [0.0, 1.0, 2.0, 3.0]), + ("worker:right", 1, [10.0, 11.0, 12.0, 13.0]), + ] + assert all(result["successes"] == expected_successes for result in results) + assert all( + result["validation_order"] == [("worker:left", 1), ("worker:right", 1)] + for result in results + ) + for key, expected in ( + ("context_error", "context"), + ("operation_error", "operation"), + ("transaction_error", "transaction"), + ("limits_error", "limits"), + ("coordinate_error", "semantic coordinate/order"), + ("callback_presence_error", "callback configuration"), + ("prepare_error", "semantic-coordinate[0]"), + ("validate_error", "receiver rejected plan"), + ("consume_error", "receiver could not stage value"), + ("budget_error", "max_collective_tensor_bytes"), + ("corruption_error", "invalid digest"), + ("status_error", "target changed"), + ): + messages = [result[key] for result in results] + assert messages[0] == messages[1] + assert expected in messages[0] From ec78af319c20b1d9b7a07f502fc1c0d459e0fa0e Mon Sep 17 00:00:00 2001 From: thad0ctor Date: Mon, 13 Jul 2026 00:39:11 -0700 Subject: [PATCH 13/52] Preserve signed zero in portable momentum --- src/gefen/portable.py | 13 +++++++++++++ tests/test_portable_state_math.py | 4 +++- 2 files changed, 16 insertions(+), 1 deletion(-) diff --git a/src/gefen/portable.py b/src/gefen/portable.py index 118ed56..23fac0c 100644 --- a/src/gefen/portable.py +++ b/src/gefen/portable.py @@ -296,6 +296,11 @@ def _recompress_dense_momentum( if momentum.device != codebook.device: raise ValueError("dense momentum and codebook must share a device") period = _validate_period(period, numel=momentum.numel()) + preserve_zero_sign = ( + period == 1 + and bool(codebook[0] == -1.0) + and bool(codebook[-1] == 1.0) + ) blocks = momentum.numel() // period indices = _new_output( @@ -322,6 +327,14 @@ def _recompress_dense_momentum( normalized.div_(magnitude_chunk) normalized.masked_fill_(~nonzero, 0.0) index_chunk = _nearest_codebook_indices(codebook, normalized) + if preserve_zero_sign: + zero_values = block_values == 0 + zero_indices = torch.where( + torch.signbit(block_values), + torch.zeros_like(index_chunk), + torch.full_like(index_chunk, codebook.numel() - 1), + ) + index_chunk = torch.where(zero_values, zero_indices, index_chunk) indices[row_start:row_stop].copy_(index_chunk) magnitudes[row_start:row_stop].copy_(magnitude_chunk) return ( diff --git a/tests/test_portable_state_math.py b/tests/test_portable_state_math.py index b7c83d0..a8e2d7b 100644 --- a/tests/test_portable_state_math.py +++ b/tests/test_portable_state_math.py @@ -142,10 +142,12 @@ def test_period_one_recompression_is_exact_for_every_finite_fp32_scale(): step=99, ) - torch.testing.assert_close(reconstructed, momentum, rtol=0, atol=0) + assert torch.equal(reconstructed.view(torch.int32), momentum.view(torch.int32)) assert torch.equal(magnitudes.reshape(-1), momentum.abs()) assert torch.equal(indices[momentum < 0], torch.zeros_like(indices[momentum < 0])) assert torch.equal(indices[momentum > 0], torch.full_like(indices[momentum > 0], 4)) + assert indices[3].item() == 0 + assert indices[4].item() == codebook.numel() - 1 def test_scalar_momentum_roundtrips_as_zero_dimensional_logical_state(): From 1aaf9876638100309872220da106eddcb9f6926f Mon Sep 17 00:00:00 2001 From: thad0ctor Date: Mon, 13 Jul 2026 00:43:43 -0700 Subject: [PATCH 14/52] Make period-one state projection bit exact --- src/gefen/portable.py | 7 +++++++ tests/test_portable_state_math.py | 21 +++++++++++++++++++++ 2 files changed, 28 insertions(+) diff --git a/src/gefen/portable.py b/src/gefen/portable.py index 23fac0c..d759a94 100644 --- a/src/gefen/portable.py +++ b/src/gefen/portable.py @@ -405,6 +405,13 @@ def _reduce_block_second_moment( dtype=torch.float32, device=dense.device, ) + if period == 1: + reduced_flat = reduced.reshape(-1) + for start, stop in _element_chunks(dense.numel()): + reduced_flat[start:stop].copy_( + _read_flat_chunk(dense, start, stop) + ) + return _finish_output(reduced, name="block second moment") for row_start, row_stop in _whole_row_chunks(blocks, period): flat_start = row_start * period flat_stop = row_stop * period diff --git a/tests/test_portable_state_math.py b/tests/test_portable_state_math.py index a8e2d7b..b5bfdea 100644 --- a/tests/test_portable_state_math.py +++ b/tests/test_portable_state_math.py @@ -150,6 +150,27 @@ def test_period_one_recompression_is_exact_for_every_finite_fp32_scale(): assert indices[4].item() == codebook.numel() - 1 +def test_period_one_block_reduction_preserves_every_fp32_bit(): + tiny = torch.nextafter(torch.tensor(0.0), torch.tensor(1.0)) + dense = torch.stack( + ( + torch.tensor(-0.0), + torch.tensor(0.0), + tiny, + torch.tensor(torch.finfo(torch.float32).tiny), + torch.tensor(19.125), + torch.tensor(torch.finfo(torch.float32).max), + ) + ) + + reduced = _reduce_block_second_moment(dense, period=1, step=99) + + assert torch.equal( + reduced.reshape(-1).view(torch.int32), + dense.view(torch.int32), + ) + + def test_scalar_momentum_roundtrips_as_zero_dimensional_logical_state(): momentum = torch.tensor(-123.75, requires_grad=True) indices, magnitudes = _recompress_dense_momentum( From 139790330454fe30adeb0f6662fe372ffaba91da Mon Sep 17 00:00:00 2001 From: thad0ctor Date: Mon, 13 Jul 2026 01:11:36 -0700 Subject: [PATCH 15/52] Add exact portable state semantics --- src/gefen/portable_state.py | 1171 ++++++++++++++++++++++++++++++++++ tests/test_portable_state.py | 874 +++++++++++++++++++++++++ 2 files changed, 2045 insertions(+) create mode 100644 src/gefen/portable_state.py create mode 100644 tests/test_portable_state.py diff --git a/src/gefen/portable_state.py b/src/gefen/portable_state.py new file mode 100644 index 0000000..e28f722 --- /dev/null +++ b/src/gefen/portable_state.py @@ -0,0 +1,1171 @@ +"""Exact semantic normalization, assembly, and projection for portable Gefen state.""" + +from __future__ import annotations + +from dataclasses import dataclass +import math +import struct + +import torch + +from gefen.contracts import ( + ParameterLayout, + ProcessGroupIdentity, + ShardIdentity, + ShardingManifest, +) +from gefen.portable import _recompress_dense_momentum, _reduce_block_second_moment, _validate_codebook +from gefen.portable_fields import _assemble_dense_logical_field, _project_dense_logical_field, _tensor_bits_equal +from gefen.portable_identity import ( + _normalize_parameter_identity, + _normalize_shard_identity, + _normalize_sharding_manifest, + _parse_parameter_identity, + _parse_shard_identity, + _parse_sharding_manifest, + _serialize_parameter_identity, + _serialize_sharding_manifest, +) +from gefen.portable_schema import build_portable_state_document, normalize_portable_state_document +from gefen.portable_wire import ( + _CanonicalWireLimits, + _parse_canonical_wire_metadata, + _prepare_canonical_wire_value, + _reconstruct_canonical_wire_value, +) + + +_IMPLEMENTATIONS = frozenset({"gefen.Gefen", "gefen.GefenMuon"}) +_FRAGMENT_FORMAT = "gefen.portable_state_fragment" +_FRAGMENT_FORMAT_VERSION = 1 +_FRAGMENT_COVERAGE = "local_logical_optimizer_fragment" +_MOMENTUM_PROJECTION = "dense_fp32_target_period_one_v1" +_SECOND_MOMENT_PROJECTION = "exact_representation_target_period_one_v1" +_MAX_EXACT_COUNTER = (1 << 53) - 1 + +_FRAGMENT_KEYS = frozenset( + { + "format", + "format_version", + "coverage", + "implementation", + "member", + "policy", + "common", + "manifest", + "catalog", + "logical_slots", + } +) +_CATALOG_KEYS = frozenset({"identity", "algorithm_options"}) +_SLOT_KEYS = frozenset( + { + "group_index", + "original_slot_index", + "compatibility_name", + "shard", + "algorithm_options", + "role", + "source_period", + "source_second_moment", + "state_variant", + "state", + } +) +_COMMON_KEYS = frozenset({"gefen_global_step", "gefen_codebook", "gefen_deterministic"}) +_POLICY_KEYS = frozenset( + { + "schema_version", + "factored_v_2d", + "force_1d_period_one", + "force_2d_period_one", + "period_one_substrings", + "codebook_refresh_every", + "stochastic_round", + "momentum_projection", + "second_moment_projection", + } +) +_PLAIN_OPTION_KEYS = frozenset({"lr", "beta1", "beta2", "eps", "weight_decay", "second_moment_policy"}) +_MUON_OPTION_KEYS = frozenset( + { + "lr", + "weight_decay", + "momentum", + "nesterov", + "ns_schedule", + "ns_eps", + "adjust_lr_fn", + "sharded_mode", + "fp8_ns", + "fp8_ns_compile", + "batched_ns", + "batched_ns_workspace_bytes", + "normuon", + "normuon_beta2", + "normuon_eps", + "cautious", + } +) +_PLAIN_HINT_KEYS = frozenset({"source_periods", "source_second_moment", "target_period"}) +_MUON_HINT_KEYS = frozenset({"source_periods", "target_period"}) +_PARAMETER_KEYS = frozenset({"identity", "algorithm_options", "state_variant", "state", "projection_hints"}) + + +@dataclass(frozen=True, slots=True) +class PortableStateLimits: + """Public resource limits shared by portable semantic and collective I/O.""" + + max_fragment_tensor_bytes: int + max_collective_tensor_bytes: int + max_collective_metadata_bytes: int + chunk_bytes: int = 8 << 20 + max_members: int = 4096 + max_metadata_bytes: int = 64 << 20 + max_tree_nodes: int = 1_000_000 + max_tree_depth: int = 64 + max_container_items: int = 1_000_000 + max_string_bytes: int = 1 << 20 + max_integer_bytes: int = 4096 + max_tensors: int = 262_144 + max_tensor_rank: int = 64 + diagnostic_bytes: int = 2048 + + def __post_init__(self) -> None: + self._wire_limits() + + def _wire_limits(self, *, collective: bool = False) -> _CanonicalWireLimits: + return _CanonicalWireLimits( + max_fragment_tensor_bytes=( + self.max_collective_tensor_bytes if collective else self.max_fragment_tensor_bytes + ), + max_collective_tensor_bytes=self.max_collective_tensor_bytes, + max_collective_metadata_bytes=self.max_collective_metadata_bytes, + chunk_bytes=self.chunk_bytes, + max_members=self.max_members, + max_metadata_bytes=(self.max_collective_metadata_bytes if collective else self.max_metadata_bytes), + max_tree_nodes=self.max_tree_nodes, + max_tree_depth=self.max_tree_depth, + max_container_items=self.max_container_items, + max_string_bytes=self.max_string_bytes, + max_integer_bytes=self.max_integer_bytes, + max_tensors=self.max_tensors, + max_tensor_rank=self.max_tensor_rank, + diagnostic_bytes=self.diagnostic_bytes, + ) + + +def _require_limits(limits) -> PortableStateLimits: + if type(limits) is not PortableStateLimits: + raise TypeError("limits must be a PortableStateLimits") + return limits + + +def _bounded_clone(value, limits: PortableStateLimits, *, collective: bool = False): + wire_limits = limits._wire_limits(collective=collective) + plan = _prepare_canonical_wire_value(value, wire_limits) + prepared = _parse_canonical_wire_metadata(plan.metadata, limits=wire_limits) + return _reconstruct_canonical_wire_value( + prepared, + plan.payload_tensors, + expected_fragment_digest=plan.fragment_digest, + ) + + +def _exact_record(value, keys, *, name: str): + if type(value) is not dict or set(value) != keys: + raise ValueError("{} has an invalid schema".format(name)) + return value + + +def _strict_int(value, *, name: str, minimum: int = 0, maximum=None) -> int: + if type(value) is not int: + raise ValueError("{} must be an int".format(name)) + if value < minimum: + raise ValueError("{} must be at least {}".format(name, minimum)) + if maximum is not None and value > maximum: + raise ValueError("{} must be at most {}".format(name, maximum)) + return value + + +def _strict_float(value, *, name: str, minimum=None, maximum=None, maximum_open=False) -> float: + if type(value) is not float or not math.isfinite(value): + raise ValueError("{} must be a finite float".format(name)) + if minimum is not None and value < minimum: + raise ValueError("{} is below its minimum".format(name)) + if maximum is not None and (value > maximum or (maximum_open and value == maximum)): + raise ValueError("{} is above its maximum".format(name)) + return value + + +def _strict_name(value, *, name: str) -> str: + if type(value) is not str or not value or value != value.strip() or "\x00" in value: + raise ValueError("{} must be a non-empty canonical string".format(name)) + return value + + +def _float_bits_equal(left: float, right: float) -> bool: + return struct.pack(">d", left) == struct.pack(">d", right) + + +def _values_equal(left, right) -> bool: + if type(left) is not type(right): + return False + if type(left) is torch.Tensor: + return left.dtype == right.dtype and tuple(left.shape) == tuple(right.shape) and _tensor_bits_equal(left, right) + if type(left) is float: + return _float_bits_equal(left, right) + if type(left) is dict: + return set(left) == set(right) and all(_values_equal(left[key], right[key]) for key in left) + if type(left) in {list, tuple}: + return len(left) == len(right) and all(_values_equal(a, b) for a, b in zip(left, right)) + return left == right + + +def _tight_fp32(value, *, name: str, shape=None, nonnegative: bool = False) -> torch.Tensor: + if ( + type(value) is not torch.Tensor + or value.layout is not torch.strided + or value.device.type != "cpu" + or value.dtype != torch.float32 + or value.is_meta + or value.is_nested + or value.is_quantized + or value.requires_grad + or not value.is_contiguous() + or value.storage_offset() != 0 + or value.untyped_storage().nbytes() != value.numel() * value.element_size() + ): + raise ValueError("{} must be a tight detached CPU fp32 tensor".format(name)) + if shape is not None and tuple(value.shape) != tuple(shape): + raise ValueError("{} has invalid shape".format(name)) + if not bool(torch.isfinite(value).all()): + raise ValueError("{} must be finite".format(name)) + if nonnegative and not bool((value >= 0).all()): + raise ValueError("{} must be nonnegative".format(name)) + return value + + +def _tight_clone(value: torch.Tensor) -> torch.Tensor: + result = torch.empty(tuple(value.shape), dtype=torch.float32, device="cpu") + result.copy_(value) + return result + + +def _normalize_codebook(value, *, global_step: int): + if value is None: + return None + value = _tight_fp32(value, name="common.gefen_codebook", shape=(256,)) + _validate_codebook(value) + if float(value[0].item()) != -1.0 or float(value[-1].item()) != 1.0: + raise ValueError("portable codebook must retain exact -1 and +1 endpoints") + return value + + +def _normalize_common(value): + value = _exact_record(value, _COMMON_KEYS, name="portable common state") + global_step = _strict_int(value["gefen_global_step"], name="gefen_global_step", maximum=_MAX_EXACT_COUNTER) + if type(value["gefen_deterministic"]) is not bool: + raise ValueError("gefen_deterministic must be a bool") + return { + "gefen_global_step": global_step, + "gefen_codebook": _normalize_codebook(value["gefen_codebook"], global_step=global_step), + "gefen_deterministic": value["gefen_deterministic"], + } + + +def _normalize_policy(value, implementation: str): + value = _exact_record(value, _POLICY_KEYS, name="portable policy") + if value["schema_version"] != 1 or type(value["schema_version"]) is not int: + raise ValueError("unsupported portable semantic policy schema_version") + _strict_int(value["codebook_refresh_every"], name="codebook_refresh_every") + if value["stochastic_round"] is not False: + raise ValueError("portable semantic state requires stochastic_round=False") + if value["momentum_projection"] != _MOMENTUM_PROJECTION: + raise ValueError("unsupported portable momentum projection") + if value["second_moment_projection"] != _SECOND_MOMENT_PROJECTION: + raise ValueError("unsupported portable second-moment projection") + for key in ("factored_v_2d", "force_1d_period_one", "force_2d_period_one"): + if type(value[key]) is not bool: + raise ValueError("{} must be a bool".format(key)) + if implementation == "gefen.GefenMuon" and value["factored_v_2d"]: + raise ValueError("Muon portable policy requires factored_v_2d=False") + substrings = value["period_one_substrings"] + if type(substrings) is not list or any(type(item) is not str or item != item.lower() for item in substrings): + raise ValueError("period_one_substrings must be a canonical lowercase string list") + return {**value, "period_one_substrings": list(substrings)} + + +def _normalize_ns_schedule(value): + if type(value) is not list or not value or len(value) >= 100: + raise ValueError("ns_schedule must contain between 1 and 99 entries") + if any(type(item) is not list or len(item) != 3 for item in value): + raise ValueError("ns_schedule must contain three-float lists") + return [[_strict_float(component, name="ns_schedule") for component in item] for item in value] + + +def _normalize_options(value, implementation: str): + keys = _PLAIN_OPTION_KEYS if implementation == "gefen.Gefen" else _MUON_OPTION_KEYS + value = _exact_record(value, keys, name="portable algorithm options") + result = dict(value) + result["lr"] = _strict_float(value["lr"], name="lr", minimum=0.0) + result["weight_decay"] = _strict_float(value["weight_decay"], name="weight_decay", minimum=0.0) + if implementation == "gefen.Gefen": + result["beta1"] = _strict_float(value["beta1"], name="beta1", minimum=0.0, maximum=1.0, maximum_open=True) + result["beta2"] = _strict_float(value["beta2"], name="beta2", minimum=0.0, maximum=1.0, maximum_open=True) + result["eps"] = _strict_float(value["eps"], name="eps", minimum=0.0) + if result["eps"] == 0.0: + raise ValueError("eps must be positive") + if value["second_moment_policy"] not in {"block", "factored"}: + raise ValueError("second_moment_policy must be 'block' or 'factored'") + return result + + result["momentum"] = _strict_float(value["momentum"], name="momentum", minimum=0.0, maximum=1.0, maximum_open=True) + for key in ("nesterov", "fp8_ns", "fp8_ns_compile", "batched_ns", "normuon", "cautious"): + if type(value[key]) is not bool: + raise ValueError("{} must be a bool".format(key)) + result["ns_schedule"] = _normalize_ns_schedule(value["ns_schedule"]) + result["ns_eps"] = _strict_float(value["ns_eps"], name="ns_eps", minimum=0.0) + if result["ns_eps"] == 0.0: + raise ValueError("ns_eps must be positive") + if value["adjust_lr_fn"] not in {None, "original", "match_rms_adamw"}: + raise ValueError("adjust_lr_fn is unsupported") + if value["sharded_mode"] not in {"exact", "approx", "distributed"}: + raise ValueError("sharded_mode is unsupported") + result["batched_ns_workspace_bytes"] = _strict_int( + value["batched_ns_workspace_bytes"], name="batched_ns_workspace_bytes", minimum=1 + ) + result["normuon_beta2"] = _strict_float( + value["normuon_beta2"], name="normuon_beta2", minimum=0.0, maximum=1.0, maximum_open=True + ) + result["normuon_eps"] = _strict_float(value["normuon_eps"], name="normuon_eps", minimum=0.0) + if result["normuon_eps"] == 0.0: + raise ValueError("normuon_eps must be positive") + return result + + +def _normalize_periods(value, *, name: str): + if type(value) is not list or any(type(period) is not int or period <= 0 for period in value): + raise ValueError("{} must be a list of positive ints".format(name)) + if value != sorted(set(value)): + raise ValueError("{} must be sorted and unique".format(name)) + if any(period != 1 for period in value): + raise ValueError("exact portable v3 state supports only source period one") + return list(value) + + +def _normalize_complete_parameter_record(fqn, value, implementation: str, global_step: int, policy): + value = _exact_record(value, _PARAMETER_KEYS, name="portable parameter {!r}".format(fqn)) + identity_record = _normalize_parameter_identity(value["identity"]) + identity = _parse_parameter_identity(identity_record) + if identity.fqn != fqn: + raise ValueError("portable parameter identity does not match its FQN key") + options = _normalize_options(value["algorithm_options"], implementation) + variant = value["state_variant"] + if type(variant) is not str: + raise ValueError("portable state_variant must be a string") + state = value["state"] + hints = value["projection_hints"] + if type(state) is not dict or type(hints) is not dict: + raise ValueError("portable parameter state and projection_hints must be dictionaries") + shape = identity.global_shape + if identity.numel == 0 and variant != "pristine": + raise ValueError("empty logical parameters must remain pristine") + + if implementation == "gefen.Gefen": + _exact_record(hints, _PLAIN_HINT_KEYS, name="plain projection_hints") + periods = _normalize_periods(hints["source_periods"], name="source_periods") + if hints["target_period"] != 1 or type(hints["target_period"]) is not int: + raise ValueError("portable target_period must be exactly one") + if hints["source_second_moment"] not in {None, "block", "factored"}: + raise ValueError("portable source_second_moment is invalid") + expected_representation = "factored" if policy["factored_v_2d"] and len(shape) == 2 else "block" + if options["second_moment_policy"] != expected_representation: + raise ValueError("parameter second_moment_policy conflicts with the optimizer policy") + if variant == "pristine": + expected_keys = frozenset() + if periods or hints["source_second_moment"] is not None: + raise ValueError("pristine portable state has invalid projection hints") + elif variant == "period_selected": + expected_keys = frozenset() + if not periods or hints["source_second_moment"] is not None: + raise ValueError("period-selected portable state has invalid projection hints") + elif variant == "initialized_dense": + expected_keys = frozenset({"step", "momentum", "second_moment", "second_moment_step"}) + if not periods or hints["source_second_moment"] != "block" or expected_representation != "block": + raise ValueError("dense block state conflicts with its policy or projection hints") + elif variant == "initialized_factored": + expected_keys = frozenset({"step", "momentum", "v_row", "v_col", "factored_step"}) + if not periods or hints["source_second_moment"] != "factored" or expected_representation != "factored": + raise ValueError("factored state conflicts with its policy or projection hints") + else: + raise ValueError("unsupported plain portable state_variant") + _exact_record(state, expected_keys, name="plain portable parameter state") + normalized_state = dict(state) + if variant.startswith("initialized_"): + if identity.numel == 0: + raise ValueError("empty logical parameters cannot carry initialized state") + step = _strict_int(state["step"], name="step", minimum=1, maximum=_MAX_EXACT_COUNTER) + if step > global_step: + raise ValueError("parameter step exceeds optimizer global step") + normalized_state["momentum"] = _tight_fp32(state["momentum"], name="momentum", shape=shape) + if variant == "initialized_dense": + second_step = _strict_int( + state["second_moment_step"], name="second_moment_step", minimum=1, maximum=_MAX_EXACT_COUNTER + ) + normalized_state["second_moment"] = _tight_fp32( + state["second_moment"], name="second_moment", shape=shape, nonnegative=True + ) + else: + if len(shape) != 2: + raise ValueError("factored portable state requires a logical matrix") + second_step = _strict_int( + state["factored_step"], name="factored_step", minimum=1, maximum=_MAX_EXACT_COUNTER + ) + normalized_state["v_row"] = _tight_fp32( + state["v_row"], name="v_row", shape=(shape[0],), nonnegative=True + ) + normalized_state["v_col"] = _tight_fp32( + state["v_col"], name="v_col", shape=(shape[1],), nonnegative=True + ) + if second_step > step: + raise ValueError("secondary parameter counter exceeds step") + return { + "identity": identity_record, + "algorithm_options": options, + "state_variant": variant, + "state": normalized_state, + "projection_hints": { + "source_periods": periods, + "source_second_moment": hints["source_second_moment"], + "target_period": 1, + }, + } + + _exact_record(hints, _MUON_HINT_KEYS, name="Muon projection_hints") + periods = _normalize_periods(hints["source_periods"], name="source_periods") + if hints["target_period"] != 1 or type(hints["target_period"]) is not int: + raise ValueError("portable target_period must be exactly one") + if len(shape) != 2: + raise ValueError("portable Muon parameters must be logical matrices") + if variant == "pristine": + expected_keys = frozenset() + if periods: + raise ValueError("pristine Muon state must not contain source periods") + elif variant == "period_selected": + expected_keys = frozenset() + if not periods: + raise ValueError("period-selected Muon state requires source periods") + elif variant == "initialized_dense": + expected_keys = frozenset({"step", "momentum"}) + if options["normuon"] or not periods: + raise ValueError("initialized Muon state conflicts with its options or periods") + elif variant == "initialized_dense_normuon": + expected_keys = frozenset({"step", "momentum", "normuon_v", "normuon_step"}) + if not options["normuon"] or not periods: + raise ValueError("initialized NorMuon state conflicts with its options or periods") + else: + raise ValueError("unsupported Muon portable state_variant") + _exact_record(state, expected_keys, name="Muon portable parameter state") + normalized_state = dict(state) + if variant.startswith("initialized_"): + if identity.numel == 0: + raise ValueError("empty logical parameters cannot carry initialized state") + step = _strict_int(state["step"], name="step", minimum=1, maximum=_MAX_EXACT_COUNTER) + if step > global_step: + raise ValueError("parameter step exceeds optimizer global step") + normalized_state["momentum"] = _tight_fp32(state["momentum"], name="momentum", shape=shape) + if variant == "initialized_dense_normuon": + normuon_step = _strict_int( + state["normuon_step"], name="normuon_step", minimum=1, maximum=_MAX_EXACT_COUNTER + ) + if normuon_step > step: + raise ValueError("normuon_step exceeds parameter step") + normalized_state["normuon_v"] = _tight_fp32( + state["normuon_v"], name="normuon_v", shape=(shape[0], 1), nonnegative=True + ) + return { + "identity": identity_record, + "algorithm_options": options, + "state_variant": variant, + "state": normalized_state, + "projection_hints": {"source_periods": periods, "target_period": 1}, + } + + +def _normalize_gefen_portable_state_document(state, *, limits, expected_implementation=None): + """Bound, clone, and semantically validate one complete portable v3 document.""" + + limits = _require_limits(limits) + if expected_implementation is not None and expected_implementation not in _IMPLEMENTATIONS: + raise ValueError("unsupported expected portable implementation") + bounded = _bounded_clone(state, limits, collective=True) + document = normalize_portable_state_document(bounded, expected_implementation=expected_implementation) + implementation = document["implementation"] + if implementation not in _IMPLEMENTATIONS: + raise ValueError("unsupported portable Gefen implementation") + if document["provenance"] is not None: + raise ValueError("Gefen portable v3 provenance must be None") + policy = _normalize_policy(document["policy"], implementation) + common = _normalize_common(document["common"]) + parameters = document["parameters"] + if type(parameters) is not dict or not parameters: + raise ValueError("portable parameters must be a non-empty FQN mapping") + normalized_parameters = {} + for fqn in sorted(parameters): + _strict_name(fqn, name="parameter FQN") + normalized_parameters[fqn] = _normalize_complete_parameter_record( + fqn, + parameters[fqn], + implementation, + common["gefen_global_step"], + policy, + ) + if common["gefen_codebook"] is None and any( + record["state_variant"] != "pristine" for record in normalized_parameters.values() + ): + raise ValueError("non-pristine portable parameter state requires a codebook") + normalized = build_portable_state_document( + implementation=implementation, + policy=policy, + common=common, + parameters=normalized_parameters, + provenance=None, + ) + if not _values_equal(normalized, document): + raise ValueError("portable state is not in canonical semantic form") + return normalized + + +def _derived_role(shard: ShardIdentity) -> str: + if shard.layout is ParameterLayout.REPLICATED: + return "live" if shard.parameter.numel else "empty_replicated" + if shard.layout is ParameterLayout.FLATTENED_ELEMENT_SHARD: + return "live" if shard.logical_slice.length else "empty_flat" + if shard.layout is ParameterLayout.WHOLE_PARAMETER_OWNER: + return "whole_owner" if shard.local_member == shard.owner else "whole_nonowner" + raise ValueError("unsupported portable fragment layout") + + +def _local_dense_shape(shard: ShardIdentity): + role = _derived_role(shard) + if role not in {"live", "whole_owner"} or shard.parameter.numel == 0: + return None + if shard.layout in {ParameterLayout.REPLICATED, ParameterLayout.WHOLE_PARAMETER_OWNER}: + return shard.parameter.global_shape + if shard.layout is ParameterLayout.FLATTENED_ELEMENT_SHARD: + return (shard.logical_slice.length,) + raise ValueError("unsupported portable fragment layout") + + +def _normalize_catalog(value, implementation: str, manifest: ShardingManifest, policy): + if type(value) is not dict: + raise ValueError("portable fragment catalog must be an FQN mapping") + manifest_fqns = {shard.parameter.fqn for shard in manifest.shards} + if set(value) != manifest_fqns: + raise ValueError("portable fragment catalog does not match the manifest") + result = {} + for fqn in sorted(value): + record = _exact_record(value[fqn], _CATALOG_KEYS, name="portable catalog entry") + identity_record = _normalize_parameter_identity(record["identity"]) + identity = _parse_parameter_identity(identity_record) + if identity.fqn != fqn or any(shard.parameter != identity for shard in manifest.for_parameter(fqn)): + raise ValueError("portable catalog identity does not match the manifest") + options = _normalize_options(record["algorithm_options"], implementation) + if implementation == "gefen.Gefen": + expected = "factored" if policy["factored_v_2d"] and len(identity.global_shape) == 2 else "block" + if options["second_moment_policy"] != expected: + raise ValueError("catalog second_moment_policy conflicts with the optimizer policy") + if expected == "factored" and any( + shard.layout is not ParameterLayout.REPLICATED for shard in manifest.for_parameter(fqn) + ): + raise ValueError("factored logical matrices require replicated portable shards") + elif len(identity.global_shape) != 2: + raise ValueError("portable Muon catalog parameters must be logical matrices") + result[fqn] = { + "identity": identity_record, + "algorithm_options": options, + } + return result + + +def _normalize_fragment_slot(value, implementation: str, common, catalog, member: str): + value = _exact_record(value, _SLOT_KEYS, name="portable logical slot") + group_index = _strict_int(value["group_index"], name="group_index") + slot_index = _strict_int(value["original_slot_index"], name="original_slot_index") + compatibility_name = _strict_name(value["compatibility_name"], name="compatibility_name") + if compatibility_name != compatibility_name.lower(): + raise ValueError("compatibility_name must be lowercase") + shard_record = _normalize_shard_identity(value["shard"]) + shard = _parse_shard_identity(shard_record) + allowed_layouts = ( + {ParameterLayout.REPLICATED, ParameterLayout.FLATTENED_ELEMENT_SHARD} + if implementation == "gefen.Gefen" + else {ParameterLayout.REPLICATED, ParameterLayout.WHOLE_PARAMETER_OWNER} + ) + if shard.layout not in allowed_layouts: + raise ValueError("portable fragment layout is unsupported for its implementation") + if shard.process_group is None or shard.local_member != member: + raise ValueError("portable fragment member does not match its shard") + role = _derived_role(shard) + if value["role"] != role: + raise ValueError("portable fragment role does not match its shard") + options = _normalize_options(value["algorithm_options"], implementation) + catalog_entry = catalog.get(shard.parameter.fqn) + if catalog_entry is None or not _values_equal(options, catalog_entry["algorithm_options"]): + raise ValueError("portable slot options do not match the catalog") + if ( + implementation == "gefen.GefenMuon" + and shard.layout is ParameterLayout.WHOLE_PARAMETER_OWNER + and options["sharded_mode"] != "distributed" + ): + raise ValueError("whole-parameter Muon state requires sharded_mode='distributed'") + variant = value["state_variant"] + if variant not in { + "pristine", + "period_selected", + "initialized_dense", + "initialized_factored", + "initialized_dense_normuon", + }: + raise ValueError("portable fragment state_variant is unsupported") + period = value["source_period"] + if period is not None: + _strict_int(period, name="source_period", minimum=1) + if period != 1: + raise ValueError("exact portable v3 fragments require source_period=1") + local_shape = _local_dense_shape(shard) + if local_shape is None or math.prod(local_shape) % period != 0: + raise ValueError("source_period does not divide the local logical payload") + source_second = value["source_second_moment"] + if source_second not in {None, "block", "factored"}: + raise ValueError("portable source_second_moment is invalid") + state = value["state"] + if type(state) is not dict: + raise ValueError("portable fragment slot state must be a dictionary") + payload_shape = _local_dense_shape(shard) + payload_role = role in {"live", "whole_owner"} and shard.parameter.numel > 0 + if not payload_role and (variant != "pristine" or period is not None or state or source_second is not None): + raise ValueError("empty and whole-nonowner slots must remain pristine") + if variant == "pristine": + _exact_record(state, frozenset(), name="pristine fragment state") + if period is not None or source_second is not None: + raise ValueError("pristine fragment state has invalid metadata") + elif variant == "period_selected": + _exact_record(state, frozenset(), name="period-selected fragment state") + if period is None or source_second is not None: + raise ValueError("period-selected fragment state has invalid metadata") + elif implementation == "gefen.Gefen" and variant == "initialized_dense": + _exact_record( + state, frozenset({"step", "momentum", "second_moment", "second_moment_step"}), name="dense fragment state" + ) + if period is None or source_second != "block": + raise ValueError("initialized block fragment has invalid metadata") + if options["second_moment_policy"] != "block": + raise ValueError("initialized block fragment conflicts with its algorithm options") + elif implementation == "gefen.Gefen" and variant == "initialized_factored": + _exact_record( + state, frozenset({"step", "momentum", "v_row", "v_col", "factored_step"}), name="factored fragment state" + ) + if period is None or source_second != "factored" or shard.layout is not ParameterLayout.REPLICATED: + raise ValueError("factored fragments require replicated storage and factored metadata") + if options["second_moment_policy"] != "factored": + raise ValueError("initialized factored fragment conflicts with its algorithm options") + elif implementation == "gefen.GefenMuon" and variant in {"initialized_dense", "initialized_dense_normuon"}: + keys = {"step", "momentum"} + if variant == "initialized_dense_normuon": + keys.update({"normuon_v", "normuon_step"}) + _exact_record(state, frozenset(keys), name="Muon fragment state") + if period is None or source_second is not None: + raise ValueError("initialized Muon fragment has invalid metadata") + if (variant == "initialized_dense_normuon") != options["normuon"]: + raise ValueError("initialized Muon fragment conflicts with its NorMuon option") + else: + raise ValueError("fragment state_variant does not match its implementation") + normalized_state = dict(state) + if variant.startswith("initialized_"): + step = _strict_int(state["step"], name="step", minimum=1, maximum=_MAX_EXACT_COUNTER) + if step > common["gefen_global_step"]: + raise ValueError("fragment parameter step exceeds optimizer global step") + normalized_state["momentum"] = _tight_fp32(state["momentum"], name="local momentum", shape=payload_shape) + if variant == "initialized_dense" and implementation == "gefen.Gefen": + second_step = _strict_int( + state["second_moment_step"], name="second_moment_step", minimum=1, maximum=_MAX_EXACT_COUNTER + ) + normalized_state["second_moment"] = _tight_fp32( + state["second_moment"], name="local second_moment", shape=payload_shape, nonnegative=True + ) + if second_step > step: + raise ValueError("fragment second_moment_step exceeds step") + elif variant == "initialized_factored": + rows, columns = shard.parameter.global_shape + factored_step = _strict_int( + state["factored_step"], name="factored_step", minimum=1, maximum=_MAX_EXACT_COUNTER + ) + normalized_state["v_row"] = _tight_fp32(state["v_row"], name="v_row", shape=(rows,), nonnegative=True) + normalized_state["v_col"] = _tight_fp32(state["v_col"], name="v_col", shape=(columns,), nonnegative=True) + if factored_step > step: + raise ValueError("fragment factored_step exceeds step") + elif variant == "initialized_dense_normuon": + normuon_step = _strict_int( + state["normuon_step"], name="normuon_step", minimum=1, maximum=_MAX_EXACT_COUNTER + ) + normalized_state["normuon_v"] = _tight_fp32( + state["normuon_v"], name="normuon_v", shape=(shard.parameter.global_shape[0], 1), nonnegative=True + ) + if normuon_step > step: + raise ValueError("fragment normuon_step exceeds step") + return { + "group_index": group_index, + "original_slot_index": slot_index, + "compatibility_name": compatibility_name, + "shard": shard_record, + "algorithm_options": options, + "role": role, + "source_period": period, + "source_second_moment": source_second, + "state_variant": variant, + "state": normalized_state, + } + + +def _normalize_portable_state_fragment(fragment, *, limits): + """Bound, clone, and validate one reversible local portable-state fragment.""" + + limits = _require_limits(limits) + fragment = _bounded_clone(fragment, limits) + _exact_record(fragment, _FRAGMENT_KEYS, name="portable state fragment") + if fragment["format"] != _FRAGMENT_FORMAT or fragment["format_version"] != _FRAGMENT_FORMAT_VERSION: + raise ValueError("unsupported portable state fragment format") + if type(fragment["format_version"]) is not int or fragment["coverage"] != _FRAGMENT_COVERAGE: + raise ValueError("unsupported portable state fragment coverage") + implementation = fragment["implementation"] + if implementation not in _IMPLEMENTATIONS: + raise ValueError("unsupported portable fragment implementation") + member = _strict_name(fragment["member"], name="portable fragment member") + policy = _normalize_policy(fragment["policy"], implementation) + common = _normalize_common(fragment["common"]) + manifest_record = _normalize_sharding_manifest(fragment["manifest"]) + manifest = _parse_sharding_manifest(manifest_record) + allowed_layouts = ( + {ParameterLayout.REPLICATED, ParameterLayout.FLATTENED_ELEMENT_SHARD} + if implementation == "gefen.Gefen" + else {ParameterLayout.REPLICATED, ParameterLayout.WHOLE_PARAMETER_OWNER} + ) + if any(shard.layout not in allowed_layouts or shard.process_group is None for shard in manifest.shards): + raise ValueError("portable fragment manifest contains an unsupported or ungrouped shard") + catalog = _normalize_catalog(fragment["catalog"], implementation, manifest, policy) + slots_value = fragment["logical_slots"] + if type(slots_value) is not list or not slots_value: + raise ValueError("portable fragment logical_slots must be a non-empty list") + slots = [_normalize_fragment_slot(slot, implementation, common, catalog, member) for slot in slots_value] + if common["gefen_codebook"] is None and any(slot["state_variant"] != "pristine" for slot in slots): + raise ValueError("non-pristine portable fragment state requires a codebook") + positions = [(slot["group_index"], slot["original_slot_index"]) for slot in slots] + if positions[0] != (0, 0) or positions != sorted(positions) or len(set(positions)) != len(positions): + raise ValueError("portable logical slot positions must start at (0, 0) and be strictly ordered") + expected_group = 0 + expected_slot = 0 + for group_index, slot_index in positions: + if group_index == expected_group and slot_index == expected_slot: + expected_slot += 1 + elif group_index == expected_group + 1 and slot_index == 0: + expected_group += 1 + expected_slot = 1 + else: + raise ValueError("portable logical slot positions must be contiguous") + identities = [_parse_shard_identity(slot["shard"]) for slot in slots] + if len(set(identities)) != len(identities): + raise ValueError("portable fragment contains duplicate shard identities") + if any(shard not in manifest.shards for shard in identities): + raise ValueError("portable fragment slot is absent from the manifest") + return { + "format": _FRAGMENT_FORMAT, + "format_version": _FRAGMENT_FORMAT_VERSION, + "coverage": _FRAGMENT_COVERAGE, + "implementation": implementation, + "member": member, + "policy": policy, + "common": common, + "manifest": manifest_record, + "catalog": catalog, + "logical_slots": slots, + } + + +def _build_portable_state_fragment(*, implementation, member, policy, common, manifest, catalog, logical_slots, limits): + """Build one strict local fragment from serialized primitives and local dense fields.""" + + if isinstance(manifest, ShardingManifest): + manifest = _serialize_sharding_manifest(manifest) + fragment = { + "format": _FRAGMENT_FORMAT, + "format_version": _FRAGMENT_FORMAT_VERSION, + "coverage": _FRAGMENT_COVERAGE, + "implementation": implementation, + "member": member, + "policy": policy, + "common": common, + "manifest": manifest, + "catalog": catalog, + "logical_slots": logical_slots, + } + return _normalize_portable_state_fragment(fragment, limits=limits) + + +def _required_payload_slots(slots): + return [ + slot + for slot in slots + if slot["role"] in {"live", "whole_owner"} and _parse_shard_identity(slot["shard"]).parameter.numel > 0 + ] + + +def _consensus(values, *, name: str): + if not values: + raise ValueError("{} requires at least one value".format(name)) + reference = values[0] + if any(not _values_equal(reference, value) for value in values[1:]): + raise ValueError("{} disagree across portable fragments".format(name)) + return reference + + +def _assemble_special_field(slots, *, key: str, shape, name: str) -> torch.Tensor: + values = [] + for slot in slots: + shard = _parse_shard_identity(slot["shard"]) + payload = slot["state"].get(key) + if shard.layout is ParameterLayout.REPLICATED: + values.append(_tight_fp32(payload, name=name, shape=shape, nonnegative=True)) + elif shard.layout is ParameterLayout.WHOLE_PARAMETER_OWNER: + if shard.local_member == shard.owner: + values.append(_tight_fp32(payload, name=name, shape=shape, nonnegative=True)) + elif payload is not None: + raise ValueError("whole-parameter nonowner carries {}".format(name)) + else: + raise ValueError("{} supports only exact replicas or a sole whole owner".format(name)) + return _tight_clone(_consensus(values, name=name)) + + +def _assemble_parameter(fqn, slots, manifest: ShardingManifest, catalog_entry, implementation, policy, global_step): + parameter = _parse_parameter_identity(catalog_entry["identity"]) + required = _required_payload_slots(slots) + variants = {slot["state_variant"] for slot in required} + if not required: + variant = "pristine" + elif variants == {"pristine"}: + variant = "pristine" + elif variants == {"period_selected"}: + variant = "period_selected" + elif len(variants) == 1 and next(iter(variants)).startswith("initialized_"): + variant = next(iter(variants)) + else: + raise ValueError("portable parameter shards disagree on initialization state") + if any(slot["state_variant"] != "pristine" for slot in slots if slot not in required): + raise ValueError("non-payload portable slots must remain pristine") + periods = sorted({slot["source_period"] for slot in required if slot["source_period"] is not None}) + if variant == "pristine" and periods: + raise ValueError("pristine portable parameter has source periods") + if variant != "pristine" and not periods: + raise ValueError("non-pristine portable parameter requires source periods") + state = {} + source_second = None + if variant.startswith("initialized_"): + step = _consensus([slot["state"]["step"] for slot in required], name="parameter step") + if step > global_step: + raise ValueError("parameter step exceeds optimizer global step") + momentum = _assemble_dense_logical_field( + manifest, + parameter, + [(_parse_shard_identity(slot["shard"]), slot["state"].get("momentum")) for slot in slots], + ) + state = {"step": step, "momentum": momentum} + if implementation == "gefen.Gefen" and variant == "initialized_dense": + source_second = "block" + second_step = _consensus( + [slot["state"]["second_moment_step"] for slot in required], name="second_moment_step" + ) + state.update( + { + "second_moment": _assemble_dense_logical_field( + manifest, + parameter, + [(_parse_shard_identity(slot["shard"]), slot["state"].get("second_moment")) for slot in slots], + ), + "second_moment_step": second_step, + } + ) + elif implementation == "gefen.Gefen" and variant == "initialized_factored": + source_second = "factored" + if any(_parse_shard_identity(slot["shard"]).layout is not ParameterLayout.REPLICATED for slot in slots): + raise ValueError("factored portable state must be assembled from exact replicas") + factored_step = _consensus([slot["state"]["factored_step"] for slot in required], name="factored_step") + state.update( + { + "v_row": _assemble_special_field( + slots, key="v_row", shape=(parameter.global_shape[0],), name="v_row" + ), + "v_col": _assemble_special_field( + slots, key="v_col", shape=(parameter.global_shape[1],), name="v_col" + ), + "factored_step": factored_step, + } + ) + elif implementation == "gefen.GefenMuon" and variant == "initialized_dense_normuon": + state.update( + { + "normuon_v": _assemble_special_field( + slots, + key="normuon_v", + shape=(parameter.global_shape[0], 1), + name="normuon_v", + ), + "normuon_step": _consensus( + [slot["state"]["normuon_step"] for slot in required], name="normuon_step" + ), + } + ) + if implementation == "gefen.Gefen": + expected = "factored" if policy["factored_v_2d"] and len(parameter.global_shape) == 2 else "block" + if variant == "initialized_dense" and expected != "block": + raise ValueError("block portable state conflicts with factored_v_2d policy") + if variant == "initialized_factored" and expected != "factored": + raise ValueError("factored portable state conflicts with factored_v_2d policy") + hints = {"source_periods": periods, "source_second_moment": source_second, "target_period": 1} + else: + hints = {"source_periods": periods, "target_period": 1} + return { + "identity": _serialize_parameter_identity(parameter), + "algorithm_options": catalog_entry["algorithm_options"], + "state_variant": variant, + "state": state, + "projection_hints": hints, + } + + +def _assemble_portable_state_fragments(fragments, *, process_group_identity, limits): + """Assemble member-ordered fragments into one complete topology-neutral document.""" + + limits = _require_limits(limits) + if not isinstance(process_group_identity, ProcessGroupIdentity): + raise TypeError("process_group_identity must be a ProcessGroupIdentity") + process_group_identity = ProcessGroupIdentity( + process_group_identity.semantic_name, + process_group_identity.ordered_members, + schema_version=process_group_identity.schema_version, + ) + members = tuple(_strict_name(member, name="ordered member") for member in process_group_identity.ordered_members) + if len(set(members)) != len(members) or len(members) > limits.max_members: + raise ValueError("process-group members must be unique and within max_members") + if isinstance(fragments, (str, bytes, bytearray)): + raise TypeError("fragments must be a sequence") + try: + iterator = iter(fragments) + except TypeError as exc: + raise TypeError("fragments must be a sequence") from exc + bounded_fragments = [] + for fragment in iterator: + bounded_fragments.append(fragment) + if len(bounded_fragments) > limits.max_members: + raise ValueError("portable fragments exceed max_members") + fragments = tuple(bounded_fragments) + if len(fragments) != len(members): + raise ValueError("portable fragments must contain every process-group member exactly once") + + normalized = [] + total_tensor_bytes = 0 + total_metadata_bytes = 0 + for fragment in fragments: + normalized_fragment = _normalize_portable_state_fragment(fragment, limits=limits) + plan = _prepare_canonical_wire_value(normalized_fragment, limits._wire_limits()) + if total_tensor_bytes > limits.max_collective_tensor_bytes - plan.total_tensor_bytes: + raise ValueError("portable fragments exceed max_collective_tensor_bytes") + if total_metadata_bytes > limits.max_collective_metadata_bytes - len(plan.metadata): + raise ValueError("portable fragments exceed max_collective_metadata_bytes") + total_tensor_bytes += plan.total_tensor_bytes + total_metadata_bytes += len(plan.metadata) + normalized.append(normalized_fragment) + if tuple(fragment["member"] for fragment in normalized) != members: + raise ValueError("portable fragments are not in process-group member order") + + implementation = _consensus([fragment["implementation"] for fragment in normalized], name="implementation") + policy = _consensus([fragment["policy"] for fragment in normalized], name="policy") + common = _consensus([fragment["common"] for fragment in normalized], name="common") + manifest_record = _consensus([fragment["manifest"] for fragment in normalized], name="manifest") + catalog = _consensus([fragment["catalog"] for fragment in normalized], name="catalog") + manifest = _parse_sharding_manifest(manifest_record) + if any(shard.process_group != process_group_identity for shard in manifest.shards): + raise ValueError("portable manifest process groups must exactly match process_group_identity") + all_slots = [slot for fragment in normalized for slot in fragment["logical_slots"]] + shard_slots = {_parse_shard_identity(slot["shard"]): slot for slot in all_slots} + if len(shard_slots) != len(all_slots) or set(shard_slots) != set(manifest.shards): + raise ValueError("portable fragments must exactly cover the complete manifest") + + parameters = {} + parameter_positions = {} + for fqn in sorted(catalog): + slots = [shard_slots[shard] for shard in manifest.for_parameter(fqn)] + parameter_positions[fqn] = _consensus( + [(slot["group_index"], slot["original_slot_index"]) for slot in slots], + name="logical slot position", + ) + _consensus([slot["compatibility_name"] for slot in slots], name="compatibility_name") + parameters[fqn] = _assemble_parameter( + fqn, + slots, + manifest, + catalog[fqn], + implementation, + policy, + common["gefen_global_step"], + ) + if len(set(parameter_positions.values())) != len(parameter_positions): + raise ValueError("portable parameters must have unique logical slot positions") + document = build_portable_state_document( + implementation=implementation, + policy=policy, + common=common, + parameters=parameters, + provenance=None, + ) + return _normalize_gefen_portable_state_document( + document, + limits=limits, + expected_implementation=implementation, + ) + + +def _project_portable_parameter_state( + document_record, + target_shard, + *, + implementation, + global_step, + codebook, + target_algorithm_options, + target_second_moment=None, +): + """Project one normalized global record to a target shard's native state fields.""" + + if implementation not in _IMPLEMENTATIONS: + raise ValueError("unsupported portable projection implementation") + global_step = _strict_int(global_step, name="gefen_global_step", maximum=_MAX_EXACT_COUNTER) + common_codebook = _normalize_codebook(codebook, global_step=global_step) + if not isinstance(target_shard, ShardIdentity): + raise TypeError("target_shard must be a ShardIdentity") + target_layouts = ( + {ParameterLayout.REPLICATED, ParameterLayout.FLATTENED_ELEMENT_SHARD} + if implementation == "gefen.Gefen" + else {ParameterLayout.REPLICATED, ParameterLayout.WHOLE_PARAMETER_OWNER} + ) + if target_shard.layout not in target_layouts: + raise ValueError("target shard layout is unsupported for the portable implementation") + if target_shard.process_group is None: + raise ValueError("portable projection requires a grouped target shard") + source_factored = ( + implementation == "gefen.Gefen" + and type(document_record) is dict + and type(document_record.get("algorithm_options")) is dict + and document_record["algorithm_options"].get("second_moment_policy") == "factored" + ) + policy = { + "schema_version": 1, + "factored_v_2d": source_factored, + "force_1d_period_one": False, + "force_2d_period_one": False, + "period_one_substrings": [], + "codebook_refresh_every": 0, + "stochastic_round": False, + "momentum_projection": _MOMENTUM_PROJECTION, + "second_moment_projection": _SECOND_MOMENT_PROJECTION, + } + fqn = target_shard.parameter.fqn + record = _normalize_complete_parameter_record(fqn, document_record, implementation, global_step, policy) + if _parse_parameter_identity(record["identity"]) != target_shard.parameter: + raise ValueError("portable parameter identity does not match the target shard") + target_options = _normalize_options(target_algorithm_options, implementation) + if not _values_equal(record["algorithm_options"], target_options): + raise ValueError("portable algorithm options do not match the target") + if implementation == "gefen.Gefen": + if target_second_moment not in {"block", "factored"}: + raise ValueError("plain Gefen projection requires target_second_moment") + if target_second_moment != target_options["second_moment_policy"]: + raise ValueError("target_second_moment conflicts with target algorithm options") + if ( + target_second_moment == "factored" + and len(target_shard.parameter.global_shape) == 2 + and target_shard.layout is not ParameterLayout.REPLICATED + ): + raise ValueError("factored logical matrices require replicated target shards") + source_second = record["projection_hints"]["source_second_moment"] + if source_second is not None and source_second != target_second_moment: + raise ValueError("portable second-moment representation migration is unsupported") + else: + if target_second_moment is not None: + raise ValueError("Muon projection does not accept target_second_moment") + if ( + target_shard.layout is ParameterLayout.WHOLE_PARAMETER_OWNER + and target_options["sharded_mode"] != "distributed" + ): + raise ValueError("whole-parameter Muon targets require sharded_mode='distributed'") + + variant = record["state_variant"] + if variant == "pristine": + return {} + if common_codebook is None: + raise ValueError("non-pristine portable projection requires a codebook") + if variant == "period_selected": + return {"automatic_period": 1} if _local_dense_shape(target_shard) is not None else {} + local_momentum = _project_dense_logical_field( + target_shard.parameter, + record["state"]["momentum"], + target_shard, + ) + if local_momentum is None: + return {} + step = record["state"]["step"] + indices, magnitudes = _recompress_dense_momentum(local_momentum, common_codebook, period=1, step=step) + result = { + "automatic_period": 1, + "step": step, + "m_codebook": indices, + "m_magnitude": magnitudes, + } + if implementation == "gefen.Gefen" and variant == "initialized_dense": + local_second = _project_dense_logical_field( + target_shard.parameter, + record["state"]["second_moment"], + target_shard, + ) + if local_second is None: + return {} + second_step = record["state"]["second_moment_step"] + result.update( + { + "vmean": _reduce_block_second_moment(local_second, period=1, step=second_step), + "vmean_step": second_step, + } + ) + elif implementation == "gefen.Gefen" and variant == "initialized_factored": + if target_shard.layout is not ParameterLayout.REPLICATED: + raise ValueError("factored portable state projects only to replicated targets") + result.update( + { + "v_row": _tight_clone(record["state"]["v_row"]), + "v_col": _tight_clone(record["state"]["v_col"]), + "factored_step": record["state"]["factored_step"], + } + ) + elif implementation == "gefen.GefenMuon" and variant == "initialized_dense_normuon": + if target_shard.layout not in {ParameterLayout.REPLICATED, ParameterLayout.WHOLE_PARAMETER_OWNER}: + raise ValueError("NorMuon state projects only to replicas or a whole-parameter owner") + result.update( + { + "normuon_v": _tight_clone(record["state"]["normuon_v"]), + "normuon_step": record["state"]["normuon_step"], + } + ) + return result + + +__all__ = ["PortableStateLimits"] diff --git a/tests/test_portable_state.py b/tests/test_portable_state.py new file mode 100644 index 0000000..1894cd5 --- /dev/null +++ b/tests/test_portable_state.py @@ -0,0 +1,874 @@ +"""Warning-strict CPU coverage for portable Gefen optimizer-state semantics.""" + +import copy +import dataclasses +import math + +import pytest +import torch + +from gefen.contracts import ( + LogicalSlice, + ParameterIdentity, + ParameterLayout, + PlacementKind, + ProcessGroupIdentity, + ShardIdentity, + ShardPlacement, + ShardingManifest, +) +from gefen.portable import _decode_quantized_momentum +from gefen.portable_fields import _project_dense_logical_field +from gefen.portable_identity import ( + _serialize_parameter_identity, + _serialize_shard_identity, +) +from gefen.portable_schema import build_portable_state_document +from gefen.portable_state import ( + PortableStateLimits, + _assemble_portable_state_fragments, + _build_portable_state_fragment, + _normalize_gefen_portable_state_document, + _normalize_portable_state_fragment, + _project_portable_parameter_state, +) + + +def _limits(**changes): + limits = PortableStateLimits( + max_fragment_tensor_bytes=4 << 20, + max_collective_tensor_bytes=16 << 20, + max_collective_metadata_bytes=64 << 20, + ) + return dataclasses.replace(limits, **changes) + + +def _codebook(): + return torch.linspace(-1.0, 1.0, 256, dtype=torch.float32) + + +def _policy(*, factored=False): + return { + "schema_version": 1, + "factored_v_2d": factored, + "force_1d_period_one": False, + "force_2d_period_one": False, + "period_one_substrings": [], + "stochastic_round": False, + "codebook_refresh_every": 0, + "momentum_projection": "dense_fp32_target_period_one_v1", + "second_moment_projection": "exact_representation_target_period_one_v1", + } + + +def _plain_options(*, second="block"): + return { + "lr": 0.001, + "beta1": 0.9, + "beta2": 0.999, + "eps": 1e-8, + "weight_decay": 0.01, + "second_moment_policy": second, + } + + +def _muon_options(*, normuon=False, mode="distributed"): + return { + "lr": 0.01, + "weight_decay": 0.1, + "momentum": 0.95, + "nesterov": True, + "ns_schedule": [[3.4445, -4.775, 2.0315]], + "ns_eps": 1e-7, + "adjust_lr_fn": "match_rms_adamw", + "sharded_mode": mode, + "fp8_ns": False, + "fp8_ns_compile": True, + "batched_ns": False, + "batched_ns_workspace_bytes": 1 << 20, + "normuon": normuon, + "normuon_beta2": 0.95, + "normuon_eps": 1e-8, + "cautious": False, + } + + +def _group(members=("rank:0", "rank:1")): + return ProcessGroupIdentity("checkpoint", members) + + +def _placement(group, member, kind): + return ShardPlacement( + "checkpoint", + kind, + group.ordered_members.index(member), + len(group.ordered_members), + ) + + +def _replicas(parameter, group): + return tuple( + ShardIdentity( + parameter, + ParameterLayout.REPLICATED, + LogicalSlice.full(parameter), + placements=(_placement(group, member, PlacementKind.REPLICATE),), + process_group=group, + local_member=member, + ) + for member in group.ordered_members + ) + + +def _flat(parameter, group, lengths): + offset = 0 + shards = [] + for member, length in zip(group.ordered_members, lengths): + shards.append( + ShardIdentity( + parameter, + ParameterLayout.FLATTENED_ELEMENT_SHARD, + LogicalSlice(offset, length), + placements=(_placement(group, member, PlacementKind.FLAT_SHARD),), + process_group=group, + local_member=member, + ) + ) + offset += length + return tuple(shards) + + +def _owners(parameter, group, owner): + return tuple( + ShardIdentity( + parameter, + ParameterLayout.WHOLE_PARAMETER_OWNER, + LogicalSlice.full(parameter) if member == owner else LogicalSlice(0, 0), + placements=(_placement(group, member, PlacementKind.WHOLE_PARAMETER_OWNER),), + process_group=group, + local_member=member, + owner=owner, + ) + for member in group.ordered_members + ) + + +def _role(shard): + if shard.layout is ParameterLayout.WHOLE_PARAMETER_OWNER: + return "whole_owner" if shard.local_member == shard.owner else "whole_nonowner" + if shard.layout is ParameterLayout.FLATTENED_ELEMENT_SHARD and shard.logical_slice.length == 0: + return "empty_flat" + if shard.parameter.numel == 0: + return "empty_replicated" + return "live" + + +def _fragments( + implementation, + parameter, + shards, + *, + options, + policy, + variant, + momentum=None, + second=None, + v_row=None, + v_col=None, + normuon_v=None, +): + manifest = ShardingManifest(shards) + common = { + "gefen_global_step": 4 if variant != "pristine" else 0, + "gefen_codebook": _codebook() if variant != "pristine" else None, + "gefen_deterministic": True, + } + catalog = { + parameter.fqn: { + "identity": _serialize_parameter_identity(parameter), + "algorithm_options": options, + } + } + result = [] + for shard in shards: + payload = _role(shard) in {"live", "whole_owner"} and parameter.numel > 0 + local_variant = variant if payload else "pristine" + state = {} + source_second = None + period = None + if local_variant == "period_selected": + period = 1 + elif local_variant.startswith("initialized_"): + period = 1 + state = { + "step": 4, + "momentum": _project_dense_logical_field(parameter, momentum, shard), + } + if implementation == "gefen.Gefen" and local_variant == "initialized_dense": + source_second = "block" + state.update( + { + "second_moment": _project_dense_logical_field(parameter, second, shard), + "second_moment_step": 3, + } + ) + elif local_variant == "initialized_factored": + source_second = "factored" + state.update({"v_row": v_row.clone(), "v_col": v_col.clone(), "factored_step": 3}) + elif local_variant == "initialized_dense_normuon": + state.update({"normuon_v": normuon_v.clone(), "normuon_step": 2}) + slot = { + "group_index": 0, + "original_slot_index": 0, + "compatibility_name": "weight", + "shard": _serialize_shard_identity(shard), + "algorithm_options": options, + "role": _role(shard), + "source_period": period, + "source_second_moment": source_second, + "state_variant": local_variant, + "state": state, + } + result.append( + _build_portable_state_fragment( + implementation=implementation, + member=shard.local_member, + policy=policy, + common=common, + manifest=manifest, + catalog=catalog, + logical_slots=[slot], + limits=_limits(), + ) + ) + return result + + +@pytest.mark.parametrize("variant", ["pristine", "period_selected"]) +def test_plain_pristine_and_period_selected_round_trip(variant): + parameter = ParameterIdentity("layer.weight", (2, 3)) + group = _group() + fragments = _fragments( + "gefen.Gefen", + parameter, + _replicas(parameter, group), + options=_plain_options(), + policy=_policy(), + variant=variant, + ) + document = _assemble_portable_state_fragments(fragments, process_group_identity=group, limits=_limits()) + record = document["parameters"][parameter.fqn] + assert record["state_variant"] == variant + assert record["state"] == {} + assert record["projection_hints"]["source_periods"] == ([] if variant == "pristine" else [1]) + normalized = _normalize_gefen_portable_state_document(document, limits=_limits()) + assert normalized["completion"] == document["completion"] + + +@pytest.mark.parametrize("layout", ["replicated", "flat"]) +def test_plain_initialized_block_assembles_and_projects_exact_period_one(layout): + parameter = ParameterIdentity("layer.weight", (2, 3)) + group = _group() + shards = _replicas(parameter, group) if layout == "replicated" else _flat(parameter, group, (2, 4)) + momentum = torch.tensor([[-0.0, 0.0, float.fromhex("0x1p-149")], [1.25, -3.5, torch.finfo(torch.float32).max]]) + second = torch.arange(1, 7, dtype=torch.float32).reshape(2, 3) + fragments = _fragments( + "gefen.Gefen", + parameter, + shards, + options=_plain_options(), + policy=_policy(), + variant="initialized_dense", + momentum=momentum, + second=second, + ) + document = _assemble_portable_state_fragments(fragments, process_group_identity=group, limits=_limits()) + record = document["parameters"][parameter.fqn] + assert torch.equal(record["state"]["second_moment"], second) + assert torch.equal(record["state"]["momentum"].view(torch.int32), momentum.view(torch.int32)) + + for target in _flat(parameter, group, (3, 3)): + projected = _project_portable_parameter_state( + record, + target, + implementation="gefen.Gefen", + global_step=4, + codebook=document["common"]["gefen_codebook"], + target_algorithm_options=_plain_options(), + target_second_moment="block", + ) + local_momentum = _decode_quantized_momentum( + document["common"]["gefen_codebook"], + projected["m_codebook"], + projected["m_magnitude"], + logical_shape=(target.logical_slice.length,), + period=1, + step=4, + ) + expected = _project_dense_logical_field(parameter, momentum, target) + assert torch.equal(local_momentum.view(torch.int32), expected.view(torch.int32)) + assert projected["automatic_period"] == 1 + assert torch.equal(projected["vmean"].reshape(-1), _project_dense_logical_field(parameter, second, target)) + + +def test_empty_flat_fragment_is_pristine_and_ownership_is_enforced(): + parameter = ParameterIdentity("layer.weight", (2,)) + group = _group(("rank:0", "rank:1", "rank:2")) + shards = _flat(parameter, group, (1, 0, 1)) + fragments = _fragments( + "gefen.Gefen", + parameter, + shards, + options=_plain_options(), + policy=_policy(), + variant="initialized_dense", + momentum=torch.tensor([-0.0, 2.0]), + second=torch.tensor([1.0, 2.0]), + ) + assert fragments[1]["logical_slots"][0]["role"] == "empty_flat" + assert fragments[1]["logical_slots"][0]["state_variant"] == "pristine" + corrupted = copy.deepcopy(fragments[1]) + corrupted["logical_slots"][0]["role"] = "live" + with pytest.raises(ValueError, match="role"): + _normalize_portable_state_fragment(corrupted, limits=_limits()) + + +def test_factored_replicas_preserve_authoritative_factor_bytes_and_reject_migration(): + parameter = ParameterIdentity("matrix.weight", (2, 3)) + group = _group() + shards = _replicas(parameter, group) + momentum = torch.arange(6, dtype=torch.float32).reshape(2, 3) + v_row = torch.tensor([-0.0, 4.0]) + v_col = torch.tensor([1.0, 2.0, 3.0]) + fragments = _fragments( + "gefen.Gefen", + parameter, + shards, + options=_plain_options(second="factored"), + policy=_policy(factored=True), + variant="initialized_factored", + momentum=momentum, + v_row=v_row, + v_col=v_col, + ) + document = _assemble_portable_state_fragments(fragments, process_group_identity=group, limits=_limits()) + record = document["parameters"][parameter.fqn] + assert torch.equal(record["state"]["v_row"].view(torch.int32), v_row.view(torch.int32)) + projected = _project_portable_parameter_state( + record, + shards[0], + implementation="gefen.Gefen", + global_step=4, + codebook=_codebook(), + target_algorithm_options=_plain_options(second="factored"), + target_second_moment="factored", + ) + assert torch.equal(projected["v_row"].view(torch.int32), v_row.view(torch.int32)) + with pytest.raises(ValueError, match="target"): + _project_portable_parameter_state( + record, + shards[0], + implementation="gefen.Gefen", + global_step=4, + codebook=_codebook(), + target_algorithm_options=_plain_options(second="factored"), + target_second_moment="block", + ) + fragments[1]["logical_slots"][0]["state"]["v_row"][0] = 0.0 + with pytest.raises(ValueError, match="v_row disagree"): + _assemble_portable_state_fragments(fragments, process_group_identity=group, limits=_limits()) + + +def test_period_selected_projection_binds_target_second_moment_policy(): + parameter = ParameterIdentity("matrix.weight", (2, 3)) + group = _group() + shards = _replicas(parameter, group) + document = _assemble_portable_state_fragments( + _fragments( + "gefen.Gefen", + parameter, + shards, + options=_plain_options(second="factored"), + policy=_policy(factored=True), + variant="period_selected", + ), + process_group_identity=group, + limits=_limits(), + ) + with pytest.raises(ValueError, match="conflicts with target algorithm options"): + _project_portable_parameter_state( + document["parameters"][parameter.fqn], + shards[0], + implementation="gefen.Gefen", + global_step=0, + codebook=document["common"]["gefen_codebook"], + target_algorithm_options=_plain_options(second="factored"), + target_second_moment="block", + ) + + +def test_factored_matrix_rejects_flat_fragments_before_initialization(): + parameter = ParameterIdentity("matrix.weight", (2, 3)) + group = _group() + with pytest.raises(ValueError, match="factored logical matrices require replicated"): + _fragments( + "gefen.Gefen", + parameter, + _flat(parameter, group, (2, 4)), + options=_plain_options(second="factored"), + policy=_policy(factored=True), + variant="period_selected", + ) + document = _assemble_portable_state_fragments( + _fragments( + "gefen.Gefen", + parameter, + _replicas(parameter, group), + options=_plain_options(second="factored"), + policy=_policy(factored=True), + variant="period_selected", + ), + process_group_identity=group, + limits=_limits(), + ) + with pytest.raises(ValueError, match="replicated target shards"): + _project_portable_parameter_state( + document["parameters"][parameter.fqn], + _flat(parameter, group, (2, 4))[0], + implementation="gefen.Gefen", + global_step=document["common"]["gefen_global_step"], + codebook=document["common"]["gefen_codebook"], + target_algorithm_options=_plain_options(second="factored"), + target_second_moment="factored", + ) + + +@pytest.mark.parametrize("normuon", [False, True]) +def test_muon_whole_owner_assembles_and_projects_after_owner_move(normuon): + parameter = ParameterIdentity("muon.weight", (2, 3)) + group = _group() + source = _owners(parameter, group, "rank:0") + momentum = torch.tensor([[-1.0, 0.0, 1.0], [2.0, 3.0, 4.0]]) + normuon_v = torch.tensor([[1.0], [2.0]]) + variant = "initialized_dense_normuon" if normuon else "initialized_dense" + fragments = _fragments( + "gefen.GefenMuon", + parameter, + source, + options=_muon_options(normuon=normuon), + policy=_policy(), + variant=variant, + momentum=momentum, + normuon_v=normuon_v, + ) + document = _assemble_portable_state_fragments(fragments, process_group_identity=group, limits=_limits()) + target = _owners(parameter, group, "rank:1") + nonowner = _project_portable_parameter_state( + document["parameters"][parameter.fqn], + target[0], + implementation="gefen.GefenMuon", + global_step=4, + codebook=_codebook(), + target_algorithm_options=_muon_options(normuon=normuon), + ) + owner = _project_portable_parameter_state( + document["parameters"][parameter.fqn], + target[1], + implementation="gefen.GefenMuon", + global_step=4, + codebook=_codebook(), + target_algorithm_options=_muon_options(normuon=normuon), + ) + assert nonowner == {} + assert owner["automatic_period"] == 1 + if normuon: + assert torch.equal(owner["normuon_v"], normuon_v) + + +def test_muon_whole_owner_projection_requires_distributed_mode(): + parameter = ParameterIdentity("muon.weight", (2, 3)) + group = _group() + document = _assemble_portable_state_fragments( + _fragments( + "gefen.GefenMuon", + parameter, + _replicas(parameter, group), + options=_muon_options(mode="exact"), + policy=_policy(), + variant="period_selected", + ), + process_group_identity=group, + limits=_limits(), + ) + target = _owners(parameter, group, "rank:0")[0] + with pytest.raises(ValueError, match="sharded_mode='distributed'"): + _project_portable_parameter_state( + document["parameters"][parameter.fqn], + target, + implementation="gefen.GefenMuon", + global_step=document["common"]["gefen_global_step"], + codebook=document["common"]["gefen_codebook"], + target_algorithm_options=_muon_options(mode="exact"), + ) + + +def test_strict_document_schema_completion_digest_counters_and_period_gate(): + parameter = ParameterIdentity("layer.weight", (2,)) + record = { + "identity": _serialize_parameter_identity(parameter), + "algorithm_options": _plain_options(), + "state_variant": "period_selected", + "state": {}, + "projection_hints": { + "source_periods": [1], + "source_second_moment": None, + "target_period": 1, + }, + } + document = build_portable_state_document( + implementation="gefen.Gefen", + policy=_policy(), + common={"gefen_global_step": 1, "gefen_codebook": _codebook(), "gefen_deterministic": False}, + parameters={parameter.fqn: record}, + provenance=None, + ) + _normalize_gefen_portable_state_document(document, limits=_limits()) + corrupt = copy.deepcopy(document) + corrupt["completion"]["digest"] = "0" * 64 + with pytest.raises(ValueError, match="digest"): + _normalize_gefen_portable_state_document(corrupt, limits=_limits()) + bad_period = copy.deepcopy(record) + bad_period["projection_hints"]["source_periods"] = [2] + with pytest.raises(ValueError, match="period one"): + _normalize_gefen_portable_state_document( + build_portable_state_document( + implementation="gefen.Gefen", + policy=_policy(), + common=document["common"], + parameters={parameter.fqn: bad_period}, + provenance=None, + ), + limits=_limits(), + ) + too_large = copy.deepcopy(record) + too_large["state_variant"] = "pristine" + too_large["projection_hints"]["source_periods"] = [] + with pytest.raises(ValueError, match="at most"): + _normalize_gefen_portable_state_document( + build_portable_state_document( + implementation="gefen.Gefen", + policy=_policy(), + common={ + "gefen_global_step": (1 << 53), + "gefen_codebook": None, + "gefen_deterministic": False, + }, + parameters={parameter.fqn: too_large}, + provenance=None, + ), + limits=_limits(), + ) + + +def test_global_step_zero_allows_period_selection_only_with_a_valid_codebook(): + parameter = ParameterIdentity("layer.weight", (2,)) + record = { + "identity": _serialize_parameter_identity(parameter), + "algorithm_options": _plain_options(), + "state_variant": "period_selected", + "state": {}, + "projection_hints": { + "source_periods": [1], + "source_second_moment": None, + "target_period": 1, + }, + } + + def document(codebook): + return build_portable_state_document( + implementation="gefen.Gefen", + policy=_policy(), + common={ + "gefen_global_step": 0, + "gefen_codebook": codebook, + "gefen_deterministic": False, + }, + parameters={parameter.fqn: record}, + provenance=None, + ) + + normalized = _normalize_gefen_portable_state_document(document(_codebook()), limits=_limits()) + assert normalized["common"]["gefen_global_step"] == 0 + with pytest.raises(ValueError, match="requires a codebook"): + _normalize_gefen_portable_state_document(document(None), limits=_limits()) + + +def test_fragment_rejects_ungrouped_shards_and_empty_period_projection_stays_empty(): + parameter = ParameterIdentity("layer.weight", (2,)) + ungrouped = ShardIdentity( + parameter, + ParameterLayout.REPLICATED, + LogicalSlice.full(parameter), + ) + catalog = { + parameter.fqn: { + "identity": _serialize_parameter_identity(parameter), + "algorithm_options": _plain_options(), + } + } + with pytest.raises(ValueError, match="ungrouped"): + _build_portable_state_fragment( + implementation="gefen.Gefen", + member="rank:0", + policy=_policy(), + common={ + "gefen_global_step": 0, + "gefen_codebook": None, + "gefen_deterministic": False, + }, + manifest=ShardingManifest((ungrouped,)), + catalog=catalog, + logical_slots=[ + { + "group_index": 0, + "original_slot_index": 0, + "compatibility_name": "weight", + "shard": _serialize_shard_identity(ungrouped), + "algorithm_options": _plain_options(), + "role": "live", + "source_period": None, + "source_second_moment": None, + "state_variant": "pristine", + "state": {}, + } + ], + limits=_limits(), + ) + + group = _group(("rank:0", "rank:1", "rank:2")) + source = _replicas(parameter, group) + document = _assemble_portable_state_fragments( + _fragments( + "gefen.Gefen", + parameter, + source, + options=_plain_options(), + policy=_policy(), + variant="period_selected", + ), + process_group_identity=group, + limits=_limits(), + ) + empty_target = _flat(parameter, group, (1, 0, 1))[1] + assert ( + _project_portable_parameter_state( + document["parameters"][parameter.fqn], + empty_target, + implementation="gefen.Gefen", + global_step=4, + codebook=document["common"]["gefen_codebook"], + target_algorithm_options=_plain_options(), + target_second_moment="block", + ) + == {} + ) + + +@pytest.mark.parametrize("implementation", ["gefen.Gefen", "gefen.GefenMuon"]) +def test_complete_zero_numel_parameter_must_remain_pristine(implementation): + parameter = ParameterIdentity("empty.weight", (0, 3)) + options = _plain_options(second="factored") if implementation == "gefen.Gefen" else _muon_options() + policy = _policy(factored=implementation == "gefen.Gefen") + hints = {"source_periods": [1], "target_period": 1} + if implementation == "gefen.Gefen": + hints["source_second_moment"] = None + record = { + "identity": _serialize_parameter_identity(parameter), + "algorithm_options": options, + "state_variant": "period_selected", + "state": {}, + "projection_hints": hints, + } + document = build_portable_state_document( + implementation=implementation, + policy=policy, + common={ + "gefen_global_step": 0, + "gefen_codebook": _codebook(), + "gefen_deterministic": False, + }, + parameters={parameter.fqn: record}, + provenance=None, + ) + with pytest.raises(ValueError, match="empty logical parameters must remain pristine"): + _normalize_gefen_portable_state_document(document, limits=_limits()) + + +def test_assembly_authenticates_full_process_group_identity(): + parameter = ParameterIdentity("layer.weight", (2,)) + source_group = _group() + fragments = _fragments( + "gefen.Gefen", + parameter, + _replicas(parameter, source_group), + options=_plain_options(), + policy=_policy(), + variant="pristine", + ) + wrong_scope = ProcessGroupIdentity("different-checkpoint", source_group.ordered_members) + with pytest.raises(ValueError, match="exactly match process_group_identity"): + _assemble_portable_state_fragments( + fragments, + process_group_identity=wrong_scope, + limits=_limits(), + ) + + +def test_muon_schedule_respects_native_length_limit(): + parameter = ParameterIdentity("muon.weight", (2, 2)) + options = _muon_options() + options["ns_schedule"] = [[1.0, 2.0, 3.0] for _ in range(100)] + record = { + "identity": _serialize_parameter_identity(parameter), + "algorithm_options": options, + "state_variant": "pristine", + "state": {}, + "projection_hints": {"source_periods": [], "target_period": 1}, + } + document = build_portable_state_document( + implementation="gefen.GefenMuon", + policy=_policy(), + common={ + "gefen_global_step": 0, + "gefen_codebook": None, + "gefen_deterministic": False, + }, + parameters={parameter.fqn: record}, + provenance=None, + ) + with pytest.raises(ValueError, match="between 1 and 99"): + _normalize_gefen_portable_state_document(document, limits=_limits()) + + +def test_limits_are_strict_and_bound_tensor_and_member_aggregation(): + with pytest.raises(TypeError): + PortableStateLimits(True, 2, 64 << 20) + with pytest.raises(ValueError): + PortableStateLimits(2, 1, 64 << 20) + parameter = ParameterIdentity("layer.weight", (2,)) + group = _group() + fragments = _fragments( + "gefen.Gefen", + parameter, + _replicas(parameter, group), + options=_plain_options(), + policy=_policy(), + variant="initialized_dense", + momentum=torch.ones(2), + second=torch.ones(2), + ) + with pytest.raises(ValueError, match="max_fragment_tensor_bytes"): + _normalize_portable_state_fragment( + fragments[0], + limits=PortableStateLimits(1, 64, 64 << 20), + ) + with pytest.raises(ValueError, match="max_members"): + _assemble_portable_state_fragments( + fragments, + process_group_identity=group, + limits=dataclasses.replace(_limits(), max_members=1), + ) + with pytest.raises(ValueError, match="max_collective_tensor_bytes"): + _assemble_portable_state_fragments( + fragments, + process_group_identity=group, + limits=PortableStateLimits( + max_fragment_tensor_bytes=2048, + max_collective_tensor_bytes=2048, + max_collective_metadata_bytes=64 << 20, + ), + ) + + +def test_signed_zero_is_bitwise_consensus_not_numeric_consensus(): + parameter = ParameterIdentity("layer.weight", (1,)) + group = _group() + fragments = _fragments( + "gefen.Gefen", + parameter, + _replicas(parameter, group), + options=_plain_options(), + policy=_policy(), + variant="initialized_dense", + momentum=torch.tensor([-0.0]), + second=torch.tensor([1.0]), + ) + fragments[1]["logical_slots"][0]["state"]["momentum"][0] = 0.0 + with pytest.raises(ValueError, match="replicas disagree"): + _assemble_portable_state_fragments(fragments, process_group_identity=group, limits=_limits()) + + +def test_logical_slot_positions_must_agree_across_members(): + first = ParameterIdentity("first.weight", (1,)) + second = ParameterIdentity("second.weight", (1,)) + group = _group() + shards = _replicas(first, group) + _replicas(second, group) + manifest = ShardingManifest(shards) + common = { + "gefen_global_step": 0, + "gefen_codebook": None, + "gefen_deterministic": False, + } + catalog = { + parameter.fqn: { + "identity": _serialize_parameter_identity(parameter), + "algorithm_options": _plain_options(), + } + for parameter in (first, second) + } + fragments = [] + for member_index, member in enumerate(group.ordered_members): + member_shards = [shard for shard in shards if shard.local_member == member] + if member_index: + member_shards.reverse() + logical_slots = [] + for slot_index, shard in enumerate(member_shards): + logical_slots.append( + { + "group_index": 0, + "original_slot_index": slot_index, + "compatibility_name": shard.parameter.fqn, + "shard": _serialize_shard_identity(shard), + "algorithm_options": _plain_options(), + "role": "live", + "source_period": None, + "source_second_moment": None, + "state_variant": "pristine", + "state": {}, + } + ) + fragments.append( + _build_portable_state_fragment( + implementation="gefen.Gefen", + member=member, + policy=_policy(), + common=common, + manifest=manifest, + catalog=catalog, + logical_slots=logical_slots, + limits=_limits(), + ) + ) + with pytest.raises(ValueError, match="logical slot position disagree"): + _assemble_portable_state_fragments( + fragments, + process_group_identity=group, + limits=_limits(), + ) + + +def test_public_limits_are_frozen_and_only_public_export(): + limits = _limits() + with pytest.raises(dataclasses.FrozenInstanceError): + limits.chunk_bytes = 1 + assert math.isfinite(float(limits.max_collective_tensor_bytes)) From b078196fa885e9a7b0d98c5d340c11a62600ddbe Mon Sep 17 00:00:00 2001 From: thad0ctor Date: Mon, 13 Jul 2026 02:19:15 -0700 Subject: [PATCH 16/52] Add collective portable state I/O --- docs/optimizer_contracts.md | 43 +- src/gefen/__init__.py | 7 + src/gefen/contracts.py | 77 +- src/gefen/gefen.py | 69 +- src/gefen/gefen_muon.py | 32 +- src/gefen/portable_runtime.py | 1908 ++++++++++++++++++++ tests/test_portable_runtime.py | 400 ++++ tests/test_portable_runtime_consensus.py | 477 +++++ tests/test_portable_runtime_distributed.py | 567 ++++++ tests/test_portable_runtime_hardening.py | 457 +++++ tests/test_portable_schema.py | 2 + 11 files changed, 4027 insertions(+), 12 deletions(-) create mode 100644 src/gefen/portable_runtime.py create mode 100644 tests/test_portable_runtime.py create mode 100644 tests/test_portable_runtime_consensus.py create mode 100644 tests/test_portable_runtime_distributed.py create mode 100644 tests/test_portable_runtime_hardening.py diff --git a/docs/optimizer_contracts.md b/docs/optimizer_contracts.md index 7826cc7..499f8a1 100644 --- a/docs/optimizer_contracts.md +++ b/docs/optimizer_contracts.md @@ -22,7 +22,7 @@ assert ParameterLayout.DTENSOR_1D_DEFAULT_WORLD in rank_local_dcp.same_topology - `OptimizerStateLayout` separates optimizer-common authoritative state, per-parameter authoritative state, derived caches, checkpoint transport fields, and composite child namespaces. - `StateVariant` identifies valid lazy, initialized, local-shard, global-parameter, owner, non-owner, and migrated state combinations using structured layout, mode, rank, extent, ownership, and inactive-field declarations. - `TrainingSupport` qualifies each validated parameter layout by process-group source, mesh dimensionality, sharded mode, and whether the update needs complete parameter storage or a transient complete logical matrix. -- `CheckpointSupport` reports same-topology, topology-changing, and fail-before-mutation load support separately for native, PyTorch rank-local, and composite checkpoint transports. +- `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. @@ -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. DTensor stable identity, Hybrid composite rebinding, topology-changing canonical checkpoint I/O, and offload remain unclaimed. +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. The portable global-state path described below can reshard supported finalized layouts; DTensor stable identity consumption, Hybrid composite rebinding, and offload remain unclaimed. ## Explicit learned-codebook process groups @@ -51,7 +51,7 @@ The optimizer owns one learned codebook and therefore accepts one scope. Histogr `initialize_codebook()` and `refresh_codebook()` expose these operations for adapters that enter their optimizers in a deterministic order; `binding.sort_key` supplies the stable process-group portion of that schedule. Normal `step()` still initializes automatically and plain Gefen still honors `codebook_refresh_every`. Every scoped step exchanges a common operation header before any rank-dependent branch, and the first step after binding or native load additionally verifies codebook bytes and the complete manifest. Scoped native AMP requires every member to select the same protocol and present identical `found_inf` and `grad_scale` values; a mismatch raises collectively and requires a group-aware gradient scaler rather than changing external scaler state behind its back. Multi-member explicit scopes reject `capturable=True` because their host validation and process-group collectives are not CUDA-graph-safe. A one-member local scope may initialize during eager warmup and then use ordinary capturable stepping, but manual codebook replacement remains rejected. Ordinary unscoped behavior remains collective-free. Explicit scope does not replace DTensor mesh collectives, AMP mesh preflights, Parallel-Muon ownership collectives, or checkpoint transport groups. -Native checkpoints store a primitive rank-neutral scope record—format version, semantic name, ordered members, and refresh schedule—at the top level and in the existing transport mirror. Scoped group metadata uses a new outer format version so an older loader rejects rather than silently discarding the scope. Runtime process-group handles and collective devices are live adapter configuration and are never serialized or reconstructed. Flattened and whole-owner local payloads additionally carry a primitive rank-local shard-identity guard. Guard format v2 records every original logical slot in group and slot order, including its lowercase compatibility name and portable shard identity, so pruned whole-owner nonowners remain bound to their original positions and equal-shaped parameters cannot be reinterpreted positionally. New checkpoints emit v2, while the loader continues to accept v1 guards by comparing their legacy live-slot and separately sorted pruned-shard projection exactly. A scoped checkpoint requires the target to complete `post_sharding` with the same stable group identity, refresh policy, and native local layout before load. A missing or different record rejects inside the optimizer's atomic shadow load. Native flattened-shard checkpoint support is declared only on a finalized scoped instance with this guard; unscoped optimizers no longer claim it. Legacy unscoped checkpoints remain loadable by unscoped optimizers. Whole-owner checkpoint completeness, scoped DTensor rank-local transport, scope migration, topology-changing canonical I/O, and Hybrid-wide coordination are not claimed. +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, scope migration, and Hybrid-wide coordination remain separate from the portable global transport and are not claimed by the native path. ## Exact-binding canonical local state @@ -59,13 +59,42 @@ Native checkpoints store a primitive rank-neutral scope record—format version, The v1 document contains only primitive containers and detached tight finite CPU leaves of the ordinary dense `torch.Tensor` type, round-trips through `torch.load(weights_only=True)`, records the complete stable manifest, and maps each locally present parameter's authoritative state and algorithm-shaping group options by exact case-preserving FQN. After finalization the exact FQN is also the name used by period-routing policy, so compatibility names, devices, runtime process-group handles, derived lookup tables, capturable buffers, checkpoint carriers, and native positional IDs are not optimizer meaning in this format; the target retains its own lowercase compatibility names after import. Codebook scope identity is recorded, but its runtime handle and collective device must already be installed through `post_sharding`. -This is an exact-binding transport-neutral local fragment, not the dense global logical state planned for portable DCP v3. Its dynamic `CANONICAL_LOCAL` checkpoint entry covers finalized plain-Gefen replicated and flattened local shards and finalized replicated GefenMuon, performs no collectives, reports atomic local import, and has an empty topology-changing set. A different member, slice, manifest, algorithm policy, group option, or declared state variant rejects. Export, preparation, and commit are quiescent checkpoint-boundary operations; prepared imports use content-bearing freshness tokens, including device counters, to reject intervening mutation. Export remains available after capturable warmup, but a capturable import target must still be fresh, before authoritative device state or a CUDA graph exists; importing replaces state identities, so an already captured graph cannot safely remain attached. Configurations with `stochastic_round=True` do not claim canonical v1 because the decomposed path intentionally lacks the fused stochastic quantizer, so changing effective fused availability would change the algorithm. DTensor, whole-owner completeness, Hybrid composition, rank-fragment gathering, resharding, world-size change, dense momentum decoding, and target-topology recompression remain unclaimed. +This is an exact-binding transport-neutral local fragment, distinct from the dense global logical state in portable v3 below. Its dynamic `CANONICAL_LOCAL` checkpoint entry covers finalized plain-Gefen replicated and flattened local shards and finalized replicated GefenMuon, performs no collectives, reports atomic local import, and has an empty topology-changing set. A different member, slice, manifest, algorithm policy, group option, or declared state variant rejects. Export, preparation, and commit are quiescent checkpoint-boundary operations; prepared imports use content-bearing freshness tokens, including device counters, to reject intervening mutation. Export remains available after capturable warmup, but a capturable import target must still be fresh, before authoritative device state or a CUDA graph exists; importing replaces state identities, so an already captured graph cannot safely remain attached. Configurations with `stochastic_round=True` do not claim canonical v1 because the decomposed path intentionally lacks the fused stochastic quantizer, so changing effective fused availability would change the algorithm. DTensor, whole-owner completeness, Hybrid composition, rank-fragment gathering, resharding, world-size change, dense momentum decoding, and target-topology recompression remain unclaimed by `CANONICAL_LOCAL`; supported collective assembly and resharding use `CANONICAL_GLOBAL`. -## Portable global-state v3 envelope +## Portable global-state v3 -The transport-neutral `gefen.portable_state` version-3 format is the structural wire envelope for a future complete `global_logical_optimizer` artifact. 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, optional source provenance, and a completion marker whose deterministic SHA-256 covers every preceding value including tensor dtype, shape, and canonical little-endian bytes. The generic builder and normalizer validate the exact envelope, canonical value grammar, parameter identity records, completion marker, and digest; they do not by themselves prove model-catalog completeness or optimizer-specific state variants, tensor geometries, and projection hints. They accept only finite weights-only-safe primitive values, stream tensor cloning, finite validation, and hashing in bounded chunks, produce detached tight CPU tensors, and reject an incomplete or corrupted digest. Runtime global ranks are deliberately 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, 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. -`CheckpointProcessGroupBinding` separately binds an adapter-defined `ProcessGroupIdentity` and semantic local member to an explicit live PyTorch process group and CPU/CUDA collective device. Multi-member scopes, including the default world, must pass an explicit handle; `None` means exactly one member. Runtime validation checks initialized membership, group size and coordinate order, and backend/device compatibility without serializing global ranks or executing a collective. The `CANONICAL_GLOBAL` transport enum is defined for this v3 path, but Gefen and GefenMuon do not yet advertise it: a positive capability requires collective fragment registration, dense aggregation, target projection/recompression, unanimous preparation, and atomic local publication to be connected end to end. +`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` and `GefenMuon` instances. The exact `CheckpointProcessGroupBinding` must match the optimizer's installed `CodebookProcessGroupBinding` in stable identity, local member, live process-group handle, and collective device; all collectives use that optimizer-owned transport. Every member must enter operations in the same deterministic order with the same trimmed transaction ID, limits, target logical slot schema, and, for import, complete document. `PortableStateLimits` bounds per-member and aggregate tensor bytes, metadata, tree shape, strings, tensor count and rank, member count, wire chunk size, and diagnostics; structural limits are checked before dense materialization. `chunk_bytes` controls wire cloning and collective transfer, runtime value validation and freshness hashing use independently fixed bounded chunks, and dense decode/projection allocations remain bounded by the declared fragment and collective tensor-byte ceilings. + +```python +from gefen import PortableStateLimits + +limits = PortableStateLimits( + max_fragment_tensor_bytes=2 << 30, + max_collective_tensor_bytes=16 << 30, + max_collective_metadata_bytes=256 << 20, +) + +document = optimizer.export_portable_state( + checkpoint_process_group=checkpoint_binding, + transaction_id="optimizer-save-0001", + limits=limits, +) + +target_optimizer.import_portable_state( + document, + checkpoint_process_group=target_checkpoint_binding, + transaction_id="optimizer-load-0001", + limits=limits, +) +``` + +The dynamic `CANONICAL_GLOBAL` checkpoint declaration appears only while the live finalized optimizer passes the exact runtime readiness checks: explicit process-group scope, stable logical slots and manifest, ordinary built-in containers, supported CPU/CUDA tensor storage, no active compilation or CUDA capture, `capturable=False`, `stochastic_round=False`, a complete declared native state variant, and period one for selected or initialized state. Plain Gefen supports replicated and contiguous flattened element shards. Block-second-moment state can reshard between replicated and flattened targets; a logical matrix using factored second moments remains replicated and same-topology because factored-to-block representation migration is not implemented. GefenMuon supports replicated matrices and whole-parameter ownership when every participating group uses `sharded_mode="distributed"`; the transport can change placement and redistribute owners across world sizes, including NorMuon row state. 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, GefenMuon modes other than `distributed` for whole-owner transport, `GefenMuonHybrid`, tied-alias expansion, and direct DCP orchestration. The optimizer-facing API is the adapter boundary for a DCP or platform integration; the core does not register a framework-specific state-dict adapter or perform storage I/O. `state_offload` remains false because a CPU portable document is a checkpoint artifact, not live optimizer state that can be stepped while offloaded. ## Quiescent optimizer-state movement diff --git a/src/gefen/__init__.py b/src/gefen/__init__.py index e52302e..87129dc 100644 --- a/src/gefen/__init__.py +++ b/src/gefen/__init__.py @@ -21,6 +21,8 @@ "CheckpointProcessGroupBinding", "CodebookProcessGroupBinding", "PreparedCanonicalStateImport", + "PortableStateLimits", + "PortableStateProvider", "OptimizerCapabilities", "OptimizerChildContract", "OptimizerContract", @@ -103,6 +105,10 @@ def __getattr__(name): from . import portable_schema return getattr(portable_schema, name) + if name == "PortableStateLimits": + from .portable_state import PortableStateLimits + + return PortableStateLimits if name in ( "CONTRACT_SCHEMA_VERSION", "IDENTITY_SCHEMA_VERSION", @@ -120,6 +126,7 @@ def __getattr__(name): "ParameterIdentity", "ParameterStateRole", "PlacementKind", + "PortableStateProvider", "Precision", "ProcessGroupScope", "ProcessGroupIdentity", diff --git a/src/gefen/contracts.py b/src/gefen/contracts.py index 7b80234..eaefe3e 100644 --- a/src/gefen/contracts.py +++ b/src/gefen/contracts.py @@ -1022,6 +1022,26 @@ def import_canonical_state(self, state) -> None: """Prepare and commit a canonical fragment atomically.""" +@runtime_checkable +class PortableStateProvider(Protocol): + """Structural protocol for collective topology-neutral state I/O.""" + + def export_portable_state( + self, *, checkpoint_process_group, transaction_id, limits + ): + """Collectively return one complete portable global-state document.""" + + def import_portable_state( + self, + state, + *, + checkpoint_process_group, + transaction_id, + limits, + ) -> None: + """Collectively stage and atomically publish portable global state.""" + + @runtime_checkable class StateMovementProvider(Protocol): """Structural protocol for quiescent atomic optimizer-state movement.""" @@ -1216,9 +1236,19 @@ def _gefen_contract( explicit_process_group_codebook_scope: bool = False, native_flattened_checkpoint: bool = False, canonical_state_layouts: AbstractSet[ParameterLayout] = frozenset(), + canonical_global_same_topology: AbstractSet[ParameterLayout] = frozenset(), + canonical_global_topology_changing: AbstractSet[ParameterLayout] = frozenset(), + canonical_global_topology_change_kinds: AbstractSet[TopologyChange] = frozenset(), atomic_state_movement: bool = False, ) -> OptimizerContract: canonical_state_layouts = _frozenset(canonical_state_layouts) + canonical_global_same_topology = _frozenset(canonical_global_same_topology) + canonical_global_topology_changing = _frozenset( + canonical_global_topology_changing + ) + canonical_global_topology_change_kinds = _frozenset( + canonical_global_topology_change_kinds + ) block_fields = ( StateField("vmean", StateScope.PARAMETER, StateGeometry.BLOCK, True), StateField("vmean_step", StateScope.PARAMETER, StateGeometry.SCALAR, True), @@ -1403,6 +1433,18 @@ def _gefen_contract( atomic_load=True, ) ) + if canonical_global_same_topology or canonical_global_topology_changing: + checkpoints.append( + CheckpointSupport( + CheckpointTransport.CANONICAL_GLOBAL, + canonical_global_same_topology, + canonical_global_topology_changing, + ProcessGroupScope.ADAPTER_DEFINED, + topology_change_kinds=canonical_global_topology_change_kinds, + requires_collective=True, + atomic_load=True, + ) + ) return OptimizerContract( implementation="gefen.Gefen", state_layout=OptimizerStateLayout(fields, tuple(variants)), @@ -1415,7 +1457,11 @@ def _gefen_contract( explicit_process_group_codebook_scope=explicit_process_group_codebook_scope, shard_rebinding=True, post_sharding=True, - canonical_state_io=bool(canonical_state_layouts), + canonical_state_io=bool( + canonical_state_layouts + or canonical_global_same_topology + or canonical_global_topology_changing + ), atomic_state_movement=atomic_state_movement, ), ) @@ -1451,9 +1497,19 @@ def _gefen_muon_contract( explicit_process_group_codebook_scope: bool = False, whole_parameter_owner: bool = False, canonical_state_layouts: AbstractSet[ParameterLayout] = frozenset(), + canonical_global_same_topology: AbstractSet[ParameterLayout] = frozenset(), + canonical_global_topology_changing: AbstractSet[ParameterLayout] = frozenset(), + canonical_global_topology_change_kinds: AbstractSet[TopologyChange] = frozenset(), atomic_state_movement: bool = False, ) -> OptimizerContract: canonical_state_layouts = _frozenset(canonical_state_layouts) + canonical_global_same_topology = _frozenset(canonical_global_same_topology) + canonical_global_topology_changing = _frozenset( + canonical_global_topology_changing + ) + canonical_global_topology_change_kinds = _frozenset( + canonical_global_topology_change_kinds + ) sharded_modes = _frozenset(sharded_modes) normuon_modes = _frozenset(normuon_modes) non_normuon_modes = _frozenset(non_normuon_modes) @@ -1744,6 +1800,19 @@ def _gefen_muon_contract( atomic_load=True, ) ) + if canonical_global_same_topology or canonical_global_topology_changing: + checkpoints.append( + CheckpointSupport( + CheckpointTransport.CANONICAL_GLOBAL, + canonical_global_same_topology, + canonical_global_topology_changing, + ProcessGroupScope.ADAPTER_DEFINED, + topology_change_kinds=canonical_global_topology_change_kinds, + required_sharded_modes=sharded_modes, + requires_collective=True, + atomic_load=True, + ) + ) return OptimizerContract( implementation="gefen.GefenMuon", state_layout=OptimizerStateLayout(fields, tuple(variants)), @@ -1756,7 +1825,11 @@ def _gefen_muon_contract( explicit_process_group_codebook_scope=explicit_process_group_codebook_scope, shard_rebinding=True, post_sharding=True, - canonical_state_io=bool(canonical_state_layouts), + canonical_state_io=bool( + canonical_state_layouts + or canonical_global_same_topology + or canonical_global_topology_changing + ), atomic_state_movement=atomic_state_movement, ), ) diff --git a/src/gefen/gefen.py b/src/gefen/gefen.py index 3d982b9..71fc121 100644 --- a/src/gefen/gefen.py +++ b/src/gefen/gefen.py @@ -43,6 +43,7 @@ ShardIdentity, ShardPlacement, ShardingManifest, + TopologyChange, _gefen_contract, ) from gefen.partitioning import find_period_by_block_variance @@ -1248,12 +1249,38 @@ def optimizer_contract(self) -> OptimizerContract: canonical_state_layouts = self._canonical_state_layouts() except Exception: canonical_state_layouts = frozenset() + try: + from gefen.portable_runtime import _portable_runtime_layouts + + canonical_global_same_topology = _portable_runtime_layouts(self) + except Exception: + canonical_global_same_topology = frozenset() + has_factored_matrix = canonical_global_same_topology and self._factored_v_2d and any( + len(slot.shard.parameter.global_shape) == 2 + for slot in self._gefen_logical_slots + ) + if canonical_global_same_topology and not has_factored_matrix: + canonical_global_topology_changing = frozenset( + { + ParameterLayout.REPLICATED, + ParameterLayout.FLATTENED_ELEMENT_SHARD, + } + ) + canonical_global_topology_change_kinds = frozenset( + {TopologyChange.PLACEMENT_RESHARD} + ) + else: + canonical_global_topology_changing = frozenset() + canonical_global_topology_change_kinds = frozenset() return _gefen_contract( factored_v_2d=self._factored_v_2d, canonical_parameter_fqns=identity_ready, stable_shard_identity=identity_ready, explicit_process_group_codebook_scope=True, canonical_state_layouts=canonical_state_layouts, + canonical_global_same_topology=canonical_global_same_topology, + canonical_global_topology_changing=canonical_global_topology_changing, + canonical_global_topology_change_kinds=canonical_global_topology_change_kinds, atomic_state_movement=self._atomic_state_movement_supported(), native_flattened_checkpoint=( self._codebook_scope_ready() @@ -6864,6 +6891,44 @@ def import_canonical_state(self, state) -> None: self.prepare_canonical_state_import(state) ) + def export_portable_state( + self, + *, + checkpoint_process_group, + transaction_id, + limits, + ): + """Collectively export exact period-one state by logical identity.""" + + from gefen.portable_runtime import _export_portable_state + + return _export_portable_state( + self, + checkpoint_process_group=checkpoint_process_group, + transaction_id=transaction_id, + limits=limits, + ) + + def import_portable_state( + self, + state, + *, + checkpoint_process_group, + transaction_id, + limits, + ) -> None: + """Collectively stage and atomically publish exact portable state.""" + + from gefen.portable_runtime import _import_portable_state + + _import_portable_state( + self, + state, + checkpoint_process_group=checkpoint_process_group, + transaction_id=transaction_id, + limits=limits, + ) + canonical_state_dict = export_canonical_state load_canonical_state_dict = import_canonical_state @@ -7690,9 +7755,9 @@ def _commit_staged_load_state_dict(self, staged) -> None: # preserve that public object identity while publishing the staged # value. Both mappings are ordinary built-in dicts created by Optimizer. live_defaults = self.defaults - live_defaults.update(staged.defaults) + dict.update(live_defaults, staged.defaults) staged.defaults = live_defaults - self.__dict__.update(staged.__dict__) + dict.update(self.__dict__, staged.__dict__) def _validate_loaded_native_state(self) -> None: """Validate the complete prepared native state before publication.""" diff --git a/src/gefen/gefen_muon.py b/src/gefen/gefen_muon.py index ffe1eb1..e3454b3 100644 --- a/src/gefen/gefen_muon.py +++ b/src/gefen/gefen_muon.py @@ -7,7 +7,12 @@ import torch import torch.nn as nn -from gefen.contracts import OptimizerContract, ParameterLayout, _gefen_muon_contract +from gefen.contracts import ( + OptimizerContract, + ParameterLayout, + TopologyChange, + _gefen_muon_contract, +) from gefen.gefen import ( Gefen, _assert_optimizer_gradients_structurally_valid, @@ -772,6 +777,28 @@ def optimizer_contract(self) -> OptimizerContract: canonical_state_layouts = self._canonical_state_layouts() except Exception: canonical_state_layouts = frozenset() + try: + from gefen.portable_runtime import _portable_runtime_layouts + + canonical_global_same_topology = _portable_runtime_layouts(self) + except Exception: + canonical_global_same_topology = frozenset() + if canonical_global_same_topology and sharded_modes == frozenset({"distributed"}): + canonical_global_topology_changing = frozenset( + { + ParameterLayout.REPLICATED, + ParameterLayout.WHOLE_PARAMETER_OWNER, + } + ) + canonical_global_topology_change_kinds = frozenset( + { + TopologyChange.PLACEMENT_RESHARD, + TopologyChange.WORLD_SIZE_OWNER_REDISTRIBUTION, + } + ) + else: + canonical_global_topology_changing = frozenset() + canonical_global_topology_change_kinds = frozenset() return _gefen_muon_contract( sharded_modes=sharded_modes, normuon_modes=normuon_modes, @@ -780,6 +807,9 @@ def optimizer_contract(self) -> OptimizerContract: stable_shard_identity=self._canonical_identity_ready(), explicit_process_group_codebook_scope=True, canonical_state_layouts=canonical_state_layouts, + canonical_global_same_topology=canonical_global_same_topology, + canonical_global_topology_changing=canonical_global_topology_changing, + canonical_global_topology_change_kinds=canonical_global_topology_change_kinds, atomic_state_movement=self._atomic_state_movement_supported(), whole_parameter_owner=( self._codebook_scope_ready() diff --git a/src/gefen/portable_runtime.py b/src/gefen/portable_runtime.py new file mode 100644 index 0000000..353fdc0 --- /dev/null +++ b/src/gefen/portable_runtime.py @@ -0,0 +1,1908 @@ +"""Optimizer-facing collective runtime for exact portable Gefen state.""" + +from __future__ import annotations + +from collections import defaultdict +import hashlib +import math + +import torch +from torch import nn + +from gefen.canonical import canonical_value_supported +from gefen.checkpoint import CheckpointProcessGroupBinding +from gefen.codebook import CodebookProcessGroupBinding +from gefen.contracts import ( + IDENTITY_SCHEMA_VERSION, + LogicalSlice, + ParameterIdentity, + ParameterLayout, + PlacementKind, + ProcessGroupIdentity, + ShardIdentity, + ShardPlacement, + ShardingManifest, +) +from gefen.portable import ( + _decode_quantized_momentum, + _expand_block_second_moment, + _read_flat_chunk, +) +from gefen.portable_collective import ( + _collective_unanimous_status, + _collective_visit_canonical_fragments, +) +from gefen.portable_identity import ( + _serialize_parameter_identity, + _serialize_process_group_identity, + _serialize_shard_identity, + _serialize_sharding_manifest, +) +from gefen.portable_schema import portable_state_digest +from gefen.portable_state import ( + PortableStateLimits, + _MOMENTUM_PROJECTION, + _SECOND_MOMENT_PROJECTION, + _assemble_portable_state_fragments, + _build_portable_state_fragment, + _derived_role, + _local_dense_shape, + _normalize_common, + _normalize_gefen_portable_state_document, + _normalize_options, + _normalize_policy, + _normalize_portable_state_fragment, + _project_portable_parameter_state, + _values_equal, +) +from gefen.rebinding import LogicalSlotBinding + + +_PLAIN_IMPLEMENTATION = "gefen.Gefen" +_MUON_IMPLEMENTATION = "gefen.GefenMuon" +_MAX_TRANSACTION_BYTES = 1024 +_EXPORT_PREFLIGHT_TRANSACTION = "gefen-portable-export-preflight-v1" +_IMPORT_PREFLIGHT_TRANSACTION = "gefen-portable-import-preflight-v1" +_STATUS_FALLBACK_LIMITS = PortableStateLimits( + max_fragment_tensor_bytes=1, + max_collective_tensor_bytes=1, + max_collective_metadata_bytes=1, + max_metadata_bytes=1, +)._wire_limits() + +_PLAIN_GROUP_REQUIRED = frozenset({"params", "param_names", "lr", "beta1", "beta2", "eps", "weight_decay"}) +_PLAIN_GROUP_ALLOWED = _PLAIN_GROUP_REQUIRED | {"name"} +_MUON_GROUP_REQUIRED = _PLAIN_GROUP_REQUIRED | frozenset( + { + "momentum", + "nesterov", + "ns_coefficients", + "ns_steps", + "adjust_lr_fn", + "sharded_mode", + "fp8_ns", + "fp8_ns_compile", + "batched_ns", + "batched_ns_workspace_bytes", + "normuon", + "normuon_beta2", + "normuon_eps", + "cautious", + } +) +_MUON_GROUP_ALLOWED = _MUON_GROUP_REQUIRED | {"name", "ns_schedule"} + +_AUTHORITATIVE_STATE_KEYS = frozenset( + { + "name", + "automatic_period", + "step", + "m_codebook", + "m_magnitude", + "vmean", + "vmean_step", + "v_row", + "v_col", + "factored_step", + "normuon_v", + "normuon_step", + } +) +_IGNORED_DERIVED_STATE_KEYS = frozenset({"stepsize", "_h_buf", "m_codebook_shape"}) +_PLAIN_BLOCK_KEYS = frozenset( + { + "name", + "automatic_period", + "step", + "m_codebook", + "m_magnitude", + "vmean", + "vmean_step", + } +) +_PLAIN_FACTORED_KEYS = frozenset( + { + "name", + "automatic_period", + "step", + "m_codebook", + "m_magnitude", + "v_row", + "v_col", + "factored_step", + } +) +_MUON_KEYS = frozenset({"name", "automatic_period", "step", "m_codebook", "m_magnitude"}) +_NORMUON_KEYS = _MUON_KEYS | {"normuon_v", "normuon_step"} + + +def _bounded_utf8_length(value: str, *, limit: int, name: str) -> int: + total = 0 + for start in range(0, len(value), 4096): + chunk = value[start : start + 4096].encode("utf-8") + if total > limit - len(chunk): + raise ValueError("{} exceeds its UTF-8 byte limit".format(name)) + total += len(chunk) + return total + + +class _PreflightTensor: + __slots__ = ("shape",) + + def __init__(self, shape): + self.shape = tuple(shape) + + +def _preflight_portable_value(value, limits: PortableStateLimits) -> None: + """Apply exact wire structural limits without allocating tensor payloads.""" + + wire = limits._wire_limits() + nodes = 0 + tensor_bytes = 0 + tensor_shapes = [] + metadata_bytes = 16 + + def add_metadata(count): + nonlocal metadata_bytes + metadata_bytes += count + if metadata_bytes > wire.max_metadata_bytes: + raise ValueError("portable fragment exceeds max_metadata_bytes before materialization") + + def visit(item, *, depth): + nonlocal nodes, tensor_bytes + if depth > wire.max_tree_depth: + raise ValueError("portable fragment exceeds max_tree_depth before materialization") + nodes += 1 + if nodes > wire.max_tree_nodes: + raise ValueError("portable fragment exceeds max_tree_nodes before materialization") + item_type = type(item) + if item is None or item_type is bool: + add_metadata(1) + return + if item_type is int: + magnitude_bytes = (abs(item).bit_length() + 7) // 8 + if magnitude_bytes > wire.max_integer_bytes: + raise ValueError("portable fragment integer exceeds max_integer_bytes") + add_metadata(10 + magnitude_bytes) + return + if item_type is float: + if not math.isfinite(item): + raise ValueError("portable fragment contains a nonfinite float") + add_metadata(9) + return + if item_type is str: + encoded_bytes = _bounded_utf8_length( + item, + limit=wire.max_string_bytes, + name="portable fragment string", + ) + add_metadata(9 + encoded_bytes) + return + if item_type is _PreflightTensor: + shape = item.shape + if any(type(dimension) is not int or dimension < 0 for dimension in shape): + raise ValueError("portable fragment tensor has invalid geometry") + if len(shape) > wire.max_tensor_rank: + raise ValueError("portable fragment tensor exceeds max_tensor_rank") + if len(tensor_shapes) >= wire.max_tensors: + raise ValueError("portable fragment exceeds max_tensors") + numel = math.prod(shape) + nbytes = numel * 4 + if tensor_bytes > wire.max_fragment_tensor_bytes - nbytes: + raise ValueError("portable fragment exceeds max_fragment_tensor_bytes before materialization") + tensor_bytes += nbytes + tensor_shapes.append(shape) + add_metadata(9) + return + if item_type in {list, tuple}: + if len(item) > wire.max_container_items: + raise ValueError("portable fragment container exceeds max_container_items") + add_metadata(9) + for child in item: + visit(child, depth=depth + 1) + return + if item_type is dict: + if len(item) > wire.max_container_items: + raise ValueError("portable fragment container exceeds max_container_items") + if any(type(key) is not str for key in item): + raise TypeError("portable fragment dictionary keys must be strings") + add_metadata(9) + for key in sorted(item): + visit(key, depth=depth + 1) + visit(item[key], depth=depth + 1) + return + raise TypeError("portable preflight encountered an unsupported value") + + visit(value, depth=0) + add_metadata(8) + for shape in tensor_shapes: + add_metadata(56 + 8 * len(shape)) + + +def _optimizer_implementation(optimizer) -> str: + # Keep imports lazy so Gefen may delegate to this module without creating a + # module-import cycle. + from gefen.gefen import Gefen + from gefen.gefen_muon import GefenMuon + + if type(optimizer) is Gefen: + return _PLAIN_IMPLEMENTATION + if type(optimizer) is GefenMuon: + return _MUON_IMPLEMENTATION + raise TypeError("portable state supports exact Gefen and GefenMuon instances only") + + +def _reject_method_shadows(optimizer) -> None: + if type(optimizer.__dict__) is not dict: + raise TypeError("portable state requires an exact optimizer attribute mapping") + for name in optimizer.__dict__: + descriptor = None + for owner in type(optimizer).__mro__: + if name in owner.__dict__: + descriptor = owner.__dict__[name] + break + if isinstance(descriptor, (staticmethod, classmethod)): + descriptor = descriptor.__func__ + if callable(descriptor): + raise TypeError("portable state rejects instance-level method shadows") + + +def _require_limits(limits) -> PortableStateLimits: + if type(limits) is not PortableStateLimits: + raise TypeError("limits must be a PortableStateLimits") + return limits + + +def _require_binding(binding) -> CheckpointProcessGroupBinding: + if type(binding) is not CheckpointProcessGroupBinding: + raise TypeError("checkpoint_process_group must be a CheckpointProcessGroupBinding") + return binding + + +def _validate_exact_process_group_identity(identity) -> ProcessGroupIdentity: + if ( + type(identity) is not ProcessGroupIdentity + or type(identity.semantic_name) is not str + or type(identity.ordered_members) is not tuple + or any(type(member) is not str for member in identity.ordered_members) + or type(identity.schema_version) is not int + ): + raise TypeError("portable state requires an exact ProcessGroupIdentity") + validated = ProcessGroupIdentity( + identity.semantic_name, + identity.ordered_members, + schema_version=identity.schema_version, + ) + if validated != identity: + raise ValueError("portable process-group identities must be canonical") + return validated + + +def _validate_exact_shard_identity(shard) -> ShardIdentity: + if type(shard) is not ShardIdentity: + raise TypeError("portable state requires exact ShardIdentity values") + if type(shard.parameter) is not ParameterIdentity: + raise TypeError("portable state requires exact ParameterIdentity values") + if ( + type(shard.parameter.fqn) is not str + or type(shard.parameter.global_shape) is not tuple + or any(type(dimension) is not int for dimension in shard.parameter.global_shape) + or type(shard.parameter.schema_version) is not int + ): + raise TypeError("portable parameter identities require exact primitive fields") + parameter = ParameterIdentity( + shard.parameter.fqn, + shard.parameter.global_shape, + schema_version=shard.parameter.schema_version, + ) + if type(shard.logical_slice) is not LogicalSlice: + raise TypeError("portable state requires exact LogicalSlice values") + if ( + type(shard.logical_slice.flat_offset) is not int + or type(shard.logical_slice.length) is not int + or type(shard.local_member) is not str + or (shard.owner is not None and type(shard.owner) is not str) + or type(shard.schema_version) is not int + ): + raise TypeError("portable shard identities require exact primitive fields") + logical_slice = LogicalSlice( + shard.logical_slice.flat_offset, + shard.logical_slice.length, + ) + process_group = _validate_exact_process_group_identity(shard.process_group) + if type(shard.placements) is not tuple or any( + type(placement) is not ShardPlacement for placement in shard.placements + ): + raise TypeError("portable state requires exact ShardPlacement values") + if any( + type(placement.mesh_axis) is not str + or type(placement.kind) is not PlacementKind + or type(placement.coordinate) is not int + or type(placement.parts) is not int + or (placement.parameter_dimension is not None and type(placement.parameter_dimension) is not int) + for placement in shard.placements + ): + raise TypeError("portable shard placements require exact primitive fields") + placements = tuple( + ShardPlacement( + placement.mesh_axis, + placement.kind, + placement.coordinate, + placement.parts, + placement.parameter_dimension, + ) + for placement in shard.placements + ) + if type(shard.layout) is not ParameterLayout: + raise TypeError("portable shard layouts require exact ParameterLayout values") + validated = ShardIdentity( + parameter, + shard.layout, + logical_slice, + placements=placements, + process_group=process_group, + local_member=shard.local_member, + owner=shard.owner, + schema_version=shard.schema_version, + ) + if validated != shard: + raise ValueError("portable shard identities must be canonical") + return validated + + +def _validate_exact_manifest(manifest) -> None: + if ( + type(manifest) is not ShardingManifest + or type(manifest.shards) is not tuple + or type(manifest.schema_version) is not int + ): + raise TypeError("portable state requires an exact ShardingManifest") + validated = ShardingManifest( + tuple(_validate_exact_shard_identity(shard) for shard in manifest.shards), + schema_version=manifest.schema_version, + ) + if validated != manifest: + raise ValueError("portable sharding manifests must be canonical") + + +def _require_transaction_id(transaction_id) -> str: + if ( + type(transaction_id) is not str + or not transaction_id + or transaction_id != transaction_id.strip() + or "\x00" in transaction_id + ): + raise ValueError("transaction_id must be non-empty, trimmed, and without NUL") + _bounded_utf8_length( + transaction_id, + limit=_MAX_TRANSACTION_BYTES, + name="transaction_id", + ) + return transaction_id + + +def _preflight_transport_binding(optimizer, supplied): + scope = getattr(optimizer, "_gefen_codebook_process_group", None) + if type(scope) is CodebookProcessGroupBinding: + identity = ProcessGroupIdentity( + str(scope.identity.semantic_name), + tuple(str(member) for member in scope.identity.ordered_members), + schema_version=IDENTITY_SCHEMA_VERSION, + ) + return CheckpointProcessGroupBinding( + identity, + str(scope.local_member), + scope.process_group, + scope.collective_device, + ) + if isinstance(supplied, CheckpointProcessGroupBinding): + return supplied + raise TypeError("portable preflight requires a transport-usable process group") + + +def _validate_supplied_binding(supplied, transport) -> None: + supplied = _require_binding(supplied) + _validate_exact_process_group_identity(supplied.identity) + _validate_exact_process_group_identity(transport.identity) + if ( + type(supplied.local_member) is not str + or type(supplied.collective_device) is not torch.device + or type(transport.local_member) is not str + or type(transport.collective_device) is not torch.device + ): + raise TypeError("portable checkpoint bindings require exact primitive fields") + if supplied.identity != transport.identity or supplied.local_member != transport.local_member: + raise ValueError("checkpoint and codebook process-group identities must match") + if ( + supplied.process_group is not transport.process_group + or supplied.collective_device != transport.collective_device + ): + raise ValueError("checkpoint_process_group must exactly match the optimizer-owned runtime transport") + + +def _validate_context_identity_bounds( + binding: CheckpointProcessGroupBinding, + limits: PortableStateLimits, +) -> None: + identity = binding.identity + member_count = len(identity.ordered_members) + if member_count > limits.max_members or member_count > limits.max_container_items: + raise ValueError("checkpoint process-group identity exceeds member limits") + _bounded_utf8_length( + identity.semantic_name, + limit=limits.max_string_bytes, + name="process-group semantic name", + ) + for member in identity.ordered_members: + _bounded_utf8_length( + member, + limit=limits.max_string_bytes, + name="process-group member", + ) + + +def _context_digest(value) -> bytes: + return bytes.fromhex(portable_state_digest(value)) + + +def _base_context(binding: CheckpointProcessGroupBinding, implementation: str): + return { + "format": "gefen.portable_runtime_context", + "format_version": 1, + "implementation": implementation, + "checkpoint_process_group": _serialize_process_group_identity(binding.identity), + } + + +def _plain_tensor( + value, + *, + name: str, + dtype=None, + shape=None, + nonnegative=False, + validate_values=True, +): + if ( + type(value) is not torch.Tensor + or value.layout is not torch.strided + or value.is_meta + or value.is_nested + or value.is_quantized + or value.requires_grad + or value.is_conj() + or value.is_neg() + ): + raise TypeError("{} must be a plain materialized strided tensor".format(name)) + if dtype is not None and value.dtype != dtype: + raise TypeError("{} has an invalid dtype".format(name)) + if shape is not None and tuple(value.shape) != tuple(shape): + raise ValueError("{} has invalid geometry".format(name)) + if value.is_floating_point() and validate_values: + for start in range(0, value.numel(), 1 << 20): + chunk = _read_flat_chunk( + value, + start, + min(start + (1 << 20), value.numel()), + ) + if not bool(torch.isfinite(chunk).all()): + raise ValueError("{} must be finite".format(name)) + if nonnegative and not bool((chunk >= 0).all()): + raise ValueError("{} must be nonnegative".format(name)) + return value + + +def _tight_cpu_fp32(value, *, name: str, shape=None, nonnegative=False): + value = _plain_tensor( + value, + name=name, + dtype=torch.float32, + shape=shape, + nonnegative=nonnegative, + ) + result = torch.empty(tuple(value.shape), dtype=torch.float32, device="cpu") + result.copy_(value.detach()) + return result + + +def _strict_counter(value, *, name: str, minimum=0): + if type(value) is not int or value < minimum or value > (1 << 53) - 1: + raise ValueError("{} must be an exact bounded host int".format(name)) + return value + + +def _live_policy(optimizer, implementation: str): + raw = optimizer._canonical_policy() + if type(raw) is not dict: + raise ValueError("live portable policy must be a plain dictionary") + policy = { + "schema_version": 1, + "factored_v_2d": raw.get("factored_v_2d"), + "force_1d_period_one": raw.get("force_1d_period_one"), + "force_2d_period_one": raw.get("force_2d_period_one"), + "period_one_substrings": raw.get("period_one_substrings"), + "codebook_refresh_every": raw.get("codebook_refresh_every"), + "stochastic_round": raw.get("stochastic_round"), + "momentum_projection": _MOMENTUM_PROJECTION, + "second_moment_projection": _SECOND_MOMENT_PROJECTION, + } + return _normalize_policy(policy, implementation) + + +def _normalized_ns_schedule(group): + from gefen.gefen_muon import _normalize_ns_schedule + + schedule = _normalize_ns_schedule(group["ns_coefficients"], group["ns_steps"]) + return [[float(a), float(b), float(c)] for a, b, c in schedule] + + +def _group_options(group, implementation: str, *, second_moment_policy=None): + if type(group) is not dict: + raise TypeError("portable parameter groups must be plain dictionaries") + allowed = _PLAIN_GROUP_ALLOWED if implementation == _PLAIN_IMPLEMENTATION else _MUON_GROUP_ALLOWED + required = _PLAIN_GROUP_REQUIRED if implementation == _PLAIN_IMPLEMENTATION else _MUON_GROUP_REQUIRED + if not required.issubset(group) or not set(group).issubset(allowed): + raise ValueError("portable parameter group contains missing or unknown keys") + if type(group["params"]) is not list or type(group["param_names"]) is not list: + raise TypeError("portable parameter group params and names must be plain lists") + + if implementation == _PLAIN_IMPLEMENTATION: + return _normalize_options( + { + "lr": group["lr"], + "beta1": group["beta1"], + "beta2": group["beta2"], + "eps": group["eps"], + "weight_decay": group["weight_decay"], + "second_moment_policy": second_moment_policy, + }, + implementation, + ) + + if any(type(group[key]) is not float for key in ("beta1", "beta2", "momentum")): + raise ValueError("Muon inherited betas and momentum must be exact floats") + if group["beta1"] != group["momentum"] or group["beta2"] != 0.0: + raise ValueError("Muon inherited beta values disagree with its momentum policy") + if "ns_schedule" in group and not canonical_value_supported(group["ns_schedule"], finite_tensors=True): + raise ValueError("Muon raw ns_schedule must remain weights-only-safe metadata") + return _normalize_options( + { + "lr": group["lr"], + "weight_decay": group["weight_decay"], + "momentum": group["momentum"], + "nesterov": group["nesterov"], + "ns_schedule": _normalized_ns_schedule(group), + "ns_eps": group["eps"], + "adjust_lr_fn": group["adjust_lr_fn"], + "sharded_mode": group["sharded_mode"], + "fp8_ns": group["fp8_ns"], + "fp8_ns_compile": group["fp8_ns_compile"], + "batched_ns": group["batched_ns"], + "batched_ns_workspace_bytes": group["batched_ns_workspace_bytes"], + "normuon": group["normuon"], + "normuon_beta2": group["normuon_beta2"], + "normuon_eps": group["normuon_eps"], + "cautious": group["cautious"], + }, + implementation, + ) + + +def _parameter_supported(parameter, shard: ShardIdentity) -> bool: + if type(parameter) not in {torch.Tensor, nn.Parameter} or not ( + parameter.layout is torch.strided + and parameter.device.type in {"cpu", "cuda"} + and parameter.dtype in {torch.float16, torch.bfloat16, torch.float32, torch.float64} + and not torch.is_complex(parameter) + and not parameter.is_meta + and not parameter.is_nested + and not parameter.is_quantized + and getattr(parameter, "fake_mode", None) is None + ): + return False + if shard.layout in { + ParameterLayout.REPLICATED, + ParameterLayout.WHOLE_PARAMETER_OWNER, + }: + return ( + tuple(parameter.shape) == shard.parameter.global_shape and parameter.numel() == shard.logical_slice.length + ) + if shard.layout is ParameterLayout.FLATTENED_ELEMENT_SHARD: + return parameter.ndim == 1 and parameter.is_contiguous() and parameter.numel() == shard.logical_slice.length + return False + + +def _validate_optimizer_shell( + optimizer, + implementation: str, + binding: CheckpointProcessGroupBinding, +) -> None: + _reject_method_shadows(optimizer) + _validate_exact_manifest(optimizer._gefen_sharding_manifest) + optimizer._assert_finalized_binding_layout() + if type(optimizer.defaults) is not dict: + raise TypeError("portable state requires exact built-in optimizer defaults") + if type(optimizer.param_groups) is not list: + raise TypeError("portable state requires an exact parameter-group list") + if optimizer.capturable is not False: + raise RuntimeError("portable state does not support capturable optimizers") + if optimizer._stochastic_round is not False: + raise RuntimeError("portable state does not support stochastic rounding") + if optimizer._capt_stacks is not None: + raise RuntimeError("portable state requires inactive capturable stacks") + if optimizer._gefen_global_step_by_device or optimizer._sr_seed_by_device: + raise RuntimeError("portable state does not support device-authoritative counters") + if torch.compiler.is_compiling(): + raise RuntimeError("portable state cannot run while compiling") + if torch.cuda.is_available() and torch.cuda.is_current_stream_capturing(): + raise RuntimeError("portable state cannot run during CUDA capture") + if type(optimizer._gefen_logical_slots) is not tuple or not optimizer._gefen_logical_slots: + raise ValueError("portable state requires immutable logical slots") + if type(optimizer._gefen_local_shard_bindings) is not tuple: + raise TypeError("portable state requires immutable local shard bindings") + + scope = optimizer._gefen_codebook_process_group + if type(scope) is not CodebookProcessGroupBinding: + raise RuntimeError("portable state requires an explicit codebook process group") + _validate_exact_process_group_identity(scope.identity) + if type(scope.local_member) is not str or type(scope.collective_device) is not torch.device: + raise TypeError("portable codebook bindings require exact primitive fields") + if scope.identity != binding.identity or scope.local_member != binding.local_member: + raise ValueError("checkpoint and codebook process-group identities must match") + optimizer._validate_codebook_runtime_binding(scope) + for shard in optimizer._gefen_sharding_manifest.shards: + if shard.process_group != binding.identity: + raise ValueError("portable manifest process groups must match the checkpoint binding") + if implementation == _PLAIN_IMPLEMENTATION: + if shard.layout not in { + ParameterLayout.REPLICATED, + ParameterLayout.FLATTENED_ELEMENT_SHARD, + }: + raise ValueError("plain portable state has an unsupported layout") + elif shard.layout not in { + ParameterLayout.REPLICATED, + ParameterLayout.WHOLE_PARAMETER_OWNER, + }: + raise ValueError("Muon portable state has an unsupported layout") + for slot in optimizer._gefen_logical_slots: + if type(slot) is not LogicalSlotBinding: + raise TypeError("portable logical slots must be exact LogicalSlotBinding values") + _validate_exact_shard_identity(slot.shard) + if slot.shard.local_member != binding.local_member: + raise ValueError("local logical slots must match the checkpoint member") + + +def _live_maps(optimizer, binding: CheckpointProcessGroupBinding): + local_by_fqn = {} + for item in optimizer._gefen_local_shard_bindings: + if type(item) is not tuple or len(item) != 2: + raise TypeError("portable local bindings must be parameter/shard tuples") + parameter, shard = item + _validate_exact_shard_identity(shard) + if shard.process_group != binding.identity: + raise ValueError("portable local shard process groups must exactly match the checkpoint binding") + if shard.local_member != binding.local_member: + raise ValueError("portable local shards must match the checkpoint member") + if shard.parameter.fqn in local_by_fqn: + raise ValueError("portable local bindings contain duplicate FQNs") + if parameter is not None and not _parameter_supported(parameter, shard): + raise TypeError("portable state requires supported live shard geometry and dtype") + local_by_fqn[shard.parameter.fqn] = (parameter, shard) + + expected_live = { + parameter + for parameter, shard in local_by_fqn.values() + if not (shard.layout is ParameterLayout.WHOLE_PARAMETER_OWNER and shard.local_member != shard.owner) + } + if None in expected_live: + raise TypeError("portable state requires plain materialized local parameters") + state_type = type(optimizer.state) + if not (state_type is dict or (state_type is defaultdict and optimizer.state.default_factory is dict)): + raise TypeError("portable optimizer state must use a standard mapping") + if set(optimizer.state) != expected_live: + raise ValueError("portable optimizer state contains missing or foreign parameter keys") + if type(optimizer._param_names) is not dict or set(optimizer._param_names) != expected_live: + raise ValueError("portable parameter-name state does not match live parameters") + return local_by_fqn + + +def _parameter_storage_token(parameter): + try: + version = parameter._version + except RuntimeError: + version = None + return ( + id(parameter), + version, + str(parameter.device), + str(parameter.dtype), + str(parameter.layout), + tuple(parameter.shape), + tuple(parameter.stride()), + parameter.storage_offset(), + parameter.numel(), + parameter.is_contiguous(), + parameter.untyped_storage().data_ptr(), + parameter.untyped_storage().nbytes(), + ) + + +def _portable_value_token(value): + if type(value) is torch.Tensor: + hasher = hashlib.sha256() + element_size = value.element_size() + chunk_elements = max(1, (1 << 20) // max(1, element_size)) + for start in range(0, value.numel(), chunk_elements): + stop = min(start + chunk_elements, value.numel()) + chunk = _read_flat_chunk(value, start, stop).contiguous().cpu() + hasher.update(memoryview(chunk.view(torch.uint8).numpy())) + try: + version = value._version + except RuntimeError: + version = None + return ( + "tensor", + id(value), + version, + hasher.digest(), + str(value.device), + str(value.dtype), + str(value.layout), + tuple(value.shape), + tuple(value.stride()), + value.storage_offset(), + value.requires_grad, + value.is_conj(), + value.is_neg(), + ) + if type(value) is dict: + return ( + "dict", + tuple((key, _portable_value_token(value[key])) for key in sorted(value, key=repr)), + ) + if type(value) in {list, tuple}: + return ( + type(value).__name__, + tuple(_portable_value_token(item) for item in value), + ) + try: + hash(value) + token = value + except TypeError: + token = (id(value), repr(value)) + return (type(value).__name__, token) + + +def _portable_live_token(optimizer): + _reject_method_shadows(optimizer) + + def live_group_options(group): + options = optimizer._canonical_group_options_value(group) + if type(options) is dict: + options = dict(options) + options.pop("ns_schedule", None) + return options + + groups = tuple( + ( + id(group), + tuple(id(parameter) for parameter in group["params"]), + tuple(str(name) for name, _parameter in optimizer._iter_group_params_with_names(group)), + _portable_value_token(live_group_options(group)), + ) + for group in optimizer.param_groups + ) + state_entries = tuple( + sorted( + ( + id(parameter), + type(parameter), + id(state), + _portable_value_token(state), + ) + for parameter, state in optimizer.state.items() + ) + ) + local_storage = tuple( + None if parameter is None else _parameter_storage_token(parameter) + for parameter, _shard in optimizer._gefen_local_shard_bindings + ) + return ( + id(optimizer.param_groups), + id(optimizer.state), + id(optimizer.defaults), + type(optimizer.defaults), + type(optimizer.param_groups), + type(optimizer.state), + _portable_value_token(optimizer.defaults), + groups, + state_entries, + _portable_value_token(optimizer._param_names), + local_storage, + optimizer._gefen_global_step, + _portable_value_token(optimizer._gefen_codebook), + _portable_value_token(optimizer._gefen_global_step_by_device), + _portable_value_token(optimizer._sr_seed_by_device), + _portable_value_token(optimizer._canonical_policy()), + optimizer._deterministic, + optimizer.capturable, + optimizer.fused, + optimizer.verbose, + optimizer._fused_build_ok, + id(optimizer._gefen_codebook_process_group), + _portable_value_token(optimizer._serialized_codebook_scope()), + id(optimizer._gefen_sharding_manifest), + _portable_value_token(_serialize_sharding_manifest(optimizer._gefen_sharding_manifest)), + id(optimizer._gefen_logical_slots), + tuple( + ( + id(slot), + slot.group_index, + slot.original_slot_index, + slot.compatibility_name, + _portable_value_token(_serialize_shard_identity(slot.shard)), + ) + for slot in optimizer._gefen_logical_slots + ), + tuple( + ( + None if parameter is None else id(parameter), + shard.sort_key, + ) + for parameter, shard in optimizer._gefen_local_shard_bindings + ), + ) + + +def _catalog_and_options(optimizer, implementation: str, policy): + group_options = [] + for group in optimizer.param_groups: + # Plain options depend on each logical parameter rank, so validate the + # group shell here and construct per-slot values below. + if implementation == _PLAIN_IMPLEMENTATION: + _group_options(group, implementation, second_moment_policy="block") + group_options.append(group) + else: + group_options.append(_group_options(group, implementation)) + + catalog = {} + options_by_fqn = {} + for slot in optimizer._gefen_logical_slots: + identity = slot.shard.parameter + if slot.group_index >= len(group_options): + raise ValueError("portable logical slot group index is out of range") + if implementation == _PLAIN_IMPLEMENTATION: + second = "factored" if policy["factored_v_2d"] and len(identity.global_shape) == 2 else "block" + options = _group_options( + group_options[slot.group_index], + implementation, + second_moment_policy=second, + ) + if second == "factored" and slot.shard.layout is not ParameterLayout.REPLICATED: + raise ValueError("factored portable state supports replicated matrices only") + else: + options = group_options[slot.group_index] + if slot.shard.layout is ParameterLayout.WHOLE_PARAMETER_OWNER and options["sharded_mode"] != "distributed": + raise ValueError("Muon whole-owner portable state requires distributed mode") + fqn = identity.fqn + if fqn in catalog: + raise ValueError("portable logical slots contain duplicate FQNs") + options_by_fqn[fqn] = options + catalog[fqn] = { + "identity": _serialize_parameter_identity(identity), + "algorithm_options": options, + } + return catalog, options_by_fqn + + +def _common_state(optimizer): + global_step = _strict_counter(optimizer._canonical_common_global_step(), name="gefen_global_step") + codebook = optimizer._gefen_codebook + codebook_cpu = None + if codebook is not None: + codebook_cpu = _tight_cpu_fp32( + codebook, + name="gefen_codebook", + shape=(256,), + ) + return _normalize_common( + { + "gefen_global_step": global_step, + "gefen_codebook": codebook_cpu, + "gefen_deterministic": optimizer._deterministic, + } + ) + + +def _preflight_common_state(optimizer, *, validate_values=True): + global_step = _strict_counter( + optimizer._canonical_common_global_step(), + name="gefen_global_step", + ) + codebook = optimizer._gefen_codebook + codebook_preview = None + if codebook is not None: + codebook = _plain_tensor( + codebook, + name="gefen_codebook", + dtype=torch.float32, + shape=(256,), + validate_values=validate_values, + ) + if validate_values: + if float(codebook[0].item()) != -1.0 or float(codebook[-1].item()) != 1.0: + raise ValueError("portable codebook must retain exact endpoints") + if not bool(torch.all(codebook.detach()[1:] >= codebook.detach()[:-1])): + raise ValueError("portable codebook must be sorted") + codebook_preview = _PreflightTensor((256,)) + if type(optimizer._deterministic) is not bool: + raise TypeError("gefen_deterministic must be a bool") + return { + "gefen_global_step": global_step, + "gefen_codebook": codebook_preview, + "gefen_deterministic": optimizer._deterministic, + } + + +def _parameter_state_core(optimizer, parameter, compatibility_name: str): + state = optimizer.state[parameter] + if type(state) is not dict: + raise TypeError("portable per-parameter state must be a plain dictionary") + if any(key not in _AUTHORITATIVE_STATE_KEYS and key not in _IGNORED_DERIVED_STATE_KEYS for key in state): + raise ValueError("portable parameter state contains an unknown key") + core = {key: value for key, value in state.items() if key not in _IGNORED_DERIVED_STATE_KEYS} + if core.get("name") != compatibility_name: + raise ValueError("portable parameter state name does not match its logical slot") + return core + + +def _decode_local_momentum(core, shard: ShardIdentity, codebook): + local_shape = _local_dense_shape(shard) + if local_shape is None: + raise ValueError("a non-payload shard cannot carry initialized momentum") + indices = _plain_tensor(core["m_codebook"], name="m_codebook", dtype=torch.uint8) + magnitudes = _plain_tensor( + core["m_magnitude"], + name="m_magnitude", + dtype=torch.float32, + nonnegative=True, + ) + if indices.device != magnitudes.device: + raise ValueError("momentum indices and magnitudes must share a device") + local_codebook = codebook.detach().to(device=indices.device).contiguous() + dense = _decode_quantized_momentum( + local_codebook, + indices, + magnitudes, + logical_shape=local_shape, + period=core["automatic_period"], + step=core["step"], + ) + return _tight_cpu_fp32(dense, name="dense momentum", shape=local_shape) + + +def _slot_record( + optimizer, + implementation: str, + slot: LogicalSlotBinding, + parameter, + shard: ShardIdentity, + options, + common, +): + role = _derived_role(shard) + record = { + "group_index": slot.group_index, + "original_slot_index": slot.original_slot_index, + "compatibility_name": slot.compatibility_name, + "shard": _serialize_shard_identity(shard), + "algorithm_options": options, + "role": role, + "source_period": None, + "source_second_moment": None, + "state_variant": "pristine", + "state": {}, + } + payload_role = role in {"live", "whole_owner"} and shard.parameter.numel > 0 + if not payload_role: + if ( + parameter is not None + and shard.layout is ParameterLayout.WHOLE_PARAMETER_OWNER + and shard.local_member != shard.owner + ): + raise ValueError("whole-owner nonowners must not retain parameter storage") + if parameter is not None: + core = _parameter_state_core(optimizer, parameter, slot.compatibility_name) + if frozenset(core) != {"name"}: + raise ValueError("empty portable shards must remain pristine") + return record + if parameter is None: + raise ValueError("a portable payload shard requires local parameter storage") + + core = _parameter_state_core(optimizer, parameter, slot.compatibility_name) + keys = frozenset(core) + if keys == {"name"}: + return record + if keys == {"name", "automatic_period"}: + if type(core["automatic_period"]) is not int or core["automatic_period"] != 1: + raise ValueError("exact portable period-selected state requires period one") + record["source_period"] = 1 + record["state_variant"] = "period_selected" + return record + if common["gefen_codebook"] is None: + raise ValueError("initialized portable state requires a codebook") + if type(core.get("automatic_period")) is not int or core["automatic_period"] != 1: + raise ValueError("exact portable initialized state requires period one") + step = _strict_counter(core.get("step"), name="step", minimum=1) + if step > common["gefen_global_step"]: + raise ValueError("parameter step exceeds optimizer global step") + dense_momentum = _decode_local_momentum( + core, + shard, + optimizer._gefen_codebook, + ) + record["source_period"] = 1 + + if implementation == _PLAIN_IMPLEMENTATION: + if options["second_moment_policy"] == "factored": + if keys != _PLAIN_FACTORED_KEYS: + raise ValueError("plain factored state is incomplete or contains inactive fields") + if shard.layout is not ParameterLayout.REPLICATED: + raise ValueError("factored state requires a replicated shard") + factored_step = _strict_counter(core["factored_step"], name="factored_step", minimum=1) + if factored_step > step: + raise ValueError("factored_step exceeds parameter step") + record.update( + { + "source_second_moment": "factored", + "state_variant": "initialized_factored", + "state": { + "step": step, + "momentum": dense_momentum, + "v_row": _tight_cpu_fp32( + core["v_row"], + name="v_row", + shape=(shard.parameter.global_shape[0],), + nonnegative=True, + ), + "v_col": _tight_cpu_fp32( + core["v_col"], + name="v_col", + shape=(shard.parameter.global_shape[1],), + nonnegative=True, + ), + "factored_step": factored_step, + }, + } + ) + return record + + if keys != _PLAIN_BLOCK_KEYS: + raise ValueError("plain block state is incomplete or contains inactive fields") + vmean_step = _strict_counter(core["vmean_step"], name="vmean_step", minimum=1) + if vmean_step > step: + raise ValueError("vmean_step exceeds parameter step") + local_shape = _local_dense_shape(shard) + vmean = _plain_tensor( + core["vmean"], + name="vmean", + dtype=torch.float32, + nonnegative=True, + ) + second = _expand_block_second_moment( + vmean, + logical_shape=local_shape, + period=1, + step=vmean_step, + ) + record.update( + { + "source_second_moment": "block", + "state_variant": "initialized_dense", + "state": { + "step": step, + "momentum": dense_momentum, + "second_moment": _tight_cpu_fp32( + second, + name="dense second moment", + shape=local_shape, + nonnegative=True, + ), + "second_moment_step": vmean_step, + }, + } + ) + return record + + expected = _NORMUON_KEYS if options["normuon"] else _MUON_KEYS + if keys != expected: + raise ValueError("Muon state is incomplete or conflicts with its NorMuon policy") + state = {"step": step, "momentum": dense_momentum} + variant = "initialized_dense" + if options["normuon"]: + normuon_step = _strict_counter(core["normuon_step"], name="normuon_step", minimum=1) + if normuon_step > step: + raise ValueError("normuon_step exceeds parameter step") + state.update( + { + "normuon_v": _tight_cpu_fp32( + core["normuon_v"], + name="normuon_v", + shape=(shard.parameter.global_shape[0], 1), + nonnegative=True, + ), + "normuon_step": normuon_step, + } + ) + variant = "initialized_dense_normuon" + record["state_variant"] = variant + record["state"] = state + return record + + +def _prepare_local_structure( + optimizer, + implementation: str, + binding: CheckpointProcessGroupBinding, + limits: PortableStateLimits, + *, + include_payload: bool, +): + for name, value in ( + ("logical slots", optimizer._gefen_logical_slots), + ("manifest shards", optimizer._gefen_sharding_manifest.shards), + ("parameter groups", optimizer.param_groups), + ): + if type(value) not in {list, tuple}: + continue + if len(value) > limits.max_container_items or len(value) > limits.max_tree_nodes: + raise ValueError("portable {} exceed structural limits".format(name)) + for group in optimizer.param_groups: + if type(group) is not dict: + continue + for key in ("params", "param_names"): + value = group.get(key) + if type(value) is list and (len(value) > limits.max_container_items or len(value) > limits.max_tree_nodes): + raise ValueError("portable group slots exceed structural limits") + _validate_optimizer_shell(optimizer, implementation, binding) + local_by_fqn = _live_maps(optimizer, binding) + policy = _live_policy(optimizer, implementation) + catalog, options_by_fqn = _catalog_and_options(optimizer, implementation, policy) + common_preview = _preflight_common_state( + optimizer, + validate_values=False, + ) + preview_slots = [] + for logical_slot in optimizer._gefen_logical_slots: + fqn = logical_slot.shard.parameter.fqn + parameter, shard = local_by_fqn[fqn] + if shard != logical_slot.shard: + raise ValueError("logical-slot and local shard identities disagree") + preview_slots.append( + _preflight_slot_record( + optimizer, + implementation, + logical_slot, + parameter, + shard, + options_by_fqn[fqn], + common_preview, + ) + ) + manifest = _serialize_sharding_manifest(optimizer._gefen_sharding_manifest) + target_slots = [ + { + "fqn": slot["shard"]["parameter"]["fqn"], + "group_index": slot["group_index"], + "original_slot_index": slot["original_slot_index"], + "compatibility_name": slot["compatibility_name"], + } + for slot in preview_slots + ] + if include_payload: + bounded_value = { + "format": "gefen.portable_state_fragment", + "format_version": 1, + "coverage": "local_logical_optimizer_fragment", + "implementation": implementation, + "member": binding.local_member, + "policy": policy, + "common": common_preview, + "manifest": manifest, + "catalog": catalog, + "logical_slots": preview_slots, + } + else: + bounded_value = { + "policy": policy, + "manifest": manifest, + "catalog": catalog, + "logical_slots": target_slots, + } + _preflight_portable_value(bounded_value, limits) + return { + "local_by_fqn": local_by_fqn, + "policy": policy, + "catalog": catalog, + "options_by_fqn": options_by_fqn, + "target_descriptor": { + "policy": policy, + "manifest": manifest, + "catalog": catalog, + "logical_slots": target_slots, + }, + } + + +def _validate_prepared_local_state(optimizer, implementation: str, prepared) -> None: + common = _preflight_common_state(optimizer) + local_by_fqn = prepared["local_by_fqn"] + options_by_fqn = prepared["options_by_fqn"] + for slot in optimizer._gefen_logical_slots: + fqn = slot.shard.parameter.fqn + parameter, shard = local_by_fqn[fqn] + if shard != slot.shard: + raise ValueError("logical-slot and local shard identities disagree") + _validate_readiness_slot( + optimizer, + implementation, + slot, + parameter, + shard, + options_by_fqn[fqn], + common, + ) + + +def _materialize_local_fragment( + optimizer, + implementation: str, + binding: CheckpointProcessGroupBinding, + limits: PortableStateLimits, + prepared, +): + local_by_fqn = prepared["local_by_fqn"] + policy = prepared["policy"] + catalog = prepared["catalog"] + options_by_fqn = prepared["options_by_fqn"] + common = _common_state(optimizer) + slots = [] + for logical_slot in optimizer._gefen_logical_slots: + fqn = logical_slot.shard.parameter.fqn + parameter, shard = local_by_fqn[fqn] + if shard != logical_slot.shard: + raise ValueError("logical-slot and local shard identities disagree") + slots.append( + _slot_record( + optimizer, + implementation, + logical_slot, + parameter, + shard, + options_by_fqn[fqn], + common, + ) + ) + return _build_portable_state_fragment( + implementation=implementation, + member=binding.local_member, + policy=policy, + common=common, + manifest=optimizer._gefen_sharding_manifest, + catalog=catalog, + logical_slots=slots, + limits=limits, + ) + + +def _target_context(base, fragment): + return { + **base, + "target_policy": fragment["policy"], + "target_manifest": fragment["manifest"], + "target_catalog": fragment["catalog"], + "target_logical_slots": fragment["logical_slots"], + } + + +def _export_portable_state( + optimizer, + *, + checkpoint_process_group, + transaction_id, + limits, +): + """Collectively export one complete exact portable v3 optimizer document.""" + + transport_binding = _preflight_transport_binding( + optimizer, + checkpoint_process_group, + ) + binding = None + normalized_limits = None + normalized_transaction = None + implementation = None + context = bytes(32) + wire_limits = _STATUS_FALLBACK_LIMITS + local_fragment = None + live_token = None + error = None + try: + _validate_supplied_binding(checkpoint_process_group, transport_binding) + binding = transport_binding + normalized_limits = _require_limits(limits) + wire_limits = normalized_limits._wire_limits() + normalized_transaction = _require_transaction_id(transaction_id) + implementation = _optimizer_implementation(optimizer) + _validate_context_identity_bounds(binding, normalized_limits) + base_context = _base_context(binding, implementation) + _preflight_portable_value(base_context, normalized_limits) + context = _context_digest(base_context) + _prepare_local_structure( + optimizer, + implementation, + binding, + normalized_limits, + include_payload=True, + ) + live_token = _portable_live_token(optimizer) + prepared_local = _prepare_local_structure( + optimizer, + implementation, + binding, + normalized_limits, + include_payload=True, + ) + _validate_prepared_local_state( + optimizer, + implementation, + prepared_local, + ) + local_fragment = _materialize_local_fragment( + optimizer, + implementation, + binding, + normalized_limits, + prepared_local, + ) + if live_token != _portable_live_token(optimizer): + raise RuntimeError("live optimizer state changed during portable export preparation") + except Exception as exc: + error = exc + _collective_unanimous_status( + transport_binding, + error, + operation="portable_export_prepare", + transaction_id=_EXPORT_PREFLIGHT_TRANSACTION, + context_digest=context, + limits=wire_limits, + ) + assert ( + normalized_limits is not None + and normalized_transaction is not None + and implementation is not None + and binding is not None + and local_fragment is not None + ) + + fragments = [] + _collective_visit_canonical_fragments( + binding, + local_fragment, + operation="portable_export_fragments", + transaction_id=normalized_transaction, + context_digest=context, + limits=wire_limits, + consume=lambda _member, value: fragments.append( + _normalize_portable_state_fragment( + value, + limits=normalized_limits, + ) + ), + ) + + document = None + error = None + try: + document = _assemble_portable_state_fragments( + fragments, + process_group_identity=binding.identity, + limits=normalized_limits, + ) + _validate_live_readiness(optimizer, implementation, binding) + if live_token != _portable_live_token(optimizer): + raise RuntimeError("live optimizer state changed during portable export") + except Exception as exc: + error = exc + _collective_unanimous_status( + binding, + error, + operation="portable_export_finalize", + transaction_id=normalized_transaction, + context_digest=context, + limits=wire_limits, + ) + assert document is not None + return document + + +def _stage_portable_import( + optimizer, + implementation: str, + binding: CheckpointProcessGroupBinding, + limits: PortableStateLimits, + document, +): + _prepare_local_structure( + optimizer, + implementation, + binding, + limits, + include_payload=False, + ) + live_token = _portable_live_token(optimizer) + prepared_local = _prepare_local_structure( + optimizer, + implementation, + binding, + limits, + include_payload=False, + ) + _validate_prepared_local_state( + optimizer, + implementation, + prepared_local, + ) + target_fragment = prepared_local["target_descriptor"] + if not _values_equal(document["policy"], target_fragment["policy"]): + raise ValueError("portable document policy does not match the target") + common = document["common"] + if set(document["parameters"]) != set(target_fragment["catalog"]): + raise ValueError("portable parameter catalog does not match the target") + + local_by_fqn = _live_maps(optimizer, binding) + canonical_parameters = {} + for fqn, (parameter, shard) in local_by_fqn.items(): + if parameter is None: + continue + target_options = target_fragment["catalog"][fqn]["algorithm_options"] + target_second = target_options["second_moment_policy"] if implementation == _PLAIN_IMPLEMENTATION else None + projected = _project_portable_parameter_state( + document["parameters"][fqn], + shard, + implementation=implementation, + global_step=common["gefen_global_step"], + codebook=common["gefen_codebook"], + target_algorithm_options=target_options, + target_second_moment=target_second, + ) + canonical_parameters[fqn] = {"state": projected} + + canonical_state = { + "common": { + "gefen_global_step": common["gefen_global_step"], + "gefen_codebook": common["gefen_codebook"], + "gefen_deterministic": common["gefen_deterministic"], + "gefen_codebook_scope": optimizer._serialized_codebook_scope(), + }, + "parameters": canonical_parameters, + } + native = optimizer._canonical_native_state_dict(canonical_state) + # Portable state intentionally carries the source replica-determinism + # setting. Native loading ordinarily rejects a policy change, so perform + # the same complete staging path through an isolated owner configured with + # the document value. The live optimizer remains untouched until commit. + staging_owner = object.__new__(type(optimizer)) + staging_owner.__dict__ = optimizer.__dict__.copy() + staging_owner._deterministic = common["gefen_deterministic"] + staged = staging_owner._stage_load_state_dict(native) + optimizer._preserve_canonical_target_configuration(staged) + if ( + type(staged.__dict__) is not dict + or type(staged.defaults) is not dict + or set(staged.defaults) != set(optimizer.defaults) + or staged.param_groups is not optimizer.param_groups + ): + raise TypeError("portable import staging produced unsafe publication containers") + if live_token != _portable_live_token(optimizer): + raise RuntimeError("live optimizer state changed during portable import preparation") + return staged, live_token, target_fragment + + +def _import_portable_state( + optimizer, + state, + *, + checkpoint_process_group, + transaction_id, + limits, +) -> None: + """Collectively stage, vote, and atomically publish portable v3 state.""" + + transport_binding = _preflight_transport_binding( + optimizer, + checkpoint_process_group, + ) + binding = None + normalized_limits = None + normalized_transaction = None + implementation = None + base = None + base_digest = bytes(32) + wire_limits = _STATUS_FALLBACK_LIMITS + document = None + error = None + try: + _validate_supplied_binding(checkpoint_process_group, transport_binding) + binding = transport_binding + normalized_limits = _require_limits(limits) + wire_limits = normalized_limits._wire_limits() + normalized_transaction = _require_transaction_id(transaction_id) + implementation = _optimizer_implementation(optimizer) + _validate_context_identity_bounds(binding, normalized_limits) + base = _base_context(binding, implementation) + _preflight_portable_value(base, normalized_limits) + base_digest = _context_digest(base) + document = _normalize_gefen_portable_state_document( + state, + limits=normalized_limits, + expected_implementation=implementation, + ) + except Exception as exc: + error = exc + _collective_unanimous_status( + transport_binding, + error, + operation="portable_import_document", + transaction_id=_IMPORT_PREFLIGHT_TRANSACTION, + context_digest=base_digest, + limits=wire_limits, + ) + assert ( + normalized_limits is not None + and normalized_transaction is not None + and implementation is not None + and binding is not None + and base is not None + and document is not None + ) + + staged = None + live_token = None + target_fragment = None + error = None + try: + staged, live_token, target_fragment = _stage_portable_import( + optimizer, + implementation, + binding, + normalized_limits, + document, + ) + target_context = _target_context(base, target_fragment) + target_context["document_digest"] = document["completion"]["digest"] + _preflight_portable_value(target_context, normalized_limits) + import_digest = _context_digest(target_context) + except Exception as exc: + error = exc + import_digest = _context_digest({**base, "document_digest": document["completion"]["digest"]}) + _collective_unanimous_status( + binding, + error, + operation="portable_import_prepare", + transaction_id=normalized_transaction, + context_digest=import_digest, + limits=wire_limits, + ) + assert staged is not None and live_token is not None and target_fragment is not None + + error = None + try: + _validate_live_readiness(optimizer, implementation, binding) + if live_token != _portable_live_token(optimizer): + raise RuntimeError("live optimizer state changed after portable import preparation") + except Exception as exc: + error = exc + _collective_unanimous_status( + binding, + error, + operation="portable_import_freshness", + transaction_id=normalized_transaction, + context_digest=import_digest, + limits=wire_limits, + ) + + # This is the established non-throwing publication primitive. No collective + # or semantic validation is allowed after this point. + from gefen.gefen import Gefen + + Gefen._commit_staged_load_state_dict(optimizer, staged) + + +def _validate_readiness_slot( + optimizer, + implementation: str, + slot: LogicalSlotBinding, + parameter, + shard: ShardIdentity, + options, + common, + *, + validate_values=True, +) -> None: + role = _derived_role(shard) + payload_role = role in {"live", "whole_owner"} and shard.parameter.numel > 0 + if not payload_role: + if parameter is None: + if shard.layout is not ParameterLayout.WHOLE_PARAMETER_OWNER: + raise ValueError("only whole-parameter nonowners may omit storage") + return + core = _parameter_state_core(optimizer, parameter, slot.compatibility_name) + if frozenset(core) != {"name"}: + raise ValueError("empty portable shards must remain pristine") + return + if parameter is None: + raise ValueError("a portable payload shard requires local storage") + + core = _parameter_state_core(optimizer, parameter, slot.compatibility_name) + keys = frozenset(core) + if keys == {"name"}: + return + if keys == {"name", "automatic_period"}: + if type(core["automatic_period"]) is not int or core["automatic_period"] != 1: + raise ValueError("exact portable period-selected state requires period one") + if common["gefen_codebook"] is None: + raise ValueError("period-selected portable state requires a codebook") + return + if common["gefen_codebook"] is None: + raise ValueError("initialized portable state requires a codebook") + if type(core.get("automatic_period")) is not int or core["automatic_period"] != 1: + raise ValueError("exact portable initialized state requires period one") + step = _strict_counter(core.get("step"), name="step", minimum=1) + if step > common["gefen_global_step"]: + raise ValueError("parameter step exceeds optimizer global step") + local_shape = _local_dense_shape(shard) + if local_shape is None: + raise ValueError("initialized portable state requires a local payload") + blocks = math.prod(local_shape) + indices = _plain_tensor( + core.get("m_codebook"), + name="m_codebook", + dtype=torch.uint8, + shape=(blocks, 1), + validate_values=validate_values, + ) + magnitudes = _plain_tensor( + core.get("m_magnitude"), + name="m_magnitude", + dtype=torch.float32, + shape=(blocks, 1), + nonnegative=True, + validate_values=validate_values, + ) + if indices.device != magnitudes.device: + raise ValueError("momentum indices and magnitudes must share a device") + if indices.device != parameter.device: + raise ValueError("portable parameter state must share its parameter device") + + if implementation == _PLAIN_IMPLEMENTATION: + if options["second_moment_policy"] == "factored": + if keys != _PLAIN_FACTORED_KEYS: + raise ValueError("plain factored state is incomplete or contains inactive fields") + factored_step = _strict_counter(core["factored_step"], name="factored_step", minimum=1) + if factored_step > step: + raise ValueError("factored_step exceeds parameter step") + rows, columns = shard.parameter.global_shape + v_row = _plain_tensor( + core["v_row"], + name="v_row", + dtype=torch.float32, + shape=(rows,), + nonnegative=True, + validate_values=validate_values, + ) + v_col = _plain_tensor( + core["v_col"], + name="v_col", + dtype=torch.float32, + shape=(columns,), + nonnegative=True, + validate_values=validate_values, + ) + if v_row.device != parameter.device or v_col.device != parameter.device: + raise ValueError("factored portable state must share its parameter device") + return + if keys != _PLAIN_BLOCK_KEYS: + raise ValueError("plain block state is incomplete or contains inactive fields") + vmean_step = _strict_counter(core["vmean_step"], name="vmean_step", minimum=1) + if vmean_step > step: + raise ValueError("vmean_step exceeds parameter step") + vmean = _plain_tensor( + core["vmean"], + name="vmean", + dtype=torch.float32, + shape=(blocks, 1), + nonnegative=True, + validate_values=validate_values, + ) + if vmean.device != parameter.device: + raise ValueError("block portable state must share its parameter device") + return + + expected = _NORMUON_KEYS if options["normuon"] else _MUON_KEYS + if keys != expected: + raise ValueError("Muon state is incomplete or conflicts with its NorMuon policy") + if options["normuon"]: + normuon_step = _strict_counter(core["normuon_step"], name="normuon_step", minimum=1) + if normuon_step > step: + raise ValueError("normuon_step exceeds parameter step") + normuon_v = _plain_tensor( + core["normuon_v"], + name="normuon_v", + dtype=torch.float32, + shape=(shard.parameter.global_shape[0], 1), + nonnegative=True, + validate_values=validate_values, + ) + if normuon_v.device != parameter.device: + raise ValueError("NorMuon portable state must share its parameter device") + + +def _preflight_slot_record( + optimizer, + implementation: str, + slot: LogicalSlotBinding, + parameter, + shard: ShardIdentity, + options, + common, +): + _validate_readiness_slot( + optimizer, + implementation, + slot, + parameter, + shard, + options, + common, + validate_values=False, + ) + role = _derived_role(shard) + record = { + "group_index": slot.group_index, + "original_slot_index": slot.original_slot_index, + "compatibility_name": slot.compatibility_name, + "shard": _serialize_shard_identity(shard), + "algorithm_options": options, + "role": role, + "source_period": None, + "source_second_moment": None, + "state_variant": "pristine", + "state": {}, + } + payload_role = role in {"live", "whole_owner"} and shard.parameter.numel > 0 + if not payload_role: + return record + core = _parameter_state_core(optimizer, parameter, slot.compatibility_name) + keys = frozenset(core) + if keys == {"name"}: + return record + record["source_period"] = 1 + if keys == {"name", "automatic_period"}: + record["state_variant"] = "period_selected" + return record + + local_shape = _local_dense_shape(shard) + state = { + "step": core["step"], + "momentum": _PreflightTensor(local_shape), + } + if implementation == _PLAIN_IMPLEMENTATION: + if options["second_moment_policy"] == "factored": + rows, columns = shard.parameter.global_shape + record.update( + { + "source_second_moment": "factored", + "state_variant": "initialized_factored", + "state": { + **state, + "v_row": _PreflightTensor((rows,)), + "v_col": _PreflightTensor((columns,)), + "factored_step": core["factored_step"], + }, + } + ) + return record + record.update( + { + "source_second_moment": "block", + "state_variant": "initialized_dense", + "state": { + **state, + "second_moment": _PreflightTensor(local_shape), + "second_moment_step": core["vmean_step"], + }, + } + ) + return record + + if options["normuon"]: + state.update( + { + "normuon_v": _PreflightTensor((shard.parameter.global_shape[0], 1)), + "normuon_step": core["normuon_step"], + } + ) + record["state_variant"] = "initialized_dense_normuon" + else: + record["state_variant"] = "initialized_dense" + record["state"] = state + return record + + +def _validate_live_readiness( + optimizer, + implementation: str, + binding: CheckpointProcessGroupBinding, +): + _validate_optimizer_shell(optimizer, implementation, binding) + local_by_fqn = _live_maps(optimizer, binding) + policy = _live_policy(optimizer, implementation) + _, options_by_fqn = _catalog_and_options(optimizer, implementation, policy) + common = _preflight_common_state(optimizer) + for slot in optimizer._gefen_logical_slots: + parameter, shard = local_by_fqn[slot.shard.parameter.fqn] + if shard != slot.shard: + raise ValueError("logical-slot and local shard identities disagree") + _validate_readiness_slot( + optimizer, + implementation, + slot, + parameter, + shard, + options_by_fqn[slot.shard.parameter.fqn], + common, + ) + return frozenset(slot.shard.layout for slot in optimizer._gefen_logical_slots) + + +def _portable_runtime_layouts(optimizer): + """Return dynamically eligible layouts without executing a collective.""" + + try: + implementation = _optimizer_implementation(optimizer) + scope = optimizer._gefen_codebook_process_group + if type(scope) is not CodebookProcessGroupBinding: + return frozenset() + binding = CheckpointProcessGroupBinding( + scope.identity, + scope.local_member, + scope.process_group, + scope.collective_device, + ) + return _validate_live_readiness(optimizer, implementation, binding) + except Exception: + return frozenset() + + +__all__ = [] diff --git a/tests/test_portable_runtime.py b/tests/test_portable_runtime.py new file mode 100644 index 0000000..df657bf --- /dev/null +++ b/tests/test_portable_runtime.py @@ -0,0 +1,400 @@ +"""Warning-strict CPU coverage for the optimizer-facing portable runtime.""" + +import copy + +import pytest +import torch + +from gefen.checkpoint import CheckpointProcessGroupBinding +from gefen.codebook import CodebookProcessGroupBinding +from gefen.contracts import ( + CheckpointTransport, + LogicalSlice, + ParameterIdentity, + ParameterLayout, + PlacementKind, + PortableStateProvider, + ProcessGroupScope, + ProcessGroupIdentity, + ShardIdentity, + ShardPlacement, + ShardingManifest, + TopologyChange, +) +from gefen.gefen import Gefen +from gefen.gefen_muon import GefenMuon +from gefen.portable import _decode_quantized_momentum +from gefen.portable_runtime import ( + _export_portable_state, + _import_portable_state, + _optimizer_implementation, + _portable_runtime_layouts, +) +from gefen.portable_state import PortableStateLimits +from gefen.rebinding import ParameterRebinding + + +_MEMBER = "rank:0" + + +def _limits(): + return PortableStateLimits( + max_fragment_tensor_bytes=4 << 20, + max_collective_tensor_bytes=16 << 20, + max_collective_metadata_bytes=64 << 20, + ) + + +def _bindings(group): + return ( + CodebookProcessGroupBinding(group, _MEMBER, None, torch.device("cpu")), + CheckpointProcessGroupBinding(group, _MEMBER, None, torch.device("cpu")), + ) + + +def _finalize(optimizer, parameter, *, layout): + group = ProcessGroupIdentity("checkpoint", (_MEMBER,)) + identity = ParameterIdentity("layer.weight", tuple(parameter.shape)) + if layout is ParameterLayout.REPLICATED: + kind = PlacementKind.REPLICATE + owner = None + else: + kind = PlacementKind.WHOLE_PARAMETER_OWNER + owner = _MEMBER + shard = ShardIdentity( + identity, + layout, + LogicalSlice.full(identity), + placements=(ShardPlacement("checkpoint", kind, 0, 1),), + process_group=group, + local_member=_MEMBER, + owner=owner, + ) + codebook_binding, checkpoint_binding = _bindings(group) + optimizer.post_sharding( + (ParameterRebinding(parameter, parameter, shard),), + manifest=ShardingManifest((shard,)), + codebook_process_group=codebook_binding, + ) + return checkpoint_binding + + +def _plain(*, factored, deterministic): + parameter = torch.nn.Parameter(torch.arange(1, 7, dtype=torch.float32).reshape(2, 3)) + optimizer = Gefen( + [("weight", parameter)], + fused=False, + factored_v_2d=factored, + force_2d_period_one=True, + deterministic=deterministic, + ) + binding = _finalize(optimizer, parameter, layout=ParameterLayout.REPLICATED) + return optimizer, parameter, binding + + +def _muon(*, deterministic): + parameter = torch.nn.Parameter(torch.arange(1, 7, dtype=torch.float32).reshape(2, 3)) + optimizer = GefenMuon( + [("weight", parameter)], + fused=False, + sharded_mode="distributed", + normuon=True, + deterministic=deterministic, + ) + binding = _finalize( + optimizer, + parameter, + layout=ParameterLayout.WHOLE_PARAMETER_OWNER, + ) + return optimizer, parameter, binding + + +def _initialize(optimizer, parameter, *, variant): + optimizer._gefen_global_step = 4 + optimizer._gefen_codebook = torch.linspace(-1.0, 1.0, 256, dtype=torch.float32) + state = optimizer.state[parameter] + state.update( + { + "automatic_period": 1, + "step": 4, + "m_codebook": torch.tensor([[0], [255], [128], [64], [192], [0]], dtype=torch.uint8), + "m_magnitude": torch.tensor([[1.0], [2.0], [3.0], [4.0], [5.0], [6.0]], dtype=torch.float32), + } + ) + if variant == "block": + state.update( + { + "vmean": torch.arange(1, 7, dtype=torch.float32).reshape(6, 1), + "vmean_step": 3, + } + ) + elif variant == "factored": + state.update( + { + "v_row": torch.tensor([2.0, 3.0], dtype=torch.float32), + "v_col": torch.tensor([5.0, 7.0, 11.0], dtype=torch.float32), + "factored_step": 3, + } + ) + elif variant == "normuon": + state.update( + { + "normuon_v": torch.tensor([[2.0], [3.0]], dtype=torch.float32), + "normuon_step": 2, + } + ) + else: + raise AssertionError("unknown test variant") + + +@pytest.mark.parametrize("variant", ["block", "factored", "normuon"]) +def test_initialized_singleton_export_import_round_trip(variant): + if variant == "normuon": + source, source_parameter, source_binding = _muon(deterministic=True) + target, target_parameter, target_binding = _muon(deterministic=False) + else: + source, source_parameter, source_binding = _plain( + factored=variant == "factored", + deterministic=True, + ) + target, target_parameter, target_binding = _plain( + factored=variant == "factored", + deterministic=False, + ) + _initialize(source, source_parameter, variant=variant) + assert isinstance(source, PortableStateProvider) + support = next( + item + for item in source.optimizer_contract().capabilities.checkpoints + if item.transport is CheckpointTransport.CANONICAL_GLOBAL + ) + assert support.requires_collective + assert support.atomic_load + assert support.process_group_scope is ProcessGroupScope.ADAPTER_DEFINED + expected_layout = ParameterLayout.WHOLE_PARAMETER_OWNER if variant == "normuon" else ParameterLayout.REPLICATED + assert expected_layout in support.same_topology + if variant == "factored": + assert not support.topology_changing + elif variant == "block": + assert support.topology_changing == frozenset( + { + ParameterLayout.REPLICATED, + ParameterLayout.FLATTENED_ELEMENT_SHARD, + } + ) + assert support.topology_change_kinds == frozenset({TopologyChange.PLACEMENT_RESHARD}) + else: + assert support.topology_changing == frozenset( + { + ParameterLayout.REPLICATED, + ParameterLayout.WHOLE_PARAMETER_OWNER, + } + ) + assert support.topology_change_kinds == frozenset( + { + TopologyChange.PLACEMENT_RESHARD, + TopologyChange.WORLD_SIZE_OWNER_REDISTRIBUTION, + } + ) + + document = _export_portable_state( + source, + checkpoint_process_group=source_binding, + transaction_id="export-{}".format(variant), + limits=_limits(), + ) + _import_portable_state( + target, + document, + checkpoint_process_group=target_binding, + transaction_id="import-{}".format(variant), + limits=_limits(), + ) + + assert target._gefen_global_step == 4 + assert target._deterministic is True + assert torch.equal(target._gefen_codebook, document["common"]["gefen_codebook"]) + target_state = target.state[target_parameter] + target_momentum = _decode_quantized_momentum( + target._gefen_codebook, + target_state["m_codebook"], + target_state["m_magnitude"], + logical_shape=tuple(target_parameter.shape), + period=1, + step=target_state["step"], + ) + record = document["parameters"]["layer.weight"] + assert torch.equal(target_momentum.view(torch.int32), record["state"]["momentum"].view(torch.int32)) + if variant == "block": + assert torch.equal(target_state["vmean"].reshape(2, 3), record["state"]["second_moment"]) + assert target_state["vmean_step"] == 3 + elif variant == "factored": + assert torch.equal(target_state["v_row"], record["state"]["v_row"]) + assert torch.equal(target_state["v_col"], record["state"]["v_col"]) + assert target_state["factored_step"] == 3 + else: + assert torch.equal(target_state["normuon_v"], record["state"]["normuon_v"]) + assert target_state["normuon_step"] == 2 + + +def test_import_rejects_corrupt_document_without_live_mutation(): + source, source_parameter, source_binding = _plain(factored=False, deterministic=True) + target, _, target_binding = _plain(factored=False, deterministic=False) + _initialize(source, source_parameter, variant="block") + document = _export_portable_state( + source, + checkpoint_process_group=source_binding, + transaction_id="export-corruption-source", + limits=_limits(), + ) + corrupt = copy.deepcopy(document) + corrupt["completion"]["digest"] = "0" * 64 + before = target._canonical_import_live_token() + + with pytest.raises(RuntimeError, match="digest"): + _import_portable_state( + target, + corrupt, + checkpoint_process_group=target_binding, + transaction_id="import-corruption", + limits=_limits(), + ) + + assert target._canonical_import_live_token() == before + + +def test_export_rejects_non_period_one_and_unknown_group_state(): + optimizer, parameter, binding = _plain(factored=False, deterministic=False) + optimizer._gefen_global_step = 2 + optimizer._gefen_codebook = torch.linspace(-1.0, 1.0, 256, dtype=torch.float32) + optimizer.state[parameter].update( + { + "automatic_period": 2, + "step": 2, + "m_codebook": torch.zeros((3, 2), dtype=torch.uint8), + "m_magnitude": torch.ones((3, 1), dtype=torch.float32), + "vmean": torch.ones((3, 1), dtype=torch.float32), + "vmean_step": 2, + } + ) + with pytest.raises(RuntimeError, match="period one"): + _export_portable_state( + optimizer, + checkpoint_process_group=binding, + transaction_id="export-period-two", + limits=_limits(), + ) + + optimizer.state[parameter] = {"name": "weight"} + optimizer.param_groups[0]["extension"] = 1 + with pytest.raises(RuntimeError, match="unknown keys"): + _export_portable_state( + optimizer, + checkpoint_process_group=binding, + transaction_id="export-unknown-group-option", + limits=_limits(), + ) + + +def test_exact_type_and_full_process_group_identity_are_required(): + class GefenSubclass(Gefen): + pass + + parameter = torch.nn.Parameter(torch.ones(2, dtype=torch.float32)) + optimizer = GefenSubclass([("weight", parameter)], fused=False) + with pytest.raises(TypeError, match="exact Gefen"): + _optimizer_implementation(optimizer) + + exact, _, _ = _plain(factored=False, deterministic=False) + wrong_group = ProcessGroupIdentity("different-checkpoint", (_MEMBER,)) + _, wrong_binding = _bindings(wrong_group) + with pytest.raises(RuntimeError, match="identities must match"): + _export_portable_state( + exact, + checkpoint_process_group=wrong_binding, + transaction_id="export-wrong-process-group", + limits=_limits(), + ) + + +def test_runtime_readiness_fails_closed_on_state_group_and_subclass_mutation(): + optimizer, parameter, binding = _plain(factored=False, deterministic=False) + assert _portable_runtime_layouts(optimizer) == frozenset({ParameterLayout.REPLICATED}) + + optimizer.state[parameter]["automatic_period"] = 2 + assert not _portable_runtime_layouts(optimizer) + assert all( + support.transport is not CheckpointTransport.CANONICAL_GLOBAL + for support in optimizer.optimizer_contract().capabilities.checkpoints + ) + optimizer.state[parameter] = {"name": "weight"} + optimizer.param_groups[0]["extension"] = 1 + assert not _portable_runtime_layouts(optimizer) + optimizer.param_groups[0].pop("extension") + + _initialize(optimizer, parameter, variant="block") + optimizer.state[parameter]["m_magnitude"] = optimizer.state[parameter]["m_magnitude"].requires_grad_() + assert not _portable_runtime_layouts(optimizer) + with pytest.raises(RuntimeError, match="plain materialized"): + _export_portable_state( + optimizer, + checkpoint_process_group=binding, + transaction_id="export-grad-state", + limits=_limits(), + ) + + class GefenSubclass(Gefen): + pass + + subclass_parameter = torch.nn.Parameter(torch.ones(2)) + subclassed = GefenSubclass([("weight", subclass_parameter)], fused=False) + assert not _portable_runtime_layouts(subclassed) + + +def test_muon_readiness_rejects_nonexact_beta_and_unsafe_raw_schedule(): + optimizer, _, binding = _muon(deterministic=False) + optimizer.param_groups[0]["beta2"] = torch.tensor(0.0) + assert not _portable_runtime_layouts(optimizer) + with pytest.raises(RuntimeError, match="exact floats"): + _export_portable_state( + optimizer, + checkpoint_process_group=binding, + transaction_id="export-tensor-beta", + limits=_limits(), + ) + + optimizer.param_groups[0]["beta2"] = 0.0 + optimizer.param_groups[0]["ns_schedule"] = lambda: None + assert not _portable_runtime_layouts(optimizer) + with pytest.raises(RuntimeError, match="weights-only-safe"): + _export_portable_state( + optimizer, + checkpoint_process_group=binding, + transaction_id="export-callable-schedule", + limits=_limits(), + ) + + +def test_invalid_collective_arguments_fail_through_the_status_gate(): + optimizer, _, binding = _plain(factored=False, deterministic=False) + with pytest.raises(RuntimeError, match="PortableStateLimits"): + _export_portable_state( + optimizer, + checkpoint_process_group=binding, + transaction_id="invalid-limits", + limits=object(), + ) + with pytest.raises(RuntimeError, match="transaction_id"): + _export_portable_state( + optimizer, + checkpoint_process_group=binding, + transaction_id="", + limits=_limits(), + ) + with pytest.raises(RuntimeError, match="trimmed"): + _export_portable_state( + optimizer, + checkpoint_process_group=binding, + transaction_id=" invalid-transaction ", + limits=_limits(), + ) diff --git a/tests/test_portable_runtime_consensus.py b/tests/test_portable_runtime_consensus.py new file mode 100644 index 0000000..6af5ec0 --- /dev/null +++ b/tests/test_portable_runtime_consensus.py @@ -0,0 +1,477 @@ +"""Warning-strict Gloo regressions for portable runtime consensus gates.""" + +from datetime import timedelta +import multiprocessing as mp +import os +import queue as queue_module +import tempfile +import traceback +import warnings + +import pytest +import torch +import torch.distributed as dist + +from gefen.checkpoint import CheckpointProcessGroupBinding +from gefen.codebook import CodebookProcessGroupBinding +from gefen.contracts import ( + LogicalSlice, + ParameterIdentity, + ParameterLayout, + PlacementKind, + ProcessGroupIdentity, + ShardIdentity, + ShardPlacement, + ShardingManifest, +) +from gefen.gefen import Gefen +from gefen.portable_state import PortableStateLimits +from gefen.rebinding import ParameterRebinding + + +_WORLD = 2 +_SINGLE_FQN = "model.weight" + + +class _CheckpointBindingSubclass(CheckpointProcessGroupBinding): + pass + + +def _members(): + return tuple("rank:{}".format(rank) for rank in range(_WORLD)) + + +def _limits(): + return PortableStateLimits( + max_fragment_tensor_bytes=1 << 20, + max_collective_tensor_bytes=4 << 20, + max_collective_metadata_bytes=4 << 20, + chunk_bytes=17, + max_members=4, + max_metadata_bytes=1 << 20, + max_tree_nodes=10_000, + max_tree_depth=32, + max_container_items=10_000, + max_string_bytes=16 << 10, + max_integer_bytes=128, + max_tensors=256, + max_tensor_rank=8, + diagnostic_bytes=1024, + ) + + +def _bindings(group, rank): + member = _members()[rank] + codebook = CodebookProcessGroupBinding( + group, + member, + dist.group.WORLD, + torch.device("cpu"), + ) + checkpoint = CheckpointProcessGroupBinding( + group, + member, + dist.group.WORLD, + torch.device("cpu"), + ) + return codebook, checkpoint + + +def _replicated_shards(identity, group): + return tuple( + ShardIdentity( + identity, + ParameterLayout.REPLICATED, + LogicalSlice.full(identity), + placements=( + ShardPlacement( + "checkpoint", + PlacementKind.REPLICATE, + coordinate, + len(group.ordered_members), + ), + ), + process_group=group, + local_member=member, + ) + for coordinate, member in enumerate(group.ordered_members) + ) + + +def _new_optimizer(entries, *, deterministic): + return Gefen( + entries, + lr=2.5e-3, + betas=(0.8, 0.97), + eps=2.0e-8, + weight_decay=0.03, + fused=False, + force_2d_period_one=True, + factored_v_2d=False, + deterministic=deterministic, + ) + + +def _make_single(rank, group, *, compatibility_name, deterministic): + identity = ParameterIdentity(_SINGLE_FQN, (2, 3)) + parameter = torch.nn.Parameter(torch.zeros(identity.global_shape, dtype=torch.float32)) + optimizer = _new_optimizer( + [(compatibility_name, parameter)], + deterministic=deterministic, + ) + shards = _replicated_shards(identity, group) + codebook_binding, checkpoint_binding = _bindings(group, rank) + optimizer.post_sharding( + (ParameterRebinding(parameter, parameter, shards[rank]),), + manifest=ShardingManifest(shards), + codebook_process_group=codebook_binding, + ) + return optimizer, checkpoint_binding + + +def _make_two(rank, group, *, reverse, deterministic): + identities = ( + ParameterIdentity("model.alpha", (2, 2)), + ParameterIdentity("model.beta", (2, 2)), + ) + alpha = torch.nn.Parameter(torch.zeros(identities[0].global_shape, dtype=torch.float32)) + beta = torch.nn.Parameter(torch.zeros(identities[1].global_shape, dtype=torch.float32)) + entries = [("alpha", alpha), ("beta", beta)] + if reverse: + entries.reverse() + optimizer = _new_optimizer(entries, deterministic=deterministic) + shards_by_identity = tuple(_replicated_shards(identity, group) for identity in identities) + manifest = ShardingManifest(tuple(shard for shards in shards_by_identity for shard in shards)) + codebook_binding, checkpoint_binding = _bindings(group, rank) + optimizer.post_sharding( + ( + ParameterRebinding(alpha, alpha, shards_by_identity[0][rank]), + ParameterRebinding(beta, beta, shards_by_identity[1][rank]), + ), + manifest=manifest, + codebook_process_group=codebook_binding, + ) + return optimizer, checkpoint_binding + + +def _subclass_binding(binding): + return _CheckpointBindingSubclass( + binding.identity, + binding.local_member, + binding.process_group, + binding.collective_device, + ) + + +def _binding_with_process_group(binding, process_group): + return CheckpointProcessGroupBinding( + binding.identity, + binding.local_member, + process_group, + binding.collective_device, + ) + + +def _attempt_failure(optimizer, operation): + state_mapping = optimizer.state + param_groups = optimizer.param_groups + defaults = optimizer.defaults + parameters = tuple(parameter for group in optimizer.param_groups for parameter in group["params"]) + state_objects = tuple(optimizer.state[parameter] for parameter in parameters) + token = optimizer._canonical_import_live_token() + global_step = optimizer._gefen_global_step + codebook = optimizer._gefen_codebook + deterministic = optimizer._deterministic + try: + operation() + except RuntimeError as exc: + message = str(exc) + else: + message = None + return { + "message": message, + "token_unchanged": optimizer._canonical_import_live_token() == token, + "containers_unchanged": ( + optimizer.state is state_mapping + and optimizer.param_groups is param_groups + and optimizer.defaults is defaults + ), + "state_objects_unchanged": all( + optimizer.state[parameter] is state for parameter, state in zip(parameters, state_objects) + ), + "states_pristine": all( + optimizer.state[parameter] == {"name": optimizer._param_names[parameter]} for parameter in parameters + ), + "common_unchanged": ( + optimizer._gefen_global_step == global_step + and optimizer._gefen_codebook is codebook + and optimizer._deterministic is deterministic + ), + } + + +def _worker(rank, init_file, result_queue): + warnings.simplefilter("error") + alternate_process_group = None + try: + dist.init_process_group( + "gloo", + init_method="file://{}".format(init_file), + rank=rank, + world_size=_WORLD, + timeout=timedelta(seconds=20), + ) + group = ProcessGroupIdentity("portable_consensus", _members()) + alternate_process_group = dist.new_group( + ranks=list(range(_WORLD)), + backend="gloo", + timeout=timedelta(seconds=20), + ) + + source, exact_binding = _make_single( + rank, + group, + compatibility_name="weight", + deterministic=True, + ) + document = source.export_portable_state( + checkpoint_process_group=exact_binding, + transaction_id="consensus-healthy-single-export", + limits=_limits(), + ) + asymmetric_binding = _subclass_binding(exact_binding) if rank == 0 else exact_binding + binding_export = _attempt_failure( + source, + lambda: source.export_portable_state( + checkpoint_process_group=asymmetric_binding, + transaction_id="consensus-subclass-export", + limits=_limits(), + ), + ) + dist.barrier() + object_binding = object() if rank == 0 else exact_binding + object_export = _attempt_failure( + source, + lambda: source.export_portable_state( + checkpoint_process_group=object_binding, + transaction_id="consensus-object-export", + limits=_limits(), + ), + ) + dist.barrier() + alternate_binding = ( + _binding_with_process_group(exact_binding, alternate_process_group) if rank == 0 else exact_binding + ) + alternate_export = _attempt_failure( + source, + lambda: source.export_portable_state( + checkpoint_process_group=alternate_binding, + transaction_id="consensus-alternate-handle-export", + limits=_limits(), + ), + ) + dist.barrier() + + binding_target, target_binding = _make_single( + rank, + group, + compatibility_name="weight", + deterministic=False, + ) + asymmetric_target_binding = _subclass_binding(target_binding) if rank == 0 else target_binding + binding_import = _attempt_failure( + binding_target, + lambda: binding_target.import_portable_state( + document, + checkpoint_process_group=asymmetric_target_binding, + transaction_id="consensus-subclass-import", + limits=_limits(), + ), + ) + dist.barrier() + object_target, object_target_binding = _make_single( + rank, + group, + compatibility_name="weight", + deterministic=False, + ) + asymmetric_object_binding = object() if rank == 0 else object_target_binding + object_import = _attempt_failure( + object_target, + lambda: object_target.import_portable_state( + document, + checkpoint_process_group=asymmetric_object_binding, + transaction_id="consensus-object-import", + limits=_limits(), + ), + ) + dist.barrier() + alternate_target, alternate_target_binding = _make_single( + rank, + group, + compatibility_name="weight", + deterministic=False, + ) + asymmetric_alternate_binding = ( + _binding_with_process_group( + alternate_target_binding, + alternate_process_group, + ) + if rank == 0 + else alternate_target_binding + ) + alternate_import = _attempt_failure( + alternate_target, + lambda: alternate_target.import_portable_state( + document, + checkpoint_process_group=asymmetric_alternate_binding, + transaction_id="consensus-alternate-handle-import", + limits=_limits(), + ), + ) + dist.barrier() + + name_target, name_binding = _make_single( + rank, + group, + compatibility_name="weight" if rank == 0 else "other_weight", + deterministic=False, + ) + compatibility_import = _attempt_failure( + name_target, + lambda: name_target.import_portable_state( + document, + checkpoint_process_group=name_binding, + transaction_id="consensus-compatibility-import", + limits=_limits(), + ), + ) + dist.barrier() + + two_source, two_source_binding = _make_two( + rank, + group, + reverse=False, + deterministic=True, + ) + two_document = two_source.export_portable_state( + checkpoint_process_group=two_source_binding, + transaction_id="consensus-healthy-two-export", + limits=_limits(), + ) + position_target, position_binding = _make_two( + rank, + group, + reverse=rank == 1, + deterministic=False, + ) + position_import = _attempt_failure( + position_target, + lambda: position_target.import_portable_state( + two_document, + checkpoint_process_group=position_binding, + transaction_id="consensus-position-import", + limits=_limits(), + ), + ) + dist.barrier() + + result_queue.put( + { + "rank": rank, + "binding_export": binding_export, + "binding_import": binding_import, + "object_export": object_export, + "object_import": object_import, + "alternate_export": alternate_export, + "alternate_import": alternate_import, + "compatibility_import": compatibility_import, + "position_import": position_import, + } + ) + except Exception: + result_queue.put({"rank": rank, "fatal_error": traceback.format_exc()}) + finally: + if dist.is_available() and dist.is_initialized(): + if alternate_process_group is not None: + dist.destroy_process_group(alternate_process_group) + dist.destroy_process_group() + + +def _run_workers(): + context = mp.get_context("spawn") + result_queue = context.Queue() + descriptor, init_file = tempfile.mkstemp(prefix="gefen-portable-consensus-") + os.close(descriptor) + os.unlink(init_file) + processes = [context.Process(target=_worker, args=(rank, init_file, result_queue)) for rank in range(_WORLD)] + results = [] + try: + for process in processes: + process.start() + try: + for _ in processes: + results.append(result_queue.get(timeout=60)) + except queue_module.Empty: + pass + for process in processes: + process.join(timeout=5) + for process in processes: + if process.is_alive(): + process.terminate() + process.join(timeout=5) + finally: + for process in processes: + if process.is_alive(): + process.terminate() + process.join(timeout=5) + result_queue.close() + result_queue.join_thread() + if os.path.exists(init_file): + os.unlink(init_file) + assert len(results) == _WORLD, (results, [process.exitcode for process in processes]) + assert all(process.exitcode == 0 for process in processes), [process.exitcode for process in processes] + return sorted(results, key=lambda item: item["rank"]) + + +@pytest.mark.skipif( + not dist.is_available() or not dist.is_gloo_available(), + reason="portable runtime consensus coverage requires Gloo", +) +def test_two_process_portable_runtime_rejects_asymmetric_consensus_inputs(): + results = _run_workers() + assert all("fatal_error" not in result for result in results), results + + expected_messages = { + "binding_export": "CheckpointProcessGroupBinding", + "binding_import": "CheckpointProcessGroupBinding", + "object_export": "CheckpointProcessGroupBinding", + "object_import": "CheckpointProcessGroupBinding", + "alternate_export": "optimizer-owned runtime transport", + "alternate_import": "optimizer-owned runtime transport", + "compatibility_import": "context", + "position_import": "context", + } + failures = [] + for case, expected_message in expected_messages.items(): + messages = [result[case]["message"] for result in results] + if messages[0] != messages[1]: + failures.append((case, "messages diverged", messages)) + if messages[0] is None or expected_message not in messages[0]: + failures.append((case, "missing expected failure", messages)) + mutation_free = all( + all( + result[case][key] + for key in ( + "token_unchanged", + "containers_unchanged", + "state_objects_unchanged", + "states_pristine", + "common_unchanged", + ) + ) + for result in results + ) + if not mutation_free: + failures.append((case, "live state changed", [result[case] for result in results])) + assert not failures, (failures, results) diff --git a/tests/test_portable_runtime_distributed.py b/tests/test_portable_runtime_distributed.py new file mode 100644 index 0000000..1996c37 --- /dev/null +++ b/tests/test_portable_runtime_distributed.py @@ -0,0 +1,567 @@ +"""Warning-strict distributed coverage for topology-neutral optimizer state.""" + +from datetime import timedelta +import multiprocessing as mp +import os +import queue as queue_module +import tempfile +import traceback +import warnings + +import pytest +import torch +import torch.distributed as dist + +from gefen.checkpoint import CheckpointProcessGroupBinding +from gefen.codebook import CodebookProcessGroupBinding +from gefen.contracts import ( + LogicalSlice, + ParameterIdentity, + ParameterLayout, + PlacementKind, + ProcessGroupIdentity, + ShardIdentity, + ShardPlacement, + ShardingManifest, +) +from gefen.gefen import Gefen +from gefen.gefen_muon import GefenMuon +from gefen.portable import _decode_quantized_momentum +from gefen.portable_state import PortableStateLimits +from gefen.rebinding import ParameterRebinding + + +_WORLD = 2 +_PLAIN_FQN = "model.weight" +_MUON_FQN = "model.matrix" + + +def _limits(): + return PortableStateLimits( + max_fragment_tensor_bytes=1 << 20, + max_collective_tensor_bytes=4 << 20, + max_collective_metadata_bytes=4 << 20, + chunk_bytes=7, + max_members=4, + max_metadata_bytes=1 << 20, + max_tree_nodes=10_000, + max_tree_depth=32, + max_container_items=10_000, + max_string_bytes=16 << 10, + max_integer_bytes=128, + max_tensors=256, + max_tensor_rank=8, + diagnostic_bytes=1024, + ) + + +def _members(): + return tuple("rank:{}".format(rank) for rank in range(_WORLD)) + + +def _bindings(group, rank): + member = _members()[rank] + return ( + CodebookProcessGroupBinding(group, member, dist.group.WORLD, torch.device("cpu")), + CheckpointProcessGroupBinding(group, member, dist.group.WORLD, torch.device("cpu")), + ) + + +def _replicated_manifest(identity, group): + members = group.ordered_members + shards = tuple( + ShardIdentity( + identity, + ParameterLayout.REPLICATED, + LogicalSlice.full(identity), + placements=( + ShardPlacement( + "checkpoint", + PlacementKind.REPLICATE, + coordinate, + len(members), + ), + ), + process_group=group, + local_member=member, + ) + for coordinate, member in enumerate(members) + ) + return ShardingManifest(shards), shards + + +def _flat_manifest(identity, group, lengths): + members = group.ordered_members + if len(lengths) != len(members) or sum(lengths) != identity.numel: + raise AssertionError("invalid test partition") + offset = 0 + shards = [] + for coordinate, (member, length) in enumerate(zip(members, lengths)): + shards.append( + ShardIdentity( + identity, + ParameterLayout.FLATTENED_ELEMENT_SHARD, + LogicalSlice(offset, length), + placements=( + ShardPlacement( + "checkpoint", + PlacementKind.FLAT_SHARD, + coordinate, + len(members), + ), + ), + process_group=group, + local_member=member, + ) + ) + offset += length + return ShardingManifest(tuple(shards)), tuple(shards) + + +def _owner_manifest(identity, group, owner): + members = group.ordered_members + shards = tuple( + ShardIdentity( + identity, + ParameterLayout.WHOLE_PARAMETER_OWNER, + LogicalSlice.full(identity) if member == owner else LogicalSlice(0, 0), + placements=( + ShardPlacement( + "checkpoint", + PlacementKind.WHOLE_PARAMETER_OWNER, + coordinate, + len(members), + ), + ), + process_group=group, + local_member=member, + owner=owner, + ) + for coordinate, member in enumerate(members) + ) + return ShardingManifest(shards), shards + + +def _codebook(): + return torch.linspace(-1.0, 1.0, 256, dtype=torch.float32) + + +def _momentum_values(shape): + bits = torch.tensor( + [ + -2147483648, + 0, + 1, + -2147483647, + 1056964608, + -1090519040, + 1078984704, + -1058013184, + ], + dtype=torch.int32, + ) + return bits.view(torch.float32).reshape(shape).clone() + + +def _second_moment_values(shape): + bits = torch.tensor( + [ + -2147483648, + 0, + 1, + 8388608, + 1048576000, + 1065353216, + 1073741824, + 2139095039, + ], + dtype=torch.int32, + ) + return bits.view(torch.float32).reshape(shape).clone() + + +def _quantized_period_one(momentum): + indices = torch.where( + torch.signbit(momentum.reshape(-1)), + torch.zeros(momentum.numel(), dtype=torch.uint8), + torch.full((momentum.numel(),), 255, dtype=torch.uint8), + ) + return indices.reshape(-1, 1), momentum.abs().reshape(-1, 1).clone() + + +def _bits_equal(left, right): + return ( + left.dtype == right.dtype == torch.float32 + and tuple(left.shape) == tuple(right.shape) + and torch.equal(left.contiguous().view(torch.int32), right.contiguous().view(torch.int32)) + ) + + +def _make_plain_replicated(rank, group, *, deterministic): + identity = ParameterIdentity(_PLAIN_FQN, (2, 4)) + parameter = torch.nn.Parameter(torch.zeros(identity.global_shape, dtype=torch.float32)) + optimizer = Gefen( + [("weight", parameter)], + lr=2.5e-3, + betas=(0.8, 0.97), + eps=2.0e-8, + weight_decay=0.03, + fused=False, + force_2d_period_one=True, + factored_v_2d=False, + deterministic=deterministic, + ) + manifest, shards = _replicated_manifest(identity, group) + codebook_binding, checkpoint_binding = _bindings(group, rank) + optimizer.post_sharding( + (ParameterRebinding(parameter, parameter, shards[rank]),), + manifest=manifest, + codebook_process_group=codebook_binding, + ) + return optimizer, parameter, checkpoint_binding + + +def _make_plain_flat(rank, group, lengths, *, deterministic): + identity = ParameterIdentity(_PLAIN_FQN, (2, 4)) + original = torch.nn.Parameter(torch.zeros(identity.global_shape, dtype=torch.float32)) + local = torch.nn.Parameter(torch.zeros(lengths[rank], dtype=torch.float32)) + optimizer = Gefen( + [("weight", original)], + lr=2.5e-3, + betas=(0.8, 0.97), + eps=2.0e-8, + weight_decay=0.03, + fused=False, + force_2d_period_one=True, + factored_v_2d=False, + deterministic=deterministic, + ) + manifest, shards = _flat_manifest(identity, group, lengths) + codebook_binding, checkpoint_binding = _bindings(group, rank) + optimizer.post_sharding( + (ParameterRebinding(original, local, shards[rank]),), + manifest=manifest, + codebook_process_group=codebook_binding, + ) + return optimizer, local, shards[rank], checkpoint_binding + + +def _seed_plain_state(optimizer, parameter): + momentum = _momentum_values((2, 4)) + second_moment = _second_moment_values((2, 4)) + indices, magnitudes = _quantized_period_one(momentum) + optimizer._gefen_global_step = 9 + optimizer._gefen_codebook = _codebook() + optimizer.state[parameter].update( + { + "automatic_period": 1, + "step": 7, + "m_codebook": indices, + "m_magnitude": magnitudes, + "vmean": second_moment.reshape(-1, 1).clone(), + "vmean_step": 6, + } + ) + return momentum, second_moment + + +def _make_muon_owner(rank, group, *, owner, deterministic): + identity = ParameterIdentity(_MUON_FQN, (3, 2)) + original = torch.nn.Parameter(torch.zeros(identity.global_shape, dtype=torch.float32)) + optimizer = GefenMuon( + [("matrix", original)], + lr=3.0e-3, + weight_decay=0.02, + momentum=0.85, + nesterov=False, + ns_steps=2, + fused=False, + sharded_mode="distributed", + deterministic=deterministic, + normuon=True, + normuon_beta2=0.9, + normuon_eps=3.0e-8, + ) + manifest, shards = _owner_manifest(identity, group, owner) + local = original if _members()[rank] == owner else None + codebook_binding, checkpoint_binding = _bindings(group, rank) + optimizer.post_sharding( + (ParameterRebinding(original, local, shards[rank]),), + manifest=manifest, + codebook_process_group=codebook_binding, + ) + return optimizer, local, checkpoint_binding + + +def _seed_muon_state(optimizer, parameter): + optimizer._gefen_global_step = 13 + optimizer._gefen_codebook = _codebook() + if parameter is None: + return + momentum = _momentum_values((2, 4)).reshape(-1)[:6].reshape(3, 2).clone() + indices, magnitudes = _quantized_period_one(momentum) + normuon_bits = torch.tensor([-2147483648, 1, 1082130432], dtype=torch.int32) + optimizer.state[parameter].update( + { + "automatic_period": 1, + "step": 11, + "m_codebook": indices, + "m_magnitude": magnitudes, + "normuon_v": normuon_bits.view(torch.float32).reshape(3, 1).clone(), + "normuon_step": 10, + } + ) + + +def _plain_result(rank, group): + source, source_parameter, source_binding = _make_plain_replicated(rank, group, deterministic=True) + expected_momentum, expected_second = _seed_plain_state(source, source_parameter) + document = source.export_portable_state( + checkpoint_process_group=source_binding, + transaction_id="plain-replicated-export-v1", + limits=_limits(), + ) + + lengths = (3, 5) + target, target_parameter, target_shard, target_binding = _make_plain_flat( + rank, + group, + lengths, + deterministic=False, + ) + target.import_portable_state( + document, + checkpoint_process_group=target_binding, + transaction_id="plain-flat-import-v1", + limits=_limits(), + ) + state = target.state[target_parameter] + decoded = _decode_quantized_momentum( + target._gefen_codebook, + state["m_codebook"], + state["m_magnitude"], + logical_shape=(lengths[rank],), + period=1, + step=state["step"], + ) + start = target_shard.logical_slice.flat_offset + stop = start + target_shard.logical_slice.length + expected_local_momentum = expected_momentum.reshape(-1)[start:stop].clone() + expected_local_second = expected_second.reshape(-1)[start:stop].clone() + parameter_record = document["parameters"][_PLAIN_FQN]["state"] + return { + "digest": document["completion"]["digest"], + "document_momentum_exact": _bits_equal(parameter_record["momentum"], expected_momentum), + "document_second_exact": _bits_equal(parameter_record["second_moment"], expected_second), + "local_momentum_exact": _bits_equal(decoded, expected_local_momentum), + "local_second_exact": _bits_equal(state["vmean"].reshape(-1), expected_local_second), + "counters_exact": target._gefen_global_step == 9 and state["step"] == 7 and state["vmean_step"] == 6, + "deterministic_imported": target._deterministic is True, + "document": document, + } + + +def _asymmetric_failure_result(rank, group, document): + target, target_parameter, _, binding = _make_plain_flat(rank, group, (5, 3), deterministic=False) + if rank == 0: + target.param_groups[0]["rank_local_extension"] = "reject-me" + token_before = target._canonical_import_live_token() + state_object_before = target.state[target_parameter] + try: + target.import_portable_state( + document, + checkpoint_process_group=binding, + transaction_id="plain-asymmetric-import-failure-v1", + limits=_limits(), + ) + except RuntimeError as exc: + message = str(exc) + else: + message = None + return { + "message": message, + "token_unchanged": target._canonical_import_live_token() == token_before, + "state_object_unchanged": target.state[target_parameter] is state_object_before, + "still_pristine": target.state[target_parameter] == {"name": "weight"}, + "common_unchanged": target._gefen_global_step == 0 and target._gefen_codebook is None, + } + + +def _muon_result(rank, group): + source_owner = _members()[0] + target_owner = _members()[1] + source, source_parameter, source_binding = _make_muon_owner( + rank, + group, + owner=source_owner, + deterministic=True, + ) + _seed_muon_state(source, source_parameter) + document = source.export_portable_state( + checkpoint_process_group=source_binding, + transaction_id="muon-owner-export-v1", + limits=_limits(), + ) + + target, target_parameter, target_binding = _make_muon_owner( + rank, + group, + owner=target_owner, + deterministic=False, + ) + target.import_portable_state( + document, + checkpoint_process_group=target_binding, + transaction_id="muon-owner-import-v1", + limits=_limits(), + ) + expected_momentum = _momentum_values((2, 4)).reshape(-1)[:6].reshape(3, 2).clone() + expected_normuon = torch.tensor([-2147483648, 1, 1082130432], dtype=torch.int32).view(torch.float32).reshape(3, 1) + record = document["parameters"][_MUON_FQN]["state"] + if target_parameter is None: + local_exact = len(target.state) == 0 and len(target.param_groups[0]["params"]) == 0 + counters_exact = target._gefen_global_step == 13 + else: + state = target.state[target_parameter] + decoded = _decode_quantized_momentum( + target._gefen_codebook, + state["m_codebook"], + state["m_magnitude"], + logical_shape=(3, 2), + period=1, + step=state["step"], + ) + local_exact = _bits_equal(decoded, expected_momentum) and _bits_equal(state["normuon_v"], expected_normuon) + counters_exact = target._gefen_global_step == 13 and state["step"] == 11 and state["normuon_step"] == 10 + return { + "digest": document["completion"]["digest"], + "document_momentum_exact": _bits_equal(record["momentum"], expected_momentum), + "document_normuon_exact": _bits_equal(record["normuon_v"], expected_normuon), + "local_exact": local_exact, + "counters_exact": counters_exact, + "deterministic_imported": target._deterministic is True, + "target_has_storage": target_parameter is not None, + } + + +def _distributed_worker(rank, init_file, result_queue): + warnings.simplefilter("error") + try: + dist.init_process_group( + "gloo", + init_method="file://{}".format(init_file), + rank=rank, + world_size=_WORLD, + timeout=timedelta(seconds=90), + ) + group = ProcessGroupIdentity("portable_checkpoint", _members()) + plain = _plain_result(rank, group) + dist.barrier() + asymmetric = _asymmetric_failure_result(rank, group, plain.pop("document")) + dist.barrier() + muon = _muon_result(rank, group) + dist.barrier() + result_queue.put( + { + "rank": rank, + "plain": plain, + "asymmetric": asymmetric, + "muon": muon, + } + ) + except Exception: + result_queue.put({"rank": rank, "fatal_error": traceback.format_exc()}) + finally: + if dist.is_available() and dist.is_initialized(): + dist.destroy_process_group() + + +def _run_distributed_workers(): + context = mp.get_context("spawn") + result_queue = context.Queue() + descriptor, init_file = tempfile.mkstemp(prefix="gefen-portable-runtime-") + os.close(descriptor) + os.unlink(init_file) + processes = [ + context.Process(target=_distributed_worker, args=(rank, init_file, result_queue)) for rank in range(_WORLD) + ] + results = [] + try: + for process in processes: + process.start() + try: + for _ in processes: + results.append(result_queue.get(timeout=180)) + except queue_module.Empty: + pass + for process in processes: + process.join(timeout=10) + for process in processes: + if process.is_alive(): + process.terminate() + process.join(timeout=5) + finally: + for process in processes: + if process.is_alive(): + process.terminate() + process.join(timeout=5) + result_queue.close() + result_queue.join_thread() + if os.path.exists(init_file): + os.unlink(init_file) + assert len(results) == _WORLD, (results, [process.exitcode for process in processes]) + assert all(process.exitcode == 0 for process in processes), [process.exitcode for process in processes] + return sorted(results, key=lambda item: item["rank"]) + + +@pytest.mark.skipif( + not dist.is_available() or not dist.is_gloo_available(), + reason="portable runtime topology-change coverage requires Gloo", +) +def test_two_process_portable_runtime_reshards_relocates_and_fails_atomically(): + results = _run_distributed_workers() + + assert all("fatal_error" not in result for result in results), results + assert len({result["plain"]["digest"] for result in results}) == 1 + assert len({result["muon"]["digest"] for result in results}) == 1 + for result in results: + plain = result["plain"] + assert all( + plain[key] + for key in ( + "document_momentum_exact", + "document_second_exact", + "local_momentum_exact", + "local_second_exact", + "counters_exact", + "deterministic_imported", + ) + ), result + muon = result["muon"] + assert all( + muon[key] + for key in ( + "document_momentum_exact", + "document_normuon_exact", + "local_exact", + "counters_exact", + "deterministic_imported", + ) + ), result + assert muon["target_has_storage"] is (result["rank"] == 1) + + messages = [result["asymmetric"]["message"] for result in results] + assert messages[0] == messages[1] + assert messages[0] is not None and "unknown keys" in messages[0] + assert all( + all( + result["asymmetric"][key] + for key in ( + "token_unchanged", + "state_object_unchanged", + "still_pristine", + "common_unchanged", + ) + ) + for result in results + ), results diff --git a/tests/test_portable_runtime_hardening.py b/tests/test_portable_runtime_hardening.py new file mode 100644 index 0000000..9d1ccbf --- /dev/null +++ b/tests/test_portable_runtime_hardening.py @@ -0,0 +1,457 @@ +"""Warning-strict CPU regressions for portable runtime publication guards.""" + +import pytest +import torch + +import gefen.portable_runtime as portable_runtime +from gefen.checkpoint import CheckpointProcessGroupBinding +from gefen.codebook import CodebookProcessGroupBinding +from gefen.contracts import ( + CheckpointTransport, + LogicalSlice, + ParameterIdentity, + ParameterLayout, + PlacementKind, + ProcessGroupIdentity, + ShardIdentity, + ShardPlacement, + ShardingManifest, +) +from gefen.gefen import Gefen +from gefen.portable_state import PortableStateLimits +from gefen.rebinding import ParameterRebinding + + +_MEMBER = "rank:0" + + +class _ParameterIdentitySubclass(ParameterIdentity): + pass + + +class _ProcessGroupIdentitySubclass(ProcessGroupIdentity): + pass + + +class _ShardPlacementSubclass(ShardPlacement): + pass + + +class _StringSubclass(str): + pass + + +def _limits(*, max_fragment_tensor_bytes=4 << 20): + return PortableStateLimits( + max_fragment_tensor_bytes=max_fragment_tensor_bytes, + max_collective_tensor_bytes=16 << 20, + max_collective_metadata_bytes=64 << 20, + ) + + +@pytest.mark.parametrize("subclass_kind", ["parameter", "process_group", "placement"]) +def test_nested_identity_subclasses_are_rejected(subclass_kind): + parameter_type = _ParameterIdentitySubclass if subclass_kind == "parameter" else ParameterIdentity + group_type = _ProcessGroupIdentitySubclass if subclass_kind == "process_group" else ProcessGroupIdentity + placement_type = _ShardPlacementSubclass if subclass_kind == "placement" else ShardPlacement + identity = parameter_type("layer.weight", (2, 3)) + group = group_type("checkpoint", (_MEMBER,)) + shard = ShardIdentity( + identity, + ParameterLayout.REPLICATED, + LogicalSlice.full(identity), + placements=( + placement_type( + "checkpoint", + PlacementKind.REPLICATE, + 0, + 1, + ), + ), + process_group=group, + local_member=_MEMBER, + ) + + with pytest.raises(TypeError, match="exact"): + portable_runtime._validate_exact_shard_identity(shard) + + +def _optimizer(*, layout, deterministic): + identity = ParameterIdentity("layer.weight", (2, 3)) + if layout is ParameterLayout.REPLICATED: + parameter = torch.nn.Parameter(torch.arange(1, 7, dtype=torch.float32).reshape(identity.global_shape)) + placement_kind = PlacementKind.REPLICATE + elif layout is ParameterLayout.FLATTENED_ELEMENT_SHARD: + parameter = torch.nn.Parameter(torch.arange(1, 7, dtype=torch.float32)) + placement_kind = PlacementKind.FLAT_SHARD + else: + raise AssertionError("unsupported test layout") + + optimizer = Gefen( + [("weight", parameter)], + fused=False, + factored_v_2d=False, + force_2d_period_one=True, + deterministic=deterministic, + ) + group = ProcessGroupIdentity("checkpoint", (_MEMBER,)) + shard = ShardIdentity( + identity, + layout, + LogicalSlice.full(identity), + placements=(ShardPlacement("checkpoint", placement_kind, 0, 1),), + process_group=group, + local_member=_MEMBER, + ) + optimizer.post_sharding( + (ParameterRebinding(parameter, parameter, shard),), + manifest=ShardingManifest((shard,)), + codebook_process_group=CodebookProcessGroupBinding( + group, + _MEMBER, + None, + torch.device("cpu"), + ), + ) + binding = CheckpointProcessGroupBinding( + group, + _MEMBER, + None, + torch.device("cpu"), + ) + return optimizer, parameter, binding + + +def _initialize(optimizer, parameter): + optimizer._gefen_global_step = 4 + optimizer._gefen_codebook = torch.linspace(-1.0, 1.0, 256, dtype=torch.float32) + optimizer.state[parameter].update( + { + "automatic_period": 1, + "step": 4, + "m_codebook": torch.tensor( + [[0], [255], [128], [64], [192], [0]], + dtype=torch.uint8, + ), + "m_magnitude": torch.tensor( + [[1.0], [2.0], [3.0], [4.0], [5.0], [6.0]], + dtype=torch.float32, + ), + "vmean": torch.arange(1, 7, dtype=torch.float32).reshape(6, 1), + "vmean_step": 3, + } + ) + + +def _has_canonical_global(optimizer): + return any( + support.transport is CheckpointTransport.CANONICAL_GLOBAL + for support in optimizer.optimizer_contract().capabilities.checkpoints + ) + + +@pytest.mark.parametrize( + ("descriptor", "field", "value"), + [ + ("manifest", "schema_version", 2), + ("shard", "schema_version", 2), + ("parameter", "schema_version", 2), + ("process_group", "schema_version", 2), + ("process_group", "semantic_name", ""), + ("process_group", "semantic_name", " checkpoint"), + ("process_group", "semantic_name", "check\x00point"), + ("parameter", "fqn", ""), + ("parameter", "fqn", " layer.weight"), + ("parameter", "fqn", "layer..weight"), + ("logical_slice", "flat_offset", -1), + ("placement", "mesh_axis", " checkpoint"), + ("placement", "coordinate", -1), + ("placement", "parts", 0), + ], +) +def test_mutated_descriptor_invariants_disable_portable_capability( + descriptor, + field, + value, +): + optimizer, _, _ = _optimizer( + layout=ParameterLayout.REPLICATED, + deterministic=True, + ) + manifest = optimizer._gefen_sharding_manifest + shard = manifest.shards[0] + target = { + "manifest": manifest, + "shard": shard, + "parameter": shard.parameter, + "process_group": shard.process_group, + "logical_slice": shard.logical_slice, + "placement": shard.placements[0], + }[descriptor] + object.__setattr__(target, field, value) + + with pytest.raises((TypeError, ValueError)): + portable_runtime._validate_exact_manifest(manifest) + assert not _has_canonical_global(optimizer) + + +def test_codebook_binding_local_member_subclass_disables_portable_capability(): + optimizer, _, _ = _optimizer( + layout=ParameterLayout.REPLICATED, + deterministic=True, + ) + scope = optimizer._gefen_codebook_process_group + object.__setattr__(scope, "local_member", _StringSubclass(scope.local_member)) + + assert not _has_canonical_global(optimizer) + + +def _drift_storage(parameter, drift): + if drift == "replicated_shape": + parameter.data = parameter.detach().reshape(3, 2) + elif drift == "flat_shape": + parameter.data = parameter.detach().reshape(2, 3) + elif drift == "flat_noncontiguous": + parameter.data = torch.arange(12, dtype=torch.float32)[::2] + elif drift == "complex_dtype": + parameter.data = parameter.detach().to(dtype=torch.complex64) + else: + raise AssertionError("unknown storage drift") + + +@pytest.mark.parametrize( + ("layout", "drift"), + [ + (ParameterLayout.REPLICATED, "replicated_shape"), + (ParameterLayout.FLATTENED_ELEMENT_SHARD, "flat_shape"), + (ParameterLayout.FLATTENED_ELEMENT_SHARD, "flat_noncontiguous"), + (ParameterLayout.REPLICATED, "complex_dtype"), + ], +) +def test_finalized_parameter_storage_drift_disables_and_rejects_portable_io( + layout, + drift, +): + source, source_parameter, source_binding = _optimizer( + layout=layout, + deterministic=True, + ) + _initialize(source, source_parameter) + document = source.export_portable_state( + checkpoint_process_group=source_binding, + transaction_id="hardening-source-export-{}".format(drift), + limits=_limits(), + ) + + _drift_storage(source_parameter, drift) + assert not _has_canonical_global(source) + with pytest.raises(RuntimeError, match="geometry|dtype"): + source.export_portable_state( + checkpoint_process_group=source_binding, + transaction_id="hardening-drifted-export-{}".format(drift), + limits=_limits(), + ) + + target, target_parameter, target_binding = _optimizer( + layout=layout, + deterministic=False, + ) + _drift_storage(target_parameter, drift) + assert not _has_canonical_global(target) + state_mapping = target.state + parameter_state = target.state[target_parameter] + defaults = target.defaults + with pytest.raises(RuntimeError, match="geometry|dtype"): + target.import_portable_state( + document, + checkpoint_process_group=target_binding, + transaction_id="hardening-drifted-import-{}".format(drift), + limits=_limits(), + ) + assert target.state is state_mapping + assert target.state[target_parameter] is parameter_state + assert parameter_state == {"name": "weight"} + assert target.defaults is defaults + assert target._gefen_global_step == 0 + assert target._gefen_codebook is None + assert target._deterministic is False + + +class _UpdateBombDict(dict): + update_calls = 0 + + def update(self, *args, **kwargs): + self.update_calls += 1 + self["publication_started"] = True + raise AssertionError("defaults publication must not start") + + +def test_defaults_dict_subclass_is_rejected_before_publication(): + source, source_parameter, source_binding = _optimizer( + layout=ParameterLayout.REPLICATED, + deterministic=True, + ) + _initialize(source, source_parameter) + document = source.export_portable_state( + checkpoint_process_group=source_binding, + transaction_id="hardening-defaults-source", + limits=_limits(), + ) + target, target_parameter, target_binding = _optimizer( + layout=ParameterLayout.REPLICATED, + deterministic=False, + ) + hostile_defaults = dict.__new__(_UpdateBombDict) + dict.update(hostile_defaults, target.defaults) + hostile_defaults.update_calls = 0 + target.defaults = hostile_defaults + before_defaults = dict(hostile_defaults) + state_mapping = target.state + parameter_state = target.state[target_parameter] + + assert not _has_canonical_global(target) + with pytest.raises(RuntimeError, match="exact built-in optimizer defaults"): + target.import_portable_state( + document, + checkpoint_process_group=target_binding, + transaction_id="hardening-defaults-import", + limits=_limits(), + ) + + assert target.defaults is hostile_defaults + assert dict(hostile_defaults) == before_defaults + assert hostile_defaults.update_calls == 0 + assert "publication_started" not in hostile_defaults + assert target.state is state_mapping + assert target.state[target_parameter] is parameter_state + assert parameter_state == {"name": "weight"} + assert target._gefen_global_step == 0 + assert target._gefen_codebook is None + assert target._deterministic is False + + +def test_foreign_state_inserted_after_prepare_fails_freshness_without_publication( + monkeypatch, +): + source, source_parameter, source_binding = _optimizer( + layout=ParameterLayout.REPLICATED, + deterministic=True, + ) + _initialize(source, source_parameter) + document = source.export_portable_state( + checkpoint_process_group=source_binding, + transaction_id="hardening-freshness-source", + limits=_limits(), + ) + target, target_parameter, target_binding = _optimizer( + layout=ParameterLayout.REPLICATED, + deterministic=False, + ) + assert _has_canonical_global(target) + state_mapping = target.state + parameter_state = target.state[target_parameter] + foreign_parameter = torch.nn.Parameter(torch.zeros(1, dtype=torch.float32)) + foreign_state = {"name": "foreign"} + injected = False + original_status = portable_runtime._collective_unanimous_status + + def status_with_concurrent_insertion(*args, **kwargs): + nonlocal injected + result = original_status(*args, **kwargs) + if kwargs["operation"] == "portable_import_prepare" and not injected: + target.state[foreign_parameter] = foreign_state + injected = True + return result + + monkeypatch.setattr( + portable_runtime, + "_collective_unanimous_status", + status_with_concurrent_insertion, + ) + with pytest.raises(RuntimeError, match="foreign parameter keys|changed after"): + target.import_portable_state( + document, + checkpoint_process_group=target_binding, + transaction_id="hardening-freshness-import", + limits=_limits(), + ) + + assert injected + assert target.state is state_mapping + assert target.state[target_parameter] is parameter_state + assert parameter_state == {"name": "weight"} + assert target.state[foreign_parameter] is foreign_state + assert target._gefen_global_step == 0 + assert target._gefen_codebook is None + assert target._deterministic is False + + +def test_fragment_limit_rejects_before_dense_momentum_materialization(monkeypatch): + optimizer, parameter, binding = _optimizer( + layout=ParameterLayout.REPLICATED, + deterministic=True, + ) + _initialize(optimizer, parameter) + reached_decode = False + reached_live_token = False + + def fail_if_decoded(*args, **kwargs): + nonlocal reached_decode + reached_decode = True + raise AssertionError("dense momentum decode was reached") + + def fail_if_tokenized(*args, **kwargs): + nonlocal reached_live_token + reached_live_token = True + raise AssertionError("content-bearing live token was reached") + + monkeypatch.setattr( + portable_runtime, + "_decode_local_momentum", + fail_if_decoded, + ) + monkeypatch.setattr( + portable_runtime, + "_portable_live_token", + fail_if_tokenized, + ) + with pytest.raises(RuntimeError, match="max_fragment_tensor_bytes"): + optimizer.export_portable_state( + checkpoint_process_group=binding, + transaction_id="hardening-fragment-budget", + limits=_limits(max_fragment_tensor_bytes=1), + ) + assert not reached_decode + assert not reached_live_token + + +def test_pristine_import_does_not_bound_or_materialize_obsolete_target_payload(): + source, _source_parameter, source_binding = _optimizer( + layout=ParameterLayout.REPLICATED, + deterministic=True, + ) + tiny_limits = _limits(max_fragment_tensor_bytes=1) + document = source.export_portable_state( + checkpoint_process_group=source_binding, + transaction_id="hardening-pristine-source", + limits=tiny_limits, + ) + + target, target_parameter, target_binding = _optimizer( + layout=ParameterLayout.REPLICATED, + deterministic=False, + ) + _initialize(target, target_parameter) + old_state = target.state[target_parameter] + target.import_portable_state( + document, + checkpoint_process_group=target_binding, + transaction_id="hardening-pristine-target", + limits=tiny_limits, + ) + + assert target.state[target_parameter] is not old_state + assert target.state[target_parameter] == {"name": "weight"} + assert target._gefen_global_step == 0 + assert target._gefen_codebook is None + assert target._deterministic is True diff --git a/tests/test_portable_schema.py b/tests/test_portable_schema.py index 6dc3115..e8ba823 100644 --- a/tests/test_portable_schema.py +++ b/tests/test_portable_schema.py @@ -282,4 +282,6 @@ def test_portable_schema_and_checkpoint_binding_exports_are_public(): assert gefen.normalize_portable_state_document is normalize_portable_state_document assert gefen.portable_state_digest is portable_state_digest assert gefen.CheckpointProcessGroupBinding.__module__ == "gefen.checkpoint" + assert gefen.PortableStateLimits.__module__ == "gefen.portable_state" + assert gefen.PortableStateProvider.__module__ == "gefen.contracts" assert gefen.LogicalRegion.__module__ == "gefen.contracts" From 14279e916638e9df35ed7a24e10ea11991e2d703 Mon Sep 17 00:00:00 2001 From: thad0ctor Date: Mon, 13 Jul 2026 02:56:42 -0700 Subject: [PATCH 17/52] Add portable DCP persistence --- COMPATIBILITY.md | 2 + docs/optimizer_contracts.md | 26 +- src/gefen/__init__.py | 6 + src/gefen/portable_dcp.py | 530 ++++++++++++++++++++ tests/test_portable_dcp.py | 637 +++++++++++++++++++++++++ tests/test_portable_dcp_distributed.py | 530 ++++++++++++++++++++ tests/test_portable_schema.py | 2 + 7 files changed, 1732 insertions(+), 1 deletion(-) create mode 100644 src/gefen/portable_dcp.py create mode 100644 tests/test_portable_dcp.py create mode 100644 tests/test_portable_dcp_distributed.py diff --git a/COMPATIBILITY.md b/COMPATIBILITY.md index a831707..b457dbd 100644 --- a/COMPATIBILITY.md +++ b/COMPATIBILITY.md @@ -95,6 +95,8 @@ Native single-process optimizer `state_dict()`/`load_state_dict()` and the expli This format is deliberately same-topology only and currently requires one 1-D DeviceMesh spanning the default process-group world. Every rank must participate in both save and restore; each process temporarily holds all serialized rank payloads on CPU, so the leading checkpoint-time CPU cost approaches `world_size ×` its local optimizer-state size plus local serialization scratch. Loading validates world size, parameter order and names, global and local shapes and dtypes, mesh membership and names, structural placements, rank coordinates, global step, deterministic policy, frozen codebook, and sharded mode before mutation. Multidimensional meshes, subgroups, pipeline-local optimizers, world-size/topology changes, and old unsafe untagged full checkpoints fail closed rather than silently applying rank 0's state to every shard. No optimizer-state reshard portability is claimed, model-only DCP is unaffected, and the full-state DCP support described here does not extend beyond plain Gefen and Muon `approx`. +Finalized exact period-one plain Gefen and `GefenMuon(sharded_mode="distributed")` 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. 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. DTensor targets, Hybrid composition, asynchronous saving, mixed model/optimizer `Stateful` composition, singleton checkpoint scopes inside a larger initialized default world, and multi-member checkpoint groups whose coordinate zero is not global rank zero are not supported by this path. + ## Transformers Trainer DDP The `benchmarks.trainer_resume` gate exercises plain Gefen, GefenMuon+AdamW, and GefenMuon+Gefen through Trainer's internal Accelerate wrapper with tied weights, gradient accumulation, a changing scheduler, native Trainer checkpoint files, BF16 fused updates, and two-rank DDP replica hashes. All three recipes have passed its deterministic fused-BF16 two-rank configuration on homogeneous GPUs, which requires exact model, optimizer, scheduler, LR, and logged-loss agreement between uninterrupted and resumed runs. Run the gate with: diff --git a/docs/optimizer_contracts.md b/docs/optimizer_contracts.md index 499f8a1..bc4fa13 100644 --- a/docs/optimizer_contracts.md +++ b/docs/optimizer_contracts.md @@ -90,11 +90,35 @@ 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. + +```python +import torch.distributed.checkpoint as dcp + +from gefen import load_portable_dcp, save_portable_dcp + +save_portable_dcp( + optimizer, + checkpoint_process_group=checkpoint_binding, + storage_writer=dcp.FileSystemWriter("checkpoint/optimizer"), + transaction_id="optimizer-save-0001", + limits=limits, +) + +load_portable_dcp( + target_optimizer, + checkpoint_process_group=target_checkpoint_binding, + storage_reader=dcp.FileSystemReader("checkpoint/optimizer"), + transaction_id="optimizer-load-0001", + limits=limits, +) +``` + The dynamic `CANONICAL_GLOBAL` checkpoint declaration appears only while the live finalized optimizer passes the exact runtime readiness checks: explicit process-group scope, stable logical slots and manifest, ordinary built-in containers, supported CPU/CUDA tensor storage, no active compilation or CUDA capture, `capturable=False`, `stochastic_round=False`, a complete declared native state variant, and period one for selected or initialized state. Plain Gefen supports replicated and contiguous flattened element shards. Block-second-moment state can reshard between replicated and flattened targets; a logical matrix using factored second moments remains replicated and same-topology because factored-to-block representation migration is not implemented. GefenMuon supports replicated matrices and whole-parameter ownership when every participating group uses `sharded_mode="distributed"`; the transport can change placement and redistribute owners across world sizes, including NorMuon row state. 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, GefenMuon modes other than `distributed` for whole-owner transport, `GefenMuonHybrid`, tied-alias expansion, and direct DCP orchestration. The optimizer-facing API is the adapter boundary for a DCP or platform integration; the core does not register a framework-specific state-dict adapter or perform storage I/O. `state_offload` remains false because a CPU portable document is a checkpoint artifact, not live optimizer state that can be stepped while offloaded. +Portable v3 currently excludes non-period-one initialized state, second-moment representation conversion, DTensor layouts, stochastic rounding, capturable/device-authoritative state, GefenMuon modes other than `distributed` for whole-owner transport, `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. `state_offload` remains false because a CPU portable document is a checkpoint artifact, not live optimizer state that can be stepped while offloaded. ## Quiescent optimizer-state movement diff --git a/src/gefen/__init__.py b/src/gefen/__init__.py index 87129dc..9f0988f 100644 --- a/src/gefen/__init__.py +++ b/src/gefen/__init__.py @@ -53,6 +53,8 @@ "build_portable_state_document", "normalize_portable_state_document", "portable_state_digest", + "load_portable_dcp", + "save_portable_dcp", "split_params_for_muon", "validate_split", "kernels", @@ -109,6 +111,10 @@ def __getattr__(name): from .portable_state import PortableStateLimits return PortableStateLimits + if name in ("load_portable_dcp", "save_portable_dcp"): + from . import portable_dcp + + return getattr(portable_dcp, name) if name in ( "CONTRACT_SCHEMA_VERSION", "IDENTITY_SCHEMA_VERSION", diff --git a/src/gefen/portable_dcp.py b/src/gefen/portable_dcp.py new file mode 100644 index 0000000..9b8dd1d --- /dev/null +++ b/src/gefen/portable_dcp.py @@ -0,0 +1,530 @@ +"""Synchronous PyTorch DCP persistence for collective portable Gefen state.""" + +from __future__ import annotations + +import math +import os + +import torch + +from gefen.checkpoint import CheckpointProcessGroupBinding +from gefen.portable_state import PortableStateLimits +from gefen.portable_wire import ( + _DTYPE_BY_VALUE, + _CanonicalWireLimits, + _CanonicalWirePlan, + _parse_canonical_wire_metadata, + _prepare_canonical_wire_value, + _reconstruct_canonical_wire_value, +) + + +_DCP_METADATA_KEY = "__gefen_portable_metadata_v1__" +_DCP_PAYLOAD_PREFIX = "__gefen_portable_tensor_" +_DCP_PAYLOAD_SUFFIX = "__" +_DCP_PAYLOAD_DIGITS = 16 +_DCP_SAVE_PREFLIGHT_TRANSACTION = "gefen-portable-dcp-save-preflight-v1" +_DCP_LOAD_PREFLIGHT_TRANSACTION = "gefen-portable-dcp-load-preflight-v1" + + +def _require_namespace(namespace, limits: PortableStateLimits) -> str: + if type(namespace) is not str or not namespace or namespace != namespace.strip(): + raise ValueError("namespace must be a non-empty trimmed string") + if "." in namespace or "\x00" in namespace: + raise ValueError("namespace must not contain dots or NUL") + if any(not (character.isascii() and (character.isalnum() or character in {"_", "-"})) for character in namespace): + raise ValueError("namespace must contain only ASCII letters, digits, underscores, and hyphens") + from gefen.portable_runtime import _bounded_utf8_length + + _bounded_utf8_length( + namespace, + limit=limits.max_string_bytes, + name="DCP namespace", + ) + return namespace + + +def _payload_key(index: int) -> str: + if type(index) is not int or index < 0 or index >= 10**_DCP_PAYLOAD_DIGITS: + raise ValueError("portable DCP payload index is out of range") + return "{}{:0{}d}{}".format( + _DCP_PAYLOAD_PREFIX, + index, + _DCP_PAYLOAD_DIGITS, + _DCP_PAYLOAD_SUFFIX, + ) + + +def _flat_key(namespace: str, inner_key: str) -> str: + return "{}.{}".format(namespace, inner_key) + + +def _validate_dcp_runtime(binding: CheckpointProcessGroupBinding) -> None: + import torch.distributed as dist + + members = binding.identity.ordered_members + if len(members) == 1: + if not dist.is_available() or not dist.is_initialized(): + return + if dist.get_world_size() > 1: + raise RuntimeError("a singleton portable DCP binding cannot run inside a larger initialized default world") + backend = str(dist.get_backend()).lower() + else: + if dist.get_global_rank(binding.process_group, 0) != 0: + raise RuntimeError("portable DCP requires checkpoint group coordinate zero to be global rank zero") + backend = str(dist.get_backend(binding.process_group)).lower() + if "nccl" in backend: + if binding.collective_device.type != "cuda": + raise ValueError("portable DCP requires a CUDA checkpoint device with NCCL") + expected = binding.collective_device.index + if expected is None or torch.cuda.current_device() != expected: + raise ValueError("portable DCP requires the current CUDA device to match the checkpoint binding") + elif "gloo" in backend or "mpi" in backend: + if binding.collective_device.type != "cpu": + raise ValueError("portable DCP requires a CPU checkpoint device with this backend") + + +def _storage_identity(storage, limits: PortableStateLimits): + storage_class = type(storage) + module = storage_class.__module__ + qualname = storage_class.__qualname__ + if type(module) is not str or type(qualname) is not str: + raise TypeError("DCP storage classes require exact string identities") + checkpoint_id = getattr(storage, "checkpoint_id", None) + try: + checkpoint_id = os.fspath(checkpoint_id) + except TypeError as exc: + raise TypeError("DCP storage must expose a string or path-like checkpoint_id") from exc + if type(checkpoint_id) is not str or not checkpoint_id or "\x00" in checkpoint_id: + raise ValueError("DCP storage checkpoint_id must be a non-empty string without NUL") + from torch.distributed.checkpoint import FileSystemReader, FileSystemWriter + + if isinstance(storage, (FileSystemReader, FileSystemWriter)): + checkpoint_id = os.path.abspath(checkpoint_id) + from gefen.portable_runtime import _bounded_utf8_length + + for name, value in ( + ("DCP storage class module", module), + ("DCP storage class qualname", qualname), + ("DCP storage checkpoint_id", checkpoint_id), + ): + _bounded_utf8_length( + value, + limit=limits.max_string_bytes, + name=name, + ) + return { + "module": module, + "qualname": qualname, + "checkpoint_id": checkpoint_id, + } + + +def _preflight_dcp_operation( + optimizer, + *, + checkpoint_process_group, + transaction_id, + limits, + namespace, + storage, + storage_type, + operation: str, +): + from gefen import portable_runtime as runtime + from gefen.portable_collective import _collective_unanimous_status + + transport = runtime._preflight_transport_binding( + optimizer, + checkpoint_process_group, + ) + binding = None + normalized_limits = None + normalized_transaction = None + normalized_namespace = None + implementation = None + context_digest = bytes(32) + wire_limits = runtime._STATUS_FALLBACK_LIMITS + error = None + try: + runtime._validate_supplied_binding(checkpoint_process_group, transport) + binding = transport + normalized_limits = runtime._require_limits(limits) + wire_limits = normalized_limits._wire_limits() + normalized_transaction = runtime._require_transaction_id(transaction_id) + normalized_namespace = _require_namespace(namespace, normalized_limits) + if not isinstance(storage, storage_type): + raise TypeError("storage must be a {}".format(storage_type.__name__)) + storage_identity = _storage_identity(storage, normalized_limits) + implementation = runtime._optimizer_implementation(optimizer) + runtime._validate_context_identity_bounds(binding, normalized_limits) + _validate_dcp_runtime(binding) + prepared = runtime._prepare_local_structure( + optimizer, + implementation, + binding, + normalized_limits, + include_payload=False, + ) + runtime._validate_prepared_local_state( + optimizer, + implementation, + prepared, + ) + context = { + **runtime._base_context(binding, implementation), + "dcp_envelope_version": 1, + "dcp_namespace": normalized_namespace, + "dcp_storage": storage_identity, + "transaction_id": normalized_transaction, + } + runtime._preflight_portable_value(context, normalized_limits) + context_digest = runtime._context_digest(context) + except Exception as exc: + error = exc + _collective_unanimous_status( + transport, + error, + operation="portable_dcp_{}_preflight".format(operation), + transaction_id=(_DCP_SAVE_PREFLIGHT_TRANSACTION if operation == "save" else _DCP_LOAD_PREFLIGHT_TRANSACTION), + context_digest=context_digest, + limits=wire_limits, + ) + assert ( + binding is not None + and normalized_limits is not None + and normalized_transaction is not None + and normalized_namespace is not None + and implementation is not None + ) + return ( + binding, + normalized_transaction, + normalized_limits, + normalized_namespace, + ) + + +def _metadata_tensor(metadata: bytes) -> torch.Tensor: + if type(metadata) is not bytes: + raise TypeError("portable DCP metadata must be bytes") + return torch.frombuffer(bytearray(metadata), dtype=torch.uint8) + + +def _state_from_plan( + namespace: str, + plan: _CanonicalWirePlan, + limits: _CanonicalWireLimits, +): + if type(plan) is not _CanonicalWirePlan: + raise TypeError("plan must be a canonical wire plan") + if type(limits) is not _CanonicalWireLimits: + raise TypeError("limits must be canonical wire limits") + if len(plan.payload_tensors) + 1 > limits.max_container_items: + raise ValueError("portable DCP envelope exceeds max_container_items") + envelope = {_DCP_METADATA_KEY: _metadata_tensor(plan.metadata)} + envelope.update({_payload_key(index): tensor for index, tensor in enumerate(plan.payload_tensors)}) + return {namespace: envelope} + + +def _validate_tensor_metadata(value, *, name: str, limits: _CanonicalWireLimits): + from torch.distributed.checkpoint.metadata import ( + ChunkStorageMetadata, + TensorProperties, + TensorStorageMetadata, + ) + + if type(value) is not TensorStorageMetadata: + raise TypeError("{} must be full-tensor DCP metadata".format(name)) + properties = value.properties + if ( + type(properties) is not TensorProperties + or type(value.size) is not torch.Size + or properties.dtype not in _DTYPE_BY_VALUE + or properties.layout is not torch.strided + or properties.requires_grad is not False + or properties.memory_format is not torch.contiguous_format + or properties.pin_memory is not False + ): + raise TypeError("{} has unsupported tensor properties".format(name)) + shape = tuple(value.size) + if len(shape) > limits.max_tensor_rank or any( + type(dimension) is not int or dimension < 0 or dimension > (1 << 63) - 1 for dimension in shape + ): + raise ValueError("{} has invalid tensor geometry".format(name)) + if type(value.chunks) is not list or len(value.chunks) != 1: + raise ValueError("{} must contain exactly one full DCP chunk".format(name)) + chunk = value.chunks[0] + if ( + type(chunk) is not ChunkStorageMetadata + or type(chunk.offsets) is not torch.Size + or type(chunk.sizes) is not torch.Size + or tuple(chunk.offsets) != (0,) * len(shape) + or tuple(chunk.sizes) != shape + ): + raise ValueError("{} must contain one complete unsharded tensor".format(name)) + numel = math.prod(shape) + element_size = _DTYPE_BY_VALUE[properties.dtype][1] + nbytes = numel * element_size + return properties.dtype, shape, nbytes + + +def _allocate_dcp_state( + metadata, + *, + namespace: str, + limits: _CanonicalWireLimits, +): + from torch.distributed.checkpoint.metadata import Metadata + + if type(metadata) is not Metadata: + raise TypeError("portable DCP requires exact checkpoint Metadata") + entries = metadata.state_dict_metadata + if type(entries) is not dict: + raise TypeError("portable DCP metadata entries must be a dict") + if len(entries) < 1 or len(entries) > limits.max_tensors + 1 or len(entries) > limits.max_container_items: + raise ValueError("portable DCP tensor count exceeds limits") + metadata_key = _flat_key(namespace, _DCP_METADATA_KEY) + if metadata_key not in entries: + raise ValueError("portable DCP metadata tensor is missing") + payload_count = len(entries) - 1 + expected_payload_keys = tuple(_flat_key(namespace, _payload_key(index)) for index in range(payload_count)) + expected_keys = {metadata_key, *expected_payload_keys} + if set(entries) != expected_keys or any(type(key) is not str for key in entries): + raise ValueError("portable DCP checkpoint has unexpected tensor keys") + planner_data = metadata.planner_data + expected_planner_data = {key: tuple(key.split(".", 1)) for key in expected_keys} + if ( + type(planner_data) is not dict + or planner_data != expected_planner_data + or any( + type(key) is not str + or type(path) is not tuple + or len(path) != 2 + or any(type(component) is not str for component in path) + for key, path in planner_data.items() + ) + ): + raise ValueError("portable DCP checkpoint has incompatible planner paths") + + metadata_dtype, metadata_shape, metadata_nbytes = _validate_tensor_metadata( + entries[metadata_key], + name="portable DCP wire metadata", + limits=limits, + ) + if ( + metadata_dtype is not torch.uint8 + or len(metadata_shape) != 1 + or metadata_nbytes == 0 + or metadata_nbytes > limits.max_metadata_bytes + ): + raise ValueError("portable DCP wire metadata exceeds its byte limit") + + payload_specs = [] + total_payload_bytes = 0 + for index, key in enumerate(expected_payload_keys): + dtype, shape, nbytes = _validate_tensor_metadata( + entries[key], + name="portable DCP payload {}".format(index), + limits=limits, + ) + if total_payload_bytes > limits.max_fragment_tensor_bytes - nbytes: + raise ValueError("portable DCP payload tensors exceed their byte limit") + total_payload_bytes += nbytes + payload_specs.append((_payload_key(index), dtype, shape)) + + envelope = { + _DCP_METADATA_KEY: torch.empty( + metadata_shape, + dtype=metadata_dtype, + device="cpu", + ) + } + for key, dtype, shape in payload_specs: + envelope[key] = torch.empty(shape, dtype=dtype, device="cpu") + return {namespace: envelope} + + +def _load_planner(namespace: str, limits: _CanonicalWireLimits): + import torch.distributed.checkpoint as dcp + + class _PortableDCPDynamicLoadPlanner(dcp.DefaultLoadPlanner): + def set_up_planner( + self, + state_dict, + metadata=None, + is_coordinator=False, + ) -> None: + if ( + type(state_dict) is not dict + or set(state_dict) != {namespace} + or type(state_dict[namespace]) is not dict + or state_dict[namespace] + ): + raise ValueError("portable DCP load planner requires one empty exact namespace") + allocated = _allocate_dcp_state( + metadata, + namespace=namespace, + limits=limits, + ) + dict.update(state_dict[namespace], allocated[namespace]) + super().set_up_planner(state_dict, metadata, is_coordinator) + + return _PortableDCPDynamicLoadPlanner() + + +def _decode_dcp_state( + state, + *, + namespace: str, + limits: _CanonicalWireLimits, +): + if type(state) is not dict or set(state) != {namespace}: + raise ValueError("portable DCP root state has an invalid schema") + envelope = state[namespace] + if type(envelope) is not dict or _DCP_METADATA_KEY not in envelope: + raise ValueError("portable DCP envelope has an invalid schema") + metadata_tensor = envelope[_DCP_METADATA_KEY] + if ( + type(metadata_tensor) is not torch.Tensor + or metadata_tensor.dtype is not torch.uint8 + or metadata_tensor.device.type != "cpu" + or metadata_tensor.ndim != 1 + or not metadata_tensor.is_contiguous() + or metadata_tensor.numel() > limits.max_metadata_bytes + ): + raise ValueError("portable DCP wire metadata tensor is invalid") + metadata = bytes(memoryview(metadata_tensor.numpy())) + prepared = _parse_canonical_wire_metadata(metadata, limits=limits) + expected_keys = { + _DCP_METADATA_KEY, + *(_payload_key(index) for index in range(len(prepared.tensor_specs))), + } + if set(envelope) != expected_keys: + raise ValueError("portable DCP payload key count disagrees with wire metadata") + payloads = tuple(envelope[_payload_key(index)] for index in range(len(prepared.tensor_specs))) + document = _reconstruct_canonical_wire_value(prepared, payloads) + return document, prepared.fragment_digest + + +def save_portable_dcp( + optimizer, + *, + checkpoint_process_group, + storage_writer, + transaction_id, + limits, + namespace="optimizer", +): + """Collectively save one tensor-only portable v3 document through DCP.""" + + import torch.distributed.checkpoint as dcp + from torch.distributed.checkpoint.storage import StorageWriter + + binding, transaction_id, limits, namespace = _preflight_dcp_operation( + optimizer, + checkpoint_process_group=checkpoint_process_group, + transaction_id=transaction_id, + limits=limits, + namespace=namespace, + storage=storage_writer, + storage_type=StorageWriter, + operation="save", + ) + document = optimizer.export_portable_state( + checkpoint_process_group=binding, + transaction_id=transaction_id, + limits=limits, + ) + wire_limits = limits._wire_limits(collective=True) + plan = None + state = None + digest = bytes(32) + error = None + try: + plan = _prepare_canonical_wire_value(document, wire_limits) + state = _state_from_plan(namespace, plan, wire_limits) + digest = plan.fragment_digest + except Exception as exc: + error = exc + from gefen.portable_collective import _collective_unanimous_status + + _collective_unanimous_status( + binding, + error, + operation="portable_dcp_save_encode", + transaction_id=transaction_id, + context_digest=digest, + limits=wire_limits, + ) + assert plan is not None and state is not None + return dcp.save( + state, + storage_writer=storage_writer, + process_group=binding.process_group, + ) + + +def load_portable_dcp( + optimizer, + *, + checkpoint_process_group, + storage_reader, + transaction_id, + limits, + namespace="optimizer", +) -> None: + """Collectively load, verify, and atomically import portable v3 from DCP.""" + + import torch.distributed.checkpoint as dcp + from torch.distributed.checkpoint.storage import StorageReader + + binding, transaction_id, limits, namespace = _preflight_dcp_operation( + optimizer, + checkpoint_process_group=checkpoint_process_group, + transaction_id=transaction_id, + limits=limits, + namespace=namespace, + storage=storage_reader, + storage_type=StorageReader, + operation="load", + ) + wire_limits = limits._wire_limits(collective=True) + state = {namespace: {}} + dcp.load( + state, + storage_reader=storage_reader, + planner=_load_planner(namespace, wire_limits), + process_group=binding.process_group, + ) + document = None + digest = bytes(32) + error = None + try: + document, digest = _decode_dcp_state( + state, + namespace=namespace, + limits=wire_limits, + ) + except Exception as exc: + error = exc + from gefen.portable_collective import _collective_unanimous_status + + _collective_unanimous_status( + binding, + error, + operation="portable_dcp_load_decode", + transaction_id=transaction_id, + context_digest=digest, + limits=wire_limits, + ) + assert document is not None + optimizer.import_portable_state( + document, + checkpoint_process_group=binding, + transaction_id=transaction_id, + limits=limits, + ) + + +__all__ = ["load_portable_dcp", "save_portable_dcp"] diff --git a/tests/test_portable_dcp.py b/tests/test_portable_dcp.py new file mode 100644 index 0000000..ac0fd8a --- /dev/null +++ b/tests/test_portable_dcp.py @@ -0,0 +1,637 @@ +"""Warning-strict CPU coverage for tensor-only portable DCP persistence.""" + +from __future__ import annotations + +from types import SimpleNamespace +import warnings + +import pytest +import torch +import torch.distributed.checkpoint as dcp +from torch.distributed.checkpoint.api import CheckpointException +from torch.distributed.checkpoint.metadata import ( + BytesStorageMetadata, + ChunkStorageMetadata, + Metadata, + TensorProperties, + TensorStorageMetadata, +) + +import gefen +import gefen.portable_dcp as portable_dcp +from gefen import load_portable_dcp, save_portable_dcp +from gefen.checkpoint import CheckpointProcessGroupBinding +from gefen.codebook import CodebookProcessGroupBinding +from gefen.contracts import ( + LogicalSlice, + ParameterIdentity, + ParameterLayout, + PlacementKind, + ProcessGroupIdentity, + ShardIdentity, + ShardPlacement, + ShardingManifest, +) +from gefen.gefen import Gefen +from gefen.gefen_muon import GefenMuon +from gefen.portable import _decode_quantized_momentum +from gefen.portable_state import PortableStateLimits +from gefen.rebinding import ParameterRebinding + + +_MEMBER = "rank:0" +_SINGLE_PROCESS_DCP_WARNING = ( + r"^torch\.distributed is (?:disabled, )?unavailable or uninitialized, " + r"assuming the intent is to (?:save|load) in a single process\.$" +) +_TYPED_STORAGE_DCP_WARNING = ( + r"^TypedStorage is deprecated\. It will be removed in the future and UntypedStorage will be the only storage class\. " + r"This should only matter to you if you are using storages directly\. To access UntypedStorage directly, use " + r"tensor\.untyped_storage\(\) instead of tensor\.storage\(\)$" +) + + +class _TensorPropertiesSubclass(TensorProperties): + pass + + +@pytest.fixture(autouse=True) +def _warning_strict(): + with warnings.catch_warnings(): + warnings.simplefilter("error") + warnings.filterwarnings( + "ignore", + message=_SINGLE_PROCESS_DCP_WARNING, + category=UserWarning, + ) + warnings.filterwarnings( + "ignore", + message=_TYPED_STORAGE_DCP_WARNING, + category=UserWarning, + ) + yield + + +def _limits( + *, + max_fragment_tensor_bytes=4 << 20, + max_metadata_bytes=64 << 20, +): + return PortableStateLimits( + max_fragment_tensor_bytes=max_fragment_tensor_bytes, + max_collective_tensor_bytes=16 << 20, + max_collective_metadata_bytes=64 << 20, + max_metadata_bytes=max_metadata_bytes, + ) + + +def _bindings(group): + return ( + CodebookProcessGroupBinding( + group, + _MEMBER, + None, + torch.device("cpu"), + ), + CheckpointProcessGroupBinding( + group, + _MEMBER, + None, + torch.device("cpu"), + ), + ) + + +def _finalize(optimizer, parameter, *, layout): + group = ProcessGroupIdentity("checkpoint", (_MEMBER,)) + identity = ParameterIdentity("layer.weight", tuple(parameter.shape)) + if layout is ParameterLayout.REPLICATED: + kind = PlacementKind.REPLICATE + owner = None + else: + kind = PlacementKind.WHOLE_PARAMETER_OWNER + owner = _MEMBER + shard = ShardIdentity( + identity, + layout, + LogicalSlice.full(identity), + placements=(ShardPlacement("checkpoint", kind, 0, 1),), + process_group=group, + local_member=_MEMBER, + owner=owner, + ) + codebook_binding, checkpoint_binding = _bindings(group) + optimizer.post_sharding( + (ParameterRebinding(parameter, parameter, shard),), + manifest=ShardingManifest((shard,)), + codebook_process_group=codebook_binding, + ) + return checkpoint_binding + + +def _plain(*, factored, deterministic): + parameter = torch.nn.Parameter(torch.arange(1, 7, dtype=torch.float32).reshape(2, 3)) + optimizer = Gefen( + [("weight", parameter)], + fused=False, + factored_v_2d=factored, + force_2d_period_one=True, + deterministic=deterministic, + ) + binding = _finalize( + optimizer, + parameter, + layout=ParameterLayout.REPLICATED, + ) + return optimizer, parameter, binding + + +def _muon(*, deterministic): + parameter = torch.nn.Parameter(torch.arange(1, 7, dtype=torch.float32).reshape(2, 3)) + optimizer = GefenMuon( + [("weight", parameter)], + fused=False, + sharded_mode="distributed", + normuon=True, + deterministic=deterministic, + ) + binding = _finalize( + optimizer, + parameter, + layout=ParameterLayout.WHOLE_PARAMETER_OWNER, + ) + return optimizer, parameter, binding + + +def _optimizer_pair(variant): + if variant == "muon-normuon": + source, source_parameter, source_binding = _muon(deterministic=True) + target, target_parameter, target_binding = _muon(deterministic=False) + else: + factored = variant == "plain-factored" + source, source_parameter, source_binding = _plain( + factored=factored, + deterministic=True, + ) + target, target_parameter, target_binding = _plain( + factored=factored, + deterministic=False, + ) + return ( + source, + source_parameter, + source_binding, + target, + target_parameter, + target_binding, + ) + + +def _initialize(optimizer, parameter, *, variant): + optimizer._gefen_global_step = 4 + optimizer._gefen_codebook = torch.linspace( + -1.0, + 1.0, + 256, + dtype=torch.float32, + ) + state = optimizer.state[parameter] + state.update( + { + "automatic_period": 1, + "step": 4, + "m_codebook": torch.tensor( + [[0], [255], [128], [64], [192], [0]], + dtype=torch.uint8, + ), + "m_magnitude": torch.tensor( + [[1.0], [2.0], [3.0], [4.0], [5.0], [6.0]], + dtype=torch.float32, + ), + } + ) + if variant == "plain-block": + state.update( + { + "vmean": torch.arange(1, 7, dtype=torch.float32).reshape(6, 1), + "vmean_step": 3, + } + ) + elif variant == "plain-factored": + state.update( + { + "v_row": torch.tensor([2.0, 3.0], dtype=torch.float32), + "v_col": torch.tensor([5.0, 7.0, 11.0], dtype=torch.float32), + "factored_step": 3, + } + ) + elif variant == "muon-normuon": + state.update( + { + "normuon_v": torch.tensor([[2.0], [3.0]], dtype=torch.float32), + "normuon_step": 2, + } + ) + else: + raise AssertionError("unknown test variant") + + +def _assert_loaded_state(target, target_parameter, expected, *, variant): + assert target._gefen_global_step == 4 + assert target._deterministic is True + assert torch.equal( + target._gefen_codebook, + expected["common"]["gefen_codebook"], + ) + target_state = target.state[target_parameter] + target_momentum = _decode_quantized_momentum( + target._gefen_codebook, + target_state["m_codebook"], + target_state["m_magnitude"], + logical_shape=tuple(target_parameter.shape), + period=1, + step=target_state["step"], + ) + record = expected["parameters"]["layer.weight"] + assert torch.equal( + target_momentum.view(torch.int32), + record["state"]["momentum"].view(torch.int32), + ) + if variant == "plain-block": + assert torch.equal( + target_state["vmean"].reshape(2, 3), + record["state"]["second_moment"], + ) + assert target_state["vmean_step"] == 3 + elif variant == "plain-factored": + assert torch.equal(target_state["v_row"], record["state"]["v_row"]) + assert torch.equal(target_state["v_col"], record["state"]["v_col"]) + assert target_state["factored_step"] == 3 + else: + assert torch.equal( + target_state["normuon_v"], + record["state"]["normuon_v"], + ) + assert target_state["normuon_step"] == 2 + + +@pytest.mark.parametrize( + "variant", + ["plain-block", "plain-factored", "muon-normuon"], +) +def test_filesystem_dcp_round_trip_is_tensor_only(tmp_path, variant): + ( + source, + source_parameter, + source_binding, + target, + target_parameter, + target_binding, + ) = _optimizer_pair(variant) + _initialize(source, source_parameter, variant=variant) + limits = _limits() + expected = source.export_portable_state( + checkpoint_process_group=source_binding, + transaction_id="dcp-expected-{}".format(variant), + limits=limits, + ) + checkpoint = tmp_path / variant + + metadata = save_portable_dcp( + source, + checkpoint_process_group=source_binding, + storage_writer=dcp.FileSystemWriter(checkpoint), + transaction_id="dcp-save-{}".format(variant), + limits=limits, + ) + + assert type(metadata) is Metadata + assert metadata.state_dict_metadata + assert all(type(value) is TensorStorageMetadata for value in metadata.state_dict_metadata.values()) + assert not any(type(value) is BytesStorageMetadata for value in metadata.state_dict_metadata.values()) + on_disk_metadata = dcp.FileSystemReader(checkpoint).read_metadata() + assert set(on_disk_metadata.state_dict_metadata) == set(metadata.state_dict_metadata) + assert all(type(value) is TensorStorageMetadata for value in on_disk_metadata.state_dict_metadata.values()) + + load_portable_dcp( + target, + checkpoint_process_group=target_binding, + storage_reader=dcp.FileSystemReader(checkpoint), + transaction_id="dcp-load-{}".format(variant), + limits=limits, + ) + + _assert_loaded_state( + target, + target_parameter, + expected, + variant=variant, + ) + + +@pytest.mark.parametrize( + "namespace", + ["", " optimizer", "optimizer.", "optim\x00izer", "optim/izer", "optimé"], +) +def test_namespace_validation_precedes_dcp_io(tmp_path, namespace): + optimizer, parameter, binding = _plain( + factored=False, + deterministic=True, + ) + _initialize(optimizer, parameter, variant="plain-block") + checkpoint = tmp_path / "invalid-namespace" + + with pytest.raises(RuntimeError, match="namespace"): + save_portable_dcp( + optimizer, + checkpoint_process_group=binding, + storage_writer=dcp.FileSystemWriter(checkpoint), + transaction_id="dcp-invalid-namespace", + limits=_limits(), + namespace=namespace, + ) + + assert not (checkpoint / ".metadata").exists() + + +@pytest.mark.parametrize( + ("members", "world_size", "group_zero", "match"), + [ + (("rank:0",), 2, 0, "singleton"), + (("rank:1", "rank:2"), 3, 1, "global rank zero"), + ], +) +def test_dcp_runtime_rejects_unsafe_default_world_or_subgroup( + monkeypatch, + members, + world_size, + group_zero, + match, +): + binding = SimpleNamespace( + identity=SimpleNamespace(ordered_members=members), + process_group=object(), + collective_device=torch.device("cpu"), + ) + monkeypatch.setattr(torch.distributed, "is_available", lambda: True) + monkeypatch.setattr(torch.distributed, "is_initialized", lambda: True) + monkeypatch.setattr(torch.distributed, "get_world_size", lambda: world_size) + monkeypatch.setattr( + torch.distributed, + "get_global_rank", + lambda _group, _rank: group_zero, + ) + + with pytest.raises(RuntimeError, match=match): + portable_dcp._validate_dcp_runtime(binding) + + +def test_load_rejects_namespace_mismatch_without_target_mutation(tmp_path): + ( + source, + source_parameter, + source_binding, + target, + _, + target_binding, + ) = _optimizer_pair("plain-block") + _initialize(source, source_parameter, variant="plain-block") + checkpoint = tmp_path / "namespace-mismatch" + save_portable_dcp( + source, + checkpoint_process_group=source_binding, + storage_writer=dcp.FileSystemWriter(checkpoint), + transaction_id="dcp-save-namespace-mismatch", + limits=_limits(), + namespace="optimizer", + ) + before = target._canonical_import_live_token() + + with pytest.raises(CheckpointException): + load_portable_dcp( + target, + checkpoint_process_group=target_binding, + storage_reader=dcp.FileSystemReader(checkpoint), + transaction_id="dcp-load-namespace-mismatch", + limits=_limits(), + namespace="other", + ) + + assert target._canonical_import_live_token() == before + + +def _tensor_metadata(dtype, shape, *, offsets=None, sizes=None): + shape = torch.Size(shape) + if offsets is None: + offsets = (0,) * len(shape) + if sizes is None: + sizes = shape + return TensorStorageMetadata( + properties=TensorProperties(dtype=dtype), + size=shape, + chunks=[ + ChunkStorageMetadata( + offsets=torch.Size(offsets), + sizes=torch.Size(sizes), + ) + ], + ) + + +def _planner_metadata(case, *, namespace, limits): + metadata_key = portable_dcp._flat_key( + namespace, + portable_dcp._DCP_METADATA_KEY, + ) + payload_key = portable_dcp._flat_key( + namespace, + portable_dcp._payload_key(0), + ) + entries = {metadata_key: _tensor_metadata(torch.uint8, (1,))} + if case == "oversized-wire-metadata": + entries[metadata_key] = _tensor_metadata( + torch.uint8, + (limits.max_metadata_bytes + 1,), + ) + elif case == "oversized-payload": + entries[payload_key] = _tensor_metadata( + torch.uint8, + (limits.max_fragment_tensor_bytes + 1,), + ) + elif case == "extra-key": + entries[portable_dcp._flat_key(namespace, "extra")] = _tensor_metadata( + torch.uint8, + (1,), + ) + elif case == "byte-io": + entries[metadata_key] = BytesStorageMetadata() + elif case == "partial-tensor": + entries[metadata_key] = _tensor_metadata( + torch.uint8, + (8,), + offsets=(0,), + sizes=(4,), + ) + elif case == "partial-planner-data": + entries[payload_key] = _tensor_metadata(torch.float32, (1,)) + elif case == "properties-subclass": + entry = _tensor_metadata(torch.uint8, (1,)) + entry.properties = _TensorPropertiesSubclass(dtype=torch.uint8) + entries[metadata_key] = entry + elif case == "list-size": + entry = _tensor_metadata(torch.uint8, (1,)) + entry.size = [1] + entries[metadata_key] = entry + elif case == "list-chunk-offsets": + entry = _tensor_metadata(torch.uint8, (1,)) + entry.chunks[0].offsets = [0] + entries[metadata_key] = entry + else: + raise AssertionError("unknown planner metadata case") + planner_data = {key: tuple(key.split(".", 1)) for key in entries} + if case == "partial-planner-data": + planner_data.pop(payload_key) + return Metadata( + state_dict_metadata=entries, + planner_data=planner_data, + ) + + +@pytest.mark.parametrize( + ("case", "error_type", "match"), + [ + ("oversized-wire-metadata", ValueError, "metadata exceeds"), + ("oversized-payload", ValueError, "payload tensors exceed"), + ("extra-key", ValueError, "unexpected tensor keys"), + ("byte-io", TypeError, "full-tensor DCP metadata"), + ("partial-tensor", ValueError, "complete unsharded tensor"), + ("partial-planner-data", ValueError, "incompatible planner paths"), + ("properties-subclass", TypeError, "unsupported tensor properties"), + ("list-size", TypeError, "unsupported tensor properties"), + ("list-chunk-offsets", ValueError, "complete unsharded tensor"), + ], +) +def test_load_planner_rejects_untrusted_metadata_before_allocation( + monkeypatch, + case, + error_type, + match, +): + namespace = "optimizer" + limits = _limits( + max_fragment_tensor_bytes=64, + max_metadata_bytes=64, + )._wire_limits(collective=True) + metadata = _planner_metadata( + case, + namespace=namespace, + limits=limits, + ) + planner = portable_dcp._load_planner(namespace, limits) + allocations = [] + + def fail_if_allocated(*args, **kwargs): + allocations.append((args, kwargs)) + raise AssertionError("untrusted DCP metadata reached tensor allocation") + + monkeypatch.setattr(torch, "empty", fail_if_allocated) + + with pytest.raises(error_type, match=match): + planner.set_up_planner( + {namespace: {}}, + metadata, + is_coordinator=True, + ) + + assert allocations == [] + + +def test_dcp_envelope_container_limit_is_symmetric_before_save(): + limits = PortableStateLimits( + max_fragment_tensor_bytes=64, + max_collective_tensor_bytes=64, + max_collective_metadata_bytes=4096, + max_metadata_bytes=4096, + max_container_items=3, + max_tensors=8, + )._wire_limits(collective=True) + plan = portable_dcp._prepare_canonical_wire_value( + ( + torch.ones(1), + torch.ones(1), + torch.ones(1), + ), + limits, + ) + + with pytest.raises(ValueError, match="max_container_items"): + portable_dcp._state_from_plan("optimizer", plan, limits) + + +def test_corrupted_tensor_payload_leaves_target_unchanged(tmp_path): + ( + source, + source_parameter, + source_binding, + target, + target_parameter, + target_binding, + ) = _optimizer_pair("plain-block") + _initialize(source, source_parameter, variant="plain-block") + _initialize(target, target_parameter, variant="plain-block") + target._deterministic = False + target._gefen_global_step = 9 + limits = _limits() + original = tmp_path / "original" + corrupted = tmp_path / "corrupted" + save_portable_dcp( + source, + checkpoint_process_group=source_binding, + storage_writer=dcp.FileSystemWriter(original), + transaction_id="dcp-save-corruption-source", + limits=limits, + ) + + wire_limits = limits._wire_limits(collective=True) + state = {"optimizer": {}} + dcp.load( + state, + storage_reader=dcp.FileSystemReader(original), + planner=portable_dcp._load_planner("optimizer", wire_limits), + process_group=None, + ) + payload_keys = sorted(key for key in state["optimizer"] if key.startswith(portable_dcp._DCP_PAYLOAD_PREFIX)) + assert payload_keys + payload = next(state["optimizer"][key] for key in payload_keys if state["optimizer"][key].numel()) + raw = payload.view(torch.uint8).reshape(-1) + raw[0] = int(raw[0]) ^ 1 + dcp.save( + state, + storage_writer=dcp.FileSystemWriter(corrupted), + process_group=None, + ) + before = target._canonical_import_live_token() + + with pytest.raises(RuntimeError, match="invalid digest"): + load_portable_dcp( + target, + checkpoint_process_group=target_binding, + storage_reader=dcp.FileSystemReader(corrupted), + transaction_id="dcp-load-corrupted-payload", + limits=limits, + ) + + assert target._canonical_import_live_token() == before + + +def test_portable_dcp_helpers_are_public_exports(): + assert gefen.save_portable_dcp is portable_dcp.save_portable_dcp + assert gefen.load_portable_dcp is portable_dcp.load_portable_dcp + assert save_portable_dcp is portable_dcp.save_portable_dcp + assert load_portable_dcp is portable_dcp.load_portable_dcp + assert "save_portable_dcp" in gefen.__all__ + assert "load_portable_dcp" in gefen.__all__ + assert portable_dcp.__all__ == [ + "load_portable_dcp", + "save_portable_dcp", + ] diff --git a/tests/test_portable_dcp_distributed.py b/tests/test_portable_dcp_distributed.py new file mode 100644 index 0000000..8c44e84 --- /dev/null +++ b/tests/test_portable_dcp_distributed.py @@ -0,0 +1,530 @@ +"""CPU-only distributed coverage for portable optimizer DCP persistence.""" + +from datetime import timedelta +import multiprocessing as mp +import os +import queue as queue_module +import tempfile +import traceback +import warnings + +import pytest +import torch +import torch.distributed as dist +from torch.distributed.checkpoint import FileSystemReader, FileSystemWriter + +from gefen import load_portable_dcp, save_portable_dcp +from gefen.checkpoint import CheckpointProcessGroupBinding +from gefen.codebook import CodebookProcessGroupBinding +from gefen.contracts import ( + LogicalSlice, + ParameterIdentity, + ParameterLayout, + PlacementKind, + ProcessGroupIdentity, + ShardIdentity, + ShardPlacement, + ShardingManifest, +) +from gefen.gefen import Gefen +from gefen.portable import _decode_quantized_momentum +from gefen.portable_state import PortableStateLimits +from gefen.rebinding import ParameterRebinding + + +_SAVE_WORLD = 2 +_FQN = "model.weight" + + +def _limits(): + return PortableStateLimits( + max_fragment_tensor_bytes=1 << 20, + max_collective_tensor_bytes=4 << 20, + max_collective_metadata_bytes=4 << 20, + chunk_bytes=11, + max_members=4, + max_metadata_bytes=1 << 20, + max_tree_nodes=10_000, + max_tree_depth=32, + max_container_items=10_000, + max_string_bytes=16 << 10, + max_integer_bytes=128, + max_tensors=256, + max_tensor_rank=8, + diagnostic_bytes=1024, + ) + + +def _save_members(): + return tuple("save:{}".format(rank) for rank in range(_SAVE_WORLD)) + + +def _optimizer(parameter, *, deterministic): + return Gefen( + [("weight", parameter)], + lr=2.5e-3, + betas=(0.8, 0.97), + eps=2.0e-8, + weight_decay=0.03, + fused=False, + force_2d_period_one=True, + factored_v_2d=False, + deterministic=deterministic, + ) + + +def _bindings(identity, local_member, process_group): + return ( + CodebookProcessGroupBinding( + identity, + local_member, + process_group, + torch.device("cpu"), + ), + CheckpointProcessGroupBinding( + identity, + local_member, + process_group, + torch.device("cpu"), + ), + ) + + +def _flat_manifest(identity, group, lengths): + if len(lengths) != len(group.ordered_members) or sum(lengths) != identity.numel: + raise AssertionError("invalid flat test partition") + shards = [] + offset = 0 + for coordinate, (member, length) in enumerate(zip(group.ordered_members, lengths)): + shards.append( + ShardIdentity( + identity, + ParameterLayout.FLATTENED_ELEMENT_SHARD, + LogicalSlice(offset, length), + placements=( + ShardPlacement( + "checkpoint", + PlacementKind.FLAT_SHARD, + coordinate, + len(group.ordered_members), + ), + ), + process_group=group, + local_member=member, + ) + ) + offset += length + return ShardingManifest(tuple(shards)), tuple(shards) + + +def _replicated_manifest(identity, group): + shard = ShardIdentity( + identity, + ParameterLayout.REPLICATED, + LogicalSlice.full(identity), + placements=( + ShardPlacement( + "checkpoint", + PlacementKind.REPLICATE, + 0, + 1, + ), + ), + process_group=group, + local_member=group.ordered_members[0], + ) + return ShardingManifest((shard,)), shard + + +def _make_flat_source(rank, group): + identity = ParameterIdentity(_FQN, (2, 4)) + lengths = (3, 5) + original = torch.nn.Parameter(torch.zeros(identity.global_shape, dtype=torch.float32)) + local = torch.nn.Parameter(torch.zeros(lengths[rank], dtype=torch.float32)) + optimizer = _optimizer(original, deterministic=True) + manifest, shards = _flat_manifest(identity, group, lengths) + codebook_binding, checkpoint_binding = _bindings(group, _save_members()[rank], dist.group.WORLD) + optimizer.post_sharding( + (ParameterRebinding(original, local, shards[rank]),), + manifest=manifest, + codebook_process_group=codebook_binding, + ) + return optimizer, local, shards[rank], checkpoint_binding + + +def _make_replicated_target(initial_parameter, *, deterministic): + identity = ParameterIdentity(_FQN, (2, 4)) + group = ProcessGroupIdentity("portable_dcp_singleton", ("load:0",)) + parameter = torch.nn.Parameter(initial_parameter.clone()) + optimizer = _optimizer(parameter, deterministic=deterministic) + manifest, shard = _replicated_manifest(identity, group) + codebook_binding, checkpoint_binding = _bindings(group, "load:0", None) + optimizer.post_sharding( + (ParameterRebinding(parameter, parameter, shard),), + manifest=manifest, + codebook_process_group=codebook_binding, + ) + return optimizer, parameter, checkpoint_binding + + +def _codebook(): + return torch.linspace(-1.0, 1.0, 256, dtype=torch.float32) + + +def _momentum_values(): + return ( + torch.tensor( + [ + -2147483648, + 0, + 1, + -2147483647, + 1056964608, + -1090519040, + 1078984704, + -1058013184, + ], + dtype=torch.int32, + ) + .view(torch.float32) + .reshape(2, 4) + .clone() + ) + + +def _second_moment_values(): + return ( + torch.tensor( + [ + -2147483648, + 0, + 1, + 8388608, + 1048576000, + 1065353216, + 1073741824, + 2139095039, + ], + dtype=torch.int32, + ) + .view(torch.float32) + .reshape(2, 4) + .clone() + ) + + +def _quantized_period_one(momentum): + flat = momentum.reshape(-1) + indices = torch.where( + torch.signbit(flat), + torch.zeros(flat.numel(), dtype=torch.uint8), + torch.full((flat.numel(),), 255, dtype=torch.uint8), + ) + return indices.reshape(-1, 1), flat.abs().reshape(-1, 1).clone() + + +def _seed_state(optimizer, parameter, momentum, second_moment): + indices, magnitudes = _quantized_period_one(momentum) + optimizer._gefen_global_step = 9 + optimizer._gefen_codebook = _codebook() + optimizer.state[parameter].update( + { + "automatic_period": 1, + "step": 7, + "m_codebook": indices, + "m_magnitude": magnitudes, + "vmean": second_moment.reshape(-1, 1).clone(), + "vmean_step": 6, + } + ) + + +def _bits_equal(left, right): + return ( + type(left) is torch.Tensor + and type(right) is torch.Tensor + and left.dtype == right.dtype + and tuple(left.shape) == tuple(right.shape) + and torch.equal(left.contiguous().view(torch.uint8), right.contiguous().view(torch.uint8)) + ) + + +def _capture_runtime_error(operation): + try: + operation() + except RuntimeError as exc: + return str(exc) + raise AssertionError("portable DCP operation unexpectedly succeeded") + + +def _strict_worker_warnings(): + warnings.simplefilter("error") + warnings.filterwarnings( + "ignore", + message="TypedStorage is deprecated.*", + category=UserWarning, + ) + + +def _save_worker( + rank, + init_file, + checkpoint_dir, + mismatch_dir, + divergent_dir, + invalid_dir, + result_queue, +): + _strict_worker_warnings() + try: + dist.init_process_group( + "gloo", + init_method="file://{}".format(init_file), + rank=rank, + world_size=_SAVE_WORLD, + timeout=timedelta(seconds=60), + ) + group = ProcessGroupIdentity("portable_dcp_save", _save_members()) + optimizer, parameter, shard, checkpoint_binding = _make_flat_source(rank, group) + momentum = _momentum_values().reshape(-1) + second_moment = _second_moment_values().reshape(-1) + start = shard.logical_slice.flat_offset + stop = start + shard.logical_slice.length + _seed_state( + optimizer, + parameter, + momentum[start:stop].clone(), + second_moment[start:stop].clone(), + ) + + save_portable_dcp( + optimizer, + checkpoint_process_group=checkpoint_binding, + storage_writer=FileSystemWriter(checkpoint_dir), + transaction_id="portable-dcp-save-2-to-1", + limits=_limits(), + namespace="optimizer", + ) + dist.barrier() + + token = optimizer._canonical_import_live_token() + namespace_message = _capture_runtime_error( + lambda: save_portable_dcp( + optimizer, + checkpoint_process_group=checkpoint_binding, + storage_writer=FileSystemWriter(mismatch_dir), + transaction_id="portable-dcp-mismatched-namespace", + limits=_limits(), + namespace="optimizer-a" if rank == 0 else "optimizer-b", + ) + ) + namespace_unchanged = optimizer._canonical_import_live_token() == token + dist.barrier() + + token = optimizer._canonical_import_live_token() + storage_message = _capture_runtime_error( + lambda: save_portable_dcp( + optimizer, + checkpoint_process_group=checkpoint_binding, + storage_writer=FileSystemWriter("{}-{}".format(divergent_dir, rank)), + transaction_id="portable-dcp-mismatched-storage", + limits=_limits(), + namespace="optimizer", + ) + ) + storage_unchanged = optimizer._canonical_import_live_token() == token + dist.barrier() + + token = optimizer._canonical_import_live_token() + storage = object() if rank == 0 else FileSystemWriter(invalid_dir) + invalid_message = _capture_runtime_error( + lambda: save_portable_dcp( + optimizer, + checkpoint_process_group=checkpoint_binding, + storage_writer=storage, + transaction_id="portable-dcp-invalid-writer", + limits=_limits(), + namespace="optimizer", + ) + ) + invalid_unchanged = optimizer._canonical_import_live_token() == token + dist.barrier() + + result_queue.put( + { + "rank": rank, + "saved": True, + "namespace_message": namespace_message, + "namespace_unchanged": namespace_unchanged, + "storage_message": storage_message, + "storage_unchanged": storage_unchanged, + "invalid_message": invalid_message, + "invalid_unchanged": invalid_unchanged, + } + ) + except BaseException: + result_queue.put({"rank": rank, "fatal_error": traceback.format_exc()}) + finally: + if dist.is_available() and dist.is_initialized(): + dist.destroy_process_group() + + +def _loaded_state_is_exact(optimizer, parameter): + state = optimizer.state[parameter] + decoded = _decode_quantized_momentum( + optimizer._gefen_codebook, + state["m_codebook"], + state["m_magnitude"], + logical_shape=(2, 4), + period=1, + step=state["step"], + ) + return ( + _bits_equal(decoded, _momentum_values()) + and _bits_equal(state["vmean"].reshape(2, 4), _second_moment_values()) + and _bits_equal(optimizer._gefen_codebook, _codebook()) + and optimizer._gefen_global_step == 9 + and state["automatic_period"] == 1 + and state["step"] == 7 + and state["vmean_step"] == 6 + and optimizer._deterministic is True + ) + + +def _next_step_is_exact(optimizer, parameter, initial_parameter): + oracle, oracle_parameter, _ = _make_replicated_target(initial_parameter, deterministic=True) + _seed_state(oracle, oracle_parameter, _momentum_values(), _second_moment_values()) + gradient = torch.tensor( + [0.25, -0.5, 0.75, -1.0, 1.25, -1.5, 1.75, -2.0], + dtype=torch.float32, + ).reshape(2, 4) + parameter.grad = gradient.clone() + oracle_parameter.grad = gradient.clone() + optimizer.step() + oracle.step() + state = optimizer.state[parameter] + oracle_state = oracle.state[oracle_parameter] + return ( + _bits_equal(parameter.detach(), oracle_parameter.detach()) + and optimizer._gefen_global_step == oracle._gefen_global_step == 10 + and state["step"] == oracle_state["step"] == 8 + and state["vmean_step"] == oracle_state["vmean_step"] == 7 + and all(_bits_equal(state[key], oracle_state[key]) for key in ("m_codebook", "m_magnitude", "vmean")) + ) + + +def _load_worker(rank, init_file, checkpoint_dir, result_queue): + _strict_worker_warnings() + try: + dist.init_process_group( + "gloo", + init_method="file://{}".format(init_file), + rank=rank, + world_size=1, + timeout=timedelta(seconds=60), + ) + initial_parameter = torch.linspace(-0.4, 0.3, 8, dtype=torch.float32).reshape(2, 4) + optimizer, parameter, checkpoint_binding = _make_replicated_target( + initial_parameter, + deterministic=False, + ) + load_portable_dcp( + optimizer, + checkpoint_process_group=checkpoint_binding, + storage_reader=FileSystemReader(checkpoint_dir), + transaction_id="portable-dcp-load-2-to-1", + limits=_limits(), + namespace="optimizer", + ) + restored_exact = _loaded_state_is_exact(optimizer, parameter) + next_step_exact = _next_step_is_exact(optimizer, parameter, initial_parameter) + result_queue.put( + { + "rank": rank, + "restored_exact": restored_exact, + "next_step_exact": next_step_exact, + } + ) + except BaseException: + result_queue.put({"rank": rank, "fatal_error": traceback.format_exc()}) + finally: + if dist.is_available() and dist.is_initialized(): + dist.destroy_process_group() + + +def _run_phase(worker, world_size, *worker_args): + context = mp.get_context("spawn") + result_queue = context.Queue() + descriptor, init_file = tempfile.mkstemp(prefix="gefen-portable-dcp-") + os.close(descriptor) + os.unlink(init_file) + processes = [ + context.Process( + target=worker, + args=(rank, init_file, *worker_args, result_queue), + ) + for rank in range(world_size) + ] + results = [] + try: + for process in processes: + process.start() + try: + for _ in processes: + results.append(result_queue.get(timeout=120)) + except queue_module.Empty: + pass + for process in processes: + process.join(timeout=10) + for process in processes: + if process.is_alive(): + process.terminate() + process.join(timeout=5) + finally: + for process in processes: + if process.is_alive(): + process.terminate() + process.join(timeout=5) + result_queue.close() + result_queue.join_thread() + if os.path.exists(init_file): + os.unlink(init_file) + assert len(results) == world_size, (results, [process.exitcode for process in processes]) + assert all(process.exitcode == 0 for process in processes), [process.exitcode for process in processes] + return sorted(results, key=lambda item: item["rank"]) + + +@pytest.mark.skipif( + not dist.is_available() or not dist.is_gloo_available(), + reason="portable DCP topology-change coverage requires Gloo", +) +def test_portable_dcp_saves_flattened_world_and_loads_replicated_singleton(tmp_path): + checkpoint_dir = str(tmp_path / "checkpoint") + save_results = _run_phase( + _save_worker, + _SAVE_WORLD, + checkpoint_dir, + str(tmp_path / "mismatched-namespace"), + str(tmp_path / "mismatched-storage"), + str(tmp_path / "invalid-writer"), + ) + + assert all("fatal_error" not in result for result in save_results), save_results + assert all(result["saved"] for result in save_results) + namespace_messages = [result["namespace_message"] for result in save_results] + assert namespace_messages[0] == namespace_messages[1] + assert "context" in namespace_messages[0] + storage_messages = [result["storage_message"] for result in save_results] + assert storage_messages[0] == storage_messages[1] + assert "context" in storage_messages[0] + invalid_messages = [result["invalid_message"] for result in save_results] + assert invalid_messages[0] == invalid_messages[1] + assert "StorageWriter" in invalid_messages[0] + assert all( + result["namespace_unchanged"] and result["storage_unchanged"] and result["invalid_unchanged"] + for result in save_results + ) + + load_results = _run_phase(_load_worker, 1, checkpoint_dir) + assert all("fatal_error" not in result for result in load_results), load_results + assert load_results == [{"rank": 0, "restored_exact": True, "next_step_exact": True}] diff --git a/tests/test_portable_schema.py b/tests/test_portable_schema.py index e8ba823..535d109 100644 --- a/tests/test_portable_schema.py +++ b/tests/test_portable_schema.py @@ -284,4 +284,6 @@ def test_portable_schema_and_checkpoint_binding_exports_are_public(): assert gefen.CheckpointProcessGroupBinding.__module__ == "gefen.checkpoint" assert gefen.PortableStateLimits.__module__ == "gefen.portable_state" assert gefen.PortableStateProvider.__module__ == "gefen.contracts" + assert gefen.load_portable_dcp.__module__ == "gefen.portable_dcp" + assert gefen.save_portable_dcp.__module__ == "gefen.portable_dcp" assert gefen.LogicalRegion.__module__ == "gefen.contracts" From c3bb4b222977e9aa2d88be811f8b90a2a3cac078 Mon Sep 17 00:00:00 2001 From: thad0ctor Date: Mon, 13 Jul 2026 04:25:08 -0700 Subject: [PATCH 18/52] Complete hybrid portability and state offload --- COMPATIBILITY.md | 2 +- README.md | 7 +- docs/optimizer_contracts.md | 22 +- src/gefen/__init__.py | 2 + src/gefen/contracts.py | 304 +++-- src/gefen/gefen.py | 1357 ++++++++++++++------- src/gefen/hybrid.py | 637 ++++++++-- src/gefen/portable_dcp.py | 83 +- src/gefen/portable_hybrid.py | 672 ++++++++++ src/gefen/portable_runtime.py | 6 + tests/test_hybrid_rebinding.py | 487 ++++++++ tests/test_optimizer_contracts.py | 59 +- tests/test_portable_dcp_nccl.py | 364 ++++++ tests/test_portable_dcp_topologies.py | 920 ++++++++++++++ tests/test_portable_hybrid.py | 197 +++ tests/test_portable_hybrid_distributed.py | 679 +++++++++++ tests/test_portable_hybrid_runtime.py | 391 ++++++ tests/test_portable_schema.py | 45 +- tests/test_rebinding_cpu.py | 9 +- tests/test_state_offload.py | 602 +++++++++ 20 files changed, 6157 insertions(+), 688 deletions(-) create mode 100644 src/gefen/portable_hybrid.py create mode 100644 tests/test_hybrid_rebinding.py create mode 100644 tests/test_portable_dcp_nccl.py create mode 100644 tests/test_portable_dcp_topologies.py create mode 100644 tests/test_portable_hybrid.py create mode 100644 tests/test_portable_hybrid_distributed.py create mode 100644 tests/test_portable_hybrid_runtime.py create mode 100644 tests/test_state_offload.py diff --git a/COMPATIBILITY.md b/COMPATIBILITY.md index b457dbd..ecb244f 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 and `GefenMuon(sharded_mode="distributed")` 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. 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. DTensor targets, Hybrid composition, 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 gathers and verifies complete logical optimizer state, stores a tensor-only bounded canonical-wire envelope through PyTorch DCP, and projects it onto supported replicated, flattened, or whole-owner targets after load. It can change flattened placement and redistribute Muon owners across checkpoint world sizes; factored second moments remain replicated and same-topology. Gefen-backed Hybrid uses a separately versioned, digested composite wrapper around unchanged child v3 documents, exact disjoint FQN routing, and one all-child freshness/commit boundary. Every checkpoint member temporarily holds the complete dense global document and encoded CPU payload. The target import is fail-before-local-mutation after DCP has completed the read, but DCP storage publication itself is not transactionally atomic. AdamW-backed Hybrid, DTensor targets, asynchronous saving, mixed model/optimizer `Stateful` composition, singleton checkpoint scopes inside a larger initialized default world, and multi-member checkpoint groups whose coordinate zero is not global rank zero are not supported by this path. ## Transformers Trainer DDP diff --git a/README.md b/README.md index 21f676b..a57330c 100644 --- a/README.md +++ b/README.md @@ -632,10 +632,15 @@ 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. + +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 `state_dict()` uses its own nested `{"muon": ..., "backup": ..., "backup_optimizer": "gefen" | "adamw"}` layout. Resume from a checkpoint the hybrid itself saved—not one consolidated or converted to the flat torch `{state, param_groups}` layout. Cross-backend loads are rejected before either child is mutated; legacy untagged hybrid checkpoints are interpreted as Gefen-backed. +- **Hybrid checkpoint schema.** `GefenMuonHybrid`'s ordinary `state_dict()` uses its own nested `{"muon": ..., "backup": ..., "backup_optimizer": "gefen" | "adamw"}` layout. Resume from a checkpoint the hybrid itself saved—not one consolidated or converted to the flat torch `{state, param_groups}` layout. Cross-backend loads are rejected before either child is mutated; legacy untagged hybrid checkpoints are interpreted as Gefen-backed. The separate topology-neutral DCP path above supports only a finalized Gefen-backed Hybrid; AdamW-backed Hybrid remains same-topology through its ordinary nested checkpoint. - **FSDP2 optimizer checkpoints don't reshard.** Plain Gefen and Muon `approx` resume only on the same GPU count and layout; changing either refuses to load. Model weights are unaffected — [details](https://github.com/thad0ctor/Gefen-X/blob/main/COMPATIBILITY.md#optimizer-checkpoint-scope). +- **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. ## Troubleshooting diff --git a/docs/optimizer_contracts.md b/docs/optimizer_contracts.md index bc4fa13..7ff685e 100644 --- a/docs/optimizer_contracts.md +++ b/docs/optimizer_contracts.md @@ -39,19 +39,19 @@ 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. The portable global-state path described below can reshard supported finalized layouts; DTensor stable identity consumption, Hybrid composite rebinding, and offload remain unclaimed. +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. ## Explicit learned-codebook process groups `CodebookProcessGroupBinding` maps one stable `ProcessGroupIdentity` and local semantic member to an opaque PyTorch process-group handle plus an explicit collective device. It is accepted only through the complete `post_sharding(..., codebook_process_group=...)` transaction: every manifest shard must use that one semantic group, each local shard must name the binding's local member, the runtime group size and coordinate must match `ordered_members`, and the backend must support the supplied device. A one-member scope uses `process_group=None`; multi-member scopes must pass a real handle, including `dist.group.WORLD` when the default world is intentionally the semantic scope. Gefen never treats `None` as an implicit default-world selection. -The `explicit_process_group_codebook_scope` capability reports that Gefen or GefenMuon implements this API; `codebook_process_group_binding()` separately reports whether a particular finalized instance has an active binding. GefenMuonHybrid remains negative at the composite level because its independent children do not form one atomically coordinated codebook scope. +The `explicit_process_group_codebook_scope` capability reports that Gefen or GefenMuon implements this API; `codebook_process_group_binding()` separately reports whether a particular finalized instance has an active binding. A Gefen-backed `GefenMuonHybrid` implements one atomic composite transaction that installs the same exact binding object on every present child while retaining independent child codebooks. AdamW-backed Hybrid remains negative because AdamW has no corresponding stable rebinding and staged portable-state contract. The optimizer owns one learned codebook and therefore accepts one scope. Histogram accounting represents each logical parameter once: the first ordered member contributes a replicated parameter after all members agree on gradient presence and its automatic period, every flattened shard contributes its local logical slice after all nonempty slices agree on gradient presence, and only the declared whole-parameter owner contributes an owned matrix. Local inputs are visited in canonical manifest order. Members first stage periods and an integer histogram without touching live optimizer state, exchange operation, step, scope, manifest, active-slice, period, policy, and old-codebook controls, sum the fixed-size `int64` histogram in the supplied group, solve the same exact-DP problem, verify codebook agreement, and only then publish periods and the codebook. Refresh additionally stages every replacement momentum-index tensor in canonical order and exchanges readiness before replacing indices or invalidating derived codebook/LUT caches. A local preparation, solve, or requantization failure is reported by every participant before optimizer-state commit; this is not rollback after process death or a collective-backend failure during the final commit window. `initialize_codebook()` and `refresh_codebook()` expose these operations for adapters that enter their optimizers in a deterministic order; `binding.sort_key` supplies the stable process-group portion of that schedule. Normal `step()` still initializes automatically and plain Gefen still honors `codebook_refresh_every`. Every scoped step exchanges a common operation header before any rank-dependent branch, and the first step after binding or native load additionally verifies codebook bytes and the complete manifest. Scoped native AMP requires every member to select the same protocol and present identical `found_inf` and `grad_scale` values; a mismatch raises collectively and requires a group-aware gradient scaler rather than changing external scaler state behind its back. Multi-member explicit scopes reject `capturable=True` because their host validation and process-group collectives are not CUDA-graph-safe. A one-member local scope may initialize during eager warmup and then use ordinary capturable stepping, but manual codebook replacement remains rejected. Ordinary unscoped behavior remains collective-free. Explicit scope does not replace DTensor mesh collectives, AMP mesh preflights, Parallel-Muon ownership collectives, or checkpoint transport groups. -Native checkpoints store a primitive rank-neutral scope record—format version, semantic name, ordered members, and refresh schedule—at the top level and in the existing transport mirror. Scoped group metadata uses a new outer format version so an older loader rejects rather than silently discarding the scope. Runtime process-group handles and collective devices are live adapter configuration and are never serialized or reconstructed. Flattened and whole-owner local payloads additionally carry a primitive rank-local shard-identity guard. Guard format v2 records every original logical slot in group and slot order, including its lowercase compatibility name and portable shard identity, so pruned whole-owner nonowners remain bound to their original positions and equal-shaped parameters cannot be reinterpreted positionally. New checkpoints emit v2, while the loader continues to accept v1 guards by comparing their legacy live-slot and separately sorted pruned-shard projection exactly. A scoped checkpoint requires the target to complete `post_sharding` with the same stable group identity, refresh policy, and native local layout before load. A missing or different record rejects inside the optimizer's atomic shadow load. Native flattened-shard checkpoint support is declared only on a finalized scoped instance with this guard; unscoped optimizers no longer claim it. Legacy unscoped checkpoints remain loadable by unscoped optimizers. Native whole-owner completeness, scoped DTensor rank-local transport, scope migration, and Hybrid-wide coordination remain separate from the portable global transport and are not claimed by the native 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. Native whole-owner completeness, scoped DTensor rank-local transport, and scope migration remain separate from the portable global transport and are not claimed by the native path. Hybrid's ordinary nested native loader also remains distinct from its atomic composite portable path. ## Exact-binding canonical local state @@ -65,7 +65,7 @@ This is an exact-binding transport-neutral local fragment, distinct from the den The transport-neutral `gefen.portable_state` version-3 format is a complete `global_logical_optimizer` artifact for the supported exact period-one configurations. Its schema carries the implementation and algorithm policy, optimizer-common state, an exact FQN-keyed parameter catalog with global identities, algorithm options, state variants, dense authoritative state, projection hints, source provenance, and a completion marker whose deterministic SHA-256 covers every preceding value including tensor dtype, shape, and canonical little-endian bytes. Optimizer export validates local native state, decodes quantized momentum and block second moments to dense logical fp32 fields, collectively assembles complete parameters by stable shard identity, verifies replicated and owner consensus bit-for-bit, and returns the same normalized complete document to every participant. Import validates the complete document and digest before projecting dense fields to the target shard, recompressing momentum at period one, staging a native shadow load, exchanging unanimous readiness and freshness, and publishing locally only after the final vote. Runtime global ranks are absent from the durable identity. -`export_portable_state(checkpoint_process_group=..., transaction_id=..., limits=...)` and `import_portable_state(state, checkpoint_process_group=..., transaction_id=..., limits=...)` implement this collective path on exact `Gefen` and `GefenMuon` 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 this collective path on exact `Gefen`, `GefenMuon`, and supported Gefen-backed `GefenMuonHybrid` instances. The exact `CheckpointProcessGroupBinding` must match the optimizer's installed `CodebookProcessGroupBinding` in stable identity, local member, live process-group handle, and collective device; all collectives use that optimizer-owned transport. Every member must enter operations in the same deterministic order with the same trimmed transaction ID, limits, target logical slot schema, and, for import, complete document. `PortableStateLimits` bounds per-member and aggregate tensor bytes, metadata, tree shape, strings, tensor count and rank, member count, wire chunk size, and diagnostics; structural limits are checked before dense materialization. `chunk_bytes` controls wire cloning and collective transfer, runtime value validation and freshness hashing use independently fixed bounded chunks, and dense decode/projection allocations remain bounded by the declared fragment and collective tensor-byte ceilings. ```python from gefen import PortableStateLimits @@ -92,6 +92,8 @@ 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. +A finalized Gefen-backed `GefenMuonHybrid` uses a separate `gefen.portable_composite_state` version-1 wrapper rather than pretending its children share one portable-v3 policy or codebook. The wrapper embeds unchanged Muon and Gefen v3 documents, records the exact disjoint FQN-to-role routing and backup policy, requires both child global steps and deterministic settings to agree, and covers the complete wrapper with its own digest. Export agrees on child presence and routing before entering child collectives in fixed role order. Import stages every child, exchanges one composite target and freshness vote, and then performs only the prevalidated nonthrowing child publications; failure before that vote leaves both children unchanged. The DCP tensor envelope is tree-agnostic and stores this wrapper without an opaque payload. Muon-only and Gefen-backup-only Hybrids are supported, while AdamW-backed Hybrid remains negative because AdamW lacks stable FQN rebinding, a topology-neutral moment codec, and a staged commit primitive. + ```python import torch.distributed.checkpoint as dcp @@ -114,19 +116,23 @@ 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`, 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. 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 and contiguous flattened element shards. Block-second-moment state can reshard between replicated and flattened targets; a logical matrix using factored second moments remains replicated and same-topology because factored-to-block representation migration is not implemented. GefenMuon supports replicated matrices and whole-parameter ownership when every participating group uses `sharded_mode="distributed"`; the transport can change placement and redistribute owners across world sizes, including NorMuon row state. A ready Hybrid declaration is the union of its heterogeneous child layouts and change kinds; an adapter must inspect each role's child contract together with the finalized immutable `optimizer.parameter_routing()` result rather than apply that union indiscriminately to every parameter. Pristine and period-selected states are supported under the same policy rules, and zero-element parameters remain pristine. The collective protocol exchanges fixed-size preparation headers before payload movement, visits member fragments in stable semantic order, bounds metadata and tensor chunks, propagates asymmetric local failures to every participant, and performs no semantic checks after the final freshness vote. Import preserves the target's parameter groups, defaults, parameters, compatibility names, and runtime process-group configuration while restoring portable common state, including the source deterministic setting. The atomic claim is fail-before-local-mutation for live, quiescent optimizer instances; it is not rollback after process death, backend failure, or concurrent mutation after the final vote. Ordinary state-dict hooks are bypassed. Adapters must quiesce training, avoid retaining state-container identities across a successful import, and persist the returned weights-only-safe CPU document with their checkpoint system. -Portable v3 currently excludes non-period-one initialized state, second-moment representation conversion, DTensor layouts, stochastic rounding, capturable/device-authoritative state, GefenMuon modes other than `distributed` for whole-owner transport, `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. `state_offload` remains false because a CPU portable document is a checkpoint artifact, not live optimizer state that can be stepped while offloaded. +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. -## Quiescent optimizer-state movement +## Quiescent optimizer-state movement and offload `StateMovementProvider.move_state_(device=None)` performs blocking CPU/CUDA co-location movement for Gefen and GefenMuon. With `device=None`, each declared authoritative per-parameter tensor moves to that parameter's current local device, including declared state attached to wrapper-orphaned parameter keys, while the canonical learned codebook moves to the first live local parameter device in parameter-group order. An optimizer with no local parameter storage keeps common state on CPU. An explicit device is accepted only after every live local parameter already resides there; an unindexed CUDA target is resolved from the one co-located live parameter device. The adapter must therefore move the model parameters first and invoke `move_state_` at a quiescent boundary before the next optimizer step. The core validates the finalized binding and complete declared state representation, allocates detached tight copies of the codebook and every authoritative tensor, and waits for all participating CUDA devices before one local publication. A preparation, transfer, synchronization, or validation failure leaves the exact live optimizer state, caches, parameters, gradients, groups, tensor learning rates, names, bindings, and metadata unchanged. Successful movement replaces the public `optimizer.state` mapping, every reachable or orphan per-parameter state dictionary, the canonical codebook identity, and every moved tensor identity; preserved non-tensor values and rank-local carrier tensors retain their identities, so adapters must not retain the replaced containers. Successful publication preserves host counters and metadata plus rank-local checkpoint carriers, discards only the rebuildable `stepsize` and `_h_buf` buffers, and invalidates per-device codebook/LUT copies, codebook-scope validation, the compiled static-address signature, and the tensor-learning-rate scalar cache. Preserved extension metadata is limited to provably tensor-free trees made from `None`, exact `bool`, `int`, `float`, `complex`, `str`, `bytes`, `torch.device`, `torch.dtype`, `torch.layout`, or `torch.memory_format` leaves and exact `dict`, `list`, `tuple`, `set`, `frozenset`, `deque`, or `torch.Size` containers. Cyclic or multiply referenced container graphs are outside that tree form. Arbitrary opaque objects, `defaultdict` or `OrderedDict` extension values, custom container subclasses, and non-dictionary per-parameter state mappings are rejected even when they appear tensor-free; the optimizer-owned top-level `state` may use its normal exact `defaultdict(dict)` representation. Undeclared tensor-bearing state, meta/nested/subclassed state tensors, FakeTensor parameters, capturable state, active compilation, and CUDA graph capture are likewise rejected rather than partially moved. -`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` remains false because Gefen does not support stepping while authoritative state is parked away from its parameter, asynchronous paging, or transparent CPU offload. +`StateOffloadProvider.offload_state_(device="cpu")` enables synchronous CPU-authoritative per-parameter state for an exact plain `Gefen` instance with ordinary replicated CUDA parameters. Activation first validates the complete declared state, stages tight detached CPU copies, waits for CUDA transfers, and publishes the policy and replacement state mapping together. At each eager step, Gefen copies only the current parameter's persistent tensor state to that parameter's CUDA device, runs the ordinary fused or non-fused block or factored update against a private runtime dictionary, synchronously copies the updated persistent state back to CPU, publishes that one dictionary, and releases the device temporaries. The optimizer-common learned codebook remains resident on CUDA and its normal per-device caches remain available. `restore_state_()` atomically co-locates all state with the parameters and disables offload; `move_state_()` has the same policy-disabling effect after its requested movement succeeds. + +Activation and restore are fail-before-mutation. Activation also rejects persistent state tensors whose storage overlaps another persistent field, a parameter, or the common codebook because independent parameter paging cannot preserve such aliasing. If the update itself raises, Gefen attempts to preserve the resulting runtime state on CPU before propagating the original error. If copyback fails after a parameter may have changed, the optimizer is marked poisoned and refuses subsequent steps or native, canonical, and portable exports until a complete successful native `load_state_dict()` establishes known-good state. An active offload policy is target-local runtime configuration and is preserved across such a load rather than serialized as checkpoint meaning; the active loader maps parameter state directly to CPU and never accumulates the checkpoint's full parameter state on CUDA. State offload is implemented only for an exact plain `Gefen` instance; the composite Hybrid API, `GefenMuon`, nonreplicated finalized layouts, DTensor or tensor-subclass parameters, opaque extension state, multi-member explicit codebook scopes, capturable/device-authoritative state, compilation, and CUDA graph capture are excluded. The multi-member exclusion prevents one rank's copyback poison from bypassing the next scoped collective while peers enter it. Offload must be restored before post-sharding rebinding. It is blocking parameter-scoped paging, not asynchronous prefetch, overlap, or a distributed offload engine, and portable global-state I/O remains unavailable while its authoritative tensors are parked on CPU. + +`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. diff --git a/src/gefen/__init__.py b/src/gefen/__init__.py index 9f0988f..e66eef9 100644 --- a/src/gefen/__init__.py +++ b/src/gefen/__init__.py @@ -46,6 +46,7 @@ "StateGeometry", "StateKeyMatch", "StateMovementProvider", + "StateOffloadProvider", "StateScope", "StateVariant", "TopologyChange", @@ -144,6 +145,7 @@ def __getattr__(name): "StateGeometry", "StateKeyMatch", "StateMovementProvider", + "StateOffloadProvider", "StateScope", "StateVariant", "TopologyChange", diff --git a/src/gefen/contracts.py b/src/gefen/contracts.py index eaefe3e..4718165 100644 --- a/src/gefen/contracts.py +++ b/src/gefen/contracts.py @@ -140,7 +140,9 @@ def _validate_dimensions(name, values, *, positive): def _validate_identity_name(name, value): if not isinstance(value, str) or not value or value != value.strip(): - raise ValueError("{} must be a non-empty string without outer whitespace".format(name)) + raise ValueError( + "{} must be a non-empty string without outer whitespace".format(name) + ) if "\x00" in value: raise ValueError("{} must not contain NUL".format(name)) @@ -161,12 +163,18 @@ class ParameterIdentity: def __post_init__(self) -> None: _validate_identity_name("ParameterIdentity.fqn", self.fqn) if self.fqn.startswith(".") or self.fqn.endswith(".") or ".." in self.fqn: - raise ValueError("ParameterIdentity.fqn must contain non-empty dot-separated components") + raise ValueError( + "ParameterIdentity.fqn must contain non-empty dot-separated components" + ) if isinstance(self.global_shape, (str, bytes, bytearray)): - raise TypeError("ParameterIdentity.global_shape must be a sequence of dimensions") + raise TypeError( + "ParameterIdentity.global_shape must be a sequence of dimensions" + ) object.__setattr__(self, "global_shape", _tuple(self.global_shape)) if any(type(dim) is not int or dim < 0 for dim in self.global_shape): - raise ValueError("ParameterIdentity.global_shape must contain nonnegative integers") + raise ValueError( + "ParameterIdentity.global_shape must contain nonnegative integers" + ) _validate_identity_schema_version("parameter identity", self.schema_version) @property @@ -185,9 +193,13 @@ class ProcessGroupIdentity: schema_version: int = IDENTITY_SCHEMA_VERSION def __post_init__(self) -> None: - _validate_identity_name("ProcessGroupIdentity.semantic_name", self.semantic_name) + _validate_identity_name( + "ProcessGroupIdentity.semantic_name", self.semantic_name + ) if isinstance(self.ordered_members, (str, bytes)): - raise TypeError("ProcessGroupIdentity.ordered_members must be a sequence of member IDs") + raise TypeError( + "ProcessGroupIdentity.ordered_members must be a sequence of member IDs" + ) object.__setattr__(self, "ordered_members", _tuple(self.ordered_members)) if not self.ordered_members: raise ValueError("ProcessGroupIdentity.ordered_members must be non-empty") @@ -214,13 +226,24 @@ def __post_init__(self) -> None: raise TypeError("ShardPlacement.kind must be a PlacementKind") if type(self.parts) is not int or self.parts <= 0: raise ValueError("ShardPlacement.parts must be a positive integer") - if type(self.coordinate) is not int or self.coordinate < 0 or self.coordinate >= self.parts: + if ( + type(self.coordinate) is not int + or self.coordinate < 0 + or self.coordinate >= self.parts + ): raise ValueError("ShardPlacement.coordinate must be within parts") if self.kind is PlacementKind.DIMENSION_SHARD: - if type(self.parameter_dimension) is not int or self.parameter_dimension < 0: - raise ValueError("dimension-shard placement requires a nonnegative parameter dimension") + if ( + type(self.parameter_dimension) is not int + or self.parameter_dimension < 0 + ): + raise ValueError( + "dimension-shard placement requires a nonnegative parameter dimension" + ) elif self.parameter_dimension is not None: - raise ValueError("only a dimension-shard placement may name a parameter dimension") + raise ValueError( + "only a dimension-shard placement may name a parameter dimension" + ) @dataclass(frozen=True) @@ -266,9 +289,7 @@ def __post_init__(self) -> None: "{} must be a sequence of dimensions".format(name) ) from exc if any(type(value) is not int or value < 0 for value in normalized): - raise ValueError( - "{} must contain nonnegative integers".format(name) - ) + raise ValueError("{} must contain nonnegative integers".format(name)) object.__setattr__(self, name.rsplit(".", 1)[1], normalized) if len(self.offsets) != len(self.lengths): raise ValueError( @@ -301,9 +322,7 @@ def validate_bounds(self, parameter: ParameterIdentity) -> None: if not isinstance(parameter, ParameterIdentity): raise TypeError("parameter must be a ParameterIdentity") if self.rank != len(parameter.global_shape): - raise ValueError( - "LogicalRegion rank must match the global parameter rank" - ) + raise ValueError("LogicalRegion rank must match the global parameter rank") if any( offset + length > dimension for offset, length, dimension in zip( @@ -320,8 +339,7 @@ def intersection(self, other: "LogicalRegion") -> "LogicalRegion": if self.rank != other.rank: raise ValueError("LogicalRegion intersection requires equal ranks") offsets = tuple( - max(left, right) - for left, right in zip(self.offsets, other.offsets) + max(left, right) for left, right in zip(self.offsets, other.offsets) ) ends = tuple( min(left_offset + left_length, right_offset + right_length) @@ -366,9 +384,7 @@ def validate_exact_coverage( if any(region.overlaps(other) for other in regions[index + 1 :]): raise ValueError("LogicalRegions must not overlap") if sum(region.numel for region in regions) != parameter.numel: - raise ValueError( - "LogicalRegions must exactly cover the global parameter" - ) + raise ValueError("LogicalRegions must exactly cover the global parameter") @dataclass(frozen=True) @@ -393,8 +409,7 @@ def __post_init__(self) -> None: if self.layout is ParameterLayout.DTENSOR_1D_DEFAULT_WORLD: if not uses_logical_region: raise ValueError( - "stable DTensor identity requires a logical-region descriptor " - "(LogicalRegion)" + "stable DTensor identity requires a logical-region descriptor (LogicalRegion)" ) self.logical_slice.validate_bounds(self.parameter) elif not isinstance(self.logical_slice, LogicalSlice): @@ -403,7 +418,9 @@ def __post_init__(self) -> None: ) placements = _tuple(self.placements) if any(not isinstance(item, ShardPlacement) for item in placements): - raise TypeError("ShardIdentity.placements must contain ShardPlacement values") + raise TypeError( + "ShardIdentity.placements must contain ShardPlacement values" + ) axes = tuple(item.mesh_axis for item in placements) if len(set(axes)) != len(axes): raise ValueError("ShardIdentity placement mesh axes must be unique") @@ -420,16 +437,30 @@ def __post_init__(self) -> None: raise ValueError("ShardIdentity.logical_slice exceeds the global parameter") if self.process_group is None: if self.local_member is not None or self.owner is not None: - raise ValueError("ShardIdentity members and owners require a process-group identity") + raise ValueError( + "ShardIdentity members and owners require a process-group identity" + ) else: if not isinstance(self.process_group, ProcessGroupIdentity): - raise TypeError("ShardIdentity.process_group must be a ProcessGroupIdentity") + raise TypeError( + "ShardIdentity.process_group must be a ProcessGroupIdentity" + ) if self.local_member not in self.process_group.ordered_members: - raise ValueError("ShardIdentity.local_member must belong to the process group") - if self.owner is not None and self.owner not in self.process_group.ordered_members: + raise ValueError( + "ShardIdentity.local_member must belong to the process group" + ) + if ( + self.owner is not None + and self.owner not in self.process_group.ordered_members + ): raise ValueError("ShardIdentity.owner must belong to the process group") - if self.layout is not ParameterLayout.WHOLE_PARAMETER_OWNER and self.owner is not None: - raise ValueError("ShardIdentity.owner is valid only for whole-parameter ownership") + if ( + self.layout is not ParameterLayout.WHOLE_PARAMETER_OWNER + and self.owner is not None + ): + raise ValueError( + "ShardIdentity.owner is valid only for whole-parameter ownership" + ) full = ( self.logical_slice == LogicalRegion.full(self.parameter) @@ -440,13 +471,21 @@ def __post_init__(self) -> None: if self.process_group is not None: member_index = self.process_group.ordered_members.index(self.local_member) for placement in self.placements: - if placement.parts != len(self.process_group.ordered_members) or placement.coordinate != member_index: - raise ValueError("ShardIdentity placement coordinates must match the ordered process-group members") + if ( + placement.parts != len(self.process_group.ordered_members) + or placement.coordinate != member_index + ): + raise ValueError( + "ShardIdentity placement coordinates must match the ordered process-group members" + ) for placement in self.placements: - if placement.parameter_dimension is not None and placement.parameter_dimension >= len( - self.parameter.global_shape + if ( + placement.parameter_dimension is not None + and placement.parameter_dimension >= len(self.parameter.global_shape) ): - raise ValueError("ShardIdentity placement dimension exceeds parameter rank") + raise ValueError( + "ShardIdentity placement dimension exceeds parameter rank" + ) if self.layout is ParameterLayout.REPLICATED: if not full or self.owner is not None: raise ValueError("replicated identity must cover the full parameter") @@ -455,22 +494,34 @@ def __post_init__(self) -> None: if self.process_group is None and self.placements: raise ValueError("an ungrouped replicated identity has no placements") if self.process_group is not None and len(self.placements) != 1: - raise ValueError("a process-group replicated identity requires one placement") + raise ValueError( + "a process-group replicated identity requires one placement" + ) elif self.layout is ParameterLayout.FLATTENED_ELEMENT_SHARD: if self.process_group is None or self.owner is not None: - raise ValueError("flattened element shards require a process group and no owner") + raise ValueError( + "flattened element shards require a process group and no owner" + ) if len(kinds) != 1 or kinds[0] is not PlacementKind.FLAT_SHARD: - raise ValueError("flattened element shards require one flat-shard placement") + raise ValueError( + "flattened element shards require one flat-shard placement" + ) elif self.layout is ParameterLayout.WHOLE_PARAMETER_OWNER: if self.process_group is None or self.owner is None: - raise ValueError("whole-parameter ownership requires a process group and owner") + raise ValueError( + "whole-parameter ownership requires a process group and owner" + ) owns_parameter = self.local_member == self.owner if owns_parameter and not full: - raise ValueError("the owner must carry the full whole-parameter logical slice") + raise ValueError( + "the owner must carry the full whole-parameter logical slice" + ) if not owns_parameter and self.logical_slice != LogicalSlice(0, 0): raise ValueError("a non-owner whole-parameter slice must be empty") if len(kinds) != 1 or kinds[0] is not PlacementKind.WHOLE_PARAMETER_OWNER: - raise ValueError("whole-parameter ownership requires one owner placement") + raise ValueError( + "whole-parameter ownership requires one owner placement" + ) elif self.layout is ParameterLayout.DTENSOR_1D_DEFAULT_WORLD: if self.process_group is None or self.owner is not None: raise ValueError( @@ -481,8 +532,7 @@ def __post_init__(self) -> None: PlacementKind.REPLICATE, }: raise ValueError( - "one-dimensional DTensor identities require one dimension-shard " - "or replicate placement" + "one-dimensional DTensor identities require one dimension-shard or replicate placement" ) placement = self.placements[0] if placement.kind is PlacementKind.REPLICATE: @@ -503,8 +553,7 @@ def __post_init__(self) -> None: offset != 0 or length != global_length ): raise ValueError( - "a dimension-sharded DTensor region must cover every " - "unsharded parameter dimension" + "a dimension-sharded DTensor region must cover every unsharded parameter dimension" ) _validate_identity_schema_version("shard identity", self.schema_version) @@ -599,19 +648,33 @@ def __post_init__(self) -> None: for fqn, parameter_shards in by_fqn.items(): parameter = parameter_shards[0].parameter if any(item.parameter != parameter for item in parameter_shards[1:]): - raise ValueError("manifest shards for {!r} disagree on parameter identity".format(fqn)) + raise ValueError( + "manifest shards for {!r} disagree on parameter identity".format( + fqn + ) + ) layouts = {item.layout for item in parameter_shards} groups = {item.process_group for item in parameter_shards} if len(layouts) != 1 or len(groups) != 1: - raise ValueError("manifest shards for {!r} disagree on layout or process group".format(fqn)) + raise ValueError( + "manifest shards for {!r} disagree on layout or process group".format( + fqn + ) + ) layout = parameter_shards[0].layout group = parameter_shards[0].process_group members = tuple(item.local_member for item in parameter_shards) if group is None: if len(parameter_shards) != 1: - raise ValueError("an ungrouped parameter must have exactly one manifest shard") - elif set(members) != set(group.ordered_members) or len(members) != len(group.ordered_members): - raise ValueError("manifest shards must contain each process-group member exactly once") + raise ValueError( + "an ungrouped parameter must have exactly one manifest shard" + ) + elif set(members) != set(group.ordered_members) or len(members) != len( + group.ordered_members + ): + raise ValueError( + "manifest shards must contain each process-group member exactly once" + ) if group is not None: placement_shapes = { tuple( @@ -626,7 +689,9 @@ def __post_init__(self) -> None: for item in parameter_shards } if len(placement_shapes) != 1: - raise ValueError("manifest member placements must agree on one topology") + raise ValueError( + "manifest member placements must agree on one topology" + ) if layout is ParameterLayout.FLATTENED_ELEMENT_SHARD: cursor = 0 @@ -643,22 +708,31 @@ def __post_init__(self) -> None: empty_offsets.append(item.logical_slice.flat_offset) continue if item.logical_slice.flat_offset != cursor: - raise ValueError("flattened manifest slices must be gapless and non-overlapping") + raise ValueError( + "flattened manifest slices must be gapless and non-overlapping" + ) cursor += item.logical_slice.length boundaries.add(cursor) if cursor != parameter.numel: - raise ValueError("flattened manifest slices must cover the global parameter") + raise ValueError( + "flattened manifest slices must cover the global parameter" + ) if any(offset not in boundaries for offset in empty_offsets): raise ValueError( "empty flattened manifest slices must use a partition boundary" ) elif layout is ParameterLayout.REPLICATED: - if any(item.logical_slice != LogicalSlice.full(parameter) for item in parameter_shards): + if any( + item.logical_slice != LogicalSlice.full(parameter) + for item in parameter_shards + ): raise ValueError("replicated manifest shards must all be complete") elif layout is ParameterLayout.WHOLE_PARAMETER_OWNER: owners = {item.owner for item in parameter_shards} if len(owners) != 1: - raise ValueError("whole-parameter manifest shards must agree on one owner") + raise ValueError( + "whole-parameter manifest shards must agree on one owner" + ) elif layout is ParameterLayout.DTENSOR_1D_DEFAULT_WORLD: placements = tuple(item.placements[0] for item in parameter_shards) placement_kind = placements[0].kind @@ -688,8 +762,7 @@ def __post_init__(self) -> None: cursor += region.lengths[shard_dimension] if cursor != parameter.global_shape[shard_dimension]: raise ValueError( - "dimension-sharded DTensor manifest regions must cover " - "the global parameter" + "dimension-sharded DTensor manifest regions must cover the global parameter" ) def for_parameter(self, fqn: str) -> Tuple[ShardIdentity, ...]: @@ -782,12 +855,14 @@ def __post_init__(self) -> None: raise TypeError("StateVariant.role must be a ParameterStateRole") if not set(self.inactive_fields).issubset(self.fields): raise ValueError("StateVariant.inactive_fields must be present in fields") - if self.parameter_ranks is not None and set( - self.parameter_ranks - ) & set(self.excluded_parameter_ranks): + if self.parameter_ranks is not None and set(self.parameter_ranks) & set( + self.excluded_parameter_ranks + ): raise ValueError("included and excluded parameter ranks must be disjoint") if self.parameter_ranks is not None: - _validate_dimensions("parameter_ranks", self.parameter_ranks, positive=False) + _validate_dimensions( + "parameter_ranks", self.parameter_ranks, positive=False + ) _validate_dimensions( "excluded_parameter_ranks", self.excluded_parameter_ranks, @@ -806,7 +881,9 @@ class OptimizerStateLayout: def __post_init__(self) -> None: object.__setattr__(self, "fields", _tuple(self.fields)) object.__setattr__(self, "parameter_variants", _tuple(self.parameter_variants)) - object.__setattr__(self, "composite_namespaces", _tuple(self.composite_namespaces)) + object.__setattr__( + self, "composite_namespaces", _tuple(self.composite_namespaces) + ) names = tuple(field.name for field in self.fields) if len(set(names)) != len(names): raise ValueError("OptimizerStateLayout fields must have unique names") @@ -856,9 +933,7 @@ class TrainingSupport: def __post_init__(self) -> None: if self.mesh_dimensions is not None: object.__setattr__(self, "mesh_dimensions", _tuple(self.mesh_dimensions)) - _validate_dimensions( - "mesh_dimensions", self.mesh_dimensions, positive=True - ) + _validate_dimensions("mesh_dimensions", self.mesh_dimensions, positive=True) if not isinstance(self.layout, ParameterLayout): raise TypeError("TrainingSupport.layout must be a ParameterLayout") if not isinstance(self.process_group_scope, ProcessGroupScope): @@ -896,7 +971,9 @@ class CheckpointSupport: 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_changing", _frozenset(self.topology_changing) + ) object.__setattr__( self, "topology_change_kinds", _frozenset(self.topology_change_kinds) ) @@ -907,9 +984,7 @@ def __post_init__(self) -> None: ) if self.mesh_dimensions is not None: object.__setattr__(self, "mesh_dimensions", _tuple(self.mesh_dimensions)) - _validate_dimensions( - "mesh_dimensions", self.mesh_dimensions, positive=True - ) + _validate_dimensions("mesh_dimensions", self.mesh_dimensions, positive=True) if not isinstance(self.transport, CheckpointTransport): raise TypeError("CheckpointSupport.transport must be a CheckpointTransport") if not isinstance(self.process_group_scope, ProcessGroupScope): @@ -1050,6 +1125,21 @@ def move_state_(self, device=None) -> None: """Co-locate authoritative state with the optimizer's live parameters.""" +@runtime_checkable +class StateOffloadProvider(Protocol): + """Structural protocol for persistent live optimizer-state offload.""" + + @property + def state_offload_device(self): + """Return the active offload device, or ``None`` while resident.""" + + def offload_state_(self, device="cpu") -> None: + """Atomically park supported state and enable transparent step paging.""" + + def restore_state_(self) -> None: + """Atomically co-locate state with parameters and disable offload.""" + + _ALL_PRECISIONS = frozenset( {Precision.FLOAT32, Precision.BFLOAT16, Precision.FLOAT16, Precision.FLOAT64} ) @@ -1094,7 +1184,9 @@ def _common_fields() -> Tuple[StateField, ...]: def _base_parameter_fields() -> Tuple[StateField, ...]: return ( StateField("name", StateScope.PARAMETER, StateGeometry.OPAQUE, True), - StateField("automatic_period", StateScope.PARAMETER, StateGeometry.SCALAR, True), + StateField( + "automatic_period", StateScope.PARAMETER, StateGeometry.SCALAR, True + ), StateField("step", StateScope.PARAMETER, StateGeometry.SCALAR, True), StateField("m_codebook", StateScope.PARAMETER, StateGeometry.PARAMETER, True), StateField("m_magnitude", StateScope.PARAMETER, StateGeometry.BLOCK, True), @@ -1147,7 +1239,9 @@ def _derived_fields() -> Tuple[StateField, ...]: False, description="Rebuildable cross-member scope-agreement cache.", ), - StateField("_sr_seed_by_device", StateScope.DERIVED, StateGeometry.SCALAR, False), + StateField( + "_sr_seed_by_device", StateScope.DERIVED, StateGeometry.SCALAR, False + ), StateField( "_gefen_global_step_by_device", StateScope.DERIVED, @@ -1210,6 +1304,7 @@ def _negative_capabilities( post_sharding: bool = False, canonical_state_io: bool = False, atomic_state_movement: bool = False, + state_offload: bool = False, ) -> OptimizerCapabilities: return OptimizerCapabilities( training=training, @@ -1224,7 +1319,7 @@ def _negative_capabilities( post_sharding=post_sharding, canonical_state_io=canonical_state_io, atomic_state_movement=atomic_state_movement, - state_offload=False, + state_offload=state_offload, ) @@ -1240,12 +1335,11 @@ def _gefen_contract( canonical_global_topology_changing: AbstractSet[ParameterLayout] = frozenset(), canonical_global_topology_change_kinds: AbstractSet[TopologyChange] = frozenset(), atomic_state_movement: bool = False, + state_offload: bool = False, ) -> OptimizerContract: canonical_state_layouts = _frozenset(canonical_state_layouts) canonical_global_same_topology = _frozenset(canonical_global_same_topology) - canonical_global_topology_changing = _frozenset( - canonical_global_topology_changing - ) + canonical_global_topology_changing = _frozenset(canonical_global_topology_changing) canonical_global_topology_change_kinds = _frozenset( canonical_global_topology_change_kinds ) @@ -1388,10 +1482,7 @@ def _gefen_contract( ) ) fields = ( - _common_fields() - + _base_parameter_fields() - + block_fields - + factored_fields + _common_fields() + _base_parameter_fields() + block_fields + factored_fields ) fields += _derived_fields() training = _base_training() + ( @@ -1463,6 +1554,7 @@ def _gefen_contract( or canonical_global_topology_changing ), atomic_state_movement=atomic_state_movement, + state_offload=state_offload, ), ) @@ -1504,9 +1596,7 @@ def _gefen_muon_contract( ) -> OptimizerContract: canonical_state_layouts = _frozenset(canonical_state_layouts) canonical_global_same_topology = _frozenset(canonical_global_same_topology) - canonical_global_topology_changing = _frozenset( - canonical_global_topology_changing - ) + canonical_global_topology_changing = _frozenset(canonical_global_topology_changing) canonical_global_topology_change_kinds = _frozenset( canonical_global_topology_change_kinds ) @@ -1757,10 +1847,7 @@ def _gefen_muon_contract( atomic_load=True, ) ] - if ( - "approx" in sharded_modes - and sharded_modes.issubset({"approx", "distributed"}) - ): + if "approx" in sharded_modes and sharded_modes.issubset({"approx", "distributed"}): checkpoints.append( CheckpointSupport( CheckpointTransport.PYTORCH_RANK_LOCAL, @@ -1840,6 +1927,14 @@ def _hybrid_contract( muon: Optional[OptimizerContract], backup: Optional[OptimizerContract], backup_implementation: str, + canonical_parameter_fqns: bool = False, + stable_shard_identity: bool = False, + explicit_process_group_codebook_scope: bool = False, + shard_rebinding: bool = False, + post_sharding: bool = False, + canonical_global_same_topology: AbstractSet[ParameterLayout] = frozenset(), + canonical_global_topology_changing: AbstractSet[ParameterLayout] = frozenset(), + canonical_global_topology_change_kinds: AbstractSet[TopologyChange] = frozenset(), ) -> OptimizerContract: fields = ( StateField( @@ -1853,21 +1948,36 @@ def _hybrid_contract( if muon is not None: children.append(OptimizerChildContract("muon", muon.implementation, muon)) if backup_implementation: - children.append( - OptimizerChildContract("backup", backup_implementation, backup) - ) + children.append(OptimizerChildContract("backup", backup_implementation, backup)) if muon is None: training = _base_training() else: training = muon.capabilities.training - checkpoints = ( + canonical_global_same_topology = _frozenset(canonical_global_same_topology) + canonical_global_topology_changing = _frozenset(canonical_global_topology_changing) + canonical_global_topology_change_kinds = _frozenset( + canonical_global_topology_change_kinds + ) + checkpoints = [ CheckpointSupport( CheckpointTransport.COMPOSITE_NATIVE, frozenset({ParameterLayout.REPLICATED}), frozenset(), ProcessGroupScope.NONE, ), - ) + ] + if canonical_global_same_topology or canonical_global_topology_changing: + checkpoints.append( + CheckpointSupport( + CheckpointTransport.CANONICAL_GLOBAL, + canonical_global_same_topology, + canonical_global_topology_changing, + ProcessGroupScope.ADAPTER_DEFINED, + topology_change_kinds=canonical_global_topology_change_kinds, + requires_collective=True, + atomic_load=True, + ) + ) return OptimizerContract( implementation="gefen.GefenMuonHybrid", state_layout=OptimizerStateLayout( @@ -1877,8 +1987,16 @@ def _hybrid_contract( ), capabilities=_negative_capabilities( training=training, - checkpoints=checkpoints, + checkpoints=tuple(checkpoints), supported_parameter_ranks=None, + canonical_parameter_fqns=canonical_parameter_fqns, + stable_shard_identity=stable_shard_identity, + explicit_process_group_codebook_scope=explicit_process_group_codebook_scope, + shard_rebinding=shard_rebinding, + post_sharding=post_sharding, + canonical_state_io=bool( + canonical_global_same_topology or canonical_global_topology_changing + ), ), children=tuple(children), ) @@ -1900,6 +2018,7 @@ def _hybrid_contract( "ParameterLayout", "ParameterIdentity", "ParameterStateRole", + "PortableStateProvider", "PlacementKind", "Precision", "ProcessGroupScope", @@ -1912,6 +2031,7 @@ def _hybrid_contract( "StateGeometry", "StateKeyMatch", "StateMovementProvider", + "StateOffloadProvider", "StateScope", "StateVariant", "TrainingSupport", diff --git a/src/gefen/gefen.py b/src/gefen/gefen.py index 71fc121..c8e7df1 100644 --- a/src/gefen/gefen.py +++ b/src/gefen/gefen.py @@ -17,6 +17,7 @@ import os import warnings from collections import defaultdict, deque, OrderedDict +from copy import deepcopy from itertools import chain from typing import Iterable, Optional, Tuple, Union @@ -170,6 +171,7 @@ def _rank_local_payload_key(global_rank: int) -> str: return "{}{}".format(_RANK_LOCAL_PAYLOAD_KEY_PREFIX, int(global_rank)) + # Optional explicit override for the period-search backend ("cuda_kernel" / "cpu" # / "gpu"). None (default) means resolve per call from each tensor's own device # -- see _resolve_find_period_backend. NOT used as an auto-populated cache. @@ -528,9 +530,7 @@ def _gefen_exact_histogram_from_grad_periods( torch.long ) bin_indices = bin_indices.clamp(0, histogram_bins - 1) - local_counts = torch.bincount( - bin_indices, minlength=histogram_bins - ) + local_counts = torch.bincount(bin_indices, minlength=histogram_bins) bin_counts_cpu.add_(local_counts.cpu()) finally: torch.set_deterministic_debug_mode(prev_mode) @@ -669,8 +669,9 @@ def _amp_prepare_optimizer_step(optimizer) -> bool: if torch.is_tensor(found_inf): if found_inf.numel() != 1: raise RuntimeError( - "GradScaler supplied a non-scalar found_inf tensor with " - "shape {}".format(tuple(found_inf.shape)) + "GradScaler supplied a non-scalar found_inf tensor with shape {}".format( + tuple(found_inf.shape) + ) ) overflow = bool(found_inf.detach().item()) else: @@ -686,8 +687,9 @@ def _amp_prepare_optimizer_step(optimizer) -> bool: if torch.is_tensor(grad_scale): if grad_scale.numel() != 1: raise RuntimeError( - "GradScaler supplied a non-scalar grad_scale tensor with shape " - "{}".format(tuple(grad_scale.shape)) + "GradScaler supplied a non-scalar grad_scale tensor with shape {}".format( + tuple(grad_scale.shape) + ) ) scale = grad_scale.detach() # Match torch.amp.GradScaler.unscale_: computing the reciprocal in @@ -752,8 +754,7 @@ def _amp_dtensor_protocol_preflight(optimizer) -> bool: params = [param for group in optimizer.param_groups for param in group["params"]] has_fp16_storage = any( - param.dtype == torch.float16 - and not getattr(param, "_is_flat_param", False) + param.dtype == torch.float16 and not getattr(param, "_is_flat_param", False) for param in params ) if not has_fp16_storage: @@ -777,8 +778,7 @@ def _amp_dtensor_protocol_preflight(optimizer) -> bool: str(mesh.device_type), tuple(int(item) for item in mesh.shape), tuple( - int(item) - for item in mesh.mesh.detach().cpu().reshape(-1).tolist() + int(item) for item in mesh.mesh.detach().cpu().reshape(-1).tolist() ), tuple(str(group.group_name) for group in process_groups), ) @@ -793,9 +793,7 @@ def _amp_dtensor_protocol_preflight(optimizer) -> bool: name = ( names[index] if isinstance(names, (list, tuple)) and index < len(names) - else getattr(optimizer, "_param_name", lambda _: "parameter")( - param - ) + else getattr(optimizer, "_param_name", lambda _: "parameter")(param) ) entry["items"].append((str(name), param, param.grad is not None)) @@ -931,9 +929,7 @@ def _assert_optimizer_gradients_structurally_valid( else getattr(optimizer, "_param_name", lambda _: "parameter")(param) ) if require_2d_params and param.ndim != 2: - error_factory = getattr( - optimizer, "_step_non_2d_parameter_error", None - ) + error_factory = getattr(optimizer, "_step_non_2d_parameter_error", None) if callable(error_factory): raise ValueError(error_factory(param)) raise ValueError( @@ -950,19 +946,17 @@ def _assert_optimizer_gradients_structurally_valid( raise RuntimeError( "Gefen does not support sparse gradients or other " "non-strided layouts; parameter {!r} has gradient layout " - "{}.".format( - str(name), layout - ) + "{}.".format(str(name), layout) ) if torch.is_complex(grad): raise RuntimeError( - "Gefen optimizers do not support complex gradients, but " - "parameter {!r} has dtype {}.".format(str(name), grad.dtype) + "Gefen optimizers do not support complex gradients, but parameter {!r} has dtype {}.".format( + str(name), grad.dtype + ) ) if tuple(grad.shape) != tuple(param.shape): raise RuntimeError( - "Gefen gradient shape {} for parameter {!r} does not match " - "parameter shape {}.".format( + "Gefen gradient shape {} for parameter {!r} does not match parameter shape {}.".format( tuple(grad.shape), str(name), tuple(param.shape) ) ) @@ -1088,8 +1082,7 @@ def __init__( ) if fused and not torch.cuda.is_available(): warnings.warn( - "Gefen optimizer got fused=True, but CUDA is not available. " - "Changing fused to False.", + "Gefen optimizer got fused=True, but CUDA is not available. Changing fused to False.", UserWarning, stacklevel=2, ) @@ -1109,9 +1102,7 @@ def __init__( # partitioning.memory_safe_fallback_period. self._force_1d_period_one = force_1d_period_one self._force_2d_period_one = force_2d_period_one - self._period_one_substrings = tuple( - s.lower() for s in period_one_substrings - ) + self._period_one_substrings = tuple(s.lower() for s in period_one_substrings) self._factored_v_2d = factored_v_2d self._deterministic = deterministic if type(codebook_refresh_every) is not int: @@ -1201,6 +1192,17 @@ def __init__( # one device counter per parameter device and advance it in the captured # step tail; state_dict synchronizes the host mirror before serializing. self._gefen_global_step_by_device = {} + # Native parameter-state offload is an eager, target-local runtime + # policy rather than optimizer checkpoint meaning. None keeps the + # historical co-located path. A CPU device means declared persistent + # per-parameter tensors are CPU-authoritative between steps while the + # small optimizer-common codebook remains resident/cached normally. + self._gefen_state_offload_device = None + # A failed CUDA-to-CPU copyback can leave a parameter updated while its + # last published CPU state is stale. Preserve that diagnosis across + # restore/movement and reject later steps until a successful load + # establishes a complete known-good state again. + self._gefen_state_offload_poisoned = False defaults = dict( lr=lr, @@ -1255,9 +1257,13 @@ 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 and self._factored_v_2d and any( - len(slot.shard.parameter.global_shape) == 2 - for slot in self._gefen_logical_slots + has_factored_matrix = ( + canonical_global_same_topology + and self._factored_v_2d + and any( + len(slot.shard.parameter.global_shape) == 2 + for slot in self._gefen_logical_slots + ) ) if canonical_global_same_topology and not has_factored_matrix: canonical_global_topology_changing = frozenset( @@ -1282,6 +1288,7 @@ def optimizer_contract(self) -> OptimizerContract: canonical_global_topology_changing=canonical_global_topology_changing, canonical_global_topology_change_kinds=canonical_global_topology_change_kinds, atomic_state_movement=self._atomic_state_movement_supported(), + state_offload=self._state_offload_supported(), native_flattened_checkpoint=( self._codebook_scope_ready() and any( @@ -1317,6 +1324,8 @@ def _canonical_group_options_value(group): } def _canonical_state_layouts(self): + if self.state_offload_poisoned: + return frozenset() if not self._canonical_identity_ready(): return frozenset() if self._stochastic_round: @@ -1389,9 +1398,7 @@ def _canonical_state_layouts(self): ) except (TypeError, ValueError, RuntimeError): return frozenset() - return frozenset( - shard.layout for _, shard in self._gefen_local_shard_bindings - ) + return frozenset(shard.layout for _, shard in self._gefen_local_shard_bindings) def _codebook_scope_ready(self) -> bool: return ( @@ -1500,9 +1507,7 @@ def _finalized_binding_layout_matches(self) -> bool: live_group = [] name_group = [] for logical_slot in logical_group: - parameter, shard = local_bindings[ - logical_slot.shard.parameter.fqn - ] + parameter, shard = local_bindings[logical_slot.shard.parameter.fqn] if shard != logical_slot.shard: return False pruned = ( @@ -1538,9 +1543,7 @@ def _finalized_binding_layout_matches(self) -> bool: if len(self._param_names) != len(expected_live_bindings): return False - if { - id(parameter) for parameter in self._param_names - } != expected_live_ids: + if {id(parameter) for parameter in self._param_names} != expected_live_ids: return False for group_index, (group, expected_params, expected_names) in enumerate( zip(self.param_groups, expected_live_groups, expected_name_groups) @@ -1588,7 +1591,10 @@ def _finalized_binding_layout_matches(self) -> bool: return False def _assert_finalized_binding_layout(self) -> None: - if self._gefen_post_sharding_finalized and not self._finalized_binding_layout_matches(): + if ( + self._gefen_post_sharding_finalized + and not self._finalized_binding_layout_matches() + ): raise RuntimeError( "Gefen finalized parameter layout changed outside post_sharding" ) @@ -1795,9 +1801,7 @@ def _normalize_serialized_canonical_shard(cls, record): ) placements = [] for placement_record in record["placements"]: - if not isinstance(placement_record, dict) or set( - placement_record - ) != { + if not isinstance(placement_record, dict) or set(placement_record) != { "mesh_axis", "kind", "coordinate", @@ -1817,9 +1821,7 @@ def _normalize_serialized_canonical_shard(cls, record): shard = ShardIdentity( parameter, ParameterLayout(record["layout"]), - LogicalSlice( - slice_record["flat_offset"], slice_record["length"] - ), + LogicalSlice(slice_record["flat_offset"], slice_record["length"]), process_group=process_group, local_member=record["local_member"], owner=record["owner"], @@ -1901,9 +1903,7 @@ def _normalize_serialized_native_local_shard(cls, record): "placements", } if not isinstance(record, dict) or set(record) != expected: - raise ValueError( - "Gefen native local-shard metadata has an invalid schema" - ) + raise ValueError("Gefen native local-shard metadata has an invalid schema") group_record = record["process_group"] if not isinstance(group_record, dict) or set(group_record) != { "semantic_name", @@ -1913,9 +1913,7 @@ def _normalize_serialized_native_local_shard(cls, record): "Gefen native local-shard process-group metadata is invalid" ) if not isinstance(group_record["ordered_members"], list): - raise ValueError( - "Gefen native local-shard ordered_members must be a list" - ) + raise ValueError("Gefen native local-shard ordered_members must be a list") if not isinstance(record["global_shape"], list) or not isinstance( record["placements"], list ): @@ -1923,18 +1921,14 @@ def _normalize_serialized_native_local_shard(cls, record): "Gefen native local-shard shapes and placements must be lists" ) try: - parameter = ParameterIdentity( - record["fqn"], tuple(record["global_shape"]) - ) + parameter = ParameterIdentity(record["fqn"], tuple(record["global_shape"])) group = ProcessGroupIdentity( group_record["semantic_name"], tuple(group_record["ordered_members"]), ) placements = [] for placement_record in record["placements"]: - if not isinstance(placement_record, dict) or set( - placement_record - ) != { + if not isinstance(placement_record, dict) or set(placement_record) != { "mesh_axis", "kind", "coordinate", @@ -1961,9 +1955,7 @@ def _normalize_serialized_native_local_shard(cls, record): placements=tuple(placements), ) except (TypeError, ValueError) as exc: - raise ValueError( - "Gefen native local-shard metadata is invalid" - ) from exc + raise ValueError("Gefen native local-shard metadata is invalid") from exc if shard.layout not in { ParameterLayout.REPLICATED, ParameterLayout.FLATTENED_ELEMENT_SHARD, @@ -1981,9 +1973,7 @@ def _normalize_serialized_native_local_shards_v1(cls, value): "param_groups", "pruned_shards", }: - raise ValueError( - "Gefen native local-shard metadata has an invalid schema" - ) + raise ValueError("Gefen native local-shard metadata has an invalid schema") format_version = value["format_version"] if ( type(format_version) is not int @@ -2035,9 +2025,7 @@ def _normalize_serialized_native_local_shards_v2(cls, value): "format_version", "logical_slots", }: - raise ValueError( - "Gefen native local-shard metadata has an invalid schema" - ) + raise ValueError("Gefen native local-shard metadata has an invalid schema") format_version = value["format_version"] if ( type(format_version) is not int @@ -2136,9 +2124,7 @@ def _normalize_serialized_native_local_shards(cls, value): if value is None: return None if type(value) is not dict or "format_version" not in value: - raise ValueError( - "Gefen native local-shard metadata has an invalid schema" - ) + raise ValueError("Gefen native local-shard metadata has an invalid schema") format_version = value["format_version"] if type(format_version) is not int: raise ValueError( @@ -2215,6 +2201,10 @@ def _target_may_have_internal_storage_overlap(parameter) -> bool: def _assert_rebinding_pristine(self, rebindings) -> None: if self._gefen_post_sharding_finalized: raise RuntimeError("Gefen post-sharding identity is already finalized") + if self.state_offload_active or self.state_offload_poisoned: + raise RuntimeError( + "Gefen parameter rebinding requires resident known-good state; restore state first" + ) if ( self._gefen_shard_bindings or self._gefen_local_shard_bindings @@ -2325,8 +2315,7 @@ def _validate_rebinding_layout(self, rebinding: ParameterRebinding) -> None: ) else: raise ValueError( - "plain Gefen rebinding supports replicated or flattened element " - "shards only" + "plain Gefen rebinding supports replicated or flattened element shards only" ) if target.numel() != shard.logical_slice.length: raise ValueError( @@ -2366,8 +2355,7 @@ def _validate_codebook_process_group_binding( ParameterLayout.WHOLE_PARAMETER_OWNER, }: raise ValueError( - "explicit codebook scope supports replicated, flattened, or " - "whole-parameter owner identities" + "explicit codebook scope supports replicated, flattened, or whole-parameter owner identities" ) self._validate_codebook_runtime_binding(binding) @@ -2390,7 +2378,10 @@ def _validate_codebook_runtime_binding(self, binding) -> None: raise ValueError( "a multi-member codebook scope requires an explicit runtime process group" ) - if not torch.distributed.is_available() or not torch.distributed.is_initialized(): + if ( + not torch.distributed.is_available() + or not torch.distributed.is_initialized() + ): raise RuntimeError( "a multi-member codebook scope requires initialized torch.distributed" ) @@ -2398,9 +2389,7 @@ def _validate_codebook_runtime_binding(self, binding) -> None: try: world = dist.get_world_size(binding.process_group) - group_rank = dist.get_group_rank( - binding.process_group, dist.get_rank() - ) + group_rank = dist.get_group_rank(binding.process_group, dist.get_rank()) backend = dist.get_backend(binding.process_group) except Exception as exc: raise ValueError( @@ -2451,9 +2440,7 @@ def _codebook_parameter_contributes(self, parameter) -> bool: return shard.local_member == binding.identity.ordered_members[0] return True - def _stage_post_sharding( - self, rebindings, manifest, codebook_process_group=None - ): + def _stage_post_sharding(self, rebindings, manifest, codebook_process_group=None): self._assert_rebinding_pristine(rebindings) live_slots = [] for group_index, group in enumerate(self.param_groups): @@ -2462,17 +2449,13 @@ def _stage_post_sharding( if len(names) != len(params): names = [self._param_name(param) for param in params] for parameter_index, (parameter, name) in enumerate(zip(params, names)): - live_slots.append( - (group_index, parameter_index, parameter, str(name)) - ) + live_slots.append((group_index, parameter_index, parameter, str(name))) if len(rebindings) != len(live_slots): raise ValueError( "post_sharding requires exactly one rebinding for every optimizer slot" ) - manifest_fqns = { - shard.parameter.fqn for shard in manifest.shards - } + manifest_fqns = {shard.parameter.fqn for shard in manifest.shards} local_fqns = [item.shard.parameter.fqn for item in rebindings] if len(set(local_fqns)) != len(local_fqns): raise ValueError("local canonical parameter FQNs must be unique") @@ -2507,11 +2490,12 @@ def _stage_post_sharding( 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" + "stable DTensor rebinding requires the deferred logical-region identity schema" ) if torch.is_complex(target): - raise ValueError("Gefen does not support complex rebound parameters") + raise ValueError( + "Gefen does not support complex rebound parameters" + ) if target.layout is not torch.strided or target.dtype not in ( torch.float16, torch.bfloat16, @@ -2524,8 +2508,7 @@ def _stage_post_sharding( if target.is_meta: raise ValueError("rebound parameters require materialized storage") if ( - rebinding.shard.layout - is ParameterLayout.FLATTENED_ELEMENT_SHARD + rebinding.shard.layout is ParameterLayout.FLATTENED_ELEMENT_SHARD and not target.is_contiguous() ): raise ValueError( @@ -2533,8 +2516,7 @@ def _stage_post_sharding( ) if self._target_may_have_internal_storage_overlap(target): raise ValueError( - "rebound target storage must be provably free of internal " - "storage overlap" + "rebound target storage must be provably free of internal storage overlap" ) if not target.is_leaf and not target.retains_grad: raise ValueError("can't optimize a non-leaf rebound Tensor") @@ -2565,9 +2547,7 @@ def _stage_post_sharding( target_positions = [] if target is not None: target_positions = [ - index - for index, slot in enumerate(live_slots) - if slot[2] is target + index for index, slot in enumerate(live_slots) if slot[2] is target ] if len(source_positions) == 1: position = source_positions[0] @@ -2708,9 +2688,7 @@ def post_sharding( if self._parameter_in(sources, rebinding.old_parameter): raise ValueError("rebinding source tensors must be unique") sources.append(rebinding.old_parameter) - staged = self._stage_post_sharding( - rebindings, manifest, codebook_process_group - ) + staged = self._stage_post_sharding(rebindings, manifest, codebook_process_group) self.__dict__.update(staged.__dict__) def rebind_shard( @@ -2755,10 +2733,7 @@ def _state_value_is_movement_safe_metadata(value, seen=None) -> bool: if value is None or value_type in _STATE_MOVEMENT_METADATA_LEAF_TYPES: return True is_mapping = value_type in _STATE_MOVEMENT_METADATA_MAPPING_TYPES - if ( - not is_mapping - and value_type not in _STATE_MOVEMENT_METADATA_SEQUENCE_TYPES - ): + if not is_mapping and value_type not in _STATE_MOVEMENT_METADATA_SEQUENCE_TYPES: return False if seen is None: seen = set() @@ -2774,18 +2749,14 @@ def _state_value_is_movement_safe_metadata(value, seen=None) -> bool: for key, item in items ) return all( - Gefen._state_value_is_movement_safe_metadata(item, seen) - for item in items + Gefen._state_value_is_movement_safe_metadata(item, seen) for item in items ) @staticmethod def _state_movement_tensor_supported(value) -> bool: return ( type(value) is torch.Tensor - and not ( - hasattr(value, "to_local") - and hasattr(value, "placements") - ) + and not (hasattr(value, "to_local") and hasattr(value, "placements")) and value.layout is torch.strided and not value.is_nested and not value.is_quantized @@ -2793,6 +2764,450 @@ def _state_movement_tensor_supported(value) -> bool: and value.device.type in _STATE_MOVEMENT_DEVICE_TYPES ) + @property + def state_offload_active(self) -> bool: + """Whether native parameter-state CPU offload is currently enabled.""" + + return getattr(self, "_gefen_state_offload_device", None) is not None + + @property + def state_offload_device(self): + """Return the active state-offload device, or ``None`` when disabled.""" + + return getattr(self, "_gefen_state_offload_device", None) + + @property + def state_offload_poisoned(self) -> bool: + """Whether a failed copyback made further stepping unsafe.""" + + return bool(getattr(self, "_gefen_state_offload_poisoned", False)) + + def _assert_state_export_safe(self) -> None: + if self.state_offload_poisoned: + raise RuntimeError( + "Gefen cannot export optimizer state after a failed state-offload copyback; " + "load a known-good checkpoint first" + ) + + @staticmethod + def _state_offload_parameter_supported(parameter) -> bool: + return ( + type(parameter) in {torch.Tensor, nn.Parameter} + and not ( + hasattr(parameter, "to_local") and hasattr(parameter, "placements") + ) + and parameter.layout is torch.strided + and parameter.device.type == "cuda" + and parameter.dtype + in {torch.float16, torch.bfloat16, torch.float32, torch.float64} + and not torch.is_complex(parameter) + and not parameter.is_meta + and not parameter.is_nested + and not parameter.is_quantized + and getattr(parameter, "fake_mode", None) is None + ) + + @classmethod + def _state_offload_cpu_tensor_supported(cls, value) -> bool: + return ( + cls._state_movement_tensor_supported(value) + and value.device.type == "cpu" + and not value.requires_grad + and value.is_contiguous() + and value.storage_offset() == 0 + and value.untyped_storage().nbytes() == value.numel() * value.element_size() + ) + + @classmethod + def _state_offload_storage_disjoint(cls, parameters, state, codebook) -> bool: + """Conservatively reject aliases that parameter-scoped paging cannot preserve.""" + + tensors = [("parameter", parameter) for parameter in parameters] + if torch.is_tensor(codebook): + tensors.append(("common_state", codebook)) + for parameter_state in state.values(): + for key, value in parameter_state.items(): + if key in _STATE_MOVEMENT_TENSOR_KEYS and torch.is_tensor(value): + tensors.append(("parameter_state", value)) + + storage_ranges = [] + try: + for kind, tensor in tensors: + if tensor.numel() == 0: + continue + if ( + kind != "parameter" + and cls._target_may_have_internal_storage_overlap(tensor) + ): + return False + storage = tensor.untyped_storage() + storage_id = (str(tensor.device), storage.data_ptr()) + if tensor.is_contiguous(): + start = tensor.storage_offset() * tensor.element_size() + end = start + tensor.numel() * tensor.element_size() + else: + start = end = None + for other_kind, other_id, other_start, other_end in storage_ranges: + if other_id != storage_id or (kind == other_kind == "parameter"): + continue + if ( + start is None + or other_start is None + or max(start, other_start) < min(end, other_end) + ): + return False + storage_ranges.append((kind, storage_id, start, end)) + except Exception: + return False + return True + + @staticmethod + def _state_offload_capturing_on_parameter_device(parameters) -> bool: + if torch.cuda.is_available() and torch.cuda.is_current_stream_capturing(): + return True + devices = { + parameter.device + for parameter in parameters + if torch.is_tensor(parameter) and parameter.device.type == "cuda" + } + devices = sorted( + devices, + key=lambda device: -1 if device.index is None else device.index, + ) + for device in devices: + with torch.cuda.device(device): + if torch.cuda.is_current_stream_capturing(): + return True + return False + + def _state_offload_rejection_reason( + self, *, require_cpu_state: bool, allow_poisoned: bool = False + ): + if type(self) is not Gefen: + return "native state offload is implemented only by plain Gefen" + if self.state_offload_active and self.state_offload_device != torch.device( + "cpu" + ): + return "the active state-offload policy has an invalid device" + if self.state_offload_poisoned and not allow_poisoned: + return "a previous state copyback failed; load a known-good checkpoint" + scope = self._gefen_codebook_process_group + if scope is not None: + if type(scope) is not CodebookProcessGroupBinding: + return "the explicit codebook process-group binding is invalid" + if len(scope.identity.ordered_members) > 1: + return "state offload does not support multi-member explicit codebook scopes" + if self.capturable: + return "capturable optimizers have device-authoritative replay state" + try: + parameters = [ + parameter + for group in self.param_groups + for parameter in group["params"] + ] + except (KeyError, TypeError): + return "parameter groups have an invalid structure" + if torch.compiler.is_compiling(): + return "state offload cannot run during torch.compile" + if torch.cuda.is_available(): + try: + if self._state_offload_capturing_on_parameter_device(parameters): + return "state offload cannot run during CUDA graph capture" + except RuntimeError: + return "the CUDA graph-capture state could not be inspected" + + movement_reason = self._state_movement_rejection_reason() + if movement_reason is not None: + return movement_reason + + if not parameters: + return "state offload requires at least one CUDA parameter" + if any( + not self._state_offload_parameter_supported(parameter) + for parameter in parameters + ): + return "state offload requires ordinary replicated CUDA parameters" + + if self._gefen_post_sharding_finalized: + if len(self._gefen_local_shard_bindings) != len(parameters): + return "the finalized local shard registry is incomplete" + if any( + parameter is None or shard.layout is not ParameterLayout.REPLICATED + for parameter, shard in self._gefen_local_shard_bindings + ): + return "state offload supports only finalized replicated layouts" + + state_type = type(self.state) + if not ( + state_type is dict + or (state_type is defaultdict and self.state.default_factory is dict) + ): + return "optimizer state must use a supported standard mapping" + if set(self.state) != set(parameters): + return "state offload requires exactly one state entry per live parameter" + + allowed_keys = ( + _CANONICAL_PARAMETER_STATE_KEYS | _CANONICAL_DERIVED_PARAMETER_STATE_KEYS + ) + for parameter in parameters: + parameter_state = self.state.get(parameter) + if type(parameter_state) is not dict: + return "per-parameter optimizer state must use a plain dictionary" + if any(key not in allowed_keys for key in parameter_state): + return "state offload does not support custom per-parameter state" + if require_cpu_state and any( + key in _STATE_MOVEMENT_SCRATCH_KEYS for key in parameter_state + ): + return "offloaded state contains device-side runtime scratch" + for key, value in parameter_state.items(): + if key in _STATE_MOVEMENT_COUNTER_KEYS: + if type(value) is not int: + return ( + "noncapturable offloaded counters must be Python integers" + ) + continue + if key not in _STATE_MOVEMENT_TENSOR_KEYS: + continue + if not torch.is_tensor(value): + return "authoritative offloaded state has an invalid value" + if require_cpu_state: + if not self._state_offload_cpu_tensor_supported(value): + return ( + "authoritative offloaded tensors must be tight CPU tensors" + ) + elif not self._state_movement_tensor_supported(value): + return ( + "authoritative tensor state has an unsupported representation" + ) + + if not self._state_offload_storage_disjoint( + parameters, + self.state, + self._gefen_codebook, + ): + return "persistent optimizer-state storage aliases another live tensor" + + try: + self._validate_loaded_native_state() + except Exception: + return "optimizer state does not match Gefen's declared native schema" + if ( + require_cpu_state + and self._gefen_codebook is not None + and self._gefen_codebook.device.type != "cuda" + ): + return "the optimizer-common codebook must remain CUDA-resident" + return None + + def _state_offload_supported(self) -> bool: + """Return whether this live optimizer can safely enter or retain offload.""" + + try: + reason = self._state_offload_rejection_reason( + require_cpu_state=self.state_offload_active + ) + except Exception: + return False + return reason is None + + @staticmethod + def _normalize_state_offload_target(device) -> torch.device: + try: + target = torch.device(device) + except (TypeError, RuntimeError) as exc: + raise TypeError("state offload device must be CPU") from exc + if target.type != "cpu": + raise ValueError("Gefen native state offload currently supports only CPU") + return torch.device("cpu") + + def _copy_state_tensor_to_offload_cpu(self, tensor: torch.Tensor) -> torch.Tensor: + return self._copy_state_tensor_for_move(tensor, torch.device("cpu")) + + def _stage_state_offload_resident_codebook(self): + codebook = self._gefen_codebook + if codebook is None or codebook.device.type == "cuda": + return codebook + parameter = next( + parameter for group in self.param_groups for parameter in group["params"] + ) + target = parameter.device + staged = self._copy_state_tensor_for_move(codebook, target) + self._validate_staged_state_tensor(codebook, staged, target) + torch.cuda.synchronize(target) + return staged + + def _prepare_offloaded_cpu_parameter_state(self, parameter_state): + if type(parameter_state) is not dict: + raise RuntimeError("offloaded runtime state must use a plain dictionary") + result = {} + cuda_devices = set() + allowed_keys = ( + _CANONICAL_PARAMETER_STATE_KEYS | _CANONICAL_DERIVED_PARAMETER_STATE_KEYS + ) + for key, value in parameter_state.items(): + if key not in allowed_keys: + raise RuntimeError( + "offloaded stepping produced unsupported state key {!r}".format(key) + ) + if key in _STATE_MOVEMENT_SCRATCH_KEYS: + continue + if key in _STATE_MOVEMENT_CAPTURABLE_KEYS: + raise RuntimeError( + "offloaded stepping produced capturable runtime state" + ) + if key in _STATE_MOVEMENT_COUNTER_KEYS: + if type(value) is not int: + raise RuntimeError( + "offloaded stepping produced a device-authoritative counter" + ) + result[key] = value + continue + if key in _STATE_MOVEMENT_TENSOR_KEYS: + if not self._state_movement_tensor_supported(value): + raise RuntimeError( + "offloaded stepping produced unsupported tensor state" + ) + staged = self._copy_state_tensor_to_offload_cpu(value) + self._validate_staged_state_tensor(value, staged, torch.device("cpu")) + result[key] = staged + if value.device.type == "cuda": + cuda_devices.add(value.device) + continue + result[key] = value + + for cuda_device in sorted( + cuda_devices, key=lambda item: -1 if item.index is None else item.index + ): + torch.cuda.synchronize(cuda_device) + return result + + def _stage_all_parameter_state_to_cpu(self): + staged_state = defaultdict(dict) + for parameter, parameter_state in self.state.items(): + staged_state[parameter] = self._prepare_offloaded_cpu_parameter_state( + parameter_state + ) + return staged_state + + def _stage_offloaded_parameter_state(self, parameter): + cpu_state = self.state[parameter] + if type(cpu_state) is not dict: + raise RuntimeError("offloaded parameter state must use a plain dictionary") + target = parameter.device + runtime_state = {} + for key, value in cpu_state.items(): + if key in _STATE_MOVEMENT_COUNTER_KEYS: + if type(value) is not int: + raise RuntimeError( + "offloaded parameter counters must remain Python integers" + ) + runtime_state[key] = value + elif key in _STATE_MOVEMENT_TENSOR_KEYS: + if not self._state_offload_cpu_tensor_supported(value): + raise RuntimeError( + "offloaded parameter tensors must remain tight CPU tensors" + ) + staged = self._copy_state_tensor_for_move(value, target) + self._validate_staged_state_tensor(value, staged, target) + runtime_state[key] = staged + else: + runtime_state[key] = value + torch.cuda.synchronize(target) + return runtime_state + + def _step_with_offloaded_parameter_state( + self, update, group, param_name, parameter, grad + ) -> None: + runtime_state = self._stage_offloaded_parameter_state(parameter) + try: + update(group, param_name, parameter, grad, state=runtime_state) + except BaseException as operation_error: + try: + cpu_state = self._prepare_offloaded_cpu_parameter_state(runtime_state) + except BaseException as copyback_error: + self._gefen_state_offload_poisoned = True + if hasattr(operation_error, "add_note"): + operation_error.add_note( + "Gefen state copyback also failed; the optimizer is poisoned" + ) + raise operation_error from copyback_error + self.state[parameter] = cpu_state + raise + + try: + cpu_state = self._prepare_offloaded_cpu_parameter_state(runtime_state) + except BaseException as exc: + self._gefen_state_offload_poisoned = True + raise RuntimeError( + "Gefen state copyback failed after a parameter update; load a " + "known-good checkpoint before stepping again" + ) from exc + self.state[parameter] = cpu_state + + def _assert_state_offload_step_ready(self) -> None: + if self.state_offload_poisoned: + raise RuntimeError( + "Gefen state offload is poisoned after a failed copyback; load a " + "known-good checkpoint before stepping again" + ) + if not self.state_offload_active: + return + reason = self._state_offload_rejection_reason(require_cpu_state=True) + if reason is not None: + raise RuntimeError("Gefen state offload cannot step: {}".format(reason)) + + @torch.no_grad() + def offload_state_(self, device="cpu") -> None: + """Atomically enable synchronous CPU-authoritative parameter state.""" + + target = self._normalize_state_offload_target(device) + self._assert_finalized_binding_layout() + reason = self._state_offload_rejection_reason(require_cpu_state=False) + if reason is not None: + raise RuntimeError("Gefen state offload is unavailable: {}".format(reason)) + staged_state = self._stage_all_parameter_state_to_cpu() + staged_codebook = self._stage_state_offload_resident_codebook() + updates = { + "state": staged_state, + "_gefen_state_offload_device": target, + "_static_mark_sig": None, + "_lr_scalar_cache": None, + } + if staged_codebook is not self._gefen_codebook: + updates.update( + { + "_gefen_codebook": staged_codebook, + "_gefen_codebook_by_device": {}, + "_gefen_codebook_lut_by_device": {}, + "_gefen_codebook_scope_validated": False, + } + ) + self.__dict__.update(updates) + + @torch.no_grad() + def restore_state_(self) -> None: + """Atomically co-locate parameter state and disable native offload.""" + + if not self.state_offload_active: + return + parameters = [ + parameter for group in self.param_groups for parameter in group["params"] + ] + if torch.cuda.is_available(): + try: + capturing = self._state_offload_capturing_on_parameter_device( + parameters + ) + except RuntimeError as exc: + raise RuntimeError( + "Gefen state restore could not inspect CUDA graph-capture state" + ) from exc + if capturing: + raise RuntimeError( + "Gefen state restore cannot run during CUDA graph capture" + ) + self.move_state_() + def _state_movement_rejection_reason(self): if ( self._gefen_post_sharding_finalized @@ -2827,10 +3242,7 @@ def _state_movement_rejection_reason(self): state_type = type(self.state) if not ( state_type is dict - or ( - state_type is defaultdict - and self.state.default_factory is dict - ) + or (state_type is defaultdict and self.state.default_factory is dict) ): return "optimizer state must use a supported standard mapping" for parameter, parameter_state in self.state.items(): @@ -2855,7 +3267,10 @@ def _state_movement_rejection_reason(self): if torch.is_tensor(value): if not self._state_movement_tensor_supported(value): return "authoritative tensor state has an unsupported representation" - elif key not in _STATE_MOVEMENT_COUNTER_KEYS or type(value) is not int: + elif ( + key not in _STATE_MOVEMENT_COUNTER_KEYS + or type(value) is not int + ): return "authoritative tensor state has an invalid value" continue if isinstance(key, str) and key.startswith( @@ -2865,7 +3280,9 @@ def _state_movement_rejection_reason(self): if not self._state_movement_tensor_supported(value): return "rank-local transport state has an unsupported representation" elif not self._state_value_is_movement_safe_metadata(value): - return "rank-local transport state contains unsupported metadata" + return ( + "rank-local transport state contains unsupported metadata" + ) continue if not self._state_value_is_movement_safe_metadata(value): return "undeclared optimizer state is not provably tensor-free metadata" @@ -2886,7 +3303,12 @@ def _atomic_state_movement_supported(self) -> bool: return False if torch.cuda.is_available(): try: - if torch.cuda.is_current_stream_capturing(): + parameters = [ + parameter + for group in self.param_groups + for parameter in group["params"] + ] + if self._state_offload_capturing_on_parameter_device(parameters): return False except RuntimeError: return False @@ -2923,7 +3345,9 @@ def _normalize_state_move_target(device, live_devices) -> torch.device: target = torch.device("cpu") else: if not torch.cuda.is_available(): - raise ValueError("CUDA state movement requires an available CUDA device") + raise ValueError( + "CUDA state movement requires an available CUDA device" + ) if target.index is None: unique_devices = set(live_devices) if len(unique_devices) != 1: @@ -2932,10 +3356,14 @@ def _normalize_state_move_target(device, live_devices) -> torch.device: ) candidate = next(iter(unique_devices)) if candidate.type != "cuda": - raise ValueError("CUDA state movement requires CUDA-resident parameters") + raise ValueError( + "CUDA state movement requires CUDA-resident parameters" + ) target = candidate if target.index < 0 or target.index >= torch.cuda.device_count(): - raise ValueError("CUDA state movement target is not an available device") + raise ValueError( + "CUDA state movement target is not an available device" + ) if not live_devices: if target.type != "cpu": @@ -2944,8 +3372,9 @@ def _normalize_state_move_target(device, live_devices) -> torch.device: ) elif any(parameter_device != target for parameter_device in live_devices): raise ValueError( - "explicit state movement requires every live parameter to already be " - "co-located on {}".format(target) + "explicit state movement requires every live parameter to already be co-located on {}".format( + target + ) ) return target @@ -2981,12 +3410,11 @@ def _validate_staged_state_tensor( def _stage_state_move(self, device): live_parameters = [ - parameter - for group in self.param_groups - for parameter in group["params"] + parameter for group in self.param_groups for parameter in group["params"] ] live_devices = [ - self._state_move_parameter_device(parameter) for parameter in live_parameters + self._state_move_parameter_device(parameter) + for parameter in live_parameters ] live_parameter_ids = {id(parameter) for parameter in live_parameters} explicit_target = ( @@ -3054,11 +3482,29 @@ def move_state_(self, device=None) -> None: "Gefen atomic state movement could not inspect live optimizer state" ) from exc if reason is not None: - raise RuntimeError("Gefen atomic state movement is unavailable: {}".format(reason)) + raise RuntimeError( + "Gefen atomic state movement is unavailable: {}".format(reason) + ) if torch.compiler.is_compiling(): raise RuntimeError("Gefen state movement cannot run during torch.compile") - if torch.cuda.is_available() and torch.cuda.is_current_stream_capturing(): - raise RuntimeError("Gefen state movement cannot run during CUDA graph capture") + if torch.cuda.is_available(): + parameters = [ + parameter + for group in self.param_groups + for parameter in group["params"] + ] + try: + capturing = self._state_offload_capturing_on_parameter_device( + parameters + ) + except RuntimeError as exc: + raise RuntimeError( + "Gefen state movement could not inspect CUDA graph-capture state" + ) from exc + if capturing: + raise RuntimeError( + "Gefen state movement cannot run during CUDA graph capture" + ) staged_state, staged_codebook = self._stage_state_move(device) self.__dict__.update( @@ -3070,6 +3516,7 @@ def move_state_(self, device=None) -> None: "_gefen_codebook_scope_validated": False, "_static_mark_sig": None, "_lr_scalar_cache": None, + "_gefen_state_offload_device": None, } ) @@ -3094,13 +3541,9 @@ def _validate_group_options(lr, betas, eps, weight_decay): if not 0.0 <= lr: raise ValueError("Invalid learning rate: {}".format(lr)) elif not 0.0 <= betas[0] < 1.0: - raise ValueError( - "Invalid beta parameter at index 0: {}".format(betas[0]) - ) + raise ValueError("Invalid beta parameter at index 0: {}".format(betas[0])) elif not 0.0 <= betas[1] < 1.0: - raise ValueError( - "Invalid beta parameter at index 1: {}".format(betas[1]) - ) + raise ValueError("Invalid beta parameter at index 1: {}".format(betas[1])) elif not 0.0 <= weight_decay: raise ValueError("Invalid weight_decay value: {}".format(weight_decay)) elif not math.isfinite(eps) or eps <= 0.0: @@ -3126,8 +3569,9 @@ def _iter_params_with_names(group_params): ) if torch.is_complex(param): raise ValueError( - "Gefen does not support complex parameters, but got a " - "parameter with dtype {}".format(param.dtype) + "Gefen does not support complex parameters, but got a parameter with dtype {}".format( + param.dtype + ) ) yield param_name, param @@ -3184,8 +3628,7 @@ def _sync_param_names_to_state(self) -> None: logical_groups = [[] for _ in self.param_groups] for logical_slot in self._gefen_logical_slots: if ( - logical_slot.shard.layout - is ParameterLayout.WHOLE_PARAMETER_OWNER + logical_slot.shard.layout is ParameterLayout.WHOLE_PARAMETER_OWNER and logical_slot.shard.local_member != logical_slot.shard.owner ): continue @@ -3202,8 +3645,7 @@ def _sync_param_names_to_state(self) -> None: names = logical_groups[group_index] if len(names) != len(params): raise RuntimeError( - "Gefen finalized logical-slot registry does not match the " - "loaded layout" + "Gefen finalized logical-slot registry does not match the loaded layout" ) normalized_names = [] for param, name in zip(params, names): @@ -3223,15 +3665,17 @@ def add_param_group(self, param_group): stable lowercase name is stored in its per-param state and in the group's ``param_names`` list for introspection. """ + if getattr(self, "_gefen_state_offload_device", None) is not None: + raise RuntimeError( + "Gefen cannot add parameter groups while state offload is active; restore state first" + ) if getattr(self, "_gefen_post_sharding_finalized", False): raise RuntimeError( "Gefen cannot add parameter groups after post_sharding finalization" ) if not isinstance(param_group, dict): raise TypeError( - "param_group must be a dict, got {}".format( - type(param_group).__name__ - ) + "param_group must be a dict, got {}".format(type(param_group).__name__) ) if "params" not in param_group: raise ValueError("Gefen parameter group is missing the 'params' key.") @@ -3435,9 +3879,7 @@ def _assert_capturable_if_capturing(self) -> None: ) if capturing: devices = { - param.device - for group in self.param_groups - for param in group["params"] + param.device for group in self.param_groups for param in group["params"] } capture_device = torch.device("cuda", torch.cuda.current_device()) if devices != {capture_device}: @@ -3470,9 +3912,7 @@ def _static_mark_signature(self): acc += sum(map(id, params)) acc ^= sum(id(p.grad) for p in params) state = self.state - acc += sum( - map(id, chain.from_iterable(map(dict.values, state.values()))) - ) + acc += sum(map(id, chain.from_iterable(map(dict.values, state.values())))) stacks = self._capt_stacks return ( acc, @@ -3525,9 +3965,7 @@ def _mark_state_static_for_compile(self) -> None: mark(self._gefen_codebook) lut = _codebook_search_lut(self._gefen_codebook) mark(lut) - self._gefen_codebook_lut_by_device[ - self._gefen_codebook.device - ] = lut + self._gefen_codebook_lut_by_device[self._gefen_codebook.device] = lut for device, cached in self._gefen_codebook_by_device.items(): mark(cached) lut = _codebook_search_lut(cached) @@ -3600,8 +4038,10 @@ def _sr_seed_on(self, device: torch.device): else int(global_counter.item()) ) seed = torch.full( - (), initial_step, - dtype=torch.int64, device=device, + (), + initial_step, + dtype=torch.int64, + device=device, ) self._sr_seed_by_device[device] = seed return seed @@ -3626,8 +4066,9 @@ def _device_gefen_global_step(self): ] if any(step != steps[0] for step in steps[1:]): raise RuntimeError( - "Gefen capturable global-step counters disagree across devices: " - "{}".format(steps) + "Gefen capturable global-step counters disagree across devices: {}".format( + steps + ) ) return steps[0] @@ -3643,7 +4084,9 @@ def _ensure_gefen_global_step_devices(self) -> None: ] if missing: device_step = self._device_gefen_global_step() - initial_step = self._gefen_global_step if device_step is None else device_step + initial_step = ( + self._gefen_global_step if device_step is None else device_step + ) self._gefen_global_step = initial_step for device in missing: self._gefen_global_step_by_device[device] = torch.full( @@ -3667,13 +4110,10 @@ def _synchronize_gefen_global_step_for_checkpoint(self) -> None: if device_step is not None: self._gefen_global_step = device_step if self.capturable and self._stochastic_round: - seed_steps = [ - int(seed.item()) for seed in self._sr_seed_by_device.values() - ] + seed_steps = [int(seed.item()) for seed in self._sr_seed_by_device.values()] if any(seed_step != self._gefen_global_step for seed_step in seed_steps): raise RuntimeError( - "Gefen capturable stochastic-round seeds disagree with the " - "optimizer global step: {} != {}".format( + "Gefen capturable stochastic-round seeds disagree with the optimizer global step: {} != {}".format( seed_steps, self._gefen_global_step ) ) @@ -3928,10 +4368,7 @@ def _capt_build_stacks(self) -> bool: steps2d = torch.empty((2, n), dtype=torch.float32, device=device) steps2d[0].copy_( torch.stack( - [ - self.state[p][key].detach() - for p, key in zip(rows, d["keys"]) - ] + [self.state[p][key].detach() for p, key in zip(rows, d["keys"])] ) ) steps2d[1].copy_( @@ -4057,9 +4494,7 @@ def _capt_batched_ready(self) -> bool: # per-row alias/routing guard, but validate group-owned hypers once per # contiguous device run instead of once per parameter. Heterogeneous # registries take the fully general path below. - if stacks and all( - len(device_stacks) == 1 for device_stacks in stacks.values() - ): + if stacks and all(len(device_stacks) == 1 for device_stacks in stacks.values()): if self._capt_single_cohort_ready(stacks): return True if capturing: @@ -4100,8 +4535,7 @@ def _capt_batched_ready(self) -> bool: or state["step"] is not stack["step_views"][i] or state.get(info[0]) is not stack["bc2_views"][i] or state.get("_capt_scalars") is not stack["scalar_views"][i] - or stack["betas"][i] - != (group["beta1"], group["beta2"]) + or stack["betas"][i] != (group["beta1"], group["beta2"]) or group["weight_decay"] != stack["weight_decay0"] ): ok = False @@ -4290,9 +4724,7 @@ def _gefen_codebook_on(self, device: torch.device) -> Optional[torch.Tensor]: self._gefen_codebook_by_device[device] = cached return cached - def _gefen_codebook_lut_on( - self, device: torch.device - ) -> Optional[torch.Tensor]: + def _gefen_codebook_lut_on(self, device: torch.device) -> Optional[torch.Tensor]: # Resolve once per device and retain the tensor on the optimizer. The # eager Muon momentum path calls this per matrix; returning a stable # cached tensor avoids rebuilding the data_ptr/version cache key in the @@ -4421,9 +4853,7 @@ def _predict_period_from_grad_sq( ) if backend == "cpu": - period_input = ( - grad_work.detach().float().square().reshape(-1).cpu().numpy() - ) + period_input = grad_work.detach().float().square().reshape(-1).cpu().numpy() elif backend == "gpu": if grad_work.device.type != "cuda": raise ValueError("FIND_PERIOD_BACKEND='gpu' requires a CUDA tensor") @@ -4566,8 +4996,9 @@ def _prepare_scoped_amp_optimizer_step(self) -> bool: elif torch.is_tensor(found_inf): if found_inf.numel() != 1: raise RuntimeError( - "GradScaler supplied a non-scalar found_inf tensor with " - "shape {}".format(tuple(found_inf.shape)) + "GradScaler supplied a non-scalar found_inf tensor with shape {}".format( + tuple(found_inf.shape) + ) ) local_overflow = bool(found_inf.detach().item()) else: @@ -4583,17 +5014,16 @@ def _prepare_scoped_amp_optimizer_step(self) -> bool: elif torch.is_tensor(grad_scale): if grad_scale.numel() != 1: raise RuntimeError( - "GradScaler supplied a non-scalar grad_scale tensor with " - "shape {}".format(tuple(grad_scale.shape)) + "GradScaler supplied a non-scalar grad_scale tensor with shape {}".format( + tuple(grad_scale.shape) + ) ) scale_present = True local_scale = float(grad_scale.detach().item()) else: scale_present = True local_scale = float(grad_scale) - if scale_present and ( - not math.isfinite(local_scale) or local_scale <= 0.0 - ): + if scale_present and (not math.isfinite(local_scale) or local_scale <= 0.0): raise RuntimeError( "GradScaler supplied a non-finite or non-positive grad_scale" ) @@ -4603,9 +5033,7 @@ def _prepare_scoped_amp_optimizer_step(self) -> bool: scale_present = False local_scale = 0.0 local_error = exc - self._synchronize_codebook_scope_failure( - local_error, "AMP overflow preflight" - ) + self._synchronize_codebook_scope_failure(local_error, "AMP overflow preflight") if len(binding.identity.ordered_members) > 1: import torch.distributed as dist @@ -4615,8 +5043,7 @@ def _prepare_scoped_amp_optimizer_step(self) -> bool: device=binding.collective_device, ) controls = [ - torch.empty_like(amp_control) - for _ in binding.identity.ordered_members + torch.empty_like(amp_control) for _ in binding.identity.ordered_members ] dist.all_gather(controls, amp_control, group=binding.process_group) if any(not torch.equal(item, controls[0]) for item in controls[1:]): @@ -4686,11 +5113,7 @@ def _codebook_value_fingerprint(self): if self._gefen_codebook is None: return (0, 0, 0, 0) raw = bytes( - self._gefen_codebook.detach() - .cpu() - .contiguous() - .view(torch.uint8) - .tolist() + self._gefen_codebook.detach().cpu().contiguous().view(torch.uint8).tolist() ) return self._sha256_int64(raw) @@ -4804,8 +5227,7 @@ 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 parameters require every nonempty shard to agree on gradient presence" ) continue replicated_controls = { @@ -4825,9 +5247,11 @@ def _verify_codebook_scope_agreement(self, codebook: torch.Tensor) -> None: self._assert_runtime_codebook_process_group() import torch.distributed as dist - local = codebook.detach().to( - device=binding.collective_device, dtype=torch.float32 - ).contiguous() + local = ( + codebook.detach() + .to(device=binding.collective_device, dtype=torch.float32) + .contiguous() + ) gathered = [torch.empty_like(local) for _ in binding.identity.ordered_members] dist.all_gather(gathered, local, group=binding.process_group) if any(not torch.equal(item, gathered[0]) for item in gathered[1:]): @@ -4859,9 +5283,7 @@ def _ensure_codebook_scope_agreement(self) -> None: dtype=torch.int64, device=binding.collective_device, ) - controls = [ - torch.empty_like(control) for _ in binding.identity.ordered_members - ] + controls = [torch.empty_like(control) for _ in binding.identity.ordered_members] dist.all_gather(controls, control, group=binding.process_group) if any(not torch.equal(item, controls[0]) for item in controls[1:]): raise RuntimeError( @@ -4987,7 +5409,13 @@ def _ensure_gefen_codebook(self, reuse_existing_periods: bool = False) -> None: ) def _step_automatic_factored( - self, group, param_name: str, p: torch.Tensor, grad: torch.Tensor + self, + group, + param_name: str, + p: torch.Tensor, + grad: torch.Tensor, + *, + state=None, ) -> None: # Adafactor-style factored second moment for a 2D param (opt-in via # factored_v_2d). The quantized-momentum machinery is byte-identical to @@ -5002,7 +5430,7 @@ def _step_automatic_factored( grad = grad.to_local() if hasattr(grad, "wait"): grad = grad.wait() - state = self.state[p] + state = self.state[p] if state is None else state beta1 = group["beta1"] beta2 = group["beta2"] lr = group["lr"] @@ -5018,9 +5446,7 @@ def _step_automatic_factored( elif flat_grad.numel() == 1: automatic_period = 1 else: - automatic_period = self._resolve_automatic_period( - param_name, p, grad - ) + automatic_period = self._resolve_automatic_period(param_name, p, grad) if flat_grad.numel() % automatic_period != 0: raise ValueError( "Automatic partition period {} does not divide parameter {} with numel {}".format( @@ -5078,8 +5504,7 @@ def _step_automatic_factored( codebook = self._gefen_codebook_on(p.device) if codebook is None: raise ValueError( - "Expected Gefen codebook to be initialized before the " - "factored update." + "Expected Gefen codebook to be initialized before the factored update." ) # The raw kernel flat-indexes a contiguous matrix. A strided leaf # Parameter is uncommon but valid (and can arise from model surgery @@ -5237,14 +5662,10 @@ def _refresh_codebook_with_requant(self) -> None: return staged_indices = [] if self._gefen_codebook_process_group is None: - parameters = [ - p for pgroup in self.param_groups for p in pgroup["params"] - ] + parameters = [p for pgroup in self.param_groups for p in pgroup["params"]] else: parameters = [ - p - for p, _ in self._gefen_local_shard_bindings - if p is not None + p for p, _ in self._gefen_local_shard_bindings if p is not None ] try: for p in parameters: @@ -5300,9 +5721,8 @@ def _maybe_refresh_gefen_codebook(self) -> None: binding = self._gefen_codebook_process_group if ( - (binding is None or len(binding.identity.ordered_members) == 1) - and not self._has_local_codebook_gradients() - ): + binding is None or len(binding.identity.ordered_members) == 1 + ) and not self._has_local_codebook_gradients(): # No local histogram exists, so learning would return None. This # fast no-op is also what keeps a no-gradient capturable step free # of the host-driven exact-DP preparation path. Multi-member scopes @@ -5338,9 +5758,7 @@ def initialize_codebook(self) -> bool: except Exception as exc: local_error = exc if self._gefen_codebook_process_group is not None: - self._synchronize_codebook_scope_failure( - local_error, "gradient preflight" - ) + self._synchronize_codebook_scope_failure(local_error, "gradient preflight") elif local_error is not None: raise local_error self._ensure_codebook_scope_agreement() @@ -5376,9 +5794,7 @@ def refresh_codebook(self) -> bool: except Exception as exc: local_error = exc if self._gefen_codebook_process_group is not None: - self._synchronize_codebook_scope_failure( - local_error, "gradient preflight" - ) + self._synchronize_codebook_scope_failure(local_error, "gradient preflight") elif local_error is not None: raise local_error self._ensure_codebook_scope_agreement() @@ -5694,9 +6110,7 @@ def _automatic_momentum_update( * m_magnitude[r0:r1] ) # lerp() needs matching dtypes, so promote the (bf16) grad rows to fp32. - updated_m = current_m.lerp( - grad_view[r0:r1].to(current_m.dtype), 1 - beta1 - ) + updated_m = current_m.lerp(grad_view[r0:r1].to(current_m.dtype), 1 - beta1) # New per-block magnitude = max |updated_m| over each row, written back # into the persistent fp32 m_magnitude state in place. @@ -5780,7 +6194,13 @@ def _automatic_momentum_update_merged( return quantized_m * state["m_magnitude"] def _step_automatic( - self, group, param_name: str, p: torch.Tensor, grad: torch.Tensor + self, + group, + param_name: str, + p: torch.Tensor, + grad: torch.Tensor, + *, + state=None, ) -> None: # Unwrap the gradient to its local shard for both the fused and # non-fused paths. Under FSDP2 the gradient arrives as a sharded DTensor; @@ -5792,7 +6212,7 @@ def _step_automatic( if hasattr(grad, "wait"): grad = grad.wait() - state = self.state[p] + state = self.state[p] if state is None else state beta1 = group["beta1"] beta2 = group["beta2"] lr = group["lr"] @@ -5823,9 +6243,7 @@ def _step_automatic( elif local_numel == 1: automatic_period = 1 elif local_numel > 1: - automatic_period = self._resolve_automatic_period( - param_name, p, grad - ) + automatic_period = self._resolve_automatic_period(param_name, p, grad) else: raise ValueError( "Automatic partition received an empty parameter {}".format( @@ -5875,9 +6293,7 @@ def _step_automatic( # A capturable (tensor) step must be CLONED -- assigning the same # 0-dim tensor would alias the two counters and double-increment. step = state["step"] - state["vmean_step"] = ( - step.clone() if torch.is_tensor(step) else step - ) + state["vmean_step"] = step.clone() if torch.is_tensor(step) else step # Tier-1 uses a separately calibrated predicate for the current full # v1/v2 kernels. Both fold the vmean EMA, per-block stepsize math, and @@ -5905,8 +6321,12 @@ def _step_automatic( # Deterministic mode keeps the fused update and selects v1-full, whose # fixed thread tree produces bit-identical vmean/parameter writes for # identical inputs on a homogeneous architecture. - route_v2 = False if self._deterministic else _should_use_v2_full( - grad_view.shape[0], automatic_period, grad_view.device + route_v2 = ( + False + if self._deterministic + else _should_use_v2_full( + grad_view.shape[0], automatic_period, grad_view.device + ) ) if ( route_v2 @@ -5938,9 +6358,7 @@ def _step_automatic( # local shard's storage. DTensor.to_local() returns a view of that # storage, so in-place ops propagate back to the param. p_local = ( - p.to_local() - if (hasattr(p, "to_local") and hasattr(p, "placements")) - else p + p.to_local() if (hasattr(p, "to_local") and hasattr(p, "placements")) else p ) if self.capturable and (use_full_fused or use_v2_full): @@ -5978,8 +6396,7 @@ def _step_automatic( # contiguous; a non-contiguous shard is updated on a contiguous # copy that is copied back to preserve the in-place semantics. weight_decay_factor = ( - 1.0 if step_scalars is not None - else 1.0 - lr * group["weight_decay"] + 1.0 if step_scalars is not None else 1.0 - lr * group["weight_decay"] ) if p_local.is_contiguous(): self._automatic_gefen_fused_full_update( @@ -6019,8 +6436,7 @@ def _step_automatic( # (see _automatic_gefen_fused_update_v2_full). Same contiguity # handling as the v1-full branch. weight_decay_factor = ( - 1.0 if step_scalars is not None - else 1.0 - lr * group["weight_decay"] + 1.0 if step_scalars is not None else 1.0 - lr * group["weight_decay"] ) if p_local.is_contiguous(): self._automatic_gefen_fused_update_v2_full( @@ -6254,7 +6670,7 @@ def _step_automatic_merged(self, items) -> None: if weight_decay > 0.0: torch._foreach_mul_(params, 1 - lr * weight_decay) - bias_correction_1 = 1 - beta1 ** step_count + bias_correction_1 = 1 - beta1**step_count bias_correction_2 = 1 - beta2 ** first_state["vmean_step"] h = (merged_vmean.sqrt() / math.sqrt(bias_correction_2)).add_(eps) stepsize = (1 / bias_correction_1) / h @@ -6277,7 +6693,9 @@ def _step_automatic_merged(self, items) -> None: update_views = update.reshape(k, nblocks, period) param_updates = [update_views[i].reshape(params[i].shape) for i in range(k)] torch._foreach_add_(params, param_updates, alpha=-1.0) - torch._foreach_copy_(vmeans, list(merged_vmean.reshape(k, nblocks, 1).unbind(0))) + torch._foreach_copy_( + vmeans, list(merged_vmean.reshape(k, nblocks, 1).unbind(0)) + ) torch._foreach_copy_(mags, list(merged_mag.reshape(k, nblocks, 1).unbind(0))) torch._foreach_copy_( codebooks, list(merged_codebook.reshape(k, nblocks, period).unbind(0)) @@ -6297,13 +6715,10 @@ def _canonical_common_global_step(self) -> int: device_step = self._device_gefen_global_step() step = self._gefen_global_step if device_step is None else device_step if self.capturable and self._stochastic_round: - seed_steps = [ - int(seed.item()) for seed in self._sr_seed_by_device.values() - ] + seed_steps = [int(seed.item()) for seed in self._sr_seed_by_device.values()] if any(seed_step != step for seed_step in seed_steps): raise RuntimeError( - "Gefen capturable stochastic-round seeds disagree with the " - "canonical optimizer global step" + "Gefen capturable stochastic-round seeds disagree with the canonical optimizer global step" ) return step @@ -6385,12 +6800,9 @@ def _canonical_import_live_token(self): id(group), tuple(id(parameter) for parameter in group["params"]), tuple( - str(name) - for name, _ in self._iter_group_params_with_names(group) - ), - self._canonical_value_token( - self._canonical_group_options_value(group) + str(name) for name, _ in self._iter_group_params_with_names(group) ), + self._canonical_value_token(self._canonical_group_options_value(group)), ) for group in self.param_groups ) @@ -6421,6 +6833,8 @@ def _canonical_import_live_token(self): self.fused, self.verbose, self._fused_build_ok, + self.state_offload_device, + self.state_offload_poisoned, id(self._gefen_codebook_process_group), self._canonical_value_token(self._serialized_codebook_scope()), id(self._gefen_sharding_manifest), @@ -6485,9 +6899,7 @@ def _canonical_period_must_be_one(self, parameter, shard) -> bool: if self._force_2d_period_one and parameter.ndim == 2: return True fqn = shard.parameter.fqn.lower() - return any( - substring in fqn for substring in self._period_one_substrings - ) + return any(substring in fqn for substring in self._period_one_substrings) def _validate_canonical_parameter_semantics( self, @@ -6519,8 +6931,9 @@ def _validate_canonical_parameter_semantics( counters[counter_key] = counter if counter > global_step: raise ValueError( - "Gefen canonical parameter counter {} exceeds the optimizer " - "global step".format(counter_key) + "Gefen canonical parameter counter {} exceeds the optimizer global step".format( + counter_key + ) ) if "step" in counters and any( counters[counter_key] > counters["step"] @@ -6556,17 +6969,14 @@ def _validate_canonical_parameter_semantics( shard, group_options, parameter_state, state_layout ): raise ValueError( - "Gefen canonical parameter state does not match a declared state " - "variant" + "Gefen canonical parameter state does not match a declared state variant" ) @staticmethod def _assert_canonical_state_outside_cuda_capture(operation) -> None: if torch.cuda.is_available() and torch.cuda.is_current_stream_capturing(): raise RuntimeError( - "canonical state {} cannot run during CUDA capture".format( - operation - ) + "canonical state {} cannot run during CUDA capture".format(operation) ) def _assert_canonical_import_target_safe(self) -> None: @@ -6594,11 +7004,11 @@ def export_canonical_state(self): """Export an exact-binding, device-neutral local state fragment.""" self._assert_finalized_binding_layout() + self._assert_state_export_safe() self._assert_canonical_state_outside_cuda_capture("export") if not self._canonical_state_layouts(): raise RuntimeError( - "canonical state export requires a supported finalized identity " - "layout and primitive algorithm policy" + "canonical state export requires a supported finalized identity layout and primitive algorithm policy" ) entries = self._canonical_live_entries() parameters = {} @@ -6612,8 +7022,9 @@ def export_canonical_state(self): continue if key not in _CANONICAL_PARAMETER_STATE_KEYS: raise RuntimeError( - "canonical state export found undeclared parameter state key " - "{!r}".format(key) + "canonical state export found undeclared parameter state key {!r}".format( + key + ) ) state[key] = clone_canonical_value( value, path="parameters.{}.state.{}".format(fqn, key) @@ -6631,9 +7042,7 @@ def export_canonical_state(self): "format_version": CANONICAL_STATE_FORMAT_VERSION, "coverage": "local_optimizer_fragment", "implementation": self.optimizer_contract().implementation, - "policy": clone_canonical_value( - self._canonical_policy(), path="policy" - ), + "policy": clone_canonical_value(self._canonical_policy(), path="policy"), "common": { "gefen_global_step": self._canonical_common_global_step(), "gefen_codebook": clone_canonical_value( @@ -6721,8 +7130,7 @@ def _normalize_canonical_state_import(self, state): if type(manifest) is not list: raise ValueError("Gefen canonical manifest must be a list") normalized_manifest = [ - self._normalize_serialized_canonical_shard(record) - for record in manifest + self._normalize_serialized_canonical_shard(record) for record in manifest ] if normalized_manifest != self._serialized_sharding_manifest(): raise ValueError( @@ -6747,14 +7155,10 @@ def _normalize_canonical_state_import(self, state): "state", }: raise ValueError( - "Gefen canonical parameter {!r} has an invalid schema".format( - fqn - ) + "Gefen canonical parameter {!r} has an invalid schema".format(fqn) ) if type(record["compatibility_name"]) is not str: - raise ValueError( - "Gefen canonical compatibility names must be strings" - ) + raise ValueError("Gefen canonical compatibility names must be strings") parameter, shard, _, live_options = live_entries[fqn] normalized_shard = self._normalize_serialized_canonical_shard( record["shard"] @@ -6887,9 +7291,7 @@ def commit_canonical_state_import(self, prepared) -> None: def import_canonical_state(self, state) -> None: """Atomically prepare and commit an exact-binding canonical fragment.""" - self.commit_canonical_state_import( - self.prepare_canonical_state_import(state) - ) + self.commit_canonical_state_import(self.prepare_canonical_state_import(state)) def export_portable_state( self, @@ -6936,9 +7338,11 @@ def state_dict(self): """Run optimizer state-dict hooks around Gefen's complete schema.""" self._assert_finalized_binding_layout() + self._assert_state_export_safe() for pre_hook in self._optimizer_state_dict_pre_hooks.values(): pre_hook(self) self._assert_finalized_binding_layout() + 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) @@ -7027,10 +7431,12 @@ def reject(message): raise RuntimeError(message) return None - if not torch.distributed.is_available() or not torch.distributed.is_initialized(): + if ( + not torch.distributed.is_available() + or not torch.distributed.is_initialized() + ): return reject( - "Gefen rank-local DTensor checkpointing requires an initialized " - "default distributed process group" + "Gefen rank-local DTensor checkpointing requires an initialized default distributed process group" ) import torch.distributed as dist @@ -7187,9 +7593,7 @@ def _rank_local_sharded_signature(self, context=None): for item in p.placements ], "mesh": self._device_mesh_checkpoint_signature(mesh), - "coordinate": None - if coordinate is None - else list(coordinate), + "coordinate": None if coordinate is None else list(coordinate), "sharded_mode": group.get("sharded_mode"), } ) @@ -7281,13 +7685,15 @@ def _unpack_checkpoint_payload(cls, packed, tensors): def _serialize_rank_local_payload(payload) -> torch.Tensor: buffer = io.BytesIO() torch.save(payload, buffer) - return torch.frombuffer( - bytearray(buffer.getvalue()), dtype=torch.uint8 - ).clone() + return torch.frombuffer(bytearray(buffer.getvalue()), dtype=torch.uint8).clone() @staticmethod def _deserialize_rank_local_payload(payload: torch.Tensor): - if not torch.is_tensor(payload) or payload.dtype != torch.uint8 or payload.dim() != 1: + if ( + not torch.is_tensor(payload) + or payload.dtype != torch.uint8 + or payload.dim() != 1 + ): raise ValueError( "Gefen rank-local checkpoint payload must be a 1-D uint8 tensor" ) @@ -7369,24 +7775,27 @@ def _consolidate_rank_local_sharded_state( ) if any(control["manifest"] != local_manifest for control in controls): raise RuntimeError( - "Gefen DTensor checkpoint parameter identifiers, names, dtypes, " - "or global shapes differ across ranks" + "Gefen DTensor checkpoint parameter identifiers, names, dtypes, or global shapes differ across ranks" ) global_steps = [control["global_step"] for control in controls] - if any(type(step) is not int for step in global_steps) or len( - set(global_steps) - ) != 1: + if ( + any(type(step) is not int for step in global_steps) + or len(set(global_steps)) != 1 + ): raise RuntimeError( - "Gefen rank-local checkpoint global_step differs across ranks: " - "{}".format(global_steps) + "Gefen rank-local checkpoint global_step differs across ranks: {}".format( + global_steps + ) ) deterministic_values = [control["deterministic"] for control in controls] - if any(type(value) is not bool for value in deterministic_values) or len( - set(deterministic_values) - ) != 1: + if ( + any(type(value) is not bool for value in deterministic_values) + or len(set(deterministic_values)) != 1 + ): raise RuntimeError( - "Gefen rank-local checkpoint deterministic policy differs across " - "ranks: {}".format(deterministic_values) + "Gefen rank-local checkpoint deterministic policy differs across ranks: {}".format( + deterministic_values + ) ) signatures = { @@ -7426,8 +7835,9 @@ def _consolidate_rank_local_sharded_state( dist.all_gather_object(serialization_statuses, serialization_status) if any(not item.get("ok", False) for item in serialization_statuses): raise RuntimeError( - "Gefen rank-local checkpoint payload serialization failed across " - "ranks: {}".format(serialization_statuses) + "Gefen rank-local checkpoint payload serialization failed across ranks: {}".format( + serialization_statuses + ) ) serialized_payloads = { str(src): self._broadcast_checkpoint_payload( @@ -7506,9 +7916,7 @@ def _state_dict_impl(self, *, consolidate_rank_local: bool = True): # membership matches the id()-keyed param_mappings the base class # builds. The no-orphan path (every non-wrapped run) takes the plain # super() call, bit-for-bit unchanged. - reachable = { - param for group in self.param_groups for param in group["params"] - } + reachable = {param for group in self.param_groups for param in group["params"]} orphaned = [param for param in self.state if param not in reachable] if orphaned: original_order = list(self.state) @@ -7557,8 +7965,7 @@ def _state_dict_impl(self, *, consolidate_rank_local: bool = True): def _is_scratch_key(key): return key in scratch_keys or ( - isinstance(key, str) - and key.startswith(_RANK_LOCAL_PAYLOAD_KEY_PREFIX) + isinstance(key, str) and key.startswith(_RANK_LOCAL_PAYLOAD_KEY_PREFIX) ) def _compact(value): @@ -7579,11 +7986,7 @@ def _compact(value): state_dict["state"] = { pid: ( - { - k: _compact(v) - for k, v in pstate.items() - if not _is_scratch_key(k) - } + {k: _compact(v) for k, v in pstate.items() if not _is_scratch_key(k)} if isinstance(pstate, dict) else pstate ) @@ -7615,9 +8018,7 @@ def _compact(value): # handing groups to torch, so it never leaks into live scheduler groups. checkpoint_metadata = { "format_version": ( - _SCOPED_NATIVE_METADATA_VERSION - if codebook_scope is not None - else 1 + _SCOPED_NATIVE_METADATA_VERSION if codebook_scope is not None else 1 ), "global_step": self._gefen_global_step, "codebook": self._gefen_codebook, @@ -7635,9 +8036,7 @@ def _compact(value): if native_local_shards is not None: checkpoint_metadata["native_local_shards"] = native_local_shards if consolidate_rank_local: - self._consolidate_rank_local_sharded_state( - state_dict, checkpoint_metadata - ) + self._consolidate_rank_local_sharded_state(state_dict, checkpoint_metadata) for group in state_dict["param_groups"]: group["_gefen_checkpoint_metadata"] = checkpoint_metadata return state_dict @@ -7745,7 +8144,23 @@ def _stage_load_state_dict(self, state_dict): staged._capt_stacks = None staged._load_state_dict_impl(state_dict) + if staged.state_offload_active: + reason = staged._state_offload_rejection_reason( + require_cpu_state=False, + allow_poisoned=True, + ) + if reason is not None: + raise RuntimeError( + "Gefen could not preserve active state offload while loading: {}".format( + reason + ) + ) + staged.state = staged._stage_all_parameter_state_to_cpu() + staged._gefen_codebook = staged._stage_state_offload_resident_codebook() staged._validate_loaded_native_state() + # A complete successfully validated load is the only operation that can + # re-establish known-good optimizer state after a failed copyback. + staged._gefen_state_offload_poisoned = False return staged def _commit_staged_load_state_dict(self, staged) -> None: @@ -7762,9 +8177,7 @@ def _commit_staged_load_state_dict(self, staged) -> None: def _validate_loaded_native_state(self) -> None: """Validate the complete prepared native state before publication.""" - self._validate_rank_local_counter( - "gefen_global_step", self._gefen_global_step - ) + self._validate_rank_local_counter("gefen_global_step", self._gefen_global_step) signature = self._rank_local_sharded_signature(context={}) states = [ self.state[param] @@ -7777,8 +8190,7 @@ def _validate_loaded_native_state(self) -> None: self._gefen_codebook, allow_legacy_vmean_counter=True, allow_preinitialized_periods=( - self._gefen_global_step == 0 - and self._gefen_codebook is not None + self._gefen_global_step == 0 and self._gefen_codebook is not None ), ) for group in self.param_groups: @@ -7800,6 +8212,70 @@ def _base_load_state_dict_without_hooks(self, state_dict): self._optimizer_load_state_dict_pre_hooks = pre_hooks self._optimizer_load_state_dict_post_hooks = post_hooks + def _base_load_state_dict_to_offload_cpu(self, state_dict) -> None: + """Apply the base optimizer mapping while keeping parameter state on CPU.""" + + groups = self.param_groups + saved_groups = deepcopy(state_dict["param_groups"]) + if len(groups) != len(saved_groups): + raise ValueError( + "loaded state dict has a different number of parameter groups" + ) + if any( + len(group["params"]) != len(saved_group["params"]) + for group, saved_group in zip(groups, saved_groups) + ): + raise ValueError( + "loaded state dict contains a parameter group that doesn't match the size of optimizer's group" + ) + + id_map = dict( + zip( + chain.from_iterable(group["params"] for group in saved_groups), + chain.from_iterable(group["params"] for group in groups), + ) + ) + + def clone_to_cpu(value): + if torch.is_tensor(value): + return value.to( + device="cpu", + dtype=value.dtype, + non_blocking=False, + copy=True, + memory_format=torch.contiguous_format, + ).detach() + if isinstance(value, dict): + return {key: clone_to_cpu(item) for key, item in value.items()} + if type(value) is list: + return [clone_to_cpu(item) for item in value] + if type(value) is tuple: + return tuple(clone_to_cpu(item) for item in value) + if type(value) is set: + return {clone_to_cpu(item) for item in value} + if type(value) is frozenset: + return frozenset(clone_to_cpu(item) for item in value) + if type(value) is deque: + return deque( + (clone_to_cpu(item) for item in value), maxlen=value.maxlen + ) + return deepcopy(value) + + loaded_state = defaultdict(dict) + for key, value in state_dict["state"].items(): + if key in id_map: + loaded_state[id_map[key]] = clone_to_cpu(value) + else: + loaded_state[key] = value + + param_groups = [] + for live_group, saved_group in zip(groups, saved_groups): + saved_group["params"] = live_group["params"] + if "param_names" in live_group and "param_names" not in saved_group: + saved_group["param_names"] = live_group["param_names"] + param_groups.append(saved_group) + self.__setstate__({"state": loaded_state, "param_groups": param_groups}) + @staticmethod def _validate_rank_local_codebook(codebook, *, required: bool) -> None: if codebook is None: @@ -7880,17 +8356,14 @@ def _validate_rank_local_states( expected_name = param_signature["name"] if pstate.get("name") != expected_name: raise ValueError( - "Gefen rank-local checkpoint state name/order differs: " - "checkpoint={!r} expected={!r}".format( + "Gefen rank-local checkpoint state name/order differs: checkpoint={!r} expected={!r}".format( pstate.get("name"), expected_name ) ) counters = {} for key in counter_keys: if key in pstate: - counters[key] = self._validate_rank_local_counter( - key, pstate[key] - ) + counters[key] = self._validate_rank_local_counter(key, pstate[key]) local_shape = param_signature["local_shape"] state_shape = ( @@ -7907,8 +8380,7 @@ def _validate_rank_local_states( ) if state_numel == 0 or state_numel % period != 0: raise ValueError( - "Gefen rank-local checkpoint automatic_period does not divide " - "the parameter state geometry" + "Gefen rank-local checkpoint automatic_period does not divide the parameter state geometry" ) momentum_keys = ("m_codebook", "m_magnitude") @@ -7925,8 +8397,7 @@ def _validate_rank_local_states( ) if any(key in pstate for key in initialized_keys) and not carries_momentum: raise ValueError( - "Gefen rank-local checkpoint initialized state is missing " - "quantized momentum" + "Gefen rank-local checkpoint initialized state is missing quantized momentum" ) if ( period is not None @@ -7941,8 +8412,7 @@ def _validate_rank_local_states( ) ): raise ValueError( - "Gefen rank-local checkpoint automatic_period is invalid without " - "initialized momentum" + "Gefen rank-local checkpoint automatic_period is invalid without initialized momentum" ) if carries_momentum: has_quantized_momentum = True @@ -7977,8 +8447,7 @@ def _validate_rank_local_states( or not bool((magnitude >= 0).all()) ): raise ValueError( - "Gefen rank-local checkpoint m_magnitude geometry/dtype/values " - "are invalid" + "Gefen rank-local checkpoint m_magnitude geometry/dtype/values are invalid" ) vmean = pstate.get("vmean") if vmean is not None and ( @@ -7989,8 +8458,7 @@ def _validate_rank_local_states( or not bool((vmean >= 0).all()) ): raise ValueError( - "Gefen rank-local checkpoint vmean geometry/dtype/values " - "are invalid" + "Gefen rank-local checkpoint vmean geometry/dtype/values are invalid" ) has_vmean = "vmean" in pstate @@ -7999,14 +8467,9 @@ def _validate_rank_local_states( raise ValueError( "Gefen rank-local checkpoint block second moment is incomplete" ) - if ( - has_vmean - and not has_vmean_step - and not allow_legacy_vmean_counter - ): + if has_vmean and not has_vmean_step and not allow_legacy_vmean_counter: raise ValueError( - "Gefen rank-local checkpoint block second moment is missing " - "vmean_step" + "Gefen rank-local checkpoint block second moment is missing vmean_step" ) # GefenMuon groups carry a sharded_mode and intentionally use @@ -8017,26 +8480,19 @@ def _validate_rank_local_states( if carries_factored: if "factored_step" not in pstate: raise ValueError( - "Gefen rank-local checkpoint factored state is missing " - "factored_step" + "Gefen rank-local checkpoint factored state is missing factored_step" ) if counters["factored_step"] < 1: raise ValueError( - "Gefen rank-local checkpoint initialized factored " - "state requires factored_step >= 1" + "Gefen rank-local checkpoint initialized factored state requires factored_step >= 1" ) elif "vmean" not in pstate: raise ValueError( - "Gefen rank-local checkpoint plain momentum is missing a " - "second moment" + "Gefen rank-local checkpoint plain momentum is missing a second moment" ) - elif ( - "vmean_step" in counters - and counters["vmean_step"] < 1 - ): + elif "vmean_step" in counters and counters["vmean_step"] < 1: raise ValueError( - "Gefen rank-local checkpoint initialized vmean requires " - "vmean_step >= 1" + "Gefen rank-local checkpoint initialized vmean requires vmean_step >= 1" ) factored = (pstate.get("v_row"), pstate.get("v_col")) @@ -8067,8 +8523,9 @@ def _validate_rank_local_states( or not bool((tensor >= 0).all()) ): raise ValueError( - "Gefen rank-local checkpoint {} geometry/dtype/values " - "are invalid".format(key) + "Gefen rank-local checkpoint {} geometry/dtype/values are invalid".format( + key + ) ) has_normuon_v = "normuon_v" in pstate @@ -8080,18 +8537,15 @@ def _validate_rank_local_states( if has_normuon_v: if param_signature.get("sharded_mode") is None: raise ValueError( - "Gefen rank-local checkpoint NorMuon state is invalid for " - "a plain Gefen parameter" + "Gefen rank-local checkpoint NorMuon state is invalid for a plain Gefen parameter" ) if not carries_momentum: raise ValueError( - "Gefen rank-local checkpoint NorMuon state is missing " - "initialized momentum" + "Gefen rank-local checkpoint NorMuon state is missing initialized momentum" ) if len(state_shape) != 2: raise ValueError( - "Gefen rank-local checkpoint NorMuon state requires a 2-D " - "parameter" + "Gefen rank-local checkpoint NorMuon state requires a 2-D parameter" ) normuon_v = pstate["normuon_v"] if ( @@ -8102,18 +8556,14 @@ def _validate_rank_local_states( or not bool((normuon_v >= 0).all()) ): raise ValueError( - "Gefen rank-local checkpoint normuon_v geometry/dtype/values " - "are invalid" + "Gefen rank-local checkpoint normuon_v geometry/dtype/values are invalid" ) if counters["normuon_step"] < 1: raise ValueError( - "Gefen rank-local checkpoint initialized NorMuon state " - "requires normuon_step >= 1" + "Gefen rank-local checkpoint initialized NorMuon state requires normuon_step >= 1" ) - self._validate_rank_local_codebook( - codebook, required=has_quantized_momentum - ) + self._validate_rank_local_codebook(codebook, required=has_quantized_momentum) def _unwrap_rank_local_sharded_checkpoint(self, state_dict) -> None: groups = state_dict.get("param_groups", ()) @@ -8145,8 +8595,7 @@ def _unwrap_rank_local_sharded_checkpoint(self, state_dict) -> None: if not markers: if wrapped_state: raise ValueError( - "Gefen checkpoint carries rank-local DTensor payloads but " - "is missing their topology metadata" + "Gefen checkpoint carries rank-local DTensor payloads but is missing their topology metadata" ) has_quantized_momentum = any( isinstance(pstate, dict) and "m_codebook" in pstate @@ -8158,15 +8607,12 @@ def _unwrap_rank_local_sharded_checkpoint(self, state_dict) -> None: "rank-local DTensor optimizer state. Older generic FSDP/DCP " "full optimizer checkpoints kept rank 0's codebook and block " "state for every rank. Re-save from the original run with a " - "Gefen version that writes {} payloads.".format( - _RANK_LOCAL_FORMAT - ) + "Gefen version that writes {} payloads.".format(_RANK_LOCAL_FORMAT) ) return if len(markers) != len(groups) or any(item is None for item in metadata): raise ValueError( - "Gefen rank-local DTensor checkpoint metadata is present on only " - "some parameter groups" + "Gefen rank-local DTensor checkpoint metadata is present on only some parameter groups" ) marker = markers[0] if not isinstance(marker, dict): @@ -8183,8 +8629,7 @@ def _unwrap_rank_local_sharded_checkpoint(self, state_dict) -> None: ) if not self._uses_rank_local_sharded_state(): raise ValueError( - "A rank-local DTensor Gefen checkpoint can only load into the " - "same sharded optimizer topology" + "A rank-local DTensor Gefen checkpoint can only load into the same sharded optimizer topology" ) try: context = self._rank_local_checkpoint_context() @@ -8194,8 +8639,7 @@ def _unwrap_rank_local_sharded_checkpoint(self, state_dict) -> None: rank = context["global_rank"] if marker.get("world_size") != world: raise ValueError( - "Gefen rank-local DTensor checkpoints require the same world " - "size; checkpoint={} current={}".format( + "Gefen rank-local DTensor checkpoints require the same world size; checkpoint={} current={}".format( marker.get("world_size"), world ) ) @@ -8212,8 +8656,7 @@ def _unwrap_rank_local_sharded_checkpoint(self, state_dict) -> None: current_signature = self._rank_local_sharded_signature(context) if signatures[str(rank)] != current_signature: raise ValueError( - "Gefen rank-local DTensor checkpoint topology differs on rank {}: " - "checkpoint={!r} current={!r}".format( + "Gefen rank-local DTensor checkpoint topology differs on rank {}: checkpoint={!r} current={!r}".format( rank, signatures[str(rank)], current_signature ) ) @@ -8261,14 +8704,19 @@ def _unwrap_rank_local_sharded_checkpoint(self, state_dict) -> None: if isinstance(key, str) and key.startswith(_RANK_LOCAL_PAYLOAD_KEY_PREFIX) } - extras = set(pstate) - carrier_keys - { - _RANK_LOCAL_MEMBER_KEY, - "name", - } + extras = ( + set(pstate) + - carrier_keys + - { + _RANK_LOCAL_MEMBER_KEY, + "name", + } + ) if extras: raise ValueError( - "Gefen rank-local checkpoint parameter state has an invalid " - "schema: {!r}".format(pstate) + "Gefen rank-local checkpoint parameter state has an invalid schema: {!r}".format( + pstate + ) ) if pstate.get("name") != 0: raise ValueError( @@ -8283,17 +8731,14 @@ def _unwrap_rank_local_sharded_checkpoint(self, state_dict) -> None: if has_carrier: if carrier_keys != expected_carrier_keys: raise ValueError( - "Gefen rank-local checkpoint carrier has missing or " - "unexpected per-rank payload keys" + "Gefen rank-local checkpoint carrier has missing or unexpected per-rank payload keys" ) if serialized_payloads is not None: raise ValueError( "Gefen rank-local checkpoint repeats its payload transport" ) serialized_payloads = { - str(global_rank): pstate[ - _rank_local_payload_key(global_rank) - ] + str(global_rank): pstate[_rank_local_payload_key(global_rank)] for global_rank in context["world_ranks"] } elif pstate.get(_RANK_LOCAL_MEMBER_KEY) is not True: @@ -8323,8 +8768,13 @@ def _unwrap_rank_local_sharded_checkpoint(self, state_dict) -> None: "states", "codebook", } - if not isinstance(selected_payload, dict) or set(selected_payload) != expected_payload_keys: - raise ValueError("Gefen rank-local checkpoint payload has an invalid schema") + if ( + not isinstance(selected_payload, dict) + or set(selected_payload) != expected_payload_keys + ): + raise ValueError( + "Gefen rank-local checkpoint payload has an invalid schema" + ) payload_identity = ( selected_payload["format"] == _RANK_LOCAL_FORMAT and selected_payload["global_rank"] == rank @@ -8396,12 +8846,8 @@ def _load_state_dict_impl(self, state_dict): gefen_codebook_scope = self._normalize_serialized_codebook_scope( gefen_codebook_scope ) - has_top_level_native_local_shards = ( - "gefen_native_local_shards" in state_dict - ) - native_local_shards = state_dict.pop( - "gefen_native_local_shards", None - ) + has_top_level_native_local_shards = "gefen_native_local_shards" in state_dict + native_local_shards = state_dict.pop("gefen_native_local_shards", None) if has_top_level_native_local_shards: native_local_shards = self._normalize_serialized_native_local_shards( native_local_shards @@ -8410,8 +8856,9 @@ def _load_state_dict_impl(self, state_dict): gefen_deterministic = state_dict.pop("gefen_deterministic", None) if has_top_level_deterministic and type(gefen_deterministic) is not bool: raise ValueError( - "Gefen checkpoint top-level deterministic policy must be a bool, " - "got {!r}".format(gefen_deterministic) + "Gefen checkpoint top-level deterministic policy must be a bool, got {!r}".format( + gefen_deterministic + ) ) if group_metadata: if len(group_metadata) != len(state_dict["param_groups"]): @@ -8446,32 +8893,28 @@ def _load_state_dict_impl(self, state_dict): and type(metadata["deterministic"]) is not bool ): raise ValueError( - "Gefen checkpoint parameter-group deterministic policy " - "must be a bool, got {!r}".format(metadata["deterministic"]) + "Gefen checkpoint parameter-group deterministic policy must be a bool, got {!r}".format( + metadata["deterministic"] + ) ) normalized_metadata_native_local_shards.append( self._normalize_serialized_native_local_shards( metadata.get("native_local_shards") ) ) - for metadata_index, metadata in enumerate( - group_metadata[1:], start=1 - ): - same_version = metadata.get( + for metadata_index, metadata in enumerate(group_metadata[1:], start=1): + same_version = metadata.get("format_version") == first_metadata.get( "format_version" - ) == first_metadata.get("format_version") + ) same_step = metadata.get("global_step") == first_metadata.get( "global_step" ) left_codebook = metadata.get("codebook") right_codebook = first_metadata.get("codebook") - same_codebook = ( - left_codebook is right_codebook - or ( - torch.is_tensor(left_codebook) - and torch.is_tensor(right_codebook) - and torch.equal(left_codebook, right_codebook) - ) + same_codebook = left_codebook is right_codebook or ( + torch.is_tensor(left_codebook) + and torch.is_tensor(right_codebook) + and torch.equal(left_codebook, right_codebook) ) same_deterministic = metadata.get( "deterministic" @@ -8492,8 +8935,7 @@ def _load_state_dict_impl(self, state_dict): or not same_native_local_shards ): raise ValueError( - "Gefen checkpoint parameter groups carry inconsistent " - "optimizer metadata" + "Gefen checkpoint parameter groups carry inconsistent optimizer metadata" ) metadata_step = first_metadata.get("global_step", 0) metadata_codebook = first_metadata.get("codebook") @@ -8501,9 +8943,7 @@ def _load_state_dict_impl(self, state_dict): metadata_codebook_scope = self._normalize_serialized_codebook_scope( first_metadata.get("codebook_scope") ) - metadata_native_local_shards = normalized_metadata_native_local_shards[ - 0 - ] + metadata_native_local_shards = normalized_metadata_native_local_shards[0] if gefen_global_step is None: gefen_global_step = metadata_step elif gefen_global_step != metadata_step: @@ -8530,8 +8970,7 @@ def _load_state_dict_impl(self, state_dict): and gefen_deterministic != metadata_deterministic ): raise ValueError( - "Gefen checkpoint top-level and parameter-group deterministic " - "policies disagree" + "Gefen checkpoint top-level and parameter-group deterministic policies disagree" ) if not has_top_level_codebook_scope: gefen_codebook_scope = metadata_codebook_scope @@ -8557,8 +8996,9 @@ def _load_state_dict_impl(self, state_dict): if gefen_deterministic is not None: if type(gefen_deterministic) is not bool: raise ValueError( - "Gefen checkpoint deterministic policy must be a bool, got " - "{!r}".format(gefen_deterministic) + "Gefen checkpoint deterministic policy must be a bool, got {!r}".format( + gefen_deterministic + ) ) if gefen_deterministic != self._deterministic: raise ValueError( @@ -8604,7 +9044,10 @@ def _load_state_dict_impl(self, state_dict): # and re-aliases lazily). self._capt_invalidate() - self._base_load_state_dict_without_hooks(state_dict) + if self.state_offload_active: + self._base_load_state_dict_to_offload_cpu(state_dict) + else: + self._base_load_state_dict_without_hooks(state_dict) self._gefen_global_step = gefen_global_step # Capturable SR seeds are optimizer-level scratch (a device mirror of # gefen_global_step): drop them so the first post-load SR kernel call @@ -8648,7 +9091,10 @@ def _load_state_dict_impl(self, state_dict): if not torch.is_tensor(saved_value): continue live_value = live_state.get(key) - if torch.is_tensor(live_value) and live_value.dtype == saved_value.dtype: + if ( + torch.is_tensor(live_value) + and live_value.dtype == saved_value.dtype + ): continue device = ( live_value.device if torch.is_tensor(live_value) else param.device @@ -8701,6 +9147,7 @@ def step(self, closure=None): closure feed the first step's codebook learning correctly. The returned loss is passed through. """ + self._assert_state_offload_step_ready() self._assert_finalized_binding_layout() self._assert_runtime_codebook_process_group() self._assert_capturable_if_capturing() @@ -8711,6 +9158,7 @@ def step(self, closure=None): with torch.enable_grad(): loss = closure() + self._assert_state_offload_step_ready() self._assert_finalized_binding_layout() self._assert_runtime_codebook_process_group() try: @@ -8778,6 +9226,7 @@ def step(self, closure=None): not self._use_fused_gefen_automatic_step() and not self._use_fused_automatic_vmean() and not self.capturable + and not self.state_offload_active ) # Collect parameters that share block geometry + hyperparameters so they @@ -8801,12 +9250,17 @@ def step(self, closure=None): # path) instead of any vmean path. Sharded (DTensor) params # fall through to the standard path: local-shard row/col # statistics would be wrong under sharding. - if ( - self._factored_v_2d - and p.ndim == 2 - and not hasattr(p, "placements") - ): - self._step_automatic_factored(group, name, p, grad) + if self._factored_v_2d and p.ndim == 2 and not hasattr(p, "placements"): + if self.state_offload_active: + self._step_with_offloaded_parameter_state( + self._step_automatic_factored, + group, + name, + p, + grad, + ) + else: + self._step_automatic_factored(group, name, p, grad) continue if ( batch_nonfused @@ -8819,7 +9273,16 @@ def step(self, closure=None): (group, name, p, grad) ) continue - self._step_automatic(group, name, p, grad) + if self.state_offload_active: + self._step_with_offloaded_parameter_state( + self._step_automatic, + group, + name, + p, + grad, + ) + else: + self._step_automatic(group, name, p, grad) for items in nonfused_groups.values(): if len(items) == 1: diff --git a/src/gefen/hybrid.py b/src/gefen/hybrid.py index 72e3b50..a2757e6 100644 --- a/src/gefen/hybrid.py +++ b/src/gefen/hybrid.py @@ -53,6 +53,7 @@ period-one off, stochastic rounding off), with ``ns_schedule="tuned3"`` and ``normuon=True`` as the hybrid-specific defaults. """ + import copy import logging from collections import OrderedDict @@ -60,7 +61,16 @@ import torch import torch.nn as nn -from gefen.contracts import OptimizerContract, _hybrid_contract +from gefen.codebook import CodebookProcessGroupBinding +from gefen.contracts import ( + LogicalSlice, + OptimizerContract, + ParameterIdentity, + ParameterLayout, + ShardIdentity, + ShardingManifest, + _hybrid_contract, +) from gefen.gefen import ( Gefen, _amp_native_scaling_required, @@ -74,6 +84,7 @@ split_params_for_muon, validate_split, ) +from gefen.rebinding import ParameterRebinding logger = logging.getLogger(__name__) @@ -337,9 +348,7 @@ def __init__( if backup_named_params is None: # Single-argument convenience form: the first arg is a model or a # named-param iterable to split internally. - muon_named_params, backup_named_params = self._auto_split( - muon_named_params, backup_substrings - ) + muon_named_params, backup_named_params = self._auto_split(muon_named_params, backup_substrings) elif backup_substrings is not None: raise TypeError( "backup_substrings is only valid with the single-argument " @@ -359,9 +368,7 @@ def __init__( # 1D-only model (all params with real storage) still constructs a # backup-only hybrid exactly as before. if not muon_named_params and backup_named_params: - zero_numel = [ - name for name, p in backup_named_params if p.numel() == 0 - ] + zero_numel = [name for name, p in backup_named_params if p.numel() == 0] # Only the ALL-placeholder set is the ZeRO-3 signature; a mixed # set (e.g. an intentionally empty slot among real 1-D weights) # still constructs a backup-only hybrid as before. @@ -383,11 +390,7 @@ def __init__( ) ) if backup_optimizer not in ("gefen", "adamw"): - raise ValueError( - "backup_optimizer must be 'gefen' or 'adamw' but is: {!r}".format( - backup_optimizer - ) - ) + raise ValueError("backup_optimizer must be 'gefen' or 'adamw' but is: {!r}".format(backup_optimizer)) self.backup_optimizer = backup_optimizer # Catch the silent footguns: a param routed to both halves (stepped # twice) or a duplicate name (codebook cache key collision). Completeness @@ -408,9 +411,7 @@ def __init__( "GefenMuonHybrid: adjust_lr_fn={!r} uses Muon-native LR scaling, so " "an AdamW-scale lr will mis-scale the 2D Muon matrices relative to " "the backup half. Use adjust_lr_fn='match_rms_adamw' (the default) " - "to share one AdamW-scale lr, or set muon_lr explicitly.".format( - adjust_lr_fn - ), + "to share one AdamW-scale lr, or set muon_lr explicitly.".format(adjust_lr_fn), stacklevel=2, ) @@ -430,9 +431,7 @@ def __init__( # matrices while leaving the backup (norms/biases/embeddings/head) on a # different schedule. Both default to the shared weight_decay. muon_weight_decay = weight_decay if muon_weight_decay is None else muon_weight_decay - backup_weight_decay = ( - weight_decay if backup_weight_decay is None else backup_weight_decay - ) + backup_weight_decay = weight_decay if backup_weight_decay is None else backup_weight_decay self.muon = ( GefenMuon( @@ -566,6 +565,17 @@ def _is_no_decay(name): for p in group["params"]: self._state_param_owner[id(p)] = (p, o) + # Composite stable identity is installed only after every present Gefen + # child has staged the same complete post-sharding transaction. These + # fields remain separate from either child's rank-local metadata so a + # failed second-child preparation cannot leave the Hybrid half-bound. + self._hybrid_post_sharding_finalized = False + self._hybrid_sharding_manifest = None + self._hybrid_local_shard_bindings = () + self._hybrid_fqn_roles = () + self._hybrid_codebook_process_group = None + self._hybrid_finalized_slots = () + # Deliberately do NOT call super().__init__(): we expose each # sub-optimizer's real param_groups/state via properties (shared dict # refs), so the LR scheduler's in-place ``group["lr"] = ...`` updates @@ -624,16 +634,470 @@ def state(self): # unknown key (DeepSpeed ZeRO's flattened 1-D fp32 partitions) fails # fast with a self-explanatory KeyError instead of silently landing in # a rebuilt-per-access throwaway dict. + self._assert_finalized_binding_layout() return _HybridMergedState(self._subopts, self._state_param_owner) + def _gefen_rebinding_children(self): + children = [] + if self.muon is not None: + if type(self.muon) is not GefenMuon: + raise TypeError("GefenMuonHybrid post_sharding requires an exact GefenMuon child") + children.append(("muon", self.muon)) + if self.backup is not None: + if self.backup_optimizer != "gefen" or type(self.backup) is not Gefen: + raise NotImplementedError( + "GefenMuonHybrid post_sharding does not yet support an AdamW " + "backup; use backup_optimizer='gefen' until AdamW has stable " + "rebinding identity and atomic staged state I/O" + ) + children.append(("backup", self.backup)) + if not children: + raise RuntimeError("GefenMuonHybrid has no child optimizer to rebind") + return tuple(children) + + @staticmethod + def _same_local_binding(left, right) -> bool: + return left[0] is right[0] and left[1] == right[1] + + @staticmethod + def _reject_rebinding_method_shadows(value) -> None: + if type(value.__dict__) is not dict: + raise TypeError("GefenMuonHybrid rebinding requires exact attribute dictionaries") + for name in value.__dict__: + descriptor = None + for owner in type(value).__mro__: + if name in owner.__dict__: + descriptor = owner.__dict__[name] + break + if isinstance(descriptor, (staticmethod, classmethod)): + descriptor = descriptor.__func__ + if callable(descriptor): + raise TypeError("GefenMuonHybrid rebinding rejects instance-level method shadows") + + def _hybrid_identity_metadata_empty(self) -> bool: + return ( + self._hybrid_sharding_manifest is None + and self._hybrid_local_shard_bindings == () + and self._hybrid_fqn_roles == () + and self._hybrid_codebook_process_group is None + and self._hybrid_finalized_slots == () + ) + + def _assert_composite_rebinding_pristine(self, children) -> None: + if self._hybrid_post_sharding_finalized: + raise RuntimeError("GefenMuonHybrid post-sharding identity is already finalized") + if not self._hybrid_identity_metadata_empty(): + raise RuntimeError("GefenMuonHybrid parameter rebinding found an incomplete prior identity plan") + if type(self._subopts) is not list or len(self._subopts) != len(children): + raise RuntimeError("GefenMuonHybrid child optimizer order changed before post_sharding") + if any(live is not expected for live, (_role, expected) in zip(self._subopts, children)): + raise RuntimeError("GefenMuonHybrid child optimizer order changed before post_sharding") + if type(self._state_param_owner) is not dict: + raise TypeError("GefenMuonHybrid parameter ownership must use an exact dictionary") + child_set = {child for _role, child in children} + owner_counts = {child: 0 for child in child_set} + for parameter_id, live in self._state_param_owner.items(): + if ( + type(live) is not tuple + or len(live) != 2 + or not isinstance(live[0], torch.Tensor) + or parameter_id != id(live[0]) + or live[1] not in child_set + ): + raise RuntimeError("GefenMuonHybrid parameter ownership changed before post_sharding") + owner_counts[live[1]] += 1 + for _role, child in children: + live_slot_count = sum(len(group["params"]) for group in child.param_groups) + if owner_counts[child] != live_slot_count: + raise RuntimeError("GefenMuonHybrid child slot counts changed before post_sharding") + + def _stage_post_sharding( + self, + rebindings, + manifest: ShardingManifest, + codebook_process_group=None, + ): + GefenMuonHybrid._reject_rebinding_method_shadows(self) + children = GefenMuonHybrid._gefen_rebinding_children(self) + for _role, child in children: + GefenMuonHybrid._reject_rebinding_method_shadows(child) + GefenMuonHybrid._assert_composite_rebinding_pristine(self, children) + if type(manifest) is not ShardingManifest: + raise TypeError("manifest must be an exact ShardingManifest") + validated_manifest = ShardingManifest( + manifest.shards, + schema_version=manifest.schema_version, + ) + if validated_manifest != manifest: + raise ValueError("GefenMuonHybrid requires a canonical manifest") + if type(rebindings) is not tuple or not rebindings: + raise TypeError("rebindings must be a non-empty tuple of ParameterRebinding values") + if any(type(item) is not ParameterRebinding for item in rebindings): + raise TypeError("rebindings must contain exact ParameterRebinding values") + if codebook_process_group is not None and type(codebook_process_group) is not CodebookProcessGroupBinding: + raise TypeError("codebook_process_group must be an exact CodebookProcessGroupBinding") + + source_entries = {} + source_roles = {} + child_by_role = dict(children) + for parameter_id, (parameter, child) in self._state_param_owner.items(): + if parameter_id != id(parameter): + raise RuntimeError("GefenMuonHybrid parameter ownership contains a stale key") + role = "muon" if child is self.muon else "backup" + if role not in child_by_role or child_by_role[role] is not child: + raise RuntimeError("GefenMuonHybrid parameter ownership names a foreign child") + source_entries[parameter_id] = parameter + source_roles[parameter_id] = role + if len(rebindings) != len(source_entries): + raise ValueError( + "GefenMuonHybrid post_sharding requires exactly one rebinding for every original child slot" + ) + + seen_sources = set() + seen_targets = set() + seen_fqns = set() + targets = [] + by_role = {role: [] for role, _child in children} + fqn_roles = {} + for rebinding in rebindings: + source = rebinding.old_parameter + source_id = id(source) + if source_id not in source_entries or source_entries[source_id] is not source: + raise ValueError("GefenMuonHybrid rebinding source is not an original child slot") + if source_id in seen_sources: + raise ValueError("GefenMuonHybrid rebinding source tensors must be unique") + seen_sources.add(source_id) + target = rebinding.new_parameter + if target is not None: + if not isinstance(target, torch.Tensor): + raise TypeError("GefenMuonHybrid rebinding target must be a Tensor or None") + target_id = id(target) + if target_id in seen_targets: + raise ValueError("GefenMuonHybrid rebound target tensors must be unique") + seen_targets.add(target_id) + original_target = source_entries.get(target_id) + if original_target is not None and original_target is not source: + raise ValueError("GefenMuonHybrid rebound targets cannot steal another original child slot") + targets.append(target) + fqn = rebinding.shard.parameter.fqn + if fqn in seen_fqns: + raise ValueError("GefenMuonHybrid local canonical parameter FQNs must be unique") + seen_fqns.add(fqn) + role = source_roles[source_id] + fqn_roles[fqn] = role + by_role[role].append(rebinding) + + if seen_sources != set(source_entries): + raise ValueError("GefenMuonHybrid post_sharding did not bind every original child slot") + Gefen._assert_rebound_storage_disjoint(targets) + + manifest_fqns = {shard.parameter.fqn for shard in manifest.shards} + if seen_fqns != manifest_fqns: + raise ValueError("GefenMuonHybrid manifest FQNs must exactly match all child slots") + if codebook_process_group is not None: + if any(shard.process_group != codebook_process_group.identity for shard in manifest.shards): + raise ValueError("every Hybrid manifest shard must use the shared codebook process-group identity") + if any(item.shard.local_member != codebook_process_group.local_member for item in rebindings): + raise ValueError("every Hybrid local shard must match the shared codebook member") + + staged_children = [] + for role, child in children: + child_rebindings = tuple(by_role[role]) + if not child_rebindings: + raise ValueError("GefenMuonHybrid child {!r} has no routed rebinding".format(role)) + child_fqns = {item.shard.parameter.fqn for item in child_rebindings} + child_manifest = ShardingManifest( + tuple(shard for shard in manifest.shards if shard.parameter.fqn in child_fqns), + schema_version=manifest.schema_version, + ) + staged = child._stage_post_sharding( + child_rebindings, + child_manifest, + codebook_process_group, + ) + if ( + type(child.__dict__) is not dict + or type(staged.__dict__) is not dict + or staged._gefen_codebook_process_group is not codebook_process_group + or not staged._finalized_binding_layout_matches() + ): + raise TypeError("GefenMuonHybrid child staging produced an unsafe finalized optimizer") + staged_children.append((role, child, staged, child_manifest)) + + local_bindings = [] + local_fqns = set() + new_state_param_owner = {} + finalized_slots = [] + for role, child, staged, _child_manifest in staged_children: + finalized_slots.append( + ( + role, + tuple(tuple(group["params"]) for group in staged.param_groups), + ) + ) + for parameter, shard in staged._gefen_local_shard_bindings: + fqn = shard.parameter.fqn + if fqn in local_fqns or fqn_roles.get(fqn) != role: + raise ValueError("GefenMuonHybrid staged child identities overlap or changed routing") + local_fqns.add(fqn) + local_bindings.append((parameter, shard)) + for group in staged.param_groups: + for parameter in group["params"]: + parameter_id = id(parameter) + if parameter_id in new_state_param_owner: + raise ValueError("GefenMuonHybrid staged children share a live parameter") + new_state_param_owner[parameter_id] = (parameter, child) + if local_fqns != seen_fqns: + raise ValueError("GefenMuonHybrid staged children do not cover the full manifest") + local_bindings.sort(key=lambda item: item[1].sort_key) + if type(self.__dict__) is not dict: + raise TypeError("GefenMuonHybrid publication requires an exact attribute dictionary") + return { + "children": tuple(staged_children), + "state_param_owner": new_state_param_owner, + "manifest": manifest, + "local_bindings": tuple(local_bindings), + "fqn_roles": tuple(sorted(fqn_roles.items())), + "codebook_process_group": codebook_process_group, + "finalized_slots": tuple(finalized_slots), + } + + def post_sharding( + self, + rebindings, + *, + manifest: ShardingManifest, + codebook_process_group=None, + ) -> None: + """Atomically finalize every present Gefen child after sharding.""" + + if isinstance(rebindings, (str, bytes)): + raise TypeError("rebindings must be a sequence") + try: + rebindings = tuple(rebindings) + except TypeError as exc: + raise TypeError("rebindings must be a sequence") from exc + staged = GefenMuonHybrid._stage_post_sharding( + self, + rebindings, + manifest, + codebook_process_group, + ) + for _role, child, staged_child, _child_manifest in staged["children"]: + dict.update(child.__dict__, staged_child.__dict__) + dict.update( + self.__dict__, + { + "_state_param_owner": staged["state_param_owner"], + "_hybrid_post_sharding_finalized": True, + "_hybrid_sharding_manifest": staged["manifest"], + "_hybrid_local_shard_bindings": staged["local_bindings"], + "_hybrid_fqn_roles": staged["fqn_roles"], + "_hybrid_codebook_process_group": staged["codebook_process_group"], + "_hybrid_finalized_slots": staged["finalized_slots"], + }, + ) + + def rebind_shard( + self, + old_parameter, + new_parameter, + *, + shard: ShardIdentity, + manifest: ShardingManifest, + ) -> None: + """Apply a complete one-slot Hybrid shard plan.""" + + self.post_sharding( + (ParameterRebinding(old_parameter, new_parameter, shard),), + manifest=manifest, + ) + + def rebind_parameter( + self, + old_parameter, + new_parameter, + *, + identity: ParameterIdentity, + ) -> None: + """Bind one complete replicated parameter on a one-slot Hybrid.""" + + if type(identity) is not ParameterIdentity: + raise TypeError("identity must be an exact ParameterIdentity") + shard = ShardIdentity( + identity, + ParameterLayout.REPLICATED, + LogicalSlice.full(identity), + ) + self.rebind_shard( + old_parameter, + new_parameter, + shard=shard, + manifest=ShardingManifest((shard,)), + ) + + def _finalized_binding_layout_matches(self) -> bool: + try: + if ( + not self._hybrid_post_sharding_finalized + or type(self._hybrid_sharding_manifest) is not ShardingManifest + or type(self._hybrid_local_shard_bindings) is not tuple + or type(self._hybrid_fqn_roles) is not tuple + or type(self._hybrid_finalized_slots) is not tuple + or type(self._state_param_owner) is not dict + ): + return False + children = self._gefen_rebinding_children() + if ( + type(self._subopts) is not list + or len(self._subopts) != len(children) + or any(live is not child for live, (_role, child) in zip(self._subopts, children)) + or self.defaults is not children[0][1].defaults + ): + return False + role_by_fqn = dict(self._hybrid_fqn_roles) + if ( + len(role_by_fqn) != len(self._hybrid_fqn_roles) + or tuple(sorted(role_by_fqn.items())) != self._hybrid_fqn_roles + or any(role not in {"muon", "backup"} for role in role_by_fqn.values()) + ): + return False + manifest_fqns = {shard.parameter.fqn for shard in self._hybrid_sharding_manifest.shards} + if manifest_fqns != set(role_by_fqn): + return False + binding = self._hybrid_codebook_process_group + if binding is not None: + if type(binding) is not CodebookProcessGroupBinding: + return False + if any(shard.process_group != binding.identity for shard in self._hybrid_sharding_manifest.shards): + return False + + expected_local = [] + expected_owner = {} + expected_slots = [] + for role, child in children: + if ( + not child._finalized_binding_layout_matches() + or child._gefen_codebook_process_group is not self._hybrid_codebook_process_group + ): + return False + expected_child_shards = tuple( + shard + for shard in self._hybrid_sharding_manifest.shards + if role_by_fqn.get(shard.parameter.fqn) == role + ) + if child._gefen_sharding_manifest.shards != expected_child_shards: + return False + expected_slots.append( + ( + role, + tuple(tuple(group["params"]) for group in child.param_groups), + ) + ) + for parameter, shard in child._gefen_local_shard_bindings: + if role_by_fqn.get(shard.parameter.fqn) != role: + return False + if binding is not None and ( + shard.process_group != binding.identity or shard.local_member != binding.local_member + ): + return False + expected_local.append((parameter, shard)) + for group in child.param_groups: + for parameter in group["params"]: + parameter_id = id(parameter) + if parameter_id in expected_owner: + return False + expected_owner[parameter_id] = (parameter, child) + expected_local.sort(key=lambda item: item[1].sort_key) + if len(expected_local) != len(self._hybrid_local_shard_bindings): + return False + if any( + not self._same_local_binding(live, expected) + for live, expected in zip( + self._hybrid_local_shard_bindings, + expected_local, + ) + ): + return False + if tuple(expected_slots) != self._hybrid_finalized_slots: + return False + if set(expected_owner) != set(self._state_param_owner): + return False + for parameter_id, (parameter, child) in expected_owner.items(): + live = self._state_param_owner[parameter_id] + if type(live) is not tuple or len(live) != 2 or live[0] is not parameter or live[1] is not child: + return False + return True + except ( + AttributeError, + KeyError, + NotImplementedError, + RuntimeError, + TypeError, + ValueError, + ): + return False + + def _canonical_identity_ready(self) -> bool: + return self._hybrid_post_sharding_finalized and self._finalized_binding_layout_matches() + + def _codebook_scope_ready(self) -> bool: + return self._hybrid_codebook_process_group is not None and self._canonical_identity_ready() + + def _assert_finalized_binding_layout(self) -> None: + if self._hybrid_post_sharding_finalized: + if not self._finalized_binding_layout_matches(): + raise RuntimeError("GefenMuonHybrid finalized parameter layout changed outside post_sharding") + elif not self._hybrid_identity_metadata_empty(): + raise RuntimeError("GefenMuonHybrid found an incomplete post_sharding identity plan") + + def parameter_identity(self, parameter) -> ParameterIdentity: + """Return the canonical identity bound to one live Hybrid parameter.""" + + return self.shard_identity(parameter).parameter + + def shard_identity(self, parameter) -> ShardIdentity: + """Return the stable shard identity bound to one live parameter.""" + + self._assert_finalized_binding_layout() + entry = self._state_param_owner.get(id(parameter)) + if entry is None or entry[0] is not parameter: + raise KeyError("parameter has no finalized GefenMuonHybrid shard identity") + return entry[1].shard_identity(parameter) + + def shard_bindings(self): + """Return all local tensor/identity pairs in canonical order.""" + + self._assert_finalized_binding_layout() + return self._hybrid_local_shard_bindings + + def parameter_routing(self): + """Return the immutable canonical FQN-to-child-role routing.""" + + self._assert_finalized_binding_layout() + if not self._canonical_identity_ready(): + raise RuntimeError("GefenMuonHybrid parameter routing is not finalized") + return self._hybrid_fqn_roles + + def sharding_manifest(self): + """Return the complete composite manifest, or ``None``.""" + + self._assert_finalized_binding_layout() + return self._hybrid_sharding_manifest + + def codebook_process_group_binding(self): + """Return the one binding shared by every present Gefen child.""" + + self._assert_finalized_binding_layout() + return self._hybrid_codebook_process_group + def zero_grad(self, set_to_none: bool = True): + self._assert_finalized_binding_layout() for o in self._subopts: o.zero_grad(set_to_none=set_to_none) def _assert_capturable_devices_if_capturing(self) -> None: - capturing = ( - torch.cuda.is_available() and torch.cuda.is_current_stream_capturing() - ) + capturing = torch.cuda.is_available() and torch.cuda.is_current_stream_capturing() if not capturing: return if not self.capturable: @@ -643,21 +1107,17 @@ def _assert_capturable_devices_if_capturing(self) -> None: "to make step() graph-safe." ) devices = { - param.device - for optimizer in self._subopts - for group in optimizer.param_groups - for param in group["params"] + param.device for optimizer in self._subopts for group in optimizer.param_groups for param in group["params"] } capture_device = torch.device("cuda", torch.cuda.current_device()) if devices != {capture_device}: raise RuntimeError( "CUDA graph capture requires every GefenMuonHybrid parameter on " - "the current capture device {}; found {}".format( - capture_device, sorted(map(str, devices)) - ) + "the current capture device {}; found {}".format(capture_device, sorted(map(str, devices))) ) def step(self, closure=None): + self._assert_finalized_binding_layout() self._assert_capturable_devices_if_capturing() # Dispatch the INSTANCE step hooks around the composite step, mirroring # torch.optim.Optimizer.profile_hook_step exactly: hooks receive @@ -687,17 +1147,14 @@ def step(self, closure=None): if closure is not None: with torch.enable_grad(): loss = closure() + self._assert_finalized_binding_layout() for child in self._subopts: - _assert_optimizer_gradients_structurally_valid( - child, require_2d_params=child is self.muon - ) + _assert_optimizer_gradients_structurally_valid(child, require_2d_params=child is self.muon) # A non-finite gradient in either half skips BOTH children before their # codebooks, states, counters, or parameters can move. Explicit # scaler.unscale_(hybrid) is detected by grad_scale=None and is not # repeated; automatic unscale covers every child parameter exactly once. - if ( - hasattr(self, "found_inf") or hasattr(self, "grad_scale") - ) and not _amp_prepare_optimizer_step(self): + if (hasattr(self, "found_inf") or hasattr(self, "grad_scale")) and not _amp_prepare_optimizer_step(self): for post_hook in self._optimizer_step_post_hooks.values(): post_hook(self, args, kwargs) return loss @@ -710,11 +1167,13 @@ def step(self, closure=None): return loss def state_dict(self): + self._assert_finalized_binding_layout() # Instance state-dict pre/post hooks, mirroring Optimizer.state_dict: # pre-hooks take (optimizer) and return nothing; a post-hook may return # a replacement state_dict. for pre_hook in self._optimizer_state_dict_pre_hooks.values(): pre_hook(self) + self._assert_finalized_binding_layout() state_dict = { "muon": self.muon.state_dict() if self.muon is not None else None, "backup": self.backup.state_dict() if self.backup is not None else None, @@ -729,9 +1188,26 @@ def state_dict(self): def optimizer_contract(self) -> OptimizerContract: """Return the composite contract without flattening either child schema.""" - muon_contract = ( - self.muon.optimizer_contract() if self.muon is not None else None - ) + try: + GefenMuonHybrid._reject_rebinding_method_shadows(self) + rebinding_ready = bool(GefenMuonHybrid._gefen_rebinding_children(self)) + identity_ready = GefenMuonHybrid._canonical_identity_ready(self) + except Exception: + rebinding_ready = False + identity_ready = False + try: + from gefen.portable_hybrid import _hybrid_portable_contract_support + + ( + canonical_global_same_topology, + canonical_global_topology_changing, + canonical_global_topology_change_kinds, + ) = _hybrid_portable_contract_support(self) + except Exception: + canonical_global_same_topology = frozenset() + canonical_global_topology_changing = frozenset() + canonical_global_topology_change_kinds = frozenset() + muon_contract = self.muon.optimizer_contract() if self.muon is not None else None backup_contract = ( self.backup.optimizer_contract() if self.backup is not None and hasattr(self.backup, "optimizer_contract") @@ -748,9 +1224,56 @@ def optimizer_contract(self) -> OptimizerContract: muon=muon_contract, backup=backup_contract, backup_implementation=backup_implementation, + canonical_parameter_fqns=identity_ready, + stable_shard_identity=identity_ready, + explicit_process_group_codebook_scope=rebinding_ready, + shard_rebinding=rebinding_ready, + post_sharding=rebinding_ready, + canonical_global_same_topology=canonical_global_same_topology, + canonical_global_topology_changing=canonical_global_topology_changing, + canonical_global_topology_change_kinds=canonical_global_topology_change_kinds, + ) + + def export_portable_state( + self, + *, + checkpoint_process_group, + transaction_id, + limits, + ): + """Collectively export one complete Gefen-backed composite document.""" + + from gefen.portable_hybrid import _export_hybrid_portable_state + + return _export_hybrid_portable_state( + self, + checkpoint_process_group=checkpoint_process_group, + transaction_id=transaction_id, + limits=limits, + ) + + def import_portable_state( + self, + state, + *, + checkpoint_process_group, + transaction_id, + limits, + ) -> None: + """Collectively stage and atomically publish composite portable state.""" + + from gefen.portable_hybrid import _import_hybrid_portable_state + + _import_hybrid_portable_state( + self, + state, + checkpoint_process_group=checkpoint_process_group, + transaction_id=transaction_id, + limits=limits, ) def load_state_dict(self, state_dict): + self._assert_finalized_binding_layout() # Instance load pre-hooks first (a pre-hook may return a replacement # dict -- e.g. one that converts a foreign schema), mirroring # Optimizer.load_state_dict's shallow copy + hook pass. @@ -759,6 +1282,7 @@ def load_state_dict(self, state_dict): hook_result = pre_hook(self, state_dict) if hook_result is not None: state_dict = hook_result + self._assert_finalized_binding_layout() # Schema guard: this used to silently skip loading whenever the keys # were absent, so resuming from a standard {"state", "param_groups"} @@ -799,9 +1323,7 @@ def load_state_dict(self, state_dict): # optimizer loader otherwise accepts the other backend's group/state # layout and fails only on a later step (or, worse, misinterprets state). if self.backup is not None: - checkpoint_backup_optimizer = state_dict.get( - "backup_optimizer", "gefen" - ) + checkpoint_backup_optimizer = state_dict.get("backup_optimizer", "gefen") if checkpoint_backup_optimizer != self.backup_optimizer: raise ValueError( "GefenMuonHybrid.load_state_dict: checkpoint backup_optimizer " @@ -841,15 +1363,9 @@ def _auto_split(params_or_model, backup_substrings): T5 ``shared``, would slip to Muon, so a model is preferred). Bare tensors raise, since they carry neither names nor module type to route on. """ - subs = ( - DEFAULT_BACKUP_SUBSTRINGS - if backup_substrings is None - else tuple(backup_substrings) - ) + subs = DEFAULT_BACKUP_SUBSTRINGS if backup_substrings is None else tuple(backup_substrings) if isinstance(params_or_model, nn.Module): - muon_named, backup_named = split_params_for_muon( - params_or_model, backup_substrings=subs - ) + muon_named, backup_named = split_params_for_muon(params_or_model, backup_substrings=subs) validate_split(muon_named, backup_named, model=params_or_model) return muon_named, backup_named items = list(params_or_model) @@ -865,9 +1381,7 @@ def _auto_split(params_or_model, backup_substrings): "an iterable of (name, param) pairs (e.g. model.named_parameters()); " "got a bare {}. Bare tensors can't be routed (embeddings/heads " "would go to Muon) -- pass the model or use " - "GefenMuonHybrid.from_model(model, ...).".format( - type(item).__name__ - ) + "GefenMuonHybrid.from_model(model, ...).".format(type(item).__name__) ) muon_named, backup_named = [], [] for name, param in items: @@ -899,22 +1413,9 @@ def from_model(cls, model, *, backup_substrings=None, **kwargs): return cls(muon_named, backup_named, **kwargs) def add_param_group(self, param_group): - raise NotImplementedError( - "GefenMuonHybrid splits params at construction; add_param_group is unsupported" - ) + raise NotImplementedError("GefenMuonHybrid splits params at construction; add_param_group is unsupported") def __repr__(self): - nm = ( - sum(len(group["params"]) for group in self.muon.param_groups) - if self.muon is not None - else 0 - ) - nb = ( - sum(len(group["params"]) for group in self.backup.param_groups) - if self.backup is not None - else 0 - ) - return ( - f"GefenMuonHybrid(muon_params={nm}, backup_params={nb}, " - f"backup_optimizer={self.backup_optimizer!r})" - ) + nm = sum(len(group["params"]) for group in self.muon.param_groups) if self.muon is not None else 0 + nb = sum(len(group["params"]) for group in self.backup.param_groups) if self.backup is not None else 0 + return f"GefenMuonHybrid(muon_params={nm}, backup_params={nb}, backup_optimizer={self.backup_optimizer!r})" diff --git a/src/gefen/portable_dcp.py b/src/gefen/portable_dcp.py index 9b8dd1d..c698bcb 100644 --- a/src/gefen/portable_dcp.py +++ b/src/gefen/portable_dcp.py @@ -133,11 +133,20 @@ def _preflight_dcp_operation( ): from gefen import portable_runtime as runtime from gefen.portable_collective import _collective_unanimous_status + from gefen.hybrid import GefenMuonHybrid - transport = runtime._preflight_transport_binding( - optimizer, - checkpoint_process_group, - ) + if type(optimizer) is GefenMuonHybrid: + from gefen.portable_hybrid import _hybrid_transport_binding + + transport = _hybrid_transport_binding( + optimizer, + checkpoint_process_group, + ) + else: + transport = runtime._preflight_transport_binding( + optimizer, + checkpoint_process_group, + ) binding = None normalized_limits = None normalized_transaction = None @@ -147,32 +156,56 @@ def _preflight_dcp_operation( wire_limits = runtime._STATUS_FALLBACK_LIMITS error = None try: - runtime._validate_supplied_binding(checkpoint_process_group, transport) - binding = transport - normalized_limits = runtime._require_limits(limits) + if type(optimizer) is GefenMuonHybrid: + from gefen.portable_hybrid import ( + HYBRID_PORTABLE_STATE_IMPLEMENTATION, + _preflight_hybrid_portable_local, + ) + + ( + binding, + normalized_transaction, + normalized_limits, + _children, + _routing, + base_context, + _base_digest, + _live_token, + ) = _preflight_hybrid_portable_local( + optimizer, + checkpoint_process_group=checkpoint_process_group, + transaction_id=transaction_id, + limits=limits, + ) + implementation = HYBRID_PORTABLE_STATE_IMPLEMENTATION + else: + runtime._validate_supplied_binding(checkpoint_process_group, transport) + binding = transport + normalized_limits = runtime._require_limits(limits) + normalized_transaction = runtime._require_transaction_id(transaction_id) + implementation = runtime._optimizer_implementation(optimizer) + runtime._validate_context_identity_bounds(binding, normalized_limits) + prepared = runtime._prepare_local_structure( + optimizer, + implementation, + binding, + normalized_limits, + include_payload=False, + ) + runtime._validate_prepared_local_state( + optimizer, + implementation, + prepared, + ) + base_context = runtime._base_context(binding, implementation) wire_limits = normalized_limits._wire_limits() - normalized_transaction = runtime._require_transaction_id(transaction_id) normalized_namespace = _require_namespace(namespace, normalized_limits) if not isinstance(storage, storage_type): raise TypeError("storage must be a {}".format(storage_type.__name__)) storage_identity = _storage_identity(storage, normalized_limits) - implementation = runtime._optimizer_implementation(optimizer) - runtime._validate_context_identity_bounds(binding, normalized_limits) _validate_dcp_runtime(binding) - prepared = runtime._prepare_local_structure( - optimizer, - implementation, - binding, - normalized_limits, - include_payload=False, - ) - runtime._validate_prepared_local_state( - optimizer, - implementation, - prepared, - ) context = { - **runtime._base_context(binding, implementation), + **base_context, "dcp_envelope_version": 1, "dcp_namespace": normalized_namespace, "dcp_storage": storage_identity, @@ -416,7 +449,7 @@ def save_portable_dcp( limits, namespace="optimizer", ): - """Collectively save one tensor-only portable v3 document through DCP.""" + """Collectively save one tensor-only portable optimizer document through DCP.""" import torch.distributed.checkpoint as dcp from torch.distributed.checkpoint.storage import StorageWriter @@ -474,7 +507,7 @@ def load_portable_dcp( limits, namespace="optimizer", ) -> None: - """Collectively load, verify, and atomically import portable v3 from DCP.""" + """Collectively load, verify, and atomically import portable optimizer state.""" import torch.distributed.checkpoint as dcp from torch.distributed.checkpoint.storage import StorageReader diff --git a/src/gefen/portable_hybrid.py b/src/gefen/portable_hybrid.py new file mode 100644 index 0000000..3cb0ffc --- /dev/null +++ b/src/gefen/portable_hybrid.py @@ -0,0 +1,672 @@ +"""Strict composite envelope for Gefen-backed Hybrid portable state.""" + +import hashlib +import hmac + +from gefen.portable_schema import ( + PORTABLE_STATE_DIGEST_ALGORITHM, + normalize_portable_state_document, + portable_state_digest, +) + + +HYBRID_PORTABLE_STATE_FORMAT = "gefen.portable_composite_state" +HYBRID_PORTABLE_STATE_FORMAT_VERSION = 1 +HYBRID_PORTABLE_STATE_COVERAGE = "global_logical_composite_optimizer" +HYBRID_PORTABLE_STATE_IMPLEMENTATION = "gefen.GefenMuonHybrid" + +_HYBRID_PORTABLE_ROLES = ("muon", "backup") +_HYBRID_PORTABLE_CHILD_IMPLEMENTATIONS = { + "muon": "gefen.GefenMuon", + "backup": "gefen.Gefen", +} +_HYBRID_PORTABLE_TOP_LEVEL_KEYS = frozenset( + { + "format", + "format_version", + "coverage", + "implementation", + "backup_optimizer", + "routing", + "children", + "completion", + } +) +_HYBRID_PORTABLE_COMPLETION_KEYS = frozenset({"status", "digest_algorithm", "digest"}) +_HYBRID_EXPORT_PREFLIGHT_TRANSACTION = "gefen-hybrid-portable-export-preflight-v1" +_HYBRID_IMPORT_PREFLIGHT_TRANSACTION = "gefen-hybrid-portable-import-preflight-v1" + + +def _hybrid_portable_payload(document): + return {key: document[key] for key in sorted(_HYBRID_PORTABLE_TOP_LEVEL_KEYS - {"completion"})} + + +def _normalize_routing(routing): + if type(routing) is not dict: + raise ValueError("portable Hybrid routing must be an FQN dictionary") + if any(type(fqn) is not str for fqn in routing): + raise ValueError("portable Hybrid routing keys must be strings") + normalized = {} + for fqn in sorted(routing): + role = routing[fqn] + if not fqn or fqn != fqn.strip(): + raise ValueError("portable Hybrid routing keys must be non-empty trimmed FQNs") + if type(role) is not str or role not in _HYBRID_PORTABLE_ROLES: + raise ValueError("portable Hybrid routing values must name exact child roles") + normalized[fqn] = role + return normalized + + +def _normalize_children(children): + if type(children) is not dict or set(children) != set(_HYBRID_PORTABLE_ROLES): + raise ValueError("portable Hybrid children must contain exact muon and backup roles") + normalized = {} + for role in _HYBRID_PORTABLE_ROLES: + child = children[role] + normalized[role] = ( + None + if child is None + else normalize_portable_state_document( + child, + expected_implementation=_HYBRID_PORTABLE_CHILD_IMPLEMENTATIONS[role], + ) + ) + if all(child is None for child in normalized.values()): + raise ValueError("portable Hybrid state must contain at least one child") + return normalized + + +def _validate_child_routing(children, routing) -> None: + expected_routing = {} + common_step = None + deterministic = None + for role in _HYBRID_PORTABLE_ROLES: + child = children[role] + if child is None: + continue + for fqn in child["parameters"]: + if fqn in expected_routing: + raise ValueError("portable Hybrid child parameter FQNs must be disjoint") + expected_routing[fqn] = role + child_common = child["common"] + if type(child_common) is not dict: + raise ValueError("portable Hybrid child common state must be a dictionary") + child_step = child_common.get("gefen_global_step") + child_deterministic = child_common.get("gefen_deterministic") + if type(child_step) is not int or child_step < 0: + raise ValueError("portable Hybrid children require nonnegative exact global steps") + if type(child_deterministic) is not bool: + raise ValueError("portable Hybrid children require exact deterministic policies") + if common_step is None: + common_step = child_step + deterministic = child_deterministic + elif child_step != common_step: + raise ValueError("portable Hybrid child global steps must agree") + elif child_deterministic is not deterministic: + raise ValueError("portable Hybrid child deterministic policies must agree") + if routing != {fqn: expected_routing[fqn] for fqn in sorted(expected_routing)}: + raise ValueError("portable Hybrid routing does not exactly match child parameters") + + +def _normalize_hybrid_portable_payload(state): + if state["format"] != HYBRID_PORTABLE_STATE_FORMAT: + raise ValueError("unsupported portable Hybrid state format") + if type(state["format_version"]) is not int or state["format_version"] != HYBRID_PORTABLE_STATE_FORMAT_VERSION: + raise ValueError("unsupported portable Hybrid state format_version: {}".format(state["format_version"])) + if state["coverage"] != HYBRID_PORTABLE_STATE_COVERAGE: + raise ValueError("unsupported portable Hybrid state coverage") + if state["implementation"] != HYBRID_PORTABLE_STATE_IMPLEMENTATION: + raise ValueError("portable Hybrid state implementation does not match the target") + if state["backup_optimizer"] != "gefen": + raise ValueError("portable Hybrid state requires a Gefen backup policy") + routing = _normalize_routing(state["routing"]) + children = _normalize_children(state["children"]) + _validate_child_routing(children, routing) + return { + "format": HYBRID_PORTABLE_STATE_FORMAT, + "format_version": HYBRID_PORTABLE_STATE_FORMAT_VERSION, + "coverage": HYBRID_PORTABLE_STATE_COVERAGE, + "implementation": HYBRID_PORTABLE_STATE_IMPLEMENTATION, + "backup_optimizer": "gefen", + "routing": routing, + "children": children, + } + + +def _validate_hybrid_portable_completion(state, payload) -> None: + completion = state["completion"] + if type(completion) is not dict or set(completion) != _HYBRID_PORTABLE_COMPLETION_KEYS: + raise ValueError("portable Hybrid completion marker has an invalid schema") + if completion["status"] != "complete": + raise ValueError("portable Hybrid state is not marked complete") + if completion["digest_algorithm"] != PORTABLE_STATE_DIGEST_ALGORITHM: + raise ValueError("unsupported portable Hybrid state digest algorithm") + digest = completion["digest"] + if ( + type(digest) is not str + or len(digest) != hashlib.sha256().digest_size * 2 + or any(character not in "0123456789abcdef" for character in digest) + ): + raise ValueError("portable Hybrid completion digest is invalid") + if not hmac.compare_digest(digest, portable_state_digest(payload)): + raise ValueError("portable Hybrid completion digest does not match its payload") + + +def build_hybrid_portable_state_document(*, backup_optimizer, routing, children): + """Build one complete composite wrapper around unchanged portable v3 children.""" + + state = { + "format": HYBRID_PORTABLE_STATE_FORMAT, + "format_version": HYBRID_PORTABLE_STATE_FORMAT_VERSION, + "coverage": HYBRID_PORTABLE_STATE_COVERAGE, + "implementation": HYBRID_PORTABLE_STATE_IMPLEMENTATION, + "backup_optimizer": backup_optimizer, + "routing": routing, + "children": children, + } + payload = _normalize_hybrid_portable_payload(state) + payload["completion"] = { + "status": "complete", + "digest_algorithm": PORTABLE_STATE_DIGEST_ALGORITHM, + "digest": portable_state_digest(payload), + } + return payload + + +def normalize_hybrid_portable_state_document(state): + """Clone and strictly validate one complete portable Hybrid wrapper.""" + + if type(state) is not dict or set(state) != _HYBRID_PORTABLE_TOP_LEVEL_KEYS: + raise ValueError("portable Hybrid state has an invalid top-level schema") + payload = _normalize_hybrid_portable_payload(state) + _validate_hybrid_portable_completion(state, payload) + payload["completion"] = dict(state["completion"]) + return payload + + +def _validate_hybrid_portable_limits(state, limits) -> None: + from gefen.portable_wire import _prepare_canonical_wire_value + + _prepare_canonical_wire_value(state, limits._wire_limits(collective=True)) + + +def _hybrid_children(optimizer): + from gefen.gefen import Gefen + from gefen.gefen_muon import GefenMuon + from gefen.hybrid import GefenMuonHybrid + + if type(optimizer) is not GefenMuonHybrid: + raise TypeError("portable Hybrid state requires an exact GefenMuonHybrid") + if optimizer.backup_optimizer != "gefen": + raise NotImplementedError("portable Hybrid state does not support an AdamW backup") + children = [] + if optimizer.muon is not None: + if type(optimizer.muon) is not GefenMuon: + raise TypeError("portable Hybrid state requires an exact GefenMuon child") + children.append(("muon", optimizer.muon)) + if optimizer.backup is not None: + if type(optimizer.backup) is not Gefen: + raise TypeError("portable Hybrid state requires an exact Gefen backup child") + children.append(("backup", optimizer.backup)) + if not children: + raise RuntimeError("portable Hybrid state requires at least one child") + return tuple(children) + + +def _hybrid_transport_binding(optimizer, supplied): + from gefen import portable_runtime as runtime + + # The caller-provided exact binding is the only safe first-status + # transport: child presence/type may differ across faulty ranks and must be + # voted before any rank tries to enter a child collective. + del optimizer + return runtime._require_binding(supplied) + + +def _hybrid_portable_live_token(optimizer): + from gefen import portable_runtime as runtime + + children = _hybrid_children(optimizer) + return ( + optimizer.backup_optimizer, + optimizer._deterministic, + optimizer._hybrid_post_sharding_finalized, + id(optimizer._hybrid_sharding_manifest), + id(optimizer._hybrid_codebook_process_group), + id(optimizer._state_param_owner), + runtime._portable_value_token(optimizer._hybrid_fqn_roles), + runtime._portable_value_token(optimizer._hybrid_finalized_slots), + tuple( + ( + None if parameter is None else id(parameter), + shard.sort_key, + ) + for parameter, shard in optimizer._hybrid_local_shard_bindings + ), + tuple((role, id(child), runtime._portable_live_token(child)) for role, child in children), + ) + + +def _validate_hybrid_portable_readiness(optimizer, binding): + from gefen import portable_runtime as runtime + from gefen.checkpoint import CheckpointProcessGroupBinding + from gefen.hybrid import GefenMuonHybrid + + if type(binding) is not CheckpointProcessGroupBinding: + raise TypeError("portable Hybrid readiness requires an exact checkpoint binding") + GefenMuonHybrid._reject_rebinding_method_shadows(optimizer) + children = _hybrid_children(optimizer) + if not optimizer._canonical_identity_ready(): + raise RuntimeError("portable Hybrid state requires a finalized composite binding") + if optimizer._hybrid_codebook_process_group is None: + raise RuntimeError("portable Hybrid state requires one shared codebook scope") + routing = dict(optimizer._hybrid_fqn_roles) + if ( + type(optimizer._hybrid_fqn_roles) is not tuple + or len(routing) != len(optimizer._hybrid_fqn_roles) + or tuple(sorted(routing.items())) != optimizer._hybrid_fqn_roles + ): + raise RuntimeError("portable Hybrid routing metadata is not canonical") + expected_routing = {} + layouts = {} + common_step = None + deterministic = None + for role, child in children: + if child._gefen_codebook_process_group is not optimizer._hybrid_codebook_process_group: + raise RuntimeError("portable Hybrid children must share one exact codebook binding") + child_binding = runtime._preflight_transport_binding(child, binding) + runtime._validate_supplied_binding(binding, child_binding) + layouts[role] = runtime._validate_live_readiness( + child, + runtime._optimizer_implementation(child), + binding, + ) + for slot in child._gefen_logical_slots: + fqn = slot.shard.parameter.fqn + if fqn in expected_routing: + raise RuntimeError("portable Hybrid child FQNs must be disjoint") + expected_routing[fqn] = role + child_step = child._canonical_common_global_step() + child_deterministic = child._deterministic + if common_step is None: + common_step = child_step + deterministic = child_deterministic + elif child_step != common_step: + raise RuntimeError("portable Hybrid child global steps must agree") + elif child_deterministic is not deterministic: + raise RuntimeError("portable Hybrid child deterministic policies must agree") + if routing != {fqn: expected_routing[fqn] for fqn in sorted(expected_routing)}: + raise RuntimeError("portable Hybrid routing does not match its child slots") + if optimizer._deterministic is not deterministic: + raise RuntimeError("portable Hybrid deterministic policy disagrees with its children") + return children, routing, layouts + + +def _hybrid_child_transaction_id(parent: str, role: str, operation: str) -> str: + if role not in _HYBRID_PORTABLE_ROLES or operation not in {"export", "import"}: + raise ValueError("invalid portable Hybrid child transaction domain") + digest = hashlib.sha256( + b"gefen.portable_hybrid.child.v1\0" + + operation.encode("ascii") + + b"\0" + + role.encode("ascii") + + b"\0" + + parent.encode("utf-8") + ).hexdigest() + return "hybrid-{}-{}-{}".format(operation, role, digest) + + +def _preflight_hybrid_portable_local( + optimizer, + *, + checkpoint_process_group, + transaction_id, + limits, +): + from gefen import portable_runtime as runtime + + binding = _hybrid_transport_binding(optimizer, checkpoint_process_group) + runtime._validate_supplied_binding(checkpoint_process_group, binding) + normalized_limits = runtime._require_limits(limits) + normalized_transaction = runtime._require_transaction_id(transaction_id) + runtime._validate_context_identity_bounds(binding, normalized_limits) + children, routing, layouts = _validate_hybrid_portable_readiness( + optimizer, + binding, + ) + context = { + **runtime._base_context(binding, HYBRID_PORTABLE_STATE_IMPLEMENTATION), + "composite_format": HYBRID_PORTABLE_STATE_FORMAT, + "composite_format_version": HYBRID_PORTABLE_STATE_FORMAT_VERSION, + "backup_optimizer": optimizer.backup_optimizer, + "roles": tuple(role for role, _child in children), + "routing": routing, + "layouts": { + role: tuple(sorted(layout.value for layout in role_layouts)) for role, role_layouts in layouts.items() + }, + "transaction_id": normalized_transaction, + } + runtime._preflight_portable_value(context, normalized_limits) + return ( + binding, + normalized_transaction, + normalized_limits, + children, + routing, + context, + runtime._context_digest(context), + _hybrid_portable_live_token(optimizer), + ) + + +def _hybrid_portable_contract_support(optimizer): + """Return dynamic composite CANONICAL_GLOBAL support without collectives.""" + + try: + from gefen.checkpoint import CheckpointProcessGroupBinding + from gefen.codebook import CodebookProcessGroupBinding + from gefen.contracts import CheckpointTransport + + scope = optimizer._hybrid_codebook_process_group + if type(scope) is not CodebookProcessGroupBinding: + return frozenset(), frozenset(), frozenset() + binding = CheckpointProcessGroupBinding( + scope.identity, + scope.local_member, + scope.process_group, + scope.collective_device, + ) + children, _routing, _layouts = _validate_hybrid_portable_readiness( + optimizer, + binding, + ) + same_topology = set() + topology_changing = set() + topology_change_kinds = set() + for _role, child in children: + supports = tuple( + support + for support in child.optimizer_contract().capabilities.checkpoints + if support.transport is CheckpointTransport.CANONICAL_GLOBAL + ) + if len(supports) != 1: + return frozenset(), frozenset(), frozenset() + support = supports[0] + same_topology.update(support.same_topology) + topology_changing.update(support.topology_changing) + topology_change_kinds.update(support.topology_change_kinds) + return ( + frozenset(same_topology), + frozenset(topology_changing), + frozenset(topology_change_kinds), + ) + except Exception: + return frozenset(), frozenset(), frozenset() + + +def _export_hybrid_portable_state( + optimizer, + *, + checkpoint_process_group, + transaction_id, + limits, +): + from gefen import portable_runtime as runtime + from gefen.portable_collective import _collective_unanimous_status + + transport = _hybrid_transport_binding(optimizer, checkpoint_process_group) + wire_limits = runtime._STATUS_FALLBACK_LIMITS + binding = None + normalized_transaction = None + normalized_limits = None + children = None + routing = None + context_digest = bytes(32) + live_token = None + error = None + try: + ( + binding, + normalized_transaction, + normalized_limits, + children, + routing, + _context, + context_digest, + live_token, + ) = _preflight_hybrid_portable_local( + optimizer, + checkpoint_process_group=checkpoint_process_group, + transaction_id=transaction_id, + limits=limits, + ) + wire_limits = normalized_limits._wire_limits() + except Exception as exc: + error = exc + _collective_unanimous_status( + transport, + error, + operation="hybrid_portable_export_preflight", + transaction_id=_HYBRID_EXPORT_PREFLIGHT_TRANSACTION, + context_digest=context_digest, + limits=wire_limits, + ) + assert ( + binding is not None + and normalized_transaction is not None + and normalized_limits is not None + and children is not None + and routing is not None + and live_token is not None + ) + + child_documents = {"muon": None, "backup": None} + for role, child in children: + child_documents[role] = child.export_portable_state( + checkpoint_process_group=binding, + transaction_id=_hybrid_child_transaction_id( + normalized_transaction, + role, + "export", + ), + limits=normalized_limits, + ) + + document = None + error = None + try: + document = build_hybrid_portable_state_document( + backup_optimizer=optimizer.backup_optimizer, + routing=routing, + children=child_documents, + ) + _validate_hybrid_portable_limits(document, normalized_limits) + _validate_hybrid_portable_readiness(optimizer, binding) + if live_token != _hybrid_portable_live_token(optimizer): + raise RuntimeError("live Hybrid state changed during portable export") + context_digest = bytes.fromhex(document["completion"]["digest"]) + except Exception as exc: + error = exc + _collective_unanimous_status( + binding, + error, + operation="hybrid_portable_export_finalize", + transaction_id=normalized_transaction, + context_digest=context_digest, + limits=wire_limits, + ) + assert document is not None + return document + + +def _import_hybrid_portable_state( + optimizer, + state, + *, + checkpoint_process_group, + transaction_id, + limits, +) -> None: + from gefen import portable_runtime as runtime + from gefen.portable_collective import _collective_unanimous_status + + transport = _hybrid_transport_binding(optimizer, checkpoint_process_group) + wire_limits = runtime._STATUS_FALLBACK_LIMITS + binding = None + normalized_transaction = None + normalized_limits = None + children = None + routing = None + base_context = None + context_digest = bytes(32) + composite_live_token = None + document = None + error = None + try: + ( + binding, + normalized_transaction, + normalized_limits, + children, + routing, + base_context, + context_digest, + composite_live_token, + ) = _preflight_hybrid_portable_local( + optimizer, + checkpoint_process_group=checkpoint_process_group, + transaction_id=transaction_id, + limits=limits, + ) + wire_limits = normalized_limits._wire_limits() + from gefen.portable_state import ( + _bounded_clone, + _normalize_gefen_portable_state_document, + ) + + document = normalize_hybrid_portable_state_document(_bounded_clone(state, normalized_limits, collective=True)) + for role in _HYBRID_PORTABLE_ROLES: + child_document = document["children"][role] + if child_document is not None: + document["children"][role] = _normalize_gefen_portable_state_document( + child_document, + limits=normalized_limits, + expected_implementation=_HYBRID_PORTABLE_CHILD_IMPLEMENTATIONS[role], + ) + if document["routing"] != routing: + raise ValueError("portable Hybrid document routing does not match the target") + expected_presence = { + role: child is not None for role, child in (("muon", optimizer.muon), ("backup", optimizer.backup)) + } + if any( + (document["children"][role] is not None) is not expected_presence[role] for role in _HYBRID_PORTABLE_ROLES + ): + raise ValueError("portable Hybrid child presence does not match the target") + context_digest = bytes.fromhex(document["completion"]["digest"]) + except Exception as exc: + error = exc + _collective_unanimous_status( + transport, + error, + operation="hybrid_portable_import_document", + transaction_id=_HYBRID_IMPORT_PREFLIGHT_TRANSACTION, + context_digest=context_digest, + limits=wire_limits, + ) + assert ( + binding is not None + and normalized_transaction is not None + and normalized_limits is not None + and children is not None + and routing is not None + and base_context is not None + and composite_live_token is not None + and document is not None + ) + + staged_children = [] + target_children = {} + error = None + try: + for role, child in children: + implementation = runtime._optimizer_implementation(child) + staged, live_token, target_fragment = runtime._stage_portable_import( + child, + implementation, + binding, + normalized_limits, + document["children"][role], + ) + staged_children.append((role, child, implementation, staged, live_token)) + target_children[role] = runtime._target_context( + runtime._base_context(binding, implementation), + target_fragment, + ) + target_context = { + **base_context, + "document_digest": document["completion"]["digest"], + "target_children": target_children, + } + runtime._preflight_portable_value(target_context, normalized_limits) + context_digest = runtime._context_digest(target_context) + except Exception as exc: + error = exc + context_digest = runtime._context_digest( + { + **base_context, + "document_digest": document["completion"]["digest"], + } + ) + _collective_unanimous_status( + binding, + error, + operation="hybrid_portable_import_prepare", + transaction_id=normalized_transaction, + context_digest=context_digest, + limits=wire_limits, + ) + assert len(staged_children) == len(children) + + target_deterministic = next( + document["children"][role]["common"]["gefen_deterministic"] for role, _child in children + ) + from gefen.gefen import Gefen + + commit_staged = Gefen._commit_staged_load_state_dict + + error = None + try: + _validate_hybrid_portable_readiness(optimizer, binding) + if composite_live_token != _hybrid_portable_live_token(optimizer): + raise RuntimeError("live Hybrid state changed after portable import preparation") + for _role, child, _implementation, _staged, live_token in staged_children: + if live_token != runtime._portable_live_token(child): + raise RuntimeError("live Hybrid child changed after portable import preparation") + except Exception as exc: + error = exc + _collective_unanimous_status( + binding, + error, + operation="hybrid_portable_import_freshness", + transaction_id=normalized_transaction, + context_digest=context_digest, + limits=wire_limits, + ) + + for _role, child, _implementation, staged, _live_token in staged_children: + commit_staged(child, staged) + dict.__setitem__( + optimizer.__dict__, + "_deterministic", + target_deterministic, + ) + + +__all__ = [ + "HYBRID_PORTABLE_STATE_COVERAGE", + "HYBRID_PORTABLE_STATE_FORMAT", + "HYBRID_PORTABLE_STATE_FORMAT_VERSION", + "HYBRID_PORTABLE_STATE_IMPLEMENTATION", + "build_hybrid_portable_state_document", + "normalize_hybrid_portable_state_document", +] diff --git a/src/gefen/portable_runtime.py b/src/gefen/portable_runtime.py index 353fdc0..84f1843 100644 --- a/src/gefen/portable_runtime.py +++ b/src/gefen/portable_runtime.py @@ -652,6 +652,10 @@ def _validate_optimizer_shell( raise RuntimeError("portable state requires inactive capturable stacks") if optimizer._gefen_global_step_by_device or optimizer._sr_seed_by_device: raise RuntimeError("portable state does not support device-authoritative counters") + if optimizer.state_offload_poisoned: + raise RuntimeError("portable state does not support poisoned optimizer state") + if optimizer.state_offload_active: + raise RuntimeError("portable state does not support active optimizer-state offload") if torch.compiler.is_compiling(): raise RuntimeError("portable state cannot run while compiling") if torch.cuda.is_available() and torch.cuda.is_current_stream_capturing(): @@ -849,6 +853,8 @@ def live_group_options(group): optimizer.fused, optimizer.verbose, optimizer._fused_build_ok, + optimizer.state_offload_device, + optimizer.state_offload_poisoned, id(optimizer._gefen_codebook_process_group), _portable_value_token(optimizer._serialized_codebook_scope()), id(optimizer._gefen_sharding_manifest), diff --git a/tests/test_hybrid_rebinding.py b/tests/test_hybrid_rebinding.py new file mode 100644 index 0000000..284f30c --- /dev/null +++ b/tests/test_hybrid_rebinding.py @@ -0,0 +1,487 @@ +"""Focused CPU coverage for atomic Gefen-backed Hybrid rebinding.""" + +import pytest +import torch + +from gefen import GefenMuonHybrid +from gefen.codebook import CodebookProcessGroupBinding +from gefen.contracts import ( + LogicalSlice, + ParameterIdentity, + ParameterLayout, + PlacementKind, + ProcessGroupIdentity, + ShardIdentity, + ShardPlacement, + ShardingManifest, +) +from gefen.rebinding import ParameterRebinding + + +_MEMBER = "rank:0" + + +def _optimizer(*, backup_optimizer="gefen", sharded_mode="distributed"): + matrix = torch.nn.Parameter(torch.arange(4, dtype=torch.float32).reshape(2, 2)) + bias = torch.nn.Parameter(torch.arange(4, dtype=torch.float32)) + optimizer = GefenMuonHybrid( + [("layer.weight", matrix)], + [("layer.bias", bias)], + lr=1e-3, + fused=False, + sharded_mode=sharded_mode, + backup_optimizer=backup_optimizer, + ) + return optimizer, matrix, bias + + +def _ungrouped_replicated(fqn, shape): + identity = ParameterIdentity(fqn, shape) + return ShardIdentity( + identity, + ParameterLayout.REPLICATED, + LogicalSlice.full(identity), + ) + + +def _grouped_replicated(fqn, shape, group, member): + identity = ParameterIdentity(fqn, shape) + coordinate = group.ordered_members.index(member) + return ShardIdentity( + identity, + ParameterLayout.REPLICATED, + LogicalSlice.full(identity), + placements=( + ShardPlacement( + "checkpoint", + PlacementKind.REPLICATE, + coordinate, + len(group.ordered_members), + ), + ), + process_group=group, + local_member=member, + ) + + +def _flat_shards(fqn, shape, group, lengths): + identity = ParameterIdentity(fqn, shape) + offset = 0 + shards = [] + for coordinate, (member, length) in enumerate(zip(group.ordered_members, lengths)): + shards.append( + ShardIdentity( + identity, + ParameterLayout.FLATTENED_ELEMENT_SHARD, + LogicalSlice(offset, length), + placements=( + ShardPlacement( + "checkpoint", + PlacementKind.FLAT_SHARD, + coordinate, + len(group.ordered_members), + ), + ), + process_group=group, + local_member=member, + ) + ) + offset += length + return tuple(shards) + + +def _owner_shards(fqn, shape, group, owner): + identity = ParameterIdentity(fqn, shape) + shards = [] + for coordinate, member in enumerate(group.ordered_members): + shards.append( + ShardIdentity( + identity, + ParameterLayout.WHOLE_PARAMETER_OWNER, + (LogicalSlice.full(identity) if member == owner else LogicalSlice(0, 0)), + placements=( + ShardPlacement( + "checkpoint", + PlacementKind.WHOLE_PARAMETER_OWNER, + coordinate, + len(group.ordered_members), + ), + ), + process_group=group, + local_member=member, + owner=owner, + ) + ) + return tuple(shards) + + +def _snapshot(optimizer): + return { + "children": tuple((child, child.__dict__.copy()) for child in optimizer._subopts), + "owner": optimizer._state_param_owner, + "finalized": optimizer._hybrid_post_sharding_finalized, + "manifest": optimizer._hybrid_sharding_manifest, + "local": optimizer._hybrid_local_shard_bindings, + "roles": optimizer._hybrid_fqn_roles, + "binding": optimizer._hybrid_codebook_process_group, + "slots": optimizer._hybrid_finalized_slots, + } + + +def _assert_snapshot(optimizer, snapshot): + assert optimizer._state_param_owner is snapshot["owner"] + assert optimizer._hybrid_post_sharding_finalized is snapshot["finalized"] + assert optimizer._hybrid_sharding_manifest is snapshot["manifest"] + assert optimizer._hybrid_local_shard_bindings is snapshot["local"] + assert optimizer._hybrid_fqn_roles is snapshot["roles"] + assert optimizer._hybrid_codebook_process_group is snapshot["binding"] + assert optimizer._hybrid_finalized_slots is snapshot["slots"] + assert len(optimizer._subopts) == len(snapshot["children"]) + for live, (expected_child, expected_attributes) in zip(optimizer._subopts, snapshot["children"]): + assert live is expected_child + assert set(live.__dict__) == set(expected_attributes) + for name, expected in expected_attributes.items(): + assert live.__dict__[name] is expected + + +def test_composite_post_sharding_publishes_children_and_rebuilds_routing(): + optimizer, old_matrix, old_bias = _optimizer() + matrix = torch.nn.Parameter(torch.full((2, 2), 7.0)) + bias = torch.nn.Parameter(torch.full((4,), 11.0)) + group = ProcessGroupIdentity("checkpoint", (_MEMBER,)) + binding = CodebookProcessGroupBinding( + group, + _MEMBER, + None, + torch.device("cpu"), + ) + matrix_shard = _grouped_replicated( + "Model.Layer.Weight", + (2, 2), + group, + _MEMBER, + ) + bias_shard = _grouped_replicated( + "Model.Layer.Bias", + (4,), + group, + _MEMBER, + ) + manifest = ShardingManifest((matrix_shard, bias_shard)) + + optimizer.post_sharding( + ( + ParameterRebinding(old_bias, bias, bias_shard), + ParameterRebinding(old_matrix, matrix, matrix_shard), + ), + manifest=manifest, + codebook_process_group=binding, + ) + + assert optimizer._canonical_identity_ready() + assert optimizer._codebook_scope_ready() + assert optimizer.sharding_manifest() is manifest + assert optimizer.codebook_process_group_binding() is binding + assert optimizer.muon.codebook_process_group_binding() is binding + assert optimizer.backup.codebook_process_group_binding() is binding + assert optimizer._hybrid_fqn_roles == ( + ("Model.Layer.Bias", "backup"), + ("Model.Layer.Weight", "muon"), + ) + assert optimizer.parameter_routing() == optimizer._hybrid_fqn_roles + assert optimizer.muon.sharding_manifest().shards == (matrix_shard,) + assert optimizer.backup.sharding_manifest().shards == (bias_shard,) + assert tuple(shard.parameter.fqn for _, shard in optimizer.shard_bindings()) == ( + "Model.Layer.Bias", + "Model.Layer.Weight", + ) + assert optimizer.parameter_identity(matrix) == matrix_shard.parameter + assert optimizer.shard_identity(bias) == bias_shard + with pytest.raises(KeyError, match="no finalized"): + optimizer.shard_identity(old_matrix) + + optimizer.state[bias]["sentinel"] = 13 + optimizer.state[matrix]["sentinel"] = 17 + assert optimizer.backup.state[bias]["sentinel"] == 13 + assert optimizer.muon.state[matrix]["sentinel"] == 17 + assert set(optimizer._state_param_owner) == {id(matrix), id(bias)} + + +def test_composite_supports_muon_nonowner_and_flat_backup_without_binding(): + optimizer, old_matrix, old_bias = _optimizer() + group = ProcessGroupIdentity("data_parallel", ("rank:0", "rank:1")) + matrix_shards = _owner_shards( + "Model.Layer.Weight", + (2, 2), + group, + owner="rank:1", + ) + bias_shards = _flat_shards( + "Model.Layer.Bias", + (4,), + group, + lengths=(2, 2), + ) + local_matrix = matrix_shards[0] + local_bias = bias_shards[0] + bias = torch.nn.Parameter(torch.full((2,), 5.0)) + manifest = ShardingManifest(matrix_shards + bias_shards) + + optimizer.post_sharding( + ( + ParameterRebinding(old_matrix, None, local_matrix), + ParameterRebinding(old_bias, bias, local_bias), + ), + manifest=manifest, + ) + + assert optimizer._canonical_identity_ready() + assert not optimizer._codebook_scope_ready() + assert optimizer.muon.param_groups[0]["params"] == [] + assert optimizer.muon.state == {} + assert optimizer.backup.param_groups[0]["params"] == [bias] + assert optimizer.codebook_process_group_binding() is None + assert optimizer.shard_bindings() == ( + (bias, local_bias), + (None, local_matrix), + ) + assert set(optimizer._state_param_owner) == {id(bias)} + + +def test_second_child_staging_failure_leaves_both_children_unchanged(): + optimizer, old_matrix, old_bias = _optimizer() + matrix = torch.nn.Parameter(torch.full((2, 2), 3.0)) + invalid_bias = torch.nn.Parameter(torch.full((3,), 4.0)) + matrix_shard = _ungrouped_replicated("Model.Layer.Weight", (2, 2)) + bias_shard = _ungrouped_replicated("Model.Layer.Bias", (4,)) + manifest = ShardingManifest((matrix_shard, bias_shard)) + snapshot = _snapshot(optimizer) + matrix_before = matrix.detach().clone() + bias_before = invalid_bias.detach().clone() + + with pytest.raises(ValueError, match="complete parameter storage"): + optimizer.post_sharding( + ( + ParameterRebinding(old_matrix, matrix, matrix_shard), + ParameterRebinding(old_bias, invalid_bias, bias_shard), + ), + manifest=manifest, + ) + + _assert_snapshot(optimizer, snapshot) + assert torch.equal(matrix, matrix_before) + assert torch.equal(invalid_bias, bias_before) + assert not optimizer._canonical_identity_ready() + + +def test_external_child_parameter_replacement_can_be_finalized_atomically(): + optimizer, old_matrix, old_bias = _optimizer() + matrix = torch.nn.Parameter(torch.full((2, 2), 3.0)) + bias = torch.nn.Parameter(torch.full((4,), 4.0)) + optimizer.muon.param_groups[0]["params"][0] = matrix + optimizer.backup.param_groups[0]["params"][0] = bias + matrix_shard = _ungrouped_replicated("Model.Layer.Weight", (2, 2)) + bias_shard = _ungrouped_replicated("Model.Layer.Bias", (4,)) + + optimizer.post_sharding( + ( + ParameterRebinding(old_matrix, matrix, matrix_shard), + ParameterRebinding(old_bias, bias, bias_shard), + ), + manifest=ShardingManifest((matrix_shard, bias_shard)), + ) + + assert optimizer._canonical_identity_ready() + assert optimizer.muon.param_groups[0]["params"] == [matrix] + assert optimizer.backup.param_groups[0]["params"] == [bias] + assert set(optimizer._state_param_owner) == {id(matrix), id(bias)} + + +def test_cross_child_target_storage_overlap_is_rejected_atomically(): + optimizer, old_matrix, old_bias = _optimizer() + storage = torch.arange(8, dtype=torch.float32) + matrix = torch.nn.Parameter(storage[:4].view(2, 2)) + bias = torch.nn.Parameter(storage[2:6]) + assert matrix.untyped_storage().data_ptr() == bias.untyped_storage().data_ptr() + matrix_shard = _ungrouped_replicated("Model.Layer.Weight", (2, 2)) + bias_shard = _ungrouped_replicated("Model.Layer.Bias", (4,)) + snapshot = _snapshot(optimizer) + + with pytest.raises(ValueError, match="storage ranges must not overlap"): + optimizer.post_sharding( + ( + ParameterRebinding(old_matrix, matrix, matrix_shard), + ParameterRebinding(old_bias, bias, bias_shard), + ), + manifest=ShardingManifest((matrix_shard, bias_shard)), + ) + + _assert_snapshot(optimizer, snapshot) + + +def test_global_source_target_and_fqn_guards_run_before_publication(): + optimizer, old_matrix, old_bias = _optimizer() + matrix = torch.nn.Parameter(torch.ones(2, 2)) + bias = torch.nn.Parameter(torch.ones(4)) + matrix_shard = _ungrouped_replicated("Model.Shared", (2, 2)) + duplicate_fqn = _ungrouped_replicated("Model.Shared", (4,)) + snapshot = _snapshot(optimizer) + + with pytest.raises(ValueError, match="FQNs must be unique"): + optimizer.post_sharding( + ( + ParameterRebinding(old_matrix, matrix, matrix_shard), + ParameterRebinding(old_bias, bias, duplicate_fqn), + ), + manifest=ShardingManifest((matrix_shard,)), + ) + _assert_snapshot(optimizer, snapshot) + + matrix_shard = _ungrouped_replicated("Model.Layer.Weight", (2, 2)) + bias_shard = _ungrouped_replicated("Model.Layer.Bias", (4,)) + with pytest.raises(ValueError, match="source tensors must be unique"): + optimizer.post_sharding( + ( + ParameterRebinding(old_matrix, matrix, matrix_shard), + ParameterRebinding(old_matrix, bias, bias_shard), + ), + manifest=ShardingManifest((matrix_shard, bias_shard)), + ) + _assert_snapshot(optimizer, snapshot) + + with pytest.raises(ValueError, match="cannot steal"): + optimizer.post_sharding( + ( + ParameterRebinding(old_matrix, old_bias, matrix_shard), + ParameterRebinding(old_bias, bias, bias_shard), + ), + manifest=ShardingManifest((matrix_shard, bias_shard)), + ) + _assert_snapshot(optimizer, snapshot) + + +def test_shared_binding_must_match_every_child_shard_before_staging(): + optimizer, old_matrix, old_bias = _optimizer() + shard_group = ProcessGroupIdentity("shards", (_MEMBER,)) + binding_group = ProcessGroupIdentity("checkpoint", (_MEMBER,)) + binding = CodebookProcessGroupBinding( + binding_group, + _MEMBER, + None, + torch.device("cpu"), + ) + matrix_shard = _grouped_replicated( + "Model.Layer.Weight", + (2, 2), + shard_group, + _MEMBER, + ) + bias_shard = _grouped_replicated( + "Model.Layer.Bias", + (4,), + shard_group, + _MEMBER, + ) + snapshot = _snapshot(optimizer) + + with pytest.raises(ValueError, match="shared codebook"): + optimizer.post_sharding( + ( + ParameterRebinding( + old_matrix, + torch.nn.Parameter(torch.ones(2, 2)), + matrix_shard, + ), + ParameterRebinding( + old_bias, + torch.nn.Parameter(torch.ones(4)), + bias_shard, + ), + ), + manifest=ShardingManifest((matrix_shard, bias_shard)), + codebook_process_group=binding, + ) + + _assert_snapshot(optimizer, snapshot) + + +@pytest.mark.parametrize("role", ["muon", "backup"]) +def test_one_child_hybrids_support_convenience_rebinding(role): + if role == "muon": + old = torch.nn.Parameter(torch.ones(2, 2)) + new = torch.nn.Parameter(torch.full((2, 2), 2.0)) + optimizer = GefenMuonHybrid( + [("layer.weight", old)], + [], + lr=1e-3, + fused=False, + ) + identity = ParameterIdentity("Model.Layer.Weight", (2, 2)) + else: + old = torch.nn.Parameter(torch.ones(4)) + new = torch.nn.Parameter(torch.full((4,), 2.0)) + optimizer = GefenMuonHybrid( + [], + [("layer.bias", old)], + lr=1e-3, + fused=False, + ) + identity = ParameterIdentity("Model.Layer.Bias", (4,)) + + optimizer.rebind_parameter(old, new, identity=identity) + + assert optimizer._canonical_identity_ready() + assert optimizer.parameter_identity(new) == identity + assert optimizer.shard_bindings() == ((new, optimizer.shard_identity(new)),) + assert set(optimizer._state_param_owner) == {id(new)} + + +def test_adamw_backup_rejects_before_any_child_or_hybrid_mutation(): + optimizer, old_matrix, old_bias = _optimizer(backup_optimizer="adamw") + matrix_shard = _ungrouped_replicated("Model.Layer.Weight", (2, 2)) + bias_shard = _ungrouped_replicated("Model.Layer.Bias", (4,)) + snapshot = _snapshot(optimizer) + + with pytest.raises(NotImplementedError, match="AdamW"): + optimizer.post_sharding( + ( + ParameterRebinding( + old_matrix, + torch.nn.Parameter(torch.ones(2, 2)), + matrix_shard, + ), + ParameterRebinding( + old_bias, + torch.nn.Parameter(torch.ones(4)), + bias_shard, + ), + ), + manifest=ShardingManifest((matrix_shard, bias_shard)), + ) + + _assert_snapshot(optimizer, snapshot) + assert not optimizer._canonical_identity_ready() + + +def test_composite_guard_fails_closed_after_owner_metadata_corruption(): + optimizer, old_matrix, old_bias = _optimizer() + matrix = torch.nn.Parameter(torch.ones(2, 2)) + bias = torch.nn.Parameter(torch.ones(4)) + matrix_shard = _ungrouped_replicated("Model.Layer.Weight", (2, 2)) + bias_shard = _ungrouped_replicated("Model.Layer.Bias", (4,)) + optimizer.post_sharding( + ( + ParameterRebinding(old_matrix, matrix, matrix_shard), + ParameterRebinding(old_bias, bias, bias_shard), + ), + manifest=ShardingManifest((matrix_shard, bias_shard)), + ) + assert optimizer._canonical_identity_ready() + + optimizer._state_param_owner = {} + + assert not optimizer._canonical_identity_ready() + with pytest.raises(RuntimeError, match="changed outside post_sharding"): + optimizer.sharding_manifest() + with pytest.raises(RuntimeError, match="changed outside post_sharding"): + _ = optimizer.state diff --git a/tests/test_optimizer_contracts.py b/tests/test_optimizer_contracts.py index e7348e9..9504f1c 100644 --- a/tests/test_optimizer_contracts.py +++ b/tests/test_optimizer_contracts.py @@ -25,6 +25,7 @@ StateGeometry, StateKeyMatch, StateMovementProvider, + StateOffloadProvider, StateScope, StateVariant, TopologyChange, @@ -36,8 +37,10 @@ def _values_equal(left, right): if torch.is_tensor(left) or torch.is_tensor(right): - return torch.is_tensor(left) and torch.is_tensor(right) and torch.equal( - left, right + return ( + torch.is_tensor(left) + and torch.is_tensor(right) + and torch.equal(left, right) ) if type(left) is not type(right): return False @@ -92,8 +95,7 @@ def _matching_variants( if set(variant.fields) == keys and layout in variant.layouts and ( - variant.parameter_ranks is None - or parameter_rank in variant.parameter_ranks + variant.parameter_ranks is None or parameter_rank in variant.parameter_ranks ) and parameter_rank not in variant.excluded_parameter_ranks and variant.sharded_mode == sharded_mode @@ -105,8 +107,7 @@ def _training_support(contract, layout, mode=None): matches = [ item for item in contract.capabilities.training - if item.layout is layout - and (mode is None or item.sharded_mode == mode) + if item.layout is layout and (mode is None or item.sharded_mode == mode) ] assert len(matches) == 1 return matches[0] @@ -136,16 +137,16 @@ def test_plain_contract_matches_live_persistent_state(factored_v_2d): assert contract.schema_version == CONTRACT_SCHEMA_VERSION assert contract.implementation == "gefen.Gefen" assert contract.capabilities.supported_parameter_ranks is None - assert { - item.layout for item in contract.capabilities.training - } == { + assert {item.layout for item in contract.capabilities.training} == { ParameterLayout.REPLICATED, ParameterLayout.FLATTENED_ELEMENT_SHARD, _DTENSOR, } dtensor_training = _training_support(contract, _DTENSOR) assert dtensor_training.mesh_dimensions == (1,) - assert dtensor_training.process_group_scope is ProcessGroupScope.INFERRED_DEVICE_MESH + assert ( + dtensor_training.process_group_scope is ProcessGroupScope.INFERRED_DEVICE_MESH + ) dcp = _checkpoint_support( contract, CheckpointTransport.PYTORCH_RANK_LOCAL, _DTENSOR ) @@ -172,9 +173,7 @@ def test_plain_contract_matches_live_persistent_state(factored_v_2d): assert contract.capabilities.atomic_state_movement assert not contract.capabilities.state_offload assert Precision.FLOAT64 in contract.capabilities.precisions - flattened = _training_support( - contract, ParameterLayout.FLATTENED_ELEMENT_SHARD - ) + flattened = _training_support(contract, ParameterLayout.FLATTENED_ELEMENT_SHARD) assert flattened.process_group_scope is ProcessGroupScope.NONE assert optimizer.state[param] == {"name": "layer.weight"} @@ -208,9 +207,7 @@ def test_plain_contract_matches_live_persistent_state(factored_v_2d): state_dict = optimizer.state_dict() common = { field.name - for field in contract.state_layout.fields_for_scope( - StateScope.OPTIMIZER_COMMON - ) + for field in contract.state_layout.fields_for_scope(StateScope.OPTIMIZER_COMMON) } assert common == { "gefen_global_step", @@ -220,9 +217,7 @@ def test_plain_contract_matches_live_persistent_state(factored_v_2d): } required_common = { field.name - for field in contract.state_layout.fields_for_scope( - StateScope.OPTIMIZER_COMMON - ) + for field in contract.state_layout.fields_for_scope(StateScope.OPTIMIZER_COMMON) if not field.optional } assert required_common.issubset(state_dict) @@ -408,9 +403,7 @@ def test_muon_contract_separates_mode_topology_and_state_extent( common_names = { field.name - for field in contract.state_layout.fields_for_scope( - StateScope.OPTIMIZER_COMMON - ) + for field in contract.state_layout.fields_for_scope(StateScope.OPTIMIZER_COMMON) } assert "gefen_muon_distributed" not in common_names if sharded_mode == "distributed": @@ -491,12 +484,20 @@ 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") - assert not contract.capabilities.explicit_process_group_codebook_scope + gefen_backed = backup_optimizer == "gefen" + assert contract.capabilities.explicit_process_group_codebook_scope is gefen_backed + assert contract.capabilities.shard_rebinding is gefen_backed + assert contract.capabilities.post_sharding is gefen_backed + assert not contract.capabilities.canonical_state_io assert not contract.capabilities.atomic_state_movement assert not contract.capabilities.state_offload - assert contract.children[0].contract.capabilities.explicit_process_group_codebook_scope + assert contract.children[ + 0 + ].contract.capabilities.explicit_process_group_codebook_scope if backup_optimizer == "gefen": - assert contract.children[1].contract.capabilities.explicit_process_group_codebook_scope + assert contract.children[ + 1 + ].contract.capabilities.explicit_process_group_codebook_scope assert {field.name for field in contract.state_layout.fields} == { "backup_optimizer" } @@ -520,7 +521,8 @@ def test_muon_contract_keeps_mixed_normuon_variants_in_one_mode(): sharded_mode="exact", ) variants = { - item.name for item in optimizer.optimizer_contract().state_layout.parameter_variants + item.name + for item in optimizer.optimizer_contract().state_layout.parameter_variants } assert "quantized_muon_replicated_exact" in variants assert "quantized_normuon_replicated_exact" in variants @@ -622,9 +624,7 @@ def test_muon_mixed_approx_distributed_checkpoint_is_same_topology_only(): _DTENSOR, ) assert len(support) == 1 - assert support[0].required_sharded_modes == frozenset( - {"approx", "distributed"} - ) + assert support[0].required_sharded_modes == frozenset({"approx", "distributed"}) assert not support[0].topology_changing @@ -747,6 +747,7 @@ def test_all_public_contract_exports_resolve(): assert all(getattr(gefen, name) is not None for name in contracts.__all__) assert gefen.StateMovementProvider is StateMovementProvider + assert gefen.StateOffloadProvider is StateOffloadProvider def test_portable_global_transport_is_defined_but_not_claimed_before_integration(): diff --git a/tests/test_portable_dcp_nccl.py b/tests/test_portable_dcp_nccl.py new file mode 100644 index 0000000..e543e4e --- /dev/null +++ b/tests/test_portable_dcp_nccl.py @@ -0,0 +1,364 @@ +"""Warning-strict CUDA/NCCL coverage for portable DCP persistence.""" + +from datetime import timedelta +import multiprocessing as mp +import os +import queue as queue_module +import tempfile +import traceback +import warnings + +import pytest +import torch +import torch.distributed as dist +from torch.distributed.checkpoint import FileSystemReader, FileSystemWriter + +from gefen import load_portable_dcp, save_portable_dcp +from gefen.checkpoint import CheckpointProcessGroupBinding +from gefen.codebook import CodebookProcessGroupBinding +from gefen.contracts import ( + LogicalSlice, + ParameterIdentity, + ParameterLayout, + PlacementKind, + ProcessGroupIdentity, + ShardIdentity, + ShardPlacement, + ShardingManifest, +) +from gefen.gefen import Gefen +from gefen.portable_state import PortableStateLimits +from gefen.rebinding import ParameterRebinding + + +_WORLD = 2 +_FQN = "model.weight" + + +def _limits(): + return PortableStateLimits( + max_fragment_tensor_bytes=1 << 20, + max_collective_tensor_bytes=4 << 20, + max_collective_metadata_bytes=4 << 20, + chunk_bytes=13, + max_members=4, + max_metadata_bytes=1 << 20, + max_tree_nodes=10_000, + max_tree_depth=32, + max_container_items=10_000, + max_string_bytes=16 << 10, + max_integer_bytes=128, + max_tensors=256, + max_tensor_rank=8, + diagnostic_bytes=1024, + ) + + +def _members(): + return tuple("rank:{}".format(rank) for rank in range(_WORLD)) + + +def _replicated_manifest(identity, group): + shards = tuple( + ShardIdentity( + identity, + ParameterLayout.REPLICATED, + LogicalSlice.full(identity), + placements=( + ShardPlacement( + "checkpoint", + PlacementKind.REPLICATE, + coordinate, + len(group.ordered_members), + ), + ), + process_group=group, + local_member=member, + ) + for coordinate, member in enumerate(group.ordered_members) + ) + return ShardingManifest(shards), shards + + +def _optimizer(rank, group, *, deterministic): + device = torch.device("cuda", rank) + parameter = torch.nn.Parameter(torch.linspace(-0.4, 0.3, 8, dtype=torch.float32, device=device).reshape(2, 4)) + optimizer = Gefen( + [("weight", parameter)], + lr=2.5e-3, + betas=(0.8, 0.97), + eps=2.0e-8, + weight_decay=0.03, + fused=False, + force_2d_period_one=True, + factored_v_2d=False, + deterministic=deterministic, + ) + manifest, shards = _replicated_manifest(ParameterIdentity(_FQN, (2, 4)), group) + member = _members()[rank] + codebook_binding = CodebookProcessGroupBinding( + group, + member, + dist.group.WORLD, + device, + ) + checkpoint_binding = CheckpointProcessGroupBinding( + group, + member, + dist.group.WORLD, + device, + ) + optimizer.post_sharding( + (ParameterRebinding(parameter, parameter, shards[rank]),), + manifest=manifest, + codebook_process_group=codebook_binding, + ) + return optimizer, parameter, checkpoint_binding + + +def _codebook(device): + return torch.linspace(-1.0, 1.0, 256, dtype=torch.float32, device=device) + + +def _momentum(device): + return ( + torch.tensor( + [ + -2147483648, + 0, + 1, + -2147483647, + 1056964608, + -1090519040, + 1078984704, + -1058013184, + ], + dtype=torch.int32, + device=device, + ) + .view(torch.float32) + .reshape(2, 4) + .clone() + ) + + +def _second_moment(device): + return ( + torch.tensor( + [ + -2147483648, + 0, + 1, + 8388608, + 1048576000, + 1065353216, + 1073741824, + 2139095039, + ], + dtype=torch.int32, + device=device, + ) + .view(torch.float32) + .reshape(2, 4) + .clone() + ) + + +def _seed_state(optimizer, parameter): + momentum = _momentum(parameter.device) + flat = momentum.reshape(-1) + indices = torch.where( + torch.signbit(flat), + torch.zeros(flat.numel(), dtype=torch.uint8, device=parameter.device), + torch.full((flat.numel(),), 255, dtype=torch.uint8, device=parameter.device), + ) + optimizer._gefen_global_step = 9 + optimizer._gefen_codebook = _codebook(parameter.device) + optimizer.state[parameter].update( + { + "automatic_period": 1, + "step": 7, + "m_codebook": indices.reshape(-1, 1), + "m_magnitude": flat.abs().reshape(-1, 1).clone(), + "vmean": _second_moment(parameter.device).reshape(-1, 1), + "vmean_step": 6, + } + ) + + +def _tensor_bits_equal(left, right): + return ( + torch.is_tensor(left) + and torch.is_tensor(right) + and left.dtype == right.dtype + and tuple(left.shape) == tuple(right.shape) + and torch.equal( + left.detach().cpu().contiguous().view(torch.uint8), + right.detach().cpu().contiguous().view(torch.uint8), + ) + ) + + +def _state_equal(left, right): + if set(left) != set(right): + return False + for key, expected in right.items(): + actual = left[key] + if torch.is_tensor(expected): + if not _tensor_bits_equal(actual, expected): + return False + elif actual != expected: + return False + return True + + +def _optimizer_state_equal(left, left_parameter, right, right_parameter): + return ( + left._gefen_global_step == right._gefen_global_step + and left._deterministic is right._deterministic + and _tensor_bits_equal(left._gefen_codebook, right._gefen_codebook) + and _state_equal(left.state[left_parameter], right.state[right_parameter]) + ) + + +def _strict_worker_warnings(): + warnings.simplefilter("error") + warnings.filterwarnings( + "ignore", + message="TypedStorage is deprecated.*", + category=UserWarning, + ) + + +def _worker(rank, init_file, checkpoint_dir, result_queue): + _strict_worker_warnings() + try: + torch.cuda.set_device(rank) + dist.init_process_group( + "nccl", + init_method="file://{}".format(init_file), + rank=rank, + world_size=_WORLD, + timeout=timedelta(seconds=90), + ) + group = ProcessGroupIdentity("portable_dcp_nccl", _members()) + source, source_parameter, source_binding = _optimizer( + rank, + group, + deterministic=True, + ) + _seed_state(source, source_parameter) + + save_portable_dcp( + source, + checkpoint_process_group=source_binding, + storage_writer=FileSystemWriter(checkpoint_dir), + transaction_id="portable-dcp-nccl-save-v1", + limits=_limits(), + ) + dist.barrier(device_ids=[rank]) + + target, target_parameter, target_binding = _optimizer( + rank, + group, + deterministic=False, + ) + load_portable_dcp( + target, + checkpoint_process_group=target_binding, + storage_reader=FileSystemReader(checkpoint_dir), + transaction_id="portable-dcp-nccl-load-v1", + limits=_limits(), + ) + restored_exact = _tensor_bits_equal(target_parameter, source_parameter) and _optimizer_state_equal( + target, + target_parameter, + source, + source_parameter, + ) + + gradient = torch.tensor( + [0.25, -0.5, 0.75, -1.0, 1.25, -1.5, 1.75, -2.0], + dtype=torch.float32, + device=source_parameter.device, + ).reshape(2, 4) + source_parameter.grad = gradient.clone() + target_parameter.grad = gradient.clone() + source.step() + target.step() + continuation = { + "parameter": _tensor_bits_equal(target_parameter, source_parameter), + "global_step": target._gefen_global_step == source._gefen_global_step, + "deterministic": target._deterministic is source._deterministic, + "codebook": _tensor_bits_equal(target._gefen_codebook, source._gefen_codebook), + "state": _state_equal(target.state[target_parameter], source.state[source_parameter]), + } + next_step_exact = all(continuation.values()) + result_queue.put( + { + "rank": rank, + "restored_exact": restored_exact, + "next_step_exact": next_step_exact, + "continuation": continuation, + } + ) + except BaseException: + result_queue.put({"rank": rank, "fatal_error": traceback.format_exc()}) + finally: + if dist.is_available() and dist.is_initialized(): + dist.destroy_process_group() + + +def _run_workers(checkpoint_dir): + context = mp.get_context("spawn") + result_queue = context.Queue() + descriptor, init_file = tempfile.mkstemp(prefix="gefen-portable-dcp-nccl-") + os.close(descriptor) + os.unlink(init_file) + processes = [ + context.Process( + target=_worker, + args=(rank, init_file, checkpoint_dir, result_queue), + ) + for rank in range(_WORLD) + ] + results = [] + try: + for process in processes: + process.start() + try: + for _ in processes: + results.append(result_queue.get(timeout=180)) + except queue_module.Empty: + pass + for process in processes: + process.join(timeout=10) + for process in processes: + if process.is_alive(): + process.terminate() + process.join(timeout=5) + finally: + for process in processes: + if process.is_alive(): + process.terminate() + process.join(timeout=5) + result_queue.close() + result_queue.join_thread() + if os.path.exists(init_file): + os.unlink(init_file) + assert len(results) == _WORLD, (results, [process.exitcode for process in processes]) + assert all(process.exitcode == 0 for process in processes), [process.exitcode for process in processes] + return sorted(results, key=lambda item: item["rank"]) + + +@pytest.mark.skipif( + not dist.is_available() or not dist.is_nccl_available() or torch.cuda.device_count() < _WORLD, + reason="portable DCP NCCL coverage requires NCCL and two CUDA devices", +) +def test_portable_dcp_round_trip_and_next_step_are_exact_on_nccl(tmp_path): + results = _run_workers(str(tmp_path / "checkpoint")) + + assert all("fatal_error" not in result for result in results), results + assert [result["rank"] for result in results] == [0, 1] + assert all(result["restored_exact"] and result["next_step_exact"] for result in results), results + assert all(all(result["continuation"].values()) for result in results), results diff --git a/tests/test_portable_dcp_topologies.py b/tests/test_portable_dcp_topologies.py new file mode 100644 index 0000000..9dc6a44 --- /dev/null +++ b/tests/test_portable_dcp_topologies.py @@ -0,0 +1,920 @@ +"""Warning-strict CPU/Gloo topology coverage for portable DCP.""" + +from datetime import timedelta +import multiprocessing as mp +import os +import queue as queue_module +import tempfile +import traceback +import warnings + +import pytest +import torch +import torch.distributed as dist +from torch.distributed.checkpoint import FileSystemReader, FileSystemWriter + +from gefen import load_portable_dcp, save_portable_dcp +from gefen.checkpoint import CheckpointProcessGroupBinding +from gefen.codebook import CodebookProcessGroupBinding +from gefen.contracts import ( + LogicalSlice, + ParameterIdentity, + ParameterLayout, + PlacementKind, + ProcessGroupIdentity, + ShardIdentity, + ShardPlacement, + ShardingManifest, +) +from gefen.gefen import Gefen +from gefen.gefen_muon import GefenMuon +from gefen.portable import _decode_quantized_momentum +from gefen.portable_state import PortableStateLimits +from gefen.rebinding import ParameterRebinding + + +_PLAIN_FQN = "model.weight" +_MUON_FQN = "model.matrix" + + +def _limits(): + return PortableStateLimits( + max_fragment_tensor_bytes=1 << 20, + max_collective_tensor_bytes=4 << 20, + max_collective_metadata_bytes=4 << 20, + chunk_bytes=13, + max_members=4, + max_metadata_bytes=1 << 20, + max_tree_nodes=10_000, + max_tree_depth=32, + max_container_items=10_000, + max_string_bytes=16 << 10, + max_integer_bytes=128, + max_tensors=256, + max_tensor_rank=8, + diagnostic_bytes=1024, + ) + + +def _strict_worker_warnings(): + warnings.simplefilter("error") + warnings.filterwarnings( + "ignore", + message="TypedStorage is deprecated.*", + category=UserWarning, + ) + + +def _bindings(group, local_member, process_group): + return ( + CodebookProcessGroupBinding( + group, + local_member, + process_group, + torch.device("cpu"), + ), + CheckpointProcessGroupBinding( + group, + local_member, + process_group, + torch.device("cpu"), + ), + ) + + +def _replicated_manifest(identity, group): + shards = tuple( + ShardIdentity( + identity, + ParameterLayout.REPLICATED, + LogicalSlice.full(identity), + placements=( + ShardPlacement( + "checkpoint", + PlacementKind.REPLICATE, + coordinate, + len(group.ordered_members), + ), + ), + process_group=group, + local_member=member, + ) + for coordinate, member in enumerate(group.ordered_members) + ) + return ShardingManifest(shards), shards + + +def _flat_manifest(identity, group, lengths): + if len(lengths) != len(group.ordered_members) or sum(lengths) != identity.numel: + raise AssertionError("invalid flattened test partition") + offset = 0 + shards = [] + for coordinate, (member, length) in enumerate(zip(group.ordered_members, lengths)): + shards.append( + ShardIdentity( + identity, + ParameterLayout.FLATTENED_ELEMENT_SHARD, + LogicalSlice(offset, length), + placements=( + ShardPlacement( + "checkpoint", + PlacementKind.FLAT_SHARD, + coordinate, + len(group.ordered_members), + ), + ), + process_group=group, + local_member=member, + ) + ) + offset += length + shards = tuple(shards) + return ShardingManifest(shards), shards + + +def _owner_manifest(identity, group, owner): + shards = tuple( + ShardIdentity( + identity, + ParameterLayout.WHOLE_PARAMETER_OWNER, + (LogicalSlice.full(identity) if member == owner else LogicalSlice(0, 0)), + placements=( + ShardPlacement( + "checkpoint", + PlacementKind.WHOLE_PARAMETER_OWNER, + coordinate, + len(group.ordered_members), + ), + ), + process_group=group, + local_member=member, + owner=owner, + ) + for coordinate, member in enumerate(group.ordered_members) + ) + return ShardingManifest(shards), shards + + +def _plain_optimizer(parameter, *, deterministic): + return Gefen( + [("weight", parameter)], + lr=2.5e-3, + betas=(0.8, 0.97), + eps=2.0e-8, + weight_decay=0.03, + fused=False, + force_2d_period_one=True, + factored_v_2d=False, + deterministic=deterministic, + ) + + +def _muon_optimizer(parameter, *, deterministic, normuon): + return GefenMuon( + [("matrix", parameter)], + lr=3.0e-3, + weight_decay=0.02, + momentum=0.85, + nesterov=False, + ns_steps=2, + fused=False, + sharded_mode="distributed", + deterministic=deterministic, + normuon=normuon, + normuon_beta2=0.9, + normuon_eps=3.0e-8, + ) + + +def _codebook(): + return torch.linspace(-1.0, 1.0, 256, dtype=torch.float32) + + +def _plain_initial_parameter(): + return torch.linspace(-0.4, 0.3, 8, dtype=torch.float32).reshape(2, 4) + + +def _plain_momentum(): + return ( + torch.tensor( + [ + -2147483648, + 0, + 1, + -2147483647, + 1056964608, + -1090519040, + 1078984704, + -1058013184, + ], + dtype=torch.int32, + ) + .view(torch.float32) + .reshape(2, 4) + .clone() + ) + + +def _plain_second_moment(): + return ( + torch.tensor( + [ + -2147483648, + 0, + 1, + 8388608, + 1048576000, + 1065353216, + 1073741824, + 2139095039, + ], + dtype=torch.int32, + ) + .view(torch.float32) + .reshape(2, 4) + .clone() + ) + + +def _plain_gradient(): + return torch.tensor( + [0.25, -0.5, 0.75, -1.0, 1.25, -1.5, 1.75, -2.0], + dtype=torch.float32, + ).reshape(2, 4) + + +def _muon_initial_parameter(): + return torch.linspace(-0.3, 0.2, 6, dtype=torch.float32).reshape(3, 2) + + +def _muon_momentum(): + return torch.tensor( + [[-0.75, 0.5], [1.25, -1.5], [2.0, -2.5]], + dtype=torch.float32, + ) + + +def _muon_normuon_v(): + return torch.tensor([[0.5], [1.5], [2.5]], dtype=torch.float32) + + +def _muon_gradient(): + return torch.tensor( + [[0.4, -0.7], [1.1, -1.3], [1.7, -1.9]], + dtype=torch.float32, + ) + + +def _quantized_period_one(momentum): + flat = momentum.reshape(-1) + indices = torch.where( + torch.signbit(flat), + torch.zeros(flat.numel(), dtype=torch.uint8), + torch.full((flat.numel(),), 255, dtype=torch.uint8), + ) + return indices.reshape(-1, 1), flat.abs().reshape(-1, 1).clone() + + +def _seed_plain_state(optimizer, parameter, momentum, second_moment): + indices, magnitudes = _quantized_period_one(momentum) + optimizer._gefen_global_step = 9 + optimizer._gefen_codebook = _codebook() + optimizer.state[parameter].update( + { + "automatic_period": 1, + "step": 7, + "m_codebook": indices, + "m_magnitude": magnitudes, + "vmean": second_moment.reshape(-1, 1).clone(), + "vmean_step": 6, + } + ) + + +def _seed_muon_state(optimizer, parameter, *, normuon): + indices, magnitudes = _quantized_period_one(_muon_momentum()) + optimizer._gefen_global_step = 13 + optimizer._gefen_codebook = _codebook() + state = optimizer.state[parameter] + state.update( + { + "automatic_period": 1, + "step": 11, + "m_codebook": indices, + "m_magnitude": magnitudes, + } + ) + if normuon: + state.update( + { + "normuon_v": _muon_normuon_v(), + "normuon_step": 10, + } + ) + + +def _bits_equal(left, right): + return ( + torch.is_tensor(left) + and torch.is_tensor(right) + and left.dtype == right.dtype + and tuple(left.shape) == tuple(right.shape) + and torch.equal( + left.detach().contiguous().view(torch.uint8), + right.detach().contiguous().view(torch.uint8), + ) + ) + + +def _state_equal(left, right): + if set(left) != set(right): + return False + for key in left: + left_value = left[key] + right_value = right[key] + if torch.is_tensor(left_value) or torch.is_tensor(right_value): + if not _bits_equal(left_value, right_value): + return False + elif type(left_value) is not type(right_value) or left_value != right_value: + return False + return True + + +def _plain_reference(local_parameter, momentum, second_moment): + parameter = torch.nn.Parameter(local_parameter.clone()) + optimizer = _plain_optimizer(parameter, deterministic=True) + _seed_plain_state(optimizer, parameter, momentum, second_moment) + return optimizer, parameter + + +def _muon_reference(*, normuon): + parameter = torch.nn.Parameter(_muon_initial_parameter().clone()) + optimizer = _muon_optimizer( + parameter, + deterministic=True, + normuon=normuon, + ) + _seed_muon_state(optimizer, parameter, normuon=normuon) + return optimizer, parameter + + +def _run_phase(worker, world_size, *worker_args, timeout=180): + context = mp.get_context("spawn") + result_queue = context.Queue() + descriptor, init_file = tempfile.mkstemp(prefix="gefen-portable-dcp-topology-") + os.close(descriptor) + os.unlink(init_file) + processes = [ + context.Process( + target=worker, + args=(rank, world_size, init_file, *worker_args, result_queue), + ) + for rank in range(world_size) + ] + results = [] + try: + for process in processes: + process.start() + try: + for _ in processes: + results.append(result_queue.get(timeout=timeout)) + except queue_module.Empty: + pass + for process in processes: + process.join(timeout=10) + for process in processes: + if process.is_alive(): + process.terminate() + process.join(timeout=5) + finally: + for process in processes: + if process.is_alive(): + process.terminate() + process.join(timeout=5) + result_queue.close() + result_queue.join_thread() + if os.path.exists(init_file): + os.unlink(init_file) + assert len(results) == world_size, (results, [process.exitcode for process in processes]) + assert all(process.exitcode == 0 for process in processes), [process.exitcode for process in processes] + return sorted(results, key=lambda item: item["rank"]) + + +def _plain_singleton_save_worker( + rank, + world_size, + init_file, + checkpoint_dir, + result_queue, +): + _strict_worker_warnings() + try: + dist.init_process_group( + "gloo", + init_method="file://{}".format(init_file), + rank=rank, + world_size=world_size, + timeout=timedelta(seconds=60), + ) + member = "save:0" + group = ProcessGroupIdentity("plain_singleton_save", (member,)) + identity = ParameterIdentity(_PLAIN_FQN, (2, 4)) + parameter = torch.nn.Parameter(_plain_initial_parameter().clone()) + optimizer = _plain_optimizer(parameter, deterministic=True) + manifest, shards = _replicated_manifest(identity, group) + codebook_binding, checkpoint_binding = _bindings(group, member, None) + optimizer.post_sharding( + (ParameterRebinding(parameter, parameter, shards[0]),), + manifest=manifest, + codebook_process_group=codebook_binding, + ) + _seed_plain_state( + optimizer, + parameter, + _plain_momentum(), + _plain_second_moment(), + ) + save_portable_dcp( + optimizer, + checkpoint_process_group=checkpoint_binding, + storage_writer=FileSystemWriter(checkpoint_dir), + transaction_id="plain-singleton-to-flat-save", + limits=_limits(), + ) + result_queue.put({"rank": rank, "saved": True}) + except BaseException: + result_queue.put({"rank": rank, "fatal_error": traceback.format_exc()}) + finally: + if dist.is_available() and dist.is_initialized(): + dist.destroy_process_group() + + +def _plain_flat_load_worker( + rank, + world_size, + init_file, + checkpoint_dir, + result_queue, +): + _strict_worker_warnings() + try: + dist.init_process_group( + "gloo", + init_method="file://{}".format(init_file), + rank=rank, + world_size=world_size, + timeout=timedelta(seconds=60), + ) + members = tuple("load:{}".format(index) for index in range(world_size)) + group = ProcessGroupIdentity("plain_flat_load", members) + identity = ParameterIdentity(_PLAIN_FQN, (2, 4)) + lengths = (3, 5) + manifest, shards = _flat_manifest(identity, group, lengths) + shard = shards[rank] + start = shard.logical_slice.flat_offset + stop = start + shard.logical_slice.length + original = torch.nn.Parameter(_plain_initial_parameter().clone()) + local_initial = _plain_initial_parameter().reshape(-1)[start:stop].clone() + local = torch.nn.Parameter(local_initial.clone()) + optimizer = _plain_optimizer(original, deterministic=False) + codebook_binding, checkpoint_binding = _bindings( + group, + members[rank], + dist.group.WORLD, + ) + optimizer.post_sharding( + (ParameterRebinding(original, local, shard),), + manifest=manifest, + codebook_process_group=codebook_binding, + ) + load_portable_dcp( + optimizer, + checkpoint_process_group=checkpoint_binding, + storage_reader=FileSystemReader(checkpoint_dir), + transaction_id="plain-singleton-to-flat-load", + limits=_limits(), + ) + + expected_momentum = _plain_momentum().reshape(-1)[start:stop].clone() + expected_second = _plain_second_moment().reshape(-1)[start:stop].clone() + state = optimizer.state[local] + decoded = _decode_quantized_momentum( + optimizer._gefen_codebook, + state["m_codebook"], + state["m_magnitude"], + logical_shape=(lengths[rank],), + period=1, + step=state["step"], + ) + restored_exact = ( + _bits_equal(decoded, expected_momentum) + and _bits_equal(state["vmean"].reshape(-1), expected_second) + and state["vmean_step"] == 6 + and state["step"] == 7 + and optimizer._gefen_global_step == 9 + and optimizer._deterministic is True + ) + + reference, reference_parameter = _plain_reference( + local_initial, + expected_momentum, + expected_second, + ) + gradient = _plain_gradient().reshape(-1)[start:stop].clone() + local.grad = gradient.clone() + reference_parameter.grad = gradient.clone() + optimizer.step() + reference.step() + continuation = { + "parameter": _bits_equal(local, reference_parameter), + "state": _state_equal( + optimizer.state[local], + reference.state[reference_parameter], + ), + "global_step": optimizer._gefen_global_step == reference._gefen_global_step, + "codebook": _bits_equal( + optimizer._gefen_codebook, + reference._gefen_codebook, + ), + } + next_step_exact = all(continuation.values()) + result_queue.put( + { + "rank": rank, + "restored_exact": restored_exact, + "next_step_exact": next_step_exact, + "continuation": continuation, + "length": lengths[rank], + } + ) + except BaseException: + result_queue.put({"rank": rank, "fatal_error": traceback.format_exc()}) + finally: + if dist.is_available() and dist.is_initialized(): + dist.destroy_process_group() + + +def _muon_singleton_save_worker( + rank, + world_size, + init_file, + checkpoint_dir, + normuon, + result_queue, +): + _strict_worker_warnings() + try: + dist.init_process_group( + "gloo", + init_method="file://{}".format(init_file), + rank=rank, + world_size=world_size, + timeout=timedelta(seconds=60), + ) + member = "save:0" + group = ProcessGroupIdentity("muon_singleton_save", (member,)) + identity = ParameterIdentity(_MUON_FQN, (3, 2)) + parameter = torch.nn.Parameter(_muon_initial_parameter().clone()) + optimizer = _muon_optimizer( + parameter, + deterministic=True, + normuon=normuon, + ) + manifest, shards = _replicated_manifest(identity, group) + codebook_binding, checkpoint_binding = _bindings(group, member, None) + optimizer.post_sharding( + (ParameterRebinding(parameter, parameter, shards[0]),), + manifest=manifest, + codebook_process_group=codebook_binding, + ) + _seed_muon_state(optimizer, parameter, normuon=normuon) + save_portable_dcp( + optimizer, + checkpoint_process_group=checkpoint_binding, + storage_writer=FileSystemWriter(checkpoint_dir), + transaction_id="muon-singleton-to-owner-save-{}".format(int(normuon)), + limits=_limits(), + ) + result_queue.put({"rank": rank, "saved": True}) + except BaseException: + result_queue.put({"rank": rank, "fatal_error": traceback.format_exc()}) + finally: + if dist.is_available() and dist.is_initialized(): + dist.destroy_process_group() + + +def _muon_owner_load_worker( + rank, + world_size, + init_file, + checkpoint_dir, + normuon, + result_queue, +): + _strict_worker_warnings() + try: + dist.init_process_group( + "gloo", + init_method="file://{}".format(init_file), + rank=rank, + world_size=world_size, + timeout=timedelta(seconds=60), + ) + members = tuple("load:{}".format(index) for index in range(world_size)) + owner = members[-1] + group = ProcessGroupIdentity("muon_owner_load", members) + identity = ParameterIdentity(_MUON_FQN, (3, 2)) + manifest, shards = _owner_manifest(identity, group, owner) + original = torch.nn.Parameter(_muon_initial_parameter().clone()) + optimizer = _muon_optimizer( + original, + deterministic=False, + normuon=normuon, + ) + local = original if members[rank] == owner else None + codebook_binding, checkpoint_binding = _bindings( + group, + members[rank], + dist.group.WORLD, + ) + optimizer.post_sharding( + (ParameterRebinding(original, local, shards[rank]),), + manifest=manifest, + codebook_process_group=codebook_binding, + ) + load_portable_dcp( + optimizer, + checkpoint_process_group=checkpoint_binding, + storage_reader=FileSystemReader(checkpoint_dir), + transaction_id="muon-singleton-to-owner-load-{}".format(int(normuon)), + limits=_limits(), + ) + + if local is None: + restored_exact = ( + not optimizer.param_groups[0]["params"] + and not optimizer.state + and optimizer._gefen_global_step == 13 + and optimizer._deterministic is True + and _bits_equal(optimizer._gefen_codebook, _codebook()) + ) + else: + state = optimizer.state[local] + decoded = _decode_quantized_momentum( + optimizer._gefen_codebook, + state["m_codebook"], + state["m_magnitude"], + logical_shape=(3, 2), + period=1, + step=state["step"], + ) + restored_exact = ( + _bits_equal(decoded, _muon_momentum()) + and state["step"] == 11 + and optimizer._gefen_global_step == 13 + and optimizer._deterministic is True + and ( + not normuon or (_bits_equal(state["normuon_v"], _muon_normuon_v()) and state["normuon_step"] == 10) + ) + ) + + reference = None + reference_parameter = None + if local is not None: + reference, reference_parameter = _muon_reference(normuon=normuon) + local.grad = _muon_gradient().clone() + reference_parameter.grad = _muon_gradient().clone() + optimizer.step() + if reference is not None: + reference.step() + continuation = { + "parameter": _bits_equal(local, reference_parameter), + "state": _state_equal( + optimizer.state[local], + reference.state[reference_parameter], + ), + "global_step": optimizer._gefen_global_step == reference._gefen_global_step, + "codebook": _bits_equal( + optimizer._gefen_codebook, + reference._gefen_codebook, + ), + } + next_step_exact = all(continuation.values()) + else: + continuation = { + "empty_nonowner": not optimizer.param_groups[0]["params"] and not optimizer.state, + "global_step": optimizer._gefen_global_step == 14, + } + next_step_exact = all(continuation.values()) + result_queue.put( + { + "rank": rank, + "owner": members[rank] == owner, + "restored_exact": restored_exact, + "next_step_exact": next_step_exact, + "continuation": continuation, + } + ) + except BaseException: + result_queue.put({"rank": rank, "fatal_error": traceback.format_exc()}) + finally: + if dist.is_available() and dist.is_initialized(): + dist.destroy_process_group() + + +def _subgroup_worker( + rank, + world_size, + init_file, + checkpoint_dir, + result_queue, +): + _strict_worker_warnings() + try: + dist.init_process_group( + "gloo", + init_method="file://{}".format(init_file), + rank=rank, + world_size=world_size, + timeout=timedelta(seconds=90), + ) + subgroup_ranks = (0, 2) + subgroup = dist.new_group( + ranks=list(subgroup_ranks), + backend="gloo", + timeout=timedelta(seconds=60), + ) + if rank not in subgroup_ranks: + dist.barrier() + result_queue.put({"rank": rank, "excluded": True}) + return + + coordinate = subgroup_ranks.index(rank) + members = tuple("subgroup:{}".format(item) for item in subgroup_ranks) + group = ProcessGroupIdentity("portable_dcp_subgroup", members) + identity = ParameterIdentity(_PLAIN_FQN, (2, 4)) + manifest, shards = _replicated_manifest(identity, group) + source_parameter = torch.nn.Parameter(_plain_initial_parameter().clone()) + source = _plain_optimizer(source_parameter, deterministic=True) + source_codebook, source_checkpoint = _bindings( + group, + members[coordinate], + subgroup, + ) + source.post_sharding( + ( + ParameterRebinding( + source_parameter, + source_parameter, + shards[coordinate], + ), + ), + manifest=manifest, + codebook_process_group=source_codebook, + ) + _seed_plain_state( + source, + source_parameter, + _plain_momentum(), + _plain_second_moment(), + ) + save_portable_dcp( + source, + checkpoint_process_group=source_checkpoint, + storage_writer=FileSystemWriter(checkpoint_dir), + transaction_id="supported-subgroup-save", + limits=_limits(), + ) + + target_parameter = torch.nn.Parameter(_plain_initial_parameter().clone()) + target = _plain_optimizer(target_parameter, deterministic=False) + target_codebook, target_checkpoint = _bindings( + group, + members[coordinate], + subgroup, + ) + target.post_sharding( + ( + ParameterRebinding( + target_parameter, + target_parameter, + shards[coordinate], + ), + ), + manifest=manifest, + codebook_process_group=target_codebook, + ) + load_portable_dcp( + target, + checkpoint_process_group=target_checkpoint, + storage_reader=FileSystemReader(checkpoint_dir), + transaction_id="supported-subgroup-load", + limits=_limits(), + ) + restored_exact = ( + _state_equal(target.state[target_parameter], source.state[source_parameter]) + and target._gefen_global_step == source._gefen_global_step + and target._deterministic is True + and _bits_equal(target._gefen_codebook, source._gefen_codebook) + ) + + reference, reference_parameter = _plain_reference( + _plain_initial_parameter(), + _plain_momentum(), + _plain_second_moment(), + ) + target_parameter.grad = _plain_gradient().clone() + reference_parameter.grad = _plain_gradient().clone() + target.step() + reference.step() + continuation = { + "parameter": _bits_equal(target_parameter, reference_parameter), + "state": _state_equal( + target.state[target_parameter], + reference.state[reference_parameter], + ), + "global_step": target._gefen_global_step == reference._gefen_global_step, + } + next_step_exact = all(continuation.values()) + dist.barrier() + result_queue.put( + { + "rank": rank, + "excluded": False, + "restored_exact": restored_exact, + "next_step_exact": next_step_exact, + "continuation": continuation, + } + ) + except BaseException: + result_queue.put({"rank": rank, "fatal_error": traceback.format_exc()}) + finally: + if dist.is_available() and dist.is_initialized(): + dist.destroy_process_group() + + +@pytest.mark.skipif( + not dist.is_available() or not dist.is_gloo_available(), + reason="portable DCP topology coverage requires Gloo", +) +def test_portable_dcp_replicated_singleton_to_uneven_flattened_world(tmp_path): + checkpoint_dir = str(tmp_path / "plain-singleton-to-flat") + save_results = _run_phase( + _plain_singleton_save_worker, + 1, + checkpoint_dir, + ) + assert save_results == [{"rank": 0, "saved": True}] + + load_results = _run_phase( + _plain_flat_load_worker, + 2, + checkpoint_dir, + ) + assert all("fatal_error" not in result for result in load_results), load_results + assert [result["length"] for result in load_results] == [3, 5] + assert all(result["restored_exact"] and result["next_step_exact"] for result in load_results), load_results + + +@pytest.mark.parametrize("normuon", [False, True]) +@pytest.mark.skipif( + not dist.is_available() or not dist.is_gloo_available(), + reason="portable DCP topology coverage requires Gloo", +) +def test_portable_dcp_muon_replicated_singleton_to_world_owner(tmp_path, normuon): + checkpoint_dir = str(tmp_path / "muon-singleton-to-owner-{}".format(int(normuon))) + save_results = _run_phase( + _muon_singleton_save_worker, + 1, + checkpoint_dir, + normuon, + ) + assert save_results == [{"rank": 0, "saved": True}] + + load_results = _run_phase( + _muon_owner_load_worker, + 2, + checkpoint_dir, + normuon, + ) + assert all("fatal_error" not in result for result in load_results), load_results + assert [result["owner"] for result in load_results] == [False, True] + assert all(result["restored_exact"] and result["next_step_exact"] for result in load_results), load_results + + +@pytest.mark.skipif( + not dist.is_available() or not dist.is_gloo_available(), + reason="portable DCP subgroup coverage requires Gloo", +) +def test_portable_dcp_supported_adapter_defined_subgroup(tmp_path): + results = _run_phase( + _subgroup_worker, + 3, + str(tmp_path / "supported-subgroup"), + timeout=240, + ) + assert all("fatal_error" not in result for result in results), results + assert results[1] == {"rank": 1, "excluded": True} + participants = (results[0], results[2]) + assert all(not result["excluded"] for result in participants) + assert all(result["restored_exact"] and result["next_step_exact"] for result in participants), results diff --git a/tests/test_portable_hybrid.py b/tests/test_portable_hybrid.py new file mode 100644 index 0000000..44466c3 --- /dev/null +++ b/tests/test_portable_hybrid.py @@ -0,0 +1,197 @@ +"""Strict composite-envelope coverage for Gefen-backed Hybrid portable state.""" + +import copy +import io + +import pytest +import torch + +from gefen.portable_hybrid import ( + HYBRID_PORTABLE_STATE_COVERAGE, + HYBRID_PORTABLE_STATE_FORMAT, + HYBRID_PORTABLE_STATE_FORMAT_VERSION, + HYBRID_PORTABLE_STATE_IMPLEMENTATION, + build_hybrid_portable_state_document, + normalize_hybrid_portable_state_document, +) +from gefen.portable_schema import build_portable_state_document + + +def _parameter_record(fqn, tensor): + return { + "identity": { + "schema_version": 1, + "fqn": fqn, + "global_shape": list(tensor.shape), + }, + "algorithm_options": {"period": 1}, + "state_variant": "period_selected", + "state": {"step": 3, "momentum": tensor}, + "projection_hints": {}, + } + + +def _child(role, *, step=7, deterministic=True, fqn=None, tensor=None): + if fqn is None: + fqn = "Model.Weight" if role == "muon" else "Model.Bias" + if tensor is None: + tensor = torch.arange(6, dtype=torch.float32).reshape(2, 3) + implementation = "gefen.GefenMuon" if role == "muon" else "gefen.Gefen" + return build_portable_state_document( + implementation=implementation, + policy={"role": role}, + common={ + "gefen_global_step": step, + "gefen_deterministic": deterministic, + }, + parameters={fqn: _parameter_record(fqn, tensor)}, + provenance={"source_layouts": ["replicated"]}, + ) + + +def _document(*, muon=None, backup=None): + if muon is None: + muon = _child("muon") + if backup is None: + backup = _child("backup") + return build_hybrid_portable_state_document( + backup_optimizer="gefen", + routing={ + "Model.Weight": "muon", + "Model.Bias": "backup", + }, + children={"muon": muon, "backup": backup}, + ) + + +def test_hybrid_portable_document_is_complete_owned_and_weights_only_safe(): + source = torch.arange(30, dtype=torch.float32)[3:15:2].reshape(2, 3) + muon = _child("muon", tensor=source) + document = _document(muon=muon) + + assert document["format"] == HYBRID_PORTABLE_STATE_FORMAT + assert document["format_version"] == HYBRID_PORTABLE_STATE_FORMAT_VERSION + assert document["coverage"] == HYBRID_PORTABLE_STATE_COVERAGE + assert document["implementation"] == HYBRID_PORTABLE_STATE_IMPLEMENTATION + assert document["completion"]["status"] == "complete" + nested = document["children"]["muon"]["parameters"]["Model.Weight"]["state"]["momentum"] + assert nested.device.type == "cpu" + assert nested.is_contiguous() + assert nested.storage_offset() == 0 + assert nested.untyped_storage().nbytes() == nested.numel() * nested.element_size() + assert torch.equal(nested, source) + assert nested is not source + + buffer = io.BytesIO() + torch.save(document, buffer) + buffer.seek(0) + normalized = normalize_hybrid_portable_state_document(torch.load(buffer, weights_only=True)) + assert normalized["completion"] == document["completion"] + + +@pytest.mark.parametrize( + ("key", "value", "match"), + [ + ("format", "other", "format"), + ("format_version", 2, "format_version"), + ("coverage", "partial", "coverage"), + ("implementation", "gefen.Gefen", "implementation"), + ("backup_optimizer", "adamw", "Gefen backup"), + ], +) +def test_hybrid_portable_document_rejects_wrong_envelope_identity(key, value, match): + document = _document() + document[key] = value + + with pytest.raises(ValueError, match=match): + normalize_hybrid_portable_state_document(document) + + +def test_hybrid_portable_document_rejects_outer_and_nested_digest_corruption(): + document = _document() + outer = copy.deepcopy(document) + outer["routing"]["Model.Bias"] = "muon" + with pytest.raises(ValueError, match="routing"): + normalize_hybrid_portable_state_document(outer) + + nested = copy.deepcopy(document) + nested["children"]["muon"]["parameters"]["Model.Weight"]["state"]["step"] += 1 + with pytest.raises(ValueError, match="digest"): + normalize_hybrid_portable_state_document(nested) + + completion = copy.deepcopy(document) + completion["completion"]["digest"] = "0" * 64 + with pytest.raises(ValueError, match="digest"): + normalize_hybrid_portable_state_document(completion) + + +def test_hybrid_portable_document_rejects_schema_and_routing_mismatches(): + document = _document() + extra = copy.deepcopy(document) + extra["extra"] = None + with pytest.raises(ValueError, match="top-level"): + normalize_hybrid_portable_state_document(extra) + + missing_role = copy.deepcopy(document) + del missing_role["children"]["backup"] + with pytest.raises(ValueError, match="children"): + normalize_hybrid_portable_state_document(missing_role) + + missing_fqn = copy.deepcopy(document) + del missing_fqn["routing"]["Model.Bias"] + with pytest.raises(ValueError, match="routing"): + normalize_hybrid_portable_state_document(missing_fqn) + + wrong_child = copy.deepcopy(document) + wrong_child["children"]["muon"] = _child("backup", fqn="Model.Weight") + with pytest.raises(ValueError, match="implementation"): + normalize_hybrid_portable_state_document(wrong_child) + + +@pytest.mark.parametrize( + ("backup", "match"), + [ + (_child("backup", step=8), "global steps"), + (_child("backup", deterministic=False), "deterministic"), + (_child("backup", fqn="Model.Weight"), "disjoint"), + ], +) +def test_hybrid_portable_document_rejects_incompatible_children(backup, match): + with pytest.raises(ValueError, match=match): + build_hybrid_portable_state_document( + backup_optimizer="gefen", + routing={"Model.Weight": "muon", next(iter(backup["parameters"])): "backup"}, + children={"muon": _child("muon"), "backup": backup}, + ) + + +@pytest.mark.parametrize("role", ["muon", "backup"]) +def test_hybrid_portable_document_supports_one_present_child(role): + child = _child(role) + fqn = next(iter(child["parameters"])) + children = {"muon": None, "backup": None} + children[role] = child + + document = build_hybrid_portable_state_document( + backup_optimizer="gefen", + routing={fqn: role}, + children=children, + ) + + assert document["children"][role] is not None + assert document["children"]["backup" if role == "muon" else "muon"] is None + + +def test_hybrid_portable_document_rejects_no_children_and_noncanonical_routing(): + with pytest.raises(ValueError, match="at least one child"): + build_hybrid_portable_state_document( + backup_optimizer="gefen", + routing={}, + children={"muon": None, "backup": None}, + ) + with pytest.raises(ValueError, match="trimmed FQNs"): + build_hybrid_portable_state_document( + backup_optimizer="gefen", + routing={" Model.Weight": "muon"}, + children={"muon": _child("muon"), "backup": None}, + ) diff --git a/tests/test_portable_hybrid_distributed.py b/tests/test_portable_hybrid_distributed.py new file mode 100644 index 0000000..a239593 --- /dev/null +++ b/tests/test_portable_hybrid_distributed.py @@ -0,0 +1,679 @@ +"""Warning-strict distributed topology coverage for portable Hybrid state.""" + +from datetime import timedelta +import multiprocessing as mp +import os +import queue as queue_module +import tempfile +import traceback +import warnings + +import pytest +import torch +import torch.distributed as dist +from torch.distributed.checkpoint import FileSystemReader, FileSystemWriter + +from gefen import GefenMuonHybrid, load_portable_dcp, save_portable_dcp +from gefen.checkpoint import CheckpointProcessGroupBinding +from gefen.codebook import CodebookProcessGroupBinding +from gefen.contracts import ( + LogicalSlice, + ParameterIdentity, + ParameterLayout, + PlacementKind, + ProcessGroupIdentity, + ShardIdentity, + ShardPlacement, + ShardingManifest, +) +from gefen.portable_hybrid import _hybrid_portable_live_token +from gefen.portable_state import PortableStateLimits +from gefen.rebinding import ParameterRebinding + + +_WORLD = 2 +_MUON_FQN = "model.matrix" +_BACKUP_FQN = "model.vector" +_SOURCE_OWNER = "rank:0" +_TARGET_OWNER = "rank:1" +_SOURCE_LENGTHS = (3, 5) +_TARGET_LENGTHS = (5, 3) + + +def _limits(): + return PortableStateLimits( + max_fragment_tensor_bytes=1 << 20, + max_collective_tensor_bytes=4 << 20, + max_collective_metadata_bytes=4 << 20, + chunk_bytes=11, + max_members=4, + max_metadata_bytes=1 << 20, + max_tree_nodes=10_000, + max_tree_depth=32, + max_container_items=10_000, + max_string_bytes=16 << 10, + max_integer_bytes=128, + max_tensors=256, + max_tensor_rank=8, + diagnostic_bytes=1024, + ) + + +def _members(): + return tuple("rank:{}".format(rank) for rank in range(_WORLD)) + + +def _bindings(group, rank): + member = _members()[rank] + return ( + CodebookProcessGroupBinding( + group, + member, + dist.group.WORLD, + torch.device("cpu"), + ), + CheckpointProcessGroupBinding( + group, + member, + dist.group.WORLD, + torch.device("cpu"), + ), + ) + + +def _owner_shards(identity, group, owner): + return tuple( + ShardIdentity( + identity, + ParameterLayout.WHOLE_PARAMETER_OWNER, + (LogicalSlice.full(identity) if member == owner else LogicalSlice(0, 0)), + placements=( + ShardPlacement( + "checkpoint", + PlacementKind.WHOLE_PARAMETER_OWNER, + coordinate, + len(group.ordered_members), + ), + ), + process_group=group, + local_member=member, + owner=owner, + ) + for coordinate, member in enumerate(group.ordered_members) + ) + + +def _flat_shards(identity, group, lengths): + if len(lengths) != len(group.ordered_members) or sum(lengths) != identity.numel: + raise AssertionError("invalid flattened test partition") + offset = 0 + shards = [] + for coordinate, (member, length) in enumerate(zip(group.ordered_members, lengths)): + shards.append( + ShardIdentity( + identity, + ParameterLayout.FLATTENED_ELEMENT_SHARD, + LogicalSlice(offset, length), + placements=( + ShardPlacement( + "checkpoint", + PlacementKind.FLAT_SHARD, + coordinate, + len(group.ordered_members), + ), + ), + process_group=group, + local_member=member, + ) + ) + offset += length + return tuple(shards) + + +def _muon_initial(): + return torch.linspace(-0.3, 0.2, 6, dtype=torch.float32).reshape(3, 2) + + +def _backup_initial(): + return torch.linspace(-0.4, 0.3, 8, dtype=torch.float32) + + +def _muon_momentum(): + return torch.tensor( + [[-0.75, 0.5], [1.25, -1.5], [2.0, -2.5]], + dtype=torch.float32, + ) + + +def _muon_normuon_v(): + return torch.tensor([[0.5], [1.5], [2.5]], dtype=torch.float32) + + +def _backup_momentum(): + return torch.tensor( + [-0.25, 0.5, -0.75, 1.0, -1.25, 1.5, -1.75, 2.0], + dtype=torch.float32, + ) + + +def _backup_second_moment(): + return torch.linspace(0.2, 0.9, 8, dtype=torch.float32) + + +def _muon_gradient(): + return torch.tensor( + [[0.4, -0.7], [1.1, -1.3], [1.7, -1.9]], + dtype=torch.float32, + ) + + +def _backup_gradient(): + return torch.tensor( + [0.25, -0.5, 0.75, -1.0, 1.25, -1.5, 1.75, -2.0], + dtype=torch.float32, + ) + + +def _codebook(): + return torch.linspace(-1.0, 1.0, 256, dtype=torch.float32) + + +def _quantized_period_one(momentum): + flat = momentum.reshape(-1) + indices = torch.where( + torch.signbit(flat), + torch.zeros(flat.numel(), dtype=torch.uint8), + torch.full((flat.numel(),), 255, dtype=torch.uint8), + ) + return indices.reshape(-1, 1), flat.abs().reshape(-1, 1).clone() + + +def _make_hybrid(rank, group, *, owner, lengths, deterministic): + muon_identity = ParameterIdentity(_MUON_FQN, (3, 2)) + backup_identity = ParameterIdentity(_BACKUP_FQN, (8,)) + old_muon = torch.nn.Parameter(_muon_initial().clone()) + old_backup = torch.nn.Parameter(_backup_initial().clone()) + optimizer = GefenMuonHybrid( + [("matrix", old_muon)], + [("vector", old_backup)], + lr=2.5e-3, + muon_lr=3.0e-3, + backup_lr=2.5e-3, + weight_decay=0.03, + muon_weight_decay=0.02, + backup_weight_decay=0.03, + backup_optimizer="gefen", + backup_1d_period_one=True, + betas=(0.8, 0.97), + eps=2.0e-8, + fused=False, + momentum=0.85, + nesterov=False, + ns_steps=2, + sharded_mode="distributed", + deterministic=deterministic, + normuon=True, + normuon_beta2=0.9, + normuon_eps=3.0e-8, + ) + muon_shards = _owner_shards(muon_identity, group, owner) + backup_shards = _flat_shards(backup_identity, group, lengths) + muon_local = old_muon if _members()[rank] == owner else None + backup_shard = backup_shards[rank] + start = backup_shard.logical_slice.flat_offset + stop = start + backup_shard.logical_slice.length + backup_local = torch.nn.Parameter(_backup_initial()[start:stop].clone()) + codebook_binding, checkpoint_binding = _bindings(group, rank) + optimizer.post_sharding( + ( + ParameterRebinding(old_muon, muon_local, muon_shards[rank]), + ParameterRebinding(old_backup, backup_local, backup_shard), + ), + manifest=ShardingManifest(muon_shards + backup_shards), + codebook_process_group=codebook_binding, + ) + return ( + optimizer, + muon_local, + backup_local, + backup_shard, + checkpoint_binding, + ) + + +def _seed_hybrid_state(optimizer, muon_parameter, backup_parameter, backup_shard): + for child in (optimizer.muon, optimizer.backup): + child._gefen_global_step = 13 + child._gefen_codebook = _codebook() + if muon_parameter is not None: + indices, magnitudes = _quantized_period_one(_muon_momentum()) + optimizer.muon.state[muon_parameter].update( + { + "automatic_period": 1, + "step": 11, + "m_codebook": indices, + "m_magnitude": magnitudes, + "normuon_v": _muon_normuon_v(), + "normuon_step": 10, + } + ) + start = backup_shard.logical_slice.flat_offset + stop = start + backup_shard.logical_slice.length + momentum = _backup_momentum()[start:stop] + indices, magnitudes = _quantized_period_one(momentum) + optimizer.backup.state[backup_parameter].update( + { + "automatic_period": 1, + "step": 11, + "m_codebook": indices, + "m_magnitude": magnitudes, + "vmean": _backup_second_moment()[start:stop].reshape(-1, 1).clone(), + "vmean_step": 10, + } + ) + + +def _bits_equal(left, right): + return ( + torch.is_tensor(left) + and torch.is_tensor(right) + and left.dtype == right.dtype + and tuple(left.shape) == tuple(right.shape) + and torch.equal( + left.detach().contiguous().view(torch.uint8), + right.detach().contiguous().view(torch.uint8), + ) + ) + + +def _state_equal(left, right): + if set(left) != set(right): + return False + for key in left: + left_value = left[key] + right_value = right[key] + if torch.is_tensor(left_value) or torch.is_tensor(right_value): + if not _bits_equal(left_value, right_value): + return False + elif type(left_value) is not type(right_value) or left_value != right_value: + return False + return True + + +def _local_child_exact(left, left_parameter, right, right_parameter): + if left_parameter is None or right_parameter is None: + return ( + left_parameter is right_parameter is None + and not left.param_groups[0]["params"] + and not right.param_groups[0]["params"] + and not left.state + and not right.state + ) + return _state_equal( + left.state[left_parameter], + right.state[right_parameter], + ) + + +def _hybrids_exact( + left, + left_muon, + left_backup, + right, + right_muon, + right_backup, +): + return ( + left._canonical_identity_ready() + and right._canonical_identity_ready() + and left._hybrid_fqn_roles == right._hybrid_fqn_roles == ((_MUON_FQN, "muon"), (_BACKUP_FQN, "backup")) + and left._deterministic is right._deterministic is True + and left.muon._deterministic is right.muon._deterministic is True + and left.backup._deterministic is right.backup._deterministic is True + and left.muon._gefen_global_step == right.muon._gefen_global_step + and left.backup._gefen_global_step == right.backup._gefen_global_step + and _bits_equal( + left.muon._gefen_codebook, + right.muon._gefen_codebook, + ) + and _bits_equal( + left.backup._gefen_codebook, + right.backup._gefen_codebook, + ) + and _local_child_exact( + left.muon, + left_muon, + right.muon, + right_muon, + ) + and _local_child_exact( + left.backup, + left_backup, + right.backup, + right_backup, + ) + ) + + +def _step_against_reference( + target, + target_muon, + target_backup, + target_backup_shard, + reference, + reference_muon, + reference_backup, +): + if target_muon is not None: + target_muon.grad = _muon_gradient().clone() + reference_muon.grad = _muon_gradient().clone() + start = target_backup_shard.logical_slice.flat_offset + stop = start + target_backup_shard.logical_slice.length + gradient = _backup_gradient()[start:stop] + target_backup.grad = gradient.clone() + reference_backup.grad = gradient.clone() + target.step() + reference.step() + parameters_exact = ( + (target_muon is reference_muon is None) or _bits_equal(target_muon, reference_muon) + ) and _bits_equal(target_backup, reference_backup) + return parameters_exact and _hybrids_exact( + target, + target_muon, + target_backup, + reference, + reference_muon, + reference_backup, + ) + + +def _fresh_target_and_reference(rank, group): + target = _make_hybrid( + rank, + group, + owner=_TARGET_OWNER, + lengths=_TARGET_LENGTHS, + deterministic=False, + ) + reference = _make_hybrid( + rank, + group, + owner=_TARGET_OWNER, + lengths=_TARGET_LENGTHS, + deterministic=True, + ) + _seed_hybrid_state( + reference[0], + reference[1], + reference[2], + reference[3], + ) + return target, reference + + +def _successful_import_result(rank, group, document): + target, reference = _fresh_target_and_reference(rank, group) + target[0].import_portable_state( + document, + checkpoint_process_group=target[4], + transaction_id="hybrid-direct-import-v1", + limits=_limits(), + ) + restored_exact = _hybrids_exact( + target[0], + target[1], + target[2], + reference[0], + reference[1], + reference[2], + ) + next_step_exact = _step_against_reference( + target[0], + target[1], + target[2], + target[3], + reference[0], + reference[1], + reference[2], + ) + return { + "restored_exact": restored_exact, + "next_step_exact": next_step_exact, + "target_owns_muon": target[1] is not None, + "target_backup_length": target[3].logical_slice.length, + } + + +def _asymmetric_failure_result(rank, group, document): + target = _make_hybrid( + rank, + group, + owner=_TARGET_OWNER, + lengths=_TARGET_LENGTHS, + deterministic=False, + ) + optimizer, muon_parameter, backup_parameter, _shard, binding = target + if rank == 0: + optimizer.backup.param_groups[0]["rank_local_extension"] = "reject-me" + token_before = _hybrid_portable_live_token(optimizer) + child_state_objects = (optimizer.muon.state, optimizer.backup.state) + parameter_values = ( + None if muon_parameter is None else muon_parameter.detach().clone(), + backup_parameter.detach().clone(), + ) + try: + optimizer.import_portable_state( + document, + checkpoint_process_group=binding, + transaction_id="hybrid-asymmetric-import-failure-v1", + limits=_limits(), + ) + except RuntimeError as exc: + message = str(exc) + else: + message = None + parameters_unchanged = ( + (muon_parameter is None and parameter_values[0] is None) or _bits_equal(muon_parameter, parameter_values[0]) + ) and _bits_equal(backup_parameter, parameter_values[1]) + return { + "message": message, + "token_unchanged": _hybrid_portable_live_token(optimizer) == token_before, + "state_objects_unchanged": ( + optimizer.muon.state is child_state_objects[0] and optimizer.backup.state is child_state_objects[1] + ), + "parameters_unchanged": parameters_unchanged, + "common_unchanged": ( + optimizer._deterministic is False + and optimizer.muon._gefen_global_step == 0 + and optimizer.backup._gefen_global_step == 0 + and optimizer.muon._gefen_codebook is None + and optimizer.backup._gefen_codebook is None + ), + } + + +def _dcp_result(rank, group, checkpoint_dir): + target, reference = _fresh_target_and_reference(rank, group) + load_portable_dcp( + target[0], + checkpoint_process_group=target[4], + storage_reader=FileSystemReader(checkpoint_dir), + transaction_id="hybrid-dcp-load-v1", + limits=_limits(), + ) + restored_exact = _hybrids_exact( + target[0], + target[1], + target[2], + reference[0], + reference[1], + reference[2], + ) + next_step_exact = _step_against_reference( + target[0], + target[1], + target[2], + target[3], + reference[0], + reference[1], + reference[2], + ) + return { + "restored_exact": restored_exact, + "next_step_exact": next_step_exact, + } + + +def _strict_worker_warnings(): + warnings.simplefilter("error") + warnings.filterwarnings( + "ignore", + message="TypedStorage is deprecated.*", + category=UserWarning, + ) + + +def _distributed_worker(rank, init_file, checkpoint_dir, result_queue): + _strict_worker_warnings() + try: + dist.init_process_group( + "gloo", + init_method="file://{}".format(init_file), + rank=rank, + world_size=_WORLD, + timeout=timedelta(seconds=120), + ) + group = ProcessGroupIdentity("hybrid_portable_checkpoint", _members()) + source = _make_hybrid( + rank, + group, + owner=_SOURCE_OWNER, + lengths=_SOURCE_LENGTHS, + deterministic=True, + ) + _seed_hybrid_state(source[0], source[1], source[2], source[3]) + document = source[0].export_portable_state( + checkpoint_process_group=source[4], + transaction_id="hybrid-direct-export-v1", + limits=_limits(), + ) + direct = _successful_import_result(rank, group, document) + dist.barrier() + asymmetric = _asymmetric_failure_result(rank, group, document) + dist.barrier() + save_portable_dcp( + source[0], + checkpoint_process_group=source[4], + storage_writer=FileSystemWriter(checkpoint_dir), + transaction_id="hybrid-dcp-save-v1", + limits=_limits(), + ) + dist.barrier() + dcp = _dcp_result(rank, group, checkpoint_dir) + dist.barrier() + result_queue.put( + { + "rank": rank, + "digest": document["completion"]["digest"], + "document_exact": ( + document["routing"] == {_MUON_FQN: "muon", _BACKUP_FQN: "backup"} + and all( + document["children"][role]["common"]["gefen_global_step"] == 13 + and document["children"][role]["common"]["gefen_deterministic"] is True + and _bits_equal( + document["children"][role]["common"]["gefen_codebook"], + _codebook(), + ) + for role in ("muon", "backup") + ) + ), + "direct": direct, + "asymmetric": asymmetric, + "dcp": dcp, + } + ) + except BaseException: + result_queue.put({"rank": rank, "fatal_error": traceback.format_exc()}) + finally: + if dist.is_available() and dist.is_initialized(): + dist.destroy_process_group() + + +def _run_distributed_workers(checkpoint_dir): + context = mp.get_context("spawn") + result_queue = context.Queue() + descriptor, init_file = tempfile.mkstemp(prefix="gefen-portable-hybrid-") + os.close(descriptor) + os.unlink(init_file) + processes = [ + context.Process( + target=_distributed_worker, + args=(rank, init_file, checkpoint_dir, result_queue), + ) + for rank in range(_WORLD) + ] + results = [] + try: + for process in processes: + process.start() + try: + for _ in processes: + results.append(result_queue.get(timeout=240)) + except queue_module.Empty: + pass + for process in processes: + process.join(timeout=10) + for process in processes: + if process.is_alive(): + process.terminate() + process.join(timeout=5) + finally: + for process in processes: + if process.is_alive(): + process.terminate() + process.join(timeout=5) + result_queue.close() + result_queue.join_thread() + if os.path.exists(init_file): + os.unlink(init_file) + assert len(results) == _WORLD, (results, [process.exitcode for process in processes]) + assert all(process.exitcode == 0 for process in processes), [process.exitcode for process in processes] + return sorted(results, key=lambda item: item["rank"]) + + +@pytest.mark.skipif( + not dist.is_available() or not dist.is_gloo_available(), + reason="portable Hybrid topology coverage requires Gloo", +) +def test_two_process_portable_hybrid_topology_change_and_dcp(tmp_path): + results = _run_distributed_workers(str(tmp_path / "hybrid-portable-dcp")) + + assert all("fatal_error" not in result for result in results), results + assert len({result["digest"] for result in results}) == 1 + assert all(result["document_exact"] for result in results) + assert [result["direct"]["target_owns_muon"] for result in results] == [ + False, + True, + ] + assert [result["direct"]["target_backup_length"] for result in results] == [ + 5, + 3, + ] + assert all( + result["direct"]["restored_exact"] + and result["direct"]["next_step_exact"] + and result["dcp"]["restored_exact"] + and result["dcp"]["next_step_exact"] + for result in results + ), results + + messages = [result["asymmetric"]["message"] for result in results] + assert messages[0] == messages[1] + assert messages[0] is not None and "unknown keys" in messages[0] + assert all( + result["asymmetric"]["token_unchanged"] + and result["asymmetric"]["state_objects_unchanged"] + and result["asymmetric"]["parameters_unchanged"] + and result["asymmetric"]["common_unchanged"] + for result in results + ), results diff --git a/tests/test_portable_hybrid_runtime.py b/tests/test_portable_hybrid_runtime.py new file mode 100644 index 0000000..f9e8a16 --- /dev/null +++ b/tests/test_portable_hybrid_runtime.py @@ -0,0 +1,391 @@ +"""Warning-strict runtime coverage for Gefen-backed Hybrid portable state.""" + +import copy +import warnings + +import pytest +import torch +import torch.distributed.checkpoint as dcp + +from gefen import ( + CheckpointProcessGroupBinding, + CheckpointTransport, + GefenMuonHybrid, + PortableStateLimits, + PortableStateProvider, + load_portable_dcp, + save_portable_dcp, +) +from gefen.codebook import CodebookProcessGroupBinding +from gefen.contracts import ( + LogicalSlice, + ParameterIdentity, + ParameterLayout, + PlacementKind, + ProcessGroupIdentity, + ShardIdentity, + ShardPlacement, + ShardingManifest, +) +from gefen.rebinding import ParameterRebinding +from gefen.portable_schema import portable_state_digest + + +_MEMBER = "rank:0" +_SINGLE_PROCESS_DCP_WARNING = ( + r"^torch\.distributed is (?:disabled, )?unavailable or uninitialized, " + r"assuming the intent is to (?:save|load) in a single process\.$" +) +_TYPED_STORAGE_DCP_WARNING = ( + r"^TypedStorage is deprecated\. It will be removed in the future and UntypedStorage will be the only storage class\. " + r"This should only matter to you if you are using storages directly\. To access UntypedStorage directly, use " + r"tensor\.untyped_storage\(\) instead of tensor\.storage\(\)$" +) + + +@pytest.fixture(autouse=True) +def _warning_strict(): + with warnings.catch_warnings(): + warnings.simplefilter("error") + warnings.filterwarnings( + "ignore", + message=_SINGLE_PROCESS_DCP_WARNING, + category=UserWarning, + ) + warnings.filterwarnings( + "ignore", + message=_TYPED_STORAGE_DCP_WARNING, + category=UserWarning, + ) + yield + + +def _limits(): + return PortableStateLimits( + max_fragment_tensor_bytes=4 << 20, + max_collective_tensor_bytes=16 << 20, + max_collective_metadata_bytes=64 << 20, + ) + + +def _grouped_replicated(fqn, shape, group): + identity = ParameterIdentity(fqn, shape) + return ShardIdentity( + identity, + ParameterLayout.REPLICATED, + LogicalSlice.full(identity), + placements=( + ShardPlacement( + "checkpoint", + PlacementKind.REPLICATE, + 0, + 1, + ), + ), + process_group=group, + local_member=_MEMBER, + ) + + +def _make_hybrid(matrix, bias, *, deterministic): + optimizer = GefenMuonHybrid( + [("layer.weight", matrix)], + [("layer.bias", bias)], + lr=1e-3, + fused=False, + sharded_mode="distributed", + backup_optimizer="gefen", + backup_1d_period_one=True, + normuon=True, + deterministic=deterministic, + ) + group = ProcessGroupIdentity("hybrid_checkpoint", (_MEMBER,)) + matrix_shard = _grouped_replicated( + "Model.Layer.Weight", + tuple(matrix.shape), + group, + ) + bias_shard = _grouped_replicated( + "Model.Layer.Bias", + tuple(bias.shape), + group, + ) + manifest = ShardingManifest((matrix_shard, bias_shard)) + codebook_binding = CodebookProcessGroupBinding( + group, + _MEMBER, + None, + torch.device("cpu"), + ) + checkpoint_binding = CheckpointProcessGroupBinding( + group, + _MEMBER, + None, + torch.device("cpu"), + ) + optimizer.post_sharding( + ( + ParameterRebinding(matrix, matrix, matrix_shard), + ParameterRebinding(bias, bias, bias_shard), + ), + manifest=manifest, + codebook_process_group=codebook_binding, + ) + return optimizer, checkpoint_binding + + +def _set_grads(matrix, bias, offset): + matrix.grad = ( + torch.tensor( + [[0.25, -0.5], [0.75, -1.0]], + dtype=matrix.dtype, + ) + + offset + ) + bias.grad = torch.tensor([0.2, -0.4, 0.6, -0.8], dtype=bias.dtype) - offset + + +def _equal(left, right): + if torch.is_tensor(left) or torch.is_tensor(right): + return ( + torch.is_tensor(left) + and torch.is_tensor(right) + and left.dtype == right.dtype + and tuple(left.shape) == tuple(right.shape) + and torch.equal(left.cpu(), right.cpu()) + ) + if type(left) is not type(right): + return False + if type(left) is dict: + return set(left) == set(right) and all(_equal(left[key], right[key]) for key in left) + if type(left) in {list, tuple}: + return len(left) == len(right) and all(_equal(a, b) for a, b in zip(left, right)) + return left == right + + +def _redigest(document): + payload = {key: document[key] for key in document if key != "completion"} + document["completion"]["digest"] = portable_state_digest(payload) + + +def _assert_same_optimizer_state(left, right): + assert _equal(left.muon.state_dict(), right.muon.state_dict()) + assert _equal(left.backup.state_dict(), right.backup.state_dict()) + assert left._deterministic is right._deterministic + + +def _initialized_source(): + matrix = torch.nn.Parameter(torch.tensor([[1.0, -2.0], [3.0, -4.0]])) + bias = torch.nn.Parameter(torch.tensor([0.5, -1.5, 2.5, -3.5])) + optimizer, binding = _make_hybrid(matrix, bias, deterministic=True) + _set_grads(matrix, bias, 0.0) + optimizer.step() + return optimizer, matrix, bias, binding + + +def _target_from_source(matrix, bias, *, deterministic=False): + target_matrix = torch.nn.Parameter(matrix.detach().clone()) + target_bias = torch.nn.Parameter(bias.detach().clone()) + optimizer, binding = _make_hybrid( + target_matrix, + target_bias, + deterministic=deterministic, + ) + return optimizer, target_matrix, target_bias, binding + + +def test_hybrid_portable_round_trip_and_exact_next_step(): + source, source_matrix, source_bias, source_binding = _initialized_source() + document = source.export_portable_state( + checkpoint_process_group=source_binding, + transaction_id="hybrid-export-v1", + limits=_limits(), + ) + target, target_matrix, target_bias, target_binding = _target_from_source( + source_matrix, + source_bias, + ) + + target.import_portable_state( + document, + checkpoint_process_group=target_binding, + transaction_id="hybrid-import-v1", + limits=_limits(), + ) + + assert isinstance(target, PortableStateProvider) + assert document["routing"] == { + "Model.Layer.Bias": "backup", + "Model.Layer.Weight": "muon", + } + assert document["children"]["muon"]["implementation"] == "gefen.GefenMuon" + assert document["children"]["backup"]["implementation"] == "gefen.Gefen" + assert target._deterministic is True + _assert_same_optimizer_state(source, target) + support = tuple( + item + for item in target.optimizer_contract().capabilities.checkpoints + if item.transport is CheckpointTransport.CANONICAL_GLOBAL + ) + assert len(support) == 1 + assert support[0].atomic_load + + _set_grads(source_matrix, source_bias, 0.125) + _set_grads(target_matrix, target_bias, 0.125) + source.step() + target.step() + + assert torch.equal(source_matrix, target_matrix) + assert torch.equal(source_bias, target_bias) + _assert_same_optimizer_state(source, target) + + +def test_hybrid_portable_import_stages_every_child_before_mutation(): + source, source_matrix, source_bias, source_binding = _initialized_source() + document = source.export_portable_state( + checkpoint_process_group=source_binding, + transaction_id="hybrid-atomic-export-v1", + limits=_limits(), + ) + target, _target_matrix, _target_bias, target_binding = _target_from_source( + source_matrix, + source_bias, + ) + target.backup.param_groups[0]["eps"] = 1e-7 + muon_before = copy.deepcopy(target.muon.state_dict()) + backup_before = copy.deepcopy(target.backup.state_dict()) + muon_state_object = target.muon.state + backup_state_object = target.backup.state + + with pytest.raises((RuntimeError, ValueError), match="policy|option|portable"): + target.import_portable_state( + document, + checkpoint_process_group=target_binding, + transaction_id="hybrid-atomic-import-v1", + limits=_limits(), + ) + + assert target.muon.state is muon_state_object + assert target.backup.state is backup_state_object + assert _equal(target.muon.state_dict(), muon_before) + assert _equal(target.backup.state_dict(), backup_before) + + +def test_hybrid_portable_import_rejects_semantically_noncanonical_child(): + source, source_matrix, source_bias, source_binding = _initialized_source() + document = source.export_portable_state( + checkpoint_process_group=source_binding, + transaction_id="hybrid-semantic-export-v1", + limits=_limits(), + ) + document["children"]["backup"]["provenance"] = {} + _redigest(document["children"]["backup"]) + _redigest(document) + target, _target_matrix, _target_bias, target_binding = _target_from_source( + source_matrix, + source_bias, + ) + muon_before = copy.deepcopy(target.muon.state_dict()) + backup_before = copy.deepcopy(target.backup.state_dict()) + + with pytest.raises(RuntimeError, match="provenance"): + target.import_portable_state( + document, + checkpoint_process_group=target_binding, + transaction_id="hybrid-semantic-import-v1", + limits=_limits(), + ) + + assert _equal(target.muon.state_dict(), muon_before) + assert _equal(target.backup.state_dict(), backup_before) + + +def test_hybrid_portable_composite_uses_collective_tensor_limit(): + source, source_matrix, source_bias, source_binding = _initialized_source() + limits = PortableStateLimits( + max_fragment_tensor_bytes=1536, + max_collective_tensor_bytes=4096, + max_collective_metadata_bytes=64 << 20, + ) + + document = source.export_portable_state( + checkpoint_process_group=source_binding, + transaction_id="hybrid-aggregate-export-v1", + limits=limits, + ) + target, _target_matrix, _target_bias, target_binding = _target_from_source( + source_matrix, + source_bias, + ) + target.import_portable_state( + document, + checkpoint_process_group=target_binding, + transaction_id="hybrid-aggregate-import-v1", + limits=limits, + ) + + _assert_same_optimizer_state(source, target) + + +def test_hybrid_portable_dcp_round_trip(tmp_path): + source, source_matrix, source_bias, source_binding = _initialized_source() + checkpoint = tmp_path / "hybrid-portable" + save_portable_dcp( + source, + checkpoint_process_group=source_binding, + storage_writer=dcp.FileSystemWriter(checkpoint), + transaction_id="hybrid-dcp-save-v1", + limits=_limits(), + ) + target, target_matrix, target_bias, target_binding = _target_from_source( + source_matrix, + source_bias, + ) + + load_portable_dcp( + target, + checkpoint_process_group=target_binding, + storage_reader=dcp.FileSystemReader(checkpoint), + transaction_id="hybrid-dcp-load-v1", + limits=_limits(), + ) + + _assert_same_optimizer_state(source, target) + _set_grads(source_matrix, source_bias, -0.0625) + _set_grads(target_matrix, target_bias, -0.0625) + source.step() + target.step() + assert torch.equal(source_matrix, target_matrix) + assert torch.equal(source_bias, target_bias) + _assert_same_optimizer_state(source, target) + + +def test_adamw_backed_hybrid_remains_explicitly_nonportable(): + matrix = torch.nn.Parameter(torch.ones(2, 2)) + bias = torch.nn.Parameter(torch.ones(4)) + optimizer = GefenMuonHybrid( + [("layer.weight", matrix)], + [("layer.bias", bias)], + lr=1e-3, + fused=False, + sharded_mode="distributed", + backup_optimizer="adamw", + ) + assert not optimizer.optimizer_contract().capabilities.canonical_state_io + assert all( + item.transport is not CheckpointTransport.CANONICAL_GLOBAL + for item in optimizer.optimizer_contract().capabilities.checkpoints + ) + group = ProcessGroupIdentity("hybrid_checkpoint", (_MEMBER,)) + binding = CheckpointProcessGroupBinding( + group, + _MEMBER, + None, + torch.device("cpu"), + ) + with pytest.raises(RuntimeError, match="AdamW"): + optimizer.export_portable_state( + checkpoint_process_group=binding, + transaction_id="adamw-hybrid-reject-v1", + limits=_limits(), + ) diff --git a/tests/test_portable_schema.py b/tests/test_portable_schema.py index 535d109..ebd746b 100644 --- a/tests/test_portable_schema.py +++ b/tests/test_portable_schema.py @@ -7,6 +7,7 @@ import torch import gefen +import gefen.contracts as contracts_module import gefen.portable_schema as portable_schema_module from gefen.portable_schema import ( PORTABLE_STATE_COVERAGE, @@ -57,15 +58,15 @@ def test_portable_document_is_complete_device_neutral_and_weights_only_safe(): assert document["format_version"] == PORTABLE_STATE_FORMAT_VERSION assert document["coverage"] == PORTABLE_STATE_COVERAGE assert document["completion"]["status"] == "complete" - assert ( - document["completion"]["digest_algorithm"] - == PORTABLE_STATE_DIGEST_ALGORITHM - ) + assert document["completion"]["digest_algorithm"] == PORTABLE_STATE_DIGEST_ALGORITHM momentum = document["parameters"]["Model.Weight"]["state"]["momentum"] assert momentum.device.type == "cpu" assert momentum.is_contiguous() assert momentum.storage_offset() == 0 - assert momentum.untyped_storage().nbytes() == momentum.numel() * momentum.element_size() + assert ( + momentum.untyped_storage().nbytes() + == momentum.numel() * momentum.element_size() + ) assert torch.equal(momentum, view) assert momentum is not view @@ -83,10 +84,18 @@ def test_portable_digest_is_deterministic_and_type_shape_dtype_sensitive(): baseline = portable_state_digest(left) assert baseline == portable_state_digest(right) - assert baseline != portable_state_digest({"a": torch.tensor([[1.0, 2.0]]), "b": [1, 2]}) - assert baseline != portable_state_digest({"a": torch.tensor([1.0, 2.0], dtype=torch.float64), "b": [1, 2]}) - assert baseline != portable_state_digest({"a": torch.tensor([1.0, 3.0]), "b": [1, 2]}) - assert baseline != portable_state_digest({"a": torch.tensor([1.0, 2.0]), "b": (1, 2)}) + assert baseline != portable_state_digest( + {"a": torch.tensor([[1.0, 2.0]]), "b": [1, 2]} + ) + assert baseline != portable_state_digest( + {"a": torch.tensor([1.0, 2.0], dtype=torch.float64), "b": [1, 2]} + ) + assert baseline != portable_state_digest( + {"a": torch.tensor([1.0, 3.0]), "b": [1, 2]} + ) + assert baseline != portable_state_digest( + {"a": torch.tensor([1.0, 2.0]), "b": (1, 2)} + ) def test_portable_digest_streams_tensor_bytes_without_changing_the_digest(monkeypatch): @@ -113,8 +122,7 @@ def test_portable_digest_v3_grammar_has_a_cross_version_golden_vector(): } assert portable_state_digest(payload) == ( - "348e09a77f1b3eae286adda1573722df9" - "38be6d8100631b272665f8c22c6e23a" + "348e09a77f1b3eae286adda1573722df938be6d8100631b272665f8c22c6e23a" ) @@ -149,7 +157,9 @@ def tracked(value): calls.append(value) return original(value) - monkeypatch.setattr(portable_schema_module, "_canonical_portable_state_digest", tracked) + monkeypatch.setattr( + portable_schema_module, "_canonical_portable_state_digest", tracked + ) document = _document() assert len(calls) == 1 @@ -272,8 +282,14 @@ def test_builder_and_normalizer_do_not_alias_or_mutate_inputs(): assert document is not normalized assert document["parameters"] is not parameters assert document["parameters"]["Model.Weight"] is not record - assert document["parameters"]["Model.Weight"]["state"]["momentum"] is not source_momentum - assert normalized["parameters"]["Model.Weight"]["state"]["momentum"] is not document["parameters"]["Model.Weight"]["state"]["momentum"] + assert ( + document["parameters"]["Model.Weight"]["state"]["momentum"] + is not source_momentum + ) + assert ( + normalized["parameters"]["Model.Weight"]["state"]["momentum"] + is not document["parameters"]["Model.Weight"]["state"]["momentum"] + ) def test_portable_schema_and_checkpoint_binding_exports_are_public(): @@ -284,6 +300,7 @@ def test_portable_schema_and_checkpoint_binding_exports_are_public(): assert gefen.CheckpointProcessGroupBinding.__module__ == "gefen.checkpoint" assert gefen.PortableStateLimits.__module__ == "gefen.portable_state" assert gefen.PortableStateProvider.__module__ == "gefen.contracts" + assert "PortableStateProvider" in contracts_module.__all__ assert gefen.load_portable_dcp.__module__ == "gefen.portable_dcp" assert gefen.save_portable_dcp.__module__ == "gefen.portable_dcp" assert gefen.LogicalRegion.__module__ == "gefen.contracts" diff --git a/tests/test_rebinding_cpu.py b/tests/test_rebinding_cpu.py index b852a69..0fd1d2e 100644 --- a/tests/test_rebinding_cpu.py +++ b/tests/test_rebinding_cpu.py @@ -1012,7 +1012,7 @@ def test_muon_flattened_rebinding_rejects_without_mutation(): _assert_snapshot(optimizer, snapshot) -def test_hybrid_contract_does_not_claim_composite_rebinding(): +def test_gefen_backed_hybrid_contract_claims_composite_rebinding(): matrix = torch.nn.Parameter(torch.ones(4, 4)) bias = torch.nn.Parameter(torch.ones(4)) optimizer = GefenMuonHybrid( @@ -1022,6 +1022,9 @@ def test_hybrid_contract_does_not_claim_composite_rebinding(): fused=False, ) contract = optimizer.optimizer_contract() - assert not contract.capabilities.shard_rebinding - assert not contract.capabilities.post_sharding + assert contract.capabilities.shard_rebinding + assert contract.capabilities.post_sharding + assert contract.capabilities.explicit_process_group_codebook_scope assert not contract.capabilities.canonical_parameter_fqns + assert not contract.capabilities.stable_shard_identity + assert optimizer.state_dict()["backup_optimizer"] == "gefen" diff --git a/tests/test_state_offload.py b/tests/test_state_offload.py new file mode 100644 index 0000000..3a0a6bb --- /dev/null +++ b/tests/test_state_offload.py @@ -0,0 +1,602 @@ +"""Focused CPU/CUDA coverage for native plain-Gefen parameter-state offload.""" + +import copy +import io +import types + +import pytest +import torch + +from gefen import ( + CheckpointProcessGroupBinding, + CheckpointTransport, + Gefen, + GefenMuon, + PortableStateLimits, + StateOffloadProvider, +) +from gefen.codebook import CodebookProcessGroupBinding +from gefen.contracts import ( + LogicalSlice, + ParameterIdentity, + ParameterLayout, + PlacementKind, + ProcessGroupIdentity, + ShardIdentity, + ShardPlacement, + ShardingManifest, +) +from gefen.rebinding import ParameterRebinding + + +_PERSISTENT_TENSOR_KEYS = frozenset({"m_codebook", "m_magnitude", "vmean", "v_row", "v_col"}) +_CUDA_REQUIRED = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") + + +def _make_gefen(parameter, *, factored=False, fused=False): + optimizer = Gefen( + [("layer.weight", parameter)], + lr=2.0e-3, + fused=fused, + factored_v_2d=factored, + deterministic=True, + ) + optimizer._resolve_automatic_period = lambda *args, **kwargs: 4 + return optimizer + + +def _persistent_snapshot(optimizer, parameter): + result = {} + for key, value in optimizer.state[parameter].items(): + if key in _PERSISTENT_TENSOR_KEYS: + result[key] = value.detach().cpu().clone() + elif key in {"automatic_period", "step", "vmean_step", "factored_step"}: + result[key] = value + return result + + +def _assert_persistent_equal(left, right): + assert set(left) == set(right) + for key in left: + if torch.is_tensor(left[key]): + torch.testing.assert_close(left[key], right[key], rtol=0, atol=0) + else: + assert left[key] == right[key] + + +def _assert_cpu_boundary(optimizer): + assert optimizer.state_offload_active + assert optimizer.state_offload_device == torch.device("cpu") + assert not optimizer.state_offload_poisoned + for parameter_state in optimizer.state.values(): + assert "stepsize" not in parameter_state + assert "_h_buf" not in parameter_state + for key, value in parameter_state.items(): + if key in _PERSISTENT_TENSOR_KEYS: + assert type(value) is torch.Tensor + assert value.device.type == "cpu" + assert value.is_contiguous() + assert value.storage_offset() == 0 + assert value.untyped_storage().nbytes() == value.numel() * value.element_size() + + +def test_state_offload_visibility_is_read_only_and_cpu_parameters_fail_closed(): + parameter = torch.nn.Parameter(torch.ones(8)) + optimizer = _make_gefen(parameter) + + assert not optimizer.state_offload_active + assert optimizer.state_offload_device is None + assert not optimizer.state_offload_poisoned + assert isinstance(optimizer, StateOffloadProvider) + assert not optimizer.optimizer_contract().capabilities.state_offload + with pytest.raises(AttributeError): + optimizer.state_offload_active = True + with pytest.raises(AttributeError): + optimizer.state_offload_device = torch.device("cpu") + with pytest.raises(AttributeError): + optimizer.state_offload_poisoned = True + with pytest.raises(RuntimeError, match="ordinary replicated CUDA"): + optimizer.offload_state_() + + +def test_state_offload_rejects_non_cpu_targets_and_muon(): + parameter = torch.nn.Parameter(torch.ones(2, 4)) + optimizer = _make_gefen(parameter) + muon = GefenMuon( + [("layer.weight", parameter)], + fused=False, + ns_steps=1, + ) + + with pytest.raises(ValueError, match="only CPU"): + optimizer.offload_state_("meta") + with pytest.raises(RuntimeError, match="only by plain Gefen"): + muon.offload_state_() + + +@_CUDA_REQUIRED +@pytest.mark.parametrize("factored", [False, True]) +@pytest.mark.parametrize("fused", [False, True]) +def test_pristine_offload_multistep_is_exact_and_cpu_authoritative(factored, fused): + shape = (2, 4) if factored else (8,) + initial = torch.linspace(-0.7, 0.6, 8, device="cuda").reshape(shape) + reference_parameter = torch.nn.Parameter(initial.clone()) + offloaded_parameter = torch.nn.Parameter(initial.clone()) + reference = _make_gefen(reference_parameter, factored=factored, fused=fused) + offloaded = _make_gefen(offloaded_parameter, factored=factored, fused=fused) + + assert offloaded.optimizer_contract().capabilities.state_offload + offloaded.offload_state_() + assert offloaded.optimizer_contract().capabilities.state_offload + _assert_cpu_boundary(offloaded) + assert offloaded._gefen_codebook is None + + for step in range(3): + grad = torch.linspace( + -1.1 + 0.2 * step, + 0.9 - 0.1 * step, + 8, + device="cuda", + ).reshape(shape) + reference_parameter.grad = grad.clone() + offloaded_parameter.grad = grad.clone() + reference.step() + offloaded.step() + + torch.testing.assert_close(offloaded_parameter, reference_parameter, rtol=0, atol=0) + _assert_persistent_equal( + _persistent_snapshot(offloaded, offloaded_parameter), + _persistent_snapshot(reference, reference_parameter), + ) + assert offloaded._gefen_global_step == reference._gefen_global_step + _assert_cpu_boundary(offloaded) + assert offloaded._gefen_codebook.device.type == "cuda" + + +@_CUDA_REQUIRED +def test_initialized_activation_keeps_common_codebook_resident_and_move_disables(): + parameter = torch.nn.Parameter(torch.linspace(-0.4, 0.8, 8, device="cuda")) + optimizer = _make_gefen(parameter) + parameter.grad = torch.linspace(-1.0, 1.0, 8, device="cuda") + optimizer.step() + codebook = optimizer._gefen_codebook + cache = optimizer._gefen_codebook_by_device + cache[parameter.device] = codebook + + optimizer.offload_state_("cpu:0") + + _assert_cpu_boundary(optimizer) + assert optimizer._gefen_codebook is codebook + assert optimizer._gefen_codebook_by_device is cache + assert optimizer._gefen_codebook_by_device[parameter.device] is codebook + + optimizer.move_state_() + + assert not optimizer.state_offload_active + assert optimizer.state_offload_device is None + assert not optimizer.state_offload_poisoned + assert all( + value.device == parameter.device + for key, value in optimizer.state[parameter].items() + if key in _PERSISTENT_TENSOR_KEYS + ) + + +@_CUDA_REQUIRED +def test_runtime_state_is_private_and_only_one_parameter_is_staged(): + first = torch.nn.Parameter(torch.arange(8, device="cuda", dtype=torch.float32)) + second = torch.nn.Parameter(torch.arange(8, 16, device="cuda", dtype=torch.float32)) + optimizer = Gefen( + [("first", first), ("second", second)], + fused=False, + factored_v_2d=False, + ) + optimizer._resolve_automatic_period = lambda *args, **kwargs: 4 + optimizer.offload_state_() + observed = [] + original = optimizer._step_automatic + + def inspected(self, group, name, parameter, grad, *, state=None): + assert state is not self.state[parameter] + assert all( + value.device.type == "cpu" + for published in self.state.values() + for key, value in published.items() + if key in _PERSISTENT_TENSOR_KEYS + ) + assert all(value.device == parameter.device for key, value in state.items() if key in _PERSISTENT_TENSOR_KEYS) + observed.append(parameter) + return original(group, name, parameter, grad, state=state) + + optimizer._step_automatic = types.MethodType(inspected, optimizer) + first.grad = torch.linspace(-1.0, 1.0, 8, device="cuda") + second.grad = torch.linspace(1.0, -1.0, 8, device="cuda") + optimizer.step() + + assert observed == [first, second] + _assert_cpu_boundary(optimizer) + + +@_CUDA_REQUIRED +def test_active_load_preserves_target_policy_and_exact_continuation(monkeypatch): + source_parameter = torch.nn.Parameter(torch.linspace(-0.5, 0.5, 8, device="cuda")) + source = _make_gefen(source_parameter) + source_parameter.grad = torch.linspace(-1.0, 0.7, 8, device="cuda") + source.step() + checkpoint = copy.deepcopy(source.state_dict()) + + target_parameter = torch.nn.Parameter(source_parameter.detach().clone()) + target = _make_gefen(target_parameter) + target.offload_state_() + + def reject_parameter_device_cast(*_args, **_kwargs): + raise AssertionError("active offload load used PyTorch's parameter-device cast") + + monkeypatch.setattr( + torch.optim.Optimizer, + "_process_value_according_to_param_policy", + reject_parameter_device_cast, + ) + target.load_state_dict(checkpoint) + + _assert_cpu_boundary(target) + _assert_persistent_equal( + _persistent_snapshot(source, source_parameter), + _persistent_snapshot(target, target_parameter), + ) + + continuation = torch.linspace(0.8, -0.6, 8, device="cuda") + source_parameter.grad = continuation.clone() + target_parameter.grad = continuation.clone() + source.step() + target.step() + + torch.testing.assert_close(target_parameter, source_parameter, rtol=0, atol=0) + _assert_persistent_equal( + _persistent_snapshot(source, source_parameter), + _persistent_snapshot(target, target_parameter), + ) + _assert_cpu_boundary(target) + + +@_CUDA_REQUIRED +@pytest.mark.parametrize("activate_before_load", [False, True]) +def test_cpu_mapped_checkpoint_keeps_common_codebook_cuda_resident( + activate_before_load, +): + source_parameter = torch.nn.Parameter(torch.linspace(-0.5, 0.5, 8, device="cuda")) + source = _make_gefen(source_parameter) + source_parameter.grad = torch.linspace(-1.0, 0.7, 8, device="cuda") + source.step() + serialized = io.BytesIO() + torch.save(source.state_dict(), serialized) + serialized.seek(0) + checkpoint = torch.load( + serialized, + map_location="cpu", + weights_only=False, + ) + assert checkpoint["gefen_codebook"].device.type == "cpu" + + target_parameter = torch.nn.Parameter(source_parameter.detach().clone()) + target = _make_gefen(target_parameter) + if activate_before_load: + target.offload_state_() + target.load_state_dict(checkpoint) + if not activate_before_load: + assert target._gefen_codebook.device.type == "cpu" + target.offload_state_() + + _assert_cpu_boundary(target) + assert target._gefen_codebook.device == target_parameter.device + torch.testing.assert_close( + target._gefen_codebook, + source._gefen_codebook, + rtol=0, + atol=0, + ) + + continuation = torch.linspace(0.8, -0.6, 8, device="cuda") + source_parameter.grad = continuation.clone() + target_parameter.grad = continuation.clone() + source.step() + target.step() + torch.testing.assert_close(target_parameter, source_parameter, rtol=0, atol=0) + _assert_cpu_boundary(target) + + +@_CUDA_REQUIRED +def test_activation_and_restore_copy_failures_are_atomic(monkeypatch): + parameter = torch.nn.Parameter(torch.linspace(-0.5, 0.5, 8, device="cuda")) + optimizer = _make_gefen(parameter) + parameter.grad = torch.linspace(-1.0, 1.0, 8, device="cuda") + optimizer.step() + state_before = optimizer.state + parameter_state_before = optimizer.state[parameter] + codebook_before = optimizer._gefen_codebook + + def fail_cpu_copy(_tensor): + raise RuntimeError("injected activation copy failure") + + monkeypatch.setattr(optimizer, "_copy_state_tensor_to_offload_cpu", fail_cpu_copy) + with pytest.raises(RuntimeError, match="injected activation"): + optimizer.offload_state_() + assert optimizer.state is state_before + assert optimizer.state[parameter] is parameter_state_before + assert optimizer._gefen_codebook is codebook_before + assert not optimizer.state_offload_active + + monkeypatch.undo() + optimizer.offload_state_() + state_before = optimizer.state + parameter_state_before = optimizer.state[parameter] + + def fail_move(_tensor, _device): + raise RuntimeError("injected restore copy failure") + + monkeypatch.setattr(optimizer, "_copy_state_tensor_for_move", fail_move) + with pytest.raises(RuntimeError, match="injected restore"): + optimizer.restore_state_() + assert optimizer.state is state_before + assert optimizer.state[parameter] is parameter_state_before + assert optimizer.state_offload_active + _assert_cpu_boundary(optimizer) + + +@_CUDA_REQUIRED +def test_copyback_failure_poison_is_sticky_until_successful_load(monkeypatch): + parameter = torch.nn.Parameter(torch.linspace(-0.4, 0.7, 8, device="cuda")) + optimizer = _make_gefen(parameter) + parameter.grad = torch.linspace(-1.0, 1.0, 8, device="cuda") + optimizer.step() + checkpoint = copy.deepcopy(optimizer.state_dict()) + optimizer.offload_state_() + published_state = optimizer.state[parameter] + parameter.grad = torch.linspace(0.9, -0.8, 8, device="cuda") + + def fail_copyback(_tensor): + raise RuntimeError("injected copyback failure") + + monkeypatch.setattr(optimizer, "_copy_state_tensor_to_offload_cpu", fail_copyback) + with pytest.raises(RuntimeError, match="known-good checkpoint"): + optimizer.step() + + assert optimizer.state_offload_active + assert optimizer.state_offload_poisoned + assert not optimizer.optimizer_contract().capabilities.state_offload + with pytest.raises(RuntimeError, match="cannot export optimizer state"): + optimizer.state_dict() + assert optimizer.state[parameter] is published_state + assert all(value.device.type == "cpu" for key, value in published_state.items() if key in _PERSISTENT_TENSOR_KEYS) + parameter_after_failure = parameter.detach().clone() + with pytest.raises(RuntimeError, match="poisoned"): + optimizer.step() + torch.testing.assert_close(parameter, parameter_after_failure, rtol=0, atol=0) + + monkeypatch.undo() + optimizer.load_state_dict(checkpoint) + assert optimizer.state_offload_active + assert not optimizer.state_offload_poisoned + assert optimizer.optimizer_contract().capabilities.state_offload + _assert_cpu_boundary(optimizer) + + +@_CUDA_REQUIRED +def test_active_offload_blocks_rebinding_and_portable_global_state(): + parameter = torch.nn.Parameter(torch.ones(8, device="cuda")) + optimizer = _make_gefen(parameter) + del optimizer._resolve_automatic_period + identity = ParameterIdentity("Model.Weight", (8,)) + group = ProcessGroupIdentity("singleton", ("rank:0",)) + shard = ShardIdentity( + identity, + ParameterLayout.REPLICATED, + LogicalSlice.full(identity), + placements=(ShardPlacement("checkpoint", PlacementKind.REPLICATE, 0, 1),), + process_group=group, + local_member="rank:0", + ) + manifest = ShardingManifest((shard,)) + codebook_binding = CodebookProcessGroupBinding( + group, + "rank:0", + None, + parameter.device, + ) + checkpoint_binding = CheckpointProcessGroupBinding( + group, + "rank:0", + None, + parameter.device, + ) + + optimizer.offload_state_() + with pytest.raises(RuntimeError, match="restore state first"): + optimizer.post_sharding( + (ParameterRebinding(parameter, parameter, shard),), + manifest=manifest, + codebook_process_group=codebook_binding, + ) + assert not optimizer._canonical_identity_ready() + + optimizer.restore_state_() + optimizer.post_sharding( + (ParameterRebinding(parameter, parameter, shard),), + manifest=manifest, + codebook_process_group=codebook_binding, + ) + optimizer.offload_state_() + assert all( + support.transport is not CheckpointTransport.CANONICAL_GLOBAL + for support in optimizer.optimizer_contract().capabilities.checkpoints + ) + with pytest.raises(RuntimeError, match="active optimizer-state offload"): + optimizer.export_portable_state( + checkpoint_process_group=checkpoint_binding, + transaction_id="active-offload-reject-v1", + limits=PortableStateLimits( + max_fragment_tensor_bytes=1 << 20, + max_collective_tensor_bytes=4 << 20, + max_collective_metadata_bytes=4 << 20, + max_metadata_bytes=1 << 20, + ), + ) + + +@_CUDA_REQUIRED +def test_state_offload_rejects_persistent_aliases_and_multi_member_scope(): + parameter = torch.nn.Parameter(torch.linspace(-0.5, 0.5, 8, device="cuda")) + optimizer = _make_gefen(parameter) + parameter.grad = torch.linspace(-1.0, 1.0, 8, device="cuda") + optimizer.step() + parameter_state = optimizer.state[parameter] + parameter_state["vmean"] = parameter_state["m_magnitude"] + state_before = optimizer.state + + assert not optimizer.optimizer_contract().capabilities.state_offload + with pytest.raises(RuntimeError, match="storage aliases"): + optimizer.offload_state_() + assert optimizer.state is state_before + assert parameter_state["vmean"] is parameter_state["m_magnitude"] + + parameter_state["vmean"] = parameter_state["m_magnitude"].clone() + group = ProcessGroupIdentity("multi", ("rank:0", "rank:1")) + optimizer._gefen_codebook_process_group = CodebookProcessGroupBinding( + group, + "rank:0", + object(), + parameter.device, + ) + with pytest.raises(RuntimeError, match="multi-member"): + optimizer.offload_state_() + + +@_CUDA_REQUIRED +@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="two CUDA devices are required") +def test_state_offload_checks_capture_on_every_parameter_device(monkeypatch): + first = torch.nn.Parameter(torch.ones(8, device="cuda:0")) + second = torch.nn.Parameter(torch.ones(8, device="cuda:1")) + optimizer = Gefen( + [("first", first), ("second", second)], + fused=False, + factored_v_2d=False, + ) + monkeypatch.setattr( + torch.cuda, + "is_current_stream_capturing", + lambda: torch.cuda.current_device() == second.device.index, + ) + + assert not optimizer.optimizer_contract().capabilities.atomic_state_movement + with pytest.raises(RuntimeError, match="CUDA graph capture"): + optimizer.move_state_() + with pytest.raises(RuntimeError, match="CUDA graph capture"): + optimizer.offload_state_() + assert not optimizer.state_offload_active + + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False) + optimizer.offload_state_() + monkeypatch.setattr( + torch.cuda, + "is_current_stream_capturing", + lambda: torch.cuda.current_device() == second.device.index, + ) + with pytest.raises(RuntimeError, match="CUDA graph capture"): + optimizer.restore_state_() + assert optimizer.state_offload_active + + +@_CUDA_REQUIRED +def test_operation_error_copies_state_back_without_poisoning(monkeypatch): + parameter = torch.nn.Parameter(torch.linspace(-0.4, 0.7, 8, device="cuda")) + optimizer = _make_gefen(parameter) + parameter.grad = torch.linspace(-1.0, 1.0, 8, device="cuda") + optimizer.step() + optimizer.offload_state_() + original = optimizer._step_automatic + + def update_then_fail(group, name, live_parameter, grad, *, state=None): + original(group, name, live_parameter, grad, state=state) + raise RuntimeError("injected update failure") + + monkeypatch.setattr(optimizer, "_step_automatic", update_then_fail) + parameter.grad = torch.linspace(0.8, -0.6, 8, device="cuda") + with pytest.raises(RuntimeError, match="injected update"): + optimizer.step() + + assert not optimizer.state_offload_poisoned + _assert_cpu_boundary(optimizer) + assert optimizer.state[parameter]["step"] == 2 + + +@_CUDA_REQUIRED +def test_custom_dtensor_like_and_finalized_nonreplicated_state_fail_closed(): + custom_parameter = torch.nn.Parameter(torch.ones(8, device="cuda")) + custom = _make_gefen(custom_parameter) + custom.state[custom_parameter]["extension"] = {"value": 1} + state_before = custom.state + with pytest.raises(RuntimeError, match="custom per-parameter state"): + custom.offload_state_() + assert custom.state is state_before + assert not custom.state_offload_active + + dtensor_like_parameter = torch.nn.Parameter(torch.ones(8, device="cuda")) + dtensor_like_parameter.to_local = lambda: dtensor_like_parameter + dtensor_like_parameter.placements = () + dtensor_like = _make_gefen(dtensor_like_parameter) + with pytest.raises(RuntimeError, match="ordinary replicated CUDA"): + dtensor_like.offload_state_() + + flat_parameter = torch.nn.Parameter(torch.ones(8, device="cuda")) + flattened = _make_gefen(flat_parameter) + identity = ParameterIdentity("layer.weight", (16,)) + process_group = ProcessGroupIdentity("data_parallel", ("rank:0", "rank:1")) + shards = tuple( + ShardIdentity( + identity, + ParameterLayout.FLATTENED_ELEMENT_SHARD, + LogicalSlice(index * 8, 8), + placements=( + ShardPlacement( + "data_parallel", + PlacementKind.FLAT_SHARD, + index, + 2, + ), + ), + process_group=process_group, + local_member=member, + ) + for index, member in enumerate(process_group.ordered_members) + ) + shard = shards[0] + flattened.rebind_shard( + flat_parameter, + flat_parameter, + shard=shard, + manifest=ShardingManifest(shards), + ) + with pytest.raises(RuntimeError, match="finalized replicated"): + flattened.offload_state_() + + +@_CUDA_REQUIRED +def test_capturable_compile_and_capture_states_fail_closed(monkeypatch): + parameter = torch.nn.Parameter(torch.ones(8, device="cuda")) + capturable = Gefen( + [("layer.weight", parameter)], + fused=False, + factored_v_2d=False, + capturable=True, + ) + with pytest.raises(RuntimeError, match="capturable"): + capturable.offload_state_() + + optimizer = _make_gefen(torch.nn.Parameter(torch.ones(8, device="cuda"))) + monkeypatch.setattr(torch.compiler, "is_compiling", lambda: True) + with pytest.raises(RuntimeError, match="torch.compile"): + optimizer.offload_state_() + monkeypatch.setattr(torch.compiler, "is_compiling", lambda: False) + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: True) + with pytest.raises(RuntimeError, match="CUDA graph capture"): + optimizer.offload_state_() From 7e2aaec4ee9342ea4264bb83305a59e471c330d8 Mon Sep 17 00:00:00 2001 From: thad0ctor Date: Mon, 13 Jul 2026 08:11:08 -0700 Subject: [PATCH 19/52] Type-validate OptimizerCapabilities and OptimizerChildContract contents OptimizerCapabilities accepted arbitrary objects in training, checkpoints, and precisions and non-bool capability flags, and OptimizerChildContract accepted any object as its child contract, so an invalid declaration constructed silently and failed far from where it was built. Check element and flag types in __post_init__ like every sibling descriptor. --- src/gefen/contracts.py | 33 +++++++++++++++++++++++++++++ tests/test_optimizer_contracts.py | 35 ++++++++++++++++++++++++++++++- 2 files changed, 67 insertions(+), 1 deletion(-) diff --git a/src/gefen/contracts.py b/src/gefen/contracts.py index 4718165..0e9cd32 100644 --- a/src/gefen/contracts.py +++ b/src/gefen/contracts.py @@ -1019,6 +1019,33 @@ def __post_init__(self) -> None: object.__setattr__(self, "training", _tuple(self.training)) object.__setattr__(self, "checkpoints", _tuple(self.checkpoints)) object.__setattr__(self, "precisions", _frozenset(self.precisions)) + if any(not isinstance(item, TrainingSupport) for item in self.training): + raise TypeError( + "OptimizerCapabilities.training must contain TrainingSupport values" + ) + if any(not isinstance(item, CheckpointSupport) for item in self.checkpoints): + raise TypeError( + "OptimizerCapabilities.checkpoints must contain CheckpointSupport values" + ) + if any(not isinstance(item, Precision) for item in self.precisions): + raise TypeError( + "OptimizerCapabilities.precisions must contain Precision values" + ) + for name in ( + "accepts_semantic_parameter_names", + "canonical_parameter_fqns", + "stable_shard_identity", + "explicit_process_group_codebook_scope", + "shard_rebinding", + "post_sharding", + "canonical_state_io", + "atomic_state_movement", + "state_offload", + ): + if type(getattr(self, name)) is not bool: + raise TypeError( + "OptimizerCapabilities.{} must be a bool".format(name) + ) if self.supported_parameter_ranks is not None: object.__setattr__( self, @@ -1045,6 +1072,12 @@ def __post_init__(self) -> None: raise ValueError("OptimizerChildContract.role must be non-empty") if not self.implementation: raise ValueError("OptimizerChildContract.implementation must be non-empty") + if self.contract is not None and not isinstance( + self.contract, OptimizerContract + ): + raise TypeError( + "OptimizerChildContract.contract must be an OptimizerContract or None" + ) @dataclass(frozen=True) diff --git a/tests/test_optimizer_contracts.py b/tests/test_optimizer_contracts.py index 9504f1c..77a635f 100644 --- a/tests/test_optimizer_contracts.py +++ b/tests/test_optimizer_contracts.py @@ -2,7 +2,7 @@ import copy from collections import defaultdict, OrderedDict -from dataclasses import FrozenInstanceError, dataclass +from dataclasses import FrozenInstanceError, dataclass, replace import pytest import torch @@ -14,6 +14,7 @@ Gefen, GefenMuon, GefenMuonHybrid, + OptimizerChildContract, OptimizerContractProvider, OptimizerStateLayout, ParameterLayout, @@ -758,3 +759,35 @@ def test_portable_global_transport_is_defined_but_not_claimed_before_integration support.transport is not CheckpointTransport.CANONICAL_GLOBAL for support in optimizer.optimizer_contract().capabilities.checkpoints ) + + +def test_capabilities_reject_untyped_entries_and_flags(): + parameter = torch.nn.Parameter(torch.ones(4)) + optimizer = Gefen([("parameter", parameter)], fused=False) + capabilities = optimizer.optimizer_contract().capabilities + with pytest.raises(TypeError, match="training must contain TrainingSupport"): + replace(capabilities, training=[{"junk": 1}]) + with pytest.raises(TypeError, match="checkpoints must contain CheckpointSupport"): + replace(capabilities, checkpoints=["nonsense"]) + with pytest.raises(TypeError, match="precisions must contain Precision"): + replace(capabilities, precisions={"float32"}) + for name in ( + "accepts_semantic_parameter_names", + "canonical_parameter_fqns", + "stable_shard_identity", + "explicit_process_group_codebook_scope", + "shard_rebinding", + "post_sharding", + "canonical_state_io", + "atomic_state_movement", + "state_offload", + ): + with pytest.raises(TypeError, match="{} must be a bool".format(name)): + replace(capabilities, **{name: "yes"}) + + +def test_child_contract_rejects_untyped_contract_payload(): + with pytest.raises( + TypeError, match="contract must be an OptimizerContract or None" + ): + OptimizerChildContract("backup", "torch.optim.adamw.AdamW", {"junk": 1}) From 6668940cbe4cea772e89768d686a7c52a88d8b45 Mon Sep 17 00:00:00 2001 From: thad0ctor Date: Mon, 13 Jul 2026 08:11:31 -0700 Subject: [PATCH 20/52] Compose hybrid training claims from both routed children _hybrid_contract copied the muon child's training claims verbatim and fell back to the static base tuple when muon was absent, never consulting the backup child. A Gefen-backed backup-only hybrid therefore omitted the flattened element-shard layout its backup validates, while an AdamW-backed hybrid claimed DTensor training no code validates for that child. Declare the ordered union of the present children's own claims instead; a backup without a contract contributes only plain replicated training. --- src/gefen/contracts.py | 38 ++++++++++++++++++--- tests/test_optimizer_contracts.py | 55 +++++++++++++++++++++++++++++++ 2 files changed, 89 insertions(+), 4 deletions(-) diff --git a/src/gefen/contracts.py b/src/gefen/contracts.py index 0e9cd32..37144fa 100644 --- a/src/gefen/contracts.py +++ b/src/gefen/contracts.py @@ -1955,6 +1955,39 @@ def _gefen_muon_contract( ) +def _hybrid_training( + muon: Optional[OptimizerContract], + backup: Optional[OptimizerContract], + backup_present: bool, +) -> Tuple[TrainingSupport, ...]: + """Compose hybrid training claims from the routed children's own claims. + + The hybrid routes each parameter to exactly one child, so the composite + declares the ordered union of the present children's validated claims; an + adapter must read each role's child contract together with the frozen + parameter routing rather than apply one entry to every parameter. A backup + child without a contract of its own (AdamW) contributes only the plain + replicated layout the composite actually exercises for it, because no code + validates any other layout for that child. + """ + + child_claims = [] + if muon is not None: + child_claims.append(muon.capabilities.training) + if backup is not None: + child_claims.append(backup.capabilities.training) + elif backup_present: + child_claims.append( + (TrainingSupport(ParameterLayout.REPLICATED, ProcessGroupScope.NONE),) + ) + training = [] + for claims in child_claims: + for support in claims: + if support not in training: + training.append(support) + return tuple(training) + + def _hybrid_contract( *, muon: Optional[OptimizerContract], @@ -1982,10 +2015,7 @@ def _hybrid_contract( children.append(OptimizerChildContract("muon", muon.implementation, muon)) if backup_implementation: children.append(OptimizerChildContract("backup", backup_implementation, backup)) - if muon is None: - training = _base_training() - else: - training = muon.capabilities.training + training = _hybrid_training(muon, backup, bool(backup_implementation)) canonical_global_same_topology = _frozenset(canonical_global_same_topology) canonical_global_topology_changing = _frozenset(canonical_global_topology_changing) canonical_global_topology_change_kinds = _frozenset( diff --git a/tests/test_optimizer_contracts.py b/tests/test_optimizer_contracts.py index 77a635f..5be414a 100644 --- a/tests/test_optimizer_contracts.py +++ b/tests/test_optimizer_contracts.py @@ -30,6 +30,7 @@ StateScope, StateVariant, TopologyChange, + TrainingSupport, ) @@ -761,6 +762,60 @@ def test_portable_global_transport_is_defined_but_not_claimed_before_integration ) +@pytest.mark.parametrize("backup_optimizer", ["gefen", "adamw"]) +def test_backup_only_hybrid_training_claims_come_from_backup_child(backup_optimizer): + bias = torch.nn.Parameter(torch.arange(4, dtype=torch.float32)) + optimizer = GefenMuonHybrid( + [], + [("layer.bias", bias)], + lr=1e-3, + fused=False, + backup_optimizer=backup_optimizer, + ) + assert optimizer.muon is None + contract = optimizer.optimizer_contract() + assert tuple(child.role for child in contract.children) == ("backup",) + if backup_optimizer == "gefen": + backup_training = contract.children[0].contract.capabilities.training + assert contract.capabilities.training == backup_training + assert any( + item.layout is ParameterLayout.FLATTENED_ELEMENT_SHARD + for item in contract.capabilities.training + ) + else: + # AdamW publishes no contract, so the composite keeps only the plain + # replicated layout the hybrid actually exercises for that child; a + # DTensor claim no code validated would be an over-claim. + assert contract.children[0].contract is None + assert contract.capabilities.training == ( + TrainingSupport(ParameterLayout.REPLICATED, ProcessGroupScope.NONE), + ) + + +@pytest.mark.parametrize("backup_optimizer", ["gefen", "adamw"]) +def test_hybrid_training_claims_are_the_union_of_routed_children(backup_optimizer): + matrix = torch.nn.Parameter(torch.ones(4, 4)) + bias = torch.nn.Parameter(torch.ones(4)) + optimizer = GefenMuonHybrid( + [("layer.weight", matrix)], + [("layer.bias", bias)], + lr=1e-3, + fused=False, + backup_optimizer=backup_optimizer, + ) + contract = optimizer.optimizer_contract() + training = contract.capabilities.training + muon_training = contract.children[0].contract.capabilities.training + if backup_optimizer == "gefen": + backup_training = contract.children[1].contract.capabilities.training + else: + backup_training = ( + TrainingSupport(ParameterLayout.REPLICATED, ProcessGroupScope.NONE), + ) + assert set(training) == set(muon_training) | set(backup_training) + assert len(set(training)) == len(training) + + def test_capabilities_reject_untyped_entries_and_flags(): parameter = torch.nn.Parameter(torch.ones(4)) optimizer = Gefen([("parameter", parameter)], fused=False) From 9f23f135ce00797881f5a3790feb572db094adaf Mon Sep 17 00:00:00 2001 From: thad0ctor Date: Mon, 13 Jul 2026 08:11:42 -0700 Subject: [PATCH 21/52] Deepen fail-before-mutation snapshots in canonical and hybrid tests The canonical-import and hybrid-rebinding atomicity tests asserted only top-level object identity, so a regression that published staged state by mutating live containers in place (state[param].update, copy_() into an existing state tensor, group option or group['params'] element edits) would pass every assertion. Add tests/_state_snapshot.py, a shared deep snapshot helper that captures container identities plus bitwise clones of every reachable tensor, per-parameter state dict, counter, and param-group entry, and rewire both files' snapshot helpers onto it. Verified that all eight simulated in-place regressions now fail the assertions while the current implementation still passes. --- tests/_state_snapshot.py | 126 ++++++++++++++++++++++++++++++ tests/test_canonical_state_cpu.py | 15 ++-- tests/test_hybrid_rebinding.py | 14 ++-- 3 files changed, 142 insertions(+), 13 deletions(-) create mode 100644 tests/_state_snapshot.py diff --git a/tests/_state_snapshot.py b/tests/_state_snapshot.py new file mode 100644 index 0000000..777ab9e --- /dev/null +++ b/tests/_state_snapshot.py @@ -0,0 +1,126 @@ +"""Shared deep fail-before-mutation snapshot helpers for optimizer tests. + +``deep_state_snapshot`` captures both the container identities and bitwise +clones of every tensor reachable from an optimizer (per-parameter state, +codebook tensors and per-device caches, counters, param-group contents). A +regression that publishes staged data by mutating live containers in place — +``state[param].update(...)``, ``copy_()`` into an existing state tensor, an +in-place ``group['params']`` element swap, or an in-place option edit — +preserves every top-level object identity and is only caught by comparing the +live values against these clones. ``assert_deep_state_snapshot`` therefore +checks identity AND bitwise value equality together. +""" + +import copy + +import torch + + +def _cloned(value): + if torch.is_tensor(value): + return value.detach().clone() + if type(value) is dict: + return {key: _cloned(item) for key, item in value.items()} + if type(value) in (list, tuple): + return type(value)(_cloned(item) for item in value) + return copy.deepcopy(value) + + +def _nested_equal(live, expected): + if torch.is_tensor(expected): + assert torch.is_tensor(live) + assert torch.equal(live, expected) + return + assert type(live) is type(expected) + if isinstance(expected, dict): + assert set(live) == set(expected) + for key in expected: + _nested_equal(live[key], expected[key]) + elif isinstance(expected, (list, tuple)): + assert len(live) == len(expected) + for live_item, expected_item in zip(live, expected): + _nested_equal(live_item, expected_item) + else: + assert live == expected + + +def _tensor_value_pairs(optimizer): + """Every tensor reachable from ``optimizer.__dict__`` with a clone of it.""" + + pairs = [] + visited = set() + + def visit(value): + if torch.is_tensor(value): + if id(value) not in visited: + visited.add(id(value)) + pairs.append((value, value.detach().clone())) + elif isinstance(value, dict): + if id(value) in visited: + return + visited.add(id(value)) + for key, item in value.items(): + visit(key) + visit(item) + elif isinstance(value, (list, tuple, set, frozenset)): + if id(value) in visited: + return + visited.add(id(value)) + for item in value: + visit(item) + + for value in optimizer.__dict__.values(): + visit(value) + return tuple(pairs) + + +def deep_state_snapshot(optimizer): + return { + "attributes": optimizer.__dict__.copy(), + "state": optimizer.state, + "state_items": tuple( + (parameter, state, _cloned(dict(state))) + for parameter, state in optimizer.state.items() + ), + "param_groups": optimizer.param_groups, + "groups": tuple( + ( + group, + group["params"], + tuple(group["params"]), + _cloned({key: value for key, value in group.items() if key != "params"}), + ) + for group in optimizer.param_groups + ), + "tensors": _tensor_value_pairs(optimizer), + } + + +def assert_deep_state_snapshot(optimizer, snapshot): + assert optimizer.__dict__.keys() == snapshot["attributes"].keys() + for name, value in snapshot["attributes"].items(): + assert optimizer.__dict__[name] is value + assert optimizer.state is snapshot["state"] + assert optimizer.param_groups is snapshot["param_groups"] + assert len(optimizer.state) == len(snapshot["state_items"]) + for live_parameter, (parameter, state, expected) in zip( + optimizer.state, snapshot["state_items"] + ): + assert live_parameter is parameter + assert optimizer.state[parameter] is state + _nested_equal(dict(state), expected) + assert len(optimizer.param_groups) == len(snapshot["groups"]) + for live_group, (group, params, expected_params, expected_options) in zip( + optimizer.param_groups, snapshot["groups"] + ): + assert live_group is group + assert group["params"] is params + assert len(params) == len(expected_params) + for live_param, expected_param in zip(params, expected_params): + assert live_param is expected_param + _nested_equal( + {key: value for key, value in group.items() if key != "params"}, + expected_options, + ) + for tensor, expected in snapshot["tensors"]: + assert torch.equal(tensor, expected) diff --git a/tests/test_canonical_state_cpu.py b/tests/test_canonical_state_cpu.py index 3e5becc..cd1cfe7 100644 --- a/tests/test_canonical_state_cpu.py +++ b/tests/test_canonical_state_cpu.py @@ -24,6 +24,8 @@ ) from gefen.rebinding import LogicalSlotBinding +from _state_snapshot import assert_deep_state_snapshot, deep_state_snapshot + def _replicated(identity, group=None, member=None): placements = () @@ -74,23 +76,20 @@ def _finalize(optimizer, bindings, manifest): def _snapshot(optimizer): + # Deep snapshot: container identities plus bitwise clones of every state + # tensor, codebook tensor, counter, and param-group entry, so a rejected + # import that mutated live containers in place cannot pass unnoticed. return { - "dict": optimizer.__dict__.copy(), - "groups": optimizer.param_groups, - "state": optimizer.state, + "deep": deep_state_snapshot(optimizer), "codebook": optimizer._gefen_codebook, "global_step": optimizer._gefen_global_step, } def _assert_snapshot_identity(optimizer, snapshot): - assert optimizer.param_groups is snapshot["groups"] - assert optimizer.state is snapshot["state"] assert optimizer._gefen_codebook is snapshot["codebook"] assert optimizer._gefen_global_step is snapshot["global_step"] - assert optimizer.__dict__.keys() == snapshot["dict"].keys() - for key, value in snapshot["dict"].items(): - assert optimizer.__dict__[key] is value + assert_deep_state_snapshot(optimizer, snapshot["deep"]) def _two_parameter_source(): diff --git a/tests/test_hybrid_rebinding.py b/tests/test_hybrid_rebinding.py index 284f30c..f74e01d 100644 --- a/tests/test_hybrid_rebinding.py +++ b/tests/test_hybrid_rebinding.py @@ -17,6 +17,8 @@ ) from gefen.rebinding import ParameterRebinding +from _state_snapshot import assert_deep_state_snapshot, deep_state_snapshot + _MEMBER = "rank:0" @@ -116,8 +118,12 @@ def _owner_shards(fqn, shape, group, owner): def _snapshot(optimizer): + # Each child is snapshotted deeply (container identities plus bitwise + # clones of per-parameter state, group['params'] entries, and options) + # so a failed composite transaction that partially rebound a child in + # place cannot pass on top-level identity alone. return { - "children": tuple((child, child.__dict__.copy()) for child in optimizer._subopts), + "children": tuple((child, deep_state_snapshot(child)) for child in optimizer._subopts), "owner": optimizer._state_param_owner, "finalized": optimizer._hybrid_post_sharding_finalized, "manifest": optimizer._hybrid_sharding_manifest, @@ -137,11 +143,9 @@ def _assert_snapshot(optimizer, snapshot): assert optimizer._hybrid_codebook_process_group is snapshot["binding"] assert optimizer._hybrid_finalized_slots is snapshot["slots"] assert len(optimizer._subopts) == len(snapshot["children"]) - for live, (expected_child, expected_attributes) in zip(optimizer._subopts, snapshot["children"]): + for live, (expected_child, expected_snapshot) in zip(optimizer._subopts, snapshot["children"]): assert live is expected_child - assert set(live.__dict__) == set(expected_attributes) - for name, expected in expected_attributes.items(): - assert live.__dict__[name] is expected + assert_deep_state_snapshot(live, expected_snapshot) def test_composite_post_sharding_publishes_children_and_rebuilds_routing(): From 1336281498a6bb5d1e059343356215f0a2ffdc56 Mon Sep 17 00:00:00 2001 From: thad0ctor Date: Mon, 13 Jul 2026 08:11:54 -0700 Subject: [PATCH 22/52] Assert persistent state values in offload failure-atomicity test test_activation_and_restore_copy_failures_are_atomic only checked that optimizer.state, the per-parameter dict, and the codebook were the same objects after an injected copy failure; an offload/restore regression that partially overwrote m_codebook/m_magnitude/vmean in place before failing would keep every identity intact and pass. Snapshot the persistent tensors and counters with the file's existing _persistent_snapshot helper before each injected failure and assert bitwise equality (plus codebook value and global step) afterwards. Verified the new assertions catch a simulated in-place corruption that the old identity checks missed. --- tests/test_state_offload.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/tests/test_state_offload.py b/tests/test_state_offload.py index 3a0a6bb..7235aa9 100644 --- a/tests/test_state_offload.py +++ b/tests/test_state_offload.py @@ -314,6 +314,9 @@ def test_activation_and_restore_copy_failures_are_atomic(monkeypatch): state_before = optimizer.state parameter_state_before = optimizer.state[parameter] codebook_before = optimizer._gefen_codebook + codebook_value_before = optimizer._gefen_codebook.detach().cpu().clone() + persistent_before = _persistent_snapshot(optimizer, parameter) + global_step_before = optimizer._gefen_global_step def fail_cpu_copy(_tensor): raise RuntimeError("injected activation copy failure") @@ -325,11 +328,15 @@ def fail_cpu_copy(_tensor): assert optimizer.state[parameter] is parameter_state_before assert optimizer._gefen_codebook is codebook_before assert not optimizer.state_offload_active + _assert_persistent_equal(_persistent_snapshot(optimizer, parameter), persistent_before) + assert torch.equal(optimizer._gefen_codebook.detach().cpu(), codebook_value_before) + assert optimizer._gefen_global_step == global_step_before monkeypatch.undo() optimizer.offload_state_() state_before = optimizer.state parameter_state_before = optimizer.state[parameter] + persistent_before = _persistent_snapshot(optimizer, parameter) def fail_move(_tensor, _device): raise RuntimeError("injected restore copy failure") @@ -341,6 +348,9 @@ def fail_move(_tensor, _device): assert optimizer.state[parameter] is parameter_state_before assert optimizer.state_offload_active _assert_cpu_boundary(optimizer) + _assert_persistent_equal(_persistent_snapshot(optimizer, parameter), persistent_before) + assert torch.equal(optimizer._gefen_codebook.detach().cpu(), codebook_value_before) + assert optimizer._gefen_global_step == global_step_before @_CUDA_REQUIRED From 7920078c0f0f7d3efdc9075adfde34b3e20eb97d Mon Sep 17 00:00:00 2001 From: thad0ctor Date: Mon, 13 Jul 2026 08:11:54 -0700 Subject: [PATCH 23/52] Cover zero-length flattened shard in scoped codebook collectives Every flattened shard in the scoped-collective tests had nonzero length, leaving the dedicated empty-shard logic (nonempty_activity filtering of the gradient-presence consensus and the empty-slice no-op paths) entirely unexercised. Add a two-member gloo test where one member binds a legal zero-length flattened slice to a live numel-0 tensor and assert it joins initialization, step, refresh, and failure-sync collectives symmetrically: codebooks agree with the nonempty-only oracle on both members, presence asymmetry on the empty member is accepted, the empty member's state stays inert, a nonempty-member exact-DP failure is synchronized to the empty member and leaves both atomic and retryable. --- tests/test_codebook_scope_distributed.py | 201 +++++++++++++++++++++++ 1 file changed, 201 insertions(+) diff --git a/tests/test_codebook_scope_distributed.py b/tests/test_codebook_scope_distributed.py index 4120be4..f0fe80c 100644 --- a/tests/test_codebook_scope_distributed.py +++ b/tests/test_codebook_scope_distributed.py @@ -935,6 +935,207 @@ def test_explicit_gloo_subgroups_are_isolated_from_default_world(): assert first[0] != second[0] +def _zero_length_flat_worker(rank, world, init_file, queue): + try: + dist.init_process_group( + "gloo", + init_method="file://{}".format(init_file), + rank=rank, + world_size=world, + timeout=timedelta(seconds=45), + ) + members = tuple("rank:{}".format(index) for index in range(world)) + group = ProcessGroupIdentity("data_parallel", members) + runtime_group = dist.group.WORLD + + # Uneven sharding of a small parameter: member one holds a legal + # zero-length flattened slice bound to a live numel-0 tensor. The + # empty member creates no per-parameter state and its gradient + # presence is excluded from the nonempty-shard consensus, but it must + # still join every scoped collective symmetrically. + identity = ParameterIdentity("ZeroFlat", (4,)) + records = ( + _flat(identity, group, "rank:0", 0, 4), + _flat(identity, group, "rank:1", 4, 0), + ) + parameter = torch.nn.Parameter(torch.zeros(4 if rank == 0 else 0)) + optimizer = Gefen([("zero_flat", parameter)], fused=False, factored_v_2d=False) + _finalize( + optimizer, + parameter, + records[rank], + ShardingManifest(records), + _binding(group, rank, runtime_group), + ) + optimizer._resolve_automatic_period = lambda *args: 4 + nonempty_grad = torch.tensor([-4.0, -1.0, 2.0, 8.0]) + if rank == 0: + parameter.grad = nonempty_grad.clone() + initialized = optimizer.initialize_codebook() + init_oracle = learn_gefen_exact_codebook_from_grad_periods( + grad_periods=(("ZeroFlat", nonempty_grad, 4, nonempty_grad),), + codebook_device=torch.device("cpu"), + num_codebooks=256, + force_endpoints=True, + verbose=False, + compute_mse_logging=False, + use_fused_histogram=False, + ) + init_matches_oracle = torch.equal(optimizer._gefen_codebook, init_oracle) + init_codebooks = [torch.empty_like(optimizer._gefen_codebook) for _ in range(world)] + dist.all_gather(init_codebooks, optimizer._gefen_codebook) + init_agreement = all(torch.equal(item, init_codebooks[0]) for item in init_codebooks[1:]) + + optimizer.step() + step_valid = optimizer._gefen_global_step == 1 and ( + rank == 0 or optimizer.state[parameter] == {"name": "zero_flat"} + ) + + continuation_grad = nonempty_grad.flip(0) + if rank == 0: + parameter.grad = continuation_grad.clone() + refreshed = optimizer.refresh_codebook() + refresh_oracle = learn_gefen_exact_codebook_from_grad_periods( + grad_periods=(("ZeroFlat", continuation_grad, 4, continuation_grad),), + codebook_device=torch.device("cpu"), + num_codebooks=256, + force_endpoints=True, + verbose=False, + compute_mse_logging=False, + use_fused_histogram=False, + ) + refresh_matches_oracle = torch.equal(optimizer._gefen_codebook, refresh_oracle) + refreshed_codebooks = [torch.empty_like(optimizer._gefen_codebook) for _ in range(world)] + dist.all_gather(refreshed_codebooks, optimizer._gefen_codebook) + refresh_agreement = all(torch.equal(item, refreshed_codebooks[0]) for item in refreshed_codebooks[1:]) + empty_state_inert = rank == 0 or optimizer.state[parameter] == {"name": "zero_flat"} + + # A failure on the nonempty member must be observed by the empty + # member through the failure-sync collective, leave both members + # unchanged, and keep the pair retryable. + failure_param = torch.nn.Parameter(torch.zeros(4 if rank == 0 else 0)) + failure_optimizer = Gefen( + [("zero_failure", failure_param)], + fused=False, + factored_v_2d=False, + ) + failure_identity = ParameterIdentity("ZeroFailure", (4,)) + failure_records = ( + _flat(failure_identity, group, "rank:0", 0, 4), + _flat(failure_identity, group, "rank:1", 4, 0), + ) + _finalize( + failure_optimizer, + failure_param, + failure_records[rank], + ShardingManifest(failure_records), + _binding(group, rank, runtime_group), + ) + failure_optimizer._resolve_automatic_period = lambda *args: 4 + if rank == 0: + failure_param.grad = nonempty_grad.clone() + original_exact_dp = gefen_module.quantization_module.exact_dp + if rank == 0: + + def fail_exact_dp(*args, **kwargs): + raise RuntimeError("rank-local exact-DP failure") + + gefen_module.quantization_module.exact_dp = fail_exact_dp + try: + failure_optimizer.initialize_codebook() + failure_seen = False + except RuntimeError as exc: + failure_seen = "exact-DP" in str(exc) + finally: + gefen_module.quantization_module.exact_dp = original_exact_dp + failure_atomic = ( + failure_optimizer._gefen_codebook is None + and failure_optimizer._gefen_global_step == 0 + and failure_optimizer.state[failure_param] == {"name": "zero_failure"} + ) + retry_succeeded = failure_optimizer.initialize_codebook() + retry_codebooks = [torch.empty_like(failure_optimizer._gefen_codebook) for _ in range(world)] + dist.all_gather(retry_codebooks, failure_optimizer._gefen_codebook) + retry_agreement = all(torch.equal(item, retry_codebooks[0]) for item in retry_codebooks[1:]) + + queue.put( + { + "rank": rank, + "initialized": initialized, + "init_matches_oracle": init_matches_oracle, + "init_agreement": init_agreement, + "step_valid": step_valid, + "refreshed": refreshed, + "refresh_matches_oracle": refresh_matches_oracle, + "refresh_agreement": refresh_agreement, + "empty_state_inert": empty_state_inert, + "failure_seen": failure_seen, + "failure_atomic": failure_atomic, + "retry_succeeded": retry_succeeded, + "retry_agreement": retry_agreement, + } + ) + except Exception as exc: + queue.put({"rank": rank, "error": repr(exc)}) + finally: + if dist.is_available() and dist.is_initialized(): + dist.destroy_process_group() + + +def _run_zero_length_flat_workers(world=2): + context = mp.get_context("spawn") + queue = context.Queue() + fd, init_file = tempfile.mkstemp(prefix="gefen-codebook-zero-flat-") + os.close(fd) + os.unlink(init_file) + processes = [ + context.Process( + target=_zero_length_flat_worker, + args=(rank, world, init_file, queue), + ) + for rank in range(world) + ] + try: + for process in processes: + process.start() + results = [queue.get(timeout=60) for _ in processes] + for process in processes: + process.join(timeout=10) + for process in processes: + if process.is_alive(): + process.terminate() + process.join(timeout=5) + pytest.fail("zero-length flattened shard worker hung") + assert process.exitcode == 0 + return sorted(results, key=lambda item: item["rank"]) + finally: + for process in processes: + if process.is_alive(): + process.terminate() + process.join(timeout=5) + if os.path.exists(init_file): + os.unlink(init_file) + + +def test_zero_length_flattened_shard_member_joins_every_scoped_collective(): + results = _run_zero_length_flat_workers() + + assert all("error" not in item for item in results), results + for item in results: + assert item["initialized"], item + assert item["init_matches_oracle"], item + assert item["init_agreement"], item + assert item["step_valid"], item + assert item["refreshed"], item + assert item["refresh_matches_oracle"], item + assert item["refresh_agreement"], item + assert item["empty_state_inert"], item + assert item["failure_seen"], item + assert item["failure_atomic"], item + assert item["retry_succeeded"], item + assert item["retry_agreement"], item + + def _nccl_empty_owner_worker(rank, world, init_file, queue): try: torch.cuda.set_device(rank) From e73d4ac512e76941954f5922470bb0ae3edf520c Mon Sep 17 00:00:00 2001 From: thad0ctor Date: Mon, 13 Jul 2026 08:14:15 -0700 Subject: [PATCH 24/52] Scope-agree checkpoint-resume period reuse before codebook collectives The reuse_existing_periods flag that gates the conditional "initialize" operation header in _prepare_gefen_exact_codebook was derived from the rank-local _resuming_from_checkpoint() predicate on both hot paths (_maybe_refresh_gefen_codebook and initialize_codebook). Per-parameter state presence is legitimately rank-asymmetric under an explicit scope (empty non-owner slots and zero-length flattened shards never create state), so after a consolidation-style resume that strips the common codebook, state-bearing members computed reuse=True and skipped the header all_gather while empty members computed reuse=False and issued it, mismatching the group's collective schedules. Resolve one group-wide decision before any member branches on it: _scope_agreed_resuming_from_checkpoint all_reduces (MAX) the local predicate over the codebook scope, so any member that restored periods makes the whole scope reuse them. Both call sites are reached by every member of a multi-member scope in the same order, and empty members are unaffected behaviorally by reuse=True because parameters without gradients or elements never enter the period iterator. --- src/gefen/gefen.py | 26 ++++++++++++++++++++++++-- 1 file changed, 24 insertions(+), 2 deletions(-) diff --git a/src/gefen/gefen.py b/src/gefen/gefen.py index c8e7df1..8b94763 100644 --- a/src/gefen/gefen.py +++ b/src/gefen/gefen.py @@ -5708,6 +5708,28 @@ def _resuming_from_checkpoint(self) -> bool: for p in group["params"] ) + def _scope_agreed_resuming_from_checkpoint(self) -> bool: + # Per-parameter state presence is legitimately rank-asymmetric under an + # explicit scope (empty non-owner slots and zero-length shards never + # create state), so the rank-local _resuming_from_checkpoint() predicate + # must not gate collectives directly: resolve one group-wide decision + # before any member branches on it. Any member that restored periods + # makes the whole scope reuse them. + resuming = self._resuming_from_checkpoint() + binding = self._gefen_codebook_process_group + if binding is None or len(binding.identity.ordered_members) == 1: + return resuming + self._assert_runtime_codebook_process_group() + import torch.distributed as dist + + control = torch.tensor( + int(resuming), + dtype=torch.int32, + device=binding.collective_device, + ) + dist.all_reduce(control, op=dist.ReduceOp.MAX, group=binding.process_group) + return bool(int(control.item())) + def _maybe_refresh_gefen_codebook(self) -> None: if self._gefen_codebook is not None: # A codebook restored from a checkpoint may land on CPU (depending on @@ -5739,7 +5761,7 @@ def _maybe_refresh_gefen_codebook(self) -> None: # the refreshed periods desync from the restored vmean/m_codebook block # geometry and the vmean kernel aborts on a block-count mismatch. self._ensure_gefen_codebook( - reuse_existing_periods=self._resuming_from_checkpoint() + reuse_existing_periods=self._scope_agreed_resuming_from_checkpoint() ) @torch.no_grad() @@ -5765,7 +5787,7 @@ def initialize_codebook(self) -> bool: if self._gefen_codebook is not None: return False self._ensure_gefen_codebook( - reuse_existing_periods=self._resuming_from_checkpoint() + reuse_existing_periods=self._scope_agreed_resuming_from_checkpoint() ) return self._gefen_codebook is not None From 22e81cbae6298450cf332badcc7dbbafa8f438ec Mon Sep 17 00:00:00 2001 From: thad0ctor Date: Mon, 13 Jul 2026 08:14:42 -0700 Subject: [PATCH 25/52] Re-validate codebook scope agreement on every member together _ensure_codebook_scope_agreement early-returned on the rank-local _gefen_codebook_scope_validated flag. Collective-free rank-local operations (move_state_, offload_state_, staged and native checkpoint loads) reset that flag on the member that ran them only, while every value the step operation header fingerprints stays identical (the codebook fingerprint is computed on CPU bytes and is device independent). One member then entered the agreement all_gather while the others early-returned, diverging the scoped collective schedule directly behind a passing header exchange. Fold the flag into the always-exchanged operation header as a trailing decision bit that is excluded from the equality check, and lower the flag on every member when any member reports it cleared, so the group re-validates together and then proceeds with identical collectives. This resolution was chosen over the two alternatives deliberately: raising on a flag mismatch would turn the advertised collective-free rank-local operations into scope-wide faults on legitimately asymmetric use, and dropping the reset for movement/offload would keep the divergence for staged and native loads, which genuinely can change scope-agreement-relevant state and must keep resetting the flag. --- src/gefen/gefen.py | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) diff --git a/src/gefen/gefen.py b/src/gefen/gefen.py index 8b94763..3ed4330 100644 --- a/src/gefen/gefen.py +++ b/src/gefen/gefen.py @@ -5144,17 +5144,28 @@ def _validate_codebook_scope_operation_header(self, operation: str) -> None: *self._codebook_scope_fingerprint(), *self._codebook_manifest_fingerprint(), *self._codebook_value_fingerprint(), + # Trailing decision bit, deliberately excluded from the equality + # check below: collective-free rank-local operations + # (move_state_, offload_state_, staged/native checkpoint loads) + # legitimately reset _gefen_codebook_scope_validated on a subset + # of members without changing any fingerprinted value. The group + # resolves "does any member need re-validation" here so + # _ensure_codebook_scope_agreement re-validates on every member + # together instead of diverging in front of its collectives. + int(self._gefen_codebook_scope_validated), ], dtype=torch.int64, device=binding.collective_device, ) headers = [torch.empty_like(header) for _ in binding.identity.ordered_members] dist.all_gather(headers, header, group=binding.process_group) - if any(not torch.equal(item, headers[0]) for item in headers[1:]): + if any(not torch.equal(item[:-1], headers[0][:-1]) for item in headers[1:]): raise RuntimeError( "scoped Gefen codebook operation, step, scope, manifest, policy, " "or old codebook differs across process-group members" ) + if any(int(item[-1].item()) == 0 for item in headers): + self._gefen_codebook_scope_validated = False if operation != "step" and self._gefen_codebook is not None: self._verify_codebook_scope_agreement(self._gefen_codebook) From 5aea07a7ef9441ae60bbb9518421b5d09b1e3b8f Mon Sep 17 00:00:00 2001 From: thad0ctor Date: Mon, 13 Jul 2026 08:14:59 -0700 Subject: [PATCH 26/52] Include parameter storage identity in canonical import live token _canonical_import_live_token identified locally bound parameters by id() only, so commit_canonical_state_import could not detect that a parameter's storage was retargeted (or mutated in place) between prepare and commit, even though the staged shadow was device-cast and geometry-validated against the parameters as they existed at prepare time. The sibling portable path already folds _parameter_storage_token (device, dtype, layout, shape, stride, storage pointer/size, version) into its live token for every local binding; mirror it here so a stale prepared import is refused with the existing freshness error. --- src/gefen/gefen.py | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) diff --git a/src/gefen/gefen.py b/src/gefen/gefen.py index 3ed4330..e450386 100644 --- a/src/gefen/gefen.py +++ b/src/gefen/gefen.py @@ -6828,6 +6828,8 @@ def _canonical_value_token(value): return (type(value).__name__, token) def _canonical_import_live_token(self): + from gefen.portable_runtime import _parameter_storage_token + groups = tuple( ( id(group), @@ -6885,7 +6887,16 @@ def _canonical_import_live_token(self): for logical_slot in self._gefen_logical_slots ), tuple( - (id(parameter), shard.sort_key) + ( + id(parameter), + shard.sort_key, + # Mirror the portable live token: a prepared import must go + # stale when a locally bound parameter's storage is + # retargeted or mutated between prepare and commit. + None + if parameter is None + else _parameter_storage_token(parameter), + ) for parameter, shard in self._gefen_local_shard_bindings ), ) From b25756a683313a10815fc03458946a5d9a63fb7a Mon Sep 17 00:00:00 2001 From: thad0ctor Date: Mon, 13 Jul 2026 08:15:08 -0700 Subject: [PATCH 27/52] Add regression tests for scope-agreed decisions and import freshness Cover the three review findings: a consolidation-style resume with a whole-parameter owner and an empty non-owner must gate the conditional initialize header on one group-wide reuse decision and then complete the resume step collectively with restored periods intact; a collective-free move_state_ on one member only must lead every member to the same scope re-validation decision on the next step; and a prepared canonical import must go stale when a bound parameter's storage is retargeted or mutated between prepare and commit, while an undisturbed prepare/commit round still succeeds. --- .../test_scoped_collective_agreement_fixes.py | 420 ++++++++++++++++++ 1 file changed, 420 insertions(+) create mode 100644 tests/test_scoped_collective_agreement_fixes.py diff --git a/tests/test_scoped_collective_agreement_fixes.py b/tests/test_scoped_collective_agreement_fixes.py new file mode 100644 index 0000000..8c50a71 --- /dev/null +++ b/tests/test_scoped_collective_agreement_fixes.py @@ -0,0 +1,420 @@ +"""Scope-agreed codebook decisions and canonical import parameter freshness.""" + +from datetime import timedelta +import multiprocessing as mp +import os +import tempfile + +import pytest +import torch +import torch.distributed as dist + +from gefen import ( + CodebookProcessGroupBinding, + Gefen, + GefenMuon, + LogicalSlice, + ParameterIdentity, + ParameterLayout, + ParameterRebinding, + PlacementKind, + ProcessGroupIdentity, + ShardIdentity, + ShardPlacement, + ShardingManifest, +) + + +_WORLD = 2 + + +def _replicated(parameter, group, member): + coordinate = group.ordered_members.index(member) + return ShardIdentity( + parameter, + ParameterLayout.REPLICATED, + LogicalSlice.full(parameter), + process_group=group, + local_member=member, + placements=( + ShardPlacement( + "dp", + PlacementKind.REPLICATE, + coordinate, + len(group.ordered_members), + ), + ), + ) + + +def _whole_owner(parameter, group, member, owner): + coordinate = group.ordered_members.index(member) + return ShardIdentity( + parameter, + ParameterLayout.WHOLE_PARAMETER_OWNER, + LogicalSlice.full(parameter) if member == owner else LogicalSlice(0, 0), + process_group=group, + local_member=member, + owner=owner, + placements=( + ShardPlacement( + "dp", + PlacementKind.WHOLE_PARAMETER_OWNER, + coordinate, + len(group.ordered_members), + ), + ), + ) + + +def _init_worker(rank, init_file): + dist.init_process_group( + "gloo", + init_method="file://{}".format(init_file), + rank=rank, + world_size=_WORLD, + timeout=timedelta(seconds=90), + ) + members = tuple("rank:{}".format(index) for index in range(_WORLD)) + return ProcessGroupIdentity("data_parallel", members), members + + +def _reuse_decision_worker(rank, init_file, queue): + """Rank-asymmetric automatic_period presence must not gate collectives. + + Reproduces the documented FSDP optim-state consolidation resume: the + loader restores per-parameter state but strips the optimizer-common + codebook on every member. A whole-parameter owner then holds restored + periods while empty non-owner members hold no state at all, so the + rank-local _resuming_from_checkpoint() predicate diverges. The group + must still resolve ONE reuse_existing_periods decision before the + conditionally issued "initialize" operation header. + """ + + try: + group, members = _init_worker(rank, init_file) + runtime_group = dist.group.WORLD + + owner_source = torch.nn.Parameter(torch.ones(2, 2)) + optimizer = GefenMuon([("matrix", owner_source)], fused=False) + identity = ParameterIdentity("Matrix", (2, 2)) + records = tuple( + _whole_owner(identity, group, member, "rank:0") for member in members + ) + optimizer.post_sharding( + ( + ParameterRebinding( + owner_source, + owner_source if rank == 0 else None, + records[rank], + ), + ), + manifest=ShardingManifest(records), + codebook_process_group=CodebookProcessGroupBinding( + group, members[rank], runtime_group, torch.device("cpu") + ), + ) + if rank == 0: + owner_source.grad = torch.tensor([[1.0, -2.0], [3.0, -4.0]]) + optimizer.step() + first_period = ( + int(optimizer.state[owner_source]["automatic_period"]) + if rank == 0 + else None + ) + state_asymmetric = ( + bool(optimizer.state) if rank == 0 else not optimizer.state + ) + + # Simulate the consolidation loader on every member: per-parameter + # state survives, the optimizer-common codebook does not. + optimizer._gefen_codebook = None + optimizer._gefen_codebook_by_device.clear() + optimizer._gefen_codebook_lut_by_device.clear() + optimizer._gefen_codebook_scope_validated = False + + if rank == 0: + owner_source.grad = torch.tensor([[0.5, -1.5], [2.5, -3.5]]) + + # Capture the reuse decision each member would gate collectives on, + # then abort symmetrically before any divergent collective runs. + captured = [] + + def capture(reuse_existing_periods=False): + captured.append(bool(reuse_existing_periods)) + raise RuntimeError("captured reuse decision") + + optimizer._ensure_gefen_codebook = capture + try: + optimizer.step() + raise AssertionError("sentinel resume step unexpectedly completed") + except RuntimeError as exc: + if "captured reuse decision" not in str(exc): + raise + finally: + del optimizer._ensure_gefen_codebook + + decisions = [None] * _WORLD + dist.all_gather_object(decisions, captured[0]) + + resume_completed = False + codebook_agreement = False + period_preserved = False + if len(set(decisions)) == 1: + # The schedule is agreed; the real resume step must complete + # collectively and keep the restored periods on the owner. + step_before = optimizer._gefen_global_step + optimizer.step() + resume_completed = ( + optimizer._gefen_codebook is not None + and optimizer._gefen_global_step == step_before + 1 + ) + codebooks = [ + torch.empty_like(optimizer._gefen_codebook) for _ in range(_WORLD) + ] + dist.all_gather(codebooks, optimizer._gefen_codebook) + codebook_agreement = all( + torch.equal(item, codebooks[0]) for item in codebooks[1:] + ) + if rank == 0: + period_preserved = ( + int(optimizer.state[owner_source]["automatic_period"]) + == first_period + ) + else: + period_preserved = not optimizer.state + queue.put( + { + "rank": rank, + "state_asymmetric": state_asymmetric, + "decisions": decisions, + "resume_completed": resume_completed, + "codebook_agreement": codebook_agreement, + "period_preserved": period_preserved, + } + ) + except Exception as exc: + queue.put({"rank": rank, "error": repr(exc)}) + finally: + if dist.is_available() and dist.is_initialized(): + dist.destroy_process_group() + + +def _revalidation_decision_worker(rank, init_file, queue): + """A rank-local scope-validated reset must re-validate on every member. + + move_state_ is collective-free and resets _gefen_codebook_scope_validated + on the member that called it only, while every fingerprint the step + header exchanges stays identical. The next step must make the SAME + early-return decision inside _ensure_codebook_scope_agreement on every + member. + """ + + try: + group, members = _init_worker(rank, init_file) + runtime_group = dist.group.WORLD + + parameter = torch.nn.Parameter(torch.zeros(4)) + optimizer = Gefen([("weight", parameter)], fused=False, factored_v_2d=False) + identity = ParameterIdentity("Weight", (4,)) + records = tuple(_replicated(identity, group, member) for member in members) + optimizer.post_sharding( + (ParameterRebinding(parameter, parameter, records[rank]),), + manifest=ShardingManifest(records), + codebook_process_group=CodebookProcessGroupBinding( + group, members[rank], runtime_group, torch.device("cpu") + ), + ) + optimizer._resolve_automatic_period = lambda *args: 4 + parameter.grad = torch.tensor([-1.0, -0.5, 0.25, 1.0]) + optimizer.step() + validated_after_first = bool(optimizer._gefen_codebook_scope_validated) + + # Advertised collective-free state movement on ONE member only; every + # state value stays bit-identical, so the step header fingerprints + # still agree across members. + if rank == 0: + optimizer.move_state_(torch.device("cpu")) + reset_asymmetric = ( + not optimizer._gefen_codebook_scope_validated + if rank == 0 + else bool(optimizer._gefen_codebook_scope_validated) + ) + + parameter.grad = torch.tensor([0.75, -0.25, 0.5, -1.0]) + captured = [] + + def capture(): + captured.append(not optimizer._gefen_codebook_scope_validated) + raise RuntimeError("captured revalidation decision") + + optimizer._ensure_codebook_scope_agreement = capture + try: + optimizer.step() + raise AssertionError("sentinel step unexpectedly completed") + except RuntimeError as exc: + if "captured revalidation decision" not in str(exc): + raise + finally: + del optimizer._ensure_codebook_scope_agreement + + decisions = [None] * _WORLD + dist.all_gather_object(decisions, captured[0]) + + second_completed = False + revalidated = False + parameters_agree = False + if len(set(decisions)) == 1: + optimizer.step() + second_completed = True + revalidated = bool(optimizer._gefen_codebook_scope_validated) + gathered = [torch.empty_like(parameter.detach()) for _ in range(_WORLD)] + dist.all_gather(gathered, parameter.detach()) + parameters_agree = all( + torch.equal(item, gathered[0]) for item in gathered[1:] + ) + queue.put( + { + "rank": rank, + "validated_after_first": validated_after_first, + "reset_asymmetric": reset_asymmetric, + "decisions": decisions, + "second_completed": second_completed, + "revalidated": revalidated, + "parameters_agree": parameters_agree, + } + ) + except Exception as exc: + queue.put({"rank": rank, "error": repr(exc)}) + finally: + if dist.is_available() and dist.is_initialized(): + dist.destroy_process_group() + + +def _run_workers(target): + context = mp.get_context("spawn") + queue = context.Queue() + fd, init_file = tempfile.mkstemp(prefix="gefen-scope-agreement-") + os.close(fd) + os.unlink(init_file) + processes = [ + context.Process(target=target, args=(rank, init_file, queue)) + for rank in range(_WORLD) + ] + try: + for process in processes: + process.start() + results = [queue.get(timeout=120) for _ in processes] + for process in processes: + process.join(timeout=10) + for process in processes: + if process.is_alive(): + process.terminate() + process.join(timeout=5) + pytest.fail("scoped agreement worker hung") + assert process.exitcode == 0 + return sorted(results, key=lambda item: item["rank"]) + finally: + for process in processes: + if process.is_alive(): + process.terminate() + process.join(timeout=5) + if os.path.exists(init_file): + os.unlink(init_file) + + +def test_resume_reuse_decision_is_scope_agreed_for_empty_members(): + results = _run_workers(_reuse_decision_worker) + + assert all("error" not in item for item in results), results + for item in results: + assert item["state_asymmetric"], item + # Every member must gate the conditional "initialize" header on the + # same group-wide decision; a resumed scope reuses restored periods. + assert len(set(item["decisions"])) == 1, item + assert all(item["decisions"]), item + assert item["resume_completed"], item + assert item["codebook_agreement"], item + assert item["period_preserved"], item + + +def test_rank_local_scope_validated_reset_revalidates_on_every_member(): + results = _run_workers(_revalidation_decision_worker) + + assert all("error" not in item for item in results), results + for item in results: + assert item["validated_after_first"], item + assert item["reset_asymmetric"], item + # Every member must make the same re-validation decision after a + # collective-free rank-local reset. + assert len(set(item["decisions"])) == 1, item + assert item["second_completed"], item + assert item["revalidated"], item + assert item["parameters_agree"], item + + +def _build_finalized_canonical(): + parameter = torch.nn.Parameter(torch.linspace(-0.5, 0.7, 8)) + optimizer = Gefen( + [("layer.weight", parameter)], + fused=False, + factored_v_2d=False, + ) + optimizer.rebind_parameter( + parameter, + parameter, + identity=ParameterIdentity("Layer.Weight", (8,)), + ) + optimizer._resolve_automatic_period = lambda *args, **kwargs: 4 + return optimizer, parameter + + +def _exported_canonical_state(): + source, source_parameter = _build_finalized_canonical() + source_parameter.grad = torch.linspace(-1.0, 0.8, 8) + source.step() + return source.export_canonical_state() + + +def test_parameter_storage_retarget_invalidates_prepared_canonical_import(): + exported = _exported_canonical_state() + target, target_parameter = _build_finalized_canonical() + + prepared = target.prepare_canonical_state_import(exported) + with torch.no_grad(): + target_parameter.data = target_parameter.data.clone() + + with pytest.raises( + RuntimeError, match="changed after canonical import preparation" + ): + target.commit_canonical_state_import(prepared) + + # The refusal must keep canonical I/O available with a fresh preparation. + target.import_canonical_state(exported) + assert "automatic_period" in target.state[target_parameter] + + +def test_parameter_inplace_mutation_invalidates_prepared_canonical_import(): + exported = _exported_canonical_state() + target, target_parameter = _build_finalized_canonical() + + prepared = target.prepare_canonical_state_import(exported) + with torch.no_grad(): + target_parameter.mul_(2.0) + + with pytest.raises( + RuntimeError, match="changed after canonical import preparation" + ): + target.commit_canonical_state_import(prepared) + + target.import_canonical_state(exported) + assert "automatic_period" in target.state[target_parameter] + + +def test_undisturbed_prepared_canonical_import_still_commits(): + exported = _exported_canonical_state() + target, target_parameter = _build_finalized_canonical() + + prepared = target.prepare_canonical_state_import(exported) + target.commit_canonical_state_import(prepared) + assert "automatic_period" in target.state[target_parameter] From c0d6509fdeb7349becc0d44d312fd641b9e151f8 Mon Sep 17 00:00:00 2001 From: thad0ctor Date: Mon, 13 Jul 2026 08:20:51 -0700 Subject: [PATCH 28/52] Synchronize hybrid step AMP and preflight failures across the codebook scope GefenMuonHybrid.step decided the GradScaler overflow skip with the rank-local _amp_prepare_optimizer_step and raised structural gradient preflight errors rank-locally, even when post_sharding installed one multi-member codebook process-group binding shared by both children. A rank whose found_inf was set (or whose local gradients were malformed) then skipped or raised alone while its peers entered the children's scoped step collectives, hanging the group. Route both decisions through the children's scoped protocol before any child steps: the composite runs Gefen._prepare_scoped_amp_optimizer_step on itself (found_inf/grad_scale agreement is validated collectively and an overflow skip is a group-wide decision, entered and exited symmetrically on every member) and synchronizes the structural preflight through the children's _synchronize_codebook_scope_failure so every scope member raises together, keeping the atomic both-children-skip semantics. Without a binding the local behavior is unchanged. --- src/gefen/hybrid.py | 68 ++- tests/test_hybrid_scoped_failure_protocol.py | 602 +++++++++++++++++++ 2 files changed, 663 insertions(+), 7 deletions(-) create mode 100644 tests/test_hybrid_scoped_failure_protocol.py diff --git a/src/gefen/hybrid.py b/src/gefen/hybrid.py index a2757e6..a8f3592 100644 --- a/src/gefen/hybrid.py +++ b/src/gefen/hybrid.py @@ -74,7 +74,6 @@ from gefen.gefen import ( Gefen, _amp_native_scaling_required, - _amp_prepare_optimizer_step, _assert_optimizer_gradients_structurally_valid, ) from gefen.gefen_muon import GefenMuon @@ -1091,6 +1090,40 @@ def codebook_process_group_binding(self): self._assert_finalized_binding_layout() return self._hybrid_codebook_process_group + @property + def _gefen_codebook_process_group(self): + # Read-only alias: the scoped-step protocol borrowed from Gefen + # (_prepare_scoped_amp_optimizer_step and the failure synchronization + # it drives) looks the binding up under the child attribute name. The + # composite's one shared binding lives in + # _hybrid_codebook_process_group; nothing may store under this name. + return self._hybrid_codebook_process_group + + def _synchronize_codebook_scope_failure(self, error, phase: str) -> None: + # Composite-level preflight/AMP failures synchronize on the one + # binding shared by every child, so every scope member raises + # together instead of stranding peers inside a child's scoped step + # collectives. The finalized layout guarantees the first child holds + # the identical binding; without a binding, the local error is raised + # unchanged, exactly as the children do. + binding = self._hybrid_codebook_process_group + if binding is None: + if error is not None: + raise error + return + self._subopts[0]._synchronize_codebook_scope_failure(error, phase) + + def _prepare_scoped_amp_optimizer_step(self) -> bool: + # GradScaler attaches found_inf/grad_scale to the composite, never to + # the children, so the children's own scoped AMP gates cannot see an + # overflow. Run the base scoped protocol once with the composite as + # the optimizer: under a multi-member codebook binding it validates + # found_inf/grad_scale agreement collectively and makes an overflow + # skip a group-wide decision, and on a finite step it unscales the + # union of both children's gradients exactly once. Without a binding + # this is exactly _amp_prepare_optimizer_step(self). + return Gefen._prepare_scoped_amp_optimizer_step(self) + def zero_grad(self, set_to_none: bool = True): self._assert_finalized_binding_layout() for o in self._subopts: @@ -1148,13 +1181,34 @@ def step(self, closure=None): with torch.enable_grad(): loss = closure() self._assert_finalized_binding_layout() - for child in self._subopts: - _assert_optimizer_gradients_structurally_valid(child, require_2d_params=child is self.muon) + # Composite structural preflight, atomically over BOTH children before + # either child steps. Under a shared codebook scope the failure is + # synchronized on the binding first (mirroring Gefen.step and + # GefenMuon.step), so a rank-local structural error raises on every + # scope member together instead of stranding peers inside a child's + # scoped step collectives. + try: + for child in self._subopts: + _assert_optimizer_gradients_structurally_valid(child, require_2d_params=child is self.muon) + local_preflight_error = None + except Exception as exc: + local_preflight_error = exc + if self._hybrid_codebook_process_group is not None: + self._synchronize_codebook_scope_failure( + local_preflight_error, "gradient preflight" + ) + elif local_preflight_error is not None: + raise local_preflight_error # A non-finite gradient in either half skips BOTH children before their - # codebooks, states, counters, or parameters can move. Explicit - # scaler.unscale_(hybrid) is detected by grad_scale=None and is not - # repeated; automatic unscale covers every child parameter exactly once. - if (hasattr(self, "found_inf") or hasattr(self, "grad_scale")) and not _amp_prepare_optimizer_step(self): + # codebooks, states, counters, or parameters can move. GradScaler + # attaches found_inf/grad_scale to the composite, so under a shared + # multi-member codebook scope the overflow skip is the children's + # scoped AMP protocol run here -- collective found_inf/grad_scale + # agreement, then a group-wide skip -- before any child enters its + # scoped step collectives. Explicit scaler.unscale_(hybrid) is + # detected by grad_scale=None and is not repeated; automatic unscale + # covers every child parameter exactly once. + if (hasattr(self, "found_inf") or hasattr(self, "grad_scale")) and not self._prepare_scoped_amp_optimizer_step(): for post_hook in self._optimizer_step_post_hooks.values(): post_hook(self, args, kwargs) return loss diff --git a/tests/test_hybrid_scoped_failure_protocol.py b/tests/test_hybrid_scoped_failure_protocol.py new file mode 100644 index 0000000..4daa84f --- /dev/null +++ b/tests/test_hybrid_scoped_failure_protocol.py @@ -0,0 +1,602 @@ +"""Scoped failure-protocol coverage for GefenMuonHybrid.step. + +GradScaler attaches found_inf/grad_scale to the composite (never to the +children), and the hybrid's structural gradient preflight runs before either +child's scoped step collectives. Under the one multi-member codebook binding +shared by both children, both decisions must therefore be group decisions: + +* a rank-local AMP overflow must not skip the children on one member while the + peers enter the children's scoped step collectives -- found_inf/grad_scale + agreement is validated collectively and an overflow skip is entered/exited + symmetrically on every member; +* a rank-local structural preflight failure must raise on every scope member + together via the children's synchronized failure protocol instead of + stranding peers inside a child's scope header all_gather. + +Single-process behavior (no binding, and the collective-free one-member +binding) keeps the local semantics. +""" + +from datetime import timedelta +import multiprocessing as mp +import os +import queue as queue_module +import tempfile +import traceback + +import pytest +import torch +import torch.distributed as dist + +from gefen import GefenMuonHybrid +from gefen.codebook import CodebookProcessGroupBinding +from gefen.contracts import ( + LogicalSlice, + ParameterIdentity, + ParameterLayout, + PlacementKind, + ProcessGroupIdentity, + ShardIdentity, + ShardPlacement, + ShardingManifest, +) +from gefen.rebinding import ParameterRebinding + + +_WORLD = 2 +_MUON_FQN = "model.matrix" +_BACKUP_FQN = "model.vector" +_MUON_OWNER = "rank:1" +_BACKUP_LENGTHS = (3, 5) + + +def _members(): + return tuple("rank:{}".format(rank) for rank in range(_WORLD)) + + +def _owner_shards(identity, group, owner): + return tuple( + ShardIdentity( + identity, + ParameterLayout.WHOLE_PARAMETER_OWNER, + (LogicalSlice.full(identity) if member == owner else LogicalSlice(0, 0)), + placements=( + ShardPlacement( + "dp", + PlacementKind.WHOLE_PARAMETER_OWNER, + coordinate, + len(group.ordered_members), + ), + ), + process_group=group, + local_member=member, + owner=owner, + ) + for coordinate, member in enumerate(group.ordered_members) + ) + + +def _flat_shards(identity, group, lengths): + offset = 0 + shards = [] + for coordinate, (member, length) in enumerate(zip(group.ordered_members, lengths)): + shards.append( + ShardIdentity( + identity, + ParameterLayout.FLATTENED_ELEMENT_SHARD, + LogicalSlice(offset, length), + placements=( + ShardPlacement( + "dp", + PlacementKind.FLAT_SHARD, + coordinate, + len(group.ordered_members), + ), + ), + process_group=group, + local_member=member, + ) + ) + offset += length + return tuple(shards) + + +def _muon_initial(): + return torch.linspace(-0.3, 0.2, 6, dtype=torch.float32).reshape(3, 2) + + +def _backup_initial(): + return torch.linspace(-0.4, 0.3, 8, dtype=torch.float32) + + +def _muon_gradient(): + return torch.tensor( + [[0.4, -0.7], [1.1, -1.3], [1.7, -1.9]], + dtype=torch.float32, + ) + + +def _backup_gradient(): + return torch.tensor( + [0.25, -0.5, 0.75, -1.0, 1.25, -1.5, 1.75, -2.0], + dtype=torch.float32, + ) + + +def _codebook(): + return torch.linspace(-1.0, 1.0, 256, dtype=torch.float32) + + +def _quantized_period_one(momentum): + flat = momentum.reshape(-1) + indices = torch.where( + torch.signbit(flat), + torch.zeros(flat.numel(), dtype=torch.uint8), + torch.full((flat.numel(),), 255, dtype=torch.uint8), + ) + return indices.reshape(-1, 1), flat.abs().reshape(-1, 1).clone() + + +def _hybrid_kwargs(): + return dict( + lr=2.5e-3, + muon_lr=3.0e-3, + backup_lr=2.5e-3, + weight_decay=0.03, + muon_weight_decay=0.02, + backup_weight_decay=0.03, + backup_optimizer="gefen", + backup_1d_period_one=True, + betas=(0.8, 0.97), + eps=2.0e-8, + fused=False, + momentum=0.85, + nesterov=False, + ns_steps=2, + deterministic=True, + normuon=True, + normuon_beta2=0.9, + normuon_eps=3.0e-8, + ) + + +def _make_scoped_hybrid(rank, group): + muon_identity = ParameterIdentity(_MUON_FQN, (3, 2)) + backup_identity = ParameterIdentity(_BACKUP_FQN, (8,)) + old_muon = torch.nn.Parameter(_muon_initial().clone()) + old_backup = torch.nn.Parameter(_backup_initial().clone()) + optimizer = GefenMuonHybrid( + [("matrix", old_muon)], + [("vector", old_backup)], + sharded_mode="distributed", + **_hybrid_kwargs(), + ) + muon_shards = _owner_shards(muon_identity, group, _MUON_OWNER) + backup_shards = _flat_shards(backup_identity, group, _BACKUP_LENGTHS) + muon_local = old_muon if _members()[rank] == _MUON_OWNER else None + backup_shard = backup_shards[rank] + start = backup_shard.logical_slice.flat_offset + stop = start + backup_shard.logical_slice.length + backup_local = torch.nn.Parameter(_backup_initial()[start:stop].clone()) + binding = CodebookProcessGroupBinding( + group, + _members()[rank], + dist.group.WORLD, + torch.device("cpu"), + ) + optimizer.post_sharding( + ( + ParameterRebinding(old_muon, muon_local, muon_shards[rank]), + ParameterRebinding(old_backup, backup_local, backup_shard), + ), + manifest=ShardingManifest(muon_shards + backup_shards), + codebook_process_group=binding, + ) + return optimizer, muon_local, backup_local, backup_shard + + +def _seed_hybrid_state(optimizer, muon_parameter, backup_parameter, backup_shard): + for child in (optimizer.muon, optimizer.backup): + child._gefen_global_step = 13 + child._gefen_codebook = _codebook() + if muon_parameter is not None: + momentum = torch.tensor( + [[-0.75, 0.5], [1.25, -1.5], [2.0, -2.5]], + dtype=torch.float32, + ) + indices, magnitudes = _quantized_period_one(momentum) + optimizer.muon.state[muon_parameter].update( + { + "automatic_period": 1, + "step": 11, + "m_codebook": indices, + "m_magnitude": magnitudes, + "normuon_v": torch.tensor([[0.5], [1.5], [2.5]], dtype=torch.float32), + "normuon_step": 10, + } + ) + start = backup_shard.logical_slice.flat_offset + stop = start + backup_shard.logical_slice.length + momentum = torch.tensor( + [-0.25, 0.5, -0.75, 1.0, -1.25, 1.5, -1.75, 2.0], + dtype=torch.float32, + )[start:stop] + indices, magnitudes = _quantized_period_one(momentum) + optimizer.backup.state[backup_parameter].update( + { + "automatic_period": 1, + "step": 11, + "m_codebook": indices, + "m_magnitude": magnitudes, + "vmean": torch.linspace(0.2, 0.9, 8, dtype=torch.float32)[start:stop].reshape(-1, 1).clone(), + "vmean_step": 10, + } + ) + + +def _bits_equal(left, right): + return ( + torch.is_tensor(left) + and torch.is_tensor(right) + and left.dtype == right.dtype + and tuple(left.shape) == tuple(right.shape) + and torch.equal( + left.detach().contiguous().view(torch.uint8), + right.detach().contiguous().view(torch.uint8), + ) + ) + + +def _set_local_gradients(muon_parameter, backup_parameter, backup_shard, scale=1.0): + if muon_parameter is not None: + muon_parameter.grad = _muon_gradient() * scale + start = backup_shard.logical_slice.flat_offset + stop = start + backup_shard.logical_slice.length + backup_parameter.grad = _backup_gradient()[start:stop] * scale + + +def _untouched(optimizer, muon_parameter, backup_parameter): + return ( + optimizer.muon._gefen_global_step == 0 + and optimizer.backup._gefen_global_step == 0 + and optimizer.muon._gefen_codebook is None + and optimizer.backup._gefen_codebook is None + and (muon_parameter is None or _bits_equal(muon_parameter, _muon_initial())) + and ( + muon_parameter is None + or optimizer.muon.state[muon_parameter] == {"name": "matrix"} + ) + and optimizer.backup.state[backup_parameter] == {"name": "vector"} + ) + + +def _amp_divergent_overflow_result(rank, group): + optimizer, muon_parameter, backup_parameter, backup_shard = _make_scoped_hybrid(rank, group) + _set_local_gradients(muon_parameter, backup_parameter, backup_shard) + grad_before = backup_parameter.grad.detach().clone() + optimizer.found_inf = torch.tensor(float(rank == 0)) + optimizer.grad_scale = torch.tensor(8.0) + try: + optimizer.step() + message = None + except RuntimeError as exc: + message = str(exc) + return { + "message": message, + "untouched": _untouched(optimizer, muon_parameter, backup_parameter), + "grads_untouched": _bits_equal(backup_parameter.grad, grad_before), + } + + +def _amp_group_wide_overflow_result(rank, group): + optimizer, muon_parameter, backup_parameter, backup_shard = _make_scoped_hybrid(rank, group) + _set_local_gradients(muon_parameter, backup_parameter, backup_shard) + grad_before = backup_parameter.grad.detach().clone() + optimizer.found_inf = torch.tensor(1.0) + optimizer.grad_scale = torch.tensor(8.0) + post_hook_calls = [] + optimizer.register_step_post_hook( + lambda _optimizer, _args, _kwargs: post_hook_calls.append(True) + ) + try: + optimizer.step() + skipped = True + except RuntimeError: + skipped = False + return { + "skipped": skipped, + "post_hook_fired": len(post_hook_calls) == 1, + "untouched": _untouched(optimizer, muon_parameter, backup_parameter), + "grads_untouched": _bits_equal(backup_parameter.grad, grad_before), + } + + +def _amp_agreeing_finite_result(rank, group): + target, target_muon, target_backup, target_shard = _make_scoped_hybrid(rank, group) + reference, reference_muon, reference_backup, reference_shard = _make_scoped_hybrid(rank, group) + _seed_hybrid_state(target, target_muon, target_backup, target_shard) + _seed_hybrid_state(reference, reference_muon, reference_backup, reference_shard) + _set_local_gradients(target_muon, target_backup, target_shard, scale=8.0) + _set_local_gradients(reference_muon, reference_backup, reference_shard) + target.found_inf = torch.tensor(0.0) + target.grad_scale = torch.tensor(8.0) + target.step() + reference.step() + return { + "stepped": ( + target.muon._gefen_global_step == 14 + and target.backup._gefen_global_step == 14 + ), + "unscaled_once_exactly": ( + (target_muon is None or _bits_equal(target_muon, reference_muon)) + and _bits_equal(target_backup, reference_backup) + ), + } + + +def _preflight_divergent_result(rank, group): + optimizer, muon_parameter, backup_parameter, backup_shard = _make_scoped_hybrid(rank, group) + _set_local_gradients(muon_parameter, backup_parameter, backup_shard) + if rank == 0: + backup_parameter.grad = torch.sparse_coo_tensor( + torch.tensor([[0]]), + torch.tensor([1.0]), + size=(backup_shard.logical_slice.length,), + ) + try: + optimizer.step() + message = None + except RuntimeError as exc: + message = str(exc) + return { + "message": message, + "untouched": _untouched(optimizer, muon_parameter, backup_parameter), + } + + +def _distributed_worker(rank, init_file, result_queue): + try: + dist.init_process_group( + "gloo", + init_method="file://{}".format(init_file), + rank=rank, + world_size=_WORLD, + timeout=timedelta(seconds=45), + ) + group = ProcessGroupIdentity("hybrid_scoped_failure", _members()) + amp_divergent = _amp_divergent_overflow_result(rank, group) + dist.barrier() + amp_group_wide = _amp_group_wide_overflow_result(rank, group) + dist.barrier() + amp_finite = _amp_agreeing_finite_result(rank, group) + dist.barrier() + preflight = _preflight_divergent_result(rank, group) + dist.barrier() + result_queue.put( + { + "rank": rank, + "amp_divergent": amp_divergent, + "amp_group_wide": amp_group_wide, + "amp_finite": amp_finite, + "preflight": preflight, + } + ) + except BaseException: + result_queue.put({"rank": rank, "fatal_error": traceback.format_exc()}) + finally: + if dist.is_available() and dist.is_initialized(): + dist.destroy_process_group() + + +def _run_distributed_workers(): + context = mp.get_context("spawn") + result_queue = context.Queue() + descriptor, init_file = tempfile.mkstemp(prefix="gefen-hybrid-scoped-failure-") + os.close(descriptor) + os.unlink(init_file) + processes = [ + context.Process( + target=_distributed_worker, + args=(rank, init_file, result_queue), + ) + for rank in range(_WORLD) + ] + results = [] + try: + for process in processes: + process.start() + try: + for _ in processes: + results.append(result_queue.get(timeout=240)) + except queue_module.Empty: + pass + for process in processes: + process.join(timeout=10) + finally: + for process in processes: + if process.is_alive(): + process.terminate() + process.join(timeout=5) + result_queue.close() + result_queue.join_thread() + if os.path.exists(init_file): + os.unlink(init_file) + assert len(results) == _WORLD, (results, [process.exitcode for process in processes]) + return sorted(results, key=lambda item: item["rank"]) + + +@pytest.mark.skipif( + not dist.is_available() or not dist.is_gloo_available(), + reason="hybrid scoped failure-protocol coverage requires Gloo", +) +def test_hybrid_step_synchronizes_amp_and_preflight_across_the_scope(): + results = _run_distributed_workers() + assert all("fatal_error" not in result for result in results), results + + # Rank-divergent found_inf must raise the children's collective AMP + # agreement error on EVERY member (never a one-sided silent skip that + # strands the peer inside a child's scope header collective), leaving + # both children and the gradients untouched on both ranks. + divergent = [result["amp_divergent"] for result in results] + assert all( + item["message"] is not None and "group-aware gradient scaler" in item["message"] + for item in divergent + ), divergent + assert all(item["untouched"] and item["grads_untouched"] for item in divergent), divergent + + # A group-wide overflow skips BOTH children symmetrically on every + # member: step returns through the post-hook path with no mutation and + # no gradient unscale, and every member exits together (the barriers + # after this scenario would otherwise desynchronize). + group_wide = [result["amp_group_wide"] for result in results] + assert all( + item["skipped"] and item["post_hook_fired"] and item["untouched"] and item["grads_untouched"] + for item in group_wide + ), group_wide + + # An agreeing finite scale steps both children on every member and + # unscales the union of both children's gradients exactly once: the AMP + # step over 8x-scaled gradients is bit-identical to the plain step. + finite = [result["amp_finite"] for result in results] + assert all(item["stepped"] and item["unscaled_once_exactly"] for item in finite), finite + + # A rank-local structural preflight failure raises on every scope member + # together through the children's synchronized failure protocol, keeping + # the atomic both-children-skip semantics on both ranks. + preflight = [result["preflight"] for result in results] + assert preflight[0]["message"] is not None and preflight[1]["message"] is not None, preflight + assert "gradient preflight failed on local member rank:0" in preflight[0]["message"] + assert "sparse gradients" in preflight[0]["message"] + assert "gradient preflight failed on another process-group member" in preflight[1]["message"] + assert all(item["untouched"] for item in preflight), preflight + + +def _make_plain_hybrid(): + matrix = torch.nn.Parameter(_muon_initial().clone()) + vector = torch.nn.Parameter(_backup_initial().clone()) + optimizer = GefenMuonHybrid( + [("matrix", matrix)], + [("vector", vector)], + **_hybrid_kwargs(), + ) + return optimizer, matrix, vector + + +def test_unscoped_amp_overflow_and_finite_step_behavior_is_preserved(): + optimizer, matrix, vector = _make_plain_hybrid() + matrix.grad = _muon_gradient().clone() + vector.grad = _backup_gradient().clone() + optimizer.found_inf = torch.tensor(1.0) + optimizer.grad_scale = torch.tensor(4.0) + assert optimizer.step() is None + assert optimizer.muon._gefen_global_step == 0 + assert optimizer.backup._gefen_global_step == 0 + assert _bits_equal(matrix, _muon_initial()) + assert _bits_equal(vector, _backup_initial()) + assert _bits_equal(vector.grad, _backup_gradient()) + + matrix.grad = _muon_gradient() * 4.0 + vector.grad = _backup_gradient() * 4.0 + optimizer.found_inf = torch.tensor(0.0) + optimizer.grad_scale = torch.tensor(4.0) + optimizer.step() + assert optimizer.muon._gefen_global_step == 1 + assert optimizer.backup._gefen_global_step == 1 + assert not _bits_equal(matrix, _muon_initial()) + assert not _bits_equal(vector, _backup_initial()) + assert _bits_equal(vector.grad, _backup_gradient()) + + +def test_unscoped_preflight_failure_raises_locally_and_skips_both_children(): + optimizer, matrix, vector = _make_plain_hybrid() + matrix.grad = _muon_gradient().clone() + vector.grad = torch.sparse_coo_tensor( + torch.tensor([[0]]), + torch.tensor([1.0]), + size=(8,), + ) + with pytest.raises(RuntimeError) as excinfo: + optimizer.step() + assert "sparse gradients" in str(excinfo.value) + assert "scoped Gefen codebook" not in str(excinfo.value) + assert optimizer.muon._gefen_global_step == 0 + assert optimizer.backup._gefen_global_step == 0 + assert optimizer.muon._gefen_codebook is None + assert optimizer.backup._gefen_codebook is None + assert _bits_equal(matrix, _muon_initial()) + assert _bits_equal(vector, _backup_initial()) + + +def _make_one_member_scoped_hybrid(): + group = ProcessGroupIdentity("solo", ("rank:0",)) + muon_identity = ParameterIdentity(_MUON_FQN, (3, 2)) + backup_identity = ParameterIdentity(_BACKUP_FQN, (8,)) + matrix = torch.nn.Parameter(_muon_initial().clone()) + vector = torch.nn.Parameter(_backup_initial().clone()) + optimizer = GefenMuonHybrid( + [("matrix", matrix)], + [("vector", vector)], + **_hybrid_kwargs(), + ) + + def _replicated(identity): + return ShardIdentity( + identity, + ParameterLayout.REPLICATED, + LogicalSlice.full(identity), + placements=(ShardPlacement("dp", PlacementKind.REPLICATE, 0, 1),), + process_group=group, + local_member="rank:0", + ) + + muon_shard = _replicated(muon_identity) + backup_shard = _replicated(backup_identity) + binding = CodebookProcessGroupBinding(group, "rank:0", None, torch.device("cpu")) + optimizer.post_sharding( + ( + ParameterRebinding(matrix, matrix, muon_shard), + ParameterRebinding(vector, vector, backup_shard), + ), + manifest=ShardingManifest((muon_shard, backup_shard)), + codebook_process_group=binding, + ) + return optimizer, matrix, vector + + +def test_one_member_scope_amp_skip_and_preflight_raise_stay_collective_free(): + optimizer, matrix, vector = _make_one_member_scoped_hybrid() + matrix.grad = _muon_gradient().clone() + vector.grad = _backup_gradient().clone() + optimizer.found_inf = torch.tensor(1.0) + optimizer.grad_scale = torch.tensor(2.0) + assert optimizer.step() is None + assert optimizer.muon._gefen_global_step == 0 + assert optimizer.backup._gefen_global_step == 0 + assert _bits_equal(matrix, _muon_initial()) + assert _bits_equal(vector, _backup_initial()) + + vector.grad = torch.sparse_coo_tensor( + torch.tensor([[0]]), + torch.tensor([1.0]), + size=(8,), + ) + del optimizer.found_inf + del optimizer.grad_scale + # A one-member binding synchronizes without collectives and re-raises the + # local structural error unchanged, exactly like the children. + with pytest.raises(RuntimeError) as excinfo: + optimizer.step() + assert "sparse gradients" in str(excinfo.value) + assert optimizer.muon._gefen_global_step == 0 + assert optimizer.backup._gefen_global_step == 0 + + matrix.grad = _muon_gradient() * 2.0 + vector.grad = _backup_gradient() * 2.0 + optimizer.found_inf = torch.tensor(0.0) + optimizer.grad_scale = torch.tensor(2.0) + optimizer.step() + assert optimizer.muon._gefen_global_step == 1 + assert optimizer.backup._gefen_global_step == 1 + assert not _bits_equal(matrix, _muon_initial()) + assert not _bits_equal(vector, _backup_initial()) + assert _bits_equal(vector.grad, _backup_gradient()) From 16410fa180d52e5f06843843cf3168d91f299aec Mon Sep 17 00:00:00 2001 From: thad0ctor Date: Mon, 13 Jul 2026 08:47:05 -0700 Subject: [PATCH 29/52] Cache layout forensics behind O(params) step tokens step() re-ran the complete O(params x world) finalized-layout forensic rebuild on every guard call (2 passes per unscoped step, up to 7 under a multi-member codebook scope) and recomputed the manifest sha256 fingerprint inside every scoped operation header, costing seconds of host time per step at large scale. Cache one forensic verdict as an O(local params) identity-token snapshot (finalized registries by object identity, every live group container, parameter and compatibility name, plus a version counter bumped by every legitimate mutating API), and compute the manifest shard set and digest once per finalized manifest at post_sharding. Steady-state step guards now reuse the verdict; checkpoint prepare/commit, canonical export/import, rebinding, state movement/offload, codebook initialize/refresh, scope re-validation, and contract readiness still run the full forensic rebuild. The offload step-readiness scan is deduped under the same scheme. Caches live in __slots__ so staged __dict__ copies and fail-before-mutation snapshots never see them. Public-container tampering (group params/param_names slots, state names, the name cache), including closure-time mutation, still fails the step guard before any state mutation; in-place edits inside the private registries move to detection at the next full-forensics boundary, as now documented in docs/optimizer_contracts.md. --- docs/optimizer_contracts.md | 4 +- src/gefen/gefen.py | 236 ++++++++++++++++++++++++++++++++---- 2 files changed, 214 insertions(+), 26 deletions(-) diff --git a/docs/optimizer_contracts.md b/docs/optimizer_contracts.md index 7ff685e..46fa1ec 100644 --- a/docs/optimizer_contracts.md +++ b/docs/optimizer_contracts.md @@ -41,6 +41,8 @@ Rebinding is allowed only while the entire optimizer is pristine: global step ze Plain Gefen currently rebinds complete replicated parameters and contiguous physical 1-D flattened element shards. A flattened logical matrix requires `factored_v_2d=False` because canonical row/column factored-state projection is not implemented. GefenMuon rebinds complete replicated matrices and can finalize/prune whole-parameter owner manifests. A whole-owner instance may step only when the same transaction installs the explicit codebook process group described below; the owner performs the update and the machine-readable training declaration requires the adapter to synchronize the updated matrix afterward. A Gefen-backed `GefenMuonHybrid` atomically partitions one complete manifest and rebinding plan by its frozen exact FQN routing, stages both children, validates cross-child storage disjointness, rebuilds composite state routing, and publishes only after every child succeeds. AdamW-backed Hybrid and DTensor composite rebinding remain unsupported. The portable global-state path described below can reshard supported finalized layouts. +After finalization every entry point re-validates the published layout, at two explicit costs. Steps and identity queries use an O(local params) fast path: one complete forensic rebuild caches an identity-token verdict — the finalized registries by object identity, every live group container, parameter, and compatibility name, and a version counter bumped by every legitimate mutating API — and the full rebuild re-runs only when a token changes. Checkpoint prepare/commit, canonical export/import, rebinding, state movement and offload activation, collective codebook initialize/refresh, scope re-validation, and contract readiness always re-run the complete forensic rebuild, and `post_sharding` computes the manifest shard set and sha256 digest exactly once per finalized manifest for the scoped operation headers. Consequently, corruption that preserves every fast-path token — in-place value replacement inside the private finalized registries or `object.__setattr__` on frozen identity records — is detected at the next full-forensics boundary rather than at the next step, while anything reachable through the public containers (group `params`/`param_names` slots, per-parameter state names, the compatibility-name cache) still fails the step guard itself before any state mutation, including mutations made by a closure between the pre- and post-closure guard blocks. + ## Explicit learned-codebook process groups `CodebookProcessGroupBinding` maps one stable `ProcessGroupIdentity` and local semantic member to an opaque PyTorch process-group handle plus an explicit collective device. It is accepted only through the complete `post_sharding(..., codebook_process_group=...)` transaction: every manifest shard must use that one semantic group, each local shard must name the binding's local member, the runtime group size and coordinate must match `ordered_members`, and the backend must support the supplied device. A one-member scope uses `process_group=None`; multi-member scopes must pass a real handle, including `dist.group.WORLD` when the default world is intentionally the semantic scope. Gefen never treats `None` as an implicit default-world selection. @@ -130,7 +132,7 @@ The core validates the finalized binding and complete declared state representat `StateOffloadProvider.offload_state_(device="cpu")` enables synchronous CPU-authoritative per-parameter state for an exact plain `Gefen` instance with ordinary replicated CUDA parameters. Activation first validates the complete declared state, stages tight detached CPU copies, waits for CUDA transfers, and publishes the policy and replacement state mapping together. At each eager step, Gefen copies only the current parameter's persistent tensor state to that parameter's CUDA device, runs the ordinary fused or non-fused block or factored update against a private runtime dictionary, synchronously copies the updated persistent state back to CPU, publishes that one dictionary, and releases the device temporaries. The optimizer-common learned codebook remains resident on CUDA and its normal per-device caches remain available. `restore_state_()` atomically co-locates all state with the parameters and disables offload; `move_state_()` has the same policy-disabling effect after its requested movement succeeds. -Activation and restore are fail-before-mutation. Activation also rejects persistent state tensors whose storage overlaps another persistent field, a parameter, or the common codebook because independent parameter paging cannot preserve such aliasing. If the update itself raises, Gefen attempts to preserve the resulting runtime state on CPU before propagating the original error. If copyback fails after a parameter may have changed, the optimizer is marked poisoned and refuses subsequent steps or native, canonical, and portable exports until a complete successful native `load_state_dict()` establishes known-good state. An active offload policy is target-local runtime configuration and is preserved across such a load rather than serialized as checkpoint meaning; the active loader maps parameter state directly to CPU and never accumulates the checkpoint's full parameter state on CUDA. State offload is implemented only for an exact plain `Gefen` instance; the composite Hybrid API, `GefenMuon`, nonreplicated finalized layouts, DTensor or tensor-subclass parameters, opaque extension state, multi-member explicit codebook scopes, capturable/device-authoritative state, compilation, and CUDA graph capture are excluded. The multi-member exclusion prevents one rank's copyback poison from bypassing the next scoped collective while peers enter it. Offload must be restored before post-sharding rebinding. It is blocking parameter-scoped paging, not asynchronous prefetch, overlap, or a distributed offload engine, and portable global-state I/O remains unavailable while its authoritative tensors are parked on CPU. +Activation and restore are fail-before-mutation. Activation also rejects persistent state tensors whose storage overlaps another persistent field, a parameter, or the common codebook because independent parameter paging cannot preserve such aliasing. The per-step offload readiness check reuses one cached verdict under the same identity-token scheme as the layout guard: the complete scan (per-tensor storage validation, pairwise disjointness, native-schema validation) re-runs at activation, movement, staged checkpoint loads, and whenever a mutating API bumps the layout version, so external in-place edits that corrupt already-validated offloaded state tensors while preserving container identities are detected at the next such boundary or by the step's own staging/copyback validation rather than by the step-entry check. If the update itself raises, Gefen attempts to preserve the resulting runtime state on CPU before propagating the original error. If copyback fails after a parameter may have changed, the optimizer is marked poisoned and refuses subsequent steps or native, canonical, and portable exports until a complete successful native `load_state_dict()` establishes known-good state. An active offload policy is target-local runtime configuration and is preserved across such a load rather than serialized as checkpoint meaning; the active loader maps parameter state directly to CPU and never accumulates the checkpoint's full parameter state on CUDA. State offload is implemented only for an exact plain `Gefen` instance; the composite Hybrid API, `GefenMuon`, nonreplicated finalized layouts, DTensor or tensor-subclass parameters, opaque extension state, multi-member explicit codebook scopes, capturable/device-authoritative state, compilation, and CUDA graph capture are excluded. The multi-member exclusion prevents one rank's copyback poison from bypassing the next scoped collective while peers enter it. Offload must be restored before post-sharding rebinding. It is blocking parameter-scoped paging, not asynchronous prefetch, overlap, or a distributed offload engine, and portable global-state I/O remains unavailable while its authoritative tensors are parked on CPU. `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. diff --git a/src/gefen/gefen.py b/src/gefen/gefen.py index e450386..8a55f5a 100644 --- a/src/gefen/gefen.py +++ b/src/gefen/gefen.py @@ -1052,6 +1052,19 @@ class Gefen(torch.optim.Optimizer): learning and period prediction. """ + # Pure rebuildable memoization for the layout-forensics guards. These live + # in ``__slots__`` rather than the instance ``__dict__`` so that staged + # ``__dict__`` copies, staged commits, and fail-before-mutation attribute + # snapshots treat them as what they are — caches, never optimizer state. + # Every accessor reads them through ``getattr`` with a cold default, so an + # object whose slots were never populated simply revalidates fully. + __slots__ = ( + "_gefen_layout_version", + "_gefen_layout_forensics_verdict", + "_gefen_manifest_forensics_cache", + "_gefen_state_offload_step_verdict", + ) + def __init__( self, params: Iterable[Union[nn.Parameter, Tuple[str, nn.Parameter]]], @@ -1221,6 +1234,21 @@ def __init__( self._gefen_post_sharding_finalized = False self._gefen_finalized_slots = () self._gefen_logical_slots = () + # Layout-forensics fast path. After one full structural pass succeeds, + # the exact container identities it validated are remembered together + # with a version counter that every legitimate mutating API bumps + # (post_sharding, staged checkpoint load commits, state movement and + # offload). Per-step guards accept only that unchanged token set; + # boundary operations (checkpoint prepare/commit, rebinding, movement + # and offload, collective codebook initialize/refresh, contract + # readiness) always rerun the complete forensic rebuild. + self._gefen_layout_version = 0 + self._gefen_layout_forensics_verdict = None + # (manifest, frozenset(manifest.shards), sha256 digest) computed once + # per finalized manifest object instead of rehashing every identity on + # every scoped step header. + self._gefen_manifest_forensics_cache = None + self._gefen_state_offload_step_verdict = 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 @@ -1412,9 +1440,117 @@ def _canonical_identity_ready(self) -> bool: or self._gefen_sharding_manifest is None ): return False - return self._finalized_binding_layout_matches() + # Contract readiness is an honest external claim, so it never trusts + # the per-step fast-path verdict. + return self._finalized_binding_layout_matches(full=True) + + def _invalidate_layout_forensics_caches(self) -> None: + self._gefen_layout_version = getattr(self, "_gefen_layout_version", 0) + 1 + self._gefen_layout_forensics_verdict = None + self._gefen_state_offload_step_verdict = None + + @staticmethod + def _forensics_tokens_match(cached, live) -> bool: + # Containers, tensors, and strings compare by object identity — the + # cached tuple holds strong references, so an unchanged attribute is + # the same object. Plain ints (version counter, lengths) compare by + # value because CPython interns only small integers. + return len(cached) == len(live) and all( + cached_item is live_item + or ( + type(cached_item) is int + and type(live_item) is int + and cached_item == live_item + ) + for cached_item, live_item in zip(cached, live) + ) - def _finalized_binding_layout_matches(self) -> bool: + def _layout_forensics_fast_tokens(self): + # O(local params) identity snapshot of everything user code can reach + # between two guard calls: the finalized registries by object identity + # (they are replaced, never mutated, by legitimate APIs) plus every + # live group container, parameter, and compatibility name element, so + # in-place mutation of the public groups — e.g. a closure swapping one + # ``group["params"]`` slot — still falls back to the full forensic + # pass before any state mutation. + live = [ + getattr(self, "_gefen_layout_version", 0), + self._gefen_post_sharding_finalized, + self._gefen_sharding_manifest, + self._gefen_logical_slots, + self._gefen_finalized_slots, + self._gefen_local_shard_bindings, + self._gefen_shard_bindings, + len(self._gefen_shard_bindings), + self._param_names, + len(self._param_names), + self.param_groups, + len(self.param_groups), + self.state, + ] + live.extend(self._param_names.keys()) + live.extend(self._param_names.values()) + for group in self.param_groups: + params = group.get("params") if type(group) is dict else None + names = group.get("param_names") if type(group) is dict else None + live.append(group) + live.append(params) + live.append(names) + if isinstance(params, (list, tuple)): + live.append(len(params)) + live.extend(params) + for parameter in params: + parameter_state = self.state.get(parameter) + live.append(type(parameter_state)) + live.append( + parameter_state.get("name") + if type(parameter_state) is dict + else None + ) + if isinstance(names, (list, tuple)): + live.append(len(names)) + live.extend(names) + return tuple(live) + + def _manifest_layout_forensics(self, *, refresh: bool = False): + """Return the cached ``(shard set, digest)`` for the live manifest. + + Both values are pure functions of one immutable ``ShardingManifest``, + so they are computed once per manifest object (at post_sharding + finalization) instead of rehashing every ``ShardIdentity`` — including + its world-size ``ordered_members`` tuple — on every step. Full + forensic passes recompute with ``refresh=True`` so boundary checks + never trust a stale cache. + """ + + manifest = self._gefen_sharding_manifest + cache = getattr(self, "_gefen_manifest_forensics_cache", None) + if not refresh and cache is not None and cache[0] is manifest: + return cache[1], cache[2] + shards = frozenset(manifest.shards) + digest = self._compute_codebook_manifest_fingerprint(manifest) + self._gefen_manifest_forensics_cache = (manifest, shards, digest) + return shards, digest + + def _finalized_binding_layout_matches(self, *, full: bool = False) -> bool: + try: + verdict = getattr(self, "_gefen_layout_forensics_verdict", None) + if not full and verdict is not None: + if self._forensics_tokens_match( + verdict, self._layout_forensics_fast_tokens() + ): + return True + self._gefen_layout_forensics_verdict = None + matches = self._finalized_binding_layout_matches_full() + if matches: + self._gefen_layout_forensics_verdict = ( + self._layout_forensics_fast_tokens() + ) + return matches + except AttributeError: + return False + + def _finalized_binding_layout_matches_full(self) -> bool: try: if ( type(self._gefen_logical_slots) is not tuple @@ -1471,7 +1607,7 @@ def _finalized_binding_layout_matches(self) -> bool: if any(not group for group in logical_groups): return False - manifest_shards = frozenset(self._gefen_sharding_manifest.shards) + manifest_shards = self._manifest_layout_forensics(refresh=True)[0] if { shard.parameter.fqn for shard in manifest_shards } != logical_fqns or any( @@ -1590,10 +1726,10 @@ def _finalized_binding_layout_matches(self) -> bool: ): return False - def _assert_finalized_binding_layout(self) -> None: + def _assert_finalized_binding_layout(self, *, full: bool = False) -> None: if ( self._gefen_post_sharding_finalized - and not self._finalized_binding_layout_matches() + and not self._finalized_binding_layout_matches(full=full) ): raise RuntimeError( "Gefen finalized parameter layout changed outside post_sharding" @@ -2420,11 +2556,11 @@ def _validate_codebook_runtime_binding(self, binding) -> None: raise ValueError("codebook collective CUDA device is unavailable") @torch._dynamo.disable - def _assert_runtime_codebook_process_group(self) -> None: + def _assert_runtime_codebook_process_group(self, *, full: bool = False) -> None: binding = self._gefen_codebook_process_group if binding is None: return - self._assert_finalized_binding_layout() + self._assert_finalized_binding_layout(full=full) self._validate_codebook_runtime_binding(binding) def _codebook_parameter_contributes(self, parameter) -> bool: @@ -2690,6 +2826,11 @@ def post_sharding( sources.append(rebinding.old_parameter) staged = self._stage_post_sharding(rebindings, manifest, codebook_process_group) self.__dict__.update(staged.__dict__) + self._invalidate_layout_forensics_caches() + # Compute the manifest shard set and digest exactly once at + # finalization; every later step guard and scoped operation header + # reuses this cache instead of rehashing every shard identity. + self._manifest_layout_forensics(refresh=True) def rebind_shard( self, @@ -3144,6 +3285,18 @@ def _step_with_offloaded_parameter_state( ) from exc self.state[parameter] = cpu_state + def _state_offload_step_tokens(self): + return ( + getattr(self, "_gefen_layout_version", 0), + self.state, + self.param_groups, + len(self.param_groups), + self._gefen_state_offload_device, + self._gefen_codebook, + self._gefen_codebook_process_group, + self.capturable, + ) + def _assert_state_offload_step_ready(self) -> None: if self.state_offload_poisoned: raise RuntimeError( @@ -3152,16 +3305,30 @@ def _assert_state_offload_step_ready(self) -> None: ) if not self.state_offload_active: return + # The complete offload scan (per-tensor storage checks, pairwise + # disjointness, native-schema validation) runs once after activation + # or any mutating API bumps the layout version; unchanged token + # identities reuse that verdict on the step hot path. Every boundary + # (activation, movement, staged checkpoint load) still runs the full + # scan directly through _state_offload_rejection_reason. + verdict = getattr(self, "_gefen_state_offload_step_verdict", None) + if verdict is not None: + if self._forensics_tokens_match( + verdict, self._state_offload_step_tokens() + ): + return + self._gefen_state_offload_step_verdict = None reason = self._state_offload_rejection_reason(require_cpu_state=True) if reason is not None: raise RuntimeError("Gefen state offload cannot step: {}".format(reason)) + self._gefen_state_offload_step_verdict = self._state_offload_step_tokens() @torch.no_grad() def offload_state_(self, device="cpu") -> None: """Atomically enable synchronous CPU-authoritative parameter state.""" target = self._normalize_state_offload_target(device) - self._assert_finalized_binding_layout() + self._assert_finalized_binding_layout(full=True) reason = self._state_offload_rejection_reason(require_cpu_state=False) if reason is not None: raise RuntimeError("Gefen state offload is unavailable: {}".format(reason)) @@ -3183,6 +3350,7 @@ def offload_state_(self, device="cpu") -> None: } ) self.__dict__.update(updates) + self._invalidate_layout_forensics_caches() @torch.no_grad() def restore_state_(self) -> None: @@ -3211,7 +3379,7 @@ def restore_state_(self) -> None: def _state_movement_rejection_reason(self): if ( self._gefen_post_sharding_finalized - and not self._finalized_binding_layout_matches() + and not self._finalized_binding_layout_matches(full=True) ): return "the finalized parameter binding no longer matches live groups" if self.capturable: @@ -3474,7 +3642,7 @@ def stage_tensor(value, target): def move_state_(self, device=None) -> None: """Atomically co-locate authoritative optimizer state with live parameters.""" - self._assert_finalized_binding_layout() + self._assert_finalized_binding_layout(full=True) try: reason = self._state_movement_rejection_reason() except Exception as exc: @@ -3519,6 +3687,7 @@ def move_state_(self, device=None) -> None: "_gefen_state_offload_device": None, } ) + self._invalidate_layout_forensics_caches() @staticmethod def _normalize_param_groups(params): @@ -5069,16 +5238,30 @@ def _reduce_codebook_scope_histogram(self, histogram: torch.Tensor) -> torch.Ten return reduced.cpu() def _codebook_manifest_fingerprint(self): - manifest = self._gefen_sharding_manifest + return self._manifest_layout_forensics()[1] + + def _compute_codebook_manifest_fingerprint(self, manifest): + # Scoped manifests always carry contiguous slices and one process-group + # identity, so their payload (and digest) is unchanged; the ungrouped + # and logical-region forms only appear for unscoped finalized layouts, + # whose digest is cached but never exchanged in a scope header. payload = tuple( ( shard.parameter.fqn, shard.parameter.global_shape, shard.layout.value, - shard.logical_slice.flat_offset, - shard.logical_slice.length, - shard.process_group.semantic_name, - shard.process_group.ordered_members, + shard.logical_slice.flat_offset + if isinstance(shard.logical_slice, LogicalSlice) + else tuple(shard.logical_slice.offsets), + shard.logical_slice.length + if isinstance(shard.logical_slice, LogicalSlice) + else tuple(shard.logical_slice.lengths), + None + if shard.process_group is None + else shard.process_group.semantic_name, + None + if shard.process_group is None + else shard.process_group.ordered_members, shard.local_member, shard.owner, tuple( @@ -5274,7 +5457,9 @@ def _ensure_codebook_scope_agreement(self) -> None: binding = self._gefen_codebook_process_group if binding is None or self._gefen_codebook_scope_validated: return - self._assert_runtime_codebook_process_group() + # Scope re-validation runs only after a mutating API reset the flag, + # so its collectives already amortize one full forensic rebuild. + self._assert_runtime_codebook_process_group(full=True) if len(binding.identity.ordered_members) == 1: self._gefen_codebook_scope_validated = self._gefen_codebook is not None return @@ -5779,7 +5964,7 @@ def _maybe_refresh_gefen_codebook(self) -> None: def initialize_codebook(self) -> bool: """Collectively initialize the learned codebook without taking a step.""" - self._assert_finalized_binding_layout() + self._assert_finalized_binding_layout(full=True) self._assert_runtime_codebook_process_group() self._assert_codebook_capture_ready() self._validate_codebook_scope_operation_header("initialize") @@ -5806,7 +5991,7 @@ def initialize_codebook(self) -> bool: def refresh_codebook(self) -> bool: """Collectively relearn and atomically requantize the current codebook.""" - self._assert_finalized_binding_layout() + self._assert_finalized_binding_layout(full=True) self._assert_runtime_codebook_process_group() if torch.cuda.is_available() and torch.cuda.is_current_stream_capturing(): raise RuntimeError( @@ -7047,7 +7232,7 @@ def _assert_canonical_import_target_safe(self) -> None: def export_canonical_state(self): """Export an exact-binding, device-neutral local state fragment.""" - self._assert_finalized_binding_layout() + self._assert_finalized_binding_layout(full=True) self._assert_state_export_safe() self._assert_canonical_state_outside_cuda_capture("export") if not self._canonical_state_layouts(): @@ -7296,7 +7481,7 @@ def _preserve_canonical_target_configuration(self, staged) -> None: def prepare_canonical_state_import(self, state): """Validate and stage a canonical local import without live mutation.""" - self._assert_finalized_binding_layout() + self._assert_finalized_binding_layout(full=True) self._assert_canonical_state_outside_cuda_capture("import preparation") self._assert_canonical_import_target_safe() if not self._canonical_state_layouts(): @@ -7322,7 +7507,7 @@ def commit_canonical_state_import(self, prepared) -> None: raise ValueError("prepared canonical state belongs to another optimizer") if prepared._consumed: raise RuntimeError("prepared canonical state import was already consumed") - self._assert_finalized_binding_layout() + self._assert_finalized_binding_layout(full=True) self._assert_canonical_state_outside_cuda_capture("import commit") self._assert_canonical_import_target_safe() if prepared._live_token != self._canonical_import_live_token(): @@ -7381,11 +7566,11 @@ def import_portable_state( def state_dict(self): """Run optimizer state-dict hooks around Gefen's complete schema.""" - self._assert_finalized_binding_layout() + 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() + 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(): @@ -8150,13 +8335,13 @@ def _pack_legacy_param_groups_for_load(self, state_dict): def load_state_dict(self, state_dict): """Atomically restore Gefen state between the public load hooks.""" - self._assert_finalized_binding_layout() + 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() + self._assert_finalized_binding_layout(full=True) staged = self._stage_load_state_dict(state_dict) self._commit_staged_load_state_dict(staged) for post_hook in self._optimizer_load_state_dict_post_hooks.values(): @@ -8217,6 +8402,7 @@ def _commit_staged_load_state_dict(self, staged) -> None: dict.update(live_defaults, staged.defaults) staged.defaults = live_defaults dict.update(self.__dict__, staged.__dict__) + self._invalidate_layout_forensics_caches() def _validate_loaded_native_state(self) -> None: """Validate the complete prepared native state before publication.""" From af1a46529f63dd5e7ad9b227788673572e5c7a18 Mon Sep 17 00:00:00 2001 From: thad0ctor Date: Mon, 13 Jul 2026 08:47:05 -0700 Subject: [PATCH 30/52] Add layout-guard cost tests and microbenchmark tests/test_layout_guard_cost.py pins the guard-cost contract: zero full forensic passes in steady-state steps, exactly one after finalization or any mutating API, the manifest digest computed once per finalized manifest, warm-verdict detection of closure/public-container tampering, boundary detection of private-registry in-place tampering, offload scan dedupe, and the caches staying out of the public attribute namespace. benchmarks/microbench/bench_layout_guard.py builds a synthetic 512 member x 300 parameter (153,600 shard) manifest with no process group and compares old per-step guard cost (7 forensic passes + 2 digest computes scoped; 2 passes unscoped) against the warm guard sequence: 42.6s -> 171us scoped (~250,000x), 9.6s -> 170us unscoped, far above the required 100x. --- benchmarks/microbench/bench_layout_guard.py | 200 +++++++++++ tests/test_layout_guard_cost.py | 362 ++++++++++++++++++++ 2 files changed, 562 insertions(+) create mode 100644 benchmarks/microbench/bench_layout_guard.py create mode 100644 tests/test_layout_guard_cost.py diff --git a/benchmarks/microbench/bench_layout_guard.py b/benchmarks/microbench/bench_layout_guard.py new file mode 100644 index 0000000..bcffbf6 --- /dev/null +++ b/benchmarks/microbench/bench_layout_guard.py @@ -0,0 +1,200 @@ +"""Benchmark the per-step layout-forensics guard cost on a synthetic manifest. + +Builds a finalized plain Gefen over a large synthetic flattened-shard +``ShardingManifest`` (defaults: 512 process-group members x 300 parameters = +153,600 global ``ShardIdentity`` records) with no real process group, then +compares: + + * old per-step guard cost — the pre-fix step() sequence re-ran the complete + O(params x world) forensic rebuild on every guard call (2 passes per + unscoped step, up to 7 under an explicit multi-member codebook scope) and + recomputed the manifest sha256 fingerprint inside every scoped operation + header (2 per step); + * new per-step guard cost — the exact warm step() guard sequence, which + reuses one cached forensic verdict through O(local params) identity + tokens and the manifest digest computed once at post_sharding + finalization. + +The headline is the scoped worst case (7 forensic passes + 2 digest +computes); the unscoped floor (2 passes, no digests) is printed alongside. +Exits nonzero if the scoped per-step reduction is below the required 100x. + +Run from the repo root (CPU only, no distributed init needed): + + PYTHONPATH=src python benchmarks/microbench/bench_layout_guard.py +""" + +from __future__ import annotations + +import argparse +import sys +import time + +import torch + +from gefen import ( + Gefen, + LogicalSlice, + ParameterIdentity, + ParameterLayout, + ParameterRebinding, + PlacementKind, + ProcessGroupIdentity, + ShardIdentity, + ShardPlacement, + ShardingManifest, +) + + +def build_finalized_optimizer(members_count: int, params_count: int, local_length: int): + members = tuple("member{:04d}".format(index) for index in range(members_count)) + group = ProcessGroupIdentity("data_parallel", members) + local_member = members[0] + manifest_shards = [] + local_shards = [] + for index in range(params_count): + identity = ParameterIdentity( + "Model.Block{}.Weight".format(index), (members_count * local_length,) + ) + offset = 0 + for coordinate, member in enumerate(members): + shard = ShardIdentity( + identity, + ParameterLayout.FLATTENED_ELEMENT_SHARD, + LogicalSlice(offset, local_length), + placements=( + ShardPlacement( + "data_parallel", + PlacementKind.FLAT_SHARD, + coordinate, + members_count, + ), + ), + process_group=group, + local_member=member, + ) + manifest_shards.append(shard) + if member == local_member: + local_shards.append(shard) + offset += local_length + + start = time.perf_counter() + manifest = ShardingManifest(tuple(manifest_shards)) + manifest_seconds = time.perf_counter() - start + + parameters = [ + torch.nn.Parameter(torch.randn(local_length)) for _ in range(params_count) + ] + optimizer = Gefen( + [ + ("model.block{}.weight".format(index), parameters[index]) + for index in range(params_count) + ], + fused=False, + factored_v_2d=False, + ) + start = time.perf_counter() + optimizer.post_sharding( + tuple( + ParameterRebinding(parameters[index], parameters[index], local_shards[index]) + for index in range(params_count) + ), + manifest=manifest, + ) + finalize_seconds = time.perf_counter() - start + return optimizer, parameters, manifest, manifest_seconds, finalize_seconds + + +def timed(callable_, repeats: int) -> float: + start = time.perf_counter() + for _ in range(repeats): + callable_() + return (time.perf_counter() - start) / repeats + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--members", type=int, default=512) + parser.add_argument("--params", type=int, default=300) + parser.add_argument("--local-length", type=int, default=4) + parser.add_argument("--full-repeats", type=int, default=3) + parser.add_argument("--fast-repeats", type=int, default=200) + args = parser.parse_args() + + optimizer, parameters, manifest, manifest_seconds, finalize_seconds = ( + build_finalized_optimizer(args.members, args.params, args.local_length) + ) + print( + "manifest: {} shards ({} members x {} params), built in {:.2f}s; " + "post_sharding finalized in {:.2f}s".format( + len(manifest.shards), + args.members, + args.params, + manifest_seconds, + finalize_seconds, + ) + ) + + # Warm the runtime and the forensic verdict exactly the way training does. + for parameter in parameters: + parameter.grad = torch.full_like(parameter, 0.5) + optimizer.step() + + full_pass = timed( + lambda: optimizer._finalized_binding_layout_matches(full=True), + args.full_repeats, + ) + digest_compute = timed( + lambda: optimizer._compute_codebook_manifest_fingerprint(manifest), + args.full_repeats, + ) + cached_digest = timed( + optimizer._codebook_manifest_fingerprint, args.fast_repeats + ) + + def warm_step_guards(): + # The complete step() guard sequence (both the pre-closure and the + # post-closure blocks), on a warm verdict. + optimizer._assert_state_offload_step_ready() + optimizer._assert_finalized_binding_layout() + optimizer._assert_runtime_codebook_process_group() + optimizer._assert_state_offload_step_ready() + optimizer._assert_finalized_binding_layout() + optimizer._assert_runtime_codebook_process_group() + + warm_guards = timed(warm_step_guards, args.fast_repeats) + + # Pre-fix per-step guard cost. Unscoped step(): 2 complete forensic + # passes. Scoped step(): up to 7 complete passes (step entry/re-entry, + # scope asserts, operation headers, failure synchronization, scope + # agreement) plus 2 manifest fingerprint recomputes in the exchanged + # "step" and "periodic_step" headers. + old_unscoped = 2 * full_pass + old_scoped = 7 * full_pass + 2 * digest_compute + new_unscoped = warm_guards + new_scoped = warm_guards + 2 * cached_digest + + print("one full forensic pass: {:>12.6f}s".format(full_pass)) + print("one manifest digest compute: {:>12.6f}s".format(digest_compute)) + print("one cached digest fetch: {:>12.6f}s".format(cached_digest)) + print("warm step guard sequence: {:>12.6f}s".format(warm_guards)) + print( + "old per-step guards (unscoped): {:>11.6f}s -> new: {:.6f}s ({:.0f}x)".format( + old_unscoped, new_unscoped, old_unscoped / new_unscoped + ) + ) + scoped_ratio = old_scoped / new_scoped + print( + "old per-step guards (scoped): {:>11.6f}s -> new: {:.6f}s ({:.0f}x)".format( + old_scoped, new_scoped, scoped_ratio + ) + ) + if scoped_ratio < 100.0: + print("FAIL: scoped per-step guard reduction is below 100x") + return 1 + print("PASS: scoped per-step guard reduction is >= 100x") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tests/test_layout_guard_cost.py b/tests/test_layout_guard_cost.py new file mode 100644 index 0000000..73f8ae9 --- /dev/null +++ b/tests/test_layout_guard_cost.py @@ -0,0 +1,362 @@ +"""Per-step layout/offload guard cost: fast-path tokens, dedupe, boundaries. + +The finalized-layout guards run on every ``step()``; these tests pin the +contract that steady-state steps reuse one cached forensic verdict (an +O(local params) identity token check) while every legitimate mutating API and +every boundary operation (checkpoint prepare/commit, rebinding, state +movement/offload, contract readiness) still runs the complete O(params x +world) forensic rebuild. Detection-before-mutation is preserved: anything a +closure or adapter can corrupt through the public optimizer containers still +raises at the step guard itself. +""" + +import time + +import pytest +import torch + +from gefen import ( + Gefen, + LogicalSlice, + ParameterIdentity, + ParameterLayout, + ParameterRebinding, + PlacementKind, + ProcessGroupIdentity, + ShardIdentity, + ShardPlacement, + ShardingManifest, +) + + +def _replicated_shard(fqn, shape): + identity = ParameterIdentity(fqn, shape) + return ShardIdentity( + identity, + ParameterLayout.REPLICATED, + LogicalSlice.full(identity), + ) + + +def _finalized_replicated_optimizer(count=3): + parameters = [ + torch.nn.Parameter(torch.full((4,), float(index + 1))) + for index in range(count) + ] + optimizer = Gefen( + [("weight{}".format(index), parameter) for index, parameter in enumerate(parameters)], + fused=False, + factored_v_2d=False, + ) + shards = tuple( + _replicated_shard("Model.Weight{}".format(index), (4,)) + for index in range(count) + ) + optimizer.post_sharding( + tuple( + ParameterRebinding(parameter, parameter, shard) + for parameter, shard in zip(parameters, shards) + ), + manifest=ShardingManifest(shards), + ) + return optimizer, parameters + + +def _flat_sharded_optimizer(members_count, params_count, local_length=4): + members = tuple("m{:04d}".format(index) for index in range(members_count)) + group = ProcessGroupIdentity("data_parallel", members) + local_member = members[0] + manifest_shards = [] + local_shards = [] + for index in range(params_count): + identity = ParameterIdentity( + "Model.Block{}.Weight".format(index), (members_count * local_length,) + ) + offset = 0 + for coordinate, member in enumerate(members): + shard = ShardIdentity( + identity, + ParameterLayout.FLATTENED_ELEMENT_SHARD, + LogicalSlice(offset, local_length), + placements=( + ShardPlacement( + "data_parallel", + PlacementKind.FLAT_SHARD, + coordinate, + members_count, + ), + ), + process_group=group, + local_member=member, + ) + manifest_shards.append(shard) + if member == local_member: + local_shards.append(shard) + offset += local_length + manifest = ShardingManifest(tuple(manifest_shards)) + parameters = [ + torch.nn.Parameter(torch.randn(local_length)) for _ in range(params_count) + ] + optimizer = Gefen( + [ + ("model.block{}.weight".format(index), parameters[index]) + for index in range(params_count) + ], + fused=False, + factored_v_2d=False, + ) + optimizer.post_sharding( + tuple( + ParameterRebinding(parameters[index], parameters[index], local_shards[index]) + for index in range(params_count) + ), + manifest=manifest, + ) + return optimizer, parameters + + +def _count_full_layout_passes(monkeypatch): + calls = {"count": 0} + original = Gefen._finalized_binding_layout_matches_full + + def counted(self): + calls["count"] += 1 + return original(self) + + monkeypatch.setattr(Gefen, "_finalized_binding_layout_matches_full", counted) + return calls + + +def _count_manifest_digest_computes(monkeypatch): + calls = {"count": 0} + original = Gefen._compute_codebook_manifest_fingerprint + + def counted(self, manifest): + calls["count"] += 1 + return original(self, manifest) + + monkeypatch.setattr( + Gefen, "_compute_codebook_manifest_fingerprint", counted + ) + return calls + + +def _step_with_grads(optimizer, parameters): + for parameter in parameters: + parameter.grad = torch.full_like(parameter, 0.5) + optimizer.step() + + +def test_steady_state_step_runs_no_full_layout_forensics(monkeypatch): + optimizer, parameters = _finalized_replicated_optimizer() + _step_with_grads(optimizer, parameters) + + calls = _count_full_layout_passes(monkeypatch) + for _ in range(3): + _step_with_grads(optimizer, parameters) + assert calls["count"] == 0 + + +def test_first_step_after_finalization_validates_fully_once(monkeypatch): + optimizer, parameters = _finalized_replicated_optimizer() + calls = _count_full_layout_passes(monkeypatch) + _step_with_grads(optimizer, parameters) + assert calls["count"] == 1 + + +def test_manifest_digest_is_computed_once_at_finalization(monkeypatch): + calls = _count_manifest_digest_computes(monkeypatch) + optimizer, parameters = _finalized_replicated_optimizer() + finalize_computes = calls["count"] + assert finalize_computes >= 1 + + first = optimizer._codebook_manifest_fingerprint() + second = optimizer._codebook_manifest_fingerprint() + assert calls["count"] == finalize_computes + assert first == second == optimizer._compute_codebook_manifest_fingerprint( + optimizer._gefen_sharding_manifest + ) + + _step_with_grads(optimizer, parameters) + _step_with_grads(optimizer, parameters) + assert calls["count"] > finalize_computes # the deliberate compare above + steady = calls["count"] + optimizer._codebook_manifest_fingerprint() + assert calls["count"] == steady + + +def test_closure_layout_mutation_is_detected_with_a_warm_verdict(): + optimizer, parameters = _finalized_replicated_optimizer() + _step_with_grads(optimizer, parameters) # warm the cached verdict + + rogue = torch.nn.Parameter(torch.ones(4) * 7) + rogue_before = rogue.detach().clone() + + def mutate_layout(): + optimizer.param_groups[0]["params"][0] = rogue + rogue.grad = torch.ones_like(rogue) + return torch.tensor(1.0) + + with pytest.raises(RuntimeError, match="changed outside post_sharding"): + optimizer.step(mutate_layout) + assert torch.equal(rogue, rogue_before) + + +@pytest.mark.parametrize( + "corruption", + ["local_bindings", "logical_slots", "group_params", "name_cache", "state_name"], +) +def test_step_detects_public_and_replaced_layout_tampering_after_warmup(corruption): + optimizer, parameters = _finalized_replicated_optimizer() + _step_with_grads(optimizer, parameters) # warm the cached verdict + + if corruption == "local_bindings": + optimizer._gefen_local_shard_bindings = tuple( + reversed(optimizer._gefen_local_shard_bindings) + ) + elif corruption == "logical_slots": + optimizer._gefen_logical_slots = tuple( + reversed(optimizer._gefen_logical_slots) + ) + elif corruption == "group_params": + optimizer.param_groups[0]["params"].reverse() + elif corruption == "name_cache": + optimizer._param_names[parameters[0]] = "changed" + else: + optimizer.state[parameters[0]]["name"] = "changed" + + with pytest.raises(RuntimeError, match="changed outside post_sharding"): + _step_with_grads(optimizer, parameters) + + +def test_private_registry_inplace_tamper_is_detected_at_the_next_boundary(): + # In-place value swaps inside the private finalized registries preserve + # every container identity the per-step tokens can see; per the documented + # contract they are detected by the next full forensic boundary + # (checkpoint prepare here, and contract readiness below) instead of the + # next step. + optimizer, parameters = _finalized_replicated_optimizer() + _step_with_grads(optimizer, parameters) + + first, second = parameters[0], parameters[1] + bindings = optimizer._gefen_shard_bindings + bindings[first], bindings[second] = bindings[second], bindings[first] + + _step_with_grads(optimizer, parameters) # fast tokens cannot see this + assert not optimizer.optimizer_contract().capabilities.stable_shard_identity + with pytest.raises(RuntimeError, match="changed outside post_sharding"): + optimizer.state_dict() + + +def test_mutating_apis_bump_the_layout_version_and_force_revalidation(monkeypatch): + optimizer, parameters = _finalized_replicated_optimizer() + _step_with_grads(optimizer, parameters) + checkpoint = optimizer.state_dict() + + calls = _count_full_layout_passes(monkeypatch) + + version = optimizer._gefen_layout_version + optimizer.move_state_() + assert optimizer._gefen_layout_version == version + 1 + assert optimizer._gefen_layout_forensics_verdict is None + + _step_with_grads(optimizer, parameters) + after_move = calls["count"] + assert after_move >= 1 + _step_with_grads(optimizer, parameters) + assert calls["count"] == after_move # steady again + + version = optimizer._gefen_layout_version + optimizer.load_state_dict(checkpoint) + assert optimizer._gefen_layout_version > version + assert optimizer._gefen_layout_forensics_verdict is None + + +def test_state_offload_step_scan_is_deduped_until_invalidated(monkeypatch): + optimizer, parameters = _finalized_replicated_optimizer() + _step_with_grads(optimizer, parameters) + + calls = {"count": 0} + + def counted(self, *, require_cpu_state, allow_poisoned=False): + calls["count"] += 1 + return None + + monkeypatch.setattr(Gefen, "_state_offload_rejection_reason", counted) + monkeypatch.setattr( + Gefen, + "state_offload_active", + property(lambda self: True), + ) + + optimizer._assert_state_offload_step_ready() + optimizer._assert_state_offload_step_ready() + optimizer._assert_state_offload_step_ready() + assert calls["count"] == 1 + + optimizer._invalidate_layout_forensics_caches() + optimizer._assert_state_offload_step_ready() + assert calls["count"] == 2 + + optimizer._gefen_state_offload_poisoned = True + with pytest.raises(RuntimeError, match="poisoned"): + optimizer._assert_state_offload_step_ready() + + +def test_offload_verdict_is_not_cached_when_the_scan_rejects(monkeypatch): + optimizer, parameters = _finalized_replicated_optimizer() + + calls = {"count": 0} + + def counted(self, *, require_cpu_state, allow_poisoned=False): + calls["count"] += 1 + return "authoritative offloaded tensors must be tight CPU tensors" + + monkeypatch.setattr(Gefen, "_state_offload_rejection_reason", counted) + monkeypatch.setattr( + Gefen, + "state_offload_active", + property(lambda self: True), + ) + + for _ in range(2): + with pytest.raises(RuntimeError, match="cannot step"): + optimizer._assert_state_offload_step_ready() + assert calls["count"] == 2 + + +def test_forensics_caches_stay_out_of_the_public_attribute_namespace(): + optimizer, parameters = _finalized_replicated_optimizer() + _step_with_grads(optimizer, parameters) + assert optimizer._gefen_layout_forensics_verdict is not None + for name in Gefen.__slots__: + assert name not in optimizer.__dict__ + + +def test_warm_step_guards_are_far_cheaper_than_one_forensic_pass(): + # Relative timing with a deliberately generous margin: the old behavior + # ran (at least) two full forensic passes inside every step, so the warm + # per-step guard sequence must be much cheaper than even one full pass on + # a moderately sized manifest. Identity-token checks are microseconds + # while the full pass rehashes 96-member identities for every shard, so + # the 10x assertion holds with orders of magnitude to spare. + optimizer, parameters = _flat_sharded_optimizer(96, 60) + _step_with_grads(optimizer, parameters) # warm caches + + iterations = 50 + start = time.perf_counter() + for _ in range(iterations): + optimizer._assert_state_offload_step_ready() + optimizer._assert_finalized_binding_layout() + optimizer._assert_runtime_codebook_process_group() + optimizer._assert_state_offload_step_ready() + optimizer._assert_finalized_binding_layout() + optimizer._assert_runtime_codebook_process_group() + warm_guard = (time.perf_counter() - start) / iterations + + start = time.perf_counter() + for _ in range(iterations): + assert optimizer._finalized_binding_layout_matches(full=True) + full_pass = (time.perf_counter() - start) / iterations + + assert warm_guard * 10 < full_pass From 56803c444740be612921b49c2e0205aabd659c8f Mon Sep 17 00:00:00 2001 From: thad0ctor Date: Mon, 13 Jul 2026 12:48:08 -0700 Subject: [PATCH 31/52] Run the offload readiness scan on every step The perf caching applied a cached identity-token verdict to the per-step offload readiness check. Its token captured the state and param-group containers by identity but not the per-parameter offloaded state tensors, which are legitimately replaced every step. A token-preserving in-place corruption of a later parameter's offloaded state therefore slipped past step entry and was only caught mid-step, after earlier parameters had already been updated and committed, breaking fail-before-mutation. The readiness scan is O(local params) and was never the layout-forensics cost the cache targeted, so drop the offload verdict cache and run the full scan before any parameter is staged. The layout-manifest digest cache is unchanged. Update the layout-guard tests to assert every-step scanning and add a CUDA regression that corrupts a later parameter and checks the first is left byte-for-byte untouched. --- src/gefen/gefen.py | 38 +++++++---------------------- tests/test_layout_guard_cost.py | 16 ++++++------ tests/test_state_offload.py | 43 +++++++++++++++++++++++++++++++++ 3 files changed, 61 insertions(+), 36 deletions(-) diff --git a/src/gefen/gefen.py b/src/gefen/gefen.py index 8a55f5a..3517fd2 100644 --- a/src/gefen/gefen.py +++ b/src/gefen/gefen.py @@ -1062,7 +1062,6 @@ class Gefen(torch.optim.Optimizer): "_gefen_layout_version", "_gefen_layout_forensics_verdict", "_gefen_manifest_forensics_cache", - "_gefen_state_offload_step_verdict", ) def __init__( @@ -1248,7 +1247,6 @@ def __init__( # per finalized manifest object instead of rehashing every identity on # every scoped step header. self._gefen_manifest_forensics_cache = None - self._gefen_state_offload_step_verdict = 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 @@ -1447,7 +1445,6 @@ def _canonical_identity_ready(self) -> bool: def _invalidate_layout_forensics_caches(self) -> None: self._gefen_layout_version = getattr(self, "_gefen_layout_version", 0) + 1 self._gefen_layout_forensics_verdict = None - self._gefen_state_offload_step_verdict = None @staticmethod def _forensics_tokens_match(cached, live) -> bool: @@ -3285,18 +3282,6 @@ def _step_with_offloaded_parameter_state( ) from exc self.state[parameter] = cpu_state - def _state_offload_step_tokens(self): - return ( - getattr(self, "_gefen_layout_version", 0), - self.state, - self.param_groups, - len(self.param_groups), - self._gefen_state_offload_device, - self._gefen_codebook, - self._gefen_codebook_process_group, - self.capturable, - ) - def _assert_state_offload_step_ready(self) -> None: if self.state_offload_poisoned: raise RuntimeError( @@ -3305,23 +3290,18 @@ def _assert_state_offload_step_ready(self) -> None: ) if not self.state_offload_active: return - # The complete offload scan (per-tensor storage checks, pairwise - # disjointness, native-schema validation) runs once after activation - # or any mutating API bumps the layout version; unchanged token - # identities reuse that verdict on the step hot path. Every boundary - # (activation, movement, staged checkpoint load) still runs the full - # scan directly through _state_offload_rejection_reason. - verdict = getattr(self, "_gefen_state_offload_step_verdict", None) - if verdict is not None: - if self._forensics_tokens_match( - verdict, self._state_offload_step_tokens() - ): - return - self._gefen_state_offload_step_verdict = None + # Run the complete offload scan (per-tensor storage checks, pairwise + # disjointness, native-schema validation) on every step before any + # parameter is staged. Unlike the finalized-layout manifest, the + # per-parameter offloaded state tensors are legitimately replaced each + # step, so a cached verdict cannot represent them; a token-preserving + # in-place corruption of a later parameter's state would otherwise slip + # past step entry and only be caught mid-step, after earlier parameters + # were already updated and committed. The scan is O(local params) and + # was never the layout-forensics cost this cache was introduced for. reason = self._state_offload_rejection_reason(require_cpu_state=True) if reason is not None: raise RuntimeError("Gefen state offload cannot step: {}".format(reason)) - self._gefen_state_offload_step_verdict = self._state_offload_step_tokens() @torch.no_grad() def offload_state_(self, device="cpu") -> None: diff --git a/tests/test_layout_guard_cost.py b/tests/test_layout_guard_cost.py index 73f8ae9..0449716 100644 --- a/tests/test_layout_guard_cost.py +++ b/tests/test_layout_guard_cost.py @@ -272,7 +272,13 @@ def test_mutating_apis_bump_the_layout_version_and_force_revalidation(monkeypatc assert optimizer._gefen_layout_forensics_verdict is None -def test_state_offload_step_scan_is_deduped_until_invalidated(monkeypatch): +def test_state_offload_step_scan_runs_on_every_step(monkeypatch): + # The offload readiness scan is intentionally NOT cached: the per-parameter + # offloaded state tensors are legitimately replaced on every step, so a + # cached verdict cannot represent them, and a token-preserving in-place + # corruption of a later parameter would otherwise slip past step entry and + # only be caught mid-step, after earlier parameters were already mutated. + # The scan is O(local params) and must run before any parameter is staged. optimizer, parameters = _finalized_replicated_optimizer() _step_with_grads(optimizer, parameters) @@ -292,18 +298,14 @@ def counted(self, *, require_cpu_state, allow_poisoned=False): optimizer._assert_state_offload_step_ready() optimizer._assert_state_offload_step_ready() optimizer._assert_state_offload_step_ready() - assert calls["count"] == 1 - - optimizer._invalidate_layout_forensics_caches() - optimizer._assert_state_offload_step_ready() - assert calls["count"] == 2 + assert calls["count"] == 3 optimizer._gefen_state_offload_poisoned = True with pytest.raises(RuntimeError, match="poisoned"): optimizer._assert_state_offload_step_ready() -def test_offload_verdict_is_not_cached_when_the_scan_rejects(monkeypatch): +def test_offload_scan_re_rejects_on_every_step(monkeypatch): optimizer, parameters = _finalized_replicated_optimizer() calls = {"count": 0} diff --git a/tests/test_state_offload.py b/tests/test_state_offload.py index 7235aa9..a951def 100644 --- a/tests/test_state_offload.py +++ b/tests/test_state_offload.py @@ -217,6 +217,49 @@ def inspected(self, group, name, parameter, grad, *, state=None): _assert_cpu_boundary(optimizer) +@_CUDA_REQUIRED +def test_later_parameter_corruption_is_caught_before_earlier_parameter_mutates(): + # Regression: the offloaded step processes parameters sequentially + # (stage -> update -> copyback -> commit per parameter). If a later + # parameter's offloaded state is corrupted in a way that preserves the + # state/param-group container identities and the layout version, a cached + # step-readiness verdict would let step() begin, mutate the first + # parameter, and only reject when the second parameter is staged -- + # violating fail-before-mutation. The readiness scan therefore runs in + # full on every step, so the corruption is caught at step entry and the + # first parameter is left byte-for-byte untouched. + first = torch.nn.Parameter(torch.arange(8, device="cuda", dtype=torch.float32)) + second = torch.nn.Parameter(torch.arange(8, 16, device="cuda", dtype=torch.float32)) + optimizer = Gefen( + [("first", first), ("second", second)], + fused=False, + factored_v_2d=False, + ) + optimizer._resolve_automatic_period = lambda *args, **kwargs: 4 + optimizer.offload_state_() + + first.grad = torch.linspace(-1.0, 1.0, 8, device="cuda") + second.grad = torch.linspace(1.0, -1.0, 8, device="cuda") + optimizer.step() + + first_state_before = _persistent_snapshot(optimizer, first) + first_value_before = first.detach().clone() + + # Corrupt the SECOND parameter's offloaded state with a non-tight CPU view. + # self.state, self.state[second], and the layout version are all unchanged. + corrupt = optimizer.state[second]["m_magnitude"].repeat_interleave(2)[::2] + assert not corrupt.is_contiguous() + optimizer.state[second]["m_magnitude"] = corrupt + + first.grad = torch.linspace(-1.0, 1.0, 8, device="cuda") + second.grad = torch.linspace(1.0, -1.0, 8, device="cuda") + with pytest.raises(RuntimeError, match="cannot step"): + optimizer.step() + + _assert_persistent_equal(_persistent_snapshot(optimizer, first), first_state_before) + assert torch.equal(first.detach(), first_value_before) + + @_CUDA_REQUIRED def test_active_load_preserves_target_policy_and_exact_continuation(monkeypatch): source_parameter = torch.nn.Parameter(torch.linspace(-0.5, 0.5, 8, device="cuda")) From a9db94f4d2f993f5b8acde88ace0f7fdea3a54ef Mon Sep 17 00:00:00 2001 From: thad0ctor Date: Mon, 13 Jul 2026 12:48:08 -0700 Subject: [PATCH 32/52] Clarify layout-guard and offload paragraphs in contract docs Split the two run-on paragraphs into readable sentences, state the step-vs-boundary detection split as an explicit consequence for integrators, and list the full-rebuild boundaries plainly. Correct the offload paragraph: the readiness scan now runs in full on every step rather than reusing a cached verdict, and it does not run at state movement (movement disables offload). --- docs/optimizer_contracts.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/docs/optimizer_contracts.md b/docs/optimizer_contracts.md index 46fa1ec..270ce89 100644 --- a/docs/optimizer_contracts.md +++ b/docs/optimizer_contracts.md @@ -41,7 +41,7 @@ Rebinding is allowed only while the entire optimizer is pristine: global step ze Plain Gefen currently rebinds complete replicated parameters and contiguous physical 1-D flattened element shards. A flattened logical matrix requires `factored_v_2d=False` because canonical row/column factored-state projection is not implemented. GefenMuon rebinds complete replicated matrices and can finalize/prune whole-parameter owner manifests. A whole-owner instance may step only when the same transaction installs the explicit codebook process group described below; the owner performs the update and the machine-readable training declaration requires the adapter to synchronize the updated matrix afterward. A Gefen-backed `GefenMuonHybrid` atomically partitions one complete manifest and rebinding plan by its frozen exact FQN routing, stages both children, validates cross-child storage disjointness, rebuilds composite state routing, and publishes only after every child succeeds. AdamW-backed Hybrid and DTensor composite rebinding remain unsupported. The portable global-state path described below can reshard supported finalized layouts. -After finalization every entry point re-validates the published layout, at two explicit costs. Steps and identity queries use an O(local params) fast path: one complete forensic rebuild caches an identity-token verdict — the finalized registries by object identity, every live group container, parameter, and compatibility name, and a version counter bumped by every legitimate mutating API — and the full rebuild re-runs only when a token changes. Checkpoint prepare/commit, canonical export/import, rebinding, state movement and offload activation, collective codebook initialize/refresh, scope re-validation, and contract readiness always re-run the complete forensic rebuild, and `post_sharding` computes the manifest shard set and sha256 digest exactly once per finalized manifest for the scoped operation headers. Consequently, corruption that preserves every fast-path token — in-place value replacement inside the private finalized registries or `object.__setattr__` on frozen identity records — is detected at the next full-forensics boundary rather than at the next step, while anything reachable through the public containers (group `params`/`param_names` slots, per-parameter state names, the compatibility-name cache) still fails the step guard itself before any state mutation, including mutations made by a closure between the pre- and post-closure guard blocks. +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. ## Explicit learned-codebook process groups @@ -132,7 +132,7 @@ The core validates the finalized binding and complete declared state representat `StateOffloadProvider.offload_state_(device="cpu")` enables synchronous CPU-authoritative per-parameter state for an exact plain `Gefen` instance with ordinary replicated CUDA parameters. Activation first validates the complete declared state, stages tight detached CPU copies, waits for CUDA transfers, and publishes the policy and replacement state mapping together. At each eager step, Gefen copies only the current parameter's persistent tensor state to that parameter's CUDA device, runs the ordinary fused or non-fused block or factored update against a private runtime dictionary, synchronously copies the updated persistent state back to CPU, publishes that one dictionary, and releases the device temporaries. The optimizer-common learned codebook remains resident on CUDA and its normal per-device caches remain available. `restore_state_()` atomically co-locates all state with the parameters and disables offload; `move_state_()` has the same policy-disabling effect after its requested movement succeeds. -Activation and restore are fail-before-mutation. Activation also rejects persistent state tensors whose storage overlaps another persistent field, a parameter, or the common codebook because independent parameter paging cannot preserve such aliasing. The per-step offload readiness check reuses one cached verdict under the same identity-token scheme as the layout guard: the complete scan (per-tensor storage validation, pairwise disjointness, native-schema validation) re-runs at activation, movement, staged checkpoint loads, and whenever a mutating API bumps the layout version, so external in-place edits that corrupt already-validated offloaded state tensors while preserving container identities are detected at the next such boundary or by the step's own staging/copyback validation rather than by the step-entry check. If the update itself raises, Gefen attempts to preserve the resulting runtime state on CPU before propagating the original error. If copyback fails after a parameter may have changed, the optimizer is marked poisoned and refuses subsequent steps or native, canonical, and portable exports until a complete successful native `load_state_dict()` establishes known-good state. An active offload policy is target-local runtime configuration and is preserved across such a load rather than serialized as checkpoint meaning; the active loader maps parameter state directly to CPU and never accumulates the checkpoint's full parameter state on CUDA. State offload is implemented only for an exact plain `Gefen` instance; the composite Hybrid API, `GefenMuon`, nonreplicated finalized layouts, DTensor or tensor-subclass parameters, opaque extension state, multi-member explicit codebook scopes, capturable/device-authoritative state, compilation, and CUDA graph capture are excluded. The multi-member exclusion prevents one rank's copyback poison from bypassing the next scoped collective while peers enter it. Offload must be restored before post-sharding rebinding. It is blocking parameter-scoped paging, not asynchronous prefetch, overlap, or a distributed offload engine, and portable global-state I/O remains unavailable while its authoritative tensors are parked on CPU. +Activation and restore are fail-before-mutation. Activation also rejects persistent state tensors whose storage overlaps another persistent field, a parameter, or the common codebook because independent parameter paging cannot preserve such aliasing. Every offloaded step then re-runs the complete readiness scan — per-tensor storage validation, pairwise disjointness, and native-schema validation — before any parameter is staged, so an external in-place edit that corrupts an already-validated offloaded state tensor is caught at step entry, before the step mutates any parameter. This scan is O(local params) and is deliberately not cached across steps: unlike the finalized layout, the per-parameter offloaded state tensors are legitimately replaced on every step, so no cached verdict could stand in for them. If the update itself raises, Gefen attempts to preserve the resulting runtime state on CPU before propagating the original error. If copyback fails after a parameter may have changed, the optimizer is marked poisoned and refuses subsequent steps or native, canonical, and portable exports until a complete successful native `load_state_dict()` establishes known-good state. An active offload policy is target-local runtime configuration and is preserved across such a load rather than serialized as checkpoint meaning; the active loader maps parameter state directly to CPU and never accumulates the checkpoint's full parameter state on CUDA. State offload is implemented only for an exact plain `Gefen` instance; the composite Hybrid API, `GefenMuon`, nonreplicated finalized layouts, DTensor or tensor-subclass parameters, opaque extension state, multi-member explicit codebook scopes, capturable/device-authoritative state, compilation, and CUDA graph capture are excluded. The multi-member exclusion prevents one rank's copyback poison from bypassing the next scoped collective while peers enter it. Offload must be restored before post-sharding rebinding. It is blocking parameter-scoped paging, not asynchronous prefetch, overlap, or a distributed offload engine, and portable global-state I/O remains unavailable while its authoritative tensors are parked on CPU. `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. From 4dee59cb2daca19533bbff353b8fb0b90d49107b Mon Sep 17 00:00:00 2001 From: thad0ctor Date: Mon, 13 Jul 2026 12:59:33 -0700 Subject: [PATCH 33/52] Cache the composite hybrid finalized-layout forensics behind O(local) step tokens GefenMuonHybrid.step() calls its finalized-layout guard twice per step and each call ran the full O(params) composite forensic rebuild -- routing, ownership, per-child manifest partitioning, and local-binding sort -- so a finalized hybrid paid the reconstruction Gefen/GefenMuon already avoid. Memoize one composite verdict as an O(local params) identity-token snapshot, mirroring the base scheme. The token captures every field the rebuild reads by identity (the finalized flags, the composite manifest/roles/local bindings/finalized slots/owner registry/codebook binding, the subopt list and each child's private manifest, local bindings, codebook binding and defaults, and every live child group['params'] container) and folds in each child's own already-cached fast-path verdict, so any child-level change the children can detect propagates automatically and any composite container replacement is caught directly. post_sharding -- the only hybrid API that reassigns the composite fields, and the one every rebind helper routes through -- bumps a version counter and clears the verdict. Contract readiness still runs the full rebuild (full=True), and the hybrid guards now accept the same full= kwarg the base scoped path passes through _assert_runtime_codebook_process_group. Detection stays before any mutation: an in-place child param slot swap, a swapped child binding, closure-time public-container tampering, and composite registry replacement all still raise at the step guard. --- src/gefen/hybrid.py | 109 +++++++++++++++++++- tests/test_hybrid_layout_cache.py | 162 ++++++++++++++++++++++++++++++ 2 files changed, 267 insertions(+), 4 deletions(-) create mode 100644 tests/test_hybrid_layout_cache.py diff --git a/src/gefen/hybrid.py b/src/gefen/hybrid.py index a8f3592..b6a72f4 100644 --- a/src/gefen/hybrid.py +++ b/src/gefen/hybrid.py @@ -575,6 +575,15 @@ def _is_no_decay(name): self._hybrid_codebook_process_group = None self._hybrid_finalized_slots = () + # Composite finalized-layout forensics cache, mirroring the O(local + # params) scheme Gefen/GefenMuon use for their own step guards. The + # version counter is bumped by the only hybrid API that reassigns any + # of the finalized composite fields (post_sharding, which every rebind + # helper routes through); the cached verdict is one identity-token + # snapshot of everything the composite forensic rebuild reads. + self._hybrid_layout_version = 0 + self._hybrid_layout_forensics_verdict = None + # Deliberately do NOT call super().__init__(): we expose each # sub-optimizer's real param_groups/state via properties (shared dict # refs), so the LR scheduler's in-place ``group["lr"] = ...`` updates @@ -894,6 +903,13 @@ def post_sharding( "_hybrid_fqn_roles": staged["fqn_roles"], "_hybrid_codebook_process_group": staged["codebook_process_group"], "_hybrid_finalized_slots": staged["finalized_slots"], + # Reassigning the composite fields invalidates any warm verdict + # (there is none from a pristine hybrid, but bumping keeps the + # counter honest and forces the next guard through a full + # rebuild before it caches a fresh verdict). + "_hybrid_layout_version": getattr(self, "_hybrid_layout_version", 0) + + 1, + "_hybrid_layout_forensics_verdict": None, }, ) @@ -935,7 +951,88 @@ def rebind_parameter( manifest=ShardingManifest((shard,)), ) - def _finalized_binding_layout_matches(self) -> bool: + @staticmethod + def _hybrid_child_param_group_tokens(child, tokens): + # Append the child's live group containers, their ``params`` list + # objects, and every parameter in them so an in-place slot swap in a + # child's public param_groups is visible to the composite fast path. + tokens.append(child) + tokens.append(getattr(child, "_gefen_sharding_manifest", None)) + tokens.append(getattr(child, "_gefen_codebook_process_group", None)) + tokens.append(getattr(child, "_gefen_local_shard_bindings", None)) + tokens.append(getattr(child, "defaults", None)) + groups = child.param_groups + tokens.append(groups) + tokens.append(len(groups)) + for group in groups: + params = group.get("params") if type(group) is dict else None + tokens.append(group) + tokens.append(params) + if isinstance(params, (list, tuple)): + tokens.append(len(params)) + tokens.extend(params) + + def _hybrid_layout_forensics_fast_tokens(self): + # O(local params) identity snapshot of everything the composite + # forensic rebuild reads by identity, plus each child's own O(local) + # cached fast-path verdict. Calling the children's fast path means any + # child-level change they can detect (an in-place ``group['params']`` + # slot swap, a replaced private child registry, a mutated name cache) + # flips a bool in this token and forces the composite through a full + # rebuild; the composite-owned fields are captured directly so a + # replaced manifest/roles/local-bindings/slots/owner/binding container + # is caught even if no child noticed. Legitimate mutating APIs replace + # these containers (never mutate them in place) and bump the version + # counter, so an unchanged attribute is the same object. + live = [ + getattr(self, "_hybrid_layout_version", 0), + self._hybrid_post_sharding_finalized, + self._hybrid_sharding_manifest, + self._hybrid_local_shard_bindings, + self._hybrid_fqn_roles, + self._hybrid_finalized_slots, + self._hybrid_codebook_process_group, + self._state_param_owner, + len(self._state_param_owner), + self.defaults, + self._subopts, + len(self._subopts), + ] + for child in self._subopts: + live.append(child._finalized_binding_layout_matches()) + GefenMuonHybrid._hybrid_child_param_group_tokens(child, live) + return tuple(live) + + def _finalized_binding_layout_matches(self, *, full: bool = False) -> bool: + # Steady-state guard: reuse one cached verdict validated by an + # O(local params) identity-token check. Any real finalized-layout + # change either bumps the hybrid version, replaces a composite + # container, or flips a child's own fast-path verdict -- all captured + # by the token -- so a stale hit is impossible; a full boundary check + # (``full=True`` from base contract-readiness call sites) always + # rebuilds. Detection stays before any mutation. + verdict = getattr(self, "_hybrid_layout_forensics_verdict", None) + if not full and verdict is not None: + try: + fast_tokens = self._hybrid_layout_forensics_fast_tokens() + except (AttributeError, KeyError, TypeError, ValueError, RuntimeError): + fast_tokens = None + if fast_tokens is not None and Gefen._forensics_tokens_match( + verdict, fast_tokens + ): + return True + self._hybrid_layout_forensics_verdict = None + matches = self._finalized_binding_layout_matches_full() + if matches: + try: + self._hybrid_layout_forensics_verdict = ( + self._hybrid_layout_forensics_fast_tokens() + ) + except (AttributeError, KeyError, TypeError, ValueError, RuntimeError): + self._hybrid_layout_forensics_verdict = None + return matches + + def _finalized_binding_layout_matches_full(self) -> bool: try: if ( not self._hybrid_post_sharding_finalized @@ -1038,14 +1135,18 @@ def _finalized_binding_layout_matches(self) -> bool: return False def _canonical_identity_ready(self) -> bool: - return self._hybrid_post_sharding_finalized and self._finalized_binding_layout_matches() + # Contract readiness is an honest external claim, so it never trusts the + # per-step fast-path verdict and always runs the full forensic rebuild. + return self._hybrid_post_sharding_finalized and self._finalized_binding_layout_matches( + full=True + ) def _codebook_scope_ready(self) -> bool: return self._hybrid_codebook_process_group is not None and self._canonical_identity_ready() - def _assert_finalized_binding_layout(self) -> None: + def _assert_finalized_binding_layout(self, *, full: bool = False) -> None: if self._hybrid_post_sharding_finalized: - if not self._finalized_binding_layout_matches(): + if not self._finalized_binding_layout_matches(full=full): raise RuntimeError("GefenMuonHybrid finalized parameter layout changed outside post_sharding") elif not self._hybrid_identity_metadata_empty(): raise RuntimeError("GefenMuonHybrid found an incomplete post_sharding identity plan") diff --git a/tests/test_hybrid_layout_cache.py b/tests/test_hybrid_layout_cache.py new file mode 100644 index 0000000..f6bbf2a --- /dev/null +++ b/tests/test_hybrid_layout_cache.py @@ -0,0 +1,162 @@ +"""Per-step composite layout-guard cost for a finalized GefenMuonHybrid. + +GefenMuonHybrid.step() calls its finalized-layout guard twice per step, and +each call used to run the full O(params) composite forensic rebuild (routing, +ownership, per-child manifest partitioning, local-binding sort). These tests +pin the contract that a steady-state finalized composite reuses one cached +verdict (an O(local params) identity-token check that folds in each child's +own cached fast-path verdict) instead of rebuilding, while every real +finalized-layout change -- an in-place child ``group['params']`` slot swap, a +replaced composite registry, a closure that tampers between the pre- and +post-closure guards -- is still detected before any state mutation. +""" + +import pytest +import torch + +from gefen import GefenMuonHybrid +from gefen.contracts import ( + LogicalSlice, + ParameterIdentity, + ParameterLayout, + ShardIdentity, + ShardingManifest, +) +from gefen.rebinding import ParameterRebinding + + +def _replicated_shard(fqn, shape): + identity = ParameterIdentity(fqn, shape) + return ShardIdentity( + identity, + ParameterLayout.REPLICATED, + LogicalSlice.full(identity), + ) + + +def _finalized_hybrid(): + matrix = torch.nn.Parameter(torch.randn(2, 2)) + bias = torch.nn.Parameter(torch.randn(4)) + optimizer = GefenMuonHybrid( + [("layer.weight", matrix)], + [("layer.bias", bias)], + lr=1e-3, + fused=False, + sharded_mode="distributed", + ) + matrix_shard = _replicated_shard("Model.Layer.Weight", (2, 2)) + bias_shard = _replicated_shard("Model.Layer.Bias", (4,)) + optimizer.post_sharding( + ( + ParameterRebinding(matrix, matrix, matrix_shard), + ParameterRebinding(bias, bias, bias_shard), + ), + manifest=ShardingManifest((matrix_shard, bias_shard)), + ) + return optimizer, matrix, bias + + +def _count_full_composite_rebuilds(monkeypatch): + # The full composite rebuild is the only step-time caller of + # _gefen_rebinding_children; the cached fast path never touches it. This + # distinguishes a real rebuild in both the pre- and post-cache code. + calls = {"count": 0} + original = GefenMuonHybrid._gefen_rebinding_children + + def counted(self): + calls["count"] += 1 + return original(self) + + monkeypatch.setattr(GefenMuonHybrid, "_gefen_rebinding_children", counted) + return calls + + +def _step_with_grads(optimizer, params): + for parameter in params: + parameter.grad = torch.full_like(parameter, 0.5) + optimizer.step() + + +def test_steady_state_finalized_hybrid_step_does_not_rebuild(monkeypatch): + optimizer, matrix, bias = _finalized_hybrid() + _step_with_grads(optimizer, (matrix, bias)) # warm the cached verdict + + calls = _count_full_composite_rebuilds(monkeypatch) + for _ in range(3): + _step_with_grads(optimizer, (matrix, bias)) + # Pre-cache each step's two guard calls each ran a full rebuild (6 total); + # with the verdict cache a steady-state step rebuilds zero times. + assert calls["count"] == 0 + + +def test_first_finalized_hybrid_step_rebuilds_then_caches(monkeypatch): + optimizer, matrix, bias = _finalized_hybrid() + calls = _count_full_composite_rebuilds(monkeypatch) + _step_with_grads(optimizer, (matrix, bias)) + # The first guard rebuilds once and caches; the second guard in the same + # step reuses the warm verdict. + assert calls["count"] == 1 + assert optimizer._hybrid_layout_forensics_verdict is not None + + +def test_finalized_hybrid_detects_child_param_slot_swap_with_warm_verdict(): + optimizer, matrix, bias = _finalized_hybrid() + _step_with_grads(optimizer, (matrix, bias)) # warm the cached verdict + assert optimizer._hybrid_layout_forensics_verdict is not None + + rogue = torch.nn.Parameter(torch.full((2, 2), 7.0)) + rogue_before = rogue.detach().clone() + optimizer.muon.param_groups[0]["params"][0] = rogue + + with pytest.raises(RuntimeError, match="finalized parameter layout changed"): + _step_with_grads(optimizer, (bias,)) + assert torch.equal(rogue, rogue_before) + + +def test_finalized_hybrid_detects_swapped_child_binding_with_warm_verdict(): + optimizer, matrix, bias = _finalized_hybrid() + _step_with_grads(optimizer, (matrix, bias)) # warm the cached verdict + + # Swap a child's private finalized manifest for a foreign one after + # finalization: the composite fast token captures the child manifest by + # identity, so the swap forces a full rebuild that rejects the layout. + optimizer.muon._gefen_sharding_manifest = optimizer.backup._gefen_sharding_manifest + + with pytest.raises(RuntimeError, match="finalized parameter layout changed"): + _step_with_grads(optimizer, (matrix, bias)) + + +def test_finalized_hybrid_rechecks_layout_after_closure(): + optimizer, matrix, bias = _finalized_hybrid() + _step_with_grads(optimizer, (matrix, bias)) # warm the cached verdict + + rogue = torch.nn.Parameter(torch.full((2, 2), 3.0)) + rogue_before = rogue.detach().clone() + matrix_before = matrix.detach().clone() + + def mutate_layout(): + # Tamper with a public child container between the pre- and + # post-closure guards; the post-closure guard must catch it before any + # child steps. + optimizer.muon.param_groups[0]["params"][0] = rogue + rogue.grad = torch.ones_like(rogue) + bias.grad = torch.ones_like(bias) + return torch.tensor(1.0) + + with pytest.raises(RuntimeError, match="finalized parameter layout changed"): + optimizer.step(mutate_layout) + assert torch.equal(rogue, rogue_before) + assert torch.equal(matrix, matrix_before) + + +def test_composite_registry_replacement_is_detected_with_warm_verdict(): + optimizer, matrix, bias = _finalized_hybrid() + _step_with_grads(optimizer, (matrix, bias)) # warm the cached verdict + + # Replacing the composite ownership registry preserves no container + # identity the fast token recorded, so the next step guard rebuilds and + # rejects. + optimizer._state_param_owner = {} + + with pytest.raises(RuntimeError, match="finalized parameter layout changed"): + _step_with_grads(optimizer, (matrix, bias)) From 4cfb0691cd656ea9d022ce11b52b483e6b98663c Mon Sep 17 00:00:00 2001 From: thad0ctor Date: Mon, 13 Jul 2026 14:09:41 -0700 Subject: [PATCH 34/52] Resolve value exports in the CI import smoke MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The import smoke derived optimizer names from __all__ and read __name__ on each, skipping only kernels and __version__. The contracts layer adds public integer schema constants (CONTRACT_SCHEMA_VERSION, CANONICAL_STATE_FORMAT_VERSION, PORTABLE_STATE_FORMAT_VERSION, IDENTITY_SCHEMA_VERSION), which have no __name__, so the step raised AttributeError. Resolve every export instead — which also exercises the lazy __getattr__ hooks — and fall back to repr for value exports. --- .github/workflows/ci.yml | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 38e1060..155776c 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -118,11 +118,12 @@ jobs: run: | python - <<'PY' import gefen - # Derive the public optimizers from __all__ so this adapts per branch - # (e.g. main has no GefenMuonHybrid). Skip non-class exports: "kernels" - # is a submodule and "__version__" is a string, neither has __name__. - names = [n for n in gefen.__all__ if n not in ("kernels", "__version__")] - loaded = {n: getattr(gefen, n).__name__ for n in names} + # Resolve every public export in __all__ so this adapts per branch and + # exercises the lazy __getattr__ hooks. Exports are a mix of classes, + # a submodule ("kernels"), the version string, and integer schema + # constants (e.g. CONTRACT_SCHEMA_VERSION), so report __name__ where it + # exists and fall back to repr for value exports. + loaded = {n: getattr(getattr(gefen, n), "__name__", repr(getattr(gefen, n))) for n in gefen.__all__} assert "Gefen" in loaded, "Gefen must always be importable" print("import OK ->", loaded) PY From 993ec96b5a03754b32cdf62de8ec077edac93672 Mon Sep 17 00:00:00 2001 From: thad0ctor Date: Mon, 13 Jul 2026 14:37:50 -0700 Subject: [PATCH 35/52] Run GPU test job on PRs to main with fail-fast timeout --- .github/workflows/ci.yml | 14 +++++++++++--- 1 file changed, 11 insertions(+), 3 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 155776c..8cd3773 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -136,10 +136,18 @@ jobs: gpu-tests: name: GPU kernel parity tests - # Only when explicitly requested via the Actions tab. Skipped on push/PR so the - # job never sits pending forever when no GPU runner is connected. - if: github.event_name == 'workflow_dispatch' && inputs.run_gpu_tests + # Runs on every PR targeting main so the CUDA-only fail-before-mutation / + # kernel-parity tests actually gate merges, plus on-demand via the Actions tab. + # (Pushes to other branches and non-main PRs still skip it.) + if: >- + (github.event_name == 'pull_request' && github.base_ref == 'main') || + (github.event_name == 'workflow_dispatch' && inputs.run_gpu_tests) runs-on: [self-hosted, gpu] + # Tradeoff: the self-hosted GPU runner must be online for PRs to main. If it + # is offline the job is not picked up and, after this timeout, fails fast so + # the PR shows a failed GPU check (a visible, honest signal) instead of either + # sitting pending forever or silently skipping GPU coverage. + timeout-minutes: 30 steps: - uses: actions/checkout@34e114876b0b11c390a56381ad16ebd13914f8d5 # v4.3.1 with: From ba5e0905f7dbb84c9637899135690d3bbfb04f3d Mon Sep 17 00:00:00 2001 From: thad0ctor Date: Mon, 13 Jul 2026 14:46:41 -0700 Subject: [PATCH 36/52] Make GefenMuonHybrid.load_state_dict two-phase fail-before-mutation The composite native load committed the muon child (self.muon.load_state_dict) before the backup half was validated, then recovered a backup rejection with a second full self.muon.load_state_dict(snapshot) reload. That rollback could itself raise (e.g. CUDA OOM re-staging every muon state tensor), masking the real backup error and leaving a half-loaded hybrid (new muon, old backup) -- unlike every other mutating entry point on this branch, which stages fully then publishes through non-throwing swaps. Give the composite genuine two-phase semantics: validate and stage BOTH halves before publishing EITHER, so a rejection on either half leaves both live children byte-for-byte untouched. Each Gefen child already loads atomically; expose that as a reusable primitive by splitting Gefen.load_state_dict into _prepare_load_state_dict (runs the load pre-hooks and stages an isolated shadow, mutating nothing) and _publish_load_state_dict (commits through the existing non-throwing dict swaps and runs post-hooks). The hybrid stages the muon child, then: * Gefen backup: stages the backup too (all rejection happens here, before any mutation), then publishes both through non-throwing swaps; * AdamW backup: torch's Optimizer.load_state_dict validates group structure and casts every tensor before its single __setstate__, so it is itself fail-before-mutation -- commit it first while the muon child is only staged, then publish the muon swap (which cannot fail). Chosen mechanism is process-group-safe: staging goes through Gefen._stage_load_state_dict, which 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. The removed snapshot path's copy.deepcopy is gone (import dropped); we deliberately avoid deepcopy of a child or its process groups. tests/test_hybrid_load_atomicity.py pins the contract: a backup rejection (mismatched param-group count / foreign schema) must not reload the muon child (byte-for-byte and object-identity unchanged), and an original backup error must surface over a failing rollback reload rather than be masked. Covers Gefen and AdamW backups; all three fail on the pre-fix code. --- src/gefen/gefen.py | 20 +++- src/gefen/hybrid.py | 62 ++++++++--- tests/test_hybrid_load_atomicity.py | 167 ++++++++++++++++++++++++++++ 3 files changed, 231 insertions(+), 18 deletions(-) create mode 100644 tests/test_hybrid_load_atomicity.py diff --git a/src/gefen/gefen.py b/src/gefen/gefen.py index 3517fd2..36e924f 100644 --- a/src/gefen/gefen.py +++ b/src/gefen/gefen.py @@ -8315,6 +8315,20 @@ def _pack_legacy_param_groups_for_load(self, state_dict): def load_state_dict(self, state_dict): """Atomically restore Gefen state between the public load hooks.""" + staged = self._prepare_load_state_dict(state_dict) + self._publish_load_state_dict(staged) + + def _prepare_load_state_dict(self, state_dict): + """Validate and stage a restore without mutating live state (phase one). + + Runs the load pre-hooks and stages the complete restore, returning an + isolated shadow. Nothing on the live optimizer is mutated, so a rejection + here (schema/layout mismatch, foreign backend, corrupted state) is a true + no-op. Composite owners (GefenMuonHybrid) call this to validate every + child before publishing any of them; ``load_state_dict`` pairs it with + ``_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(): @@ -8322,7 +8336,11 @@ def load_state_dict(self, state_dict): if hook_result is not None: state_dict = hook_result self._assert_finalized_binding_layout(full=True) - staged = self._stage_load_state_dict(state_dict) + return self._stage_load_state_dict(state_dict) + + 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) diff --git a/src/gefen/hybrid.py b/src/gefen/hybrid.py index b6a72f4..88d3c27 100644 --- a/src/gefen/hybrid.py +++ b/src/gefen/hybrid.py @@ -54,7 +54,6 @@ ``normuon=True`` as the hybrid-specific defaults. """ -import copy import logging from collections import OrderedDict @@ -1487,23 +1486,52 @@ def load_state_dict(self, state_dict): checkpoint_backup_optimizer, self.backup_optimizer ) ) - # Load the two children atomically: torch's optimizer loader can still - # reject a child on its internal layout (e.g. a different param count), - # and a backup failure after the muon child loaded would leave the - # hybrid half-loaded (new muon, old backup). Snapshot the muon child's - # state first and roll it back if the backup load raises. - muon_snapshot = None + # 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: - if self.backup is not None: - muon_snapshot = copy.deepcopy(self.muon.state_dict()) - self.muon.load_state_dict(state_dict["muon"]) - if self.backup is not None: - try: - self.backup.load_state_dict(state_dict["backup"]) - except Exception: - if muon_snapshot is not None: - self.muon.load_state_dict(muon_snapshot) - raise + 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) diff --git a/tests/test_hybrid_load_atomicity.py b/tests/test_hybrid_load_atomicity.py new file mode 100644 index 0000000..e71431e --- /dev/null +++ b/tests/test_hybrid_load_atomicity.py @@ -0,0 +1,167 @@ +"""Fail-before-mutation tests for GefenMuonHybrid.load_state_dict. + +The composite load restores two children (the GefenMuon half and the backup +half). A rejection on *either* half must leave *both* live children +byte-for-byte untouched -- the same two-phase (stage-then-publish) contract +every other mutating entry point on this branch honours. + +These tests pin the failure mode of the previous implementation, which committed +the muon child first and only then validated/loaded the backup, recovering via a +second full ``muon.load_state_dict(snapshot)`` reload. That reload could itself +raise (e.g. CUDA OOM re-staging every muon tensor), masking the real backup error +and leaving a half-loaded hybrid. They run on CPU (fused=False). +""" +import copy +import warnings + +import pytest +import torch +import torch.nn as nn + +from gefen import GefenMuonHybrid + + +class TinyLM(nn.Module): + def __init__(self, vocab=32, dim=16): + super().__init__() + self.embed = nn.Embedding(vocab, dim) + self.hidden = nn.Linear(dim, dim, bias=True) + self.norm = nn.LayerNorm(dim) + self.lm_head = nn.Linear(dim, vocab, bias=False) + + def forward(self, idx): + return self.lm_head(self.norm(self.hidden(self.embed(idx)))) + + +def _build(backup_optimizer="gefen"): + torch.manual_seed(0) + model = TinyLM() + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + opt = GefenMuonHybrid( + model, lr=1e-3, fused=False, backup_optimizer=backup_optimizer + ) + return opt, model + + +def _step(opt, model): + idx = torch.randint(0, 32, (4, 6)) + tgt = torch.randint(0, 32, (4, 6)) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + opt.zero_grad() + nn.functional.cross_entropy( + model(idx).reshape(-1, 32), tgt.reshape(-1) + ).backward() + opt.step() + + +def _muon_param_and_state(opt): + param = opt.muon.param_groups[0]["params"][0] + state = opt.muon.state[param] + tensors = {k: v for k, v in state.items() if torch.is_tensor(v)} + assert tensors, "expected the stepped muon child to hold momentum tensors" + return param, tensors + + +def _backup_with_mismatched_group_count(state_dict): + """A backup half the live backup's loader rejects (extra param group).""" + bad = copy.deepcopy(state_dict) + bad["backup"]["param_groups"].append( + copy.deepcopy(bad["backup"]["param_groups"][0]) + ) + return bad + + +def test_backup_rejection_leaves_muon_child_byte_for_byte_untouched(): + opt, model = _build() + _step(opt, model) + bad = _backup_with_mismatched_group_count(opt.state_dict()) + + param, before = _muon_param_and_state(opt) + before_values = {k: v.detach().clone() for k, v in before.items()} + before_ids = {k: id(v) for k, v in before.items()} + + # Detect any reload of the muon child. The two-phase load validates the + # backup half *before* publishing muon, so muon must never be loaded (nor + # loaded-then-rolled-back, as the previous snapshot/reload path did). + calls = {"n": 0} + original_load = opt.muon.load_state_dict + + def counting_load(sd): + calls["n"] += 1 + return original_load(sd) + + opt.muon.load_state_dict = counting_load + + with pytest.raises(ValueError, match="different number of parameter groups"): + opt.load_state_dict(bad) + + assert calls["n"] == 0, "backup rejection must not reload the muon child" + + param, after = _muon_param_and_state(opt) + for key, value in before_values.items(): + assert torch.equal(value, after[key]), f"muon state {key!r} changed value" + # Deep value comparison is the untouched contract; object identity is an + # extra witness that no restage happened (a rollback reload swaps objects). + assert {k: id(v) for k, v in after.items()} == before_ids + + +def test_backup_rejection_surfaces_over_a_failing_rollback_reload(): + opt, model = _build() + _step(opt, model) + bad = _backup_with_mismatched_group_count(opt.state_dict()) + + param, before = _muon_param_and_state(opt) + before_values = {k: v.detach().clone() for k, v in before.items()} + + # Simulate the previous implementation's rollback masking: the muon child's + # first load (the real one) succeeds, but the second (the rollback reload) + # blows up -- as re-staging every muon tensor on CUDA can. The *original* + # backup error must still be what surfaces, and muon must stay untouched. + calls = {"n": 0} + original_load = opt.muon.load_state_dict + + def boom_on_rollback(sd): + calls["n"] += 1 + if calls["n"] >= 2: + raise RuntimeError("rollback reload OOM") + return original_load(sd) + + opt.muon.load_state_dict = boom_on_rollback + + with pytest.raises(ValueError, match="different number of parameter groups"): + opt.load_state_dict(bad) + + param, after = _muon_param_and_state(opt) + for key, value in before_values.items(): + assert torch.equal(value, after[key]), f"muon state {key!r} changed value" + + +def test_adamw_backup_rejection_leaves_muon_child_untouched(): + opt, model = _build(backup_optimizer="adamw") + _step(opt, model) + bad = _backup_with_mismatched_group_count(opt.state_dict()) + + param, before = _muon_param_and_state(opt) + before_values = {k: v.detach().clone() for k, v in before.items()} + before_ids = {k: id(v) for k, v in before.items()} + + calls = {"n": 0} + original_load = opt.muon.load_state_dict + + def counting_load(sd): + calls["n"] += 1 + return original_load(sd) + + opt.muon.load_state_dict = counting_load + + with pytest.raises(ValueError, match="different number of parameter groups"): + opt.load_state_dict(bad) + + assert calls["n"] == 0, "backup rejection must not reload the muon child" + + param, after = _muon_param_and_state(opt) + for key, value in before_values.items(): + assert torch.equal(value, after[key]), f"muon state {key!r} changed value" + assert {k: id(v) for k, v in after.items()} == before_ids From 1135298b2d84da90d56365ab62cba6ac29475589 Mon Sep 17 00:00:00 2001 From: thad0ctor Date: Mon, 13 Jul 2026 14:50:23 -0700 Subject: [PATCH 37/52] Revert "Run GPU test job on PRs to main with fail-fast timeout" This reverts commit 993ec96b5a03754b32cdf62de8ec077edac93672. --- .github/workflows/ci.yml | 14 +++----------- 1 file changed, 3 insertions(+), 11 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 8cd3773..155776c 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -136,18 +136,10 @@ jobs: gpu-tests: name: GPU kernel parity tests - # Runs on every PR targeting main so the CUDA-only fail-before-mutation / - # kernel-parity tests actually gate merges, plus on-demand via the Actions tab. - # (Pushes to other branches and non-main PRs still skip it.) - if: >- - (github.event_name == 'pull_request' && github.base_ref == 'main') || - (github.event_name == 'workflow_dispatch' && inputs.run_gpu_tests) + # Only when explicitly requested via the Actions tab. Skipped on push/PR so the + # job never sits pending forever when no GPU runner is connected. + if: github.event_name == 'workflow_dispatch' && inputs.run_gpu_tests runs-on: [self-hosted, gpu] - # Tradeoff: the self-hosted GPU runner must be online for PRs to main. If it - # is offline the job is not picked up and, after this timeout, fails fast so - # the PR shows a failed GPU check (a visible, honest signal) instead of either - # sitting pending forever or silently skipping GPU coverage. - timeout-minutes: 30 steps: - uses: actions/checkout@34e114876b0b11c390a56381ad16ebd13914f8d5 # v4.3.1 with: From 6d6501dc55543da26e6958d7d74c526881b218a2 Mon Sep 17 00:00:00 2001 From: thad0ctor Date: Mon, 13 Jul 2026 14:51:28 -0700 Subject: [PATCH 38/52] Cover state offload and movement in the local GPU release gate MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The state-offload and state-movement features are exercised only by CUDA-gated tests, which skip in the CPU-only push/PR CI. Their fail-before- mutation and device-transfer invariants were therefore verified by no gate. Add tests/test_state_offload.py and tests/test_state_movement.py to the mandatory two-GPU release gate list (72 passed, 0 skips on two devices), so the release gate — not a hosted CI runner — covers them. --- scripts/release_gpu_gate.sh | 2 ++ 1 file changed, 2 insertions(+) diff --git a/scripts/release_gpu_gate.sh b/scripts/release_gpu_gate.sh index 8fb15fc..faa4a7b 100755 --- a/scripts/release_gpu_gate.sh +++ b/scripts/release_gpu_gate.sh @@ -216,6 +216,8 @@ GEFEN_VERBOSE_BUILD=1 \ tests/test_gefen_fsdp2_checkpoint.py \ tests/test_muon_distributed_checkpoint_safety.py \ tests/test_muon_grad_presence.py \ + tests/test_state_offload.py \ + tests/test_state_movement.py \ tests/test_step_preflight_atomicity.py \ tests/test_muon_fsdp2_parity.py::test_muon_fsdp2_fused_cuda_parity \ tests/test_muon_fsdp2_parity.py::test_muon_fsdp2_fused_multirank_parity \ From c0e2fb30067b1c782789d7cc4bcc04c944a54f06 Mon Sep 17 00:00:00 2001 From: thad0ctor Date: Mon, 13 Jul 2026 15:56:32 -0700 Subject: [PATCH 39/52] Synchronize pre-collective step failures across the codebook scope GefenMuon.step and GefenMuonHybrid.step ran the capturable-readiness checks, the instance step pre-hooks, and the closure BEFORE any scope synchronization. Because capture readiness depends on rank-local gradients and the closure is user code, any of these can fail on only a subset of scope members. When it did, the failing member exited step() while its peers proceeded into the scoped operation-header all_gather, the composite preflight synchronization, or a child's scoped step collectives -- and hung. Capture these preamble failures and route them through _synchronize_codebook_scope_failure on the shared codebook binding before any member enters a later scoped collective, so the failure raises on every member together (mirroring the existing gradient-preflight synchronization). The finalized-layout and runtime-process-group guards stay local: they establish the very binding used to synchronize. Add asymmetric two-rank gloo regressions (a rank-0 closure failure) for both GefenMuon and the hybrid; without the fix the peer strands in a scoped collective until the gloo application timeout. Also fail the hybrid scoped-failure harness when a worker hangs or exits nonzero so an unclean distributed shutdown can no longer pass silently. --- src/gefen/gefen_muon.py | 28 ++++- src/gefen/hybrid.py | 54 ++++++--- tests/test_codebook_scope_distributed.py | 109 +++++++++++++++++++ tests/test_hybrid_scoped_failure_protocol.py | 47 +++++++- 4 files changed, 214 insertions(+), 24 deletions(-) diff --git a/src/gefen/gefen_muon.py b/src/gefen/gefen_muon.py index e3454b3..e967794 100644 --- a/src/gefen/gefen_muon.py +++ b/src/gefen/gefen_muon.py @@ -2864,12 +2864,30 @@ def step(self, closure=None): "GefenMuon whole-parameter owner stepping requires the separate " "explicit process-group codebook scope" ) - self._assert_capturable_if_capturing() - self._assert_codebook_capture_ready() + # Capture-readiness depends on rank-local gradients and the closure is + # user code, so both can fail on only a subset of scope members. Capture + # such a failure and synchronize it on the codebook binding BEFORE any + # member enters the scoped operation-header or later collectives, so it + # raises on every member together instead of stranding peers inside a + # collective. The finalized-layout and runtime-process-group guards above + # stay local: they establish the very binding used to synchronize. loss = None - if closure is not None: - with torch.enable_grad(): - loss = closure() + try: + self._assert_capturable_if_capturing() + self._assert_codebook_capture_ready() + if closure is not None: + with torch.enable_grad(): + loss = closure() + local_preamble_error = None + except Exception as exc: + loss = None + local_preamble_error = exc + if self._gefen_codebook_process_group is not None: + self._synchronize_codebook_scope_failure( + local_preamble_error, "step preamble" + ) + elif local_preamble_error is not None: + raise local_preamble_error self._assert_finalized_binding_layout() self._assert_runtime_codebook_process_group() diff --git a/src/gefen/hybrid.py b/src/gefen/hybrid.py index 88d3c27..e023f9c 100644 --- a/src/gefen/hybrid.py +++ b/src/gefen/hybrid.py @@ -1251,7 +1251,6 @@ def _assert_capturable_devices_if_capturing(self) -> None: def step(self, closure=None): self._assert_finalized_binding_layout() - self._assert_capturable_devices_if_capturing() # Dispatch the INSTANCE step hooks around the composite step, mirroring # torch.optim.Optimizer.profile_hook_step exactly: hooks receive # (optimizer, args, kwargs) where args are the raw step() call args @@ -1260,26 +1259,45 @@ def step(self, closure=None): # _patch_step_function); without this, registered hooks would silently # never fire. GLOBAL step hooks are NOT dispatched here -- they fire on # each child sub-optimizer's (wrapped) step; see the class docstring. + # + # The capturable-device check, the user step pre-hooks, and the closure + # can all fail on only a subset of scope members. Capture such a failure + # and synchronize it on the one shared codebook binding BEFORE any member + # enters the composite preflight synchronize or a child's scoped step + # collectives, so it raises on every member together instead of stranding + # peers. The finalized-layout guard above stays local: it establishes the + # binding used to synchronize. args = (self, closure) if closure is not None else (self,) kwargs = {} - for pre_hook in self._optimizer_step_pre_hooks.values(): - result = pre_hook(self, args, kwargs) - if result is not None: - if isinstance(result, tuple) and len(result) == 2: - args, kwargs = result - else: - raise RuntimeError( - f"{self.__class__.__name__}.step pre hook must return None " - f"or a tuple of (new_args, new_kwargs), but got {result}." - ) - # Re-read the closure from the (possibly hook-rewritten) call args, as - # torch's wrapper would by calling step(*args, **kwargs). - closure = args[1] if len(args) > 1 else kwargs.get("closure") - loss = None - if closure is not None: - with torch.enable_grad(): - loss = closure() + try: + self._assert_capturable_devices_if_capturing() + for pre_hook in self._optimizer_step_pre_hooks.values(): + result = pre_hook(self, args, kwargs) + if result is not None: + if isinstance(result, tuple) and len(result) == 2: + args, kwargs = result + else: + raise RuntimeError( + f"{self.__class__.__name__}.step pre hook must return None " + f"or a tuple of (new_args, new_kwargs), but got {result}." + ) + # Re-read the closure from the (possibly hook-rewritten) call args, as + # torch's wrapper would by calling step(*args, **kwargs). + closure = args[1] if len(args) > 1 else kwargs.get("closure") + if closure is not None: + with torch.enable_grad(): + loss = closure() + local_preamble_error = None + except Exception as exc: + loss = None + local_preamble_error = exc + if self._hybrid_codebook_process_group is not None: + self._synchronize_codebook_scope_failure( + local_preamble_error, "step preamble" + ) + elif local_preamble_error is not None: + raise local_preamble_error self._assert_finalized_binding_layout() # Composite structural preflight, atomically over BOTH children before # either child steps. Under a shared codebook scope the failure is diff --git a/tests/test_codebook_scope_distributed.py b/tests/test_codebook_scope_distributed.py index f0fe80c..a36699d 100644 --- a/tests/test_codebook_scope_distributed.py +++ b/tests/test_codebook_scope_distributed.py @@ -1239,3 +1239,112 @@ def test_nccl_scope_uses_explicit_collective_device_with_empty_nonowner(): process.join(timeout=5) if os.path.exists(init_file): os.unlink(init_file) + + +def _closure_preamble_worker(rank, world, init_file, queue): + # GefenMuon.step runs the closure BEFORE any scope synchronization. A + # rank-local closure failure must raise on every scope member together + # instead of leaving the failing rank to exit step() while the peer enters + # the scoped operation-header / synchronization collectives and hangs. + try: + dist.init_process_group( + "gloo", + init_method="file://{}".format(init_file), + rank=rank, + world_size=world, + timeout=timedelta(seconds=45), + ) + members = tuple("rank:{}".format(index) for index in range(world)) + group = ProcessGroupIdentity("data_parallel", members) + runtime_group = dist.group.WORLD + + matrix = torch.nn.Parameter(torch.zeros(2, 2)) + optimizer = GefenMuon([("matrix", matrix)], fused=False) + identity = ParameterIdentity("Matrix", (2, 2)) + records = tuple(_replicated(identity, group, member) for member in members) + _finalize( + optimizer, + matrix, + records[rank], + ShardingManifest(records), + _binding(group, rank, runtime_group), + ) + matrix.grad = torch.tensor([[1.0, -2.0], [3.0, -4.0]]) + + def closure(): + if rank == 0: + raise RuntimeError("closure boom on rank:0") + return torch.tensor(1.0) + + try: + optimizer.step(closure) + message = None + except RuntimeError as exc: + message = str(exc) + # The synchronized failure must leave the step fully un-run on both ranks. + untouched = ( + optimizer._gefen_global_step == 0 + and optimizer._gefen_codebook is None + and optimizer.state[matrix] == {"name": "matrix"} + and torch.equal(matrix, torch.zeros(2, 2)) + ) + queue.put({"rank": rank, "message": message, "untouched": untouched}) + except Exception as exc: + queue.put({"rank": rank, "error": repr(exc)}) + finally: + if dist.is_available() and dist.is_initialized(): + dist.destroy_process_group() + + +def _run_closure_preamble_workers(world=2): + context = mp.get_context("spawn") + queue = context.Queue() + fd, init_file = tempfile.mkstemp(prefix="gefen-codebook-closure-preamble-") + os.close(fd) + os.unlink(init_file) + processes = [ + context.Process( + target=_closure_preamble_worker, + args=(rank, world, init_file, queue), + ) + for rank in range(world) + ] + try: + for process in processes: + process.start() + results = [] + try: + for _ in processes: + results.append(queue.get(timeout=120)) + except Exception: + pass + for process in processes: + process.join(timeout=10) + for process in processes: + if process.is_alive(): + process.terminate() + process.join(timeout=5) + pytest.fail("closure-preamble worker hung") + assert process.exitcode == 0 + return sorted(results, key=lambda item: item["rank"]) + finally: + for process in processes: + if process.is_alive(): + process.terminate() + process.join(timeout=5) + if os.path.exists(init_file): + os.unlink(init_file) + + +@pytest.mark.skipif( + not dist.is_available() or not dist.is_gloo_available(), + reason="scoped closure-preamble coverage requires Gloo", +) +def test_scoped_step_closure_failure_raises_symmetrically_across_the_scope(): + results = _run_closure_preamble_workers() + assert len(results) == 2, results + assert all("error" not in item for item in results), results + assert results[0]["message"] is not None and results[1]["message"] is not None, results + assert "closure boom on rank:0" in results[0]["message"] + assert "step preamble failed on another process-group member" in results[1]["message"] + assert all(item["untouched"] for item in results), results diff --git a/tests/test_hybrid_scoped_failure_protocol.py b/tests/test_hybrid_scoped_failure_protocol.py index 4daa84f..dc0f73e 100644 --- a/tests/test_hybrid_scoped_failure_protocol.py +++ b/tests/test_hybrid_scoped_failure_protocol.py @@ -354,6 +354,30 @@ def _preflight_divergent_result(rank, group): } +def _closure_divergent_result(rank, group): + # The closure runs BEFORE any scope synchronization. A rank-local closure + # failure must raise on every scope member together instead of leaving the + # failing rank to exit step() while the peer enters the composite preflight + # synchronize / a child's scoped step collectives and hangs. + optimizer, muon_parameter, backup_parameter, backup_shard = _make_scoped_hybrid(rank, group) + _set_local_gradients(muon_parameter, backup_parameter, backup_shard) + + def closure(): + if rank == 0: + raise RuntimeError("closure boom on rank:0") + return torch.tensor(1.0) + + try: + optimizer.step(closure) + message = None + except RuntimeError as exc: + message = str(exc) + return { + "message": message, + "untouched": _untouched(optimizer, muon_parameter, backup_parameter), + } + + def _distributed_worker(rank, init_file, result_queue): try: dist.init_process_group( @@ -372,6 +396,8 @@ def _distributed_worker(rank, init_file, result_queue): dist.barrier() preflight = _preflight_divergent_result(rank, group) dist.barrier() + closure_divergent = _closure_divergent_result(rank, group) + dist.barrier() result_queue.put( { "rank": rank, @@ -379,6 +405,7 @@ def _distributed_worker(rank, init_file, result_queue): "amp_group_wide": amp_group_wide, "amp_finite": amp_finite, "preflight": preflight, + "closure_divergent": closure_divergent, } ) except BaseException: @@ -412,6 +439,12 @@ def _run_distributed_workers(): pass for process in processes: process.join(timeout=10) + # A worker that queued its result but then hangs in + # destroy_process_group() (or exits nonzero) means the release-critical + # synchronization did not shut down cleanly; fail rather than silently + # terminate it below. + hung = [index for index, process in enumerate(processes) if process.is_alive()] + exit_codes = [process.exitcode for process in processes] finally: for process in processes: if process.is_alive(): @@ -421,7 +454,9 @@ def _run_distributed_workers(): result_queue.join_thread() if os.path.exists(init_file): os.unlink(init_file) - assert len(results) == _WORLD, (results, [process.exitcode for process in processes]) + assert not hung, ("hybrid scoped-failure workers hung", hung, exit_codes) + assert all(code == 0 for code in exit_codes), exit_codes + assert len(results) == _WORLD, (results, exit_codes) return sorted(results, key=lambda item: item["rank"]) @@ -470,6 +505,16 @@ def test_hybrid_step_synchronizes_amp_and_preflight_across_the_scope(): assert "gradient preflight failed on another process-group member" in preflight[1]["message"] assert all(item["untouched"] for item in preflight), preflight + # A rank-local closure failure (before any scope synchronization) raises on + # every scope member together through the composite preamble synchronization + # instead of stranding the peer inside a later scoped collective. Both ranks + # exit with an error and leave both children untouched. + closure_divergent = [result["closure_divergent"] for result in results] + assert closure_divergent[0]["message"] is not None and closure_divergent[1]["message"] is not None, closure_divergent + assert "closure boom on rank:0" in closure_divergent[0]["message"] + assert "step preamble failed on another process-group member" in closure_divergent[1]["message"] + assert all(item["untouched"] for item in closure_divergent), closure_divergent + def _make_plain_hybrid(): matrix = torch.nn.Parameter(_muon_initial().clone()) From 181c3ecf85e1ed65a0297ce49d714bd538f33709 Mon Sep 17 00:00:00 2001 From: thad0ctor Date: Mon, 13 Jul 2026 15:56:38 -0700 Subject: [PATCH 40/52] Intersect composite checkpoint guarantees across children A composite canonical-global save/load processes EVERY child, so the Hybrid may advertise a checkpoint topology/exactness guarantee only when both children support it. _hybrid_portable_contract_support unioned the children's same_topology, topology_changing, and topology_change_kinds, which could advertise a guarantee that one child does not honor. Intersect them instead. (This is the opposite of the per-routed-parameter TRAINING claims, which remain legitimately unioned.) Add a regression that a Hybrid with children of differing canonical-global support advertises only the intersection. --- src/gefen/portable_hybrid.py | 37 +++++++++++---- tests/test_portable_hybrid_runtime.py | 65 +++++++++++++++++++++++++++ 2 files changed, 93 insertions(+), 9 deletions(-) diff --git a/src/gefen/portable_hybrid.py b/src/gefen/portable_hybrid.py index 3cb0ffc..bf6b117 100644 --- a/src/gefen/portable_hybrid.py +++ b/src/gefen/portable_hybrid.py @@ -380,9 +380,15 @@ def _hybrid_portable_contract_support(optimizer): optimizer, binding, ) - same_topology = set() - topology_changing = set() - topology_change_kinds = set() + # A composite canonical-global save/load processes EVERY child, so the + # Hybrid may only advertise a checkpoint guarantee that every child + # supports: intersect (not union) the children's guarantee sets. (This + # is the opposite of the per-routed-parameter TRAINING claims, which are + # legitimately unioned.) None marks "no child seen yet" so the first + # child seeds each set and later children narrow it. + same_topology = None + topology_changing = None + topology_change_kinds = None for _role, child in children: supports = tuple( support @@ -392,13 +398,26 @@ def _hybrid_portable_contract_support(optimizer): if len(supports) != 1: return frozenset(), frozenset(), frozenset() support = supports[0] - same_topology.update(support.same_topology) - topology_changing.update(support.topology_changing) - topology_change_kinds.update(support.topology_change_kinds) + child_same = set(support.same_topology) + child_changing = set(support.topology_changing) + child_kinds = set(support.topology_change_kinds) + same_topology = ( + child_same if same_topology is None else same_topology & child_same + ) + topology_changing = ( + child_changing + if topology_changing is None + else topology_changing & child_changing + ) + topology_change_kinds = ( + child_kinds + if topology_change_kinds is None + else topology_change_kinds & child_kinds + ) return ( - frozenset(same_topology), - frozenset(topology_changing), - frozenset(topology_change_kinds), + frozenset(same_topology or ()), + frozenset(topology_changing or ()), + frozenset(topology_change_kinds or ()), ) except Exception: return frozenset(), frozenset(), frozenset() diff --git a/tests/test_portable_hybrid_runtime.py b/tests/test_portable_hybrid_runtime.py index f9e8a16..0d43303 100644 --- a/tests/test_portable_hybrid_runtime.py +++ b/tests/test_portable_hybrid_runtime.py @@ -389,3 +389,68 @@ def test_adamw_backed_hybrid_remains_explicitly_nonportable(): transaction_id="adamw-hybrid-reject-v1", limits=_limits(), ) + + +def test_hybrid_canonical_global_checkpoint_guarantees_are_intersected(monkeypatch): + # A composite canonical-global save/load processes EVERY child, so the + # Hybrid may advertise a checkpoint guarantee only when both children + # support it. The children's guarantee sets must intersect, not union. + from types import SimpleNamespace + + from gefen.contracts import ( + CheckpointSupport, + ProcessGroupScope, + TopologyChange, + ) + from gefen.portable_hybrid import _hybrid_portable_contract_support + + source, _matrix, _bias, _binding = _initialized_source() + + def _support(same_topology, kinds): + return CheckpointSupport( + transport=CheckpointTransport.CANONICAL_GLOBAL, + same_topology=same_topology, + topology_changing=frozenset({ParameterLayout.REPLICATED}), + process_group_scope=ProcessGroupScope.DEFAULT_WORLD, + topology_change_kinds=kinds, + atomic_load=True, + ) + + muon_support = _support( + frozenset( + {ParameterLayout.REPLICATED, ParameterLayout.FLATTENED_ELEMENT_SHARD} + ), + frozenset( + { + TopologyChange.WORLD_SIZE_OWNER_REDISTRIBUTION, + TopologyChange.PLACEMENT_RESHARD, + } + ), + ) + backup_support = _support( + frozenset({ParameterLayout.REPLICATED}), + frozenset({TopologyChange.WORLD_SIZE_OWNER_REDISTRIBUTION}), + ) + + def _stub(support): + # Patch at the class level: portable readiness rejects instance-level + # method shadows. The muon and backup children are distinct classes. + return lambda _self: SimpleNamespace( + capabilities=SimpleNamespace(checkpoints=(support,)) + ) + + assert type(source.muon) is not type(source.backup) + monkeypatch.setattr(type(source.muon), "optimizer_contract", _stub(muon_support)) + monkeypatch.setattr( + type(source.backup), "optimizer_contract", _stub(backup_support) + ) + + same_topology, topology_changing, topology_change_kinds = ( + _hybrid_portable_contract_support(source) + ) + + assert same_topology == frozenset({ParameterLayout.REPLICATED}) + assert topology_changing == frozenset({ParameterLayout.REPLICATED}) + assert topology_change_kinds == frozenset( + {TopologyChange.WORLD_SIZE_OWNER_REDISTRIBUTION} + ) From 006054f60246795a12124dfd07cefaf793992697 Mon Sep 17 00:00:00 2001 From: thad0ctor Date: Mon, 13 Jul 2026 15:56:43 -0700 Subject: [PATCH 41/52] Bound non-contiguous portable clone scratch to the chunk budget _portable_tensor_chunk_elements sized chunks on the element size alone, but the strided read path materializes an int64 linear index, its running quotient, and one int64 coordinate tensor per dimension. On a high-rank non-contiguous tensor the scratch oversubscribed the fixed clone budget and could OOM. Account for that per-element scratch in the chunk size, mirroring portable_wire._clone_tensor. Add a high-rank regression asserting the chunk stays within the budget. --- src/gefen/portable_schema.py | 10 +++++++++- tests/test_portable_schema.py | 31 +++++++++++++++++++++++++++++++ 2 files changed, 40 insertions(+), 1 deletion(-) diff --git a/src/gefen/portable_schema.py b/src/gefen/portable_schema.py index a50244b..1ffa415 100644 --- a/src/gefen/portable_schema.py +++ b/src/gefen/portable_schema.py @@ -53,7 +53,15 @@ def _portable_tensor_chunk_elements(value: torch.Tensor) -> int: budget = _PORTABLE_CLONE_CHUNK_BYTES if type(budget) is not int or budget <= 0: raise RuntimeError("portable clone chunk budget must be a positive int") - return max(1, budget // value.element_size()) + scratch_bytes_per_element = value.element_size() + if not value.is_contiguous(): + # The non-contiguous read path (_read_portable_tensor_chunk) materializes + # one int64 coordinate vector per dimension plus the int64 linear index + # and its running quotient. Size the chunk for that scratch too, so a + # high-rank strided tensor stays within the fixed clone budget instead of + # oversubscribing it. Mirrors portable_wire._clone_tensor. + scratch_bytes_per_element += 8 * (value.ndim + 2) + return max(1, budget // scratch_bytes_per_element) def _read_portable_tensor_chunk( diff --git a/tests/test_portable_schema.py b/tests/test_portable_schema.py index ebd746b..67c02b5 100644 --- a/tests/test_portable_schema.py +++ b/tests/test_portable_schema.py @@ -149,6 +149,37 @@ def tracked(value, start, stop): assert max(stop - start for start, stop in calls) <= 4 +def test_noncontiguous_chunk_elements_stay_within_the_clone_budget(): + budget = portable_schema_module._PORTABLE_CLONE_CHUNK_BYTES + # A high-rank non-contiguous view: the strided read path materializes an + # int64 linear index, its running quotient, and one int64 coordinate tensor + # per dimension. The chunk size must keep that scratch within the fixed clone + # budget rather than sizing on the element size alone (which would + # oversubscribe it and risk OOM on high-rank tensors). + base = torch.arange(2 * 3 * 4 * 5 * 6 * 7, dtype=torch.float32).reshape( + 2, 3, 4, 5, 6, 7 + ) + view = base.permute(5, 4, 3, 2, 1, 0) + assert not view.is_contiguous() + + chunk_elements = portable_schema_module._portable_tensor_chunk_elements(view) + scratch_per_element = view.element_size() + 8 * (view.ndim + 2) + assert chunk_elements >= 1 + assert chunk_elements * scratch_per_element <= budget + # Element-size-only sizing (the prior behavior) would blow past the budget. + assert (budget // view.element_size()) * scratch_per_element > budget + + contiguous = base.reshape(-1) + assert portable_schema_module._portable_tensor_chunk_elements(contiguous) == max( + 1, budget // contiguous.element_size() + ) + + # End-to-end the clone still reproduces a high-rank strided tensor exactly. + cloned = portable_schema_module._clone_portable_value(view, path="tensor") + assert torch.equal(cloned, view) + assert cloned.is_contiguous() + + def test_builder_and_normalizer_each_hash_the_tensor_tree_once(monkeypatch): calls = [] original = portable_schema_module._canonical_portable_state_digest From 5430974d420d261dc8ce158b177185f1c1ed3525 Mon Sep 17 00:00:00 2001 From: thad0ctor Date: Mon, 13 Jul 2026 15:56:51 -0700 Subject: [PATCH 42/52] Strengthen fail-before-mutation snapshots against in-place mutation The shared deep-snapshot helper compared tensors with torch.equal, which reports +0.0/-0.0 as equal and every NaN as unequal; compare dtype/layout/shape then raw bytes so a rejected transaction that flips a sign bit is caught and an unchanged NaN payload is not a false failure. Several atomicity snapshots retained mutable containers only by reference, so a failed transaction that mutated contents in place (preserving object identity) would still pass. Deep-snapshot and compare the contents: optimizer state / param groups / registries in the codebook-scope atomicity check, the hybrid _state_param_owner mapping, and the Gefen _gefen_shard_bindings registry and per-device caches. Also assert the rejected canonical-import commits are no-ops immediately after each RuntimeError -- before the recovery import -- so a fresh import can no longer mask a fail-before-mutation regression. The real implementation passes every strengthened assertion. --- tests/_state_snapshot.py | 20 +++++++++-- tests/test_codebook_scope_cpu.py | 7 ++++ tests/test_hybrid_rebinding.py | 9 +++++ tests/test_rebinding_cpu.py | 34 +++++++++++++++++++ .../test_scoped_collective_agreement_fixes.py | 8 +++++ 5 files changed, 76 insertions(+), 2 deletions(-) diff --git a/tests/_state_snapshot.py b/tests/_state_snapshot.py index 777ab9e..795e981 100644 --- a/tests/_state_snapshot.py +++ b/tests/_state_snapshot.py @@ -26,10 +26,26 @@ def _cloned(value): return copy.deepcopy(value) +def _tensors_bitwise_equal(live, expected): + # ``torch.equal`` reports +0.0/-0.0 as equal and every NaN as unequal, so it + # can miss a failed transaction that flips a sign bit and can spuriously fail + # on unchanged NaN payloads. A fail-before-mutation snapshot needs exact byte + # preservation, so compare dtype/layout/shape and then the raw bytes. + if ( + live.dtype != expected.dtype + or live.layout != expected.layout + or tuple(live.shape) != tuple(expected.shape) + ): + return False + live_bytes = live.detach().contiguous().view(torch.uint8) + expected_bytes = expected.detach().contiguous().view(torch.uint8) + return bool(torch.equal(live_bytes, expected_bytes)) + + def _nested_equal(live, expected): if torch.is_tensor(expected): assert torch.is_tensor(live) - assert torch.equal(live, expected) + assert _tensors_bitwise_equal(live, expected) return assert type(live) is type(expected) if isinstance(expected, dict): @@ -123,4 +139,4 @@ def assert_deep_state_snapshot(optimizer, snapshot): expected_options, ) for tensor, expected in snapshot["tensors"]: - assert torch.equal(tensor, expected) + assert _tensors_bitwise_equal(tensor, expected) diff --git a/tests/test_codebook_scope_cpu.py b/tests/test_codebook_scope_cpu.py index e3c7c13..b427d4a 100644 --- a/tests/test_codebook_scope_cpu.py +++ b/tests/test_codebook_scope_cpu.py @@ -24,6 +24,8 @@ ) import gefen.gefen as gefen_module +from _state_snapshot import assert_deep_state_snapshot, deep_state_snapshot + def _replicated_shard(parameter, group, member): coordinate = group.ordered_members.index(member) @@ -108,6 +110,10 @@ def _snapshot(optimizer): "codebook": optimizer._gefen_codebook, "binding": optimizer._gefen_codebook_process_group, "validated": optimizer._gefen_codebook_scope_validated, + # Bitwise clones of every reachable state tensor, param-group option, and + # mutable registry so a rejected transaction that mutates contents in + # place (without swapping the top-level container) is still caught. + "deep": deep_state_snapshot(optimizer), } @@ -120,6 +126,7 @@ def _assert_snapshot_identity(optimizer, snapshot): assert optimizer.__dict__.keys() == snapshot["dict"].keys() for key, value in snapshot["dict"].items(): assert optimizer.__dict__[key] is value + assert_deep_state_snapshot(optimizer, snapshot["deep"]) def _replace_native_guard(checkpoint, guard): diff --git a/tests/test_hybrid_rebinding.py b/tests/test_hybrid_rebinding.py index f74e01d..4902bda 100644 --- a/tests/test_hybrid_rebinding.py +++ b/tests/test_hybrid_rebinding.py @@ -125,6 +125,9 @@ def _snapshot(optimizer): return { "children": tuple((child, deep_state_snapshot(child)) for child in optimizer._subopts), "owner": optimizer._state_param_owner, + # id(param) -> (param, owning child); copy the mapping so an in-place + # remap that keeps the registry object identity is still caught. + "owner_contents": dict(optimizer._state_param_owner), "finalized": optimizer._hybrid_post_sharding_finalized, "manifest": optimizer._hybrid_sharding_manifest, "local": optimizer._hybrid_local_shard_bindings, @@ -136,6 +139,12 @@ def _snapshot(optimizer): def _assert_snapshot(optimizer, snapshot): assert optimizer._state_param_owner is snapshot["owner"] + expected_owner = snapshot["owner_contents"] + assert optimizer._state_param_owner.keys() == expected_owner.keys() + for key, (expected_param, expected_child) in expected_owner.items(): + live_param, live_child = optimizer._state_param_owner[key] + assert live_param is expected_param + assert live_child is expected_child assert optimizer._hybrid_post_sharding_finalized is snapshot["finalized"] assert optimizer._hybrid_sharding_manifest is snapshot["manifest"] assert optimizer._hybrid_local_shard_bindings is snapshot["local"] diff --git a/tests/test_rebinding_cpu.py b/tests/test_rebinding_cpu.py index 0fd1d2e..72e7ec6 100644 --- a/tests/test_rebinding_cpu.py +++ b/tests/test_rebinding_cpu.py @@ -137,6 +137,9 @@ def _snapshot(optimizer): "param_names": optimizer._param_names, "param_names_value": dict(optimizer._param_names), "bindings": optimizer._gefen_shard_bindings, + # Copy the mapping so an in-place add/remove/replace of a binding that + # keeps the registry object identity is still caught. + "bindings_contents": dict(optimizer._gefen_shard_bindings), "local_bindings": optimizer._gefen_local_shard_bindings, "manifest": optimizer._gefen_sharding_manifest, "finalized": optimizer._gefen_post_sharding_finalized, @@ -151,6 +154,23 @@ def _snapshot(optimizer): "_gefen_global_step_by_device", ) ), + # Bitwise clone of each device-cache entry so an in-place copy_() into a + # cached tensor that preserves the cache identity is still caught. + "caches_contents": tuple( + ( + name, + { + device: value.detach().clone() if torch.is_tensor(value) else copy.deepcopy(value) + for device, value in getattr(optimizer, name).items() + }, + ) + for name in ( + "_gefen_codebook_by_device", + "_gefen_codebook_lut_by_device", + "_sr_seed_by_device", + "_gefen_global_step_by_device", + ) + ), "capt_stacks": optimizer._capt_stacks, "static_mark_sig": optimizer._static_mark_sig, "global_step": optimizer._gefen_global_step, @@ -164,6 +184,10 @@ def _assert_snapshot(optimizer, snapshot): assert optimizer._param_names is snapshot["param_names"] assert optimizer._param_names == snapshot["param_names_value"] assert optimizer._gefen_shard_bindings is snapshot["bindings"] + expected_bindings = snapshot["bindings_contents"] + assert optimizer._gefen_shard_bindings.keys() == expected_bindings.keys() + for key, expected_binding in expected_bindings.items(): + assert optimizer._gefen_shard_bindings[key] is expected_binding assert optimizer._gefen_local_shard_bindings is snapshot["local_bindings"] assert optimizer._gefen_sharding_manifest is snapshot["manifest"] assert optimizer._gefen_post_sharding_finalized is snapshot["finalized"] @@ -186,6 +210,16 @@ def _assert_snapshot(optimizer, snapshot): ) for name, cache_ref in snapshot["caches"]: assert getattr(optimizer, name) is cache_ref + for name, expected_cache in snapshot["caches_contents"]: + live_cache = getattr(optimizer, name) + assert live_cache.keys() == expected_cache.keys() + for device, expected_value in expected_cache.items(): + live_value = live_cache[device] + if torch.is_tensor(expected_value): + assert torch.is_tensor(live_value) + assert torch.equal(live_value, expected_value) + else: + assert live_value == expected_value @pytest.mark.parametrize( diff --git a/tests/test_scoped_collective_agreement_fixes.py b/tests/test_scoped_collective_agreement_fixes.py index 8c50a71..e69491a 100644 --- a/tests/test_scoped_collective_agreement_fixes.py +++ b/tests/test_scoped_collective_agreement_fixes.py @@ -24,6 +24,8 @@ ShardingManifest, ) +from _state_snapshot import assert_deep_state_snapshot, deep_state_snapshot + _WORLD = 2 @@ -384,10 +386,13 @@ def test_parameter_storage_retarget_invalidates_prepared_canonical_import(): with torch.no_grad(): target_parameter.data = target_parameter.data.clone() + before_rejection = deep_state_snapshot(target) with pytest.raises( RuntimeError, match="changed after canonical import preparation" ): target.commit_canonical_state_import(prepared) + # The rejected commit must be a no-op before any recovery import masks it. + assert_deep_state_snapshot(target, before_rejection) # The refusal must keep canonical I/O available with a fresh preparation. target.import_canonical_state(exported) @@ -402,10 +407,13 @@ def test_parameter_inplace_mutation_invalidates_prepared_canonical_import(): with torch.no_grad(): target_parameter.mul_(2.0) + before_rejection = deep_state_snapshot(target) with pytest.raises( RuntimeError, match="changed after canonical import preparation" ): target.commit_canonical_state_import(prepared) + # The rejected commit must be a no-op before any recovery import masks it. + assert_deep_state_snapshot(target, before_rejection) target.import_canonical_state(exported) assert "automatic_period" in target.state[target_parameter] From 9134162a37af83e91657c623712af03f209c323a Mon Sep 17 00:00:00 2001 From: thad0ctor Date: Mon, 13 Jul 2026 15:56:58 -0700 Subject: [PATCH 43/52] Assert steady-state steps never recompute the manifest digest The manifest-digest recompute assertion was trivially satisfied by the deliberate compute earlier in the test, so a per-step recompute regression would still pass. Warm past the first step (which validates the layout fully once, a separately guarded behavior) and then require the computation count to stay fixed across further steps and a cache-backed fingerprint read. --- tests/test_layout_guard_cost.py | 19 +++++++++++++++---- 1 file changed, 15 insertions(+), 4 deletions(-) diff --git a/tests/test_layout_guard_cost.py b/tests/test_layout_guard_cost.py index 0449716..00cc207 100644 --- a/tests/test_layout_guard_cost.py +++ b/tests/test_layout_guard_cost.py @@ -176,13 +176,24 @@ def test_manifest_digest_is_computed_once_at_finalization(monkeypatch): assert first == second == optimizer._compute_codebook_manifest_fingerprint( optimizer._gefen_sharding_manifest ) - + # The deliberate recompute above must not mask a per-step regression, so + # anchor the count separately from it. + after_deliberate_compute = calls["count"] + assert after_deliberate_compute > finalize_computes + + # The first step after finalization validates the layout fully exactly once + # (see test_first_step_after_finalization_validates_fully_once), which + # recomputes the digest a single time. Warm past it, then require that every + # subsequent steady-state step recomputes nothing -- a per-step recompute + # regression would keep growing the count here. + _step_with_grads(optimizer, parameters) + warm = calls["count"] + assert warm == after_deliberate_compute + 1 _step_with_grads(optimizer, parameters) _step_with_grads(optimizer, parameters) - assert calls["count"] > finalize_computes # the deliberate compare above - steady = calls["count"] + assert calls["count"] == warm optimizer._codebook_manifest_fingerprint() - assert calls["count"] == steady + assert calls["count"] == warm def test_closure_layout_mutation_is_detected_with_a_warm_verdict(): From be797af109ec8795d193822607504e1c87119697 Mon Sep 17 00:00:00 2001 From: thad0ctor Date: Mon, 13 Jul 2026 15:56:58 -0700 Subject: [PATCH 44/52] Use a race-free file:// rendezvous in the state-movement distributed test _free_port() released the probed socket before the workers bound it, so another process could steal the port and flake CI. Rendezvous over a temporary file:// init_method that stays valid until every rank initializes, mirroring the DCP distributed test. --- tests/test_state_movement_distributed.py | 26 +++++++++++------------- 1 file changed, 12 insertions(+), 14 deletions(-) diff --git a/tests/test_state_movement_distributed.py b/tests/test_state_movement_distributed.py index c109ce7..7801c61 100644 --- a/tests/test_state_movement_distributed.py +++ b/tests/test_state_movement_distributed.py @@ -5,7 +5,7 @@ from datetime import timedelta import os import queue -import socket +import tempfile import traceback import pytest @@ -28,12 +28,6 @@ ) -def _free_port() -> str: - with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: - sock.bind(("127.0.0.1", 0)) - return str(sock.getsockname()[1]) - - def _oversized_copy(tensor: torch.Tensor) -> torch.Tensor: backing = torch.empty( tensor.numel() + 13, @@ -508,17 +502,14 @@ def assign_grad(parameters, seed): ) -def _distributed_worker(rank, world, port, result_queue) -> None: +def _distributed_worker(rank, world, init_file, result_queue) -> None: import torch.distributed as dist from torch.distributed.tensor import 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", + init_method="file://{}".format(init_file), rank=rank, world_size=world, timeout=timedelta(seconds=90), @@ -546,11 +537,16 @@ def test_atomic_state_movement_across_distributed_cpu_representations(): world = 2 context = mp.get_context("spawn") result_queue = context.Queue() - port = _free_port() + # A file:// rendezvous stays valid until every rank has initialized, unlike a + # pre-probed free TCP port that another process can steal before the workers + # bind it. Mirrors the DCP distributed test. + descriptor, init_file = tempfile.mkstemp(prefix="gefen-state-movement-") + os.close(descriptor) + os.unlink(init_file) processes = [ context.Process( target=_distributed_worker, - args=(rank, world, port, result_queue), + args=(rank, world, init_file, result_queue), ) for rank in range(world) ] @@ -570,6 +566,8 @@ def test_atomic_state_movement_across_distributed_cpu_representations(): if process.is_alive(): process.terminate() process.join(timeout=5) + if os.path.exists(init_file): + os.unlink(init_file) assert all(process.exitcode == 0 for process in processes), [ process.exitcode for process in processes From 081e13d277ac880c1ea4c2f45a51ad9a07c3585d Mon Sep 17 00:00:00 2001 From: thad0ctor Date: Mon, 13 Jul 2026 18:08:12 -0700 Subject: [PATCH 45/52] Value-compare child registries in deep fail-before-mutation snapshots The deep snapshot helper only identity-checked top-level optimizer __dict__ attributes, so a post_sharding regression that populated a live child's per-parameter registries (_param_names, _gefen_shard_bindings) in place before a later child raised would leave a half-populated dict that kept its object identity and slipped past every atomicity test. Capture the key set (by parameter identity) and cloned values of those registries in deep_state_snapshot and value-compare them in assert_deep_state_snapshot, handling optimizer types that do not stage them. Add hybrid-rebinding tests covering injected in-place adds and value changes to both registries. --- tests/_state_snapshot.py | 50 ++++++++++++++++++++++++++++++++++ tests/test_hybrid_rebinding.py | 49 +++++++++++++++++++++++++++++++++ 2 files changed, 99 insertions(+) diff --git a/tests/_state_snapshot.py b/tests/_state_snapshot.py index 795e981..08c1ecf 100644 --- a/tests/_state_snapshot.py +++ b/tests/_state_snapshot.py @@ -9,6 +9,17 @@ preserves every top-level object identity and is only caught by comparing the live values against these clones. ``assert_deep_state_snapshot`` therefore checks identity AND bitwise value equality together. + +The staged post_sharding transaction also publishes two per-parameter registry +dicts on each child — ``_param_names`` (parameter -> compatibility FQN string) +and ``_gefen_shard_bindings`` (parameter -> ``ShardIdentity``) — by rebuilding +them on a staged shadow and swapping the shadow in wholesale. Those dicts keep +their live object identity across a transaction, so a regression that wrote +``self._param_names[target] = ...`` on the LIVE child before a later child +raised would leave a half-populated registry that top-level identity checks +miss. ``deep_state_snapshot`` therefore also records the key set (by parameter +identity) and cloned values of these registries, and +``assert_deep_state_snapshot`` value-compares them after the transaction. """ import copy @@ -16,6 +27,13 @@ import torch +# Per-parameter registry dicts each Gefen/GefenMuon child rebuilds on a staged +# shadow during a post_sharding transaction. Keys are live parameters (compared +# by identity); values are FQN strings / ``ShardIdentity`` objects (compared by +# value). Absent on optimizer types that do not stage these (handled below). +_CHILD_REGISTRY_ATTRS = ("_param_names", "_gefen_shard_bindings") + + def _cloned(value): if torch.is_tensor(value): return value.detach().clone() @@ -90,9 +108,31 @@ def visit(value): return tuple(pairs) +def _registry_snapshot(optimizer): + """Cloned contents of each per-parameter registry dict that staging swaps. + + Returns ``name -> ((param, cloned_value), ...)`` for every registry that is + present as a ``dict`` on ``optimizer``; registries absent on this optimizer + type are skipped so the helper stays usable across optimizer kinds. Keys are + kept by reference for identity comparison; values are cloned so a later + in-place value mutation cannot alias the snapshot. + """ + + registries = {} + for name in _CHILD_REGISTRY_ATTRS: + registry = getattr(optimizer, name, None) + if type(registry) is not dict: + continue + registries[name] = tuple( + (key, _cloned(value)) for key, value in registry.items() + ) + return registries + + def deep_state_snapshot(optimizer): return { "attributes": optimizer.__dict__.copy(), + "registries": _registry_snapshot(optimizer), "state": optimizer.state, "state_items": tuple( (parameter, state, _cloned(dict(state))) @@ -140,3 +180,13 @@ def assert_deep_state_snapshot(optimizer, snapshot): ) for tensor, expected in snapshot["tensors"]: assert _tensors_bitwise_equal(tensor, expected) + for name, expected_entries in snapshot["registries"].items(): + live = getattr(optimizer, name, None) + assert type(live) is dict + live_entries = tuple(live.items()) + assert len(live_entries) == len(expected_entries) + for (live_key, live_value), (expected_key, expected_value) in zip( + live_entries, expected_entries + ): + assert live_key is expected_key + _nested_equal(live_value, expected_value) diff --git a/tests/test_hybrid_rebinding.py b/tests/test_hybrid_rebinding.py index 4902bda..1d91932 100644 --- a/tests/test_hybrid_rebinding.py +++ b/tests/test_hybrid_rebinding.py @@ -157,6 +157,55 @@ def _assert_snapshot(optimizer, snapshot): assert_deep_state_snapshot(live, expected_snapshot) +@pytest.mark.parametrize("registry", ["_param_names", "_gefen_shard_bindings"]) +def test_deep_snapshot_catches_inplace_child_registry_addition(registry): + # The staged post_sharding transaction rebuilds each child's per-parameter + # registries (``_param_names`` / ``_gefen_shard_bindings``) on a shadow and + # swaps it in, so a regression that instead wrote onto the LIVE registry + # before a later child raised would leave a half-populated dict. The dict + # keeps its object identity, so only value comparison of its contents + # catches the partial publication. + optimizer, matrix, bias = _optimizer() + for child in optimizer._subopts: + live = getattr(child, registry) + snapshot = deep_state_snapshot(child) + # The live registry object identity is preserved by an in-place add, so + # the top-level identity checks stay green; only content comparison sees + # the injected key. + assert getattr(child, registry) is live + injected = torch.nn.Parameter(torch.zeros(1)) + if registry == "_param_names": + live[injected] = "injected" + else: + live[injected] = _ungrouped_replicated("Injected.Param", (1,)) + with pytest.raises(AssertionError): + assert_deep_state_snapshot(child, snapshot) + del live[injected] + # With the injection undone the child is byte-for-byte its snapshot. + assert_deep_state_snapshot(child, snapshot) + + +def test_deep_snapshot_catches_inplace_child_registry_value_change(): + # A value edit that keeps the key set intact (e.g. rebinding a child's shard + # in place onto the live registry) must also be caught by content + # comparison, not just added/removed keys. + optimizer, matrix, bias = _optimizer() + child = optimizer._subopts[0] + (param,) = child._param_names + child._gefen_shard_bindings[param] = _ungrouped_replicated("Seed.Param", (2, 2)) + snapshot = deep_state_snapshot(child) + + child._param_names[param] = "mutated-name" + with pytest.raises(AssertionError): + assert_deep_state_snapshot(child, snapshot) + child._param_names[param] = snapshot["registries"]["_param_names"][0][1] + assert_deep_state_snapshot(child, snapshot) + + child._gefen_shard_bindings[param] = _ungrouped_replicated("Other.Param", (4,)) + with pytest.raises(AssertionError): + assert_deep_state_snapshot(child, snapshot) + + def test_composite_post_sharding_publishes_children_and_rebuilds_routing(): optimizer, old_matrix, old_bias = _optimizer() matrix = torch.nn.Parameter(torch.full((2, 2), 7.0)) From 4837a2d1f289de43010bdce7d633ff77e3d4b250 Mon Sep 17 00:00:00 2001 From: thad0ctor Date: Mon, 13 Jul 2026 18:12:14 -0700 Subject: [PATCH 46/52] Reuse the cached layout verdict on the per-step offload readiness path Under active state offload the per-step readiness guard (_assert_state_offload_step_ready, run ~2x/step) reached _state_movement_rejection_reason, which always consulted the finalized layout with full=True. That forced an uncached full layout rebuild plus a manifest sha256 recompute on every step (two full passes and two digest recomputes per step), bypassing the per-step layout-forensics memoization the non-offload guards already use. The finalized binding layout is immutable across steps, so thread a require_full_layout flag through _state_offload_rejection_reason and _state_movement_rejection_reason. The per-step offload readiness path uses the memoized fast path (full=False, reusing _gefen_layout_forensics_verdict), while the move_state_ / _atomic_state_movement_supported boundaries, activation, load, and contract readiness keep the full rebuild. The per-tensor offload state scan (CPU tightness/dtype, pairwise disjointness, native-schema validation) is untouched and still runs on every step, so a token-preserving corruption of a later parameter's offloaded state is still caught before any earlier parameter mutates. --- src/gefen/gefen.py | 37 +++++++++-- tests/test_state_offload.py | 119 ++++++++++++++++++++++++++++++++++++ 2 files changed, 150 insertions(+), 6 deletions(-) diff --git a/src/gefen/gefen.py b/src/gefen/gefen.py index 36e924f..1240682 100644 --- a/src/gefen/gefen.py +++ b/src/gefen/gefen.py @@ -3019,8 +3019,20 @@ def _state_offload_capturing_on_parameter_device(parameters) -> bool: return False def _state_offload_rejection_reason( - self, *, require_cpu_state: bool, allow_poisoned: bool = False + self, + *, + require_cpu_state: bool, + allow_poisoned: bool = False, + require_full_layout: bool = False, ): + # ``require_full_layout`` controls only the immutable finalized-LAYOUT + # forensics reached through ``_state_movement_rejection_reason``. It + # defaults to False so the per-step offload-readiness caller reuses the + # memoized layout verdict; boundary callers (activation, load, contract + # readiness) pass True for a full rebuild. Every per-tensor offload + # check below still runs on each call regardless of this flag, because + # offloaded per-parameter state tensors are legitimately replaced each + # step and a cached verdict there would miss corruption. if type(self) is not Gefen: return "native state offload is implemented only by plain Gefen" if self.state_offload_active and self.state_offload_device != torch.device( @@ -3054,7 +3066,9 @@ def _state_offload_rejection_reason( except RuntimeError: return "the CUDA graph-capture state could not be inspected" - movement_reason = self._state_movement_rejection_reason() + movement_reason = self._state_movement_rejection_reason( + require_full_layout=require_full_layout + ) if movement_reason is not None: return movement_reason @@ -3142,7 +3156,8 @@ def _state_offload_supported(self) -> bool: try: reason = self._state_offload_rejection_reason( - require_cpu_state=self.state_offload_active + require_cpu_state=self.state_offload_active, + require_full_layout=True, ) except Exception: return False @@ -3309,7 +3324,9 @@ def offload_state_(self, device="cpu") -> None: target = self._normalize_state_offload_target(device) self._assert_finalized_binding_layout(full=True) - reason = self._state_offload_rejection_reason(require_cpu_state=False) + reason = self._state_offload_rejection_reason( + require_cpu_state=False, require_full_layout=True + ) if reason is not None: raise RuntimeError("Gefen state offload is unavailable: {}".format(reason)) staged_state = self._stage_all_parameter_state_to_cpu() @@ -3356,10 +3373,17 @@ def restore_state_(self) -> None: ) self.move_state_() - def _state_movement_rejection_reason(self): + def _state_movement_rejection_reason(self, *, require_full_layout: bool = True): + # The finalized layout is immutable across steps, so the per-step + # offload-readiness caller passes ``require_full_layout=False`` to reuse + # the memoized layout-forensics verdict (an O(local params) identity + # token check) instead of forcing an uncached full rebuild + manifest + # digest recompute every step. Boundary callers (``move_state_``, + # ``_atomic_state_movement_supported``) keep the default full rebuild. + # The per-tensor state checks below always run regardless. if ( self._gefen_post_sharding_finalized - and not self._finalized_binding_layout_matches(full=True) + and not self._finalized_binding_layout_matches(full=require_full_layout) ): return "the finalized parameter binding no longer matches live groups" if self.capturable: @@ -8375,6 +8399,7 @@ def _stage_load_state_dict(self, state_dict): reason = staged._state_offload_rejection_reason( require_cpu_state=False, allow_poisoned=True, + require_full_layout=True, ) if reason is not None: raise RuntimeError( diff --git a/tests/test_state_offload.py b/tests/test_state_offload.py index a951def..1ecf49f 100644 --- a/tests/test_state_offload.py +++ b/tests/test_state_offload.py @@ -653,3 +653,122 @@ def test_capturable_compile_and_capture_states_fail_closed(monkeypatch): monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: True) with pytest.raises(RuntimeError, match="CUDA graph capture"): optimizer.offload_state_() + + +def _replicated_shard(fqn, shape): + identity = ParameterIdentity(fqn, shape) + return ShardIdentity( + identity, + ParameterLayout.REPLICATED, + LogicalSlice.full(identity), + ) + + +def _finalized_replicated_cuda_optimizer(count=3): + parameters = [ + torch.nn.Parameter(torch.full((4,), float(index + 1), device="cuda")) + for index in range(count) + ] + optimizer = Gefen( + [ + ("weight{}".format(index), parameter) + for index, parameter in enumerate(parameters) + ], + fused=False, + factored_v_2d=False, + ) + optimizer._resolve_automatic_period = lambda *args, **kwargs: 4 + shards = tuple( + _replicated_shard("Model.Weight{}".format(index), (4,)) + for index in range(count) + ) + optimizer.post_sharding( + tuple( + ParameterRebinding(parameter, parameter, shard) + for parameter, shard in zip(parameters, shards) + ), + manifest=ShardingManifest(shards), + ) + return optimizer, parameters + + +def _step_offload(optimizer, parameters): + for parameter in parameters: + parameter.grad = torch.full_like(parameter, 0.5) + optimizer.step() + + +def _count_full_layout_passes(monkeypatch): + calls = {"count": 0} + original = Gefen._finalized_binding_layout_matches_full + + def counted(self): + calls["count"] += 1 + return original(self) + + monkeypatch.setattr(Gefen, "_finalized_binding_layout_matches_full", counted) + return calls + + +def _count_manifest_digest_computes(monkeypatch): + calls = {"count": 0} + original = Gefen._compute_codebook_manifest_fingerprint + + def counted(self, manifest): + calls["count"] += 1 + return original(self, manifest) + + monkeypatch.setattr(Gefen, "_compute_codebook_manifest_fingerprint", counted) + return calls + + +@_CUDA_REQUIRED +def test_offload_steady_state_reuses_cached_layout_forensics(monkeypatch): + # The finalized layout is immutable across steps, so the per-step offload + # readiness path (called ~2x/step) must reuse the memoized layout-forensics + # verdict instead of rebuilding the full layout + recomputing the manifest + # digest on every step. The per-tensor offload scan still runs every step; + # only the immutable layout forensics is cached. + optimizer, parameters = _finalized_replicated_cuda_optimizer() + optimizer.offload_state_() + + layout = _count_full_layout_passes(monkeypatch) + digest = _count_manifest_digest_computes(monkeypatch) + + # The first offloaded step still fully validates the layout once, because + # offload_state_ invalidated the cached verdict; that single full pass + # recomputes the manifest digest once. + _step_offload(optimizer, parameters) + first_layout = layout["count"] + first_digest = digest["count"] + assert first_layout >= 1 + assert first_digest >= 1 + + # Steady state: neither the full layout rebuild nor the manifest digest is + # recomputed again. On the pre-fix code each of the two per-step readiness + # calls forced a full rebuild + digest recompute, so these counts grew by + # four per step. + for _ in range(3): + _step_offload(optimizer, parameters) + assert layout["count"] == first_layout + assert digest["count"] == first_digest + + +@_CUDA_REQUIRED +def test_offload_layout_change_is_still_caught_with_a_warm_verdict(): + optimizer, parameters = _finalized_replicated_cuda_optimizer() + optimizer.offload_state_() + _step_offload(optimizer, parameters) # warm the cached layout verdict + + # Replace a finalized registry container so the fast-path tokens diverge: + # the offload readiness path must fall back to the full forensic rebuild and + # reject the step before any parameter is staged or mutated. + first_before = parameters[0].detach().clone() + optimizer._gefen_local_shard_bindings = tuple( + reversed(optimizer._gefen_local_shard_bindings) + ) + for parameter in parameters: + parameter.grad = torch.full_like(parameter, 0.5) + with pytest.raises(RuntimeError, match="cannot step"): + optimizer.step() + assert torch.equal(parameters[0].detach(), first_before) From 6feefa6dd53bb9259fb35baafa716f387a436e0a Mon Sep 17 00:00:00 2001 From: thad0ctor Date: Mon, 13 Jul 2026 18:23:43 -0700 Subject: [PATCH 47/52] Flatten scalars before the bytewise snapshot comparison The fail-before-mutation snapshot compared raw bytes via view(torch.uint8), which PyTorch rejects on a 0-dim tensor ("self.dim() cannot be 0 to view Float as Byte"). Capturable device-resident scalar counters hit this, failing the CUDA-only capturable canonical-import tests (which the CPU CI skips). Flatten to 1-D before reinterpreting; shapes are already checked equal, so both sides flatten identically. --- tests/_state_snapshot.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/tests/_state_snapshot.py b/tests/_state_snapshot.py index 08c1ecf..7bb0115 100644 --- a/tests/_state_snapshot.py +++ b/tests/_state_snapshot.py @@ -55,8 +55,12 @@ def _tensors_bitwise_equal(live, expected): or tuple(live.shape) != tuple(expected.shape) ): return False - live_bytes = live.detach().contiguous().view(torch.uint8) - expected_bytes = expected.detach().contiguous().view(torch.uint8) + # ``view(torch.uint8)`` rejects a 0-dim tensor (a scalar cannot hold the + # several uint8 elements one value reinterprets to), which capturable + # device-resident scalar counters hit. Flatten to 1-D first; shapes were + # already checked equal above, so both sides flatten identically. + live_bytes = live.detach().contiguous().reshape(-1).view(torch.uint8) + expected_bytes = expected.detach().contiguous().reshape(-1).view(torch.uint8) return bool(torch.equal(live_bytes, expected_bytes)) From a114fcb0b6e04105462aafd4b52f42330eb15aff Mon Sep 17 00:00:00 2001 From: thad0ctor Date: Wed, 15 Jul 2026 09:49:09 -0700 Subject: [PATCH 48/52] Re-point codebook-scope failure sync onto the upstreamed primitive (Tier-2 convention side) (#72) * Preserve legacy vmean loads on the rank-local sharded path (#70) * Preserve legacy vmean loads on the rank-local sharded path _unwrap_rank_local_sharded_checkpoint validated the payload with the strict default (allow_legacy_vmean_counter=False), so a rank-local (DTensor/FSDP) checkpoint carrying vmean without the separate vmean_step counter -- a pre-counter state the step-time resume path backfills from step -- was rejected at load. The native load path already opts into that tolerance (allow_legacy_vmean_counter=True); the rank-local path did not, so the two disagreed and an otherwise valid legacy resume failed only on the sharded path. This regressed against pre-load-atomicity behavior, which had no such check. Pass allow_legacy_vmean_counter=True at the rank-local call site to match the native path. This only relaxes the vmean-without-vmean_step case; a current checkpoint (which carries vmean_step) is unaffected, and the inverse vmean_step-without-vmean corruption check is unchanged. tests/test_cpu_step_checkpoint.py::test_rank_local_validator_tolerates_legacy_vmean_without_step pins the validator tolerance both ways (the strict default rejects the pre-counter payload; the tolerance the rank-local path now opts into accepts it). * Guard the rank-local unwrap path directly, not just the validator Add test_rank_local_unwrap_tolerates_and_backfills_legacy_vmean: a single-rank gloo world builds a rank-local (rank_local_dtensor_v2) checkpoint, drops the separate vmean_step counter to simulate a pre-counter state, loads it through _unwrap_rank_local_sharded_checkpoint (the call site the fix touches), and asserts the first resumed step backfills vmean_step from step. The existing validator test calls _validate_rank_local_states directly, so it stays green if the unwrap call site regresses to strict validation. This test fails without allow_legacy_vmean_counter=True on the rank-local path (verified: reverting the fix raises "block second moment is missing vmean_step" at load). CPU-only, no multiprocessing (~2s). * Synchronize pre-collective Muon step failures * Disable Dynamo tracing on the pre-collective sync wrappers Match the mainline fix (#71): the sharded sync wrappers (_synchronize_sharded_step_flag/_error/_control_range, _prepare_synchronized_amp_step) and their codebook-scope analogs (_synchronize_codebook_scope_failure, _prepare_scoped_amp_optimizer_step) all branch/raise on dist.all_reduce results, so they carry @torch._dynamo.disable like the sibling _step_failure_process_groups and _assert_sharded_grad_presence_consistent. Tracing-only; no runtime change (precollective + scoped-agreement tests still 8/8). * Close pre-collective sync gaps on the convention side Four distributed-correctness fixes (found in review), on top of re-pointing the codebook scope onto the mainline primitive: * Overlapping-mesh deadlock + wrong flag device: same _step_failure_process_- groups rewrite as the mainline PR -- sync over every participated mesh (deduped, ordered) with the shard device, mirroring _assert_sharded_grad_presence_consistent. Threads collective_device through _synchronize_step_control_range too. * Scoped Gefen AMP presence hang: plain Gefen.step gated _prepare_scoped_amp_optimizer_step() on local found_inf/grad_scale, so with a multi-member codebook scope a rank with those attributes entered the presence all-gather while a rank without them skipped it -- deadlock. Run the scoped AMP agreement on EVERY member of a multi-member scope, matching the unconditional Muon/Hybrid preflight. * Dropped post-closure binding recheck (the CI failure): the merged Tier-2 step rewrite lost the convention's _assert_finalized_binding_layout() recheck after closure() on the mesh (non-scoped) branch, so a closure that rebinds param_groups was not caught -- test_finalized_layout_guard_rechecks_after_- closure[muon] failed. Re-assert it inside the preflight try (caught + synchronized) so a closure rebind fails before mutation and on every member. Full convention suite: 1095 passed, 0 failed. * Preserve mesh dimension order + close hybrid preflight gaps Three more review fixes: * DeviceMesh dimension order (same as the mainline PR): mirror _assert_sharded_grad_presence_consistent -- dedup meshes by content key, sorted-key order, get_all_groups() dimension order per mesh -- instead of flattening and sorting all groups by name, which could reorder a 2-D HSDP/TP mesh's row/column groups per rank and lock-order deadlock the preflight. * Hybrid post-closure binding recheck (unscoped path): the binding-is-None branch stepped children without re-asserting _assert_finalized_binding_layout() after the closure/pre-hooks, so a rank-local closure swapping a same-shaped backup parameter passed the gradient scan and stepped under stale routing. The scoped path already rechecked; re-assert in the unscoped preflight try (caught + synchronized). * Hybrid capture readiness: the _gefen_hybrid_precollective_preflight marker suppresses the muon child's own _assert_codebook_capture_ready() guard, and the hybrid preamble only checked devices, so a hybrid capturing before its Muon codebook initialized could run host-driven codebook init during capture. Call the muon capture/codebook readiness guard in the hybrid preflight. Full convention suite with GPU (2x3090 NCCL): 1420 passed. --- src/gefen/gefen.py | 247 ++++++++++---- src/gefen/gefen_muon.py | 299 ++++++++++++++--- src/gefen/hybrid.py | 103 ++++-- tests/test_codebook_scope_distributed.py | 6 + tests/test_cpu_step_checkpoint.py | 41 +++ tests/test_gefen_fsdp2_checkpoint.py | 82 +++++ tests/test_precollective_failure_sync.py | 402 +++++++++++++++++++++++ 7 files changed, 1035 insertions(+), 145 deletions(-) create mode 100644 tests/test_precollective_failure_sync.py diff --git a/src/gefen/gefen.py b/src/gefen/gefen.py index 18797e7..7bddd3f 100644 --- a/src/gefen/gefen.py +++ b/src/gefen/gefen.py @@ -648,6 +648,123 @@ def _resolve_find_period_backend(grad: torch.Tensor) -> str: return "cuda_kernel" if grad_work.device.type == "cuda" else "cpu" +def _step_failure_collective_device( + process_group, *, collective_device=None +) -> torch.device: + """Choose a backend-compatible device for a process-group control flag.""" + import torch.distributed as dist + + if collective_device is not None: + return torch.device(collective_device) + + group = process_group if process_group is not None else dist.group.WORLD + bound_device = getattr(group, "bound_device_id", None) + if bound_device is not None: + return torch.device(bound_device) + + backend = str(dist.get_backend(process_group)).lower() + if "nccl" in backend: + return torch.device("cuda", torch.cuda.current_device()) + if "xccl" in backend: + return torch.device("xpu", torch.xpu.current_device()) + # Gloo, MPI, UCC, and PyTorch's multi-backend default process group all + # accept CPU control tensors. Keeping their flag on CPU also avoids an + # unrelated accelerator initialization in spawned Gloo workers. + return torch.device("cpu") + + +@torch.no_grad() +@torch._dynamo.disable +def _synchronize_step_failure( + local_failed, process_group, *, collective_device=None +) -> bool: + """Return whether any process-group member reported a step failure.""" + failed = bool(local_failed) + if not torch.distributed.is_available() or not torch.distributed.is_initialized(): + return failed + + import torch.distributed as dist + + if dist.get_world_size(process_group) < 2: + return failed + flag = torch.tensor( + int(failed), + dtype=torch.int32, + device=_step_failure_collective_device( + process_group, collective_device=collective_device + ), + ) + dist.all_reduce(flag, op=dist.ReduceOp.MAX, group=process_group) + return bool(flag.item()) + + +@torch.no_grad() +@torch._dynamo.disable +def _synchronize_step_control_range( + local_minimum, local_maximum, process_group, *, collective_device=None +): + """Expand control-value bounds across one process group.""" + minimum = tuple(float(item) for item in local_minimum) + maximum = tuple(float(item) for item in local_maximum) + if len(minimum) != len(maximum): + raise ValueError("step-control bounds must have equal lengths") + if not torch.distributed.is_available() or not torch.distributed.is_initialized(): + return minimum, maximum + + import torch.distributed as dist + + if dist.get_world_size(process_group) < 2: + return minimum, maximum + bounds = torch.tensor( + minimum + tuple(-item for item in maximum), + dtype=torch.float64, + device=_step_failure_collective_device( + process_group, collective_device=collective_device + ), + ) + dist.all_reduce(bounds, op=dist.ReduceOp.MIN, group=process_group) + values = bounds.cpu().tolist() + width = len(minimum) + return tuple(values[:width]), tuple(-item for item in values[width:]) + + +def _amp_optimizer_step_controls(optimizer): + """Parse GradScaler's temporary controls without mutating gradients.""" + found_inf = getattr(optimizer, "found_inf", None) + if found_inf is not None: + if torch.is_tensor(found_inf): + if found_inf.numel() != 1: + raise RuntimeError( + "GradScaler supplied a non-scalar found_inf tensor with " + "shape {}".format(tuple(found_inf.shape)) + ) + overflow = bool(found_inf.detach().item()) + else: + overflow = bool(found_inf) + else: + overflow = False + + grad_scale = getattr(optimizer, "grad_scale", None) + if torch.is_tensor(grad_scale): + if grad_scale.numel() != 1: + raise RuntimeError( + "GradScaler supplied a non-scalar grad_scale tensor with shape " + "{}".format(tuple(grad_scale.shape)) + ) + scale_value = float(grad_scale.detach().item()) + elif grad_scale is not None: + scale_value = float(grad_scale) + else: + scale_value = 0.0 + if grad_scale is not None and ( + not math.isfinite(scale_value) or scale_value <= 0.0 + ): + raise RuntimeError( + "GradScaler supplied a non-finite or non-positive grad_scale" + ) + return overflow, grad_scale, scale_value + + @torch.no_grad() def _amp_prepare_optimizer_step(optimizer) -> bool: """Honor PyTorch's native ``GradScaler`` optimizer-step protocol. @@ -664,33 +781,14 @@ def _amp_prepare_optimizer_step(optimizer) -> bool: keeps every existing fused/unfused update path operating on ordinary unscaled gradients and also works for local DTensor shards. """ - found_inf = getattr(optimizer, "found_inf", None) - if found_inf is not None: - if torch.is_tensor(found_inf): - if found_inf.numel() != 1: - raise RuntimeError( - "GradScaler supplied a non-scalar found_inf tensor with shape {}".format( - tuple(found_inf.shape) - ) - ) - overflow = bool(found_inf.detach().item()) - else: - overflow = bool(found_inf) - if overflow: - return False - - grad_scale = getattr(optimizer, "grad_scale", None) + overflow, grad_scale, _ = _amp_optimizer_step_controls(optimizer) + if overflow: + return False if grad_scale is None: # Explicit scaler.unscale_(optimizer), or an ordinary non-AMP step. return True if torch.is_tensor(grad_scale): - if grad_scale.numel() != 1: - raise RuntimeError( - "GradScaler supplied a non-scalar grad_scale tensor with shape {}".format( - tuple(grad_scale.shape) - ) - ) scale = grad_scale.detach() # Match torch.amp.GradScaler.unscale_: computing the reciprocal in # fp64 avoids compile-option-dependent fp32 division differences. @@ -5128,6 +5226,7 @@ def _iter_gefen_grad_periods( yield param_name, flat, period, grad + @torch._dynamo.disable def _synchronize_codebook_scope_failure(self, error, phase: str) -> None: binding = self._gefen_codebook_process_group if binding is None or len(binding.identity.ordered_members) == 1: @@ -5135,15 +5234,12 @@ def _synchronize_codebook_scope_failure(self, error, phase: str) -> None: raise error return self._assert_runtime_codebook_process_group() - import torch.distributed as dist - - failed = torch.tensor( - int(error is not None), - dtype=torch.int32, - device=binding.collective_device, + failed = _synchronize_step_failure( + error is not None, + binding.process_group, + collective_device=binding.collective_device, ) - dist.all_reduce(failed, op=dist.ReduceOp.MAX, group=binding.process_group) - if int(failed.item()) == 0: + if not failed: return if error is not None: raise RuntimeError( @@ -5157,49 +5253,29 @@ def _synchronize_codebook_scope_failure(self, error, phase: str) -> None: ) ) + @torch._dynamo.disable def _prepare_scoped_amp_optimizer_step(self) -> bool: binding = self._gefen_codebook_process_group if binding is None: return _amp_prepare_optimizer_step(self) found_inf = getattr(self, "found_inf", None) - grad_scale = getattr(self, "grad_scale", None) + local_present = hasattr(self, "found_inf") or hasattr( + self, "grad_scale" + ) try: - if found_inf is None: - local_overflow = False - elif torch.is_tensor(found_inf): - if found_inf.numel() != 1: - raise RuntimeError( - "GradScaler supplied a non-scalar found_inf tensor with shape {}".format( - tuple(found_inf.shape) - ) - ) - local_overflow = bool(found_inf.detach().item()) - else: - if len(binding.identity.ordered_members) > 1: - raise RuntimeError( - "a multi-member scoped optimizer requires a group-aware " - "gradient scaler to provide tensor found_inf on every member" - ) - local_overflow = bool(found_inf) - if grad_scale is None: - scale_present = False - local_scale = 0.0 - elif torch.is_tensor(grad_scale): - if grad_scale.numel() != 1: - raise RuntimeError( - "GradScaler supplied a non-scalar grad_scale tensor with shape {}".format( - tuple(grad_scale.shape) - ) - ) - scale_present = True - local_scale = float(grad_scale.detach().item()) - else: - scale_present = True - local_scale = float(grad_scale) - if scale_present and (not math.isfinite(local_scale) or local_scale <= 0.0): + local_overflow, grad_scale, local_scale = ( + _amp_optimizer_step_controls(self) + ) + if ( + found_inf is not None + and not torch.is_tensor(found_inf) + and len(binding.identity.ordered_members) > 1 + ): raise RuntimeError( - "GradScaler supplied a non-finite or non-positive grad_scale" + "a multi-member scoped optimizer requires a group-aware " + "gradient scaler to provide tensor found_inf on every member" ) + scale_present = grad_scale is not None local_error = None except Exception as exc: local_overflow = False @@ -5211,7 +5287,12 @@ def _prepare_scoped_amp_optimizer_step(self) -> bool: import torch.distributed as dist amp_control = torch.tensor( - [int(local_overflow), int(scale_present), local_scale], + [ + int(local_present), + int(local_overflow), + int(scale_present), + local_scale, + ], dtype=torch.float64, device=binding.collective_device, ) @@ -5221,8 +5302,9 @@ def _prepare_scoped_amp_optimizer_step(self) -> bool: dist.all_gather(controls, amp_control, group=binding.process_group) if any(not torch.equal(item, controls[0]) for item in controls[1:]): raise RuntimeError( - "scoped Gefen AMP requires identical found_inf and grad_scale " - "on every process-group member; use a group-aware gradient scaler" + "scoped Gefen AMP requires identical protocol presence, " + "found_inf, and grad_scale on every process-group member; " + "use a group-aware gradient scaler" ) if local_overflow: return False @@ -9058,8 +9140,17 @@ def _unwrap_rank_local_sharded_checkpoint(self, state_dict) -> None: ) selected_states = selected_payload["states"] selected_codebook = selected_payload["codebook"] + # Match the native load path's legacy tolerance: an older rank-local + # checkpoint can carry vmean without vmean_step (a pre-counter state, as + # old as step), and the step-time resume path backfills vmean_step from + # step. Rejecting it here (the default is strict) would block that + # backfill and break otherwise valid legacy DTensor/FSDP resumes, while + # the native path already accepts them. self._validate_rank_local_states( - selected_states, current_signature, selected_codebook + selected_states, + current_signature, + selected_codebook, + allow_legacy_vmean_counter=True, ) state_dict["state"] = dict(zip(saved_ids, selected_states)) state_dict["gefen_global_step"] = marker_step @@ -9442,11 +9533,21 @@ def step(self, closure=None): # GradScaler invokes native-AMP optimizers even on overflow. Decide # before codebook learning, periodic refresh, capturable counters, or - # parameter/state mutation. The attribute gate compiles away on the + # parameter/state mutation. A multi-member codebook scope must run the + # AMP protocol agreement on EVERY member -- including one with no local + # GradScaler attributes -- otherwise the presence gather inside + # _prepare_scoped_amp_optimizer_step is entered by only some members and + # deadlocks (mirrors the unconditional Muon/Hybrid AMP preflight). Off a + # multi-member scope the attribute gate still compiles away on the # ordinary BF16/FP32 path where GradScaler attaches nothing. - if ( - hasattr(self, "found_inf") or hasattr(self, "grad_scale") - ) and not self._prepare_scoped_amp_optimizer_step(): + scope = self._gefen_codebook_process_group + if scope is not None and len(scope.identity.ordered_members) > 1: + should_step = self._prepare_scoped_amp_optimizer_step() + elif hasattr(self, "found_inf") or hasattr(self, "grad_scale"): + should_step = self._prepare_scoped_amp_optimizer_step() + else: + should_step = True + if not should_step: return loss if ( diff --git a/src/gefen/gefen_muon.py b/src/gefen/gefen_muon.py index e967794..7f68f24 100644 --- a/src/gefen/gefen_muon.py +++ b/src/gefen/gefen_muon.py @@ -15,7 +15,11 @@ ) from gefen.gefen import ( Gefen, + _amp_optimizer_step_controls, + _amp_prepare_optimizer_step, _assert_optimizer_gradients_structurally_valid, + _synchronize_step_control_range, + _synchronize_step_failure, ) EPS = 1e-7 @@ -1440,6 +1444,197 @@ def _dist_available() -> bool: return False return torch.distributed.is_initialized() + @torch._dynamo.disable + def _step_failure_process_groups(self): + """Return the control scope for eager exact/distributed Muon steps. + + Enter EXACTLY the collectives ``_assert_sharded_grad_presence_consistent`` + enters, in the same order: every sharded (non-approx) mesh this rank + participates in, keyed and deduplicated by mesh content and iterated in + sorted-key order, and within each mesh its ``get_all_groups()`` dimension + order. Preserving the per-mesh dimension order matters for multi-dim + (HSDP/TP) meshes -- flattening and sorting all groups by name could make + one rank enter a row all-reduce while a peer is already blocked in its + column all-reduce, deadlocking the preflight itself. Each group carries + the local shard device so the control flag lands on the same device as + the real collectives rather than the ambient ``current_device``. + + Like the grad-presence preflight, this is a single pass per group; a + failure only reaches ranks sharing a mesh with the failing rank. Deeply + overlapping non-enclosing meshes therefore keep that preflight's existing + cross-mesh limitation. + """ + if not self._dist_available(): + return () + + params = [param for group in self.param_groups for param in group["params"]] + # The protocol returns host-readable flags and therefore cannot be + # captured. Captured steps already require an eager warmup with a fixed + # control-flow and gradient-presence pattern. + if any(param.device.type == "cuda" for param in params) and ( + torch.cuda.is_current_stream_capturing() + ): + return () + + import torch.distributed as dist + + by_mesh = {} + for group in self.param_groups: + if group["sharded_mode"] == "approx": + continue + for param in group["params"]: + if not self._is_sharded(param): + continue + mesh = param.device_mesh + if mesh.get_coordinate() is None or mesh.size() < 2: + continue + process_groups = tuple(mesh.get_all_groups()) + members = tuple( + int(item) + for item in mesh.mesh.detach().cpu().reshape(-1).tolist() + ) + key = ( + str(mesh.device_type), + tuple(int(item) for item in mesh.shape), + members, + tuple(str(pg.group_name) for pg in process_groups), + ) + by_mesh.setdefault( + key, (self._state_tensor_device(param), process_groups) + ) + result = [] + for key in sorted(by_mesh): + device, process_groups = by_mesh[key] + for process_group in process_groups: + if dist.get_world_size(process_group) > 1: + result.append((process_group, device)) + return tuple(result) + + @staticmethod + @torch._dynamo.disable + def _synchronize_sharded_step_flag(local_value, process_groups) -> bool: + synchronized = bool(local_value) + for process_group, collective_device in process_groups: + synchronized = _synchronize_step_failure( + synchronized, process_group, collective_device=collective_device + ) + return synchronized + + @torch._dynamo.disable + def _synchronize_sharded_step_error( + self, error, phase: str, process_groups + ) -> None: + if not process_groups: + if error is not None: + raise error + return + failed = self._synchronize_sharded_step_flag( + error is not None, process_groups + ) + if not failed: + return + + import torch.distributed as dist + + if error is not None: + raise RuntimeError( + "GefenMuon {} failed on local rank {}: {}".format( + phase, dist.get_rank(), error + ) + ) from error + raise RuntimeError( + "GefenMuon {} failed on another process-group member".format( + phase + ) + ) + + @staticmethod + @torch._dynamo.disable + def _synchronize_sharded_step_control_range(local_control, process_groups): + minimum = tuple(float(item) for item in local_control) + maximum = minimum + for process_group, collective_device in process_groups: + minimum, maximum = _synchronize_step_control_range( + minimum, maximum, process_group, collective_device=collective_device + ) + return minimum, maximum + + @torch._dynamo.disable + def _prepare_synchronized_amp_step(self, optimizer, process_groups) -> bool: + """Agree on AMP controls before unscaling or entering Muon collectives.""" + local_present = hasattr(optimizer, "found_inf") or hasattr( + optimizer, "grad_scale" + ) + if not process_groups: + if not local_present: + return True + return _amp_prepare_optimizer_step(optimizer) + + try: + if local_present: + local_overflow, grad_scale, scale_value = ( + _amp_optimizer_step_controls(optimizer) + ) + local_scale_present = grad_scale is not None + else: + local_overflow = False + local_scale_present = False + scale_value = 0.0 + local_amp_error = None + except Exception as exc: + local_overflow = False + local_scale_present = False + scale_value = 0.0 + local_amp_error = exc + self._synchronize_sharded_step_error( + local_amp_error, "AMP control preflight", process_groups + ) + + minimum, maximum = self._synchronize_sharded_step_control_range( + ( + int(local_present), + int(local_overflow), + int(local_scale_present), + scale_value, + ), + process_groups, + ) + if minimum[0] != maximum[0]: + raise RuntimeError( + "GefenMuon AMP protocol presence differs across process-group " + "members; use the same group-aware gradient scaler on every member" + ) + if not bool(maximum[0]): + return True + if minimum[1] != maximum[1]: + raise RuntimeError( + "GefenMuon AMP found_inf differs across process-group members; " + "use a group-aware gradient scaler" + ) + if minimum[2] != maximum[2]: + raise RuntimeError( + "GefenMuon AMP grad_scale presence differs across process-group " + "members; use the same group-aware gradient scaler on every member" + ) + if bool(maximum[2]) and minimum[3] != maximum[3]: + raise RuntimeError( + "GefenMuon AMP grad_scale differs across process-group members; " + "use a group-aware gradient scaler" + ) + if bool(maximum[1]): + return False + + try: + should_step = _amp_prepare_optimizer_step(optimizer) + local_amp_error = None + except Exception as exc: + should_step = False + local_amp_error = exc + self._synchronize_sharded_step_error( + local_amp_error, "AMP preparation", process_groups + ) + return should_step + @torch._dynamo.disable def _assert_sharded_grad_presence_consistent(self) -> None: """Fail before collectives when mesh ranks disagree on the step inputs. @@ -2864,57 +3059,77 @@ def step(self, closure=None): "GefenMuon whole-parameter owner stepping requires the separate " "explicit process-group codebook scope" ) - # Capture-readiness depends on rank-local gradients and the closure is - # user code, so both can fail on only a subset of scope members. Capture - # such a failure and synchronize it on the codebook binding BEFORE any - # member enters the scoped operation-header or later collectives, so it - # raises on every member together instead of stranding peers inside a - # collective. The finalized-layout and runtime-process-group guards above - # stay local: they establish the very binding used to synchronize. loss = None - try: - self._assert_capturable_if_capturing() - self._assert_codebook_capture_ready() - if closure is not None: - with torch.enable_grad(): - loss = closure() - local_preamble_error = None - except Exception as exc: - loss = None - local_preamble_error = exc - if self._gefen_codebook_process_group is not None: + hybrid_preflight_complete = bool( + getattr(self, "_gefen_hybrid_precollective_preflight", False) + ) + if not hybrid_preflight_complete and ( + self._gefen_codebook_process_group is not None + ): + # The explicit convention binding is the exclusive control scope: + # synchronize on it before any operation header or mesh collective, + # never in addition to the mainline DTensor-derived scope. + try: + self._assert_capturable_if_capturing() + self._assert_codebook_capture_ready() + if closure is not None: + with torch.enable_grad(): + loss = closure() + local_preamble_error = None + except Exception as exc: + loss = None + local_preamble_error = exc self._synchronize_codebook_scope_failure( local_preamble_error, "step preamble" ) - elif local_preamble_error is not None: - raise local_preamble_error - self._assert_finalized_binding_layout() - self._assert_runtime_codebook_process_group() - try: - _assert_optimizer_gradients_structurally_valid( - self, require_2d_params=True - ) - local_preflight_error = None - except Exception as exc: - local_preflight_error = exc - self._validate_codebook_scope_operation_header("step") - if self._gefen_codebook_process_group is not None: + self._assert_finalized_binding_layout() + self._assert_runtime_codebook_process_group() + try: + _assert_optimizer_gradients_structurally_valid( + self, require_2d_params=True + ) + local_preflight_error = None + except Exception as exc: + local_preflight_error = exc + self._validate_codebook_scope_operation_header("step") self._synchronize_codebook_scope_failure( local_preflight_error, "gradient preflight" ) - elif local_preflight_error is not None: - raise local_preflight_error - self._ensure_codebook_scope_agreement() + self._ensure_codebook_scope_agreement() - # Native GradScaler calls AMP-aware optimizers even when its non-finite - # scan found an overflow. Skip before the sharded preflight, first-step - # codebook learning, or any optimizer/parameter mutation; finite scaled - # gradients are unscaled once here for every existing Muon path. - if ( - hasattr(self, "found_inf") or hasattr(self, "grad_scale") - ) and not self._prepare_scoped_amp_optimizer_step(): - return loss + if not self._prepare_scoped_amp_optimizer_step(): + return loss + elif not hybrid_preflight_complete: + process_groups = self._step_failure_process_groups() + try: + self._assert_capturable_if_capturing() + self._assert_codebook_capture_ready() + if closure is not None: + with torch.enable_grad(): + loss = closure() + # A closure can rebind param_groups (e.g. swap in a rogue + # parameter), so re-assert the finalized binding layout after it + # runs -- before any state mutation -- matching the scoped path + # and the pre-Tier-2 step. The raise is caught below and + # synchronized so every mesh member fails together. + self._assert_finalized_binding_layout() + _assert_optimizer_gradients_structurally_valid( + self, require_2d_params=True + ) + local_preflight_error = None + except Exception as exc: + loss = None + local_preflight_error = exc + self._synchronize_sharded_step_error( + local_preflight_error, "step preflight", process_groups + ) + + # Every mesh member enters the AMP control agreement, including a + # member with no local GradScaler attributes. This prevents the + # protocol-presence decision itself from becoming rank-divergent. + if not self._prepare_synchronized_amp_step(self, process_groups): + return loss # Partition the work once so distributed-mode sharded params can take the # stable-owner Parallel-Muon path while every other param keeps the normal diff --git a/src/gefen/hybrid.py b/src/gefen/hybrid.py index bb823ab..1530cc5 100644 --- a/src/gefen/hybrid.py +++ b/src/gefen/hybrid.py @@ -1251,6 +1251,12 @@ def _assert_capturable_devices_if_capturing(self) -> None: def step(self, closure=None): self._assert_finalized_binding_layout() + binding = self._hybrid_codebook_process_group + process_groups = ( + self.muon._step_failure_process_groups() + if self.muon is not None and binding is None + else () + ) # Dispatch the INSTANCE step hooks around the composite step, mirroring # torch.optim.Optimizer.profile_hook_step exactly: hooks receive # (optimizer, args, kwargs) where args are the raw step() call args @@ -1292,47 +1298,84 @@ def step(self, closure=None): except Exception as exc: loss = None local_preamble_error = exc - if self._hybrid_codebook_process_group is not None: + + if binding is not None: self._synchronize_codebook_scope_failure( local_preamble_error, "step preamble" ) - elif local_preamble_error is not None: - raise local_preamble_error - self._assert_finalized_binding_layout() - # Composite structural preflight, atomically over BOTH children before - # either child steps. Under a shared codebook scope the failure is - # synchronized on the binding first (mirroring Gefen.step and - # GefenMuon.step), so a rank-local structural error raises on every - # scope member together instead of stranding peers inside a child's - # scoped step collectives. - try: - for child in self._subopts: - _assert_optimizer_gradients_structurally_valid(child, require_2d_params=child is self.muon) - local_preflight_error = None - except Exception as exc: - local_preflight_error = exc - if self._hybrid_codebook_process_group is not None: + self._assert_finalized_binding_layout() + try: + for child in self._subopts: + _assert_optimizer_gradients_structurally_valid( + child, require_2d_params=child is self.muon + ) + local_preflight_error = None + except Exception as exc: + local_preflight_error = exc self._synchronize_codebook_scope_failure( local_preflight_error, "gradient preflight" ) - elif local_preflight_error is not None: - raise local_preflight_error - # A non-finite gradient in either half skips BOTH children before their - # codebooks, states, counters, or parameters can move. GradScaler - # attaches found_inf/grad_scale to the composite, so under a shared - # multi-member codebook scope the overflow skip is the children's - # scoped AMP protocol run here -- collective found_inf/grad_scale - # agreement, then a group-wide skip -- before any child enters its - # scoped step collectives. Explicit scaler.unscale_(hybrid) is - # detected by grad_scale=None and is not repeated; automatic unscale - # covers every child parameter exactly once. - if (hasattr(self, "found_inf") or hasattr(self, "grad_scale")) and not self._prepare_scoped_amp_optimizer_step(): + should_step = self._prepare_scoped_amp_optimizer_step() + else: + if local_preamble_error is None: + try: + # A closure/pre-hook can rebind param_groups after the entry + # guard (e.g. swap a same-shaped backup parameter), so + # re-assert the finalized layout before stepping children; + # and the muon child's own capture/codebook readiness guard + # is suppressed below by the preflight marker, so run it here. + # Both raises are caught and synchronized so every mesh member + # fails together. + self._assert_finalized_binding_layout() + if self.muon is not None: + self.muon._assert_codebook_capture_ready() + for child in self._subopts: + _assert_optimizer_gradients_structurally_valid( + child, require_2d_params=child is self.muon + ) + local_preflight_error = None + except Exception as exc: + local_preflight_error = exc + else: + local_preflight_error = local_preamble_error + if self.muon is not None: + self.muon._synchronize_sharded_step_error( + local_preflight_error, + "hybrid step preflight", + process_groups, + ) + should_step = self.muon._prepare_synchronized_amp_step( + self, process_groups + ) + else: + if local_preflight_error is not None: + raise local_preflight_error + should_step = self._prepare_scoped_amp_optimizer_step() + if not should_step: for post_hook in self._optimizer_step_post_hooks.values(): post_hook(self, args, kwargs) return loss + with torch.no_grad(): for o in self._subopts: - o.step() + if o is not self.muon: + o.step() + continue + if binding is not None: + o.step() + continue + marker = object() + previous = getattr( + o, "_gefen_hybrid_precollective_preflight", marker + ) + o._gefen_hybrid_precollective_preflight = True + try: + o.step() + finally: + if previous is marker: + del o._gefen_hybrid_precollective_preflight + else: + o._gefen_hybrid_precollective_preflight = previous for post_hook in self._optimizer_step_post_hooks.values(): post_hook(self, args, kwargs) diff --git a/tests/test_codebook_scope_distributed.py b/tests/test_codebook_scope_distributed.py index a36699d..36a8160 100644 --- a/tests/test_codebook_scope_distributed.py +++ b/tests/test_codebook_scope_distributed.py @@ -1175,6 +1175,10 @@ def _nccl_empty_owner_worker(rank, world, init_file, queue): gathered = [torch.empty_like(local) for _ in range(world)] dist.all_gather(gathered, local) agreement = all(torch.equal(item, gathered[0]) for item in gathered[1:]) + # The codebook binding, not ambient CUDA state, owns the device used by + # every scoped control collective. Deliberately make current_device + # disagree before the step preamble exercises its shared failure flag. + torch.cuda.set_device(1 - rank) optimizer.step() queue.put( { @@ -1188,6 +1192,7 @@ def _nccl_empty_owner_worker(rank, world, init_file, queue): and optimizer._gefen_codebook.device.type == "cpu" ), "step": optimizer._gefen_global_step, + "ambient_device_differed": torch.cuda.current_device() != rank, } ) except Exception as exc: @@ -1232,6 +1237,7 @@ def test_nccl_scope_uses_explicit_collective_device_with_empty_nonowner(): assert all(item["agreement"] for item in results) assert all(item["empty_nonowner"] for item in results) assert all(item["step"] == 1 for item in results) + assert all(item["ambient_device_differed"] for item in results) finally: for process in processes: if process.is_alive(): diff --git a/tests/test_cpu_step_checkpoint.py b/tests/test_cpu_step_checkpoint.py index 03fdbc7..989fb0c 100644 --- a/tests/test_cpu_step_checkpoint.py +++ b/tests/test_cpu_step_checkpoint.py @@ -320,6 +320,47 @@ def test_load_rejects_quantized_momentum_without_codebook(): fresh_opt.load_state_dict(sd) +def test_rank_local_validator_tolerates_legacy_vmean_without_step(): + # An older rank-local (DTensor/FSDP) checkpoint can carry ``vmean`` without + # the separate ``vmean_step`` counter (a pre-counter state, "vmean as old as + # step"); the step-time resume path backfills ``vmean_step`` from ``step``. + # The native load path accepts these (``allow_legacy_vmean_counter=True``), + # and ``_unwrap_rank_local_sharded_checkpoint`` must too -- otherwise the + # validator rejects the payload before the backfill can run, breaking an + # otherwise valid legacy resume. Pin the validator tolerance both ways. + model = _small_model() + opt = Gefen(list(model.named_parameters()), lr=1e-3, fused=False) + for step_grads in _synthetic_grads(model, 3): + _apply_grads(model, step_grads) + opt.step() + + signature = opt._rank_local_sharded_signature(context={}) + states = [opt.state[p] for group in opt.param_groups for p in group["params"]] + legacy = [dict(state) for state in states] + dropped = 0 + for state in legacy: + if "vmean_step" in state: + del state["vmean_step"] + dropped += 1 + assert dropped, "expected the stepped optimizer to carry vmean_step" + + # Strict (the historical rank-local default) rejects the pre-counter payload. + with pytest.raises(ValueError, match="block second moment is missing"): + opt._validate_rank_local_states( + legacy, + signature, + opt._gefen_codebook, + allow_legacy_vmean_counter=False, + ) + # The tolerance the rank-local load path now opts into accepts it (no raise). + opt._validate_rank_local_states( + legacy, + signature, + opt._gefen_codebook, + allow_legacy_vmean_counter=True, + ) + + def test_load_rejects_partial_or_conflicting_group_metadata(): _, opt = _run_and_save(factored=False) sd = copy.deepcopy(opt.state_dict()) diff --git a/tests/test_gefen_fsdp2_checkpoint.py b/tests/test_gefen_fsdp2_checkpoint.py index 2fcbbb0..ef65c3a 100644 --- a/tests/test_gefen_fsdp2_checkpoint.py +++ b/tests/test_gefen_fsdp2_checkpoint.py @@ -736,3 +736,85 @@ 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 + + +@pytest.mark.skipif( + not torch.distributed.is_available() + or not torch.distributed.is_gloo_available(), + reason="rank-local unwrap needs a (gloo) process group and DTensor", +) +def test_rank_local_unwrap_tolerates_and_backfills_legacy_vmean(): + # Exercise the actual load path -- _unwrap_rank_local_sharded_checkpoint -- + # not just the validator: a pre-counter rank-local checkpoint (vmean without + # the separate vmean_step counter) must load, and the first resumed step must + # backfill vmean_step from step. A single-rank gloo world drives the same + # rank_local_dtensor_v2 format the multi-rank path uses, so it stays a + # CPU-only regression guard for the load-path tolerance (rejected pre-fix). + import torch.distributed as dist + import torch.nn as nn + from torch.distributed.device_mesh import init_device_mesh + from torch.distributed.tensor import Shard, distribute_tensor + + os.environ.setdefault("MASTER_ADDR", "127.0.0.1") + os.environ["MASTER_PORT"] = _free_port() + os.environ["RANK"] = "0" + os.environ["WORLD_SIZE"] = "1" + dist.init_process_group("gloo", rank=0, world_size=1) + try: + from gefen import Gefen + + mesh = init_device_mesh("cpu", (1,), mesh_dim_names=("dp",)) + + def build(): + model = nn.Module() + model.register_parameter( + "w", + nn.Parameter( + distribute_tensor( + torch.linspace(-1, 1, 64).reshape(8, 8), mesh, [Shard(0)] + ) + ), + ) + optimizer = Gefen( + [{"params": [("w", model.w)]}], + lr=1e-3, + fused=False, + factored_v_2d=False, + ) + return model, optimizer + + def apply_grad(model): + model.w.grad = distribute_tensor( + torch.arange(64).reshape(8, 8).float().cos(), mesh, [Shard(0)] + ) + + model, optimizer = build() + for _ in range(2): + apply_grad(model) + optimizer.step() + + # Save a pre-counter rank-local checkpoint: drop the separate vmean_step + # counter from the live state before serializing. + param = optimizer.param_groups[0]["params"][0] + assert "vmean_step" in optimizer.state[param] + optimizer.state[param].pop("vmean_step") + checkpoint = copy.deepcopy(optimizer.state_dict()) + assert any( + "rank_local_sharded_state" in group.get("_gefen_checkpoint_metadata", {}) + for group in checkpoint["param_groups"] + ), "expected a rank-local sharded checkpoint" + + # Load through the rank-local unwrap path (rejected before the fix). The + # pre-counter payload must be accepted with vmean_step still absent... + target_model, target_optimizer = build() + target_optimizer.load_state_dict(checkpoint) + target_param = target_optimizer.param_groups[0]["params"][0] + assert "vmean" in target_optimizer.state[target_param] + assert "vmean_step" not in target_optimizer.state[target_param] + + # ...and the first resumed step must backfill vmean_step from step. + apply_grad(target_model) + target_optimizer.step() + assert "vmean_step" in target_optimizer.state[target_param] + finally: + dist.destroy_process_group() diff --git a/tests/test_precollective_failure_sync.py b/tests/test_precollective_failure_sync.py new file mode 100644 index 0000000..91ff3b7 --- /dev/null +++ b/tests/test_precollective_failure_sync.py @@ -0,0 +1,402 @@ +"""CPU/Gloo coverage for failures before sharded Muon step collectives.""" + +from __future__ import annotations + +import copy +from datetime import timedelta +import os +from pathlib import Path +import queue as queue_module +import tempfile +import time +import traceback + +import pytest +import torch +import torch.distributed as dist +from torch import nn + +from gefen import GefenMuon, GefenMuonHybrid + + +_WORLD_SIZE = 2 +_PROCESS_GROUP_TIMEOUT_SECONDS = 20 +_STEP_DEADLINE_SECONDS = 10.0 +_WORKER_DEADLINE_SECONDS = 60.0 +_CASES = ("muon_exact", "muon_distributed", "hybrid_exact") + + +def _local_clone(value: torch.Tensor) -> torch.Tensor: + value = value.to_local() if hasattr(value, "to_local") else value + value = value.wait() if hasattr(value, "wait") else value + return value.detach().clone() + + +def _clone_tree(value): + if torch.is_tensor(value): + return _local_clone(value) + if isinstance(value, dict): + return {key: _clone_tree(item) for key, item in value.items()} + if isinstance(value, list): + return [_clone_tree(item) for item in value] + if isinstance(value, tuple): + return tuple(_clone_tree(item) for item in value) + return copy.deepcopy(value) + + +def _trees_equal(left, right) -> bool: + if torch.is_tensor(left) or torch.is_tensor(right): + return ( + torch.is_tensor(left) + and torch.is_tensor(right) + and left.dtype == right.dtype + and tuple(left.shape) == tuple(right.shape) + and torch.equal(left, right) + ) + if isinstance(left, dict) or isinstance(right, dict): + return ( + isinstance(left, dict) + and isinstance(right, dict) + and left.keys() == right.keys() + and all(_trees_equal(left[key], right[key]) for key in left) + ) + if isinstance(left, (list, tuple)) or isinstance(right, (list, tuple)): + return ( + type(left) is type(right) + and len(left) == len(right) + and all(_trees_equal(a, b) for a, b in zip(left, right)) + ) + return type(left) is type(right) and left == right + + +def _attribute_snapshot(owner, name: str): + if not hasattr(owner, name): + return (False, None, None) + value = getattr(owner, name) + return (True, id(value), _clone_tree(value)) + + +def _optimizer_snapshot(optimizer, parameters): + children = tuple(getattr(optimizer, "_subopts", (optimizer,))) + child_states = [] + codebooks = [] + global_steps = [] + for child in children: + child_states.append( + ( + id(child.state), + tuple( + (id(parameter), id(state), _clone_tree(state)) + for parameter, state in child.state.items() + ), + ) + ) + codebooks.append( + tuple( + _attribute_snapshot(child, name) + for name in ( + "_gefen_codebook", + "_gefen_codebook_by_device", + "_gefen_codebook_lut_by_device", + ) + ) + ) + global_steps.append( + tuple( + _attribute_snapshot(child, name) + for name in ( + "_gefen_global_step", + "_gefen_global_step_by_device", + ) + ) + ) + return { + "parameters": tuple(_local_clone(parameter) for parameter in parameters), + "gradients": tuple( + None if parameter.grad is None else _local_clone(parameter.grad) + for parameter in parameters + ), + "state": tuple(child_states), + "codebook": tuple(codebooks), + "global_step": tuple(global_steps), + } + + +def _snapshot_comparison(before, after): + return { + name: _trees_equal(before[name], after[name]) + for name in ("parameters", "gradients", "state", "codebook", "global_step") + } + + +def _make_optimizer(mesh, case: str): + from torch.distributed.tensor import Shard, distribute_tensor + + full_weight = torch.linspace(-0.8, 0.9, 64, dtype=torch.float32).reshape(8, 8) + full_weight_grad = torch.linspace(0.6, -0.7, 64, dtype=torch.float32).reshape(8, 8) + weight = nn.Parameter(distribute_tensor(full_weight.clone(), mesh, [Shard(0)])) + + if case == "hybrid_exact": + bias = nn.Parameter(torch.linspace(-0.4, 0.3, 8, dtype=torch.float32)) + optimizer = GefenMuonHybrid( + [("weight", weight)], + [("bias", bias)], + lr=1e-3, + fused=False, + ns_steps=1, + ns_schedule="standard", + sharded_mode="exact", + backup_optimizer="gefen", + normuon=False, + ) + parameters = (weight, bias) + bias.grad = torch.linspace(0.2, -0.3, 8, dtype=torch.float32) + else: + mode = case.removeprefix("muon_") + optimizer = GefenMuon( + [("weight", weight)], + lr=1e-3, + fused=False, + ns_steps=1, + ns_schedule="standard", + sharded_mode=mode, + ) + parameters = (weight,) + + weight.grad = distribute_tensor(full_weight_grad.clone(), mesh, [Shard(0)]) + return optimizer, parameters + + +def _closure_failure_result(rank: int, mesh, case: str): + optimizer, parameters = _make_optimizer(mesh, case) + before = _optimizer_snapshot(optimizer, parameters) + + def closure(): + if rank == 0: + raise RuntimeError("rank-zero closure failure") + return torch.tensor(1.0) + + started = time.monotonic() + try: + optimizer.step(closure) + message = None + error_traceback = None + except BaseException as exc: + message = "{}: {}".format(type(exc).__name__, exc) + error_traceback = traceback.format_exc() + elapsed = time.monotonic() - started + after = _optimizer_snapshot(optimizer, parameters) + return { + "message": message, + "traceback": error_traceback, + "elapsed": elapsed, + "unchanged": _snapshot_comparison(before, after), + } + + +def _amp_control_result(rank: int, mesh, case: str, scenario: str): + optimizer, parameters = _make_optimizer(mesh, case) + if scenario == "divergent_found_inf": + optimizer.found_inf = torch.tensor(float(rank == 0)) + optimizer.grad_scale = torch.tensor(8.0) + elif scenario == "protocol_presence": + if rank == 0: + optimizer.found_inf = torch.tensor(0.0) + optimizer.grad_scale = torch.tensor(8.0) + elif scenario == "scale_disagreement": + optimizer.found_inf = torch.tensor(0.0) + optimizer.grad_scale = torch.tensor(8.0 if rank == 0 else 4.0) + elif scenario == "group_wide_overflow": + optimizer.found_inf = torch.tensor(1.0) + optimizer.grad_scale = torch.tensor(8.0) + else: + raise AssertionError("unknown AMP-control scenario: {}".format(scenario)) + + post_hook_calls = [] + hook_handle = None + if case == "hybrid_exact" and scenario == "group_wide_overflow": + hook_handle = optimizer.register_step_post_hook( + lambda _optimizer, _args, _kwargs: post_hook_calls.append(True) + ) + before = _optimizer_snapshot(optimizer, parameters) + + started = time.monotonic() + try: + optimizer.step() + message = None + error_traceback = None + except BaseException as exc: + message = "{}: {}".format(type(exc).__name__, exc) + error_traceback = traceback.format_exc() + elapsed = time.monotonic() - started + after = _optimizer_snapshot(optimizer, parameters) + if hook_handle is not None: + hook_handle.remove() + return { + "message": message, + "traceback": error_traceback, + "elapsed": elapsed, + "controls_present": hasattr(optimizer, "found_inf") or hasattr(optimizer, "grad_scale"), + "local_found_inf": ( + float(optimizer.found_inf.item()) if hasattr(optimizer, "found_inf") else None + ), + "local_grad_scale": ( + float(optimizer.grad_scale.item()) if hasattr(optimizer, "grad_scale") else None + ), + "post_hook_calls": ( + len(post_hook_calls) + if case == "hybrid_exact" and scenario == "group_wide_overflow" + else None + ), + "unchanged": _snapshot_comparison(before, after), + } + + +def _distributed_worker(rank: int, init_method: str, case: str, result_queue) -> None: + try: + torch.set_num_threads(1) + dist.init_process_group( + "gloo", + init_method=init_method, + rank=rank, + world_size=_WORLD_SIZE, + timeout=timedelta(seconds=_PROCESS_GROUP_TIMEOUT_SECONDS), + ) + from torch.distributed.tensor import init_device_mesh + + mesh = init_device_mesh("cpu", (_WORLD_SIZE,), mesh_dim_names=("dp",)) + closure_failure = _closure_failure_result(rank, mesh, case) + dist.barrier() + divergent_found_inf = _amp_control_result(rank, mesh, case, "divergent_found_inf") + dist.barrier() + protocol_presence = _amp_control_result(rank, mesh, case, "protocol_presence") + dist.barrier() + scale_disagreement = _amp_control_result(rank, mesh, case, "scale_disagreement") + dist.barrier() + group_wide_overflow = _amp_control_result(rank, mesh, case, "group_wide_overflow") + dist.barrier() + result_queue.put( + { + "rank": rank, + "closure_failure": closure_failure, + "divergent_found_inf": divergent_found_inf, + "protocol_presence": protocol_presence, + "scale_disagreement": scale_disagreement, + "group_wide_overflow": group_wide_overflow, + } + ) + except BaseException: + result_queue.put({"rank": rank, "fatal_error": traceback.format_exc()}) + finally: + if dist.is_available() and dist.is_initialized(): + dist.destroy_process_group() + + +def _run_distributed_case(case: str): + context = torch.multiprocessing.get_context("spawn") + result_queue = context.Queue() + descriptor, rendezvous_path = tempfile.mkstemp(prefix="gefen-precollective-sync-") + os.close(descriptor) + os.unlink(rendezvous_path) + init_method = Path(rendezvous_path).resolve().as_uri() + processes = [ + context.Process( + target=_distributed_worker, + args=(rank, init_method, case, result_queue), + ) + for rank in range(_WORLD_SIZE) + ] + results = [] + hung_ranks = [] + exit_codes = [] + try: + for process in processes: + process.start() + deadline = time.monotonic() + _WORKER_DEADLINE_SECONDS + for process in processes: + process.join(timeout=max(0.0, deadline - time.monotonic())) + hung_ranks = [rank for rank, process in enumerate(processes) if process.is_alive()] + exit_codes = [process.exitcode for process in processes] + while len(results) < _WORLD_SIZE: + try: + results.append(result_queue.get(timeout=0.5)) + except queue_module.Empty: + break + finally: + for process in processes: + if process.is_alive(): + process.terminate() + process.join(timeout=5) + result_queue.close() + result_queue.join_thread() + try: + os.unlink(rendezvous_path) + except FileNotFoundError: + pass + + assert not hung_ranks, (case, "workers hung", hung_ranks, exit_codes, results) + assert all(code == 0 for code in exit_codes), (case, "nonzero worker exit", exit_codes, results) + assert len(results) == _WORLD_SIZE, (case, "missing worker result", exit_codes, results) + return sorted(results, key=lambda result: result["rank"]) + + +@pytest.mark.skipif( + not dist.is_available() or not dist.is_gloo_available(), + reason="pre-collective failure synchronization coverage requires Gloo", +) +@pytest.mark.parametrize("case", _CASES) +def test_rank_local_closure_and_amp_controls_are_synchronized(case): + results = _run_distributed_case(case) + assert all("fatal_error" not in result for result in results), results + + closure_results = [result["closure_failure"] for result in results] + assert all(item["message"] is not None for item in closure_results), closure_results + phase = "hybrid step preflight" if case == "hybrid_exact" else "step preflight" + assert "GefenMuon {} failed on local rank 0".format(phase) in closure_results[0]["message"] + assert "rank-zero closure failure" in closure_results[0]["message"] + assert "GefenMuon {} failed on another process-group member".format(phase) in closure_results[1]["message"] + assert all(item["elapsed"] < _STEP_DEADLINE_SECONDS for item in closure_results), closure_results + assert all(all(item["unchanged"].values()) for item in closure_results), closure_results + + overflow_results = [result["divergent_found_inf"] for result in results] + assert [item["local_found_inf"] for item in overflow_results] == [1.0, 0.0] + assert all( + item["message"] is not None + and "AMP found_inf differs across process-group members" in item["message"] + and "group-aware gradient scaler" in item["message"] + for item in overflow_results + ), overflow_results + assert all(item["elapsed"] < _STEP_DEADLINE_SECONDS for item in overflow_results), overflow_results + assert all(all(item["unchanged"].values()) for item in overflow_results), overflow_results + + presence_results = [result["protocol_presence"] for result in results] + assert [item["controls_present"] for item in presence_results] == [True, False] + assert all( + item["message"] is not None + and "AMP protocol presence differs across process-group members" in item["message"] + and "group-aware gradient scaler" in item["message"] + for item in presence_results + ), presence_results + assert all(item["elapsed"] < _STEP_DEADLINE_SECONDS for item in presence_results), presence_results + assert all(all(item["unchanged"].values()) for item in presence_results), presence_results + + scale_results = [result["scale_disagreement"] for result in results] + assert [item["local_found_inf"] for item in scale_results] == [0.0, 0.0] + assert [item["local_grad_scale"] for item in scale_results] == [8.0, 4.0] + assert all( + item["message"] is not None + and "AMP grad_scale differs across process-group members" in item["message"] + and "group-aware gradient scaler" in item["message"] + for item in scale_results + ), scale_results + assert all(item["elapsed"] < _STEP_DEADLINE_SECONDS for item in scale_results), scale_results + assert all(all(item["unchanged"].values()) for item in scale_results), scale_results + + skip_results = [result["group_wide_overflow"] for result in results] + assert all(item["message"] is None for item in skip_results), skip_results + assert all(item["elapsed"] < _STEP_DEADLINE_SECONDS for item in skip_results), skip_results + assert all(all(item["unchanged"].values()) for item in skip_results), skip_results + if case == "hybrid_exact": + assert [item["post_hook_calls"] for item in skip_results] == [1, 1] + else: + assert [item["post_hook_calls"] for item in skip_results] == [None, None] From 3b9478ec786020c932ada8b4b66c41f3c3d2126c Mon Sep 17 00:00:00 2001 From: thad0ctor Date: Wed, 15 Jul 2026 11:32:34 -0700 Subject: [PATCH 49/52] Slim optimizer convention to contracts, portable state, and DCP --- README.md | 3 - benchmarks/microbench/bench_layout_guard.py | 6 +- docs/optimizer_contracts.md | 22 +- scripts/release_gpu_gate.sh | 2 - src/gefen/__init__.py | 4 - src/gefen/contracts.py | 37 +- src/gefen/gefen.py | 1109 +---------------- src/gefen/gefen_muon.py | 58 +- src/gefen/hybrid.py | 32 +- src/gefen/portable_runtime.py | 6 - tests/test_canonical_state_cpu.py | 4 - tests/test_codebook_scope_distributed.py | 83 +- tests/test_hybrid_scoped_failure_protocol.py | 80 ++ tests/test_layout_guard_cost.py | 81 +- tests/test_optimizer_contracts.py | 94 +- .../test_scoped_collective_agreement_fixes.py | 14 +- tests/test_state_movement.py | 889 ------------- tests/test_state_movement_distributed.py | 577 --------- tests/test_state_offload.py | 774 ------------ 19 files changed, 312 insertions(+), 3563 deletions(-) delete mode 100644 tests/test_state_movement.py delete mode 100644 tests/test_state_movement_distributed.py delete mode 100644 tests/test_state_offload.py diff --git a/README.md b/README.md index a57330c..dc1efd8 100644 --- a/README.md +++ b/README.md @@ -634,13 +634,10 @@ Measured (Qwen3-0.6B, 2 and 4 GPUs): `"distributed"` matched `"exact"` in the re 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. -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. - **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. ## Troubleshooting diff --git a/benchmarks/microbench/bench_layout_guard.py b/benchmarks/microbench/bench_layout_guard.py index bcffbf6..33bd8b7 100644 --- a/benchmarks/microbench/bench_layout_guard.py +++ b/benchmarks/microbench/bench_layout_guard.py @@ -153,12 +153,10 @@ def main() -> int: ) def warm_step_guards(): - # The complete step() guard sequence (both the pre-closure and the - # post-closure blocks), on a warm verdict. - optimizer._assert_state_offload_step_ready() + # The finalized-layout/process-group step guard sequence (both the + # pre-closure and post-closure blocks), on a warm verdict. optimizer._assert_finalized_binding_layout() optimizer._assert_runtime_codebook_process_group() - optimizer._assert_state_offload_step_ready() optimizer._assert_finalized_binding_layout() optimizer._assert_runtime_codebook_process_group() diff --git a/docs/optimizer_contracts.md b/docs/optimizer_contracts.md index 270ce89..62f9403 100644 --- a/docs/optimizer_contracts.md +++ b/docs/optimizer_contracts.md @@ -23,7 +23,7 @@ assert ParameterLayout.DTENSOR_1D_DEFAULT_WORLD in rank_local_dcp.same_topology - `StateVariant` identifies valid lazy, initialized, local-shard, global-parameter, owner, non-owner, and migrated state combinations using structured layout, mode, rank, extent, ownership, and inactive-field declarations. - `TrainingSupport` qualifies each validated parameter layout by process-group source, mesh dimensionality, sharded mode, and whether the update needs complete parameter storage or a transient complete logical matrix. - `CheckpointSupport` reports same-topology, topology-changing, and fail-before-mutation load support separately for native, PyTorch rank-local, canonical local, canonical global, and composite checkpoint transports. -- Precision, canonical parameter identity, stable shard identity, explicit process-group-scoped codebooks, shard rebinding, post-sharding, canonical state I/O, 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. +- Precision, canonical parameter identity, stable shard identity, explicit process-group-scoped codebooks, shard rebinding, post-sharding, canonical state I/O, and the retained `atomic_state_movement` and `state_offload` fields are independent capability declarations. A false field is an explicit unsupported contract, not an invitation for an adapter to infer support from internal state. The current DTensor declaration is deliberately narrow: `DTENSOR_1D_DEFAULT_WORLD` means one shared one-dimensional mesh spanning the default world. Multidimensional meshes, subgroups, and placement-changing loads are not implied by that declaration. @@ -31,7 +31,7 @@ The current DTensor declaration is deliberately narrow: `DTENSOR_1D_DEFAULT_WORL `ParameterIdentity` records an exact, case-preserving model FQN and global logical shape independently of any live tensor object. `ProcessGroupIdentity` records an adapter-defined semantic group name and authoritative ordered member IDs without importing a framework process-group type. `ShardIdentity` combines those values with a contiguous row-major `LogicalSlice` for replicated, flattened, and whole-owner layouts or an axis-aligned `LogicalRegion` for the narrow one-dimensional DTensor layout, plus structured placements, the local member, and an optional whole-parameter owner. `ShardingManifest` validates and deterministically orders the complete identity set; flattened slices and dimension-sharded regions must cover each logical parameter exactly once without gaps or overlaps, replicated manifests carry one complete identity per declared member, and whole-parameter manifests identify one complete owner while retaining empty non-owner records. DTensor regions currently describe one default-world mesh axis with either replication or one parameter-dimension shard, including uneven and empty shards. This identity vocabulary does not by itself claim DTensor post-sharding rebinding or portable checkpoint support; those remain negative until the optimizer data plane consumes the regions. -These descriptors do not treat legacy `param_names`, generated names, Python tensor identity, rank-local parameter IDs, devices, or dtypes as canonical identity. They also do not contain runtime collective handles. An adapter remains responsible for mapping a stable `ProcessGroupIdentity` to its framework process group and for canonicalizing tied aliases to one primary FQN and one optimizer slot; alias-rich identity is not part of schema version 1. Declaring identity metadata alone does not enable rebinding, canonical checkpoint I/O, topology-changing load, codebook scoping, state movement, or offload; those capabilities remain separate. +These descriptors do not treat legacy `param_names`, generated names, Python tensor identity, rank-local parameter IDs, devices, or dtypes as canonical identity. They also do not contain runtime collective handles. An adapter remains responsible for mapping a stable `ProcessGroupIdentity` to its framework process group and for canonicalizing tied aliases to one primary FQN and one optimizer slot; alias-rich identity is not part of schema version 1. Declaring identity metadata alone does not enable rebinding, canonical checkpoint I/O, topology-changing load, or codebook scoping; those capabilities remain separate. ## Atomic post-sharding rebinding @@ -41,7 +41,7 @@ Rebinding is allowed only while the entire optimizer is pristine: global step ze Plain Gefen currently rebinds complete replicated parameters and contiguous physical 1-D flattened element shards. A flattened logical matrix requires `factored_v_2d=False` because canonical row/column factored-state projection is not implemented. GefenMuon rebinds complete replicated matrices and can finalize/prune whole-parameter owner manifests. A whole-owner instance may step only when the same transaction installs the explicit codebook process group described below; the owner performs the update and the machine-readable training declaration requires the adapter to synchronize the updated matrix afterward. A Gefen-backed `GefenMuonHybrid` atomically partitions one complete manifest and rebinding plan by its frozen exact FQN routing, stages both children, validates cross-child storage disjointness, rebuilds composite state routing, and publishes only after every child succeeds. AdamW-backed Hybrid and DTensor composite rebinding remain unsupported. The portable global-state path described below can reshard supported finalized layouts. -After finalization, every entry point re-validates the published layout, and this has two costs. Steps and identity queries take an O(local params) fast path: the first complete forensic rebuild caches a verdict keyed by cheap identity tokens — the finalized registries by object identity, every live group container, parameter, and compatibility name, and a version counter that every mutating API bumps — and the full rebuild re-runs only when one of those tokens changes. A fixed set of boundaries always runs the complete rebuild regardless of the cache: checkpoint save and load (`state_dict` / `load_state_dict`), canonical export, import prepare, and import commit; `post_sharding` rebinding; 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. +After finalization, every entry point re-validates the published layout, and this has two costs. Steps and identity queries take an O(local params) fast path: the first complete forensic rebuild caches a verdict keyed by cheap identity tokens — the finalized registries by object identity, every live group container, parameter, and compatibility name, and a version counter that every mutating API bumps — and the full rebuild re-runs only when one of those tokens changes. A fixed set of boundaries always runs the complete rebuild regardless of the cache: checkpoint save and load (`state_dict` / `load_state_dict`), canonical export, import prepare, and import commit; `post_sharding` rebinding; collective codebook initialize and refresh; codebook-scope re-validation; and external contract-readiness queries. `post_sharding` additionally computes the manifest shard set and its sha256 digest once per finalized manifest, for reuse by the scoped operation headers. The practical consequence for an integrator is a clean split: any layout corruption reachable through the public containers — the group `params` / `param_names` slots, per-parameter state names, or the compatibility-name cache — still fails the step guard before any state is mutated, including corruption a closure introduces between the pre- and post-closure guards. Only corruption that leaves every fast-path token intact — in-place value replacement inside the private finalized registries, or an `object.__setattr__` on a frozen identity record — waits until the next boundary above to be caught rather than being caught at the next step. ## Explicit learned-codebook process groups @@ -118,23 +118,15 @@ 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`, a complete declared native state variant, and period one for selected or initialized state. Plain Gefen supports replicated and contiguous flattened element shards. Block-second-moment state can reshard between replicated and flattened targets; a logical matrix using factored second moments remains replicated and same-topology because factored-to-block representation migration is not implemented. GefenMuon supports replicated matrices and whole-parameter ownership when every participating group uses `sharded_mode="distributed"`; the transport can change placement and redistribute owners across world sizes, including NorMuon row state. A ready Hybrid declaration is the union of its heterogeneous child layouts and change kinds; an adapter must inspect each role's child contract together with the finalized immutable `optimizer.parameter_routing()` result rather than apply that union indiscriminately to every parameter. Pristine and period-selected states are supported under the same policy rules, and zero-element parameters remain pristine. The collective protocol exchanges fixed-size preparation headers before payload movement, visits member fragments in stable semantic order, bounds metadata and tensor chunks, propagates asymmetric local failures to every participant, and performs no semantic checks after the final freshness vote. Import preserves the target's parameter groups, defaults, parameters, compatibility names, and runtime process-group configuration while restoring portable common state, including the source deterministic setting. The atomic claim is fail-before-local-mutation for live, quiescent optimizer instances; it is not rollback after process death, backend failure, or concurrent mutation after the final vote. Ordinary state-dict hooks are bypassed. Adapters must quiesce training, avoid retaining state-container identities across a successful import, and persist the returned weights-only-safe CPU document with their checkpoint system. -Portable v3 currently excludes non-period-one initialized state, second-moment representation conversion, DTensor layouts, stochastic rounding, capturable/device-authoritative state, retained or migrated legacy parameter-state variants, tensor-valued or callable algorithm options, GefenMuon modes other than `distributed` for whole-owner transport, AdamW-backed `GefenMuonHybrid`, tied-alias expansion, asynchronous DCP, and mixed model/optimizer `Stateful` composition. The dedicated DCP helpers require every checkpoint member to enter synchronously and use the exact optimizer-owned process group. A singleton binding is rejected inside an initialized default world larger than one because the common PyTorch 2.5–2.12 `dcp.save/load` API interprets `process_group=None` as that world. A multi-member checkpoint group must have global rank zero at group coordinate zero; PyTorch 2.5's DCP coordinator path can otherwise address group coordinate zero as global rank zero and hang, so the adapter applies this compatibility restriction on every supported version. DCP storage publication is not transactionally atomic; the optimizer load remains fail-before-local-mutation after a complete successful read and verification. A CPU portable document remains a detached checkpoint artifact and is not evidence of live state offload. +Portable v3 currently excludes non-period-one initialized state, second-moment representation conversion, DTensor layouts, stochastic rounding, capturable/device-authoritative state, retained or migrated legacy parameter-state variants, tensor-valued or callable algorithm options, GefenMuon modes other than `distributed` for whole-owner transport, AdamW-backed `GefenMuonHybrid`, tied-alias expansion, asynchronous DCP, and mixed model/optimizer `Stateful` composition. The dedicated DCP helpers require every checkpoint member to enter synchronously and use the exact optimizer-owned process group. A singleton binding is rejected inside an initialized default world larger than one because the common PyTorch 2.5–2.12 `dcp.save/load` API interprets `process_group=None` as that world. A multi-member checkpoint group must have global rank zero at group coordinate zero; PyTorch 2.5's DCP coordinator path can otherwise address group coordinate zero as global rank zero and hang, so the adapter applies this compatibility restriction on every supported version. DCP storage publication is not transactionally atomic; the optimizer load remains fail-before-local-mutation after a complete successful read and verification. A CPU portable document remains a detached checkpoint artifact and does not change the live optimizer's declared capabilities. -## Quiescent optimizer-state movement and offload +## Optimizer-state placement capability fields -`StateMovementProvider.move_state_(device=None)` performs blocking CPU/CUDA co-location movement for Gefen and GefenMuon. With `device=None`, each declared authoritative per-parameter tensor moves to that parameter's current local device, including declared state attached to wrapper-orphaned parameter keys, while the canonical learned codebook moves to the first live local parameter device in parameter-group order. An optimizer with no local parameter storage keeps common state on CPU. An explicit device is accepted only after every live local parameter already resides there; an unindexed CUDA target is resolved from the one co-located live parameter device. The adapter must therefore move the model parameters first and invoke `move_state_` at a quiescent boundary before the next optimizer step. - -The core validates the finalized binding and complete declared state representation, allocates detached tight copies of the codebook and every authoritative tensor, and waits for all participating CUDA devices before one local publication. A preparation, transfer, synchronization, or validation failure leaves the exact live optimizer state, caches, parameters, gradients, groups, tensor learning rates, names, bindings, and metadata unchanged. Successful movement replaces the public `optimizer.state` mapping, every reachable or orphan per-parameter state dictionary, the canonical codebook identity, and every moved tensor identity; preserved non-tensor values and rank-local carrier tensors retain their identities, so adapters must not retain the replaced containers. Successful publication preserves host counters and metadata plus rank-local checkpoint carriers, discards only the rebuildable `stepsize` and `_h_buf` buffers, and invalidates per-device codebook/LUT copies, codebook-scope validation, the compiled static-address signature, and the tensor-learning-rate scalar cache. Preserved extension metadata is limited to provably tensor-free trees made from `None`, exact `bool`, `int`, `float`, `complex`, `str`, `bytes`, `torch.device`, `torch.dtype`, `torch.layout`, or `torch.memory_format` leaves and exact `dict`, `list`, `tuple`, `set`, `frozenset`, `deque`, or `torch.Size` containers. Cyclic or multiply referenced container graphs are outside that tree form. Arbitrary opaque objects, `defaultdict` or `OrderedDict` extension values, custom container subclasses, and non-dictionary per-parameter state mappings are rejected even when they appear tensor-free; the optimizer-owned top-level `state` may use its normal exact `defaultdict(dict)` representation. Undeclared tensor-bearing state, meta/nested/subclassed state tensors, FakeTensor parameters, capturable state, active compilation, and CUDA graph capture are likewise rejected rather than partially moved. - -`StateOffloadProvider.offload_state_(device="cpu")` enables synchronous CPU-authoritative per-parameter state for an exact plain `Gefen` instance with ordinary replicated CUDA parameters. Activation first validates the complete declared state, stages tight detached CPU copies, waits for CUDA transfers, and publishes the policy and replacement state mapping together. At each eager step, Gefen copies only the current parameter's persistent tensor state to that parameter's CUDA device, runs the ordinary fused or non-fused block or factored update against a private runtime dictionary, synchronously copies the updated persistent state back to CPU, publishes that one dictionary, and releases the device temporaries. The optimizer-common learned codebook remains resident on CUDA and its normal per-device caches remain available. `restore_state_()` atomically co-locates all state with the parameters and disables offload; `move_state_()` has the same policy-disabling effect after its requested movement succeeds. - -Activation and restore are fail-before-mutation. Activation also rejects persistent state tensors whose storage overlaps another persistent field, a parameter, or the common codebook because independent parameter paging cannot preserve such aliasing. Every offloaded step then re-runs the complete readiness scan — per-tensor storage validation, pairwise disjointness, and native-schema validation — before any parameter is staged, so an external in-place edit that corrupts an already-validated offloaded state tensor is caught at step entry, before the step mutates any parameter. This scan is O(local params) and is deliberately not cached across steps: unlike the finalized layout, the per-parameter offloaded state tensors are legitimately replaced on every step, so no cached verdict could stand in for them. If the update itself raises, Gefen attempts to preserve the resulting runtime state on CPU before propagating the original error. If copyback fails after a parameter may have changed, the optimizer is marked poisoned and refuses subsequent steps or native, canonical, and portable exports until a complete successful native `load_state_dict()` establishes known-good state. An active offload policy is target-local runtime configuration and is preserved across such a load rather than serialized as checkpoint meaning; the active loader maps parameter state directly to CPU and never accumulates the checkpoint's full parameter state on CUDA. State offload is implemented only for an exact plain `Gefen` instance; the composite Hybrid API, `GefenMuon`, nonreplicated finalized layouts, DTensor or tensor-subclass parameters, opaque extension state, multi-member explicit codebook scopes, capturable/device-authoritative state, compilation, and CUDA graph capture are excluded. The multi-member exclusion prevents one rank's copyback poison from bypassing the next scoped collective while peers enter it. Offload must be restored before post-sharding rebinding. It is blocking parameter-scoped paging, not asynchronous prefetch, overlap, or a distributed offload engine, and portable global-state I/O remains unavailable while its authoritative tensors are parked on CPU. - -`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. +`atomic_state_movement` and `state_offload` remain in schema version 1 so adapters receive explicit negative declarations instead of inferring support from device placement, checkpoint contents, or private implementation details. They are false for live optimizers and do not expose movement or offload operations through this contract. Plain Gefen declares replicated, flattened element-shard, and the narrow DTensor training layouts. Its PyTorch rank-local checkpoint transport is same-topology only. GefenMuon declares replicated and narrow DTensor training, with mode-specific state extents: `approx` state is local, `exact` state is logically global, and `distributed` momentum is held by the parameter owner while non-owners retain metadata only. Native Parallel-Muon checkpoints separately declare world-size owner redistribution, not placement-changing resharding. `GefenMuonHybrid` retains its nested child namespaces and does not flatten AdamW or Gefen child state into a fabricated common schema. diff --git a/scripts/release_gpu_gate.sh b/scripts/release_gpu_gate.sh index faa4a7b..8fb15fc 100755 --- a/scripts/release_gpu_gate.sh +++ b/scripts/release_gpu_gate.sh @@ -216,8 +216,6 @@ GEFEN_VERBOSE_BUILD=1 \ tests/test_gefen_fsdp2_checkpoint.py \ tests/test_muon_distributed_checkpoint_safety.py \ tests/test_muon_grad_presence.py \ - tests/test_state_offload.py \ - tests/test_state_movement.py \ tests/test_step_preflight_atomicity.py \ tests/test_muon_fsdp2_parity.py::test_muon_fsdp2_fused_cuda_parity \ tests/test_muon_fsdp2_parity.py::test_muon_fsdp2_fused_multirank_parity \ diff --git a/src/gefen/__init__.py b/src/gefen/__init__.py index e66eef9..80075ac 100644 --- a/src/gefen/__init__.py +++ b/src/gefen/__init__.py @@ -45,8 +45,6 @@ "StateField", "StateGeometry", "StateKeyMatch", - "StateMovementProvider", - "StateOffloadProvider", "StateScope", "StateVariant", "TopologyChange", @@ -144,8 +142,6 @@ def __getattr__(name): "StateField", "StateGeometry", "StateKeyMatch", - "StateMovementProvider", - "StateOffloadProvider", "StateScope", "StateVariant", "TopologyChange", diff --git a/src/gefen/contracts.py b/src/gefen/contracts.py index 37144fa..7b5159f 100644 --- a/src/gefen/contracts.py +++ b/src/gefen/contracts.py @@ -1150,29 +1150,6 @@ def import_portable_state( """Collectively stage and atomically publish portable global state.""" -@runtime_checkable -class StateMovementProvider(Protocol): - """Structural protocol for quiescent atomic optimizer-state movement.""" - - def move_state_(self, device=None) -> None: - """Co-locate authoritative state with the optimizer's live parameters.""" - - -@runtime_checkable -class StateOffloadProvider(Protocol): - """Structural protocol for persistent live optimizer-state offload.""" - - @property - def state_offload_device(self): - """Return the active offload device, or ``None`` while resident.""" - - def offload_state_(self, device="cpu") -> None: - """Atomically park supported state and enable transparent step paging.""" - - def restore_state_(self) -> None: - """Atomically co-locate state with parameters and disable offload.""" - - _ALL_PRECISIONS = frozenset( {Precision.FLOAT32, Precision.BFLOAT16, Precision.FLOAT16, Precision.FLOAT64} ) @@ -1336,8 +1313,6 @@ def _negative_capabilities( shard_rebinding: bool = False, post_sharding: bool = False, canonical_state_io: bool = False, - atomic_state_movement: bool = False, - state_offload: bool = False, ) -> OptimizerCapabilities: return OptimizerCapabilities( training=training, @@ -1351,8 +1326,8 @@ def _negative_capabilities( shard_rebinding=shard_rebinding, post_sharding=post_sharding, canonical_state_io=canonical_state_io, - atomic_state_movement=atomic_state_movement, - state_offload=state_offload, + atomic_state_movement=False, + state_offload=False, ) @@ -1367,8 +1342,6 @@ 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(), - atomic_state_movement: bool = False, - state_offload: bool = False, ) -> OptimizerContract: canonical_state_layouts = _frozenset(canonical_state_layouts) canonical_global_same_topology = _frozenset(canonical_global_same_topology) @@ -1586,8 +1559,6 @@ def _gefen_contract( or canonical_global_same_topology or canonical_global_topology_changing ), - atomic_state_movement=atomic_state_movement, - state_offload=state_offload, ), ) @@ -1625,7 +1596,6 @@ def _gefen_muon_contract( canonical_global_same_topology: AbstractSet[ParameterLayout] = frozenset(), canonical_global_topology_changing: AbstractSet[ParameterLayout] = frozenset(), canonical_global_topology_change_kinds: AbstractSet[TopologyChange] = frozenset(), - atomic_state_movement: bool = False, ) -> OptimizerContract: canonical_state_layouts = _frozenset(canonical_state_layouts) canonical_global_same_topology = _frozenset(canonical_global_same_topology) @@ -1950,7 +1920,6 @@ def _gefen_muon_contract( or canonical_global_same_topology or canonical_global_topology_changing ), - atomic_state_movement=atomic_state_movement, ), ) @@ -2093,8 +2062,6 @@ def _hybrid_contract( "StateField", "StateGeometry", "StateKeyMatch", - "StateMovementProvider", - "StateOffloadProvider", "StateScope", "StateVariant", "TrainingSupport", diff --git a/src/gefen/gefen.py b/src/gefen/gefen.py index 7bddd3f..82de0ab 100644 --- a/src/gefen/gefen.py +++ b/src/gefen/gefen.py @@ -16,8 +16,7 @@ import math import os import warnings -from collections import defaultdict, deque, OrderedDict -from copy import deepcopy +from collections import defaultdict, OrderedDict from itertools import chain from typing import Iterable, Optional, Tuple, Union @@ -117,57 +116,6 @@ "m_codebook_shape", } ) -_STATE_MOVEMENT_TENSOR_KEYS = frozenset( - { - "step", - "m_codebook", - "m_magnitude", - "vmean", - "vmean_step", - "v_row", - "v_col", - "factored_step", - "normuon_v", - "normuon_step", - } -) -_STATE_MOVEMENT_COUNTER_KEYS = frozenset( - {"step", "vmean_step", "factored_step", "normuon_step"} -) -_STATE_MOVEMENT_SCRATCH_KEYS = frozenset({"stepsize", "_h_buf"}) -_STATE_MOVEMENT_CAPTURABLE_KEYS = frozenset( - { - "_capt_scalars", - "_capt_consts", - "_capt_consts_key", - "_capt_stack", - "_capt_row", - } -) -_STATE_MOVEMENT_DEVICE_TYPES = frozenset({"cpu", "cuda"}) -_STATE_MOVEMENT_METADATA_LEAF_TYPES = ( - bool, - int, - float, - complex, - str, - bytes, - torch.device, - torch.dtype, - torch.layout, - torch.memory_format, -) -_STATE_MOVEMENT_METADATA_MAPPING_TYPES = (dict,) -_STATE_MOVEMENT_METADATA_SEQUENCE_TYPES = ( - list, - tuple, - set, - frozenset, - deque, - torch.Size, -) - - def _rank_local_payload_key(global_rank: int) -> str: return "{}{}".format(_RANK_LOCAL_PAYLOAD_KEY_PREFIX, int(global_rank)) @@ -1302,18 +1250,6 @@ def __init__( # one device counter per parameter device and advance it in the captured # step tail; state_dict synchronizes the host mirror before serializing. self._gefen_global_step_by_device = {} - # Native parameter-state offload is an eager, target-local runtime - # policy rather than optimizer checkpoint meaning. None keeps the - # historical co-located path. A CPU device means declared persistent - # per-parameter tensors are CPU-authoritative between steps while the - # small optimizer-common codebook remains resident/cached normally. - self._gefen_state_offload_device = None - # A failed CUDA-to-CPU copyback can leave a parameter updated while its - # last published CPU state is stale. Preserve that diagnosis across - # restore/movement and reject later steps until a successful load - # establishes a complete known-good state again. - self._gefen_state_offload_poisoned = False - defaults = dict( lr=lr, beta1=betas[0], @@ -1334,11 +1270,10 @@ def __init__( # Layout-forensics fast path. After one full structural pass succeeds, # the exact container identities it validated are remembered together # with a version counter that every legitimate mutating API bumps - # (post_sharding, staged checkpoint load commits, state movement and - # offload). Per-step guards accept only that unchanged token set; - # boundary operations (checkpoint prepare/commit, rebinding, movement - # and offload, collective codebook initialize/refresh, contract - # readiness) always rerun the complete forensic rebuild. + # (post_sharding and staged checkpoint load commits). Per-step guards + # accept only that unchanged token set; boundary operations (checkpoint + # prepare/commit, rebinding, collective codebook initialize/refresh, + # and contract readiness) always rerun the complete forensic rebuild. self._gefen_layout_version = 0 self._gefen_layout_forensics_verdict = None # (manifest, frozenset(manifest.shards), sha256 digest) computed once @@ -1371,10 +1306,7 @@ def optimizer_contract(self) -> OptimizerContract: """Return the immutable state-layout and integration capability contract.""" identity_ready = self._canonical_identity_ready() - try: - canonical_state_layouts = self._canonical_state_layouts() - except Exception: - canonical_state_layouts = frozenset() + canonical_state_layouts = self._canonical_state_layouts() try: from gefen.portable_runtime import _portable_runtime_layouts @@ -1411,8 +1343,6 @@ 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, - atomic_state_movement=self._atomic_state_movement_supported(), - state_offload=self._state_offload_supported(), native_flattened_checkpoint=( self._codebook_scope_ready() and any( @@ -1448,8 +1378,6 @@ def _canonical_group_options_value(group): } def _canonical_state_layouts(self): - if self.state_offload_poisoned: - return frozenset() if not self._canonical_identity_ready(): return frozenset() if self._stochastic_round: @@ -2432,10 +2360,6 @@ def _target_may_have_internal_storage_overlap(parameter) -> bool: def _assert_rebinding_pristine(self, rebindings) -> None: if self._gefen_post_sharding_finalized: raise RuntimeError("Gefen post-sharding identity is already finalized") - if self.state_offload_active or self.state_offload_poisoned: - raise RuntimeError( - "Gefen parameter rebinding requires resident known-good state; restore state first" - ) if ( self._gefen_shard_bindings or self._gefen_local_shard_bindings @@ -2658,6 +2582,17 @@ def _assert_runtime_codebook_process_group(self, *, full: bool = False) -> None: self._assert_finalized_binding_layout(full=full) self._validate_codebook_runtime_binding(binding) + def _capture_codebook_scope_binding_for_step(self): + """Return a usable failure-vote binding before inspecting live layout.""" + + binding = self._gefen_codebook_process_group + if binding is None: + return None + if not isinstance(binding, CodebookProcessGroupBinding): + raise TypeError("the runtime codebook process-group binding is invalid") + self._validate_codebook_runtime_binding(binding) + return binding + def _codebook_parameter_contributes(self, parameter) -> bool: binding = self._gefen_codebook_process_group if binding is None: @@ -2963,834 +2898,6 @@ def rebind_parameter( manifest=ShardingManifest((shard,)), ) - @staticmethod - def _state_value_is_movement_safe_metadata(value, seen=None) -> bool: - value_type = type(value) - if value is None or value_type in _STATE_MOVEMENT_METADATA_LEAF_TYPES: - return True - is_mapping = value_type in _STATE_MOVEMENT_METADATA_MAPPING_TYPES - if not is_mapping and value_type not in _STATE_MOVEMENT_METADATA_SEQUENCE_TYPES: - return False - if seen is None: - seen = set() - value_id = id(value) - if value_id in seen: - return False - seen.add(value_id) - items = value.items() if is_mapping else value - if is_mapping: - return all( - Gefen._state_value_is_movement_safe_metadata(key, seen) - and Gefen._state_value_is_movement_safe_metadata(item, seen) - for key, item in items - ) - return all( - Gefen._state_value_is_movement_safe_metadata(item, seen) for item in items - ) - - @staticmethod - def _state_movement_tensor_supported(value) -> bool: - return ( - type(value) is torch.Tensor - and not (hasattr(value, "to_local") and hasattr(value, "placements")) - and value.layout is torch.strided - and not value.is_nested - and not value.is_quantized - and getattr(value, "fake_mode", None) is None - and value.device.type in _STATE_MOVEMENT_DEVICE_TYPES - ) - - @property - def state_offload_active(self) -> bool: - """Whether native parameter-state CPU offload is currently enabled.""" - - return getattr(self, "_gefen_state_offload_device", None) is not None - - @property - def state_offload_device(self): - """Return the active state-offload device, or ``None`` when disabled.""" - - return getattr(self, "_gefen_state_offload_device", None) - - @property - def state_offload_poisoned(self) -> bool: - """Whether a failed copyback made further stepping unsafe.""" - - return bool(getattr(self, "_gefen_state_offload_poisoned", False)) - - def _assert_state_export_safe(self) -> None: - if self.state_offload_poisoned: - raise RuntimeError( - "Gefen cannot export optimizer state after a failed state-offload copyback; " - "load a known-good checkpoint first" - ) - - @staticmethod - def _state_offload_parameter_supported(parameter) -> bool: - return ( - type(parameter) in {torch.Tensor, nn.Parameter} - and not ( - hasattr(parameter, "to_local") and hasattr(parameter, "placements") - ) - and parameter.layout is torch.strided - and parameter.device.type == "cuda" - and parameter.dtype - in {torch.float16, torch.bfloat16, torch.float32, torch.float64} - and not torch.is_complex(parameter) - and not parameter.is_meta - and not parameter.is_nested - and not parameter.is_quantized - and getattr(parameter, "fake_mode", None) is None - ) - - @classmethod - def _state_offload_cpu_tensor_supported(cls, value) -> bool: - return ( - cls._state_movement_tensor_supported(value) - and value.device.type == "cpu" - and not value.requires_grad - and value.is_contiguous() - and value.storage_offset() == 0 - and value.untyped_storage().nbytes() == value.numel() * value.element_size() - ) - - @classmethod - def _state_offload_storage_disjoint(cls, parameters, state, codebook) -> bool: - """Conservatively reject aliases that parameter-scoped paging cannot preserve.""" - - tensors = [("parameter", parameter) for parameter in parameters] - if torch.is_tensor(codebook): - tensors.append(("common_state", codebook)) - for parameter_state in state.values(): - for key, value in parameter_state.items(): - if key in _STATE_MOVEMENT_TENSOR_KEYS and torch.is_tensor(value): - tensors.append(("parameter_state", value)) - - storage_ranges = [] - try: - for kind, tensor in tensors: - if tensor.numel() == 0: - continue - if ( - kind != "parameter" - and cls._target_may_have_internal_storage_overlap(tensor) - ): - return False - storage = tensor.untyped_storage() - storage_id = (str(tensor.device), storage.data_ptr()) - if tensor.is_contiguous(): - start = tensor.storage_offset() * tensor.element_size() - end = start + tensor.numel() * tensor.element_size() - else: - start = end = None - for other_kind, other_id, other_start, other_end in storage_ranges: - if other_id != storage_id or (kind == other_kind == "parameter"): - continue - if ( - start is None - or other_start is None - or max(start, other_start) < min(end, other_end) - ): - return False - storage_ranges.append((kind, storage_id, start, end)) - except Exception: - return False - return True - - @staticmethod - def _state_offload_capturing_on_parameter_device(parameters) -> bool: - if torch.cuda.is_available() and torch.cuda.is_current_stream_capturing(): - return True - devices = { - parameter.device - for parameter in parameters - if torch.is_tensor(parameter) and parameter.device.type == "cuda" - } - devices = sorted( - devices, - key=lambda device: -1 if device.index is None else device.index, - ) - for device in devices: - with torch.cuda.device(device): - if torch.cuda.is_current_stream_capturing(): - return True - return False - - def _state_offload_rejection_reason( - self, - *, - require_cpu_state: bool, - allow_poisoned: bool = False, - require_full_layout: bool = False, - ): - # ``require_full_layout`` controls only the immutable finalized-LAYOUT - # forensics reached through ``_state_movement_rejection_reason``. It - # defaults to False so the per-step offload-readiness caller reuses the - # memoized layout verdict; boundary callers (activation, load, contract - # readiness) pass True for a full rebuild. Every per-tensor offload - # check below still runs on each call regardless of this flag, because - # offloaded per-parameter state tensors are legitimately replaced each - # step and a cached verdict there would miss corruption. - if type(self) is not Gefen: - return "native state offload is implemented only by plain Gefen" - if self.state_offload_active and self.state_offload_device != torch.device( - "cpu" - ): - return "the active state-offload policy has an invalid device" - if self.state_offload_poisoned and not allow_poisoned: - return "a previous state copyback failed; load a known-good checkpoint" - scope = self._gefen_codebook_process_group - if scope is not None: - if type(scope) is not CodebookProcessGroupBinding: - return "the explicit codebook process-group binding is invalid" - if len(scope.identity.ordered_members) > 1: - return "state offload does not support multi-member explicit codebook scopes" - if self.capturable: - return "capturable optimizers have device-authoritative replay state" - try: - parameters = [ - parameter - for group in self.param_groups - for parameter in group["params"] - ] - except (KeyError, TypeError): - return "parameter groups have an invalid structure" - if torch.compiler.is_compiling(): - return "state offload cannot run during torch.compile" - if torch.cuda.is_available(): - try: - if self._state_offload_capturing_on_parameter_device(parameters): - return "state offload cannot run during CUDA graph capture" - except RuntimeError: - return "the CUDA graph-capture state could not be inspected" - - movement_reason = self._state_movement_rejection_reason( - require_full_layout=require_full_layout - ) - if movement_reason is not None: - return movement_reason - - if not parameters: - return "state offload requires at least one CUDA parameter" - if any( - not self._state_offload_parameter_supported(parameter) - for parameter in parameters - ): - return "state offload requires ordinary replicated CUDA parameters" - - if self._gefen_post_sharding_finalized: - if len(self._gefen_local_shard_bindings) != len(parameters): - return "the finalized local shard registry is incomplete" - if any( - parameter is None or shard.layout is not ParameterLayout.REPLICATED - for parameter, shard in self._gefen_local_shard_bindings - ): - return "state offload supports only finalized replicated layouts" - - state_type = type(self.state) - if not ( - state_type is dict - or (state_type is defaultdict and self.state.default_factory is dict) - ): - return "optimizer state must use a supported standard mapping" - if set(self.state) != set(parameters): - return "state offload requires exactly one state entry per live parameter" - - allowed_keys = ( - _CANONICAL_PARAMETER_STATE_KEYS | _CANONICAL_DERIVED_PARAMETER_STATE_KEYS - ) - for parameter in parameters: - parameter_state = self.state.get(parameter) - if type(parameter_state) is not dict: - return "per-parameter optimizer state must use a plain dictionary" - if any(key not in allowed_keys for key in parameter_state): - return "state offload does not support custom per-parameter state" - if require_cpu_state and any( - key in _STATE_MOVEMENT_SCRATCH_KEYS for key in parameter_state - ): - return "offloaded state contains device-side runtime scratch" - for key, value in parameter_state.items(): - if key in _STATE_MOVEMENT_COUNTER_KEYS: - if type(value) is not int: - return ( - "noncapturable offloaded counters must be Python integers" - ) - continue - if key not in _STATE_MOVEMENT_TENSOR_KEYS: - continue - if not torch.is_tensor(value): - return "authoritative offloaded state has an invalid value" - if require_cpu_state: - if not self._state_offload_cpu_tensor_supported(value): - return ( - "authoritative offloaded tensors must be tight CPU tensors" - ) - elif not self._state_movement_tensor_supported(value): - return ( - "authoritative tensor state has an unsupported representation" - ) - - if not self._state_offload_storage_disjoint( - parameters, - self.state, - self._gefen_codebook, - ): - return "persistent optimizer-state storage aliases another live tensor" - - try: - self._validate_loaded_native_state() - except Exception: - return "optimizer state does not match Gefen's declared native schema" - if ( - require_cpu_state - and self._gefen_codebook is not None - and self._gefen_codebook.device.type != "cuda" - ): - return "the optimizer-common codebook must remain CUDA-resident" - return None - - def _state_offload_supported(self) -> bool: - """Return whether this live optimizer can safely enter or retain offload.""" - - try: - reason = self._state_offload_rejection_reason( - require_cpu_state=self.state_offload_active, - require_full_layout=True, - ) - except Exception: - return False - return reason is None - - @staticmethod - def _normalize_state_offload_target(device) -> torch.device: - try: - target = torch.device(device) - except (TypeError, RuntimeError) as exc: - raise TypeError("state offload device must be CPU") from exc - if target.type != "cpu": - raise ValueError("Gefen native state offload currently supports only CPU") - return torch.device("cpu") - - def _copy_state_tensor_to_offload_cpu(self, tensor: torch.Tensor) -> torch.Tensor: - return self._copy_state_tensor_for_move(tensor, torch.device("cpu")) - - def _stage_state_offload_resident_codebook(self): - codebook = self._gefen_codebook - if codebook is None or codebook.device.type == "cuda": - return codebook - parameter = next( - parameter for group in self.param_groups for parameter in group["params"] - ) - target = parameter.device - staged = self._copy_state_tensor_for_move(codebook, target) - self._validate_staged_state_tensor(codebook, staged, target) - torch.cuda.synchronize(target) - return staged - - def _prepare_offloaded_cpu_parameter_state(self, parameter_state): - if type(parameter_state) is not dict: - raise RuntimeError("offloaded runtime state must use a plain dictionary") - result = {} - cuda_devices = set() - allowed_keys = ( - _CANONICAL_PARAMETER_STATE_KEYS | _CANONICAL_DERIVED_PARAMETER_STATE_KEYS - ) - for key, value in parameter_state.items(): - if key not in allowed_keys: - raise RuntimeError( - "offloaded stepping produced unsupported state key {!r}".format(key) - ) - if key in _STATE_MOVEMENT_SCRATCH_KEYS: - continue - if key in _STATE_MOVEMENT_CAPTURABLE_KEYS: - raise RuntimeError( - "offloaded stepping produced capturable runtime state" - ) - if key in _STATE_MOVEMENT_COUNTER_KEYS: - if type(value) is not int: - raise RuntimeError( - "offloaded stepping produced a device-authoritative counter" - ) - result[key] = value - continue - if key in _STATE_MOVEMENT_TENSOR_KEYS: - if not self._state_movement_tensor_supported(value): - raise RuntimeError( - "offloaded stepping produced unsupported tensor state" - ) - staged = self._copy_state_tensor_to_offload_cpu(value) - self._validate_staged_state_tensor(value, staged, torch.device("cpu")) - result[key] = staged - if value.device.type == "cuda": - cuda_devices.add(value.device) - continue - result[key] = value - - for cuda_device in sorted( - cuda_devices, key=lambda item: -1 if item.index is None else item.index - ): - torch.cuda.synchronize(cuda_device) - return result - - def _stage_all_parameter_state_to_cpu(self): - staged_state = defaultdict(dict) - for parameter, parameter_state in self.state.items(): - staged_state[parameter] = self._prepare_offloaded_cpu_parameter_state( - parameter_state - ) - return staged_state - - def _stage_offloaded_parameter_state(self, parameter): - cpu_state = self.state[parameter] - if type(cpu_state) is not dict: - raise RuntimeError("offloaded parameter state must use a plain dictionary") - target = parameter.device - runtime_state = {} - for key, value in cpu_state.items(): - if key in _STATE_MOVEMENT_COUNTER_KEYS: - if type(value) is not int: - raise RuntimeError( - "offloaded parameter counters must remain Python integers" - ) - runtime_state[key] = value - elif key in _STATE_MOVEMENT_TENSOR_KEYS: - if not self._state_offload_cpu_tensor_supported(value): - raise RuntimeError( - "offloaded parameter tensors must remain tight CPU tensors" - ) - staged = self._copy_state_tensor_for_move(value, target) - self._validate_staged_state_tensor(value, staged, target) - runtime_state[key] = staged - else: - runtime_state[key] = value - torch.cuda.synchronize(target) - return runtime_state - - def _step_with_offloaded_parameter_state( - self, update, group, param_name, parameter, grad - ) -> None: - runtime_state = self._stage_offloaded_parameter_state(parameter) - try: - update(group, param_name, parameter, grad, state=runtime_state) - except BaseException as operation_error: - try: - cpu_state = self._prepare_offloaded_cpu_parameter_state(runtime_state) - except BaseException as copyback_error: - self._gefen_state_offload_poisoned = True - if hasattr(operation_error, "add_note"): - operation_error.add_note( - "Gefen state copyback also failed; the optimizer is poisoned" - ) - raise operation_error from copyback_error - self.state[parameter] = cpu_state - raise - - try: - cpu_state = self._prepare_offloaded_cpu_parameter_state(runtime_state) - except BaseException as exc: - self._gefen_state_offload_poisoned = True - raise RuntimeError( - "Gefen state copyback failed after a parameter update; load a " - "known-good checkpoint before stepping again" - ) from exc - self.state[parameter] = cpu_state - - def _assert_state_offload_step_ready(self) -> None: - if self.state_offload_poisoned: - raise RuntimeError( - "Gefen state offload is poisoned after a failed copyback; load a " - "known-good checkpoint before stepping again" - ) - if not self.state_offload_active: - return - # Run the complete offload scan (per-tensor storage checks, pairwise - # disjointness, native-schema validation) on every step before any - # parameter is staged. Unlike the finalized-layout manifest, the - # per-parameter offloaded state tensors are legitimately replaced each - # step, so a cached verdict cannot represent them; a token-preserving - # in-place corruption of a later parameter's state would otherwise slip - # past step entry and only be caught mid-step, after earlier parameters - # were already updated and committed. The scan is O(local params) and - # was never the layout-forensics cost this cache was introduced for. - reason = self._state_offload_rejection_reason(require_cpu_state=True) - if reason is not None: - raise RuntimeError("Gefen state offload cannot step: {}".format(reason)) - - @torch.no_grad() - def offload_state_(self, device="cpu") -> None: - """Atomically enable synchronous CPU-authoritative parameter state.""" - - target = self._normalize_state_offload_target(device) - self._assert_finalized_binding_layout(full=True) - reason = self._state_offload_rejection_reason( - require_cpu_state=False, require_full_layout=True - ) - if reason is not None: - raise RuntimeError("Gefen state offload is unavailable: {}".format(reason)) - staged_state = self._stage_all_parameter_state_to_cpu() - staged_codebook = self._stage_state_offload_resident_codebook() - updates = { - "state": staged_state, - "_gefen_state_offload_device": target, - "_static_mark_sig": None, - "_lr_scalar_cache": None, - } - if staged_codebook is not self._gefen_codebook: - updates.update( - { - "_gefen_codebook": staged_codebook, - "_gefen_codebook_by_device": {}, - "_gefen_codebook_lut_by_device": {}, - "_gefen_codebook_scope_validated": False, - } - ) - self.__dict__.update(updates) - self._invalidate_layout_forensics_caches() - - @torch.no_grad() - def restore_state_(self) -> None: - """Atomically co-locate parameter state and disable native offload.""" - - if not self.state_offload_active: - return - parameters = [ - parameter for group in self.param_groups for parameter in group["params"] - ] - if torch.cuda.is_available(): - try: - capturing = self._state_offload_capturing_on_parameter_device( - parameters - ) - except RuntimeError as exc: - raise RuntimeError( - "Gefen state restore could not inspect CUDA graph-capture state" - ) from exc - if capturing: - raise RuntimeError( - "Gefen state restore cannot run during CUDA graph capture" - ) - self.move_state_() - - def _state_movement_rejection_reason(self, *, require_full_layout: bool = True): - # The finalized layout is immutable across steps, so the per-step - # offload-readiness caller passes ``require_full_layout=False`` to reuse - # the memoized layout-forensics verdict (an O(local params) identity - # token check) instead of forcing an uncached full rebuild + manifest - # digest recompute every step. Boundary callers (``move_state_``, - # ``_atomic_state_movement_supported``) keep the default full rebuild. - # The per-tensor state checks below always run regardless. - if ( - self._gefen_post_sharding_finalized - and not self._finalized_binding_layout_matches(full=require_full_layout) - ): - return "the finalized parameter binding no longer matches live groups" - if self.capturable: - return "capturable optimizers have device-authoritative replay state" - if self._capt_stacks is not None: - return "capturable state stacks are active" - if self._gefen_global_step_by_device or self._sr_seed_by_device: - return "capturable device counters or stochastic-rounding seeds are active" - - try: - parameters = [ - parameter - for group in self.param_groups - for parameter in group["params"] - ] - except (KeyError, TypeError): - return "parameter groups have an invalid structure" - for parameter in parameters: - if not torch.is_tensor(parameter): - return "parameter groups contain a non-tensor value" - if getattr(parameter, "fake_mode", None) is not None: - return "FakeTensor parameters do not have movable storage" - try: - self._state_move_parameter_device(parameter) - except Exception: - return "parameters must use CPU or CUDA local storage" - - state_type = type(self.state) - if not ( - state_type is dict - or (state_type is defaultdict and self.state.default_factory is dict) - ): - return "optimizer state must use a supported standard mapping" - for parameter, parameter_state in self.state.items(): - if not torch.is_tensor(parameter): - return "optimizer state is keyed by a non-tensor value" - if getattr(parameter, "fake_mode", None) is not None: - return "FakeTensor state keys do not have movable storage" - try: - self._state_move_parameter_device(parameter) - except Exception: - return "optimizer state is keyed by a parameter without CPU or CUDA local storage" - if type(parameter_state) is not dict: - return "per-parameter optimizer state must use a plain dictionary" - for key, value in parameter_state.items(): - if not self._state_value_is_movement_safe_metadata(key): - return "optimizer state contains an unsupported key" - if key in _STATE_MOVEMENT_SCRATCH_KEYS: - continue - if key in _STATE_MOVEMENT_CAPTURABLE_KEYS: - return "capturable per-parameter state is active" - if key in _STATE_MOVEMENT_TENSOR_KEYS: - if torch.is_tensor(value): - if not self._state_movement_tensor_supported(value): - return "authoritative tensor state has an unsupported representation" - elif ( - key not in _STATE_MOVEMENT_COUNTER_KEYS - or type(value) is not int - ): - return "authoritative tensor state has an invalid value" - continue - if isinstance(key, str) and key.startswith( - _RANK_LOCAL_PAYLOAD_KEY_PREFIX - ): - if torch.is_tensor(value): - if not self._state_movement_tensor_supported(value): - return "rank-local transport state has an unsupported representation" - elif not self._state_value_is_movement_safe_metadata(value): - return ( - "rank-local transport state contains unsupported metadata" - ) - continue - if not self._state_value_is_movement_safe_metadata(value): - return "undeclared optimizer state is not provably tensor-free metadata" - - codebook = self._gefen_codebook - if codebook is not None and not self._state_movement_tensor_supported(codebook): - return "the canonical codebook has an unsupported tensor representation" - return None - - def _atomic_state_movement_supported(self) -> bool: - try: - reason = self._state_movement_rejection_reason() - except Exception: - return False - if reason is not None: - return False - if torch.compiler.is_compiling(): - return False - if torch.cuda.is_available(): - try: - parameters = [ - parameter - for group in self.param_groups - for parameter in group["params"] - ] - if self._state_offload_capturing_on_parameter_device(parameters): - return False - except RuntimeError: - return False - return True - - @staticmethod - def _state_tensor_device(parameter: torch.Tensor) -> torch.device: - local = parameter.to_local() if hasattr(parameter, "to_local") else parameter - if hasattr(local, "wait"): - local = local.wait() - return local.device - - @classmethod - def _state_move_parameter_device(cls, parameter: torch.Tensor) -> torch.device: - local = parameter.to_local() if hasattr(parameter, "to_local") else parameter - if hasattr(local, "wait"): - local = local.wait() - if not torch.is_tensor(local): - raise RuntimeError("parameter local storage must be a tensor") - device = local.device - if device.type not in _STATE_MOVEMENT_DEVICE_TYPES: - raise RuntimeError("state movement supports only CPU and CUDA parameters") - return torch.device("cpu") if device.type == "cpu" else device - - @staticmethod - def _normalize_state_move_target(device, live_devices) -> torch.device: - try: - target = torch.device(device) - except (TypeError, RuntimeError) as exc: - raise TypeError("device must identify a CPU or CUDA device") from exc - if target.type not in _STATE_MOVEMENT_DEVICE_TYPES: - raise ValueError("Gefen state movement supports only CPU and CUDA devices") - if target.type == "cpu": - target = torch.device("cpu") - else: - if not torch.cuda.is_available(): - raise ValueError( - "CUDA state movement requires an available CUDA device" - ) - if target.index is None: - unique_devices = set(live_devices) - if len(unique_devices) != 1: - raise ValueError( - "an unindexed CUDA target requires one co-located live parameter device" - ) - candidate = next(iter(unique_devices)) - if candidate.type != "cuda": - raise ValueError( - "CUDA state movement requires CUDA-resident parameters" - ) - target = candidate - if target.index < 0 or target.index >= torch.cuda.device_count(): - raise ValueError( - "CUDA state movement target is not an available device" - ) - - if not live_devices: - if target.type != "cpu": - raise ValueError( - "an optimizer without local parameters keeps common state on CPU" - ) - elif any(parameter_device != target for parameter_device in live_devices): - raise ValueError( - "explicit state movement requires every live parameter to already be co-located on {}".format( - target - ) - ) - return target - - @staticmethod - def _copy_state_tensor_for_move( - tensor: torch.Tensor, device: torch.device - ) -> torch.Tensor: - return tensor.to( - device=device, - dtype=tensor.dtype, - non_blocking=False, - copy=True, - memory_format=torch.contiguous_format, - ).detach() - - @classmethod - def _validate_staged_state_tensor( - cls, source: torch.Tensor, staged, device: torch.device - ) -> None: - if ( - not cls._state_movement_tensor_supported(staged) - or staged is source - or staged.device != device - or staged.dtype != source.dtype - or tuple(staged.shape) != tuple(source.shape) - or staged.requires_grad - or not staged.is_contiguous() - or staged.storage_offset() != 0 - or staged.untyped_storage().nbytes() - != staged.numel() * staged.element_size() - ): - raise RuntimeError("state movement produced an invalid staged tensor") - - def _stage_state_move(self, device): - live_parameters = [ - parameter for group in self.param_groups for parameter in group["params"] - ] - live_devices = [ - self._state_move_parameter_device(parameter) - for parameter in live_parameters - ] - live_parameter_ids = {id(parameter) for parameter in live_parameters} - explicit_target = ( - None - if device is None - else self._normalize_state_move_target(device, live_devices) - ) - codebook_target = ( - explicit_target - if explicit_target is not None - else (live_devices[0] if live_devices else torch.device("cpu")) - ) - staged_state = defaultdict(dict) - cuda_devices = set() - - def stage_tensor(value, target): - staged = self._copy_state_tensor_for_move(value, target) - self._validate_staged_state_tensor(value, staged, target) - if value.device.type == "cuda": - cuda_devices.add(value.device) - if target.type == "cuda": - cuda_devices.add(target) - return staged - - staged_codebook = ( - None - if self._gefen_codebook is None - else stage_tensor(self._gefen_codebook, codebook_target) - ) - for parameter, parameter_state in self.state.items(): - staged_parameter_state = {} - parameter_target = ( - explicit_target if id(parameter) in live_parameter_ids else None - ) - for key, value in parameter_state.items(): - if key in _STATE_MOVEMENT_SCRATCH_KEYS: - continue - if key in _STATE_MOVEMENT_CAPTURABLE_KEYS: - raise RuntimeError( - "Gefen state movement cannot migrate capturable scratch state" - ) - if key in _STATE_MOVEMENT_TENSOR_KEYS and torch.is_tensor(value): - if parameter_target is None: - parameter_target = self._state_move_parameter_device(parameter) - staged_parameter_state[key] = stage_tensor(value, parameter_target) - else: - staged_parameter_state[key] = value - staged_state[parameter] = staged_parameter_state - - for cuda_device in sorted( - cuda_devices, key=lambda item: -1 if item.index is None else item.index - ): - torch.cuda.synchronize(cuda_device) - return staged_state, staged_codebook - - @torch.no_grad() - def move_state_(self, device=None) -> None: - """Atomically co-locate authoritative optimizer state with live parameters.""" - - self._assert_finalized_binding_layout(full=True) - try: - reason = self._state_movement_rejection_reason() - except Exception as exc: - raise RuntimeError( - "Gefen atomic state movement could not inspect live optimizer state" - ) from exc - if reason is not None: - raise RuntimeError( - "Gefen atomic state movement is unavailable: {}".format(reason) - ) - if torch.compiler.is_compiling(): - raise RuntimeError("Gefen state movement cannot run during torch.compile") - if torch.cuda.is_available(): - parameters = [ - parameter - for group in self.param_groups - for parameter in group["params"] - ] - try: - capturing = self._state_offload_capturing_on_parameter_device( - parameters - ) - except RuntimeError as exc: - raise RuntimeError( - "Gefen state movement could not inspect CUDA graph-capture state" - ) from exc - if capturing: - raise RuntimeError( - "Gefen state movement cannot run during CUDA graph capture" - ) - - staged_state, staged_codebook = self._stage_state_move(device) - self.__dict__.update( - { - "state": staged_state, - "_gefen_codebook": staged_codebook, - "_gefen_codebook_by_device": {}, - "_gefen_codebook_lut_by_device": {}, - "_gefen_codebook_scope_validated": False, - "_static_mark_sig": None, - "_lr_scalar_cache": None, - "_gefen_state_offload_device": None, - } - ) - self._invalidate_layout_forensics_caches() - @staticmethod def _normalize_param_groups(params): if isinstance(params, torch.Tensor): @@ -3936,10 +3043,6 @@ def add_param_group(self, param_group): stable lowercase name is stored in its per-param state and in the group's ``param_names`` list for introspection. """ - if getattr(self, "_gefen_state_offload_device", None) is not None: - raise RuntimeError( - "Gefen cannot add parameter groups while state offload is active; restore state first" - ) if getattr(self, "_gefen_post_sharding_finalized", False): raise RuntimeError( "Gefen cannot add parameter groups after post_sharding finalization" @@ -5234,6 +4337,19 @@ def _synchronize_codebook_scope_failure(self, error, phase: str) -> None: raise error return self._assert_runtime_codebook_process_group() + self._synchronize_prevalidated_codebook_scope_failure(error, phase, binding) + + @staticmethod + @torch._dynamo.disable + def _synchronize_prevalidated_codebook_scope_failure( + error, phase: str, binding + ) -> None: + """Synchronize through a binding validated before rank-local user code.""" + + if binding is None or len(binding.identity.ordered_members) == 1: + if error is not None: + raise error + return failed = _synchronize_step_failure( error is not None, binding.process_group, @@ -5415,7 +4531,7 @@ def _validate_codebook_scope_operation_header(self, operation: str) -> None: *self._codebook_value_fingerprint(), # Trailing decision bit, deliberately excluded from the equality # check below: collective-free rank-local operations - # (move_state_, offload_state_, staged/native checkpoint loads) + # (post-sharding, staged/native checkpoint loads) # legitimately reset _gefen_codebook_scope_validated on a subset # of members without changing any fingerprinted value. The group # resolves "does any member need re-validation" here so @@ -7139,8 +6255,6 @@ def _canonical_import_live_token(self): self.fused, self.verbose, self._fused_build_ok, - self.state_offload_device, - self.state_offload_poisoned, id(self._gefen_codebook_process_group), self._canonical_value_token(self._serialized_codebook_scope()), id(self._gefen_sharding_manifest), @@ -7319,7 +6433,6 @@ def export_canonical_state(self): """Export an exact-binding, device-neutral local state fragment.""" self._assert_finalized_binding_layout(full=True) - self._assert_state_export_safe() self._assert_canonical_state_outside_cuda_capture("export") if not self._canonical_state_layouts(): raise RuntimeError( @@ -7653,11 +6766,9 @@ 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) @@ -8489,24 +7600,7 @@ def _stage_load_state_dict(self, state_dict): staged._capt_stacks = None staged._load_state_dict_impl(state_dict) - if staged.state_offload_active: - reason = staged._state_offload_rejection_reason( - require_cpu_state=False, - allow_poisoned=True, - require_full_layout=True, - ) - if reason is not None: - raise RuntimeError( - "Gefen could not preserve active state offload while loading: {}".format( - reason - ) - ) - staged.state = staged._stage_all_parameter_state_to_cpu() - staged._gefen_codebook = staged._stage_state_offload_resident_codebook() staged._validate_loaded_native_state() - # A complete successfully validated load is the only operation that can - # re-establish known-good optimizer state after a failed copyback. - staged._gefen_state_offload_poisoned = False return staged def _commit_staged_load_state_dict(self, staged) -> None: @@ -8559,70 +7653,6 @@ def _base_load_state_dict_without_hooks(self, state_dict): self._optimizer_load_state_dict_pre_hooks = pre_hooks self._optimizer_load_state_dict_post_hooks = post_hooks - def _base_load_state_dict_to_offload_cpu(self, state_dict) -> None: - """Apply the base optimizer mapping while keeping parameter state on CPU.""" - - groups = self.param_groups - saved_groups = deepcopy(state_dict["param_groups"]) - if len(groups) != len(saved_groups): - raise ValueError( - "loaded state dict has a different number of parameter groups" - ) - if any( - len(group["params"]) != len(saved_group["params"]) - for group, saved_group in zip(groups, saved_groups) - ): - raise ValueError( - "loaded state dict contains a parameter group that doesn't match the size of optimizer's group" - ) - - id_map = dict( - zip( - chain.from_iterable(group["params"] for group in saved_groups), - chain.from_iterable(group["params"] for group in groups), - ) - ) - - def clone_to_cpu(value): - if torch.is_tensor(value): - return value.to( - device="cpu", - dtype=value.dtype, - non_blocking=False, - copy=True, - memory_format=torch.contiguous_format, - ).detach() - if isinstance(value, dict): - return {key: clone_to_cpu(item) for key, item in value.items()} - if type(value) is list: - return [clone_to_cpu(item) for item in value] - if type(value) is tuple: - return tuple(clone_to_cpu(item) for item in value) - if type(value) is set: - return {clone_to_cpu(item) for item in value} - if type(value) is frozenset: - return frozenset(clone_to_cpu(item) for item in value) - if type(value) is deque: - return deque( - (clone_to_cpu(item) for item in value), maxlen=value.maxlen - ) - return deepcopy(value) - - loaded_state = defaultdict(dict) - for key, value in state_dict["state"].items(): - if key in id_map: - loaded_state[id_map[key]] = clone_to_cpu(value) - else: - loaded_state[key] = value - - param_groups = [] - for live_group, saved_group in zip(groups, saved_groups): - saved_group["params"] = live_group["params"] - if "param_names" in live_group and "param_names" not in saved_group: - saved_group["param_names"] = live_group["param_names"] - param_groups.append(saved_group) - self.__setstate__({"state": loaded_state, "param_groups": param_groups}) - @staticmethod def _validate_rank_local_codebook(codebook, *, required: bool) -> None: if codebook is None: @@ -9400,10 +8430,7 @@ def _load_state_dict_impl(self, state_dict): # and re-aliases lazily). self._capt_invalidate() - if self.state_offload_active: - self._base_load_state_dict_to_offload_cpu(state_dict) - else: - self._base_load_state_dict_without_hooks(state_dict) + self._base_load_state_dict_without_hooks(state_dict) self._gefen_global_step = gefen_global_step # Capturable SR seeds are optimizer-level scratch (a device mirror of # gefen_global_step): drop them so the first post-load SR kernel call @@ -9503,32 +8530,49 @@ def step(self, closure=None): closure feed the first step's codebook learning correctly. The returned loss is passed through. """ - self._assert_state_offload_step_ready() - self._assert_finalized_binding_layout() - self._assert_runtime_codebook_process_group() - self._assert_capturable_if_capturing() - self._assert_codebook_capture_ready() + # Capture and validate the control scope independently of public + # parameter containers. A rank-local mutation made between steps can + # then be reported through the last known-good binding instead of + # stranding peers after a one-sided entry-guard failure. + scope_binding = self._capture_codebook_scope_binding_for_step() loss = None - if closure is not None: - with torch.enable_grad(): - loss = closure() - - self._assert_state_offload_step_ready() - self._assert_finalized_binding_layout() - self._assert_runtime_codebook_process_group() try: + self._assert_finalized_binding_layout() + self._assert_runtime_codebook_process_group() + self._assert_capturable_if_capturing() + self._assert_codebook_capture_ready() + if closure is not None: + with torch.enable_grad(): + loss = closure() + local_preamble_error = None + except Exception as exc: + loss = None + local_preamble_error = exc + if scope_binding is not None: + self._synchronize_prevalidated_codebook_scope_failure( + local_preamble_error, "step preamble", scope_binding + ) + elif local_preamble_error is not None: + raise local_preamble_error + + # The closure can replace a finalized parameter or otherwise invalidate + # the runtime binding. Recheck before the operation header and synchronize + # structural failures before any peer enters a scoped codebook collective. + try: + self._assert_finalized_binding_layout() + self._assert_runtime_codebook_process_group() _assert_optimizer_gradients_structurally_valid(self) local_preflight_error = None except Exception as exc: local_preflight_error = exc - self._validate_codebook_scope_operation_header("step") - if self._gefen_codebook_process_group is not None: - self._synchronize_codebook_scope_failure( - local_preflight_error, "gradient preflight" + if scope_binding is not None: + self._synchronize_prevalidated_codebook_scope_failure( + local_preflight_error, "gradient preflight", scope_binding ) elif local_preflight_error is not None: raise local_preflight_error + self._validate_codebook_scope_operation_header("step") self._ensure_codebook_scope_agreement() # GradScaler invokes native-AMP optimizers even on overflow. Decide @@ -9592,7 +8636,6 @@ def step(self, closure=None): not self._use_fused_gefen_automatic_step() and not self._use_fused_automatic_vmean() and not self.capturable - and not self.state_offload_active ) # Collect parameters that share block geometry + hyperparameters so they @@ -9617,16 +8660,7 @@ def step(self, closure=None): # fall through to the standard path: local-shard row/col # statistics would be wrong under sharding. if self._factored_v_2d and p.ndim == 2 and not hasattr(p, "placements"): - if self.state_offload_active: - self._step_with_offloaded_parameter_state( - self._step_automatic_factored, - group, - name, - p, - grad, - ) - else: - self._step_automatic_factored(group, name, p, grad) + self._step_automatic_factored(group, name, p, grad) continue if ( batch_nonfused @@ -9639,16 +8673,7 @@ def step(self, closure=None): (group, name, p, grad) ) continue - if self.state_offload_active: - self._step_with_offloaded_parameter_state( - self._step_automatic, - group, - name, - p, - grad, - ) - else: - self._step_automatic(group, name, p, grad) + self._step_automatic(group, name, p, grad) for items in nonfused_groups.values(): if len(items) == 1: diff --git a/src/gefen/gefen_muon.py b/src/gefen/gefen_muon.py index 7f68f24..8330b3d 100644 --- a/src/gefen/gefen_muon.py +++ b/src/gefen/gefen_muon.py @@ -777,10 +777,7 @@ def optimizer_contract(self) -> OptimizerContract: for group in self.param_groups if not group.get("normuon", False) ) - try: - canonical_state_layouts = self._canonical_state_layouts() - except Exception: - canonical_state_layouts = frozenset() + canonical_state_layouts = self._canonical_state_layouts() try: from gefen.portable_runtime import _portable_runtime_layouts @@ -814,7 +811,6 @@ 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, - atomic_state_movement=self._atomic_state_movement_supported(), whole_parameter_owner=( self._codebook_scope_ready() and any( @@ -2192,6 +2188,15 @@ def _step_automatic( update = self._compute_muon_update(group, param_name, p, grad, eff_numel) self._apply_muon_update(group, p, update, is_sharded, approx) + @staticmethod + def _state_tensor_device(p: torch.Tensor) -> torch.device: + if hasattr(p, "to_local"): + local = p.to_local() + if hasattr(local, "wait"): + local = local.wait() + return local.device + return p.device + def _distributed_state_items(self, state_dict): saved_ids = [] for saved_group in state_dict["param_groups"]: @@ -3052,24 +3057,31 @@ def _load_state_dict_impl(self, state_dict): @torch.no_grad() def step(self, closure=None): - self._assert_finalized_binding_layout() - self._assert_runtime_codebook_process_group() - if self._has_unscoped_whole_owner_bindings(): - raise RuntimeError( - "GefenMuon whole-parameter owner stepping requires the separate " - "explicit process-group codebook scope" - ) - loss = None + scope_binding = self._capture_codebook_scope_binding_for_step() hybrid_preflight_complete = bool( getattr(self, "_gefen_hybrid_precollective_preflight", False) ) - if not hybrid_preflight_complete and ( - self._gefen_codebook_process_group is not None - ): + if hybrid_preflight_complete or scope_binding is None: + self._assert_finalized_binding_layout() + self._assert_runtime_codebook_process_group() + if self._has_unscoped_whole_owner_bindings(): + raise RuntimeError( + "GefenMuon whole-parameter owner stepping requires the separate " + "explicit process-group codebook scope" + ) + loss = None + if not hybrid_preflight_complete and scope_binding is not None: # The explicit convention binding is the exclusive control scope: # synchronize on it before any operation header or mesh collective, # never in addition to the mainline DTensor-derived scope. try: + self._assert_finalized_binding_layout() + self._assert_runtime_codebook_process_group() + if self._has_unscoped_whole_owner_bindings(): + raise RuntimeError( + "GefenMuon whole-parameter owner stepping requires the separate " + "explicit process-group codebook scope" + ) self._assert_capturable_if_capturing() self._assert_codebook_capture_ready() if closure is not None: @@ -3079,23 +3091,23 @@ def step(self, closure=None): except Exception as exc: loss = None local_preamble_error = exc - self._synchronize_codebook_scope_failure( - local_preamble_error, "step preamble" + self._synchronize_prevalidated_codebook_scope_failure( + local_preamble_error, "step preamble", scope_binding ) - self._assert_finalized_binding_layout() - self._assert_runtime_codebook_process_group() try: + self._assert_finalized_binding_layout() + self._assert_runtime_codebook_process_group() _assert_optimizer_gradients_structurally_valid( self, require_2d_params=True ) local_preflight_error = None except Exception as exc: local_preflight_error = exc - self._validate_codebook_scope_operation_header("step") - self._synchronize_codebook_scope_failure( - local_preflight_error, "gradient preflight" + self._synchronize_prevalidated_codebook_scope_failure( + local_preflight_error, "gradient preflight", scope_binding ) + self._validate_codebook_scope_operation_header("step") self._ensure_codebook_scope_agreement() if not self._prepare_scoped_amp_optimizer_step(): diff --git a/src/gefen/hybrid.py b/src/gefen/hybrid.py index 1530cc5..04f7db2 100644 --- a/src/gefen/hybrid.py +++ b/src/gefen/hybrid.py @@ -1213,6 +1213,15 @@ def _synchronize_codebook_scope_failure(self, error, phase: str) -> None: return self._subopts[0]._synchronize_codebook_scope_failure(error, phase) + def _synchronize_prevalidated_codebook_scope_failure( + self, error, phase: str, binding + ) -> None: + """Synchronize through the shared binding validated at step entry.""" + + self._subopts[0]._synchronize_prevalidated_codebook_scope_failure( + error, phase, binding + ) + def _prepare_scoped_amp_optimizer_step(self) -> bool: # GradScaler attaches found_inf/grad_scale to the composite, never to # the children, so the children's own scoped AMP gates cannot see an @@ -1250,8 +1259,13 @@ def _assert_capturable_devices_if_capturing(self) -> None: ) def step(self, closure=None): - self._assert_finalized_binding_layout() binding = self._hybrid_codebook_process_group + if binding is None: + self._assert_finalized_binding_layout() + else: + if not isinstance(binding, CodebookProcessGroupBinding): + raise TypeError("the runtime codebook process-group binding is invalid") + self._subopts[0]._validate_codebook_runtime_binding(binding) process_groups = ( self.muon._step_failure_process_groups() if self.muon is not None and binding is None @@ -1271,12 +1285,14 @@ def step(self, closure=None): # and synchronize it on the one shared codebook binding BEFORE any member # enters the composite preflight synchronize or a child's scoped step # collectives, so it raises on every member together instead of stranding - # peers. The finalized-layout guard above stays local: it establishes the - # binding used to synchronize. + # peers. A present runtime binding is validated independently first, so + # the finalized-layout entry guard can join this synchronized phase. args = (self, closure) if closure is not None else (self,) kwargs = {} loss = None try: + if binding is not None: + self._assert_finalized_binding_layout() self._assert_capturable_devices_if_capturing() for pre_hook in self._optimizer_step_pre_hooks.values(): result = pre_hook(self, args, kwargs) @@ -1300,11 +1316,11 @@ def step(self, closure=None): local_preamble_error = exc if binding is not None: - self._synchronize_codebook_scope_failure( - local_preamble_error, "step preamble" + self._synchronize_prevalidated_codebook_scope_failure( + local_preamble_error, "step preamble", binding ) - self._assert_finalized_binding_layout() try: + self._assert_finalized_binding_layout() for child in self._subopts: _assert_optimizer_gradients_structurally_valid( child, require_2d_params=child is self.muon @@ -1312,8 +1328,8 @@ def step(self, closure=None): local_preflight_error = None except Exception as exc: local_preflight_error = exc - self._synchronize_codebook_scope_failure( - local_preflight_error, "gradient preflight" + self._synchronize_prevalidated_codebook_scope_failure( + local_preflight_error, "gradient preflight", binding ) should_step = self._prepare_scoped_amp_optimizer_step() else: diff --git a/src/gefen/portable_runtime.py b/src/gefen/portable_runtime.py index 84f1843..353fdc0 100644 --- a/src/gefen/portable_runtime.py +++ b/src/gefen/portable_runtime.py @@ -652,10 +652,6 @@ def _validate_optimizer_shell( raise RuntimeError("portable state requires inactive capturable stacks") if optimizer._gefen_global_step_by_device or optimizer._sr_seed_by_device: raise RuntimeError("portable state does not support device-authoritative counters") - if optimizer.state_offload_poisoned: - raise RuntimeError("portable state does not support poisoned optimizer state") - if optimizer.state_offload_active: - raise RuntimeError("portable state does not support active optimizer-state offload") if torch.compiler.is_compiling(): raise RuntimeError("portable state cannot run while compiling") if torch.cuda.is_available() and torch.cuda.is_current_stream_capturing(): @@ -853,8 +849,6 @@ def live_group_options(group): optimizer.fused, optimizer.verbose, optimizer._fused_build_ok, - optimizer.state_offload_device, - optimizer.state_offload_poisoned, id(optimizer._gefen_codebook_process_group), _portable_value_token(optimizer._serialized_codebook_scope()), id(optimizer._gefen_sharding_manifest), diff --git a/tests/test_canonical_state_cpu.py b/tests/test_canonical_state_cpu.py index cd1cfe7..6e98c0c 100644 --- a/tests/test_canonical_state_cpu.py +++ b/tests/test_canonical_state_cpu.py @@ -311,10 +311,6 @@ def test_initialized_import_maps_by_fqn_across_order_and_group_boundaries(): assert torch.equal(target_second, source_second) assert torch.equal(target._gefen_codebook, source._gefen_codebook) - target.move_state_("cpu") - assert target._gefen_logical_slots is logical_slots - assert target.optimizer_contract().capabilities.canonical_state_io - with pytest.raises(RuntimeError, match="already consumed"): target.commit_canonical_state_import(prepared) diff --git a/tests/test_codebook_scope_distributed.py b/tests/test_codebook_scope_distributed.py index 36a8160..b55c458 100644 --- a/tests/test_codebook_scope_distributed.py +++ b/tests/test_codebook_scope_distributed.py @@ -1247,11 +1247,10 @@ def test_nccl_scope_uses_explicit_collective_device_with_empty_nonowner(): os.unlink(init_file) -def _closure_preamble_worker(rank, world, init_file, queue): - # GefenMuon.step runs the closure BEFORE any scope synchronization. A - # rank-local closure failure must raise on every scope member together - # instead of leaving the failing rank to exit step() while the peer enters - # the scoped operation-header / synchronization collectives and hangs. +def _closure_preamble_worker( + rank, world, init_file, queue, optimizer_kind, failure_mode +): + # Optimizer.step runs the closure before any later scope collectives. A rank-local closure failure must raise on every scope member together instead of leaving the failing rank to exit step() while the peer enters the scoped operation-header / synchronization collectives and hangs. try: dist.init_process_group( "gloo", @@ -1265,7 +1264,14 @@ def _closure_preamble_worker(rank, world, init_file, queue): runtime_group = dist.group.WORLD matrix = torch.nn.Parameter(torch.zeros(2, 2)) - optimizer = GefenMuon([("matrix", matrix)], fused=False) + if optimizer_kind == "muon": + optimizer = GefenMuon([("matrix", matrix)], fused=False) + elif optimizer_kind == "gefen": + optimizer = Gefen( + [("matrix", matrix)], fused=False, factored_v_2d=False + ) + else: + raise AssertionError("unknown closure-preamble optimizer kind") identity = ParameterIdentity("Matrix", (2, 2)) records = tuple(_replicated(identity, group, member) for member in members) _finalize( @@ -1276,10 +1282,18 @@ def _closure_preamble_worker(rank, world, init_file, queue): _binding(group, rank, runtime_group), ) matrix.grad = torch.tensor([[1.0, -2.0], [3.0, -4.0]]) + rogue = torch.nn.Parameter(torch.ones(2, 2)) + rogue_before = rogue.detach().clone() + if rank == 0 and failure_mode == "replace_before": + optimizer.param_groups[0]["params"][0] = rogue + rogue.grad = torch.ones_like(rogue) def closure(): - if rank == 0: + if rank == 0 and failure_mode == "raise": raise RuntimeError("closure boom on rank:0") + if rank == 0 and failure_mode == "replace": + optimizer.param_groups[0]["params"][0] = rogue + rogue.grad = torch.ones_like(rogue) return torch.tensor(1.0) try: @@ -1293,6 +1307,7 @@ def closure(): and optimizer._gefen_codebook is None and optimizer.state[matrix] == {"name": "matrix"} and torch.equal(matrix, torch.zeros(2, 2)) + and torch.equal(rogue, rogue_before) ) queue.put({"rank": rank, "message": message, "untouched": untouched}) except Exception as exc: @@ -1302,7 +1317,9 @@ def closure(): dist.destroy_process_group() -def _run_closure_preamble_workers(world=2): +def _run_closure_preamble_workers( + optimizer_kind="muon", failure_mode="raise", world=2 +): context = mp.get_context("spawn") queue = context.Queue() fd, init_file = tempfile.mkstemp(prefix="gefen-codebook-closure-preamble-") @@ -1311,7 +1328,7 @@ def _run_closure_preamble_workers(world=2): processes = [ context.Process( target=_closure_preamble_worker, - args=(rank, world, init_file, queue), + args=(rank, world, init_file, queue, optimizer_kind, failure_mode), ) for rank in range(world) ] @@ -1354,3 +1371,51 @@ def test_scoped_step_closure_failure_raises_symmetrically_across_the_scope(): assert "closure boom on rank:0" in results[0]["message"] assert "step preamble failed on another process-group member" in results[1]["message"] assert all(item["untouched"] for item in results), results + + +@pytest.mark.skipif( + not dist.is_available() or not dist.is_gloo_available(), + reason="plain Gefen scoped closure-preamble coverage requires Gloo", +) +def test_plain_gefen_scoped_step_closure_failure_raises_symmetrically_across_the_scope(): + results = _run_closure_preamble_workers("gefen") + assert len(results) == 2, results + assert all("error" not in item for item in results), results + assert results[0]["message"] is not None and results[1]["message"] is not None, results + assert "closure boom on rank:0" in results[0]["message"] + assert "step preamble failed on another process-group member" in results[1]["message"] + assert all(item["untouched"] for item in results), results + + +@pytest.mark.skipif( + not dist.is_available() or not dist.is_gloo_available(), + reason="scoped closure layout-mutation coverage requires Gloo", +) +@pytest.mark.parametrize("optimizer_kind", ["gefen", "muon"]) +def test_scoped_step_closure_layout_mutation_raises_symmetrically_across_the_scope( + optimizer_kind, +): + results = _run_closure_preamble_workers(optimizer_kind, "replace") + assert len(results) == 2, results + assert all("error" not in item for item in results), results + assert results[0]["message"] is not None and results[1]["message"] is not None, results + assert "changed outside post_sharding" in results[0]["message"] + assert "gradient preflight failed on another process-group member" in results[1]["message"] + assert all(item["untouched"] for item in results), results + + +@pytest.mark.skipif( + not dist.is_available() or not dist.is_gloo_available(), + reason="scoped step-entry layout-mutation coverage requires Gloo", +) +@pytest.mark.parametrize("optimizer_kind", ["gefen", "muon"]) +def test_scoped_step_entry_layout_mutation_raises_symmetrically_across_the_scope( + optimizer_kind, +): + results = _run_closure_preamble_workers(optimizer_kind, "replace_before") + assert len(results) == 2, results + assert all("error" not in item for item in results), results + assert results[0]["message"] is not None and results[1]["message"] is not None, results + assert "changed outside post_sharding" in results[0]["message"] + assert "step preamble failed on another process-group member" in results[1]["message"] + assert all(item["untouched"] for item in results), results diff --git a/tests/test_hybrid_scoped_failure_protocol.py b/tests/test_hybrid_scoped_failure_protocol.py index dc0f73e..68e59a6 100644 --- a/tests/test_hybrid_scoped_failure_protocol.py +++ b/tests/test_hybrid_scoped_failure_protocol.py @@ -378,6 +378,58 @@ def closure(): } +def _closure_layout_divergent_result(rank, group): + optimizer, muon_parameter, backup_parameter, backup_shard = _make_scoped_hybrid( + rank, group + ) + _set_local_gradients(muon_parameter, backup_parameter, backup_shard) + rogue = torch.nn.Parameter(torch.ones_like(backup_parameter)) + rogue_before = rogue.detach().clone() + + def closure(): + if rank == 0: + optimizer.backup.param_groups[0]["params"][0] = rogue + rogue.grad = torch.ones_like(rogue) + return torch.tensor(1.0) + + try: + optimizer.step(closure) + message = None + except RuntimeError as exc: + message = str(exc) + return { + "message": message, + "untouched": ( + _untouched(optimizer, muon_parameter, backup_parameter) + and torch.equal(rogue, rogue_before) + ), + } + + +def _step_entry_layout_divergent_result(rank, group): + optimizer, muon_parameter, backup_parameter, backup_shard = _make_scoped_hybrid( + rank, group + ) + _set_local_gradients(muon_parameter, backup_parameter, backup_shard) + rogue = torch.nn.Parameter(torch.ones_like(backup_parameter)) + rogue_before = rogue.detach().clone() + if rank == 0: + optimizer.backup.param_groups[0]["params"][0] = rogue + rogue.grad = torch.ones_like(rogue) + try: + optimizer.step() + message = None + except RuntimeError as exc: + message = str(exc) + return { + "message": message, + "untouched": ( + _untouched(optimizer, muon_parameter, backup_parameter) + and torch.equal(rogue, rogue_before) + ), + } + + def _distributed_worker(rank, init_file, result_queue): try: dist.init_process_group( @@ -398,6 +450,10 @@ def _distributed_worker(rank, init_file, result_queue): dist.barrier() closure_divergent = _closure_divergent_result(rank, group) dist.barrier() + closure_layout_divergent = _closure_layout_divergent_result(rank, group) + dist.barrier() + step_entry_layout_divergent = _step_entry_layout_divergent_result(rank, group) + dist.barrier() result_queue.put( { "rank": rank, @@ -406,6 +462,8 @@ def _distributed_worker(rank, init_file, result_queue): "amp_finite": amp_finite, "preflight": preflight, "closure_divergent": closure_divergent, + "closure_layout_divergent": closure_layout_divergent, + "step_entry_layout_divergent": step_entry_layout_divergent, } ) except BaseException: @@ -515,6 +573,28 @@ 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 + # A closure can silently replace one finalized child parameter. Synchronize + # that post-closure layout failure through the binding captured at step entry, + # before either member revalidates the now rank-divergent live binding. + closure_layout_divergent = [ + result["closure_layout_divergent"] for result in results + ] + assert closure_layout_divergent[0]["message"] is not None and closure_layout_divergent[1]["message"] is not None, closure_layout_divergent + assert "changed outside post_sharding" in closure_layout_divergent[0]["message"] + assert "gradient preflight failed on another process-group member" in closure_layout_divergent[1]["message"] + assert all(item["untouched"] for item in closure_layout_divergent), closure_layout_divergent + + # A mutation made between steps fails the first finalized-layout guard. The + # explicit binding is captured independently so that guard joins the same + # preamble vote instead of failing on only one member. + step_entry_layout_divergent = [ + result["step_entry_layout_divergent"] for result in results + ] + assert step_entry_layout_divergent[0]["message"] is not None and step_entry_layout_divergent[1]["message"] is not None, step_entry_layout_divergent + assert "changed outside post_sharding" in step_entry_layout_divergent[0]["message"] + assert "step preamble failed on another process-group member" in step_entry_layout_divergent[1]["message"] + assert all(item["untouched"] for item in step_entry_layout_divergent), step_entry_layout_divergent + def _make_plain_hybrid(): matrix = torch.nn.Parameter(_muon_initial().clone()) diff --git a/tests/test_layout_guard_cost.py b/tests/test_layout_guard_cost.py index 00cc207..a2b17c8 100644 --- a/tests/test_layout_guard_cost.py +++ b/tests/test_layout_guard_cost.py @@ -1,11 +1,11 @@ -"""Per-step layout/offload guard cost: fast-path tokens, dedupe, boundaries. +"""Per-step layout guard cost: fast-path tokens, dedupe, and boundaries. The finalized-layout guards run on every ``step()``; these tests pin the contract that steady-state steps reuse one cached forensic verdict (an O(local params) identity token check) while every legitimate mutating API and -every boundary operation (checkpoint prepare/commit, rebinding, state -movement/offload, contract readiness) still runs the complete O(params x -world) forensic rebuild. Detection-before-mutation is preserved: anything a +every boundary operation (checkpoint prepare/commit, rebinding, and contract +readiness) still runs the complete O(params x world) forensic rebuild. +Detection-before-mutation is preserved: anything a closure or adapter can corrupt through the public optimizer containers still raises at the step guard itself. """ @@ -259,83 +259,24 @@ def test_private_registry_inplace_tamper_is_detected_at_the_next_boundary(): optimizer.state_dict() -def test_mutating_apis_bump_the_layout_version_and_force_revalidation(monkeypatch): +def test_load_state_dict_bumps_the_layout_version_and_forces_revalidation(monkeypatch): optimizer, parameters = _finalized_replicated_optimizer() _step_with_grads(optimizer, parameters) checkpoint = optimizer.state_dict() calls = _count_full_layout_passes(monkeypatch) - version = optimizer._gefen_layout_version - optimizer.move_state_() - assert optimizer._gefen_layout_version == version + 1 - assert optimizer._gefen_layout_forensics_verdict is None - - _step_with_grads(optimizer, parameters) - after_move = calls["count"] - assert after_move >= 1 - _step_with_grads(optimizer, parameters) - assert calls["count"] == after_move # steady again - version = optimizer._gefen_layout_version optimizer.load_state_dict(checkpoint) assert optimizer._gefen_layout_version > version assert optimizer._gefen_layout_forensics_verdict is None - -def test_state_offload_step_scan_runs_on_every_step(monkeypatch): - # The offload readiness scan is intentionally NOT cached: the per-parameter - # offloaded state tensors are legitimately replaced on every step, so a - # cached verdict cannot represent them, and a token-preserving in-place - # corruption of a later parameter would otherwise slip past step entry and - # only be caught mid-step, after earlier parameters were already mutated. - # The scan is O(local params) and must run before any parameter is staged. - optimizer, parameters = _finalized_replicated_optimizer() + after_load = calls["count"] _step_with_grads(optimizer, parameters) - - calls = {"count": 0} - - def counted(self, *, require_cpu_state, allow_poisoned=False): - calls["count"] += 1 - return None - - monkeypatch.setattr(Gefen, "_state_offload_rejection_reason", counted) - monkeypatch.setattr( - Gefen, - "state_offload_active", - property(lambda self: True), - ) - - optimizer._assert_state_offload_step_ready() - optimizer._assert_state_offload_step_ready() - optimizer._assert_state_offload_step_ready() - assert calls["count"] == 3 - - optimizer._gefen_state_offload_poisoned = True - with pytest.raises(RuntimeError, match="poisoned"): - optimizer._assert_state_offload_step_ready() - - -def test_offload_scan_re_rejects_on_every_step(monkeypatch): - optimizer, parameters = _finalized_replicated_optimizer() - - calls = {"count": 0} - - def counted(self, *, require_cpu_state, allow_poisoned=False): - calls["count"] += 1 - return "authoritative offloaded tensors must be tight CPU tensors" - - monkeypatch.setattr(Gefen, "_state_offload_rejection_reason", counted) - monkeypatch.setattr( - Gefen, - "state_offload_active", - property(lambda self: True), - ) - - for _ in range(2): - with pytest.raises(RuntimeError, match="cannot step"): - optimizer._assert_state_offload_step_ready() - assert calls["count"] == 2 + assert calls["count"] > after_load + after_revalidation = calls["count"] + _step_with_grads(optimizer, parameters) + assert calls["count"] == after_revalidation # steady again def test_forensics_caches_stay_out_of_the_public_attribute_namespace(): @@ -359,10 +300,8 @@ def test_warm_step_guards_are_far_cheaper_than_one_forensic_pass(): iterations = 50 start = time.perf_counter() for _ in range(iterations): - optimizer._assert_state_offload_step_ready() optimizer._assert_finalized_binding_layout() optimizer._assert_runtime_codebook_process_group() - optimizer._assert_state_offload_step_ready() optimizer._assert_finalized_binding_layout() optimizer._assert_runtime_codebook_process_group() warm_guard = (time.perf_counter() - start) / iterations diff --git a/tests/test_optimizer_contracts.py b/tests/test_optimizer_contracts.py index 5be414a..4412acc 100644 --- a/tests/test_optimizer_contracts.py +++ b/tests/test_optimizer_contracts.py @@ -1,8 +1,7 @@ """CPU coverage for public optimizer capability and state-layout contracts.""" import copy -from collections import defaultdict, OrderedDict -from dataclasses import FrozenInstanceError, dataclass, replace +from dataclasses import FrozenInstanceError, replace import pytest import torch @@ -25,8 +24,6 @@ StateField, StateGeometry, StateKeyMatch, - StateMovementProvider, - StateOffloadProvider, StateScope, StateVariant, TopologyChange, @@ -133,7 +130,6 @@ def test_plain_contract_matches_live_persistent_state(factored_v_2d): factored_v_2d=factored_v_2d, ) assert isinstance(optimizer, OptimizerContractProvider) - assert isinstance(optimizer, StateMovementProvider) contract = optimizer.optimizer_contract() assert contract.schema_version == CONTRACT_SCHEMA_VERSION @@ -172,7 +168,7 @@ def test_plain_contract_matches_live_persistent_state(factored_v_2d): assert contract.capabilities.shard_rebinding assert contract.capabilities.post_sharding assert not contract.capabilities.canonical_state_io - assert contract.capabilities.atomic_state_movement + assert not contract.capabilities.atomic_state_movement assert not contract.capabilities.state_offload assert Precision.FLOAT64 in contract.capabilities.precisions flattened = _training_support(contract, ParameterLayout.FLATTENED_ELEMENT_SHARD) @@ -348,13 +344,12 @@ def test_muon_contract_separates_mode_topology_and_state_extent( sharded_mode=sharded_mode, normuon=normuon, ) - assert isinstance(optimizer, StateMovementProvider) contract = optimizer.optimizer_contract() assert contract.implementation == "gefen.GefenMuon" assert contract.capabilities.supported_parameter_ranks == (2,) assert contract.capabilities.explicit_process_group_codebook_scope - assert contract.capabilities.atomic_state_movement + assert not contract.capabilities.atomic_state_movement assert not contract.capabilities.state_offload native = next( item @@ -474,7 +469,6 @@ def test_hybrid_contract_preserves_child_namespaces(backup_optimizer): fused=False, backup_optimizer=backup_optimizer, ) - assert not isinstance(optimizer, StateMovementProvider) contract = optimizer.optimizer_contract() assert contract.implementation == "gefen.GefenMuonHybrid" @@ -530,86 +524,6 @@ def test_muon_contract_keeps_mixed_normuon_variants_in_one_mode(): assert "quantized_normuon_replicated_exact" in variants -@pytest.mark.parametrize("implementation", ["gefen", "muon"]) -def test_capturable_contract_declines_atomic_state_movement(implementation): - shape = (4,) if implementation == "gefen" else (2, 2) - parameter = torch.nn.Parameter(torch.ones(shape)) - optimizer_type = Gefen if implementation == "gefen" else GefenMuon - optimizer = optimizer_type( - [("parameter", parameter)], - fused=False, - capturable=True, - ) - - capabilities = optimizer.optimizer_contract().capabilities - assert not capabilities.atomic_state_movement - assert not capabilities.state_offload - - -@pytest.mark.parametrize("implementation", ["gefen", "muon"]) -def test_undeclared_tensor_state_disables_atomic_state_movement(implementation): - shape = (4,) if implementation == "gefen" else (2, 2) - parameter = torch.nn.Parameter(torch.ones(shape)) - optimizer_type = Gefen if implementation == "gefen" else GefenMuon - optimizer = optimizer_type([("parameter", parameter)], fused=False) - assert optimizer.optimizer_contract().capabilities.atomic_state_movement - - optimizer.state[parameter]["undeclared_tensor"] = torch.ones(1) - capabilities = optimizer.optimizer_contract().capabilities - assert not capabilities.atomic_state_movement - assert not capabilities.state_offload - - -@pytest.mark.parametrize("implementation", ["gefen", "muon"]) -@pytest.mark.parametrize("contains_tensor", [False, True]) -def test_opaque_extension_state_disables_atomic_state_movement( - implementation, contains_tensor -): - @dataclass - class ExtensionState: - payload: object - - shape = (4,) if implementation == "gefen" else (2, 2) - parameter = torch.nn.Parameter(torch.ones(shape)) - optimizer_type = Gefen if implementation == "gefen" else GefenMuon - optimizer = optimizer_type([("parameter", parameter)], fused=False) - payload = torch.ones(1) if contains_tensor else "tensor-free" - optimizer.state[parameter]["extension"] = ExtensionState(payload) - - capabilities = optimizer.optimizer_contract().capabilities - assert not capabilities.atomic_state_movement - assert not capabilities.state_offload - - -@pytest.mark.parametrize("implementation", ["gefen", "muon"]) -@pytest.mark.parametrize("container_type", ["defaultdict", "ordered_dict"]) -def test_extension_mapping_with_hidden_tensor_disables_atomic_state_movement( - implementation, container_type -): - class TensorFactory: - def __init__(self, tensor): - self.tensor = tensor - - def __call__(self): - return self.tensor - - shape = (4,) if implementation == "gefen" else (2, 2) - parameter = torch.nn.Parameter(torch.ones(shape)) - optimizer_type = Gefen if implementation == "gefen" else GefenMuon - optimizer = optimizer_type([("parameter", parameter)], fused=False) - hidden_tensor = torch.ones(1) - if container_type == "defaultdict": - extension = defaultdict(TensorFactory(hidden_tensor)) - else: - extension = OrderedDict() - extension.hidden_tensor = hidden_tensor - optimizer.state[parameter]["extension"] = extension - - capabilities = optimizer.optimizer_contract().capabilities - assert not capabilities.atomic_state_movement - assert not capabilities.state_offload - - def test_muon_mixed_approx_distributed_checkpoint_is_same_topology_only(): first = torch.nn.Parameter(torch.ones(4, 4)) second = torch.nn.Parameter(torch.ones(4, 4)) @@ -748,8 +662,6 @@ def test_all_public_contract_exports_resolve(): from gefen import contracts assert all(getattr(gefen, name) is not None for name in contracts.__all__) - assert gefen.StateMovementProvider is StateMovementProvider - assert gefen.StateOffloadProvider is StateOffloadProvider def test_portable_global_transport_is_defined_but_not_claimed_before_integration(): diff --git a/tests/test_scoped_collective_agreement_fixes.py b/tests/test_scoped_collective_agreement_fixes.py index e69491a..51a27a4 100644 --- a/tests/test_scoped_collective_agreement_fixes.py +++ b/tests/test_scoped_collective_agreement_fixes.py @@ -205,8 +205,9 @@ def capture(reuse_existing_periods=False): def _revalidation_decision_worker(rank, init_file, queue): """A rank-local scope-validated reset must re-validate on every member. - move_state_ is collective-free and resets _gefen_codebook_scope_validated - on the member that called it only, while every fingerprint the step + A native checkpoint reload is rank-local and resets + _gefen_codebook_scope_validated on the member that loaded it only, while + every fingerprint the step header exchanges stays identical. The next step must make the SAME early-return decision inside _ensure_codebook_scope_agreement on every member. @@ -232,11 +233,12 @@ def _revalidation_decision_worker(rank, init_file, queue): optimizer.step() validated_after_first = bool(optimizer._gefen_codebook_scope_validated) - # Advertised collective-free state movement on ONE member only; every - # state value stays bit-identical, so the step header fingerprints - # still agree across members. + # Reload native state on ONE member only; every state value stays + # bit-identical, so the step header fingerprints still agree across + # members. + checkpoint = optimizer.state_dict() if rank == 0: - optimizer.move_state_(torch.device("cpu")) + optimizer.load_state_dict(checkpoint) reset_asymmetric = ( not optimizer._gefen_codebook_scope_validated if rank == 0 diff --git a/tests/test_state_movement.py b/tests/test_state_movement.py deleted file mode 100644 index a290acc..0000000 --- a/tests/test_state_movement.py +++ /dev/null @@ -1,889 +0,0 @@ -"""Atomic optimizer-state movement coverage for Gefen and GefenMuon.""" - -import copy -from collections import defaultdict, OrderedDict -import warnings - -import pytest -import torch - -from gefen import Gefen, GefenMuon, ParameterIdentity - - -_KINDS = ("block", "factored", "muon", "normuon") -_MOVABLE_STATE_KEYS = frozenset( - { - "step", - "m_codebook", - "m_magnitude", - "vmean", - "vmean_step", - "v_row", - "v_col", - "factored_step", - "normuon_v", - "normuon_step", - } -) -_SCRATCH_KEYS = frozenset({"stepsize", "_h_buf"}) -_CACHE_ATTRS = ( - "_gefen_codebook_by_device", - "_gefen_codebook_lut_by_device", - "_sr_seed_by_device", - "_gefen_global_step_by_device", -) - - -def _build_initialized(kind): - shape = (8,) if kind == "block" else (2, 4) - values = torch.linspace(-0.4, 0.7, 8, dtype=torch.float32).reshape(shape) - parameter = torch.nn.Parameter(values.clone()) - tensor_lr = torch.tensor(2.0e-3, dtype=torch.float32) - group_metadata = {"labels": ["preserve", kind]} - group = { - "params": [("layer.weight", parameter)], - "lr": tensor_lr, - "movement_metadata": group_metadata, - } - if kind == "block": - optimizer = Gefen( - [group], - lr=tensor_lr, - fused=False, - factored_v_2d=False, - ) - elif kind == "factored": - optimizer = Gefen( - [group], - lr=tensor_lr, - fused=False, - factored_v_2d=True, - ) - elif kind in ("muon", "normuon"): - optimizer = GefenMuon( - [group], - lr=tensor_lr, - weight_decay=0.0, - fused=False, - ns_steps=1, - normuon=kind == "normuon", - ) - else: - raise AssertionError("unknown optimizer kind: {}".format(kind)) - - optimizer._resolve_automatic_period = lambda *args, **kwargs: 4 - optimizer._predict_period_from_grad_sq = lambda *args, **kwargs: 4 - parameter.grad = torch.linspace(-1.25, 0.75, 8).reshape_as(parameter) - optimizer.step() - return optimizer, parameter, tensor_lr, group_metadata - - -def _oversized_copy(tensor): - flat_size = tensor.numel() - backing = torch.empty( - flat_size + 11, - dtype=tensor.dtype, - device=tensor.device, - ) - result = backing.narrow(0, 5, flat_size).view(tensor.shape) - result.copy_(tensor) - assert result.untyped_storage().nbytes() > flat_size * tensor.element_size() - return result - - -def _make_persistent_state_oversized(optimizer, parameter): - state = optimizer.state[parameter] - for key, value in tuple(state.items()): - if key in _MOVABLE_STATE_KEYS and torch.is_tensor(value): - state[key] = _oversized_copy(value) - optimizer._gefen_codebook = _oversized_copy(optimizer._gefen_codebook) - - -def _persistent_tensor_snapshot(optimizer, parameter): - state = optimizer.state[parameter] - result = { - key: (value, value.detach().clone()) - for key, value in state.items() - if key in _MOVABLE_STATE_KEYS and torch.is_tensor(value) - } - result["_gefen_codebook"] = ( - optimizer._gefen_codebook, - optimizer._gefen_codebook.detach().clone(), - ) - return result - - -def _assert_fresh_tight_copy(actual, old, expected, device): - assert actual is not old - assert actual.device == device - assert actual.dtype == expected.dtype - assert actual.shape == expected.shape - assert actual.layout == torch.strided - assert actual.is_contiguous() - assert actual.storage_offset() == 0 - assert actual.untyped_storage().nbytes() == actual.numel() * actual.element_size() - torch.testing.assert_close(actual, expected, rtol=0, atol=0, equal_nan=True) - - -def _seed_discardable_runtime_state(optimizer, parameter, tensor_lr): - state = optimizer.state[parameter] - state["stepsize"] = _oversized_copy(torch.tensor([91.0])) - state["_h_buf"] = _oversized_copy(torch.tensor([92.0])) - carrier = _oversized_copy(torch.tensor([7, 11, 13], dtype=torch.int64)) - state["_gefen_rank_local_payload_0"] = carrier - state_metadata = { - "labels": ["persistent", "metadata"], - "device": torch.device("cpu"), - "dtype": torch.float32, - "layout": torch.strided, - "memory_format": torch.contiguous_format, - "shape": torch.Size((2, 3)), - "complex": 1.0 + 2.0j, - "bytes": b"metadata", - "optional": None, - } - state["movement_metadata"] = state_metadata - - cpu = torch.device("cpu") - optimizer._gefen_codebook_by_device[cpu] = optimizer._gefen_codebook - optimizer._gefen_codebook_lut_by_device[cpu] = torch.tensor([19.0]) - optimizer._gefen_codebook_scope_validated = True - optimizer._static_mark_sig = ("stale",) - optimizer._lr_scalar_cache = ( - tensor_lr, - tensor_lr._version, - float(tensor_lr.item()), - ) - return carrier, state_metadata - - -def _tensor_payload(tensor): - if tensor.device.type == "meta": - return None - if tensor.is_nested: - return tuple(item.detach().clone() for item in tensor.unbind()) - return tensor.detach().clone() - - -def _snapshot_tree(value): - if torch.is_tensor(value): - return ( - "tensor", - value, - value._version, - value.dtype, - value.device, - value.layout, - _tensor_payload(value), - ) - if isinstance(value, dict): - return ( - "dict", - value, - tuple((key, _snapshot_tree(item)) for key, item in value.items()), - ) - if isinstance(value, list): - return ("list", value, tuple(_snapshot_tree(item) for item in value)) - if isinstance(value, tuple): - return ("tuple", value, tuple(_snapshot_tree(item) for item in value)) - return ("leaf", value, copy.deepcopy(value)) - - -def _assert_tensor_payload(actual, payload): - if payload is None: - return - if actual.is_nested: - actual_items = actual.unbind() - assert len(actual_items) == len(payload) - for item, expected in zip(actual_items, payload): - torch.testing.assert_close(item, expected, rtol=0, atol=0, equal_nan=True) - return - torch.testing.assert_close(actual, payload, rtol=0, atol=0, equal_nan=True) - - -def _assert_tree_exact(actual, snapshot): - kind = snapshot[0] - if kind == "tensor": - _, reference, version, dtype, device, layout, payload = snapshot - assert actual is reference - assert actual._version == version - assert actual.dtype == dtype - assert actual.device == device - assert actual.layout == layout - _assert_tensor_payload(actual, payload) - return - if kind == "dict": - _, reference, entries = snapshot - assert actual is reference - assert tuple(actual) == tuple(key for key, _ in entries) - for key, child in entries: - _assert_tree_exact(actual[key], child) - return - if kind in ("list", "tuple"): - _, reference, entries = snapshot - assert actual is reference - assert len(actual) == len(entries) - for item, child in zip(actual, entries): - _assert_tree_exact(item, child) - return - _, reference, expected = snapshot - assert actual is reference - assert actual == expected - - -def _snapshot_exact_optimizer(optimizer): - parameters = tuple( - parameter - for group in optimizer.param_groups - for parameter in group["params"] - ) - return { - "top": optimizer.__dict__.copy(), - "state": _snapshot_tree(optimizer.state), - "groups": _snapshot_tree(optimizer.param_groups), - "defaults": _snapshot_tree(optimizer.defaults), - "param_names": _snapshot_tree(optimizer._param_names), - "codebook": _snapshot_tree(optimizer._gefen_codebook), - "caches": { - name: _snapshot_tree(getattr(optimizer, name)) for name in _CACHE_ATTRS - }, - "grads": tuple( - (parameter, _snapshot_tree(parameter.grad)) for parameter in parameters - ), - } - - -def _assert_exact_optimizer_snapshot(optimizer, snapshot): - assert optimizer.__dict__.keys() == snapshot["top"].keys() - for key, value in snapshot["top"].items(): - assert optimizer.__dict__[key] is value - _assert_tree_exact(optimizer.state, snapshot["state"]) - _assert_tree_exact(optimizer.param_groups, snapshot["groups"]) - _assert_tree_exact(optimizer.defaults, snapshot["defaults"]) - _assert_tree_exact(optimizer._param_names, snapshot["param_names"]) - _assert_tree_exact(optimizer._gefen_codebook, snapshot["codebook"]) - for name, expected in snapshot["caches"].items(): - _assert_tree_exact(getattr(optimizer, name), expected) - for parameter, expected in snapshot["grads"]: - _assert_tree_exact(parameter.grad, expected) - - -def _movement_candidates(optimizer, parameter): - candidates = {_tensor_storage_token(optimizer._gefen_codebook)} - candidates.update( - _tensor_storage_token(value) - for key, value in optimizer.state[parameter].items() - if key in _MOVABLE_STATE_KEYS and torch.is_tensor(value) - ) - return candidates - - -def _tensor_storage_token(tensor): - return ( - tensor.device, - tensor.untyped_storage().data_ptr(), - tensor.storage_offset(), - tensor.numel(), - ) - - -def _install_late_to_failure(monkeypatch, candidates, *, destination_type): - original_to = torch.Tensor.to - completed_copies = [] - - def flaky_to(tensor, *args, **kwargs): - result = original_to(tensor, *args, **kwargs) - if ( - _tensor_storage_token(tensor) in candidates - and result is not tensor - and result.device.type == destination_type - ): - completed_copies.append(result) - if len(completed_copies) == 3: - raise RuntimeError("injected late state-copy failure") - return result - - monkeypatch.setattr(torch.Tensor, "to", flaky_to) - return completed_copies - - -def _move_parameter_module(parameter, device): - module = torch.nn.Module() - module.register_parameter("weight", parameter) - module.to(device) - assert module.weight is parameter - return module - - -def _persistent_values(optimizer, parameter): - values = { - key: value.detach().cpu().clone() if torch.is_tensor(value) else copy.deepcopy(value) - for key, value in optimizer.state[parameter].items() - if key in _MOVABLE_STATE_KEYS - } - values["_gefen_codebook"] = optimizer._gefen_codebook.detach().cpu().clone() - return values - - -def _assert_persistent_values_equal(left, right): - assert set(left) == set(right) - for key in left: - if torch.is_tensor(left[key]): - torch.testing.assert_close(left[key], right[key], rtol=0, atol=0, equal_nan=True) - else: - assert left[key] == right[key] - - -@pytest.mark.parametrize("kind", _KINDS) -@pytest.mark.parametrize("explicit_destination", [False, True]) -def test_cpu_state_movement_is_fresh_tight_and_preserves_live_training_objects( - kind, explicit_destination -): - optimizer, parameter, tensor_lr, group_metadata = _build_initialized(kind) - assert optimizer.optimizer_contract().capabilities.atomic_state_movement - assert not optimizer.optimizer_contract().capabilities.state_offload - - _make_persistent_state_oversized(optimizer, parameter) - carrier, state_metadata = _seed_discardable_runtime_state( - optimizer, parameter, tensor_lr - ) - tensors_before = _persistent_tensor_snapshot(optimizer, parameter) - state_before = optimizer.state[parameter] - non_tensor_before = { - key: value - for key, value in state_before.items() - if not torch.is_tensor(value) and key not in _SCRATCH_KEYS - } - groups_before = optimizer.param_groups - group_before = optimizer.param_groups[0] - group_params_before = group_before["params"] - defaults_before = optimizer.defaults - param_names_before = optimizer._param_names - grad_before = parameter.grad - grad_value_before = grad_before.detach().clone() - parameter_value_before = parameter.detach().clone() - global_step_before = optimizer._gefen_global_step - - destination = torch.device("cpu") if explicit_destination else None - optimizer.move_state_(destination) - - assert optimizer.param_groups is groups_before - assert optimizer.param_groups[0] is group_before - assert optimizer.param_groups[0]["params"] is group_params_before - assert optimizer.param_groups[0]["params"][0] is parameter - assert optimizer.param_groups[0]["lr"] is tensor_lr - assert optimizer.param_groups[0]["movement_metadata"] is group_metadata - assert optimizer.defaults is defaults_before - assert optimizer.defaults["lr"] is tensor_lr - assert optimizer._param_names is param_names_before - assert parameter.grad is grad_before - torch.testing.assert_close(parameter.grad, grad_value_before, rtol=0, atol=0) - torch.testing.assert_close(parameter, parameter_value_before, rtol=0, atol=0) - assert optimizer._gefen_global_step == global_step_before - - state = optimizer.state[parameter] - for key, value in non_tensor_before.items(): - assert state[key] is value - assert state["movement_metadata"] is state_metadata - assert state["_gefen_rank_local_payload_0"] is carrier - assert "stepsize" not in state - assert "_h_buf" not in state - - for key, (old, expected) in tensors_before.items(): - actual = optimizer._gefen_codebook if key == "_gefen_codebook" else state[key] - _assert_fresh_tight_copy(actual, old, expected, torch.device("cpu")) - - assert optimizer._gefen_codebook_by_device == {} - assert optimizer._gefen_codebook_lut_by_device == {} - assert optimizer._sr_seed_by_device == {} - assert optimizer._gefen_global_step_by_device == {} - assert optimizer._gefen_codebook_scope_validated is False - assert optimizer._capt_stacks is None - assert optimizer._static_mark_sig is None - assert optimizer._lr_scalar_cache is None - assert optimizer.optimizer_contract().capabilities.atomic_state_movement - assert not optimizer.optimizer_contract().capabilities.state_offload - - -def test_pristine_state_movement_is_repeatable_and_returns_none(): - parameter = torch.nn.Parameter(torch.arange(8, dtype=torch.float32)) - optimizer = Gefen( - [("layer.weight", parameter)], - fused=False, - factored_v_2d=False, - ) - original_state = optimizer.state - original_parameter_state = optimizer.state[parameter] - parameter_value = parameter.detach().clone() - - assert optimizer.move_state_() is None - first_state = optimizer.state - first_parameter_state = optimizer.state[parameter] - assert first_state is not original_state - assert first_parameter_state is not original_parameter_state - assert first_parameter_state == {"name": "layer.weight"} - assert optimizer._gefen_codebook is None - - assert optimizer.move_state_(torch.device("cpu")) is None - assert optimizer.state is not first_state - assert optimizer.state[parameter] is not first_parameter_state - assert optimizer.state[parameter] == {"name": "layer.weight"} - torch.testing.assert_close(parameter, parameter_value, rtol=0, atol=0) - assert optimizer.optimizer_contract().capabilities.atomic_state_movement - - -def test_codebook_preinitialized_state_moves_without_initializing_parameter_tensors(): - parameter = torch.nn.Parameter(torch.arange(1, 9, dtype=torch.float32)) - tensor_lr = torch.tensor(3.0e-3) - optimizer = Gefen( - [("layer.weight", parameter)], - lr=tensor_lr, - fused=False, - factored_v_2d=False, - ) - optimizer._resolve_automatic_period = lambda *args, **kwargs: 4 - parameter.grad = torch.linspace(-2.0, 1.0, parameter.numel()) - optimizer._ensure_gefen_codebook(reuse_existing_periods=False) - assert set(optimizer.state[parameter]) == {"name", "automatic_period"} - optimizer._gefen_codebook = _oversized_copy(optimizer._gefen_codebook) - old_codebook = optimizer._gefen_codebook - expected_codebook = old_codebook.detach().clone() - grad_before = parameter.grad - state_metadata = {"preinitialized": True} - optimizer.state[parameter]["movement_metadata"] = state_metadata - - optimizer.move_state_() - - assert set(optimizer.state[parameter]) == { - "name", - "automatic_period", - "movement_metadata", - } - assert optimizer.state[parameter]["movement_metadata"] is state_metadata - assert parameter.grad is grad_before - assert optimizer._gefen_global_step == 0 - _assert_fresh_tight_copy( - optimizer._gefen_codebook, - old_codebook, - expected_codebook, - torch.device("cpu"), - ) - - -def test_noncontiguous_declared_state_is_normalized_to_a_tight_copy(): - optimizer, parameter, _, _ = _build_initialized("block") - source = torch.arange(12, dtype=torch.float32).reshape(3, 4).t() - assert not source.is_contiguous() - optimizer.state[parameter]["m_magnitude"] = source - expected = source.detach().clone() - - optimizer.move_state_() - - _assert_fresh_tight_copy( - optimizer.state[parameter]["m_magnitude"], - source, - expected, - torch.device("cpu"), - ) - - -def test_wrapper_orphan_state_is_preserved_and_co_located_with_its_key(): - optimizer, _, _, _ = _build_initialized("block") - orphan = torch.nn.Parameter(torch.arange(6, dtype=torch.float32)) - orphan_tensor = _oversized_copy(torch.linspace(1.0, 2.0, 3)) - orphan_carrier = _oversized_copy(torch.tensor([3, 5, 8], dtype=torch.uint8)) - orphan_metadata = {"owner": "wrapper"} - old_orphan_state = { - "name": "orphan.weight", - "automatic_period": 2, - "step": 7, - "m_magnitude": orphan_tensor, - "stepsize": torch.tensor([99.0]), - "_gefen_rank_local_payload_0": orphan_carrier, - "movement_metadata": orphan_metadata, - } - optimizer.state[orphan] = old_orphan_state - expected_tensor = orphan_tensor.detach().clone() - - optimizer.move_state_(torch.device("cpu")) - - assert orphan in optimizer.state - assert optimizer.state[orphan] is not old_orphan_state - moved = optimizer.state[orphan] - assert moved["name"] == "orphan.weight" - assert moved["automatic_period"] == 2 - assert moved["step"] == 7 - assert moved["movement_metadata"] is orphan_metadata - assert moved["_gefen_rank_local_payload_0"] is orphan_carrier - assert "stepsize" not in moved - _assert_fresh_tight_copy( - moved["m_magnitude"], - orphan_tensor, - expected_tensor, - torch.device("cpu"), - ) - - -@pytest.mark.parametrize("kind", _KINDS) -def test_native_checkpoint_continuation_remains_exact_after_movement(kind): - optimizer, parameter, _, _ = _build_initialized(kind) - resumed, resumed_parameter, _, _ = _build_initialized(kind) - - optimizer.move_state_() - resumed.load_state_dict(copy.deepcopy(optimizer.state_dict())) - _assert_persistent_values_equal( - _persistent_values(optimizer, parameter), - _persistent_values(resumed, resumed_parameter), - ) - - continuation_grad = torch.linspace(0.45, -0.85, parameter.numel()).reshape_as( - parameter - ) - parameter.grad = continuation_grad.clone() - resumed_parameter.grad = continuation_grad.clone() - optimizer.step() - resumed.step() - - torch.testing.assert_close(parameter, resumed_parameter, rtol=0, atol=0) - _assert_persistent_values_equal( - _persistent_values(optimizer, parameter), - _persistent_values(resumed, resumed_parameter), - ) - - -def test_movement_invalidates_prepared_canonical_import_but_keeps_io_available(): - def build_finalized(): - parameter = torch.nn.Parameter(torch.linspace(-0.5, 0.7, 8)) - optimizer = Gefen( - [("layer.weight", parameter)], - fused=False, - factored_v_2d=False, - ) - optimizer.rebind_parameter( - parameter, - parameter, - identity=ParameterIdentity("Layer.Weight", (8,)), - ) - optimizer._resolve_automatic_period = lambda *args, **kwargs: 4 - return optimizer, parameter - - source, source_parameter = build_finalized() - source_parameter.grad = torch.linspace(-1.0, 0.8, 8) - source.step() - source.move_state_() - exported = source.export_canonical_state() - - target, target_parameter = build_finalized() - prepared = target.prepare_canonical_state_import(exported) - target.move_state_() - with pytest.raises(RuntimeError, match="changed after canonical import preparation"): - target.commit_canonical_state_import(prepared) - - assert target.optimizer_contract().capabilities.canonical_state_io - assert target.optimizer_contract().capabilities.atomic_state_movement - target.import_canonical_state(exported) - _assert_persistent_values_equal( - _persistent_values(source, source_parameter), - _persistent_values(target, target_parameter), - ) - - -@pytest.mark.parametrize("kind", _KINDS) -def test_late_cpu_copy_failure_is_exactly_atomic(kind, monkeypatch): - optimizer, parameter, _, _ = _build_initialized(kind) - _make_persistent_state_oversized(optimizer, parameter) - candidates = _movement_candidates(optimizer, parameter) - assert len(candidates) >= 3 - snapshot = _snapshot_exact_optimizer(optimizer) - completed = _install_late_to_failure( - monkeypatch, - candidates, - destination_type="cpu", - ) - - with pytest.raises(RuntimeError, match="injected late state-copy failure"): - optimizer.move_state_() - - assert len(completed) == 3 - _assert_exact_optimizer_snapshot(optimizer, snapshot) - - -def _build_rejection_case(case): - if case == "capturable": - parameter = torch.nn.Parameter(torch.ones(8)) - optimizer = Gefen( - [("layer.weight", parameter)], - fused=False, - factored_v_2d=False, - capturable=True, - ) - return optimizer, parameter, None - if case == "meta_parameter": - parameter = torch.nn.Parameter(torch.empty(8, device="meta")) - optimizer = Gefen( - [("layer.weight", parameter)], - fused=False, - factored_v_2d=False, - ) - return optimizer, parameter, None - - optimizer, parameter, _, _ = _build_initialized("block") - if case == "meta_destination": - return optimizer, parameter, torch.device("meta") - if case == "mismatched_destination": - return optimizer, parameter, torch.device("cuda:0") - if case == "meta_state": - optimizer.state[parameter]["m_magnitude"] = optimizer.state[parameter][ - "m_magnitude" - ].to("meta") - elif case == "undeclared_tensor": - optimizer.state[parameter]["extension"] = torch.ones(2) - elif case == "tensor_in_extension_container": - optimizer.state[parameter]["extension"] = { - "nested": [torch.ones(2)] - } - elif case == "tensor_in_ordered_extension": - optimizer.state[parameter]["extension"] = OrderedDict( - (("nested", torch.ones(2)),) - ) - elif case == "tensor_in_rank_local_carrier": - optimizer.state[parameter]["_gefen_rank_local_payload_0"] = { - "nested": torch.ones(2) - } - elif case == "nested_tensor_layout": - with warnings.catch_warnings(): - warnings.simplefilter("ignore", UserWarning) - optimizer.state[parameter]["m_magnitude"] = torch.nested.nested_tensor( - [torch.ones(1), torch.ones(2)] - ) - elif case == "tensor_subclass_state": - class StateTensor(torch.Tensor): - pass - - optimizer.state[parameter]["m_magnitude"] = optimizer.state[parameter][ - "m_magnitude" - ].as_subclass(StateTensor) - elif case in {"opaque_tensor_extension", "opaque_scalar_extension"}: - class OpaqueExtension: - def __init__(self, payload): - self.payload = payload - - def __eq__(self, other): - if not isinstance(other, OpaqueExtension): - return False - if torch.is_tensor(self.payload): - return torch.equal(self.payload, other.payload) - return self.payload == other.payload - - payload = torch.ones(1) if case == "opaque_tensor_extension" else "metadata" - optimizer.state[parameter]["extension"] = OpaqueExtension(payload) - elif case == "defaultdict_factory_extension": - hidden = torch.ones(1) - optimizer.state[parameter]["extension"] = defaultdict( - lambda: hidden, - {"metadata": "value"}, - ) - elif case == "ordered_hidden_extension": - extension = OrderedDict((("metadata", "value"),)) - extension.hidden_tensor = torch.ones(1) - optimizer.state[parameter]["extension"] = extension - elif case == "custom_parameter_state_mapping": - custom_state = OrderedDict(optimizer.state[parameter]) - custom_state.hidden_tensor = torch.ones(1) - optimizer.state[parameter] = custom_state - elif case == "custom_top_level_state_mapping": - custom_state = OrderedDict(optimizer.state) - custom_state.hidden_tensor = torch.ones(1) - optimizer.state = custom_state - else: - raise AssertionError("unknown rejection case: {}".format(case)) - return optimizer, parameter, None - - -@pytest.mark.parametrize( - "case", - ( - "capturable", - "meta_parameter", - "meta_destination", - "mismatched_destination", - "meta_state", - "undeclared_tensor", - "tensor_in_extension_container", - "tensor_in_ordered_extension", - "tensor_in_rank_local_carrier", - "nested_tensor_layout", - "tensor_subclass_state", - "opaque_tensor_extension", - "opaque_scalar_extension", - "defaultdict_factory_extension", - "ordered_hidden_extension", - "custom_parameter_state_mapping", - "custom_top_level_state_mapping", - ), -) -def test_invalid_state_movement_is_rejected_before_any_live_mutation(case): - optimizer, _, destination = _build_rejection_case(case) - snapshot = _snapshot_exact_optimizer(optimizer) - - if case not in {"meta_destination", "mismatched_destination"}: - assert not optimizer.optimizer_contract().capabilities.atomic_state_movement - - with pytest.raises((TypeError, ValueError, RuntimeError)): - optimizer.move_state_(destination) - - _assert_exact_optimizer_snapshot(optimizer, snapshot) - - -@pytest.mark.parametrize("graph_type", ("cycle", "shared_container")) -def test_non_tree_extension_metadata_is_rejected_without_mutation(graph_type): - optimizer, parameter, _, _ = _build_initialized("block") - if graph_type == "cycle": - extension = [] - extension.append(extension) - else: - shared = ["metadata"] - extension = [shared, shared] - optimizer.state[parameter]["extension"] = extension - state_before = optimizer.state - parameter_state_before = optimizer.state[parameter] - items_before = tuple(parameter_state_before.items()) - codebook_before = optimizer._gefen_codebook - - assert not optimizer.optimizer_contract().capabilities.atomic_state_movement - with pytest.raises(RuntimeError, match="not provably tensor-free metadata"): - optimizer.move_state_() - - assert optimizer.state is state_before - assert optimizer.state[parameter] is parameter_state_before - assert tuple(optimizer.state[parameter]) == tuple(key for key, _ in items_before) - for key, value in items_before: - assert optimizer.state[parameter][key] is value - assert optimizer._gefen_codebook is codebook_before - - -@pytest.mark.parametrize( - "runtime_state", - ("compile", "capture", "capturable_stacks", "device_counter", "sr_seed"), -) -def test_active_runtime_state_disables_movement_without_mutation( - runtime_state, monkeypatch -): - optimizer, _, _, _ = _build_initialized("block") - if runtime_state == "compile": - monkeypatch.setattr(torch.compiler, "is_compiling", lambda: True) - elif runtime_state == "capture": - monkeypatch.setattr(torch.cuda, "is_available", lambda: True) - monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: True) - elif runtime_state == "capturable_stacks": - optimizer._capt_stacks = {} - elif runtime_state == "device_counter": - optimizer._gefen_global_step_by_device[torch.device("cpu")] = torch.tensor(1) - elif runtime_state == "sr_seed": - optimizer._sr_seed_by_device[torch.device("cpu")] = torch.tensor(1) - else: - raise AssertionError("unknown runtime state: {}".format(runtime_state)) - - assert not optimizer.optimizer_contract().capabilities.atomic_state_movement - snapshot = _snapshot_exact_optimizer(optimizer) - with pytest.raises(RuntimeError): - optimizer.move_state_() - _assert_exact_optimizer_snapshot(optimizer, snapshot) - - -def test_stale_finalized_binding_disables_movement_without_mutation(): - parameter = torch.nn.Parameter(torch.ones(8)) - optimizer = Gefen( - [("layer.weight", parameter)], - fused=False, - factored_v_2d=False, - ) - optimizer.rebind_parameter( - parameter, - parameter, - identity=ParameterIdentity("Layer.Weight", (8,)), - ) - replacement = torch.nn.Parameter(torch.full((8,), 2.0)) - optimizer.param_groups[0]["params"][0] = replacement - - assert not optimizer.optimizer_contract().capabilities.atomic_state_movement - snapshot = _snapshot_exact_optimizer(optimizer) - with pytest.raises(RuntimeError, match="finalized parameter layout changed"): - optimizer.move_state_() - _assert_exact_optimizer_snapshot(optimizer, snapshot) - - -def test_movement_does_not_narrow_muons_generic_local_state_device_helper(): - meta_parameter = torch.empty(4, device="meta") - - assert GefenMuon._state_tensor_device(meta_parameter) == torch.device("meta") - with pytest.raises(RuntimeError, match="only CPU and CUDA"): - GefenMuon._state_move_parameter_device(meta_parameter) - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") -@pytest.mark.parametrize("kind", _KINDS) -def test_cpu_cuda_cpu_state_round_trip_preserves_cpu_continuation(kind): - optimizer, parameter, tensor_lr, _ = _build_initialized(kind) - reference, reference_parameter, reference_lr, _ = _build_initialized(kind) - initial_reference = _persistent_values(reference, reference_parameter) - _assert_persistent_values_equal( - _persistent_values(optimizer, parameter), initial_reference - ) - module = _move_parameter_module(parameter, "cuda") - cuda_grad = parameter.grad - cuda_lr = optimizer.param_groups[0]["lr"] - - optimizer.move_state_() - - assert parameter.grad is cuda_grad - assert optimizer.param_groups[0]["lr"] is cuda_lr is tensor_lr - assert optimizer._gefen_codebook.device.type == "cuda" - for key, value in optimizer.state[parameter].items(): - if key in _MOVABLE_STATE_KEYS and torch.is_tensor(value): - assert value.device.type == "cuda" - - module.to("cpu") - assert module.weight is parameter - cpu_grad = parameter.grad - optimizer.move_state_(torch.device("cpu")) - assert parameter.grad is cpu_grad - assert optimizer.param_groups[0]["lr"] is tensor_lr - assert optimizer._gefen_codebook.device.type == "cpu" - - continuation_grad = torch.linspace(0.6, -0.9, parameter.numel()).reshape_as(parameter) - parameter.grad = continuation_grad.clone() - reference_parameter.grad = continuation_grad.clone() - optimizer.step() - reference.step() - - assert optimizer.param_groups[0]["lr"] is tensor_lr - assert reference.param_groups[0]["lr"] is reference_lr - torch.testing.assert_close(parameter, reference_parameter, rtol=0, atol=0) - _assert_persistent_values_equal( - _persistent_values(optimizer, parameter), - _persistent_values(reference, reference_parameter), - ) - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") -def test_real_cpu_to_cuda_copy_failure_is_exactly_atomic(monkeypatch): - optimizer, parameter, _, _ = _build_initialized("block") - module = _move_parameter_module(parameter, "cuda") - assert optimizer._gefen_codebook.device.type == "cpu" - assert optimizer.state[parameter]["m_codebook"].device.type == "cpu" - candidates = _movement_candidates(optimizer, parameter) - snapshot = _snapshot_exact_optimizer(optimizer) - completed = _install_late_to_failure( - monkeypatch, - candidates, - destination_type="cuda", - ) - - with pytest.raises(RuntimeError, match="injected late state-copy failure"): - optimizer.move_state_() - - assert len(completed) == 3 - assert all(tensor.device.type == "cuda" for tensor in completed) - _assert_exact_optimizer_snapshot(optimizer, snapshot) - module.to("cpu") diff --git a/tests/test_state_movement_distributed.py b/tests/test_state_movement_distributed.py deleted file mode 100644 index 7801c61..0000000 --- a/tests/test_state_movement_distributed.py +++ /dev/null @@ -1,577 +0,0 @@ -"""Real-process-group coverage for atomic optimizer-state movement.""" - -from __future__ import annotations - -from datetime import timedelta -import os -import queue -import tempfile -import traceback - -import pytest -import torch - - -_AUTHORITATIVE_TENSOR_KEYS = frozenset( - { - "step", - "m_codebook", - "m_magnitude", - "vmean", - "vmean_step", - "v_row", - "v_col", - "factored_step", - "normuon_v", - "normuon_step", - } -) - - -def _oversized_copy(tensor: torch.Tensor) -> torch.Tensor: - backing = torch.empty( - tensor.numel() + 13, - dtype=tensor.dtype, - device=tensor.device, - ) - result = backing.narrow(0, 7, tensor.numel()).view(tensor.shape) - result.copy_(tensor) - assert result.untyped_storage().nbytes() > tensor.numel() * tensor.element_size() - return result - - -def _assert_fresh_tight_copy( - actual: torch.Tensor, - old: torch.Tensor, - expected: torch.Tensor, -) -> None: - assert type(actual) is torch.Tensor - assert actual is not old - assert actual.device == torch.device("cpu") - assert actual.dtype == expected.dtype - assert actual.shape == expected.shape - assert actual.layout is torch.strided - assert actual.is_contiguous() - assert actual.storage_offset() == 0 - assert actual.untyped_storage().nbytes() == actual.numel() * actual.element_size() - torch.testing.assert_close(actual, expected, rtol=0, atol=0, equal_nan=True) - - -def _seed_and_snapshot_movable_state(optimizer, parameters, rank: int, label: str): - first_state = optimizer.state[parameters[0]] - carrier = torch.tensor([rank, len(label), 41], dtype=torch.int64) - metadata = {"label": label, "members": ["preserve", rank]} - first_state["_gefen_rank_local_payload_{}".format(rank)] = carrier - first_state["movement_metadata"] = metadata - - records = [] - for parameter_index, parameter in enumerate(parameters): - for key, value in tuple(optimizer.state[parameter].items()): - if key not in _AUTHORITATIVE_TENSOR_KEYS or not torch.is_tensor(value): - continue - assert type(value) is torch.Tensor - oversized = _oversized_copy(value) - optimizer.state[parameter][key] = oversized - records.append( - ( - parameter_index, - key, - oversized, - oversized.detach().clone(), - ) - ) - - assert type(optimizer._gefen_codebook) is torch.Tensor - optimizer._gefen_codebook = _oversized_copy(optimizer._gefen_codebook) - codebook_record = ( - optimizer._gefen_codebook, - optimizer._gefen_codebook.detach().clone(), - ) - return records, codebook_record, carrier, metadata - - -def _assert_authoritative_state_equal(optimizer, reference, parameters, reference_parameters): - for parameter, reference_parameter in zip(parameters, reference_parameters): - state = optimizer.state[parameter] - reference_state = reference.state[reference_parameter] - keys = { - key - for key in state - if key in _AUTHORITATIVE_TENSOR_KEYS - } | { - key - for key in reference_state - if key in _AUTHORITATIVE_TENSOR_KEYS - } - for key in keys: - assert key in state and key in reference_state - actual = state[key] - expected = reference_state[key] - if torch.is_tensor(expected): - assert torch.is_tensor(actual) - torch.testing.assert_close(actual, expected, rtol=0, atol=0, equal_nan=True) - else: - assert actual == expected - torch.testing.assert_close( - optimizer._gefen_codebook, - reference._gefen_codebook, - rtol=0, - atol=0, - equal_nan=True, - ) - - -def _dtensor_case(rank, world, mesh, kind: str) -> None: - import torch.nn as nn - from torch.distributed.tensor import Shard, distribute_tensor - - from gefen import Gefen, GefenMuon - - shapes = ((4, 4), (4, 6)) if kind == "gefen" else ((4, 4), (1, 4)) - generator = torch.Generator().manual_seed(7100 + sum(ord(item) for item in kind)) - initial = [torch.randn(shape, generator=generator) * 0.05 for shape in shapes] - - def build(): - parameters = [ - nn.Parameter(distribute_tensor(value.clone(), mesh, [Shard(0)])) - for value in initial - ] - tensor_lr = torch.tensor(2.0e-3) - group_metadata = {"kind": kind, "ordered_members": list(range(world))} - group = { - "params": [ - ("{}.{}".format(kind, index), parameter) - for index, parameter in enumerate(parameters) - ], - "lr": tensor_lr, - "movement_metadata": group_metadata, - } - if kind == "gefen": - optimizer = Gefen( - [group], - lr=tensor_lr, - fused=False, - factored_v_2d=False, - ) - else: - optimizer = GefenMuon( - [group], - lr=tensor_lr, - fused=False, - ns_steps=1, - weight_decay=0.0, - sharded_mode=kind, - ) - optimizer._resolve_automatic_period = lambda *args, **kwargs: 4 - optimizer._predict_period_from_grad_sq = lambda *args, **kwargs: 4 - return optimizer, parameters, tensor_lr, group_metadata - - def assign_grads(parameters, seed): - grad_generator = torch.Generator().manual_seed(seed) - for parameter, shape in zip(parameters, shapes): - full_grad = torch.randn(shape, generator=grad_generator) * 0.01 - parameter.grad = distribute_tensor(full_grad, mesh, [Shard(0)]) - - optimizer, parameters, tensor_lr, group_metadata = build() - reference, reference_parameters, _, _ = build() - assign_grads(parameters, 7200) - assign_grads(reference_parameters, 7200) - optimizer.step() - reference.step() - - if kind == "distributed": - stateful = { - index - for index, parameter in enumerate(parameters) - if any( - key in _AUTHORITATIVE_TENSOR_KEYS - for key in optimizer.state[parameter] - ) - } - assert stateful == {rank} - if rank == 1: - assert parameters[1].to_local().numel() == 0 - assert optimizer.state[parameters[1]]["m_codebook"].numel() == shapes[1][0] * shapes[1][1] - elif kind == "approx": - stateful = { - index - for index, parameter in enumerate(parameters) - if any( - key in _AUTHORITATIVE_TENSOR_KEYS - for key in optimizer.state[parameter] - ) - } - assert stateful == ({0, 1} if rank == 0 else {0}) - else: - assert all( - any(key in _AUTHORITATIVE_TENSOR_KEYS for key in optimizer.state[parameter]) - for parameter in parameters - ) - - records, codebook_record, carrier, state_metadata = _seed_and_snapshot_movable_state( - optimizer, - parameters, - rank, - kind, - ) - assert records - contract_before = optimizer.optimizer_contract() - assert contract_before.capabilities.atomic_state_movement - assert not contract_before.capabilities.state_offload - - groups_before = optimizer.param_groups - group_before = optimizer.param_groups[0] - group_params_before = group_before["params"] - defaults_before = optimizer.defaults - names_before = optimizer._param_names - grads_before = [parameter.grad for parameter in parameters] - local_values_before = [parameter.detach().to_local().clone() for parameter in parameters] - meshes_before = [parameter.device_mesh for parameter in parameters] - placements_before = [parameter.placements for parameter in parameters] - mesh_groups_before = [parameter.device_mesh.get_group() for parameter in parameters] - codebook_binding_before = optimizer._gefen_codebook_process_group - shard_bindings_before = optimizer._gefen_shard_bindings - local_bindings_before = optimizer._gefen_local_shard_bindings - manifest_before = optimizer._gefen_sharding_manifest - - optimizer.move_state_() - - assert optimizer.optimizer_contract() == contract_before - assert optimizer.param_groups is groups_before - assert optimizer.param_groups[0] is group_before - assert optimizer.param_groups[0]["params"] is group_params_before - assert optimizer.param_groups[0]["lr"] is tensor_lr - assert optimizer.param_groups[0]["movement_metadata"] is group_metadata - assert optimizer.defaults is defaults_before - assert optimizer.defaults["lr"] is tensor_lr - assert optimizer._param_names is names_before - assert optimizer._gefen_codebook_process_group is codebook_binding_before - assert optimizer._gefen_shard_bindings is shard_bindings_before - assert optimizer._gefen_local_shard_bindings is local_bindings_before - assert optimizer._gefen_sharding_manifest is manifest_before - - for index, parameter in enumerate(parameters): - assert optimizer.param_groups[0]["params"][index] is parameter - assert parameter.grad is grads_before[index] - assert parameter.device_mesh is meshes_before[index] - assert parameter.placements == placements_before[index] - assert parameter.device_mesh.get_group() is mesh_groups_before[index] - torch.testing.assert_close( - parameter.detach().to_local(), - local_values_before[index], - rtol=0, - atol=0, - ) - - for parameter_index, key, old, expected in records: - _assert_fresh_tight_copy( - optimizer.state[parameters[parameter_index]][key], - old, - expected, - ) - _assert_fresh_tight_copy( - optimizer._gefen_codebook, - codebook_record[0], - codebook_record[1], - ) - assert optimizer.state[parameters[0]]["_gefen_rank_local_payload_{}".format(rank)] is carrier - assert optimizer.state[parameters[0]]["movement_metadata"] is state_metadata - - assign_grads(parameters, 7300) - assign_grads(reference_parameters, 7300) - optimizer.step() - reference.step() - for parameter, reference_parameter in zip(parameters, reference_parameters): - torch.testing.assert_close( - parameter.detach().to_local(), - reference_parameter.detach().to_local(), - rtol=0, - atol=0, - ) - _assert_authoritative_state_equal( - optimizer, - reference, - parameters, - reference_parameters, - ) - - -def _whole_owner_case(rank, world) -> None: - import torch.distributed as dist - import torch.nn as nn - - from gefen import ( - CodebookProcessGroupBinding, - GefenMuon, - LogicalSlice, - ParameterIdentity, - ParameterLayout, - ParameterRebinding, - PlacementKind, - ProcessGroupIdentity, - ShardIdentity, - ShardPlacement, - ShardingManifest, - ) - - members = tuple("rank:{}".format(index) for index in range(world)) - group_identity = ProcessGroupIdentity("movement_owner", members) - identities = ( - ParameterIdentity("Owner.First", (4, 4)), - ParameterIdentity("Owner.Second", (4, 4)), - ) - - def owner_shard(identity, member, owner): - return ShardIdentity( - identity, - ParameterLayout.WHOLE_PARAMETER_OWNER, - LogicalSlice.full(identity) if member == owner else LogicalSlice(0, 0), - process_group=group_identity, - local_member=member, - owner=owner, - placements=( - ShardPlacement( - "dp", - PlacementKind.WHOLE_PARAMETER_OWNER, - members.index(member), - world, - ), - ), - ) - - records_by_parameter = tuple( - tuple( - owner_shard(identity, member, members[index]) - for member in members - ) - for index, identity in enumerate(identities) - ) - manifest = ShardingManifest( - tuple( - shard - for parameter_records in records_by_parameter - for shard in parameter_records - ) - ) - - def build(): - generator = torch.Generator().manual_seed(7400) - old_parameters = [ - nn.Parameter(torch.randn(identity.global_shape, generator=generator) * 0.05) - for identity in identities - ] - tensor_lr = torch.tensor(2.0e-3) - group_metadata = {"layout": "whole-owner", "rank": rank} - optimizer = GefenMuon( - [ - { - "params": [ - ("owner.{}".format(index), parameter) - for index, parameter in enumerate(old_parameters) - ], - "lr": tensor_lr, - "movement_metadata": group_metadata, - } - ], - lr=tensor_lr, - fused=False, - ns_steps=1, - weight_decay=0.0, - ) - local_member = members[rank] - local_records = [ - next( - shard - for shard in parameter_records - if shard.local_member == local_member - ) - for parameter_records in records_by_parameter - ] - binding = CodebookProcessGroupBinding( - group_identity, - local_member, - dist.group.WORLD, - torch.device("cpu"), - ) - optimizer.post_sharding( - tuple( - ParameterRebinding( - parameter, - parameter if local_record.owner == local_member else None, - local_record, - ) - for parameter, local_record in zip(old_parameters, local_records) - ), - manifest=manifest, - codebook_process_group=binding, - ) - optimizer._resolve_automatic_period = lambda *args, **kwargs: 4 - optimizer._predict_period_from_grad_sq = lambda *args, **kwargs: 4 - return ( - optimizer, - list(optimizer.param_groups[0]["params"]), - tensor_lr, - group_metadata, - binding, - ) - - def assign_grad(parameters, seed): - generator = torch.Generator().manual_seed(seed + rank) - assert len(parameters) == 1 - parameters[0].grad = torch.randn(parameters[0].shape, generator=generator) * 0.01 - - optimizer, parameters, tensor_lr, group_metadata, binding = build() - reference, reference_parameters, _, _, _ = build() - assert len(parameters) == 1 - assert len(optimizer.shard_bindings()) == 2 - assert sum(parameter is None for parameter, _ in optimizer.shard_bindings()) == 1 - - assign_grad(parameters, 7500) - assign_grad(reference_parameters, 7500) - optimizer.step() - reference.step() - - records, codebook_record, carrier, state_metadata = _seed_and_snapshot_movable_state( - optimizer, - parameters, - rank, - "whole-owner", - ) - assert records - contract_before = optimizer.optimizer_contract() - assert contract_before.capabilities.atomic_state_movement - assert not contract_before.capabilities.state_offload - bindings_before = optimizer.shard_bindings() - manifest_before = optimizer.sharding_manifest() - binding_before = optimizer.codebook_process_group_binding() - groups_before = optimizer.param_groups - group_before = optimizer.param_groups[0] - group_params_before = group_before["params"] - defaults_before = optimizer.defaults - names_before = optimizer._param_names - grad_before = parameters[0].grad - value_before = parameters[0].detach().clone() - - optimizer.move_state_() - - assert optimizer.optimizer_contract() == contract_before - assert optimizer.param_groups is groups_before - assert optimizer.param_groups[0] is group_before - assert optimizer.param_groups[0]["params"] is group_params_before - assert optimizer.param_groups[0]["params"][0] is parameters[0] - assert optimizer.param_groups[0]["lr"] is tensor_lr - assert optimizer.param_groups[0]["movement_metadata"] is group_metadata - assert optimizer.defaults is defaults_before - assert optimizer.defaults["lr"] is tensor_lr - assert optimizer._param_names is names_before - assert optimizer.shard_bindings() is bindings_before - assert optimizer.sharding_manifest() is manifest_before - assert optimizer.codebook_process_group_binding() is binding_before is binding - assert parameters[0].grad is grad_before - torch.testing.assert_close(parameters[0], value_before, rtol=0, atol=0) - - for parameter_index, key, old, expected in records: - _assert_fresh_tight_copy( - optimizer.state[parameters[parameter_index]][key], - old, - expected, - ) - _assert_fresh_tight_copy( - optimizer._gefen_codebook, - codebook_record[0], - codebook_record[1], - ) - assert optimizer.state[parameters[0]]["_gefen_rank_local_payload_{}".format(rank)] is carrier - assert optimizer.state[parameters[0]]["movement_metadata"] is state_metadata - - assign_grad(parameters, 7600) - assign_grad(reference_parameters, 7600) - optimizer.step() - reference.step() - torch.testing.assert_close( - parameters[0], - reference_parameters[0], - rtol=0, - atol=0, - ) - _assert_authoritative_state_equal( - optimizer, - reference, - parameters, - reference_parameters, - ) - - -def _distributed_worker(rank, world, init_file, result_queue) -> None: - import torch.distributed as dist - 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",)) - for kind in ("gefen", "exact", "approx", "distributed"): - _dtensor_case(rank, world, mesh, kind) - _whole_owner_case(rank, world) - result_queue.put(("result", rank)) - 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="distributed state movement coverage needs Gloo", -) -def test_atomic_state_movement_across_distributed_cpu_representations(): - import torch.multiprocessing as mp - - world = 2 - context = mp.get_context("spawn") - result_queue = context.Queue() - # A file:// rendezvous stays valid until every rank has initialized, unlike a - # pre-probed free TCP port that another process can steal before the workers - # bind it. Mirrors the DCP distributed test. - descriptor, init_file = tempfile.mkstemp(prefix="gefen-state-movement-") - os.close(descriptor) - os.unlink(init_file) - processes = [ - context.Process( - target=_distributed_worker, - args=(rank, world, init_file, 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=180)) - 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) - if os.path.exists(init_file): - os.unlink(init_file) - - assert all(process.exitcode == 0 for process in processes), [ - process.exitcode for process in processes - ] - errors = [item[2] for item in messages if item[0] == "error"] - assert not errors, "\n".join(errors) - assert {item[1] for item in messages if item[0] == "result"} == set(range(world)) diff --git a/tests/test_state_offload.py b/tests/test_state_offload.py deleted file mode 100644 index 1ecf49f..0000000 --- a/tests/test_state_offload.py +++ /dev/null @@ -1,774 +0,0 @@ -"""Focused CPU/CUDA coverage for native plain-Gefen parameter-state offload.""" - -import copy -import io -import types - -import pytest -import torch - -from gefen import ( - CheckpointProcessGroupBinding, - CheckpointTransport, - Gefen, - GefenMuon, - PortableStateLimits, - StateOffloadProvider, -) -from gefen.codebook import CodebookProcessGroupBinding -from gefen.contracts import ( - LogicalSlice, - ParameterIdentity, - ParameterLayout, - PlacementKind, - ProcessGroupIdentity, - ShardIdentity, - ShardPlacement, - ShardingManifest, -) -from gefen.rebinding import ParameterRebinding - - -_PERSISTENT_TENSOR_KEYS = frozenset({"m_codebook", "m_magnitude", "vmean", "v_row", "v_col"}) -_CUDA_REQUIRED = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") - - -def _make_gefen(parameter, *, factored=False, fused=False): - optimizer = Gefen( - [("layer.weight", parameter)], - lr=2.0e-3, - fused=fused, - factored_v_2d=factored, - deterministic=True, - ) - optimizer._resolve_automatic_period = lambda *args, **kwargs: 4 - return optimizer - - -def _persistent_snapshot(optimizer, parameter): - result = {} - for key, value in optimizer.state[parameter].items(): - if key in _PERSISTENT_TENSOR_KEYS: - result[key] = value.detach().cpu().clone() - elif key in {"automatic_period", "step", "vmean_step", "factored_step"}: - result[key] = value - return result - - -def _assert_persistent_equal(left, right): - assert set(left) == set(right) - for key in left: - if torch.is_tensor(left[key]): - torch.testing.assert_close(left[key], right[key], rtol=0, atol=0) - else: - assert left[key] == right[key] - - -def _assert_cpu_boundary(optimizer): - assert optimizer.state_offload_active - assert optimizer.state_offload_device == torch.device("cpu") - assert not optimizer.state_offload_poisoned - for parameter_state in optimizer.state.values(): - assert "stepsize" not in parameter_state - assert "_h_buf" not in parameter_state - for key, value in parameter_state.items(): - if key in _PERSISTENT_TENSOR_KEYS: - assert type(value) is torch.Tensor - assert value.device.type == "cpu" - assert value.is_contiguous() - assert value.storage_offset() == 0 - assert value.untyped_storage().nbytes() == value.numel() * value.element_size() - - -def test_state_offload_visibility_is_read_only_and_cpu_parameters_fail_closed(): - parameter = torch.nn.Parameter(torch.ones(8)) - optimizer = _make_gefen(parameter) - - assert not optimizer.state_offload_active - assert optimizer.state_offload_device is None - assert not optimizer.state_offload_poisoned - assert isinstance(optimizer, StateOffloadProvider) - assert not optimizer.optimizer_contract().capabilities.state_offload - with pytest.raises(AttributeError): - optimizer.state_offload_active = True - with pytest.raises(AttributeError): - optimizer.state_offload_device = torch.device("cpu") - with pytest.raises(AttributeError): - optimizer.state_offload_poisoned = True - with pytest.raises(RuntimeError, match="ordinary replicated CUDA"): - optimizer.offload_state_() - - -def test_state_offload_rejects_non_cpu_targets_and_muon(): - parameter = torch.nn.Parameter(torch.ones(2, 4)) - optimizer = _make_gefen(parameter) - muon = GefenMuon( - [("layer.weight", parameter)], - fused=False, - ns_steps=1, - ) - - with pytest.raises(ValueError, match="only CPU"): - optimizer.offload_state_("meta") - with pytest.raises(RuntimeError, match="only by plain Gefen"): - muon.offload_state_() - - -@_CUDA_REQUIRED -@pytest.mark.parametrize("factored", [False, True]) -@pytest.mark.parametrize("fused", [False, True]) -def test_pristine_offload_multistep_is_exact_and_cpu_authoritative(factored, fused): - shape = (2, 4) if factored else (8,) - initial = torch.linspace(-0.7, 0.6, 8, device="cuda").reshape(shape) - reference_parameter = torch.nn.Parameter(initial.clone()) - offloaded_parameter = torch.nn.Parameter(initial.clone()) - reference = _make_gefen(reference_parameter, factored=factored, fused=fused) - offloaded = _make_gefen(offloaded_parameter, factored=factored, fused=fused) - - assert offloaded.optimizer_contract().capabilities.state_offload - offloaded.offload_state_() - assert offloaded.optimizer_contract().capabilities.state_offload - _assert_cpu_boundary(offloaded) - assert offloaded._gefen_codebook is None - - for step in range(3): - grad = torch.linspace( - -1.1 + 0.2 * step, - 0.9 - 0.1 * step, - 8, - device="cuda", - ).reshape(shape) - reference_parameter.grad = grad.clone() - offloaded_parameter.grad = grad.clone() - reference.step() - offloaded.step() - - torch.testing.assert_close(offloaded_parameter, reference_parameter, rtol=0, atol=0) - _assert_persistent_equal( - _persistent_snapshot(offloaded, offloaded_parameter), - _persistent_snapshot(reference, reference_parameter), - ) - assert offloaded._gefen_global_step == reference._gefen_global_step - _assert_cpu_boundary(offloaded) - assert offloaded._gefen_codebook.device.type == "cuda" - - -@_CUDA_REQUIRED -def test_initialized_activation_keeps_common_codebook_resident_and_move_disables(): - parameter = torch.nn.Parameter(torch.linspace(-0.4, 0.8, 8, device="cuda")) - optimizer = _make_gefen(parameter) - parameter.grad = torch.linspace(-1.0, 1.0, 8, device="cuda") - optimizer.step() - codebook = optimizer._gefen_codebook - cache = optimizer._gefen_codebook_by_device - cache[parameter.device] = codebook - - optimizer.offload_state_("cpu:0") - - _assert_cpu_boundary(optimizer) - assert optimizer._gefen_codebook is codebook - assert optimizer._gefen_codebook_by_device is cache - assert optimizer._gefen_codebook_by_device[parameter.device] is codebook - - optimizer.move_state_() - - assert not optimizer.state_offload_active - assert optimizer.state_offload_device is None - assert not optimizer.state_offload_poisoned - assert all( - value.device == parameter.device - for key, value in optimizer.state[parameter].items() - if key in _PERSISTENT_TENSOR_KEYS - ) - - -@_CUDA_REQUIRED -def test_runtime_state_is_private_and_only_one_parameter_is_staged(): - first = torch.nn.Parameter(torch.arange(8, device="cuda", dtype=torch.float32)) - second = torch.nn.Parameter(torch.arange(8, 16, device="cuda", dtype=torch.float32)) - optimizer = Gefen( - [("first", first), ("second", second)], - fused=False, - factored_v_2d=False, - ) - optimizer._resolve_automatic_period = lambda *args, **kwargs: 4 - optimizer.offload_state_() - observed = [] - original = optimizer._step_automatic - - def inspected(self, group, name, parameter, grad, *, state=None): - assert state is not self.state[parameter] - assert all( - value.device.type == "cpu" - for published in self.state.values() - for key, value in published.items() - if key in _PERSISTENT_TENSOR_KEYS - ) - assert all(value.device == parameter.device for key, value in state.items() if key in _PERSISTENT_TENSOR_KEYS) - observed.append(parameter) - return original(group, name, parameter, grad, state=state) - - optimizer._step_automatic = types.MethodType(inspected, optimizer) - first.grad = torch.linspace(-1.0, 1.0, 8, device="cuda") - second.grad = torch.linspace(1.0, -1.0, 8, device="cuda") - optimizer.step() - - assert observed == [first, second] - _assert_cpu_boundary(optimizer) - - -@_CUDA_REQUIRED -def test_later_parameter_corruption_is_caught_before_earlier_parameter_mutates(): - # Regression: the offloaded step processes parameters sequentially - # (stage -> update -> copyback -> commit per parameter). If a later - # parameter's offloaded state is corrupted in a way that preserves the - # state/param-group container identities and the layout version, a cached - # step-readiness verdict would let step() begin, mutate the first - # parameter, and only reject when the second parameter is staged -- - # violating fail-before-mutation. The readiness scan therefore runs in - # full on every step, so the corruption is caught at step entry and the - # first parameter is left byte-for-byte untouched. - first = torch.nn.Parameter(torch.arange(8, device="cuda", dtype=torch.float32)) - second = torch.nn.Parameter(torch.arange(8, 16, device="cuda", dtype=torch.float32)) - optimizer = Gefen( - [("first", first), ("second", second)], - fused=False, - factored_v_2d=False, - ) - optimizer._resolve_automatic_period = lambda *args, **kwargs: 4 - optimizer.offload_state_() - - first.grad = torch.linspace(-1.0, 1.0, 8, device="cuda") - second.grad = torch.linspace(1.0, -1.0, 8, device="cuda") - optimizer.step() - - first_state_before = _persistent_snapshot(optimizer, first) - first_value_before = first.detach().clone() - - # Corrupt the SECOND parameter's offloaded state with a non-tight CPU view. - # self.state, self.state[second], and the layout version are all unchanged. - corrupt = optimizer.state[second]["m_magnitude"].repeat_interleave(2)[::2] - assert not corrupt.is_contiguous() - optimizer.state[second]["m_magnitude"] = corrupt - - first.grad = torch.linspace(-1.0, 1.0, 8, device="cuda") - second.grad = torch.linspace(1.0, -1.0, 8, device="cuda") - with pytest.raises(RuntimeError, match="cannot step"): - optimizer.step() - - _assert_persistent_equal(_persistent_snapshot(optimizer, first), first_state_before) - assert torch.equal(first.detach(), first_value_before) - - -@_CUDA_REQUIRED -def test_active_load_preserves_target_policy_and_exact_continuation(monkeypatch): - source_parameter = torch.nn.Parameter(torch.linspace(-0.5, 0.5, 8, device="cuda")) - source = _make_gefen(source_parameter) - source_parameter.grad = torch.linspace(-1.0, 0.7, 8, device="cuda") - source.step() - checkpoint = copy.deepcopy(source.state_dict()) - - target_parameter = torch.nn.Parameter(source_parameter.detach().clone()) - target = _make_gefen(target_parameter) - target.offload_state_() - - def reject_parameter_device_cast(*_args, **_kwargs): - raise AssertionError("active offload load used PyTorch's parameter-device cast") - - monkeypatch.setattr( - torch.optim.Optimizer, - "_process_value_according_to_param_policy", - reject_parameter_device_cast, - ) - target.load_state_dict(checkpoint) - - _assert_cpu_boundary(target) - _assert_persistent_equal( - _persistent_snapshot(source, source_parameter), - _persistent_snapshot(target, target_parameter), - ) - - continuation = torch.linspace(0.8, -0.6, 8, device="cuda") - source_parameter.grad = continuation.clone() - target_parameter.grad = continuation.clone() - source.step() - target.step() - - torch.testing.assert_close(target_parameter, source_parameter, rtol=0, atol=0) - _assert_persistent_equal( - _persistent_snapshot(source, source_parameter), - _persistent_snapshot(target, target_parameter), - ) - _assert_cpu_boundary(target) - - -@_CUDA_REQUIRED -@pytest.mark.parametrize("activate_before_load", [False, True]) -def test_cpu_mapped_checkpoint_keeps_common_codebook_cuda_resident( - activate_before_load, -): - source_parameter = torch.nn.Parameter(torch.linspace(-0.5, 0.5, 8, device="cuda")) - source = _make_gefen(source_parameter) - source_parameter.grad = torch.linspace(-1.0, 0.7, 8, device="cuda") - source.step() - serialized = io.BytesIO() - torch.save(source.state_dict(), serialized) - serialized.seek(0) - checkpoint = torch.load( - serialized, - map_location="cpu", - weights_only=False, - ) - assert checkpoint["gefen_codebook"].device.type == "cpu" - - target_parameter = torch.nn.Parameter(source_parameter.detach().clone()) - target = _make_gefen(target_parameter) - if activate_before_load: - target.offload_state_() - target.load_state_dict(checkpoint) - if not activate_before_load: - assert target._gefen_codebook.device.type == "cpu" - target.offload_state_() - - _assert_cpu_boundary(target) - assert target._gefen_codebook.device == target_parameter.device - torch.testing.assert_close( - target._gefen_codebook, - source._gefen_codebook, - rtol=0, - atol=0, - ) - - continuation = torch.linspace(0.8, -0.6, 8, device="cuda") - source_parameter.grad = continuation.clone() - target_parameter.grad = continuation.clone() - source.step() - target.step() - torch.testing.assert_close(target_parameter, source_parameter, rtol=0, atol=0) - _assert_cpu_boundary(target) - - -@_CUDA_REQUIRED -def test_activation_and_restore_copy_failures_are_atomic(monkeypatch): - parameter = torch.nn.Parameter(torch.linspace(-0.5, 0.5, 8, device="cuda")) - optimizer = _make_gefen(parameter) - parameter.grad = torch.linspace(-1.0, 1.0, 8, device="cuda") - optimizer.step() - state_before = optimizer.state - parameter_state_before = optimizer.state[parameter] - codebook_before = optimizer._gefen_codebook - codebook_value_before = optimizer._gefen_codebook.detach().cpu().clone() - persistent_before = _persistent_snapshot(optimizer, parameter) - global_step_before = optimizer._gefen_global_step - - def fail_cpu_copy(_tensor): - raise RuntimeError("injected activation copy failure") - - monkeypatch.setattr(optimizer, "_copy_state_tensor_to_offload_cpu", fail_cpu_copy) - with pytest.raises(RuntimeError, match="injected activation"): - optimizer.offload_state_() - assert optimizer.state is state_before - assert optimizer.state[parameter] is parameter_state_before - assert optimizer._gefen_codebook is codebook_before - assert not optimizer.state_offload_active - _assert_persistent_equal(_persistent_snapshot(optimizer, parameter), persistent_before) - assert torch.equal(optimizer._gefen_codebook.detach().cpu(), codebook_value_before) - assert optimizer._gefen_global_step == global_step_before - - monkeypatch.undo() - optimizer.offload_state_() - state_before = optimizer.state - parameter_state_before = optimizer.state[parameter] - persistent_before = _persistent_snapshot(optimizer, parameter) - - def fail_move(_tensor, _device): - raise RuntimeError("injected restore copy failure") - - monkeypatch.setattr(optimizer, "_copy_state_tensor_for_move", fail_move) - with pytest.raises(RuntimeError, match="injected restore"): - optimizer.restore_state_() - assert optimizer.state is state_before - assert optimizer.state[parameter] is parameter_state_before - assert optimizer.state_offload_active - _assert_cpu_boundary(optimizer) - _assert_persistent_equal(_persistent_snapshot(optimizer, parameter), persistent_before) - assert torch.equal(optimizer._gefen_codebook.detach().cpu(), codebook_value_before) - assert optimizer._gefen_global_step == global_step_before - - -@_CUDA_REQUIRED -def test_copyback_failure_poison_is_sticky_until_successful_load(monkeypatch): - parameter = torch.nn.Parameter(torch.linspace(-0.4, 0.7, 8, device="cuda")) - optimizer = _make_gefen(parameter) - parameter.grad = torch.linspace(-1.0, 1.0, 8, device="cuda") - optimizer.step() - checkpoint = copy.deepcopy(optimizer.state_dict()) - optimizer.offload_state_() - published_state = optimizer.state[parameter] - parameter.grad = torch.linspace(0.9, -0.8, 8, device="cuda") - - def fail_copyback(_tensor): - raise RuntimeError("injected copyback failure") - - monkeypatch.setattr(optimizer, "_copy_state_tensor_to_offload_cpu", fail_copyback) - with pytest.raises(RuntimeError, match="known-good checkpoint"): - optimizer.step() - - assert optimizer.state_offload_active - assert optimizer.state_offload_poisoned - assert not optimizer.optimizer_contract().capabilities.state_offload - with pytest.raises(RuntimeError, match="cannot export optimizer state"): - optimizer.state_dict() - assert optimizer.state[parameter] is published_state - assert all(value.device.type == "cpu" for key, value in published_state.items() if key in _PERSISTENT_TENSOR_KEYS) - parameter_after_failure = parameter.detach().clone() - with pytest.raises(RuntimeError, match="poisoned"): - optimizer.step() - torch.testing.assert_close(parameter, parameter_after_failure, rtol=0, atol=0) - - monkeypatch.undo() - optimizer.load_state_dict(checkpoint) - assert optimizer.state_offload_active - assert not optimizer.state_offload_poisoned - assert optimizer.optimizer_contract().capabilities.state_offload - _assert_cpu_boundary(optimizer) - - -@_CUDA_REQUIRED -def test_active_offload_blocks_rebinding_and_portable_global_state(): - parameter = torch.nn.Parameter(torch.ones(8, device="cuda")) - optimizer = _make_gefen(parameter) - del optimizer._resolve_automatic_period - identity = ParameterIdentity("Model.Weight", (8,)) - group = ProcessGroupIdentity("singleton", ("rank:0",)) - shard = ShardIdentity( - identity, - ParameterLayout.REPLICATED, - LogicalSlice.full(identity), - placements=(ShardPlacement("checkpoint", PlacementKind.REPLICATE, 0, 1),), - process_group=group, - local_member="rank:0", - ) - manifest = ShardingManifest((shard,)) - codebook_binding = CodebookProcessGroupBinding( - group, - "rank:0", - None, - parameter.device, - ) - checkpoint_binding = CheckpointProcessGroupBinding( - group, - "rank:0", - None, - parameter.device, - ) - - optimizer.offload_state_() - with pytest.raises(RuntimeError, match="restore state first"): - optimizer.post_sharding( - (ParameterRebinding(parameter, parameter, shard),), - manifest=manifest, - codebook_process_group=codebook_binding, - ) - assert not optimizer._canonical_identity_ready() - - optimizer.restore_state_() - optimizer.post_sharding( - (ParameterRebinding(parameter, parameter, shard),), - manifest=manifest, - codebook_process_group=codebook_binding, - ) - optimizer.offload_state_() - assert all( - support.transport is not CheckpointTransport.CANONICAL_GLOBAL - for support in optimizer.optimizer_contract().capabilities.checkpoints - ) - with pytest.raises(RuntimeError, match="active optimizer-state offload"): - optimizer.export_portable_state( - checkpoint_process_group=checkpoint_binding, - transaction_id="active-offload-reject-v1", - limits=PortableStateLimits( - max_fragment_tensor_bytes=1 << 20, - max_collective_tensor_bytes=4 << 20, - max_collective_metadata_bytes=4 << 20, - max_metadata_bytes=1 << 20, - ), - ) - - -@_CUDA_REQUIRED -def test_state_offload_rejects_persistent_aliases_and_multi_member_scope(): - parameter = torch.nn.Parameter(torch.linspace(-0.5, 0.5, 8, device="cuda")) - optimizer = _make_gefen(parameter) - parameter.grad = torch.linspace(-1.0, 1.0, 8, device="cuda") - optimizer.step() - parameter_state = optimizer.state[parameter] - parameter_state["vmean"] = parameter_state["m_magnitude"] - state_before = optimizer.state - - assert not optimizer.optimizer_contract().capabilities.state_offload - with pytest.raises(RuntimeError, match="storage aliases"): - optimizer.offload_state_() - assert optimizer.state is state_before - assert parameter_state["vmean"] is parameter_state["m_magnitude"] - - parameter_state["vmean"] = parameter_state["m_magnitude"].clone() - group = ProcessGroupIdentity("multi", ("rank:0", "rank:1")) - optimizer._gefen_codebook_process_group = CodebookProcessGroupBinding( - group, - "rank:0", - object(), - parameter.device, - ) - with pytest.raises(RuntimeError, match="multi-member"): - optimizer.offload_state_() - - -@_CUDA_REQUIRED -@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="two CUDA devices are required") -def test_state_offload_checks_capture_on_every_parameter_device(monkeypatch): - first = torch.nn.Parameter(torch.ones(8, device="cuda:0")) - second = torch.nn.Parameter(torch.ones(8, device="cuda:1")) - optimizer = Gefen( - [("first", first), ("second", second)], - fused=False, - factored_v_2d=False, - ) - monkeypatch.setattr( - torch.cuda, - "is_current_stream_capturing", - lambda: torch.cuda.current_device() == second.device.index, - ) - - assert not optimizer.optimizer_contract().capabilities.atomic_state_movement - with pytest.raises(RuntimeError, match="CUDA graph capture"): - optimizer.move_state_() - with pytest.raises(RuntimeError, match="CUDA graph capture"): - optimizer.offload_state_() - assert not optimizer.state_offload_active - - monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False) - optimizer.offload_state_() - monkeypatch.setattr( - torch.cuda, - "is_current_stream_capturing", - lambda: torch.cuda.current_device() == second.device.index, - ) - with pytest.raises(RuntimeError, match="CUDA graph capture"): - optimizer.restore_state_() - assert optimizer.state_offload_active - - -@_CUDA_REQUIRED -def test_operation_error_copies_state_back_without_poisoning(monkeypatch): - parameter = torch.nn.Parameter(torch.linspace(-0.4, 0.7, 8, device="cuda")) - optimizer = _make_gefen(parameter) - parameter.grad = torch.linspace(-1.0, 1.0, 8, device="cuda") - optimizer.step() - optimizer.offload_state_() - original = optimizer._step_automatic - - def update_then_fail(group, name, live_parameter, grad, *, state=None): - original(group, name, live_parameter, grad, state=state) - raise RuntimeError("injected update failure") - - monkeypatch.setattr(optimizer, "_step_automatic", update_then_fail) - parameter.grad = torch.linspace(0.8, -0.6, 8, device="cuda") - with pytest.raises(RuntimeError, match="injected update"): - optimizer.step() - - assert not optimizer.state_offload_poisoned - _assert_cpu_boundary(optimizer) - assert optimizer.state[parameter]["step"] == 2 - - -@_CUDA_REQUIRED -def test_custom_dtensor_like_and_finalized_nonreplicated_state_fail_closed(): - custom_parameter = torch.nn.Parameter(torch.ones(8, device="cuda")) - custom = _make_gefen(custom_parameter) - custom.state[custom_parameter]["extension"] = {"value": 1} - state_before = custom.state - with pytest.raises(RuntimeError, match="custom per-parameter state"): - custom.offload_state_() - assert custom.state is state_before - assert not custom.state_offload_active - - dtensor_like_parameter = torch.nn.Parameter(torch.ones(8, device="cuda")) - dtensor_like_parameter.to_local = lambda: dtensor_like_parameter - dtensor_like_parameter.placements = () - dtensor_like = _make_gefen(dtensor_like_parameter) - with pytest.raises(RuntimeError, match="ordinary replicated CUDA"): - dtensor_like.offload_state_() - - flat_parameter = torch.nn.Parameter(torch.ones(8, device="cuda")) - flattened = _make_gefen(flat_parameter) - identity = ParameterIdentity("layer.weight", (16,)) - process_group = ProcessGroupIdentity("data_parallel", ("rank:0", "rank:1")) - shards = tuple( - ShardIdentity( - identity, - ParameterLayout.FLATTENED_ELEMENT_SHARD, - LogicalSlice(index * 8, 8), - placements=( - ShardPlacement( - "data_parallel", - PlacementKind.FLAT_SHARD, - index, - 2, - ), - ), - process_group=process_group, - local_member=member, - ) - for index, member in enumerate(process_group.ordered_members) - ) - shard = shards[0] - flattened.rebind_shard( - flat_parameter, - flat_parameter, - shard=shard, - manifest=ShardingManifest(shards), - ) - with pytest.raises(RuntimeError, match="finalized replicated"): - flattened.offload_state_() - - -@_CUDA_REQUIRED -def test_capturable_compile_and_capture_states_fail_closed(monkeypatch): - parameter = torch.nn.Parameter(torch.ones(8, device="cuda")) - capturable = Gefen( - [("layer.weight", parameter)], - fused=False, - factored_v_2d=False, - capturable=True, - ) - with pytest.raises(RuntimeError, match="capturable"): - capturable.offload_state_() - - optimizer = _make_gefen(torch.nn.Parameter(torch.ones(8, device="cuda"))) - monkeypatch.setattr(torch.compiler, "is_compiling", lambda: True) - with pytest.raises(RuntimeError, match="torch.compile"): - optimizer.offload_state_() - monkeypatch.setattr(torch.compiler, "is_compiling", lambda: False) - monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: True) - with pytest.raises(RuntimeError, match="CUDA graph capture"): - optimizer.offload_state_() - - -def _replicated_shard(fqn, shape): - identity = ParameterIdentity(fqn, shape) - return ShardIdentity( - identity, - ParameterLayout.REPLICATED, - LogicalSlice.full(identity), - ) - - -def _finalized_replicated_cuda_optimizer(count=3): - parameters = [ - torch.nn.Parameter(torch.full((4,), float(index + 1), device="cuda")) - for index in range(count) - ] - optimizer = Gefen( - [ - ("weight{}".format(index), parameter) - for index, parameter in enumerate(parameters) - ], - fused=False, - factored_v_2d=False, - ) - optimizer._resolve_automatic_period = lambda *args, **kwargs: 4 - shards = tuple( - _replicated_shard("Model.Weight{}".format(index), (4,)) - for index in range(count) - ) - optimizer.post_sharding( - tuple( - ParameterRebinding(parameter, parameter, shard) - for parameter, shard in zip(parameters, shards) - ), - manifest=ShardingManifest(shards), - ) - return optimizer, parameters - - -def _step_offload(optimizer, parameters): - for parameter in parameters: - parameter.grad = torch.full_like(parameter, 0.5) - optimizer.step() - - -def _count_full_layout_passes(monkeypatch): - calls = {"count": 0} - original = Gefen._finalized_binding_layout_matches_full - - def counted(self): - calls["count"] += 1 - return original(self) - - monkeypatch.setattr(Gefen, "_finalized_binding_layout_matches_full", counted) - return calls - - -def _count_manifest_digest_computes(monkeypatch): - calls = {"count": 0} - original = Gefen._compute_codebook_manifest_fingerprint - - def counted(self, manifest): - calls["count"] += 1 - return original(self, manifest) - - monkeypatch.setattr(Gefen, "_compute_codebook_manifest_fingerprint", counted) - return calls - - -@_CUDA_REQUIRED -def test_offload_steady_state_reuses_cached_layout_forensics(monkeypatch): - # The finalized layout is immutable across steps, so the per-step offload - # readiness path (called ~2x/step) must reuse the memoized layout-forensics - # verdict instead of rebuilding the full layout + recomputing the manifest - # digest on every step. The per-tensor offload scan still runs every step; - # only the immutable layout forensics is cached. - optimizer, parameters = _finalized_replicated_cuda_optimizer() - optimizer.offload_state_() - - layout = _count_full_layout_passes(monkeypatch) - digest = _count_manifest_digest_computes(monkeypatch) - - # The first offloaded step still fully validates the layout once, because - # offload_state_ invalidated the cached verdict; that single full pass - # recomputes the manifest digest once. - _step_offload(optimizer, parameters) - first_layout = layout["count"] - first_digest = digest["count"] - assert first_layout >= 1 - assert first_digest >= 1 - - # Steady state: neither the full layout rebuild nor the manifest digest is - # recomputed again. On the pre-fix code each of the two per-step readiness - # calls forced a full rebuild + digest recompute, so these counts grew by - # four per step. - for _ in range(3): - _step_offload(optimizer, parameters) - assert layout["count"] == first_layout - assert digest["count"] == first_digest - - -@_CUDA_REQUIRED -def test_offload_layout_change_is_still_caught_with_a_warm_verdict(): - optimizer, parameters = _finalized_replicated_cuda_optimizer() - optimizer.offload_state_() - _step_offload(optimizer, parameters) # warm the cached layout verdict - - # Replace a finalized registry container so the fast-path tokens diverge: - # the offload readiness path must fall back to the full forensic rebuild and - # reject the step before any parameter is staged or mutated. - first_before = parameters[0].detach().clone() - optimizer._gefen_local_shard_bindings = tuple( - reversed(optimizer._gefen_local_shard_bindings) - ) - for parameter in parameters: - parameter.grad = torch.full_like(parameter, 0.5) - with pytest.raises(RuntimeError, match="cannot step"): - optimizer.step() - assert torch.equal(parameters[0].detach(), first_before) From a1bd0d133d16d52e30f55b6576ebfd342a36b5a2 Mon Sep 17 00:00:00 2001 From: thad0ctor Date: Wed, 15 Jul 2026 15:02:20 -0700 Subject: [PATCH 50/52] Harden convention codebook-scope preambles and contract validation (#76) * Harden convention codebook-scope preambles and contract validation Address a round of Codex + CodeRabbit findings on the slimmed optimizer convention (PR #67). All findings are convention-specific. Codebook-scope collective symmetry: - initialize_codebook() now captures the failure-vote binding and synchronizes a one-sided preamble failure (finalized layout / runtime binding / capture-readiness) through the scope before the scoped "initialize" operation header, so a member that fails its preamble no longer strands peers inside the header all_gather. - plain Gefen.step() rejects a closure that replaces or clears the runtime codebook binding between capture and the header: the captured-binding recheck raises and is synchronized through the captured scope, so every rank follows the same header collective instead of one rank skipping it. Checkpoint and contract validation: - CheckpointProcessGroupBinding no longer rejects a CUDA collective device on Gloo/MPI (both support CUDA tensors); such a device now falls through to the existing device-availability check. - Contract dataclasses enforce exact runtime types: StateField.checkpointed, StateVariant.initialized/migration_only, and CheckpointSupport.requires_collective/atomic_load must be real bools, and OptimizerContract.schema_version must be a real int, so truthy strings or bool/float look-alikes can no longer advertise false guarantees. - The hybrid finalized-layout fast token folds in the _state_param_owner registry contents (keys and values), so an in-place owner replacement that preserves the dict identity and length invalidates the cached verdict. Tests: - Snapshot per-device cache membership so a cache clear/removal is detected. - Verify the backup parameter is untouched in hybrid scoped-failure tests. - Guard the remaining Gloo-only distributed tests with the dist/Gloo skipif. - Documented that only ordinary nested GefenMuonHybrid.load_state_dict() reports atomic_load=False; portable composite imports keep the guarantee. - New gloo regression tests for the initialize-preamble and step group-swap synchronization, plus unit tests for the checkpoint device, contract type, and hybrid owner-token fixes. * Keep MPI checkpoint bindings CPU-only in collective-device validation Narrow the F4 relaxation: Gloo genuinely supports CUDA tensors, but MPI moves GPU tensors only when built CUDA-aware, which PyTorch cannot reliably detect. Reject a CUDA collective device on the MPI backend during validate_runtime() instead of deferring a backend error to the portable collective's all_gather/broadcast (Codex P2 on #76). --- docs/optimizer_contracts.md | 2 +- src/gefen/checkpoint.py | 8 +- src/gefen/contracts.py | 13 +- src/gefen/gefen.py | 34 ++++- src/gefen/hybrid.py | 8 + tests/_state_snapshot.py | 38 +++++ tests/test_checkpoint_binding.py | 63 +++++++- tests/test_codebook_scope_cpu.py | 17 +++ tests/test_codebook_scope_distributed.py | 151 +++++++++++++++++++ tests/test_hybrid_layout_cache.py | 22 +++ tests/test_hybrid_scoped_failure_protocol.py | 17 ++- tests/test_optimizer_contracts.py | 40 +++++ 12 files changed, 391 insertions(+), 22 deletions(-) diff --git a/docs/optimizer_contracts.md b/docs/optimizer_contracts.md index 62f9403..464c9d8 100644 --- a/docs/optimizer_contracts.md +++ b/docs/optimizer_contracts.md @@ -130,7 +130,7 @@ Portable v3 currently excludes non-period-one initialized state, second-moment r Plain Gefen declares replicated, flattened element-shard, and the narrow DTensor training layouts. Its PyTorch rank-local checkpoint transport is same-topology only. GefenMuon declares replicated and narrow DTensor training, with mode-specific state extents: `approx` state is local, `exact` state is logically global, and `distributed` momentum is held by the parameter owner while non-owners retain metadata only. Native Parallel-Muon checkpoints separately declare world-size owner redistribution, not placement-changing resharding. `GefenMuonHybrid` retains its nested child namespaces and does not flatten AdamW or Gefen child state into a fabricated common schema. -Gefen and GefenMuon native loads, rank-local payload restoration through `load_state_dict`, and distributed-owner payload restoration prepare their complete core restore before changing local live optimizer state. Their transport entries report `atomic_load=True`. This is a per-optimizer-instance fail-before-mutation guarantee at the optimizer load boundary, not a coordinated all-rank commit or a guarantee over work an external checkpoint orchestrator performs before calling the optimizer. Load pre-hooks run before that boundary and load post-hooks run afterward, so arbitrary side effects in user hooks are also outside the guarantee. Hybrid composite loads do not yet provide the same guarantee and report `atomic_load=False`. +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. Ordinary nested `GefenMuonHybrid.load_state_dict()` loads do not yet provide the same guarantee and report `atomic_load=False`; portable composite imports retain the separate fail-before-mutation guarantee described above. ## Adapter requirements diff --git a/src/gefen/checkpoint.py b/src/gefen/checkpoint.py index 945d8f3..6779543 100644 --- a/src/gefen/checkpoint.py +++ b/src/gefen/checkpoint.py @@ -122,11 +122,17 @@ def _validate_collective_device(self, backend: object) -> None: raise ValueError( "checkpoint collective device is incompatible with the runtime backend" ) - elif "gloo" in backend_name or "mpi" in backend_name: + elif "mpi" in backend_name: + # MPI moves GPU tensors only when built CUDA-aware, which PyTorch + # cannot reliably detect at runtime. Keep MPI CPU-only so a CUDA + # binding is rejected here rather than deferring a backend error to + # the later portable-collective all_gather/broadcast. if self.collective_device.type != "cpu": raise ValueError( "checkpoint collective device is incompatible with the runtime backend" ) + # Gloo supports CUDA tensors in addition to CPU, so a CUDA collective + # device on Gloo still falls through to the availability check below. if self.collective_device.type == "cuda": index = self.collective_device.index diff --git a/src/gefen/contracts.py b/src/gefen/contracts.py index 7b5159f..cb20439 100644 --- a/src/gefen/contracts.py +++ b/src/gefen/contracts.py @@ -798,6 +798,8 @@ def __post_init__(self) -> None: raise TypeError("StateField.geometry must be a StateGeometry") if not isinstance(self.key_match, StateKeyMatch): raise TypeError("StateField.key_match must be a StateKeyMatch") + if type(self.checkpointed) is not bool: + raise TypeError("StateField.checkpointed must be a bool") if type(self.optional) is not bool: raise TypeError("StateField.optional must be a bool") @@ -853,6 +855,9 @@ def __post_init__(self) -> None: raise TypeError("StateVariant.extent must be a StateExtent") if not isinstance(self.role, ParameterStateRole): raise TypeError("StateVariant.role must be a ParameterStateRole") + for name in ("initialized", "migration_only"): + if type(getattr(self, name)) is not bool: + raise TypeError("StateVariant.{} must be a bool".format(name)) if not set(self.inactive_fields).issubset(self.fields): raise ValueError("StateVariant.inactive_fields must be present in fields") if self.parameter_ranks is not None and set(self.parameter_ranks) & set( @@ -991,6 +996,9 @@ def __post_init__(self) -> None: raise TypeError( "CheckpointSupport.process_group_scope must be a ProcessGroupScope" ) + for name in ("requires_collective", "atomic_load"): + if type(getattr(self, name)) is not bool: + raise TypeError("CheckpointSupport.{} must be a bool".format(name)) if bool(self.topology_changing) != bool(self.topology_change_kinds): raise ValueError( "topology-changing layouts and change kinds must be declared together" @@ -1094,7 +1102,10 @@ def __post_init__(self) -> None: object.__setattr__(self, "children", _tuple(self.children)) if not self.implementation: raise ValueError("OptimizerContract.implementation must be non-empty") - if self.schema_version != CONTRACT_SCHEMA_VERSION: + if ( + type(self.schema_version) is not int + or self.schema_version != CONTRACT_SCHEMA_VERSION + ): raise ValueError( "unsupported optimizer contract schema version: {}".format( self.schema_version diff --git a/src/gefen/gefen.py b/src/gefen/gefen.py index 82de0ab..088d74a 100644 --- a/src/gefen/gefen.py +++ b/src/gefen/gefen.py @@ -5166,9 +5166,25 @@ def _maybe_refresh_gefen_codebook(self) -> None: def initialize_codebook(self) -> bool: """Collectively initialize the learned codebook without taking a step.""" - self._assert_finalized_binding_layout(full=True) - self._assert_runtime_codebook_process_group() - self._assert_codebook_capture_ready() + # Capture a usable failure-vote binding before inspecting live layout so + # a one-sided preamble failure -- for example capture-readiness raising + # only on the gradient-owning member while non-owners proceed -- is + # reported through the scope instead of stranding peers inside the + # scoped operation-header collective below. + scope_binding = self._capture_codebook_scope_binding_for_step() + try: + self._assert_finalized_binding_layout(full=True) + self._assert_runtime_codebook_process_group() + self._assert_codebook_capture_ready() + local_preamble_error = None + except Exception as exc: + local_preamble_error = exc + if scope_binding is not None: + self._synchronize_prevalidated_codebook_scope_failure( + local_preamble_error, "initialize preamble", scope_binding + ) + elif local_preamble_error is not None: + raise local_preamble_error self._validate_codebook_scope_operation_header("initialize") try: _assert_optimizer_gradients_structurally_valid( @@ -8557,9 +8573,17 @@ def step(self, closure=None): raise local_preamble_error # The closure can replace a finalized parameter or otherwise invalidate - # the runtime binding. Recheck before the operation header and synchronize - # structural failures before any peer enters a scoped codebook collective. + # the runtime binding. It can also rebind (or clear) the runtime + # process-group between capture and the operation header; a rank that + # silently swapped to None or a different binding would enter a + # different header collective than its peers and deadlock. Recheck + # against the captured binding and synchronize structural failures + # before any peer enters a scoped codebook collective. try: + if self._gefen_codebook_process_group is not scope_binding: + raise RuntimeError( + "Gefen codebook process-group binding changed during the step preamble" + ) self._assert_finalized_binding_layout() self._assert_runtime_codebook_process_group() _assert_optimizer_gradients_structurally_valid(self) diff --git a/src/gefen/hybrid.py b/src/gefen/hybrid.py index 04f7db2..074b5c5 100644 --- a/src/gefen/hybrid.py +++ b/src/gefen/hybrid.py @@ -997,6 +997,14 @@ def _hybrid_layout_forensics_fast_tokens(self): self._subopts, len(self._subopts), ] + # Capture the owner registry contents by identity (keys, then the + # (parameter, child) values) so an in-place entry replacement that + # preserves the dict object and its length -- which the identity/len + # tokens above cannot see -- flips a token and forces a full rebuild + # instead of reusing the stale verdict. Mirrors the _param_names + # snapshot in the base fast-token path. + live.extend(self._state_param_owner.keys()) + live.extend(self._state_param_owner.values()) for child in self._subopts: live.append(child._finalized_binding_layout_matches()) GefenMuonHybrid._hybrid_child_param_group_tokens(child, live) diff --git a/tests/_state_snapshot.py b/tests/_state_snapshot.py index 7bb0115..9a58c47 100644 --- a/tests/_state_snapshot.py +++ b/tests/_state_snapshot.py @@ -34,6 +34,21 @@ _CHILD_REGISTRY_ATTRS = ("_param_names", "_gefen_shard_bindings") +# Per-device cache dicts each Gefen/GefenMuon child keeps live across steps. +# Their tensors are also reachable from ``__dict__``, but the tensor-value +# snapshot keeps a strong reference to each tensor, so an in-place ``clear()`` or +# entry removal leaves every retained tensor still matching its clone while the +# cache membership silently vanishes. Recording the dict identity plus its keyed +# contents makes that membership change fail the assertion. Absent on optimizer +# types that never build these caches (skipped below). +_CHILD_CACHE_ATTRS = ( + "_gefen_codebook_by_device", + "_gefen_codebook_lut_by_device", + "_sr_seed_by_device", + "_gefen_global_step_by_device", +) + + def _cloned(value): if torch.is_tensor(value): return value.detach().clone() @@ -133,10 +148,29 @@ def _registry_snapshot(optimizer): return registries +def _cache_snapshot(optimizer): + """Live dict identity plus cloned contents of each per-device cache. + + Returns ``name -> (live_dict, cloned_contents)`` for every cache present as a + ``dict`` on ``optimizer``. The live dict is kept by reference so the assertion + can confirm it was not replaced, and its contents are cloned so a later + ``clear()`` or entry removal (or an in-place value edit) produces a mismatch. + """ + + caches = {} + for name in _CHILD_CACHE_ATTRS: + cache = getattr(optimizer, name, None) + if type(cache) is not dict: + continue + caches[name] = (cache, _cloned(cache)) + return caches + + def deep_state_snapshot(optimizer): return { "attributes": optimizer.__dict__.copy(), "registries": _registry_snapshot(optimizer), + "caches": _cache_snapshot(optimizer), "state": optimizer.state, "state_items": tuple( (parameter, state, _cloned(dict(state))) @@ -194,3 +228,7 @@ def assert_deep_state_snapshot(optimizer, snapshot): ): assert live_key is expected_key _nested_equal(live_value, expected_value) + for name, (expected_ref, expected_contents) in snapshot["caches"].items(): + live = getattr(optimizer, name, None) + assert live is expected_ref + _nested_equal(live, expected_contents) diff --git a/tests/test_checkpoint_binding.py b/tests/test_checkpoint_binding.py index e1c0e49..510d1ff 100644 --- a/tests/test_checkpoint_binding.py +++ b/tests/test_checkpoint_binding.py @@ -46,6 +46,48 @@ def test_checkpoint_binding_validates_descriptor_types_membership_and_device(): CheckpointProcessGroupBinding(identity, "worker:0", None, torch.device("meta")) +def test_validate_collective_device_allows_gloo_cuda_but_requires_available_device(): + identity = ProcessGroupIdentity("local", ("worker:0",)) + unavailable_index = ( + torch.cuda.device_count() if torch.cuda.is_available() else 0 + ) + cuda_binding = CheckpointProcessGroupBinding( + identity, "worker:0", None, torch.device("cuda", unavailable_index) + ) + # Gloo supports CUDA tensors, so a CUDA collective device is not rejected on + # backend grounds; it must instead fail the later availability check when the + # configured device does not exist. + with pytest.raises(ValueError, match="CUDA device is unavailable"): + cuda_binding._validate_collective_device("gloo") + + cpu_binding = CheckpointProcessGroupBinding( + identity, "worker:0", None, torch.device("cpu") + ) + # NCCL still requires a CUDA collective device. + with pytest.raises(ValueError, match="runtime backend"): + cpu_binding._validate_collective_device("nccl") + # Gloo with a CPU device remains valid. + cpu_binding._validate_collective_device("gloo") + + +def test_validate_collective_device_keeps_mpi_cpu_only(): + identity = ProcessGroupIdentity("local", ("worker:0",)) + # MPI moves CUDA tensors only when built CUDA-aware, which PyTorch cannot + # reliably detect, so a CUDA collective device must be rejected on backend + # grounds rather than deferring a backend error to the later collective. + cuda_binding = CheckpointProcessGroupBinding( + identity, "worker:0", None, torch.device("cuda", 0) + ) + with pytest.raises(ValueError, match="runtime backend"): + cuda_binding._validate_collective_device("mpi") + + # MPI with a CPU device remains valid. + cpu_binding = CheckpointProcessGroupBinding( + identity, "worker:0", None, torch.device("cpu") + ) + cpu_binding._validate_collective_device("mpi") + + def test_checkpoint_binding_requires_explicit_multi_member_handle_and_local_singleton(): singleton = ProcessGroupIdentity("local", ("worker:0",)) multiple = ProcessGroupIdentity("data", ("worker:0", "worker:1")) @@ -111,16 +153,23 @@ def _distributed_binding_worker(rank, world_size, init_file, queue): "world size", ) - wrong_device_binding = CheckpointProcessGroupBinding( + # Gloo accepts CUDA tensors, so a CUDA collective device is not rejected + # on backend grounds; it must instead fail the later availability check + # when the configured device does not exist. An index at or beyond the + # visible device count (0 when no CUDA is present) is always unavailable. + unavailable_index = ( + torch.cuda.device_count() if torch.cuda.is_available() else 0 + ) + unavailable_device_binding = CheckpointProcessGroupBinding( world_identity, world_members[rank], dist.group.WORLD, - torch.device("cuda:0"), + torch.device("cuda", unavailable_index), ) - device_mismatch_rejected = _expect_rejection( - wrong_device_binding.validate_runtime, + cuda_unavailable_rejected = _expect_rejection( + unavailable_device_binding.validate_runtime, ValueError, - "runtime backend", + "CUDA device is unavailable", ) subgroup_global_ranks = (0, 2) @@ -162,7 +211,7 @@ def _distributed_binding_worker(rank, world_size, init_file, queue): "world_validated": world_validated, "order_mismatch_rejected": order_mismatch_rejected, "size_mismatch_rejected": size_mismatch_rejected, - "device_mismatch_rejected": device_mismatch_rejected, + "cuda_unavailable_rejected": cuda_unavailable_rejected, "subgroup_validated": subgroup_validated, "nonmember_rejected": nonmember_rejected, "uninitialized_rejected": uninitialized_rejected, @@ -221,7 +270,7 @@ def test_checkpoint_binding_validates_real_world_and_subgroup_membership(): assert all(item["world_validated"] for item in results) assert all(item["order_mismatch_rejected"] for item in results) assert all(item["size_mismatch_rejected"] for item in results) - assert all(item["device_mismatch_rejected"] for item in results) + assert all(item["cuda_unavailable_rejected"] for item in results) assert all(item["subgroup_validated"] for item in results if item["rank"] in (0, 2)) assert results[1]["nonmember_rejected"] assert all(item["uninitialized_rejected"] for item in results) diff --git a/tests/test_codebook_scope_cpu.py b/tests/test_codebook_scope_cpu.py index b427d4a..9f1dec8 100644 --- a/tests/test_codebook_scope_cpu.py +++ b/tests/test_codebook_scope_cpu.py @@ -135,6 +135,23 @@ def _replace_native_guard(checkpoint, guard): group["_gefen_checkpoint_metadata"]["native_local_shards"] = copy.deepcopy(guard) +def test_deep_snapshot_detects_per_device_cache_membership_clear(): + # A per-device cache clear preserves the dict attribute identity, and the + # retained cache tensors still match their tensor-value clones, so only a + # membership-aware snapshot catches the removal. + parameter = torch.nn.Parameter(torch.randn(4, 4)) + optimizer = Gefen([("Layer.Weight", parameter)], fused=False) + device = torch.device("cpu") + optimizer._gefen_codebook_by_device[device] = torch.randn(8) + + snapshot = deep_state_snapshot(optimizer) + assert_deep_state_snapshot(optimizer, snapshot) # unchanged membership passes + + optimizer._gefen_codebook_by_device.clear() + with pytest.raises(AssertionError): + assert_deep_state_snapshot(optimizer, snapshot) + + def test_codebook_process_group_binding_is_public_frozen_and_ordered(): group = ProcessGroupIdentity("replica", ("worker:b", "worker:a")) binding = CodebookProcessGroupBinding(group, "worker:b", object(), torch.device("cpu")) diff --git a/tests/test_codebook_scope_distributed.py b/tests/test_codebook_scope_distributed.py index b55c458..94cc68a 100644 --- a/tests/test_codebook_scope_distributed.py +++ b/tests/test_codebook_scope_distributed.py @@ -783,6 +783,10 @@ def _run_workers(world=2): os.unlink(init_file) +@pytest.mark.skipif( + not dist.is_available() or not dist.is_gloo_available(), + reason="explicit Gloo scope coverage requires Gloo", +) def test_explicit_gloo_scope_aggregates_logical_state_and_fails_atomically(): results = _run_workers() @@ -924,6 +928,10 @@ def _run_subgroup_workers(): os.unlink(init_file) +@pytest.mark.skipif( + not dist.is_available() or not dist.is_gloo_available(), + reason="explicit Gloo subgroup coverage requires Gloo", +) def test_explicit_gloo_subgroups_are_isolated_from_default_world(): results = _run_subgroup_workers() @@ -1117,6 +1125,10 @@ def _run_zero_length_flat_workers(world=2): os.unlink(init_file) +@pytest.mark.skipif( + not dist.is_available() or not dist.is_gloo_available(), + reason="zero-length flattened shard coverage requires Gloo", +) def test_zero_length_flattened_shard_member_joins_every_scoped_collective(): results = _run_zero_length_flat_workers() @@ -1294,6 +1306,12 @@ def closure(): if rank == 0 and failure_mode == "replace": optimizer.param_groups[0]["params"][0] = rogue rogue.grad = torch.ones_like(rogue) + if rank == 0 and failure_mode == "swap_group": + # Clear the runtime codebook binding after it was captured. The + # local preamble still succeeds, so without a captured-binding + # recheck this rank would skip the scoped step header while its + # peer entered the all_gather and hung. + optimizer._gefen_codebook_process_group = None return torch.tensor(1.0) try: @@ -1419,3 +1437,136 @@ def test_scoped_step_entry_layout_mutation_raises_symmetrically_across_the_scope assert "changed outside post_sharding" in results[0]["message"] assert "step preamble failed on another process-group member" in results[1]["message"] assert all(item["untouched"] for item in results), results + + +@pytest.mark.skipif( + not dist.is_available() or not dist.is_gloo_available(), + reason="plain Gefen scoped closure group-swap coverage requires Gloo", +) +def test_plain_gefen_scoped_step_closure_group_swap_raises_symmetrically_across_the_scope(): + # A closure that clears the captured runtime binding on one rank must not let + # that rank skip the scoped step header while its peer enters the all_gather. + # The captured-binding recheck raises on the swapping rank and the failure is + # synchronized through the captured scope so BOTH ranks raise fast. + results = _run_closure_preamble_workers("gefen", "swap_group") + assert len(results) == 2, results + assert all("error" not in item for item in results), results + assert results[0]["message"] is not None and results[1]["message"] is not None, results + assert "binding changed during the step preamble" in results[0]["message"] + assert ( + "gradient preflight failed on another process-group member" + in results[1]["message"] + ) + assert all(item["untouched"] for item in results), results + + +def _initialize_preamble_worker(rank, world, init_file, queue): + # initialize_codebook() runs its finalized-layout / runtime-binding / + # capture-readiness preamble before the scoped operation-header collective. + # A rank-local preamble failure must raise on every scope member together + # instead of leaving the failing rank to exit while the peer enters the + # scoped "initialize" all_gather and hangs. + try: + dist.init_process_group( + "gloo", + init_method="file://{}".format(init_file), + rank=rank, + world_size=world, + timeout=timedelta(seconds=45), + ) + members = tuple("rank:{}".format(index) for index in range(world)) + group = ProcessGroupIdentity("data_parallel", members) + runtime_group = dist.group.WORLD + + matrix = torch.nn.Parameter(torch.zeros(2, 2)) + optimizer = Gefen([("matrix", matrix)], fused=False, factored_v_2d=False) + identity = ParameterIdentity("Matrix", (2, 2)) + records = tuple(_replicated(identity, group, member) for member in members) + _finalize( + optimizer, + matrix, + records[rank], + ShardingManifest(records), + _binding(group, rank, runtime_group), + ) + matrix.grad = torch.tensor([[1.0, -2.0], [3.0, -4.0]]) + rogue = torch.nn.Parameter(torch.ones(2, 2)) + rogue_before = rogue.detach().clone() + if rank == 0: + # Break the finalized layout on one rank before the preamble runs. + optimizer.param_groups[0]["params"][0] = rogue + rogue.grad = torch.ones_like(rogue) + + try: + optimizer.initialize_codebook() + message = None + except RuntimeError as exc: + message = str(exc) + untouched = ( + optimizer._gefen_global_step == 0 + and optimizer._gefen_codebook is None + and torch.equal(rogue, rogue_before) + ) + queue.put({"rank": rank, "message": message, "untouched": untouched}) + except Exception as exc: + queue.put({"rank": rank, "error": repr(exc)}) + finally: + if dist.is_available() and dist.is_initialized(): + dist.destroy_process_group() + + +def _run_initialize_preamble_workers(world=2): + context = mp.get_context("spawn") + queue = context.Queue() + fd, init_file = tempfile.mkstemp(prefix="gefen-codebook-init-preamble-") + os.close(fd) + os.unlink(init_file) + processes = [ + context.Process( + target=_initialize_preamble_worker, + args=(rank, world, init_file, queue), + ) + for rank in range(world) + ] + try: + for process in processes: + process.start() + results = [] + try: + for _ in processes: + results.append(queue.get(timeout=120)) + except Exception: + pass + for process in processes: + process.join(timeout=10) + for process in processes: + if process.is_alive(): + process.terminate() + process.join(timeout=5) + pytest.fail("initialize-preamble worker hung") + assert process.exitcode == 0 + return sorted(results, key=lambda item: item["rank"]) + finally: + for process in processes: + if process.is_alive(): + process.terminate() + process.join(timeout=5) + if os.path.exists(init_file): + os.unlink(init_file) + + +@pytest.mark.skipif( + not dist.is_available() or not dist.is_gloo_available(), + reason="scoped initialize-preamble coverage requires Gloo", +) +def test_initialize_codebook_preamble_failure_raises_symmetrically_across_the_scope(): + results = _run_initialize_preamble_workers() + assert len(results) == 2, results + assert all("error" not in item for item in results), results + assert results[0]["message"] is not None and results[1]["message"] is not None, results + assert "changed outside post_sharding" in results[0]["message"] + assert ( + "initialize preamble failed on another process-group member" + in results[1]["message"] + ) + assert all(item["untouched"] for item in results), results diff --git a/tests/test_hybrid_layout_cache.py b/tests/test_hybrid_layout_cache.py index f6bbf2a..de06917 100644 --- a/tests/test_hybrid_layout_cache.py +++ b/tests/test_hybrid_layout_cache.py @@ -160,3 +160,25 @@ def test_composite_registry_replacement_is_detected_with_warm_verdict(): with pytest.raises(RuntimeError, match="finalized parameter layout changed"): _step_with_grads(optimizer, (matrix, bias)) + + +def test_composite_registry_in_place_entry_swap_is_detected_with_warm_verdict(): + optimizer, matrix, bias = _finalized_hybrid() + _step_with_grads(optimizer, (matrix, bias)) # warm the cached verdict + assert optimizer._hybrid_layout_forensics_verdict is not None + + # Replace ONE ownership entry in place, preserving the dict object identity + # and its length. The identity/len tokens cannot see this, so the fast token + # must fold in the owner-registry contents; otherwise the stale True verdict + # is reused and the mismatched owner routing slips past the guard. + key = next(iter(optimizer._state_param_owner)) + _parameter, child = optimizer._state_param_owner[key] + rogue = torch.nn.Parameter(torch.full((2, 2), 5.0)) + rogue_before = rogue.detach().clone() + optimizer._state_param_owner[key] = (rogue, child) + assert optimizer._state_param_owner[key] is not _parameter + assert len(optimizer._state_param_owner) == 2 + + with pytest.raises(RuntimeError, match="finalized parameter layout changed"): + _step_with_grads(optimizer, (matrix, bias)) + assert torch.equal(rogue, rogue_before) diff --git a/tests/test_hybrid_scoped_failure_protocol.py b/tests/test_hybrid_scoped_failure_protocol.py index 68e59a6..ab46808 100644 --- a/tests/test_hybrid_scoped_failure_protocol.py +++ b/tests/test_hybrid_scoped_failure_protocol.py @@ -255,7 +255,9 @@ def _set_local_gradients(muon_parameter, backup_parameter, backup_shard, scale=1 backup_parameter.grad = _backup_gradient()[start:stop] * scale -def _untouched(optimizer, muon_parameter, backup_parameter): +def _untouched(optimizer, muon_parameter, backup_parameter, backup_shard): + start = backup_shard.logical_slice.flat_offset + stop = start + backup_shard.logical_slice.length return ( optimizer.muon._gefen_global_step == 0 and optimizer.backup._gefen_global_step == 0 @@ -266,6 +268,7 @@ def _untouched(optimizer, muon_parameter, backup_parameter): muon_parameter is None or optimizer.muon.state[muon_parameter] == {"name": "matrix"} ) + and _bits_equal(backup_parameter, _backup_initial()[start:stop]) and optimizer.backup.state[backup_parameter] == {"name": "vector"} ) @@ -283,7 +286,7 @@ def _amp_divergent_overflow_result(rank, group): message = str(exc) return { "message": message, - "untouched": _untouched(optimizer, muon_parameter, backup_parameter), + "untouched": _untouched(optimizer, muon_parameter, backup_parameter, backup_shard), "grads_untouched": _bits_equal(backup_parameter.grad, grad_before), } @@ -306,7 +309,7 @@ def _amp_group_wide_overflow_result(rank, group): return { "skipped": skipped, "post_hook_fired": len(post_hook_calls) == 1, - "untouched": _untouched(optimizer, muon_parameter, backup_parameter), + "untouched": _untouched(optimizer, muon_parameter, backup_parameter, backup_shard), "grads_untouched": _bits_equal(backup_parameter.grad, grad_before), } @@ -350,7 +353,7 @@ def _preflight_divergent_result(rank, group): message = str(exc) return { "message": message, - "untouched": _untouched(optimizer, muon_parameter, backup_parameter), + "untouched": _untouched(optimizer, muon_parameter, backup_parameter, backup_shard), } @@ -374,7 +377,7 @@ def closure(): message = str(exc) return { "message": message, - "untouched": _untouched(optimizer, muon_parameter, backup_parameter), + "untouched": _untouched(optimizer, muon_parameter, backup_parameter, backup_shard), } @@ -400,7 +403,7 @@ def closure(): return { "message": message, "untouched": ( - _untouched(optimizer, muon_parameter, backup_parameter) + _untouched(optimizer, muon_parameter, backup_parameter, backup_shard) and torch.equal(rogue, rogue_before) ), } @@ -424,7 +427,7 @@ def _step_entry_layout_divergent_result(rank, group): return { "message": message, "untouched": ( - _untouched(optimizer, muon_parameter, backup_parameter) + _untouched(optimizer, muon_parameter, backup_parameter, backup_shard) and torch.equal(rogue, rogue_before) ), } diff --git a/tests/test_optimizer_contracts.py b/tests/test_optimizer_contracts.py index 4412acc..31351ff 100644 --- a/tests/test_optimizer_contracts.py +++ b/tests/test_optimizer_contracts.py @@ -753,6 +753,46 @@ def test_capabilities_reject_untyped_entries_and_flags(): replace(capabilities, **{name: "yes"}) +def test_contract_flags_and_schema_version_reject_non_exact_types(): + # Boolean flags must reject truthy non-bools (e.g. "false" reads truthy and + # would falsely advertise a guarantee); schema_version must reject bools and + # floats that merely compare equal to the supported integer version. + field = StateField("name", StateScope.PARAMETER, StateGeometry.OPAQUE, True) + with pytest.raises(TypeError, match="StateField.checkpointed must be a bool"): + replace(field, checkpointed="false") + + variant = StateVariant( + "name_only", + ("name",), + frozenset({ParameterLayout.REPLICATED}), + StateExtent.METADATA_ONLY, + initialized=False, + ) + for name in ("initialized", "migration_only"): + with pytest.raises( + TypeError, match="StateVariant.{} must be a bool".format(name) + ): + replace(variant, **{name: 1}) + + optimizer = Gefen([("parameter", torch.nn.Parameter(torch.ones(4)))], fused=False) + contract = optimizer.optimizer_contract() + checkpoint_support = contract.capabilities.checkpoints[0] + for name in ("requires_collective", "atomic_load"): + with pytest.raises( + TypeError, match="CheckpointSupport.{} must be a bool".format(name) + ): + replace(checkpoint_support, **{name: "false"}) + + for bad_version in (True, float(CONTRACT_SCHEMA_VERSION)): + with pytest.raises(ValueError, match="schema version"): + replace(contract, schema_version=bad_version) + # The exact supported integer version still validates. + assert ( + replace(contract, schema_version=CONTRACT_SCHEMA_VERSION).schema_version + == CONTRACT_SCHEMA_VERSION + ) + + def test_child_contract_rejects_untyped_contract_payload(): with pytest.raises( TypeError, match="contract must be an OptimizerContract or None" From 01fd99830823b8a8d3d4265d9bc600796afbbf6b Mon Sep 17 00:00:00 2001 From: thad0ctor Date: Wed, 15 Jul 2026 15:31:21 -0700 Subject: [PATCH 51/52] Address CodeRabbit review on the optimizer-contracts convention Mirror the plain-Gefen scoped step-preamble hardening into GefenMuon.step: recheck the runtime codebook process-group against the captured binding in the gradient preflight so a closure that clears or replaces it is caught and synchronized before any peer enters the scoped operation-header collective, instead of one rank silently skipping the header and stranding its peers. Add a GefenMuon group-swap regression (fails without the recheck: the swapping rank runs the step while its peer hangs in the all_gather). Fix a tautological assertion in the composite owner in-place-swap layout cache test: the registry value is a (parameter, child) tuple, so unpack and assert the replaced members instead of comparing the tuple to a parameter. Qualify the optimizer-contracts private-registry deferral note: in-place value replacement is deferred to a boundary only for registries not folded into the content/fast token; _state_param_owner contents do participate, so such a replacement is rejected by the next step guard. --- docs/optimizer_contracts.md | 2 +- src/gefen/gefen_muon.py | 11 +++++++++++ tests/test_codebook_scope_distributed.py | 23 +++++++++++++++++++++++ tests/test_hybrid_layout_cache.py | 5 ++++- 4 files changed, 39 insertions(+), 2 deletions(-) diff --git a/docs/optimizer_contracts.md b/docs/optimizer_contracts.md index 464c9d8..61dcf44 100644 --- a/docs/optimizer_contracts.md +++ b/docs/optimizer_contracts.md @@ -41,7 +41,7 @@ Rebinding is allowed only while the entire optimizer is pristine: global step ze Plain Gefen currently rebinds complete replicated parameters and contiguous physical 1-D flattened element shards. A flattened logical matrix requires `factored_v_2d=False` because canonical row/column factored-state projection is not implemented. GefenMuon rebinds complete replicated matrices and can finalize/prune whole-parameter owner manifests. A whole-owner instance may step only when the same transaction installs the explicit codebook process group described below; the owner performs the update and the machine-readable training declaration requires the adapter to synchronize the updated matrix afterward. A Gefen-backed `GefenMuonHybrid` atomically partitions one complete manifest and rebinding plan by its frozen exact FQN routing, stages both children, validates cross-child storage disjointness, rebuilds composite state routing, and publishes only after every child succeeds. AdamW-backed Hybrid and DTensor composite rebinding remain unsupported. The portable global-state path described below can reshard supported finalized layouts. -After finalization, every entry point re-validates the published layout, and this has two costs. Steps and identity queries take an O(local params) fast path: the first complete forensic rebuild caches a verdict keyed by cheap identity tokens — the finalized registries by object identity, every live group container, parameter, and compatibility name, and a version counter that every mutating API bumps — and the full rebuild re-runs only when one of those tokens changes. A fixed set of boundaries always runs the complete rebuild regardless of the cache: checkpoint save and load (`state_dict` / `load_state_dict`), canonical export, import prepare, and import commit; `post_sharding` rebinding; collective codebook initialize and refresh; codebook-scope re-validation; and external contract-readiness queries. `post_sharding` additionally computes the manifest shard set and its sha256 digest once per finalized manifest, for reuse by the scoped operation headers. The practical consequence for an integrator is a clean split: any layout corruption reachable through the public containers — the group `params` / `param_names` slots, per-parameter state names, or the compatibility-name cache — still fails the step guard before any state is mutated, including corruption a closure introduces between the pre- and post-closure guards. Only corruption that leaves every fast-path token intact — in-place value replacement inside the private finalized registries, or an `object.__setattr__` on a frozen identity record — waits until the next boundary above to be caught rather than being caught at the next step. +After finalization, every entry point re-validates the published layout, and this has two costs. Steps and identity queries take an O(local params) fast path: the first complete forensic rebuild caches a verdict keyed by cheap identity tokens — the finalized registries by object identity, every live group container, parameter, and compatibility name, and a version counter that every mutating API bumps — and the full rebuild re-runs only when one of those tokens changes. A fixed set of boundaries always runs the complete rebuild regardless of the cache: checkpoint save and load (`state_dict` / `load_state_dict`), canonical export, import prepare, and import commit; `post_sharding` rebinding; collective codebook initialize and refresh; codebook-scope re-validation; and external contract-readiness queries. `post_sharding` additionally computes the manifest shard set and its sha256 digest once per finalized manifest, for reuse by the scoped operation headers. The practical consequence for an integrator is a clean split: any layout corruption reachable through the public containers — the group `params` / `param_names` slots, per-parameter state names, or the compatibility-name cache — still fails the step guard before any state is mutated, including corruption a closure introduces between the pre- and post-closure guards. Only corruption that leaves every fast-path token intact — an `object.__setattr__` on a frozen identity record, or an in-place value replacement inside a private finalized registry whose contents are not folded into that optimizer's content/fast token — waits until the next boundary above to be caught rather than being caught at the next step. The composite owner registry `_state_param_owner` is not such a registry: its contents do participate in the fast token, so an in-place entry replacement there changes the token and is rejected by the very next step guard before any state mutation, rather than being deferred to a boundary. ## Explicit learned-codebook process groups diff --git a/src/gefen/gefen_muon.py b/src/gefen/gefen_muon.py index 0cef5ae..d794003 100644 --- a/src/gefen/gefen_muon.py +++ b/src/gefen/gefen_muon.py @@ -3116,7 +3116,18 @@ def step(self, closure=None): local_preamble_error, "step preamble", scope_binding ) + # The closure can replace a finalized parameter or otherwise + # invalidate the runtime binding. It can also rebind (or clear) the + # runtime process-group between capture and the operation header; a + # rank that silently swapped to None or a different binding would + # enter a different header collective than its peers and deadlock. + # Recheck against the captured binding and synchronize structural + # failures before any peer enters a scoped codebook collective. try: + if self._gefen_codebook_process_group is not scope_binding: + raise RuntimeError( + "GefenMuon codebook process-group binding changed during the step preamble" + ) self._assert_finalized_binding_layout() self._assert_runtime_codebook_process_group() _assert_optimizer_gradients_structurally_valid( diff --git a/tests/test_codebook_scope_distributed.py b/tests/test_codebook_scope_distributed.py index 94cc68a..bf3fb77 100644 --- a/tests/test_codebook_scope_distributed.py +++ b/tests/test_codebook_scope_distributed.py @@ -1460,6 +1460,29 @@ def test_plain_gefen_scoped_step_closure_group_swap_raises_symmetrically_across_ assert all(item["untouched"] for item in results), results +@pytest.mark.skipif( + not dist.is_available() or not dist.is_gloo_available(), + reason="GefenMuon scoped closure group-swap coverage requires Gloo", +) +def test_gefen_muon_scoped_step_closure_group_swap_raises_symmetrically_across_the_scope(): + # GefenMuon twin of the plain-Gefen group-swap coverage. A closure that + # clears the captured runtime binding on one rank keeps the rank-local + # preamble valid, so without the captured-binding recheck the swapping rank + # would skip the scoped step header while its peer entered the all_gather and + # hung. The recheck raises on the swapping rank and the failure is + # synchronized through the captured scope so BOTH ranks raise fast. + results = _run_closure_preamble_workers("muon", "swap_group") + assert len(results) == 2, results + assert all("error" not in item for item in results), results + assert results[0]["message"] is not None and results[1]["message"] is not None, results + assert "binding changed during the step preamble" in results[0]["message"] + assert ( + "gradient preflight failed on another process-group member" + in results[1]["message"] + ) + assert all(item["untouched"] for item in results), results + + def _initialize_preamble_worker(rank, world, init_file, queue): # initialize_codebook() runs its finalized-layout / runtime-binding / # capture-readiness preamble before the scoped operation-header collective. diff --git a/tests/test_hybrid_layout_cache.py b/tests/test_hybrid_layout_cache.py index de06917..e4d8fd7 100644 --- a/tests/test_hybrid_layout_cache.py +++ b/tests/test_hybrid_layout_cache.py @@ -176,7 +176,10 @@ def test_composite_registry_in_place_entry_swap_is_detected_with_warm_verdict(): rogue = torch.nn.Parameter(torch.full((2, 2), 5.0)) rogue_before = rogue.detach().clone() optimizer._state_param_owner[key] = (rogue, child) - assert optimizer._state_param_owner[key] is not _parameter + replaced_parameter, replaced_child = optimizer._state_param_owner[key] + assert replaced_parameter is rogue + assert replaced_parameter is not _parameter + assert replaced_child is child assert len(optimizer._state_param_owner) == 2 with pytest.raises(RuntimeError, match="finalized parameter layout changed"): From 97a63df7f8b2b49c9116908a6669966c614ee548 Mon Sep 17 00:00:00 2001 From: thad0ctor Date: Sat, 18 Jul 2026 07:11:20 -0700 Subject: [PATCH 52/52] Reconcile portable state with DCP name provenance --- src/gefen/gefen.py | 9 +++++++++ src/gefen/portable_runtime.py | 8 ++++++-- 2 files changed, 15 insertions(+), 2 deletions(-) diff --git a/src/gefen/gefen.py b/src/gefen/gefen.py index ee2dc17..d43f2c3 100644 --- a/src/gefen/gefen.py +++ b/src/gefen/gefen.py @@ -1444,6 +1444,7 @@ def _canonical_group_options_value(group): not in { "params", "param_names", + _SYNTHESIZED_PARAM_NAMES_KEY, "name", "_gefen_checkpoint_metadata", } @@ -2841,6 +2842,7 @@ def _stage_post_sharding(self, rebindings, manifest, codebook_process_group=None staged.param_groups = [] staged.state = defaultdict(dict) staged._param_names = {} + staged._synthesized_param_names = {} staged._gefen_shard_bindings = {} local_bindings = [] logical_slots = [] @@ -2850,6 +2852,7 @@ def _stage_post_sharding(self, rebindings, manifest, codebook_process_group=None staged_group.pop("_gefen_checkpoint_metadata", None) staged_params = [] staged_names = [] + staged_synthesized_names = [] names = list(group.get("param_names", ())) if len(names) != len(group["params"]): names = [self._param_name(param) for param in group["params"]] @@ -2873,10 +2876,16 @@ def _stage_post_sharding(self, rebindings, manifest, codebook_process_group=None staged_params.append(target) staged_names.append(compatibility_name) staged._param_names[target] = compatibility_name + synthesized_name = self._synthesized_param_names.get( + rebinding.old_parameter, True + ) + staged_synthesized_names.append(synthesized_name) + staged._synthesized_param_names[target] = synthesized_name staged.state[target]["name"] = compatibility_name staged._gefen_shard_bindings[target] = rebinding.shard staged_group["params"] = staged_params staged_group["param_names"] = staged_names + staged_group[_SYNTHESIZED_PARAM_NAMES_KEY] = staged_synthesized_names staged.param_groups.append(staged_group) staged._gefen_local_shard_bindings = tuple( diff --git a/src/gefen/portable_runtime.py b/src/gefen/portable_runtime.py index 353fdc0..193f166 100644 --- a/src/gefen/portable_runtime.py +++ b/src/gefen/portable_runtime.py @@ -71,7 +71,7 @@ )._wire_limits() _PLAIN_GROUP_REQUIRED = frozenset({"params", "param_names", "lr", "beta1", "beta2", "eps", "weight_decay"}) -_PLAIN_GROUP_ALLOWED = _PLAIN_GROUP_REQUIRED | {"name"} +_PLAIN_GROUP_ALLOWED = _PLAIN_GROUP_REQUIRED | {"name", "param_names_synthesized"} _MUON_GROUP_REQUIRED = _PLAIN_GROUP_REQUIRED | frozenset( { "momentum", @@ -90,7 +90,11 @@ "cautious", } ) -_MUON_GROUP_ALLOWED = _MUON_GROUP_REQUIRED | {"name", "ns_schedule"} +_MUON_GROUP_ALLOWED = _MUON_GROUP_REQUIRED | { + "name", + "ns_schedule", + "param_names_synthesized", +} _AUTHORITATIVE_STATE_KEYS = frozenset( {