Reject incomplete DASC layer sets - #2384
Conversation
Signed-off-by: Kai Xu <kaix@nvidia.com>
|
/claude review |
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (2)
Included review availability: Your plan provides up to 12 included reviews per hour; 8 remain after this review. 📝 WalkthroughWalkthroughThe change separates GDN identity discovery from decay-parameter validation. It reports missing ChangesGDN discovery validation
Priority: ⬇️ Low Estimated code review effort: 2 (Simple) | ~10 minutes Merge Risk: ⚪ Minimal · up to The updated GDN validation rejects incomplete supported layers and regression coverage exercises the mixed valid and malformed-layer path. No merge-blocking risk remains. 🚥 Pre-merge checks | ✅ 6✅ Passed checks (6 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
| partially_valid.good = GatedDeltaNet() | ||
| partially_valid.bad = GatedDeltaNet() | ||
| del partially_valid.bad.dt_bias | ||
| with pytest.raises(ApplyModeError, match="bad"): |
There was a problem hiding this comment.
[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:
| with pytest.raises(ApplyModeError, match="bad"): | |
| with pytest.raises(ApplyModeError, match=r"without A_log and dt_bias tensors at: bad$"): |
There was a problem hiding this comment.
Addressed in #2385. The regression now anchors the full diagnostic suffix and therefore proves only bad, not its valid sibling, is reported.
| 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" |
There was a problem hiding this comment.
[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.
| f"{', '.join(unsupported_subclasses)}; convert the module with ModelOpt or use a " | ||
| "supported class directly" | ||
| ) | ||
| modules = dict(identity_modules) |
There was a problem hiding this comment.
[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.
There was a problem hiding this comment.
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.
There was a problem hiding this comment.
Claude review passed — no blocking issues found. LGTM
Findings: CRITICAL: 0, IMPORTANT: 0, SUGGESTION: 3
Scope: full review — 2 changed files (+39/-37), both opened (modelopt/torch/sparsity/state_sparsity/policy.py, tests/unit/torch/sparsity/state_sparsity/test_dasc.py).
Correctness assessment
I traced _get_gdn_modules against all four call sites (analyze_gdn_decay, calibrate, validate_dasc_model_structure, validate_dasc_decay_parameters) and enumerated the resolution cases. The refactor is semantically sound:
- Accepted set is unchanged.
modules = dict(identity_modules)drops the oldA_log/dt_biaspredicate, but themissing_decay_parametersguard above it raises before that point, so every surviving entry still satisfies the tensor contract._analyze_gdn_modulesand_model_structurealready dereferencemodule.A_logunconditionally on the resolved set, so the new check strictly upgrades a would-beAttributeErrorinto a cleanApplyModeError— no path loses a layer it previously kept. - The two rejection lists are disjoint.
identity_modulesrequirestype in supported_classesorDynamicModule;unsupported_subclassesrequires the negation of both. No module can be simultaneously accepted and rejected, andnamed_modules()guarantees unique keys sodict(identity_modules)cannot collide. - No false failures on ModelOpt-converted models.
DynamicModule.convertsynthesizes a class deriving from(DynamicModule, cls), so converted GDN layers areDynamicModuleinstances and are excluded from the subclass check while still matching_has_supported_gdn_identityvia MRO. Layer names are unaffected by conversion, somodel_structure_sha256still round-trips pre/post-convert. - Fail-closed direction is right. For a model whose GDN layers are uniformly a non-dynamic subclass, old and new code both raise; the only newly-failing case is the mixed one, which is precisely the incomplete layer set the PR title targets.
_is_gdn_moduleremoval is clean — no remaining references anywhere in the tree.
No mode-registration, config-schema, or public-API surface changed (__all__ untouched), and the modelopt_state contract is unaffected, so there is no restore-migration concern. Backward compatibility is a non-issue here: the base is feature/dasc-state-sparsity-storage-contract, so any checkpoint written under the permissive behavior is within the same unreleased stack.
Most impactful finding
The one thing worth acting on before merge is coverage, not correctness: hoisting the unsupported_subclasses check out of the if not modules: branch is a second behavior change that the PR description does not mention and no test exercises. The existing UnsupportedSubclass case is an all-subclass model, which the old code already rejected via the empty-modules path — so the 100% statement coverage cited in the description is satisfied by the pre-existing test while the newly reachable mixed-model path stays untested. A mixed case mirroring the new partially_valid block closes it. The other two suggestions are a loose match="bad" assertion and shedding the now-redundant modules alias / splitting the 40-line resolver.
Risk: Low. Small, well-targeted change that tightens validation in the fail-closed direction with no widening of accepted inputs.
Note on verification: findings are from static analysis and cross-file tracing — sandbox approval for pytest was declined, so I did not independently re-run the suite. The 23-passed / 134-passed figures are the author's reported numbers, unverified by me.
🤖 Generated with Claude Code
|
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## feature/dasc-state-sparsity-storage-contract #2384 +/- ##
================================================================================
- Coverage 78.76% 78.76% -0.01%
================================================================================
Files 548 548
Lines 64190 64189 -1
================================================================================
- Hits 50562 50561 -1
Misses 13628 13628
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
## Summary Consolidates the complete reviewed fix stack for #2375 into one DCO-safe commit: - harden package exports, measurement semantics, wrapper handling, and actionable calibration errors - validate exact installed GDN identities plus ModelOpt dynamic subclasses; reject lookalikes, ordinary subclasses, incomplete layers, and partial layer sets - make stale checkpoints saveable and restorable while keeping deployment export strict - make DASC recalibration replace and deduplicate existing mode state without stale-metadata refresh - record the declared decay-parameter checkpoint storage dtype and use derived FP16/BF16/FP32 rounding bounds - preserve BF16/FP16 storage and wider/cross-dtype reload compatibility without globally widening FP32 tolerance - add installed Transformers path coverage, optional Megatron gating, lifecycle, tamper, lossy-cast, and mixed-layer regressions - document the explicit storage-dtype contract This consolidated PR supersedes the mechanically stacked review-fix PRs #2377, #2378, #2379, #2380, #2382, #2383, #2384, and #2385. Its tree is byte-identical to the independently reviewed leaf commit from #2386. ## 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 - DASC package coverage: 408/408 statements, 100% - full pre-commit on all touched files: passed - real Transformers Qwen3NextGatedDeltaNet BF16 storage to FP32 reload smoke: passed - commit author and Signed-off-by identity both use kaix-nv <kaix@nvidia.com> <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **New Features** * Added support for configuring decay-parameter storage precision with FP16, BF16, or FP32. * Added safer recalibration that replaces existing DASC state. * Expanded compatibility with supported GDN adapter classes and model wrappers. * Added improved validation for sparsity policies, measurements, model structure, and decay parameters. * Added support for perplexity-retention values above 1. * **Documentation** * Clarified evaluation responsibilities, recalibration behavior, stale-policy handling, supported adapters, and dtype requirements. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: kaix-nv <kaix@nvidia.com>
1 participant
Summary
Follow-up to #2383 addressing its final Claude suggestion:
This PR is intentionally stacked on #2383 because repository rules protect a PR head branch after creation.
Validation
Summary by CodeRabbit
The fail-closed coverage applies symmetrically to mixed models containing a valid GDN layer plus either a malformed supported identity or an unconverted ordinary GDN subclass; neither layer may be silently omitted.