Coverage for transformer_lens/model_bridge/supported_architectures/internlm2.py: 65%
89 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-09-21 19:27 +0000
1"""InternLM2 architecture adapter."""
3from typing import Any
5import torch
6import torch.nn as nn
8from transformer_lens.conversion_utils.conversion_steps import RearrangeTensorConversion
9from transformer_lens.conversion_utils.param_processing_conversion import (
10 ParamProcessingConversion,
11)
12from transformer_lens.model_bridge.architecture_adapter import ArchitectureAdapter
13from transformer_lens.model_bridge.buffer_restore import restore_rotary_inv_freq
14from transformer_lens.model_bridge.compat import patch_dynamic_cache_v5
15from transformer_lens.model_bridge.generalized_components import (
16 BlockBridge,
17 EmbeddingBridge,
18 JointQKVPositionEmbeddingsAttentionBridge,
19 LinearBridge,
20 RMSNormalizationBridge,
21 UnembeddingBridge,
22)
23from transformer_lens.model_bridge.supported_architectures._remote_code_compat import (
24 force_import_remote_class,
25 iter_remote_modeling_modules,
26 patch_init_weights_skip_loaded,
27)
30class _InternLM2AttentionBridge(JointQKVPositionEmbeddingsAttentionBridge):
31 """Attention bridge returning 3-tuple for InternLM2's decoder layer contract.
33 InternLM2's decoder layer unpacks (hidden_states, attn_weights, present_key_value)
34 from self.attention(), but the base bridge returns only (output, weights).
35 """
37 def forward(self, *args: Any, **kwargs: Any) -> Any:
38 """Supply position_embeddings from this layer's own rotary module -- InternLM2
39 has per-attention rotary and never passes them down, so RoPE is otherwise skipped."""
40 if kwargs.get("position_embeddings") is None: 40 ↛ 52line 40 didn't jump to line 52 because the condition on line 40 was always true
41 rotary = getattr(self.original_component, "rotary_emb", None)
42 hidden_states = kwargs.get("hidden_states")
43 if hidden_states is None and args and isinstance(args[0], torch.Tensor): 43 ↛ 44line 43 didn't jump to line 44 because the condition on line 43 was never true
44 hidden_states = args[0]
45 if rotary is not None and isinstance(hidden_states, torch.Tensor): 45 ↛ 52line 45 didn't jump to line 52 because the condition on line 45 was always true
46 position_ids = kwargs.get("position_ids")
47 if position_ids is None:
48 position_ids = torch.arange(
49 hidden_states.shape[1], device=hidden_states.device
50 ).unsqueeze(0)
51 kwargs["position_embeddings"] = rotary(hidden_states, position_ids)
52 return super().forward(*args, **kwargs)
54 def _reconstruct_attention(
55 self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, **kwargs
56 ) -> tuple:
57 attn_output, attn_weights = super()._reconstruct_attention(q, k, v, **kwargs)
58 past_key_value = kwargs.get("past_key_values", kwargs.get("past_key_value", None))
59 return (attn_output, attn_weights, past_key_value)
62class InternLM2ArchitectureAdapter(ArchitectureAdapter):
63 """Architecture adapter for InternLM2 models.
65 InternLM2 uses remote code (trust_remote_code=True) and differs from Llama in:
66 - Fused interleaved GQA wqkv weight (not standard [Q|K|V] split). The attention
67 bridge splits it at load and drops the fused key from the state dict, so
68 fold_ln sees ordinary q/k/v keys and needs no special handling here.
69 - Non-standard module names: tok_embeddings, output, attention, feed_forward,
70 wqkv/wo, w1(gate)/w3(up)/w2(down), attention_norm, ffn_norm
71 - Per-layer rotary_emb (no model-level shared instance)
73 Optional parameters (may not exist in state_dict):
74 - blocks.{i}.attn.b_Q / b_K / b_V / b_O — config.bias=False on shipped models
75 - blocks.{i}.mlp.b_gate / b_in / b_out — MLP always bias=False
76 - blocks.{i}.ln1.b / ln2.b / ln_final.b — RMSNorm has no bias
77 """
79 def __init__(self, cfg: Any) -> None:
80 super().__init__(cfg)
82 self._set_rms_rotary_defaults()
84 n_kv_heads = getattr(cfg, "n_key_value_heads", None) or cfg.n_heads
86 self.weight_processing_conversions = {
87 "blocks.{i}.attn.q.weight": ParamProcessingConversion(
88 tensor_conversion=RearrangeTensorConversion("(n h) m -> n m h", n=cfg.n_heads),
89 ),
90 "blocks.{i}.attn.k.weight": ParamProcessingConversion(
91 tensor_conversion=RearrangeTensorConversion("(n h) m -> n m h", n=n_kv_heads),
92 ),
93 "blocks.{i}.attn.v.weight": ParamProcessingConversion(
94 tensor_conversion=RearrangeTensorConversion("(n h) m -> n m h", n=n_kv_heads),
95 ),
96 "blocks.{i}.attn.o.weight": ParamProcessingConversion(
97 tensor_conversion=RearrangeTensorConversion("m (n h) -> n h m", n=cfg.n_heads),
98 ),
99 }
101 self.component_mapping = {
102 "embed": EmbeddingBridge(name="model.tok_embeddings"),
103 "blocks": BlockBridge(
104 name="model.layers",
105 submodules={
106 "ln1": RMSNormalizationBridge(name="attention_norm", config=self.cfg),
107 "ln2": RMSNormalizationBridge(name="ffn_norm", config=self.cfg),
108 "attn": _InternLM2AttentionBridge(
109 name="attention",
110 config=self.cfg,
111 split_qkv_matrix=self._split_internlm2_wqkv,
112 submodules={
113 "qkv": LinearBridge(name="wqkv"),
114 "o": LinearBridge(name="wo"),
115 },
116 ),
117 "mlp": self._gated_mlp(name="feed_forward", gate="w1", up="w3", down="w2"),
118 },
119 ),
120 "ln_final": RMSNormalizationBridge(name="model.norm", config=self.cfg),
121 "unembed": UnembeddingBridge(name="output", config=self.cfg),
122 }
124 def _split_internlm2_wqkv(
125 self, attention_component: Any
126 ) -> tuple[nn.Linear, nn.Linear, nn.Linear]:
127 """Split InternLM2's interleaved wqkv into separate Q, K, V linear modules.
129 InternLM2 uses an interleaved GQA layout rather than the standard [Q_all|K_all|V_all].
130 For each of n_kv_heads groups, the weight rows are:
131 [q0, q1, ..., q(n_kv_groups-1), k, v] (each slot = head_dim rows)
132 i.e. gs = n_kv_groups + 2 slots per kv-head group.
133 """
134 wqkv = attention_component.wqkv
135 w = wqkv.weight.data
136 d_model = w.shape[1]
137 has_bias = wqkv.bias is not None
139 n_kv_heads = getattr(self.cfg, "n_key_value_heads", None) or self.cfg.n_heads
140 n_kv_groups = self.cfg.n_heads // n_kv_heads
141 head_dim = self.cfg.d_model // self.cfg.n_heads
142 gs = n_kv_groups + 2
144 w_grouped = w.reshape(n_kv_heads, gs, head_dim, d_model)
145 q_w = w_grouped[:, :n_kv_groups, :, :].reshape(self.cfg.n_heads * head_dim, d_model)
146 k_w = w_grouped[:, n_kv_groups, :, :].reshape(n_kv_heads * head_dim, d_model)
147 v_w = w_grouped[:, n_kv_groups + 1, :, :].reshape(n_kv_heads * head_dim, d_model)
149 q_b: torch.Tensor | None = None
150 k_b: torch.Tensor | None = None
151 v_b: torch.Tensor | None = None
152 if has_bias: 152 ↛ 159line 152 didn't jump to line 159 because the condition on line 152 was always true
153 b = wqkv.bias.data
154 b_grouped = b.reshape(n_kv_heads, gs, head_dim)
155 q_b = b_grouped[:, :n_kv_groups, :].reshape(self.cfg.n_heads * head_dim)
156 k_b = b_grouped[:, n_kv_groups, :].reshape(n_kv_heads * head_dim)
157 v_b = b_grouped[:, n_kv_groups + 1, :].reshape(n_kv_heads * head_dim)
159 def _make_linear(weight: torch.Tensor, bias: torch.Tensor | None) -> nn.Linear:
160 lin = nn.Linear(d_model, weight.shape[0], bias=bias is not None)
161 lin.weight = nn.Parameter(weight)
162 if bias is not None:
163 lin.bias = nn.Parameter(bias)
164 return lin
166 return _make_linear(q_w, q_b), _make_linear(k_w, k_b), _make_linear(v_w, v_b)
168 def setup_component_testing(self, hf_model: Any, bridge_model: Any = None) -> None:
169 """Inject per-layer rotary embedding for component testing."""
170 try:
171 rotary_emb = hf_model.model.layers[0].attention.rotary_emb
172 except (AttributeError, IndexError):
173 return
175 if bridge_model is not None and hasattr(bridge_model, "blocks"):
176 for block in bridge_model.blocks:
177 if hasattr(block, "attn"):
178 block.attn.set_rotary_emb(rotary_emb)
180 attn_bridge = self.get_generalized_component("blocks.0.attn")
181 attn_bridge.set_rotary_emb(rotary_emb)
183 def prepare_model(self, hf_model: Any) -> None:
184 """Restore per-layer rotary ``inv_freq`` lost to meta-device loading -- this
185 remote code predates HF's ``original_inv_freq`` auto-restore, so positions
186 would otherwise rotate by random values."""
187 super().prepare_model(hf_model)
188 restore_rotary_inv_freq(hf_model)
190 def prepare_loading(self, model_name: str, model_kwargs: dict) -> None:
191 """Patch transformers v5 incompatibilities before from_pretrained runs."""
192 config = model_kwargs.get("config")
193 if config is not None:
194 tp = getattr(config, "pretraining_tp", 1)
195 if tp > 1:
196 raise ValueError(
197 f"InternLM2 adapter does not support pretraining_tp={tp}; "
198 "only pretraining_tp=1 is supported for logit correctness."
199 )
201 patch_dynamic_cache_v5()
203 # Force-import the remote modeling module so we can patch _init_weights.
204 force_import_remote_class(model_name, "modeling_internlm2.InternLM2ForCausalLM")
206 # v5 calls _init_weights on all modules after weight materialization;
207 # skip already-real tensors so checkpoint values survive.
208 for module in iter_remote_modeling_modules("internlm2"):
209 pretrained_cls = getattr(module, "InternLM2PreTrainedModel", None)
210 if pretrained_cls is not None:
211 patch_init_weights_skip_loaded(pretrained_cls)