Coverage for transformer_lens/model_bridge/sources/vllm/driver.py: 66%
352 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-09-21 19:27 +0000
1"""vLLM Driver: forward dispatches via ``llm.generate``; captures via ``collective_rpc``."""
2from __future__ import annotations
4import gc
5import logging
6import threading
7import warnings
8from typing import Any, Mapping, Optional, Sequence, Union
10import torch
12from transformer_lens.model_bridge.driver_protocol import (
13 ForwardResult,
14 Intervention,
15 TensorLike,
16)
17from transformer_lens.model_bridge.sources._driver_base import DriverBase
19from .intervention_specs import validate_spec
20from .worker_extension import _TL_TENSOR_KEY, decode_tensor
22# vLLM's distributed teardown operates on process-wide globals, so close() may only run
23# it when no other VLLMDriver-owned engine is alive (notebook re-binding boots B before
24# dropping A; A's __del__ must not destroy the process groups B is using).
25_LIVE_DRIVERS = 0
26_LIVE_DRIVERS_LOCK = threading.Lock()
29class VLLMDriver(DriverBase):
30 """Driver wrapping a vLLM ``LLM``; captures via ``collective_rpc``."""
32 # vLLM owns the model in a worker — no torch module surface (parameters/state_dict/
33 # grads). Named-weight reads ARE served: get_param() returns CPU clones via the
34 # tl_get_param RPC (what logit reconstruction and direct-logit-attribution use).
35 _supported_features = frozenset({"weight_access"})
36 # Full-sequence logits reconstructed host-side (ln_final @ lm_head.weight.T); vLLM's
37 # sampler only hands back the final position, so the driver rebuilds the rest.
38 provides_sequence_logits = True
40 # Post-weight final-norm capture that lm_head consumes — reconstruction reads it.
41 _LN_FINAL = "ln_final.hook_normalized"
43 def __init__(
44 self,
45 llm: Any,
46 adapter: Any,
47 tokenizer: Any,
48 overlay: Any,
49 hf_config: Any,
50 max_num_batched_tokens: int,
51 enable_batching: bool = False,
52 enable_position_interventions: bool = False,
53 tensor_parallel_size: int = 1,
54 pipeline_parallel_size: int = 1,
55 ) -> None:
56 super().__init__(adapter.cfg, tokenizer)
57 self._llm = llm
58 self._max_num_batched_tokens = max_num_batched_tokens
59 self._enable_batching = enable_batching
60 # Rank layout: PP stages own disjoint layer (hence hook) subsets; the TP
61 # ranks within a stage hold replicas of that stage's post-all-reduce hook
62 # points. Reads merge across stages and take the first replica within one;
63 # the first capture-bearing forward cross-checks the whole layout
64 # (see _verify_rank_layout) and then trusts it.
65 self._tp_size = tensor_parallel_size
66 self._layout_verified = tensor_parallel_size == 1 and pipeline_parallel_size == 1
67 # Position-scoped 'pos' interventions need (max_n, width) affine buffers,
68 # allocated at boot only when this is set (see plugin.patched_load_model).
69 self._enable_position_interventions = enable_position_interventions
70 # Logprobs per forward = real vocab (boot's max_logprobs). d_vocab can be
71 # padded larger, which vLLM would reject; the logits tensor stays d_vocab.
72 self._n_logprobs = int(getattr(hf_config, "vocab_size", self.bridge_config.d_vocab))
73 # Unembedding cache: (weight_fp32, bias_fp32|None) once probe_logit_reconstruction
74 # runs — a per-forward re-fetch clones a d_vocab×d_model tensor to CPU every call.
75 self._unembed: tuple[torch.Tensor, Any] | None = None
76 self._unembed_probed = False
77 # fp32 reciprocal of the guarded final-norm weight; None after a failed probe.
78 self._lnf_inv_denom: torch.Tensor | None = None
79 self._lnf_probed = False
81 capture_specs = overlay.capture_specs(hf_config)
82 self._hook_widths = {name: width for name, (_path, width) in capture_specs.items()}
83 self._capture_paths = {name: path for name, (path, _width) in capture_specs.items()}
84 self.supported_hook_points = frozenset(capture_specs.keys())
86 global _LIVE_DRIVERS
87 with _LIVE_DRIVERS_LOCK:
88 _LIVE_DRIVERS += 1
90 n_layers = getattr(hf_config, "num_hidden_layers", 0)
91 if not isinstance(n_layers, int) or n_layers <= 0:
92 # A raw "{i}" template would land unexpanded in non_fireable_hook_points —
93 # a broken config should fail at boot, not surface as a garbled hook name.
94 raise ValueError(
95 f"VLLMDriver: hf_config.num_hidden_layers={n_layers!r} — expected a "
96 "positive int; the config is missing or malformed."
97 )
98 nonfiring: list[str] = []
99 for tmpl in overlay.nonfiring_hooks():
100 if "{i}" in tmpl:
101 nonfiring.extend(tmpl.replace("{i}", str(i)) for i in range(n_layers))
102 else:
103 nonfiring.append(tmpl)
104 self.non_fireable_hook_points = frozenset(nonfiring)
106 def forward(
107 self,
108 # Wider than the protocol's TensorLike: the batched path documents plain
109 # (ragged) list[int] / list[list[int]] prompts, which have no shape/dtype.
110 input_ids: Optional[Union[TensorLike, Sequence[Any]]] = None,
111 *,
112 capture: tuple[str, ...] = (),
113 intervene: Mapping[str, Intervention] | None = None,
114 max_new_tokens: int = 1,
115 return_logits: bool = True,
116 **kwargs: Any,
117 ) -> ForwardResult:
118 if input_ids is None: 118 ↛ 119line 118 didn't jump to line 119 because the condition on line 118 was never true
119 raise ValueError("VLLMDriver requires input_ids")
120 if int(max_new_tokens) != 1:
121 raise NotImplementedError(
122 "VLLMDriver supports max_new_tokens=1 only — decode-step writes "
123 "overwrite the prefill buffer; multi-step capture is multi-buffer work."
124 )
125 intervene_specs = self._validate_interventions(intervene or {})
126 # Pad tokens fed to vLLM are real content to it (no mask concept) — honor a
127 # caller-supplied mask by trimming rows to their true lengths, never swallow it.
128 attention_mask = kwargs.pop("attention_mask", None)
130 # capture is authoritative — the bridge sends exactly the hooked names, so ()
131 # means "capture nothing" and a plain forward(tokens) skips the GPU→CPU copy
132 # entirely. (The worker's None-means-all convention is never triggered from here;
133 # an empty tuple used to collapse to None and silently copy every buffer.)
134 names = list(capture)
136 if self._enable_batching: 136 ↛ 137line 136 didn't jump to line 137 because the condition on line 136 was never true
137 return self._forward_batched(
138 input_ids, intervene_specs, return_logits, names, attention_mask
139 )
141 ids_list = self._normalize_input_ids(input_ids)
142 if attention_mask is not None: 142 ↛ 143line 142 didn't jump to line 143 because the condition on line 142 was never true
143 n_real = int(torch.as_tensor(attention_mask).sum())
144 ids_list = ids_list[:n_real]
145 if len(ids_list) > self._max_num_batched_tokens:
146 # Worker buffers silently clamp on overflow — fail loud here instead.
147 raise ValueError(
148 f"Prompt length {len(ids_list)} exceeds max_num_batched_tokens="
149 f"{self._max_num_batched_tokens}; raise the boot_vllm kwarg or "
150 "shorten the prompt."
151 )
152 # 'pos'-scoped edits target affine-buffer rows, but the compiled hook only reads
153 # rows [0, len(ids_list)); a pos past the prompt length would be a silent no-op.
154 self._reject_pos_beyond_seq(intervene_specs, len(ids_list))
156 from vllm import SamplingParams
157 from vllm.inputs import TokensPrompt
159 # Push intervention state (possibly empty) before generate — this also
160 # resets stale interventions from prior forwards.
161 self._llm.collective_rpc("tl_set_interventions", args=(intervene_specs,))
162 # Open per-hook capture gates; first-write-wins means a fresh prefill writes and
163 # subsequent forwards self-copy. Without this, repeated bridge.forward calls would
164 # see the gate closed from the prior call and read stale buffers.
165 self._llm.collective_rpc("tl_reset_capture_flags")
166 outputs = self._llm.generate(
167 prompts=[TokensPrompt(prompt_token_ids=ids_list)],
168 sampling_params=SamplingParams(
169 max_tokens=int(max_new_tokens),
170 temperature=0.0,
171 logprobs=self._sampler_logprobs(return_logits),
172 ),
173 )
175 n_tokens = len(ids_list)
176 read_names = self._read_names(names, return_logits)
177 # collective_rpc returns one result per worker; each returns only the hooks
178 # it owns (PP shards) with TP replicas within a stage — merge across ranks.
179 # Nothing to read (no captures, logits off) → skip the crossing altogether.
180 if read_names:
181 per_rank_raw = self._llm.collective_rpc(
182 "tl_read_captures", args=([n_tokens], read_names)
183 )
184 if not self._layout_verified:
185 self._verify_rank_layout([self._rpc_captures(caps) for caps in per_rank_raw])
186 self._layout_verified = True
187 # Merge on the wire dicts, decode only the survivors — decoding every TP
188 # replica just to discard all but the first is wasted CPU per forward.
189 worker_captures = self._rpc_captures(
190 self._merge_rank_captures(per_rank_raw, read_names)
191 )
192 else:
193 worker_captures = {}
195 logits: torch.Tensor | None = None
196 if return_logits:
197 recon = self._reconstruct_logits(worker_captures.get(self._LN_FINAL))
198 # Fall back to final-position log-probs if the unembedding isn't fetchable.
199 logits = (
200 recon.unsqueeze(0)
201 if recon is not None
202 else self._synthesize_logits(outputs[0], n_tokens, self.bridge_config.d_vocab)
203 )
205 captured = self._expose_captured(
206 {name: t.unsqueeze(0) for name, t in worker_captures.items()}, names
207 )
208 return ForwardResult(logits=logits, captured=captured, raw_output=outputs[0])
210 def _sampler_logprobs(self, return_logits: bool) -> int | None:
211 # Sampler logprobs are the reconstruction fallback only; skipping them avoids
212 # marshaling a d_vocab-entry Logprob dict host-side per request.
213 return self._n_logprobs if return_logits and self._unembed is None else None
215 def _read_names(self, names: list[str], return_logits: bool) -> list[str]:
216 # Reconstruction needs ln_final even when uncaptured; skip once probed-unavailable.
217 recon_possible = self._unembed is not None or not self._unembed_probed
218 if return_logits and recon_possible and self._LN_FINAL not in names:
219 return names + [self._LN_FINAL]
220 return names
222 def _expose_captured(
223 self, captured: Mapping[str, torch.Tensor], names: list[str]
224 ) -> dict[str, torch.Tensor]:
225 """Filter to the caller's requested hooks (dropping any forced ln_final) and
226 convert the exposed ln_final to the pre-weight convention its name promises —
227 reconstruction consumed the raw post-weight value already."""
228 out = {name: t for name, t in captured.items() if name in names}
229 if self._LN_FINAL in out:
230 out[self._LN_FINAL] = self._unfold_ln_final(out[self._LN_FINAL])
231 return out
233 def _forward_batched(
234 self,
235 input_ids: Union[TensorLike, Sequence[Any]],
236 intervene_specs: dict,
237 return_logits: bool,
238 names: list[str],
239 attention_mask: Any = None,
240 ) -> ForwardResult:
241 """Eager batched path: per-request capture, right-padded to (B, S, W).
243 No per-prompt length gate — chunked prefill accumulates long prompts
244 across forwards. Interventions are global across the batch. ``names`` is
245 authoritative: exactly the hooks to return (empty = none).
246 """
247 from vllm import SamplingParams
248 from vllm.inputs import TokensPrompt
250 prompts_ids = self._normalize_input_ids_batched(input_ids)
251 if attention_mask is not None:
252 # Right-padded tensor batches carry pad ids vLLM would treat as content;
253 # trim each row to its masked length so per-row final positions are real.
254 mask = torch.as_tensor(attention_mask)
255 if mask.dim() == 1:
256 mask = mask.unsqueeze(0)
257 if mask.shape[0] != len(prompts_ids):
258 raise ValueError(
259 f"attention_mask batch dim {mask.shape[0]} != number of prompts "
260 f"{len(prompts_ids)}."
261 )
262 prompts_ids = [ids[: int(row.sum())] for ids, row in zip(prompts_ids, mask)]
263 prompt_lens = [len(ids) for ids in prompts_ids]
265 # Reset accumulators so prior-forward chunks don't leak into the cat.
266 self._llm.collective_rpc("tl_reset_accumulators")
267 self._llm.collective_rpc("tl_set_batched_interventions", args=(intervene_specs,))
269 d_vocab = self.bridge_config.d_vocab
270 outputs = self._llm.generate(
271 prompts=[TokensPrompt(prompt_token_ids=ids) for ids in prompts_ids],
272 sampling_params=SamplingParams(
273 max_tokens=1,
274 temperature=0.0,
275 logprobs=self._sampler_logprobs(return_logits),
276 ),
277 )
279 read_names = self._read_names(names, return_logits)
280 # Keyed by req_id (no guaranteed order) — _assemble_padded joins to slot
281 # k via outputs[k].request_id, not by position. Empty read → skip both the
282 # crossing AND the join (_assemble_padded requires one worker key per request,
283 # so it can't run on an empty dict — mirror the single path's empty captured).
284 if read_names:
285 worker_captures = self._llm.collective_rpc(
286 "tl_read_batched_captures", args=(read_names,)
287 )[0]
288 # Batched runs single-rank today, but coerce anyway — multiproc RPC
289 # serializes tensors to lists (see _rpc_tensor).
290 worker_captures = {
291 req_id: self._rpc_captures(caps) for req_id, caps in worker_captures.items()
292 }
293 captured = self._assemble_padded(outputs, worker_captures, prompt_lens)
294 else:
295 captured = {}
297 logits: torch.Tensor | None = None
298 if return_logits:
299 recon = self._reconstruct_logits(
300 captured.get(self._LN_FINAL)
301 ) # (batch, max_seq, d_vocab)
302 if recon is not None:
303 # Pad rows reconstruct from zero-filled ln_final into finite garbage
304 # (0 @ W = plausible uniform logits) — mask them to the -inf convention
305 # the fallback and per-row consumers rely on.
306 for k, n in enumerate(prompt_lens):
307 recon[k, n:] = float("-inf")
308 logits = (
309 recon
310 if recon is not None
311 else self._synthesize_logits_batched(outputs, prompt_lens, d_vocab)
312 )
313 captured = self._expose_captured(captured, names)
315 return ForwardResult(logits=logits, captured=captured, raw_output=outputs)
317 @staticmethod
318 def _assemble_padded(
319 outputs: list,
320 worker_captures: Mapping[str, Mapping[str, torch.Tensor]],
321 prompt_lens: list[int],
322 ) -> dict[str, torch.Tensor]:
323 """Stack per-request captures into right-padded ``(batch, max_seq, width)``.
325 Pad is a cache-assembly artifact only: vLLM computes each request
326 independently, so real-token activations don't depend on the padding.
327 """
328 batch = len(outputs)
329 max_seq = max(prompt_lens) if prompt_lens else 0
330 # Worker keys are engine-internal req_ids ("10-83c3532c"); RequestOutput
331 # carries only the public id ("10"). Join exact-or-prefix; the "-" keeps
332 # "1" from matching "10-...".
333 worker_keys = list(worker_captures.keys())
335 def _captures_for(public_rid: str) -> Mapping[str, torch.Tensor]:
336 if public_rid in worker_captures: 336 ↛ 337line 336 didn't jump to line 337 because the condition on line 336 was never true
337 return worker_captures[public_rid]
338 matches = [k for k in worker_keys if k.startswith(f"{public_rid}-")]
339 # Raise, never silently zero-fill the row: a missing or ambiguous join
340 # is indistinguishable from a genuine zero activation, which is silent
341 # data loss on the collection path.
342 if len(matches) != 1: 342 ↛ 348line 342 didn't jump to line 348 because the condition on line 342 was always true
343 raise RuntimeError(
344 f"Cannot join request {public_rid!r} to worker captures: found "
345 f"{len(matches)} key(s) in {worker_keys}. Expected exactly one "
346 f"(exact or '{public_rid}-<hash>')."
347 )
348 return worker_captures[matches[0]]
350 per_slot = [_captures_for(o.request_id) for o in outputs]
352 hook_names: set[str] = set()
353 for caps in per_slot:
354 hook_names |= set(caps.keys())
356 assembled: dict[str, torch.Tensor] = {}
357 for name in hook_names:
358 sample = next(caps[name] for caps in per_slot if name in caps)
359 buf = torch.zeros(batch, max_seq, sample.shape[-1], dtype=sample.dtype)
360 for k, caps in enumerate(per_slot):
361 t = caps.get(name)
362 if t is not None:
363 buf[k, : t.shape[0]] = t
364 assembled[name] = buf
365 return assembled
367 @staticmethod
368 def _synthesize_logits_batched(
369 outputs: list, prompt_lens: list[int], d_vocab: int
370 ) -> torch.Tensor:
371 """Build ``(batch, max_seq, d_vocab)`` logits; next-token dist at each row's
372 ``prompt_lens[k] - 1``, never ``-1`` (a pad position for shorter prompts)."""
373 batch = len(outputs)
374 max_seq = max(prompt_lens) if prompt_lens else 0
375 logits = torch.full((batch, max_seq, d_vocab), float("-inf"), dtype=torch.float16)
376 for k, request_output in enumerate(outputs):
377 gen = request_output.outputs[0] if request_output.outputs else None
378 if gen is None:
379 continue
380 pos = prompt_lens[k] - 1
381 if gen.logprobs:
382 for token_id, lp_obj in gen.logprobs[0].items():
383 logits[k, pos, int(token_id)] = float(lp_obj.logprob)
384 elif gen.token_ids:
385 logits[k, pos, int(gen.token_ids[0])] = 0.0
386 return logits
388 def get_param(self, dotted_name: str) -> torch.Tensor | None:
389 """Fetch a named model tensor (e.g. ``model.norm.weight``) for conversions
390 the bridge can't otherwise do (ln_final post→pre-weight; see the overlay).
391 Gathered across ranks: replicated params return one copy, vocab-sharded
392 weights concatenate, stage-local params (PP) come from whichever rank owns
393 them. None if closed or the path resolves nowhere."""
394 return self._gather_param(dotted_name)
396 @staticmethod
397 def _rpc_tensor(value: Any) -> torch.Tensor:
398 """Decode a non-None collective_rpc payload back to a tensor.
400 Worker methods return the explicit wire format (see
401 ``worker_extension.encode_tensor``) because vLLM's multiproc RPC can't
402 round-trip raw tensors. Raw tensors still pass through for mocks and any
403 legacy in-process payloads; anything else is best-effort coerced."""
404 if isinstance(value, torch.Tensor):
405 return value
406 if isinstance(value, Mapping) and value.get(_TL_TENSOR_KEY): 406 ↛ 407line 406 didn't jump to line 407 because the condition on line 406 was never true
407 return decode_tensor(dict(value))
408 return torch.as_tensor(value)
410 @classmethod
411 def _rpc_captures(cls, captures: Mapping[str, Any]) -> dict[str, torch.Tensor]:
412 return {name: cls._rpc_tensor(value) for name, value in captures.items()}
414 def _gather_param(self, dotted_name: str, dim: int = 0) -> torch.Tensor | None:
415 """Fetch a param across all TP ranks: replicated → rank 0; sharded → concat.
417 Vocab-parallel weights (``lm_head.weight``, ``embed_tokens.weight``) hold
418 contiguous per-rank slices in rank order, which is also collective_rpc's
419 result order — concatenation along ``dim`` reassembles the full tensor
420 (vocab padding lands in the tail rows, sliced off by reconstruction).
421 """
422 if self._llm is None:
423 return None
424 shards = [
425 self._rpc_tensor(s)
426 for s in self._llm.collective_rpc("tl_get_param", args=(dotted_name,))
427 if s is not None
428 ]
429 if not shards:
430 return None
431 if len(shards) == 1:
432 return shards[0]
433 if all(s.shape == shards[0].shape and torch.equal(s, shards[0]) for s in shards[1:]):
434 return shards[0] # replicated (norm weights, biases on some archs)
435 return torch.cat(shards, dim=dim)
437 def _merge_rank_captures(self, per_rank_captures: list, requested: list[str]) -> dict[str, Any]:
438 """Merge per-rank capture dicts into one (values stay opaque — callers decode):
439 PP stages own disjoint hook subsets; TP replicas within a stage agree (verified
440 once by _verify_rank_layout), so the first copy wins. A requested, supported
441 hook served by NO rank fails loud — zero-filling would be silent data loss
442 (same policy as the batched join)."""
443 merged: dict[str, Any] = {}
444 for captures in per_rank_captures:
445 for name, tensor in captures.items():
446 if name not in merged:
447 merged[name] = tensor
448 missing = [
449 name for name in requested if name not in merged and name in self.supported_hook_points
450 ]
451 if missing:
452 raise RuntimeError(
453 f"Capture(s) {missing} returned by no rank — a hook fired nowhere "
454 "despite passing the boot-time coverage check. Report this."
455 )
456 return merged
458 def _verify_rank_layout(self, per_rank_captures: list) -> None:
459 """One-time cross-rank check of the capture layout under TP/PP.
461 Invariant: each hook lives on exactly one PP stage and is replicated across
462 that stage's ``tp_size`` ranks (every overlay hook point is post-all-reduce).
463 A wrong copy count means installation drift; divergent replicas mean a hook
464 point moved pre-all-reduce — either way rank-merged reads would be silently
465 wrong, so fail loud on the first capture-bearing forward. Ranks serving the
466 same hook set are TP replicas of one stage and must report identical fire
467 counters; counters are NOT comparable across stages — per-rank installed
468 counts overcount modules that exist but never run (tied-embedding aliases,
469 archs that instantiate norm on every rank).
470 """
471 all_names: set[str] = set()
472 for captures in per_rank_captures:
473 all_names |= set(captures.keys())
474 for name in sorted(all_names):
475 copies = [
476 (rank, caps[name]) for rank, caps in enumerate(per_rank_captures) if name in caps
477 ]
478 if len(copies) != self._tp_size:
479 raise RuntimeError(
480 f"Rank-layout check failed for {name!r}: found on {len(copies)} rank(s) "
481 f"but expected exactly tp_size={self._tp_size} (one PP stage × its TP "
482 "replicas). Per-rank installation drifted — report this."
483 )
484 rank0, t0 = copies[0]
485 for rank, t_r in copies[1:]:
486 if t0.shape != t_r.shape or not torch.allclose(
487 t0.float(), t_r.float(), atol=1e-5, rtol=1e-5
488 ):
489 diff = (
490 (t0.float() - t_r.float()).abs().max().item()
491 if t0.shape == t_r.shape
492 else float("inf")
493 )
494 raise RuntimeError(
495 f"TP replication check failed for {name!r}: rank {rank0} vs rank "
496 f"{rank} max abs diff {diff:.3e}. This hook point is no longer "
497 "replicated across ranks (vLLM may have moved it pre-all-reduce) — "
498 "merged capture reads would be silently wrong. Report this."
499 )
500 counters = [int(c) for c in self._llm.collective_rpc("tl_read_counter")]
501 stages: dict = {}
502 for rank, captures in enumerate(per_rank_captures):
503 stages.setdefault(frozenset(captures.keys()), []).append(rank)
504 for ranks in stages.values():
505 group = [counters[r] for r in ranks]
506 if len(set(group)) > 1:
507 raise RuntimeError(
508 f"Fire-counter mismatch within a TP replica group: ranks {ranks} "
509 f"serve the same hooks but fired {group} times. Replicated forwards "
510 "diverged — merged capture reads would be silently wrong. Report this."
511 )
513 def probe_logit_reconstruction(self) -> bool:
514 """One-time unembedding fetch + cache; downgrades ``provides_sequence_logits``
515 honestly when no unembedding is reachable (the fallback path is final-position
516 log-probs, which cannot back a loss). Idempotent — later calls return the
517 cached availability."""
518 if self._unembed_probed:
519 return self._unembed is not None
520 self._unembed_probed = True
521 # Gathered reads: under TP the unembedding is vocab-sharded per rank.
522 weight = self._gather_param("lm_head.weight")
523 if weight is None: # tied embeddings expose no separate lm_head
524 weight = self._gather_param("model.embed_tokens.weight")
525 if weight is None:
526 self.provides_sequence_logits = False
527 return False
528 bias = self._gather_param("lm_head.bias")
529 d_vocab = int(self.bridge_config.d_vocab)
530 # Slice vLLM's vocab-pad rows at cache time; fp32 residency (~2× checkpoint
531 # dtype on CPU) trades memory for skipping a full-matrix upcast per forward.
532 self._unembed = (
533 weight.to(torch.float32)[:d_vocab],
534 bias.to(torch.float32)[:d_vocab] if bias is not None else None,
535 )
536 self.provides_sequence_logits = True
537 return True
539 def _reconstruct_logits(self, ln_final: Any) -> torch.Tensor | None:
540 """Rebuild real logits from the captured post-weight ln_final:
541 ``ln_final @ lm_head.weight.T`` (+ bias, + Gemma-family tanh soft-cap).
543 vLLM's ``ln_final.hook_normalized`` is the POST-weight RMSNorm value lm_head
544 consumes (verified empirically: it equals HF's pre-weight value times the norm
545 weight), so no un-fold is needed here — the raw worker value feeds this directly.
546 Accepts any ``(..., d_model)`` tensor and returns ``(..., d_vocab)`` on CPU.
547 ``None`` if ln_final wasn't captured or no unembedding weight is fetchable —
548 the caller then falls back to the sampler's log-probs.
549 """
550 if ln_final is None or not self.probe_logit_reconstruction():
551 return None
552 assert self._unembed is not None # probe returned True
553 weight32, bias32 = self._unembed
554 lf = ln_final.to(device=weight32.device, dtype=torch.float32)
555 logits = lf @ weight32.T
556 if bias32 is not None: 556 ↛ 557line 556 didn't jump to line 557 because the condition on line 556 was never true
557 logits = logits + bias32.to(device=logits.device)
558 d_vocab = int(self.bridge_config.d_vocab)
559 if logits.shape[-1] > d_vocab:
560 # vLLM pads vocab embeddings to a multiple of 64; its own sampler slices to
561 # org_vocab_size before sampling — mirror that, or pad columns (zero-filled
562 # at load) become phantom argmax candidates and bias softmax denominators.
563 logits = logits[..., :d_vocab]
564 cap = getattr(self.bridge_config, "output_logits_soft_cap", None)
565 if cap is not None and cap > 0: # Gemma-family cap; -1.0 is the "disabled" sentinel 565 ↛ 566line 565 didn't jump to line 566 because the condition on line 565 was never true
566 logits = float(cap) * torch.tanh(logits / float(cap))
567 if logits.shape[-1] < d_vocab: # pad the padded-vocab tail (never predicted) 567 ↛ 568line 567 didn't jump to line 568 because the condition on line 567 was never true
568 pad = logits.new_full((*logits.shape[:-1], d_vocab - logits.shape[-1]), float("-inf"))
569 logits = torch.cat([logits, pad], dim=-1)
570 return logits.cpu()
572 def _unfold_ln_final(self, t: torch.Tensor) -> torch.Tensor:
573 """Convert vLLM's post-weight RMSNorm capture to the pre-weight value the
574 canonical hook name promises (÷ weight; Gemma folds ``1 + weight``). Warns
575 once and serves the raw value when the norm weight is unreachable — loud
576 beats silent cross-backend mismatch."""
577 if not self._lnf_probed:
578 self._lnf_probed = True
579 # The overlay's capture spec owns the module path; derive the weight from
580 # it. Gathered read: under PP the final norm lives only on the last stage,
581 # so a rank-0 get_param would miss it.
582 path = self._capture_paths.get(self._LN_FINAL, "model.norm")
583 weight = self._gather_param(f"{path}.weight")
584 if weight is None:
585 warnings.warn(
586 "ln_final.hook_normalized: norm weight unreachable — the captured "
587 "value stays POST-weight and will not match boot_transformers.",
588 UserWarning,
589 stacklevel=3,
590 )
591 else:
592 w = weight.detach().to(torch.float32)
593 denom = (1.0 + w) if "gemma" in self.architecture.lower() else w
594 # Near-zero weight entries would blow up the division; identity beats inf.
595 denom = torch.where(denom.abs() < 1e-6, torch.ones_like(denom), denom)
596 self._lnf_inv_denom = denom.reciprocal()
597 if self._lnf_inv_denom is None:
598 return t
599 inv = self._lnf_inv_denom.to(device=t.device)
600 return (t.to(torch.float32) * inv).to(t.dtype)
602 def close(self) -> None:
603 if self._llm is None: # already closed — keep the refcount single-shot
604 return
605 # Detach hooks before dropping the LLM so they don't stay registered on
606 # worker modules for the life of the process (long-running notebooks).
607 log = logging.getLogger("transformer_lens.vllm")
608 try:
609 self._llm.collective_rpc("tl_remove_hooks")
610 except Exception as e:
611 # Best-effort: engine may already be torn down or the RPC surface
612 # gone. Log so hook-leak debugging has a thread to pull.
613 log.debug("tl_remove_hooks failed during close(): %s", e)
614 self._llm = None
615 self._unembed = None
616 self._lnf_inv_denom = None
618 global _LIVE_DRIVERS
619 with _LIVE_DRIVERS_LOCK:
620 _LIVE_DRIVERS -= 1
621 last_driver = _LIVE_DRIVERS <= 0
622 # vLLM 0.20.2 has no LLM.shutdown(); the distributed teardown below hits
623 # process-global state, so it runs only for the last live driver (see
624 # _LIVE_DRIVERS). Both calls are best-effort no-ops otherwise.
625 if last_driver:
626 try:
627 from vllm.distributed.parallel_state import (
628 destroy_distributed_environment,
629 destroy_model_parallel,
630 )
632 destroy_model_parallel()
633 destroy_distributed_environment()
634 except Exception as e:
635 log.debug("vLLM distributed teardown failed during close(): %s", e)
636 # Free the caching allocator's blocks here so the caller doesn't have to.
637 gc.collect()
638 try:
639 if torch.cuda.is_available(): 639 ↛ 640line 639 didn't jump to line 640 because the condition on line 639 was never true
640 torch.cuda.empty_cache()
641 except Exception as e:
642 log.debug("torch.cuda.empty_cache failed during close(): %s", e)
644 @staticmethod
645 def _synthesize_logits(request_output: Any, n_tokens: int, d_vocab: int) -> torch.Tensor:
646 """Build a (1, n_tokens, d_vocab) logits-like tensor from vLLM's sampler output.
648 Values are **log-probs**, not raw logits (vLLM returns log_softmax): fine
649 for argmax/next-token, wrong for absolute scale (temperature, logit-lens).
650 lm_head is bypassed so only position -1 is filled (next-token); earlier
651 positions stay -inf — populating them needs prompt_logprobs per call.
652 """
653 logits = torch.full((1, n_tokens, d_vocab), float("-inf"), dtype=torch.float16)
654 gen = request_output.outputs[0] if request_output.outputs else None
655 if gen is None:
656 return logits
657 # Prefer real logprobs; fall back to the generated token id (one-hot-ish)
658 # if logprobs weren't requested (e.g. return_logits=False elsewhere).
659 if gen.logprobs:
660 for token_id, lp_obj in gen.logprobs[0].items():
661 logits[0, -1, int(token_id)] = float(lp_obj.logprob)
662 elif gen.token_ids:
663 logits[0, -1, int(gen.token_ids[0])] = 0.0
664 return logits
666 def _validate_interventions(self, intervene: Mapping[str, Any]) -> dict:
667 """Shared spec validation plus driver-side gating (hook membership, pos support)."""
668 out: dict = {}
669 for hook_name, spec in intervene.items():
670 validated = validate_spec(hook_name, spec, width=self._hook_widths.get(hook_name))
671 if hook_name not in self.supported_hook_points:
672 raise ValueError(
673 f"Cannot intervene on {hook_name!r}: not in supported_hook_points."
674 )
675 if validated.get("pos") is not None:
676 if self._enable_batching:
677 # The batched/eager path applies ops to the raw tensor, not the
678 # (max_n, width) affine buffers, so it has no position surface.
679 raise NotImplementedError(
680 f"Intervention {hook_name!r}: per-position 'pos' is not supported on the "
681 "batched/eager path. Boot the compiled path (enable_batching=False) with "
682 "enable_position_interventions=True."
683 )
684 if not self._enable_position_interventions:
685 # Default affine buffers are (width,) and broadcast across every position;
686 # honoring 'pos' needs the (max_n, width) buffers allocated at boot.
687 raise NotImplementedError(
688 f"Intervention {hook_name!r}: per-position 'pos' requires "
689 "boot_vllm(enable_position_interventions=True) (its default affine "
690 "buffers broadcast across all positions). Use the Inspect/HF backend for "
691 "position-scoped patching, or drop 'pos' for a whole-sequence edit."
692 )
693 out[hook_name] = validated
694 return out
696 @staticmethod
697 def _reject_pos_beyond_seq(specs: Mapping[str, Any], seq_len: int) -> None:
698 """Fail loud if a spec's 'pos' targets a row past the actual prompt length.
700 The compiled hook applies the affine over rows ``[0, seq_len)`` only, so a ``pos``
701 in ``[seq_len, max_num_batched_tokens)`` clears the driver's non-negativity check
702 and the worker's buffer-capacity check yet is never read — a silent no-op. Bound it
703 against the real sequence length (known once ``ids_list`` exists) and raise instead.
704 """
705 for hook_name, spec in specs.items():
706 pos = spec.get("pos")
707 if pos is None: 707 ↛ 708line 707 didn't jump to line 708 because the condition on line 707 was never true
708 continue
709 idx = [pos] if isinstance(pos, int) else list(pos)
710 bad = [p for p in idx if p >= seq_len]
711 if bad:
712 raise ValueError(
713 f"Intervention {hook_name!r}: 'pos' {bad} is beyond the prompt length "
714 f"{seq_len} (positions are 0-indexed); the edit would be silently ignored."
715 )
717 @staticmethod
718 def _normalize_input_ids(input_ids: Any) -> list:
719 """Coerce input_ids to a flat list[int] for ``TokensPrompt``; batch_size=1 only."""
720 if isinstance(input_ids, torch.Tensor): 720 ↛ 723line 720 didn't jump to line 723 because the condition on line 720 was always true
721 ids_list = input_ids.tolist()
722 else:
723 ids_list = list(input_ids)
724 if ids_list and isinstance(ids_list[0], list): 724 ↛ 728line 724 didn't jump to line 728 because the condition on line 724 was always true
725 if len(ids_list) != 1:
726 raise NotImplementedError("VLLMDriver supports batch_size=1 only.")
727 ids_list = ids_list[0]
728 return ids_list
730 @staticmethod
731 def _normalize_input_ids_batched(input_ids: Any) -> list[list[int]]:
732 """Coerce to ``list[list[int]]`` (one per prompt); accepts 1-D/2-D tensor,
733 flat list (single prompt), or ragged list-of-lists."""
734 if isinstance(input_ids, torch.Tensor):
735 if input_ids.dim() == 1:
736 return [input_ids.tolist()]
737 return [row.tolist() for row in input_ids]
738 seq = list(input_ids)
739 if seq and isinstance(seq[0], (list, tuple)):
740 return [list(row) for row in seq]
741 return [list(seq)]