Skip to content

Pin DASC policy validation to CPU #2399

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
kaix-nv merged 2 commits into from
Sep 11, 2026
Merged

Pin DASC policy validation to CPU #2399

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
48daa95
Pin DASC policy validation to CPU
kaix-nv Sep 11, 2026
431bb04
Strengthen DASC device lifecycle regression (#2400)
kaix-nv Sep 11, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion modelopt/torch/sparsity/state_sparsity/policy.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -520,7 +520,7 @@ def validate_dasc_decay_parameters(model: nn.Module, policy: DASCPolicy) -> None
"DASC policy head mask does not match current decay parameters in layer "
f"{name!r}"
)
stored = torch.tensor(layer.static_horizons, dtype=torch.float64)
stored = torch.tensor(layer.static_horizons, device="cpu", dtype=torch.float64)
numerical_slack = 32.0 * torch.finfo(torch.float64).eps
if torch.any(stored < lower * (1.0 - numerical_slack)) or torch.any(
stored > upper * (1.0 + numerical_slack)
Expand Down
11 changes: 11 additions & 0 deletions tests/unit/torch/sparsity/state_sparsity/test_dasc.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -508,6 +508,17 @@ def test_horizon_computation_ignores_the_default_device():
assert horizons.device.type == "cpu"


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"):
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"
Comment on lines +514 to +518

@claude claude Bot Sep 11, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[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:

  1. 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"
  1. Or drive the changed line through export_policy (api.py:86-87), the one caller that propagates ApplyModeError instead 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.

@kaix-nv kaix-nv Sep 11, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

assert policy["layers"]["linear_attn"]["static_horizons"]


def test_bf16_storage_round_trip_loaded_in_fp32_preserves_policy():
"""Accept BF16-rounded values after a checkpoint loader materializes FP32 tensors."""
model = mtss.calibrate(
Expand Down
Loading