Tighten DASC calibration contract - #2391
Conversation
Signed-off-by: kaix-nv <kaix@nvidia.com>
|
/claude 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 (3)
📝 WalkthroughWalkthroughThe change centralizes GDN decay tensor validation, adds a shared storage-dtype alias, removes redundant widening-cast bounds, updates documentation, and expands tests for rounding and non-finite parameters. ChangesGDN decay validation and dtype bounds
Estimated code review effort: 3 (Moderate) | ~20 minutes Change: Bug fix Merge Risk: ⚪ Minimal · up to The updated validation and dtype-contract paths have focused coverage, with no current merge-blocking issue identified. 🚥 Pre-merge checks | ✅ 6✅ Passed checks (6 passed)
✨ Finishing Touches 💡 1🛠️ Fix failing CI checks 💡
📝 Generate docstrings
🧪 Generate unit tests (beta)
Comment |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
| def _validated_gdn_decay_tensors( | ||
| a_log: torch.Tensor, dt_bias: torch.Tensor | ||
| ) -> tuple[torch.Tensor, torch.Tensor]: | ||
| """Validate decay tensors and return deterministic CPU float64 values.""" | ||
| _validate_gdn_decay_tensors(a_log, dt_bias) | ||
| return ( | ||
| a_log.detach().to(device="cpu", dtype=torch.float64), | ||
| dt_bias.detach().to(device="cpu", dtype=torch.float64), | ||
| ) | ||
|
|
There was a problem hiding this comment.
[SUGGESTION] The split leaves two private helpers whose names differ only by a trailing d on the verb — _validate_gdn_decay_tensors (returns None) and _validated_gdn_decay_tensors (returns CPU float64 tensors) — and the latter now has exactly one caller (compute_gdn_decay_horizons at line 124).
Why it matters: at a call site the two names are visually indistinguishable, so which one is invoked is no longer obvious from reading the surrounding code, and a future edit can silently pick the wrong one. Calling _validated_* where _validate_* was meant reintroduces exactly the wasted float64 materialization this PR set out to remove (it wouldn't fail — just quietly allocates), which is the failure mode most likely to survive review.
Suggestion: drop the single-use wrapper and do the cast inline in its one caller, keeping _validate_gdn_decay_tensors as the sole validation entry point:
def compute_gdn_decay_horizons(...):
"""Compute one static retention horizon per GDN head in CPU float64."""
if not 0.0 < epsilon < 1.0:
raise ValueError("epsilon must be in (0, 1)")
_validate_gdn_decay_tensors(a_log, dt_bias)
a_log_cpu = a_log.detach().to(device="cpu", dtype=torch.float64)
dt_bias_cpu = dt_bias.detach().to(device="cpu", dtype=torch.float64)If you'd rather keep the wrapper, renaming it to say what it returns (e.g. _gdn_decay_tensors_as_cpu_float64) gets the same disambiguation.
There was a problem hiding this comment.
Addressed in #2392. The single-use converting wrapper is removed; compute_gdn_decay_horizons now validates once and performs the CPU float64 casts inline.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Claude review passed — no blocking issues found. LGTM
Scope
Full review of all 4 changed files (+53/-31): modelopt/torch/sparsity/state_sparsity/config.py and policy.py, docs/source/guides/6_sparsity.rst, and tests/unit/torch/sparsity/state_sparsity/test_dasc.py. Traced every caller of the split validation helpers and every reference to the new type alias across modelopt/ and tests/.
Findings — CRITICAL: 0 | IMPORTANT: 0 | SUGGESTION: 1
The one suggestion is a naming/shape point on the _validate_gdn_decay_tensors / _validated_gdn_decay_tensors pair — see the inline comment. Non-blocking.
What I verified
- Validation split is behavior-preserving. Moving the
isfinitecheck off the CPU float64 copies onto the live tensors is equivalent: widening any float dtype to float64 is exact, so it can neither create nor destroy a non-finite value. Call sites are exhaustively covered —_analyze_gdn_modules(policy.py:215),validate_dasc_decay_parameters(policy.py:483), andcompute_gdn_decay_horizonsvia the wrapper (policy.py:124) — no path lost a check. - No storage-dtype validation gap after the cast. In
_analyze_gdn_modulesthe pre-cast call is now validation-only, but the tensors are re-validated post-cast insidecompute_gdn_decay_horizons, so an FP16 downcast that overflowsA_logtoinfstill fails closed rather than silently yielding a zero horizon. The pre-cast call is not redundant either: it is what rejects a non-floating-pointA_log, which the post-cast check would let through. - The new NaN test's claim holds. With a NaN in
dt_bias, both bounds from_storage_cast_horizon_boundsgo NaN, and every comparison invalidate_dasc_decay_parameters(head_upper <= wmax,head_lower > wmax, and both stored-horizon bound checks) evaluatesFalse— so absent the finiteness check the head really would be silently accepted.test_non_finite_decay_parameters_are_rejected_on_exportnames and isolates that regression accurately, and thematchstring still matches through theApplyModeErrorwrapper. - No config or checkpoint compatibility impact.
_DecayParameterStorageDtypehas identical members to the inlineLiteralit replaces on bothDASCConfigandDASCPolicy, so the Pydantic schema andformat_version: Literal[1]are unchanged and existing serialized policies round-trip.analyze_gdn_decay's public signature is likewise unchanged semantically — the alias is a plainLiteralbinding, soautodoc_typehints = "description"still resolves it to the three concrete options in the rendered docs.Literalremains correctly imported in config.py, where it is still used. - The docs correction is accurate. With storage
float32and a BF16/FP16 live dtype,_storage_rounding_radiusdedups viadict.fromkeysand composes both casts, so "the default float32 storage dtype contributes only a negligible float32 step beyond it" is right — and strictly more accurate than the removed "adds no storage slack of its own," which understated the same-dtype case. - Forwarding
epsilon/static_gate_inputin the docs recipe is the right fix, not an API change.build_dasc_policyrecomputes horizons fromconfig.epsilon/config.static_gate_input, and_validate_measurement_geometrybinds the caller-reportedretained_heads/total_headsto them, so a caller who analyzed under different values already fails closed. The gap was documentation-only.
Risk: low. Docs, a type-alias consolidation, a behavior-preserving helper split, and a test reorganization; no algorithmic, mode-registration, state-schema, or export-path change.
Note: I reviewed statically — running pytest tests/unit/torch/sparsity/state_sparsity/test_dasc.py was not permitted in this environment, so I did not independently reproduce the 25-passed/1-skipped result reported in the PR description.
🤖 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-review-validation #2391 +/- ##
==================================================================================
+ Coverage 68.74% 78.78% +10.03%
==================================================================================
Files 548 548
Lines 64220 64227 +7
==================================================================================
+ Hits 44151 50599 +6448
+ Misses 20069 13628 -6441
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:
|
Addresses the final numerical-bound and cleanup review findings. Changes: - exclude exact widening dtypes from composed inverse rounding bounds - retain both bounds only when neither dtype exactly contains the other, such as FP16 versus BF16 - add direct FP16/BF16-with-FP32 widening regressions - remove the ambiguous single-use validated/converting helper and cast inline after validation - align the guide with the tightened exact-widening contract Validation: - focused DASC tests: 27 passed, 1 skipped (optional Megatron dependency) - pre-commit on all touched files: passed Commit is ED25519-signed and carries a matching Signed-off-by trailer. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit - **Bug Fixes** - Improved decay-policy validation for sequential rounding operations. - Prevented redundant precision conversions from adding incorrect rounding allowances. - Preserved accurate validation for BF16 and FP16 decay tensors. - **Tests** - Added coverage confirming consistent rounding-radius calculations across FP16, BF16, and FP32 comparisons. - **Documentation** - Updated sparsity guidance to clarify decay-policy rounding behavior. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: kaix-nv <kaix@nvidia.com>
02d8cfe
into
feature/dasc-state-sparsity-review-validation
1 participant
Closes the remaining review-contract and cleanup findings from #2389 and #2390.
Changes:
Validation:
Commit is ED25519-signed and carries a matching Signed-off-by trailer.
Summary by CodeRabbit
Bug Fixes
Documentation
Refactor