registry.py 25.3 KB
Newer Older
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
5
6
"""
Whenever you add an architecture to this page, please also update
`tests/models/registry.py` with example HuggingFace models for it.
"""
7
import importlib
8
import os
9
import pickle
10
11
import subprocess
import sys
12
import tempfile
13
from abc import ABC, abstractmethod
14
from collections.abc import Set
15
16
from dataclasses import dataclass, field
from functools import lru_cache
17
from typing import Callable, Optional, TypeVar, Union
18

19
import cloudpickle
20
21
22
23
import torch.nn as nn

from vllm.logger import init_logger

24
25
26
27
from .interfaces import (has_inner_state, has_noops, is_attention_free,
                         is_hybrid, supports_cross_encoding,
                         supports_multimodal, supports_pp,
                         supports_transcription, supports_v0_only)
28
from .interfaces_base import is_text_generation_model
29
30
31

logger = init_logger(__name__)

32
# yapf: disable
33
34
_TEXT_GENERATION_MODELS = {
    # [Decoder-only]
35
36
37
    "AquilaModel": ("llama", "LlamaForCausalLM"),
    "AquilaForCausalLM": ("llama", "LlamaForCausalLM"),  # AquilaChat2
    "ArcticForCausalLM": ("arctic", "ArcticForCausalLM"),
38
    "MiniMaxText01ForCausalLM": ("minimax_text_01", "MiniMaxText01ForCausalLM"),
39
40
41
42
    # baichuan-7b, upper case 'C' in the class name
    "BaiChuanForCausalLM": ("baichuan", "BaiChuanForCausalLM"),
    # baichuan-13b, lower case 'c' in the class name
    "BaichuanForCausalLM": ("baichuan", "BaichuanForCausalLM"),
Yu Chin Fabian Lim's avatar
Yu Chin Fabian Lim committed
43
    "BambaForCausalLM": ("bamba", "BambaForCausalLM"),
44
    "BloomForCausalLM": ("bloom", "BloomForCausalLM"),
45
    "ChatGLMModel": ("chatglm", "ChatGLMForCausalLM"),
46
    "ChatGLMForConditionalGeneration": ("chatglm", "ChatGLMForCausalLM"),
47
    "CohereForCausalLM": ("commandr", "CohereForCausalLM"),
48
    "Cohere2ForCausalLM": ("commandr", "CohereForCausalLM"),
49
    "DbrxForCausalLM": ("dbrx", "DbrxForCausalLM"),
50
    "DeciLMForCausalLM": ("nemotron_nas", "DeciLMForCausalLM"),
51
52
    "DeepseekForCausalLM": ("deepseek", "DeepseekForCausalLM"),
    "DeepseekV2ForCausalLM": ("deepseek_v2", "DeepseekV2ForCausalLM"),
53
    "DeepseekV3ForCausalLM": ("deepseek_v2", "DeepseekV3ForCausalLM"),
54
55
    "ExaoneForCausalLM": ("exaone", "ExaoneForCausalLM"),
    "FalconForCausalLM": ("falcon", "FalconForCausalLM"),
56
    "Fairseq2LlamaForCausalLM": ("fairseq2_llama", "Fairseq2LlamaForCausalLM"),
57
58
    "GemmaForCausalLM": ("gemma", "GemmaForCausalLM"),
    "Gemma2ForCausalLM": ("gemma2", "Gemma2ForCausalLM"),
59
    "Gemma3ForCausalLM": ("gemma3", "Gemma3ForCausalLM"),
60
    "GlmForCausalLM": ("glm", "GlmForCausalLM"),
Yuxuan Zhang's avatar
Yuxuan Zhang committed
61
    "Glm4ForCausalLM": ("glm4", "Glm4ForCausalLM"),
62
63
64
65
66
67
    "GPT2LMHeadModel": ("gpt2", "GPT2LMHeadModel"),
    "GPTBigCodeForCausalLM": ("gpt_bigcode", "GPTBigCodeForCausalLM"),
    "GPTJForCausalLM": ("gpt_j", "GPTJForCausalLM"),
    "GPTNeoXForCausalLM": ("gpt_neox", "GPTNeoXForCausalLM"),
    "GraniteForCausalLM": ("granite", "GraniteForCausalLM"),
    "GraniteMoeForCausalLM": ("granitemoe", "GraniteMoeForCausalLM"),
68
    "GraniteMoeHybridForCausalLM": ("granitemoehybrid", "GraniteMoeHybridForCausalLM"),   # noqa: E501
69
    "GraniteMoeSharedForCausalLM": ("granitemoeshared", "GraniteMoeSharedForCausalLM"),   # noqa: E501
70
    "GritLM": ("gritlm", "GritLM"),
Michael Goin's avatar
Michael Goin committed
71
    "Grok1ModelForCausalLM": ("grok1", "Grok1ForCausalLM"),
72
73
    "InternLMForCausalLM": ("llama", "LlamaForCausalLM"),
    "InternLM2ForCausalLM": ("internlm2", "InternLM2ForCausalLM"),
74
    "InternLM2VEForCausalLM": ("internlm2_ve", "InternLM2VEForCausalLM"),
75
    "InternLM3ForCausalLM": ("llama", "LlamaForCausalLM"),
76
77
78
79
80
    "JAISLMHeadModel": ("jais", "JAISLMHeadModel"),
    "JambaForCausalLM": ("jamba", "JambaForCausalLM"),
    "LlamaForCausalLM": ("llama", "LlamaForCausalLM"),
    # For decapoda-research/llama-*
    "LLaMAForCausalLM": ("llama", "LlamaForCausalLM"),
81
    "MambaForCausalLM": ("mamba", "MambaForCausalLM"),
82
    "FalconMambaForCausalLM": ("mamba", "MambaForCausalLM"),
Dhia Eddine Rhaiem's avatar
Dhia Eddine Rhaiem committed
83
    "FalconH1ForCausalLM":("falcon_h1", "FalconH1ForCausalLM"),
84
    "Mamba2ForCausalLM": ("mamba2", "Mamba2ForCausalLM"),
85
86
    "MiniCPMForCausalLM": ("minicpm", "MiniCPMForCausalLM"),
    "MiniCPM3ForCausalLM": ("minicpm3", "MiniCPM3ForCausalLM"),
87
88
89
90
91
92
    "MistralForCausalLM": ("llama", "LlamaForCausalLM"),
    "MixtralForCausalLM": ("mixtral", "MixtralForCausalLM"),
    "QuantMixtralForCausalLM": ("mixtral_quant", "MixtralForCausalLM"),
    # transformers's mpt class has lower case
    "MptForCausalLM": ("mpt", "MPTForCausalLM"),
    "MPTForCausalLM": ("mpt", "MPTForCausalLM"),
93
    "MiMoForCausalLM": ("mimo", "MiMoForCausalLM"),
94
    "NemotronForCausalLM": ("nemotron", "NemotronForCausalLM"),
Luis Vega's avatar
Luis Vega committed
95
    "NemotronHForCausalLM": ("nemotron_h", "NemotronHForCausalLM"),
96
    "OlmoForCausalLM": ("olmo", "OlmoForCausalLM"),
97
    "Olmo2ForCausalLM": ("olmo2", "Olmo2ForCausalLM"),
98
99
100
101
102
103
104
105
    "OlmoeForCausalLM": ("olmoe", "OlmoeForCausalLM"),
    "OPTForCausalLM": ("opt", "OPTForCausalLM"),
    "OrionForCausalLM": ("orion", "OrionForCausalLM"),
    "PersimmonForCausalLM": ("persimmon", "PersimmonForCausalLM"),
    "PhiForCausalLM": ("phi", "PhiForCausalLM"),
    "Phi3ForCausalLM": ("phi3", "Phi3ForCausalLM"),
    "Phi3SmallForCausalLM": ("phi3_small", "Phi3SmallForCausalLM"),
    "PhiMoEForCausalLM": ("phimoe", "PhiMoEForCausalLM"),
Shinichi Hemmi's avatar
Shinichi Hemmi committed
106
    "Plamo2ForCausalLM": ("plamo2", "Plamo2ForCausalLM"),
107
    "QWenLMHeadModel": ("qwen", "QWenLMHeadModel"),
108
109
    "Qwen2ForCausalLM": ("qwen2", "Qwen2ForCausalLM"),
    "Qwen2MoeForCausalLM": ("qwen2_moe", "Qwen2MoeForCausalLM"),
110
111
    "Qwen3ForCausalLM": ("qwen3", "Qwen3ForCausalLM"),
    "Qwen3MoeForCausalLM": ("qwen3_moe", "Qwen3MoeForCausalLM"),
112
113
114
115
116
    "RWForCausalLM": ("falcon", "FalconForCausalLM"),
    "StableLMEpochForCausalLM": ("stablelm", "StablelmForCausalLM"),
    "StableLmForCausalLM": ("stablelm", "StablelmForCausalLM"),
    "Starcoder2ForCausalLM": ("starcoder2", "Starcoder2ForCausalLM"),
    "SolarForCausalLM": ("solar", "SolarForCausalLM"),
117
    "TeleChat2ForCausalLM": ("telechat2", "TeleChat2ForCausalLM"),
118
    "TeleFLMForCausalLM": ("teleflm", "TeleFLMForCausalLM"),
119
    "XverseForCausalLM": ("llama", "LlamaForCausalLM"),
120
    "Zamba2ForCausalLM": ("zamba2", "Zamba2ForCausalLM"),
121
122
123
    # [Encoder-decoder]
    "BartModel": ("bart", "BartForConditionalGeneration"),
    "BartForConditionalGeneration": ("bart", "BartForConditionalGeneration"),
124
125
126
}

