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(