Coverage for transformer_lens/model_bridge/sources/inspect/vllm

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

88 supports_attention_mask = False

90 def __init__(

91 self,

92 model_name: str,

93 base_url: str | None = None,

94 api_key: str | None = None,

95 config: GenerateConfig = GenerateConfig(),

96 **model_args: Any,

97 ) -> None:

98 super().__init__(model_name, base_url, api_key, [], config)

99 from transformers import AutoConfig, AutoTokenizer

100

101 try:

102 import vllm # noqa: F401 — import check; construction is in construct_instrumented_llm

103 except ImportError as exc:

104 raise ImportError(

105 "The tl_bridge_vllm provider requires vLLM (Linux + CUDA). Install with "

106 'pip install "transformer-lens[vllm]" or uv sync --extra vllm; '

107 "validated against vllm 0.20.x."

108 ) from exc

109

110 from transformer_lens.model_bridge.sources._hf_format import (

111 determine_architecture_from_hf_config,

112 )

113 from transformer_lens.utilities.hf_utils import get_hf_token

114

115 from ..vllm.internals import extract_hf_config

116 from ..vllm.overlays import DEFAULT_VLLM_OVERLAY, get_overlay

117 from ..vllm.source import _dtype_from_hf_config, construct_instrumented_llm

118 from ..vllm.worker_extension import dtype_name

119

120 # Caller-overridable LLM kwargs go through ``vllm_kwargs``; the locked set above

121 # may not be overridden (multi-device / vLLM-owned tokenizer break our wire path).

122 vllm_kwargs = model_args.pop("vllm_kwargs", {})

123 for key, locked in _LOCKED_VLLM_KWARGS.items():

124 if key in vllm_kwargs and vllm_kwargs[key] != locked:

125 raise ValueError(

126 f"tl_bridge_vllm forces {key}={locked}; caller passed "

127 f"{key}={vllm_kwargs[key]}."

128 )

129 gpu_memory_utilization = model_args.pop("gpu_memory_utilization", 0.5)

130 max_model_len = model_args.pop("max_model_len", None)

131 max_num_batched_tokens = int(model_args.pop("max_num_batched_tokens", 2048))

132 dtype = model_args.pop("dtype", None)

133

134 # vLLM is GPU-only in production; "device" stays caller-overridable mainly so

135 # mocked unit tests on CPU machines can build the prompt tensor without CUDA

136 # (the base's _messages_to_ids honors self._device; we drop to a list before

137 # handing the prompt to vLLM, so the device only governs the intermediate tensor).

138 self._device = model_args.pop("device", "cuda")

139 # skip_tokenizer_init=True ⇒ vLLM has no tokenizer; we own one for prompt rendering

140 # (via the base) and for decoding generated token ids back to strings.

141 hf_token = get_hf_token()

142 self._tokenizer = AutoTokenizer.from_pretrained(model_name, token=hf_token)

143

144 # Pre-LLM: resolve architecture WITHOUT loading weights, then prime the plugin

145 # so its monkey-patched Worker.load_model installs capture hooks pre-compile.

146 hf_config_preview = AutoConfig.from_pretrained(model_name, token=hf_token)

147 # Shared resolution (not architectures[0]) handles architectures=None configs

148 # via model_type. Unlike boot_vllm, an arch unknown to TL's registry must not

149 # abort here: get_overlay serves the decoder-only default for any name, so a

150 # resolver miss falls back to that same default and the model stays loadable.

151 try:

152 overlay = get_overlay(determine_architecture_from_hf_config(hf_config_preview))

153 except ValueError:

154 overlay = DEFAULT_VLLM_OVERLAY

155 resolved_dtype = dtype if dtype is not None else _dtype_from_hf_config(hf_config_preview)

156 capture_specs = overlay.capture_specs(hf_config_preview)

157 self._llm = construct_instrumented_llm(

158 model_name,

159 capture_specs=capture_specs,

160 max_num_batched_tokens=max_num_batched_tokens,

161 dtype=resolved_dtype,

162 enable_batching=False, # eager batched path is a later increment

163 llm_kwargs={

164 "gpu_memory_utilization": gpu_memory_utilization,

165 "max_model_len": max_model_len,

166 # Full-vocab logprobs so _generate_capture can synthesize logits at the

167 # generated position (vLLM caps logprobs to this value; default 20 is too

168 # small for mech interp).

169 "max_logprobs": int(hf_config_preview.vocab_size),

170 "dtype": dtype_name(resolved_dtype) if dtype is not None else "auto",

171 "skip_tokenizer_init": True,

172 "disable_log_stats": True,

173 **vllm_kwargs,

174 },

175 )

176 hf_config = extract_hf_config(self._llm)

177

178 # Capture-relevant constants used by _generate_capture.

179 self._d_vocab = int(hf_config.vocab_size)

