From f4e9bd6162dde0464034fb1b989c7146b64f841e Mon Sep 17 00:00:00 2001 From: Kai Xu Date: Thu, 10 Sep 2026 16:07:29 -0700 Subject: [PATCH 1/2] Add DASC recurrent state sparsity policy Signed-off-by: Kai Xu --- CHANGELOG.rst | 1 + docs/source/guides/6_sparsity.rst | 57 ++++ modelopt/torch/sparsity/__init__.py | 4 +- .../torch/sparsity/state_sparsity/__init__.py | 21 ++ modelopt/torch/sparsity/state_sparsity/api.py | 80 ++++++ .../torch/sparsity/state_sparsity/config.py | 232 ++++++++++++++++ .../sparsity/state_sparsity/conversion.py | 107 ++++++++ .../torch/sparsity/state_sparsity/mode.py | 74 +++++ .../torch/sparsity/state_sparsity/policy.py | 258 ++++++++++++++++++ .../sparsity/state_sparsity/test_dasc.py | 204 ++++++++++++++ 10 files changed, 1037 insertions(+), 1 deletion(-) create mode 100644 modelopt/torch/sparsity/state_sparsity/__init__.py create mode 100644 modelopt/torch/sparsity/state_sparsity/api.py create mode 100644 modelopt/torch/sparsity/state_sparsity/config.py create mode 100644 modelopt/torch/sparsity/state_sparsity/conversion.py create mode 100644 modelopt/torch/sparsity/state_sparsity/mode.py create mode 100644 modelopt/torch/sparsity/state_sparsity/policy.py create mode 100644 tests/unit/torch/sparsity/state_sparsity/test_dasc.py diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 59e999ea2ee..30c8d480c4d 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -13,6 +13,7 @@ Changelog *Misc* - A tracked ``examples/hf_ptq/hf_ptq.py`` run now writes ``.experiment.json`` into ``--export_path`` and uploads the same file with the run, so a checkpoint on disk names the experiment and MLflow run id that produced it. The pointer is written only once the export completes, and an export that is not tracked removes one it would otherwise inherit from a reused ``--export_path`` or from a quantized source checkpoint. +- Add an experimental GDN DASC state-sparsity API that derives whole-head decay policies and selects a checkpoint-specific recovery window from caller-supplied quality and storage measurements. The first release exports restorable policy metadata only; serving runtimes must implement checkpoint packing and DASC-NR or DASC-WR recovery. **Backward Breaking Changes** diff --git a/docs/source/guides/6_sparsity.rst b/docs/source/guides/6_sparsity.rst index 12a17b3c66c..d507dd8ba3f 100644 --- a/docs/source/guides/6_sparsity.rst +++ b/docs/source/guides/6_sparsity.rst @@ -114,6 +114,63 @@ To restore the saved sparse model you can use Please see :ref:`saving and restoring of ModelOpt-modified models ` to learn about all the available options for saving and restoring. +Decay-aware recurrent-state sparsity (experimental) +--------------------------------------------------- + +The :mod:`modelopt.torch.sparsity.state_sparsity` package calibrates `DASC +`_ policies for persisted Gated DeltaNet (GDN) prefix state. It +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. + +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 +must implement storage and recovery while preserving convolution state exactly and materializing +ordinary dense recurrent state before continuation. + +.. code-block:: python + + import modelopt.torch.sparsity.state_sparsity as mtss + + config = { + "variant": "dasc_wr", # "dasc_nr" uses zero recovery instead + "wmax_candidates": [32], + "model_id": "org/model", + "model_revision": "immutable-model-revision", + "model_config_id": "sha256:", + "calibration_data_id": "sha256:", + } + measurements = [ + { + "variant": "dasc_wr", + "wmax": 32, + "retained_heads": 40, + "total_heads": 96, + "checkpoint_savings": 0.21, + "quality": [ + { + "slice_id": "validation-context-1024", + "perplexity_retention": 0.999, + "top1_agreement": 0.99, + "finite_continuation_logits": True, + "retained_state_exact": True, + "omitted_state_matches_recovery": True, + "convolution_state_exact": True, + } + ], + }, + ] + + model = mtss.calibrate(model, config, measurements) + policy = mtss.export_policy(model) + +``dasc_nr`` and ``dasc_wr`` remain explicit deployment contracts: DASC-NR restores omitted heads +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. + .. _sparsity-concepts: Sparsity Concepts diff --git a/modelopt/torch/sparsity/__init__.py b/modelopt/torch/sparsity/__init__.py index 2013fded1ae..9196d84e3cc 100644 --- a/modelopt/torch/sparsity/__init__.py +++ b/modelopt/torch/sparsity/__init__.py @@ -15,10 +15,12 @@ """API for sparsification algorithms. -This module provides access to both weight sparsity and attention sparsity algorithms. +This module provides access to weight, attention, and recurrent-state sparsity algorithms. For backward compatibility, weight sparsity APIs are re-exported at the module level. """ +from . import state_sparsity + # Import weight sparsity for backward compatibility from .weight_sparsity import mode, module, plugins from .weight_sparsity.sparsification import * diff --git a/modelopt/torch/sparsity/state_sparsity/__init__.py b/modelopt/torch/sparsity/state_sparsity/__init__.py new file mode 100644 index 00000000000..97fe2d89b13 --- /dev/null +++ b/modelopt/torch/sparsity/state_sparsity/__init__.py @@ -0,0 +1,21 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Decay-aware sparsity policies for persisted recurrent state.""" + +from . import mode +from .api import * +from .config import * +from .policy import * diff --git a/modelopt/torch/sparsity/state_sparsity/api.py b/modelopt/torch/sparsity/state_sparsity/api.py new file mode 100644 index 00000000000..60d05620e46 --- /dev/null +++ b/modelopt/torch/sparsity/state_sparsity/api.py @@ -0,0 +1,80 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Public DASC state-sparsity APIs.""" + +import copy +from collections.abc import Iterable +from typing import Any + +from torch import nn + +from modelopt.torch.opt.conversion import apply_mode + +from .config import DASCCalibrationMeasurement, DASCConfig +from .conversion import get_attached_dasc_policy +from .mode import DASCModeRegistry +from .policy import validate_dasc_decay_parameters, validate_dasc_model_structure + +__all__ = ["calibrate", "export_policy"] + + +def calibrate( + model: nn.Module, + config: dict[str, Any] | DASCConfig, + measurements: Iterable[DASCCalibrationMeasurement | dict], +) -> nn.Module: + """Calibrate and attach a DASC policy without changing model execution. + + ``measurements`` must contain exactly one entry for every configured ``Wmax`` candidate. + The largest candidate passing every quality, lifecycle, and storage gate is selected. + + Example:: + + import modelopt.torch.sparsity.state_sparsity as mtss + + model = mtss.calibrate(model, config, measurements) + deployment_policy = mtss.export_policy(model) + + Args: + model: Model containing GatedDeltaNet modules with one-dimensional ``A_log`` and + ``dt_bias`` tensors. + config: Checkpoint provenance, candidate windows, and quality gates. + measurements: Quality and checkpoint-storage results produced by the caller's paired + dense-versus-DASC calibration workflow. + + Returns: + The input model with a serializable DASC policy attached through ModelOpt state. + """ + config_dict = config.model_dump() if isinstance(config, DASCConfig) else config + return apply_mode( + model, + mode=[("dasc", config_dict)], + registry=DASCModeRegistry, + mode_kwargs={"measurements": measurements}, + ) + + +def export_policy(model: nn.Module) -> dict[str, Any]: + """Export a JSON-safe DASC policy after validating model structure and decay parameters. + + This policy does not implement checkpoint packing or recovery. A serving backend must preserve + convolution state, store retained complete GDN heads, recover omitted heads according to the + declared variant, and materialize the ordinary dense runtime state before continuation. + """ + policy = get_attached_dasc_policy(model) + validate_dasc_model_structure(model, policy) + validate_dasc_decay_parameters(model, policy) + return copy.deepcopy(policy.model_dump(mode="json")) diff --git a/modelopt/torch/sparsity/state_sparsity/config.py b/modelopt/torch/sparsity/state_sparsity/config.py new file mode 100644 index 00000000000..c7bca8be9f3 --- /dev/null +++ b/modelopt/torch/sparsity/state_sparsity/config.py @@ -0,0 +1,232 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Configuration and result schemas for DASC state sparsity.""" + +import math +from typing import Literal + +from pydantic import Field, field_validator, model_validator + +from modelopt.torch.opt.config import ModeloptBaseConfig, ModeloptField + +__all__ = [ + "DASCCalibrationMeasurement", + "DASCConfig", + "DASCPolicy", + "DASCQualityMeasurement", +] + + +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) + 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) + omitted_state_matches_recovery: bool = Field(strict=True) + convolution_state_exact: bool = Field(strict=True) + + +class DASCCalibrationMeasurement(ModeloptBaseConfig): + """Caller-supplied measurements for one candidate recovery window.""" + + variant: Literal["dasc_nr", "dasc_wr"] + wmax: int = Field(strict=True, gt=0) + retained_heads: int = Field(strict=True, ge=0) + total_heads: int = Field(strict=True, gt=0) + checkpoint_savings: float = Field(ge=0.0, lt=1.0, allow_inf_nan=False) + quality: list[DASCQualityMeasurement] = Field(min_length=1) + + @field_validator("quality") + @classmethod + def validate_unique_slices( + cls, quality: list[DASCQualityMeasurement] + ) -> list[DASCQualityMeasurement]: + """Require one result per named calibration slice.""" + slice_ids = [measurement.slice_id for measurement in quality] + if len(slice_ids) != len(set(slice_ids)): + raise ValueError("quality slice_id values must be unique") + return quality + + +class DASCConfig(ModeloptBaseConfig): + """Configuration for GDN decay-aware state checkpoint sparsity.""" + + variant: Literal["dasc_nr", "dasc_wr"] = ModeloptField( + default="dasc_wr", + description="Use zero recovery (DASC-NR) or suffix replay recovery (DASC-WR).", + ) + epsilon: float = ModeloptField( + default=1e-3, + description="Retained contribution threshold used to derive static decay horizons.", + ) + static_gate_input: float = ModeloptField( + default=-0.3, + description="Static gate input added to each GDN head's dt_bias.", + ) + wmax_candidates: list[int] = ModeloptField( + default=[8, 16, 32, 64, 128, 256], + description="Positive candidate windows evaluated during offline calibration.", + ) + min_perplexity_retention: float = ModeloptField(default=0.995) + min_top1_agreement: float = ModeloptField(default=0.98) + min_checkpoint_savings: float = ModeloptField(default=0.2) + model_id: str = ModeloptField(default="", validate_default=True) + model_revision: str = ModeloptField(default="", validate_default=True) + model_config_id: str = ModeloptField( + default="", + description="Immutable identifier or digest for the model architecture configuration.", + validate_default=True, + ) + calibration_data_id: str = ModeloptField( + default="", + description="Immutable identifier or digest for the calibration data and protocol.", + validate_default=True, + ) + granularity: Literal["gdn_head"] = ModeloptField(default="gdn_head") + preserve_convolution_state: Literal[True] = ModeloptField(default=True) + + @field_validator("epsilon") + @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)") + 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") + return value + + @field_validator("wmax_candidates", mode="before") + @classmethod + def validate_wmax_candidates(cls, candidates: object) -> object: + """Require unique positive integer windows without power-of-two restrictions.""" + if not isinstance(candidates, list) or not candidates: + raise ValueError("wmax_candidates must be a non-empty list") + if any( + isinstance(value, bool) or not isinstance(value, int) or value <= 0 + for value in candidates + ): + raise ValueError("wmax_candidates must contain only positive integers") + if len(candidates) != len(set(candidates)): + raise ValueError("wmax_candidates must be unique") + return sorted(candidates) + + @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]") + return value + + @field_validator("min_top1_agreement") + @classmethod + def validate_top1_gate(cls, value: float) -> float: + """Require a finite agreement gate in [0, 1].""" + if not math.isfinite(value) or not 0.0 <= value <= 1.0: + raise ValueError("min_top1_agreement must be finite and in [0, 1]") + return value + + @field_validator("min_checkpoint_savings") + @classmethod + def validate_savings_gate(cls, value: float) -> float: + """Require a finite physical checkpoint-savings gate in [0, 1).""" + if not math.isfinite(value) or not 0.0 <= value < 1.0: + raise ValueError("min_checkpoint_savings must be finite and in [0, 1)") + return value + + @field_validator("model_id", "model_revision", "model_config_id", "calibration_data_id") + @classmethod + def validate_provenance(cls, value: str) -> str: + """Require explicit immutable provenance instead of inferred defaults.""" + if not value.strip(): + raise ValueError("DASC provenance fields must be non-empty") + return value + + +class DASCLayerPolicy(ModeloptBaseConfig): + """Serializable whole-head policy for one GDN layer.""" + + num_heads: int = Field(strict=True, gt=0) + static_horizons: list[float] = Field(min_length=1) + retained_heads: list[int] + omitted_heads: list[int] + + @model_validator(mode="after") + def validate_partition(self) -> "DASCLayerPolicy": + """Validate that retained and omitted indices partition every GDN head.""" + if len(self.static_horizons) != self.num_heads or not all( + math.isfinite(value) and value > 0.0 for value in self.static_horizons + ): + raise ValueError("static_horizons must contain one finite positive value per head") + if sorted(self.retained_heads + self.omitted_heads) != list(range(self.num_heads)): + raise ValueError("retained_heads and omitted_heads must partition all head indices") + return self + + +class DASCPolicy(ModeloptBaseConfig): + """Standalone JSON-safe DASC deployment policy.""" + + format_version: Literal[1] = 1 + variant: Literal["dasc_nr", "dasc_wr"] + recovery: Literal["zero", "suffix_replay"] + epsilon: float + static_gate_input: float + selected_wmax: int = Field(strict=True, gt=0) + wmax_candidates: list[int] = Field(min_length=1) + quality_gates: dict[str, float] + model_id: str + model_revision: str + model_config_id: str + calibration_data_id: str + granularity: Literal["gdn_head"] + 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}$") + layers: dict[str, DASCLayerPolicy] = Field(min_length=1) + measurements: list[DASCCalibrationMeasurement] = Field(min_length=1) + + @model_validator(mode="after") + def validate_policy(self) -> "DASCPolicy": + """Reject inconsistent variant, candidate, measurement, or mask metadata.""" + expected_recovery = "zero" if self.variant == "dasc_nr" else "suffix_replay" + if self.recovery != expected_recovery: + raise ValueError(f"{self.variant} requires recovery={expected_recovery!r}") + if self.selected_wmax not in self.wmax_candidates: + raise ValueError("selected_wmax must be present in wmax_candidates") + measured = [measurement.wmax for measurement in self.measurements] + if sorted(measured) != sorted(self.wmax_candidates) or len(measured) != len(set(measured)): + raise ValueError("measurements must contain exactly one result per wmax candidate") + if any(measurement.variant != self.variant for measurement in self.measurements): + raise ValueError("measurement variants must match the policy variant") + for layer in self.layers.values(): + expected_retained = [ + head + for head, horizon in enumerate(layer.static_horizons) + if horizon > self.selected_wmax + ] + if layer.retained_heads != expected_retained: + raise ValueError("retained_heads do not match the selected decay threshold") + return self diff --git a/modelopt/torch/sparsity/state_sparsity/conversion.py b/modelopt/torch/sparsity/state_sparsity/conversion.py new file mode 100644 index 00000000000..5e756a2587d --- /dev/null +++ b/modelopt/torch/sparsity/state_sparsity/conversion.py @@ -0,0 +1,107 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""ModelOpt conversion and restoration for DASC policy metadata.""" + +import copy +from collections.abc import Iterable + +from pydantic import ValidationError +from torch import nn + +from modelopt.torch.opt.conversion import ApplyModeError +from modelopt.torch.opt.mode import ConvertReturnType, MetadataDict + +from .config import DASCCalibrationMeasurement, DASCConfig, DASCPolicy +from .policy import build_dasc_policy, validate_dasc_decay_parameters, validate_dasc_model_structure + +__all__ = [] + +_DASC_POLICY_ATTRIBUTE = "_modelopt_dasc_policy" + + +def _attach_policy(model: nn.Module, policy: DASCPolicy) -> None: + setattr(model, _DASC_POLICY_ATTRIBUTE, policy.model_dump(mode="json")) + + +def convert_dasc_model( + model: nn.Module, + config: DASCConfig, + *, + measurements: Iterable[DASCCalibrationMeasurement | dict], +) -> ConvertReturnType: + """Analyze GDN decay and attach the selected DASC policy without changing execution.""" + policy = build_dasc_policy(model, config, measurements) + _attach_policy(model, policy) + return model, {"policy": policy.model_dump(mode="json")} + + +def restore_dasc_model(model: nn.Module, config: DASCConfig, metadata: MetadataDict) -> nn.Module: + """Restore and structurally validate a serialized DASC policy.""" + if set(metadata) != {"policy"}: + raise ApplyModeError("DASC metadata must contain only the policy field") + try: + policy = DASCPolicy(**metadata["policy"]) + except (TypeError, ValidationError) as error: + raise ApplyModeError(f"Invalid DASC policy metadata: {error}") from error + + expected_config = { + "variant": config.variant, + "epsilon": config.epsilon, + "static_gate_input": config.static_gate_input, + "wmax_candidates": config.wmax_candidates, + "quality_gates": { + "min_perplexity_retention": config.min_perplexity_retention, + "min_top1_agreement": config.min_top1_agreement, + "min_checkpoint_savings": config.min_checkpoint_savings, + }, + "model_id": config.model_id, + "model_revision": config.model_revision, + "model_config_id": config.model_config_id, + "calibration_data_id": config.calibration_data_id, + "granularity": config.granularity, + "preserve_convolution_state": config.preserve_convolution_state, + } + mismatched = { + key: (value, getattr(policy, key)) + for key, value in expected_config.items() + if value != getattr(policy, key) + } + if mismatched: + raise ApplyModeError(f"DASC policy metadata does not match its mode config: {mismatched}") + + validate_dasc_model_structure(model, policy) + _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.""" + 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) + metadata.clear() + metadata["policy"] = copy.deepcopy(policy.model_dump(mode="json")) + + +def get_attached_dasc_policy(model: nn.Module) -> DASCPolicy: + """Return the validated policy attached by conversion or restoration.""" + try: + return DASCPolicy(**getattr(model, _DASC_POLICY_ATTRIBUTE)) + except (AttributeError, TypeError, ValidationError) as error: + raise ApplyModeError("Model has no valid attached DASC policy") from error diff --git a/modelopt/torch/sparsity/state_sparsity/mode.py b/modelopt/torch/sparsity/state_sparsity/mode.py new file mode 100644 index 00000000000..23b3a17ffc3 --- /dev/null +++ b/modelopt/torch/sparsity/state_sparsity/mode.py @@ -0,0 +1,74 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""DASC state-sparsity mode descriptor.""" + +from typing import cast + +from modelopt.torch.opt.config import ModeloptBaseConfig +from modelopt.torch.opt.mode import ( + ConvertEntrypoint, + ModeDescriptor, + RestoreEntrypoint, + UpdateEntrypoint, + _ModeRegistryCls, +) + +from .config import DASCConfig +from .conversion import convert_dasc_model, restore_dasc_model, update_dasc_metadata + +__all__ = ["DASCModeRegistry"] + +DASCModeRegistry = _ModeRegistryCls("state_sparsity") + + +@DASCModeRegistry.register_mode +class DASCModeDescriptor(ModeDescriptor): + """Describe checkpoint-specific GDN DASC policy calibration.""" + + @property + def name(self) -> str: + """Return the mode name.""" + return "dasc" + + @property + def config_class(self) -> type[ModeloptBaseConfig]: + """Return the validated DASC configuration class.""" + return DASCConfig + + @property + def next_prohibited_modes(self) -> set[str]: + """Prevent applying DASC twice to the same model state.""" + return {"dasc"} + + @property + def convert(self) -> ConvertEntrypoint: + """Return the DASC calibration entrypoint.""" + return cast("ConvertEntrypoint", convert_dasc_model) + + @property + def restore(self) -> RestoreEntrypoint: + """Return the DASC restore entrypoint.""" + return restore_dasc_model + + @property + def update_for_save(self) -> UpdateEntrypoint: + """Return the metadata refresh entrypoint.""" + return update_dasc_metadata + + @property + def update_for_new_mode(self) -> UpdateEntrypoint: + """Return the metadata refresh entrypoint.""" + return update_dasc_metadata diff --git a/modelopt/torch/sparsity/state_sparsity/policy.py b/modelopt/torch/sparsity/state_sparsity/policy.py new file mode 100644 index 00000000000..7342a66bf17 --- /dev/null +++ b/modelopt/torch/sparsity/state_sparsity/policy.py @@ -0,0 +1,258 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Decay analysis and policy selection for GDN state sparsity.""" + +import hashlib +import json +from collections.abc import Iterable + +import torch +import torch.nn.functional as F +from torch import nn + +from modelopt.torch.opt.conversion import ApplyModeError + +from .config import DASCCalibrationMeasurement, DASCConfig, DASCLayerPolicy, DASCPolicy + +__all__ = ["analyze_gdn_decay", "compute_gdn_decay_horizons"] + + +def compute_gdn_decay_horizons( + a_log: torch.Tensor, + dt_bias: torch.Tensor, + *, + epsilon: float = 1e-3, + 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 0.0 < epsilon < 1.0: + raise ValueError("epsilon must be in (0, 1)") + + a_log_cpu = a_log.detach().to(device="cpu", dtype=torch.float64) + dt_bias_cpu = dt_bias.detach().to(device="cpu", dtype=torch.float64) + if not torch.isfinite(a_log_cpu).all() or not torch.isfinite(dt_bias_cpu).all(): + raise ValueError("GDN decay parameters must be finite") + + decay = -torch.exp(a_log_cpu) * F.softplus(dt_bias_cpu + static_gate_input) + horizons = torch.log(torch.tensor(epsilon, dtype=torch.float64)) / 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 + + +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 _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())) + + +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.""" + horizons = {} + for name, module in _get_gdn_modules(model).items(): + try: + layer_horizons = compute_gdn_decay_horizons( + module.A_log, + module.dt_bias, + epsilon=epsilon, + static_gate_input=static_gate_input, + ) + except ValueError as error: + raise ApplyModeError( + f"Invalid GDN decay parameters in module {name!r}: {error}" + ) from error + horizons[name] = layer_horizons.tolist() + return horizons + + +def _canonical_sha256(value: object) -> str: + 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]]: + return [ + { + "name": name, + "num_heads": int(module.A_log.numel()), + } + for name, module in modules.items() + ] + + +def _decay_parameters(modules: dict[str, nn.Module]) -> list[dict[str, object]]: + 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(), + } + for name, module in modules.items() + ] + + +def _validate_measurement_coverage( + config: DASCConfig, measurements: list[DASCCalibrationMeasurement] +) -> None: + measured = [measurement.wmax for measurement in measurements] + unexpected = sorted(set(measured) - set(config.wmax_candidates)) + missing = sorted(set(config.wmax_candidates) - set(measured)) + if unexpected or missing or len(measured) != len(set(measured)): + raise ApplyModeError( + "DASC measurements must contain exactly one result per configured candidate; " + f"missing={missing}, unexpected={unexpected}, duplicates={len(measured) != len(set(measured))}" + ) + if any(measurement.variant != config.variant for measurement in measurements): + raise ApplyModeError( + f"All DASC measurements must use configured variant {config.variant!r}" + ) + + +def _validate_measurement_geometry( + horizons: dict[str, list[float]], measurements: list[DASCCalibrationMeasurement] +) -> None: + total_heads = sum(len(layer_horizons) for layer_horizons in horizons.values()) + for measurement in measurements: + retained_heads = sum( + horizon > measurement.wmax + for layer_horizons in horizons.values() + for horizon in layer_horizons + ) + if measurement.total_heads != total_heads or measurement.retained_heads != retained_heads: + raise ApplyModeError( + f"DASC measurement geometry for Wmax={measurement.wmax} does not match the model; " + f"expected retained/total={retained_heads}/{total_heads}, got " + f"{measurement.retained_heads}/{measurement.total_heads}" + ) + + +def _candidate_passes(config: DASCConfig, measurement: DASCCalibrationMeasurement) -> bool: + 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 + and result.finite_continuation_logits + and result.retained_state_exact + and result.omitted_state_matches_recovery + and result.convolution_state_exact + for result in measurement.quality + ) + + +def build_dasc_policy( + model: nn.Module, + config: DASCConfig, + measurements: Iterable[DASCCalibrationMeasurement | dict], +) -> DASCPolicy: + """Build a checkpoint-specific policy from decay parameters and measured quality.""" + try: + validated_measurements = [ + measurement + if isinstance(measurement, DASCCalibrationMeasurement) + else DASCCalibrationMeasurement(**measurement) + for measurement in measurements + ] + except (TypeError, ValueError) as error: + raise ApplyModeError(f"Invalid DASC calibration measurements: {error}") from error + _validate_measurement_coverage(config, validated_measurements) + 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 + ) + _validate_measurement_geometry(horizons, validated_measurements) + + passing = [ + measurement.wmax + for measurement in validated_measurements + if _candidate_passes(config, measurement) + ] + if not passing: + raise ApplyModeError( + "No DASC Wmax candidate passed every configured quality and storage gate" + ) + selected_wmax = max(passing) + + layers = {} + for name, values in horizons.items(): + retained = [head for head, horizon in enumerate(values) if horizon > selected_wmax] + layers[name] = DASCLayerPolicy( + num_heads=len(values), + static_horizons=values, + retained_heads=retained, + omitted_heads=[head for head in range(len(values)) if head not in retained], + ) + + return DASCPolicy( + variant=config.variant, + recovery="zero" if config.variant == "dasc_nr" else "suffix_replay", + epsilon=config.epsilon, + static_gate_input=config.static_gate_input, + selected_wmax=selected_wmax, + wmax_candidates=config.wmax_candidates, + quality_gates={ + "min_perplexity_retention": config.min_perplexity_retention, + "min_top1_agreement": config.min_top1_agreement, + "min_checkpoint_savings": config.min_checkpoint_savings, + }, + model_id=config.model_id, + model_revision=config.model_revision, + model_config_id=config.model_config_id, + calibration_data_id=config.calibration_data_id, + 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)), + layers=layers, + measurements=validated_measurements, + ) + + +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: + 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.""" + 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") diff --git a/tests/unit/torch/sparsity/state_sparsity/test_dasc.py b/tests/unit/torch/sparsity/state_sparsity/test_dasc.py new file mode 100644 index 00000000000..c212808c046 --- /dev/null +++ b/tests/unit/torch/sparsity/state_sparsity/test_dasc.py @@ -0,0 +1,204 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""CPU tests for DASC state-sparsity policy calibration.""" + +import copy +import json + +import pytest +import torch +from pydantic import ValidationError +from torch import nn + +import modelopt.torch.opt as mto +import modelopt.torch.sparsity.state_sparsity as mtss +from modelopt.torch.opt.conversion import ApplyModeError + + +class TinyGatedDeltaNet(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) + self.A_log = nn.Parameter(a_log) + self.dt_bias = nn.Parameter(dt_bias) + + def forward(self, inputs): + return inputs + + +class TinyGatedDeltaNetForCausalLM(nn.Module): + """Parent model whose name must not be mistaken for a GDN layer.""" + + def __init__(self, num_heads: int = 2): + super().__init__() + self.linear_attn = TinyGatedDeltaNet(num_heads) + + def forward(self, inputs): + return self.linear_attn(inputs) + + +def _config(**overrides): + config = { + "variant": "dasc_wr", + "epsilon": 1e-3, + "static_gate_input": 0.0, + "wmax_candidates": [7, 11], + "min_perplexity_retention": 0.995, + "min_top1_agreement": 0.98, + "min_checkpoint_savings": 0.2, + "model_id": "tiny-gdn", + "model_revision": "revision-1", + "model_config_id": "sha256:config", + "calibration_data_id": "sha256:calibration", + } + config.update(overrides) + return config + + +def _candidate(wmax, *, variant="dasc_wr", top1=0.99, convolution_state_exact=True): + return { + "variant": variant, + "wmax": wmax, + "retained_heads": 1, + "total_heads": 2, + "checkpoint_savings": 0.3, + "quality": [ + { + "slice_id": "validation-0", + "perplexity_retention": 0.999, + "top1_agreement": top1, + "finite_continuation_logits": True, + "retained_state_exact": True, + "omitted_state_matches_recovery": True, + "convolution_state_exact": convolution_state_exact, + } + ], + } + + +def test_calibrate_selects_largest_passing_candidate_and_round_trips(): + model = TinyGatedDeltaNetForCausalLM() + original_state = {name: value.clone() for name, value in model.state_dict().items()} + + model = mtss.calibrate(model, _config(), [_candidate(11, top1=0.9), _candidate(7)]) + policy = mtss.export_policy(model) + + assert policy["selected_wmax"] == 7 + assert policy["variant"] == "dasc_wr" + assert policy["recovery"] == "suffix_replay" + assert policy["granularity"] == "gdn_head" + assert policy["preserve_convolution_state"] is True + assert policy["active_runtime_state"] == "dense" + 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] + assert all( + torch.equal(original_state[name], value) for name, value in model.state_dict().items() + ) + json.dumps(policy) + + restored = mto.restore_from_modelopt_state( + TinyGatedDeltaNetForCausalLM(), mto.modelopt_state(model) + ) + assert mtss.export_policy(restored) == policy + + policy["selected_wmax"] = 999 + assert mtss.export_policy(model)["selected_wmax"] == 7 + + +@pytest.mark.parametrize( + ("variant", "recovery"), [("dasc_nr", "zero"), ("dasc_wr", "suffix_replay")] +) +def test_variant_is_explicit_in_exported_policy(variant, recovery): + model = mtss.calibrate( + TinyGatedDeltaNetForCausalLM(), + _config(variant=variant, wmax_candidates=[7]), + [_candidate(7, variant=variant)], + ) + policy = mtss.export_policy(model) + assert policy["variant"] == variant + assert policy["recovery"] == recovery + + +@pytest.mark.parametrize( + "override", + [ + {"wmax_candidates": []}, + {"wmax_candidates": [0]}, + {"wmax_candidates": [7, 7]}, + {"model_revision": ""}, + {"preserve_convolution_state": False}, + ], +) +def test_config_fails_closed(override): + with pytest.raises(ValidationError): + mtss.DASCConfig(**_config(**override)) + + assert mtss.DASCConfig(**_config(wmax_candidates=[7])).wmax_candidates == [7] + + +def test_calibration_fails_closed_on_measurements_and_model_mismatch(): + with pytest.raises(ApplyModeError, match="exactly one result"): + mtss.calibrate(TinyGatedDeltaNetForCausalLM(), _config(), [_candidate(7)]) + + with pytest.raises(ApplyModeError, match="No DASC Wmax candidate"): + mtss.calibrate( + TinyGatedDeltaNetForCausalLM(), + _config(wmax_candidates=[7]), + [_candidate(7, convolution_state_exact=False)], + ) + + mismatched_geometry = _candidate(7) + mismatched_geometry["retained_heads"] = 2 + with pytest.raises(ApplyModeError, match="geometry"): + mtss.calibrate( + TinyGatedDeltaNetForCausalLM(), + _config(wmax_candidates=[7]), + [mismatched_geometry], + ) + + class NotGDN(nn.Module): + pass + + with pytest.raises(ApplyModeError, match="no GatedDeltaNet modules"): + mtss.calibrate(NotGDN(), _config(wmax_candidates=[7]), [_candidate(7)]) + + +def test_export_rejects_changed_decay_parameters_and_restore_rejects_structure(): + model = mtss.calibrate( + TinyGatedDeltaNetForCausalLM(), _config(wmax_candidates=[7]), [_candidate(7)] + ) + state = mto.modelopt_state(model) + + with torch.no_grad(): + model.linear_attn.A_log.add_(1.0) + with pytest.raises(ApplyModeError, match="decay parameters"): + mtss.export_policy(model) + with pytest.raises(ApplyModeError, match="decay parameters"): + mto.modelopt_state(model) + + with pytest.raises(ApplyModeError, match="module structure"): + mto.restore_from_modelopt_state(TinyGatedDeltaNetForCausalLM(num_heads=3), state) + + tampered_state = copy.deepcopy(state) + tampered_state["modelopt_state_dict"][0][1]["metadata"]["policy"]["quality_gates"][ + "min_top1_agreement" + ] = 0.5 + with pytest.raises(ApplyModeError, match="does not match its mode config"): + mto.restore_from_modelopt_state(TinyGatedDeltaNetForCausalLM(), tampered_state) From 0859c13d0f827c479a512bd370c2b0b5528a7037 Mon Sep 17 00:00:00 2001 From: kaix-nv Date: Thu, 10 Sep 2026 21:51:51 -0700 Subject: [PATCH 2/2] Address DASC review feedback (#2387) ## Summary Consolidates the complete reviewed fix stack for #2375 into one DCO-safe commit: - harden package exports, measurement semantics, wrapper handling, and actionable calibration errors - validate exact installed GDN identities plus ModelOpt dynamic subclasses; reject lookalikes, ordinary subclasses, incomplete layers, and partial layer sets - make stale checkpoints saveable and restorable while keeping deployment export strict - make DASC recalibration replace and deduplicate existing mode state without stale-metadata refresh - record the declared decay-parameter checkpoint storage dtype and use derived FP16/BF16/FP32 rounding bounds - preserve BF16/FP16 storage and wider/cross-dtype reload compatibility without globally widening FP32 tolerance - add installed Transformers path coverage, optional Megatron gating, lifecycle, tamper, lossy-cast, and mixed-layer regressions - document the explicit storage-dtype contract This consolidated PR supersedes the mechanically stacked review-fix PRs #2377, #2378, #2379, #2380, #2382, #2383, #2384, and #2385. Its tree is byte-identical to the independently reviewed leaf commit from #2386. ## Validation - focused DASC suite: 23 passed, 1 absent optional Megatron skip - DASC plus weight sparsity plus attention sparsity compatibility suite: 134 passed, 1 optional skip - DASC package coverage: 408/408 statements, 100% - full pre-commit on all touched files: passed - real Transformers Qwen3NextGatedDeltaNet BF16 storage to FP32 reload smoke: passed - commit author and Signed-off-by identity both use kaix-nv ## Summary by CodeRabbit * **New Features** * Added support for configuring decay-parameter storage precision with FP16, BF16, or FP32. * Added safer recalibration that replaces existing DASC state. * Expanded compatibility with supported GDN adapter classes and model wrappers. * Added improved validation for sparsity policies, measurements, model structure, and decay parameters. * Added support for perplexity-retention values above 1. * **Documentation** * Clarified evaluation responsibilities, recalibration behavior, stale-policy handling, supported adapters, and dtype requirements. --------- Signed-off-by: kaix-nv --- docs/source/guides/6_sparsity.rst | 33 +- .../torch/sparsity/state_sparsity/__init__.py | 26 +- modelopt/torch/sparsity/state_sparsity/api.py | 16 +- .../torch/sparsity/state_sparsity/config.py | 73 ++- .../sparsity/state_sparsity/conversion.py | 81 ++- .../torch/sparsity/state_sparsity/mode.py | 2 +- .../torch/sparsity/state_sparsity/policy.py | 362 ++++++++++-- .../sparsity/state_sparsity/test_dasc.py | 556 +++++++++++++++++- 8 files changed, 1060 insertions(+), 89 deletions(-) diff --git a/docs/source/guides/6_sparsity.rst b/docs/source/guides/6_sparsity.rst index d507dd8ba3f..70166aecac4 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 @@ -136,12 +138,24 @@ 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", "model_id": "org/model", "model_revision": "immutable-model-revision", "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", @@ -170,6 +184,23 @@ 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 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. 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/__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..1a8f87e0236 100644 --- a/modelopt/torch/sparsity/state_sparsity/config.py +++ b/modelopt/torch/sparsity/state_sparsity/config.py @@ -16,25 +16,64 @@ """Configuration and result schemas for DASC state sparsity.""" import math +from numbers import Real 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", ] +_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.""" 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,18 +106,24 @@ 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).", ) 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: _DecayParameterStorageDtype = 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.", @@ -105,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") @@ -135,9 +178,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 +231,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: _DecayParameterStorageDtype = "float32" selected_wmax: int = Field(strict=True, gt=0) wmax_candidates: list[int] = Field(min_length=1) quality_gates: dict[str, float] @@ -204,7 +250,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..213cc351f45 100644 --- a/modelopt/torch/sparsity/state_sparsity/conversion.py +++ b/modelopt/torch/sparsity/state_sparsity/conversion.py @@ -16,16 +16,23 @@ """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 +from .policy import ( + _DASCRecoverableStalenessError, + build_dasc_policy, + validate_dasc_decay_parameters, + validate_dasc_model_structure, +) __all__ = [] @@ -33,6 +40,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 +48,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 +74,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 +96,72 @@ 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 _DASCRecoverableStalenessError 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 while allowing recoverable policy staleness to remain serializable.""" + 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) + 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")) +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..7baf53e9586 100644 --- a/modelopt/torch/sparsity/state_sparsity/policy.py +++ b/modelopt/torch/sparsity/state_sparsity/policy.py @@ -16,19 +16,97 @@ """Decay analysis and policy selection for GDN state sparsity.""" import hashlib +import importlib +import importlib.util 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 .config import DASCCalibrationMeasurement, DASCConfig, DASCLayerPolicy, DASCPolicy +from modelopt.torch.opt.dynamic import DynamicModule +from modelopt.torch.utils import unwrap_model + +from .config import ( + DASCCalibrationMeasurement, + DASCConfig, + DASCLayerPolicy, + DASCPolicy, + _DecayParameterStorageDtype, + _validate_analysis_arguments, +) __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: 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.""" + 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, @@ -38,56 +116,106 @@ 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 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 -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 _DASCDecayParametersUnavailableError( + "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, +def _analyze_gdn_modules( + modules: dict[str, nn.Module], *, - epsilon: float = 1e-3, - static_gate_input: float = -0.3, + epsilon: float, + static_gate_input: float, + storage_dtype: torch.dtype | None = None, ) -> dict[str, list[float]]: - """Return deterministic per-head horizons for every GDN module in a model.""" + """Compute horizons, optionally from checkpoint-storage-canonical parameters.""" horizons = {} - for name, module in _get_gdn_modules(model).items(): + 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, ) @@ -99,12 +227,40 @@ def analyze_gdn_decay( return horizons +def analyze_gdn_decay( + model: nn.Module, + *, + epsilon: float = 1e-3, + static_gate_input: float = -0.3, + decay_parameter_storage_dtype: _DecayParameterStorageDtype | None = None, +) -> dict[str, list[float]]: + """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, + storage_dtype=storage_dtype, + ) + + 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 +270,82 @@ 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 _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 + 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 + 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 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) + 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 +363,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 +380,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 +411,11 @@ 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, + storage_dtype=_STORAGE_DTYPES[config.decay_parameter_storage_dtype], ) _validate_measurement_geometry(horizons, validated_measurements) @@ -222,6 +445,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 +460,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 +471,60 @@ 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: - raise ApplyModeError("DASC policy does not match the model's GDN module structure") + 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 _DASCModelStructureMismatchError( + "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 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) - 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") + for name, module in modules.items(): + layer = policy.layers[name] + 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, + static_gate_input=policy.static_gate_input, + storage_dtype=_STORAGE_DTYPES[policy.decay_parameter_storage_dtype], + ) + declared_retained = set(layer.retained_heads) + for head, (head_lower, head_upper) in enumerate(zip(lower, upper)): + retained_is_impossible = ( + head in declared_retained and head_upper <= policy.selected_wmax + ) + omitted_is_impossible = ( + head not in declared_retained and head_lower > policy.selected_wmax + ) + if retained_is_impossible or omitted_is_impossible: + raise ApplyModeError( + "DASC policy head mask does not match current decay parameters in layer " + f"{name!r}" + ) + stored = torch.tensor(layer.static_horizons, 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) + ): + 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..cad44b1501d 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,367 @@ 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_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( + 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 + + +@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( TinyGatedDeltaNetForCausalLM(), _config(wmax_candidates=[7]), [_candidate(7)] ) @@ -188,13 +604,48 @@ 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) + + 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"][ @@ -202,3 +653,84 @@ 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_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( + 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)