From 48daa9584d5840440fb9f75ad23b4f6042f99d4c Mon Sep 17 00:00:00 2001 From: kaix-nv Date: Thu, 10 Sep 2026 21:38:26 -0700 Subject: [PATCH 1/2] Pin DASC policy validation to CPU Signed-off-by: kaix-nv --- modelopt/torch/sparsity/state_sparsity/policy.py | 2 +- tests/unit/torch/sparsity/state_sparsity/test_dasc.py | 9 +++++++++ 2 files changed, 10 insertions(+), 1 deletion(-) diff --git a/modelopt/torch/sparsity/state_sparsity/policy.py b/modelopt/torch/sparsity/state_sparsity/policy.py index 52097f05e76..7baf53e9586 100644 --- a/modelopt/torch/sparsity/state_sparsity/policy.py +++ b/modelopt/torch/sparsity/state_sparsity/policy.py @@ -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) diff --git a/tests/unit/torch/sparsity/state_sparsity/test_dasc.py b/tests/unit/torch/sparsity/state_sparsity/test_dasc.py index ddf4023d6f2..0f3ca8974fc 100644 --- a/tests/unit/torch/sparsity/state_sparsity/test_dasc.py +++ b/tests/unit/torch/sparsity/state_sparsity/test_dasc.py @@ -508,6 +508,15 @@ 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) + assert state["modelopt_state_dict"][0][0] == "dasc" + + 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( From 431bb04d8c18fa2cd2fc3de468ef5e2592a6448a Mon Sep 17 00:00:00 2001 From: kaix-nv Date: Thu, 10 Sep 2026 21:49:21 -0700 Subject: [PATCH 2/2] Strengthen DASC device lifecycle regression (#2400) ## 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 ## 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. Signed-off-by: kaix-nv --- tests/unit/torch/sparsity/state_sparsity/test_dasc.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/unit/torch/sparsity/state_sparsity/test_dasc.py b/tests/unit/torch/sparsity/state_sparsity/test_dasc.py index 0f3ca8974fc..cad44b1501d 100644 --- a/tests/unit/torch/sparsity/state_sparsity/test_dasc.py +++ b/tests/unit/torch/sparsity/state_sparsity/test_dasc.py @@ -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"] def test_bf16_storage_round_trip_loaded_in_fp32_preserves_policy():