From 9a7720f2ff32024a5ca43cd03bacbb7cfce5a6b8 Mon Sep 17 00:00:00 2001 From: Kai Xu Date: Thu, 10 Sep 2026 19:22:00 -0700 Subject: [PATCH] Record DASC decay storage dtype Signed-off-by: Kai Xu --- docs/source/guides/6_sparsity.rst | 5 ++ .../torch/sparsity/state_sparsity/config.py | 7 +- .../sparsity/state_sparsity/conversion.py | 1 + .../torch/sparsity/state_sparsity/policy.py | 86 +++++++++++++------ .../sparsity/state_sparsity/test_dasc.py | 57 ++++++++++-- 5 files changed, 120 insertions(+), 36 deletions(-) diff --git a/docs/source/guides/6_sparsity.rst b/docs/source/guides/6_sparsity.rst index 2f5d7b85ef2..691b637f022 100644 --- a/docs/source/guides/6_sparsity.rst +++ b/docs/source/guides/6_sparsity.rst @@ -139,6 +139,8 @@ ordinary dense recurrent state before continuation. config = { "variant": "dasc_wr", # "dasc_nr" uses zero recovery instead "wmax_candidates": [32], + # Set this to the dtype used to store A_log and dt_bias in the checkpoint. + "decay_parameter_storage_dtype": "bfloat16", "model_id": "org/model", "model_revision": "immutable-model-revision", "model_config_id": "sha256:", @@ -179,6 +181,9 @@ Re-running :func:`~modelopt.torch.sparsity.state_sparsity.calibrate` replaces th 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. +Set ``decay_parameter_storage_dtype`` to the checkpoint dtype for ``A_log`` and ``dt_bias`` before +calibration. Policy validation allows only the rounding introduced by that declared storage dtype +and the live tensor dtype; the default ``float32`` keeps unconfigured policies strict. .. _sparsity-concepts: diff --git a/modelopt/torch/sparsity/state_sparsity/config.py b/modelopt/torch/sparsity/state_sparsity/config.py index 96df5202a18..7f67d7f7d47 100644 --- a/modelopt/torch/sparsity/state_sparsity/config.py +++ b/modelopt/torch/sparsity/state_sparsity/config.py @@ -86,6 +86,10 @@ class DASCConfig(ModeloptBaseConfig): default=-0.3, description="Static gate input added to each GDN head's dt_bias.", ) + decay_parameter_storage_dtype: Literal["float16", "bfloat16", "float32"] = ModeloptField( + default="float32", + description="Expected checkpoint storage dtype for GDN A_log and dt_bias.", + ) wmax_candidates: list[int] = ModeloptField( default=[8, 16, 32, 64, 128, 256], description="Positive candidate windows evaluated during offline calibration.", @@ -202,6 +206,7 @@ class DASCPolicy(ModeloptBaseConfig): recovery: Literal["zero", "suffix_replay"] epsilon: float static_gate_input: float + decay_parameter_storage_dtype: Literal["float16", "bfloat16", "float32"] = "float32" selected_wmax: int = Field(strict=True, gt=0) wmax_candidates: list[int] = Field(min_length=1) quality_gates: dict[str, float] @@ -215,7 +220,7 @@ class DASCPolicy(ModeloptBaseConfig): model_structure_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.", + description="Storage-dtype-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 a748cc813b6..f826f09b9df 100644 --- a/modelopt/torch/sparsity/state_sparsity/conversion.py +++ b/modelopt/torch/sparsity/state_sparsity/conversion.py @@ -69,6 +69,7 @@ def restore_dasc_model(model: nn.Module, config: DASCConfig, metadata: MetadataD "variant": config.variant, "epsilon": config.epsilon, "static_gate_input": config.static_gate_input, + "decay_parameter_storage_dtype": config.decay_parameter_storage_dtype, "wmax_candidates": config.wmax_candidates, "quality_gates": { "min_perplexity_retention": config.min_perplexity_retention, diff --git a/modelopt/torch/sparsity/state_sparsity/policy.py b/modelopt/torch/sparsity/state_sparsity/policy.py index d80c7fddcdc..d72ac620b5e 100644 --- a/modelopt/torch/sparsity/state_sparsity/policy.py +++ b/modelopt/torch/sparsity/state_sparsity/policy.py @@ -39,7 +39,11 @@ ("megatron.core.ssm.gated_delta_net", "GatedDeltaNet"), ("transformers.models.qwen3_next.modeling_qwen3_next", "Qwen3NextGatedDeltaNet"), ) -_SUPPORTED_STORAGE_DTYPES = (torch.float16, torch.bfloat16) +_STORAGE_DTYPES = { + "float16": torch.float16, + "bfloat16": torch.bfloat16, + "float32": torch.float32, +} @lru_cache(maxsize=1) @@ -84,6 +88,8 @@ def compute_gdn_decay_horizons( raise ValueError( "GDN A_log and dt_bias must be non-empty one-dimensional tensors of equal shape" ) + if not a_log.dtype.is_floating_point or not dt_bias.dtype.is_floating_point: + raise ValueError("GDN A_log and dt_bias must use floating-point dtypes") if not 0.0 < epsilon < 1.0: raise ValueError("epsilon must be in (0, 1)") @@ -101,16 +107,20 @@ def compute_gdn_decay_horizons( def _is_gdn_module(module: nn.Module, supported_classes: tuple[type[nn.Module], ...]) -> bool: """Accept supported GDN implementations and their ModelOpt dynamic subclasses.""" + return _has_supported_gdn_identity(module, supported_classes) and all( + isinstance(getattr(module, name, None), torch.Tensor) for name in ("A_log", "dt_bias") + ) + + +def _has_supported_gdn_identity( + module: nn.Module, supported_classes: tuple[type[nn.Module], ...] +) -> bool: + """Return whether a module has an exact or ModelOpt-generated supported GDN identity.""" module_class = type(module) - is_supported_class = module_class in supported_classes or ( + return module_class in supported_classes or ( isinstance(module, DynamicModule) and any(base in supported_classes for base in module_class.__mro__) ) - return ( - is_supported_class - and isinstance(getattr(module, "A_log", None), torch.Tensor) - and isinstance(getattr(module, "dt_bias", None), torch.Tensor) - ) def _get_gdn_modules(model: nn.Module) -> dict[str, nn.Module]: @@ -122,6 +132,20 @@ def _get_gdn_modules(model: nn.Module) -> dict[str, nn.Module]: name: module for name, module in named_modules if _is_gdn_module(module, supported_classes) } if not modules: + missing_decay_parameters = [ + name or "" + for name, module in named_modules + if _has_supported_gdn_identity(module, supported_classes) + and not all( + isinstance(getattr(module, parameter, None), torch.Tensor) + for parameter in ("A_log", "dt_bias") + ) + ] + if missing_decay_parameters: + raise ApplyModeError( + "DASC found supported GDN modules without A_log and dt_bias tensors at: " + f"{', '.join(missing_decay_parameters)}" + ) unsupported_subclasses = [ name or "" for name, module in named_modules @@ -132,7 +156,8 @@ def _get_gdn_modules(model: nn.Module) -> dict[str, nn.Module]: 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" + f"{', '.join(unsupported_subclasses)}; convert the module with ModelOpt or use a " + "supported class directly" ) supported = ", ".join( f"{module_name}.{class_name}" for module_name, class_name in _SUPPORTED_GDN_CLASS_PATHS @@ -191,17 +216,19 @@ 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 a compact BF16-canonicalized calibration snapshot for provenance.""" +def _decay_parameters( + modules: dict[str, nn.Module], storage_dtype: torch.dtype +) -> list[dict[str, object]]: + """Serialize a storage-dtype-canonicalized calibration snapshot for provenance.""" return [ { "name": name, "A_log": module.A_log.detach() - .to(device="cpu", dtype=torch.bfloat16) + .to(device="cpu", dtype=storage_dtype) .to(dtype=torch.float32) .tolist(), "dt_bias": module.dt_bias.detach() - .to(device="cpu", dtype=torch.bfloat16) + .to(device="cpu", dtype=storage_dtype) .to(dtype=torch.float32) .tolist(), } @@ -209,30 +236,31 @@ def _decay_parameters(modules: dict[str, nn.Module]) -> list[dict[str, object]]: ] -def _storage_rounding_radius(tensor: torch.Tensor) -> torch.Tensor: - """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") +def _storage_rounding_radius(tensor: torch.Tensor, storage_dtype: torch.dtype) -> torch.Tensor: + """Compose inverse error bounds for storage and live-dtype materialization casts.""" values = tensor.detach().to(device="cpu", dtype=torch.float64).abs() - radius = torch.zeros_like(values) - for dtype in (tensor.dtype, *_SUPPORTED_STORAGE_DTYPES): + upper = values + cast_dtypes = tuple(dict.fromkeys((storage_dtype, tensor.dtype))) + for dtype in reversed(cast_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 + smallest_subnormal = dtype_info.tiny * dtype_info.eps + upper = (upper + smallest_subnormal) / (1.0 - unit_roundoff) + return upper - values def _storage_cast_horizon_bounds( - module: nn.Module, *, epsilon: float, static_gate_input: float + module: nn.Module, + *, + epsilon: float, + static_gate_input: float, + storage_dtype: torch.dtype, ) -> tuple[torch.Tensor, torch.Tensor]: """Bound horizons compatible with the current parameters before one storage cast.""" a_log = module.A_log.detach().to(device="cpu", dtype=torch.float64) dt_bias = module.dt_bias.detach().to(device="cpu", dtype=torch.float64) - a_radius = _storage_rounding_radius(module.A_log) - dt_radius = _storage_rounding_radius(module.dt_bias) + a_radius = _storage_rounding_radius(module.A_log, storage_dtype) + dt_radius = _storage_rounding_radius(module.dt_bias, storage_dtype) scale = -math.log(epsilon) lower = scale / ( torch.exp(a_log + a_radius) * F.softplus(dt_bias + dt_radius + static_gate_input) @@ -343,6 +371,7 @@ def build_dasc_policy( recovery="zero" if config.variant == "dasc_nr" else "suffix_replay", epsilon=config.epsilon, static_gate_input=config.static_gate_input, + decay_parameter_storage_dtype=config.decay_parameter_storage_dtype, selected_wmax=selected_wmax, wmax_candidates=config.wmax_candidates, quality_gates={ @@ -357,7 +386,9 @@ def build_dasc_policy( granularity=config.granularity, preserve_convolution_state=config.preserve_convolution_state, model_structure_sha256=_canonical_sha256(_model_structure(modules)), - decay_parameters_sha256=_canonical_sha256(_decay_parameters(modules)), + decay_parameters_sha256=_canonical_sha256( + _decay_parameters(modules, _STORAGE_DTYPES[config.decay_parameter_storage_dtype]) + ), layers=layers, measurements=validated_measurements, ) @@ -401,6 +432,7 @@ def validate_dasc_decay_parameters(model: nn.Module, policy: DASCPolicy) -> None modules[name], epsilon=policy.epsilon, static_gate_input=policy.static_gate_input, + storage_dtype=_STORAGE_DTYPES[policy.decay_parameter_storage_dtype], ) stored = torch.tensor(layer.static_horizons, dtype=torch.float64) numerical_slack = 32.0 * torch.finfo(torch.float64).eps diff --git a/tests/unit/torch/sparsity/state_sparsity/test_dasc.py b/tests/unit/torch/sparsity/state_sparsity/test_dasc.py index aca9c692775..18c3fe15796 100644 --- a/tests/unit/torch/sparsity/state_sparsity/test_dasc.py +++ b/tests/unit/torch/sparsity/state_sparsity/test_dasc.py @@ -121,6 +121,7 @@ def test_calibrate_selects_largest_passing_candidate_and_round_trips(): assert policy["granularity"] == "gdn_head" assert policy["preserve_convolution_state"] is True assert policy["active_runtime_state"] == "dense" + assert policy["decay_parameter_storage_dtype"] == "float32" assert policy["layers"]["linear_attn"]["retained_heads"] == [0] assert policy["layers"]["linear_attn"]["omitted_heads"] == [1] assert [measurement["wmax"] for measurement in policy["measurements"]] == [7, 11] @@ -129,11 +130,18 @@ def test_calibrate_selects_largest_passing_candidate_and_round_trips(): ) json.dumps(policy) - restored = mto.restore_from_modelopt_state( - TinyGatedDeltaNetForCausalLM(), mto.modelopt_state(model) - ) + state = mto.modelopt_state(model) + restored = mto.restore_from_modelopt_state(TinyGatedDeltaNetForCausalLM(), state) assert mtss.export_policy(restored) == policy + legacy_state = copy.deepcopy(state) + del legacy_state["modelopt_state_dict"][0][1]["config"]["decay_parameter_storage_dtype"] + del legacy_state["modelopt_state_dict"][0][1]["metadata"]["policy"][ + "decay_parameter_storage_dtype" + ] + legacy_restored = mto.restore_from_modelopt_state(TinyGatedDeltaNetForCausalLM(), legacy_state) + assert mtss.export_policy(legacy_restored)["decay_parameter_storage_dtype"] == "float32" + policy["selected_wmax"] = 999 assert mtss.export_policy(model)["selected_wmax"] == 7 @@ -161,6 +169,7 @@ def test_variant_is_explicit_in_exported_policy(variant, recovery): {"wmax_candidates": [7, 7]}, {"model_revision": ""}, {"preserve_convolution_state": False}, + {"decay_parameter_storage_dtype": "float8"}, ], ) def test_config_fails_closed(override): @@ -251,6 +260,11 @@ class UnsupportedSubclass(GatedDeltaNet): with pytest.raises(ApplyModeError, match="no supported GDN modules"): mtss.calibrate(same_name_lookalike(), _config(wmax_candidates=[7]), [_candidate(7)]) + missing_decay = TinyGatedDeltaNetForCausalLM() + del missing_decay.linear_attn.A_log + with pytest.raises(ApplyModeError, match="without A_log and dt_bias tensors"): + mtss.calibrate(missing_decay, _config(wmax_candidates=[7]), [_candidate(7)]) + invalid_decay = TinyGatedDeltaNetForCausalLM() invalid_decay.linear_attn.dt_bias = nn.Parameter(torch.zeros(3)) with pytest.raises(ApplyModeError, match="Invalid GDN decay parameters"): @@ -295,8 +309,7 @@ def import_module(name): 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") + pytest.importorskip(root_module) module = dasc_policy.importlib.import_module(module_name) assert issubclass(getattr(module, class_name), nn.Module) @@ -361,8 +374,11 @@ def test_public_exports_and_wrapped_model_export(): @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.""" + storage_dtype = "bfloat16" if dtype == torch.bfloat16 else "float16" model = mtss.calibrate( - TinyGatedDeltaNetForCausalLM(), _config(wmax_candidates=[7]), [_candidate(7)] + TinyGatedDeltaNetForCausalLM(), + _config(wmax_candidates=[7], decay_parameter_storage_dtype=storage_dtype), + [_candidate(7)], ) policy = mtss.export_policy(model) original_decay = torch.cat( @@ -381,7 +397,9 @@ def test_dtype_cast_preserves_policy_when_the_selected_mask_is_unchanged(dtype): 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)] + TinyGatedDeltaNetForCausalLM(), + _config(wmax_candidates=[7], decay_parameter_storage_dtype="bfloat16"), + [_candidate(7)], ) policy = mtss.export_policy(model) original_decay = torch.cat( @@ -406,6 +424,29 @@ def test_bf16_storage_round_trip_loaded_in_fp32_preserves_policy(): mtss.export_policy(model) +@pytest.mark.parametrize( + ("storage_name", "storage_dtype", "live_dtype"), + [ + ("float16", torch.float16, torch.bfloat16), + ("bfloat16", torch.bfloat16, torch.float16), + ], +) +def test_cross_dtype_reload_accumulates_both_rounding_bounds( + storage_name, storage_dtype, live_dtype +): + """Accept two distinct declared-storage and live-materialization rounding steps.""" + model = mtss.calibrate( + TinyGatedDeltaNetForCausalLM(), + _config(wmax_candidates=[7], decay_parameter_storage_dtype=storage_name), + [_candidate(7)], + ) + policy = mtss.export_policy(model) + + model.to(storage_dtype).to(live_dtype) + + assert mtss.export_policy(model) == policy + + def test_export_rejects_changed_decay_parameters_and_restore_rejects_structure(): """Keep saving recoverable while rejecting stale or tampered deployment policies.""" model = mtss.calibrate( @@ -469,7 +510,7 @@ def test_export_rejects_changed_decay_parameters_and_restore_rejects_structure() layer = tampered_state["modelopt_state_dict"][0][1]["metadata"]["policy"]["layers"][ "linear_attn" ] - layer["static_horizons"][0] *= 1.01 + layer["static_horizons"][0] *= 1.0001 restored = mto.restore_from_modelopt_state(TinyGatedDeltaNetForCausalLM(), tampered_state) with pytest.raises(ApplyModeError, match="horizons do not match"): mtss.export_policy(restored)