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