Record DASC decay storage dtype #2383
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change | ||||
|---|---|---|---|---|---|---|
|
|
@@ -39,7 +39,11 @@ | |||||
| ("megatron.core.ssm.gated_delta_net", "GatedDeltaNet"), | ||||||
| ("transformers.models.qwen3_next.modeling_qwen3_next", "Qwen3NextGatedDeltaNet"), | ||||||
| ) | ||||||
| _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) | ||||||
| return module_class in supported_classes or ( | ||||||
| isinstance(module, DynamicModule) | ||||||
| and any(base in supported_classes for base in module_class.__mro__) | ||||||
| ) | ||||||
|
|
||||||
|
|
||||||
| 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 "<root>" | ||||||
| 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") | ||||||
| ) | ||||||
| ] | ||||||
|
Comment on lines
+135
to
+143
There was a problem hiding this comment. [SUGGESTION] The This block only runs inside Suggested simplification: if not modules:
missing_decay_parameters = [
name or "<root>"
for name, module in named_modules
if _has_supported_gdn_identity(module, supported_classes)
]Behavior is identical, and the tensor predicate stays in one place. Separately (not a regression, pre-existing): when some GDN modules have the tensors and one does not, the one without is silently dropped from the policy rather than reported here. That is self-consistent with
Contributor
Author
There was a problem hiding this comment. Addressed in #2384. Supported identities are now collected once and every matched identity is checked for A_log and dt_bias before any layer is accepted. The duplicated dead predicate is removed, and a mixed valid-plus-malformed model test proves the malformed layer cannot be silently omitted. |
||||||
| if missing_decay_parameters: | ||||||
| raise ApplyModeError( | ||||||
| "DASC found supported GDN modules without A_log and dt_bias tensors at: " | ||||||
| f"{', '.join(missing_decay_parameters)}" | ||||||
| ) | ||||||
|
Comment on lines
+135
to
+148
Contributor
There was a problem hiding this comment. 🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win Validate all supported GDN modules before filtering valid modules.
🤖 Prompt for AI Agents |
||||||
| unsupported_subclasses = [ | ||||||
| name or "<root>" | ||||||
| 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)}; 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,48 +216,51 @@ def _model_structure(modules: dict[str, nn.Module]) -> list[dict[str, object]]: | |||||
| ] | ||||||
|
|
||||||
|
|
||||||
| 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=storage_dtype) | ||||||
| .to(dtype=torch.float32) | ||||||
| .tolist(), | ||||||
| "dt_bias": module.dt_bias.detach() | ||||||
| .to(device="cpu", dtype=storage_dtype) | ||||||
| .to(dtype=torch.float32) | ||||||
| .tolist(), | ||||||
| } | ||||||
| for name, module in modules.items() | ||||||
| ] | ||||||
|
|
||||||
|
|
||||||
| 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() | ||||||
| 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 | ||||||
| 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, | ||||||
| storage_dtype: torch.dtype, | ||||||
| ) -> tuple[torch.Tensor, torch.Tensor]: | ||||||
| """Bound horizons compatible with the current parameters before one storage cast.""" | ||||||
|
There was a problem hiding this comment. [SUGGESTION] This docstring is now stale: the function no longer bounds "one storage cast".
Suggested change
Also worth double-checking the ordering intent is documented somewhere: |
||||||
| 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, 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, _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], | ||||||
|
There was a problem hiding this comment. [SUGGESTION] The rejection message a few lines below ( Before this PR the bound unconditionally granted FP16+BF16 slack; now a policy that omits raise ApplyModeError(
f"DASC policy horizons do not match current decay parameters in layer {name!r} "
f"within the declared {policy.decay_parameter_storage_dtype} storage and "
f"{modules[name].A_log.dtype} live rounding bounds"
)The math itself checks out: for the documented flow (declare the checkpoint dtype, then round-trip) the composed bound covers the true error in every combination I traced — |
||||||
| ) | ||||||
| stored = torch.tensor(layer.static_horizons, dtype=torch.float64) | ||||||
| numerical_slack = 32.0 * torch.finfo(torch.float64).eps | ||||||
|
|
||||||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
[SUGGESTION] The allowed dtype set is now spelled out in three places that must stay in lockstep: this
Literal, the identicalLiteralonDASCConfig.decay_parameter_storage_dtype(line 89), and_STORAGE_DTYPESinpolicy.py. If a future dtype is added to eitherLiteralbut not to the dict,_STORAGE_DTYPES[policy.decay_parameter_storage_dtype]raises a bareKeyErrorfrom insidevalidate_dasc_decay_parameters/build_dasc_policyinstead of a config-validation error — a drift that pydantic can't catch.A single alias in
config.pyused by both fields removes two of the three copies:and
policy.pycan then key_STORAGE_DTYPESofftyping.get_args(DecayParameterStorageDtype)(or simply assert the sets match at import) so the mapping cannot silently fall behind.