Coverage for transformer_lens/model_bridge/generalized_components/joint_gate_up_mlp.py: 94%

1"""Bridge component for MLP layers with fused gate+up projections (e.g., Phi-3).""" 

2from __future__ import annotations 

3 

4from collections.abc import Callable 

5from typing import Any, Dict, Optional 

6 

7import torch 

8 

9from transformer_lens.model_bridge.generalized_components.base import ( 

10 GeneralizedComponent, 

11 align_offloaded_subtree, 

12) 

13from transformer_lens.model_bridge.generalized_components.gated_mlp import ( 

14 GatedMLPBridge, 

15 resolve_activation_fn, 

16) 

17from transformer_lens.model_bridge.generalized_components.linear import LinearBridge 

18from transformer_lens.utilities.quantization import require_readable_weight 

19 

20 

21class JointGateUpMLPBridge(GatedMLPBridge): 

22 """Bridge for MLPs with fused gate+up projections (e.g., Phi-3's gate_up_proj). 

23 

24 Splits the fused projection into separate LinearBridges and reconstructs 

25 the gated MLP forward pass, allowing individual hook access to gate and up 

26 activations. Follows the same pattern as JointQKVAttentionBridge for fused QKV. 

27 

28 Hook interface matches GatedMLPBridge: hook_pre (gate), hook_pre_linear (up), 

29 hook_post (before down_proj). 

30 """ 

31 

32 def __init__( 

33 self, 

34 name: str, 

35 config: Optional[Any] = None, 

36 submodules: Optional[Dict[str, GeneralizedComponent]] = None, 

37 split_gate_up_matrix: Optional[Callable] = None, 

38 fused_attr: str = "gate_up_proj", 

39 ): 

40 super().__init__(name, config, submodules=submodules) 

41 # HF names the fused module differently per family (gate_up_proj, 

42 # input_linear, proj_1); parametrizing the default splitter keeps one 

43 # guarded, bias-aware implementation instead of per-adapter copies. 

44 self.fused_attr = fused_attr 

45 self.split_gate_up_matrix = ( 

46 split_gate_up_matrix 

47 if split_gate_up_matrix is not None 

48 else self._default_split_gate_up 

49 ) 

50 

51 # Up projection registered as "in" to match GatedMLPBridge convention 

52 # (hook_aliases, property_aliases, and weight keys all use "in"). 

53 self.gate = LinearBridge(name="gate") 

54 _up_bridge = LinearBridge(name="in") 

55 setattr(self, "in", _up_bridge) # "in" is a keyword; use setattr 

56 

57 self.submodules["gate"] = self.gate 

58 self.submodules["in"] = _up_bridge 

59 

60 self.real_components["gate"] = ("gate", self.gate) 

61 self.real_components["in"] = ("in", _up_bridge) 

62 if hasattr(self, "out"): 62 ↛ 63line 62 didn't jump to line 63 because the condition on line 62 was never true

63 self.real_components["out"] = ("out", self.out) 

64 

65 # Typed as Any: HF exposes activation_fn as nn.Module (e.g. nn.SiLU) 

66 self._activation_fn: Any = None 

67 

68 self._register_state_dict_hook(JointGateUpMLPBridge._filter_gate_up_state_dict) 

69 self.register_load_state_dict_pre_hook( 

70 JointGateUpMLPBridge._restore_filtered_gate_up_state_dict 

71 ) 

72 

73 @staticmethod 

74 def _filter_gate_up_state_dict( 

75 module: torch.nn.Module, 

76 state_dict: Dict[str, Any], 

77 prefix: str, 

78 local_metadata: Dict[str, Any], 

79 ) -> None: 

80 """State dict hook that removes stale combined gate_up entries.""" 

81 gate_up_prefix = prefix + "gate_up." 

82 keys_to_remove = [k for k in state_dict if k.startswith(gate_up_prefix)] 

83 for k in keys_to_remove: 

84 del state_dict[k] 

85 

86 @staticmethod 

87 def _restore_filtered_gate_up_state_dict( 

88 module: torch.nn.Module, 

89 state_dict: Dict[str, Any], 

90 prefix: str, 

91 local_metadata: Dict[str, Any], 

92 strict: bool, 

93 missing_keys: list[str], 

94 unexpected_keys: list[str], 

95 error_msgs: list[str], 

96 ) -> None: 

97 """Insert current combined weights only to satisfy strict key matching. 

98 

99 Production checkpoints restore authoritative values through the unfiltered 

100 Hugging Face ``_original_component`` path. 

101 """ 

102 del local_metadata, strict, missing_keys, unexpected_keys, error_msgs 

