Coverage for transformer_lens/model_bridge/generalized_components/joint_gate_up_mlp.py: 94%
98 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-09-21 19:27 +0000
1"""Bridge component for MLP layers with fused gate+up projections (e.g., Phi-3)."""
2from __future__ import annotations
4from collections.abc import Callable
5from typing import Any, Dict, Optional
7import torch
9from transformer_lens.model_bridge._relevance_rules import half_rule, identity_rule
10from transformer_lens.model_bridge.generalized_components.base import (
11 GeneralizedComponent,
12 align_offloaded_subtree,
13)
14from transformer_lens.model_bridge.generalized_components.gated_mlp import (
15 GatedMLPBridge,
16 identity_rule_supports_activation,
17 resolve_activation_fn,
18)
19from transformer_lens.model_bridge.generalized_components.linear import LinearBridge
20from transformer_lens.utilities.quantization import require_readable_weight
23class JointGateUpMLPBridge(GatedMLPBridge):
24 """Bridge for MLPs with fused gate+up projections (e.g., Phi-3's gate_up_proj).
26 Splits the fused projection into separate LinearBridges and reconstructs
27 the gated MLP forward pass, allowing individual hook access to gate and up
28 activations. Follows the same pattern as JointQKVAttentionBridge for fused QKV.
30 Hook interface matches GatedMLPBridge: hook_pre (gate), hook_pre_linear (up),
31 hook_post (before down_proj).
32 """
34 def __init__(
35 self,
36 name: str,
37 config: Optional[Any] = None,
38 submodules: Optional[Dict[str, GeneralizedComponent]] = None,
39 split_gate_up_matrix: Optional[Callable] = None,
40 fused_attr: str = "gate_up_proj",
41 ):
42 super().__init__(name, config, submodules=submodules)
43 # HF names the fused module differently per family (gate_up_proj,
44 # input_linear, proj_1); parametrizing the default splitter keeps one
45 # guarded, bias-aware implementation instead of per-adapter copies.
46 self.fused_attr = fused_attr
47 self.split_gate_up_matrix = (
48 split_gate_up_matrix
49 if split_gate_up_matrix is not None
50 else self._default_split_gate_up
51 )
53 # Up projection registered as "in" to match GatedMLPBridge convention
54 # (hook_aliases, property_aliases, and weight keys all use "in").
55 self.gate = LinearBridge(name="gate")
56 _up_bridge = LinearBridge(name="in")
57 setattr(self, "in", _up_bridge) # "in" is a keyword; use setattr
59 self.submodules["gate"] = self.gate
60 self.submodules["in"] = _up_bridge
62 self.real_components["gate"] = ("gate", self.gate)
63 self.real_components["in"] = ("in", _up_bridge)
64 if hasattr(self, "out"): 64 ↛ 65line 64 didn't jump to line 65 because the condition on line 64 was never true
65 self.real_components["out"] = ("out", self.out)
67 # Typed as Any: HF exposes activation_fn as nn.Module (e.g. nn.SiLU)
68 self._activation_fn: Any = None
70 self._register_state_dict_hook(JointGateUpMLPBridge._filter_gate_up_state_dict)
71 self.register_load_state_dict_pre_hook(
72 JointGateUpMLPBridge._restore_filtered_gate_up_state_dict
73 )
75 @staticmethod
76 def _filter_gate_up_state_dict(
77 module: torch.nn.Module,
78 state_dict: Dict[str, Any],
79 prefix: str,
80 local_metadata: Dict[str, Any],
81 ) -> None:
82 """State dict hook that removes stale combined gate_up entries."""
83 gate_up_prefix = prefix + "gate_up."
84 keys_to_remove = [k for k in state_dict if k.startswith(gate_up_prefix)]
85 for k in keys_to_remove:
86 del state_dict[k]
88 @staticmethod
89 def _restore_filtered_gate_up_state_dict(
90 module: torch.nn.Module,
91 state_dict: Dict[str, Any],
92 prefix: str,
93 local_metadata: Dict[str, Any],
94 strict: bool,
95 missing_keys: list[str],
96 unexpected_keys: list[str],
97 error_msgs: list[str],
98 ) -> None:
99 """Insert current combined weights only to satisfy strict key matching.
101 Production checkpoints restore authoritative values through the unfiltered
102 Hugging Face ``_original_component`` path.
103 """
104 del local_metadata, strict, missing_keys, unexpected_keys, error_msgs
105 gate_up = module._modules.get("gate_up")
106 if gate_up is None:
107 return
108 for key, value in gate_up.state_dict(prefix=f"{prefix}gate_up.").items():
109 state_dict.setdefault(key, value)
111 def _default_split_gate_up(
112 self,
113 original_mlp_component: Any,
114 ) -> tuple[torch.nn.Module, torch.nn.Module]:
115 """Split the fused [2*d_mlp, d_model] projection into (gate, up) Linears."""
116 fused_module = getattr(original_mlp_component, self.fused_attr)
117 # Guard before the split: float8 slices and survives nn.Parameter()
118 # silently, producing scale-less projections with no error anywhere.
119 fused_weight = require_readable_weight(
120 fused_module.weight,
121 operation="split a fused gate/up projection at boot",
122 owner=fused_module,
123 )
124 gate_w, up_w = torch.tensor_split(fused_weight, 2, dim=0)
125 d_model = fused_weight.shape[1]
126 d_mlp = gate_w.shape[0]
128 has_bias = getattr(fused_module, "bias", None) is not None
129 gate_b: torch.Tensor | None = None
130 up_b: torch.Tensor | None = None
131 if has_bias:
132 gate_b, up_b = torch.tensor_split(fused_module.bias, 2, dim=0)
134 # skip_init: these Linears exist only to carry the split views, so a
135 # kaiming init would be waste and would advance the global RNG.
136 gate_proj = torch.nn.utils.skip_init(torch.nn.Linear, d_model, d_mlp, bias=has_bias)
137 gate_proj.weight = torch.nn.Parameter(gate_w)
138 if gate_b is not None:
139 gate_proj.bias = torch.nn.Parameter(gate_b)
141 up_proj = torch.nn.utils.skip_init(torch.nn.Linear, d_model, d_mlp, bias=has_bias)
142 up_proj.weight = torch.nn.Parameter(up_w)
143 if up_b is not None:
144 up_proj.bias = torch.nn.Parameter(up_b)
146 return gate_proj, up_proj
148 def set_original_component(self, original_component: torch.nn.Module) -> None:
149 """Set the original MLP component and split fused projections."""
150 super().set_original_component(original_component)
152 # Same setup-time raw-weight read as JointQKVAttentionBridge.set_original_component:
153 # split_gate_up_matrix reads original_component.gate_up_proj directly, once, to build
154 # independent gate/up slices - needs real (not meta) data materialized for that one read
155 # under Accelerate offload. See align_offloaded_subtree's docstring for why the whole
156 # subtree, not just original_component itself, needs materializing here.
157 with align_offloaded_subtree(original_component):
158 gate_proj, up_proj = self.split_gate_up_matrix(original_component)
159 self.gate.set_original_component(gate_proj)
160 getattr(self, "in").set_original_component(up_proj)
162 # Capture activation function from original component. Attr names vary:
163 # activation_fn/act_fn (phi3, glm), act (OpenELM), activation
164 # (GraniteMoeHybrid shared_mlp). Missing one silently falls back to
165 # cfg.act_fn — which is how OpenELM ran ReLU instead of SiLU.
166 for attr in ("activation_fn", "act_fn", "act", "activation"): 166 ↛ exitline 166 didn't return from function 'set_original_component' because the loop on line 166 didn't complete
167 candidate = getattr(original_component, attr, None)
168 if callable(candidate):
169 self._activation_fn = candidate
170 break
172 def _resolve_activation_fn(self) -> Callable:
173 """Resolve the activation function for the reconstructed forward pass."""
174 if self._activation_fn is not None: 174 ↛ 176line 174 didn't jump to line 176 because the condition on line 174 was always true
175 return self._activation_fn
176 return resolve_activation_fn(self.config)
178 def _activation_rule_installable(self) -> bool:
179 """The reconstructed forward calls the activation itself, so only the config
180 activation form gates the Identity-rule; the opaque-path requirement of a
181 wrappable activation callable does not apply here."""
182 return identity_rule_supports_activation(self.config)
184 # The reconstructed forward applies both rules inline off the boolean flags set
185 # by the base ``_enable_relevance_rule``/``_disable_relevance_rule``. The
186 # opaque-path installers must stay disabled: the gate hook lives on the shared
187 # down projection this forward also calls, so leaving it active would halve the
188 # gate*up gradient a second time on top of the inline ``half_rule``.
189 def _install_activation_rule(self) -> None:
190 return None
192 def _teardown_activation_rule(self) -> None:
193 return None
195 def _install_gate_rule(self) -> None:
196 return None
198 def _teardown_gate_rule(self) -> None:
199 return None
201 def forward(self, *args: Any, **kwargs: Any) -> torch.Tensor:
202 """Reconstructed gated MLP forward with individual hook access."""
203 # Delegate to GatedMLPBridge's processed-weights path only when ALL
204 # processed weights exist; its fallback bypasses intermediate hooks.
205 if (
206 hasattr(self, "_use_processed_weights")
207 and self._use_processed_weights
208 and hasattr(self, "_processed_W_gate")
209 and hasattr(self, "_processed_W_in")
210 ):
211 return super().forward(*args, **kwargs)
213 hidden_states = self.hook_in(args[0])
215 gate_output = self.gate(hidden_states)
216 up_output = getattr(self, "in")(hidden_states)
218 act_fn = self._resolve_activation_fn()
219 activated = (
220 identity_rule(gate_output, act_fn)
221 if self._relevance_rule_activation_active
222 else act_fn(gate_output)
223 )
224 gated = (
225 half_rule(activated, up_output)
226 if self._relevance_rule_gate_active
227 else activated * up_output
228 )
230 if hasattr(self, "out") and self.out is not None: 230 ↛ 233line 230 didn't jump to line 233 because the condition on line 230 was always true
231 output = self.out(gated)
232 else:
233 raise RuntimeError(
234 f"No 'out' (down_proj) submodule found in {self.__class__.__name__}. "
235 "Ensure 'out' is provided in submodules."
236 )
238 return self.hook_out(output)