_EMBEDDING_MODELS = {
127
    # [Text-only]
128
    "BertModel": ("bert", "BertEmbeddingModel"),
129
    "DeciLMForCausalLM": ("nemotron_nas", "DeciLMForCausalLM"),
130
    "Gemma2Model": ("gemma2", "Gemma2ForCausalLM"),
131
    "GlmForCausalLM": ("glm", "GlmForCausalLM"),
132
    "GritLM": ("gritlm", "GritLM"),
133
134
    "GteModel": ("bert_with_rope", "SnowflakeGteNewModel"),
    "GteNewModel": ("bert_with_rope", "GteNewModel"),
135
    "InternLM2ForRewardModel": ("internlm2", "InternLM2ForRewardModel"),
136
    "JambaForSequenceClassification": ("jamba", "JambaForSequenceClassification"),  # noqa: E501
137
    "LlamaModel": ("llama", "LlamaForCausalLM"),
138
139
140
141
142
    **{
        # Multiple models share the same architecture, so we include them all
        k: (mod, arch) for k, (mod, arch) in _TEXT_GENERATION_MODELS.items()
        if arch == "LlamaForCausalLM"
    },
143
    "MistralModel": ("llama", "LlamaForCausalLM"),
144
    "ModernBertModel": ("modernbert", "ModernBertModel"),
145
    "NomicBertModel": ("bert_with_rope", "NomicBertModel"),
146
    "Phi3ForCausalLM": ("phi3", "Phi3ForCausalLM"),
147
    "Qwen2Model": ("qwen2", "Qwen2ForCausalLM"),
148
    "Qwen2ForCausalLM": ("qwen2", "Qwen2ForCausalLM"),
149
    "Qwen2ForRewardModel": ("qwen2_rm", "Qwen2ForRewardModel"),
150
    "Qwen2ForProcessRewardModel": ("qwen2_rm", "Qwen2ForProcessRewardModel"),
151
152
    "RobertaForMaskedLM": ("roberta", "RobertaEmbeddingModel"),
    "RobertaModel": ("roberta", "RobertaEmbeddingModel"),
153
    "TeleChat2ForCausalLM": ("telechat2", "TeleChat2ForCausalLM"),
154
    "XLMRobertaModel": ("roberta", "RobertaEmbeddingModel"),
155
    # [Multimodal]
Cyrus Leung's avatar
Cyrus Leung committed
156
    "LlavaNextForConditionalGeneration": ("llava_next", "LlavaNextForConditionalGeneration"),  # noqa: E501
157
    "Phi3VForCausalLM": ("phi3v", "Phi3VForCausalLM"),
158
    "Qwen2VLForConditionalGeneration": ("qwen2_vl", "Qwen2VLForConditionalGeneration"),  # noqa: E501
159
160
    # [Auto-converted (see adapters.py)]
    "Qwen2ForSequenceClassification": ("qwen2", "Qwen2ForCausalLM"),
161
162
163
164
    # Technically PrithviGeoSpatialMAE is a model that works on images, both in
    # input and output. I am adding it here because it piggy-backs on embedding
    # models for the time being.
    "PrithviGeoSpatialMAE": ("prithvi_geospatial_mae", "PrithviGeoSpatialMAE"),
165
166
}

