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"