A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 13:30 UTC
[Misc] Make cached tokenizer pickle-compatible (#17048)
Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk>
Cyrus Leung committed
Apr 27, 2025 at 13:05 UTC
93a126fbc726ce2b79cec905bb750276d7ddcd76
5 files changed
+80
-56
benchmarks/benchmark_prefix_caching.py
+8
-6
@@ -63,14 +63,16 @@ class Request:
63
output_len: int
64
65
66
-def sample_tokens(tokenizer: PreTrainedTokenizerBase, length: int) -> str:
66
+def sample_tokens(tokenizer: PreTrainedTokenizerBase,
67
+ length: int) -> list[int]:
68
vocab = tokenizer.get_vocab()
69
+ all_special_ids = set(tokenizer.all_special_ids)
70
+
71
# Remove the special tokens.
69
- vocab = {
70
- k: v
71
- for k, v in vocab.items() if k not in tokenizer.all_special_ids
72
- }
73
- return random.choices(list(vocab.values()), k=length)
72
+ return random.choices(
73
+ [v for k, v in vocab.items() if k not in all_special_ids],
74
+ k=length,
75
+ )
76
77
78
def sample_requests_from_dataset(
tests/tokenization/test_cached_tokenizer.py
+31
-12
@@ -1,24 +1,43 @@
1
# SPDX-License-Identifier: Apache-2.0
2
-
2
+import pickle
3
from copy import deepcopy
4
5
+import pytest
6
from transformers import AutoTokenizer
7
7
-from vllm.transformers_utils.tokenizer import get_cached_tokenizer
8
+from vllm.transformers_utils.tokenizer import (AnyTokenizer,
9
+ get_cached_tokenizer)
10
11
10
-def test_cached_tokenizer():
11
- reference_tokenizer = AutoTokenizer.from_pretrained("gpt2")
12
+@pytest.mark.parametrize("model_id", ["gpt2", "THUDM/chatglm3-6b"])
13
+def test_cached_tokenizer(model_id: str):
14
+ reference_tokenizer = AutoTokenizer.from_pretrained(model_id,
15
+ trust_remote_code=True)
16
reference_tokenizer.add_special_tokens({"cls_token": "<CLS>"})
17
reference_tokenizer.add_special_tokens(
18
{"additional_special_tokens": ["<SEP>"]})
19
+
20
cached_tokenizer = get_cached_tokenizer(deepcopy(reference_tokenizer))
21
+ _check_consistency(cached_tokenizer, reference_tokenizer)
22
+
23
+ pickled_tokenizer = pickle.dumps(cached_tokenizer)
24
+ unpickled_tokenizer = pickle.loads(pickled_tokenizer)
25
+ _check_consistency(unpickled_tokenizer, reference_tokenizer)
26
+
27
+
28
+def _check_consistency(target: AnyTokenizer, expected: AnyTokenizer):
29
+ assert isinstance(target, type(expected))
30
+
31
+ # Cached attributes
32
+ assert target.all_special_ids == expected.all_special_ids
33
+ assert target.all_special_tokens == expected.all_special_tokens
34
+ assert (target.all_special_tokens_extended ==
35
+ expected.all_special_tokens_extended)
36
+ assert target.get_vocab() == expected.get_vocab()
37
+ assert len(target) == len(expected)
38
+
39
+ # Other attributes
40
+ assert getattr(target, "padding_side",
41
+ None) == getattr(expected, "padding_side", None)
42
17
- assert reference_tokenizer.encode("prompt") == cached_tokenizer.encode(
18
- "prompt")
19
- assert set(reference_tokenizer.all_special_ids) == set(
20
- cached_tokenizer.all_special_ids)
21
- assert set(reference_tokenizer.all_special_tokens) == set(
22
- cached_tokenizer.all_special_tokens)
23
- assert set(reference_tokenizer.all_special_tokens_extended) == set(
24
- cached_tokenizer.all_special_tokens_extended)
43
+ assert target.encode("prompt") == expected.encode("prompt")
vllm/transformers_utils/tokenizer.py
+19
-16
@@ -1,6 +1,7 @@
1
# SPDX-License-Identifier: Apache-2.0
2
3
import contextlib
4
+import copy
5
import os
6
import warnings
7
from functools import lru_cache
@@ -70,18 +71,17 @@ def encode_tokens(
71
72
73
def get_cached_tokenizer(tokenizer: AnyTokenizer) -> AnyTokenizer:
73
- """Get tokenizer with cached properties.
74
-
75
- This will patch the tokenizer object in place.
76
-
74
+ """
75
By default, transformers will recompute multiple tokenizer properties
78
- each time they are called, leading to a significant slowdown. This
79
- function caches these properties for faster access."""
76
+ each time they are called, leading to a significant slowdown.
77
+ This proxy caches these properties for faster access.
78
+ """
79
+ cached_tokenizer = copy.copy(tokenizer)
80
81
- tokenizer_all_special_ids = set(tokenizer.all_special_ids)
81
+ tokenizer_all_special_ids = tokenizer.all_special_ids
82
+ tokenizer_all_special_tokens = tokenizer.all_special_tokens
83
tokenizer_all_special_tokens_extended = (
84
tokenizer.all_special_tokens_extended)
84
- tokenizer_all_special_tokens = set(tokenizer.all_special_tokens)
85
tokenizer_vocab = tokenizer.get_vocab()
86
tokenizer_len = len(tokenizer)
87
@@ -97,31 +97,34 @@ def get_cached_tokenizer(tokenizer: AnyTokenizer) -> AnyTokenizer:
97
class CachedTokenizer(tokenizer.__class__): # type: ignore
98
99
@property
100
- def all_special_ids(self):
100
+ def all_special_ids(self) -> list[int]:
101
return tokenizer_all_special_ids
102
103
@property
104
- def all_special_tokens(self):
104
+ def all_special_tokens(self) -> list[str]:
105
return tokenizer_all_special_tokens
106
107
@property
108
- def all_special_tokens_extended(self):
108
+ def all_special_tokens_extended(self) -> list[str]:
109
return tokenizer_all_special_tokens_extended
110
111
@property
112
- def max_token_id(self):
112
+ def max_token_id(self) -> int:
113
return max_token_id
114
115
- def get_vocab(self):
115
+ def get_vocab(self) -> dict[str, int]:
116
return tokenizer_vocab
117
118
- def __len__(self):
118
+ def __len__(self) -> int:
119
return tokenizer_len
120
121
+ def __reduce__(self):
122
+ return get_cached_tokenizer, (tokenizer, )
123
+
124
CachedTokenizer.__name__ = f"Cached{tokenizer.__class__.__name__}"
125
123
- tokenizer.__class__ = CachedTokenizer
124
- return tokenizer
126
+ cached_tokenizer.__class__ = CachedTokenizer
127
+ return cached_tokenizer
128
129
130
def patch_padding_side(tokenizer: PreTrainedTokenizer) -> None:
vllm/transformers_utils/tokenizer_base.py
+17
-17
@@ -2,7 +2,7 @@
2
3
import importlib
4
from abc import ABC, abstractmethod
5
-from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union
5
+from typing import TYPE_CHECKING, Any, Optional, Union
6
7
if TYPE_CHECKING:
8
from vllm.entrypoints.chat_utils import ChatCompletionMessageParam
@@ -12,17 +12,17 @@ class TokenizerBase(ABC):
12
13
@property
14
@abstractmethod
15
- def all_special_tokens_extended(self) -> List[str]:
15
+ def all_special_tokens_extended(self) -> list[str]:
16
raise NotImplementedError()
17
18
@property
19
@abstractmethod
20
- def all_special_tokens(self) -> List[str]:
20
+ def all_special_tokens(self) -> list[str]:
21
raise NotImplementedError()
22
23
@property
24
@abstractmethod
25
- def all_special_ids(self) -> List[int]:
25
+ def all_special_ids(self) -> list[int]:
26
raise NotImplementedError()
27
28
@property
@@ -66,7 +66,7 @@ class TokenizerBase(ABC):
66
@abstractmethod
67
def __call__(
68
self,
69
- text: Union[str, List[str], List[int]],
69
+ text: Union[str, list[str], list[int]],
70
text_pair: Optional[str] = None,
71
add_special_tokens: bool = False,
72
truncation: bool = False,
@@ -75,11 +75,11 @@ class TokenizerBase(ABC):
75
raise NotImplementedError()
76
77
@abstractmethod
78
- def get_vocab(self) -> Dict[str, int]:
78
+ def get_vocab(self) -> dict[str, int]:
79
raise NotImplementedError()
80
81
@abstractmethod
82
- def get_added_vocab(self) -> Dict[str, int]:
82
+ def get_added_vocab(self) -> dict[str, int]:
83
raise NotImplementedError()
84
85
@abstractmethod
@@ -88,44 +88,44 @@ class TokenizerBase(ABC):
88
text: str,
89
truncation: bool = False,
90
max_length: Optional[int] = None,
91
- ) -> List[int]:
91
+ ) -> list[int]:
92
raise NotImplementedError()
93
94
@abstractmethod
95
def encode(self,
96
text: str,
97
- add_special_tokens: Optional[bool] = None) -> List[int]:
97
+ add_special_tokens: Optional[bool] = None) -> list[int]:
98
raise NotImplementedError()
99
100
@abstractmethod
101
def apply_chat_template(self,
102
- messages: List["ChatCompletionMessageParam"],
103
- tools: Optional[List[Dict[str, Any]]] = None,
104
- **kwargs) -> List[int]:
102
+ messages: list["ChatCompletionMessageParam"],
103
+ tools: Optional[list[dict[str, Any]]] = None,
104
+ **kwargs) -> list[int]:
105
raise NotImplementedError()
106
107
@abstractmethod
108
- def convert_tokens_to_string(self, tokens: List[str]) -> str:
108
+ def convert_tokens_to_string(self, tokens: list[str]) -> str:
109
raise NotImplementedError()
110
111
@abstractmethod
112
def decode(self,
113
- ids: Union[List[int], int],
113
+ ids: Union[list[int], int],
114
skip_special_tokens: bool = True) -> str:
115
raise NotImplementedError()
116
117
@abstractmethod
118
def convert_ids_to_tokens(
119
self,
120
- ids: List[int],
120
+ ids: list[int],
121
skip_special_tokens: bool = True,
122
- ) -> List[str]:
122
+ ) -> list[str]:
123
raise NotImplementedError()
124
125
126
class TokenizerRegistry:
127
# Tokenizer name -> (tokenizer module, tokenizer class)
128
- REGISTRY: Dict[str, Tuple[str, str]] = {}
128
+ REGISTRY: dict[str, tuple[str, str]] = {}
129
130
@staticmethod
131
def register(name: str, module: str, class_name: str) -> None:
vllm/transformers_utils/tokenizers/mistral.py
+5
-5
@@ -257,7 +257,7 @@ class MistralTokenizer(TokenizerBase):
257
# the following attributes are set to fit vLLM's design and are used
258
# by the guided structured output backends.
259
@property
260
- def all_special_tokens_extended(self) -> List[str]:
260
+ def all_special_tokens_extended(self) -> list[str]:
261
from mistral_common.tokens.tokenizers.base import SpecialTokens
262
263
# tekken defines its own extended special tokens list
@@ -271,11 +271,11 @@ class MistralTokenizer(TokenizerBase):
271
]
272
273
@property
274
- def all_special_tokens(self) -> List[str]:
274
+ def all_special_tokens(self) -> list[str]:
275
return self.all_special_tokens_extended
276
277
@property
278
- def all_special_ids(self) -> List[int]:
278
+ def all_special_ids(self) -> list[int]:
279
return [
280
self.all_special_tokens.index(t) for t in self.all_special_tokens
281
]
@@ -335,12 +335,12 @@ class MistralTokenizer(TokenizerBase):
335
input_ids = self.encode_one(text, truncation, max_length)
336
return Encoding(input_ids=input_ids)
337
338
- def get_vocab(self) -> Dict[str, int]:
338
+ def get_vocab(self) -> dict[str, int]:
339
# NB: the dictionary form of the vocabulary collapses token ids that map
340
# to the same string but have different bytes
341
return self._vocab_dict
342
343
- def get_added_vocab(self) -> Dict[str, int]:
343
+ def get_added_vocab(self) -> dict[str, int]:
344
# Mistral tokenizers have no added vocabulary
345
return {}
346