Skip to content

Reject incomplete DASC layer sets #2384

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
kaix-nv wants to merge 1 commit into from
Closed

Reject incomplete DASC layer sets #2384

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
69 changes: 32 additions & 37 deletions modelopt/torch/sparsity/state_sparsity/policy.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand All @@ -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"
Comment on lines +142 to +153

@claude claude Bot Sep 11, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[SUGGESTION] Hoisting this block out of the if not modules: branch is a second, unannounced behavior change, and it's the one half of the diff with no test.

Previously the unconverted-subclass check only ran when zero GDN layers resolved; now it runs on every _get_gdn_modules() call. So a mixed model — some direct Qwen3NextGatedDeltaNet instances plus some non-dynamic subclass instances — used to succeed while silently dropping the subclass layers, and now hard-fails. That's the same class of bug as the A_log/dt_bias case and the change looks right, but the PR description only lists the decay-parameter fix, and test_calibration_fails_closed_on_measurements_and_model_mismatch only covers a model that is entirely UnsupportedSubclass (line 253) — a case the old code already rejected via the empty-modules path.

The 100% statement coverage cited in the PR description is therefore satisfied by the pre-existing test; the newly reachable path is uncovered. Worth a mixed-model case mirroring the partially_valid one:

mixed_subclass = nn.Module()
mixed_subclass.good = GatedDeltaNet()
mixed_subclass.stale = UnsupportedSubclass()
with pytest.raises(ApplyModeError, match=r"not ModelOpt dynamic modules at: stale$"):
    mtss.analyze_gdn_decay(mixed_subclass)

and a line in the PR description noting the broadened rejection.

@kaix-nv kaix-nv Sep 11, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Addressed in #2385. A mixed model containing good plus an ordinary stale subclass now asserts the exact unconverted-subclass diagnostic. The broadened fail-closed behavior is also called out in #2384's description.

)
modules = dict(identity_modules)

@claude claude Bot Sep 11, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[SUGGESTION] modules is now a pure alias — this dict is built only to be .items()'d back apart in the return dict(sorted(modules.items())) two lines below. Dropping the intermediate covers both:

    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))

More broadly (non-blocking, fine to defer): this refactor leaves _get_gdn_modules at ~40 lines doing resolution plus two independent rejection passes, with three inline comprehensions that each take a couple of reads to parse. Per CONTRIBUTING's "put high-level behavior first, hide lower-level details behind well-named helpers", the top level reads better as intent:

    identity_modules = [...]
    _reject_incomplete_gdn_modules(identity_modules)
    _reject_unconverted_gdn_subclasses(named_modules, supported_classes)

with the comprehensions and their ApplyModeError messages inside those helpers. That also gives the newly-unconditional subclass check a name explaining why it now runs on every call.

@kaix-nv kaix-nv Sep 11, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Addressed in #2385. The redundant modules alias is removed, and the incomplete-identity and unconverted-subclass passes are now named helpers so _get_gdn_modules reads as the intended high-level validation sequence.

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
)
Expand Down
7 changes: 7 additions & 0 deletions tests/unit/torch/sparsity/state_sparsity/test_dasc.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"):

@claude claude Bot Sep 11, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[SUGGESTION] This is the PR's headline regression test, but match="bad" is loose enough that it doesn't pin the behavior being fixed.

The point of the change is that a malformed layer is reported rather than silently dropped, and that the well-formed sibling is not falsely blamed. re.search("bad", ...) passes for any ApplyModeError whose text happens to contain that substring, and it would still pass if the code regressed to listing every identity module (good, bad) instead of only the incomplete one.

Anchoring on the full message makes both properties explicit:

Suggested change
with pytest.raises(ApplyModeError, match="bad"):
with pytest.raises(ApplyModeError, match=r"without A_log and dt_bias tensors at: bad$"):

@kaix-nv kaix-nv Sep 11, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Addressed in #2385. The regression now anchors the full diagnostic suffix and therefore proves only bad, not its valid sibling, is reported.

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"):
Expand Down
Loading