registry.py 11.9 KB
Newer Older
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
from collections.abc import Mapping
4
from dataclasses import dataclass
5
from typing import TYPE_CHECKING, Generic, Protocol, TypeVar
6

7
8
import torch.nn as nn

9
from vllm.config.multimodal import BaseDummyOptions
10
from vllm.logger import init_logger
11
from vllm.transformers_utils.tokenizer import AnyTokenizer, cached_tokenizer_from_config
12
from vllm.utils.collection_utils import ClassRegistry
13

14
from .cache import BaseMultiModalProcessorCache
15
16
17
18
19
20
21
22
23
24
25
from .processing import (
    BaseMultiModalProcessor,
    BaseProcessingInfo,
    InputProcessingContext,
)
from .profiling import (
    BaseDummyInputsBuilder,
    DummyDecoderData,
    DummyEncoderData,
    MultiModalProfiler,
)
26

27
28
29
if TYPE_CHECKING:
    from vllm.config import ModelConfig

30
31
logger = init_logger(__name__)

32
N = TypeVar("N", bound=type[nn.Module])
33
34
_I = TypeVar("_I", bound=BaseProcessingInfo)
_I_co = TypeVar("_I_co", bound=BaseProcessingInfo, covariant=True)
35
36


37
class ProcessingInfoFactory(Protocol[_I_co]):
38
39
40
41
42
    """
    Constructs a
    [`BaseMultiModalProcessor`][vllm.multimodal.processing.BaseMultiModalProcessor]
    instance from the context.
    """
43
44
45
46

    def __call__(
        self,
        ctx: InputProcessingContext,
47
    ) -> _I_co: ...
48
49


50
class DummyInputsBuilderFactory(Protocol[_I]):  # type: ignore[misc]
51
    """
52
53
54
    Constructs a
    [`BaseDummyInputsBuilder`][vllm.multimodal.profiling.BaseDummyInputsBuilder]
    instance from the context.
55
56
    """

57
    def __call__(self, info: _I) -> BaseDummyInputsBuilder[_I]: ...
58
59


60
class MultiModalProcessorFactory(Protocol[_I]):  # type: ignore[misc]
61
62
63
64
65
    """
    Constructs a
    [`BaseMultiModalProcessor`][vllm.multimodal.processing.BaseMultiModalProcessor]
    instance from the context.
    """
66
67
68
69
70

    def __call__(
        self,
        info: _I,
        dummy_inputs: BaseDummyInputsBuilder[_I],
71
        *,
72
        cache: BaseMultiModalProcessorCache | None = None,
73
    ) -> BaseMultiModalProcessor[_I]: ...
74

75

76
77
78
79
80
81
82
83
84
85
@dataclass(frozen=True)
class _ProcessorFactories(Generic[_I]):
    info: ProcessingInfoFactory[_I]
    processor: MultiModalProcessorFactory[_I]
    dummy_inputs: DummyInputsBuilderFactory[_I]

    def build_processor(
        self,
        ctx: InputProcessingContext,
        *,
86
        cache: BaseMultiModalProcessorCache | None = None,
87
88
89
90
91
92
    ):
        info = self.info(ctx)
        dummy_inputs_builder = self.dummy_inputs(info)
        return self.processor(info, dummy_inputs_builder, cache=cache)


93
94
class MultiModalRegistry:
    """
95
    A registry that dispatches data processing according to the model.
96
97
    """

98
    def __init__(self) -> None:
99
        self._processor_factories = ClassRegistry[nn.Module, _ProcessorFactories]()
100

101
102
103
    def _extract_mm_options(
        self,
        model_config: "ModelConfig",
104
    ) -> Mapping[str, BaseDummyOptions] | None:
105
106
107
108
109
110
111
112
113
114
115
116
        """
        Extract multimodal dummy options from model config.

        Returns None if no configurable options are found, otherwise returns
        a mapping of modality names to their dummy options.
        """
        if not model_config.multimodal_config:
            return None

        mm_options = {
            m: opt
            for m in model_config.multimodal_config.limit_per_prompt
117
            if (opt := model_config.multimodal_config.get_dummy_options(m)) is not None
118
119
120
121
        }

        return mm_options if len(mm_options) > 0 else None

