Align DASC analysis and lifecycle contracts

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,14 +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, so the effective tolerance is the wider of the two. The default
-``float32`` adds no storage slack of its own; decay tensors that are live in BF16 or FP16 are still
-validated against that live dtype's rounding.
+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

@@ -28,7 +28,7 @@
 
 from .config import DASCCalibrationMeasurement, DASCConfig, DASCPolicy
 from .policy import (
-    _DASCModelStructureMismatchError,
+    _DASCRecoverableStalenessError,
     build_dasc_policy,
     validate_dasc_decay_parameters,
     validate_dasc_model_structure,
@@ -98,7 +98,7 @@ def restore_dasc_model(model: nn.Module, config: DASCConfig, metadata: MetadataD
 
     try:
         validate_dasc_model_structure(model, policy)
-    except _DASCModelStructureMismatchError as error:
+    except _DASCRecoverableStalenessError as error:
         warnings.warn(
             f"{error}. The restored DASC policy is stale; re-run calibrate() before deployment",
             stacklevel=2,
@@ -108,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"))
 

modelopt/torch/sparsity/state_sparsity/policy.py

@@ -32,25 +32,52 @@
 from modelopt.torch.opt.dynamic import DynamicModule
 from modelopt.torch.utils import unwrap_model
 
-from .config import DASCCalibrationMeasurement, DASCConfig, DASCLayerPolicy, DASCPolicy
+from .config import (
+    DASCCalibrationMeasurement,
+    DASCConfig,
+    DASCLayerPolicy,
+    DASCPolicy,
+    _DecayParameterStorageDtype,
+    _validate_analysis_arguments,
+)
 
 __all__ = ["analyze_gdn_decay", "compute_gdn_decay_horizons"]
 
 _SUPPORTED_GDN_CLASS_PATHS = (
     ("megatron.core.ssm.gated_delta_net", "GatedDeltaNet"),
     ("transformers.models.qwen3_next.modeling_qwen3_next", "Qwen3NextGatedDeltaNet"),
 )
-_STORAGE_DTYPES = {
+_STORAGE_DTYPES: dict[_DecayParameterStorageDtype, torch.dtype] = {
     "float16": torch.float16,
     "bfloat16": torch.bfloat16,
     "float32": torch.float32,
 }
 
 
-class _DASCModelStructureMismatchError(ApplyModeError):
+class _DASCRecoverableStalenessError(ApplyModeError):
+    """Identify DASC state that may become valid after model rematerialization."""
+
+
+class _DASCModelStructureMismatchError(_DASCRecoverableStalenessError):
     """Identify recoverable policy-versus-GDN-geometry drift during restore."""
 
 
+class _DASCDecayParametersUnavailableError(_DASCRecoverableStalenessError):
+    """Identify supported GDN modules whose decay tensors are temporarily unavailable."""
+
+
+def _validate_gdn_decay_tensors(a_log: torch.Tensor, dt_bias: torch.Tensor) -> None:
+    """Reject decay tensors that cannot produce well-defined horizons."""
+    if a_log.ndim != 1 or dt_bias.ndim != 1 or a_log.shape != dt_bias.shape or not a_log.numel():
+        raise ValueError(
+            "GDN A_log and dt_bias must be non-empty one-dimensional tensors of equal shape"
+        )
+    if not a_log.dtype.is_floating_point or not dt_bias.dtype.is_floating_point:
+        raise ValueError("GDN A_log and dt_bias must use floating-point dtypes")
+    if not torch.isfinite(a_log).all() or not torch.isfinite(dt_bias).all():
+        raise ValueError("GDN decay parameters must be finite")
+
+
 @lru_cache(maxsize=1)
 def _supported_gdn_classes() -> tuple[type[nn.Module], ...]:
     """Resolve installed GDN implementations without making either framework mandatory."""
@@ -89,22 +116,13 @@ def compute_gdn_decay_horizons(
     static_gate_input: float = -0.3,
 ) -> torch.Tensor:
     """Compute one static retention horizon per GDN head in CPU float64."""
-    if a_log.ndim != 1 or dt_bias.ndim != 1 or a_log.shape != dt_bias.shape or not a_log.numel():
-        raise ValueError(
-            "GDN A_log and dt_bias must be non-empty one-dimensional tensors of equal shape"
-        )
-    if not a_log.dtype.is_floating_point or not dt_bias.dtype.is_floating_point:
-        raise ValueError("GDN A_log and dt_bias must use floating-point dtypes")
-    if not 0.0 < epsilon < 1.0:
-        raise ValueError("epsilon must be in (0, 1)")
-
+    _validate_analysis_arguments(epsilon, static_gate_input)
+    _validate_gdn_decay_tensors(a_log, dt_bias)
     a_log_cpu = a_log.detach().to(device="cpu", dtype=torch.float64)
     dt_bias_cpu = dt_bias.detach().to(device="cpu", dtype=torch.float64)
-    if not torch.isfinite(a_log_cpu).all() or not torch.isfinite(dt_bias_cpu).all():
-        raise ValueError("GDN decay parameters must be finite")
 
     decay = -torch.exp(a_log_cpu) * F.softplus(dt_bias_cpu + static_gate_input)
-    horizons = torch.log(torch.tensor(epsilon, dtype=torch.float64)) / decay
+    horizons = math.log(epsilon) / decay
     if not torch.isfinite(horizons).all() or not torch.all(horizons > 0):
         raise ValueError("GDN decay parameters produced non-finite or non-positive horizons")
     return horizons
@@ -132,7 +150,7 @@ def _reject_incomplete_gdn_modules(identity_modules: list[tuple[str, nn.Module]]
         )
     ]
     if missing_decay_parameters:
-        raise ApplyModeError(
+        raise _DASCDecayParametersUnavailableError(
             "DASC found supported GDN modules without A_log and dt_bias tensors at: "
             f"{', '.join(missing_decay_parameters)}"
         )
@@ -191,9 +209,8 @@ def _analyze_gdn_modules(
         a_log = module.A_log
         dt_bias = module.dt_bias
         try:
-            if not a_log.dtype.is_floating_point or not dt_bias.dtype.is_floating_point:
-                raise ValueError("GDN A_log and dt_bias must use floating-point dtypes")
             if storage_dtype is not None:
+                _validate_gdn_decay_tensors(a_log, dt_bias)
                 a_log = a_log.detach().to(device="cpu", dtype=storage_dtype)
                 dt_bias = dt_bias.detach().to(device="cpu", dtype=storage_dtype)
             layer_horizons = compute_gdn_decay_horizons(
@@ -215,10 +232,24 @@ def analyze_gdn_decay(
     *,
     epsilon: float = 1e-3,
     static_gate_input: float = -0.3,
+    decay_parameter_storage_dtype: _DecayParameterStorageDtype | None = None,
 ) -> dict[str, list[float]]:
-    """Return deterministic per-head horizons for every GDN module in a model."""
+    """Return per-head horizons, optionally canonicalized to a checkpoint storage dtype."""
+    _validate_analysis_arguments(epsilon, static_gate_input)
+    storage_dtype = None
+    if decay_parameter_storage_dtype is not None:
+        if (
+            not isinstance(decay_parameter_storage_dtype, str)
+            or decay_parameter_storage_dtype not in _STORAGE_DTYPES
+        ):
+            supported = ", ".join(_STORAGE_DTYPES)
+            raise ValueError(f"decay_parameter_storage_dtype must be one of: {supported}")
+        storage_dtype = _STORAGE_DTYPES[decay_parameter_storage_dtype]
     return _analyze_gdn_modules(
-        _get_gdn_modules(model), epsilon=epsilon, static_gate_input=static_gate_input
+        _get_gdn_modules(model),
+        epsilon=epsilon,
+        static_gate_input=static_gate_input,
+        storage_dtype=storage_dtype,
     )
 
 
@@ -259,11 +290,28 @@ def _decay_parameters(
     ]
 
 
+def _dtype_exactly_contains(source: torch.dtype, target: torch.dtype) -> bool:
+    """Return whether every finite source value is exactly representable in the target dtype."""
+    source_info = torch.finfo(source)
+    target_info = torch.finfo(target)
+    return (
+        target_info.max >= source_info.max
+        and target_info.eps <= source_info.eps
+        and target_info.tiny * target_info.eps <= source_info.tiny * source_info.eps
+    )
+
+
 def _storage_rounding_radius(tensor: torch.Tensor, storage_dtype: torch.dtype) -> torch.Tensor:
     """Compose inverse error bounds for storage and live-dtype materialization casts."""
     values = tensor.detach().to(device="cpu", dtype=torch.float64).abs()
     upper = values
-    cast_dtypes = tuple(dict.fromkeys((storage_dtype, tensor.dtype)))
+    live_dtype = tensor.dtype
+    if _dtype_exactly_contains(live_dtype, storage_dtype):
+        cast_dtypes = (live_dtype,)
+    elif _dtype_exactly_contains(storage_dtype, live_dtype):
+        cast_dtypes = (storage_dtype,)
+    else:
+        cast_dtypes = (storage_dtype, live_dtype)
     for dtype in reversed(cast_dtypes):
         dtype_info = torch.finfo(dtype)
         unit_roundoff = dtype_info.eps / 2.0
@@ -279,7 +327,7 @@ def _storage_cast_horizon_bounds(
     static_gate_input: float,
     storage_dtype: torch.dtype,
 ) -> tuple[torch.Tensor, torch.Tensor]:
-    """Bound horizons compatible with the current parameters before one storage cast."""
+    """Bound horizons compatible with current parameters under storage and live-dtype casts."""
     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, storage_dtype)
@@ -447,8 +495,12 @@ def validate_dasc_decay_parameters(model: nn.Module, policy: DASCPolicy) -> None
     modules = _get_gdn_modules(model)
     for name, module in modules.items():
         layer = policy.layers[name]
-        if not module.A_log.dtype.is_floating_point or not module.dt_bias.dtype.is_floating_point:
-            raise ApplyModeError("GDN A_log and dt_bias must use floating-point dtypes")
+        try:
+            _validate_gdn_decay_tensors(module.A_log, module.dt_bias)
+        except ValueError as error:
+            raise ApplyModeError(
+                f"Invalid GDN decay parameters in module {name!r}: {error}"
+            ) from error
         lower, upper = _storage_cast_horizon_bounds(
             module,
             epsilon=policy.epsilon,
@@ -468,7 +520,7 @@ def validate_dasc_decay_parameters(model: nn.Module, policy: DASCPolicy) -> None
                     "DASC policy head mask does not match current decay parameters in layer "
                     f"{name!r}"
                 )
-        stored = torch.tensor(layer.static_horizons, dtype=torch.float64)
+        stored = torch.tensor(layer.static_horizons, device="cpu", 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)