Skip to content

Strengthen DASC device lifecycle regression #2400

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 1 commit into from
Sep 11, 2026
+2 −0
Merged

Strengthen DASC device lifecycle regression #2400

Changes from all commits
Commits
Show all changes
1 commit
Select commit
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: 2 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 @@ -514,7 +514,9 @@ def test_policy_lifecycle_ignores_the_default_device():
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"]

@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 assertion can never fail on its own, so it carries less signal than it looks like it does.

DASCLayerPolicy.static_horizons is declared list[float] = Field(min_length=1) and validate_partition additionally requires len(static_horizons) == num_heads with every value finite and > 0.0 (modelopt/torch/sparsity/state_sparsity/config.py:215-225). By the time export_policy() returns a dumped policy, a non-empty list is guaranteed by the schema — a truthiness check is tautological beyond confirming the layers/linear_attn/static_horizons key path exists and that export_policy() didn't raise.

That said, the placement is the valuable part and it is correct: every horizon path pins CPU explicitly (policy.py:121-122, 214-215, 331-332, 523), so if a future change dropped one of those pins, the horizons/bounds tensors would materialize on meta and either the CPU-vs-meta comparison in validate_dasc_decay_parameters or the .tolist() in the dump would raise inside the torch.device("meta") block. This does close the gap the PR describes.

To make the assertion itself meaningful, tie the exported horizons to values computed on the normal CPU path — that catches a silently-wrong-but-non-empty policy, not just an exception:

def test_policy_lifecycle_ignores_the_default_device():
    """Keep calibration and checkpoint metadata validation on their declared CPU path."""
    model = TinyGatedDeltaNetForCausalLM()
    expected_horizons = mtss.analyze_gdn_decay(model)["linear_attn"].tolist()
    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"] == pytest.approx(expected_horizons)

Non-blocking — the current form is a net improvement over not calling export_policy() under the ambient device at all.

@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.

Thanks. No further code change here: the load-bearing guard is that export_policy() executes inside the meta default-device context and propagates any validation failure. The truthiness check is retained only as a structural/schema assertion; exact horizon numerics and storage-canonical equality are already covered by dedicated tests, so duplicating that comparison here would mix concerns.



def test_bf16_storage_round_trip_loaded_in_fp32_preserves_policy():
Expand Down
Loading