Coverage for transformer_lens/model_bridge/generalized_components/gated_mlp.py: 81%
108 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-10-23 20:45 +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
7import torch
9from transformer_lens.model_bridge.generalized_components.base import (
10 GeneralizedComponent,
11)
12from transformer_lens.model_bridge.generalized_components.mlp import MLPBridge
15def _resolve_activation_fn_name(config: Any) -> Optional[str]:
16 """The raw activation-name attribute a config exposes, in adapter priority order."""
17 if config is None:
18 return None
19 for attr in ("activation_function", "hidden_activation", "hidden_act", "act_fn"): 19 ↛ 23line 19 didn't jump to line 23 because the loop on line 19 didn't complete
20 name = getattr(config, attr, None)
21 if name is not None:
22 return str(name)
23 return None
26def resolve_activation_fn(config: Any) -> Callable:
27 """Resolve activation function from a model config.
29 Checks config attributes in order: activation_function, hidden_activation,
30 hidden_act, act_fn. Maps common aliases to torch.nn.functional callables.
31 """
32 act_fn_name = _resolve_activation_fn_name(config)
34 if act_fn_name is None or act_fn_name in ("silu", "swish"):
35 return torch.nn.functional.silu
36 if act_fn_name == "gelu":
37 return torch.nn.functional.gelu
38 if act_fn_name in ("gelu_new", "gelu_pytorch_tanh"): 38 ↛ 40line 38 didn't jump to line 40 because the condition on line 38 was never true
40 def gelu_tanh(x: torch.Tensor) -> torch.Tensor:
41 return torch.nn.functional.gelu(x, approximate="tanh")
43 return gelu_tanh
44 if act_fn_name == "relu":
45 return torch.nn.functional.relu
46 if act_fn_name in ("relu2", "relu_2", "relu_squared"): 46 ↛ 48line 46 didn't jump to line 48 because the condition on line 46 was never true
48 def relu_squared(x: torch.Tensor) -> torch.Tensor:
49 return torch.nn.functional.relu(x).square()
51 return relu_squared
52 return torch.nn.functional.silu
55class GatedMLPBridge(MLPBridge):
56 """Bridge component for gated MLP layers.
58 This component wraps a gated MLP layer from a remote model (e.g., LLaMA, Gemma)
59 and provides a consistent interface for accessing its weights and performing MLP operations.
61 Gated MLPs have the structure:
62 output = down_proj(act_fn(gate_proj(x)) * up_proj(x))
64 Where:
65 - gate_proj: The gating projection (produces the activation to be gated)
66 - up_proj (in): The input projection (produces the linear component)
67 - down_proj (out): The output projection
68 """
70 hook_aliases = {
71 "hook_pre": "gate.hook_out",
72 "hook_pre_linear": "in.hook_out",
73 "hook_post": "out.hook_in",
74 }
75 # property_aliases inherited from MLPBridge (W_gate, b_gate, W_in, b_in, W_out, b_out)
77 def __init__(
78 self,
79 name: Optional[str],
80 config: Optional[Any] = None,
81 submodules: Optional[Dict[str, GeneralizedComponent]] = None,
82 optional: bool = False,
83 ):
84 """Initialize the gated MLP bridge.
86 Args:
87 name: The name of the component in the model (None if no container exists)
88 config: Optional configuration (unused for GatedMLPBridge)
89 submodules: Dictionary of submodules to register (e.g., gate_proj, up_proj, down_proj)
90 optional: If True, setup skips this bridge when absent (hybrid architectures).
91 """
92 super().__init__(name, config, submodules=submodules or {}, optional=optional)
94 def forward(self, *args, **kwargs) -> torch.Tensor:
95 """Forward pass through the gated MLP bridge.
97 Intermediate hooks (gate.hook_out, in.hook_out, out.hook_in) only fire in
98 compatibility mode with processed weights enabled. In non-compatibility mode,
99 the HF component is called as an opaque forward and only hook_in/hook_out fire.
101 Args:
102 *args: Positional arguments for the original component
103 **kwargs: Keyword arguments for the original component
105 Returns:
106 Output hidden states
107 """
108 if hasattr(self, "_use_processed_weights") and self._use_processed_weights:
109 W_gate, b_gate = self._processed_projection("gate")
110 W_in, b_in = self._processed_projection("in")
111 W_out, b_out = self._processed_projection("out")
112 hidden_states = args[0]
113 hidden_states = self.hook_in(hidden_states)
114 gate_output = torch.nn.functional.linear(hidden_states, W_gate, b_gate)
115 if hasattr(self, "gate") and hasattr(self.gate, "hook_out"): 115 ↛ 117line 115 didn't jump to line 117 because the condition on line 115 was always true
116 gate_output = self.gate.hook_out(gate_output)
117 linear_output = torch.nn.functional.linear(hidden_states, W_in, b_in)
118 in_module = getattr(self, "in", None)
119 if in_module is not None and hasattr(in_module, "hook_out"): 119 ↛ 121line 119 didn't jump to line 121 because the condition on line 119 was always true
120 linear_output = in_module.hook_out(linear_output) # type: ignore[misc]
121 act_fn = resolve_activation_fn(self.config)
122 activated = act_fn(gate_output)
123 hidden = activated * linear_output
124 if hasattr(self, "out") and hasattr(self.out, "hook_in"): 124 ↛ 126line 124 didn't jump to line 126 because the condition on line 124 was always true
125 hidden = self.out.hook_in(hidden)
126 output = torch.nn.functional.linear(hidden, W_out, b_out)
127 # The functional path bypasses the wrapped `out` projection, so fire
128 # its hook_out here — it is the down-projection output.
129 if hasattr(self, "out") and hasattr(self.out, "hook_out"): 129 ↛ 131line 129 didn't jump to line 131 because the condition on line 129 was always true
130 output = self.out.hook_out(output)
131 output = self.hook_out(output)
132 return output
133 if self.original_component is None: 133 ↛ 134line 133 didn't jump to line 134 because the condition on line 133 was never true
134 raise RuntimeError(
135 f"Original component not set for {self.name}. Call set_original_component() first."
136 )
137 hidden_states = args[0]
138 hidden_states = self.hook_in(hidden_states)
139 new_args = (hidden_states,) + args[1:]
140 output = self.original_component(*new_args, **kwargs)
141 output = self.hook_out(output)
142 return output
144 def _processed_projection(self, name: str) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
145 """Current weight and bias of a projection, still attached to autograd so they train."""
146 projection = getattr(self, name)
147 return projection.weight, projection.bias
149 def set_processed_weights(
150 self, weights: Mapping[str, torch.Tensor | None], verbose: bool = False
151 ) -> None:
152 """Set the processed weights to use when layer norm is folded.
154 Args:
155 W_gate: The processed MLP gate weight tensor
156 W_in: The processed MLP input weight tensor
157 W_out: The processed MLP output weight tensor
158 b_gate: The processed MLP gate bias tensor (optional)
159 b_in: The processed MLP input bias tensor (optional)
160 b_out: The processed MLP output bias tensor (optional)
161 verbose: If True, print detailed information about weight setting
162 """
163 if verbose: 163 ↛ 164line 163 didn't jump to line 164 because the condition on line 163 was never true
164 print(
165 f"\n set_processed_weights: GatedMLPBridge (name={getattr(self, 'name', 'unknown')})"
166 )
167 print(f" Received {len(weights)} weight keys")
169 super().set_processed_weights(weights, verbose=verbose) # type: ignore[arg-type]
170 W_gate = weights.get("gate.weight")
171 if W_gate is None: 171 ↛ 172line 171 didn't jump to line 172 because the condition on line 171 was never true
172 return
173 b_gate = weights.get("gate.bias")
175 W_in = weights.get("in.weight")
176 b_in = weights.get("in.bias")
177 W_out = weights.get("out.weight")
178 b_out = weights.get("out.bias")
180 if verbose: 180 ↛ 181line 180 didn't jump to line 181 because the condition on line 180 was never true
181 print(f" Setting W_gate with shape: {W_gate.shape}")
182 if b_gate is not None:
183 print(f" Setting b_gate with shape: {b_gate.shape}")
184 if W_in is not None:
185 print(f" Setting W_in with shape: {W_in.shape}")
186 if W_out is not None:
187 print(f" Setting W_out with shape: {W_out.shape}")
189 self._use_processed_weights = True
191 # Distribute to submodules if they support it
192 gate_module = getattr(self, "gate", None)
193 if gate_module and hasattr(gate_module, "set_processed_weights"): 193 ↛ 199line 193 didn't jump to line 199 because the condition on line 193 was always true
194 gate_weights: Dict[str, torch.Tensor] = {"weight": W_gate}
195 if b_gate is not None:
196 gate_weights["bias"] = b_gate
197 gate_module.set_processed_weights(gate_weights, verbose=verbose)
199 in_module = getattr(self, "in", None)
200 if in_module and hasattr(in_module, "set_processed_weights") and W_in is not None:
201 in_weights: Dict[str, torch.Tensor] = {"weight": W_in}
202 if b_in is not None:
203 in_weights["bias"] = b_in
204 in_module.set_processed_weights(in_weights, verbose=verbose)
206 out_module = getattr(self, "out", None)
207 if out_module and hasattr(out_module, "set_processed_weights") and W_out is not None:
208 out_weights: Dict[str, torch.Tensor] = {"weight": W_out}
209 if b_out is not None:
210 out_weights["bias"] = b_out
211 out_module.set_processed_weights(out_weights, verbose=verbose)