File size: 4,067 Bytes
8101696 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 | """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())
|