167
168
169
170
171
172
_CROSS_ENCODER_MODELS = {
    "BertForSequenceClassification": ("bert", "BertForSequenceClassification"),
    "RobertaForSequenceClassification": ("roberta",
                                         "RobertaForSequenceClassification"),
    "XLMRobertaForSequenceClassification": ("roberta",
                                            "RobertaForSequenceClassification"),
xsank's avatar
xsank committed
173
174
    "ModernBertForSequenceClassification": ("modernbert",
                                            "ModernBertForSequenceClassification"),
175
    "Qwen3ForSequenceClassification": ("qwen3", "Qwen3ForSequenceClassification"), # noqa: E501
176
177
}

178
_MULTIMODAL_MODELS = {
179
    # [Decoder-only]
180
    "AriaForConditionalGeneration": ("aria", "AriaForConditionalGeneration"),
Jennifer Zhao's avatar
Jennifer Zhao committed
181
    "AyaVisionForConditionalGeneration": ("aya_vision", "AyaVisionForConditionalGeneration"),  # noqa: E501
182
183
    "Blip2ForConditionalGeneration": ("blip2", "Blip2ForConditionalGeneration"),
    "ChameleonForConditionalGeneration": ("chameleon", "ChameleonForConditionalGeneration"),  # noqa: E501
184
    "DeepseekVLV2ForCausalLM": ("deepseek_vl2", "DeepseekVLV2ForCausalLM"),
185
    "FuyuForCausalLM": ("fuyu", "FuyuForCausalLM"),
186
    "Gemma3ForConditionalGeneration": ("gemma3_mm", "Gemma3ForConditionalGeneration"),  # noqa: E501
187
    "GLM4VForCausalLM": ("glm4v", "GLM4VForCausalLM"),
188
    "GraniteSpeechForConditionalGeneration": ("granite_speech", "GraniteSpeechForConditionalGeneration"),  # noqa: E501
189
    "H2OVLChatModel": ("h2ovl", "H2OVLChatModel"),
190
    "InternVLChatModel": ("internvl", "InternVLChatModel"),
191
    "Idefics3ForConditionalGeneration":("idefics3","Idefics3ForConditionalGeneration"),
192
    "SmolVLMForConditionalGeneration": ("smolvlm","SmolVLMForConditionalGeneration"),  # noqa: E501
193
    "KimiVLForConditionalGeneration": ("kimi_vl", "KimiVLForConditionalGeneration"),  # noqa: E501
194
195
196
197
    "LlavaForConditionalGeneration": ("llava", "LlavaForConditionalGeneration"),
    "LlavaNextForConditionalGeneration": ("llava_next", "LlavaNextForConditionalGeneration"),  # noqa: E501
    "LlavaNextVideoForConditionalGeneration": ("llava_next_video", "LlavaNextVideoForConditionalGeneration"),  # noqa: E501
    "LlavaOnevisionForConditionalGeneration": ("llava_onevision", "LlavaOnevisionForConditionalGeneration"),  # noqa: E501
198
    "MantisForConditionalGeneration": ("llava", "MantisForConditionalGeneration"),  # noqa: E501
199
    "MiniMaxVL01ForConditionalGeneration": ("minimax_vl_01", "MiniMaxVL01ForConditionalGeneration"),  # noqa: E501
200
    "MiniCPMO": ("minicpmo", "MiniCPMO"),
201
    "MiniCPMV": ("minicpmv", "MiniCPMV"),
202
    "Mistral3ForConditionalGeneration": ("mistral3", "Mistral3ForConditionalGeneration"),  # noqa: E501
203
    "MolmoForCausalLM": ("molmo", "MolmoForCausalLM"),
204
    "NVLM_D": ("nvlm_d", "NVLM_D_Model"),
205
    "Ovis": ("ovis", "Ovis"),
206
    "PaliGemmaForConditionalGeneration": ("paligemma", "PaliGemmaForConditionalGeneration"),  # noqa: E501
207
    "Phi3VForCausalLM": ("phi3v", "Phi3VForCausalLM"),
208
    "PixtralForConditionalGeneration": ("pixtral", "PixtralForConditionalGeneration"),  # noqa: E501
209
    "QwenVLForConditionalGeneration": ("qwen_vl", "QwenVLForConditionalGeneration"),  # noqa: E501
210
    "Qwen2VLForConditionalGeneration": ("qwen2_vl", "Qwen2VLForConditionalGeneration"),  # noqa: E501
Roger Wang's avatar
Roger Wang committed
211
    "Qwen2_5_VLForConditionalGeneration": ("qwen2_5_vl", "Qwen2_5_VLForConditionalGeneration"),  # noqa: E501
212
    "Qwen2AudioForConditionalGeneration": ("qwen2_audio", "Qwen2AudioForConditionalGeneration"),  # noqa: E501
213
    "Qwen2_5OmniModel": ("qwen2_5_omni_thinker", "Qwen2_5OmniThinkerForConditionalGeneration"),  # noqa: E501
214
    "Qwen2_5OmniForConditionalGeneration": ("qwen2_5_omni_thinker", "Qwen2_5OmniThinkerForConditionalGeneration"),  # noqa: E501
215
    "UltravoxModel": ("ultravox", "UltravoxModel"),
216
    "Phi4MMForCausalLM": ("phi4mm", "Phi4MMForCausalLM"),
汪志鹏's avatar
汪志鹏 committed
217
    "TarsierForConditionalGeneration": ("tarsier", "TarsierForConditionalGeneration"),  # noqa: E501
218
    # [Encoder-decoder]
219
    "Florence2ForConditionalGeneration": ("florence2", "Florence2ForConditionalGeneration"),  # noqa: E501
220
    "MllamaForConditionalGeneration": ("mllama", "MllamaForConditionalGeneration"),  # noqa: E501
221
    "Llama4ForConditionalGeneration": ("mllama4", "Llama4ForConditionalGeneration"),  # noqa: E501
222
    "SkyworkR1VChatModel": ("skyworkr1v", "SkyworkR1VChatModel"),
223
    "WhisperForConditionalGeneration": ("whisper", "WhisperForConditionalGeneration"),  # noqa: E501
224
}
225
226

