Coverage for transformer_lens/model_bridge/supported_architectures/opt.py: 92%

1"""OPT architecture adapter.""" 

2 

3from typing import Any, Iterator 

4 

5import torch 

6 

7from transformer_lens.conversion_utils.conversion_steps import RearrangeTensorConversion 

8from transformer_lens.conversion_utils.conversion_steps.base_tensor_conversion import ( 

9 BaseTensorConversion, 

10) 

11from transformer_lens.conversion_utils.param_processing_conversion import ( 

12 ParamProcessingConversion, 

13) 

14from transformer_lens.model_bridge.architecture_adapter import ArchitectureAdapter 

15from transformer_lens.model_bridge.generalized_components import ( 

16 AttentionBridge, 

17 BlockBridge, 

18 EmbeddingBridge, 

19 LinearBridge, 

20 MLPBridge, 

21 NormalizationBridge, 

22 PosEmbedBridge, 

23 UnembeddingBridge, 

24) 

25 

26 

27class _UnflattenTokens(BaseTensorConversion): 

28 """Restore [batch, seq, d] on hooks inside OPT's flattened region. 

29 

30 OPTDecoderLayer reshapes hidden states to [batch*seq, d] before 

31 final_layer_norm/fc1/fc2 and back only after the residual add, so the MLP 

32 and ln2 hooks would otherwise fire 2D — silently wrong for 

33 position-indexed patching, crashing for `b s d` einops. The block stamps 

34 the live (batch, seq) at each forward entry; runs only while user hooks 

35 are attached. 

36 """ 

37 

38 def __init__(self) -> None: 

39 super().__init__() 

40 self.batch_seq: tuple[int, int] | None = None 

41 

42 def handle_conversion(self, input_value, *full_context): 

43 bs = self.batch_seq 

44 if ( 44 ↛ 51line 44 didn't jump to line 51 because the condition on line 44 was always true

45 bs is not None 

46 and isinstance(input_value, torch.Tensor) 

47 and input_value.dim() == 2 

48 and input_value.shape[0] == bs[0] * bs[1] 

49 ): 

50 return input_value.view(bs[0], bs[1], input_value.shape[-1]) 

51 return input_value 

52 

53 def revert(self, input_value, *full_context): 

54 bs = self.batch_seq 

55 if ( 55 ↛ 63line 55 didn't jump to line 63 because the condition on line 55 was always true

56 bs is not None 

57 and isinstance(input_value, torch.Tensor) 

58 and input_value.dim() == 3 

59 and tuple(input_value.shape[:2]) == bs 

60 ): 

61 # reshape (not view) — hooks may return non-contiguous tensors 

62 return input_value.reshape(-1, input_value.shape[-1]) 

63 return input_value 

64 

65 

66class _OptBlockBridge(BlockBridge): 

67 """BlockBridge that stamps (batch, seq) onto the unflatten conversions.""" 

68 

69 def __init__(self, *args: Any, **kwargs: Any) -> None: 

70 super().__init__(*args, **kwargs) 

71 # hook_mlp_in is the block's own HookPoint, fired from a ln2 pre-hook — 

72 # inside OPT's flattened region, so it needs the conversion too. 

73 self.hook_mlp_in.hook_conversion = _UnflattenTokens() 

74 

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

76 hidden = args[0] if args else kwargs.get("hidden_states") 

77 if isinstance(hidden, torch.Tensor) and hidden.dim() == 3: 77 ↛ 81line 77 didn't jump to line 81 because the condition on line 77 was always true

78 batch_seq = (hidden.shape[0], hidden.shape[1]) 

79 for conversion in self._unflatten_conversions(): 

80 conversion.batch_seq = batch_seq 

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

82 

83 def _unflatten_conversions(self) -> Iterator[_UnflattenTokens]: 

84 block_conversion = getattr(self.hook_mlp_in, "hook_conversion", None) 

85 if isinstance(block_conversion, _UnflattenTokens): 85 ↛ 87line 85 didn't jump to line 87 because the condition on line 85 was always true

86 yield block_conversion 

87 for key in ("mlp", "ln2"): 

88 component = self.submodules.get(key) 

89 if component is None: 89 ↛ 90line 89 didn't jump to line 90 because the condition on line 89 was never true

90 continue 

91 members = [component, *getattr(component, "submodules", {}).values()] 

92 for member in members: 

93 for hook_name in ("hook_in", "hook_out"): 

94 hook_point = getattr(member, hook_name, None) 

95 conversion = getattr(hook_point, "hook_conversion", None) 

96 if isinstance(conversion, _UnflattenTokens): 96 ↛ 93line 96 didn't jump to line 93 because the condition on line 96 was always true

97 yield conversion 

98 

99 

100class OptArchitectureAdapter(ArchitectureAdapter): 

101 """Architecture adapter for OPT models.""" 

102 

103 @staticmethod 

104 def _with_unflatten(component: Any) -> Any: 

105 """Attach _UnflattenTokens to the component's (and submodules') hooks.""" 

106 members = [component, *getattr(component, "submodules", {}).values()] 

107 for member in members: 

108 for hook_name in ("hook_in", "hook_out"): 

