diff --git a/docs/source/guides/6_sparsity.rst b/docs/source/guides/6_sparsity.rst index ec83d51fb13..2f5d7b85ef2 100644 --- a/docs/source/guides/6_sparsity.rst +++ b/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: diff --git a/modelopt/torch/sparsity/state_sparsity/api.py b/modelopt/torch/sparsity/state_sparsity/api.py index 60d05620e46..1469e27cc63 100644 --- a/modelopt/torch/sparsity/state_sparsity/api.py +++ b/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}, ) diff --git a/modelopt/torch/sparsity/state_sparsity/config.py b/modelopt/torch/sparsity/state_sparsity/config.py index 26c52e9dcf7..96df5202a18 100644 --- a/modelopt/torch/sparsity/state_sparsity/config.py +++ b/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) diff --git a/modelopt/torch/sparsity/state_sparsity/conversion.py b/modelopt/torch/sparsity/state_sparsity/conversion.py index 95d50a18223..3b6a44bef41 100644 --- a/modelopt/torch/sparsity/state_sparsity/conversion.py +++ b/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) diff --git a/modelopt/torch/sparsity/state_sparsity/mode.py b/modelopt/torch/sparsity/state_sparsity/mode.py index 36b0a881dac..96379f7f7db 100644 --- a/modelopt/torch/sparsity/state_sparsity/mode.py +++ b/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.""" diff --git a/modelopt/torch/sparsity/state_sparsity/policy.py b/modelopt/torch/sparsity/state_sparsity/policy.py index c5f6f9b1838..611b3618f0e 100644 --- a/modelopt/torch/sparsity/state_sparsity/policy.py +++ b/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, diff --git a/tests/unit/torch/sparsity/state_sparsity/test_dasc.py b/tests/unit/torch/sparsity/state_sparsity/test_dasc.py index 570e65fc24e..07782a53dea 100644 --- a/tests/unit/torch/sparsity/state_sparsity/test_dasc.py +++ b/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,15 +239,25 @@ 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)] ) @@ -248,18 +265,26 @@ def test_public_exports_and_wrapped_model_export(): 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(