Coverage for transformer_lens/model_bridge/architecture_adapter.py: 75%
504 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-09-01 16:23 +0000
1"""Architecture adapter base class.
3This module contains the base class for architecture adapters that map between different model architectures.
4"""
5from typing import Any, Dict, Optional, cast
7import einops
8import torch
10from transformer_lens.config import TransformerBridgeConfig
11from transformer_lens.conversion_utils.conversion_steps.rearrange_tensor_conversion import (
12 RearrangeTensorConversion,
13)
14from transformer_lens.conversion_utils.param_processing_conversion import (
15 ParamProcessingConversion,
16)
17from transformer_lens.model_bridge.generalized_components.base import (
18 GeneralizedComponent,
19)
20from transformer_lens.model_bridge.generalized_components.gated_mlp import (
21 GatedMLPBridge,
22)
23from transformer_lens.model_bridge.generalized_components.linear import LinearBridge
24from transformer_lens.model_bridge.generalized_components.mlp import MLPBridge
25from transformer_lens.model_bridge.types import (
26 ComponentMapping,
27 RemoteComponent,
28 RemoteModel,
29 RemotePath,
30 TransformerLensPath,
31)
32from transformer_lens.utilities.activation_functions import apply_softcap
35class ArchitectureAdapter:
36 """Base class for architecture adapters.
38 This class provides the interface for adapting between different model architectures.
39 It handles both component mapping (for accessing model parts) and weight conversion
40 (for initializing weights from one format to another).
41 """
43 default_cfg: dict[str, Any] = {}
45 # verify_models phase applicability. Architectures that cannot participate
46 # in specific phases (e.g. SSMs don't have the transformer-shaped hooks/
47 # weights the benchmark phases assume) should override. An empty list
48 # means "skip verify_models entirely; verification lives in integration
49 # tests."
50 applicable_phases: list[int] = [1, 2, 3, 4]
52 # Whether this architecture supports text generation via generate().
53 # Encoder-only models (e.g. BERT, HuBERT) should set this to False.
54 supports_generation: bool = True
56 # Whether the wrapped forward speaks the HF KV-cache protocol. False for
57 # recurrent/convolutional decoders (RWKV's bespoke `state`, HyenaDNA's
58 # kwarg-less FFT forward): generation then recomputes the full prefix each
59 # step, which is exact but O(n^2).
60 supports_kv_cache: bool = True
62 # Whether batched generation is sound. False where the wrapped forward
63 # ignores attention_mask (RWKV) or rejects it (HyenaDNA), so left-padding
64 # would silently corrupt results rather than be masked out.
65 supports_batched_generation: bool = True
67 # Name of a native non-autoregressive sampler on the wrapped model, exposed
68 # via bridge.diffusion_generate() (e.g. Dream's "diffusion_generate",
69 # LLaDA2's/Gidd's "generate"). None means the architecture has none.
70 native_sampler: Optional[str] = None
72 def native_sampler_kwargs(self, max_new_tokens: int, prompt_len: int) -> Dict[str, Any]:
73 """Map a token budget onto the native sampler's own parameter names
74 (diffusion samplers each spell the budget differently)."""
75 return {"max_new_tokens": max_new_tokens}
77 # Runtime gate for enable_compatibility_mode(): set False when the stored-
78 # processed-weights forward is known to diverge (e.g. VaultGemma).
79 supports_compatibility_mode: bool = True
80 # Component paths (by suffix) whose isolated forward cannot run on the
81 # component harness's synthesized probes — e.g. fused top-k routers whose
82 # forward sorts/scatters. They stay hookable at runtime.
83 component_test_skip_suffixes: tuple = ()
85 # Whether run_with_cache should request attention tensors from Hugging Face.
86 # Set False when an adapter reconstructs those tensors but the HF wrapper
87 # rejects output_attentions=True.
88 supports_hf_output_attentions: bool = True
90 # Whether Bridge's shifted next-token cross-entropy is meaningful.
91 supports_causal_loss: bool = True
93 # Optional libraries this adapter needs at load time (e.g. the multimodal group's timm).
94 # Checked at construction so a missing one raises a clear error, not a deep HF failure.
95 required_libraries: list[str] = []
96 # Dependency group that ships required_libraries (named in the error); empty on the base.
97 required_libraries_group: str = ""
99 def __init__(self, cfg: TransformerBridgeConfig) -> None:
100 """Initialize the architecture adapter.
102 Args:
103 cfg: The configuration object.
104 """
105 self._check_required_libraries()
106 self.cfg = cfg
107 self.component_mapping: ComponentMapping | None = None
108 self.weight_processing_conversions: Dict[str, ParamProcessingConversion | str] | None = None
109 self.uses_split_attention: bool = getattr(cfg, "uses_split_attention", False)
110 self._fold_ln_requested: bool = True
111 self._merge_default_config()
113 def _check_required_libraries(self) -> None:
114 """Raise a clear error if an optional library this adapter needs is not installed."""
115 import importlib.util
117 missing = [lib for lib in self.required_libraries if importlib.util.find_spec(lib) is None]
118 if missing:
119 joined = ", ".join(missing)
120 plural = "y" if len(missing) == 1 else "ies"
121 group = self.required_libraries_group
122 group_clause = f" from the '{group}' dependency group" if group else ""
123 contrib = f" (contributors: `uv sync --group {group}`)" if group else ""
124 raise ImportError(
125 f"{type(self).__name__} needs the optional {joined} librar{plural}{group_clause}. "
126 f"Install with `pip install {' '.join(missing)}`{contrib}."
127 )
129 def _merge_default_config(self) -> None:
130 """Merge default_cfg into cfg for variables that don't exist in cfg."""
131 for key, value in self.default_cfg.items():
132 if not hasattr(self.cfg, key): 132 ↛ 131line 132 didn't jump to line 131 because the condition on line 132 was always true
133 setattr(self.cfg, key, value)
135 def _gated_mlp(
136 self,
137 name: str = "mlp",
138 *,
139 gate: str = "gate_proj",
140 up: str = "up_proj",
141 down: str = "down_proj",
142 optional: bool = False,
143 ) -> GatedMLPBridge:
144 """GatedMLPBridge with the standard gate/in/out LinearBridge submodules
145 (up->in, down->out)."""
146 return GatedMLPBridge(
147 name=name,
148 config=self.cfg,
149 submodules={
150 "gate": LinearBridge(name=gate),
151 "in": LinearBridge(name=up),
152 "out": LinearBridge(name=down),
153 },
154 optional=optional,
155 )
157 @property
158 def components(self) -> ComponentMapping:
159 """component_mapping, asserted built — for subclasses that extend or edit
160 a parent's mapping without per-file None narrowing."""
161 assert self.component_mapping is not None, "component_mapping has not been built"
162 return self.component_mapping
164 def _ungated_mlp(
165 self,
166 name: str = "mlp",
167 *,
168 up: str = "up_proj",
169 down: str = "down_proj",
170 optional: bool = False,
171 ) -> MLPBridge:
172 """MLPBridge with the standard in/out LinearBridge submodules (up->in, down->out)."""
173 return MLPBridge(
174 name=name,
175 config=self.cfg,
176 submodules={
177 "in": LinearBridge(name=up),
178 "out": LinearBridge(name=down),
179 },
180 optional=optional,
181 )
183 def _canonical_layer_types(self, cfg: Any) -> list[str]:
184 """Per-layer mixer-type list, normalized to canonical TL names
185 (mamba->linear_attention, attention->full_attention; others pass through)."""
186 aliases = {"mamba": "linear_attention", "attention": "full_attention"}
187 raw = getattr(cfg, "layers_block_type", None) or getattr(cfg, "layer_types", None) or []
188 return [aliases.get(t, t) for t in raw]
190 def _extract_vision_dims(self, cfg: Any) -> None:
191 """Copy vision-tower dims onto cfg, handling both HF-standard naming
192 (num_hidden_layers/num_attention_heads) and Qwen-style (depth/num_heads)."""
193 vision_cfg = getattr(cfg, "vision_config", None)
194 if vision_cfg is None:
195 return
196 self.cfg.vision_hidden_size = getattr(vision_cfg, "hidden_size", None)
197 self.cfg.vision_num_layers = getattr(
198 vision_cfg, "num_hidden_layers", getattr(vision_cfg, "depth", None)
199 )
200 self.cfg.vision_num_heads = getattr(
201 vision_cfg, "num_attention_heads", getattr(vision_cfg, "num_heads", None)
202 )
204 def _set_rms_rotary_defaults(self, *, final_rms: bool = True) -> None:
205 """Set the Llama-family config flags: RMS norms, rotary positions, gated MLP
206 (final_rms is per-architecture -- Mistral/Mixtral/OLMoE set False)."""
207 self.cfg.normalization_type = "RMS"
208 self.cfg.positional_embedding_type = "rotary"
209 self.cfg.final_rms = final_rms
210 self.cfg.gated_mlp = True
211 self.cfg.attn_only = False
212 self.cfg.uses_rms_norm = True
214 def apply_output_logits_transform(self, logits: torch.Tensor) -> torch.Tensor:
215 """Apply the architecture's declared post-unembedding transform.
217 Most causal decoders only apply the shared optional Gemma-style soft cap.
218 Adapters with additional scaling override this method instead of relying
219 on similarly named configuration fields used by unrelated architectures.
220 """
221 return apply_softcap(logits, getattr(self.cfg, "output_logits_soft_cap", None))
223 def validate_output_logits_transform(self) -> None:
224 """Validate that the adapter can reproduce its post-unembedding path."""
226 def _qkvo_weight_conversions(
227 self, n_kv_heads: Optional[int] = None, include_biases: bool = False
228 ) -> Dict[str, ParamProcessingConversion]:
229 """Standard Q/K/V/O weight rearrangement conversions.
231 Most decoder-only models use the same rearrange patterns for attention
232 weights. Override only when your model's layout differs.
234 Args:
235 n_kv_heads: Number of KV heads for GQA. If None, falls back to n_heads.
236 include_biases: Also emit Q/K/V bias reshapes. K/V use the kv-head
237 count — a hand-rolled n_heads reshape breaks on GQA checkpoints.
238 """
239 if n_kv_heads is None:
240 n_kv_heads = getattr(self.cfg, "n_key_value_heads", None) or self.cfg.n_heads
241 conversions = {
242 "blocks.{i}.attn.q.weight": ParamProcessingConversion(
243 tensor_conversion=RearrangeTensorConversion("(n h) m -> n m h", n=self.cfg.n_heads),
244 ),
245 "blocks.{i}.attn.k.weight": ParamProcessingConversion(
246 tensor_conversion=RearrangeTensorConversion("(n h) m -> n m h", n=n_kv_heads),
247 ),
248 "blocks.{i}.attn.v.weight": ParamProcessingConversion(
249 tensor_conversion=RearrangeTensorConversion("(n h) m -> n m h", n=n_kv_heads),
250 ),
251 "blocks.{i}.attn.o.weight": ParamProcessingConversion(
252 tensor_conversion=RearrangeTensorConversion("m (n h) -> n h m", n=self.cfg.n_heads),
253 ),
254 }
255 if include_biases:
256 conversions.update(
257 {
258 "blocks.{i}.attn.q.bias": ParamProcessingConversion(
259 tensor_conversion=RearrangeTensorConversion(
260 "(h d_head) -> h d_head", h=self.cfg.n_heads
261 ),
262 ),
263 "blocks.{i}.attn.k.bias": ParamProcessingConversion(
264 tensor_conversion=RearrangeTensorConversion(
265 "(h d_head) -> h d_head", h=n_kv_heads
266 ),
267 ),
268 "blocks.{i}.attn.v.bias": ParamProcessingConversion(
269 tensor_conversion=RearrangeTensorConversion(
270 "(h d_head) -> h d_head", h=n_kv_heads
271 ),
272 ),
273 }
274 )
275 return conversions
277 def _reprefix_components(self, old: str, new: str) -> None:
278 """Rewrite component names starting with ``old`` to start with ``new``.
280 For adapters that reuse a parent mapping under a different module
281 nesting (e.g. a multimodal wrapper's ``model.language_model.``).
282 """
283 assert self.component_mapping is not None
284 for component in self.component_mapping.values():
285 if component.name and component.name.startswith(old):
286 component.name = new + component.name[len(old) :]
288 def preprocess_weights(self, state_dict: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
289 """Apply architecture-specific weight transformations before ProcessWeights.
291 This method allows architectures to apply custom transformations to weights
292 before standard weight processing (fold_layer_norm, center_writing_weights, etc.).
293 For example, Gemma models scale embeddings by sqrt(d_model).
295 Args:
296 state_dict: The state dictionary with HuggingFace format keys
298 Returns:
299 The modified state dictionary (default implementation returns unchanged)
300 """
301 return state_dict
303 def get_component_mapping(self) -> ComponentMapping:
304 """Get the full component mapping.
306 Returns:
307 The component mapping dictionary
309 Raises:
310 ValueError: If the component mapping is not set
311 """
312 if self.component_mapping is None: 312 ↛ 313line 312 didn't jump to line 313 because the condition on line 312 was never true
313 raise ValueError("component_mapping must be set before calling get_component_mapping")
314 return self.component_mapping
316 def get_remote_component(self, model: RemoteModel, path: RemotePath) -> RemoteComponent:
317 """Get a component from a remote model by its path.
319 This method should be overridden by subclasses to provide the logic for
320 accessing components in a specific model architecture.
322 Args:
323 model: The remote model
324 path: The path to the component in the remote model's format
326 Returns:
327 The component (e.g., a PyTorch module)
329 Raises:
330 AttributeError: If a component in the path doesn't exist
331 IndexError: If an invalid index is accessed
332 ValueError: If the path is empty or invalid
334 Examples:
335 Get an embedding component:
337 >>> # adapter.get_remote_component(model, "model.embed_tokens")
338 >>> # <Embedding>
340 Get a transformer block:
342 >>> # adapter.get_remote_component(model, "model.layers.0")
343 >>> # <TransformerBlock> # type: ignore[index]
345 Get a layer norm component:
347 >>> # adapter.get_remote_component(model, "model.layers.0.ln1")
348 >>> # <LayerNorm>
349 """
350 current = model
351 parent_stack: list[RemoteComponent] = [] # Track parent components for .. navigation
353 # Handle ../ pattern by replacing with a marker before splitting
354 # This is needed because "../output.dense".split(".") gives ['', '', '/output', 'dense']
355 path_with_markers = path.replace("../", "##PARENT##.")
357 for part in path_with_markers.split("."):
358 # If current is a GeneralizedComponent bridge, unwrap to get the original HF component
359 if (
360 isinstance(current, GeneralizedComponent)
361 and hasattr(current, "original_component")
362 and current.original_component is not None
363 ):
364 current = current.original_component
366 if part == "##PARENT##": 366 ↛ 368line 366 didn't jump to line 368 because the condition on line 366 was never true
367 # Navigate to parent component (from ../ syntax)
368 if not parent_stack:
369 raise ValueError(f"Cannot navigate above root in path: {path}")
370 current = parent_stack.pop()
371 elif part == "..": 371 ↛ 373line 371 didn't jump to line 373 because the condition on line 371 was never true
372 # Navigate to parent component (from plain .. syntax)
373 if not parent_stack:
374 raise ValueError(f"Cannot navigate above root in path: {path}")
375 current = parent_stack.pop()
376 elif part.isdigit():
377 parent_stack.append(current)
378 current = current[int(part)] # type: ignore[index]
379 else:
380 parent_stack.append(current)
381 current = getattr(current, part)
382 return current
384 def get_component_from_list_module(
385 self, list_module: RemoteComponent, bridge_component: GeneralizedComponent, parts: list[str]
386 ) -> RemoteComponent:
387 """Get a component from a list module using the bridge component and the transformer lens path.
388 Args:
389 list_module: The remote list module to get the component from
390 bridge_component: The bridge component
391 parts: The parts of the transformer lens path to navigate
392 Returns:
393 The requested component from the list module described by the path
394 """
395 item_index = parts[1]
396 if not item_index.isdigit():
397 raise ValueError(f"Expected item index, got {item_index}")
398 if not hasattr(list_module, "__getitem__"): 398 ↛ 399line 398 didn't jump to line 399 because the condition on line 398 was never true
399 raise TypeError(f"Component {bridge_component.name} is not indexable")
400 indexable_container = cast(Any, list_module)
401 item = indexable_container[int(item_index)]
402 if len(parts) == 2:
403 return item
404 else:
405 subcomponent_name = parts[2]
406 if subcomponent_name in bridge_component.submodules: 406 ↛ 436line 406 didn't jump to line 436 because the condition on line 406 was always true
407 subcomponent_bridge = bridge_component.submodules[subcomponent_name]
408 if len(parts) > 3:
409 current_bridge = subcomponent_bridge
410 if subcomponent_bridge.name is None: 410 ↛ 411line 410 didn't jump to line 411 because the condition on line 410 was never true
411 current = item
412 else:
413 current = self.get_remote_component(item, subcomponent_bridge.name)
414 for i in range(3, len(parts)):
415 deeper_component_name = parts[i]
416 if deeper_component_name.isdigit() and current_bridge.is_list_item: 416 ↛ 417line 416 didn't jump to line 417 because the condition on line 416 was never true
417 return self.get_component_from_list_module(
418 current, current_bridge, parts[i - 1 :]
419 )
420 if deeper_component_name in current_bridge.submodules: 420 ↛ 427line 420 didn't jump to line 427 because the condition on line 420 was always true
421 current_bridge = current_bridge.submodules[deeper_component_name]
422 if current_bridge.name is None: 422 ↛ 423line 422 didn't jump to line 423 because the condition on line 422 was never true
423 pass
424 else:
425 current = self.get_remote_component(current, current_bridge.name)
426 else:
427 raise ValueError(
428 f"Component {deeper_component_name} not found in {'.'.join(parts[:i])} components"
429 )
430 return current
431 elif subcomponent_bridge.name is None: 431 ↛ 432line 431 didn't jump to line 432 because the condition on line 431 was never true
432 return item
433 else:
434 return self.get_remote_component(item, subcomponent_bridge.name)
435 else:
436 raise ValueError(
437 f"Component {subcomponent_name} not found in {parts[0]} components"
438 )
440 def get_generalized_component(self, path: TransformerLensPath) -> GeneralizedComponent:
441 """Get the generalized component (bridge component) for a given TransformerLens path.
443 Args:
444 path: The TransformerLens path to get the component for
446 Returns:
447 The generalized component that handles this path
449 Raises:
450 ValueError: If component_mapping is not set or if the component is not found
452 Examples:
453 Get the embedding bridge component:
455 >>> # adapter.get_generalized_component("embed")
456 >>> # <EmbeddingBridge>
458 Get the attention bridge component:
460 >>> # adapter.get_generalized_component("blocks.0.attn")
461 >>> # <AttentionBridge>
462 """
463 if self.component_mapping is None:
464 raise ValueError(
465 "component_mapping must be set before calling get_generalized_component"
466 )
467 component_path, _ = self._preprocess_parameter_path(path)
468 parts = component_path.split(".")
469 if not parts: 469 ↛ 470line 469 didn't jump to line 470 because the condition on line 469 was never true
470 raise ValueError("Empty path")
471 if parts[0] not in self.component_mapping:
472 raise ValueError(f"Component {parts[0]} not found in component mapping")
473 bridge_component = self.component_mapping[parts[0]]
474 if len(parts) == 1:
475 return bridge_component
476 current_component = bridge_component
477 for i in range(1, len(parts)):
478 part = parts[i]
479 if part.isdigit():
480 continue
481 if hasattr(current_component, "submodules") and part in current_component.submodules:
482 current_component = current_component.submodules[part]
483 elif ( 483 ↛ 488line 483 didn't jump to line 488 because the condition on line 483 was never true
484 hasattr(current_component, "__class__")
485 and "AttentionBridge" in current_component.__class__.__name__
486 and (part in ["q", "k", "v", "o"])
487 ):
488 if "JointQKV" in current_component.__class__.__name__:
489 continue
490 elif (
491 hasattr(current_component, "submodules")
492 and part in current_component.submodules
493 ):
494 current_component = current_component.submodules[part]
495 continue
496 elif ( 496 ↛ 501line 496 didn't jump to line 501 because the condition on line 496 was never true
497 hasattr(current_component, "__class__")
498 and "MLPBridge" in current_component.__class__.__name__
499 and (part in ["in", "out", "gate"])
500 ):
501 if (
502 hasattr(current_component, "submodules")
503 and part in current_component.submodules
504 ):
505 current_component = current_component.submodules[part]
506 continue
507 else:
508 continue
509 else:
510 raise ValueError(f"Component {part} not found in {'.'.join(parts[:i])} components")
511 return current_component
513 def get_component(self, model: RemoteModel, path: TransformerLensPath) -> RemoteComponent:
514 """Get a component from the model using the component_mapping.
516 Args:
517 model: The model to extract components from
518 path: The path of the component to get, as defined in component_mapping
520 Returns:
521 The requested component from the model
523 Raises:
524 ValueError: If component_mapping is not set or if the component is not found
525 AttributeError: If a component in the path doesn't exist
526 IndexError: If an invalid index is accessed
528 Examples:
529 Get an embedding component:
531 >>> # adapter.get_component(model, "embed")
532 >>> # <Embedding>
534 Get a transformer block:
536 >>> # adapter.get_component(model, "blocks.0")
537 >>> # <TransformerBlock>
539 Get a layer norm component:
541 >>> # adapter.get_component(model, "blocks.0.ln1")
542 >>> # <LayerNorm>
543 """
544 if self.component_mapping is None: 544 ↛ 545line 544 didn't jump to line 545 because the condition on line 544 was never true
545 raise ValueError("component_mapping must be set before calling get_component")
546 parts = path.split(".")
547 if not parts: 547 ↛ 548line 547 didn't jump to line 548 because the condition on line 547 was never true
548 raise ValueError("Empty path")
549 if self.component_mapping is None or parts[0] not in self.component_mapping:
550 raise ValueError(f"Component {parts[0]} not found in component mapping")
551 bridge_component = self.component_mapping[parts[0]]
552 if len(parts) == 1:
553 if bridge_component.name is None: 553 ↛ 554line 553 didn't jump to line 554 because the condition on line 553 was never true
554 return model
555 return self.get_remote_component(model, bridge_component.name)
556 if bridge_component.is_list_item and len(parts) >= 2: 556 ↛ 561line 556 didn't jump to line 561 because the condition on line 556 was always true
557 if bridge_component.name is None: 557 ↛ 558line 557 didn't jump to line 558 because the condition on line 557 was never true
558 raise ValueError(f"List component {parts[0]} must have a name")
559 list_module = self.get_remote_component(model, bridge_component.name)
560 return self.get_component_from_list_module(list_module, bridge_component, parts)
561 remote_path = bridge_component.name
562 if remote_path is None:
563 raise ValueError(f"Component {parts[0]} must have a name for nested paths")
564 if len(parts) > 1:
565 remote_path = f"{remote_path}.{'.'.join(parts[1:])}"
566 return self.get_remote_component(model, remote_path)
568 def translate_transformer_lens_path(
569 self, path: TransformerLensPath, last_component_only: bool = False
570 ) -> RemotePath:
571 """Translate a TransformerLens path to a remote model path.
573 Args:
574 path: The TransformerLens path to translate
575 last_component_only: If True, return only the last component of the path
577 Returns:
578 The corresponding remote model path
580 Raises:
581 ValueError: If the path is not found in the component mapping
582 """
583 if self.component_mapping is None: 583 ↛ 584line 583 didn't jump to line 584 because the condition on line 583 was never true
584 raise ValueError(
585 "component_mapping must be set before calling translate_transformer_lens_path"
586 )
587 path, param_suffix = self._preprocess_parameter_path(path)
588 parts = path.split(".")
589 if not parts: 589 ↛ 590line 589 didn't jump to line 590 because the condition on line 589 was never true
590 raise ValueError("Empty path")
591 if parts[0] not in self.component_mapping:
592 raise ValueError(f"Component {parts[0]} not found in component mapping")
593 bridge_component = self.component_mapping[parts[0]]
594 if len(parts) == 1:
595 remote_path = bridge_component.name
596 if remote_path is None: 596 ↛ 597line 596 didn't jump to line 597 because the condition on line 596 was never true
597 raise ValueError(f"Component {parts[0]} must have a name for path translation")
598 if param_suffix:
599 remote_path = remote_path + param_suffix
600 if last_component_only:
601 return remote_path.split(".")[-1]
602 return remote_path
603 if bridge_component.is_list_item and len(parts) >= 2: 603 ↛ 659line 603 didn't jump to line 659 because the condition on line 603 was always true
604 item_index = parts[1]
605 if not item_index.isdigit():
606 raise ValueError(f"Expected item index, got {item_index}")
607 items_path = bridge_component.name
608 if items_path is None: 608 ↛ 609line 608 didn't jump to line 609 because the condition on line 608 was never true
609 raise ValueError(f"List component {parts[0]} must have a name for path translation")
610 if len(parts) == 2:
611 remote_path = f"{items_path}.{item_index}"
612 if param_suffix: 612 ↛ 613line 612 didn't jump to line 613 because the condition on line 612 was never true
613 remote_path = remote_path + param_suffix
614 if last_component_only:
615 return remote_path.split(".")[-1]
616 return remote_path
617 else:
618 subcomponent_name = parts[2]
619 if subcomponent_name in bridge_component.submodules:
620 subcomponent_bridge = bridge_component.submodules[subcomponent_name]
621 if len(parts) > 3:
622 current_bridge = subcomponent_bridge
623 remote_path_parts = [items_path, item_index]
624 if subcomponent_bridge.name is not None:
625 remote_path_parts.append(subcomponent_bridge.name)
626 for i in range(3, len(parts)):
627 deeper_component_name = parts[i]
628 if deeper_component_name in current_bridge.submodules: 628 ↛ 634line 628 didn't jump to line 634 because the condition on line 628 was always true
629 current_bridge = current_bridge.submodules[deeper_component_name]
630 deeper_name = current_bridge.name
631 if deeper_name is not None: 631 ↛ 626line 631 didn't jump to line 626 because the condition on line 631 was always true
632 remote_path_parts.append(deeper_name)
633 else:
634 raise ValueError(
635 f"Component {deeper_component_name} not found in {'.'.join(parts[:i])} components"
636 )
637 remote_path = ".".join(remote_path_parts)
638 if param_suffix:
639 remote_path = remote_path + param_suffix
640 if last_component_only:
641 return remote_path.split(".")[-1]
642 return remote_path
643 else:
644 subcomponent_name_str = subcomponent_bridge.name
645 if subcomponent_name_str is None:
646 raise ValueError(
647 f"Synthetic component {subcomponent_name} has no remote path"
648 )
649 remote_path = f"{items_path}.{item_index}.{subcomponent_name_str}"
650 if param_suffix:
651 remote_path = remote_path + param_suffix
652 if last_component_only:
653 return remote_path.split(".")[-1]
654 return remote_path
655 else:
656 raise ValueError(
657 f"Component {subcomponent_name} not found in {parts[0]} components"
658 )
659 remote_path = bridge_component.name
660 if remote_path is None:
661 raise ValueError(f"Component {parts[0]} must have a name for path translation")
662 if len(parts) > 1:
663 remote_path = f"{remote_path}.{'.'.join(parts[1:])}"
664 if param_suffix:
665 remote_path = remote_path + param_suffix
666 if last_component_only:
667 return remote_path.split(".")[-1]
668 return remote_path
670 def _preprocess_parameter_path(self, path: str) -> tuple[str, str]:
671 """Preprocess TransformerLens path to map parameter names to component names.
673 Args:
674 path: The original TransformerLens path
676 Returns:
677 Tuple of (preprocessed_path, parameter_suffix)
678 """
679 param_suffix = ""
680 if path.endswith(
681 (
682 ".W_Q",
683 ".W_K",
684 ".W_V",
685 ".W_O",
686 ".W_in",
687 ".W_out",
688 ".W_gate",
689 ".W_E",
690 ".W_U",
691 ".W_pos",
692 ".w",
693 "._W_K",
694 "._W_V",
695 )
696 ):
697 param_suffix = ".weight"
698 elif path.endswith(
699 (
700 ".b_Q",
701 ".b_K", # type: ignore[assignment]
702 ".b_V",
703 ".b_O",
704 ".b_in",
705 ".b_out",
706 ".b_gate",
707 ".b_E",
708 ".b_U",
709 ".b_pos",
710 ".b",
711 "._b_K",
712 "._b_V",
713 )
714 ):
715 param_suffix = ".bias"
716 if any(
717 (
718 path.endswith(suffix)
719 for suffix in [
720 ".W_Q",
721 ".W_K",
722 ".W_V",
723 ".b_Q",
724 ".b_K",
725 ".b_V",
726 "._W_K",
727 "._W_V",
728 "._b_K",
729 "._b_V",
730 ]
731 )
732 ):
733 attn_path_parts = path.split(".")
734 if len(attn_path_parts) >= 3 and attn_path_parts[-2] == "attn": 734 ↛ 761line 734 didn't jump to line 761 because the condition on line 734 was always true
735 attn_component_path = ".".join(attn_path_parts[:-1])
736 try:
737 if self.component_mapping: 737 ↛ 761line 737 didn't jump to line 761 because the condition on line 737 was always true
738 current_mapping = self.component_mapping
739 for part in attn_component_path.split("."):
740 if (
741 hasattr(current_mapping, "submodules")
742 and part in current_mapping.submodules
743 ):
744 current_mapping = current_mapping.submodules[part]
745 elif hasattr(current_mapping, "__getitem__"):
746 current_mapping = current_mapping[part] # type: ignore[assignment]
747 if hasattr(current_mapping, "submodules"): 747 ↛ 761line 747 didn't jump to line 761 because the condition on line 747 was always true
748 attn_components = list(current_mapping.submodules.keys())
749 path = path.replace(".W_Q", ".q")
750 path = path.replace(".W_K", ".k")
751 path = path.replace(".W_V", ".v")
752 path = path.replace(".b_Q", ".q")
753 path = path.replace(".b_K", ".k")
754 path = path.replace(".b_V", ".v")
755 path = path.replace("._W_K", ".k")
756 path = path.replace("._W_V", ".v")
757 path = path.replace("._b_K", ".k")
758 path = path.replace("._b_V", ".v")
759 except Exception:
760 pass
761 if any( 761 ↛ 764line 761 didn't jump to line 764 because the condition on line 761 was never true
762 (path.endswith(suffix) for suffix in [".W_Q", ".W_K", ".W_V", ".b_Q", ".b_K", ".b_V"])
763 ):
764 path = path.replace(".W_Q", ".q")
765 path = path.replace(".W_K", ".k")
766 path = path.replace(".W_V", ".v")
767 path = path.replace(".b_Q", ".q")
768 path = path.replace(".b_K", ".k")
769 path = path.replace(".b_V", ".v")
770 path = path.replace(".W_O", ".o")
771 path = path.replace(".b_O", ".o")
772 if any(
773 (
774 path.endswith(suffix)
775 for suffix in [".W_in", ".W_out", ".b_in", ".b_out", ".ln.w", ".ln.b"]
776 )
777 ):
778 mlp_path_parts = path.split(".")
779 if len(mlp_path_parts) >= 3 and mlp_path_parts[-2] == "mlp": 779 ↛ 814line 779 didn't jump to line 814 because the condition on line 779 was always true
780 mlp_component_path = ".".join(mlp_path_parts[:-1])
781 try:
782 if self.component_mapping: 782 ↛ 814line 782 didn't jump to line 814 because the condition on line 782 was always true
783 current_mapping = self.component_mapping
784 for part in mlp_component_path.split("."):
785 if (
786 hasattr(current_mapping, "submodules")
787 and part in current_mapping.submodules
788 ):
789 current_mapping = current_mapping.submodules[part]
790 elif hasattr(current_mapping, "__getitem__"):
791 current_mapping = current_mapping[part] # type: ignore[assignment]
792 if hasattr(current_mapping, "submodules"): 792 ↛ 814line 792 didn't jump to line 814 because the condition on line 792 was always true
793 mlp_components = list(current_mapping.submodules.keys())
794 if "input" in mlp_components and "out" in mlp_components: 794 ↛ 795line 794 didn't jump to line 795 because the condition on line 794 was never true
795 path = path.replace(".W_in", ".input")
796 path = path.replace(".b_in", ".input")
797 path = path.replace(".W_out", ".out")
798 path = path.replace(".b_out", ".out")
799 elif "in" in mlp_components and "out" in mlp_components: 799 ↛ 804line 799 didn't jump to line 804 because the condition on line 799 was always true
800 path = path.replace(".W_in", ".in")
801 path = path.replace(".b_in", ".in")
802 path = path.replace(".W_out", ".out")
803 path = path.replace(".b_out", ".out")
804 elif "fc_in" in mlp_components and "fc_out" in mlp_components:
805 path = path.replace(".W_in", ".fc_in")
806 path = path.replace(".b_in", ".fc_in")
807 path = path.replace(".W_out", ".fc_out")
808 path = path.replace(".b_out", ".fc_out")
809 if "ln" in mlp_components: 809 ↛ 810line 809 didn't jump to line 810 because the condition on line 809 was never true
810 path = path.replace(".ln.w", ".ln")
811 path = path.replace(".ln.b", ".ln")
812 except Exception:
813 pass
814 if any((path.endswith(suffix) for suffix in [".W_in", ".W_out", ".b_in", ".b_out"])): 814 ↛ 815line 814 didn't jump to line 815 because the condition on line 814 was never true
815 path = path.replace(".W_in", ".in")
816 path = path.replace(".b_in", ".in")
817 path = path.replace(".W_out", ".out")
818 path = path.replace(".b_out", ".out")
819 path = path.replace(".W_gate", ".gate")
820 path = path.replace(".b_gate", ".gate")
821 if not (path.endswith(".weight") or path.endswith(".bias")): 821 ↛ 830line 821 didn't jump to line 830 because the condition on line 821 was always true
822 path = path.replace(".W_E", "")
823 path = path.replace(".b_E", "")
824 path = path.replace(".W_U", "")
825 path = path.replace(".b_U", "")
826 path = path.replace(".W_pos", "")
827 path = path.replace(".b_pos", "")
828 path = path.replace(".w", "")
829 path = path.replace(".b", "")
830 return (path, param_suffix)
832 def convert_hf_key_to_tl_key(self, hf_key: str) -> str:
833 """Convert a HuggingFace-style key to TransformerLens format key using component mapping.
835 The component mapping keys ARE the TL format names (e.g., "embed", "pos_embed", "blocks").
836 The component.name is the HF path (e.g., "transformer.wte").
838 Args:
839 hf_key: The HuggingFace-style key (e.g., "transformer.wte.weight")
841 Returns:
842 The TransformerLens format key (e.g., "embed.weight")
843 """
844 if self.component_mapping is None: 844 ↛ 845line 844 didn't jump to line 845 because the condition on line 844 was never true
845 return hf_key
846 for tl_name, component in self.component_mapping.items():
847 if tl_name in ("blocks", "L_blocks", "H_blocks"):
848 continue
849 hf_path = component.name
850 if hf_path is not None and hf_key.startswith(hf_path + "."):
851 param = hf_key[len(hf_path) + 1 :]
852 return f"{tl_name}.{param}"
853 for bl_tl_name in ("blocks", "L_blocks", "H_blocks"):
854 blocks_component = self.component_mapping.get(bl_tl_name)
855 if blocks_component:
856 hf_blocks_prefix = blocks_component.name
857 if hf_blocks_prefix is not None and hf_key.startswith(hf_blocks_prefix + "."):
858 rest = hf_key[len(hf_blocks_prefix) + 1 :]
859 parts = rest.split(".", 1)
860 if len(parts) >= 2 and parts[0].isdigit(): 860 ↛ 853line 860 didn't jump to line 853 because the condition on line 860 was always true
861 layer_idx = parts[0]
862 subkey = parts[1]
863 if hasattr(blocks_component, "submodules"): 863 ↛ 853line 863 didn't jump to line 853 because the condition on line 863 was always true
864 for tl_subname, subcomponent in blocks_component.submodules.items(): 864 ↛ 853line 864 didn't jump to line 853 because the loop on line 864 didn't complete
865 hf_subpath = subcomponent.name
866 if hasattr(subcomponent, "submodules"): 866 ↛ 882line 866 didn't jump to line 882 because the condition on line 866 was always true
867 for (
868 tl_nested_name,
869 nested_comp,
870 ) in subcomponent.submodules.items():
871 if hf_subpath is not None:
872 hf_nested_path: Optional[
873 str
874 ] = f"{hf_subpath}.{nested_comp.name}"
875 else:
876 hf_nested_path = nested_comp.name
877 if hf_nested_path is not None and subkey.startswith(
878 hf_nested_path + "."
879 ):
880 param = subkey[len(hf_nested_path) + 1 :]
881 return f"{bl_tl_name}.{layer_idx}.{tl_subname}.{tl_nested_name}.{param}"
882 if hf_subpath is not None and subkey.startswith(hf_subpath + "."):
883 param = subkey[len(hf_subpath) + 1 :]
884 return f"{bl_tl_name}.{layer_idx}.{tl_subname}.{param}"
885 if hf_subpath is None and subkey.startswith(tl_subname + "."):
886 param = subkey[len(tl_subname) + 1 :]
887 return f"{bl_tl_name}.{layer_idx}.{tl_subname}.{param}"
888 return hf_key
890 def prepare_loading(self, model_name: str, model_kwargs: dict) -> None:
891 """Called before HuggingFace model loading to apply architecture-specific patches.
893 Override this to patch HF model classes before from_pretrained() is called.
894 For example, patching custom model code that is incompatible with transformers v5
895 meta device initialization. Overrides that also want the base eager-forcing
896 must call super().prepare_loading(...).
898 The base implementation forces eager attention on the loading config when
899 cfg.attn_implementation == "eager": composite configs (vision towers) don't
900 reliably inherit the from_pretrained attn_implementation kwarg.
902 Args:
903 model_name: The HuggingFace model name/path
904 model_kwargs: The kwargs dict that will be passed to from_pretrained()
905 """
906 if getattr(self.cfg, "attn_implementation", None) == "eager":
907 config = model_kwargs.get("config")
908 if config is not None and hasattr(config, "_attn_implementation"): 908 ↛ exitline 908 didn't return from function 'prepare_loading' because the condition on line 908 was always true
909 config._attn_implementation = "eager"
911 def prepare_model(self, hf_model: Any) -> None:
912 """Called after HuggingFace model loading but before bridge creation.
914 Override this to fix up the loaded model (e.g., create synthetic modules,
915 re-initialize deferred computations, apply post-load patches). Overrides
916 that also want the base eager-forcing must call super().prepare_model(...).
918 The base implementation mirrors prepare_loading's eager-forcing onto the
919 loaded model's config when cfg.attn_implementation == "eager".
921 Args:
922 hf_model: The loaded HuggingFace model instance
923 """
924 if getattr(self.cfg, "attn_implementation", None) == "eager":
925 if hasattr(hf_model, "config"): 925 ↛ exitline 925 didn't return from function 'prepare_model' because the condition on line 925 was always true
926 hf_model.config._attn_implementation = "eager"
928 def create_stateful_cache(
929 self,
930 hf_model: Any,
931 batch_size: int,
932 device: Any,
933 dtype: torch.dtype,
934 ) -> Any:
935 """Build the HF cache object for a stateful (SSM) generation loop.
937 Called by ``TransformerBridge.generate()`` once before the token loop
938 when ``cfg.is_stateful`` is True. The returned object is threaded
939 through each forward call as ``cache_params=...`` and is expected to
940 mutate itself in-place.
942 Subclasses for SSM architectures (Mamba, Mamba-2, etc.) must override
943 this. The base raises to catch adapters that set ``is_stateful=True``
944 without providing a cache implementation.
946 Args:
947 hf_model: The wrapped HF model (source of ``.config``).
948 batch_size: Number of sequences generated in parallel.
949 device: Device for cache tensors.
950 dtype: Cache tensor dtype (usually the model's param dtype).
951 """
952 raise NotImplementedError(
953 f"{type(self).__name__}.create_stateful_cache is not implemented. "
954 "If this adapter represents a stateful model (cfg.is_stateful=True), "
955 "it must override create_stateful_cache to return the appropriate "
956 "HF cache object."
957 )
959 # setup_component_testing knobs; adapters override these one-line class
960 # attributes instead of re-declaring the standard wiring method.
961 _testing_lm_attr: str = "model"
962 _testing_eager: Optional[str] = "layers"
963 _testing_hybrid: bool = False
964 _testing_rotary_attr: str = "rotary_emb"
965 _testing_wire_rotary: bool = True
967 def setup_component_testing(self, hf_model: Any, bridge_model: Any = None) -> None:
968 """Wire model-specific references for component testing (eager-forcing + rotary
969 per the ``_testing_*`` attributes); override and call super() for extra needs."""
970 self._wire_rotary_for_testing(
971 hf_model,
972 bridge_model,
973 lm_attr=self._testing_lm_attr,
974 hybrid=self._testing_hybrid,
975 eager=self._testing_eager,
976 rotary_attr=self._testing_rotary_attr,
977 wire_rotary=self._testing_wire_rotary,
978 )
980 def _wire_rotary_for_testing(
981 self,
982 hf_model: Any,
983 bridge_model: Any = None,
984 *,
985 lm_attr: str = "model",
986 hybrid: bool = False,
987 eager: Optional[str] = "layers",
988 rotary_attr: str = "rotary_emb",
989 wire_rotary: bool = True,
990 ) -> None:
991 """Force eager attention and set the model's shared rotary_emb on each attention
992 bridge and the template (``wire_rotary=False`` for delegated attention, which has
993 no set_rotary_emb; ``hybrid`` tolerates attention-less blocks)."""
994 lm = hf_model
995 for segment in lm_attr.split("."):
996 lm = getattr(lm, segment, None)
997 if lm is None:
998 break
1000 if (
1001 eager is not None
1002 and hasattr(hf_model, "config")
1003 and hasattr(hf_model.config, "_attn_implementation")
1004 ):
1005 hf_model.config._attn_implementation = "eager"
1006 # Nested multimodal configs carry their own attn implementation.
1007 text_config = getattr(hf_model.config, "text_config", None)
1008 if text_config is not None: 1008 ↛ 1009line 1008 didn't jump to line 1009 because the condition on line 1008 was never true
1009 text_config._attn_implementation = "eager"
1010 if eager == "layers" and lm is not None and hasattr(lm, "layers"):
1011 for layer in lm.layers:
1012 if hasattr(layer, "self_attn") and hasattr(layer.self_attn, "config"): 1012 ↛ 1011line 1012 didn't jump to line 1011 because the condition on line 1012 was always true
1013 layer.self_attn.config._attn_implementation = "eager"
1015 if not wire_rotary or lm is None or not hasattr(lm, rotary_attr):
1016 return
1017 rotary_emb = getattr(lm, rotary_attr)
1019 # No attn submodule in the blocks template (fully delegated attention,
1020 # e.g. Zamba2's shared blocks): wiring is inapplicable, not an error.
1021 blocks_template = (self.component_mapping or {}).get("blocks")
1022 if blocks_template is not None and "attn" not in getattr(blocks_template, "submodules", {}):
1023 return
1025 if bridge_model is not None and hasattr(bridge_model, "blocks"):
1026 for block in bridge_model.blocks:
1027 has_attn = ("attn" in block._modules) if hybrid else hasattr(block, "attn")
1028 if has_attn and hasattr(block.attn, "set_rotary_emb"):
1029 block.attn.set_rotary_emb(rotary_emb)
1031 try:
1032 template = self.get_generalized_component("blocks.0.attn")
1033 except (ValueError, AttributeError, KeyError):
1034 if not hybrid: 1034 ↛ 1035line 1034 didn't jump to line 1035 because the condition on line 1034 was never true
1035 raise
1036 template = None
1037 if template is not None and hasattr(template, "set_rotary_emb"):
1038 template.set_rotary_emb(rotary_emb)
1040 def _enable_ht_attention(self, attn_bridge, hf_attn):
1041 """Enable HT computation for attention (architecture-agnostic).
1043 Detects the architecture by checking which weight attributes exist.
1044 """
1045 n_heads = getattr(
1046 self.cfg,
1047 "n_heads",
1048 getattr(self.cfg, "n_head", getattr(self.cfg, "num_attention_heads", None)),
1049 )
1050 d_model = getattr(
1051 self.cfg, "d_model", getattr(self.cfg, "n_embd", getattr(self.cfg, "hidden_size", None))
1052 )
1053 if n_heads is None or d_model is None:
1054 raise RuntimeError(f"Could not determine n_heads or d_model from config: {self.cfg}")
1055 d_head = d_model // n_heads
1056 if hasattr(hf_attn, "c_attn"):
1057 W_Q, W_K, W_V, b_Q, b_K, b_V = self._extract_qkv_gpt2_style(
1058 hf_attn.c_attn, n_heads, d_model, d_head
1059 )
1060 W_O, b_O = self._extract_output_proj(hf_attn.c_proj, n_heads, d_head, d_model)
1061 elif (
1062 hasattr(hf_attn, "q_proj") and hasattr(hf_attn, "k_proj") and hasattr(hf_attn, "v_proj")
1063 ):
1064 W_Q, b_Q = self._extract_linear_ht_format(hf_attn.q_proj, n_heads, d_head, d_model) # type: ignore[attr-defined]
1065 W_K, b_K = self._extract_linear_ht_format(hf_attn.k_proj, n_heads, d_head, d_model) # type: ignore[attr-defined]
1066 W_V, b_V = self._extract_linear_ht_format(hf_attn.v_proj, n_heads, d_head, d_model) # type: ignore[attr-defined]
1067 out_proj = hf_attn.out_proj if hasattr(hf_attn, "out_proj") else hf_attn.o_proj
1068 W_O, b_O = self._extract_output_proj(out_proj, n_heads, d_head, d_model)
1069 elif hasattr(hf_attn, "query_key_value"):
1070 W_Q, W_K, W_V, b_Q, b_K, b_V = self._extract_qkv_neox_style( # type: ignore[attr-defined]
1071 hf_attn.query_key_value, n_heads, d_model, d_head
1072 )
1073 W_O, b_O = self._extract_output_proj(hf_attn.dense, n_heads, d_head, d_model)
1074 else:
1075 raise ValueError(
1076 f"Unsupported attention architecture. Module has attributes: {dir(hf_attn)}"
1077 )
1078 attn_bridge.set_processed_weights(
1079 {
1080 "W_Q": W_Q,
1081 "W_K": W_K,
1082 "W_V": W_V,
1083 "W_O": W_O,
1084 "b_Q": b_Q,
1085 "b_K": b_K,
1086 "b_V": b_V,
1087 "b_O": b_O,
1088 }
1089 )
1090 self._disable_hook_conversions(attn_bridge)
1092 def _extract_qkv_gpt2_style(self, c_attn, n_heads, d_model, d_head):
1093 """Extract Q, K, V weights from GPT-2 style combined c_attn.
1095 GPT-2 uses Conv1D which stores weights as [in_features, out_features] = [d_model, 3*d_model].
1096 We need to split and reshape to [n_heads, d_model, d_head] format for HookedTransformer.
1097 """
1098 W = c_attn.weight.data
1099 W_Q, W_K, W_V = torch.tensor_split(W, 3, dim=1)
1100 W_Q = einops.rearrange(W_Q, "m (i h)->i m h", i=n_heads)
1101 W_K = einops.rearrange(W_K, "m (i h)->i m h", i=n_heads)
1102 W_V = einops.rearrange(W_V, "m (i h)->i m h", i=n_heads)
1103 qkv_bias = c_attn.bias.data
1104 qkv_bias = einops.rearrange(
1105 qkv_bias, "(qkv index head)->qkv index head", qkv=3, index=n_heads, head=d_head
1106 )
1107 b_Q = qkv_bias[0]
1108 b_K = qkv_bias[1]
1109 b_V = qkv_bias[2]
1110 return (W_Q, W_K, W_V, b_Q, b_K, b_V)
1112 def _extract_output_proj(self, out_proj, n_heads, d_head, d_model):
1113 """Extract output projection weights in HT format.
1115 Returns W_O in [n_heads, d_head, d_model] format for HookedTransformer compatibility.
1117 For Conv1D (GPT-2), weight is stored as [d_model, d_model] = [nx, nf].
1118 For Linear, weight is stored as [d_model, d_model] = [out_features, in_features].
1119 """
1120 weight = out_proj.weight.data
1121 bias = out_proj.bias.data if hasattr(out_proj, "bias") else None
1122 W_O = weight.view(n_heads, d_head, d_model).contiguous()
1123 b_O = bias.contiguous() if bias is not None else None
1124 return (W_O, b_O)
1126 def _disable_hook_conversions(self, attn_bridge):
1127 """Disable hook conversions for attention submodules.
1129 Note: In no_processing mode, we DON'T disable conversions because Q/K/V hooks need
1130 to convert from 3D [batch, seq, d_model] to 4D [batch, seq, n_heads, d_head].
1131 We also preserve o.hook_in.hook_conversion (hook_z).
1133 This method is kept for potential future use but currently does nothing in no_processing mode.
1134 """
1135 pass