Coverage for transformer_lens/model_bridge/generalized_components/mlp.py: 87%

1"""MLP bridge component. 

2 

3This module contains the bridge component for MLP layers. 

4""" 

5from typing import Any, Dict, Optional 

6 

7import torch 

8 

9from transformer_lens.model_bridge.generalized_components.base import ( 

10 GeneralizedComponent, 

11) 

12from transformer_lens.utilities.quantization import require_readable_weight 

13 

14 

15def weight_layout_in_out(proj: Any) -> Optional[bool]: 

16 """Whether proj's wrapped module stores its weight as [in, out]. 

17 

18 Conv1D (GPT-2 style) stores [in_features, out_features]; nn.Linear stores 

19 [out_features, in_features]. Returns None when the wrapped module is 

20 neither, so callers fall back to in_features/out_features or a shape heuristic. 

21 """ 

22 from transformers.pytorch_utils import Conv1D 

23 

24 component = getattr(proj, "original_component", None) 

25 if isinstance(component, Conv1D): 

26 return True 

27 if isinstance(component, torch.nn.Linear): 

28 return False 

29 return None 

30 

31 

32def normalize_mlp_weight( 

33 weight: torch.Tensor, layout: Optional[bool], proj: Any, pattern: str = "in" 

34) -> torch.Tensor: 

35 """Normalize an MLP projection weight to TL orientation ([d_model, d_mlp] 

36 for "in"/W_gate, [d_mlp, d_model] for "out").""" 

37 if layout is None: 

38 component = getattr(proj, "original_component", None) 

39 in_f = getattr(component, "in_features", None) 

40 out_f = getattr(component, "out_features", None) 

41 if in_f is not None and out_f is not None: 41 ↛ 46line 41 didn't jump to line 46 because the condition on line 41 was always true

42 layout = weight.shape[0] == in_f 

43 else: 

44 # Last resort. WARNING: assumes d_model < d_mlp (false for GIDD's 

45 # ScaledLinear, which is why in_features/out_features go first). 

46 if pattern == "in": 

47 layout = weight.shape[0] < weight.shape[1] 

48 else: 

49 layout = weight.shape[0] > weight.shape[1] 

50 if layout: 

51 return weight # Conv1D-style: already in TL orientation 

52 return weight.T # nn.Linear-style: transpose to TL orientation 

53 

54 

55class MLPBridge(GeneralizedComponent): 

56 """Bridge component for MLP layers. 

57 

58 This component wraps an MLP layer from a remote model and provides a consistent interface 

59 for accessing its weights and performing MLP operations. 

60 """ 

61 

62 hook_aliases = {"hook_pre": "in.hook_out", "hook_post": "out.hook_in"} 

63 # Containerless (name=None) instances refuse forward, so hook_in/hook_out 

64 # must be mirrored from the in/out subcomponents at setup. 

65 mirror_placeholder_hooks = True 

66 # W_* are real properties below (layout-aware); only the 1-D biases are 

67 # orientation-free enough for raw passthrough aliases. 

68 property_aliases = { 

69 "b_gate": "gate.bias", 

70 "b_in": "in.bias", 

71 "b_out": "out.bias", 

72 } 

73 

74 def __init__( 

75 self, 

76 name: Optional[str], 

77 config: Optional[Any] = None, 

78 submodules: Optional[Dict[str, GeneralizedComponent]] = {}, 

79 optional: bool = False, 

80 ): 

81 """Initialize the MLP bridge. 

82 

83 Args: 

84 name: The name of the component in the model (None if no container exists) 

85 config: Optional configuration (unused for MLPBridge) 

86 submodules: Dictionary of submodules to register (e.g., gate_proj, up_proj, down_proj) 

87 optional: If True, setup skips this bridge when absent (hybrid architectures). 

88 """ 

89 super().__init__(name, config, submodules=submodules, optional=optional) 

90 

91 def forward(self, *args, **kwargs) -> Any: 

92 """Forward pass through the MLP bridge. 

93 

94 Returns a tensor, or the component's own (hidden, ...) tuple re-packed 

95 with hooked hidden states for recurrent MLPs. 

96 

97 Args: 

98 *args: Positional arguments for the original component 

99 **kwargs: Keyword arguments for the original component 

100 

101 Returns: 

102 Output hidden states 

103 """ 

104 if self.name is None: 

105 # Containerless (fc1/fc2 sit on the decoder layer): the PARENT 

106 # LAYER is bound as original_component, so delegating would 

107 # silently run the whole layer — attention included. 

108 raise RuntimeError( 

109 f"{type(self).__name__} is containerless (name=None) — calling it " 

110 "directly would execute the whole parent layer. Call the block, " 

111 "or use the in/out submodules and their hooks." 

112 ) 

113 hidden_states = args[0] 

114 hidden_states = self.hook_in(hidden_states) 

115 in_module = getattr(self, "in", None) or getattr(self, "input", None) 

