Coverage for transformer_lens/model_bridge/generalized_components/gated_mlp.py: 79%
209 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-09-21 19:27 +0000
1"""Gated MLP bridge component.
3This module contains the bridge component for gated MLP layers (e.g., LLaMA, Gemma).
4"""
5from typing import Any, Callable, Dict, Mapping, Optional, Tuple, cast
7import torch
8import torch.nn as nn
10from transformer_lens.model_bridge._relevance_rules import (
11 half_rule,
12 identity_rule,
13 scale_gradient,
14)
15from transformer_lens.model_bridge.generalized_components.base import (
16 GeneralizedComponent,
17)
18from transformer_lens.model_bridge.generalized_components.mlp import MLPBridge
21def _resolve_activation_fn_name(config: Any) -> Optional[str]:
22 """The raw activation-name attribute a config exposes, in adapter priority order."""
23 if config is None:
24 return None
25 for attr in ("activation_function", "hidden_activation", "hidden_act", "act_fn"): 25 ↛ 29line 25 didn't jump to line 29 because the loop on line 25 didn't complete
26 name = getattr(config, attr, None)
27 if name is not None:
28 return str(name)
29 return None
32_IDENTITY_RULE_UNSUPPORTED_ACTIVATIONS = {"relu", "relu2", "relu_2", "relu_squared"}
35def identity_rule_supports_activation(config: Any) -> bool:
36 """Whether the config's resolved activation form is safe for the Identity-rule.
38 The Identity-rule's backward multiplier is ``f(x) / x`` (the removable-singularity
39 limit filled in at zero) rather than the ordinary derivative -- the correct
40 LRP-style rule for SiLU and both GELU variants, but not for the relu family:
41 relu-squared's ratio reduces to ``relu(x)``, not its true derivative
42 ``2 * relu(x)``, and plain relu has no smooth two-sided derivative for the ratio
43 to represent at the removable singularity either. Both are therefore excluded
44 rather than silently applying a rule that does not hold for them.
45 """
46 return _resolve_activation_fn_name(config) not in _IDENTITY_RULE_UNSUPPORTED_ACTIVATIONS
49def resolve_activation_fn(config: Any) -> Callable:
50 """Resolve activation function from a model config.
52 Checks config attributes in order: activation_function, hidden_activation,
53 hidden_act, act_fn. Maps common aliases to torch.nn.functional callables.
54 """
55 act_fn_name = _resolve_activation_fn_name(config)
57 if act_fn_name is None or act_fn_name in ("silu", "swish"):
58 return torch.nn.functional.silu
59 if act_fn_name == "gelu":
60 return torch.nn.functional.gelu
61 if act_fn_name in ("gelu_new", "gelu_pytorch_tanh"): 61 ↛ 63line 61 didn't jump to line 63 because the condition on line 61 was never true
63 def gelu_tanh(x: torch.Tensor) -> torch.Tensor:
64 return torch.nn.functional.gelu(x, approximate="tanh")
66 return gelu_tanh
67 if act_fn_name == "relu":
68 return torch.nn.functional.relu
69 if act_fn_name in ("relu2", "relu_2", "relu_squared"): 69 ↛ 71line 69 didn't jump to line 71 because the condition on line 69 was never true
71 def relu_squared(x: torch.Tensor) -> torch.Tensor:
72 return torch.nn.functional.relu(x).square()
74 return relu_squared
75 return torch.nn.functional.silu
78class _IdentityRuleActivation(nn.Module):
79 """Route a wrapped activation through the Identity-rule for a scope's duration.
81 The opaque gated-MLP path keeps the HF module's own forward intact and installs
82 the Identity-rule by swapping the module's activation callable for this wrapper.
83 Its forward returns ``act_fn(x)`` unchanged, so the native forward value is
84 preserved, while the backward follows the Identity-rule VJP. Storing the wrapped
85 activation as an attribute registers it as a child module when it is itself an
86 ``nn.Module`` (the common ``ACT2FN`` case), so its parameters, if any, stay live.
87 """
89 def __init__(self, wrapped: Callable[[torch.Tensor], torch.Tensor]):
90 super().__init__()
91 self._wrapped_activation = wrapped
93 def forward(self, x: torch.Tensor) -> torch.Tensor:
94 return identity_rule(x, self._wrapped_activation)
97class GatedMLPBridge(MLPBridge):
98 """Bridge component for gated MLP layers.
100 This component wraps a gated MLP layer from a remote model (e.g., LLaMA, Gemma)
101 and provides a consistent interface for accessing its weights and performing MLP operations.
103 Gated MLPs have the structure:
104 output = down_proj(act_fn(gate_proj(x)) * up_proj(x))
106 Where:
107 - gate_proj: The gating projection (produces the activation to be gated)
108 - up_proj (in): The input projection (produces the linear component)
109 - down_proj (out): The output projection
110 """
112 hook_aliases = {
113 "hook_pre": "gate.hook_out",
114 "hook_pre_linear": "in.hook_out",
115 "hook_post": "out.hook_in",
116 }
117 # property_aliases inherited from MLPBridge (W_gate, b_gate, W_in, b_in, W_out, b_out)
119 def __init__(
120 self,
121 name: Optional[str],
122 config: Optional[Any] = None,
123 submodules: Optional[Dict[str, GeneralizedComponent]] = None,
124 optional: bool = False,
125 ):
126 """Initialize the gated MLP bridge.
128 Args:
129 name: The name of the component in the model (None if no container exists)
130 config: Optional configuration (unused for GatedMLPBridge)
131 submodules: Dictionary of submodules to register (e.g., gate_proj, up_proj, down_proj)
132 optional: If True, setup skips this bridge when absent (hybrid architectures).
133 """
134 super().__init__(name, config, submodules=submodules or {}, optional=optional)
135 self._relevance_rule_activation_active = False
136 self._relevance_rule_gate_active = False
137 # Opaque-path rule installers hold their teardown state here. The activation
138 # wrap records (attr_name, original_value, was_child_module) so the swapped
139 # activation callable can be restored exactly; the gate handle is the
140 # forward-pre-hook that scales the gradient entering the down projection.
141 self._relevance_activation_wrap: Optional[Tuple[str, Any, bool]] = None
142 self._relevance_gate_hook_handle: Optional[Any] = None
144 def _is_gated_mlp_shaped(self) -> bool:
145 """Whether this instance has the gate/up/down submodules a gated MLP needs.
147 A container missing one of these was never wired up as a gated-MLP node at
148 all (a different architecture at this mount), which is benign
149 non-applicability rather than an unsupported configuration of a gated-MLP
150 node -- unlike an activation form the Identity-rule cannot honor, which
151 occupies exactly this node's shape but cannot be honored correctly.
152 """
153 if self.original_component is None: 153 ↛ 154line 153 didn't jump to line 154 because the condition on line 153 was never true
154 return False
155 gate_module = getattr(self, "gate", None)
156 in_module = getattr(self, "in", None)
157 out_module = getattr(self, "out", None)
158 return gate_module is not None and in_module is not None and out_module is not None
160 def _find_activation_attr(self) -> Optional[str]:
161 """The attribute name under which the HF module holds its activation callable.
163 The opaque path installs the Identity-rule by swapping this attribute, so the
164 activation must be reachable as a callable attribute the native forward calls
165 (the ``ACT2FN`` module the gated-MLP families store as ``act_fn``). Returns
166 ``None`` when no such attribute exists, in which case the Identity-rule cannot
167 be wrapped in and ``"activation"`` is reported unsupported rather than
168 installed as a silent no-op.
169 """
170 component = self.original_component
171 if component is None: 171 ↛ 172line 171 didn't jump to line 172 because the condition on line 171 was never true
172 return None
173 for attr in ("act_fn", "activation_fn", "act", "activation"):
174 if callable(getattr(component, attr, None)):
175 return attr
176 return None
178 def _activation_rule_installable(self) -> bool:
179 """Whether the Identity-rule can be installed on this instance's activation.
181 Requires both a config activation form the ratio rule is valid for (the
182 relu family is excluded) and, on the opaque path, an activation callable the
183 bridge can wrap in place. Subclasses that reconstruct the forward themselves
184 override this, since they call the activation directly and never wrap it.
185 """
186 return (
187 identity_rule_supports_activation(self.config)
188 and self._find_activation_attr() is not None
189 )
191 @property
192 def _relevance_rule_kinds(self) -> Tuple[str, ...]:
193 """The relevance-rule kinds this instance can currently honor.
195 Empty when this is not a gated-MLP-shaped node. Otherwise always includes
196 ``"multiplicative_gate"`` (the Half-rule is a gradient scale at the gate*up
197 product and needs no weight or activation access) and includes
198 ``"activation"`` only when the Identity-rule can be installed on the
199 configured activation, so a relu-family activation, or one the bridge cannot
200 reach to wrap, is excluded.
201 """
202 if not self._is_gated_mlp_shaped(): 202 ↛ 203line 202 didn't jump to line 203 because the condition on line 202 was never true
203 return ()
204 kinds: Tuple[str, ...] = ("multiplicative_gate",)
205 if self._activation_rule_installable():
206 kinds = ("activation",) + kinds
207 return kinds
209 @property
210 def _relevance_rule_unsupported_kinds(self) -> Tuple[str, ...]:
211 """Kinds this gated-MLP node is expected to honor but currently cannot.
213 Unlike a kind simply absent from ``_relevance_rule_kinds`` because this is
214 not a gated-MLP-shaped node at all (benign non-applicability, reported
215 skipped), a gated-MLP node whose activation form or activation callable the
216 Identity-rule cannot honor is exactly the kind of component a caller expects
217 the rule to work on. Requesting ``"activation"`` there raises instead of
218 silently reporting the mount skipped. The Half-rule applies to every
219 gated-MLP node, so ``"multiplicative_gate"`` is never reported unsupported.
220 """
221 if not self._is_gated_mlp_shaped(): 221 ↛ 222line 221 didn't jump to line 222 because the condition on line 221 was never true
222 return ()
223 if "activation" not in self._relevance_rule_kinds: 223 ↛ 225line 223 didn't jump to line 225 because the condition on line 223 was always true
224 return ("activation",)
225 return ()
227 def _install_activation_rule(self) -> None:
228 """Swap the HF module's activation callable for the Identity-rule wrapper.
230 No-op when the activation callable cannot be located; requesting the
231 activation rule in that case is refused earlier through
232 ``_relevance_rule_unsupported_kinds``. The original value and whether it was
233 a registered child module are recorded so teardown restores it exactly.
234 """
235 component = self.original_component
236 attr = self._find_activation_attr()
237 if component is None or attr is None: 237 ↛ 238line 237 didn't jump to line 238 because the condition on line 237 was never true
238 return
239 was_child_module = attr in component._modules
240 original = component._modules[attr] if was_child_module else getattr(component, attr, None)
241 # _find_activation_attr only returns an attribute whose value is callable.
242 wrapper = _IdentityRuleActivation(cast(Callable[[torch.Tensor], torch.Tensor], original))
243 if not was_child_module: 243 ↛ 244line 243 didn't jump to line 244 because the condition on line 243 was never true
244 component.__dict__.pop(attr, None)
245 component._modules[attr] = wrapper
246 self._relevance_activation_wrap = (attr, original, was_child_module)
248 def _teardown_activation_rule(self) -> None:
249 """Restore the activation callable swapped in by ``_install_activation_rule``."""
250 if self._relevance_activation_wrap is None: 250 ↛ 251line 250 didn't jump to line 251 because the condition on line 250 was never true
251 return
252 attr, original, was_child_module = self._relevance_activation_wrap
253 component = self.original_component
254 if component is not None: 254 ↛ 260line 254 didn't jump to line 260 because the condition on line 254 was always true
255 component._modules.pop(attr, None)
256 if was_child_module: 256 ↛ 259line 256 didn't jump to line 259 because the condition on line 256 was always true
257 component._modules[attr] = original
258 else:
259 component.__dict__[attr] = original
260 self._relevance_activation_wrap = None
262 def _install_gate_rule(self) -> None:
263 """Halve the gradient entering the down projection to reproduce the Half-rule.
265 The gate*up product is the down projection's input, so a forward-pre-hook
266 that routes that input through ``scale_gradient(..., 0.5)`` halves the single
267 gradient feeding the product before it splits, which matches halving both
268 product-rule terms. The native forward value is unchanged, and the down
269 projection's own weight gradient stays ordinary because it is taken against
270 the unscaled downstream gradient.
271 """
272 out_module = getattr(self, "out", None)
273 down_component = getattr(out_module, "original_component", None)
274 if down_component is None: 274 ↛ anywhereline 274 didn't jump anywhere: it always raised an exception.
275 return
277 def _scale_product_gradient(
278 module: nn.Module, args: Tuple[Any, ...]
279 ) -> Optional[Tuple[Any, ...]]:
280 if not args: 280 ↛ 281line 280 didn't jump to line 281 because the condition on line 280 was never true
281 return None
282 return (scale_gradient(args[0], 0.5),) + tuple(args[1:])
284 self._relevance_gate_hook_handle = down_component.register_forward_pre_hook(
285 _scale_product_gradient
286 )
288 def _teardown_gate_rule(self) -> None:
289 """Remove the down-projection gradient-scale hook."""
290 if self._relevance_gate_hook_handle is not None: 290 ↛ exitline 290 didn't return from function '_teardown_gate_rule' because the condition on line 290 was always true
291 self._relevance_gate_hook_handle.remove()
292 self._relevance_gate_hook_handle = None
294 def _enable_relevance_rule(self, kind: str) -> None:
295 """Activate the named rule and install its opaque-path hook.
297 The boolean flag drives the reconstructed forward paths (compatibility mode
298 here, and the inline forward of subclasses that override it). The install
299 step additionally attaches the rule to the live HF submodules for the opaque
300 native forward, which a flag alone cannot alter.
301 """
302 if kind == "activation":
303 self._relevance_rule_activation_active = True
304 self._install_activation_rule()
305 elif kind == "multiplicative_gate": 305 ↛ exitline 305 didn't return from function '_enable_relevance_rule' because the condition on line 305 was always true
306 self._relevance_rule_gate_active = True
307 self._install_gate_rule()
309 def _disable_relevance_rule(self, kind: str) -> None:
310 """Deactivate the named rule and tear down its opaque-path hook."""
311 if kind == "activation":
312 self._teardown_activation_rule()
313 self._relevance_rule_activation_active = False
314 elif kind == "multiplicative_gate": 314 ↛ exitline 314 didn't return from function '_disable_relevance_rule' because the condition on line 314 was always true
315 self._teardown_gate_rule()
316 self._relevance_rule_gate_active = False
318 def forward(self, *args, **kwargs) -> torch.Tensor:
319 """Forward pass through the gated MLP bridge.
321 Intermediate hooks (gate.hook_out, in.hook_out, out.hook_in) only fire in
322 compatibility mode with processed weights enabled. In non-compatibility mode,
323 the HF component is called as an opaque forward and only hook_in/hook_out fire.
325 Args:
326 *args: Positional arguments for the original component
327 **kwargs: Keyword arguments for the original component
329 Returns:
330 Output hidden states
331 """
332 if hasattr(self, "_use_processed_weights") and self._use_processed_weights:
333 assert hasattr(self, "_processed_W_gate") and hasattr(self, "_processed_W_in"), (
334 "Processed weights flag is set but weights are missing. "
335 "This indicates a bug in set_processed_weights()."
336 )
337 assert self._processed_W_in is not None
338 assert self._processed_W_out is not None
339 hidden_states = args[0]
340 hidden_states = self.hook_in(hidden_states)
341 gate_output = torch.nn.functional.linear(
342 hidden_states, self._processed_W_gate, self._processed_b_gate
343 )
344 if hasattr(self, "gate") and hasattr(self.gate, "hook_out"): 344 ↛ 346line 344 didn't jump to line 346 because the condition on line 344 was always true
345 gate_output = self.gate.hook_out(gate_output)
346 linear_output = torch.nn.functional.linear(
347 hidden_states, self._processed_W_in, self._processed_b_in
348 )
349 in_module = getattr(self, "in", None)
350 if in_module is not None and hasattr(in_module, "hook_out"): 350 ↛ 352line 350 didn't jump to line 352 because the condition on line 350 was always true
351 linear_output = in_module.hook_out(linear_output) # type: ignore[misc]
352 act_fn = resolve_activation_fn(self.config)
353 activated = (
354 identity_rule(gate_output, act_fn)
355 if self._relevance_rule_activation_active
356 else act_fn(gate_output)
357 )
358 hidden = (
359 half_rule(activated, linear_output)
360 if self._relevance_rule_gate_active
361 else activated * linear_output
362 )
363 if hasattr(self, "out") and hasattr(self.out, "hook_in"): 363 ↛ 365line 363 didn't jump to line 365 because the condition on line 363 was always true
364 hidden = self.out.hook_in(hidden)
365 output = torch.nn.functional.linear(
366 hidden, self._processed_W_out, self._processed_b_out
367 )
368 # The functional path bypasses the wrapped `out` projection, so fire
369 # its hook_out here — it is the down-projection output.
370 if hasattr(self, "out") and hasattr(self.out, "hook_out"): 370 ↛ 372line 370 didn't jump to line 372 because the condition on line 370 was always true
371 output = self.out.hook_out(output)
372 output = self.hook_out(output)
373 return output
374 if self.original_component is None: 374 ↛ 375line 374 didn't jump to line 375 because the condition on line 374 was never true
375 raise RuntimeError(
376 f"Original component not set for {self.name}. Call set_original_component() first."
377 )
378 hidden_states = args[0]
379 hidden_states = self.hook_in(hidden_states)
380 new_args = (hidden_states,) + args[1:]
381 # The active relevance rules are attached to the live HF submodules (the
382 # activation callable and the down projection) by _enable_relevance_rule,
383 # so the native forward runs unchanged and its own internal hooks, gate
384 # multipliers, and activation sparsity all remain in the backward graph.
385 output = self.original_component(*new_args, **kwargs)
386 output = self.hook_out(output)
387 return output
389 def set_processed_weights(
390 self, weights: Mapping[str, torch.Tensor | None], verbose: bool = False
391 ) -> None:
392 """Set the processed weights to use when layer norm is folded.
394 Args:
395 W_gate: The processed MLP gate weight tensor
396 W_in: The processed MLP input weight tensor
397 W_out: The processed MLP output weight tensor
398 b_gate: The processed MLP gate bias tensor (optional)
399 b_in: The processed MLP input bias tensor (optional)
400 b_out: The processed MLP output bias tensor (optional)
401 verbose: If True, print detailed information about weight setting
402 """
403 if verbose: 403 ↛ 404line 403 didn't jump to line 404 because the condition on line 403 was never true
404 print(
405 f"\n set_processed_weights: GatedMLPBridge (name={getattr(self, 'name', 'unknown')})"
406 )
407 print(f" Received {len(weights)} weight keys")
409 super().set_processed_weights(weights, verbose=verbose) # type: ignore[arg-type]
410 W_gate = weights.get("gate.weight")
411 if W_gate is None: 411 ↛ 412line 411 didn't jump to line 412 because the condition on line 411 was never true
412 return
413 b_gate = weights.get("gate.bias")
415 W_in = weights.get("in.weight")
416 b_in = weights.get("in.bias")
417 W_out = weights.get("out.weight")
418 b_out = weights.get("out.bias")
420 if verbose: 420 ↛ 421line 420 didn't jump to line 421 because the condition on line 420 was never true
421 print(f" Setting W_gate with shape: {W_gate.shape}")
422 if b_gate is not None:
423 print(f" Setting b_gate with shape: {b_gate.shape}")
424 if W_in is not None:
425 print(f" Setting W_in with shape: {W_in.shape}")
426 if W_out is not None:
427 print(f" Setting W_out with shape: {W_out.shape}")
429 self._use_processed_weights = True
430 self._processed_W_gate = W_gate
431 self._processed_b_gate = b_gate
432 self._processed_W_in = W_in
433 self._processed_b_in = b_in
434 self._processed_W_out = W_out
435 self._processed_b_out = b_out
437 # Distribute to submodules if they support it
438 gate_module = getattr(self, "gate", None)
439 if gate_module and hasattr(gate_module, "set_processed_weights"): 439 ↛ 445line 439 didn't jump to line 445 because the condition on line 439 was always true
440 gate_weights: Dict[str, torch.Tensor] = {"weight": W_gate}
441 if b_gate is not None: 441 ↛ 442line 441 didn't jump to line 442 because the condition on line 441 was never true
442 gate_weights["bias"] = b_gate
443 gate_module.set_processed_weights(gate_weights, verbose=verbose)
445 in_module = getattr(self, "in", None)
446 if in_module and hasattr(in_module, "set_processed_weights") and W_in is not None:
447 in_weights: Dict[str, torch.Tensor] = {"weight": W_in}
448 if b_in is not None: 448 ↛ 449line 448 didn't jump to line 449 because the condition on line 448 was never true
449 in_weights["bias"] = b_in
450 in_module.set_processed_weights(in_weights, verbose=verbose)
452 out_module = getattr(self, "out", None)
453 if out_module and hasattr(out_module, "set_processed_weights") and W_out is not None:
454 out_weights: Dict[str, torch.Tensor] = {"weight": W_out}
455 if b_out is not None: 455 ↛ 456line 455 didn't jump to line 456 because the condition on line 455 was never true
456 out_weights["bias"] = b_out
457 out_module.set_processed_weights(out_weights, verbose=verbose)