Coverage for transformer_lens/model_bridge/sources/vllm/source.py: 99%
90 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-10-10 05:10 +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
12# Module import (not a name import): architecture_adapter_factory imports model_bridge,
13# which imports this module, so its names may not be defined yet at this point.
14from transformer_lens.factories import architecture_adapter_factory
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), followed by the engine's softcap and scale,
175 from the captured
176 final-norm activation — valid at every position, so ``return_type`` in
177 ``{"loss", "both"}`` works. If the unembedding weight is unreachable the
178 driver falls back to the sampler's final-position log-probs (earlier
179 positions ``-inf``), declares ``provides_sequence_logits=False``, and the
180 bridge then rejects loss.
182 GPU memory cost: each capture buffer is ``max_num_batched_tokens × width`` at
183 the model's dtype. For Llama-3.2-1B at fp16 with ``max_num_batched_tokens=2048``,
184 the unembed buffer alone is ~525 MB (2048 × 128256 × 2 bytes); residual-stream
185 buffers add ~8 MB per hook. The affine intervention hook also allocates a
186 transient output-shape tensor per forward (even in identity mode), so peak
187 forward memory is ~1.5× the capture buffers' resident size.
189 KV-cache footprint: vLLM reserves KV cache sized for ``max_model_len`` × layers
190 × heads × head_dim. If ``max_model_len`` is left as ``None``, vLLM uses the
191 model's native context (e.g. 131072 for Llama-3.2-1B) — easily 4+ GiB even
192 on a 1B model. Pass an explicit ``max_model_len`` (e.g. ``2048`` for typical
193 mech-interp prompts) to keep the budget on smaller GPUs.
195 ``enable_batching`` switches to the eager batched path (``enforce_eager``,
196 ``batch_size > 1``) — the throughput path for SAE/probe data collection.
197 Default ``False`` keeps the compile-validated single-prompt path. Batched
198 caches are right-padded with zeros to the longest sequence.
200 ``enable_position_interventions`` widens each hook's affine scale/bias buffers
201 from ``(width,)`` to ``(max_num_batched_tokens, width)`` so an intervention spec
202 can carry a ``pos`` field (int or list[int]) that scopes the edit to specific
203 sequence positions — position-scoped activation patching / tensor injection.
204 Costs ~2× extra resident GPU memory across all hooks (the scale and bias buffers
205 join the already-``(max_n, width)`` capture buffer), so it is opt-in and defaults
206 ``False``. Compiled-path only — incompatible with ``enable_batching``.
208 ``tensor_parallel_size`` > 1 enables single-node tensor parallelism. Every
209 capture point the overlay hooks is post-all-reduce and therefore replicated
210 across a stage's TP ranks; ``pipeline_parallel_size`` > 1 shards layers across
211 stages, each owning the hooks that actually fire on it.
212 Capture reads merge across ranks with a first-forward layout check that fails
213 loud on installation drift or non-replicated hook points; sharded/stage-local
214 weights (``lm_head``, final norm) are gathered for logit reconstruction and the
215 ln_final un-fold. Single-node only (the spec channel is an env var, which never
216 reaches Ray remote workers); both are incompatible with ``enable_batching``
217 (per-rank chunk boundaries are unvalidated).
218 """
219 if enable_position_interventions and enable_batching:
220 raise ValueError(
221 "enable_position_interventions requires the compiled path and is incompatible "
222 "with enable_batching=True (the batched/eager path has no affine buffers)."
223 )
224 for kwarg_name, size in (
225 ("tensor_parallel_size", tensor_parallel_size),
226 ("pipeline_parallel_size", pipeline_parallel_size),
227 ):
228 if not isinstance(size, int) or size < 1:
229 raise ValueError(f"{kwarg_name} must be a positive int; got {size!r}.")
230 if (tensor_parallel_size > 1 or pipeline_parallel_size > 1) and enable_batching:
231 raise ValueError(
232 "enable_batching=True with tensor/pipeline parallelism is unsupported: the "
233 "eager batched path's per-rank chunk boundaries are unvalidated. "
234 "Use the compiled single-prompt path."
235 )
236 _reject_locked_overrides(vllm_kwargs)
237 # Import-check first: fail with an actionable message before any network I/O
238 # or plugin state mutation.
239 try:
240 from vllm import LLM
241 except ImportError as exc:
242 raise ImportError(
243 "boot_vllm requires vLLM (Linux + CUDA). Install with "
244 'pip install "transformer-lens[vllm]" or uv sync --extra vllm; '
245 "the driver is validated against vllm 0.20.x."
246 ) from exc
248 from transformers import AutoConfig, AutoTokenizer
250 # Resolve architecture WITHOUT loading weights so we can tell the plugin
251 # which dot-paths to hook before LLM(...) constructs the worker. Shared
252 # resolution (not architectures[0]) handles architectures=None configs via
253 # model_type and rejects unsupported archs before the expensive engine boot.
254 hf_token = get_hf_token()
255 hf_config_preview = AutoConfig.from_pretrained(model_name, token=hf_token)
256 architecture = determine_architecture_from_hf_config(hf_config_preview)
257 overlay = get_overlay(architecture)
259 resolved_dtype = dtype or _dtype_from_hf_config(hf_config_preview)
261 # Batched capture reads query_start_loc from the forward context, untraceable
262 # under torch.compile — so the batched path must run eager.
263 eager_kwargs: Dict[str, Any] = {"enforce_eager": True} if enable_batching else {}
265 llm = construct_instrumented_llm(
266 model_name,
267 capture_specs=overlay.capture_specs(hf_config_preview),
268 max_num_batched_tokens=max_num_batched_tokens,
269 dtype=resolved_dtype,
270 enable_batching=enable_batching,
271 enable_position_interventions=enable_position_interventions,
272 tensor_parallel_size=tensor_parallel_size,
273 pipeline_parallel_size=pipeline_parallel_size,
274 llm_kwargs={
275 "gpu_memory_utilization": gpu_memory_utilization,
276 "max_model_len": max_model_len,
277 # Allow full-vocab logprobs so the fallback path can synthesize logits when
278 # host-side reconstruction is unavailable (vLLM caps logprobs to this value;
279 # default 20 is too small for mech-interp).
280 "max_logprobs": hf_config_preview.vocab_size,
281 # Always explicit — vLLM's "auto" downcasts fp32 checkpoints to fp16, which
282 # would leave the capture/affine buffers (allocated at resolved_dtype) at a
283 # different dtype than the engine's activations.
284 "dtype": dtype_name(resolved_dtype),
285 **eager_kwargs,
286 **_LOCKED_KWARGS,
287 **vllm_kwargs,
288 },
289 )
290 hf_config = extract_hf_config(llm)
292 # Build the adapter (RemoteBridge skips adapter.prepare_model — there's no
293 # local model tree to walk). Use the shared HF→TL config builder so
294 # bridge_config is a real TransformerBridgeConfig with all dataclass
295 # defaults (d_vocab_out=-1, etc.) — not a deep-copied HF config with
296 # extra fields, which is missing the TL-only attributes BridgeCore reads.
297 bridge_config = build_bridge_config_from_hf(hf_config, architecture, model_name, resolved_dtype)
298 adapter = architecture_adapter_factory.ArchitectureAdapterFactory.select_architecture_adapter(
299 bridge_config
300 )
302 if tokenizer is None and not skip_tokenizer_for_modality(adapter.cfg):
303 tokenizer = AutoTokenizer.from_pretrained(model_name, token=hf_token)
304 if tokenizer is not None:
305 # Match boot_transformers' tokenizer setup so to_tokens(str) is token-identical.
306 tokenizer = configure_tokenizer(tokenizer, adapter.cfg)
308 driver = VLLMDriver(
309 llm=llm,
310 adapter=adapter,
311 tokenizer=tokenizer,
312 overlay=overlay,
313 hf_config=hf_config,
314 max_num_batched_tokens=max_num_batched_tokens,
315 enable_batching=enable_batching,
316 enable_position_interventions=enable_position_interventions,
317 tensor_parallel_size=tensor_parallel_size,
318 pipeline_parallel_size=pipeline_parallel_size,
319 )
320 # One-time unembedding fetch: caches the weight for per-forward reconstruction and
321 # downgrades provides_sequence_logits honestly if no unembedding is reachable.
322 driver.probe_logit_reconstruction()
323 bridge = RemoteBridge(adapter=adapter, tokenizer=tokenizer, driver=driver)
324 _log_hook_summary(model_name, architecture, driver)
325 return bridge
328def _reject_locked_overrides(vllm_kwargs: Dict[str, Any]) -> None:
329 for key, locked in _LOCKED_KWARGS.items():
330 if key in vllm_kwargs and vllm_kwargs[key] != locked:
331 raise ValueError(
332 f"boot_vllm forces {key}={locked}; caller passed {key}={vllm_kwargs[key]}. "
333 "Prefix caching, continuous batching, and vLLM-owned tokenizers are "
334 "unsupported — each breaks the capture-read invariants. (TP/PP are "
335 "supported via the tensor_parallel_size/pipeline_parallel_size kwargs.)"
336 )
339def _dtype_from_hf_config(hf_config: Any) -> torch.dtype:
340 raw = getattr(hf_config, "torch_dtype", None)
341 if isinstance(raw, torch.dtype):
342 return raw
343 if isinstance(raw, str):
344 return getattr(torch, raw, torch.float16)
345 return torch.float16
348def _log_hook_summary(model_name: str, architecture: str, driver: VLLMDriver) -> None:
349 """Log the fireable and non-fireable hook sets so users don't have to grep."""
350 log = logging.getLogger("transformer_lens.vllm")
351 fireable = sorted(driver.supported_hook_points)
352 log.info(
353 "vLLM source on %s (%s) captures %d hook(s): %s",
354 model_name,
355 architecture,
356 len(fireable),
357 ", ".join(fireable),
358 )
359 nonfiring = sorted(driver.non_fireable_hook_points)
360 if nonfiring: 360 ↛ exitline 360 didn't return from function '_log_hook_summary' because the condition on line 360 was always true
361 log.info(
362 "vLLM source on %s (%s) cannot fire %d hook(s) (vLLM fuses these): %s. "
363 "Use boot_transformers() if you need them.",
364 model_name,
365 architecture,
366 len(nonfiring),
367 ", ".join(nonfiring),
368 )
371__all__ = ["boot_vllm", "construct_instrumented_llm"]