122
123
124
    def supports_multimodal_inputs(self, model_config: "ModelConfig") -> bool:
        """
        Checks if the model supports multimodal inputs.
125
126
        Returns True if the model is multimodal with any non-zero supported
        modalities, otherwise returns False, effectively running in
127
128
129
130
131
        text-only mode.
        """
        if not model_config.is_multimodal_model:
            return False

132
133
        info = self._create_processing_info(model_config, tokenizer=None)
        supported_modalities = info.get_supported_mm_limits()
134
135
136
137
138

        mm_config = model_config.get_multimodal_config()

        # Check if all supported modalities have limit == 0
        if all(
139
140
141
            mm_config.get_limit_per_prompt(modality) == 0
            for modality in supported_modalities
        ):
142
143
            logger.info_once(
                "All limits of multimodal modalities supported by the model "
144
145
                "are set to 0, running in text-only mode."
            )
146
147
148
149
            return False

        return True

150
151
152
    def get_max_tokens_per_item_by_modality(
        self,
        model_config: "ModelConfig",
153
        *,
154
        cache: BaseMultiModalProcessorCache | None = None,
155
        profiler_limits: Mapping[str, int] | None = None,
156
157
    ) -> Mapping[str, int]:
        """
158
        Get the maximum number of tokens per data item from each modality based
159
        on underlying model configuration.
160
        """
161
162
        if not model_config.is_multimodal_model:
            return {}
163

164
        processor = self.create_processor(model_config, cache=cache)
165
        profiler: MultiModalProfiler = MultiModalProfiler(processor)
166
167

        seq_len = model_config.max_model_len
168
169
170
        profiler_limits = (
            profiler.get_mm_limits() if profiler_limits is None else profiler_limits
        )
171

172
        return profiler.get_mm_max_contiguous_tokens(
173
            seq_len,
174
            {modality: 1 for modality, limit in profiler_limits.items() if limit > 0},
175
        )
176

177
178
    def get_mm_limits_per_prompt(
        self,
179
        model_config: "ModelConfig",
180
        *,
181
        cache: BaseMultiModalProcessorCache | None = None,
182
183
184
185
    ) -> Mapping[str, int]:
        """
        Get the maximum number of multi-modal input instances for each modality
        that are allowed per prompt for a model class.
186
        """
187
188
        if not model_config.is_multimodal_model:
            return {}
189

190
        processor = self.create_processor(model_config, cache=cache)
191
        profiler: MultiModalProfiler = MultiModalProfiler(processor)
192
        return profiler.get_mm_limits()
193
194
195

    def register_processor(
        self,
196
197
198
199
        processor: MultiModalProcessorFactory[_I],
        *,
        info: ProcessingInfoFactory[_I],
        dummy_inputs: DummyInputsBuilderFactory[_I],
200
201
    ):
        """
202
203
        Register a multi-modal processor to a model class. The processor
        is constructed lazily, hence a factory method should be passed.
204
205
206
207
208
209

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

        def wrapper(model_cls: N) -> N:
210
            if self._processor_factories.contains(model_cls, strict=True):
211
                logger.warning(
212
                    "Model class %s already has a multi-modal processor "
213
                    "registered to %s. It is overwritten by the new one.",
214
215
216
                    model_cls,
                    self,
                )
217

218
219
220
221
222
            self._processor_factories[model_cls] = _ProcessorFactories(
                info=info,
                dummy_inputs=dummy_inputs,
                processor=processor,
            )
223
224
225
226
227

            return model_cls

        return wrapper

228
    def _get_model_cls(self, model_config: "ModelConfig"):
229
230
231
232
        # Avoid circular import
        from vllm.model_executor.model_loader import get_model_architecture

        model_cls, _ = get_model_architecture(model_config)
233
234
        return model_cls

235
236
237
    def _create_processing_ctx(
        self,
        model_config: "ModelConfig",
238
        tokenizer: AnyTokenizer | None = None,
239
240
241
242
243
244
245
246
247
    ) -> InputProcessingContext:
        if tokenizer is None and not model_config.skip_tokenizer_init:
            tokenizer = cached_tokenizer_from_config(model_config)
        return InputProcessingContext(model_config, tokenizer)

    def _create_processing_info(
        self,
        model_config: "ModelConfig",
        *,
248
        tokenizer: AnyTokenizer | None = None,
249
250
251
252
253
254
    ) -> BaseProcessingInfo:
        model_cls = self._get_model_cls(model_config)
        factories = self._processor_factories[model_cls]
        ctx = self._create_processing_ctx(model_config, tokenizer)
        return factories.info(ctx)

255
256
257
    def create_processor(
        self,
        model_config: "ModelConfig",
258
        *,
259
260
        tokenizer: AnyTokenizer | None = None,
        cache: BaseMultiModalProcessorCache | None = None,
261
    ) -> BaseMultiModalProcessor[BaseProcessingInfo]:
262
263
264
        """
        Create a multi-modal processor for a specific model and tokenizer.
        """
265
266
267
        if not model_config.is_multimodal_model:
            raise ValueError(f"{model_config.model} is not a multimodal model")

268
        model_cls = self._get_model_cls(model_config)
269
        factories = self._processor_factories[model_cls]
270

271
        ctx = self._create_processing_ctx(model_config, tokenizer)
272

273
        return factories.build_processor(ctx, cache=cache)
274
275
276
277
278

    def get_decoder_dummy_data(
        self,
        model_config: "ModelConfig",
        seq_len: int,
279
        mm_counts: Mapping[str, int] | None = None,
280
        *,
281
        cache: BaseMultiModalProcessorCache | None = None,
282
283
284
285
    ) -> DummyDecoderData:
        """
        Create dummy data for profiling the memory usage of a model.

