Coverage for transformer_lens/model_bridge/sources/vllm/source.py: 99%
90 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-09-21 19:27 +0000
1"""``boot_vllm`` — construct a vLLM LLM, wrap it in a RemoteBridge via VLLMDriver."""
2from __future__ import annotations
4import logging
5import os
6import threading
7import warnings
8from typing import Any, Dict, Optional
10import torch
12from transformer_lens.factories.architecture_adapter_factory import (
13 ArchitectureAdapterFactory,
14)
15from transformer_lens.model_bridge.remote_bridge import RemoteBridge
16from transformer_lens.model_bridge.sources._bridge_builder import (
17 build_bridge_config_from_hf,
18 configure_tokenizer,
19 skip_tokenizer_for_modality,
20)
21from transformer_lens.model_bridge.sources._hf_format import (
22 determine_architecture_from_hf_config,
23)
24from transformer_lens.utilities.hf_utils import get_hf_token
26from . import plugin
27from .driver import VLLMDriver
28from .internals import extract_hf_config, verify_hook_coverage
29from .overlays import get_overlay
30from .worker_extension import dtype_name
32# Forced LLM(...) kwargs that the capture-hook design depends on. Caller override → ValueError.
33# enable_prefix_caching MUST stay off: a prefix-cache hit makes the prefill compute only the
34# uncached suffix, so hooks fire for a subset of positions — captures land row-misaligned,
35# interventions skip cached positions, and an intervened forward writes poisoned K/V that a
36# later clean forward on the same prompt would silently reuse.
37_LOCKED_KWARGS = {
38 "skip_tokenizer_init": True,
39 "disable_log_stats": True,
40 "enable_prefix_caching": False,
41}
43# Serializes the configure() → LLM(...) → clear() handoff: the spec channel is a
44# process-wide env var, so interleaved boots would cross-wire capture specs between engines.
45_BOOT_LOCK = threading.Lock()
47_WORKER_EXTENSION_CLS = (
48 "transformer_lens.model_bridge.sources.vllm.worker_extension.TLWorkerExtension"
49)
52def construct_instrumented_llm(
53 model_name: str,
54 *,
55 capture_specs: Dict[str, Any],
56 max_num_batched_tokens: int,
57 dtype: torch.dtype,
58 enable_batching: bool = False,
59 enable_position_interventions: bool = False,
60 tensor_parallel_size: int = 1,
61 pipeline_parallel_size: int = 1,
62 llm_kwargs: Dict[str, Any],
63) -> Any:
64 """The one place a TL-instrumented ``vllm.LLM`` is constructed — shared by
65 ``boot_vllm`` and the Inspect vLLM provider so the capture contract can't drift
66 between them. Owns everything hook correctness depends on:
68 - ``VLLM_DISABLE_COMPILE_CACHE=1``: our hooks are traced INTO vLLM's compiled
69 graph, but its compile cache is keyed only on its own config — a cached
70 artifact from a differently-instrumented process either crashes at AOT load
71 (bytecode binds the hook closures) or silently serves a hookless graph.
72 - ``VLLM_ENABLE_V1_MULTIPROCESSING``: forced ``"0"`` single-rank (in-process
73 worker, the historical GPU-validated path); parallel boots spawn workers,
74 which read the specs via the plugin's env channel — and must not inherit a
75 stale ``"0"`` that would force the uni-process executor.
76 - ``enable_prefix_caching=False``: a prefix-cache hit computes only the uncached
77 suffix, so captures land row-misaligned, interventions skip cached positions,
78 and an intervened forward writes poisoned K/V a later clean forward reuses.
79 - configure → register → ``LLM(...)`` → clear under ``_BOOT_LOCK``, then a
80 hook-coverage check so a spec that landed on no rank fails here, not as zeros.
82 ``llm_kwargs`` carries the caller's remaining ``LLM(...)`` arguments; restating a
83 contract kwarg is allowed only at the same value (callers validate overrides).
84 """
85 from vllm import LLM
87 os.environ["VLLM_DISABLE_COMPILE_CACHE"] = "1"
88 if tensor_parallel_size == 1 and pipeline_parallel_size == 1:
89 existing_mp = os.environ.get("VLLM_ENABLE_V1_MULTIPROCESSING")
90 if existing_mp not in (None, "0"):
91 warnings.warn(
92 f"VLLM_ENABLE_V1_MULTIPROCESSING={existing_mp!r} overridden to '0' — "
93 "single-rank TL boots keep the worker in-process.",
94 UserWarning,
95 stacklevel=3,
96 )
97 os.environ["VLLM_ENABLE_V1_MULTIPROCESSING"] = "0"
98 elif os.environ.get("VLLM_ENABLE_V1_MULTIPROCESSING") == "0":
99 os.environ.pop("VLLM_ENABLE_V1_MULTIPROCESSING")
101 with _BOOT_LOCK:
102 plugin.configure(
103 capture_specs=capture_specs,
104 max_num_batched_tokens=max_num_batched_tokens,
105 dtype=dtype,
106 enable_batching=enable_batching,
107 enable_position_interventions=enable_position_interventions,
108 )
109 plugin.register()
110 try:
111 llm = LLM(
112 **{
113 "model": model_name,
114 # vLLM defaults this to 8192 for chunked prefill; the capture buffers
115 # are sized to it, so the compiled dynamic-shape range must match —
116 # otherwise Dynamo's symbolic hint exceeds the buffer dim at compile.
117 "max_num_batched_tokens": max_num_batched_tokens,
118 # Mixes tl_* methods into the Worker for collective_rpc (multiple
119 # inheritance; the tl_ prefix avoids attribute collisions).
120 "worker_extension_cls": _WORKER_EXTENSION_CLS,
121 "enable_prefix_caching": False,
122 "tensor_parallel_size": tensor_parallel_size,
123 "pipeline_parallel_size": pipeline_parallel_size,
124 **llm_kwargs,
125 }
126 )
127 finally:
128 # Always clear, even on a failed boot: stale specs would make the next
129 # in-process vllm.LLM(...) walk our dot-paths on a foreign model.
130 plugin.clear_config()
132 # Hook installation skips modules absent on a rank (PP shards); a spec that
133 # landed on NO rank is a broken dot-path and must fail here, not read zeros.
134 verify_hook_coverage(llm)
135 return llm
138def boot_vllm(
139 model_name: str,
140 tokenizer: Optional[Any] = None,
141 dtype: Optional[torch.dtype] = None,
142 gpu_memory_utilization: float = 0.5,
143 max_model_len: Optional[int] = None,
144 max_num_batched_tokens: int = 2048,
145 enable_batching: bool = False,
146 enable_position_interventions: bool = False,
147 tensor_parallel_size: int = 1,
148 pipeline_parallel_size: int = 1,
149 **vllm_kwargs: Any,
150) -> RemoteBridge:
151 """Boot a model via vLLM and wrap it in a :class:`RemoteBridge` via :class:`VLLMDriver`.
153 vLLM drives the forward pass (PagedAttention + ``torch.compile`` + CUDA graphs).
154 Capture buffers are populated by hooks the plugin installs pre-compile inside
155 the worker; they come back via ``collective_rpc`` and replay through the
156 bridge's HookPoint tree.
158 **Scope vs vllm-lens:** vllm-lens is observation-only. This source extends to
159 observation + spec-vocabulary *mutation* — each capture hook also applies an
160 affine transform ``output = output * scale + bias`` (default identity), so
161 interventions (``suppress`` / ``scale`` / ``add`` / ``set``) propagate to
162 downstream layers. The hook's return value replaces the module output per
163 PyTorch ``register_forward_hook`` semantics. The mutation path under
164 torch.compile + CUDA graphs is exercised end-to-end by
165 ``demos/vLLM_Bridge_Integration_Test.ipynb`` (a manual GPU run, not CI);
166 unit tests cover the dispatch protocol only.
168 Some captures use vLLM-native conventions that differ from HF/HT; see
169 :mod:`transformer_lens.model_bridge.sources.vllm.overlays.decoder_only` for
170 which hooks diverge and the conversion to apply for HT-equivalent values.
172 **Returned logits are reconstructed full-sequence logits.** vLLM's sampler
173 bypasses ``lm_head``, so the driver rebuilds real logits host-side as
174 ``ln_final @ lm_head.weight.T`` (+ bias, + Gemma soft-cap) from the captured
175 final-norm activation — valid at every position, so ``return_type`` in
176 ``{"loss", "both"}`` works. If the unembedding weight is unreachable the
177 driver falls back to the sampler's final-position log-probs (earlier
178 positions ``-inf``), declares ``provides_sequence_logits=False``, and the
179 bridge then rejects loss.
181 GPU memory cost: each capture buffer is ``max_num_batched_tokens × width`` at
182 the model's dtype. For Llama-3.2-1B at fp16 with ``max_num_batched_tokens=2048``,
183 the unembed buffer alone is ~525 MB (2048 × 128256 × 2 bytes); residual-stream
184 buffers add ~8 MB per hook. The affine intervention hook also allocates a
185 transient output-shape tensor per forward (even in identity mode), so peak
186 forward memory is ~1.5× the capture buffers' resident size.
188 KV-cache footprint: vLLM reserves KV cache sized for ``max_model_len`` × layers
189 × heads × head_dim. If ``max_model_len`` is left as ``None``, vLLM uses the
190 model's native context (e.g. 131072 for Llama-3.2-1B) — easily 4+ GiB even
191 on a 1B model. Pass an explicit ``max_model_len`` (e.g. ``2048`` for typical
192 mech-interp prompts) to keep the budget on smaller GPUs.
194 ``enable_batching`` switches to the eager batched path (``enforce_eager``,
195 ``batch_size > 1``) — the throughput path for SAE/probe data collection.
196 Default ``False`` keeps the compile-validated single-prompt path. Batched
197 caches are right-padded with zeros to the longest sequence.
199 ``enable_position_interventions`` widens each hook's affine scale/bias buffers
200 from ``(width,)`` to ``(max_num_batched_tokens, width)`` so an intervention spec
201 can carry a ``pos`` field (int or list[int]) that scopes the edit to specific
202 sequence positions — position-scoped activation patching / tensor injection.
203 Costs ~2× extra resident GPU memory across all hooks (the scale and bias buffers
204 join the already-``(max_n, width)`` capture buffer), so it is opt-in and defaults
205 ``False``. Compiled-path only — incompatible with ``enable_batching``.
207 ``tensor_parallel_size`` > 1 enables single-node tensor parallelism. Every
208 capture point the overlay hooks is post-all-reduce and therefore replicated
209 across a stage's TP ranks; ``pipeline_parallel_size`` > 1 shards layers across
210 stages, each owning the hooks that actually fire on it.
211 Capture reads merge across ranks with a first-forward layout check that fails
212 loud on installation drift or non-replicated hook points; sharded/stage-local
213 weights (``lm_head``, final norm) are gathered for logit reconstruction and the
214 ln_final un-fold. Single-node only (the spec channel is an env var, which never
215 reaches Ray remote workers); both are incompatible with ``enable_batching``
216 (per-rank chunk boundaries are unvalidated).
217 """
218 if enable_position_interventions and enable_batching:
219 raise ValueError(
220 "enable_position_interventions requires the compiled path and is incompatible "
221 "with enable_batching=True (the batched/eager path has no affine buffers)."
222 )
223 for kwarg_name, size in (
224 ("tensor_parallel_size", tensor_parallel_size),
225 ("pipeline_parallel_size", pipeline_parallel_size),
226 ):
227 if not isinstance(size, int) or size < 1:
228 raise ValueError(f"{kwarg_name} must be a positive int; got {size!r}.")
229 if (tensor_parallel_size > 1 or pipeline_parallel_size > 1) and enable_batching:
230 raise ValueError(
231 "enable_batching=True with tensor/pipeline parallelism is unsupported: the "
232 "eager batched path's per-rank chunk boundaries are unvalidated. "
233 "Use the compiled single-prompt path."
234 )
235 _reject_locked_overrides(vllm_kwargs)
236 # Import-check first: fail with an actionable message before any network I/O
237 # or plugin state mutation.
238 try:
239 from vllm import LLM
240 except ImportError as exc:
241 raise ImportError(
242 "boot_vllm requires vLLM (Linux + CUDA). Install with "
243 'pip install "transformer-lens[vllm]" or uv sync --extra vllm; '
244 "the driver is validated against vllm 0.20.x."
245 ) from exc
247 from transformers import AutoConfig, AutoTokenizer
249 # Resolve architecture WITHOUT loading weights so we can tell the plugin
250 # which dot-paths to hook before LLM(...) constructs the worker. Shared
251 # resolution (not architectures[0]) handles architectures=None configs via
252 # model_type and rejects unsupported archs before the expensive engine boot.
253 hf_token = get_hf_token()
254 hf_config_preview = AutoConfig.from_pretrained(model_name, token=hf_token)
255 architecture = determine_architecture_from_hf_config(hf_config_preview)
256 overlay = get_overlay(architecture)
258 resolved_dtype = dtype or _dtype_from_hf_config(hf_config_preview)
260 # Batched capture reads query_start_loc from the forward context, untraceable
261 # under torch.compile — so the batched path must run eager.
262 eager_kwargs: Dict[str, Any] = {"enforce_eager": True} if enable_batching else {}
264 llm = construct_instrumented_llm(
265 model_name,
266 capture_specs=overlay.capture_specs(hf_config_preview),
267 max_num_batched_tokens=max_num_batched_tokens,
268 dtype=resolved_dtype,
269 enable_batching=enable_batching,
270 enable_position_interventions=enable_position_interventions,
271 tensor_parallel_size=tensor_parallel_size,
272 pipeline_parallel_size=pipeline_parallel_size,
273 llm_kwargs={
274 "gpu_memory_utilization": gpu_memory_utilization,
275 "max_model_len": max_model_len,
276 # Allow full-vocab logprobs so the fallback path can synthesize logits when
277 # host-side reconstruction is unavailable (vLLM caps logprobs to this value;
278 # default 20 is too small for mech-interp).
279 "max_logprobs": hf_config_preview.vocab_size,
280 # Always explicit — vLLM's "auto" downcasts fp32 checkpoints to fp16, which
281 # would leave the capture/affine buffers (allocated at resolved_dtype) at a
282 # different dtype than the engine's activations.
283 "dtype": dtype_name(resolved_dtype),
284 **eager_kwargs,
285 **_LOCKED_KWARGS,
286 **vllm_kwargs,
287 },
288 )
289 hf_config = extract_hf_config(llm)
291 # Build the adapter (RemoteBridge skips adapter.prepare_model — there's no
292 # local model tree to walk). Use the shared HF→TL config builder so
293 # bridge_config is a real TransformerBridgeConfig with all dataclass
294 # defaults (d_vocab_out=-1, etc.) — not a deep-copied HF config with
295 # extra fields, which is missing the TL-only attributes BridgeCore reads.
296 bridge_config = build_bridge_config_from_hf(hf_config, architecture, model_name, resolved_dtype)
297 adapter = ArchitectureAdapterFactory.select_architecture_adapter(bridge_config)
299 if tokenizer is None and not skip_tokenizer_for_modality(adapter.cfg):
300 tokenizer = AutoTokenizer.from_pretrained(model_name, token=hf_token)
301 if tokenizer is not None:
302 # Match boot_transformers' tokenizer setup so to_tokens(str) is token-identical.
303 tokenizer = configure_tokenizer(tokenizer, adapter.cfg)
305 driver = VLLMDriver(
306 llm=llm,
307 adapter=adapter,
308 tokenizer=tokenizer,
309 overlay=overlay,
310 hf_config=hf_config,
311 max_num_batched_tokens=max_num_batched_tokens,
312 enable_batching=enable_batching,
313 enable_position_interventions=enable_position_interventions,
314 tensor_parallel_size=tensor_parallel_size,
315 pipeline_parallel_size=pipeline_parallel_size,
316 )
317 # One-time unembedding fetch: caches the weight for per-forward reconstruction and
318 # downgrades provides_sequence_logits honestly if no unembedding is reachable.
319 driver.probe_logit_reconstruction()
320 bridge = RemoteBridge(adapter=adapter, tokenizer=tokenizer, driver=driver)
321 _log_hook_summary(model_name, architecture, driver)
322 return bridge
325def _reject_locked_overrides(vllm_kwargs: Dict[str, Any]) -> None:
326 for key, locked in _LOCKED_KWARGS.items():
327 if key in vllm_kwargs and vllm_kwargs[key] != locked:
328 raise ValueError(
329 f"boot_vllm forces {key}={locked}; caller passed {key}={vllm_kwargs[key]}. "
330 "Prefix caching, continuous batching, and vLLM-owned tokenizers are "
331 "unsupported — each breaks the capture-read invariants. (TP/PP are "
332 "supported via the tensor_parallel_size/pipeline_parallel_size kwargs.)"
333 )
336def _dtype_from_hf_config(hf_config: Any) -> torch.dtype:
337 raw = getattr(hf_config, "torch_dtype", None)
338 if isinstance(raw, torch.dtype):
339 return raw
340 if isinstance(raw, str):
341 return getattr(torch, raw, torch.float16)
342 return torch.float16
345def _log_hook_summary(model_name: str, architecture: str, driver: VLLMDriver) -> None:
346 """Log the fireable and non-fireable hook sets so users don't have to grep."""
347 log = logging.getLogger("transformer_lens.vllm")
348 fireable = sorted(driver.supported_hook_points)
349 log.info(
350 "vLLM source on %s (%s) captures %d hook(s): %s",
351 model_name,
352 architecture,
353 len(fireable),
354 ", ".join(fireable),
355 )
356 nonfiring = sorted(driver.non_fireable_hook_points)
357 if nonfiring: 357 ↛ exitline 357 didn't return from function '_log_hook_summary' because the condition on line 357 was always true
358 log.info(
359 "vLLM source on %s (%s) cannot fire %d hook(s) (vLLM fuses these): %s. "
360 "Use boot_transformers() if you need them.",
361 model_name,
362 architecture,
363 len(nonfiring),
364 ", ".join(nonfiring),
365 )
368__all__ = ["boot_vllm", "construct_instrumented_llm"]