From 4d05a1b013fe58613efa4511f1ed50d2bb71ead5 Mon Sep 17 00:00:00 2001 From: kaix-nv Date: Thu, 10 Sep 2026 21:01:28 -0700 Subject: [PATCH 1/2] Harden DASC analysis and restore boundaries Signed-off-by: kaix-nv --- docs/source/guides/6_sparsity.rst | 5 +- .../sparsity/state_sparsity/conversion.py | 6 +-- .../torch/sparsity/state_sparsity/policy.py | 39 +++++++++++---- .../sparsity/state_sparsity/test_dasc.py | 49 ++++++++++++++++++- 4 files changed, 83 insertions(+), 16 deletions(-) diff --git a/docs/source/guides/6_sparsity.rst b/docs/source/guides/6_sparsity.rst index f02962e1878..70166aecac4 100644 --- a/docs/source/guides/6_sparsity.rst +++ b/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 diff --git a/modelopt/torch/sparsity/state_sparsity/conversion.py b/modelopt/torch/sparsity/state_sparsity/conversion.py index f5977d6c293..213cc351f45 100644 --- a/modelopt/torch/sparsity/state_sparsity/conversion.py +++ b/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, diff --git a/modelopt/torch/sparsity/state_sparsity/policy.py b/modelopt/torch/sparsity/state_sparsity/policy.py index 2b7df10aa4b..27914bc2436 100644 --- a/modelopt/torch/sparsity/state_sparsity/policy.py +++ b/modelopt/torch/sparsity/state_sparsity/policy.py @@ -53,10 +53,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(): @@ -143,7 +151,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)}" ) @@ -202,8 +210,9 @@ def _analyze_gdn_modules( a_log = module.A_log dt_bias = module.dt_bias try: + # Check the original tensors before a storage cast can hide an invalid integer dtype. + _validate_gdn_decay_tensors(a_log, dt_bias) 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( @@ -228,15 +237,27 @@ 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.""" + try: + epsilon_is_valid = math.isfinite(epsilon) and 0.0 < epsilon < 1.0 + except TypeError: + 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: + static_gate_input_is_valid = False + if not static_gate_input_is_valid: + raise ValueError("static_gate_input must be finite") 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, diff --git a/tests/unit/torch/sparsity/state_sparsity/test_dasc.py b/tests/unit/torch/sparsity/state_sparsity/test_dasc.py index 6fdb0074ad4..e5362d7c9e4 100644 --- a/tests/unit/torch/sparsity/state_sparsity/test_dasc.py +++ b/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,34 @@ 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, []]) +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, []]) +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] + ) + + 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 +652,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( From c2da5fa38237210a58f5e097706fcda03ec6a424 Mon Sep 17 00:00:00 2001 From: kaix-nv Date: Thu, 10 Sep 2026 21:50:47 -0700 Subject: [PATCH 2/2] Unify DASC public analysis validation (#2395) ## Summary Addresses the two follow-up findings on #2394: - use one argument validator for both exported analysis entry points - normalize invalid epsilon and static gate inputs to ValueError before tensor/model work - remove the redundant validation hoist and retain the existing pre-cast check only where a storage cast occurs - add direct public-API regressions for `compute_gdn_decay_horizons` ## Validation - `python -m pytest -q tests/unit/torch/sparsity/state_sparsity/test_dasc.py` (36 passed, 1 optional Megatron skip) - pre-commit hooks on both changed files - signed commit with DCO sign-off ## Summary by CodeRabbit - **Bug Fixes** - Improved validation for sparsity analysis inputs, including finite numeric values and supported epsilon ranges. - Ensured invalid decay-tensor values are detected consistently across analysis workflows. - Standardized error handling for direct decay-horizon calculations. - **Tests** - Added coverage for invalid epsilon and static gate input values, including verification of consistent error messages. --------- Signed-off-by: kaix-nv --- .../torch/sparsity/state_sparsity/config.py | 42 ++++++++++++--- .../torch/sparsity/state_sparsity/policy.py | 25 +++------ .../sparsity/state_sparsity/test_dasc.py | 52 ++++++++++++++++++- 3 files changed, 92 insertions(+), 27 deletions(-) diff --git a/modelopt/torch/sparsity/state_sparsity/config.py b/modelopt/torch/sparsity/state_sparsity/config.py index b5c47e49a63..1a8f87e0236 100644 --- a/modelopt/torch/sparsity/state_sparsity/config.py +++ b/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") diff --git a/modelopt/torch/sparsity/state_sparsity/policy.py b/modelopt/torch/sparsity/state_sparsity/policy.py index 27914bc2436..7baf53e9586 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"] @@ -115,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 @@ -210,9 +209,8 @@ def _analyze_gdn_modules( a_log = module.A_log dt_bias = module.dt_bias try: - # Check the original tensors before a storage cast can hide an invalid integer dtype. - _validate_gdn_decay_tensors(a_log, dt_bias) 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( @@ -237,18 +235,7 @@ 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.""" - try: - epsilon_is_valid = math.isfinite(epsilon) and 0.0 < epsilon < 1.0 - except TypeError: - 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: - static_gate_input_is_valid = False - if not static_gate_input_is_valid: - raise ValueError("static_gate_input must be finite") + _validate_analysis_arguments(epsilon, static_gate_input) storage_dtype = None if decay_parameter_storage_dtype is not None: if ( @@ -533,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 e5362d7c9e4..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, []]) +@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, []]) +@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"): @@ -471,6 +473,52 @@ def test_analysis_rejects_invalid_static_gate_input_at_the_public_boundary(stati ) +@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(