registry.py 8.91 KB
Newer Older
1
import functools
2
from collections import UserDict
3
from typing import Any, Dict, Mapping, Optional, Sequence
4

5
from vllm.config import ModelConfig
6
7
from vllm.logger import init_logger

8
from .audio import AudioPlugin
9
from .base import (MultiModalDataDict, MultiModalInputMapper, MultiModalInputs,
10
                   MultiModalPlugin, MultiModalTokensCalc, NestedTensors)
11
from .image import ImagePlugin
12
from .video import VideoPlugin
13
14
15
16

logger = init_logger(__name__)


17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
class _MultiModalLimits(UserDict):
    """
    Wraps `_limits_by_model` for a more informative error message
    when attempting to access a model that does not exist.
    """

    def __getitem__(self, key: ModelConfig) -> Dict[str, int]:
        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


32
33
class MultiModalRegistry:
    """
34
35
    A registry that dispatches data processing to the
    :class:`~vllm.multimodal.MultiModalPlugin` for each modality.
36
37
    """

38
    DEFAULT_PLUGINS = (ImagePlugin(), AudioPlugin(), VideoPlugin())
39

40
    def __init__(
41
42
43
44
            self,
            *,
            plugins: Sequence[MultiModalPlugin] = DEFAULT_PLUGINS) -> None:
        self._plugins = {p.get_data_key(): p for p in plugins}
45

46
47
48
49
50
        # This is used for non-multimodal models
        self._disabled_limits_per_plugin = {k: 0 for k in self._plugins}

        self._limits_by_model = _MultiModalLimits()

51
    def register_plugin(self, plugin: MultiModalPlugin) -> None:
52
53
54
55
56
57
        """
        Register a multi-modal plugin so it can be recognized by vLLM.

        See also:
            :ref:`adding_multimodal_plugin`
        """
58
        data_type_key = plugin.get_data_key()
59

60
        if data_type_key in self._plugins:
61
62
            logger.warning(
                "A plugin is already registered for data type %s, "
63
                "and will be overwritten by the new plugin %s.", data_type_key,
64
65
                plugin)

66
        self._plugins[data_type_key] = plugin
67

68
69
70
71
    def _get_plugin(self, data_type_key: str):
        plugin = self._plugins.get(data_type_key)
        if plugin is not None:
            return plugin
72

73
        msg = f"Unknown multi-modal data type: {data_type_key}"
74
75
        raise NotImplementedError(msg)

76
    def register_input_mapper(
77
        self,
78
        data_type_key: str,
79
        mapper: Optional[MultiModalInputMapper] = None,
80
    ):
81
        """
82
        Register an input mapper for a specific modality to a model class.
83

84
        See :meth:`MultiModalPlugin.register_input_mapper` for more details.
85
        """
86
        return self._get_plugin(data_type_key).register_input_mapper(mapper)
87

88
    def register_image_input_mapper(
89
        self,
90
        mapper: Optional[MultiModalInputMapper] = None,
91
    ):
92
        """
93
        Register an input mapper for image data to a model class.
94

95
        See :meth:`MultiModalPlugin.register_input_mapper` for more details.
96
        """
97
        return self.register_input_mapper("image", mapper)
98

99
100
101
102
103
104
    def map_input(
        self,
        model_config: ModelConfig,
        data: MultiModalDataDict,
        mm_processor_kwargs: Optional[Dict[str, Any]] = None,
    ) -> MultiModalInputs:
105
        """
106
        Apply an input mapper to the data passed to the model.
107
108
109
110
111

        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.

112
        See :meth:`MultiModalPlugin.map_input` for more details.
113
114
115

        Note:
            This should be called after :meth:`init_mm_limits_per_prompt`.
116
        """
117
        merged_dict: Dict[str, NestedTensors] = {}
118
119

        for data_key, data_value in data.items():
120
            plugin = self._get_plugin(data_key)
121

122
123
124
125
126
127
128
129
            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.")

130
131
            input_dict = plugin.map_input(model_config, data_value,
                                          mm_processor_kwargs)
132
133
134
135
136
137
138
139
140
            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

        return MultiModalInputs(merged_dict)
141

142
    def create_input_mapper(self, model_config: ModelConfig):
143
        """
144
        Create an input mapper (see :meth:`map_input`) for a specific model.
145
        """
146
147
148
149
150
151
152
153
154
        # 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
155
        return functools.partial(self.map_input, model_config)
156

157
158
159
160
161
    def register_max_multimodal_tokens(
        self,
        data_type_key: str,
        max_mm_tokens: Optional[MultiModalTokensCalc] = None,
    ):
162
        """
163
164
165
        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.
166
167
168
169
170
171
172
173
174
        """
        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,
    ):
        """
175
176
        Register the maximum number of image tokens, corresponding to a single
        image, that are passed to the language model for a model class.
177
178
179
180
181
182
183
        """
        return self.register_max_multimodal_tokens("image", max_mm_tokens)

    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.
184

185
        See :meth:`MultiModalPlugin.get_max_multimodal_tokens` for more details.
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208

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

        return sum((limits_per_plugin[key] *
                    plugin.get_max_multimodal_tokens(model_config))
                   for key, plugin in self._plugins.items())

    def init_mm_limits_per_prompt(
        self,
        model_config: ModelConfig,
    ) -> 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)

209
        multimodal_config = model_config.multimodal_config
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
        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,
        model_config: ModelConfig,
    ) -> 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`.
242
        """
243
        return self._limits_by_model[model_config]