Address DASC policy 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
@@ -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:
 

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/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"]

modelopt/torch/sparsity/state_sparsity/conversion.py

@@ -16,13 +16,15 @@
 """ModelOpt conversion and restoration for DASC policy metadata."""
 
 import copy
+import warnings
 from collections.abc import Iterable
 
 from pydantic import ValidationError
 from torch import nn
 
 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,16 +35,22 @@
 
 
 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"))
 
 
 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:

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."""

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}"
+            )