Coverage for transformer_lens/model_bridge/supported_architectures/internlm2.py: 69%
168 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-09-01 16:23 +0000
1"""InternLM2 architecture adapter."""
3import sys
4from typing import Any
6import torch
7import torch.nn as nn
9from transformer_lens.conversion_utils.conversion_steps import RearrangeTensorConversion
10from transformer_lens.conversion_utils.param_processing_conversion import (
11 ParamProcessingConversion,
12)
13from transformer_lens.model_bridge.architecture_adapter import ArchitectureAdapter
14from transformer_lens.model_bridge.buffer_restore import restore_rotary_inv_freq
15from transformer_lens.model_bridge.compat import patch_dynamic_cache_v5
16from transformer_lens.model_bridge.generalized_components import (
17 BlockBridge,
18 EmbeddingBridge,
19 JointQKVPositionEmbeddingsAttentionBridge,
20 LinearBridge,
21 RMSNormalizationBridge,
22 UnembeddingBridge,
23)
26class _InternLM2AttentionBridge(JointQKVPositionEmbeddingsAttentionBridge):
27 """Attention bridge returning 3-tuple for InternLM2's decoder layer contract.
29 InternLM2's decoder layer unpacks (hidden_states, attn_weights, present_key_value)
30 from self.attention(), but the base bridge returns only (output, weights).
31 """
33 def forward(self, *args: Any, **kwargs: Any) -> Any:
34 """Supply position_embeddings from this layer's own rotary module -- InternLM2
35 has per-attention rotary and never passes them down, so RoPE is otherwise skipped."""
36 if kwargs.get("position_embeddings") is None: 36 ↛ 48line 36 didn't jump to line 48 because the condition on line 36 was always true
37 rotary = getattr(self.original_component, "rotary_emb", None)
38 hidden_states = kwargs.get("hidden_states")
39 if hidden_states is None and args and isinstance(args[0], torch.Tensor): 39 ↛ 40line 39 didn't jump to line 40 because the condition on line 39 was never true
40 hidden_states = args[0]
41 if rotary is not None and isinstance(hidden_states, torch.Tensor): 41 ↛ 48line 41 didn't jump to line 48 because the condition on line 41 was always true
42 position_ids = kwargs.get("position_ids")
43 if position_ids is None: 43 ↛ 47line 43 didn't jump to line 47 because the condition on line 43 was always true
44 position_ids = torch.arange(
45 hidden_states.shape[1], device=hidden_states.device
46 ).unsqueeze(0)
47 kwargs["position_embeddings"] = rotary(hidden_states, position_ids)
48 return super().forward(*args, **kwargs)
50 def _reconstruct_attention(
51 self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, **kwargs
52 ) -> tuple:
53 attn_output, attn_weights = super()._reconstruct_attention(q, k, v, **kwargs)
54 past_key_value = kwargs.get("past_key_values", kwargs.get("past_key_value", None))
55 return (attn_output, attn_weights, past_key_value)
58def _patch_init_weights_for_internlm2() -> None:
59 """Prevent _init_weights from re-randomizing loaded checkpoint weights.
61 Transformers v5 calls _init_weights on all modules after weight
62 materialization. For modules with real (non-meta) tensors, we must
63 skip re-initialization to preserve the loaded checkpoint values.
64 Same approach as openelm.py.
65 """
66 for key in list(sys.modules.keys()):
67 if "internlm2" not in key.lower() or "modeling" not in key.lower():
68 continue
69 module = sys.modules[key]
70 pretrained_cls = getattr(module, "InternLM2PreTrainedModel", None)
71 if pretrained_cls is None or getattr(pretrained_cls, "_tl_patched", False):
72 continue
74 original_init_weights = pretrained_cls._init_weights
76 def safe_init_weights(self, mod, _original=original_init_weights): # type: ignore[no-untyped-def]
77 first_param = next(mod.parameters(), None)
78 if first_param is not None and first_param.device.type != "meta":
79 return
80 _original(self, mod)
82 pretrained_cls._init_weights = safe_init_weights
83 pretrained_cls._tl_patched = True
86class InternLM2ArchitectureAdapter(ArchitectureAdapter):
87 """Architecture adapter for InternLM2 models.
89 InternLM2 uses remote code (trust_remote_code=True) and differs from Llama in:
90 - Fused interleaved GQA wqkv weight (not standard [Q|K|V] split)
91 - Non-standard module names: tok_embeddings, output, attention, feed_forward,
92 wqkv/wo, w1(gate)/w3(up)/w2(down), attention_norm, ffn_norm
93 - Per-layer rotary_emb (no model-level shared instance)
94 - supports_fold_ln=False: fold_ln is done manually in preprocess_weights because
95 the bridge state dict has the fused qkv key, not split q/k/v keys, so
96 fold_layer_norm's extract_attention_tensors_for_folding would silently skip attn.
98 Optional parameters (may not exist in state_dict):
99 - blocks.{i}.attn.b_Q / b_K / b_V / b_O — config.bias=False on shipped models
100 - blocks.{i}.mlp.b_gate / b_in / b_out — MLP always bias=False
101 - blocks.{i}.ln1.b / ln2.b / ln_final.b — RMSNorm has no bias
102 """
104 def __init__(self, cfg: Any) -> None:
105 super().__init__(cfg)
107 self._set_rms_rotary_defaults()
109 # Standard fold_ln silently skips attention when wqkv is fused (see class docstring).
110 # preprocess_weights() handles it instead — same approach as phi3.py.
111 self.supports_fold_ln = False
113 n_kv_heads = getattr(cfg, "n_key_value_heads", None) or cfg.n_heads
115 self.weight_processing_conversions = {
116 "blocks.{i}.attn.q.weight": ParamProcessingConversion(
117 tensor_conversion=RearrangeTensorConversion("(n h) m -> n m h", n=cfg.n_heads),
118 ),
119 "blocks.{i}.attn.k.weight": ParamProcessingConversion(
120 tensor_conversion=RearrangeTensorConversion("(n h) m -> n m h", n=n_kv_heads),
121 ),
122 "blocks.{i}.attn.v.weight": ParamProcessingConversion(
123 tensor_conversion=RearrangeTensorConversion("(n h) m -> n m h", n=n_kv_heads),
124 ),
125 "blocks.{i}.attn.o.weight": ParamProcessingConversion(
126 tensor_conversion=RearrangeTensorConversion("m (n h) -> n h m", n=cfg.n_heads),
127 ),
128 }
130 self.component_mapping = {
131 "embed": EmbeddingBridge(name="model.tok_embeddings"),
132 "blocks": BlockBridge(
133 name="model.layers",
134 submodules={
135 "ln1": RMSNormalizationBridge(name="attention_norm", config=self.cfg),
136 "ln2": RMSNormalizationBridge(name="ffn_norm", config=self.cfg),
137 "attn": _InternLM2AttentionBridge(
138 name="attention",
139 config=self.cfg,
140 split_qkv_matrix=self._split_internlm2_wqkv,
141 submodules={
142 "qkv": LinearBridge(name="wqkv"),
143 "o": LinearBridge(name="wo"),
144 },
145 ),
146 "mlp": self._gated_mlp(name="feed_forward", gate="w1", up="w3", down="w2"),
147 },
148 ),
149 "ln_final": RMSNormalizationBridge(name="model.norm", config=self.cfg),
150 "unembed": UnembeddingBridge(name="output", config=self.cfg),
151 }
153 def _split_internlm2_wqkv(
154 self, attention_component: Any
155 ) -> tuple[nn.Linear, nn.Linear, nn.Linear]:
156 """Split InternLM2's interleaved wqkv into separate Q, K, V linear modules.
158 InternLM2 uses an interleaved GQA layout rather than the standard [Q_all|K_all|V_all].
159 For each of n_kv_heads groups, the weight rows are:
160 [q0, q1, ..., q(n_kv_groups-1), k, v] (each slot = head_dim rows)
161 i.e. gs = n_kv_groups + 2 slots per kv-head group.
162 """
163 wqkv = attention_component.wqkv
164 w = wqkv.weight.data
165 d_model = w.shape[1]
166 has_bias = wqkv.bias is not None
168 n_kv_heads = getattr(self.cfg, "n_key_value_heads", None) or self.cfg.n_heads
169 n_kv_groups = self.cfg.n_heads // n_kv_heads
170 head_dim = self.cfg.d_model // self.cfg.n_heads
171 gs = n_kv_groups + 2
173 w_grouped = w.reshape(n_kv_heads, gs, head_dim, d_model)
174 q_w = w_grouped[:, :n_kv_groups, :, :].reshape(self.cfg.n_heads * head_dim, d_model)
175 k_w = w_grouped[:, n_kv_groups, :, :].reshape(n_kv_heads * head_dim, d_model)
176 v_w = w_grouped[:, n_kv_groups + 1, :, :].reshape(n_kv_heads * head_dim, d_model)
178 q_b: torch.Tensor | None = None
179 k_b: torch.Tensor | None = None
180 v_b: torch.Tensor | None = None
181 if has_bias:
182 b = wqkv.bias.data
183 b_grouped = b.reshape(n_kv_heads, gs, head_dim)
184 q_b = b_grouped[:, :n_kv_groups, :].reshape(self.cfg.n_heads * head_dim)
185 k_b = b_grouped[:, n_kv_groups, :].reshape(n_kv_heads * head_dim)
186 v_b = b_grouped[:, n_kv_groups + 1, :].reshape(n_kv_heads * head_dim)
188 def _make_linear(weight: torch.Tensor, bias: torch.Tensor | None) -> nn.Linear:
189 lin = nn.Linear(d_model, weight.shape[0], bias=bias is not None)
190 lin.weight = nn.Parameter(weight)
191 if bias is not None:
192 lin.bias = nn.Parameter(bias)
193 return lin
195 return _make_linear(q_w, q_b), _make_linear(k_w, k_b), _make_linear(v_w, v_b)
197 def setup_component_testing(self, hf_model: Any, bridge_model: Any = None) -> None:
198 """Inject per-layer rotary embedding for component testing."""
199 try:
200 rotary_emb = hf_model.model.layers[0].attention.rotary_emb
201 except (AttributeError, IndexError):
202 return
204 if bridge_model is not None and hasattr(bridge_model, "blocks"):
205 for block in bridge_model.blocks:
206 if hasattr(block, "attn"):
207 block.attn.set_rotary_emb(rotary_emb)
209 attn_bridge = self.get_generalized_component("blocks.0.attn")
210 attn_bridge.set_rotary_emb(rotary_emb)
212 def prepare_model(self, hf_model: Any) -> None:
213 """Restore per-layer rotary ``inv_freq`` lost to meta-device loading -- this
214 remote code predates HF's ``original_inv_freq`` auto-restore, so positions
215 would otherwise rotate by random values."""
216 super().prepare_model(hf_model)
217 restore_rotary_inv_freq(hf_model)
219 def prepare_loading(self, model_name: str, model_kwargs: dict) -> None:
220 """Patch transformers v5 incompatibilities before from_pretrained runs."""
221 config = model_kwargs.get("config")
222 if config is not None:
223 tp = getattr(config, "pretraining_tp", 1)
224 if tp > 1:
225 raise ValueError(
226 f"InternLM2 adapter does not support pretraining_tp={tp}; "
227 "only pretraining_tp=1 is supported for logit correctness."
228 )
230 patch_dynamic_cache_v5()
232 # Force-import the remote modeling module so we can patch _init_weights.
233 try:
234 from transformers.dynamic_module_utils import get_class_from_dynamic_module
236 get_class_from_dynamic_module(
237 "modeling_internlm2.InternLM2ForCausalLM",
238 model_name,
239 )
240 except Exception:
241 pass
243 _patch_init_weights_for_internlm2()
245 def preprocess_weights(self, state_dict: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
246 """Fold layer norms into QKV and MLP weights.
248 Standard fold_ln can't reach split Q/K/V when wqkv is fused in the bridge state dict.
249 We extract and fold here, then write split keys so RearrangeTensorConversion can follow.
250 MLP projections (w1/w2/w3) are separate linears so they fold normally.
251 Mirrors phi3.py.preprocess_weights, adapted for InternLM2's layout.
252 """
253 fold_ln = getattr(self, "_fold_ln_requested", True)
254 if not fold_ln:
255 return state_dict
257 n_kv_heads = getattr(self.cfg, "n_key_value_heads", None) or self.cfg.n_heads
258 n_kv_groups = self.cfg.n_heads // n_kv_heads
259 head_dim = self.cfg.d_model // self.cfg.n_heads
260 gs = n_kv_groups + 2
262 for i in range(self.cfg.n_layers):
263 # --- Fold ln1 into Q/K/V (extracted from interleaved wqkv) ---
264 qkv_key = f"blocks.{i}.attn.qkv.weight"
265 ln1_key = f"blocks.{i}.ln1.weight"
266 if qkv_key in state_dict and ln1_key in state_dict:
267 ln1_w = state_dict[ln1_key].float()
268 qkv_w = state_dict[qkv_key].float()
269 d_model = qkv_w.shape[1]
270 orig_dtype = state_dict[qkv_key].dtype
272 w_grouped = qkv_w.reshape(n_kv_heads, gs, head_dim, d_model)
273 q_w = w_grouped[:, :n_kv_groups, :, :].reshape(self.cfg.n_heads * head_dim, d_model)
274 k_w = w_grouped[:, n_kv_groups, :, :].reshape(n_kv_heads * head_dim, d_model)
275 v_w = w_grouped[:, n_kv_groups + 1, :, :].reshape(n_kv_heads * head_dim, d_model)
277 state_dict[f"blocks.{i}.attn.q.weight"] = (q_w * ln1_w[None, :]).to(orig_dtype)
278 state_dict[f"blocks.{i}.attn.k.weight"] = (k_w * ln1_w[None, :]).to(orig_dtype)
279 state_dict[f"blocks.{i}.attn.v.weight"] = (v_w * ln1_w[None, :]).to(orig_dtype)
280 del state_dict[qkv_key]
281 state_dict[ln1_key] = torch.ones_like(state_dict[ln1_key])
283 qkv_bias_key = f"blocks.{i}.attn.qkv.bias"
284 if qkv_bias_key in state_dict:
285 b = state_dict[qkv_bias_key]
286 expected_len = (self.cfg.n_heads + 2 * n_kv_heads) * head_dim
287 if b.shape[0] != expected_len: 287 ↛ 288line 287 didn't jump to line 288 because the condition on line 287 was never true
288 raise ValueError(
289 f"Unexpected wqkv bias shape at layer {i}: {b.shape[0]} "
290 f"(expected {expected_len}). Cannot split interleaved bias."
291 )
292 orig_dtype = b.dtype
293 b_f = b.float()
294 b_grouped = b_f.reshape(n_kv_heads, gs, head_dim)
295 q_b = b_grouped[:, :n_kv_groups, :].reshape(self.cfg.n_heads * head_dim)
296 k_b = b_grouped[:, n_kv_groups, :].reshape(n_kv_heads * head_dim)
297 v_b = b_grouped[:, n_kv_groups + 1, :].reshape(n_kv_heads * head_dim)
298 state_dict[f"blocks.{i}.attn.q.bias"] = q_b.to(orig_dtype)
299 state_dict[f"blocks.{i}.attn.k.bias"] = k_b.to(orig_dtype)
300 state_dict[f"blocks.{i}.attn.v.bias"] = v_b.to(orig_dtype)
301 del state_dict[qkv_bias_key]
303 # --- Fold ln2 into MLP gate (w1) and up (w3) projections ---
304 ln2_key = f"blocks.{i}.ln2.weight"
305 if ln2_key in state_dict:
306 ln2_w = state_dict[ln2_key].float()
307 for mlp_key in [
308 f"blocks.{i}.mlp.gate.weight",
309 f"blocks.{i}.mlp.in.weight",
310 ]:
311 if mlp_key in state_dict: 311 ↛ 307line 311 didn't jump to line 307 because the condition on line 311 was always true
312 orig_dtype = state_dict[mlp_key].dtype
313 state_dict[mlp_key] = (state_dict[mlp_key].float() * ln2_w[None, :]).to(
314 orig_dtype
315 )
316 state_dict[ln2_key] = torch.ones_like(state_dict[ln2_key])
318 # --- Fold ln_final into unembed ---
319 ln_final_key = "ln_final.weight"
320 unembed_key = "unembed.weight"
321 if ln_final_key in state_dict and unembed_key in state_dict: 321 ↛ 331line 321 didn't jump to line 331 because the condition on line 321 was always true
322 ln_w = state_dict[ln_final_key].float()
323 u_w = state_dict[unembed_key].float()
324 orig_dtype = state_dict[unembed_key].dtype
325 if u_w.shape[-1] == ln_w.shape[0]: 325 ↛ 327line 325 didn't jump to line 327 because the condition on line 325 was always true
326 state_dict[unembed_key] = (u_w * ln_w[None, :]).to(orig_dtype)
327 elif u_w.shape[0] == ln_w.shape[0]:
328 state_dict[unembed_key] = (u_w * ln_w[:, None]).to(orig_dtype)
329 state_dict[ln_final_key] = torch.ones_like(state_dict[ln_final_key])
331 return state_dict