From 487de7ad6efad6265ea8a656f74663ad3c9e3118 Mon Sep 17 00:00:00 2001 From: kaix-nv Date: Thu, 10 Sep 2026 21:18:19 -0700 Subject: [PATCH 1/2] Preserve DASC numeric argument compatibility Signed-off-by: kaix-nv --- .../torch/sparsity/state_sparsity/config.py | 34 +++++++++++++++---- .../torch/sparsity/state_sparsity/policy.py | 11 ++---- .../sparsity/state_sparsity/test_dasc.py | 20 +++++++++-- 3 files changed, 47 insertions(+), 18 deletions(-) diff --git a/modelopt/torch/sparsity/state_sparsity/config.py b/modelopt/torch/sparsity/state_sparsity/config.py index b5c47e49a63..3f2bc98b72d 100644 --- a/modelopt/torch/sparsity/state_sparsity/config.py +++ b/modelopt/torch/sparsity/state_sparsity/config.py @@ -16,7 +16,7 @@ """Configuration and result schemas for DASC state sparsity.""" import math -from typing import Literal +from typing import Any, Literal from pydantic import ConfigDict, Field, field_validator, model_validator @@ -31,6 +31,28 @@ ] _DecayParameterStorageDtype = Literal["float16", "bfloat16", "float32"] +_DEFAULT_EPSILON = 1e-3 +_DEFAULT_STATIC_GATE_INPUT = -0.3 + + +def _validate_analysis_arguments( + epsilon: Any = _DEFAULT_EPSILON, + static_gate_input: Any = _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 = 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 = False + if not static_gate_input_is_valid: + raise ValueError("static_gate_input must be finite") class DASCQualityMeasurement(ModeloptBaseConfig): @@ -81,11 +103,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 +140,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") diff --git a/modelopt/torch/sparsity/state_sparsity/policy.py b/modelopt/torch/sparsity/state_sparsity/policy.py index ba812906bb7..c278d155364 100644 --- a/modelopt/torch/sparsity/state_sparsity/policy.py +++ b/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"] @@ -65,14 +66,6 @@ class _DASCDecayParametersUnavailableError(_DASCRecoverableStalenessError): """Identify supported GDN modules whose decay tensors are temporarily unavailable.""" -def _validate_analysis_arguments(epsilon: float, static_gate_input: float) -> None: - """Normalize invalid public analysis arguments to the ValueError contract.""" - if not isinstance(epsilon, float | int) or not (math.isfinite(epsilon) and 0.0 < epsilon < 1.0): - raise ValueError("epsilon must be finite and in (0, 1)") - if not isinstance(static_gate_input, float | int) or not math.isfinite(static_gate_input): - raise ValueError("static_gate_input must be finite") - - 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(): @@ -129,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.tensor(epsilon, dtype=torch.float64)) / decay + horizons = torch.log(torch.as_tensor(epsilon, device="cpu", dtype=torch.float64)) / 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 diff --git a/tests/unit/torch/sparsity/state_sparsity/test_dasc.py b/tests/unit/torch/sparsity/state_sparsity/test_dasc.py index 259d7c2f0f2..ef620ac8968 100644 --- a/tests/unit/torch/sparsity/state_sparsity/test_dasc.py +++ b/tests/unit/torch/sparsity/state_sparsity/test_dasc.py @@ -454,14 +454,14 @@ def test_analysis_arguments_fail_at_the_public_boundary(invalid_storage_dtype): ) -@pytest.mark.parametrize("epsilon", [1.0, []]) +@pytest.mark.parametrize("epsilon", [1.0, [], 10**1000]) 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, []]) +@pytest.mark.parametrize("static_gate_input", [torch.nan, [], 10**1000]) 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,7 +475,11 @@ 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", torch.nan, r"epsilon must be finite and in \(0, 1\)"), + ("epsilon", 10**1000, r"epsilon must be finite and in \(0, 1\)"), ("static_gate_input", [], "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"), ], ) def test_horizon_computation_rejects_invalid_public_arguments(argument, value, message): @@ -489,6 +493,18 @@ 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_bf16_storage_round_trip_loaded_in_fp32_preserves_policy(): """Accept BF16-rounded values after a checkpoint loader materializes FP32 tensors.""" model = mtss.calibrate( From 2f033a1aee5034f4ca6b542ed8c101d80c994ce8 Mon Sep 17 00:00:00 2001 From: kaix-nv Date: Thu, 10 Sep 2026 21:50:17 -0700 Subject: [PATCH 2/2] Align DASC numeric validation with scalar API (#2397) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary Addresses both Claude findings on #2396 while retaining CodeRabbit’s oversized-integer fix: - match the annotated scalar-float API by accepting `numbers.Real` inputs, including NumPy real scalars - reject all tensor objects before scalar conversion, so multi-element tensors cannot leak RuntimeError - keep OverflowError normalization for oversized integer inputs - remove the tensor-scalar compatibility claim and restore the original scalar conversion - cover multi-element tensors through both exported public entry points ## Validation - focused DASC suite: 46 passed, 1 optional Megatron skip - pre-commit hooks on all three changed files - signed commit with DCO sign-off ## Summary by CodeRabbit - **Bug Fixes** - Improved validation for decay-analysis inputs by requiring real-number values for epsilon and static gate parameters. - Invalid non-numeric and tensor-based inputs are now rejected consistently, including zero-dimensional tensor values. - Improved numerical consistency when calculating decay horizons by preserving CPU double-precision behavior. - Added clearer safeguards against invalid or unsupported input types before analysis calculations are performed. --------- Signed-off-by: kaix-nv --- .../torch/sparsity/state_sparsity/config.py | 24 ++++++++---- .../torch/sparsity/state_sparsity/policy.py | 4 +- .../sparsity/state_sparsity/test_dasc.py | 38 +++++++++++++------ 3 files changed, 45 insertions(+), 21 deletions(-) diff --git a/modelopt/torch/sparsity/state_sparsity/config.py b/modelopt/torch/sparsity/state_sparsity/config.py index 3f2bc98b72d..1a8f87e0236 100644 --- a/modelopt/torch/sparsity/state_sparsity/config.py +++ b/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") diff --git a/modelopt/torch/sparsity/state_sparsity/policy.py b/modelopt/torch/sparsity/state_sparsity/policy.py index c278d155364..7baf53e9586 100644 --- a/modelopt/torch/sparsity/state_sparsity/policy.py +++ b/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) diff --git a/tests/unit/torch/sparsity/state_sparsity/test_dasc.py b/tests/unit/torch/sparsity/state_sparsity/test_dasc.py index ef620ac8968..cad44b1501d 100644 --- a/tests/unit/torch/sparsity/state_sparsity/test_dasc.py +++ b/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():