180 self._max_logprobs = int(hf_config_preview.vocab_size)

181 self._max_num_batched_tokens = max_num_batched_tokens

182

183 # Boundary kinds served by the vLLM overlay (decoder-only: resid_post / attn_out /

184 # mlp_out). vLLM's fused execution doesn't expose block input, so resid_pre and

185 # the derived resid_mid are gated — the InspectDriver consults this via the profile.

186 self._kinds = _kinds_from_specs(capture_specs)

187 self._capability_note = (

188 "tl_bridge_vllm: vLLM's fused execution gates resid_pre (no block-input hook) "

189 "and the derived resid_mid. Use boot_inspect(provider='tl_bridge') for those."

190 )

191 self._eval_capture = self._parse_eval_capture(model_args)

192

193 def _generate_capture(self, input: Any, extra_args: Mapping[str, Any], config: GenerateConfig):

194 """TL-driven single-token capture: push interventions to the worker, run a

195 single-token generate (vLLM's prefill populates the capture buffers), read them

196 back via ``collective_rpc``, and return the wire-format ``metadata["activations"]``

197 + synthesized ``tl_logits`` the InspectDriver expects."""

198 if extra_args.get("attention_mask") is not None:

199 raise NotImplementedError("tl_bridge_vllm does not support attention_mask.")

200 from vllm import SamplingParams

201 from vllm.inputs import TokensPrompt

202

203 input_ids = extra_args.get("input_ids")

204 if input_ids is None: 204 ↛ 205line 204 didn't jump to line 205 because the condition on line 204 was never true

205 input_ids = self._messages_to_ids(input)[0].tolist()

206 n_tokens = len(input_ids)

207 if n_tokens > self._max_num_batched_tokens:

208 raise ValueError(

209 f"Prompt length {n_tokens} exceeds max_num_batched_tokens="

210 f"{self._max_num_batched_tokens}; raise via the model_args kwarg "

211 "or shorten the prompt."

212 )

213

214 # Validate capture kinds against the structural self-check (same protection the

215 # HF provider gives — driver-path AND eval/extra_args entry points).

216 capture_keys = list(extra_args.get("capture", []))

217 for key in capture_keys:

218 _, _, kind = key.partition(":")

219 _require_served(kind, self._kinds, self._capability_note, f"capture {key!r}")

220

221 # Translate wire keys ↔ TL hook names. The worker extension is keyed by hook name

222 # (e.g. "blocks.0.hook_out"); the wire format uses "<layer>:<kind>".

223 name_by_wire = {wk: hooks.name_from_wire_key(wk) for wk in capture_keys}

224 if any(name is None for name in name_by_wire.values()):

225 unknown = sorted(wk for wk, name in name_by_wire.items() if name is None)

226 raise ValueError(f"unrecognised wire keys: {unknown}")

227 capture_names = list(name_by_wire.values())

228

229 interventions: Mapping[str, Any] = extra_args.get("interventions", {})

230 intervention_specs: dict[str, Any] = {}

231 for wk, spec in interventions.items():

232 # extra_args is a documented surface: gated/capture-only kinds must fail

233 # here, not deep in the worker.

234 _, _, kind = wk.partition(":")

235 _require_interveneable(kind, self._kinds, self._capability_note, f"intervention {wk!r}")

236 if isinstance(spec, Mapping) and spec.get("pos") is not None:

237 raise ValueError(

238 f"intervention {wk!r}: per-position 'pos' is not supported on the "

239 "tl_bridge_vllm provider (its worker runs without position "

240 "interventions). Use boot_inspect(provider='tl_bridge') for "

241 "position-targeted patching."

242 )

243 name = hooks.name_from_wire_key(wk)

244 if name is None:

245 raise ValueError(f"intervention wire key {wk!r} is not a fireable hook.")

246 intervention_specs[name] = spec

247

248 want_logits = bool(extra_args.get("return_logits", True))

249

250 # Push intervention state (possibly empty — also resets stale interventions from

251 # a prior call), open the per-hook capture gates (so the prefill below writes,

252 # and any later forward — should this driver be reused — would self-copy until

253 # the next explicit reset). Then run a single-token prefill; vLLM's prefill

254 # populates the capture buffers we registered via plugin.configure.

255 self._llm.collective_rpc("tl_set_interventions", args=(intervention_specs,))

256 self._llm.collective_rpc("tl_reset_capture_flags")

257 outputs = self._llm.generate(

258 prompts=[TokensPrompt(prompt_token_ids=list(input_ids))],

259 sampling_params=SamplingParams(

260 max_tokens=1,

261 temperature=0.0,

262 logprobs=self._max_logprobs if want_logits else None,

263 ),

264 )

265 # collective_rpc returns one result per worker; single-rank ⇒ [0].

