From d512da24f83077136e43a2fc45a5debb1aa40798 Mon Sep 17 00:00:00 2001 From: kaix-nv Date: Thu, 10 Sep 2026 20:10:49 -0700 Subject: [PATCH 1/2] Fix DASC storage-boundary validation Signed-off-by: kaix-nv --- docs/source/guides/6_sparsity.rst | 4 +- .../sparsity/state_sparsity/conversion.py | 9 ++- .../torch/sparsity/state_sparsity/policy.py | 67 +++++++++++++------ .../sparsity/state_sparsity/test_dasc.py | 33 +++++++++ 4 files changed, 90 insertions(+), 23 deletions(-) diff --git a/docs/source/guides/6_sparsity.rst b/docs/source/guides/6_sparsity.rst index 691b637f022..979d354cd8c 100644 --- a/docs/source/guides/6_sparsity.rst +++ b/docs/source/guides/6_sparsity.rst @@ -183,7 +183,9 @@ policy remains serializable so it cannot block saving or composing other ModelOp :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. Policy validation allows only the rounding introduced by that declared storage dtype -and the live tensor dtype; the default ``float32`` keeps unconfigured policies strict. +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. .. _sparsity-concepts: diff --git a/modelopt/torch/sparsity/state_sparsity/conversion.py b/modelopt/torch/sparsity/state_sparsity/conversion.py index f826f09b9df..bf8a607ce1b 100644 --- a/modelopt/torch/sparsity/state_sparsity/conversion.py +++ b/modelopt/torch/sparsity/state_sparsity/conversion.py @@ -27,7 +27,12 @@ from modelopt.torch.utils import unwrap_model from .config import DASCCalibrationMeasurement, DASCConfig, DASCPolicy -from .policy import build_dasc_policy, validate_dasc_decay_parameters, validate_dasc_model_structure +from .policy import ( + _DASCModelStructureMismatchError, + build_dasc_policy, + validate_dasc_decay_parameters, + validate_dasc_model_structure, +) __all__ = [] @@ -93,7 +98,7 @@ def restore_dasc_model(model: nn.Module, config: DASCConfig, metadata: MetadataD try: validate_dasc_model_structure(model, policy) - except ApplyModeError as error: + except _DASCModelStructureMismatchError as error: warnings.warn( f"{error}. The restored 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 098dbcb4a59..987f66d7961 100644 --- a/modelopt/torch/sparsity/state_sparsity/policy.py +++ b/modelopt/torch/sparsity/state_sparsity/policy.py @@ -17,6 +17,7 @@ import hashlib import importlib +import importlib.util import json import math import warnings @@ -46,6 +47,10 @@ } +class _DASCModelStructureMismatchError(ApplyModeError): + """Identify recoverable policy-versus-GDN-geometry drift during restore.""" + + @lru_cache(maxsize=1) def _supported_gdn_classes() -> tuple[type[nn.Module], ...]: """Resolve installed GDN implementations without making either framework mandatory.""" @@ -174,15 +179,26 @@ def _get_gdn_modules(model: nn.Module) -> dict[str, nn.Module]: def _analyze_gdn_modules( - modules: dict[str, nn.Module], *, epsilon: float, static_gate_input: float + modules: dict[str, nn.Module], + *, + epsilon: float, + static_gate_input: float, + storage_dtype: torch.dtype | None = None, ) -> dict[str, list[float]]: - """Compute deterministic per-head horizons for already-resolved GDN modules.""" + """Compute horizons, optionally from checkpoint-storage-canonical parameters.""" horizons = {} for name, module in modules.items(): + 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: + 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( - module.A_log, - module.dt_bias, + a_log, + dt_bias, epsilon=epsilon, static_gate_input=static_gate_input, ) @@ -348,7 +364,10 @@ def build_dasc_policy( modules = _get_gdn_modules(model) horizons = _analyze_gdn_modules( - modules, epsilon=config.epsilon, static_gate_input=config.static_gate_input + modules, + epsilon=config.epsilon, + static_gate_input=config.static_gate_input, + storage_dtype=_STORAGE_DTYPES[config.decay_parameter_storage_dtype], ) _validate_measurement_geometry(horizons, validated_measurements) @@ -413,34 +432,42 @@ def validate_dasc_model_structure(model: nn.Module, policy: DASCPolicy) -> None: actual_structure != policy_structure or _canonical_sha256(actual_structure) != policy.model_structure_sha256 ): - raise ApplyModeError("DASC policy does not match the model's GDN module structure") + raise _DASCModelStructureMismatchError( + "DASC policy does not match the model's GDN module structure" + ) def validate_dasc_decay_parameters(model: nn.Module, policy: DASCPolicy) -> None: """Reject deployment when current decay parameters no longer derive the stored policy. - Numerical validation deliberately uses re-derived horizons and the exact selected mask rather - than the provenance digest because an FP16 or BF16 storage cast is lossy. + Numerical validation uses inverse cast bounds rather than the provenance digest because an + FP16 or BF16 storage cast is lossy. A stored mask is rejected only when its head's complete + admissible horizon interval lies on the opposite side of the strict ``horizon > Wmax`` rule. """ modules = _get_gdn_modules(model) - current_horizons = _analyze_gdn_modules( - modules, - epsilon=policy.epsilon, - static_gate_input=policy.static_gate_input, - ) - for name, values in current_horizons.items(): + for name, module in modules.items(): layer = policy.layers[name] - retained = [head for head, horizon in enumerate(values) if horizon > policy.selected_wmax] - if retained != layer.retained_heads: - raise ApplyModeError( - f"DASC policy head mask does not match current decay parameters in layer {name!r}" - ) + 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") lower, upper = _storage_cast_horizon_bounds( - modules[name], + module, epsilon=policy.epsilon, static_gate_input=policy.static_gate_input, storage_dtype=_STORAGE_DTYPES[policy.decay_parameter_storage_dtype], ) + declared_retained = set(layer.retained_heads) + for head, (head_lower, head_upper) in enumerate(zip(lower, upper)): + retained_is_impossible = ( + head in declared_retained and head_upper <= policy.selected_wmax + ) + omitted_is_impossible = ( + head not in declared_retained and head_lower > policy.selected_wmax + ) + if retained_is_impossible or omitted_is_impossible: + raise ApplyModeError( + "DASC policy head mask does not match current decay parameters in layer " + f"{name!r}" + ) stored = torch.tensor(layer.static_horizons, dtype=torch.float64) numerical_slack = 32.0 * torch.finfo(torch.float64).eps if torch.any(stored < lower * (1.0 - numerical_slack)) or torch.any( diff --git a/tests/unit/torch/sparsity/state_sparsity/test_dasc.py b/tests/unit/torch/sparsity/state_sparsity/test_dasc.py index 06379e856bf..343b3d804cf 100644 --- a/tests/unit/torch/sparsity/state_sparsity/test_dasc.py +++ b/tests/unit/torch/sparsity/state_sparsity/test_dasc.py @@ -413,6 +413,36 @@ def test_dtype_cast_preserves_policy_when_the_selected_mask_is_unchanged(dtype): assert mtss.export_policy(model) == policy +def test_calibration_uses_storage_canonical_mask_at_wmax_boundary(): + """Derive the mask from stored decay values and accept their explained boundary flip.""" + model = TinyGatedDeltaNetForCausalLM() + with torch.no_grad(): + model.linear_attn.A_log[0] = 0.0 + model.linear_attn.dt_bias[0] = 0.520263671875 + live_horizon = mtss.compute_gdn_decay_horizons( + model.linear_attn.A_log, model.linear_attn.dt_bias, static_gate_input=0.0 + )[0] + stored_horizon = mtss.compute_gdn_decay_horizons( + model.linear_attn.A_log.to(torch.float16), + model.linear_attn.dt_bias.to(torch.float16), + static_gate_input=0.0, + )[0] + assert live_horizon > 7 + assert stored_horizon < 7 + + measurement = _candidate(7) + measurement["retained_heads"] = 0 + model = mtss.calibrate( + model, + _config(wmax_candidates=[7], decay_parameter_storage_dtype="float16"), + [measurement], + ) + policy = mtss.export_policy(model) + + assert policy["layers"]["linear_attn"]["retained_heads"] == [] + assert policy["layers"]["linear_attn"]["static_horizons"][0] < 7 + + 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( @@ -508,6 +538,9 @@ def test_export_rejects_changed_decay_parameters_and_restore_rejects_structure() with pytest.raises(ApplyModeError, match="module structure"): mtss.export_policy(mismatched) + with pytest.raises(ApplyModeError, match="no supported GDN modules"): + mto.restore_from_modelopt_state(nn.Linear(2, 2), state) + tampered_state = copy.deepcopy(state) tampered_state["modelopt_state_dict"][0][1]["metadata"]["policy"]["quality_gates"][ "min_top1_agreement" From d50bd5bd6a9f2f7c1c5e26ca3288e6004ce02cb5 Mon Sep 17 00:00:00 2001 From: kaix-nv Date: Thu, 10 Sep 2026 21:51:22 -0700 Subject: [PATCH 2/2] Align DASC analysis and lifecycle contracts (#2389) Addresses the follow-up review findings on #2388. Changes: - expose optional `decay_parameter_storage_dtype` canonicalization through public `analyze_gdn_decay`, so callers can derive exactly the geometry consumed by calibration - document that evaluated masks and measurement head counts must use the same declared storage dtype - make unsupported-architecture save and restore both fail closed - keep GDN geometry/decay drift serializable with a warning and deployment-blocked until recalibration - add public analysis and save/restore symmetry regressions Validation: - focused DASC tests: 24 passed, 1 skipped (optional Megatron dependency) - state + weight + attention sparsity tests: 302 passed, 1 skipped - pre-commit on all touched files: passed Commit is ED25519-signed and carries a matching Signed-off-by trailer. ## Summary by CodeRabbit - **New Features** - Calibration now accounts for checkpoint storage formats (float16, bfloat16, and float32) when analyzing decay behavior. - Redundant precision-rounding slack is excluded when storage formats are compatible. - Invalid decay parameters and storage-format selections are clearly rejected. - **Bug Fixes** - Recoverable structure or decay mismatches now allow metadata updates to continue with warnings. - Checkpoint serialization fails safely when a calibrated GDN layer is replaced with an unsupported module. - **Documentation** - Clarified calibration parameter requirements and sequential handling of rounding bounds. --------- Signed-off-by: kaix-nv --- docs/source/guides/6_sparsity.rst | 27 +++- .../torch/sparsity/state_sparsity/config.py | 48 +++++-- .../sparsity/state_sparsity/conversion.py | 17 ++- .../torch/sparsity/state_sparsity/policy.py | 102 ++++++++++---- .../sparsity/state_sparsity/test_dasc.py | 133 +++++++++++++++++- 5 files changed, 279 insertions(+), 48 deletions(-) diff --git a/docs/source/guides/6_sparsity.rst b/docs/source/guides/6_sparsity.rst index 979d354cd8c..70166aecac4 100644 --- a/docs/source/guides/6_sparsity.rst +++ b/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", @@ -146,6 +148,14 @@ ordinary dense recurrent state before continuation. "model_config_id": "sha256:", "calibration_data_id": "sha256:", } + # Use these storage-canonical horizons to derive the evaluated mask and the + # 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 = [ { "variant": "dasc_wr", @@ -178,14 +188,19 @@ The initial GDN adapter accepts the ``GatedDeltaNet`` and ``Qwen3NextGatedDeltaN including ModelOpt-generated dynamic subclasses, and fails closed for unrelated implementations 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 stale -policy remains serializable so it cannot block saving or composing other ModelOpt modes, but +mode-state entry and supersedes its stale policy without growing the checkpoint history. A policy +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. Policy validation allows only the rounding introduced by that 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. +calibration. Derive the evaluated head masks and reported ``retained_heads``/``total_heads`` from +: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: diff --git a/modelopt/torch/sparsity/state_sparsity/config.py b/modelopt/torch/sparsity/state_sparsity/config.py index 7f67d7f7d47..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 @@ -30,6 +31,39 @@ "DASCQualityMeasurement", ] +_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): """Quality and lifecycle measurements for one calibration slice.""" @@ -79,14 +113,14 @@ 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: 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.", ) @@ -116,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") @@ -206,7 +238,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] diff --git a/modelopt/torch/sparsity/state_sparsity/conversion.py b/modelopt/torch/sparsity/state_sparsity/conversion.py index bf8a607ce1b..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, @@ -108,16 +108,23 @@ def restore_dasc_model(model: nn.Module, config: DASCConfig, metadata: MetadataD def update_dasc_metadata(model: nn.Module, config: DASCConfig, metadata: MetadataDict) -> None: - """Refresh metadata without making unrelated ModelOpt save or compose paths unusable.""" + """Refresh metadata while allowing recoverable policy staleness to remain serializable.""" policy = get_attached_dasc_policy(model) try: validate_dasc_model_structure(model, policy) - validate_dasc_decay_parameters(model, policy) - except ApplyModeError as error: + except _DASCRecoverableStalenessError as error: warnings.warn( f"{error}. The saved DASC policy is stale; re-run calibrate() before deployment", stacklevel=2, ) + else: + try: + validate_dasc_decay_parameters(model, policy) + except ApplyModeError as error: + warnings.warn( + f"{error}. The saved DASC policy is stale; re-run calibrate() before deployment", + stacklevel=2, + ) metadata.clear() metadata["policy"] = copy.deepcopy(policy.model_dump(mode="json")) diff --git a/modelopt/torch/sparsity/state_sparsity/policy.py b/modelopt/torch/sparsity/state_sparsity/policy.py index 987f66d7961..7baf53e9586 100644 --- a/modelopt/torch/sparsity/state_sparsity/policy.py +++ b/modelopt/torch/sparsity/state_sparsity/policy.py @@ -32,7 +32,14 @@ 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, + _validate_analysis_arguments, +) __all__ = ["analyze_gdn_decay", "compute_gdn_decay_horizons"] @@ -40,17 +47,37 @@ ("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, } -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(): + 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.""" @@ -89,22 +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 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_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) - 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 + 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 @@ -132,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)}" ) @@ -191,9 +209,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( @@ -215,10 +232,24 @@ def analyze_gdn_decay( *, epsilon: float = 1e-3, static_gate_input: float = -0.3, + decay_parameter_storage_dtype: _DecayParameterStorageDtype | None = None, ) -> dict[str, list[float]]: - """Return deterministic per-head horizons for every GDN module in a model.""" + """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: + 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}") + storage_dtype = _STORAGE_DTYPES[decay_parameter_storage_dtype] return _analyze_gdn_modules( - _get_gdn_modules(model), epsilon=epsilon, static_gate_input=static_gate_input + _get_gdn_modules(model), + epsilon=epsilon, + static_gate_input=static_gate_input, + storage_dtype=storage_dtype, ) @@ -259,11 +290,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 @@ -279,7 +327,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) @@ -447,8 +495,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, @@ -468,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 343b3d804cf..cad44b1501d 100644 --- a/tests/unit/torch/sparsity/state_sparsity/test_dasc.py +++ b/tests/unit/torch/sparsity/state_sparsity/test_dasc.py @@ -422,11 +422,11 @@ def test_calibration_uses_storage_canonical_mask_at_wmax_boundary(): live_horizon = mtss.compute_gdn_decay_horizons( model.linear_attn.A_log, model.linear_attn.dt_bias, static_gate_input=0.0 )[0] - stored_horizon = mtss.compute_gdn_decay_horizons( - model.linear_attn.A_log.to(torch.float16), - model.linear_attn.dt_bias.to(torch.float16), + stored_horizon = mtss.analyze_gdn_decay( + model, static_gate_input=0.0, - )[0] + decay_parameter_storage_dtype="float16", + )["linear_attn"][0] assert live_horizon > 7 assert stored_horizon < 7 @@ -443,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( @@ -496,6 +572,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( @@ -541,6 +640,13 @@ def test_export_rejects_changed_decay_parameters_and_restore_rejects_structure() with pytest.raises(ApplyModeError, match="no supported GDN modules"): mto.restore_from_modelopt_state(nn.Linear(2, 2), state) + unsupported = mtss.calibrate( + TinyGatedDeltaNetForCausalLM(), _config(wmax_candidates=[7]), [_candidate(7)] + ) + unsupported.linear_attn = nn.Linear(2, 2) + with pytest.raises(ApplyModeError, match="no supported GDN modules"): + mto.modelopt_state(unsupported) + tampered_state = copy.deepcopy(state) tampered_state["modelopt_state_dict"][0][1]["metadata"]["policy"]["quality_gates"][ "min_top1_agreement" @@ -594,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(