Pin DASC policy validation to CPU - #2399
Conversation
Signed-off-by: kaix-nv <kaix@nvidia.com>
|
/claude review |
|
@coderabbitai full review |
|
Warning Review limit reachedNext included review available in 1 minute. 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 (1)
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; 0 remain after this review. 📝 WalkthroughWalkthroughDASC decay-parameter validation now creates the stored horizon tensor on the CPU. A regression test covers calibration and ModelOpt state extraction when PyTorch uses ChangesDASC CPU placement
Priority: ⬇️ Low Estimated code review effort: 2 (Simple) | ~10 minutes Change: Bug fix Merge Risk: ⚪ Minimal · up to The DASC change preserves calibration and state extraction under a non-CPU default-device context, with focused validation passing. 🚥 Pre-merge checks | ✅ 6✅ Passed checks (6 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Comment |
✅ Action performedFull review finished. |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
| with torch.device("meta"): | ||
| calibrated = mtss.calibrate(model, _config(wmax_candidates=[7]), [_candidate(7)]) | ||
| state = mto.modelopt_state(calibrated) | ||
| assert state["modelopt_state_dict"][0][0] == "dasc" |
There was a problem hiding this comment.
[SUGGESTION] This regression guard only works by accident of the error type that escapes.
The line being fixed lives in validate_dasc_decay_parameters, and the only path from mto.modelopt_state() to it is update_dasc_metadata (conversion.py:118-127), which deliberately downgrades validation failures to warnings:
else:
try:
validate_dasc_decay_parameters(model, policy)
except ApplyModeError as error:
warnings.warn(f"{error}. The restored DASC policy is stale; ...")Today the un-pinned torch.tensor(...) fails with a device-mismatch RuntimeError, which slips past that except ApplyModeError and fails the test — so the test does currently catch the regression. But nothing in the assertion depends on validation having succeeded: state["modelopt_state_dict"][0][0] == "dasc" is true whether the horizon check passed or was swallowed into a staleness warning. If a future change wraps device/dtype errors in ApplyModeError (which the file already does for ValueError at policy.py:500-503), this test goes green while the device bug is back.
Two ways to make the guard load-bearing:
- Assert the validator was actually satisfied, by rejecting the staleness warning:
def test_policy_lifecycle_ignores_the_default_device():
"""Keep calibration and checkpoint metadata validation on their declared CPU path."""
model = TinyGatedDeltaNetForCausalLM()
with torch.device("meta"), warnings.catch_warnings():
warnings.simplefilter("error")
calibrated = mtss.calibrate(model, _config(wmax_candidates=[7]), [_candidate(7)])
state = mto.modelopt_state(calibrated)
assert state["modelopt_state_dict"][0][0] == "dasc"- Or drive the changed line through
export_policy(api.py:86-87), the one caller that propagatesApplyModeErrorinstead of warning — so any failure of the horizon check surfaces regardless of exception type:
with torch.device("meta"):
calibrated = mtss.calibrate(model, _config(wmax_candidates=[7]), [_candidate(7)])
state = mto.modelopt_state(calibrated)
policy = mtss.export_policy(calibrated)
assert state["modelopt_state_dict"][0][0] == "dasc"
assert policy["layers"]["linear_attn"]["static_horizons"]Either keeps the higher-level lifecycle coverage you wanted while making the assertion fail for the reason the test names. Non-blocking.
There was a problem hiding this comment.
Fixed in #2400. The default-device lifecycle regression now calls export_policy() and asserts validated static horizons, so an ApplyModeError cannot be silently downgraded by the checkpoint metadata path.
![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; the PR is 2 files, +10/-1. Both files reviewed in full context (modelopt/torch/sparsity/state_sparsity/policy.py, tests/unit/torch/sparsity/state_sparsity/test_dasc.py), plus the surrounding call graph needed to judge the fix: api.py, conversion.py (restore / update_dasc_metadata paths), and the rest of policy.py.
Findings
- CRITICAL: 0
- IMPORTANT: 0
- SUGGESTION: 1 — the new lifecycle test's guard is only load-bearing by accident of exception type (inline)
Assessment
The fix is correct and correctly scoped. layer.static_horizons is a plain Python list, so torch.tensor(...) was the one allocation in the module that honored PyTorch's ambient default-device context — under a non-CPU torch.device() context the resulting tensor would not be comparable with the CPU lower/upper bounds from _storage_cast_horizon_bounds, making validate_dasc_decay_parameters fail for a reason unrelated to policy validity. Pinning device="cpu" matches the function's documented CPU-float64 contract and the explicit .to(device="cpu", ...) idiom used everywhere else in the file.
I independently confirmed the PR's completeness claim rather than taking it on trust. Scanning state_sparsity/ for tensor factory calls (torch.tensor/zeros/ones/empty/full/arange/eye/rand/randn/linspace/as_tensor/from_numpy) and for other device-sensitive construction (torch.Tensor(, new_*, scalar_tensor, nn.Parameter() returns line 523 as the only factory call in the package. Every other tensor in the module derives from a module parameter and is explicitly moved with .to(device="cpu", ...) (policy.py:121-122, 214-215, 281-285, 306, 331-332), so it inherits nothing from the ambient device. The device-independence work on this file is complete.
Also checked and found clean: no modelopt_state schema change (metadata is still {"policy": ...}, so existing DASC checkpoints round-trip unaffected), no public API or config surface change, no mode-registration change, and no new CPU-GPU sync — the pinned tensor and its comparison operands are all CPU-resident by construction. No CHANGELOG.rst entry is warranted: this is a fix within the same unreleased cycle as the feature it touches.
Note on verification: I was not able to execute the test suite in this environment (Python execution was declined by the sandbox), so the 52-passed result in the PR description is unverified by me; my analysis of the fix and of the test's coverage is from reading the code paths. Tracing them, the new test does reach the changed line — mto.modelopt_state() -> update_dasc_metadata -> validate_dasc_decay_parameters — and a device-mismatch RuntimeError escapes the except ApplyModeError there, so the test does fail without the fix. The inline SUGGESTION is about making that dependence explicit rather than incidental.
Risk: low.
🤖 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-device-independent #2399 +/- ##
==========================================================================================
+ Coverage 68.75% 78.78% +10.02%
==========================================================================================
Files 548 548
Lines 64244 64243 -1
==========================================================================================
+ Hits 44173 50615 +6442
+ Misses 20071 13628 -6443
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 Addresses the sole non-blocking suggestion on approved PR #2399: - call `export_policy()` under the non-CPU ambient default-device context - assert validated static horizons are present - make the regression fail even if a future checkpoint-save path downgrades validation failures to warnings ## Validation - targeted lifecycle regression passed - pre-commit hooks passed - signed commit with DCO sign-off <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Tests** * Extended device lifecycle coverage to verify policy export while operating on the `meta` device. * Added validation that static horizon metadata is included in exported policies. <!-- end of auto-generated comment: release notes by coderabbit.ai --> Signed-off-by: kaix-nv <kaix@nvidia.com>
1dcafa5
into
feature/dasc-state-sparsity-review-device-independent
1 participant
Summary
Closes the remaining full-package device-independence finding on #2398:
modelopt_state()under a non-CPU ambient default-device contextA full factory-call scan confirms this was the only remaining ambient-device tensor allocation in
state_sparsity.Validation
Summary by CodeRabbit
Bug Fixes
meta.Tests