266 worker_captures = self._llm.collective_rpc(

267 "tl_read_captures", args=([n_tokens], capture_names)

268 )[0]

269

270 # Convert TL-name-keyed (n_tokens, width) tensors → wire-key-keyed numpy arrays,

271 # then encode in the same envelope wire.decode_activations consumes on the driver.

272 captured_wire: dict[str, np.ndarray] = {}

273 for wk, name in name_by_wire.items():

274 tensor = worker_captures.get(name)

275 if tensor is not None: 275 ↛ 273line 275 didn't jump to line 273 because the condition on line 275 was always true

276 captured_wire[wk] = tensor.detach().float().cpu().numpy()

277 metadata: dict[str, Any] = {"activations": wire.encode_activations(captured_wire)}

278

279 if want_logits:

280 logits = _synthesize_logits(outputs[0], n_tokens, self._d_vocab)

281 metadata["tl_logits"] = wire.encode_array(logits[0].cpu().numpy())

282

283 # The completion is the single generated token (matches the HF provider's shape).

284 next_id = int(outputs[0].outputs[0].token_ids[0])

285 return ModelOutput(

286 model=self.model_name,

287 choices=[

288 ChatCompletionChoice(

289 message=ChatMessageAssistant(content=str(self._tokenizer.decode([next_id]))),

290 stop_reason="stop",

291 )

292 ],

293 metadata=metadata,

294 )

295

296 def _generate_eval(self, input: Any, config: GenerateConfig, tools: Any):

297 """vLLM generation: chat input → ``llm.generate`` → completion + Logprobs + usage.

298 If ``model_args['capture']`` was set, opens per-hook capture gates before the eval

299 generate so prefill captures the prompt activations and decode steps self-copy

300 (first-write-wins on the worker; no separate forward — see plugin._gated_capture)."""

301 from vllm import SamplingParams

302 from vllm.inputs import TokensPrompt

303

304 _warn_unsupported_config(config, PROVIDER_NAME)

305

306 # Worker-side intervention buffers are persistent: any prior capture-path call that

307 # pushed specs (e.g. bridge.forward(intervene=...)) would still be applied here

308 # without a reset. The HF provider installs hooks per-call so it's leak-immune.

309 self._llm.collective_rpc("tl_set_interventions", args=({},))

310 ids = self._messages_to_ids(input, tools)[0].tolist()

311 prompt_len = len(ids)

312 if self._eval_capture:

313 if prompt_len > self._max_num_batched_tokens: 313 ↛ 314line 313 didn't jump to line 314 because the condition on line 313 was never true

314 raise ValueError(

315 f"Prompt length {prompt_len} exceeds max_num_batched_tokens="

316 f"{self._max_num_batched_tokens}; per-turn capture cannot snapshot it."

317 )

318 self._llm.collective_rpc("tl_reset_capture_flags")

319 max_new = int(config.max_tokens) if config.max_tokens else 16

320 temperature = float(config.temperature) if config.temperature is not None else 0.0

321

322 sp_kwargs: dict[str, Any] = {"max_tokens": max_new, "temperature": temperature}

323 if config.stop_seqs:

324 sp_kwargs["stop"] = list(config.stop_seqs)

325 if temperature > 0:

326 if config.top_p is not None: 326 ↛ 328line 326 didn't jump to line 328 because the condition on line 326 was always true

327 sp_kwargs["top_p"] = float(config.top_p)

328 if config.top_k is not None: 328 ↛ 329line 328 didn't jump to line 329 because the condition on line 328 was never true

329 sp_kwargs["top_k"] = int(config.top_k)

330 if config.seed is not None:

331 sp_kwargs["seed"] = int(config.seed)

332 if config.logprobs:

333 # vLLM returns this many top logprobs per generated token (incl. the chosen).

334 sp_kwargs["logprobs"] = int(config.top_logprobs) if config.top_logprobs else 1

335

336 outputs = self._llm.generate(

337 prompts=[TokensPrompt(prompt_token_ids=ids)],

338 sampling_params=SamplingParams(**sp_kwargs),

339 )

340 request_output = outputs[0]

341 output = request_output.outputs[0] # one prompt, one sample

342 new_ids = list(output.token_ids)

343 n_new = len(new_ids)

344 completion = str(self._tokenizer.decode(new_ids, skip_special_tokens=True))

345

346 logprobs = None

347 if config.logprobs and output.logprobs:

348 logprobs = Logprobs(

349 content=[

350 self._logprob_from_dict(int(tid), step, config.top_logprobs)

351 for tid, step in zip(new_ids, output.logprobs)

352 ]

353 )

354

355 tool_calls = _parse_tool_calls(completion) if len(tools) else None

356 finish = (output.finish_reason or "").lower()

357 stop_reason: StopReason

358 if tool_calls: 358 ↛ 359line 358 didn't jump to line 359 because the condition on line 358 was never true

