From 2da09a64650860f5d4cb9fff64746c0f173930a4 Mon Sep 17 00:00:00 2001 From: Kai Xu Date: Thu, 10 Sep 2026 19:40:49 -0700 Subject: [PATCH 1/2] Clarify DASC layer validation Signed-off-by: Kai Xu --- .../torch/sparsity/state_sparsity/policy.py | 38 ++++++++++++------- .../sparsity/state_sparsity/test_dasc.py | 14 ++++++- 2 files changed, 38 insertions(+), 14 deletions(-) diff --git a/modelopt/torch/sparsity/state_sparsity/policy.py b/modelopt/torch/sparsity/state_sparsity/policy.py index c335410524d..e7d84351b57 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)) 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"): From 1b2d49415bcdde3622b8ace1cf15ee7c830d586c Mon Sep 17 00:00:00 2001 From: kaix-nv Date: Thu, 10 Sep 2026 19:49:38 -0700 Subject: [PATCH 2/2] Sort DASC modules by explicit name (#2386) ## Summary Follow-up to #2385 addressing its sole defensive Claude suggestion: - sort discovered GDN modules with an explicit qualified-name key - prevent any hypothetical duplicate-name fallback from comparing nn.Module objects This PR is intentionally stacked on #2385 because repository rules protect a PR head branch after creation. ## 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 - full pre-commit on the touched file: passed - prior full-package coverage remains 408/408 statements; this change adds no statements and the line is exercised by the focused suite ## Summary by CodeRabbit * **Bug Fixes** * GDN modules are now returned in a consistent, predictable order. Signed-off-by: Kai Xu --- modelopt/torch/sparsity/state_sparsity/policy.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/modelopt/torch/sparsity/state_sparsity/policy.py b/modelopt/torch/sparsity/state_sparsity/policy.py index e7d84351b57..098dbcb4a59 100644 --- a/modelopt/torch/sparsity/state_sparsity/policy.py +++ b/modelopt/torch/sparsity/state_sparsity/policy.py @@ -170,7 +170,7 @@ def _get_gdn_modules(model: nn.Module) -> dict[str, nn.Module]: 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)) + return dict(sorted(identity_modules, key=lambda item: item[0])) def _analyze_gdn_modules(