From 680cfb0fe945f7893f452d3d593b6745629abcc5 Mon Sep 17 00:00:00 2001
From: Kai Xu <kaix@nvidia.com>
Date: Thu, 10 Sep 2026 19:31:28 -0700
Subject: [PATCH] Reject incomplete DASC layer sets

Signed-off-by: Kai Xu <kaix@nvidia.com>
---
 .../torch/sparsity/state_sparsity/policy.py   | 69 +++++++++----------
 .../sparsity/state_sparsity/test_dasc.py      |  7 ++
 2 files changed, 39 insertions(+), 37 deletions(-)

diff --git a/modelopt/torch/sparsity/state_sparsity/policy.py b/modelopt/torch/sparsity/state_sparsity/policy.py
index d72ac620b5e..c335410524d 100644
--- a/modelopt/torch/sparsity/state_sparsity/policy.py
+++ b/modelopt/torch/sparsity/state_sparsity/policy.py
@@ -105,13 +105,6 @@ def compute_gdn_decay_horizons(
     return horizons
 
 
-def _is_gdn_module(module: nn.Module, supported_classes: tuple[type[nn.Module], ...]) -> bool:
-    """Accept supported GDN implementations and their ModelOpt dynamic subclasses."""
-    return _has_supported_gdn_identity(module, supported_classes) and all(
-        isinstance(getattr(module, name, None), torch.Tensor) for name in ("A_log", "dt_bias")
-    )
-
-
 def _has_supported_gdn_identity(
     module: nn.Module, supported_classes: tuple[type[nn.Module], ...]
 ) -> bool:
@@ -128,37 +121,39 @@ def _get_gdn_modules(model: nn.Module) -> dict[str, nn.Module]:
     model = unwrap_model(model, force_unwrap=True)
     supported_classes = _supported_gdn_classes()
     named_modules = list(model.named_modules())
-    modules = {
-        name: module for name, module in named_modules if _is_gdn_module(module, supported_classes)
-    }
+    identity_modules = [
+        (name, module)
+        for name, module in named_modules
+        if _has_supported_gdn_identity(module, supported_classes)
+    ]
+    missing_decay_parameters = [
+        name or "<root>"
+        for name, module in identity_modules
+        if not all(
+            isinstance(getattr(module, parameter, None), torch.Tensor)
+            for parameter in ("A_log", "dt_bias")
+        )
+    ]
+    if missing_decay_parameters:
+        raise ApplyModeError(
+            "DASC found supported GDN modules without A_log and dt_bias tensors at: "
+            f"{', '.join(missing_decay_parameters)}"
+        )
+    unsupported_subclasses = [
+        name or "<root>"
+        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"
+        )
+    modules = dict(identity_modules)
     if not modules:
-        missing_decay_parameters = [
-            name or "<root>"
-            for name, module in named_modules
-            if _has_supported_gdn_identity(module, supported_classes)
-            and not all(
-                isinstance(getattr(module, parameter, None), torch.Tensor)
-                for parameter in ("A_log", "dt_bias")
-            )
-        ]
-        if missing_decay_parameters:
-            raise ApplyModeError(
-                "DASC found supported GDN modules without A_log and dt_bias tensors at: "
-                f"{', '.join(missing_decay_parameters)}"
-            )
-        unsupported_subclasses = [
-            name or "<root>"
-            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"
-            )
         supported = ", ".join(
             f"{module_name}.{class_name}" for module_name, class_name in _SUPPORTED_GDN_CLASS_PATHS
         )
diff --git a/tests/unit/torch/sparsity/state_sparsity/test_dasc.py b/tests/unit/torch/sparsity/state_sparsity/test_dasc.py
index 18c3fe15796..61c6e5cdf28 100644
--- a/tests/unit/torch/sparsity/state_sparsity/test_dasc.py
+++ b/tests/unit/torch/sparsity/state_sparsity/test_dasc.py
@@ -265,6 +265,13 @@ class UnsupportedSubclass(GatedDeltaNet):
     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="bad"):
+        mtss.analyze_gdn_decay(partially_valid)
+
     invalid_decay = TinyGatedDeltaNetForCausalLM()
     invalid_decay.linear_attn.dt_bias = nn.Parameter(torch.zeros(3))
     with pytest.raises(ApplyModeError, match="Invalid GDN decay parameters"):