From 6e86fd30f9219896f64bbdebe87aedf007138a4f Mon Sep 17 00:00:00 2001 From: kaix-nv Date: Thu, 10 Sep 2026 20:22:27 -0700 Subject: [PATCH 1/2] Validate DASC decay tensors consistently Signed-off-by: kaix-nv --- docs/source/guides/6_sparsity.rst | 6 +-- .../torch/sparsity/state_sparsity/policy.py | 40 ++++++++++++------- .../sparsity/state_sparsity/test_dasc.py | 8 ++++ 3 files changed, 37 insertions(+), 17 deletions(-) diff --git a/docs/source/guides/6_sparsity.rst b/docs/source/guides/6_sparsity.rst index 0e018a622e6..5843073bcf0 100644 --- a/docs/source/guides/6_sparsity.rst +++ b/docs/source/guides/6_sparsity.rst @@ -192,9 +192,9 @@ Set ``decay_parameter_storage_dtype`` to the checkpoint dtype for ``A_log`` and calibration. Derive the evaluated head masks and reported ``retained_heads``/``total_heads`` from :func:`~modelopt.torch.sparsity.state_sparsity.analyze_gdn_decay` using that same storage dtype. Policy validation allows only the rounding introduced by the 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. +tensor dtype. When they differ, their inverse rounding bounds are composed in sequence; when they +match, the duplicate cast is counted once. 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/policy.py b/modelopt/torch/sparsity/state_sparsity/policy.py index 92e401e157f..111205ff4e5 100644 --- a/modelopt/torch/sparsity/state_sparsity/policy.py +++ b/modelopt/torch/sparsity/state_sparsity/policy.py @@ -52,6 +52,24 @@ class _DASCModelStructureMismatchError(ApplyModeError): """Identify recoverable policy-versus-GDN-geometry drift during restore.""" +def _validated_gdn_decay_tensors( + a_log: torch.Tensor, dt_bias: torch.Tensor +) -> tuple[torch.Tensor, torch.Tensor]: + """Validate decay tensors and return deterministic CPU float64 values.""" + 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") + + 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") + return a_log_cpu, dt_bias_cpu + + @lru_cache(maxsize=1) def _supported_gdn_classes() -> tuple[type[nn.Module], ...]: """Resolve installed GDN implementations without making either framework mandatory.""" @@ -90,19 +108,10 @@ 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)") - 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") + a_log_cpu, dt_bias_cpu = _validated_gdn_decay_tensors(a_log, dt_bias) decay = -torch.exp(a_log_cpu) * F.softplus(dt_bias_cpu + static_gate_input) horizons = torch.log(torch.tensor(epsilon, dtype=torch.float64)) / decay @@ -192,9 +201,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: + _validated_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( @@ -461,8 +469,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: + _validated_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, diff --git a/tests/unit/torch/sparsity/state_sparsity/test_dasc.py b/tests/unit/torch/sparsity/state_sparsity/test_dasc.py index 8f1f0fa00cf..17f85a3da5a 100644 --- a/tests/unit/torch/sparsity/state_sparsity/test_dasc.py +++ b/tests/unit/torch/sparsity/state_sparsity/test_dasc.py @@ -474,6 +474,14 @@ def test_bf16_storage_round_trip_loaded_in_fp32_preserves_policy(): with pytest.raises(ApplyModeError, match="floating-point dtype"): mtss.export_policy(model) + invalid = mtss.calibrate( + TinyGatedDeltaNetForCausalLM(), _config(wmax_candidates=[7]), [_candidate(7)] + ) + with torch.no_grad(): + invalid.linear_attn.dt_bias[0] = torch.nan + with pytest.raises(ApplyModeError, match="GDN decay parameters must be finite"): + mtss.export_policy(invalid) + @pytest.mark.parametrize( ("storage_name", "storage_dtype", "live_dtype"), From 02d8cfe7332525e22dae57bdcfd0381b46cb4c78 Mon Sep 17 00:00:00 2001 From: kaix-nv Date: Thu, 10 Sep 2026 20:50:32 -0700 Subject: [PATCH 2/2] Tighten DASC calibration contract (#2391) Closes the remaining review-contract and cleanup findings from #2389 and #2390. Changes: - make the documented measurement recipe forward all three horizon inputs: epsilon, static gate input, and decay-parameter storage dtype - use one shared storage-dtype type alias across config, policy metadata, public analysis, and runtime dtype mapping - split validation-only checks from CPU float64 materialization - finish the storage/live cast-bound documentation and clarify float32 contribution - give the NaN export regression its own accurately named test Validation: - focused DASC tests: 25 passed, 1 skipped (optional Megatron dependency) - pre-commit on all touched files: passed Commit is ED25519-signed and carries a matching Signed-off-by trailer. ## Summary by CodeRabbit - **Bug Fixes** - Improved decay-parameter validation across storage and live tensor data types, including more accurate rounding bounds. - Added validation for non-finite decay parameters during policy export. - **Documentation** - Updated sparsity guidance with epsilon and static gate input configuration examples. - Clarified how dtype conversions and rounding affect policy validation. - **Refactor** - Improved consistency and reliability of decay analysis and policy validation without changing supported dtype options or defaults. --------- Signed-off-by: kaix-nv --- docs/source/guides/6_sparsity.rst | 15 ++++-- .../torch/sparsity/state_sparsity/config.py | 6 ++- .../torch/sparsity/state_sparsity/policy.py | 54 ++++++++++++------- .../sparsity/state_sparsity/test_dasc.py | 31 ++++++++--- 4 files changed, 73 insertions(+), 33 deletions(-) diff --git a/docs/source/guides/6_sparsity.rst b/docs/source/guides/6_sparsity.rst index 5843073bcf0..f02962e1878 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", @@ -150,6 +152,8 @@ ordinary dense recurrent state before continuation. # 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 = [ @@ -190,11 +194,12 @@ supported GDN architecture fails closed on both save and restore. In every stale :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 -:func:`~modelopt.torch.sparsity.state_sparsity.analyze_gdn_decay` using that same storage dtype. -Policy validation allows only the rounding introduced by the declared storage dtype and the live -tensor dtype. When they differ, their inverse rounding bounds are composed in sequence; when they -match, the duplicate cast is counted once. 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. +: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..b5c47e49a63 100644 --- a/modelopt/torch/sparsity/state_sparsity/config.py +++ b/modelopt/torch/sparsity/state_sparsity/config.py @@ -30,6 +30,8 @@ "DASCQualityMeasurement", ] +_DecayParameterStorageDtype = Literal["float16", "bfloat16", "float32"] + class DASCQualityMeasurement(ModeloptBaseConfig): """Quality and lifecycle measurements for one calibration slice.""" @@ -86,7 +88,7 @@ class DASCConfig(ModeloptBaseConfig): default=-0.3, 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.", ) @@ -206,7 +208,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/policy.py b/modelopt/torch/sparsity/state_sparsity/policy.py index 111205ff4e5..2b7df10aa4b 100644 --- a/modelopt/torch/sparsity/state_sparsity/policy.py +++ b/modelopt/torch/sparsity/state_sparsity/policy.py @@ -23,7 +23,6 @@ import warnings from collections.abc import Iterable from functools import lru_cache -from typing import Literal import torch import torch.nn.functional as F @@ -33,7 +32,13 @@ 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, +) __all__ = ["analyze_gdn_decay", "compute_gdn_decay_horizons"] @@ -41,7 +46,7 @@ ("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, @@ -52,22 +57,16 @@ class _DASCModelStructureMismatchError(ApplyModeError): """Identify recoverable policy-versus-GDN-geometry drift during restore.""" -def _validated_gdn_decay_tensors( - a_log: torch.Tensor, dt_bias: torch.Tensor -) -> tuple[torch.Tensor, torch.Tensor]: - """Validate decay tensors and return deterministic CPU float64 values.""" +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") - - 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(): + if not torch.isfinite(a_log).all() or not torch.isfinite(dt_bias).all(): raise ValueError("GDN decay parameters must be finite") - return a_log_cpu, dt_bias_cpu @lru_cache(maxsize=1) @@ -111,7 +110,9 @@ def compute_gdn_decay_horizons( if not 0.0 < epsilon < 1.0: raise ValueError("epsilon must be in (0, 1)") - a_log_cpu, dt_bias_cpu = _validated_gdn_decay_tensors(a_log, dt_bias) + _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 @@ -202,7 +203,7 @@ def _analyze_gdn_modules( dt_bias = module.dt_bias try: if storage_dtype is not None: - _validated_gdn_decay_tensors(a_log, dt_bias) + _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( @@ -224,7 +225,7 @@ def analyze_gdn_decay( *, epsilon: float = 1e-3, static_gate_input: float = -0.3, - decay_parameter_storage_dtype: Literal["float16", "bfloat16", "float32"] | None = None, + decay_parameter_storage_dtype: _DecayParameterStorageDtype | None = None, ) -> dict[str, list[float]]: """Return per-head horizons, optionally canonicalized to a checkpoint storage dtype.""" storage_dtype = None @@ -281,11 +282,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 @@ -301,7 +319,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) @@ -470,7 +488,7 @@ def validate_dasc_decay_parameters(model: nn.Module, policy: DASCPolicy) -> None for name, module in modules.items(): layer = policy.layers[name] try: - _validated_gdn_decay_tensors(module.A_log, module.dt_bias) + _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}" diff --git a/tests/unit/torch/sparsity/state_sparsity/test_dasc.py b/tests/unit/torch/sparsity/state_sparsity/test_dasc.py index 17f85a3da5a..6fdb0074ad4 100644 --- a/tests/unit/torch/sparsity/state_sparsity/test_dasc.py +++ b/tests/unit/torch/sparsity/state_sparsity/test_dasc.py @@ -474,14 +474,6 @@ def test_bf16_storage_round_trip_loaded_in_fp32_preserves_policy(): with pytest.raises(ApplyModeError, match="floating-point dtype"): mtss.export_policy(model) - invalid = mtss.calibrate( - TinyGatedDeltaNetForCausalLM(), _config(wmax_candidates=[7]), [_candidate(7)] - ) - with torch.no_grad(): - invalid.linear_attn.dt_bias[0] = torch.nan - with pytest.raises(ApplyModeError, match="GDN decay parameters must be finite"): - mtss.export_policy(invalid) - @pytest.mark.parametrize( ("storage_name", "storage_dtype", "live_dtype"), @@ -506,6 +498,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(