From e8edcf21736c446174d9e067fedb2c3e06febce1 Mon Sep 17 00:00:00 2001 From: kaix-nv Date: Thu, 10 Sep 2026 19:51:31 -0700 Subject: [PATCH 1/2] Address DASC review feedback Signed-off-by: kaix-nv --- docs/source/guides/6_sparsity.rst | 16 +- .../torch/sparsity/state_sparsity/__init__.py | 26 +- modelopt/torch/sparsity/state_sparsity/api.py | 16 +- .../torch/sparsity/state_sparsity/config.py | 29 +- .../sparsity/state_sparsity/conversion.py | 67 ++- .../torch/sparsity/state_sparsity/mode.py | 2 +- .../torch/sparsity/state_sparsity/policy.py | 255 +++++++++-- .../sparsity/state_sparsity/test_dasc.py | 398 +++++++++++++++++- 8 files changed, 741 insertions(+), 68 deletions(-) diff --git a/docs/source/guides/6_sparsity.rst b/docs/source/guides/6_sparsity.rst index d507dd8ba3f..691b637f022 100644 --- a/docs/source/guides/6_sparsity.rst +++ b/docs/source/guides/6_sparsity.rst @@ -123,7 +123,9 @@ derives one static decay horizon per complete GDN head from ``A_log`` and ``dt_bias``, then selects the largest caller-evaluated ``Wmax`` that passes every configured quality, lifecycle, and physical-storage gate. ``Wmax`` may be any positive integer. The measurements are evidence inputs produced by a caller-owned paired evaluation; this API does -not run the dense-versus-recovery suffix evaluation itself. +not run the dense-versus-recovery suffix evaluation itself. ``perplexity_retention`` is defined as +``dense_perplexity / DASC_perplexity`` (equivalently ``exp(dense_NLL - DASC_NLL)``), so higher is +better and values above one are valid. The initial API exports policy metadata only. It does not change model execution, quantize state, pack ragged checkpoints, replay a suffix, or add a linear-attention kernel. A serving integration @@ -137,6 +139,8 @@ ordinary dense recurrent state before continuation. config = { "variant": "dasc_wr", # "dasc_nr" uses zero recovery instead "wmax_candidates": [32], + # Set this to the dtype used to store A_log and dt_bias in the checkpoint. + "decay_parameter_storage_dtype": "bfloat16", "model_id": "org/model", "model_revision": "immutable-model-revision", "model_config_id": "sha256:", @@ -170,6 +174,16 @@ ordinary dense recurrent state before continuation. from zero, while DASC-WR reconstructs them from a zero-initialized suffix replay of at most the selected ``Wmax`` tokens. Both retain whole GDN heads, preserve convolution state, and resume with dense recurrence. KDA and serving-runtime integration are not supported by this initial API. +The initial GDN adapter accepts the ``GatedDeltaNet`` and ``Qwen3NextGatedDeltaNet`` base classes, +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 +: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. .. _sparsity-concepts: diff --git a/modelopt/torch/sparsity/state_sparsity/__init__.py b/modelopt/torch/sparsity/state_sparsity/__init__.py index 97fe2d89b13..b0d1df3cfe0 100644 --- a/modelopt/torch/sparsity/state_sparsity/__init__.py +++ b/modelopt/torch/sparsity/state_sparsity/__init__.py @@ -15,7 +15,25 @@ """Decay-aware sparsity policies for persisted recurrent state.""" -from . import mode -from .api import * -from .config import * -from .policy import * +from . import mode # imported for mode-registration side effects +from .api import calibrate, export_policy +from .config import ( + DASCCalibrationMeasurement, + DASCConfig, + DASCLayerPolicy, + DASCPolicy, + DASCQualityMeasurement, +) +from .policy import analyze_gdn_decay, compute_gdn_decay_horizons + +__all__ = [ + "DASCCalibrationMeasurement", + "DASCConfig", + "DASCLayerPolicy", + "DASCPolicy", + "DASCQualityMeasurement", + "analyze_gdn_decay", + "calibrate", + "compute_gdn_decay_horizons", + "export_policy", +] diff --git a/modelopt/torch/sparsity/state_sparsity/api.py b/modelopt/torch/sparsity/state_sparsity/api.py index 60d05620e46..1469e27cc63 100644 --- a/modelopt/torch/sparsity/state_sparsity/api.py +++ b/modelopt/torch/sparsity/state_sparsity/api.py @@ -21,10 +21,11 @@ from torch import nn -from modelopt.torch.opt.conversion import apply_mode +from modelopt.torch.opt.conversion import ModeloptStateManager, apply_mode +from modelopt.torch.utils import unwrap_model from .config import DASCCalibrationMeasurement, DASCConfig -from .conversion import get_attached_dasc_policy +from .conversion import get_attached_dasc_policy, replace_dasc_mode from .mode import DASCModeRegistry from .policy import validate_dasc_decay_parameters, validate_dasc_model_structure @@ -40,6 +41,7 @@ def calibrate( ``measurements`` must contain exactly one entry for every configured ``Wmax`` candidate. The largest candidate passing every quality, lifecycle, and storage gate is selected. + Recalibrating replaces the existing DASC mode-state entry in place. Example:: @@ -58,10 +60,16 @@ def calibrate( Returns: The input model with a serializable DASC policy attached through ModelOpt state. """ - config_dict = config.model_dump() if isinstance(config, DASCConfig) else config + model = unwrap_model(model, force_unwrap=True) + config_object = config if isinstance(config, DASCConfig) else DASCConfig(**config) + if ModeloptStateManager.is_converted(model, is_root=True) and any( + mode == "dasc" for mode, _ in ModeloptStateManager(model).state_dict() + ): + return replace_dasc_mode(model, config_object, measurements) + return apply_mode( model, - mode=[("dasc", config_dict)], + mode=[("dasc", config_object.model_dump())], registry=DASCModeRegistry, mode_kwargs={"measurements": measurements}, ) diff --git a/modelopt/torch/sparsity/state_sparsity/config.py b/modelopt/torch/sparsity/state_sparsity/config.py index c7bca8be9f3..7f67d7f7d47 100644 --- a/modelopt/torch/sparsity/state_sparsity/config.py +++ b/modelopt/torch/sparsity/state_sparsity/config.py @@ -18,13 +18,14 @@ import math from typing import Literal -from pydantic import Field, field_validator, model_validator +from pydantic import ConfigDict, Field, field_validator, model_validator from modelopt.torch.opt.config import ModeloptBaseConfig, ModeloptField __all__ = [ "DASCCalibrationMeasurement", "DASCConfig", + "DASCLayerPolicy", "DASCPolicy", "DASCQualityMeasurement", ] @@ -34,7 +35,11 @@ class DASCQualityMeasurement(ModeloptBaseConfig): """Quality and lifecycle measurements for one calibration slice.""" slice_id: str = Field(min_length=1) - perplexity_retention: float = Field(gt=0.0, le=1.0, allow_inf_nan=False) + perplexity_retention: float = Field( + gt=0.0, + allow_inf_nan=False, + description="Dense perplexity divided by DASC perplexity; values above one are valid.", + ) top1_agreement: float = Field(ge=0.0, le=1.0, allow_inf_nan=False) finite_continuation_logits: bool = Field(strict=True) retained_state_exact: bool = Field(strict=True) @@ -67,6 +72,8 @@ def validate_unique_slices( class DASCConfig(ModeloptBaseConfig): """Configuration for GDN decay-aware state checkpoint sparsity.""" + model_config = ConfigDict(protected_namespaces=()) + variant: Literal["dasc_nr", "dasc_wr"] = ModeloptField( default="dasc_wr", description="Use zero recovery (DASC-NR) or suffix replay recovery (DASC-WR).", @@ -79,6 +86,10 @@ 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( + default="float32", + description="Expected checkpoint storage dtype for GDN A_log and dt_bias.", + ) wmax_candidates: list[int] = ModeloptField( default=[8, 16, 32, 64, 128, 256], description="Positive candidate windows evaluated during offline calibration.", @@ -135,9 +146,9 @@ def validate_wmax_candidates(cls, candidates: object) -> object: @field_validator("min_perplexity_retention") @classmethod def validate_perplexity_gate(cls, value: float) -> float: - """Require a finite retention gate in (0, 1].""" - if not math.isfinite(value) or not 0.0 < value <= 1.0: - raise ValueError("min_perplexity_retention must be finite and in (0, 1]") + """Require a finite positive retention gate.""" + if not math.isfinite(value) or value <= 0.0: + raise ValueError("min_perplexity_retention must be finite and positive") return value @field_validator("min_top1_agreement") @@ -188,11 +199,14 @@ def validate_partition(self) -> "DASCLayerPolicy": class DASCPolicy(ModeloptBaseConfig): """Standalone JSON-safe DASC deployment policy.""" + model_config = ConfigDict(protected_namespaces=()) + format_version: Literal[1] = 1 variant: Literal["dasc_nr", "dasc_wr"] recovery: Literal["zero", "suffix_replay"] epsilon: float static_gate_input: float + decay_parameter_storage_dtype: Literal["float16", "bfloat16", "float32"] = "float32" selected_wmax: int = Field(strict=True, gt=0) wmax_candidates: list[int] = Field(min_length=1) quality_gates: dict[str, float] @@ -204,7 +218,10 @@ class DASCPolicy(ModeloptBaseConfig): preserve_convolution_state: Literal[True] active_runtime_state: Literal["dense"] = "dense" model_structure_sha256: str = Field(pattern=r"^[0-9a-f]{64}$") - decay_parameters_sha256: str = Field(pattern=r"^[0-9a-f]{64}$") + decay_parameters_sha256: str = Field( + pattern=r"^[0-9a-f]{64}$", + description="Storage-dtype-canonicalized calibration snapshot retained for provenance.", + ) layers: dict[str, DASCLayerPolicy] = Field(min_length=1) measurements: list[DASCCalibrationMeasurement] = Field(min_length=1) diff --git a/modelopt/torch/sparsity/state_sparsity/conversion.py b/modelopt/torch/sparsity/state_sparsity/conversion.py index 5e756a2587d..f826f09b9df 100644 --- a/modelopt/torch/sparsity/state_sparsity/conversion.py +++ b/modelopt/torch/sparsity/state_sparsity/conversion.py @@ -16,13 +16,15 @@ """ModelOpt conversion and restoration for DASC policy metadata.""" import copy +import warnings from collections.abc import Iterable from pydantic import ValidationError from torch import nn -from modelopt.torch.opt.conversion import ApplyModeError +from modelopt.torch.opt.conversion import ApplyModeError, ModeloptStateManager from modelopt.torch.opt.mode import ConvertReturnType, MetadataDict +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 @@ -33,6 +35,7 @@ def _attach_policy(model: nn.Module, policy: DASCPolicy) -> None: + """Attach a JSON-safe DASC policy to an unwrapped model.""" setattr(model, _DASC_POLICY_ATTRIBUTE, policy.model_dump(mode="json")) @@ -40,9 +43,14 @@ def convert_dasc_model( model: nn.Module, config: DASCConfig, *, - measurements: Iterable[DASCCalibrationMeasurement | dict], + measurements: Iterable[DASCCalibrationMeasurement | dict] | None = None, ) -> ConvertReturnType: """Analyze GDN decay and attach the selected DASC policy without changing execution.""" + if measurements is None: + raise ApplyModeError( + "DASC requires calibration measurements; use " + "modelopt.torch.sparsity.state_sparsity.calibrate(model, config, measurements)" + ) policy = build_dasc_policy(model, config, measurements) _attach_policy(model, policy) return model, {"policy": policy.model_dump(mode="json")} @@ -61,6 +69,7 @@ def restore_dasc_model(model: nn.Module, config: DASCConfig, metadata: MetadataD "variant": config.variant, "epsilon": config.epsilon, "static_gate_input": config.static_gate_input, + "decay_parameter_storage_dtype": config.decay_parameter_storage_dtype, "wmax_candidates": config.wmax_candidates, "quality_gates": { "min_perplexity_retention": config.min_perplexity_retention, @@ -82,25 +91,65 @@ def restore_dasc_model(model: nn.Module, config: DASCConfig, metadata: MetadataD if mismatched: raise ApplyModeError(f"DASC policy metadata does not match its mode config: {mismatched}") - validate_dasc_model_structure(model, policy) + try: + validate_dasc_model_structure(model, policy) + except ApplyModeError as error: + warnings.warn( + f"{error}. The restored DASC policy is stale; re-run calibrate() before deployment", + stacklevel=2, + ) _attach_policy(model, policy) return model def update_dasc_metadata(model: nn.Module, config: DASCConfig, metadata: MetadataDict) -> None: - """Refresh serialized metadata from the immutable attached DASC policy.""" + """Refresh metadata without making unrelated ModelOpt save or compose paths unusable.""" + policy = get_attached_dasc_policy(model) try: - policy = DASCPolicy(**getattr(model, _DASC_POLICY_ATTRIBUTE)) - except (AttributeError, TypeError, ValidationError) as error: - raise ApplyModeError("Model has no valid attached DASC policy") from error - validate_dasc_model_structure(model, policy) - validate_dasc_decay_parameters(model, policy) + validate_dasc_model_structure(model, policy) + 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")) +def replace_dasc_mode( + model: nn.Module, + config: DASCConfig, + measurements: Iterable[DASCCalibrationMeasurement | dict], +) -> nn.Module: + """Replace existing DASC mode state in place with a newly derived policy.""" + model = unwrap_model(model, force_unwrap=True) + manager = ModeloptStateManager(model) + state = manager.state_dict() + dasc_indices = [index for index, (mode, _) in enumerate(state) if mode == "dasc"] + if not dasc_indices: + raise ApplyModeError("Cannot replace DASC mode because the model has no DASC state") + policy = build_dasc_policy(model, config, measurements) + if dasc_indices[-1] != len(state) - 1: + manager.update_last_state_before_new_mode(model) + + first_index = dasc_indices[0] + state[first_index] = ( + "dasc", + { + "config": config.model_dump(), + "metadata": {"policy": policy.model_dump(mode="json")}, + }, + ) + for index in reversed(dasc_indices[1:]): + del state[index] + _attach_policy(model, policy) + return model + + def get_attached_dasc_policy(model: nn.Module) -> DASCPolicy: """Return the validated policy attached by conversion or restoration.""" + model = unwrap_model(model, force_unwrap=True) try: return DASCPolicy(**getattr(model, _DASC_POLICY_ATTRIBUTE)) except (AttributeError, TypeError, ValidationError) as error: diff --git a/modelopt/torch/sparsity/state_sparsity/mode.py b/modelopt/torch/sparsity/state_sparsity/mode.py index 23b3a17ffc3..96379f7f7db 100644 --- a/modelopt/torch/sparsity/state_sparsity/mode.py +++ b/modelopt/torch/sparsity/state_sparsity/mode.py @@ -50,7 +50,7 @@ def config_class(self) -> type[ModeloptBaseConfig]: @property def next_prohibited_modes(self) -> set[str]: - """Prevent applying DASC twice to the same model state.""" + """Route repeat calibration through the replacing public API.""" return {"dasc"} @property diff --git a/modelopt/torch/sparsity/state_sparsity/policy.py b/modelopt/torch/sparsity/state_sparsity/policy.py index 7342a66bf17..098dbcb4a59 100644 --- a/modelopt/torch/sparsity/state_sparsity/policy.py +++ b/modelopt/torch/sparsity/state_sparsity/policy.py @@ -16,19 +16,65 @@ """Decay analysis and policy selection for GDN state sparsity.""" import hashlib +import importlib import json +import math +import warnings from collections.abc import Iterable +from functools import lru_cache import torch import torch.nn.functional as F from torch import nn from modelopt.torch.opt.conversion import ApplyModeError +from modelopt.torch.opt.dynamic import DynamicModule +from modelopt.torch.utils import unwrap_model from .config import DASCCalibrationMeasurement, DASCConfig, DASCLayerPolicy, DASCPolicy __all__ = ["analyze_gdn_decay", "compute_gdn_decay_horizons"] +_SUPPORTED_GDN_CLASS_PATHS = ( + ("megatron.core.ssm.gated_delta_net", "GatedDeltaNet"), + ("transformers.models.qwen3_next.modeling_qwen3_next", "Qwen3NextGatedDeltaNet"), +) +_STORAGE_DTYPES = { + "float16": torch.float16, + "bfloat16": torch.bfloat16, + "float32": torch.float32, +} + + +@lru_cache(maxsize=1) +def _supported_gdn_classes() -> tuple[type[nn.Module], ...]: + """Resolve installed GDN implementations without making either framework mandatory.""" + classes = [] + for module_name, class_name in _SUPPORTED_GDN_CLASS_PATHS: + try: + candidate = getattr(importlib.import_module(module_name), class_name) + except ModuleNotFoundError as error: + root_module = module_name.partition(".")[0] + if importlib.util.find_spec(root_module) is None: + continue + warnings.warn( + f"DASC could not resolve {module_name}.{class_name}: {error!r}", stacklevel=2 + ) + continue + except Exception as error: + warnings.warn( + f"DASC could not resolve {module_name}.{class_name}: {error!r}", stacklevel=2 + ) + continue + if isinstance(candidate, type) and issubclass(candidate, nn.Module): + classes.append(candidate) + else: + warnings.warn( + f"DASC resolved {module_name}.{class_name}, but it is not an nn.Module class", + stacklevel=2, + ) + return tuple(classes) + def compute_gdn_decay_horizons( a_log: torch.Tensor, @@ -42,6 +88,8 @@ def compute_gdn_decay_horizons( 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)") @@ -57,33 +105,80 @@ def compute_gdn_decay_horizons( return horizons -def _is_gdn_module(module: nn.Module) -> bool: - class_name = "".join( - character for character in type(module).__name__.lower() if character.isalnum() - ) - return ( - "gateddeltanet" in class_name - and isinstance(getattr(module, "A_log", None), torch.Tensor) - and isinstance(getattr(module, "dt_bias", None), torch.Tensor) +def _has_supported_gdn_identity( + module: nn.Module, supported_classes: tuple[type[nn.Module], ...] +) -> bool: + """Return whether a module has an exact or ModelOpt-generated supported GDN identity.""" + module_class = type(module) + return module_class in supported_classes or ( + isinstance(module, DynamicModule) + and any(base in supported_classes for base in module_class.__mro__) ) +def _reject_incomplete_gdn_modules(identity_modules: list[tuple[str, nn.Module]]) -> None: + """Reject supported identities that do not expose both required decay tensors.""" + missing_decay_parameters = [ + name or "" + for name, module in identity_modules + if not all( + isinstance(getattr(module, parameter, None), torch.Tensor) + for parameter in ("A_log", "dt_bias") + ) + ] + if missing_decay_parameters: + raise ApplyModeError( + "DASC found supported GDN modules without A_log and dt_bias tensors at: " + f"{', '.join(missing_decay_parameters)}" + ) + + +def _reject_unconverted_gdn_subclasses( + named_modules: list[tuple[str, nn.Module]], + supported_classes: tuple[type[nn.Module], ...], +) -> None: + """Reject ordinary subclasses that would otherwise be silently omitted from the policy.""" + unsupported_subclasses = [ + name or "" + for name, module in named_modules + if not isinstance(module, DynamicModule) + and type(module) not in supported_classes + and any(base in supported_classes for base in type(module).__mro__[1:]) + ] + if unsupported_subclasses: + raise ApplyModeError( + "DASC found GDN subclasses that are not ModelOpt dynamic modules at: " + f"{', '.join(unsupported_subclasses)}; convert the module with ModelOpt or use a " + "supported class directly" + ) + + def _get_gdn_modules(model: nn.Module) -> dict[str, nn.Module]: - modules = {name: module for name, module in model.named_modules() if _is_gdn_module(module)} - if not modules: - raise ApplyModeError("DASC found no GatedDeltaNet modules; only GDN is supported") - return dict(sorted(modules.items())) + """Find supported GDN layers after removing a recognized model wrapper.""" + model = unwrap_model(model, force_unwrap=True) + supported_classes = _supported_gdn_classes() + named_modules = list(model.named_modules()) + identity_modules = [ + (name, module) + for name, module in named_modules + if _has_supported_gdn_identity(module, supported_classes) + ] + _reject_incomplete_gdn_modules(identity_modules) + _reject_unconverted_gdn_subclasses(named_modules, supported_classes) + if not identity_modules: + supported = ", ".join( + f"{module_name}.{class_name}" for module_name, class_name in _SUPPORTED_GDN_CLASS_PATHS + ) + raise ApplyModeError(f"DASC found no supported GDN modules; expected one of: {supported}") + return dict(sorted(identity_modules, key=lambda item: item[0])) -def analyze_gdn_decay( - model: nn.Module, - *, - epsilon: float = 1e-3, - static_gate_input: float = -0.3, +def _analyze_gdn_modules( + modules: dict[str, nn.Module], *, epsilon: float, static_gate_input: float ) -> dict[str, list[float]]: - """Return deterministic per-head horizons for every GDN module in a model.""" + """Compute deterministic per-head horizons for already-resolved GDN modules.""" horizons = {} - for name, module in _get_gdn_modules(model).items(): + for name, module in modules.items(): try: layer_horizons = compute_gdn_decay_horizons( module.A_log, @@ -99,12 +194,26 @@ def analyze_gdn_decay( return horizons +def analyze_gdn_decay( + model: nn.Module, + *, + epsilon: float = 1e-3, + static_gate_input: float = -0.3, +) -> dict[str, list[float]]: + """Return deterministic per-head horizons for every GDN module in a model.""" + return _analyze_gdn_modules( + _get_gdn_modules(model), epsilon=epsilon, static_gate_input=static_gate_input + ) + + def _canonical_sha256(value: object) -> str: + """Hash a JSON value with deterministic ordering and no non-finite numbers.""" payload = json.dumps(value, sort_keys=True, separators=(",", ":"), allow_nan=False) return hashlib.sha256(payload.encode()).hexdigest() def _model_structure(modules: dict[str, nn.Module]) -> list[dict[str, object]]: + """Describe the layer names and head counts that define policy geometry.""" return [ { "name": name, @@ -114,20 +223,65 @@ def _model_structure(modules: dict[str, nn.Module]) -> list[dict[str, object]]: ] -def _decay_parameters(modules: dict[str, nn.Module]) -> list[dict[str, object]]: +def _decay_parameters( + modules: dict[str, nn.Module], storage_dtype: torch.dtype +) -> list[dict[str, object]]: + """Serialize a storage-dtype-canonicalized calibration snapshot for provenance.""" return [ { "name": name, - "A_log": module.A_log.detach().to(device="cpu", dtype=torch.float64).tolist(), - "dt_bias": module.dt_bias.detach().to(device="cpu", dtype=torch.float64).tolist(), + "A_log": module.A_log.detach() + .to(device="cpu", dtype=storage_dtype) + .to(dtype=torch.float32) + .tolist(), + "dt_bias": module.dt_bias.detach() + .to(device="cpu", dtype=storage_dtype) + .to(dtype=torch.float32) + .tolist(), } for name, module in modules.items() ] +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))) + for dtype in reversed(cast_dtypes): + dtype_info = torch.finfo(dtype) + unit_roundoff = dtype_info.eps / 2.0 + smallest_subnormal = dtype_info.tiny * dtype_info.eps + upper = (upper + smallest_subnormal) / (1.0 - unit_roundoff) + return upper - values + + +def _storage_cast_horizon_bounds( + module: nn.Module, + *, + epsilon: float, + static_gate_input: float, + storage_dtype: torch.dtype, +) -> tuple[torch.Tensor, torch.Tensor]: + """Bound horizons compatible with the current parameters before one storage cast.""" + 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) + dt_radius = _storage_rounding_radius(module.dt_bias, storage_dtype) + scale = -math.log(epsilon) + lower = scale / ( + torch.exp(a_log + a_radius) * F.softplus(dt_bias + dt_radius + static_gate_input) + ) + upper = scale / ( + torch.exp(a_log - a_radius) * F.softplus(dt_bias - dt_radius + static_gate_input) + ) + return lower, upper + + def _validate_measurement_coverage( config: DASCConfig, measurements: list[DASCCalibrationMeasurement] ) -> None: + """Require exactly one matching measurement for every configured window.""" measured = [measurement.wmax for measurement in measurements] unexpected = sorted(set(measured) - set(config.wmax_candidates)) missing = sorted(set(config.wmax_candidates) - set(measured)) @@ -145,6 +299,7 @@ def _validate_measurement_coverage( def _validate_measurement_geometry( horizons: dict[str, list[float]], measurements: list[DASCCalibrationMeasurement] ) -> None: + """Bind caller-reported retained and total head counts to the analyzed model.""" total_heads = sum(len(layer_horizons) for layer_horizons in horizons.values()) for measurement in measurements: retained_heads = sum( @@ -161,6 +316,7 @@ def _validate_measurement_geometry( def _candidate_passes(config: DASCConfig, measurement: DASCCalibrationMeasurement) -> bool: + """Return whether one candidate passes every configured evidence gate.""" return measurement.checkpoint_savings >= config.min_checkpoint_savings and all( result.perplexity_retention >= config.min_perplexity_retention and result.top1_agreement >= config.min_top1_agreement @@ -191,8 +347,8 @@ def build_dasc_policy( validated_measurements.sort(key=lambda measurement: measurement.wmax) modules = _get_gdn_modules(model) - horizons = analyze_gdn_decay( - model, epsilon=config.epsilon, static_gate_input=config.static_gate_input + horizons = _analyze_gdn_modules( + modules, epsilon=config.epsilon, static_gate_input=config.static_gate_input ) _validate_measurement_geometry(horizons, validated_measurements) @@ -222,6 +378,7 @@ def build_dasc_policy( recovery="zero" if config.variant == "dasc_nr" else "suffix_replay", epsilon=config.epsilon, static_gate_input=config.static_gate_input, + decay_parameter_storage_dtype=config.decay_parameter_storage_dtype, selected_wmax=selected_wmax, wmax_candidates=config.wmax_candidates, quality_gates={ @@ -236,7 +393,9 @@ def build_dasc_policy( granularity=config.granularity, preserve_convolution_state=config.preserve_convolution_state, model_structure_sha256=_canonical_sha256(_model_structure(modules)), - decay_parameters_sha256=_canonical_sha256(_decay_parameters(modules)), + decay_parameters_sha256=_canonical_sha256( + _decay_parameters(modules, _STORAGE_DTYPES[config.decay_parameter_storage_dtype]) + ), layers=layers, measurements=validated_measurements, ) @@ -245,14 +404,48 @@ def build_dasc_policy( def validate_dasc_model_structure(model: nn.Module, policy: DASCPolicy) -> None: """Reject restoring a policy onto a different GDN module structure.""" modules = _get_gdn_modules(model) - actual = _canonical_sha256(_model_structure(modules)) - if actual != policy.model_structure_sha256: + actual_structure = _model_structure(modules) + policy_structure = [ + {"name": name, "num_heads": layer.num_heads} + for name, layer in sorted(policy.layers.items()) + ] + if ( + 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") def validate_dasc_decay_parameters(model: nn.Module, policy: DASCPolicy) -> None: - """Reject exporting a policy for different GDN decay parameters.""" + """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. + """ modules = _get_gdn_modules(model) - actual = _canonical_sha256(_decay_parameters(modules)) - if actual != policy.decay_parameters_sha256: - raise ApplyModeError("DASC policy does not match the model's GDN decay parameters") + current_horizons = _analyze_gdn_modules( + modules, + epsilon=policy.epsilon, + static_gate_input=policy.static_gate_input, + ) + for name, values in current_horizons.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}" + ) + lower, upper = _storage_cast_horizon_bounds( + modules[name], + epsilon=policy.epsilon, + static_gate_input=policy.static_gate_input, + storage_dtype=_STORAGE_DTYPES[policy.decay_parameter_storage_dtype], + ) + 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( + stored > upper * (1.0 + numerical_slack) + ): + raise ApplyModeError( + f"DASC policy horizons do not match current decay parameters in layer {name!r}" + ) diff --git a/tests/unit/torch/sparsity/state_sparsity/test_dasc.py b/tests/unit/torch/sparsity/state_sparsity/test_dasc.py index c212808c046..06379e856bf 100644 --- a/tests/unit/torch/sparsity/state_sparsity/test_dasc.py +++ b/tests/unit/torch/sparsity/state_sparsity/test_dasc.py @@ -16,6 +16,7 @@ """CPU tests for DASC state-sparsity policy calibration.""" import copy +import io import json import pytest @@ -25,16 +26,22 @@ import modelopt.torch.opt as mto import modelopt.torch.sparsity.state_sparsity as mtss -from modelopt.torch.opt.conversion import ApplyModeError +import modelopt.torch.sparsity.state_sparsity.policy as dasc_policy +from modelopt.torch.opt.conversion import ApplyModeError, ModeloptStateManager +from modelopt.torch.opt.dynamic import DynamicModule +from modelopt.torch.sparsity.state_sparsity.conversion import replace_dasc_mode +from modelopt.torch.sparsity.state_sparsity.mode import DASCModeRegistry +_resolve_supported_gdn_classes = dasc_policy._supported_gdn_classes.__wrapped__ -class TinyGatedDeltaNet(nn.Module): + +class GatedDeltaNet(nn.Module): """Minimal GDN-shaped module for framework-independent tests.""" def __init__(self, num_heads: int = 2): super().__init__() - a_log = torch.zeros(num_heads) - dt_bias = torch.tensor([-2.0, 2.0]) if num_heads == 2 else torch.zeros(num_heads) + a_log = torch.tensor([0.1, 0.7]) if num_heads == 2 else torch.zeros(num_heads) + dt_bias = torch.tensor([-2.3, 1.7]) if num_heads == 2 else torch.zeros(num_heads) self.A_log = nn.Parameter(a_log) self.dt_bias = nn.Parameter(dt_bias) @@ -47,13 +54,20 @@ class TinyGatedDeltaNetForCausalLM(nn.Module): def __init__(self, num_heads: int = 2): super().__init__() - self.linear_attn = TinyGatedDeltaNet(num_heads) + self.linear_attn = GatedDeltaNet(num_heads) def forward(self, inputs): return self.linear_attn(inputs) +@pytest.fixture(autouse=True) +def _register_test_gdn_class(monkeypatch): + """Use the exact toy GDN identity without weakening production class checks.""" + monkeypatch.setattr(dasc_policy, "_supported_gdn_classes", lambda: (GatedDeltaNet,)) + + def _config(**overrides): + """Return a complete test configuration with selected overrides.""" config = { "variant": "dasc_wr", "epsilon": 1e-3, @@ -72,6 +86,7 @@ def _config(**overrides): def _candidate(wmax, *, variant="dasc_wr", top1=0.99, convolution_state_exact=True): + """Return passing caller-supplied evidence for one recovery window.""" return { "variant": variant, "wmax": wmax, @@ -93,6 +108,7 @@ def _candidate(wmax, *, variant="dasc_wr", top1=0.99, convolution_state_exact=Tr def test_calibrate_selects_largest_passing_candidate_and_round_trips(): + """Select the largest passing window and preserve weights and ModelOpt state.""" model = TinyGatedDeltaNetForCausalLM() original_state = {name: value.clone() for name, value in model.state_dict().items()} @@ -105,6 +121,7 @@ def test_calibrate_selects_largest_passing_candidate_and_round_trips(): assert policy["granularity"] == "gdn_head" assert policy["preserve_convolution_state"] is True assert policy["active_runtime_state"] == "dense" + assert policy["decay_parameter_storage_dtype"] == "float32" assert policy["layers"]["linear_attn"]["retained_heads"] == [0] assert policy["layers"]["linear_attn"]["omitted_heads"] == [1] assert [measurement["wmax"] for measurement in policy["measurements"]] == [7, 11] @@ -113,11 +130,18 @@ def test_calibrate_selects_largest_passing_candidate_and_round_trips(): ) json.dumps(policy) - restored = mto.restore_from_modelopt_state( - TinyGatedDeltaNetForCausalLM(), mto.modelopt_state(model) - ) + state = mto.modelopt_state(model) + restored = mto.restore_from_modelopt_state(TinyGatedDeltaNetForCausalLM(), state) assert mtss.export_policy(restored) == policy + legacy_state = copy.deepcopy(state) + del legacy_state["modelopt_state_dict"][0][1]["config"]["decay_parameter_storage_dtype"] + del legacy_state["modelopt_state_dict"][0][1]["metadata"]["policy"][ + "decay_parameter_storage_dtype" + ] + legacy_restored = mto.restore_from_modelopt_state(TinyGatedDeltaNetForCausalLM(), legacy_state) + assert mtss.export_policy(legacy_restored)["decay_parameter_storage_dtype"] == "float32" + policy["selected_wmax"] = 999 assert mtss.export_policy(model)["selected_wmax"] == 7 @@ -126,6 +150,7 @@ def test_calibrate_selects_largest_passing_candidate_and_round_trips(): ("variant", "recovery"), [("dasc_nr", "zero"), ("dasc_wr", "suffix_replay")] ) def test_variant_is_explicit_in_exported_policy(variant, recovery): + """Keep zero and suffix-replay recovery as explicit deployment contracts.""" model = mtss.calibrate( TinyGatedDeltaNetForCausalLM(), _config(variant=variant, wmax_candidates=[7]), @@ -144,16 +169,35 @@ def test_variant_is_explicit_in_exported_policy(variant, recovery): {"wmax_candidates": [7, 7]}, {"model_revision": ""}, {"preserve_convolution_state": False}, + {"decay_parameter_storage_dtype": "float8"}, ], ) def test_config_fails_closed(override): + """Reject incomplete provenance and invalid window or lifecycle settings.""" with pytest.raises(ValidationError): mtss.DASCConfig(**_config(**override)) assert mtss.DASCConfig(**_config(wmax_candidates=[7])).wmax_candidates == [7] +def test_perplexity_retention_accepts_parity_improvements(): + """Allow improvement measurements and thresholds instead of requiring clamping.""" + candidate = _candidate(7) + candidate["quality"][0]["perplexity_retention"] = 1.0004 + measurement = mtss.DASCCalibrationMeasurement(**candidate) + + assert measurement.quality[0].perplexity_retention == 1.0004 + + model = mtss.calibrate( + TinyGatedDeltaNetForCausalLM(), + _config(wmax_candidates=[7], min_perplexity_retention=1.0002), + [candidate], + ) + assert mtss.export_policy(model)["selected_wmax"] == 7 + + def test_calibration_fails_closed_on_measurements_and_model_mismatch(): + """Reject incomplete evidence, failing gates, wrong geometry, and unsupported layers.""" with pytest.raises(ApplyModeError, match="exactly one result"): mtss.calibrate(TinyGatedDeltaNetForCausalLM(), _config(), [_candidate(7)]) @@ -164,6 +208,22 @@ def test_calibration_fails_closed_on_measurements_and_model_mismatch(): [_candidate(7, convolution_state_exact=False)], ) + with pytest.raises(ApplyModeError, match="configured variant"): + mtss.calibrate( + TinyGatedDeltaNetForCausalLM(), + _config(wmax_candidates=[7]), + [_candidate(7, variant="dasc_nr")], + ) + + invalid_measurement = _candidate(7) + del invalid_measurement["quality"] + with pytest.raises(ApplyModeError, match="Invalid DASC calibration measurements"): + mtss.calibrate( + TinyGatedDeltaNetForCausalLM(), + _config(wmax_candidates=[7]), + [invalid_measurement], + ) + mismatched_geometry = _candidate(7) mismatched_geometry["retained_heads"] = 2 with pytest.raises(ApplyModeError, match="geometry"): @@ -176,11 +236,238 @@ def test_calibration_fails_closed_on_measurements_and_model_mismatch(): class NotGDN(nn.Module): pass - with pytest.raises(ApplyModeError, match="no GatedDeltaNet modules"): + with pytest.raises(ApplyModeError, match="no supported GDN modules"): mtss.calibrate(NotGDN(), _config(wmax_candidates=[7]), [_candidate(7)]) + class UnsupportedGatedDeltaNet(nn.Module): + """Unrelated lookalike must not satisfy the supported-base-class contract.""" + + def __init__(self): + super().__init__() + self.A_log = nn.Parameter(torch.zeros(2)) + self.dt_bias = nn.Parameter(torch.zeros(2)) + + with pytest.raises(ApplyModeError, match="no supported GDN modules"): + mtss.calibrate(UnsupportedGatedDeltaNet(), _config(wmax_candidates=[7]), [_candidate(7)]) + + class UnsupportedSubclass(GatedDeltaNet): + pass + + with pytest.raises(ApplyModeError, match="subclasses that are not ModelOpt dynamic modules"): + mtss.calibrate(UnsupportedSubclass(), _config(wmax_candidates=[7]), [_candidate(7)]) + + same_name_lookalike = type("GatedDeltaNet", (UnsupportedGatedDeltaNet,), {}) + with pytest.raises(ApplyModeError, match="no supported GDN modules"): + mtss.calibrate(same_name_lookalike(), _config(wmax_candidates=[7]), [_candidate(7)]) + + missing_decay = TinyGatedDeltaNetForCausalLM() + del missing_decay.linear_attn.A_log + with pytest.raises(ApplyModeError, match="without A_log and dt_bias tensors"): + mtss.calibrate(missing_decay, _config(wmax_candidates=[7]), [_candidate(7)]) + + partially_valid = nn.Module() + partially_valid.good = GatedDeltaNet() + partially_valid.bad = GatedDeltaNet() + del partially_valid.bad.dt_bias + with pytest.raises(ApplyModeError, match=r"without A_log and dt_bias tensors at: bad$"): + mtss.analyze_gdn_decay(partially_valid) + + mixed_subclass = nn.Module() + mixed_subclass.good = GatedDeltaNet() + mixed_subclass.stale = UnsupportedSubclass() + with pytest.raises( + ApplyModeError, + match=( + r"not ModelOpt dynamic modules at: stale; convert the module with ModelOpt or use a " + r"supported class directly$" + ), + ): + mtss.analyze_gdn_decay(mixed_subclass) + + invalid_decay = TinyGatedDeltaNetForCausalLM() + invalid_decay.linear_attn.dt_bias = nn.Parameter(torch.zeros(3)) + with pytest.raises(ApplyModeError, match="Invalid GDN decay parameters"): + mtss.analyze_gdn_decay(invalid_decay) + + +def test_supported_class_resolution_uses_imported_module_identities(monkeypatch): + """Ignore absent and non-module symbols while retaining exact supported identities.""" + module_paths = ( + ("valid", "GatedDeltaNet"), + ("invalid", "NotAModule"), + ("missing", "Missing"), + ("installed", "Missing"), + ("broken", "Broken"), + ) + modules = { + "valid": type("ValidModule", (), {"GatedDeltaNet": GatedDeltaNet}), + "invalid": type("InvalidModule", (), {"NotAModule": object()}), + } + + def import_module(name): + if name in {"missing", "installed"}: + raise ModuleNotFoundError(name) + if name == "broken": + raise RuntimeError(name) + return modules[name] + + monkeypatch.setattr(dasc_policy, "_SUPPORTED_GDN_CLASS_PATHS", module_paths) + monkeypatch.setattr(dasc_policy.importlib, "import_module", import_module) + monkeypatch.setattr( + dasc_policy.importlib.util, + "find_spec", + lambda name: object() if name == "installed" else None, + ) + + with pytest.warns(UserWarning) as caught: + assert _resolve_supported_gdn_classes() == (GatedDeltaNet,) + assert len(caught) == 3 + + +@pytest.mark.parametrize(("module_name", "class_name"), dasc_policy._SUPPORTED_GDN_CLASS_PATHS) +def test_declared_gdn_paths_resolve_when_framework_is_installed(module_name, class_name): + """Guard supported identities against upstream dependency path drift.""" + root_module = module_name.partition(".")[0] + pytest.importorskip(root_module) + module = dasc_policy.importlib.import_module(module_name) + assert issubclass(getattr(module, class_name), nn.Module) + + +def test_generic_mode_application_reports_missing_measurements(monkeypatch): + """Give generic apply_mode callers an actionable calibration-evidence error.""" + assert DASCModeRegistry["dasc"].next_prohibited_modes == {"dasc"} + assert DASCModeRegistry["dasc"].update_for_new_mode is not None + with pytest.raises(ApplyModeError, match="requires calibration measurements"): + mto.apply_mode( + TinyGatedDeltaNetForCausalLM(), + mode=[("dasc", _config(wmax_candidates=[7]))], + ) + + model = TinyGatedDeltaNetForCausalLM() + ModeloptStateManager(model, init_state=True) + with pytest.raises(ApplyModeError, match="model has no DASC state"): + replace_dasc_mode( + model, + mtss.DASCConfig(**_config(wmax_candidates=[7])), + [_candidate(7)], + ) + + model = mtss.calibrate( + TinyGatedDeltaNetForCausalLM(), _config(wmax_candidates=[7]), [_candidate(7)] + ) + ModeloptStateManager(model).state_dict().append(("trailing-mode", {})) + refreshed = [] + monkeypatch.setattr( + ModeloptStateManager, + "update_last_state_before_new_mode", + lambda _manager, current_model: refreshed.append(current_model), + ) + replace_dasc_mode( + model, + mtss.DASCConfig(**_config(wmax_candidates=[7])), + [_candidate(7)], + ) + assert refreshed == [model] + + +def test_public_exports_and_wrapped_model_export(): + """Expose only supported symbols and accept wrappers and ModelOpt subclasses.""" + model = mtss.calibrate( + TinyGatedDeltaNetForCausalLM(), _config(wmax_candidates=[7]), [_candidate(7)] + ) + + assert "DASCLayerPolicy" in mtss.__all__ + assert "mode" not in mtss.__all__ + assert mtss.export_policy(nn.DataParallel(model)) == mtss.export_policy(model) + + dynamic_class = type("_DynamicGatedDeltaNet", (DynamicModule, GatedDeltaNet), {}) + dynamic_module = GatedDeltaNet() + dynamic_module.__class__ = dynamic_class + model.linear_attn = dynamic_module + assert mtss.export_policy(model)["layers"]["linear_attn"]["num_heads"] == 2 + + with pytest.raises(ApplyModeError, match="no valid attached DASC policy"): + mtss.export_policy(TinyGatedDeltaNetForCausalLM()) + + +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) +def test_dtype_cast_preserves_policy_when_the_selected_mask_is_unchanged(dtype): + """Treat ordinary low-precision casts as equivalent when they preserve the policy mask.""" + storage_dtype = "bfloat16" if dtype == torch.bfloat16 else "float16" + model = mtss.calibrate( + TinyGatedDeltaNetForCausalLM(), + _config(wmax_candidates=[7], decay_parameter_storage_dtype=storage_dtype), + [_candidate(7)], + ) + policy = mtss.export_policy(model) + original_decay = torch.cat( + [model.linear_attn.A_log.detach(), model.linear_attn.dt_bias.detach()] + ) + + model.to(dtype) + + cast_decay = torch.cat( + [model.linear_attn.A_log.detach().float(), model.linear_attn.dt_bias.detach().float()] + ) + assert not torch.equal(original_decay, cast_decay) + assert mtss.export_policy(model) == policy + + +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( + TinyGatedDeltaNetForCausalLM(), + _config(wmax_candidates=[7], decay_parameter_storage_dtype="bfloat16"), + [_candidate(7)], + ) + policy = mtss.export_policy(model) + original_decay = torch.cat( + [model.linear_attn.A_log.detach(), model.linear_attn.dt_bias.detach()] + ) + + with torch.no_grad(): + model.linear_attn.A_log.copy_(model.linear_attn.A_log.to(torch.bfloat16).float()) + model.linear_attn.dt_bias.copy_(model.linear_attn.dt_bias.to(torch.bfloat16).float()) + + reloaded_decay = torch.cat( + [model.linear_attn.A_log.detach(), model.linear_attn.dt_bias.detach()] + ) + assert reloaded_decay.dtype == torch.float32 + assert not torch.equal(original_decay, reloaded_decay) + assert mtss.export_policy(model) == policy + + model.linear_attn.A_log = nn.Parameter( + model.linear_attn.A_log.detach().to(torch.int64), requires_grad=False + ) + with pytest.raises(ApplyModeError, match="floating-point dtype"): + mtss.export_policy(model) + + +@pytest.mark.parametrize( + ("storage_name", "storage_dtype", "live_dtype"), + [ + ("float16", torch.float16, torch.bfloat16), + ("bfloat16", torch.bfloat16, torch.float16), + ], +) +def test_cross_dtype_reload_accumulates_both_rounding_bounds( + storage_name, storage_dtype, live_dtype +): + """Accept two distinct declared-storage and live-materialization rounding steps.""" + model = mtss.calibrate( + TinyGatedDeltaNetForCausalLM(), + _config(wmax_candidates=[7], decay_parameter_storage_dtype=storage_name), + [_candidate(7)], + ) + policy = mtss.export_policy(model) + + model.to(storage_dtype).to(live_dtype) + + assert mtss.export_policy(model) == policy + def test_export_rejects_changed_decay_parameters_and_restore_rejects_structure(): + """Keep saving recoverable while rejecting stale or tampered deployment policies.""" model = mtss.calibrate( TinyGatedDeltaNetForCausalLM(), _config(wmax_candidates=[7]), [_candidate(7)] ) @@ -188,13 +475,38 @@ def test_export_rejects_changed_decay_parameters_and_restore_rejects_structure() with torch.no_grad(): model.linear_attn.A_log.add_(1.0) - with pytest.raises(ApplyModeError, match="decay parameters"): + with pytest.raises(ApplyModeError, match="horizons"): mtss.export_policy(model) - with pytest.raises(ApplyModeError, match="decay parameters"): + with pytest.warns(UserWarning, match="saved DASC policy is stale"): mto.modelopt_state(model) + checkpoint = io.BytesIO() + with pytest.warns(UserWarning, match="saved DASC policy is stale"): + mto.save(model, checkpoint) + assert checkpoint.tell() > 0 + + manager_state = ModeloptStateManager(model).state_dict() + manager_state.append(copy.deepcopy(manager_state[0])) + delattr(model, "_modelopt_dasc_policy") + model = mtss.calibrate( + model, + _config(wmax_candidates=[7], model_revision="revision-2"), + [_candidate(7)], + ) + policy = mtss.export_policy(model) + assert policy["model_revision"] == "revision-2" + model_state = copy.deepcopy(model.state_dict()) + recalibrated_state = mto.modelopt_state(model) + assert [mode for mode, _ in recalibrated_state["modelopt_state_dict"]] == ["dasc"] + restored = mto.restore_from_modelopt_state(TinyGatedDeltaNetForCausalLM(), recalibrated_state) + restored.load_state_dict(model_state) + assert mtss.export_policy(restored) == policy + with pytest.warns(UserWarning, match="restored DASC policy is stale"): + mismatched = mto.restore_from_modelopt_state( + TinyGatedDeltaNetForCausalLM(num_heads=3), state + ) with pytest.raises(ApplyModeError, match="module structure"): - mto.restore_from_modelopt_state(TinyGatedDeltaNetForCausalLM(num_heads=3), state) + mtss.export_policy(mismatched) tampered_state = copy.deepcopy(state) tampered_state["modelopt_state_dict"][0][1]["metadata"]["policy"]["quality_gates"][ @@ -202,3 +514,65 @@ def test_export_rejects_changed_decay_parameters_and_restore_rejects_structure() ] = 0.5 with pytest.raises(ApplyModeError, match="does not match its mode config"): mto.restore_from_modelopt_state(TinyGatedDeltaNetForCausalLM(), tampered_state) + + tampered_state = copy.deepcopy(state) + tampered_state["modelopt_state_dict"][0][1]["metadata"]["unexpected"] = True + with pytest.raises(ApplyModeError, match="only the policy field"): + mto.restore_from_modelopt_state(TinyGatedDeltaNetForCausalLM(), tampered_state) + + tampered_state = copy.deepcopy(state) + del tampered_state["modelopt_state_dict"][0][1]["metadata"]["policy"]["variant"] + with pytest.raises(ApplyModeError, match="Invalid DASC policy metadata"): + mto.restore_from_modelopt_state(TinyGatedDeltaNetForCausalLM(), tampered_state) + + tampered_state = copy.deepcopy(state) + layer = tampered_state["modelopt_state_dict"][0][1]["metadata"]["policy"]["layers"][ + "linear_attn" + ] + layer["static_horizons"][0] *= 1.0001 + restored = mto.restore_from_modelopt_state(TinyGatedDeltaNetForCausalLM(), tampered_state) + with pytest.raises(ApplyModeError, match="horizons do not match"): + mtss.export_policy(restored) + + +def test_structure_staleness_does_not_block_checkpoint_save(): + """Keep checkpoint save and restore available after a calibrated GDN structure changes.""" + model = mtss.calibrate( + TinyGatedDeltaNetForCausalLM(), _config(wmax_candidates=[7]), [_candidate(7)] + ) + model.linear_attn = GatedDeltaNet(num_heads=3) + + checkpoint = io.BytesIO() + with pytest.warns(UserWarning, match="saved DASC policy is stale"): + mto.save(model, checkpoint) + + assert checkpoint.tell() > 0 + with pytest.raises(ApplyModeError, match="module structure"): + mtss.export_policy(model) + + with pytest.warns(UserWarning, match="saved DASC policy is stale"): + stale_state = mto.modelopt_state(model) + with pytest.warns(UserWarning, match="restored DASC policy is stale"): + restored = mto.restore_from_modelopt_state( + TinyGatedDeltaNetForCausalLM(num_heads=3), stale_state + ) + restored.load_state_dict(model.state_dict()) + with pytest.raises(ApplyModeError, match="module structure"): + 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( + TinyGatedDeltaNetForCausalLM(), _config(wmax_candidates=[54]), [_candidate(54)] + ) + state = mto.modelopt_state(model) + layer = state["modelopt_state_dict"][0][1]["metadata"]["policy"]["layers"]["linear_attn"] + layer["static_horizons"][0] = 53.0 + layer["retained_heads"] = [] + layer["omitted_heads"] = [0, 1] + + restored = mto.restore_from_modelopt_state(TinyGatedDeltaNetForCausalLM(), state) + + with pytest.raises(ApplyModeError, match="head mask does not match"): + mtss.export_policy(restored) From bd276a1d9a6140c5b6444861292a6b2df911cddc Mon Sep 17 00:00:00 2001 From: kaix-nv Date: Thu, 10 Sep 2026 21:51:37 -0700 Subject: [PATCH 2/2] Fix DASC storage-boundary validation (#2388) Addresses the fresh full-diff review findings on #2387. Changes: - derive calibration horizons, geometry, and masks from decay parameters canonicalized to the declared checkpoint storage dtype - validate boundary masks with inverse storage/live-dtype intervals while preserving the strict `horizon > Wmax` rule - fail closed when restoring DASC metadata onto a model with no supported GDN modules, while retaining stale-policy warnings for recoverable GDN geometry drift - import `importlib.util` explicitly and correct the storage-tolerance documentation - add FP16 boundary-flip and wrong-architecture restore regressions Validation: - `python -m pytest -q tests/unit/torch/sparsity/state_sparsity/test_dasc.py`: 24 passed, 1 skipped (optional Megatron dependency) - state + weight + attention sparsity tests: 302 passed, 1 skipped - pre-commit on all touched files: passed - real Transformers `Qwen3NextGatedDeltaNet` CPU smoke: BF16-declared calibration, BF16 reload, and policy export round-trip passed Commit is ED25519-signed and carries a matching Signed-off-by trailer. ## Summary by CodeRabbit * **Bug Fixes** * Improved sparsity policy validation to account for rounding in both stored and live tensor data types. * Policies now tolerate valid decay-related rounding differences while rejecting masks that are inconsistent with current parameters. * Improved handling and messaging for unsupported model structures during state restoration. * **Tests** * Added coverage for storage-precision threshold behavior and unsupported model structures. --------- Signed-off-by: kaix-nv --- docs/source/guides/6_sparsity.rst | 25 ++- .../torch/sparsity/state_sparsity/config.py | 48 +++++- .../sparsity/state_sparsity/conversion.py | 22 ++- .../torch/sparsity/state_sparsity/policy.py | 157 ++++++++++++----- .../sparsity/state_sparsity/test_dasc.py | 158 ++++++++++++++++++ 5 files changed, 354 insertions(+), 56 deletions(-) diff --git a/docs/source/guides/6_sparsity.rst b/docs/source/guides/6_sparsity.rst index 691b637f022..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,12 +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; the default ``float32`` keeps unconfigured policies strict. +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 f826f09b9df..213cc351f45 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 ( + _DASCRecoverableStalenessError, + 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 _DASCRecoverableStalenessError as error: warnings.warn( f"{error}. The restored DASC policy is stale; re-run calibrate() before deployment", stacklevel=2, @@ -103,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 098dbcb4a59..7baf53e9586 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 @@ -31,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"] @@ -39,13 +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 _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.""" @@ -84,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 @@ -127,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)}" ) @@ -174,15 +197,25 @@ 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 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( - module.A_log, - module.dt_bias, + a_log, + dt_bias, epsilon=epsilon, static_gate_input=static_gate_input, ) @@ -199,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, ) @@ -243,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 @@ -263,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) @@ -348,7 +412,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,35 +480,47 @@ 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: + try: + _validate_gdn_decay_tensors(module.A_log, module.dt_bias) + except ValueError as error: raise ApplyModeError( - f"DASC policy head mask does not match current decay parameters in layer {name!r}" - ) + f"Invalid GDN decay parameters in module {name!r}: {error}" + ) from error 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], ) - stored = torch.tensor(layer.static_horizons, dtype=torch.float64) + 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, 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 06379e856bf..cad44b1501d 100644 --- a/tests/unit/torch/sparsity/state_sparsity/test_dasc.py +++ b/tests/unit/torch/sparsity/state_sparsity/test_dasc.py @@ -413,6 +413,112 @@ 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.analyze_gdn_decay( + model, + static_gate_input=0.0, + decay_parameter_storage_dtype="float16", + )["linear_attn"][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 + + +@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( @@ -466,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( @@ -508,6 +637,16 @@ 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) + + 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" @@ -561,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(