diff --git a/docs/source/guides/6_sparsity.rst b/docs/source/guides/6_sparsity.rst index 0e018a622e6..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, 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. +: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 92e401e157f..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,6 +57,18 @@ class _DASCModelStructureMismatchError(ApplyModeError): """Identify recoverable policy-versus-GDN-geometry drift during restore.""" +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.""" @@ -90,19 +107,12 @@ 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_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 @@ -192,9 +202,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( @@ -216,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 @@ -273,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 @@ -293,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) @@ -461,8 +487,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, diff --git a/tests/unit/torch/sparsity/state_sparsity/test_dasc.py b/tests/unit/torch/sparsity/state_sparsity/test_dasc.py index 8f1f0fa00cf..6fdb0074ad4 100644 --- a/tests/unit/torch/sparsity/state_sparsity/test_dasc.py +++ b/tests/unit/torch/sparsity/state_sparsity/test_dasc.py @@ -498,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(