Coverage for transformer_lens/model_bridge/generalized_components/joint_gate_up_mlp.py: 93%
85 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-09-01 16:23 +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.generalized_components.base import (
10 GeneralizedComponent,
11)
12from transformer_lens.model_bridge.generalized_components.gated_mlp import (
13 GatedMLPBridge,
14 resolve_activation_fn,
15)
16from transformer_lens.model_bridge.generalized_components.linear import LinearBridge
17from transformer_lens.utilities.quantization import require_readable_weight
20class JointGateUpMLPBridge(GatedMLPBridge):
21 """Bridge for MLPs with fused gate+up projections (e.g., Phi-3's gate_up_proj).
23 Splits the fused projection into separate LinearBridges and reconstructs
24 the gated MLP forward pass, allowing individual hook access to gate and up
25 activations. Follows the same pattern as JointQKVAttentionBridge for fused QKV.
27 Hook interface matches GatedMLPBridge: hook_pre (gate), hook_pre_linear (up),
28 hook_post (before down_proj).
29 """
31 def __init__(
32 self,
33 name: str,
34 config: Optional[Any] = None,
35 submodules: Optional[Dict[str, GeneralizedComponent]] = None,
36 split_gate_up_matrix: Optional[Callable] = None,
37 fused_attr: str = "gate_up_proj",
38 ):
39 super().__init__(name, config, submodules=submodules)
40 # HF names the fused module differently per family (gate_up_proj,
41 # input_linear, proj_1); parametrizing the default splitter keeps one
42 # guarded, bias-aware implementation instead of per-adapter copies.
43 self.fused_attr = fused_attr
44 self.split_gate_up_matrix = (
45 split_gate_up_matrix
46 if split_gate_up_matrix is not None
47 else self._default_split_gate_up
48 )
50 # Up projection registered as "in" to match GatedMLPBridge convention
51 # (hook_aliases, property_aliases, and weight keys all use "in").
52 self.gate = LinearBridge(name="gate")
53 _up_bridge = LinearBridge(name="in")
54 setattr(self, "in", _up_bridge) # "in" is a keyword; use setattr
56 self.submodules["gate"] = self.gate
57 self.submodules["in"] = _up_bridge
59 self.real_components["gate"] = ("gate", self.gate)
60 self.real_components["in"] = ("in", _up_bridge)
61 if hasattr(self, "out"): 61 ↛ 62line 61 didn't jump to line 62 because the condition on line 61 was never true
62 self.real_components["out"] = ("out", self.out)
64 # Typed as Any: HF exposes activation_fn as nn.Module (e.g. nn.SiLU)
65 self._activation_fn: Any = None
67 self._register_state_dict_hook(JointGateUpMLPBridge._filter_gate_up_state_dict)
68 self.register_load_state_dict_pre_hook(
69 JointGateUpMLPBridge._restore_filtered_gate_up_state_dict
70 )
72 @staticmethod
73 def _filter_gate_up_state_dict(
74 module: torch.nn.Module,
75 state_dict: Dict[str, Any],
76 prefix: str,
77 local_metadata: Dict[str, Any],
78 ) -> None:
79 """State dict hook that removes stale combined gate_up entries."""
80 gate_up_prefix = prefix + "gate_up."
81 keys_to_remove = [k for k in state_dict if k.startswith(gate_up_prefix)]
82 for k in keys_to_remove:
83 del state_dict[k]
85 @staticmethod
86 def _restore_filtered_gate_up_state_dict(
87 module: torch.nn.Module,
88 state_dict: Dict[str, Any],
89 prefix: str,
90 local_metadata: Dict[str, Any],
91 strict: bool,
92 missing_keys: list[str],
93 unexpected_keys: list[str],
94 error_msgs: list[str],
95 ) -> None:
96 """Insert current combined weights only to satisfy strict key matching.
98 Production checkpoints restore authoritative values through the unfiltered
99 Hugging Face ``_original_component`` path.
100 """
101 del local_metadata, strict, missing_keys, unexpected_keys, error_msgs
102 gate_up = module._modules.get("gate_up")
103 if gate_up is None: 103 ↛ 104line 103 didn't jump to line 104 because the condition on line 103 was never true
104 return
105 for key, value in gate_up.state_dict(prefix=f"{prefix}gate_up.").items():
106 state_dict.setdefault(key, value)
108 def _default_split_gate_up(
109 self,
110 original_mlp_component: Any,
111 ) -> tuple[torch.nn.Module, torch.nn.Module]:
112 """Split the fused [2*d_mlp, d_model] projection into (gate, up) Linears."""
113 fused_module = getattr(original_mlp_component, self.fused_attr)
114 # Guard before the split: float8 slices and survives nn.Parameter()
115 # silently, producing scale-less projections with no error anywhere.
116 fused_weight = require_readable_weight(
117 fused_module.weight,
118 operation="split a fused gate/up projection at boot",
119 owner=fused_module,
120 )
121 gate_w, up_w = torch.tensor_split(fused_weight, 2, dim=0)
122 d_model = fused_weight.shape[1]
123 d_mlp = gate_w.shape[0]
125 has_bias = getattr(fused_module, "bias", None) is not None
126 gate_b: torch.Tensor | None = None
127 up_b: torch.Tensor | None = None
128 if has_bias:
129 gate_b, up_b = torch.tensor_split(fused_module.bias, 2, dim=0)
131 # skip_init: these Linears exist only to carry the split views, so a
132 # kaiming init would be waste and would advance the global RNG.
133 gate_proj = torch.nn.utils.skip_init(torch.nn.Linear, d_model, d_mlp, bias=has_bias)
134 gate_proj.weight = torch.nn.Parameter(gate_w)
135 if gate_b is not None:
136 gate_proj.bias = torch.nn.Parameter(gate_b)
138 up_proj = torch.nn.utils.skip_init(torch.nn.Linear, d_model, d_mlp, bias=has_bias)
139 up_proj.weight = torch.nn.Parameter(up_w)
140 if up_b is not None:
141 up_proj.bias = torch.nn.Parameter(up_b)
143 return gate_proj, up_proj
145 def set_original_component(self, original_component: torch.nn.Module) -> None:
146 """Set the original MLP component and split fused projections."""
147 super().set_original_component(original_component)
149 gate_proj, up_proj = self.split_gate_up_matrix(original_component)
150 self.gate.set_original_component(gate_proj)
151 getattr(self, "in").set_original_component(up_proj)
153 # Capture activation function from original component. Attr names vary:
154 # activation_fn/act_fn (phi3, glm), act (OpenELM), activation
155 # (GraniteMoeHybrid shared_mlp). Missing one silently falls back to
156 # cfg.act_fn — which is how OpenELM ran ReLU instead of SiLU.
157 for attr in ("activation_fn", "act_fn", "act", "activation"): 157 ↛ exitline 157 didn't return from function 'set_original_component' because the loop on line 157 didn't complete
158 candidate = getattr(original_component, attr, None)
159 if callable(candidate):
160 self._activation_fn = candidate
161 break
163 def _resolve_activation_fn(self) -> Callable:
164 """Resolve the activation function for the reconstructed forward pass."""
165 if self._activation_fn is not None: 165 ↛ 167line 165 didn't jump to line 167 because the condition on line 165 was always true
166 return self._activation_fn
167 return resolve_activation_fn(self.config)
169 def forward(self, *args: Any, **kwargs: Any) -> torch.Tensor:
170 """Reconstructed gated MLP forward with individual hook access."""
171 # Delegate to GatedMLPBridge's processed-weights path only when ALL
172 # processed weights exist; its fallback bypasses intermediate hooks.
173 if (
174 hasattr(self, "_use_processed_weights")
175 and self._use_processed_weights
176 and hasattr(self, "_processed_W_gate")
177 and hasattr(self, "_processed_W_in")
178 ):
179 return super().forward(*args, **kwargs)
181 hidden_states = self.hook_in(args[0])
183 gate_output = self.gate(hidden_states)
184 up_output = getattr(self, "in")(hidden_states)
186 act_fn = self._resolve_activation_fn()
187 gated = act_fn(gate_output) * up_output
189 if hasattr(self, "out") and self.out is not None: 189 ↛ 192line 189 didn't jump to line 192 because the condition on line 189 was always true
190 output = self.out(gated)
191 else:
192 raise RuntimeError(
193 f"No 'out' (down_proj) submodule found in {self.__class__.__name__}. "
194 "Ensure 'out' is provided in submodules."
195 )
197 return self.hook_out(output)