Harden DASC analysis and restore boundaries

docs/source/guides/6_sparsity.rst

@@ -189,8 +189,9 @@ including ModelOpt-generated dynamic subclasses, and fails closed for unrelated
 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 or decay drift remains serializable, but removing or replacing the
-supported GDN architecture fails closed on both save and restore. In every stale-policy case,
+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

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
@@ -31,6 +32,37 @@
 ]
 
 _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):
@@ -81,11 +113,11 @@ 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: _DecayParameterStorageDtype = ModeloptField(
@@ -118,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")

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,
@@ -112,7 +112,7 @@ def update_dasc_metadata(model: nn.Module, config: DASCConfig, metadata: Metadat
     policy = get_attached_dasc_policy(model)
     try:
         validate_dasc_model_structure(model, policy)
-    except _DASCModelStructureMismatchError as error:
+    except _DASCRecoverableStalenessError as error:
         warnings.warn(
             f"{error}. The saved DASC policy is stale; re-run calibrate() before deployment",
             stacklevel=2,

modelopt/torch/sparsity/state_sparsity/policy.py

@@ -38,6 +38,7 @@
     DASCLayerPolicy,
     DASCPolicy,
     _DecayParameterStorageDtype,
+    _validate_analysis_arguments,
 )
 
 __all__ = ["analyze_gdn_decay", "compute_gdn_decay_horizons"]
@@ -53,10 +54,18 @@
 }
 
 
-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():
@@ -107,15 +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 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)
 
     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
@@ -143,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)}"
         )
