Coverage for transformer_lens/model_bridge/generalized

1"""MLP bridge component.

3This module contains the bridge component for MLP layers.

4"""

5from typing import Any, Dict, Optional

7import torch

9from transformer_lens.model_bridge.generalized_components.base import (

10 GeneralizedComponent,

11)

12from transformer_lens.utilities.quantization import require_readable_weight

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

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

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

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

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

55class MLPBridge(GeneralizedComponent):

56 """Bridge component for MLP layers.

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 """

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

63 property_aliases = {

64 "b_gate": "gate.bias",

65 "b_in": "in.bias",

66 "b_out": "out.bias",

67 }

69 def __init__(

70 self,

71 name: Optional[str],

72 config: Optional[Any] = None,

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

74 optional: bool = False,

75 ):

76 """Initialize the MLP bridge.

78 Args:

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

80 config: Optional configuration (unused for MLPBridge)

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

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

83 """

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

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

87 """Forward pass through the MLP bridge.

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

90 with hooked hidden states for recurrent MLPs.

92 Args:

93 *args: Positional arguments for the original component

94 **kwargs: Keyword arguments for the original component

96 Returns:

97 Output hidden states

98 """

99 if self.name is None:

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

101 # LAYER is bound as original_component, so delegating would

102 # silently run the whole layer — attention included.

103 raise RuntimeError(

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

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

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

107 )

108 hidden_states = args[0]

109 hidden_states = self.hook_in(hidden_states)

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

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

112 in_module = None

113 if in_module is not None:

114 hidden_states = in_module.hook_in(hidden_states)

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

116 original_component = self.original_component

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

118 raise RuntimeError(

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

120 )

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

122 if out_module is not None:

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

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

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

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

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

128 if in_module is not None:

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

130 try:

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

132 finally:

133 if in_module is not None:

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

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

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

137 # interventions on it would be silently dropped.

138 if isinstance(output, tuple):

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

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

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

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

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

144 output = self.hook_out(output)

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

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

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

148 # output over the additive contribution.

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

150 output = out_module.hook_out(output)

151 return output

152

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

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

155 return weight_layout_in_out(proj)

156

157 def _normalize_mlp_weight(

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

159 ) -> torch.Tensor:

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

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

162

163 @property

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

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

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

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

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

169 weight = require_readable_weight(

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

171 )

172 layout = self._weight_layout_in_out(in_module)

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

174

175 @property

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

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

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

179 if gate_module is None:

180 return None

181 weight = require_readable_weight(

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

183 )

184 layout = self._weight_layout_in_out(gate_module)

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

186

187 @property

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

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

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

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

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

193 weight = require_readable_weight(

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

195 )

196 layout = self._weight_layout_in_out(out_module)

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

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

87 statements