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