359 stop_reason = "tool_calls"

360 elif finish == "length":

361 stop_reason = "max_tokens"

362 elif finish == "stop": 362 ↛ 365line 362 didn't jump to line 365 because the condition on line 362 was always true

363 stop_reason = "stop"

364 else:

365 stop_reason = "unknown"

366

367 # Per-turn capture lands in metadata. First-write-wins gating made prefill the

368 # only forward that wrote to the capture buffer, so we read it now (decode steps

369 # left rows 1..prompt_len-1 untouched and row 0 self-copied).

370 eval_metadata: dict[str, Any] = {}

371 if self._eval_capture:

372 capture_names = list(self._eval_capture.values())

373 worker_captures = self._llm.collective_rpc(

374 "tl_read_captures", args=([prompt_len], capture_names)

375 )[0]

376 captured_wire: dict[str, np.ndarray] = {}

377 for wk, name in self._eval_capture.items():

378 tensor = worker_captures.get(name)

379 if tensor is not None: 379 ↛ 377line 379 didn't jump to line 377 because the condition on line 379 was always true

380 captured_wire[wk] = tensor.detach().float().cpu().numpy()

381 eval_metadata = {"activations": wire.encode_activations(captured_wire)}

382

383 return ModelOutput(

384 model=self.model_name,

385 choices=[

386 ChatCompletionChoice(

387 message=ChatMessageAssistant(content=completion, tool_calls=tool_calls),

388 stop_reason=stop_reason,

389 logprobs=logprobs,

390 )

391 ],

392 usage=ModelUsage(

393 input_tokens=prompt_len, output_tokens=n_new, total_tokens=prompt_len + n_new

394 ),

395 metadata=eval_metadata or None,

396 )

397

398 def _logprob_from_dict(self, token_id: int, step_logprobs: Any, top_n: Any) -> Logprob:

399 """vLLM per-step ``{token_id: Logprob(logprob, rank, decoded_token)}`` →

400 :class:`inspect_ai.model.Logprob` with optional top-k alternatives."""

401 chosen = step_logprobs.get(token_id)

402 chosen_lp = float(chosen.logprob) if chosen is not None else float("-inf")

403 top: list[TopLogprob] = []

404 if top_n: 404 ↛ 409line 404 didn't jump to line 409 because the condition on line 404 was always true

405 ranked = sorted(step_logprobs.items(), key=lambda kv: -float(kv[1].logprob))

406 for tid, lp in ranked[: int(top_n)]:

407 token = lp.decoded_token or self._tokenizer.decode([int(tid)])

408 top.append(TopLogprob(token=str(token), logprob=float(lp.logprob), bytes=None))

409 return Logprob(

410 token=str(self._tokenizer.decode([int(token_id)])),

411 logprob=chosen_lp,

412 bytes=None,

413 top_logprobs=top,

414 )

415

416 def close(self) -> None:

417 """Best-effort vLLM teardown — vLLM 0.20.2 has no ``LLM.shutdown()``, so weights

418 + KV cache stay resident until process exit unless we destroy the distributed env."""

419 self._llm = None

420 try:

421 from vllm.distributed.parallel_state import (

422 destroy_distributed_environment,

423 destroy_model_parallel,

424 )

425

426 destroy_model_parallel()

427 destroy_distributed_environment()

428 except Exception:

429 pass

430 gc.collect()

431 if torch.cuda.is_available():

432 try:

433 torch.cuda.empty_cache()

434 except Exception:

435 pass

436

437

438def _synthesize_logits(request_output: Any, n_tokens: int, d_vocab: int) -> torch.Tensor:

439 """Build a ``(1, n_tokens, d_vocab)`` logits-like tensor from vLLM's sampler output —

440 log-probs (not raw logits), with earlier positions ``-inf`` (lm_head is bypassed so

441 only the generated position is populated). Matches the HF provider's ``tl_logits``

442 shape so the InspectDriver consumes both the same way."""

443 logits = torch.full((1, n_tokens, d_vocab), float("-inf"), dtype=torch.float32)

444 gen = request_output.outputs[0] if request_output.outputs else None

445 if gen is None: 445 ↛ 446line 445 didn't jump to line 446 because the condition on line 445 was never true

446 return logits

447 if gen.logprobs: 447 ↛ 450line 447 didn't jump to line 450 because the condition on line 447 was always true

448 for token_id, lp in gen.logprobs[0].items():

449 logits[0, -1, int(token_id)] = float(lp.logprob)

450 elif gen.token_ids:

451 logits[0, -1, int(gen.token_ids[0])] = 0.0

452 return logits

453

454

455__all__ = ["TransformerLensVLLMModelAPI"]

Coverage for transformer_lens/model_bridge/sources/inspect/vllm_provider.py: 85%

199 statements