Coverage for transformer_lens/model_bridge/supported_architectures/baichuan.py: 77%
213 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-09-01 16:23 +0000
1"""Baichuan architecture adapter.
3Supports both BaiChuanForCausalLM (v1) and BaichuanForCausalLM (v2).
4Both use combined QKV via W_pack with RoPE, RMSNorm, and gated MLP.
5"""
7import importlib.util
8import sys
9from typing import Any
11import torch
12import torch.nn as nn
14from transformer_lens.conversion_utils.conversion_steps import RearrangeTensorConversion
15from transformer_lens.conversion_utils.param_processing_conversion import (
16 ParamProcessingConversion,
17)
18from transformer_lens.model_bridge.architecture_adapter import ArchitectureAdapter
19from transformer_lens.model_bridge.compat import patch_dynamic_cache_v5
20from transformer_lens.model_bridge.generalized_components import (
21 BlockBridge,
22 EmbeddingBridge,
23 JointQKVPositionEmbeddingsAttentionBridge,
24 LinearBridge,
25 RMSNormalizationBridge,
26 UnembeddingBridge,
27)
30class _BaichuanAttentionBridge(JointQKVPositionEmbeddingsAttentionBridge):
31 """Attention bridge for Baichuan's v4-era decoder-layer contract.
33 Baichuan predates HF's Cache API and differs from the base bridge in two
34 ways we have to own:
36 1. **Rotary from position_ids**: HF passes `position_ids` (not a
37 pre-computed `position_embeddings` tuple), so we call the per-layer
38 `rotary_emb(v, seq_len=kv_seq_len)` ourselves and slice cos/sin by
39 `position_ids`.
40 2. **Legacy (k, v) cache tuple**: HF's DecoderLayer passes
41 `past_key_value=(k, v)` (singular, per-layer legacy tuple) and expects
42 `self_attn(...)` to return a matching `(k_full, v_full)` as
43 `present_key_value` so Model.forward's `next_decoder_cache` accumulates
44 real tensors. The base bridge's `_update_kv_cache` only handles the
45 Cache-object plural path, so we reimplement the attention body here
46 (mirroring HF's own Attention.forward).
47 """
49 def _reconstruct_attention(
50 self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, **kwargs
51 ) -> tuple:
52 assert self.original_component is not None
53 assert self.config is not None
54 num_heads = self.config.n_heads
55 num_kv_heads = getattr(self.config, "n_key_value_heads", None) or num_heads
57 q, k, v, batch_size, seq_len, head_dim = self._reshape_qkv_to_heads(
58 q, k, v, num_heads, num_kv_heads
59 )
61 past_kv_raw = kwargs.get("past_key_value")
62 past_key_value: tuple[torch.Tensor, torch.Tensor] | None = None
63 if (
64 isinstance(past_kv_raw, tuple)
65 and len(past_kv_raw) >= 2
66 and isinstance(past_kv_raw[0], torch.Tensor)
67 and isinstance(past_kv_raw[1], torch.Tensor)
68 ):
69 past_key_value = (past_kv_raw[0], past_kv_raw[1])
70 past_len = past_key_value[0].shape[-2] if past_key_value is not None else 0
72 # Rotary: derive cos/sin over the full kv_seq_len, index by position_ids.
73 if "position_embeddings" not in kwargs:
74 rotary_emb = getattr(self.original_component, "rotary_emb", None)
75 position_ids = kwargs.get("position_ids")
76 if rotary_emb is not None and position_ids is not None: 76 ↛ 83line 76 didn't jump to line 83 because the condition on line 76 was always true
77 kv_seq_len = seq_len + past_len
78 cos, sin = rotary_emb(v, seq_len=kv_seq_len)
79 cos = cos.squeeze(1).squeeze(0)[position_ids]
80 sin = sin.squeeze(1).squeeze(0)[position_ids]
81 kwargs["position_embeddings"] = (cos, sin)
83 position_embeddings = kwargs.get("position_embeddings")
84 if position_embeddings is not None and isinstance(position_embeddings, tuple): 84 ↛ 89line 84 didn't jump to line 89 because the condition on line 84 was always true
85 cos, sin = self._apply_position_embedding_hooks(position_embeddings)
86 q, k = self._apply_rotary_pos_emb(q, k, cos, sin)
88 # Concat prior (k, v) — already rotary-applied from its own step.
89 if past_key_value is not None:
90 k = torch.cat([past_key_value[0], k], dim=-2)
91 v = torch.cat([past_key_value[1], v], dim=-2)
93 # Build present cache from pre-GQA-expansion (k, v) so downstream
94 # steps don't pay for duplicated heads.
95 use_cache = bool(kwargs.get("use_cache", False))
96 present_key_value = (k, v) if use_cache else None
98 if num_kv_heads != num_heads: 98 ↛ 99line 98 didn't jump to line 99 because the condition on line 98 was never true
99 n_rep = num_heads // num_kv_heads
100 k = k.repeat_interleave(n_rep, dim=1)
101 v = v.repeat_interleave(n_rep, dim=1)
103 kv_seq_len = k.shape[-2]
104 attn_scores = torch.matmul(q, k.transpose(-2, -1)) * (head_dim ** (-0.5))
105 attention_mask = kwargs.get("attention_mask", None)
106 attn_scores = self._apply_reconstruct_attention_mask(
107 attn_scores=attn_scores,
108 attention_mask=attention_mask,
109 seq_len=kv_seq_len,
110 q_seq_len=seq_len,
111 )
112 attn_scores = self.hook_attn_scores(attn_scores)
113 attn_weights = self._softmax_dropout_pattern(attn_scores)
114 attn_output = torch.matmul(attn_weights, v)
115 attn_output = self._reshape_attn_output(
116 attn_output, batch_size, seq_len, num_heads, head_dim
117 )
118 if ( 118 ↛ 123line 118 didn't jump to line 123 because the condition on line 118 was never true
119 bool(getattr(self.config, "use_attn_result", False))
120 and hasattr(self, "o")
121 and self.o.original_component is not None
122 ):
123 attn_output = self.o.hook_in(attn_output)
124 z_4d = attn_output.view(batch_size, seq_len, num_heads, head_dim)
125 attn_output = self._compute_per_head_result(z_4d, num_heads, head_dim)
126 else:
127 attn_output = self._apply_output_projection(attn_output)
129 return (attn_output, attn_weights, present_key_value)
132from transformers import PreTrainedModel as _HFPreTrainedModel
135def _patch_init_weights_for_baichuan() -> None:
136 """Prevent _init_weights from re-randomizing loaded checkpoint weights.
138 Transformers v5 calls _init_weights on all modules after weight
139 materialization. For modules with real (non-meta) tensors, we must
140 skip re-initialization to preserve the loaded checkpoint values.
141 """
142 for key in list(sys.modules.keys()):
143 if "baichuan" not in key.lower() or "modeling" not in key.lower():
144 continue
145 module = sys.modules[key]
146 # Both v1 (BaiChuan) and v2 (Baichuan) define a PreTrainedModel subclass
147 for cls_name in ("BaiChuanPreTrainedModel", "BaichuanPreTrainedModel", "PreTrainedModel"):
148 pretrained_cls = getattr(module, cls_name, None)
149 if pretrained_cls is None or getattr(pretrained_cls, "_tl_patched", False):
150 continue
151 # The remote module does `from transformers import PreTrainedModel`,
152 # so the "PreTrainedModel" name can resolve to the real base class.
153 # Patching that would disable _init_weights — including HF's rotary
154 # buffer restoration — for every model loaded later in the process.
155 if pretrained_cls is _HFPreTrainedModel: 155 ↛ 158line 155 didn't jump to line 158 because the condition on line 155 was always true
156 continue
157 # Only patch classes that define their own _init_weights
158 if "_init_weights" not in pretrained_cls.__dict__:
159 continue
161 original_init_weights = pretrained_cls._init_weights
163 def safe_init_weights(self, mod, _original=original_init_weights): # type: ignore[no-untyped-def]
164 first_param = next(mod.parameters(), None)
165 if first_param is not None and first_param.device.type != "meta":
166 return
167 _original(self, mod)
169 pretrained_cls._init_weights = safe_init_weights
170 pretrained_cls._tl_patched = True
173class BaichuanArchitectureAdapter(ArchitectureAdapter):
174 """Architecture adapter for Baichuan models (v1 and v2).
176 Baichuan uses combined QKV via W_pack (nn.Linear(h, 3*h)) with RoPE,
177 RMSNorm, and gated MLP (SwiGLU). Per-layer rotary embeddings.
179 Optional Parameters (may not exist in state_dict):
180 -------------------------------------------------
181 Baichuan models do NOT have biases on any projection:
183 - blocks.{i}.attn.b_Q / b_K / b_V / b_O — no bias
184 - blocks.{i}.mlp.b_gate / b_in / b_out — no bias
185 - blocks.{i}.ln1.b / ln2.b / ln_final.b — RMSNorm has no bias
186 """
188 def __init__(self, cfg: Any) -> None:
189 super().__init__(cfg)
191 self._set_rms_rotary_defaults()
193 # Fused W_pack prevents standard fold_ln from reaching Q/K/V separately.
194 # preprocess_weights() handles it instead.
195 self.supports_fold_ln = False
197 self.weight_processing_conversions = {
198 "blocks.{i}.attn.q.weight": ParamProcessingConversion(
199 tensor_conversion=RearrangeTensorConversion("(n h) m -> n m h", n=cfg.n_heads),
200 ),
201 "blocks.{i}.attn.k.weight": ParamProcessingConversion(
202 tensor_conversion=RearrangeTensorConversion("(n h) m -> n m h", n=cfg.n_heads),
203 ),
204 "blocks.{i}.attn.v.weight": ParamProcessingConversion(
205 tensor_conversion=RearrangeTensorConversion("(n h) m -> n m h", n=cfg.n_heads),
206 ),
207 "blocks.{i}.attn.o.weight": ParamProcessingConversion(
208 tensor_conversion=RearrangeTensorConversion("m (n h) -> n h m", n=cfg.n_heads),
209 ),
210 }
212 self.component_mapping = {
213 "embed": EmbeddingBridge(name="model.embed_tokens"),
214 "blocks": BlockBridge(
215 name="model.layers",
216 submodules={
217 "ln1": RMSNormalizationBridge(name="input_layernorm", config=self.cfg),
218 "ln2": RMSNormalizationBridge(name="post_attention_layernorm", config=self.cfg),
219 "attn": _BaichuanAttentionBridge(
220 name="self_attn",
221 config=self.cfg,
222 split_qkv_matrix=self._split_baichuan_w_pack,
223 submodules={
224 "qkv": LinearBridge(name="W_pack"),
225 "o": LinearBridge(name="o_proj"),
226 },
227 ),
228 "mlp": self._gated_mlp(),
229 },
230 ),
231 "ln_final": RMSNormalizationBridge(name="model.norm", config=self.cfg),
232 "unembed": UnembeddingBridge(name="lm_head", config=self.cfg),
233 }
235 def _split_baichuan_w_pack(
236 self, attention_component: Any
237 ) -> tuple[nn.Linear, nn.Linear, nn.Linear]:
238 """Split Baichuan's W_pack into separate Q, K, V linear modules.
240 W_pack is a simple concatenation: [Q | K | V], each of size hidden_size.
241 No interleaving, no GQA — all three chunks are equal size.
242 """
243 w_pack = attention_component.W_pack
244 weight = w_pack.weight.data
245 d_model = weight.shape[1]
246 hidden_size = d_model # Q, K, V each have hidden_size output features
248 q_w = weight[:hidden_size, :]
249 k_w = weight[hidden_size : 2 * hidden_size, :]
250 v_w = weight[2 * hidden_size :, :]
252 def _make_linear(w: torch.Tensor) -> nn.Linear:
253 lin = nn.Linear(d_model, hidden_size, bias=False)
254 lin.weight = nn.Parameter(w)
255 return lin
257 return _make_linear(q_w), _make_linear(k_w), _make_linear(v_w)
259 def setup_component_testing(self, hf_model: Any, bridge_model: Any = None) -> None:
260 """Inject per-layer rotary embedding for component testing."""
261 try:
262 rotary_emb = hf_model.model.layers[0].self_attn.rotary_emb
263 except (AttributeError, IndexError):
264 return
266 if bridge_model is not None and hasattr(bridge_model, "blocks"):
267 for block in bridge_model.blocks:
268 if hasattr(block, "attn"):
269 block.attn.set_rotary_emb(rotary_emb)
271 attn_bridge = self.get_generalized_component("blocks.0.attn")
272 attn_bridge.set_rotary_emb(rotary_emb)
274 def prepare_loading(self, model_name: str, model_kwargs: dict) -> None:
275 """Patch transformers v5 incompatibilities before from_pretrained runs."""
276 patch_dynamic_cache_v5()
278 # Force-import the remote modeling module so we can patch _init_weights.
279 # Baichuan2 variants ship quantizer.py which imports bitsandbytes;
280 # transformers' check_imports scans every .py file in the repo and
281 # raises ImportError if bitsandbytes is missing, even though quantizer
282 # is not used in normal inference. Catch that case and tell the user
283 # how to install the optional dependency group.
284 try:
285 from transformers.dynamic_module_utils import get_class_from_dynamic_module
287 last_exc: Exception | None = None
288 # Try both class names (v1 and v2)
289 for cls_name in (
290 "modeling_baichuan.BaichuanForCausalLM",
291 "modeling_baichuan.BaiChuanForCausalLM",
292 ):
293 try:
294 get_class_from_dynamic_module(cls_name, model_name)
295 last_exc = None
296 break
297 except Exception as exc:
298 last_exc = exc
299 continue
300 if last_exc is not None and "bitsandbytes" in str(last_exc):
301 if importlib.util.find_spec("bitsandbytes") is None: 301 ↛ 313line 301 didn't jump to line 313 because the condition on line 301 was always true
302 raise ImportError(
303 "Baichuan2 variants require `bitsandbytes` for "
304 "trust_remote_code loading (their shipped quantizer.py "
305 "imports it). Install the quantization extras: "
306 "`uv sync --group quantization`."
307 ) from last_exc
308 except ImportError:
309 raise
310 except Exception:
311 pass
313 _patch_init_weights_for_baichuan()
315 def prepare_model(self, hf_model: Any) -> None:
316 """Fix rotary caches and normalize NormHead weights before bridge creation.
318 RotaryEmbedding differs between v1 and v2:
319 - v1 (Baichuan-7B): `inv_freq` is a persistent buffer, loaded from the
320 checkpoint as bfloat16, but `cos_cached`/`sin_cached` are non-persistent
321 and materialize as garbage under meta-init.
322 - v2 (Baichuan2-*): `inv_freq`, `cos_cached`, `sin_cached` are all plain
323 attributes (no `register_buffer`). v5's meta-init materializes them on
324 meta, and nothing in the checkpoint overwrites them.
326 Both cases are resolved by computing inv_freq + caches from scratch at
327 float32 using config-derived head_dim and base=10000. Recomputing v1 at
328 float32 is also an upgrade over its bfloat16 checkpoint values.
330 Baichuan2 Chat also uses NormHead which row-normalizes lm_head during
331 forward. We apply that once here so the bridge sees the normalized
332 weights directly without needing NormHead's forward path.
333 """
334 # Pick a real device/dtype by scanning real (non-meta) parameters.
335 target_device = torch.device("cpu")
336 params_fn = getattr(hf_model, "parameters", None)
337 if callable(params_fn): 337 ↛ 338line 337 didn't jump to line 338 because the condition on line 337 was never true
338 for param in params_fn():
339 if param.device.type != "meta":
340 target_device = param.device
341 break
343 head_dim = self.cfg.d_model // self.cfg.n_heads
344 base = 10000.0
346 model_core = getattr(hf_model, "model", None)
347 if model_core is not None:
348 for layer in getattr(model_core, "layers", []):
349 rotary = getattr(getattr(layer, "self_attn", None), "rotary_emb", None)
350 if rotary is None: 350 ↛ 351line 350 didn't jump to line 351 because the condition on line 350 was never true
351 continue
352 max_seq = getattr(rotary, "max_seq_len_cached", self.cfg.n_ctx or 4096)
353 inv_freq = 1.0 / (
354 base
355 ** (
356 torch.arange(0, head_dim, 2, device=target_device, dtype=torch.float32)
357 / head_dim
358 )
359 )
360 t = torch.arange(max_seq, device=target_device, dtype=torch.float32)
361 freqs = torch.einsum("i,j->ij", t, inv_freq)
362 emb = torch.cat((freqs, freqs), dim=-1)
363 rotary.inv_freq = inv_freq
364 rotary.cos_cached = emb.cos()[None, None, :, :]
365 rotary.sin_cached = emb.sin()[None, None, :, :]
367 # Normalize NormHead weights (Baichuan2 Chat)
368 lm_head = getattr(hf_model, "lm_head", None)
369 if lm_head is not None and hasattr(lm_head, "first_flag"):
370 w = lm_head.weight.data
371 lm_head.weight.data = torch.nn.functional.normalize(w, dim=-1)
373 def preprocess_weights(self, state_dict: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
374 """Split fused W_pack QKV and optionally fold layer norms."""
375 fold_ln = getattr(self, "_fold_ln_requested", True)
376 if not fold_ln:
377 # Still need to split W_pack into Q/K/V for weight conversions
378 for i in range(self.cfg.n_layers):
379 qkv_key = f"blocks.{i}.attn.qkv.weight"
380 if qkv_key not in state_dict: 380 ↛ 381line 380 didn't jump to line 381 because the condition on line 380 was never true
381 continue
382 w = state_dict[qkv_key]
383 hidden_size = w.shape[1]
384 q_w = w[:hidden_size, :]
385 k_w = w[hidden_size : 2 * hidden_size, :]
386 v_w = w[2 * hidden_size :, :]
387 state_dict[f"blocks.{i}.attn.q.weight"] = q_w
388 state_dict[f"blocks.{i}.attn.k.weight"] = k_w
389 state_dict[f"blocks.{i}.attn.v.weight"] = v_w
390 del state_dict[qkv_key]
391 return state_dict
393 for i in range(self.cfg.n_layers):
394 # --- Fold ln1 into Q/K/V (split from W_pack) ---
395 qkv_key = f"blocks.{i}.attn.qkv.weight"
396 ln1_key = f"blocks.{i}.ln1.weight"
397 if qkv_key in state_dict and ln1_key in state_dict: 397 ↛ 414line 397 didn't jump to line 414 because the condition on line 397 was always true
398 ln1_w = state_dict[ln1_key].float()
399 w = state_dict[qkv_key].float()
400 orig_dtype = state_dict[qkv_key].dtype
401 hidden_size = w.shape[1]
403 q_w = w[:hidden_size, :]
404 k_w = w[hidden_size : 2 * hidden_size, :]
405 v_w = w[2 * hidden_size :, :]
407 state_dict[f"blocks.{i}.attn.q.weight"] = (q_w * ln1_w[None, :]).to(orig_dtype)
408 state_dict[f"blocks.{i}.attn.k.weight"] = (k_w * ln1_w[None, :]).to(orig_dtype)
409 state_dict[f"blocks.{i}.attn.v.weight"] = (v_w * ln1_w[None, :]).to(orig_dtype)
410 del state_dict[qkv_key]
411 state_dict[ln1_key] = torch.ones_like(state_dict[ln1_key])
413 # --- Fold ln2 into MLP gate and up projections ---
414 ln2_key = f"blocks.{i}.ln2.weight"
415 if ln2_key in state_dict: 415 ↛ 393line 415 didn't jump to line 393 because the condition on line 415 was always true
416 ln2_w = state_dict[ln2_key].float()
417 for mlp_key in [
418 f"blocks.{i}.mlp.gate.weight",
419 f"blocks.{i}.mlp.in.weight",
420 ]:
421 if mlp_key in state_dict: 421 ↛ 417line 421 didn't jump to line 417 because the condition on line 421 was always true
422 orig_dtype = state_dict[mlp_key].dtype
423 state_dict[mlp_key] = (state_dict[mlp_key].float() * ln2_w[None, :]).to(
424 orig_dtype
425 )
426 state_dict[ln2_key] = torch.ones_like(state_dict[ln2_key])
428 # --- Fold ln_final into unembed ---
429 ln_final_key = "ln_final.weight"
430 unembed_key = "unembed.weight"
431 if ln_final_key in state_dict and unembed_key in state_dict: 431 ↛ 441line 431 didn't jump to line 441 because the condition on line 431 was always true
432 ln_w = state_dict[ln_final_key].float()
433 u_w = state_dict[unembed_key].float()
434 orig_dtype = state_dict[unembed_key].dtype
435 if u_w.shape[-1] == ln_w.shape[0]: 435 ↛ 437line 435 didn't jump to line 437 because the condition on line 435 was always true
436 state_dict[unembed_key] = (u_w * ln_w[None, :]).to(orig_dtype)
437 elif u_w.shape[0] == ln_w.shape[0]:
438 state_dict[unembed_key] = (u_w * ln_w[:, None]).to(orig_dtype)
439 state_dict[ln_final_key] = torch.ones_like(state_dict[ln_final_key])
441 return state_dict