109 hook_point = getattr(member, hook_name, None) 

110 if hook_point is not None and hook_point.hook_conversion is None: 110 ↛ 108line 110 didn't jump to line 108 because the condition on line 110 was always true

111 hook_point.hook_conversion = _UnflattenTokens() 

112 return component 

113 

114 def __init__(self, cfg: Any) -> None: 

115 """Initialize the OPT architecture adapter.""" 

116 super().__init__(cfg) 

117 

118 # Set config variables for weight processing 

119 self.cfg.normalization_type = "LN" 

120 self.cfg.positional_embedding_type = "standard" 

121 self.cfg.final_rms = False 

122 self.cfg.gated_mlp = False 

123 self.cfg.attn_only = False 

124 

125 # OPT models were trained with BOS tokens (inherits default_prepend_bos = True) 

126 

127 # Post-norm: disable fold_ln and center_writing_weights (pre-norm only). 

128 is_post_norm = not getattr(self.cfg, "do_layer_norm_before", True) 

129 if is_post_norm: 

130 self.supports_fold_ln = False 

131 self.supports_center_writing_weights = False 

132 

133 self.weight_processing_conversions = { 

134 "blocks.{i}.attn.q.weight": ParamProcessingConversion( 

135 tensor_conversion=RearrangeTensorConversion("(n h) m -> n m h", n=self.cfg.n_heads), 

136 ), 

137 "blocks.{i}.attn.k.weight": ParamProcessingConversion( 

138 tensor_conversion=RearrangeTensorConversion("(n h) m -> n m h", n=self.cfg.n_heads), 

139 ), 

140 "blocks.{i}.attn.v.weight": ParamProcessingConversion( 

141 tensor_conversion=RearrangeTensorConversion("(n h) m -> n m h", n=self.cfg.n_heads), 

142 ), 

143 "blocks.{i}.attn.o.weight": ParamProcessingConversion( 

144 tensor_conversion=RearrangeTensorConversion("m (n h) -> n h m", n=self.cfg.n_heads), 

145 ), 

146 } 

147 

148 # OPT-350m is uniquely the only OPT size where word_embed_proj_dim (512) 

149 # != hidden_size (1024). It uses project_in/project_out linear layers 

150 # instead of a final_layer_norm. Detect this and conditionally include 

151 # ln_final only when the model actually has one. 

152 word_embed_proj_dim = getattr(self.cfg, "word_embed_proj_dim", self.cfg.d_model) 

153 has_final_layer_norm = word_embed_proj_dim == self.cfg.d_model 

154 

155 self.component_mapping = { 

156 "embed": EmbeddingBridge(name="model.decoder.embed_tokens"), 

157 "pos_embed": PosEmbedBridge(name="model.decoder.embed_positions"), 

158 "blocks": _OptBlockBridge( 

159 name="model.decoder.layers", 

160 # fc2 IS the mlp output (no container fires hook_out). No 

161 # hook_mlp_in override: pre-norm, the block already provides it. 

162 hook_alias_overrides={"hook_mlp_out": "mlp.out.hook_out"}, 

163 submodules={ 

164 "ln1": NormalizationBridge( 

165 name="self_attn_layer_norm", 

166 config=self.cfg, 

167 use_native_layernorm_autograd=True, 

168 ), 

169 "attn": AttentionBridge( 

170 name="self_attn", 

171 config=self.cfg, 

172 requires_attention_mask=True, # OPT requires attention_mask 

173 attention_mask_4d=True, # OPT expects 4D mask [batch, 1, tgt_len, src_len] 

174 submodules={ 

175 "q": LinearBridge(name="q_proj"), 

176 "k": LinearBridge(name="k_proj"), 

177 "v": LinearBridge(name="v_proj"), 

178 "o": LinearBridge(name="out_proj"), 

179 }, 

180 ), 

181 "ln2": self._with_unflatten( 

182 NormalizationBridge( 

183 name="final_layer_norm", 

184 config=self.cfg, 

185 use_native_layernorm_autograd=True, 

186 ) 

187 ), 

188 # Containerless fc1/fc2, as BERT. ln2/fc1/fc2 run inside 

189 # HF's [batch*seq, d] region — hence the unflatten wrap. 

190 "mlp": self._with_unflatten( 

191 MLPBridge( 

192 name=None, 

193 config=self.cfg, 

194 submodules={ 

195 "in": LinearBridge(name="fc1"), 

196 "out": LinearBridge(name="fc2"), 

197 }, 

198 ) 

199 ), 

200 }, 

201 ), 

202 "unembed": UnembeddingBridge(name="lm_head"), 

203 } 

204 if has_final_layer_norm: 

205 self.component_mapping["ln_final"] = NormalizationBridge( 

206 name="model.decoder.final_layer_norm", 

207 config=self.cfg, 

208 use_native_layernorm_autograd=True, 

209 ) 

210 # project_in/project_out bridge word_embed_proj_dim <-> hidden_size. 

211 if not has_final_layer_norm: 

212 self.component_mapping["project_in"] = LinearBridge(name="model.decoder.project_in") 

213 self.component_mapping["project_out"] = LinearBridge(name="model.decoder.project_out")