116 if in_module is not None and not hasattr(in_module, "hook_in"): 116 ↛ 117line 116 didn't jump to line 117 because the condition on line 116 was never true

117 in_module = None 

118 if in_module is not None: 

119 hidden_states = in_module.hook_in(hidden_states) 

120 new_args = (hidden_states,) + args[1:] 

121 original_component = self.original_component 

122 if original_component is None: 122 ↛ 123line 122 didn't jump to line 123 because the condition on line 122 was never true

123 raise RuntimeError( 

124 f"Original component not set for {self.name}. Call set_original_component() first." 

125 ) 

126 out_module = getattr(self, "out", None) 

127 if out_module is not None: 

128 object.__setattr__(out_module, "_fired_hook_out", False) 

129 # The pre-fire above already hooked the tensor entering the module; tell 

130 # the replaced `in` projection to skip its own hook_in once so the same 

131 # tensor is not double-hooked (wrapped forwards that bypass the 

132 # projection leave the flag set; it is cleared below). 

133 if in_module is not None: 

134 object.__setattr__(in_module, "_suppress_next_hook_in", True) 

135 try: 

136 output = original_component(*new_args, **kwargs) 

137 finally: 

138 if in_module is not None: 

139 object.__setattr__(in_module, "_suppress_next_hook_in", False) 

140 # Recurrent MLPs (RWKV's channel-mix) return (hidden, state). Hook the 

141 # hidden states and re-pack, or hook_out would hand users a tuple and 

142 # interventions on it would be silently dropped. 

143 if isinstance(output, tuple): 

144 # Only the block-level hook: the wrapped `out` projection already 

145 # fired its own hook_out inside the forward, and for gated recurrent 

146 # MLPs (RWKV channel-mix) output[0] is the post-gate product — a 

147 # different tensor, so re-firing would double-apply interventions. 

148 return (self.hook_out(output[0]),) + output[1:] 

149 output = self.hook_out(output) 

150 # Fallback only, for wrapped forwards that bypass the replaced `out` 

151 # projection: if it did run, re-firing would double-apply interventions 

152 # and, on residual-inside MLPs (MPT), stamp the residual-added module 

153 # output over the additive contribution. 

154 if out_module is not None and not out_module._fired_hook_out: 154 ↛ 155line 154 didn't jump to line 155 because the condition on line 154 was never true

155 output = out_module.hook_out(output) 

156 return output 

157 

158 def _weight_layout_in_out(self, proj: Any) -> Optional[bool]: 

159 """Whether proj's wrapped module stores its weight as [in, out].""" 

160 return weight_layout_in_out(proj) 

161 

162 def _normalize_mlp_weight( 

163 self, weight: torch.Tensor, layout: Optional[bool], proj: Any, pattern: str = "in" 

164 ) -> torch.Tensor: 

165 """Normalize MLP weight to TL orientation.""" 

166 return normalize_mlp_weight(weight, layout, proj, pattern=pattern) 

167 

168 @property 

169 def W_in(self) -> torch.Tensor: 

170 """MLP input weight in TL orientation [d_model, d_mlp].""" 

171 in_module = getattr(self, "in", None) 

172 if in_module is None: 172 ↛ 173line 172 didn't jump to line 173 because the condition on line 172 was never true

173 raise AttributeError("No 'in' submodule on this MLP bridge") 

174 weight = require_readable_weight( 

175 in_module.weight, operation=f"read W_in from {self.name}", owner=in_module 

176 ) 

177 layout = self._weight_layout_in_out(in_module) 

178 return self._normalize_mlp_weight(weight, layout, in_module, pattern="in") 

179 

180 @property 

181 def W_gate(self) -> Optional[torch.Tensor]: 

182 """MLP gate weight in TL orientation [d_model, d_mlp], or None if ungated.""" 

183 gate_module = getattr(self, "gate", None) 

184 if gate_module is None: 

185 return None 

186 weight = require_readable_weight( 

187 gate_module.weight, operation=f"read W_gate from {self.name}", owner=gate_module 

188 ) 

189 layout = self._weight_layout_in_out(gate_module) 

190 return self._normalize_mlp_weight(weight, layout, gate_module, pattern="in") 

191 

192 @property 

193 def W_out(self) -> torch.Tensor: 

194 """MLP output weight in TL orientation [d_mlp, d_model].""" 

195 out_module = getattr(self, "out", None) 

196 if out_module is None: 196 ↛ 197line 196 didn't jump to line 197 because the condition on line 196 was never true

197 raise AttributeError("No 'out' submodule on this MLP bridge") 

198 weight = require_readable_weight( 

199 out_module.weight, operation=f"read W_out from {self.name}", owner=out_module 

200 ) 

201 layout = self._weight_layout_in_out(out_module) 

202 return self._normalize_mlp_weight(weight, layout, out_module, pattern="out")