From 81cc09eef38565435709564b431984397fded12a Mon Sep 17 00:00:00 2001 From: kaix-nv Date: Thu, 10 Sep 2026 21:32:03 -0700 Subject: [PATCH 1/2] Make DASC scalar analysis device independent Signed-off-by: kaix-nv --- .../torch/sparsity/state_sparsity/config.py | 15 ++++++++++----- .../torch/sparsity/state_sparsity/policy.py | 2 +- .../torch/sparsity/state_sparsity/test_dasc.py | 17 +++++++++++++++-- 3 files changed, 26 insertions(+), 8 deletions(-) diff --git a/modelopt/torch/sparsity/state_sparsity/config.py b/modelopt/torch/sparsity/state_sparsity/config.py index 91ca8dc1261..1a8f87e0236 100644 --- a/modelopt/torch/sparsity/state_sparsity/config.py +++ b/modelopt/torch/sparsity/state_sparsity/config.py @@ -43,18 +43,23 @@ def _validate_analysis_arguments( """Reject decay-analysis arguments that cannot produce well-defined horizons.""" try: epsilon_is_valid = ( - isinstance(epsilon, Real) and math.isfinite(epsilon) and 0.0 < epsilon < 1.0 + isinstance(epsilon, Real) + and not isinstance(epsilon, bool) + and math.isfinite(epsilon) + and 0.0 < epsilon < 1.0 ) - except (TypeError, ValueError, OverflowError): + except OverflowError: epsilon_is_valid = False if not epsilon_is_valid: raise ValueError("epsilon must be finite and in (0, 1)") try: - static_gate_input_is_valid = isinstance(static_gate_input, Real) and math.isfinite( - static_gate_input + static_gate_input_is_valid = ( + isinstance(static_gate_input, Real) + and not isinstance(static_gate_input, bool) + and math.isfinite(static_gate_input) ) - except (TypeError, ValueError, OverflowError): + except OverflowError: static_gate_input_is_valid = False if not static_gate_input_is_valid: raise ValueError("static_gate_input must be finite") diff --git a/modelopt/torch/sparsity/state_sparsity/policy.py b/modelopt/torch/sparsity/state_sparsity/policy.py index 9d22d0ff0b5..52097f05e76 100644 --- a/modelopt/torch/sparsity/state_sparsity/policy.py +++ b/modelopt/torch/sparsity/state_sparsity/policy.py @@ -122,7 +122,7 @@ def compute_gdn_decay_horizons( dt_bias_cpu = dt_bias.detach().to(device="cpu", dtype=torch.float64) decay = -torch.exp(a_log_cpu) * F.softplus(dt_bias_cpu + static_gate_input) - horizons = torch.log(torch.tensor(epsilon, dtype=torch.float64)) / decay + horizons = math.log(epsilon) / decay if not torch.isfinite(horizons).all() or not torch.all(horizons > 0): raise ValueError("GDN decay parameters produced non-finite or non-positive horizons") return horizons diff --git a/tests/unit/torch/sparsity/state_sparsity/test_dasc.py b/tests/unit/torch/sparsity/state_sparsity/test_dasc.py index e7c31eec228..ddf4023d6f2 100644 --- a/tests/unit/torch/sparsity/state_sparsity/test_dasc.py +++ b/tests/unit/torch/sparsity/state_sparsity/test_dasc.py @@ -454,14 +454,16 @@ def test_analysis_arguments_fail_at_the_public_boundary(invalid_storage_dtype): ) -@pytest.mark.parametrize("epsilon", [1.0, [], 10**1000, torch.tensor([1e-3, 2e-3])]) +@pytest.mark.parametrize("epsilon", [1.0, True, [], 10**1000, torch.tensor([1e-3, 2e-3])]) def test_analysis_rejects_invalid_epsilon_at_the_public_boundary(epsilon): """Normalize invalid epsilon values to the public ValueError contract.""" with pytest.raises(ValueError, match=r"epsilon must be finite and in \(0, 1\)"): mtss.analyze_gdn_decay(TinyGatedDeltaNetForCausalLM(), epsilon=epsilon) # type: ignore[arg-type] -@pytest.mark.parametrize("static_gate_input", [torch.nan, [], 10**1000, torch.tensor([-0.3, -0.2])]) +@pytest.mark.parametrize( + "static_gate_input", [torch.nan, True, [], 10**1000, torch.tensor([-0.3, -0.2])] +) def test_analysis_rejects_invalid_static_gate_input_at_the_public_boundary(static_gate_input): """Normalize invalid static gate values to the public ValueError contract.""" with pytest.raises(ValueError, match="static_gate_input must be finite"): @@ -475,10 +477,12 @@ def test_analysis_rejects_invalid_static_gate_input_at_the_public_boundary(stati ("argument", "value", "message"), [ ("epsilon", [], r"epsilon must be finite and in \(0, 1\)"), + ("epsilon", True, r"epsilon must be finite and in \(0, 1\)"), ("epsilon", torch.nan, r"epsilon must be finite and in \(0, 1\)"), ("epsilon", 10**1000, r"epsilon must be finite and in \(0, 1\)"), ("epsilon", torch.tensor([1e-3, 2e-3]), r"epsilon must be finite and in \(0, 1\)"), ("static_gate_input", [], "static_gate_input must be finite"), + ("static_gate_input", True, "static_gate_input must be finite"), ("static_gate_input", torch.nan, "static_gate_input must be finite"), ("static_gate_input", 10**1000, "static_gate_input must be finite"), ("static_gate_input", torch.tensor([-0.3, -0.2]), "static_gate_input must be finite"), @@ -495,6 +499,15 @@ def test_horizon_computation_rejects_invalid_public_arguments(argument, value, m ) +def test_horizon_computation_ignores_the_default_device(): + """Keep CPU horizon analysis independent of PyTorch's ambient allocation device.""" + a_log = torch.tensor([0.0]) + dt_bias = torch.tensor([0.0]) + with torch.device("meta"): + horizons = mtss.compute_gdn_decay_horizons(a_log, dt_bias) + assert horizons.device.type == "cpu" + + 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 1dcafa5dc6ff2a70773d2b0654df9d49141fa0cf Mon Sep 17 00:00:00 2001 From: kaix-nv Date: Thu, 10 Sep 2026 21:49:50 -0700 Subject: [PATCH 2/2] Pin DASC policy validation to CPU (#2399) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## 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 ## 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. --------- Signed-off-by: kaix-nv --- modelopt/torch/sparsity/state_sparsity/policy.py | 2 +- tests/unit/torch/sparsity/state_sparsity/test_dasc.py | 11 +++++++++++ 2 files changed, 12 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..cad44b1501d 100644 --- a/tests/unit/torch/sparsity/state_sparsity/test_dasc.py +++ b/tests/unit/torch/sparsity/state_sparsity/test_dasc.py @@ -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" + 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(