Coverage for transformer_lens/weight_processing.py: 79%
876 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-10-23 20:45 +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:
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: 599 ↛ 601line 599 didn't jump to line 601 because the condition on line 599 was never true
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:
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:
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:
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:
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 # RMSNorm outputs are not zero-mean, so centering would change predictions.
1042 if center_weights and not getattr(cfg, "final_rms", False):
1043 unembed_weight_centered = ProcessWeights.convert_tensor_to_tl_format(
1044 unembed_W_U_key, state_dict, state_dict.get(unembed_W_U_key), cfg, adapter, None
1045 )
1046 assert (
1047 unembed_weight_centered is not None
1048 ), f"Unembed weight not found at key {unembed_W_U_key}"
1049 if len(unembed_weight_centered.shape) == 2: 1049 ↛ 1070line 1049 didn't jump to line 1070 because the condition on line 1049 was always true
1050 # Center along d_model: detect TL vs HF format
1051 d_vocab = getattr(cfg, "d_vocab", None) if cfg is not None else None
1052 if (
1053 d_vocab is not None
1054 and unembed_weight_centered.shape[0] == d_vocab
1055 and unembed_weight_centered.shape[-1] != d_vocab
1056 ):
1057 # HF format [d_vocab, d_model] — center along dim=-1
1058 unembed_weight_centered = (
1059 unembed_weight_centered - unembed_weight_centered.mean(-1, keepdim=True)
1060 )
1061 else:
1062 # TL format [d_model, d_vocab] — center along dim=0
1063 unembed_weight_centered = (
1064 unembed_weight_centered - unembed_weight_centered.mean(0, keepdim=True)
1065 )
1066 state_dict[unembed_W_U_key] = ProcessWeights.convert_tensor_to_hf_format(
1067 unembed_W_U_key, unembed_weight_centered, cfg, adapter, None
1068 )
1069 else:
1070 raise ValueError(
1071 f"Unexpected unembedding weight shape: {unembed_weight_centered.shape}"
1072 )
1074 return state_dict
1076 @staticmethod
1077 def _fold_final_rms_bias(
1078 state_dict: Dict[str, torch.Tensor], cfg, fold_biases: bool, adapter
1079 ) -> Dict[str, torch.Tensor]:
1080 """Fold final RMS bias into unembedding (separate from regular unembed folding).
1082 Args:
1083 state_dict: The state dictionary to process (modified in place)
1084 cfg: Model configuration object
1085 fold_biases: Whether to fold LayerNorm biases
1086 adapter: Optional architecture adapter for parameter key translation
1087 """
1088 unembed_b_U_key = ProcessWeights._get_param_key("unembed.b_U", adapter)
1089 unembed_W_U_key = ProcessWeights._get_param_key("unembed.W_U", adapter)
1090 ln_final_b_key = ProcessWeights._get_param_key("ln_final.b", adapter)
1091 has_unembed_bias = unembed_b_U_key in state_dict
1092 has_ln_final_bias = ln_final_b_key in state_dict
1093 if (
1094 not getattr(cfg, "final_rms", False)
1095 and fold_biases
1096 and has_unembed_bias
1097 and has_ln_final_bias
1098 ):
1099 unembed_weight = ProcessWeights.convert_tensor_to_tl_format(
1100 unembed_W_U_key, state_dict, state_dict.get(unembed_W_U_key), cfg, adapter, None
1101 )
1102 ln_bias = state_dict[ln_final_b_key]
1103 assert unembed_weight is not None, f"Unembed weight not found at key {unembed_W_U_key}"
1104 if len(unembed_weight.shape) == 2 and len(ln_bias.shape) == 1: 1104 ↛ 1114line 1104 didn't jump to line 1114 because the condition on line 1104 was always true
1105 if unembed_weight.shape[1] == ln_bias.shape[0]:
1106 bias_contribution = (unembed_weight * ln_bias[None, :]).sum(dim=-1)
1107 elif unembed_weight.shape[0] == ln_bias.shape[0]: 1107 ↛ 1110line 1107 didn't jump to line 1110 because the condition on line 1107 was always true
1108 bias_contribution = (unembed_weight * ln_bias[:, None]).sum(dim=-2)
1109 else:
1110 raise ValueError(
1111 f"Cannot broadcast unembedding weight {unembed_weight.shape} with layer norm bias {ln_bias.shape}"
1112 )
1113 else:
1114 raise ValueError(
1115 f"Unexpected tensor shapes: unembedding {unembed_weight.shape}, layer norm bias {ln_bias.shape}"
1116 )
1117 unembed_b_U = ProcessWeights.convert_tensor_to_tl_format(
1118 unembed_b_U_key, state_dict, state_dict.get(unembed_b_U_key), cfg, adapter, None
1119 )
1120 assert unembed_b_U is not None, f"Unembed bias not found at key {unembed_b_U_key}"
1121 new_unembed_b_U = unembed_b_U + bias_contribution
1122 state_dict[unembed_b_U_key] = ProcessWeights.convert_tensor_to_hf_format(
1123 unembed_b_U_key, new_unembed_b_U, cfg, adapter, None
1124 )
1125 if ln_final_b_key in state_dict: 1125 ↛ 1127line 1125 didn't jump to line 1127 because the condition on line 1125 was always true
1126 state_dict[ln_final_b_key] = torch.zeros_like(ln_bias)
1127 alternate_final_b_key = (
1128 ln_final_b_key.replace("ln_f", "ln_final")
1129 if "ln_f" in ln_final_b_key
1130 else ln_final_b_key.replace("ln_final", "ln_f")
1131 )
1132 if alternate_final_b_key != ln_final_b_key and alternate_final_b_key in state_dict: 1132 ↛ 1133line 1132 didn't jump to line 1133 because the condition on line 1132 was never true
1133 state_dict[alternate_final_b_key] = torch.zeros_like(ln_bias)
1135 return state_dict
1137 @staticmethod
1138 def fold_layer_norm(
1139 state_dict: Dict[str, torch.Tensor],
1140 cfg,
1141 fold_biases: bool = True,
1142 center_weights: bool = True,
1143 adapter=None,
1144 ) -> Dict[str, torch.Tensor]:
1145 """Fold Layer Norm. Can also be used to fold RMS Norm, when fold_biases and center_weights are set to False.
1147 Takes in a state dict from a pretrained model, formatted to be consistent with
1148 HookedTransformer but with LayerNorm weights and biases. Folds these into the neighbouring
1149 weights. See further_comments.md for more details.
1151 Args:
1152 state_dict (Dict[str, torch.Tensor]): State dict of pretrained model.
1153 cfg: Model configuration object with n_layers, n_key_value_heads, etc.
1154 fold_biases (bool): Enables folding of LN biases. Should be disabled when RMS Norm is used.
1155 center_weights (bool): Enables the centering of weights after folding in LN. Should be disabled when RMS Norm is used.
1156 adapter: Optional architecture adapter for parameter key translation.
1158 Returns:
1159 Dict[str, torch.Tensor]: Modified state dict with LayerNorm folded into linear layers.
1160 """
1161 # Make a deep copy to avoid modifying the original
1162 state_dict = {
1163 k: v.clone() if isinstance(v, torch.Tensor) else v for k, v in state_dict.items()
1164 }
1165 gqa = "" if getattr(cfg, "n_key_value_heads", None) is None else "_"
1166 for l in range(cfg.n_layers):
1167 state_dict = ProcessWeights._fold_layer(
1168 state_dict, cfg, l, fold_biases, center_weights, adapter, gqa
1169 )
1170 state_dict = ProcessWeights._fold_final_rms_bias(state_dict, cfg, fold_biases, adapter)
1171 state_dict = ProcessWeights._fold_unembed_layer_norm(
1172 state_dict, cfg, fold_biases, center_weights, adapter
1173 )
1174 return state_dict
1176 @staticmethod
1177 def center_writing_weights(
1178 state_dict: Dict[str, torch.Tensor], cfg, adapter=None
1179 ) -> Dict[str, torch.Tensor]:
1180 """Center Writing Weights.
1182 Centers the weights of the model that write to the residual stream - W_out, W_E, W_pos and
1183 W_out. This is done by subtracting the mean of the weights from the weights themselves. This
1184 is done in-place. See fold_layer_norm for more details.
1186 Args:
1187 state_dict (Dict[str, torch.Tensor]): State dict of the model.
1188 cfg: Model configuration object.
1189 adapter: Optional architecture adapter for parameter key translation.
1191 Returns:
1192 Dict[str, torch.Tensor]: Modified state dict with centered writing weights.
1193 """
1194 # Post-norm models leave the first attention's input un-normed, so centering
1195 # the embedding would shift a residual stream nothing re-normalizes.
1196 architecture = getattr(cfg, "original_architecture", None)
1197 if architecture in POST_NORM_ARCHITECTURES:
1198 print(f"Not centering embedding weights for {architecture}")
1199 else:
1200 # Make a deep copy to avoid modifying the original
1201 embed_W_E_key = ProcessWeights._get_param_key("embed.W_E", adapter)
1202 try:
1203 pos_embed_W_pos_key = (
1204 ProcessWeights._get_param_key("pos_embed.W_pos", adapter)
1205 if getattr(cfg, "positional_embedding_type", "standard")
1206 not in ("rotary", "alibi", "none")
1207 else None
1208 )
1209 except ValueError:
1210 pos_embed_W_pos_key = None
1211 if embed_W_E_key not in state_dict: 1211 ↛ 1212line 1211 didn't jump to line 1212 because the condition on line 1211 was never true
1212 raise KeyError(
1213 f"Expected embedding key '{embed_W_E_key}' not found in state_dict. Available keys: {list(state_dict.keys())[:10]}..."
1214 )
1215 embed_W_E = ProcessWeights.convert_tensor_to_tl_format(
1216 embed_W_E_key, state_dict, state_dict.get(embed_W_E_key), cfg, adapter, None
1217 )
1218 assert embed_W_E is not None, f"Embedding not found at key {embed_W_E_key}"
1219 embed_W_E = embed_W_E - embed_W_E.mean(-1, keepdim=True)
1220 state_dict[embed_W_E_key] = ProcessWeights.convert_tensor_to_hf_format(
1221 embed_W_E_key, embed_W_E, cfg, adapter, None
1222 )
1224 if (
1225 getattr(cfg, "positional_embedding_type", "standard")
1226 not in ("rotary", "alibi", "none")
1227 and pos_embed_W_pos_key is not None
1228 ):
1229 if pos_embed_W_pos_key not in state_dict: 1229 ↛ 1230line 1229 didn't jump to line 1230 because the condition on line 1229 was never true
1230 raise KeyError(
1231 f"Expected positional embedding key '{pos_embed_W_pos_key}' not found in state_dict. Available keys: {list(state_dict.keys())[:10]}..."
1232 )
1233 pos_embed_W_pos = ProcessWeights.convert_tensor_to_tl_format(
1234 pos_embed_W_pos_key,
1235 state_dict,
1236 state_dict.get(pos_embed_W_pos_key),
1237 cfg,
1238 adapter,
1239 None,
1240 )
1241 assert (
1242 pos_embed_W_pos is not None
1243 ), f"Positional embedding not found at key {pos_embed_W_pos_key}"
1244 pos_embed_W_pos = pos_embed_W_pos - pos_embed_W_pos.mean(-1, keepdim=True)
1245 state_dict[pos_embed_W_pos_key] = ProcessWeights.convert_tensor_to_hf_format(
1246 pos_embed_W_pos_key, pos_embed_W_pos, cfg, adapter, None
1247 )
1248 for l in range(cfg.n_layers):
1249 attn_W_O_key = ProcessWeights._get_param_key(f"blocks.{l}.attn.W_O", adapter)
1250 attn_b_O_key = ProcessWeights._get_param_key(f"blocks.{l}.attn.b_O", adapter)
1251 try:
1252 mlp_W_out_key = ProcessWeights._resolve_state_dict_key(
1253 state_dict, ProcessWeights._get_param_key(f"blocks.{l}.mlp.W_out", adapter), l
1254 )
1255 mlp_b_out_key = ProcessWeights._resolve_state_dict_key(
1256 state_dict, ProcessWeights._get_param_key(f"blocks.{l}.mlp.b_out", adapter), l
1257 )
1258 except ValueError:
1259 mlp_W_out_key = None
1260 mlp_b_out_key = None
1261 if attn_W_O_key in state_dict:
1262 attn_W_O = ProcessWeights.convert_tensor_to_tl_format(
1263 attn_W_O_key, state_dict, state_dict.get(attn_W_O_key), cfg, adapter, l
1264 )
1265 assert attn_W_O is not None, f"Attention W_O not found at key {attn_W_O_key}"
1266 attn_W_O = attn_W_O - attn_W_O.mean(-1, keepdim=True)
1267 state_dict[attn_W_O_key] = ProcessWeights.convert_tensor_to_hf_format(
1268 attn_W_O_key, attn_W_O, cfg, adapter, l
1269 )
1270 if attn_b_O_key in state_dict:
1271 attn_b_O = ProcessWeights.convert_tensor_to_tl_format(
1272 attn_b_O_key, state_dict, state_dict.get(attn_b_O_key), cfg, adapter, l
1273 )
1274 assert attn_b_O is not None, f"Attention b_O not found at key {attn_b_O_key}"
1275 attn_b_O = attn_b_O - attn_b_O.mean()
1276 state_dict[attn_b_O_key] = ProcessWeights.convert_tensor_to_hf_format(
1277 attn_b_O_key, attn_b_O, cfg, adapter, l
1278 )
1279 if not getattr(cfg, "attn_only", False):
1280 is_moe = getattr(cfg, "num_experts", None) is not None and cfg.num_experts > 0
1281 if is_moe:
1282 num_experts = cfg.num_experts
1283 for e in range(num_experts):
1284 expert_W_out_key = None
1285 expert_b_out_key = None
1286 expert_W_out_patterns = [
1287 f"blocks.{l}.mlp.experts.{e}.W_out",
1288 f"blocks.{l}.mlp.experts.{e}.W_out.weight",
1289 ]
1290 for pattern in expert_W_out_patterns:
1291 if pattern in state_dict: 1291 ↛ 1292line 1291 didn't jump to line 1292 because the condition on line 1291 was never true
1292 expert_W_out_key = pattern
1293 break
1294 if expert_W_out_key is None and adapter: 1294 ↛ 1295line 1294 didn't jump to line 1295 because the condition on line 1294 was never true
1295 try:
1296 candidate = ProcessWeights._get_param_key(
1297 f"blocks.{l}.mlp.experts.{e}.W_out", adapter
1298 )
1299 expert_W_out_key = ProcessWeights._resolve_state_dict_key(
1300 state_dict, candidate, l
1301 )
1302 except ValueError:
1303 pass
1304 if expert_W_out_key and expert_W_out_key in state_dict: 1304 ↛ 1305line 1304 didn't jump to line 1305 because the condition on line 1304 was never true
1305 expert_W_out = ProcessWeights.convert_tensor_to_tl_format(
1306 expert_W_out_key,
1307 state_dict,
1308 state_dict.get(expert_W_out_key),
1309 cfg,
1310 adapter,
1311 l,
1312 )
1313 assert (
1314 expert_W_out is not None
1315 ), f"Expert W_out not found at key {expert_W_out_key}"
1316 expert_W_out = expert_W_out - expert_W_out.mean(-1, keepdim=True)
1317 state_dict[
1318 expert_W_out_key
1319 ] = ProcessWeights.convert_tensor_to_hf_format(
1320 expert_W_out_key, expert_W_out, cfg, adapter, l
1321 )
1322 expert_b_out_patterns = [
1323 f"blocks.{l}.mlp.experts.{e}.b_out",
1324 f"blocks.{l}.mlp.experts.{e}.b_out.bias",
1325 ]
1326 for pattern in expert_b_out_patterns:
1327 if pattern in state_dict: 1327 ↛ 1328line 1327 didn't jump to line 1328 because the condition on line 1327 was never true
1328 expert_b_out_key = pattern
1329 break
1330 if expert_b_out_key is None and adapter: 1330 ↛ 1331line 1330 didn't jump to line 1331 because the condition on line 1330 was never true
1331 try:
1332 candidate = ProcessWeights._get_param_key(
1333 f"blocks.{l}.mlp.experts.{e}.b_out", adapter
1334 )
1335 expert_b_out_key = ProcessWeights._resolve_state_dict_key(
1336 state_dict, candidate, l
1337 )
1338 except ValueError:
1339 pass
1340 if expert_b_out_key and expert_b_out_key in state_dict: 1340 ↛ 1341line 1340 didn't jump to line 1341 because the condition on line 1340 was never true
1341 expert_b_out = ProcessWeights.convert_tensor_to_tl_format(
1342 expert_b_out_key,
1343 state_dict,
1344 state_dict.get(expert_b_out_key),
1345 cfg,
1346 adapter,
1347 l,
1348 )
1349 assert (
1350 expert_b_out is not None
1351 ), f"Expert b_out not found at key {expert_b_out_key}"
1352 expert_b_out = expert_b_out - expert_b_out.mean()
1353 state_dict[
1354 expert_b_out_key
1355 ] = ProcessWeights.convert_tensor_to_hf_format(
1356 expert_b_out_key, expert_b_out, cfg, adapter, l
1357 )
1358 elif mlp_W_out_key is not None and mlp_W_out_key in state_dict:
1359 mlp_W_out = ProcessWeights.convert_tensor_to_tl_format(
1360 mlp_W_out_key, state_dict, state_dict.get(mlp_W_out_key), cfg, adapter, l
1361 )
1362 assert mlp_W_out is not None, f"MLP W_out not found at key {mlp_W_out_key}"
1363 # Center along d_model dimension. In TL format W_out is [d_mlp, d_model]
1364 # so d_model is dim=-1. But bridge adapters may keep HF format
1365 # [d_model, d_mlp] where d_model is dim=0. Detect via cfg.d_model.
1366 if mlp_W_out.shape[-1] == cfg.d_model:
1367 mlp_W_out = mlp_W_out - mlp_W_out.mean(-1, keepdim=True)
1368 elif mlp_W_out.shape[0] == cfg.d_model: 1368 ↛ 1371line 1368 didn't jump to line 1371 because the condition on line 1368 was always true
1369 mlp_W_out = mlp_W_out - mlp_W_out.mean(0, keepdim=True)
1370 else:
1371 mlp_W_out = mlp_W_out - mlp_W_out.mean(-1, keepdim=True)
1372 state_dict[mlp_W_out_key] = ProcessWeights.convert_tensor_to_hf_format(
1373 mlp_W_out_key, mlp_W_out, cfg, adapter, l
1374 )
1375 if mlp_b_out_key is not None and mlp_b_out_key in state_dict:
1376 mlp_b_out = ProcessWeights.convert_tensor_to_tl_format(
1377 mlp_b_out_key,
1378 state_dict,
1379 state_dict.get(mlp_b_out_key),
1380 cfg,
1381 adapter,
1382 l,
1383 )
1384 assert mlp_b_out is not None, f"MLP b_out not found at key {mlp_b_out_key}"
1385 mlp_b_out = mlp_b_out - mlp_b_out.mean()
1386 state_dict[mlp_b_out_key] = ProcessWeights.convert_tensor_to_hf_format(
1387 mlp_b_out_key, mlp_b_out, cfg, adapter, l
1388 )
1389 return state_dict
1391 @staticmethod
1392 def center_unembed(
1393 state_dict: Dict[str, torch.Tensor], cfg=None, adapter=None
1394 ) -> Dict[str, torch.Tensor]:
1395 """Center the unembedding weights W_U.
1397 This is done by subtracting the mean of the weights from the weights themselves. This is
1398 done in-place. As softmax is translation invariant, this changes the logits but not the log
1399 probs, and makes the model logits (slightly) more interpretable - when trying to understand
1400 how components contribute to the logits, we'll be less misled by components that just add
1401 something to every logit.
1403 Args:
1404 state_dict (Dict[str, torch.Tensor]): State dict of the model.
1405 cfg: Model configuration (used to determine d_vocab for correct centering dimension).
1406 adapter: Optional architecture adapter for parameter key translation.
1408 Returns:
1409 Dict[str, torch.Tensor]: Modified state dict with centered unembedding weights.
1410 """
1411 # Make a deep copy to avoid modifying the original
1412 state_dict = {
1413 k: v.clone() if isinstance(v, torch.Tensor) else v for k, v in state_dict.items()
1414 }
1415 unembed_W_U_key = ProcessWeights._get_param_key("unembed.W_U", adapter)
1416 unembed_b_U_key = ProcessWeights._get_param_key("unembed.b_U", adapter)
1417 if unembed_W_U_key not in state_dict:
1418 raise KeyError(
1419 f"Expected unembedding weight key '{unembed_W_U_key}' not found in state_dict. Available keys: {list(state_dict.keys())[:10]}..."
1420 )
1421 W_U = ProcessWeights.convert_tensor_to_tl_format(
1422 unembed_W_U_key, state_dict, state_dict.get(unembed_W_U_key), None, adapter, None
1423 )
1424 assert W_U is not None, f"Unembed weight not found at key {unembed_W_U_key}"
1426 # Detect W_U format to center along correct dim (wrong dim corrupts output)
1427 vocab_dim = -1 # Default: TL format [d_model, d_vocab]
1428 if cfg is not None:
1429 d_vocab = getattr(cfg, "d_vocab", None)
1430 d_model = getattr(cfg, "d_model", None)
1431 if d_vocab is not None and W_U.shape[0] == d_vocab and W_U.shape[-1] != d_vocab:
1432 # HF format [d_vocab, d_model] — center along dim=0
1433 vocab_dim = 0
1434 elif d_model is not None and W_U.shape[-1] == d_model and W_U.shape[0] != d_model:
1435 # Padded-vocab checkpoints (e.g. HyenaDNA pads 12→16 rows) defeat
1436 # the d_vocab match; the d_model axis is unambiguous.
1437 vocab_dim = 0
1438 W_U = W_U - W_U.mean(vocab_dim, keepdim=True)
1439 state_dict[unembed_W_U_key] = ProcessWeights.convert_tensor_to_hf_format(
1440 unembed_W_U_key, W_U, None, adapter, None
1441 )
1442 if unembed_b_U_key in state_dict:
1443 unembed_b_U = ProcessWeights.convert_tensor_to_tl_format(
1444 unembed_b_U_key, state_dict, state_dict.get(unembed_b_U_key), None, adapter, None
1445 )
1446 assert unembed_b_U is not None, f"Unembed bias not found at key {unembed_b_U_key}"
1447 unembed_b_U = unembed_b_U - unembed_b_U.mean()
1448 state_dict[unembed_b_U_key] = ProcessWeights.convert_tensor_to_hf_format(
1449 unembed_b_U_key, unembed_b_U, None, adapter, None
1450 )
1451 return state_dict
1453 @staticmethod
1454 def fold_value_biases(
1455 state_dict: Dict[str, torch.Tensor], cfg, adapter=None
1456 ) -> Dict[str, torch.Tensor]:
1457 """Fold the value biases into the output bias.
1459 Because attention patterns add up to 1, the value biases always have a constant effect on a
1460 head's output. Further, as the outputs of each head in a layer add together, each head's
1461 value bias has a constant effect on the *layer's* output, which can make it harder to
1462 interpret the effect of any given head, and it doesn't matter which head a bias is
1463 associated with. We can factor this all into a single output bias to the layer, and make it
1464 easier to interpret the head's output. Formally, we take b_O_new = b_O_original +
1465 sum_head(b_V_head @ W_O_head).
1467 Args:
1468 state_dict (Dict[str, torch.Tensor]): State dict of the model.
1469 cfg: Model configuration object.
1470 adapter: Optional architecture adapter for parameter key translation.
1472 Returns:
1473 Dict[str, torch.Tensor]: Modified state dict with value biases folded into output bias.
1474 """
1475 # Make a deep copy to avoid modifying the original
1476 state_dict = {
1477 k: v.clone() if isinstance(v, torch.Tensor) else v for k, v in state_dict.items()
1478 }
1479 layer = 0
1480 for layer in range(cfg.n_layers):
1481 split_v_bias_key = f"blocks.{layer}.attn.v.bias"
1482 if split_v_bias_key in state_dict:
1483 b_V_key = split_v_bias_key
1484 W_O_key = ProcessWeights._get_param_key(f"blocks.{layer}.attn.W_O", adapter)
1485 b_O_key = ProcessWeights._get_param_key(f"blocks.{layer}.attn.b_O", adapter)
1486 else:
1487 if getattr(cfg, "n_key_value_heads", None) is None:
1488 b_V_key = ProcessWeights._get_param_key(f"blocks.{layer}.attn.b_V", adapter)
1489 else:
1490 b_V_key = ProcessWeights._get_param_key(f"blocks.{layer}.attn._b_V", adapter)
1491 W_O_key = ProcessWeights._get_param_key(f"blocks.{layer}.attn.W_O", adapter)
1492 b_O_key = ProcessWeights._get_param_key(f"blocks.{layer}.attn.b_O", adapter)
1493 if b_V_key in state_dict:
1494 b_V = ProcessWeights.convert_tensor_to_tl_format(
1495 b_V_key, state_dict, state_dict.get(b_V_key), cfg, adapter, layer
1496 )
1497 assert b_V is not None, f"Value bias not found at key {b_V_key}"
1498 if b_V.numel() == 0: 1498 ↛ 1499line 1498 didn't jump to line 1499 because the condition on line 1498 was never true
1499 continue
1500 W_O = ProcessWeights.convert_tensor_to_tl_format(
1501 W_O_key, state_dict, state_dict.get(W_O_key), cfg, adapter, layer
1502 )
1503 assert W_O is not None, f"Attention W_O not found at key {W_O_key}"
1504 if b_O_key not in state_dict:
1505 # Create zero b_O to absorb the folded value bias
1506 b_O_original = torch.zeros(cfg.d_model, dtype=b_V.dtype, device=b_V.device)
1507 state_dict[b_O_key] = b_O_original
1508 else:
1509 b_O_original_maybe = ProcessWeights.convert_tensor_to_tl_format(
1510 b_O_key, state_dict, state_dict.get(b_O_key), cfg, adapter, layer
1511 )
1512 assert (
1513 b_O_original_maybe is not None
1514 ), f"Attention b_O not found at key {b_O_key}"
1515 b_O_original = b_O_original_maybe
1516 # Align W_O / b_O to b_V's device.
1517 if W_O.device != b_V.device: 1517 ↛ 1518line 1517 didn't jump to line 1518 because the condition on line 1517 was never true
1518 W_O = W_O.to(b_V.device)
1519 if b_O_original.device != b_V.device: 1519 ↛ 1520line 1519 didn't jump to line 1520 because the condition on line 1519 was never true
1520 b_O_original = b_O_original.to(b_V.device)
1521 is_split_format = ".attn.v.bias" in b_V_key or ".attn.k.bias" in b_V_key
1522 if is_split_format and len(b_V.shape) == 1 and (len(W_O.shape) == 2): 1522 ↛ 1523line 1522 didn't jump to line 1523 because the condition on line 1522 was never true
1523 n_heads = cfg.n_heads
1524 d_head = cfg.d_head
1525 d_model = cfg.d_model
1526 b_V_only = b_V
1527 b_V_reshaped = b_V_only.reshape(n_heads, d_head)
1528 W_O_reshaped = einops.rearrange(W_O, "(i h) m -> i h m", i=n_heads)
1529 folded_b_O = b_O_original + (b_V_reshaped[:, :, None] * W_O_reshaped).sum(
1530 [0, 1]
1531 )
1532 state_dict[b_O_key] = ProcessWeights.convert_tensor_to_hf_format(
1533 b_O_key, folded_b_O, cfg, adapter, layer
1534 )
1535 tl_b_O_key = f"blocks.{layer}.attn.b_O"
1536 if tl_b_O_key in state_dict:
1537 state_dict[tl_b_O_key] = ProcessWeights.convert_tensor_to_hf_format(
1538 tl_b_O_key, folded_b_O, cfg, adapter, layer
1539 )
1540 state_dict[b_V_key] = ProcessWeights.convert_tensor_to_hf_format(
1541 b_V_key, torch.zeros_like(b_V), cfg, adapter, layer
1542 )
1543 elif len(b_V.shape) == 1 and len(W_O.shape) == 2: 1543 ↛ 1544line 1543 didn't jump to line 1544 because the condition on line 1543 was never true
1544 n_heads = cfg.n_heads
1545 d_head = cfg.d_head
1546 d_model = cfg.d_model
1547 v_bias_start = 2 * n_heads * d_head
1548 v_bias_end = 3 * n_heads * d_head
1549 b_V_only = b_V[v_bias_start:v_bias_end]
1550 if b_V_only.numel() == 0:
1551 continue
1552 b_V_reshaped = b_V_only.reshape(n_heads, d_head)
1553 W_O_reshaped = einops.rearrange(W_O, "(i h) m -> i h m", i=n_heads)
1554 folded_b_O = b_O_original + (b_V_reshaped[:, :, None] * W_O_reshaped).sum(
1555 [0, 1]
1556 )
1557 new_b_V = b_V.clone()
1558 new_b_V[v_bias_start:v_bias_end] = 0
1559 state_dict[b_V_key] = ProcessWeights.convert_tensor_to_hf_format(
1560 b_V_key, new_b_V, cfg, adapter, layer
1561 )
1562 elif is_split_format and len(b_V.shape) == 1 and len(W_O.shape) == 3: 1562 ↛ 1564line 1562 didn't jump to line 1564 because the condition on line 1562 was never true
1563 # Split bias [n_heads * d_head] with W_O already in TL format [n_heads, d_head, d_model]
1564 n_heads = cfg.n_heads
1565 d_head = cfg.d_head
1566 b_V_reshaped = b_V.reshape(n_heads, d_head)
1567 if getattr(cfg, "n_key_value_heads", None) is not None:
1568 b_V_reshaped = torch.repeat_interleave(
1569 b_V_reshaped, dim=0, repeats=cfg.n_heads // cfg.n_key_value_heads
1570 )
1571 folded_b_O = b_O_original + (b_V_reshaped[:, :, None] * W_O).sum([0, 1])
1572 state_dict[b_V_key] = ProcessWeights.convert_tensor_to_hf_format(
1573 b_V_key, torch.zeros_like(b_V), cfg, adapter, layer
1574 )
1575 elif len(b_V.shape) == 2 and len(W_O.shape) == 3: 1575 ↛ 1589line 1575 didn't jump to line 1589 because the condition on line 1575 was always true
1576 b_V_original_shape = b_V.shape
1577 if getattr(cfg, "n_key_value_heads", None) is not None:
1578 b_V = torch.repeat_interleave(
1579 b_V, dim=0, repeats=cfg.n_heads // cfg.n_key_value_heads
1580 )
1581 folded_b_O = b_O_original + (b_V[:, :, None] * W_O).sum([0, 1])
1582 state_dict[b_V_key] = ProcessWeights.convert_tensor_to_hf_format(
1583 b_V_key,
1584 torch.zeros(b_V_original_shape, dtype=b_V.dtype, device=b_V.device),
1585 cfg,
1586 adapter,
1587 layer,
1588 )
1589 elif len(b_V.shape) == 2 and len(W_O.shape) == 2:
1590 n_heads = cfg.n_heads
1591 d_head = cfg.d_head
1592 d_model = cfg.d_model
1593 b_V_original_shape = b_V.shape
1595 # Handle split QKV format where bias might be [1, d_model] or [n_heads, d_head]
1596 is_split_format = ".attn.v.bias" in b_V_key or ".attn.k.bias" in b_V_key
1597 if is_split_format and b_V.shape[0] == 1 and b_V.shape[1] == n_heads * d_head:
1598 # Reshape [1, n_heads * d_head] to [n_heads, d_head]
1599 b_V = b_V.reshape(n_heads, d_head)
1600 elif b_V.shape != (n_heads, d_head):
1601 # If not already [n_heads, d_head], try to reshape
1602 if b_V.numel() == n_heads * d_head:
1603 b_V = b_V.reshape(n_heads, d_head)
1605 if getattr(cfg, "n_key_value_heads", None) is not None:
1606 b_V = torch.repeat_interleave(
1607 b_V, dim=0, repeats=cfg.n_heads // cfg.n_key_value_heads
1608 )
1610 W_O_reshaped = einops.rearrange(W_O, "(i h) m -> i h m", i=n_heads)
1611 folded_b_O = b_O_original + (b_V[:, :, None] * W_O_reshaped).sum([0, 1])
1612 state_dict[b_V_key] = ProcessWeights.convert_tensor_to_hf_format(
1613 b_V_key,
1614 torch.zeros(b_V_original_shape, dtype=b_V.dtype, device=b_V.device),
1615 cfg,
1616 adapter,
1617 layer,
1618 )
1619 else:
1620 raise ValueError(f"Unexpected tensor shapes: b_V {b_V.shape}, W_O {W_O.shape}")
1621 state_dict[b_O_key] = ProcessWeights.convert_tensor_to_hf_format(
1622 b_O_key, folded_b_O, cfg, adapter, layer
1623 )
1624 return state_dict
1626 @staticmethod
1627 def process_weights(
1628 state_dict: Dict[str, torch.Tensor],
1629 cfg,
1630 fold_ln: bool = True,
1631 center_writing_weights: bool = True,
1632 center_unembed: bool = True,
1633 fold_value_biases: bool = True,
1634 refactor_factored_attn_matrices: bool = False,
1635 adapter=None,
1636 ) -> Dict[str, torch.Tensor]:
1637 """Apply all weight processing transformations in the correct order.
1639 This is a convenience function that applies all the weight processing steps
1640 in the same order as the legacy HookedTransformer load path.
1642 Args:
1643 state_dict (Dict[str, torch.Tensor]): State dict of the model.
1644 cfg: Model configuration object.
1645 fold_ln (bool): Whether to fold LayerNorm weights into subsequent layers.
1646 center_writing_weights (bool): Whether to center weights writing to residual stream.
1647 center_unembed (bool): Whether to center unembedding weights.
1648 fold_value_biases (bool): Whether to fold value biases into output bias.
1649 refactor_factored_attn_matrices (bool): Whether to refactor attention matrices.
1650 adapter: Optional architecture adapter for parameter key translation.
1652 Returns:
1653 Dict[str, torch.Tensor]: Fully processed state dict.
1654 """
1655 # Upcast to float32 for weight processing to avoid precision loss in
1656 # reduced-precision dtypes (bfloat16, float16). Operations like LayerNorm
1657 # folding involve multiplications that accumulate rounding errors when
1658 # performed in low precision.
1659 original_dtypes: Dict[str, torch.dtype] = {}
1660 original_keys = set(state_dict)
1661 for k, v in state_dict.items():
1662 if isinstance(v, torch.Tensor) and v.is_floating_point() and v.dtype != torch.float32:
1663 original_dtypes[k] = v.dtype
1664 state_dict[k] = v.float()
1666 # Skip fold_ln for adapters that don't support it (e.g., post-LN architectures
1667 # like BERT where LN placement means folding goes into the wrong sublayer).
1668 if fold_ln and adapter and not getattr(adapter, "supports_fold_ln", True):
1669 import warnings
1671 warnings.warn(
1672 f"{type(adapter).__name__} does not support fold_ln; norm weights "
1673 "stay unfolded and analyses assuming folded LN will be off.",
1674 UserWarning,
1675 )
1676 fold_ln = False
1677 if fold_ln:
1678 if getattr(cfg, "normalization_type", "LN") in ["LN", "LNPre"]:
1679 state_dict = ProcessWeights.fold_layer_norm(
1680 state_dict, cfg, fold_biases=True, center_weights=True, adapter=adapter
1681 )
1682 elif getattr(cfg, "normalization_type", "LN") in ["RMS", "RMSPre"]: 1682 ↛ 1693line 1682 didn't jump to line 1693 because the condition on line 1682 was always true
1683 state_dict = ProcessWeights.fold_layer_norm(
1684 state_dict, cfg, fold_biases=False, center_weights=False, adapter=adapter
1685 )
1686 # Note: Each folding function (_fold_layer for attention, _fold_mlp_layer_norm
1687 # for MLP) sets its own LN weights to 1.0 after successful folding.
1688 # We must NOT unconditionally set all LN weights to 1.0 here, because
1689 # models with combined QKV projections (e.g., OpenELM's qkv_proj) may
1690 # not be able to fold attention LN — setting ln1.w=1.0 without folding
1691 # destroys the RMS scaling.
1692 # Some adapters (e.g., post-LN) don't support center_writing_weights.
1693 if ( 1693 ↛ 1698line 1693 didn't jump to line 1698 because the condition on line 1693 was never true
1694 center_writing_weights
1695 and adapter
1696 and not getattr(adapter, "supports_center_writing_weights", True)
1697 ):
1698 import warnings
1700 warnings.warn(
1701 f"{type(adapter).__name__} does not support center_writing_weights; "
1702 "writing weights stay uncentered.",
1703 UserWarning,
1704 )
1705 center_writing_weights = False
1706 should_center_writing_weights = (
1707 center_writing_weights
1708 and getattr(cfg, "normalization_type", "LN") in ["LN", "LNPre"]
1709 and not getattr(cfg, "final_rms", False)
1710 )
1711 if should_center_writing_weights:
1712 state_dict = ProcessWeights.center_writing_weights(state_dict, cfg, adapter=adapter)
1713 if center_unembed:
1714 state_dict = ProcessWeights.center_unembed(state_dict, cfg=cfg, adapter=adapter)
1715 if fold_value_biases:
1716 state_dict = ProcessWeights.fold_value_biases(state_dict, cfg, adapter=adapter)
1717 if should_center_writing_weights:
1718 for layer_idx in range(cfg.n_layers):
1719 b_O_key = ProcessWeights._get_param_key(f"blocks.{layer_idx}.attn.b_O", adapter)
1720 if b_O_key in state_dict:
1721 b_O = ProcessWeights.convert_tensor_to_tl_format(
1722 b_O_key, state_dict, state_dict.get(b_O_key), cfg, adapter, layer_idx
1723 )
1724 assert b_O is not None, f"Attention b_O not found at key {b_O_key}"
1725 b_O = b_O - b_O.mean()
1726 state_dict[b_O_key] = ProcessWeights.convert_tensor_to_hf_format(
1727 b_O_key, b_O, cfg, adapter, layer_idx
1728 )
1729 if refactor_factored_attn_matrices:
1730 state_dict = ProcessWeights.refactor_factored_attn_matrices(
1731 state_dict, cfg, adapter=adapter
1732 )
1734 # Downcast back to original dtypes; biases created by folding take their weight's dtype
1735 for k in state_dict.keys() - original_keys:
1736 weight_dtype = original_dtypes.get(k.removesuffix(".bias") + ".weight")
1737 if weight_dtype is not None:
1738 original_dtypes[k] = weight_dtype
1739 for k, orig_dtype in original_dtypes.items():
1740 if k in state_dict and isinstance(state_dict[k], torch.Tensor): 1740 ↛ 1739line 1740 didn't jump to line 1739 because the condition on line 1740 was always true
1741 state_dict[k] = state_dict[k].to(orig_dtype)
1743 return state_dict
1745 @staticmethod
1746 def refactor_factored_attn_matrices(
1747 state_dict: Dict[str, torch.Tensor], cfg, adapter=None
1748 ) -> Dict[str, torch.Tensor]:
1749 """Experimental method for managing queries, keys and values.
1751 As argued in [A Mathematical Framework for Transformer
1752 Circuits](https://transformer-circuits.pub/2021/framework/index.html), queries, keys and
1753 values are somewhat arbitrary intermediate terms when computing with the low rank factored
1754 matrices W_QK = W_Q @ W_K.T and W_OV = W_V @ W_O, and these matrices are the only thing
1755 determining head behaviour. But there are many ways to find a low rank factorization to a
1756 given matrix, and hopefully some of these are more interpretable than others! This method is
1757 one attempt, which makes all of the matrices have orthogonal rows or columns, W_O into a
1758 rotation and W_Q and W_K having the nth column in each having the same norm. The formula is
1759 $W_V = U @ S,W_O=Vh.T,W_Q=U@S.sqrt(),W_K=Vh@S.sqrt()$.
1761 More details:
1763 If W_OV = U @ S @ Vh.T in its singular value decomposition, (where S is in R^d_head not
1764 R^d_model, as W_OV is low rank), W_OV = (U @ S) @ (Vh.T) is an equivalent low rank
1765 factorisation, where rows/columns of each matrix are orthogonal! So setting $W_V=US$ and
1766 $W_O=Vh.T$ works just as well. I *think* this is a more interpretable setup, because now
1767 $W_O$ is just a rotation, and doesn't change the norm, so $z$ has the same norm as the
1768 result of the head.
1770 For $W_QK = W_Q @ W_K.T$ we use the refactor $W_Q = U @ S.sqrt()$ and $W_K = Vh @ S.sqrt()$,
1771 which is also equivalent ($S==S.sqrt() @ S.sqrt()$ as $S$ is diagonal). Here we keep the
1772 matrices as having the same norm, since there's not an obvious asymmetry between the keys
1773 and queries.
1775 Biases are more fiddly to deal with. For OV it's pretty easy - we just need (x @ W_V + b_V)
1776 @ 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
1777 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
1778 the head_index dimension too).
1780 For QK it's messy - we need to preserve the bilinear form of (x @ W_Q + b_Q) * (y @ W_K +
1781 b_K), which is fairly messy. To deal with the biases, we concatenate them to W_Q and W_K to
1782 simulate a d_model+1 dimensional input (whose final coordinate is always 1), do the SVD
1783 factorization on this effective matrix, then separate out into final weights and biases.
1785 Args:
1786 state_dict (Dict[str, torch.Tensor]): State dict of the model.
1787 cfg: Model configuration object.
1788 adapter: Optional architecture adapter for parameter key translation.
1790 Returns:
1791 Dict[str, torch.Tensor]: Modified state dict with refactored attention matrices.
1792 """
1793 # Make a deep copy to avoid modifying the original
1794 state_dict = {
1795 k: v.clone() if isinstance(v, torch.Tensor) else v for k, v in state_dict.items()
1796 }
1797 assert (
1798 getattr(cfg, "positional_embedding_type", "standard") != "rotary"
1799 ), "You can't refactor the QK circuit when using rotary embeddings (as the QK matrix depends on the position of the query and key)"
1801 for l in range(cfg.n_layers):
1802 W_Q_key = ProcessWeights._get_param_key(f"blocks.{l}.attn.W_Q", adapter)
1803 b_Q_key = ProcessWeights._get_param_key(f"blocks.{l}.attn.b_Q", adapter)
1804 W_K_key = ProcessWeights._get_param_key(f"blocks.{l}.attn.W_K", adapter)
1805 b_K_key = ProcessWeights._get_param_key(f"blocks.{l}.attn.b_K", adapter)
1806 W_V_key = ProcessWeights._get_param_key(f"blocks.{l}.attn.W_V", adapter)
1807 W_O_key = ProcessWeights._get_param_key(f"blocks.{l}.attn.W_O", adapter)
1808 b_V_key = ProcessWeights._get_param_key(f"blocks.{l}.attn.b_V", adapter)
1809 b_O_key = ProcessWeights._get_param_key(f"blocks.{l}.attn.b_O", adapter)
1811 # Skip hybrid layers without attention (other loops already guard individually)
1812 if W_Q_key not in state_dict:
1813 continue
1814 # If Q is present, K/V/O must be too
1815 for _required_key in [W_K_key, W_V_key, W_O_key]:
1816 if _required_key not in state_dict:
1817 raise ValueError(
1818 f"Inconsistent attention weights at layer {l}: "
1819 f"'{W_Q_key}' found but '{_required_key}' missing. "
1820 f"All of W_Q, W_K, W_V, W_O must be present together."
1821 )
1823 # W_QK = W_Q @ W_K.T
1824 # Concatenate biases to make a d_model+1 input dimension
1825 W_Q = ProcessWeights.convert_tensor_to_tl_format(
1826 W_Q_key, state_dict, state_dict.get(W_Q_key), cfg, adapter, l
1827 )
1828 b_Q = ProcessWeights.convert_tensor_to_tl_format(
1829 b_Q_key, state_dict, state_dict.get(b_Q_key), cfg, adapter, l
1830 )
1831 W_K = ProcessWeights.convert_tensor_to_tl_format(
1832 W_K_key, state_dict, state_dict.get(W_K_key), cfg, adapter, l
1833 )
1834 b_K = ProcessWeights.convert_tensor_to_tl_format(
1835 b_K_key, state_dict, state_dict.get(b_K_key), cfg, adapter, l
1836 )
1837 assert W_Q is not None, f"W_Q not found at key {W_Q_key}"
1838 assert b_Q is not None, f"b_Q not found at key {b_Q_key}"
1839 assert W_K is not None, f"W_K not found at key {W_K_key}"
1840 assert b_K is not None, f"b_K not found at key {b_K_key}"
1842 W_Q_eff = torch.cat([W_Q, b_Q[:, None, :]], dim=1)
1843 W_K_eff = torch.cat([W_K, b_K[:, None, :]], dim=1)
1845 W_Q_eff_even, W_K_eff_even_T = (
1846 FactoredMatrix(W_Q_eff, W_K_eff.transpose(-1, -2)).make_even().pair
1847 )
1848 W_K_eff_even = W_K_eff_even_T.transpose(-1, -2)
1850 state_dict[W_Q_key] = ProcessWeights.convert_tensor_to_hf_format(
1851 W_Q_key, W_Q_eff_even[:, :-1, :], cfg, adapter, l
1852 )
1853 state_dict[b_Q_key] = ProcessWeights.convert_tensor_to_hf_format(
1854 b_Q_key, W_Q_eff_even[:, -1, :], cfg, adapter, l
1855 )
1856 state_dict[W_K_key] = ProcessWeights.convert_tensor_to_hf_format(
1857 W_K_key, W_K_eff_even[:, :-1, :], cfg, adapter, l
1858 )
1859 state_dict[b_K_key] = ProcessWeights.convert_tensor_to_hf_format(
1860 b_K_key, W_K_eff_even[:, -1, :], cfg, adapter, l
1861 )
1863 # W_OV = W_V @ W_O
1864 W_V = ProcessWeights.convert_tensor_to_tl_format(
1865 W_V_key, state_dict, state_dict.get(W_V_key), cfg, adapter, l
1866 )
1867 W_O = ProcessWeights.convert_tensor_to_tl_format(
1868 W_O_key, state_dict, state_dict.get(W_O_key), cfg, adapter, l
1869 )
1871 # Factors the bias to be consistent.
1872 b_V = ProcessWeights.convert_tensor_to_tl_format(
1873 b_V_key, state_dict, state_dict.get(b_V_key), cfg, adapter, l
1874 )
1875 b_O = ProcessWeights.convert_tensor_to_tl_format(
1876 b_O_key, state_dict, state_dict.get(b_O_key), cfg, adapter, l
1877 )
1878 assert W_V is not None, f"W_V not found at key {W_V_key}"
1879 assert W_O is not None, f"W_O not found at key {W_O_key}"
1880 assert b_V is not None, f"b_V not found at key {b_V_key}"
1881 assert b_O is not None, f"b_O not found at key {b_O_key}"
1883 # Add singleton dimension for broadcasting
1884 b_V_expanded = einops.rearrange(b_V, "head_index d_head -> head_index d_head 1")
1886 b_V_times_W_O = b_V_expanded * W_O
1888 # Sum over d_head and head_index dimensions
1889 b_V_contribution = b_V_times_W_O.sum(1).sum(0)
1891 effective_bias = b_O + b_V_contribution
1892 state_dict[b_V_key] = ProcessWeights.convert_tensor_to_hf_format(
1893 b_V_key, torch.zeros_like(b_V), cfg, adapter, l
1894 )
1895 state_dict[b_O_key] = ProcessWeights.convert_tensor_to_hf_format(
1896 b_O_key, effective_bias, cfg, adapter, l
1897 )
1899 # Helper class to efficiently deal with low rank factored matrices.
1900 W_OV = FactoredMatrix(W_V, W_O)
1901 U, S, Vh = W_OV.svd()
1902 state_dict[W_V_key] = ProcessWeights.convert_tensor_to_hf_format(
1903 W_V_key, U @ S.diag_embed(), cfg, adapter, l
1904 )
1905 state_dict[W_O_key] = ProcessWeights.convert_tensor_to_hf_format(
1906 W_O_key, utils.transpose(Vh), cfg, adapter, l
1907 )
1909 return state_dict
1911 @overload
1912 @staticmethod
1913 def convert_tensor_to_tl_format(
1914 param_name: str,
1915 model_state_dict: Dict[str, torch.Tensor],
1916 tensor: torch.Tensor,
1917 cfg: Optional["TransformerLensConfig"],
1918 adapter: Optional["ArchitectureAdapter"] = None,
1919 layer_idx: Optional[int] = None,
1920 ) -> torch.Tensor:
1921 ...
1923 @overload
1924 @staticmethod
1925 def convert_tensor_to_tl_format(
1926 param_name: str,
1927 model_state_dict: Dict[str, torch.Tensor],
1928 tensor: None,
1929 cfg: Optional["TransformerLensConfig"],
1930 adapter: Optional["ArchitectureAdapter"] = None,
1931 layer_idx: Optional[int] = None,
1932 ) -> None:
1933 ...
1935 @staticmethod
1936 def convert_tensor_to_tl_format(
1937 param_name: str,
1938 model_state_dict: Dict[str, torch.Tensor],
1939 tensor: Optional[torch.Tensor],
1940 cfg: Optional["TransformerLensConfig"],
1941 adapter: Optional["ArchitectureAdapter"] = None,
1942 layer_idx: Optional[int] = None,
1943 ) -> Optional[torch.Tensor]:
1944 """Convert a tensor from its original format to TransformerLens format.
1946 Args:
1947 param_name: The parameter name in TransformerLens format (e.g., "blocks.0.attn.W_Q")
1948 model_state_dict: The model's state dictionary containing the actual tensors
1949 tensor: The tensor to convert, or None for optional parameters
1950 cfg: Model configuration
1951 adapter: Optional architecture adapter for component retrieval and key translation.
1952 If None, the tensor is returned unchanged.
1953 layer_idx: Layer index (required for layer-specific parameters)
1955 Returns:
1956 The tensor converted to TransformerLens format, or None if the parameter doesn't exist
1957 (which is valid for optional parameters like biases in models that don't use them).
1958 If adapter is None, returns the tensor unchanged.
1959 """
1960 # If no adapter provided, return tensor unchanged (handle None gracefully)
1961 if adapter is None:
1962 return tensor
1964 if ( 1964 ↛ 2045line 1964 didn't jump to line 2045 because the condition on line 1964 was always true
1965 hasattr(adapter, "weight_processing_conversions")
1966 and adapter.weight_processing_conversions is not None
1967 ):
1968 # Create placeholder param name by replacing layer index with {i}
1969 placeholder_param_name = param_name
1970 if "blocks." in param_name:
1971 placeholder_param_name = re.sub(r"blocks\.(\d+)\.", "blocks.{i}.", param_name)
1973 # Check if we have a conversion for this parameter.
1974 # Try exact match first, then strip .weight suffix for adapters
1975 # that define conversions without the suffix (e.g. Pythia's "blocks.{i}.attn.q").
1976 # NOTE: Only strip .weight, NOT .bias — stripping .bias would incorrectly
1977 # match bias keys against weight conversions (e.g. "blocks.{i}.attn.q.bias"
1978 # would match the weight conversion for "blocks.{i}.attn.q").
1979 matched_key = None
1980 if placeholder_param_name in adapter.weight_processing_conversions:
1981 matched_key = placeholder_param_name
1982 elif placeholder_param_name.endswith(".weight"):
1983 stripped = placeholder_param_name[: -len(".weight")]
1984 if stripped in adapter.weight_processing_conversions:
1985 matched_key = stripped
1987 if matched_key is not None:
1988 param_conversion = adapter.weight_processing_conversions[matched_key]
1990 # Handle both ParamProcessingConversion objects and legacy string mappings
1991 if isinstance(param_conversion, str): 1991 ↛ 1994line 1991 didn't jump to line 1994 because the condition on line 1991 was never true
1992 # Legacy string mapping - just return the tensor as-is
1993 # (string mappings are handled elsewhere in the architecture adapter)
1994 return tensor
1995 else:
1996 # Skip conversion for optional parameters that don't exist (e.g. biases)
1997 if tensor is None and param_name not in model_state_dict:
1998 return None
1999 # Try ParamProcessingConversion.convert() first (uses source_key
2000 # to fetch from state dict — needed for split conversions like
2001 # GPT-2's QKV). If source_key resolves to a missing key and we
2002 # already have the tensor, fall back to applying the tensor
2003 # conversion directly (needed for adapters like GPT-Neo whose
2004 # source_key references HF keys not in the bridge state dict).
2005 if (
2006 hasattr(param_conversion, "source_key")
2007 and param_conversion.source_key is not None
2008 ):
2009 resolved_key = param_conversion._resolve_key(
2010 param_name, param_conversion.source_key
2011 )
2012 if resolved_key not in model_state_dict and tensor is not None: 2012 ↛ 2039line 2012 didn't jump to line 2039 because the condition on line 2012 was always true
2013 # Source key not in state dict — the tensor is already in
2014 # bridge format (e.g. already split from combined QKV).
2015 # If the conversion is a ChainTensorConversion that includes
2016 # a SplitTensorConversion, skip the split step since
2017 # it was already applied during bridge construction.
2018 from transformer_lens.conversion_utils.conversion_steps.chain_tensor_conversion import (
2019 ChainTensorConversion,
2020 )
2021 from transformer_lens.conversion_utils.conversion_steps.split_tensor_conversion import (
2022 SplitTensorConversion,
2023 )
2025 tc = param_conversion.tensor_conversion
2026 if isinstance(tc, ChainTensorConversion): 2026 ↛ 2027line 2026 didn't jump to line 2027 because the condition on line 2026 was never true
2027 non_split = [
2028 c
2029 for c in tc.conversions
2030 if not isinstance(c, SplitTensorConversion)
2031 ]
2032 if len(non_split) < len(tc.conversions):
2033 # Apply only the non-split conversions
2034 result = tensor
2035 for conv in non_split:
2036 result = conv.handle_conversion(result, model_state_dict)
2037 return result
2038 return tc.convert(tensor, model_state_dict)
2039 return param_conversion.convert(model_state_dict, param_name)
2040 else:
2041 # No conversion defined, return tensor as-is (may be None for optional params)
2042 return tensor
2043 else:
2044 # No conversions defined, return tensor as-is (may be None for optional params)
2045 return tensor
2047 @overload
2048 @staticmethod
2049 def convert_tensor_to_hf_format(
2050 param_name: str,
2051 tensor: torch.Tensor,
2052 cfg: Optional["TransformerLensConfig"],
2053 adapter: Optional["ArchitectureAdapter"] = None,
2054 layer_idx: Optional[int] = None,
2055 ) -> torch.Tensor:
2056 ...
2058 @overload
2059 @staticmethod
2060 def convert_tensor_to_hf_format(
2061 param_name: str,
2062 tensor: None,
2063 cfg: Optional["TransformerLensConfig"],
2064 adapter: Optional["ArchitectureAdapter"] = None,
2065 layer_idx: Optional[int] = None,
2066 ) -> None:
2067 ...
2069 @staticmethod
2070 def convert_tensor_to_hf_format(
2071 param_name: str,
2072 tensor: Optional[torch.Tensor],
2073 cfg: Optional["TransformerLensConfig"],
2074 adapter: Optional["ArchitectureAdapter"] = None,
2075 layer_idx: Optional[int] = None,
2076 ) -> Optional[torch.Tensor]:
2077 """Convert a tensor from TransformerLens format back to its original format.
2079 Args:
2080 param_name: The parameter name in TransformerLens format (e.g., "blocks.0.attn.W_Q")
2081 tensor: The tensor to convert (in TransformerLens format), or None if parameter is optional
2082 cfg: Model configuration
2083 adapter: Optional architecture adapter for component retrieval and key translation.
2084 If None, the tensor is returned unchanged.
2085 layer_idx: Layer index (required for layer-specific parameters)
2087 Returns:
2088 The tensor converted back to original format, or None if tensor was None.
2089 If adapter is None, returns the tensor unchanged.
2090 """
2091 # Handle None tensors (optional parameters)
2092 if tensor is None: 2092 ↛ 2093line 2092 didn't jump to line 2093 because the condition on line 2092 was never true
2093 return None
2095 # If no adapter provided, return tensor unchanged
2096 if adapter is None:
2097 return tensor
2099 if ( 2099 ↛ 2151line 2099 didn't jump to line 2151 because the condition on line 2099 was always true
2100 hasattr(adapter, "weight_processing_conversions")
2101 and adapter.weight_processing_conversions is not None
2102 ):
2103 # Create placeholder param name by replacing layer index with {i}
2104 placeholder_param_name = param_name
2105 if "blocks." in param_name:
2106 placeholder_param_name = re.sub(r"blocks\.(\d+)\.", "blocks.{i}.", param_name)
2108 # Check if we have a conversion for this parameter.
2109 # Try exact match first, then strip .weight suffix (not .bias — see convert_tensor_to_tl_format).
2110 matched_key = None
2111 if placeholder_param_name in adapter.weight_processing_conversions:
2112 matched_key = placeholder_param_name
2113 elif placeholder_param_name.endswith(".weight"):
2114 stripped = placeholder_param_name[: -len(".weight")]
2115 if stripped in adapter.weight_processing_conversions:
2116 matched_key = stripped
2118 if matched_key is not None:
2119 param_conversion = adapter.weight_processing_conversions[matched_key]
2121 # Handle both ParamProcessingConversion objects and legacy string mappings
2122 if isinstance(param_conversion, str): 2122 ↛ 2124line 2122 didn't jump to line 2124 because the condition on line 2122 was never true
2123 # Legacy string mapping - just return the tensor as-is
2124 return tensor
2125 else:
2126 # Revert the conversion. For ChainTensorConversions that include
2127 # SplitTensorConversion, skip the split revert step (which is a
2128 # no-op anyway) to match the forward conversion path.
2129 from transformer_lens.conversion_utils.conversion_steps.chain_tensor_conversion import (
2130 ChainTensorConversion,
2131 )
2132 from transformer_lens.conversion_utils.conversion_steps.split_tensor_conversion import (
2133 SplitTensorConversion,
2134 )
2136 tc = param_conversion.tensor_conversion
2137 if isinstance(tc, ChainTensorConversion): 2137 ↛ 2138line 2137 didn't jump to line 2138 because the condition on line 2137 was never true
2138 non_split = [
2139 c for c in tc.conversions if not isinstance(c, SplitTensorConversion)
2140 ]
2141 if len(non_split) < len(tc.conversions):
2142 # Revert only the non-split conversions in reverse order
2143 result = tensor
2144 for conv in reversed(non_split):
2145 result = conv.revert(result)
2146 return result
2147 return param_conversion.revert(tensor)
2148 else:
2149 return tensor
2150 else:
2151 return tensor
2153 @staticmethod
2154 def distribute_weights_to_components(
2155 state_dict: Dict[str, torch.Tensor],
2156 component_mapping: Dict[str, Any],
2157 verbose: bool = False,
2158 ) -> None:
2159 """Distribute processed weights from state_dict to generalized components.
2161 This function loops through the component_mapping and extracts relevant weights
2162 for each component using filter_dict_by_prefix, then calls set_processed_weights
2163 on each component. For list components (like blocks), it determines the number
2164 of items and distributes weights to each indexed component.
2166 Args:
2167 state_dict: Dictionary of processed weights in MODERN TransformerLens format
2168 (e.g., blocks.0.attn.q.weight, not transformer.h.0.attn.q.weight)
2169 component_mapping: Dictionary (real_components) mapping TL keys to tuples of
2170 (remote_path, component_instance), where component_instance can be either
2171 a single component or a list of components
2172 verbose: If True, print detailed information about weight distribution
2174 Example:
2175 For a real_components mapping like:
2176 {
2177 "embed": ("transformer.wte", <EmbeddingBridge instance>),
2178 "blocks": ("transformer.h", [<BlockBridge 0>, <BlockBridge 1>, ...]),
2179 "unembed": ("lm_head", <UnembeddingBridge instance>)
2180 }
2182 With modern TL keys in state_dict like "embed.weight", "blocks.0.attn.q.weight":
2183 1. Extract weights starting with "embed" and pass to embed component
2184 2. For blocks, extract all "blocks.*" weights, determine the number of blocks,
2185 then for each block index, extract weights for that specific block
2186 3. Extract "unembed" weights and pass to unembed component
2187 """
2188 if verbose: 2188 ↛ 2189line 2188 didn't jump to line 2189 because the condition on line 2188 was never true
2189 print(f"\n{'='*80}")
2190 print(f"distribute_weights_to_components: Starting weight distribution")
2191 print(f"State dict has {len(state_dict)} keys")
2192 print(f"Component mapping has {len(component_mapping)} components")
2193 print(f"{'='*80}\n")
2195 for component_name, component_tuple in component_mapping.items():
2196 # component_mapping is real_components format: (remote_path, instance)
2197 # instance can be either a single component or a list of components
2198 if not isinstance(component_tuple, tuple): 2198 ↛ 2199line 2198 didn't jump to line 2199 because the condition on line 2198 was never true
2199 raise ValueError(
2200 f"Expected tuple for component '{component_name}' in real_components, "
2201 f"but got {type(component_tuple).__name__}: {component_tuple}"
2202 )
2203 remote_key, component = component_tuple
2204 is_list = isinstance(component, list)
2206 # Use the component_name (TL format) as prefix instead of remote_key (HF format)
2207 # since state_dict now has modern TL keys
2208 tl_prefix = component_name
2210 if verbose: 2210 ↛ 2211line 2210 didn't jump to line 2211 because the condition on line 2210 was never true
2211 print(f"\nProcessing component: {component_name}")
2212 print(f" Remote key (HF): {remote_key}")
2213 print(f" TL prefix: {tl_prefix}")
2214 print(f" Is list: {is_list}")
2216 if is_list:
2217 # This is a list component like "blocks"
2218 # Extract all weights that start with this prefix
2219 all_list_weights = filter_dict_by_prefix(state_dict, tl_prefix)
2221 if verbose: 2221 ↛ 2222line 2221 didn't jump to line 2222 because the condition on line 2221 was never true
2222 print(f" Found {len(all_list_weights)} weights for list component")
2223 print(f" List has {len(component)} instances")
2225 # Component is a list of actual instances
2226 for i, instance in enumerate(component):
2227 # Extract weights for this specific index
2228 # This will get keys like "0.attn.q.weight" and strip the "0." to get "attn.q.weight"
2229 indexed_weights = filter_dict_by_prefix(all_list_weights, str(i))
2231 if verbose: 2231 ↛ 2232line 2231 didn't jump to line 2232 because the condition on line 2231 was never true
2232 print(f" Instance {i}: Found {len(indexed_weights)} weights")
2233 for key in indexed_weights.keys():
2234 print(f" - {key}")
2236 # Skip if no weights found for this component (e.g., Q/K/V Linear sub-components
2237 # that get their weights from parent JointQKVAttentionBridge)
2238 if len(indexed_weights) == 0: 2238 ↛ 2239line 2238 didn't jump to line 2239 because the condition on line 2238 was never true
2239 if verbose:
2240 print(f" Skipping instance {i} - no weights found")
2241 continue
2243 instance.set_processed_weights(indexed_weights, verbose=verbose)
2244 else:
2245 # This is a single component (not a list)
2246 component_weights = filter_dict_by_prefix(state_dict, tl_prefix)
2248 if verbose: 2248 ↛ 2249line 2248 didn't jump to line 2249 because the condition on line 2248 was never true
2249 print(f" Found {len(component_weights)} weights for single component")
2250 for key in component_weights.keys():
2251 print(f" - {key}")
2253 # Skip if no weights found for this component
2254 if len(component_weights) == 0:
2255 if verbose: 2255 ↛ 2256line 2255 didn't jump to line 2256 because the condition on line 2255 was never true
2256 print(f" Skipping component - no weights found")
2257 continue
2259 component.set_processed_weights(component_weights, verbose=verbose)