Fix DASC storage-boundary validation

docs/source/guides/6_sparsity.rst

@@ -138,6 +138,8 @@ ordinary dense recurrent state before continuation.
 
     config = {
         "variant": "dasc_wr",  # "dasc_nr" uses zero recovery instead
+        "epsilon": 1e-3,
+        "static_gate_input": -0.3,
         "wmax_candidates": [32],
         # Set this to the dtype used to store A_log and dt_bias in the checkpoint.
         "decay_parameter_storage_dtype": "bfloat16",
@@ -146,6 +148,14 @@ ordinary dense recurrent state before continuation.
         "model_config_id": "sha256:<config-digest>",
         "calibration_data_id": "sha256:<dataset-and-protocol-digest>",
     }
+    # Use these storage-canonical horizons to derive the evaluated mask and the
+    # retained_heads/total_heads measurement geometry for every Wmax candidate.
+    horizons = mtss.analyze_gdn_decay(
+        model,
+        epsilon=config["epsilon"],
+        static_gate_input=config["static_gate_input"],
+        decay_parameter_storage_dtype=config["decay_parameter_storage_dtype"],
+    )
     measurements = [
         {
             "variant": "dasc_wr",
@@ -178,12 +188,19 @@ The initial GDN adapter accepts the ``GatedDeltaNet`` and ``Qwen3NextGatedDeltaN
 including ModelOpt-generated dynamic subclasses, and fails closed for unrelated implementations
 even when they expose similarly named decay tensors.
 Re-running :func:`~modelopt.torch.sparsity.state_sparsity.calibrate` replaces the existing DASC
-mode-state entry and supersedes its stale policy without growing the checkpoint history. A stale
-policy remains serializable so it cannot block saving or composing other ModelOpt modes, but
+mode-state entry and supersedes its stale policy without growing the checkpoint history. A policy
+with recoverable GDN geometry drift, decay drift, or temporarily unavailable decay tensors remains
+serializable, but removing or replacing the supported GDN architecture fails closed on both save
+and restore. In every stale-policy case,
 :func:`~modelopt.torch.sparsity.state_sparsity.export_policy` rejects it until recalibration.
 Set ``decay_parameter_storage_dtype`` to the checkpoint dtype for ``A_log`` and ``dt_bias`` before
-calibration. Policy validation allows only the rounding introduced by that declared storage dtype
-and the live tensor dtype; the default ``float32`` keeps unconfigured policies strict.
+calibration. Derive the evaluated head masks and reported ``retained_heads``/``total_heads`` from
+:func:`~modelopt.torch.sparsity.state_sparsity.analyze_gdn_decay` using the same ``epsilon``,
+``static_gate_input``, and storage dtype passed to calibration. Policy validation allows only the
+rounding introduced by the declared storage dtype and the live tensor dtype. When they differ,
+distinct lossy inverse rounding bounds are composed in sequence; a duplicate dtype or an exact
+widening cast contributes no additional slack. Decay tensors that are live in BF16 or FP16 are
+still validated against that live dtype's rounding.
 
 .. _sparsity-concepts:
 

modelopt/torch/sparsity/state_sparsity/config.py

@@ -16,6 +16,7 @@
 """Configuration and result schemas for DASC state sparsity."""
 
 import math
+from numbers import Real
 from typing import Literal
 
 from pydantic import ConfigDict, Field, field_validator, model_validator
@@ -30,6 +31,39 @@
     "DASCQualityMeasurement",
 ]
 
+_DecayParameterStorageDtype = Literal["float16", "bfloat16", "float32"]
+_DEFAULT_EPSILON = 1e-3
+_DEFAULT_STATIC_GATE_INPUT = -0.3
+
+
+def _validate_analysis_arguments(
+    epsilon: object = _DEFAULT_EPSILON,
+    static_gate_input: object = _DEFAULT_STATIC_GATE_INPUT,
+) -> None:
+    """Reject decay-analysis arguments that cannot produce well-defined horizons."""
+    try:
+        epsilon_is_valid = (
+            isinstance(epsilon, Real)
+            and not isinstance(epsilon, bool)
+            and math.isfinite(epsilon)
+            and 0.0 < epsilon < 1.0
+        )
+    except OverflowError:
+        epsilon_is_valid = False
+    if not epsilon_is_valid:
+        raise ValueError("epsilon must be finite and in (0, 1)")
+
+    try:
+        static_gate_input_is_valid = (
+            isinstance(static_gate_input, Real)
+            and not isinstance(static_gate_input, bool)
+            and math.isfinite(static_gate_input)
+        )
+    except OverflowError:
+        static_gate_input_is_valid = False
+    if not static_gate_input_is_valid:
+        raise ValueError("static_gate_input must be finite")
+
 
 class DASCQualityMeasurement(ModeloptBaseConfig):
     """Quality and lifecycle measurements for one calibration slice."""
@@ -79,14 +113,14 @@ class DASCConfig(ModeloptBaseConfig):
         description="Use zero recovery (DASC-NR) or suffix replay recovery (DASC-WR).",
     )
     epsilon: float = ModeloptField(
-        default=1e-3,
+        default=_DEFAULT_EPSILON,
         description="Retained contribution threshold used to derive static decay horizons.",
     )
     static_gate_input: float = ModeloptField(
-        default=-0.3,
+        default=_DEFAULT_STATIC_GATE_INPUT,
         description="Static gate input added to each GDN head's dt_bias.",
     )
