Handle DASC storage round trips and imports

modelopt/torch/sparsity/state_sparsity/policy.py

@@ -19,6 +19,7 @@
 import importlib
 import json
 import math
+import warnings
 from collections.abc import Iterable
 from functools import lru_cache
 
@@ -38,6 +39,7 @@
     ("megatron.core.ssm.gated_delta_net", "GatedDeltaNet"),
     ("transformers.models.qwen3_next.modeling_qwen3_next", "Qwen3NextGatedDeltaNet"),
 )
+_SUPPORTED_STORAGE_DTYPES = (torch.float16, torch.bfloat16)
 
 
 @lru_cache(maxsize=1)
@@ -47,10 +49,26 @@ def _supported_gdn_classes() -> tuple[type[nn.Module], ...]:
     for module_name, class_name in _SUPPORTED_GDN_CLASS_PATHS:
         try:
             candidate = getattr(importlib.import_module(module_name), class_name)
-        except (AttributeError, ImportError):
+        except ModuleNotFoundError as error:
+            root_module = module_name.partition(".")[0]
+            if importlib.util.find_spec(root_module) is None:
+                continue
+            warnings.warn(
+                f"DASC could not resolve {module_name}.{class_name}: {error!r}", stacklevel=2
+            )
+            continue
+        except Exception as error:
+            warnings.warn(
+                f"DASC could not resolve {module_name}.{class_name}: {error!r}", stacklevel=2
+            )
             continue
         if isinstance(candidate, type) and issubclass(candidate, nn.Module):
             classes.append(candidate)
+        else:
+            warnings.warn(
+                f"DASC resolved {module_name}.{class_name}, but it is not an nn.Module class",
+                stacklevel=2,
+            )
     return tuple(classes)
 
 
