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