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