@@ -81,9 +99,8 @@ def compute_gdn_decay_horizons(
     return horizons
 
 
-def _is_gdn_module(module: nn.Module) -> bool:
+def _is_gdn_module(module: nn.Module, supported_classes: tuple[type[nn.Module], ...]) -> bool:
     """Accept supported GDN implementations and their ModelOpt dynamic subclasses."""
-    supported_classes = _supported_gdn_classes()
     module_class = type(module)
     is_supported_class = module_class in supported_classes or (
         isinstance(module, DynamicModule)
@@ -99,22 +116,37 @@ def _is_gdn_module(module: nn.Module) -> bool:
 def _get_gdn_modules(model: nn.Module) -> dict[str, nn.Module]:
     """Find supported GDN layers after removing a recognized model wrapper."""
     model = unwrap_model(model, force_unwrap=True)
-    modules = {name: module for name, module in model.named_modules() if _is_gdn_module(module)}
+    supported_classes = _supported_gdn_classes()
+    named_modules = list(model.named_modules())
+    modules = {
+        name: module for name, module in named_modules if _is_gdn_module(module, supported_classes)
+    }
     if not modules:
-        supported = ", ".join(class_name for _, class_name in _SUPPORTED_GDN_CLASS_PATHS)
+        unsupported_subclasses = [
+            name or "<root>"
+            for name, module in named_modules
+            if not isinstance(module, DynamicModule)
+            and type(module) not in supported_classes
+            and any(base in supported_classes for base in type(module).__mro__[1:])
+        ]
+        if unsupported_subclasses:
+            raise ApplyModeError(
+                "DASC found GDN subclasses that are not ModelOpt dynamic modules at: "
+                f"{', '.join(unsupported_subclasses)}; use an exact supported class"
+            )
+        supported = ", ".join(
+            f"{module_name}.{class_name}" for module_name, class_name in _SUPPORTED_GDN_CLASS_PATHS
+        )
         raise ApplyModeError(f"DASC found no supported GDN modules; expected one of: {supported}")
     return dict(sorted(modules.items()))
 
 
-def analyze_gdn_decay(
-    model: nn.Module,
-    *,
-    epsilon: float = 1e-3,
-    static_gate_input: float = -0.3,
+def _analyze_gdn_modules(
+    modules: dict[str, nn.Module], *, epsilon: float, static_gate_input: float
 ) -> dict[str, list[float]]:
-    """Return deterministic per-head horizons for every GDN module in a model."""
+    """Compute deterministic per-head horizons for already-resolved GDN modules."""
     horizons = {}
-    for name, module in _get_gdn_modules(model).items():
+    for name, module in modules.items():
         try:
             layer_horizons = compute_gdn_decay_horizons(
                 module.A_log,
@@ -130,6 +162,18 @@ def analyze_gdn_decay(
     return horizons
 
 
+def analyze_gdn_decay(
+    model: nn.Module,
+    *,
+    epsilon: float = 1e-3,
+    static_gate_input: float = -0.3,
+) -> dict[str, list[float]]:
+    """Return deterministic per-head horizons for every GDN module in a model."""
+    return _analyze_gdn_modules(
+        _get_gdn_modules(model), epsilon=epsilon, static_gate_input=static_gate_input
+    )
+
+
 def _canonical_sha256(value: object) -> str:
     """Hash a JSON value with deterministic ordering and no non-finite numbers."""
     payload = json.dumps(value, sort_keys=True, separators=(",", ":"), allow_nan=False)
@@ -166,12 +210,19 @@ def _decay_parameters(modules: dict[str, nn.Module]) -> list[dict[str, object]]:
 
 
 def _storage_rounding_radius(tensor: torch.Tensor) -> torch.Tensor:
-    """Bound one cast-to-storage rounding step around the represented tensor values."""
-    values = tensor.detach().to(device="cpu", dtype=torch.float64)
-    dtype_info = torch.finfo(tensor.dtype)
-    unit_roundoff = dtype_info.eps / 2.0
-    subnormal_slack = dtype_info.tiny * dtype_info.eps
-    return values.abs() * (unit_roundoff / (1.0 - unit_roundoff)) + subnormal_slack
+    """Bound one FP16/BF16 storage cast even after values are reloaded in a wider dtype."""
+    if not tensor.dtype.is_floating_point:
+        raise ApplyModeError("DASC decay parameters must use a floating-point dtype")
+    values = tensor.detach().to(device="cpu", dtype=torch.float64).abs()
+    radius = torch.zeros_like(values)
+    for dtype in (tensor.dtype, *_SUPPORTED_STORAGE_DTYPES):
+        dtype_info = torch.finfo(dtype)
+        unit_roundoff = dtype_info.eps / 2.0
+        candidate = (
+            values * (unit_roundoff / (1.0 - unit_roundoff)) + dtype_info.tiny * dtype_info.eps
+        )
+        radius = torch.maximum(radius, candidate)
+    return radius
 
 
 def _storage_cast_horizon_bounds(
@@ -261,8 +312,8 @@ def build_dasc_policy(
     validated_measurements.sort(key=lambda measurement: measurement.wmax)
 
     modules = _get_gdn_modules(model)
-    horizons = analyze_gdn_decay(
-        model, epsilon=config.epsilon, static_gate_input=config.static_gate_input
+    horizons = _analyze_gdn_modules(
+        modules, epsilon=config.epsilon, static_gate_input=config.static_gate_input
     )
     _validate_measurement_geometry(horizons, validated_measurements)
 
@@ -333,12 +384,12 @@ def validate_dasc_decay_parameters(model: nn.Module, policy: DASCPolicy) -> None
     Numerical validation deliberately uses re-derived horizons and the exact selected mask rather
     than the provenance digest because an FP16 or BF16 storage cast is lossy.
     """
-    current_horizons = analyze_gdn_decay(
-        model,
+    modules = _get_gdn_modules(model)
+    current_horizons = _analyze_gdn_modules(
+        modules,
         epsilon=policy.epsilon,
         static_gate_input=policy.static_gate_input,
     )
-    modules = _get_gdn_modules(model)
     for name, values in current_horizons.items():
         layer = policy.layers[name]
         retained = [head for head, horizon in enumerate(values) if horizon > policy.selected_wmax]

tests/unit/torch/sparsity/state_sparsity/test_dasc.py

@@ -244,7 +244,7 @@ def __init__(self):
     class UnsupportedSubclass(GatedDeltaNet):
         pass
 
-    with pytest.raises(ApplyModeError, match="no supported GDN modules"):
+    with pytest.raises(ApplyModeError, match="subclasses that are not ModelOpt dynamic modules"):
         mtss.calibrate(UnsupportedSubclass(), _config(wmax_candidates=[7]), [_candidate(7)])
 
     same_name_lookalike = type("GatedDeltaNet", (UnsupportedGatedDeltaNet,), {})
@@ -263,21 +263,42 @@ def test_supported_class_resolution_uses_imported_module_identities(monkeypatch)
         ("valid", "GatedDeltaNet"),
         ("invalid", "NotAModule"),
         ("missing", "Missing"),
+        ("installed", "Missing"),
+        ("broken", "Broken"),
     )
     modules = {
         "valid": type("ValidModule", (), {"GatedDeltaNet": GatedDeltaNet}),
         "invalid": type("InvalidModule", (), {"NotAModule": object()}),
     }
 
     def import_module(name):
-        if name == "missing":
-            raise ImportError(name)
+        if name in {"missing", "installed"}:
+            raise ModuleNotFoundError(name)
+        if name == "broken":
+            raise RuntimeError(name)
         return modules[name]
 
     monkeypatch.setattr(dasc_policy, "_SUPPORTED_GDN_CLASS_PATHS", module_paths)
     monkeypatch.setattr(dasc_policy.importlib, "import_module", import_module)
+    monkeypatch.setattr(
+        dasc_policy.importlib.util,
+        "find_spec",
+        lambda name: object() if name == "installed" else None,
+    )
 
-    assert _resolve_supported_gdn_classes() == (GatedDeltaNet,)
+    with pytest.warns(UserWarning) as caught:
+        assert _resolve_supported_gdn_classes() == (GatedDeltaNet,)
+    assert len(caught) == 3
+
+
+@pytest.mark.parametrize(("module_name", "class_name"), dasc_policy._SUPPORTED_GDN_CLASS_PATHS)
+def test_declared_gdn_paths_resolve_when_framework_is_installed(module_name, class_name):
+    """Guard supported identities against upstream dependency path drift."""
+    root_module = module_name.partition(".")[0]
+    if dasc_policy.importlib.util.find_spec(root_module) is None:
+        pytest.skip(f"optional framework {root_module!r} is not installed")
+    module = dasc_policy.importlib.import_module(module_name)
+    assert issubclass(getattr(module, class_name), nn.Module)
 
 
 def test_generic_mode_application_reports_missing_measurements(monkeypatch):
@@ -357,6 +378,34 @@ def test_dtype_cast_preserves_policy_when_the_selected_mask_is_unchanged(dtype):
     assert mtss.export_policy(model) == policy
 
 
+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(
+        TinyGatedDeltaNetForCausalLM(), _config(wmax_candidates=[7]), [_candidate(7)]
+    )
+    policy = mtss.export_policy(model)
+    original_decay = torch.cat(
+        [model.linear_attn.A_log.detach(), model.linear_attn.dt_bias.detach()]
+    )
+
+    with torch.no_grad():
+        model.linear_attn.A_log.copy_(model.linear_attn.A_log.to(torch.bfloat16).float())
+        model.linear_attn.dt_bias.copy_(model.linear_attn.dt_bias.to(torch.bfloat16).float())
+
+    reloaded_decay = torch.cat(
+        [model.linear_attn.A_log.detach(), model.linear_attn.dt_bias.detach()]
+    )
+    assert reloaded_decay.dtype == torch.float32
+    assert not torch.equal(original_decay, reloaded_decay)
+    assert mtss.export_policy(model) == policy
+
+    model.linear_attn.A_log = nn.Parameter(
+        model.linear_attn.A_log.detach().to(torch.int64), requires_grad=False
+    )
+    with pytest.raises(ApplyModeError, match="floating-point dtype"):
+        mtss.export_policy(model)
+
+
 def test_export_rejects_changed_decay_parameters_and_restore_rejects_structure():
     """Keep saving recoverable while rejecting stale or tampered deployment policies."""
     model = mtss.calibrate(