Coverage for transformer_lens/model_bridge/get_params_util.py: 98%
120 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-09-01 16:23 +0000
1"""Utility function for getting model parameters in TransformerLens format."""
2import logging
3from typing import Dict, Optional
5import torch
7logger = logging.getLogger(__name__)
10def _tensor_attr(obj, *names: str) -> Optional[torch.Tensor]:
11 """First attribute of ``obj`` among ``names`` that is an actual tensor, else None.
13 NotImplementedError counts as absent: MLA attention raises it from W_Q/W_K/W_V/W_O
14 (compressed projections have no standard per-head form).
15 """
16 for name in names:
17 try:
18 value = getattr(obj, name)
19 except (AttributeError, TypeError, NotImplementedError):
20 continue
21 if isinstance(value, torch.Tensor):
22 return value
23 return None
26def get_bridge_params(bridge) -> Dict[str, torch.Tensor]:
27 """Model parameters in SVDInterpreter format.
29 Reads the bridge components' TL-layout weight properties (``W_Q``,
30 ``W_in``, ...), which already account for layout conversion and weight
31 processing. For missing weights, returns zero tensors of appropriate shape
32 instead of raising exceptions. Skips attn keys for non-attention layers.
33 LayerNorm params (``blocks.{i}.ln1.w`` etc.) are included when the modules
34 still carry them (i.e. before folding) so consumers can detect fold state.
36 Returns:
37 dict: Dictionary of parameter tensors with TransformerLens naming convention
39 Raises:
40 ValueError: If configuration is inconsistent (e.g., cfg.n_layers != len(blocks))
41 """
42 cfg = bridge.cfg
43 params_dict: Dict[str, torch.Tensor] = {}
45 def _get_device_dtype():
46 """Infer device/dtype from the first available model parameter."""
47 device = getattr(cfg, "device", None) or torch.device("cpu")
48 dtype = torch.float32
49 try:
50 first_param = next(bridge.parameters())
51 device = first_param.device
52 dtype = first_param.dtype
53 except (StopIteration, TypeError, AttributeError):
54 pass
55 return (device, dtype)
57 def _zeros(*shape) -> torch.Tensor:
58 device, dtype = _get_device_dtype()
59 return torch.zeros(*shape, device=device, dtype=dtype)
61 embed = _tensor_attr(getattr(bridge, "embed", None), "W_E", "weight")
62 params_dict["embed.W_E"] = embed if embed is not None else _zeros(cfg.d_vocab, cfg.d_model)
64 pos = _tensor_attr(getattr(bridge, "pos_embed", None), "W_pos", "weight")
65 params_dict["pos_embed.W_pos"] = pos if pos is not None else _zeros(cfg.n_ctx, cfg.d_model)
67 for layer_idx in range(cfg.n_layers):
68 if layer_idx >= len(bridge.blocks):
69 raise ValueError(
70 f"Configuration mismatch: cfg.n_layers={cfg.n_layers} but only "
71 f"{len(bridge.blocks)} blocks found. Layer {layer_idx} does not exist."
72 )
73 block = bridge.blocks[layer_idx]
75 # Skip non-attention layers entirely (no zero-fill — prevents SVDInterpreter garbage)
76 try:
77 has_attn = "attn" in block._modules
78 except (TypeError, AttributeError):
79 has_attn = hasattr(block, "attn") # Mock fallback
80 if has_attn:
81 attn = block.attn
82 w_q = _tensor_attr(attn, "W_Q")
83 w_k = _tensor_attr(attn, "W_K")
84 w_v = _tensor_attr(attn, "W_V")
85 w_o = _tensor_attr(attn, "W_O")
86 if w_q is None or w_k is None or w_v is None or w_o is None:
87 logger.debug(
88 "Block %d has 'attn' but no TL-layout W_Q/W_K/W_V/W_O properties — "
89 "skipping attention weights for this layer",
90 layer_idx,
91 )
92 else:
93 # GQA: expand grouped K/V (and their biases below) to n_heads so
94 # per-head pairings like SVDInterpreter's OV = W_V[h] @ W_O[h]
95 # line up — the legacy HT convention repeat_interleaved these.
96 n_kv_heads = w_k.shape[0]
97 if w_k.ndim == 3 and 0 < n_kv_heads < cfg.n_heads:
98 if cfg.n_heads % n_kv_heads != 0: 98 ↛ 99line 98 didn't jump to line 99 because the condition on line 98 was never true
99 raise ValueError(
100 f"blocks.{layer_idx}.attn: n_heads ({cfg.n_heads}) is not "
101 f"divisible by n_kv_heads ({n_kv_heads}); cannot expand "
102 "grouped K/V to per-query heads."
103 )
104 repeats = cfg.n_heads // n_kv_heads
105 w_k = torch.repeat_interleave(w_k, repeats, dim=0)
106 w_v = torch.repeat_interleave(w_v, repeats, dim=0)
107 params_dict[f"blocks.{layer_idx}.attn.W_Q"] = w_q
108 params_dict[f"blocks.{layer_idx}.attn.W_K"] = w_k
109 params_dict[f"blocks.{layer_idx}.attn.W_V"] = w_v
110 params_dict[f"blocks.{layer_idx}.attn.W_O"] = w_o
111 for bias_name in ("b_Q", "b_K", "b_V"):
112 bias = _tensor_attr(attn, bias_name)
113 if bias is None:
114 bias = _zeros(cfg.n_heads, cfg.d_head)
115 elif bias.ndim == 2 and 0 < bias.shape[0] < cfg.n_heads:
116 bias = torch.repeat_interleave(bias, cfg.n_heads // bias.shape[0], dim=0)
117 params_dict[f"blocks.{layer_idx}.attn.{bias_name}"] = bias
118 b_O = _tensor_attr(attn, "b_O")
119 params_dict[f"blocks.{layer_idx}.attn.b_O"] = (
120 b_O if b_O is not None else _zeros(cfg.d_model)
121 )
123 d_mlp = cfg.d_mlp if cfg.d_mlp is not None else 4 * cfg.d_model
124 mlp = getattr(block, "mlp", None)
125 w_in = _tensor_attr(mlp, "W_in")
126 if w_in is None:
127 if mlp is not None: 127 ↛ 138line 127 didn't jump to line 138 because the condition on line 127 was always true
128 # Zero-filling a real MLP silently yields wrong numbers downstream
129 # (SVD/weight analyses decompose zeros). Say so — the fill stays for
130 # architectures that genuinely have no MLP under this name.
131 logger.warning(
132 "Block %d MLP weights could not be extracted — emitting ZEROS "
133 "for blocks.%d.mlp.W_in/W_out/b_in/b_out. Any weight-space "
134 "analysis of this layer will be meaningless.",
135 layer_idx,
136 layer_idx,
137 )
138 params_dict[f"blocks.{layer_idx}.mlp.W_in"] = _zeros(cfg.d_model, d_mlp)
139 params_dict[f"blocks.{layer_idx}.mlp.W_out"] = _zeros(d_mlp, cfg.d_model)
140 params_dict[f"blocks.{layer_idx}.mlp.b_in"] = _zeros(d_mlp)
141 params_dict[f"blocks.{layer_idx}.mlp.b_out"] = _zeros(cfg.d_model)
142 else:
143 params_dict[f"blocks.{layer_idx}.mlp.W_in"] = w_in
144 w_out = _tensor_attr(mlp, "W_out")
145 params_dict[f"blocks.{layer_idx}.mlp.W_out"] = (
146 w_out if w_out is not None else _zeros(d_mlp, cfg.d_model)
147 )
148 b_in = _tensor_attr(mlp, "b_in")
149 params_dict[f"blocks.{layer_idx}.mlp.b_in"] = (
150 b_in if b_in is not None else _zeros(d_mlp)
151 )
152 b_out = _tensor_attr(mlp, "b_out")
153 params_dict[f"blocks.{layer_idx}.mlp.b_out"] = (
154 b_out if b_out is not None else _zeros(cfg.d_model)
155 )
156 w_gate = _tensor_attr(mlp, "W_gate")
157 # Raw-attribute fallback is for plain gated MLPs only: `gate` on an
158 # interleaved-MoE component (anything exposing bound_dense) is the
159 # sparse layers' ROUTER, never a gate projection.
160 is_moe = getattr(type(mlp), "bound_dense", None) is not None
161 if w_gate is None and not is_moe:
162 w_gate = _tensor_attr(getattr(mlp, "gate", None), "weight")
163 if w_gate is not None:
164 params_dict[f"blocks.{layer_idx}.mlp.W_gate"] = w_gate
165 b_gate = _tensor_attr(mlp, "b_gate")
166 if b_gate is None and not is_moe:
167 b_gate = _tensor_attr(getattr(mlp, "gate", None), "bias")
168 if b_gate is not None:
169 params_dict[f"blocks.{layer_idx}.mlp.b_gate"] = b_gate
171 # LN params (present pre-folding; folded models carry identities or none).
172 for ln_name in ("ln1", "ln2"):
173 ln = getattr(block, ln_name, None)
174 ln_w = _tensor_attr(ln, "w", "weight")
175 if ln_w is not None:
176 params_dict[f"blocks.{layer_idx}.{ln_name}.w"] = ln_w
177 ln_b = _tensor_attr(ln, "b", "bias")
178 if ln_b is not None:
179 params_dict[f"blocks.{layer_idx}.{ln_name}.b"] = ln_b
181 ln_final_w = _tensor_attr(getattr(bridge, "ln_final", None), "w", "weight")
182 if ln_final_w is not None:
183 params_dict["ln_final.w"] = ln_final_w
184 ln_final_b = _tensor_attr(getattr(bridge, "ln_final", None), "b", "bias")
185 if ln_final_b is not None:
186 params_dict["ln_final.b"] = ln_final_b
188 unembed = getattr(bridge, "unembed", None)
189 w_u = _tensor_attr(unembed, "W_U")
190 if w_u is None:
191 raw = _tensor_attr(unembed, "weight")
192 w_u = raw.T if raw is not None else _zeros(cfg.d_model, cfg.d_vocab)
193 params_dict["unembed.W_U"] = w_u
194 b_u = _tensor_attr(unembed, "b_U")
195 params_dict["unembed.b_U"] = b_u if b_u is not None else _zeros(cfg.d_vocab)
197 return params_dict