Coverage for transformer_lens/model_bridge/get_params_util.py: 96%
124 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-09-21 19:27 +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 absent optional weights, returns zero tensors of appropriate
32 shape instead of raising exceptions. Real components whose weights cannot be
33 represented in the dense TL layout are omitted. Skips attn keys for
34 non-attention layers, and
35 omits ``pos_embed.W_pos`` for rotary models (which have no learned position
36 table), matching HookedTransformer's parameter set.
37 LayerNorm params (``blocks.{i}.ln1.w`` etc.) are included when the modules
38 still carry them (i.e. before folding) so consumers can detect fold state.
40 Returns:
41 dict: Dictionary of parameter tensors with TransformerLens naming convention
43 Raises:
44 ValueError: If configuration is inconsistent (e.g., cfg.n_layers != len(blocks))
45 """
46 cfg = bridge.cfg
47 params_dict: Dict[str, torch.Tensor] = {}
49 def _get_device_dtype():
50 """Infer device/dtype from the first available model parameter."""
51 device = getattr(cfg, "device", None) or torch.device("cpu")
52 dtype = torch.float32
53 try:
54 first_param = next(bridge.parameters())
55 device = first_param.device
56 dtype = first_param.dtype
57 except (StopIteration, TypeError, AttributeError):
58 pass
59 return (device, dtype)
61 def _zeros(*shape) -> torch.Tensor:
62 device, dtype = _get_device_dtype()
63 return torch.zeros(*shape, device=device, dtype=dtype)
65 embed = _tensor_attr(getattr(bridge, "embed", None), "W_E", "weight")
66 params_dict["embed.W_E"] = embed if embed is not None else _zeros(cfg.d_vocab, cfg.d_model)
68 pos = _tensor_attr(getattr(bridge, "pos_embed", None), "W_pos", "weight")
69 if pos is not None:
70 params_dict["pos_embed.W_pos"] = pos
71 elif getattr(cfg, "positional_embedding_type", "standard") != "rotary":
72 # Rotary models have no learned position table, and HookedTransformer never
73 # registers pos_embed for them — fabricating one would advertise a weight the
74 # architecture lacks and allocate n_ctx x d_model (GBs on an 8B model).
75 params_dict["pos_embed.W_pos"] = _zeros(cfg.n_ctx, cfg.d_model)
77 for layer_idx in range(cfg.n_layers):
78 if layer_idx >= len(bridge.blocks):
79 raise ValueError(
80 f"Configuration mismatch: cfg.n_layers={cfg.n_layers} but only "
81 f"{len(bridge.blocks)} blocks found. Layer {layer_idx} does not exist."
82 )
83 block = bridge.blocks[layer_idx]
85 # Skip non-attention layers entirely (no zero-fill — prevents SVDInterpreter garbage)
86 try:
87 has_attn = "attn" in block._modules
88 except (TypeError, AttributeError):
89 has_attn = hasattr(block, "attn") # Mock fallback
90 if has_attn:
91 attn = block.attn
92 w_q = _tensor_attr(attn, "W_Q")
93 w_k = _tensor_attr(attn, "W_K")
94 w_v = _tensor_attr(attn, "W_V")
95 w_o = _tensor_attr(attn, "W_O")
96 if w_q is None or w_k is None or w_v is None or w_o is None:
97 logger.debug(
98 "Block %d has 'attn' but no TL-layout W_Q/W_K/W_V/W_O properties — "
99 "skipping attention weights for this layer",
100 layer_idx,
101 )
102 else:
103 # GQA: expand grouped K/V (and their biases below) to n_heads so
104 # per-head pairings like SVDInterpreter's OV = W_V[h] @ W_O[h]
105 # line up — the legacy HT convention repeat_interleaved these.
106 n_kv_heads = w_k.shape[0]
107 if w_k.ndim == 3 and 0 < n_kv_heads < cfg.n_heads:
108 if cfg.n_heads % n_kv_heads != 0: 108 ↛ 109line 108 didn't jump to line 109 because the condition on line 108 was never true
109 raise ValueError(
110 f"blocks.{layer_idx}.attn: n_heads ({cfg.n_heads}) is not "
111 f"divisible by n_kv_heads ({n_kv_heads}); cannot expand "
112 "grouped K/V to per-query heads."
113 )
114 repeats = cfg.n_heads // n_kv_heads
115 w_k = torch.repeat_interleave(w_k, repeats, dim=0)
116 w_v = torch.repeat_interleave(w_v, repeats, dim=0)
117 params_dict[f"blocks.{layer_idx}.attn.W_Q"] = w_q
118 params_dict[f"blocks.{layer_idx}.attn.W_K"] = w_k
119 params_dict[f"blocks.{layer_idx}.attn.W_V"] = w_v
120 params_dict[f"blocks.{layer_idx}.attn.W_O"] = w_o
121 for bias_name in ("b_Q", "b_K", "b_V"):
122 bias = _tensor_attr(attn, bias_name)
123 if bias is None:
124 bias = _zeros(cfg.n_heads, cfg.d_head)
125 elif bias.ndim == 2 and 0 < bias.shape[0] < cfg.n_heads:
126 bias = torch.repeat_interleave(bias, cfg.n_heads // bias.shape[0], dim=0)
127 params_dict[f"blocks.{layer_idx}.attn.{bias_name}"] = bias
128 b_O = _tensor_attr(attn, "b_O")
129 params_dict[f"blocks.{layer_idx}.attn.b_O"] = (
130 b_O if b_O is not None else _zeros(cfg.d_model)
131 )
133 d_mlp = cfg.d_mlp if cfg.d_mlp is not None else 4 * cfg.d_model
134 mlp = getattr(block, "mlp", None)
135 w_in = _tensor_attr(mlp, "W_in")
136 if w_in is None:
137 if mlp is None: 137 ↛ 138line 137 didn't jump to line 138 because the condition on line 137 was never true
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 logger.warning(
144 "Block %d MLP does not expose a single dense W_in/W_out; "
145 "omitting its MLP parameters from the TL-style parameter dictionary.",
146 layer_idx,
147 )
148 else:
149 params_dict[f"blocks.{layer_idx}.mlp.W_in"] = w_in
150 w_out = _tensor_attr(mlp, "W_out")
151 if w_out is not None:
152 params_dict[f"blocks.{layer_idx}.mlp.W_out"] = w_out
153 b_in = _tensor_attr(mlp, "b_in")
154 params_dict[f"blocks.{layer_idx}.mlp.b_in"] = (
155 b_in if b_in is not None else _zeros(d_mlp)
156 )
157 b_out = _tensor_attr(mlp, "b_out")
158 params_dict[f"blocks.{layer_idx}.mlp.b_out"] = (
159 b_out if b_out is not None else _zeros(cfg.d_model)
160 )
161 w_gate = _tensor_attr(mlp, "W_gate")
162 # Raw-attribute fallback is for plain gated MLPs only: `gate` on an
163 # interleaved-MoE component (anything exposing bound_dense) is the
164 # sparse layers' ROUTER, never a gate projection.
165 is_moe = getattr(type(mlp), "bound_dense", None) is not None
166 if w_gate is None and not is_moe:
167 w_gate = _tensor_attr(getattr(mlp, "gate", None), "weight")
168 if w_gate is not None:
169 params_dict[f"blocks.{layer_idx}.mlp.W_gate"] = w_gate
170 b_gate = _tensor_attr(mlp, "b_gate")
171 if b_gate is None and not is_moe:
172 b_gate = _tensor_attr(getattr(mlp, "gate", None), "bias")
173 if b_gate is not None:
174 params_dict[f"blocks.{layer_idx}.mlp.b_gate"] = b_gate
176 # LN params (present pre-folding; folded models carry identities or none).
177 for ln_name in ("ln1", "ln2"):
178 ln = getattr(block, ln_name, None)
179 ln_w = _tensor_attr(ln, "w", "weight")
180 if ln_w is not None:
181 params_dict[f"blocks.{layer_idx}.{ln_name}.w"] = ln_w
182 ln_b = _tensor_attr(ln, "b", "bias")
183 if ln_b is not None:
184 params_dict[f"blocks.{layer_idx}.{ln_name}.b"] = ln_b
186 ln_final_w = _tensor_attr(getattr(bridge, "ln_final", None), "w", "weight")
187 if ln_final_w is not None:
188 params_dict["ln_final.w"] = ln_final_w
189 ln_final_b = _tensor_attr(getattr(bridge, "ln_final", None), "b", "bias")
190 if ln_final_b is not None:
191 params_dict["ln_final.b"] = ln_final_b
193 unembed = getattr(bridge, "unembed", None)
194 w_u = _tensor_attr(unembed, "W_U")
195 if w_u is None:
196 raw = _tensor_attr(unembed, "weight")
197 w_u = raw.T if raw is not None else _zeros(cfg.d_model, cfg.d_vocab)
198 params_dict["unembed.W_U"] = w_u
199 b_u = _tensor_attr(unembed, "b_U")
200 params_dict["unembed.b_U"] = b_u if b_u is not None else _zeros(cfg.d_vocab)
202 return params_dict