Address DASC review feedback

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
@@ -136,12 +138,24 @@ 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",
         "model_id": "org/model",
         "model_revision": "immutable-model-revision",
         "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",
@@ -170,6 +184,23 @@ 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`` base classes,
+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 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. 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/__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",
+]

modelopt/torch/sparsity/state_sparsity/api.py

@@ -21,10 +21,11 @@
 
 from torch import nn
 
-from modelopt.torch.opt.conversion import apply_mode
+from modelopt.torch.opt.conversion import ModeloptStateManager, apply_mode
+from modelopt.torch.utils import unwrap_model
 
 from .config import DASCCalibrationMeasurement, DASCConfig
-from .conversion import get_attached_dasc_policy
+from .conversion import get_attached_dasc_policy, replace_dasc_mode
 from .mode import DASCModeRegistry
 from .policy import validate_dasc_decay_parameters, validate_dasc_model_structure
 
@@ -40,6 +41,7 @@ def calibrate(
 
     ``measurements`` must contain exactly one entry for every configured ``Wmax`` candidate.
     The largest candidate passing every quality, lifecycle, and storage gate is selected.
+    Recalibrating replaces the existing DASC mode-state entry in place.
 
     Example::
 
@@ -58,10 +60,16 @@ def calibrate(
     Returns:
         The input model with a serializable DASC policy attached through ModelOpt state.
     """
-    config_dict = config.model_dump() if isinstance(config, DASCConfig) else config
+    model = unwrap_model(model, force_unwrap=True)
+    config_object = config if isinstance(config, DASCConfig) else DASCConfig(**config)
+    if ModeloptStateManager.is_converted(model, is_root=True) and any(
+        mode == "dasc" for mode, _ in ModeloptStateManager(model).state_dict()
+    ):
+        return replace_dasc_mode(model, config_object, measurements)
+
     return apply_mode(
         model,
-        mode=[("dasc", config_dict)],
+        mode=[("dasc", config_object.model_dump())],
         registry=DASCModeRegistry,
         mode_kwargs={"measurements": measurements},
     )

modelopt/torch/sparsity/state_sparsity/config.py

@@ -16,25 +16,64 @@
 """Configuration and result schemas for DASC state sparsity."""
 
 import math
+from numbers import Real
 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",
 ]
 
+_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."""
 
     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,18 +106,24 @@ def validate_unique_slices(
 class DASCConfig(ModeloptBaseConfig):
     """Configuration for GDN decay-aware state checkpoint sparsity."""
 
+    model_config = ConfigDict(protected_namespaces=())
+
     variant: Literal["dasc_nr", "dasc_wr"] = ModeloptField(
         default="dasc_wr",
         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: _DecayParameterStorageDtype = ModeloptField(
+        default="float32",
+        description="Expected checkpoint storage dtype for GDN A_log and dt_bias.",
+    )
     wmax_candidates: list[int] = ModeloptField(
         default=[8, 16, 32, 64, 128, 256],
         description="Positive candidate windows evaluated during offline calibration.",
@@ -105,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")
@@ -135,9 +178,9 @@ def validate_wmax_candidates(cls, candidates: object) -> object:
     @field_validator("min_perplexity_retention")
     @classmethod
     def validate_perplexity_gate(cls, value: float) -> float:
-        """Require a finite retention gate in (0, 1]."""
-        if not math.isfinite(value) or not 0.0 < value <= 1.0:
-            raise ValueError("min_perplexity_retention must be finite and in (0, 1]")
+        """Require a finite positive retention gate."""
+        if not math.isfinite(value) or value <= 0.0:
+            raise ValueError("min_perplexity_retention must be finite and positive")
         return value
 
     @field_validator("min_top1_agreement")
@@ -188,11 +231,14 @@ def validate_partition(self) -> "DASCLayerPolicy":
 class DASCPolicy(ModeloptBaseConfig):
     """Standalone JSON-safe DASC deployment policy."""
 
+    model_config = ConfigDict(protected_namespaces=())
+
     format_version: Literal[1] = 1
     variant: Literal["dasc_nr", "dasc_wr"]
     recovery: Literal["zero", "suffix_replay"]
     epsilon: float
     static_gate_input: float
+    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]
@@ -204,7 +250,10 @@ class DASCPolicy(ModeloptBaseConfig):
     preserve_convolution_state: Literal[True]
     active_runtime_state: Literal["dense"] = "dense"
     model_structure_sha256: str = Field(pattern=r"^[0-9a-f]{64}$")
-    decay_parameters_sha256: str = Field(pattern=r"^[0-9a-f]{64}$")
+    decay_parameters_sha256: str = Field(
+        pattern=r"^[0-9a-f]{64}$",
+        description="Storage-dtype-canonicalized calibration snapshot retained for provenance.",
+    )
     layers: dict[str, DASCLayerPolicy] = Field(min_length=1)
     measurements: list[DASCCalibrationMeasurement] = Field(min_length=1)