Harden DASC recalibration lifecycle

docs/source/guides/6_sparsity.rst

@@ -172,10 +172,12 @@ ordinary dense recurrent state before continuation.
 from zero, while DASC-WR reconstructs them from a zero-initialized suffix replay of at most the
 selected ``Wmax`` tokens. Both retain whole GDN heads, preserve convolution state, and resume with
 dense recurrence. KDA and serving-runtime integration are not supported by this initial API.
-The initial GDN adapter accepts the ``GatedDeltaNet`` and ``Qwen3NextGatedDeltaNet`` class names
-and fails closed for other implementations, even when they expose similarly named decay tensors.
-Re-running :func:`~modelopt.torch.sparsity.state_sparsity.calibrate` supersedes a stale policy. A
-stale policy remains serializable so it cannot block saving or composing other ModelOpt modes, but
+The initial GDN adapter accepts the ``GatedDeltaNet`` and ``Qwen3NextGatedDeltaNet`` base classes,
+including ModelOpt-generated dynamic subclasses, and fails closed for unrelated implementations
+even when they expose similarly named decay tensors.
+Re-running :func:`~modelopt.torch.sparsity.state_sparsity.calibrate` replaces the existing DASC
+mode-state entry and supersedes its stale policy without growing the checkpoint history. A stale
+policy remains serializable so it cannot block saving or composing other ModelOpt modes, but
 :func:`~modelopt.torch.sparsity.state_sparsity.export_policy` rejects it until recalibration.
 
 .. _sparsity-concepts:

modelopt/torch/sparsity/state_sparsity/api.py

@@ -21,10 +21,11 @@
 
 from torch import nn
 
-from modelopt.torch.opt.conversion import apply_mode
+from modelopt.torch.opt.conversion import ModeloptStateManager, apply_mode
+from modelopt.torch.utils import unwrap_model
 
 from .config import DASCCalibrationMeasurement, DASCConfig
-from .conversion import get_attached_dasc_policy
+from .conversion import get_attached_dasc_policy, replace_dasc_mode
 from .mode import DASCModeRegistry
 from .policy import validate_dasc_decay_parameters, validate_dasc_model_structure
 