@@ -228,15 +235,16 @@ def analyze_gdn_decay(
     decay_parameter_storage_dtype: _DecayParameterStorageDtype | None = None,
 ) -> dict[str, list[float]]:
     """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:
-        try:
-            storage_dtype = _STORAGE_DTYPES[decay_parameter_storage_dtype]
-        except KeyError as error:
+        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}"
-            ) from error
+            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,
@@ -512,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)

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

@@ -429,8 +429,6 @@ def test_calibration_uses_storage_canonical_mask_at_wmax_boundary():
     )["linear_attn"][0]
     assert live_horizon > 7
     assert stored_horizon < 7
-    with pytest.raises(ValueError, match="decay_parameter_storage_dtype must be one of"):
-        mtss.analyze_gdn_decay(model, decay_parameter_storage_dtype="float8")
 
     measurement = _candidate(7)
     measurement["retained_heads"] = 0
@@ -445,6 +443,82 @@ def test_calibration_uses_storage_canonical_mask_at_wmax_boundary():
     assert policy["layers"]["linear_attn"]["static_horizons"][0] < 7
 
 
+@pytest.mark.parametrize("invalid_storage_dtype", ["float8", []])
+def test_analysis_arguments_fail_at_the_public_boundary(invalid_storage_dtype):
+    """Report invalid analysis arguments uniformly without blaming a GDN module."""
+    model = TinyGatedDeltaNetForCausalLM()
+    with pytest.raises(ValueError, match="decay_parameter_storage_dtype must be one of"):
+        mtss.analyze_gdn_decay(
+            model,
+            decay_parameter_storage_dtype=invalid_storage_dtype,  # type: ignore[arg-type]
+        )
+
+
+@pytest.mark.parametrize("epsilon", [1.0, True, [], 10**1000, torch.tensor([1e-3, 2e-3])])
+def test_analysis_rejects_invalid_epsilon_at_the_public_boundary(epsilon):
+    """Normalize invalid epsilon values to the public ValueError contract."""
+    with pytest.raises(ValueError, match=r"epsilon must be finite and in \(0, 1\)"):
+        mtss.analyze_gdn_decay(TinyGatedDeltaNetForCausalLM(), epsilon=epsilon)  # type: ignore[arg-type]
+
+
+@pytest.mark.parametrize(
+    "static_gate_input", [torch.nan, True, [], 10**1000, torch.tensor([-0.3, -0.2])]
+)
+def test_analysis_rejects_invalid_static_gate_input_at_the_public_boundary(static_gate_input):
+    """Normalize invalid static gate values to the public ValueError contract."""
+    with pytest.raises(ValueError, match="static_gate_input must be finite"):
+        mtss.analyze_gdn_decay(
+            TinyGatedDeltaNetForCausalLM(),
+            static_gate_input=static_gate_input,  # type: ignore[arg-type]
+        )
+
+
+@pytest.mark.parametrize(
+    ("argument", "value", "message"),
+    [
+        ("epsilon", [], r"epsilon must be finite and in \(0, 1\)"),
+        ("epsilon", True, r"epsilon must be finite and in \(0, 1\)"),
+        ("epsilon", torch.nan, r"epsilon must be finite and in \(0, 1\)"),
+        ("epsilon", 10**1000, r"epsilon must be finite and in \(0, 1\)"),
+        ("epsilon", torch.tensor([1e-3, 2e-3]), r"epsilon must be finite and in \(0, 1\)"),
+        ("static_gate_input", [], "static_gate_input must be finite"),
+        ("static_gate_input", True, "static_gate_input must be finite"),
+        ("static_gate_input", torch.nan, "static_gate_input must be finite"),
+        ("static_gate_input", 10**1000, "static_gate_input must be finite"),
+        ("static_gate_input", torch.tensor([-0.3, -0.2]), "static_gate_input must be finite"),
+    ],
+)
+def test_horizon_computation_rejects_invalid_public_arguments(argument, value, message):
+    """Use the same public argument contract for direct horizon computation."""
+    kwargs = {argument: value}
+    with pytest.raises(ValueError, match=message):
+        mtss.compute_gdn_decay_horizons(
+            torch.tensor([0.0]),
+            torch.tensor([0.0]),
+            **kwargs,  # type: ignore[arg-type]
+        )
+
+
+def test_horizon_computation_ignores_the_default_device():
+    """Keep CPU horizon analysis independent of PyTorch's ambient allocation device."""
+    a_log = torch.tensor([0.0])
+    dt_bias = torch.tensor([0.0])
+    with torch.device("meta"):
+        horizons = mtss.compute_gdn_decay_horizons(a_log, dt_bias)
+    assert horizons.device.type == "cpu"
+
+
+def test_policy_lifecycle_ignores_the_default_device():
+    """Keep calibration and checkpoint metadata validation on their declared CPU path."""
+    model = TinyGatedDeltaNetForCausalLM()
+    with torch.device("meta"):
+        calibrated = mtss.calibrate(model, _config(wmax_candidates=[7]), [_candidate(7)])
+        state = mto.modelopt_state(calibrated)
+        policy = mtss.export_policy(calibrated)
+    assert state["modelopt_state_dict"][0][0] == "dasc"
+    assert policy["layers"]["linear_attn"]["static_horizons"]
+
+
 def test_bf16_storage_round_trip_loaded_in_fp32_preserves_policy():
     """Accept BF16-rounded values after a checkpoint loader materializes FP32 tensors."""
     model = mtss.calibrate(
@@ -626,6 +700,25 @@ def test_structure_staleness_does_not_block_checkpoint_save():
         mtss.export_policy(restored)
 
 
+def test_temporarily_unavailable_decay_tensors_are_recoverable_staleness():
+    """Keep save and restore symmetric when a supported GDN is temporarily flattened."""
+    model = mtss.calibrate(
+        TinyGatedDeltaNetForCausalLM(), _config(wmax_candidates=[7]), [_candidate(7)]
+    )
+    state = mto.modelopt_state(model)
+    model.linear_attn.A_log = None
+
+    with pytest.warns(UserWarning, match="saved DASC policy is stale"):
+        mto.modelopt_state(model)
+
+    target = TinyGatedDeltaNetForCausalLM()
+    target.linear_attn.A_log = None
+    with pytest.warns(UserWarning, match="restored DASC policy is stale"):
+        restored = mto.restore_from_modelopt_state(target, state)
+    with pytest.raises(ApplyModeError, match="without A_log and dt_bias tensors"):
+        mtss.export_policy(restored)
+
+
 def test_export_rederives_the_selected_head_mask():
     """Reject a self-consistent stored mask that current decay parameters do not derive."""
     model = mtss.calibrate(