Handle DASC storage round trips and imports - #2382
Conversation
Signed-off-by: Kai Xu <kaix@nvidia.com>
|
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 policy now resolves optional GDN classes with warnings, reuses resolved modules during analysis and policy construction, rejects unsupported subclasses, and calculates storage rounding bounds across parameter, FP16, and BF16 dtypes. ChangesGDN policy updates
Priority: ⬇️ Low Estimated code review effort: 3 (Moderate) | ~20 minutes Merge Risk: ⚪ Minimal · up to The updated policy behavior is covered by focused compatibility and regression tests, with no actionable merge risk identified. 🚥 Pre-merge checks | ✅ 6✅ Passed checks (6 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Comment |
|
/claude review |
![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}
| values = tensor.detach().to(device="cpu", dtype=torch.float64).abs() | ||
| radius = torch.zeros_like(values) | ||
| for dtype in (tensor.dtype, *_SUPPORTED_STORAGE_DTYPES): | ||
| dtype_info = torch.finfo(dtype) | ||
| unit_roundoff = dtype_info.eps / 2.0 | ||
| candidate = ( | ||
| values * (unit_roundoff / (1.0 - unit_roundoff)) + dtype_info.tiny * dtype_info.eps | ||
| ) | ||
| radius = torch.maximum(radius, candidate) | ||
| return radius |
There was a problem hiding this comment.
[SUGGESTION] Taking the max over (tensor.dtype, fp16, bf16) unconditionally applies the coarsest storage tolerance to every model, including ones that were never round-tripped through a narrow dtype. For an FP32 A_log/dt_bias, the relative radius goes from fp32_eps/2 ≈ 6e-8 to bf16_eps/2 ≈ 3.9e-3 — roughly a 65,000× widening of the acceptance band in _storage_cast_horizon_bounds.
Why it matters: the resulting [lower, upper] window on policy.static_horizons is the only check on the horizon values themselves (the head mask is checked exactly at line 396, and DASCPolicy cross-validates mask ↔ horizons at config.py:236-243, so behavior-changing tampering is still caught). But provenance-level drift is now tolerated up to ~0.5% of each horizon instead of ~1e-7. Concretely, tests/.../test_dasc.py:472 tampers a horizon by *= 1.01 and expects rejection — that test now clears the bound by only ~2×, where before the margin was ~5 orders of magnitude. It will silently stop being a meaningful tamper test if the fixture's A_log/dt_bias magnitudes ever change.
Suggested fix: record the storage dtype used at calibration in DASCPolicy (a new field with a permissive default keeps old JSON policies loadable) and apply only that dtype's radius plus the live tensor's own, instead of the union of all candidates.
Second, narrower point: the max-over-dtypes bounds a single cast. A checkpoint saved in FP16 and reloaded as BF16 incurs two lossy roundings (FP16 ULP/2 + BF16 ULP/2 ≈ 1.13× the BF16-only radius), so that path can be spuriously rejected — numerical_slack at line 406 (32·fp64_eps) is far too small to absorb it. The reverse direction (BF16 → FP16) is safe because a BF16 value is exactly representable in FP16. If you want to cover both, sum the current dtype's radius and the storage dtype's rather than taking the max.
There was a problem hiding this comment.
Addressed in #2383. DASCConfig and DASCPolicy now record decay_parameter_storage_dtype, defaulting to strict float32 for backward-compatible parsing. Bounds compose only that declared storage cast and any distinct live-dtype materialization cast. Provenance hashing uses the same declared dtype. Tests cover strict 1e-4 horizon tampering, BF16 to FP32 reload, FP16 to BF16, BF16 to FP16, and legacy metadata.
| if not tensor.dtype.is_floating_point: | ||
| raise ApplyModeError("DASC decay parameters must use a floating-point dtype") |
There was a problem hiding this comment.
[SUGGESTION] The new floating-point guard sits too late in the validation order to be the error the user actually sees. In validate_dasc_decay_parameters, _analyze_gdn_modules runs first (line 388) and happily casts an integer A_log to float64, then the exact mask comparison at line 396 runs, and only after that does _storage_cast_horizon_bounds reach this check. For any integer-dtype parameter whose truncation shifts a horizon across selected_wmax, the user gets "DASC policy head mask does not match current decay parameters" — which points at the wrong cause.
tests/.../test_dasc.py (the new test_bf16_storage_round_trip_loaded_in_fp32_preserves_policy, pytest.raises(..., match="floating-point dtype")) only passes because that fixture's retained-head set happens to survive .to(torch.int64) truncation; it's order-dependent rather than testing the guard directly.
Suggest moving the dtype check up to where the parameters are first read — _analyze_gdn_modules or compute_gdn_decay_horizons — so it fires before any horizon math, and asserting it in the test against that entry point.
There was a problem hiding this comment.
Addressed in #2383. Floating-point dtype validation now occurs in compute_gdn_decay_horizons before conversion or mask derivation, so integer decay tensors produce the direct dtype error independent of the resulting mask.
| if not modules: | ||
| supported = ", ".join(class_name for _, class_name in _SUPPORTED_GDN_CLASS_PATHS) | ||
| 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)}; use an exact supported class" | ||
| ) | ||
| 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}") |
There was a problem hiding this comment.
[SUGGESTION] Two things about the new diagnostic split:
-
The remediation text
"use an exact supported class"is misleading —_is_gdn_module(line 105-108) also accepts anyDynamicModulewhose MRO contains a supported class, so an exact class is not the only valid option. Something like "convert the module with ModelOpt or use a supported class directly" matches the actual acceptance rule. -
There's a third failure mode that still lands in the generic "no supported GDN modules" branch: a module that is an exact supported class (or a proper dynamic subclass) but is missing the
A_log/dt_biastensors that_is_gdn_modulealso requires. Theunsupported_subclassesfilter excludes it (type(module) not in supported_classesis False), so the user is told no GDN module was found when one clearly was, just without the expected attributes. Since you're already paying for thenamed_modules()scan on this error path, a second bucket for "class matched but decay parameters are missing" would make that case self-explanatory.
There was a problem hiding this comment.
Addressed in #2383. The subclass message now says to convert with ModelOpt or use a supported class directly. A separate diagnostic detects an exact or ModelOpt-dynamic supported identity missing A_log or dt_bias tensors, with regression coverage.
| root_module = module_name.partition(".")[0] | ||
| if dasc_policy.importlib.util.find_spec(root_module) is None: | ||
| pytest.skip(f"optional framework {root_module!r} is not installed") | ||
| module = dasc_policy.importlib.import_module(module_name) | ||
| assert issubclass(getattr(module, class_name), nn.Module) |
There was a problem hiding this comment.
[SUGGESTION] This is the first test under tests/unit/ that performs a real import megatron.core...; every other unit test either monkeypatches Megatron or avoids it (tests/unit/torch/export/test_get_quantization.py patches the exporter). The guard is find_spec(root_module) is None, which only proves the distribution is present, not importable — on a runner with megatron-core installed but without transformer_engine/CUDA, import_module("megatron.core.ssm.gated_delta_net") raises and this becomes a hard unit-test failure rather than a skip. That's precisely the "installed but unresolvable" case the production code at policy.py:52-64 now degrades to a warning for.
pytest.importorskip(module_name) would fix the CI fragility but defeats the test's purpose — path drift would silently skip instead of failing. Gating on the root package instead keeps both properties:
@pytest.mark.parametrize(("module_name", "class_name"), dasc_policy._SUPPORTED_GDN_CLASS_PATHS)
def test_declared_gdn_paths_resolve_when_framework_is_installed(module_name, class_name):
"""Guard supported identities against upstream dependency path drift."""
pytest.importorskip(module_name.partition(".")[0])
module = dasc_policy.importlib.import_module(module_name)
assert issubclass(getattr(module, class_name), nn.Module)Root unimportable (missing TE, no CUDA) → skip; root importable but the leaf path moved → the test fails, which is the drift you want to catch. It also drops the reach through dasc_policy.importlib.util.
There was a problem hiding this comment.
Addressed in #2383. The real-path guard now uses pytest.importorskip on the framework root. An unimportable optional root skips, while an importable root with a moved leaf path or class still fails.
There was a problem hiding this comment.
Claude review passed — no blocking issues found. LGTM
Summary
CRITICAL: 0 · IMPORTANT: 0 · SUGGESTION: 5
Scope: full review of both changed files (modelopt/torch/sparsity/state_sparsity/policy.py, tests/unit/torch/sparsity/state_sparsity/test_dasc.py; 127 additions / 27 deletions). Reviewed the whole diff plus surrounding context in policy.py and state_sparsity/config.py.
What I verified
- Storage-rounding math is sound.
_storage_rounding_radiuscorrectly uses theu/(1-u)form, which is the right bound when the radius is expressed relative to the rounded value rather than the original — that is exactly what the FP32-reload scenario needs. The near-zero slack (tiny * eps) is 2× the true half-subnormal-spacing bound, i.e. conservative in the safe direction. The union over dtypes picks up the BF16 relative term for large values and the FP16 subnormal term near zero, which is the correct combination. - Monotonicity in
_storage_cast_horizon_boundsis correct.horizon = -log(eps)/(exp(A_log)·softplus(dt_bias+g)); perturbing both terms upward maximises the denominator and therefore yields the lower horizon bound, and vice versa. No sign error. - The BF16 → FP16 reload path is actually safe, contrary to what a quick reading suggests: BF16 carries 7 explicit mantissa bits versus FP16's 10, so a BF16 value is exactly representable in FP16 and the second cast contributes no error.
- The
_get_gdn_modules/_analyze_gdn_modulesrefactor is a genuine fix.build_dasc_policyandvalidate_dasc_decay_parameterspreviously resolved GDN modules twice (once directly, once insideanalyze_gdn_decay); they now resolve once and share the result. Publicanalyze_gdn_decaykeeps its signature and behavior, and_is_gdn_module's new required parameter is private with only one call site in the repo — no backward-compat concern. - No mode/state or export surface is touched.
modelopt_stateschema, mode registration, andDASCPolicy's serialized fields are unchanged, so existing checkpoints and exported policies still round-trip. - No
-W errorfilter is configured repo-wide, so the newwarnings.warncalls in the resolver cannot fail unrelated CI;lru_cachealso bounds them to one emission per process. - Test-isolation is handled correctly —
_resolve_supported_gdn_classes = dasc_policy._supported_gdn_classes.__wrapped__(line 35) bypasses the cache, so the monkeypatched-paths test does not poison other tests.
Most impactful findings
- Storage tolerance widens ~65,000× for FP32 models (
policy.py:216-225). Taking the max over(tensor.dtype, fp16, bf16)applies the coarsest tolerance unconditionally, including to models never round-tripped through a narrow dtype. Runtime behavior is still protected — the head mask is compared exactly andDASCPolicycross-validates mask ↔ horizons — so this is provenance fidelity, not correctness. Buttest_dasc.py:472(static_horizons[0] *= 1.01) now clears the bound by only ~2× where it previously had ~5 orders of magnitude, so it will quietly stop being a meaningful tamper test if the fixture's parameter magnitudes change. Recording the calibration storage dtype inDASCPolicy(new field with a permissive default, keeping old policies loadable) would let validation apply only the tolerance the checkpoint actually needs. - New unit test performs a real Megatron import (
test_dasc.py:297-301).find_speconly proves the distribution is present, not importable — on a runner withmegatron-corebut notransformer_engine/CUDA this becomes a hard failure instead of a skip. Gating onpytest.importorskip(root_module)and then importing the leaf path preserves the drift-detection intent while degrading gracefully, mirroring howpolicy.py:52-64now handles the same situation. importlib.utilused without an explicitimport importlib.util(policy.py:54), working only becausemodelopt/torch/__init__.py:19imports it first. Additionally,find_speccan itself raise, and doing so from inside thisexcepthandler escapes_supported_gdn_classes()— turning the graceful-degradation path into a hard failure.- Diagnostic-message polish in
_get_gdn_modulesand dtype-guard placement in the validation order — details inline.
Risk assessment
Low. Tightly scoped to the DASC policy module, no public API or serialized-schema change, and the refactor is behavior-preserving. Every finding is non-blocking; the FP32 tolerance widening is the one worth a deliberate decision before merge rather than an automatic accept, since the PR description frames it purely as a compatibility fix without noting the strictness cost.
🤖 Generated with Claude Code
| except (AttributeError, ImportError): | ||
| except ModuleNotFoundError as error: | ||
| root_module = module_name.partition(".")[0] | ||
| if importlib.util.find_spec(root_module) is None: |
There was a problem hiding this comment.
[SUGGESTION] importlib.util is used here but only import importlib is at the top of the file (line 19). Accessing a submodule attribute off a package works only if something else already imported it — here it happens to work because modelopt/torch/__init__.py:19 does import importlib.util, and importing this module always executes that first. That's an invisible dependency: if that line in modelopt/torch/__init__.py is ever dropped, this raises AttributeError on the first ModuleNotFoundError path. modelopt/torch/speculative/utils.py:20 and modelopt/onnx/quantization/autotune/benchmark.py:30 use the explicit form.
Also worth hardening: find_spec itself can raise (ValueError when a name is in sys.modules with __spec__ is None, ImportError while resolving a parent). Raising from inside this except handler escapes _supported_gdn_classes() entirely, which contradicts the docstring's promise of resolving "without making either framework mandatory" — the graceful-degradation path becomes a hard failure. Wrapping the probe in try/except Exception and treating a failure as "not installed" keeps that contract.
Line 19 fix:
| if importlib.util.find_spec(root_module) is None: | |
| import importlib.util |
(applied to the import line, not this one)
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## feature/dasc-state-sparsity-identity-restore #2382 +/- ##
=============================================================================
Coverage 78.76% 78.76%
=============================================================================
Files 548 548
Lines 64163 64184 +21
=============================================================================
+ Hits 50535 50556 +21
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 #2380 that addresses its complete automated review pass:
This PR is intentionally stacked on #2380 because repository rules protect a PR head branch after creation.
Validation
Summary by CodeRabbit
Bug Fixes
Tests