registry.py 12.9 KB
Newer Older
1
import functools
2
from collections import UserDict
3
from typing import (TYPE_CHECKING, Any, Dict, Mapping, Optional, Protocol,
4
                    Sequence, Type, TypeVar)
5

6
7
8
import torch.nn as nn

from vllm.inputs import InputProcessingContext
9
from vllm.logger import init_logger
10
from vllm.transformers_utils.tokenizer import AnyTokenizer
11
from vllm.utils import ClassRegistry
12

13
from .audio import AudioPlugin
14
from .base import MultiModalInputMapper, MultiModalPlugin, MultiModalTokensCalc
15
from .image import ImagePlugin
16
from .inputs import MultiModalDataDict, MultiModalKwargs, NestedTensors
17
from .processing import BaseMultiModalProcessor, ProcessingCache
18
from .video import VideoPlugin
19

20
21
22
if TYPE_CHECKING:
    from vllm.config import ModelConfig

23
24
logger = init_logger(__name__)

25
26
27
# TODO: Tune the MM cache size
MM_CACHE_SIZE = 256

28
29
30
N = TypeVar("N", bound=Type[nn.Module])


31
32
33
34
35
36
37
38
39
40
class MultiModalProcessorFactory(Protocol):
    """Constructs a :class:`MultiModalProcessor` instance from the context."""

    def __call__(
        self,
        ctx: InputProcessingContext,
        *,
        cache: Optional[ProcessingCache] = None,
    ) -> BaseMultiModalProcessor:
        ...
41

42

43
class _MultiModalLimits(UserDict["ModelConfig", Dict[str, int]]):
44
45
46
47
48
    """
    Wraps `_limits_by_model` for a more informative error message
    when attempting to access a model that does not exist.
    """

49
    def __getitem__(self, key: "ModelConfig") -> Dict[str, int]:
50
51
52
53
54
55
56
57
        try:
            return super().__getitem__(key)
        except KeyError as exc:
            msg = (f"Cannot find `mm_limits` for model={key.model}. Did you "
                   "forget to call `init_mm_limits_per_prompt`?")
            raise KeyError(msg) from exc


58
59
class MultiModalRegistry:
    """
60
61
    A registry that dispatches data processing to the
    :class:`~vllm.multimodal.MultiModalPlugin` for each modality.
62
63
    """

64
    DEFAULT_PLUGINS = (ImagePlugin(), AudioPlugin(), VideoPlugin())
65

66
    def __init__(
67
68
69
70
            self,
            *,
            plugins: Sequence[MultiModalPlugin] = DEFAULT_PLUGINS) -> None:
        self._plugins = {p.get_data_key(): p for p in plugins}
71

72
73
        self._processor_factories = ClassRegistry[nn.Module,
                                                  MultiModalProcessorFactory]()
74

75
76
77
78
79
        # This is used for non-multimodal models
        self._disabled_limits_per_plugin = {k: 0 for k in self._plugins}

        self._limits_by_model = _MultiModalLimits()

80
81
        self._processing_cache = ProcessingCache(MM_CACHE_SIZE)

82
    def register_plugin(self, plugin: MultiModalPlugin) -> None:
83
84
85
86
        """
        Register a multi-modal plugin so it can be recognized by vLLM.

        See also:
87
            :ref:`adding-multimodal-plugin`
88
        """
89
        data_type_key = plugin.get_data_key()
90

91
        if data_type_key in self._plugins:
92
93
            logger.warning(
                "A plugin is already registered for data type %s, "
94
                "and will be overwritten by the new plugin %s.", data_type_key,
95
96
                plugin)

97
        self._plugins[data_type_key] = plugin
98

99
100
101
102
    def _get_plugin(self, data_type_key: str):
        plugin = self._plugins.get(data_type_key)
        if plugin is not None:
            return plugin
103

104
        msg = f"Unknown multi-modal data type: {data_type_key}"
105
106
        raise NotImplementedError(msg)

107
    def register_input_mapper(
108
        self,
109
        data_type_key: str,
110
        mapper: Optional[MultiModalInputMapper] = None,
111
    ):
112
        """
113
        Register an input mapper for a specific modality to a model class.
114

115
        See :meth:`MultiModalPlugin.register_input_mapper` for more details.
116
        """
117
        return self._get_plugin(data_type_key).register_input_mapper(mapper)
118

119
    def register_image_input_mapper(
120
        self,
121
        mapper: Optional[MultiModalInputMapper] = None,
122
    ):
