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