Make DASC scalar analysis device independent - #2398
Conversation
Signed-off-by: kaix-nv <kaix@nvidia.com>
|
/claude review |
|
@coderabbitai full review |
|
Warning Review limit reachedNext included review available in 57 seconds. View limit detailsLimit details: You’ve used all 12 included reviews currently available. Enable usage-based reviews in Billing to review now. Otherwise, wait until the next included review is available. Review configuration: ⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (2)
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 (3)
Included review availability: Your plan provides up to 12 included reviews per hour; 0 remain after this review. 📝 WalkthroughWalkthroughThe change tightens boolean argument validation and narrows exception handling. Decay-horizon computation now uses ChangesState sparsity behavior
Priority: ⬇️ Low Estimated code review effort: 2 (Simple) | ~10 minutes Change: Bug fix Merge Risk: ⚪ Minimal · up to The PR adds the intended boolean validation and CPU-safe scalar horizon computation without an identified merge-blocking risk. 🚥 Pre-merge checks | ✅ 6✅ Passed checks (6 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Comment |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
There was a problem hiding this comment.
Claude review — PR #2398 (Make DASC scalar analysis device independent)
Scope: full review. 3 files changed (26 additions / 8 deletions); reviewed all three — modelopt/torch/sparsity/state_sparsity/config.py, modelopt/torch/sparsity/state_sparsity/policy.py, and tests/unit/torch/sparsity/state_sparsity/test_dasc.py.
Findings: CRITICAL: 1, IMPORTANT: 0, SUGGESTION: 0
The two changes in this PR are correct on their own terms:
math.log(epsilon) / decayatpolicy.py:125removes the last ambient-device allocation fromcompute_gdn_decay_horizons, keeps the float64 precision of the previoustorch.log(both are Cdouble), and now matches the already-scalarscale = -math.log(epsilon)in_storage_cast_horizon_bounds(policy.py:335) — so the two horizon paths are consistent.mathis already imported atpolicy.py:22.- The
not isinstance(..., bool)guards in_validate_analysis_argumentsare the right fix:boolsubclassesint, which is registered withnumbers.Real, soTruepreviously passed asepsilon(and0.0 < True < 1.0isFalse, so it happened to be rejected, butstatic_gate_input=Truewas silently accepted). Narrowingexcept (TypeError, ValueError, OverflowError)toexcept OverflowErrorholds for every reachable input, since non-Realtypes short-circuit at theisinstancecheck beforemath.isfiniteis reached. - The
torch.device("meta")regression genuinely exercises the fix — on the old code the ambient-devicetorch.tensor(epsilon)produced a meta result andnot torch.isfinite(horizons).all()would have raised on thebool()conversion. It also runs on CPU-only CI.
[CRITICAL ModeState] modelopt/torch/sparsity/state_sparsity/policy.py:523 — the same ambient-default-device bug remains in the save/apply path
validate_dasc_decay_parameters still builds a tensor with an ambient-device factory call:
stored = torch.tensor(layer.static_horizons, dtype=torch.float64)What the issue is. lower and upper come from _storage_cast_horizon_bounds, which explicitly forces .to(device="cpu", ...) and (after this PR) uses only scalar math.log, so both are unconditionally CPU float64. stored, by contrast, honors PyTorch's ambient default device. Under torch.set_default_device("cuda") — or inside a with torch.device(...) block, which is how HF and Megatron model construction/loading is commonly wrapped — stored lands on CUDA and the very next comparison
if torch.any(stored < lower * (1.0 - numerical_slack)) or torch.any(
stored > upper * (1.0 + numerical_slack)
):raises RuntimeError: Expected all tensors to be on the same device.
Why it matters. This is the same defect class the PR set out to eliminate, but on a path that is worse to break than the analysis helper:
conversion.py:122calls it fromupdate_dasc_metadata, guarded byexcept ApplyModeError. A device-mismatchRuntimeErroris not anApplyModeError, so it escapes the staleness-warning handler and propagates straight out ofmto.save()/mto.modelopt_state(). Saving a DASC checkpoint fails hard instead of emitting the intended "re-run calibrate()" warning.api.py:87calls it on the calibrate/convert path, where it likewise turns into an unhandledRuntimeError.
Under a meta default device the failure mode is different but equally fatal: stored becomes a meta tensor and torch.any(...) cannot be converted to a Python bool.
Suggested fix — pin the tensor to CPU, matching the explicit device="cpu" convention used everywhere else in this module:
stored = torch.tensor(layer.static_horizons, dtype=torch.float64, device="cpu")This is now the only remaining device-sensitive factory call in modelopt/torch/sparsity/state_sparsity/ (verified by grepping torch.tensor/zeros/ones/arange/empty/full/as_tensor across the sub-package), so fixing it completes the PR's stated goal.
It would also be worth extending the new regression beyond compute_gdn_decay_horizons — e.g. a calibrate() + mto.modelopt_state() round trip inside with torch.device("meta") (or torch.set_default_device) would have caught this line, whereas test_horizon_computation_ignores_the_default_device cannot reach it.
Risk assessment
Low risk as written, but the fix is incomplete. Nothing in the diff regresses existing behavior, and modelopt_state schema, mode registration, and public signatures are untouched, so there are no backward-compatibility or restore-fidelity concerns from the changed lines themselves. The single blocking item is the leftover ambient-device allocation on the save path — a one-line change that keeps the PR's device-independence guarantee from having a hole in the path users hit when they persist a DASC checkpoint.
🤖 Generated with Claude Code
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
|
|
Addressed the remaining full-package device-independence finding in #2399. |
|
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## feature/dasc-state-sparsity-review-real-scalars #2398 +/- ##
====================================================================================
+ Coverage 68.76% 78.78% +10.02%
====================================================================================
Files 548 548
Lines 64244 64243 -1
====================================================================================
+ Hits 44175 50615 +6440
+ 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:
|
## Summary Closes the remaining full-package device-independence finding on #2398: - explicitly allocate stored policy horizons on CPU before comparing with CPU bounds - exercise calibration and `modelopt_state()` under a non-CPU ambient default-device context - retain the existing low-level horizon default-device regression A full factory-call scan confirms this was the only remaining ambient-device tensor allocation in `state_sparsity`. ## Validation - focused DASC suite: 52 passed, 1 optional Megatron skip - pre-commit hooks on both changed files - signed commit with DCO sign-off <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Bug Fixes** * Improved DASC calibration and checkpoint metadata generation when PyTorch’s default device is set to `meta`. * Ensured decay-parameter validation remains CPU-backed for reliable processing. * **Tests** * Added regression coverage for CPU-backed DASC calibration and checkpoint metadata. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: kaix-nv <kaix@nvidia.com>
e6da8d5
into
feature/dasc-state-sparsity-review-real-scalars
1 participant
Summary
Addresses every finding on #2397:
math.log, avoiding ambient PyTorch default-device allocation entirelynumbers.RealinheritanceOverflowErrorValidation
Summary by CodeRabbit
Bug Fixes
Tests