Merge pull request #589 from Butanium/feat/vllm-token-input-compat

feat(vllm): accept token lists and HuggingFace tokenizer results

Author: JadenFiotto-Kaufman

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 = []
 

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 (
@@ -118,7 +119,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 +153,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 +200,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)
 

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"