Coverage for transformer_lens/model_bridge/supported_architectures/native.py: 99%

1"""Architecture adapter for TL-native models built via ``boot_native``. 

2 

3Component mapping adapts to cfg: gated MLP → ``GatedMLPBridge``, RMS norm → 

4``RMSNormalizationBridge``, param-free pre-norm (LNPre / RMSPre) → the 

5``*PreBridge`` pair, rotary drops ``pos_embed``, ``attn_only`` drops MLP. 

6""" 

7 

8from typing import Any 

9 

10from transformer_lens.model_bridge.architecture_adapter import ArchitectureAdapter 

11from transformer_lens.model_bridge.generalized_components import ( 

12 AttentionBridge, 

13 BlockBridge, 

14 EmbeddingBridge, 

15 GatedMLPBridge, 

16 LayerNormPreBridge, 

17 LinearBridge, 

18 MLPBridge, 

19 NormalizationBridge, 

20 PosEmbedBridge, 

21 RMSNormalizationBridge, 

22 RMSNormPreBridge, 

23 UnembeddingBridge, 

24) 

25from transformer_lens.model_bridge.generalized_components.base import ( 

26 GeneralizedComponent, 

27) 

28 

29 

30def _norm_type(cfg: Any) -> str: 

31 return (getattr(cfg, "normalization_type", None) or "LN").upper() 

32 

33 

34def _uses_rms(cfg: Any) -> bool: 

35 return _norm_type(cfg) in ("RMS", "RMSPRE") 

36 

37 

38def _uses_param_free_norm(cfg: Any) -> bool: 

39 return _norm_type(cfg) in ("RMSPRE", "LNPRE") 

40 

41 

42def _uses_no_norm(cfg: Any) -> bool: 

43 return getattr(cfg, "normalization_type", None) is None 

44 

45 

46def _is_rotary(cfg: Any) -> bool: 

47 return (getattr(cfg, "positional_embedding_type", None) or "standard").lower() == "rotary" 

48 

49 

50def _make_norm_bridge(name: str, cfg: Any, *, force_rms: bool = False): 

51 param_free = _uses_param_free_norm(cfg) 

52 if force_rms or _uses_rms(cfg): 

53 # Mirrors _make_norm: final_rms must not reintroduce a scale. 

54 if param_free: 

55 return RMSNormPreBridge(name=name, config=cfg) 

56 return RMSNormalizationBridge(name=name, config=cfg) 

57 if _norm_type(cfg) == "LNPRE": 

58 return LayerNormPreBridge(name=name, config=cfg) 

59 if _uses_no_norm(cfg): 

60 from transformer_lens.model_bridge.generalized_components.base import ( 

61 GeneralizedComponent, 

62 ) 

63 

64 return GeneralizedComponent(name=name, config=cfg) 

65 return NormalizationBridge(name=name, config=cfg) 

66 

67 

68def _make_mlp_bridge(cfg: Any): 

69 if cfg.gated_mlp: 

70 return GatedMLPBridge( 

71 name="mlp", 

72 config=cfg, 

73 submodules={ 

74 "gate": LinearBridge(name="gate"), 

75 "in": LinearBridge(name="in"), 

76 "out": LinearBridge(name="out"), 

77 }, 

78 ) 

79 submodules: dict[str, GeneralizedComponent] = { 

80 "in": LinearBridge(name="fc_in"), 

81 "out": LinearBridge(name="fc_out"), 

82 } 

83 if (cfg.act_fn or "").lower() == "solu_ln": 

84 # SoLU-LN checkpoints carry a mid-MLP LayerNorm (NativeMLP.ln); expose 

85 # it so its params load under blocks.{i}.mlp.ln.* and its hooks fire. 

86 submodules["ln"] = NormalizationBridge(name="ln", config=cfg) 

87 return MLPBridge( 

88 name="mlp", 

89 submodules=submodules, 

90 ) 

91 

92 

93def _make_block_submodules(cfg: Any) -> dict: 

