Conversation
Signed-off-by: Kai Xu <kaix@nvidia.com>
|
/claude review |
|
@coderabbitai review |
|
Caution Review failedThe pull request is closed. ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (5)
📝 WalkthroughWalkthroughThe DASC updates validate policy fields, restrict Megatron GDN analysis to single-process model parallelism, preserve unrelated mode state during recalibration, and consistently unwrap models during policy attachment and restoration. Documentation and unit tests cover the new behavior. ChangesDASC policy updates
Estimated code review effort: 3 (Moderate) | ~25 minutes Change: Bug fix ✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Comment |
|
|
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
| ] | ||
| _reject_incomplete_gdn_modules(identity_modules) | ||
| _reject_unconverted_gdn_subclasses(named_modules, supported_classes) | ||
| _reject_distributed_megatron_gdn(identity_modules, supported_classes) |
There was a problem hiding this comment.
[CRITICAL ModeState] Putting the distributed guard in _get_gdn_modules makes it fire on the restore and save paths, not just policy derivation — which turns an existing recoverable-staleness case into an unloadable checkpoint.
_get_gdn_modules has four callers: analyze_gdn_decay (L297), build_dasc_policy (L461), validate_dasc_model_structure (L521) and validate_dasc_decay_parameters (L543). The last two are what restore_dasc_model and update_dasc_metadata call, and both of those only tolerate _DASCRecoverableStalenessError. The new check raises a plain ApplyModeError, so it propagates.
Failure scenario: calibrate a single-process Megatron GDN model, mto.save(). Later mto.restore() the same checkpoint into a TP=2 (or PP=2) rematerialization of the model:
restore_dasc_model→validate_dasc_model_structure→_get_gdn_modules→ApplyModeError("DASC policy export supports Megatron GDN only with tensor_model_parallel_size=1 …"), uncaught → restore aborts.- Before this PR, the sharded
A_log.numel()/ missing PP layers produced_DASCModelStructureMismatchError, a subclass of_DASCRecoverableStalenessError, so restore warned "the restored DASC policy is stale; re-run calibrate()" and proceeded. DASC is inert at runtime (convert_dasc_model: "attach the selected DASC policy without changing execution"), so proceeding was safe and intentional. - The same applies to saving:
update_dasc_metadatacatches recoverable staleness fromvalidate_dasc_model_structurebut not this new error, somto.save()on a distributed model with DASC state now hard-fails.
The error text ("policy export supports …") is also misleading when it surfaces during restore.
Fix: call the guard only where a policy is actually derived, and leave the read-only validators alone:
def _get_gdn_modules(model: nn.Module, *, reject_distributed_megatron: bool = False) -> dict[str, nn.Module]:
...
_reject_incomplete_gdn_modules(identity_modules)
_reject_unconverted_gdn_subclasses(named_modules, supported_classes)
if reject_distributed_megatron:
_reject_distributed_megatron_gdn(identity_modules, supported_classes)with reject_distributed_megatron=True passed from analyze_gdn_decay and build_dasc_policy only. If you do want restore to notice the geometry, raise _DASCModelStructureMismatchError there instead so it stays on the documented recoverable-staleness path. Worth a test that restores a TP=1-calibrated policy onto a TP=2 model and asserts a warning rather than a raise.
| megatron_class = next( | ||
| ( | ||
| candidate | ||
| for candidate in supported_classes | ||
| if candidate.__module__ == module_name and candidate.__name__ == class_name | ||
| ), | ||
| None, | ||
| ) | ||
| if megatron_class is None: | ||
| return |
There was a problem hiding this comment.
[IMPORTANT Compatibility] This lookup identifies the Megatron class by string-matching candidate.__module__ / __name__, and silently disables the whole guard when it doesn't match (if megatron_class is None: return). That is a fail-open path in a check whose entire purpose is to fail closed.
_supported_gdn_classes() resolves the class with getattr(importlib.import_module("megatron.core.ssm.gated_delta_net"), "GatedDeltaNet"). __module__ is the class's defining module, not the import path — so if Megatron ever defines GatedDeltaNet in a private submodule and re-exports it from megatron.core.ssm.gated_delta_net (a routine refactor upstream, and Megatron pins here are loose), megatron_class becomes None, _reject_distributed_megatron_gdn returns without checking anything, and a TP=2 / PP=2 policy exports with rank-local head and layer indices — exactly the corruption this PR adds. Nothing fails loudly, and no test catches it: test_declared_gdn_paths_resolve_when_framework_is_installed only asserts the path resolves to an nn.Module subclass, and the two new Megatron tests use a fake whose __module__ is hand-set to match.
Two ways to remove the string coupling:
- Have the resolver return the declared path alongside the class, so the Megatron class is identified by which entry it was resolved from rather than by where it happens to be defined — e.g.
_supported_gdn_classes()returnstuple[tuple[tuple[str, str], type[nn.Module]], ...], and this function picks the entry whose key is_MEGATRON_GDN_CLASS_PATH. - Or classify by package over the MRO, which survives re-exports:
megatron_classes = tuple(
candidate
for candidate in supported_classes
if candidate.__module__.partition(".")[0] == "megatron"
)
if not megatron_classes:
returnEither way, please also make the resolution-check test assert the identity actually used here (cls.__module__ == _MEGATRON_GDN_CLASS_PATH[0]) so a silent upstream move breaks a test instead of the guard.
| widening cast contributes no additional slack. Decay tensors that are live in BF16 or FP16 are | ||
| still validated against that live dtype's rounding. | ||
| Choose ``static_gate_input`` as a conservative lower bound, such as a low percentile measured on the | ||
| calibration slices. Increasing it shortens the derived horizons and omits more heads; a value above |
There was a problem hiding this comment.
[SUGGESTION] "omits more heads" is ambiguous in a way that can invert the reader's understanding. Every head still appears in the policy with its static_horizons; what shrinks is retained_heads (retained = [head for head, horizon in enumerate(values) if horizon > selected_wmax], L482-483 of policy.py), i.e. more head state gets discarded. A reader who parses "omits heads" as "omits heads from sparsification" will conclude the opposite of the next clause.
The direction itself is right — larger static_gate_input → larger softplus(dt_bias + g) → more negative decay → shorter log(eps)/decay horizon → fewer heads clear horizon > Wmax. Only the noun is off:
| calibration slices. Increasing it shortens the derived horizons and omits more heads; a value above | |
| calibration slices. Increasing it shortens the derived horizons and retains fewer heads; a value above |
Same phrasing in config.py's static_gate_input description ("shorten horizons and omit more heads") is worth aligning.
There was a problem hiding this comment.
Claude review — 5 files changed, all reviewed (+171/-15)
Scope: full review (the trigger comment carried no extra scoping). Reviewed all 5 changed files: modelopt/torch/sparsity/state_sparsity/{policy,config,conversion}.py, docs/source/guides/6_sparsity.rst, tests/unit/torch/sparsity/state_sparsity/test_dasc.py.
Findings: CRITICAL: 1, IMPORTANT: 1, SUGGESTION: 1
Most impactful
-
[CRITICAL ModeState] The new distributed guard also fires on restore and save (policy.py:238).
_reject_distributed_megatron_gdnwas added inside_get_gdn_modules, which is called byvalidate_dasc_model_structureandvalidate_dasc_decay_parameters— the restore and metadata-refresh paths — not only byanalyze_gdn_decay/build_dasc_policy. Those paths only tolerate_DASCRecoverableStalenessError, so the plainApplyModeErrorpropagates: a TP=1-calibrated checkpoint restored into a TP=2 or PP=2 rematerialization now abortsmto.restore(), where it previously warned about a stale policy and proceeded (DASC is inert at runtime by design).mto.save()on such a model breaks the same way, and the message reads "policy export supports …" while surfacing during restore. Gating the call behind a keyword argument passed only from the two derivation entrypoints — or raising_DASCModelStructureMismatchErroron the validation path — keeps the export guard without regressing checkpoint loading. -
[IMPORTANT Compatibility] The guard can silently disable itself (policy.py:187-196). The Megatron class is located by comparing
candidate.__module__/__name__to the declared path, andmegatron_class is Nonereturns without checking anything.__module__is the defining module, so an upstream re-export makes this fail-closed check a no-op and a distributed policy exports rank-local indices with no error. No test covers that: the new Megatron tests use a fake whose__module__is hand-set to match, and the resolution test only asserts the path yields annn.Modulesubclass. Identify the class by the declared path it was resolved from, or classify over the MRO bymegatronpackage prefix.
Verified correct (no action needed)
- Removing
update_last_state_before_new_modefromreplace_dasc_modeis right: that hook exists to refresh the previous last mode before a new mode is stacked on, and recalibration replaces DASC state in place. DASC's own entry is fully rewritten, so nothing is left stale, and the new test pins the trailing mode's metadata rather than asserting the call. _attach_policyunwrapping withforce_unwrap=Trueis symmetric withget_attached_dasc_policyandreplace_dasc_mode, so wrapped-model attach/lookup round-trips; the newnn.DataParallelrestore test covers it.- New
DASCPolicyvalidators forepsilon/static_gate_inputreuse the same_validate_analysis_argumentsalready used byDASCConfigandanalyze_gdn_decay, so the schema boundary and the analysis boundary cannot disagree. No circular-import concern — the helper lives inconfig.py. - Gate values added to the docs snippet (0.995 / 0.98 / 0.2) match the actual
DASCConfigdefaults. static_gate_inputmonotonicity in the docs is correct: larger gate input → largersoftplus→ more negative decay → shorterlog(eps)/decayhorizon → fewer heads clearhorizon > Wmax. Only the wording is worth tightening (see the SUGGESTION on6_sparsity.rst).
Risk
Moderate, and concentrated in one place. Both blocking findings are about where the new fail-closed check lives rather than what it computes: as placed it is simultaneously too broad (breaks restore/save on distributed geometry) and too narrow (silently skippable via an upstream re-export). Everything else — the in-place recalibration fix, the wrapper-safe attachment, the policy-schema validators, and the docs — looks sound and is well covered by the new tests.
🤖 Generated with Claude Code
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## feature/dasc-state-sparsity #2408 +/- ##
===============================================================
- Coverage 78.78% 78.77% -0.01%
===============================================================
Files 548 549 +1
Lines 64243 66562 +2319
===============================================================
+ Hits 50615 52436 +1821
- Misses 13628 14126 +498
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:
|
1 participant
What does this PR do?
Type of change: Bug fix and documentation.
Closes the remaining review gaps on #2375. Distributed Megatron GDN policy export now fails closed for tensor or pipeline parallel configurations whose head and layer indices would be rank-local. The patch also validates standalone policy decay inputs, keeps policy attachment wrapper-safe, avoids refreshing unrelated mode metadata during recalibration, clarifies subclass remediation, and documents gate direction and active quality defaults.
Two review suggestions are intentionally not implemented: a head-count ratio is not a universal byte-savings bound for heterogeneous state geometry, and nonlinear measured quality is not guaranteed to be monotonic with DASC-NR window size.
Usage
No API changes. Megatron GDN calibration remains supported for TP=1 and PP=1 and fails closed otherwise.
Testing
Before your PR is Ready for review
Summary by CodeRabbit
New Features
Bug Fixes
Documentation