123
        """
124
        Register an input mapper for image data to a model class.
125

126
        See :meth:`MultiModalPlugin.register_input_mapper` for more details.
127
        """
128
        return self.register_input_mapper("image", mapper)
129

130
131
    def map_input(
        self,
132
        model_config: "ModelConfig",
133
134
        data: MultiModalDataDict,
        mm_processor_kwargs: Optional[Dict[str, Any]] = None,
135
    ) -> MultiModalKwargs:
136
        """
137
        Apply an input mapper to the data passed to the model.
138
139
140
141
142

        The data belonging to each modality is passed to the corresponding
        plugin which in turn converts the data into into keyword arguments
        via the input mapper registered for that model.

143
        See :meth:`MultiModalPlugin.map_input` for more details.
144
145
146

        Note:
            This should be called after :meth:`init_mm_limits_per_prompt`.
147
        """
148
        merged_dict: Dict[str, NestedTensors] = {}
149
150

        for data_key, data_value in data.items():
151
            plugin = self._get_plugin(data_key)
152

153
154
155
156
157
158
159
160
            num_items = len(data_value) if isinstance(data_value, list) else 1
            max_items = self._limits_by_model[model_config][data_key]
            if num_items > max_items:
                raise ValueError(
                    f"You set {data_key}={max_items} (or defaulted to 1) in "
                    f"`--limit-mm-per-prompt`, but found {num_items} items "
                    "in the same prompt.")

161
162
            input_dict = plugin.map_input(model_config, data_value,
                                          mm_processor_kwargs)
163
164
165
166
167
168
169
170
            for input_key, input_tensor in input_dict.items():
                if input_key in merged_dict:
                    raise ValueError(f"The input mappers (keys={set(data)}) "
                                     f"resulted in a conflicting keyword "
                                     f"argument to `forward()`: {input_key}")

                merged_dict[input_key] = input_tensor

171
        return MultiModalKwargs(merged_dict)
172

173
    def create_input_mapper(self, model_config: "ModelConfig"):
174
        """
175
        Create an input mapper (see :meth:`map_input`) for a specific model.
176
        """
177
178
179
180
181
182
183
184
185
        # NOTE - we currently make the assumption that if a model has multiple
        # supported modalities, they take the same kwargs. For the default,
        # this could be an issue in the future if it falls back to two HF
        # resources and we can't inspect the signature easily since it's
        # getting initialized through the autoclass.
        #
        # If this is a problem in the future, we should revisit it, but since
        # it potentially introduces a lot of complexity for a currently
        # uncommon case, we do not for simplicity of both use & implementation
186
        return functools.partial(self.map_input, model_config)
187

188
189
190
191
192
    def register_max_multimodal_tokens(
        self,
        data_type_key: str,
        max_mm_tokens: Optional[MultiModalTokensCalc] = None,
    ):
193
        """
194
195
196
        Register the maximum number of tokens, corresponding to a single
        instance of multimodal data belonging to a specific modality, that are
        passed to the language model for a model class.
197
198
199
200
201
202
203
204
205
        """
        return self._get_plugin(data_type_key) \
            .register_max_multimodal_tokens(max_mm_tokens)

    def register_max_image_tokens(
        self,
        max_mm_tokens: Optional[MultiModalTokensCalc] = None,
    ):
        """
206
207
        Register the maximum number of image tokens, corresponding to a single
        image, that are passed to the language model for a model class.
208
209
210
        """
        return self.register_max_multimodal_tokens("image", max_mm_tokens)

211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
    def get_max_tokens_per_item_by_modality(
        self,
        model_config: "ModelConfig",
    ) -> Mapping[str, int]:
        """
        Get the maximum number of tokens per data item from each modality
        for profiling the memory usage of a model.

        Note:
            This is currently directly used only in V1.
        """

        return {
            key: plugin.get_max_multimodal_tokens(model_config)
            for key, plugin in self._plugins.items()
        }

228
229
230
231
    def get_max_tokens_by_modality(
        self,
        model_config: "ModelConfig",
    ) -> Mapping[str, int]:
232
        """
233
        Get the maximum number of tokens from each modality
234
        for profiling the memory usage of a model.
235

236
        See :meth:`MultiModalPlugin.get_max_multimodal_tokens` for more details.
237
238
239
240
241
242

        Note:
            This should be called after :meth:`init_mm_limits_per_prompt`.
        """
        limits_per_plugin = self._limits_by_model[model_config]

