Coverage for transformer_lens/weight_processing.py: 77%
844 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-09-01 16:23 +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: 346 ↛ 347line 346 didn't jump to line 347 because the condition on line 346 was never true
347 W_Q_key = W_Q_key.replace(".W_Q", "._W_Q")
348 if W_K_key not in state_dict: 348 ↛ 349line 348 didn't jump to line 349 because the condition on line 348 was never true
349 W_K_key = W_K_key.replace(".W_K", "._W_K")
350 if W_V_key not in state_dict: 350 ↛ 351line 350 didn't jump to line 351 because the condition on line 350 was never true
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_layer(
424 state_dict: Dict[str, torch.Tensor],
425 cfg,
426 layer_idx: int,
427 fold_biases: bool,
428 center_weights: bool,
429 adapter,
430 gqa: str,
431 ) -> Dict[str, torch.Tensor]:
432 """Fold LayerNorm for a single layer.
434 Args:
435 state_dict: The state dictionary to process (modified in place)
436 cfg: Model configuration object
437 layer_idx: The layer index to process
438 fold_biases: Whether to fold LayerNorm biases
439 center_weights: Whether to center weights after folding
440 adapter: Optional architecture adapter for parameter key translation
441 gqa: GQA prefix string (empty or "_")
442 """
443 layer = layer_idx
444 tensors = ProcessWeights.extract_attention_tensors_for_folding(
445 state_dict, cfg, layer, adapter
446 )
447 wq_tensor = tensors["wq"]
448 wk_tensor = tensors["wk"]
449 wv_tensor = tensors["wv"]
450 bq_tensor = tensors["bq"]
451 bk_tensor = tensors["bk"]
452 bv_tensor = tensors["bv"]
453 ln1_b = tensors["ln1_b"]
454 ln1_w = tensors["ln1_w"]
455 keys = tensors["keys"]
457 # Fold LN into QKV (skip if combined QKV, e.g., OpenELM)
458 if wq_tensor is not None: 458 ↛ 533line 458 didn't jump to line 533 because the condition on line 458 was always true
459 assert isinstance(wq_tensor, torch.Tensor)
460 assert isinstance(keys, dict)
461 if wk_tensor is not None: 461 ↛ 463line 461 didn't jump to line 463 because the condition on line 461 was always true
462 assert isinstance(wk_tensor, torch.Tensor)
463 if wv_tensor is not None: 463 ↛ 465line 463 didn't jump to line 465 because the condition on line 463 was always true
464 assert isinstance(wv_tensor, torch.Tensor)
465 if bq_tensor is not None:
466 assert isinstance(bq_tensor, torch.Tensor)
467 if bk_tensor is not None:
468 assert isinstance(bk_tensor, torch.Tensor)
469 if bv_tensor is not None:
470 assert isinstance(bv_tensor, torch.Tensor)
471 # RMS norm (Gemma): ln1_b may be None, only ln1_w required
472 if ln1_w is not None:
473 assert isinstance(ln1_w, torch.Tensor)
474 # Fold biases if present (RMS norm has none; missing QKV biases get zeros)
475 if fold_biases and ln1_b is not None:
476 assert isinstance(ln1_b, torch.Tensor)
477 assert wq_tensor is not None
478 assert wk_tensor is not None
479 assert wv_tensor is not None
480 bq_tensor, bk_tensor, bv_tensor = ProcessWeights.fold_layer_norm_biases(
481 wq_tensor, wk_tensor, wv_tensor, bq_tensor, bk_tensor, bv_tensor, ln1_b
482 )
483 if keys["ln1_b"] in state_dict: 483 ↛ 485line 483 didn't jump to line 485 because the condition on line 483 was always true
484 state_dict[keys["ln1_b"]] = torch.zeros_like(ln1_b)
485 alternate_b_key = (
486 keys["ln1_b"].replace("ln_1", "ln1")
487 if "ln_1" in keys["ln1_b"]
488 else keys["ln1_b"].replace("ln1", "ln_1")
489 )
490 if alternate_b_key != keys["ln1_b"] and alternate_b_key in state_dict: 490 ↛ 491line 490 didn't jump to line 491 because the condition on line 490 was never true
491 state_dict[alternate_b_key] = torch.zeros_like(ln1_b)
492 # Fold ln1_w; use (1+w) for rmsnorm_uses_offset (Gemma), then set to identity
493 rmsnorm_uses_offset = getattr(cfg, "rmsnorm_uses_offset", False)
494 effective_ln1_w = (1.0 + ln1_w) if rmsnorm_uses_offset else ln1_w
495 if wk_tensor is not None and wv_tensor is not None: 495 ↛ 500line 495 didn't jump to line 500 because the condition on line 495 was always true
496 wq_tensor, wk_tensor, wv_tensor = ProcessWeights.fold_layer_norm_weights(
497 wq_tensor, wk_tensor, wv_tensor, effective_ln1_w
498 )
499 # Set ln1.w to identity: ones (standard) or zeros (rmsnorm_uses_offset)
500 identity_val = (
501 torch.zeros_like(ln1_w) if rmsnorm_uses_offset else torch.ones_like(ln1_w)
502 )
503 if keys["ln1_w"] in state_dict: 503 ↛ 505line 503 didn't jump to line 505 because the condition on line 503 was always true
504 state_dict[keys["ln1_w"]] = identity_val
505 alternate_w_key = (
506 keys["ln1_w"].replace("ln_1", "ln1")
507 if "ln_1" in keys["ln1_w"]
508 else keys["ln1_w"].replace("ln1", "ln_1")
509 )
510 if alternate_w_key != keys["ln1_w"] and alternate_w_key in state_dict: 510 ↛ 511line 510 didn't jump to line 511 because the condition on line 510 was never true
511 state_dict[alternate_w_key] = identity_val
512 if center_weights and wk_tensor is not None and (wv_tensor is not None):
513 wq_tensor, wk_tensor, wv_tensor = ProcessWeights.center_attention_weights(
514 wq_tensor, wk_tensor, wv_tensor
515 )
516 state_dict = ProcessWeights._store_processed_attention_tensors(
517 state_dict,
518 keys,
519 wq_tensor,
520 wk_tensor,
521 wv_tensor,
522 bq_tensor,
523 bk_tensor,
524 bv_tensor,
525 adapter,
526 cfg,
527 layer,
528 )
530 # ln1_post.w (Gemma 2/3): keep original; independent post-attention normalization
532 # Fold MLP LN: shared ln1 (Phi-2, GPT-J) or separate ln2 (Pythia)
533 if getattr(cfg, "parallel_attn_mlp", False) and ln1_w is not None:
534 # Check if a separate ln2 exists for this layer
535 ln2_check_key = ProcessWeights._resolve_state_dict_key(
536 state_dict,
537 ProcessWeights._get_param_key(f"blocks.{layer_idx}.ln2.w", adapter),
538 layer_idx,
539 )
540 if ln2_check_key in state_dict: 540 ↛ 547line 540 didn't jump to line 547 because the condition on line 540 was always true
541 # Separate ln2 (e.g., GPT-NeoX/Pythia) — fold ln2 → MLP normally
542 state_dict = ProcessWeights._fold_mlp_layer_norm(
543 state_dict, cfg, layer, fold_biases, center_weights, adapter
544 )
545 else:
546 # Shared ln1 (e.g., Phi-2, GPT-J) — fold ln1 → MLP via override
547 assert isinstance(ln1_w, torch.Tensor)
548 assert ln1_b is None or isinstance(ln1_b, torch.Tensor)
549 state_dict = ProcessWeights._fold_mlp_layer_norm(
550 state_dict,
551 cfg,
552 layer,
553 fold_biases,
554 center_weights,
555 adapter,
556 override_ln_w=ln1_w,
557 override_ln_b=ln1_b,
558 )
559 else:
560 state_dict = ProcessWeights._fold_mlp_layer_norm(
561 state_dict, cfg, layer, fold_biases, center_weights, adapter
562 )
564 return state_dict
566 @staticmethod
567 def _fold_mlp_layer_norm(
568 state_dict: Dict[str, torch.Tensor],
569 cfg,
570 layer: int,
571 fold_biases: bool,
572 center_weights: bool,
573 adapter,
574 override_ln_w: Optional[torch.Tensor] = None,
575 override_ln_b: Optional[torch.Tensor] = None,
576 ) -> Dict[str, torch.Tensor]:
577 """Fold LayerNorm into MLP layer.
579 Args:
580 state_dict: The state dictionary to process (modified in place)
581 cfg: Model configuration object
582 layer: The layer index to process
583 fold_biases: Whether to fold LayerNorm biases
584 center_weights: Whether to center weights after folding
585 adapter: Optional architecture adapter for parameter key translation
586 override_ln_w: Override LN weight tensor. Used for parallel architectures
587 where MLP reads from ln1 (same as attention) instead of a separate ln2.
588 override_ln_b: Override LN bias tensor. Used with override_ln_w.
589 """
590 if getattr(cfg, "attn_only", False):
591 return state_dict
593 mlp_b_in_key = ProcessWeights._resolve_state_dict_key(
594 state_dict, ProcessWeights._get_param_key(f"blocks.{layer}.mlp.b_in", adapter), layer
595 )
596 mlp_W_in_key = ProcessWeights._resolve_state_dict_key(
597 state_dict, ProcessWeights._get_param_key(f"blocks.{layer}.mlp.W_in", adapter), layer
598 )
599 mlp_W_gate_key = (
600 ProcessWeights._resolve_state_dict_key(
601 state_dict,
602 ProcessWeights._get_param_key(f"blocks.{layer}.mlp.W_gate", adapter),
603 layer,
604 )
605 if getattr(cfg, "gated_mlp", False)
606 else None
607 )
608 mlp_b_gate_key = (
609 ProcessWeights._resolve_state_dict_key(
610 state_dict,
611 ProcessWeights._get_param_key(f"blocks.{layer}.mlp.b_gate", adapter),
612 layer,
613 )
614 if getattr(cfg, "gated_mlp", False)
615 else None
616 )
618 # For parallel architectures, ln1 values are passed via override params.
619 # Otherwise, look up ln2 from the state dict.
620 ln2_w: Optional[torch.Tensor]
621 ln2_b: Optional[torch.Tensor]
622 if override_ln_w is not None: 622 ↛ 623line 622 didn't jump to line 623 because the condition on line 622 was never true
623 ln2_w = override_ln_w
624 ln2_b = override_ln_b
625 ln2_w_key = None # No state dict key to zero out (already done by attention folding)
626 ln2_b_key = None
627 has_ln = True
628 else:
629 ln2_b_key = ProcessWeights._resolve_state_dict_key(
630 state_dict, ProcessWeights._get_param_key(f"blocks.{layer}.ln2.b", adapter), layer
631 )
632 ln2_w_key = ProcessWeights._resolve_state_dict_key(
633 state_dict, ProcessWeights._get_param_key(f"blocks.{layer}.ln2.w", adapter), layer
634 )
635 has_ln = ln2_w_key in state_dict
636 ln2_w = state_dict.get(ln2_w_key) if has_ln else None
637 ln2_b = state_dict.get(ln2_b_key) if has_ln else None
639 # CRITICAL FIX: For RMS norm (Gemma), ln2_b doesn't exist. Only require ln2_w!
640 if has_ln and ln2_w is not None:
641 # MoE layers: fold ln2 into router gate and each expert's W_in/W_gate
642 if getattr(cfg, "num_experts", None) is not None and cfg.num_experts > 0:
643 # MoE: fold into router + experts; skip identity if wrapped
644 expert_fold_count = 0
645 expected_expert_folds = cfg.num_experts * 2 # W_in + W_gate per expert
647 # Fold into router gate
648 router_key = ProcessWeights._resolve_state_dict_key(
649 state_dict, f"blocks.{layer}.mlp.W_gate.weight", layer
650 )
651 if router_key in state_dict:
652 state_dict[router_key] = state_dict[router_key] * ln2_w[None, :]
653 # Fold into each expert's W_in and W_gate (SwiGLU gate)
654 for e in range(cfg.num_experts):
655 for suffix in ("W_in.weight", "W_gate.weight"):
656 key = ProcessWeights._resolve_state_dict_key(
657 state_dict,
658 f"blocks.{layer}.mlp.experts.{e}.{suffix}",
659 layer,
660 )
661 if key in state_dict:
662 state_dict[key] = state_dict[key] * ln2_w[None, :]
663 expert_fold_count += 1
665 # Only set ln2 to identity if we actually folded into expert weights.
666 if expert_fold_count > 0:
667 if ln2_w_key is not None: 667 ↛ 680line 667 didn't jump to line 680 because the condition on line 667 was always true
668 state_dict[ln2_w_key] = torch.ones_like(ln2_w)
669 alternate_ln2_w_key = (
670 ln2_w_key.replace("ln_2", "ln2")
671 if "ln_2" in ln2_w_key
672 else ln2_w_key.replace("ln2", "ln_2")
673 )
674 if alternate_ln2_w_key != ln2_w_key and alternate_ln2_w_key in state_dict: 674 ↛ 675line 674 didn't jump to line 675 because the condition on line 674 was never true
675 state_dict[alternate_ln2_w_key] = torch.ones_like(ln2_w)
676 else:
677 # No expert weights found — undo router gate fold for consistency.
678 if router_key in state_dict: 678 ↛ 679line 678 didn't jump to line 679 because the condition on line 678 was never true
679 state_dict[router_key] = state_dict[router_key] / ln2_w[None, :]
680 return state_dict
682 mlp_W_in = ProcessWeights.convert_tensor_to_tl_format(
683 mlp_W_in_key, state_dict, state_dict.get(mlp_W_in_key), cfg, adapter, layer
684 )
685 assert mlp_W_in is not None, f"MLP W_in not found at key {mlp_W_in_key}"
686 # rmsnorm_uses_offset: effective scale is (1+w), identity is 0.0
687 rmsnorm_uses_offset = getattr(cfg, "rmsnorm_uses_offset", False)
688 effective_ln2_w = (1.0 + ln2_w) if rmsnorm_uses_offset else ln2_w
689 if mlp_W_in.shape[1] == effective_ln2_w.shape[0]:
690 ln2_w_broadcast = effective_ln2_w[None, :]
691 sum_dim = -1
692 if ln2_b is not None:
693 ln2_b_broadcast = ln2_b[None, :]
694 elif mlp_W_in.shape[0] == effective_ln2_w.shape[0]: 694 ↛ 700line 694 didn't jump to line 700 because the condition on line 694 was always true
695 ln2_w_broadcast = effective_ln2_w[:, None]
696 sum_dim = -2
697 if ln2_b is not None:
698 ln2_b_broadcast = ln2_b[:, None]
699 else:
700 raise ValueError(
701 f"Cannot broadcast MLP weight {mlp_W_in.shape} with layer norm weight {effective_ln2_w.shape}"
702 )
703 # Only fold biases if they exist (LayerNorm). RMS norm has no biases.
704 if fold_biases and ln2_b is not None:
705 mlp_b_in = ProcessWeights.convert_tensor_to_tl_format(
706 mlp_b_in_key, state_dict, state_dict.get(mlp_b_in_key), cfg, adapter, layer
707 )
708 ln2_b_folded = (mlp_W_in * ln2_b_broadcast).sum(sum_dim)
709 if mlp_b_in is not None: 709 ↛ 713line 709 didn't jump to line 713 because the condition on line 709 was always true
710 new_mlp_b_in = mlp_b_in + ln2_b_folded
711 else:
712 # MLP has no bias — create one from the folded LN bias
713 new_mlp_b_in = ln2_b_folded
714 state_dict[mlp_b_in_key] = ProcessWeights.convert_tensor_to_hf_format(
715 mlp_b_in_key, new_mlp_b_in, cfg, adapter, layer
716 )
717 # Set ln2.b to zero (skip for parallel override — ln1 already zeroed)
718 if ln2_b_key is not None: 718 ↛ 727line 718 didn't jump to line 727 because the condition on line 718 was always true
719 state_dict[ln2_b_key] = torch.zeros_like(ln2_b)
720 alternate_ln2_b_key = (
721 ln2_b_key.replace("ln_2", "ln2")
722 if "ln_2" in ln2_b_key
723 else ln2_b_key.replace("ln2", "ln_2")
724 )
725 if alternate_ln2_b_key != ln2_b_key and alternate_ln2_b_key in state_dict: 725 ↛ 726line 725 didn't jump to line 726 because the condition on line 725 was never true
726 state_dict[alternate_ln2_b_key] = torch.zeros_like(ln2_b)
727 new_mlp_W_in = mlp_W_in * ln2_w_broadcast
728 state_dict[mlp_W_in_key] = ProcessWeights.convert_tensor_to_hf_format(
729 mlp_W_in_key, new_mlp_W_in, cfg, adapter, layer
730 )
731 if getattr(cfg, "gated_mlp", False) and mlp_W_gate_key is not None:
732 mlp_W_gate = ProcessWeights.convert_tensor_to_tl_format(
733 mlp_W_gate_key, state_dict, state_dict.get(mlp_W_gate_key), cfg, adapter, layer
734 )
735 # Combined gate+up (OpenELM): no separate gate, already folded above
736 if mlp_W_gate is not None: 736 ↛ 761line 736 didn't jump to line 761 because the condition on line 736 was always true
737 new_mlp_W_gate = mlp_W_gate * ln2_w_broadcast
738 state_dict[mlp_W_gate_key] = ProcessWeights.convert_tensor_to_hf_format(
739 mlp_W_gate_key, new_mlp_W_gate, cfg, adapter, layer
740 )
741 # Also fold ln2 bias into gate bias (mirrors the in-proj bias folding above)
742 if fold_biases and ln2_b is not None and mlp_b_gate_key is not None:
743 mlp_b_gate = ProcessWeights.convert_tensor_to_tl_format(
744 mlp_b_gate_key,
745 state_dict,
746 state_dict.get(mlp_b_gate_key),
747 cfg,
748 adapter,
749 layer,
750 )
751 ln2_b_gate_folded = (mlp_W_gate * ln2_b_broadcast).sum(sum_dim)
752 if mlp_b_gate is not None: 752 ↛ 753line 752 didn't jump to line 753 because the condition on line 752 was never true
753 new_mlp_b_gate = mlp_b_gate + ln2_b_gate_folded
754 else:
755 new_mlp_b_gate = ln2_b_gate_folded
756 state_dict[mlp_b_gate_key] = ProcessWeights.convert_tensor_to_hf_format(
757 mlp_b_gate_key, new_mlp_b_gate, cfg, adapter, layer
758 )
759 # After folding, set ln2.w to identity (skip for parallel override —
760 # ln1 was already set to identity by the attention folding code).
761 if ln2_w_key is not None: 761 ↛ 773line 761 didn't jump to line 773 because the condition on line 761 was always true
762 identity_ln2 = (
763 torch.zeros_like(ln2_w) if rmsnorm_uses_offset else torch.ones_like(ln2_w)
764 )
765 state_dict[ln2_w_key] = identity_ln2
766 alternate_ln2_w_key = (
767 ln2_w_key.replace("ln_2", "ln2")
768 if "ln_2" in ln2_w_key
769 else ln2_w_key.replace("ln2", "ln_2")
770 )
771 if alternate_ln2_w_key != ln2_w_key and alternate_ln2_w_key in state_dict: 771 ↛ 772line 771 didn't jump to line 772 because the condition on line 771 was never true
772 state_dict[alternate_ln2_w_key] = identity_ln2
773 if center_weights and mlp_W_in_key in state_dict:
774 mlp_W_in_centered = ProcessWeights.convert_tensor_to_tl_format(
775 mlp_W_in_key, state_dict, state_dict.get(mlp_W_in_key), cfg, adapter, layer
776 )
777 assert mlp_W_in_centered is not None, f"MLP W_in not found at key {mlp_W_in_key}"
778 # Center along d_model: TL [d_model, d_mlp] or HF [d_mlp, d_model]
779 d_model = cfg.d_model if cfg is not None else None
780 if (
781 d_model is not None
782 and mlp_W_in_centered.shape[0] == d_model
783 and mlp_W_in_centered.shape[-1] != d_model
784 ):
785 # TL format [d_model, d_mlp]
786 mlp_W_in_centered = mlp_W_in_centered - mlp_W_in_centered.mean(0, keepdim=True)
787 elif ( 787 ↛ 796line 787 didn't jump to line 796 because the condition on line 787 was always true
788 d_model is not None
789 and mlp_W_in_centered.shape[-1] == d_model
790 and mlp_W_in_centered.shape[0] != d_model
791 ):
792 # HF format [d_mlp, d_model]
793 mlp_W_in_centered = mlp_W_in_centered - mlp_W_in_centered.mean(-1, keepdim=True)
794 else:
795 # Fallback: assume TL format
796 mlp_W_in_centered = mlp_W_in_centered - mlp_W_in_centered.mean(0, keepdim=True)
797 state_dict[mlp_W_in_key] = ProcessWeights.convert_tensor_to_hf_format(
798 mlp_W_in_key, mlp_W_in_centered, cfg, adapter, layer
799 )
800 if getattr(cfg, "act_fn", None) is not None and cfg.act_fn.startswith("solu"):
801 mlp_b_out_key = ProcessWeights._get_param_key(f"blocks.{layer}.mlp.b_out", adapter)
802 mlp_W_out_key = ProcessWeights._get_param_key(f"blocks.{layer}.mlp.W_out", adapter)
803 mlp_ln_b_key = ProcessWeights._get_param_key(f"blocks.{layer}.mlp.ln.b", adapter)
804 mlp_ln_w_key = ProcessWeights._get_param_key(f"blocks.{layer}.mlp.ln.w", adapter)
806 mlp_b_out = ProcessWeights.convert_tensor_to_tl_format(
807 mlp_b_out_key, state_dict, state_dict.get(mlp_b_out_key), cfg, adapter, layer
808 )
809 mlp_W_out = ProcessWeights.convert_tensor_to_tl_format(
810 mlp_W_out_key, state_dict, state_dict.get(mlp_W_out_key), cfg, adapter, layer
811 )
812 mlp_ln_b = state_dict.get(mlp_ln_b_key)
813 mlp_ln_w = state_dict.get(mlp_ln_w_key)
814 assert mlp_b_out is not None, f"MLP b_out not found at key {mlp_b_out_key}"
815 assert mlp_W_out is not None, f"MLP W_out not found at key {mlp_W_out_key}"
816 assert mlp_ln_b is not None, f"MLP ln.b not found at key {mlp_ln_b_key}"
817 assert mlp_ln_w is not None, f"MLP ln.w not found at key {mlp_ln_w_key}"
819 if fold_biases: 819 ↛ 827line 819 didn't jump to line 827 because the condition on line 819 was always true
820 new_mlp_b_out = mlp_b_out + (mlp_W_out * mlp_ln_b[:, None]).sum(-2)
821 state_dict[mlp_b_out_key] = ProcessWeights.convert_tensor_to_hf_format(
822 mlp_b_out_key, new_mlp_b_out, cfg, adapter, layer
823 )
824 if mlp_ln_b_key in state_dict: 824 ↛ 827line 824 didn't jump to line 827 because the condition on line 824 was always true
825 state_dict[mlp_ln_b_key] = torch.zeros_like(mlp_ln_b)
827 new_mlp_W_out = mlp_W_out * mlp_ln_w[:, None]
829 if center_weights: 829 ↛ 849line 829 didn't jump to line 849 because the condition on line 829 was always true
830 # Center along d_mlp dimension. Detect format:
831 # TL format [d_mlp, d_model] -> center along dim=0
832 # HF format [d_model, d_mlp] -> center along dim=-1
833 d_model_val = cfg.d_model if cfg is not None else None
834 if ( 834 ↛ 840line 834 didn't jump to line 840 because the condition on line 834 was always true
835 d_model_val is not None
836 and new_mlp_W_out.shape[-1] == d_model_val
837 and new_mlp_W_out.shape[0] != d_model_val
838 ):
839 new_mlp_W_out = new_mlp_W_out - new_mlp_W_out.mean(0, keepdim=True)
840 elif (
841 d_model_val is not None
842 and new_mlp_W_out.shape[0] == d_model_val
843 and new_mlp_W_out.shape[-1] != d_model_val
844 ):
845 new_mlp_W_out = new_mlp_W_out - new_mlp_W_out.mean(-1, keepdim=True)
846 else:
847 new_mlp_W_out = new_mlp_W_out - new_mlp_W_out.mean(0, keepdim=True)
849 state_dict[mlp_W_out_key] = ProcessWeights.convert_tensor_to_hf_format(
850 mlp_W_out_key, new_mlp_W_out, cfg, adapter, layer
851 )
853 if mlp_ln_w_key in state_dict: 853 ↛ 858line 853 didn't jump to line 858 because the condition on line 853 was always true
854 state_dict[mlp_ln_w_key] = torch.ones_like(mlp_ln_w)
856 # ln2_post.w (Gemma 2/3): keep original; independent post-MLP normalization
858 return state_dict
860 @staticmethod
861 def _store_processed_attention_tensors(
862 state_dict: Dict[str, torch.Tensor],
863 keys: Dict[str, str],
864 wq_tensor: Optional[torch.Tensor],
865 wk_tensor: Optional[torch.Tensor],
866 wv_tensor: Optional[torch.Tensor],
867 bq_tensor: Optional[torch.Tensor],
868 bk_tensor: Optional[torch.Tensor],
869 bv_tensor: Optional[torch.Tensor],
870 adapter,
871 cfg,
872 layer: int,
873 ) -> Dict[str, torch.Tensor]:
874 """Store processed attention tensors back to state dict in appropriate format.
876 Args:
877 state_dict: The state dictionary to update (modified in place)
878 keys: Dictionary mapping tensor names to state dict keys
879 wq_tensor, wk_tensor, wv_tensor: Processed attention weight tensors
880 bq_tensor, bk_tensor, bv_tensor: Processed attention bias tensors
881 adapter: Optional architecture adapter for parameter key translation
882 cfg: Model configuration object
883 layer: The layer index
884 """
885 if wq_tensor is None: 885 ↛ 886line 885 didn't jump to line 886 because the condition on line 885 was never true
886 return state_dict
887 wq_key = keys["W_Q"]
888 wk_key = keys["W_K"]
889 wv_key = keys["W_V"]
890 bq_key = keys["b_Q"]
891 bk_key = keys["b_K"]
892 bv_key = keys["b_V"]
894 # Store processed tensors directly in 3D format (set_processed_weights will flatten to 2D)
895 if wq_tensor is None or wk_tensor is None or wv_tensor is None: 895 ↛ 896line 895 didn't jump to line 896 because the condition on line 895 was never true
896 raise ValueError(f"Required attention weights missing for layer {layer}")
897 state_dict[wq_key] = ProcessWeights.convert_tensor_to_hf_format(
898 wq_key, wq_tensor, cfg, adapter, layer_idx=layer
899 )
900 state_dict[wk_key] = ProcessWeights.convert_tensor_to_hf_format(
901 wk_key, wk_tensor, cfg, adapter, layer_idx=layer
902 )
903 state_dict[wv_key] = ProcessWeights.convert_tensor_to_hf_format(
904 wv_key, wv_tensor, cfg, adapter, layer_idx=layer
905 )
906 if bq_tensor is not None:
907 state_dict[bq_key] = ProcessWeights.convert_tensor_to_hf_format(
908 bq_key, bq_tensor, cfg, adapter, layer_idx=layer
909 )
910 if bk_tensor is not None:
911 state_dict[bk_key] = ProcessWeights.convert_tensor_to_hf_format(
912 bk_key, bk_tensor, cfg, adapter, layer_idx=layer
913 )
914 if bv_tensor is not None:
915 state_dict[bv_key] = ProcessWeights.convert_tensor_to_hf_format(
916 bv_key, bv_tensor, cfg, adapter, layer_idx=layer
917 )
919 return state_dict
921 @staticmethod
922 def _fold_unembed_layer_norm(
923 state_dict: Dict[str, torch.Tensor], cfg, fold_biases: bool, center_weights: bool, adapter
924 ) -> Dict[str, torch.Tensor]:
925 """Fold LayerNorm into unembedding layer.
927 Args:
928 state_dict: The state dictionary to process (modified in place)
929 cfg: Model configuration object
930 fold_biases: Whether to fold LayerNorm biases
931 center_weights: Whether to center weights after folding
932 adapter: Optional architecture adapter for parameter key translation
933 """
934 unembed_b_U_key = ProcessWeights._get_param_key("unembed.b_U", adapter)
935 unembed_W_U_key = ProcessWeights._get_param_key("unembed.W_U", adapter)
936 ln_final_b_key = ProcessWeights._get_param_key("ln_final.b", adapter)
937 ln_final_w_key = ProcessWeights._get_param_key("ln_final.w", adapter)
939 # Skip layer norm folding if ln_final doesn't exist
940 # (e.g., encoder-decoder models like T5 have encoder_ln_final/decoder_ln_final instead)
941 if ln_final_w_key not in state_dict: 941 ↛ 942line 941 didn't jump to line 942 because the condition on line 941 was never true
942 return state_dict
944 has_unembed_bias = unembed_b_U_key in state_dict
945 unembed_weight = ProcessWeights.convert_tensor_to_tl_format(
946 unembed_W_U_key, state_dict, state_dict.get(unembed_W_U_key), cfg, adapter, None
947 )
948 ln_weight = state_dict[ln_final_w_key]
949 assert unembed_weight is not None, f"Unembed weight not found at key {unembed_W_U_key}"
950 # rmsnorm_uses_offset: effective scale is (1+w), identity is 0.0
951 rmsnorm_uses_offset = getattr(cfg, "rmsnorm_uses_offset", False)
952 effective_ln_weight = (1.0 + ln_weight) if rmsnorm_uses_offset else ln_weight
953 if len(unembed_weight.shape) == 2 and len(ln_weight.shape) == 1: 953 ↛ 963line 953 didn't jump to line 963 because the condition on line 953 was always true
954 if unembed_weight.shape[1] == ln_weight.shape[0]:
955 new_unembed_weight = unembed_weight * effective_ln_weight[None, :]
956 elif unembed_weight.shape[0] == ln_weight.shape[0]: 956 ↛ 959line 956 didn't jump to line 959 because the condition on line 956 was always true
957 new_unembed_weight = unembed_weight * effective_ln_weight[:, None]
958 else:
959 raise ValueError(
960 f"Cannot broadcast unembedding weight {unembed_weight.shape} with layer norm weight {ln_weight.shape}"
961 )
962 else:
963 raise ValueError(
964 f"Unexpected tensor shapes: unembedding {unembed_weight.shape}, layer norm {ln_weight.shape}"
965 )
966 state_dict[unembed_W_U_key] = ProcessWeights.convert_tensor_to_hf_format(
967 unembed_W_U_key, new_unembed_weight, cfg, adapter, None
968 )
969 # Set ln_final.w to identity: zeros (rmsnorm_uses_offset) or ones (standard)
970 identity_val = (
971 torch.zeros_like(ln_weight) if rmsnorm_uses_offset else torch.ones_like(ln_weight)
972 )
973 if ln_final_w_key in state_dict: 973 ↛ 975line 973 didn't jump to line 975 because the condition on line 973 was always true
974 state_dict[ln_final_w_key] = identity_val
975 alternate_final_w_key = (
976 ln_final_w_key.replace("ln_f", "ln_final")
977 if "ln_f" in ln_final_w_key
978 else ln_final_w_key.replace("ln_final", "ln_f")
979 )
980 if alternate_final_w_key != ln_final_w_key and alternate_final_w_key in state_dict: 980 ↛ 981line 980 didn't jump to line 981 because the condition on line 980 was never true
981 state_dict[alternate_final_w_key] = identity_val
982 if center_weights:
983 unembed_weight_centered = ProcessWeights.convert_tensor_to_tl_format(
984 unembed_W_U_key, state_dict, state_dict.get(unembed_W_U_key), cfg, adapter, None
985 )
986 assert (
987 unembed_weight_centered is not None
988 ), f"Unembed weight not found at key {unembed_W_U_key}"
989 if len(unembed_weight_centered.shape) == 2: 989 ↛ 1010line 989 didn't jump to line 1010 because the condition on line 989 was always true
990 # Center along d_model: detect TL vs HF format
991 d_vocab = getattr(cfg, "d_vocab", None) if cfg is not None else None
992 if (
993 d_vocab is not None
994 and unembed_weight_centered.shape[0] == d_vocab
995 and unembed_weight_centered.shape[-1] != d_vocab
996 ):
997 # HF format [d_vocab, d_model] — center along dim=-1
998 unembed_weight_centered = (
999 unembed_weight_centered - unembed_weight_centered.mean(-1, keepdim=True)
1000 )
1001 else:
1002 # TL format [d_model, d_vocab] — center along dim=0
1003 unembed_weight_centered = (
1004 unembed_weight_centered - unembed_weight_centered.mean(0, keepdim=True)
1005 )
1006 state_dict[unembed_W_U_key] = ProcessWeights.convert_tensor_to_hf_format(
1007 unembed_W_U_key, unembed_weight_centered, cfg, adapter, None
1008 )
1009 else:
1010 raise ValueError(
1011 f"Unexpected unembedding weight shape: {unembed_weight_centered.shape}"
1012 )
1014 return state_dict
1016 @staticmethod
1017 def _fold_final_rms_bias(
1018 state_dict: Dict[str, torch.Tensor], cfg, fold_biases: bool, adapter
1019 ) -> Dict[str, torch.Tensor]:
1020 """Fold final RMS bias into unembedding (separate from regular unembed folding).
1022 Args:
1023 state_dict: The state dictionary to process (modified in place)
1024 cfg: Model configuration object
1025 fold_biases: Whether to fold LayerNorm biases
1026 adapter: Optional architecture adapter for parameter key translation
1027 """
1028 unembed_b_U_key = ProcessWeights._get_param_key("unembed.b_U", adapter)
1029 unembed_W_U_key = ProcessWeights._get_param_key("unembed.W_U", adapter)
1030 ln_final_b_key = ProcessWeights._get_param_key("ln_final.b", adapter)
1031 has_unembed_bias = unembed_b_U_key in state_dict
1032 has_ln_final_bias = ln_final_b_key in state_dict
1033 if (
1034 not getattr(cfg, "final_rms", False)
1035 and fold_biases
1036 and has_unembed_bias
1037 and has_ln_final_bias
1038 ):
1039 unembed_weight = ProcessWeights.convert_tensor_to_tl_format(
1040 unembed_W_U_key, state_dict, state_dict.get(unembed_W_U_key), cfg, adapter, None
1041 )
1042 ln_bias = state_dict[ln_final_b_key]
1043 assert unembed_weight is not None, f"Unembed weight not found at key {unembed_W_U_key}"
1044 if len(unembed_weight.shape) == 2 and len(ln_bias.shape) == 1: 1044 ↛ 1054line 1044 didn't jump to line 1054 because the condition on line 1044 was always true
1045 if unembed_weight.shape[1] == ln_bias.shape[0]:
1046 bias_contribution = (unembed_weight * ln_bias[None, :]).sum(dim=-1)
1047 elif unembed_weight.shape[0] == ln_bias.shape[0]: 1047 ↛ 1050line 1047 didn't jump to line 1050 because the condition on line 1047 was always true
1048 bias_contribution = (unembed_weight * ln_bias[:, None]).sum(dim=-2)
1049 else:
1050 raise ValueError(
1051 f"Cannot broadcast unembedding weight {unembed_weight.shape} with layer norm bias {ln_bias.shape}"
1052 )
1053 else:
1054 raise ValueError(
1055 f"Unexpected tensor shapes: unembedding {unembed_weight.shape}, layer norm bias {ln_bias.shape}"
1056 )
1057 unembed_b_U = ProcessWeights.convert_tensor_to_tl_format(
1058 unembed_b_U_key, state_dict, state_dict.get(unembed_b_U_key), cfg, adapter, None
1059 )
1060 assert unembed_b_U is not None, f"Unembed bias not found at key {unembed_b_U_key}"
1061 new_unembed_b_U = unembed_b_U + bias_contribution
1062 state_dict[unembed_b_U_key] = ProcessWeights.convert_tensor_to_hf_format(
1063 unembed_b_U_key, new_unembed_b_U, cfg, adapter, None
1064 )
1065 if ln_final_b_key in state_dict: 1065 ↛ 1067line 1065 didn't jump to line 1067 because the condition on line 1065 was always true
1066 state_dict[ln_final_b_key] = torch.zeros_like(ln_bias)
1067 alternate_final_b_key = (
1068 ln_final_b_key.replace("ln_f", "ln_final")
1069 if "ln_f" in ln_final_b_key
1070 else ln_final_b_key.replace("ln_final", "ln_f")
1071 )
1072 if alternate_final_b_key != ln_final_b_key and alternate_final_b_key in state_dict: 1072 ↛ 1073line 1072 didn't jump to line 1073 because the condition on line 1072 was never true
1073 state_dict[alternate_final_b_key] = torch.zeros_like(ln_bias)
1075 return state_dict
1077 @staticmethod
1078 def fold_layer_norm(
1079 state_dict: Dict[str, torch.Tensor],
1080 cfg,
1081 fold_biases: bool = True,
1082 center_weights: bool = True,
1083 adapter=None,
1084 ) -> Dict[str, torch.Tensor]:
1085 """Fold Layer Norm. Can also be used to fold RMS Norm, when fold_biases and center_weights are set to False.
1087 Takes in a state dict from a pretrained model, formatted to be consistent with
1088 HookedTransformer but with LayerNorm weights and biases. Folds these into the neighbouring
1089 weights. See further_comments.md for more details.
1091 Args:
1092 state_dict (Dict[str, torch.Tensor]): State dict of pretrained model.
1093 cfg: Model configuration object with n_layers, n_key_value_heads, etc.
1094 fold_biases (bool): Enables folding of LN biases. Should be disabled when RMS Norm is used.
1095 center_weights (bool): Enables the centering of weights after folding in LN. Should be disabled when RMS Norm is used.
1096 adapter: Optional architecture adapter for parameter key translation.
1098 Returns:
1099 Dict[str, torch.Tensor]: Modified state dict with LayerNorm folded into linear layers.
1100 """
1101 # Make a deep copy to avoid modifying the original
1102 state_dict = {
1103 k: v.clone() if isinstance(v, torch.Tensor) else v for k, v in state_dict.items()
1104 }
1105 gqa = "" if getattr(cfg, "n_key_value_heads", None) is None else "_"
1106 for l in range(cfg.n_layers):
1107 state_dict = ProcessWeights._fold_layer(
1108 state_dict, cfg, l, fold_biases, center_weights, adapter, gqa
1109 )
1110 state_dict = ProcessWeights._fold_final_rms_bias(state_dict, cfg, fold_biases, adapter)
1111 state_dict = ProcessWeights._fold_unembed_layer_norm(
1112 state_dict, cfg, fold_biases, center_weights, adapter
1113 )
1114 return state_dict
1116 @staticmethod
1117 def center_writing_weights(
1118 state_dict: Dict[str, torch.Tensor], cfg, adapter=None
1119 ) -> Dict[str, torch.Tensor]:
1120 """Center Writing Weights.
1122 Centers the weights of the model that write to the residual stream - W_out, W_E, W_pos and
1123 W_out. This is done by subtracting the mean of the weights from the weights themselves. This
1124 is done in-place. See fold_layer_norm for more details.
1126 Args:
1127 state_dict (Dict[str, torch.Tensor]): State dict of the model.
1128 cfg: Model configuration object.
1129 adapter: Optional architecture adapter for parameter key translation.
1131 Returns:
1132 Dict[str, torch.Tensor]: Modified state dict with centered writing weights.
1133 """
1134 # Post-norm models leave the first attention's input un-normed, so centering
1135 # the embedding would shift a residual stream nothing re-normalizes.
1136 architecture = getattr(cfg, "original_architecture", None)
1137 if architecture in POST_NORM_ARCHITECTURES:
1138 print(f"Not centering embedding weights for {architecture}")
1139 else:
1140 # Make a deep copy to avoid modifying the original
1141 embed_W_E_key = ProcessWeights._get_param_key("embed.W_E", adapter)
1142 try:
1143 pos_embed_W_pos_key = (
1144 ProcessWeights._get_param_key("pos_embed.W_pos", adapter)
1145 if getattr(cfg, "positional_embedding_type", "standard")
1146 not in ("rotary", "alibi", "none")
1147 else None
1148 )
1149 except ValueError:
1150 pos_embed_W_pos_key = None
1151 if embed_W_E_key not in state_dict: 1151 ↛ 1152line 1151 didn't jump to line 1152 because the condition on line 1151 was never true
1152 raise KeyError(
1153 f"Expected embedding key '{embed_W_E_key}' not found in state_dict. Available keys: {list(state_dict.keys())[:10]}..."
1154 )
1155 embed_W_E = ProcessWeights.convert_tensor_to_tl_format(
1156 embed_W_E_key, state_dict, state_dict.get(embed_W_E_key), cfg, adapter, None
1157 )
1158 assert embed_W_E is not None, f"Embedding not found at key {embed_W_E_key}"
1159 embed_W_E = embed_W_E - embed_W_E.mean(-1, keepdim=True)
1160 state_dict[embed_W_E_key] = ProcessWeights.convert_tensor_to_hf_format(
1161 embed_W_E_key, embed_W_E, cfg, adapter, None
1162 )
1164 if (
1165 getattr(cfg, "positional_embedding_type", "standard")
1166 not in ("rotary", "alibi", "none")
1167 and pos_embed_W_pos_key is not None
1168 ):
1169 if pos_embed_W_pos_key not in state_dict: 1169 ↛ 1170line 1169 didn't jump to line 1170 because the condition on line 1169 was never true
1170 raise KeyError(
1171 f"Expected positional embedding key '{pos_embed_W_pos_key}' not found in state_dict. Available keys: {list(state_dict.keys())[:10]}..."
1172 )
1173 pos_embed_W_pos = ProcessWeights.convert_tensor_to_tl_format(
1174 pos_embed_W_pos_key,
1175 state_dict,
1176 state_dict.get(pos_embed_W_pos_key),
1177 cfg,
1178 adapter,
1179 None,
1180 )
1181 assert (
1182 pos_embed_W_pos is not None
1183 ), f"Positional embedding not found at key {pos_embed_W_pos_key}"
1184 pos_embed_W_pos = pos_embed_W_pos - pos_embed_W_pos.mean(-1, keepdim=True)
1185 state_dict[pos_embed_W_pos_key] = ProcessWeights.convert_tensor_to_hf_format(
1186 pos_embed_W_pos_key, pos_embed_W_pos, cfg, adapter, None
1187 )
1188 for l in range(cfg.n_layers):
1189 attn_W_O_key = ProcessWeights._get_param_key(f"blocks.{l}.attn.W_O", adapter)
1190 attn_b_O_key = ProcessWeights._get_param_key(f"blocks.{l}.attn.b_O", adapter)
1191 try:
1192 mlp_W_out_key = ProcessWeights._resolve_state_dict_key(
1193 state_dict, ProcessWeights._get_param_key(f"blocks.{l}.mlp.W_out", adapter), l
1194 )
1195 mlp_b_out_key = ProcessWeights._resolve_state_dict_key(
1196 state_dict, ProcessWeights._get_param_key(f"blocks.{l}.mlp.b_out", adapter), l
1197 )
1198 except ValueError:
1199 mlp_W_out_key = None
1200 mlp_b_out_key = None
1201 if attn_W_O_key in state_dict: 1201 ↛ 1219line 1201 didn't jump to line 1219 because the condition on line 1201 was always true
1202 attn_W_O = ProcessWeights.convert_tensor_to_tl_format(
1203 attn_W_O_key, state_dict, state_dict.get(attn_W_O_key), cfg, adapter, l
1204 )
1205 assert attn_W_O is not None, f"Attention W_O not found at key {attn_W_O_key}"
1206 attn_W_O = attn_W_O - attn_W_O.mean(-1, keepdim=True)
1207 state_dict[attn_W_O_key] = ProcessWeights.convert_tensor_to_hf_format(
1208 attn_W_O_key, attn_W_O, cfg, adapter, l
1209 )
1210 if attn_b_O_key in state_dict: 1210 ↛ 1219line 1210 didn't jump to line 1219 because the condition on line 1210 was always true
1211 attn_b_O = ProcessWeights.convert_tensor_to_tl_format(
1212 attn_b_O_key, state_dict, state_dict.get(attn_b_O_key), cfg, adapter, l
1213 )
1214 assert attn_b_O is not None, f"Attention b_O not found at key {attn_b_O_key}"
1215 attn_b_O = attn_b_O - attn_b_O.mean()
1216 state_dict[attn_b_O_key] = ProcessWeights.convert_tensor_to_hf_format(
1217 attn_b_O_key, attn_b_O, cfg, adapter, l
1218 )
1219 if not getattr(cfg, "attn_only", False):
1220 is_moe = getattr(cfg, "num_experts", None) is not None and cfg.num_experts > 0
1221 if is_moe:
1222 num_experts = cfg.num_experts
1223 for e in range(num_experts):
1224 expert_W_out_key = None
1225 expert_b_out_key = None
1226 expert_W_out_patterns = [
1227 f"blocks.{l}.mlp.experts.{e}.W_out",
1228 f"blocks.{l}.mlp.experts.{e}.W_out.weight",
1229 ]
1230 for pattern in expert_W_out_patterns:
1231 if pattern in state_dict: 1231 ↛ 1232line 1231 didn't jump to line 1232 because the condition on line 1231 was never true
1232 expert_W_out_key = pattern
1233 break
1234 if expert_W_out_key is None and adapter: 1234 ↛ 1235line 1234 didn't jump to line 1235 because the condition on line 1234 was never true
1235 try:
1236 candidate = ProcessWeights._get_param_key(
1237 f"blocks.{l}.mlp.experts.{e}.W_out", adapter
1238 )
1239 expert_W_out_key = ProcessWeights._resolve_state_dict_key(
1240 state_dict, candidate, l
1241 )
1242 except ValueError:
1243 pass
1244 if expert_W_out_key and expert_W_out_key in state_dict: 1244 ↛ 1245line 1244 didn't jump to line 1245 because the condition on line 1244 was never true
1245 expert_W_out = ProcessWeights.convert_tensor_to_tl_format(
1246 expert_W_out_key,
1247 state_dict,
1248 state_dict.get(expert_W_out_key),
1249 cfg,
1250 adapter,
1251 l,
1252 )
1253 assert (
1254 expert_W_out is not None
1255 ), f"Expert W_out not found at key {expert_W_out_key}"
1256 expert_W_out = expert_W_out - expert_W_out.mean(-1, keepdim=True)
1257 state_dict[
1258 expert_W_out_key
1259 ] = ProcessWeights.convert_tensor_to_hf_format(
1260 expert_W_out_key, expert_W_out, cfg, adapter, l
1261 )
1262 expert_b_out_patterns = [
1263 f"blocks.{l}.mlp.experts.{e}.b_out",
1264 f"blocks.{l}.mlp.experts.{e}.b_out.bias",
1265 ]
1266 for pattern in expert_b_out_patterns:
1267 if pattern in state_dict: 1267 ↛ 1268line 1267 didn't jump to line 1268 because the condition on line 1267 was never true
1268 expert_b_out_key = pattern
1269 break
1270 if expert_b_out_key is None and adapter: 1270 ↛ 1271line 1270 didn't jump to line 1271 because the condition on line 1270 was never true
1271 try:
1272 candidate = ProcessWeights._get_param_key(
1273 f"blocks.{l}.mlp.experts.{e}.b_out", adapter
1274 )
1275 expert_b_out_key = ProcessWeights._resolve_state_dict_key(
1276 state_dict, candidate, l
1277 )
1278 except ValueError:
1279 pass
1280 if expert_b_out_key and expert_b_out_key in state_dict: 1280 ↛ 1281line 1280 didn't jump to line 1281 because the condition on line 1280 was never true
1281 expert_b_out = ProcessWeights.convert_tensor_to_tl_format(
1282 expert_b_out_key,
1283 state_dict,
1284 state_dict.get(expert_b_out_key),
1285 cfg,
1286 adapter,
1287 l,
1288 )
1289 assert (
1290 expert_b_out is not None
1291 ), f"Expert b_out not found at key {expert_b_out_key}"
1292 expert_b_out = expert_b_out - expert_b_out.mean()
1293 state_dict[
1294 expert_b_out_key
1295 ] = ProcessWeights.convert_tensor_to_hf_format(
1296 expert_b_out_key, expert_b_out, cfg, adapter, l
1297 )
1298 elif mlp_W_out_key is not None and mlp_W_out_key in state_dict: 1298 ↛ 1188line 1298 didn't jump to line 1188 because the condition on line 1298 was always true
1299 mlp_W_out = ProcessWeights.convert_tensor_to_tl_format(
1300 mlp_W_out_key, state_dict, state_dict.get(mlp_W_out_key), cfg, adapter, l
1301 )
1302 assert mlp_W_out is not None, f"MLP W_out not found at key {mlp_W_out_key}"
1303 # Center along d_model dimension. In TL format W_out is [d_mlp, d_model]
1304 # so d_model is dim=-1. But bridge adapters may keep HF format
1305 # [d_model, d_mlp] where d_model is dim=0. Detect via cfg.d_model.
1306 if mlp_W_out.shape[-1] == cfg.d_model:
1307 mlp_W_out = mlp_W_out - mlp_W_out.mean(-1, keepdim=True)
1308 elif mlp_W_out.shape[0] == cfg.d_model: 1308 ↛ 1311line 1308 didn't jump to line 1311 because the condition on line 1308 was always true
1309 mlp_W_out = mlp_W_out - mlp_W_out.mean(0, keepdim=True)
1310 else:
1311 mlp_W_out = mlp_W_out - mlp_W_out.mean(-1, keepdim=True)
1312 state_dict[mlp_W_out_key] = ProcessWeights.convert_tensor_to_hf_format(
1313 mlp_W_out_key, mlp_W_out, cfg, adapter, l
1314 )
1315 if mlp_b_out_key is not None and mlp_b_out_key in state_dict: 1315 ↛ 1188line 1315 didn't jump to line 1188 because the condition on line 1315 was always true
1316 mlp_b_out = ProcessWeights.convert_tensor_to_tl_format(
1317 mlp_b_out_key,
1318 state_dict,
1319 state_dict.get(mlp_b_out_key),
1320 cfg,
1321 adapter,
1322 l,
1323 )
1324 assert mlp_b_out is not None, f"MLP b_out not found at key {mlp_b_out_key}"
1325 mlp_b_out = mlp_b_out - mlp_b_out.mean()
1326 state_dict[mlp_b_out_key] = ProcessWeights.convert_tensor_to_hf_format(
1327 mlp_b_out_key, mlp_b_out, cfg, adapter, l
1328 )
1329 return state_dict
1331 @staticmethod
1332 def center_unembed(
1333 state_dict: Dict[str, torch.Tensor], cfg=None, adapter=None
1334 ) -> Dict[str, torch.Tensor]:
1335 """Center the unembedding weights W_U.
1337 This is done by subtracting the mean of the weights from the weights themselves. This is
1338 done in-place. As softmax is translation invariant, this changes the logits but not the log
1339 probs, and makes the model logits (slightly) more interpretable - when trying to understand
1340 how components contribute to the logits, we'll be less misled by components that just add
1341 something to every logit.
1343 Args:
1344 state_dict (Dict[str, torch.Tensor]): State dict of the model.
1345 cfg: Model configuration (used to determine d_vocab for correct centering dimension).
1346 adapter: Optional architecture adapter for parameter key translation.
1348 Returns:
1349 Dict[str, torch.Tensor]: Modified state dict with centered unembedding weights.
1350 """
1351 # Make a deep copy to avoid modifying the original
1352 state_dict = {
1353 k: v.clone() if isinstance(v, torch.Tensor) else v for k, v in state_dict.items()
1354 }
1355 unembed_W_U_key = ProcessWeights._get_param_key("unembed.W_U", adapter)
1356 unembed_b_U_key = ProcessWeights._get_param_key("unembed.b_U", adapter)
1357 if unembed_W_U_key not in state_dict:
1358 raise KeyError(
1359 f"Expected unembedding weight key '{unembed_W_U_key}' not found in state_dict. Available keys: {list(state_dict.keys())[:10]}..."
1360 )
1361 W_U = ProcessWeights.convert_tensor_to_tl_format(
1362 unembed_W_U_key, state_dict, state_dict.get(unembed_W_U_key), None, adapter, None
1363 )
1364 assert W_U is not None, f"Unembed weight not found at key {unembed_W_U_key}"
1366 # Detect W_U format to center along correct dim (wrong dim corrupts output)
1367 vocab_dim = -1 # Default: TL format [d_model, d_vocab]
1368 if cfg is not None:
1369 d_vocab = getattr(cfg, "d_vocab", None)
1370 d_model = getattr(cfg, "d_model", None)
1371 if d_vocab is not None and W_U.shape[0] == d_vocab and W_U.shape[-1] != d_vocab:
1372 # HF format [d_vocab, d_model] — center along dim=0
1373 vocab_dim = 0
1374 elif d_model is not None and W_U.shape[-1] == d_model and W_U.shape[0] != d_model:
1375 # Padded-vocab checkpoints (e.g. HyenaDNA pads 12→16 rows) defeat
1376 # the d_vocab match; the d_model axis is unambiguous.
1377 vocab_dim = 0
1378 W_U = W_U - W_U.mean(vocab_dim, keepdim=True)
1379 state_dict[unembed_W_U_key] = ProcessWeights.convert_tensor_to_hf_format(
1380 unembed_W_U_key, W_U, None, adapter, None
1381 )
1382 if unembed_b_U_key in state_dict:
1383 unembed_b_U = ProcessWeights.convert_tensor_to_tl_format(
1384 unembed_b_U_key, state_dict, state_dict.get(unembed_b_U_key), None, adapter, None
1385 )
1386 assert unembed_b_U is not None, f"Unembed bias not found at key {unembed_b_U_key}"
1387 unembed_b_U = unembed_b_U - unembed_b_U.mean()
1388 state_dict[unembed_b_U_key] = ProcessWeights.convert_tensor_to_hf_format(
1389 unembed_b_U_key, unembed_b_U, None, adapter, None
1390 )
1391 return state_dict
1393 @staticmethod
1394 def fold_value_biases(
1395 state_dict: Dict[str, torch.Tensor], cfg, adapter=None
1396 ) -> Dict[str, torch.Tensor]:
1397 """Fold the value biases into the output bias.
1399 Because attention patterns add up to 1, the value biases always have a constant effect on a
1400 head's output. Further, as the outputs of each head in a layer add together, each head's
1401 value bias has a constant effect on the *layer's* output, which can make it harder to
1402 interpret the effect of any given head, and it doesn't matter which head a bias is
1403 associated with. We can factor this all into a single output bias to the layer, and make it
1404 easier to interpret the head's output. Formally, we take b_O_new = b_O_original +
1405 sum_head(b_V_head @ W_O_head).
1407 Args:
1408 state_dict (Dict[str, torch.Tensor]): State dict of the model.
1409 cfg: Model configuration object.
1410 adapter: Optional architecture adapter for parameter key translation.
1412 Returns:
1413 Dict[str, torch.Tensor]: Modified state dict with value biases folded into output bias.
1414 """
1415 # Make a deep copy to avoid modifying the original
1416 state_dict = {
1417 k: v.clone() if isinstance(v, torch.Tensor) else v for k, v in state_dict.items()
1418 }
1419 layer = 0
1420 for layer in range(cfg.n_layers):
1421 split_v_bias_key = f"blocks.{layer}.attn.v.bias"
1422 if split_v_bias_key in state_dict:
1423 b_V_key = split_v_bias_key
1424 W_O_key = ProcessWeights._get_param_key(f"blocks.{layer}.attn.W_O", adapter)
1425 b_O_key = ProcessWeights._get_param_key(f"blocks.{layer}.attn.b_O", adapter)
1426 else:
1427 if getattr(cfg, "n_key_value_heads", None) is None:
1428 b_V_key = ProcessWeights._get_param_key(f"blocks.{layer}.attn.b_V", adapter)
1429 else:
1430 b_V_key = ProcessWeights._get_param_key(f"blocks.{layer}.attn._b_V", adapter)
1431 W_O_key = ProcessWeights._get_param_key(f"blocks.{layer}.attn.W_O", adapter)
1432 b_O_key = ProcessWeights._get_param_key(f"blocks.{layer}.attn.b_O", adapter)
1433 if b_V_key in state_dict:
1434 b_V = ProcessWeights.convert_tensor_to_tl_format(
1435 b_V_key, state_dict, state_dict.get(b_V_key), cfg, adapter, layer
1436 )
1437 assert b_V is not None, f"Value bias not found at key {b_V_key}"
1438 if b_V.numel() == 0: 1438 ↛ 1439line 1438 didn't jump to line 1439 because the condition on line 1438 was never true
1439 continue
1440 W_O = ProcessWeights.convert_tensor_to_tl_format(
1441 W_O_key, state_dict, state_dict.get(W_O_key), cfg, adapter, layer
1442 )
1443 assert W_O is not None, f"Attention W_O not found at key {W_O_key}"
1444 if b_O_key not in state_dict:
1445 # Create zero b_O to absorb the folded value bias
1446 b_O_original = torch.zeros(cfg.d_model, dtype=b_V.dtype, device=b_V.device)
1447 state_dict[b_O_key] = b_O_original
1448 else:
1449 b_O_original_maybe = ProcessWeights.convert_tensor_to_tl_format(
1450 b_O_key, state_dict, state_dict.get(b_O_key), cfg, adapter, layer
1451 )
1452 assert (
1453 b_O_original_maybe is not None
1454 ), f"Attention b_O not found at key {b_O_key}"
1455 b_O_original = b_O_original_maybe
1456 # Align W_O / b_O to b_V's device.
1457 if W_O.device != b_V.device: 1457 ↛ 1458line 1457 didn't jump to line 1458 because the condition on line 1457 was never true
1458 W_O = W_O.to(b_V.device)
1459 if b_O_original.device != b_V.device: 1459 ↛ 1460line 1459 didn't jump to line 1460 because the condition on line 1459 was never true
1460 b_O_original = b_O_original.to(b_V.device)
1461 is_split_format = ".attn.v.bias" in b_V_key or ".attn.k.bias" in b_V_key
1462 if is_split_format and len(b_V.shape) == 1 and (len(W_O.shape) == 2): 1462 ↛ 1463line 1462 didn't jump to line 1463 because the condition on line 1462 was never true
1463 n_heads = cfg.n_heads
1464 d_head = cfg.d_head
1465 d_model = cfg.d_model
1466 b_V_only = b_V
1467 b_V_reshaped = b_V_only.reshape(n_heads, d_head)
1468 W_O_reshaped = einops.rearrange(W_O, "(i h) m -> i h m", i=n_heads)
1469 folded_b_O = b_O_original + (b_V_reshaped[:, :, None] * W_O_reshaped).sum(
1470 [0, 1]
1471 )
1472 state_dict[b_O_key] = ProcessWeights.convert_tensor_to_hf_format(
1473 b_O_key, folded_b_O, cfg, adapter, layer
1474 )
1475 tl_b_O_key = f"blocks.{layer}.attn.b_O"
1476 if tl_b_O_key in state_dict:
1477 state_dict[tl_b_O_key] = ProcessWeights.convert_tensor_to_hf_format(
1478 tl_b_O_key, folded_b_O, cfg, adapter, layer
1479 )
1480 state_dict[b_V_key] = ProcessWeights.convert_tensor_to_hf_format(
1481 b_V_key, torch.zeros_like(b_V), cfg, adapter, layer
1482 )
1483 elif len(b_V.shape) == 1 and len(W_O.shape) == 2: 1483 ↛ 1484line 1483 didn't jump to line 1484 because the condition on line 1483 was never true
1484 n_heads = cfg.n_heads
1485 d_head = cfg.d_head
1486 d_model = cfg.d_model
1487 v_bias_start = 2 * n_heads * d_head
1488 v_bias_end = 3 * n_heads * d_head
1489 b_V_only = b_V[v_bias_start:v_bias_end]
1490 if b_V_only.numel() == 0:
1491 continue
1492 b_V_reshaped = b_V_only.reshape(n_heads, d_head)
1493 W_O_reshaped = einops.rearrange(W_O, "(i h) m -> i h m", i=n_heads)
1494 folded_b_O = b_O_original + (b_V_reshaped[:, :, None] * W_O_reshaped).sum(
1495 [0, 1]
1496 )
1497 new_b_V = b_V.clone()
1498 new_b_V[v_bias_start:v_bias_end] = 0
1499 state_dict[b_V_key] = ProcessWeights.convert_tensor_to_hf_format(
1500 b_V_key, new_b_V, cfg, adapter, layer
1501 )
1502 elif is_split_format and len(b_V.shape) == 1 and len(W_O.shape) == 3: 1502 ↛ 1504line 1502 didn't jump to line 1504 because the condition on line 1502 was never true
1503 # Split bias [n_heads * d_head] with W_O already in TL format [n_heads, d_head, d_model]
1504 n_heads = cfg.n_heads
1505 d_head = cfg.d_head
1506 b_V_reshaped = b_V.reshape(n_heads, d_head)
1507 if getattr(cfg, "n_key_value_heads", None) is not None:
1508 b_V_reshaped = torch.repeat_interleave(
1509 b_V_reshaped, dim=0, repeats=cfg.n_heads // cfg.n_key_value_heads
1510 )
1511 folded_b_O = b_O_original + (b_V_reshaped[:, :, None] * W_O).sum([0, 1])
1512 state_dict[b_V_key] = ProcessWeights.convert_tensor_to_hf_format(
1513 b_V_key, torch.zeros_like(b_V), cfg, adapter, layer
1514 )
1515 elif len(b_V.shape) == 2 and len(W_O.shape) == 3: 1515 ↛ 1529line 1515 didn't jump to line 1529 because the condition on line 1515 was always true
1516 b_V_original_shape = b_V.shape
1517 if getattr(cfg, "n_key_value_heads", None) is not None:
1518 b_V = torch.repeat_interleave(
1519 b_V, dim=0, repeats=cfg.n_heads // cfg.n_key_value_heads
1520 )
1521 folded_b_O = b_O_original + (b_V[:, :, None] * W_O).sum([0, 1])
1522 state_dict[b_V_key] = ProcessWeights.convert_tensor_to_hf_format(
1523 b_V_key,
1524 torch.zeros(b_V_original_shape, dtype=b_V.dtype, device=b_V.device),
1525 cfg,
1526 adapter,
1527 layer,
1528 )
1529 elif len(b_V.shape) == 2 and len(W_O.shape) == 2:
1530 n_heads = cfg.n_heads
1531 d_head = cfg.d_head
1532 d_model = cfg.d_model
1533 b_V_original_shape = b_V.shape
1535 # Handle split QKV format where bias might be [1, d_model] or [n_heads, d_head]
1536 is_split_format = ".attn.v.bias" in b_V_key or ".attn.k.bias" in b_V_key
1537 if is_split_format and b_V.shape[0] == 1 and b_V.shape[1] == n_heads * d_head:
1538 # Reshape [1, n_heads * d_head] to [n_heads, d_head]
1539 b_V = b_V.reshape(n_heads, d_head)
1540 elif b_V.shape != (n_heads, d_head):
1541 # If not already [n_heads, d_head], try to reshape
1542 if b_V.numel() == n_heads * d_head:
1543 b_V = b_V.reshape(n_heads, d_head)
1545 if getattr(cfg, "n_key_value_heads", None) is not None:
1546 b_V = torch.repeat_interleave(
1547 b_V, dim=0, repeats=cfg.n_heads // cfg.n_key_value_heads
1548 )
1550 W_O_reshaped = einops.rearrange(W_O, "(i h) m -> i h m", i=n_heads)
1551 folded_b_O = b_O_original + (b_V[:, :, None] * W_O_reshaped).sum([0, 1])
1552 state_dict[b_V_key] = ProcessWeights.convert_tensor_to_hf_format(
1553 b_V_key,
1554 torch.zeros(b_V_original_shape, dtype=b_V.dtype, device=b_V.device),
1555 cfg,
1556 adapter,
1557 layer,
1558 )
1559 else:
1560 raise ValueError(f"Unexpected tensor shapes: b_V {b_V.shape}, W_O {W_O.shape}")
1561 state_dict[b_O_key] = ProcessWeights.convert_tensor_to_hf_format(
1562 b_O_key, folded_b_O, cfg, adapter, layer
1563 )
1564 return state_dict
1566 @staticmethod
1567 def process_weights(
1568 state_dict: Dict[str, torch.Tensor],
1569 cfg,
1570 fold_ln: bool = True,
1571 center_writing_weights: bool = True,
1572 center_unembed: bool = True,
1573 fold_value_biases: bool = True,
1574 refactor_factored_attn_matrices: bool = False,
1575 adapter=None,
1576 ) -> Dict[str, torch.Tensor]:
1577 """Apply all weight processing transformations in the correct order.
1579 This is a convenience function that applies all the weight processing steps
1580 in the same order as HookedTransformer.load_and_process_state_dict().
1582 Args:
1583 state_dict (Dict[str, torch.Tensor]): State dict of the model.
1584 cfg: Model configuration object.
1585 fold_ln (bool): Whether to fold LayerNorm weights into subsequent layers.
1586 center_writing_weights (bool): Whether to center weights writing to residual stream.
1587 center_unembed (bool): Whether to center unembedding weights.
1588 fold_value_biases (bool): Whether to fold value biases into output bias.
1589 refactor_factored_attn_matrices (bool): Whether to refactor attention matrices.
1590 adapter: Optional architecture adapter for parameter key translation.
1592 Returns:
1593 Dict[str, torch.Tensor]: Fully processed state dict.
1594 """
1595 # Upcast to float32 for weight processing to avoid precision loss in
1596 # reduced-precision dtypes (bfloat16, float16). Operations like LayerNorm
1597 # folding involve multiplications that accumulate rounding errors when
1598 # performed in low precision.
1599 original_dtypes: Dict[str, torch.dtype] = {}
1600 for k, v in state_dict.items():
1601 if isinstance(v, torch.Tensor) and v.is_floating_point() and v.dtype != torch.float32: 1601 ↛ 1602line 1601 didn't jump to line 1602 because the condition on line 1601 was never true
1602 original_dtypes[k] = v.dtype
1603 state_dict[k] = v.float()
1605 # Skip fold_ln for adapters that don't support it (e.g., post-LN architectures
1606 # like BERT where LN placement means folding goes into the wrong sublayer).
1607 if fold_ln and adapter and not getattr(adapter, "supports_fold_ln", True):
1608 import warnings
1610 warnings.warn(
1611 f"{type(adapter).__name__} does not support fold_ln; norm weights "
1612 "stay unfolded and analyses assuming folded LN will be off.",
1613 UserWarning,
1614 )
1615 fold_ln = False
1616 if fold_ln:
1617 if getattr(cfg, "normalization_type", "LN") in ["LN", "LNPre"]:
1618 state_dict = ProcessWeights.fold_layer_norm(
1619 state_dict, cfg, fold_biases=True, center_weights=True, adapter=adapter
1620 )
1621 elif getattr(cfg, "normalization_type", "LN") in ["RMS", "RMSPre"]: 1621 ↛ 1632line 1621 didn't jump to line 1632 because the condition on line 1621 was always true
1622 state_dict = ProcessWeights.fold_layer_norm(
1623 state_dict, cfg, fold_biases=False, center_weights=False, adapter=adapter
1624 )
1625 # Note: Each folding function (_fold_layer for attention, _fold_mlp_layer_norm
1626 # for MLP) sets its own LN weights to 1.0 after successful folding.
1627 # We must NOT unconditionally set all LN weights to 1.0 here, because
1628 # models with combined QKV projections (e.g., OpenELM's qkv_proj) may
1629 # not be able to fold attention LN — setting ln1.w=1.0 without folding
1630 # destroys the RMS scaling.
1631 # Some adapters (e.g., post-LN) don't support center_writing_weights.
1632 if ( 1632 ↛ 1637line 1632 didn't jump to line 1637 because the condition on line 1632 was never true
1633 center_writing_weights
1634 and adapter
1635 and not getattr(adapter, "supports_center_writing_weights", True)
1636 ):
1637 import warnings
1639 warnings.warn(
1640 f"{type(adapter).__name__} does not support center_writing_weights; "
1641 "writing weights stay uncentered.",
1642 UserWarning,
1643 )
1644 center_writing_weights = False
1645 if center_writing_weights:
1646 if getattr(cfg, "normalization_type", "LN") in ["LN", "LNPre"] and (
1647 not getattr(cfg, "final_rms", False)
1648 ):
1649 state_dict = ProcessWeights.center_writing_weights(state_dict, cfg, adapter=adapter)
1650 if center_unembed:
1651 state_dict = ProcessWeights.center_unembed(state_dict, cfg=cfg, adapter=adapter)
1652 if fold_value_biases:
1653 state_dict = ProcessWeights.fold_value_biases(state_dict, cfg, adapter=adapter)
1654 if center_writing_weights and getattr(cfg, "normalization_type", "LN") in [
1655 "LN",
1656 "LNPre",
1657 ]:
1658 for layer_idx in range(cfg.n_layers):
1659 b_O_key = ProcessWeights._get_param_key(f"blocks.{layer_idx}.attn.b_O", adapter)
1660 if b_O_key in state_dict: 1660 ↛ 1658line 1660 didn't jump to line 1658 because the condition on line 1660 was always true
1661 b_O = ProcessWeights.convert_tensor_to_tl_format(
1662 b_O_key, state_dict, state_dict.get(b_O_key), cfg, adapter, layer_idx
1663 )
1664 assert b_O is not None, f"Attention b_O not found at key {b_O_key}"
1665 b_O = b_O - b_O.mean()
1666 state_dict[b_O_key] = ProcessWeights.convert_tensor_to_hf_format(
1667 b_O_key, b_O, cfg, adapter, layer_idx
1668 )
1669 if refactor_factored_attn_matrices:
1670 state_dict = ProcessWeights.refactor_factored_attn_matrices(
1671 state_dict, cfg, adapter=adapter
1672 )
1674 # Downcast back to original dtypes
1675 for k, orig_dtype in original_dtypes.items(): 1675 ↛ 1676line 1675 didn't jump to line 1676 because the loop on line 1675 never started
1676 if k in state_dict and isinstance(state_dict[k], torch.Tensor):
1677 state_dict[k] = state_dict[k].to(orig_dtype)
1679 return state_dict
1681 @staticmethod
1682 def refactor_factored_attn_matrices(
1683 state_dict: Dict[str, torch.Tensor], cfg, adapter=None
1684 ) -> Dict[str, torch.Tensor]:
1685 """Experimental method for managing queries, keys and values.
1687 As argued in [A Mathematical Framework for Transformer
1688 Circuits](https://transformer-circuits.pub/2021/framework/index.html), queries, keys and
1689 values are somewhat arbitrary intermediate terms when computing with the low rank factored
1690 matrices W_QK = W_Q @ W_K.T and W_OV = W_V @ W_O, and these matrices are the only thing
1691 determining head behaviour. But there are many ways to find a low rank factorization to a
1692 given matrix, and hopefully some of these are more interpretable than others! This method is
1693 one attempt, which makes all of the matrices have orthogonal rows or columns, W_O into a
1694 rotation and W_Q and W_K having the nth column in each having the same norm. The formula is
1695 $W_V = U @ S,W_O=Vh.T,W_Q=U@S.sqrt(),W_K=Vh@S.sqrt()$.
1697 More details:
1699 If W_OV = U @ S @ Vh.T in its singular value decomposition, (where S is in R^d_head not
1700 R^d_model, as W_OV is low rank), W_OV = (U @ S) @ (Vh.T) is an equivalent low rank
1701 factorisation, where rows/columns of each matrix are orthogonal! So setting $W_V=US$ and
1702 $W_O=Vh.T$ works just as well. I *think* this is a more interpretable setup, because now
1703 $W_O$ is just a rotation, and doesn't change the norm, so $z$ has the same norm as the
1704 result of the head.
1706 For $W_QK = W_Q @ W_K.T$ we use the refactor $W_Q = U @ S.sqrt()$ and $W_K = Vh @ S.sqrt()$,
1707 which is also equivalent ($S==S.sqrt() @ S.sqrt()$ as $S$ is diagonal). Here we keep the
1708 matrices as having the same norm, since there's not an obvious asymmetry between the keys
1709 and queries.
1711 Biases are more fiddly to deal with. For OV it's pretty easy - we just need (x @ W_V + b_V)
1712 @ 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
1713 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
1714 the head_index dimension too).
1716 For QK it's messy - we need to preserve the bilinear form of (x @ W_Q + b_Q) * (y @ W_K +
1717 b_K), which is fairly messy. To deal with the biases, we concatenate them to W_Q and W_K to
1718 simulate a d_model+1 dimensional input (whose final coordinate is always 1), do the SVD
1719 factorization on this effective matrix, then separate out into final weights and biases.
1721 Args:
1722 state_dict (Dict[str, torch.Tensor]): State dict of the model.
1723 cfg: Model configuration object.
1724 adapter: Optional architecture adapter for parameter key translation.
1726 Returns:
1727 Dict[str, torch.Tensor]: Modified state dict with refactored attention matrices.
1728 """
1729 # Make a deep copy to avoid modifying the original
1730 state_dict = {
1731 k: v.clone() if isinstance(v, torch.Tensor) else v for k, v in state_dict.items()
1732 }
1733 assert (
1734 getattr(cfg, "positional_embedding_type", "standard") != "rotary"
1735 ), "You can't refactor the QK circuit when using rotary embeddings (as the QK matrix depends on the position of the query and key)"
1737 for l in range(cfg.n_layers):
1738 W_Q_key = ProcessWeights._get_param_key(f"blocks.{l}.attn.W_Q", adapter)
1739 b_Q_key = ProcessWeights._get_param_key(f"blocks.{l}.attn.b_Q", adapter)
1740 W_K_key = ProcessWeights._get_param_key(f"blocks.{l}.attn.W_K", adapter)
1741 b_K_key = ProcessWeights._get_param_key(f"blocks.{l}.attn.b_K", adapter)
1742 W_V_key = ProcessWeights._get_param_key(f"blocks.{l}.attn.W_V", adapter)
1743 W_O_key = ProcessWeights._get_param_key(f"blocks.{l}.attn.W_O", adapter)
1744 b_V_key = ProcessWeights._get_param_key(f"blocks.{l}.attn.b_V", adapter)
1745 b_O_key = ProcessWeights._get_param_key(f"blocks.{l}.attn.b_O", adapter)
1747 # Skip hybrid layers without attention (other loops already guard individually)
1748 if W_Q_key not in state_dict:
1749 continue
1750 # If Q is present, K/V/O must be too
1751 for _required_key in [W_K_key, W_V_key, W_O_key]:
1752 if _required_key not in state_dict:
1753 raise ValueError(
1754 f"Inconsistent attention weights at layer {l}: "
1755 f"'{W_Q_key}' found but '{_required_key}' missing. "
1756 f"All of W_Q, W_K, W_V, W_O must be present together."
1757 )
1759 # W_QK = W_Q @ W_K.T
1760 # Concatenate biases to make a d_model+1 input dimension
1761 W_Q = ProcessWeights.convert_tensor_to_tl_format(
1762 W_Q_key, state_dict, state_dict.get(W_Q_key), cfg, adapter, l
1763 )
1764 b_Q = ProcessWeights.convert_tensor_to_tl_format(
1765 b_Q_key, state_dict, state_dict.get(b_Q_key), cfg, adapter, l
1766 )
1767 W_K = ProcessWeights.convert_tensor_to_tl_format(
1768 W_K_key, state_dict, state_dict.get(W_K_key), cfg, adapter, l
1769 )
1770 b_K = ProcessWeights.convert_tensor_to_tl_format(
1771 b_K_key, state_dict, state_dict.get(b_K_key), cfg, adapter, l
1772 )
1773 assert W_Q is not None, f"W_Q not found at key {W_Q_key}"
1774 assert b_Q is not None, f"b_Q not found at key {b_Q_key}"
1775 assert W_K is not None, f"W_K not found at key {W_K_key}"
1776 assert b_K is not None, f"b_K not found at key {b_K_key}"
1778 W_Q_eff = torch.cat([W_Q, b_Q[:, None, :]], dim=1)
1779 W_K_eff = torch.cat([W_K, b_K[:, None, :]], dim=1)
1781 W_Q_eff_even, W_K_eff_even_T = (
1782 FactoredMatrix(W_Q_eff, W_K_eff.transpose(-1, -2)).make_even().pair
1783 )
1784 W_K_eff_even = W_K_eff_even_T.transpose(-1, -2)
1786 state_dict[W_Q_key] = ProcessWeights.convert_tensor_to_hf_format(
1787 W_Q_key, W_Q_eff_even[:, :-1, :], cfg, adapter, l
1788 )
1789 state_dict[b_Q_key] = ProcessWeights.convert_tensor_to_hf_format(
1790 b_Q_key, W_Q_eff_even[:, -1, :], cfg, adapter, l
1791 )
1792 state_dict[W_K_key] = ProcessWeights.convert_tensor_to_hf_format(
1793 W_K_key, W_K_eff_even[:, :-1, :], cfg, adapter, l
1794 )
1795 state_dict[b_K_key] = ProcessWeights.convert_tensor_to_hf_format(
1796 b_K_key, W_K_eff_even[:, -1, :], cfg, adapter, l
1797 )
1799 # W_OV = W_V @ W_O
1800 W_V = ProcessWeights.convert_tensor_to_tl_format(
1801 W_V_key, state_dict, state_dict.get(W_V_key), cfg, adapter, l
1802 )
1803 W_O = ProcessWeights.convert_tensor_to_tl_format(
1804 W_O_key, state_dict, state_dict.get(W_O_key), cfg, adapter, l
1805 )
1807 # Factors the bias to be consistent.
1808 b_V = ProcessWeights.convert_tensor_to_tl_format(
1809 b_V_key, state_dict, state_dict.get(b_V_key), cfg, adapter, l
1810 )
1811 b_O = ProcessWeights.convert_tensor_to_tl_format(
1812 b_O_key, state_dict, state_dict.get(b_O_key), cfg, adapter, l
1813 )
1814 assert W_V is not None, f"W_V not found at key {W_V_key}"
1815 assert W_O is not None, f"W_O not found at key {W_O_key}"
1816 assert b_V is not None, f"b_V not found at key {b_V_key}"
1817 assert b_O is not None, f"b_O not found at key {b_O_key}"
1819 # Add singleton dimension for broadcasting
1820 b_V_expanded = einops.rearrange(b_V, "head_index d_head -> head_index d_head 1")
1822 b_V_times_W_O = b_V_expanded * W_O
1824 # Sum over d_head and head_index dimensions
1825 b_V_contribution = b_V_times_W_O.sum(1).sum(0)
1827 effective_bias = b_O + b_V_contribution
1828 state_dict[b_V_key] = ProcessWeights.convert_tensor_to_hf_format(
1829 b_V_key, torch.zeros_like(b_V), cfg, adapter, l
1830 )
1831 state_dict[b_O_key] = ProcessWeights.convert_tensor_to_hf_format(
1832 b_O_key, effective_bias, cfg, adapter, l
1833 )
1835 # Helper class to efficiently deal with low rank factored matrices.
1836 W_OV = FactoredMatrix(W_V, W_O)
1837 U, S, Vh = W_OV.svd()
1838 state_dict[W_V_key] = ProcessWeights.convert_tensor_to_hf_format(
1839 W_V_key, U @ S.diag_embed(), cfg, adapter, l
1840 )
1841 state_dict[W_O_key] = ProcessWeights.convert_tensor_to_hf_format(
1842 W_O_key, utils.transpose(Vh), cfg, adapter, l
1843 )
1845 return state_dict
1847 @overload
1848 @staticmethod
1849 def convert_tensor_to_tl_format(
1850 param_name: str,
1851 model_state_dict: Dict[str, torch.Tensor],
1852 tensor: torch.Tensor,
1853 cfg: Optional["TransformerLensConfig"],
1854 adapter: Optional["ArchitectureAdapter"] = None,
1855 layer_idx: Optional[int] = None,
1856 ) -> torch.Tensor:
1857 ...
1859 @overload
1860 @staticmethod
1861 def convert_tensor_to_tl_format(
1862 param_name: str,
1863 model_state_dict: Dict[str, torch.Tensor],
1864 tensor: None,
1865 cfg: Optional["TransformerLensConfig"],
1866 adapter: Optional["ArchitectureAdapter"] = None,
1867 layer_idx: Optional[int] = None,
1868 ) -> None:
1869 ...
1871 @staticmethod
1872 def convert_tensor_to_tl_format(
1873 param_name: str,
1874 model_state_dict: Dict[str, torch.Tensor],
1875 tensor: Optional[torch.Tensor],
1876 cfg: Optional["TransformerLensConfig"],
1877 adapter: Optional["ArchitectureAdapter"] = None,
1878 layer_idx: Optional[int] = None,
1879 ) -> Optional[torch.Tensor]:
1880 """Convert a tensor from its original format to TransformerLens format.
1882 Args:
1883 param_name: The parameter name in TransformerLens format (e.g., "blocks.0.attn.W_Q")
1884 model_state_dict: The model's state dictionary containing the actual tensors
1885 tensor: The tensor to convert, or None for optional parameters
1886 cfg: Model configuration
1887 adapter: Optional architecture adapter for component retrieval and key translation.
1888 If None, the tensor is returned unchanged.
1889 layer_idx: Layer index (required for layer-specific parameters)
1891 Returns:
1892 The tensor converted to TransformerLens format, or None if the parameter doesn't exist
1893 (which is valid for optional parameters like biases in models that don't use them).
1894 If adapter is None, returns the tensor unchanged.
1895 """
1896 # If no adapter provided, return tensor unchanged (handle None gracefully)
1897 if adapter is None:
1898 return tensor
1900 if ( 1900 ↛ 1981line 1900 didn't jump to line 1981 because the condition on line 1900 was always true
1901 hasattr(adapter, "weight_processing_conversions")
1902 and adapter.weight_processing_conversions is not None
1903 ):
1904 # Create placeholder param name by replacing layer index with {i}
1905 placeholder_param_name = param_name
1906 if "blocks." in param_name:
1907 placeholder_param_name = re.sub(r"blocks\.(\d+)\.", "blocks.{i}.", param_name)
1909 # Check if we have a conversion for this parameter.
1910 # Try exact match first, then strip .weight suffix for adapters
1911 # that define conversions without the suffix (e.g. Pythia's "blocks.{i}.attn.q").
1912 # NOTE: Only strip .weight, NOT .bias — stripping .bias would incorrectly
1913 # match bias keys against weight conversions (e.g. "blocks.{i}.attn.q.bias"
1914 # would match the weight conversion for "blocks.{i}.attn.q").
1915 matched_key = None
1916 if placeholder_param_name in adapter.weight_processing_conversions:
1917 matched_key = placeholder_param_name
1918 elif placeholder_param_name.endswith(".weight"):
1919 stripped = placeholder_param_name[: -len(".weight")]
1920 if stripped in adapter.weight_processing_conversions:
1921 matched_key = stripped
1923 if matched_key is not None:
1924 param_conversion = adapter.weight_processing_conversions[matched_key]
1926 # Handle both ParamProcessingConversion objects and legacy string mappings
1927 if isinstance(param_conversion, str): 1927 ↛ 1930line 1927 didn't jump to line 1930 because the condition on line 1927 was never true
1928 # Legacy string mapping - just return the tensor as-is
1929 # (string mappings are handled elsewhere in the architecture adapter)
1930 return tensor
1931 else:
1932 # Skip conversion for optional parameters that don't exist (e.g. biases)
1933 if tensor is None and param_name not in model_state_dict: 1933 ↛ 1934line 1933 didn't jump to line 1934 because the condition on line 1933 was never true
1934 return None
1935 # Try ParamProcessingConversion.convert() first (uses source_key
1936 # to fetch from state dict — needed for split conversions like
1937 # GPT-2's QKV). If source_key resolves to a missing key and we
1938 # already have the tensor, fall back to applying the tensor
1939 # conversion directly (needed for adapters like GPT-Neo whose
1940 # source_key references HF keys not in the bridge state dict).
1941 if (
1942 hasattr(param_conversion, "source_key")
1943 and param_conversion.source_key is not None
1944 ):
1945 resolved_key = param_conversion._resolve_key(
1946 param_name, param_conversion.source_key
1947 )
1948 if resolved_key not in model_state_dict and tensor is not None: 1948 ↛ 1975line 1948 didn't jump to line 1975 because the condition on line 1948 was always true
1949 # Source key not in state dict — the tensor is already in
1950 # bridge format (e.g. already split from combined QKV).
1951 # If the conversion is a ChainTensorConversion that includes
1952 # a SplitTensorConversion, skip the split step since
1953 # it was already applied during bridge construction.
1954 from transformer_lens.conversion_utils.conversion_steps.chain_tensor_conversion import (
1955 ChainTensorConversion,
1956 )
1957 from transformer_lens.conversion_utils.conversion_steps.split_tensor_conversion import (
1958 SplitTensorConversion,
1959 )
1961 tc = param_conversion.tensor_conversion
1962 if isinstance(tc, ChainTensorConversion): 1962 ↛ 1963line 1962 didn't jump to line 1963 because the condition on line 1962 was never true
1963 non_split = [
1964 c
1965 for c in tc.conversions
1966 if not isinstance(c, SplitTensorConversion)
1967 ]
1968 if len(non_split) < len(tc.conversions):
1969 # Apply only the non-split conversions
1970 result = tensor
1971 for conv in non_split:
1972 result = conv.handle_conversion(result, model_state_dict)
1973 return result
1974 return tc.convert(tensor, model_state_dict)
1975 return param_conversion.convert(model_state_dict, param_name)
1976 else:
1977 # No conversion defined, return tensor as-is (may be None for optional params)
1978 return tensor
1979 else:
1980 # No conversions defined, return tensor as-is (may be None for optional params)
1981 return tensor
1983 @overload
1984 @staticmethod
1985 def convert_tensor_to_hf_format(
1986 param_name: str,
1987 tensor: torch.Tensor,
1988 cfg: Optional["TransformerLensConfig"],
1989 adapter: Optional["ArchitectureAdapter"] = None,
1990 layer_idx: Optional[int] = None,
1991 ) -> torch.Tensor:
1992 ...
1994 @overload
1995 @staticmethod
1996 def convert_tensor_to_hf_format(
1997 param_name: str,
1998 tensor: None,
1999 cfg: Optional["TransformerLensConfig"],
2000 adapter: Optional["ArchitectureAdapter"] = None,
2001 layer_idx: Optional[int] = None,
2002 ) -> None:
2003 ...
2005 @staticmethod
2006 def convert_tensor_to_hf_format(
2007 param_name: str,
2008 tensor: Optional[torch.Tensor],
2009 cfg: Optional["TransformerLensConfig"],
2010 adapter: Optional["ArchitectureAdapter"] = None,
2011 layer_idx: Optional[int] = None,
2012 ) -> Optional[torch.Tensor]:
2013 """Convert a tensor from TransformerLens format back to its original format.
2015 Args:
2016 param_name: The parameter name in TransformerLens format (e.g., "blocks.0.attn.W_Q")
2017 tensor: The tensor to convert (in TransformerLens format), or None if parameter is optional
2018 cfg: Model configuration
2019 adapter: Optional architecture adapter for component retrieval and key translation.
2020 If None, the tensor is returned unchanged.
2021 layer_idx: Layer index (required for layer-specific parameters)
2023 Returns:
2024 The tensor converted back to original format, or None if tensor was None.
2025 If adapter is None, returns the tensor unchanged.
2026 """
2027 # Handle None tensors (optional parameters)
2028 if tensor is None: 2028 ↛ 2029line 2028 didn't jump to line 2029 because the condition on line 2028 was never true
2029 return None
2031 # If no adapter provided, return tensor unchanged
2032 if adapter is None:
2033 return tensor
2035 if ( 2035 ↛ 2087line 2035 didn't jump to line 2087 because the condition on line 2035 was always true
2036 hasattr(adapter, "weight_processing_conversions")
2037 and adapter.weight_processing_conversions is not None
2038 ):
2039 # Create placeholder param name by replacing layer index with {i}
2040 placeholder_param_name = param_name
2041 if "blocks." in param_name:
2042 placeholder_param_name = re.sub(r"blocks\.(\d+)\.", "blocks.{i}.", param_name)
2044 # Check if we have a conversion for this parameter.
2045 # Try exact match first, then strip .weight suffix (not .bias — see convert_tensor_to_tl_format).
2046 matched_key = None
2047 if placeholder_param_name in adapter.weight_processing_conversions:
2048 matched_key = placeholder_param_name
2049 elif placeholder_param_name.endswith(".weight"):
2050 stripped = placeholder_param_name[: -len(".weight")]
2051 if stripped in adapter.weight_processing_conversions:
2052 matched_key = stripped
2054 if matched_key is not None:
2055 param_conversion = adapter.weight_processing_conversions[matched_key]
2057 # Handle both ParamProcessingConversion objects and legacy string mappings
2058 if isinstance(param_conversion, str): 2058 ↛ 2060line 2058 didn't jump to line 2060 because the condition on line 2058 was never true
2059 # Legacy string mapping - just return the tensor as-is
2060 return tensor
2061 else:
2062 # Revert the conversion. For ChainTensorConversions that include
2063 # SplitTensorConversion, skip the split revert step (which is a
2064 # no-op anyway) to match the forward conversion path.
2065 from transformer_lens.conversion_utils.conversion_steps.chain_tensor_conversion import (
2066 ChainTensorConversion,
2067 )
2068 from transformer_lens.conversion_utils.conversion_steps.split_tensor_conversion import (
2069 SplitTensorConversion,
2070 )
2072 tc = param_conversion.tensor_conversion
2073 if isinstance(tc, ChainTensorConversion): 2073 ↛ 2074line 2073 didn't jump to line 2074 because the condition on line 2073 was never true
2074 non_split = [
2075 c for c in tc.conversions if not isinstance(c, SplitTensorConversion)
2076 ]
2077 if len(non_split) < len(tc.conversions):
2078 # Revert only the non-split conversions in reverse order
2079 result = tensor
2080 for conv in reversed(non_split):
2081 result = conv.revert(result)
2082 return result
2083 return param_conversion.revert(tensor)
2084 else:
2085 return tensor
2086 else:
2087 return tensor
2089 @staticmethod
2090 def distribute_weights_to_components(
2091 state_dict: Dict[str, torch.Tensor],
2092 component_mapping: Dict[str, Any],
2093 verbose: bool = False,
2094 ) -> None:
2095 """Distribute processed weights from state_dict to generalized components.
2097 This function loops through the component_mapping and extracts relevant weights
2098 for each component using filter_dict_by_prefix, then calls set_processed_weights
2099 on each component. For list components (like blocks), it determines the number
2100 of items and distributes weights to each indexed component.
2102 Args:
2103 state_dict: Dictionary of processed weights in MODERN TransformerLens format
2104 (e.g., blocks.0.attn.q.weight, not transformer.h.0.attn.q.weight)
2105 component_mapping: Dictionary (real_components) mapping TL keys to tuples of
2106 (remote_path, component_instance), where component_instance can be either
2107 a single component or a list of components
2108 verbose: If True, print detailed information about weight distribution
2110 Example:
2111 For a real_components mapping like:
2112 {
2113 "embed": ("transformer.wte", <EmbeddingBridge instance>),
2114 "blocks": ("transformer.h", [<BlockBridge 0>, <BlockBridge 1>, ...]),
2115 "unembed": ("lm_head", <UnembeddingBridge instance>)
2116 }
2118 With modern TL keys in state_dict like "embed.weight", "blocks.0.attn.q.weight":
2119 1. Extract weights starting with "embed" and pass to embed component
2120 2. For blocks, extract all "blocks.*" weights, determine the number of blocks,
2121 then for each block index, extract weights for that specific block
2122 3. Extract "unembed" weights and pass to unembed component
2123 """
2124 if verbose: 2124 ↛ 2125line 2124 didn't jump to line 2125 because the condition on line 2124 was never true
2125 print(f"\n{'='*80}")
2126 print(f"distribute_weights_to_components: Starting weight distribution")
2127 print(f"State dict has {len(state_dict)} keys")
2128 print(f"Component mapping has {len(component_mapping)} components")
2129 print(f"{'='*80}\n")
2131 for component_name, component_tuple in component_mapping.items():
2132 # component_mapping is real_components format: (remote_path, instance)
2133 # instance can be either a single component or a list of components
2134 if not isinstance(component_tuple, tuple): 2134 ↛ 2135line 2134 didn't jump to line 2135 because the condition on line 2134 was never true
2135 raise ValueError(
2136 f"Expected tuple for component '{component_name}' in real_components, "
2137 f"but got {type(component_tuple).__name__}: {component_tuple}"
2138 )
2139 remote_key, component = component_tuple
2140 is_list = isinstance(component, list)
2142 # Use the component_name (TL format) as prefix instead of remote_key (HF format)
2143 # since state_dict now has modern TL keys
2144 tl_prefix = component_name
2146 if verbose: 2146 ↛ 2147line 2146 didn't jump to line 2147 because the condition on line 2146 was never true
2147 print(f"\nProcessing component: {component_name}")
2148 print(f" Remote key (HF): {remote_key}")
2149 print(f" TL prefix: {tl_prefix}")
2150 print(f" Is list: {is_list}")
2152 if is_list:
2153 # This is a list component like "blocks"
2154 # Extract all weights that start with this prefix
2155 all_list_weights = filter_dict_by_prefix(state_dict, tl_prefix)
2157 if verbose: 2157 ↛ 2158line 2157 didn't jump to line 2158 because the condition on line 2157 was never true
2158 print(f" Found {len(all_list_weights)} weights for list component")
2159 print(f" List has {len(component)} instances")
2161 # Component is a list of actual instances
2162 for i, instance in enumerate(component):
2163 # Extract weights for this specific index
2164 # This will get keys like "0.attn.q.weight" and strip the "0." to get "attn.q.weight"
2165 indexed_weights = filter_dict_by_prefix(all_list_weights, str(i))
2167 if verbose: 2167 ↛ 2168line 2167 didn't jump to line 2168 because the condition on line 2167 was never true
2168 print(f" Instance {i}: Found {len(indexed_weights)} weights")
2169 for key in indexed_weights.keys():
2170 print(f" - {key}")
2172 # Skip if no weights found for this component (e.g., Q/K/V Linear sub-components
2173 # that get their weights from parent JointQKVAttentionBridge)
2174 if len(indexed_weights) == 0: 2174 ↛ 2175line 2174 didn't jump to line 2175 because the condition on line 2174 was never true
2175 if verbose:
2176 print(f" Skipping instance {i} - no weights found")
2177 continue
2179 instance.set_processed_weights(indexed_weights, verbose=verbose)
2180 else:
2181 # This is a single component (not a list)
2182 component_weights = filter_dict_by_prefix(state_dict, tl_prefix)
2184 if verbose: 2184 ↛ 2185line 2184 didn't jump to line 2185 because the condition on line 2184 was never true
2185 print(f" Found {len(component_weights)} weights for single component")
2186 for key in component_weights.keys():
2187 print(f" - {key}")
2189 # Skip if no weights found for this component
2190 if len(component_weights) == 0:
2191 if verbose: 2191 ↛ 2192line 2191 didn't jump to line 2192 because the condition on line 2191 was never true
2192 print(f" Skipping component - no weights found")
2193 continue
2195 component.set_processed_weights(component_weights, verbose=verbose)