_SPECULATIVE_DECODING_MODELS = {
227
    "MiMoMTPModel": ("mimo_mtp", "MiMoMTP"),
228
    "EAGLEModel": ("eagle", "EAGLE"),
229
    "EagleLlamaForCausalLM": ("llama_eagle", "EagleLlamaForCausalLM"),
230
    "EagleMiniCPMForCausalLM": ("minicpm_eagle", "EagleMiniCPMForCausalLM"),
231
    "Eagle3LlamaForCausalLM": ("llama_eagle3", "Eagle3LlamaForCausalLM"),
232
    "DeepSeekMTPModel": ("deepseek_mtp", "DeepSeekMTP"),
233
234
    "MedusaModel": ("medusa", "Medusa"),
    "MLPSpeculatorPreTrainedModel": ("mlp_speculator", "MLPSpeculator"),
235
}
236

237
_TRANSFORMERS_MODELS = {
238
    "TransformersForCausalLM": ("transformers", "TransformersForCausalLM"),
239
}
240
# yapf: enable
241

242
_VLLM_MODELS = {
243
    **_TEXT_GENERATION_MODELS,
244
    **_EMBEDDING_MODELS,
245
    **_CROSS_ENCODER_MODELS,
246
    **_MULTIMODAL_MODELS,
247
    **_SPECULATIVE_DECODING_MODELS,
248
    **_TRANSFORMERS_MODELS,
249
250
}

