Coverage for transformer_lens/model_bridge/_relevance_rules.py: 99%
121 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-09-21 19:27 +0000
1"""Forward-equivalent relevance-rule primitives, plus the scoped context that installs them.
3Each primitive is a ``torch.autograd.Function`` that reproduces its native forward
4value exactly while replacing the backward pass with the rule's closed-form VJP.
5``use_relevance_rules`` installs these rules on a model's canonical mount points only
6for the duration of a ``with`` block, targeting components positionally (by mount
7name, never by class) and reporting which mounts were installed versus skipped.
8"""
10import dataclasses
11from contextlib import contextmanager
12from typing import (
13 Any,
14 Callable,
15 Dict,
16 Iterator,
17 List,
18 Mapping,
19 Protocol,
20 Tuple,
21 runtime_checkable,
22)
24import torch
25import torch.nn as nn
28def ln_rule_grad(grad_output: torch.Tensor, denom: torch.Tensor) -> torch.Tensor:
29 """Core LN-rule VJP: divide by ``denom`` without differentiating through it.
31 Shared by the ``ln_rule`` primitive below and by any integration (such as
32 ``NormalizationBridge``) that wraps a component's own native forward call
33 instead of reproducing the division itself.
34 """
35 return grad_output / denom
38class _LNRule(torch.autograd.Function):
39 """LN-rule: forward is ``numerator / denom``; the VJP treats ``denom`` as constant."""
41 @staticmethod
42 def forward(ctx: Any, numerator: torch.Tensor, denom: torch.Tensor) -> torch.Tensor:
43 ctx.save_for_backward(denom)
44 return numerator / denom
46 @staticmethod
47 def backward(ctx: Any, grad_output: torch.Tensor) -> Tuple[torch.Tensor, None]:
48 (denom,) = ctx.saved_tensors
49 return ln_rule_grad(grad_output, denom), None
52def ln_rule(numerator: torch.Tensor, denom: torch.Tensor) -> torch.Tensor:
53 """Apply the LN-rule: native division forward, denom-as-constant backward."""
54 result: torch.Tensor = _LNRule.apply(numerator, denom)
55 return result
58class _IdentityRule(torch.autograd.Function):
59 """Identity-rule: forward is the native activation; the VJP is ``grad_out * phi(x)``.
61 ``phi`` is ``f(x) / x``, the removable singularity at ``x == 0`` filled in with its
62 limit ``0.5``. This holds for any elementwise activation with ``f(0) == 0`` and a
63 well-defined derivative at zero, which covers SiLU and both GELU variants.
64 """
66 @staticmethod
67 def forward(
68 ctx: Any, x: torch.Tensor, act_fn: Callable[[torch.Tensor], torch.Tensor]
69 ) -> torch.Tensor:
70 y = act_fn(x)
71 ctx.save_for_backward(x, y)
72 return y
74 @staticmethod
75 def backward(ctx: Any, grad_output: torch.Tensor) -> Tuple[torch.Tensor, None]:
76 x, y = ctx.saved_tensors
77 safe_x = torch.where(x == 0, torch.ones_like(x), x)
78 phi = torch.where(x == 0, torch.full_like(x, 0.5), y / safe_x)
79 return grad_output * phi, None
82def identity_rule(x: torch.Tensor, act_fn: Callable[[torch.Tensor], torch.Tensor]) -> torch.Tensor:
83 """Apply the Identity-rule for an elementwise activation with ``f(0) == 0``."""
84 result: torch.Tensor = _IdentityRule.apply(x, act_fn)
85 return result
88class _HalfRule(torch.autograd.Function):
89 """Half-rule: forward is ``u * v``; the VJP halves each ordinary product-rule term."""
91 @staticmethod
92 def forward(ctx: Any, u: torch.Tensor, v: torch.Tensor) -> torch.Tensor:
93 ctx.save_for_backward(u, v)
94 return u * v
96 @staticmethod
97 def backward(ctx: Any, grad_output: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
98 u, v = ctx.saved_tensors
99 return 0.5 * grad_output * v, 0.5 * grad_output * u
102def half_rule(u: torch.Tensor, v: torch.Tensor) -> torch.Tensor:
103 """Apply the Half-rule: native product forward, evenly split backward."""
104 result: torch.Tensor = _HalfRule.apply(u, v)
105 return result
108class _ScaleGradient(torch.autograd.Function):
109 """Identity forward; the VJP scales the incoming gradient by a constant factor."""
111 @staticmethod
112 def forward(ctx: Any, x: torch.Tensor, factor: float) -> torch.Tensor:
113 ctx.factor = factor
114 return x
116 @staticmethod
117 def backward(ctx: Any, grad_output: torch.Tensor) -> Tuple[torch.Tensor, None]:
118 return ctx.factor * grad_output, None
121def scale_gradient(x: torch.Tensor, factor: float) -> torch.Tensor:
122 """Pass ``x`` through unchanged while scaling its gradient by ``factor``.
124 The Half-rule on a product ``u * v`` halves each ordinary product-rule term,
125 which is the same as halving the single gradient that enters the product before
126 it splits. When the product is computed inside an opaque module the bridge cannot
127 reach term by term (its native forward is called as one unit), scaling the
128 gradient entering the product by ``0.5`` reproduces the Half-rule at that point
129 without altering the native forward value.
130 """
131 result: torch.Tensor = _ScaleGradient.apply(x, factor)
132 return result
135class RelevanceRuleConflictError(RuntimeError):
136 """A hook would silently break a rule-active forward/backward invariant.
138 Raised instead of the ordinary warn-and-fall-back a component would use when
139 no rule is active, since falling back while a rule is active would compose the
140 rule with the hook edit and break the bit-identical-forward guarantee.
141 """
144class RelevanceRuleUnsupportedError(RuntimeError):
145 """A requested relevance rule cannot be installed on an otherwise-capable component.
147 Raised at ``use_relevance_rules`` entry, before any forward or backward pass, when
148 a component reports the requested kind in its own ``_relevance_rule_unsupported_kinds``
149 -- for example a gated-MLP recompute path backed by an unrecognized weight-orientation
150 class, or an activation form the Identity-rule does not support. Distinct from a
151 kind that is simply absent from ``_relevance_rule_kinds`` without being named there
152 (reported ``skipped``, not raised): that covers a component not implementing the
153 protocol at all, or one whose mount genuinely never deals with the kind (for example
154 normalization on a dispatch path the LN-rule does not wrap), both benign
155 non-applicability rather than a rule request the component was expected to honor.
156 """
159@dataclasses.dataclass(frozen=True)
160class RelevanceRules:
161 """Which relevance rules to request for the duration of a ``use_relevance_rules`` scope.
163 Each field names a rule kind. Setting it ``True`` requests that rule wherever a
164 component at that kind's canonical mount point implements ``_RelevanceRuleCapable``.
165 Unset fields (the default) leave the corresponding components untouched.
166 """
168 normalization: bool = False
169 activation: bool = False
170 multiplicative_gate: bool = False
171 attention: bool = False
174@dataclasses.dataclass
175class RelevanceRuleCoverage:
176 """Which canonical mounts a ``use_relevance_rules`` scope installed versus skipped.
178 ``installed`` holds the dotted path of every mount where a requested rule kind was
179 actually enabled. ``skipped`` holds the dotted path of every mount that matched a
180 requested kind's canonical mount name but did not implement the relevance-rule
181 protocol there, so no rule could be installed.
182 """
184 installed: Tuple[str, ...]
185 skipped: Tuple[str, ...]
188@runtime_checkable
189class _RelevanceRuleCapable(Protocol):
190 """Structural contract a component must satisfy to accept a relevance rule.
192 ``_relevance_rule_kinds`` names every ``RelevanceRules`` field the component
193 answers to at its current mount -- a gated-MLP node answers to both
194 "activation" (Identity-rule on its activation function) and
195 "multiplicative_gate" (Half-rule on its gate*up product) independently, since
196 either can be requested without the other. ``_enable_relevance_rule``/
197 ``_disable_relevance_rule`` take the specific kind being toggled and touch
198 only that kind's state, without touching model configuration, so the
199 component's own state is the only thing that changes and only for the
200 scope's duration.
202 A component may optionally also define ``_relevance_rule_unsupported_kinds``
203 (a ``Tuple[str, ...]``, not part of this structural protocol so components that
204 omit it stay isinstance-compatible) naming kinds it is expected to honor at its
205 mount but currently cannot -- ``use_relevance_rules`` raises
206 ``RelevanceRuleUnsupportedError`` for those instead of reporting them skipped.
207 """
209 _relevance_rule_kinds: Tuple[str, ...]
211 def _enable_relevance_rule(self, kind: str) -> None:
212 ...
214 def _disable_relevance_rule(self, kind: str) -> None:
215 ...
218# Canonical mount name per rule kind. Targeting is positional: a component is only
219# considered for a kind when it sits at that kind's mount name, never by isinstance,
220# so a same-class component mounted elsewhere (for example a q_norm sharing
221# NormalizationBridge's class) is left untouched. The normalization kind lists both
222# the pre-norm mounts (ln1, ln2) and the sandwich post-norm mounts (ln1_post,
223# ln2_post) so the LN-rule reaches the post-attention/post-MLP norms that
224# sandwich-norm architectures mount there, matching the pinned RelP reference.
225_CANONICAL_MOUNTS: Mapping[str, Tuple[str, ...]] = {
226 "normalization": ("ln1", "ln2", "ln1_post", "ln2_post"),
227 "activation": ("mlp",),
228 "multiplicative_gate": ("mlp",),
229}
232def _iter_canonical_mount_candidates(
233 model: nn.Module, mount_names: Tuple[str, ...]
234) -> Iterator[Tuple[str, _RelevanceRuleCapable]]:
235 """Yield each distinct module reachable at one of ``mount_names``, once.
237 A bridge component reachable at a canonical mount name (for example
238 ``blocks.0.ln1``) is also reachable, under the same parent, through the
239 raw HF module tree the bridge wraps in place (for example
240 ``blocks.0._original_component.input_layernorm``) -- both names resolve to
241 the identical object. ``nn.Module.named_modules()`` deduplicates by object
242 identity and keeps only whichever path it visits first, which is the raw
243 HF-attribute path (registered before the canonical alias), so on a real
244 assembled model the canonical name is silently never seen. Walking with
245 ``remove_duplicate=False`` restores every path so the canonical name is
246 visible, and picking the fewest-dot-separated-segments path per object
247 (breaking a tie between two paths that both happen to end in a mount name,
248 such as ``mlp``, which HF's own attribute name also frequently matches)
249 reports the shallower, canonical-looking path rather than an internal one.
250 """
251 best_by_id: Dict[int, Tuple[str, _RelevanceRuleCapable]] = {}
252 for name, module in model.named_modules(remove_duplicate=False):
253 if name.rsplit(".", 1)[-1] not in mount_names:
254 continue
255 existing = best_by_id.get(id(module))
256 if existing is None or name.count(".") < existing[0].count("."): 256 ↛ 252line 256 didn't jump to line 252 because the condition on line 256 was always true
257 best_by_id[id(module)] = (name, module)
258 yield from best_by_id.values()
261def _acquire_rule(module: _RelevanceRuleCapable, kind: str) -> None:
262 """Enable ``module``'s ``kind`` rule only on the outermost scope that requests it.
264 Refcounted per kind, not per module: a gated-MLP node can have its
265 "activation" rule and "multiplicative_gate" rule independently nested to
266 different depths, so one kind's inner exit must never disable the other.
267 """
268 counts: Dict[str, int] = getattr(module, "_relevance_rule_refcounts", None) or {}
269 count = counts.get(kind, 0)
270 if count == 0:
271 module._enable_relevance_rule(kind)
272 counts[kind] = count + 1
273 setattr(module, "_relevance_rule_refcounts", counts)
276def _release_rule(module: _RelevanceRuleCapable, kind: str) -> None:
277 """Disable ``module``'s ``kind`` rule only once its innermost scope exits."""
278 counts: Dict[str, int] = getattr(module, "_relevance_rule_refcounts", None) or {}
279 count = counts.get(kind, 0) - 1
280 counts[kind] = max(count, 0)
281 setattr(module, "_relevance_rule_refcounts", counts)
282 if count <= 0:
283 module._disable_relevance_rule(kind)
286@contextmanager
287def use_relevance_rules(model: nn.Module, rules: RelevanceRules) -> Iterator[RelevanceRuleCoverage]:
288 """Install the requested relevance rules on ``model`` only for this scope.
290 Targeting is positional: a component is considered for a rule kind only when it
291 sits at that kind's canonical mount name (never by class). A canonical mount
292 occupied by a component that does not implement ``_RelevanceRuleCapable``, or
293 whose ``_relevance_rule_kinds`` simply omits the requested kind, is reported as
294 skipped -- both are benign non-applicability, covering a structurally different
295 architecture or a mount whose current dispatch path the rule does not wrap. A
296 component that additionally names the requested kind in its own
297 ``_relevance_rule_unsupported_kinds`` raises ``RelevanceRuleUnsupportedError``
298 instead: that names a kind the component is expected to honor at this mount but
299 cannot given its current configuration, so silently skipping it would let
300 analysis proceed as if the caller had never asked. Requesting a kind that has no
301 canonical mount at all (for example "attention", a valid ``RelevanceRules`` field
302 with no mount defined) raises ``ValueError`` before the install loop, since there
303 is no mount to target and the scope would otherwise return empty coverage as
304 though the request had succeeded. Scopes over the same model
305 are reference-counted, so an inner scope's exit never disables a rule an outer
306 scope still needs. No model configuration is mutated; the only state that
307 changes lives on the participating components, and only for the scope's
308 duration.
309 """
310 requested_kinds = [
311 field.name for field in dataclasses.fields(rules) if getattr(rules, field.name)
312 ]
314 unmapped_kinds = [kind for kind in requested_kinds if kind not in _CANONICAL_MOUNTS]
315 if unmapped_kinds:
316 joined = ", ".join(repr(kind) for kind in unmapped_kinds)
317 raise ValueError(
318 f"No canonical mount is defined for relevance-rule kind(s) {joined}. "
319 "Such a kind has no mount to target, so it would install nothing and "
320 "silently return an empty coverage report instead of applying the rule."
321 )
323 installed: List[Tuple[str, _RelevanceRuleCapable, str]] = []
324 skipped: List[str] = []
325 for kind in requested_kinds:
326 mount_names = _CANONICAL_MOUNTS.get(kind, ())
327 for name, module in _iter_canonical_mount_candidates(model, mount_names):
328 if isinstance(module, _RelevanceRuleCapable) and kind in module._relevance_rule_kinds:
329 installed.append((name, module, kind))
330 continue
331 unsupported_kinds = getattr(module, "_relevance_rule_unsupported_kinds", ())
332 if isinstance(module, _RelevanceRuleCapable) and kind in unsupported_kinds:
333 raise RelevanceRuleUnsupportedError(
334 f"{name!r} ({type(module).__name__}) cannot install the {kind!r} "
335 "relevance rule: unsupported configuration for this component."
336 )
337 skipped.append(name)
339 for _, module, kind in installed:
340 _acquire_rule(module, kind)
341 try:
342 yield RelevanceRuleCoverage(
343 installed=tuple(name for name, _, _ in installed),
344 skipped=tuple(skipped),
345 )
346 finally:
347 for _, module, kind in installed:
348 _release_rule(module, kind)