94 submods: dict = { 

95 "ln1": _make_norm_bridge("ln1", cfg), 

96 "attn": AttentionBridge( 

97 name="attn", 

98 config=cfg, 

99 submodules={ 

100 "q": LinearBridge(name="q"), 

101 "k": LinearBridge(name="k"), 

102 "v": LinearBridge(name="v"), 

103 "o": LinearBridge(name="o"), 

104 }, 

105 ), 

106 } 

107 if not cfg.attn_only: 

108 submods["ln2"] = _make_norm_bridge("ln2", cfg) 

109 submods["mlp"] = _make_mlp_bridge(cfg) 

110 return submods 

111 

112 

113class NativeArchitectureAdapter(ArchitectureAdapter): 

114 """Adapter for ``NativeModel`` — TL-native, split-QKV, pre-LN; feature set 

115 driven by cfg (gated MLP, RMS norm, rotary, GQA, soft-cap, attn_only).""" 

116 

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

118 super().__init__(cfg) 

119 

120 self.supports_fold_ln = True 

121 self.supports_center_writing_weights = True 

122 # Native Q/K/V/O are nn.Linear [out, in]; the fold_layer_norm formulas 

123 # index [head, d_model, d_head]. Without these rearranges folding either 

124 # raises inside einops or silently mis-places the scale. 

125 self.weight_processing_conversions = { 

126 **self._qkvo_weight_conversions(include_biases=True), 

127 } 

128 

129 # Internal attribute names avoid collisions with bridge slot names 

130 # ("embed", "blocks", "ln_final", "unembed") — the bridge's __getattr__ 

131 # forwards to original_model and would shadow add_module otherwise. 

132 mapping: dict = { 

133 "embed": EmbeddingBridge(name="tok_embed"), 

134 } 

135 if not _is_rotary(cfg): 

136 mapping["pos_embed"] = PosEmbedBridge(name="pos") 

137 block_bridge = BlockBridge( 

138 name="layers", 

139 config=self.cfg, 

140 submodules=_make_block_submodules(self.cfg), 

141 ) 

142 # Under attn_only there's no ln2 / mlp to point at; drop the aliases 

143 # that would otherwise warn during _register_aliases. 

144 if self.cfg.attn_only: 

145 if block_bridge.hook_aliases is BlockBridge.hook_aliases: 145 ↛ 147line 145 didn't jump to line 147 because the condition on line 145 was always true

146 block_bridge.hook_aliases = dict(block_bridge.hook_aliases) 

147 block_bridge.hook_aliases.pop("hook_resid_mid", None) 

148 block_bridge.hook_aliases.pop("hook_mlp_out", None) 

149 mapping["blocks"] = block_bridge 

150 # final_rms forces RMS on the final norm independent of block norm — 

151 # matches Llama's TL config semantic. 

152 mapping["ln_final"] = _make_norm_bridge( 

153 "ln_out", self.cfg, force_rms=bool(getattr(self.cfg, "final_rms", False)) 

154 ) 

155 mapping["unembed"] = UnembeddingBridge(name="head") 

156 self.component_mapping = mapping 

157 

158 def prepare_model(self, model: Any) -> None: 

159 """Reject modules whose attribute names collide with bridge slots. 

160 

161 Bridge's ``__getattr__`` falls back to ``getattr(original_model, name)`` 

162 for unknown attrs, so a name match — submodule, buffer, plain tensor, 

163 or property — makes ``add_module`` raise mid-setup with an opaque 

164 message. Failing here points at the real cause. Reserved set is derived 

165 from ``component_mapping.keys()`` so adapter variants stay in sync. 

166 """ 

167 reserved = set(self.component_mapping.keys()) if self.component_mapping else set() 

168 collisions = sorted(name for name in reserved if hasattr(model, name)) 

169 if collisions: 

170 raise ValueError( 

171 f"{type(model).__name__} cannot be wrapped by NativeArchitectureAdapter: " 

172 f"attribute names {collisions} collide with bridge component slots " 

173 f"({sorted(reserved)}). Rename these attributes to non-colliding names " 

174 f"(e.g. tok_embed, layers, ln_out, head) and update the adapter's " 

175 f"component_mapping ``name=`` fields to match." 

176 )