251
252
253
254
255
256
257
258
# This variable is used as the args for subprocess.run(). We
# can modify  this variable to alter the args if needed. e.g.
# when we use par format to pack things together, sys.executable
# might not be the target we want to run.
_SUBPROCESS_COMMAND = [
    sys.executable, "-m", "vllm.model_executor.models.registry"
]

259

260
261
@dataclass(frozen=True)
class _ModelInfo:
262
    architecture: str
263
    is_text_generation_model: bool
264
    is_pooling_model: bool
265
    supports_cross_encoding: bool
266
267
    supports_multimodal: bool
    supports_pp: bool
268
269
    has_inner_state: bool
    is_attention_free: bool
270
    is_hybrid: bool
271
    has_noops: bool
272
    supports_transcription: bool
273
    supports_v0_only: bool
274
275

    @staticmethod
276
    def from_model_cls(model: type[nn.Module]) -> "_ModelInfo":
277
        return _ModelInfo(
278
            architecture=model.__name__,
279
            is_text_generation_model=is_text_generation_model(model),
280
            is_pooling_model=True,  # Can convert any model into a pooling model
281
            supports_cross_encoding=supports_cross_encoding(model),
282
283
            supports_multimodal=supports_multimodal(model),
            supports_pp=supports_pp(model),
284
285
            has_inner_state=has_inner_state(model),
            is_attention_free=is_attention_free(model),
286
            is_hybrid=is_hybrid(model),
287
288
            supports_transcription=supports_transcription(model),
            supports_v0_only=supports_v0_only(model),
289
            has_noops=has_noops(model),
290
        )