@@ -40,6 +41,7 @@ def calibrate(
 
     ``measurements`` must contain exactly one entry for every configured ``Wmax`` candidate.
     The largest candidate passing every quality, lifecycle, and storage gate is selected.
+    Recalibrating replaces the existing DASC mode-state entry in place.
 
     Example::
 
@@ -58,10 +60,16 @@ def calibrate(
     Returns:
         The input model with a serializable DASC policy attached through ModelOpt state.
     """
-    config_dict = config.model_dump() if isinstance(config, DASCConfig) else config
+    model = unwrap_model(model, force_unwrap=True)
+    config_object = config if isinstance(config, DASCConfig) else DASCConfig(**config)
+    if ModeloptStateManager.is_converted(model, is_root=True) and any(
+        mode == "dasc" for mode, _ in ModeloptStateManager(model).state_dict()
+    ):
+        return replace_dasc_mode(model, config_object, measurements)
+
     return apply_mode(
         model,
-        mode=[("dasc", config_dict)],
+        mode=[("dasc", config_object.model_dump())],
         registry=DASCModeRegistry,
         mode_kwargs={"measurements": measurements},
     )

modelopt/torch/sparsity/state_sparsity/config.py

@@ -72,7 +72,7 @@ def validate_unique_slices(
 class DASCConfig(ModeloptBaseConfig):
     """Configuration for GDN decay-aware state checkpoint sparsity."""
 
-    model_config = ConfigDict(extra="forbid", validate_assignment=True, protected_namespaces=())
+    model_config = ConfigDict(protected_namespaces=())
 
     variant: Literal["dasc_nr", "dasc_wr"] = ModeloptField(
         default="dasc_wr",
@@ -195,7 +195,7 @@ def validate_partition(self) -> "DASCLayerPolicy":
 class DASCPolicy(ModeloptBaseConfig):
     """Standalone JSON-safe DASC deployment policy."""
 
-    model_config = ConfigDict(extra="forbid", validate_assignment=True, protected_namespaces=())
+    model_config = ConfigDict(protected_namespaces=())
 
     format_version: Literal[1] = 1
     variant: Literal["dasc_nr", "dasc_wr"]
@@ -213,7 +213,10 @@ class DASCPolicy(ModeloptBaseConfig):
     preserve_convolution_state: Literal[True]
     active_runtime_state: Literal["dense"] = "dense"
     model_structure_sha256: str = Field(pattern=r"^[0-9a-f]{64}$")
-    decay_parameters_sha256: str = Field(pattern=r"^[0-9a-f]{64}$")
+    decay_parameters_sha256: str = Field(
+        pattern=r"^[0-9a-f]{64}$",
+        description="BF16-canonicalized calibration snapshot retained for provenance.",
+    )
     layers: dict[str, DASCLayerPolicy] = Field(min_length=1)
     measurements: list[DASCCalibrationMeasurement] = Field(min_length=1)
 

modelopt/torch/sparsity/state_sparsity/conversion.py

@@ -22,7 +22,7 @@
 from pydantic import ValidationError
 from torch import nn
 
-from modelopt.torch.opt.conversion import ApplyModeError
+from modelopt.torch.opt.conversion import ApplyModeError, ModeloptStateManager
 from modelopt.torch.opt.mode import ConvertReturnType, MetadataDict
 from modelopt.torch.utils import unwrap_model
 
@@ -98,8 +98,8 @@ def restore_dasc_model(model: nn.Module, config: DASCConfig, metadata: MetadataD
 def update_dasc_metadata(model: nn.Module, config: DASCConfig, metadata: MetadataDict) -> None:
     """Refresh metadata without making unrelated ModelOpt save or compose paths unusable."""
     policy = get_attached_dasc_policy(model)
-    validate_dasc_model_structure(model, policy)
     try:
+        validate_dasc_model_structure(model, policy)
         validate_dasc_decay_parameters(model, policy)
     except ApplyModeError as error:
         warnings.warn(
@@ -110,6 +110,35 @@ def update_dasc_metadata(model: nn.Module, config: DASCConfig, metadata: Metadat
     metadata["policy"] = copy.deepcopy(policy.model_dump(mode="json"))
 
 
+def replace_dasc_mode(
+    model: nn.Module,
+    config: DASCConfig,
+    measurements: Iterable[DASCCalibrationMeasurement | dict],
+) -> nn.Module:
+    """Replace existing DASC mode state in place with a newly derived policy."""
+    model = unwrap_model(model, force_unwrap=True)
+    policy = build_dasc_policy(model, config, measurements)
+    manager = ModeloptStateManager(model)
+    manager.update_last_state_before_new_mode(model)
+    state = manager.state_dict()
+    dasc_indices = [index for index, (mode, _) in enumerate(state) if mode == "dasc"]
+    if not dasc_indices:
+        raise ApplyModeError("Cannot replace DASC mode because the model has no DASC state")
+
+    first_index = dasc_indices[0]
+    state[first_index] = (
+        "dasc",
+        {
+            "config": config.model_dump(),
+            "metadata": {"policy": policy.model_dump(mode="json")},
+        },
+    )
+    for index in reversed(dasc_indices[1:]):
+        del state[index]
+    _attach_policy(model, policy)
+    return model
+
+
 def get_attached_dasc_policy(model: nn.Module) -> DASCPolicy:
     """Return the validated policy attached by conversion or restoration."""
     model = unwrap_model(model, force_unwrap=True)

modelopt/torch/sparsity/state_sparsity/mode.py

@@ -48,6 +48,11 @@ def config_class(self) -> type[ModeloptBaseConfig]:
         """Return the validated DASC configuration class."""
         return DASCConfig
 
+    @property
+    def next_prohibited_modes(self) -> set[str]:
+        """Route repeat calibration through the replacing public API."""
+        return {"dasc"}
+
     @property
     def convert(self) -> ConvertEntrypoint:
         """Return the DASC calibration entrypoint."""

modelopt/torch/sparsity/state_sparsity/policy.py

@@ -31,6 +31,7 @@
 __all__ = ["analyze_gdn_decay", "compute_gdn_decay_horizons"]
 
 _SUPPORTED_GDN_CLASS_NAMES = frozenset({"GatedDeltaNet", "Qwen3NextGatedDeltaNet"})
+# Covers ordinary FP16/BF16 storage casts; the exact retained-head mask below is the semantic gate.
 _HORIZON_DTYPE_CAST_RTOL = 0.05
 
 
@@ -62,9 +63,9 @@ def compute_gdn_decay_horizons(
 
 
 def _is_gdn_module(module: nn.Module) -> bool:
-    """Return whether a module has one of the explicitly supported GDN implementations."""
+    """Accept supported GDN implementations and their ModelOpt dynamic subclasses."""
     return (
-        type(module).__name__ in _SUPPORTED_GDN_CLASS_NAMES
+        any(base.__name__ in _SUPPORTED_GDN_CLASS_NAMES for base in type(module).__mro__)
         and isinstance(getattr(module, "A_log", None), torch.Tensor)
         and isinstance(getattr(module, "dt_bias", None), torch.Tensor)
     )
@@ -122,7 +123,7 @@ def _model_structure(modules: dict[str, nn.Module]) -> list[dict[str, object]]:
 
 
 def _decay_parameters(modules: dict[str, nn.Module]) -> list[dict[str, object]]:
-    """Serialize decay parameters at canonical BF16 precision for dtype-stable hashing."""
+    """Serialize a compact BF16-canonicalized calibration snapshot for provenance."""
     return [
         {
             "name": name,
@@ -275,12 +276,11 @@ def validate_dasc_model_structure(model: nn.Module, policy: DASCPolicy) -> None:
 
 
 def validate_dasc_decay_parameters(model: nn.Module, policy: DASCPolicy) -> None:
-    """Reject deployment when current decay parameters no longer derive the stored policy."""
-    modules = _get_gdn_modules(model)
-    actual = _canonical_sha256(_decay_parameters(modules))
-    if actual != policy.decay_parameters_sha256:
-        raise ApplyModeError("DASC policy does not match the model's GDN decay parameters")
+    """Reject deployment when current decay parameters no longer derive the stored policy.
 
+    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,
         epsilon=policy.epsilon,

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

@@ -26,7 +26,9 @@
 
 import modelopt.torch.opt as mto
 import modelopt.torch.sparsity.state_sparsity as mtss
-from modelopt.torch.opt.conversion import ApplyModeError
+from modelopt.torch.opt.conversion import ApplyModeError, ModeloptStateManager
+from modelopt.torch.sparsity.state_sparsity.conversion import replace_dasc_mode
+from modelopt.torch.sparsity.state_sparsity.mode import DASCModeRegistry
 
 
 class GatedDeltaNet(nn.Module):
@@ -218,8 +220,13 @@ class NotGDN(nn.Module):
     with pytest.raises(ApplyModeError, match="no supported GDN modules"):
         mtss.calibrate(NotGDN(), _config(wmax_candidates=[7]), [_candidate(7)])
 
-    class UnsupportedGatedDeltaNet(GatedDeltaNet):
-        pass
+    class UnsupportedGatedDeltaNet(nn.Module):
+        """Unrelated lookalike must not satisfy the supported-base-class contract."""
+
+        def __init__(self):
+            super().__init__()
+            self.A_log = nn.Parameter(torch.zeros(2))
+            self.dt_bias = nn.Parameter(torch.zeros(2))
 
     with pytest.raises(ApplyModeError, match="no supported GDN modules"):
         mtss.calibrate(UnsupportedGatedDeltaNet(), _config(wmax_candidates=[7]), [_candidate(7)])
@@ -232,34 +239,52 @@ class UnsupportedGatedDeltaNet(GatedDeltaNet):
 
 def test_generic_mode_application_reports_missing_measurements():
     """Give generic apply_mode callers an actionable calibration-evidence error."""
+    assert DASCModeRegistry["dasc"].next_prohibited_modes == {"dasc"}
     with pytest.raises(ApplyModeError, match="requires calibration measurements"):
         mto.apply_mode(
             TinyGatedDeltaNetForCausalLM(),
             mode=[("dasc", _config(wmax_candidates=[7]))],
         )
 
+    model = TinyGatedDeltaNetForCausalLM()
+    ModeloptStateManager(model, init_state=True)
+    with pytest.raises(ApplyModeError, match="model has no DASC state"):
+        replace_dasc_mode(
+            model,
+            mtss.DASCConfig(**_config(wmax_candidates=[7])),
+            [_candidate(7)],
+        )
+
 
 def test_public_exports_and_wrapped_model_export():
-    """Expose only supported symbols and unwrap recognized parallel wrappers."""
+    """Expose only supported symbols and accept wrappers and ModelOpt subclasses."""
     model = mtss.calibrate(
         TinyGatedDeltaNetForCausalLM(), _config(wmax_candidates=[7]), [_candidate(7)]
     )
 
     assert "DASCLayerPolicy" in mtss.__all__
     assert "mode" not in mtss.__all__
     assert mtss.export_policy(nn.DataParallel(model)) == mtss.export_policy(model)
+
+    class _DynamicGatedDeltaNet(GatedDeltaNet):
+        """Stand in for the subclass synthesized by ModelOpt dynamic conversion."""
+
+    model.linear_attn = _DynamicGatedDeltaNet()
+    assert mtss.export_policy(model)["layers"]["linear_attn"]["num_heads"] == 2
+
     with pytest.raises(ApplyModeError, match="no valid attached DASC policy"):
         mtss.export_policy(TinyGatedDeltaNetForCausalLM())
 
 
-def test_dtype_cast_preserves_policy_when_the_selected_mask_is_unchanged():
-    """Treat an ordinary BF16 cast as equivalent when it preserves the policy mask."""
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
+def test_dtype_cast_preserves_policy_when_the_selected_mask_is_unchanged(dtype):
+    """Treat ordinary low-precision casts as equivalent when they preserve the policy mask."""
     model = mtss.calibrate(
         TinyGatedDeltaNetForCausalLM(), _config(wmax_candidates=[7]), [_candidate(7)]
     )
     policy = mtss.export_policy(model)
 
-    model.to(torch.bfloat16)
+    model.to(dtype)
 
     assert mtss.export_policy(model) == policy
 
@@ -273,7 +298,7 @@ def test_export_rejects_changed_decay_parameters_and_restore_rejects_structure()
 
     with torch.no_grad():
         model.linear_attn.A_log.add_(1.0)
-    with pytest.raises(ApplyModeError, match="decay parameters"):
+    with pytest.raises(ApplyModeError, match="horizons"):
         mtss.export_policy(model)
     with pytest.warns(UserWarning, match="saved DASC policy is stale"):
         mto.modelopt_state(model)
@@ -282,11 +307,19 @@ def test_export_rejects_changed_decay_parameters_and_restore_rejects_structure()
         mto.save(model, checkpoint)
     assert checkpoint.tell() > 0
 
+    manager_state = ModeloptStateManager(model).state_dict()
+    manager_state.append(copy.deepcopy(manager_state[0]))
     with pytest.warns(UserWarning, match="saved DASC policy is stale"):
-        model = mtss.calibrate(model, _config(wmax_candidates=[7]), [_candidate(7)])
+        model = mtss.calibrate(
+            model,
+            _config(wmax_candidates=[7], model_revision="revision-2"),
+            [_candidate(7)],
+        )
     policy = mtss.export_policy(model)
+    assert policy["model_revision"] == "revision-2"
     model_state = copy.deepcopy(model.state_dict())
     recalibrated_state = mto.modelopt_state(model)
+    assert [mode for mode, _ in recalibrated_state["modelopt_state_dict"]] == ["dasc"]
     restored = mto.restore_from_modelopt_state(TinyGatedDeltaNetForCausalLM(), recalibrated_state)
     restored.load_state_dict(model_state)
     assert mtss.export_policy(restored) == policy
@@ -323,6 +356,22 @@ def test_export_rejects_changed_decay_parameters_and_restore_rejects_structure()
         mtss.export_policy(restored)
 
 
+def test_structure_staleness_does_not_block_checkpoint_save():
+    """Keep checkpoint saving available after a calibrated GDN structure changes."""
+    model = mtss.calibrate(
+        TinyGatedDeltaNetForCausalLM(), _config(wmax_candidates=[7]), [_candidate(7)]
+    )
+    model.linear_attn = GatedDeltaNet(num_heads=3)
+
+    checkpoint = io.BytesIO()
+    with pytest.warns(UserWarning, match="saved DASC policy is stale"):
+        mto.save(model, checkpoint)
+
+    assert checkpoint.tell() > 0
+    with pytest.raises(ApplyModeError, match="module structure"):
+        mtss.export_policy(model)
+
+
 def test_export_rederives_the_selected_head_mask():
     """Reject a self-consistent stored mask that current decay parameters do not derive."""
     model = mtss.calibrate(