Coverage for transformer_lens/model_bridge/sources/inspect/source.py: 93%
60 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-09-21 19:27 +0000
1"""``boot_inspect`` — wrap an ``inspect_ai`` provider in a RemoteBridge via InspectDriver."""
2from __future__ import annotations
4import logging
5import warnings
6from typing import Any, Optional
8import torch
10from transformer_lens.factories.architecture_adapter_factory import (
11 ArchitectureAdapterFactory,
12)
13from transformer_lens.model_bridge.remote_bridge import RemoteBridge
14from transformer_lens.model_bridge.sources._bridge_builder import (
15 build_bridge_config_from_hf,
16 configure_tokenizer,
17 skip_tokenizer_for_modality,
18)
19from transformer_lens.model_bridge.sources._hf_format import (
20 determine_architecture_from_hf_config,
21)
22from transformer_lens.utilities.hf_utils import get_hf_token
24from . import profiles
25from .driver import InspectDriver
27# Providers that expose the structural self-check + capture wire format the InspectDriver
28# consumes. boot_inspect queries supported_kinds on these; others route via for_provider.
29_TL_BRIDGE_PROVIDERS = {"tl_bridge", "tl_bridge_vllm"}
32def boot_inspect(
33 model_name: str,
34 tokenizer: Optional[Any] = None,
35 dtype: Optional[torch.dtype] = None,
36 provider: str = "tl_bridge",
37 **inspect_kwargs: Any,
38) -> RemoteBridge:
39 """Boot a model via an ``inspect_ai`` provider and wrap it in a :class:`RemoteBridge`.
41 The driver is provider-agnostic: ``provider`` defaults to our own HF-backed
42 ``tl_bridge`` provider (residual/attn/mlp capture + full affine interventions +
43 full-sequence logits); ``"vllm-lens"`` targets a running vllm-lens vLLM provider
44 (residual-only, additive-steering-only) — wire-aligned with its documented format,
45 but not yet verified against a live provider.
47 Fireable hooks (``tl_bridge``, TransformerBridge-native names): ``blocks.{i}.hook_in``
48 (resid_pre) / ``ln2.hook_in`` (resid_mid) /
49 ``hook_out`` (resid_post) / ``attn.hook_out`` / ``mlp.hook_out``, plus the head-split
50 attention hooks where the structural probe finds them: ``attn.hook_q/k/v`` (pre-RoPE
51 projection outputs; separate-projection archs only — fused qkv gates them),
52 ``attn.hook_z`` (out-projection input), and ``attn.hook_pattern`` (post-softmax,
53 capture-only, eager attention required). The provider runs a structural self-check per
54 model and gates any boundary it can't serve faithfully: ``resid_mid`` for
55 parallel-residual or norm-variant blocks, ``attn_out``/``mlp_out`` when their submodule
56 isn't locatable (it warns when it gates one). ``embed``, ``ln_final``, and
57 ``attn.hook_attn_scores`` are always non-fireable — use ``boot_transformers()`` for
58 those.
60 For parity with ``boot_transformers`` the provider loads with the same dtype (fp32 by
61 default) and eager attention. Full-sequence logits ride on ``return_logits=True`` (the
62 default); pass ``return_logits=False`` to skip the (seq × d_vocab) payload for pure
63 activation capture (``run_with_cache`` keeps them since it returns logits).
64 """
65 from inspect_ai.model import get_model
66 from transformers import AutoConfig, AutoTokenizer
68 from . import ( # noqa: F401 — import registers @modelapi
69 transformers_provider as _provider,
70 )
72 hf_token = get_hf_token()
73 hf_config = AutoConfig.from_pretrained(model_name, token=hf_token)
74 # Shared resolution (not architectures[0]) handles architectures=None configs via
75 # model_type and rejects unsupported archs before the provider loads weights.
76 architecture = determine_architecture_from_hf_config(hf_config)
77 # Default fp32 to match boot_transformers (which loads/casts fp32 regardless of the
78 # config's native dtype); an explicit dtype still wins.
79 resolved_dtype = dtype if dtype is not None else torch.float32
81 bridge_config = build_bridge_config_from_hf(hf_config, architecture, model_name, resolved_dtype)
82 adapter = ArchitectureAdapterFactory.select_architecture_adapter(bridge_config)
83 if tokenizer is None and not skip_tokenizer_for_modality(adapter.cfg):
84 tokenizer = AutoTokenizer.from_pretrained(model_name, token=hf_token)
85 if tokenizer is not None:
86 # Match boot_transformers' tokenizer setup so to_tokens(str) is token-identical.
87 tokenizer = configure_tokenizer(tokenizer, adapter.cfg)
89 if provider == "tl_bridge":
90 # The provider's raw HF forward must match boot_transformers' load: same dtype,
91 # eager attention (TL forces eager — SDPA/flash diverge and accumulate with depth),
92 # and auth/remote-code so gated/custom models load at all.
93 inspect_kwargs["model_kwargs"] = _provider_model_kwargs(
94 dict(inspect_kwargs.get("model_kwargs", {})), adapter, resolved_dtype, hf_token
95 )
96 elif provider == "tl_bridge_vllm": 96 ↛ 104line 96 didn't jump to line 104 because the condition on line 96 was always true
97 # Otherwise the provider defaults to the HF-config dtype and bridge_config.dtype
98 # lies about what the engine actually loaded.
99 inspect_kwargs["dtype"] = resolved_dtype
101 # memoize=False: inspect_ai caches get_model by name, which would (a) return a stale
102 # model ignoring a changed dtype/kwargs on re-boot and (b) keep weights resident past
103 # close(). Each boot must honor its own args and own its model's lifecycle.
104 model = get_model(f"{provider}/{model_name}", memoize=False, **inspect_kwargs)
105 # Both TL-bridge providers (HF + vLLM) restrict their profile to the boundaries the
106 # provider's structural self-check / overlay found this model can serve; warn only
107 # if it gated something. Third-party providers (e.g. vllm-lens) route via for_provider.
108 if provider in _TL_BRIDGE_PROVIDERS: 108 ↛ 123line 108 didn't jump to line 123 because the condition on line 108 was always true
109 api = getattr(model, "api", None)
110 kinds = None
111 note = ""
112 # Default True for back-compat; vLLM provider sets False so RemoteBridge.forward
113 # rejects loss/both (otherwise loss over -inf earlier positions silently NaNs).
114 psl = True
115 if api is not None: 115 ↛ 119line 115 didn't jump to line 119 because the condition on line 115 was always true
116 kinds = api.supported_kinds() if hasattr(api, "supported_kinds") else None
117 note = api.capability_note() if hasattr(api, "capability_note") else ""
118 psl = bool(getattr(api, "provides_sequence_logits", True))
119 profile = profiles.TLBridgeProfile(supported_kinds=kinds, provides_sequence_logits=psl)
120 if note:
121 warnings.warn(note, UserWarning, stacklevel=2)
122 else:
123 profile = profiles.for_provider(provider)
124 driver = InspectDriver(model=model, adapter=adapter, tokenizer=tokenizer, profile=profile)
125 bridge = RemoteBridge(adapter=adapter, tokenizer=tokenizer, driver=driver)
126 _log_hook_summary(model_name, architecture, provider, driver)
127 return bridge
130def _provider_model_kwargs(
131 model_kwargs: dict[str, Any], adapter: Any, dtype: torch.dtype, hf_token: Optional[str]
132) -> dict[str, Any]:
133 """Load kwargs for the HF provider that mirror boot_transformers, so the provider's
134 raw forward matches the bridge. Caller-supplied keys win (setdefault)."""
135 model_kwargs.setdefault("torch_dtype", dtype)
136 # boot_transformers forces eager unless the adapter pins an implementation.
137 model_kwargs.setdefault(
138 "attn_implementation", getattr(adapter.cfg, "attn_implementation", None) or "eager"
139 )
140 if hf_token: 140 ↛ 142line 140 didn't jump to line 142 because the condition on line 140 was always true
141 model_kwargs.setdefault("token", hf_token)
142 return model_kwargs
145def _log_hook_summary(
146 model_name: str, architecture: str, provider: str, driver: InspectDriver
147) -> None:
148 log = logging.getLogger("transformer_lens.inspect")
149 fireable = sorted(driver.supported_hook_points)
150 log.info(
151 "Inspect source on %s (%s) via provider %r serves %d fireable hook(s).",
152 model_name,
153 architecture,
154 provider,
155 len(fireable),
156 )
159__all__ = ["boot_inspect"]