transformer_lens.model_bridge.sources.vllm package

Subpackages

Submodules

Module contents

vLLM source for TransformerBridge.

Provides boot_vllm(), which constructs a vLLM LLM and wraps its inner nn.Module in a TransformerBridge. vLLM drives the forward pass (PagedAttention, torch.compile, CUDA graphs); the bridge surface is populated from GPU buffers written by hooks the plugin installs pre-compile.

transformer_lens.model_bridge.sources.vllm.boot_vllm(model_name: str, tokenizer: Any | None = None, dtype: dtype | None = None, gpu_memory_utilization: float = 0.5, max_model_len: int | None = None, max_num_batched_tokens: int = 2048, enable_batching: bool = False, enable_position_interventions: bool = False, tensor_parallel_size: int = 1, pipeline_parallel_size: int = 1, **vllm_kwargs: Any) → RemoteBridge

Boot a model via vLLM and wrap it in a RemoteBridge via VLLMDriver.

vLLM drives the forward pass (PagedAttention + torch.compile + CUDA graphs). Capture buffers are populated by hooks the plugin installs pre-compile inside the worker; they come back via collective_rpc and replay through the bridge’s HookPoint tree.

Scope vs vllm-lens: vllm-lens is observation-only. This source extends to observation + spec-vocabulary mutation — each capture hook also applies an affine transform output = output * scale + bias (default identity), so interventions (suppress / scale / add / set) propagate to downstream layers. The hook’s return value replaces the module output per PyTorch register_forward_hook semantics. The mutation path under torch.compile + CUDA graphs is exercised end-to-end by demos/vLLM_Bridge_Integration_Test.ipynb (a manual GPU run, not CI); unit tests cover the dispatch protocol only.

Some captures use vLLM-native conventions that differ from HF/HT; see transformer_lens.model_bridge.sources.vllm.overlays.decoder_only for which hooks diverge and the conversion to apply for HT-equivalent values.

Returned logits are reconstructed full-sequence logits. vLLM’s sampler bypasses lm_head, so the driver rebuilds real logits host-side as ln_final @ lm_head.weight.T (+ bias, + Gemma soft-cap) from the captured final-norm activation — valid at every position, so return_type in {"loss", "both"} works. If the unembedding weight is unreachable the driver falls back to the sampler’s final-position log-probs (earlier positions -inf), declares provides_sequence_logits=False, and the bridge then rejects loss.

GPU memory cost: each capture buffer is max_num_batched_tokens × width at the model’s dtype. For Llama-3.2-1B at fp16 with max_num_batched_tokens=2048, the unembed buffer alone is ~525 MB (2048 × 128256 × 2 bytes); residual-stream buffers add ~8 MB per hook. The affine intervention hook also allocates a transient output-shape tensor per forward (even in identity mode), so peak forward memory is ~1.5× the capture buffers’ resident size.

KV-cache footprint: vLLM reserves KV cache sized for max_model_len × layers × heads × head_dim. If max_model_len is left as None, vLLM uses the model’s native context (e.g. 131072 for Llama-3.2-1B) — easily 4+ GiB even on a 1B model. Pass an explicit max_model_len (e.g. 2048 for typical mech-interp prompts) to keep the budget on smaller GPUs.

enable_batching switches to the eager batched path (enforce_eager, batch_size > 1) — the throughput path for SAE/probe data collection. Default False keeps the compile-validated single-prompt path. Batched caches are right-padded with zeros to the longest sequence.

enable_position_interventions widens each hook’s affine scale/bias buffers from (width,) to (max_num_batched_tokens, width) so an intervention spec can carry a pos field (int or list[int]) that scopes the edit to specific sequence positions — position-scoped activation patching / tensor injection. Costs ~2× extra resident GPU memory across all hooks (the scale and bias buffers join the already-(max_n, width) capture buffer), so it is opt-in and defaults False. Compiled-path only — incompatible with enable_batching.

tensor_parallel_size > 1 enables single-node tensor parallelism. Every capture point the overlay hooks is post-all-reduce and therefore replicated across a stage’s TP ranks; pipeline_parallel_size > 1 shards layers across stages, each owning the hooks that actually fire on it. Capture reads merge across ranks with a first-forward layout check that fails loud on installation drift or non-replicated hook points; sharded/stage-local weights (lm_head, final norm) are gathered for logit reconstruction and the ln_final un-fold. Single-node only (the spec channel is an env var, which never reaches Ray remote workers); both are incompatible with enable_batching (per-rank chunk boundaries are unvalidated).