diff --git a/modelopt/torch/sparsity/state_sparsity/policy.py b/modelopt/torch/sparsity/state_sparsity/policy.py index c335410524d..098dbcb4a59 100644 --- a/modelopt/torch/sparsity/state_sparsity/policy.py +++ b/modelopt/torch/sparsity/state_sparsity/policy.py @@ -116,16 +116,8 @@ def _has_supported_gdn_identity( ) -def _get_gdn_modules(model: nn.Module) -> dict[str, nn.Module]: - """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) - ] +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 @@ -139,6 +131,13 @@ def _get_gdn_modules(model: nn.Module) -> dict[str, nn.Module]: "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 @@ -152,13 +151,26 @@ def _get_gdn_modules(model: nn.Module) -> dict[str, nn.Module]: f"{', '.join(unsupported_subclasses)}; convert the module with ModelOpt or use a " "supported class directly" ) - modules = dict(identity_modules) - if not modules: + + +def _get_gdn_modules(model: nn.Module) -> dict[str, nn.Module]: + """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(modules.items())) + return dict(sorted(identity_modules, key=lambda item: item[0])) def _analyze_gdn_modules( diff --git a/tests/unit/torch/sparsity/state_sparsity/test_dasc.py b/tests/unit/torch/sparsity/state_sparsity/test_dasc.py index 61c6e5cdf28..06379e856bf 100644 --- a/tests/unit/torch/sparsity/state_sparsity/test_dasc.py +++ b/tests/unit/torch/sparsity/state_sparsity/test_dasc.py @@ -269,9 +269,21 @@ class UnsupportedSubclass(GatedDeltaNet): partially_valid.good = GatedDeltaNet() partially_valid.bad = GatedDeltaNet() del partially_valid.bad.dt_bias - with pytest.raises(ApplyModeError, match="bad"): + 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"):