diff --git a/modelopt/torch/sparsity/state_sparsity/conversion.py b/modelopt/torch/sparsity/state_sparsity/conversion.py index 3b6a44bef41..a748cc813b6 100644 --- a/modelopt/torch/sparsity/state_sparsity/conversion.py +++ b/modelopt/torch/sparsity/state_sparsity/conversion.py @@ -90,7 +90,13 @@ def restore_dasc_model(model: nn.Module, config: DASCConfig, metadata: MetadataD if mismatched: raise ApplyModeError(f"DASC policy metadata does not match its mode config: {mismatched}") - validate_dasc_model_structure(model, policy) + try: + validate_dasc_model_structure(model, policy) + except ApplyModeError as error: + warnings.warn( + f"{error}. The restored DASC policy is stale; re-run calibrate() before deployment", + stacklevel=2, + ) _attach_policy(model, policy) return model @@ -117,13 +123,14 @@ def replace_dasc_mode( ) -> nn.Module: """Replace existing DASC mode state in place with a newly derived policy.""" model = unwrap_model(model, force_unwrap=True) - policy = build_dasc_policy(model, config, measurements) manager = ModeloptStateManager(model) - manager.update_last_state_before_new_mode(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] = ( diff --git a/modelopt/torch/sparsity/state_sparsity/policy.py b/modelopt/torch/sparsity/state_sparsity/policy.py index 611b3618f0e..a5a9f6c5d6e 100644 --- a/modelopt/torch/sparsity/state_sparsity/policy.py +++ b/modelopt/torch/sparsity/state_sparsity/policy.py @@ -16,23 +16,42 @@ """Decay analysis and policy selection for GDN state sparsity.""" import hashlib +import importlib import json +import math from collections.abc import Iterable +from functools import lru_cache import torch import torch.nn.functional as F from torch import nn from modelopt.torch.opt.conversion import ApplyModeError +from modelopt.torch.opt.dynamic import DynamicModule from modelopt.torch.utils import unwrap_model from .config import DASCCalibrationMeasurement, DASCConfig, DASCLayerPolicy, DASCPolicy __all__ = ["analyze_gdn_decay", "compute_gdn_decay_horizons"] -_SUPPORTED_GDN_CLASS_NAMES = frozenset({"GatedDeltaNet", "Qwen3NextGatedDeltaNet"}) -# Covers ordinary FP16/BF16 storage casts; the exact retained-head mask below is the semantic gate. -_HORIZON_DTYPE_CAST_RTOL = 0.05 +_SUPPORTED_GDN_CLASS_PATHS = ( + ("megatron.core.ssm.gated_delta_net", "GatedDeltaNet"), + ("transformers.models.qwen3_next.modeling_qwen3_next", "Qwen3NextGatedDeltaNet"), +) + + +@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 (AttributeError, ImportError): + continue + if isinstance(candidate, type) and issubclass(candidate, nn.Module): + classes.append(candidate) + return tuple(classes) def compute_gdn_decay_horizons( @@ -64,8 +83,14 @@ def compute_gdn_decay_horizons( def _is_gdn_module(module: nn.Module) -> bool: """Accept supported GDN implementations and their ModelOpt dynamic subclasses.""" + supported_classes = _supported_gdn_classes() + module_class = type(module) + is_supported_class = module_class in supported_classes or ( + isinstance(module, DynamicModule) + and any(base in supported_classes for base in module_class.__mro__) + ) return ( - any(base.__name__ in _SUPPORTED_GDN_CLASS_NAMES for base in type(module).__mro__) + is_supported_class and isinstance(getattr(module, "A_log", None), torch.Tensor) and isinstance(getattr(module, "dt_bias", None), torch.Tensor) ) @@ -76,7 +101,7 @@ def _get_gdn_modules(model: nn.Module) -> dict[str, nn.Module]: model = unwrap_model(model, force_unwrap=True) modules = {name: module for name, module in model.named_modules() if _is_gdn_module(module)} if not modules: - supported = ", ".join(sorted(_SUPPORTED_GDN_CLASS_NAMES)) + supported = ", ".join(class_name for _, class_name in _SUPPORTED_GDN_CLASS_PATHS) raise ApplyModeError(f"DASC found no supported GDN modules; expected one of: {supported}") return dict(sorted(modules.items())) @@ -140,6 +165,33 @@ def _decay_parameters(modules: dict[str, nn.Module]) -> list[dict[str, object]]: ] +def _storage_rounding_radius(tensor: torch.Tensor) -> torch.Tensor: + """Bound one cast-to-storage rounding step around the represented tensor values.""" + values = tensor.detach().to(device="cpu", dtype=torch.float64) + dtype_info = torch.finfo(tensor.dtype) + unit_roundoff = dtype_info.eps / 2.0 + subnormal_slack = dtype_info.tiny * dtype_info.eps + return values.abs() * (unit_roundoff / (1.0 - unit_roundoff)) + subnormal_slack + + +def _storage_cast_horizon_bounds( + module: nn.Module, *, epsilon: float, static_gate_input: float +) -> tuple[torch.Tensor, torch.Tensor]: + """Bound horizons compatible with the current parameters before one storage cast.""" + a_log = module.A_log.detach().to(device="cpu", dtype=torch.float64) + dt_bias = module.dt_bias.detach().to(device="cpu", dtype=torch.float64) + a_radius = _storage_rounding_radius(module.A_log) + dt_radius = _storage_rounding_radius(module.dt_bias) + 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: @@ -286,19 +338,24 @@ def validate_dasc_decay_parameters(model: nn.Module, policy: DASCPolicy) -> None epsilon=policy.epsilon, static_gate_input=policy.static_gate_input, ) + modules = _get_gdn_modules(model) for name, values in current_horizons.items(): layer = policy.layers[name] - if not torch.allclose( - torch.tensor(values, dtype=torch.float64), - torch.tensor(layer.static_horizons, dtype=torch.float64), - rtol=_HORIZON_DTYPE_CAST_RTOL, - atol=0.0, - ): - raise ApplyModeError( - f"DASC policy horizons do not match current decay parameters in layer {name!r}" - ) retained = [head for head, horizon in enumerate(values) if horizon > policy.selected_wmax] if retained != layer.retained_heads: raise ApplyModeError( f"DASC policy head mask does not match current decay parameters in layer {name!r}" ) + lower, upper = _storage_cast_horizon_bounds( + modules[name], + epsilon=policy.epsilon, + static_gate_input=policy.static_gate_input, + ) + stored = torch.tensor(layer.static_horizons, dtype=torch.float64) + numerical_slack = 32.0 * torch.finfo(torch.float64).eps + if torch.any(stored < lower * (1.0 - numerical_slack)) or torch.any( + stored > upper * (1.0 + numerical_slack) + ): + raise ApplyModeError( + f"DASC policy horizons do not match current decay parameters in layer {name!r}" + ) diff --git a/tests/unit/torch/sparsity/state_sparsity/test_dasc.py b/tests/unit/torch/sparsity/state_sparsity/test_dasc.py index 07782a53dea..68dcaee6b4f 100644 --- a/tests/unit/torch/sparsity/state_sparsity/test_dasc.py +++ b/tests/unit/torch/sparsity/state_sparsity/test_dasc.py @@ -26,18 +26,22 @@ import modelopt.torch.opt as mto import modelopt.torch.sparsity.state_sparsity as mtss +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 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) @@ -56,6 +60,12 @@ 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 = { @@ -231,15 +241,49 @@ def __init__(self): 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="no supported GDN 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)]) + 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_generic_mode_application_reports_missing_measurements(): +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"), + ) + modules = { + "valid": type("ValidModule", (), {"GatedDeltaNet": GatedDeltaNet}), + "invalid": type("InvalidModule", (), {"NotAModule": object()}), + } + + def import_module(name): + if name == "missing": + raise ImportError(name) + return modules[name] + + monkeypatch.setattr(dasc_policy, "_SUPPORTED_GDN_CLASS_PATHS", module_paths) + monkeypatch.setattr(dasc_policy.importlib, "import_module", import_module) + + assert _resolve_supported_gdn_classes() == (GatedDeltaNet,) + + +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(), @@ -255,6 +299,23 @@ def test_generic_mode_application_reports_missing_measurements(): [_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.""" @@ -266,10 +327,10 @@ def test_public_exports_and_wrapped_model_export(): assert "mode" not in mtss.__all__ assert mtss.export_policy(nn.DataParallel(model)) == mtss.export_policy(model) - class _DynamicGatedDeltaNet(GatedDeltaNet): - """Stand in for the subclass synthesized by ModelOpt dynamic conversion.""" - - model.linear_attn = _DynamicGatedDeltaNet() + 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"): @@ -283,9 +344,16 @@ def test_dtype_cast_preserves_policy_when_the_selected_mask_is_unchanged(dtype): TinyGatedDeltaNetForCausalLM(), _config(wmax_candidates=[7]), [_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 @@ -309,12 +377,12 @@ def test_export_rejects_changed_decay_parameters_and_restore_rejects_structure() manager_state = ModeloptStateManager(model).state_dict() manager_state.append(copy.deepcopy(manager_state[0])) - with pytest.warns(UserWarning, match="saved DASC policy is stale"): - model = mtss.calibrate( - model, - _config(wmax_candidates=[7], model_revision="revision-2"), - [_candidate(7)], - ) + 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()) @@ -324,8 +392,12 @@ def test_export_rejects_changed_decay_parameters_and_restore_rejects_structure() restored.load_state_dict(model_state) assert mtss.export_policy(restored) == policy + with pytest.warns(UserWarning, match="restored DASC policy is stale"): + mismatched = mto.restore_from_modelopt_state( + TinyGatedDeltaNetForCausalLM(num_heads=3), state + ) with pytest.raises(ApplyModeError, match="module structure"): - mto.restore_from_modelopt_state(TinyGatedDeltaNetForCausalLM(num_heads=3), state) + mtss.export_policy(mismatched) tampered_state = copy.deepcopy(state) tampered_state["modelopt_state_dict"][0][1]["metadata"]["policy"]["quality_gates"][ @@ -348,16 +420,14 @@ def test_export_rejects_changed_decay_parameters_and_restore_rejects_structure() layer = tampered_state["modelopt_state_dict"][0][1]["metadata"]["policy"]["layers"][ "linear_attn" ] - layer["static_horizons"][0] = 1.0 - layer["retained_heads"] = [] - layer["omitted_heads"] = [0, 1] + layer["static_horizons"][0] *= 1.01 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 saving available after a calibrated GDN structure changes.""" + """Keep checkpoint save and restore available after a calibrated GDN structure changes.""" model = mtss.calibrate( TinyGatedDeltaNetForCausalLM(), _config(wmax_candidates=[7]), [_candidate(7)] ) @@ -371,6 +441,16 @@ def test_structure_staleness_does_not_block_checkpoint_save(): with pytest.raises(ApplyModeError, match="module structure"): mtss.export_policy(model) + with pytest.warns(UserWarning, match="saved DASC policy is stale"): + stale_state = mto.modelopt_state(model) + with pytest.warns(UserWarning, match="restored DASC policy is stale"): + restored = mto.restore_from_modelopt_state( + TinyGatedDeltaNetForCausalLM(num_heads=3), stale_state + ) + restored.load_state_dict(model.state_dict()) + with pytest.raises(ApplyModeError, match="module structure"): + mtss.export_policy(restored) + def test_export_rederives_the_selected_head_mask(): """Reject a self-consistent stored mask that current decay parameters do not derive."""