Validate DASC decay tensors consistently

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",
@@ -150,6 +152,8 @@ ordinary dense recurrent state before continuation.
     # 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 = [
@@ -190,11 +194,12 @@ supported GDN architecture fails closed on both save and restore. In every stale
 :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 that same storage dtype.
-Policy validation allows only the rounding introduced by the 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.
+: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

@@ -30,6 +30,8 @@
     "DASCQualityMeasurement",
 ]
 
+_DecayParameterStorageDtype = Literal["float16", "bfloat16", "float32"]
+
 
 class DASCQualityMeasurement(ModeloptBaseConfig):
     """Quality and lifecycle measurements for one calibration slice."""
@@ -86,7 +88,7 @@ class DASCConfig(ModeloptBaseConfig):
         default=-0.3,
         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.",
     )
@@ -206,7 +208,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/policy.py

@@ -23,7 +23,6 @@
 import warnings
 from collections.abc import Iterable
 from functools import lru_cache
-from typing import Literal
 
 import torch
 import torch.nn.functional as F
@@ -33,15 +32,21 @@
 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,
+)
 
 __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,
@@ -52,6 +57,18 @@ class _DASCModelStructureMismatchError(ApplyModeError):
     """Identify recoverable policy-versus-GDN-geometry drift during restore."""
 
 
+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."""
@@ -90,19 +107,12 @@ 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_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
@@ -192,9 +202,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(
@@ -216,7 +225,7 @@ def analyze_gdn_decay(
     *,
     epsilon: float = 1e-3,
     static_gate_input: float = -0.3,
-    decay_parameter_storage_dtype: Literal["float16", "bfloat16", "float32"] | None = None,
+    decay_parameter_storage_dtype: _DecayParameterStorageDtype | None = None,
 ) -> dict[str, list[float]]:
     """Return per-head horizons, optionally canonicalized to a checkpoint storage dtype."""
     storage_dtype = None
@@ -273,11 +282,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
@@ -293,7 +319,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)
@@ -461,8 +487,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,

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

@@ -498,6 +498,29 @@ def test_cross_dtype_reload_accumulates_both_rounding_bounds(
     assert mtss.export_policy(model) == policy
 
 
+@pytest.mark.parametrize("live_dtype", [torch.float16, torch.bfloat16])
+def test_storage_rounding_excludes_exact_fp32_widening(live_dtype):
+    """Do not add FP32 slack when only the low-precision cast can round values."""
+    tensor = torch.tensor([1.25], dtype=live_dtype)
+
+    assert torch.equal(
+        dasc_policy._storage_rounding_radius(tensor, torch.float32),
+        dasc_policy._storage_rounding_radius(tensor, live_dtype),
+    )
+
+
+def test_non_finite_decay_parameters_are_rejected_on_export():
+    """Reject NaNs before interval comparisons can silently accept them."""
+    model = mtss.calibrate(
+        TinyGatedDeltaNetForCausalLM(), _config(wmax_candidates=[7]), [_candidate(7)]
+    )
+    with torch.no_grad():
+        model.linear_attn.dt_bias[0] = torch.nan
+
+    with pytest.raises(ApplyModeError, match="GDN decay parameters must be finite"):
+        mtss.export_policy(model)
+
+
 def test_export_rejects_changed_decay_parameters_and_restore_rejects_structure():
     """Keep saving recoverable while rejecting stale or tampered deployment policies."""
     model = mtss.calibrate(