Coverage for transformer_lens/weight_processing.py: 77%
871 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-09-21 19:27 +0000
1"""
2Weight Processing Functions for Transformer Models.
4This module contains all the weight processing functions extracted from HookedTransformer,
5organized into a single ProcessWeights class with static methods. These functions are used
6to modify transformer model weights for better interpretability and analysis.
7"""
8import re
9from typing import Any, Dict, Optional, Union, overload
11import einops
12import torch
14import transformer_lens.utilities as utils
15from transformer_lens.config.transformer_lens_config import TransformerLensConfig
16from transformer_lens.FactoredMatrix import FactoredMatrix
17from transformer_lens.model_bridge.architecture_adapter import ArchitectureAdapter
18from transformer_lens.utilities import filter_dict_by_prefix
19from transformer_lens.utilities.architectures import POST_NORM_ARCHITECTURES
22class ProcessWeights:
23 """
24 A collection of static methods for processing transformer model weights.
26 These methods are extracted from HookedTransformer and provide various weight
27 transformations for improved model interpretability:
28 - LayerNorm folding: Merges LayerNorm parameters into subsequent linear layers
29 - Weight centering: Centers weights that write to the residual stream
30 - Unembed centering: Centers unembedding weights (translation invariant)
31 - Value bias folding: Consolidates value biases into output biases
32 - Attention matrix refactoring: Experimental QK/OV matrix factorization
34 When an architecture adapter is provided, the methods will translate TransformerLens
35 parameter names to the target format (e.g., HuggingFace) for processing.
36 """
38 @staticmethod
39 def _get_param_key(tl_key: str, adapter=None) -> str:
40 """Convert legacy TL key format (W_Q, b_Q) to component-based format (q.weight, q.bias).
42 Args:
43 tl_key: TransformerLens format parameter key (e.g., "blocks.0.attn.W_Q")
44 adapter: Architecture adapter for translating paths
46 Returns:
47 The component-based key (e.g., "blocks.0.attn.q.weight")
48 """
49 if adapter is None:
50 return tl_key
52 return ProcessWeights._prepare_component_path(tl_key)
54 @staticmethod
55 def _prepare_component_path(tl_key: str) -> str:
56 """Map a TransformerLens key to bridge-style component path.
58 Converts TransformerLens weight names (like "W_Q", "b_in") to bridge-style
59 paths (like "q.weight", "in.bias"). The full path is assembled before being
60 passed to the architecture adapter for translation.
62 Args:
63 tl_key: TransformerLens key like "blocks.0.attn.W_Q"
65 Returns:
66 Full path like "blocks.0.attn.q.weight"
67 """
68 suffix_map: Dict[str, str] = {
69 "W_Q": "q.weight",
70 "_W_Q": "q.weight",
71 "b_Q": "q.bias",
72 "_b_Q": "q.bias",
73 "W_K": "k.weight",
74 "_W_K": "k.weight",
75 "b_K": "k.bias",
76 "_b_K": "k.bias",
77 "W_V": "v.weight",
78 "_W_V": "v.weight",
79 "b_V": "v.bias",
80 "_b_V": "v.bias",
81 "W_O": "o.weight",
82 "b_O": "o.bias",
83 "W_in": "in.weight",
84 "b_in": "in.bias",
85 "W_gate": "gate.weight",
86 "b_gate": "gate.bias",
87 "W_out": "out.weight",
88 "b_out": "out.bias",
89 "W_E": "weight",
90 "b_E": "bias",
91 "W_pos": "weight",
92 "b_pos": "bias",
93 "W_U": "weight",
94 "b_U": "bias",
95 "w": "weight",
96 "b": "bias",
97 "weight": "weight",
98 "bias": "bias",
99 }
100 if "." not in tl_key: 100 ↛ 101line 100 didn't jump to line 101 because the condition on line 100 was never true
101 return tl_key
102 base_path, suffix = tl_key.rsplit(".", 1)
103 if suffix in suffix_map: 103 ↛ 106line 103 didn't jump to line 106 because the condition on line 103 was always true
104 replacement = suffix_map[suffix]
105 return f"{base_path}.{replacement}"
106 return tl_key
108 @staticmethod
109 def _resolve_state_dict_key(
110 state_dict: Dict[str, torch.Tensor],
111 key: str,
112 layer: Optional[int] = None,
113 ) -> str:
114 """Resolve a bridge-style key to the actual key in the state_dict.
116 Some architectures (e.g., BD3LM's symbolic attention) store parameters
117 with HF-style prefixes instead of bridge-style prefixes. This method
118 handles the key resolution by falling back to a suffix search.
120 Args:
121 state_dict: Model state dictionary
122 key: The expected key (e.g., "blocks.0.mlp.in.weight")
123 layer: Optional layer index for layer-specific searches
125 Returns:
126 The actual key found in state_dict, or the original key if no match
127 """
128 if key in state_dict:
129 return key
131 # Extract the component path after "blocks.{i}."
132 import re
134 match = re.match(r"blocks\.(\d+)\.(.*)", key)
135 if match: 135 ↛ 151line 135 didn't jump to line 151 because the condition on line 135 was always true
136 layer_idx = match.group(1)
137 component_suffix = match.group(2)
138 # Search for keys ending with the component suffix that include the layer index
139 for sd_key in state_dict:
140 if sd_key.endswith(f".{component_suffix}") and f".{layer_idx}." in sd_key: 140 ↛ 141line 140 didn't jump to line 141 because the condition on line 140 was never true
141 return sd_key
143 # Mixed MoE/dense adapters (Llama4, Laguna, LLaDA2-MoE) put a non-MoE
144 # layer's gated MLP in the MoE slot under dense_gate/dense_in/dense_out
145 # so it coexists with the router/experts. When the standard block-level
146 # mlp.{gate,in,out} key is absent, fall back to that dense_ variant.
147 # Fails closed: the standard resolution above already ran (so a real
148 # mlp.gate router is never shadowed), the match is anchored to a block's
149 # top-level MLP (not any nested `.mlp.`), and the dense key is returned
150 # only if it exists — models without the convention are untouched.
151 proj_match = re.match(r"(blocks\.\d+\.mlp\.)(in|gate|out)(\..+)$", key)
152 if proj_match:
153 dense_key = f"{proj_match.group(1)}dense_{proj_match.group(2)}{proj_match.group(3)}"
154 if dense_key in state_dict:
155 return dense_key
156 dm = re.match(r"blocks\.(\d+)\.(.*)", dense_key)
157 if dm: 157 ↛ 162line 157 didn't jump to line 162 because the condition on line 157 was always true
158 for sd_key in state_dict:
159 if sd_key.endswith(f".{dm.group(2)}") and f".{dm.group(1)}." in sd_key: 159 ↛ 160line 159 didn't jump to line 160 because the condition on line 159 was never true
160 return sd_key
162 return key
164 @staticmethod
165 def _safe_get_tensor(
166 state_dict: Dict[str, torch.Tensor],
167 tl_key: str,
168 adapter=None,
169 default: Optional[torch.Tensor] = None,
170 ) -> Optional[torch.Tensor]:
171 """Safely get a tensor from state_dict, handling optional parameters.
173 This is the recommended way to access parameters that may not exist in all architectures
174 (e.g., biases in Qwen2/LLaMA/Gemma). Returns None if the parameter doesn't exist,
175 rather than raising a KeyError.
177 Args:
178 state_dict: Model state dictionary
179 tl_key: TransformerLens format parameter key (e.g., "blocks.0.attn.b_Q")
180 adapter: Optional architecture adapter for key translation
181 default: Optional default value to return if key not found (defaults to None)
183 Returns:
184 The tensor if found, otherwise the default value (None if not specified)
186 Examples:
187 # Get optional bias (may be None for Qwen2/LLaMA)
188 b_Q = ProcessWeights._safe_get_tensor(state_dict, "blocks.0.attn.b_Q", adapter)
190 # Get required weight (will be None if missing, can check explicitly)
191 W_Q = ProcessWeights._safe_get_tensor(state_dict, "blocks.0.attn.W_Q", adapter)
192 if W_Q is None:
193 raise ValueError("Required weight W_Q not found")
194 """
195 actual_key = ProcessWeights._get_param_key(tl_key, adapter)
196 return state_dict.get(actual_key, default)
198 @staticmethod
199 def fold_layer_norm_bias_single(
200 w_tensor: torch.Tensor, b_tensor: torch.Tensor, ln_bias: torch.Tensor
201 ) -> torch.Tensor:
202 """Fold LayerNorm bias into a single attention bias.
204 Args:
205 w_tensor: Weight tensor [n_heads, d_model, d_head]
206 b_tensor: Bias tensor [n_heads, d_head]
207 ln_bias: LayerNorm bias [d_model]
209 Returns:
210 New bias tensor with folded LayerNorm bias
211 """
212 return b_tensor + (w_tensor * ln_bias[None, :, None]).sum(-2)
214 @staticmethod
215 def fold_layer_norm_weight_single(
216 w_tensor: torch.Tensor, ln_weight: torch.Tensor
217 ) -> torch.Tensor:
218 """Fold LayerNorm weight into a single attention weight.
220 Args:
221 w_tensor: Weight tensor [n_heads, d_model, d_head]
222 ln_weight: LayerNorm weight [d_model]
224 Returns:
225 New weight tensor with folded LayerNorm weight
226 """
227 return w_tensor * ln_weight[None, :, None]
229 @staticmethod
230 def center_weight_single(w_tensor: torch.Tensor) -> torch.Tensor:
231 """Center a single attention weight by subtracting the mean.
233 Args:
234 w_tensor: Weight tensor [n_heads, d_model, d_head]
236 Returns:
237 Centered weight tensor
238 """
239 return w_tensor - einops.reduce(
240 w_tensor, "head_index d_model d_head -> head_index 1 d_head", "mean"
241 )
243 @staticmethod
244 def fold_layer_norm_biases(
245 wq_tensor: torch.Tensor,
246 wk_tensor: torch.Tensor,
247 wv_tensor: torch.Tensor,
248 bq_tensor: Optional[torch.Tensor],
249 bk_tensor: Optional[torch.Tensor],
250 bv_tensor: Optional[torch.Tensor],
251 ln_bias: torch.Tensor,
252 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
253 """Fold LayerNorm bias into attention biases.
255 When QKV biases don't exist (e.g., GPT-Neo), creates zero-initialized biases
256 to absorb the LN bias contribution, similar to how MLP folding handles missing biases.
258 Args:
259 wq_tensor, wk_tensor, wv_tensor: Weight tensors [n_heads, d_model, d_head]
260 bq_tensor, bk_tensor, bv_tensor: Bias tensors [n_heads, d_head] or None if no bias
261 ln_bias: LayerNorm bias [d_model]
263 Returns:
264 Tuple of (new_bq, new_bk, new_bv) with folded biases (always non-None)
265 """
267 def _zero_bias(w: torch.Tensor) -> torch.Tensor:
268 return torch.zeros(w.shape[0], w.shape[2], dtype=w.dtype, device=w.device)
270 new_bq = ProcessWeights.fold_layer_norm_bias_single(
271 wq_tensor, bq_tensor if bq_tensor is not None else _zero_bias(wq_tensor), ln_bias
272 )
273 new_bk = ProcessWeights.fold_layer_norm_bias_single(
274 wk_tensor, bk_tensor if bk_tensor is not None else _zero_bias(wk_tensor), ln_bias
275 )
276 new_bv = ProcessWeights.fold_layer_norm_bias_single(
277 wv_tensor, bv_tensor if bv_tensor is not None else _zero_bias(wv_tensor), ln_bias
278 )
279 return (new_bq, new_bk, new_bv)
281 @staticmethod
282 def fold_layer_norm_weights(
283 wq_tensor: torch.Tensor,
284 wk_tensor: torch.Tensor,
285 wv_tensor: torch.Tensor,
286 ln_weight: torch.Tensor,
287 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
288 """Fold LayerNorm weight into attention weights.
290 Args:
291 wq_tensor, wk_tensor, wv_tensor: Weight tensors [n_heads, d_model, d_head]
292 ln_weight: LayerNorm weight [d_model]
294 Returns:
295 Tuple of (new_wq, new_wk, new_wv) with folded weights
296 """
297 new_wq = ProcessWeights.fold_layer_norm_weight_single(wq_tensor, ln_weight)
298 new_wk = ProcessWeights.fold_layer_norm_weight_single(wk_tensor, ln_weight)
299 new_wv = ProcessWeights.fold_layer_norm_weight_single(wv_tensor, ln_weight)
300 return (new_wq, new_wk, new_wv)
302 @staticmethod
303 def center_attention_weights(
304 wq_tensor: torch.Tensor, wk_tensor: torch.Tensor, wv_tensor: torch.Tensor
305 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
306 """Center attention weights by subtracting the mean.
308 Args:
309 wq_tensor, wk_tensor, wv_tensor: Weight tensors [n_heads, d_model, d_head]
311 Returns:
312 Tuple of (centered_wq, centered_wk, centered_wv)
313 """
314 centered_wq = ProcessWeights.center_weight_single(wq_tensor)
315 centered_wk = ProcessWeights.center_weight_single(wk_tensor)
316 centered_wv = ProcessWeights.center_weight_single(wv_tensor)
317 return (centered_wq, centered_wk, centered_wv)
319 @staticmethod
320 def extract_attention_tensors_for_folding(
321 state_dict: Dict[str, torch.Tensor], cfg, layer: int, adapter
322 ) -> Dict[str, Union[torch.Tensor, None, Dict[str, str]]]:
323 """Extract attention tensors in TransformerLens format for layer norm folding.
325 Args:
326 state_dict: The state dictionary containing tensors
327 cfg: Model configuration object
328 layer: Layer index
329 adapter: Optional architecture adapter for parameter key translation
331 Returns:
332 Dictionary with keys: 'wq', 'wk', 'wv', 'bq', 'bk', 'bv', 'ln1_b', 'ln1_w'
333 All tensors are in TransformerLens format for consistent processing
334 """
335 b_Q_key = ProcessWeights._get_param_key(f"blocks.{layer}.attn.b_Q", adapter)
336 W_Q_key = ProcessWeights._get_param_key(f"blocks.{layer}.attn.W_Q", adapter)
337 b_K_key = ProcessWeights._get_param_key(f"blocks.{layer}.attn.b_K", adapter)
338 W_K_key = ProcessWeights._get_param_key(f"blocks.{layer}.attn.W_K", adapter)
339 b_V_key = ProcessWeights._get_param_key(f"blocks.{layer}.attn.b_V", adapter)
340 W_V_key = ProcessWeights._get_param_key(f"blocks.{layer}.attn.W_V", adapter)
341 ln1_b_key = ProcessWeights._get_param_key(f"blocks.{layer}.ln1.b", adapter)
342 ln1_w_key = ProcessWeights._get_param_key(f"blocks.{layer}.ln1.w", adapter)
344 # For GQA models, Q, K and V weights may use underscore prefix (_W_Q, _W_K, _W_V)
345 # Check if standard keys exist, otherwise update to use underscore-prefixed versions
346 if W_Q_key not in state_dict:
347 W_Q_key = W_Q_key.replace(".W_Q", "._W_Q")
348 if W_K_key not in state_dict:
349 W_K_key = W_K_key.replace(".W_K", "._W_K")
350 if W_V_key not in state_dict:
351 W_V_key = W_V_key.replace(".W_V", "._W_V")
352 if b_Q_key not in state_dict:
353 b_Q_key = b_Q_key.replace(".b_Q", "._b_Q")
354 if b_K_key not in state_dict:
355 b_K_key = b_K_key.replace(".b_K", "._b_K")
356 if b_V_key not in state_dict:
357 b_V_key = b_V_key.replace(".b_V", "._b_V")
359 wq_tensor: Optional[torch.Tensor] = state_dict.get(W_Q_key)
360 wk_tensor: Optional[torch.Tensor] = state_dict.get(W_K_key)
361 wv_tensor: Optional[torch.Tensor] = state_dict.get(W_V_key)
362 bq_tensor: Optional[torch.Tensor] = state_dict.get(b_Q_key)
363 bk_tensor: Optional[torch.Tensor] = state_dict.get(b_K_key)
364 bv_tensor: Optional[torch.Tensor] = state_dict.get(b_V_key)
365 ln1_b = state_dict.get(ln1_b_key, None)
366 ln1_w = state_dict.get(ln1_w_key, None)
367 if adapter:
368 wq_tensor = ProcessWeights.convert_tensor_to_tl_format(
369 W_Q_key, state_dict, wq_tensor, cfg, adapter, layer
370 )
371 wk_tensor = ProcessWeights.convert_tensor_to_tl_format(
372 W_K_key, state_dict, wk_tensor, cfg, adapter, layer
373 )
374 wv_tensor = ProcessWeights.convert_tensor_to_tl_format(
375 W_V_key, state_dict, wv_tensor, cfg, adapter, layer
376 )
377 bq_tensor = ProcessWeights.convert_tensor_to_tl_format(
378 b_Q_key, state_dict, bq_tensor, cfg, adapter, layer
379 )
380 bk_tensor = ProcessWeights.convert_tensor_to_tl_format(
381 b_K_key, state_dict, bk_tensor, cfg, adapter, layer
382 )
383 bv_tensor = ProcessWeights.convert_tensor_to_tl_format(
384 b_V_key, state_dict, bv_tensor, cfg, adapter, layer
385 )
387 # Auto-reshape 1D biases for 3D weights (e.g., OPT)
388 def _reshape_bias_if_needed(bias, weight):
389 if bias is not None and weight is not None:
390 if len(weight.shape) == 3 and len(bias.shape) == 1:
391 n_heads = weight.shape[0]
392 d_head = weight.shape[2]
393 if bias.shape[0] == n_heads * d_head: 393 ↛ 395line 393 didn't jump to line 395 because the condition on line 393 was always true
394 return bias.reshape(n_heads, d_head)
395 return bias
397 bq_tensor = _reshape_bias_if_needed(bq_tensor, wq_tensor)
398 bk_tensor = _reshape_bias_if_needed(bk_tensor, wk_tensor)
399 bv_tensor = _reshape_bias_if_needed(bv_tensor, wv_tensor)
401 return {
402 "wq": wq_tensor,
403 "wk": wk_tensor,
404 "wv": wv_tensor,
405 "bq": bq_tensor,
406 "bk": bk_tensor,
407 "bv": bv_tensor,
408 "ln1_b": ln1_b,
409 "ln1_w": ln1_w,
410 "keys": {
411 "W_Q": W_Q_key,
412 "W_K": W_K_key,
413 "W_V": W_V_key,
414 "b_Q": b_Q_key,
415 "b_K": b_K_key,
416 "b_V": b_V_key,
417 "ln1_b": ln1_b_key,
418 "ln1_w": ln1_w_key,
419 },
420 }
422 @staticmethod
423 def _fold_short_conv_layer_norm(
424 state_dict: Dict[str, torch.Tensor],
425 cfg,
426 layer: int,
427 fold_biases: bool,
428 adapter,
429 ln1_w: torch.Tensor,
430 ln1_b: Optional[torch.Tensor],
431 keys: Dict[str, str],
432 ) -> bool:
433 """Fold ln1 into a short-conv mixer's in-projection. Returns whether it folded.
435 Hybrid stacks (LFM2) interleave attention layers with short-conv layers that
436 have no q_proj, so the attention path skips them and their gain would survive
437 into an otherwise-folded model — one basis per layer type. The in-projection
438 is the mixer's only reader of the norm output (the conv and the gates consume
439 its outputs), so scaling its input columns is exact despite the mixer being
440 quadratic in that projection.
441 """
442 w_key = ProcessWeights._get_param_key(f"blocks.{layer}.conv.W_in", adapter)
443 conv_W_in = state_dict.get(w_key)
444 if conv_W_in is None: 444 ↛ 445line 444 didn't jump to line 445 because the condition on line 444 was never true
445 return False
447 rmsnorm_uses_offset = getattr(cfg, "rmsnorm_uses_offset", False)
448 effective_ln1_w = (1.0 + ln1_w) if rmsnorm_uses_offset else ln1_w
449 if conv_W_in.shape[-1] == effective_ln1_w.shape[0]: 449 ↛ 452line 449 didn't jump to line 452 because the condition on line 449 was always true
450 ln1_w_broadcast = effective_ln1_w[None, :]
451 sum_dim = -1
452 elif conv_W_in.shape[0] == effective_ln1_w.shape[0]:
453 ln1_w_broadcast = effective_ln1_w[:, None]
454 sum_dim = -2
455 else:
456 return False
458 b_key = ProcessWeights._get_param_key(f"blocks.{layer}.conv.b_in", adapter)
459 if fold_biases and ln1_b is not None: 459 ↛ 462line 459 didn't jump to line 462 because the condition on line 459 was never true
460 # Absorbing the shift needs somewhere to put it; fabricating a bias the
461 # mixer does not have would not survive distribution to the component.
462 if b_key not in state_dict:
463 return False
464 ln1_b_broadcast = ln1_b[None, :] if sum_dim == -1 else ln1_b[:, None]
465 state_dict[b_key] = state_dict[b_key] + (conv_W_in * ln1_b_broadcast).sum(sum_dim)
466 state_dict[keys["ln1_b"]] = torch.zeros_like(ln1_b)
468 state_dict[w_key] = conv_W_in * ln1_w_broadcast
469 state_dict[keys["ln1_w"]] = (
470 torch.zeros_like(ln1_w) if rmsnorm_uses_offset else torch.ones_like(ln1_w)
471 )
472 return True
474 @staticmethod
475 def _fold_layer(
476 state_dict: Dict[str, torch.Tensor],
477 cfg,
478 layer_idx: int,
479 fold_biases: bool,
480 center_weights: bool,
481 adapter,
482 gqa: str,
483 ) -> Dict[str, torch.Tensor]:
484 """Fold LayerNorm for a single layer.
486 Args:
487 state_dict: The state dictionary to process (modified in place)
488 cfg: Model configuration object
489 layer_idx: The layer index to process
490 fold_biases: Whether to fold LayerNorm biases
491 center_weights: Whether to center weights after folding
492 adapter: Optional architecture adapter for parameter key translation
493 gqa: GQA prefix string (empty or "_")
494 """
495 layer = layer_idx
496 tensors = ProcessWeights.extract_attention_tensors_for_folding(
497 state_dict, cfg, layer, adapter
498 )
499 wq_tensor = tensors["wq"]
500 wk_tensor = tensors["wk"]
501 wv_tensor = tensors["wv"]
502 bq_tensor = tensors["bq"]
503 bk_tensor = tensors["bk"]
504 bv_tensor = tensors["bv"]
505 ln1_b = tensors["ln1_b"]
506 ln1_w = tensors["ln1_w"]
507 keys = tensors["keys"]
509 # Fold LN into QKV (skip if combined QKV, e.g., OpenELM)
510 if wq_tensor is not None:
511 assert isinstance(wq_tensor, torch.Tensor)
512 assert isinstance(keys, dict)
513 if wk_tensor is not None: 513 ↛ 515line 513 didn't jump to line 515 because the condition on line 513 was always true
514 assert isinstance(wk_tensor, torch.Tensor)
515 if wv_tensor is not None: 515 ↛ 517line 515 didn't jump to line 517 because the condition on line 515 was always true
516 assert isinstance(wv_tensor, torch.Tensor)
517 if bq_tensor is not None:
518 assert isinstance(bq_tensor, torch.Tensor)
519 if bk_tensor is not None:
520 assert isinstance(bk_tensor, torch.Tensor)
521 if bv_tensor is not None:
522 assert isinstance(bv_tensor, torch.Tensor)
523 # RMS norm (Gemma): ln1_b may be None, only ln1_w required
524 if ln1_w is not None:
525 assert isinstance(ln1_w, torch.Tensor)
526 # Fold biases if present (RMS norm has none; missing QKV biases get zeros)
527 if fold_biases and ln1_b is not None:
528 assert isinstance(ln1_b, torch.Tensor)
529 assert wq_tensor is not None
530 assert wk_tensor is not None
531 assert wv_tensor is not None
532 bq_tensor, bk_tensor, bv_tensor = ProcessWeights.fold_layer_norm_biases(
533 wq_tensor, wk_tensor, wv_tensor, bq_tensor, bk_tensor, bv_tensor, ln1_b
534 )
535 if keys["ln1_b"] in state_dict: 535 ↛ 537line 535 didn't jump to line 537 because the condition on line 535 was always true
536 state_dict[keys["ln1_b"]] = torch.zeros_like(ln1_b)
537 alternate_b_key = (
538 keys["ln1_b"].replace("ln_1", "ln1")
539 if "ln_1" in keys["ln1_b"]
540 else keys["ln1_b"].replace("ln1", "ln_1")
541 )
542 if alternate_b_key != keys["ln1_b"] and alternate_b_key in state_dict: 542 ↛ 543line 542 didn't jump to line 543 because the condition on line 542 was never true
543 state_dict[alternate_b_key] = torch.zeros_like(ln1_b)
544 # Fold ln1_w; use (1+w) for rmsnorm_uses_offset (Gemma), then set to identity
545 rmsnorm_uses_offset = getattr(cfg, "rmsnorm_uses_offset", False)
546 effective_ln1_w = (1.0 + ln1_w) if rmsnorm_uses_offset else ln1_w
547 if wk_tensor is not None and wv_tensor is not None: 547 ↛ 552line 547 didn't jump to line 552 because the condition on line 547 was always true
548 wq_tensor, wk_tensor, wv_tensor = ProcessWeights.fold_layer_norm_weights(
549 wq_tensor, wk_tensor, wv_tensor, effective_ln1_w
550 )
551 # Set ln1.w to identity: ones (standard) or zeros (rmsnorm_uses_offset)
552 identity_val = (
553 torch.zeros_like(ln1_w) if rmsnorm_uses_offset else torch.ones_like(ln1_w)
554 )
555 if keys["ln1_w"] in state_dict: 555 ↛ 557line 555 didn't jump to line 557 because the condition on line 555 was always true
556 state_dict[keys["ln1_w"]] = identity_val
557 alternate_w_key = (
558 keys["ln1_w"].replace("ln_1", "ln1")
559 if "ln_1" in keys["ln1_w"]
560 else keys["ln1_w"].replace("ln1", "ln_1")
561 )
562 if alternate_w_key != keys["ln1_w"] and alternate_w_key in state_dict: 562 ↛ 563line 562 didn't jump to line 563 because the condition on line 562 was never true
563 state_dict[alternate_w_key] = identity_val
564 if center_weights and wk_tensor is not None and (wv_tensor is not None):
565 wq_tensor, wk_tensor, wv_tensor = ProcessWeights.center_attention_weights(
566 wq_tensor, wk_tensor, wv_tensor
567 )
568 state_dict = ProcessWeights._store_processed_attention_tensors(
569 state_dict,
570 keys,
571 wq_tensor,
572 wk_tensor,
573 wv_tensor,
574 bq_tensor,
575 bk_tensor,
576 bv_tensor,
577 adapter,
578 cfg,
579 layer,
580 )
581 elif ln1_w is not None: 581 ↛ 592line 581 didn't jump to line 592 because the condition on line 581 was always true
582 assert isinstance(ln1_w, torch.Tensor)
583 assert ln1_b is None or isinstance(ln1_b, torch.Tensor)
584 assert isinstance(keys, dict)
585 ProcessWeights._fold_short_conv_layer_norm(
586 state_dict, cfg, layer, fold_biases, adapter, ln1_w, ln1_b, keys
587 )
589 # ln1_post.w (Gemma 2/3): keep original; independent post-attention normalization
591 # Fold MLP LN: shared ln1 (Phi-2, GPT-J) or separate ln2 (Pythia)
592 if getattr(cfg, "parallel_attn_mlp", False) and ln1_w is not None: 592 ↛ 594line 592 didn't jump to line 594 because the condition on line 592 was never true
593 # Check if a separate ln2 exists for this layer
594 ln2_check_key = ProcessWeights._resolve_state_dict_key(
595 state_dict,
596 ProcessWeights._get_param_key(f"blocks.{layer_idx}.ln2.w", adapter),
597 layer_idx,
598 )
599 if ln2_check_key in state_dict:
600 # Separate ln2 (e.g., GPT-NeoX/Pythia) — fold ln2 → MLP normally
601 state_dict = ProcessWeights._fold_mlp_layer_norm(
602 state_dict, cfg, layer, fold_biases, center_weights, adapter
603 )
604 else:
605 # Shared ln1 (e.g., Phi-2, GPT-J) — fold ln1 → MLP via override
606 assert isinstance(ln1_w, torch.Tensor)
607 assert ln1_b is None or isinstance(ln1_b, torch.Tensor)
608 state_dict = ProcessWeights._fold_mlp_layer_norm(
609 state_dict,
610 cfg,
611 layer,
612 fold_biases,
613 center_weights,
614 adapter,
615 override_ln_w=ln1_w,
616 override_ln_b=ln1_b,
617 )
618 else:
619 state_dict = ProcessWeights._fold_mlp_layer_norm(
620 state_dict, cfg, layer, fold_biases, center_weights, adapter
621 )
623 return state_dict
625 @staticmethod
626 def _fold_mlp_layer_norm(
627 state_dict: Dict[str, torch.Tensor],
628 cfg,
629 layer: int,
630 fold_biases: bool,
631 center_weights: bool,
632 adapter,
633 override_ln_w: Optional[torch.Tensor] = None,
634 override_ln_b: Optional[torch.Tensor] = None,
635 ) -> Dict[str, torch.Tensor]:
636 """Fold LayerNorm into MLP layer.
638 Args:
639 state_dict: The state dictionary to process (modified in place)
640 cfg: Model configuration object
641 layer: The layer index to process
642 fold_biases: Whether to fold LayerNorm biases
643 center_weights: Whether to center weights after folding
644 adapter: Optional architecture adapter for parameter key translation
645 override_ln_w: Override LN weight tensor. Used for parallel architectures
646 where MLP reads from ln1 (same as attention) instead of a separate ln2.
647 override_ln_b: Override LN bias tensor. Used with override_ln_w.
648 """
649 if getattr(cfg, "attn_only", False):
650 return state_dict
652 mlp_b_in_key = ProcessWeights._resolve_state_dict_key(
653 state_dict, ProcessWeights._get_param_key(f"blocks.{layer}.mlp.b_in", adapter), layer
654 )
655 mlp_W_in_key = ProcessWeights._resolve_state_dict_key(
656 state_dict, ProcessWeights._get_param_key(f"blocks.{layer}.mlp.W_in", adapter), layer
657 )
658 mlp_W_gate_key = (
659 ProcessWeights._resolve_state_dict_key(
660 state_dict,
661 ProcessWeights._get_param_key(f"blocks.{layer}.mlp.W_gate", adapter),
662 layer,
663 )
664 if getattr(cfg, "gated_mlp", False)
665 else None
666 )
667 mlp_b_gate_key = (
668 ProcessWeights._resolve_state_dict_key(
669 state_dict,
670 ProcessWeights._get_param_key(f"blocks.{layer}.mlp.b_gate", adapter),
671 layer,
672 )
673 if getattr(cfg, "gated_mlp", False)
674 else None
675 )
677 # For parallel architectures, ln1 values are passed via override params.
678 # Otherwise, look up ln2 from the state dict.
679 ln2_w: Optional[torch.Tensor]
680 ln2_b: Optional[torch.Tensor]
681 if override_ln_w is not None: 681 ↛ 682line 681 didn't jump to line 682 because the condition on line 681 was never true
682 ln2_w = override_ln_w
683 ln2_b = override_ln_b
684 ln2_w_key = None # No state dict key to zero out (already done by attention folding)
685 ln2_b_key = None
686 has_ln = True
687 else:
688 ln2_b_key = ProcessWeights._resolve_state_dict_key(
689 state_dict, ProcessWeights._get_param_key(f"blocks.{layer}.ln2.b", adapter), layer
690 )
691 ln2_w_key = ProcessWeights._resolve_state_dict_key(
692 state_dict, ProcessWeights._get_param_key(f"blocks.{layer}.ln2.w", adapter), layer
693 )
694 has_ln = ln2_w_key in state_dict
695 ln2_w = state_dict.get(ln2_w_key) if has_ln else None
696 ln2_b = state_dict.get(ln2_b_key) if has_ln else None
698 # CRITICAL FIX: For RMS norm (Gemma), ln2_b doesn't exist. Only require ln2_w!
699 if has_ln and ln2_w is not None:
700 # MoE layers: fold ln2 into router gate and each expert's W_in/W_gate
701 if getattr(cfg, "num_experts", None) is not None and cfg.num_experts > 0:
702 # Every expert reads ln2's output, so folding into only some of them
703 # would change the model — collect them all before touching anything.
704 expert_keys = [
705 key
706 for e in range(cfg.num_experts)
707 for suffix in ("W_in.weight", "W_gate.weight")
708 for key in (
709 ProcessWeights._resolve_state_dict_key(
710 state_dict, f"blocks.{layer}.mlp.experts.{e}.{suffix}", layer
711 ),
712 )
713 if key in state_dict
714 ]
715 if expert_keys:
716 router_key = ProcessWeights._resolve_state_dict_key(
717 state_dict, f"blocks.{layer}.mlp.W_gate.weight", layer
718 )
719 if router_key in state_dict: 719 ↛ 721line 719 didn't jump to line 721 because the condition on line 719 was always true
720 state_dict[router_key] = state_dict[router_key] * ln2_w[None, :]
721 for key in expert_keys:
722 state_dict[key] = state_dict[key] * ln2_w[None, :]
723 if ln2_w_key is not None: 723 ↛ 732line 723 didn't jump to line 732 because the condition on line 723 was always true
724 state_dict[ln2_w_key] = torch.ones_like(ln2_w)
725 alternate_ln2_w_key = (
726 ln2_w_key.replace("ln_2", "ln2")
727 if "ln_2" in ln2_w_key
728 else ln2_w_key.replace("ln2", "ln_2")
729 )
730 if alternate_ln2_w_key != ln2_w_key and alternate_ln2_w_key in state_dict: 730 ↛ 731line 730 didn't jump to line 731 because the condition on line 730 was never true
731 state_dict[alternate_ln2_w_key] = torch.ones_like(ln2_w)
732 return state_dict
734 # A dense layer of an interleaved stack (Llama4, Laguna, LLaDA2-MoE,
735 # LFM2-MoE) parks its gated MLP under dense_in/dense_gate and has no
736 # experts of its own, so it folds exactly like a non-MoE layer. Anything
737 # else here is a sparse layer whose expert weights this state dict does
738 # not carry — leave its ln2 alone rather than fold half the readers.
739 if not mlp_W_in_key.endswith(".mlp.dense_in.weight"):
740 return state_dict
742 mlp_W_in = ProcessWeights.convert_tensor_to_tl_format(
743 mlp_W_in_key, state_dict, state_dict.get(mlp_W_in_key), cfg, adapter, layer
744 )
745 assert mlp_W_in is not None, f"MLP W_in not found at key {mlp_W_in_key}"
746 # rmsnorm_uses_offset: effective scale is (1+w), identity is 0.0
747 rmsnorm_uses_offset = getattr(cfg, "rmsnorm_uses_offset", False)
748 effective_ln2_w = (1.0 + ln2_w) if rmsnorm_uses_offset else ln2_w
749 if mlp_W_in.shape[1] == effective_ln2_w.shape[0]:
750 ln2_w_broadcast = effective_ln2_w[None, :]
751 sum_dim = -1
752 if ln2_b is not None:
753 ln2_b_broadcast = ln2_b[None, :]
754 elif mlp_W_in.shape[0] == effective_ln2_w.shape[0]: 754 ↛ 760line 754 didn't jump to line 760 because the condition on line 754 was always true
755 ln2_w_broadcast = effective_ln2_w[:, None]
756 sum_dim = -2
757 if ln2_b is not None:
758 ln2_b_broadcast = ln2_b[:, None]
759 else:
760 raise ValueError(
761 f"Cannot broadcast MLP weight {mlp_W_in.shape} with layer norm weight {effective_ln2_w.shape}"
762 )
763 # Only fold biases if they exist (LayerNorm). RMS norm has no biases.
764 if fold_biases and ln2_b is not None:
765 mlp_b_in = ProcessWeights.convert_tensor_to_tl_format(
766 mlp_b_in_key, state_dict, state_dict.get(mlp_b_in_key), cfg, adapter, layer
767 )
768 ln2_b_folded = (mlp_W_in * ln2_b_broadcast).sum(sum_dim)
769 if mlp_b_in is not None: 769 ↛ 773line 769 didn't jump to line 773 because the condition on line 769 was always true
770 new_mlp_b_in = mlp_b_in + ln2_b_folded
771 else:
772 # MLP has no bias — create one from the folded LN bias
773 new_mlp_b_in = ln2_b_folded
774 state_dict[mlp_b_in_key] = ProcessWeights.convert_tensor_to_hf_format(
775 mlp_b_in_key, new_mlp_b_in, cfg, adapter, layer
776 )
777 # Set ln2.b to zero (skip for parallel override — ln1 already zeroed)
778 if ln2_b_key is not None: 778 ↛ 787line 778 didn't jump to line 787 because the condition on line 778 was always true
779 state_dict[ln2_b_key] = torch.zeros_like(ln2_b)
780 alternate_ln2_b_key = (
781 ln2_b_key.replace("ln_2", "ln2")
782 if "ln_2" in ln2_b_key
783 else ln2_b_key.replace("ln2", "ln_2")
784 )
785 if alternate_ln2_b_key != ln2_b_key and alternate_ln2_b_key in state_dict: 785 ↛ 786line 785 didn't jump to line 786 because the condition on line 785 was never true
786 state_dict[alternate_ln2_b_key] = torch.zeros_like(ln2_b)
787 new_mlp_W_in = mlp_W_in * ln2_w_broadcast
788 state_dict[mlp_W_in_key] = ProcessWeights.convert_tensor_to_hf_format(
789 mlp_W_in_key, new_mlp_W_in, cfg, adapter, layer
790 )
791 if getattr(cfg, "gated_mlp", False) and mlp_W_gate_key is not None:
792 mlp_W_gate = ProcessWeights.convert_tensor_to_tl_format(
793 mlp_W_gate_key, state_dict, state_dict.get(mlp_W_gate_key), cfg, adapter, layer
794 )
795 # Combined gate+up (OpenELM): no separate gate, already folded above
796 if mlp_W_gate is not None: 796 ↛ 821line 796 didn't jump to line 821 because the condition on line 796 was always true
797 new_mlp_W_gate = mlp_W_gate * ln2_w_broadcast
798 state_dict[mlp_W_gate_key] = ProcessWeights.convert_tensor_to_hf_format(
799 mlp_W_gate_key, new_mlp_W_gate, cfg, adapter, layer
800 )
801 # Also fold ln2 bias into gate bias (mirrors the in-proj bias folding above)
802 if fold_biases and ln2_b is not None and mlp_b_gate_key is not None:
803 mlp_b_gate = ProcessWeights.convert_tensor_to_tl_format(
804 mlp_b_gate_key,
805 state_dict,
806 state_dict.get(mlp_b_gate_key),
807 cfg,
808 adapter,
809 layer,
810 )
811 ln2_b_gate_folded = (mlp_W_gate * ln2_b_broadcast).sum(sum_dim)
812 if mlp_b_gate is not None: 812 ↛ 813line 812 didn't jump to line 813 because the condition on line 812 was never true
813 new_mlp_b_gate = mlp_b_gate + ln2_b_gate_folded
814 else:
815 new_mlp_b_gate = ln2_b_gate_folded
816 state_dict[mlp_b_gate_key] = ProcessWeights.convert_tensor_to_hf_format(
817 mlp_b_gate_key, new_mlp_b_gate, cfg, adapter, layer
818 )
819 # After folding, set ln2.w to identity (skip for parallel override —
820 # ln1 was already set to identity by the attention folding code).
821 if ln2_w_key is not None: 821 ↛ 833line 821 didn't jump to line 833 because the condition on line 821 was always true
822 identity_ln2 = (
823 torch.zeros_like(ln2_w) if rmsnorm_uses_offset else torch.ones_like(ln2_w)
824 )
825 state_dict[ln2_w_key] = identity_ln2
826 alternate_ln2_w_key = (
827 ln2_w_key.replace("ln_2", "ln2")
828 if "ln_2" in ln2_w_key
829 else ln2_w_key.replace("ln2", "ln_2")
830 )
831 if alternate_ln2_w_key != ln2_w_key and alternate_ln2_w_key in state_dict: 831 ↛ 832line 831 didn't jump to line 832 because the condition on line 831 was never true
832 state_dict[alternate_ln2_w_key] = identity_ln2
833 if center_weights and mlp_W_in_key in state_dict:
834 mlp_W_in_centered = ProcessWeights.convert_tensor_to_tl_format(
835 mlp_W_in_key, state_dict, state_dict.get(mlp_W_in_key), cfg, adapter, layer
836 )
837 assert mlp_W_in_centered is not None, f"MLP W_in not found at key {mlp_W_in_key}"
838 # Center along d_model: TL [d_model, d_mlp] or HF [d_mlp, d_model]
839 d_model = cfg.d_model if cfg is not None else None
840 if (
841 d_model is not None
842 and mlp_W_in_centered.shape[0] == d_model
843 and mlp_W_in_centered.shape[-1] != d_model
844 ):
845 # TL format [d_model, d_mlp]
846 mlp_W_in_centered = mlp_W_in_centered - mlp_W_in_centered.mean(0, keepdim=True)
847 elif ( 847 ↛ 856line 847 didn't jump to line 856 because the condition on line 847 was always true
848 d_model is not None
849 and mlp_W_in_centered.shape[-1] == d_model
850 and mlp_W_in_centered.shape[0] != d_model
851 ):
852 # HF format [d_mlp, d_model]
853 mlp_W_in_centered = mlp_W_in_centered - mlp_W_in_centered.mean(-1, keepdim=True)
854 else:
855 # Fallback: assume TL format
856 mlp_W_in_centered = mlp_W_in_centered - mlp_W_in_centered.mean(0, keepdim=True)
857 state_dict[mlp_W_in_key] = ProcessWeights.convert_tensor_to_hf_format(
858 mlp_W_in_key, mlp_W_in_centered, cfg, adapter, layer
859 )
860 if getattr(cfg, "act_fn", None) is not None and cfg.act_fn.startswith("solu"):
861 mlp_b_out_key = ProcessWeights._get_param_key(f"blocks.{layer}.mlp.b_out", adapter)
862 mlp_W_out_key = ProcessWeights._get_param_key(f"blocks.{layer}.mlp.W_out", adapter)
863 mlp_ln_b_key = ProcessWeights._get_param_key(f"blocks.{layer}.mlp.ln.b", adapter)
864 mlp_ln_w_key = ProcessWeights._get_param_key(f"blocks.{layer}.mlp.ln.w", adapter)
866 mlp_b_out = ProcessWeights.convert_tensor_to_tl_format(
867 mlp_b_out_key, state_dict, state_dict.get(mlp_b_out_key), cfg, adapter, layer
868 )
869 mlp_W_out = ProcessWeights.convert_tensor_to_tl_format(
870 mlp_W_out_key, state_dict, state_dict.get(mlp_W_out_key), cfg, adapter, layer
871 )
872 mlp_ln_b = state_dict.get(mlp_ln_b_key)
873 mlp_ln_w = state_dict.get(mlp_ln_w_key)
874 assert mlp_b_out is not None, f"MLP b_out not found at key {mlp_b_out_key}"
875 assert mlp_W_out is not None, f"MLP W_out not found at key {mlp_W_out_key}"
876 assert mlp_ln_b is not None, f"MLP ln.b not found at key {mlp_ln_b_key}"
877 assert mlp_ln_w is not None, f"MLP ln.w not found at key {mlp_ln_w_key}"
879 # TL keys hold W_out as [d_mlp, d_model]; nn.Linear bridges hold [d_model, d_mlp].
880 # The fold and the centering must use the same neuron axis, so resolve it once.
881 d_mlp = mlp_ln_w.shape[0]
882 if mlp_W_out.shape[0] == d_mlp and mlp_W_out.shape[-1] != d_mlp:
883 neuron_dim = 0
884 elif mlp_W_out.shape[-1] == d_mlp and mlp_W_out.shape[0] != d_mlp: 884 ↛ 887line 884 didn't jump to line 887 because the condition on line 884 was always true
885 neuron_dim = -1
886 else:
887 raise ValueError(
888 f"Cannot resolve the neuron axis of MLP W_out {tuple(mlp_W_out.shape)} "
889 f"against the mid-MLP LayerNorm (d_mlp={d_mlp}) at layer {layer}."
890 )
891 ln_shape = (-1, 1) if neuron_dim == 0 else (1, -1)
893 if fold_biases: 893 ↛ 901line 893 didn't jump to line 901 because the condition on line 893 was always true
894 new_mlp_b_out = mlp_b_out + (mlp_W_out * mlp_ln_b.reshape(ln_shape)).sum(neuron_dim)
895 state_dict[mlp_b_out_key] = ProcessWeights.convert_tensor_to_hf_format(
896 mlp_b_out_key, new_mlp_b_out, cfg, adapter, layer
897 )
898 if mlp_ln_b_key in state_dict: 898 ↛ 901line 898 didn't jump to line 901 because the condition on line 898 was always true
899 state_dict[mlp_ln_b_key] = torch.zeros_like(mlp_ln_b)
901 new_mlp_W_out = mlp_W_out * mlp_ln_w.reshape(ln_shape)
903 if center_weights: 903 ↛ 908line 903 didn't jump to line 908 because the condition on line 903 was always true
904 # The folded mid-MLP LayerNorm emits activations that are mean-zero across
905 # neurons, so a per-row constant along that axis never reaches the output.
906 new_mlp_W_out = new_mlp_W_out - new_mlp_W_out.mean(neuron_dim, keepdim=True)
908 state_dict[mlp_W_out_key] = ProcessWeights.convert_tensor_to_hf_format(
909 mlp_W_out_key, new_mlp_W_out, cfg, adapter, layer
910 )
912 if mlp_ln_w_key in state_dict: 912 ↛ 917line 912 didn't jump to line 917 because the condition on line 912 was always true
913 state_dict[mlp_ln_w_key] = torch.ones_like(mlp_ln_w)
915 # ln2_post.w (Gemma 2/3): keep original; independent post-MLP normalization
917 return state_dict
919 @staticmethod
920 def _store_processed_attention_tensors(
921 state_dict: Dict[str, torch.Tensor],
922 keys: Dict[str, str],
923 wq_tensor: Optional[torch.Tensor],
924 wk_tensor: Optional[torch.Tensor],
925 wv_tensor: Optional[torch.Tensor],
926 bq_tensor: Optional[torch.Tensor],
927 bk_tensor: Optional[torch.Tensor],
928 bv_tensor: Optional[torch.Tensor],
929 adapter,
930 cfg,
931 layer: int,
932 ) -> Dict[str, torch.Tensor]:
933 """Store processed attention tensors back to state dict in appropriate format.
935 Args:
936 state_dict: The state dictionary to update (modified in place)
937 keys: Dictionary mapping tensor names to state dict keys
938 wq_tensor, wk_tensor, wv_tensor: Processed attention weight tensors
939 bq_tensor, bk_tensor, bv_tensor: Processed attention bias tensors
940 adapter: Optional architecture adapter for parameter key translation
941 cfg: Model configuration object
942 layer: The layer index
943 """
944 if wq_tensor is None: 944 ↛ 945line 944 didn't jump to line 945 because the condition on line 944 was never true
945 return state_dict
946 wq_key = keys["W_Q"]
947 wk_key = keys["W_K"]
948 wv_key = keys["W_V"]
949 bq_key = keys["b_Q"]
950 bk_key = keys["b_K"]
951 bv_key = keys["b_V"]
953 # Store processed tensors directly in 3D format (set_processed_weights will flatten to 2D)
954 if wq_tensor is None or wk_tensor is None or wv_tensor is None: 954 ↛ 955line 954 didn't jump to line 955 because the condition on line 954 was never true
955 raise ValueError(f"Required attention weights missing for layer {layer}")
956 state_dict[wq_key] = ProcessWeights.convert_tensor_to_hf_format(
957 wq_key, wq_tensor, cfg, adapter, layer_idx=layer
958 )
959 state_dict[wk_key] = ProcessWeights.convert_tensor_to_hf_format(
960 wk_key, wk_tensor, cfg, adapter, layer_idx=layer
961 )
962 state_dict[wv_key] = ProcessWeights.convert_tensor_to_hf_format(
963 wv_key, wv_tensor, cfg, adapter, layer_idx=layer
964 )
965 if bq_tensor is not None:
966 state_dict[bq_key] = ProcessWeights.convert_tensor_to_hf_format(
967 bq_key, bq_tensor, cfg, adapter, layer_idx=layer
968 )
969 if bk_tensor is not None:
970 state_dict[bk_key] = ProcessWeights.convert_tensor_to_hf_format(
971 bk_key, bk_tensor, cfg, adapter, layer_idx=layer
972 )
973 if bv_tensor is not None:
974 state_dict[bv_key] = ProcessWeights.convert_tensor_to_hf_format(
975 bv_key, bv_tensor, cfg, adapter, layer_idx=layer
976 )
978 return state_dict
980 @staticmethod
981 def _fold_unembed_layer_norm(
982 state_dict: Dict[str, torch.Tensor], cfg, fold_biases: bool, center_weights: bool, adapter
983 ) -> Dict[str, torch.Tensor]:
984 """Fold LayerNorm into unembedding layer.
986 Args:
987 state_dict: The state dictionary to process (modified in place)
988 cfg: Model configuration object
989 fold_biases: Whether to fold LayerNorm biases
990 center_weights: Whether to center weights after folding
991 adapter: Optional architecture adapter for parameter key translation
992 """
993 unembed_b_U_key = ProcessWeights._get_param_key("unembed.b_U", adapter)
994 unembed_W_U_key = ProcessWeights._get_param_key("unembed.W_U", adapter)
995 ln_final_b_key = ProcessWeights._get_param_key("ln_final.b", adapter)
996 ln_final_w_key = ProcessWeights._get_param_key("ln_final.w", adapter)
998 # Skip layer norm folding if ln_final doesn't exist
999 # (e.g., encoder-decoder models like T5 have encoder_ln_final/decoder_ln_final instead)
1000 if ln_final_w_key not in state_dict:
1001 return state_dict
1003 has_unembed_bias = unembed_b_U_key in state_dict
1004 unembed_weight = ProcessWeights.convert_tensor_to_tl_format(
1005 unembed_W_U_key, state_dict, state_dict.get(unembed_W_U_key), cfg, adapter, None
1006 )
1007 ln_weight = state_dict[ln_final_w_key]
1008 assert unembed_weight is not None, f"Unembed weight not found at key {unembed_W_U_key}"
1009 # rmsnorm_uses_offset: effective scale is (1+w), identity is 0.0
1010 rmsnorm_uses_offset = getattr(cfg, "rmsnorm_uses_offset", False)
1011 effective_ln_weight = (1.0 + ln_weight) if rmsnorm_uses_offset else ln_weight
1012 if len(unembed_weight.shape) == 2 and len(ln_weight.shape) == 1: 1012 ↛ 1022line 1012 didn't jump to line 1022 because the condition on line 1012 was always true
1013 if unembed_weight.shape[1] == ln_weight.shape[0]:
1014 new_unembed_weight = unembed_weight * effective_ln_weight[None, :]
1015 elif unembed_weight.shape[0] == ln_weight.shape[0]: 1015 ↛ 1018line 1015 didn't jump to line 1018 because the condition on line 1015 was always true
1016 new_unembed_weight = unembed_weight * effective_ln_weight[:, None]
1017 else:
1018 raise ValueError(
1019 f"Cannot broadcast unembedding weight {unembed_weight.shape} with layer norm weight {ln_weight.shape}"
1020 )
1021 else:
1022 raise ValueError(
1023 f"Unexpected tensor shapes: unembedding {unembed_weight.shape}, layer norm {ln_weight.shape}"
1024 )
1025 state_dict[unembed_W_U_key] = ProcessWeights.convert_tensor_to_hf_format(
1026 unembed_W_U_key, new_unembed_weight, cfg, adapter, None
1027 )
1028 # Set ln_final.w to identity: zeros (rmsnorm_uses_offset) or ones (standard)
1029 identity_val = (
1030 torch.zeros_like(ln_weight) if rmsnorm_uses_offset else torch.ones_like(ln_weight)
1031 )
1032 if ln_final_w_key in state_dict: 1032 ↛ 1034line 1032 didn't jump to line 1034 because the condition on line 1032 was always true
1033 state_dict[ln_final_w_key] = identity_val
1034 alternate_final_w_key = (
1035 ln_final_w_key.replace("ln_f", "ln_final")
1036 if "ln_f" in ln_final_w_key
1037 else ln_final_w_key.replace("ln_final", "ln_f")
1038 )
1039 if alternate_final_w_key != ln_final_w_key and alternate_final_w_key in state_dict: 1039 ↛ 1040line 1039 didn't jump to line 1040 because the condition on line 1039 was never true
1040 state_dict[alternate_final_w_key] = identity_val
1041 if center_weights:
1042 unembed_weight_centered = ProcessWeights.convert_tensor_to_tl_format(
1043 unembed_W_U_key, state_dict, state_dict.get(unembed_W_U_key), cfg, adapter, None
1044 )
1045 assert (
1046 unembed_weight_centered is not None
1047 ), f"Unembed weight not found at key {unembed_W_U_key}"
1048 if len(unembed_weight_centered.shape) == 2: 1048 ↛ 1069line 1048 didn't jump to line 1069 because the condition on line 1048 was always true
1049 # Center along d_model: detect TL vs HF format
1050 d_vocab = getattr(cfg, "d_vocab", None) if cfg is not None else None
1051 if (
1052 d_vocab is not None
1053 and unembed_weight_centered.shape[0] == d_vocab
1054 and unembed_weight_centered.shape[-1] != d_vocab
1055 ):
1056 # HF format [d_vocab, d_model] — center along dim=-1
1057 unembed_weight_centered = (
1058 unembed_weight_centered - unembed_weight_centered.mean(-1, keepdim=True)
1059 )
1060 else:
1061 # TL format [d_model, d_vocab] — center along dim=0
1062 unembed_weight_centered = (
1063 unembed_weight_centered - unembed_weight_centered.mean(0, keepdim=True)
1064 )
1065 state_dict[unembed_W_U_key] = ProcessWeights.convert_tensor_to_hf_format(
1066 unembed_W_U_key, unembed_weight_centered, cfg, adapter, None
1067 )
1068 else:
1069 raise ValueError(
1070 f"Unexpected unembedding weight shape: {unembed_weight_centered.shape}"
1071 )
1073 return state_dict
1075 @staticmethod
1076 def _fold_final_rms_bias(
1077 state_dict: Dict[str, torch.Tensor], cfg, fold_biases: bool, adapter
1078 ) -> Dict[str, torch.Tensor]:
1079 """Fold final RMS bias into unembedding (separate from regular unembed folding).
1081 Args:
1082 state_dict: The state dictionary to process (modified in place)
1083 cfg: Model configuration object
1084 fold_biases: Whether to fold LayerNorm biases
1085 adapter: Optional architecture adapter for parameter key translation
1086 """
1087 unembed_b_U_key = ProcessWeights._get_param_key("unembed.b_U", adapter)
1088 unembed_W_U_key = ProcessWeights._get_param_key("unembed.W_U", adapter)
1089 ln_final_b_key = ProcessWeights._get_param_key("ln_final.b", adapter)
1090 has_unembed_bias = unembed_b_U_key in state_dict
1091 has_ln_final_bias = ln_final_b_key in state_dict
1092 if (
1093 not getattr(cfg, "final_rms", False)
1094 and fold_biases
1095 and has_unembed_bias
1096 and has_ln_final_bias
1097 ):
1098 unembed_weight = ProcessWeights.convert_tensor_to_tl_format(
1099 unembed_W_U_key, state_dict, state_dict.get(unembed_W_U_key), cfg, adapter, None
1100 )
1101 ln_bias = state_dict[ln_final_b_key]
1102 assert unembed_weight is not None, f"Unembed weight not found at key {unembed_W_U_key}"
1103 if len(unembed_weight.shape) == 2 and len(ln_bias.shape) == 1: 1103 ↛ 1113line 1103 didn't jump to line 1113 because the condition on line 1103 was always true
1104 if unembed_weight.shape[1] == ln_bias.shape[0]:
1105 bias_contribution = (unembed_weight * ln_bias[None, :]).sum(dim=-1)
1106 elif unembed_weight.shape[0] == ln_bias.shape[0]: 1106 ↛ 1109line 1106 didn't jump to line 1109 because the condition on line 1106 was always true
1107 bias_contribution = (unembed_weight * ln_bias[:, None]).sum(dim=-2)
1108 else:
1109 raise ValueError(
1110 f"Cannot broadcast unembedding weight {unembed_weight.shape} with layer norm bias {ln_bias.shape}"
1111 )
1112 else:
1113 raise ValueError(
1114 f"Unexpected tensor shapes: unembedding {unembed_weight.shape}, layer norm bias {ln_bias.shape}"
1115 )
1116 unembed_b_U = ProcessWeights.convert_tensor_to_tl_format(
1117 unembed_b_U_key, state_dict, state_dict.get(unembed_b_U_key), cfg, adapter, None
1118 )
1119 assert unembed_b_U is not None, f"Unembed bias not found at key {unembed_b_U_key}"
1120 new_unembed_b_U = unembed_b_U + bias_contribution
1121 state_dict[unembed_b_U_key] = ProcessWeights.convert_tensor_to_hf_format(
1122 unembed_b_U_key, new_unembed_b_U, cfg, adapter, None
1123 )
1124 if ln_final_b_key in state_dict: 1124 ↛ 1126line 1124 didn't jump to line 1126 because the condition on line 1124 was always true
1125 state_dict[ln_final_b_key] = torch.zeros_like(ln_bias)
1126 alternate_final_b_key = (
1127 ln_final_b_key.replace("ln_f", "ln_final")
1128 if "ln_f" in ln_final_b_key
1129 else ln_final_b_key.replace("ln_final", "ln_f")
1130 )
1131 if alternate_final_b_key != ln_final_b_key and alternate_final_b_key in state_dict: 1131 ↛ 1132line 1131 didn't jump to line 1132 because the condition on line 1131 was never true
1132 state_dict[alternate_final_b_key] = torch.zeros_like(ln_bias)
1134 return state_dict
1136 @staticmethod
1137 def fold_layer_norm(
1138 state_dict: Dict[str, torch.Tensor],
1139 cfg,
1140 fold_biases: bool = True,
1141 center_weights: bool = True,
1142 adapter=None,
1143 ) -> Dict[str, torch.Tensor]:
1144 """Fold Layer Norm. Can also be used to fold RMS Norm, when fold_biases and center_weights are set to False.
1146 Takes in a state dict from a pretrained model, formatted to be consistent with
1147 HookedTransformer but with LayerNorm weights and biases. Folds these into the neighbouring
1148 weights. See further_comments.md for more details.
1150 Args:
1151 state_dict (Dict[str, torch.Tensor]): State dict of pretrained model.
1152 cfg: Model configuration object with n_layers, n_key_value_heads, etc.
1153 fold_biases (bool): Enables folding of LN biases. Should be disabled when RMS Norm is used.
1154 center_weights (bool): Enables the centering of weights after folding in LN. Should be disabled when RMS Norm is used.
1155 adapter: Optional architecture adapter for parameter key translation.
1157 Returns:
1158 Dict[str, torch.Tensor]: Modified state dict with LayerNorm folded into linear layers.
1159 """
1160 # Make a deep copy to avoid modifying the original
1161 state_dict = {
1162 k: v.clone() if isinstance(v, torch.Tensor) else v for k, v in state_dict.items()
1163 }
1164 gqa = "" if getattr(cfg, "n_key_value_heads", None) is None else "_"
1165 for l in range(cfg.n_layers):
1166 state_dict = ProcessWeights._fold_layer(
1167 state_dict, cfg, l, fold_biases, center_weights, adapter, gqa
1168 )
1169 state_dict = ProcessWeights._fold_final_rms_bias(state_dict, cfg, fold_biases, adapter)
1170 state_dict = ProcessWeights._fold_unembed_layer_norm(
1171 state_dict, cfg, fold_biases, center_weights, adapter
1172 )
1173 return state_dict
1175 @staticmethod
1176 def center_writing_weights(
1177 state_dict: Dict[str, torch.Tensor], cfg, adapter=None
1178 ) -> Dict[str, torch.Tensor]:
1179 """Center Writing Weights.
1181 Centers the weights of the model that write to the residual stream - W_out, W_E, W_pos and
1182 W_out. This is done by subtracting the mean of the weights from the weights themselves. This
1183 is done in-place. See fold_layer_norm for more details.
1185 Args:
1186 state_dict (Dict[str, torch.Tensor]): State dict of the model.
1187 cfg: Model configuration object.
1188 adapter: Optional architecture adapter for parameter key translation.
1190 Returns:
1191 Dict[str, torch.Tensor]: Modified state dict with centered writing weights.
1192 """
1193 # Post-norm models leave the first attention's input un-normed, so centering
1194 # the embedding would shift a residual stream nothing re-normalizes.
1195 architecture = getattr(cfg, "original_architecture", None)
1196 if architecture in POST_NORM_ARCHITECTURES:
1197 print(f"Not centering embedding weights for {architecture}")
1198 else:
1199 # Make a deep copy to avoid modifying the original
1200 embed_W_E_key = ProcessWeights._get_param_key("embed.W_E", adapter)
1201 try:
1202 pos_embed_W_pos_key = (
1203 ProcessWeights._get_param_key("pos_embed.W_pos", adapter)
1204 if getattr(cfg, "positional_embedding_type", "standard")
1205 not in ("rotary", "alibi", "none")
1206 else None
1207 )
1208 except ValueError:
1209 pos_embed_W_pos_key = None
1210 if embed_W_E_key not in state_dict: 1210 ↛ 1211line 1210 didn't jump to line 1211 because the condition on line 1210 was never true
1211 raise KeyError(
1212 f"Expected embedding key '{embed_W_E_key}' not found in state_dict. Available keys: {list(state_dict.keys())[:10]}..."
1213 )
1214 embed_W_E = ProcessWeights.convert_tensor_to_tl_format(
1215 embed_W_E_key, state_dict, state_dict.get(embed_W_E_key), cfg, adapter, None
1216 )
1217 assert embed_W_E is not None, f"Embedding not found at key {embed_W_E_key}"
1218 embed_W_E = embed_W_E - embed_W_E.mean(-1, keepdim=True)
1219 state_dict[embed_W_E_key] = ProcessWeights.convert_tensor_to_hf_format(
1220 embed_W_E_key, embed_W_E, cfg, adapter, None
1221 )
1223 if (
1224 getattr(cfg, "positional_embedding_type", "standard")
1225 not in ("rotary", "alibi", "none")
1226 and pos_embed_W_pos_key is not None
1227 ):
1228 if pos_embed_W_pos_key not in state_dict: 1228 ↛ 1229line 1228 didn't jump to line 1229 because the condition on line 1228 was never true
1229 raise KeyError(
1230 f"Expected positional embedding key '{pos_embed_W_pos_key}' not found in state_dict. Available keys: {list(state_dict.keys())[:10]}..."
1231 )
1232 pos_embed_W_pos = ProcessWeights.convert_tensor_to_tl_format(
1233 pos_embed_W_pos_key,
1234 state_dict,
1235 state_dict.get(pos_embed_W_pos_key),
1236 cfg,
1237 adapter,
1238 None,
1239 )
1240 assert (
1241 pos_embed_W_pos is not None
1242 ), f"Positional embedding not found at key {pos_embed_W_pos_key}"
1243 pos_embed_W_pos = pos_embed_W_pos - pos_embed_W_pos.mean(-1, keepdim=True)
1244 state_dict[pos_embed_W_pos_key] = ProcessWeights.convert_tensor_to_hf_format(
1245 pos_embed_W_pos_key, pos_embed_W_pos, cfg, adapter, None
1246 )
1247 for l in range(cfg.n_layers):
1248 attn_W_O_key = ProcessWeights._get_param_key(f"blocks.{l}.attn.W_O", adapter)
1249 attn_b_O_key = ProcessWeights._get_param_key(f"blocks.{l}.attn.b_O", adapter)
1250 try:
1251 mlp_W_out_key = ProcessWeights._resolve_state_dict_key(
1252 state_dict, ProcessWeights._get_param_key(f"blocks.{l}.mlp.W_out", adapter), l
1253 )
1254 mlp_b_out_key = ProcessWeights._resolve_state_dict_key(
1255 state_dict, ProcessWeights._get_param_key(f"blocks.{l}.mlp.b_out", adapter), l
1256 )
1257 except ValueError:
1258 mlp_W_out_key = None
1259 mlp_b_out_key = None
1260 if attn_W_O_key in state_dict:
1261 attn_W_O = ProcessWeights.convert_tensor_to_tl_format(
1262 attn_W_O_key, state_dict, state_dict.get(attn_W_O_key), cfg, adapter, l
1263 )
1264 assert attn_W_O is not None, f"Attention W_O not found at key {attn_W_O_key}"
1265 attn_W_O = attn_W_O - attn_W_O.mean(-1, keepdim=True)
1266 state_dict[attn_W_O_key] = ProcessWeights.convert_tensor_to_hf_format(
1267 attn_W_O_key, attn_W_O, cfg, adapter, l
1268 )
1269 if attn_b_O_key in state_dict: 1269 ↛ 1278line 1269 didn't jump to line 1278 because the condition on line 1269 was always true
1270 attn_b_O = ProcessWeights.convert_tensor_to_tl_format(
1271 attn_b_O_key, state_dict, state_dict.get(attn_b_O_key), cfg, adapter, l
1272 )
1273 assert attn_b_O is not None, f"Attention b_O not found at key {attn_b_O_key}"
1274 attn_b_O = attn_b_O - attn_b_O.mean()
1275 state_dict[attn_b_O_key] = ProcessWeights.convert_tensor_to_hf_format(
1276 attn_b_O_key, attn_b_O, cfg, adapter, l
1277 )
1278 if not getattr(cfg, "attn_only", False):
1279 is_moe = getattr(cfg, "num_experts", None) is not None and cfg.num_experts > 0
1280 if is_moe:
1281 num_experts = cfg.num_experts
1282 for e in range(num_experts):
1283 expert_W_out_key = None
1284 expert_b_out_key = None
1285 expert_W_out_patterns = [
1286 f"blocks.{l}.mlp.experts.{e}.W_out",
1287 f"blocks.{l}.mlp.experts.{e}.W_out.weight",
1288 ]
1289 for pattern in expert_W_out_patterns:
1290 if pattern in state_dict: 1290 ↛ 1291line 1290 didn't jump to line 1291 because the condition on line 1290 was never true
1291 expert_W_out_key = pattern
1292 break
1293 if expert_W_out_key is None and adapter: 1293 ↛ 1294line 1293 didn't jump to line 1294 because the condition on line 1293 was never true
1294 try:
1295 candidate = ProcessWeights._get_param_key(
1296 f"blocks.{l}.mlp.experts.{e}.W_out", adapter
1297 )
1298 expert_W_out_key = ProcessWeights._resolve_state_dict_key(
1299 state_dict, candidate, l
1300 )
1301 except ValueError:
1302 pass
1303 if expert_W_out_key and expert_W_out_key in state_dict: 1303 ↛ 1304line 1303 didn't jump to line 1304 because the condition on line 1303 was never true
1304 expert_W_out = ProcessWeights.convert_tensor_to_tl_format(
1305 expert_W_out_key,
1306 state_dict,
1307 state_dict.get(expert_W_out_key),
1308 cfg,
1309 adapter,
1310 l,
1311 )
1312 assert (
1313 expert_W_out is not None
1314 ), f"Expert W_out not found at key {expert_W_out_key}"
1315 expert_W_out = expert_W_out - expert_W_out.mean(-1, keepdim=True)
1316 state_dict[
1317 expert_W_out_key
1318 ] = ProcessWeights.convert_tensor_to_hf_format(
1319 expert_W_out_key, expert_W_out, cfg, adapter, l
1320 )
1321 expert_b_out_patterns = [
1322 f"blocks.{l}.mlp.experts.{e}.b_out",
1323 f"blocks.{l}.mlp.experts.{e}.b_out.bias",
1324 ]
1325 for pattern in expert_b_out_patterns:
1326 if pattern in state_dict: 1326 ↛ 1327line 1326 didn't jump to line 1327 because the condition on line 1326 was never true
1327 expert_b_out_key = pattern
1328 break
1329 if expert_b_out_key is None and adapter: 1329 ↛ 1330line 1329 didn't jump to line 1330 because the condition on line 1329 was never true
1330 try:
1331 candidate = ProcessWeights._get_param_key(
1332 f"blocks.{l}.mlp.experts.{e}.b_out", adapter
1333 )
1334 expert_b_out_key = ProcessWeights._resolve_state_dict_key(
1335 state_dict, candidate, l
1336 )
1337 except ValueError:
1338 pass
1339 if expert_b_out_key and expert_b_out_key in state_dict: 1339 ↛ 1340line 1339 didn't jump to line 1340 because the condition on line 1339 was never true
1340 expert_b_out = ProcessWeights.convert_tensor_to_tl_format(
1341 expert_b_out_key,
1342 state_dict,
1343 state_dict.get(expert_b_out_key),
1344 cfg,
1345 adapter,
1346 l,
1347 )
1348 assert (
1349 expert_b_out is not None
1350 ), f"Expert b_out not found at key {expert_b_out_key}"
1351 expert_b_out = expert_b_out - expert_b_out.mean()
1352 state_dict[
1353 expert_b_out_key
1354 ] = ProcessWeights.convert_tensor_to_hf_format(
1355 expert_b_out_key, expert_b_out, cfg, adapter, l
1356 )
1357 elif mlp_W_out_key is not None and mlp_W_out_key in state_dict:
1358 mlp_W_out = ProcessWeights.convert_tensor_to_tl_format(
1359 mlp_W_out_key, state_dict, state_dict.get(mlp_W_out_key), cfg, adapter, l
1360 )
1361 assert mlp_W_out is not None, f"MLP W_out not found at key {mlp_W_out_key}"
1362 # Center along d_model dimension. In TL format W_out is [d_mlp, d_model]
1363 # so d_model is dim=-1. But bridge adapters may keep HF format
1364 # [d_model, d_mlp] where d_model is dim=0. Detect via cfg.d_model.
1365 if mlp_W_out.shape[-1] == cfg.d_model:
1366 mlp_W_out = mlp_W_out - mlp_W_out.mean(-1, keepdim=True)
1367 elif mlp_W_out.shape[0] == cfg.d_model: 1367 ↛ 1370line 1367 didn't jump to line 1370 because the condition on line 1367 was always true
1368 mlp_W_out = mlp_W_out - mlp_W_out.mean(0, keepdim=True)
1369 else:
1370 mlp_W_out = mlp_W_out - mlp_W_out.mean(-1, keepdim=True)
1371 state_dict[mlp_W_out_key] = ProcessWeights.convert_tensor_to_hf_format(
1372 mlp_W_out_key, mlp_W_out, cfg, adapter, l
1373 )
1374 if mlp_b_out_key is not None and mlp_b_out_key in state_dict: 1374 ↛ 1247line 1374 didn't jump to line 1247 because the condition on line 1374 was always true
1375 mlp_b_out = ProcessWeights.convert_tensor_to_tl_format(
1376 mlp_b_out_key,
1377 state_dict,
1378 state_dict.get(mlp_b_out_key),
1379 cfg,
1380 adapter,
1381 l,
1382 )
1383 assert mlp_b_out is not None, f"MLP b_out not found at key {mlp_b_out_key}"
1384 mlp_b_out = mlp_b_out - mlp_b_out.mean()
1385 state_dict[mlp_b_out_key] = ProcessWeights.convert_tensor_to_hf_format(
1386 mlp_b_out_key, mlp_b_out, cfg, adapter, l
1387 )
1388 return state_dict
1390 @staticmethod
1391 def center_unembed(
1392 state_dict: Dict[str, torch.Tensor], cfg=None, adapter=None
1393 ) -> Dict[str, torch.Tensor]:
1394 """Center the unembedding weights W_U.
1396 This is done by subtracting the mean of the weights from the weights themselves. This is
1397 done in-place. As softmax is translation invariant, this changes the logits but not the log
1398 probs, and makes the model logits (slightly) more interpretable - when trying to understand
1399 how components contribute to the logits, we'll be less misled by components that just add
1400 something to every logit.
1402 Args:
1403 state_dict (Dict[str, torch.Tensor]): State dict of the model.
1404 cfg: Model configuration (used to determine d_vocab for correct centering dimension).
1405 adapter: Optional architecture adapter for parameter key translation.
1407 Returns:
1408 Dict[str, torch.Tensor]: Modified state dict with centered unembedding weights.
1409 """
1410 # Make a deep copy to avoid modifying the original
1411 state_dict = {
1412 k: v.clone() if isinstance(v, torch.Tensor) else v for k, v in state_dict.items()
1413 }
1414 unembed_W_U_key = ProcessWeights._get_param_key("unembed.W_U", adapter)
1415 unembed_b_U_key = ProcessWeights._get_param_key("unembed.b_U", adapter)
1416 if unembed_W_U_key not in state_dict:
1417 raise KeyError(
1418 f"Expected unembedding weight key '{unembed_W_U_key}' not found in state_dict. Available keys: {list(state_dict.keys())[:10]}..."
1419 )
1420 W_U = ProcessWeights.convert_tensor_to_tl_format(
1421 unembed_W_U_key, state_dict, state_dict.get(unembed_W_U_key), None, adapter, None
1422 )
1423 assert W_U is not None, f"Unembed weight not found at key {unembed_W_U_key}"
1425 # Detect W_U format to center along correct dim (wrong dim corrupts output)
1426 vocab_dim = -1 # Default: TL format [d_model, d_vocab]
1427 if cfg is not None:
1428 d_vocab = getattr(cfg, "d_vocab", None)
1429 d_model = getattr(cfg, "d_model", None)
1430 if d_vocab is not None and W_U.shape[0] == d_vocab and W_U.shape[-1] != d_vocab:
1431 # HF format [d_vocab, d_model] — center along dim=0
1432 vocab_dim = 0
1433 elif d_model is not None and W_U.shape[-1] == d_model and W_U.shape[0] != d_model:
1434 # Padded-vocab checkpoints (e.g. HyenaDNA pads 12→16 rows) defeat
1435 # the d_vocab match; the d_model axis is unambiguous.
1436 vocab_dim = 0
1437 W_U = W_U - W_U.mean(vocab_dim, keepdim=True)
1438 state_dict[unembed_W_U_key] = ProcessWeights.convert_tensor_to_hf_format(
1439 unembed_W_U_key, W_U, None, adapter, None
1440 )
1441 if unembed_b_U_key in state_dict:
1442 unembed_b_U = ProcessWeights.convert_tensor_to_tl_format(
1443 unembed_b_U_key, state_dict, state_dict.get(unembed_b_U_key), None, adapter, None
1444 )
1445 assert unembed_b_U is not None, f"Unembed bias not found at key {unembed_b_U_key}"
1446 unembed_b_U = unembed_b_U - unembed_b_U.mean()
1447 state_dict[unembed_b_U_key] = ProcessWeights.convert_tensor_to_hf_format(
1448 unembed_b_U_key, unembed_b_U, None, adapter, None
1449 )
1450 return state_dict
1452 @staticmethod
1453 def fold_value_biases(
1454 state_dict: Dict[str, torch.Tensor], cfg, adapter=None
1455 ) -> Dict[str, torch.Tensor]:
1456 """Fold the value biases into the output bias.
1458 Because attention patterns add up to 1, the value biases always have a constant effect on a
1459 head's output. Further, as the outputs of each head in a layer add together, each head's
1460 value bias has a constant effect on the *layer's* output, which can make it harder to
1461 interpret the effect of any given head, and it doesn't matter which head a bias is
1462 associated with. We can factor this all into a single output bias to the layer, and make it
1463 easier to interpret the head's output. Formally, we take b_O_new = b_O_original +
1464 sum_head(b_V_head @ W_O_head).
1466 Args:
1467 state_dict (Dict[str, torch.Tensor]): State dict of the model.
1468 cfg: Model configuration object.
1469 adapter: Optional architecture adapter for parameter key translation.
1471 Returns:
1472 Dict[str, torch.Tensor]: Modified state dict with value biases folded into output bias.
1473 """
1474 # Make a deep copy to avoid modifying the original
1475 state_dict = {
1476 k: v.clone() if isinstance(v, torch.Tensor) else v for k, v in state_dict.items()
1477 }
1478 layer = 0
1479 for layer in range(cfg.n_layers):
1480 split_v_bias_key = f"blocks.{layer}.attn.v.bias"
1481 if split_v_bias_key in state_dict:
1482 b_V_key = split_v_bias_key
1483 W_O_key = ProcessWeights._get_param_key(f"blocks.{layer}.attn.W_O", adapter)
1484 b_O_key = ProcessWeights._get_param_key(f"blocks.{layer}.attn.b_O", adapter)
1485 else:
1486 if getattr(cfg, "n_key_value_heads", None) is None:
1487 b_V_key = ProcessWeights._get_param_key(f"blocks.{layer}.attn.b_V", adapter)
1488 else:
1489 b_V_key = ProcessWeights._get_param_key(f"blocks.{layer}.attn._b_V", adapter)
1490 W_O_key = ProcessWeights._get_param_key(f"blocks.{layer}.attn.W_O", adapter)
1491 b_O_key = ProcessWeights._get_param_key(f"blocks.{layer}.attn.b_O", adapter)
1492 if b_V_key in state_dict:
1493 b_V = ProcessWeights.convert_tensor_to_tl_format(
1494 b_V_key, state_dict, state_dict.get(b_V_key), cfg, adapter, layer
1495 )
1496 assert b_V is not None, f"Value bias not found at key {b_V_key}"
1497 if b_V.numel() == 0: 1497 ↛ 1498line 1497 didn't jump to line 1498 because the condition on line 1497 was never true
1498 continue
1499 W_O = ProcessWeights.convert_tensor_to_tl_format(
1500 W_O_key, state_dict, state_dict.get(W_O_key), cfg, adapter, layer
1501 )
1502 assert W_O is not None, f"Attention W_O not found at key {W_O_key}"
1503 if b_O_key not in state_dict:
1504 # Create zero b_O to absorb the folded value bias
1505 b_O_original = torch.zeros(cfg.d_model, dtype=b_V.dtype, device=b_V.device)
1506 state_dict[b_O_key] = b_O_original
1507 else:
1508 b_O_original_maybe = ProcessWeights.convert_tensor_to_tl_format(
1509 b_O_key, state_dict, state_dict.get(b_O_key), cfg, adapter, layer
1510 )
1511 assert (
1512 b_O_original_maybe is not None
1513 ), f"Attention b_O not found at key {b_O_key}"
1514 b_O_original = b_O_original_maybe
1515 # Align W_O / b_O to b_V's device.
1516 if W_O.device != b_V.device: 1516 ↛ 1517line 1516 didn't jump to line 1517 because the condition on line 1516 was never true
1517 W_O = W_O.to(b_V.device)
1518 if b_O_original.device != b_V.device: 1518 ↛ 1519line 1518 didn't jump to line 1519 because the condition on line 1518 was never true
1519 b_O_original = b_O_original.to(b_V.device)
1520 is_split_format = ".attn.v.bias" in b_V_key or ".attn.k.bias" in b_V_key
1521 if is_split_format and len(b_V.shape) == 1 and (len(W_O.shape) == 2): 1521 ↛ 1522line 1521 didn't jump to line 1522 because the condition on line 1521 was never true
1522 n_heads = cfg.n_heads
1523 d_head = cfg.d_head
1524 d_model = cfg.d_model
1525 b_V_only = b_V
1526 b_V_reshaped = b_V_only.reshape(n_heads, d_head)
1527 W_O_reshaped = einops.rearrange(W_O, "(i h) m -> i h m", i=n_heads)
1528 folded_b_O = b_O_original + (b_V_reshaped[:, :, None] * W_O_reshaped).sum(
1529 [0, 1]
1530 )
1531 state_dict[b_O_key] = ProcessWeights.convert_tensor_to_hf_format(
1532 b_O_key, folded_b_O, cfg, adapter, layer
1533 )
1534 tl_b_O_key = f"blocks.{layer}.attn.b_O"
1535 if tl_b_O_key in state_dict:
1536 state_dict[tl_b_O_key] = ProcessWeights.convert_tensor_to_hf_format(
1537 tl_b_O_key, folded_b_O, cfg, adapter, layer
1538 )
1539 state_dict[b_V_key] = ProcessWeights.convert_tensor_to_hf_format(
1540 b_V_key, torch.zeros_like(b_V), cfg, adapter, layer
1541 )
1542 elif len(b_V.shape) == 1 and len(W_O.shape) == 2: 1542 ↛ 1543line 1542 didn't jump to line 1543 because the condition on line 1542 was never true
1543 n_heads = cfg.n_heads
1544 d_head = cfg.d_head
1545 d_model = cfg.d_model
1546 v_bias_start = 2 * n_heads * d_head
1547 v_bias_end = 3 * n_heads * d_head
1548 b_V_only = b_V[v_bias_start:v_bias_end]
1549 if b_V_only.numel() == 0:
1550 continue
1551 b_V_reshaped = b_V_only.reshape(n_heads, d_head)
1552 W_O_reshaped = einops.rearrange(W_O, "(i h) m -> i h m", i=n_heads)
1553 folded_b_O = b_O_original + (b_V_reshaped[:, :, None] * W_O_reshaped).sum(
1554 [0, 1]
1555 )
1556 new_b_V = b_V.clone()
1557 new_b_V[v_bias_start:v_bias_end] = 0
1558 state_dict[b_V_key] = ProcessWeights.convert_tensor_to_hf_format(
1559 b_V_key, new_b_V, cfg, adapter, layer
1560 )
1561 elif is_split_format and len(b_V.shape) == 1 and len(W_O.shape) == 3: 1561 ↛ 1563line 1561 didn't jump to line 1563 because the condition on line 1561 was never true
1562 # Split bias [n_heads * d_head] with W_O already in TL format [n_heads, d_head, d_model]
1563 n_heads = cfg.n_heads
1564 d_head = cfg.d_head
1565 b_V_reshaped = b_V.reshape(n_heads, d_head)
1566 if getattr(cfg, "n_key_value_heads", None) is not None:
1567 b_V_reshaped = torch.repeat_interleave(
1568 b_V_reshaped, dim=0, repeats=cfg.n_heads // cfg.n_key_value_heads
1569 )
1570 folded_b_O = b_O_original + (b_V_reshaped[:, :, None] * W_O).sum([0, 1])
1571 state_dict[b_V_key] = ProcessWeights.convert_tensor_to_hf_format(
1572 b_V_key, torch.zeros_like(b_V), cfg, adapter, layer
1573 )
1574 elif len(b_V.shape) == 2 and len(W_O.shape) == 3: 1574 ↛ 1588line 1574 didn't jump to line 1588 because the condition on line 1574 was always true
1575 b_V_original_shape = b_V.shape
1576 if getattr(cfg, "n_key_value_heads", None) is not None:
1577 b_V = torch.repeat_interleave(
1578 b_V, dim=0, repeats=cfg.n_heads // cfg.n_key_value_heads
1579 )
1580 folded_b_O = b_O_original + (b_V[:, :, None] * W_O).sum([0, 1])
1581 state_dict[b_V_key] = ProcessWeights.convert_tensor_to_hf_format(
1582 b_V_key,
1583 torch.zeros(b_V_original_shape, dtype=b_V.dtype, device=b_V.device),
1584 cfg,
1585 adapter,
1586 layer,
1587 )
1588 elif len(b_V.shape) == 2 and len(W_O.shape) == 2:
1589 n_heads = cfg.n_heads
1590 d_head = cfg.d_head
1591 d_model = cfg.d_model
1592 b_V_original_shape = b_V.shape
1594 # Handle split QKV format where bias might be [1, d_model] or [n_heads, d_head]
1595 is_split_format = ".attn.v.bias" in b_V_key or ".attn.k.bias" in b_V_key
1596 if is_split_format and b_V.shape[0] == 1 and b_V.shape[1] == n_heads * d_head:
1597 # Reshape [1, n_heads * d_head] to [n_heads, d_head]
1598 b_V = b_V.reshape(n_heads, d_head)
1599 elif b_V.shape != (n_heads, d_head):
1600 # If not already [n_heads, d_head], try to reshape
1601 if b_V.numel() == n_heads * d_head:
1602 b_V = b_V.reshape(n_heads, d_head)
1604 if getattr(cfg, "n_key_value_heads", None) is not None:
1605 b_V = torch.repeat_interleave(
1606 b_V, dim=0, repeats=cfg.n_heads // cfg.n_key_value_heads
1607 )
1609 W_O_reshaped = einops.rearrange(W_O, "(i h) m -> i h m", i=n_heads)
1610 folded_b_O = b_O_original + (b_V[:, :, None] * W_O_reshaped).sum([0, 1])
1611 state_dict[b_V_key] = ProcessWeights.convert_tensor_to_hf_format(
1612 b_V_key,
1613 torch.zeros(b_V_original_shape, dtype=b_V.dtype, device=b_V.device),
1614 cfg,
1615 adapter,
1616 layer,
1617 )
1618 else:
1619 raise ValueError(f"Unexpected tensor shapes: b_V {b_V.shape}, W_O {W_O.shape}")
1620 state_dict[b_O_key] = ProcessWeights.convert_tensor_to_hf_format(
1621 b_O_key, folded_b_O, cfg, adapter, layer
1622 )
1623 return state_dict
1625 @staticmethod
1626 def process_weights(
1627 state_dict: Dict[str, torch.Tensor],
1628 cfg,
1629 fold_ln: bool = True,
1630 center_writing_weights: bool = True,
1631 center_unembed: bool = True,
1632 fold_value_biases: bool = True,
1633 refactor_factored_attn_matrices: bool = False,
1634 adapter=None,
1635 ) -> Dict[str, torch.Tensor]:
1636 """Apply all weight processing transformations in the correct order.
1638 This is a convenience function that applies all the weight processing steps
1639 in the same order as the legacy HookedTransformer load path.
1641 Args:
1642 state_dict (Dict[str, torch.Tensor]): State dict of the model.
1643 cfg: Model configuration object.
1644 fold_ln (bool): Whether to fold LayerNorm weights into subsequent layers.
1645 center_writing_weights (bool): Whether to center weights writing to residual stream.
1646 center_unembed (bool): Whether to center unembedding weights.
1647 fold_value_biases (bool): Whether to fold value biases into output bias.
1648 refactor_factored_attn_matrices (bool): Whether to refactor attention matrices.
1649 adapter: Optional architecture adapter for parameter key translation.
1651 Returns:
1652 Dict[str, torch.Tensor]: Fully processed state dict.
1653 """
1654 # Upcast to float32 for weight processing to avoid precision loss in
1655 # reduced-precision dtypes (bfloat16, float16). Operations like LayerNorm
1656 # folding involve multiplications that accumulate rounding errors when
1657 # performed in low precision.
1658 original_dtypes: Dict[str, torch.dtype] = {}
1659 for k, v in state_dict.items():
1660 if isinstance(v, torch.Tensor) and v.is_floating_point() and v.dtype != torch.float32: 1660 ↛ 1661line 1660 didn't jump to line 1661 because the condition on line 1660 was never true
1661 original_dtypes[k] = v.dtype
1662 state_dict[k] = v.float()
1664 # Skip fold_ln for adapters that don't support it (e.g., post-LN architectures
1665 # like BERT where LN placement means folding goes into the wrong sublayer).
1666 if fold_ln and adapter and not getattr(adapter, "supports_fold_ln", True):
1667 import warnings
1669 warnings.warn(
1670 f"{type(adapter).__name__} does not support fold_ln; norm weights "
1671 "stay unfolded and analyses assuming folded LN will be off.",
1672 UserWarning,
1673 )
1674 fold_ln = False
1675 if fold_ln:
1676 if getattr(cfg, "normalization_type", "LN") in ["LN", "LNPre"]:
1677 state_dict = ProcessWeights.fold_layer_norm(
1678 state_dict, cfg, fold_biases=True, center_weights=True, adapter=adapter
1679 )
1680 elif getattr(cfg, "normalization_type", "LN") in ["RMS", "RMSPre"]: 1680 ↛ 1691line 1680 didn't jump to line 1691 because the condition on line 1680 was always true
1681 state_dict = ProcessWeights.fold_layer_norm(
1682 state_dict, cfg, fold_biases=False, center_weights=False, adapter=adapter
1683 )
1684 # Note: Each folding function (_fold_layer for attention, _fold_mlp_layer_norm
1685 # for MLP) sets its own LN weights to 1.0 after successful folding.
1686 # We must NOT unconditionally set all LN weights to 1.0 here, because
1687 # models with combined QKV projections (e.g., OpenELM's qkv_proj) may
1688 # not be able to fold attention LN — setting ln1.w=1.0 without folding
1689 # destroys the RMS scaling.
1690 # Some adapters (e.g., post-LN) don't support center_writing_weights.
1691 if ( 1691 ↛ 1696line 1691 didn't jump to line 1696 because the condition on line 1691 was never true
1692 center_writing_weights
1693 and adapter
1694 and not getattr(adapter, "supports_center_writing_weights", True)
1695 ):
1696 import warnings
1698 warnings.warn(
1699 f"{type(adapter).__name__} does not support center_writing_weights; "
1700 "writing weights stay uncentered.",
1701 UserWarning,
1702 )
1703 center_writing_weights = False
1704 if center_writing_weights:
1705 if getattr(cfg, "normalization_type", "LN") in ["LN", "LNPre"] and (
1706 not getattr(cfg, "final_rms", False)
1707 ):
1708 state_dict = ProcessWeights.center_writing_weights(state_dict, cfg, adapter=adapter)
1709 if center_unembed:
1710 state_dict = ProcessWeights.center_unembed(state_dict, cfg=cfg, adapter=adapter)
1711 if fold_value_biases:
1712 state_dict = ProcessWeights.fold_value_biases(state_dict, cfg, adapter=adapter)
1713 if center_writing_weights and getattr(cfg, "normalization_type", "LN") in [
1714 "LN",
1715 "LNPre",
1716 ]:
1717 for layer_idx in range(cfg.n_layers):
1718 b_O_key = ProcessWeights._get_param_key(f"blocks.{layer_idx}.attn.b_O", adapter)
1719 if b_O_key in state_dict: 1719 ↛ 1717line 1719 didn't jump to line 1717 because the condition on line 1719 was always true
1720 b_O = ProcessWeights.convert_tensor_to_tl_format(
1721 b_O_key, state_dict, state_dict.get(b_O_key), cfg, adapter, layer_idx
1722 )
1723 assert b_O is not None, f"Attention b_O not found at key {b_O_key}"
1724 b_O = b_O - b_O.mean()
1725 state_dict[b_O_key] = ProcessWeights.convert_tensor_to_hf_format(
1726 b_O_key, b_O, cfg, adapter, layer_idx
1727 )
1728 if refactor_factored_attn_matrices:
1729 state_dict = ProcessWeights.refactor_factored_attn_matrices(
1730 state_dict, cfg, adapter=adapter
1731 )
1733 # Downcast back to original dtypes
1734 for k, orig_dtype in original_dtypes.items(): 1734 ↛ 1735line 1734 didn't jump to line 1735 because the loop on line 1734 never started
1735 if k in state_dict and isinstance(state_dict[k], torch.Tensor):
1736 state_dict[k] = state_dict[k].to(orig_dtype)
1738 return state_dict
1740 @staticmethod
1741 def refactor_factored_attn_matrices(
1742 state_dict: Dict[str, torch.Tensor], cfg, adapter=None
1743 ) -> Dict[str, torch.Tensor]:
1744 """Experimental method for managing queries, keys and values.
1746 As argued in [A Mathematical Framework for Transformer
1747 Circuits](https://transformer-circuits.pub/2021/framework/index.html), queries, keys and
1748 values are somewhat arbitrary intermediate terms when computing with the low rank factored
1749 matrices W_QK = W_Q @ W_K.T and W_OV = W_V @ W_O, and these matrices are the only thing
1750 determining head behaviour. But there are many ways to find a low rank factorization to a
1751 given matrix, and hopefully some of these are more interpretable than others! This method is
1752 one attempt, which makes all of the matrices have orthogonal rows or columns, W_O into a
1753 rotation and W_Q and W_K having the nth column in each having the same norm. The formula is
1754 $W_V = U @ S,W_O=Vh.T,W_Q=U@S.sqrt(),W_K=Vh@S.sqrt()$.
1756 More details:
1758 If W_OV = U @ S @ Vh.T in its singular value decomposition, (where S is in R^d_head not
1759 R^d_model, as W_OV is low rank), W_OV = (U @ S) @ (Vh.T) is an equivalent low rank
1760 factorisation, where rows/columns of each matrix are orthogonal! So setting $W_V=US$ and
1761 $W_O=Vh.T$ works just as well. I *think* this is a more interpretable setup, because now
1762 $W_O$ is just a rotation, and doesn't change the norm, so $z$ has the same norm as the
1763 result of the head.
1765 For $W_QK = W_Q @ W_K.T$ we use the refactor $W_Q = U @ S.sqrt()$ and $W_K = Vh @ S.sqrt()$,
1766 which is also equivalent ($S==S.sqrt() @ S.sqrt()$ as $S$ is diagonal). Here we keep the
1767 matrices as having the same norm, since there's not an obvious asymmetry between the keys
1768 and queries.
1770 Biases are more fiddly to deal with. For OV it's pretty easy - we just need (x @ W_V + b_V)
1771 @ W_O + b_O to be preserved, so we can set b_V' = 0. and b_O' = b_V @ W_O + b_O (note that
1772 b_V in R^{head_index x d_head} while b_O in R^{d_model}, so we need to sum b_V @ W_O along
1773 the head_index dimension too).
1775 For QK it's messy - we need to preserve the bilinear form of (x @ W_Q + b_Q) * (y @ W_K +
1776 b_K), which is fairly messy. To deal with the biases, we concatenate them to W_Q and W_K to
1777 simulate a d_model+1 dimensional input (whose final coordinate is always 1), do the SVD
1778 factorization on this effective matrix, then separate out into final weights and biases.
1780 Args:
1781 state_dict (Dict[str, torch.Tensor]): State dict of the model.
1782 cfg: Model configuration object.
1783 adapter: Optional architecture adapter for parameter key translation.
1785 Returns:
1786 Dict[str, torch.Tensor]: Modified state dict with refactored attention matrices.
1787 """
1788 # Make a deep copy to avoid modifying the original
1789 state_dict = {
1790 k: v.clone() if isinstance(v, torch.Tensor) else v for k, v in state_dict.items()
1791 }
1792 assert (
1793 getattr(cfg, "positional_embedding_type", "standard") != "rotary"
1794 ), "You can't refactor the QK circuit when using rotary embeddings (as the QK matrix depends on the position of the query and key)"
1796 for l in range(cfg.n_layers):
1797 W_Q_key = ProcessWeights._get_param_key(f"blocks.{l}.attn.W_Q", adapter)
1798 b_Q_key = ProcessWeights._get_param_key(f"blocks.{l}.attn.b_Q", adapter)
1799 W_K_key = ProcessWeights._get_param_key(f"blocks.{l}.attn.W_K", adapter)
1800 b_K_key = ProcessWeights._get_param_key(f"blocks.{l}.attn.b_K", adapter)
1801 W_V_key = ProcessWeights._get_param_key(f"blocks.{l}.attn.W_V", adapter)
1802 W_O_key = ProcessWeights._get_param_key(f"blocks.{l}.attn.W_O", adapter)
1803 b_V_key = ProcessWeights._get_param_key(f"blocks.{l}.attn.b_V", adapter)
1804 b_O_key = ProcessWeights._get_param_key(f"blocks.{l}.attn.b_O", adapter)
1806 # Skip hybrid layers without attention (other loops already guard individually)
1807 if W_Q_key not in state_dict:
1808 continue
1809 # If Q is present, K/V/O must be too
1810 for _required_key in [W_K_key, W_V_key, W_O_key]:
1811 if _required_key not in state_dict:
1812 raise ValueError(
1813 f"Inconsistent attention weights at layer {l}: "
1814 f"'{W_Q_key}' found but '{_required_key}' missing. "
1815 f"All of W_Q, W_K, W_V, W_O must be present together."
1816 )
1818 # W_QK = W_Q @ W_K.T
1819 # Concatenate biases to make a d_model+1 input dimension
1820 W_Q = ProcessWeights.convert_tensor_to_tl_format(
1821 W_Q_key, state_dict, state_dict.get(W_Q_key), cfg, adapter, l
1822 )
1823 b_Q = ProcessWeights.convert_tensor_to_tl_format(
1824 b_Q_key, state_dict, state_dict.get(b_Q_key), cfg, adapter, l
1825 )
1826 W_K = ProcessWeights.convert_tensor_to_tl_format(
1827 W_K_key, state_dict, state_dict.get(W_K_key), cfg, adapter, l
1828 )
1829 b_K = ProcessWeights.convert_tensor_to_tl_format(
1830 b_K_key, state_dict, state_dict.get(b_K_key), cfg, adapter, l
1831 )
1832 assert W_Q is not None, f"W_Q not found at key {W_Q_key}"
1833 assert b_Q is not None, f"b_Q not found at key {b_Q_key}"
1834 assert W_K is not None, f"W_K not found at key {W_K_key}"
1835 assert b_K is not None, f"b_K not found at key {b_K_key}"
1837 W_Q_eff = torch.cat([W_Q, b_Q[:, None, :]], dim=1)
1838 W_K_eff = torch.cat([W_K, b_K[:, None, :]], dim=1)
1840 W_Q_eff_even, W_K_eff_even_T = (
1841 FactoredMatrix(W_Q_eff, W_K_eff.transpose(-1, -2)).make_even().pair
1842 )
1843 W_K_eff_even = W_K_eff_even_T.transpose(-1, -2)
1845 state_dict[W_Q_key] = ProcessWeights.convert_tensor_to_hf_format(
1846 W_Q_key, W_Q_eff_even[:, :-1, :], cfg, adapter, l
1847 )
1848 state_dict[b_Q_key] = ProcessWeights.convert_tensor_to_hf_format(
1849 b_Q_key, W_Q_eff_even[:, -1, :], cfg, adapter, l
1850 )
1851 state_dict[W_K_key] = ProcessWeights.convert_tensor_to_hf_format(
1852 W_K_key, W_K_eff_even[:, :-1, :], cfg, adapter, l
1853 )
1854 state_dict[b_K_key] = ProcessWeights.convert_tensor_to_hf_format(
1855 b_K_key, W_K_eff_even[:, -1, :], cfg, adapter, l
1856 )
1858 # W_OV = W_V @ W_O
1859 W_V = ProcessWeights.convert_tensor_to_tl_format(
1860 W_V_key, state_dict, state_dict.get(W_V_key), cfg, adapter, l
1861 )
1862 W_O = ProcessWeights.convert_tensor_to_tl_format(
1863 W_O_key, state_dict, state_dict.get(W_O_key), cfg, adapter, l
1864 )
1866 # Factors the bias to be consistent.
1867 b_V = ProcessWeights.convert_tensor_to_tl_format(
1868 b_V_key, state_dict, state_dict.get(b_V_key), cfg, adapter, l
1869 )
1870 b_O = ProcessWeights.convert_tensor_to_tl_format(
1871 b_O_key, state_dict, state_dict.get(b_O_key), cfg, adapter, l
1872 )
1873 assert W_V is not None, f"W_V not found at key {W_V_key}"
1874 assert W_O is not None, f"W_O not found at key {W_O_key}"
1875 assert b_V is not None, f"b_V not found at key {b_V_key}"
1876 assert b_O is not None, f"b_O not found at key {b_O_key}"
1878 # Add singleton dimension for broadcasting
1879 b_V_expanded = einops.rearrange(b_V, "head_index d_head -> head_index d_head 1")
1881 b_V_times_W_O = b_V_expanded * W_O
1883 # Sum over d_head and head_index dimensions
1884 b_V_contribution = b_V_times_W_O.sum(1).sum(0)
1886 effective_bias = b_O + b_V_contribution
1887 state_dict[b_V_key] = ProcessWeights.convert_tensor_to_hf_format(
1888 b_V_key, torch.zeros_like(b_V), cfg, adapter, l
1889 )
1890 state_dict[b_O_key] = ProcessWeights.convert_tensor_to_hf_format(
1891 b_O_key, effective_bias, cfg, adapter, l
1892 )
1894 # Helper class to efficiently deal with low rank factored matrices.
1895 W_OV = FactoredMatrix(W_V, W_O)
1896 U, S, Vh = W_OV.svd()
1897 state_dict[W_V_key] = ProcessWeights.convert_tensor_to_hf_format(
1898 W_V_key, U @ S.diag_embed(), cfg, adapter, l
1899 )
1900 state_dict[W_O_key] = ProcessWeights.convert_tensor_to_hf_format(
1901 W_O_key, utils.transpose(Vh), cfg, adapter, l
1902 )
1904 return state_dict
1906 @overload
1907 @staticmethod
1908 def convert_tensor_to_tl_format(
1909 param_name: str,
1910 model_state_dict: Dict[str, torch.Tensor],
1911 tensor: torch.Tensor,
1912 cfg: Optional["TransformerLensConfig"],
1913 adapter: Optional["ArchitectureAdapter"] = None,
1914 layer_idx: Optional[int] = None,
1915 ) -> torch.Tensor:
1916 ...
1918 @overload
1919 @staticmethod
1920 def convert_tensor_to_tl_format(
1921 param_name: str,
1922 model_state_dict: Dict[str, torch.Tensor],
1923 tensor: None,
1924 cfg: Optional["TransformerLensConfig"],
1925 adapter: Optional["ArchitectureAdapter"] = None,
1926 layer_idx: Optional[int] = None,
1927 ) -> None:
1928 ...
1930 @staticmethod
1931 def convert_tensor_to_tl_format(
1932 param_name: str,
1933 model_state_dict: Dict[str, torch.Tensor],
1934 tensor: Optional[torch.Tensor],
1935 cfg: Optional["TransformerLensConfig"],
1936 adapter: Optional["ArchitectureAdapter"] = None,
1937 layer_idx: Optional[int] = None,
1938 ) -> Optional[torch.Tensor]:
1939 """Convert a tensor from its original format to TransformerLens format.
1941 Args:
1942 param_name: The parameter name in TransformerLens format (e.g., "blocks.0.attn.W_Q")
1943 model_state_dict: The model's state dictionary containing the actual tensors
1944 tensor: The tensor to convert, or None for optional parameters
1945 cfg: Model configuration
1946 adapter: Optional architecture adapter for component retrieval and key translation.
1947 If None, the tensor is returned unchanged.
1948 layer_idx: Layer index (required for layer-specific parameters)
1950 Returns:
1951 The tensor converted to TransformerLens format, or None if the parameter doesn't exist
1952 (which is valid for optional parameters like biases in models that don't use them).
1953 If adapter is None, returns the tensor unchanged.
1954 """
1955 # If no adapter provided, return tensor unchanged (handle None gracefully)
1956 if adapter is None:
1957 return tensor
1959 if ( 1959 ↛ 2040line 1959 didn't jump to line 2040 because the condition on line 1959 was always true
1960 hasattr(adapter, "weight_processing_conversions")
1961 and adapter.weight_processing_conversions is not None
1962 ):
1963 # Create placeholder param name by replacing layer index with {i}
1964 placeholder_param_name = param_name
1965 if "blocks." in param_name:
1966 placeholder_param_name = re.sub(r"blocks\.(\d+)\.", "blocks.{i}.", param_name)
1968 # Check if we have a conversion for this parameter.
1969 # Try exact match first, then strip .weight suffix for adapters
1970 # that define conversions without the suffix (e.g. Pythia's "blocks.{i}.attn.q").
1971 # NOTE: Only strip .weight, NOT .bias — stripping .bias would incorrectly
1972 # match bias keys against weight conversions (e.g. "blocks.{i}.attn.q.bias"
1973 # would match the weight conversion for "blocks.{i}.attn.q").
1974 matched_key = None
1975 if placeholder_param_name in adapter.weight_processing_conversions:
1976 matched_key = placeholder_param_name
1977 elif placeholder_param_name.endswith(".weight"):
1978 stripped = placeholder_param_name[: -len(".weight")]
1979 if stripped in adapter.weight_processing_conversions:
1980 matched_key = stripped
1982 if matched_key is not None:
1983 param_conversion = adapter.weight_processing_conversions[matched_key]
1985 # Handle both ParamProcessingConversion objects and legacy string mappings
1986 if isinstance(param_conversion, str): 1986 ↛ 1989line 1986 didn't jump to line 1989 because the condition on line 1986 was never true
1987 # Legacy string mapping - just return the tensor as-is
1988 # (string mappings are handled elsewhere in the architecture adapter)
1989 return tensor
1990 else:
1991 # Skip conversion for optional parameters that don't exist (e.g. biases)
1992 if tensor is None and param_name not in model_state_dict:
1993 return None
1994 # Try ParamProcessingConversion.convert() first (uses source_key
1995 # to fetch from state dict — needed for split conversions like
1996 # GPT-2's QKV). If source_key resolves to a missing key and we
1997 # already have the tensor, fall back to applying the tensor
1998 # conversion directly (needed for adapters like GPT-Neo whose
1999 # source_key references HF keys not in the bridge state dict).
2000 if (
2001 hasattr(param_conversion, "source_key")
2002 and param_conversion.source_key is not None
2003 ):
2004 resolved_key = param_conversion._resolve_key(
2005 param_name, param_conversion.source_key
2006 )
2007 if resolved_key not in model_state_dict and tensor is not None: 2007 ↛ 2034line 2007 didn't jump to line 2034 because the condition on line 2007 was always true
2008 # Source key not in state dict — the tensor is already in
2009 # bridge format (e.g. already split from combined QKV).
2010 # If the conversion is a ChainTensorConversion that includes
2011 # a SplitTensorConversion, skip the split step since
2012 # it was already applied during bridge construction.
2013 from transformer_lens.conversion_utils.conversion_steps.chain_tensor_conversion import (
2014 ChainTensorConversion,
2015 )
2016 from transformer_lens.conversion_utils.conversion_steps.split_tensor_conversion import (
2017 SplitTensorConversion,
2018 )
2020 tc = param_conversion.tensor_conversion
2021 if isinstance(tc, ChainTensorConversion): 2021 ↛ 2022line 2021 didn't jump to line 2022 because the condition on line 2021 was never true
2022 non_split = [
2023 c
2024 for c in tc.conversions
2025 if not isinstance(c, SplitTensorConversion)
2026 ]
2027 if len(non_split) < len(tc.conversions):
2028 # Apply only the non-split conversions
2029 result = tensor
2030 for conv in non_split:
2031 result = conv.handle_conversion(result, model_state_dict)
2032 return result
2033 return tc.convert(tensor, model_state_dict)
2034 return param_conversion.convert(model_state_dict, param_name)
2035 else:
2036 # No conversion defined, return tensor as-is (may be None for optional params)
2037 return tensor
2038 else:
2039 # No conversions defined, return tensor as-is (may be None for optional params)
2040 return tensor
2042 @overload
2043 @staticmethod
2044 def convert_tensor_to_hf_format(
2045 param_name: str,
2046 tensor: torch.Tensor,
2047 cfg: Optional["TransformerLensConfig"],
2048 adapter: Optional["ArchitectureAdapter"] = None,
2049 layer_idx: Optional[int] = None,
2050 ) -> torch.Tensor:
2051 ...
2053 @overload
2054 @staticmethod
2055 def convert_tensor_to_hf_format(
2056 param_name: str,
2057 tensor: None,
2058 cfg: Optional["TransformerLensConfig"],
2059 adapter: Optional["ArchitectureAdapter"] = None,
2060 layer_idx: Optional[int] = None,
2061 ) -> None:
2062 ...
2064 @staticmethod
2065 def convert_tensor_to_hf_format(
2066 param_name: str,
2067 tensor: Optional[torch.Tensor],
2068 cfg: Optional["TransformerLensConfig"],
2069 adapter: Optional["ArchitectureAdapter"] = None,
2070 layer_idx: Optional[int] = None,
2071 ) -> Optional[torch.Tensor]:
2072 """Convert a tensor from TransformerLens format back to its original format.
2074 Args:
2075 param_name: The parameter name in TransformerLens format (e.g., "blocks.0.attn.W_Q")
2076 tensor: The tensor to convert (in TransformerLens format), or None if parameter is optional
2077 cfg: Model configuration
2078 adapter: Optional architecture adapter for component retrieval and key translation.
2079 If None, the tensor is returned unchanged.
2080 layer_idx: Layer index (required for layer-specific parameters)
2082 Returns:
2083 The tensor converted back to original format, or None if tensor was None.
2084 If adapter is None, returns the tensor unchanged.
2085 """
2086 # Handle None tensors (optional parameters)
2087 if tensor is None: 2087 ↛ 2088line 2087 didn't jump to line 2088 because the condition on line 2087 was never true
2088 return None
2090 # If no adapter provided, return tensor unchanged
2091 if adapter is None:
2092 return tensor
2094 if ( 2094 ↛ 2146line 2094 didn't jump to line 2146 because the condition on line 2094 was always true
2095 hasattr(adapter, "weight_processing_conversions")
2096 and adapter.weight_processing_conversions is not None
2097 ):
2098 # Create placeholder param name by replacing layer index with {i}
2099 placeholder_param_name = param_name
2100 if "blocks." in param_name:
2101 placeholder_param_name = re.sub(r"blocks\.(\d+)\.", "blocks.{i}.", param_name)
2103 # Check if we have a conversion for this parameter.
2104 # Try exact match first, then strip .weight suffix (not .bias — see convert_tensor_to_tl_format).
2105 matched_key = None
2106 if placeholder_param_name in adapter.weight_processing_conversions:
2107 matched_key = placeholder_param_name
2108 elif placeholder_param_name.endswith(".weight"):
2109 stripped = placeholder_param_name[: -len(".weight")]
2110 if stripped in adapter.weight_processing_conversions:
2111 matched_key = stripped
2113 if matched_key is not None:
2114 param_conversion = adapter.weight_processing_conversions[matched_key]
2116 # Handle both ParamProcessingConversion objects and legacy string mappings
2117 if isinstance(param_conversion, str): 2117 ↛ 2119line 2117 didn't jump to line 2119 because the condition on line 2117 was never true
2118 # Legacy string mapping - just return the tensor as-is
2119 return tensor
2120 else:
2121 # Revert the conversion. For ChainTensorConversions that include
2122 # SplitTensorConversion, skip the split revert step (which is a
2123 # no-op anyway) to match the forward conversion path.
2124 from transformer_lens.conversion_utils.conversion_steps.chain_tensor_conversion import (
2125 ChainTensorConversion,
2126 )
2127 from transformer_lens.conversion_utils.conversion_steps.split_tensor_conversion import (
2128 SplitTensorConversion,
2129 )
2131 tc = param_conversion.tensor_conversion
2132 if isinstance(tc, ChainTensorConversion): 2132 ↛ 2133line 2132 didn't jump to line 2133 because the condition on line 2132 was never true
2133 non_split = [
2134 c for c in tc.conversions if not isinstance(c, SplitTensorConversion)
2135 ]
2136 if len(non_split) < len(tc.conversions):
2137 # Revert only the non-split conversions in reverse order
2138 result = tensor
2139 for conv in reversed(non_split):
2140 result = conv.revert(result)
2141 return result
2142 return param_conversion.revert(tensor)
2143 else:
2144 return tensor
2145 else:
2146 return tensor
2148 @staticmethod
2149 def distribute_weights_to_components(
2150 state_dict: Dict[str, torch.Tensor],
2151 component_mapping: Dict[str, Any],
2152 verbose: bool = False,
2153 ) -> None:
2154 """Distribute processed weights from state_dict to generalized components.
2156 This function loops through the component_mapping and extracts relevant weights
2157 for each component using filter_dict_by_prefix, then calls set_processed_weights
2158 on each component. For list components (like blocks), it determines the number
2159 of items and distributes weights to each indexed component.
2161 Args:
2162 state_dict: Dictionary of processed weights in MODERN TransformerLens format
2163 (e.g., blocks.0.attn.q.weight, not transformer.h.0.attn.q.weight)
2164 component_mapping: Dictionary (real_components) mapping TL keys to tuples of
2165 (remote_path, component_instance), where component_instance can be either
2166 a single component or a list of components
2167 verbose: If True, print detailed information about weight distribution
2169 Example:
2170 For a real_components mapping like:
2171 {
2172 "embed": ("transformer.wte", <EmbeddingBridge instance>),
2173 "blocks": ("transformer.h", [<BlockBridge 0>, <BlockBridge 1>, ...]),
2174 "unembed": ("lm_head", <UnembeddingBridge instance>)
2175 }
2177 With modern TL keys in state_dict like "embed.weight", "blocks.0.attn.q.weight":
2178 1. Extract weights starting with "embed" and pass to embed component
2179 2. For blocks, extract all "blocks.*" weights, determine the number of blocks,
2180 then for each block index, extract weights for that specific block
2181 3. Extract "unembed" weights and pass to unembed component
2182 """
2183 if verbose: 2183 ↛ 2184line 2183 didn't jump to line 2184 because the condition on line 2183 was never true
2184 print(f"\n{'='*80}")
2185 print(f"distribute_weights_to_components: Starting weight distribution")
2186 print(f"State dict has {len(state_dict)} keys")
2187 print(f"Component mapping has {len(component_mapping)} components")
2188 print(f"{'='*80}\n")
2190 for component_name, component_tuple in component_mapping.items():
2191 # component_mapping is real_components format: (remote_path, instance)
2192 # instance can be either a single component or a list of components
2193 if not isinstance(component_tuple, tuple): 2193 ↛ 2194line 2193 didn't jump to line 2194 because the condition on line 2193 was never true
2194 raise ValueError(
2195 f"Expected tuple for component '{component_name}' in real_components, "
2196 f"but got {type(component_tuple).__name__}: {component_tuple}"
2197 )
2198 remote_key, component = component_tuple
2199 is_list = isinstance(component, list)
2201 # Use the component_name (TL format) as prefix instead of remote_key (HF format)
2202 # since state_dict now has modern TL keys
2203 tl_prefix = component_name
2205 if verbose: 2205 ↛ 2206line 2205 didn't jump to line 2206 because the condition on line 2205 was never true
2206 print(f"\nProcessing component: {component_name}")
2207 print(f" Remote key (HF): {remote_key}")
2208 print(f" TL prefix: {tl_prefix}")
2209 print(f" Is list: {is_list}")
2211 if is_list:
2212 # This is a list component like "blocks"
2213 # Extract all weights that start with this prefix
2214 all_list_weights = filter_dict_by_prefix(state_dict, tl_prefix)
2216 if verbose: 2216 ↛ 2217line 2216 didn't jump to line 2217 because the condition on line 2216 was never true
2217 print(f" Found {len(all_list_weights)} weights for list component")
2218 print(f" List has {len(component)} instances")
2220 # Component is a list of actual instances
2221 for i, instance in enumerate(component):
2222 # Extract weights for this specific index
2223 # This will get keys like "0.attn.q.weight" and strip the "0." to get "attn.q.weight"
2224 indexed_weights = filter_dict_by_prefix(all_list_weights, str(i))
2226 if verbose: 2226 ↛ 2227line 2226 didn't jump to line 2227 because the condition on line 2226 was never true
2227 print(f" Instance {i}: Found {len(indexed_weights)} weights")
2228 for key in indexed_weights.keys():
2229 print(f" - {key}")
2231 # Skip if no weights found for this component (e.g., Q/K/V Linear sub-components
2232 # that get their weights from parent JointQKVAttentionBridge)
2233 if len(indexed_weights) == 0: 2233 ↛ 2234line 2233 didn't jump to line 2234 because the condition on line 2233 was never true
2234 if verbose:
2235 print(f" Skipping instance {i} - no weights found")
2236 continue
2238 instance.set_processed_weights(indexed_weights, verbose=verbose)
2239 else:
2240 # This is a single component (not a list)
2241 component_weights = filter_dict_by_prefix(state_dict, tl_prefix)
2243 if verbose: 2243 ↛ 2244line 2243 didn't jump to line 2244 because the condition on line 2243 was never true
2244 print(f" Found {len(component_weights)} weights for single component")
2245 for key in component_weights.keys():
2246 print(f" - {key}")
2248 # Skip if no weights found for this component
2249 if len(component_weights) == 0:
2250 if verbose: 2250 ↛ 2251line 2250 didn't jump to line 2251 because the condition on line 2250 was never true
2251 print(f" Skipping component - no weights found")
2252 continue
2254 component.set_processed_weights(component_weights, verbose=verbose)