Coverage for transformer_lens/model_bridge/supported_architectures/baichuan.py: 78%
147 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-09-21 19:27 +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
8from typing import Any
10import torch
11import torch.nn as nn
13from transformer_lens.conversion_utils.conversion_steps import RearrangeTensorConversion
14from transformer_lens.conversion_utils.param_processing_conversion import (
15 ParamProcessingConversion,
16)
17from transformer_lens.model_bridge.architecture_adapter import ArchitectureAdapter
18from transformer_lens.model_bridge.compat import patch_dynamic_cache_v5
19from transformer_lens.model_bridge.generalized_components import (
20 BlockBridge,
21 EmbeddingBridge,
22 JointQKVPositionEmbeddingsAttentionBridge,
23 LinearBridge,
24 RMSNormalizationBridge,
25 UnembeddingBridge,
26)
27from transformer_lens.model_bridge.supported_architectures._remote_code_compat import (
28 iter_remote_modeling_modules,
29 patch_init_weights_skip_loaded,
30)
33class _BaichuanAttentionBridge(JointQKVPositionEmbeddingsAttentionBridge):
34 """Attention bridge for Baichuan's v4-era decoder-layer contract.
36 Baichuan predates HF's Cache API and differs from the base bridge in two
37 ways we have to own:
39 1. **Rotary from position_ids**: HF passes `position_ids` (not a
40 pre-computed `position_embeddings` tuple), so we call the per-layer
41 `rotary_emb(v, seq_len=kv_seq_len)` ourselves and slice cos/sin by
42 `position_ids`.
43 2. **Legacy (k, v) cache tuple**: HF's DecoderLayer passes
44 `past_key_value=(k, v)` (singular, per-layer legacy tuple) and expects
45 `self_attn(...)` to return a matching `(k_full, v_full)` as
46 `present_key_value` so Model.forward's `next_decoder_cache` accumulates
47 real tensors. The base bridge's `_update_kv_cache` only handles the
48 Cache-object plural path, so we reimplement the attention body here
49 (mirroring HF's own Attention.forward).
50 """
52 def _reconstruct_attention(
53 self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, **kwargs
54 ) -> tuple:
55 assert self.original_component is not None
56 assert self.config is not None
57 num_heads = self.config.n_heads
58 num_kv_heads = getattr(self.config, "n_key_value_heads", None) or num_heads
60 q, k, v, batch_size, seq_len, head_dim = self._reshape_qkv_to_heads(
61 q, k, v, num_heads, num_kv_heads
62 )
64 past_kv_raw = kwargs.get("past_key_value")
65 past_key_value: tuple[torch.Tensor, torch.Tensor] | None = None
66 if (
67 isinstance(past_kv_raw, tuple)
68 and len(past_kv_raw) >= 2
69 and isinstance(past_kv_raw[0], torch.Tensor)
70 and isinstance(past_kv_raw[1], torch.Tensor)
71 ):
72 past_key_value = (past_kv_raw[0], past_kv_raw[1])
73 past_len = past_key_value[0].shape[-2] if past_key_value is not None else 0
75 # Rotary: derive cos/sin over the full kv_seq_len, index by position_ids.
76 if "position_embeddings" not in kwargs:
77 rotary_emb = getattr(self.original_component, "rotary_emb", None)
78 position_ids = kwargs.get("position_ids")
79 if rotary_emb is not None and position_ids is not None: 79 ↛ 86line 79 didn't jump to line 86 because the condition on line 79 was always true
80 kv_seq_len = seq_len + past_len
81 cos, sin = rotary_emb(v, seq_len=kv_seq_len)
82 cos = cos.squeeze(1).squeeze(0)[position_ids]
83 sin = sin.squeeze(1).squeeze(0)[position_ids]
84 kwargs["position_embeddings"] = (cos, sin)
86 position_embeddings = kwargs.get("position_embeddings")
87 if position_embeddings is not None and isinstance(position_embeddings, tuple): 87 ↛ 92line 87 didn't jump to line 92 because the condition on line 87 was always true
88 cos, sin = self._apply_position_embedding_hooks(position_embeddings)
89 q, k = self._apply_rotary_pos_emb(q, k, cos, sin)
91 # Concat prior (k, v) — already rotary-applied from its own step.
92 if past_key_value is not None:
93 k = torch.cat([past_key_value[0], k], dim=-2)
94 v = torch.cat([past_key_value[1], v], dim=-2)
96 # Build present cache from pre-GQA-expansion (k, v) so downstream
97 # steps don't pay for duplicated heads.
98 use_cache = bool(kwargs.get("use_cache", False))
99 present_key_value = (k, v) if use_cache else None
101 if num_kv_heads != num_heads: 101 ↛ 102line 101 didn't jump to line 102 because the condition on line 101 was never true
102 n_rep = num_heads // num_kv_heads
103 k = k.repeat_interleave(n_rep, dim=1)
104 v = v.repeat_interleave(n_rep, dim=1)
106 kv_seq_len = k.shape[-2]
107 attn_scores = torch.matmul(q, k.transpose(-2, -1)) * (head_dim ** (-0.5))
108 attention_mask = kwargs.get("attention_mask", None)
109 attn_scores = self._apply_reconstruct_attention_mask(
110 attn_scores=attn_scores,
111 attention_mask=attention_mask,
112 seq_len=kv_seq_len,
113 q_seq_len=seq_len,
114 )
115 attn_scores = self.hook_attn_scores(attn_scores)
116 attn_weights = self._softmax_dropout_pattern(attn_scores)
117 attn_output = torch.matmul(attn_weights, v)
118 attn_output = self._reshape_attn_output(
119 attn_output, batch_size, seq_len, num_heads, head_dim
120 )
121 if ( 121 ↛ 126line 121 didn't jump to line 126 because the condition on line 121 was never true
122 bool(getattr(self.config, "use_attn_result", False))
123 and hasattr(self, "o")
124 and self.o.original_component is not None
125 ):
126 attn_output = self.o.hook_in(attn_output)
127 z_4d = attn_output.view(batch_size, seq_len, num_heads, head_dim)
128 attn_output = self._compute_per_head_result(z_4d, num_heads, head_dim)
129 else:
130 attn_output = self._apply_output_projection(attn_output)
132 return (attn_output, attn_weights, present_key_value)
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; the helper skips modules with real (non-meta) tensors so
140 loaded checkpoint values survive.
141 """
142 for module in iter_remote_modeling_modules("baichuan"):
143 # Both v1 (BaiChuan) and v2 (Baichuan) define a PreTrainedModel subclass.
144 # The remote module also does `from transformers import PreTrainedModel`,
145 # so the last name can resolve to the real HF base class — the helper
146 # refuses that (and anything without its own _init_weights); skip and
147 # keep scanning.
148 for cls_name in ("BaiChuanPreTrainedModel", "BaichuanPreTrainedModel", "PreTrainedModel"):
149 pretrained_cls = getattr(module, cls_name, None)
150 if pretrained_cls is None:
151 continue
152 try:
153 patch_init_weights_skip_loaded(pretrained_cls)
154 except ValueError:
155 continue
158class BaichuanArchitectureAdapter(ArchitectureAdapter):
159 """Architecture adapter for Baichuan models (v1 and v2).
161 Baichuan uses combined QKV via W_pack (nn.Linear(h, 3*h)) with RoPE,
162 RMSNorm, and gated MLP (SwiGLU). Per-layer rotary embeddings.
164 The attention bridge splits W_pack at load and drops the fused key from the
165 state dict, so fold_ln sees ordinary q/k/v keys and needs no help here.
167 Optional Parameters (may not exist in state_dict):
168 -------------------------------------------------
169 Baichuan models do NOT have biases on any projection:
171 - blocks.{i}.attn.b_Q / b_K / b_V / b_O — no bias
172 - blocks.{i}.mlp.b_gate / b_in / b_out — no bias
173 - blocks.{i}.ln1.b / ln2.b / ln_final.b — RMSNorm has no bias
174 """
176 def __init__(self, cfg: Any) -> None:
177 super().__init__(cfg)
179 self._set_rms_rotary_defaults()
181 self.weight_processing_conversions = {
182 "blocks.{i}.attn.q.weight": ParamProcessingConversion(
183 tensor_conversion=RearrangeTensorConversion("(n h) m -> n m h", n=cfg.n_heads),
184 ),
185 "blocks.{i}.attn.k.weight": ParamProcessingConversion(
186 tensor_conversion=RearrangeTensorConversion("(n h) m -> n m h", n=cfg.n_heads),
187 ),
188 "blocks.{i}.attn.v.weight": ParamProcessingConversion(
189 tensor_conversion=RearrangeTensorConversion("(n h) m -> n m h", n=cfg.n_heads),
190 ),
191 "blocks.{i}.attn.o.weight": ParamProcessingConversion(
192 tensor_conversion=RearrangeTensorConversion("m (n h) -> n h m", n=cfg.n_heads),
193 ),
194 }
196 self.component_mapping = {
197 "embed": EmbeddingBridge(name="model.embed_tokens"),
198 "blocks": BlockBridge(
199 name="model.layers",
200 submodules={
201 "ln1": RMSNormalizationBridge(name="input_layernorm", config=self.cfg),
202 "ln2": RMSNormalizationBridge(name="post_attention_layernorm", config=self.cfg),
203 "attn": _BaichuanAttentionBridge(
204 name="self_attn",
205 config=self.cfg,
206 split_qkv_matrix=self._split_baichuan_w_pack,
207 submodules={
208 "qkv": LinearBridge(name="W_pack"),
209 "o": LinearBridge(name="o_proj"),
210 },
211 ),
212 "mlp": self._gated_mlp(),
213 },
214 ),
215 "ln_final": RMSNormalizationBridge(name="model.norm", config=self.cfg),
216 "unembed": UnembeddingBridge(name="lm_head", config=self.cfg),
217 }
219 def _split_baichuan_w_pack(
220 self, attention_component: Any
221 ) -> tuple[nn.Linear, nn.Linear, nn.Linear]:
222 """Split Baichuan's W_pack into separate Q, K, V linear modules.
224 W_pack is a simple concatenation: [Q | K | V], each of size hidden_size.
225 No interleaving, no GQA — all three chunks are equal size.
226 """
227 w_pack = attention_component.W_pack
228 weight = w_pack.weight.data
229 d_model = weight.shape[1]
230 hidden_size = d_model # Q, K, V each have hidden_size output features
232 q_w = weight[:hidden_size, :]
233 k_w = weight[hidden_size : 2 * hidden_size, :]
234 v_w = weight[2 * hidden_size :, :]
236 def _make_linear(w: torch.Tensor) -> nn.Linear:
237 lin = nn.Linear(d_model, hidden_size, bias=False)
238 lin.weight = nn.Parameter(w)
239 return lin
241 return _make_linear(q_w), _make_linear(k_w), _make_linear(v_w)
243 def setup_component_testing(self, hf_model: Any, bridge_model: Any = None) -> None:
244 """Inject per-layer rotary embedding for component testing."""
245 try:
246 rotary_emb = hf_model.model.layers[0].self_attn.rotary_emb
247 except (AttributeError, IndexError):
248 return
250 if bridge_model is not None and hasattr(bridge_model, "blocks"):
251 for block in bridge_model.blocks:
252 if hasattr(block, "attn"):
253 block.attn.set_rotary_emb(rotary_emb)
255 attn_bridge = self.get_generalized_component("blocks.0.attn")
256 attn_bridge.set_rotary_emb(rotary_emb)
258 def prepare_loading(self, model_name: str, model_kwargs: dict) -> None:
259 """Patch transformers v5 incompatibilities before from_pretrained runs."""
260 patch_dynamic_cache_v5()
262 # Force-import the remote modeling module so we can patch _init_weights.
263 # Baichuan2 variants ship quantizer.py which imports bitsandbytes;
264 # transformers' check_imports scans every .py file in the repo and
265 # raises ImportError if bitsandbytes is missing, even though quantizer
266 # is not used in normal inference. Catch that case and tell the user
267 # how to install the optional dependency group.
268 try:
269 from transformers.dynamic_module_utils import get_class_from_dynamic_module
271 last_exc: Exception | None = None
272 # Try both class names (v1 and v2)
273 for cls_name in (
274 "modeling_baichuan.BaichuanForCausalLM",
275 "modeling_baichuan.BaiChuanForCausalLM",
276 ):
277 try:
278 get_class_from_dynamic_module(cls_name, model_name)
279 last_exc = None
280 break
281 except Exception as exc:
282 last_exc = exc
283 continue
284 if last_exc is not None and "bitsandbytes" in str(last_exc):
285 if importlib.util.find_spec("bitsandbytes") is None: 285 ↛ 297line 285 didn't jump to line 297 because the condition on line 285 was always true
286 raise ImportError(
287 "Baichuan2 variants require `bitsandbytes` for "
288 "trust_remote_code loading (their shipped quantizer.py "
289 "imports it). Install the quantization extras: "
290 "`uv sync --group quantization`."
291 ) from last_exc
292 except ImportError:
293 raise
294 except Exception:
295 pass
297 _patch_init_weights_for_baichuan()
299 def prepare_model(self, hf_model: Any) -> None:
300 """Fix rotary caches and normalize NormHead weights before bridge creation.
302 RotaryEmbedding differs between v1 and v2:
303 - v1 (Baichuan-7B): `inv_freq` is a persistent buffer, loaded from the
304 checkpoint as bfloat16, but `cos_cached`/`sin_cached` are non-persistent
305 and materialize as garbage under meta-init.
306 - v2 (Baichuan2-*): `inv_freq`, `cos_cached`, `sin_cached` are all plain
307 attributes (no `register_buffer`). v5's meta-init materializes them on
308 meta, and nothing in the checkpoint overwrites them.
310 Both cases are resolved by computing inv_freq + caches from scratch at
311 float32 using config-derived head_dim and base=10000. Recomputing v1 at
312 float32 is also an upgrade over its bfloat16 checkpoint values.
314 Baichuan2 Chat also uses NormHead which row-normalizes lm_head during
315 forward. We apply that once here so the bridge sees the normalized
316 weights directly without needing NormHead's forward path.
317 """
318 # Pick a real device/dtype by scanning real (non-meta) parameters.
319 target_device = torch.device("cpu")
320 params_fn = getattr(hf_model, "parameters", None)
321 if callable(params_fn): 321 ↛ 322line 321 didn't jump to line 322 because the condition on line 321 was never true
322 for param in params_fn():
323 if param.device.type != "meta":
324 target_device = param.device
325 break
327 head_dim = self.cfg.d_model // self.cfg.n_heads
328 base = 10000.0
330 model_core = getattr(hf_model, "model", None)
331 if model_core is not None:
332 for layer in getattr(model_core, "layers", []):
333 rotary = getattr(getattr(layer, "self_attn", None), "rotary_emb", None)
334 if rotary is None: 334 ↛ 335line 334 didn't jump to line 335 because the condition on line 334 was never true
335 continue
336 max_seq = getattr(rotary, "max_seq_len_cached", self.cfg.n_ctx or 4096)
337 inv_freq = 1.0 / (
338 base
339 ** (
340 torch.arange(0, head_dim, 2, device=target_device, dtype=torch.float32)
341 / head_dim
342 )
343 )
344 t = torch.arange(max_seq, device=target_device, dtype=torch.float32)
345 freqs = torch.einsum("i,j->ij", t, inv_freq)
346 emb = torch.cat((freqs, freqs), dim=-1)
347 rotary.inv_freq = inv_freq
348 rotary.cos_cached = emb.cos()[None, None, :, :]
349 rotary.sin_cached = emb.sin()[None, None, :, :]
351 # Normalize NormHead weights (Baichuan2 Chat)
352 lm_head = getattr(hf_model, "lm_head", None)
353 if lm_head is not None and hasattr(lm_head, "first_flag"):
354 w = lm_head.weight.data
355 lm_head.weight.data = torch.nn.functional.normalize(w, dim=-1)