Coverage for transformer_lens/model_bridge/bridge.py: 83%
2343 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-09-01 16:23 +0000
1"""Bridge module for connecting different model architectures.
3This module provides the bridge components that wrap remote model components and provide
4a consistent interface for accessing their weights and performing operations.
5"""
7import inspect
8import logging
9import re
10import warnings
11from collections.abc import Generator
12from contextlib import contextmanager
13from functools import lru_cache
14from typing import (
15 TYPE_CHECKING,
16 Any,
17 Callable,
18 Dict,
19 FrozenSet,
20 Iterable,
21 Iterator,
22 List,
23 Literal,
24 Optional,
25 Tuple,
26 Union,
27 cast,
28 overload,
29)
31import einops
32import numpy as np
33import torch
34import tqdm
35from torch import nn
36from torch.nn import functional as F
37from transformers.tokenization_utils_base import PreTrainedTokenizerBase
39from transformer_lens import utilities as utils
40from transformer_lens.ActivationCache import ActivationCache
41from transformer_lens.config import TransformerBridgeConfig
42from transformer_lens.FactoredMatrix import FactoredMatrix
43from transformer_lens.hook_points import HookIntrospectionMixin, HookPoint
44from transformer_lens.model_bridge.architecture_adapter import ArchitectureAdapter
45from transformer_lens.model_bridge.component_setup import (
46 refresh_container_state_owners,
47 set_original_components,
48)
49from transformer_lens.model_bridge.composition_scores import CompositionScores
50from transformer_lens.model_bridge.exceptions import StopAtLayerException
51from transformer_lens.model_bridge.generalized_components.base import (
52 GeneralizedComponent,
53 alias_generation,
54)
55from transformer_lens.model_bridge.generalized_components.block import (
56 _BLOCK_INTERNAL_MODULES,
57 _NORM_PREFIXES,
58 _VARIANT_SUBMODULE_SET,
59 VARIANT_SUBMODULE_NAMES,
60)
61from transformer_lens.model_bridge.get_params_util import get_bridge_params
62from transformer_lens.utilities.activation_functions import softcap_enabled
63from transformer_lens.utilities.aliases import resolve_alias
64from transformer_lens.utilities.devices import move_to_and_update_config
65from transformer_lens.utilities.lm_utils import lm_cross_entropy_loss
66from transformer_lens.utilities.quantization import require_readable_weight
67from transformer_lens.utilities.slice import Slice, SliceInput
69if TYPE_CHECKING:
70 from transformer_lens.ActivationCache import ActivationCache
72_BLOCK_PATTERN = re.compile("blocks\\.(\\d+)")
74# Block-list container attributes a bridge may expose.
75_BLOCK_LIST_ATTRS = ("blocks", "encoder_blocks", "decoder_blocks", "L_blocks", "H_blocks")
78def _resolve_attr_path(obj: nn.Module, attr_path: str) -> torch.Tensor:
79 """Walk a dot-separated attribute path and return the final tensor."""
80 result = obj
81 for attr in attr_path.split("."):
82 result = getattr(result, attr)
83 return cast(torch.Tensor, result)
86def build_alias_to_canonical_map(hook_dict, prefix=""):
87 """Build a mapping from alias hook names to their canonical names.
89 Args:
90 hook_dict: Dictionary mapping hook names to HookPoint objects
91 prefix: Prefix for nested keys
93 Returns:
94 Dictionary mapping alias names to canonical names
96 Example:
97 If hook_dict contains:
98 - "blocks.0.hook_q" -> HookPoint(name="blocks.0.attn.q.hook_out")
100 Returns:
101 - {"blocks.0.hook_q": "blocks.0.attn.q.hook_out"}
102 """
103 aliases = {}
104 for key, value in hook_dict.items():
105 full_key = f"{prefix}.{key}" if prefix else key
106 if isinstance(value, dict): 106 ↛ 107line 106 didn't jump to line 107 because the condition on line 106 was never true
107 aliases.update(build_alias_to_canonical_map(value, full_key))
108 elif hasattr(value, "name"): 108 ↛ 104line 108 didn't jump to line 104 because the condition on line 108 was always true
109 if key != value.name:
110 aliases[full_key] = value.name
111 return aliases
114def _pos_axis_for_hook(hook_name: str, hook_point: HookPoint) -> int:
115 """Return the axis `pos_slice` applies to for a hook's cached activation.
117 Head-split tensors are [batch, pos, head, d_head], so their position axis is two from the
118 end; everything else (residual stream, MLP, attention patterns keyed by destination
119 position) has it one from the end. HookedTransformer-style names are recognised by
120 suffix, bridge-native ones (`attn.q.hook_out`, `attn.o.hook_in`, ...) by the head-splitting
121 conversion installed on their hook point.
122 """
123 if hook_name.endswith(("hook_q", "hook_k", "hook_v", "hook_z", "hook_result")):
124 return -3
125 conversion = getattr(hook_point, "hook_conversion", None)
126 if getattr(conversion, "splits_attention_heads", False):
127 return -3
128 return -2
131class TransformerBridge(HookIntrospectionMixin, nn.Module):
132 """Bridge between HuggingFace and TransformerLens models.
134 This class provides a standardized interface to access components of a transformer
135 model, regardless of the underlying architecture. It uses an architecture adapter
136 to map between the TransformerLens and HuggingFace model structures.
138 Tokenization notes
139 ------------------
141 :meth:`to_tokens`, :meth:`to_str_tokens`, :meth:`get_token_position`,
142 :meth:`forward` (string input), and :meth:`generate` accept ``prepend_bos``
143 to control BOS prepending. Resolution: explicit arg →
144 ``cfg.default_prepend_bos`` (defaults ``True``, even for non-BOS-trained
145 models — attention heads tend to use position 0 as a resting state).
146 **Pass ``prepend_bos=False`` when tokenizing a fragment of a larger
147 prompt** — off-by-one position errors usually trace back here.
149 Reconciliation with ``cfg.tokenizer_prepends_bos`` (tokenizers that add
150 BOS automatically) is handled internally — pass the value you want;
151 the bridge adds or strips manually as needed. When
152 ``cfg.tokenizer_appends_eos=True`` (OLMo, Apertus, etc.),
153 :meth:`to_tokens` also strips trailing EOS tokens so the model receives
154 a continuation rather than a terminated sequence; this path is
155 bridge-specific.
157 BPE/SentencePiece tokenizers treat ``"hello"``, ``" hello"``, and
158 ``"Hello"`` as distinct tokens. Concatenated prompts may not tokenize
159 as the sum of parts — inspect with :meth:`to_str_tokens` when in doubt.
161 BOS token and chat templates
162 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~
164 ``model.tokenizer`` is configured with ``add_bos_token=True`` and is
165 **not** the stock HuggingFace tokenizer. Direct ``.encode()`` calls
166 will prepend BOS automatically.
168 When passing pre-applied chat-template text (i.e., the output of
169 ``tokenizer.apply_chat_template(..., tokenize=False)``), pass
170 ``prepend_bos=False`` to :meth:`to_tokens` to avoid a double BOS::
172 # Correct pattern for chat templates:
173 text = model.tokenizer.apply_chat_template(messages, tokenize=False)
174 tokens = model.to_tokens(text, prepend_bos=False)
176 The chat template already embeds the model's expected BOS token in
177 the rendered text; letting :meth:`to_tokens` add another would produce
178 a malformed sequence like ``[BOS, BOS, ...]``.
180 To inspect what tokens will actually be fed to the model during
181 generation, use :meth:`to_tokens` directly or pass
182 ``return_input_tokens=True`` to :meth:`generate`.
183 """
185 hook_aliases: Dict[str, Union[str, List[str]]] = {
186 # Prefer embed_ln.hook_out for post-LN models (Bloom, BERT)
187 "hook_embed": ["embed_ln.hook_out", "embed.hook_out"],
188 "hook_pos_embed": ["pos_embed.hook_out", "rotary_emb.hook_out"],
189 "hook_unembed": "unembed.hook_out",
190 }
192 def __init__(self, model: nn.Module, adapter: ArchitectureAdapter, tokenizer: Any):
193 """Initialize the bridge.
195 Args:
196 model: The model to bridge (must be a PyTorch nn.Module or PreTrainedModel)
197 adapter: The architecture adapter to use
198 tokenizer: The tokenizer to use (required)
199 """
200 super().__init__()
201 self._n_params_total = sum(parameter.numel() for parameter in model.parameters())
202 self.__dict__["original_model"] = model
203 self.adapter = adapter
204 self.cfg = adapter.cfg
205 self._tokenizer = None
206 if tokenizer is not None:
207 self.tokenizer = tokenizer # Use the property setter
208 if self.cfg.d_vocab_out == -1:
209 self.cfg.d_vocab_out = self.cfg.d_vocab
210 self.compatibility_mode = False
211 self._weights_processed = False
212 self._hook_cache = None
213 # Nesting depth of hook-adding contexts; hooks are tagged with it so teardown removes
214 # only the ones this call/context added, leaving the caller's own hooks in place.
215 self.context_level = 0
216 self._hook_registry: Dict[str, HookPoint] = {}
217 self._hook_registry_initialized = False
218 self._hook_alias_registry: Dict[str, Union[str, List[str]]] = {}
219 self._block_alias_cache: Optional[Tuple[Tuple[int, int], Dict[str, str]]] = None
220 self._property_alias_registry: Dict[str, str] = {}
221 # real_components maps TL keys to (remote_path, actual_instance) tuples
222 # For list components, actual_instance will be a list of component instances
223 self.real_components: Dict[str, tuple] = {}
224 if not hasattr(self.cfg, "device") or self.cfg.device is None: 224 ↛ 225line 224 didn't jump to line 225 because the condition on line 224 was never true
225 try:
226 self.cfg.device = str(next(self.original_model.parameters()).device)
227 except StopIteration:
228 self.cfg.device = "cpu"
229 if not hasattr(adapter, "component_mapping") or adapter.component_mapping is None: 229 ↛ 230line 229 didn't jump to line 230 because the condition on line 229 was never true
230 raise ValueError("Adapter must have a component_mapping attribute")
231 original_model = self.__dict__["original_model"]
232 set_original_components(self, self.adapter, original_model)
233 self._initialize_hook_registry()
234 self._register_aliases()
235 self._register_all_aliases_recursive()
236 self._setup_hook_compatibility()
237 self._initialize_hooks_to_cache()
238 self.processor = None
239 # Bridge wrappers are inserted into the HF module tree after
240 # from_pretrained's eval(), and nn.Module defaults to training=True —
241 # without re-syncing, reconstruction paths apply dropout at inference.
242 # train() recurses, so this stamps the wrappers with the model's mode.
243 original_model.train(original_model.training)
244 self.train(original_model.training)
245 self.cfg._bind_bridge(self)
247 def __setstate__(self, state: dict[str, Any]) -> None:
248 """Restore runtime config routing after deepcopy or deserialization."""
249 super().__setstate__(state)
250 self.cfg._bind_bridge(self)
252 @property
253 def tokenizer(self) -> Any:
254 """The tokenizer used for encoding/decoding text."""
255 return self._tokenizer
257 @tokenizer.setter
258 def tokenizer(self, value: Any) -> None:
259 """Set tokenizer and re-run wiring (d_vocab, BOS/EOS detection, padding).
261 On initial assignment (during __init__), the boot path has already run
262 setup_tokenizer, so we skip calling it again. However, we still infer
263 d_vocab if it wasn't set from the model config (d_vocab == -1).
265 On reassignment, we re-run the tokenizer wiring and update d_vocab to
266 keep cfg in sync with the new tokenizer.
267 """
268 is_reassignment = getattr(self, "_tokenizer", None) is not None
269 cfg = getattr(self, "cfg", None)
270 if value is not None and cfg is not None:
271 if is_reassignment:
272 from transformer_lens.model_bridge.sources._bridge_builder import (
273 detect_tokenizer_bos_eos,
274 )
275 from transformer_lens.model_bridge.sources.transformers import (
276 setup_tokenizer,
277 )
279 value = setup_tokenizer(
280 value, default_padding_side=getattr(cfg, "default_padding_side", None)
281 )
282 cfg.tokenizer_prepends_bos, cfg.tokenizer_appends_eos = detect_tokenizer_bos_eos(
283 value
284 )
286 # Infer d_vocab: on initial assignment only if not set (-1),
287 # on reassignment always update to match new tokenizer.
288 # Use getattr for cfg attributes since tests may use SimpleNamespace.
289 d_vocab = getattr(cfg, "d_vocab", None)
290 if d_vocab == -1 or is_reassignment:
291 if hasattr(value, "get_vocab"): 291 ↛ 294line 291 didn't jump to line 294 because the condition on line 291 was always true
292 vocab = value.get_vocab()
293 cfg.d_vocab = max(vocab.values()) + 1
294 elif hasattr(value, "vocab"):
295 cfg.d_vocab = max(value.vocab.values()) + 1
296 else:
297 cfg.d_vocab = getattr(value, "vocab_size", 50257)
298 d_vocab_out = getattr(cfg, "d_vocab_out", None)
299 if d_vocab_out == -1 or is_reassignment:
300 cfg.d_vocab_out = getattr(cfg, "d_vocab", d_vocab_out)
301 self._tokenizer = value
303 @classmethod
304 def boot_transformers(
305 cls,
306 model_name: str,
307 hf_config_overrides: Optional[dict] = None,
308 device: Optional[Union[str, torch.device]] = None,
309 dtype: torch.dtype = torch.float32,
310 tokenizer: Optional[Any] = None,
311 load_weights: bool = True,
312 trust_remote_code: bool = False,
313 model_class: Optional[type] = None,
314 hf_model: Optional[Any] = None,
315 device_map: Optional[Union[str, Dict[str, Union[str, int]]]] = None,
316 n_devices: Optional[int] = None,
317 max_memory: Optional[Dict[Union[str, int], Union[str, int]]] = None,
318 n_ctx: Optional[int] = None,
319 revision: Optional[str] = None,
320 checkpoint_index: Optional[int] = None,
321 checkpoint_value: Optional[int] = None,
322 ) -> "TransformerBridge":
323 """Boot a model from HuggingFace (alias for sources.transformers.boot).
325 Returns raw HF weights by default — logits/activations match HF, *not*
326 legacy ``HookedTransformer`` (which folds LayerNorm + centers weights).
327 Call ``enable_compatibility_mode()`` on the result for HookedTransformer-
328 equivalent numerics. Generation, argmax, and CE loss are unaffected.
330 Attention implementation is forced to ``"eager"`` so hooks can capture scores
331 and patterns. For an apples-to-apples HF comparison, load the HF model with
332 ``attn_implementation="eager"`` too; comparing against the default ``"sdpa"``
333 shows ~1e-3 fp32 drift from kernel-level op reordering, not a bridge bug.
335 Args:
336 model_name: The name of the model to load.
337 hf_config_overrides: Optional overrides applied to the HuggingFace config before model load.
338 device: The device to use. If None, will be determined automatically. Mutually exclusive
339 with ``device_map``.
340 dtype: The dtype to use for the model.
341 tokenizer: Optional pre-initialized tokenizer to use; if not provided one will be created.
342 load_weights: If False, load model without weights (on meta device) for config inspection only.
343 trust_remote_code: Whether to trust remote code for custom model architectures.
344 model_class: Optional HuggingFace model class to use instead of the default
345 auto-detected class (e.g., BertForNextSentencePrediction).
346 hf_model: Optional pre-loaded HuggingFace model to use instead of loading one. Useful
347 for models loaded with custom configurations (e.g., quantization via
348 BitsAndBytesConfig). When provided, load_weights is ignored. If the pre-loaded
349 model was built with a ``device_map``, ``cfg.device`` and ``cfg.n_devices`` are
350 derived from its ``hf_device_map`` automatically.
351 device_map: HuggingFace-style device map for dispatched inference. Pass ``"auto"``,
352 ``"balanced"``, ``"sequential"``, or an explicit ``{submodule_path: device}``
353 dict. Explicit maps may include CPU targets; disk / meta offload targets are
354 still rejected because Bridge component wrappers need additional offload-hook
355 routing work. Mutually exclusive with ``device``.
356 n_devices: Convenience shortcut: split the model across this many CUDA devices.
357 Translated to a ``max_memory`` dict over devices 0..n_devices-1 and passed as
358 ``device_map`` to HF. Requires CUDA with at least this many visible devices.
359 max_memory: Optional per-device memory budget, passed through to HF's dispatcher.
360 Only used when ``device_map`` or ``n_devices`` is in effect.
361 n_ctx: Optional context length override. Writes to the appropriate HF config field
362 for this model automatically (callers don't need to know the field name).
363 Warns if larger than the model's default context length.
364 revision: Optional HF revision (branch, tag, or commit). Forwarded to the underlying
365 ``AutoConfig.from_pretrained`` and ``AutoModelForCausalLM.from_pretrained`` calls.
366 Mutually exclusive with ``checkpoint_index`` / ``checkpoint_value``.
367 checkpoint_index: Index into the available training checkpoints for the model family
368 (currently ``EleutherAI/pythia*`` and ``stanford-crfm/*``). Resolved to a revision
369 string via known per-family naming conventions.
370 checkpoint_value: Training step or token count of the desired checkpoint. Alternative
371 to ``checkpoint_index``; must match an entry in the family's checkpoint label list.
373 Returns:
374 The bridge to the loaded model.
375 """
376 from transformer_lens.model_bridge.sources.transformers import boot
378 return boot(
379 model_name=model_name,
380 hf_config_overrides=hf_config_overrides,
381 device=device,
382 dtype=dtype,
383 tokenizer=tokenizer,
384 load_weights=load_weights,
385 trust_remote_code=trust_remote_code,
386 model_class=model_class,
387 hf_model=hf_model,
388 device_map=device_map,
389 n_devices=n_devices,
390 max_memory=max_memory,
391 n_ctx=n_ctx,
392 revision=revision,
393 checkpoint_index=checkpoint_index,
394 checkpoint_value=checkpoint_value,
395 )
397 @overload
398 @classmethod
399 def boot_native(
400 cls,
401 config: TransformerBridgeConfig,
402 tokenizer: Optional[Any] = None,
403 device: Optional[Union[str, torch.device]] = None,
404 dtype: Optional[torch.dtype] = None,
405 model_name: str = "native",
406 ) -> "TransformerBridge":
407 ...
409 @overload
410 @classmethod
411 def boot_native(
412 cls,
413 config: Dict[str, Any],
414 tokenizer: Optional[Any] = None,
415 device: Optional[Union[str, torch.device]] = None,
416 dtype: Optional[torch.dtype] = None,
417 model_name: str = "native",
418 ) -> "TransformerBridge":
419 ...
421 @classmethod
422 def boot_native(
423 cls,
424 config: Any,
425 tokenizer: Optional[Any] = None,
426 device: Optional[Union[str, torch.device]] = None,
427 dtype: Optional[torch.dtype] = None,
428 model_name: str = "native",
429 ) -> "TransformerBridge":
430 """Build a bridge around a small, randomly-initialized TL-native model.
432 No HuggingFace Hub call, no ``transformers`` import. ``config.init_mode``
433 and ``config.seed`` control reproducibility.
434 """
435 # Impl signature stays Any so this guard is reachable — a Union hint
436 # would have beartype reject foreign configs with its own error first.
437 if not isinstance(config, (TransformerBridgeConfig, dict)):
438 raise TypeError(
439 "boot_native expected a TransformerBridgeConfig or dict, "
440 f"got {type(config).__name__}. Construct a TransformerBridgeConfig "
441 "with the same fields."
442 )
444 import copy as _copy
446 from transformer_lens.config import TransformerBridgeConfig as _Cfg
447 from transformer_lens.model_bridge.sources._bridge_builder import (
448 build_bridge_from_module,
449 )
450 from transformer_lens.model_bridge.sources.native import (
451 NativeModel,
452 initialize_native_model,
453 )
455 cfg: TransformerBridgeConfig
456 if isinstance(config, dict):
457 cfg = _Cfg.from_dict(config)
458 else:
459 # Deep-copy so NativeModel's default-resolution writes don't land
460 # on the caller's config.
461 cfg = _copy.deepcopy(config)
463 # Foreign architecture strings would dispatch to the wrong adapter and
464 # crash deep in prepare_model. Refuse them with a pointing message.
465 if cfg.architecture not in (None, "TransformerLensNative"):
466 raise ValueError(
467 f"boot_native cannot build a {cfg.architecture!r} model — "
468 f"it only constructs the TL-native architecture. Either clear "
469 f"config.architecture or set it to 'TransformerLensNative', "
470 f"or use boot_transformers / build_bridge_from_module for "
471 f"non-native architectures."
472 )
473 architecture = "TransformerLensNative"
475 # Fork RNG around construction + init when seeded so neither nn.Linear's
476 # default reset_parameters nor our scoped init perturb the caller's RNG.
477 # When custom init is disabled, construction keeps PyTorch's normal global
478 # RNG semantics and cfg.seed has no initialization work to control.
479 if cfg.init_weights and cfg.seed is not None:
480 with torch.random.fork_rng(devices=[]):
481 model = NativeModel(cfg)
482 initialize_native_model(model, cfg)
483 else:
484 model = NativeModel(cfg)
485 if cfg.init_weights:
486 initialize_native_model(model, cfg)
488 if device is not None:
489 model = model.to(device)
490 if dtype is not None: 490 ↛ 491line 490 didn't jump to line 491 because the condition on line 490 was never true
491 model = model.to(dtype=dtype)
493 return build_bridge_from_module(
494 model,
495 architecture=architecture,
496 tl_config=cfg,
497 tokenizer=tokenizer,
498 dtype=dtype,
499 device=device,
500 model_name=model_name,
501 )
503 def init_weights(self) -> None:
504 """Reinitialize a TL-native model in place using the bridge config."""
505 from transformer_lens.model_bridge.sources.native.init import (
506 initialize_native_model,
507 )
508 from transformer_lens.model_bridge.sources.native.model import NativeModel
510 model = self.original_model
511 if not isinstance(model, NativeModel):
512 raise RuntimeError(
513 "TransformerBridge.init_weights() is only supported for TL-native "
514 "bridges created with TransformerBridge.boot_native(...); this bridge "
515 f"wraps {type(model).__name__}."
516 )
517 initialize_native_model(model, self.cfg)
519 @property
520 def original_model(self) -> nn.Module:
521 """Return the wrapped underlying model; raises AttributeError if it was never set."""
522 if "original_model" not in self.__dict__:
523 raise AttributeError("original_model has not been set")
524 return self.__dict__["original_model"]
526 @original_model.setter
527 def original_model(self, value: nn.Module) -> None:
528 """Store the model in __dict__ so nn.Module does not register it as a submodule."""
529 self.__dict__["original_model"] = value
531 def _register_aliases(self) -> None:
532 """Register bridge-level aliases.
534 This is called at the END of __init__ when all components are set up.
535 It registers the top-level bridge aliases (hook_embed, hook_pos_embed, etc.)
536 and creates direct attribute references.
537 """
538 if self.hook_aliases: 538 ↛ exitline 538 didn't return from function '_register_aliases' because the condition on line 538 was always true
539 self._hook_alias_registry.update(self.hook_aliases)
540 for alias_name, target_path in self.hook_aliases.items():
541 try:
542 if isinstance(target_path, list):
543 for single_target in target_path:
544 try:
545 target_obj = self
546 for part in single_target.split("."):
547 target_obj = getattr(target_obj, part)
548 object.__setattr__(self, alias_name, target_obj)
549 break
550 except AttributeError:
551 continue
552 else:
553 target_obj = self
554 for part in target_path.split("."):
555 target_obj = getattr(target_obj, part)
556 object.__setattr__(self, alias_name, target_obj)
557 except AttributeError:
558 pass
560 def _set_processed_weight_attributes(self) -> None:
561 """Create 3D processed weight attributes for attention components.
563 For each attention component, if it has 2D weights (q.weight, k.weight, v.weight),
564 reshape them to 3D format [n_heads, d_model, d_head] and set as:
565 - _processed_W_Q
566 - _processed_W_K
567 - _processed_W_V
568 - _processed_b_Q
569 - _processed_b_K
570 - _processed_b_V
572 This allows property aliases (W_Q, W_K, W_V) to return 3D format for
573 HookedTransformer compatibility while keeping 2D format for calculations.
574 """
576 n_heads = self.cfg.n_heads
577 d_head = self.cfg.d_head
578 d_model = self.cfg.d_model
579 blocks_iter = []
580 for bl_name in _BLOCK_LIST_ATTRS:
581 if hasattr(self, bl_name):
582 blocks_iter.append(getattr(self, bl_name))
583 if not blocks_iter:
584 return
585 for block in [b for bl in blocks_iter for b in bl]:
586 if "attn" not in block._modules:
587 continue
588 attn = block.attn
589 if not (hasattr(attn, "q") and hasattr(attn.q, "weight")):
590 continue
591 try:
592 w_q_2d = attn.q.weight.data
593 w_k_2d = attn.k.weight.data
594 w_v_2d = attn.v.weight.data
595 attn._processed_W_Q = einops.rearrange(
596 w_q_2d, "m (i h) -> i m h", i=n_heads, h=d_head
597 )
598 attn._processed_W_K = einops.rearrange(
599 w_k_2d, "m (i h) -> i m h", i=n_heads, h=d_head
600 )
601 attn._processed_W_V = einops.rearrange(
602 w_v_2d, "m (i h) -> i m h", i=n_heads, h=d_head
603 )
604 if hasattr(attn.q, "bias") and attn.q.bias is not None:
605 b_q_2d = attn.q.bias.data
606 b_k_2d = attn.k.bias.data
607 b_v_2d = attn.v.bias.data
608 attn._processed_b_Q = einops.rearrange(
609 b_q_2d, "(i h) -> i h", i=n_heads, h=d_head
610 )
611 attn._processed_b_K = einops.rearrange(
612 b_k_2d, "(i h) -> i h", i=n_heads, h=d_head
613 )
614 attn._processed_b_V = einops.rearrange(
615 b_v_2d, "(i h) -> i h", i=n_heads, h=d_head
616 )
617 if hasattr(attn, "o") and hasattr(attn.o, "weight"):
618 w_o_2d = attn.o.weight.data
619 w_o_transposed = w_o_2d.T
620 attn._processed_W_O = einops.rearrange(
621 w_o_transposed, "m (i h) -> i h m", i=n_heads, h=d_head
622 )
623 if hasattr(attn.o, "bias") and attn.o.bias is not None:
624 attn._processed_b_O = attn.o.bias.data
625 except Exception:
626 pass
628 def _register_all_aliases_recursive(self) -> None:
629 """Recursively register aliases on all bridge components.
631 This walks through all components and calls _register_aliases() on each one.
632 Used after weight processing to ensure aliases point to processed weights.
633 """
634 if hasattr(self, "_register_aliases"): 634 ↛ 636line 634 didn't jump to line 636 because the condition on line 634 was always true
635 self._register_aliases()
636 for module in self.modules():
637 if module is not self and hasattr(module, "_register_aliases"):
638 getattr(module, "_register_aliases")()
640 def __setattr__(self, name: str, value: Any) -> None:
641 """Override setattr to track HookPoint objects dynamically."""
642 super().__setattr__(name, value)
643 if isinstance(value, HookPoint): 643 ↛ 644line 643 didn't jump to line 644 because the condition on line 643 was never true
644 value.name = name
645 self._hook_registry[name] = value
646 elif hasattr(value, "get_hooks") and callable(getattr(value, "get_hooks")):
647 component_hooks = value.get_hooks()
648 for hook_name, hook in component_hooks.items():
649 full_name = f"{name}.{hook_name}"
650 hook.name = full_name
651 self._hook_registry[full_name] = hook
653 def _initialize_hook_registry(self) -> None:
654 """Initialize the hook registry by scanning existing components."""
655 if self._hook_registry_initialized: 655 ↛ 656line 655 didn't jump to line 656 because the condition on line 655 was never true
656 return
657 self._scan_existing_hooks(self, "")
658 self._hook_registry_initialized = True
660 def _collect_component_aliases(self, component_mapping, prefix="", _ancestors=frozenset()):
661 """Recursively collect aliases from the architecture's component templates.
663 ``_ancestors`` is path-scoped, not globally-visited: a cycle is cut
664 (else RecursionError at boot) while a diamond-shared component still
665 contributes aliases under both names.
666 """
667 aliases: Dict[str, str] = {}
668 if id(component_mapping) in _ancestors:
669 return aliases
670 _ancestors = _ancestors | {id(component_mapping)}
671 if isinstance(component_mapping, dict):
672 for name, component in component_mapping.items():
673 sub_prefix = f"{prefix}.{name}" if prefix else name
674 aliases.update(self._collect_component_aliases(component, sub_prefix, _ancestors))
675 else:
676 if hasattr(component_mapping, "hook_aliases") and component_mapping.hook_aliases:
677 for alias_name, target in component_mapping.hook_aliases.items():
678 # Skip fallback-list targets: the lru_cached consumer
679 # endswith-matches strings, so a list is neither matchable
680 # nor hashable; lists resolve via the instance walks instead.
681 if not isinstance(target, str):
682 continue
683 full_alias = f"{prefix}.{alias_name}" if prefix else alias_name
684 full_target = f"{prefix}.{target}" if prefix else target
685 aliases[full_alias] = full_target
686 if hasattr(component_mapping, "submodules") and component_mapping.submodules:
687 for sub_name, sub_component in component_mapping.submodules.items():
688 sub_prefix = f"{prefix}.{sub_name}" if prefix else sub_name
689 aliases.update(
690 self._collect_component_aliases(sub_component, sub_prefix, _ancestors)
691 )
692 return aliases
694 @staticmethod
695 @lru_cache(maxsize=128)
696 def _compute_hook_aliases_cached(
697 hook_names_tuple: Tuple[str, ...], component_aliases_tuple: Tuple[Tuple[str, str], ...]
698 ) -> Tuple[Tuple[str, str], ...]:
699 """Cached computation of hook aliases. Takes immutable inputs for caching."""
700 aliases = {}
701 component_aliases = dict(component_aliases_tuple)
702 for hook_name in hook_names_tuple:
703 for alias_pattern, target_pattern in component_aliases.items():
704 if "blocks." in target_pattern and "blocks." in hook_name:
705 block_match = _BLOCK_PATTERN.search(hook_name)
706 if block_match: 706 ↛ 703line 706 didn't jump to line 703 because the condition on line 706 was always true
707 block_num = block_match.group(1)
708 dynamic_alias_pattern = alias_pattern.replace(
709 "blocks.", f"blocks.{block_num}."
710 )
711 dynamic_target_pattern = target_pattern.replace(
712 "blocks.", f"blocks.{block_num}."
713 )
714 if hook_name.endswith(dynamic_target_pattern):
715 target_len = len(dynamic_target_pattern)
716 alias_name = hook_name[:-target_len] + dynamic_alias_pattern
717 aliases[alias_name] = hook_name
718 elif hook_name.endswith(target_pattern):
719 target_len = len(target_pattern)
720 alias_name = hook_name[:-target_len] + alias_pattern
721 aliases[alias_name] = hook_name
722 return tuple(aliases.items())
724 def _collect_hook_aliases_from_registry(self):
725 """Collect aliases based on existing hooks in the registry."""
726 if hasattr(self.adapter, "component_mapping"): 726 ↛ 736line 726 didn't jump to line 736 because the condition on line 726 was always true
727 component_aliases = self._collect_component_aliases(self.adapter.component_mapping)
728 hook_names_tuple = tuple(sorted(self._hook_registry.keys()))
729 component_aliases_tuple = tuple(sorted(component_aliases.items())) # type: ignore[operator]
730 aliases_tuple = self._compute_hook_aliases_cached(
731 hook_names_tuple, component_aliases_tuple
732 )
733 aliases = dict(aliases_tuple)
734 aliases.update(self._collect_block_instance_aliases())
735 return aliases
736 return {}
738 def _collect_block_instance_aliases(self) -> Dict[str, str]:
739 """Collect per-block-instance aliases, overriding template-derived ones.
741 Templates cannot see per-layer rebinds (OlmoHybrid, MoE dense/sparse).
742 Memoized on (registry size, alias generation): size alone misses a
743 dense<->sparse rebind, which changes targets without changing size.
744 """
745 cache_key = (len(self._hook_registry), alias_generation())
746 cached = self._block_alias_cache
747 if cached is not None and cached[0] == cache_key:
748 return cached[1]
749 aliases: Dict[str, str] = {}
750 unresolved: List[str] = []
751 for bl_name in _BLOCK_LIST_ATTRS:
752 block_list = getattr(self, bl_name, None)
753 if block_list is None:
754 continue
755 for i, block in enumerate(block_list):
756 # A block with no registered hooks means the registry hasn't
757 # scanned it yet — unresolved aliases there are timing, not drops.
758 block_prefix = f"{bl_name}.{i}."
759 if f"{block_prefix}hook_in" not in self._hook_registry:
760 continue
761 # Walk the block and its submodule tree: components rebind
762 # aliases per layer at bind time either at block level
763 # (OlmoHybrid) or one level down (MoEBridge's dense/sparse
764 # dispatch). id()-seen guards against shared/cyclic submodule
765 # references, which would otherwise hang boot.
766 # (prefix, component, ids-on-this-path): path-scoped rather than
767 # globally-visited so a cycle is cut while a component shared
768 # under two names still contributes aliases at both.
769 stack: List[Tuple[str, Any, FrozenSet[int]]] = [("", block, frozenset())]
770 while stack:
771 sub_prefix, component, ancestors = stack.pop()
772 if id(component) in ancestors:
773 continue
774 ancestors = ancestors | {id(component)}
775 component_aliases = getattr(component, "hook_aliases", None)
776 if component_aliases:
777 for alias_name, target in component_aliases.items():
778 targets = target if isinstance(target, list) else [target]
779 for single_target in targets:
780 full_target = f"{block_prefix}{sub_prefix}{single_target}"
781 if full_target in self._hook_registry:
782 aliases[f"{block_prefix}{sub_prefix}{alias_name}"] = full_target
783 break
784 else:
785 unresolved.append(f"{block_prefix}{sub_prefix}{alias_name}")
786 for nested_name, nested in (
787 getattr(component, "submodules", None) or {}
788 ).items():
789 stack.append((f"{sub_prefix}{nested_name}.", nested, ancestors))
790 if unresolved:
791 # Surface drops instead of silently swallowing, mirroring
792 # GeneralizedComponent._register_aliases.
793 warnings.warn(
794 f"{len(unresolved)} block hook alias(es) did not resolve to a "
795 f"registered hook (e.g. '{unresolved[0]}'). Any such alias falls "
796 "back to the architecture template's mapping, which for a "
797 "per-layer rebind is the wrong tensor for this layer.",
798 stacklevel=2,
799 )
800 self._block_alias_cache = (cache_key, aliases)
801 return aliases
803 def _add_aliases_to_hooks(self, hooks: Dict[str, HookPoint]) -> None:
804 """Add aliases to hooks in place."""
805 component_aliases = self._collect_hook_aliases_from_registry()
806 all_aliases = {**self.hook_aliases, **component_aliases}
807 if not all_aliases: 807 ↛ 808line 807 didn't jump to line 808 because the condition on line 807 was never true
808 return
809 for alias_name, target in all_aliases.items():
810 if isinstance(target, list):
811 for single_target in target:
812 try:
813 target_hook = resolve_alias(self, alias_name, {alias_name: single_target})
814 if target_hook is not None: 814 ↛ 811line 814 didn't jump to line 811 because the condition on line 814 was always true
815 hooks[alias_name] = target_hook
816 break
817 except AttributeError:
818 continue
819 else:
820 try:
821 target_hook = resolve_alias(self, alias_name, {alias_name: target})
822 if target_hook is not None: 822 ↛ 809line 822 didn't jump to line 809 because the condition on line 822 was always true
823 hooks[alias_name] = target_hook
824 except AttributeError:
825 continue
827 def _scan_existing_hooks(self, module: nn.Module, prefix: str = "") -> None:
828 """Scan existing modules for hooks and add them to registry."""
829 visited = set()
830 # Protect canonical HookPoint names from alias overwrites
831 named_hook_ids: set = set()
833 def scan_module(mod: nn.Module, path: str = "") -> None:
834 obj_id = id(mod)
835 if obj_id in visited:
836 return
837 visited.add(obj_id)
838 if hasattr(mod, "get_hooks") and callable(getattr(mod, "get_hooks")):
839 component_hooks = mod.get_hooks() # type: ignore[operator]
840 if isinstance(component_hooks, dict): 840 ↛ 849line 840 didn't jump to line 849 because the condition on line 840 was always true
841 hooks_dict = cast(Dict[str, HookPoint], component_hooks)
842 for hook_name, hook in hooks_dict.items():
843 full_name = f"{path}.{hook_name}" if path else hook_name
844 hook_id = id(hook)
845 if hook_id not in named_hook_ids:
846 hook.name = full_name
847 named_hook_ids.add(hook_id)
848 self._hook_registry[full_name] = hook
849 for attr_name in dir(mod):
850 if attr_name.startswith("_"):
851 continue
852 if attr_name == "original_component" or attr_name == "original_model":
853 continue
854 if attr_name in [
855 "OV",
856 "QK",
857 "W_V",
858 "W_O",
859 "W_Q",
860 "W_K",
861 "W_in",
862 "W_gate",
863 "W_out",
864 "b_V",
865 "b_O",
866 "b_Q",
867 "b_K",
868 "b_in",
869 "b_out",
870 ]:
871 continue
872 try:
873 attr = getattr(mod, attr_name)
874 except (AttributeError, NameError, RuntimeError, TypeError):
875 continue
876 name = f"{path}.{attr_name}" if path else attr_name
877 if isinstance(attr, HookPoint):
878 hook_id = id(attr)
879 if hook_id not in named_hook_ids:
880 attr.name = name
881 named_hook_ids.add(hook_id)
882 self._hook_registry[name] = attr
883 for child_name, child_module in mod.named_children():
884 if (
885 child_name == "original_component"
886 or child_name == "_original_component"
887 or child_name == "original_model"
888 ):
889 continue
890 child_path = f"{path}.{child_name}" if path else child_name
891 scan_module(child_module, child_path)
893 scan_module(module, prefix)
895 @property
896 def hook_dict(self) -> dict[str, HookPoint]:
897 """Get all HookPoint objects in the model for compatibility with TransformerLens."""
898 hooks = self._hook_registry.copy()
899 self._add_aliases_to_hooks(hooks)
900 return hooks
902 @property
903 def n_params_total(self) -> int:
904 """Number of parameters in the wrapped model before bridge instrumentation.
906 This follows PyTorch's parameter iteration semantics, counting tied
907 parameters once. Bridge-created split views and synthetic zero tensors
908 are excluded, so the result can differ from
909 :attr:`HookedTransformer.n_params_total` and :meth:`tl_parameters`.
911 Returns:
912 int: Parameter count of the uninstrumented wrapped model.
913 """
914 return self._n_params_total
916 def clear_hook_registry(self) -> None:
917 """Clear the hook registry and force re-initialization."""
918 self._hook_registry.clear()
919 self._hook_registry_initialized = False
921 def _initialize_hooks_to_cache(self) -> None:
922 """Initialize the hooks to cache when running the model with cache."""
923 self.hooks_to_cache = {}
924 default_cached_hooks_names = [
925 "embed.hook_in",
926 "embed.hook_out",
927 "pos_embed.hook_in",
928 "pos_embed.hook_out",
929 "rotary_embed.hook_in",
930 "rotary_embed.hook_out",
931 "ln_final.hook_in",
932 "ln_final.hook_scale",
933 "ln_final.hook_normalized",
934 "ln_final.hook_out",
935 "unembed.hook_in",
936 "unembed.hook_out",
937 ]
938 for block_idx in range(self.cfg.n_layers):
939 default_cached_hooks_names.append(f"blocks.{block_idx}.hook_in")
940 default_cached_hooks_names.append(f"blocks.{block_idx}.ln1.hook_in")
941 default_cached_hooks_names.append(f"blocks.{block_idx}.ln1.hook_scale")
942 default_cached_hooks_names.append(f"blocks.{block_idx}.ln1.hook_normalized")
943 default_cached_hooks_names.append(f"blocks.{block_idx}.ln1.hook_out")
944 default_cached_hooks_names.append(f"blocks.{block_idx}.ln1_post.hook_in")
945 default_cached_hooks_names.append(f"blocks.{block_idx}.ln1_post.hook_scale")
946 default_cached_hooks_names.append(f"blocks.{block_idx}.ln1_post.hook_normalized")
947 default_cached_hooks_names.append(f"blocks.{block_idx}.ln1_post.hook_out")
948 default_cached_hooks_names.append(f"blocks.{block_idx}.attn.hook_in")
949 default_cached_hooks_names.append(f"blocks.{block_idx}.attn.q.hook_in")
950 default_cached_hooks_names.append(f"blocks.{block_idx}.attn.q.hook_out")
951 default_cached_hooks_names.append(f"blocks.{block_idx}.attn.q_norm.hook_in")
952 default_cached_hooks_names.append(f"blocks.{block_idx}.attn.q_norm.hook_out")
953 default_cached_hooks_names.append(f"blocks.{block_idx}.attn.k.hook_in")
954 default_cached_hooks_names.append(f"blocks.{block_idx}.attn.k.hook_out")
955 default_cached_hooks_names.append(f"blocks.{block_idx}.attn.k_norm.hook_in")
956 default_cached_hooks_names.append(f"blocks.{block_idx}.attn.k_norm.hook_out")
957 default_cached_hooks_names.append(f"blocks.{block_idx}.attn.v.hook_in")
958 default_cached_hooks_names.append(f"blocks.{block_idx}.attn.v.hook_out")
959 default_cached_hooks_names.append(f"blocks.{block_idx}.attn.o.hook_in")
960 default_cached_hooks_names.append(f"blocks.{block_idx}.attn.o.hook_out")
961 default_cached_hooks_names.append(f"blocks.{block_idx}.attn.hook_attn_scores")
962 default_cached_hooks_names.append(f"blocks.{block_idx}.attn.hook_pattern") # type: ignore[operator]
963 default_cached_hooks_names.append(f"blocks.{block_idx}.attn.hook_hidden_states")
964 default_cached_hooks_names.append(f"blocks.{block_idx}.attn.hook_out")
965 default_cached_hooks_names.append(f"blocks.{block_idx}.ln2.hook_in")
966 default_cached_hooks_names.append(f"blocks.{block_idx}.ln2.hook_scale")
967 default_cached_hooks_names.append(f"blocks.{block_idx}.ln2.hook_normalized")
968 default_cached_hooks_names.append(f"blocks.{block_idx}.ln2.hook_out")
969 default_cached_hooks_names.append(f"blocks.{block_idx}.ln2_post.hook_in") # type: ignore[operator]
970 default_cached_hooks_names.append(f"blocks.{block_idx}.ln2_post.hook_scale")
971 default_cached_hooks_names.append(f"blocks.{block_idx}.ln2_post.hook_normalized")
972 default_cached_hooks_names.append(f"blocks.{block_idx}.ln2_post.hook_out")
973 default_cached_hooks_names.append(f"blocks.{block_idx}.mlp.hook_in") # type: ignore[operator]
974 default_cached_hooks_names.append(f"blocks.{block_idx}.mlp.in.hook_in")
975 default_cached_hooks_names.append(f"blocks.{block_idx}.mlp.in.hook_out") # type: ignore[operator]
976 default_cached_hooks_names.append(f"blocks.{block_idx}.mlp.out.hook_in")
977 default_cached_hooks_names.append(f"blocks.{block_idx}.mlp.out.hook_out")
978 default_cached_hooks_names.append(f"blocks.{block_idx}.mlp.gate.hook_in")
979 default_cached_hooks_names.append(f"blocks.{block_idx}.mlp.gate.hook_out")
980 default_cached_hooks_names.append(f"blocks.{block_idx}.mlp.hook_out")
981 default_cached_hooks_names.append(f"blocks.{block_idx}.hook_out")
982 for hook_name in default_cached_hooks_names:
983 if hook_name in self._hook_registry:
984 self.hooks_to_cache[hook_name] = self._hook_registry[hook_name] # type: ignore[arg-type]
986 def __getattr__(self, name: str) -> Any:
987 """Provide a clear error message for missing attributes."""
988 if name in self.__dict__: # type: ignore[arg-type] 988 ↛ 989line 988 didn't jump to line 989 because the condition on line 988 was never true
989 return self.__dict__[name]
990 # Use __dict__ directly to avoid recursion
991 if "_modules" in self.__dict__ and name in self.__dict__["_modules"]: # type: ignore[arg-type]
992 return self.__dict__["_modules"][name]
993 adapter = self.__dict__.get("adapter")
994 component_mapping = getattr(adapter, "component_mapping", None)
995 if component_mapping is not None and name in component_mapping:
996 raise AttributeError(f"'{type(self).__name__}' object has no attribute '{name}'")
997 if "original_model" in self.__dict__ and self.__dict__["original_model"] is not None:
998 try:
999 name_split = name.split(".")
1000 if len(name_split) > 1: 1000 ↛ 1001line 1000 didn't jump to line 1001 because the condition on line 1000 was never true
1001 current = getattr(self.__dict__["original_model"], name_split[0])
1002 for part in name_split[1:]: # type: ignore[operator]
1003 current = getattr(current, part)
1004 return current
1005 else:
1006 return getattr(self.__dict__["original_model"], name)
1007 except AttributeError:
1008 pass # type: ignore[operator,assignment]
1009 raise AttributeError(f"'{type(self).__name__}' object has no attribute '{name}'")
1011 def __str__(self) -> str:
1012 """Get a string representation of the bridge.
1013 # type: ignore[operator]
1014 Returns:
1015 A string describing the bridge's components # type: ignore[operator]
1016 """
1017 lines = ["TransformerBridge:"]
1018 mapping = self.adapter.get_component_mapping()
1019 lines.extend(self._format_component_mapping(mapping, indent=1))
1020 return "\n".join(lines)
1022 def enable_compatibility_mode(
1023 self,
1024 disable_warnings: bool = False,
1025 no_processing: bool = False,
1026 fold_ln: bool = True,
1027 center_writing_weights: bool = True,
1028 center_unembed: bool = True,
1029 fold_value_biases: bool = True,
1030 refactor_factored_attn_matrices: bool = False,
1031 ) -> None:
1032 """Apply HookedTransformer-equivalent weight processing and legacy hook compatibility.
1034 Defaults match HookedTransformer's load-time processing (fold_ln + weight
1035 centering) — required for analyses that reason in HookedTransformer's
1036 post-processed coordinate system: logit lens, direct logit attribution,
1037 residual-stream norms. Also enables legacy hook/component name aliases.
1039 Hook semantic parity (issue #1317): ``hook_q_input``, ``hook_k_input``,
1040 ``hook_v_input``, ``hook_attn_in``, and ``hook_mlp_in`` fire on the
1041 pre-norm residual. Carve-outs: post-norm architectures (OLMo 2,
1042 BERT-style) read the post-attention residual instead, and MLA blocks
1043 (DeepSeek V2/V3/R1) do not expose the split-qkv aliases. ``hook_mlp_in``
1044 is gated on ``cfg.use_hook_mlp_in``; toggle it via
1045 :py:meth:`set_use_hook_mlp_in`.
1047 Args:
1048 disable_warnings: Whether to disable warnings about legacy components/hooks
1049 no_processing: Whether to disable ALL pre-processing steps of the model.
1050 If True, overrides fold_ln, center_writing_weights, and center_unembed to False.
1051 fold_ln: Whether to fold layer norm weights into the subsequent linear layers.
1052 Default: True. Ignored if no_processing=True.
1053 center_writing_weights: Whether to center the writing weights (W_out in attention and MLPs).
1054 Default: True. Ignored if no_processing=True.
1055 center_unembed: Whether to center the unembedding matrix.
1056 Default: True. Ignored if no_processing=True.
1057 fold_value_biases: Whether to fold value biases into output bias.
1058 Default: True. Ignored if no_processing=True.
1059 refactor_factored_attn_matrices: Whether to refactor factored attention matrices.
1060 Default: False. Ignored if no_processing=True.
1061 """
1062 from transformer_lens.utilities.bridge_components import (
1063 apply_fn_to_all_components,
1064 )
1066 if not getattr(self.adapter, "supports_compatibility_mode", True):
1067 raise RuntimeError(
1068 f"{type(self.adapter).__name__} does not support compatibility mode: "
1069 "its stored-processed-weights is known to diverge from the "
1070 "reference model. Use the default bridge forward instead."
1071 )
1073 self.compatibility_mode = True
1075 def set_compatibility_mode(component: Any) -> None:
1076 """Set compatibility mode on a component."""
1077 component.compatibility_mode = True
1078 component.disable_warnings = disable_warnings
1080 apply_fn_to_all_components(self, set_compatibility_mode)
1081 self.clear_hook_registry()
1082 # Drop block capture-hook handles from any prior call so they don't accumulate.
1083 if hasattr(self, "blocks"): 1083 ↛ 1087line 1083 didn't jump to line 1087 because the condition on line 1083 was always true
1084 for block in self.blocks:
1085 if hasattr(block, "_teardown_capture_hooks"): 1085 ↛ 1084line 1085 didn't jump to line 1084 because the condition on line 1085 was always true
1086 block._teardown_capture_hooks()
1087 try:
1088 if not no_processing:
1089 self.process_weights(
1090 fold_ln=fold_ln,
1091 center_writing_weights=center_writing_weights,
1092 center_unembed=center_unembed,
1093 fold_value_biases=fold_value_biases,
1094 refactor_factored_attn_matrices=refactor_factored_attn_matrices,
1095 )
1096 finally:
1097 # Re-initialize hooks even on failure so bridge stays usable
1098 self._initialize_hook_registry()
1099 self._setup_hook_compatibility()
1100 self._register_all_aliases_recursive()
1102 def _setup_hook_compatibility(self) -> None:
1103 """Setup hook compatibility transformations to match HookedTransformer behavior.
1105 This method sets up hook conversions and wrappers that ensure Bridge hooks
1106 have the same shapes and behavior as HookedTransformer hooks. This includes:
1107 1. hook_z reshaping from [batch, seq, d_model] to [batch, seq, n_heads, d_head]
1108 2. Wrapping HF attention forward to inject position embeddings/attention masks
1109 3. Architecture-specific setup (e.g., rotary embedding references)
1111 This is called during __init__ and should always be run, regardless of whether
1112 compatibility mode or weight processing is enabled.
1114 Note: This method is idempotent - can be called multiple times safely.
1115 """
1116 if hasattr(self.adapter, "setup_hook_compatibility"):
1117 self.adapter.setup_hook_compatibility(self)
1118 elif hasattr(self.adapter, "setup_no_processing_hooks"): 1118 ↛ 1119line 1118 didn't jump to line 1119 because the condition on line 1118 was never true
1119 self.adapter.setup_no_processing_hooks(self)
1120 blocks_to_process = []
1121 for block_list_name in (
1122 "blocks",
1123 "encoder_blocks",
1124 "decoder_blocks",
1125 "L_blocks",
1126 "H_blocks",
1127 ):
1128 if hasattr(self, block_list_name):
1129 blocks_to_process.extend(getattr(self, block_list_name))
1130 for block in blocks_to_process:
1131 for attn_name in ["attn", "self_attn", "cross_attn"]:
1132 if hasattr(block, attn_name):
1133 attn = getattr(block, attn_name)
1134 if hasattr(attn, "setup_hook_compatibility"):
1135 attn.setup_hook_compatibility()
1136 elif hasattr(attn, "setup_no_processing_hooks"): 1136 ↛ 1137line 1136 didn't jump to line 1137 because the condition on line 1136 was never true
1137 attn.setup_no_processing_hooks()
1139 def process_weights(
1140 self,
1141 verbose: bool = False,
1142 fold_ln: bool = True,
1143 center_writing_weights: bool = True,
1144 center_unembed: bool = True,
1145 fold_value_biases: bool = True,
1146 refactor_factored_attn_matrices: bool = False,
1147 ) -> None:
1148 """Process weights directly using ProcessWeights and architecture adapter.
1150 This method applies weight processing transformations to improve model interpretability
1151 without requiring a reference HookedTransformer model. Works with all architectures
1152 supported by TransformerBridge, including GPT-OSS and other new models.
1154 Args:
1155 verbose: If True, print detailed progress messages. Default: False
1156 fold_ln: Fold LayerNorm weights/biases into subsequent layers. Default: True
1157 center_writing_weights: Center weights that write to residual stream. Default: True
1158 center_unembed: Center unembedding weights (translation invariant). Default: True
1159 fold_value_biases: Fold value biases into output bias. Default: True
1160 refactor_factored_attn_matrices: Experimental QK/OV factorization. Default: False
1161 """
1162 # Folding and centering do arithmetic on raw weights, so packed or
1163 # scale-separated storage would produce silent garbage. The forward
1164 # path stays usable when quantized; only this transformation does not.
1165 for name, param in self.original_model.named_parameters():
1166 # No name filter: modern batched-MoE experts are Parameters that do
1167 # NOT end in .weight (mlp.experts.gate_up_proj), and those are the
1168 # very tensors the converters had to guard. unreadable_weight_reason
1169 # is dtype-driven, so every full-width float parameter still passes.
1170 # Routed through the shared helper so meta gets its own "load with
1171 # real weights" message instead of being silently skipped — folding
1172 # on meta tensors yields meta tensors, which is the same
1173 # silent-garbage failure this guard exists to stop.
1174 require_readable_weight(
1175 param,
1176 operation=f"process weights ({name})",
1177 owner=self.original_model,
1178 remedy=(
1179 "Load the model dequantized, or use the bridge without weight "
1180 "processing (enable_compatibility_mode(no_processing=True))."
1181 ),
1182 )
1184 # A failed or partial processing attempt is no longer guaranteed to retain
1185 # the raw HuggingFace basis, so invalidate that contract before any work.
1186 self._weights_processed = True
1187 from transformer_lens.weight_processing import ProcessWeights
1189 if verbose: 1189 ↛ 1190line 1189 didn't jump to line 1190 because the condition on line 1189 was never true
1190 print(f"Processing weights for {self.cfg.model_name}...")
1192 # Soft capping (tanh) is not translation-invariant; centering would change output.
1193 if center_unembed and softcap_enabled(getattr(self.cfg, "output_logits_soft_cap", None)): 1193 ↛ 1194line 1193 didn't jump to line 1194 because the condition on line 1193 was never true
1194 import logging
1196 logging.warning(
1197 "center_unembed=True is incompatible with logit softcapping "
1198 "(output_logits_soft_cap=%.1f). Disabling center_unembed.",
1199 self.cfg.output_logits_soft_cap,
1200 )
1201 center_unembed = False
1203 if verbose: 1203 ↛ 1204line 1203 didn't jump to line 1204 because the condition on line 1203 was never true
1204 print(" Extracting state dict from existing model...")
1205 state_dict = self.state_dict()
1206 adapter = self.adapter
1208 # Untie embed/unembed weights (GPT-2) so centering affects only unembed
1209 embed_key = "embed.weight"
1210 unembed_key = "unembed.weight"
1212 if embed_key in state_dict and unembed_key in state_dict: 1212 ↛ 1220line 1212 didn't jump to line 1220 because the condition on line 1212 was always true
1213 # Check if they point to the same tensor (weight tying)
1214 if state_dict[embed_key].data_ptr() == state_dict[unembed_key].data_ptr():
1215 if verbose: 1215 ↛ 1216line 1215 didn't jump to line 1216 because the condition on line 1215 was never true
1216 print(" Breaking weight tying between embed and unembed in state dict...")
1217 # Clone the unembed weight to break the tie
1218 state_dict[unembed_key] = state_dict[unembed_key].clone()
1220 if adapter and hasattr(adapter, "preprocess_weights"): 1220 ↛ 1226line 1220 didn't jump to line 1226 because the condition on line 1220 was always true
1221 adapter._fold_ln_requested = fold_ln # type: ignore[union-attr]
1222 state_dict = adapter.preprocess_weights(state_dict)
1224 # Use unified ProcessWeights.process_weights() like HookedTransformer does.
1225 # Float32 upcasting for precision is handled centrally in process_weights().
1226 if verbose: 1226 ↛ 1227line 1226 didn't jump to line 1227 because the condition on line 1226 was never true
1227 print(" Processing weights (fold_ln, center_writing_weights, etc.)...")
1228 state_dict = ProcessWeights.process_weights(
1229 state_dict,
1230 self.cfg,
1231 fold_ln=fold_ln,
1232 center_writing_weights=center_writing_weights,
1233 center_unembed=center_unembed,
1234 fold_value_biases=fold_value_biases,
1235 refactor_factored_attn_matrices=refactor_factored_attn_matrices,
1236 adapter=adapter,
1237 )
1239 # Normalize HF-prefix keys to TL format for weight routing
1240 import re
1242 hf_to_tl_prefix = {}
1243 for tl_name, (remote_path, _component) in self.real_components.items():
1244 if remote_path and remote_path != tl_name: 1244 ↛ 1243line 1244 didn't jump to line 1243 because the condition on line 1244 was always true
1245 hf_to_tl_prefix[remote_path] = tl_name
1247 normalized_state_dict = {}
1248 for key, value in state_dict.items():
1249 new_key = key
1250 for hf_prefix, tl_prefix in hf_to_tl_prefix.items():
1251 if key.startswith(hf_prefix + "."): 1251 ↛ 1252line 1251 didn't jump to line 1252 because the condition on line 1251 was never true
1252 suffix = key[len(hf_prefix) + 1 :]
1253 new_key = f"{tl_prefix}.{suffix}"
1254 break
1255 normalized_state_dict[new_key] = value
1256 state_dict = normalized_state_dict
1258 if verbose: 1258 ↛ 1259line 1258 didn't jump to line 1259 because the condition on line 1258 was never true
1259 print(" Distributing weights to generalized components...")
1260 ProcessWeights.distribute_weights_to_components(
1261 state_dict=state_dict,
1262 component_mapping=self.real_components,
1263 )
1265 def _calculate_loss(self, logits, tokens, loss_per_token=False):
1266 """Calculate cross-entropy loss."""
1267 shift_logits = logits[..., :-1, :].contiguous()
1268 shift_labels = tokens[..., 1:].contiguous()
1269 loss_fct = torch.nn.CrossEntropyLoss(reduction="none" if loss_per_token else "mean")
1270 flat_logits = shift_logits.view(-1, shift_logits.size(-1))
1271 flat_labels = shift_labels.view(-1)
1272 loss = loss_fct(flat_logits, flat_labels)
1273 if loss_per_token:
1274 return loss.view(shift_labels.shape)
1275 else:
1276 return loss
1278 def _extract_hf_weights(self):
1279 """Extract weights from the original HuggingFace model."""
1280 hf_state_dict = self.state_dict()
1281 for layer_idx in range(self.cfg.n_layers):
1282 combined_qkv_key = f"transformer.h.{layer_idx}.attn.c_attn.weight"
1283 combined_qkv_bias_key = f"transformer.h.{layer_idx}.attn.c_attn.bias"
1284 if combined_qkv_key in hf_state_dict:
1285 separate_keys_to_remove = [
1286 f"transformer.h.{layer_idx}.attn.q.weight",
1287 f"transformer.h.{layer_idx}.attn.q.bias",
1288 f"transformer.h.{layer_idx}.attn.k.weight",
1289 f"transformer.h.{layer_idx}.attn.k.bias",
1290 f"transformer.h.{layer_idx}.attn.v.weight",
1291 f"transformer.h.{layer_idx}.attn.v.bias",
1292 ]
1293 for key_to_remove in separate_keys_to_remove:
1294 if key_to_remove in hf_state_dict:
1295 del hf_state_dict[key_to_remove]
1296 return hf_state_dict
1298 def to_tokens(
1299 self,
1300 input: Union[str, List[str]],
1301 prepend_bos: Optional[bool] = None,
1302 padding_side: Optional[str] = None,
1303 move_to_device: bool = True,
1304 truncate: bool = True,
1305 ) -> torch.Tensor:
1306 """Converts a string to a tensor of tokens.
1308 See the class-level "Tokenization notes" for full ``prepend_bos``
1309 semantics, the ``default_prepend_bos`` /
1310 ``tokenizer_prepends_bos`` interaction, and the whitespace-
1311 sensitivity gotcha. **Pass ``prepend_bos=False`` whenever you're
1312 tokenizing only part of a prompt.**
1314 Args:
1315 input: The input to tokenize.
1316 prepend_bos: Overrides ``self.cfg.default_prepend_bos``. Defaults
1317 to ``None`` (use the cfg setting). Pass ``True`` or ``False``
1318 to override locally.
1319 padding_side: Which side to pad on when tokenizing multiple
1320 strings of different lengths. Defaults to the tokenizer's
1321 ``padding_side``.
1322 move_to_device: Whether to move the result to ``cfg.device``.
1323 truncate: Whether to truncate inputs longer than ``cfg.n_ctx``.
1325 Returns:
1326 Token tensor of shape ``[batch, pos]``.
1327 """
1328 assert self.tokenizer is not None, "Cannot use to_tokens without a tokenizer"
1329 if prepend_bos is None:
1330 prepend_bos = getattr(self.cfg, "default_prepend_bos", True)
1331 if padding_side is None:
1332 padding_side = getattr(self.tokenizer, "padding_side", "right")
1333 tokenizer_prepends_bos = getattr(self.cfg, "tokenizer_prepends_bos", True)
1334 if prepend_bos and (not tokenizer_prepends_bos):
1335 bos = self.tokenizer.bos_token
1336 encodes_atomically = (
1337 bos is not None
1338 and len(self.tokenizer(bos, add_special_tokens=False)["input_ids"]) == 1
1339 )
1340 if encodes_atomically:
1341 input = utils.get_input_with_manually_prepended_bos(bos, input)
1342 # else: the fallback BOS is not an atom in this vocab (e.g.
1343 # '<|endoftext|>' installed on BERT); prepending the string would
1344 # tokenize to subword garbage, so skip rather than pollute the input.
1345 if isinstance(input, str):
1346 input = [input]
1347 tokens = self.tokenizer(
1348 input,
1349 return_tensors="pt",
1350 padding=True,
1351 padding_side=padding_side,
1352 truncation=truncate,
1353 max_length=self.cfg.n_ctx if truncate else None,
1354 )["input_ids"]
1355 # Strip auto-appended EOS tokens (e.g., OLMo)
1356 if (
1357 getattr(self.cfg, "tokenizer_appends_eos", False)
1358 and self.tokenizer.eos_token_id is not None
1359 ):
1360 # Remove trailing EOS, keep at least 1 token
1361 while tokens.shape[-1] > 1 and (tokens[:, -1] == self.tokenizer.eos_token_id).all():
1362 tokens = tokens[:, :-1]
1363 if not prepend_bos and tokenizer_prepends_bos:
1364 tokens = utils.get_tokens_with_bos_removed(
1365 self.tokenizer, tokens, padding_side=padding_side
1366 )
1367 if move_to_device:
1368 tokens = tokens.to(self.cfg.device)
1369 return tokens
1371 def to_string(
1372 self, tokens: Union[List[int], torch.Tensor, np.ndarray]
1373 ) -> Union[str, List[str]]:
1374 """Convert tokens to string(s).
1376 Args:
1377 tokens: Tokens to convert
1379 Returns:
1380 Decoded string(s)
1381 """
1382 if not isinstance(tokens, torch.Tensor): 1382 ↛ 1383line 1382 didn't jump to line 1383 because the condition on line 1382 was never true
1383 tokens = torch.tensor(tokens)
1384 if len(tokens.shape) == 2:
1385 return self.tokenizer.batch_decode(tokens, clean_up_tokenization_spaces=False)
1386 elif len(tokens.shape) <= 1: 1386 ↛ 1389line 1386 didn't jump to line 1389 because the condition on line 1386 was always true
1387 return self.tokenizer.decode(tokens, clean_up_tokenization_spaces=False)
1388 else:
1389 raise ValueError(f"Invalid shape passed in: {tokens.shape}")
1391 def to_str_tokens(
1392 self,
1393 input: Union[str, torch.Tensor, np.ndarray, List],
1394 prepend_bos: Optional[bool] = None,
1395 padding_side: Optional[str] = None,
1396 ) -> Union[List[str], List[List[str]]]:
1397 """Map text or tokens to a list of tokens as strings.
1399 See the class-level "Tokenization notes" for full ``prepend_bos``
1400 semantics. **Pass ``prepend_bos=False`` whenever you're tokenizing
1401 only part of a prompt.** When ``input`` is already a tensor or
1402 array, ``prepend_bos`` and ``padding_side`` are ignored.
1404 Args:
1405 input: A string, list of strings, or tensor/array of token IDs.
1406 prepend_bos: Overrides ``self.cfg.default_prepend_bos``. Only
1407 applies when ``input`` is a string. Defaults to ``None``
1408 (use the cfg setting).
1409 padding_side: Which side to pad on. Only applies when ``input``
1410 is a string.
1412 Returns:
1413 List of token strings.
1414 """
1415 if isinstance(input, list): 1415 ↛ 1416line 1415 didn't jump to line 1416 because the condition on line 1415 was never true
1416 return cast(
1417 List[List[str]],
1418 [self.to_str_tokens(item, prepend_bos, padding_side) for item in input],
1419 )
1420 elif isinstance(input, str): 1420 ↛ 1422line 1420 didn't jump to line 1422 because the condition on line 1420 was always true
1421 tokens = self.to_tokens(input, prepend_bos=prepend_bos, padding_side=padding_side)[0]
1422 elif isinstance(input, torch.Tensor):
1423 tokens = input.squeeze()
1424 if tokens.dim() == 0:
1425 tokens = tokens.unsqueeze(0)
1426 assert (
1427 tokens.dim() == 1
1428 ), f"Invalid tokens input to to_str_tokens, has shape: {tokens.shape}"
1429 elif isinstance(input, np.ndarray):
1430 tokens_np = input.squeeze()
1431 if tokens_np.ndim == 0:
1432 tokens_np = np.expand_dims(tokens_np, axis=0)
1433 assert (
1434 tokens_np.ndim == 1
1435 ), f"Invalid tokens input to to_str_tokens, has shape: {tokens_np.shape}"
1436 tokens = torch.tensor(tokens_np)
1437 else:
1438 raise ValueError(f"Invalid input type to to_str_tokens: {type(input)}")
1439 # v5 compat: wrap each token so batch_decode decodes them individually
1440 tokens_list = [[int(t)] for t in tokens.tolist()]
1441 str_tokens = self.tokenizer.batch_decode(tokens_list, clean_up_tokenization_spaces=False)
1442 return str_tokens
1444 def to_single_token(self, string: str) -> int:
1445 """Map a string that makes up a single token to the id for that token.
1447 Args:
1448 string: The string to convert
1450 Returns:
1451 Token ID
1453 Raises:
1454 AssertionError: If string is not a single token
1455 """
1456 token = self.to_tokens(string, prepend_bos=False).squeeze()
1457 if token.numel() != 1: 1457 ↛ 1458line 1457 didn't jump to line 1458 because the condition on line 1457 was never true
1458 raise AssertionError(f"Input string: {string} is not a single token!")
1459 return int(token.item())
1461 def get_token_position(
1462 self,
1463 single_token: Union[str, int],
1464 input: Union[str, torch.Tensor],
1465 mode="first",
1466 prepend_bos: Optional[Union[bool, None]] = None,
1467 padding_side: Optional[Union[Literal["left", "right"], None]] = None,
1468 ):
1469 """Get the position of a single_token in a string or sequence of tokens.
1471 Raises an error if the token is not present.
1473 When ``input`` is a string it's tokenized internally — see the
1474 class-level "Tokenization notes" for ``prepend_bos`` semantics.
1475 Off-by-one position errors usually mean ``prepend_bos`` is on
1476 when it shouldn't be (or vice versa); pass ``prepend_bos=False``
1477 when ``input`` is a fragment of a larger prompt.
1479 Args:
1480 single_token (Union[str, int]): The token to search for. Can
1481 be a token index, or a string (but the string must correspond to a single token).
1482 input (Union[str, torch.Tensor]): The sequence to
1483 search in. Can be a string or a rank 1 tensor of tokens or a rank 2 tensor of tokens
1484 with a dummy batch dimension.
1485 mode (str, optional): If there are multiple matches, which match to return. Supports
1486 "first" or "last". Defaults to "first".
1487 prepend_bos (bool, optional): Overrides ``self.cfg.default_prepend_bos``. Only
1488 applies when ``input`` is a string. Defaults to ``None`` (use the cfg setting).
1489 padding_side (Union[Literal["left", "right"], None], optional): Specifies which
1490 side to pad when tokenizing multiple strings of different lengths.
1491 """
1492 if isinstance(input, str):
1493 tokens = self.to_tokens(input, prepend_bos=prepend_bos, padding_side=padding_side)
1494 else:
1495 tokens = input
1496 if len(tokens.shape) == 2:
1497 assert (
1498 tokens.shape[0] == 1
1499 ), f"If tokens are rank two, they must have shape [1, seq_len], not {tokens.shape}"
1500 tokens = tokens[0]
1501 if isinstance(single_token, str):
1502 single_token = self.to_single_token(single_token)
1503 elif isinstance(single_token, torch.Tensor): 1503 ↛ 1504line 1503 didn't jump to line 1504 because the condition on line 1503 was never true
1504 single_token = single_token.item()
1505 indices = torch.arange(len(tokens), device=tokens.device)[tokens == single_token]
1506 assert len(indices) > 0, "The token does not occur in the prompt"
1507 if mode == "first":
1508 return indices[0].item()
1509 elif mode == "last": 1509 ↛ 1512line 1509 didn't jump to line 1512 because the condition on line 1509 was always true
1510 return indices[-1].item()
1511 else:
1512 raise ValueError(f"mode must be 'first' or 'last', not {mode}")
1514 def to_single_str_token(self, int_token: int) -> str:
1515 """Get the single token corresponding to an int in string form.
1517 Args:
1518 int_token: The token ID
1520 Returns:
1521 The token string
1522 """
1523 assert isinstance(int_token, int)
1524 token = self.to_str_tokens(torch.tensor([int_token]))
1525 if isinstance(token, list) and len(token) == 1:
1526 return str(token[0])
1527 raise AssertionError("Expected a single string token.")
1529 def blocks_with(self, submodule: str) -> List[Tuple[int, "GeneralizedComponent"]]:
1530 """Return (index, block) pairs for blocks with the named bridged submodule.
1532 Checks _modules (not hasattr) so HF-internal attrs don't match.
1533 Use instead of assuming blocks[0] is representative on hybrid models.
1534 """
1535 if not hasattr(self, "blocks"):
1536 return []
1537 return [(i, block) for i, block in enumerate(self.blocks) if submodule in block._modules]
1539 def stack_params_for(
1540 self, submodule: str, attr_path: str, reshape_fn: Optional[Callable] = None
1541 ) -> Tuple[List[int], torch.Tensor]:
1542 """Stack a parameter across matching blocks only. Returns (layer_indices, tensor).
1544 Use for hybrid models where not all blocks have the submodule.
1545 """
1546 matching = self.blocks_with(submodule)
1547 if not matching:
1548 raise ValueError(
1549 f"No blocks have submodule '{submodule}'. "
1550 f"Available submodules can be checked with blocks_with()."
1551 )
1552 indices: List[int] = []
1553 weights: List[torch.Tensor] = []
1554 for idx, block in matching:
1555 w = _resolve_attr_path(block, attr_path)
1556 if reshape_fn is not None:
1557 w = reshape_fn(w)
1558 weights.append(w)
1559 indices.append(idx)
1560 return indices, torch.stack(weights, dim=0)
1562 def _stack_block_params(
1563 self, attr_path: str, reshape_fn: Optional[Callable] = None
1564 ) -> torch.Tensor:
1565 """Stack a parameter across all blocks; falls back to matching-only on hybrids.
1567 Filters on FULL-path resolution, not the first segment: on interleaved
1568 MoE, every block has `mlp` but only dense layers expose `mlp.W_in`, so
1569 a first-segment filter matched everything and the sparse layer's
1570 AttributeError killed the accessor for the whole model.
1571 """
1572 first_attr = attr_path.split(".")[0]
1573 matching_blocks = []
1574 for i, block in enumerate(self.blocks):
1575 if first_attr not in block._modules:
1576 continue
1577 try:
1578 weight = _resolve_attr_path(block, attr_path)
1579 except AttributeError:
1580 continue
1581 matching_blocks.append((i, weight))
1583 if len(matching_blocks) == 0:
1584 raise AttributeError(
1585 f"No blocks resolve '{attr_path}'. "
1586 f"Use bridge.blocks_with('{first_attr}') to check availability."
1587 )
1589 if len(matching_blocks) < len(self.blocks):
1590 indices = [i for i, _ in matching_blocks]
1591 logging.warning(
1592 "Hybrid model: only %d/%d blocks resolve '%s'. Returning stacked tensor "
1593 "for layers %s only. Tensor index i corresponds to original layer "
1594 "indices[i], not layer i. For explicit index mapping, use "
1595 "bridge.stack_params_for('%s', '%s').",
1596 len(matching_blocks),
1597 len(self.blocks),
1598 attr_path,
1599 indices,
1600 first_attr,
1601 attr_path,
1602 )
1604 weights: List[torch.Tensor] = []
1605 for _, weight in matching_blocks:
1606 if reshape_fn is not None:
1607 weight = reshape_fn(weight)
1608 weights.append(weight)
1609 # Under a device_map split, per-block tensors live on different devices.
1610 # torch.stack requires a common device; gather onto cfg.device (the embedding /
1611 # input device — a natural "home" for cross-layer reductions).
1612 if getattr(self.cfg, "n_devices", 1) > 1 and weights and self.cfg.device is not None:
1613 target_device = torch.device(self.cfg.device)
1614 weights = [w.to(target_device) for w in weights]
1615 return torch.stack(weights, dim=0)
1617 def _reshape_qkv(self, w: torch.Tensor) -> torch.Tensor:
1618 """Reshape 2D [d_model, d_model] QKV weight to 3D [n_heads, d_model, d_head]."""
1619 if w.shape == (self.cfg.d_model, self.cfg.d_model): 1619 ↛ 1620line 1619 didn't jump to line 1620 because the condition on line 1619 was never true
1620 d_head = self.cfg.d_model // self.cfg.n_heads
1621 return w.reshape(self.cfg.n_heads, self.cfg.d_model, d_head)
1622 return w
1624 def _reshape_o(self, w: torch.Tensor) -> torch.Tensor:
1625 """Reshape 2D [d_model, d_model] O weight to 3D [n_heads, d_head, d_model]."""
1626 if w.shape == (self.cfg.d_model, self.cfg.d_model): 1626 ↛ 1627line 1626 didn't jump to line 1627 because the condition on line 1626 was never true
1627 d_head = self.cfg.d_model // self.cfg.n_heads
1628 return w.reshape(self.cfg.n_heads, d_head, self.cfg.d_model)
1629 return w
1631 def _expand_kv_heads(self, w: torch.Tensor) -> torch.Tensor:
1632 """Expand stacked grouped K/V weights along the head axis to n_heads.
1634 GQA models store one K/V projection per key-value head while W_Q/W_O are
1635 per-query-head, so weight circuits must repeat the grouped K/V up to
1636 n_heads before factoring: query head h reads kv head
1637 h // (n_heads // n_kv_heads), i.e. repeat_interleave — the same layout
1638 GroupedQueryAttention.W_K/W_V expose on HookedTransformer. No-op for MHA,
1639 where the head axes already match.
1640 """
1641 if w.ndim != 4 or w.shape[1] == self.cfg.n_heads:
1642 return w
1643 n_kv_heads = w.shape[1]
1644 if self.cfg.n_heads % n_kv_heads != 0:
1645 raise ValueError(
1646 f"Cannot expand {n_kv_heads} key-value heads to {self.cfg.n_heads} "
1647 f"query heads: n_heads must be a multiple of n_kv_heads."
1648 )
1649 return w.repeat_interleave(self.cfg.n_heads // n_kv_heads, dim=1)
1651 @property
1652 def W_K(self) -> torch.Tensor:
1653 """Stack the key weights across all layers."""
1654 return self._stack_block_params("attn.W_K", self._reshape_qkv)
1656 @property
1657 def W_Q(self) -> torch.Tensor:
1658 """Stack the query weights across all layers."""
1659 return self._stack_block_params("attn.W_Q", self._reshape_qkv)
1661 @property
1662 def W_V(self) -> torch.Tensor:
1663 """Stack the value weights across all layers."""
1664 return self._stack_block_params("attn.W_V", self._reshape_qkv)
1666 @property
1667 def W_O(self) -> torch.Tensor:
1668 """Stack the attn output weights across all layers."""
1669 return self._stack_block_params("attn.W_O", self._reshape_o)
1671 @property
1672 def W_in(self) -> torch.Tensor:
1673 """Stack the MLP input weights across all layers."""
1674 return self._stack_block_params("mlp.W_in")
1676 @property
1677 def W_gate(self) -> Union[torch.Tensor, None]:
1678 """Stack the MLP gate weights across all layers (gated MLPs only)."""
1679 if getattr(self.cfg, "gated_mlp", False):
1680 return self._stack_block_params("mlp.W_gate")
1681 return None
1683 @property
1684 def W_out(self) -> torch.Tensor:
1685 """Stack the MLP output weights across all layers."""
1686 return self._stack_block_params("mlp.W_out")
1688 @property
1689 def b_K(self) -> torch.Tensor:
1690 """Stack the key biases across all layers."""
1691 return self._stack_block_params("attn.b_K")
1693 @property
1694 def b_Q(self) -> torch.Tensor:
1695 """Stack the query biases across all layers."""
1696 return self._stack_block_params("attn.b_Q")
1698 @property
1699 def b_V(self) -> torch.Tensor:
1700 """Stack the value biases across all layers."""
1701 return self._stack_block_params("attn.b_V")
1703 @property
1704 def b_O(self) -> torch.Tensor:
1705 """Stack the attn output biases across all layers."""
1706 return self._stack_block_params("attn.b_O")
1708 @property
1709 def b_in(self) -> torch.Tensor:
1710 """Stack the MLP input biases across all layers."""
1711 return self._stack_block_params("mlp.b_in")
1713 @property
1714 def b_out(self) -> torch.Tensor:
1715 """Stack the MLP output biases across all layers."""
1716 return self._stack_block_params("mlp.b_out")
1718 @property
1719 def W_U(self) -> torch.Tensor:
1720 """Unembedding matrix (d_model, d_vocab). Maps residual stream to logits."""
1721 return self.unembed.W_U
1723 @property
1724 def b_U(self) -> torch.Tensor:
1725 """Unembedding bias (d_vocab)."""
1726 return self.unembed.b_U
1728 @property
1729 def W_E(self) -> torch.Tensor:
1730 """Token embedding matrix (d_vocab, d_model)."""
1731 return self.embed.W_E
1733 @property
1734 def QK(self):
1735 """QK circuit. On hybrids, returns attn layers only (with warning). See QK_for_attn_layers()."""
1736 return FactoredMatrix(self.W_Q, self._expand_kv_heads(self.W_K).transpose(-2, -1))
1738 @property
1739 def OV(self):
1740 """OV circuit. On hybrids, returns attn layers only (with warning). See OV_for_attn_layers()."""
1741 return FactoredMatrix(self._expand_kv_heads(self.W_V), self.W_O)
1743 def QK_for_attn_layers(self) -> Tuple[List[int], FactoredMatrix]:
1744 """QK circuit for attention layers only. Returns (layer_indices, FactoredMatrix)."""
1745 q_indices, W_Q = self.stack_params_for("attn", "attn.W_Q", self._reshape_qkv)
1746 _, W_K = self.stack_params_for("attn", "attn.W_K", self._reshape_qkv)
1747 return q_indices, FactoredMatrix(W_Q, self._expand_kv_heads(W_K).transpose(-2, -1))
1749 def OV_for_attn_layers(self) -> Tuple[List[int], FactoredMatrix]:
1750 """OV circuit for attention layers only. Returns (layer_indices, FactoredMatrix)."""
1751 v_indices, W_V = self.stack_params_for("attn", "attn.W_V", self._reshape_qkv)
1752 _, W_O = self.stack_params_for("attn", "attn.W_O", self._reshape_o)
1753 return v_indices, FactoredMatrix(self._expand_kv_heads(W_V), W_O)
1755 # ------------------------------------------------------------------
1756 # Mechanistic interpretability analysis methods
1757 # ------------------------------------------------------------------
1759 def tokens_to_residual_directions(
1760 self,
1761 tokens: Union[str, int, torch.Tensor],
1762 ) -> torch.Tensor:
1763 """Map tokens to their unembedding vectors (residual stream directions).
1765 Returns the columns of W_U corresponding to the given tokens — i.e. the
1766 directions in the residual stream that the model dots with to produce the
1767 logit for each token.
1769 WARNING: If you use this without folding in LayerNorm (compatibility mode),
1770 the results will be misleading because LN weights change the unembed map.
1772 Args:
1773 tokens: A single token (str, int, or scalar tensor), a 1-D tensor of
1774 token IDs, or a 2-D batch of token IDs.
1776 Returns:
1777 Tensor of unembedding vectors with shape matching the input token shape
1778 plus a trailing d_model dimension.
1779 """
1780 if isinstance(tokens, torch.Tensor) and tokens.numel() > 1:
1781 residual_directions = self.W_U[:, tokens]
1782 residual_directions = einops.rearrange(
1783 residual_directions, "d_model ... -> ... d_model"
1784 )
1785 return residual_directions
1786 else:
1787 if isinstance(tokens, str):
1788 token = self.to_single_token(tokens)
1789 elif isinstance(tokens, int): 1789 ↛ 1791line 1789 didn't jump to line 1791 because the condition on line 1789 was always true
1790 token = tokens
1791 elif isinstance(tokens, torch.Tensor) and tokens.numel() == 1:
1792 token = int(tokens.item())
1793 else:
1794 raise ValueError(f"Invalid token type: {type(tokens)}")
1795 residual_direction = self.W_U[:, token]
1796 return residual_direction
1798 # Variant → attr paths for the output bias that feeds the residual stream.
1799 _VARIANT_OUTPUT_BIAS_ATTRS: Dict[str, tuple] = {
1800 "attn": ("b_O",),
1801 "linear_attn": ("out_proj.bias",),
1802 "mamba": ("out_proj.bias",),
1803 "mixer": ("out_proj.bias",),
1804 "ssm": ("out_proj.bias",),
1805 }
1807 def _get_block_variant_bias(self, block: "GeneralizedComponent") -> Optional[torch.Tensor]:
1808 """Return the output bias from this block's variant submodule, or None."""
1809 for name in VARIANT_SUBMODULE_NAMES:
1810 if name not in block._modules:
1811 continue
1812 variant = block._modules[name]
1813 for attr_path in self._VARIANT_OUTPUT_BIAS_ATTRS.get(name, ()): 1813 ↛ 1809line 1813 didn't jump to line 1809 because the loop on line 1813 didn't complete
1814 obj = variant
1815 try:
1816 for attr in attr_path.split("."):
1817 obj = getattr(obj, attr)
1818 except AttributeError:
1819 continue
1820 if obj is not None and isinstance(obj, torch.Tensor): 1820 ↛ 1813line 1820 didn't jump to line 1813 because the condition on line 1820 was always true
1821 return obj
1822 return None
1824 def accumulated_bias(
1825 self,
1826 layer: int,
1827 mlp_input: bool = False,
1828 include_mlp_biases: bool = True,
1829 ) -> torch.Tensor:
1830 """Sum of variant + MLP output biases through the residual stream up to `layer`.
1832 Includes all layer types (attn, SSM, linear-attn). Set mlp_input=True
1833 to include the variant bias of the target layer itself.
1834 """
1835 accumulated = torch.zeros(self.cfg.d_model, device=self.cfg.device)
1836 for i in range(layer):
1837 block = self.blocks[i]
1838 b_O = self._get_block_variant_bias(block)
1839 if b_O is not None:
1840 accumulated = accumulated + b_O.to(accumulated.device)
1841 if include_mlp_biases and "mlp" in block._modules:
1842 b_out = getattr(block.mlp, "b_out", None)
1843 if b_out is not None: 1843 ↛ 1836line 1843 didn't jump to line 1836 because the condition on line 1843 was always true
1844 accumulated = accumulated + b_out.to(accumulated.device)
1845 if mlp_input:
1846 assert layer < self.cfg.n_layers, "Cannot include attn_bias from beyond the final layer"
1847 block = self.blocks[layer]
1848 b_O = self._get_block_variant_bias(block)
1849 if b_O is not None:
1850 accumulated = accumulated + b_O.to(accumulated.device)
1851 return accumulated
1853 def all_composition_scores(self, mode: str) -> CompositionScores:
1854 """Composition scores for all attention head pairs. Returns CompositionScores.
1856 See https://transformer-circuits.pub/2021/framework/index.html
1857 On hybrid models, only attention layers are included; layer_indices
1858 maps tensor position i to original layer number.
1859 """
1860 attn_blocks = self.blocks_with("attn")
1861 if not attn_blocks: 1861 ↛ 1862line 1861 didn't jump to line 1862 because the condition on line 1861 was never true
1862 raise ValueError("No attention layers found — cannot compute composition scores.")
1864 indices = [idx for idx, _ in attn_blocks]
1865 blocks_list = [block for _, block in attn_blocks]
1867 def _stack(attr_path: str, reshape_fn: Optional[Callable] = None) -> torch.Tensor:
1868 weights: List[torch.Tensor] = []
1869 for block in blocks_list:
1870 w = _resolve_attr_path(block, attr_path)
1871 if reshape_fn is not None: 1871 ↛ 1873line 1871 didn't jump to line 1873 because the condition on line 1871 was always true
1872 w = reshape_fn(w)
1873 weights.append(w)
1874 # See _stack_block_params: gather per-block tensors onto cfg.device when split.
1875 if getattr(self.cfg, "n_devices", 1) > 1 and weights and self.cfg.device is not None: 1875 ↛ 1876line 1875 didn't jump to line 1876 because the condition on line 1875 was never true
1876 target_device = torch.device(self.cfg.device)
1877 weights = [w.to(target_device) for w in weights]
1878 return torch.stack(weights, dim=0)
1880 W_V = self._expand_kv_heads(_stack("attn.W_V", self._reshape_qkv))
1881 W_O = _stack("attn.W_O", self._reshape_o)
1882 left = FactoredMatrix(W_V, W_O)
1884 if mode == "Q":
1885 W_Q = _stack("attn.W_Q", self._reshape_qkv)
1886 W_K = self._expand_kv_heads(_stack("attn.W_K", self._reshape_qkv))
1887 right = FactoredMatrix(W_Q, W_K.transpose(-2, -1))
1888 elif mode == "K":
1889 W_Q = _stack("attn.W_Q", self._reshape_qkv)
1890 W_K = self._expand_kv_heads(_stack("attn.W_K", self._reshape_qkv))
1891 right = FactoredMatrix(W_Q, W_K.transpose(-2, -1)).T
1892 elif mode == "V":
1893 right = left
1894 else:
1895 raise ValueError(f"mode must be one of ['Q', 'K', 'V'] not {mode}")
1897 scores = utils.composition_scores(left, right, broadcast_dims=True)
1898 n_attn = len(indices)
1899 idx_tensor = torch.arange(n_attn, device=self.cfg.device)
1900 mask = idx_tensor[:, None, None, None] < idx_tensor[None, None, :, None]
1901 scores = torch.where(mask, scores, torch.zeros_like(scores))
1903 labels = [f"L{l}H{h}" for l in indices for h in range(self.cfg.n_heads)]
1904 return CompositionScores(scores=scores, layer_indices=indices, head_labels=labels)
1906 def composition_layer_indices(self) -> List[int]:
1907 """Original layer indices for attention layers (maps composition score positions)."""
1908 return [idx for idx, _ in self.blocks_with("attn")]
1910 def block_hooks(self, layer_idx: int) -> List[str]:
1911 """Sorted hook names available on block `layer_idx` (block-relative paths)."""
1912 prefix = f"blocks.{layer_idx}."
1913 return sorted(name[len(prefix) :] for name in self.hook_dict if name.startswith(prefix))
1915 def block_submodules(self, layer_idx: int) -> List[str]:
1916 """Return bridged submodule names on block `layer_idx`."""
1917 block = self.blocks[layer_idx]
1918 return [name for name in block._modules if name not in _BLOCK_INTERNAL_MODULES]
1920 def layer_types(self) -> List[str]:
1921 """Per-block type labels, e.g. ["attn+mlp", "ssm+mlp", ...]. Deterministic order."""
1922 types = []
1923 for block in self.blocks:
1924 variants = [n for n in VARIANT_SUBMODULE_NAMES if n in block._modules]
1925 universals = sorted(
1926 n
1927 for n in block._modules
1928 if n not in _VARIANT_SUBMODULE_SET
1929 and n not in _BLOCK_INTERNAL_MODULES
1930 and not n.startswith(_NORM_PREFIXES)
1931 )
1932 parts = variants + universals
1933 types.append("+".join(parts) if parts else "unknown")
1934 return types
1936 @property
1937 def all_head_labels(self) -> list[str]:
1938 """Human-readable labels for all attention heads, e.g. ['L0H0', 'L0H1', ...]."""
1939 return [f"L{l}H{h}" for l in range(self.cfg.n_layers) for h in range(self.cfg.n_heads)]
1941 @property
1942 def attn_head_labels(self) -> list[str]:
1943 """Head labels for attention layers only — matches all_composition_scores() dims."""
1944 return [
1945 f"L{l}H{h}" for l in self.composition_layer_indices() for h in range(self.cfg.n_heads)
1946 ]
1948 def parameters(self, recurse: bool = True) -> Iterator[nn.Parameter]:
1949 """Returns parameters following standard PyTorch semantics.
1951 This method delegates to the underlying HuggingFace model's parameters().
1952 For TransformerLens-style parameter generator, use tl_parameters() instead.
1954 Args:
1955 recurse: If True, yields parameters of this module and all submodules
1957 Returns:
1958 Iterator of nn.Parameter objects
1959 """
1960 return self.original_model.parameters(recurse=recurse)
1962 def named_parameters(
1963 self, prefix: str = "", recurse: bool = True, remove_duplicate: bool = True
1964 ) -> Iterator[tuple[str, nn.Parameter]]:
1965 """Returns named parameters following standard PyTorch semantics.
1967 This method delegates to the underlying HuggingFace model's named_parameters().
1968 For TransformerLens-style generator, use tl_named_parameters() instead.
1970 Args:
1971 prefix: Prefix to prepend to all parameter names
1972 recurse: If True, yields parameters of this module and all submodules
1973 remove_duplicate: If True, removes duplicate parameters
1975 Returns:
1976 Iterator of (name, parameter) tuples
1977 """
1978 return self.original_model.named_parameters(prefix, recurse, remove_duplicate)
1980 def tl_parameters(self) -> dict[str, torch.Tensor]:
1981 """Returns TransformerLens-style parameter dictionary.
1983 Parameter names follow TransformerLens conventions (e.g., 'blocks.0.attn.W_Q') and may
1984 include processed weights (non-leaf tensors). This format is expected by SVDInterpreter
1985 among other analysis tools.
1987 Returns:
1988 Dictionary mapping TransformerLens parameter names to tensors
1990 Example:
1991 >>> bridge = TransformerBridge.boot_transformers("gpt2")
1992 >>> tl_params = bridge.tl_parameters()
1993 >>> W_Q = tl_params["blocks.0.attn.W_Q"] # Shape: [n_heads, d_model, d_head]
1994 """
1995 return self.get_params()
1997 def tl_named_parameters(self) -> Iterator[tuple[str, torch.Tensor]]:
1998 """Returns iterator of TransformerLens-style named parameters.
2000 This provides the same parameters as tl_parameters() but as an iterator
2001 for consistency with PyTorch's named_parameters() API pattern.
2003 Returns:
2004 Iterator of (name, tensor) tuples with TransformerLens naming conventions
2006 Example:
2007 >>> bridge = TransformerBridge.boot_transformers("gpt2")
2008 >>> for name, param in bridge.tl_named_parameters():
2009 ... if "attn.W_Q" in name:
2010 ... print(f"{name}: {param.shape}") # doctest: +ELLIPSIS
2011 blocks.0.attn.W_Q: torch.Size([12, 768, 64])
2012 ...
2013 """
2014 return iter(self.get_params().items())
2016 def _accepts_derived_position_ids(self) -> bool:
2017 """Whether it is safe to hand the wrapped model a mask-derived ``position_ids``.
2019 Two families of model must be left alone, so the injection below is
2020 gated on the target the same way ``output_attentions`` is in
2021 :meth:`run_with_cache`:
2023 * **Fixed-signature models.** Remote-code forwards such as
2024 ``LLaDAModelLM.forward`` take neither ``position_ids`` nor
2025 ``**kwargs``, so passing it raises ``TypeError`` where the model
2026 previously returned logits.
2027 * **Models that own their position derivation.** mRoPE architectures
2028 (Qwen2-VL, Qwen2.5-VL, Qwen3-VL, GLM-4V) build a 3-D temporal /
2029 height / width index in ``get_rope_index``, and only while
2030 ``position_ids is None``; a supplied 2-D tensor is silently expanded
2031 across all three streams instead. Their derivation already scatters
2032 positions onto attended slots only, so it handles left padding
2033 correctly on its own and needs no help from us.
2034 * **Mask-consuming positional embeddings.** OPT's
2035 ``OPTLearnedPositionalEmbedding.forward`` takes the mask and derives
2036 the same positions we would, so injection buys nothing — but it does
2037 replace the model's own padding-slot convention with ours, which
2038 shows up as a whole-tensor diff.
2039 """
2040 underlying = getattr(self, "original_model", None)
2041 if underlying is None:
2042 return False
2044 cached = self.__dict__.get("_derived_position_ids_ok")
2045 if cached is not None and cached[0] is underlying:
2046 return bool(cached[1])
2048 def verdict() -> bool:
2049 fwd_params = inspect.signature(underlying.forward).parameters
2050 if "position_ids" not in fwd_params and not any(
2051 p.kind is inspect.Parameter.VAR_KEYWORD for p in fwd_params.values()
2052 ):
2053 return False
2055 # ``get_rope_index`` lives on the inner text model, not the
2056 # ForConditionalGeneration wrapper that is usually original_model.
2057 for module in (
2058 underlying,
2059 getattr(underlying, "model", None),
2060 getattr(underlying, "language_model", None),
2061 ):
2062 if module is not None and hasattr(module, "get_rope_index"):
2063 return False
2065 # Config-level backstop for mRoPE models that spell the derivation
2066 # differently; the section list is what makes positions 3-D.
2067 config = getattr(underlying, "config", None)
2068 for candidate in (config, getattr(config, "text_config", None)):
2069 scaling = getattr(candidate, "rope_scaling", None)
2070 if isinstance(scaling, dict) and "mrope_section" in scaling:
2071 return False
2073 # A positional embedding that takes the mask derives positions for
2074 # itself. Only embeddings that override nn.Embedding.forward are
2075 # worth inspecting, which keeps this to a handful per model.
2076 for module in underlying.modules():
2077 if not isinstance(module, nn.Embedding):
2078 continue
2079 if type(module).forward is nn.Embedding.forward:
2080 continue
2081 if "attention_mask" in inspect.signature(module.forward).parameters:
2082 return False
2083 return True
2085 accepts = verdict()
2086 self.__dict__["_derived_position_ids_ok"] = (underlying, accepts)
2087 return accepts
2089 def forward(
2090 self,
2091 input: Union[str, List[str], torch.Tensor],
2092 return_type: Optional[str] = "logits",
2093 loss_per_token: bool = False,
2094 prepend_bos: Optional[bool] = None,
2095 padding_side: Optional[str] = None,
2096 attention_mask: Optional[torch.Tensor] = None,
2097 labels: Optional[torch.Tensor] = None,
2098 start_at_layer: Optional[int] = None,
2099 stop_at_layer: Optional[int] = None,
2100 pixel_values: Optional[torch.Tensor] = None,
2101 input_values: Optional[torch.Tensor] = None,
2102 **kwargs,
2103 ) -> Any:
2104 """Forward pass through the model.
2106 Args:
2107 input: Input to the model
2108 return_type: Type of output to return ('logits', 'loss', 'both', 'predictions', None)
2109 loss_per_token: Whether to return loss per token
2110 prepend_bos: Whether to prepend BOS token
2111 padding_side: Which side to pad on
2112 labels: Explicit language-model targets. Encoder-decoder models require
2113 labels for loss; decoder-only models fall back to input IDs when omitted.
2114 start_at_layer: Not implemented in TransformerBridge. The bridge delegates
2115 to HuggingFace's model.forward() which owns the layer iteration loop,
2116 making start_at_layer infeasible without monkey-patching HF internals
2117 (fragile across HF versions) or exception-based layer skipping (corrupts
2118 model state). Raises NotImplementedError if a non-None value is passed.
2119 stop_at_layer: Layer to stop forward pass at
2120 pixel_values: Optional image tensor for multimodal models (e.g., LLaVA, Gemma3)
2121 and vision models (eg. ViT, DeiT).
2122 The tensor is passed directly to the underlying HuggingFace model.
2123 Only valid when cfg.is_multimodal is True or cfg.is_visual_model is True.
2124 input_values: Optional audio waveform tensor for audio models (e.g., HuBERT).
2125 The tensor is passed directly to the underlying HuggingFace model.
2126 Only valid when cfg.is_audio_model is True.
2127 **kwargs: Additional arguments passed to model
2129 Returns:
2130 Model output based on return_type
2131 """
2133 model_config = getattr(self.original_model, "config", None)
2134 is_encoder_decoder = bool(getattr(model_config, "is_encoder_decoder", False))
2135 if return_type in ("loss", "both") and is_encoder_decoder and labels is None:
2136 raise ValueError(
2137 "labels are required for seq2seq return_type='loss' or 'both'; "
2138 "encoder input_ids are not decoder targets"
2139 )
2140 if (
2141 return_type in ("loss", "both")
2142 and not is_encoder_decoder
2143 and not self.adapter.supports_causal_loss
2144 ):
2145 architecture = self.cfg.architecture or type(self.adapter).__name__
2146 raise NotImplementedError(
2147 f"{architecture} does not support TransformerBridge's shifted causal "
2148 "loss. Request return_type='logits' and compute the architecture-specific "
2149 "masked-token objective explicitly."
2150 )
2152 if labels is not None:
2153 kwargs["labels"] = labels
2155 if start_at_layer is not None: 2155 ↛ 2156line 2155 didn't jump to line 2156 because the condition on line 2155 was never true
2156 raise NotImplementedError(
2157 "start_at_layer is not supported in TransformerBridge. "
2158 "The bridge delegates to HuggingFace's model.forward() which controls "
2159 "the layer iteration loop. See the TransformerBridge review plan for a "
2160 "detailed analysis of implementation approaches and their tradeoffs."
2161 )
2163 # Set stop_at_layer flag on all blocks if requested
2164 if stop_at_layer is not None:
2165 if ( 2165 ↛ 2171line 2165 didn't jump to line 2171 because the condition on line 2165 was never true
2166 hasattr(self, "L_blocks")
2167 or hasattr(self, "H_blocks")
2168 or hasattr(self, "encoder_blocks")
2169 or hasattr(self, "decoder_blocks")
2170 ):
2171 raise NotImplementedError(
2172 "stop_at_layer is not supported on non-standard block list "
2173 "names (L_blocks, H_blocks, encoder_blocks, decoder_blocks). "
2174 "The bridge only supports stop_at_layer on 'blocks'."
2175 )
2176 if hasattr(self, "blocks"): 2176 ↛ 2184line 2176 didn't jump to line 2184 because the condition on line 2176 was always true
2177 effective_stop_at_layer = (
2178 len(self.blocks) + stop_at_layer if stop_at_layer < 0 else stop_at_layer
2179 )
2180 for block in self.blocks:
2181 block._stop_at_layer_idx = effective_stop_at_layer
2183 # Map HookedEncoderDecoder-style kwargs to HF-compatible names
2184 if "decoder_input" in kwargs:
2185 kwargs["decoder_input_ids"] = kwargs.pop("decoder_input")
2186 if "one_zero_attention_mask" in kwargs: 2186 ↛ 2187line 2186 didn't jump to line 2187 because the condition on line 2186 was never true
2187 if attention_mask is None:
2188 attention_mask = kwargs.pop("one_zero_attention_mask")
2189 else:
2190 kwargs.pop("one_zero_attention_mask")
2192 # Detect batched list input that may need padding. Forward follows the
2193 # requested/tokenizer side; generation separately forces left-padding.
2194 _is_batched_list = (
2195 isinstance(input, list)
2196 and len(input) > 1
2197 and not getattr(self.cfg, "is_audio_model", False)
2198 and not getattr(self.cfg, "is_visual_model", False)
2199 )
2200 _resolved_padding_side = padding_side
2201 if _resolved_padding_side is None and self.tokenizer is not None:
2202 _resolved_padding_side = getattr(self.tokenizer, "padding_side", "right")
2204 try:
2205 if isinstance(input, (str, list)):
2206 if getattr(self.cfg, "is_audio_model", False): 2206 ↛ 2207line 2206 didn't jump to line 2207 because the condition on line 2206 was never true
2207 raise ValueError(
2208 "Audio models require tensor input (raw waveform), not text. "
2209 "Pass a torch.Tensor or use the input_values parameter."
2210 )
2211 if getattr(self.cfg, "is_visual_model", False): 2211 ↛ 2212line 2211 didn't jump to line 2212 because the condition on line 2211 was never true
2212 raise ValueError(
2213 "Visual models require tensor input (pixel values), not text. "
2214 "Pass a torch.Tensor or use the pixel_values parameter."
2215 )
2216 input_ids = self.to_tokens(
2217 input, prepend_bos=prepend_bos, padding_side=padding_side
2218 )
2219 else:
2220 input_ids = input
2221 # Promote 1D integer token tensors to 2D [batch=1, seq] to match
2222 # HookedTransformer's contract. Float tensors (inputs_embeds,
2223 # audio waveforms) are passed through unchanged.
2224 if (
2225 isinstance(input_ids, torch.Tensor)
2226 and input_ids.ndim == 1
2227 and not input_ids.is_floating_point()
2228 ):
2229 input_ids = input_ids.unsqueeze(0)
2231 # Detect inputs_embeds: if the tensor is floating point, it's pre-computed
2232 # embeddings (e.g., from multimodal models) rather than token IDs.
2233 _is_inputs_embeds = (
2234 isinstance(input_ids, torch.Tensor) and input_ids.is_floating_point()
2235 )
2237 # Left padding needs a mask and corrected positions. Right padding is
2238 # harmless for causal real-token positions and remains unmasked to
2239 # match HookedTransformer; bidirectional/encoder inputs still need it.
2240 if (
2241 _is_batched_list
2242 and attention_mask is None
2243 and self.tokenizer is not None
2244 and self.tokenizer.pad_token_id is not None
2245 and not _is_inputs_embeds
2246 and (
2247 _resolved_padding_side == "left"
2248 or is_encoder_decoder
2249 or not self.adapter.supports_causal_loss
2250 )
2251 ):
2252 attention_mask = utils.get_attention_mask(
2253 self.tokenizer,
2254 input_ids,
2255 prepend_bos=getattr(self.cfg, "default_prepend_bos", True),
2256 padding_side=_resolved_padding_side,
2257 ).to(self.cfg.device)
2258 # Gated on the target for the same reason the derivation below is:
2259 # a fixed-signature forward raises TypeError on the kwarg, and a
2260 # model that owns its own position derivation is overridden by it
2261 # (#1626).
2262 if "position_ids" not in kwargs and self._accepts_derived_position_ids():
2263 position_ids = attention_mask.long().cumsum(-1) - 1
2264 position_ids.masked_fill_(attention_mask == 0, 1)
2265 kwargs["position_ids"] = position_ids
2267 # Any masked-out token shifts the absolute position of every real token
2268 # after it, so positions must be derived from the mask rather than left
2269 # to HF's default arange. This is the same derivation HookedTransformer
2270 # applies in pos_embed; without it the bridge silently returns wrong
2271 # logits. An all-ones mask reduces to arange, so this is a no-op there.
2272 #
2273 # The mask spans any cached prefix as well as the new tokens, so it is
2274 # offset back to just the tokens actually being passed — matching how
2275 # AbstractAttention/PosEmbed use past_kv_pos_offset.
2276 if (
2277 attention_mask is not None
2278 and "position_ids" not in kwargs
2279 and not _is_inputs_embeds
2280 and attention_mask.ndim == 2
2281 and isinstance(input_ids, torch.Tensor)
2282 and input_ids.ndim == 2
2283 and attention_mask.shape[1] >= input_ids.shape[1]
2284 and self._accepts_derived_position_ids()
2285 ):
2286 # .long() because callers may hand in a float 0/1 mask, and
2287 # positions index an embedding table.
2288 _derived = utils.get_offset_position_ids(0, attention_mask.long())
2289 _arange = torch.arange(attention_mask.shape[1], device=_derived.device)
2290 # Decide per row, not per batch. A row only needs the derived
2291 # positions when its mask actually moves one of its attended
2292 # tokens off the default position — i.e. a masked token precedes
2293 # a real one (left padding, or an interior gap). Rows that are
2294 # unpadded or purely right-padded keep arange verbatim, so one
2295 # left-padded row in a batch cannot perturb its neighbours.
2296 _needs = ((_derived != _arange) & (attention_mask != 0)).any(dim=1, keepdim=True)
2297 if bool(_needs.any()):
2298 _positions = torch.where(_needs, _derived, _arange.expand_as(_derived))
2299 kwargs["position_ids"] = _positions[
2300 :, attention_mask.shape[1] - input_ids.shape[1] :
2301 ]
2303 if attention_mask is not None:
2304 kwargs["attention_mask"] = attention_mask
2305 if kwargs.pop("use_past_kv_cache", False) or kwargs.get("use_cache", False):
2306 kwargs["use_cache"] = True
2307 # Auto-generate decoder_input_ids for encoder-decoder models
2308 if "decoder_input_ids" not in kwargs and labels is None and is_encoder_decoder:
2309 decoder_start_token_id = getattr(
2310 self.original_model.config, "decoder_start_token_id", None
2311 )
2312 if decoder_start_token_id is not None:
2313 shifted = input_ids[:, :-1]
2314 start_tokens = torch.full(
2315 (input_ids.shape[0], 1),
2316 decoder_start_token_id,
2317 dtype=input_ids.dtype,
2318 device=input_ids.device,
2319 )
2320 kwargs["decoder_input_ids"] = torch.cat([start_tokens, shifted], dim=1)
2321 else:
2322 kwargs["decoder_input_ids"] = input_ids
2324 # Tell PosEmbedBridge to expand batch=1 position_ids to full batch.
2325 if hasattr(self, "pos_embed"):
2326 self.pos_embed._current_batch_size = input_ids.shape[0]
2328 # Handle pixel_values for multimodal models
2329 if pixel_values is not None:
2330 if not (
2331 getattr(self.cfg, "is_multimodal", False)
2332 or getattr(self.cfg, "is_visual_model", False)
2333 ):
2334 raise ValueError(
2335 "pixel_values can only be passed to multimodal or vision models "
2336 "(cfg.is_multimodal or cfg.is_visual_model must be True)"
2337 )
2338 kwargs["pixel_values"] = pixel_values
2340 # Handle input_values for audio models
2341 if input_values is not None: 2341 ↛ 2342line 2341 didn't jump to line 2342 because the condition on line 2341 was never true
2342 if not getattr(self.cfg, "is_audio_model", False):
2343 raise ValueError(
2344 "input_values can only be passed to audio models "
2345 "(cfg.is_audio_model must be True)"
2346 )
2347 kwargs["input_values"] = input_values
2349 # Audio models use input_values (waveform), not input_ids
2350 if getattr(self.cfg, "is_audio_model", False):
2351 if input_values is not None: 2351 ↛ 2352line 2351 didn't jump to line 2352 because the condition on line 2351 was never true
2352 output = self.original_model(**kwargs)
2353 elif isinstance(input, torch.Tensor): 2353 ↛ 2357line 2353 didn't jump to line 2357 because the condition on line 2353 was always true
2354 kwargs["input_values"] = input
2355 output = self.original_model(**kwargs)
2356 else:
2357 raise ValueError(
2358 "Audio models require tensor input (raw waveform). "
2359 "Pass a torch.Tensor or use input_values parameter."
2360 )
2361 elif getattr(self.cfg, "is_visual_model", False):
2362 # "pixel_values" may already be in kwargs from the "if pixel_values is not None:"
2363 # gate above (explicit pixel_values=... call); otherwise treat `input` itself as
2364 # the image tensor, matching how the audio branch treats `input` as the waveform.
2365 if "pixel_values" not in kwargs: 2365 ↛ 2373line 2365 didn't jump to line 2373 because the condition on line 2365 was always true
2366 if isinstance(input, torch.Tensor): 2366 ↛ 2369line 2366 didn't jump to line 2369 because the condition on line 2366 was always true
2367 kwargs["pixel_values"] = input
2368 else:
2369 raise ValueError(
2370 "Visual models require tensor input (pixel values). "
2371 "Pass a torch.Tensor as `input` or use the pixel_values parameter."
2372 )
2373 output = self.original_model(**kwargs)
2374 elif _is_inputs_embeds:
2375 output = self.original_model(inputs_embeds=input_ids, **kwargs)
2376 else:
2377 output = self.original_model(input_ids, **kwargs)
2378 # Stash only the cache object (not the full output) for generate().
2379 if getattr(self, "_capture_hf_cache", False):
2380 self._last_hf_cache = getattr(output, "past_key_values", None)
2381 if hasattr(output, "logits"):
2382 logits = output.logits
2383 elif isinstance(output, tuple) and len(output) > 0:
2384 # With labels forwarded, HF tuple outputs are (loss, logits, ...).
2385 if labels is not None and len(output) > 1: 2385 ↛ 2388line 2385 didn't jump to line 2388 because the condition on line 2385 was always true
2386 logits = output[1]
2387 else:
2388 logits = output[0]
2389 elif hasattr(output, "last_hidden_state"):
2390 # Bare encoder models (ViTModel, DeiTModel, BertModel, etc. without
2391 # a task head) return e.g. BaseModelOutput/BaseModelOutputWithPooling,
2392 # which has neither `.logits` nor tuple semantics. Fall back to
2393 # `last_hidden_state` so return_type="logits" still yields a plain
2394 # tensor rather than silently handing back the raw HF output object.
2395 logits = output.last_hidden_state
2396 else:
2397 logits = output
2398 if return_type == "logits":
2399 return logits
2400 elif return_type == "logits_and_cache":
2401 past_key_values = getattr(output, "past_key_values", None)
2402 return (logits, past_key_values)
2403 elif is_encoder_decoder and return_type in ("loss", "both"):
2404 assert isinstance(
2405 logits, torch.Tensor
2406 ), f"Expected seq2seq logits tensor, got {type(logits)}"
2407 assert isinstance(labels, torch.Tensor)
2408 return self._finalize_seq2seq_return(
2409 return_type,
2410 logits,
2411 labels,
2412 output,
2413 loss_per_token=loss_per_token,
2414 )
2415 elif return_type == "loss":
2416 if getattr(self.cfg, "is_audio_model", False): 2416 ↛ 2417line 2416 didn't jump to line 2417 because the condition on line 2416 was never true
2417 raise ValueError(
2418 "Audio models do not support return_type='loss'. "
2419 "CTC loss requires aligned frame-level labels."
2420 )
2421 if getattr(self.cfg, "is_visual_model", False): 2421 ↛ 2422line 2421 didn't jump to line 2422 because the condition on line 2421 was never true
2422 raise ValueError(
2423 "Vision classification models do not support return_type='loss' "
2424 "via this path (no next-token LM target exists for image "
2425 "classification). Compute cross-entropy against `labels` "
2426 "yourself from the returned logits, or use hf_generate()-style "
2427 "direct access to self.original_model for HF's own loss."
2428 )
2429 if _is_inputs_embeds and labels is None: 2429 ↛ 2430line 2429 didn't jump to line 2430 because the condition on line 2429 was never true
2430 raise ValueError(
2431 "Cannot compute loss with inputs_embeds — token IDs required for labels."
2432 )
2433 # Always use self.loss_fn for consistency with HT's formula
2434 # (log_softmax + gather). HF's output.loss uses F.cross_entropy
2435 # which gives different results in bfloat16.
2436 assert isinstance(
2437 logits, torch.Tensor
2438 ), f"Expected logits tensor, got {type(logits)}"
2439 if labels is not None:
2440 return self._causal_labels_loss(
2441 logits,
2442 labels,
2443 attention_mask=attention_mask,
2444 per_token=loss_per_token,
2445 )
2446 return self.loss_fn(
2447 logits,
2448 input_ids,
2449 attention_mask=attention_mask,
2450 per_token=loss_per_token,
2451 )
2452 elif return_type == "both":
2453 if getattr(self.cfg, "is_audio_model", False): 2453 ↛ 2454line 2453 didn't jump to line 2454 because the condition on line 2453 was never true
2454 raise ValueError(
2455 "Audio models do not support return_type='both'. "
2456 "CTC loss requires aligned frame-level labels."
2457 )
2458 if _is_inputs_embeds and labels is None: 2458 ↛ 2459line 2458 didn't jump to line 2459 because the condition on line 2458 was never true
2459 raise ValueError(
2460 "Cannot compute loss with inputs_embeds — token IDs required for labels."
2461 )
2462 assert isinstance(
2463 logits, torch.Tensor
2464 ), f"Expected logits tensor, got {type(logits)}"
2465 if labels is not None:
2466 loss = self._causal_labels_loss(
2467 logits,
2468 labels,
2469 attention_mask=attention_mask,
2470 per_token=loss_per_token,
2471 )
2472 else:
2473 loss = self.loss_fn(
2474 logits,
2475 input_ids,
2476 attention_mask=attention_mask,
2477 per_token=loss_per_token,
2478 )
2479 return (logits, loss)
2480 elif return_type == "predictions": 2480 ↛ 2481line 2480 didn't jump to line 2481 because the condition on line 2480 was never true
2481 assert (
2482 self.tokenizer is not None
2483 ), "Must have a tokenizer to use return_type='predictions'"
2484 if logits.shape[-1] == 2:
2485 # Next Sentence Prediction — 2-class output
2486 logprobs = logits.log_softmax(dim=-1)
2487 predictions = [
2488 "The sentences are sequential",
2489 "The sentences are NOT sequential",
2490 ]
2491 return predictions[logprobs.argmax(dim=-1).item()]
2492 else:
2493 # Masked Language Modeling — decode [MASK] tokens
2494 logprobs = logits[input_ids == self.tokenizer.mask_token_id].log_softmax(dim=-1)
2495 predictions = self.tokenizer.decode(logprobs.argmax(dim=-1))
2496 if " " in predictions:
2497 predictions = predictions.split(" ")
2498 predictions = [f"Prediction {i}: {p}" for i, p in enumerate(predictions)]
2499 return predictions
2500 elif return_type is None: 2500 ↛ 2503line 2500 didn't jump to line 2503 because the condition on line 2500 was always true
2501 return None
2502 else:
2503 raise ValueError(f"Invalid return_type: {return_type}")
2504 except StopAtLayerException as e:
2505 # Execution stopped at the requested layer
2506 return e.layer_output
2507 finally:
2508 # Clean up state that may be inconsistent after StopAtLayerException
2509 if stop_at_layer is not None:
2510 for bl_name in (
2511 "blocks",
2512 "encoder_blocks",
2513 "decoder_blocks",
2514 "L_blocks",
2515 "H_blocks",
2516 ):
2517 if hasattr(self, bl_name):
2518 for block in getattr(self, bl_name):
2519 block._stop_at_layer_idx = None
2521 # Clear any stale KV cache — layers after the stop point didn't
2522 # execute, so the cache is incomplete and would corrupt subsequent
2523 # generate() calls that expect a full cache.
2524 if hasattr(self, "_last_hf_cache"):
2525 del self._last_hf_cache
2527 def get_hook_point(self, hook_name: str) -> Optional[HookPoint]:
2528 """Get a hook point by name from the bridge's hook system."""
2529 if hook_name in self._hook_registry:
2530 return self._hook_registry[hook_name]
2531 try:
2532 parts = hook_name.split(".")
2533 current = self
2534 for part in parts:
2535 current = getattr(current, part)
2536 if isinstance(current, HookPoint):
2537 return current
2538 except AttributeError:
2539 pass
2540 return None
2542 def loss_fn(
2543 self,
2544 logits: torch.Tensor,
2545 tokens: torch.Tensor,
2546 attention_mask: Optional[torch.Tensor] = None,
2547 per_token: bool = False,
2548 ) -> torch.Tensor:
2549 """Calculate cross-entropy loss.
2551 Uses the same formula as HookedTransformer (log_softmax + gather) to ensure
2552 numerically identical results when logits match.
2554 Args:
2555 logits: Model logits
2556 tokens: Target tokens
2557 attention_mask: Optional attention mask for padding
2558 per_token: Whether to return per-token loss
2560 Returns:
2561 Loss tensor
2562 """
2563 if tokens.device != logits.device: 2563 ↛ 2564line 2563 didn't jump to line 2564 because the condition on line 2563 was never true
2564 tokens = tokens.to(logits.device)
2565 if attention_mask is not None:
2566 if attention_mask.device != logits.device: 2566 ↛ 2567line 2566 didn't jump to line 2567 because the condition on line 2566 was never true
2567 attention_mask = attention_mask.to(logits.device)
2568 attention_mask = self._prepare_loss_attention_mask(attention_mask, tokens)
2569 return lm_cross_entropy_loss(logits, tokens, attention_mask, per_token)
2571 def _causal_labels_loss(
2572 self,
2573 logits: torch.Tensor,
2574 labels: torch.Tensor,
2575 attention_mask: Optional[torch.Tensor] = None,
2576 per_token: bool = False,
2577 ) -> torch.Tensor:
2578 """Compute shifted causal loss against explicit labels, ignoring ``-100``."""
2579 if labels.device != logits.device: 2579 ↛ 2580line 2579 didn't jump to line 2580 because the condition on line 2579 was never true
2580 labels = labels.to(logits.device)
2581 if labels.shape != logits.shape[:-1]: 2581 ↛ 2582line 2581 didn't jump to line 2582 because the condition on line 2581 was never true
2582 raise ValueError(
2583 "causal labels must match the logits batch and position dimensions, "
2584 f"got labels {tuple(labels.shape)} and logits {tuple(logits.shape)}"
2585 )
2587 losses = F.cross_entropy(
2588 logits[:, :-1].flatten(0, 1),
2589 labels[:, 1:].flatten(),
2590 reduction="none",
2591 ignore_index=-100,
2592 ).view_as(labels[:, 1:])
2593 valid_targets = labels[:, 1:] != -100
2594 if attention_mask is not None:
2595 if attention_mask.device != logits.device: 2595 ↛ 2596line 2595 didn't jump to line 2596 because the condition on line 2595 was never true
2596 attention_mask = attention_mask.to(logits.device)
2597 token_mask = self._prepare_loss_attention_mask(attention_mask, labels)
2598 valid_targets &= token_mask[:, :-1] & token_mask[:, 1:]
2599 losses = losses.masked_fill(~valid_targets, 0.0)
2600 return losses if per_token else losses.sum() / valid_targets.sum()
2602 @staticmethod
2603 def _seq2seq_loss(
2604 logits: torch.Tensor,
2605 labels: torch.Tensor,
2606 native_loss: Any,
2607 *,
2608 per_token: bool,
2609 ) -> torch.Tensor:
2610 """Return encoder-decoder loss without the causal LM token shift."""
2611 if labels.device != logits.device: 2611 ↛ 2612line 2611 didn't jump to line 2612 because the condition on line 2611 was never true
2612 labels = labels.to(logits.device)
2613 if labels.shape != logits.shape[:-1]: 2613 ↛ 2614line 2613 didn't jump to line 2614 because the condition on line 2613 was never true
2614 raise ValueError(
2615 "seq2seq labels must match the decoder logits batch and position "
2616 f"dimensions, got labels {tuple(labels.shape)} and logits "
2617 f"{tuple(logits.shape)}"
2618 )
2619 if not per_token and isinstance(native_loss, torch.Tensor):
2620 return native_loss
2622 losses = F.cross_entropy(
2623 logits.flatten(0, 1),
2624 labels.flatten(),
2625 reduction="none" if per_token else "mean",
2626 ignore_index=-100,
2627 )
2628 return losses.view_as(labels) if per_token else losses
2630 def _finalize_seq2seq_return(
2631 self,
2632 return_type: str,
2633 logits: torch.Tensor,
2634 labels: torch.Tensor,
2635 native_output: Any,
2636 *,
2637 loss_per_token: bool,
2638 ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
2639 loss = self._seq2seq_loss(
2640 logits,
2641 labels,
2642 getattr(native_output, "loss", None),
2643 per_token=loss_per_token,
2644 )
2645 return (logits, loss) if return_type == "both" else loss
2647 @staticmethod
2648 def _prepare_loss_attention_mask(
2649 attention_mask: torch.Tensor, tokens: torch.Tensor
2650 ) -> torch.Tensor:
2651 """Reduce a forward attention mask to the token window scored by the loss."""
2652 batch, pos = tokens.shape
2653 if attention_mask.ndim not in (2, 4): 2653 ↛ 2654line 2653 didn't jump to line 2654 because the condition on line 2653 was never true
2654 raise ValueError(
2655 "attention_mask must be 2D [batch, key_pos] or 4D "
2656 f"[batch, *, query_pos, key_pos], got shape {tuple(attention_mask.shape)}"
2657 )
2658 if attention_mask.shape[0] != batch: 2658 ↛ 2659line 2658 didn't jump to line 2659 because the condition on line 2658 was never true
2659 raise ValueError(
2660 "attention_mask batch dimension must match tokens, "
2661 f"got {attention_mask.shape[0]} and {batch}"
2662 )
2664 if attention_mask.ndim == 2:
2665 if attention_mask.shape[1] < pos: 2665 ↛ 2666line 2665 didn't jump to line 2666 because the condition on line 2665 was never true
2666 raise ValueError(
2667 "attention_mask must cover every scored token, "
2668 f"got length {attention_mask.shape[1]} for {pos} tokens"
2669 )
2670 return attention_mask[:, -pos:].bool()
2672 query_pos, key_pos = attention_mask.shape[-2:]
2673 if key_pos < pos: 2673 ↛ 2674line 2673 didn't jump to line 2674 because the condition on line 2673 was never true
2674 raise ValueError(
2675 "attention_mask must cover every scored token, "
2676 f"got key length {key_pos} for {pos} tokens"
2677 )
2679 blocked = attention_mask if attention_mask.dtype is torch.bool else attention_mask < -1.0
2680 if query_pos == 1:
2681 # Broadcast key-only masks use one query row for the full sequence.
2682 keep = ~blocked[..., 0, -pos:]
2683 else:
2684 if query_pos < pos: 2684 ↛ 2685line 2684 didn't jump to line 2685 because the condition on line 2684 was never true
2685 raise ValueError(
2686 "attention_mask must contain a query row for every scored token, "
2687 f"got {query_pos} rows for {pos} tokens"
2688 )
2689 # The aligned diagonal excludes causal masking while retaining padding.
2690 diagonal = torch.diagonal(
2691 blocked,
2692 offset=key_pos - query_pos,
2693 dim1=-2,
2694 dim2=-1,
2695 )
2696 if diagonal.shape[-1] < pos: 2696 ↛ 2697line 2696 didn't jump to line 2697 because the condition on line 2696 was never true
2697 raise ValueError(
2698 "attention_mask diagonal must cover every scored token, "
2699 f"got length {diagonal.shape[-1]} for {pos} tokens"
2700 )
2701 keep = ~diagonal[..., -pos:]
2703 # A token is padding only when every broadcast/head mask blocks its key.
2704 return keep.reshape(batch, -1, pos).any(dim=1)
2706 @overload
2707 def run_with_cache(
2708 self,
2709 input: Union[str, List[str], torch.Tensor],
2710 return_cache_object: Literal[True] = True,
2711 remove_batch_dim: bool = False,
2712 names_filter: Optional[Union[str, List[str], Callable[[str], bool]]] = None,
2713 stop_at_layer: Optional[int] = None,
2714 incl_bwd: bool = False,
2715 pos_slice: Optional[Union[Slice, SliceInput]] = None,
2716 reset_hooks_end: bool = True,
2717 clear_contexts: bool = False,
2718 **kwargs,
2719 ) -> Tuple[Any, ActivationCache]:
2720 """Run with cache - placeholder implementation."""
2721 pass
2723 @overload
2724 def run_with_cache(
2725 self,
2726 input: Union[str, List[str], torch.Tensor],
2727 return_cache_object: Literal[False],
2728 remove_batch_dim: bool = False,
2729 names_filter: Optional[Union[str, List[str], Callable[[str], bool]]] = None,
2730 stop_at_layer: Optional[int] = None,
2731 incl_bwd: bool = False,
2732 pos_slice: Optional[Union[Slice, SliceInput]] = None,
2733 reset_hooks_end: bool = True,
2734 clear_contexts: bool = False,
2735 **kwargs,
2736 ) -> Tuple[Any, Dict[str, torch.Tensor]]:
2737 """Run with cache - placeholder implementation."""
2738 pass
2740 def run_with_cache(
2741 self,
2742 input: Union[str, List[str], torch.Tensor],
2743 return_cache_object: bool = True,
2744 remove_batch_dim: bool = False,
2745 names_filter: Optional[Union[str, List[str], Callable[[str], bool]]] = None,
2746 stop_at_layer: Optional[int] = None,
2747 incl_bwd: bool = False,
2748 pos_slice: Optional[Union[Slice, SliceInput]] = None,
2749 reset_hooks_end: bool = True,
2750 clear_contexts: bool = False,
2751 **kwargs,
2752 ) -> Tuple[Any, Union[ActivationCache, Dict[str, torch.Tensor]]]:
2753 """Run the model and cache all activations.
2755 Args:
2756 input: Input to the model
2757 return_cache_object: Whether to return ActivationCache object
2758 remove_batch_dim: Whether to remove batch dimension
2759 names_filter: Filter for which activations to cache (str, list of str, or callable)
2760 stop_at_layer: Layer to stop forward pass at (uses StopAtLayerException; cleans up KV cache on stop)
2761 incl_bwd: If True, also caches gradients under a ``_grad`` suffix by calling
2762 ``backward()`` on the output, matching HookedRootModule.run_with_cache.
2763 Requires a scalar output, so pass ``return_type="loss"``.
2764 pos_slice: Slice applied to the position axis of every cached tensor
2765 (and of the gradients when ``incl_bwd``). Defaults to None, do nothing.
2766 reset_hooks_end: If True, removes the hooks this call added when it finishes.
2767 Hooks the caller added beforehand are left alone either way.
2768 clear_contexts: If True, clears the hook contexts of the touched hook points
2769 when the hooks are removed.
2770 device: Where to store cached activations (matches ActivationCache.to;
2771 does not move the model). Defaults to per-layer storage.
2772 **kwargs: Additional arguments
2773 # type: ignore[name-defined]
2774 Returns:
2775 Tuple of (output, cache)
2776 """
2777 if incl_bwd and stop_at_layer is not None:
2778 raise ValueError(
2779 "incl_bwd=True cannot be combined with stop_at_layer: the run returns the "
2780 "intermediate activation at that layer, not a scalar to call backward() on."
2781 )
2782 if incl_bwd and not torch.is_grad_enabled():
2783 raise ValueError(
2784 "incl_bwd=True needs autograd, but gradient tracking is off "
2785 "(torch.no_grad(), set_grad_enabled(False) or inference mode). "
2786 "Run the call with gradients enabled."
2787 )
2789 aliases = build_alias_to_canonical_map(self.hook_dict)
2791 def create_names_filter_fn(filter_input):
2792 if filter_input is None:
2793 return lambda name: True
2794 elif isinstance(filter_input, str):
2795 mapped_name = aliases.get(filter_input, None)
2796 if mapped_name:
2797 return lambda name: name == mapped_name or name == filter_input
2798 else:
2799 return lambda name: name == filter_input
2800 elif isinstance(filter_input, list):
2801 mapped_list = []
2802 for item in filter_input:
2803 mapped_list.append(item)
2804 mapped_name = aliases.get(item, None)
2805 if mapped_name:
2806 mapped_list.append(mapped_name)
2807 return lambda name: name in mapped_list
2808 elif callable(filter_input): 2808 ↛ 2811line 2808 didn't jump to line 2811 because the condition on line 2808 was always true
2809 return filter_input
2810 else:
2811 raise ValueError("names_filter must be a string, list of strings, or callable")
2813 names_filter_fn = create_names_filter_fn(names_filter)
2814 cache: Dict[str, torch.Tensor] = {}
2815 hooks: List[Tuple[HookPoint, str]] = []
2817 # None → no-op .to(None), tensors stay on their current device.
2818 cache_device = kwargs.pop("device", None)
2819 resolved_pos_slice = Slice.unwrap(pos_slice)
2821 def make_cache_hook(name: str, pos_dim: int = -2, is_backward: bool = False):
2822 key = f"{name}_grad" if is_backward else name
2824 def store(tensor: torch.Tensor) -> None:
2825 tensor = tensor.detach().to(cache_device)
2826 if pos_slice is not None and isinstance(tensor, torch.Tensor) and tensor.dim() >= 2:
2827 # Tensors too flat for the hook's layout (2D token ids at embed.hook_in)
2828 # still carry position on axis 1, right after batch. Deliberate divergence
2829 # from HookedRootModule, which would slice such a tensor on batch.
2830 axis = pos_dim if tensor.dim() >= 1 - pos_dim else 1
2831 tensor = resolved_pos_slice.apply(tensor, dim=axis)
2832 cache[key] = tensor
2834 def cache_hook(tensor: torch.Tensor, *, hook: Any) -> Optional[torch.Tensor]:
2835 if tensor is None: 2835 ↛ 2836line 2835 didn't jump to line 2836 because the condition on line 2835 was never true
2836 cache[key] = None
2837 elif isinstance(tensor, torch.Tensor): 2837 ↛ 2839line 2837 didn't jump to line 2839 because the condition on line 2837 was always true
2838 store(tensor)
2839 elif isinstance(tensor, tuple):
2840 if len(tensor) > 0 and isinstance(tensor[0], torch.Tensor):
2841 store(tensor[0])
2842 else:
2843 pass
2844 else:
2845 try:
2846 if hasattr(tensor, "detach"):
2847 store(tensor)
2848 except Exception:
2849 pass
2850 # A non-None return from a backward hook replaces grad_input; stay read-only.
2851 return None if is_backward else tensor
2853 return cache_hook
2855 hook_dict = self.hook_dict
2856 effective_stop_layer = None
2857 if stop_at_layer is not None and hasattr(self, "blocks"):
2858 if stop_at_layer < 0:
2859 effective_stop_layer = len(self.blocks) + stop_at_layer
2860 else:
2861 effective_stop_layer = stop_at_layer
2862 gated_names_skipped: List[str] = []
2863 for hook_name, hook in hook_dict.items():
2864 if names_filter_fn(hook_name):
2865 if effective_stop_layer is not None:
2866 if hook_name.startswith("blocks."):
2867 try:
2868 layer_num = int(hook_name.split(".")[1])
2869 if layer_num >= effective_stop_layer:
2870 continue
2871 except (IndexError, ValueError):
2872 pass
2874 # Only validate gated hooks when the caller explicitly supplied
2875 # a names_filter. The default filter matches every hook and must
2876 # not cause gated hooks to be treated as explicitly requested.
2877 if names_filter is not None:
2878 try:
2879 self.check_hooks_to_add(hook_name)
2880 except ValueError:
2881 gated_names_skipped.append(hook_name)
2882 continue
2884 hooks.append((hook, hook_name))
2886 if names_filter is not None and gated_names_skipped:
2887 warnings.warn(
2888 f"run_with_cache: skipped {len(gated_names_skipped)} gated-off hook name(s) "
2889 f"that will never be cached: {gated_names_skipped}. Call the relevant "
2890 "set_use_*(True) setter first to enable them.",
2891 stacklevel=2,
2892 )
2893 self.context_level += 1
2894 context_level = self.context_level
2895 try:
2896 for hp, name in hooks:
2897 pos_dim = _pos_axis_for_hook(name, hp)
2898 hp.add_hook(make_cache_hook(name, pos_dim), level=context_level)
2899 if incl_bwd:
2900 hp.add_hook(
2901 make_cache_hook(name, pos_dim, is_backward=True),
2902 dir="bwd",
2903 level=context_level,
2904 )
2905 processed_args = [input]
2906 if processed_args and isinstance(processed_args[0], str):
2907 assert self.tokenizer is not None, "Tokenizer must be set to pass string input."
2908 prepend_bos = kwargs.pop("prepend_bos", None)
2909 input_ids = self.to_tokens(processed_args[0], prepend_bos=prepend_bos)
2910 input_ids = input_ids.to(next(self.original_model.parameters()).device)
2911 kwargs["input_ids"] = input_ids
2912 processed_args = processed_args[1:]
2913 elif "input" in kwargs and isinstance(kwargs["input"], str): 2913 ↛ 2914line 2913 didn't jump to line 2914 because the condition on line 2913 was never true
2914 assert self.tokenizer is not None, "Tokenizer must be set to pass string input."
2915 prepend_bos = kwargs.pop("prepend_bos", None)
2916 input_ids = self.to_tokens(kwargs["input"], prepend_bos=prepend_bos)
2917 input_ids = input_ids.to(next(self.original_model.parameters()).device)
2918 kwargs["input_ids"] = input_ids
2919 del kwargs["input"]
2920 if stop_at_layer is not None and hasattr(self, "blocks"):
2921 if stop_at_layer < 0:
2922 stop_at_layer = len(self.blocks) + stop_at_layer
2923 last_layer_to_process = stop_at_layer - 1
2925 def stop_hook(tensor: torch.Tensor, *, hook: Any) -> torch.Tensor:
2926 raise StopAtLayerException(tensor)
2928 if stop_at_layer >= 0 and stop_at_layer < len(self.blocks): 2928 ↛ 2935line 2928 didn't jump to line 2935 because the condition on line 2928 was always true
2929 # Stop at the beginning of the specified block, not at the end of the previous block
2930 block_hook_name = f"blocks.{stop_at_layer}.hook_in"
2931 hook_dict = self.hook_dict
2932 if block_hook_name in hook_dict: 2932 ↛ 2935line 2932 didn't jump to line 2935 because the condition on line 2932 was always true
2933 hook_dict[block_hook_name].add_hook(stop_hook, level=context_level)
2934 hooks.append((hook_dict[block_hook_name], block_hook_name))
2935 filtered_kwargs = kwargs.copy()
2936 # `cache_device` is honored by `make_cache_hook` above (`tensor.detach().to(cache_device)`);
2937 # the model and inputs stay where the caller put them, matching `ActivationCache.to`.
2938 if cache_device is not None and getattr(self.cfg, "n_devices", 1) > 1:
2939 # Moving a dispatched model to a single device collapses accelerate's
2940 # split and breaks its routing hooks. The cache will stay spread across
2941 # the per-layer devices; callers can .to(cache_device) on cache entries
2942 # after the fact if they need a single-device cache.
2943 warnings.warn(
2944 f"run_with_cache(device={cache_device!r}) ignored: model is dispatched "
2945 f"across {self.cfg.n_devices} devices via device_map. Cached activations "
2946 "will remain on their per-layer devices.",
2947 stacklevel=2,
2948 )
2949 if (
2950 "output_attentions" not in filtered_kwargs
2951 and self.adapter.supports_hf_output_attentions
2952 ):
2953 # Attention-free remote models (e.g. HyenaDNA) reject the kwarg
2954 # outright; only pass it when the HF forward actually accepts it.
2955 fwd_params = inspect.signature(self.original_model.forward).parameters
2956 if "output_attentions" in fwd_params or any( 2956 ↛ 2960line 2956 didn't jump to line 2960 because the condition on line 2956 was always true
2957 p.kind is inspect.Parameter.VAR_KEYWORD for p in fwd_params.values()
2958 ):
2959 filtered_kwargs["output_attentions"] = True
2960 if processed_args:
2961 output = self.forward(processed_args[0], **filtered_kwargs)
2962 elif "input_ids" in filtered_kwargs: 2962 ↛ 2968line 2962 didn't jump to line 2968 because the condition on line 2962 was always true
2963 output = self.forward(
2964 filtered_kwargs["input_ids"],
2965 **{k: v for k, v in filtered_kwargs.items() if k != "input_ids"},
2966 )
2967 else:
2968 output = self.forward(**filtered_kwargs)
2969 if hasattr(output, "logits"): 2969 ↛ 2970line 2969 didn't jump to line 2970 because the condition on line 2969 was never true
2970 output = output.logits
2971 if incl_bwd:
2972 # Gradients land in the cache via the bwd hooks, which the finally below removes,
2973 # so the backward pass has to happen inside this try.
2974 if not isinstance(output, torch.Tensor) or output.numel() != 1:
2975 shape = tuple(output.shape) if isinstance(output, torch.Tensor) else None
2976 raise ValueError(
2977 "incl_bwd=True needs a scalar output to call backward() on, got "
2978 f"{type(output).__name__}{f' of shape {shape}' if shape else ''}. "
2979 'Pass return_type="loss".'
2980 )
2981 if not output.requires_grad:
2982 raise ValueError(
2983 "incl_bwd=True got an output with no grad_fn — the model's parameters "
2984 "have requires_grad=False, so there is nothing to differentiate."
2985 )
2986 output.backward()
2987 except StopAtLayerException as e:
2988 output = e.layer_output
2989 except Exception as e:
2990 raise e
2991 finally:
2992 if reset_hooks_end:
2993 for hp, _ in hooks:
2994 # `level` keeps this to the hooks added above — hooks the caller
2995 # attached before the call survive.
2996 hp.remove_hooks(dir="fwd", level=context_level)
2997 if incl_bwd:
2998 hp.remove_hooks(dir="bwd", level=context_level)
2999 if clear_contexts:
3000 hp.clear_context()
3001 self.context_level -= 1
3002 if self.compatibility_mode == True:
3003 reverse_aliases = {}
3004 for old_name, new_name in aliases.items():
3005 if isinstance(new_name, list): 3005 ↛ 3006line 3005 didn't jump to line 3006 because the condition on line 3005 was never true
3006 for single_new_name in new_name:
3007 reverse_aliases[single_new_name] = old_name
3008 else:
3009 reverse_aliases[new_name] = old_name
3010 # Gradient entries are keyed "<hook_name>_grad", so alias lookups run on the
3011 # base name and the suffix is re-attached to the aliased key.
3012 suffixes = ("", "_grad") if incl_bwd else ("",)
3013 cache_items_to_add = {}
3014 for cache_name, cached_value in cache.items():
3015 base_name, suffix = (
3016 (cache_name[: -len("_grad")], "_grad")
3017 if cache_name.endswith("_grad")
3018 else (cache_name, "")
3019 )
3020 old_name = reverse_aliases.get(base_name)
3021 if old_name is not None:
3022 cache_items_to_add[old_name + suffix] = cached_value
3023 cache.update(cache_items_to_add)
3024 for alias_name, target_name in aliases.items():
3025 targets = target_name if isinstance(target_name, list) else [target_name]
3026 for suffix in suffixes:
3027 if alias_name + suffix in cache:
3028 continue
3029 for single_target in targets:
3030 if single_target + suffix in cache: 3030 ↛ 3031line 3030 didn't jump to line 3031 because the condition on line 3030 was never true
3031 cache[alias_name + suffix] = cache[single_target + suffix]
3032 break
3033 if return_cache_object:
3034 activation_cache = ActivationCache(cache, self, has_batch_dim=True)
3035 if remove_batch_dim: 3035 ↛ 3036line 3035 didn't jump to line 3036 because the condition on line 3035 was never true
3036 activation_cache.remove_batch_dim()
3037 return (output, activation_cache)
3038 else:
3039 if remove_batch_dim:
3040 for key in cache:
3041 if cache[key] is not None and isinstance(cache[key], torch.Tensor): 3041 ↛ 3040line 3041 didn't jump to line 3040 because the condition on line 3041 was always true
3042 if cache[key].size(0) == 1: 3042 ↛ 3040line 3042 didn't jump to line 3040 because the condition on line 3042 was always true
3043 cache[key] = cache[key][0]
3044 return (output, cache)
3046 def run_with_hooks(
3047 self,
3048 input: Union[str, List[str], torch.Tensor],
3049 fwd_hooks: List[Tuple[Union[str, Callable], Callable]] = [],
3050 bwd_hooks: List[Tuple[Union[str, Callable], Callable]] = [],
3051 reset_hooks_end: bool = True,
3052 clear_contexts: bool = False,
3053 return_type: Optional[str] = "logits",
3054 names_filter: Optional[Union[str, List[str], Callable[[str], bool]]] = None,
3055 stop_at_layer: Optional[int] = None,
3056 remove_batch_dim: bool = False,
3057 **kwargs,
3058 ) -> Any:
3059 """Run the model with specified forward and backward hooks.
3061 Args:
3062 input: Input to the model
3063 fwd_hooks: Forward hooks to apply
3064 bwd_hooks: Backward hooks to apply
3065 reset_hooks_end: Whether to reset hooks at the end
3066 clear_contexts: Whether to clear hook contexts
3067 return_type: What to return ("logits", "loss", etc.)
3068 names_filter: Filter for hook names (not used directly, for compatibility)
3069 stop_at_layer: Layer to stop at (uses StopAtLayerException; cleans up KV cache on stop)
3070 remove_batch_dim: Whether to remove batch dimension from hook inputs (only works for batch_size==1)
3071 **kwargs: Additional arguments
3073 Returns:
3074 Model output
3075 """
3076 added_hooks: List[Tuple[HookPoint, Literal["fwd", "bwd"]]] = []
3077 # Claimed here for the closures below, but only committed to self inside the try that
3078 # decrements it, so a raise in between can't leave the counter incremented.
3079 context_level = self.context_level + 1
3080 effective_stop_layer = None
3081 if stop_at_layer is not None and hasattr(self, "blocks"):
3082 if stop_at_layer < 0: 3082 ↛ 3083line 3082 didn't jump to line 3083 because the condition on line 3082 was never true
3083 effective_stop_layer = len(self.blocks) + stop_at_layer
3084 else:
3085 effective_stop_layer = stop_at_layer
3087 def add_hook_to_point(
3088 hook_point: HookPoint,
3089 hook_fn: Callable,
3090 name: str,
3091 dir: Literal["fwd", "bwd"] = "fwd",
3092 *,
3093 is_explicit: bool = True,
3094 ):
3095 if effective_stop_layer is not None and name.startswith("blocks."):
3096 try:
3097 layer_num = int(name.split(".")[1])
3098 if layer_num >= effective_stop_layer:
3099 return
3100 except (IndexError, ValueError):
3101 pass
3102 if is_explicit:
3103 self.check_hooks_to_add(name)
3104 elif self._gated_hook_reason(name) is not None: 3104 ↛ 3105line 3104 didn't jump to line 3105 because the condition on line 3104 was never true
3105 warnings.warn(
3106 f"run_with_hooks(): filter matched gated-off hook name '{name}', skipped. "
3107 "Call the relevant set_use_*(True) setter first to enable it.",
3108 stacklevel=2,
3109 )
3110 return
3111 if self.compatibility_mode and name != hook_point.name: 3111 ↛ 3112line 3111 didn't jump to line 3112 because the condition on line 3111 was never true
3112 alias_names_list: list[str] = []
3113 if hook_point.name is not None:
3114 alias_names_list.append(hook_point.name)
3115 alias_names_list.append(name)
3116 hook_point.add_hook(
3117 hook_fn, dir=dir, alias_names=alias_names_list, level=context_level
3118 )
3119 else:
3120 hook_point.add_hook(hook_fn, dir=dir, level=context_level)
3121 added_hooks.append((hook_point, dir))
3123 def apply_hooks(hooks: List[Tuple[Union[str, Callable], Callable]], is_fwd: bool):
3124 direction: Literal["fwd", "bwd"] = "fwd" if is_fwd else "bwd"
3125 aliases = build_alias_to_canonical_map(self.hook_dict)
3126 for hook_name_or_filter, hook_fn in hooks:
3127 if remove_batch_dim: 3127 ↛ 3128line 3127 didn't jump to line 3128 because the condition on line 3127 was never true
3128 original_hook_fn = hook_fn
3130 # Default arg captures hook_fn by value (avoids closure issue)
3131 def wrapped_hook_fn(tensor, hook, _orig_fn=original_hook_fn):
3132 if tensor.shape[0] == 1:
3133 tensor_no_batch = tensor.squeeze(0)
3134 result = _orig_fn(tensor_no_batch, hook)
3135 if result.dim() == tensor_no_batch.dim():
3136 result = result.unsqueeze(0)
3137 return result
3138 else:
3139 return _orig_fn(tensor, hook)
3141 hook_fn = wrapped_hook_fn
3142 if isinstance(hook_name_or_filter, str):
3143 hook_dict = self.hook_dict
3144 actual_hook_name = hook_name_or_filter
3145 if hook_name_or_filter in aliases:
3146 actual_hook_name = aliases[hook_name_or_filter]
3147 if actual_hook_name in hook_dict: 3147 ↛ 3126line 3147 didn't jump to line 3126 because the condition on line 3147 was always true
3148 add_hook_to_point(
3149 hook_dict[actual_hook_name],
3150 hook_fn,
3151 actual_hook_name,
3152 direction,
3153 is_explicit=True,
3154 )
3155 else:
3156 hook_dict = self.hook_dict
3157 seen_hooks = set()
3158 for name, hook_point in hook_dict.items():
3159 if hook_name_or_filter(name):
3160 hook_id = id(hook_point)
3161 if hook_id in seen_hooks: 3161 ↛ 3162line 3161 didn't jump to line 3162 because the condition on line 3161 was never true
3162 continue
3163 seen_hooks.add(hook_id)
3164 hook_name_to_use = hook_point.name if hook_point.name else name
3165 add_hook_to_point(
3166 hook_point,
3167 hook_fn,
3168 hook_name_to_use,
3169 direction,
3170 is_explicit=False,
3171 )
3173 try:
3174 self.context_level = context_level
3175 if stop_at_layer is not None and hasattr(self, "blocks"):
3176 if stop_at_layer < 0: 3176 ↛ 3177line 3176 didn't jump to line 3177 because the condition on line 3176 was never true
3177 stop_at_layer = len(self.blocks) + stop_at_layer
3178 last_layer_to_process = stop_at_layer - 1
3180 def stop_hook(tensor: torch.Tensor, *, hook: Any) -> torch.Tensor:
3181 raise StopAtLayerException(tensor)
3183 if stop_at_layer >= 0 and stop_at_layer < len(self.blocks): 3183 ↛ 3192line 3183 didn't jump to line 3192 because the condition on line 3183 was always true
3184 # Stop at the beginning of the specified block, not at the end of the previous block
3185 block_hook_name = f"blocks.{stop_at_layer}.hook_in"
3186 hook_dict = self.hook_dict
3187 if block_hook_name in hook_dict: 3187 ↛ 3192line 3187 didn't jump to line 3192 because the condition on line 3187 was always true
3188 add_hook_to_point(
3189 hook_dict[block_hook_name], stop_hook, block_hook_name, "fwd"
3190 )
3192 apply_hooks(fwd_hooks, True)
3193 apply_hooks(bwd_hooks, False)
3194 try:
3195 output = self.forward(
3196 input, return_type=return_type, stop_at_layer=stop_at_layer, **kwargs
3197 )
3198 except StopAtLayerException as e:
3199 output = e.layer_output
3200 return output
3201 finally:
3202 if reset_hooks_end: 3202 ↛ 3209line 3202 didn't jump to line 3209 because the condition on line 3202 was always true
3203 for hook_point, direction in added_hooks:
3204 # `level` keeps this to the hooks added above — hooks the caller
3205 # attached before the call survive.
3206 hook_point.remove_hooks(dir=direction, level=context_level)
3207 if clear_contexts:
3208 hook_point.clear_context()
3209 self.context_level -= 1
3211 def _resolve_stopping_criteria(
3212 self,
3213 stop_strings: Optional[Union[str, List[str]]],
3214 stopping_criteria: Optional[Any],
3215 ) -> Optional[Any]:
3216 """Combine ``stop_strings`` and ``stopping_criteria`` into one StoppingCriteriaList.
3218 Returns ``None`` when neither is supplied (or both reduce to no-ops),
3219 so callers can cheaply check whether any extra stop signal is active.
3220 ``stop_strings`` is turned into a HuggingFace ``StopStringCriteria`` (which reproduces
3221 HF's exact partial-token-aware, end-anchored matching: it fires when the stop string
3222 ends the generated text, even if the string straddles token boundaries) and therefore
3223 requires a tokenizer.
3224 A user-supplied ``stopping_criteria`` may be a single ``StoppingCriteria``,
3225 a list of them, or a ``StoppingCriteriaList``.
3227 Raises:
3228 ValueError: if ``stop_strings`` is supplied without a tokenizer.
3229 TypeError: if ``stopping_criteria`` is not a ``StoppingCriteria``, a
3230 list/tuple of them, or a ``StoppingCriteriaList``.
3231 """
3232 if stop_strings is None and stopping_criteria is None:
3233 return None
3235 from transformers import ( # local import: matches the file's transformers usage
3236 StoppingCriteria,
3237 StoppingCriteriaList,
3238 StopStringCriteria,
3239 )
3241 criteria = StoppingCriteriaList()
3243 if stop_strings is not None:
3244 strings = [stop_strings] if isinstance(stop_strings, str) else list(stop_strings)
3245 strings = [s for s in strings if s] # drop empty strings (HF errors on them)
3246 if strings:
3247 if self.tokenizer is None:
3248 raise ValueError(
3249 "stop_strings requires a tokenizer (stop strings are detected by "
3250 "matching against the tokenizer vocabulary), but this TransformerBridge "
3251 "has no tokenizer. Pass a stopping_criteria callable that operates on "
3252 "token ids instead, or use hf_generate()."
3253 )
3254 criteria.append(StopStringCriteria(tokenizer=self.tokenizer, stop_strings=strings))
3256 if stopping_criteria is not None:
3257 if isinstance(stopping_criteria, StoppingCriteriaList):
3258 criteria.extend(stopping_criteria)
3259 elif isinstance(stopping_criteria, (list, tuple)):
3260 criteria.extend(stopping_criteria)
3261 elif isinstance(stopping_criteria, StoppingCriteria):
3262 criteria.append(stopping_criteria)
3263 else:
3264 raise TypeError(
3265 "stopping_criteria must be a transformers.StoppingCriteria, a list of "
3266 f"them, or a StoppingCriteriaList, but got {type(stopping_criteria).__name__}."
3267 )
3269 return criteria if len(criteria) > 0 else None
3271 def _encdec_ngram_processor(self) -> Optional[Any]:
3272 """generation_config.no_repeat_ngram_size as transformers' own
3273 processor, or None. HF applies it by default; parity for models whose
3274 greedy decode needs it to escape token attractors."""
3275 size = getattr(
3276 getattr(self.original_model, "generation_config", None),
3277 "no_repeat_ngram_size",
3278 None,
3279 )
3280 if not size:
3281 return None
3282 from transformers.generation.logits_process import NoRepeatNGramLogitsProcessor
3284 return NoRepeatNGramLogitsProcessor(size)
3286 def _generate_tokens(
3287 self,
3288 current_tokens: torch.Tensor,
3289 input_tokens: torch.Tensor,
3290 batch_size: int,
3291 *,
3292 max_new_tokens: int,
3293 do_sample: bool,
3294 top_k: Optional[int],
3295 top_p: Optional[float],
3296 temperature: float,
3297 freq_penalty: float,
3298 repetition_penalty: float,
3299 stop_at_eos: bool,
3300 stop_tokens: List[int],
3301 eos_token_for_padding: int,
3302 finished_sequences: torch.Tensor,
3303 use_past_kv_cache: bool,
3304 use_stateful_cache: bool,
3305 mamba_cache: Any,
3306 mamba_conv_kernel: int,
3307 is_encoder_decoder: bool,
3308 _is_batched_list: bool,
3309 _generate_from_embeds: bool,
3310 encoder_input: Optional[torch.Tensor],
3311 decoder_tokens: Optional[torch.Tensor],
3312 generated_token_ids: Optional[List[torch.Tensor]],
3313 pixel_values: Optional[torch.Tensor],
3314 multimodal_kwargs: Dict[str, Any],
3315 verbose: bool,
3316 stopping_criteria_list: Optional[Any] = None,
3317 initial_attention_mask: Optional[torch.Tensor] = None,
3318 min_decoder_length: Optional[int] = None,
3319 ngram_processor: Optional[Any] = None,
3320 encoder_attention_mask: Optional[torch.Tensor] = None,
3321 ) -> Generator[Tuple[torch.Tensor, torch.Tensor, bool], None, None]:
3322 """Core generation loop. Yields (sampled_tokens, final_logits, all_finished) per step.
3324 Owns the forward pass, sampling, stop handling (EOS and any
3325 ``stopping_criteria_list``), token accumulation, and KV cache management. Callers
3326 are responsible for try/finally cleanup of ``_capture_hf_cache``.
3328 ``stopping_criteria_list`` (from ``_resolve_stopping_criteria``) is evaluated on
3329 the running sequence each step and folded into the finished-sequence mask alongside
3330 EOS, so when it is ``None`` the loop runs the EOS-only path unchanged.
3331 """
3332 _hf_kv_cache = None
3333 # A row may finish via EOS and/or any of the configured stopping criteria.
3334 any_stop_active = stop_at_eos or stopping_criteria_list is not None
3336 # Models that own their position derivation (the gate refuses them) cache
3337 # mRoPE deltas on the module between calls; a text-only prefill never
3338 # refreshes them, so a stale delta from an earlier multimodal forward gets
3339 # added to every cached-step position. HF's generate recomputes them at
3340 # prefill via prepare_inputs_for_generation, which this loop bypasses —
3341 # so match it by clearing before the prompt pass. A multimodal prefill
3342 # recomputes its own fresh deltas regardless.
3343 if not self._accepts_derived_position_ids():
3344 underlying = getattr(self, "original_model", None)
3345 for module in (
3346 underlying,
3347 getattr(underlying, "model", None),
3348 getattr(underlying, "language_model", None),
3349 ):
3350 if module is not None and hasattr(module, "rope_deltas"):
3351 module.rope_deltas = None
3353 # Pure-SSM models (Mamba-1/2) take the stateful cache as `cache_params`;
3354 # modern hybrids (Bamba, NemotronH, FalconH1) take `past_key_values` and
3355 # would receive a duplicate cache_params via **kwargs cascade otherwise.
3356 stateful_cache_kwarg = "cache_params"
3357 if use_stateful_cache:
3358 forward_params = inspect.signature(self.original_model.forward).parameters
3359 if "cache_params" not in forward_params:
3360 stateful_cache_kwarg = "past_key_values"
3362 for gen_step_idx in tqdm.tqdm(range(max_new_tokens), disable=not verbose):
3363 with torch.no_grad():
3364 if is_encoder_decoder:
3365 assert encoder_input is not None
3366 encdec_kwargs: Dict[str, Any] = {}
3367 if encoder_attention_mask is not None:
3368 encdec_kwargs["attention_mask"] = encoder_attention_mask.to(
3369 encoder_input.device
3370 )
3371 logits = self(
3372 encoder_input,
3373 return_type="logits",
3374 decoder_input=decoder_tokens,
3375 **encdec_kwargs,
3376 )
3377 else:
3378 forward_kwargs: Dict[str, Any] = {}
3379 # A prompt mask covers only the prompt, so extend it by one
3380 # attended column per token generated so far. position_ids are
3381 # left to forward(), which derives them from the mask for the
3382 # models that can take them.
3383 running_attention_mask: Optional[torch.Tensor] = None
3384 if initial_attention_mask is not None:
3385 n_generated = current_tokens.shape[1] - initial_attention_mask.shape[1]
3386 running_attention_mask = torch.cat(
3387 [
3388 initial_attention_mask.to(current_tokens.device),
3389 torch.ones(
3390 (current_tokens.shape[0], n_generated),
3391 dtype=initial_attention_mask.dtype,
3392 device=current_tokens.device,
3393 ),
3394 ],
3395 dim=1,
3396 )
3397 forward_kwargs["attention_mask"] = running_attention_mask
3398 # Compute attention mask and position_ids for batched
3399 # inputs with padding.
3400 if (
3401 initial_attention_mask is None
3402 and _is_batched_list
3403 and self.tokenizer is not None
3404 and self.tokenizer.pad_token_id is not None
3405 ):
3406 _prev_side = self.tokenizer.padding_side
3407 self.tokenizer.padding_side = "left"
3408 attn_mask = utils.get_attention_mask(
3409 self.tokenizer,
3410 current_tokens,
3411 prepend_bos=getattr(self.cfg, "default_prepend_bos", True),
3412 ).to(self.cfg.device)
3413 self.tokenizer.padding_side = _prev_side
3414 forward_kwargs["attention_mask"] = attn_mask
3415 # Same target gate as the forward() path: the mask is safe
3416 # for every model, the derived positions are not (#1626).
3417 if self._accepts_derived_position_ids():
3418 position_ids = attn_mask.long().cumsum(-1) - 1
3419 position_ids.masked_fill_(attn_mask == 0, 1)
3420 forward_kwargs["position_ids"] = position_ids
3421 if gen_step_idx == 0:
3422 if pixel_values is not None:
3423 forward_kwargs["pixel_values"] = pixel_values
3424 if multimodal_kwargs:
3425 forward_kwargs.update(multimodal_kwargs)
3426 if use_stateful_cache:
3427 forward_kwargs[stateful_cache_kwarg] = mamba_cache
3428 forward_kwargs["use_cache"] = True
3429 if gen_step_idx == 0:
3430 # Mamba's conv-window warmup positions vs standard
3431 # full-prompt positions for past_key_values hybrids.
3432 prefill_len = (
3433 mamba_conv_kernel
3434 if stateful_cache_kwarg == "cache_params"
3435 else current_tokens.shape[1]
3436 )
3437 cache_position = torch.arange(0, prefill_len, device=self.cfg.device)
3438 forward_kwargs["cache_position"] = cache_position
3439 logits = self(
3440 current_tokens,
3441 return_type="logits",
3442 **forward_kwargs,
3443 )
3444 else:
3445 input_seq_pos = input_tokens.shape[1] + gen_step_idx - 1
3446 cache_position = torch.tensor([input_seq_pos], device=self.cfg.device)
3447 forward_kwargs["cache_position"] = cache_position
3448 if "position_ids" in forward_kwargs: 3448 ↛ 3449line 3448 didn't jump to line 3449 because the condition on line 3448 was never true
3449 forward_kwargs["position_ids"] = forward_kwargs["position_ids"][
3450 :, -1:
3451 ]
3452 logits = self(
3453 current_tokens[:, -1:],
3454 return_type="logits",
3455 **forward_kwargs,
3456 )
3457 elif use_past_kv_cache:
3458 forward_kwargs["use_cache"] = True
3459 if _hf_kv_cache is not None:
3460 forward_kwargs["past_key_values"] = _hf_kv_cache
3461 # HF v5 + macOS-arm64 NaNs when these are inferred
3462 # from cache state alone. Mirror HF generate(): pass
3463 # both an (batch, total_len) attention_mask and a
3464 # (batch, 1) position_ids for the new token.
3465 batch_size = current_tokens.shape[0]
3466 total_len = current_tokens.shape[1]
3467 device = current_tokens.device
3468 if "attention_mask" not in forward_kwargs:
3469 forward_kwargs["attention_mask"] = torch.ones(
3470 (batch_size, total_len),
3471 dtype=torch.long,
3472 device=device,
3473 )
3474 # Gated as a whole (#1626): every branch below supplies
3475 # position_ids, so gating only the prompt derivation
3476 # above would divert a refused model into the
3477 # total_len - 1 fallback, which counts pad slots and is
3478 # wrong per row for a left-padded batch. A model that
3479 # owns its position derivation gets the mask alone,
3480 # matching the uncached path.
3481 if self._accepts_derived_position_ids():
3482 if "position_ids" in forward_kwargs:
3483 forward_kwargs["position_ids"] = forward_kwargs["position_ids"][
3484 :, -1:
3485 ]
3486 elif running_attention_mask is not None:
3487 # total_len - 1 counts pad slots, so it is wrong
3488 # for a left-padded prompt. Derive the new token's
3489 # position from the mask instead.
3490 forward_kwargs["position_ids"] = utils.get_offset_position_ids(
3491 0, running_attention_mask.long()
3492 )[:, -1:]
3493 else:
3494 forward_kwargs["position_ids"] = torch.full(
3495 (batch_size, 1),
3496 total_len - 1,
3497 dtype=torch.long,
3498 device=device,
3499 )
3500 logits = self(
3501 current_tokens[:, -1:],
3502 return_type="logits",
3503 **forward_kwargs,
3504 )
3505 else:
3506 logits = self(
3507 current_tokens,
3508 return_type="logits",
3509 **forward_kwargs,
3510 )
3511 else:
3512 logits = self(current_tokens, return_type="logits", **forward_kwargs)
3513 if use_past_kv_cache and hasattr(self, "_last_hf_cache"):
3514 _hf_kv_cache = self._last_hf_cache or _hf_kv_cache
3515 del self._last_hf_cache
3516 final_logits = logits[:, -1, :]
3518 # Sample next token
3519 penalty_tokens = (
3520 torch.stack(generated_token_ids, dim=1)
3521 if _generate_from_embeds and generated_token_ids
3522 else None
3523 )
3524 # transformers' own NoRepeatNGramLogitsProcessor, honoring
3525 # generation_config (bart-large-cnn pins 3; without it greedy
3526 # decoding falls into a BOS attractor and emits nothing).
3527 if ngram_processor is not None and decoder_tokens is not None:
3528 final_logits = ngram_processor(decoder_tokens, final_logits)
3529 # HF's generate() suppresses EOS below generation_config.min_length
3530 # (bart-large-cnn pins 56); without this the loop can EOS on step
3531 # one and emit an empty summary.
3532 if (
3533 min_decoder_length is not None
3534 and is_encoder_decoder
3535 and decoder_tokens is not None
3536 and decoder_tokens.shape[1] < min_decoder_length
3537 and stop_tokens
3538 ):
3539 final_logits = final_logits.clone()
3540 final_logits[:, stop_tokens] = float("-inf")
3541 if do_sample:
3542 sampled_tokens = utils.sample_logits(
3543 final_logits,
3544 top_k=top_k,
3545 top_p=top_p,
3546 temperature=temperature,
3547 freq_penalty=freq_penalty,
3548 repetition_penalty=repetition_penalty,
3549 tokens=(
3550 penalty_tokens
3551 if _generate_from_embeds
3552 else (decoder_tokens if is_encoder_decoder else current_tokens)
3553 ),
3554 ).to(self.cfg.device)
3555 else:
3556 sampled_tokens = utils.sample_logits(
3557 final_logits,
3558 temperature=0.0,
3559 repetition_penalty=repetition_penalty,
3560 tokens=(
3561 penalty_tokens
3562 if _generate_from_embeds
3563 else (decoder_tokens if is_encoder_decoder else current_tokens)
3564 ),
3565 ).to(self.cfg.device)
3567 # Freeze rows that finished on an earlier step so they stop emitting
3568 # real tokens. Applies to every active stop mechanism, not just EOS.
3569 if any_stop_active:
3570 sampled_tokens[finished_sequences] = eos_token_for_padding
3572 # Fold this step's EOS matches into the finished mask.
3573 if stop_at_eos:
3574 finished_sequences.logical_or_(
3575 torch.isin(
3576 sampled_tokens.to(self.cfg.device),
3577 torch.tensor(stop_tokens).to(self.cfg.device),
3578 )
3579 )
3581 # Update token sequences
3582 if is_encoder_decoder:
3583 assert decoder_tokens is not None
3584 decoder_tokens = torch.cat([decoder_tokens, sampled_tokens.unsqueeze(1)], dim=1)
3585 elif _generate_from_embeds:
3586 assert generated_token_ids is not None
3587 generated_token_ids.append(sampled_tokens)
3588 embed_fn = self.original_model.get_input_embeddings() # type: ignore[operator]
3589 assert embed_fn is not None
3590 new_embed = embed_fn(sampled_tokens.unsqueeze(1)).to(current_tokens.dtype)
3591 current_tokens = torch.cat([current_tokens, new_embed], dim=1)
3592 else:
3593 current_tokens = torch.cat([current_tokens, sampled_tokens.unsqueeze(1)], dim=1)
3595 # Fold stop_strings / stopping_criteria into the finished mask. They are
3596 # evaluated on the full running sequence (prompt + everything generated so
3597 # far, including the token just appended) with this step's logits as the
3598 # scores argument, matching transformers' StoppingCriteria contract. The
3599 # combined list returns a per-row bool [batch] OR-ing every criterion.
3600 # generate()/generate_stream() guarantee this is plain decoder-only token
3601 # generation, so current_tokens is the running token sequence.
3602 if stopping_criteria_list is not None:
3603 criteria_finished = stopping_criteria_list(current_tokens, final_logits).to(
3604 device=self.cfg.device, dtype=torch.bool
3605 )
3606 if criteria_finished.shape != finished_sequences.shape:
3607 raise ValueError(
3608 "A stopping criterion returned shape "
3609 f"{tuple(criteria_finished.shape)}, expected a per-row bool of "
3610 f"shape {tuple(finished_sequences.shape)} (one entry per sequence)."
3611 )
3612 finished_sequences.logical_or_(criteria_finished)
3614 all_finished = bool(any_stop_active and finished_sequences.all().item())
3616 yield sampled_tokens, final_logits, all_finished
3618 if all_finished: 3618 ↛ 3619line 3618 didn't jump to line 3619 because the condition on line 3618 was never true
3619 return
3621 def _resolve_generation_caching(self, use_past_kv_cache: bool, batched: bool) -> bool:
3622 """Honor adapter caching/batching limits (recurrent/conv decoders have no KV
3623 cache; batching is rejected where padding can't be masked, not mis-generated)."""
3624 if batched and not getattr(self.adapter, "supports_batched_generation", True):
3625 architecture = self.cfg.architecture or type(self.adapter).__name__
3626 raise NotImplementedError(
3627 f"Batched generation is not supported by {architecture}: its forward does not "
3628 "apply an attention mask, so padded rows would corrupt the output. Generate one "
3629 "sequence at a time, or pass equal-length inputs as a tensor."
3630 )
3631 if not getattr(self.adapter, "supports_kv_cache", True):
3632 return False
3633 return use_past_kv_cache
3635 def _ensure_generation_supported(self, api_name: str) -> None:
3636 """Reject autoregressive generation for forward-only architectures."""
3637 if not self.adapter.supports_generation:
3638 architecture = self.cfg.architecture or type(self.adapter).__name__
3639 hint = (
3640 " Use diffusion_generate() — this architecture samples by iterative denoising, "
3641 "not left-to-right."
3642 if getattr(self.adapter, "native_sampler", None)
3643 else ""
3644 )
3645 raise NotImplementedError(
3646 f"TransformerBridge.{api_name}() generation is not supported by "
3647 f"the {architecture} architecture.{hint}"
3648 )
3650 def generate(
3651 self,
3652 input: Union[str, List[str], torch.Tensor] = "",
3653 max_new_tokens: int = 10,
3654 stop_at_eos: bool = True,
3655 eos_token_id: Optional[int] = None,
3656 do_sample: bool = True,
3657 top_k: Optional[int] = None,
3658 top_p: Optional[float] = None,
3659 temperature: float = 1.0,
3660 freq_penalty: float = 0.0,
3661 repetition_penalty: float = 1.0,
3662 use_past_kv_cache: bool = True,
3663 prepend_bos: Optional[bool] = None,
3664 padding_side: Optional[str] = None,
3665 return_type: Optional[str] = "input",
3666 verbose: bool = True,
3667 output_logits: bool = False,
3668 return_cache: bool = False,
3669 return_input_tokens: bool = False,
3670 names_filter: Optional[Union[str, List[str], Callable[[str], bool]]] = None,
3671 device: Optional[Union[str, torch.device]] = None,
3672 pixel_values: Optional[torch.Tensor] = None,
3673 stop_strings: Optional[Union[str, List[str]]] = None,
3674 stopping_criteria: Optional[Any] = None,
3675 attention_mask: Optional[torch.Tensor] = None,
3676 forced_bos_token_id: Optional[int] = None,
3677 **multimodal_kwargs,
3678 ) -> (
3679 str
3680 | list[str]
3681 | torch.Tensor
3682 | Any
3683 | tuple[Any, ActivationCache]
3684 | tuple[Any, torch.Tensor]
3685 ): # Any for transformers.utils.ModelOutput
3686 # Any: beartype forward ref limitation (beartype#546)
3687 """Sample tokens from the model.
3689 Sample tokens from the model until the model outputs eos_token or max_new_tokens is reached.
3690 This implementation is based on HookedTransformer.generate() to ensure consistent behavior.
3692 Args:
3693 input: Text string, list of strings, or tensor of tokens
3694 max_new_tokens: Maximum number of tokens to generate
3695 stop_at_eos: If True, stop generating tokens when the model outputs eos_token
3696 eos_token_id: The token ID to use for end of sentence
3697 do_sample: If True, sample from the model's output distribution. Otherwise, use greedy search
3698 top_k: Number of tokens to sample from. If None, sample from all tokens
3699 top_p: Probability mass to sample from. If 1.0, sample from all tokens
3700 temperature: Temperature for sampling. Higher values will make the model more random
3701 freq_penalty: Frequency penalty for sampling - how much to penalise previous tokens
3702 repetition_penalty: HuggingFace-style repetition penalty. Values > 1.0 discourage
3703 repetition by dividing positive logits and multiplying negative logits for
3704 previously seen tokens. Default 1.0 (no penalty).
3705 use_past_kv_cache: If True, use KV caching for faster generation
3706 prepend_bos: Whether to prepend a BOS token when tokenizing string inputs.
3707 Defaults to None (uses ``cfg.default_prepend_bos``, typically True).
3708 Pass ``prepend_bos=False`` when the input is pre-formatted chat-template
3709 text that already contains the BOS token to avoid double-BOS.
3710 Ignored when input is already a token tensor.
3711 padding_side: Which side to pad when tokenizing multiple strings of different
3712 lengths. For batched list inputs, left-padding is forced internally for
3713 correct generation behavior. Defaults to None (tokenizer default).
3714 return_type: The type of output to return - 'input', 'str', or 'tokens'
3715 verbose: Not used in Bridge (kept for API compatibility)
3716 output_logits: If True, return a ModelOutput with sequences and logits tuple
3717 return_cache: If True, also return an ActivationCache for the full prompt +
3718 generated sequence, identical to ``run_with_cache(output)``, and the call
3719 returns an ``(output, cache)`` tuple. Implemented as one extra clean forward
3720 over the output, so the cache includes every hook point (attention patterns
3721 included). Supported only for single-sequence, decoder-only text generation;
3722 encoder-decoder, SSM, multimodal, batched, and inputs_embeds inputs raise
3723 NotImplementedError. The cache spans prompt + max_new_tokens and can be large,
3724 use ``names_filter`` to scope it and/or ``device`` to offload it.
3725 return_input_tokens: If True, return an ``(output, input_tokens)`` tuple where
3726 ``input_tokens`` is the token tensor that was actually fed to the model
3727 (after BOS handling). Useful for debugging tokenization, especially when
3728 using chat templates where BOS handling can be subtle. Can be combined
3729 with ``return_cache`` to get ``(output, cache, input_tokens)``.
3730 names_filter: Passed to ``run_with_cache`` when ``return_cache=True``; restricts
3731 which activations are cached (str, list of str, or callable).
3732 device: Passed through when ``return_cache=True`` to offload the cached tensors
3733 to this device (e.g. "cpu") to save accelerator memory.
3734 pixel_values: Optional image tensor for multimodal models. Only passed on the
3735 first generation step (the vision encoder processes the image once, then
3736 embeddings are part of the token sequence for subsequent steps).
3737 stop_strings: Optional string or list of strings. A sequence stops once its
3738 generated text ends with one of these strings, using HuggingFace's
3739 StopStringCriteria (partial-token-aware, end-anchored) matching.
3740 Requires a tokenizer (raises ValueError otherwise).
3741 Independent of stop_at_eos: either can stop a sequence.
3742 stopping_criteria: Optional HuggingFace stopping criteria, a single
3743 transformers.StoppingCriteria, a list of them, or a StoppingCriteriaList.
3744 Each is called as criterion(input_ids, scores) after every step and ORed
3745 with the other stop signals, where input_ids is the running sequence and
3746 scores is this step's logits ([batch, d_vocab]). Each criterion must return
3747 a per-row bool [batch] (or a scalar bool). stop_strings and stopping_criteria
3748 are supported only for standard decoder-only text generation. Encoder-decoder,
3749 inputs_embeds, and multimodal generation always raise NotImplementedError.
3750 Stateful/SSM models raise only when run with use_past_kv_cache=False (the
3751 default keeps them on the hooked loop). Each error names the supported
3752 alternative.
3753 attention_mask: Optional ``[batch, pos]`` 0/1 mask over the prompt, marking
3754 which prompt tokens are real. Required to generate correctly from an
3755 already-padded token tensor: without it the pad tokens are treated as
3756 real context and every real token's position is shifted, so the
3757 continuation differs from the same prompt unpadded. The mask is extended
3758 by one attended column per generated token. Takes precedence over the
3759 ``padding_side`` heuristic, and unlike it can express an interior gap or
3760 a pad id that also occurs as a real token. Passing ``padding_side``
3761 instead reads the padding off the pad token, which is enough for the
3762 common single-edge case, and raises if this bridge has no tokenizer
3763 or pad id to read it from. On the encoder-decoder and inputs_embeds
3764 paths the mask is forwarded to the model as-is rather than grown per
3765 step, which is what processors emitting one alongside
3766 ``pixel_values`` expect.
3767 forced_bos_token_id: Optional token id seeded as the first decoder token
3768 after ``decoder_start`` on encoder-decoder models. Multilingual
3769 translators (M2M100/MBart/NLLB) select their target language this way.
3770 Raises ValueError on decoder-only models.
3772 Returns:
3773 Generated sequence as string, list of strings, or tensor depending on input type and return_type.
3774 If output_logits=True, returns a ModelOutput-like object with 'sequences' and 'logits' attributes.
3775 If return_cache=True, returns an ``(output, ActivationCache)`` tuple where ``output`` is the
3776 value that would otherwise be returned and the cache equals ``run_with_cache(output)``.
3777 If return_input_tokens=True, returns an ``(output, input_tokens)`` tuple.
3778 If both return_cache and return_input_tokens are True, returns ``(output, cache, input_tokens)``.
3780 Example:
3781 ``out, cache = model.generate(prompt, max_new_tokens=20, return_cache=True)`` returns a
3782 normal ActivationCache over the full prompt + generated sequence (equivalent to
3783 ``run_with_cache(out)``).
3785 ``out, input_tokens = model.generate(prompt, return_input_tokens=True)`` returns
3786 the tokens that were fed to the model, useful for verifying BOS handling with
3787 chat templates.
3788 """
3789 self._ensure_generation_supported("generate")
3790 # padding_side is handled internally: for batched list inputs, left-padding
3791 # is forced to ensure correct generation. See _is_batched_list logic below.
3793 # Stateful dispatch is decided after input parsing so we can fall back
3794 # to hf_generate() for input types the stateful loop doesn't handle.
3795 is_stateful_model = getattr(self.cfg, "is_stateful", False)
3797 _is_batched_list = isinstance(input, list) and len(input) > 1
3798 use_past_kv_cache = self._resolve_generation_caching(use_past_kv_cache, _is_batched_list)
3800 _generate_from_embeds = False
3801 _encdec_early = hasattr(self.original_model, "config") and getattr(
3802 self.original_model.config, "is_encoder_decoder", False
3803 )
3804 if isinstance(input, str):
3805 if _encdec_early:
3806 # Deliberate divergence: prepend_bos is IGNORED for enc-dec
3807 # string/list input. Encoder input follows the tokenizer's own
3808 # recipe (lang token + trailing </s>); to_tokens' decoder-style
3809 # BOS policy corrupts it — m2m100 degenerates to loops.
3810 input_tokens = self.tokenizer(input, return_tensors="pt")["input_ids"].to(
3811 self.cfg.device
3812 )
3813 else:
3814 input_tokens = self.to_tokens(
3815 input, prepend_bos=prepend_bos, move_to_device=True, truncate=False
3816 )
3817 input_type = "str"
3818 elif isinstance(input, list):
3819 if _encdec_early:
3820 # Same native-recipe rule as the str branch: to_tokens' BOS
3821 # policy corrupts encoder inputs (stray <s>, dropped </s>).
3822 # Keep the tokenizer's mask too — unequal rows otherwise
3823 # attend over pads in the encoder.
3824 _enc_batch = self.tokenizer(input, return_tensors="pt", padding=True)
3825 input_tokens = _enc_batch["input_ids"].to(self.cfg.device)
3826 if attention_mask is None and "attention_mask" in _enc_batch: 3826 ↛ 3839line 3826 didn't jump to line 3839 because the condition on line 3826 was always true
3827 attention_mask = _enc_batch["attention_mask"].to(self.cfg.device)
3828 else:
3829 # Force left-padding for batched generation so real tokens are
3830 # flush-right and logits[:, -1, :] is always the last real token.
3831 if _is_batched_list: 3831 ↛ 3834line 3831 didn't jump to line 3834 because the condition on line 3831 was always true
3832 _orig_padding_side = self.tokenizer.padding_side
3833 self.tokenizer.padding_side = "left"
3834 input_tokens = self.to_tokens(
3835 input, prepend_bos=prepend_bos, move_to_device=True, truncate=False
3836 )
3837 if _is_batched_list: 3837 ↛ 3839line 3837 didn't jump to line 3839 because the condition on line 3837 was always true
3838 self.tokenizer.padding_side = _orig_padding_side
3839 input_type = "list"
3840 elif isinstance(input, torch.Tensor) and input.is_floating_point():
3841 # inputs_embeds: pre-computed embeddings (e.g., from multimodal models)
3842 input_tokens = input.to(self.cfg.device)
3843 input_type = "embeds"
3844 _generate_from_embeds = True
3845 else:
3846 input_tokens = input.to(self.cfg.device)
3847 input_type = "tokens"
3849 # Without one of these a pre-padded tensor generates as though its pads were
3850 # real context, shifting every real token's position (#1612). An explicit
3851 # mask wins; otherwise the padding is read off the tokens, but only when the
3852 # caller asked for that by passing padding_side. Deriving a mask on the
3853 # default path would silently change behaviour for every existing caller,
3854 # and would demand a real tokenizer where today none is required.
3855 initial_attention_mask: Optional[torch.Tensor] = attention_mask
3856 if initial_attention_mask is not None and (
3857 _generate_from_embeds
3858 or getattr(getattr(self.original_model, "config", None), "is_encoder_decoder", False)
3859 ):
3860 # Growing the mask per step only means something for decoder-only token
3861 # generation. On these paths the mask used to arrive via
3862 # **multimodal_kwargs and be forwarded to the model untouched — as
3863 # processors emit it alongside pixel_values — so keep doing that rather
3864 # than reject a call that worked before this parameter existed.
3865 multimodal_kwargs = {**multimodal_kwargs, "attention_mask": initial_attention_mask}
3866 initial_attention_mask = None
3867 if initial_attention_mask is not None:
3868 if initial_attention_mask.shape != input_tokens.shape:
3869 raise ValueError(
3870 f"attention_mask shape {tuple(initial_attention_mask.shape)} does not "
3871 f"match the prompt shape {tuple(input_tokens.shape)}. Pass a 0/1 mask "
3872 "covering exactly the prompt tokens; generate() extends it itself."
3873 )
3874 initial_attention_mask = initial_attention_mask.to(self.cfg.device)
3875 elif padding_side is not None and input_type == "tokens":
3876 # Reading the padding off the tokens needs a tokenizer with a pad id.
3877 # Without one the argument would be inert, leaving exactly the bug this
3878 # fixes — silently, on a bridge booted without a tokenizer. Say so
3879 # rather than generate something quietly wrong.
3880 if not isinstance(self.tokenizer, PreTrainedTokenizerBase):
3881 raise ValueError(
3882 "generate(padding_side=...) reads the padding off the pad token, "
3883 "which needs a tokenizer; this bridge has none. Pass "
3884 "attention_mask=... to state the padding directly instead."
3885 )
3886 if self.tokenizer.pad_token_id is None:
3887 raise ValueError(
3888 "generate(padding_side=...) reads the padding off the pad token, "
3889 "but this tokenizer has no pad_token_id. Set one, or pass "
3890 "attention_mask=... to state the padding directly instead."
3891 )
3892 _prepend = self.cfg.default_prepend_bos if prepend_bos is None else prepend_bos
3893 _orig_side = self.tokenizer.padding_side
3894 self.tokenizer.padding_side = padding_side
3895 try:
3896 initial_attention_mask = utils.get_attention_mask(
3897 self.tokenizer, input_tokens, _prepend
3898 ).to(self.cfg.device)
3899 finally:
3900 self.tokenizer.padding_side = _orig_side
3901 # An all-ones mask is what the model assumes anyway; skipping it keeps
3902 # the unpadded path byte-identical to before.
3903 if initial_attention_mask is not None and bool(initial_attention_mask.all()): 3903 ↛ 3904line 3903 didn't jump to line 3904 because the condition on line 3903 was never true
3904 initial_attention_mask = None
3906 # Determine return type
3907 if return_type == "input":
3908 if input_type in ["str", "list"]:
3909 return_type = "str"
3910 elif input_type == "embeds":
3911 return_type = "tokens"
3912 else:
3913 return_type = "tokens"
3915 batch_size = input_tokens.shape[0]
3917 # Setup EOS token handling
3918 stop_tokens = []
3919 eos_token_for_padding = 0
3920 if stop_at_eos:
3921 tokenizer_has_eos_token = (
3922 self.tokenizer is not None and self.tokenizer.eos_token_id is not None
3923 )
3924 if eos_token_id is None:
3925 # Some chat models use a turn-end token that differs from the
3926 # tokenizer's primary EOS. Let adapters provide the full stop
3927 # set via cfg.eos_token_id; otherwise fall back to the tokenizer.
3928 eos_token_id = getattr(self.cfg, "eos_token_id", None)
3929 if eos_token_id is None:
3930 assert (
3931 tokenizer_has_eos_token
3932 ), "Must pass eos_token_id if stop_at_eos is True and tokenizer is None or has no eos_token_id"
3933 assert self.tokenizer is not None
3934 eos_token_id = self.tokenizer.eos_token_id
3936 if isinstance(eos_token_id, int):
3937 stop_tokens = [eos_token_id]
3938 eos_token_for_padding = eos_token_id
3939 else:
3940 stop_tokens = list(eos_token_id)
3941 if tokenizer_has_eos_token:
3942 assert self.tokenizer is not None
3943 eos_token_for_padding = self.tokenizer.eos_token_id
3944 else:
3945 eos_token_for_padding = eos_token_id[0]
3947 # Track which sequences have finished
3948 finished_sequences = torch.zeros(batch_size, dtype=torch.bool, device=self.cfg.device)
3950 # Optionally collect logits at each generation step for downstream tooling/tests
3951 logits_seq_list: list[torch.Tensor] | None = [] if output_logits else None
3953 # Detect encoder-decoder models (T5, BART, etc.)
3954 is_encoder_decoder = hasattr(self.original_model, "config") and getattr(
3955 self.original_model.config, "is_encoder_decoder", False
3956 )
3957 if forced_bos_token_id is None and is_encoder_decoder:
3958 # HF's generate() applies generation_config defaults; bart-large-cnn
3959 # pins forced_bos_token_id=0 there and degrades without it.
3960 forced_bos_token_id = getattr(
3961 getattr(self.original_model, "generation_config", None),
3962 "forced_bos_token_id",
3963 None,
3964 )
3965 if forced_bos_token_id is not None and not is_encoder_decoder: 3965 ↛ 3968line 3965 didn't jump to line 3968 because the condition on line 3965 was never true
3966 # Raise before any state mutation (_capture_hf_cache) and before
3967 # the stateful hf_generate early-return would drop the kwarg.
3968 raise ValueError("forced_bos_token_id is only meaningful for encoder-decoder models")
3970 # return_cache recomputes run_with_cache on the generated output (see issue #697).
3971 # That is well-defined only for single-sequence, decoder-only text generation, so
3972 # reject the paths whose cache would be wrong/undefined, with a clear pointer to the
3973 # run_with_cache workaround. Fail fast here, before any generation work.
3974 if return_cache:
3975 if is_encoder_decoder: 3975 ↛ 3976line 3975 didn't jump to line 3976 because the condition on line 3975 was never true
3976 raise NotImplementedError(
3977 "generate(return_cache=True) is not supported for encoder-decoder "
3978 "models yet. Run run_with_cache on the generated output instead."
3979 )
3980 if is_stateful_model: 3980 ↛ 3981line 3980 didn't jump to line 3981 because the condition on line 3980 was never true
3981 raise NotImplementedError(
3982 "generate(return_cache=True) is not supported for stateful/SSM models "
3983 "(e.g. Mamba); they do not expose standard transformer hook points."
3984 )
3985 if pixel_values is not None or multimodal_kwargs: 3985 ↛ 3986line 3985 didn't jump to line 3986 because the condition on line 3985 was never true
3986 raise NotImplementedError(
3987 "generate(return_cache=True) is not supported for multimodal generation "
3988 "yet. Run run_with_cache on the generated output instead."
3989 )
3990 if _generate_from_embeds:
3991 raise NotImplementedError(
3992 "generate(return_cache=True) requires token input, not inputs_embeds."
3993 )
3994 if batch_size > 1:
3995 raise NotImplementedError(
3996 "generate(return_cache=True) is not supported for batched/multi-prompt "
3997 "generation yet. Pass a single prompt, or run run_with_cache on each "
3998 "output sequence."
3999 )
4001 # HF cache flows opaquely through the component chain via
4002 # _reconstruct_attention() → _update_kv_cache() on each layer.
4003 _hf_kv_cache = None
4004 if use_past_kv_cache and is_encoder_decoder:
4005 # Encoder-decoder models (T5, BART) don't support the opaque
4006 # cache path — silently disable rather than crash, since
4007 # use_past_kv_cache=True is the default.
4008 use_past_kv_cache = False
4010 # SSMs (Mamba/Mamba-2) run through a dedicated cache path so hooks
4011 # fire on every step. Unsupported input types fall back to hf_generate().
4012 use_stateful_cache = (
4013 is_stateful_model
4014 and use_past_kv_cache
4015 and not is_encoder_decoder
4016 and not _generate_from_embeds
4017 and pixel_values is None
4018 and not multimodal_kwargs
4019 )
4021 # stop_strings / stopping_criteria are applied inside the hooked _generate_tokens
4022 # loop, so they are supported only on the standard decoder-only text path. Reject
4023 # the paths that route around that loop with a clear error rather than silently
4024 # dropping the kwargs. This must run before the stateful delegation below.
4025 stopping_criteria_list = self._resolve_stopping_criteria(stop_strings, stopping_criteria)
4026 if stopping_criteria_list is not None:
4027 if is_encoder_decoder:
4028 _unsupported = "encoder-decoder models"
4029 elif _generate_from_embeds:
4030 _unsupported = "inputs_embeds generation"
4031 elif pixel_values is not None or multimodal_kwargs:
4032 _unsupported = "multimodal (pixel_values) generation"
4033 else:
4034 _unsupported = None
4035 if _unsupported is not None:
4036 raise NotImplementedError(
4037 f"stop_strings/stopping_criteria are not supported for {_unsupported} in "
4038 "TransformerBridge.generate(). Call hf_generate(...), which runs "
4039 "HuggingFace's own generation loop and supports HF-native stopping on "
4040 "those inputs."
4041 )
4042 if is_stateful_model and not use_stateful_cache:
4043 # Reached only for a stateful/SSM model with use_past_kv_cache=False: the
4044 # hooked loop needs the stateful cache, so generate() would otherwise fall
4045 # back to hf_generate() and drop these kwargs. The default cache setting
4046 # keeps generation on the hooked loop, where stopping is applied.
4047 raise NotImplementedError(
4048 "stop_strings/stopping_criteria on a stateful/SSM model require the "
4049 "stateful cache path, which runs only with use_past_kv_cache=True (the "
4050 "default). With use_past_kv_cache=False generate() falls back to "
4051 "hf_generate(). Set use_past_kv_cache=True to keep stopping on the hooked "
4052 "loop, or call hf_generate(...) directly for HF-native stopping."
4053 )
4054 # Finished rows are overwritten with this id so they stop emitting real tokens
4055 # while the rest of a batch keeps going. stop_at_eos already set a sensible
4056 # value, otherwise fall back to the tokenizer pad/eos id. (For a single
4057 # sequence this id is never read: the loop exits when the row finishes.)
4058 if not stop_at_eos:
4059 _pad_id = None
4060 if self.tokenizer is not None:
4061 _pad_id = (
4062 self.tokenizer.pad_token_id
4063 if self.tokenizer.pad_token_id is not None
4064 else self.tokenizer.eos_token_id
4065 )
4066 if _pad_id is not None:
4067 eos_token_for_padding = _pad_id
4068 elif batch_size > 1:
4069 raise ValueError(
4070 "Batched generation with stopping_criteria and stop_at_eos=False "
4071 "needs a padding token to freeze finished rows, but no tokenizer "
4072 "pad/eos id is available. Set stop_at_eos=True, use a tokenizer with "
4073 "a pad or eos token, or generate one sequence at a time."
4074 )
4076 if is_stateful_model and not use_stateful_cache: 4076 ↛ 4077line 4076 didn't jump to line 4077 because the condition on line 4076 was never true
4077 hf_kwargs: dict[str, Any] = {
4078 "max_new_tokens": max_new_tokens,
4079 "do_sample": do_sample,
4080 "temperature": temperature,
4081 }
4082 if top_k is not None:
4083 hf_kwargs["top_k"] = top_k
4084 if top_p is not None:
4085 hf_kwargs["top_p"] = top_p
4086 if eos_token_id is not None:
4087 hf_kwargs["eos_token_id"] = eos_token_id
4088 return self.hf_generate(input, **hf_kwargs)
4090 # SSM cache is built once and mutated in place across forward calls.
4091 # Adapter owns the cache-type choice; new SSMs just override
4092 # create_stateful_cache().
4093 mamba_cache: Any = None
4094 mamba_conv_kernel: int = 0
4095 if use_stateful_cache:
4096 hf_model: Any = self.original_model
4097 mamba_conv_kernel = int(getattr(hf_model.config, "conv_kernel", 4))
4098 cache_dtype = self.cfg.dtype or torch.float32
4099 mamba_cache = self.adapter.create_stateful_cache(
4100 hf_model=hf_model,
4101 batch_size=batch_size,
4102 device=self.cfg.device,
4103 dtype=cache_dtype,
4104 )
4106 if use_past_kv_cache and not use_stateful_cache:
4107 self._capture_hf_cache = True # Signal forward() to stash cache
4109 # Generate tokens
4110 current_tokens = input_tokens.clone()
4111 # For inputs_embeds generation, also track generated token IDs for decoding
4112 if _generate_from_embeds:
4113 generated_token_ids: list[torch.Tensor] = []
4114 sampled_tokens_list = []
4116 # For encoder-decoder models, keep encoder input fixed and grow decoder input
4117 if is_encoder_decoder:
4118 encoder_input = input_tokens.clone()
4119 decoder_start_token_id = getattr(
4120 self.original_model.config, "decoder_start_token_id", None
4121 )
4122 if decoder_start_token_id is None:
4123 # HF's fallback chain: bos, then eos (MBart-family checkpoints
4124 # like IndicBART leave decoder_start unset and start from EOS).
4125 fallback = getattr(self.original_model.config, "bos_token_id", None)
4126 if fallback is None:
4127 fallback = getattr(self.original_model.config, "eos_token_id", None)
4128 if isinstance(fallback, (list, tuple)): 4128 ↛ 4129line 4128 didn't jump to line 4129 because the condition on line 4128 was never true
4129 fallback = fallback[0]
4130 decoder_start_token_id = fallback if fallback is not None else 0
4131 decoder_tokens = torch.full(
4132 (batch_size, 1),
4133 decoder_start_token_id,
4134 dtype=input_tokens.dtype,
4135 device=self.cfg.device,
4136 )
4137 if forced_bos_token_id is not None:
4138 # Multilingual seq2seq (M2M100/MBart/NLLB) selects the target
4139 # language via the first decoder token after decoder_start.
4140 forced = torch.full(
4141 (batch_size, 1),
4142 forced_bos_token_id,
4143 dtype=input_tokens.dtype,
4144 device=self.cfg.device,
4145 )
4146 decoder_tokens = torch.cat([decoder_tokens, forced], dim=1)
4148 try:
4149 for sampled_tokens, final_logits, all_finished in self._generate_tokens(
4150 current_tokens,
4151 input_tokens,
4152 batch_size,
4153 max_new_tokens=max_new_tokens,
4154 do_sample=do_sample,
4155 top_k=top_k,
4156 top_p=top_p,
4157 temperature=temperature,
4158 freq_penalty=freq_penalty,
4159 repetition_penalty=repetition_penalty,
4160 stop_at_eos=stop_at_eos,
4161 stop_tokens=stop_tokens,
4162 eos_token_for_padding=eos_token_for_padding,
4163 finished_sequences=finished_sequences,
4164 use_past_kv_cache=use_past_kv_cache,
4165 use_stateful_cache=use_stateful_cache,
4166 mamba_cache=mamba_cache,
4167 mamba_conv_kernel=mamba_conv_kernel,
4168 is_encoder_decoder=is_encoder_decoder,
4169 _is_batched_list=_is_batched_list,
4170 _generate_from_embeds=_generate_from_embeds,
4171 encoder_input=encoder_input if is_encoder_decoder else None,
4172 decoder_tokens=decoder_tokens if is_encoder_decoder else None,
4173 generated_token_ids=generated_token_ids if _generate_from_embeds else None,
4174 pixel_values=pixel_values,
4175 multimodal_kwargs=multimodal_kwargs if multimodal_kwargs else {},
4176 verbose=verbose,
4177 stopping_criteria_list=stopping_criteria_list,
4178 initial_attention_mask=initial_attention_mask,
4179 min_decoder_length=(
4180 getattr(
4181 getattr(self.original_model, "generation_config", None),
4182 "min_length",
4183 None,
4184 )
4185 if is_encoder_decoder
4186 else None
4187 ),
4188 ngram_processor=(self._encdec_ngram_processor() if is_encoder_decoder else None),
4189 encoder_attention_mask=(attention_mask if is_encoder_decoder else None),
4190 ):
4191 sampled_tokens_list.append(sampled_tokens.unsqueeze(1))
4192 if logits_seq_list is not None:
4193 logits_seq_list.append(final_logits.clone())
4194 if all_finished:
4195 break
4196 finally:
4197 self._capture_hf_cache = False
4198 if hasattr(self, "_last_hf_cache"): 4198 ↛ 4199line 4198 didn't jump to line 4199 because the condition on line 4198 was never true
4199 del self._last_hf_cache
4201 # Concatenate all sampled tokens
4202 sampled_tokens = torch.cat(sampled_tokens_list, dim=1)
4203 if is_encoder_decoder:
4204 # Reconstruct full decoder sequence: start token + generated tokens
4205 decoder_seed_len = 2 if forced_bos_token_id is not None else 1
4206 output_tokens = torch.cat([decoder_tokens[:, :decoder_seed_len], sampled_tokens], dim=1)
4207 elif _generate_from_embeds:
4208 # For inputs_embeds, we only have the generated token IDs (no input token IDs)
4209 output_tokens = sampled_tokens
4210 else:
4211 output_tokens = torch.cat([input_tokens, sampled_tokens], dim=1)
4213 # Build the formatted output (shape unchanged: ModelOutput / str / list[str] / tokens).
4214 result: Any
4215 if output_logits and logits_seq_list is not None:
4216 from transformers.utils import ModelOutput # type: ignore
4218 def _logits_to_tuple(logits_list: list[torch.Tensor]) -> tuple[torch.Tensor, ...]:
4219 assert logits_list is not None
4220 # Convert list of [batch, vocab] tensors to tuple
4221 return tuple(logits_list)
4223 try:
4224 from transformers.generation.utils import GenerateDecoderOnlyOutput
4226 # HF-compatible ModelOutput structure.
4227 # GenerateDecoderOnlyOutput expects: sequences, scores (optional), logits (optional)
4228 result = GenerateDecoderOnlyOutput(
4229 sequences=cast(torch.LongTensor, output_tokens),
4230 # HF's type hint says tuple[FloatTensor] but should be tuple[FloatTensor, ...]
4231 # (variable-length tuple with one element per generated token)
4232 logits=_logits_to_tuple(logits_seq_list), # type: ignore[arg-type]
4233 )
4234 except (ImportError, AttributeError):
4235 # Fallback if GenerateDecoderOnlyOutput not available in this transformers version
4236 result = ModelOutput(
4237 sequences=output_tokens,
4238 logits=_logits_to_tuple(logits_seq_list),
4239 )
4240 elif return_type == "str":
4241 assert self.tokenizer is not None
4242 if input_type == "str":
4243 result = self.tokenizer.decode(output_tokens[0], skip_special_tokens=True)
4244 else:
4245 decoded_texts = [
4246 self.tokenizer.decode(tokens, skip_special_tokens=True)
4247 for tokens in output_tokens
4248 ]
4249 result = decoded_texts[0] if len(decoded_texts) == 1 else decoded_texts
4250 else: # return_type == "tokens"
4251 result = output_tokens
4253 if not return_cache and not return_input_tokens:
4254 return result
4256 if return_cache:
4257 # return_cache: recompute one clean forward over the full generated sequence so the
4258 # cache is identical to run_with_cache(output_tokens) - all hook points, including
4259 # attention patterns. The guards above restrict this to single-sequence, decoder-only
4260 # text generation (see issue #697).
4261 _, cache = self.run_with_cache(output_tokens, names_filter=names_filter, device=device)
4262 if return_input_tokens:
4263 return result, cache, input_tokens
4264 return result, cache
4266 # return_input_tokens only (no cache)
4267 return result, input_tokens
4269 @torch.no_grad()
4270 def diffusion_generate(
4271 self,
4272 input: Union[str, List[str], torch.Tensor],
4273 max_new_tokens: int = 32,
4274 prepend_bos: Optional[bool] = None,
4275 **kwargs: Any,
4276 ) -> Union[str, torch.Tensor]:
4277 """Sample from a non-autoregressive (diffusion) architecture.
4279 Delegates to the model's own sampler, which calls the model through
4280 ``__call__`` so bridge hooks fire on every denoising step.
4281 """
4282 sampler_name = getattr(self.adapter, "native_sampler", None)
4283 architecture = self.cfg.architecture or type(self.adapter).__name__
4284 if sampler_name is None: 4284 ↛ 4285line 4284 didn't jump to line 4285 because the condition on line 4284 was never true
4285 raise NotImplementedError(
4286 f"{architecture} has no native sampler; use generate() for autoregressive "
4287 "architectures."
4288 )
4289 sampler = getattr(self.original_model, sampler_name, None)
4290 if sampler is None: 4290 ↛ 4291line 4290 didn't jump to line 4291 because the condition on line 4290 was never true
4291 raise NotImplementedError(
4292 f"{architecture} declares native_sampler={sampler_name!r} but the loaded model "
4293 "has no such method."
4294 )
4296 was_string = isinstance(input, str)
4297 if isinstance(input, list) and len(input) > 1: 4297 ↛ 4301line 4297 didn't jump to line 4301 because the condition on line 4297 was never true
4298 # Unequal prompts would be right-padded into the sampler's canvas,
4299 # where pad tokens read as real context. generate() gates batching
4300 # for the same reason; do not silently corrupt rows here.
4301 raise NotImplementedError(
4302 f"diffusion_generate() does not support batched prompts for {architecture}: "
4303 "the native samplers condition on a padded canvas. Sample one prompt at a time."
4304 )
4305 if isinstance(input, torch.Tensor): 4305 ↛ 4310line 4305 didn't jump to line 4310 because the condition on line 4305 was always true
4306 tokens = input.to(self.cfg.device)
4307 else:
4308 # Tokenization is the bridge's concern, not the sampler's — absorb
4309 # prepend_bos here rather than forwarding it into native kwargs.
4310 tokens = self.to_tokens(
4311 input, prepend_bos=prepend_bos, move_to_device=True, truncate=False
4312 )
4314 sampler_kwargs = self.adapter.native_sampler_kwargs(max_new_tokens, tokens.shape[-1])
4315 sampler_kwargs.update(kwargs)
4316 output = sampler(tokens, **sampler_kwargs)
4317 # Samplers return either bare ids or a generation output object.
4318 sequences = getattr(output, "sequences", output)
4319 # They also disagree on whether the prompt is included: Dream and Gidd
4320 # return the whole canvas, LLaDA2 slices it off. Normalize to
4321 # generate()'s contract (prompt + continuation).
4322 if isinstance(sequences, torch.Tensor) and sequences.ndim == tokens.ndim: 4322 ↛ 4333line 4322 didn't jump to line 4333 because the condition on line 4322 was always true
4323 prompt_len = tokens.shape[-1]
4324 # Compare on one device: torch.equal raises on a device mismatch,
4325 # which some samplers produce by assembling output on CPU.
4326 sequences = sequences.to(tokens.device)
4327 includes_prompt = sequences.shape[-1] >= prompt_len and torch.equal(
4328 sequences[..., :prompt_len], tokens
4329 )
4330 if not includes_prompt:
4331 sequences = torch.cat([tokens, sequences], dim=-1)
4333 if was_string and self.tokenizer is not None: 4333 ↛ 4334line 4333 didn't jump to line 4334 because the condition on line 4333 was never true
4334 return self.tokenizer.decode(sequences[0], skip_special_tokens=True)
4335 return sequences
4337 @torch.no_grad()
4338 def generate_stream(
4339 self,
4340 input: Union[str, List[str], torch.Tensor] = "",
4341 max_new_tokens: int = 10,
4342 max_tokens_per_yield: int = 25,
4343 stop_at_eos: bool = True,
4344 eos_token_id: Optional[int] = None,
4345 do_sample: bool = True,
4346 top_k: Optional[int] = None,
4347 top_p: Optional[float] = None,
4348 temperature: float = 1.0,
4349 freq_penalty: float = 0.0,
4350 repetition_penalty: float = 1.0,
4351 use_past_kv_cache: bool = True,
4352 prepend_bos: Optional[bool] = None,
4353 padding_side: Optional[str] = None,
4354 return_type: Optional[str] = "input",
4355 verbose: bool = True,
4356 stop_strings: Optional[Union[str, List[str]]] = None,
4357 stopping_criteria: Optional[Any] = None,
4358 ) -> Generator[Union[torch.Tensor, str], None, None]:
4359 """Stream tokens from the model as they are generated.
4361 Yields batches of tokens progressively during generation rather than
4362 waiting for the entire sequence. Uses the same core loop as generate().
4364 Args:
4365 input: Text string, list of strings, or tensor of tokens.
4366 max_new_tokens: Maximum number of tokens to generate.
4367 max_tokens_per_yield: Yield accumulated tokens every this many steps.
4368 stop_at_eos: If True, stop when eos_token is produced.
4369 eos_token_id: Token ID(s) for end of sentence. Defaults to tokenizer's.
4370 do_sample: If True, sample; otherwise greedy.
4371 top_k: Top-k sampling. None means no filtering.
4372 top_p: Nucleus sampling threshold.
4373 temperature: Sampling temperature.
4374 freq_penalty: Frequency penalty for previous tokens.
4375 repetition_penalty: HF-style repetition penalty (>1.0 discourages repeats).
4376 use_past_kv_cache: Use KV caching for faster generation.
4377 prepend_bos: Whether to prepend a BOS token when tokenizing string inputs.
4378 Defaults to None (uses ``cfg.default_prepend_bos``, typically True).
4379 Pass ``prepend_bos=False`` when the input is pre-formatted chat-template
4380 text that already contains the BOS token to avoid double-BOS.
4381 Ignored when input is already a token tensor.
4382 padding_side: Which side to pad for batched list inputs. Left-padding
4383 is forced internally for batched generation.
4384 return_type: 'input' (match input type), 'str', or 'tokens'.
4385 verbose: Show progress bar.
4386 stop_strings: Optional string or list of strings. A sequence stops once its
4387 generated text ends with one of them (HF StopStringCriteria). Requires a
4388 tokenizer. See generate() for details.
4389 stopping_criteria: Optional transformers StoppingCriteria, list, or
4390 StoppingCriteriaList, called as criterion(input_ids, scores) each step
4391 (scores is the step's logits). See generate() for the full contract.
4393 Yields:
4394 Token tensors [batch, seq_len] or strings, accumulated up to
4395 max_tokens_per_yield tokens between yields. First yield includes
4396 the input tokens; subsequent yields contain only new tokens.
4397 """
4398 self._ensure_generation_supported("generate_stream")
4399 # --- Input parsing (mirrors generate()) ---
4400 _is_batched_list = isinstance(input, list) and len(input) > 1
4401 use_past_kv_cache = self._resolve_generation_caching(use_past_kv_cache, _is_batched_list)
4403 _encdec_early = hasattr(self.original_model, "config") and getattr(
4404 self.original_model.config, "is_encoder_decoder", False
4405 )
4406 if isinstance(input, str):
4407 if _encdec_early: 4407 ↛ 4409line 4407 didn't jump to line 4409 because the condition on line 4407 was never true
4408 # Native recipe: to_tokens' BOS policy corrupts encoder inputs.
4409 input_tokens = self.tokenizer(input, return_tensors="pt")["input_ids"].to(
4410 self.cfg.device
4411 )
4412 else:
4413 input_tokens = self.to_tokens(
4414 input, prepend_bos=prepend_bos, move_to_device=True, truncate=False
4415 )
4416 input_type = "str"
4417 elif isinstance(input, list): 4417 ↛ 4418line 4417 didn't jump to line 4418 because the condition on line 4417 was never true
4418 if _encdec_early:
4419 input_tokens = self.tokenizer(input, return_tensors="pt", padding=True)[
4420 "input_ids"
4421 ].to(self.cfg.device)
4422 elif _is_batched_list:
4423 _orig_ps = self.tokenizer.padding_side
4424 self.tokenizer.padding_side = "left"
4425 try:
4426 input_tokens = self.to_tokens(
4427 input, prepend_bos=prepend_bos, move_to_device=True, truncate=False
4428 )
4429 finally:
4430 self.tokenizer.padding_side = _orig_ps
4431 else:
4432 input_tokens = self.to_tokens(
4433 input, prepend_bos=prepend_bos, move_to_device=True, truncate=False
4434 )
4435 input_type = "list"
4436 else:
4437 input_tokens = input.to(self.cfg.device)
4438 input_type = "tokens"
4440 if return_type == "input": 4440 ↛ 4441line 4440 didn't jump to line 4441 because the condition on line 4440 was never true
4441 return_type = "str" if input_type in ["str", "list"] else "tokens"
4443 batch_size = input_tokens.shape[0]
4445 # --- EOS setup ---
4446 stop_tokens: List[int] = []
4447 eos_token_for_padding = 0
4448 if stop_at_eos:
4449 tokenizer_has_eos_token = (
4450 self.tokenizer is not None and self.tokenizer.eos_token_id is not None
4451 )
4452 if eos_token_id is None:
4453 # Some chat models use a turn-end token that differs from the
4454 # tokenizer's primary EOS. Let adapters provide the full stop
4455 # set via cfg.eos_token_id; otherwise fall back to the tokenizer.
4456 eos_token_id = getattr(self.cfg, "eos_token_id", None)
4457 if eos_token_id is None:
4458 assert (
4459 tokenizer_has_eos_token
4460 ), "Must pass eos_token_id if stop_at_eos is True and tokenizer is None or has no eos_token_id"
4461 assert self.tokenizer is not None
4462 eos_token_id = self.tokenizer.eos_token_id
4463 if isinstance(eos_token_id, int):
4464 stop_tokens = [eos_token_id]
4465 eos_token_for_padding = eos_token_id
4466 else:
4467 stop_tokens = list(eos_token_id)
4468 if tokenizer_has_eos_token: 4468 ↛ 4469line 4468 didn't jump to line 4469 because the condition on line 4468 was never true
4469 assert self.tokenizer is not None
4470 eos_token_for_padding = self.tokenizer.eos_token_id
4471 else:
4472 eos_token_for_padding = eos_token_id[0]
4474 finished_sequences = torch.zeros(batch_size, dtype=torch.bool, device=self.cfg.device)
4476 # stop_strings / stopping_criteria: build the combined criteria list (validates
4477 # tokenizer for stop_strings). generate_stream only runs the decoder-only text
4478 # path, so no path guards are needed here.
4479 stopping_criteria_list = self._resolve_stopping_criteria(stop_strings, stopping_criteria)
4480 if stopping_criteria_list is not None and not stop_at_eos:
4481 _pad_id = None
4482 if self.tokenizer is not None:
4483 _pad_id = (
4484 self.tokenizer.pad_token_id
4485 if self.tokenizer.pad_token_id is not None
4486 else self.tokenizer.eos_token_id
4487 )
4488 if _pad_id is not None:
4489 eos_token_for_padding = _pad_id
4490 elif batch_size > 1: 4490 ↛ 4499line 4490 didn't jump to line 4499 because the condition on line 4490 was always true
4491 raise ValueError(
4492 "Batched generate_stream with stopping_criteria and stop_at_eos=False "
4493 "needs a padding token to freeze finished rows, but no tokenizer pad/eos "
4494 "id is available. Set stop_at_eos=True or use a tokenizer with a pad/eos "
4495 "token."
4496 )
4498 # --- Cache setup ---
4499 if use_past_kv_cache:
4500 self._capture_hf_cache = True
4502 current_tokens = input_tokens.clone()
4504 # --- Streaming loop ---
4505 # All yields are token tensors [batch, seq_len]. Each yield contains
4506 # only the newly generated tokens since the previous yield (the first
4507 # yield additionally prepends the input tokens for context).
4508 accumulated_tokens: Optional[torch.Tensor] = None
4509 tokens_since_last_yield = 0
4511 def _maybe_decode(
4512 tokens: torch.Tensor,
4513 ) -> Union[torch.Tensor, str]:
4514 if return_type == "str":
4515 assert self.tokenizer is not None
4516 return self.tokenizer.decode(tokens[0], skip_special_tokens=True)
4517 return tokens
4519 try:
4520 for step_idx, (sampled_tokens, _, all_finished) in enumerate(
4521 self._generate_tokens(
4522 current_tokens,
4523 input_tokens,
4524 batch_size,
4525 max_new_tokens=max_new_tokens,
4526 do_sample=do_sample,
4527 top_k=top_k,
4528 top_p=top_p,
4529 temperature=temperature,
4530 freq_penalty=freq_penalty,
4531 repetition_penalty=repetition_penalty,
4532 stop_at_eos=stop_at_eos,
4533 stop_tokens=stop_tokens,
4534 eos_token_for_padding=eos_token_for_padding,
4535 finished_sequences=finished_sequences,
4536 use_past_kv_cache=use_past_kv_cache,
4537 use_stateful_cache=False,
4538 mamba_cache=None,
4539 mamba_conv_kernel=0,
4540 is_encoder_decoder=False,
4541 _is_batched_list=_is_batched_list,
4542 _generate_from_embeds=False,
4543 encoder_input=None,
4544 decoder_tokens=None,
4545 generated_token_ids=None,
4546 pixel_values=None,
4547 multimodal_kwargs={},
4548 verbose=verbose,
4549 stopping_criteria_list=stopping_criteria_list,
4550 )
4551 ):
4552 new_tokens = sampled_tokens.unsqueeze(-1)
4554 if step_idx == 0:
4555 accumulated_tokens = torch.cat([input_tokens, new_tokens], dim=-1)
4556 tokens_since_last_yield = accumulated_tokens.shape[1]
4557 else:
4558 if accumulated_tokens is None:
4559 accumulated_tokens = new_tokens
4560 else:
4561 accumulated_tokens = torch.cat([accumulated_tokens, new_tokens], dim=-1)
4562 tokens_since_last_yield += 1
4564 if tokens_since_last_yield >= max_tokens_per_yield:
4565 yield _maybe_decode(accumulated_tokens)
4566 tokens_since_last_yield = 0
4567 accumulated_tokens = None
4569 if all_finished:
4570 if accumulated_tokens is not None: 4570 ↛ 4573line 4570 didn't jump to line 4573 because the condition on line 4570 was always true
4571 yield _maybe_decode(accumulated_tokens)
4572 accumulated_tokens = None
4573 break
4575 # Yield remainder after loop completes without break
4576 if accumulated_tokens is not None:
4577 yield _maybe_decode(accumulated_tokens)
4578 finally:
4579 self._capture_hf_cache = False
4580 if hasattr(self, "_last_hf_cache"): 4580 ↛ 4581line 4580 didn't jump to line 4581 because the condition on line 4580 was never true
4581 del self._last_hf_cache
4583 def hf_generate(
4584 self,
4585 input: str | list[str] | torch.Tensor = "",
4586 max_new_tokens: int = 10,
4587 stop_at_eos: bool = True,
4588 eos_token_id: int | None = None,
4589 do_sample: bool = True,
4590 top_k: int | None = None,
4591 top_p: float | None = None,
4592 temperature: float = 1.0,
4593 use_past_kv_cache: bool = True,
4594 return_type: str | None = "input",
4595 pixel_values: torch.Tensor | None = None,
4596 **generation_kwargs,
4597 ) -> str | list[str] | torch.Tensor | Any: # Any for HF ModelOutput types
4598 # Any: beartype forward ref limitation (beartype#546)
4599 """Generate text using the underlying HuggingFace model with full HF API support.
4601 This method provides direct access to HuggingFace's generation API, forwarding all
4602 generation parameters (including output_scores, output_logits, output_attentions,
4603 output_hidden_states) directly to the underlying HF model. Use this when you need
4604 full HuggingFace generation features not supported by the standard generate() method.
4606 For standard generation compatible with HookedTransformer, use generate() instead.
4608 Args:
4609 input: Text string, list of strings, or tensor of tokens
4610 max_new_tokens: Maximum number of tokens to generate
4611 stop_at_eos: If True, stop generating tokens when the model outputs eos_token
4612 eos_token_id: The token ID to use for end of sentence
4613 do_sample: If True, sample from the model's output distribution
4614 top_k: Number of tokens to sample from
4615 top_p: Probability mass to sample from
4616 temperature: Temperature for sampling
4617 use_past_kv_cache: If True, use KV caching for faster generation
4618 return_type: The type of output to return - 'input', 'str', or 'tokens'
4619 **generation_kwargs: Additional HuggingFace generation parameters including:
4620 - output_scores: Return generation scores
4621 - output_logits: Return generation logits
4622 - output_attentions: Return attention weights
4623 - output_hidden_states: Return hidden states
4624 - return_dict_in_generate: Return ModelOutput object
4625 - And any other HF generation parameters
4627 Returns:
4628 Generated sequence as string, list of strings, tensor, or HF ModelOutput
4629 depending on input type, return_type, and generation_kwargs.
4631 Example::
4633 # Get full HF ModelOutput with logits and attentions
4634 from transformer_lens import HookedTransformer
4635 model = HookedTransformer.from_pretrained("tiny-stories-1M")
4636 result = model.hf_generate(
4637 "Hello world",
4638 max_new_tokens=5,
4639 output_logits=True,
4640 output_attentions=True,
4641 return_dict_in_generate=True
4642 )
4643 print(result.sequences) # Generated tokens
4644 print(result.logits) # Logits for each generation step
4645 print(result.attentions) # Attention weights
4646 """
4647 self._ensure_generation_supported("hf_generate")
4648 # Handle string input by tokenizing it
4649 if isinstance(input, str):
4650 inputs = self.tokenizer(input, return_tensors="pt", padding=False, truncation=False).to(
4651 self.cfg.device
4652 )
4653 input_ids = inputs["input_ids"]
4654 input_type = "str"
4655 elif isinstance(input, list): 4655 ↛ 4662line 4655 didn't jump to line 4662 because the condition on line 4655 was always true
4656 inputs = self.tokenizer(input, return_tensors="pt", padding=True, truncation=False).to(
4657 self.cfg.device
4658 )
4659 input_ids = inputs["input_ids"]
4660 input_type = "list"
4661 else:
4662 input_ids = input
4663 if input_ids.device != self.cfg.device:
4664 input_ids = input_ids.to(self.cfg.device)
4665 input_type = "tokens"
4667 # Build generation_kwargs from explicit args and kwargs
4668 generation_kwargs = dict(generation_kwargs) if generation_kwargs is not None else {}
4669 generation_kwargs.update(
4670 {
4671 "max_new_tokens": max_new_tokens,
4672 "do_sample": do_sample,
4673 "temperature": temperature,
4674 "pad_token_id": self.tokenizer.eos_token_id,
4675 }
4676 )
4678 if top_k is not None: 4678 ↛ 4679line 4678 didn't jump to line 4679 because the condition on line 4678 was never true
4679 generation_kwargs["top_k"] = top_k
4680 if top_p is not None: 4680 ↛ 4681line 4680 didn't jump to line 4681 because the condition on line 4680 was never true
4681 generation_kwargs["top_p"] = top_p
4682 if eos_token_id is not None: 4682 ↛ 4683line 4682 didn't jump to line 4683 because the condition on line 4682 was never true
4683 generation_kwargs["eos_token_id"] = eos_token_id
4684 elif stop_at_eos and self.tokenizer.eos_token_id is not None: 4684 ↛ 4687line 4684 didn't jump to line 4687 because the condition on line 4684 was always true
4685 generation_kwargs["eos_token_id"] = self.tokenizer.eos_token_id
4687 if pixel_values is not None: 4687 ↛ 4688line 4687 didn't jump to line 4688 because the condition on line 4687 was never true
4688 generation_kwargs["pixel_values"] = pixel_values
4690 if use_past_kv_cache: 4690 ↛ 4694line 4690 didn't jump to line 4694 because the condition on line 4690 was always true
4691 generation_kwargs["use_cache"] = True
4693 # HF dict flags that trigger ModelOutput returns
4694 hf_dict_flags = (
4695 "output_scores",
4696 "output_logits",
4697 "output_attentions",
4698 "output_hidden_states",
4699 )
4701 # If any HF-style output flags are provided, ensure return_dict_in_generate is set
4702 any_flag_set = False
4703 for f in hf_dict_flags:
4704 if generation_kwargs.get(f) is not None:
4705 generation_kwargs[f] = bool(generation_kwargs[f])
4706 any_flag_set = True
4708 if any_flag_set: 4708 ↛ 4712line 4708 didn't jump to line 4712 because the condition on line 4708 was always true
4709 generation_kwargs.setdefault("return_dict_in_generate", True)
4711 # Generate using the original HuggingFace model
4712 with torch.no_grad():
4713 outputs = self.original_model.generate(input_ids, **generation_kwargs) # type: ignore[operator]
4715 # Check if output is a ModelOutput
4716 try:
4717 from transformers.utils import ModelOutput # type: ignore
4719 is_model_output = isinstance(outputs, ModelOutput)
4720 except Exception:
4721 is_model_output = False
4723 # Return based on return_type and input format
4724 if return_type == "input" or return_type is None:
4725 if input_type == "str":
4726 # Decode the full output back to string
4727 if is_model_output and hasattr(outputs, "sequences"): 4727 ↛ 4729line 4727 didn't jump to line 4729 because the condition on line 4727 was always true
4728 return self.tokenizer.decode(outputs.sequences[0], skip_special_tokens=True)
4729 return self.tokenizer.decode(outputs[0], skip_special_tokens=True)
4730 elif input_type == "list": 4730 ↛ 4740line 4730 didn't jump to line 4740 because the condition on line 4730 was always true
4731 # Decode each sequence in the batch
4732 if is_model_output and hasattr(outputs, "sequences"): 4732 ↛ 4737line 4732 didn't jump to line 4737 because the condition on line 4732 was always true
4733 return [
4734 self.tokenizer.decode(seq, skip_special_tokens=True)
4735 for seq in outputs.sequences
4736 ]
4737 return [self.tokenizer.decode(seq, skip_special_tokens=True) for seq in outputs]
4738 else:
4739 # Return the full token sequence including input
4740 return outputs
4741 elif return_type == "tokens": 4741 ↛ 4745line 4741 didn't jump to line 4745 because the condition on line 4741 was always true
4742 return outputs
4743 else:
4744 # For other return types, default to the decoded text
4745 if input_type == "str":
4746 if is_model_output and hasattr(outputs, "sequences"):
4747 return self.tokenizer.decode(outputs.sequences[0], skip_special_tokens=True)
4748 return self.tokenizer.decode(outputs[0], skip_special_tokens=True)
4749 elif input_type == "list":
4750 if is_model_output and hasattr(outputs, "sequences"):
4751 return [
4752 self.tokenizer.decode(seq, skip_special_tokens=True)
4753 for seq in outputs.sequences
4754 ]
4755 return [self.tokenizer.decode(seq, skip_special_tokens=True) for seq in outputs]
4756 else:
4757 return outputs
4759 def prepare_multimodal_inputs(
4760 self,
4761 text: Union[str, List[str]],
4762 images: Optional[Any] = None,
4763 ) -> Dict[str, torch.Tensor]:
4764 """Prepare multimodal inputs using the model's processor.
4766 Converts text and images into model-ready tensors (input_ids, pixel_values,
4767 attention_mask, etc.) using the HuggingFace processor loaded during boot().
4769 Args:
4770 text: Text prompt(s), typically containing image placeholder tokens
4771 (e.g., "<image>" for LLaVA).
4772 images: PIL Image or list of PIL Images to process. Pass None for
4773 text-only inputs on a multimodal model.
4775 Returns:
4776 Dictionary with 'input_ids', 'pixel_values', 'attention_mask', etc.
4777 All tensors are moved to the model's device.
4779 Raises:
4780 ValueError: If model is not multimodal or processor is not available.
4781 """
4782 if not getattr(self.cfg, "is_multimodal", False):
4783 raise ValueError(
4784 "prepare_multimodal_inputs() requires a multimodal model "
4785 "(cfg.is_multimodal must be True)"
4786 )
4787 if self.processor is None:
4788 raise ValueError(
4789 "No processor available. Load model with boot_transformers() or "
4790 "set bridge.processor = AutoProcessor.from_pretrained(...) manually."
4791 )
4792 inputs = self.processor(text=text, images=images, return_tensors="pt")
4793 return {k: v.to(self.cfg.device) if hasattr(v, "to") else v for k, v in inputs.items()}
4795 def to(self, *args, **kwargs) -> "TransformerBridge":
4796 """Move model to device and/or change dtype.
4798 Args:
4799 args: Positional arguments for nn.Module.to
4800 kwargs: Keyword arguments for nn.Module.to
4801 print_details: Whether to print details about device/dtype changes (default: True)
4803 Returns:
4804 Self for chaining
4805 """
4806 # Extract print_details if provided
4807 print_details = kwargs.pop("print_details", True)
4809 # Handle both device and dtype changes
4810 # torch.nn.Module.to() supports: to(device), to(dtype), to(device, dtype),
4811 # to(device=...), to(dtype=...), to(device=..., dtype=...)
4812 target_device, target_dtype = None, None
4814 if len(args) >= 1:
4815 first_arg = args[0]
4816 if isinstance(first_arg, (torch.device, str)):
4817 target_device = first_arg
4818 elif isinstance(first_arg, torch.dtype): 4818 ↛ 4820line 4818 didn't jump to line 4820 because the condition on line 4818 was always true
4819 target_dtype = first_arg
4820 if len(args) >= 2:
4821 second_arg = args[1]
4822 if isinstance(second_arg, torch.dtype): 4822 ↛ 4826line 4822 didn't jump to line 4826 because the condition on line 4822 was always true
4823 target_dtype = second_arg
4825 # these override positional args
4826 if "device" in kwargs: 4826 ↛ 4827line 4826 didn't jump to line 4827 because the condition on line 4826 was never true
4827 target_device = kwargs["device"]
4828 if "dtype" in kwargs:
4829 target_dtype = kwargs["dtype"]
4831 # Moving a multi-device (device_map-dispatched) model to a single device would
4832 # collapse the split and break accelerate's hook routing. Warn and drop the
4833 # device move; still honor dtype changes.
4834 if target_device is not None and getattr(self.cfg, "n_devices", 1) > 1:
4835 warnings.warn(
4836 f"TransformerBridge.to({target_device!r}) ignored: model is dispatched "
4837 f"across {self.cfg.n_devices} devices via device_map. Reload with "
4838 "device=... (and no device_map/n_devices) to move to a single device.",
4839 stacklevel=2,
4840 )
4841 target_device = None
4843 if target_device is not None:
4844 move_to_and_update_config(self, target_device, print_details)
4845 if target_dtype is not None:
4846 move_to_and_update_config(self, target_dtype, print_details)
4848 # Move the original model with all original args/kwargs (with print_details removed).
4849 # When we've nulled target_device for multi-GPU safety, strip device args so the
4850 # underlying module isn't moved either.
4851 if target_device is None and (len(args) > 0 or "device" in kwargs):
4852 kwargs.pop("device", None)
4853 # Filter positional args: drop devices/strings, keep dtypes.
4854 args = tuple(a for a in args if not isinstance(a, (torch.device, str)))
4855 self.original_model = self.original_model.to(*args, **kwargs)
4856 return self
4858 def cuda(self, device: Optional[Union[int, torch.device]] = None) -> "TransformerBridge":
4859 """Move model to CUDA.
4861 Args:
4862 device: CUDA device
4864 Returns:
4865 Self for chaining
4866 """
4867 if isinstance(device, int):
4868 return self.to(f"cuda:{device}")
4869 elif device is None:
4870 return self.to("cuda")
4871 else:
4872 return self.to(device)
4874 def cpu(self) -> "TransformerBridge":
4875 """Move model to CPU.
4877 Returns:
4878 Self for chaining
4879 """
4880 return self.to(torch.device("cpu"))
4882 def mps(self) -> "TransformerBridge":
4883 """Move model to MPS.
4885 Returns:
4886 Self for chaining
4887 """
4888 return self.to(torch.device("mps"))
4890 def train(self, mode: bool = True) -> "TransformerBridge":
4891 """Set training mode, propagating to the wrapped source model."""
4892 super().train(mode)
4893 original = getattr(self, "original_model", None)
4894 if isinstance(original, torch.nn.Module): 4894 ↛ 4896line 4894 didn't jump to line 4896 because the condition on line 4894 was always true
4895 original.train(mode)
4896 return self
4898 def _gated_hook_reason(self, hook_point_name: str) -> Optional[str]:
4899 """Return the disabled setter name if hook_point_name is gated off, else None."""
4900 if hook_point_name.endswith("attn.hook_result") and not self.cfg.use_attn_result:
4901 return "use_attn_result"
4902 if (
4903 hook_point_name.endswith(("hook_q_input", "hook_k_input", "hook_v_input"))
4904 and not self.cfg.use_split_qkv_input
4905 ):
4906 return "use_split_qkv_input"
4907 if hook_point_name.endswith("mlp_in") and not self.cfg.use_hook_mlp_in:
4908 return "use_hook_mlp_in"
4909 if hook_point_name.endswith("attn_in") and not self.cfg.use_attn_in:
4910 return "use_attn_in"
4911 return None
4913 def check_hooks_to_add(self, hook_point_name: str) -> None:
4914 """Raise a clear error if a hook is being explicitly added to a gated-off hook point.
4916 Mirrors HookedTransformer.check_hooks_to_add, but raises a ValueError
4917 naming the setter to call, instead of a bare assert. Only for explicit,
4918 user-named hook points — a filter/callable matching a gated name uses
4919 _gated_hook_reason directly and skips with a warning instead, since the
4920 filter was not necessarily targeting that name on purpose.
4921 """
4922 reason = self._gated_hook_reason(hook_point_name)
4923 if reason is not None:
4924 raise ValueError(
4925 f"Cannot add hook {hook_point_name} because {reason} is False. "
4926 f"Call set_{reason}(True) first."
4927 )
4929 def add_hook(
4930 self,
4931 name: Union[str, Callable[[str], bool]],
4932 hook_fn,
4933 dir="fwd",
4934 is_permanent=False,
4935 ):
4936 """Add a hook to a specific component or to all components matching a filter.
4938 Args:
4939 name: Either a string hook point name (e.g. "blocks.0.attn.hook_q")
4940 or a callable filter ``(str) -> bool`` that is applied to every
4941 hook point name; the hook is added to each point where the filter
4942 returns True.
4943 hook_fn: The hook function ``(activation, hook) -> activation | None``.
4944 dir: Hook direction, ``"fwd"`` or ``"bwd"``.
4945 is_permanent: If True the hook survives ``reset_hooks()`` calls.
4946 """
4947 if callable(name) and not isinstance(name, str): 4947 ↛ 4948line 4947 didn't jump to line 4948 because the condition on line 4947 was never true
4948 hook_dict = self.hook_dict
4949 seen_hooks: set[int] = set()
4950 gated_names_skipped: List[str] = []
4951 for hook_name, hook_point in hook_dict.items():
4952 if name(hook_name):
4953 hook_id = id(hook_point)
4954 if hook_id in seen_hooks:
4955 continue
4956 seen_hooks.add(hook_id)
4957 if self._gated_hook_reason(hook_name) is not None:
4958 gated_names_skipped.append(hook_name)
4959 continue
4960 hook_point.add_hook(hook_fn, dir=dir, is_permanent=is_permanent)
4961 if gated_names_skipped:
4962 warnings.warn(
4963 f"add_hook: filter matched {len(gated_names_skipped)} gated-off hook "
4964 f"name(s) that were skipped: {gated_names_skipped}. Call the relevant "
4965 "set_use_*(True) setter first to enable them.",
4966 stacklevel=2,
4967 )
4968 return
4970 component = self
4971 parts = name.split(".")
4972 for part in parts[:-1]:
4973 if hasattr(component, part): 4973 ↛ 4976line 4973 didn't jump to line 4976 because the condition on line 4973 was always true
4974 component = getattr(component, part)
4975 else:
4976 raise AttributeError(f"Component path '{'.'.join(parts[:-1])}' not found")
4977 hook_name = parts[-1]
4978 if hasattr(component, hook_name): 4978 ↛ 4988line 4978 didn't jump to line 4988 because the condition on line 4978 was always true
4979 hook_point = getattr(component, hook_name)
4980 if isinstance(hook_point, HookPoint): 4980 ↛ 4984line 4980 didn't jump to line 4984 because the condition on line 4980 was always true
4981 self.check_hooks_to_add(name)
4982 hook_point.add_hook(hook_fn, dir=dir, is_permanent=is_permanent)
4983 else:
4984 raise AttributeError(
4985 f"'{hook_name}' is not a hook point. Found object of type: {type(hook_point)} with value: {hook_point}"
4986 )
4987 else:
4988 raise AttributeError(f"Hook point '{hook_name}' not found on component")
4990 def add_perma_hook(
4991 self,
4992 name: Union[str, Callable[[str], bool]],
4993 hook_fn,
4994 dir="fwd",
4995 ) -> None:
4996 """Add a permanent hook that survives ``reset_hooks()`` calls.
4998 Convenience wrapper for ``add_hook(..., is_permanent=True)``. To remove,
4999 call ``reset_hooks(including_permanent=True)`` or remove from the
5000 underlying ``HookPoint`` directly.
5001 """
5002 self.add_hook(name, hook_fn, dir=dir, is_permanent=True)
5004 def hook_points(self) -> Iterable[HookPoint]:
5005 """All registered :class:`HookPoint` instances."""
5006 return self._hook_registry.values()
5008 def clear_contexts(self) -> None:
5009 """Clear the stored ``ctx`` on every registered hook point."""
5010 for hp in self._hook_registry.values():
5011 hp.clear_context()
5013 def remove_all_hook_fns(
5014 self,
5015 direction: Literal["fwd", "bwd", "both"] = "both",
5016 including_permanent: bool = False,
5017 level: Optional[int] = None,
5018 ) -> None:
5019 """Remove hook functions from every registered hook point."""
5020 for hp in self._hook_registry.values():
5021 hp.remove_hooks(dir=direction, including_permanent=including_permanent, level=level)
5023 def reset_hooks(
5024 self,
5025 clear_contexts: bool = True,
5026 direction: Literal["fwd", "bwd", "both"] = "both",
5027 including_permanent: bool = False,
5028 level: Optional[int] = None,
5029 ) -> None:
5030 """Remove hooks from the model; mirrors ``HookedRootModule.reset_hooks``.
5032 Clears through the hook registry (which holds hook points the component
5033 walk cannot reach, e.g. alias-registered points) and, on a full reset,
5034 additionally walks the component tree — dev's registry is not asserted
5035 canonical, so both passes run belt-and-suspenders.
5036 """
5037 if clear_contexts: 5037 ↛ 5039line 5037 didn't jump to line 5039 because the condition on line 5037 was always true
5038 self.clear_contexts()
5039 self.remove_all_hook_fns(direction, including_permanent=including_permanent, level=level)
5041 if direction == "both" and level is None: 5041 ↛ exitline 5041 didn't return from function 'reset_hooks' because the condition on line 5041 was always true
5043 def remove_hooks_recursive(module):
5044 if isinstance(module, GeneralizedComponent):
5045 module.remove_hooks()
5046 for child in module.children():
5047 remove_hooks_recursive(child)
5049 remove_hooks_recursive(self)
5051 def hooks(self, fwd_hooks=[], bwd_hooks=[], reset_hooks_end=True, clear_contexts=False):
5052 """Context manager for temporarily adding hooks.
5054 Args:
5055 fwd_hooks: List of (hook_name, hook_fn) tuples for forward hooks
5056 bwd_hooks: List of (hook_name, hook_fn) tuples for backward hooks
5057 reset_hooks_end: If True, removes the hooks this context added when it exits.
5058 Hooks the caller added beforehand are left alone either way.
5059 clear_contexts: If True, clears the hook contexts of the touched hook points
5060 when the hooks are removed
5062 Example:
5063 with model.hooks(fwd_hooks=[("hook_embed", my_hook)]):
5064 output = model("Hello world")
5065 """
5067 @contextmanager
5068 def _hooks_context():
5069 added_hooks: List[Tuple[HookPoint, Literal["fwd", "bwd"]]] = []
5070 # Claimed here for the closures below, but only committed to self inside the try
5071 # that decrements it, so a raise in between can't leave the counter incremented.
5072 context_level = self.context_level + 1
5074 def add_hook_to_point(
5075 hook_point: HookPoint,
5076 hook_fn: Callable,
5077 name: str,
5078 dir: Literal["fwd", "bwd"] = "fwd",
5079 *,
5080 is_explicit: bool = True,
5081 ):
5082 if is_explicit: 5082 ↛ 5084line 5082 didn't jump to line 5084 because the condition on line 5082 was always true
5083 self.check_hooks_to_add(name)
5084 elif self._gated_hook_reason(name) is not None:
5085 warnings.warn(
5086 f"hooks(): filter matched gated-off hook name '{name}', skipped. "
5087 "Call the relevant set_use_*(True) setter first to enable it.",
5088 stacklevel=2,
5089 )
5090 return
5091 if self.compatibility_mode and name != hook_point.name: 5091 ↛ 5092line 5091 didn't jump to line 5092 because the condition on line 5091 was never true
5092 alias_names_list: list[str] = []
5093 if hook_point.name is not None:
5094 alias_names_list.append(hook_point.name)
5095 alias_names_list.append(name)
5096 hook_point.add_hook(
5097 hook_fn, dir=dir, alias_names=alias_names_list, level=context_level
5098 )
5099 else:
5100 hook_point.add_hook(hook_fn, dir=dir, level=context_level)
5101 added_hooks.append((hook_point, dir))
5103 def apply_hooks(hooks: List[Tuple[Union[str, Callable], Callable]], is_fwd: bool):
5104 direction: Literal["fwd", "bwd"] = "fwd" if is_fwd else "bwd"
5105 aliases = build_alias_to_canonical_map(self.hook_dict)
5106 for hook_name_or_filter, hook_fn in hooks:
5107 if isinstance(hook_name_or_filter, str): 5107 ↛ 5121line 5107 didn't jump to line 5121 because the condition on line 5107 was always true
5108 hook_dict = self.hook_dict
5109 actual_hook_name = hook_name_or_filter
5110 if hook_name_or_filter in aliases:
5111 actual_hook_name = aliases[hook_name_or_filter]
5112 if actual_hook_name in hook_dict: 5112 ↛ 5106line 5112 didn't jump to line 5106 because the condition on line 5112 was always true
5113 add_hook_to_point(
5114 hook_dict[actual_hook_name],
5115 hook_fn,
5116 actual_hook_name,
5117 direction,
5118 is_explicit=True,
5119 )
5120 else:
5121 hook_dict = self.hook_dict
5122 seen_hooks = set()
5123 for name, hook_point in hook_dict.items():
5124 if hook_name_or_filter(name):
5125 hook_id = id(hook_point)
5126 if hook_id in seen_hooks:
5127 continue
5128 seen_hooks.add(hook_id)
5129 hook_name_to_use = hook_point.name if hook_point.name else name
5130 add_hook_to_point(
5131 hook_point,
5132 hook_fn,
5133 hook_name_to_use,
5134 direction,
5135 is_explicit=False,
5136 )
5138 try:
5139 self.context_level = context_level
5140 apply_hooks(fwd_hooks, True)
5141 apply_hooks(bwd_hooks, False)
5142 yield self
5143 finally:
5144 if reset_hooks_end: 5144 ↛ 5151line 5144 didn't jump to line 5151 because the condition on line 5144 was always true
5145 for hook_point, direction in added_hooks:
5146 # `level` keeps this to the hooks added above — hooks the caller
5147 # attached before the context survive.
5148 hook_point.remove_hooks(dir=direction, level=context_level)
5149 if clear_contexts: 5149 ↛ 5150line 5149 didn't jump to line 5150 because the condition on line 5149 was never true
5150 hook_point.clear_context()
5151 self.context_level -= 1
5153 return _hooks_context()
5155 def set_use_attn_result(self, use_attn_result: bool):
5156 """Toggle whether to explicitly calculate and expose the result for each attention head.
5158 Useful for interpretability but can easily burn through GPU memory.
5159 """
5160 if use_attn_result:
5161 self._validate_attention_fork_supported("use_attn_result")
5162 self.cfg._set_bridge_managed_hook_flag("use_attn_result", use_attn_result)
5163 self._propagate_attention_flag("use_attn_result", use_attn_result)
5165 def set_use_split_qkv_input(self, use_split_qkv_input: bool):
5166 """Toggle independent residual copies for Q/K/V so each path can be patched alone.
5168 Mutually exclusive with `use_attn_in` — set that flag off first if it's on.
5169 """
5170 if use_split_qkv_input:
5171 if bool(getattr(self.cfg, "use_attn_in", False)):
5172 raise ValueError(
5173 "use_split_qkv_input and use_attn_in are mutually exclusive. "
5174 "Call set_use_attn_in(False) before enabling use_split_qkv_input."
5175 )
5176 self._validate_attention_fork_supported("use_split_qkv_input")
5177 self.cfg._set_bridge_managed_hook_flag("use_split_qkv_input", use_split_qkv_input)
5178 self._propagate_attention_flag("use_split_qkv_input", use_split_qkv_input)
5180 def set_use_attn_in(self, use_attn_in: bool):
5181 """Toggle a single 4D residual copy feeding all three Q/K/V projections.
5183 Mutually exclusive with `use_split_qkv_input` — set that flag off first
5184 if it's on. When on, `hook_attn_in` fires at
5185 `[batch, pos, n_heads, d_model]`, enabling coarse-grained interventions
5186 on the residual-stream copy shared across Q/K/V.
5187 """
5188 if use_attn_in:
5189 if bool(getattr(self.cfg, "use_split_qkv_input", False)):
5190 raise ValueError(
5191 "use_attn_in and use_split_qkv_input are mutually exclusive. "
5192 "Call set_use_split_qkv_input(False) before enabling use_attn_in."
5193 )
5194 self._validate_attention_fork_supported("use_attn_in")
5195 self.cfg._set_bridge_managed_hook_flag("use_attn_in", use_attn_in)
5196 self._propagate_attention_flag("use_attn_in", use_attn_in)
5198 def set_use_hook_mlp_in(self, use_hook_mlp_in: bool) -> None:
5199 """Toggle the ``hook_mlp_in`` HookPoint (the MLP-branch entry: pre-ln2, or
5200 the MLP input on post-norm blocks), matching legacy semantics.
5202 See :py:meth:`HookedTransformer.set_use_hook_mlp_in`.
5203 """
5204 self.cfg._set_bridge_managed_hook_flag("use_hook_mlp_in", use_hook_mlp_in)
5205 if not hasattr(self, "blocks"): 5205 ↛ 5206line 5205 didn't jump to line 5206 because the condition on line 5205 was never true
5206 return
5207 for block in self.blocks:
5208 block_cfg = getattr(block, "config", None)
5209 if block_cfg is not None and block_cfg is not self.cfg:
5210 try:
5211 self._write_propagated_hook_flag(block_cfg, "use_hook_mlp_in", use_hook_mlp_in)
5212 except (AttributeError, TypeError):
5213 pass
5214 block._use_hook_mlp_in = use_hook_mlp_in
5216 @staticmethod
5217 def _write_propagated_hook_flag(config: Any, flag_name: str, value: bool) -> None:
5218 """Write a cloned config flag without dispatching through its live Bridge."""
5219 if isinstance(config, TransformerBridgeConfig): 5219 ↛ 5222line 5219 didn't jump to line 5222 because the condition on line 5219 was always true
5220 config._set_bridge_managed_hook_flag(flag_name, value)
5221 else:
5222 object.__setattr__(config, flag_name, value)
5224 def _propagate_attention_flag(self, flag_name: str, value: bool) -> None:
5225 """Mirror `bridge.cfg.<flag>` onto every block's attention config.
5227 Some adapters (Llama family) deep-copy the block template during
5228 `setup_blocks_bridge`, cloning the attention bridge's config along
5229 with it. Others (Pythia, GPT-2) override `__deepcopy__` to share the
5230 config. Setting the flag only on `self.cfg` silently misses the
5231 cloned-config case. Propagating explicitly keeps both patterns
5232 honest — a no-op when configs are shared, a correctness fix when
5233 they aren't.
5234 """
5235 if not hasattr(self, "blocks"): 5235 ↛ 5236line 5235 didn't jump to line 5236 because the condition on line 5235 was never true
5236 return
5237 for block in self.blocks:
5238 attn = block._modules.get("attn") if hasattr(block, "_modules") else None
5239 if attn is None: 5239 ↛ 5240line 5239 didn't jump to line 5240 because the condition on line 5239 was never true
5240 continue
5241 attn_cfg = getattr(attn, "config", None)
5242 if attn_cfg is not None and attn_cfg is not self.cfg:
5243 try:
5244 self._write_propagated_hook_flag(attn_cfg, flag_name, value)
5245 except (AttributeError, TypeError):
5246 # Some config-like objects reject attributes even when
5247 # bypassing their custom __setattr__ implementation.
5248 pass
5250 def _validate_attention_fork_supported(self, flag_name: str) -> None:
5251 """Raise / warn if the model can't honor a fine-grained attention flag.
5253 The post-ln1 fork path lives on JointQKVAttentionBridge and
5254 PositionEmbeddingsAttentionBridge. Plain AttentionBridge delegates to
5255 HF and exposes no fork point; we raise rather than setting the flag
5256 silently. For hybrid models (some attention layers, some not), we warn
5257 and list which layers will honor the flag.
5258 """
5259 # Deferred imports: tight circular dependency with bridge setup.
5260 from transformer_lens.model_bridge.generalized_components.joint_qkv_attention import (
5261 JointQKVAttentionBridge,
5262 )
5263 from transformer_lens.model_bridge.generalized_components.position_embeddings_attention import (
5264 PositionEmbeddingsAttentionBridge,
5265 )
5267 if not hasattr(self, "blocks"): 5267 ↛ 5268line 5267 didn't jump to line 5268 because the condition on line 5267 was never true
5268 raise NotImplementedError(
5269 f"{flag_name}: this bridge has no `blocks` attribute, so no "
5270 "attention bridges to apply the flag to."
5271 )
5272 supported_classes = (JointQKVAttentionBridge, PositionEmbeddingsAttentionBridge)
5273 supporting_layers: list[int] = []
5274 attn_classes: set[str] = set()
5275 total_with_attn = 0
5276 for idx, block in enumerate(self.blocks):
5277 attn = block._modules.get("attn") if hasattr(block, "_modules") else None
5278 if attn is None: 5278 ↛ 5279line 5278 didn't jump to line 5279 because the condition on line 5278 was never true
5279 continue
5280 total_with_attn += 1
5281 attn_classes.add(type(attn).__name__)
5282 supports_flag = (
5283 bool(getattr(attn, "supports_attn_result", False))
5284 if flag_name == "use_attn_result"
5285 else isinstance(attn, supported_classes)
5286 )
5287 if supports_flag:
5288 supporting_layers.append(idx)
5289 if total_with_attn == 0: 5289 ↛ 5290line 5289 didn't jump to line 5290 because the condition on line 5289 was never true
5290 raise NotImplementedError(f"{flag_name}: no attention bridges found on self.blocks.")
5291 if not supporting_layers:
5292 if flag_name == "use_attn_result":
5293 capability_detail = "Per-head result computation is unavailable."
5294 else:
5295 capability_detail = "No hook point is available before the Q/K/V projections."
5296 raise NotImplementedError(
5297 f"{flag_name}: none of this model's attention bridges support "
5298 "the requested fine-grained attention hook. Found attention classes: "
5299 f"{sorted(attn_classes)}. Supported classes: "
5300 f"{[c.__name__ for c in supported_classes]}. {capability_detail}"
5301 )
5302 if len(supporting_layers) < total_with_attn: 5302 ↛ 5303line 5302 didn't jump to line 5303 because the condition on line 5302 was never true
5303 skipped = total_with_attn - len(supporting_layers)
5304 warnings.warn(
5305 f"{flag_name}: {skipped} of {total_with_attn} attention layers "
5306 "use an attention-bridge class that cannot honor this flag "
5307 f"(attention classes present: {sorted(attn_classes)}). "
5308 f"The flag will affect layers: {supporting_layers}.",
5309 stacklevel=3,
5310 )
5312 def _is_valid_bridge_path(self, hf_path: str) -> bool:
5313 """Check if a HuggingFace path corresponds to a valid bridge component.
5315 This validates that the path follows the bridge component structure and doesn't
5316 contain nested HuggingFace components that should have been wrapped.
5318 Args:
5319 hf_path: HuggingFace path after removing _original_component
5321 Returns:
5322 True if the path is valid, False if it contains nested HF components
5323 """
5324 # Split the path into parts
5325 parts = hf_path.split(".")
5327 # Get the component mapping for validation
5328 component_mapping = self.adapter.component_mapping
5329 if not component_mapping: 5329 ↛ 5330line 5329 didn't jump to line 5330 because the condition on line 5329 was never true
5330 return True # If no mapping, accept all keys
5332 # Walk through the path and check if each level is a registered bridge component
5333 # For example, transformer.h.0.mlp.in.weight should be valid
5334 # but transformer.h.0.mlp.c_fc.weight should be invalid (c_fc is nested HF component)
5336 # Start from the root
5337 current_component = None
5338 idx = 0
5340 # Find which top-level component this belongs to
5341 for tl_name, component in component_mapping.items():
5342 if component.name and hf_path.startswith(component.name + "."):
5343 current_component = component
5344 # Skip past the HF prefix
5345 remaining_path = hf_path[len(component.name) + 1 :]
5346 parts = remaining_path.split(".")
5347 idx = 0
5348 break
5350 if current_component is None:
5351 return True # Path doesn't match any component, let it through
5353 # Special handling for blocks
5354 if hasattr(current_component, "is_list_item") and current_component.is_list_item:
5355 # Skip the layer index
5356 if idx < len(parts) and parts[idx].isdigit(): 5356 ↛ 5360line 5356 didn't jump to line 5360 because the condition on line 5356 was always true
5357 idx += 1
5359 # Now validate the rest of the path against submodules
5360 while idx < len(parts): 5360 ↛ 5387line 5360 didn't jump to line 5387 because the condition on line 5360 was always true
5361 part = parts[idx]
5363 # If we hit 'weight' or 'bias', we're at a parameter - this is valid
5364 if part in ("weight", "bias"):
5365 return True
5367 # Check if this part is a registered submodule
5368 if hasattr(current_component, "submodules") and current_component.submodules:
5369 if part in current_component.submodules:
5370 current_component = current_component.submodules[part]
5371 idx += 1
5372 continue
5373 else:
5374 # This part is not a registered bridge component
5375 # It's likely a nested HF component (like c_fc, c_proj, c_attn)
5376 return False
5377 else:
5378 # No submodules to check, but not at a parameter yet
5379 # Check if next is weight/bias
5380 if idx + 1 < len(parts) and parts[idx + 1] in ("weight", "bias"):
5381 return True
5382 # Otherwise this is likely a nested HF component
5383 return False
5385 idx += 1
5387 return True
5389 def _normalize_bridge_key_to_hf(self, key: str) -> str:
5390 """Normalize a key that uses bridge attribute names to use HF module names.
5392 PyTorch's state_dict uses the Python attribute names (e.g., 'ln1')
5393 but the conversion logic expects HF module names (e.g., 'ln_1'). This
5394 function only replaces non-nested component names, leaving bridge
5395 subcomponents (like 'in', 'out', 'q', 'k', 'v') unchanged since they're
5396 handled by the component structure.
5398 Args:
5399 key: Key that may use bridge attribute names
5401 Returns:
5402 Key with attribute names replaced by module names where needed
5403 """
5404 component_mapping = self.adapter.component_mapping
5405 if not component_mapping: 5405 ↛ 5406line 5405 didn't jump to line 5406 because the condition on line 5405 was never true
5406 return key
5408 # Build a mapping of only the direct module attribute names to HF names
5409 # We only care about top-level and block-level component names, NOT subcomponents
5410 attr_to_hf = {}
5412 # Map top-level components
5413 block_list_names = {"blocks", "L_blocks", "H_blocks", "encoder_blocks", "decoder_blocks"}
5414 for tl_name, component in component_mapping.items():
5415 if component.name and tl_name not in block_list_names:
5416 # Skip if TL name is already a segment of its HF path (avoids doubling).
5417 if tl_name != component.name and tl_name not in component.name.split("."):
5418 attr_to_hf[tl_name] = component.name
5420 # Map block-level components (ln1, ln2, attn, mlp) for all block lists
5421 for bl_name in block_list_names:
5422 blocks_component = component_mapping.get(bl_name)
5423 if blocks_component and hasattr(blocks_component, "submodules"):
5424 for tl_subname, subcomponent in blocks_component.submodules.items():
5425 if subcomponent.name:
5426 # Only map if the names differ (e.g., ln1 -> ln_1, but attn -> attn)
5427 if tl_subname != subcomponent.name:
5428 attr_to_hf[tl_subname] = subcomponent.name
5430 # Replace only these specific attribute names in the key
5431 # We need to be careful to only replace whole path components, not substrings
5432 parts = key.split(".")
5433 result_parts = []
5435 for part in parts:
5436 if part in attr_to_hf:
5437 result_parts.append(attr_to_hf[part])
5438 else:
5439 result_parts.append(part)
5441 return ".".join(result_parts)
5443 def state_dict(self, destination=None, prefix="", keep_vars=False):
5444 """Get state dict with TransformerLens format keys.
5446 Converts HuggingFace format keys to TransformerLens format and filters out
5447 _original_component references and nested HuggingFace components.
5449 A direct no-argument call returns a clean state dict with bridge component
5450 paths converted to TL format. Calls that supply ``destination`` or
5451 ``prefix`` use standard ``nn.Module`` recursive semantics so a Bridge can
5452 compose inside a parent module.
5454 Args:
5455 destination: Optional dict to store state dict in
5456 prefix: Optional prefix to add to all keys
5457 keep_vars: Whether to keep variables as Variables instead of tensors
5459 Returns:
5460 Direct calls return TransformerLens-format keys; recursive calls
5461 return the supplied destination with standard module-tree keys.
5462 """
5463 if destination is not None or prefix:
5464 return super().state_dict(
5465 destination=destination,
5466 prefix=prefix,
5467 keep_vars=keep_vars,
5468 )
5470 raw_state_dict = self.original_model.state_dict(keep_vars=keep_vars)
5472 # Clean _original_component references and convert to TL format
5473 # Also filter out nested HuggingFace components that are wrapped by bridge components
5474 tl_state_dict = {}
5476 for key, value in raw_state_dict.items():
5477 # Skip _original_component keys
5478 if key == "_original_component" or key.startswith("_original_component."): 5478 ↛ 5479line 5478 didn't jump to line 5479 because the condition on line 5478 was never true
5479 continue
5481 # Remove all _original_component from the key
5482 clean_key = key.replace("._original_component", "")
5484 # Check if this is a valid bridge path (not a nested HF component)
5485 if not self._is_valid_bridge_path(clean_key):
5486 continue
5488 # Normalize bridge component names to HF names for conversion
5489 # (e.g., 'ln1' -> 'ln_1', 'mlp.in' -> 'mlp.c_fc')
5490 hf_key = self._normalize_bridge_key_to_hf(clean_key)
5492 # Convert to TL format - this uses the adapter's component_mapping
5493 tl_key = self.adapter.convert_hf_key_to_tl_key(hf_key)
5495 # Only add if we haven't seen this TL key yet (handles duplicates)
5496 if tl_key not in tl_state_dict:
5497 tl_state_dict[tl_key] = value
5499 return tl_state_dict
5501 def _tl_key_to_actual_keys(self) -> dict[str, list[str]]:
5502 """Inverse of the renaming state_dict() applies: map each TL-format key
5503 back to every raw parameter/buffer path that represents it.
5505 Mirrors the filtering and key-conversion in state_dict() exactly, except
5506 it keeps every raw key for a given TL key instead of only the first-seen
5507 one. Bridge components frequently expose the same underlying parameter
5508 through more than one attribute path (e.g. GPT-2's split q/k/v weights
5509 are views into the wrapped module's combined c_attn weight, reachable
5510 both via a block-level shortcut and via the nested _original_component
5511 chain) - all of those aliases must be written for the round trip to
5512 actually change what forward() reads, not just what state_dict() shows.
5513 """
5514 mapping: dict[str, list[str]] = {}
5515 for actual_key in self.original_model.state_dict():
5516 if actual_key == "_original_component" or actual_key.startswith("_original_component."): 5516 ↛ 5517line 5516 didn't jump to line 5517 because the condition on line 5516 was never true
5517 continue
5518 clean_key = actual_key.replace("._original_component", "")
5519 if not self._is_valid_bridge_path(clean_key):
5520 continue
5521 hf_key = self._normalize_bridge_key_to_hf(clean_key)
5522 tl_key = self.adapter.convert_hf_key_to_tl_key(hf_key)
5523 mapping.setdefault(tl_key, []).append(actual_key)
5524 return mapping
5526 def load_state_dict(self, state_dict, strict=True, assign=False):
5527 """Load state dict into the model, handling both clean keys and original keys with _original_component references.
5529 Accepts three key formats: TL-format keys as emitted by state_dict()
5530 (e.g. "blocks.0.attn.q.weight"), raw native parameter paths (e.g. for
5531 ``boot_native`` / tracr-style loading), and raw paths with
5532 "_original_component" segments stripped.
5534 Args:
5535 state_dict: Dictionary containing a whole state of the module
5536 strict: Whether to strictly enforce that the keys in state_dict match the keys returned by this module's state_dict() function
5537 assign: Whether to assign items in the state dictionary to their corresponding keys in the module instead of copying them
5539 Returns:
5540 NamedTuple with missing_keys and unexpected_keys fields
5541 """
5542 current_state_dict = self.original_model.state_dict()
5543 clean_to_actual = {}
5544 for actual_key in current_state_dict.keys():
5545 if actual_key != "_original_component": 5545 ↛ 5544line 5545 didn't jump to line 5544 because the condition on line 5545 was always true
5546 clean_to_actual[actual_key.replace("._original_component", "")] = actual_key
5548 tl_to_actual = self._tl_key_to_actual_keys()
5550 mapped_state_dict = {}
5551 unexpected_keys = []
5552 for input_key, value in state_dict.items():
5553 if input_key in current_state_dict:
5554 mapped_state_dict[input_key] = value
5555 elif input_key in clean_to_actual:
5556 mapped_state_dict[clean_to_actual[input_key]] = value
5557 elif input_key in tl_to_actual:
5558 for actual_key in tl_to_actual[input_key]:
5559 mapped_state_dict[actual_key] = value
5560 else:
5561 unexpected_keys.append(input_key)
5563 # A TL key's actual-key aliases share the same underlying storage (see
5564 # _tl_key_to_actual_keys), so writing any one of them already updates
5565 # what forward() reads for all of them. Treat the group as satisfied
5566 # if any alias was written -- e.g. a caller supplying clean/raw keys
5567 # (the branch above maps each clean key to exactly one actual key)
5568 # shouldn't have the *other*, unwritten aliases reported as missing.
5569 missing_keys = sorted(
5570 actual_key
5571 for actual_keys in tl_to_actual.values()
5572 if not any(k in mapped_state_dict for k in actual_keys)
5573 for actual_key in actual_keys
5574 )
5576 if strict and (missing_keys or unexpected_keys):
5577 error_msgs = []
5578 if unexpected_keys:
5579 error_msgs.append(
5580 "Unexpected key(s) in state_dict: "
5581 + ", ".join(f'"{k}"' for k in sorted(unexpected_keys))
5582 )
5583 if missing_keys:
5584 error_msgs.append(
5585 "Missing key(s) in state_dict: " + ", ".join(f'"{k}"' for k in missing_keys)
5586 )
5587 raise RuntimeError(
5588 "Error(s) in loading state_dict for {}:\n\t{}".format(
5589 type(self.original_model).__name__, "\n\t".join(error_msgs)
5590 )
5591 )
5593 result = self.original_model.load_state_dict(mapped_state_dict, strict=False, assign=assign)
5594 if assign:
5595 refresh_container_state_owners(self)
5596 return type(result)(missing_keys=missing_keys, unexpected_keys=unexpected_keys)
5598 def get_params(self):
5599 """Access to model parameters in the format expected by SVDInterpreter.
5601 For missing weights, returns zero tensors of appropriate shape instead of raising exceptions.
5602 This ensures compatibility across different model architectures.
5604 Returns:
5605 dict: Dictionary of parameter tensors with TransformerLens naming convention
5607 Raises:
5608 ValueError: If configuration is inconsistent (e.g., cfg.n_layers != len(blocks))
5609 """
5610 return get_bridge_params(self)
5612 # NOTE: list_supported_models and check_model_support are attached to this class
5613 # dynamically by transformer_lens.model_bridge.sources.transformers module.
5614 # These are HuggingFace-specific methods that belong in the transformers source module.