243
        return {
244
245
246
            key: limits_per_plugin[key] * max_tokens_per_mm_item
            for key, max_tokens_per_mm_item in
            self.get_max_tokens_per_item_by_modality(model_config).items()
247
248
249
250
251
252
253
254
255
256
257
258
259
        }

    def get_max_multimodal_tokens(self, model_config: "ModelConfig") -> int:
        """
        Get the maximum number of multi-modal tokens
        for profiling the memory usage of a model.

        See :meth:`MultiModalPlugin.get_max_multimodal_tokens` for more details.

        Note:
            This should be called after :meth:`init_mm_limits_per_prompt`.
        """
        return sum(self.get_max_tokens_by_modality(model_config).values())
260
261
262

    def init_mm_limits_per_prompt(
        self,
263
        model_config: "ModelConfig",
264
265
266
267
268
269
270
271
272
273
    ) -> None:
        """
        Initialize the maximum number of multi-modal input instances for each
        modality that are allowed per prompt for a model class.
        """
        if model_config in self._limits_by_model:
            logger.warning(
                "`mm_limits` has already been set for model=%s, and will "
                "be overwritten by the new values.", model_config.model)

274
        multimodal_config = model_config.multimodal_config
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
        if multimodal_config is None:
            limits_per_plugin = self._disabled_limits_per_plugin
        else:
            config_limits_per_plugin = multimodal_config.limit_per_prompt

            extra_keys = config_limits_per_plugin.keys() - self._plugins.keys()
            if extra_keys:
                logger.warning(
                    "Detected extra keys in `--limit-mm-per-prompt` which "
                    "are not registered as multi-modal plugins: %s. "
                    "They will be ignored.", extra_keys)

            # NOTE: Currently the default is set to 1 for each plugin
            # TODO: Automatically determine the limits based on budget
            # once more models support multi-image inputs
            limits_per_plugin = {
                key: config_limits_per_plugin.get(key, 1)
                for key in self._plugins
            }

        self._limits_by_model[model_config] = limits_per_plugin

    def get_mm_limits_per_prompt(
        self,
299
        model_config: "ModelConfig",
300
301
302
303
304
305
306
    ) -> Mapping[str, int]:
        """
        Get the maximum number of multi-modal input instances for each modality
        that are allowed per prompt for a model class.

        Note:
            This should be called after :meth:`init_mm_limits_per_prompt`.
307
        """
308
        return self._limits_by_model[model_config]
309
310
311
312
313
314

    def register_processor(
        self,
        factory: MultiModalProcessorFactory,
    ):
        """
315
316
        Register a multi-modal processor to a model class. The processor
        is constructed lazily, hence a factory method should be passed.
317
318
319
320
321

        When the model receives multi-modal data, the provided function is
        invoked to transform the data into a dictionary of model inputs.

        See also:
322
323
            - :ref:`input-processing-pipeline`
            - :ref:`enabling-multimodal-inputs`
324
325
326
        """

        def wrapper(model_cls: N) -> N:
327
            if self._processor_factories.contains(model_cls, strict=True):
328
                logger.warning(
329
                    "Model class %s already has a multi-modal processor "
330
331
332
333
334
335
336
337
338
                    "registered to %s. It is overwritten by the new one.",
                    model_cls, self)

            self._processor_factories[model_cls] = factory

            return model_cls

        return wrapper

339
    def _get_model_cls(self, model_config: "ModelConfig"):
340
341
342
343
        # Avoid circular import
        from vllm.model_executor.model_loader import get_model_architecture

        model_cls, _ = get_model_architecture(model_config)
344
345
346
347
348
349
350
        return model_cls

    def has_processor(self, model_config: "ModelConfig") -> bool:
        """
        Test whether a multi-modal processor is defined for a specific model.
        """
        return self._get_model_cls(model_config) in self._processor_factories
351
352
353
354
355

    def create_processor(
        self,
        model_config: "ModelConfig",
        tokenizer: AnyTokenizer,
356
    ) -> BaseMultiModalProcessor:
357
358
359
        """
        Create a multi-modal processor for a specific model and tokenizer.
        """
360
        model_cls = self._get_model_cls(model_config)
361
362
363
        processor_factory = self._processor_factories[model_cls]

        ctx = InputProcessingContext(model_config, tokenizer)
364
365
366
367
        cache = (None if model_config.disable_mm_preprocessor_cache else
                 self._processing_cache)

        return processor_factory(ctx, cache=cache)