Coverage for transformer_lens/model_bridge/transformer_bridge.py: 88%
1860 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-10-10 05:10 +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 typing import (
13 TYPE_CHECKING,
14 Any,
15 Callable,
16 ClassVar,
17 Dict,
18 Iterator,
19 List,
20 Literal,
21 Optional,
22 Sequence,
23 Tuple,
24 Union,
25 cast,
26)
28import einops
29import numpy as np
30import torch
31import tqdm
32from torch import nn
33from torch.nn import functional as F
34from transformers.tokenization_utils_base import PreTrainedTokenizerBase
36from transformer_lens import utilities as utils
37from transformer_lens.ActivationCache import ActivationCache
38from transformer_lens.config import TransformerBridgeConfig
39from transformer_lens.FactoredMatrix import FactoredMatrix
40from transformer_lens.hook_points import HookIntrospectionMixin, HookPoint
41from transformer_lens.model_bridge.architecture_adapter import ArchitectureAdapter
42from transformer_lens.model_bridge.bridge_core import (
43 _BLOCK_LIST_ATTRS,
44 _SELF_ATTENTION_NAMES,
45 BridgeCore,
46)
47from transformer_lens.model_bridge.component_setup import (
48 refresh_container_state_owners,
49 set_original_components,
50)
51from transformer_lens.model_bridge.composition_scores import CompositionScores
52from transformer_lens.model_bridge.driver_protocol import (
53 TensorLike,
54 to_torch,
55 validate_driver,
56)
57from transformer_lens.model_bridge.exceptions import StopAtLayerException
58from transformer_lens.model_bridge.generalized_components.base import (
59 GeneralizedComponent,
60)
61from transformer_lens.model_bridge.generalized_components.block import (
62 _BLOCK_INTERNAL_MODULES,
63 _NORM_PREFIXES,
64 _VARIANT_SUBMODULE_SET,
65 VARIANT_SUBMODULE_NAMES,
66)
67from transformer_lens.model_bridge.generalized_components.moe import (
68 fold_scale_into_moe_block,
69 has_batched_experts,
70)
71from transformer_lens.model_bridge.get_params_util import get_bridge_params
72from transformer_lens.utilities.activation_functions import softcap_enabled
73from transformer_lens.utilities.devices import move_to_and_update_config
74from transformer_lens.utilities.position_ids import accepts_mask_derived_position_ids
75from transformer_lens.utilities.quantization import require_readable_weight
77if TYPE_CHECKING:
78 pass
80_BLOCK_PATTERN = re.compile("blocks\\.(\\d+)")
83def _resolve_attr_path(obj: nn.Module, attr_path: str) -> Optional[torch.Tensor]:
84 """Walk a dot-separated attribute path and return the final tensor (None if bias-free)."""
85 result = obj
86 for attr in attr_path.split("."):
87 result = getattr(result, attr)
88 return cast(torch.Tensor, result)
91def _storage_group_keys(state_dict: dict[str, torch.Tensor]) -> set[str]:
92 """Keys whose current tensor shares underlying storage with another key's.
94 Covers every direction of the desync #1637 reported: a partial view onto
95 a combined weight (e.g. a split QKV/gate-up component's ``torch.tensor_split``
96 view into ``c_attn``/``gate_up_proj``), the combined weight itself (writing
97 it directly would silently orphan the views that share its storage), and a
98 same-size tied pair (e.g. tied embed/unembed weights, #1725) -- none of
99 which reliably show up via ``Tensor._base``, since wrapping in
100 ``nn.Parameter`` doesn't preserve that tracking here.
102 Meta tensors are excluded from the comparison: every meta tensor reports
103 ``untyped_storage().data_ptr() == 0`` (it has no real backing memory), so
104 comparing meta tensors by data pointer would spuriously group every
105 unrelated offloaded parameter in the model together. A key whose current
106 target is meta falls through to ordinary ``assign=True`` handling instead
107 (the standard way to materialize a meta tensor from a real one).
108 """
109 by_storage: dict[int, list[str]] = {}
110 for key, tensor in state_dict.items():
111 if tensor.is_meta:
112 continue
113 by_storage.setdefault(tensor.untyped_storage().data_ptr(), []).append(key)
114 return {key for keys in by_storage.values() if len(keys) > 1 for key in keys}
117class TransformerBridge(BridgeCore, HookIntrospectionMixin, nn.Module):
118 """Torch-backed bridge: HF, vLLM-via-torch, anything that wraps an ``nn.Module``.
120 Provides a standardized interface to access components of a transformer
121 model, regardless of the underlying architecture. It uses an architecture adapter
122 to map between the TransformerLens and HuggingFace model structures.
124 Stateless reparametrization is unsupported
125 ------------------------------------------
127 ``torch.func.functional_call`` (and ``torch.nn.utils.stateless``) fails to
128 restore parameters through this tree: each replaced component is registered
129 both inside the wrapped HF model and as a bridge submodule, and torch's
130 tied-weight handling double-swaps the shared slot, leaving the override
131 installed. For temporary weight edits use
132 :func:`transformer_lens.utilities.temporarily_swap_parameter`.
134 Tokenization notes
135 ------------------
137 :meth:`to_tokens`, :meth:`to_str_tokens`, :meth:`get_token_position`,
138 :meth:`forward` (string input), and :meth:`generate` accept ``prepend_bos``
139 to control BOS prepending. Resolution: explicit arg →
140 ``cfg.default_prepend_bos`` (defaults ``True``, even for non-BOS-trained
141 models — attention heads tend to use position 0 as a resting state).
142 **Pass ``prepend_bos=False`` when tokenizing a fragment of a larger
143 prompt** — off-by-one position errors usually trace back here.
145 Reconciliation with ``cfg.tokenizer_prepends_bos`` (tokenizers that add
146 BOS automatically) is handled internally — pass the value you want;
147 the bridge adds or strips manually as needed. When
148 ``cfg.tokenizer_appends_eos=True`` (OLMo, Apertus, etc.),
149 :meth:`to_tokens` also strips trailing EOS tokens so the model receives
150 a continuation rather than a terminated sequence; this path is
151 bridge-specific.
153 BPE/SentencePiece tokenizers treat ``"hello"``, ``" hello"``, and
154 ``"Hello"`` as distinct tokens. Concatenated prompts may not tokenize
155 as the sum of parts — inspect with :meth:`to_str_tokens` when in doubt.
156 """
158 # hook_aliases inherited from BridgeCore
160 def __init__(
161 self,
162 model: nn.Module,
163 adapter: ArchitectureAdapter,
164 tokenizer: Any,
165 *,
166 driver: Any = None,
167 ):
168 """Initialize the bridge.
170 Args:
171 model: The model to bridge (must be a PyTorch nn.Module or PreTrainedModel)
172 adapter: The architecture adapter to use
173 tokenizer: The tokenizer to use (required)
174 driver: Optional pre-built :class:`Driver`. Sources that construct
175 exotic drivers (vLLM, Inspect) pass them here. When ``None``,
176 a :class:`TransformersDriver` is built from the supplied
177 ``model``/``adapter``/``tokenizer`` — kept for backward
178 compatibility with direct ``TransformerBridge(...)`` callers.
179 """
180 nn.Module.__init__(self)
181 self._n_params_total = sum(parameter.numel() for parameter in model.parameters())
182 self.__dict__["original_model"] = model
183 # Production sources construct their Driver and pass it via ``driver=``.
184 # The fallback covers tests / direct callers with a hand-rolled triple.
185 if driver is None:
186 from transformer_lens.model_bridge.sources.transformers_driver import (
187 TransformersDriver,
188 )
190 driver = TransformersDriver(model, adapter, tokenizer)
191 BridgeCore.__init__(self, adapter, tokenizer, driver)
192 # real_components maps TL keys to (remote_path, actual_instance) tuples;
193 # for list components, actual_instance is a list of instances.
194 self.real_components: Dict[str, tuple] = {}
195 if not hasattr(self.cfg, "device") or self.cfg.device is None: 195 ↛ 196line 195 didn't jump to line 196 because the condition on line 195 was never true
196 try:
197 self.cfg.device = str(next(self.original_model.parameters()).device)
198 except StopIteration:
199 self.cfg.device = "cpu"
200 set_original_components(self, self.adapter, self.__dict__["original_model"])
201 self._initialize_hook_registry()
202 self._register_aliases()
203 self._register_all_aliases_recursive()
204 # Re-scan after alias registration so alias names (hook_resid_pre, …)
205 # join the registry alongside their canonical targets. Shared HookPoint
206 # instances, so no double-firing.
207 self._scan_existing_hooks(self, "")
208 self._setup_hook_compatibility()
209 # Backfill supported_hook_points = registry − non_fireable. Whitelist-
210 # semantic drivers (Inspect) declared their own non-empty set and skip.
211 if not self._driver.supported_hook_points:
212 self._driver.supported_hook_points = (
213 frozenset(self._hook_registry) - self._driver.non_fireable_hook_points
214 )
215 # Fail fast on a misshapen driver here, not at first capture.
216 validate_driver(self._driver, after_bridge_construction=True)
217 self.processor = None
218 # Bridge wrappers are inserted into the HF module tree after
219 # from_pretrained's eval(), and nn.Module defaults to training=True —
220 # without re-syncing, reconstruction paths apply dropout at inference.
221 # train() recurses, so this stamps the wrappers with the model's mode.
222 model.train(model.training)
223 self.train(model.training)
224 self.cfg._bind_bridge(self)
226 def __setstate__(self, state: dict[str, Any]) -> None:
227 """Restore runtime config routing after deepcopy or deserialization."""
228 super().__setstate__(state)
229 self.cfg._bind_bridge(self)
231 # boot_transformers / list_supported_models / check_model_support are
232 # attached by sources.transformers.__init__ (setattr on this class) so the
233 # source package owns its own boot entry point.
234 boot_transformers: ClassVar[Callable[..., "TransformerBridge"]]
236 @property
237 def original_model(self) -> nn.Module:
238 """The wrapped ``nn.Module``. Raises :class:`AttributeError` for
239 non-torch drivers (vLLM, Inspect) that don't expose a local module."""
240 driver = getattr(self, "_driver", None)
241 if driver is None:
242 # Bridges assembled without __init__ (object.__new__ scaffolds) keep the
243 # module in the __dict__ mirror the setter maintains.
244 model = self.__dict__.get("original_model")
245 if model is None:
246 raise AttributeError(f"'{type(self).__name__}' has no driver and no original_model")
247 return model
248 underlying = getattr(driver, "underlying_model", None)
249 if underlying is None:
250 raise AttributeError(
251 f"{type(driver).__name__} does not expose an nn.Module — "
252 "non-torch drivers (vLLM, Inspect) operate without a local module."
253 )
254 return underlying
256 @original_model.setter
257 def original_model(self, value: nn.Module) -> None:
258 """Used by weight-processing paths that move the model across devices."""
259 self.__dict__["original_model"] = value
260 # Sync via the driver's public API; non-torch drivers don't implement it.
261 setter = getattr(getattr(self, "_driver", None), "set_underlying_model", None)
262 if callable(setter):
263 setter(value)
265 def init_weights(self) -> None:
266 """Reinitialize a TL-native model in place using the bridge config."""
267 from transformer_lens.model_bridge.sources.native.init import (
268 initialize_native_model,
269 )
270 from transformer_lens.model_bridge.sources.native.model import NativeModel
272 model = self.original_model
273 if not isinstance(model, NativeModel):
274 raise RuntimeError(
275 "TransformerBridge.init_weights() is only supported for TL-native "
276 "bridges created with TransformerBridge.boot_native(...); this bridge "
277 f"wraps {type(model).__name__}."
278 )
279 initialize_native_model(model, self.cfg)
281 def _set_processed_weight_attributes(self) -> None:
282 """Create 3D processed weight attributes for attention components.
284 For each attention component, if it has 2D weights (q.weight, k.weight, v.weight),
285 reshape them to 3D format [n_heads, d_model, d_head] and set as:
286 - _processed_W_Q
287 - _processed_W_K
288 - _processed_W_V
289 - _processed_b_Q
290 - _processed_b_K
291 - _processed_b_V
293 This allows property aliases (W_Q, W_K, W_V) to return 3D format for
294 HookedTransformer compatibility while keeping 2D format for calculations.
295 """
297 n_heads = self.cfg.n_heads
298 d_head = self.cfg.d_head
299 d_model = self.cfg.d_model
300 blocks_iter = []
301 for bl_name in _BLOCK_LIST_ATTRS:
302 if hasattr(self, bl_name):
303 blocks_iter.append(getattr(self, bl_name))
304 if not blocks_iter:
305 return
306 for block in [b for bl in blocks_iter for b in bl]:
307 if "attn" not in block._modules:
308 continue
309 attn = block.attn
310 if not (hasattr(attn, "q") and hasattr(attn.q, "weight")):
311 continue
312 try:
313 w_q_2d = attn.q.weight.data
314 w_k_2d = attn.k.weight.data
315 w_v_2d = attn.v.weight.data
316 attn._processed_W_Q = einops.rearrange(
317 w_q_2d, "m (i h) -> i m h", i=n_heads, h=d_head
318 )
319 attn._processed_W_K = einops.rearrange(
320 w_k_2d, "m (i h) -> i m h", i=n_heads, h=d_head
321 )
322 attn._processed_W_V = einops.rearrange(
323 w_v_2d, "m (i h) -> i m h", i=n_heads, h=d_head
324 )
325 if hasattr(attn.q, "bias") and attn.q.bias is not None:
326 b_q_2d = attn.q.bias.data
327 b_k_2d = attn.k.bias.data
328 b_v_2d = attn.v.bias.data
329 attn._processed_b_Q = einops.rearrange(
330 b_q_2d, "(i h) -> i h", i=n_heads, h=d_head
331 )
332 attn._processed_b_K = einops.rearrange(
333 b_k_2d, "(i h) -> i h", i=n_heads, h=d_head
334 )
335 attn._processed_b_V = einops.rearrange(
336 b_v_2d, "(i h) -> i h", i=n_heads, h=d_head
337 )
338 if hasattr(attn, "o") and hasattr(attn.o, "weight"):
339 w_o_2d = attn.o.weight.data
340 w_o_transposed = w_o_2d.T
341 attn._processed_W_O = einops.rearrange(
342 w_o_transposed, "m (i h) -> i h m", i=n_heads, h=d_head
343 )
344 if hasattr(attn.o, "bias") and attn.o.bias is not None:
345 attn._processed_b_O = attn.o.bias.data
346 except Exception:
347 pass
349 def _register_all_aliases_recursive(self) -> None:
350 """Recursively register aliases on all bridge components.
352 This walks through all components and calls _register_aliases() on each one.
353 Used after weight processing to ensure aliases point to processed weights.
354 """
355 if hasattr(self, "_register_aliases"): 355 ↛ 357line 355 didn't jump to line 357 because the condition on line 355 was always true
356 self._register_aliases()
357 for module in self.modules():
358 if module is not self and hasattr(module, "_register_aliases"):
359 getattr(module, "_register_aliases")()
361 def __setattr__(self, name: str, value: Any) -> None:
362 """Override setattr to track HookPoint objects dynamically."""
363 # nn.Module.__setattr__ claims Module values before any data descriptor runs, so the
364 # original_model property setter would never fire. object.__setattr__ invokes it:
365 # a registered copy aliases every HF weight into state_dict and makes .to() move twice.
366 if name == "original_model":
367 object.__setattr__(self, name, value)
368 return
369 super().__setattr__(name, value)
370 if isinstance(value, HookPoint):
371 value.name = name
372 self._hook_registry[name] = value
373 elif hasattr(value, "get_hooks") and callable(getattr(value, "get_hooks")):
374 component_hooks = value.get_hooks()
375 for hook_name, hook in component_hooks.items():
376 full_name = f"{name}.{hook_name}"
377 hook.name = full_name
378 self._hook_registry[full_name] = hook
380 def _scan_existing_hooks(self, module: nn.Module, prefix: str = "") -> None:
381 """Scan existing modules for hooks and add them to registry."""
382 visited = set()
383 # Protect canonical HookPoint names from alias overwrites. Seeded with
384 # already-registered hooks so the post-alias-registration re-scan adds
385 # alias registry entries without renaming shared HookPoint instances
386 # (dir() sorts alphabetically, so block-level alias attributes like
387 # hook_mlp_out are visited before the mlp child they point into).
388 named_hook_ids: set = {id(hp) for hp in self._hook_registry.values() if hp.name is not None}
390 def scan_module(mod: nn.Module, path: str = "") -> None:
391 obj_id = id(mod)
392 if obj_id in visited:
393 return
394 visited.add(obj_id)
395 if hasattr(mod, "get_hooks") and callable(getattr(mod, "get_hooks")):
396 component_hooks = mod.get_hooks() # type: ignore[operator]
397 if isinstance(component_hooks, dict): 397 ↛ 406line 397 didn't jump to line 406 because the condition on line 397 was always true
398 hooks_dict = cast(Dict[str, HookPoint], component_hooks)
399 for hook_name, hook in hooks_dict.items():
400 full_name = f"{path}.{hook_name}" if path else hook_name
401 hook_id = id(hook)
402 if hook_id not in named_hook_ids:
403 hook.name = full_name
404 named_hook_ids.add(hook_id)
405 self._hook_registry[full_name] = hook
406 for attr_name in dir(mod):
407 if attr_name.startswith("_"):
408 continue
409 if attr_name == "original_component" or attr_name == "original_model":
410 continue
411 if attr_name in [
412 "OV",
413 "QK",
414 "W_V",
415 "W_O",
416 "W_Q",
417 "W_K",
418 "W_in",
419 "W_gate",
420 "W_out",
421 "b_V",
422 "b_O",
423 "b_Q",
424 "b_K",
425 "b_in",
426 "b_out",
427 ]:
428 continue
429 try:
430 attr = getattr(mod, attr_name)
431 except (AttributeError, NameError, RuntimeError, TypeError):
432 continue
433 name = f"{path}.{attr_name}" if path else attr_name
434 if isinstance(attr, HookPoint):
435 hook_id = id(attr)
436 if hook_id not in named_hook_ids:
437 attr.name = name
438 named_hook_ids.add(hook_id)
439 self._hook_registry[name] = attr
440 for child_name, child_module in mod.named_children():
441 if (
442 child_name == "original_component"
443 or child_name == "_original_component"
444 or child_name == "original_model"
445 ):
446 continue
447 child_path = f"{path}.{child_name}" if path else child_name
448 scan_module(child_module, child_path)
450 scan_module(module, prefix)
452 @property
453 def n_params_total(self) -> int:
454 """Number of parameters in the wrapped model before bridge instrumentation.
456 This follows PyTorch's parameter iteration semantics, counting tied
457 parameters once. Bridge-created split views and synthetic zero tensors
458 are excluded, so the result can differ from
459 :attr:`HookedTransformer.n_params_total` and :meth:`tl_parameters`.
461 Returns:
462 int: Parameter count of the uninstrumented wrapped model.
463 """
464 return self._n_params_total
466 def _has_registered_blocks(self) -> bool:
467 """Whether a ``blocks`` stack is registered as a submodule on this bridge.
469 Checks ``_modules`` directly rather than ``hasattr``: ``__getattr__`` falls
470 through to the wrapped HF model, so ``hasattr(self, "blocks")`` can be True
471 for a model that merely exposes its own ``.blocks`` attribute.
472 """
473 modules = self.__dict__.get("_modules") or {}
474 return "blocks" in modules
476 def __getattr__(self, name: str) -> Any:
477 """Provide a clear error message for missing attributes."""
478 # Re-invoke original_model's property so its descriptive AttributeError
479 # for non-torch drivers isn't shadowed by the __dict__ fallback below.
480 if name == "original_model":
481 prop = type(self).__dict__.get("original_model")
482 if isinstance(prop, property) and prop.fget is not None:
483 return prop.fget(self)
484 if name in self.__dict__: # type: ignore[arg-type] 484 ↛ 485line 484 didn't jump to line 485 because the condition on line 484 was never true
485 return self.__dict__[name]
486 # Use __dict__ directly to avoid recursion
487 if "_modules" in self.__dict__ and name in self.__dict__["_modules"]: # type: ignore[arg-type]
488 return self.__dict__["_modules"][name]
489 adapter = self.__dict__.get("adapter")
490 component_mapping = getattr(adapter, "component_mapping", None)
491 if component_mapping is not None and name in component_mapping:
492 raise AttributeError(f"'{type(self).__name__}' object has no attribute '{name}'")
493 if "original_model" in self.__dict__ and self.__dict__["original_model"] is not None:
494 try:
495 name_split = name.split(".")
496 if len(name_split) > 1: 496 ↛ 497line 496 didn't jump to line 497 because the condition on line 496 was never true
497 current = getattr(self.__dict__["original_model"], name_split[0])
498 for part in name_split[1:]: # type: ignore[operator]
499 current = getattr(current, part)
500 return current
501 else:
502 return getattr(self.__dict__["original_model"], name)
503 except AttributeError:
504 pass # type: ignore[operator,assignment]
505 # A class property whose fget raised AttributeError lands here with the
506 # informative message discarded (CPython drops it before __getattr__).
507 # Re-invoke the property so its own diagnostic (e.g. "bias-free
508 # projection") surfaces instead of a generic missing-attribute error.
509 descriptor = getattr(type(self), name, None)
510 if isinstance(descriptor, property) and descriptor.fget is not None:
511 descriptor.fget(self)
512 raise AttributeError(f"'{type(self).__name__}' object has no attribute '{name}'")
514 def __str__(self) -> str:
515 """One-line-per-component summary of the bridge.
517 Returns:
518 A string describing the bridge's components.
519 """
520 lines = ["TransformerBridge:"]
521 mapping = self.adapter.get_component_mapping()
523 def _describe(component_mapping, indent):
524 pad = " " * indent
525 for name, component in component_mapping.items():
526 lines.append(f"{pad}{name}: {type(component).__name__}")
527 submodules = getattr(component, "submodules", None)
528 if submodules:
529 _describe(submodules, indent + 1)
531 _describe(mapping, 1)
532 return "\n".join(lines)
534 def enable_compatibility_mode(
535 self,
536 disable_warnings: bool = False,
537 no_processing: bool = False,
538 fold_ln: bool = True,
539 center_writing_weights: bool = True,
540 center_unembed: bool = True,
541 fold_value_biases: bool = True,
542 refactor_factored_attn_matrices: bool = False,
543 ) -> None:
544 """Apply HookedTransformer-equivalent weight processing and legacy hook compatibility.
546 Defaults match HookedTransformer's load-time processing (fold_ln + weight
547 centering) — required for analyses that reason in HookedTransformer's
548 post-processed coordinate system: logit lens, direct logit attribution,
549 residual-stream norms. Also enables legacy hook/component name aliases.
551 Hook semantic parity (issue #1317): ``hook_q_input``, ``hook_k_input``,
552 ``hook_v_input``, ``hook_attn_in``, and ``hook_mlp_in`` fire on the
553 pre-norm residual. Carve-outs: post-norm architectures (OLMo 2,
554 BERT-style) read the post-attention residual instead, and MLA blocks
555 (DeepSeek V2/V3/R1) do not expose the split-qkv aliases. ``hook_mlp_in``
556 is gated on ``cfg.use_hook_mlp_in``; toggle it via
557 :py:meth:`set_use_hook_mlp_in`.
559 Args:
560 disable_warnings: Whether to disable warnings about legacy components/hooks
561 no_processing: Whether to disable ALL pre-processing steps of the model.
562 If True, overrides fold_ln, center_writing_weights, and center_unembed to False.
563 fold_ln: Whether to fold layer norm weights into the subsequent linear layers.
564 Default: True. Ignored if no_processing=True.
565 center_writing_weights: Whether to center the writing weights (W_out in attention and MLPs).
566 Default: True. Ignored if no_processing=True.
567 center_unembed: Whether to center the unembedding matrix.
568 Default: True. Ignored if no_processing=True.
569 fold_value_biases: Whether to fold value biases into output bias.
570 Default: True. Ignored if no_processing=True.
571 refactor_factored_attn_matrices: Whether to refactor factored attention matrices.
572 Default: False. Ignored if no_processing=True.
573 """
574 from transformer_lens.utilities.bridge_components import (
575 apply_fn_to_all_components,
576 )
578 if not getattr(self.adapter, "supports_compatibility_mode", True):
579 raise RuntimeError(
580 f"{type(self.adapter).__name__} does not support compatibility mode: "
581 "its stored-processed-weights path is known to diverge from the "
582 "reference model. Use the default bridge forward instead."
583 )
585 hf_device_map = getattr(self.original_model, "hf_device_map", None)
586 if hf_device_map and not no_processing:
587 offloaded = {k for k, v in hf_device_map.items() if str(v).lower() in ("cpu", "disk")}
588 if offloaded: 588 ↛ 602line 588 didn't jump to line 602 because the condition on line 588 was always true
589 raise RuntimeError(
590 "enable_compatibility_mode() with weight processing "
591 "(fold_ln/center_writing_weights/center_unembed/fold_value_biases) is not "
592 "supported on a bridge with an offloaded device_map "
593 f"({sorted(offloaded)} are CPU/disk-offloaded). Weight processing reads and "
594 "rewrites parameters directly across many components at once, not through a "
595 "single component's own forward() call the way the default (non-compat-mode) "
596 "bridge forward does, so it isn't covered by GeneralizedComponent's per-call "
597 "materialization and hits raw meta tensors. Load without a CPU/disk device_map "
598 "for compatibility mode, or call enable_compatibility_mode(no_processing=True) "
599 "for the hook/component compatibility layer without the weight transforms."
600 )
602 if getattr(self.cfg, "is_audio_model", False):
603 # Audio encoders have no text embed/unembed for the legacy weight
604 # processing to operate on; without this guard the processing path
605 # dies later with an opaque KeyError ('embed.weight').
606 raise NotImplementedError(
607 "enable_compatibility_mode() is not supported for audio encoder models: "
608 "the legacy weight processing (fold_ln/centering) assumes a text "
609 "embed/unembed, which audio encoders do not have. Use the bridge's "
610 "native hooks (run_with_cache / run_with_hooks) directly."
611 )
613 self.compatibility_mode = True
615 def set_compatibility_mode(component: Any) -> None:
616 """Set compatibility mode on a component."""
617 component.compatibility_mode = True
618 component.disable_warnings = disable_warnings
620 apply_fn_to_all_components(self, set_compatibility_mode)
621 self.clear_hook_registry()
622 # Drop block capture-hook handles from any prior call so they don't accumulate.
623 if hasattr(self, "blocks"): 623 ↛ 627line 623 didn't jump to line 627 because the condition on line 623 was always true
624 for block in self.blocks:
625 if hasattr(block, "_teardown_capture_hooks"): 625 ↛ 624line 625 didn't jump to line 624 because the condition on line 625 was always true
626 block._teardown_capture_hooks()
627 try:
628 if not no_processing:
629 self.process_weights(
630 fold_ln=fold_ln,
631 center_writing_weights=center_writing_weights,
632 center_unembed=center_unembed,
633 fold_value_biases=fold_value_biases,
634 refactor_factored_attn_matrices=refactor_factored_attn_matrices,
635 )
636 finally:
637 # Re-initialize hooks even on failure so bridge stays usable
638 self._initialize_hook_registry()
639 self._setup_hook_compatibility()
640 self._register_all_aliases_recursive()
642 def _setup_hook_compatibility(self) -> None:
643 """Setup hook compatibility transformations to match HookedTransformer behavior.
645 This method sets up hook conversions and wrappers that ensure Bridge hooks
646 have the same shapes and behavior as HookedTransformer hooks. This includes:
647 1. hook_z reshaping from [batch, seq, d_model] to [batch, seq, n_heads, d_head]
648 2. Wrapping HF attention forward to inject position embeddings/attention masks
649 3. Architecture-specific setup (e.g., rotary embedding references)
651 This is called during __init__ and should always be run, regardless of whether
652 compatibility mode or weight processing is enabled.
654 Note: This method is idempotent - can be called multiple times safely.
655 """
656 if hasattr(self.adapter, "setup_hook_compatibility"):
657 self.adapter.setup_hook_compatibility(self)
658 elif hasattr(self.adapter, "setup_no_processing_hooks"): 658 ↛ 659line 658 didn't jump to line 659 because the condition on line 658 was never true
659 self.adapter.setup_no_processing_hooks(self)
660 blocks_to_process = []
661 for block_list_name in (
662 "blocks",
663 "encoder_blocks",
664 "decoder_blocks",
665 "L_blocks",
666 "H_blocks",
667 ):
668 if hasattr(self, block_list_name):
669 blocks_to_process.extend(getattr(self, block_list_name))
670 # A vision tower's layers are not a top-level block list, so they were
671 # skipped here and their attention never had q/k/v/z reshaping set up --
672 # those hooks fired flat, ignoring the tower's head count.
673 vision_encoder = getattr(self, "vision_encoder", None)
674 vision_layers = getattr(vision_encoder, "encoder_layers", None)
675 if vision_layers is not None:
676 blocks_to_process.extend(vision_layers)
677 for block in blocks_to_process:
678 for attn_name in ["attn", "self_attn", "cross_attn"]:
679 if hasattr(block, attn_name):
680 attn = getattr(block, attn_name)
681 if hasattr(attn, "setup_hook_compatibility"):
682 attn.setup_hook_compatibility()
683 elif hasattr(attn, "setup_no_processing_hooks"): 683 ↛ 684line 683 didn't jump to line 684 because the condition on line 683 was never true
684 attn.setup_no_processing_hooks()
686 def process_weights(
687 self,
688 verbose: bool = False,
689 fold_ln: bool = True,
690 center_writing_weights: bool = True,
691 center_unembed: bool = True,
692 fold_value_biases: bool = True,
693 refactor_factored_attn_matrices: bool = False,
694 ) -> None:
695 """Process weights directly using ProcessWeights and architecture adapter.
697 This method applies weight processing transformations to improve model interpretability
698 without requiring a reference HookedTransformer model. Works with all architectures
699 supported by TransformerBridge, including GPT-OSS and other new models.
701 Args:
702 verbose: If True, print detailed progress messages. Default: False
703 fold_ln: Fold LayerNorm weights/biases into subsequent layers. Default: True
704 center_writing_weights: Center weights that write to residual stream. Default: True
705 center_unembed: Center unembedding weights (translation invariant). Default: True
706 fold_value_biases: Fold value biases into output bias. Default: True
707 refactor_factored_attn_matrices: Experimental QK/OV factorization. Default: False
708 """
709 # Match HookedTransformer: warn and skip rather than raise. Folding and centering
710 # read their own factors out of the weights and neutralize them, so a repeat pass
711 # is a no-op for every architecture — but adapters that fold a factor held outside
712 # the weights would double-apply it, and notebook cell re-runs make that routine.
713 if self._weights_processed:
714 logging.warning(
715 "process_weights was already applied to this bridge. Skipping: re-running "
716 "it would re-apply adapter weight folds. Boot a fresh bridge with "
717 "TransformerBridge.boot_transformers(...) to process with different options."
718 )
719 return
721 # Folding and centering do arithmetic on raw weights, so packed or
722 # scale-separated storage would produce silent garbage. The forward
723 # path stays usable when quantized; only this transformation does not.
724 for name, param in self.original_model.named_parameters():
725 # No name filter: modern batched-MoE experts are Parameters that do
726 # NOT end in .weight (mlp.experts.gate_up_proj), and those are the
727 # very tensors the converters had to guard. unreadable_weight_reason
728 # is dtype-driven, so every full-width float parameter still passes.
729 # Routed through the shared helper so meta gets its own "load with
730 # real weights" message instead of being silently skipped — folding
731 # on meta tensors yields meta tensors, which is the same
732 # silent-garbage failure this guard exists to stop.
733 require_readable_weight(
734 param,
735 operation=f"process weights ({name})",
736 owner=self.original_model,
737 remedy=(
738 "Load the model dequantized, or use the bridge without weight "
739 "processing (enable_compatibility_mode(no_processing=True))."
740 ),
741 )
743 # A failed or partial processing attempt is no longer guaranteed to retain
744 # the raw HuggingFace basis, so invalidate that contract before any work.
745 self._weights_processed = True
746 from transformer_lens.weight_processing import ProcessWeights
748 if verbose: 748 ↛ 749line 748 didn't jump to line 749 because the condition on line 748 was never true
749 print(f"Processing weights for {self.cfg.model_name}...")
751 # Soft capping (tanh) is not translation-invariant; centering would change output.
752 if center_unembed and softcap_enabled(getattr(self.cfg, "output_logits_soft_cap", None)): 752 ↛ 753line 752 didn't jump to line 753 because the condition on line 752 was never true
753 logging.warning(
754 "center_unembed=True is incompatible with logit softcapping "
755 "(output_logits_soft_cap=%.1f). Disabling center_unembed.",
756 self.cfg.output_logits_soft_cap,
757 )
758 center_unembed = False
760 if verbose: 760 ↛ 761line 760 didn't jump to line 761 because the condition on line 760 was never true
761 print(" Extracting state dict from existing model...")
762 state_dict = self.state_dict()
763 adapter = self.adapter
765 # Untie embed/unembed weights (GPT-2) so centering affects only unembed
766 embed_key = "embed.weight"
767 unembed_key = "unembed.weight"
769 if embed_key in state_dict and unembed_key in state_dict: 769 ↛ 777line 769 didn't jump to line 777 because the condition on line 769 was always true
770 # Check if they point to the same tensor (weight tying)
771 if state_dict[embed_key].data_ptr() == state_dict[unembed_key].data_ptr():
772 if verbose: 772 ↛ 773line 772 didn't jump to line 773 because the condition on line 772 was never true
773 print(" Breaking weight tying between embed and unembed in state dict...")
774 # Clone the unembed weight to break the tie
775 state_dict[unembed_key] = state_dict[unembed_key].clone()
777 if adapter and hasattr(adapter, "preprocess_weights"): 777 ↛ 783line 777 didn't jump to line 783 because the condition on line 777 was always true
778 adapter._fold_ln_requested = fold_ln # type: ignore[union-attr]
779 state_dict = adapter.preprocess_weights(state_dict)
781 # Use the unified ProcessWeights.process_weights() pipeline.
782 # Float32 upcasting for precision is handled centrally in process_weights().
783 if verbose: 783 ↛ 784line 783 didn't jump to line 784 because the condition on line 783 was never true
784 print(" Processing weights (fold_ln, center_writing_weights, etc.)...")
785 state_dict = ProcessWeights.process_weights(
786 state_dict,
787 self.cfg,
788 fold_ln=fold_ln,
789 center_writing_weights=center_writing_weights,
790 center_unembed=center_unembed,
791 fold_value_biases=fold_value_biases,
792 refactor_factored_attn_matrices=refactor_factored_attn_matrices,
793 adapter=adapter,
794 )
796 # Normalize HF-prefix keys to TL format for weight routing
797 import re
799 hf_to_tl_prefix = {}
800 for tl_name, (remote_path, _component) in self.real_components.items():
801 if remote_path and remote_path != tl_name: 801 ↛ 800line 801 didn't jump to line 800 because the condition on line 801 was always true
802 hf_to_tl_prefix[remote_path] = tl_name
804 normalized_state_dict = {}
805 for key, value in state_dict.items():
806 new_key = key
807 for hf_prefix, tl_prefix in hf_to_tl_prefix.items():
808 if key.startswith(hf_prefix + "."): 808 ↛ 809line 808 didn't jump to line 809 because the condition on line 808 was never true
809 suffix = key[len(hf_prefix) + 1 :]
810 new_key = f"{tl_prefix}.{suffix}"
811 break
812 normalized_state_dict[new_key] = value
813 state_dict = normalized_state_dict
815 if verbose: 815 ↛ 816line 815 didn't jump to line 816 because the condition on line 815 was never true
816 print(" Distributing weights to generalized components...")
817 ProcessWeights.distribute_weights_to_components(
818 state_dict=state_dict,
819 component_mapping=self.real_components,
820 )
821 if fold_ln:
822 self._fold_layer_norms_into_batched_experts()
823 if adapter is not None: 823 ↛ exitline 823 didn't return from function 'process_weights' because the condition on line 823 was always true
824 adapter.postprocess_weights(self)
826 def _fold_layer_norms_into_batched_experts(self) -> None:
827 """Fold each batched-expert MoE layer's ln2 into the weights that read it.
829 transformers stores a whole expert stack in one 3-D Parameter, which is not a
830 ``weight``/``bias`` leaf of a declared bridge submodule and so never appears in
831 the state dict ProcessWeights folds. Left alone those layers keep their learned
832 FFN gains while attention and dense layers of the same model lose theirs, and
833 the mixed basis surfaces only in DLA / logit-lens / factored-matrix reads.
835 Runs on the live model because the readers include tensors the state dict
836 cannot carry: batched experts, and the shared-expert MLPs that Qwen2-MoE and
837 GLM4-MoE hang off the same norm.
838 """
839 # An adapter that declares fold_ln unsupported cannot have its norms folded
840 # weight-preservingly. ProcessWeights already warns and skips for these; folding
841 # the expert stack anyway leaves the model diverging from the reference.
842 if not getattr(self.adapter, "supports_fold_ln", True):
843 return
845 uses_offset = bool(getattr(self.cfg, "rmsnorm_uses_offset", False))
846 skipped: list[str] = []
847 for list_name in ("blocks", "encoder_blocks", "decoder_blocks"):
848 for index, block in enumerate(getattr(self, list_name, None) or []):
849 mlp = getattr(block, "mlp", None)
850 if mlp is None or not has_batched_experts(mlp):
851 continue
852 label = f"{list_name}.{index}"
853 norm = getattr(block, "ln2", None)
854 weight = getattr(norm, "weight", None) if norm is not None else None
855 if not isinstance(weight, torch.Tensor) or weight.ndim != 1: 855 ↛ 856line 855 didn't jump to line 856 because the condition on line 855 was never true
856 skipped.append(label)
857 continue
858 if getattr(norm, "bias", None) is not None: 858 ↛ 860line 858 didn't jump to line 860 because the condition on line 858 was never true
859 # Folding the gain but not the shift would change the block's output.
860 skipped.append(label)
861 continue
862 scale = weight.detach().clone()
863 if uses_offset: 863 ↛ 864line 863 didn't jump to line 864 because the condition on line 863 was never true
864 scale += 1.0
865 if not fold_scale_into_moe_block(mlp, scale): 865 ↛ 866line 865 didn't jump to line 866 because the condition on line 865 was never true
866 skipped.append(label)
867 continue
868 with torch.no_grad():
869 weight.fill_(0.0 if uses_offset else 1.0)
870 if skipped: 870 ↛ 871line 870 didn't jump to line 871 because the condition on line 870 was never true
871 logging.warning(
872 "fold_ln could not reach the expert weights of %s, so those layers keep "
873 "their FFN norm gains while the rest of the model is folded. Direct "
874 "logit attribution and logit lens will mix two bases.",
875 ", ".join(skipped),
876 )
878 def to_tokens(
879 self,
880 input: Union[str, List[str]],
881 prepend_bos: Optional[bool] = None,
882 padding_side: Optional[str] = None,
883 move_to_device: bool = True,
884 truncate: bool = True,
885 ) -> torch.Tensor:
886 """Converts a string to a tensor of tokens.
888 See the class-level "Tokenization notes" for full ``prepend_bos``
889 semantics, the ``default_prepend_bos`` /
890 ``tokenizer_prepends_bos`` interaction, and the whitespace-
891 sensitivity gotcha. **Pass ``prepend_bos=False`` whenever you're
892 tokenizing only part of a prompt.**
894 Args:
895 input: The input to tokenize.
896 prepend_bos: Overrides ``self.cfg.default_prepend_bos``. Defaults
897 to ``None`` (use the cfg setting). Pass ``True`` or ``False``
898 to override locally.
899 padding_side: Which side to pad on when tokenizing multiple
900 strings of different lengths. Defaults to the tokenizer's
901 ``padding_side``.
902 move_to_device: Whether to move the result to ``cfg.device``.
903 truncate: Whether to truncate inputs longer than ``cfg.n_ctx``.
905 Returns:
906 Token tensor of shape ``[batch, pos]``.
907 """
908 assert self.tokenizer is not None, "Cannot use to_tokens without a tokenizer"
909 if prepend_bos is None:
910 prepend_bos = getattr(self.cfg, "default_prepend_bos", True)
911 if padding_side is None:
912 padding_side = getattr(self.tokenizer, "padding_side", "right")
913 tokenizer_prepends_bos = getattr(self.cfg, "tokenizer_prepends_bos", True)
914 if prepend_bos and (not tokenizer_prepends_bos):
915 bos = self.tokenizer.bos_token
916 encodes_atomically = (
917 bos is not None
918 and len(self.tokenizer(bos, add_special_tokens=False)["input_ids"]) == 1
919 )
920 if encodes_atomically:
921 input = utils.get_input_with_manually_prepended_bos(bos, input)
922 # else: the fallback BOS is not an atom in this vocab (e.g.
923 # '<|endoftext|>' installed on BERT); prepending the string would
924 # tokenize to subword garbage, so skip rather than pollute the input.
925 if isinstance(input, str):
926 input = [input]
927 tokens = self.tokenizer(
928 input,
929 return_tensors="pt",
930 padding=True,
931 padding_side=padding_side,
932 truncation=truncate,
933 max_length=self.cfg.n_ctx if truncate else None,
934 )["input_ids"]
935 # Strip auto-appended EOS tokens (e.g., OLMo)
936 if (
937 getattr(self.cfg, "tokenizer_appends_eos", False)
938 and self.tokenizer.eos_token_id is not None
939 ):
940 # Remove trailing EOS, keep at least 1 token
941 while tokens.shape[-1] > 1 and (tokens[:, -1] == self.tokenizer.eos_token_id).all():
942 tokens = tokens[:, :-1]
943 if not prepend_bos and tokenizer_prepends_bos:
944 tokens = utils.get_tokens_with_bos_removed(
945 self.tokenizer, tokens, padding_side=padding_side
946 )
947 if move_to_device:
948 tokens = tokens.to(self.cfg.device)
949 return tokens
951 def encoder_output(
952 self,
953 frames: torch.Tensor,
954 one_zero_attention_mask: Optional[torch.Tensor] = None,
955 ) -> torch.Tensor:
956 """Run the audio encoder from precomputed frames, skipping feature extraction.
958 The audio-path analogue of ``start_at_layer``: ``frames`` is
959 ``[batch, frames, d_model]``, the tensor the conv front end would have
960 produced (observable at ``feat_proj.hook_out``), so callers can inject or
961 reuse frames without re-running the waveform convolutions. Positional
962 convolution and the encoder layer norm are applied first, exactly as the
963 full path does, then the blocks run. Mirrors
964 ``HookedAudioEncoder.encoder_output``.
966 Hooks on the bridged components fire as usual, so this composes with
967 ``add_hook`` / ``get_caching_hooks``.
969 Args:
970 frames: ``[batch, frames, d_model]`` precomputed encoder frames.
971 one_zero_attention_mask: Optional ``[batch, frames]`` mask, 1 for
972 real frames and 0 for padding.
974 Returns:
975 The residual stream leaving the final block, ``[batch, frames, d_model]``.
976 """
977 self._require_frame_entry_support()
978 if frames.ndim != 3:
979 raise ValueError(
980 "encoder_output expects precomputed frames [batch, frames, d_model]; "
981 f"got a {frames.ndim}D tensor. Pass a waveform to forward() instead."
982 )
983 frames = frames.to(self.cfg.device)
985 if one_zero_attention_mask is not None:
986 # HF zeroes pad frames before the positional conv (kernel 128 smears
987 # pad content into real frames otherwise); masked_fill, not HF's
988 # in-place write, so the caller's tensor survives.
989 frames = frames.masked_fill(
990 ~one_zero_attention_mask.to(frames.device).bool().unsqueeze(-1), 0.0
991 )
993 resid = frames + self.conv_pos_embed(frames)
994 # Post-LN encoders (do_stable_layer_norm=False: hubert-base,
995 # wav2vec2-base) normalize before the blocks; stable-LN encoders
996 # (wav2vec2-large and kin) apply this same module AFTER the blocks —
997 # HF's Wav2Vec2EncoderStableLayerNorm.forward. Mirror the real order or
998 # every stable-LN activation is silently wrong.
999 stable_ln = self._audio_encoder_is_stable_layer_norm()
1000 if not stable_ln:
1001 resid = self.embed_ln(resid)
1003 additive_attention_mask = None
1004 if one_zero_attention_mask is not None:
1005 mask = one_zero_attention_mask.to(self.cfg.device)
1006 additive_attention_mask = torch.where(
1007 mask[:, None, None, :] == 0,
1008 torch.tensor(float("-inf"), dtype=resid.dtype, device=resid.device),
1009 torch.tensor(0.0, dtype=resid.dtype, device=resid.device),
1010 )
1012 for block in self.blocks:
1013 output = block(resid, attention_mask=additive_attention_mask)
1014 resid = output[0] if isinstance(output, tuple) else output
1015 if stable_ln:
1016 resid = self.embed_ln(resid)
1017 return resid
1019 def _audio_encoder_is_stable_layer_norm(self) -> bool:
1020 """Whether the wrapped audio encoder is the pre-LN ("stable") variant.
1022 Structural, not config-driven: build_bridge_from_module never runs the
1023 adapter's prepare_loading, so cfg.do_stable_layer_norm can be absent
1024 on that path while the module class is authoritative on both.
1025 """
1026 model = self.original_model
1027 for attr in ("encoder", "wav2vec2", "hubert", "model"): 1027 ↛ 1034line 1027 didn't jump to line 1034 because the loop on line 1027 didn't complete
1028 candidate = getattr(model, attr, None)
1029 if candidate is None: 1029 ↛ 1030line 1029 didn't jump to line 1030 because the condition on line 1029 was never true
1030 continue
1031 encoder = candidate if attr == "encoder" else getattr(candidate, "encoder", None)
1032 if encoder is not None: 1032 ↛ 1027line 1032 didn't jump to line 1027 because the condition on line 1032 was always true
1033 return type(encoder).__name__.endswith("StableLayerNorm")
1034 return bool(getattr(self.cfg, "do_stable_layer_norm", False))
1036 def _require_frame_entry_support(self) -> None:
1037 """Reject models with no conv-frame stage to re-enter."""
1038 if not getattr(self.cfg, "is_audio_model", False):
1039 raise NotImplementedError(
1040 "encoder_output is an audio-encoder entry point; this bridge is not an "
1041 "audio model. Use start_at_layer for residual re-entry on text models."
1042 )
1043 missing = [
1044 name
1045 for name in ("conv_pos_embed", "embed_ln", "blocks")
1046 if self._modules.get(name) is None
1047 ]
1048 if missing:
1049 raise NotImplementedError(
1050 "encoder_output needs a waveform encoder with a convolutional front end "
1051 f"(missing {missing}). Spectrogram encoders such as AST have no "
1052 "precomputed-frame stage to re-enter, so there is nothing to bypass."
1053 )
1055 def to_sentence_pair_tokens(
1056 self,
1057 sentence_a: str,
1058 sentence_b: str,
1059 move_to_device: bool = True,
1060 truncate: bool = True,
1061 ) -> Dict[str, torch.Tensor]:
1062 """Pair-tokenize two sentences as ``[CLS] a [SEP] b [SEP]``.
1064 Returns ``input_ids``, ``token_type_ids`` and ``attention_mask``. The
1065 segment ids are not decorative: without them a BERT NSP head sees both
1066 sentences as one segment and its logits collapse. Mirrors
1067 ``BertNextSentencePrediction.to_tokens``.
1069 Args:
1070 sentence_a: First sentence of the pair.
1071 sentence_b: Second sentence of the pair.
1072 move_to_device: Move the returned tensors to ``cfg.device``.
1073 truncate: Truncate to the model's context window.
1074 """
1075 assert self.tokenizer is not None, "Cannot pair-tokenize without a tokenizer"
1076 encodings = self.tokenizer(
1077 sentence_a,
1078 sentence_b,
1079 return_tensors="pt",
1080 padding=True,
1081 truncation=truncate,
1082 max_length=self.cfg.n_ctx if truncate else None,
1083 )
1084 if "token_type_ids" not in encodings: 1084 ↛ 1085line 1084 didn't jump to line 1085 because the condition on line 1084 was never true
1085 raise ValueError(
1086 f"{type(self.tokenizer).__name__} emits no token_type_ids, so it cannot "
1087 "express a sentence pair. Next-sentence prediction needs a "
1088 "segment-aware tokenizer (e.g. BERT's)."
1089 )
1090 keys = ("input_ids", "token_type_ids", "attention_mask")
1091 tokens = {key: encodings[key] for key in keys if key in encodings}
1092 if move_to_device: 1092 ↛ 1094line 1092 didn't jump to line 1094 because the condition on line 1092 was always true
1093 tokens = {key: value.to(self.cfg.device) for key, value in tokens.items()}
1094 return tokens
1096 def predict_next_sentence(
1097 self,
1098 sentence_a: str,
1099 sentence_b: str,
1100 return_type: Optional[str] = "predictions",
1101 truncate: bool = True,
1102 ) -> Any:
1103 """Run next-sentence prediction over a sentence pair given as strings.
1105 Owns the ``token_type_ids`` plumbing that a hand-rolled pair forward has
1106 to remember. Requires a bridge booted onto an NSP head — otherwise the
1107 model has no 2-class output to decode. Mirrors
1108 ``BertNextSentencePrediction.forward``.
1110 Args:
1111 sentence_a: First sentence of the pair.
1112 sentence_b: Second sentence of the pair.
1113 return_type: ``"predictions"`` for the decoded verdict, or
1114 ``"logits"`` for the raw 2-class scores.
1115 truncate: Truncate to the model's context window.
1116 """
1117 tokens = self.to_sentence_pair_tokens(sentence_a, sentence_b, truncate=truncate)
1118 forward_kwargs: Dict[str, Any] = {
1119 key: value for key, value in tokens.items() if key != "input_ids"
1120 }
1121 logits = self(tokens["input_ids"], return_type="logits", **forward_kwargs)
1122 if logits.shape[-1] != 2:
1123 raise ValueError(
1124 "predict_next_sentence needs a next-sentence-prediction head, but this "
1125 f"bridge produces {logits.shape[-1]} output classes. Boot it with "
1126 "model_class=BertForNextSentencePrediction."
1127 )
1128 if return_type == "logits":
1129 return logits
1130 return self._finalize_return(return_type, logits, tokens["input_ids"])
1132 def to_string(
1133 self, tokens: Union[List[int], torch.Tensor, np.ndarray]
1134 ) -> Union[str, List[str]]:
1135 """Convert tokens to string(s).
1137 Args:
1138 tokens: Tokens to convert
1140 Returns:
1141 Decoded string(s)
1142 """
1143 if not isinstance(tokens, torch.Tensor): 1143 ↛ 1144line 1143 didn't jump to line 1144 because the condition on line 1143 was never true
1144 tokens = torch.tensor(tokens)
1145 if len(tokens.shape) == 2:
1146 return self.tokenizer.batch_decode(tokens, clean_up_tokenization_spaces=False)
1147 elif len(tokens.shape) <= 1: 1147 ↛ 1150line 1147 didn't jump to line 1150 because the condition on line 1147 was always true
1148 return self.tokenizer.decode(tokens, clean_up_tokenization_spaces=False)
1149 else:
1150 raise ValueError(f"Invalid shape passed in: {tokens.shape}")
1152 def to_str_tokens(
1153 self,
1154 input: Union[str, torch.Tensor, np.ndarray, List],
1155 prepend_bos: Optional[bool] = None,
1156 padding_side: Optional[str] = None,
1157 ) -> Union[List[str], List[List[str]]]:
1158 """Map text or tokens to a list of tokens as strings.
1160 See the class-level "Tokenization notes" for full ``prepend_bos``
1161 semantics. **Pass ``prepend_bos=False`` whenever you're tokenizing
1162 only part of a prompt.** When ``input`` is already a tensor or
1163 array, ``prepend_bos`` and ``padding_side`` are ignored.
1165 Args:
1166 input: A string, list of strings, or tensor/array of token IDs.
1167 prepend_bos: Overrides ``self.cfg.default_prepend_bos``. Only
1168 applies when ``input`` is a string. Defaults to ``None``
1169 (use the cfg setting).
1170 padding_side: Which side to pad on. Only applies when ``input``
1171 is a string.
1173 Returns:
1174 List of token strings.
1175 """
1176 if isinstance(input, list):
1177 return cast(
1178 List[List[str]],
1179 [self.to_str_tokens(item, prepend_bos, padding_side) for item in input],
1180 )
1181 elif isinstance(input, str):
1182 tokens = self.to_tokens(input, prepend_bos=prepend_bos, padding_side=padding_side)[0]
1183 elif isinstance(input, torch.Tensor): 1183 ↛ 1190line 1183 didn't jump to line 1190 because the condition on line 1183 was always true
1184 tokens = input.squeeze()
1185 if tokens.dim() == 0: 1185 ↛ 1186line 1185 didn't jump to line 1186 because the condition on line 1185 was never true
1186 tokens = tokens.unsqueeze(0)
1187 assert (
1188 tokens.dim() == 1
1189 ), f"Invalid tokens input to to_str_tokens, has shape: {tokens.shape}"
1190 elif isinstance(input, np.ndarray):
1191 tokens_np = input.squeeze()
1192 if tokens_np.ndim == 0:
1193 tokens_np = np.expand_dims(tokens_np, axis=0)
1194 assert (
1195 tokens_np.ndim == 1
1196 ), f"Invalid tokens input to to_str_tokens, has shape: {tokens_np.shape}"
1197 tokens = torch.tensor(tokens_np)
1198 else:
1199 raise ValueError(f"Invalid input type to to_str_tokens: {type(input)}")
1200 # v5 compat: wrap each token so batch_decode decodes them individually
1201 tokens_list = [[int(t)] for t in tokens.tolist()]
1202 str_tokens = self.tokenizer.batch_decode(tokens_list, clean_up_tokenization_spaces=False)
1203 return str_tokens
1205 def to_single_token(self, string: str) -> int:
1206 """Map a string that makes up a single token to the id for that token.
1208 Args:
1209 string: The string to convert
1211 Returns:
1212 Token ID
1214 Raises:
1215 AssertionError: If string is not a single token
1216 """
1217 token = self.to_tokens(string, prepend_bos=False).squeeze()
1218 if token.numel() != 1: 1218 ↛ 1219line 1218 didn't jump to line 1219 because the condition on line 1218 was never true
1219 raise AssertionError(f"Input string: {string} is not a single token!")
1220 return int(token.item())
1222 def get_token_position(
1223 self,
1224 single_token: Union[str, int],
1225 input: Union[str, torch.Tensor],
1226 mode="first",
1227 prepend_bos: Optional[Union[bool, None]] = None,
1228 padding_side: Optional[Union[Literal["left", "right"], None]] = None,
1229 ):
1230 """Get the position of a single_token in a string or sequence of tokens.
1232 Raises an error if the token is not present.
1234 When ``input`` is a string it's tokenized internally — see the
1235 class-level "Tokenization notes" for ``prepend_bos`` semantics.
1236 Off-by-one position errors usually mean ``prepend_bos`` is on
1237 when it shouldn't be (or vice versa); pass ``prepend_bos=False``
1238 when ``input`` is a fragment of a larger prompt.
1240 Args:
1241 single_token (Union[str, int]): The token to search for. Can
1242 be a token index, or a string (but the string must correspond to a single token).
1243 input (Union[str, torch.Tensor]): The sequence to
1244 search in. Can be a string or a rank 1 tensor of tokens or a rank 2 tensor of tokens
1245 with a dummy batch dimension.
1246 mode (str, optional): If there are multiple matches, which match to return. Supports
1247 "first" or "last". Defaults to "first".
1248 prepend_bos (bool, optional): Overrides ``self.cfg.default_prepend_bos``. Only
1249 applies when ``input`` is a string. Defaults to ``None`` (use the cfg setting).
1250 padding_side (Union[Literal["left", "right"], None], optional): Specifies which
1251 side to pad when tokenizing multiple strings of different lengths.
1252 """
1253 if isinstance(input, str):
1254 tokens = self.to_tokens(input, prepend_bos=prepend_bos, padding_side=padding_side)
1255 else:
1256 tokens = input
1257 if len(tokens.shape) == 2:
1258 assert (
1259 tokens.shape[0] == 1
1260 ), f"If tokens are rank two, they must have shape [1, seq_len], not {tokens.shape}"
1261 tokens = tokens[0]
1262 if isinstance(single_token, str):
1263 single_token = self.to_single_token(single_token)
1264 elif isinstance(single_token, torch.Tensor): 1264 ↛ 1265line 1264 didn't jump to line 1265 because the condition on line 1264 was never true
1265 single_token = single_token.item()
1266 indices = torch.arange(len(tokens), device=tokens.device)[tokens == single_token]
1267 assert len(indices) > 0, "The token does not occur in the prompt"
1268 if mode == "first":
1269 return indices[0].item()
1270 elif mode == "last": 1270 ↛ 1273line 1270 didn't jump to line 1273 because the condition on line 1270 was always true
1271 return indices[-1].item()
1272 else:
1273 raise ValueError(f"mode must be 'first' or 'last', not {mode}")
1275 def to_single_str_token(self, int_token: int) -> str:
1276 """Get the single token corresponding to an int in string form.
1278 Args:
1279 int_token: The token ID
1281 Returns:
1282 The token string
1283 """
1284 assert isinstance(int_token, int)
1285 token = self.to_str_tokens(torch.tensor([int_token]))
1286 if isinstance(token, list) and len(token) == 1:
1287 return str(token[0])
1288 raise AssertionError("Expected a single string token.")
1290 def _enumerate_blocks(self) -> List[Tuple[int, Any]]:
1291 """(index, block) over every registered block list, encoder before decoder.
1293 Decoder-only models register a single ``blocks``, so the indices are the
1294 plain layer indices. Encoder-decoder models register ``encoder_blocks``
1295 and ``decoder_blocks`` instead; those are concatenated into one index
1296 space, matching ``HookedEncoderDecoder``'s ``chain(encoder, decoder)``.
1297 """
1298 pairs: List[Tuple[int, Any]] = []
1299 for list_name in _BLOCK_LIST_ATTRS:
1300 block_list = self._modules.get(list_name)
1301 if not isinstance(block_list, nn.ModuleList):
1302 continue
1303 for block in block_list:
1304 pairs.append((len(pairs), block))
1305 return pairs
1307 def _resolve_submodule_name(self, block: Any, submodule: str) -> Optional[str]:
1308 """The block's actual name for ``submodule``.
1310 Encoder blocks name self-attention ``attn`` while decoder blocks name it
1311 ``self_attn``, so a caller asking for ``attn`` means "this block's
1312 self-attention" on either side. Cross-attention is never resolved here:
1313 ``HookedEncoderDecoder`` omits it from stacked weights, and this mirrors
1314 that so the two stacks line up layer for layer.
1315 """
1316 for candidate in _SELF_ATTENTION_NAMES.get(submodule, (submodule,)):
1317 if candidate in block._modules:
1318 return candidate
1319 return None
1321 def _rewrite_submodule_path(self, attr_path: str, submodule: str, actual: Optional[str]) -> str:
1322 """Re-point ``attr_path``'s leading segment at the block's actual submodule."""
1323 if actual is None or actual == submodule:
1324 return attr_path
1325 if attr_path.split(".")[0] != submodule: 1325 ↛ 1326line 1325 didn't jump to line 1326 because the condition on line 1325 was never true
1326 return attr_path
1327 return actual + attr_path[len(submodule) :]
1329 def blocks_with(self, submodule: str) -> List[Tuple[int, "GeneralizedComponent"]]:
1330 """Return (index, block) pairs for blocks with the named bridged submodule.
1332 Checks _modules (not hasattr) so HF-internal attrs don't match.
1333 Use instead of assuming blocks[0] is representative on hybrid models.
1334 On encoder-decoder models the indices span encoder then decoder blocks,
1335 and ``"attn"`` matches the decoder's ``self_attn`` too.
1336 """
1337 return [
1338 (index, block)
1339 for index, block in self._enumerate_blocks()
1340 if self._resolve_submodule_name(block, submodule) is not None
1341 ]
1343 def stack_params_for(
1344 self, submodule: str, attr_path: str, reshape_fn: Optional[Callable] = None
1345 ) -> Tuple[List[int], torch.Tensor]:
1346 """Stack a parameter across matching blocks only. Returns (layer_indices, tensor).
1348 Use for hybrid models where not all blocks have the submodule.
1349 """
1350 matching = self.blocks_with(submodule)
1351 if not matching:
1352 raise ValueError(
1353 f"No blocks have submodule '{submodule}'. "
1354 f"Available submodules can be checked with blocks_with()."
1355 )
1356 indices: List[int] = []
1357 weights: List[torch.Tensor] = []
1358 for idx, block in matching:
1359 resolved = self._resolve_submodule_name(block, submodule)
1360 w = _resolve_attr_path(
1361 block, self._rewrite_submodule_path(attr_path, submodule, resolved)
1362 )
1363 if w is None: 1363 ↛ 1364line 1363 didn't jump to line 1364 because the condition on line 1363 was never true
1364 raise AttributeError(
1365 f"blocks[{idx}].{attr_path} is None — this checkpoint has no such "
1366 f"parameter (bias-free projection)."
1367 )
1368 if reshape_fn is not None:
1369 w = reshape_fn(w)
1370 weights.append(w)
1371 indices.append(idx)
1372 return indices, torch.stack(weights, dim=0)
1374 def _stack_block_params(
1375 self, attr_path: str, reshape_fn: Optional[Callable] = None
1376 ) -> torch.Tensor:
1377 """Stack a parameter across all blocks; falls back to matching-only on hybrids.
1379 Filters on FULL-path resolution, not the first segment: on interleaved
1380 MoE, every block has `mlp` but only dense layers expose `mlp.W_in`, so
1381 a first-segment filter matched everything and the sparse layer's
1382 AttributeError killed the accessor for the whole model.
1383 """
1384 first_attr = attr_path.split(".")[0]
1385 all_blocks = self._enumerate_blocks()
1386 matching_blocks: List[Tuple[int, torch.Tensor]] = []
1387 for i, block in all_blocks:
1388 resolved = self._resolve_submodule_name(block, first_attr)
1389 if resolved is None:
1390 continue
1391 try:
1392 weight = _resolve_attr_path(
1393 block, self._rewrite_submodule_path(attr_path, first_attr, resolved)
1394 )
1395 except AttributeError:
1396 continue
1397 if weight is None:
1398 # Bias-free checkpoints (e.g. qkv_bias=False) expose None here;
1399 # stacking would raise an opaque TypeError.
1400 raise AttributeError(
1401 f"blocks[{i}].{attr_path} is None — this checkpoint has no "
1402 f"such parameter (bias-free projection). Bias-free models expose no "
1403 f"stacked {attr_path.rsplit('.', 1)[-1]}; construct zeros explicitly "
1404 f"if your analysis needs one."
1405 )
1406 matching_blocks.append((i, weight))
1408 if len(matching_blocks) == 0:
1409 raise AttributeError(
1410 f"No blocks resolve '{attr_path}'. "
1411 f"Use bridge.blocks_with('{first_attr}') to check availability."
1412 )
1414 if len(matching_blocks) < len(all_blocks):
1415 indices = [i for i, _ in matching_blocks]
1416 logging.warning(
1417 "Hybrid model: only %d/%d blocks resolve '%s'. Returning stacked tensor "
1418 "for layers %s only. Tensor index i corresponds to original layer "
1419 "indices[i], not layer i. For explicit index mapping, use "
1420 "bridge.stack_params_for('%s', '%s').",
1421 len(matching_blocks),
1422 len(all_blocks),
1423 attr_path,
1424 indices,
1425 first_attr,
1426 attr_path,
1427 )
1429 weights: List[torch.Tensor] = []
1430 for _, weight in matching_blocks:
1431 w = reshape_fn(weight) if reshape_fn is not None else weight
1432 weights.append(w)
1433 # Under a device_map split, per-block tensors live on different devices.
1434 # torch.stack requires a common device; gather onto cfg.device (the embedding /
1435 # input device — a natural "home" for cross-layer reductions).
1436 if getattr(self.cfg, "n_devices", 1) > 1 and weights and self.cfg.device is not None:
1437 target_device = torch.device(self.cfg.device)
1438 weights = [w.to(target_device) for w in weights]
1439 return torch.stack(weights, dim=0)
1441 def _reshape_qkv(self, w: torch.Tensor) -> torch.Tensor:
1442 """Reshape 2D [d_model, d_model] QKV weight to 3D [n_heads, d_model, d_head]."""
1443 if w.shape == (self.cfg.d_model, self.cfg.d_model): 1443 ↛ 1444line 1443 didn't jump to line 1444 because the condition on line 1443 was never true
1444 d_head = self.cfg.d_model // self.cfg.n_heads
1445 return w.reshape(self.cfg.n_heads, self.cfg.d_model, d_head)
1446 return w
1448 def _reshape_o(self, w: torch.Tensor) -> torch.Tensor:
1449 """Reshape 2D [d_model, d_model] O weight to 3D [n_heads, d_head, d_model]."""
1450 if w.shape == (self.cfg.d_model, self.cfg.d_model): 1450 ↛ 1451line 1450 didn't jump to line 1451 because the condition on line 1450 was never true
1451 d_head = self.cfg.d_model // self.cfg.n_heads
1452 return w.reshape(self.cfg.n_heads, d_head, self.cfg.d_model)
1453 return w
1455 def _expand_kv_heads(self, w: torch.Tensor) -> torch.Tensor:
1456 """Expand stacked grouped K/V weights along the head axis to n_heads.
1458 GQA models store one K/V projection per key-value head while W_Q/W_O are
1459 per-query-head, so weight circuits must repeat the grouped K/V up to
1460 n_heads before factoring: query head h reads kv head
1461 h // (n_heads // n_kv_heads), i.e. repeat_interleave — the same layout
1462 GroupedQueryAttention.W_K/W_V expose on HookedTransformer. No-op for MHA,
1463 where the head axes already match.
1464 """
1465 if w.ndim != 4 or w.shape[1] == self.cfg.n_heads:
1466 return w
1467 n_kv_heads = w.shape[1]
1468 if self.cfg.n_heads % n_kv_heads != 0:
1469 raise ValueError(
1470 f"Cannot expand {n_kv_heads} key-value heads to {self.cfg.n_heads} "
1471 f"query heads: n_heads must be a multiple of n_kv_heads."
1472 )
1473 return w.repeat_interleave(self.cfg.n_heads // n_kv_heads, dim=1)
1475 @property
1476 def W_K(self) -> torch.Tensor:
1477 """Stack the key weights across all layers."""
1478 return self._stack_block_params("attn.W_K", self._reshape_qkv)
1480 @property
1481 def W_Q(self) -> torch.Tensor:
1482 """Stack the query weights across all layers."""
1483 return self._stack_block_params("attn.W_Q", self._reshape_qkv)
1485 @property
1486 def W_V(self) -> torch.Tensor:
1487 """Stack the value weights across all layers."""
1488 return self._stack_block_params("attn.W_V", self._reshape_qkv)
1490 @property
1491 def W_O(self) -> torch.Tensor:
1492 """Stack the attn output weights across all layers."""
1493 return self._stack_block_params("attn.W_O", self._reshape_o)
1495 @property
1496 def W_in(self) -> torch.Tensor:
1497 """Stack the MLP input weights across all layers."""
1498 return self._stack_block_params("mlp.W_in")
1500 @property
1501 def W_gate(self) -> Union[torch.Tensor, None]:
1502 """Stack the MLP gate weights across all layers (gated MLPs only)."""
1503 if getattr(self.cfg, "gated_mlp", False):
1504 return self._stack_block_params("mlp.W_gate")
1505 return None
1507 @property
1508 def W_out(self) -> torch.Tensor:
1509 """Stack the MLP output weights across all layers."""
1510 return self._stack_block_params("mlp.W_out")
1512 @property
1513 def b_K(self) -> torch.Tensor:
1514 """Stack the key biases across all layers."""
1515 return self._stack_block_params("attn.b_K")
1517 @property
1518 def b_Q(self) -> torch.Tensor:
1519 """Stack the query biases across all layers."""
1520 return self._stack_block_params("attn.b_Q")
1522 @property
1523 def b_V(self) -> torch.Tensor:
1524 """Stack the value biases across all layers."""
1525 return self._stack_block_params("attn.b_V")
1527 @property
1528 def b_O(self) -> torch.Tensor:
1529 """Stack the attn output biases across all layers."""
1530 return self._stack_block_params("attn.b_O")
1532 @property
1533 def b_in(self) -> torch.Tensor:
1534 """Stack the MLP input biases across all layers."""
1535 return self._stack_block_params("mlp.b_in")
1537 @property
1538 def b_out(self) -> torch.Tensor:
1539 """Stack the MLP output biases across all layers."""
1540 return self._stack_block_params("mlp.b_out")
1542 @property
1543 def W_U(self) -> torch.Tensor:
1544 """Unembedding matrix (d_model, d_vocab). Maps residual stream to logits."""
1545 return self.unembed.W_U
1547 @property
1548 def b_U(self) -> torch.Tensor:
1549 """Unembedding bias (d_vocab)."""
1550 return self.unembed.b_U
1552 @property
1553 def W_E(self) -> torch.Tensor:
1554 """Token embedding matrix (d_vocab, d_model)."""
1555 return self.embed.W_E
1557 @property
1558 def W_pos(self) -> torch.Tensor:
1559 """Positional embedding matrix (n_ctx, d_model).
1561 Only defined for models with learned absolute positional embeddings;
1562 rotary/ALiBi models have no such matrix. Reflects the weights as
1563 currently processed, like every other accessor — equal to
1564 ``HookedTransformer.W_pos`` only under matching processing
1565 (``from_pretrained_no_processing`` vs the bridge default, or
1566 compatibility mode vs HT defaults).
1567 """
1568 pos_embed = getattr(self, "pos_embed", None)
1569 pos_type = getattr(self.cfg, "positional_embedding_type", "unknown")
1570 if pos_embed is None or not hasattr(pos_embed, "W_pos"):
1571 raise AttributeError(
1572 "W_pos is only defined for models with a learned absolute positional "
1573 f"embedding component; this model exposes none "
1574 f"(positional_embedding_type={pos_type!r})."
1575 )
1576 w = pos_embed.W_pos
1577 # T5-family adapters map pos_embed to the relative attention bias — a
1578 # [num_buckets, n_heads] table, not a positional matrix. The legacy
1579 # accessor refused these models; keep refusing rather than hand back a
1580 # wrong-shaped tensor.
1581 if w.ndim != 2 or w.shape[-1] != self.cfg.d_model:
1582 raise AttributeError(
1583 "W_pos is only defined for learned absolute positional embeddings; "
1584 f"this model's pos_embed holds a {tuple(w.shape)} table "
1585 f"(positional_embedding_type={pos_type!r})."
1586 )
1587 # OPT/BART-family checkpoints allocate max_position_embeddings + 2 rows
1588 # (positions are offset by 2 in HF's forward). HookedTransformer's
1589 # converters slice those rows off; mirror that so shapes and values agree.
1590 if w.shape[0] == self.cfg.n_ctx + 2:
1591 return w[2:]
1592 return w
1594 @property
1595 def W_E_pos(self) -> torch.Tensor:
1596 """Concatenated ``[W_E; W_pos]`` (d_vocab + n_ctx, d_model).
1598 A full (overcomplete) basis of the input space, used for full QK/OV
1599 circuits. Mirrors ``HookedTransformer.W_E_pos``.
1600 """
1601 return torch.cat([self.W_E, self.W_pos], dim=0)
1603 @property
1604 def QK(self):
1605 """QK circuit. On hybrids, returns attn layers only (with warning). See QK_for_attn_layers()."""
1606 return FactoredMatrix(self.W_Q, self._expand_kv_heads(self.W_K).transpose(-2, -1))
1608 @property
1609 def OV(self):
1610 """OV circuit. On hybrids, returns attn layers only (with warning). See OV_for_attn_layers()."""
1611 return FactoredMatrix(self._expand_kv_heads(self.W_V), self.W_O)
1613 def QK_for_attn_layers(self) -> Tuple[List[int], FactoredMatrix]:
1614 """QK circuit for attention layers only. Returns (layer_indices, FactoredMatrix)."""
1615 q_indices, W_Q = self.stack_params_for("attn", "attn.W_Q", self._reshape_qkv)
1616 _, W_K = self.stack_params_for("attn", "attn.W_K", self._reshape_qkv)
1617 return q_indices, FactoredMatrix(W_Q, self._expand_kv_heads(W_K).transpose(-2, -1))
1619 def OV_for_attn_layers(self) -> Tuple[List[int], FactoredMatrix]:
1620 """OV circuit for attention layers only. Returns (layer_indices, FactoredMatrix)."""
1621 v_indices, W_V = self.stack_params_for("attn", "attn.W_V", self._reshape_qkv)
1622 _, W_O = self.stack_params_for("attn", "attn.W_O", self._reshape_o)
1623 return v_indices, FactoredMatrix(self._expand_kv_heads(W_V), W_O)
1625 # ------------------------------------------------------------------
1626 # Mechanistic interpretability analysis methods
1627 # ------------------------------------------------------------------
1629 def tokens_to_residual_directions(
1630 self,
1631 tokens: Union[str, int, torch.Tensor],
1632 ) -> torch.Tensor:
1633 """Map tokens to their unembedding vectors (residual stream directions).
1635 Returns the columns of W_U corresponding to the given tokens — i.e. the
1636 directions in the residual stream that the model dots with to produce the
1637 logit for each token.
1639 WARNING: If you use this without folding in LayerNorm (compatibility mode),
1640 the results will be misleading because LN weights change the unembed map.
1642 Args:
1643 tokens: A single token (str, int, or scalar tensor), a 1-D tensor of
1644 token IDs, or a 2-D batch of token IDs.
1646 Returns:
1647 Tensor of unembedding vectors with shape matching the input token shape
1648 plus a trailing d_model dimension.
1649 """
1650 if isinstance(tokens, torch.Tensor) and tokens.numel() > 1:
1651 residual_directions = self.W_U[:, tokens]
1652 residual_directions = einops.rearrange(
1653 residual_directions, "d_model ... -> ... d_model"
1654 )
1655 return residual_directions
1656 else:
1657 if isinstance(tokens, str):
1658 token = self.to_single_token(tokens)
1659 elif isinstance(tokens, int):
1660 token = tokens
1661 elif isinstance(tokens, torch.Tensor) and tokens.numel() == 1: 1661 ↛ 1664line 1661 didn't jump to line 1664 because the condition on line 1661 was always true
1662 token = int(tokens.item())
1663 else:
1664 raise ValueError(f"Invalid token type: {type(tokens)}")
1665 residual_direction = self.W_U[:, token]
1666 return residual_direction
1668 # Variant → attr paths for the output bias that feeds the residual stream.
1669 _VARIANT_OUTPUT_BIAS_ATTRS: Dict[str, tuple] = {
1670 "attn": ("b_O",),
1671 "linear_attn": ("out_proj.bias",),
1672 "mamba": ("out_proj.bias",),
1673 "mixer": ("out_proj.bias",),
1674 "ssm": ("out_proj.bias",),
1675 }
1677 def _get_block_variant_bias(self, block: "GeneralizedComponent") -> Optional[torch.Tensor]:
1678 """Return the output bias from this block's variant submodule, or None."""
1679 for name in VARIANT_SUBMODULE_NAMES:
1680 if name not in block._modules:
1681 continue
1682 variant = block._modules[name]
1683 for attr_path in self._VARIANT_OUTPUT_BIAS_ATTRS.get(name, ()): 1683 ↛ 1679line 1683 didn't jump to line 1679 because the loop on line 1683 didn't complete
1684 obj = variant
1685 try:
1686 for attr in attr_path.split("."):
1687 obj = getattr(obj, attr)
1688 except AttributeError:
1689 continue
1690 if obj is not None and isinstance(obj, torch.Tensor): 1690 ↛ 1683line 1690 didn't jump to line 1683 because the condition on line 1690 was always true
1691 return obj
1692 return None
1694 def accumulated_bias(
1695 self,
1696 layer: int,
1697 mlp_input: bool = False,
1698 include_mlp_biases: bool = True,
1699 ) -> torch.Tensor:
1700 """Sum of variant + MLP output biases through the residual stream up to `layer`.
1702 Includes all layer types (attn, SSM, linear-attn). Set mlp_input=True
1703 to include the variant bias of the target layer itself.
1704 """
1705 accumulated = torch.zeros(self.cfg.d_model, device=self.cfg.device)
1706 for i in range(layer):
1707 block = self.blocks[i]
1708 b_O = self._get_block_variant_bias(block)
1709 if b_O is not None:
1710 accumulated = accumulated + b_O.to(accumulated.device)
1711 if include_mlp_biases and "mlp" in block._modules:
1712 b_out = getattr(block.mlp, "b_out", None)
1713 if b_out is not None: 1713 ↛ 1706line 1713 didn't jump to line 1706 because the condition on line 1713 was always true
1714 accumulated = accumulated + b_out.to(accumulated.device)
1715 if mlp_input:
1716 assert layer < self.cfg.n_layers, "Cannot include attn_bias from beyond the final layer"
1717 block = self.blocks[layer]
1718 b_O = self._get_block_variant_bias(block)
1719 if b_O is not None:
1720 accumulated = accumulated + b_O.to(accumulated.device)
1721 return accumulated
1723 def all_composition_scores(self, mode: str) -> CompositionScores:
1724 """Composition scores for all attention head pairs. Returns CompositionScores.
1726 See https://transformer-circuits.pub/2021/framework/index.html
1727 On hybrid models, only attention layers are included; layer_indices
1728 maps tensor position i to original layer number.
1729 """
1730 self._reject_encoder_decoder_composition()
1731 attn_blocks = self.blocks_with("attn")
1732 if not attn_blocks: 1732 ↛ 1733line 1732 didn't jump to line 1733 because the condition on line 1732 was never true
1733 raise ValueError("No attention layers found — cannot compute composition scores.")
1735 indices = [idx for idx, _ in attn_blocks]
1736 blocks_list = [block for _, block in attn_blocks]
1738 def _stack(attr_path: str, reshape_fn: Optional[Callable] = None) -> torch.Tensor:
1739 weights: List[torch.Tensor] = []
1740 for block_idx, block in zip(indices, blocks_list):
1741 w = _resolve_attr_path(block, attr_path)
1742 if w is None: 1742 ↛ 1743line 1742 didn't jump to line 1743 because the condition on line 1742 was never true
1743 raise AttributeError(
1744 f"blocks[{block_idx}].{attr_path} is None — this checkpoint has "
1745 f"no such parameter (bias-free projection)."
1746 )
1747 if reshape_fn is not None: 1747 ↛ 1749line 1747 didn't jump to line 1749 because the condition on line 1747 was always true
1748 w = reshape_fn(w)
1749 weights.append(w)
1750 # See _stack_block_params: gather per-block tensors onto cfg.device when split.
1751 if getattr(self.cfg, "n_devices", 1) > 1 and weights and self.cfg.device is not None: 1751 ↛ 1752line 1751 didn't jump to line 1752 because the condition on line 1751 was never true
1752 target_device = torch.device(self.cfg.device)
1753 weights = [w.to(target_device) for w in weights]
1754 return torch.stack(weights, dim=0)
1756 W_V = self._expand_kv_heads(_stack("attn.W_V", self._reshape_qkv))
1757 W_O = _stack("attn.W_O", self._reshape_o)
1758 left = FactoredMatrix(W_V, W_O)
1760 if mode == "Q":
1761 W_Q = _stack("attn.W_Q", self._reshape_qkv)
1762 W_K = self._expand_kv_heads(_stack("attn.W_K", self._reshape_qkv))
1763 right = FactoredMatrix(W_Q, W_K.transpose(-2, -1))
1764 elif mode == "K":
1765 W_Q = _stack("attn.W_Q", self._reshape_qkv)
1766 W_K = self._expand_kv_heads(_stack("attn.W_K", self._reshape_qkv))
1767 right = FactoredMatrix(W_Q, W_K.transpose(-2, -1)).T
1768 elif mode == "V":
1769 right = left
1770 else:
1771 raise ValueError(f"mode must be one of ['Q', 'K', 'V'] not {mode}")
1773 scores = utils.composition_scores(left, right, broadcast_dims=True)
1774 n_attn = len(indices)
1775 idx_tensor = torch.arange(n_attn, device=self.cfg.device)
1776 mask = idx_tensor[:, None, None, None] < idx_tensor[None, None, :, None]
1777 scores = torch.where(mask, scores, torch.zeros_like(scores))
1779 labels = [f"L{l}H{h}" for l in indices for h in range(self.cfg.n_heads)]
1780 return CompositionScores(scores=scores, layer_indices=indices, head_labels=labels)
1782 def _reject_encoder_decoder_composition(self) -> None:
1783 """Composition scores live in one residual stream; enc-dec models have two."""
1784 if any(hasattr(self, a) for a in ("encoder_blocks", "decoder_blocks")):
1785 raise NotImplementedError(
1786 "Composition scores are not defined across an encoder-decoder "
1787 "model's two residual streams (HookedTransformer never "
1788 "supported them there either)."
1789 )
1791 def composition_layer_indices(self) -> List[int]:
1792 """Original layer indices for attention layers (maps composition score positions)."""
1793 self._reject_encoder_decoder_composition()
1794 return [idx for idx, _ in self.blocks_with("attn")]
1796 def block_hooks(self, layer_idx: int) -> List[str]:
1797 """Sorted hook names available on block `layer_idx` (block-relative paths)."""
1798 prefix = f"blocks.{layer_idx}."
1799 return sorted(name[len(prefix) :] for name in self.hook_dict if name.startswith(prefix))
1801 def block_submodules(self, layer_idx: int) -> List[str]:
1802 """Return bridged submodule names on block `layer_idx`."""
1803 block = self.blocks[layer_idx]
1804 return [name for name in block._modules if name not in _BLOCK_INTERNAL_MODULES]
1806 def layer_types(self) -> List[str]:
1807 """Per-block type labels, e.g. ["attn+mlp", "ssm+mlp", ...]. Deterministic order."""
1808 types = []
1809 for block in self.blocks:
1810 variants = [n for n in VARIANT_SUBMODULE_NAMES if n in block._modules]
1811 universals = sorted(
1812 n
1813 for n in block._modules
1814 if n not in _VARIANT_SUBMODULE_SET
1815 and n not in _BLOCK_INTERNAL_MODULES
1816 and not n.startswith(_NORM_PREFIXES)
1817 )
1818 parts = variants + universals
1819 types.append("+".join(parts) if parts else "unknown")
1820 return types
1822 @property
1823 def all_head_labels(self) -> list[str]:
1824 """Human-readable labels for all attention heads, e.g. ['L0H0', 'L0H1', ...].
1826 Encoder-decoder models use ``HookedEncoderDecoder``'s ``EL{l}H{h}`` /
1827 ``DL{l}H{h}`` scheme so encoder and decoder heads stay distinguishable;
1828 a plain ``L{l}H{h}`` list would name only half of them.
1829 """
1830 encoder_blocks = self._modules.get("encoder_blocks")
1831 decoder_blocks = self._modules.get("decoder_blocks")
1832 if isinstance(encoder_blocks, nn.ModuleList) and isinstance(decoder_blocks, nn.ModuleList):
1833 heads = range(self.cfg.n_heads)
1834 return [f"EL{l}H{h}" for l in range(len(encoder_blocks)) for h in heads] + [
1835 f"DL{l}H{h}" for l in range(len(decoder_blocks)) for h in heads
1836 ]
1837 return [f"L{l}H{h}" for l in range(self.cfg.n_layers) for h in range(self.cfg.n_heads)]
1839 @property
1840 def attn_head_labels(self) -> list[str]:
1841 """Head labels for attention layers only — matches all_composition_scores() dims."""
1842 self._reject_encoder_decoder_composition()
1843 return [
1844 f"L{l}H{h}" for l in self.composition_layer_indices() for h in range(self.cfg.n_heads)
1845 ]
1847 def parameters(self, recurse: bool = True) -> Iterator[nn.Parameter]:
1848 """Returns parameters following standard PyTorch semantics.
1850 This method delegates to the underlying HuggingFace model's parameters().
1851 For TransformerLens-style parameter generator, use tl_parameters() instead.
1853 Args:
1854 recurse: If True, yields parameters of this module and all submodules
1856 Returns:
1857 Iterator of nn.Parameter objects
1858 """
1859 return self.original_model.parameters(recurse=recurse)
1861 def named_parameters(
1862 self, prefix: str = "", recurse: bool = True, remove_duplicate: bool = True
1863 ) -> Iterator[tuple[str, nn.Parameter]]:
1864 """Returns named parameters following standard PyTorch semantics.
1866 This method delegates to the underlying HuggingFace model's named_parameters().
1867 For TransformerLens-style generator, use tl_named_parameters() instead.
1869 Args:
1870 prefix: Prefix to prepend to all parameter names
1871 recurse: If True, yields parameters of this module and all submodules
1872 remove_duplicate: If True, removes duplicate parameters
1874 Returns:
1875 Iterator of (name, parameter) tuples
1876 """
1877 return self.original_model.named_parameters(prefix, recurse, remove_duplicate)
1879 def tl_parameters(self) -> dict[str, torch.Tensor]:
1880 """Returns TransformerLens-style parameter dictionary.
1882 Parameter names follow TransformerLens conventions (e.g., 'blocks.0.attn.W_Q') and may
1883 include processed weights (non-leaf tensors). This format is expected by SVDInterpreter
1884 among other analysis tools.
1886 Returns:
1887 Dictionary mapping TransformerLens parameter names to tensors
1889 Example:
1890 >>> bridge = TransformerBridge.boot_transformers("gpt2")
1891 >>> tl_params = bridge.tl_parameters()
1892 >>> W_Q = tl_params["blocks.0.attn.W_Q"] # Shape: [n_heads, d_model, d_head]
1893 """
1894 return self.get_params()
1896 def tl_named_parameters(self) -> Iterator[tuple[str, torch.Tensor]]:
1897 """Returns iterator of TransformerLens-style named parameters.
1899 This provides the same parameters as tl_parameters() but as an iterator
1900 for consistency with PyTorch's named_parameters() API pattern.
1902 Returns:
1903 Iterator of (name, tensor) tuples with TransformerLens naming conventions
1905 Example:
1906 >>> bridge = TransformerBridge.boot_transformers("gpt2")
1907 >>> for name, param in bridge.tl_named_parameters():
1908 ... if "attn.W_Q" in name:
1909 ... print(f"{name}: {param.shape}") # doctest: +ELLIPSIS
1910 blocks.0.attn.W_Q: torch.Size([12, 768, 64])
1911 ...
1912 """
1913 return iter(self.get_params().items())
1915 def input_to_embed(
1916 self,
1917 input: Union[str, List[str], torch.Tensor],
1918 prepend_bos: Optional[bool] = None,
1919 padding_side: Optional[str] = None,
1920 attention_mask: Optional[torch.Tensor] = None,
1921 ) -> Tuple[torch.Tensor, torch.Tensor, None, Optional[torch.Tensor]]:
1922 """Convert input to the residual stream entering block 0 (``resid_pre[0]``).
1924 Bridge analog of :meth:`HookedTransformer.input_to_embed`. Returns
1925 ``(residual, tokens, shortformer_pos_embed, attention_mask)``; feed the
1926 residual to ``forward(..., start_at_layer=0)`` to resume the pass.
1928 ``shortformer_pos_embed`` is always ``None``: for the models the bridge
1929 supports the residual already carries positional information, so there is
1930 no separate positional stream to return.
1931 """
1932 if isinstance(input, (str, list)):
1933 assert self.tokenizer is not None, "Must provide a tokenizer for string input."
1934 tokens = self.to_tokens(input, prepend_bos=prepend_bos, padding_side=padding_side)
1935 else:
1936 tokens = input
1937 if isinstance(tokens, torch.Tensor) and tokens.ndim == 1: 1937 ↛ 1938line 1937 didn't jump to line 1938 because the condition on line 1937 was never true
1938 tokens = tokens.unsqueeze(0)
1939 if ( 1939 ↛ 1944line 1939 didn't jump to line 1944 because the condition on line 1939 was never true
1940 attention_mask is None
1941 and self.tokenizer is not None
1942 and self.tokenizer.padding_side == "left"
1943 ):
1944 _prepend = self.cfg.default_prepend_bos if prepend_bos is None else prepend_bos
1945 attention_mask = utils.get_attention_mask(self.tokenizer, tokens, _prepend)
1946 residual = self.forward(tokens, stop_at_layer=0, attention_mask=attention_mask)
1947 return residual, tokens, None, attention_mask
1949 def _accepts_derived_position_ids(self) -> bool:
1950 """Whether it is safe to hand the wrapped model a mask-derived ``position_ids``.
1952 Two families of model must be left alone, so the injection below is
1953 gated on the target the same way ``output_attentions`` is in
1954 :meth:`BridgeCore.run_with_cache`:
1956 * **Fixed-signature models.** Remote-code forwards such as
1957 ``LLaDAModelLM.forward`` take neither ``position_ids`` nor
1958 ``**kwargs``, so passing it raises ``TypeError`` where the model
1959 previously returned logits.
1960 * **Models that own their position derivation.** mRoPE architectures
1961 (Qwen2-VL, Qwen2.5-VL, Qwen3-VL, GLM-4V) build a 3-D temporal /
1962 height / width index in ``get_rope_index``, and only while
1963 ``position_ids is None``; a supplied 2-D tensor is silently expanded
1964 across all three streams instead. Their derivation already scatters
1965 positions onto attended slots only, so it handles left padding
1966 correctly on its own and needs no help from us.
1967 * **Mask-consuming positional embeddings.** OPT's
1968 ``OPTLearnedPositionalEmbedding.forward`` takes the mask and derives
1969 the same positions we would, so injection buys nothing — but it does
1970 replace the model's own padding-slot convention with ours, which
1971 shows up as a whole-tensor diff.
1973 Non-torch drivers (vLLM, Inspect) expose no module to introspect and
1974 manage positions internally, so they are excluded as well.
1975 """
1976 underlying = getattr(self._driver, "underlying_model", None)
1977 if underlying is None:
1978 return False
1980 cached = self.__dict__.get("_derived_position_ids_ok")
1981 if cached is not None and cached[0] is underlying:
1982 return bool(cached[1])
1984 accepts = accepts_mask_derived_position_ids(underlying)
1985 self.__dict__["_derived_position_ids_ok"] = (underlying, accepts)
1986 return accepts
1988 @staticmethod
1989 def _seq2seq_loss(
1990 logits: torch.Tensor,
1991 labels: torch.Tensor,
1992 native_loss: Any,
1993 *,
1994 per_token: bool,
1995 ) -> torch.Tensor:
1996 """Return encoder-decoder loss without the causal LM token shift."""
1997 if labels.device != logits.device: 1997 ↛ 1998line 1997 didn't jump to line 1998 because the condition on line 1997 was never true
1998 labels = labels.to(logits.device)
1999 if labels.shape != logits.shape[:-1]: 1999 ↛ 2000line 1999 didn't jump to line 2000 because the condition on line 1999 was never true
2000 raise ValueError(
2001 "seq2seq labels must match the decoder logits batch and position "
2002 f"dimensions, got labels {tuple(labels.shape)} and logits "
2003 f"{tuple(logits.shape)}"
2004 )
2005 if not per_token and isinstance(native_loss, torch.Tensor):
2006 return native_loss
2008 losses = F.cross_entropy(
2009 logits.flatten(0, 1),
2010 labels.flatten(),
2011 reduction="none" if per_token else "mean",
2012 ignore_index=-100,
2013 )
2014 return losses.view_as(labels) if per_token else losses
2016 def _finalize_seq2seq_return(
2017 self,
2018 return_type: str,
2019 logits: torch.Tensor,
2020 labels: torch.Tensor,
2021 native_output: Any,
2022 *,
2023 loss_per_token: bool,
2024 ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
2025 loss = self._seq2seq_loss(
2026 logits,
2027 labels,
2028 getattr(native_output, "loss", None),
2029 per_token=loss_per_token,
2030 )
2031 return (logits, loss) if return_type == "both" else loss
2033 def forward(
2034 self,
2035 input: Union[str, List[str], torch.Tensor],
2036 return_type: Optional[str] = "logits",
2037 loss_per_token: bool = False,
2038 prepend_bos: Optional[bool] = None,
2039 padding_side: Optional[str] = None,
2040 attention_mask: Optional[torch.Tensor] = None,
2041 labels: Optional[torch.Tensor] = None,
2042 start_at_layer: Optional[int] = None,
2043 stop_at_layer: Optional[int] = None,
2044 pixel_values: Optional[torch.Tensor] = None,
2045 input_values: Optional[torch.Tensor] = None,
2046 past_key_values: Optional[Any] = None,
2047 **kwargs,
2048 ) -> Any:
2049 """Forward pass through the model.
2051 Args:
2052 input: Input to the model
2053 return_type: Type of output to return ('logits', 'loss', 'both', 'predictions',
2054 'logits_and_cache', None). 'logits_and_cache' returns
2055 ``(logits, past_key_values)`` — the HuggingFace cache after this step, to
2056 feed back on the next call for incremental decoding.
2057 loss_per_token: Whether to return loss per token
2058 prepend_bos: Whether to prepend BOS token
2059 padding_side: Which side to pad on
2060 labels: Explicit language-model targets. Encoder-decoder models require
2061 labels for loss; decoder-only models fall back to input IDs when omitted.
2062 past_key_values: HuggingFace KV cache from a prior ``use_cache=True`` step
2063 (e.g. the second element of a ``return_type='logits_and_cache'`` return).
2064 When provided, KV caching is enabled automatically and only the new
2065 tokens' keys/values are computed; HF derives the position offset from the
2066 cache length. This is the bridge's manual KV-cache entry point — it uses
2067 HF's native cache object rather than a TransformerLens cache.
2068 start_at_layer: Resume the forward from block ``k``, treating ``input`` as
2069 the residual-stream tensor ``[batch, pos, d_model]`` entering that block
2070 (mirrors HookedTransformer). Blocks 0..k-1 still execute internally (their
2071 output is discarded when block k swaps in the residual) but are excluded
2072 from ``run_with_cache`` output. Requires an HF model that accepts
2073 ``inputs_embeds``; only supported on the standard ``blocks`` stack.
2074 stop_at_layer: Layer to stop forward pass at. Only supported on the
2075 standard ``blocks`` stack; architectures that register no ``blocks``
2076 (e.g. Raven's ``prelude``/``core_block``/``coda``) raise
2077 ``NotImplementedError`` rather than running to completion.
2078 pixel_values: Optional image tensor for multimodal models (e.g., LLaVA, Gemma3)
2079 and vision models (eg. ViT, DeiT).
2080 The tensor is passed directly to the underlying HuggingFace model.
2081 Only valid when cfg.is_multimodal is True or cfg.is_visual_model is True.
2082 input_values: Optional audio waveform tensor for audio models (e.g., HuBERT).
2083 The tensor is passed directly to the underlying HuggingFace model.
2084 Only valid when cfg.is_audio_model is True.
2085 **kwargs: Additional arguments passed to model
2087 Returns:
2088 Model output based on return_type
2089 """
2091 underlying_model = getattr(getattr(self, "_driver", None), "underlying_model", None)
2092 model_config = getattr(underlying_model, "config", None)
2093 is_encoder_decoder = bool(getattr(model_config, "is_encoder_decoder", False))
2094 if return_type in ("loss", "both"):
2095 if is_encoder_decoder and labels is None:
2096 raise ValueError(
2097 "labels are required for seq2seq return_type='loss' or 'both'; "
2098 "encoder input_ids are not decoder targets"
2099 )
2100 if (
2101 return_type in ("loss", "both")
2102 and not is_encoder_decoder
2103 and not self.adapter.supports_causal_loss
2104 ):
2105 architecture = self.cfg.architecture or type(self.adapter).__name__
2106 raise NotImplementedError(
2107 f"{architecture} does not support TransformerBridge's shifted causal "
2108 "loss. Request return_type='logits' and compute the architecture-specific "
2109 "masked-token objective explicitly."
2110 )
2112 if labels is not None:
2113 kwargs["labels"] = labels
2115 if start_at_layer is not None:
2116 input = self._setup_start_at_layer(input, start_at_layer)
2117 # Set stop_at_layer flag on all blocks if requested
2118 if stop_at_layer is not None:
2119 if not self._has_registered_blocks():
2120 raise NotImplementedError(
2121 "stop_at_layer requires a 'blocks' stack; this architecture "
2122 "does not register one."
2123 )
2124 effective_stop_at_layer = (
2125 len(self.blocks) + stop_at_layer if stop_at_layer < 0 else stop_at_layer
2126 )
2127 for block in self.blocks:
2128 block._stop_at_layer_idx = effective_stop_at_layer
2130 # Map HookedEncoderDecoder-style kwargs to HF-compatible names
2131 if "decoder_input" in kwargs:
2132 kwargs["decoder_input_ids"] = kwargs.pop("decoder_input")
2133 if "one_zero_attention_mask" in kwargs: 2133 ↛ 2134line 2133 didn't jump to line 2134 because the condition on line 2133 was never true
2134 if attention_mask is None:
2135 attention_mask = kwargs.pop("one_zero_attention_mask")
2136 else:
2137 kwargs.pop("one_zero_attention_mask")
2139 # Detect batched list input that may need padding. Forward follows the
2140 # requested/tokenizer side; generation separately forces left-padding.
2141 _is_batched_list = (
2142 isinstance(input, list)
2143 and len(input) > 1
2144 and not getattr(self.cfg, "is_audio_model", False)
2145 and not getattr(self.cfg, "is_visual_model", False)
2146 )
2147 _resolved_padding_side = padding_side
2148 if _resolved_padding_side is None and self.tokenizer is not None:
2149 _resolved_padding_side = getattr(self.tokenizer, "padding_side", "right")
2151 try:
2152 if isinstance(input, (str, list)):
2153 if getattr(self.cfg, "is_audio_model", False): 2153 ↛ 2154line 2153 didn't jump to line 2154 because the condition on line 2153 was never true
2154 raise ValueError(
2155 "Audio models require tensor input (raw waveform), not text. "
2156 "Pass a torch.Tensor or use the input_values parameter."
2157 )
2158 if getattr(self.cfg, "is_visual_model", False): 2158 ↛ 2159line 2158 didn't jump to line 2159 because the condition on line 2158 was never true
2159 raise ValueError(
2160 "Visual models require tensor input (pixel values), not text. "
2161 "Pass a torch.Tensor or use the pixel_values parameter."
2162 )
2163 input_ids = self.to_tokens(
2164 input, prepend_bos=prepend_bos, padding_side=padding_side
2165 )
2166 else:
2167 input_ids = input
2168 # Promote 1D integer token tensors to 2D [batch=1, seq] to match
2169 # HookedTransformer's contract. Float tensors (inputs_embeds,
2170 # audio waveforms) are passed through unchanged.
2171 if (
2172 isinstance(input_ids, torch.Tensor)
2173 and input_ids.ndim == 1
2174 and not input_ids.is_floating_point()
2175 ):
2176 input_ids = input_ids.unsqueeze(0)
2178 # Detect inputs_embeds: if the tensor is floating point, it's pre-computed
2179 # embeddings (e.g., from multimodal models) rather than token IDs.
2180 _is_inputs_embeds = (
2181 isinstance(input_ids, torch.Tensor) and input_ids.is_floating_point()
2182 )
2184 # Left padding needs a mask and corrected positions. Right padding is
2185 # harmless for causal real-token positions and remains unmasked to
2186 # match HookedTransformer; bidirectional/encoder inputs still need it.
2187 if (
2188 _is_batched_list
2189 and attention_mask is None
2190 and self.tokenizer is not None
2191 and self.tokenizer.pad_token_id is not None
2192 and not _is_inputs_embeds
2193 and (
2194 _resolved_padding_side == "left"
2195 or is_encoder_decoder
2196 or not self.adapter.supports_causal_loss
2197 )
2198 ):
2199 attention_mask = utils.get_attention_mask(
2200 self.tokenizer,
2201 input_ids,
2202 prepend_bos=getattr(self.cfg, "default_prepend_bos", True),
2203 padding_side=_resolved_padding_side,
2204 ).to(self.cfg.device)
2205 # Gated on the target for the same reason the derivation below is:
2206 # a fixed-signature forward raises TypeError on the kwarg, and a
2207 # model that owns its own position derivation is overridden by it
2208 # (#1626).
2209 if "position_ids" not in kwargs and self._accepts_derived_position_ids():
2210 position_ids = attention_mask.long().cumsum(-1) - 1
2211 position_ids.masked_fill_(attention_mask == 0, 1)
2212 kwargs["position_ids"] = position_ids
2214 # Any masked-out token shifts the absolute position of every real token
2215 # after it, so positions must be derived from the mask rather than left
2216 # to HF's default arange. This is the same derivation HookedTransformer
2217 # applies in pos_embed; without it the bridge silently returns wrong
2218 # logits. An all-ones mask reduces to arange, so this is a no-op there.
2219 #
2220 # The mask spans any cached prefix as well as the new tokens, so it is
2221 # offset back to just the tokens actually being passed — matching how
2222 # AbstractAttention/PosEmbed use past_kv_pos_offset.
2223 if (
2224 attention_mask is not None
2225 and "position_ids" not in kwargs
2226 and not _is_inputs_embeds
2227 and attention_mask.ndim == 2
2228 and isinstance(input_ids, torch.Tensor)
2229 and input_ids.ndim == 2
2230 and attention_mask.shape[1] >= input_ids.shape[1]
2231 and self._accepts_derived_position_ids()
2232 ):
2233 # .long() because callers may hand in a float 0/1 mask, and
2234 # positions index an embedding table.
2235 _derived = utils.get_offset_position_ids(0, attention_mask.long())
2236 _arange = torch.arange(attention_mask.shape[1], device=_derived.device)
2237 # Decide per row, not per batch. A row only needs the derived
2238 # positions when its mask actually moves one of its attended
2239 # tokens off the default position — i.e. a masked token precedes
2240 # a real one (left padding, or an interior gap). Rows that are
2241 # unpadded or purely right-padded keep arange verbatim, so one
2242 # left-padded row in a batch cannot perturb its neighbours.
2243 _needs = ((_derived != _arange) & (attention_mask != 0)).any(dim=1, keepdim=True)
2244 if bool(_needs.any()):
2245 _positions = torch.where(_needs, _derived, _arange.expand_as(_derived))
2246 kwargs["position_ids"] = _positions[
2247 :, attention_mask.shape[1] - input_ids.shape[1] :
2248 ]
2250 if attention_mask is not None:
2251 kwargs["attention_mask"] = attention_mask
2252 if past_key_values is not None:
2253 # Manual KV-cache injection: hand HF its own cache back and let it
2254 # extend it, computing only the new tokens' keys/values.
2255 kwargs["past_key_values"] = past_key_values
2256 kwargs["use_cache"] = True
2257 if kwargs.pop("use_past_kv_cache", False) or kwargs.get("use_cache", False):
2258 kwargs["use_cache"] = True
2259 # Auto-generate decoder_input_ids for encoder-decoder models
2260 if "decoder_input_ids" not in kwargs and labels is None and is_encoder_decoder:
2261 decoder_start_token_id = getattr(
2262 self.original_model.config, "decoder_start_token_id", None
2263 )
2264 if decoder_start_token_id is not None:
2265 shifted = input_ids[:, :-1]
2266 start_tokens = torch.full(
2267 (input_ids.shape[0], 1),
2268 decoder_start_token_id,
2269 dtype=input_ids.dtype,
2270 device=input_ids.device,
2271 )
2272 kwargs["decoder_input_ids"] = torch.cat([start_tokens, shifted], dim=1)
2273 else:
2274 kwargs["decoder_input_ids"] = input_ids
2276 # Tell PosEmbedBridge to expand batch=1 position_ids to full batch.
2277 if hasattr(self, "pos_embed"):
2278 self.pos_embed._current_batch_size = input_ids.shape[0]
2280 # Handle pixel_values for multimodal and vision models
2281 if pixel_values is not None:
2282 if not (
2283 getattr(self.cfg, "is_multimodal", False)
2284 or getattr(self.cfg, "is_visual_model", False)
2285 ):
2286 raise ValueError(
2287 "pixel_values can only be passed to multimodal or vision models "
2288 "(cfg.is_multimodal or cfg.is_visual_model must be True)"
2289 )
2290 kwargs["pixel_values"] = pixel_values
2292 # Handle input_values for audio models
2293 if input_values is not None: 2293 ↛ 2294line 2293 didn't jump to line 2294 because the condition on line 2293 was never true
2294 if not getattr(self.cfg, "is_audio_model", False):
2295 raise ValueError(
2296 "input_values can only be passed to audio models "
2297 "(cfg.is_audio_model must be True)"
2298 )
2299 kwargs["input_values"] = input_values
2301 # Audio models use input_values (waveform), not input_ids
2302 if getattr(self.cfg, "is_audio_model", False):
2303 if input_values is not None: 2303 ↛ 2304line 2303 didn't jump to line 2304 because the condition on line 2303 was never true
2304 result = self._driver.forward(**kwargs)
2305 elif isinstance(input, torch.Tensor): 2305 ↛ 2309line 2305 didn't jump to line 2309 because the condition on line 2305 was always true
2306 kwargs["input_values"] = input
2307 result = self._driver.forward(**kwargs)
2308 else:
2309 raise ValueError(
2310 "Audio models require tensor input (raw waveform). "
2311 "Pass a torch.Tensor or use input_values parameter."
2312 )
2313 elif getattr(self.cfg, "is_visual_model", False):
2314 # "pixel_values" may already be in kwargs from the "if pixel_values is not None:"
2315 # gate above (explicit pixel_values=... call); otherwise treat `input` itself as
2316 # the image tensor, matching how the audio branch treats `input` as the waveform.
2317 if "pixel_values" not in kwargs: 2317 ↛ 2325line 2317 didn't jump to line 2325 because the condition on line 2317 was always true
2318 if isinstance(input, torch.Tensor): 2318 ↛ 2321line 2318 didn't jump to line 2321 because the condition on line 2318 was always true
2319 kwargs["pixel_values"] = input
2320 else:
2321 raise ValueError(
2322 "Visual models require tensor input (pixel values). "
2323 "Pass a torch.Tensor as `input` or use the pixel_values parameter."
2324 )
2325 result = self._driver.forward(**kwargs)
2326 elif _is_inputs_embeds:
2327 result = self._driver.forward(inputs_embeds=input_ids, **kwargs)
2328 else:
2329 # By keyword so kw-only ``input_ids`` drivers don't TypeError.
2330 result = self._driver.forward(input_ids=input_ids, **kwargs)
2331 output = result.raw_output
2332 # No-op for HF (its hooks already fired); load-bearing for vLLM/Inspect.
2333 if result.captured: 2333 ↛ 2334line 2333 didn't jump to line 2334 because the condition on line 2333 was never true
2334 self._replay_captures(result.captured)
2335 # Convert TensorLike to torch at the boundary; let weird shapes
2336 # (audio/CTC dataclasses, tuple-of-tuples) pass through unchanged —
2337 # downstream return_type branches catch them with specific errors.
2338 logits = result.logits
2339 if isinstance(logits, torch.Tensor): 2339 ↛ 2341line 2339 didn't jump to line 2341 because the condition on line 2339 was always true
2340 pass
2341 elif logits is not None and isinstance(logits, TensorLike):
2342 logits = to_torch(logits)
2343 # Stash only the cache object (not the full output) for generate().
2344 if getattr(self, "_capture_hf_cache", False):
2345 self._last_hf_cache = getattr(output, "past_key_values", None)
2346 if return_type == "logits_and_cache":
2347 past_key_values = getattr(output, "past_key_values", None)
2348 return (logits, past_key_values)
2349 if is_encoder_decoder and return_type in ("loss", "both"):
2350 assert isinstance(
2351 logits, torch.Tensor
2352 ), f"Expected seq2seq logits tensor, got {type(logits)}"
2353 assert isinstance(labels, torch.Tensor)
2354 return self._finalize_seq2seq_return(
2355 return_type,
2356 logits,
2357 labels,
2358 output,
2359 loss_per_token=loss_per_token,
2360 )
2361 return self._finalize_return(
2362 return_type,
2363 logits,
2364 input_ids,
2365 attention_mask=attention_mask,
2366 labels=labels,
2367 is_audio_model=getattr(self.cfg, "is_audio_model", False),
2368 is_visual_model=getattr(self.cfg, "is_visual_model", False),
2369 inputs_embeds_was_used=_is_inputs_embeds,
2370 loss_per_token=loss_per_token,
2371 )
2372 except StopAtLayerException as e:
2373 # Execution stopped at the requested layer
2374 return e.layer_output
2375 finally:
2376 # Clean up state that may be inconsistent after StopAtLayerException
2377 if stop_at_layer is not None:
2378 for bl_name in (
2379 "blocks",
2380 "encoder_blocks",
2381 "decoder_blocks",
2382 "L_blocks",
2383 "H_blocks",
2384 ):
2385 if hasattr(self, bl_name):
2386 for block in getattr(self, bl_name):
2387 block._stop_at_layer_idx = None
2389 # Clear any stale KV cache — layers after the stop point didn't
2390 # execute, so the cache is incomplete and would corrupt subsequent
2391 # generate() calls that expect a full cache.
2392 if hasattr(self, "_last_hf_cache"): 2392 ↛ 2393line 2392 didn't jump to line 2393 because the condition on line 2392 was never true
2393 del self._last_hf_cache
2395 if start_at_layer is not None:
2396 self._teardown_start_at_layer()
2398 def _setup_start_at_layer(self, input: Any, start_at_layer: int) -> torch.Tensor:
2399 """Arm residual-stream injection at block ``start_at_layer``.
2401 ``input`` is the residual entering that block. It is returned (batch-promoted)
2402 to drive the HF forward as ``inputs_embeds`` so position ids / attention mask
2403 are computed for the right sequence length; block ``start_at_layer`` then swaps
2404 it back in for its own input, discarding whatever blocks 0..k-1 produced.
2405 """
2406 if getattr(self.cfg, "is_visual_model", False):
2407 raise NotImplementedError(
2408 "start_at_layer is not supported for vision models: the residual "
2409 "re-entry path assumes token inputs_embeds."
2410 )
2411 if getattr(self.cfg, "is_audio_model", False):
2412 raise NotImplementedError(
2413 "start_at_layer is not supported for audio models: audio encoders "
2414 "process waveforms through convolutional feature extraction before "
2415 "the transformer blocks, making residual-stream injection infeasible."
2416 )
2417 for alt in ("encoder_blocks", "decoder_blocks", "L_blocks", "H_blocks"):
2418 if hasattr(self, alt): 2418 ↛ 2419line 2418 didn't jump to line 2419 because the condition on line 2418 was never true
2419 raise NotImplementedError(
2420 "start_at_layer is only supported on the standard 'blocks' stack, "
2421 f"not {alt!r}."
2422 )
2423 if not self._has_registered_blocks():
2424 raise NotImplementedError("start_at_layer requires a 'blocks' stack.")
2426 if not (isinstance(input, torch.Tensor) and input.is_floating_point()):
2427 raise ValueError(
2428 "start_at_layer requires a residual-stream tensor [batch, pos, d_model]; "
2429 f"got {type(input).__name__}. Capture it from a prior run_with_cache "
2430 "(e.g. cache['blocks.k.hook_in'])."
2431 )
2432 residual = input if input.ndim == 3 else input.unsqueeze(0)
2433 n_blocks = len(self.blocks)
2434 if start_at_layer < 0:
2435 start_at_layer += n_blocks
2436 if not 0 <= start_at_layer < n_blocks:
2437 raise ValueError(f"start_at_layer={start_at_layer} out of range [0, {n_blocks}).")
2438 # The returned tensor is fed to HF as inputs_embeds; stash a clone so an
2439 # architecture that mutates inputs_embeds in place can't corrupt the value
2440 # block ``start_at_layer`` swaps back in.
2441 injected = residual.clone()
2442 for block in self.blocks:
2443 block._start_at_layer_idx = start_at_layer
2444 block._start_residual = injected
2445 return residual
2447 def _teardown_start_at_layer(self) -> None:
2448 """Clear the residual-injection state set by ``_setup_start_at_layer``."""
2449 if hasattr(self, "blocks"): 2449 ↛ exitline 2449 didn't return from function '_teardown_start_at_layer' because the condition on line 2449 was always true
2450 for block in self.blocks:
2451 block._start_at_layer_idx = None
2452 block._start_residual = None
2454 # loss_fn inherited from BridgeCore
2456 def _resolve_stopping_criteria(
2457 self,
2458 stop_strings: Optional[Union[str, List[str]]],
2459 stopping_criteria: Optional[Any],
2460 ) -> Optional[Any]:
2461 """Combine ``stop_strings`` and ``stopping_criteria`` into one StoppingCriteriaList.
2463 Returns ``None`` when neither is supplied (or both reduce to no-ops),
2464 so callers can cheaply check whether any extra stop signal is active.
2465 ``stop_strings`` is turned into a HuggingFace ``StopStringCriteria`` (which reproduces
2466 HF's exact partial-token-aware, end-anchored matching: it fires when the stop string
2467 ends the generated text, even if the string straddles token boundaries) and therefore
2468 requires a tokenizer.
2469 A user-supplied ``stopping_criteria`` may be a single ``StoppingCriteria``,
2470 a list of them, or a ``StoppingCriteriaList``.
2472 Raises:
2473 ValueError: if ``stop_strings`` is supplied without a tokenizer.
2474 TypeError: if ``stopping_criteria`` is not a ``StoppingCriteria``, a
2475 list/tuple of them, or a ``StoppingCriteriaList``.
2476 """
2477 if stop_strings is None and stopping_criteria is None:
2478 return None
2480 from transformers import ( # local import: matches the file's transformers usage
2481 StoppingCriteria,
2482 StoppingCriteriaList,
2483 StopStringCriteria,
2484 )
2486 criteria = StoppingCriteriaList()
2488 if stop_strings is not None:
2489 strings = [stop_strings] if isinstance(stop_strings, str) else list(stop_strings)
2490 strings = [s for s in strings if s] # drop empty strings (HF errors on them)
2491 if strings:
2492 if self.tokenizer is None:
2493 raise ValueError(
2494 "stop_strings requires a tokenizer (stop strings are detected by "
2495 "matching against the tokenizer vocabulary), but this TransformerBridge "
2496 "has no tokenizer. Pass a stopping_criteria callable that operates on "
2497 "token ids instead, or use hf_generate()."
2498 )
2499 criteria.append(StopStringCriteria(tokenizer=self.tokenizer, stop_strings=strings))
2501 if stopping_criteria is not None:
2502 if isinstance(stopping_criteria, StoppingCriteriaList):
2503 criteria.extend(stopping_criteria)
2504 elif isinstance(stopping_criteria, (list, tuple)):
2505 criteria.extend(stopping_criteria)
2506 elif isinstance(stopping_criteria, StoppingCriteria):
2507 criteria.append(stopping_criteria)
2508 else:
2509 raise TypeError(
2510 "stopping_criteria must be a transformers.StoppingCriteria, a list of "
2511 f"them, or a StoppingCriteriaList, but got {type(stopping_criteria).__name__}."
2512 )
2514 return criteria if len(criteria) > 0 else None
2516 def _encdec_ngram_processor(self) -> Optional[Any]:
2517 """generation_config.no_repeat_ngram_size as transformers' own
2518 processor, or None. HF applies it by default; parity for models whose
2519 greedy decode needs it to escape token attractors."""
2520 size = getattr(
2521 getattr(self.original_model, "generation_config", None),
2522 "no_repeat_ngram_size",
2523 None,
2524 )
2525 if not size:
2526 return None
2527 from transformers.generation.logits_process import NoRepeatNGramLogitsProcessor
2529 return NoRepeatNGramLogitsProcessor(size)
2531 def _encdec_forced_eos_token_id(self) -> Optional[int]:
2532 """generation_config.forced_eos_token_id, or None. HF forces EOS as the
2533 final token when max length is reached; BART-family configs set it by
2534 default and their summaries otherwise end mid-sentence."""
2535 value = getattr(
2536 getattr(self.original_model, "generation_config", None),
2537 "forced_eos_token_id",
2538 None,
2539 )
2540 return value if isinstance(value, int) else None
2542 def _encoder_decoder_setup(
2543 self, input_tokens: torch.Tensor, forced_bos_token_id: Optional[int]
2544 ) -> tuple[torch.Tensor, Optional[int]]:
2545 """Share decoder initialization and length defaults across generation paths."""
2546 config = self.original_model.config
2547 generation_config = getattr(self.original_model, "generation_config", None)
2548 if forced_bos_token_id is None:
2549 # HF defaults such as bart-large-cnn's forced BOS apply to both loops.
2550 forced_bos_token_id = getattr(generation_config, "forced_bos_token_id", None)
2551 min_decoder_length = getattr(generation_config, "min_length", None)
2552 start = getattr(config, "decoder_start_token_id", None)
2553 if start is None:
2554 # HF falls back to BOS, then EOS; IndicBART starts decoding from EOS.
2555 start = getattr(config, "bos_token_id", None)
2556 if start is None:
2557 start = getattr(config, "eos_token_id", None)
2558 if isinstance(start, (list, tuple)):
2559 start = start[0]
2560 if start is None:
2561 start = 0
2562 seed = torch.full(
2563 (input_tokens.shape[0], 1),
2564 start,
2565 dtype=input_tokens.dtype,
2566 device=input_tokens.device,
2567 )
2568 if forced_bos_token_id is not None:
2569 # Multilingual seq2seq selects the target language after decoder_start.
2570 forced = torch.full_like(seed, forced_bos_token_id)
2571 seed = torch.cat([seed, forced], dim=1)
2572 return seed, min_decoder_length
2574 def _generate_tokens(
2575 self,
2576 current_tokens: torch.Tensor,
2577 input_tokens: torch.Tensor,
2578 batch_size: int,
2579 *,
2580 max_new_tokens: int,
2581 do_sample: bool,
2582 top_k: Optional[int],
2583 top_p: Optional[float],
2584 temperature: float,
2585 freq_penalty: float,
2586 repetition_penalty: float,
2587 stop_at_eos: bool,
2588 stop_tokens: List[int],
2589 eos_token_for_padding: int,
2590 finished_sequences: torch.Tensor,
2591 use_past_kv_cache: bool,
2592 use_stateful_cache: bool,
2593 mamba_cache: Any,
2594 mamba_conv_kernel: int,
2595 is_encoder_decoder: bool,
2596 _is_batched_list: bool,
2597 _generate_from_embeds: bool,
2598 encoder_input: Optional[torch.Tensor],
2599 decoder_tokens: Optional[torch.Tensor],
2600 generated_token_ids: Optional[List[torch.Tensor]],
2601 pixel_values: Optional[torch.Tensor],
2602 multimodal_kwargs: Dict[str, Any],
2603 verbose: bool,
2604 stopping_criteria_list: Optional[Any] = None,
2605 initial_attention_mask: Optional[torch.Tensor] = None,
2606 min_decoder_length: Optional[int] = None,
2607 ngram_processor: Optional[Any] = None,
2608 encoder_attention_mask: Optional[torch.Tensor] = None,
2609 forced_eos_token_id: Optional[int] = None,
2610 ) -> Generator[Tuple[torch.Tensor, torch.Tensor, bool], None, None]:
2611 """Core generation loop. Yields (sampled_tokens, final_logits, all_finished) per step.
2613 Owns the forward pass, sampling, stop handling (EOS and any
2614 ``stopping_criteria_list``), token accumulation, and KV cache management. Callers
2615 are responsible for try/finally cleanup of ``_capture_hf_cache``.
2617 ``stopping_criteria_list`` (from ``_resolve_stopping_criteria``) is evaluated on
2618 the running sequence each step and folded into the finished-sequence mask alongside
2619 EOS, so when it is ``None`` the loop runs the EOS-only path unchanged.
2620 """
2621 _hf_kv_cache = None
2622 # A prior multimodal forward caches mrope deltas on the wrapped model, and a
2623 # text-only generation would inherit them through the KV-cache position path
2624 # (HF's own generate recomputes them at prefill; driving forward() directly
2625 # makes that reset this loop's job).
2626 base = getattr(self.original_model, "base_model", self.original_model)
2627 if getattr(base, "rope_deltas", None) is not None:
2628 base.rope_deltas = None
2629 # A row may finish via EOS and/or any of the configured stopping criteria.
2630 any_stop_active = stop_at_eos or stopping_criteria_list is not None
2632 # Models that own their position derivation (the gate refuses them) cache
2633 # mRoPE deltas on the module between calls; a text-only prefill never
2634 # refreshes them, so a stale delta from an earlier multimodal forward gets
2635 # added to every cached-step position. HF's generate recomputes them at
2636 # prefill via prepare_inputs_for_generation, which this loop bypasses —
2637 # so match it by clearing before the prompt pass. A multimodal prefill
2638 # recomputes its own fresh deltas regardless.
2639 if not self._accepts_derived_position_ids():
2640 underlying = getattr(self, "original_model", None)
2641 for module in (
2642 underlying,
2643 getattr(underlying, "model", None),
2644 getattr(underlying, "language_model", None),
2645 ):
2646 if module is not None and hasattr(module, "rope_deltas"):
2647 module.rope_deltas = None
2649 # Pure-SSM models (Mamba-1/2) take the stateful cache as `cache_params`;
2650 # modern hybrids (Bamba, NemotronH, FalconH1) take `past_key_values` and
2651 # would receive a duplicate cache_params via **kwargs cascade otherwise.
2652 stateful_cache_kwarg = "cache_params"
2653 if use_stateful_cache:
2654 forward_params = inspect.signature(self.original_model.forward).parameters
2655 if "cache_params" not in forward_params:
2656 stateful_cache_kwarg = "past_key_values"
2658 for gen_step_idx in tqdm.tqdm(range(max_new_tokens), disable=not verbose):
2659 with torch.no_grad():
2660 if is_encoder_decoder:
2661 assert encoder_input is not None
2662 encdec_kwargs: Dict[str, Any] = {}
2663 if encoder_attention_mask is not None:
2664 encdec_kwargs["attention_mask"] = encoder_attention_mask.to(
2665 encoder_input.device
2666 )
2667 logits = self(
2668 encoder_input,
2669 return_type="logits",
2670 decoder_input=decoder_tokens,
2671 **encdec_kwargs,
2672 )
2673 else:
2674 forward_kwargs: Dict[str, Any] = {}
2675 # A prompt mask covers only the prompt, so extend it by one
2676 # attended column per token generated so far. position_ids are
2677 # left to forward(), which derives them from the mask for the
2678 # models that can take them.
2679 running_attention_mask: Optional[torch.Tensor] = None
2680 if initial_attention_mask is not None:
2681 n_generated = current_tokens.shape[1] - initial_attention_mask.shape[1]
2682 running_attention_mask = torch.cat(
2683 [
2684 initial_attention_mask.to(current_tokens.device),
2685 torch.ones(
2686 (current_tokens.shape[0], n_generated),
2687 dtype=initial_attention_mask.dtype,
2688 device=current_tokens.device,
2689 ),
2690 ],
2691 dim=1,
2692 )
2693 forward_kwargs["attention_mask"] = running_attention_mask
2694 # Compute attention mask and position_ids for batched
2695 # inputs with padding.
2696 if (
2697 initial_attention_mask is None
2698 and _is_batched_list
2699 and self.tokenizer is not None
2700 and self.tokenizer.pad_token_id is not None
2701 ):
2702 attn_mask = utils.get_attention_mask(
2703 self.tokenizer,
2704 current_tokens,
2705 prepend_bos=getattr(self.cfg, "default_prepend_bos", True),
2706 padding_side="left",
2707 ).to(self.cfg.device)
2708 forward_kwargs["attention_mask"] = attn_mask
2709 # Same target gate as the forward() path: the mask is safe
2710 # for every model, the derived positions are not (#1626).
2711 if self._accepts_derived_position_ids():
2712 position_ids = attn_mask.long().cumsum(-1) - 1
2713 position_ids.masked_fill_(attn_mask == 0, 1)
2714 forward_kwargs["position_ids"] = position_ids
2715 if gen_step_idx == 0:
2716 if pixel_values is not None:
2717 forward_kwargs["pixel_values"] = pixel_values
2718 if multimodal_kwargs:
2719 forward_kwargs.update(multimodal_kwargs)
2720 if use_stateful_cache:
2721 forward_kwargs[stateful_cache_kwarg] = mamba_cache
2722 forward_kwargs["use_cache"] = True
2723 if gen_step_idx == 0:
2724 # Mamba's conv-window warmup positions vs standard
2725 # full-prompt positions for past_key_values hybrids.
2726 prefill_len = (
2727 mamba_conv_kernel
2728 if stateful_cache_kwarg == "cache_params"
2729 else current_tokens.shape[1]
2730 )
2731 cache_position = torch.arange(0, prefill_len, device=self.cfg.device)
2732 forward_kwargs["cache_position"] = cache_position
2733 logits = self(
2734 current_tokens,
2735 return_type="logits",
2736 **forward_kwargs,
2737 )
2738 else:
2739 input_seq_pos = input_tokens.shape[1] + gen_step_idx - 1
2740 cache_position = torch.tensor([input_seq_pos], device=self.cfg.device)
2741 forward_kwargs["cache_position"] = cache_position
2742 if "position_ids" in forward_kwargs: 2742 ↛ 2743line 2742 didn't jump to line 2743 because the condition on line 2742 was never true
2743 forward_kwargs["position_ids"] = forward_kwargs["position_ids"][
2744 :, -1:
2745 ]
2746 logits = self(
2747 current_tokens[:, -1:],
2748 return_type="logits",
2749 **forward_kwargs,
2750 )
2751 elif use_past_kv_cache:
2752 forward_kwargs["use_cache"] = True
2753 if _hf_kv_cache is not None:
2754 forward_kwargs["past_key_values"] = _hf_kv_cache
2755 # HF v5 + macOS-arm64 NaNs when these are inferred
2756 # from cache state alone. Mirror HF generate(): pass
2757 # both an (batch, total_len) attention_mask and a
2758 # (batch, 1) position_ids for the new token.
2759 batch_size = current_tokens.shape[0]
2760 total_len = current_tokens.shape[1]
2761 device = current_tokens.device
2762 if "attention_mask" not in forward_kwargs:
2763 forward_kwargs["attention_mask"] = torch.ones(
2764 (batch_size, total_len),
2765 dtype=torch.long,
2766 device=device,
2767 )
2768 # Gated as a whole (#1626): every branch below supplies
2769 # position_ids, so gating only the prompt derivation
2770 # above would divert a refused model into the
2771 # total_len - 1 fallback, which counts pad slots and is
2772 # wrong per row for a left-padded batch. A model that
2773 # owns its position derivation gets the mask alone,
2774 # matching the uncached path.
2775 if self._accepts_derived_position_ids():
2776 if "position_ids" in forward_kwargs:
2777 forward_kwargs["position_ids"] = forward_kwargs["position_ids"][
2778 :, -1:
2779 ]
2780 elif running_attention_mask is not None:
2781 # total_len - 1 counts pad slots, so it is wrong
2782 # for a left-padded prompt. Derive the new token's
2783 # position from the mask instead.
2784 forward_kwargs["position_ids"] = utils.get_offset_position_ids(
2785 0, running_attention_mask.long()
2786 )[:, -1:]
2787 else:
2788 forward_kwargs["position_ids"] = torch.full(
2789 (batch_size, 1),
2790 total_len - 1,
2791 dtype=torch.long,
2792 device=device,
2793 )
2794 logits = self(
2795 current_tokens[:, -1:],
2796 return_type="logits",
2797 **forward_kwargs,
2798 )
2799 else:
2800 logits = self(
2801 current_tokens,
2802 return_type="logits",
2803 **forward_kwargs,
2804 )
2805 else:
2806 logits = self(current_tokens, return_type="logits", **forward_kwargs)
2807 if use_past_kv_cache and hasattr(self, "_last_hf_cache"):
2808 _hf_kv_cache = self._last_hf_cache or _hf_kv_cache
2809 del self._last_hf_cache
2810 final_logits = logits[:, -1, :]
2812 # Sample next token
2813 penalty_tokens = (
2814 torch.stack(generated_token_ids, dim=1)
2815 if _generate_from_embeds and generated_token_ids
2816 else None
2817 )
2818 # transformers' own NoRepeatNGramLogitsProcessor, honoring
2819 # generation_config (bart-large-cnn pins 3; without it greedy
2820 # decoding falls into a BOS attractor and emits nothing).
2821 if ngram_processor is not None and decoder_tokens is not None:
2822 final_logits = ngram_processor(decoder_tokens, final_logits)
2823 # HF's generate() suppresses EOS below generation_config.min_length
2824 # (bart-large-cnn pins 56); without this the loop can EOS on step
2825 # one and emit an empty summary.
2826 if (
2827 min_decoder_length is not None
2828 and is_encoder_decoder
2829 and decoder_tokens is not None
2830 and decoder_tokens.shape[1] < min_decoder_length
2831 and stop_tokens
2832 ):
2833 final_logits = final_logits.clone()
2834 final_logits[:, stop_tokens] = float("-inf")
2835 # Force EOS as the last allowed token, applied after the other
2836 # processors so it wins — the same ordering HF uses for
2837 # ForcedEOSTokenLogitsProcessor.
2838 if forced_eos_token_id is not None and gen_step_idx == max_new_tokens - 1:
2839 final_logits = torch.full_like(final_logits, float("-inf"))
2840 final_logits[:, forced_eos_token_id] = 0.0
2841 if do_sample:
2842 sampled_tokens = utils.sample_logits(
2843 final_logits,
2844 top_k=top_k,
2845 top_p=top_p,
2846 temperature=temperature,
2847 freq_penalty=freq_penalty,
2848 repetition_penalty=repetition_penalty,
2849 tokens=(
2850 penalty_tokens
2851 if _generate_from_embeds
2852 else (decoder_tokens if is_encoder_decoder else current_tokens)
2853 ),
2854 ).to(self.cfg.device)
2855 else:
2856 sampled_tokens = utils.sample_logits(
2857 final_logits,
2858 temperature=0.0,
2859 repetition_penalty=repetition_penalty,
2860 tokens=(
2861 penalty_tokens
2862 if _generate_from_embeds
2863 else (decoder_tokens if is_encoder_decoder else current_tokens)
2864 ),
2865 ).to(self.cfg.device)
2867 # Freeze rows that finished on an earlier step so they stop emitting
2868 # real tokens. Applies to every active stop mechanism, not just EOS.
2869 if any_stop_active:
2870 sampled_tokens[finished_sequences] = eos_token_for_padding
2872 # Fold this step's EOS matches into the finished mask.
2873 if stop_at_eos:
2874 finished_sequences.logical_or_(
2875 torch.isin(
2876 sampled_tokens.to(self.cfg.device),
2877 torch.tensor(stop_tokens).to(self.cfg.device),
2878 )
2879 )
2881 # Update token sequences
2882 if is_encoder_decoder:
2883 assert decoder_tokens is not None
2884 decoder_tokens = torch.cat([decoder_tokens, sampled_tokens.unsqueeze(1)], dim=1)
2885 elif _generate_from_embeds:
2886 assert generated_token_ids is not None
2887 generated_token_ids.append(sampled_tokens)
2888 embed_fn = self.original_model.get_input_embeddings() # type: ignore[operator]
2889 assert embed_fn is not None
2890 new_embed = embed_fn(sampled_tokens.unsqueeze(1)).to(current_tokens.dtype)
2891 current_tokens = torch.cat([current_tokens, new_embed], dim=1)
2892 else:
2893 current_tokens = torch.cat([current_tokens, sampled_tokens.unsqueeze(1)], dim=1)
2895 # Fold stop_strings / stopping_criteria into the finished mask. They are
2896 # evaluated on the full running sequence (prompt + everything generated so
2897 # far, including the token just appended) with this step's logits as the
2898 # scores argument, matching transformers' StoppingCriteria contract. The
2899 # combined list returns a per-row bool [batch] OR-ing every criterion.
2900 # generate()/generate_stream() guarantee this is plain decoder-only token
2901 # generation, so current_tokens is the running token sequence.
2902 if stopping_criteria_list is not None:
2903 criteria_finished = stopping_criteria_list(current_tokens, final_logits).to(
2904 device=self.cfg.device, dtype=torch.bool
2905 )
2906 if criteria_finished.shape != finished_sequences.shape:
2907 raise ValueError(
2908 "A stopping criterion returned shape "
2909 f"{tuple(criteria_finished.shape)}, expected a per-row bool of "
2910 f"shape {tuple(finished_sequences.shape)} (one entry per sequence)."
2911 )
2912 finished_sequences.logical_or_(criteria_finished)
2914 all_finished = bool(any_stop_active and finished_sequences.all().item())
2916 yield sampled_tokens, final_logits, all_finished
2918 if all_finished: 2918 ↛ 2919line 2918 didn't jump to line 2919 because the condition on line 2918 was never true
2919 return
2921 def _resolve_generation_caching(self, use_past_kv_cache: bool, batched: bool) -> bool:
2922 """Honor adapter caching/batching limits (recurrent/conv decoders have no KV
2923 cache; batching is rejected where padding can't be masked, not mis-generated)."""
2924 if batched and not getattr(self.adapter, "supports_batched_generation", True):
2925 architecture = self.cfg.architecture or type(self.adapter).__name__
2926 raise NotImplementedError(
2927 f"Batched generation is not supported by {architecture}: its forward does not "
2928 "apply an attention mask, so padded rows would corrupt the output. Generate one "
2929 "sequence at a time, or pass equal-length inputs as a tensor."
2930 )
2931 if not getattr(self.adapter, "supports_kv_cache", True):
2932 return False
2933 return use_past_kv_cache
2935 def _ensure_generation_supported(self, api_name: str) -> None:
2936 """Reject autoregressive generation for forward-only architectures."""
2937 if not self.adapter.supports_generation:
2938 architecture = self.cfg.architecture or type(self.adapter).__name__
2939 hint = (
2940 " Use diffusion_generate() — this architecture samples by iterative denoising, "
2941 "not left-to-right."
2942 if getattr(self.adapter, "native_sampler", None)
2943 else ""
2944 )
2945 raise NotImplementedError(
2946 f"TransformerBridge.{api_name}() generation is not supported by "
2947 f"the {architecture} architecture.{hint}"
2948 )
2950 def generate(
2951 self,
2952 input: Union[str, List[str], torch.Tensor] = "",
2953 max_new_tokens: int = 10,
2954 stop_at_eos: bool = True,
2955 eos_token_id: Optional[int] = None,
2956 do_sample: bool = True,
2957 top_k: Optional[int] = None,
2958 top_p: Optional[float] = None,
2959 temperature: float = 1.0,
2960 freq_penalty: float = 0.0,
2961 repetition_penalty: float = 1.0,
2962 use_past_kv_cache: bool = True,
2963 prepend_bos: Optional[bool] = None,
2964 padding_side: Optional[str] = None,
2965 return_type: Optional[str] = "input",
2966 verbose: bool = True,
2967 output_logits: bool = False,
2968 return_cache: bool = False,
2969 return_input_tokens: bool = False,
2970 names_filter: Optional[Union[str, List[str], Callable[[str], bool]]] = None,
2971 device: Optional[Union[str, torch.device]] = None,
2972 pixel_values: Optional[torch.Tensor] = None,
2973 stop_strings: Optional[Union[str, List[str]]] = None,
2974 stopping_criteria: Optional[Any] = None,
2975 attention_mask: Optional[torch.Tensor] = None,
2976 forced_bos_token_id: Optional[int] = None,
2977 **multimodal_kwargs,
2978 ) -> (
2979 str
2980 | list[str]
2981 | torch.Tensor
2982 | Any
2983 | tuple[Any, ActivationCache]
2984 | tuple[Any, torch.Tensor]
2985 ): # Any for transformers.utils.ModelOutput
2986 # Any: beartype forward ref limitation (beartype#546)
2987 """Sample tokens from the model.
2989 Sample tokens from the model until the model outputs eos_token or max_new_tokens is reached.
2990 This implementation is based on HookedTransformer.generate() to ensure consistent behavior.
2992 Args:
2993 input: Text string, list of strings, or tensor of tokens
2994 max_new_tokens: Maximum number of tokens to generate
2995 stop_at_eos: If True, stop generating tokens when the model outputs eos_token
2996 eos_token_id: The token ID to use for end of sentence
2997 do_sample: If True, sample from the model's output distribution. Otherwise, use greedy search
2998 top_k: Number of tokens to sample from. If None, sample from all tokens
2999 top_p: Probability mass to sample from. If 1.0, sample from all tokens
3000 temperature: Temperature for sampling. Higher values will make the model more random
3001 freq_penalty: Frequency penalty for sampling - how much to penalise previous tokens
3002 repetition_penalty: HuggingFace-style repetition penalty. Values > 1.0 discourage
3003 repetition by dividing positive logits and multiplying negative logits for
3004 previously seen tokens. Default 1.0 (no penalty).
3005 use_past_kv_cache: If True, use KV caching for faster generation
3006 prepend_bos: Whether to prepend a BOS token when tokenizing string inputs.
3007 Defaults to None (uses ``cfg.default_prepend_bos``, typically True).
3008 Pass ``prepend_bos=False`` when the input is pre-formatted chat-template
3009 text that already contains the BOS token to avoid double-BOS.
3010 Ignored when input is already a token tensor.
3011 padding_side: Which side to pad when tokenizing multiple strings of different
3012 lengths. For batched list inputs, left-padding is forced internally for
3013 correct generation behavior. Defaults to None (tokenizer default).
3014 return_type: The type of output to return - 'input', 'str', or 'tokens'
3015 verbose: Not used in Bridge (kept for API compatibility)
3016 output_logits: If True, return a ModelOutput with sequences and logits tuple
3017 return_cache: If True, also return an ActivationCache for the full prompt +
3018 generated sequence, identical to ``run_with_cache(output)``, and the call
3019 returns an ``(output, cache)`` tuple. Implemented as one extra clean forward
3020 over the output, so the cache includes every hook point (attention patterns
3021 included). Supported only for single-sequence, decoder-only text generation;
3022 encoder-decoder, SSM, multimodal, batched, and inputs_embeds inputs raise
3023 NotImplementedError. The cache spans prompt + max_new_tokens and can be large,
3024 use ``names_filter`` to scope it and/or ``device`` to offload it.
3025 return_input_tokens: If True, return an ``(output, input_tokens)`` tuple where
3026 ``input_tokens`` is the token tensor that was actually fed to the model
3027 (after BOS handling). Useful for debugging tokenization, especially when
3028 using chat templates where BOS handling can be subtle. Can be combined
3029 with ``return_cache`` to get ``(output, cache, input_tokens)``.
3030 names_filter: Passed to ``run_with_cache`` when ``return_cache=True``; restricts
3031 which activations are cached (str, list of str, or callable).
3032 device: Passed through when ``return_cache=True`` to offload the cached tensors
3033 to this device (e.g. "cpu") to save accelerator memory.
3034 pixel_values: Optional image tensor for multimodal models. Only passed on the
3035 first generation step (the vision encoder processes the image once, then
3036 embeddings are part of the token sequence for subsequent steps).
3037 stop_strings: Optional string or list of strings. A sequence stops once its
3038 generated text ends with one of these strings, using HuggingFace's
3039 StopStringCriteria (partial-token-aware, end-anchored) matching.
3040 Requires a tokenizer (raises ValueError otherwise).
3041 Independent of stop_at_eos: either can stop a sequence.
3042 stopping_criteria: Optional HuggingFace stopping criteria, a single
3043 transformers.StoppingCriteria, a list of them, or a StoppingCriteriaList.
3044 Each is called as criterion(input_ids, scores) after every step and ORed
3045 with the other stop signals, where input_ids is the running sequence and
3046 scores is this step's logits ([batch, d_vocab]). Each criterion must return
3047 a per-row bool [batch] (or a scalar bool). stop_strings and stopping_criteria
3048 are supported only for standard decoder-only text generation. Encoder-decoder,
3049 inputs_embeds, and multimodal generation always raise NotImplementedError.
3050 Stateful/SSM models raise only when run with use_past_kv_cache=False (the
3051 default keeps them on the hooked loop). Each error names the supported
3052 alternative.
3053 attention_mask: Optional ``[batch, pos]`` 0/1 mask over the prompt, marking
3054 which prompt tokens are real. Required to generate correctly from an
3055 already-padded token tensor: without it the pad tokens are treated as
3056 real context and every real token's position is shifted, so the
3057 continuation differs from the same prompt unpadded. The mask is extended
3058 by one attended column per generated token. Takes precedence over the
3059 ``padding_side`` heuristic, and unlike it can express an interior gap or
3060 a pad id that also occurs as a real token. Passing ``padding_side``
3061 instead reads the padding off the pad token, which is enough for the
3062 common single-edge case, and raises if this bridge has no tokenizer
3063 or pad id to read it from. On the encoder-decoder and inputs_embeds
3064 paths the mask is forwarded to the model as-is rather than grown per
3065 step, which is what processors emitting one alongside
3066 ``pixel_values`` expect.
3067 forced_bos_token_id: Optional token id seeded as the first decoder token
3068 after ``decoder_start`` on encoder-decoder models. Multilingual
3069 translators (M2M100/MBart/NLLB) select their target language this way.
3070 Raises ValueError on decoder-only models.
3072 Returns:
3073 Generated sequence as string, list of strings, or tensor depending on input type and return_type.
3074 If output_logits=True, returns a ModelOutput-like object with 'sequences' and 'logits' attributes.
3075 If return_cache=True, returns an ``(output, ActivationCache)`` tuple where ``output`` is the
3076 value that would otherwise be returned and the cache equals ``run_with_cache(output)``.
3077 If return_input_tokens=True, returns an ``(output, input_tokens)`` tuple.
3078 If both return_cache and return_input_tokens are True, returns ``(output, cache, input_tokens)``.
3080 Example:
3081 ``out, cache = model.generate(prompt, max_new_tokens=20, return_cache=True)`` returns a
3082 normal ActivationCache over the full prompt + generated sequence (equivalent to
3083 ``run_with_cache(out)``).
3085 ``out, input_tokens = model.generate(prompt, return_input_tokens=True)`` returns
3086 the tokens that were fed to the model, useful for verifying BOS handling with
3087 chat templates.
3088 """
3089 self._ensure_generation_supported("generate")
3090 # padding_side is handled internally: for batched list inputs, left-padding
3091 # is forced to ensure correct generation. See _is_batched_list logic below.
3093 # Stateful dispatch is decided after input parsing so we can fall back
3094 # to hf_generate() for input types the stateful loop doesn't handle.
3095 is_stateful_model = getattr(self.cfg, "is_stateful", False)
3097 _is_batched_list = isinstance(input, list) and len(input) > 1
3098 use_past_kv_cache = self._resolve_generation_caching(use_past_kv_cache, _is_batched_list)
3100 _generate_from_embeds = False
3101 _encdec_early = hasattr(self.original_model, "config") and getattr(
3102 self.original_model.config, "is_encoder_decoder", False
3103 )
3104 if isinstance(input, str):
3105 if _encdec_early:
3106 # Deliberate divergence: prepend_bos is IGNORED for enc-dec
3107 # string/list input. Encoder input follows the tokenizer's own
3108 # recipe (lang token + trailing </s>); to_tokens' decoder-style
3109 # BOS policy corrupts it — m2m100 degenerates to loops.
3110 input_tokens = self.tokenizer(input, return_tensors="pt")["input_ids"].to(
3111 self.cfg.device
3112 )
3113 else:
3114 input_tokens = self.to_tokens(
3115 input, prepend_bos=prepend_bos, move_to_device=True, truncate=False
3116 )
3117 input_type = "str"
3118 elif isinstance(input, list):
3119 if _encdec_early:
3120 # Same native-recipe rule as the str branch: to_tokens' BOS
3121 # policy corrupts encoder inputs (stray <s>, dropped </s>).
3122 # Keep the tokenizer's mask too — unequal rows otherwise
3123 # attend over pads in the encoder.
3124 _enc_batch = self.tokenizer(input, return_tensors="pt", padding=True)
3125 input_tokens = _enc_batch["input_ids"].to(self.cfg.device)
3126 if attention_mask is None and "attention_mask" in _enc_batch: 3126 ↛ 3140line 3126 didn't jump to line 3140 because the condition on line 3126 was always true
3127 attention_mask = _enc_batch["attention_mask"].to(self.cfg.device)
3128 else:
3129 # Force left-padding for batched generation so real tokens are
3130 # flush-right and logits[:, -1, :] is always the last real token.
3131 # Passed as a kwarg rather than assigned: a raise between assignment
3132 # and restore would pin the shared tokenizer left for the session.
3133 input_tokens = self.to_tokens(
3134 input,
3135 prepend_bos=prepend_bos,
3136 padding_side="left" if _is_batched_list else None,
3137 move_to_device=True,
3138 truncate=False,
3139 )
3140 input_type = "list"
3141 elif isinstance(input, torch.Tensor) and input.is_floating_point():
3142 # inputs_embeds: pre-computed embeddings (e.g., from multimodal models)
3143 input_tokens = input.to(self.cfg.device)
3144 input_type = "embeds"
3145 _generate_from_embeds = True
3146 else:
3147 input_tokens = input.to(self.cfg.device)
3148 input_type = "tokens"
3150 # Without one of these a pre-padded tensor generates as though its pads were
3151 # real context, shifting every real token's position (#1612). An explicit
3152 # mask wins; otherwise the padding is read off the tokens, but only when the
3153 # caller asked for that by passing padding_side. Deriving a mask on the
3154 # default path would silently change behaviour for every existing caller,
3155 # and would demand a real tokenizer where today none is required.
3156 initial_attention_mask: Optional[torch.Tensor] = attention_mask
3157 if initial_attention_mask is not None and (
3158 _generate_from_embeds
3159 or getattr(getattr(self.original_model, "config", None), "is_encoder_decoder", False)
3160 ):
3161 # Growing the mask per step only means something for decoder-only token
3162 # generation. On these paths the mask used to arrive via
3163 # **multimodal_kwargs and be forwarded to the model untouched — as
3164 # processors emit it alongside pixel_values — so keep doing that rather
3165 # than reject a call that worked before this parameter existed.
3166 multimodal_kwargs = {**multimodal_kwargs, "attention_mask": initial_attention_mask}
3167 initial_attention_mask = None
3168 if initial_attention_mask is not None:
3169 if initial_attention_mask.shape != input_tokens.shape:
3170 raise ValueError(
3171 f"attention_mask shape {tuple(initial_attention_mask.shape)} does not "
3172 f"match the prompt shape {tuple(input_tokens.shape)}. Pass a 0/1 mask "
3173 "covering exactly the prompt tokens; generate() extends it itself."
3174 )
3175 initial_attention_mask = initial_attention_mask.to(self.cfg.device)
3176 elif padding_side is not None and input_type == "tokens":
3177 # Reading the padding off the tokens needs a tokenizer with a pad id.
3178 # Without one the argument would be inert, leaving exactly the bug this
3179 # fixes — silently, on a bridge booted without a tokenizer. Say so
3180 # rather than generate something quietly wrong.
3181 if not isinstance(self.tokenizer, PreTrainedTokenizerBase):
3182 raise ValueError(
3183 "generate(padding_side=...) reads the padding off the pad token, "
3184 "which needs a tokenizer; this bridge has none. Pass "
3185 "attention_mask=... to state the padding directly instead."
3186 )
3187 if self.tokenizer.pad_token_id is None:
3188 raise ValueError(
3189 "generate(padding_side=...) reads the padding off the pad token, "
3190 "but this tokenizer has no pad_token_id. Set one, or pass "
3191 "attention_mask=... to state the padding directly instead."
3192 )
3193 _prepend = self.cfg.default_prepend_bos if prepend_bos is None else prepend_bos
3194 _orig_side = self.tokenizer.padding_side
3195 self.tokenizer.padding_side = padding_side
3196 try:
3197 initial_attention_mask = utils.get_attention_mask(
3198 self.tokenizer, input_tokens, _prepend
3199 ).to(self.cfg.device)
3200 finally:
3201 self.tokenizer.padding_side = _orig_side
3202 # An all-ones mask is what the model assumes anyway; skipping it keeps
3203 # the unpadded path byte-identical to before.
3204 if initial_attention_mask is not None and bool(initial_attention_mask.all()): 3204 ↛ 3205line 3204 didn't jump to line 3205 because the condition on line 3204 was never true
3205 initial_attention_mask = None
3207 # Determine return type
3208 if return_type == "input":
3209 if input_type in ["str", "list"]:
3210 return_type = "str"
3211 elif input_type == "embeds":
3212 return_type = "tokens"
3213 else:
3214 return_type = "tokens"
3216 batch_size = input_tokens.shape[0]
3218 # Setup EOS token handling
3219 stop_tokens = []
3220 eos_token_for_padding = 0
3221 if stop_at_eos:
3222 tokenizer_has_eos_token = (
3223 self.tokenizer is not None and self.tokenizer.eos_token_id is not None
3224 )
3225 if eos_token_id is None:
3226 # Some chat models use a turn-end token that differs from the
3227 # tokenizer's primary EOS. Let adapters provide the full stop
3228 # set via cfg.eos_token_id; otherwise fall back to the tokenizer.
3229 eos_token_id = getattr(self.cfg, "eos_token_id", None)
3230 if eos_token_id is None:
3231 assert (
3232 tokenizer_has_eos_token
3233 ), "Must pass eos_token_id if stop_at_eos is True and tokenizer is None or has no eos_token_id"
3234 assert self.tokenizer is not None
3235 eos_token_id = self.tokenizer.eos_token_id
3237 if isinstance(eos_token_id, int):
3238 stop_tokens = [eos_token_id]
3239 eos_token_for_padding = eos_token_id
3240 else:
3241 stop_tokens = list(eos_token_id)
3242 if tokenizer_has_eos_token:
3243 assert self.tokenizer is not None
3244 eos_token_for_padding = self.tokenizer.eos_token_id
3245 else:
3246 eos_token_for_padding = eos_token_id[0]
3248 # Track which sequences have finished
3249 finished_sequences = torch.zeros(batch_size, dtype=torch.bool, device=self.cfg.device)
3251 # Optionally collect logits at each generation step for downstream tooling/tests
3252 logits_seq_list: list[torch.Tensor] | None = [] if output_logits else None
3254 # Detect encoder-decoder models (T5, BART, etc.)
3255 is_encoder_decoder = hasattr(self.original_model, "config") and getattr(
3256 self.original_model.config, "is_encoder_decoder", False
3257 )
3258 if forced_bos_token_id is not None and not is_encoder_decoder: 3258 ↛ 3261line 3258 didn't jump to line 3261 because the condition on line 3258 was never true
3259 # Raise before any state mutation (_capture_hf_cache) and before
3260 # the stateful hf_generate early-return would drop the kwarg.
3261 raise ValueError("forced_bos_token_id is only meaningful for encoder-decoder models")
3263 # return_cache recomputes run_with_cache on the generated output (see issue #697).
3264 # That is well-defined only for single-sequence, decoder-only text generation, so
3265 # reject the paths whose cache would be wrong/undefined, with a clear pointer to the
3266 # run_with_cache workaround. Fail fast here, before any generation work.
3267 if return_cache:
3268 if is_encoder_decoder: 3268 ↛ 3269line 3268 didn't jump to line 3269 because the condition on line 3268 was never true
3269 raise NotImplementedError(
3270 "generate(return_cache=True) is not supported for encoder-decoder "
3271 "models yet. Run run_with_cache on the generated output instead."
3272 )
3273 if is_stateful_model: 3273 ↛ 3274line 3273 didn't jump to line 3274 because the condition on line 3273 was never true
3274 raise NotImplementedError(
3275 "generate(return_cache=True) is not supported for stateful/SSM models "
3276 "(e.g. Mamba); they do not expose standard transformer hook points."
3277 )
3278 if pixel_values is not None or multimodal_kwargs: 3278 ↛ 3279line 3278 didn't jump to line 3279 because the condition on line 3278 was never true
3279 raise NotImplementedError(
3280 "generate(return_cache=True) is not supported for multimodal generation "
3281 "yet. Run run_with_cache on the generated output instead."
3282 )
3283 if _generate_from_embeds:
3284 raise NotImplementedError(
3285 "generate(return_cache=True) requires token input, not inputs_embeds."
3286 )
3287 if batch_size > 1:
3288 raise NotImplementedError(
3289 "generate(return_cache=True) is not supported for batched/multi-prompt "
3290 "generation yet. Pass a single prompt, or run run_with_cache on each "
3291 "output sequence."
3292 )
3294 # HF cache flows opaquely through the component chain via
3295 # _reconstruct_attention() → _update_kv_cache() on each layer.
3296 _hf_kv_cache = None
3297 if use_past_kv_cache and is_encoder_decoder:
3298 # Encoder-decoder models (T5, BART) don't support the opaque
3299 # cache path — silently disable rather than crash, since
3300 # use_past_kv_cache=True is the default.
3301 use_past_kv_cache = False
3303 # SSMs (Mamba/Mamba-2) run through a dedicated cache path so hooks
3304 # fire on every step. Unsupported input types fall back to hf_generate().
3305 use_stateful_cache = (
3306 is_stateful_model
3307 and use_past_kv_cache
3308 and not is_encoder_decoder
3309 and not _generate_from_embeds
3310 and pixel_values is None
3311 and not multimodal_kwargs
3312 )
3314 # stop_strings / stopping_criteria are applied inside the hooked _generate_tokens
3315 # loop, so they are supported only on the standard decoder-only text path. Reject
3316 # the paths that route around that loop with a clear error rather than silently
3317 # dropping the kwargs. This must run before the stateful delegation below.
3318 stopping_criteria_list = self._resolve_stopping_criteria(stop_strings, stopping_criteria)
3319 if stopping_criteria_list is not None:
3320 if is_encoder_decoder:
3321 _unsupported = "encoder-decoder models"
3322 elif _generate_from_embeds:
3323 _unsupported = "inputs_embeds generation"
3324 elif pixel_values is not None or multimodal_kwargs:
3325 _unsupported = "multimodal (pixel_values) generation"
3326 else:
3327 _unsupported = None
3328 if _unsupported is not None:
3329 raise NotImplementedError(
3330 f"stop_strings/stopping_criteria are not supported for {_unsupported} in "
3331 "TransformerBridge.generate(). Call hf_generate(...), which runs "
3332 "HuggingFace's own generation loop and supports HF-native stopping on "
3333 "those inputs."
3334 )
3335 if is_stateful_model and not use_stateful_cache:
3336 # Reached only for a stateful/SSM model with use_past_kv_cache=False: the
3337 # hooked loop needs the stateful cache, so generate() would otherwise fall
3338 # back to hf_generate() and drop these kwargs. The default cache setting
3339 # keeps generation on the hooked loop, where stopping is applied.
3340 raise NotImplementedError(
3341 "stop_strings/stopping_criteria on a stateful/SSM model require the "
3342 "stateful cache path, which runs only with use_past_kv_cache=True (the "
3343 "default). With use_past_kv_cache=False generate() falls back to "
3344 "hf_generate(). Set use_past_kv_cache=True to keep stopping on the hooked "
3345 "loop, or call hf_generate(...) directly for HF-native stopping."
3346 )
3347 # Finished rows are overwritten with this id so they stop emitting real tokens
3348 # while the rest of a batch keeps going. stop_at_eos already set a sensible
3349 # value, otherwise fall back to the tokenizer pad/eos id. (For a single
3350 # sequence this id is never read: the loop exits when the row finishes.)
3351 if not stop_at_eos:
3352 _pad_id = None
3353 if self.tokenizer is not None:
3354 _pad_id = (
3355 self.tokenizer.pad_token_id
3356 if self.tokenizer.pad_token_id is not None
3357 else self.tokenizer.eos_token_id
3358 )
3359 if _pad_id is not None:
3360 eos_token_for_padding = _pad_id
3361 elif batch_size > 1:
3362 raise ValueError(
3363 "Batched generation with stopping_criteria and stop_at_eos=False "
3364 "needs a padding token to freeze finished rows, but no tokenizer "
3365 "pad/eos id is available. Set stop_at_eos=True, use a tokenizer with "
3366 "a pad or eos token, or generate one sequence at a time."
3367 )
3369 if is_stateful_model and not use_stateful_cache:
3370 hf_kwargs: dict[str, Any] = {
3371 "max_new_tokens": max_new_tokens,
3372 "do_sample": do_sample,
3373 "temperature": temperature,
3374 # hf_generate defaults this to True, so omitting it would make
3375 # generate(stop_at_eos=False) silently stop at EOS on this path.
3376 "stop_at_eos": stop_at_eos,
3377 }
3378 if top_k is not None: 3378 ↛ 3379line 3378 didn't jump to line 3379 because the condition on line 3378 was never true
3379 hf_kwargs["top_k"] = top_k
3380 if top_p is not None: 3380 ↛ 3381line 3380 didn't jump to line 3381 because the condition on line 3380 was never true
3381 hf_kwargs["top_p"] = top_p
3382 if eos_token_id is not None:
3383 hf_kwargs["eos_token_id"] = eos_token_id
3384 # generate() already derived and shape-validated this mask; dropping it here
3385 # would let the model attend to pad positions on the stateful fallback.
3386 if initial_attention_mask is not None:
3387 hf_kwargs["attention_mask"] = initial_attention_mask
3388 return self.hf_generate(input, **hf_kwargs)
3390 # SSM cache is built once and mutated in place across forward calls.
3391 # Adapter owns the cache-type choice; new SSMs just override
3392 # create_stateful_cache().
3393 mamba_cache: Any = None
3394 mamba_conv_kernel: int = 0
3395 if use_stateful_cache:
3396 hf_model: Any = self.original_model
3397 mamba_conv_kernel = int(getattr(hf_model.config, "conv_kernel", 4))
3398 cache_dtype = self.cfg.dtype or torch.float32
3399 mamba_cache = self.adapter.create_stateful_cache(
3400 hf_model=hf_model,
3401 batch_size=batch_size,
3402 device=self.cfg.device,
3403 dtype=cache_dtype,
3404 )
3406 if use_past_kv_cache and not use_stateful_cache:
3407 self._capture_hf_cache = True # Signal forward() to stash cache
3409 # Generate tokens
3410 current_tokens = input_tokens.clone()
3411 # For inputs_embeds generation, also track generated token IDs for decoding
3412 if _generate_from_embeds:
3413 generated_token_ids: list[torch.Tensor] = []
3414 sampled_tokens_list = []
3416 # For encoder-decoder models, keep encoder input fixed and grow decoder input
3417 min_decoder_length = None
3418 if is_encoder_decoder:
3419 encoder_input = input_tokens.clone()
3420 decoder_tokens, min_decoder_length = self._encoder_decoder_setup(
3421 input_tokens, forced_bos_token_id
3422 )
3424 try:
3425 for sampled_tokens, final_logits, all_finished in self._generate_tokens(
3426 current_tokens,
3427 input_tokens,
3428 batch_size,
3429 max_new_tokens=max_new_tokens,
3430 do_sample=do_sample,
3431 top_k=top_k,
3432 top_p=top_p,
3433 temperature=temperature,
3434 freq_penalty=freq_penalty,
3435 repetition_penalty=repetition_penalty,
3436 stop_at_eos=stop_at_eos,
3437 stop_tokens=stop_tokens,
3438 eos_token_for_padding=eos_token_for_padding,
3439 finished_sequences=finished_sequences,
3440 use_past_kv_cache=use_past_kv_cache,
3441 use_stateful_cache=use_stateful_cache,
3442 mamba_cache=mamba_cache,
3443 mamba_conv_kernel=mamba_conv_kernel,
3444 is_encoder_decoder=is_encoder_decoder,
3445 _is_batched_list=_is_batched_list,
3446 _generate_from_embeds=_generate_from_embeds,
3447 encoder_input=encoder_input if is_encoder_decoder else None,
3448 decoder_tokens=decoder_tokens if is_encoder_decoder else None,
3449 generated_token_ids=generated_token_ids if _generate_from_embeds else None,
3450 pixel_values=pixel_values,
3451 multimodal_kwargs=multimodal_kwargs if multimodal_kwargs else {},
3452 verbose=verbose,
3453 stopping_criteria_list=stopping_criteria_list,
3454 initial_attention_mask=initial_attention_mask,
3455 min_decoder_length=min_decoder_length,
3456 ngram_processor=(self._encdec_ngram_processor() if is_encoder_decoder else None),
3457 encoder_attention_mask=(attention_mask if is_encoder_decoder else None),
3458 forced_eos_token_id=(
3459 self._encdec_forced_eos_token_id() if is_encoder_decoder else None
3460 ),
3461 ):
3462 sampled_tokens_list.append(sampled_tokens.unsqueeze(1))
3463 if logits_seq_list is not None:
3464 logits_seq_list.append(final_logits.clone())
3465 if all_finished:
3466 break
3467 finally:
3468 self._capture_hf_cache = False
3469 if hasattr(self, "_last_hf_cache"): 3469 ↛ 3470line 3469 didn't jump to line 3470 because the condition on line 3469 was never true
3470 del self._last_hf_cache
3472 sampled_tokens = torch.cat(sampled_tokens_list, dim=1)
3473 if is_encoder_decoder:
3474 # Reconstruct full decoder sequence: start token + generated tokens
3475 output_tokens = torch.cat([decoder_tokens, sampled_tokens], dim=1)
3476 elif _generate_from_embeds:
3477 # For inputs_embeds, we only have the generated token IDs (no input token IDs)
3478 output_tokens = sampled_tokens
3479 else:
3480 output_tokens = torch.cat([input_tokens, sampled_tokens], dim=1)
3482 # Build the formatted output (shape unchanged: ModelOutput / str / list[str] / tokens).
3483 result: Any
3484 if output_logits and logits_seq_list is not None:
3485 from transformers.utils import ModelOutput # type: ignore
3487 def _logits_to_tuple(logits_list: list[torch.Tensor]) -> tuple[torch.Tensor, ...]:
3488 assert logits_list is not None
3489 # Convert list of [batch, vocab] tensors to tuple
3490 return tuple(logits_list)
3492 try:
3493 from transformers.generation.utils import GenerateDecoderOnlyOutput
3495 # HF-compatible ModelOutput structure.
3496 # GenerateDecoderOnlyOutput expects: sequences, scores (optional), logits (optional)
3497 result = GenerateDecoderOnlyOutput(
3498 sequences=cast(torch.LongTensor, output_tokens),
3499 # HF's type hint says tuple[FloatTensor] but should be tuple[FloatTensor, ...]
3500 # (variable-length tuple with one element per generated token)
3501 logits=_logits_to_tuple(logits_seq_list), # type: ignore[arg-type]
3502 )
3503 except (ImportError, AttributeError):
3504 # Fallback if GenerateDecoderOnlyOutput not available in this transformers version
3505 result = ModelOutput(
3506 sequences=output_tokens,
3507 logits=_logits_to_tuple(logits_seq_list),
3508 )
3509 elif return_type == "str":
3510 assert self.tokenizer is not None
3511 if input_type == "str":
3512 result = self.tokenizer.decode(output_tokens[0], skip_special_tokens=True)
3513 else:
3514 decoded_texts = [
3515 self.tokenizer.decode(tokens, skip_special_tokens=True)
3516 for tokens in output_tokens
3517 ]
3518 result = decoded_texts[0] if len(decoded_texts) == 1 else decoded_texts
3519 else: # return_type == "tokens"
3520 result = output_tokens
3522 if not return_cache and not return_input_tokens:
3523 return result
3525 if return_cache:
3526 # return_cache: recompute one clean forward over the full generated sequence so the
3527 # cache is identical to run_with_cache(output_tokens) - all hook points, including
3528 # attention patterns. The guards above restrict this to single-sequence, decoder-only
3529 # text generation (see issue #697).
3530 _, cache = self.run_with_cache(output_tokens, names_filter=names_filter, device=device)
3531 if return_input_tokens:
3532 return result, cache, input_tokens
3533 return result, cache
3535 # return_input_tokens only (no cache)
3536 return result, input_tokens
3538 @torch.no_grad()
3539 def diffusion_generate(
3540 self,
3541 input: Union[str, List[str], torch.Tensor],
3542 max_new_tokens: int = 32,
3543 prepend_bos: Optional[bool] = None,
3544 **kwargs: Any,
3545 ) -> Union[str, torch.Tensor]:
3546 """Sample from a non-autoregressive (diffusion) architecture.
3548 Delegates to the model's own sampler, which calls the model through
3549 ``__call__`` so bridge hooks fire on every denoising step.
3550 """
3551 sampler_name = getattr(self.adapter, "native_sampler", None)
3552 architecture = self.cfg.architecture or type(self.adapter).__name__
3553 if sampler_name is None: 3553 ↛ 3554line 3553 didn't jump to line 3554 because the condition on line 3553 was never true
3554 raise NotImplementedError(
3555 f"{architecture} has no native sampler; use generate() for autoregressive "
3556 "architectures."
3557 )
3558 sampler = getattr(self.original_model, sampler_name, None)
3559 if sampler is None: 3559 ↛ 3560line 3559 didn't jump to line 3560 because the condition on line 3559 was never true
3560 raise NotImplementedError(
3561 f"{architecture} declares native_sampler={sampler_name!r} but the loaded model "
3562 "has no such method."
3563 )
3565 was_string = isinstance(input, str)
3566 if isinstance(input, list) and len(input) > 1: 3566 ↛ 3570line 3566 didn't jump to line 3570 because the condition on line 3566 was never true
3567 # Unequal prompts would be right-padded into the sampler's canvas,
3568 # where pad tokens read as real context. generate() gates batching
3569 # for the same reason; do not silently corrupt rows here.
3570 raise NotImplementedError(
3571 f"diffusion_generate() does not support batched prompts for {architecture}: "
3572 "the native samplers condition on a padded canvas. Sample one prompt at a time."
3573 )
3574 if isinstance(input, torch.Tensor): 3574 ↛ 3579line 3574 didn't jump to line 3579 because the condition on line 3574 was always true
3575 tokens = input.to(self.cfg.device)
3576 else:
3577 # Tokenization is the bridge's concern, not the sampler's — absorb
3578 # prepend_bos here rather than forwarding it into native kwargs.
3579 tokens = self.to_tokens(
3580 input, prepend_bos=prepend_bos, move_to_device=True, truncate=False
3581 )
3583 sampler_kwargs = self.adapter.native_sampler_kwargs(max_new_tokens, tokens.shape[-1])
3584 sampler_kwargs.update(kwargs)
3585 output = sampler(tokens, **sampler_kwargs)
3586 # Samplers return either bare ids or a generation output object.
3587 sequences = getattr(output, "sequences", output)
3588 # They also disagree on whether the prompt is included: Dream and Gidd
3589 # return the whole canvas, LLaDA2 slices it off. Normalize to
3590 # generate()'s contract (prompt + continuation).
3591 if isinstance(sequences, torch.Tensor) and sequences.ndim == tokens.ndim: 3591 ↛ 3602line 3591 didn't jump to line 3602 because the condition on line 3591 was always true
3592 prompt_len = tokens.shape[-1]
3593 # Compare on one device: torch.equal raises on a device mismatch,
3594 # which some samplers produce by assembling output on CPU.
3595 sequences = sequences.to(tokens.device)
3596 includes_prompt = sequences.shape[-1] >= prompt_len and torch.equal(
3597 sequences[..., :prompt_len], tokens
3598 )
3599 if not includes_prompt:
3600 sequences = torch.cat([tokens, sequences], dim=-1)
3602 if was_string and self.tokenizer is not None: 3602 ↛ 3603line 3602 didn't jump to line 3603 because the condition on line 3602 was never true
3603 return self.tokenizer.decode(sequences[0], skip_special_tokens=True)
3604 return sequences
3606 @torch.no_grad()
3607 def generate_stream(
3608 self,
3609 input: Union[str, List[str], torch.Tensor] = "",
3610 max_new_tokens: int = 10,
3611 max_tokens_per_yield: int = 25,
3612 stop_at_eos: bool = True,
3613 eos_token_id: Optional[int] = None,
3614 do_sample: bool = True,
3615 top_k: Optional[int] = None,
3616 top_p: Optional[float] = None,
3617 temperature: float = 1.0,
3618 freq_penalty: float = 0.0,
3619 repetition_penalty: float = 1.0,
3620 use_past_kv_cache: bool = True,
3621 prepend_bos: Optional[bool] = None,
3622 padding_side: Optional[str] = None,
3623 return_type: Optional[str] = "input",
3624 verbose: bool = True,
3625 stop_strings: Optional[Union[str, List[str]]] = None,
3626 stopping_criteria: Optional[Any] = None,
3627 attention_mask: Optional[torch.Tensor] = None,
3628 forced_bos_token_id: Optional[int] = None,
3629 ) -> Generator[Union[torch.Tensor, str, List[str]], None, None]:
3630 """Stream tokens from the model as they are generated.
3632 Yields batches of tokens progressively during generation rather than
3633 waiting for the entire sequence. Uses the same core loop as generate().
3635 Args:
3636 input: Text string, list of strings, or tensor of tokens.
3637 max_new_tokens: Maximum number of tokens to generate.
3638 max_tokens_per_yield: Yield accumulated tokens every this many steps.
3639 stop_at_eos: If True, stop when eos_token is produced.
3640 eos_token_id: Token ID(s) for end of sentence. Defaults to tokenizer's.
3641 do_sample: If True, sample; otherwise greedy.
3642 top_k: Top-k sampling. None means no filtering.
3643 top_p: Nucleus sampling threshold.
3644 temperature: Sampling temperature.
3645 freq_penalty: Frequency penalty for previous tokens.
3646 repetition_penalty: HF-style repetition penalty (>1.0 discourages repeats).
3647 use_past_kv_cache: Use KV caching for faster generation.
3648 prepend_bos: Whether to prepend a BOS token when tokenizing string inputs.
3649 Defaults to None (uses ``cfg.default_prepend_bos``, typically True).
3650 Pass ``prepend_bos=False`` when the input is pre-formatted chat-template
3651 text that already contains the BOS token to avoid double-BOS.
3652 Ignored when input is already a token tensor.
3653 padding_side: Which side to pad for batched list inputs. Left-padding
3654 is forced internally for batched generation.
3655 return_type: 'input' (match input type), 'str', or 'tokens'.
3656 verbose: Show progress bar.
3657 stop_strings: Optional string or list of strings. A sequence stops once its
3658 generated text ends with one of them (HF StopStringCriteria). Requires a
3659 tokenizer. See generate() for details.
3660 stopping_criteria: Optional transformers StoppingCriteria, list, or
3661 StoppingCriteriaList, called as criterion(input_ids, scores) each step
3662 (scores is the step's logits). See generate() for the full contract.
3663 attention_mask: Optional ``[batch, pos]`` 0/1 mask over the prompt,
3664 required to generate correctly from an already-padded token tensor.
3665 Extended by one attended column per generated token on decoder-only
3666 models; forwarded to the encoder as-is on encoder-decoder models.
3667 See generate() for the full contract.
3668 forced_bos_token_id: Optional token id seeded as the first decoder token
3669 after ``decoder_start`` on encoder-decoder models, overriding any
3670 generation_config default. Raises ValueError on decoder-only models.
3672 Yields:
3673 Token tensors [batch, seq_len], or decoded text when return_type='str' -
3674 a bare string for a single sequence and one string per batch row for a
3675 larger batch, matching generate(). Chunks accumulate up to
3676 max_tokens_per_yield tokens between yields; the first yield includes the
3677 input tokens and subsequent yields contain only new tokens. For encoder-decoder
3678 models, the first yield includes the decoder start token (and configured
3679 forced BOS token), not the encoder input. Encoder-decoder streaming uses
3680 the uncached loop, as generate() does; custom stopping criteria are unsupported.
3681 """
3682 self._ensure_generation_supported("generate_stream")
3683 # --- Input parsing (mirrors generate()) ---
3684 _is_batched_list = isinstance(input, list) and len(input) > 1
3685 use_past_kv_cache = self._resolve_generation_caching(use_past_kv_cache, _is_batched_list)
3687 _encdec_early = hasattr(self.original_model, "config") and getattr(
3688 self.original_model.config, "is_encoder_decoder", False
3689 )
3690 encoder_attention_mask = None
3691 if _encdec_early:
3692 use_past_kv_cache = False
3693 if isinstance(input, str):
3694 if _encdec_early:
3695 # Native recipe: to_tokens' BOS policy corrupts encoder inputs.
3696 encoded = self.tokenizer(input, return_tensors="pt")
3697 input_tokens = encoded["input_ids"].to(self.cfg.device)
3698 encoder_attention_mask = encoded.get("attention_mask")
3699 else:
3700 input_tokens = self.to_tokens(
3701 input, prepend_bos=prepend_bos, move_to_device=True, truncate=False
3702 )
3703 input_type = "str"
3704 elif isinstance(input, list):
3705 if _encdec_early:
3706 encoded = self.tokenizer(input, return_tensors="pt", padding=True)
3707 input_tokens = encoded["input_ids"].to(self.cfg.device)
3708 encoder_attention_mask = encoded.get("attention_mask")
3709 elif _is_batched_list:
3710 _orig_ps = self.tokenizer.padding_side
3711 self.tokenizer.padding_side = "left"
3712 try:
3713 input_tokens = self.to_tokens(
3714 input, prepend_bos=prepend_bos, move_to_device=True, truncate=False
3715 )
3716 finally:
3717 self.tokenizer.padding_side = _orig_ps
3718 else:
3719 input_tokens = self.to_tokens(
3720 input, prepend_bos=prepend_bos, move_to_device=True, truncate=False
3721 )
3722 input_type = "list"
3723 else:
3724 input_tokens = input.to(self.cfg.device)
3725 input_type = "tokens"
3727 if return_type == "input":
3728 return_type = "str" if input_type in ["str", "list"] else "tokens"
3730 if forced_bos_token_id is not None and not _encdec_early:
3731 raise ValueError("forced_bos_token_id is only meaningful for encoder-decoder models")
3733 # An explicit mask wins over any tokenizer-derived one (mirrors generate()):
3734 # the encoder consumes it as-is, a decoder-only prompt mask is grown per step.
3735 initial_attention_mask: Optional[torch.Tensor] = None
3736 if attention_mask is not None:
3737 if _encdec_early:
3738 encoder_attention_mask = attention_mask
3739 else:
3740 if attention_mask.shape != input_tokens.shape:
3741 raise ValueError(
3742 f"attention_mask shape {tuple(attention_mask.shape)} does not "
3743 f"match the prompt shape {tuple(input_tokens.shape)}. Pass a 0/1 "
3744 "mask covering exactly the prompt tokens; generate_stream() "
3745 "extends it itself."
3746 )
3747 initial_attention_mask = attention_mask.to(self.cfg.device)
3749 batch_size = input_tokens.shape[0]
3751 # --- EOS setup ---
3752 stop_tokens: List[int] = []
3753 eos_token_for_padding = 0
3754 if stop_at_eos:
3755 tokenizer_has_eos_token = (
3756 self.tokenizer is not None and self.tokenizer.eos_token_id is not None
3757 )
3758 if eos_token_id is None:
3759 # Some chat models use a turn-end token that differs from the
3760 # tokenizer's primary EOS. Let adapters provide the full stop
3761 # set via cfg.eos_token_id; otherwise fall back to the tokenizer.
3762 eos_token_id = getattr(self.cfg, "eos_token_id", None)
3763 if eos_token_id is None:
3764 assert (
3765 tokenizer_has_eos_token
3766 ), "Must pass eos_token_id if stop_at_eos is True and tokenizer is None or has no eos_token_id"
3767 assert self.tokenizer is not None
3768 eos_token_id = self.tokenizer.eos_token_id
3769 if isinstance(eos_token_id, int):
3770 stop_tokens = [eos_token_id]
3771 eos_token_for_padding = eos_token_id
3772 else:
3773 stop_tokens = list(eos_token_id)
3774 if tokenizer_has_eos_token: 3774 ↛ 3775line 3774 didn't jump to line 3775 because the condition on line 3774 was never true
3775 assert self.tokenizer is not None
3776 eos_token_for_padding = self.tokenizer.eos_token_id
3777 else:
3778 eos_token_for_padding = eos_token_id[0]
3780 finished_sequences = torch.zeros(batch_size, dtype=torch.bool, device=self.cfg.device)
3782 stopping_criteria_list = self._resolve_stopping_criteria(stop_strings, stopping_criteria)
3783 if _encdec_early and stopping_criteria_list is not None:
3784 raise NotImplementedError(
3785 "stop_strings/stopping_criteria are not supported for encoder-decoder "
3786 "generation in TransformerBridge.generate_stream(). Call hf_generate(...) "
3787 "for HF-native stopping on those inputs."
3788 )
3789 if stopping_criteria_list is not None and not stop_at_eos:
3790 _pad_id = None
3791 if self.tokenizer is not None:
3792 _pad_id = (
3793 self.tokenizer.pad_token_id
3794 if self.tokenizer.pad_token_id is not None
3795 else self.tokenizer.eos_token_id
3796 )
3797 if _pad_id is not None:
3798 eos_token_for_padding = _pad_id
3799 elif batch_size > 1: 3799 ↛ 3808line 3799 didn't jump to line 3808 because the condition on line 3799 was always true
3800 raise ValueError(
3801 "Batched generate_stream with stopping_criteria and stop_at_eos=False "
3802 "needs a padding token to freeze finished rows, but no tokenizer pad/eos "
3803 "id is available. Set stop_at_eos=True or use a tokenizer with a pad/eos "
3804 "token."
3805 )
3807 # --- Cache setup ---
3808 if use_past_kv_cache:
3809 self._capture_hf_cache = True
3811 current_tokens = input_tokens.clone()
3812 encoder_input = input_tokens if _encdec_early else None
3813 decoder_tokens = None
3814 stream_prefix = input_tokens
3815 min_decoder_length = None
3816 if _encdec_early:
3817 decoder_tokens, min_decoder_length = self._encoder_decoder_setup(
3818 input_tokens, forced_bos_token_id
3819 )
3820 stream_prefix = decoder_tokens
3822 # --- Streaming loop ---
3823 # All yields are token tensors [batch, seq_len]. Each yield contains
3824 # only the newly generated tokens since the previous yield (the first
3825 # yield additionally prepends the input tokens for context).
3826 accumulated_tokens: Optional[torch.Tensor] = None
3827 tokens_since_last_yield = 0
3829 # Decoding each chunk alone splits any character whose bytes straddle a yield
3830 # boundary into U+FFFD, and concatenating chunks cannot rebuild it. Decode the
3831 # whole stream each time, emit only the new text, and hold back a trailing
3832 # incomplete character until the tokens that finish it arrive.
3833 decoded_history: Optional[torch.Tensor] = None
3834 emitted_chars = [0] * batch_size
3836 def _decode_delta(hold_incomplete: bool) -> List[str]:
3837 assert self.tokenizer is not None and decoded_history is not None
3838 deltas = []
3839 for row_idx, row in enumerate(decoded_history):
3840 text = self.tokenizer.decode(row, skip_special_tokens=True)
3841 if hold_incomplete:
3842 text = text.rstrip("�")
3843 deltas.append(text[emitted_chars[row_idx] :])
3844 emitted_chars[row_idx] = len(text)
3845 return deltas
3847 def _maybe_decode(tokens: torch.Tensor) -> Union[torch.Tensor, str, List[str]]:
3848 nonlocal decoded_history
3849 if return_type != "str":
3850 return tokens
3851 decoded_history = (
3852 tokens if decoded_history is None else torch.cat([decoded_history, tokens], dim=-1)
3853 )
3854 deltas = _decode_delta(hold_incomplete=True)
3855 return deltas[0] if len(deltas) == 1 else deltas
3857 def _flush_held() -> Optional[Union[str, List[str]]]:
3858 """Emit a withheld partial character at stream end rather than dropping it."""
3859 if return_type != "str" or decoded_history is None:
3860 return None
3861 deltas = _decode_delta(hold_incomplete=False)
3862 if not any(deltas):
3863 return None
3864 return deltas[0] if len(deltas) == 1 else deltas
3866 try:
3867 for step_idx, (sampled_tokens, _, all_finished) in enumerate(
3868 self._generate_tokens(
3869 current_tokens,
3870 input_tokens,
3871 batch_size,
3872 max_new_tokens=max_new_tokens,
3873 do_sample=do_sample,
3874 top_k=top_k,
3875 top_p=top_p,
3876 temperature=temperature,
3877 freq_penalty=freq_penalty,
3878 repetition_penalty=repetition_penalty,
3879 stop_at_eos=stop_at_eos,
3880 stop_tokens=stop_tokens,
3881 eos_token_for_padding=eos_token_for_padding,
3882 finished_sequences=finished_sequences,
3883 use_past_kv_cache=use_past_kv_cache,
3884 use_stateful_cache=False,
3885 mamba_cache=None,
3886 mamba_conv_kernel=0,
3887 is_encoder_decoder=_encdec_early,
3888 _is_batched_list=_is_batched_list,
3889 _generate_from_embeds=False,
3890 encoder_input=encoder_input,
3891 decoder_tokens=decoder_tokens,
3892 generated_token_ids=None,
3893 pixel_values=None,
3894 multimodal_kwargs={},
3895 verbose=verbose,
3896 stopping_criteria_list=stopping_criteria_list,
3897 initial_attention_mask=initial_attention_mask,
3898 min_decoder_length=min_decoder_length,
3899 ngram_processor=(self._encdec_ngram_processor() if _encdec_early else None),
3900 encoder_attention_mask=encoder_attention_mask,
3901 forced_eos_token_id=(
3902 self._encdec_forced_eos_token_id() if _encdec_early else None
3903 ),
3904 )
3905 ):
3906 new_tokens = sampled_tokens.unsqueeze(-1)
3908 if step_idx == 0:
3909 accumulated_tokens = torch.cat([stream_prefix, new_tokens], dim=-1)
3910 tokens_since_last_yield = accumulated_tokens.shape[1]
3911 else:
3912 if accumulated_tokens is None:
3913 accumulated_tokens = new_tokens
3914 else:
3915 accumulated_tokens = torch.cat([accumulated_tokens, new_tokens], dim=-1)
3916 tokens_since_last_yield += 1
3918 if tokens_since_last_yield >= max_tokens_per_yield:
3919 yield _maybe_decode(accumulated_tokens)
3920 tokens_since_last_yield = 0
3921 accumulated_tokens = None
3923 if all_finished:
3924 if accumulated_tokens is not None:
3925 yield _maybe_decode(accumulated_tokens)
3926 accumulated_tokens = None
3927 break
3929 # Yield remainder after loop completes without break
3930 if accumulated_tokens is not None:
3931 yield _maybe_decode(accumulated_tokens)
3932 held_back = _flush_held()
3933 if held_back is not None:
3934 yield held_back
3935 finally:
3936 self._capture_hf_cache = False
3937 if hasattr(self, "_last_hf_cache"): 3937 ↛ 3938line 3937 didn't jump to line 3938 because the condition on line 3937 was never true
3938 del self._last_hf_cache
3940 def hf_generate(
3941 self,
3942 input: str | list[str] | torch.Tensor = "",
3943 max_new_tokens: int = 10,
3944 stop_at_eos: bool = True,
3945 eos_token_id: int | Sequence[int] | None = None,
3946 do_sample: bool = True,
3947 top_k: int | None = None,
3948 top_p: float | None = None,
3949 temperature: float = 1.0,
3950 use_past_kv_cache: bool = True,
3951 return_type: str | None = "input",
3952 pixel_values: torch.Tensor | None = None,
3953 **generation_kwargs,
3954 ) -> str | list[str] | torch.Tensor | Any: # Any for HF ModelOutput types
3955 # Any: beartype forward ref limitation (beartype#546)
3956 """Generate text using the underlying HuggingFace model with full HF API support.
3958 This method provides direct access to HuggingFace's generation API, forwarding all
3959 generation parameters (including output_scores, output_logits, output_attentions,
3960 output_hidden_states) directly to the underlying HF model. Use this when you need
3961 full HuggingFace generation features not supported by the standard generate() method.
3963 For the standard TransformerLens generation interface, use generate() instead.
3965 Args:
3966 input: Text string, list of strings, or tensor of tokens
3967 max_new_tokens: Maximum number of tokens to generate
3968 stop_at_eos: If True, stop generating tokens when the model outputs eos_token.
3969 If False, disable EOS stopping, including the model's configured EOS.
3970 Custom stopping criteria still apply.
3971 eos_token_id: The token ID (or sequence of IDs) to stop on. Ignored
3972 when stop_at_eos=False, matching generate().
3973 do_sample: If True, sample from the model's output distribution
3974 top_k: Number of tokens to sample from
3975 top_p: Probability mass to sample from
3976 temperature: Temperature for sampling
3977 use_past_kv_cache: If True, use KV caching for faster generation
3978 return_type: The type of output to return - 'input', 'str', or 'tokens'
3979 **generation_kwargs: Additional HuggingFace generation parameters including:
3980 - output_scores: Return generation scores
3981 - output_logits: Return generation logits
3982 - output_attentions: Return attention weights
3983 - output_hidden_states: Return hidden states
3984 - return_dict_in_generate: Return ModelOutput object
3985 - And any other HF generation parameters
3987 Returns:
3988 Generated sequence as string, list of strings, tensor, or HF ModelOutput
3989 depending on input type, return_type, and generation_kwargs.
3991 Example::
3993 # Get full HF ModelOutput with logits and attentions
3994 from transformer_lens import TransformerBridge
3995 model = TransformerBridge.boot_transformers("tiny-stories-1M")
3996 result = model.hf_generate(
3997 "Hello world",
3998 max_new_tokens=5,
3999 output_logits=True,
4000 output_attentions=True,
4001 return_dict_in_generate=True
4002 )
4003 print(result.sequences) # Generated tokens
4004 print(result.logits) # Logits for each generation step
4005 print(result.attentions) # Attention weights
4006 """
4007 self._ensure_generation_supported("hf_generate")
4008 input_attention_mask: torch.Tensor | None = None
4009 # Handle string input by tokenizing it
4010 if isinstance(input, str):
4011 inputs = self.tokenizer(input, return_tensors="pt", padding=False, truncation=False).to(
4012 self.cfg.device
4013 )
4014 input_ids = inputs["input_ids"]
4015 input_type = "str"
4016 elif isinstance(input, list):
4017 is_encoder_decoder = getattr(
4018 getattr(self.original_model, "config", None), "is_encoder_decoder", False
4019 )
4020 tokenizer_kwargs = {} if is_encoder_decoder else {"padding_side": "left"}
4021 inputs = self.tokenizer(
4022 input,
4023 return_tensors="pt",
4024 padding=True,
4025 truncation=False,
4026 **tokenizer_kwargs,
4027 ).to(self.cfg.device)
4028 input_ids = inputs["input_ids"]
4029 input_attention_mask = inputs["attention_mask"]
4030 input_type = "list"
4031 else:
4032 input_ids = input
4033 if input_ids.device != self.cfg.device:
4034 input_ids = input_ids.to(self.cfg.device)
4035 input_type = "tokens"
4037 # Build generation_kwargs from explicit args and kwargs
4038 generation_kwargs = dict(generation_kwargs) if generation_kwargs is not None else {}
4039 if input_attention_mask is not None:
4040 generation_kwargs["attention_mask"] = input_attention_mask
4041 generation_kwargs.update(
4042 {
4043 "max_new_tokens": max_new_tokens,
4044 "do_sample": do_sample,
4045 "temperature": temperature,
4046 "pad_token_id": self.tokenizer.eos_token_id,
4047 }
4048 )
4050 if top_k is not None: 4050 ↛ 4051line 4050 didn't jump to line 4051 because the condition on line 4050 was never true
4051 generation_kwargs["top_k"] = top_k
4052 if top_p is not None: 4052 ↛ 4053line 4052 didn't jump to line 4053 because the condition on line 4052 was never true
4053 generation_kwargs["top_p"] = top_p
4054 if not stop_at_eos:
4055 # Omitting the kwarg would restore generation_config's EOS stopping.
4056 generation_kwargs["eos_token_id"] = None
4057 elif eos_token_id is not None:
4058 # cfg.eos_token_id can carry a chat model's full stop set; HF takes
4059 # int or list, so normalize any other sequence to a list.
4060 generation_kwargs["eos_token_id"] = (
4061 eos_token_id if isinstance(eos_token_id, int) else list(eos_token_id)
4062 )
4063 elif self.tokenizer.eos_token_id is not None: 4063 ↛ 4066line 4063 didn't jump to line 4066 because the condition on line 4063 was always true
4064 generation_kwargs["eos_token_id"] = self.tokenizer.eos_token_id
4066 if pixel_values is not None: 4066 ↛ 4067line 4066 didn't jump to line 4067 because the condition on line 4066 was never true
4067 generation_kwargs["pixel_values"] = pixel_values
4069 if use_past_kv_cache: 4069 ↛ 4073line 4069 didn't jump to line 4073 because the condition on line 4069 was always true
4070 generation_kwargs["use_cache"] = True
4072 # HF dict flags that trigger ModelOutput returns
4073 hf_dict_flags = (
4074 "output_scores",
4075 "output_logits",
4076 "output_attentions",
4077 "output_hidden_states",
4078 )
4080 # If any HF-style output flags are provided, ensure return_dict_in_generate is set
4081 any_flag_set = False
4082 for f in hf_dict_flags:
4083 if generation_kwargs.get(f) is not None:
4084 generation_kwargs[f] = bool(generation_kwargs[f])
4085 any_flag_set = True
4087 if any_flag_set:
4088 generation_kwargs.setdefault("return_dict_in_generate", True)
4090 with torch.no_grad():
4091 outputs = self.original_model.generate(input_ids, **generation_kwargs) # type: ignore[operator]
4093 # Check if output is a ModelOutput
4094 try:
4095 from transformers.utils import ModelOutput # type: ignore
4097 is_model_output = isinstance(outputs, ModelOutput)
4098 except Exception:
4099 is_model_output = False
4101 # Return based on return_type and input format
4102 if return_type == "input" or return_type is None:
4103 if input_type == "str":
4104 # Decode the full output back to string
4105 if is_model_output and hasattr(outputs, "sequences"): 4105 ↛ 4107line 4105 didn't jump to line 4107 because the condition on line 4105 was always true
4106 return self.tokenizer.decode(outputs.sequences[0], skip_special_tokens=True)
4107 return self.tokenizer.decode(outputs[0], skip_special_tokens=True)
4108 elif input_type == "list":
4109 # Decode each sequence in the batch
4110 if is_model_output and hasattr(outputs, "sequences"): 4110 ↛ 4115line 4110 didn't jump to line 4115 because the condition on line 4110 was always true
4111 return [
4112 self.tokenizer.decode(seq, skip_special_tokens=True)
4113 for seq in outputs.sequences
4114 ]
4115 return [self.tokenizer.decode(seq, skip_special_tokens=True) for seq in outputs]
4116 else:
4117 # Return the full token sequence including input
4118 return outputs
4119 elif return_type == "tokens": 4119 ↛ 4123line 4119 didn't jump to line 4123 because the condition on line 4119 was always true
4120 return outputs
4121 else:
4122 # For other return types, default to the decoded text
4123 if input_type == "str":
4124 if is_model_output and hasattr(outputs, "sequences"):
4125 return self.tokenizer.decode(outputs.sequences[0], skip_special_tokens=True)
4126 return self.tokenizer.decode(outputs[0], skip_special_tokens=True)
4127 elif input_type == "list":
4128 if is_model_output and hasattr(outputs, "sequences"):
4129 return [
4130 self.tokenizer.decode(seq, skip_special_tokens=True)
4131 for seq in outputs.sequences
4132 ]
4133 return [self.tokenizer.decode(seq, skip_special_tokens=True) for seq in outputs]
4134 else:
4135 return outputs
4137 def prepare_multimodal_inputs(
4138 self,
4139 text: Union[str, List[str]],
4140 images: Optional[Any] = None,
4141 ) -> Dict[str, torch.Tensor]:
4142 """Prepare multimodal inputs using the model's processor.
4144 Converts text and images into model-ready tensors (input_ids, pixel_values,
4145 attention_mask, etc.) using the HuggingFace processor loaded during boot().
4147 Args:
4148 text: Text prompt(s), typically containing image placeholder tokens
4149 (e.g., "<image>" for LLaVA).
4150 images: PIL Image or list of PIL Images to process. Pass None for
4151 text-only inputs on a multimodal model.
4153 Returns:
4154 Dictionary with 'input_ids', 'pixel_values', 'attention_mask', etc.
4155 All tensors are moved to the model's device.
4157 Raises:
4158 ValueError: If model is not multimodal or processor is not available.
4159 """
4160 if not getattr(self.cfg, "is_multimodal", False):
4161 raise ValueError(
4162 "prepare_multimodal_inputs() requires a multimodal model "
4163 "(cfg.is_multimodal must be True)"
4164 )
4165 if self.processor is None:
4166 raise ValueError(
4167 "No processor available. Load model with boot_transformers() or "
4168 "set bridge.processor = AutoProcessor.from_pretrained(...) manually."
4169 )
4170 inputs = self.processor(text=text, images=images, return_tensors="pt")
4171 return {k: v.to(self.cfg.device) if hasattr(v, "to") else v for k, v in inputs.items()}
4173 def to(self, *args, **kwargs) -> "TransformerBridge":
4174 """Move model to device and/or change dtype.
4176 Args:
4177 args: Positional arguments for nn.Module.to
4178 kwargs: Keyword arguments for nn.Module.to
4179 print_details: Whether to print details about device/dtype changes (default: True)
4181 Returns:
4182 Self for chaining
4183 """
4184 # Extract print_details if provided
4185 print_details = kwargs.pop("print_details", True)
4187 # Handle both device and dtype changes
4188 # torch.nn.Module.to() supports: to(device), to(dtype), to(device, dtype),
4189 # to(device=...), to(dtype=...), to(device=..., dtype=...)
4190 target_device, target_dtype = None, None
4192 if len(args) >= 1:
4193 first_arg = args[0]
4194 if isinstance(first_arg, (torch.device, str)):
4195 target_device = first_arg
4196 elif isinstance(first_arg, torch.dtype): 4196 ↛ 4198line 4196 didn't jump to line 4198 because the condition on line 4196 was always true
4197 target_dtype = first_arg
4198 if len(args) >= 2:
4199 second_arg = args[1]
4200 if isinstance(second_arg, torch.dtype): 4200 ↛ 4204line 4200 didn't jump to line 4204 because the condition on line 4200 was always true
4201 target_dtype = second_arg
4203 # these override positional args
4204 if "device" in kwargs: 4204 ↛ 4205line 4204 didn't jump to line 4205 because the condition on line 4204 was never true
4205 target_device = kwargs["device"]
4206 if "dtype" in kwargs:
4207 target_dtype = kwargs["dtype"]
4209 # Moving a multi-device (device_map-dispatched) model to a single device would
4210 # collapse the split and break accelerate's hook routing. Warn and drop the
4211 # device move; still honor dtype changes.
4212 if target_device is not None and getattr(self.cfg, "n_devices", 1) > 1:
4213 warnings.warn(
4214 f"TransformerBridge.to({target_device!r}) ignored: model is dispatched "
4215 f"across {self.cfg.n_devices} devices via device_map. Reload with "
4216 "device=... (and no device_map/n_devices) to move to a single device.",
4217 stacklevel=2,
4218 )
4219 target_device = None
4221 if target_device is not None:
4222 move_to_and_update_config(self, target_device, print_details)
4223 if target_dtype is not None:
4224 move_to_and_update_config(self, target_dtype, print_details)
4226 # Move the original model with all original args/kwargs (with print_details removed).
4227 # When we've nulled target_device for multi-GPU safety, strip device args so the
4228 # underlying module isn't moved either.
4229 if target_device is None and (len(args) > 0 or "device" in kwargs):
4230 kwargs.pop("device", None)
4231 # Filter positional args: drop devices/strings, keep dtypes.
4232 args = tuple(a for a in args if not isinstance(a, (torch.device, str)))
4233 self.original_model = self.original_model.to(*args, **kwargs)
4234 return self
4236 def cuda(self, device: Optional[Union[int, torch.device]] = None) -> "TransformerBridge":
4237 """Move model to CUDA.
4239 Args:
4240 device: CUDA device
4242 Returns:
4243 Self for chaining
4244 """
4245 if isinstance(device, int):
4246 return self.to(f"cuda:{device}")
4247 elif device is None:
4248 return self.to("cuda")
4249 else:
4250 return self.to(device)
4252 def cpu(self) -> "TransformerBridge":
4253 """Move model to CPU.
4255 Returns:
4256 Self for chaining
4257 """
4258 return self.to(torch.device("cpu"))
4260 def mps(self) -> "TransformerBridge":
4261 """Move model to MPS.
4263 Returns:
4264 Self for chaining
4265 """
4266 return self.to(torch.device("mps"))
4268 def train(self, mode: bool = True) -> "TransformerBridge":
4269 """Set training mode, propagating to the wrapped source model.
4271 ``original_model`` lives in ``__dict__`` rather than the registered
4272 module tree, so the inherited ``nn.Module.train()`` recursion does not
4273 reach it.
4274 """
4275 super().train(mode)
4276 original = getattr(self, "original_model", None)
4277 if isinstance(original, torch.nn.Module):
4278 original.train(mode)
4279 return self
4281 def set_use_attn_result(self, use_attn_result: bool):
4282 """Toggle whether to explicitly calculate and expose the result for each attention head.
4284 Useful for interpretability but can easily burn through GPU memory.
4285 """
4286 if use_attn_result:
4287 self._validate_attention_fork_supported("use_attn_result")
4288 self.cfg._set_bridge_managed_hook_flag("use_attn_result", use_attn_result)
4289 self._propagate_attention_flag("use_attn_result", use_attn_result)
4291 def set_use_split_qkv_input(self, use_split_qkv_input: bool):
4292 """Toggle independent residual copies for Q/K/V so each path can be patched alone.
4294 Mutually exclusive with `use_attn_in` — set that flag off first if it's on.
4295 """
4296 if use_split_qkv_input:
4297 if bool(getattr(self.cfg, "use_attn_in", False)):
4298 raise ValueError(
4299 "use_split_qkv_input and use_attn_in are mutually exclusive. "
4300 "Call set_use_attn_in(False) before enabling use_split_qkv_input."
4301 )
4302 self._validate_attention_fork_supported("use_split_qkv_input")
4303 self.cfg._set_bridge_managed_hook_flag("use_split_qkv_input", use_split_qkv_input)
4304 self._propagate_attention_flag("use_split_qkv_input", use_split_qkv_input)
4306 def set_use_attn_in(self, use_attn_in: bool):
4307 """Toggle a single 4D residual copy feeding all three Q/K/V projections.
4309 Mutually exclusive with `use_split_qkv_input` — set that flag off first
4310 if it's on. When on, `hook_attn_in` fires at
4311 `[batch, pos, n_heads, d_model]`, enabling coarse-grained interventions
4312 on the residual-stream copy shared across Q/K/V.
4313 """
4314 if use_attn_in:
4315 if bool(getattr(self.cfg, "use_split_qkv_input", False)):
4316 raise ValueError(
4317 "use_attn_in and use_split_qkv_input are mutually exclusive. "
4318 "Call set_use_split_qkv_input(False) before enabling use_attn_in."
4319 )
4320 self._validate_attention_fork_supported("use_attn_in")
4321 self.cfg._set_bridge_managed_hook_flag("use_attn_in", use_attn_in)
4322 self._propagate_attention_flag("use_attn_in", use_attn_in)
4324 def set_use_hook_mlp_in(self, use_hook_mlp_in: bool) -> None:
4325 """Toggle the ``hook_mlp_in`` HookPoint (the MLP-branch entry: pre-ln2, or
4326 the MLP input on post-norm blocks), matching legacy semantics.
4327 """
4328 self.cfg._set_bridge_managed_hook_flag("use_hook_mlp_in", use_hook_mlp_in)
4329 if not hasattr(self, "blocks"): 4329 ↛ 4330line 4329 didn't jump to line 4330 because the condition on line 4329 was never true
4330 return
4331 for block in self.blocks:
4332 block_cfg = getattr(block, "config", None)
4333 if block_cfg is not None and block_cfg is not self.cfg:
4334 try:
4335 self._write_propagated_hook_flag(block_cfg, "use_hook_mlp_in", use_hook_mlp_in)
4336 except (AttributeError, TypeError):
4337 pass
4338 block._use_hook_mlp_in = use_hook_mlp_in
4340 @staticmethod
4341 def _write_propagated_hook_flag(config: Any, flag_name: str, value: bool) -> None:
4342 """Write a cloned config flag without dispatching through its live Bridge."""
4343 if isinstance(config, TransformerBridgeConfig):
4344 config._set_bridge_managed_hook_flag(flag_name, value)
4345 else:
4346 object.__setattr__(config, flag_name, value)
4348 def _propagate_attention_flag(self, flag_name: str, value: bool) -> None:
4349 """Mirror `bridge.cfg.<flag>` onto every block's attention config.
4351 Some adapters (Llama family) deep-copy the block template during
4352 `setup_blocks_bridge`, cloning the attention bridge's config along
4353 with it. Others (Pythia, GPT-2) override `__deepcopy__` to share the
4354 config. Setting the flag only on `self.cfg` silently misses the
4355 cloned-config case. Propagating explicitly keeps both patterns
4356 honest — a no-op when configs are shared, a correctness fix when
4357 they aren't.
4358 """
4359 if not hasattr(self, "blocks"): 4359 ↛ 4360line 4359 didn't jump to line 4360 because the condition on line 4359 was never true
4360 return
4361 for block in self.blocks:
4362 attn = block._modules.get("attn") if hasattr(block, "_modules") else None
4363 if attn is None: 4363 ↛ 4364line 4363 didn't jump to line 4364 because the condition on line 4363 was never true
4364 continue
4365 attn_cfg = getattr(attn, "config", None)
4366 if attn_cfg is not None and attn_cfg is not self.cfg:
4367 try:
4368 self._write_propagated_hook_flag(attn_cfg, flag_name, value)
4369 except (AttributeError, TypeError):
4370 # Some config-like objects reject attributes even when
4371 # bypassing their custom __setattr__ implementation.
4372 pass
4374 def _validate_attention_fork_supported(self, flag_name: str) -> None:
4375 """Raise / warn if the model can't honor a fine-grained attention flag.
4377 The post-ln1 fork path lives on JointQKVAttentionBridge and
4378 PositionEmbeddingsAttentionBridge. Plain AttentionBridge delegates to
4379 HF and exposes no fork point; we raise rather than setting the flag
4380 silently. For hybrid models (some attention layers, some not), we warn
4381 and list which layers will honor the flag.
4382 """
4383 # Deferred imports: tight circular dependency with bridge setup.
4384 from transformer_lens.model_bridge.generalized_components.joint_qkv_attention import (
4385 JointQKVAttentionBridge,
4386 )
4387 from transformer_lens.model_bridge.generalized_components.position_embeddings_attention import (
4388 PositionEmbeddingsAttentionBridge,
4389 )
4391 if not hasattr(self, "blocks"): 4391 ↛ 4392line 4391 didn't jump to line 4392 because the condition on line 4391 was never true
4392 raise NotImplementedError(
4393 f"{flag_name}: this bridge has no `blocks` attribute, so no "
4394 "attention bridges to apply the flag to."
4395 )
4396 supported_classes = (JointQKVAttentionBridge, PositionEmbeddingsAttentionBridge)
4397 supporting_layers: list[int] = []
4398 attn_classes: set[str] = set()
4399 total_with_attn = 0
4400 for idx, block in enumerate(self.blocks):
4401 attn = block._modules.get("attn") if hasattr(block, "_modules") else None
4402 if attn is None: 4402 ↛ 4403line 4402 didn't jump to line 4403 because the condition on line 4402 was never true
4403 continue
4404 total_with_attn += 1
4405 attn_classes.add(type(attn).__name__)
4406 supports_flag = (
4407 bool(getattr(attn, "supports_attn_result", False))
4408 if flag_name == "use_attn_result"
4409 else isinstance(attn, supported_classes)
4410 )
4411 if supports_flag:
4412 supporting_layers.append(idx)
4413 if total_with_attn == 0: 4413 ↛ 4414line 4413 didn't jump to line 4414 because the condition on line 4413 was never true
4414 raise NotImplementedError(f"{flag_name}: no attention bridges found on self.blocks.")
4415 if not supporting_layers:
4416 if flag_name == "use_attn_result":
4417 capability_detail = "Per-head result computation is unavailable."
4418 else:
4419 capability_detail = "No hook point is available before the Q/K/V projections."
4420 raise NotImplementedError(
4421 f"{flag_name}: none of this model's attention bridges support "
4422 "the requested fine-grained attention hook. Found attention classes: "
4423 f"{sorted(attn_classes)}. Supported classes: "
4424 f"{[c.__name__ for c in supported_classes]}. {capability_detail}"
4425 )
4426 if len(supporting_layers) < total_with_attn: 4426 ↛ 4427line 4426 didn't jump to line 4427 because the condition on line 4426 was never true
4427 skipped = total_with_attn - len(supporting_layers)
4428 warnings.warn(
4429 f"{flag_name}: {skipped} of {total_with_attn} attention layers "
4430 "use an attention-bridge class that cannot honor this flag "
4431 f"(attention classes present: {sorted(attn_classes)}). "
4432 f"The flag will affect layers: {supporting_layers}.",
4433 stacklevel=3,
4434 )
4436 def _is_valid_bridge_path(self, hf_path: str) -> bool:
4437 """Check if a HuggingFace path corresponds to a valid bridge component.
4439 This validates that the path follows the bridge component structure and doesn't
4440 contain nested HuggingFace components that should have been wrapped.
4442 Args:
4443 hf_path: HuggingFace path after removing _original_component
4445 Returns:
4446 True if the path is valid, False if it contains nested HF components
4447 """
4448 parts = hf_path.split(".")
4450 # Get the component mapping for validation
4451 component_mapping = self.adapter.component_mapping
4452 if not component_mapping: 4452 ↛ 4453line 4452 didn't jump to line 4453 because the condition on line 4452 was never true
4453 return True # If no mapping, accept all keys
4455 # Walk through the path and check if each level is a registered bridge component
4456 # For example, transformer.h.0.mlp.in.weight should be valid
4457 # but transformer.h.0.mlp.c_fc.weight should be invalid (c_fc is nested HF component)
4459 # Start from the root
4460 current_component = None
4461 idx = 0
4463 # Find which top-level component this belongs to
4464 for tl_name, component in component_mapping.items():
4465 if component.name and hf_path.startswith(component.name + "."):
4466 current_component = component
4467 # Skip past the HF prefix
4468 remaining_path = hf_path[len(component.name) + 1 :]
4469 parts = remaining_path.split(".")
4470 idx = 0
4471 break
4473 if current_component is None:
4474 return True # Path doesn't match any component, let it through
4476 # Special handling for blocks
4477 if hasattr(current_component, "is_list_item") and current_component.is_list_item:
4478 # Skip the layer index
4479 if idx < len(parts) and parts[idx].isdigit(): 4479 ↛ 4483line 4479 didn't jump to line 4483 because the condition on line 4479 was always true
4480 idx += 1
4482 # Now validate the rest of the path against submodules
4483 while idx < len(parts): 4483 ↛ 4510line 4483 didn't jump to line 4510 because the condition on line 4483 was always true
4484 part = parts[idx]
4486 # If we hit 'weight' or 'bias', we're at a parameter - this is valid
4487 if part in ("weight", "bias"):
4488 return True
4490 # Check if this part is a registered submodule
4491 if hasattr(current_component, "submodules") and current_component.submodules:
4492 if part in current_component.submodules:
4493 current_component = current_component.submodules[part]
4494 idx += 1
4495 continue
4496 else:
4497 # This part is not a registered bridge component
4498 # It's likely a nested HF component (like c_fc, c_proj, c_attn)
4499 return False
4500 else:
4501 # No submodules to check, but not at a parameter yet
4502 # Check if next is weight/bias
4503 if idx + 1 < len(parts) and parts[idx + 1] in ("weight", "bias"):
4504 return True
4505 # Otherwise this is likely a nested HF component
4506 return False
4508 idx += 1
4510 return True
4512 def _normalize_bridge_key_to_hf(self, key: str) -> str:
4513 """Normalize a key that uses bridge attribute names to use HF module names.
4515 PyTorch's state_dict uses the Python attribute names (e.g., 'ln1')
4516 but the conversion logic expects HF module names (e.g., 'ln_1'). This
4517 function only replaces non-nested component names, leaving bridge
4518 subcomponents (like 'in', 'out', 'q', 'k', 'v') unchanged since they're
4519 handled by the component structure.
4521 Args:
4522 key: Key that may use bridge attribute names
4524 Returns:
4525 Key with attribute names replaced by module names where needed
4526 """
4527 component_mapping = self.adapter.component_mapping
4528 if not component_mapping: 4528 ↛ 4529line 4528 didn't jump to line 4529 because the condition on line 4528 was never true
4529 return key
4531 # Build a mapping of only the direct module attribute names to HF names
4532 # We only care about top-level and block-level component names, NOT subcomponents
4533 attr_to_hf = {}
4535 # Map top-level components
4536 block_list_names = {"blocks", "L_blocks", "H_blocks", "encoder_blocks", "decoder_blocks"}
4537 for tl_name, component in component_mapping.items():
4538 if component.name and tl_name not in block_list_names:
4539 # Skip if TL name is already a segment of its HF path (avoids doubling).
4540 if tl_name != component.name and tl_name not in component.name.split("."):
4541 attr_to_hf[tl_name] = component.name
4543 # Map block-level components (ln1, ln2, attn, mlp) for all block lists
4544 for bl_name in block_list_names:
4545 blocks_component = component_mapping.get(bl_name)
4546 if blocks_component and hasattr(blocks_component, "submodules"):
4547 for tl_subname, subcomponent in blocks_component.submodules.items():
4548 if subcomponent.name:
4549 # Only map if the names differ (e.g., ln1 -> ln_1, but attn -> attn)
4550 if tl_subname != subcomponent.name:
4551 attr_to_hf[tl_subname] = subcomponent.name
4553 # Replace only these specific attribute names in the key
4554 # We need to be careful to only replace whole path components, not substrings
4555 parts = key.split(".")
4556 result_parts = []
4558 for part in parts:
4559 if part in attr_to_hf:
4560 result_parts.append(attr_to_hf[part])
4561 else:
4562 result_parts.append(part)
4564 return ".".join(result_parts)
4566 def state_dict(self, destination=None, prefix="", keep_vars=False):
4567 """Get state dict with TransformerLens format keys.
4569 Converts HuggingFace format keys to TransformerLens format and filters out
4570 _original_component references and nested HuggingFace components.
4572 A direct no-argument call returns a clean state dict with bridge component
4573 paths converted to TL format. Calls that supply ``destination`` or
4574 ``prefix`` use standard ``nn.Module`` recursive semantics so a Bridge can
4575 compose inside a parent module.
4577 Args:
4578 destination: Optional dict to store state dict in
4579 prefix: Optional prefix to add to all keys
4580 keep_vars: Whether to keep variables as Variables instead of tensors
4582 Returns:
4583 Direct calls return TransformerLens-format keys; recursive calls
4584 return the supplied destination with standard module-tree keys.
4585 """
4586 if destination is not None or prefix:
4587 return super().state_dict(
4588 destination=destination,
4589 prefix=prefix,
4590 keep_vars=keep_vars,
4591 )
4593 raw_state_dict = self.original_model.state_dict(keep_vars=keep_vars)
4595 # Clean _original_component references and convert to TL format
4596 # Also filter out nested HuggingFace components that are wrapped by bridge components
4597 tl_state_dict = {}
4599 for key, value in raw_state_dict.items():
4600 # Skip _original_component keys
4601 if key == "_original_component" or key.startswith("_original_component."): 4601 ↛ 4602line 4601 didn't jump to line 4602 because the condition on line 4601 was never true
4602 continue
4604 # Remove all _original_component from the key
4605 clean_key = key.replace("._original_component", "")
4607 # Check if this is a valid bridge path (not a nested HF component)
4608 if not self._is_valid_bridge_path(clean_key):
4609 continue
4611 # Normalize bridge component names to HF names for conversion
4612 # (e.g., 'ln1' -> 'ln_1', 'mlp.in' -> 'mlp.c_fc')
4613 hf_key = self._normalize_bridge_key_to_hf(clean_key)
4615 # Convert to TL format - this uses the adapter's component_mapping
4616 tl_key = self.adapter.convert_hf_key_to_tl_key(hf_key)
4618 # Only add if we haven't seen this TL key yet (handles duplicates)
4619 if tl_key not in tl_state_dict:
4620 tl_state_dict[tl_key] = value
4622 return tl_state_dict
4624 def _tl_key_to_actual_keys(self) -> dict[str, list[str]]:
4625 """Inverse of the renaming state_dict() applies: map each TL-format key
4626 back to every raw parameter/buffer path that represents it.
4628 Mirrors the filtering and key-conversion in state_dict() exactly, except
4629 it keeps every raw key for a given TL key instead of only the first-seen
4630 one. Bridge components frequently expose the same underlying parameter
4631 through more than one attribute path (e.g. GPT-2's split q/k/v weights
4632 are views into the wrapped module's combined c_attn weight, reachable
4633 both via a block-level shortcut and via the nested _original_component
4634 chain) - all of those aliases must be written for the round trip to
4635 actually change what forward() reads, not just what state_dict() shows.
4636 """
4637 mapping: dict[str, list[str]] = {}
4638 for actual_key in self.original_model.state_dict():
4639 if actual_key == "_original_component" or actual_key.startswith("_original_component."): 4639 ↛ 4640line 4639 didn't jump to line 4640 because the condition on line 4639 was never true
4640 continue
4641 clean_key = actual_key.replace("._original_component", "")
4642 if not self._is_valid_bridge_path(clean_key):
4643 continue
4644 hf_key = self._normalize_bridge_key_to_hf(clean_key)
4645 tl_key = self.adapter.convert_hf_key_to_tl_key(hf_key)
4646 mapping.setdefault(tl_key, []).append(actual_key)
4647 return mapping
4649 def load_state_dict(self, state_dict, strict=True, assign=False):
4650 """Load state dict into the model, handling both clean keys and original keys with _original_component references.
4652 Accepts three key formats: TL-format keys as emitted by state_dict()
4653 (e.g. "blocks.0.attn.q.weight"), raw native parameter paths (e.g. for
4654 ``boot_native`` / tracr-style loading), and raw paths with
4655 "_original_component" segments stripped.
4657 Args:
4658 state_dict: Dictionary containing a whole state of the module
4659 strict: Whether to strictly enforce that the keys in state_dict match the keys returned by this module's state_dict() function
4660 assign: Whether to assign items in the state dictionary to their corresponding keys in the module instead of copying them
4662 Returns:
4663 NamedTuple with missing_keys and unexpected_keys fields
4664 """
4665 current_state_dict = self.original_model.state_dict(keep_vars=True)
4666 clean_to_actual = {}
4667 for actual_key in current_state_dict.keys():
4668 if actual_key != "_original_component": 4668 ↛ 4667line 4668 didn't jump to line 4667 because the condition on line 4668 was always true
4669 clean_to_actual[actual_key.replace("._original_component", "")] = actual_key
4671 tl_to_actual = self._tl_key_to_actual_keys()
4673 mapped_state_dict = {}
4674 unexpected_keys = []
4675 for input_key, value in state_dict.items():
4676 if input_key in current_state_dict:
4677 mapped_state_dict[input_key] = value
4678 elif input_key in clean_to_actual:
4679 mapped_state_dict[clean_to_actual[input_key]] = value
4680 elif input_key in tl_to_actual:
4681 for actual_key in tl_to_actual[input_key]:
4682 mapped_state_dict[actual_key] = value
4683 else:
4684 unexpected_keys.append(input_key)
4686 # A TL key's actual-key aliases share the same underlying storage (see
4687 # _tl_key_to_actual_keys), so writing any one of them already updates
4688 # what forward() reads for all of them. Treat the group as satisfied
4689 # if any alias was written -- e.g. a caller supplying clean/raw keys
4690 # (the branch above maps each clean key to exactly one actual key)
4691 # shouldn't have the *other*, unwritten aliases reported as missing.
4692 missing_keys = sorted(
4693 actual_key
4694 for actual_keys in tl_to_actual.values()
4695 if not any(k in mapped_state_dict for k in actual_keys)
4696 for actual_key in actual_keys
4697 )
4699 if strict and (missing_keys or unexpected_keys):
4700 error_msgs = []
4701 if unexpected_keys:
4702 error_msgs.append(
4703 "Unexpected key(s) in state_dict: "
4704 + ", ".join(f'"{k}"' for k in sorted(unexpected_keys))
4705 )
4706 if missing_keys:
4707 error_msgs.append(
4708 "Missing key(s) in state_dict: " + ", ".join(f'"{k}"' for k in missing_keys)
4709 )
4710 raise RuntimeError(
4711 "Error(s) in loading state_dict for {}:\n\t{}".format(
4712 type(self.original_model).__name__, "\n\t".join(error_msgs)
4713 )
4714 )
4716 if not assign:
4717 result = self.original_model.load_state_dict(
4718 mapped_state_dict, strict=False, assign=False
4719 )
4720 return type(result)(missing_keys=missing_keys, unexpected_keys=unexpected_keys)
4722 # assign=True normally makes nn.Module.load_state_dict *replace* each
4723 # target parameter/buffer with the incoming tensor rather than copying
4724 # into existing storage. Any key whose current tensor shares storage
4725 # with another key -- a split QKV/gate-up component's view into a
4726 # combined weight, the combined weight itself (writing it directly
4727 # would silently orphan the views), or a tied pair like embed/unembed
4728 # -- would desync from whatever it shares storage with: the bridge
4729 # might keep reading a stale value, or a live view not even part of
4730 # this load would silently keep the pre-load data while
4731 # original_model.state_dict() (what save_pretrained() exports) shows
4732 # the new one. Route every key in a storage-sharing group through an
4733 # explicit in-place .data.copy_() instead, regardless of the caller's
4734 # assign=True, so those relationships survive.
4735 shared_keys = _storage_group_keys(current_state_dict)
4737 copy_items: dict[str, tuple[torch.Tensor, torch.Tensor]] = {}
4738 passthrough_items = {}
4739 errors = []
4740 for key, value in mapped_state_dict.items():
4741 target = current_state_dict.get(key)
4742 if target is None or key not in shared_keys:
4743 passthrough_items[key] = value
4744 continue
4745 problems = []
4746 if target.is_meta: 4746 ↛ 4750line 4746 didn't jump to line 4750 because the condition on line 4746 was never true
4747 # copy_() onto a meta tensor silently no-ops rather than
4748 # raising, so a meta target would otherwise report a
4749 # successful load while never actually materializing.
4750 problems.append(
4751 "the current parameter is a meta tensor; an in-place copy "
4752 "can't materialize it, and this key shares storage with "
4753 "another parameter so ordinary assign=True replacement "
4754 "isn't safe here either"
4755 )
4756 else:
4757 if tuple(target.shape) != tuple(value.shape):
4758 problems.append(f"shape {tuple(value.shape)} != expected {tuple(target.shape)}")
4759 if target.dtype != value.dtype:
4760 problems.append(f"dtype {value.dtype} != expected {target.dtype}")
4761 if target.device != value.device:
4762 problems.append(f"device {value.device} != expected {target.device}")
4763 if problems:
4764 errors.append(f'"{key}": ' + "; ".join(problems))
4765 else:
4766 copy_items[key] = (target, value)
4768 if errors:
4769 raise RuntimeError(
4770 "Cannot load the following key(s) with assign=True: each shares "
4771 "storage with another parameter/buffer (e.g. a split QKV/"
4772 "gate-up component's view into a combined weight, or a tied "
4773 "pair like embed/unembed), so it can only be loaded via an "
4774 "in-place copy, which requires an exact match.\n\t" + "\n\t".join(errors)
4775 )
4777 for target, value in copy_items.values():
4778 with torch.no_grad():
4779 target.data.copy_(value)
4781 result = self.original_model.load_state_dict(passthrough_items, strict=False, assign=True)
4782 refresh_container_state_owners(self)
4783 return type(result)(missing_keys=missing_keys, unexpected_keys=unexpected_keys)
4785 def get_params(self):
4786 """Access to model parameters in the format expected by SVDInterpreter.
4788 For missing weights, returns zero tensors of appropriate shape instead of raising exceptions.
4789 This ensures compatibility across different model architectures.
4791 Returns:
4792 dict: Dictionary of parameter tensors with TransformerLens naming convention
4794 Raises:
4795 ValueError: If configuration is inconsistent (e.g., cfg.n_layers != len(blocks))
4796 """
4797 return get_bridge_params(self)
4799 # NOTE: list_supported_models and check_model_support are attached to this class
4800 # dynamically by transformer_lens.model_bridge.sources.transformers module.
4801 # These are HuggingFace-specific methods that belong in the transformers source module.