Align DASC numeric validation with scalar API

modelopt/torch/sparsity/state_sparsity/config.py

@@ -16,7 +16,8 @@
 """Configuration and result schemas for DASC state sparsity."""
 
 import math
-from typing import Any, Literal
+from numbers import Real
+from typing import Literal
 
 from pydantic import ConfigDict, Field, field_validator, model_validator
 
@@ -36,20 +37,29 @@
 
 
 def _validate_analysis_arguments(
-    epsilon: Any = _DEFAULT_EPSILON,
-    static_gate_input: Any = _DEFAULT_STATIC_GATE_INPUT,
+    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 = math.isfinite(epsilon) and 0.0 < epsilon < 1.0
-    except (TypeError, ValueError, OverflowError):
+        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 = math.isfinite(static_gate_input)
-    except (TypeError, ValueError, OverflowError):
+        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")

modelopt/torch/sparsity/state_sparsity/policy.py

@@ -122,7 +122,7 @@ def compute_gdn_decay_horizons(
     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.as_tensor(epsilon, device="cpu", 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
@@ -520,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

@@ -454,14 +454,16 @@ def test_analysis_arguments_fail_at_the_public_boundary(invalid_storage_dtype):
         )
 
 
-@pytest.mark.parametrize("epsilon", [1.0, [], 10**1000])
+@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, [], 10**1000])
+@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"):
@@ -475,11 +477,15 @@ def test_analysis_rejects_invalid_static_gate_input_at_the_public_boundary(stati
     ("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):
@@ -493,16 +499,24 @@ def test_horizon_computation_rejects_invalid_public_arguments(argument, value, m
         )
 
 
-def test_analysis_accepts_tensor_scalar_arguments():
-    """Preserve support for real-like scalar values accepted by the numeric operations."""
-    horizons = mtss.compute_gdn_decay_horizons(
-        torch.tensor([0.0]),
-        torch.tensor([0.0]),
-        epsilon=torch.tensor(1e-3),  # type: ignore[arg-type]
-        static_gate_input=torch.tensor(-0.3),  # type: ignore[arg-type]
-    )
-    assert horizons.shape == (1,)
-    assert torch.isfinite(horizons).all()
+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():