transformer_lens.model_bridge.sources.vllm.driver module¶

vLLM Driver: forward dispatches via llm.generate; captures via collective_rpc.

class transformer_lens.model_bridge.sources.vllm.driver.VLLMDriver(llm: Any, adapter: Any, tokenizer: Any, overlay: Any, hf_config: Any, max_num_batched_tokens: int, enable_batching: bool = False, enable_position_interventions: bool = False, tensor_parallel_size: int = 1, pipeline_parallel_size: int = 1)¶

Bases: DriverBase

Driver wrapping a vLLM LLM; captures via collective_rpc.

close() → None¶: No-op default. Override when the driver owns releasable resources.

forward(input_ids: TensorLike | Sequence[Any] | None = None, *, capture: tuple[str, ...] = (), intervene: Mapping[str, Callable[[TensorLike], TensorLike] | Mapping[str, Any]] | None = None, max_new_tokens: int = 1, return_logits: bool = True, **kwargs: Any) → ForwardResult¶

get_param(dotted_name: str) → Tensor | None¶: Fetch a named model tensor (e.g. model.norm.weight) for conversions the bridge can’t otherwise do (ln_final post→pre-weight; see the overlay). Gathered across ranks: replicated params return one copy, vocab-sharded weights concatenate, stage-local params (PP) come from whichever rank owns them. None if closed or the path resolves nowhere.

probe_logit_reconstruction() → bool¶: One-time unembedding fetch + cache; downgrades provides_sequence_logits honestly when no unembedding is reachable (the fallback path is final-position log-probs, which cannot back a loss). Idempotent — later calls return the cached availability.

provides_sequence_logits: bool = True¶