291
292


293
class _BaseRegisteredModel(ABC):
294

295
296
297
    @abstractmethod
    def inspect_model_cls(self) -> _ModelInfo:
        raise NotImplementedError
298

299
    @abstractmethod
300
    def load_model_cls(self) -> type[nn.Module]:
301
        raise NotImplementedError
302
303


304
305
306
307
308
309
310
@dataclass(frozen=True)
class _RegisteredModel(_BaseRegisteredModel):
    """
    Represents a model that has already been imported in the main process.
    """

    interfaces: _ModelInfo
311
    model_cls: type[nn.Module]
312
313

    @staticmethod
314
    def from_model_cls(model_cls: type[nn.Module]):
315
316
317
318
319
320
321
322
        return _RegisteredModel(
            interfaces=_ModelInfo.from_model_cls(model_cls),
            model_cls=model_cls,
        )

    def inspect_model_cls(self) -> _ModelInfo:
        return self.interfaces

323
    def load_model_cls(self) -> type[nn.Module]:
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
        return self.model_cls


@dataclass(frozen=True)
class _LazyRegisteredModel(_BaseRegisteredModel):
    """
    Represents a model that has not been imported in the main process.
    """
    module_name: str
    class_name: str

    # Performed in another process to avoid initializing CUDA
    def inspect_model_cls(self) -> _ModelInfo:
        return _run_in_subprocess(
            lambda: _ModelInfo.from_model_cls(self.load_model_cls()))

340
    def load_model_cls(self) -> type[nn.Module]:
341
342
343
344
345
346
347
348
        mod = importlib.import_module(self.module_name)
        return getattr(mod, self.class_name)


@lru_cache(maxsize=128)
def _try_load_model_cls(
    model_arch: str,
    model: _BaseRegisteredModel,
349
) -> Optional[type[nn.Module]]:
350
    from vllm.platforms import current_platform
351
    current_platform.verify_model_arch(model_arch)
352
353
354
355
356
357
    try:
        return model.load_model_cls()
    except Exception:
        logger.exception("Error in loading model architecture '%s'",
                         model_arch)
        return None
358
359


360
361
362
363
364
365
366
367
368
369
370
@lru_cache(maxsize=128)
def _try_inspect_model_cls(
    model_arch: str,
    model: _BaseRegisteredModel,
) -> Optional[_ModelInfo]:
    try:
        return model.inspect_model_cls()
    except Exception:
        logger.exception("Error in inspecting model architecture '%s'",
                         model_arch)
        return None
371
372


373
374
375
@dataclass
class _ModelRegistry:
    # Keyed by model_arch
376
    models: dict[str, _BaseRegisteredModel] = field(default_factory=dict)
377

378
    def get_supported_archs(self) -> Set[str]:
379
        return self.models.keys()
380

381
382
383
    def register_model(
        self,
        model_arch: str,
384
        model_cls: Union[type[nn.Module], str],
385
    ) -> None:
386
387
388
        """
        Register an external model to be used in vLLM.

389
        `model_cls` can be either:
390

391
        - A [`torch.nn.Module`][] class directly referencing the model.
392
        - A string in the format `<module>:<class>` which can be used to
393
394
          lazily import the model. This is useful to avoid initializing CUDA
          when importing the model and thus the related error
395
          `RuntimeError: Cannot re-initialize CUDA in forked subprocess`.
396
        """
397
398
399
400
        if not isinstance(model_arch, str):
            msg = f"`model_arch` should be a string, not a {type(model_arch)}"
            raise TypeError(msg)

401
        if model_arch in self.models:
402
403
404
            logger.warning(
                "Model architecture %s is already registered, and will be "
                "overwritten by the new model class %s.", model_arch,
405
406
407
408
409
410
411
                model_cls)

        if isinstance(model_cls, str):
            split_str = model_cls.split(":")
            if len(split_str) != 2:
                msg = "Expected a string in the format `<module>:<class>`"
                raise ValueError(msg)
412

413
            model = _LazyRegisteredModel(*split_str)
414
        elif isinstance(model_cls, type) and issubclass(model_cls, nn.Module):
415
            model = _RegisteredModel.from_model_cls(model_cls)
416
417
418
419
        else:
            msg = ("`model_cls` should be a string or PyTorch model class, "
                   f"not a {type(model_arch)}")
            raise TypeError(msg)
420

421
        self.models[model_arch] = model
422

423
    def _raise_for_unsupported(self, architectures: list[str]):
424
        all_supported_archs = self.get_supported_archs()
425

426
427
428
429
430
        if any(arch in all_supported_archs for arch in architectures):
            raise ValueError(
                f"Model architectures {architectures} failed "
                "to be inspected. Please check the logs for more details.")

431
432
433
        raise ValueError(
            f"Model architectures {architectures} are not supported for now. "
            f"Supported architectures: {all_supported_archs}")
434

435
    def _try_load_model_cls(self,
436
                            model_arch: str) -> Optional[type[nn.Module]]:
437
438
        if model_arch not in self.models:
            return None
439

440
        return _try_load_model_cls(model_arch, self.models[model_arch])
441

442
443
444
    def _try_inspect_model_cls(self, model_arch: str) -> Optional[_ModelInfo]:
        if model_arch not in self.models:
            return None
445

446
        return _try_inspect_model_cls(model_arch, self.models[model_arch])
447

448
449
    def _normalize_archs(
        self,
450
451
        architectures: Union[str, list[str]],
    ) -> list[str]:
452
453
454
455
456
        if isinstance(architectures, str):
            architectures = [architectures]
        if not architectures:
            logger.warning("No model architectures are specified")

457
458
459
460
        # filter out support architectures
        normalized_arch = list(
            filter(lambda model: model in self.models, architectures))

461
        # make sure Transformers backend is put at the last as a fallback
462
        if len(normalized_arch) != len(architectures):
463
            normalized_arch.append("TransformersForCausalLM")
464
        return normalized_arch
465

466
467
    def inspect_model_cls(
        self,
468
469
        architectures: Union[str, list[str]],
    ) -> tuple[_ModelInfo, str]:
470
        architectures = self._normalize_archs(architectures)
471

472
473
474
        for arch in architectures:
            model_info = self._try_inspect_model_cls(arch)
            if model_info is not None:
475
                return (model_info, arch)
476

477
        return self._raise_for_unsupported(architectures)
478

479
480
    def resolve_model_cls(
        self,
481
482
        architectures: Union[str, list[str]],
    ) -> tuple[type[nn.Module], str]:
483
        architectures = self._normalize_archs(architectures)
484

485
486
487
488
        for arch in architectures:
            model_cls = self._try_load_model_cls(arch)
            if model_cls is not None:
                return (model_cls, arch)
489

490
        return self._raise_for_unsupported(architectures)
491

492
493
    def is_text_generation_model(
        self,
494
        architectures: Union[str, list[str]],
495
    ) -> bool:
496
497
        model_cls, _ = self.inspect_model_cls(architectures)
        return model_cls.is_text_generation_model
498

499
    def is_pooling_model(
500
        self,
501
        architectures: Union[str, list[str]],
502
    ) -> bool:
503
        model_cls, _ = self.inspect_model_cls(architectures)
504
        return model_cls.is_pooling_model
505

506
507
    def is_cross_encoder_model(
        self,
508
        architectures: Union[str, list[str]],
509
    ) -> bool:
510
511
        model_cls, _ = self.inspect_model_cls(architectures)
        return model_cls.supports_cross_encoding
512

513
514
    def is_multimodal_model(
        self,
515
        architectures: Union[str, list[str]],
516
    ) -> bool:
517
518
        model_cls, _ = self.inspect_model_cls(architectures)
        return model_cls.supports_multimodal
519
520
521

    def is_pp_supported_model(
        self,
522
        architectures: Union[str, list[str]],
523
    ) -> bool:
524
525
        model_cls, _ = self.inspect_model_cls(architectures)
        return model_cls.supports_pp
526

527
528
    def model_has_inner_state(
        self,
529
        architectures: Union[str, list[str]],
530
531
532
    ) -> bool:
        model_cls, _ = self.inspect_model_cls(architectures)
        return model_cls.has_inner_state
533

534
535
    def is_attention_free_model(
        self,
536
        architectures: Union[str, list[str]],
537
538
539
    ) -> bool:
        model_cls, _ = self.inspect_model_cls(architectures)
        return model_cls.is_attention_free
540

541
542
    def is_hybrid_model(
        self,
543
        architectures: Union[str, list[str]],
544
545
546
547
    ) -> bool:
        model_cls, _ = self.inspect_model_cls(architectures)
        return model_cls.is_hybrid

548
549
    def is_noops_model(
        self,
550
        architectures: Union[str, list[str]],
551
552
553
554
    ) -> bool:
        model_cls, _ = self.inspect_model_cls(architectures)
        return model_cls.has_noops

555
556
    def is_transcription_model(
        self,
557
        architectures: Union[str, list[str]],
558
559
560
561
    ) -> bool:
        model_cls, _ = self.inspect_model_cls(architectures)
        return model_cls.supports_transcription

562
563
    def is_v1_compatible(
        self,
564
        architectures: Union[str, list[str]],
565
566
567
568
    ) -> bool:
        model_cls, _ = self.inspect_model_cls(architectures)
        return not model_cls.supports_v0_only

569
570

ModelRegistry = _ModelRegistry({
571
572
    model_arch:
    _LazyRegisteredModel(
573
574
575
576
577
578
579
580
581
582
        module_name=f"vllm.model_executor.models.{mod_relname}",
        class_name=cls_name,
    )
    for model_arch, (mod_relname, cls_name) in _VLLM_MODELS.items()
})

_T = TypeVar("_T")


def _run_in_subprocess(fn: Callable[[], _T]) -> _T:
583
584
585
586
587
    # NOTE: We use a temporary directory instead of a temporary file to avoid
    # issues like https://stackoverflow.com/questions/23212435/permission-denied-to-write-to-my-temporary-file
    with tempfile.TemporaryDirectory() as tempdir:
        output_filepath = os.path.join(tempdir, "registry_output.tmp")

588
        # `cloudpickle` allows pickling lambda functions directly
589
        input_bytes = cloudpickle.dumps((fn, output_filepath))
590
591
592

        # cannot use `sys.executable __file__` here because the script
        # contains relative imports
593
594
595
        returned = subprocess.run(_SUBPROCESS_COMMAND,
                                  input=input_bytes,
                                  capture_output=True)
596
597
598
599
600
601
602
603
604

        # check if the subprocess is successful
        try:
            returned.check_returncode()
        except Exception as e:
            # wrap raised exception to provide more information
            raise RuntimeError(f"Error raised in subprocess:\n"
                               f"{returned.stderr.decode()}") from e

605
        with open(output_filepath, "rb") as f:
606
607
608
609
610
611
612
613
614
615
616
            return pickle.load(f)


def _run() -> None:
    # Setup plugins
    from vllm.plugins import load_general_plugins
    load_general_plugins()

    fn, output_file = pickle.loads(sys.stdin.buffer.read())

    result = fn()
617
618
619

    with open(output_file, "wb") as f:
        f.write(pickle.dumps(result))
620
621
622


if __name__ == "__main__":
623
    _run()