| """Parameter-group helpers for partial-gradient importance proxies.""" |
|
|
| from __future__ import annotations |
|
|
| from dataclasses import dataclass |
|
|
| import torch |
| from torch import nn |
|
|
|
|
| @dataclass(frozen=True) |
| class PartialGradientGroup: |
| """A named, non-empty subset of action-model parameters.""" |
|
|
| name: str |
| parameter_names: tuple[str, ...] |
| parameter_count: int |
|
|
|
|
| def define_partial_gradient_groups( |
| action_model: nn.Module, |
| ) -> dict[str, PartialGradientGroup]: |
| """Return the four parameter groups used for partial-gradient analysis.""" |
| named_parameters = dict(action_model.named_parameters()) |
| if not named_parameters: |
| raise ValueError("action model has no parameters") |
|
|
| transformer_blocks = getattr( |
| getattr(action_model, "model", None), |
| "transformer_blocks", |
| None, |
| ) |
| if transformer_blocks is None or len(transformer_blocks) == 0: |
| raise ValueError("action model has no DiT transformer blocks") |
| last_block_prefix = f"model.transformer_blocks.{len(transformer_blocks) - 1}." |
|
|
| selectors = { |
| "decoder": lambda name: name.startswith("action_decoder."), |
| "output_head": lambda name: ( |
| name.startswith("action_decoder.") |
| or name.startswith("model.proj_out_1.") |
| or name.startswith("model.proj_out_2.") |
| ), |
| "last_block_and_head": lambda name: ( |
| name.startswith(last_block_prefix) |
| or name.startswith("action_decoder.") |
| or name.startswith("model.proj_out_1.") |
| or name.startswith("model.proj_out_2.") |
| ), |
| "normalization": lambda name: "norm" in name.lower(), |
| } |
|
|
| groups = {} |
| for group_name, selector in selectors.items(): |
| selected_names = tuple( |
| name for name in named_parameters if selector(name) |
| ) |
| if not selected_names: |
| raise ValueError( |
| f"partial-gradient group {group_name!r} has no parameters" |
| ) |
| groups[group_name] = PartialGradientGroup( |
| name=group_name, |
| parameter_names=selected_names, |
| parameter_count=sum( |
| named_parameters[name].numel() for name in selected_names |
| ), |
| ) |
| return groups |
|
|
|
|
| def activate_partial_gradient_group( |
| action_model: nn.Module, |
| group: PartialGradientGroup, |
| ) -> list[nn.Parameter]: |
| """Freeze the action model except for one exact, validated group.""" |
| named_parameters = dict(action_model.named_parameters()) |
| requested = set(group.parameter_names) |
| missing = requested - set(named_parameters) |
| if missing: |
| raise ValueError( |
| f"group {group.name!r} contains unknown parameters: {sorted(missing)}" |
| ) |
|
|
| action_model.requires_grad_(False) |
| selected = [] |
| for name, parameter in named_parameters.items(): |
| if name in requested: |
| parameter.requires_grad_(True) |
| selected.append(parameter) |
| if not selected: |
| raise RuntimeError(f"group {group.name!r} activated no parameters") |
| observed_count = sum(parameter.numel() for parameter in selected) |
| if observed_count != group.parameter_count: |
| raise RuntimeError( |
| f"group {group.name!r} parameter count changed from " |
| f"{group.parameter_count} to {observed_count}" |
| ) |
| return selected |
|
|
|
|
| def finite_gradient_l2_norm(parameters: list[nn.Parameter]) -> float: |
| """Compute an exact finite L2 norm over the gradients that were produced.""" |
| component_norms = [] |
| for parameter in parameters: |
| if parameter.grad is None: |
| continue |
| gradient = parameter.grad.detach() |
| if not torch.isfinite(gradient).all(): |
| raise FloatingPointError("partial-gradient proxy is non-finite") |
| component_norms.append(torch.linalg.vector_norm(gradient.float())) |
| if not component_norms: |
| raise RuntimeError("partial-gradient group received no gradients") |
| return float(torch.linalg.vector_norm(torch.stack(component_norms)).item()) |
|
|
