Merge branch 'main' into dev

JadenFiotto-Kaufman · JadenFiotto-Kaufman · commit eef1026b1244 · 2026-01-12T19:38:47.000-05:00
diff --git a/llms.md b/llms.md
@@ -1026,7 +1026,7 @@ with model.trace("Hello"):
     # Layer 1's output now equals layer 0's output
     layer1_out = model.transformer.h[1].output.save()
 
-assert torch.equal(layer0_out.value[0], layer1_out[0])
+assert torch.equal(layer0_out[0], layer1_out[0])
 ```
 
 ### Skipping Constraints
diff --git a/src/nnsight/modeling/language.py b/src/nnsight/modeling/language.py
@@ -207,26 +207,24 @@ def _tokenize(
 
     def _prepare_input(
         self,
-        *inputs: Tuple[
-            Union[
-                str,
-                List[str],
-                List[List[str]],
-                List[int],
-                List[List[int]],
-                torch.Tensor,
-                List[torch.Tensor],
-                Dict[str, Any],
-                BatchEncoding,
-            ]
+        *inputs: Union[
+            str,
+            List[str],
+            List[List[str]],
+            List[int],
+            List[List[int]],
+            torch.Tensor,
+            List[torch.Tensor],
+            Dict[str, Any],
+            BatchEncoding,
         ],
         input_ids: Union[
             List[int], List[List[int]], torch.Tensor, List[torch.Tensor]
         ] = None,
         labels: Any = None,
         attention_mask: Any = None,
         **kwargs,
-    ) -> Tuple[BatchEncoding, int]:
+    ) -> Tuple[Tuple[()], Dict[str, Any]]:
         
         if input_ids is not None:
 
diff --git a/src/nnsight/modeling/vllm/model_runners/GPUModelRunner.py b/src/nnsight/modeling/vllm/model_runners/GPUModelRunner.py
@@ -3,7 +3,7 @@
 from vllm.distributed.parallel_state import get_pp_group
 from vllm.outputs import RequestOutput
 from vllm.sequence import IntermediateTensors
-from vllm.transformers_utils.tokenizer import init_tokenizer_from_configs
+from vllm.tokenizers import cached_tokenizer_from_config
 from vllm.v1.outputs import ModelRunnerOutput
 from vllm.v1.worker.gpu_model_runner import GPUModelRunner
 
@@ -186,7 +186,7 @@ def load_model(self, *args, **kwargs) -> None:
 
         self.nnsight_model = VLLM(self.model)
 
-        self.nnsight_model.tokenizer = init_tokenizer_from_configs(self.model_config)
+        self.nnsight_model.tokenizer = cached_tokenizer_from_config(self.model_config)
 
         self.nnsight_model._interleaver.mediators = []
 
diff --git a/src/nnsight/modeling/vllm/vllm.py b/src/nnsight/modeling/vllm/vllm.py
@@ -5,7 +5,8 @@
 
 from vllm.model_executor.model_loader.dummy_loader import DummyModelLoader
 from typing import TYPE_CHECKING, Any, Callable, Dict, List, Tuple, Union
-from vllm.transformers_utils.tokenizer import init_tokenizer_from_configs
+from vllm.tokenizers import cached_tokenizer_from_config
+from vllm.inputs import TokensPrompt
 
 from vllm import LLM, envs
 from vllm.distributed import (
@@ -18,6 +19,7 @@
 from vllm.entrypoints.llm import LLM
 
 from ...intervention.envoy import Envoy
+from ...intervention.tracing.tracer import ScanningTracer
 from ...intervention.tracing.util import push_variables
 from ...util import WrapperModule
 from ..mixins import RemoteableMixin
@@ -118,7 +120,9 @@ def _load_meta(self, repo_id: str, **kwargs) -> "Module":
 
         _ROPE_DICT.clear()
 
-        self.tokenizer = init_tokenizer_from_configs(vllm_config.model_config)
+        self.tokenizer = cached_tokenizer_from_config(vllm_config.model_config)
+        if getattr(self.tokenizer, "pad_token", None) is None:
+            self.tokenizer.pad_token = self.tokenizer.eos_token
 
         return model
 
@@ -150,8 +154,42 @@ def _prepare_input(
         params = []
 
         for arg in args:
-
-            if not type(arg) is list:
+            if arg == []:
+                raise ValueError("Empty list of prompts is not allowed")
+
+            if type(arg) is dict:
+                keys = set(arg.keys())
+                if "input_ids" in keys and keys.issubset(
+                    {"input_ids", "attention_mask"}
+                ):
+                    # is hf tokenizer result
+                    batch_input_ids = arg["input_ids"]
+                    batch_attention_mask = arg.get("attention_mask", None)
+                    if isinstance(batch_input_ids, torch.Tensor):
+                        batch_input_ids = batch_input_ids.tolist()
+                    if isinstance(batch_attention_mask, torch.Tensor):
+                        batch_attention_mask = batch_attention_mask.tolist()
+                    if batch_input_ids == []:
+                        raise ValueError("Empty list of token ids is not allowed")
+                    if isinstance(batch_input_ids[0], int):
+                        # list of token ids
+                        batch_input_ids = [batch_input_ids]
+                        batch_attention_mask = [batch_attention_mask]
+
+                    for input_ids, attention_mask in zip(
+                        batch_input_ids, batch_attention_mask
+                    ):
+                        prompt = TokensPrompt(
+                            prompt_token_ids=[
+                                t for t, m in zip(input_ids, attention_mask) if m != 0
+                            ]
+                        )
+                        prompts.append(prompt)
+                        params.append(NNsightSamplingParams(**kwargs))
+                    continue
+
+            if type(arg) is not list or isinstance(arg[0], int):
+                # if arg is a list of ints (token ids), we also need to wrap it in a list
                 arg = [arg]
 
             for i, prompt in enumerate(arg):
@@ -163,6 +201,9 @@ def _prepare_input(
                 if kwargs != {}:
                     param.is_default_param = False
 
+                if type(prompt) is list and isinstance(prompt[0], int):
+                    prompt = TokensPrompt(prompt_token_ids=prompt)
+
                 prompts.append(prompt)
                 params.append(param)
 
@@ -248,6 +289,9 @@ def __call__(
         push_variables(self._interleaver.mediators[0].info.frame, saves)
 
     def interleave(self, fn: Callable, *args, **kwargs):
+        """Execute the traced function with vLLM, dispatching the engine if needed."""
+        if not self.dispatched and not isinstance(self._interleaver.tracer, ScanningTracer):
+            self.dispatch()
 
         try:
             fn(*args, **kwargs)
diff --git a/tests/test_vllm.py b/tests/test_vllm.py
@@ -393,3 +393,111 @@ def test_tensor_parallelism(self, tp, vllm_gpt2, ET_prompt: str):
         assert next_token != " Paris"
         assert hs.shape == torch.Size([11, 3072])
         assert torch.all(hs[:, 2000:] == 0)
+
+
+# =============================================================================
+# Token Input Compatibility
+# =============================================================================
+
+
+class TestTokenInputs:
+    """Tests for token ID and HuggingFace tokenizer input compatibility."""
+
+    @torch.no_grad()
+    def test_single_token_list(self, vllm_gpt2, ET_prompt: str):
+        """Test passing a single list of token IDs."""
+        token_ids = vllm_gpt2.tokenizer.encode(ET_prompt)
+
+        with vllm_gpt2.trace(token_ids, temperature=0.0, top_p=1):
+            logits = vllm_gpt2.logits.output.save()
+
+        next_token = vllm_gpt2.tokenizer.decode(logits.argmax(dim=-1))
+        assert next_token == " Paris"
+
+    @torch.no_grad()
+    def test_batched_token_lists(self, vllm_gpt2, ET_prompt: str, MSG_prompt: str):
+        """Test passing multiple lists of token IDs."""
+        et_tokens = vllm_gpt2.tokenizer.encode(ET_prompt)
+        msg_tokens = vllm_gpt2.tokenizer.encode(MSG_prompt)
+
+        with vllm_gpt2.trace([et_tokens, msg_tokens], temperature=0.0, top_p=1):
+            logits = vllm_gpt2.logits.output.save()
+
+        assert logits.shape[0] == 2
+        tokens = vllm_gpt2.tokenizer.batch_decode(logits.argmax(dim=-1))
+        assert tokens == [" Paris", " New"]
+
+    @torch.no_grad()
+    def test_hf_tokenizer_dict_single(self, vllm_gpt2, ET_prompt: str):
+        """Test passing HuggingFace tokenizer output dict for single prompt."""
+        hf_output = vllm_gpt2.tokenizer(ET_prompt, return_tensors="pt")
+
+        with vllm_gpt2.trace(dict(hf_output), temperature=0.0, top_p=1):
+            logits = vllm_gpt2.logits.output.save()
+
+        next_token = vllm_gpt2.tokenizer.decode(logits.argmax(dim=-1))
+        assert next_token == " Paris"
+
+    @torch.no_grad()
+    def test_hf_tokenizer_dict_batched(
+        self, vllm_gpt2, ET_prompt: str, MSG_prompt: str
+    ):
+        """Test passing HuggingFace tokenizer output dict for batched prompts."""
+        hf_output = vllm_gpt2.tokenizer(
+            [ET_prompt, MSG_prompt], return_tensors="pt", padding=True
+        )
+
+        with vllm_gpt2.trace(dict(hf_output), temperature=0.0, top_p=1):
+            logits = vllm_gpt2.logits.output.save()
+
+        assert logits.shape[0] == 2
+        tokens = vllm_gpt2.tokenizer.batch_decode(logits.argmax(dim=-1))
+        assert tokens == [" Paris", " New"]
+
+    @torch.no_grad()
+    def test_hf_tokenizer_with_padding_mask(self, vllm_gpt2):
+        """Test that padding tokens are correctly filtered via attention_mask."""
+        short_prompt = "Hello"
+        long_prompt = "The Eiffel Tower is located in the city of"
+
+        hf_output = vllm_gpt2.tokenizer(
+            [short_prompt, long_prompt], return_tensors="pt", padding=True
+        )
+
+        with vllm_gpt2.trace(dict(hf_output), temperature=0.0, top_p=1):
+            logits = vllm_gpt2.logits.output.save()
+
+        assert logits.shape[0] == 2
+        tokens = vllm_gpt2.tokenizer.batch_decode(logits.argmax(dim=-1))
+        assert tokens[1] == " Paris"
+
+    @torch.no_grad()
+    def test_token_list_in_invoker(self, vllm_gpt2, ET_prompt: str):
+        """Test token list input within an invoker."""
+        token_ids = vllm_gpt2.tokenizer.encode(ET_prompt)
+
+        with vllm_gpt2.trace(temperature=0.0, top_p=1) as tracer:
+            with tracer.invoke(token_ids):
+                logits = vllm_gpt2.logits.output.save()
+
+        next_token = vllm_gpt2.tokenizer.decode(logits.argmax(dim=-1))
+        assert next_token == " Paris"
+
+    @torch.no_grad()
+    def test_mixed_string_and_token_invokers(
+        self, vllm_gpt2, ET_prompt: str, MSG_prompt: str
+    ):
+        """Test mixing string and token list inputs across invokers."""
+        et_tokens = vllm_gpt2.tokenizer.encode(ET_prompt)
+
+        with vllm_gpt2.trace(temperature=0.0, top_p=1) as tracer:
+            with tracer.invoke(et_tokens):
+                et_logits = vllm_gpt2.logits.output.save()
+
+            with tracer.invoke(MSG_prompt):
+                msg_logits = vllm_gpt2.logits.output.save()
+
+        et_token = vllm_gpt2.tokenizer.decode(et_logits.argmax(dim=-1))
+        msg_token = vllm_gpt2.tokenizer.decode(msg_logits.argmax(dim=-1))
+        assert et_token == " Paris"
+        assert msg_token == " New"
diff --git a/tests/test_vllm_dispatch_bug.py b/tests/test_vllm_dispatch_bug.py
@@ -0,0 +1,32 @@
+"""Test for VLLM dispatch=False tracing bug."""
+import pytest
+import torch
+
+try:
+    from nnsight.modeling.vllm import VLLM
+except Exception as e:
+    pytest.skip(f"Skipping VLLM tests: \n{e}", allow_module_level=True)
+
+
+@pytest.fixture(scope="module")
+def vllm_gpt2_no_dispatch():
+    """VLLM model initialized without dispatch=True."""
+    return VLLM("gpt2", tensor_parallel_size=1, gpu_memory_utilization=0.1)
+
+
+@torch.no_grad()
+def test_trace_without_dispatch(vllm_gpt2_no_dispatch):
+    """Tracing should work even when dispatch=False at init time."""
+    model = vllm_gpt2_no_dispatch
+
+    assert not model.dispatched, "Model should not be dispatched initially"
+    assert model.vllm_entrypoint is None, "vllm_entrypoint should be None initially"
+
+    with model.trace("The Eiffel Tower is located in the city of", temperature=0.0, top_p=1):
+        logits = model.logits.output.save()
+
+    assert model.dispatched, "Model should be dispatched after trace"
+    assert model.vllm_entrypoint is not None, "vllm_entrypoint should exist after trace"
+
+    next_token = model.tokenizer.decode(logits.argmax(dim=-1))
+    assert next_token == " Paris"
