Coverage for transformer_lens/model_bridge/sources/inspect/vllm_provider.py: 85%
196 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-09-21 19:27 +0000
1"""vLLM-backed ``inspect_ai`` model provider, registered as ``tl_bridge_vllm``.
3A sibling to the HF-backed ``tl_bridge`` provider: instead of running an HF causal LM
4locally, it generates via vLLM (PagedAttention + continuous batching) — so it scales to
5parallel-sample evals and dataset-scale workloads where the HF provider serializes.
7Inherits ``generate()`` dispatch / message-rendering / per-turn-capture validation from
8:class:`_InspectModelAPIBase`; this file owns the vLLM ``LLM`` construction (with the
9plugin + worker_extension wiring needed for capture), eval-native generation via
10``llm.generate(...)``, and the TL-driven capture path via ``collective_rpc`` to the
11worker extension. The capture wire format matches the HF provider's, so the existing
12``InspectDriver`` consumes it unchanged.
14vLLM is GPU-only and imported lazily inside ``__init__`` / ``_generate_*`` so this
15module imports cleanly in environments without vLLM (matching the HF provider pattern).
16"""
17from __future__ import annotations
19import gc
20from typing import Any, Mapping
22import numpy as np
23import torch
24from inspect_ai.model import (
25 ChatCompletionChoice,
26 ChatMessageAssistant,
27 GenerateConfig,
28 Logprob,
29 Logprobs,
30 ModelOutput,
31 ModelUsage,
32 StopReason,
33 TopLogprob,
34 modelapi,
35)
37from . import hooks, wire
38from ._provider_base import (
39 _InspectModelAPIBase,
40 _parse_tool_calls,
41 _require_interveneable,
42 _require_served,
43 _warn_unsupported_config,
44)
46# Distinct from the HF ``tl_bridge`` provider and from inspect_ai's built-in ``vllm``.
47PROVIDER_NAME = "tl_bridge_vllm"
49# Forced ``LLM(...)`` kwargs the capture-hook design depends on. Multi-device and a
50# vLLM-owned tokenizer break the wire path; prefix caching would misalign capture rows
51# (see ``construct_instrumented_llm``, which supplies the caching/parallelism values).
52_LOCKED_VLLM_KWARGS = {
53 "tensor_parallel_size": 1,
54 "pipeline_parallel_size": 1,
55 "enable_prefix_caching": False,
56 "skip_tokenizer_init": True,
57 "disable_log_stats": True,
58}
61def _kinds_from_specs(specs: dict[str, Any]) -> frozenset[str]:
62 """Boundary kinds (resid_post/attn_out/...) served by the overlay's ``capture_specs``.
64 Non-block hooks (``embed.hook_out``, ``ln_final.hook_normalized``) don't resolve to
65 a kind and don't contribute — they're already in the InspectDriver's non-fireable set.
66 """
67 kinds = set()
68 for name in specs:
69 resolved = hooks.resolve(name)
70 if resolved is not None:
71 kinds.add(resolved[1])
72 return frozenset(kinds)
75@modelapi(name=PROVIDER_NAME)
76def transformer_lens_vllm_provider():
77 """Lazy registration hook — returns the provider class on first use."""
78 return TransformerLensVLLMModelAPI
81class TransformerLensVLLMModelAPI(_InspectModelAPIBase):
82 """vLLM-backed Inspect provider. See module docstring for scope per increment."""
84 # vLLM's sampler bypasses lm_head; _synthesize_logits populates only the gen position,
85 # so earlier positions are -inf and loss would be NaN. RemoteBridge.forward must reject
86 # return_type ∈ {loss, both} — read by source.py → TLBridgeProfile → InspectDriver.
87 provides_sequence_logits = False
89 def __init__(
90 self,
91 model_name: str,
92 base_url: str | None = None,
93 api_key: str | None = None,
94 config: GenerateConfig = GenerateConfig(),
95 **model_args: Any,
96 ) -> None:
97 super().__init__(model_name, base_url, api_key, [], config)
98 from transformers import AutoConfig, AutoTokenizer
100 try:
101 import vllm # noqa: F401 — import check; construction is in construct_instrumented_llm
102 except ImportError as exc:
103 raise ImportError(
104 "The tl_bridge_vllm provider requires vLLM (Linux + CUDA). Install with "
105 'pip install "transformer-lens[vllm]" or uv sync --extra vllm; '
106 "validated against vllm 0.20.x."
107 ) from exc
109 from transformer_lens.model_bridge.sources._hf_format import (
110 determine_architecture_from_hf_config,
111 )
112 from transformer_lens.utilities.hf_utils import get_hf_token
114 from ..vllm.internals import extract_hf_config
115 from ..vllm.overlays import DEFAULT_VLLM_OVERLAY, get_overlay
116 from ..vllm.source import _dtype_from_hf_config, construct_instrumented_llm
117 from ..vllm.worker_extension import dtype_name
119 # Caller-overridable LLM kwargs go through ``vllm_kwargs``; the locked set above
120 # may not be overridden (multi-device / vLLM-owned tokenizer break our wire path).
121 vllm_kwargs = model_args.pop("vllm_kwargs", {})
122 for key, locked in _LOCKED_VLLM_KWARGS.items():
123 if key in vllm_kwargs and vllm_kwargs[key] != locked:
124 raise ValueError(
125 f"tl_bridge_vllm forces {key}={locked}; caller passed "
126 f"{key}={vllm_kwargs[key]}."
127 )
128 gpu_memory_utilization = model_args.pop("gpu_memory_utilization", 0.5)
129 max_model_len = model_args.pop("max_model_len", None)
130 max_num_batched_tokens = int(model_args.pop("max_num_batched_tokens", 2048))
131 dtype = model_args.pop("dtype", None)
133 # vLLM is GPU-only in production; "device" stays caller-overridable mainly so
134 # mocked unit tests on CPU machines can build the prompt tensor without CUDA
135 # (the base's _messages_to_ids honors self._device; we drop to a list before
136 # handing the prompt to vLLM, so the device only governs the intermediate tensor).
137 self._device = model_args.pop("device", "cuda")
138 # skip_tokenizer_init=True ⇒ vLLM has no tokenizer; we own one for prompt rendering
139 # (via the base) and for decoding generated token ids back to strings.
140 hf_token = get_hf_token()
141 self._tokenizer = AutoTokenizer.from_pretrained(model_name, token=hf_token)
143 # Pre-LLM: resolve architecture WITHOUT loading weights, then prime the plugin
144 # so its monkey-patched Worker.load_model installs capture hooks pre-compile.
145 hf_config_preview = AutoConfig.from_pretrained(model_name, token=hf_token)
146 # Shared resolution (not architectures[0]) handles architectures=None configs
147 # via model_type. Unlike boot_vllm, an arch unknown to TL's registry must not
148 # abort here: get_overlay serves the decoder-only default for any name, so a
149 # resolver miss falls back to that same default and the model stays loadable.
150 try:
151 overlay = get_overlay(determine_architecture_from_hf_config(hf_config_preview))
152 except ValueError:
153 overlay = DEFAULT_VLLM_OVERLAY
154 resolved_dtype = dtype if dtype is not None else _dtype_from_hf_config(hf_config_preview)
155 capture_specs = overlay.capture_specs(hf_config_preview)
156 self._llm = construct_instrumented_llm(
157 model_name,
158 capture_specs=capture_specs,
159 max_num_batched_tokens=max_num_batched_tokens,
160 dtype=resolved_dtype,
161 enable_batching=False, # eager batched path is a later increment
162 llm_kwargs={
163 "gpu_memory_utilization": gpu_memory_utilization,
164 "max_model_len": max_model_len,
165 # Full-vocab logprobs so _generate_capture can synthesize logits at the
166 # generated position (vLLM caps logprobs to this value; default 20 is too
167 # small for mech interp).
168 "max_logprobs": int(hf_config_preview.vocab_size),
169 "dtype": dtype_name(resolved_dtype) if dtype is not None else "auto",
170 "skip_tokenizer_init": True,
171 "disable_log_stats": True,
172 **vllm_kwargs,
173 },
174 )
175 hf_config = extract_hf_config(self._llm)
177 # Capture-relevant constants used by _generate_capture.
178 self._d_vocab = int(hf_config.vocab_size)
179 self._max_logprobs = int(hf_config_preview.vocab_size)
180 self._max_num_batched_tokens = max_num_batched_tokens
182 # Boundary kinds served by the vLLM overlay (decoder-only: resid_post / attn_out /
183 # mlp_out). vLLM's fused execution doesn't expose block input, so resid_pre and
184 # the derived resid_mid are gated — the InspectDriver consults this via the profile.
185 self._kinds = _kinds_from_specs(capture_specs)
186 self._capability_note = (
187 "tl_bridge_vllm: vLLM's fused execution gates resid_pre (no block-input hook) "
188 "and the derived resid_mid. Use boot_inspect(provider='tl_bridge') for those."
189 )
190 self._eval_capture = self._parse_eval_capture(model_args)
192 def _generate_capture(self, input: Any, extra_args: Mapping[str, Any], config: GenerateConfig):
193 """TL-driven single-token capture: push interventions to the worker, run a
194 single-token generate (vLLM's prefill populates the capture buffers), read them
195 back via ``collective_rpc``, and return the wire-format ``metadata["activations"]``
196 + synthesized ``tl_logits`` the InspectDriver expects."""
197 from vllm import SamplingParams
198 from vllm.inputs import TokensPrompt
200 input_ids = extra_args.get("input_ids")
201 if input_ids is None: 201 ↛ 202line 201 didn't jump to line 202 because the condition on line 201 was never true
202 input_ids = self._messages_to_ids(input)[0].tolist()
203 n_tokens = len(input_ids)
204 if n_tokens > self._max_num_batched_tokens:
205 raise ValueError(
206 f"Prompt length {n_tokens} exceeds max_num_batched_tokens="
207 f"{self._max_num_batched_tokens}; raise via the model_args kwarg "
208 "or shorten the prompt."
209 )
211 # Validate capture kinds against the structural self-check (same protection the
212 # HF provider gives — driver-path AND eval/extra_args entry points).
213 capture_keys = list(extra_args.get("capture", []))
214 for key in capture_keys:
215 _, _, kind = key.partition(":")
216 _require_served(kind, self._kinds, self._capability_note, f"capture {key!r}")
218 # Translate wire keys ↔ TL hook names. The worker extension is keyed by hook name
219 # (e.g. "blocks.0.hook_out"); the wire format uses "<layer>:<kind>".
220 name_by_wire = {wk: hooks.name_from_wire_key(wk) for wk in capture_keys}
221 if any(name is None for name in name_by_wire.values()):
222 unknown = sorted(wk for wk, name in name_by_wire.items() if name is None)
223 raise ValueError(f"unrecognised wire keys: {unknown}")
224 capture_names = list(name_by_wire.values())
226 interventions: Mapping[str, Any] = extra_args.get("interventions", {})
227 intervention_specs: dict[str, Any] = {}
228 for wk, spec in interventions.items():
229 # extra_args is a documented surface: gated/capture-only kinds must fail
230 # here, not deep in the worker.
231 _, _, kind = wk.partition(":")
232 _require_interveneable(kind, self._kinds, self._capability_note, f"intervention {wk!r}")
233 if isinstance(spec, Mapping) and spec.get("pos") is not None:
234 raise ValueError(
235 f"intervention {wk!r}: per-position 'pos' is not supported on the "
236 "tl_bridge_vllm provider (its worker runs without position "
237 "interventions). Use boot_inspect(provider='tl_bridge') for "
238 "position-targeted patching."
239 )
240 name = hooks.name_from_wire_key(wk)
241 if name is None:
242 raise ValueError(f"intervention wire key {wk!r} is not a fireable hook.")
243 intervention_specs[name] = spec
245 want_logits = bool(extra_args.get("return_logits", True))
247 # Push intervention state (possibly empty — also resets stale interventions from
248 # a prior call), open the per-hook capture gates (so the prefill below writes,
249 # and any later forward — should this driver be reused — would self-copy until
250 # the next explicit reset). Then run a single-token prefill; vLLM's prefill
251 # populates the capture buffers we registered via plugin.configure.
252 self._llm.collective_rpc("tl_set_interventions", args=(intervention_specs,))
253 self._llm.collective_rpc("tl_reset_capture_flags")
254 outputs = self._llm.generate(
255 prompts=[TokensPrompt(prompt_token_ids=list(input_ids))],
256 sampling_params=SamplingParams(
257 max_tokens=1,
258 temperature=0.0,
259 logprobs=self._max_logprobs if want_logits else None,
260 ),
261 )
262 # collective_rpc returns one result per worker; single-rank ⇒ [0].
263 worker_captures = self._llm.collective_rpc(
264 "tl_read_captures", args=([n_tokens], capture_names)
265 )[0]
267 # Convert TL-name-keyed (n_tokens, width) tensors → wire-key-keyed numpy arrays,
268 # then encode in the same envelope wire.decode_activations consumes on the driver.
269 captured_wire: dict[str, np.ndarray] = {}
270 for wk, name in name_by_wire.items():
271 tensor = worker_captures.get(name)
272 if tensor is not None: 272 ↛ 270line 272 didn't jump to line 270 because the condition on line 272 was always true
273 captured_wire[wk] = tensor.detach().float().cpu().numpy()
274 metadata: dict[str, Any] = {"activations": wire.encode_activations(captured_wire)}
276 if want_logits:
277 logits = _synthesize_logits(outputs[0], n_tokens, self._d_vocab)
278 metadata["tl_logits"] = wire.encode_array(logits[0].cpu().numpy())
280 # The completion is the single generated token (matches the HF provider's shape).
281 next_id = int(outputs[0].outputs[0].token_ids[0])
282 return ModelOutput(
283 model=self.model_name,
284 choices=[
285 ChatCompletionChoice(
286 message=ChatMessageAssistant(content=str(self._tokenizer.decode([next_id]))),
287 stop_reason="stop",
288 )
289 ],
290 metadata=metadata,
291 )
293 def _generate_eval(self, input: Any, config: GenerateConfig, tools: Any):
294 """vLLM generation: chat input → ``llm.generate`` → completion + Logprobs + usage.
295 If ``model_args['capture']`` was set, opens per-hook capture gates before the eval
296 generate so prefill captures the prompt activations and decode steps self-copy
297 (first-write-wins on the worker; no separate forward — see plugin._gated_capture)."""
298 from vllm import SamplingParams
299 from vllm.inputs import TokensPrompt
301 _warn_unsupported_config(config, PROVIDER_NAME)
303 # Worker-side intervention buffers are persistent: any prior capture-path call that
304 # pushed specs (e.g. bridge.forward(intervene=...)) would still be applied here
305 # without a reset. The HF provider installs hooks per-call so it's leak-immune.
306 self._llm.collective_rpc("tl_set_interventions", args=({},))
307 ids = self._messages_to_ids(input, tools)[0].tolist()
308 prompt_len = len(ids)
309 if self._eval_capture:
310 if prompt_len > self._max_num_batched_tokens: 310 ↛ 311line 310 didn't jump to line 311 because the condition on line 310 was never true
311 raise ValueError(
312 f"Prompt length {prompt_len} exceeds max_num_batched_tokens="
313 f"{self._max_num_batched_tokens}; per-turn capture cannot snapshot it."
314 )
315 self._llm.collective_rpc("tl_reset_capture_flags")
316 max_new = int(config.max_tokens) if config.max_tokens else 16
317 temperature = float(config.temperature) if config.temperature is not None else 0.0
319 sp_kwargs: dict[str, Any] = {"max_tokens": max_new, "temperature": temperature}
320 if config.stop_seqs:
321 sp_kwargs["stop"] = list(config.stop_seqs)
322 if temperature > 0:
323 if config.top_p is not None: 323 ↛ 325line 323 didn't jump to line 325 because the condition on line 323 was always true
324 sp_kwargs["top_p"] = float(config.top_p)
325 if config.top_k is not None: 325 ↛ 326line 325 didn't jump to line 326 because the condition on line 325 was never true
326 sp_kwargs["top_k"] = int(config.top_k)
327 if config.seed is not None:
328 sp_kwargs["seed"] = int(config.seed)
329 if config.logprobs:
330 # vLLM returns this many top logprobs per generated token (incl. the chosen).
331 sp_kwargs["logprobs"] = int(config.top_logprobs) if config.top_logprobs else 1
333 outputs = self._llm.generate(
334 prompts=[TokensPrompt(prompt_token_ids=ids)],
335 sampling_params=SamplingParams(**sp_kwargs),
336 )
337 request_output = outputs[0]
338 output = request_output.outputs[0] # one prompt, one sample
339 new_ids = list(output.token_ids)
340 n_new = len(new_ids)
341 completion = str(self._tokenizer.decode(new_ids, skip_special_tokens=True))
343 logprobs = None
344 if config.logprobs and output.logprobs:
345 logprobs = Logprobs(
346 content=[
347 self._logprob_from_dict(int(tid), step, config.top_logprobs)
348 for tid, step in zip(new_ids, output.logprobs)
349 ]
350 )
352 tool_calls = _parse_tool_calls(completion) if len(tools) else None
353 finish = (output.finish_reason or "").lower()
354 stop_reason: StopReason
355 if tool_calls: 355 ↛ 356line 355 didn't jump to line 356 because the condition on line 355 was never true
356 stop_reason = "tool_calls"
357 elif finish == "length":
358 stop_reason = "max_tokens"
359 elif finish == "stop": 359 ↛ 362line 359 didn't jump to line 362 because the condition on line 359 was always true
360 stop_reason = "stop"
361 else:
362 stop_reason = "unknown"
364 # Per-turn capture lands in metadata. First-write-wins gating made prefill the
365 # only forward that wrote to the capture buffer, so we read it now (decode steps
366 # left rows 1..prompt_len-1 untouched and row 0 self-copied).
367 eval_metadata: dict[str, Any] = {}
368 if self._eval_capture:
369 capture_names = list(self._eval_capture.values())
370 worker_captures = self._llm.collective_rpc(
371 "tl_read_captures", args=([prompt_len], capture_names)
372 )[0]
373 captured_wire: dict[str, np.ndarray] = {}
374 for wk, name in self._eval_capture.items():
375 tensor = worker_captures.get(name)
376 if tensor is not None: 376 ↛ 374line 376 didn't jump to line 374 because the condition on line 376 was always true
377 captured_wire[wk] = tensor.detach().float().cpu().numpy()
378 eval_metadata = {"activations": wire.encode_activations(captured_wire)}
380 return ModelOutput(
381 model=self.model_name,
382 choices=[
383 ChatCompletionChoice(
384 message=ChatMessageAssistant(content=completion, tool_calls=tool_calls),
385 stop_reason=stop_reason,
386 logprobs=logprobs,
387 )
388 ],
389 usage=ModelUsage(
390 input_tokens=prompt_len, output_tokens=n_new, total_tokens=prompt_len + n_new
391 ),
392 metadata=eval_metadata or None,
393 )
395 def _logprob_from_dict(self, token_id: int, step_logprobs: Any, top_n: Any) -> Logprob:
396 """vLLM per-step ``{token_id: Logprob(logprob, rank, decoded_token)}`` →
397 :class:`inspect_ai.model.Logprob` with optional top-k alternatives."""
398 chosen = step_logprobs.get(token_id)
399 chosen_lp = float(chosen.logprob) if chosen is not None else float("-inf")
400 top: list[TopLogprob] = []
401 if top_n: 401 ↛ 406line 401 didn't jump to line 406 because the condition on line 401 was always true
402 ranked = sorted(step_logprobs.items(), key=lambda kv: -float(kv[1].logprob))
403 for tid, lp in ranked[: int(top_n)]:
404 token = lp.decoded_token or self._tokenizer.decode([int(tid)])
405 top.append(TopLogprob(token=str(token), logprob=float(lp.logprob), bytes=None))
406 return Logprob(
407 token=str(self._tokenizer.decode([int(token_id)])),
408 logprob=chosen_lp,
409 bytes=None,
410 top_logprobs=top,
411 )
413 def close(self) -> None:
414 """Best-effort vLLM teardown — vLLM 0.20.2 has no ``LLM.shutdown()``, so weights
415 + KV cache stay resident until process exit unless we destroy the distributed env."""
416 self._llm = None
417 try:
418 from vllm.distributed.parallel_state import (
419 destroy_distributed_environment,
420 destroy_model_parallel,
421 )
423 destroy_model_parallel()
424 destroy_distributed_environment()
425 except Exception:
426 pass
427 gc.collect()
428 if torch.cuda.is_available():
429 try:
430 torch.cuda.empty_cache()
431 except Exception:
432 pass
435def _synthesize_logits(request_output: Any, n_tokens: int, d_vocab: int) -> torch.Tensor:
436 """Build a ``(1, n_tokens, d_vocab)`` logits-like tensor from vLLM's sampler output —
437 log-probs (not raw logits), with earlier positions ``-inf`` (lm_head is bypassed so
438 only the generated position is populated). Matches the HF provider's ``tl_logits``
439 shape so the InspectDriver consumes both the same way."""
440 logits = torch.full((1, n_tokens, d_vocab), float("-inf"), dtype=torch.float32)
441 gen = request_output.outputs[0] if request_output.outputs else None
442 if gen is None: 442 ↛ 443line 442 didn't jump to line 443 because the condition on line 442 was never true
443 return logits
444 if gen.logprobs: 444 ↛ 447line 444 didn't jump to line 447 because the condition on line 444 was always true
445 for token_id, lp in gen.logprobs[0].items():
446 logits[0, -1, int(token_id)] = float(lp.logprob)
447 elif gen.token_ids:
448 logits[0, -1, int(gen.token_ids[0])] = 0.0
449 return logits
452__all__ = ["TransformerLensVLLMModelAPI"]