-    decay_parameter_storage_dtype: Literal["float16", "bfloat16", "float32"] = ModeloptField(
+    decay_parameter_storage_dtype: _DecayParameterStorageDtype = ModeloptField(
         default="float32",
         description="Expected checkpoint storage dtype for GDN A_log and dt_bias.",
     )
@@ -116,16 +150,14 @@ class DASCConfig(ModeloptBaseConfig):
     @classmethod
     def validate_epsilon(cls, epsilon: float) -> float:
         """Require a finite decay threshold strictly between zero and one."""
-        if not math.isfinite(epsilon) or not 0.0 < epsilon < 1.0:
-            raise ValueError("epsilon must be finite and in (0, 1)")
+        _validate_analysis_arguments(epsilon=epsilon)
         return epsilon
 
     @field_validator("static_gate_input")
     @classmethod
     def validate_static_gate_input(cls, value: float) -> float:
         """Require a finite representative gate input."""
-        if not math.isfinite(value):
-            raise ValueError("static_gate_input must be finite")
+        _validate_analysis_arguments(static_gate_input=value)
         return value
 
     @field_validator("wmax_candidates", mode="before")
@@ -206,7 +238,7 @@ class DASCPolicy(ModeloptBaseConfig):
     recovery: Literal["zero", "suffix_replay"]
     epsilon: float
     static_gate_input: float
-    decay_parameter_storage_dtype: Literal["float16", "bfloat16", "float32"] = "float32"
+    decay_parameter_storage_dtype: _DecayParameterStorageDtype = "float32"
     selected_wmax: int = Field(strict=True, gt=0)
     wmax_candidates: list[int] = Field(min_length=1)
     quality_gates: dict[str, float]

modelopt/torch/sparsity/state_sparsity/conversion.py

@@ -27,7 +27,12 @@
 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
+from .policy import (
+    _DASCRecoverableStalenessError,
+    build_dasc_policy,
+    validate_dasc_decay_parameters,
+    validate_dasc_model_structure,
+)
 
 __all__ = []
 
@@ -93,7 +98,7 @@ def restore_dasc_model(model: nn.Module, config: DASCConfig, metadata: MetadataD
 
     try:
         validate_dasc_model_structure(model, policy)
-    except ApplyModeError as error:
+    except _DASCRecoverableStalenessError as error:
         warnings.warn(
             f"{error}. The restored DASC policy is stale; re-run calibrate() before deployment",
             stacklevel=2,
@@ -103,16 +108,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 metadata without making unrelated ModelOpt save or compose paths unusable."""
+    """Refresh metadata while allowing recoverable policy staleness to remain serializable."""
     policy = get_attached_dasc_policy(model)
     try:
         validate_dasc_model_structure(model, policy)
-        validate_dasc_decay_parameters(model, policy)
-    except ApplyModeError as error:
+    except _DASCRecoverableStalenessError as error:
         warnings.warn(
             f"{error}. The saved DASC policy is stale; re-run calibrate() before deployment",
             stacklevel=2,
         )
+    else:
+        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"))