103 gate_up = module._modules.get("gate_up") 

104 if gate_up is None: 

105 return 

106 for key, value in gate_up.state_dict(prefix=f"{prefix}gate_up.").items(): 

107 state_dict.setdefault(key, value) 

108 

109 def _default_split_gate_up( 

110 self, 

111 original_mlp_component: Any, 

112 ) -> tuple[torch.nn.Module, torch.nn.Module]: 

113 """Split the fused [2*d_mlp, d_model] projection into (gate, up) Linears.""" 

114 fused_module = getattr(original_mlp_component, self.fused_attr) 

115 # Guard before the split: float8 slices and survives nn.Parameter() 

116 # silently, producing scale-less projections with no error anywhere. 

117 fused_weight = require_readable_weight( 

118 fused_module.weight, 

119 operation="split a fused gate/up projection at boot", 

120 owner=fused_module, 

121 ) 

122 gate_w, up_w = torch.tensor_split(fused_weight, 2, dim=0) 

123 d_model = fused_weight.shape[1] 

124 d_mlp = gate_w.shape[0] 

125 

126 has_bias = getattr(fused_module, "bias", None) is not None 

127 gate_b: torch.Tensor | None = None 

128 up_b: torch.Tensor | None = None 

129 if has_bias: 

130 gate_b, up_b = torch.tensor_split(fused_module.bias, 2, dim=0) 

131 

132 # skip_init: these Linears exist only to carry the split views, so a 

133 # kaiming init would be waste and would advance the global RNG. 

134 gate_proj = torch.nn.utils.skip_init(torch.nn.Linear, d_model, d_mlp, bias=has_bias) 

135 gate_proj.weight = torch.nn.Parameter(gate_w) 

136 if gate_b is not None: 

137 gate_proj.bias = torch.nn.Parameter(gate_b) 

138 

139 up_proj = torch.nn.utils.skip_init(torch.nn.Linear, d_model, d_mlp, bias=has_bias) 

140 up_proj.weight = torch.nn.Parameter(up_w) 

141 if up_b is not None: 

142 up_proj.bias = torch.nn.Parameter(up_b) 

143 

144 return gate_proj, up_proj 

145 

146 def set_original_component(self, original_component: torch.nn.Module) -> None: 

147 """Set the original MLP component and split fused projections.""" 

148 super().set_original_component(original_component) 

149 

150 # Same setup-time raw-weight read as JointQKVAttentionBridge.set_original_component: 

151 # split_gate_up_matrix reads original_component.gate_up_proj directly, once, to build 

152 # independent gate/up slices - needs real (not meta) data materialized for that one read 

153 # under Accelerate offload. See align_offloaded_subtree's docstring for why the whole 

154 # subtree, not just original_component itself, needs materializing here. 

155 with align_offloaded_subtree(original_component): 

156 gate_proj, up_proj = self.split_gate_up_matrix(original_component) 

157 self.gate.set_original_component(gate_proj) 

158 getattr(self, "in").set_original_component(up_proj) 

159 

160 # Capture activation function from original component. Attr names vary: 

161 # activation_fn/act_fn (phi3, glm), act (OpenELM), activation 

162 # (GraniteMoeHybrid shared_mlp). Missing one silently falls back to 

163 # cfg.act_fn — which is how OpenELM ran ReLU instead of SiLU. 

164 for attr in ("activation_fn", "act_fn", "act", "activation"): 164 ↛ exitline 164 didn't return from function 'set_original_component' because the loop on line 164 didn't complete

165 candidate = getattr(original_component, attr, None) 

166 if callable(candidate): 

167 self._activation_fn = candidate 

168 break 

169 

170 def _resolve_activation_fn(self) -> Callable: 

171 """Resolve the activation function for the reconstructed forward pass.""" 

172 if self._activation_fn is not None: 172 ↛ 174line 172 didn't jump to line 174 because the condition on line 172 was always true

173 return self._activation_fn 

174 return resolve_activation_fn(self.config) 

175 

176 def forward(self, *args: Any, **kwargs: Any) -> torch.Tensor: 

177 """Reconstructed gated MLP forward with individual hook access.""" 

178 if hasattr(self, "_use_processed_weights") and self._use_processed_weights: 

179 return super().forward(*args, **kwargs) 

180 

181 hidden_states = self.hook_in(args[0]) 

182 

183 gate_output = self.gate(hidden_states) 

184 up_output = getattr(self, "in")(hidden_states) 

185 

186 act_fn = self._resolve_activation_fn() 

187 gated = act_fn(gate_output) * up_output 

188 

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 ) 

196 

197 return self.hook_out(output)