Harden DASC policy identity and restore

modelopt/torch/sparsity/state_sparsity/conversion.py

@@ -90,7 +90,13 @@ def restore_dasc_model(model: nn.Module, config: DASCConfig, metadata: MetadataD
     if mismatched:
         raise ApplyModeError(f"DASC policy metadata does not match its mode config: {mismatched}")
 
-    validate_dasc_model_structure(model, policy)
+    try:
+        validate_dasc_model_structure(model, policy)
+    except ApplyModeError as error:
+        warnings.warn(
+            f"{error}. The restored DASC policy is stale; re-run calibrate() before deployment",
+            stacklevel=2,
+        )
     _attach_policy(model, policy)
     return model
 
@@ -117,13 +123,14 @@ def replace_dasc_mode(
 ) -> nn.Module:
     """Replace existing DASC mode state in place with a newly derived policy."""
     model = unwrap_model(model, force_unwrap=True)
-    policy = build_dasc_policy(model, config, measurements)
     manager = ModeloptStateManager(model)
-    manager.update_last_state_before_new_mode(model)
     state = manager.state_dict()
     dasc_indices = [index for index, (mode, _) in enumerate(state) if mode == "dasc"]
     if not dasc_indices:
         raise ApplyModeError("Cannot replace DASC mode because the model has no DASC state")
+    policy = build_dasc_policy(model, config, measurements)
+    if dasc_indices[-1] != len(state) - 1:
+        manager.update_last_state_before_new_mode(model)
 
     first_index = dasc_indices[0]
     state[first_index] = (

modelopt/torch/sparsity/state_sparsity/policy.py

@@ -16,23 +16,42 @@
 """Decay analysis and policy selection for GDN state sparsity."""
 
 import hashlib
+import importlib
 import json
+import math
 from collections.abc import Iterable
+from functools import lru_cache
 
 import torch
 import torch.nn.functional as F
 from torch import nn
 
 from modelopt.torch.opt.conversion import ApplyModeError
+from modelopt.torch.opt.dynamic import DynamicModule
 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"})
-# Covers ordinary FP16/BF16 storage casts; the exact retained-head mask below is the semantic gate.
-_HORIZON_DTYPE_CAST_RTOL = 0.05
+_SUPPORTED_GDN_CLASS_PATHS = (
+    ("megatron.core.ssm.gated_delta_net", "GatedDeltaNet"),
+    ("transformers.models.qwen3_next.modeling_qwen3_next", "Qwen3NextGatedDeltaNet"),
+)
+
+
+@lru_cache(maxsize=1)
+def _supported_gdn_classes() -> tuple[type[nn.Module], ...]:
+    """Resolve installed GDN implementations without making either framework mandatory."""
+    classes = []
+    for module_name, class_name in _SUPPORTED_GDN_CLASS_PATHS:
+        try:
+            candidate = getattr(importlib.import_module(module_name), class_name)
+        except (AttributeError, ImportError):
+            continue
+        if isinstance(candidate, type) and issubclass(candidate, nn.Module):
+            classes.append(candidate)
+    return tuple(classes)
 
 
 def compute_gdn_decay_horizons(
@@ -64,8 +83,14 @@ def compute_gdn_decay_horizons(
 
 def _is_gdn_module(module: nn.Module) -> bool:
     """Accept supported GDN implementations and their ModelOpt dynamic subclasses."""
+    supported_classes = _supported_gdn_classes()
+    module_class = type(module)
+    is_supported_class = module_class in supported_classes or (
+        isinstance(module, DynamicModule)
+        and any(base in supported_classes for base in module_class.__mro__)
+    )
     return (
-        any(base.__name__ in _SUPPORTED_GDN_CLASS_NAMES for base in type(module).__mro__)
+        is_supported_class
         and isinstance(getattr(module, "A_log", None), torch.Tensor)
         and isinstance(getattr(module, "dt_bias", None), torch.Tensor)
     )
@@ -76,7 +101,7 @@ def _get_gdn_modules(model: nn.Module) -> dict[str, nn.Module]:
     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:
-        supported = ", ".join(sorted(_SUPPORTED_GDN_CLASS_NAMES))
+        supported = ", ".join(class_name for _, class_name in _SUPPORTED_GDN_CLASS_PATHS)
         raise ApplyModeError(f"DASC found no supported GDN modules; expected one of: {supported}")
     return dict(sorted(modules.items()))
 
@@ -140,6 +165,33 @@ def _decay_parameters(modules: dict[str, nn.Module]) -> list[dict[str, object]]:
     ]
 
 
+def _storage_rounding_radius(tensor: torch.Tensor) -> torch.Tensor:
+    """Bound one cast-to-storage rounding step around the represented tensor values."""
+    values = tensor.detach().to(device="cpu", dtype=torch.float64)
+    dtype_info = torch.finfo(tensor.dtype)
+    unit_roundoff = dtype_info.eps / 2.0
+    subnormal_slack = dtype_info.tiny * dtype_info.eps
+    return values.abs() * (unit_roundoff / (1.0 - unit_roundoff)) + subnormal_slack
+
+
+def _storage_cast_horizon_bounds(
+    module: nn.Module, *, epsilon: float, static_gate_input: float
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Bound horizons compatible with the current parameters before one storage cast."""
+    a_log = module.A_log.detach().to(device="cpu", dtype=torch.float64)
+    dt_bias = module.dt_bias.detach().to(device="cpu", dtype=torch.float64)
+    a_radius = _storage_rounding_radius(module.A_log)
+    dt_radius = _storage_rounding_radius(module.dt_bias)
+    scale = -math.log(epsilon)
+    lower = scale / (
+        torch.exp(a_log + a_radius) * F.softplus(dt_bias + dt_radius + static_gate_input)
+    )
+    upper = scale / (
+        torch.exp(a_log - a_radius) * F.softplus(dt_bias - dt_radius + static_gate_input)
+    )
+    return lower, upper
+
+
 def _validate_measurement_coverage(
     config: DASCConfig, measurements: list[DASCCalibrationMeasurement]
 ) -> None:
@@ -286,19 +338,24 @@ def validate_dasc_decay_parameters(model: nn.Module, policy: DASCPolicy) -> None
         epsilon=policy.epsilon,
         static_gate_input=policy.static_gate_input,
     )
+    modules = _get_gdn_modules(model)
     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}"
             )
+        lower, upper = _storage_cast_horizon_bounds(
+            modules[name],
+            epsilon=policy.epsilon,
+            static_gate_input=policy.static_gate_input,
+        )
+        stored = torch.tensor(layer.static_horizons, dtype=torch.float64)
+        numerical_slack = 32.0 * torch.finfo(torch.float64).eps
+        if torch.any(stored < lower * (1.0 - numerical_slack)) or torch.any(
+            stored > upper * (1.0 + numerical_slack)
+        ):
+            raise ApplyModeError(
+                f"DASC policy horizons do not match current decay parameters in layer {name!r}"
+            )

tests/unit/torch/sparsity/state_sparsity/test_dasc.py

@@ -26,18 +26,22 @@
 
 import modelopt.torch.opt as mto
 import modelopt.torch.sparsity.state_sparsity as mtss
+import modelopt.torch.sparsity.state_sparsity.policy as dasc_policy
 from modelopt.torch.opt.conversion import ApplyModeError, ModeloptStateManager
+from modelopt.torch.opt.dynamic import DynamicModule
 from modelopt.torch.sparsity.state_sparsity.conversion import replace_dasc_mode
 from modelopt.torch.sparsity.state_sparsity.mode import DASCModeRegistry
 
+_resolve_supported_gdn_classes = dasc_policy._supported_gdn_classes.__wrapped__
+
 
 class GatedDeltaNet(nn.Module):
     """Minimal GDN-shaped module for framework-independent tests."""
 
     def __init__(self, num_heads: int = 2):
         super().__init__()
-        a_log = torch.zeros(num_heads)
-        dt_bias = torch.tensor([-2.0, 2.0]) if num_heads == 2 else torch.zeros(num_heads)
+        a_log = torch.tensor([0.1, 0.7]) if num_heads == 2 else torch.zeros(num_heads)
+        dt_bias = torch.tensor([-2.3, 1.7]) if num_heads == 2 else torch.zeros(num_heads)
         self.A_log = nn.Parameter(a_log)
         self.dt_bias = nn.Parameter(dt_bias)
 
@@ -56,6 +60,12 @@ def forward(self, inputs):
         return self.linear_attn(inputs)
 
 
+@pytest.fixture(autouse=True)
+def _register_test_gdn_class(monkeypatch):
+    """Use the exact toy GDN identity without weakening production class checks."""
+    monkeypatch.setattr(dasc_policy, "_supported_gdn_classes", lambda: (GatedDeltaNet,))
+
+
 def _config(**overrides):
     """Return a complete test configuration with selected overrides."""
     config = {
@@ -231,15 +241,49 @@ def __init__(self):
     with pytest.raises(ApplyModeError, match="no supported GDN modules"):
         mtss.calibrate(UnsupportedGatedDeltaNet(), _config(wmax_candidates=[7]), [_candidate(7)])
 
+    class UnsupportedSubclass(GatedDeltaNet):
+        pass
+
+    with pytest.raises(ApplyModeError, match="no supported GDN modules"):
+        mtss.calibrate(UnsupportedSubclass(), _config(wmax_candidates=[7]), [_candidate(7)])
+
+    same_name_lookalike = type("GatedDeltaNet", (UnsupportedGatedDeltaNet,), {})
+    with pytest.raises(ApplyModeError, match="no supported GDN modules"):
+        mtss.calibrate(same_name_lookalike(), _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():
+def test_supported_class_resolution_uses_imported_module_identities(monkeypatch):
+    """Ignore absent and non-module symbols while retaining exact supported identities."""
+    module_paths = (
+        ("valid", "GatedDeltaNet"),
+        ("invalid", "NotAModule"),
+        ("missing", "Missing"),
+    )
+    modules = {
+        "valid": type("ValidModule", (), {"GatedDeltaNet": GatedDeltaNet}),
+        "invalid": type("InvalidModule", (), {"NotAModule": object()}),
+    }
+
+    def import_module(name):
+        if name == "missing":
+            raise ImportError(name)
+        return modules[name]
+
+    monkeypatch.setattr(dasc_policy, "_SUPPORTED_GDN_CLASS_PATHS", module_paths)
+    monkeypatch.setattr(dasc_policy.importlib, "import_module", import_module)
+
+    assert _resolve_supported_gdn_classes() == (GatedDeltaNet,)
+
+
+def test_generic_mode_application_reports_missing_measurements(monkeypatch):
     """Give generic apply_mode callers an actionable calibration-evidence error."""
     assert DASCModeRegistry["dasc"].next_prohibited_modes == {"dasc"}
+    assert DASCModeRegistry["dasc"].update_for_new_mode is not None
     with pytest.raises(ApplyModeError, match="requires calibration measurements"):
         mto.apply_mode(
             TinyGatedDeltaNetForCausalLM(),
@@ -255,6 +299,23 @@ def test_generic_mode_application_reports_missing_measurements():
             [_candidate(7)],
         )
 
+    model = mtss.calibrate(
+        TinyGatedDeltaNetForCausalLM(), _config(wmax_candidates=[7]), [_candidate(7)]
+    )
+    ModeloptStateManager(model).state_dict().append(("trailing-mode", {}))
+    refreshed = []
+    monkeypatch.setattr(
+        ModeloptStateManager,
+        "update_last_state_before_new_mode",
+        lambda _manager, current_model: refreshed.append(current_model),
+    )
+    replace_dasc_mode(
+        model,
+        mtss.DASCConfig(**_config(wmax_candidates=[7])),
+        [_candidate(7)],
+    )
+    assert refreshed == [model]
+
 
 def test_public_exports_and_wrapped_model_export():
     """Expose only supported symbols and accept wrappers and ModelOpt subclasses."""
@@ -266,10 +327,10 @@ def test_public_exports_and_wrapped_model_export():
     assert "mode" not in mtss.__all__
     assert mtss.export_policy(nn.DataParallel(model)) == mtss.export_policy(model)
 
-    class _DynamicGatedDeltaNet(GatedDeltaNet):
-        """Stand in for the subclass synthesized by ModelOpt dynamic conversion."""
-
-    model.linear_attn = _DynamicGatedDeltaNet()
+    dynamic_class = type("_DynamicGatedDeltaNet", (DynamicModule, GatedDeltaNet), {})
+    dynamic_module = GatedDeltaNet()
+    dynamic_module.__class__ = dynamic_class
+    model.linear_attn = dynamic_module
     assert mtss.export_policy(model)["layers"]["linear_attn"]["num_heads"] == 2
 
     with pytest.raises(ApplyModeError, match="no valid attached DASC policy"):
@@ -283,9 +344,16 @@ def test_dtype_cast_preserves_policy_when_the_selected_mask_is_unchanged(dtype):
         TinyGatedDeltaNetForCausalLM(), _config(wmax_candidates=[7]), [_candidate(7)]
     )
     policy = mtss.export_policy(model)
+    original_decay = torch.cat(
+        [model.linear_attn.A_log.detach(), model.linear_attn.dt_bias.detach()]
+    )
 
     model.to(dtype)
 
+    cast_decay = torch.cat(
+        [model.linear_attn.A_log.detach().float(), model.linear_attn.dt_bias.detach().float()]
+    )
+    assert not torch.equal(original_decay, cast_decay)
     assert mtss.export_policy(model) == policy
 
 
@@ -309,12 +377,12 @@ def test_export_rejects_changed_decay_parameters_and_restore_rejects_structure()
 
     manager_state = ModeloptStateManager(model).state_dict()
     manager_state.append(copy.deepcopy(manager_state[0]))
-    with pytest.warns(UserWarning, match="saved DASC policy is stale"):
-        model = mtss.calibrate(
-            model,
-            _config(wmax_candidates=[7], model_revision="revision-2"),
-            [_candidate(7)],
-        )
+    delattr(model, "_modelopt_dasc_policy")
+    model = mtss.calibrate(
+        model,
+        _config(wmax_candidates=[7], model_revision="revision-2"),
+        [_candidate(7)],
+    )
     policy = mtss.export_policy(model)
     assert policy["model_revision"] == "revision-2"
     model_state = copy.deepcopy(model.state_dict())
@@ -324,8 +392,12 @@ def test_export_rejects_changed_decay_parameters_and_restore_rejects_structure()
     restored.load_state_dict(model_state)
     assert mtss.export_policy(restored) == policy
 
+    with pytest.warns(UserWarning, match="restored DASC policy is stale"):
+        mismatched = mto.restore_from_modelopt_state(
+            TinyGatedDeltaNetForCausalLM(num_heads=3), state
+        )
     with pytest.raises(ApplyModeError, match="module structure"):
-        mto.restore_from_modelopt_state(TinyGatedDeltaNetForCausalLM(num_heads=3), state)
+        mtss.export_policy(mismatched)
 
     tampered_state = copy.deepcopy(state)
     tampered_state["modelopt_state_dict"][0][1]["metadata"]["policy"]["quality_gates"][
@@ -348,16 +420,14 @@ def test_export_rejects_changed_decay_parameters_and_restore_rejects_structure()
     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]
+    layer["static_horizons"][0] *= 1.01
     restored = mto.restore_from_modelopt_state(TinyGatedDeltaNetForCausalLM(), tampered_state)
     with pytest.raises(ApplyModeError, match="horizons do not match"):
         mtss.export_policy(restored)
 
 
 def test_structure_staleness_does_not_block_checkpoint_save():
-    """Keep checkpoint saving available after a calibrated GDN structure changes."""
+    """Keep checkpoint save and restore available after a calibrated GDN structure changes."""
     model = mtss.calibrate(
         TinyGatedDeltaNetForCausalLM(), _config(wmax_candidates=[7]), [_candidate(7)]
     )
@@ -371,6 +441,16 @@ def test_structure_staleness_does_not_block_checkpoint_save():
     with pytest.raises(ApplyModeError, match="module structure"):
         mtss.export_policy(model)
 
+    with pytest.warns(UserWarning, match="saved DASC policy is stale"):
+        stale_state = mto.modelopt_state(model)
+    with pytest.warns(UserWarning, match="restored DASC policy is stale"):
+        restored = mto.restore_from_modelopt_state(
+            TinyGatedDeltaNetForCausalLM(num_heads=3), stale_state
+        )
+    restored.load_state_dict(model.state_dict())
+    with pytest.raises(ApplyModeError, match="module structure"):
+        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."""