286
        The model is identified by `model_config`.
287
        """
288
        processor = self.create_processor(model_config, cache=cache)
289
290
291
292
293
294
295
        profiler: MultiModalProfiler = MultiModalProfiler(processor)

        # Extract configurable options from multimodal config.
        # Only include modalities that use advanced option types so legacy
        # count-only behavior remains unchanged.
        mm_options = self._extract_mm_options(model_config)

296
        dummy_data = profiler.get_decoder_dummy_data(seq_len, mm_counts, mm_options)
297
298
299
300
301
302

        # Having more tokens is over-conservative but otherwise fine
        token_ids = dummy_data.prompt_token_ids
        if len(token_ids) < seq_len:
            raise AssertionError(
                f"Expected at least {seq_len} dummy tokens for profiling, "
303
304
                f"but found {len(token_ids)} tokens instead."
            )
305
306
307
308
309
310
311

        return dummy_data

    def get_encoder_dummy_data(
        self,
        model_config: "ModelConfig",
        seq_len: int,
312
        mm_counts: Mapping[str, int] | None = None,
313
        *,
314
        cache: BaseMultiModalProcessorCache | None = None,
315
316
317
318
    ) -> DummyEncoderData:
        """
        Create dummy data for profiling the memory usage of a model.

319
        The model is identified by `model_config`.
320
        """
321
        processor = self.create_processor(model_config, cache=cache)
322
323
324
325
326
327
328
        profiler: MultiModalProfiler = MultiModalProfiler(processor)

        # Extract configurable options from multimodal config.
        # Only include modalities that use advanced option types so legacy
        # count-only behavior remains unchanged.
        mm_options = self._extract_mm_options(model_config)

329
        dummy_data = profiler.get_encoder_dummy_data(seq_len, mm_counts, mm_options)
330
331
332
333
334

        # Having more tokens is over-conservative but otherwise fine
        token_ids = dummy_data.prompt_token_ids
        if len(token_ids) < seq_len:
            logger.warning_once(
335
336
337
338
                "Expected at least %d dummy encoder tokens for profiling, but found %d tokens instead.",  # noqa: E501
                seq_len,
                len(token_ids),
            )
339
340

        return dummy_data
341
342
343
344
345
346
347

    def get_encdec_max_encoder_len(self, model_config: "ModelConfig") -> int:
        """
        Get the maximum length of the encoder input for encoder-decoder models.
        """
        if not model_config.is_encoder_decoder:
            return 0
348
        max_tokens = self.get_max_tokens_per_item_by_modality(model_config)
349
350
351
352
353
        if not max_tokens:
            # TODO - this function assumes encoder-decoder models are
            # multimodal. This will need to change when adding support for more
            # than whisper.
            return 0
354
355
        assert len(max_tokens) == 1, (
            "Encoder-decoder models are expected \
356
            to implement the multimodal interface with at most one modality."
357
        )
358
359
360

        first_modality = next(iter(max_tokens))
        return max_tokens[first_modality]