From edc91ec0f1eb7ec292b8fdd7c63c16b5589f62f0 Mon Sep 17 00:00:00 2001 From: Kai Xu Date: Thu, 10 Sep 2026 17:58:31 -0700 Subject: [PATCH] Address DASC policy review feedback Signed-off-by: Kai Xu --- docs/source/guides/6_sparsity.rst | 9 +- .../torch/sparsity/state_sparsity/__init__.py | 26 +++- .../torch/sparsity/state_sparsity/config.py | 13 +- .../sparsity/state_sparsity/conversion.py | 26 +++- .../torch/sparsity/state_sparsity/mode.py | 5 - .../torch/sparsity/state_sparsity/policy.py | 66 +++++++-- .../sparsity/state_sparsity/test_dasc.py | 136 +++++++++++++++++- 7 files changed, 248 insertions(+), 33 deletions(-) diff --git a/docs/source/guides/6_sparsity.rst b/docs/source/guides/6_sparsity.rst index d507dd8ba3f..ec83d51fb13 100644 --- a/docs/source/guides/6_sparsity.rst +++ b/docs/source/guides/6_sparsity.rst @@ -123,7 +123,9 @@ derives one static decay horizon per complete GDN head from ``A_log`` and ``dt_bias``, then selects the largest caller-evaluated ``Wmax`` that passes every configured quality, lifecycle, and physical-storage gate. ``Wmax`` may be any positive integer. The measurements are evidence inputs produced by a caller-owned paired evaluation; this API does -not run the dense-versus-recovery suffix evaluation itself. +not run the dense-versus-recovery suffix evaluation itself. ``perplexity_retention`` is defined as +``dense_perplexity / DASC_perplexity`` (equivalently ``exp(dense_NLL - DASC_NLL)``), so higher is +better and values above one are valid. The initial API exports policy metadata only. It does not change model execution, quantize state, pack ragged checkpoints, replay a suffix, or add a linear-attention kernel. A serving integration @@ -170,6 +172,11 @@ ordinary dense recurrent state before continuation. from zero, while DASC-WR reconstructs them from a zero-initialized suffix replay of at most the selected ``Wmax`` tokens. Both retain whole GDN heads, preserve convolution state, and resume with dense recurrence. KDA and serving-runtime integration are not supported by this initial API. +The initial GDN adapter accepts the ``GatedDeltaNet`` and ``Qwen3NextGatedDeltaNet`` class names +and fails closed for other implementations, even when they expose similarly named decay tensors. +Re-running :func:`~modelopt.torch.sparsity.state_sparsity.calibrate` supersedes a stale policy. A +stale policy remains serializable so it cannot block saving or composing other ModelOpt modes, but +:func:`~modelopt.torch.sparsity.state_sparsity.export_policy` rejects it until recalibration. .. _sparsity-concepts: diff --git a/modelopt/torch/sparsity/state_sparsity/__init__.py b/modelopt/torch/sparsity/state_sparsity/__init__.py index 97fe2d89b13..b0d1df3cfe0 100644 --- a/modelopt/torch/sparsity/state_sparsity/__init__.py +++ b/modelopt/torch/sparsity/state_sparsity/__init__.py @@ -15,7 +15,25 @@ """Decay-aware sparsity policies for persisted recurrent state.""" -from . import mode -from .api import * -from .config import * -from .policy import * +from . import mode # imported for mode-registration side effects +from .api import calibrate, export_policy +from .config import ( + DASCCalibrationMeasurement, + DASCConfig, + DASCLayerPolicy, + DASCPolicy, + DASCQualityMeasurement, +) +from .policy import analyze_gdn_decay, compute_gdn_decay_horizons + +__all__ = [ + "DASCCalibrationMeasurement", + "DASCConfig", + "DASCLayerPolicy", + "DASCPolicy", + "DASCQualityMeasurement", + "analyze_gdn_decay", + "calibrate", + "compute_gdn_decay_horizons", + "export_policy", +] diff --git a/modelopt/torch/sparsity/state_sparsity/config.py b/modelopt/torch/sparsity/state_sparsity/config.py index c7bca8be9f3..8ae62cddff6 100644 --- a/modelopt/torch/sparsity/state_sparsity/config.py +++ b/modelopt/torch/sparsity/state_sparsity/config.py @@ -18,13 +18,14 @@ import math from typing import Literal -from pydantic import Field, field_validator, model_validator +from pydantic import ConfigDict, Field, field_validator, model_validator from modelopt.torch.opt.config import ModeloptBaseConfig, ModeloptField __all__ = [ "DASCCalibrationMeasurement", "DASCConfig", + "DASCLayerPolicy", "DASCPolicy", "DASCQualityMeasurement", ] @@ -34,7 +35,11 @@ class DASCQualityMeasurement(ModeloptBaseConfig): """Quality and lifecycle measurements for one calibration slice.""" slice_id: str = Field(min_length=1) - perplexity_retention: float = Field(gt=0.0, le=1.0, allow_inf_nan=False) + perplexity_retention: float = Field( + gt=0.0, + allow_inf_nan=False, + description="Dense perplexity divided by DASC perplexity; values above one are valid.", + ) top1_agreement: float = Field(ge=0.0, le=1.0, allow_inf_nan=False) finite_continuation_logits: bool = Field(strict=True) retained_state_exact: bool = Field(strict=True) @@ -67,6 +72,8 @@ def validate_unique_slices( class DASCConfig(ModeloptBaseConfig): """Configuration for GDN decay-aware state checkpoint sparsity.""" + model_config = ConfigDict(extra="forbid", validate_assignment=True, protected_namespaces=()) + variant: Literal["dasc_nr", "dasc_wr"] = ModeloptField( default="dasc_wr", description="Use zero recovery (DASC-NR) or suffix replay recovery (DASC-WR).", @@ -188,6 +195,8 @@ def validate_partition(self) -> "DASCLayerPolicy": class DASCPolicy(ModeloptBaseConfig): """Standalone JSON-safe DASC deployment policy.""" + model_config = ConfigDict(extra="forbid", validate_assignment=True, protected_namespaces=()) + format_version: Literal[1] = 1 variant: Literal["dasc_nr", "dasc_wr"] recovery: Literal["zero", "suffix_replay"] diff --git a/modelopt/torch/sparsity/state_sparsity/conversion.py b/modelopt/torch/sparsity/state_sparsity/conversion.py index 5e756a2587d..95d50a18223 100644 --- a/modelopt/torch/sparsity/state_sparsity/conversion.py +++ b/modelopt/torch/sparsity/state_sparsity/conversion.py @@ -16,6 +16,7 @@ """ModelOpt conversion and restoration for DASC policy metadata.""" import copy +import warnings from collections.abc import Iterable from pydantic import ValidationError @@ -23,6 +24,7 @@ from modelopt.torch.opt.conversion import ApplyModeError from modelopt.torch.opt.mode import ConvertReturnType, MetadataDict +from modelopt.torch.utils import unwrap_model from .config import DASCCalibrationMeasurement, DASCConfig, DASCPolicy from .policy import build_dasc_policy, validate_dasc_decay_parameters, validate_dasc_model_structure @@ -33,6 +35,7 @@ def _attach_policy(model: nn.Module, policy: DASCPolicy) -> None: + """Attach a JSON-safe DASC policy to an unwrapped model.""" setattr(model, _DASC_POLICY_ATTRIBUTE, policy.model_dump(mode="json")) @@ -40,9 +43,14 @@ def convert_dasc_model( model: nn.Module, config: DASCConfig, *, - measurements: Iterable[DASCCalibrationMeasurement | dict], + measurements: Iterable[DASCCalibrationMeasurement | dict] | None = None, ) -> ConvertReturnType: """Analyze GDN decay and attach the selected DASC policy without changing execution.""" + if measurements is None: + raise ApplyModeError( + "DASC requires calibration measurements; use " + "modelopt.torch.sparsity.state_sparsity.calibrate(model, config, measurements)" + ) policy = build_dasc_policy(model, config, measurements) _attach_policy(model, policy) return model, {"policy": policy.model_dump(mode="json")} @@ -88,19 +96,23 @@ def restore_dasc_model(model: nn.Module, config: DASCConfig, metadata: MetadataD def update_dasc_metadata(model: nn.Module, config: DASCConfig, metadata: MetadataDict) -> None: - """Refresh serialized metadata from the immutable attached DASC policy.""" - try: - policy = DASCPolicy(**getattr(model, _DASC_POLICY_ATTRIBUTE)) - except (AttributeError, TypeError, ValidationError) as error: - raise ApplyModeError("Model has no valid attached DASC policy") from error + """Refresh metadata without making unrelated ModelOpt save or compose paths unusable.""" + policy = get_attached_dasc_policy(model) validate_dasc_model_structure(model, policy) - validate_dasc_decay_parameters(model, policy) + try: + validate_dasc_decay_parameters(model, policy) + except ApplyModeError as error: + warnings.warn( + f"{error}. The saved DASC policy is stale; re-run calibrate() before deployment", + stacklevel=2, + ) metadata.clear() metadata["policy"] = copy.deepcopy(policy.model_dump(mode="json")) def get_attached_dasc_policy(model: nn.Module) -> DASCPolicy: """Return the validated policy attached by conversion or restoration.""" + model = unwrap_model(model, force_unwrap=True) try: return DASCPolicy(**getattr(model, _DASC_POLICY_ATTRIBUTE)) except (AttributeError, TypeError, ValidationError) as error: diff --git a/modelopt/torch/sparsity/state_sparsity/mode.py b/modelopt/torch/sparsity/state_sparsity/mode.py index 23b3a17ffc3..36b0a881dac 100644 --- a/modelopt/torch/sparsity/state_sparsity/mode.py +++ b/modelopt/torch/sparsity/state_sparsity/mode.py @@ -48,11 +48,6 @@ def config_class(self) -> type[ModeloptBaseConfig]: """Return the validated DASC configuration class.""" return DASCConfig - @property - def next_prohibited_modes(self) -> set[str]: - """Prevent applying DASC twice to the same model state.""" - return {"dasc"} - @property def convert(self) -> ConvertEntrypoint: """Return the DASC calibration entrypoint.""" diff --git a/modelopt/torch/sparsity/state_sparsity/policy.py b/modelopt/torch/sparsity/state_sparsity/policy.py index 7342a66bf17..c5f6f9b1838 100644 --- a/modelopt/torch/sparsity/state_sparsity/policy.py +++ b/modelopt/torch/sparsity/state_sparsity/policy.py @@ -24,11 +24,15 @@ from torch import nn from modelopt.torch.opt.conversion import ApplyModeError +from modelopt.torch.utils import unwrap_model from .config import DASCCalibrationMeasurement, DASCConfig, DASCLayerPolicy, DASCPolicy __all__ = ["analyze_gdn_decay", "compute_gdn_decay_horizons"] +_SUPPORTED_GDN_CLASS_NAMES = frozenset({"GatedDeltaNet", "Qwen3NextGatedDeltaNet"}) +_HORIZON_DTYPE_CAST_RTOL = 0.05 + def compute_gdn_decay_horizons( a_log: torch.Tensor, @@ -58,20 +62,21 @@ def compute_gdn_decay_horizons( def _is_gdn_module(module: nn.Module) -> bool: - class_name = "".join( - character for character in type(module).__name__.lower() if character.isalnum() - ) + """Return whether a module has one of the explicitly supported GDN implementations.""" return ( - "gateddeltanet" in class_name + type(module).__name__ in _SUPPORTED_GDN_CLASS_NAMES and isinstance(getattr(module, "A_log", None), torch.Tensor) and isinstance(getattr(module, "dt_bias", None), torch.Tensor) ) def _get_gdn_modules(model: nn.Module) -> dict[str, nn.Module]: + """Find supported GDN layers after removing a recognized model wrapper.""" + model = unwrap_model(model, force_unwrap=True) modules = {name: module for name, module in model.named_modules() if _is_gdn_module(module)} if not modules: - raise ApplyModeError("DASC found no GatedDeltaNet modules; only GDN is supported") + supported = ", ".join(sorted(_SUPPORTED_GDN_CLASS_NAMES)) + raise ApplyModeError(f"DASC found no supported GDN modules; expected one of: {supported}") return dict(sorted(modules.items())) @@ -100,11 +105,13 @@ def analyze_gdn_decay( def _canonical_sha256(value: object) -> str: + """Hash a JSON value with deterministic ordering and no non-finite numbers.""" payload = json.dumps(value, sort_keys=True, separators=(",", ":"), allow_nan=False) return hashlib.sha256(payload.encode()).hexdigest() def _model_structure(modules: dict[str, nn.Module]) -> list[dict[str, object]]: + """Describe the layer names and head counts that define policy geometry.""" return [ { "name": name, @@ -115,11 +122,18 @@ def _model_structure(modules: dict[str, nn.Module]) -> list[dict[str, object]]: def _decay_parameters(modules: dict[str, nn.Module]) -> list[dict[str, object]]: + """Serialize decay parameters at canonical BF16 precision for dtype-stable hashing.""" return [ { "name": name, - "A_log": module.A_log.detach().to(device="cpu", dtype=torch.float64).tolist(), - "dt_bias": module.dt_bias.detach().to(device="cpu", dtype=torch.float64).tolist(), + "A_log": module.A_log.detach() + .to(device="cpu", dtype=torch.bfloat16) + .to(dtype=torch.float32) + .tolist(), + "dt_bias": module.dt_bias.detach() + .to(device="cpu", dtype=torch.bfloat16) + .to(dtype=torch.float32) + .tolist(), } for name, module in modules.items() ] @@ -128,6 +142,7 @@ def _decay_parameters(modules: dict[str, nn.Module]) -> list[dict[str, object]]: def _validate_measurement_coverage( config: DASCConfig, measurements: list[DASCCalibrationMeasurement] ) -> None: + """Require exactly one matching measurement for every configured window.""" measured = [measurement.wmax for measurement in measurements] unexpected = sorted(set(measured) - set(config.wmax_candidates)) missing = sorted(set(config.wmax_candidates) - set(measured)) @@ -145,6 +160,7 @@ def _validate_measurement_coverage( def _validate_measurement_geometry( horizons: dict[str, list[float]], measurements: list[DASCCalibrationMeasurement] ) -> None: + """Bind caller-reported retained and total head counts to the analyzed model.""" total_heads = sum(len(layer_horizons) for layer_horizons in horizons.values()) for measurement in measurements: retained_heads = sum( @@ -161,6 +177,7 @@ def _validate_measurement_geometry( def _candidate_passes(config: DASCConfig, measurement: DASCCalibrationMeasurement) -> bool: + """Return whether one candidate passes every configured evidence gate.""" return measurement.checkpoint_savings >= config.min_checkpoint_savings and all( result.perplexity_retention >= config.min_perplexity_retention and result.top1_agreement >= config.min_top1_agreement @@ -245,14 +262,43 @@ def build_dasc_policy( def validate_dasc_model_structure(model: nn.Module, policy: DASCPolicy) -> None: """Reject restoring a policy onto a different GDN module structure.""" modules = _get_gdn_modules(model) - actual = _canonical_sha256(_model_structure(modules)) - if actual != policy.model_structure_sha256: + actual_structure = _model_structure(modules) + policy_structure = [ + {"name": name, "num_heads": layer.num_heads} + for name, layer in sorted(policy.layers.items()) + ] + if ( + actual_structure != policy_structure + or _canonical_sha256(actual_structure) != policy.model_structure_sha256 + ): raise ApplyModeError("DASC policy does not match the model's GDN module structure") def validate_dasc_decay_parameters(model: nn.Module, policy: DASCPolicy) -> None: - """Reject exporting a policy for different GDN decay parameters.""" + """Reject deployment when current decay parameters no longer derive the stored policy.""" modules = _get_gdn_modules(model) actual = _canonical_sha256(_decay_parameters(modules)) if actual != policy.decay_parameters_sha256: raise ApplyModeError("DASC policy does not match the model's GDN decay parameters") + + current_horizons = analyze_gdn_decay( + model, + epsilon=policy.epsilon, + static_gate_input=policy.static_gate_input, + ) + for name, values in current_horizons.items(): + layer = policy.layers[name] + if not torch.allclose( + torch.tensor(values, dtype=torch.float64), + torch.tensor(layer.static_horizons, dtype=torch.float64), + rtol=_HORIZON_DTYPE_CAST_RTOL, + atol=0.0, + ): + raise ApplyModeError( + f"DASC policy horizons do not match current decay parameters in layer {name!r}" + ) + retained = [head for head, horizon in enumerate(values) if horizon > policy.selected_wmax] + if retained != layer.retained_heads: + raise ApplyModeError( + f"DASC policy head mask does not match current decay parameters in layer {name!r}" + ) diff --git a/tests/unit/torch/sparsity/state_sparsity/test_dasc.py b/tests/unit/torch/sparsity/state_sparsity/test_dasc.py index c212808c046..0ae651863c0 100644 --- a/tests/unit/torch/sparsity/state_sparsity/test_dasc.py +++ b/tests/unit/torch/sparsity/state_sparsity/test_dasc.py @@ -16,6 +16,7 @@ """CPU tests for DASC state-sparsity policy calibration.""" import copy +import io import json import pytest @@ -28,7 +29,7 @@ from modelopt.torch.opt.conversion import ApplyModeError -class TinyGatedDeltaNet(nn.Module): +class GatedDeltaNet(nn.Module): """Minimal GDN-shaped module for framework-independent tests.""" def __init__(self, num_heads: int = 2): @@ -47,13 +48,14 @@ class TinyGatedDeltaNetForCausalLM(nn.Module): def __init__(self, num_heads: int = 2): super().__init__() - self.linear_attn = TinyGatedDeltaNet(num_heads) + self.linear_attn = GatedDeltaNet(num_heads) def forward(self, inputs): return self.linear_attn(inputs) def _config(**overrides): + """Return a complete test configuration with selected overrides.""" config = { "variant": "dasc_wr", "epsilon": 1e-3, @@ -72,6 +74,7 @@ def _config(**overrides): def _candidate(wmax, *, variant="dasc_wr", top1=0.99, convolution_state_exact=True): + """Return passing caller-supplied evidence for one recovery window.""" return { "variant": variant, "wmax": wmax, @@ -93,6 +96,7 @@ def _candidate(wmax, *, variant="dasc_wr", top1=0.99, convolution_state_exact=Tr def test_calibrate_selects_largest_passing_candidate_and_round_trips(): + """Select the largest passing window and preserve weights and ModelOpt state.""" model = TinyGatedDeltaNetForCausalLM() original_state = {name: value.clone() for name, value in model.state_dict().items()} @@ -126,6 +130,7 @@ def test_calibrate_selects_largest_passing_candidate_and_round_trips(): ("variant", "recovery"), [("dasc_nr", "zero"), ("dasc_wr", "suffix_replay")] ) def test_variant_is_explicit_in_exported_policy(variant, recovery): + """Keep zero and suffix-replay recovery as explicit deployment contracts.""" model = mtss.calibrate( TinyGatedDeltaNetForCausalLM(), _config(variant=variant, wmax_candidates=[7]), @@ -147,13 +152,23 @@ def test_variant_is_explicit_in_exported_policy(variant, recovery): ], ) def test_config_fails_closed(override): + """Reject incomplete provenance and invalid window or lifecycle settings.""" with pytest.raises(ValidationError): mtss.DASCConfig(**_config(**override)) assert mtss.DASCConfig(**_config(wmax_candidates=[7])).wmax_candidates == [7] +def test_perplexity_retention_accepts_parity_improvements(): + """Allow an observed DASC perplexity improvement instead of requiring clamping.""" + measurement = mtss.DASCCalibrationMeasurement(**_candidate(7)) + measurement.quality[0].perplexity_retention = 1.0004 + + assert measurement.quality[0].perplexity_retention == 1.0004 + + def test_calibration_fails_closed_on_measurements_and_model_mismatch(): + """Reject incomplete evidence, failing gates, wrong geometry, and unsupported layers.""" with pytest.raises(ApplyModeError, match="exactly one result"): mtss.calibrate(TinyGatedDeltaNetForCausalLM(), _config(), [_candidate(7)]) @@ -164,6 +179,22 @@ def test_calibration_fails_closed_on_measurements_and_model_mismatch(): [_candidate(7, convolution_state_exact=False)], ) + with pytest.raises(ApplyModeError, match="configured variant"): + mtss.calibrate( + TinyGatedDeltaNetForCausalLM(), + _config(wmax_candidates=[7]), + [_candidate(7, variant="dasc_nr")], + ) + + invalid_measurement = _candidate(7) + del invalid_measurement["quality"] + with pytest.raises(ApplyModeError, match="Invalid DASC calibration measurements"): + mtss.calibrate( + TinyGatedDeltaNetForCausalLM(), + _config(wmax_candidates=[7]), + [invalid_measurement], + ) + mismatched_geometry = _candidate(7) mismatched_geometry["retained_heads"] = 2 with pytest.raises(ApplyModeError, match="geometry"): @@ -176,11 +207,57 @@ def test_calibration_fails_closed_on_measurements_and_model_mismatch(): class NotGDN(nn.Module): pass - with pytest.raises(ApplyModeError, match="no GatedDeltaNet modules"): + with pytest.raises(ApplyModeError, match="no supported GDN modules"): mtss.calibrate(NotGDN(), _config(wmax_candidates=[7]), [_candidate(7)]) + class UnsupportedGatedDeltaNet(GatedDeltaNet): + pass + + with pytest.raises(ApplyModeError, match="no supported GDN modules"): + mtss.calibrate(UnsupportedGatedDeltaNet(), _config(wmax_candidates=[7]), [_candidate(7)]) + + invalid_decay = TinyGatedDeltaNetForCausalLM() + invalid_decay.linear_attn.dt_bias = nn.Parameter(torch.zeros(3)) + with pytest.raises(ApplyModeError, match="Invalid GDN decay parameters"): + mtss.analyze_gdn_decay(invalid_decay) + + +def test_generic_mode_application_reports_missing_measurements(): + """Give generic apply_mode callers an actionable calibration-evidence error.""" + with pytest.raises(ApplyModeError, match="requires calibration measurements"): + mto.apply_mode( + TinyGatedDeltaNetForCausalLM(), + mode=[("dasc", _config(wmax_candidates=[7]))], + ) + + +def test_public_exports_and_wrapped_model_export(): + """Expose only supported symbols and unwrap recognized parallel wrappers.""" + model = mtss.calibrate( + TinyGatedDeltaNetForCausalLM(), _config(wmax_candidates=[7]), [_candidate(7)] + ) + + assert "DASCLayerPolicy" in mtss.__all__ + assert "mode" not in mtss.__all__ + assert mtss.export_policy(nn.DataParallel(model)) == mtss.export_policy(model) + with pytest.raises(ApplyModeError, match="no valid attached DASC policy"): + mtss.export_policy(TinyGatedDeltaNetForCausalLM()) + + +def test_dtype_cast_preserves_policy_when_the_selected_mask_is_unchanged(): + """Treat an ordinary BF16 cast as equivalent when it preserves the policy mask.""" + model = mtss.calibrate( + TinyGatedDeltaNetForCausalLM(), _config(wmax_candidates=[7]), [_candidate(7)] + ) + policy = mtss.export_policy(model) + + model.to(torch.bfloat16) + + assert mtss.export_policy(model) == policy + def test_export_rejects_changed_decay_parameters_and_restore_rejects_structure(): + """Keep saving recoverable while rejecting stale or tampered deployment policies.""" model = mtss.calibrate( TinyGatedDeltaNetForCausalLM(), _config(wmax_candidates=[7]), [_candidate(7)] ) @@ -190,8 +267,21 @@ def test_export_rejects_changed_decay_parameters_and_restore_rejects_structure() model.linear_attn.A_log.add_(1.0) with pytest.raises(ApplyModeError, match="decay parameters"): mtss.export_policy(model) - with pytest.raises(ApplyModeError, match="decay parameters"): + with pytest.warns(UserWarning, match="saved DASC policy is stale"): mto.modelopt_state(model) + checkpoint = io.BytesIO() + with pytest.warns(UserWarning, match="saved DASC policy is stale"): + mto.save(model, checkpoint) + assert checkpoint.tell() > 0 + + with pytest.warns(UserWarning, match="saved DASC policy is stale"): + model = mtss.calibrate(model, _config(wmax_candidates=[7]), [_candidate(7)]) + policy = mtss.export_policy(model) + model_state = copy.deepcopy(model.state_dict()) + recalibrated_state = mto.modelopt_state(model) + restored = mto.restore_from_modelopt_state(TinyGatedDeltaNetForCausalLM(), recalibrated_state) + restored.load_state_dict(model_state) + assert mtss.export_policy(restored) == policy with pytest.raises(ApplyModeError, match="module structure"): mto.restore_from_modelopt_state(TinyGatedDeltaNetForCausalLM(num_heads=3), state) @@ -202,3 +292,41 @@ def test_export_rejects_changed_decay_parameters_and_restore_rejects_structure() ] = 0.5 with pytest.raises(ApplyModeError, match="does not match its mode config"): mto.restore_from_modelopt_state(TinyGatedDeltaNetForCausalLM(), tampered_state) + + tampered_state = copy.deepcopy(state) + tampered_state["modelopt_state_dict"][0][1]["metadata"]["unexpected"] = True + with pytest.raises(ApplyModeError, match="only the policy field"): + mto.restore_from_modelopt_state(TinyGatedDeltaNetForCausalLM(), tampered_state) + + tampered_state = copy.deepcopy(state) + del tampered_state["modelopt_state_dict"][0][1]["metadata"]["policy"]["variant"] + with pytest.raises(ApplyModeError, match="Invalid DASC policy metadata"): + mto.restore_from_modelopt_state(TinyGatedDeltaNetForCausalLM(), tampered_state) + + tampered_state = copy.deepcopy(state) + layer = tampered_state["modelopt_state_dict"][0][1]["metadata"]["policy"]["layers"][ + "linear_attn" + ] + layer["static_horizons"][0] = 1.0 + layer["retained_heads"] = [] + layer["omitted_heads"] = [0, 1] + restored = mto.restore_from_modelopt_state(TinyGatedDeltaNetForCausalLM(), tampered_state) + with pytest.raises(ApplyModeError, match="horizons do not match"): + mtss.export_policy(restored) + + +def test_export_rederives_the_selected_head_mask(): + """Reject a self-consistent stored mask that current decay parameters do not derive.""" + model = mtss.calibrate( + TinyGatedDeltaNetForCausalLM(), _config(wmax_candidates=[54]), [_candidate(54)] + ) + state = mto.modelopt_state(model) + layer = state["modelopt_state_dict"][0][1]["metadata"]["policy"]["layers"]["linear_attn"] + layer["static_horizons"][0] = 53.0 + layer["retained_heads"] = [] + layer["omitted_heads"] = [0, 1] + + restored = mto.restore_from_modelopt_state(TinyGatedDeltaNetForCausalLM(), state) + + with pytest.raises(ApplyModeError, match="head mask does not match"): + mtss.export_policy(restored)