Coverage for transformer_lens/model_bridge/remote_bridge.py: 84%
99 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-09-21 19:27 +0000
1"""Non-torch bridge: vLLM workers, Inspect remote providers."""
2from __future__ import annotations
4import warnings
5from typing import Any, Callable, List, Optional, Tuple, Union
7from transformer_lens.hook_points import HookIntrospectionMixin, HookPoint
8from transformer_lens.model_bridge.bridge_core import BridgeCore
9from transformer_lens.model_bridge.driver_protocol import (
10 ForwardResult,
11 TensorLike,
12 to_torch,
13 validate_driver,
14)
17class RemoteBridge(BridgeCore, HookIntrospectionMixin):
18 """Bridge for backends with no local ``nn.Module`` (vLLM, Inspect).
20 No nn.Module parentage strips the torch-only surface; driver pre-declares
21 ``supported_hook_points`` (no model to walk).
22 """
24 def __init__(
25 self,
26 adapter: Any,
27 tokenizer: Any,
28 driver: Any,
29 ) -> None:
30 if not driver.supported_hook_points:
31 raise ValueError(
32 "RemoteBridge requires driver.supported_hook_points to be "
33 "non-empty: non-torch drivers own the hook namespace because "
34 "there is no local model for the bridge to walk."
35 )
36 BridgeCore.__init__(self, adapter, tokenizer, driver)
37 # No local device; tensor.to(None) is a no-op so downstream patterns degrade cleanly.
38 self.cfg.device = None
39 # HookPoint is nn.Module-backed but RemoteBridge isn't an nn.Module —
40 # named_modules() walks don't apply; only registry lookup matters.
41 for name in driver.supported_hook_points:
42 hp = HookPoint()
43 hp.name = name
44 self._hook_registry[name] = hp
45 self._hook_registry_initialized = True
46 validate_driver(self._driver, after_bridge_construction=True)
48 @staticmethod
49 def boot_vllm(*args: Any, **kwargs: Any) -> "RemoteBridge":
50 """Boot a model via vLLM. Returns a RemoteBridge wrapping a VLLMDriver.
52 Mirrors ``TransformerBridge.boot_transformers``. Lazy import so
53 ``remote_bridge`` itself stays vLLM-agnostic — only callers of this
54 method need vLLM installed. See :func:`sources.vllm.boot_vllm` for kwargs.
55 """
56 from .sources.vllm import boot_vllm as _boot_vllm
58 return _boot_vllm(*args, **kwargs)
60 @staticmethod
61 def boot_inspect(*args: Any, **kwargs: Any) -> "RemoteBridge":
62 """Boot a model via an inspect_ai provider. Returns a RemoteBridge wrapping
63 an InspectDriver. Lazy import keeps remote_bridge inspect-agnostic. See
64 :func:`sources.inspect.boot_inspect` for kwargs."""
65 from .sources.inspect import boot_inspect as _boot_inspect
67 return _boot_inspect(*args, **kwargs)
69 def _scan_existing_hooks(self, module: Any, prefix: str = "") -> None:
70 """No-op: registry built from driver declarations in __init__."""
72 def forward(
73 self,
74 input: Any = None,
75 *,
76 return_type: str | None = "logits",
77 loss_per_token: bool = False,
78 labels: Any = None,
79 **kwargs: Any,
80 ) -> Any:
81 """Tokenize → driver.forward → replay captures → finalize per return_type.
83 Explicit ``labels`` use shifted causal loss; remote encoder-decoder loss
84 is unsupported.
85 """
86 # Early copy of _finalize_return's gate — fail before the wasted remote forward.
87 self._check_loss_supported(return_type)
88 self._reject_stop_at_layer(kwargs.pop("stop_at_layer", None))
89 if isinstance(input, str):
90 kwargs["input_ids"] = self.to_tokens(input) # BOS-aware, matches boot_transformers
91 elif isinstance(input, list) and any(isinstance(item, str) for item in input):
92 # Would otherwise be treated as raw input_ids and crash deep in numpy.
93 raise TypeError(
94 "RemoteBridge.forward received a list of strings; batched string "
95 "input is unsupported here. Pass a single str, or tokenize with "
96 "to_tokens() and pass token ids."
97 )
98 elif input is not None: 98 ↛ 103line 98 didn't jump to line 103 because the condition on line 98 was always true
99 kwargs["input_ids"] = input
101 # Only request hooks with a registered handler, so a plain forward(tokens)
102 # ships logits alone instead of the full residual decomposition every call.
103 if "capture" not in kwargs: 103 ↛ 108line 103 didn't jump to line 108 because the condition on line 103 was always true
104 kwargs["capture"] = tuple(
105 name for name, hp in self._hook_registry.items() if hp.fwd_hooks
106 )
108 result: ForwardResult = self._driver.forward(**kwargs)
109 if result.captured:
110 self._replay_captures(result.captured)
112 logits: Any = result.logits
113 if logits is not None and not isinstance(logits, TensorLike): 113 ↛ 114line 113 didn't jump to line 114 because the condition on line 113 was never true
114 return logits # weird shape — let caller handle
115 if logits is not None:
116 logits = to_torch(logits)
117 if labels is not None:
118 if not isinstance(labels, TensorLike): 118 ↛ 119line 118 didn't jump to line 119 because the condition on line 118 was never true
119 raise TypeError(f"labels must be tensor-like, got {type(labels).__name__}")
120 labels = to_torch(labels)
122 return self._finalize_return(
123 return_type,
124 logits,
125 kwargs.get("input_ids"),
126 attention_mask=kwargs.get("attention_mask"),
127 labels=labels,
128 is_audio_model=getattr(self.cfg, "is_audio_model", False),
129 is_visual_model=getattr(self.cfg, "is_visual_model", False),
130 loss_per_token=loss_per_token,
131 )
133 def run_with_hooks(
134 self,
135 input: Any,
136 fwd_hooks: List[Tuple[Union[str, Callable], Callable]] = [],
137 bwd_hooks: List[Tuple[Union[str, Callable], Callable]] = [],
138 reset_hooks_end: bool = True,
139 clear_contexts: bool = False,
140 return_type: Optional[str] = "logits",
141 stop_at_layer: Optional[int] = None,
142 start_at_layer: Optional[int] = None,
143 remove_batch_dim: bool = False,
144 **kwargs: Any,
145 ) -> Any:
146 """Run with hooks. Remote fwd_hooks fire post-forward on captured
147 activations (read-only) — they can't alter the computation, so warn; use
148 ``intervene=`` specs to mutate. bwd_hooks are unsupported (no backward)."""
149 if bwd_hooks:
150 raise NotImplementedError(
151 "RemoteBridge has no backward pass; bwd_hooks are unsupported."
152 )
153 self._reject_stop_at_layer(stop_at_layer)
154 self._reject_start_at_layer(start_at_layer)
155 if fwd_hooks:
156 warnings.warn(
157 "RemoteBridge fwd_hooks fire on already-captured activations (read-only): "
158 "a hook that returns a modified tensor does NOT change the forward "
159 "computation or logits — the return is discarded. To intervene on the "
160 "computation, pass intervene={hook_name: {'op': ...}} to "
161 "forward()/run_with_cache().",
162 UserWarning,
163 stacklevel=2,
164 )
165 return super().run_with_hooks(
166 input,
167 fwd_hooks=fwd_hooks,
168 bwd_hooks=bwd_hooks,
169 reset_hooks_end=reset_hooks_end,
170 clear_contexts=clear_contexts,
171 return_type=return_type,
172 remove_batch_dim=remove_batch_dim,
173 **kwargs,
174 )
176 @staticmethod
177 def _reject_stop_at_layer(stop_at_layer: Any) -> None:
178 # BridgeCore's stop hooks need a local `blocks` tree; silently ignoring
179 # the kwarg would lie about compute cost.
180 if stop_at_layer is not None:
181 raise NotImplementedError(
182 "RemoteBridge does not support stop_at_layer: the remote engine "
183 "always runs the full forward pass."
184 )
186 @staticmethod
187 def _reject_start_at_layer(start_at_layer: Any) -> None:
188 # Residual-stream injection needs a local `blocks` tree; the remote engine
189 # only runs a full forward from tokens.
190 if start_at_layer is not None: 190 ↛ 191line 190 didn't jump to line 191 because the condition on line 190 was never true
191 raise NotImplementedError(
192 "RemoteBridge does not support start_at_layer: the remote engine "
193 "always runs the full forward pass from tokens."
194 )
196 def run_with_cache(self, *args: Any, **kwargs: Any) -> Any:
197 """Cache via driver captures; stop/start_at_layer and incl_bwd rejected."""
198 self._reject_stop_at_layer(kwargs.get("stop_at_layer"))
199 self._reject_start_at_layer(kwargs.get("start_at_layer"))
200 if kwargs.get("incl_bwd"): 200 ↛ 203line 200 didn't jump to line 203 because the condition on line 200 was never true
201 # Captures come back as detached tensors from the remote engine, so there
202 # is no local graph for backward() to walk.
203 raise NotImplementedError(
204 "RemoteBridge has no backward pass; run_with_cache(incl_bwd=True) is "
205 "unsupported."
206 )
207 return super().run_with_cache(*args, **kwargs)
209 def to_tokens(self, input: Any, prepend_bos: bool | None = None, truncate: bool = True) -> Any:
210 """Tokenize a string with the same BOS handling as ``TransformerBridge``.
212 Mirrors ``cfg.default_prepend_bos`` / ``tokenizer_prepends_bos`` so
213 ``boot_inspect(m).run_with_cache("text")`` matches ``boot_transformers(m)``
214 on the same string — a bare ``encode`` (no BOS) would silently diverge.
215 """
216 from transformer_lens import utilities as utils
218 assert self.tokenizer is not None, "Tokenizer must be set."
219 if prepend_bos is None: 219 ↛ 221line 219 didn't jump to line 221 because the condition on line 219 was always true
220 prepend_bos = getattr(self.cfg, "default_prepend_bos", True)
221 tokenizer_prepends_bos = getattr(self.cfg, "tokenizer_prepends_bos", True)
222 if prepend_bos and not tokenizer_prepends_bos: 222 ↛ 223line 222 didn't jump to line 223 because the condition on line 222 was never true
223 bos = self.tokenizer.bos_token
224 encodes_atomically = (
225 bos is not None
226 and len(self.tokenizer(bos, add_special_tokens=False)["input_ids"]) == 1
227 )
228 if encodes_atomically:
229 input = utils.get_input_with_manually_prepended_bos(bos, input)
230 # else: the fallback BOS is not an atom in this vocab (e.g.
231 # '<|endoftext|>' installed on BERT); prepending the string would
232 # tokenize to subword garbage, so skip rather than pollute the input.
233 if isinstance(input, str): 233 ↛ 237line 233 didn't jump to line 237 because the condition on line 233 was always true
234 input = [input]
235 # A single sequence never needs padding; only pad when batching (which also
236 # avoids requiring a pad_token on tokenizers that lack one, e.g. raw gpt2).
237 tokens = self.tokenizer(
238 input,
239 return_tensors="pt",
240 padding=len(input) > 1,
241 truncation=truncate,
242 max_length=self.cfg.n_ctx if truncate else None,
243 )["input_ids"]
244 if not prepend_bos and tokenizer_prepends_bos: 244 ↛ 245line 244 didn't jump to line 245 because the condition on line 244 was never true
245 tokens = utils.get_tokens_with_bos_removed(self.tokenizer, tokens)
246 return tokens
248 def __enter__(self) -> "RemoteBridge":
249 """Use as a context manager so the engine is released on exit:
250 ``with RemoteBridge.boot_vllm(...) as bridge: ...``."""
251 return self
253 def __exit__(self, *exc: Any) -> None:
254 self.close()
256 def __del__(self) -> None:
257 # Best-effort safety net for notebooks that drop the bridge without
258 # close() — repeated boot_vllm would otherwise OOM. close() is idempotent.
259 try:
260 self.close()
261 except Exception:
262 pass